diff --git a/v12x/pylib/scipy-1.18.1.dist-info/sboms/auditwheel.cdx.json b/v12x/pylib/scipy-1.18.1.dist-info/sboms/auditwheel.cdx.json new file mode 100644 index 0000000000000000000000000000000000000000..a1b46c5df859229add41ab5bbff4828c64d017be --- /dev/null +++ b/v12x/pylib/scipy-1.18.1.dist-info/sboms/auditwheel.cdx.json @@ -0,0 +1 @@ +{"bomFormat": "CycloneDX", "specVersion": "1.4", "version": 1, "metadata": {"component": {"type": "library", "bom-ref": "pkg:pypi/scipy@1.18.1?file_name=scipy-1.18.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", "name": "scipy", "version": "1.18.1", "purl": "pkg:pypi/scipy@1.18.1?file_name=scipy-1.18.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl"}, "tools": [{"name": "auditwheel", "version": "6.7.0"}]}, "components": [{"type": "library", "bom-ref": "pkg:pypi/scipy@1.18.1?file_name=scipy-1.18.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", "name": "scipy", "version": "1.18.1", "purl": "pkg:pypi/scipy@1.18.1?file_name=scipy-1.18.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl"}, {"type": "library", "bom-ref": "pkg:rpm/almalinux/libgfortran@8.5.0-28.el8_10.alma.1#3362720ce454cc402813e3cfe54f0eb4efaedcd690fa35129ddc4b3de40e0650", "name": "libgfortran", "version": "8.5.0-28.el8_10.alma.1", "purl": "pkg:rpm/almalinux/libgfortran@8.5.0-28.el8_10.alma.1"}, {"type": "library", "bom-ref": "pkg:rpm/almalinux/libquadmath@8.5.0-28.el8_10.alma.1#97b91febe2a1c7927c6482d6bd6567b83be3467f33c9d2764df1283f7ac9a5fa", "name": "libquadmath", "version": "8.5.0-28.el8_10.alma.1", "purl": "pkg:rpm/almalinux/libquadmath@8.5.0-28.el8_10.alma.1"}], "dependencies": [{"ref": "pkg:pypi/scipy@1.18.1?file_name=scipy-1.18.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", "dependsOn": ["pkg:rpm/almalinux/libgfortran@8.5.0-28.el8_10.alma.1#3362720ce454cc402813e3cfe54f0eb4efaedcd690fa35129ddc4b3de40e0650", "pkg:rpm/almalinux/libquadmath@8.5.0-28.el8_10.alma.1#97b91febe2a1c7927c6482d6bd6567b83be3467f33c9d2764df1283f7ac9a5fa"]}, {"ref": "pkg:rpm/almalinux/libgfortran@8.5.0-28.el8_10.alma.1#3362720ce454cc402813e3cfe54f0eb4efaedcd690fa35129ddc4b3de40e0650"}, {"ref": "pkg:rpm/almalinux/libquadmath@8.5.0-28.el8_10.alma.1#97b91febe2a1c7927c6482d6bd6567b83be3467f33c9d2764df1283f7ac9a5fa"}]} \ No newline at end of file diff --git a/v12x/pylib/scipy/spatial/transform/__init__.py b/v12x/pylib/scipy/spatial/transform/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..1490a18e2ef0b0435fdada41f6f083aa89a3d859 --- /dev/null +++ b/v12x/pylib/scipy/spatial/transform/__init__.py @@ -0,0 +1,31 @@ +""" +Spatial Transformations (:mod:`scipy.spatial.transform`) +======================================================== + +.. currentmodule:: scipy.spatial.transform + +This package implements various spatial transformations. For now, rotations +and rigid transforms (rotations and translations) are supported. + +Rotations in 3 dimensions +------------------------- +.. autosummary:: + :toctree: generated/ + + RigidTransform + Rotation + Slerp + RotationSpline +""" +from ._rotation import Rotation, Slerp +from ._rigid_transform import RigidTransform +from ._rotation_spline import RotationSpline + +# Deprecated namespaces, to be removed in v2.0.0 +from . import rotation + +__all__ = ['Rotation', 'Slerp', 'RotationSpline', 'RigidTransform'] + +from scipy._lib._testutils import PytestTester +test = PytestTester(__name__) +del PytestTester diff --git a/v12x/pylib/scipy/spatial/transform/_rigid_transform.py b/v12x/pylib/scipy/spatial/transform/_rigid_transform.py new file mode 100644 index 0000000000000000000000000000000000000000..c6af813b622372c2fd148f186820354878c393d9 --- /dev/null +++ b/v12x/pylib/scipy/spatial/transform/_rigid_transform.py @@ -0,0 +1,1975 @@ +from __future__ import annotations + +from collections.abc import Iterable, Iterator +from types import EllipsisType, GenericAlias, ModuleType, NotImplementedType +from collections.abc import Callable + +import numpy as np + +from scipy._lib._array_api import ( + array_namespace, + is_numpy, + ArrayLike, + Array, + xp_capabilities, +) +from scipy.spatial.transform import Rotation +from scipy.spatial.transform._rotation import _promote +import scipy.spatial.transform._rigid_transform_cy as cython_backend +import scipy.spatial.transform._rigid_transform_xp as xp_backend +import scipy._external.array_api_extra as xpx +from scipy._external.array_api_compat import device +from scipy._lib._array_api import xp_promote +from scipy._lib._util import broadcastable + + +__all__ = ["RigidTransform"] + +backend_registry = {array_namespace(np.empty(0)): cython_backend} + + +def select_backend(xp: ModuleType, cython_compatible: bool): + """Select the backend for the given array library. + + We need this selection function because the Cython backend for numpy does not + support quaternions of arbitrary dimensions. We therefore only use the Array API + backend for numpy if we are dealing with rotations of more than one leading + dimension. + """ + if is_numpy(xp) and not cython_compatible: + return xp_backend + return backend_registry.get(xp, xp_backend) + + +@xp_capabilities() +def normalize_dual_quaternion(dual_quat: ArrayLike) -> Array: + """Normalize dual quaternion.""" + xp = array_namespace(dual_quat) + dual_quat = _promote(dual_quat, xp=xp) + single = dual_quat.ndim == 1 and is_numpy(xp) + if single: + dual_quat = xpx.atleast_nd(dual_quat, ndim=2, xp=xp) + cython_compatible = dual_quat.ndim < 3 + dq = select_backend(xp, cython_compatible).normalize_dual_quaternion(dual_quat) + if single: + return dq[0] + return dq + + +rigid_transform_extra_note = ( + """The method ``apply`` is not supported with cupy<14.*. + + """) + + +@xp_capabilities( + skip_backends=[("dask.array", "missing linalg.cross/det functions")], + method_capabilities={ + "__init__": dict( + jax_jit=False, + skip_backends=[("dask.array", "missing linalg.cross/det functions")], + ), + "from_matrix": dict( + jax_jit=False, + skip_backends=[("dask.array", "missing linalg.cross/det functions")], + ), + "__getitem__": dict( + jax_jit=False, + skip_backends=[ + ("dask.array", "cannot handle zero-length rigid transforms") + ], + ), + "__setitem__": dict( + jax_jit=False, + skip_backends=[ + ("dask.array", "cannot handle zero-length rigid transforms") + ], + ), + "apply": dict( + skip_backends=[ + ("dask.array", "missing linalg.cross/det functions"), + ("cupy", "missing .mT attribute in cupy<14.*"), + ], + ), + }, + extra_note=rigid_transform_extra_note, +) +class RigidTransform: + """ + Rigid transform in 3 dimensions. + + This class provides an interface to initialize from and represent rigid + transforms (rotation and translation) in 3D space. In different fields, + this type of transform may be referred to as "*pose*" (especially in + robotics), "*extrinsic parameters*", or the "*model matrix*" (especially in + computer graphics), but the core concept is the same: a rotation and + translation describing the orientation of one 3D coordinate frame relative + to another. Mathematically, these transforms belong to the Special + Euclidean group SE(3), which encodes rotation (SO(3)) plus translation. + + The following operations on rigid transforms are supported: + + - Application on vectors + - Transformation composition + - Transformation inversion + - Transformation indexing + + Note that coordinate systems must be right-handed. Because of this, this + class more precisely represents *proper* rigid transformations in SE(3) + rather than rigid transforms in E(3) more generally [1]_. + + Indexing within a transform is supported since multiple transforms can be + stored within a single `RigidTransform` instance. + + Multiple transforms can be stored in a single `RigidTransform` object, which can be + initialized using N-dimensional arrays and supports broadcasting for all + operations. + + To create `RigidTransform` objects use ``from_...`` methods (see examples + below). ``RigidTransform(...)`` is not supposed to be instantiated directly. + + For rigorous introductions to rigid transforms, see [2]_, [3]_, and [4]_. + + Parameters + ---------- + matrix : array_like, shape (..., 4, 4) + A single transformation matrix or a stack of transformation + matrices. + normalize : bool, optional + If True, orthonormalize the rotation matrix using singular value + decomposition. If False, the rotation matrix is not checked for + orthogonality or right-handedness. + copy : bool, optional + If True, copy the input matrix. If False, a reference to the input + matrix is used. If normalize is True, the input matrix is always + copied regardless of the value of copy. + + Attributes + ---------- + single + rotation + translation + + Methods + ------- + __len__ + __getitem__ + __mul__ + __pow__ + from_matrix + from_rotation + from_translation + from_components + from_exp_coords + from_dual_quat + as_matrix + as_components + as_exp_coords + as_dual_quat + concatenate + mean + apply + inv + identity + + Notes + ----- + .. versionadded:: 1.16.0 + + References + ---------- + .. [1] https://en.wikipedia.org/wiki/Rigid_transformation + .. [2] https://motion.cs.illinois.edu/RoboticSystems/CoordinateTransformations.html + .. [3] https://www.brainvoyager.com/bv/doc/UsersGuide/CoordsAndTransforms/SpatialTransformationMatrices.html + .. [4] Kevin M. Lynch and Frank C. Park, "Modern Robotics: Mechanics, + Planning, and Control" Chapter 3.3, 2017, Cambridge University Press. + https://hades.mech.northwestern.edu/images/2/25/MR-v2.pdf#page=107.31 + .. [5] Paul Furgale, "Representing Robot Pose: The good, the bad, and the + ugly", June 9, 2014. + https://rpg.ifi.uzh.ch/docs/teaching/2024/FurgaleTutorial.pdf + + Examples + -------- + A `RigidTransform` instance can be initialized in any of the above formats + and converted to any of the others. The underlying object is independent of + the representation used for initialization. + + **Notation Conventions and Composition** + + The notation here largely follows the convention defined in [5]_. When we + name transforms, we read the subscripts from right to left. So ``tf_A_B`` + represents a transform A <- B and can be interpreted as: + + - the coordinates and orientation of B relative to A + - the transformation of points from B to A + - the pose of B described in A's coordinate system + + .. parsed-literal:: + :class: highlight-none + + tf_A_B + ^ ^ + | | + | --- from B + | + ----- to A + + When composing transforms, the order is important. Transforms are not + commutative, so in general ``tf_A_B * tf_B_C`` is not the same as + ``tf_B_C * tf_A_B``. Transforms are composed and applied to vectors + right-to-left. So ``(tf_A_B * tf_B_C).apply(p_C)`` is the same as + ``tf_A_B.apply(tf_B_C.apply(p_C))``. + + When composed, transforms should be ordered such that the multiplication + operator is surrounded by a single frame, so the frame "cancels out" and + the outside frames are left. In the example below, B cancels out and the + outside frames A and C are left. Or to put it another way, A <- C is the + same as A <- B <- C. + + .. parsed-literal:: + :class: highlight-none + + ----------- B cancels out + | | + v v + tf_A_C = tf_A_B * tf_B_C + ^ ^ + | | + ------------ to A, from C are left + + When we notate vectors, we write the subscript of the frame that the vector + is defined in. So ``p_B`` means the point ``p`` defined in frame B. To + transform this point from frame B to coordinates in frame A, we apply the + transform ``tf_A_B`` to the vector, lining things up such that the notated + B frames are next to each other and "cancel out". + + .. parsed-literal:: + :class: highlight-none + + ------------ B cancels out + | | + v v + p_A = tf_A_B.apply(p_B) + ^ + | + -------------- A is left + + **Visualization** + + >>> from scipy.spatial.transform import RigidTransform as Tf + >>> from scipy.spatial.transform import Rotation as R + >>> import numpy as np + + The following function can be used to plot transforms with Matplotlib + by showing how they transform the standard x, y, z coordinate axes: + + >>> import matplotlib.pyplot as plt + >>> colors = ("#FF6666", "#005533", "#1199EE") # Colorblind-safe RGB + >>> def plot_transformed_axes(ax, tf, name=None, scale=1): + ... r = tf.rotation + ... t = tf.translation + ... loc = np.array([t, t]) + ... for i, (axis, c) in enumerate(zip((ax.xaxis, ax.yaxis, ax.zaxis), + ... colors)): + ... axlabel = axis.axis_name + ... axis.set_label_text(axlabel) + ... axis.label.set_color(c) + ... axis.line.set_color(c) + ... axis.set_tick_params(colors=c) + ... line = np.zeros((2, 3)) + ... line[1, i] = scale + ... line_rot = r.apply(line) + ... line_plot = line_rot + loc + ... ax.plot(line_plot[:, 0], line_plot[:, 1], line_plot[:, 2], c) + ... text_loc = line[1]*1.2 + ... text_loc_rot = r.apply(text_loc) + ... text_plot = text_loc_rot + t + ... ax.text(*text_plot, axlabel.upper(), color=c, + ... va="center", ha="center") + ... ax.text(*tf.translation, name, color="k", va="center", ha="center", + ... bbox={"fc": "w", "alpha": 0.8, "boxstyle": "circle"}) + + **Defining Frames** + + Let's work through an example. + + First, define the "world frame" A, also called the "base frame". + All frames are the identity transform from their own perspective. + + >>> tf_A = Tf.identity() + + We will visualize a new frame B in A's coordinate system. So we need to + define the transform that converts coordinates from frame B to frame A + (A <- B). + + Physically, let's imagine constructing B from A by: + + 1) Rotating A by +90 degrees around its x-axis. + 2) Translating the rotated frame 2 units in A's -x direction. + + From A's perspective, B is at [-2, 0, 0] and rotated +90 degrees about the + x-axis, which is exactly the transform A <- B. + + >>> t_A_B = np.array([-2, 0, 0]) + >>> r_A_B = R.from_euler('xyz', [90, 0, 0], degrees=True) + >>> tf_A_B = Tf.from_components(t_A_B, r_A_B) + + Let's plot these frames. + + >>> fig, ax = plt.subplots(subplot_kw={"projection": "3d"}) + >>> plot_transformed_axes(ax, tf_A, name="tfA") # A plotted in A + >>> plot_transformed_axes(ax, tf_A_B, name="tfAB") # B plotted in A + >>> ax.set_title("A, B frames with respect to A") + >>> ax.set_aspect("equal") + >>> ax.figure.set_size_inches(6, 5) + >>> plt.show() + + Now let's visualize a new frame C in B's coordinate system. + Let's imagine constructing C from B by: + + 1) Translating B by 2 units in its +z direction. + 2) Rotating B by +30 degrees around its z-axis. + + >>> t_B_C = np.array([0, 0, 2]) + >>> r_B_C = R.from_euler('xyz', [0, 0, 30], degrees=True) + >>> tf_B_C = Tf.from_components(t_B_C, r_B_C) + + In order to plot these frames from a consistent perspective, we need to + calculate the transform between A and C. Note that we do not make this + transform directly, but instead compose intermediate transforms that let us + get from C to A: + + >>> tf_A_C = tf_A_B * tf_B_C # A <- B <- C + + Now we can plot these three frames from A's perspective. + + >>> fig, ax = plt.subplots(subplot_kw={"projection": "3d"}) + >>> plot_transformed_axes(ax, tf_A, name="tfA") # A plotted in A + >>> plot_transformed_axes(ax, tf_A_B, name="tfAB") # B plotted in A + >>> plot_transformed_axes(ax, tf_A_C, name="tfAC") # C plotted in A + >>> ax.set_title("A, B, C frames with respect to A") + >>> ax.set_aspect("equal") + >>> ax.figure.set_size_inches(6, 5) + >>> plt.show() + + **Transforming Vectors** + + Let's transform a vector from A, to B and C. To do this, we will first + invert the transforms we already have from B and C, to A. + + >>> tf_B_A = tf_A_B.inv() # B <- A + >>> tf_C_A = tf_A_C.inv() # C <- A + + Now we can define a point in A and use the above transforms to get its + coordinates in B and C: + + >>> p1_A = np.array([1, 0, 0]) # +1 in x_A direction + >>> p1_B = tf_B_A.apply(p1_A) + >>> p1_C = tf_C_A.apply(p1_A) + >>> print(p1_A) # Original point 1 in A + [1 0 0] + >>> print(p1_B) # Point 1 in B + [3. 0. 0.] + >>> print(p1_C) # Point 1 in C + [ 2.59807621 -1.5 -2. ] + + We can also do the reverse. We define a point in C and transform it to A: + + >>> p2_C = np.array([0, 1, 0]) # +1 in y_C direction + >>> p2_A = tf_A_C.apply(p2_C) + >>> print(p2_C) # Original point 2 in C + [0 1 0] + >>> print(p2_A) # Point 2 in A + [-2.5 -2. 0.8660254] + + Plot the frames with respect to A again, but also plot these two points: + + >>> fig, ax = plt.subplots(subplot_kw={"projection": "3d"}) + >>> plot_transformed_axes(ax, tf_A, name="tfA") # A plotted in A + >>> plot_transformed_axes(ax, tf_A_B, name="tfAB") # B plotted in A + >>> plot_transformed_axes(ax, tf_A_C, name="tfAC") # C plotted in A + >>> ax.scatter(p1_A[0], p1_A[1], p1_A[2], color=colors[0]) # +1 x_A + >>> ax.scatter(p2_A[0], p2_A[1], p2_A[2], color=colors[1]) # +1 y_C + >>> ax.set_title("A, B, C frames and points with respect to A") + >>> ax.set_aspect("equal") + >>> ax.figure.set_size_inches(6, 5) + >>> plt.show() + + **Switching Base Frames** + + Up to this point, we have been visualizing frames from A's perspective. + Let's use the transforms we defined to visualize the frames from C's + perspective. + + Now C is the "base frame" or "world frame". All frames are the identity + transform from their own perspective. + + >>> tf_C = Tf.identity() + + We've already defined the transform C <- A, and can obtain C <- B by + inverting the existing transform B <- C. + + >>> tf_C_B = tf_B_C.inv() # C <- B + + This lets us plot everything from C's perspective: + + >>> fig, ax = plt.subplots(subplot_kw={"projection": "3d"}) + >>> plot_transformed_axes(ax, tf_C, name="tfC") # C plotted in C + >>> plot_transformed_axes(ax, tf_C_B, name="tfCB") # B plotted in C + >>> plot_transformed_axes(ax, tf_C_A, name="tfCA") # A plotted in C + >>> ax.scatter(p1_C[0], p1_C[1], p1_C[2], color=colors[0]) + >>> ax.scatter(p2_C[0], p2_C[1], p2_C[2], color=colors[1]) + >>> ax.set_title("A, B, C frames and points with respect to C") + >>> ax.set_aspect("equal") + >>> ax.figure.set_size_inches(6, 5) + >>> plt.show() + """ + + # generic type compatibility with scipy-stubs + __class_getitem__: classmethod = classmethod(GenericAlias) + + def __init__(self, matrix: ArrayLike, normalize: bool = True, copy: bool = True): + """Initialize from a 4x4 transformation matrix. + + Rotations are not meant to be initialized directly. Please use one of + the `from_...` methods instead. + + Parameters + ---------- + matrix : array_like, shape (..., 4, 4) + A single transformation matrix or a stack of transformation + matrices. + normalize : bool, optional + If True, orthonormalize the rotation matrix using singular value + decomposition. If False, the rotation matrix is not checked for + orthogonality or right-handedness. + copy : bool, optional + If True, copy the input matrix. If False, a reference to the input + matrix is used. If normalize is True, the input matrix is always + copied regardless of the value of copy. + + Returns + ------- + transform : `RigidTransform` instance + """ + xp = array_namespace(matrix) + self._xp = xp + matrix = _promote(matrix, xp=xp) + if matrix.shape[-2:] != (4, 4): + raise ValueError( + f"Expected `matrix` to have shape (..., 4, 4), got {matrix.shape}." + ) + # We only need the _single flag for the cython backend. The Array API backend + # uses broadcasting by default and hence returns the correct shape without + # additional _single logic + self._single = matrix.ndim == 2 and is_numpy(xp) + if self._single: + matrix = xpx.atleast_nd(matrix, ndim=3, xp=xp) + + self._backend = select_backend(xp, matrix.ndim < 4) + self._matrix = self._backend.from_matrix(matrix, normalize, copy) + + def __repr__(self): + m = f"{self.as_matrix()!r}".splitlines() + # bump indent (+27 characters) + m[1:] = [" " * 27 + m[i] for i in range(1, len(m))] + return "RigidTransform.from_matrix(" + "\n".join(m) + ")" + + @staticmethod + def from_matrix(matrix: ArrayLike) -> RigidTransform: + """Initialize from a 4x4 transformation matrix. + + Parameters + ---------- + matrix : array_like, shape (..., 4, 4) + Transformation matrices. Each matrix[..., :, :] represents a 4x4 + transformation matrix. + + Returns + ------- + transform : `RigidTransform` instance + Object containing the rigid transform represented by the input matrix. + + Notes + ----- + 4x4 rigid transformation matrices are of the form:: + + [ tx] + [ R ty] + [ tz] + [ 0 0 0 1] + + where ``R`` is a 3x3 rotation matrix and ``t = [tx, ty, tz]`` is a 3x1 + translation vector. As rotation matrices must be proper orthogonal, the + rotation component is orthonormalized using singular value decomposition + before initialization. + + Examples + -------- + >>> from scipy.spatial.transform import RigidTransform as Tf + >>> import numpy as np + + Creating a transform from a single matrix: + + >>> m = np.array([[0, 1, 0, 2], + ... [0, 0, 1, 3], + ... [1, 0, 0, 4], + ... [0, 0, 0, 1]]) + >>> tf = Tf.from_matrix(m) + >>> tf.as_matrix() + array([[0., 1., 0., 2.], + [0., 0., 1., 3.], + [1., 0., 0., 4.], + [0., 0., 0., 1.]]) + >>> tf.single + True + + Creating a transform from an N-dimensional array of matrices: + + >>> m = np.tile(np.eye(4), (2, 5, 1, 1)) # Shape (2, 5, 4, 4) + >>> tf = Tf.from_matrix(m) + >>> tf.shape + (2, 5) + >>> tf.single + False + >>> len(tf) + 2 + + Matrices with a rotation component that is not proper orthogonal are + orthogonalized using singular value decomposition before initialization: + + >>> tf = Tf.from_matrix(np.diag([2, 2, 2, 1])) + >>> tf.as_matrix() + array([[1., 0., 0., 0.], + [0., 1., 0., 0.], + [0., 0., 1., 0.], + [0., 0., 0., 1.]]) + """ + return RigidTransform(matrix, normalize=True, copy=True) + + @staticmethod + def from_rotation(rotation: Rotation) -> RigidTransform: + """Initialize from a rotation, without a translation. + + When applying this transform to a vector ``v``, the result is the + same as if the rotation was applied to the vector. + ``Tf.from_rotation(r).apply(v) == r.apply(v)`` + + Parameters + ---------- + rotation : `Rotation` instance + A single rotation or a rotation with N leading dimensions. + + Returns + ------- + transform : `RigidTransform` instance + Object containing the rigid transform represented by the input rotation. + + Examples + -------- + >>> from scipy.spatial.transform import RigidTransform as Tf + >>> from scipy.spatial.transform import Rotation as R + >>> import numpy as np + + Creating a transform from a single rotation: + + >>> r = R.from_euler("ZYX", [90, 30, 0], degrees=True) + >>> r.apply([1, 0, 0]) + array([0. , 0.8660254, -0.5 ]) + >>> tf = Tf.from_rotation(r) + >>> tf.apply([1, 0, 0]) + array([0. , 0.8660254, -0.5 ]) + >>> tf.single + True + + The upper 3x3 submatrix of the transformation matrix is the rotation + matrix: + + >>> np.allclose(tf.as_matrix()[:3, :3], r.as_matrix(), atol=1e-12) + True + + Creating multiple transforms from a rotation with N leading dimensions: + + >>> r = R.from_euler("ZYX", [[90, 30, 0], [45, 30, 60]], degrees=True) + >>> r.apply([1, 0, 0]) + array([[0. , 0.8660254 , -0.5 ], + [0.61237244, 0.61237244, -0.5 ]]) + >>> tf = Tf.from_rotation(r) + >>> tf.apply([1, 0, 0]) + array([[0. , 0.8660254 , -0.5 ], + [0.61237244, 0.61237244, -0.5 ]]) + >>> tf.single + False + >>> len(tf) + 2 + """ + if not isinstance(rotation, Rotation): + raise TypeError( + "Expected `rotation` to be a `Rotation` instance, " + f"got {type(rotation)}." + ) + quat = rotation.as_quat() + xp = array_namespace(quat) + backend = select_backend(xp, quat.ndim < 3) + matrix = backend.from_rotation(quat) + return RigidTransform._from_raw_matrix(matrix, xp, backend) + + @staticmethod + def from_translation(translation: ArrayLike) -> RigidTransform: + """Initialize from a translation numpy array, without a rotation. + + When applying this transform to a vector ``v``, the result is the same + as if the translation and vector were added together. If ``t`` is the + displacement vector of the translation, then: + + ``Tf.from_translation(t).apply(v) == t + v`` + + Parameters + ---------- + translation : array_like, shape (..., 3) + Translation vectors. Each translation[..., :] represents a 3D + translation vector. + + Returns + ------- + transform : `RigidTransform` instance + Object containing the rigid transform represented by the input translation. + + Examples + -------- + >>> from scipy.spatial.transform import RigidTransform as Tf + >>> import numpy as np + + Creating a transform from a single translation vector: + + >>> t = np.array([2, 3, 4]) + >>> t + np.array([1, 0, 0]) + array([3, 3, 4]) + >>> tf = Tf.from_translation(t) + >>> tf.apply([1, 0, 0]) + array([3., 3., 4.]) + >>> tf.single + True + + The top 3x1 points in the rightmost column of the transformation matrix + is the translation vector: + + >>> tf.as_matrix() + array([[1., 0., 0., 2.], + [0., 1., 0., 3.], + [0., 0., 1., 4.], + [0., 0., 0., 1.]]) + >>> np.allclose(tf.as_matrix()[:3, 3], t) + True + + Creating multiple transforms from an N-dimensional array of translation + vectors: + + >>> t = np.array([[2, 3, 4], [1, 0, 0]]) + >>> t + np.array([1, 0, 0]) + array([[3, 3, 4], + [2, 0, 0]]) + >>> tf = Tf.from_translation(t) + >>> tf.apply([1, 0, 0]) + array([[3., 3., 4.], + [2., 0., 0.]]) + >>> np.allclose(tf.as_matrix()[:, :3, 3], t) + True + >>> tf.single + False + >>> len(tf) + 2 + """ + xp = array_namespace(translation) + translation = _promote(translation, xp=xp) + backend = select_backend(xp, translation.ndim < 3) + matrix = backend.from_translation(translation) + return RigidTransform._from_raw_matrix(matrix, xp, backend) + + @staticmethod + def from_components(translation: ArrayLike, rotation: Rotation) -> RigidTransform: + """Initialize a rigid transform from translation and rotation + components. + + When creating a rigid transform from a translation and rotation, the + translation is applied after the rotation, such that + ``tf = Tf.from_components(translation, rotation)`` + is equivalent to + ``tf = Tf.from_translation(translation) * Tf.from_rotation(rotation)``. + + When applying a transform to a vector ``v``, the result is the + same as if the transform was applied to the vector in the + following way: ``tf.apply(v) == translation + rotation.apply(v)`` + + Parameters + ---------- + translation : array_like, shape (..., 3) + Translation vectors. Each translation[..., :] represents a 3D + translation vector. + rotation : `Rotation` instance + Rotation objects. The shape must be broadcastable with the + translation shape. + + Returns + ------- + transform : `RigidTransform` instance + Rigid transform objects. The shape is determined by broadcasting + the translation and rotation shapes together. + + Examples + -------- + >>> from scipy.spatial.transform import RigidTransform as Tf + >>> from scipy.spatial.transform import Rotation as R + >>> import numpy as np + + Creating from a single rotation and translation: + + >>> t = np.array([2, 3, 4]) + >>> r = R.from_euler("ZYX", [90, 30, 0], degrees=True) + >>> r.as_matrix() + array([[ 0. , -1., 0. ], + [ 0.8660254, 0., 0.5 ], + [-0.5 , 0., 0.8660254 ]]) + >>> tf = Tf.from_components(t, r) + >>> tf.rotation.as_matrix() + array([[ 0. , -1., 0. ], + [ 0.8660254, 0., 0.5 ], + [-0.5 , 0., 0.8660254 ]]) + >>> tf.translation + array([2., 3., 4.]) + >>> tf.single + True + + When applying a transform to a vector ``v``, the result is the same as + if the transform was applied to the vector in the following way: + ``tf.apply(v) == translation + rotation.apply(v)`` + + >>> r.apply([1, 0, 0]) + array([0. , 0.8660254, -0.5 ]) + >>> t + r.apply([1, 0, 0]) + array([2. , 3.8660254, 3.5 ]) + >>> tf.apply([1, 0, 0]) + array([2. , 3.8660254, 3.5 ]) + """ + rotation_tf = RigidTransform.from_rotation(rotation) + return RigidTransform.from_translation(translation) * rotation_tf + + @staticmethod + def from_exp_coords(exp_coords: ArrayLike) -> RigidTransform: + r"""Initialize from exponential coordinates of transform. + + This implements the exponential map that converts 6-dimensional real + vectors to SE(3). + + An exponential coordinate vector consists of 6 elements + ``[rx, ry, rz, vx, vy, vz]``. The first 3 encode rotation (and form a + rotation vector used in `Rotation.from_rotvec`) and the last 3 encode + translation (and form a translation vector for pure translations). + The exponential mapping can be expressed as matrix exponential + ``T = exp(tau)``, where ``T`` is a 4x4 matrix representing a rigid + transform and ``tau`` is a 4x4 matrix formed from the elements of an + exponential coordinate vector:: + + tau = [ 0 -rz ry vx] + [ rz 0 -rx vy] + [-ry rx 0 vz] + [ 0 0 0 1] + + Parameters + ---------- + exp_coords : array_like, shape (..., 6) + Exponential coordinate vectors. Each exp_coords[..., :] represents + a 6D exponential coordinate vector with the expected order of + components ``[rx, ry, rz, vx, vy, vz]``. The first 3 components + encode rotation and the last 3 encode translation. + + Returns + ------- + transform : `RigidTransform` instance + Rigid transform objects with the same leading dimensions as the input. + + Examples + -------- + >>> from scipy.spatial.transform import RigidTransform as Tf + >>> import numpy as np + + Creating from a single 6d vector of exponential coordinates: + + >>> tf = Tf.from_exp_coords([ + ... -2.01041204, -0.52983629, 0.65773501, + ... 0.10386614, 0.05855009, 0.54959179]) + >>> tf.as_matrix() + array([[0.76406621, 0.10504613, -0.63652819, -0.10209961], + [0.59956454, -0.47987325, 0.64050295, 0.40158789], + [-0.2381705, -0.87102639, -0.42963687, 0.19637636], + [0., 0., 0., 1.]]) + >>> tf.single + True + + A vector of zeros represents the identity transform: + + >>> tf = Tf.from_exp_coords(np.zeros(6)) + >>> tf.as_matrix() + array([[1., 0., 0., 0.], + [0., 1., 0., 0.], + [0., 0., 1., 0.], + [0., 0., 0., 1.]]) + + The last three numbers encode translation. If the first three numbers + are zero, the last three components can be interpreted as the + translation: + + >>> tf_trans = Tf.from_exp_coords([0, 0, 0, 4.3, -2, 3.4]) + >>> tf_trans.translation + array([4.3, -2., 3.4]) + + The first three numbers encode rotation as a rotation vector: + + >>> tf_rot = Tf.from_exp_coords([0.5, 0.3, 0.1, 0, 0, 0]) + >>> tf_rot.rotation.as_rotvec() + array([0.5, 0.3, 0.1]) + + Combining translation and rotation preserves the rotation vector, + but changes the last three components as they encode translation and + rotation: + + >>> (tf_trans * tf_rot).as_exp_coords() + array([0.5, 0.3, 0.1, 3.64305882, -1.25879559, 4.46109265]) + """ + xp = array_namespace(exp_coords) + exp_coords = xp_promote(exp_coords, force_floating=True, xp=xp) + if exp_coords.shape[-1] != 6: + raise ValueError( + f"Expected `exp_coords` to have shape (..., 6), got {exp_coords.shape}." + ) + backend = select_backend(xp, cython_compatible=exp_coords.ndim < 3) + matrix = backend.from_exp_coords(exp_coords) + return RigidTransform._from_raw_matrix(matrix, xp, backend) + + @staticmethod + def from_dual_quat( + dual_quat: ArrayLike, *, scalar_first: bool = False + ) -> RigidTransform: + """Initialize from a unit dual quaternion. + + Unit dual quaternions encode orientation in a real unit quaternion + and translation in a dual quaternion. There is a double cover, i.e., + the unit dual quaternions q and -q represent the same transform. + + Unit dual quaternions must have a real quaternion with unit norm and + a dual quaternion that is orthogonal to the real quaternion to satisfy + the unit norm constraint. This function will enforce both properties + through normalization. + + Parameters + ---------- + dual_quat : array_like, shape (..., 8) + Unit dual quaternions. Each dual_quat[..., :] represents a unit + dual quaternion. The real part is stored in the first four + components and the dual part in the last four components. + scalar_first : bool, optional + Whether the scalar component goes first or last in the two + individual quaternions that represent the real and the dual part. + Default is False, i.e. the scalar-last order is used. + + Returns + ------- + transform : `RigidTransform` instance + Rigid transform objects with the same leading dimensions as the input. + + Examples + -------- + >>> from scipy.spatial.transform import RigidTransform as Tf + >>> import numpy as np + + Creating from a single unit dual quaternion: + + >>> tf = Tf.from_dual_quat([ + ... 0.0617101, -0.06483886, 0.31432811, 0.94508498, + ... 0.04985168, -0.26119618, 0.1691491, -0.07743254]) + >>> tf.as_matrix() + array([[0.79398752, -0.60213598, -0.08376202, 0.24605262], + [0.58613113, 0.79477941, -0.15740392, -0.4932833], + [0.16135089, 0.07588122, 0.98397557, 0.34262676], + [0., 0., 0., 1.]]) + >>> tf.single + True + """ + xp = array_namespace(dual_quat) + dual_quat = _promote(dual_quat, xp=xp) + backend = select_backend(xp, dual_quat.ndim < 3) + matrix = backend.from_dual_quat(dual_quat, scalar_first=scalar_first) + return RigidTransform._from_raw_matrix(matrix, xp, backend) + + @staticmethod + def identity( + num: int | None = None, *, shape: int | tuple[int, ...] | None = None + ) -> RigidTransform: + """Initialize an identity transform. + + Composition with the identity transform has no effect, and + applying the identity transform to a vector has no effect. + + Parameters + ---------- + num : int, optional + Number of identity transforms to generate. If None (default), + then a single transform is generated. + shape : int or tuple of ints, optional + Shape of the identity transforms. If specified, `num` must + be None. + + Returns + ------- + transform : `RigidTransform` instance + The identity transform. + + Examples + -------- + >>> from scipy.spatial.transform import RigidTransform as Tf + >>> from scipy.spatial.transform import Rotation as R + >>> import numpy as np + + Creating a single identity transform: + + >>> tf = Tf.identity() + >>> tf.as_matrix() + array([[1., 0., 0., 0.], + [0., 1., 0., 0.], + [0., 0., 1., 0.], + [0., 0., 0., 1.]]) + >>> tf.single + True + + The identity transform can be applied to a vector without effect: + + >>> tf.apply([1, 2, 3]) + array([1., 2., 3.]) + + The identity transform when composed with another transform has no + effect: + + >>> rng = np.random.default_rng(123) + >>> t = rng.random(3) + >>> r = R.random(rng=rng) + >>> tf = Tf.from_components(t, r) + >>> np.allclose((Tf.identity() * tf).as_matrix(), + ... tf.as_matrix(), atol=1e-12) + True + + Multiple identity transforms can be generated at once: + + >>> tf = Tf.identity(2) + >>> tf.as_matrix() + array([[[1., 0., 0., 0.], + [0., 1., 0., 0.], + [0., 0., 1., 0.], + [0., 0., 0., 1.]], + [[1., 0., 0., 0.], + [0., 1., 0., 0.], + [0., 0., 1., 0.], + [0., 0., 0., 1.]]]) + >>> tf.single + False + >>> len(tf) + 2 + """ + if num is not None and shape is not None: + raise ValueError("Only one of `num` and `shape` can be specified.") + if num is None and shape is None: + shape = () + elif num is not None: + shape = (num,) + elif isinstance(shape, int): + shape = (shape,) + elif not isinstance(shape, tuple): + raise ValueError("`shape` must be an int or a tuple of ints or None.") + matrix = np.tile(np.eye(4), shape + (1, 1)) + # No need for a backend call here since identity is easy to construct and we are + # currently not offering a backend-specific identity matrix + return RigidTransform._from_raw_matrix(matrix, array_namespace(matrix)) + + @staticmethod + def concatenate( + transforms: RigidTransform | Iterable[RigidTransform], + ) -> RigidTransform: + """ + Concatenate a sequence of `RigidTransform` objects into a + single object. + + Parameters + ---------- + transforms : sequence of `RigidTransform` + If a single `RigidTransform` instance is passed in, a copy of + it is returned. + + Returns + ------- + transform : `RigidTransform` instance + The concatenated transform. + + Examples + -------- + >>> from scipy.spatial.transform import RigidTransform as Tf + >>> tf1 = Tf.from_translation([1, 0, 0]) + >>> tf2 = Tf.from_translation([[2, 0, 0], [3, 0, 0]]) + >>> Tf.concatenate([tf1, tf2]).translation + array([[1., 0., 0.], + [2., 0., 0.], + [3., 0., 0.]]) + """ + if isinstance(transforms, RigidTransform): + return RigidTransform._from_raw_matrix( + transforms.as_matrix(), transforms._xp + ) + if not all(isinstance(x, RigidTransform) for x in transforms): + raise TypeError("input must contain RigidTransform objects only") + + xp = array_namespace(transforms[0].as_matrix()) # type:ignore[index] + matrix = xp.concat( + [xpx.atleast_nd(x.as_matrix(), ndim=3, xp=xp) for x in transforms] + ) + return RigidTransform._from_raw_matrix(matrix, xp, None) + + def mean(self, + weights: ArrayLike | None = None, + axis: None | int | tuple[int, ...] = None + ) -> RigidTransform: + """Get the mean of the transforms. + + The mean of a set of transforms is the same as the mean of its + rotation and translation components. + + The mean used for the rotation component is the chordal L2 mean (also + called the projected or induced arithmetic mean) [1]_. If ``A`` is a + set of rotation matrices, then the mean ``M`` is the rotation matrix + that minimizes the following loss function: + + .. math:: + + L(M) = \\sum_{i = 1}^{n} w_i \\lVert \\mathbf{A}_i - + \\mathbf{M} \\rVert^2 , + + where :math:`w_i`'s are the `weights` corresponding to each matrix. + + Parameters + ---------- + weights : array_like shape (..., N), optional + Weights describing the relative importance of the transforms. If + None (default), then all values in `weights` are assumed to be + equal. If given, the shape of `weights` must be broadcastable to + the transform shape. Weights must be non-negative. + axis : None, int, or tuple of ints, optional + Axis or axes along which the means are computed. The default is to + compute the mean of all transforms. + + Returns + ------- + mean : `RigidTransform` instance + Single transform containing the mean of the transforms in the + current instance. + + References + ---------- + .. [1] Hartley, Richard, et al., + "Rotation Averaging", International Journal of Computer Vision + 103, 2013, pp. 267-305. + + Examples + -------- + >>> import numpy as np + >>> from scipy.spatial.transform import RigidTransform as Tf + >>> from scipy.spatial.transform import Rotation as R + >>> rng = np.random.default_rng(seed=123) + + The mean of a set of transforms is the same as the mean of the + translation and rotation components: + + >>> t = rng.random((4, 3)) + >>> r = R.random(4, rng=rng) + >>> tf = Tf.from_components(t, r) + >>> tf.mean().as_matrix() + array([[ 0.61593485, -0.74508342, 0.25588075, 0.66999034], + [-0.59353615, -0.65246765, -0.47116962, 0.25481794], + [ 0.51801458, 0.13833531, -0.84411151, 0.52429339], + [0., 0., 0., 1.]]) + >>> Tf.from_components(t.mean(axis=0), r.mean()).as_matrix() + array([[ 0.61593485, -0.74508342, 0.25588075, 0.66999034], + [-0.59353615, -0.65246765, -0.47116962, 0.25481794], + [ 0.51801458, 0.13833531, -0.84411151, 0.52429339], + [0., 0., 0., 1.]]) + """ + mean = self._backend.mean(self._matrix, weights=weights, axis=axis) + return RigidTransform._from_raw_matrix(mean, xp=self._xp, + backend=self._backend) + + def as_matrix(self) -> Array: + """Return a copy of the matrix representation of the transform. + + 4x4 rigid transformation matrices are of the form:: + + [ tx] + [ R ty] + [ tz] + [ 0 0 0 1] + + where ``R`` is a 3x3 orthonormal rotation matrix and + ``t = [tx, ty, tz]`` is a 3x1 translation vector. + + Returns + ------- + matrix : numpy.ndarray, shape (..., 4, 4) + Transformation matrices with the same leading dimensions as the transform. + + Examples + -------- + >>> from scipy.spatial.transform import RigidTransform as Tf + >>> from scipy.spatial.transform import Rotation as R + >>> import numpy as np + + A transformation matrix is a 4x4 matrix formed from a 3x3 rotation + matrix and a 3x1 translation vector: + + >>> t = np.array([2, 3, 4]) + >>> r = R.from_matrix([[0, 0, 1], + ... [1, 0, 0], + ... [0, 1, 0]]) + >>> tf = Tf.from_components(t, r) + >>> tf.as_matrix() + array([[ 0., 0., 1., 2.], + [ 1., 0., 0., 3.], + [ 0., 1., 0., 4.], + [ 0., 0., 0., 1.]]) + + >>> Tf.identity(2).as_matrix() + array([[[1., 0., 0., 0.], + [0., 1., 0., 0.], + [0., 0., 1., 0.], + [0., 0., 0., 1.]], + [[1., 0., 0., 0.], + [0., 1., 0., 0.], + [0., 0., 1., 0.], + [0., 0., 0., 1.]]]) + """ + matrix = self._xp.asarray(self._matrix, copy=True) + if self._single: + return matrix[0, ...] + return matrix + + def as_components(self) -> tuple[Array, Rotation]: + """Return the translation and rotation components of the transform, + where the rotation is applied first, followed by the translation. + + 4x4 rigid transformation matrices are of the form:: + + [ tx] + [ R ty] + [ tz] + [ 0 0 0 1] + + Where ``R`` is a 3x3 orthonormal rotation matrix and + ``t = [tx, ty, tz]`` is a 3x1 translation vector. This function + returns the rotation corresponding to this rotation matrix + ``r = Rotation.from_matrix(R)`` and the translation vector ``t``. + + When applying a transform ``tf`` to a vector ``v``, the result is the same + as if the rotation and translation components were applied to the vector + with the following operation: + ``tf.apply(v) == translation + rotation.apply(v)``. + + Returns + ------- + translation : numpy.ndarray, shape (..., 3) + The translation of the transform. + rotation : `Rotation` instance + The rotation of the transform. + + Examples + -------- + >>> from scipy.spatial.transform import RigidTransform as Tf + >>> from scipy.spatial.transform import Rotation as R + >>> import numpy as np + + Recover the rotation and translation from a transform: + + >>> t = np.array([2, 3, 4]) + >>> r = R.from_matrix([[0, 0, 1], + ... [1, 0, 0], + ... [0, 1, 0]]) + >>> tf = Tf.from_components(t, r) + >>> t, r = tf.as_components() + >>> t + array([2., 3., 4.]) + >>> r.as_matrix() + array([[0., 0., 1.], + [1., 0., 0.], + [0., 1., 0.]]) + + The transform applied to a vector is equivalent to the rotation applied + to the vector, followed by the translation: + + >>> t + r.apply([1, 0, 0]) + array([2., 4., 4.]) + >>> tf.apply([1, 0, 0]) + array([2., 4., 4.]) + """ + return self.translation, self.rotation + + def as_exp_coords(self) -> Array: + """Return the exponential coordinates of the transform. + + This implements the logarithmic map that converts SE(3) to 6-dimensional + real vectors. + + This is an inverse of `from_exp_coords` where details on the mapping can + be found. + + Returns + ------- + exp_coords : numpy.ndarray, shape (..., 6) + Exponential coordinate vectors with the same leading dimensions as the + transform. The first three components define the rotation and the last + three components define the translation. + + Examples + -------- + >>> from scipy.spatial.transform import RigidTransform as Tf + >>> import numpy as np + + Get exponential coordinates of the identity matrix: + + >>> Tf.identity().as_exp_coords() + array([0., 0., 0., 0., 0., 0.]) + """ + exp_coords = self._backend.as_exp_coords(self._matrix) + if self._single: + exp_coords = exp_coords[0, ...] + return exp_coords + + def as_dual_quat(self, *, scalar_first: bool = False) -> Array: + """Return the dual quaternion representation of the transform. + + Unit dual quaternions encode orientation in a real unit quaternion + and translation in a dual quaternion. There is a double cover, i.e., + the unit dual quaternions q and -q represent the same transform. + + Parameters + ---------- + scalar_first : bool, optional + Whether the scalar component goes first or last in the two + individual quaternions that represent the real and the dual part. + Default is False, i.e. the scalar-last order is used. + + Returns + ------- + dual_quat : numpy.ndarray, shape (..., 8) + Unit dual quaternion vectors with the same leading dimensions as the + transform. The real part is stored in the first four components and the + dual part in the last four components. + + Examples + -------- + >>> from scipy.spatial.transform import RigidTransform as Tf + >>> import numpy as np + + Get identity dual quaternion (we use scalar-last by default): + + >>> Tf.identity().as_dual_quat() + array([0., 0., 0., 1., 0., 0., 0., 0.]) + + When we want to use the scalar-first convention, we use the argument: + + >>> Tf.identity().as_dual_quat(scalar_first=True) + array([1., 0., 0., 0., 0., 0., 0., 0.]) + """ + dual_quat = self._backend.as_dual_quat(self._matrix, scalar_first=scalar_first) + if self._single: + dual_quat = dual_quat[0] + return dual_quat + + def __len__(self) -> int: + """Return the length of the leading transform dimension. + + A transform can store an N-dimensional array of transforms. The length is the + size of the first dimension of this array. If the transform is a single + transform, the length is not defined and an error is raised. + + Returns + ------- + length : int + The number of transforms in this object. + + Raises + ------ + TypeError + If the transform is a single transform. + + Examples + -------- + >>> from scipy.spatial.transform import RigidTransform as Tf + >>> import numpy as np + >>> tf = Tf.identity(3) + >>> len(tf) + 3 + + An N-dimensional array of transforms returns its first dimension size: + + >>> t = np.ones((5, 2, 3, 1, 3)) + >>> tf = Tf.from_translation(t) + >>> len(tf) + 5 + + A single transform has no length: + + >>> tf = Tf.from_translation([1, 0, 0]) + >>> len(tf) # doctest: +IGNORE_EXCEPTION_DETAIL + Traceback (most recent call last): + ... + TypeError: Single transform has no len(). + """ + # We don't use self._single here because we also want to raise an error for + # Array API backends that call len() on single RigidTransform objects. + if self.single: + raise TypeError("Single transform has no len") + return self._matrix.shape[0] + + def __getitem__( + self, indexer: int | slice | EllipsisType | None | ArrayLike + ) -> RigidTransform: + """Extract transform(s) at given index(es) from this object. + + Creates a new `RigidTransform` instance containing a subset of + transforms stored in this object. + + Parameters + ---------- + indexer : int or slice or array_like + Specifies which transform(s) to extract. A single indexer must be + specified, i.e. as if indexing a 1 dimensional array or list. + + Returns + ------- + transform : `RigidTransform` instance + Contains + - a single transform, if `indexer` is a single index + - a stack of transform(s), if `indexer` is a slice, or an index + array. + + Raises + ------ + TypeError + If the transform is a single transform. + + Examples + -------- + >>> from scipy.spatial.transform import RigidTransform as Tf + >>> t = [[0, 0, 0], [1, 0, 0], [2, 0, 0]] # 3 translations + >>> tf = Tf.from_translation(t) + + A single index returns a single transform: + + >>> tf[0].as_matrix() + array([[1., 0., 0., 0.], + [0., 1., 0., 0.], + [0., 0., 1., 0.], + [0., 0., 0., 1.]]) + + A slice returns a stack of transforms: + + >>> tf[1:3].translation + array([[1., 0., 0.], + [2., 0., 0.]]) + + An index array returns a stack of transforms: + + >>> tf[[0, 2]].translation + array([[0., 0., 0.], + [2., 0., 0.]]) + """ + if self.single: + raise TypeError("Single transform is not subscriptable.") + + is_array = isinstance(indexer, type(self._matrix)) + xp = self._xp + # Masking is only specified in the Array API when the array is the sole index + # This special case handling is necessary to support boolean indexing and + # integer array indexing with take (see + # https://github.com/data-apis/array-api/pull/900#issuecomment-2674432480) + # Ideally we would converge to [indexer, ...] indexing, but this is not + # supported for now. + if is_array and indexer.dtype == xp.bool: # type:ignore[union-attr] + return RigidTransform(self._matrix[indexer], normalize=False) + if is_array and xp.isdtype(indexer.dtype, "integral"): # type:ignore[union-attr] + if self._matrix.shape[0] == 0: + raise IndexError("cannot take from an empty array") + return RigidTransform( + xp.take(self._matrix, indexer, axis=0), normalize=False + ) + return RigidTransform(self._matrix[indexer, ...], normalize=False) + + def __setitem__( + self, + indexer: int | slice | EllipsisType | None | ArrayLike, + value: RigidTransform, + ): + """Set transform(s) at given index(es) in this object. + + Parameters + ---------- + indexer : int or slice or array_like + Specifies which transform(s) to replace. A single indexer must be + specified, i.e. as if indexing a 1 dimensional array or list. + + value : `RigidTransform` instance + The transform(s) to set. + + Raises + ------ + TypeError + If the transform is a single transform. + + Examples + -------- + >>> from scipy.spatial.transform import RigidTransform as Tf + >>> t = [[0, 0, 0], [1, 0, 0], [2, 0, 0]] # 3 translations + >>> tf = Tf.from_translation(t) + + Set a single transform: + + >>> tf[0] = Tf.from_translation([9, 9, 9]) + >>> tf.translation + array([[9., 9., 9.], + [1., 0., 0.], + [2., 0., 0.]]) + """ + if self.single: + raise TypeError("Single transform is not subscriptable.") + + if not isinstance(value, RigidTransform): + raise TypeError("value must be a RigidTransform object") + + self._matrix = self._backend.setitem(self._matrix, indexer, value.as_matrix()) + + def __mul__( + self, other: RigidTransform | Rotation + ) -> RigidTransform | NotImplementedType: + """Compose this transform with the other. + + If ``p`` and ``q`` are two transforms, then the composition of '``q`` + followed by ``p``' is equivalent to ``p * q``. In terms of + transformation matrices, the composition can be expressed as + ``p.as_matrix() @ q.as_matrix()``. + + In terms of translations and rotations, the composition when applied to + a vector ``v`` is equivalent to + ``p.translation + p.rotation.apply(q.translation) + + (p.rotation * q.rotation).apply(v)``. + + This function supports composition of multiple transforms at a time using + broadcasting rules. The resulting shape for two `RigidTransform` instances + ``p`` and ``q`` is `np.broadcast_shapes(p.shape, q.shape)`. + + If ``other`` is a `Rotation`, it is automatically promoted to a + `RigidTransform` with zero translation. + + Parameters + ---------- + other : `RigidTransform` or `Rotation` instance + Transform(s) or rotation(s) to be composed with this one. The shapes + must be broadcastable. + + Returns + ------- + `RigidTransform` instance + The composed transform. + + Examples + -------- + >>> from scipy.spatial.transform import RigidTransform as Tf + >>> from scipy.spatial.transform import Rotation as R + >>> import numpy as np + + Compose two transforms: + + >>> tf1 = Tf.from_translation([1, 0, 0]) + >>> tf2 = Tf.from_translation([0, 1, 0]) + >>> tf = tf1 * tf2 + >>> tf.translation + array([1., 1., 0.]) + >>> tf.single + True + + When applied to a vector, the composition of two transforms is applied + in right-to-left order. + + >>> t1, r1 = [1, 2, 3], R.from_euler('z', 60, degrees=True) + >>> t2, r2 = [0, 1, 0], R.from_euler('x', 30, degrees=True) + >>> tf1 = Tf.from_components(t1, r1) + >>> tf2 = Tf.from_components(t2, r2) + >>> tf = tf1 * tf2 + >>> tf.apply([1, 0, 0]) + array([0.6339746, 3.3660254, 3. ]) + >>> tf1.apply(tf2.apply([1, 0, 0])) + array([0.6339746, 3.3660254, 3. ]) + + When at least one of the transforms is not single, the result is a stack + of transforms. + + >>> tf1 = Tf.from_translation([1, 0, 0]) + >>> tf2 = Tf.from_translation([[0, 2, 0], [0, 0, 3]]) + >>> tf = tf1 * tf2 + >>> tf.translation + array([[1., 2., 0.], + [1., 0., 3.]]) + >>> tf.single + False + >>> len(tf) + 2 + + Broadcasting rules apply when composing multiple transforms at a time. + + >>> tf1 = Tf.from_translation(np.ones((5, 1, 3))) # Shape (5, 1, 3) + >>> tf2 = Tf.from_translation(np.ones((1, 4, 3))) # Shape (1, 4, 3) + >>> tf = tf1 * tf2 # Shape (5, 4, 3) + >>> tf.translation.shape + (5, 4, 3) + """ + if isinstance(other, Rotation): + other = RigidTransform.from_rotation(other) + elif not isinstance(other, RigidTransform): + # If other is not a RigidTransform, we return NotImplemented to allow other + # types to implement __rmul__ + return NotImplemented + if not broadcastable(self._matrix.shape, other._matrix.shape): + raise ValueError( + f"Cannot broadcast {self._matrix.shape[:-2]} transforms in " + f"first to {other._matrix.shape[:-2]} transforms in second object." + ) + cython_compatible = self._matrix.ndim < 4 and other._matrix.ndim < 4 + backend = select_backend(self._xp, cython_compatible=cython_compatible) + matrix = backend.compose_transforms(self._matrix, other._matrix) + # Only necessary for cython. Array API broadcasting handles this by default + if self._single and other._single: + matrix = matrix[0, ...] + return RigidTransform(matrix, normalize=True, copy=False) + + def __rmul__(self, other: Rotation) -> RigidTransform | NotImplementedType: + """Compose a rotation with this transform (rotation applied second). + + See `__mul__` for more details. + + Parameters + ---------- + other : `Rotation` instance + The rotation to compose with this transform. The shapes must be + broadcastable. + + Returns + ------- + `RigidTransform` instance + The composed transform. + """ + if isinstance(other, Rotation): + return RigidTransform.from_rotation(other) * self + # When other is a RigidTransform __mul__ is called, so we don't handle it here + return NotImplemented + + def __pow__(self, n: float) -> RigidTransform: + """Compose this transform with itself `n` times. + + A rigid transform `p` when raised to non-integer powers can be thought + of as finding a fraction of the transformation. For example, a power of + 0.5 finds a "halfway" transform from the identity to `p`. + + This is implemented by applying screw linear interpolation (ScLERP) + between `p` and the identity transform, where the angle of the rotation + component is scaled by `n`, and the translation is proportionally + adjusted along the screw axis. + + ``q = p ** n`` can also be expressed as + ``q = RigidTransform.from_exp_coords(p.as_exp_coords() * n)``. + + If `n` is negative, then the transform is inverted before the power + is applied. In other words, ``p ** -abs(n) == p.inv() ** abs(n)``. + + Parameters + ---------- + n : float + The number of times to compose the transform with itself. + + Returns + ------- + `RigidTransform` instance + If the input Rotation `p` contains `N` multiple rotations, then + the output will contain `N` rotations where the `i` th rotation + is equal to ``p[i] ** n``. + + Notes + ----- + There are three notable cases: if ``n == 1`` then a copy of the original + transform is returned, if ``n == 0`` then the identity transform is + returned, and if ``n == -1`` then the inverse transform is returned. + + Note that fractional powers ``n`` which effectively take a root of + rotation, do so using the shortest path smallest representation of that + angle (the principal root). This means that powers of ``n`` and ``1/n`` + are not necessarily inverses of each other. For example, a 0.5 power of + a +240 degree rotation will be calculated as the 0.5 power of a -120 + degree rotation, with the result being a rotation of -60 rather than + +120 degrees. + + Examples + -------- + >>> from scipy.spatial.transform import RigidTransform as Tf + >>> import numpy as np + + A power of 2 returns the transform composed with itself: + + >>> tf = Tf.from_translation([1, 2, 3]) + >>> (tf ** 2).translation + array([2., 4., 6.]) + >>> (tf ** 2).as_matrix() + array([[1., 0., 0., 2.], + [0., 1., 0., 4.], + [0., 0., 1., 6.], + [0., 0., 0., 1.]]) + + A negative power returns the inverse of the transform raised to the + absolute value of `n`: + + >>> (tf ** -2).translation + array([-2., -4., -6.]) + >>> np.allclose((tf ** -2).as_matrix(), (tf.inv() ** 2).as_matrix(), + ... atol=1e-12) + True + + A power of 0 returns the identity transform: + + >>> (tf ** 0).as_matrix() + array([[1., 0., 0., 0.], + [0., 1., 0., 0.], + [0., 0., 1., 0.], + [0., 0., 0., 1.]]) + + A power of 1 returns a copy of the original transform: + + >>> (tf ** 1).as_matrix() + array([[1., 0., 0., 1.], + [0., 1., 0., 2.], + [0., 0., 1., 3.], + [0., 0., 0., 1.]]) + + A fractional power returns a transform with a scaled rotation and + translated along the screw axis. Here we take the square root of the + transform, which when squared recovers the original transform: + + >>> tf_half = (tf ** 0.5) + >>> tf_half.translation + array([0.5, 1., 1.5]) + >>> (tf_half ** 2).as_matrix() + array([[1., 0., 0., 1.], + [0., 1., 0., 2.], + [0., 0., 1., 3.], + [0., 0., 0., 1.]]) + """ + matrix = self._backend.pow(self._matrix, n) + if self._single: + matrix = matrix[0, ...] + return RigidTransform._from_raw_matrix(matrix, self._xp, self._backend) + + def inv(self) -> RigidTransform: + """Invert this transform. + + Composition of a transform with its inverse results in an identity + transform. + + A rigid transform is a composition of a rotation and a translation, + where the rotation is applied first, followed by the translation. So the + inverse transform is equivalent to the inverse translation followed by + the inverse rotation. + + Returns + ------- + `RigidTransform` instance + The inverse of this transform. + + Examples + -------- + >>> from scipy.spatial.transform import RigidTransform as Tf + >>> from scipy.spatial.transform import Rotation as R + >>> import numpy as np + + A transform composed with its inverse results in an identity transform: + + >>> rng = np.random.default_rng(seed=123) + >>> t = rng.random(3) + >>> r = R.random(rng=rng) + >>> tf = Tf.from_components(t, r) + >>> tf.as_matrix() + array([[-0.45431291, 0.67276178, -0.58394466, 0.68235186], + [-0.23272031, 0.54310598, 0.80676958, 0.05382102], + [ 0.85990758, 0.50242162, -0.09017473, 0.22035987], + [ 0. , 0. , 0. , 1. ]]) + + >>> (tf.inv() * tf).as_matrix() + array([[[1., 0., 0., 0.], + [0., 1., 0., 0.], + [0., 0., 1., 0.], + [0., 0., 0., 1.]]]) + + The inverse rigid transform is the same as the inverse translation + followed by the inverse rotation: + + >>> t, r = tf.as_components() + >>> r_inv = r.inv() # inverse rotation + >>> t_inv = -t # inverse translation + >>> tf_r_inv = Tf.from_rotation(r_inv) + >>> tf_t_inv = Tf.from_translation(t_inv) + >>> np.allclose((tf_r_inv * tf_t_inv).as_matrix(), + ... tf.inv().as_matrix(), + ... atol=1e-12) + True + >>> (tf_r_inv * tf_t_inv * tf).as_matrix() + array([[[1., 0., 0., 0.], + [0., 1., 0., 0.], + [0., 0., 1., 0.], + [0., 0., 0., 1.]]]) + """ + matrix = self._backend.inv(self._matrix) + if self._single: + matrix = matrix[0, ...] + return RigidTransform._from_raw_matrix(matrix, self._xp, self._backend) + + def apply(self, vector: ArrayLike, inverse: bool = False) -> Array: + """Apply the transform to a vector. + + If the original frame transforms to the final frame by this transform, + then its application to a vector can be seen in two ways: + + - As a projection of vector components expressed in the final frame + to the original frame. + - As the physical transformation of a vector being glued to the + original frame as it transforms. In this case the vector + components are expressed in the original frame before and after + the transformation. + + In terms of the rotation matrix and translation, + this application is the same as + ``self.translation + (mat @ vector[..., np.newaxis])[..., 0]``, + where ``mat = self.rotation.as_matrix()``. + + For a single transform, this is the same as + ``self.translation + vector @ self.rotation.as_matrix().T``. + + Parameters + ---------- + vector : array_like, shape (..., 3) + Vector(s) to be transformed. Each vector[..., :] represents a 3D + vector. + inverse : bool, optional + If True, the inverse of the transform is applied to the vector. + + Returns + ------- + transformed_vector : numpy.ndarray, shape (..., 3) + The transformed vector(s) with shape determined by broadcasting + the transform and vector shapes together. + + Examples + -------- + >>> from scipy.spatial.transform import RigidTransform as Tf + >>> from scipy.spatial.transform import Rotation as R + >>> import numpy as np + + Apply a single transform to a vector. Here the transform is just a + translation, so the result is the vector added to the translation + vector. + + >>> t = np.array([1, 2, 3]) + >>> tf = Tf.from_translation(t) + >>> t + np.array([1, 0, 0]) + array([2, 2, 3]) + >>> tf.apply([1, 0, 0]) + array([2., 2., 3.]) + + Apply a single transform to a stack of vectors: + + >>> tf.apply([[1, 0, 0], [0, 1, 0]]) + array([[2., 2., 3.], + [1., 3., 3.]]) + + Apply the inverse of a transform to a vector, so the result is the + negative of the translation vector added to the vector. + + >>> -t + np.array([1, 0, 0]) + array([0, -2, -3]) + >>> tf.apply([1, 0, 0], inverse=True) + array([0., -2., -3.]) + + Broadcasting is supported when applying multiple transforms to an N-dimensional + array of vectors. + + >>> tf = Tf.from_translation(np.ones((4, 5, 1, 3))) + >>> vectors = np.zeros((2, 3)) + >>> tf.apply(vectors).shape + (4, 5, 2, 3) + + For transforms which are not just pure translations, applying it to a + vector is the same as applying the rotation component to the vector and + then adding the translation component. + + >>> r = R.from_euler('z', 60, degrees=True) + >>> tf = Tf.from_components(t, r) + >>> t + r.apply([1, 0, 0]) + array([1.5, 2.8660254, 3. ]) + >>> tf.apply([1, 0, 0]) + array([1.5, 2.8660254, 3. ]) + + When applying the inverse of a transform, the result is the negative of + the translation vector added to the vector, and then rotated by the + inverse rotation. + + >>> r.inv().apply(-t + np.array([1, 0, 0])) + array([-1.73205081, -1. , -3. ]) + >>> tf.apply([1, 0, 0], inverse=True) + array([-1.73205081, -1. , -3. ]) + """ + # We do not use the cached xp here to catch cases where matrix and vector have + # different xp instances. + xp = array_namespace(self._matrix, vector) + vector = xp.asarray( + vector, dtype=self._matrix.dtype, device=device(self._matrix) + ) + cython_compatible = self._matrix.ndim < 4 and vector.ndim < 3 + backend = select_backend(xp, cython_compatible=cython_compatible) + result = backend.apply(self._matrix, vector, inverse) + if self._single and vector.ndim == 1: + result = result[0, ...] + return result + + @property + def rotation(self) -> Rotation: + """Return the rotation component of the transform. + + A transform is a composition of a rotation and a translation, such that + when applied to a vector, the vector is first rotated and then + translated. This property returns the rotation part of the transform. + + Returns + ------- + rotation : `Rotation` instance + A single rotation or a stack of rotations. + + Examples + -------- + >>> from scipy.spatial.transform import RigidTransform as Tf + >>> from scipy.spatial.transform import Rotation as R + >>> import numpy as np + + The rotation component is extracted from the transform: + + >>> t = np.array([1, 0, 0]) + >>> r = R.random(3) + >>> tf = Tf.from_components(t, r) + >>> np.allclose(tf.rotation.as_matrix(), r.as_matrix()) + True + """ + if self._single: + return Rotation.from_matrix(self._matrix[0, :3, :3]) + return Rotation.from_matrix(self._matrix[..., :3, :3]) + + @property + def translation(self) -> Array: + """Return the translation component of the transform. + + A transform is a composition of a rotation and a translation, such that + when applied to a vector, the vector is first rotated and then + translated. This property returns the translation part of the transform. + + Returns + ------- + translation : numpy.ndarray, shape (..., 3) + Translation vectors with the same leading dimensions as the transform. + + Examples + -------- + >>> from scipy.spatial.transform import RigidTransform as Tf + >>> from scipy.spatial.transform import Rotation as R + >>> import numpy as np + + The translation component is extracted from the transform: + + >>> t = np.array([[1, 0, 0], [2, 0, 0], [3, 0, 0]]) + >>> r = R.random() + >>> tf = Tf.from_components(t, r) + >>> np.allclose(tf.translation, t) + True + """ + if self._single: + return self._xp.asarray(self._matrix[0, :3, 3], copy=True) + return self._xp.asarray(self._matrix[..., :3, 3], copy=True) + + @property + def single(self) -> bool: + """Whether this instance represents a single transform. + + Single transforms are not subscriptable, and do not have a length. + + Returns + ------- + single : bool + True if this instance represents a single transform, False + otherwise. + """ + return self._single or self._matrix.ndim == 2 + + @property + def shape(self) -> tuple[int, ...]: + """The shape of the transform's leading dimensions.""" + if self._single or self._matrix.ndim == 2: + return () + return self._matrix.shape[:-2] + + def __reduce__(self) -> tuple[Callable, tuple]: + """Reduce the RigidTransform for pickling. + + We store modules inside RigidTransforms which cannot be pickled. To circumvent + this, we pickle only the matrix and restore the cached modules from the matrix + type in `from_matrix`. + """ + matrix = self._matrix + if self._single: + matrix = matrix[0, ...] + return (self.__class__.from_matrix, (matrix,)) + + def __iter__(self) -> Iterator[RigidTransform]: + """Iterate over transforms.""" + if self._single or self._matrix.ndim == 2: + raise TypeError("Single transform is not iterable.") + # We return a generator that yields a new RigidTransform object for each + # transform in the current object. We cannot rely on the default implementation + # because jax will not raise an IndexError for out-of-bounds indices. + for i in range(self._matrix.shape[0]): + yield RigidTransform._from_raw_matrix( + self._matrix[i, ...], self._xp, self._backend + ) + + @staticmethod + def _from_raw_matrix( + matrix: Array, xp: ModuleType, backend: ModuleType | None = None + ) -> RigidTransform: + """Create a RigidTransform skipping all sanitization steps. + + This method is is intended for internal, performant creation of RigidTransforms + with matrices that are guaranteed to be valid. + """ + tf = RigidTransform.__new__(RigidTransform) + tf._single = matrix.ndim == 2 and is_numpy(xp) + if tf._single: + matrix = xpx.atleast_nd(matrix, ndim=3, xp=xp) + tf._matrix = matrix + tf._xp = xp + if backend is None: + backend = select_backend(xp, matrix.ndim < 4) + tf._backend = backend + return tf diff --git a/v12x/pylib/scipy/spatial/transform/_rigid_transform_xp.py b/v12x/pylib/scipy/spatial/transform/_rigid_transform_xp.py new file mode 100644 index 0000000000000000000000000000000000000000..8700f3065283dcfdeece75c7ef1f4c18388d2757 --- /dev/null +++ b/v12x/pylib/scipy/spatial/transform/_rigid_transform_xp.py @@ -0,0 +1,448 @@ +from types import EllipsisType + +from scipy._lib._array_api import ( + array_namespace, + Array, + ArrayLike, + is_lazy_array, + xp_vector_norm, + xp_result_type, + xp_promote, + is_array_api_obj, +) +import scipy._external.array_api_extra as xpx +from scipy.spatial.transform._rotation_xp import ( + as_matrix as quat_as_matrix, + from_matrix as quat_from_matrix, + _from_matrix_orthogonal as quat_from_matrix_orthogonal, + from_rotvec as quat_from_rotvec, + as_rotvec as quat_as_rotvec, + compose_quat, + from_quat, + inv as quat_inv, + mean as quat_mean, +) +from scipy._external.array_api_compat import device as xp_device +from scipy._lib._util import broadcastable + + +def from_matrix(matrix: Array, normalize: bool = True, copy: bool = True) -> Array: + xp = array_namespace(matrix) + # Shape check should be done before calling this function + if normalize or copy: + matrix = xp.asarray(matrix, copy=True) + + last_row_ok = xp.all( + matrix[..., 3, :] == xp.asarray([0, 0, 0, 1.0], device=xp_device(matrix)), + axis=-1, + ) + lazy = is_lazy_array(matrix) + # We delay lazy branch checks until after normalization to avoid overwriting nans + # with the rotation matrix + if not lazy and xp.any(~last_row_ok): + if last_row_ok.shape == (): + idx = () + else: + idx = tuple(int(i[0]) for i in xp.nonzero(~last_row_ok)) + vals = matrix[idx + (3, ...)] + raise ValueError( + f"Expected last row of transformation matrix {idx} to be " + f"exactly [0, 0, 0, 1], got {vals}" + ) + + # The quat_from_matrix() method orthogonalizes the rotation + # component of the transformation matrix. While this does have some + # overhead in converting a rotation matrix to a quaternion and back, it + # allows for skipping singular value decomposition for near-orthogonal + # matrices, which is a computationally expensive operation. + if normalize: + rotmat = quat_as_matrix(quat_from_matrix(matrix[..., :3, :3])) + matrix = xpx.at(matrix)[..., :3, :3].set(rotmat) + # Lazy branch matrix invalidation + if lazy: + matrix = xp.where(last_row_ok[..., None, None], matrix, xp.nan) + return matrix + + +def from_rotation(quat: Array) -> Array: + xp = array_namespace(quat) + rotmat = quat_as_matrix(quat) + matrix = xp.zeros( + (*rotmat.shape[:-2], 4, 4), dtype=quat.dtype, device=xp_device(quat) + ) + matrix = xpx.at(matrix)[..., :3, :3].set(rotmat) + matrix = xpx.at(matrix)[..., 3, 3].set(1) + return matrix + + +def from_translation(translation: Array) -> Array: + xp = array_namespace(translation) + + if translation.shape[-1] != 3: + raise ValueError( + f"Expected `translation` to have shape (..., 3), got {translation.shape}." + ) + device = xp_device(translation) + dtype = xp_result_type(translation, force_floating=True, xp=xp) + eye = xp.eye(4, dtype=dtype, device=device) + + matrix = xpx.atleast_nd(eye, ndim=translation.ndim + 1, xp=xp) + matrix = xp.zeros( + (*translation.shape[:-1], 4, 4), + dtype=dtype, + device=device, + ) + matrix = xpx.at(matrix)[...].set(xp.eye(4, dtype=dtype, device=device)) + matrix = xpx.at(matrix)[..., :3, 3].set( + xp_promote(translation, force_floating=True, xp=xp) + ) + return matrix + + +def from_components(translation: Array, quat: Array) -> Array: + return compose_transforms(from_translation(translation), from_rotation(quat)) + + +def from_exp_coords(exp_coords: Array) -> Array: + rot_vec = exp_coords[..., :3] + rot_matrix = quat_as_matrix(quat_from_rotvec(rot_vec)) + translation_transform = _compute_se3_exp_translation_transform(rot_vec) + translations = (translation_transform @ exp_coords[..., 3:, None])[..., 0] + return _create_transformation_matrix(translations, rot_matrix) + + +def from_dual_quat(dual_quat: Array, *, scalar_first: bool = False) -> Array: + xp = array_namespace(dual_quat) + + if dual_quat.shape[-1] != 8: + raise ValueError( + f"Expected `dual_quat` to have shape (..., 8), got {dual_quat.shape}." + ) + + real_part = dual_quat[..., :4] + dual_part = dual_quat[..., 4:] + if scalar_first: + real_part = xp.roll(real_part, -1, axis=-1) + dual_part = xp.roll(dual_part, -1, axis=-1) + + real_part, dual_part = _normalize_dual_quaternion(real_part, dual_part) + + rot_quat = from_quat(real_part) + + translation = 2.0 * compose_quat(dual_part, quat_inv(rot_quat))[..., :3] + matrix = _create_transformation_matrix(translation, quat_as_matrix(rot_quat)) + + return matrix + + +def as_exp_coords(matrix: Array) -> Array: + xp = array_namespace(matrix) + rot_vec = quat_as_rotvec(quat_from_matrix_orthogonal(matrix[..., :3, :3])) + translation_transform = _compute_se3_log_translation_transform(rot_vec) + translations = (translation_transform @ matrix[..., :3, 3][..., None])[..., 0] + exp_coords = xp.concat([rot_vec, translations], axis=-1) + return exp_coords + + +def as_dual_quat(matrix: Array, *, scalar_first: bool = False) -> Array: + xp = array_namespace(matrix) + real_parts = quat_from_matrix_orthogonal(matrix[..., :3, :3]) + + pure_translation_quats = xp.empty( + (*matrix.shape[:-2], 4), dtype=matrix.dtype, device=xp_device(matrix) + ) + pure_translation_quats = xpx.at(pure_translation_quats)[..., :3].set( + matrix[..., :3, 3] + ) + pure_translation_quats = xpx.at(pure_translation_quats)[..., 3].set(0.0) + + dual_parts = 0.5 * compose_quat(pure_translation_quats, real_parts) + + if scalar_first: + real_parts = xp.roll(real_parts, 1, axis=-1) + dual_parts = xp.roll(dual_parts, 1, axis=-1) + + dual_quats = xp.concat([real_parts, dual_parts], axis=-1) + return dual_quats + + +def compose_transforms(tf_matrix: Array, other_tf_matrix: Array) -> Array: + if not broadcastable(tf_matrix.shape, other_tf_matrix.shape): + len_a = tf_matrix.shape[0] if tf_matrix.ndim == 3 else 1 + len_b = other_tf_matrix.shape[0] if other_tf_matrix.ndim == 3 else 1 + raise ValueError( + "Expected equal number of transforms in both or a " + "single transform in either object, got " + f"{len_a} transforms in first and {len_b}" + "transforms in second object." + ) + return tf_matrix @ other_tf_matrix + + +def inv(matrix: Array) -> Array: + xp = array_namespace(matrix) + r_inv = xp.matrix_transpose(matrix[..., :3, :3]) + # Matrix multiplication of r_inv and translation vector + t_inv = -(r_inv @ matrix[..., :3, 3][..., None])[..., 0] + matrix = xp.zeros( + (*matrix.shape[:-2], 4, 4), dtype=matrix.dtype, device=xp_device(matrix) + ) + matrix = xpx.at(matrix)[..., :3, :3].set(r_inv) + matrix = xpx.at(matrix)[..., :3, 3].set(t_inv) + matrix = xpx.at(matrix)[..., 3, 3].set(1) + return matrix + + +def apply(matrix: Array, vector: Array, inverse: bool = False) -> Array: + xp = array_namespace(matrix) + if vector.shape[-1] != 3: + raise ValueError(f"Expected vector to have shape (..., 3), got {vector.shape}.") + vec = xp.empty( + (*vector.shape[:-1], 4), dtype=vector.dtype, device=xp_device(vector) + ) + vec = xpx.at(vec)[..., :3].set(vector) + vec = xpx.at(vec)[..., 3].set(1) + vec = vec[..., None] + + if inverse: + matrix = inv(matrix) + + # We raise a ValueError manually here because letting the function run its course + # would raise heterogeneous error types and messages for different frameworks. + # However, the error only mimics numpy's error message and does not provide the + # same amount of context. + if not broadcastable(matrix.shape, vec.shape): # type:ignore[arg-type] + raise ValueError("operands could not be broadcast together") + return (matrix @ vec)[..., :3, 0] + + +def pow(matrix: Array, n: float | Array) -> Array: + xp = array_namespace(matrix) + device = xp_device(matrix) + # If n is an array, we sanitize it to a scalar and promote quat and n to + # the same dtype. + if is_array_api_obj(n): + if n.shape == (1,): # type:ignore[union-attr] + n = n[0] # type:ignore[index] + elif n.ndim != 0: # type:ignore[union-attr] + raise ValueError("Array exponent must be a scalar") + matrix, n = xp_promote(matrix, n, force_floating=True, xp=xp) + + # If n is a lazy array, we cannot take fast paths for special cases. + if is_lazy_array(n): + # Lazy execution. We compute all special cases and the general case + result = from_exp_coords(as_exp_coords(matrix) * n) + identity = xp.eye(4, dtype=matrix.dtype, device=device) + result = xp.where(n == 0, identity, result) + result = xp.where(n == -1, inv(matrix), result) + result = xp.where(n == 1, matrix, result) + return result + if n == 0: + identity = xp.eye(4, dtype=matrix.dtype, device=device) + identity = xpx.atleast_nd(identity, ndim=matrix.ndim, xp=xp) + return xp.tile(identity, (*matrix.shape[:-2], 1, 1)) + elif n == -1: + return inv(matrix) + elif n == 1: + return matrix + return from_exp_coords(as_exp_coords(matrix) * n) + + +def mean( + matrix: Array, + weights: ArrayLike | None = None, + axis: None | int | tuple[int, ...] = None +) -> Array: + xp = array_namespace(matrix) + if matrix.shape[0] == 0: + raise ValueError("Mean of an empty rotation set is undefined.") + # Axis logic: For None, we reduce over all axes. For int, we only reduce over that + # axis. For tuple, we reduce over all specified axes. + all_axes = tuple(range(matrix.ndim - 2)) + if axis is None: + axis = all_axes + elif isinstance(axis, int): + axis = (axis,) + if not isinstance(axis, tuple): + raise ValueError("`axis` must be None, int, or tuple of ints.") + # Ensure all axes are within bounds + if (axis != () and + (min(axis) < -(matrix.ndim - 2) or max(axis) > (matrix.ndim - 3)) + ): + raise ValueError( + f"axis {axis} is out of bounds for transform with shape " + f"{matrix.shape[:-2]}." + ) + # Ensure all axes are positive and unique + axis = tuple(sorted(set(x % (matrix.ndim - 2) for x in axis))) + + lazy = is_lazy_array(matrix) + quats = quat_from_matrix_orthogonal(matrix[..., :3, :3]) + if weights is None: + quats_mean = quat_mean(quats, axis=axis) + else: + neg_weights = weights < 0 # type:ignore[operator] + any_neg_weights = xp.any(neg_weights) + if not lazy and any_neg_weights: + raise ValueError("`weights` must be non-negative.") + if weights.shape != matrix.shape[:-2]: # type:ignore[union-attr] + raise ValueError( + f"Expected `weights` to match transform shape, got shape " + f"{weights.shape} for {matrix.shape[:-2]} transformations." # type:ignore[union-attr] + ) + quats_mean = quat_mean(quats, weights=weights, axis=axis) + r_mean = quat_as_matrix(quats_mean) + + t = matrix[..., :3, 3] + if weights is None: + t_mean = xp.mean(t, axis=axis) + else: + norm = xp.sum(weights[..., None], axis=axis) # type:ignore[index] + wsum = xp.sum(t * weights[..., None], axis=axis) # type:ignore[index] + t_mean = wsum / norm + + tf = _create_transformation_matrix(t_mean, r_mean) + if weights is not None and lazy: + # We cannot raise on negative weights because jit code needs to be + # non-branching. We return NaN instead + mask = xp.where(any_neg_weights, xp.nan, 1.0) # pyrefly:ignore[unbound-name] + tf = mask * tf + return tf + + +def setitem( + matrix: Array, + indexer: Array | int | tuple | slice | EllipsisType | None, + value: Array, +) -> Array: + xp = array_namespace(matrix) + if isinstance(indexer, EllipsisType): + return xpx.at(matrix)[indexer].set(value) + if is_array_api_obj(indexer) and indexer.dtype == xp.bool: # type:ignore[union-attr] + return xpx.at(matrix)[indexer].set(value) # type:ignore[index] + return xpx.at(matrix)[indexer, ...].set(value) # type:ignore[index] + + +def normalize_dual_quaternion(dual_quat: Array) -> Array: + """Normalize dual quaternion.""" + xp = array_namespace(dual_quat) + real, dual = _normalize_dual_quaternion(dual_quat[..., :4], dual_quat[..., 4:]) + return xp.concat((real, dual), axis=-1) + + +def _create_transformation_matrix( + translations: Array, rotation_matrices: Array +) -> Array: + if not translations.shape[:-1] == rotation_matrices.shape[:-2]: + raise ValueError( + "The number of rotation matrices and translations must be the same." + ) + xp = array_namespace(translations) + matrix = xp.empty( + (*translations.shape[:-1], 4, 4), + dtype=translations.dtype, + device=xp_device(translations), + ) + matrix = xpx.at(matrix)[..., :3, :3].set(rotation_matrices) + matrix = xpx.at(matrix)[..., :3, 3].set(translations) + matrix = xpx.at(matrix)[..., 3, :3].set(0) + matrix = xpx.at(matrix)[..., 3, 3].set(1) + return matrix + + +def _compute_se3_exp_translation_transform(rot_vec: Array) -> Array: + """Compute the transformation matrix from the se3 translation part to SE3 + translation. + + The transformation matrix depends on the rotation vector. + """ + xp = array_namespace(rot_vec) + device = xp_device(rot_vec) + dtype = rot_vec.dtype + angle = xp_vector_norm(rot_vec, axis=-1, keepdims=True, xp=xp) + small_scale = angle < 1e-3 + + k1_small = 0.5 - angle**2 / 24 + angle**4 / 720 + # Avoid division by zero for non-branching computations. The value will get + # discarded in the xp.where selection. + safe_angle = angle + xp.asarray(small_scale, dtype=dtype, device=device) + k1 = (1.0 - xp.cos(angle)) / safe_angle**2 + k1 = xp.where(small_scale, k1_small, k1) + + k2_small = 1 / 6 - angle**2 / 120 + angle**4 / 5040 + # Again, avoid division by zero by adding one to all near-zero angles. + safe_angle = angle + xp.asarray(small_scale, dtype=dtype, device=device) + k2 = (angle - xp.sin(angle)) / safe_angle**3 + k2 = xp.where(small_scale, k2_small, k2) + + s = _create_skew_matrix(rot_vec) + eye = xp.eye(3, dtype=dtype, device=device) + + return eye + k1[..., None] * s + k2[..., None] * s @ s + + +def _compute_se3_log_translation_transform(rot_vec: Array) -> Array: + """Compute the transformation matrix from SE3 translation to the se3 + translation part. + + It is the inverse of `_compute_se3_exp_translation_transform` in a closed + analytical form. + """ + xp = array_namespace(rot_vec) + dtype = rot_vec.dtype + device = xp_device(rot_vec) + angle = xp_vector_norm(rot_vec, axis=-1, keepdims=True, xp=xp) + mask = angle < 1e-3 + + k_small = 1 / 12 + angle**2 / 720 + angle**4 / 30240 + safe_angle = angle + xp.asarray(mask, dtype=dtype, device=device) + k = (1 - 0.5 * angle / xp.tan(0.5 * safe_angle)) / safe_angle**2 + k = xp.where(mask, k_small, k) + + s = _create_skew_matrix(rot_vec) + + return xp.eye(3, dtype=dtype, device=device) - 0.5 * s + k[..., None] * s @ s + + +def _create_skew_matrix(vec: Array) -> Array: + """Create skew-symmetric (aka cross-product) matrix for stack of vectors.""" + xp = array_namespace(vec) + result = xp.zeros((*vec.shape[:-1], 3, 3), dtype=vec.dtype, device=xp_device(vec)) + result = xpx.at(result)[..., 0, 1].set(-vec[..., 2]) + result = xpx.at(result)[..., 0, 2].set(vec[..., 1]) + result = xpx.at(result)[..., 1, 0].set(vec[..., 2]) + result = xpx.at(result)[..., 1, 2].set(-vec[..., 0]) + result = xpx.at(result)[..., 2, 0].set(-vec[..., 1]) + result = xpx.at(result)[..., 2, 1].set(vec[..., 0]) + return result + + +def _normalize_dual_quaternion( + real_part: Array, dual_part: Array +) -> tuple[Array, Array]: + """Ensure that the dual quaternion has unit norm. + + The norm is a dual number and must be 1 + 0 * epsilon, which means that + the real quaternion must have unit norm and the dual quaternion must be + orthogonal to the real quaternion. + """ + xp = array_namespace(real_part) + + real_norm = xp_vector_norm(real_part, axis=-1, keepdims=True, xp=xp) + + # special case: real quaternion is 0, we set it to identity + zero_real_mask = real_norm == 0.0 + unit_quat = xp.asarray( + [0.0, 0.0, 0.0, 1.0], dtype=real_part.dtype, device=xp_device(real_part) + ) + real_part = xp.where(zero_real_mask, unit_quat, real_part) + real_norm = xp.where(zero_real_mask, 1.0, real_norm) + + # 1. ensure unit real quaternion + real_part = real_part / real_norm + dual_part = dual_part / real_norm + + # 2. ensure orthogonality of real and dual quaternion + dual_part -= xp.sum(real_part * dual_part, axis=-1, keepdims=True) * real_part + + return real_part, dual_part diff --git a/v12x/pylib/scipy/spatial/transform/_rotation.py b/v12x/pylib/scipy/spatial/transform/_rotation.py new file mode 100644 index 0000000000000000000000000000000000000000..6765df1d79ffc6e8288d184c85dae7307052d06a --- /dev/null +++ b/v12x/pylib/scipy/spatial/transform/_rotation.py @@ -0,0 +1,2856 @@ +from __future__ import annotations + +from collections.abc import Iterator, Sequence +from types import EllipsisType, GenericAlias, ModuleType, NotImplementedType + +import numpy as np + +import scipy.spatial.transform._rotation_cy as cython_backend +import scipy.spatial.transform._rotation_xp as xp_backend +from scipy.spatial.transform._rotation_groups import create_group +from scipy._lib._array_api import ( + array_namespace, + Array, + is_numpy, + ArrayLike, + is_lazy_array, + xp_capabilities, + xp_promote, +) +from scipy._external.array_api_compat import device as xp_device +import scipy._external.array_api_extra as xpx +from scipy._lib._util import _transition_to_rng, broadcastable + +backend_registry = {array_namespace(np.empty(0)): cython_backend} + + +def select_backend(xp: ModuleType, cython_compatible: bool): + """Select the backend for the given array library. + + We need this selection function because the Cython backend for numpy does not + support quaternions of arbitrary dimensions. We therefore only use the Array API + backend for numpy if we are dealing with rotations of more than one leading + dimension. + """ + if is_numpy(xp) and not cython_compatible: + return xp_backend + return backend_registry.get(xp, xp_backend) + + +def _promote(*args: ArrayLike, xp: ModuleType) -> Array: + """Promote arrays to float64 for numpy, else according to the Array API spec. + + The return array dtype follows the following rules: + - If quat is an ArrayLike or NumPy array, we always promote to float64 + - If quat is an Array from frameworks other than NumPy, we preserve the precision + of the input array dtype. + + The first rule is required by the cython backend signatures that expect + cython.double views. The second rule is necessary to promote non-floating arrays + to the correct type in frameworks that may not support double precision (e.g. + jax by default). + """ + if is_numpy(xp): + args += (np.empty(0, dtype=np.float64),) # Force float64 conversion + out = xp_promote(*args, force_floating=True, xp=xp) + if len(args) == 2: # One argument was passed + the added empty array + return out[0] + return out[:-1] + return xp_promote(*args, force_floating=True, xp=xp) + + +rotation_extra_note = """The methods ``as_davenport``, ``apply``, and ``align_vectors`` + are not supported with cupy<14.*. + + """ + + +@xp_capabilities( + skip_backends=[("dask.array", "missing linalg.cross/det functions")], + method_capabilities={ + "__init__": dict( + jax_jit=False, + skip_backends=[("dask.array", "dask not supported")], + ), + "as_davenport": dict( + skip_backends=[ + ("dask.array", "missing linalg.cross/det functions and .mT attribute"), + ("cupy", "missing .mT attribute in cupy<14.*"), + ] + ), + "apply": dict( + skip_backends=[ + ("dask.array", "missing linalg.cross/det functions and .mT attribute"), + ("cupy", "missing .mT attribute in cupy<14.*"), + ], + ), + "__getitem__": dict( + jax_jit=False, + skip_backends=[("dask.array", "cannot handle zero-length rotations")], + ), + "__setitem__": dict( + jax_jit=False, + skip_backends=[("dask.array", "cannot handle zero-length rotations")], + ), + "align_vectors": dict( + skip_backends=[ + ("dask.array", "missing linalg.cross/det functions and .mT attribute"), + ("cupy", "missing .mT attribute in cupy<14.*"), + ], + ), + }, + extra_note=rotation_extra_note, +) +class Rotation: + """Rotation in 3 dimensions. + + This class provides an interface to initialize from and represent rotations + with: + + - Quaternions + - Rotation Matrices + - Rotation Vectors + - Modified Rodrigues Parameters + - Euler Angles + - Davenport Angles (Generalized Euler Angles) + + The following operations on rotations are supported: + + - Application on vectors + - Rotation Composition + - Rotation Inversion + - Rotation Indexing + + A `Rotation` instance can contain a single rotation transform or rotations of + multiple leading dimensions. E.g., it is possible to have an N-dimensional array of + (N, M, K) rotations. When applied to other rotations or vectors, standard + broadcasting rules apply. + + Indexing within a rotation is supported to access a subset of the rotations stored + in a `Rotation` instance. + + To create `Rotation` objects use ``from_...`` methods (see examples below). + ``Rotation(...)`` is not supposed to be instantiated directly. + + Parameters + ---------- + quat : array_like, shape (..., 4) + Quaternion representing the rotation. + normalize : bool, optional + If True, orthonormalize the rotation matrix using singular value + decomposition. If False, the rotation matrix is not checked for + orthogonality or right-handedness. + copy : bool, optional + If True, copy the input matrix. If False, a reference to the input + matrix is used. If normalize is True, the input matrix is always + copied regardless of the value of copy. + scalar_first : bool, optional + If ``True`` then `quat` is expect in scalar-first format otherwise it is + expected in scalar-last format. Defaults to ``False``. + + Attributes + ---------- + single + + Methods + ------- + __len__ + from_quat + from_matrix + from_rotvec + from_mrp + from_euler + from_davenport + as_quat + as_matrix + as_rotvec + as_mrp + as_euler + as_davenport + concatenate + apply + __mul__ + __pow__ + inv + magnitude + approx_equal + mean + reduce + create_group + __getitem__ + identity + random + align_vectors + + See Also + -------- + Slerp + + Notes + ----- + .. versionadded:: 1.2.0 + + Examples + -------- + >>> from scipy.spatial.transform import Rotation as R + >>> import numpy as np + + A `Rotation` instance can be initialized in any of the above formats and + converted to any of the others. The underlying object is independent of the + representation used for initialization. + + Consider a counter-clockwise rotation of 90 degrees about the z-axis. This + corresponds to the following quaternion (in scalar-last format): + + >>> r = R.from_quat([0, 0, np.sin(np.pi/4), np.cos(np.pi/4)]) + + The rotation can be expressed in any of the other formats: + + >>> r.as_matrix() + array([[ 2.22044605e-16, -1.00000000e+00, 0.00000000e+00], + [ 1.00000000e+00, 2.22044605e-16, 0.00000000e+00], + [ 0.00000000e+00, 0.00000000e+00, 1.00000000e+00]]) + >>> r.as_rotvec() + array([0. , 0. , 1.57079633]) + >>> r.as_euler('zyx', degrees=True) + array([90., 0., 0.]) + + The same rotation can be initialized using a rotation matrix: + + >>> r = R.from_matrix([[0, -1, 0], + ... [1, 0, 0], + ... [0, 0, 1]]) + + Representation in other formats: + + >>> r.as_quat() + array([0. , 0. , 0.70710678, 0.70710678]) + >>> r.as_rotvec() + array([0. , 0. , 1.57079633]) + >>> r.as_euler('zyx', degrees=True) + array([90., 0., 0.]) + + The rotation vector corresponding to this rotation is given by: + + >>> r = R.from_rotvec(np.pi/2 * np.array([0, 0, 1])) + + Representation in other formats: + + >>> r.as_quat() + array([0. , 0. , 0.70710678, 0.70710678]) + >>> r.as_matrix() + array([[ 2.22044605e-16, -1.00000000e+00, 0.00000000e+00], + [ 1.00000000e+00, 2.22044605e-16, 0.00000000e+00], + [ 0.00000000e+00, 0.00000000e+00, 1.00000000e+00]]) + >>> r.as_euler('zyx', degrees=True) + array([90., 0., 0.]) + + The ``from_euler`` method is quite flexible in the range of input formats + it supports. Here we initialize a single rotation about a single axis: + + >>> r = R.from_euler('z', 90, degrees=True) + + Again, the object is representation independent and can be converted to any + other format: + + >>> r.as_quat() + array([0. , 0. , 0.70710678, 0.70710678]) + >>> r.as_matrix() + array([[ 2.22044605e-16, -1.00000000e+00, 0.00000000e+00], + [ 1.00000000e+00, 2.22044605e-16, 0.00000000e+00], + [ 0.00000000e+00, 0.00000000e+00, 1.00000000e+00]]) + >>> r.as_rotvec() + array([0. , 0. , 1.57079633]) + + It is also possible to initialize multiple rotations in a single instance + using any of the ``from_...`` functions. Here we initialize a stack of 3 + rotations using the ``from_euler`` method: + + >>> r = R.from_euler('zyx', [ + ... [90, 0, 0], + ... [0, 45, 0], + ... [45, 60, 30]], degrees=True) + + The other representations also now return a stack of 3 rotations. For + example: + + >>> r.as_quat() + array([[0. , 0. , 0.70710678, 0.70710678], + [0. , 0.38268343, 0. , 0.92387953], + [0.39190384, 0.36042341, 0.43967974, 0.72331741]]) + + Applying the above rotations onto a vector: + + >>> v = [1, 2, 3] + >>> r.apply(v) + array([[-2. , 1. , 3. ], + [ 2.82842712, 2. , 1.41421356], + [ 2.24452282, 0.78093109, 2.89002836]]) + + A `Rotation` instance can be indexed and sliced as if it were an ND array: + + >>> r.as_quat() + array([[0. , 0. , 0.70710678, 0.70710678], + [0. , 0.38268343, 0. , 0.92387953], + [0.39190384, 0.36042341, 0.43967974, 0.72331741]]) + >>> p = r[0] + >>> p.as_matrix() + array([[ 2.22044605e-16, -1.00000000e+00, 0.00000000e+00], + [ 1.00000000e+00, 2.22044605e-16, 0.00000000e+00], + [ 0.00000000e+00, 0.00000000e+00, 1.00000000e+00]]) + >>> q = r[1:3] + >>> q.as_quat() + array([[0. , 0.38268343, 0. , 0.92387953], + [0.39190384, 0.36042341, 0.43967974, 0.72331741]]) + + In fact it can be converted to numpy.array: + + >>> r_array = np.asarray(r) + >>> r_array.shape + (3,) + >>> r_array[0].as_matrix() + array([[ 2.22044605e-16, -1.00000000e+00, 0.00000000e+00], + [ 1.00000000e+00, 2.22044605e-16, 0.00000000e+00], + [ 0.00000000e+00, 0.00000000e+00, 1.00000000e+00]]) + + Multiple rotations can be composed using the ``*`` operator: + + >>> r1 = R.from_euler('z', 90, degrees=True) + >>> r2 = R.from_rotvec([np.pi/4, 0, 0]) + >>> v = [1, 2, 3] + >>> r2.apply(r1.apply(v)) + array([-2. , -1.41421356, 2.82842712]) + >>> r3 = r2 * r1 # Note the order + >>> r3.apply(v) + array([-2. , -1.41421356, 2.82842712]) + + A rotation can be composed with itself using the ``**`` operator: + + >>> p = R.from_rotvec([1, 0, 0]) + >>> q = p ** 2 + >>> q.as_rotvec() + array([2., 0., 0.]) + + Finally, it is also possible to invert rotations: + + >>> r1 = R.from_euler('z', [[90], [45]], degrees=True) + >>> r2 = r1.inv() + >>> r2.as_euler('zyx', degrees=True) + array([[-90., 0., 0.], + [-45., 0., 0.]]) + + The following function can be used to plot rotations with Matplotlib by + showing how they transform the standard x, y, z coordinate axes: + + >>> import matplotlib.pyplot as plt + + >>> def plot_rotated_axes(ax, r, name=None, offset=(0, 0, 0), scale=1): + ... colors = ("#FF6666", "#005533", "#1199EE") # Colorblind-safe RGB + ... loc = np.array([offset, offset]) + ... for i, (axis, c) in enumerate(zip((ax.xaxis, ax.yaxis, ax.zaxis), + ... colors)): + ... axlabel = axis.axis_name + ... axis.set_label_text(axlabel) + ... axis.label.set_color(c) + ... axis.line.set_color(c) + ... axis.set_tick_params(colors=c) + ... line = np.zeros((2, 3)) + ... line[1, i] = scale + ... line_rot = r.apply(line) + ... line_plot = line_rot + loc + ... ax.plot(line_plot[:, 0], line_plot[:, 1], line_plot[:, 2], c) + ... text_loc = line[1]*1.2 + ... text_loc_rot = r.apply(text_loc) + ... text_plot = text_loc_rot + loc[0] + ... ax.text(*text_plot, axlabel.upper(), color=c, + ... va="center", ha="center") + ... ax.text(*offset, name, color="k", va="center", ha="center", + ... bbox={"fc": "w", "alpha": 0.8, "boxstyle": "circle"}) + + Create three rotations - the identity and two Euler rotations using + intrinsic and extrinsic conventions: + + >>> r0 = R.identity() + >>> r1 = R.from_euler("ZYX", [90, -30, 0], degrees=True) # intrinsic + >>> r2 = R.from_euler("zyx", [90, -30, 0], degrees=True) # extrinsic + + Add all three rotations to a single plot: + + >>> ax = plt.figure().add_subplot(projection="3d", proj_type="ortho") + >>> plot_rotated_axes(ax, r0, name="r0", offset=(0, 0, 0)) + >>> plot_rotated_axes(ax, r1, name="r1", offset=(3, 0, 0)) + >>> plot_rotated_axes(ax, r2, name="r2", offset=(6, 0, 0)) + >>> _ = ax.annotate( + ... "r0: Identity Rotation\\n" + ... "r1: Intrinsic Euler Rotation (ZYX)\\n" + ... "r2: Extrinsic Euler Rotation (zyx)", + ... xy=(0.6, 0.7), xycoords="axes fraction", ha="left" + ... ) + >>> ax.set(xlim=(-1.25, 7.25), ylim=(-1.25, 1.25), zlim=(-1.25, 1.25)) + >>> ax.set(xticks=range(-1, 8), yticks=[-1, 0, 1], zticks=[-1, 0, 1]) + >>> ax.set_aspect("equal", adjustable="box") + >>> ax.figure.set_size_inches(6, 5) + >>> plt.tight_layout() + + Show the plot: + + >>> plt.show() + + These examples serve as an overview into the `Rotation` class and highlight + major functionalities. For more thorough examples of the range of input and + output formats supported, consult the individual method's examples. + + """ + + # generic type compatibility with scipy-stubs + __class_getitem__: classmethod = classmethod(GenericAlias) + + def __init__( + self, + quat: ArrayLike, + normalize: bool = True, + copy: bool = True, + scalar_first: bool = False, + ): + xp = array_namespace(quat) + self._xp = xp + quat = _promote(quat, xp=xp) + if quat.shape[-1] != 4: + raise ValueError( + f"Expected `quat` to have shape (..., 4), got {quat.shape}." + ) + # Single NumPy quats or list of quats are accelerated by the cython backend. + # This backend needs inputs with fixed ndim, so we always expand to 2D and + # select the 0th element if quat was single to get the correct shape. For other + # frameworks and quaternion tensors we use the generic array API backend. + self._single = quat.ndim == 1 and is_numpy(xp) + if self._single: + quat = xpx.atleast_nd(quat, ndim=2, xp=xp) + self._backend = select_backend(xp, cython_compatible=quat.ndim < 3) + self._quat: Array = self._backend.from_quat( + quat, normalize=normalize, copy=copy, scalar_first=scalar_first + ) + + @staticmethod + def from_quat(quat: ArrayLike, *, scalar_first: bool = False) -> Rotation: + """Initialize from quaternions. + + Rotations in 3 dimensions can be represented using unit norm + quaternions [1]_. + + The 4 components of a quaternion are divided into a scalar part ``w`` + and a vector part ``(x, y, z)`` and can be expressed from the angle + ``theta`` and the axis ``n`` of a rotation as follows:: + + w = cos(theta / 2) + x = sin(theta / 2) * n_x + y = sin(theta / 2) * n_y + z = sin(theta / 2) * n_z + + There are 2 conventions to order the components in a quaternion: + + - scalar-first order -- ``(w, x, y, z)`` + - scalar-last order -- ``(x, y, z, w)`` + + The choice is controlled by `scalar_first` argument. + By default, it is False and the scalar-last order is assumed. + + Advanced users may be interested in the "double cover" of 3D space by + the quaternion representation [2]_. As of version 1.11.0, the + following subset (and only this subset) of operations on a `Rotation` + ``r`` corresponding to a quaternion ``q`` are guaranteed to preserve + the double cover property: ``r = Rotation.from_quat(q)``, + ``r.as_quat(canonical=False)``, ``r.inv()``, and composition using the + ``*`` operator such as ``r*r``. + + Parameters + ---------- + quat : array_like, shape (..., 4) + Each row is a (possibly non-unit norm) quaternion representing an + active rotation. Each quaternion will be normalized to unit norm. + scalar_first : bool, optional + Whether the scalar component goes first or last. + Default is False, i.e. the scalar-last order is assumed. + + Returns + ------- + rotation : `Rotation` instance + Object containing the rotations represented by input quaternions. + + References + ---------- + .. [1] https://en.wikipedia.org/wiki/Quaternions_and_spatial_rotation + .. [2] Hanson, Andrew J. "Visualizing quaternions." + Morgan Kaufmann Publishers Inc., San Francisco, CA. 2006. + + Examples + -------- + >>> from scipy.spatial.transform import Rotation as R + + A rotation can be initialized from a quaternion with the scalar-last + (default) or scalar-first component order as shown below: + + >>> r = R.from_quat([0, 0, 0, 1]) + >>> r.as_matrix() + array([[1., 0., 0.], + [0., 1., 0.], + [0., 0., 1.]]) + >>> r = R.from_quat([1, 0, 0, 0], scalar_first=True) + >>> r.as_matrix() + array([[1., 0., 0.], + [0., 1., 0.], + [0., 0., 1.]]) + + It is possible to initialize multiple rotations in a single object by + passing an N-dimensional array: + + >>> r = R.from_quat([[ + ... [1, 0, 0, 0], + ... [0, 0, 0, 1] + ... ]]) + >>> r.as_quat() + array([[[1., 0., 0., 0.], + [0., 0., 0., 1.]]]) + >>> r.as_quat().shape + (1, 2, 4) + + It is also possible to have a stack of a single rotation: + + >>> r = R.from_quat([[0, 0, 0, 1]]) + >>> r.as_quat() + array([[0., 0., 0., 1.]]) + >>> r.as_quat().shape + (1, 4) + + Quaternions are normalized before initialization. + + >>> r = R.from_quat([0, 0, 1, 1]) + >>> r.as_quat() + array([0. , 0. , 0.70710678, 0.70710678]) + """ + return Rotation(quat, normalize=True, scalar_first=scalar_first) + + @staticmethod + def from_matrix(matrix: ArrayLike, *, assume_valid: bool = False) -> Rotation: + """Initialize from rotation matrix. + + Rotations in 3 dimensions can be represented with 3 x 3 orthogonal + matrices [1]_. If the input is not orthogonal, an approximation is + created by orthogonalizing the input matrix using the method described + in [2]_, and then converting the orthogonal rotation matrices to + quaternions using the algorithm described in [3]_. Matrices must be + right-handed. + + Parameters + ---------- + matrix : array_like, shape (..., 3, 3) + A single matrix or an ND array of matrices, where the last two dimensions + contain the rotation matrices. + assume_valid : bool, optional + Must be False unless users can guarantee the input is a valid rotation + matrix, i.e. it is orthogonal, rows and columns have unit norm and the + determinant is 1. Setting this to True without ensuring these properties + is unsafe and will silently lead to incorrect results. If True, + normalization steps are skipped, which can improve runtime performance. + Default is False. + + Returns + ------- + rotation : `Rotation` instance + Object containing the rotations represented by the rotation + matrices. + + Notes + ----- + This function was called from_dcm before. + + .. versionadded:: 1.4.0 + + References + ---------- + .. [1] https://en.wikipedia.org/wiki/Rotation_matrix#In_three_dimensions + .. [2] https://en.wikipedia.org/wiki/Orthogonal_Procrustes_problem + .. [3] F. Landis Markley, "Unit Quaternion from Rotation Matrix", + Journal of guidance, control, and dynamics vol. 31.2, pp. + 440-442, 2008. + + Examples + -------- + >>> from scipy.spatial.transform import Rotation as R + >>> import numpy as np + + Initialize a single rotation: + + >>> r = R.from_matrix([ + ... [0, -1, 0], + ... [1, 0, 0], + ... [0, 0, 1]]) + >>> r.single + True + >>> r.as_matrix().shape + (3, 3) + + Initialize multiple rotations in a single object: + + >>> r = R.from_matrix([ + ... [ + ... [0, -1, 0], + ... [1, 0, 0], + ... [0, 0, 1], + ... ], + ... [ + ... [1, 0, 0], + ... [0, 0, -1], + ... [0, 1, 0], + ... ]]) + >>> r.as_matrix().shape + (2, 3, 3) + >>> r.single + False + >>> len(r) + 2 + + If input matrices are not special orthogonal (orthogonal with + determinant equal to +1), then a special orthogonal estimate is stored: + + >>> a = np.array([ + ... [0, -0.5, 0], + ... [0.5, 0, 0], + ... [0, 0, 0.5]]) + >>> np.linalg.det(a) + 0.125 + >>> r = R.from_matrix(a) + >>> matrix = r.as_matrix() + >>> matrix + array([[ 0., -1., 0.], + [ 1., 0., 0.], + [ 0., 0., 1.]]) + >>> np.linalg.det(matrix) + 1.0 + + It is also possible to have a stack containing a single rotation: + + >>> r = R.from_matrix([[ + ... [0, -1, 0], + ... [1, 0, 0], + ... [0, 0, 1]]]) + >>> r.as_matrix() + array([[[ 0., -1., 0.], + [ 1., 0., 0.], + [ 0., 0., 1.]]]) + >>> r.as_matrix().shape + (1, 3, 3) + + We can also create an N-dimensional array of rotations: + + >>> r = R.from_matrix(np.tile(np.eye(3), (2, 3, 1, 1))) + >>> r.shape + (2, 3) + + """ + xp = array_namespace(matrix) + matrix = _promote(matrix, xp=xp) + # Resulting quat will have 1 less dimension than matrix + backend = select_backend(xp, cython_compatible=matrix.ndim < 4) + quat = backend.from_matrix(matrix, assume_valid=assume_valid) + return Rotation._from_raw_quat(quat, xp=xp, backend=backend) + + @staticmethod + def from_rotvec(rotvec: ArrayLike, degrees: bool = False) -> Rotation: + """Initialize from rotation vectors. + + A rotation vector is a 3 dimensional vector which is co-directional to + the axis of rotation and whose norm gives the angle of rotation [1]_. + + Parameters + ---------- + rotvec : array_like, shape (..., 3) + A single vector or an ND array of vectors, where the last dimension + contains the rotation vectors. + degrees : bool, optional + If True, then the given magnitudes are assumed to be in degrees. + Default is False. + + .. versionadded:: 1.7.0 + + Returns + ------- + rotation : `Rotation` instance + Object containing the rotations represented by input rotation + vectors. + + References + ---------- + .. [1] https://en.wikipedia.org/wiki/Axis%E2%80%93angle_representation#Rotation_vector + + Examples + -------- + >>> from scipy.spatial.transform import Rotation as R + >>> import numpy as np + + Initialize a single rotation: + + >>> r = R.from_rotvec(np.pi/2 * np.array([0, 0, 1])) + >>> r.as_rotvec() + array([0. , 0. , 1.57079633]) + >>> r.as_rotvec().shape + (3,) + + Initialize a rotation in degrees, and view it in degrees: + + >>> r = R.from_rotvec(45 * np.array([0, 1, 0]), degrees=True) + >>> r.as_rotvec(degrees=True) + array([ 0., 45., 0.]) + + Initialize multiple rotations in one object: + + >>> r = R.from_rotvec([ + ... [0, 0, np.pi/2], + ... [np.pi/2, 0, 0]]) + >>> r.as_rotvec() + array([[0. , 0. , 1.57079633], + [1.57079633, 0. , 0. ]]) + >>> r.as_rotvec().shape + (2, 3) + + It is also possible to have a stack of a single rotation: + + >>> r = R.from_rotvec([[0, 0, np.pi/2]]) + >>> r.as_rotvec().shape + (1, 3) + + """ + xp = array_namespace(rotvec) + rotvec = _promote(rotvec, xp=xp) + backend = select_backend(xp, cython_compatible=rotvec.ndim < 3) + quat = backend.from_rotvec(rotvec, degrees=degrees) + return Rotation._from_raw_quat(quat, xp=xp, backend=backend) + + @staticmethod + def from_euler(seq: str, angles: ArrayLike, degrees: bool = False) -> Rotation: + """Initialize from Euler angles. + + Rotations in 3-D can be represented by a sequence of 3 + rotations around a sequence of axes. In theory, any three axes spanning + the 3-D Euclidean space are enough. In practice, the axes of rotation are + chosen to be the basis vectors. + + The three rotations can either be in a global frame of reference + (extrinsic) or in a body centred frame of reference (intrinsic), which + is attached to, and moves with, the object under rotation [1]_. + + Parameters + ---------- + seq : str + Specifies sequence of axes for rotations. Up to 3 characters + belonging to the set {'X', 'Y', 'Z'} for intrinsic rotations, or + {'x', 'y', 'z'} for extrinsic rotations. Extrinsic and intrinsic + rotations cannot be mixed in one function call. + angles : float or array_like, shape (..., [1 or 2 or 3]) + Euler angles specified in radians (`degrees` is False) or degrees + (`degrees` is True). + Each character in `seq` defines one axis around which `angles` turns. + The resulting rotation has the shape np.atleast_1d(angles).shape[:-1]. + Dimensionless angles are thus only valid for single character `seq`. + + degrees : bool, optional + If True, then the given angles are assumed to be in degrees. + Default is False. + + Returns + ------- + rotation : `Rotation` instance + Object containing the rotation represented by the sequence of + rotations around given axes with given angles. + + References + ---------- + .. [1] https://en.wikipedia.org/wiki/Euler_angles#Definition_by_intrinsic_rotations + + Examples + -------- + >>> from scipy.spatial.transform import Rotation as R + + Initialize a single rotation along a single axis: + + >>> r = R.from_euler('x', 90, degrees=True) + >>> r.as_quat().shape + (4,) + + Initialize a single rotation with a given axis sequence: + + >>> r = R.from_euler('zyx', [90, 45, 30], degrees=True) + >>> r.as_quat().shape + (4,) + + Initialize a stack with a single rotation around a single axis: + + >>> r = R.from_euler('x', [[90]], degrees=True) + >>> r.as_quat().shape + (1, 4) + + Initialize a stack with a single rotation with an axis sequence: + + >>> r = R.from_euler('zyx', [[90, 45, 30]], degrees=True) + >>> r.as_quat().shape + (1, 4) + + Initialize multiple elementary rotations in one object: + + >>> r = R.from_euler('x', [[90], [45], [30]], degrees=True) + >>> r.as_quat().shape + (3, 4) + + Initialize multiple rotations in one object: + + >>> r = R.from_euler('zyx', [[90, 45, 30], [35, 45, 90]], degrees=True) + >>> r.as_quat().shape + (2, 4) + + """ + xp = array_namespace(angles) + angles = _promote(angles, xp=xp) + backend = select_backend(xp, cython_compatible=angles.ndim < 3) + quat = backend.from_euler(seq, angles, degrees=degrees) + return Rotation._from_raw_quat(quat, xp=xp, backend=backend) + + @staticmethod + def from_davenport( + axes: ArrayLike, + order: str, + angles: ArrayLike | float, + degrees: bool = False, + ) -> Rotation: + """Initialize from Davenport angles. + + Rotations in 3-D can be represented by a sequence of 3 + rotations around a sequence of axes. + + The three rotations can either be in a global frame of reference + (extrinsic) or in a body centred frame of reference (intrinsic), which + is attached to, and moves with, the object under rotation [1]_. + + For both Euler angles and Davenport angles, consecutive axes must + be are orthogonal (``axis2`` is orthogonal to both ``axis1`` and + ``axis3``). For Euler angles, there is an additional relationship + between ``axis1`` or ``axis3``, with two possibilities: + + - ``axis1`` and ``axis3`` are also orthogonal (asymmetric sequence) + - ``axis1 == axis3`` (symmetric sequence) + + For Davenport angles, this last relationship is relaxed [2]_, and only + the consecutive orthogonal axes requirement is maintained. + + Parameters + ---------- + axes : array_like, shape (3,) or (..., [1 or 2 or 3], 3) + Axis of rotation, if one dimensional. If two or more dimensional, describes + the sequence of axes for rotations, where each axes[..., i, :] is the ith + axis. If more than one axis is given, then the second axis must be + orthogonal to both the first and third axes. + order : str + If it is equal to 'e' or 'extrinsic', the sequence will be + extrinsic. If it is equal to 'i' or 'intrinsic', sequence + will be treated as intrinsic. + angles : float or array_like, shape (..., [1 or 2 or 3]) + Angles specified in radians (`degrees` is False) or degrees + (`degrees` is True). + Each angle i in the last dimension of `angles` turns around the corresponding + axis axis[..., i, :]. The resulting rotation has the shape + np.broadcast_shapes(np.atleast_2d(axes).shape[:-2], np.atleast_1d(angles).shape[:-1]) + Dimensionless angles are thus only valid for a single axis. + + degrees : bool, optional + If True, then the given angles are assumed to be in degrees. + Default is False. + + Returns + ------- + rotation : `Rotation` instance + Object containing the rotation represented by the sequence of + rotations around given axes with given angles. + + References + ---------- + .. [1] https://en.wikipedia.org/wiki/Euler_angles#Definition_by_intrinsic_rotations + .. [2] Shuster, Malcolm & Markley, Landis. (2003). Generalization of + the Euler Angles. Journal of the Astronautical Sciences. 51. 123-132. + 10.1007/BF03546304. + + Examples + -------- + >>> from scipy.spatial.transform import Rotation as R + + Davenport angles are a generalization of Euler angles, when we use the + canonical basis axes: + + >>> ex = [1, 0, 0] + >>> ey = [0, 1, 0] + >>> ez = [0, 0, 1] + + Initialize a single rotation with a given axis sequence: + + >>> axes = [ez, ey, ex] + >>> r = R.from_davenport(axes, 'extrinsic', [90, 0, 0], degrees=True) + >>> r.as_quat().shape + (4,) + + It is equivalent to Euler angles in this case: + + >>> r.as_euler('zyx', degrees=True) + array([90., 0., -0.]) + + Initialize multiple rotations in one object: + + >>> r = R.from_davenport(axes, 'extrinsic', [[90, 45, 30], [35, 45, 90]], degrees=True) + >>> r.as_quat().shape + (2, 4) + + Using only one or two axes is also possible: + + >>> r = R.from_davenport([ez, ex], 'extrinsic', [[90, 45], [35, 45]], degrees=True) + >>> r.as_quat().shape + (2, 4) + + Non-canonical axes are possible, and they do not need to be normalized, + as long as consecutive axes are orthogonal: + + >>> e1 = [2, 0, 0] + >>> e2 = [0, 1, 0] + >>> e3 = [1, 0, 1] + >>> axes = [e1, e2, e3] + >>> r = R.from_davenport(axes, 'extrinsic', [90, 45, 30], degrees=True) + >>> r.as_quat() + [ 0.701057, 0.430459, -0.092296, 0.560986] + """ # noqa: E501 + xp = array_namespace(axes) + axes, angles = _promote(axes, angles, xp=xp) + cython_compatible = axes.ndim < 3 and angles.ndim < 2 # type:ignore[union-attr] + backend = select_backend(xp, cython_compatible=cython_compatible) + quat = backend.from_davenport(axes, order, angles, degrees) + return Rotation._from_raw_quat(quat, xp=xp, backend=backend) + + @staticmethod + def from_mrp(mrp: ArrayLike) -> Rotation: + """Initialize from Modified Rodrigues Parameters (MRPs). + + MRPs are a 3 dimensional vector co-directional to the axis of rotation and whose + magnitude is equal to ``tan(theta / 4)``, where ``theta`` is the angle of + rotation (in radians) [1]_. + + MRPs have a singularity at 360 degrees which can be avoided by ensuring the + angle of rotation does not exceed 180 degrees, i.e. switching the direction of + the rotation when it is past 180 degrees. + + Parameters + ---------- + mrp : array_like, shape (..., 3) + A single vector or an ND array of vectors, where the last dimension + contains the rotation parameters. + + Returns + ------- + rotation : `Rotation` instance + Object containing the rotations represented by input MRPs. + + Notes + ----- + + .. versionadded:: 1.6.0 + + References + ---------- + .. [1] Shuster, M. D. "A Survey of Attitude Representations", + The Journal of Astronautical Sciences, Vol. 41, No.4, 1993, + pp. 475-476 + + Examples + -------- + >>> from scipy.spatial.transform import Rotation as R + >>> import numpy as np + + Initialize a single rotation: + + >>> r = R.from_mrp([0, 0, 1]) + >>> r.as_euler('xyz', degrees=True) + array([0. , 0. , 180. ]) + >>> r.as_euler('xyz').shape + (3,) + + Initialize multiple rotations in one object: + + >>> r = R.from_mrp([ + ... [0, 0, 1], + ... [1, 0, 0]]) + >>> r.as_euler('xyz', degrees=True) + array([[0. , 0. , 180. ], + [180.0 , 0. , 0. ]]) + >>> r.as_euler('xyz').shape + (2, 3) + + It is also possible to have a stack of a single rotation: + + >>> r = R.from_mrp([[0, 0, np.pi/2]]) + >>> r.as_euler('xyz').shape + (1, 3) + + """ + xp = array_namespace(mrp) + mrp = _promote(mrp, xp=xp) + backend = select_backend(xp, cython_compatible=mrp.ndim < 3) + quat = backend.from_mrp(mrp) + return Rotation._from_raw_quat(quat, xp=xp, backend=backend) + + def as_quat(self, canonical: bool = False, *, scalar_first: bool = False) -> Array: + """Represent as quaternions. + + Rotations in 3 dimensions can be represented using unit norm + quaternions [1]_. + + The 4 components of a quaternion are divided into a scalar part ``w`` + and a vector part ``(x, y, z)`` and can be expressed from the angle + ``theta`` and the axis ``n`` of a rotation as follows:: + + w = cos(theta / 2) + x = sin(theta / 2) * n_x + y = sin(theta / 2) * n_y + z = sin(theta / 2) * n_z + + There are 2 conventions to order the components in a quaternion: + + - scalar-first order -- ``(w, x, y, z)`` + - scalar-last order -- ``(x, y, z, w)`` + + The choice is controlled by `scalar_first` argument. + By default, it is False and the scalar-last order is used. + + The mapping from quaternions to rotations is + two-to-one, i.e. quaternions ``q`` and ``-q``, where ``-q`` simply + reverses the sign of each component, represent the same spatial + rotation. + + Parameters + ---------- + canonical : `bool`, default False + Whether to map the redundant double cover of rotation space to a + unique "canonical" single cover. If True, then the quaternion is + chosen from {q, -q} such that the w term is positive. If the w term + is 0, then the quaternion is chosen such that the first nonzero + term of the x, y, and z terms is positive. + scalar_first : bool, optional + Whether the scalar component goes first or last. + Default is False, i.e. the scalar-last order is used. + + Returns + ------- + quat : `numpy.ndarray`, shape (..., 4) + Shape depends on shape of inputs used for initialization. + + References + ---------- + .. [1] https://en.wikipedia.org/wiki/Quaternions_and_spatial_rotation + + Examples + -------- + >>> from scipy.spatial.transform import Rotation as R + >>> import numpy as np + + A rotation can be represented as a quaternion with either scalar-last + (default) or scalar-first component order. + This is shown for a single rotation: + + >>> r = R.from_matrix(np.eye(3)) + >>> r.as_quat() + array([0., 0., 0., 1.]) + >>> r.as_quat(scalar_first=True) + array([1., 0., 0., 0.]) + + The resulting shape of the quaternion is always the shape of the Rotation + object with an added last dimension of size 4. E.g. when the `Rotation` object + contains an N-dimensional array (N, M, K) of rotations, the result will be a + 4-dimensional array: + + >>> r = R.from_rotvec(np.ones((2, 3, 4, 3))) + >>> r.as_quat().shape + (2, 3, 4, 4) + + Quaternions can be mapped from a redundant double cover of the + rotation space to a canonical representation with a positive w term. + + >>> r = R.from_quat([0, 0, 0, -1]) + >>> r.as_quat() + array([0. , 0. , 0. , -1.]) + >>> r.as_quat(canonical=True) + array([0. , 0. , 0. , 1.]) + """ + quat = self._backend.as_quat( + self._quat, canonical=canonical, scalar_first=scalar_first + ) + if self._single: + return quat[0, ...] + return quat + + def as_matrix(self) -> Array: + """Represent as rotation matrix. + + 3D rotations can be represented using rotation matrices, which + are 3 x 3 real orthogonal matrices with determinant equal to +1 [1]_. + + Returns + ------- + matrix : ndarray, shape (..., 3) + Shape depends on shape of inputs used for initialization. + + Notes + ----- + This function was called as_dcm before. + + .. versionadded:: 1.4.0 + + References + ---------- + .. [1] https://en.wikipedia.org/wiki/Rotation_matrix#In_three_dimensions + + Examples + -------- + >>> from scipy.spatial.transform import Rotation as R + >>> import numpy as np + + Represent a single rotation: + + >>> r = R.from_rotvec([0, 0, np.pi/2]) + >>> r.as_matrix() + array([[ 2.22044605e-16, -1.00000000e+00, 0.00000000e+00], + [ 1.00000000e+00, 2.22044605e-16, 0.00000000e+00], + [ 0.00000000e+00, 0.00000000e+00, 1.00000000e+00]]) + >>> r.as_matrix().shape + (3, 3) + + Represent a stack with a single rotation: + + >>> r = R.from_quat([[1, 1, 0, 0]]) + >>> r.as_matrix() + array([[[ 0., 1., 0.], + [ 1., 0., 0.], + [ 0., 0., -1.]]]) + >>> r.as_matrix().shape + (1, 3, 3) + + Represent multiple rotations: + + >>> r = R.from_rotvec([[np.pi/2, 0, 0], [0, 0, np.pi/2]]) + >>> r.as_matrix() + array([[[ 1.00000000e+00, 0.00000000e+00, 0.00000000e+00], + [ 0.00000000e+00, 2.22044605e-16, -1.00000000e+00], + [ 0.00000000e+00, 1.00000000e+00, 2.22044605e-16]], + [[ 2.22044605e-16, -1.00000000e+00, 0.00000000e+00], + [ 1.00000000e+00, 2.22044605e-16, 0.00000000e+00], + [ 0.00000000e+00, 0.00000000e+00, 1.00000000e+00]]]) + >>> r.as_matrix().shape + (2, 3, 3) + + """ + matrix = self._backend.as_matrix(self._quat) + if self._single: + return matrix[0, ...] + return matrix + + def as_rotvec(self, degrees: bool = False) -> Array: + """Represent as rotation vectors. + + A rotation vector is a 3 dimensional vector which is co-directional to + the axis of rotation and whose norm gives the angle of rotation [1]_. + + Parameters + ---------- + degrees : bool, optional + Returned magnitudes are in degrees if this flag is True, else they are + in radians. Default is False. + + .. versionadded:: 1.7.0 + + Returns + ------- + rotvec : ndarray, shape (..., 3) + Shape depends on shape of inputs used for initialization. + + References + ---------- + .. [1] https://en.wikipedia.org/wiki/Axis%E2%80%93angle_representation#Rotation_vector + + Examples + -------- + >>> from scipy.spatial.transform import Rotation as R + >>> import numpy as np + + Represent a single rotation: + + >>> r = R.from_euler('z', 90, degrees=True) + >>> r.as_rotvec() + array([0. , 0. , 1.57079633]) + >>> r.as_rotvec().shape + (3,) + + Represent a rotation in degrees: + + >>> r = R.from_euler('YX', (-90, -90), degrees=True) + >>> s = r.as_rotvec(degrees=True) + >>> s + array([-69.2820323, -69.2820323, -69.2820323]) + >>> np.linalg.norm(s) + 120.00000000000001 + + Represent a stack with a single rotation: + + >>> r = R.from_quat([[0, 0, 1, 1]]) + >>> r.as_rotvec() + array([[0. , 0. , 1.57079633]]) + >>> r.as_rotvec().shape + (1, 3) + + Represent multiple rotations in a single object: + + >>> r = R.from_quat([[0, 0, 1, 1], [1, 1, 0, 1]]) + >>> r.as_rotvec() + array([[0. , 0. , 1.57079633], + [1.35102172, 1.35102172, 0. ]]) + >>> r.as_rotvec().shape + (2, 3) + + """ + rotvec = self._backend.as_rotvec(self._quat, degrees=degrees) + if self._single: + return rotvec[0, ...] + return rotvec + + def as_euler( + self, seq: str, degrees: bool = False, *, suppress_warnings: bool = False + ) -> Array: + """Represent as Euler angles. + + Any orientation can be expressed as a composition of 3 elementary + rotations. Once the axis sequence has been chosen, Euler angles define + the angle of rotation around each respective axis [1]_. + + The algorithm from [2]_ has been used to calculate Euler angles for the + rotation about a given sequence of axes. + + Euler angles suffer from the problem of gimbal lock [3]_, where the + representation loses a degree of freedom and it is not possible to + determine the first and third angles uniquely. In this case, + a warning is raised (unless the ``suppress_warnings`` option is used), + and the third angle is set to zero. Note however that the returned + angles still represent the correct rotation. + + Parameters + ---------- + seq : str, length 3 + 3 characters belonging to the set {'X', 'Y', 'Z'} for intrinsic + rotations, or {'x', 'y', 'z'} for extrinsic rotations [1]_. + Adjacent axes cannot be the same. + Extrinsic and intrinsic rotations cannot be mixed in one function + call. + degrees : bool, optional + Returned angles are in degrees if this flag is True, else they are + in radians. Default is False. + suppress_warnings : bool, optional + Disable warnings about gimbal lock. Default is False. + + Returns + ------- + angles : ndarray, shape (..., 3) + Shape depends on shape of inputs used to initialize object. + The returned angles are in the range: + + - First angle belongs to [-180, 180] degrees (both inclusive) + - Third angle belongs to [-180, 180] degrees (both inclusive) + - Second angle belongs to: + + - [-90, 90] degrees if all axes are different (like xyz) + - [0, 180] degrees if first and third axes are the same + (like zxz) + + References + ---------- + .. [1] https://en.wikipedia.org/wiki/Euler_angles#Definition_by_intrinsic_rotations + .. [2] Bernardes E, Viollet S (2022) Quaternion to Euler angles + conversion: A direct, general and computationally efficient + method. PLoS ONE 17(11): e0276302. + :doi:`10.1371/journal.pone.0276302`. + .. [3] https://en.wikipedia.org/wiki/Gimbal_lock#In_applied_mathematics + + Examples + -------- + >>> from scipy.spatial.transform import Rotation as R + >>> import numpy as np + + Represent a single rotation: + + >>> r = R.from_rotvec([0, 0, np.pi/2]) + >>> r.as_euler('zxy', degrees=True) + array([90., 0., 0.]) + >>> r.as_euler('zxy', degrees=True).shape + (3,) + + Represent a stack of single rotation: + + >>> r = R.from_rotvec([[0, 0, np.pi/2]]) + >>> r.as_euler('zxy', degrees=True) + array([[90., 0., 0.]]) + >>> r.as_euler('zxy', degrees=True).shape + (1, 3) + + Represent multiple rotations in a single object: + + >>> r = R.from_rotvec([ + ... [0, 0, np.pi/2], + ... [0, -np.pi/3, 0], + ... [np.pi/4, 0, 0]]) + >>> r.as_euler('zxy', degrees=True) + array([[ 90., 0., 0.], + [ 0., 0., -60.], + [ 0., 45., 0.]]) + >>> r.as_euler('zxy', degrees=True).shape + (3, 3) + + """ + euler = self._backend.as_euler( + self._quat, seq, degrees=degrees, suppress_warnings=suppress_warnings + ) + if self._single: + return euler[0, ...] + return euler + + def as_davenport( + self, + axes: ArrayLike, + order: str, + degrees: bool = False, + *, + suppress_warnings: bool = False, + ) -> Array: + """Represent as Davenport angles. + + Any orientation can be expressed as a composition of 3 elementary + rotations. + + For both Euler angles and Davenport angles, consecutive axes must + be are orthogonal (``axis2`` is orthogonal to both ``axis1`` and + ``axis3``). For Euler angles, there is an additional relationship + between ``axis1`` or ``axis3``, with two possibilities: + + - ``axis1`` and ``axis3`` are also orthogonal (asymmetric sequence) + - ``axis1 == axis3`` (symmetric sequence) + + For Davenport angles, this last relationship is relaxed [1]_, and only + the consecutive orthogonal axes requirement is maintained. + + A slightly modified version of the algorithm from [2]_ has been used to + calculate Davenport angles for the rotation about a given sequence of + axes. + + Davenport angles, just like Euler angles, suffer from the problem of + gimbal lock [3]_, where the representation loses a degree of freedom + and it is not possible to determine the first and third angles + uniquely. In this case, a warning is raised (unless the + ``suppress_warnings`` option is used), and the third angle is set + to zero. Note however that the returned angles still represent the + correct rotation. + + Parameters + ---------- + axes : array_like, shape (..., [1 or 2 or 3], 3) or (..., 3) + Axis of rotation, if one dimensional. If N dimensional, describes the + sequence of axes for rotations, where each axes[..., i, :] is the ith + axis. If more than one axis is given, then the second axis must be + orthogonal to both the first and third axes. + order : str + If it belongs to the set {'e', 'extrinsic'}, the sequence will be + extrinsic. If it belongs to the set {'i', 'intrinsic'}, sequence + will be treated as intrinsic. + degrees : bool, optional + Returned angles are in degrees if this flag is True, else they are + in radians. Default is False. + suppress_warnings : bool, optional + Disable warnings about gimbal lock. Default is False. + + Returns + ------- + angles : ndarray, shape (..., 3) + Shape depends on shape of inputs used to initialize object. + The returned angles are in the range: + + - First angle belongs to [-180, 180] degrees (both inclusive) + - Third angle belongs to [-180, 180] degrees (both inclusive) + - Second angle belongs to a set of size 180 degrees, + given by: ``[-abs(lambda), 180 - abs(lambda)]``, where ``lambda`` + is the angle between the first and third axes. + + References + ---------- + .. [1] Shuster, Malcolm & Markley, Landis. (2003). Generalization of + the Euler Angles. Journal of the Astronautical Sciences. 51. 123-132. + 10.1007/BF03546304. + .. [2] Bernardes E, Viollet S (2022) Quaternion to Euler angles + conversion: A direct, general and computationally efficient method. + PLoS ONE 17(11): e0276302. 10.1371/journal.pone.0276302 + .. [3] https://en.wikipedia.org/wiki/Gimbal_lock#In_applied_mathematics + + Examples + -------- + >>> from scipy.spatial.transform import Rotation as R + >>> import numpy as np + + Davenport angles are a generalization of Euler angles, when we use the + canonical basis axes: + + >>> ex = [1, 0, 0] + >>> ey = [0, 1, 0] + >>> ez = [0, 0, 1] + + Represent a single rotation: + + >>> r = R.from_rotvec([0, 0, np.pi/2]) + >>> r.as_davenport([ez, ex, ey], 'extrinsic', degrees=True) + array([90., 0., 0.]) + >>> r.as_euler('zxy', degrees=True) + array([90., 0., 0.]) + >>> r.as_davenport([ez, ex, ey], 'extrinsic', degrees=True).shape + (3,) + + Represent a stack of single rotation: + + >>> r = R.from_rotvec([[0, 0, np.pi/2]]) + >>> r.as_davenport([ez, ex, ey], 'extrinsic', degrees=True) + array([[90., 0., 0.]]) + >>> r.as_davenport([ez, ex, ey], 'extrinsic', degrees=True).shape + (1, 3) + + Represent multiple rotations in a single object: + + >>> r = R.from_rotvec([ + ... [0, 0, 90], + ... [45, 0, 0]], degrees=True) + >>> r.as_davenport([ez, ex, ey], 'extrinsic', degrees=True) + array([[90., 0., 0.], + [ 0., 45., 0.]]) + >>> r.as_davenport([ez, ex, ey], 'extrinsic', degrees=True).shape + (2, 3) + """ + axes = self._xp.asarray( + axes, dtype=self._quat.dtype, device=xp_device(self._quat) + ) + davenport = self._backend.as_davenport( + self._quat, axes, order, degrees, suppress_warnings=suppress_warnings + ) + if self._single: + return davenport[0, ...] + return davenport + + def as_mrp(self) -> Array: + """Represent as Modified Rodrigues Parameters (MRPs). + + MRPs are a 3 dimensional vector co-directional to the axis of rotation and whose + magnitude is equal to ``tan(theta / 4)``, where ``theta`` is the angle of + rotation (in radians) [1]_. + + MRPs have a singularity at 360 degrees which can be avoided by ensuring the + angle of rotation does not exceed 180 degrees, i.e. switching the direction of + the rotation when it is past 180 degrees. This function will always return MRPs + corresponding to a rotation of less than or equal to 180 degrees. + + Returns + ------- + mrps : ndarray, shape (..., 3) + Shape depends on shape of inputs used for initialization. + + Notes + ----- + + .. versionadded:: 1.6.0 + + References + ---------- + .. [1] Shuster, M. D. "A Survey of Attitude Representations", + The Journal of Astronautical Sciences, Vol. 41, No.4, 1993, + pp. 475-476 + + Examples + -------- + >>> from scipy.spatial.transform import Rotation as R + >>> import numpy as np + + Represent a single rotation: + + >>> r = R.from_rotvec([0, 0, np.pi]) + >>> r.as_mrp() + array([0. , 0. , 1. ]) + >>> r.as_mrp().shape + (3,) + + Represent a stack with a single rotation: + + >>> r = R.from_euler('xyz', [[180, 0, 0]], degrees=True) + >>> r.as_mrp() + array([[1. , 0. , 0. ]]) + >>> r.as_mrp().shape + (1, 3) + + Represent multiple rotations: + + >>> r = R.from_rotvec([[np.pi/2, 0, 0], [0, 0, np.pi/2]]) + >>> r.as_mrp() + array([[0.41421356, 0. , 0. ], + [0. , 0. , 0.41421356]]) + >>> r.as_mrp().shape + (2, 3) + + """ + mrp = self._backend.as_mrp(self._quat) + if self._single: + return mrp[0, ...] + return mrp + + @staticmethod + def concatenate(rotations: Rotation | Sequence[Rotation]) -> Rotation: + """Concatenate a sequence of `Rotation` objects into a single object. + + This is useful if you want to, for example, take the mean of a set of + rotations and need to pack them into a single object to do so. + + Parameters + ---------- + rotations : sequence of `Rotation` objects + The rotations to concatenate. If a single `Rotation` object is + passed in, a copy is returned. + + Returns + ------- + concatenated : `Rotation` instance + The concatenated rotations. + + Notes + ----- + .. versionadded:: 1.8.0 + + Examples + -------- + >>> from scipy.spatial.transform import Rotation as R + >>> r1 = R.from_rotvec([0, 0, 1]) + >>> r2 = R.from_rotvec([0, 0, 2]) + >>> rc = R.concatenate([r1, r2]) + >>> rc.as_rotvec() + array([[0., 0., 1.], + [0., 0., 2.]]) + >>> rc.mean().as_rotvec() + array([0., 0., 1.5]) + + Concatenation of a split rotation recovers the original object. + + >>> rs = [r for r in rc] + >>> R.concatenate(rs).as_rotvec() + array([[0., 0., 1.], + [0., 0., 2.]]) + + Note that it may be simpler to create the desired rotations by passing + in a single list of the data during initialization, rather then by + concatenating: + + >>> R.from_rotvec([[0, 0, 1], [0, 0, 2]]).as_rotvec() + array([[0., 0., 1.], + [0., 0., 2.]]) + + """ + if isinstance(rotations, Rotation): + return Rotation(rotations.as_quat(), normalize=False, copy=True) + if not all(isinstance(x, Rotation) for x in rotations): + raise TypeError("input must contain Rotation objects only") + + xp = array_namespace(rotations[0].as_quat()) + quats = xp.concat( + [xpx.atleast_nd(x.as_quat(), ndim=2, xp=xp) for x in rotations] + ) + return Rotation._from_raw_quat(quats, xp=xp) + + def apply(self, vectors: ArrayLike, inverse: bool = False) -> Array: + """Apply this rotation to a set of vectors. + + If the original frame rotates to the final frame by this rotation, then + its application to a vector can be seen in two ways: + + - As a projection of vector components expressed in the final frame + to the original frame. + - As the physical rotation of a vector being glued to the original + frame as it rotates. In this case the vector components are + expressed in the original frame before and after the rotation. + + In terms of rotation matrices, this application is the same as + ``(self.as_matrix() @ vectors[..., np.newaxis])[..., 0]``. + For a single rotation, this is the same as + ``vectors @ self.as_matrix().T``. + + Parameters + ---------- + vectors : array_like, shape (..., 3) + Each `vectors[..., :]` represents a vector in 3D space. The shape of + rotations and shape of vectors given must follow standard numpy + broadcasting rules: either one of them equals unity or they both + equal each other. + inverse : bool, optional + If True then the inverse of the rotation(s) is applied to the input + vectors. Default is False. + + Returns + ------- + rotated_vectors : ndarray, shape (..., 3) + Result of applying rotation on input vectors. + Shape is determined according to numpy broadcasting rules. I.e., the result + will have the shape `np.broadcast_shapes(r.shape, v.shape[:-1]) + (3,)` + + Examples + -------- + >>> from scipy.spatial.transform import Rotation as R + >>> import numpy as np + + Single rotation applied on a single vector: + + >>> vector = np.array([1, 0, 0]) + >>> r = R.from_rotvec([0, 0, np.pi/2]) + >>> r.as_matrix() + array([[ 2.22044605e-16, -1.00000000e+00, 0.00000000e+00], + [ 1.00000000e+00, 2.22044605e-16, 0.00000000e+00], + [ 0.00000000e+00, 0.00000000e+00, 1.00000000e+00]]) + >>> r.apply(vector) + array([2.22044605e-16, 1.00000000e+00, 0.00000000e+00]) + >>> r.apply(vector).shape + (3,) + + Single rotation applied on multiple vectors: + + >>> vectors = np.array([ + ... [1, 0, 0], + ... [1, 2, 3]]) + >>> r = R.from_rotvec([0, 0, np.pi/4]) + >>> r.as_matrix() + array([[ 0.70710678, -0.70710678, 0. ], + [ 0.70710678, 0.70710678, 0. ], + [ 0. , 0. , 1. ]]) + >>> r.apply(vectors) + array([[ 0.70710678, 0.70710678, 0. ], + [-0.70710678, 2.12132034, 3. ]]) + >>> r.apply(vectors).shape + (2, 3) + + Multiple rotations on a single vector: + + >>> r = R.from_rotvec([[0, 0, np.pi/4], [np.pi/2, 0, 0]]) + >>> vector = np.array([1,2,3]) + >>> r.as_matrix() + array([[[ 7.07106781e-01, -7.07106781e-01, 0.00000000e+00], + [ 7.07106781e-01, 7.07106781e-01, 0.00000000e+00], + [ 0.00000000e+00, 0.00000000e+00, 1.00000000e+00]], + [[ 1.00000000e+00, 0.00000000e+00, 0.00000000e+00], + [ 0.00000000e+00, 2.22044605e-16, -1.00000000e+00], + [ 0.00000000e+00, 1.00000000e+00, 2.22044605e-16]]]) + >>> r.apply(vector) + array([[-0.70710678, 2.12132034, 3. ], + [ 1. , -3. , 2. ]]) + >>> r.apply(vector).shape + (2, 3) + + Multiple rotations on multiple vectors. Each rotation is applied on the + corresponding vector: + + >>> r = R.from_euler('zxy', [ + ... [0, 0, 90], + ... [45, 30, 60]], degrees=True) + >>> vectors = [ + ... [1, 2, 3], + ... [1, 0, -1]] + >>> r.apply(vectors) + array([[ 3. , 2. , -1. ], + [-0.09026039, 1.11237244, -0.86860844]]) + >>> r.apply(vectors).shape + (2, 3) + + Broadcasting rules apply: + + >>> r = R.from_rotvec(np.tile([0, 0, np.pi/4], (5, 1, 4, 1))) + >>> vectors = np.ones((3, 4, 3)) + >>> r.shape, vectors.shape + ((5, 1, 4), (3, 4, 3)) + >>> r.apply(vectors).shape + (5, 3, 4, 3) + + It is also possible to apply the inverse rotation: + + >>> r = R.from_euler('zxy', [ + ... [0, 0, 90], + ... [45, 30, 60]], degrees=True) + >>> vectors = [ + ... [1, 2, 3], + ... [1, 0, -1]] + >>> r.apply(vectors, inverse=True) + array([[-3. , 2. , 1. ], + [ 1.09533535, -0.8365163 , 0.3169873 ]]) + + """ + vectors = self._xp.asarray( + vectors, device=xp_device(self._quat), dtype=self._quat.dtype + ) + single_vector = vectors.ndim == 1 + # Numpy optimization: The Cython backend typing requires us to have fixed + # dimensions, so for the Numpy case we always broadcast the vector to 2D. + if vectors.shape[-1] != 3: + raise ValueError(f"Expected input of shape (..., 3), got {vectors.shape}.") + if is_numpy(self._xp): + vectors = xpx.atleast_nd(vectors, ndim=2, xp=self._xp) + cython_compatible = self._quat.ndim < 3 and vectors.ndim < 3 + backend = select_backend(self._xp, cython_compatible=cython_compatible) + result = backend.apply(self._quat, vectors, inverse=inverse) + if self._single and single_vector: + return result[0, ...] + return result + + def __mul__(self, other: Rotation) -> Rotation | NotImplementedType: + """Compose this rotation with the other. + + If `p` and `q` are two rotations, then the composition of 'q followed + by p' is equivalent to `p * q`. In terms of rotation matrices, + the composition can be expressed as + ``p.as_matrix() @ q.as_matrix()``. + + Parameters + ---------- + other : `Rotation` instance + Object containing the rotations to be composed with this one. Note + that rotation compositions are not commutative, so ``p * q`` is + generally different from ``q * p``. + + Returns + ------- + composition : `Rotation` instance + This function supports composition of multiple rotations at a time. + Composition follows standard numpy broadcasting rules. The resulting + `Rotation` object will have the shape + `np.broadcast_shapes(p.shape, q.shape)`. In dimensions with size > 1, + rotations are composed with matching indices. In dimensions with only + one rotation, the single rotation is composed with each rotation in the + other object. + + Examples + -------- + >>> from scipy.spatial.transform import Rotation as R + >>> import numpy as np + + Composition of two single rotations: + + >>> p = R.from_quat([0, 0, 1, 1]) + >>> q = R.from_quat([1, 0, 0, 1]) + >>> p.as_matrix() + array([[ 0., -1., 0.], + [ 1., 0., 0.], + [ 0., 0., 1.]]) + >>> q.as_matrix() + array([[ 1., 0., 0.], + [ 0., 0., -1.], + [ 0., 1., 0.]]) + >>> r = p * q + >>> r.as_matrix() + array([[0., 0., 1.], + [1., 0., 0.], + [0., 1., 0.]]) + + Composition of two objects containing equal number of rotations: + + >>> p = R.from_quat([[0, 0, 1, 1], [1, 0, 0, 1]]) + >>> q = R.from_rotvec([[np.pi/4, 0, 0], [-np.pi/4, 0, np.pi/4]]) + >>> p.as_quat() + array([[0. , 0. , 0.70710678, 0.70710678], + [0.70710678, 0. , 0. , 0.70710678]]) + >>> q.as_quat() + array([[ 0.38268343, 0. , 0. , 0.92387953], + [-0.37282173, 0. , 0.37282173, 0.84971049]]) + >>> r = p * q + >>> r.as_quat() + array([[ 0.27059805, 0.27059805, 0.65328148, 0.65328148], + [ 0.33721128, -0.26362477, 0.26362477, 0.86446082]]) + + Broadcasting rules apply: + >>> p = R.from_quat(np.tile(np.array([0, 0, 1, 1]), (5, 1, 1))) + >>> q = R.from_quat(np.tile(np.array([1, 0, 0, 1]), (1, 6, 1))) + >>> p.shape, q.shape + ((5, 1), (1, 6)) + >>> r = p * q + >>> r.shape + (5, 6) + """ + # Check that other is a Rotation object. We want to return NotImplemented + # instead of raising an error to allow other types to implement __rmul__. + # Python will then automatically try to delegate the multiplication to the + # other type. + # See https://github.com/scipy/scipy/issues/21541 + if not isinstance(other, Rotation): + return NotImplemented + if not broadcastable(self._quat.shape, other._quat.shape): + raise ValueError( + f"Cannot broadcast {self._quat.shape[:-1]} rotations in " + f"first to {other._quat.shape[:-1]} rotations in second object." + ) + cython_compatible = self._quat.ndim < 3 and other._quat.ndim < 3 + backend = select_backend(self._xp, cython_compatible=cython_compatible) + quat = backend.compose_quat(self._quat, other._quat) + if self._single and other._single: + quat = quat[0] + return Rotation(quat, normalize=True, copy=False) + + def __pow__(self, n: float | Array, modulus: None = None) -> Rotation: + """Compose this rotation with itself `n` times. + + Composition of a rotation ``p`` with itself can be extended to + non-integer ``n`` by considering the power ``n`` to be a scale factor + applied to the angle of rotation about the rotation's fixed axis. The + expression ``q = p ** n`` can also be expressed as + ``q = Rotation.from_rotvec(n * p.as_rotvec())``. + + If ``n`` is negative, then the rotation is inverted before the power + is applied. In other words, ``p ** -abs(n) == p.inv() ** abs(n)``. + + Parameters + ---------- + n : float | Array + The number of times to compose the rotation with itself. If `n` is + an array, then it must be 0d or 1d with shape (1,). + modulus : None + This overridden argument is not applicable to Rotations and must be + ``None``. + + Returns + ------- + power : `Rotation` instance + The resulting rotation will be of the same shape as the original rotation + object. Each element of the output is the corresponding element of the + input rotation raised to the power of ``n``. + + Notes + ----- + For example, a power of 2 will double the angle of rotation, and a + power of 0.5 will halve the angle. There are three notable cases: if + ``n == 1`` then the original rotation is returned, if ``n == 0`` + then the identity rotation is returned, and if ``n == -1`` then + ``p.inv()`` is returned. + + Note that fractional powers ``n`` which effectively take a root of + rotation, do so using the shortest path smallest representation of that + angle (the principal root). This means that powers of ``n`` and ``1/n`` + are not necessarily inverses of each other. For example, a 0.5 power of + a +240 degree rotation will be calculated as the 0.5 power of a -120 + degree rotation, with the result being a rotation of -60 rather than + +120 degrees. + + Examples + -------- + >>> from scipy.spatial.transform import Rotation as R + + Raising a rotation to a power: + + >>> p = R.from_rotvec([1, 0, 0]) + >>> q = p ** 2 + >>> q.as_rotvec() + array([2., 0., 0.]) + >>> r = p ** 0.5 + >>> r.as_rotvec() + array([0.5, 0., 0.]) + + Inverse powers do not necessarily cancel out: + + >>> p = R.from_rotvec([0, 0, 120], degrees=True) + >>> ((p ** 2) ** 0.5).as_rotvec(degrees=True) + array([ -0., -0., -60.]) + + """ + if modulus is not None: + raise NotImplementedError("modulus not supported") + quat = self._backend.pow(self._quat, n) + if self._single: + quat = quat[0] + return Rotation._from_raw_quat(quat, xp=self._xp, backend=self._backend) + + def inv(self) -> Rotation: + """Invert this rotation. + + Composition of a rotation with its inverse results in an identity + transformation. + + Returns + ------- + inverse : `Rotation` instance + Object containing inverse of the rotations in the current instance. + + Examples + -------- + >>> from scipy.spatial.transform import Rotation as R + >>> import numpy as np + + Inverting a single rotation: + + >>> p = R.from_euler('z', 45, degrees=True) + >>> q = p.inv() + >>> q.as_euler('zyx', degrees=True) + array([-45., 0., 0.]) + + Inverting multiple rotations: + + >>> p = R.from_rotvec([[0, 0, np.pi/3], [-np.pi/4, 0, 0]]) + >>> q = p.inv() + >>> q.as_rotvec() + array([[-0. , -0. , -1.04719755], + [ 0.78539816, -0. , -0. ]]) + + """ + q_inv = self._backend.inv(self._quat) + if self._single: + q_inv = q_inv[0, ...] + return Rotation._from_raw_quat(q_inv, xp=self._xp, backend=self._backend) + + def magnitude(self) -> Array: + """Get the magnitude(s) of the rotation(s). + + Returns + ------- + magnitude : ndarray or float + Angle(s) in radians, float if object contains a single rotation + and ndarray if object contains ND rotations. The magnitude + will always be in the range [0, pi]. + + Examples + -------- + >>> from scipy.spatial.transform import Rotation as R + >>> import numpy as np + >>> r = R.from_quat(np.eye(4)) + >>> r.as_quat() + array([[ 1., 0., 0., 0.], + [ 0., 1., 0., 0.], + [ 0., 0., 1., 0.], + [ 0., 0., 0., 1.]]) + >>> r.magnitude() + array([3.14159265, 3.14159265, 3.14159265, 0. ]) + + Magnitude of a single rotation: + + >>> r[0].magnitude() + 3.141592653589793 + """ + magnitude = self._backend.magnitude(self._quat) + if self._single: + # Special handling for numpy and single rotations. self._single is only set + # if xp is numpy. We therefore know that magnitude is a numpy array and + # that it contains a single element. Previously this code returned a Python + # float in that case. Here we return a numpy float64 scalar. All other + # Array API libraries return 0d arrays instead. + # See https://github.com/scipy/scipy/pull/23198#issuecomment-3003757848 + return magnitude[0] + return magnitude + + def approx_equal( + self, other: Rotation, atol: float | None = None, degrees: bool = False + ) -> Array | np.bool: + """Determine if another rotation is approximately equal to this one. + + Equality is measured by calculating the smallest angle between the + rotations, and checking to see if it is smaller than `atol`. + + Parameters + ---------- + other : `Rotation` instance + Object containing the rotations to measure against this one. + atol : float, optional + The absolute angular tolerance, below which the rotations are + considered equal. If not given, then set to 1e-8 radians by + default. + degrees : bool, optional + If True and `atol` is given, then `atol` is measured in degrees. If + False (default), then atol is measured in radians. + + Returns + ------- + approx_equal : Array or `numpy.bool` + Whether the rotations are approximately equal, `numpy.bool` if object + contains a single numpy rotation and Array if object contains multiple + rotations or is from another library. + + Examples + -------- + >>> from scipy.spatial.transform import Rotation as R + >>> import numpy as np + >>> p = R.from_quat([0, 0, 0, 1]) + >>> q = R.from_quat(np.eye(4)) + >>> p.approx_equal(q) + array([False, False, False, True]) + + Approximate equality for a single rotation: + + >>> p.approx_equal(q[0]) + np.False_ + """ + cython_compatible = self._quat.ndim < 3 and other._quat.ndim < 3 + backend = select_backend(self._xp, cython_compatible=cython_compatible) + approx_equal = backend.approx_equal( + self._quat, other._quat, atol=atol, degrees=degrees + ) + if self._single and other._single: + return approx_equal[0] + return approx_equal + + def mean( + self, + weights: ArrayLike | None = None, + axis: None | int | tuple[int, ...] = None, + ) -> Rotation: + """Get the mean of the rotations. + + The mean used is the chordal L2 mean (also called the projected or + induced arithmetic mean) [1]_. If ``A`` is a set of rotation matrices, + then the mean ``M`` is the rotation matrix that minimizes the + following loss function: + + .. math:: + + L(M) = \\sum_{i = 1}^{n} w_i \\lVert \\mathbf{A}_i - + \\mathbf{M} \\rVert^2 , + + where :math:`w_i`'s are the `weights` corresponding to each matrix. + + Parameters + ---------- + weights : array_like shape (..., N), optional + Weights describing the relative importance of the rotations. If + None (default), then all values in `weights` are assumed to be + equal. If given, the shape of `weights` must be broadcastable to + the rotation shape. Weights must be non-negative. + axis : None, int, or tuple of ints, optional + Axis or axes along which the means are computed. The default is to + compute the mean of all rotations. + + Returns + ------- + mean : `Rotation` instance + Single rotation containing the mean of the rotations in the + current instance. + + References + ---------- + .. [1] Hartley, Richard, et al., + "Rotation Averaging", International Journal of Computer Vision + 103, 2013, pp. 267-305. + + Examples + -------- + >>> from scipy.spatial.transform import Rotation as R + >>> r = R.from_euler('zyx', [[0, 0, 0], + ... [1, 0, 0], + ... [0, 1, 0], + ... [0, 0, 1]], degrees=True) + >>> r.mean().as_euler('zyx', degrees=True) + array([0.24945696, 0.25054542, 0.24945696]) + """ + mean = self._backend.mean(self._quat, weights=weights, axis=axis) + return Rotation._from_raw_quat(mean, xp=self._xp, backend=self._backend) + + def reduce( + self, + left: Rotation | None = None, + right: Rotation | None = None, + return_indices: bool = False, + ) -> Rotation | tuple[Rotation, Array, Array]: + """Reduce this rotation with the provided rotation groups. + + Reduction of a rotation ``p`` is a transformation of the form + ``q = l * p * r``, where ``l`` and ``r`` are chosen from `left` and + `right` respectively, such that rotation ``q`` has the smallest + magnitude. + + If `left` and `right` are rotation groups representing symmetries of + two objects rotated by ``p``, then ``q`` is the rotation of the + smallest magnitude to align these objects considering their symmetries. + + Parameters + ---------- + left : `Rotation` instance, optional + Object containing the left rotation(s). Default value (None) + corresponds to the identity rotation. + right : `Rotation` instance, optional + Object containing the right rotation(s). Default value (None) + corresponds to the identity rotation. + return_indices : bool, optional + Whether to return the indices of the rotations from `left` and + `right` used for reduction. + + Returns + ------- + reduced : `Rotation` instance + Object containing reduced rotations. + left_best, right_best: int ndarray + Indices of elements from `left` and `right` used for reduction. + """ + left = left.as_quat() if left is not None else None + right = right.as_quat() if right is not None else None + reduced, left_idx, right_idx = self._backend.reduce( + self._quat, left=left, right=right + ) + if self._single: + reduced = reduced[0, ...] + rot = Rotation._from_raw_quat(reduced, xp=self._xp, backend=self._backend) + if return_indices: + left_idx = left_idx if left is not None else None + right_idx = right_idx if right is not None else None + return rot, left_idx, right_idx + return rot + + @classmethod + def create_group(cls, group: str, axis: str = "Z") -> Rotation: + """Create a 3D rotation group. + + Parameters + ---------- + group : str + The name of the group. Must be one of 'I', 'O', 'T', 'Dn', 'Cn', + where `n` is a positive integer. The groups are: + + * I: Icosahedral group + * O: Octahedral group + * T: Tetrahedral group + * D: Dicyclic group + * C: Cyclic group + + axis : int + The cyclic rotation axis. Must be one of ['X', 'Y', 'Z'] (or + lowercase). Default is 'Z'. Ignored for groups 'I', 'O', and 'T'. + + Returns + ------- + rotation : `Rotation` instance + Object containing the elements of the rotation group. + + Notes + ----- + This method generates rotation groups only. The full 3-dimensional + point groups [PointGroups]_ also contain reflections. + + References + ---------- + .. [PointGroups] `Point groups + `_ + on Wikipedia. + """ + # TODO: We defer the implementation of groups for arbitrary Array API frameworks + # to the follow-up PR that adds general Array API support for Rotations. + return create_group(cls, group, axis=axis) + + def __getitem__( + self, indexer: int | slice | EllipsisType | Array | None + ) -> Rotation: + """Extract rotation(s) at given index(es) from object. + + Create a new `Rotation` instance containing a subset of rotations + stored in this object. + + Parameters + ---------- + indexer : index, slice, or index array + Specifies which rotation(s) to extract. A single indexer must be + specified, i.e. as if indexing a 1 dimensional array or list. + + Returns + ------- + rotation : `Rotation` instance + Contains + - a single rotation, if `indexer` is a single index + - a stack of rotation(s), if `indexer` is a slice, or and index + array. + + Raises + ------ + TypeError if the instance was created as a single rotation. + + Examples + -------- + >>> from scipy.spatial.transform import Rotation as R + >>> rs = R.from_quat([ + ... [1, 1, 0, 0], + ... [0, 1, 0, 1], + ... [1, 1, -1, 0]]) # These quats are normalized + >>> rs.as_quat() + array([[ 0.70710678, 0.70710678, 0. , 0. ], + [ 0. , 0.70710678, 0. , 0.70710678], + [ 0.57735027, 0.57735027, -0.57735027, 0. ]]) + + Indexing using a single index: + + >>> a = rs[0] + >>> a.as_quat() + array([0.70710678, 0.70710678, 0. , 0. ]) + + Array slicing: + + >>> b = rs[1:3] + >>> b.as_quat() + array([[ 0. , 0.70710678, 0. , 0.70710678], + [ 0.57735027, 0.57735027, -0.57735027, 0. ]]) + + List comprehension to split each rotation into its own object: + + >>> c = [r for r in rs] + >>> print([r.as_quat() for r in c]) + [array([ 0.70710678, 0.70710678, 0. , 0. ]), + array([ 0. , 0.70710678, 0. , 0.70710678]), + array([ 0.57735027, 0.57735027, -0.57735027, 0. ])] + + Concatenation of split rotations will recover the original object: + + >>> R.concatenate([a, b]).as_quat() + array([[ 0.70710678, 0.70710678, 0. , 0. ], + [ 0. , 0.70710678, 0. , 0.70710678], + [ 0.57735027, 0.57735027, -0.57735027, 0. ]]) + """ + if self._single or self._quat.ndim == 1: + raise TypeError("Single rotation is not subscriptable.") + is_array = isinstance(indexer, type(self._quat)) + # Masking is only specified in the Array API when the array is the sole index + # TODO: This special case handling is mainly a result of Array API limitations. + # Ideally we would get rid of them altogether and converge to [indexer, ...] + # indexing. + if is_array and indexer.dtype == self._xp.bool: # type:ignore[union-attr] + return Rotation(self._quat[indexer], normalize=False) + if is_array and self._xp.isdtype(indexer.dtype, "integral"): # type:ignore[union-attr] + # xp.take is implementation-defined for zero-dim arrays, hence we raise + # pre-emptively to have consistent behavior across frameworks. + if self._quat.shape[0] == 0: + raise IndexError("cannot do a non-empty take from an empty axes.") + return Rotation(self._xp.take(self._quat, indexer, axis=0), normalize=False) + return Rotation(self._quat[indexer, ...], normalize=False) + + def __setitem__(self, indexer: int | slice | EllipsisType | None, value: Rotation): + """Set rotation(s) at given index(es) from object. + + Parameters + ---------- + indexer : index, slice, or index array + Specifies which rotation(s) to replace. A single indexer must be + specified, i.e. as if indexing a 1 dimensional array or list. + + value : `Rotation` instance + The rotations to set. + + Raises + ------ + TypeError if the instance was created as a single rotation. + + Notes + ----- + + .. versionadded:: 1.8.0 + """ + if self._single or self._quat.ndim == 1: + raise TypeError("Single rotation is not subscriptable.") + + if not isinstance(value, Rotation): + raise TypeError("value must be a Rotation object") + + self._quat = self._backend.setitem(self._quat, value.as_quat(), indexer) + + @staticmethod + def identity( + num: int | None = None, *, shape: int | tuple[int, ...] | None = None + ) -> Rotation: + """Get identity rotation(s). + + Composition with the identity rotation has no effect. + + Parameters + ---------- + num : int or None, optional + Number of identity rotations to generate. If None (default), then a + single rotation is generated. + shape : int or tuple of ints, optional + Shape of identity rotations to generate. If specified, `num` must + be None. + + Returns + ------- + identity : Rotation object + The identity rotation. + """ + # TODO: We should move to one single way of specifying the output shape and + # deprecate `num`. + if num is not None and shape is not None: + raise ValueError("Only one of `num` or `shape` can be specified.") + quat = cython_backend.identity(num, shape=shape) + return Rotation._from_raw_quat(quat, xp=array_namespace(quat)) + + @staticmethod + @_transition_to_rng("random_state", position_num=2) + def random( + num: int | None = None, + rng: np.random.Generator | None = None, + *, + shape: tuple[int, ...] | None = None, + ) -> Rotation: + r"""Generate rotations that are uniformly distributed on a sphere. + + Formally, the rotations follow the Haar-uniform distribution over the SO(3) + group. + + Parameters + ---------- + num : int or None, optional + Number of random rotations to generate. If None (default), then a + single rotation is generated. + rng : `numpy.random.Generator`, optional + Pseudorandom number generator state. When `rng` is None, a new + `numpy.random.Generator` is created using entropy from the + operating system. Types other than `numpy.random.Generator` are + passed to `numpy.random.default_rng` to instantiate a `Generator`. + shape : tuple of ints, optional + Shape of random rotations to generate. If specified, `num` must be None. + + Returns + ------- + random_rotation : `Rotation` instance + Contains a single rotation if `num` is None. Otherwise contains a + stack of `num` rotations. + + Notes + ----- + This function is optimized for efficiently sampling random rotation + matrices in three dimensions. For generating random rotation matrices + in higher dimensions, see `scipy.stats.special_ortho_group`. + + Examples + -------- + >>> from scipy.spatial.transform import Rotation as R + + Sample a single rotation: + + >>> R.random().as_euler('zxy', degrees=True) + array([-110.5976185 , 55.32758512, 76.3289269 ]) # random + + Sample a stack of rotations: + + >>> R.random(5).as_euler('zxy', degrees=True) + array([[-110.5976185 , 55.32758512, 76.3289269 ], # random + [ -91.59132005, -14.3629884 , -93.91933182], + [ 25.23835501, 45.02035145, -121.67867086], + [ -51.51414184, -15.29022692, -172.46870023], + [ -81.63376847, -27.39521579, 2.60408416]]) + + See Also + -------- + scipy.stats.special_ortho_group + + """ + # TODO: We should move to one single way of specifying the output shape and + # deprecate `num`. + if num is not None and shape is not None: + raise ValueError("Only one of `num` or `shape` can be specified.") + sample = cython_backend.random(num, rng, shape=shape) + return Rotation(sample, normalize=True, copy=False) + + @staticmethod + def align_vectors( + a: ArrayLike, + b: ArrayLike, + weights: ArrayLike | None = None, + return_sensitivity: bool = False, + ) -> tuple[Rotation, float] | tuple[Rotation, float, Array]: + """Estimate a rotation to optimally align two sets of vectors. + + Find a rotation between frames A and B which best aligns a set of + vectors `a` and `b` observed in these frames. The following loss + function is minimized to solve for the rotation matrix + :math:`C`: + + .. math:: + + L(C) = \\frac{1}{2} \\sum_{i = 1}^{n} w_i \\lVert \\mathbf{a}_i - + C \\mathbf{b}_i \\rVert^2 , + + where :math:`w_i`'s are the `weights` corresponding to each vector. + + The rotation is estimated with Kabsch algorithm [1]_, and solves what + is known as the "pointing problem", or "Wahba's problem" [2]_. + + Note that the length of each vector in this formulation acts as an + implicit weight. So for use cases where all vectors need to be + weighted equally, you should normalize them to unit length prior to + calling this method. + + There are two special cases. The first is if a single vector is given + for `a` and `b`, in which the shortest distance rotation that aligns + `b` to `a` is returned. + + The second is when one of the weights is infinity. In this case, the + shortest distance rotation between the primary infinite weight vectors + is calculated as above. Then, the rotation about the aligned primary + vectors is calculated such that the secondary vectors are optimally + aligned per the above loss function. The result is the composition + of these two rotations. The result via this process is the same as the + Kabsch algorithm as the corresponding weight approaches infinity in + the limit. For a single secondary vector this is known as the + "align-constrain" algorithm [3]_. + + For both special cases (single vectors or an infinite weight), the + sensitivity matrix does not have physical meaning and an error will be + raised if it is requested. For an infinite weight, the primary vectors + act as a constraint with perfect alignment, so their contribution to + `rssd` will be forced to 0 even if they are of different lengths. + + Parameters + ---------- + a : array_like, shape (3,) or (N, 3) + Vector components observed in initial frame A. Each row of `a` + denotes a vector. + b : array_like, shape (3,) or (N, 3) + Vector components observed in another frame B. Each row of `b` + denotes a vector. + weights : array_like shape (N,), optional + Weights describing the relative importance of the vector + observations. If None (default), then all values in `weights` are + assumed to be 1. One and only one weight may be infinity, and + weights must be positive. + return_sensitivity : bool, optional + Whether to return the sensitivity matrix. See Notes for details. + Default is False. + + Returns + ------- + rotation : `Rotation` instance + Best estimate of the rotation that transforms `b` to `a`. + rssd : float + Stands for "root sum squared distance". Square root of the weighted + sum of the squared distances between the given sets of vectors + after alignment. It is equal to ``sqrt(2 * minimum_loss)``, where + ``minimum_loss`` is the loss function evaluated for the found + optimal rotation. + Note that the result will also be weighted by the vectors' + magnitudes, so perfectly aligned vector pairs will have nonzero + `rssd` if they are not of the same length. This can be avoided by + normalizing them to unit length prior to calling this method, + though note that doing this will change the resulting rotation. + sensitivity_matrix : ndarray, shape (3, 3) + Sensitivity matrix of the estimated rotation estimate as explained + in Notes. Returned only when `return_sensitivity` is True. Not + valid if aligning a single pair of vectors or if there is an + infinite weight, in which cases an error will be raised. + + Notes + ----- + The sensitivity matrix gives the sensitivity of the estimated rotation + to small perturbations of the vector measurements. Specifically we + consider the rotation estimate error as a small rotation vector of + frame A. The sensitivity matrix is proportional to the covariance of + this rotation vector assuming that the vectors in `a` was measured with + errors significantly less than their lengths. To get the true + covariance matrix, the returned sensitivity matrix must be multiplied + by harmonic mean [4]_ of variance in each observation. Note that + `weights` are supposed to be inversely proportional to the observation + variances to get consistent results. For example, if all vectors are + measured with the same accuracy of 0.01 (`weights` must be all equal), + then you should multiple the sensitivity matrix by 0.01**2 to get the + covariance. + + Refer to [5]_ for more rigorous discussion of the covariance + estimation. See [6]_ for more discussion of the pointing problem and + minimal proper pointing. + + This function does not support broadcasting or ND arrays with N > 2. + + References + ---------- + .. [1] https://en.wikipedia.org/wiki/Kabsch_algorithm + .. [2] https://en.wikipedia.org/wiki/Wahba%27s_problem + .. [3] Magner, Robert, + "Extending target tracking capabilities through trajectory and + momentum setpoint optimization." Small Satellite Conference, + 2018. + .. [4] https://en.wikipedia.org/wiki/Harmonic_mean + .. [5] F. Landis Markley, + "Attitude determination using vector observations: a fast + optimal matrix algorithm", Journal of Astronautical Sciences, + Vol. 41, No.2, 1993, pp. 261-280. + .. [6] Bar-Itzhack, Itzhack Y., Daniel Hershkowitz, and Leiba Rodman, + "Pointing in Real Euclidean Space", Journal of Guidance, + Control, and Dynamics, Vol. 20, No. 5, 1997, pp. 916-922. + + Examples + -------- + >>> import numpy as np + >>> from scipy.spatial.transform import Rotation as R + + Here we run the baseline Kabsch algorithm to best align two sets of + vectors, where there is noise on the last two vector measurements of + the ``b`` set: + + >>> a = [[0, 1, 0], [0, 1, 1], [0, 1, 1]] + >>> b = [[1, 0, 0], [1, 1.1, 0], [1, 0.9, 0]] + >>> rot, rssd, sens = R.align_vectors(a, b, return_sensitivity=True) + >>> rot.as_matrix() + array([[0., 0., 1.], + [1., 0., 0.], + [0., 1., 0.]]) + + When we apply the rotation to ``b``, we get vectors close to ``a``: + + >>> rot.apply(b) + array([[0. , 1. , 0. ], + [0. , 1. , 1.1], + [0. , 1. , 0.9]]) + + The error for the first vector is 0, and for the last two the error is + magnitude 0.1. The `rssd` is the square root of the sum of the + weighted squared errors, and the default weights are all 1, so in this + case the `rssd` is calculated as + ``sqrt(1 * 0**2 + 1 * 0.1**2 + 1 * (-0.1)**2) = 0.141421356237308`` + + >>> a - rot.apply(b) + array([[ 0., 0., 0. ], + [ 0., 0., -0.1], + [ 0., 0., 0.1]]) + >>> np.sqrt(np.sum(np.ones(3) @ (a - rot.apply(b))**2)) + 0.141421356237308 + >>> rssd + 0.141421356237308 + + The sensitivity matrix for this example is as follows: + + >>> sens + array([[0.2, 0. , 0.], + [0. , 1.5, 1.], + [0. , 1. , 1.]]) + + Special case 1: Find a minimum rotation between single vectors: + + >>> a = [1, 0, 0] + >>> b = [0, 1, 0] + >>> rot, _ = R.align_vectors(a, b) + >>> rot.as_matrix() + array([[0., 1., 0.], + [-1., 0., 0.], + [0., 0., 1.]]) + >>> rot.apply(b) + array([1., 0., 0.]) + + Special case 2: One infinite weight. Here we find a rotation between + primary and secondary vectors that can align exactly: + + >>> a = [[0, 1, 0], [0, 1, 1]] + >>> b = [[1, 0, 0], [1, 1, 0]] + >>> rot, _ = R.align_vectors(a, b, weights=[np.inf, 1]) + >>> rot.as_matrix() + array([[0., 0., 1.], + [1., 0., 0.], + [0., 1., 0.]]) + >>> rot.apply(b) + array([[0., 1., 0.], + [0., 1., 1.]]) + + Here the secondary vectors must be best-fit: + + >>> a = [[0, 1, 0], [0, 1, 1]] + >>> b = [[1, 0, 0], [1, 2, 0]] + >>> rot, _ = R.align_vectors(a, b, weights=[np.inf, 1]) + >>> rot.as_matrix() + array([[0., 0., 1.], + [1., 0., 0.], + [0., 1., 0.]]) + >>> rot.apply(b) + array([[0., 1., 0.], + [0., 1., 2.]]) + """ + xp = array_namespace(a) + a, b, weights = _promote(a, b, weights, xp=xp) + cython_compatible = ( + (a.ndim < 3) & (b.ndim < 3) & (weights is None or weights.ndim < 2) # type:ignore[union-attr] + ) + backend = select_backend(xp, cython_compatible=cython_compatible) + q, rssd, sensitivity = backend.align_vectors(a, b, weights, return_sensitivity) + if return_sensitivity: + return Rotation._from_raw_quat(q, xp=xp, backend=backend), rssd, sensitivity + return Rotation._from_raw_quat(q, xp=xp, backend=backend), rssd + + def __getstate__(self) -> tuple[Array, bool]: + return (self._quat, self._single) + + def __setstate__(self, state: tuple[Array, bool]): + quat, single = state + xp = array_namespace(quat) + self._xp = xp + self._quat = xp.asarray(quat, copy=True) + self._backend = select_backend(xp, cython_compatible=self._quat.ndim < 3) + self._single = single + + @property + def single(self) -> bool: + """Whether this instance represents a single rotation.""" + return self._single or self._quat.ndim == 1 + + @property + def shape(self) -> tuple[int, ...]: + """The shape of the rotation's leading dimensions.""" + if self._single: + return () + return self._quat.shape[:-1] + + def __bool__(self) -> bool: + """Comply with Python convention for objects to be True. + + Required because `Rotation.__len__()` is defined and not always truthy. + """ + return True + + def __len__(self) -> int: + """Number of rotations contained in this object. + + Multiple rotations can be stored in a single instance. + + Returns + ------- + length : int + Number of rotations stored in object. + + Raises + ------ + TypeError if the instance was created as a single rotation. + """ + if self._single or self._quat.ndim == 1: + raise TypeError("Single rotation has no len().") + return self._quat.shape[0] + + def __repr__(self) -> str: + m = f"{self.as_matrix()!r}".splitlines() + # bump indent (+21 characters) + m[1:] = [" " * 21 + m[i] for i in range(1, len(m))] + return "Rotation.from_matrix(" + "\n".join(m) + ")" + + def __iter__(self) -> Iterator[Rotation]: + """Iterate over rotations.""" + if self._single or self._quat.ndim == 1: + raise TypeError("Single rotation is not iterable.") + # We return a generator that yields a new Rotation object for each rotation + # in the current object. We cannot rely on the default implementation using + # __getitem__ because jax will not raise an IndexError for out-of-bounds + # indices. + for i in range(self._quat.shape[0]): + yield Rotation(self._quat[i, ...], normalize=False, copy=False) + + @staticmethod + def _from_raw_quat( + quat: Array, xp: ModuleType, backend: ModuleType | None = None + ) -> Rotation: + """Create a Rotation skipping all sanitization steps. + + This method is intended for internal, performant creation of Rotations with + quaternions that are guaranteed to be valid. + """ + rot = Rotation.__new__(Rotation) + rot._single = quat.ndim == 1 and is_numpy(xp) + if rot._single: + quat = xpx.atleast_nd(quat, ndim=2, xp=xp) + rot._quat = quat + rot._xp = xp + if backend is None: + backend = select_backend(xp, cython_compatible=quat.ndim < 3) + rot._backend = backend + return rot + + +@xp_capabilities( + skip_backends=[("dask.array", "missing linalg.cross function")], + method_capabilities={ + "__init__": dict( + jax_jit=False, + skip_backends=[("dask.array", "missing linalg.cross function")], + ), + }, +) +class Slerp: + """Spherical Linear Interpolation of Rotations. + + The interpolation between consecutive rotations is performed as a rotation + around a fixed axis with a constant angular velocity [1]_. This ensures + that the interpolated rotations follow the shortest path between initial + and final orientations. + + Parameters + ---------- + times : array_like, shape (N,) + Times of the known rotations. At least 2 times must be specified. + rotations : `Rotation` instance + Rotations to perform the interpolation between. Must contain N + rotations. + + Methods + ------- + __call__ + + See Also + -------- + Rotation + + Notes + ----- + This class only supports interpolation of rotations with a single leading + dimension. + + .. versionadded:: 1.2.0 + + References + ---------- + .. [1] https://en.wikipedia.org/wiki/Slerp#Quaternion_Slerp + + Examples + -------- + >>> from scipy.spatial.transform import Rotation as R + >>> from scipy.spatial.transform import Slerp + + Setup the fixed keyframe rotations and times: + + >>> key_rots = R.random(5, random_state=2342345) + >>> key_times = [0, 1, 2, 3, 4] + + Create the interpolator object: + + >>> slerp = Slerp(key_times, key_rots) + + Interpolate the rotations at the given times: + + >>> times = [0, 0.5, 0.25, 1, 1.5, 2, 2.75, 3, 3.25, 3.60, 4] + >>> interp_rots = slerp(times) + + The keyframe rotations expressed as Euler angles: + + >>> key_rots.as_euler('xyz', degrees=True) + array([[ 14.31443779, -27.50095894, -3.7275787 ], + [ -1.79924227, -24.69421529, 164.57701743], + [146.15020772, 43.22849451, -31.34891088], + [ 46.39959442, 11.62126073, -45.99719267], + [-88.94647804, -49.64400082, -65.80546984]]) + + The interpolated rotations expressed as Euler angles. These agree with the + keyframe rotations at both endpoints of the range of keyframe times. + + >>> interp_rots.as_euler('xyz', degrees=True) + array([[ 14.31443779, -27.50095894, -3.7275787 ], + [ 4.74588574, -32.44683966, 81.25139984], + [ 10.71094749, -31.56690154, 38.06896408], + [ -1.79924227, -24.69421529, 164.57701743], + [ 11.72796022, 51.64207311, -171.7374683 ], + [ 146.15020772, 43.22849451, -31.34891088], + [ 68.10921869, 20.67625074, -48.74886034], + [ 46.39959442, 11.62126073, -45.99719267], + [ 12.35552615, 4.21525086, -64.89288124], + [ -30.08117143, -19.90769513, -78.98121326], + [ -88.94647804, -49.64400082, -65.80546984]]) + + """ + + def __init__(self, times: ArrayLike, rotations: Rotation): + if not isinstance(rotations, Rotation): + raise TypeError("`rotations` must be a `Rotation` instance.") + + if rotations.single or len(rotations) <= 1: + raise ValueError("`rotations` must be a sequence of at least 2 rotations.") + + q = rotations.as_quat() + if q.ndim > 2: + raise ValueError( + "Rotations with more than 1 leading dimension are not supported." + ) + + xp = array_namespace(q) + times = xp.asarray(times, device=xp_device(q), dtype=q.dtype) + if times.ndim != 1: + raise ValueError( + "Expected times to be specified in a 1 dimensional array, got " + f"{times.ndim} dimensions." + ) + + if times.shape[0] != len(rotations): + raise ValueError( + "Expected number of rotations to be equal to number of timestamps " + f"given, got {len(rotations)} rotations and {times.shape[0]} " + "timestamps." + ) + self.times = times + self.timedelta = xp.diff(times) + + # We cannot check for values for lazy backends, so we cannot raise an error on + # timedelta < 0 for lazy backends. Instead, we set timedelta to nans + neg_mask = xp.any(self.timedelta <= 0) + if is_lazy_array(neg_mask): + self.timedelta = xp.where(neg_mask, xp.nan, self.timedelta) + self.times = xp.where(neg_mask, xp.nan, self.times) + elif xp.any(neg_mask): + raise ValueError("Times must be in strictly increasing order.") + + self.rotations = rotations[:-1] + self.rotvecs = (self.rotations.inv() * rotations[1:]).as_rotvec() + + def __call__(self, times: ArrayLike) -> Rotation: + """Interpolate rotations. + + Compute the interpolated rotations at the given `times`. + + Parameters + ---------- + times : array_like + Times to compute the interpolations at. Can be a scalar or + 1-dimensional. + + Returns + ------- + interpolated_rotation : `Rotation` instance + Object containing the rotations computed at given `times`. + + """ + xp = array_namespace(self.times) + device = xp_device(self.times) + # Clearly differentiate from self.times property + compute_times = xp.asarray(times, device=device, dtype=self.times.dtype) + if compute_times.ndim > 1: + raise ValueError("`times` must be at most 1-dimensional.") + + single_time = compute_times.ndim == 0 + compute_times = xpx.atleast_nd(compute_times, ndim=1, xp=xp) + + # side = 'left' (default) excludes t_min. + ind = xp.searchsorted(self.times, compute_times) - 1 + # Include t_min. Without this step, index for t_min equals -1 + ind = xpx.at(ind, compute_times == self.times[0]).set(0) + # We cannot error out on invalid indices for jit compiled code. To not produce + # an index error, we set the index to 0 in case it is out of bounds, and later + # set the result to nan. + invalid_ind: Array = (ind < 0) | (ind > len(self.rotations) - 1) + if is_lazy_array(invalid_ind): + ind = xpx.at(ind, invalid_ind).set(0) + elif xp.any(invalid_ind): + raise ValueError( + f"Interpolation times must be within the range [{self.times[0]}, " + f"{self.times[-1]}], both inclusive." + ) + + alpha = (compute_times - self.times[ind]) / self.timedelta[ind] + alpha = xpx.at(alpha, invalid_ind).set(xp.nan) + + # The array API does not support integer arrays + ellipsis indexing and won't + # stabilize this feature due to blockers in PyTorch. Therefore we need to + # construct the index for the last dimension manually. + # See https://github.com/scipy/scipy/pull/23249#discussion_r2198363047 + result = self.rotations[ind] * Rotation.from_rotvec( + self.rotvecs[ind[:, None], xp.arange(3, device=device)] * alpha[:, None] + ) + + if single_time: + result = result[0] + + return result diff --git a/v12x/pylib/scipy/spatial/transform/_rotation_groups.py b/v12x/pylib/scipy/spatial/transform/_rotation_groups.py new file mode 100644 index 0000000000000000000000000000000000000000..870e9b9e2b44bff56b8228a70607e29f8173accc --- /dev/null +++ b/v12x/pylib/scipy/spatial/transform/_rotation_groups.py @@ -0,0 +1,140 @@ +import numpy as np +from scipy.constants import golden as phi + + +def icosahedral(cls): + g1 = tetrahedral(cls).as_quat() + a = 0.5 + b = 0.5 / phi + c = phi / 2 + g2 = np.array([[+a, +b, +c, 0], + [+a, +b, -c, 0], + [+a, +c, 0, +b], + [+a, +c, 0, -b], + [+a, -b, +c, 0], + [+a, -b, -c, 0], + [+a, -c, 0, +b], + [+a, -c, 0, -b], + [+a, 0, +b, +c], + [+a, 0, +b, -c], + [+a, 0, -b, +c], + [+a, 0, -b, -c], + [+b, +a, 0, +c], + [+b, +a, 0, -c], + [+b, +c, +a, 0], + [+b, +c, -a, 0], + [+b, -a, 0, +c], + [+b, -a, 0, -c], + [+b, -c, +a, 0], + [+b, -c, -a, 0], + [+b, 0, +c, +a], + [+b, 0, +c, -a], + [+b, 0, -c, +a], + [+b, 0, -c, -a], + [+c, +a, +b, 0], + [+c, +a, -b, 0], + [+c, +b, 0, +a], + [+c, +b, 0, -a], + [+c, -a, +b, 0], + [+c, -a, -b, 0], + [+c, -b, 0, +a], + [+c, -b, 0, -a], + [+c, 0, +a, +b], + [+c, 0, +a, -b], + [+c, 0, -a, +b], + [+c, 0, -a, -b], + [0, +a, +c, +b], + [0, +a, +c, -b], + [0, +a, -c, +b], + [0, +a, -c, -b], + [0, +b, +a, +c], + [0, +b, +a, -c], + [0, +b, -a, +c], + [0, +b, -a, -c], + [0, +c, +b, +a], + [0, +c, +b, -a], + [0, +c, -b, +a], + [0, +c, -b, -a]]) + return cls.from_quat(np.concatenate((g1, g2))) + + +def octahedral(cls): + g1 = tetrahedral(cls).as_quat() + c = np.sqrt(2) / 2 + g2 = np.array([[+c, 0, 0, +c], + [0, +c, 0, +c], + [0, 0, +c, +c], + [0, 0, -c, +c], + [0, -c, 0, +c], + [-c, 0, 0, +c], + [0, +c, +c, 0], + [0, -c, +c, 0], + [+c, 0, +c, 0], + [-c, 0, +c, 0], + [+c, +c, 0, 0], + [-c, +c, 0, 0]]) + return cls.from_quat(np.concatenate((g1, g2))) + + +def tetrahedral(cls): + g1 = np.eye(4) + c = 0.5 + g2 = np.array([[c, -c, -c, +c], + [c, -c, +c, +c], + [c, +c, -c, +c], + [c, +c, +c, +c], + [c, -c, -c, -c], + [c, -c, +c, -c], + [c, +c, -c, -c], + [c, +c, +c, -c]]) + return cls.from_quat(np.concatenate((g1, g2))) + + +def dicyclic(cls, n, axis=2): + g1 = cyclic(cls, n, axis).as_rotvec() + + thetas = np.linspace(0, np.pi, n, endpoint=False) + rv = np.pi * np.vstack([np.zeros(n), np.cos(thetas), np.sin(thetas)]).T + g2 = np.roll(rv, axis, axis=1) + return cls.from_rotvec(np.concatenate((g1, g2))) + + +def cyclic(cls, n, axis=2): + thetas = np.linspace(0, 2 * np.pi, n, endpoint=False) + rv = np.vstack([thetas, np.zeros(n), np.zeros(n)]).T + return cls.from_rotvec(np.roll(rv, axis, axis=1)) + + +def create_group(cls, group, axis='Z'): + if not isinstance(group, str): + raise ValueError("`group` argument must be a string") + + permitted_axes = ['x', 'y', 'z', 'X', 'Y', 'Z'] + if axis not in permitted_axes: + raise ValueError("`axis` must be one of " + ", ".join(permitted_axes)) + + if group in ['I', 'O', 'T']: + symbol = group + order = 1 + elif group[:1] in ['C', 'D'] and group[1:].isdigit(): + symbol = group[:1] + order = int(group[1:]) + else: + raise ValueError("`group` must be one of 'I', 'O', 'T', 'Dn', 'Cn'") + + if order < 1: + raise ValueError("Group order must be positive") + + axis = 'xyz'.index(axis.lower()) + if symbol == 'I': + return icosahedral(cls) + elif symbol == 'O': + return octahedral(cls) + elif symbol == 'T': + return tetrahedral(cls) + elif symbol == 'D': + return dicyclic(cls, order, axis=axis) + elif symbol == 'C': + return cyclic(cls, order, axis=axis) + else: + assert False diff --git a/v12x/pylib/scipy/spatial/transform/_rotation_spline.py b/v12x/pylib/scipy/spatial/transform/_rotation_spline.py new file mode 100644 index 0000000000000000000000000000000000000000..9d6b99a5af898bff9ee9148b8712f861c54740fd --- /dev/null +++ b/v12x/pylib/scipy/spatial/transform/_rotation_spline.py @@ -0,0 +1,465 @@ +import numpy as np +from scipy.linalg import solve_banded +from ._rotation import Rotation + + +def _create_skew_matrix(x): + """Create skew-symmetric matrices corresponding to vectors. + + Parameters + ---------- + x : ndarray, shape (n, 3) + Set of vectors. + + Returns + ------- + ndarray, shape (n, 3, 3) + """ + result = np.zeros((len(x), 3, 3)) + result[:, 0, 1] = -x[:, 2] + result[:, 0, 2] = x[:, 1] + result[:, 1, 0] = x[:, 2] + result[:, 1, 2] = -x[:, 0] + result[:, 2, 0] = -x[:, 1] + result[:, 2, 1] = x[:, 0] + return result + + +def _matrix_vector_product_of_stacks(A, b): + """Compute the product of stack of matrices and vectors.""" + return np.einsum("ijk,ik->ij", A, b) + + +def _angular_rate_to_rotvec_dot_matrix(rotvecs): + """Compute matrices to transform angular rates to rot. vector derivatives. + + The matrices depend on the current attitude represented as a rotation + vector. + + Parameters + ---------- + rotvecs : ndarray, shape (n, 3) + Set of rotation vectors. + + Returns + ------- + ndarray, shape (n, 3, 3) + """ + norm = np.linalg.norm(rotvecs, axis=1) + k = np.empty_like(norm) + + mask = norm > 1e-4 + nm = norm[mask] + k[mask] = (1 - 0.5 * nm / np.tan(0.5 * nm)) / nm**2 + mask = ~mask + nm = norm[mask] + k[mask] = 1/12 + 1/720 * nm**2 + + skew = _create_skew_matrix(rotvecs) + + result = np.empty((len(rotvecs), 3, 3)) + result[:] = np.identity(3) + result[:] += 0.5 * skew + result[:] += k[:, None, None] * np.matmul(skew, skew) + + return result + + +def _rotvec_dot_to_angular_rate_matrix(rotvecs): + """Compute matrices to transform rot. vector derivatives to angular rates. + + The matrices depend on the current attitude represented as a rotation + vector. + + Parameters + ---------- + rotvecs : ndarray, shape (n, 3) + Set of rotation vectors. + + Returns + ------- + ndarray, shape (n, 3, 3) + """ + norm = np.linalg.norm(rotvecs, axis=1) + k1 = np.empty_like(norm) + k2 = np.empty_like(norm) + + mask = norm > 1e-4 + nm = norm[mask] + k1[mask] = (1 - np.cos(nm)) / nm ** 2 + k2[mask] = (nm - np.sin(nm)) / nm ** 3 + + mask = ~mask + nm = norm[mask] + k1[mask] = 0.5 - nm ** 2 / 24 + k2[mask] = 1 / 6 - nm ** 2 / 120 + + skew = _create_skew_matrix(rotvecs) + + result = np.empty((len(rotvecs), 3, 3)) + result[:] = np.identity(3) + result[:] -= k1[:, None, None] * skew + result[:] += k2[:, None, None] * np.matmul(skew, skew) + + return result + + +def _angular_acceleration_nonlinear_term(rotvecs, rotvecs_dot): + """Compute the non-linear term in angular acceleration. + + The angular acceleration contains a quadratic term with respect to + the derivative of the rotation vector. This function computes that. + + Parameters + ---------- + rotvecs : ndarray, shape (n, 3) + Set of rotation vectors. + rotvecs_dot : ndarray, shape (n, 3) + Set of rotation vector derivatives. + + Returns + ------- + ndarray, shape (n, 3) + """ + norm = np.linalg.norm(rotvecs, axis=1) + dp = np.sum(rotvecs * rotvecs_dot, axis=1) + cp = np.cross(rotvecs, rotvecs_dot) + ccp = np.cross(rotvecs, cp) + dccp = np.cross(rotvecs_dot, cp) + + k1 = np.empty_like(norm) + k2 = np.empty_like(norm) + k3 = np.empty_like(norm) + + mask = norm > 1e-4 + nm = norm[mask] + k1[mask] = (-nm * np.sin(nm) - 2 * (np.cos(nm) - 1)) / nm ** 4 + k2[mask] = (-2 * nm + 3 * np.sin(nm) - nm * np.cos(nm)) / nm ** 5 + k3[mask] = (nm - np.sin(nm)) / nm ** 3 + + mask = ~mask + nm = norm[mask] + k1[mask] = 1/12 - nm ** 2 / 180 + k2[mask] = -1/60 + nm ** 2 / 12604 + k3[mask] = 1/6 - nm ** 2 / 120 + + dp = dp[:, None] + k1 = k1[:, None] + k2 = k2[:, None] + k3 = k3[:, None] + + return dp * (k1 * cp + k2 * ccp) + k3 * dccp + + +def _compute_angular_rate(rotvecs, rotvecs_dot): + """Compute angular rates given rotation vectors and its derivatives. + + Parameters + ---------- + rotvecs : ndarray, shape (n, 3) + Set of rotation vectors. + rotvecs_dot : ndarray, shape (n, 3) + Set of rotation vector derivatives. + + Returns + ------- + ndarray, shape (n, 3) + """ + return _matrix_vector_product_of_stacks( + _rotvec_dot_to_angular_rate_matrix(rotvecs), rotvecs_dot) + + +def _compute_angular_acceleration(rotvecs, rotvecs_dot, rotvecs_dot_dot): + """Compute angular acceleration given rotation vector and its derivatives. + + Parameters + ---------- + rotvecs : ndarray, shape (n, 3) + Set of rotation vectors. + rotvecs_dot : ndarray, shape (n, 3) + Set of rotation vector derivatives. + rotvecs_dot_dot : ndarray, shape (n, 3) + Set of rotation vector second derivatives. + + Returns + ------- + ndarray, shape (n, 3) + """ + return (_compute_angular_rate(rotvecs, rotvecs_dot_dot) + + _angular_acceleration_nonlinear_term(rotvecs, rotvecs_dot)) + + +def _create_block_3_diagonal_matrix(A, B, d): + """Create a 3-diagonal block matrix as banded. + + The matrix has the following structure: + + DB... + ADB.. + .ADB. + ..ADB + ...AD + + The blocks A, B and D are 3-by-3 matrices. The D matrices has the form + d * I. + + Parameters + ---------- + A : ndarray, shape (n, 3, 3) + Stack of A blocks. + B : ndarray, shape (n, 3, 3) + Stack of B blocks. + d : ndarray, shape (n + 1,) + Values for diagonal blocks. + + Returns + ------- + ndarray, shape (11, 3 * (n + 1)) + Matrix in the banded form as used by `scipy.linalg.solve_banded`. + """ + ind = np.arange(3) + ind_blocks = np.arange(len(A)) + + A_i = np.empty_like(A, dtype=int) + A_i[:] = ind[:, None] + A_i += 3 * (1 + ind_blocks[:, None, None]) + + A_j = np.empty_like(A, dtype=int) + A_j[:] = ind + A_j += 3 * ind_blocks[:, None, None] + + B_i = np.empty_like(B, dtype=int) + B_i[:] = ind[:, None] + B_i += 3 * ind_blocks[:, None, None] + + B_j = np.empty_like(B, dtype=int) + B_j[:] = ind + B_j += 3 * (1 + ind_blocks[:, None, None]) + + diag_i = diag_j = np.arange(3 * len(d)) + i = np.hstack((A_i.ravel(), B_i.ravel(), diag_i)) + j = np.hstack((A_j.ravel(), B_j.ravel(), diag_j)) + values = np.hstack((A.ravel(), B.ravel(), np.repeat(d, 3))) + + u = 5 + l = 5 + result = np.zeros((u + l + 1, 3 * len(d))) + result[u + i - j, j] = values + return result + + +class RotationSpline: + """Interpolate rotations with continuous angular rate and acceleration. + + The rotation vectors between each consecutive orientation are cubic + functions of time and it is guaranteed that angular rate and acceleration + are continuous. Such interpolation are analogous to cubic spline + interpolation. + + Refer to [1]_ for math and implementation details. + + Parameters + ---------- + times : array_like, shape (N,) + Times of the known rotations. At least 2 times must be specified. + rotations : `Rotation` instance + Rotations to perform the interpolation between. Must contain N + rotations. + + Methods + ------- + __call__ + + References + ---------- + .. [1] `Smooth Attitude Interpolation + `_ + + Examples + -------- + >>> from scipy.spatial.transform import Rotation, RotationSpline + >>> import numpy as np + + Define the sequence of times and rotations from the Euler angles: + + >>> times = [0, 10, 20, 40] + >>> angles = [[-10, 20, 30], [0, 15, 40], [-30, 45, 30], [20, 45, 90]] + >>> rotations = Rotation.from_euler('XYZ', angles, degrees=True) + + Create the interpolator object: + + >>> spline = RotationSpline(times, rotations) + + Interpolate the Euler angles, angular rate and acceleration: + + >>> angular_rate = np.rad2deg(spline(times, 1)) + >>> angular_acceleration = np.rad2deg(spline(times, 2)) + >>> times_plot = np.linspace(times[0], times[-1], 100) + >>> angles_plot = spline(times_plot).as_euler('XYZ', degrees=True) + >>> angular_rate_plot = np.rad2deg(spline(times_plot, 1)) + >>> angular_acceleration_plot = np.rad2deg(spline(times_plot, 2)) + + On this plot you see that Euler angles are continuous and smooth: + + >>> import matplotlib.pyplot as plt + >>> plt.plot(times_plot, angles_plot) + >>> plt.plot(times, angles, 'x') + >>> plt.title("Euler angles") + >>> plt.show() + + The angular rate is also smooth: + + >>> plt.plot(times_plot, angular_rate_plot) + >>> plt.plot(times, angular_rate, 'x') + >>> plt.title("Angular rate") + >>> plt.show() + + The angular acceleration is continuous, but not smooth. Also note that + the angular acceleration is not a piecewise-linear function, because + it is different from the second derivative of the rotation vector (which + is a piecewise-linear function as in the cubic spline). + + >>> plt.plot(times_plot, angular_acceleration_plot) + >>> plt.plot(times, angular_acceleration, 'x') + >>> plt.title("Angular acceleration") + >>> plt.show() + """ + # Parameters for the solver for angular rate. + MAX_ITER = 10 + TOL = 1e-9 + + def _solve_for_angular_rates(self, dt, angular_rates, rotvecs): + angular_rate_first = angular_rates[0].copy() + + A = _angular_rate_to_rotvec_dot_matrix(rotvecs) + A_inv = _rotvec_dot_to_angular_rate_matrix(rotvecs) + M = _create_block_3_diagonal_matrix( + 2 * A_inv[1:-1] / dt[1:-1, None, None], + 2 * A[1:-1] / dt[1:-1, None, None], + 4 * (1 / dt[:-1] + 1 / dt[1:])) + + b0 = 6 * (rotvecs[:-1] * dt[:-1, None] ** -2 + + rotvecs[1:] * dt[1:, None] ** -2) + b0[0] -= 2 / dt[0] * A_inv[0].dot(angular_rate_first) + b0[-1] -= 2 / dt[-1] * A[-1].dot(angular_rates[-1]) + + for iteration in range(self.MAX_ITER): + rotvecs_dot = _matrix_vector_product_of_stacks(A, angular_rates) + delta_beta = _angular_acceleration_nonlinear_term( + rotvecs[:-1], rotvecs_dot[:-1]) + b = b0 - delta_beta + angular_rates_new = solve_banded((5, 5), M, b.ravel()) + angular_rates_new = angular_rates_new.reshape((-1, 3)) + + delta = np.abs(angular_rates_new - angular_rates[:-1]) + angular_rates[:-1] = angular_rates_new + if np.all(delta < self.TOL * (1 + np.abs(angular_rates_new))): + break + + rotvecs_dot = _matrix_vector_product_of_stacks(A, angular_rates) + angular_rates = np.vstack((angular_rate_first, angular_rates[:-1])) + + return angular_rates, rotvecs_dot + + def __init__(self, times, rotations): + from scipy.interpolate import PPoly + + if rotations.single: + raise ValueError("`rotations` must be a sequence of rotations.") + + if rotations.as_quat().ndim > 2: + raise ValueError( + "Rotations with more than 1 leading dimension are not supported." + ) + + if len(rotations) == 1: + raise ValueError("`rotations` must contain at least 2 rotations.") + + times = np.asarray(times, dtype=float) + if times.ndim != 1: + raise ValueError("`times` must be 1-dimensional.") + + if len(times) != len(rotations): + raise ValueError("Expected number of rotations to be equal to " + "number of timestamps given, " + f"got {len(rotations)} rotations " + f"and {len(times)} timestamps.") + + dt = np.diff(times) + if np.any(dt <= 0): + raise ValueError("Values in `times` must be in a strictly " + "increasing order.") + + rotvecs = (rotations[:-1].inv() * rotations[1:]).as_rotvec() + angular_rates = rotvecs / dt[:, None] + + if len(rotations) == 2: + rotvecs_dot = angular_rates + else: + angular_rates, rotvecs_dot = self._solve_for_angular_rates( + dt, angular_rates, rotvecs) + + dt = dt[:, None] + coeff = np.empty((4, len(times) - 1, 3)) + coeff[0] = (-2 * rotvecs + dt * angular_rates + + dt * rotvecs_dot) / dt ** 3 + coeff[1] = (3 * rotvecs - 2 * dt * angular_rates + - dt * rotvecs_dot) / dt ** 2 + coeff[2] = angular_rates + coeff[3] = 0 + + self.times = times + self.rotations = rotations + self.interpolator = PPoly(coeff, times) + + def __call__(self, times, order=0): + """Compute interpolated values. + + Parameters + ---------- + times : float or array_like + Times of interest. + order : {0, 1, 2}, optional + Order of differentiation: + + * 0 (default) : return Rotation + * 1 : return the angular rate in rad/sec + * 2 : return the angular acceleration in rad/sec/sec + + Returns + ------- + Interpolated Rotation, angular rate or acceleration. + """ + if order not in [0, 1, 2]: + raise ValueError("`order` must be 0, 1 or 2.") + + times = np.asarray(times, dtype=float) + if times.ndim > 1: + raise ValueError("`times` must be at most 1-dimensional.") + + singe_time = times.ndim == 0 + times = np.atleast_1d(times) + + rotvecs = self.interpolator(times) + if order == 0: + index = np.searchsorted(self.times, times, side='right') + index -= 1 + index[index < 0] = 0 + n_segments = len(self.times) - 1 + index[index > n_segments - 1] = n_segments - 1 + result = self.rotations[index] * Rotation.from_rotvec(rotvecs) + elif order == 1: + rotvecs_dot = self.interpolator(times, 1) + result = _compute_angular_rate(rotvecs, rotvecs_dot) + elif order == 2: + rotvecs_dot = self.interpolator(times, 1) + rotvecs_dot_dot = self.interpolator(times, 2) + result = _compute_angular_acceleration(rotvecs, rotvecs_dot, + rotvecs_dot_dot) + else: + assert False + + if singe_time: + result = result[0] + + return result diff --git a/v12x/pylib/scipy/spatial/transform/_rotation_xp.py b/v12x/pylib/scipy/spatial/transform/_rotation_xp.py new file mode 100644 index 0000000000000000000000000000000000000000..5eea00637acab3d8d914ae3bbd3bdae2a87eb9c8 --- /dev/null +++ b/v12x/pylib/scipy/spatial/transform/_rotation_xp.py @@ -0,0 +1,1140 @@ +"""Array API backend for the `Rotation` class. + +This module provides generic, functional implementations of the `Rotation` class methods +that work with any Array API-compatible backend. +""" + +# Parts of the implementation are adapted from the cython backend and +# https://github.com/jax-ml/jax/blob/d695aa4c63ffcebefce52794427c46bad576680c/jax/_src/scipy/spatial/transform.py. +import re +import warnings +from types import EllipsisType + +import numpy as np +from scipy._lib._array_api import ( + array_namespace, + Array, + ArrayLike, + is_lazy_array, + xp_vector_norm, + xp_result_type, + xp_promote, + is_jax, +) +from scipy._lib._util import broadcastable +from scipy._external.array_api_compat import device as xp_device +from scipy._external.array_api_compat import is_array_api_obj +import scipy._external.array_api_extra as xpx + +# mypy: disable-error-code=index +# mypy: disable-error-code=operator +# mypy: disable-error-code=union-attr + +def from_quat( + quat: Array, + normalize: bool = True, + copy: bool = True, + *, + scalar_first: bool = False, +) -> Array: + xp = array_namespace(quat) + if scalar_first: + quat = xp.roll(quat, -1, axis=-1) + # Normalize will always copy, so we avoid the extra copy if we normalize + if copy and not normalize: + quat = xp.asarray(quat, copy=True) + if normalize: + quat = _normalize_quaternion(quat) + return quat + + +def from_matrix(matrix: Array, assume_valid: bool = False) -> Array: + xp = array_namespace(matrix) + device = xp_device(matrix) + + if not assume_valid: + mask = xp.linalg.det(matrix) <= 0 + lazy = is_lazy_array(mask) + # Only non-lazy backends raise an error for non-positive determinants. + if not lazy and xp.any(mask): + ind = int(xp.nonzero(xpx.atleast_nd(mask, ndim=1, xp=xp))[0][0]) + raise ValueError( + "Non-positive determinant (left-handed or null coordinate frame) in " + f"rotation matrix {ind}: {matrix[ind, ...]}." + ) + elif lazy: + matrix = xp.where(mask[..., None, None], xp.nan, matrix) + + gramians = matrix @ xp.matrix_transpose(matrix) + eye = xp.eye(3, dtype=matrix.dtype, device=device) + is_orthogonal = xp.all( + xpx.isclose(gramians, eye, atol=1e-12, xp=xp), axis=(-2, -1) + ) + + if lazy: + # Lazy backends do not support non-concrete boolean indexing or any form of + # computation without statically known shapes, so we always compute SVD and + # use xp.where to select the result. + U, _, Vt = xp.linalg.svd(matrix, full_matrices=False) + matrix = xp.where(is_orthogonal[..., None, None], matrix, U @ Vt) + elif not xp.all(is_orthogonal): + # For eager frameworks, only compute SVD if needed. + is_not_orthogonal = ~is_orthogonal + U, _, Vt = xp.linalg.svd(matrix[is_not_orthogonal], full_matrices=False) + matrix = xpx.at(matrix)[is_not_orthogonal].set(U @ Vt) + + return _from_matrix_orthogonal(matrix) + + +def _from_matrix_orthogonal(matrix: Array) -> Array: + """Convert known orthogonal rotation matrix to quaternion""" + xp = array_namespace(matrix) + device = xp_device(matrix) + + matrix_trace = matrix[..., 0, 0] + matrix[..., 1, 1] + matrix[..., 2, 2] + decision = xp.stack( + [matrix[..., 0, 0], matrix[..., 1, 1], matrix[..., 2, 2], matrix_trace], + axis=-1, + ) + choice = xp.argmax(decision, axis=-1, keepdims=True) + quat = xp.empty((*matrix.shape[:-2], 4), dtype=matrix.dtype, device=device) + + # The Array API does not support mixing integer indexing with ellipsis, so we cannot + # follow the same pattern as the cython backend. Instead, we compute each case + # explicitly and assemble the final result with `xp.where`. + # TODO: Revisit this implementation if the array API supports mixed integer and + # ellipsis indexing. + + # Case 0 + quat_0 = xp.stack( + [ + 1 - matrix_trace[...] + 2 * matrix[..., 0, 0], + matrix[..., 1, 0] + matrix[..., 0, 1], + matrix[..., 2, 0] + matrix[..., 0, 2], + matrix[..., 2, 1] - matrix[..., 1, 2], + ], + axis=-1, + ) + quat = xp.where(choice == 0, quat_0, quat) + + # Case 1 + quat_1 = xp.stack( + [ + matrix[..., 1, 0] + matrix[..., 0, 1], + 1 - matrix_trace[...] + 2 * matrix[..., 1, 1], + matrix[..., 2, 1] + matrix[..., 1, 2], + matrix[..., 0, 2] - matrix[..., 2, 0], + ], + axis=-1, + ) + quat = xp.where(choice == 1, quat_1, quat) + + # Case 2 + quat_2 = xp.stack( + [ + matrix[..., 2, 0] + matrix[..., 0, 2], + matrix[..., 2, 1] + matrix[..., 1, 2], + 1 - matrix_trace[...] + 2 * matrix[..., 2, 2], + matrix[..., 1, 0] - matrix[..., 0, 1], + ], + axis=-1, + ) + quat = xp.where(choice == 2, quat_2, quat) + + # Case 3 + quat_3 = xp.stack( + [ + matrix[..., 2, 1] - matrix[..., 1, 2], + matrix[..., 0, 2] - matrix[..., 2, 0], + matrix[..., 1, 0] - matrix[..., 0, 1], + 1 + matrix_trace[...], + ], + axis=-1, + ) + quat = xp.where(choice == 3, quat_3, quat) + + return _normalize_quaternion(quat) + + +def from_rotvec(rotvec: Array, degrees: bool = False) -> Array: + xp = array_namespace(rotvec) + if rotvec.shape[-1] != 3: + raise ValueError( + f"Expected `rot_vec` to have shape (..., 3), got {rotvec.shape}" + ) + rotvec = _deg2rad(rotvec) if degrees else rotvec + + angle = xp_vector_norm(rotvec, axis=-1, keepdims=True, xp=xp) + small_angle = angle <= 1e-3 + angle2 = angle**2 + small_scale = 0.5 - angle2 / 48 + angle2**2 / 3840 + # We need to handle the case where angle is 0 to avoid division by zero. We use the + # value of the Taylor series approximation, but non-branching operations require + # that we still divide by the angle. Since we do not use the result where the angle + # is close to 0, this is safe. + div_angle = angle + xp.asarray(small_angle, dtype=angle.dtype) + large_scale = xp.sin(angle / 2) / div_angle + scale = xp.where(small_angle, small_scale, large_scale) + quat = xp.concat([rotvec * scale, xp.cos(angle / 2)], axis=-1) + return quat + + +def from_mrp(mrp: Array) -> Array: + xp = array_namespace(mrp) + if mrp.shape[-1] != 3: + raise ValueError(f"Expected `mrp` to have shape (..., 3), got {mrp.shape}") + mrp2_plus_1 = xp.linalg.vecdot(mrp, mrp, axis=-1)[..., None] + 1 + q_no_norm = xp.concat([2 * mrp[..., :3], (2 - mrp2_plus_1)], axis=-1) + quat = q_no_norm / mrp2_plus_1 + return quat + + +def from_euler(seq: str, angles: Array, degrees: bool = False) -> Array: + xp = array_namespace(angles) + num_axes = len(seq) + if num_axes < 1 or num_axes > 3: + raise ValueError( + "Expected axis specification to be a non-empty " + f"string of up to 3 characters, got {seq}" + ) + + intrinsic = re.match(r"^[XYZ]{1,3}$", seq) is not None + extrinsic = re.match(r"^[xyz]{1,3}$", seq) is not None + if not (intrinsic or extrinsic): + raise ValueError( + "Expected axes from `seq` to be from ['x', 'y', " + f"'z'] or ['X', 'Y', 'Z'], got {seq}" + ) + + if any(seq[i] == seq[i + 1] for i in range(num_axes - 1)): + raise ValueError(f"Expected consecutive axes to be different, got {seq}") + + if degrees: + angles = _deg2rad(angles) + + angles = xpx.atleast_nd(angles, ndim=1, xp=xp) + + if angles.shape[-1] != num_axes: + raise ValueError( + "Expected last dimension of `angles` to match number of sequence axes " + f"specified, got {angles.shape[-1]}." + ) + axes = [_elementary_basis_index(x) for x in seq.lower()] + q = _elementary_quat_compose(axes, angles, intrinsic) + return q + + +def from_davenport( + axes: Array, order: str, angles: Array, degrees: bool = False +) -> Array: + xp = array_namespace(axes) + device = xp_device(axes) + if order in ["e", "extrinsic"]: # Must be static, cannot be jitted + extrinsic = True + elif order in ["i", "intrinsic"]: + extrinsic = False + else: + raise ValueError( + "order should be 'e'/'extrinsic' for extrinsic sequences or 'i'/" + f"'intrinsic' for intrinsic sequences, got {order}" + ) + + if axes.shape[-1] != 3: + raise ValueError("Axes must be vectors of length 3.") + + axes = xpx.atleast_nd(axes, ndim=2, xp=xp) + angles = xpx.atleast_nd(angles, ndim=1, xp=xp) + num_axes = axes.shape[-2] + if num_axes < 1 or num_axes > 3: + raise ValueError(f"Expected up to 3 axes, got {num_axes}") + + axes = axes / xp_vector_norm(axes, axis=-1, keepdims=True, xp=xp) + + # Check if axes are orthogonal. Shape checks also work for lazy backends. + axes_not_orthogonal = xp.zeros(axes.shape[:-2], dtype=xp.bool, device=device) + if num_axes > 1: + # Cannot be True yet, so we do not need to use xp.logical_or + axes_not_orthogonal = axes_not_orthogonal | ( + xp.abs(xp.vecdot(axes[..., 0, :], axes[..., 1, :])) > 1e-7 + ) + if num_axes > 2: + axes_not_orthogonal = axes_not_orthogonal | ( + xp.abs(xp.vecdot(axes[..., 1, :], axes[..., 2, :])) > 1e-7 + ) + if not is_lazy_array(axes_not_orthogonal) and xp.any(axes_not_orthogonal): + raise ValueError("Consecutive axes must be orthogonal.") + else: + axes = xp.where(axes_not_orthogonal[..., None, None], xp.nan, axes) + + if degrees: + angles = _deg2rad(angles) + + if ( + not broadcastable(axes.shape[:-1], angles.shape) + or axes.shape[-2] != angles.shape[-1] + ): + raise ValueError( + f"Expected `angles` to match number of axes, got {angles.shape} angles " + f"and {axes.shape} axes." + ) + + q_shape = angles.shape[:-1] + (4,) + q = xp.zeros(q_shape, dtype=angles.dtype, device=xp_device(angles)) + q = xpx.at(q)[..., 3].set(1) + + for i in range(num_axes): + qi = from_rotvec(angles[..., i, None] * axes[..., i, :]) + q = compose_quat(qi, q) if extrinsic else compose_quat(q, qi) + return q + + +def as_quat( + quat: Array, canonical: bool = False, *, scalar_first: bool = False +) -> Array: + xp = array_namespace(quat) + if canonical: + quat = _quat_canonical(quat) + if scalar_first: + quat = xp.roll(quat, 1, axis=-1) + return quat + + +def as_matrix(quat: Array) -> Array: + xp = array_namespace(quat) + x = quat[..., 0] + y = quat[..., 1] + z = quat[..., 2] + w = quat[..., 3] + + x2 = x * x + y2 = y * y + z2 = z * z + w2 = w * w + + xy = x * y + zw = z * w + xz = x * z + yw = y * w + yz = y * z + xw = x * w + + matrix_elements = [ + x2 - y2 - z2 + w2, + 2 * (xy - zw), + 2 * (xz + yw), + 2 * (xy + zw), + -x2 + y2 - z2 + w2, + 2 * (yz - xw), + 2 * (xz - yw), + 2 * (yz + xw), + -x2 - y2 + z2 + w2, + ] + matrix = xp.reshape(xp.stack(matrix_elements, axis=-1), (*quat.shape[:-1], 3, 3)) + return matrix + + +def as_rotvec(quat: Array, degrees: bool = False) -> Array: + xp = array_namespace(quat) + quat = _quat_canonical(quat) + ax_norm = xp_vector_norm(quat[..., :3], axis=-1, keepdims=True, xp=xp) + angle = 2 * xp.atan2(ax_norm, quat[..., 3][..., None]) + small_angle = angle <= 1e-3 + angle2 = angle**2 + small_scale = 2 + angle2 / 12 + 7 * angle2**2 / 2880 + # We need to handle the case where sin(angle/2) is 0 to avoid division by zero. We + # use the value of the Taylor series approximation, but non-branching operations + # require that we still divide by the sin. Since we do not use the result where the + # angle is close to 0, adding one to the sin where we discard the result is safe. + div_sin = xp.sin(angle / 2.0) + xp.asarray(small_angle, dtype=angle.dtype) + large_scale = angle / div_sin + scale = xp.where(small_angle, small_scale, large_scale) + if degrees: + scale = _rad2deg(scale) + rotvec = scale * quat[..., :3] + return rotvec + + +def as_mrp(quat: Array) -> Array: + xp = array_namespace(quat) + one = xp.asarray(1.0, device=xp_device(quat), dtype=quat.dtype) + sign = xp.where(quat[..., 3, None] < 0, -1, one) + denominator = 1.0 + sign * quat[..., 3, None] + return sign * quat[..., :3] / denominator + + +def as_euler( + quat: Array, seq: str, degrees: bool = False, *, suppress_warnings: bool = False +) -> Array: + xp = array_namespace(quat) + + # Sanitize the sequence + if len(seq) != 3: + raise ValueError(f"Expected 3 axes, got {seq}.") + intrinsic = re.match(r"^[XYZ]{1,3}$", seq) is not None + extrinsic = re.match(r"^[xyz]{1,3}$", seq) is not None + if not (intrinsic or extrinsic): + raise ValueError( + "Expected axes from `seq` to be from ['x', 'y', 'z'] or ['X', 'Y', 'Z'], " + f"got {seq}" + ) + if any(seq[i] == seq[i + 1] for i in range(2)): + raise ValueError(f"Expected consecutive axes to be different, got {seq}") + + device = xp_device(quat) + axes = [_elementary_basis_index(x) for x in seq.lower()] + axes = axes if extrinsic else axes[::-1] + i, j, k = axes + symmetric = i == k + k = 3 - i - j if symmetric else k + + mask = xp.asarray(symmetric, device=device) + sign = xp.asarray((i - j) * (j - k) * (k - i) // 2, dtype=quat.dtype, device=device) + # Permute quaternion elements + a = xp.where(mask, quat[..., 3], quat[..., 3] - quat[..., j]) + b = xp.where(mask, quat[..., i], quat[..., i] + quat[..., k] * sign) + c = xp.where(mask, quat[..., j], quat[..., j] + quat[..., 3]) + d = xp.where(mask, quat[..., k] * sign, quat[..., k] * sign - quat[..., i]) + + angles = _get_angles( + extrinsic, symmetric, sign, xp.pi / 2, a, b, c, d, suppress_warnings + ) + return _rad2deg(angles) if degrees else angles + + +def as_davenport( + quat: Array, + axes: Array, + order: str, + degrees: bool = False, + *, + suppress_warnings: bool = False, +) -> Array: + xp = array_namespace(quat) + + # Check argument validity + if order in ["e", "extrinsic"]: + extrinsic = True + elif order in ["i", "intrinsic"]: + extrinsic = False + else: + raise ValueError( + "order should be 'e'/'extrinsic' for extrinsic sequences or 'i'/'intrinsic'" + f" for intrinsic sequences, got {order}" + ) + if axes.shape[-2] != 3: + raise ValueError(f"Expected 3 axes, got {axes.shape}.") + if axes.shape[-1] != 3: + raise ValueError("Axes must be vectors of length 3.") + if not broadcastable(axes.shape[:-2], quat.shape[:-1]): + raise ValueError( + f"Expected `axes` to match number of rotations, got {axes.shape} axes " + f"and {quat.shape} rotations." + ) + + # normalize axes + axes = axes / xp_vector_norm(axes, axis=-1, keepdims=True, xp=xp) + vdot_ax0_ax1 = xp.vecdot(axes[..., 0, :], axes[..., 1, :]) + vdot_ax1_ax2 = xp.vecdot(axes[..., 1, :], axes[..., 2, :]) + is_invalid = (vdot_ax0_ax1 >= 1e-7) | (vdot_ax1_ax2 >= 1e-7) + if is_lazy_array(is_invalid): + axes = xp.where(is_invalid[..., None, None], xp.nan, axes) + elif xp.any(is_invalid): + raise ValueError("Consecutive axes must be orthogonal.") + + angles = _compute_davenport_from_quat( + quat, + axes[..., 0, :], + axes[..., 1, :], + axes[..., 2, :], + extrinsic, + suppress_warnings, + ) + if degrees: + angles = _rad2deg(angles) + return angles + + +def inv(quat: Array) -> Array: + return xpx.at(quat)[..., :3].multiply(-1, copy=True) + + +def magnitude(quat: Array) -> Array: + xp = array_namespace(quat) + sin_q = xp_vector_norm(quat[..., :3], axis=-1, xp=xp) + cos_q = xp.abs(quat[..., 3]) + angles = 2 * xp.atan2(sin_q, cos_q) + return angles + + +def approx_equal( + quat: Array, other: Array, atol: float | None = None, degrees: bool = False +) -> Array: + if atol is None: + if degrees: + warnings.warn( + "atol must be set to use the degrees flag, defaulting to 1e-8 radians.", + stacklevel=2, + ) + atol = 1e-8 + elif degrees: + atol = _deg2rad(atol) + + if not broadcastable(quat.shape, other.shape): + raise ValueError( + f"Expected broadcastable shapes in both rotations, got {quat.shape[:-1]} " + f"rotations in first and {other.shape[:-1]} rotations in second object." + ) + + quat_result = compose_quat(other, inv(quat)) + angles = magnitude(quat_result) + return angles < atol + + +def mean( + quat: Array, + weights: ArrayLike | None = None, + axis: None | int | tuple[int, ...] = None, +) -> Array: + xp = array_namespace(quat) + device = xp_device(quat) + dtype = xp_result_type(quat, force_floating=True, xp=xp) + if quat.shape[0] == 0: + raise ValueError("Mean of an empty rotation set is undefined.") + # Axis logic: For None, we reduce over all axes. For int, we only reduce over that + # axis. For tuple, we reduce over all specified axes. + all_axes = tuple(range(quat.ndim - 1)) + if axis is None: + axis = all_axes + elif isinstance(axis, int): + axis = (axis,) + if not isinstance(axis, tuple): + raise ValueError("`axis` must be None, int, or tuple of ints.") + # Ensure all axes are within bounds + if axis != () and (min(axis) < -(quat.ndim - 1) or max(axis) > (quat.ndim - 2)): + raise ValueError( + f"axis {axis} is out of bounds for rotation with shape {quat.shape[:-1]}." + ) + # Ensure all axes are positive and unique + axis = tuple(sorted(set(x % (quat.ndim - 1) for x in axis))) + + lazy = is_lazy_array(quat) + # Branching code is okay for checks that include meta info such as shapes and types + quat_expand = quat[..., None, :] + if weights is None: + K = xp.matrix_transpose(quat_expand) @ quat_expand + else: + weights = xp.asarray(weights, dtype=dtype, device=device) + neg_weights = weights < 0 + if not lazy and xp.any(neg_weights): + raise ValueError("`weights` must be non-negative.") + elif lazy: + # We cannot check for negative weights because jit code needs to be + # non-branching. We return NaN instead + weights = xp.where(neg_weights, xp.nan, weights) + + if not broadcastable(quat.shape[:-1], weights.shape): + raise ValueError( + "Expected `weights` to be broadcastable to rotation shape, got shape " + f"{weights.shape} for {quat.shape[:-1]} rotations." + ) + + # Make sure we can transpose quat + weighted_quat = weights[..., None, None] * quat_expand + K = xp.matrix_transpose(weighted_quat) @ quat_expand + + # Move reduction axes to the end + keep_axes = tuple(i for i in all_axes if i not in axis) + axes_order = keep_axes + axis + K_reordered = xp.moveaxis(K, axes_order, all_axes) + # Reshape to flatten reduction axes + new_shape = K_reordered.shape[: len(keep_axes)] + (-1, 4, 4) + K = xp.mean(xp.reshape(K_reordered, new_shape), axis=-3) + _, v = xp.linalg.eigh(K) + return v[..., -1] + + +def reduce( + quat: Array, + left: Array | None = None, + right: Array | None = None, +) -> tuple[Array, Array | None, Array | None]: + if left is None and right is None: + return quat, None, None + # DECISION: We cannot have variable number of return arguments for jit compiled + # functions. We therefore always return the indices, and filter out later. + # TOOD: Properly support broadcasting. + xp = array_namespace(quat) + quat = xpx.atleast_nd(quat, ndim=2, xp=xp) + if left is None: + left = xp.ones_like(quat) + if right is None: + right = xp.ones_like(quat) + + # We want to calculate the real components of q = l * p * r. It can + # be shown that: + # qs = ls * ps * rs - ls * dot(pv, rv) - ps * dot(lv, rv) + # - rs * dot(lv, pv) - dot(cross(lv, pv), rv) + # where ls and lv denote the scalar and vector components of l. + + p = quat + ps, pv = _split_rotation(p, xp) + ls, lv = _split_rotation(left, xp) + rs, rv = _split_rotation(right, xp) + + # Compute each term without einsum (not accessible in the Array API) + # First term: np.einsum("i,j,k", ls, ps, rs) + term1 = ls[..., :, None, None] * ps[..., None, :, None] * rs[..., None, None, :] + # Second term: np.einsum('i,jx,kx', ls, pv, rv) + prv = xp.sum(pv[..., :, None, :] * rv[..., None, :, :], axis=-1) + term2 = ls[..., :, None, None] * prv[..., None, :, :] + # Third term: np.einsum('ix,j,kx', lv, ps, rv) + lrv = xp.sum(lv[..., :, None, :] * rv[..., None, :, :], axis=-1) + term3 = ps[..., None, :, None] * lrv[..., :, None, :] + # Fourth term: np.einsum('ix,jx,k', lv, pv, rs) + lpv = xp.sum(lv[..., :, None, :] * pv[..., None, :, :], axis=-1) + term4 = rs[..., None, None, :] * lpv[..., :, :, None] + # Fifth term: np.einsum('xyz,ix,jy,kz', e, lv, pv, rv). We want to avoid expanding + # the einsum into a 6D tensor to avoid excessive memory usage. Instead, we compute + # the cross product between lv and pv and then compute the dot product with rv. + # First compute cross products between lv and pv + lv_expanded = lv[..., :, None, :] + pv_expanded = pv[..., None, :, :] + cross_lp = xp.linalg.cross(lv_expanded, pv_expanded) + # Then compute dot product with rv + term5 = xp.sum(cross_lp[..., :, :, None, :] * rv[..., None, None, :, :], axis=-1) + # Combine all terms with proper shape alignment + qs = xp.abs(term1 - term2 - term3 - term4 - term5) + qs = xp.reshape(xp.moveaxis(qs, 1, 0), (qs.shape[1], -1)) + + # Find best indices from scalar components + max_ind = xp.argmax(xp.reshape(qs, (qs.shape[0], -1)), axis=1) + left_best = max_ind // rv.shape[0] + right_best = max_ind % rv.shape[0] + # Array API limitation: Integer index arrays are only allowed with integer indices + # TODO: Can we somehow avoid this? + all_idx = xp.reshape(xp.arange(left.shape[-1]), (1, -1)) + left_idx = xp.reshape(left_best, (-1, 1)) + left = left[left_idx, all_idx] + right_idx = xp.reshape(right_best, (-1, 1)) + right = right[right_idx, all_idx] + + # Reduce the rotation using the best indices + reduced = compose_quat(left, compose_quat(p, right)) + + if left is None: + left_best = None + if right is None: + right_best = None + return reduced, left_best, right_best + + +def apply(quat: Array, points: Array, inverse: bool = False) -> Array: + xp = array_namespace(quat) + mat = as_matrix(quat) + # We do not have access to einsum. To avoid broadcasting issues, we add a singleton + # dimension to the points array and remove it after the operation. + points = points[..., None] + if not broadcastable(mat.shape, points.shape): + raise ValueError( + f"Cannot broadcast {quat.shape[:-1]} rotations to {points.shape[:-1]} " + "vectors." + ) + if inverse: + # TODO: Replace with .mT once numpy 2.0 is the minimum supported version + return (xp.matrix_transpose(mat) @ points)[..., 0] + return (mat @ points)[..., 0] + + +def setitem( + quat: Array, value: Array, indexer: int | slice | EllipsisType | None +) -> Array: + return xpx.at(quat)[indexer, ...].set(value) + + +def align_vectors( + a: Array, b: Array, weights: Array | None = None, return_sensitivity: bool = False +) -> tuple[Array, Array, Array]: + xp = array_namespace(a) + # Check input vectors + dtype = xp_result_type(a, b, force_floating=True, xp=xp) + a_original = xp.asarray(a, dtype=dtype) + b_original = xp.asarray(b, dtype=dtype) + a = xpx.atleast_nd(a_original, ndim=2, xp=xp) + b = xpx.atleast_nd(b_original, ndim=2, xp=xp) + if a.shape[-1] != 3: + raise ValueError( + f"Expected input `a` to have shape (3,) or (N, 3), got {a_original.shape}" + ) + if b.shape[-1] != 3: + raise ValueError( + f"Expected input `b` to have shape (3,) or (N, 3), got {b_original.shape}" + ) + if a.shape != b.shape: + raise ValueError( + "Expected inputs `a` and `b` to have same shapes" + f", got {a_original.shape} and {b_original.shape} respectively." + ) + if a.ndim > 2 or b.ndim > 2: # This function does not support broadcasting + raise ValueError( + "Expected inputs `a` and `b` to have shape (3,) or (N, 3), got " + f"{a_original.shape} and {b_original.shape} respectively." + ) + N = a.shape[0] + + # Check weights + if weights is None: + weights = xp.ones(N, device=xp_device(a), dtype=a.dtype) + else: + weights = xp.asarray(weights, device=xp_device(a), dtype=a.dtype) + if weights.ndim != 1: + raise ValueError( + f"Expected `weights` to be 1 dimensional, got shape {weights.shape}." + ) + if N > 1 and (weights.shape[0] != N): + raise ValueError( + "Expected `weights` to have number of values equal to number of input " + f"vectors, got {weights.shape[0]} values and {N} vectors." + ) + # We can only check for negative weights in eager execution models. Lazy + # backends will return NaNs instead. + negative_weights = weights < 0 + if not is_lazy_array(negative_weights) and xp.any(negative_weights): + raise ValueError("`weights` may not contain negative values") + weights = xp.where(negative_weights, xp.nan, weights) + + # For the special case of a single vector pair, we use the infinite + # weight code path + weight_is_inf = xp.asarray([True]) if N == 1 else weights == xp.inf + n_inf = xp.sum(xp.astype(weight_is_inf, a.dtype)) + # We can only error out on multiple infinite weights or sensitivity return with + # infinite weights in eager execution models. Lazy backends will return NaNs. + if not is_lazy_array(n_inf): + if n_inf > 1: + raise ValueError("Only one infinite weight is allowed") + if n_inf == 1 and return_sensitivity: + raise ValueError( + "Cannot return sensitivity matrix with an " + "infinite weight or one vector pair" + ) + + weights = xp.where(n_inf > 1, xp.nan, weights) + + inf_branch = xp.any(weight_is_inf, axis=-1) + # DECISION: We cannot compute both branches for all frameworks. There are two main + # reasons: + # 1. Computing both for eager execution models is expensive. + # 2. Some operations will fail when running the unused branch because of numerical + # and algorithmical issues. Numpy e.g. will raise an exception when trying to + # compute the svd of a matrix with infinite weights. To prevent this, we only + # compute the branch that is needed. Lazy backends however require us to take the + # full compute graph. Therefore, we use xp.where for lazy backends and a branching + # version for eager frameworks. + # + # Note that we could also solve this by exploiting the externals of xpx.apply_where. + # However, we'd have to rely on the implementation details of apply_where, which is + # something we should avoid. + # See https://github.com/scipy/scipy/pull/22777#discussion_r2028868364 + if is_lazy_array(inf_branch): + q_opt, rssd, sensitivity = _align_vectors(a, b, weights) + q_opt_inf, rssd_inf, sensitivity_inf = _align_vectors_fixed(a, b, weights) + q_opt = xp.where(inf_branch, q_opt_inf, q_opt) + rssd = xp.where(inf_branch, rssd_inf, rssd) + sensitivity = xp.where(inf_branch, sensitivity_inf, sensitivity) + else: + if xp.any(inf_branch): + q_opt, rssd, sensitivity = _align_vectors_fixed(a, b, weights) + else: + q_opt, rssd, sensitivity = _align_vectors(a, b, weights) + return q_opt, rssd, sensitivity + + +def _align_vectors(a: Array, b: Array, weights: Array) -> tuple[Array, Array, Array]: + xp = array_namespace(a) + device = xp_device(a) + B = (weights[:, None] * a).mT @ b + u, s, vh = xp.linalg.svd(B) + + # Correct improper rotation if necessary (as in Kabsch algorithm) + neg_det = xp.linalg.det(u @ vh) < 0 + s = xpx.at(s)[..., -1].set(xp.where(neg_det, -s[..., -1], s[..., -1])) + u = xpx.at(u)[..., :, -1].set(xp.where(neg_det, -u[..., :, -1], u[..., :, -1])) + + C = u @ vh + + # DECISION: We cannot branch on the condition because jit code needs to be + # non-branching. Hence, we omit the check for uniqueness + # (s[1] + s[2] < 1e-16 * s[0]) + ssd = xp.sum(weights * xp.sum(b**2 + a**2, axis=-1), axis=-1) - 2 * xp.sum( + s, axis=-1 + ) + rssd = xp.sqrt(xp.maximum(ssd, xp.zeros(1, device=device)))[..., 0] + + # TODO: We currently need to always compute the sensitivity matrix because lazy code + # needs to be non-branching. We should check if compilers can optimize the where + # statement (e.g. in jax) and check if we can have an eager version that only + # evaluates the branch that is needed. + # See xpx.apply_where, issue: https://github.com/data-apis/array-api-extra/pull/141 + zeta = (s[..., 0] + s[..., 1]) * (s[..., 1] + s[..., 2]) * (s[..., 2] + s[..., 0]) + kappa = s[..., 0] * s[..., 1] + s[..., 1] * s[..., 2] + s[..., 2] * s[..., 0] + eye = xp.eye(3, dtype=a.dtype, device=device) + sensitivity = xp.mean(weights) / zeta * (kappa * eye + B @ B.mT) + q_opt = _from_matrix_orthogonal(C) + return q_opt, rssd, sensitivity + + +def _align_vectors_fixed( + a: Array, b: Array, weights: Array +) -> tuple[Array, Array, Array]: + xp = array_namespace(a) + device = xp_device(a) + N = a.shape[0] + weight_is_inf = xp.asarray([True], device=device) if N == 1 else weights == xp.inf + # We cannot use boolean masks for indexing because of jax. For the same reason, we + # also cannot use dynamic slices. As a workaround, we roll the array so that the + # infinitely weighted vector is at index 0. We then use static slices to get the + # primary and secondary vectors. + # + # Note that argmax fulfils a secondary purpose here: + # When we trace this function with jax, this function might get executed even if + # weight_is_inf does not have a single valid entry. Argmax will still give us a + # valid index which allows us to proceed with the function (the results are + # discarded anyways), whereas boolean indexing would result in invalid, zero-shaped + # arrays. + + inf_idx = xp.argmax(xp.astype(weight_is_inf, xp.uint8)) + # xp.argmax returns an Array, but the specification of xp.roll does not support + # Arrays as shifts. For lazy execution models we cannot convert to int because this + # raises a concretization error. However, jax does accept Arrays as shifts which + # allows jit compiling the function. Hence we do not convert to int for jax. + # See https://github.com/data-apis/array-api/issues/914#issuecomment-2918959918 + if not is_jax(xp): + inf_idx = int(inf_idx) + a_sorted = xp.roll(a, shift=-inf_idx, axis=0) + b_sorted = xp.roll(b, shift=-inf_idx, axis=0) + weights_sorted = xp.roll(weights, shift=-inf_idx, axis=0) + + a_pri = a_sorted[0, ...][None, ...] # Effectively [[0], ...] + b_pri = b_sorted[0, ...][None, ...] + a_pri_norm = xp_vector_norm(a_pri, axis=-1, keepdims=True, xp=xp) + b_pri_norm = xp_vector_norm(b_pri, axis=-1, keepdims=True, xp=xp) + if not is_lazy_array(a_pri_norm): + if xp.any(a_pri_norm == 0) or xp.any(b_pri_norm == 0): + raise ValueError("Cannot align zero length primary vectors") + + # We cannot error out on zero length vectors. We set the norm to NaN to avoid + # division by zero and mark the result as invalid. + a_pri_norm = xpx.at(a_pri_norm)[a_pri_norm == 0].set(xp.nan) + b_pri_norm = xpx.at(b_pri_norm)[b_pri_norm == 0].set(xp.nan) + + a_pri = a_pri / a_pri_norm + b_pri = b_pri / b_pri_norm + + # We first find the minimum angle rotation between the primary + # vectors. + cross = xp.linalg.cross(b_pri[..., 0, :], a_pri[..., 0, :]) + cross_norm = xp_vector_norm(cross, axis=-1, xp=xp) + theta = xp.atan2(cross_norm, xp.sum(a_pri[..., 0, :] * b_pri[..., 0, :], axis=-1)) + tolerance = 1e-3 # tolerance for small angle approximation (rad) + q_flip = xp.asarray([0.0, 0.0, 0.0, 1.0], dtype=a_pri.dtype, device=device) + + # Near pi radians, the Taylor series approximation of x/sin(x) diverges, so for + # numerical stability we flip pi and then rotate back by the small angle pi - theta + flip = xp.pi - theta < tolerance + # For antiparallel vectors, cross = [0, 0, 0] so we need to manually set an + # arbitrary orthogonal axis of rotation + i = xp.argmin(xp.abs(a_pri[..., 0, :])) + # TODO: Array API does not support fancy indexing __setitem__. The code is + # equivalent to doing the following: + # r_components = xp.zeros(3) + # r_components = xpx.at(r_components)[i - 1].set(a_pri[0, i - 2]) + # r_components = xpx.at(r_components)[i - 2].set(-a_pri[0, i - 1]) + r_components = xp.stack( + [ + xp.where(i == 1, a_pri[0, 2], xp.where(i == 2, -a_pri[0, 1], 0)), + xp.where(i == 2, a_pri[0, 0], xp.where(i == 0, -a_pri[0, 2], 0)), + xp.where(i == 0, a_pri[0, 1], xp.where(i == 1, -a_pri[0, 0], 0)), + ] + ) + r = xp.where(cross_norm == 0, r_components, cross) + + q_flip = xp.where(flip, from_rotvec(r / xp_vector_norm(r, xp=xp) * xp.pi), q_flip) + theta = xp.where(flip, xp.pi - theta, theta) + cross = xp.where(flip, -cross, cross) + + # Small angle Taylor series approximation for numerical stability + theta2 = theta * theta + small_scale = xp.abs(theta) < tolerance + r_small_scale = cross * (1 + theta2 / 6 + theta2 * theta2 * 7 / 360) + # We need to handle the case where theta is 0 to avoid division by zero. We use the + # value of the Taylor series approximation, but non-branching operations require + # that we still divide by the angle. Since we do not use the result where the angle + # is close to 0, this is safe. + theta = theta + xp.asarray(small_scale, dtype=theta.dtype) + r_large_scale = cross * theta / xp.sin(theta) + r = xp.where(small_scale, r_small_scale, r_large_scale) + q_pri = compose_quat(from_rotvec(r), q_flip) + + # Case 1): No secondary vectors, q_opt is q_pri -> Immediately done + # Case 2): Secondary vectors exist + # We cannot use boolean masks here because of jax + a_sec = a_sorted[1:, ...] + b_sec = b_sorted[1:, ...] + weights_sec = weights_sorted[1:] + + # We apply the first rotation to the b vectors to align the + # primary vectors, resulting in vectors c. The secondary + # vectors must now be rotated about that primary vector to best + # align c to a. + c_sec = apply(q_pri, b_sec) + + # Calculate vector components for the angle calculation. The + # angle phi to rotate a single 2D vector C to align to 2D + # vector A in the xy plane can be found with the equation + # phi = atan2(cross(C, A), dot(C, A)) + # = atan2(|C|*|A|*sin(phi), |C|*|A|*cos(phi)) + # The below equations perform the same operation, but with the + # 3D rotation restricted to the 2D plane normal to a_pri, where + # the secondary vectors are projected into that plane. We then + # find the composite angle from the weighted sum of the + # axial components in that plane, minimizing the 2D alignment + # problem. + sin_term = xp.linalg.vecdot(xp.linalg.cross(c_sec, a_sec), a_pri) + cos_term = xp.linalg.vecdot(c_sec, a_sec) - ( + xp.linalg.vecdot(c_sec, a_pri) * xp.linalg.vecdot(a_sec, a_pri) + ) + phi = xp.atan2(xp.sum(weights_sec * sin_term), xp.sum(weights_sec * cos_term)) + q_sec = from_rotvec(phi * a_pri[0, ...]) + + # Compose these to get the optimal rotation + q_opt = q_pri if N == 1 else compose_quat(q_sec, q_pri) + + # Calculate the root sum squared distance. We force the error to + # be zero for the infinite weight vectors since they will align + # exactly. + mask = ((N > 1) | (weights[0] == xp.inf)) & weight_is_inf + # Skip non-infinite weight single vectors pairs, we used the + # infinite weight code path but don't want to zero that weight + weights_inf_zero = xpx.at(weights, mask).set(0, copy=True, xp=xp) + a_est = apply(q_opt, b) + rssd = xp.sqrt(xp.sum(weights_inf_zero @ (a - a_est) ** 2)) + + mask = xp.any(xp.isnan(weights), axis=-1) + q_opt = xp.where(mask, xp.nan, q_opt) + return q_opt, rssd, xp.asarray(xp.nan, device=device) + + +def pow(quat: Array, n: float | Array) -> Array: + xp = array_namespace(quat) + device = xp_device(quat) + # If n is an array, we sanitize it to a scalar and promote quat and n to + # the same dtype. + if is_array_api_obj(n): + if n.shape == (1,): # pyrefly:ignore[missing-attribute] + n = n[0] # pyrefly:ignore[bad-index] + elif n.ndim != 0: # pyrefly:ignore[missing-attribute] + raise ValueError("Array exponent must be a scalar") + quat, n = xp_promote(quat, n, force_floating=True, xp=xp) + + # If n is a lazy array, we cannot take fast paths for special cases. + if is_lazy_array(n): + result = from_rotvec(n * as_rotvec(quat)) # general scaling of rotation angle + # Special cases 0 -> identity, -1 -> inv, 1 -> copy + identity = xp.zeros((*quat.shape[:-1], 4), dtype=quat.dtype, device=device) + identity = xpx.at(identity)[..., 3].set(1) + result = xp.where(n == 0, identity, result) + result = xp.where(n == -1, inv(quat), result) + result = xp.where(n == 1, quat, result) + return result + if n == 0: + identity = xp.zeros((*quat.shape[:-1], 4), dtype=quat.dtype, device=device) + return xpx.at(identity)[..., 3].set(1) + if n == -1: + return inv(quat) + if n == 1: + return quat + return from_rotvec(n * as_rotvec(quat)) + + +def _normalize_quaternion(quat: Array) -> Array: + xp = array_namespace(quat) + quat_norm = xp_vector_norm(quat, axis=-1, keepdims=True, xp=xp) + zero_norm = quat_norm == 0 + if is_lazy_array(quat_norm): + quat = xp.where(zero_norm, xp.nan, quat) + elif xp.any(zero_norm): + raise ValueError("Found zero norm quaternions in `quat`.") + return quat / quat_norm + + +def _quat_canonical(quat: Array) -> Array: + xp = array_namespace(quat) + mask = quat[..., 3] < 0 + zero_w = quat[..., 3] == 0 + mask = mask | (zero_w & (quat[..., 0] < 0)) + zero_wx = zero_w & (quat[..., 0] == 0) + mask = mask | (zero_wx & (quat[..., 1] < 0)) + zero_wxy = zero_wx & (quat[..., 1] == 0) + mask = mask | (zero_wxy & (quat[..., 2] < 0)) + return xp.where(mask[..., None], -quat, quat) + + +def _elementary_basis_index(axis: str) -> int: + if axis == "x": + return 0 + elif axis == "y": + return 1 + elif axis == "z": + return 2 + raise ValueError(f"Expected axis to be from ['x', 'y', 'z'], got {axis}") + + +def _compute_davenport_from_quat( + quat: Array, + n1: Array, + n2: Array, + n3: Array, + extrinsic: bool, + suppress_warnings: bool, +) -> Array: + # The algorithm assumes extrinsic frame transformations. The algorithm + # in the paper is formulated for rotation quaternions, which are stored + # directly by Rotation. + # Adapt the algorithm for our case by reversing both axis sequence and + # angles for intrinsic rotations when needed + xp = array_namespace(quat) + n1, n3 = (n1, n3) if extrinsic else (n3, n1) + + n_cross = xp.linalg.cross(n1, n2) + lamb = xp.atan2(xp.vecdot(n3, n_cross), xp.vecdot(n3, n1)) + + # alternative set of angles compatible with as_euler implementation + mask = lamb < 0 + n2 = xp.where(mask[..., None], -n2, n2) + lamb = xp.where(mask, -lamb, lamb) + n_cross = xp.where(mask[..., None], -n_cross, n_cross) + correct_set = mask + + quat_lamb = xp.concat( + (xp.sin(lamb / 2)[..., None] * n2, xp.cos(lamb / 2)[..., None]), axis=-1 + ) + + q_trans = compose_quat(quat_lamb, quat) + a = q_trans[..., 3] + b = xp.linalg.vecdot(q_trans[..., :3], n1) + c = xp.linalg.vecdot(q_trans[..., :3], n2) + d = xp.linalg.vecdot(q_trans[..., :3], n_cross) + + angles = _get_angles(extrinsic, False, 1, lamb, a, b, c, d, suppress_warnings) + angles = xpx.at(angles)[..., 1].set( + xp.where(correct_set, -angles[..., 1], angles[..., 1]) + ) + + return angles + + +def _elementary_quat_compose(axes: list[int], angles: Array, intrinsic: bool) -> Array: + xp = array_namespace(angles) + device = xp_device(angles) + quat = _make_elementary_quat(axes[0], angles[..., 0], device=device, xp=xp) + for i in range(1, len(axes)): + ax_quat = _make_elementary_quat(axes[i], angles[..., i], device=device, xp=xp) + quat = compose_quat(quat, ax_quat) if intrinsic else compose_quat(ax_quat, quat) + return quat + + +def _make_elementary_quat(axis: int, angle: Array, device, xp) -> Array: + quat = xp.zeros((*angle.shape, 4), dtype=angle.dtype, device=device) + quat = xpx.at(quat)[..., 3].set(xp.cos(angle / 2.0)) + quat = xpx.at(quat)[..., axis].set(xp.sin(angle / 2.0)) + return quat + + +def _get_angles( + extrinsic: bool, + symmetric: bool, + sign: int, + lamb: float, + a: Array, + b: Array, + c: Array, + d: Array, + suppress_warnings: bool, +) -> Array: + xp = array_namespace(a) + device = xp_device(a) + eps = 1e-7 + half_sum = xp.atan2(b, a) + half_diff = xp.atan2(d, c) + # We zero-initialize to automatically cover singular cases where the second angle is + # not defined uniquely. + angles = xp.zeros((*a.shape, 3), dtype=a.dtype, device=device) + + angles = xpx.at(angles)[..., 1].set(2 * xp.atan2(xp.hypot(c, d), xp.hypot(a, b))) + + angle_first = 0 if extrinsic else 2 + angle_third = 2 if extrinsic else 0 + + # Check if the second angle is close to 0 or pi, causing a singularity. + # - Case 0: Second angle is neither close to 0 nor pi. + # - Case 1: Second angle is close to 0. + # - Case 2: Second angle is close to pi. + case1 = xp.abs(angles[..., 1]) <= eps + case2 = xp.abs(angles[..., 1] - xp.pi) <= eps + case0 = ~(case1 | case2) + if not suppress_warnings and not is_lazy_array(case0) and xp.any(~case0): + warnings.warn( + "Gimbal lock detected. Setting third angle to zero " + "since it is not possible to uniquely determine " + "all angles.", + stacklevel=3, + ) + + # This writes case1 into a0 where True and case2 everywhere else. This is sound + # because we later overwrite any values without singularity with the regular value + # of case0. The second angle is covered by default since we zero-initialized the + # second dimension. + a0 = xp.where(case1, 2 * half_sum, 2 * half_diff * (-1 if extrinsic else 1)) + angles = xpx.at(angles)[..., 0].set(a0) + + # We overwrite the values of angles without singularities (case0) + a1 = xp.where(case0, half_sum - half_diff, angles[..., angle_first]) + angles = xpx.at(angles)[..., angle_first].set(a1) + + # Same as above but for the third angle. We overwrite the non-singular case0 values + a3 = xp.where(case0, half_sum + half_diff, angles[..., angle_third]) + if not symmetric: + a3 = a3 * sign + angles = xpx.at(angles)[..., 1].set(angles[..., 1] - lamb) + angles = xpx.at(angles)[..., angle_third].set(a3) + + angles = (angles + xp.pi) % (2 * xp.pi) - xp.pi + return angles + + +def compose_quat(p: Array, q: Array) -> Array: + xp = array_namespace(p) + cross = xp.linalg.cross(p[..., :3], q[..., :3]) + qx = p[..., 3] * q[..., 0] + q[..., 3] * p[..., 0] + cross[..., 0] + qy = p[..., 3] * q[..., 1] + q[..., 3] * p[..., 1] + cross[..., 1] + qz = p[..., 3] * q[..., 2] + q[..., 3] * p[..., 2] + cross[..., 2] + qw = ( + p[..., 3] * q[..., 3] + - p[..., 0] * q[..., 0] + - p[..., 1] * q[..., 1] + - p[..., 2] * q[..., 2] + ) + quat = xp.stack([qx, qy, qz, qw], axis=-1) + return quat + + +def _split_rotation(q: Array, xp) -> tuple[Array, Array]: + q = xpx.atleast_nd(q, ndim=2, xp=xp) + return q[..., -1], q[..., :-1] + + +def _deg2rad(angles: Array) -> Array: + return angles * (np.pi / 180.0) + + +def _rad2deg(angles: Array) -> Array: + return angles * (180.0 / np.pi) diff --git a/v12x/pylib/scipy/spatial/transform/rotation.py b/v12x/pylib/scipy/spatial/transform/rotation.py new file mode 100644 index 0000000000000000000000000000000000000000..a719437415124f4afc842b98c708836c1fe68f22 --- /dev/null +++ b/v12x/pylib/scipy/spatial/transform/rotation.py @@ -0,0 +1,21 @@ +# This file is not meant for public use and will be removed in SciPy v2.0.0. +# Use the `scipy.spatial` namespace for importing the functions +# included below. + +from scipy._lib.deprecation import _sub_module_deprecation + + +__all__ = [ # noqa: F822 + 'Rotation', + 'Slerp', +] + + +def __dir__(): + return __all__ + + +def __getattr__(name): + return _sub_module_deprecation(sub_package="spatial.transform", module="rotation", + private_modules=["_rotation"], all=__all__, + attribute=name) diff --git a/v12x/pylib/scipy/spatial/transform/tests/__init__.py b/v12x/pylib/scipy/spatial/transform/tests/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/v12x/pylib/scipy/spatial/transform/tests/test_rigid_transform.py b/v12x/pylib/scipy/spatial/transform/tests/test_rigid_transform.py new file mode 100644 index 0000000000000000000000000000000000000000..23d38c12987415a366533f160910ab3006031564 --- /dev/null +++ b/v12x/pylib/scipy/spatial/transform/tests/test_rigid_transform.py @@ -0,0 +1,1531 @@ +import pickle +from itertools import product + +import pytest + +import numpy as np +from scipy.spatial.transform import Rotation, RigidTransform +from scipy.spatial.transform._rigid_transform import normalize_dual_quaternion +from scipy._lib._array_api import ( + is_lazy_array, + xp_vector_norm, + is_numpy, + xp_assert_close, + make_xp_test_case, + xp_assert_equal, + xp_promote +) +import scipy._external.array_api_extra as xpx + + +def rotation_to_xp(r: Rotation, xp): + dtype = xpx.default_dtype(xp) + return Rotation.from_quat(xp.asarray(r.as_quat(), dtype=dtype)) + + +def rigid_transform_to_xp(r: RigidTransform, xp): + dtype = xpx.default_dtype(xp) + return RigidTransform.from_matrix(xp.asarray(r.as_matrix(), dtype=dtype)) + + +@make_xp_test_case((RigidTransform, "as_matrix")) +def test_repr(xp): + actual = repr(RigidTransform.from_matrix(xp.eye(4))) + expected = """\ +RigidTransform.from_matrix(array([[1., 0., 0., 0.], + [0., 1., 0., 0.], + [0., 0., 1., 0.], + [0., 0., 0., 1.]]))""" + if is_numpy(xp): + assert actual == expected + else: + assert actual.startswith("RigidTransform.from_matrix(") + + tf = RigidTransform.from_matrix(xp.asarray(RigidTransform.identity(2).as_matrix())) + actual = repr(tf) + expected = """\ +RigidTransform.from_matrix(array([[[1., 0., 0., 0.], + [0., 1., 0., 0.], + [0., 0., 1., 0.], + [0., 0., 0., 1.]], + + [[1., 0., 0., 0.], + [0., 1., 0., 0.], + [0., 0., 1., 0.], + [0., 0., 0., 1.]]]))""" + if is_numpy(xp): + assert actual == expected + else: + assert actual.startswith("RigidTransform.from_matrix(") + + +@make_xp_test_case((RigidTransform, "from_rotation")) +@pytest.mark.parametrize("ndim", range(1, 4)) +def test_from_rotation(xp, ndim: int): + atol = 1e-12 + rng = np.random.default_rng(0) + shape = (ndim,) * (ndim - 1) + (4,) + r = rotation_to_xp(Rotation.from_quat(rng.normal(size=shape)), xp=xp) + tf = RigidTransform.from_rotation(r) + xp_assert_close(tf.as_matrix()[..., :3, :3], r.as_matrix(), atol=atol) + xp_assert_close(tf.as_matrix()[..., :3, 3], xp.zeros(shape[:-1] + (3,)), atol=atol) + xp_assert_close(tf.as_matrix()[..., 3, :3], xp.zeros(shape[:-1] + (3,)), atol=atol) + xp_assert_close(tf.as_matrix()[..., 3, 3], xp.ones(shape[:-1]), atol=atol) + assert tf.single == (ndim == 1) + + +@make_xp_test_case((RigidTransform, "from_translation")) +@pytest.mark.parametrize("ndim", range(1, 4)) +def test_from_translation(xp, ndim: int): + shape = (ndim,) * (ndim - 1) + t = xp.reshape(xp.arange(ndim ** (ndim-1) * 3), shape + (3,)) + tf = RigidTransform.from_translation(t) + + expected = xp.tile(xp.eye(4), shape + (1, 1)) + t_float = xp_promote(t, force_floating=True, xp=xp) + expected = xpx.at(expected)[..., :3, 3].set(t_float) + xp_assert_close(tf.as_matrix(), expected) + assert tf.single == (ndim == 1) + + +def test_from_translation_array_like(): + # Test single translation + t = [1, 2, 3] + tf = RigidTransform.from_translation(t) + tf_expected = RigidTransform.from_translation(np.array(t)) + xp_assert_close(tf.as_matrix(), tf_expected.as_matrix()) + assert tf.single + + # Test multiple translations + t = [[1, 2, 3], [4, 5, 6]] + tf = RigidTransform.from_translation(t) + tf_expected = RigidTransform.from_translation(np.array(t)) + xp_assert_close(tf.as_matrix(), tf_expected.as_matrix()) + assert not tf.single + + +@make_xp_test_case((RigidTransform, "from_matrix"), (RigidTransform, "as_matrix")) +@pytest.mark.parametrize("ndim", range(1, 4)) +def test_from_matrix(xp, ndim: int): + atol = 1e-12 + shape = (ndim,) * (ndim - 1) + (4, 4) + dtype = xpx.default_dtype(xp) + + matrix = xp.tile(xp.eye(4), shape[:-2] + (1, 1)) + t = xp.reshape(xp.arange(ndim ** (ndim-1) * 3, dtype=dtype), shape[:-2] + (3,)) + matrix = xpx.at(matrix)[..., :3, 3].set(t) + + tf = RigidTransform.from_matrix(matrix) + xp_assert_close(tf.as_matrix(), matrix, atol=atol) + assert tf.single == (ndim == 1) + + # Test non-1 determinant + matrix = xp.tile(xp.eye(4), shape[:-2] + (1, 1)) + matrix = xpx.at(matrix)[..., :3, :3].set(xp.eye(3) * 2.0) + tf = RigidTransform.from_matrix(matrix) + expected = xp.tile(xp.eye(4), shape[:-2] + (1, 1)) + xp_assert_close(tf.as_matrix(), expected, atol=atol) + + # Test non-orthogonal rotation matrix + matrix = xp.tile(xp.eye(4), shape[:-2] + (1, 1)) + # matrix is equivalent to [[1, 1, 0, 0], + # [0, 1, 0, 0], + # [0, 0, 1, 0], + # [0, 0, 0, 1]] + matrix = xpx.at(matrix)[..., 0, 1].set(1.0) + tf = RigidTransform.from_matrix(matrix) + expected = xp.tile(xp.eye(4), shape[:-2] + (1, 1)) + expected = xpx.at(expected)[..., 0, 0].set(0.894427) + expected = xpx.at(expected)[..., 0, 1].set(0.447214) + expected = xpx.at(expected)[..., 1, 0].set(-0.447214) + expected = xpx.at(expected)[..., 1, 1].set(0.894427) + xp_assert_close(tf.as_matrix(), expected, atol=1e-6) + + # Test invalid matrix + invalid = xp.tile(xp.eye(4), shape[:-2] + (1, 1)) + invalid = xpx.at(invalid)[..., 3, 3].set(2) # Invalid last row + if is_lazy_array(invalid): + tf = RigidTransform.from_matrix(invalid) + assert xp.all(xp.isnan(tf.as_matrix())) + else: + with pytest.raises(ValueError): + RigidTransform.from_matrix(invalid) + + +def test_from_matrix_array_like(): + # Test single transform matrix + matrix = [[1, 0, 0, 0], + [0, 1, 0, 0], + [0, 0, 1, 0], + [0, 0, 0, 1]] + expected = np.eye(4) + tf = RigidTransform.from_matrix(matrix) + xp_assert_close(tf.as_matrix(), expected) + assert tf.single + + # Test multiple transform matrices + matrices = [matrix, matrix] + tf = RigidTransform.from_matrix(matrices) + for i in range(len(matrices)): + xp_assert_close(tf.as_matrix()[i, ...], expected) + assert not tf.single + + +@make_xp_test_case((RigidTransform, "from_components")) +@pytest.mark.parametrize("r_ndim", range(1, 3)) +@pytest.mark.parametrize("t_ndim", range(1, 3)) +def test_from_components(xp, r_ndim: int, t_ndim: int): + atol = 1e-12 + dims = (6, 5, 4, 3) # Common shape + q_shape = dims[:r_ndim - 1][::-1] + (4,) + t_shape = dims[:t_ndim - 1][::-1] + (3,) + tf_shape = np.broadcast_shapes(q_shape[:-1], t_shape[:-1]) + (4, 4) + rng = np.random.default_rng(0) + + t = xp.reshape(xp.arange(np.prod(t_shape[:-1]) * 3), t_shape) + r = rotation_to_xp(Rotation.from_quat(rng.random(size=q_shape)), xp=xp) + tf = RigidTransform.from_components(t, r) + + expected = xp.zeros(tf_shape) + expected = xpx.at(expected)[..., :3, :3].set(r.as_matrix()) + t_float = xp_promote(t, force_floating=True, xp=xp) + expected = xpx.at(expected)[..., :3, 3].set(t_float) + expected = xpx.at(expected)[..., 3, 3].set(1) + xp_assert_close(tf.as_matrix(), expected, atol=atol) + assert tf.single == (r_ndim == 1 and t_ndim == 1) + + +def test_from_components_array_like(): + rng = np.random.default_rng(123) + # Test single rotation and translation + t = [1, 2, 3] + r = Rotation.random(rng=rng) + tf = RigidTransform.from_components(t, r) + tf_expected = RigidTransform.from_components(np.array(t), r) + xp_assert_close(tf.as_matrix(), tf_expected.as_matrix(), atol=1e-12) + assert tf.single + + # Test multiple rotations and translations + t = [[1, 2, 3], [4, 5, 6]] + r = Rotation.random(len(t), rng=rng) + tf = RigidTransform.from_components(t, r) + tf_expected = RigidTransform.from_components(np.array(t), r) + xp_assert_close(tf.as_matrix(), tf_expected.as_matrix(), atol=1e-12) + assert not tf.single + + +@make_xp_test_case((RigidTransform, "as_components")) +@pytest.mark.parametrize("ndim", range(1, 4)) +def test_as_components(xp, ndim): + dtype = xpx.default_dtype(xp) + atol = 1e-12 if dtype == xp.float64 else 1e-6 + shape = (ndim,) * (ndim - 1) + rng = np.random.default_rng(123) + t = xp.asarray(rng.normal(size=shape + (3,)), dtype=dtype) + r = rotation_to_xp(Rotation.from_quat(rng.random(shape + (4,))), xp=xp) + tf = RigidTransform.from_components(t, r) + new_t, new_r = tf.as_components() + assert xp.all(new_r.approx_equal(r, atol=atol)) + xp_assert_close(new_t, t, atol=atol) + + +@make_xp_test_case((RigidTransform, "from_exp_coords")) +@pytest.mark.parametrize("dim", range(1, 4)) +def test_from_exp_coords(xp, dim: int): + shape = (dim,) * (dim - 1) + # example from 3.3 of + # https://hades.mech.northwestern.edu/images/2/25/MR-v2.pdf + dtype = xpx.default_dtype(xp) + angle1 = np.deg2rad(30.0) + mat = xp.asarray([ + [np.cos(angle1), -np.sin(angle1), 0.0, 1.0], + [np.sin(angle1), np.cos(angle1), 0.0, 2.0], + [0.0, 0.0, 1.0, 0.0], + [0.0, 0.0, 0.0, 1.0] + ], dtype=dtype) + mat = xp.tile(mat, shape + (1, 1)) + tf1 = RigidTransform.from_matrix(mat) + angle2 = np.deg2rad(60.0) + mat = xp.asarray([ + [np.cos(angle2), -np.sin(angle2), 0.0, 2.0], + [np.sin(angle2), np.cos(angle2), 0.0, 1.0], + [0.0, 0.0, 1.0, 0.0], + [0.0, 0.0, 0.0, 1.0] + ], dtype=dtype) + mat = xp.tile(mat, shape + (1, 1)) + tf2 = RigidTransform.from_matrix(mat) + expected = tf2 * tf1.inv() + deg2rag = xp.asarray(np.deg2rad(30.0)) + exp_coords = deg2rag * xp.asarray([0.0, 0.0, 1.0, 3.37, -3.37, 0.0]) + exp_coords = xp.tile(exp_coords, shape + (1,)) + actual = RigidTransform.from_exp_coords(exp_coords) + xp_assert_close(actual.as_matrix(), expected.as_matrix(), atol=1e-2) + + # test cases generated by comparison to pytransform3d + exp_coords = xp.asarray([ + [-2.01041204, -0.52983629, 0.65773501, + 0.10386614, 0.05855009, 0.54959179], + [-0.22537438, -0.24132627, -2.4747121, + -0.09158594, 1.88075832, -0.03197204] + ]) + exp_coords = xp.tile(exp_coords, shape + (1,)) + expected_matrix = xp.asarray([ + [[0.76406621, 0.10504613, -0.63652819, -0.10209961], + [0.59956454, -0.47987325, 0.64050295, 0.40158789], + [-0.2381705, -0.87102639, -0.42963687, 0.19637636], + [0., 0., 0., 1.]], + [[-0.78446989, 0.61157488, 0.10287448, 1.33330055], + [-0.58017785, -0.78232107, 0.22664378, 0.52660831], + [0.21909052, 0.11810973, 0.96852952, -0.02968529], + [0., 0., 0., 1.]] + ]) + expected_matrix = xp.tile(expected_matrix, shape + (1, 1)) + xp_assert_close( + RigidTransform.from_exp_coords(exp_coords).as_matrix(), + expected_matrix, atol=1e-8) + + # identity + expected_matrix = xp.tile(xp.eye(4), shape + (1, 1)) + exp_coords = xp.zeros(shape + (6,), dtype=dtype) + xp_assert_close( + RigidTransform.from_exp_coords(exp_coords).as_matrix(), + expected_matrix, atol=1e-12) + + # only translation + expected_matrix = xp.asarray([ + [[1.0, 0.0, 0.0, 3.0], + [0.0, 1.0, 0.0, -5.4], + [0.0, 0.0, 1.0, 100.2], + [0.0, 0.0, 0.0, 1.0]], + [[1.0, 0.0, 0.0, -3.0], + [0.0, 1.0, 0.0, 13.3], + [0.0, 0.0, 1.0, 1.3], + [0.0, 0.0, 0.0, 1.0]] + ]) + expected_matrix = xp.tile(expected_matrix, shape + (1, 1, 1)) + exp_coords = xp.asarray([ + [0.0, 0.0, 0.0, 3.0, -5.4, 100.2], + [0.0, 0.0, 0.0, -3.0, 13.3, 1.3], + ]) + exp_coords = xp.tile(exp_coords, shape + (1, 1)) + actual = RigidTransform.from_exp_coords(exp_coords) + xp_assert_close(actual.as_matrix(), expected_matrix, atol=1e-12) + + # only rotation + angles = xp.asarray([[34, -12, 0.5], [-102, -55, 30]]) + angles = xp.tile(angles, shape + (1, 1)) + rot = Rotation.from_euler('zyx', angles, degrees=True) + rotvec = rot.as_rotvec() + expected_matrix = xp.tile(xp.eye(4), shape + (2, 1, 1)) + expected_matrix = xpx.at(expected_matrix)[..., :3, :3].set(rot.as_matrix()) + exp_coords = xp.concat((rotvec, xp.zeros_like(rotvec)), axis=-1) + actual = RigidTransform.from_exp_coords(exp_coords) + xp_assert_close(actual.as_matrix(), expected_matrix, atol=1e-12) + + +def test_from_exp_coords_array_like(): + rng = np.random.default_rng(123) + # Test single transform + t = np.array([1, 2, 3]) + r = Rotation.random(rng=rng) + tf_expected = RigidTransform.from_components(t, r) + exp_coords = tf_expected.as_exp_coords().tolist() + assert isinstance(exp_coords, list) + tf = RigidTransform.from_exp_coords(exp_coords) + xp_assert_close(tf.as_matrix(), tf_expected.as_matrix(), atol=1e-12) + + # Test multiple transforms + t = [[1, 2, 3], [4, 5, 6]] + r = Rotation.random(len(t), rng=rng) + tf_expected = RigidTransform.from_components(t, r) + exp_coords = tf_expected.as_exp_coords().tolist() + assert isinstance(exp_coords, list) + tf = RigidTransform.from_exp_coords(exp_coords) + xp_assert_close(tf.as_matrix(), tf_expected.as_matrix(), atol=1e-12) + + +@make_xp_test_case((RigidTransform, "as_exp_coords")) +@pytest.mark.parametrize("ndim", range(1, 4)) +def test_as_exp_coords(xp, ndim: int): + shape = (ndim,) * (ndim - 1) + # identity + expected = xp.zeros(shape + (6,)) + actual = RigidTransform.from_exp_coords(expected).as_exp_coords() + xp_assert_close(actual, expected, atol=1e-12) + + rng = np.random.default_rng(10) + + # pure rotation + rot_vec = xp.asarray(rng.normal(scale=0.1, size=shape + (1000, 3))) + tf = RigidTransform.from_rotation(Rotation.from_rotvec(rot_vec)) + exp_coords = tf.as_exp_coords() + xp_assert_close(exp_coords[..., :3], rot_vec, atol=1e-12) + expected = xp.zeros_like(rot_vec) + xp_assert_close(exp_coords[..., 3:], expected, atol=1e-16) + + # pure translation + translation = xp.asarray(rng.normal(scale=100.0, size=shape + (1000, 3))) + tf = RigidTransform.from_translation(translation) + exp_coords = tf.as_exp_coords() + xp_assert_close(exp_coords[..., :3], expected, atol=1e-16) + xp_assert_close(exp_coords[..., 3:], translation, atol=1e-15) + + +@make_xp_test_case((RigidTransform, "from_dual_quat")) +@pytest.mark.parametrize("ndim", range(1, 4)) +def test_from_dual_quat(xp, ndim: int): + dtype = xpx.default_dtype(xp) + atol = 1e-12 if dtype == xp.float64 else 1e-7 + shape = (ndim,) * (ndim - 1) + + # identity + dq = xp.asarray([0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0], dtype=dtype) + dq = xp.tile(dq, shape + (1,)) + expected = xp.tile(xp.eye(4), shape + (1, 1)) + xp_assert_close(RigidTransform.from_dual_quat(dq).as_matrix(), expected, atol=atol) + dq = xp.asarray([1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], dtype=dtype) + dq = xp.tile(dq, shape + (1,)) + xp_assert_close(RigidTransform.from_dual_quat(dq, scalar_first=True).as_matrix(), + expected, atol=atol) + + # only translation + dq = xp.asarray([0, 0, 0, 1, 0.25, 0.15, -0.7, 0], dtype=dtype) + dq = xp.tile(dq, shape + (1,)) + actual = RigidTransform.from_dual_quat(dq) + expected_matrix = xp.asarray([ + [1, 0, 0, 0.5], + [0, 1, 0, 0.3], + [0, 0, 1, -1.4], + [0, 0, 0, 1] + ]) + expected_matrix = xp.tile(expected_matrix, shape + (1, 1)) + xp_assert_close(actual.as_matrix(), expected_matrix, atol=atol) + dq = xp.asarray([1, 0, 0, 0, 0, 0.25, 0.15, -0.7], dtype=dtype) + dq = xp.tile(dq, shape + (1,)) + actual = RigidTransform.from_dual_quat(dq, scalar_first=True) + xp_assert_close(actual.as_matrix(), expected_matrix, atol=atol) + + # only rotation + angles = xp.asarray([65, -13, 90], dtype=dtype) + angles = xp.tile(angles, shape + (1,)) + actual_rot = Rotation.from_euler("xyz", angles, degrees=True) + qrot = actual_rot.as_quat() + dq = xp.concat((qrot, xp.zeros_like(qrot)), axis=-1) + actual = RigidTransform.from_dual_quat(dq) + expected_matrix = xp.tile(xp.eye(4), shape + (1, 1)) + expected_matrix = xpx.at(expected_matrix)[..., :3, :3].set(actual_rot.as_matrix()) + xp_assert_close(actual.as_matrix(), expected_matrix, atol=atol) + + qrot = actual_rot.as_quat(scalar_first=True) + dq = xp.concat((qrot, xp.zeros_like(qrot)), axis=-1) + actual = RigidTransform.from_dual_quat(dq, scalar_first=True) + expected_matrix = xp.tile(xp.eye(4), shape + (1, 1)) + expected_matrix = xpx.at(expected_matrix)[..., :3, :3].set(actual_rot.as_matrix()) + xp_assert_close(actual.as_matrix(), expected_matrix, atol=atol) + + # rotation and translation + # rtol is set to 1e-7 because xp_assert_close deviates from + # np.testing.assert_allclose in that it does not automatically default to 1e-7 for + # floating point inputs. + # See https://numpy.org/doc/2.2/reference/generated/numpy.testing.assert_allclose.html + dq = xp.asarray( + [[0.0617101, -0.06483886, 0.31432811, 0.94508498, + 0.04985168, -0.26119618, 0.1691491, -0.07743254], + [0.19507259, 0.49404931, -0.06091285, 0.8450749, + 0.65049656, -0.30782513, 0.16566752, 0.04174544]]) + dq = xp.tile(dq, shape + (1, 1)) + actual = RigidTransform.from_dual_quat(dq) + expected_matrix = xp.asarray( + [[[0.79398752, -0.60213598, -0.08376202, 0.24605262], + [0.58613113, 0.79477941, -0.15740392, -0.4932833], + [0.16135089, 0.07588122, 0.98397557, 0.34262676], + [0., 0., 0., 1.]], + [[0.50440981, 0.2957028, 0.81125249, 1.20934468], + [0.08979911, 0.91647262, -0.3898898, -0.70540077], + [-0.8587822, 0.26951399, 0.43572393, -0.47776265], + [0., 0., 0., 1.]]]) + expected_matrix = xp.tile(expected_matrix, shape + (1, 1, 1)) + xp_assert_close(actual.as_matrix(), expected_matrix, atol=atol, rtol=1e-7) + + dq = xp.asarray( + [[0.94508498, 0.0617101, -0.06483886, 0.31432811, + -0.07743254, 0.04985168, -0.26119618, 0.1691491], + [0.8450749, 0.19507259, 0.49404931, -0.06091285, + 0.04174544, 0.65049656, -0.30782513, 0.16566752]]) + dq = xp.tile(dq, shape + (1, 1)) + actual = RigidTransform.from_dual_quat(dq, scalar_first=True) + xp_assert_close(actual.as_matrix(), expected_matrix, atol=atol, rtol=1e-7) + + # unnormalized dual quaternions + + # invalid real quaternion with norm 0 + dq = xp.zeros(shape + (8,)) + actual = RigidTransform.from_dual_quat(dq) + expected = xp.tile(xp.eye(4), shape + (1, 1)) + xp_assert_close(actual.as_matrix(), expected, atol=atol) + + # real quaternion with norm != 1 + unnormalized_dual_quat = xp.asarray( + [-0.2547655, 1.23506123, 0.20230088, 0.24247194, # norm 1.3 + 0.38559628, 0.08184063, 0.1755943, -0.1582222] # orthogonal + ) + xp_assert_close(xp_vector_norm(unnormalized_dual_quat[:4]), xp.asarray(1.3)[()], + atol=atol) + xp_assert_close(xp.vecdot(unnormalized_dual_quat[:4], + unnormalized_dual_quat[4:])[()], + xp.asarray(0.0)[()], atol=1e-8) + + dq = xp.tile(unnormalized_dual_quat, shape + (1,)) + dual_quat = RigidTransform.from_dual_quat(dq).as_dual_quat() + + expected_ones = xp.ones(shape) if shape != () else xp.asarray(1.0)[()] + expected_zeros = xp.zeros(shape) if shape != () else xp.asarray(0.0)[()] + xp_assert_close(xp_vector_norm(dual_quat[..., :4], axis=-1), expected_ones, + atol=1e-12) + vecdot = xp.vecdot(dual_quat[..., :4], dual_quat[..., 4:]) + vecdot = vecdot[()] if vecdot.shape == () else vecdot + xp_assert_close(vecdot, expected_zeros, atol=atol) + + # real and dual quaternion are not orthogonal + unnormalized_dual_quat = xp.asarray( + [0.20824458, 0.75098079, 0.54542913, -0.30849493, # unit norm + -0.16051025, 0.10742978, 0.21277201, 0.20596935] # not orthogonal + ) + xp_assert_close(xp_vector_norm(unnormalized_dual_quat[:4]), xp.asarray(1.0)[()], + atol=atol) + assert xp.vecdot(unnormalized_dual_quat[:4], unnormalized_dual_quat[4:]) != 0.0 + dq = xp.tile(unnormalized_dual_quat, shape + (1,)) + dual_quat = RigidTransform.from_dual_quat(dq).as_dual_quat() + + xp_assert_close(xp_vector_norm(dual_quat[..., :4], axis=-1), expected_ones, + atol=1e-12) + vecdot = xp.vecdot(dual_quat[..., :4], dual_quat[..., 4:]) + vecdot = vecdot[()] if vecdot.shape == () else vecdot + xp_assert_close(vecdot, expected_zeros, atol=atol) + + # invalid real quaternion with norm 0, non-orthogonal dual quaternion + unnormalized_dual_quat = xp.asarray( + [0.0, 0.0, 0.0, 0.0, -0.16051025, 0.10742978, 0.21277201, 0.20596935]) + assert xp.vecdot(xp.asarray([0.0, 0, 0, 1]), unnormalized_dual_quat[4:]) != 0.0 + dq = xp.tile(unnormalized_dual_quat, shape + (1,)) + dual_quat = RigidTransform.from_dual_quat(dq).as_dual_quat() + + xp_assert_close(xp_vector_norm(dual_quat[..., :4], axis=-1), expected_ones, + atol=1e-12) + vecdot = xp.vecdot(dual_quat[..., :4], dual_quat[..., 4:]) + vecdot = vecdot[()] if vecdot.shape == () else vecdot + xp_assert_close(vecdot, expected_zeros, atol=atol) + + # compensation for precision loss in real quaternion + rng = np.random.default_rng(1000) + t = xp.asarray(rng.normal(size=shape + (3,)), dtype=dtype) + q = xp.asarray(rng.normal(size=shape + (4,)), dtype=dtype) + r = Rotation.from_quat(q) + random_dual_quats = RigidTransform.from_components(t, r).as_dual_quat() + + # ensure that random quaternions are not normalized + random_dual_quats = xpx.at(random_dual_quats)[..., :4].add(0.01) + assert not xp.any(xpx.isclose(xp_vector_norm(random_dual_quats[..., :4], axis=-1), + 1.0, atol=0.0001)) + dual_quat_norm = RigidTransform.from_dual_quat( + random_dual_quats).as_dual_quat() + xp_assert_close(xp_vector_norm(dual_quat_norm[..., :4], axis=-1), expected_ones, + atol=atol) + + # compensation for precision loss in dual quaternion, results in violation + # of orthogonality constraint + t = xp.asarray(rng.normal(size=shape + (3,)), dtype=dtype) + q = xp.asarray(rng.normal(size=shape + (4,)), dtype=dtype) + r = Rotation.from_quat(q) + random_dual_quats = RigidTransform.from_components(t, r).as_dual_quat() + + # ensure that random quaternions are not normalized + random_dual_quats = xpx.at(random_dual_quats)[..., 4:].add(0.1) + q_norm = xp.vecdot(random_dual_quats[..., :4], random_dual_quats[..., 4:]) + assert not xp.any(xpx.isclose(q_norm, 0.0, atol=0.0001)) + dual_quat_norm = RigidTransform.from_dual_quat( + random_dual_quats).as_dual_quat() + vecdot = xp.vecdot(dual_quat[..., :4], dual_quat[..., 4:]) + vecdot = vecdot[()] if vecdot.shape == () else vecdot + xp_assert_close(vecdot, expected_zeros, atol=atol) + xp_assert_close(random_dual_quats[..., :4], dual_quat_norm[..., :4], atol=atol) + + +def test_from_dual_quat_array_like(): + rng = np.random.default_rng(123) + # Test single transform + t = np.array([1, 2, 3]) + r = Rotation.random(rng=rng) + tf_expected = RigidTransform.from_components(t, r) + dual_quat = tf_expected.as_dual_quat().tolist() + assert isinstance(dual_quat, list) + tf = RigidTransform.from_dual_quat(dual_quat) + xp_assert_close(tf.as_matrix(), tf_expected.as_matrix(), atol=1e-12) + + # Test multiple transforms + t = [[1, 2, 3], [4, 5, 6]] + r = Rotation.random(len(t), rng=rng) + tf_expected = RigidTransform.from_components(t, r) + dual_quat = tf_expected.as_dual_quat().tolist() + assert isinstance(dual_quat, list) + tf = RigidTransform.from_dual_quat(dual_quat) + xp_assert_close(tf.as_matrix(), tf_expected.as_matrix(), atol=1e-12) + + +@make_xp_test_case((RigidTransform, "as_dual_quat")) +@pytest.mark.parametrize("ndim", range(1, 4)) +def test_as_dual_quat(xp, ndim: int): + dtype = xpx.default_dtype(xp) + shape = (ndim,) * (ndim - 1) + # identity + expected = xp.asarray([0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0], dtype=dtype) + actual = rigid_transform_to_xp(RigidTransform.identity(), xp).as_dual_quat() + xp_assert_close(actual, expected, atol=1e-12) + + expected = xp.asarray([1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]) + tf = rigid_transform_to_xp(RigidTransform.identity(), xp) + actual = tf.as_dual_quat(scalar_first=True) + xp_assert_close(actual, expected, atol=1e-12) + + rng = np.random.default_rng(10) + + # only rotation + for _ in range(10): + q = xp.asarray(rng.normal(size=shape + (4,)), dtype=dtype) + real_part = Rotation.from_quat(q).as_quat() + dual_part = xp.zeros_like(real_part) + expected = xp.concat((real_part, dual_part), axis=-1) + actual = RigidTransform.from_dual_quat(expected).as_dual_quat() + # because of double cover: + actual = actual * xp.sign(actual[..., 0, None]) + expected = expected * xp.sign(expected[..., 0, None]) + xp_assert_close(actual, expected, atol=1e-12) + + # only translation + for _ in range(10): + tf = 0.5 * xp.asarray(rng.normal(size=shape + (3,)), dtype=dtype) + expected = xp.zeros(shape + (8,), dtype=dtype) + expected = xpx.at(expected)[..., 3].set(1.0) + expected = xpx.at(expected)[..., 4:7].set(tf) + actual = RigidTransform.from_dual_quat(expected).as_dual_quat() + # because of double cover: + actual = actual * xp.sign(actual[..., 0, None]) + expected = expected * xp.sign(expected[..., 0, None]) + xp_assert_close(actual, expected, atol=1e-12) + + # rotation and translation + for _ in range(10): + t = xp.asarray(rng.normal(size=shape + (3,)), dtype=dtype) + r = Rotation.from_quat(xp.asarray(rng.normal(size=shape + (4,)), dtype=dtype)) + expected = RigidTransform.from_components(t, r).as_dual_quat() + actual = RigidTransform.from_dual_quat(expected).as_dual_quat() + # because of double cover: + actual = actual * xp.sign(actual[..., 0, None]) + expected = expected * xp.sign(expected[..., 0, None]) + xp_assert_close(actual, expected, atol=1e-12) + + +@make_xp_test_case( + (RigidTransform, "from_components"), (RigidTransform, "as_components"), + (RigidTransform, "from_exp_coords"), (RigidTransform, "as_exp_coords"), + (RigidTransform, "from_matrix"), (RigidTransform, "as_matrix"), + (RigidTransform, "from_dual_quat"), (RigidTransform, "as_dual_quat") +) +@pytest.mark.parametrize("ndim", range(1, 4)) +def test_from_as_internal_consistency(xp, ndim: int): + dtype = xpx.default_dtype(xp) + atol = 1e-12 + n = 10 + rng = np.random.default_rng(10) + shape = (n,) + (ndim,) * (ndim - 1) + + t = xp.asarray(rng.normal(size=shape + (3,)), dtype=dtype) + r = Rotation.from_quat(xp.asarray(rng.normal(size=shape + (4,)), dtype=dtype)) + tf0 = RigidTransform.from_components(t, r) + + tf1 = RigidTransform.from_components(*tf0.as_components()) + xp_assert_close(tf0.as_matrix(), tf1.as_matrix(), atol=atol) + + tf1 = RigidTransform.from_components(tf0.translation, tf0.rotation) + xp_assert_close(tf0.as_matrix(), tf1.as_matrix(), atol=atol) + + tf1 = RigidTransform.from_exp_coords(tf0.as_exp_coords()) + xp_assert_close(tf0.as_matrix(), tf1.as_matrix(), atol=atol) + + tf1 = RigidTransform.from_matrix(tf0.as_matrix()) + xp_assert_close(tf0.as_matrix(), tf1.as_matrix(), atol=atol) + + tf1 = RigidTransform.from_dual_quat(tf0.as_dual_quat()) + xp_assert_close(tf0.as_matrix(), tf1.as_matrix(), atol=atol) + + # exp_coords small rotation + t = xp.asarray(rng.normal(scale=1000.0, size=shape + (3,)), dtype=dtype) + rotvec = xp.asarray(rng.normal(scale=1e-10, size=shape + (3,)), dtype=dtype) + r = Rotation.from_rotvec(rotvec) + tf0 = RigidTransform.from_components(t, r) + tf1 = RigidTransform.from_exp_coords(tf0.as_exp_coords()) + xp_assert_close(tf0.as_matrix(), tf1.as_matrix(), atol=atol) + + +def test_identity(): + # We do not use xp here because identity always returns numpy arrays + atol = 1e-12 + # Test single identity + tf = RigidTransform.identity() + xp_assert_close(tf.as_matrix(), np.eye(4), atol=atol) + # Test multiple identities + tf = RigidTransform.identity(5) + xp_assert_close(tf.as_matrix(), np.array([np.eye(4)] * 5), atol=atol) + # Test shape + tf = RigidTransform.identity(shape=3) + expected = np.tile(np.eye(4), (3, 1, 1)) + xp_assert_close(tf.as_matrix(), expected, atol=atol) + tf = RigidTransform.identity(shape=(2, 3)) + expected = np.tile(np.eye(4), (2, 3, 1, 1)) + xp_assert_close(tf.as_matrix(), expected, atol=atol) + # Test errors + with pytest.raises(ValueError, match="Only one of `num` and `shape` can be."): + RigidTransform.identity(10, shape=(2, 3)) + with pytest.raises(TypeError, match="takes from 0 to 1 positional arguments"): + RigidTransform.identity(None, (-1, 3)) # Shape is kwarg only + with pytest.raises(ValueError, match="`shape` must be an int or a tuple of ints"): + RigidTransform.identity(shape="invalid") + + +@make_xp_test_case((RigidTransform, "apply")) +def test_apply(xp): + atol = 1e-12 + # Broadcast shape: (6, 5, 4, 2) ( + (3,) for vectors, + (4,) for rotations) + vector_shapes = [(), (1,), (2,), (1, 2), (5, 1, 2)] + tf_shapes = [(), (1,), (2,), (1, 2), (4, 2), (1, 4, 2), (5, 4, 2), (6, 5, 4, 2)] + rng = np.random.default_rng(123) + + for tf_shape, v_shape in product(tf_shapes, vector_shapes): + # Random rotation and translation + t = xp.asarray(rng.normal(size=tf_shape + (3,))) + q = xp.asarray(rng.normal(size=tf_shape + (4,))) + r = Rotation.from_quat(q) + tf = RigidTransform.from_components(t, r) + + vecs = xp.asarray(rng.normal(size=v_shape + (3,))) + expected = t + r.apply(vecs) + res = tf.apply(vecs) + assert res.shape == np.broadcast_shapes(tf_shape, v_shape) + (3,) + xp_assert_close(res, expected, atol=atol) + + +def test_apply_array_like(): + rng = np.random.default_rng(123) + # Single vector + t = np.array([1, 2, 3]) + r = Rotation.random(rng=rng) + tf = RigidTransform.from_components(t, r) + vec = [1, 0, 0] + expected = t + r.apply(vec) + xp_assert_close(tf.apply(vec), expected, atol=1e-12) + + # Multiple vectors + t = np.array([[1, 2, 3], [4, 5, 6]]) + r = Rotation.random(len(t), rng=rng) + tf = RigidTransform.from_components(t, r) + vec = [[1, 0, 0], [0, 1, 0]] + expected = t + r.apply(vec) + xp_assert_close(tf.apply(vec), expected, atol=1e-12) + + +def test_apply_matrix_equivalence(): + """Test documented equivalence for single transform: + `apply(vector) == translation + vector @ rotation.as_matrix().T.`""" + t = np.array([1.0, 2.0, 3.0]) + r = Rotation.from_rotvec([0, 0, 1]) + tf = RigidTransform.from_components(t, r) + # Single vector (3,) + v = np.array([1.0, 0.0, 0.0]) + xp_assert_close(tf.apply(v), t + v @ r.as_matrix().T) + # Multiple vectors (P, 3) + arr = np.array([[1, 0, 0], [0, 1, 0]], dtype=float) + xp_assert_close(tf.apply(arr), t + arr @ r.as_matrix().T) + + +@make_xp_test_case((RigidTransform, "apply")) +def test_inverse_apply(xp): + atol = 1e-12 + # Broadcast shape: (6, 5, 4, 2) ( + (3,) for vectors, + (4,) for rotations) + vector_shapes = [(), (1,), (2,), (1, 2), (5, 1, 2)] + tf_shapes = [(), (1,), (2,), (1, 2), (4, 2), (1, 4, 2), (5, 4, 2), (6, 5, 4, 2)] + rng = np.random.default_rng(123) + + for tf_shape, v_shape in product(tf_shapes, vector_shapes): + # Random rotation and translation + t = xp.asarray(rng.normal(size=tf_shape + (3,))) + q = xp.asarray(rng.normal(size=tf_shape + (4,))) + r = Rotation.from_quat(q) + tf = RigidTransform.from_components(t, r) + + vecs = xp.asarray(rng.normal(size=v_shape + (3,))) + expected = tf.inv().apply(vecs) + res = tf.apply(vecs, inverse=True) + assert res.shape == np.broadcast_shapes(tf_shape, v_shape) + (3,) + xp_assert_close(res, expected, atol=atol) + + +@make_xp_test_case((RigidTransform, "apply")) +def test_rotation_alone(xp): + atol = 1e-12 + + r = Rotation.from_euler('z', xp.asarray(90), degrees=True) + tf = RigidTransform.from_rotation(r) + vec = xp.asarray([1, 0, 0]) + expected = r.apply(vec) + xp_assert_close(tf.apply(vec), expected, atol=atol) + + +@make_xp_test_case((RigidTransform, "apply")) +def test_translation_alone(xp): + atol = 1e-12 + t = xp.asarray([1.0, 2, 3]) + tf = RigidTransform.from_translation(t) + vec = xp.asarray([5.0, 6, 7]) + expected = t + vec + xp_assert_close(tf.apply(vec), expected, atol=atol) + + +@make_xp_test_case((RigidTransform, "apply"), (RigidTransform, "__mul__")) +def test_composition(xp): + atol = 1e-12 + tf_shapes = [(), (1,), (2,), (1, 2), (4, 2), (5, 4, 2)] + dtype = xpx.default_dtype(xp) + rng = np.random.default_rng(123) + + for tf_shape1, tf_shape2 in product(tf_shapes, repeat=2): + # Random rotation and translation + t1 = xp.asarray(rng.normal(size=tf_shape1 + (3,))) + q1 = xp.asarray(rng.normal(size=tf_shape1 + (4,))) + r1 = Rotation.from_quat(q1) + tf1 = RigidTransform.from_components(t1, r1) + + t2 = xp.asarray(rng.normal(size=tf_shape2 + (3,))) + q2 = xp.asarray(rng.normal(size=tf_shape2 + (4,))) + r2 = Rotation.from_quat(q2) + tf2 = RigidTransform.from_components(t2, r2) + + composed = tf2 * tf1 + vec = xp.asarray(rng.normal(size=(3,)), dtype=dtype) + expected = tf2.apply(tf1.apply(vec)) + res = composed.apply(vec) + assert res.shape == np.broadcast_shapes(tf_shape1, tf_shape2) + (3,) + xp_assert_close(res, expected, atol=atol) + + expected = t2 + r2.apply(t1 + r1.apply(vec)) + xp_assert_close(composed.apply(vec), expected, atol=atol) + assert composed.single == (tf1.single and tf2.single) + + +@make_xp_test_case((RigidTransform, "__pow__"), (RigidTransform, "__mul__")) +@pytest.mark.parametrize("ndim", range(1, 4)) +def test_pow(xp, ndim: int): + dtype = xpx.default_dtype(xp) + atol = 1e-12 if dtype == xp.float64 else 1e-6 + num = 10 + rng = np.random.default_rng(100) + shape = (num,) + (ndim,) * (ndim - 1) + t = xp.asarray(rng.normal(size=shape + (3,)), dtype=dtype) + q = xp.asarray(rng.normal(size=shape + (4,)), dtype=dtype) + r = Rotation.from_quat(q) + p = RigidTransform.from_components(t, r) + p_inv = p.inv() + + # Test the short-cuts and other integers + for n in [-5, -2, -1, 0, 1, 2, 5]: + q = p**n + # Regression test for gh-24436 + assert isinstance(q._matrix, type(p._matrix)) + r = RigidTransform.from_matrix(xp.tile(xp.eye(4), shape + (1, 1))) + for _ in range(abs(n)): + if n > 0: + r = r * p + else: + r = r * p_inv + xp_assert_close(q.as_matrix(), r.as_matrix(), atol=atol) + + # Test shape preservation of single + single_tf = RigidTransform.identity() + assert (single_tf**n).as_matrix().shape == (4, 4) + + # Test fractional powers + q = p**0.5 + xp_assert_close((q * q).as_matrix(), p.as_matrix(), atol=atol) + q = p**-0.5 + xp_assert_close((q * q).as_matrix(), p.inv().as_matrix(), atol=atol) + q = p** 1.5 + xp_assert_close((q * q).as_matrix(), (p**3).as_matrix(), atol=atol) + q = p** -1.5 + xp_assert_close((q * q).as_matrix(), (p**-3).as_matrix(), atol=atol) + + # pow function + tf = pow(RigidTransform.from_matrix(xp.eye(4)), 2) + xp_assert_close(tf.as_matrix(), xp.eye(4), atol=atol) + + +@make_xp_test_case((RigidTransform, "__pow__")) +@pytest.mark.parametrize("ndim", range(1, 4)) +def test_pow_equivalence_with_rotation(xp, ndim: int): + atol = 1e-12 + num = 10 + rng = np.random.default_rng(100) + dtype = xpx.default_dtype(xp) + shape = (num,) + (ndim,) * (ndim - 1) + + r = Rotation.from_quat(xp.asarray(rng.normal(size=shape + (4,)), dtype=dtype)) + p = RigidTransform.from_rotation(r) + for n in [-5, -2, -1.5, -1, -0.5, 0.0, 0.5, 1, 1.5, 2, 5]: + xp_assert_close((p**n).rotation.as_matrix(), (r**n).as_matrix(), atol=atol) + + +@make_xp_test_case((RigidTransform, "inv"), (RigidTransform, "__mul__")) +@pytest.mark.parametrize("ndim", range(1, 4)) +def test_inverse(xp, ndim: int): + dtype = xpx.default_dtype(xp) + atol = 1e-12 if dtype == xp.float64 else 1e-6 + rng = np.random.default_rng(100) + shape = (ndim,) * (ndim - 1) + + # Test inverse transform + r = Rotation.from_quat(xp.asarray(rng.normal(size=shape + (4,)), dtype=dtype)) + t = xp.asarray(rng.normal(size=shape + (3,)), dtype=dtype) + tf = RigidTransform.from_components(t, r) + + # Test that tf * tf.inv() equals identity + tf_inv = tf.inv() + composed = tf * tf_inv + expected = xp.tile(xp.eye(4), shape + (1, 1)) + xp_assert_close(composed.as_matrix(), expected, atol=atol) + + +@make_xp_test_case((RigidTransform, "as_matrix")) +@pytest.mark.parametrize("ndim", range(1, 4)) +def test_properties(xp, ndim: int): + atol = 1e-12 if xpx.default_dtype(xp) == xp.float64 else 1e-6 + shape = (ndim,) * (ndim - 1) + dtype = xpx.default_dtype(xp) + rng = np.random.default_rng(100) + + # Test rotation and translation properties + r = Rotation.from_quat(xp.asarray(rng.normal(size=shape + (4,)), dtype=dtype)) + t = xp.asarray(rng.normal(size=shape + (3,)), dtype=dtype) + tf = RigidTransform.from_components(t, r) + + xp_assert_close(tf.rotation.as_matrix(), r.as_matrix(), atol=atol) + assert xp.all(tf.rotation.approx_equal(r, atol=atol)) + xp_assert_close(tf.translation, t, atol=atol) + # Test that we don't return views that would modify the original array + xpx.at(tf.translation)[..., 0].set(0.0) + xp_assert_close(tf.translation, t, atol=atol) + assert tf.single == (shape == ()) + + +@make_xp_test_case((RigidTransform, "__getitem__")) +def test_indexing(xp): + atol = 1e-12 + + # Test indexing for multiple transforms + r = Rotation.from_euler('zyx', xp.asarray([[90, 0, 0], [0, 90, 0]]), degrees=True) + t = xp.asarray([[1.0, 2, 3], [4, 5, 6]]) + tf = RigidTransform.from_components(t, r) + + # Test single index + xp_assert_close(tf[0].as_matrix()[:3, :3], r[0].as_matrix(), atol=atol) + xp_assert_close(tf[0].as_matrix()[:3, 3], t[0, ...], atol=atol) + + # Test slice + tf_slice = tf[0:2] + xp_assert_close(tf_slice.as_matrix()[:, :3, :3], r[0:2].as_matrix(), atol=atol) + xp_assert_close(tf_slice.as_matrix()[:, :3, 3], t[0:2, ...], atol=atol) + + # Test boolean indexing + tf_masked = tf[xp.asarray([True, True])] + xp_assert_close(tf_masked.as_matrix()[:, :3, :3], r.as_matrix(), atol=atol) + xp_assert_close(tf_masked.as_matrix()[:, :3, 3], t, atol=atol) + + tf_masked = tf[xp.asarray([False, True])] + xp_assert_close(tf_masked.as_matrix()[:, :3, :3], + r[xp.asarray([False, True])].as_matrix(), atol=atol) + xp_assert_close(tf_masked.as_matrix()[:, :3, 3], t[xp.asarray([False, True])], + atol=atol) + + tf_masked = tf[xp.asarray([False, False])] + assert len(tf_masked) == 0 + + # Test integer array indexing + idx = xp.asarray([0, 1]) + xp_assert_close(tf[idx].as_matrix()[:, :3, :3], r[idx].as_matrix(), atol=atol) + xp_assert_close(tf[idx].as_matrix()[:, :3, 3], t, atol=atol) + + +def test_indexing_array_like(): + atol = 1e-12 + + r = Rotation.from_euler('zyx', np.array([[90, 0, 0], [0, 90, 0]]), degrees=True) + t = np.array([[1.0, 2, 3], [4, 5, 6]]) + tf = RigidTransform.from_components(t, r) + + tf_masked = tf[[False, True]] + xp_assert_close(tf_masked.as_matrix()[:, :3, :3], r[[False, True]].as_matrix(), + atol=atol) + xp_assert_close(tf_masked.as_matrix()[:, :3, 3], t[[False, True]], atol=atol) + tf_masked = tf[[False, False]] + assert len(tf_masked) == 0 + + +@make_xp_test_case((RigidTransform, "concatenate")) +def test_concatenate(xp): + atol = 1e-12 + + # Test concatenation of transforms + t1 = xp.asarray([1, 0, 0]) + r1 = Rotation.from_euler('z', xp.asarray(90), degrees=True) + tf1 = RigidTransform.from_components(t1, r1) + + t2 = xp.asarray([0, 1, 0]) + r2 = Rotation.from_euler('x', xp.asarray(90), degrees=True) + tf2 = RigidTransform.from_components(t2, r2) + + # Concatenate single transforms + concatenated1 = RigidTransform.concatenate([tf1, tf2]) + xp_assert_close(concatenated1[0].as_matrix(), tf1.as_matrix(), atol=atol) + xp_assert_close(concatenated1[1].as_matrix(), tf2.as_matrix(), atol=atol) + + # Concatenate multiple transforms + concatenated2 = RigidTransform.concatenate([tf1, concatenated1]) + xp_assert_close(concatenated2[0].as_matrix(), tf1.as_matrix(), atol=atol) + xp_assert_close(concatenated2[1].as_matrix(), tf1.as_matrix(), atol=atol) + xp_assert_close(concatenated2[2].as_matrix(), tf2.as_matrix(), atol=atol) + + # Test ND concatenation + tf3 = RigidTransform.from_translation(xp.reshape(xp.arange(18), (3, 2, 3))) + tf4 = RigidTransform.from_translation(xp.reshape(xp.arange(18) + 18, (3, 2, 3))) + concatenated3 = RigidTransform.concatenate([tf3, tf4]) + xp_assert_close(concatenated3.as_matrix()[:3, ...], tf3.as_matrix(), atol=atol) + xp_assert_close(concatenated3.as_matrix()[3:, ...], tf4.as_matrix(), atol=atol) + + +@make_xp_test_case((RigidTransform, "from_matrix")) +def test_input_validation(xp): + # Test invalid matrix shapes + inputs = [xp.eye(3), xp.zeros((4, 3)), []] + for input in inputs: + with pytest.raises(ValueError, match="Expected `matrix` to have shape"): + RigidTransform.from_matrix(input) + + # Test invalid last row + for ndim in range(3): + shape = (ndim,) * (ndim - 1) + matrix = xp.zeros(shape + (4, 4)) + matrix = xpx.at(matrix)[...].set(xp.eye(4)) + matrix = xpx.at(matrix)[..., 3, :].set(xp.asarray([1.0, 0, 0, 1])) + if is_lazy_array(matrix): + matrix = RigidTransform.from_matrix(matrix).as_matrix() + assert xp.all(xp.isnan(matrix[..., 3, :])) + else: + with pytest.raises(ValueError, match="last row of transformation matrix"): + RigidTransform.from_matrix(matrix) + + # Test left handed rotation matrix + matrix = xp.eye(4) + matrix = xpx.at(matrix)[0, 0].set(-1) + if is_lazy_array(matrix): + matrix = RigidTransform.from_matrix(matrix).as_matrix() + assert xp.all(xp.isnan(matrix[..., :3, :3])) + else: + with pytest.raises(ValueError, match="Non-positive determinant"): + RigidTransform(matrix, normalize=True) + + # Test non-Rotation input + with pytest.raises(TypeError, + match="Expected `rotation` to be a `Rotation` instance"): + RigidTransform.from_rotation(xp.eye(3)) + + +@make_xp_test_case((RigidTransform, "mean")) +@pytest.mark.parametrize("ndim", range(1, 4)) +def test_mean(xp, ndim: int): + atol = 1e-12 + rng = np.random.default_rng(123) + + dtype = xpx.default_dtype(xp) + t = xp.asarray(rng.normal(size=(ndim,) * (ndim - 1) + (3,)), dtype=dtype) + q = xp.asarray(rng.normal(size=(ndim,) * (ndim - 1) + (4,)), dtype=dtype) + r = Rotation.from_quat(q) + tf = RigidTransform.from_components(t, r) + + # Unweighted mean + axis = tuple(range(t.ndim - 1)) + t_mean = xp.mean(t, axis=axis) + r_mean = r.mean() + tf_mean = tf.mean() + assert tf_mean.shape == () + xp_assert_close(tf_mean.as_matrix(), + RigidTransform.from_components(t_mean, r_mean).as_matrix(), + atol=atol) + + # Weighted mean + if ndim == 1: + weights = None + t_mean = t + else: + weights = xp.asarray(rng.random(size=(ndim,) * (ndim - 1)), dtype=dtype) + norm = xp.sum(weights[..., None], axis=axis) + wsum = xp.sum(t * weights[..., None], axis=axis) + t_mean = wsum/norm + r_mean = r.mean(weights=weights) + tf_mean = tf.mean(weights=weights) + assert tf_mean.shape == () + xp_assert_close(tf_mean.as_matrix(), + RigidTransform.from_components(t_mean, r_mean).as_matrix(), + atol=atol) + + +@make_xp_test_case( + (RigidTransform, "from_rotation"), (RigidTransform, "mean"), (Rotation, "magnitude") +) +@pytest.mark.parametrize("ndim", range(1, 5)) +def test_mean_axis(xp, ndim: int): + axes = xp.tile(xp.concat((-xp.eye(3), xp.eye(3))), (3,) * (ndim - 1) + (1, 1)) + theta = xp.pi / 4 + r = Rotation.from_rotvec(theta * axes) + tf = RigidTransform.from_rotation(r) + + # Test mean over last axis + desired = xp.full(axes.shape[:-2], 0.0) + if ndim == 1: + desired = desired[()] + atol = 1e-6 if xpx.default_dtype(xp) is xp.float32 else 1e-10 + xp_assert_close(tf.mean(axis=-1).rotation.magnitude(), desired, atol=atol) + + # Test tuple axes + desired = xp.full(axes.shape[1:-2], 0.0) + if ndim < 3: + desired = desired[()] + xp_assert_close(tf.mean(axis=(0, -1)).rotation.magnitude(), desired, atol=atol) + + # Empty axis tuple should return RigidTransform unchanged + tf_mean = tf.mean(axis=()) + xp_assert_close(tf_mean.as_matrix(), tf.as_matrix(), atol=atol) + + +@make_xp_test_case((RigidTransform, "mean"), (Rotation, "magnitude")) +def test_mean_compare_axis(xp): + # Create a random set of transforms and compare the mean over an axis with + # the mean without axis of the sliced transform + atol = 1e-10 if xpx.default_dtype(xp) == xp.float64 else 1e-6 + rng = np.random.default_rng(0) + q = xp.asarray(rng.normal(size=(4, 5, 6, 4)), dtype=xpx.default_dtype(xp)) + r = Rotation.from_quat(q) + t = xp.asarray(rng.normal(size=(4, 5, 6, 3)), dtype=xpx.default_dtype(xp)) + tf = RigidTransform.from_components(t, r) + + mean_0 = tf.mean(axis=0) + for i in range(q.shape[1]): + for j in range(q.shape[2]): + r_slice = Rotation.from_quat(q[:, i, j, ...]) + t_slice = t[:, i, j, ...] + mean_slice_tf = RigidTransform.from_components(t_slice, r_slice).mean() + xp_assert_close( + (mean_0[i][j].rotation * mean_slice_tf.rotation.inv()).magnitude(), + xp.asarray(0.0)[()], atol=atol, + ) + xp_assert_close( + mean_0[i][j].translation, mean_slice_tf.translation, atol=atol, + ) + mean_1_2 = tf.mean(axis=(1, 2)) + for i in range(q.shape[0]): + r_slice = Rotation.from_quat(q[i, ...]) + t_slice = t[i, ...] + mean_slice_tf = RigidTransform.from_components(t_slice, r_slice).mean() + xp_assert_close( + (mean_1_2[i].rotation * mean_slice_tf.rotation.inv()).magnitude(), + xp.asarray(0.0)[()], atol=atol, + ) + xp_assert_close( + mean_1_2[i].translation, mean_slice_tf.translation, atol=atol, + ) + + +@make_xp_test_case((RigidTransform, "mean")) +def test_mean_invalid_weights(xp): + tf = RigidTransform.from_matrix(xp.tile(xp.eye(4), (4, 1, 1))) + if is_lazy_array(tf.as_matrix()): + m = tf.mean(weights=-xp.ones(4)) + assert xp.all(xp.isnan(m.as_matrix())) + else: + with pytest.raises(ValueError, match="non-negative"): + tf.mean(weights=-xp.ones(4)) + + # Test weight shape mismatch + tf = RigidTransform.from_matrix(xp.eye(4)) + with pytest.raises(ValueError, match="Expected `weights` to"): + tf.mean(weights=xp.ones((2,))) + tf = RigidTransform.from_matrix(xp.tile(xp.eye(4), (3, 2, 1, 1, 1))) + with pytest.raises(ValueError, match="Expected `weights` to"): + tf.mean(weights=xp.ones((2, 1))) + + +@make_xp_test_case((RigidTransform, "from_translation")) +def test_translation_validation(xp): + # Test invalid translation shapes + with pytest.raises(ValueError, match="Expected `translation` to have shape"): + RigidTransform.from_translation(xp.asarray([1, 2])) + + with pytest.raises(ValueError, match="Expected `translation` to have shape"): + RigidTransform.from_translation(xp.zeros((2, 2))) + + +@make_xp_test_case((RigidTransform, "apply")) +def test_vector_validation(xp): + tf = rigid_transform_to_xp(RigidTransform.identity(2), xp=xp) + + # Test invalid vector shapes + with pytest.raises(ValueError, match="Expected vector to have shape"): + tf.apply(xp.asarray([1, 2])) + + with pytest.raises(ValueError, match="Expected vector to have shape"): + tf.apply(xp.zeros((2, 2))) + + with pytest.raises(ValueError, match="operands could not be broadcast"): + tf.apply(xp.zeros((1, 4, 3))) + + +@make_xp_test_case((RigidTransform, "__getitem__")) +def test_indexing_validation(xp): + tf = RigidTransform.from_matrix(xp.eye(4)) + + # Test indexing on single transform + with pytest.raises(TypeError, match="Single transform is not subscriptable"): + tf[0] + + with pytest.raises(TypeError, match="Single transform is not subscriptable"): + tf[0:1] + + # Test length on single transform + with pytest.raises(TypeError, match="Single transform has no len"): + len(tf) + + +@make_xp_test_case((RigidTransform, "__mul__")) +def test_composition_validation(xp): + tf2 = RigidTransform.from_translation(xp.asarray([[1, 2, 3], [4, 5, 6]])) + tf3 = RigidTransform.from_translation(xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]])) + + # Test incompatible shapes + with pytest.raises(ValueError, match="Cannot broadcast"): + tf2 * tf3 + + tf4 = RigidTransform.from_matrix(xp.tile(xp.eye(4), (1, 4, 1, 1))) + # Test invalid broadcasting shape + with pytest.raises(ValueError, match="Cannot broadcast"): + tf2 * tf4 + + +@make_xp_test_case((RigidTransform, "__mul__"), (RigidTransform, "__rmul__")) +def test_rotation_promotion(xp): + """Test that Rotation is promoted to RigidTransform in composition.""" + atol = 1e-12 + dtype = xpx.default_dtype(xp) + rng = np.random.default_rng(0) + + t = xp.asarray(rng.normal(size=(3,)), dtype=dtype) + r1 = rotation_to_xp(Rotation.random(rng=rng), xp) + r2 = rotation_to_xp(Rotation.random(rng=rng), xp) + tf = RigidTransform.from_components(t, r2) + + # RigidTransform * Rotation + result = tf * r1 + expected = tf * RigidTransform.from_rotation(r1) + assert isinstance(result, RigidTransform) + xp_assert_close(result.as_matrix(), expected.as_matrix(), atol=atol) + + # Rotation * RigidTransform + result = r1 * tf + expected = RigidTransform.from_rotation(r1) * tf + assert isinstance(result, RigidTransform) + xp_assert_close(result.as_matrix(), expected.as_matrix(), atol=atol) + + +@make_xp_test_case((RigidTransform, "concatenate")) +def test_concatenate_validation(xp): + tf = RigidTransform.from_matrix(xp.eye(4)) + + # Test invalid inputs + with pytest.raises(TypeError, + match="input must contain RigidTransform objects"): + RigidTransform.concatenate([tf, xp.eye(4)]) + + # Test incompatible shapes + tf2 = RigidTransform.from_translation(xp.ones((1, 1, 1, 3))) + # Frameworks have a highly heterogeneous way of reporting errors for this case + with pytest.raises((ValueError, TypeError, RuntimeError)): + RigidTransform.concatenate([tf, tf2]) + + +@make_xp_test_case((RigidTransform, "__setitem__")) +def test_setitem(xp): + tf = RigidTransform.from_translation(xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]])) + single = RigidTransform.from_translation(xp.asarray([1, 1, 1])) + double = RigidTransform.from_translation(xp.asarray([[2, 2, 2], [3, 3, 3]])) + triple = RigidTransform.from_translation(xp.asarray([[3, 3, 3], + [4, 4, 4], + [5, 5, 5]])) + + # Test indexing with integer index + tf[0] = single + xp_assert_close(tf.translation, xp.asarray([[1.0, 1, 1], [4, 5, 6], [7, 8, 9]])) + + # Test indexing with slice + tf = RigidTransform.from_translation(xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]])) + tf[:2] = double + xp_assert_close(tf.translation, xp.asarray([[2.0, 2, 2], [3, 3, 3], [7, 8, 9]])) + + # Test indexing with ellipsis + tf = RigidTransform.from_translation(xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]])) + tf[...] = triple + xp_assert_close(tf.translation, xp.asarray([[3.0, 3, 3], [4, 4, 4], [5, 5, 5]])) + + # Test indexing with boolean array + tf = RigidTransform.from_translation(xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]])) + mask = xp.asarray([True, False, True]) + tf[mask] = double + xp_assert_close(tf.translation, xp.asarray([[2.0, 2, 2], [4, 5, 6], [3, 3, 3]])) + + +@make_xp_test_case((RigidTransform, "__setitem__")) +@pytest.mark.skip_xp_backends("array_api_strict", + reason="doesn't support fancy indexing __setitem__") +def test_setitem_fancy_indexing(xp): + double = RigidTransform.from_translation(xp.asarray([[2, 2, 2], [3, 3, 3]])) + tf = RigidTransform.from_translation(xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]])) + idx = xp.asarray([0, 2]) + tf[idx] = double + xp_assert_close(tf.translation, xp.asarray([[2.0, 2, 2], [4, 5, 6], [3, 3, 3]])) + + +@make_xp_test_case((RigidTransform, "__setitem__")) +def test_setitem_validation(xp): + tf = RigidTransform.from_translation(xp.asarray([[1, 2, 3], [4, 5, 6]])) + single = RigidTransform.from_matrix(xp.eye(4)) + + # Test setting item on single transform + with pytest.raises(TypeError, match="Single transform is not subscriptable"): + single[0] = tf + + # Test invalid value type + with pytest.raises(TypeError, match="value must be a RigidTransform"): + tf[0] = xp.eye(4) + + +@pytest.mark.skip_xp_backends("jax.numpy", + reason="JAX does not support memory sharing") +@make_xp_test_case((RigidTransform, "as_matrix")) +def test_copy_flag(xp): + # Test that copy=True creates new memory + matrix = xp.eye(4) + tf = RigidTransform(matrix, normalize=False, copy=True) + matrix[0, 0] = 2 + assert tf.as_matrix()[0, 0] == 1 + + # Test that copy=False shares memory + matrix = xp.eye(4) + tf = RigidTransform(matrix, normalize=False, copy=False) + matrix[0, 0] = 2 + assert tf.as_matrix()[0, 0] == 2 + + +@make_xp_test_case(normalize_dual_quaternion) +@pytest.mark.parametrize("ndim", range(1, 4)) +def test_normalize_dual_quaternion(xp, ndim: int): + dtype = xpx.default_dtype(xp) + atol = 1e-12 if dtype == xp.float64 else 1e-6 + rng = np.random.default_rng(100) + shape = (ndim,) * (ndim - 1) + + dual_quat = normalize_dual_quaternion(xp.zeros((1, 8))) + xp_assert_close(xp_vector_norm(dual_quat[0, :4], axis=-1), xp.asarray(1.0)[()], + atol=1e-12) + xp_assert_close(xp.vecdot(dual_quat[0, :4], dual_quat[0, 4:])[()], + xp.asarray(0.0)[()], atol=1e-12) + + dual_quat = xp.asarray(rng.normal(size=shape + (8,)), dtype=dtype) + dual_quat = normalize_dual_quaternion(dual_quat) + expected = xp.ones(shape) if shape != () else xp.asarray(1.0)[()] + xp_assert_close(xp_vector_norm(dual_quat[..., :4], axis=-1), expected, atol=atol) + expected = xp.zeros(shape) if shape != () else xp.asarray(0.0)[()] + vecdot = xp.vecdot(dual_quat[..., :4], dual_quat[..., 4:]) + vecdot = vecdot[()] if vecdot.shape == () else vecdot + xp_assert_close(vecdot, expected, atol=atol) + + +@make_xp_test_case( + (RigidTransform, "from_matrix"), (RigidTransform, "from_rotation"), + (RigidTransform, "from_translation"), (RigidTransform, "from_components"), + (RigidTransform, "from_exp_coords"), (RigidTransform, "from_dual_quat") +) +def test_empty_transform_construction(xp): + tf = RigidTransform.from_matrix(xp.empty((0, 4, 4))) + assert len(tf) == 0 + assert not tf.single + + tf = RigidTransform.from_rotation(Rotation.from_quat(xp.zeros((0, 4)))) + assert len(tf) == 0 + assert not tf.single + + tf = RigidTransform.from_translation(xp.empty((0, 3))) + assert len(tf) == 0 + assert not tf.single + + empty_rot = Rotation.from_quat(xp.zeros((0, 4))) + tf = RigidTransform.from_components(xp.empty((0, 3)), empty_rot) + assert len(tf) == 0 + assert not tf.single + + tf = RigidTransform.from_exp_coords(xp.empty((0, 6))) + assert len(tf) == 0 + assert not tf.single + + tf = RigidTransform.from_dual_quat(xp.empty((0, 8))) + assert len(tf) == 0 + assert not tf.single + + tf = RigidTransform.identity(0) + assert len(tf) == 0 + assert not tf.single + + +@make_xp_test_case((RigidTransform, "from_matrix"), (RigidTransform, "as_components"), + (RigidTransform, "as_exp_coords"), (RigidTransform, "as_dual_quat")) +def test_empty_transform_representation(xp): + tf = RigidTransform.from_matrix(xp.empty((0, 4, 4))) + + assert len(tf.rotation) == 0 + assert tf.translation.shape == (0, 3) + + t, r = tf.as_components() + assert t.shape == (0, 3) + assert len(r) == 0 + + assert tf.as_matrix().shape == (0, 4, 4) + assert tf.as_exp_coords().shape == (0, 6) + assert tf.as_dual_quat().shape == (0, 8) + + +@make_xp_test_case((RigidTransform, "from_matrix"), (RigidTransform, "apply")) +def test_empty_transform_application(xp): + tf = RigidTransform.from_matrix(xp.empty((0, 4, 4))) + + assert tf.apply(xp.zeros((3,))).shape == (0, 3) + assert tf.apply(xp.empty((0, 3))).shape == (0, 3) + + with pytest.raises(ValueError, match="operands could not be broadcast together"): + tf.apply(xp.zeros((2, 3))) + + +@make_xp_test_case((RigidTransform, "from_matrix"), (RigidTransform, "__mul__")) +def test_empty_transform_composition(xp): + tf_empty = RigidTransform.from_matrix(xp.empty((0, 4, 4))) + tf_single = RigidTransform.from_matrix(xp.eye(4)) + tf_many = rigid_transform_to_xp(RigidTransform.identity(3), xp=xp) + + assert len(tf_empty * tf_empty) == 0 + assert len(tf_empty * tf_single) == 0 + assert len(tf_single * tf_empty) == 0 + + with pytest.raises(ValueError, match="Cannot broadcast"): + tf_many * tf_empty + + with pytest.raises(ValueError, match="Cannot broadcast"): + tf_empty * tf_many + + +@make_xp_test_case((RigidTransform, "from_matrix"), (RigidTransform, "concatenate")) +def test_empty_transform_concatenation(xp): + tf_empty = RigidTransform.from_matrix(xp.empty((0, 4, 4))) + tf_single = RigidTransform.from_matrix(xp.eye(4)) + tf_many = rigid_transform_to_xp(RigidTransform.identity(2), xp=xp) + + assert len(RigidTransform.concatenate([tf_empty, tf_empty])) == 0 + assert len(RigidTransform.concatenate([tf_empty, tf_single])) == 1 + assert len(RigidTransform.concatenate([tf_single, tf_empty])) == 1 + assert len(RigidTransform.concatenate([tf_empty, tf_many])) == 2 + assert len(RigidTransform.concatenate([tf_many, tf_empty])) == 2 + assert len(RigidTransform.concatenate([tf_many, tf_empty, tf_single])) == 3 + + +@make_xp_test_case((RigidTransform, "from_matrix"), (RigidTransform, "inv"), + (RigidTransform, "__pow__")) +def test_empty_transform_inv_and_pow(xp): + tf = RigidTransform.from_matrix(xp.empty((0, 4, 4))) + assert len(tf.inv()) == 0 + assert len(tf ** 0) == 0 + assert len(tf ** 1) == 0 + assert len(tf ** -1) == 0 + assert len(tf ** 0.5) == 0 + + +@make_xp_test_case((RigidTransform, "__getitem__")) +def test_empty_transform_indexing(xp): + tf_many = rigid_transform_to_xp(RigidTransform.identity(3), xp=xp) + tf_zero = tf_many[xp.asarray([], dtype=xp.int32)] + assert len(tf_zero) == 0 + + # Array API does not specify out-of-bounds indexing. Only check for numpy. + if is_numpy(xp): + assert len(tf_zero[:5]) == 0 # Slices can go out of bounds. + + with pytest.raises(IndexError): + tf_zero[0] + + with pytest.raises(IndexError): + tf_zero[xp.asarray([0, 2])] + + with pytest.raises(IndexError): + tf_zero[xp.asarray([False, True])] + + +@make_xp_test_case((RigidTransform, "from_matrix")) +@pytest.mark.skip_xp_backends("array_api_strict", + reason="array API doesn't support pickling") +def test_pickling(xp): + # Note: Array API makes no provision for arrays to be pickleable, so + # it's OK to skip this test for the backends that don't support it + mat = xp.eye(4) + mat = xpx.at(mat)[0, 3].set(2.0) + tf = RigidTransform.from_matrix(mat) + pkl = pickle.dumps(tf) + unpickled = pickle.loads(pkl) + xp_assert_close(tf.as_matrix(), unpickled.as_matrix(), atol=1e-15) + + +@make_xp_test_case( + (RigidTransform, "as_matrix"), + (RigidTransform, "__iter__"), + (RigidTransform, "identity"), +) +def test_rigid_transform_iter(xp): + r = rigid_transform_to_xp(RigidTransform.identity(3), xp) + for i, r_i in enumerate(r): + assert isinstance(r_i, RigidTransform) + xp_assert_equal(r_i.as_matrix(), r[i].as_matrix()) + if i > len(r): + raise RuntimeError("Iteration exceeded length of transforms") + + +@make_xp_test_case((RigidTransform, "from_translation")) +@pytest.mark.parametrize("dim", range(1, 5)) +def test_shape_property(xp, dim: int): + shape = (dim,) * (dim - 1) + tf = RigidTransform.from_translation(xp.zeros(shape + (3,))) + assert tf.shape == shape + + +def test_non_writeable(): + mat = np.eye(4) + mat.flags.writeable = False + RigidTransform.from_matrix(mat) # Regression test against gh-24378 diff --git a/v12x/pylib/scipy/spatial/transform/tests/test_rotation.py b/v12x/pylib/scipy/spatial/transform/tests/test_rotation.py new file mode 100644 index 0000000000000000000000000000000000000000..b2f96b93c20e41389ef64912285147e77a52cddd --- /dev/null +++ b/v12x/pylib/scipy/spatial/transform/tests/test_rotation.py @@ -0,0 +1,3223 @@ +import math + +import pytest + +import numpy as np +from numpy.testing import assert_equal +from scipy.spatial.transform import Rotation, Slerp +import scipy.spatial.transform._rotation_cy as cython_backend +import scipy.spatial.transform._rotation_xp as xp_backend +from scipy.stats import special_ortho_group +from itertools import permutations, product +from contextlib import contextmanager +import warnings +from scipy._lib._array_api import ( + xp_assert_equal, + array_namespace, + is_numpy, + is_lazy_array, + xp_vector_norm, + xp_assert_close, + eager_warns, + xp_default_dtype, + make_xp_test_case, + make_xp_pytest_marks, + xp_device_type, +) +import scipy._external.array_api_extra as xpx + +import pickle +import copy + + +# from_quat and as_quat are used in almost all tests, so we mark them module-wide +pytestmark = make_xp_pytest_marks((Rotation, "as_quat"), (Rotation, "from_quat")) + + +def basis_vec(axis): + if axis == 'x': + return [1, 0, 0] + elif axis == 'y': + return [0, 1, 0] + elif axis == 'z': + return [0, 0, 1] + + +def rotation_to_xp(r: Rotation, xp): + dtype = xpx.default_dtype(xp) + return Rotation.from_quat(xp.asarray(r.as_quat(), dtype=dtype)) + + +def test_init_non_array(): + Rotation((0, 0, 0, 1)) + Rotation([0, 0, 0, 1]) + Rotation([[[0, 0, 0, 1]]]) + + +def test_cython_backend_selection(): + r = Rotation.from_quat(np.array([0, 0, 0, 1])) + assert r._backend is cython_backend + r = Rotation.from_quat(np.array([[0, 0, 0, 1]])) + assert r._backend is cython_backend + r = Rotation.from_quat(np.array([[[0, 0, 0, 1]]])) + assert r._backend is xp_backend + + +def test_numpy_float32_inputs(): + Rotation.from_quat(np.array([1, 0, 0, 0], dtype=np.float32)) + + +def test_generic_quat_matrix(xp): + x = xp.asarray([[3.0, 4, 0, 0], [5, 12, 0, 0]]) + r = Rotation.from_quat(x) + expected_quat = x / xp.asarray([[5.0], [13.0]]) + xp_assert_close(r.as_quat(), expected_quat) + + +@pytest.mark.parametrize("ndim", list(range(1, 6))) +def test_from_single_nd_quaternion(xp, ndim: int): + x = xp.asarray([3.0, 4, 0, 0]) + x = xp.reshape(x, (1,) * (ndim - 1) + (4,)) + r = Rotation.from_quat(x) + expected_quat = x / 5 + xp_assert_close(r.as_quat(), expected_quat) + + +@make_xp_test_case((Rotation, "as_matrix")) +def test_from_quat_scalar_first(xp): + rng = np.random.RandomState(0) + + r = Rotation.from_quat(xp.asarray([1, 0, 0, 0]), scalar_first=True) + xp_assert_close(r.as_matrix(), xp.eye(3), rtol=1e-15, atol=1e-16) + + q = xp.tile(xp.asarray([1, 0, 0, 0]), (10, 1)) + r = Rotation.from_quat(q, scalar_first=True) + xp_assert_close( + r.as_matrix(), xp.tile(xp.eye(3), (10, 1, 1)), rtol=1e-15, atol=1e-16 + ) + + q = xp.asarray(rng.randn(100, 4)) + q /= xp_vector_norm(q, axis=1)[:, None] + for i in range(q.shape[0]): # Array API conforming loop + qi = q[i, ...] + r = Rotation.from_quat(qi, scalar_first=True) + xp_assert_close(xp.roll(r.as_quat(), 1), qi, rtol=1e-15) + + r = Rotation.from_quat(q, scalar_first=True) + xp_assert_close(xp.roll(r.as_quat(), 1, axis=1), q, rtol=1e-15) + + +def test_from_quat_array_like(): + rng = np.random.default_rng(123) + # Single rotation + r_expected = Rotation.random(rng=rng) + r = Rotation.from_quat(r_expected.as_quat().tolist()) + assert r_expected.approx_equal(r, atol=1e-12) + + # Multiple rotations + r_expected = Rotation.random(3, rng=rng) + r = Rotation.from_quat(r_expected.as_quat().tolist()) + assert np.all(r_expected.approx_equal(r, atol=1e-12)) + + # Tensor of rotations + q = rng.normal(size=(3, 4, 5, 4)) + q /= xp_vector_norm(q, axis=-1)[..., None] + r_expected = Rotation.from_quat(q) + r = Rotation.from_quat(q) + assert np.all(r_expected.approx_equal(r, atol=1e-12)) + + +def test_from_quat_int_dtype(xp): + r = Rotation.from_quat(xp.asarray([1, 0, 0, 0])) + assert r.as_quat().dtype == xp_default_dtype(xp) + + +def test_quat_canonical(xp): + # Case 0: w < 0 + q = xp.asarray([0.0, 0, 0, -1]) + xp_assert_close(Rotation.from_quat(q).as_quat(canonical=True), -q) + # Case 1: w == 0, x < 0 + q = xp.asarray([-1.0, 0, 0, 0]) + xp_assert_close(Rotation.from_quat(q).as_quat(canonical=True), -q) + # Case 2: w == 0, x == 0, y < 0 + q = xp.asarray([0.0, -1, 0, 0]) + xp_assert_close(Rotation.from_quat(q).as_quat(canonical=True), -q) + # Case 3: w == 0, x == 0, y == 0, z < 0 + q = xp.asarray([0.0, 0, -1, 0]) + xp_assert_close(Rotation.from_quat(q).as_quat(canonical=True), -q) + # Other cases: w > 0, y < 0 + q = xp.asarray([0.0, -0.1, 0, 0.9]) + q = q / xp_vector_norm(q) + xp_assert_close(Rotation.from_quat(q).as_quat(canonical=True), q) + # Other cases: w > 0, z < 0 + q = xp.asarray([0.0, 0.0, -0.1, 0.9]) + q = q / xp_vector_norm(q) + xp_assert_close(Rotation.from_quat(q).as_quat(canonical=True), q) + + +@make_xp_test_case((Rotation, "from_euler")) +def test_as_quat_scalar_first(xp): + rng = np.random.RandomState(0) + + r = Rotation.from_euler('xyz', xp.zeros(3)) + xp_assert_close(r.as_quat(scalar_first=True), xp.asarray([1.0, 0, 0, 0]), + rtol=1e-15, atol=1e-16) + + r = Rotation.from_euler('xyz', xp.zeros((10, 3))) + xp_assert_close(r.as_quat(scalar_first=True), + xp.tile(xp.asarray([1.0, 0, 0, 0]), (10, 1)), + rtol=1e-15, atol=1e-16) + + q = xp.asarray(rng.randn(100, 4)) + q /= xp_vector_norm(q, axis=1)[:, None] + for i in range(q.shape[0]): # Array API conforming loop + qi = q[i, ...] + r = Rotation.from_quat(qi) + xp_assert_close(r.as_quat(scalar_first=True), xp.roll(qi, 1), + rtol=1e-15) + + xp_assert_close(r.as_quat(canonical=True, scalar_first=True), + xp.roll(r.as_quat(canonical=True), 1), + rtol=1e-15) + + r = Rotation.from_quat(q) + xp_assert_close(r.as_quat(scalar_first=True), xp.roll(q, 1, axis=1), + rtol=1e-15) + + xp_assert_close(r.as_quat(canonical=True, scalar_first=True), + xp.roll(r.as_quat(canonical=True), 1, axis=1), rtol=1e-15) + + +def test_from_square_quat_matrix(xp): + # Ensure proper norm array broadcasting + x = xp.asarray([ + [3.0, 0, 0, 4], + [5, 0, 12, 0], + [0, 0, 0, 1], + [-1, -1, -1, 1], + [0, 0, 0, -1], # Check double cover + [-1, -1, -1, -1] # Check double cover + ]) + r = Rotation.from_quat(x) + expected_quat = x / xp.asarray([[5.0], [13], [1], [2], [1], [2]]) + xp_assert_close(r.as_quat(), expected_quat) + + +def test_quat_double_to_canonical_single_cover(xp): + x = xp.asarray([ + [-1.0, 0, 0, 0], + [0, -1, 0, 0], + [0, 0, -1, 0], + [0, 0, 0, -1], + [-1, -1, -1, -1] + ]) + r = Rotation.from_quat(x) + expected_quat = xp.abs(x) / xp_vector_norm(x, axis=1)[:, None] + xp_assert_close(r.as_quat(canonical=True), expected_quat) + + +@make_xp_test_case((Rotation, "inv"), (Rotation, "__mul__")) +def test_quat_double_cover(xp): + # See the Rotation.from_quat() docstring for scope of the quaternion + # double cover property. + # Check from_quat and as_quat(canonical=False) + q = xp.asarray([0.0, 0, 0, -1]) + r = Rotation.from_quat(q) + xp_assert_equal(q, r.as_quat(canonical=False)) + # Check composition and inverse + q = xp.asarray([1.0, 0, 0, 1])/math.sqrt(2) # 90 deg rotation about x + r = Rotation.from_quat(q) + r3 = r*r*r + xp_assert_close(r.as_quat(canonical=False)*math.sqrt(2), + xp.asarray([1.0, 0, 0, 1])) + xp_assert_close(r.inv().as_quat(canonical=False)*math.sqrt(2), + xp.asarray([-1.0, 0, 0, 1])) + xp_assert_close(r3.as_quat(canonical=False)*math.sqrt(2), + xp.asarray([1.0, 0, 0, -1])) + xp_assert_close(r3.inv().as_quat(canonical=False)*math.sqrt(2), + xp.asarray([-1.0, 0, 0, -1])) + + # More sanity checks + xp_assert_close((r*r.inv()).as_quat(canonical=False), + xp.asarray([0.0, 0, 0, 1]), atol=2e-16) + xp_assert_close((r3*r3.inv()).as_quat(canonical=False), + xp.asarray([0.0, 0, 0, 1]), atol=2e-16) + xp_assert_close((r*r3).as_quat(canonical=False), + xp.asarray([0.0, 0, 0, -1]), atol=2e-16) + xp_assert_close((r.inv() * r3.inv()).as_quat(canonical=False), + xp.asarray([0.0, 0, 0, -1]), atol=2e-16) + + +@pytest.mark.parametrize("ndim", list(range(1, 6))) +def test_from_quat_wrong_shape(xp, ndim: int): + quat = xp.zeros((*((1,) * ndim), 5)) + with pytest.raises(ValueError, match="Expected `quat` to have shape"): + Rotation.from_quat(quat) + + +def test_zero_norms_from_quat(xp): + x = xp.asarray([ + [3, 4, 0, 0], + [0, 0, 0, 0], + [5, 0, 12, 0] + ]) + if is_lazy_array(x): + assert xp.all(xp.isnan(Rotation.from_quat(x).as_quat()[1, ...])) + else: + with pytest.raises(ValueError): + Rotation.from_quat(x) + + +@make_xp_test_case((Rotation, "as_matrix")) +@pytest.mark.parametrize("ndim", list(range(1, 6))) +def test_as_matrix_single_nd_quaternion(xp, ndim: int): + quat = xp.asarray([0, 0, 1, 1]) + quat = xp.reshape(quat, (1,) * (ndim - 1) + (4,)) + mat = Rotation.from_quat(quat).as_matrix() + expected_mat = xp.asarray([ + [0.0, -1, 0], + [1, 0, 0], + [0, 0, 1] + ]) + expected_mat = xp.reshape(expected_mat, (1,) * (ndim - 1) + (3, 3)) + xp_assert_close(mat, expected_mat, atol=1e-16) + + +@make_xp_test_case((Rotation, "as_matrix")) +def test_as_matrix_from_square_input(xp): + quats = xp.asarray([ + [0, 0, 1, 1], + [0, 1, 0, 1], + [0, 0, 0, 1], + [0, 0, 0, -1] + ]) + mat = Rotation.from_quat(quats).as_matrix() + assert_equal(mat.shape, (4, 3, 3)) + + expected0 = xp.asarray([ + [0.0, -1, 0], + [1, 0, 0], + [0, 0, 1] + ]) + xp_assert_close(mat[0, ...], expected0, atol=1e-16) + + expected1 = xp.asarray([ + [0.0, 0, 1], + [0, 1, 0], + [-1, 0, 0] + ]) + xp_assert_close(mat[1, ...], expected1, atol=1e-16) + xp_assert_close(mat[2, ...], xp.eye(3)) + xp_assert_close(mat[3, ...], xp.eye(3)) + + +@make_xp_test_case((Rotation, "as_matrix")) +def test_as_matrix_from_generic_input(xp): + quats = xp.asarray([ + [0, 0, 1, 1], + [0, 1, 0, 1], + [1, 2, 3, 4] + ]) + mat = Rotation.from_quat(quats).as_matrix() + assert_equal(mat.shape, (3, 3, 3)) + + expected0 = xp.asarray([ + [0.0, -1, 0], + [1, 0, 0], + [0, 0, 1] + ]) + xp_assert_close(mat[0, ...], expected0, atol=1e-16) + + expected1 = xp.asarray([ + [0.0, 0, 1], + [0, 1, 0], + [-1, 0, 0] + ]) + xp_assert_close(mat[1, ...], expected1, atol=1e-16) + + expected2 = xp.asarray([ + [0.4, -2, 2.2], + [2.8, 1, 0.4], + [-1, 2, 2] + ]) / 3 + xp_assert_close(mat[2, ...], expected2) + + +@make_xp_test_case((Rotation, "from_matrix")) +@pytest.mark.parametrize("ndim", list(range(1, 6))) +def test_from_single_nd_matrix(xp, ndim: int): + mat = xp.asarray([ + [0, 0, 1], + [1, 0, 0], + [0, 1, 0] + ]) + mat = xp.reshape(mat, (1,) * (ndim - 1) + (3, 3)) + expected_quat = xp.asarray([0.5, 0.5, 0.5, 0.5]) + expected_quat = xp.reshape(expected_quat, (1,) * (ndim - 1) + (4,)) + xp_assert_close(Rotation.from_matrix(mat).as_quat(), expected_quat) + + +@make_xp_test_case((Rotation, "from_matrix")) +def test_from_matrix_calculation(xp): + atol = 1e-8 + expected_quat = xp.asarray([1.0, 1, 6, 1]) / math.sqrt(39) + mat = xp.asarray([ + [-0.8974359, -0.2564103, 0.3589744], + [0.3589744, -0.8974359, 0.2564103], + [0.2564103, 0.3589744, 0.8974359] + ]) + xp_assert_close(Rotation.from_matrix(mat).as_quat(), expected_quat, atol=atol) + xp_assert_close(Rotation.from_matrix(xp.reshape(mat, (1, 3, 3))).as_quat(), + xp.reshape(expected_quat, (1, 4)), + atol=atol) + + +@make_xp_test_case((Rotation, "from_matrix"), (Rotation, "as_matrix")) +def test_matrix_calculation_pipeline(xp): + mat = xp.asarray(special_ortho_group.rvs(3, size=10, random_state=0)) + xp_assert_close(Rotation.from_matrix(mat).as_matrix(), mat) + + +@make_xp_test_case((Rotation, "from_matrix"), (Rotation, "as_matrix")) +def test_from_matrix_ortho_output(xp): + dtype = xpx.default_dtype(xp) + atol = 1e-12 if dtype == xp.float64 else 1e-6 + rnd = np.random.RandomState(0) + mat = xp.asarray(rnd.random_sample((100, 3, 3)), dtype=dtype) + dets = xp.linalg.det(mat) + for i in range(dets.shape[0]): + # Make sure we have a right-handed rotation matrix + if dets[i] < 0: + mat = xpx.at(mat)[i, ...].set(-mat[i, ...]) + ortho_mat = Rotation.from_matrix(mat).as_matrix() + + mult_result = xp.matmul(ortho_mat, xp.matrix_transpose(ortho_mat)) + + eye3d = xp.zeros((100, 3, 3)) + xp.eye(3) + xp_assert_close(mult_result, eye3d, atol=atol) + + +@make_xp_test_case((Rotation, "from_matrix"), (Rotation, "as_matrix")) +def test_from_matrix_normalize(xp): + mat = xp.asarray([ + [1, 1, 0], + [0, 1, 0], + [0, 0, 1]]) + expected = xp.asarray([[ 0.894427, 0.447214, 0.0], + [-0.447214, 0.894427, 0.0], + [ 0.0, 0.0, 1.0]]) + xp_assert_close(Rotation.from_matrix(mat).as_matrix(), expected, atol=1e-6) + + mat = xp.asarray([ + [0, -0.5, 0 ], + [0.5, 0 , 0 ], + [0, 0 , 0.5]]) + expected = xp.asarray([[0.0, -1, 0], + [ 1, 0, 0], + [ 0, 0, 1]]) + xp_assert_close(Rotation.from_matrix(mat).as_matrix(), expected, atol=1e-6) + + # Test a mix of normalized and non-normalized matrices + mat = xp.stack([mat, xp.eye(3)]) + expected = xp.stack([expected, xp.eye(3)]) + xp_assert_close(Rotation.from_matrix(mat).as_matrix(), expected, atol=1e-6) + + +@make_xp_test_case((Rotation, "from_matrix"), (Rotation, "as_matrix")) +def test_from_matrix_assume_valid(xp): + rng = np.random.default_rng(0) + dtype = xpx.default_dtype(xp) + atol = 1e-12 if dtype == xp.float64 else 1e-6 + # Test that normal matrices remain unchanged + rot = rotation_to_xp(Rotation.from_quat(rng.normal(size=(10, 4))), xp) + rot_no_norm = Rotation.from_matrix(rot.as_matrix(), assume_valid=True) + assert xp.all(rot.approx_equal(rot_no_norm, atol=atol)) + # We make no guarantees about matrices that are not orthogonal or do not + # have unit norm + + +@make_xp_test_case((Rotation, "from_matrix"), (Rotation, "as_matrix")) +def test_from_matrix_non_positive_determinant(xp): + mat = xp.eye(3) + mat = xpx.at(mat)[0, 0].set(0) + if is_lazy_array(mat): + assert xp.all(xp.isnan(Rotation.from_matrix(mat).as_matrix())) + else: + with pytest.raises(ValueError, match="Non-positive determinant"): + Rotation.from_matrix(mat) + + mat = xpx.at(mat)[0, 0].set(-1) + if is_lazy_array(mat): + assert xp.all(xp.isnan(Rotation.from_matrix(mat).as_matrix())) + else: + with pytest.raises(ValueError, match="Non-positive determinant"): + Rotation.from_matrix(mat) + + +def test_from_matrix_array_like(): + rng = np.random.default_rng(123) + # Single rotation + r_expected = Rotation.random(rng=rng) + r = Rotation.from_matrix(r_expected.as_matrix().tolist()) + assert r_expected.approx_equal(r, atol=1e-12) + + # Multiple rotations + r_expected = Rotation.random(3, rng=rng) + r = Rotation.from_matrix(r_expected.as_matrix().tolist()) + assert np.all(r_expected.approx_equal(r, atol=1e-12)) + + +@make_xp_test_case((Rotation, "from_matrix")) +def test_from_matrix_int_dtype(xp): + mat = xp.asarray([[1, 0, 0], [0, 1, 0], [0, 0, 1]]) + r = Rotation.from_matrix(mat) + assert r.as_quat().dtype == xp_default_dtype(xp) + + +@make_xp_test_case((Rotation, "from_rotvec")) +@pytest.mark.parametrize("ndim", list(range(1, 6))) +def test_from_nd_single_rotvec(xp, ndim: int): + atol = 1e-7 + rotvec = xp.asarray([1, 0, 0]) + rotvec = xp.reshape(rotvec, (1,) * (ndim - 1) + (3,)) + expected_quat = xp.asarray([0.4794255, 0, 0, 0.8775826]) + expected_quat = xp.reshape(expected_quat, (1,) * (ndim - 1) + (4,)) + result = Rotation.from_rotvec(rotvec) + xp_assert_close(result.as_quat(), expected_quat, atol=atol) + + +@make_xp_test_case((Rotation, "from_rotvec")) +def test_from_generic_rotvec(xp): + atol = 1e-7 + rotvec = xp.asarray([ + [1, 2, 2], + [1, -1, 0.5], + [0, 0, 0]]) + expected_quat = xp.asarray([ + [0.3324983, 0.6649967, 0.6649967, 0.0707372], + [0.4544258, -0.4544258, 0.2272129, 0.7316889], + [0, 0, 0, 1] + ]) + xp_assert_close(Rotation.from_rotvec(rotvec).as_quat(), expected_quat, atol=atol) + + +@make_xp_test_case((Rotation, "from_rotvec")) +def test_from_rotvec_small_angle(xp): + rotvec = xp.asarray([ + [5e-4 / math.sqrt(3), -5e-4 / math.sqrt(3), 5e-4 / math.sqrt(3)], + [0.2, 0.3, 0.4], + [0, 0, 0] + ]) + + quat = Rotation.from_rotvec(rotvec).as_quat() + # cos(theta/2) ~~ 1 for small theta + xp_assert_close(quat[0, 3], xp.asarray(1.0)[()]) + # sin(theta/2) / theta ~~ 0.5 for small theta + xp_assert_close(quat[0, :3], rotvec[0, ...] * 0.5) + + xp_assert_close(quat[1, 3], xp.asarray(0.9639685)[()]) + xp_assert_close(quat[1, :3], + xp.asarray([ + 0.09879603932153465, + 0.14819405898230198, + 0.19759207864306931])) + + xp_assert_equal(quat[2, ...], xp.asarray([0.0, 0, 0, 1])) + + +@make_xp_test_case((Rotation, "from_rotvec"), (Rotation, "as_rotvec")) +def test_from_rotvec_array_like(): + rng = np.random.default_rng(123) + # Single rotation + r_expected = Rotation.random(rng=rng) + r = Rotation.from_rotvec(r_expected.as_rotvec().tolist()) + assert r_expected.approx_equal(r, atol=1e-12) + + # Multiple rotations + r_expected = Rotation.random(3, rng=rng) + r = Rotation.from_rotvec(r_expected.as_rotvec().tolist()) + assert np.all(r_expected.approx_equal(r, atol=1e-12)) + + +@make_xp_test_case((Rotation, "from_rotvec")) +def test_from_rotvec_int_dtype(xp): + rotvec = xp.asarray([1, 0, 0]) + r = Rotation.from_rotvec(rotvec) + assert r.as_quat().dtype == xp_default_dtype(xp) + + +@make_xp_test_case((Rotation, "from_rotvec")) +def test_degrees_from_rotvec(xp): + rotvec1 = xp.asarray([1 / 3 ** (1/3)] * 3) + rot1 = Rotation.from_rotvec(rotvec1, degrees=True) + quat1 = rot1.as_quat() + + # deg2rad is not implemented in Array API -> / 180 * xp.pi + rotvec2 = xp.asarray(rotvec1 / 180 * xp.pi) + rot2 = Rotation.from_rotvec(rotvec2) + quat2 = rot2.as_quat() + + xp_assert_close(quat1, quat2) + + +@make_xp_test_case((Rotation, "from_rotvec")) +def test_malformed_1d_from_rotvec(xp): + with pytest.raises(ValueError, match='Expected `rot_vec` to have shape'): + Rotation.from_rotvec(xp.asarray([1, 2])) + + +@make_xp_test_case((Rotation, "from_rotvec")) +@pytest.mark.parametrize("ndim", list(range(1, 6))) +def test_malformed_nd_from_rotvec(xp, ndim: int): + shape = (1,) * (ndim - 1) + (2,) + with pytest.raises(ValueError, match='Expected `rot_vec` to have shape'): + Rotation.from_rotvec(xp.ones(shape)) + + +@make_xp_test_case((Rotation, "as_rotvec")) +def test_as_generic_rotvec(xp): + dtype = xpx.default_dtype(xp) + atol = 1e-15 if dtype == xp.float64 else 1e-7 + quat = xp.asarray([ + [1, 2, -1, 0.5], + [1, -1, 1, 0.0003], + [0, 0, 0, 1] + ]) + quat /= xp_vector_norm(quat, axis=-1, keepdims=True) + + rotvec = Rotation.from_quat(quat).as_rotvec() + angle = xp_vector_norm(rotvec, axis=-1) + + xp_assert_close(quat[:, 3], xp.cos(angle / 2)) + xp_assert_close(xp.linalg.cross(rotvec, quat[:, :3]), xp.zeros((3, 3)), atol=atol) + + +@make_xp_test_case((Rotation, "as_rotvec")) +@pytest.mark.parametrize("ndim", list(range(1, 4))) +def test_as_rotvec_single_nd_input(xp, ndim: int): + quat = xp.asarray([1, 2, -3, 2]) + quat = xp.reshape(quat, (1,) * (ndim - 1) + (4,)) + expected_rotvec = xp.asarray([0.5772381, 1.1544763, -1.7317144]) + expected_rotvec = xp.reshape(expected_rotvec, (1,) * (ndim - 1) + (3,)) + actual_rotvec = Rotation.from_quat(quat).as_rotvec() + + assert_equal(actual_rotvec.shape, expected_rotvec.shape) + xp_assert_close(actual_rotvec, expected_rotvec) + + +@make_xp_test_case((Rotation, "from_matrix"), (Rotation, "as_rotvec")) +def test_as_rotvec_degrees(xp): + # x->y, y->z, z->x + mat = xp.asarray([[0, 0, 1], [1, 0, 0], [0, 1, 0]]) + rot = Rotation.from_matrix(mat) + rotvec = rot.as_rotvec(degrees=True) + angle = xp_vector_norm(rotvec, axis=-1) + xp_assert_close(angle, xp.asarray(120.0)[()]) + xp_assert_close(rotvec[0], rotvec[1]) + xp_assert_close(rotvec[1], rotvec[2]) + + +@make_xp_test_case((Rotation, "from_rotvec"), (Rotation, "as_rotvec")) +def test_rotvec_calc_pipeline(xp): + # Include small angles + rotvec = xp.asarray([ + [0, 0, 0], + [1, -1, 2], + [-3e-4, 3.5e-4, 7.5e-5] + ]) + xp_assert_close(Rotation.from_rotvec(rotvec).as_rotvec(), rotvec) + xp_assert_close(Rotation.from_rotvec(rotvec, degrees=True).as_rotvec(degrees=True), + rotvec) + + +@make_xp_test_case((Rotation, "from_mrp")) +@pytest.mark.parametrize("ndim", range(1, 4)) +def test_from_mrp_single_nd_input(xp, ndim: int): + mrp = xp.asarray([0, 0, 1.0]) + mrp = xp.reshape(mrp, (1,) * (ndim - 1) + (3,)) + expected_quat = xp.asarray([0.0, 0, 1, 0]) + expected_quat = xp.reshape(expected_quat, (1,) * (ndim - 1) + (4,)) + result = Rotation.from_mrp(mrp) + xp_assert_close(result.as_quat(), expected_quat, atol=1e-12) + # Regression test for gh-24555 + assert isinstance(result._quat, type(array_namespace(mrp).empty(0))) + + +def test_from_mrp_array_like(): + rng = np.random.default_rng(123) + # Single rotation + r_expected = Rotation.random(rng=rng) + r = Rotation.from_mrp(r_expected.as_mrp().tolist()) + assert r_expected.approx_equal(r, atol=1e-12) + + # Multiple rotations + r_expected = Rotation.random(3, rng=rng) + r = Rotation.from_mrp(r_expected.as_mrp().tolist()) + assert np.all(r_expected.approx_equal(r, atol=1e-12)) + + +@make_xp_test_case((Rotation, "from_mrp")) +def test_from_mrp_int_dtype(xp): + mrp = xp.asarray([0, 0, 1]) + r = Rotation.from_mrp(mrp) + assert r.as_quat().dtype == xp_default_dtype(xp) + + +@make_xp_test_case((Rotation, "from_mrp")) +def test_from_generic_mrp(xp): + mrp = xp.asarray([ + [1, 2, 2], + [1, -1, 0.5], + [0, 0, 0]]) + expected_quat = xp.asarray([ + [0.2, 0.4, 0.4, -0.8], + [0.61538462, -0.61538462, 0.30769231, -0.38461538], + [0, 0, 0, 1]]) + xp_assert_close(Rotation.from_mrp(mrp).as_quat(), expected_quat) + + +@make_xp_test_case((Rotation, "from_mrp")) +@pytest.mark.parametrize("ndim", range(1, 4)) +def test_malformed_nd_from_mrp(xp, ndim: int): + shape = (1,) * (ndim - 1) + (2,) + with pytest.raises(ValueError, match='Expected `mrp` to have shape'): + Rotation.from_mrp(xp.ones(shape)) + + +@make_xp_test_case((Rotation, "as_mrp")) +def test_as_generic_mrp(xp): + quat = xp.asarray([ + [1, 2, -1, 0.5], + [1, -1, 1, 0.0003], + [0, 0, 0, 1]]) + quat /= xp_vector_norm(quat, axis=1)[:, None] + + expected_mrp = xp.asarray([ + [0.33333333, 0.66666667, -0.33333333], + [0.57725028, -0.57725028, 0.57725028], + [0, 0, 0]]) + xp_assert_close(Rotation.from_quat(quat).as_mrp(), expected_mrp) + + +@make_xp_test_case((Rotation, "from_euler"), (Rotation, "as_mrp")) +def test_past_180_degree_rotation(xp): + # ensure that a > 180 degree rotation is returned as a <180 rotation in MRPs + # in this case 270 should be returned as -90 + expected_mrp = xp.asarray([-math.tan(xp.pi / 2 / 4), 0.0, 0]) + xp_assert_close( + Rotation.from_euler('xyz', xp.asarray([270, 0, 0]), degrees=True).as_mrp(), + expected_mrp, + ) + + +@make_xp_test_case((Rotation, "as_mrp")) +@pytest.mark.parametrize("ndim", list(range(1, 4))) +def test_as_mrp_single_nd_input(xp, ndim: int): + quat = xp.asarray([1, 2, -3, 2]) + quat = xp.reshape(quat, (1,) * (ndim - 1) + (4,)) + expected_mrp = xp.asarray([0.16018862, 0.32037724, -0.48056586]) + expected_mrp = xp.reshape(expected_mrp, (1,) * (ndim - 1) + (3,)) + actual_mrp = Rotation.from_quat(quat).as_mrp() + + assert_equal(actual_mrp.shape, expected_mrp.shape) + xp_assert_close(actual_mrp, expected_mrp) + + +@make_xp_test_case((Rotation, "from_mrp"), (Rotation, "as_mrp")) +def test_mrp_calc_pipeline(xp): + actual_mrp = xp.asarray([ + [0, 0, 0], + [1, -1, 2], + [0.41421356, 0, 0], + [0.1, 0.2, 0.1]]) + expected_mrp = xp.asarray([ + [0, 0, 0], + [-0.16666667, 0.16666667, -0.33333333], + [0.41421356, 0, 0], + [0.1, 0.2, 0.1]]) + xp_assert_close(Rotation.from_mrp(actual_mrp).as_mrp(), expected_mrp) + + +@make_xp_test_case((Rotation, "from_euler")) +def test_from_euler_single_rotation(xp): + quat = Rotation.from_euler("z", xp.asarray(90), degrees=True).as_quat() + expected_quat = xp.asarray([0.0, 0, 1, 1]) / math.sqrt(2) + xp_assert_close(quat, expected_quat) + + +@make_xp_test_case((Rotation, "from_euler")) +def test_from_euler_input_validation(xp): + # Single sequence with multiple angles + with pytest.raises(ValueError, match="Expected last dimension of `angles` to"): + Rotation.from_euler("X", xp.asarray([0, 90])) + # Multiple sequences with single angle + with pytest.raises(ValueError, match="Expected last dimension of `angles` to"): + Rotation.from_euler("XYZ", xp.asarray([90])) + + +@make_xp_test_case((Rotation, "from_euler")) +@pytest.mark.parametrize("ndim", list(range(1, 4))) +def test_from_euler_nd_rotation(xp, ndim: int): + angles = xp.reshape(xp.asarray([0, 0, 90]), (1,) * (ndim - 1) + (3,)) + quat = Rotation.from_euler("xyz", angles, degrees=True).as_quat() + expected_quat = xp.asarray([0.0, 0, 1, 1]) / math.sqrt(2) + expected_quat = xp.reshape(expected_quat, (1,) * (ndim - 1) + (4,)) + xp_assert_close(quat, expected_quat) + + +@make_xp_test_case((Rotation, "from_euler"), (Rotation, "as_matrix")) +def test_single_intrinsic_extrinsic_rotation(xp): + extrinsic = Rotation.from_euler('z', xp.asarray(90), degrees=True).as_matrix() + intrinsic = Rotation.from_euler('Z', xp.asarray(90), degrees=True).as_matrix() + xp_assert_close(extrinsic, intrinsic) + + +@make_xp_test_case((Rotation, "from_euler")) +def test_from_euler_rotation_order(xp): + # Intrinsic rotation is same as extrinsic with order reversed + rnd = np.random.RandomState(0) + a = xp.asarray(rnd.randint(low=0, high=180, size=(6, 3))) + b = xp.flip(a, axis=-1) + x = Rotation.from_euler('xyz', a, degrees=True).as_quat() + y = Rotation.from_euler('ZYX', b, degrees=True).as_quat() + xp_assert_close(x, y) + + +@make_xp_test_case((Rotation, "from_euler"), (Rotation, "as_matrix")) +def test_from_euler_elementary_extrinsic_rotation(xp): + atol = 1e-12 + # Simple test to check if extrinsic rotations are implemented correctly + mat = Rotation.from_euler('zx', xp.asarray([90, 90]), degrees=True).as_matrix() + expected_mat = xp.asarray([ + [0.0, -1, 0], + [0, 0, -1], + [1, 0, 0] + ]) + xp_assert_close(mat, expected_mat, atol=atol) + + +@make_xp_test_case((Rotation, "from_euler"), (Rotation, "as_matrix")) +def test_from_euler_intrinsic_rotation_312(xp): + atol = 1e-7 + angles = xp.asarray([ + [30, 60, 45], + [30, 60, 30], + [45, 30, 60] + ]) + mat = Rotation.from_euler('ZXY', angles, degrees=True).as_matrix() + + xp_assert_close(mat[0, ...], xp.asarray([ + [0.3061862, -0.2500000, 0.9185587], + [0.8838835, 0.4330127, -0.1767767], + [-0.3535534, 0.8660254, 0.3535534] + ]), atol=atol) + + xp_assert_close(mat[1, ...], xp.asarray([ + [0.5334936, -0.2500000, 0.8080127], + [0.8080127, 0.4330127, -0.3995191], + [-0.2500000, 0.8660254, 0.4330127] + ]), atol=atol) + + xp_assert_close(mat[2, ...], xp.asarray([ + [0.0473672, -0.6123725, 0.7891491], + [0.6597396, 0.6123725, 0.4355958], + [-0.7500000, 0.5000000, 0.4330127] + ]), atol=atol) + + +@make_xp_test_case((Rotation, "from_euler"), (Rotation, "as_matrix")) +def test_from_euler_intrinsic_rotation_313(xp): + angles = xp.asarray([ + [30, 60, 45], + [30, 60, 30], + [45, 30, 60] + ]) + mat = Rotation.from_euler('ZXZ', angles, degrees=True).as_matrix() + + xp_assert_close(mat[0, ...], xp.asarray([ + [0.43559574, -0.78914913, 0.4330127], + [0.65973961, -0.04736717, -0.750000], + [0.61237244, 0.61237244, 0.500000] + ])) + + xp_assert_close(mat[1, ...], xp.asarray([ + [0.6250000, -0.64951905, 0.4330127], + [0.64951905, 0.1250000, -0.750000], + [0.4330127, 0.750000, 0.500000] + ])) + + xp_assert_close(mat[2, ...], xp.asarray([ + [-0.1767767, -0.91855865, 0.35355339], + [0.88388348, -0.30618622, -0.35355339], + [0.4330127, 0.25000000, 0.8660254] + ])) + + +@make_xp_test_case((Rotation, "from_euler"), (Rotation, "as_matrix")) +def test_from_euler_extrinsic_rotation_312(xp): + angles = xp.asarray([ + [30, 60, 45], + [30, 60, 30], + [45, 30, 60] + ]) + mat = Rotation.from_euler('zxy', angles, degrees=True).as_matrix() + + xp_assert_close(mat[0, ...], xp.asarray([ + [0.91855865, 0.1767767, 0.35355339], + [0.25000000, 0.4330127, -0.8660254], + [-0.30618622, 0.88388348, 0.35355339] + ])) + + xp_assert_close(mat[1, ...], xp.asarray([ + [0.96650635, -0.0580127, 0.2500000], + [0.25000000, 0.4330127, -0.8660254], + [-0.0580127, 0.89951905, 0.4330127] + ])) + + xp_assert_close(mat[2, ...], xp.asarray([ + [0.65973961, -0.04736717, 0.7500000], + [0.61237244, 0.61237244, -0.5000000], + [-0.43559574, 0.78914913, 0.4330127] + ])) + + +@make_xp_test_case((Rotation, "from_euler"), (Rotation, "as_matrix")) +def test_from_euler_extrinsic_rotation_313(xp): + angles = xp.asarray([ + [30, 60, 45], + [30, 60, 30], + [45, 30, 60] + ]) + mat = Rotation.from_euler('zxz', angles, degrees=True).as_matrix() + + xp_assert_close(mat[0, ...], xp.asarray([ + [0.43559574, -0.65973961, 0.61237244], + [0.78914913, -0.04736717, -0.61237244], + [0.4330127, 0.75000000, 0.500000] + ])) + + xp_assert_close(mat[1, ...], xp.asarray([ + [0.62500000, -0.64951905, 0.4330127], + [0.64951905, 0.12500000, -0.750000], + [0.4330127, 0.75000000, 0.500000] + ])) + + xp_assert_close(mat[2, ...], xp.asarray([ + [-0.1767767, -0.88388348, 0.4330127], + [0.91855865, -0.30618622, -0.250000], + [0.35355339, 0.35355339, 0.8660254] + ])) + + +def test_from_euler_array_like(): + rng = np.random.default_rng(123) + order = "xyz" + # Single rotation + r_expected = Rotation.random(rng=rng) + r = Rotation.from_euler(order, r_expected.as_euler(order).tolist()) + assert r_expected.approx_equal(r, atol=1e-12) + + # Multiple rotations + r_expected = Rotation.random(3, rng=rng) + r = Rotation.from_euler(order, r_expected.as_euler(order).tolist()) + assert np.all(r_expected.approx_equal(r, atol=1e-12)) + + +def test_from_euler_scalar(): + rng = np.random.default_rng(123) + deg = rng.uniform(low=-180, high=180) + r_expected = Rotation.from_euler("x", deg, degrees=True) + r = Rotation.from_euler("x", float(deg), degrees=True) + assert r_expected.approx_equal(r, atol=1e-12) + + +@make_xp_test_case((Rotation, "from_euler"), (Rotation, "as_euler")) +@pytest.mark.parametrize("seq_tuple", list(permutations("xyz"))) +@pytest.mark.parametrize("intrinsic", (False, True)) +def test_as_euler_asymmetric_axes(xp, seq_tuple, intrinsic): + # helper function for mean error tests + def test_stats(error, mean_max, rms_max): + mean = xp.mean(error, axis=0) + std = xp.std(error, axis=0) + rms = xp.hypot(mean, std) + assert xp.all(xp.abs(mean) < mean_max) + assert xp.all(rms < rms_max) + + rnd = np.random.RandomState(0) + n = 1000 + angles = np.empty((n, 3)) + angles[:, 0] = rnd.uniform(low=-np.pi, high=np.pi, size=(n,)) + angles[:, 1] = rnd.uniform(low=-np.pi / 2, high=np.pi / 2, size=(n,)) + angles[:, 2] = rnd.uniform(low=-np.pi, high=np.pi, size=(n,)) + angles = xp.asarray(angles) + + seq = "".join(seq_tuple) + if intrinsic: + # Extrinsic rotation (w.r.t. global world) at lower case + # intrinsic (WRT the object itself) lower case. + seq = seq.upper() + rotation = Rotation.from_euler(seq, angles) + angles_quat = rotation.as_euler(seq) + xp_assert_close(angles, angles_quat, atol=0, rtol=1e-12) + test_stats(angles_quat - angles, 1e-15, 1e-14) + + +@make_xp_test_case((Rotation, "from_euler"), (Rotation, "as_euler")) +@pytest.mark.parametrize("seq_tuple", list(permutations("xyz"))) +@pytest.mark.parametrize("intrinsic", (False, True)) +def test_as_euler_symmetric_axes(xp, seq_tuple, intrinsic): + # helper function for mean error tests + def test_stats(error, mean_max, rms_max): + mean = xp.mean(error, axis=0) + std = xp.std(error, axis=0) + rms = xp.hypot(mean, std) + assert xp.all(xp.abs(mean) < mean_max) + assert xp.all(rms < rms_max) + + rnd = np.random.RandomState(0) + n = 1000 + angles = np.empty((n, 3)) + angles[:, 0] = rnd.uniform(low=-np.pi, high=np.pi, size=(n,)) + angles[:, 1] = rnd.uniform(low=0, high=np.pi, size=(n,)) + angles[:, 2] = rnd.uniform(low=-np.pi, high=np.pi, size=(n,)) + angles = xp.asarray(angles) + + # Rotation of the form A/B/A are rotation around symmetric axes + seq = "".join([seq_tuple[0], seq_tuple[1], seq_tuple[0]]) + if intrinsic: + seq = seq.upper() + rotation = Rotation.from_euler(seq, angles) + angles_quat = rotation.as_euler(seq) + xp_assert_close(angles, angles_quat, atol=0, rtol=1.1e-13) + test_stats(angles_quat - angles, 1e-16, 1e-14) + + +@contextmanager +def maybe_warn_gimbal_lock(should_warn, xp): + if should_warn: + # We can only warn on non-lazy backends because we'd need to condition on + # traced booleans + with eager_warns(UserWarning, match="Gimbal lock", xp=xp): + yield + + else: + with warnings.catch_warnings(): + warnings.simplefilter("error") + yield + + +@make_xp_test_case( + (Rotation, "from_euler"), (Rotation, "as_matrix"), (Rotation, "as_euler") +) +@pytest.mark.parametrize("seq_tuple", list(permutations("xyz"))) +@pytest.mark.parametrize("intrinsic", (False, True)) +@pytest.mark.parametrize("suppress_warnings", (False, True)) +def test_as_euler_degenerate_asymmetric_axes( + xp, seq_tuple, intrinsic, suppress_warnings +): + dtype = xpx.default_dtype(xp) + atol = 1e-12 if dtype == xp.float64 else 1e-6 + # Since we cannot check for angle equality, we check for rotation matrix + # equality + angles = xp.asarray([ + [45, 90, 35], + [35, -90, 20], + [35, 90, 25], + [25, -90, 15]]) + + seq = "".join(seq_tuple) + if intrinsic: + # Extrinsic rotation (w.r.t. global world) at lower case + # Intrinsic (WRT the object itself) upper case. + seq = seq.upper() + rotation = Rotation.from_euler(seq, angles, degrees=True) + mat_expected = rotation.as_matrix() + + with maybe_warn_gimbal_lock(not suppress_warnings, xp): + angle_estimates = rotation.as_euler( + seq, degrees=True, suppress_warnings=suppress_warnings + ) + mat_estimated = Rotation.from_euler(seq, angle_estimates, degrees=True).as_matrix() + + xp_assert_close(mat_expected, mat_estimated, atol=atol) + + +@make_xp_test_case( + (Rotation, "from_euler"), (Rotation, "as_matrix"), (Rotation, "as_euler") +) +@pytest.mark.parametrize("seq_tuple", list(permutations("xyz"))) +@pytest.mark.parametrize("intrinsic", (False, True)) +@pytest.mark.parametrize("suppress_warnings", (False, True)) +def test_as_euler_degenerate_symmetric_axes( + xp, seq_tuple, intrinsic, suppress_warnings +): + dtype = xpx.default_dtype(xp) + atol = 1e-12 if dtype == xp.float64 else 1e-6 + # Since we cannot check for angle equality, we check for rotation matrix + # equality + angles = xp.asarray([ + [15, 0, 60], + [35, 0, 75], + [60, 180, 35], + [15, -180, 25]]) + + # Rotation of the form A/B/A are rotation around symmetric axes + seq = "".join([seq_tuple[0], seq_tuple[1], seq_tuple[0]]) + if intrinsic: + # Extrinsic rotation (w.r.t. global world) at lower case + # Intrinsic (WRT the object itself) upper case. + seq = seq.upper() + rotation = Rotation.from_euler(seq, angles, degrees=True) + mat_expected = rotation.as_matrix() + + with maybe_warn_gimbal_lock(not suppress_warnings, xp): + angle_estimates = rotation.as_euler( + seq, degrees=True, suppress_warnings=suppress_warnings + ) + mat_estimated = Rotation.from_euler(seq, angle_estimates, degrees=True).as_matrix() + + xp_assert_close(mat_expected, mat_estimated, atol=atol) + + +@make_xp_test_case((Rotation, "from_euler"), (Rotation, "as_euler")) +@pytest.mark.parametrize("ndim", list(range(1, 4))) +def test_as_euler_nd_rotation(xp, ndim: int): + mat = xp.asarray([ + [0.0, -1, 0], + [1, 0, 0], + [0, 0, 1] + ]) + mat = xp.reshape(mat, (1,) * (ndim - 1) + (3, 3)) + angles = Rotation.from_matrix(mat).as_euler("xyz", degrees=True) + expected_angles = xp.asarray([0, 0, 90.0]) + expected_angles = xp.reshape(expected_angles, (1,) * (ndim - 1) + (3,)) + xp_assert_close(angles, expected_angles, atol=1e-12) + + +@make_xp_test_case((Rotation, "as_matrix"), (Rotation, "inv")) +def test_inv(xp): + dtype = xpx.default_dtype(xp) + atol = 1e-12 if dtype == xp.float64 else 1e-7 + rnd = np.random.RandomState(0) + n = 10 + # preserve use of old random_state during SPEC 7 transition + p = Rotation.random(num=n, random_state=rnd) + p = rotation_to_xp(p, xp) + p_mat = p.as_matrix() + q_mat = p.inv().as_matrix() + + result1 = p_mat @ q_mat + result2 = q_mat @ p_mat + + eye3d = xp.empty((n, 3, 3)) + eye3d = xpx.at(eye3d)[..., :3, :3].set(xp.eye(3)) + + xp_assert_close(result1, eye3d, atol=atol) + xp_assert_close(result2, eye3d, atol=atol) + + # Batched version + batch_shape = (10, 3, 7) + atol = 1e-12 if dtype == xp.float64 else 1e-6 + quat = xp.asarray(rnd.normal(size=batch_shape + (4,)), dtype=dtype) + r = Rotation.from_quat(quat) + p_mat = r.as_matrix() + q_mat = r.inv().as_matrix() + result1 = p_mat @ q_mat + result2 = q_mat @ p_mat + eye_nd = xp.empty(batch_shape + (3, 3)) + eye_nd = xpx.at(eye_nd)[..., :3, :3].set(xp.eye(3)) + xp_assert_close(result1, eye_nd, atol=atol) + xp_assert_close(result2, eye_nd, atol=atol) + + +@make_xp_test_case((Rotation, "inv"), (Rotation, "as_matrix")) +def test_inv_single_rotation(xp): + dtype = xpx.default_dtype(xp) + atol = 1e-12 if dtype == xp.float64 else 1e-7 + rng = np.random.default_rng(146972845698875399755764481408308808739) + p = rotation_to_xp(Rotation.random(rng=rng), xp) + q = p.inv() + + p_mat = p.as_matrix() + q_mat = q.as_matrix() + res1 = xp.matmul(p_mat, q_mat) + res2 = xp.matmul(q_mat, p_mat) + + eye = xp.eye(3) + + xp_assert_close(res1, eye, atol=atol) + xp_assert_close(res2, eye, atol=atol) + + x = rotation_to_xp(Rotation.random(num=1, rng=rng), xp) + y = x.inv() + + x_matrix = x.as_matrix() + y_matrix = y.as_matrix() + result1 = xp.linalg.matmul(x_matrix, y_matrix) + result2 = xp.linalg.matmul(y_matrix, x_matrix) + + eye3d = xp.empty((1, 3, 3)) + eye3d = xpx.at(eye3d)[..., :3, :3].set(xp.eye(3)) + + xp_assert_close(result1, eye3d, atol=atol) + xp_assert_close(result2, eye3d, atol=atol) + + +@make_xp_test_case((Rotation, "magnitude"), (Rotation, "inv")) +def test_identity_magnitude(xp): + n = 10 + r = rotation_to_xp(Rotation.identity(n), xp) + expected = xp.zeros(n) + xp_assert_close(r.magnitude(), expected) + xp_assert_close(r.inv().magnitude(), expected) + + +@make_xp_test_case((Rotation, "magnitude"), (Rotation, "inv")) +def test_single_identity_magnitude(xp): + r = rotation_to_xp(Rotation.identity(), xp) + assert r.magnitude() == 0 + assert r.inv().magnitude() == 0 + + +@make_xp_test_case((Rotation, "inv"), (Rotation, "magnitude")) +def test_identity_invariance(xp): + dtype = xpx.default_dtype(xp) + atol = 1e-12 if dtype == xp.float64 else 1e-7 + n = 10 + p = rotation_to_xp(Rotation.random(n, rng=0), xp) + q = rotation_to_xp(Rotation.identity(n), xp) + result = p * q + xp_assert_close(p.as_quat(), result.as_quat()) + + result = result * p.inv() + xp_assert_close(result.magnitude(), xp.zeros(n), atol=atol) + + +@make_xp_test_case((Rotation, "inv"), (Rotation, "magnitude")) +def test_single_identity_invariance(xp): + dtype = xpx.default_dtype(xp) + atol = 1e-12 if dtype == xp.float64 else 1e-7 + n = 10 + p = rotation_to_xp(Rotation.random(n, rng=0), xp) + q = rotation_to_xp(Rotation.identity(), xp) + result = p * q + xp_assert_close(p.as_quat(), result.as_quat()) + + result = result * p.inv() + xp_assert_close(result.magnitude(), xp.zeros(n), atol=atol) + + +def test_identity_shape(): # Not an xp test, identity is using numpy only for now + r = Rotation.identity(shape=()) + assert r.as_quat().shape == (4,) + r = Rotation.identity(shape=5) # Shape can be int + assert r.as_quat().shape == (5, 4) + r = Rotation.identity(shape=(2, 3)) + assert r.as_quat().shape == (2, 3, 4) + # Test values + r = Rotation.identity(shape=(2, 2, 3)) + xp_assert_equal(r.as_quat().reshape(-1, 4), np.tile(np.eye(4)[-1], (2 * 2 * 3, 1))) + # Errors + with pytest.raises(ValueError, match="`shape` must be an int or a tuple of ints"): + Rotation.identity(shape=2.5) + with pytest.raises(ValueError, match="Only one of `num` or `shape` can be"): + Rotation.identity(num=3, shape=(2, 2)) + with pytest.raises(TypeError, match="takes from 0 to 1 positional arguments"): + Rotation.identity(3, 3) + + +@make_xp_test_case((Rotation, "magnitude")) +@pytest.mark.parametrize("ndim", list(range(1, 4))) +def test_magnitude(xp, ndim: int): + quat_shape = (1,) * (ndim - 1) + (4,) + quat = xp.reshape(xp.eye(4), quat_shape + (4,)) + r = Rotation.from_quat(quat) + result = r.magnitude() + expected_result = xp.asarray([xp.pi, xp.pi, xp.pi, 0]) + expected_result = xp.reshape(expected_result, quat_shape) + xp_assert_close(result, expected_result) + + r = Rotation.from_quat(-quat) + result = r.magnitude() + xp_assert_close(result, expected_result) + + +@make_xp_test_case((Rotation, "magnitude")) +def test_magnitude_single_rotation(xp): + r = Rotation.from_quat(xp.eye(4)) + result1 = r[0].magnitude() + xp_assert_close(result1, xp.asarray(xp.pi)[()]) + + result2 = r[3].magnitude() + xp_assert_close(result2, xp.asarray(0.0)[()]) + + +@make_xp_test_case( + (Rotation, "inv"), (Rotation, "magnitude"), (Rotation, "approx_equal") +) +def test_approx_equal(xp): + rng = np.random.default_rng(146972845698875399755764481408308808739) + p = Rotation.random(10, rng=rng) + q = Rotation.random(10, rng=rng) + r_mag = (p * q.inv()).magnitude() + p = rotation_to_xp(p, xp) + q = rotation_to_xp(q, xp) + # ensure we get mix of Trues and Falses + atol = xp.asarray(np.median(r_mag)) + xp_assert_equal(p.approx_equal(q, atol), (xp.asarray(r_mag) < atol)) + + +@make_xp_test_case((Rotation, "from_rotvec"), (Rotation, "approx_equal")) +def test_approx_equal_single_rotation(xp): + # also tests passing single argument to approx_equal + p = Rotation.from_rotvec(xp.asarray([0, 0, 1e-9])) # less than default atol of 1e-8 + q = Rotation.from_quat(xp.eye(4)) + assert p.approx_equal(q[3]) + assert not p.approx_equal(q[0]) + # Regression test for gh-24769: single approx_equal should return a bool + assert isinstance(p.approx_equal(q[0]), np.bool_) or not is_numpy(xp) + + # test passing atol and using degrees + assert not p.approx_equal(q[3], atol=1e-10) + assert not p.approx_equal(q[3], atol=1e-8, degrees=True) + with pytest.warns(UserWarning, match="atol must be set"): + assert p.approx_equal(q[3], degrees=True) + + +@make_xp_test_case( + (Rotation, "inv"), (Rotation, "magnitude"), (Rotation, "approx_equal") +) +def test_approx_equal_batched(xp): + # Same shapes + batch_shape = (2, 10, 3) + rng = np.random.default_rng(0) + p = Rotation.from_quat(rng.normal(size=batch_shape + (4,))) + q = Rotation.from_quat(rng.normal(size=batch_shape + (4,))) + r_mag = (p * q.inv()).magnitude() # Must be computed as numpy array for np.median + p = rotation_to_xp(p, xp) + q = rotation_to_xp(q, xp) + assert r_mag.shape == batch_shape + # ensure we get mix of Trues and Falses + atol = xp.asarray(np.median(r_mag)) + xp_assert_equal(p.approx_equal(q, atol), (xp.asarray(r_mag) < atol)) + + # Broadcastable shapes of same length + p = Rotation.from_quat(rng.normal(size=batch_shape + (4,))) + q = Rotation.from_quat(rng.normal(size=(1, 10, 1, 4))) + r_mag = (p * q.inv()).magnitude() + p = rotation_to_xp(p, xp) + q = rotation_to_xp(q, xp) + assert r_mag.shape == batch_shape + atol = xp.asarray(np.median(r_mag)) + xp_assert_equal(p.approx_equal(q, atol), (xp.asarray(r_mag) < atol)) + + # Broadcastable shapes of different length + p = Rotation.from_quat(rng.normal(size=batch_shape + (4,))) + q = Rotation.from_quat(rng.normal(size=(1, 3, 4))) + r_mag = (p * q.inv()).magnitude() + p = rotation_to_xp(p, xp) + q = rotation_to_xp(q, xp) + assert r_mag.shape == batch_shape + atol = xp.asarray(np.median(r_mag)) + xp_assert_equal(p.approx_equal(q, atol), (xp.asarray(r_mag) < atol)) + + +@make_xp_test_case((Rotation, "approx_equal")) +def test_approx_equal_batched_input_validation(xp): + p = Rotation.from_quat(xp.ones((2, 3, 4))) + q = Rotation.from_quat(xp.ones((3, 2, 4))) + with pytest.raises(ValueError, match="broadcastable shapes"): + p.approx_equal(q) + + p = Rotation.from_quat(xp.ones((2, 4))) + q = Rotation.from_quat(xp.ones((3, 4))) + with pytest.raises(ValueError, match="broadcastable shapes"): + p.approx_equal(q) + + +@make_xp_test_case( + (Rotation, "from_rotvec"), (Rotation, "mean"), (Rotation, "magnitude") +) +@pytest.mark.parametrize("ndim", list(range(1, 4))) +def test_mean(xp, ndim: int): + axes = xp.concat((-xp.eye(3), xp.eye(3))) + axes = xp.reshape(axes, (1,) * (ndim - 1) + (6, 3)) + thetas = xp.linspace(0, xp.pi / 2, 100) + desired = xp.asarray(0.0)[()] + atol = 1e-6 if xp_default_dtype(xp) is xp.float32 else 1e-10 + for t in thetas: + r_mean = Rotation.from_rotvec(t * axes).mean() + assert r_mean.shape == () + xp_assert_close(r_mean.magnitude(), desired, atol=atol) + + +@make_xp_test_case( + (Rotation, "from_rotvec"), (Rotation, "mean"), (Rotation, "magnitude") +) +@pytest.mark.parametrize("ndim", list(range(1, 5))) +def test_mean_axis(xp, ndim: int): + axes = xp.tile(xp.concat((-xp.eye(3), xp.eye(3))), (3,) * (ndim - 1) + (1, 1)) + theta = xp.pi / 4 + r = Rotation.from_rotvec(theta * axes) + + # Test mean over last axis + desired = xp.full(axes.shape[:-2], 0.0) + if ndim == 1: + desired = desired[()] + atol = 1e-6 if xp_default_dtype(xp) is xp.float32 else 1e-10 + xp_assert_close(r.mean(axis=-1).magnitude(), desired, atol=atol) + + # Test tuple axes + desired = xp.full(axes.shape[1:-2], 0.0) + if ndim < 3: + desired = desired[()] + xp_assert_close(r.mean(axis=(0, -1)).magnitude(), desired, atol=atol) + + # Empty axis tuple should return Rotation unchanged + r_mean = r.mean(axis=()) + xp_assert_close(r_mean.as_quat(canonical=True), r.as_quat(canonical=True), + atol=atol) + + +@make_xp_test_case((Rotation, "mean"), (Rotation, "magnitude")) +def test_mean_compare_axis(xp): + # Create a random set of rotations and compare the mean over an axis with the + # mean without axis of the sliced quaternion + atol = 1e-10 if xpx.default_dtype(xp) == xp.float64 else 1e-6 + rng = np.random.default_rng(0) + q = xp.asarray(rng.normal(size=(4, 5, 6, 4)), dtype=xpx.default_dtype(xp)) + r = Rotation.from_quat(q) + + mean_0 = r.mean(axis=0) + for i in range(q.shape[1]): + for j in range(q.shape[2]): + mean_slice = Rotation.from_quat(q[:, i, j, ...]).mean() + xp_assert_close((mean_0[i][j] * mean_slice.inv()).magnitude(), + xp.asarray(0.0)[()], atol=atol) + mean_1_2 = r.mean(axis=(1, 2)) + for i in range(q.shape[0]): + mean_slice = Rotation.from_quat(q[i, ...]).mean() + xp_assert_close((mean_1_2[i] * mean_slice.inv()).magnitude(), + xp.asarray(0.0)[()], atol=atol) + + +@make_xp_test_case((Rotation, "from_rotvec"), (Rotation, "mean"), (Rotation, "inv"), + (Rotation, "magnitude")) +@pytest.mark.parametrize("ndim", list(range(1, 4))) +def test_weighted_mean(xp, ndim: int): + # test that doubling a weight is equivalent to including a rotation twice. + thetas = xp.linspace(0, xp.pi / 2, 100) + + # Create batched copies of the same setup + batch_shape = (ndim,) * (ndim - 1) + axes = xp.asarray([[0.0, 0, 0], [1, 0, 0], [1, 0, 0]]) + weights = xp.asarray([1, 2]) + axes = xp.tile(axes, batch_shape + (1, 1)) + weights = xp.tile(weights, batch_shape + (1,)) + + expected = xp.asarray(0.0)[()] + for t in thetas: + rw = Rotation.from_rotvec(t * axes[..., :2, :]) + mw = rw.mean(weights=weights) + + r = Rotation.from_rotvec(t * axes) + m = r.mean() + assert m.shape == () + xp_assert_close((m * mw.inv()).magnitude(), expected, atol=1e-6) + + +@make_xp_test_case((Rotation, "mean")) +def test_mean_input_validation(xp): + r = Rotation.from_quat(xp.eye(4)) + if is_lazy_array(r.as_quat()): + m = r.mean(weights=-xp.ones(4)) + assert xp.all(xp.isnan(m._quat)) + else: + with pytest.raises(ValueError, match="non-negative"): + r.mean(weights=-xp.ones(4)) + + # Test weight shape mismatch + r = Rotation.from_quat(xp.ones((3, 4))) + with pytest.raises(ValueError, match="Expected `weights` to"): + r.mean(weights=xp.ones((2,))) + r = Rotation.from_quat(xp.ones((2, 3, 4))) + with pytest.raises(ValueError, match="Expected `weights` to"): + r.mean(weights=xp.ones((2, 2))) + + # Test wrong axis + with pytest.raises(ValueError, match=r"axis .* is out of bounds"): + r.mean(axis=3) + with pytest.raises(ValueError, match=r"axis .* is out of bounds"): + r.mean(axis=(-1, 2)) + with pytest.raises(ValueError, match="`axis` must be None, int, or tuple of ints."): + r.mean(axis="0") + with pytest.raises(ValueError, match=r"axis .* is out of bounds"): + r.mean(axis=-12) + + +@make_xp_test_case((Rotation, "reduce")) +def test_reduction_no_indices(xp): + r = Rotation.from_quat(xp.asarray([0.0, 0.0, 0.0, 1.0])) + result = r.reduce(return_indices=False) + assert isinstance(result, Rotation) + + +@make_xp_test_case((Rotation, "reduce")) +def test_reduction_none_indices(xp): + r = Rotation.from_quat(xp.asarray([0.0, 0.0, 0.0, 1.0])) + result = r.reduce(return_indices=True) + assert type(result) is tuple + assert len(result) == 3 + + reduced, left_best, right_best = result + assert left_best is None + assert right_best is None + + +@make_xp_test_case((Rotation, "reduce"), (Rotation, "inv"), (Rotation, "magnitude")) +def test_reduction_scalar_calculation(xp): + dtype = xpx.default_dtype(xp) + atol = 1e-12 if dtype == xp.float64 else 1e-6 + rng = np.random.default_rng(146972845698875399755764481408308808739) + l_np = Rotation.random(5, rng=rng) + r_np = Rotation.random(10, rng=rng) + p_np = Rotation.random(7, rng=rng) + l = rotation_to_xp(l_np, xp) + r = rotation_to_xp(r_np, xp) + p = rotation_to_xp(p_np, xp) + reduced, left_best, right_best = p.reduce(l, r, return_indices=True) + + # Loop implementation of the vectorized calculation in Rotation.reduce + scalars = np.zeros((len(l_np), len(p_np), len(r_np))) + for i, li in enumerate(l_np): + for j, pj in enumerate(p_np): + for k, rk in enumerate(r_np): + scalars[i, j, k] = np.abs((li * pj * rk).as_quat()[3]) + scalars = np.reshape(np.moveaxis(scalars, 1, 0), (scalars.shape[1], -1)) + + max_ind = np.argmax(np.reshape(scalars, (len(p), -1)), axis=1) + left_best_check = xp.asarray(max_ind // len(r)) + right_best_check = xp.asarray(max_ind % len(r)) + assert xp.all(left_best == left_best_check) + assert xp.all(right_best == right_best_check) + + reduced_check = l[left_best_check] * p * r[right_best_check] + mag = (reduced.inv() * reduced_check).magnitude() + xp_assert_close(mag, xp.zeros(len(p)), atol=atol) + + +@make_xp_test_case((Rotation, "from_matrix"), (Rotation, "apply")) +def test_apply_single_rotation_single_point(xp): + dtype = xpx.default_dtype(xp) + mat = xp.asarray([ + [0, -1, 0], + [1, 0, 0], + [0, 0, 1] + ]) + r_1d = Rotation.from_matrix(mat) + r_2d = Rotation.from_matrix(xp.expand_dims(mat, axis=0)) + + v_1d = xp.asarray([1.0, 2, 3], dtype=dtype) + v_2d = xp.expand_dims(v_1d, axis=0) + v1d_rotated = xp.asarray([-2.0, 1, 3], dtype=dtype) + v2d_rotated = xp.expand_dims(v1d_rotated, axis=0) + + xp_assert_close(r_1d.apply(v_1d), v1d_rotated) + xp_assert_close(r_1d.apply(v_2d), v2d_rotated) + xp_assert_close(r_2d.apply(v_1d), v2d_rotated) + xp_assert_close(r_2d.apply(v_2d), v2d_rotated) + + v1d_inverse = xp.asarray([2.0, -1, 3], dtype=dtype) + v2d_inverse = xp.expand_dims(v1d_inverse, axis=0) + + xp_assert_close(r_1d.apply(v_1d, inverse=True), v1d_inverse) + xp_assert_close(r_1d.apply(v_2d, inverse=True), v2d_inverse) + xp_assert_close(r_2d.apply(v_1d, inverse=True), v2d_inverse) + xp_assert_close(r_2d.apply(v_2d, inverse=True), v2d_inverse) + + +@make_xp_test_case((Rotation, "from_matrix"), (Rotation, "apply")) +@pytest.mark.parametrize("ndim", list(range(1, 4))) +def test_apply_single_rotation_multiple_points(xp, ndim: int): + dtype = xpx.default_dtype(xp) + mat = xp.asarray([ + [0, -1, 0], + [1, 0, 0], + [0, 0, 1] + ]) + r1 = Rotation.from_matrix(mat) + r2 = Rotation.from_matrix(xp.expand_dims(mat, axis=0)) + + rng = np.random.default_rng(0) + batch_shape = (ndim,) * (ndim - 1) + v = xp.asarray(rng.normal(size=batch_shape + (2, 3)), dtype=dtype) + v_rotated = xp.stack([-v[..., 1], v[..., 0], v[..., 2]], axis=-1) + + xp_assert_close(r1.apply(v), v_rotated) + xp_assert_close(r2.apply(v), v_rotated) + + v_inverse = xp.stack([v[..., 1], -v[..., 0], v[..., 2]], axis=-1) + + xp_assert_close(r1.apply(v, inverse=True), v_inverse) + xp_assert_close(r2.apply(v, inverse=True), v_inverse) + + +@make_xp_test_case((Rotation, "from_matrix"), (Rotation, "apply")) +@pytest.mark.parametrize("ndim", list(range(1, 5))) +def test_apply_multiple_rotations_single_point(xp, ndim: int): + dtype = xpx.default_dtype(xp) + mat = np.empty((2, 3, 3)) + mat[0] = np.array([ + [0, -1, 0], + [1, 0, 0], + [0, 0, 1] + ]) + mat[1] = np.array([ + [1, 0, 0], + [0, 0, -1], + [0, 1, 0] + ]) + mat = xp.asarray(mat, dtype=dtype) + batch_shape = (ndim,) * (ndim - 1) + mat = xp.tile(mat, batch_shape + (1, 1, 1)) + r = Rotation.from_matrix(mat) + + v1 = xp.asarray([1, 2, 3]) + v2 = xp.expand_dims(v1, axis=0) + + v_rotated = xp.asarray([[-2.0, 1, 3], [1, -3, 2]]) + v_rotated = xp.tile(v_rotated, batch_shape + (1, 1)) + + xp_assert_close(r.apply(v1), v_rotated) + xp_assert_close(r.apply(v2), v_rotated) + + v_inverse = xp.asarray([[2.0, -1, 3], [1, 3, -2]]) + v_inverse = xp.tile(v_inverse, batch_shape + (1, 1)) + + xp_assert_close(r.apply(v1, inverse=True), v_inverse) + xp_assert_close(r.apply(v2, inverse=True), v_inverse) + + +@make_xp_test_case((Rotation, "from_matrix"), (Rotation, "apply")) +@pytest.mark.parametrize("ndim", list(range(1, 5))) +def test_apply_multiple_rotations_multiple_points(xp, ndim: int): + dtype = xpx.default_dtype(xp) + mat = np.empty((2, 3, 3)) + mat[0] = np.array([ + [0, -1, 0], + [1, 0, 0], + [0, 0, 1] + ]) + mat[1] = np.array([ + [1, 0, 0], + [0, 0, -1], + [0, 1, 0] + ]) + mat = xp.asarray(mat, dtype=dtype) + batch_shape = (ndim,) * (ndim - 1) + mat = xp.tile(mat, batch_shape + (1, 1, 1)) + r = Rotation.from_matrix(mat) + + v = xp.asarray([[1, 2, 3], [4, 5, 6]], dtype=dtype) + v_rotated = xp.asarray([[-2.0, 1, 3], [4, -6, 5]], dtype=dtype) + v_rotated = xp.tile(v_rotated, batch_shape + (1, 1)) + xp_assert_close(r.apply(v), v_rotated) + + v_inverse = xp.asarray([[2.0, -1, 3], [4, 6, -5]], dtype=dtype) + v_inverse = xp.tile(v_inverse, batch_shape + (1, 1)) + xp_assert_close(r.apply(v, inverse=True), v_inverse) + + +@make_xp_test_case((Rotation, "apply")) +def test_apply_shapes(xp): + rng = np.random.default_rng(0) + # Broadcast shape: (6, 5, 4, 2) ( + (3,) for vectors, + (4,) for rotations) + vector_shapes = [(), (1,), (2,), (1, 2), (5, 1, 2)] + rot_shapes = [(), (1,), (2,), (1, 2), (4, 2), (1, 4, 2), (5, 4, 2), (6, 5, 4, 2)] + + for q_shape, v_shape in product(rot_shapes, vector_shapes): + v = xp.asarray(rng.normal(size=v_shape + (3,))) + q = xp.asarray(rng.normal(size=q_shape + (4,))) + r = Rotation.from_quat(q) + shape = np.broadcast_shapes(q_shape, v_shape) + (3,) + x = r.apply(v) + assert x.shape == shape + x = r.apply(v, inverse=True) + assert x.shape == shape + + +def test_apply_array_like(): + rng = np.random.default_rng(123) + # Single vector + r = Rotation.random(rng=rng) + t = rng.uniform(-100, 100, size=(3,)) + v = r.apply(t.tolist()) + v_expected = r.apply(t) + xp_assert_close(v, v_expected, atol=1e-12) + # Multiple vectors + t = rng.uniform(-100, 100, size=(2, 3)) + v = r.apply(t.tolist()) + v_expected = r.apply(t) + xp_assert_close(v, v_expected, atol=1e-12) + + +def test_apply_matrix_equivalence(): + """Test documented equivalence for single rotation: + `apply(vectors) == vectors @ as_matrix().T.`""" + r = Rotation.from_rotvec([0, 0, 1]) + # Single vector (3,) + v = np.array([1.0, 0.0, 0.0]) + xp_assert_close(r.apply(v), v @ r.as_matrix().T) + # Multiple vectors (P, 3) + arr = np.array([[1, 0, 0], [1, 2, 3]], dtype=float) + xp_assert_close(r.apply(arr), arr @ r.as_matrix().T) + # (3, 3) case: `as_matrix() @ vectors` would not error but give wrong result + arr33 = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=float) + xp_assert_close(r.apply(arr33), arr33 @ r.as_matrix().T) + wrong_result = r.as_matrix() @ arr33 + assert not np.allclose(r.apply(arr33), wrong_result) + + +@make_xp_test_case((Rotation, "apply")) +def test_apply_input_validation(xp): + r = Rotation.from_quat(xp.ones(4)) + with pytest.raises(ValueError, match="Expected input of shape"): + r.apply(xp.ones(2)) + with pytest.raises(ValueError, match="Expected input of shape"): + r.apply(xp.ones((2, 2))) + r = Rotation.from_quat(xp.ones((2, 4))) + with pytest.raises(ValueError, match="Cannot broadcast"): + r.apply(xp.ones((3, 3))) + r = Rotation.from_quat(xp.ones((1, 7, 2, 4))) + with pytest.raises(ValueError, match="Cannot broadcast"): + r.apply(xp.ones((2, 2, 3))) + + +@make_xp_test_case( + (Rotation, "from_matrix"), (Rotation, "as_matrix"), (Rotation, "__getitem__") +) +@pytest.mark.parametrize("ndim", list(range(1, 4))) +def test_getitem(xp, ndim: int): + rng = np.random.default_rng(0) + quat = rng.normal(size=(2, ) + (ndim,) * (ndim - 1) + (4,)) + mat = xp.asarray(Rotation.from_quat(quat).as_matrix()) + r = Rotation.from_matrix(mat) + + xp_assert_close(r[0].as_matrix(), mat[0, ...], atol=1e-15) + xp_assert_close(r[1].as_matrix(), mat[1, ...], atol=1e-15) + xp_assert_close(r[:-1].as_matrix(), xp.expand_dims(mat[0, ...], axis=0), atol=1e-15) + + +@make_xp_test_case((Rotation, "__getitem__")) +def test_getitem_single(xp): + with pytest.raises(TypeError, match='not subscriptable'): + Rotation.from_quat(xp.asarray([0, 0, 0, 1]))[0] + + +@make_xp_test_case( + (Rotation, "from_matrix"), (Rotation, "__getitem__"), (Rotation, "as_matrix") +) +def test_getitem_array_like(): + mat = np.array([[[0.0, -1, 0], + [1, 0, 0], + [0, 0, 1]], + [[1, 0, 0], + [0, 0, -1], + [0, 1, 0]]]) + r = Rotation.from_matrix(mat) + xp_assert_close(r[[0]].as_matrix(), mat[[0]], atol=1e-15) + xp_assert_close(r[[0, 1]].as_matrix(), mat[[0, 1]], atol=1e-15) + + +@make_xp_test_case((Rotation, "__setitem__")) +def test_setitem_single(xp): + r = Rotation.from_quat(xp.asarray([0, 0, 0, 1])) + with pytest.raises(TypeError, match='not subscriptable'): + r[0] = Rotation.from_quat(xp.asarray([0, 0, 0, 1])) + + +@make_xp_test_case((Rotation, "__setitem__")) +def test_setitem_slice(xp): + rng = np.random.default_rng(146972845698875399755764481408308808739) + r1 = rotation_to_xp(Rotation.random(10, rng=rng), xp) + r2 = rotation_to_xp(Rotation.random(5, rng=rng), xp) + r1[1:6] = r2 + xp_assert_equal(r1[1:6].as_quat(), r2.as_quat()) + + # Multiple dimensions + r1 = Rotation.from_quat(xp.asarray(rng.normal(size=(3, 5, 4)))) + r2 = Rotation.from_quat(xp.asarray(rng.normal(size=(2, 5, 4)))) + r1[1:3] = r2 + xp_assert_equal(r1[1:3].as_quat(), r2.as_quat()) + + +@make_xp_test_case((Rotation, "__setitem__")) +def test_setitem_integer(xp): + rng = np.random.default_rng(146972845698875399755764481408308808739) + r1 = rotation_to_xp(Rotation.random(10, rng=rng), xp) + r2 = rotation_to_xp(Rotation.random(rng=rng), xp) + r1[1] = r2 + xp_assert_equal(r1[1].as_quat(), r2.as_quat()) + + # Multiple dimensions + r1 = Rotation.from_quat(xp.asarray(rng.normal(size=(3, 5, 4)))) + r2 = Rotation.from_quat(xp.asarray(rng.normal(size=(5, 4)))) + r1[1] = r2 + xp_assert_equal(r1[1].as_quat(), r2.as_quat()) + + +@make_xp_test_case((Rotation, "__setitem__")) +def test_setitem_wrong_type(xp): + r = rotation_to_xp(Rotation.random(10, rng=0), xp) + with pytest.raises(TypeError, match='Rotation object'): + r[0] = 1 + + +@make_xp_test_case((Rotation, "from_matrix")) +def test_n_rotations(xp): + mat = np.empty((2, 3, 3)) + mat[0] = np.array([ + [0, -1, 0], + [1, 0, 0], + [0, 0, 1] + ]) + mat[1] = np.array([ + [1, 0, 0], + [0, 0, -1], + [0, 1, 0] + ]) + mat = xp.asarray(mat) + r = Rotation.from_matrix(mat) + + assert_equal(len(r), 2) + assert_equal(len(r[:-1]), 1) + + +def test_random_rotation(): + # No xp testing since random rotations are always using NumPy + rng = np.random.default_rng(0) + assert_equal(Rotation.random(rng=rng).as_quat().shape, (4,)) + assert_equal(Rotation.random(None, rng=rng).as_quat().shape, (4,)) + assert_equal(Rotation.random(1, rng=rng).as_quat().shape, (1, 4)) + assert_equal(Rotation.random(5, rng=rng).as_quat().shape, (5, 4)) + # Shape argument + assert_equal(Rotation.random(rng=rng, shape=()).as_quat().shape, (4,)) + assert_equal(Rotation.random(rng=rng, shape=(3,)).as_quat().shape, (3, 4)) + assert_equal(Rotation.random(rng=rng, shape=(2, 3)).as_quat().shape, (2, 3, 4)) + # Values should be the same for num=prod(shape) + rng1, rng2 = np.random.default_rng(42), np.random.default_rng(42) + r_num = Rotation.random(6, rng=rng1) + r_shape = Rotation.random(rng=rng2, shape=(2, 3)) + xp_assert_close(r_num.as_quat(), r_shape.as_quat().reshape(6, 4), atol=1e-12) + # Errors + with pytest.raises(ValueError, match="Only one of `num` or `shape` can be"): + Rotation.random(num=3,rng=rng, shape=(2, 2)) + with pytest.raises(ValueError, match="`shape` must be an int or a tuple of ints"): + Rotation.random(rng=rng, shape=2.5) + with pytest.raises(TypeError, match="takes from 0 to 2 positional arguments"): + Rotation.random(1, rng, None) # Shape should be kwarg only + + +@make_xp_test_case((Rotation, "align_vectors"), (Rotation, "as_matrix")) +def test_align_vectors_no_rotation(xp): + dtype = xpx.default_dtype(xp) + atol = 1e-12 if dtype == xp.float64 else 1e-5 + x = xp.asarray([[1, 2, 3], [4, 5, 6]], dtype=dtype) + y = xp.asarray(x, copy=True) + + r, rssd = Rotation.align_vectors(x, y) + xp_assert_close(r.as_matrix(), xp.eye(3), atol=atol) + xp_assert_close(rssd, xp.asarray(0.0)[()], check_shape=False, atol=1e-6) + + +@make_xp_test_case((Rotation, "apply"), (Rotation, "align_vectors")) +def test_align_vectors_no_noise(xp): + dtype = xpx.default_dtype(xp) + atol = 1e-7 if dtype == xp.float64 else 2e-3 + rng = np.random.default_rng(14697284569885399755764481408308808739) + c = rotation_to_xp(Rotation.random(rng=rng), xp) + b = xp.asarray(rng.normal(size=(5, 3)), dtype=dtype) + a = c.apply(b) + + est, rssd = Rotation.align_vectors(a, b) + xp_assert_close(c.as_quat(), est.as_quat()) + xp_assert_close(rssd, xp.asarray(0.0)[()], check_shape=False, atol=atol) + + +@make_xp_test_case((Rotation, "align_vectors"), (Rotation, "apply")) +def test_align_vectors_improper_rotation(xp): + dtype = xpx.default_dtype(xp) + atol = 1e-7 if dtype == xp.float64 else 1e-3 + # Tests correct logic for issue #10444 + x = xp.asarray([[0.89299824, -0.44372674, 0.0752378], + [0.60221789, -0.47564102, -0.6411702]]) + y = xp.asarray([[0.02386536, -0.82176463, 0.5693271], + [-0.27654929, -0.95191427, -0.1318321]]) + + est, rssd = Rotation.align_vectors(x, y) + xp_assert_close(x, est.apply(y), atol=1e-6) + xp_assert_close(rssd, xp.asarray(0.0)[()], check_shape=False, atol=atol) + + +@make_xp_test_case((Rotation, "align_vectors")) +def test_align_vectors_rssd_sensitivity(xp): + rssd_expected = xp.asarray(0.141421356237308)[()] + sens_expected = xp.asarray([[0.2, 0. , 0.], + [0. , 1.5, 1.], + [0. , 1. , 1.]]) + atol = 1e-6 + a = xp.asarray([[0, 1, 0], [0, 1, 1], [0, 1, 1]]) + b = xp.asarray([[1, 0, 0], [1, 1.1, 0], [1, 0.9, 0]]) + rot, rssd, sens = Rotation.align_vectors(a, b, return_sensitivity=True) + xp_assert_close(rssd, rssd_expected, atol=atol) + xp_assert_close(sens, sens_expected, atol=atol) + + +@make_xp_test_case((Rotation, "align_vectors"), (Rotation, "as_matrix")) +def test_align_vectors_scaled_weights(xp): + n = 10 + a = xp.asarray(Rotation.random(n, rng=0).apply([1, 0, 0])) + b = xp.asarray(Rotation.random(n, rng=1).apply([1, 0, 0])) + scale = 2 + + est1, rssd1, cov1 = Rotation.align_vectors(a, b, xp.ones(n), True) + est2, rssd2, cov2 = Rotation.align_vectors(a, b, scale * xp.ones(n), True) + + xp_assert_close(est1.as_matrix(), est2.as_matrix()) + xp_assert_close(math.sqrt(scale) * rssd1, rssd2, atol=1e-6) + xp_assert_close(cov1, cov2) + + +@make_xp_test_case( + (Rotation, "apply"), (Rotation, "from_rotvec"), (Rotation, "align_vectors") +) +def test_align_vectors_noise(xp): + dtype = xpx.default_dtype(xp) + rng = np.random.default_rng(146972845698875399755764481408308808739) + n_vectors = 100 + rot = rotation_to_xp(Rotation.random(rng=rng), xp) + vectors = xp.asarray(rng.normal(size=(n_vectors, 3)), dtype=dtype) + result = rot.apply(vectors) + + # The paper adds noise as independently distributed angular errors + sigma = np.deg2rad(1) + tolerance = 1.5 * sigma + noise = Rotation.from_rotvec( + xp.asarray(rng.normal(size=(n_vectors, 3), scale=sigma), dtype=dtype) + ) + + # Attitude errors must preserve norm. Hence apply individual random + # rotations to each vector. + noisy_result = noise.apply(result) + + est, rssd, cov = Rotation.align_vectors(noisy_result, vectors, + return_sensitivity=True) + + # Use rotation compositions to find out closeness + error_vector = (rot * est.inv()).as_rotvec() + xp_assert_close(error_vector[0], xp.asarray(0.0)[()], atol=tolerance) + xp_assert_close(error_vector[1], xp.asarray(0.0)[()], atol=tolerance) + xp_assert_close(error_vector[2], xp.asarray(0.0)[()], atol=tolerance) + + # Check error bounds using covariance matrix + cov *= xp.asarray(sigma) + xp_assert_close(cov[0, 0], xp.asarray(0.0)[()], atol=tolerance) + xp_assert_close(cov[1, 1], xp.asarray(0.0)[()], atol=tolerance) + xp_assert_close(cov[2, 2], xp.asarray(0.0)[()], atol=tolerance) + + rssd_check = xp.sum((noisy_result - est.apply(vectors)) ** 2) ** 0.5 + xp_assert_close(rssd, rssd_check, check_shape=False) + + +@make_xp_test_case((Rotation, "align_vectors")) +def test_align_vectors_invalid_input(xp): + with pytest.raises(ValueError, match="Expected input `a` to have shape"): + a, b = xp.asarray([1, 2, 3, 4]), xp.asarray([1, 2, 3]) + Rotation.align_vectors(a, b) + + with pytest.raises(ValueError, match="Expected input `b` to have shape"): + a, b = xp.asarray([1, 2, 3]), xp.asarray([1, 2, 3, 4]) + Rotation.align_vectors(a, b) + + with pytest.raises(ValueError, match="Expected inputs `a` and `b` " + "to have same shapes"): + a, b = xp.asarray([[1, 2, 3], [4, 5, 6]]), xp.asarray([[1, 2, 3]]) + Rotation.align_vectors(a, b) + + with pytest.raises(ValueError, + match="Expected `weights` to be 1 dimensional"): + a, b = xp.asarray([[1, 2, 3]]), xp.asarray([[1, 2, 3]]) + weights = xp.asarray([[1]]) + Rotation.align_vectors(a, b, weights) + + with pytest.raises(ValueError, + match="Expected `weights` to have number of values"): + a, b = xp.asarray([[1, 2, 3], [4, 5, 6]]), xp.asarray([[1, 2, 3], [4, 5, 6]]) + weights = xp.asarray([1, 2, 3]) + Rotation.align_vectors(a, b, weights) + + a, b = xp.asarray([[1, 2, 3]]), xp.asarray([[1, 2, 3]]) + weights = xp.asarray([-1]) + if is_lazy_array(weights): + r, rssd = Rotation.align_vectors(a, b, weights) + assert xp.all(xp.isnan(r.as_quat())), "Quaternion should be nan" + assert xp.isnan(rssd), "RSSD should be nan" + else: + with pytest.raises(ValueError, + match="`weights` may not contain negative values"): + Rotation.align_vectors(a, b, weights) + + a, b = xp.asarray([[1, 2, 3], [4, 5, 6]]), xp.asarray([[1, 2, 3], [4, 5, 6]]) + weights = xp.asarray([xp.inf, xp.inf]) + if is_lazy_array(weights): + r, rssd = Rotation.align_vectors(a, b, weights) + assert xp.all(xp.isnan(r.as_quat())), "Quaternion should be nan" + assert xp.isnan(rssd), "RSSD should be nan" + else: + with pytest.raises(ValueError, + match="Only one infinite weight is allowed"): + Rotation.align_vectors(a, b, weights) + + a, b = xp.asarray([[0, 0, 0]]), xp.asarray([[1, 2, 3]]) + if is_lazy_array(a): + r, rssd = Rotation.align_vectors(a, b) + assert xp.all(xp.isnan(r.as_quat())), "Quaternion should be nan" + assert xp.isnan(rssd), "RSSD should be nan" + else: + with pytest.raises(ValueError, + match="Cannot align zero length primary vectors"): + Rotation.align_vectors(a, b) + + a, b = xp.asarray([[1, 2, 3], [4, 5, 6]]), xp.asarray([[1, 2, 3], [4, 5, 6]]) + weights = xp.asarray([xp.inf, 1]) + if is_lazy_array(a): + r, rssd, sens = Rotation.align_vectors(a, b, weights, return_sensitivity=True) + assert xp.all(xp.isnan(sens)), "Sensitivity matrix should be nan" + else: + with pytest.raises(ValueError, + match="Cannot return sensitivity matrix"): + Rotation.align_vectors(a, b, weights, return_sensitivity=True) + + a, b = xp.asarray([[1, 2, 3]]), xp.asarray([[1, 2, 3]]) + if is_lazy_array(a): + r, rssd, sens = Rotation.align_vectors(a, b, return_sensitivity=True) + assert xp.all(xp.isnan(sens)), "Sensitivity matrix should be nan" + else: + with pytest.raises(ValueError, + match="Cannot return sensitivity matrix"): + Rotation.align_vectors(a, b, return_sensitivity=True) + + # No broadcast support for align_vectors + a, b = xp.asarray([[[1, 2, 3]]]), xp.asarray([[[1, 2, 3]]]) + with pytest.raises(ValueError, + match="Expected inputs `a` and `b` to have shape"): + Rotation.align_vectors(a, b) + + +@make_xp_test_case( + (Rotation, "align_vectors"), (Rotation, "as_matrix"), (Rotation, "apply") +) +def test_align_vectors_align_constrain(xp): + # Align the primary +X B axis with the primary +Y A axis, and rotate about + # it such that the +Y B axis (residual of the [1, 1, 0] secondary b vector) + # is aligned with the +Z A axis (residual of the [0, 1, 1] secondary a + # vector) + dtype = xpx.default_dtype(xp) + atol = 1e-12 if dtype == xp.float64 else 1e-6 + b = xp.asarray([[1, 0, 0], [1, 1, 0]]) + a = xp.asarray([[0.0, 1, 0], [0, 1, 1]]) + m_expected = xp.asarray([[0.0, 0, 1], + [1, 0, 0], + [0, 1, 0]]) + R, rssd = Rotation.align_vectors(a, b, weights=xp.asarray([xp.inf, 1])) + xp_assert_close(R.as_matrix(), m_expected, atol=atol) + xp_assert_close(R.apply(b), a, atol=atol) # Pri and sec align exactly + xp_assert_close(rssd, xp.asarray(0.0)[()], atol=atol) + + # Do the same but with an inexact secondary rotation + b = xp.asarray([[1, 0, 0], [1, 2, 0]]) + rssd_expected = 1.0 + R, rssd = Rotation.align_vectors(a, b, weights=xp.asarray([xp.inf, 1])) + xp_assert_close(R.as_matrix(), m_expected, atol=atol) + xp_assert_close(R.apply(b)[0, ...], a[0, ...], atol=atol) # Only pri aligns exactly + assert xpx.isclose(rssd, rssd_expected, atol=atol, xp=xp) + a_expected = xp.asarray([[0.0, 1, 0], [0, 1, 2]]) + xp_assert_close(R.apply(b), a_expected, atol=atol) + + # Check random vectors + b = xp.asarray([[1, 2, 3], [-2, 3, -1]]) + a = xp.asarray([[-1.0, 3, 2], [1, -1, 2]]) + rssd_expected = 1.3101595297515016 + R, rssd = Rotation.align_vectors(a, b, weights=xp.asarray([xp.inf, 1])) + xp_assert_close(R.apply(b)[0, ...], a[0, ...], atol=atol) # Only pri aligns exactly + assert xpx.isclose(rssd, rssd_expected, atol=atol, xp=xp) + + +@make_xp_test_case((Rotation, "align_vectors"), (Rotation, "as_matrix")) +def test_align_vectors_near_inf(xp): + # align_vectors should return near the same result for high weights as for + # infinite weights. rssd will be different with floating point error on the + # exactly aligned vector being multiplied by a large non-infinite weight + dtype = xpx.default_dtype(xp) + if dtype == xp.float32: + pytest.skip("Align vectors near inf is numerically unstable in float32") + n = 100 + mats = [] + for i in range(6): + mats.append(Rotation.random(n, rng=10 + i).as_matrix()) + + for i in range(n): + # Get random pairs of 3-element vectors + a = xp.asarray(np.array([1 * mats[0][i][0], 2 * mats[1][i][0]]), dtype=dtype) + b = xp.asarray(np.array([3 * mats[2][i][0], 4 * mats[3][i][0]]), dtype=dtype) + + R, _ = Rotation.align_vectors(a, b, weights=[1e10, 1]) + R2, _ = Rotation.align_vectors(a, b, weights=[xp.inf, 1]) + xp_assert_close(R.as_matrix(), R2.as_matrix(), atol=1e-4) + + for i in range(n): + # Get random triplets of 3-element vectors + a = xp.asarray(np.array([1*mats[0][i][0], 2*mats[1][i][0], 3*mats[2][i][0]]), + dtype=dtype) + b = xp.asarray(np.array([4*mats[3][i][0], 5*mats[4][i][0], 6*mats[5][i][0]]), + dtype=dtype) + + R, _ = Rotation.align_vectors(a, b, weights=[1e10, 2, 1]) + R2, _ = Rotation.align_vectors(a, b, weights=[xp.inf, 2, 1]) + xp_assert_close(R.as_matrix(), R2.as_matrix(), atol=1e-4) + + +@make_xp_test_case((Rotation, "align_vectors"), (Rotation, "as_matrix")) +def test_align_vectors_parallel(xp): + atol = 1e-12 + a = xp.asarray([[1.0, 0, 0], [0, 1, 0]]) + b = xp.asarray([[0.0, 1, 0], [0, 1, 0]]) + m_expected = xp.asarray([[0.0, 1, 0], + [-1, 0, 0], + [0, 0, 1]]) + R, _ = Rotation.align_vectors(a, b, weights=[xp.inf, 1]) + xp_assert_close(R.as_matrix(), m_expected, atol=atol) + R, _ = Rotation.align_vectors(a[0, ...], b[0, ...]) + xp_assert_close(R.as_matrix(), m_expected, atol=atol) + xp_assert_close(R.apply(b[0, ...]), a[0, ...], atol=atol) + + b = xp.asarray([[1, 0, 0], [1, 0, 0]]) + m_expected = xp.asarray([[1.0, 0, 0], + [0, 1, 0], + [0, 0, 1]]) + R, _ = Rotation.align_vectors(a, b, weights=[xp.inf, 1]) + xp_assert_close(R.as_matrix(), m_expected, atol=atol) + R, _ = Rotation.align_vectors(a[0, ...], b[0, ...]) + xp_assert_close(R.as_matrix(), m_expected, atol=atol) + xp_assert_close(R.apply(b[0, ...]), a[0, ...], atol=atol) + + +@make_xp_test_case( + (Rotation, "align_vectors"), (Rotation, "magnitude"), (Rotation, "apply"), + (Rotation, "from_rotvec"), (Rotation, "as_rotvec"), (Rotation, "as_matrix") +) +def test_align_vectors_antiparallel(xp): + dtype = xpx.default_dtype(xp) + # Test exact 180 deg rotation + atol = 1e-12 if dtype == xp.float64 else 1e-7 + + as_to_test = np.array([[[1.0, 0, 0], [0, 1, 0]], + [[0, 1, 0], [1, 0, 0]], + [[0, 0, 1], [0, 1, 0]]]) + + bs_to_test = np.array([[-a[0], a[1]] for a in as_to_test]) + for a, b in zip(as_to_test, bs_to_test): + a, b = xp.asarray(a, dtype=dtype), xp.asarray(b, dtype=dtype) + R, _ = Rotation.align_vectors(a, b, weights=[xp.inf, 1]) + xp_assert_close(R.magnitude(), xp.asarray(xp.pi)[()], atol=atol) + xp_assert_close(R.apply(b[0, ...]), a[0, ...], atol=atol) + + # Test exact rotations near 180 deg + Rs = Rotation.random(100, rng=0) + dRs = Rotation.from_rotvec(Rs.as_rotvec()*1e-4) # scale down to small angle + a = [[ 1, 0, 0], [0, 1, 0]] + b = [[-1, 0, 0], [0, 1, 0]] + as_to_test = [] + for dR in dRs: + as_to_test.append(np.array([dR.apply(a[0]), a[1]])) + + # GPU computations may be less accurate. See e.g. + # https://github.com/jax-ml/jax/issues/18934 and + # https://docs.pytorch.org/docs/stable/generated/torch.set_float32_matmul_precision.html + # We currently have no unified way to check which device is being used. Cupy will + # always run on the GPU regardless of SCIPY_DEVICE, hence the explicit check for + # cupy. Note that the current implementation lets other frameworks, e.g. numpy, run + # on the CPU regardless of SCIPY_DEVICE but with increased GPU tolerances. + if xp_device_type(xp.asarray(0)) == "cuda": + atol = 1e-7 + + for a in as_to_test: + a, b = xp.asarray(a), xp.asarray(b) + R, _ = Rotation.align_vectors(a, b, weights=[xp.inf, 1]) + R2, _ = Rotation.align_vectors(a, b, weights=[1e10, 1]) + xp_assert_close(R.as_matrix(), R2.as_matrix(), atol=atol) + + +@make_xp_test_case((Rotation, "align_vectors"), (Rotation, "apply")) +def test_align_vectors_primary_only(xp): + dtype = xpx.default_dtype(xp) + atol = 1e-12 if dtype == xp.float64 else 1e-5 + mats_a = Rotation.random(100, rng=0).as_matrix() + mats_b = Rotation.random(100, rng=1).as_matrix() + + for mat_a, mat_b in zip(mats_a, mats_b): + # Get random 3-element unit vectors + a = xp.asarray(mat_a[0], dtype=dtype) + b = xp.asarray(mat_b[0], dtype=dtype) + + # Compare to align_vectors with primary only + R, rssd = Rotation.align_vectors(a, b) + xp_assert_close(R.apply(b), a, atol=atol) + xp_assert_close(rssd, xp.asarray(0.0)[()], atol=atol) + + +def test_align_vectors_array_like(): + rng = np.random.default_rng(123) + c = Rotation.random(rng=rng) + b = rng.normal(size=(5, 3)) + a = c.apply(b) + + est_expected, rssd_expected = Rotation.align_vectors(a, b) + est, rssd = Rotation.align_vectors(a.tolist(), b.tolist()) + xp_assert_close(est_expected.as_quat(), est.as_quat()) + xp_assert_close(rssd, rssd_expected) + + +@make_xp_test_case((Rotation, "apply"), (Rotation, "align_vectors")) +def test_align_vectors_mixed_dtypes(xp): + dtype = xpx.default_dtype(xp) + rng = np.random.default_rng(123) + c = rotation_to_xp(Rotation.random(rng=rng), xp) + b = xp.asarray(rng.normal(size=(5, 3)), dtype=dtype) + a = xp.asarray(c.apply(b), dtype=xp.float32) # Intentionally float32 + # Check that the dtype of the output is the result type of a and b + est, _ = Rotation.align_vectors(a, b) + xp_assert_close(est.as_quat(), c.as_quat()) + + +@make_xp_test_case((Rotation, "__repr__")) +def test_repr_single_rotation(xp): + q = xp.asarray([0, 0, 0, 1]) + actual = repr(Rotation.from_quat(q)) + if is_numpy(xp): + expected = """\ +Rotation.from_matrix(array([[1., 0., 0.], + [0., 1., 0.], + [0., 0., 1.]]))""" + assert actual == expected + else: + assert actual.startswith("Rotation.from_matrix(") + + +def test_repr_rotation_sequence(xp): + q = xp.asarray([[0.0, 1, 0, 1], [0, 0, 1, 1]]) / math.sqrt(2) + actual = f"{Rotation.from_quat(q)!r}" + if is_numpy(xp): + expected = """\ +Rotation.from_matrix(array([[[ 0., 0., 1.], + [ 0., 1., 0.], + [-1., 0., 0.]], + + [[ 0., -1., 0.], + [ 1., 0., 0.], + [ 0., 0., 1.]]]))""" + def stripped(s: str) -> str: + # don't fail due to leading whitespace differences + return "\n".join(map(str.lstrip, s.splitlines())) + + assert stripped(actual) == stripped(expected) + else: + assert actual.startswith("Rotation.from_matrix(") + + +@make_xp_test_case((Slerp, "__init__"), (Slerp, "__call__")) +def test_slerp(xp): + rnd = np.random.RandomState(0) + + key_rots = Rotation.from_quat(xp.asarray(rnd.uniform(size=(5, 4)))) + key_quats = key_rots.as_quat() + + key_times = [0, 1, 2, 3, 4] + interpolator = Slerp(key_times, key_rots) + assert isinstance(interpolator.times, type(xp.asarray(0))) + + times = [0, 0.5, 0.25, 1, 1.5, 2, 2.75, 3, 3.25, 3.60, 4] + interp_rots = interpolator(times) + interp_quats = interp_rots.as_quat() + + # Dot products are affected by sign of quaternions + mask = (interp_quats[:, -1] < 0)[:, None] + interp_quats = xp.where(mask, -interp_quats, interp_quats) + # Checking for quaternion equality, perform same operation + mask = (key_quats[:, -1] < 0)[:, None] + key_quats = xp.where(mask, -key_quats, key_quats) + + # Equality at keyframes, including both endpoints + xp_assert_close(interp_quats[0, ...], key_quats[0, ...]) + xp_assert_close(interp_quats[3, ...], key_quats[1, ...]) + xp_assert_close(interp_quats[5, ...], key_quats[2, ...]) + xp_assert_close(interp_quats[7, ...], key_quats[3, ...]) + xp_assert_close(interp_quats[10, ...], key_quats[4, ...]) + + # Constant angular velocity between keyframes. Check by equating + # cos(theta) between quaternion pairs with equal time difference. + cos_theta1 = xp.sum(interp_quats[0, ...] * interp_quats[2, ...]) + cos_theta2 = xp.sum(interp_quats[2, ...] * interp_quats[1, ...]) + xp_assert_close(cos_theta1, cos_theta2) + + cos_theta4 = xp.sum(interp_quats[3, ...] * interp_quats[4, ...]) + cos_theta5 = xp.sum(interp_quats[4, ...] * interp_quats[5, ...]) + xp_assert_close(cos_theta4, cos_theta5) + + # theta1: 0 -> 0.25, theta3 : 0.5 -> 1 + # Use double angle formula for double the time difference + cos_theta3 = xp.sum(interp_quats[1, ...] * interp_quats[3, ...]) + xp_assert_close(cos_theta3, 2 * (cos_theta1**2) - 1) + + # Miscellaneous checks + assert_equal(len(interp_rots), len(times)) + + +@make_xp_test_case((Slerp, "__init__")) +def test_slerp_rot_is_rotation(xp): + with pytest.raises(TypeError, match="must be a `Rotation` instance"): + r = xp.asarray([[1,2,3,4], + [0,0,0,1]]) + t = xp.asarray([0, 1]) + Slerp(t, r) + + +SLERP_EXCEPTION_MESSAGE = "must be a sequence of at least 2 rotations" + + +@make_xp_test_case((Slerp, "__init__")) +def test_slerp_single_rot(xp): + r = Rotation.from_quat(xp.asarray([[1.0, 2, 3, 4]])) + with pytest.raises(ValueError, match=SLERP_EXCEPTION_MESSAGE): + Slerp([1], r) + + +@make_xp_test_case((Slerp, "__init__")) +def test_slerp_rot_len0(xp): + r = Rotation.random() + r = Rotation.from_quat(xp.asarray(r.as_quat())) + with pytest.raises(ValueError, match=SLERP_EXCEPTION_MESSAGE): + Slerp([], r) + + +@make_xp_test_case((Slerp, "__init__")) +def test_slerp_rot_len1(xp): + r = Rotation.random(1) + r = Rotation.from_quat(xp.asarray(r.as_quat())) + with pytest.raises(ValueError, match=SLERP_EXCEPTION_MESSAGE): + Slerp([1], r) + + +@make_xp_test_case((Slerp, "__init__")) +def test_slerp_tensor_rot(xp): + r = Rotation.from_quat(xp.ones((2, 2, 4))) + with pytest.raises(ValueError, match="Rotations with more than 1 leading"): + Slerp([1, 2], r) + + +@make_xp_test_case((Slerp, "__init__")) +def test_slerp_time_dim_mismatch(xp): + with pytest.raises(ValueError, + match="times to be specified in a 1 dimensional array"): + rnd = np.random.RandomState(0) + r = Rotation.from_quat(xp.asarray(rnd.uniform(size=(2, 4)))) + t = xp.asarray([[1], + [2]]) + Slerp(t, r) + + +@make_xp_test_case((Slerp, "__init__")) +def test_slerp_num_rotations_mismatch(xp): + with pytest.raises(ValueError, match="number of rotations to be equal to " + "number of timestamps"): + rnd = np.random.RandomState(0) + r = Rotation.from_quat(xp.asarray(rnd.uniform(size=(5, 4)))) + t = xp.arange(7) + Slerp(t, r) + + +@make_xp_test_case((Slerp, "__init__")) +def test_slerp_equal_times(xp): + rnd = np.random.RandomState(0) + q = xp.asarray(rnd.uniform(size=(5, 4))) + r = Rotation.from_quat(q) + t = [0, 1, 2, 2, 4] + if is_lazy_array(q): + s = Slerp(t, r) + assert xp.all(xp.isnan(s.times)) + else: + with pytest.raises(ValueError, match="strictly increasing order"): + Slerp(t, r) + + +@make_xp_test_case((Slerp, "__init__")) +def test_slerp_decreasing_times(xp): + rnd = np.random.RandomState(0) + q = xp.asarray(rnd.uniform(size=(5, 4))) + r = Rotation.from_quat(q) + t = [0, 1, 3, 2, 4] + if is_lazy_array(q): + s = Slerp(t, r) + assert xp.all(xp.isnan(s.times)) + else: + with pytest.raises(ValueError, match="strictly increasing order"): + Slerp(t, r) + + +@make_xp_test_case((Slerp, "__init__"), (Slerp, "__call__")) +def test_slerp_call_time_dim_mismatch(xp): + rnd = np.random.RandomState(0) + r = Rotation.from_quat(xp.asarray(rnd.uniform(size=(5, 4)))) + t = xp.arange(5) + s = Slerp(t, r) + + with pytest.raises(ValueError, + match="`times` must be at most 1-dimensional."): + interp_times = xp.asarray([[3.5], + [4.2]]) + s(interp_times) + + +@make_xp_test_case((Slerp, "__init__"), (Slerp, "__call__")) +def test_slerp_call_time_out_of_range(xp): + rnd = np.random.RandomState(0) + r = Rotation.from_quat(xp.asarray(rnd.uniform(size=(5, 4)))) + t = xp.arange(5) + 1 + s = Slerp(t, r) + + times_low = xp.asarray([0, 1, 2]) + times_high = xp.asarray([1, 2, 6]) + if is_lazy_array(times_low): + q = s(times_low).as_quat() + in_range = xp.logical_and(times_low >= xp.min(t), times_low <= xp.max(t)) + assert xp.all(xp.isnan(q[~in_range, ...])) + assert xp.all(~xp.isnan(q[in_range, ...])) + q = s(times_high).as_quat() + in_range = xp.logical_and(times_high >= xp.min(t), times_high <= xp.max(t)) + assert xp.all(xp.isnan(q[~in_range, ...])) + assert xp.all(~xp.isnan(q[in_range, ...])) + else: + with pytest.raises(ValueError, match="times must be within the range"): + s(times_low) + with pytest.raises(ValueError, match="times must be within the range"): + s(times_high) + + +@make_xp_test_case( + (Slerp, "__init__"), (Slerp, "__call__"), (Rotation, "from_euler"), + (Rotation, "inv"), (Rotation, "magnitude") +) +def test_slerp_call_scalar_time(xp): + dtype = xpx.default_dtype(xp) + atol = 1e-16 if dtype == xp.float64 else 1e-7 + r = Rotation.from_euler('X', xp.asarray([[0], [80]]), degrees=True) + s = Slerp([0, 1], r) + + r_interpolated = s(0.25) + r_interpolated_expected = Rotation.from_euler('X', xp.asarray(20), degrees=True) + + delta = r_interpolated * r_interpolated_expected.inv() + + xp_assert_close(delta.magnitude(), xp.asarray(0.0)[()], atol=atol) + + +@make_xp_test_case((Rotation, "__mul__")) +def test_multiplication(xp): + r1 = Rotation.from_quat(xp.asarray([0, 0, 0, 1])) + r2 = Rotation.from_quat(xp.asarray([0, 0, 0, 1])) + r3 = r1 * r2 + xp_assert_close(r3.as_quat(), xp.asarray([0.0, 0, 0, 1])) + + # Check that multiplication with other types fails + with pytest.raises(TypeError, match="unsupported operand type"): + r1 * 2 + # Check that __mul__ returns NotImplemented so that other types can implement + # __rmul__. See https://github.com/scipy/scipy/issues/21541 + assert r1.__mul__(1) is NotImplemented + + +@make_xp_test_case((Rotation, "__mul__")) +def test_multiplication_nd(xp): + # multiple dimensions + rng = np.random.default_rng(0) + r1 = Rotation.from_quat(xp.asarray(rng.normal(size=(2, 3, 4)))) + r2 = Rotation.from_quat(xp.asarray(rng.normal(size=(2, 3, 4)))) + r3 = r1 * r2 + assert r3.as_quat().shape == (2, 3, 4) + + # same shape len, different dimensions + r1 = Rotation.from_quat(xp.asarray(rng.normal(size=(1, 3, 4)))) + r2 = Rotation.from_quat(xp.asarray(rng.normal(size=(2, 1, 4)))) + r3 = r1 * r2 + assert r3.as_quat().shape == (2, 3, 4) + + # different shape len, different dimensions + r1 = Rotation.from_quat(xp.asarray(rng.normal(size=(3, 1, 4, 4)))) + r2 = Rotation.from_quat(xp.asarray(rng.normal(size=(2, 4, 4)))) + r3 = r1 * r2 + assert r3.as_quat().shape == (3, 2, 4, 4) + + # transition between 2D and 3D with 2D rotation as first argument. This needs to + # choose the xp_backend even though r1's backend is cython + r1 = Rotation.from_quat(xp.asarray(rng.normal(size=(2, 4)))) + r2 = Rotation.from_quat(xp.asarray(rng.normal(size=(2, 2, 4)))) + r3 = r1 * r2 + assert r3.as_quat().shape == (2, 2, 4) + + +@make_xp_test_case((Rotation, "__mul__")) +def test_multiplication_errors(xp): + rng = np.random.default_rng(0) + r1 = Rotation.from_quat(xp.asarray(rng.normal(size=(2, 4)))) + r2 = Rotation.from_quat(xp.asarray(rng.normal(size=(1, 4, 4)))) + with pytest.raises(ValueError, match="Cannot broadcast"): + r1 * r2 + + +@make_xp_test_case((Rotation, "__mul__")) +def test_multiplication_stability(xp): + qs = rotation_to_xp(Rotation.random(50, rng=0), xp) + rs = rotation_to_xp(Rotation.random(1000, rng=1), xp) + expected = xp.ones(len(rs)) + for r in qs: + rs = rs * r * rs + xp_assert_close(xp_vector_norm(rs.as_quat(), axis=1), expected) + + +@make_xp_test_case( + (Rotation, "inv"), (Rotation, "__pow__"), (Rotation, "inv"), + (Rotation, "magnitude"), (Rotation, "from_rotvec"), (Rotation, "as_rotvec") +) +@pytest.mark.parametrize("ndim", list(range(1, 4))) +def test_pow(xp, ndim: int): + dtype = xpx.default_dtype(xp) + atol = 1e-14 if dtype == xp.float64 else 1e-6 + rng = np.random.default_rng(0) + batch_shape = (ndim,) * (ndim - 1) + quat = rng.normal(size=batch_shape + (4,)) + p = Rotation.from_quat(xp.asarray(quat)) + p_inv = p.inv() + # Test the short-cuts and other integers + for n in [-5, -2, -1, 0, 1, 2, 5]: + # Test accuracy + q = p ** n + q_identity = xp.asarray([0., 0, 0, 1]) + # Regression test for gh-24436 + assert isinstance(q._quat, type(q_identity)) + r = Rotation.from_quat(xp.tile(q_identity, batch_shape + (1,))) + for _ in range(abs(n)): + if n > 0: + r = r * p + else: + r = r * p_inv + ang = (q * r.inv()).magnitude() + assert xp.all(ang < atol) + + # Test shape preservation + r = Rotation.from_quat(xp.tile(q_identity, batch_shape + (1,))) + assert (r**n).as_quat().shape == batch_shape + (4,) + + # Large angle fractional + for n in [-1.5, -0.5, -0.0, 0.0, 0.5, 1.5]: + q = p ** n + r = Rotation.from_rotvec(n * p.as_rotvec()) + xp_assert_close(q.as_quat(), r.as_quat(), atol=atol) + + # Array exponent + n = [-5, -2, -1.5, -1, -0.5, -0.0, 0, 0.0, 0.5, 1.0, 1.5, 2] + for exponent in n: + r = p ** exponent + r_array = p ** xp.asarray([exponent]) # Test with 1D array + xp_assert_close(r.as_quat(), r_array.as_quat()) + r_array = p ** xp.asarray(exponent) # Test with scalar + xp_assert_close(r.as_quat(), r_array.as_quat()) + + # Small angle + rotvec = xp.zeros(batch_shape + (3,)) + rotvec = xpx.at(rotvec)[..., 0].set(1e-12) + p = Rotation.from_rotvec(rotvec) + n = 3 + q = p ** n + r = Rotation.from_rotvec(n * p.as_rotvec()) + xp_assert_close(q.as_quat(), r.as_quat(), atol=atol) + + # Array exponent + q = p ** xp.asarray([n]) # Test with 1D array + r = Rotation.from_rotvec(n * p.as_rotvec()) + xp_assert_close(q.as_quat(), r.as_quat(), atol=atol) + q = p ** xp.asarray(n) # Test with scalar + r = Rotation.from_rotvec(n * p.as_rotvec()) + xp_assert_close(q.as_quat(), r.as_quat(), atol=atol) + + +@make_xp_test_case((Rotation, "__pow__")) +def test_pow_errors(xp): + p = rotation_to_xp(Rotation.random(rng=0), xp) + with pytest.raises(NotImplementedError, match='modulus not supported'): + pow(p, 1, 1) + with pytest.raises(ValueError, match="Array exponent must be a scalar"): + p ** xp.asarray([1, 2]) + with pytest.raises(ValueError, match="Array exponent must be a scalar"): + p ** xp.asarray([[1], [2]]) + + +def test_rotation_within_numpy_array(): + # TODO: Do we want to support this for all Array API frameworks? + single = Rotation.random(rng=0) + multiple = Rotation.random(2, rng=1) + + array = np.array(single) + assert_equal(array.shape, ()) + + array = np.array(multiple) + assert_equal(array.shape, (2,)) + xp_assert_close(array[0].as_matrix(), multiple[0].as_matrix()) + xp_assert_close(array[1].as_matrix(), multiple[1].as_matrix()) + + array = np.array([single]) + assert_equal(array.shape, (1,)) + assert_equal(array[0], single) + + array = np.array([multiple]) + assert_equal(array.shape, (1, 2)) + xp_assert_close(array[0, 0].as_matrix(), multiple[0].as_matrix()) + xp_assert_close(array[0, 1].as_matrix(), multiple[1].as_matrix()) + + array = np.array([single, multiple], dtype=object) + assert_equal(array.shape, (2,)) + assert_equal(array[0], single) + assert_equal(array[1], multiple) + + array = np.array([multiple, multiple, multiple]) + assert_equal(array.shape, (3, 2)) + + +@make_xp_test_case((Rotation, "as_matrix")) +@pytest.mark.skip_xp_backends("array_api_strict", + reason="array API doesn't support pickling") +def test_pickling(xp): + r = Rotation.from_quat(xp.asarray([0, 0, math.sin(np.pi/4), math.cos(np.pi/4)])) + pkl = pickle.dumps(r) + unpickled = pickle.loads(pkl) + xp_assert_close(r.as_matrix(), unpickled.as_matrix(), atol=1e-15) + + +@make_xp_test_case((Rotation, "as_matrix")) +@pytest.mark.skip_xp_backends("array_api_strict", + reason="array API doesn't support deepcopy") +def test_deepcopy(xp): + r = Rotation.from_quat(xp.asarray([0, 0, math.sin(np.pi/4), math.cos(np.pi/4)])) + r1 = copy.deepcopy(r) + xp_assert_close(r.as_matrix(), r1.as_matrix(), atol=1e-15) + + +def test_as_euler_contiguous(): + # The Array API does not specify contiguous arrays, so we can only check for NumPy + r = Rotation.from_quat([0, 0, 0, 1]) + e1 = r.as_euler('xyz') # extrinsic euler rotation + e2 = r.as_euler('XYZ') # intrinsic + assert e1.flags['C_CONTIGUOUS'] is True + assert e2.flags['C_CONTIGUOUS'] is True + assert all(i >= 0 for i in e1.strides) + assert all(i >= 0 for i in e2.strides) + + +@make_xp_test_case((Rotation, "concatenate")) +def test_concatenate(xp): + rotation = rotation_to_xp(Rotation.random(10, rng=0), xp) + sizes = [1, 2, 3, 1, 3] + starts = [0] + list(np.cumsum(sizes)) + split = [rotation[i:i + n] for i, n in zip(starts, sizes)] + result = Rotation.concatenate(split) + xp_assert_equal(rotation.as_quat(), result.as_quat()) + + # Test Rotation input for multiple rotations + result = Rotation.concatenate(rotation) + xp_assert_equal(rotation.as_quat(), result.as_quat()) + + # Test that a copy is returned + assert rotation is not result + + # Test Rotation input for single rotations + rng = np.random.default_rng(0) + quat = xp.asarray(rng.normal(size=(5, 2, 4))) + rotation = Rotation.from_quat(quat) + r1 = Rotation.from_quat(quat[:3, ...]) + r2 = Rotation.from_quat(quat[3:, ...]) + result = Rotation.concatenate([r1, r2]) + xp_assert_equal(rotation.as_quat(), result.as_quat()) + + +@make_xp_test_case((Rotation, "concatenate")) +def test_concatenate_wrong_type(xp): + with pytest.raises(TypeError, match='Rotation objects only'): + rot = Rotation(xp.asarray(Rotation.identity().as_quat())) + Rotation.concatenate([rot, 1, None]) + + +@make_xp_test_case((Rotation, "concatenate")) +def test_concatenate_wrong_shape(xp): + r1 = Rotation.from_quat(xp.ones((5, 2, 4))) + r2 = Rotation.from_quat(xp.ones((1, 4))) + # Frameworks throw inconsistent error types on concat failures + with pytest.raises((ValueError, RuntimeError, TypeError)): + Rotation.concatenate([r1, r2]) + + +# Regression test for gh-16663 +@make_xp_test_case() +def test_len_and_bool(xp): + rotation_multi_one = Rotation(xp.asarray([[0, 0, 0, 1]])) + rotation_multi = Rotation(xp.asarray([[0, 0, 0, 1], [0, 0, 0, 1]])) + rotation_single = Rotation(xp.asarray([0, 0, 0, 1])) + + assert len(rotation_multi_one) == 1 + assert len(rotation_multi) == 2 + with pytest.raises(TypeError, match="Single rotation has no len()."): + len(rotation_single) + + rotation_batched = Rotation.from_quat(xp.ones((3, 2, 4))) + assert len(rotation_batched) == 3 + + # Rotation should always be truthy. See gh-16663 + assert rotation_multi_one + assert rotation_multi + assert rotation_single + + +@make_xp_test_case((Rotation, "from_davenport")) +def test_from_davenport_single_rotation(xp): + axis = xp.asarray([0, 0, 1]) + quat = Rotation.from_davenport(axis, 'extrinsic', 90, + degrees=True).as_quat() + expected_quat = xp.asarray([0.0, 0, 1, 1]) / math.sqrt(2) + xp_assert_close(quat, expected_quat) + + +@make_xp_test_case((Rotation, "from_rotvec"), (Rotation, "from_davenport")) +def test_from_davenport_one_or_two_axes(xp): + ez = xp.asarray([0.0, 0, 1]) + ey = xp.asarray([0.0, 1, 0]) + + # Single rotation, single axis, axes.shape == (3, ) + rot = Rotation.from_rotvec(ez * xp.pi/4) + rot_dav = Rotation.from_davenport(ez, 'e', xp.pi/4) + xp_assert_close(rot.as_quat(canonical=True), rot_dav.as_quat(canonical=True)) + + # Single rotation, single axis, axes.shape == (1, 3), angles.shape == (1, ) + # -> Still single rotation + axes = xp.reshape(ez, (1, 3)) # Torch can't create tensors from xp.asarray([ez]) + rot = Rotation.from_rotvec(ez * xp.pi/4) + rot_dav = Rotation.from_davenport(axes, 'e', [xp.pi/4]) + xp_assert_close(rot.as_quat(canonical=True), rot_dav.as_quat(canonical=True)) + + # Single rotation, two axes, axes.shape == (2, 3) + axes = xp.stack([ez, ey], axis=0) + rot = Rotation.from_rotvec(axes * xp.asarray([[xp.pi/4], [xp.pi/6]])) + rot = rot[0] * rot[1] + axes_dav = xp.stack([ey, ez], axis=0) + rot_dav = Rotation.from_davenport(axes_dav, 'e', [xp.pi/6, xp.pi/4]) + xp_assert_close(rot.as_quat(canonical=True), rot_dav.as_quat(canonical=True)) + + # Two rotations, single axis, axes.shape == (3, ) + axes = xp.stack([ez, ez], axis=0) + rot = Rotation.from_rotvec(axes * xp.asarray([[xp.pi/6], [xp.pi/4]])) + axes_dav = xp.reshape(ez, (1, 3)) + rot_dav = Rotation.from_davenport(axes_dav, 'e', [[xp.pi/6], [xp.pi/4]]) + xp_assert_close(rot.as_quat(canonical=True), rot_dav.as_quat(canonical=True)) + + +@make_xp_test_case((Rotation, "from_davenport")) +@pytest.mark.parametrize("ndim", list(range(1, 4))) +def test_from_davenport_shapes(xp, ndim: int): + # The shape rules for ND rotations are as follows: + # axes.shape[-2] must be angles.shape[-1] + # Resulting shape is np.broadcast_shapes(axes.shape[:-2], angles.shape[:-1]) + (4,) + rng = np.random.default_rng(0) + batch_shape = (ndim,) * (ndim - 1) + # Create random, orthogonal axes + r = Rotation.from_quat(xp.asarray(rng.normal(size=(4,)))) + axes = r.as_matrix() + # axes = (3,) + angles = xp.asarray(rng.normal(size=batch_shape + (1,))) + rot = Rotation.from_davenport(axes[0, ...], 'e', angles) + assert rot.as_quat().shape == batch_shape + (4,) + # axes = (1, 3) + angles = xp.asarray(rng.normal(size=batch_shape + (1,))) + rot = Rotation.from_davenport(axes[0, None, ...], 'e', angles) + assert rot.as_quat().shape == batch_shape + (4,) + # axes = (2, 3) + angles = xp.asarray(rng.normal(size=batch_shape + (2,))) + rot = Rotation.from_davenport(axes[:2, ...], 'e', angles) + assert rot.as_quat().shape == batch_shape + (4,) + + # axes = (...,3, 3) + r = Rotation.from_quat(xp.asarray(rng.normal(size=batch_shape + (4,)))) + axes = r.as_matrix() + angles = xp.asarray(rng.normal(size=batch_shape + (3,))) + rot = Rotation.from_davenport(axes, 'e', angles) + assert rot.as_quat().shape == batch_shape + (4,) + + +@make_xp_test_case((Rotation, "from_davenport")) +def test_from_davenport_broadcast(xp): + rng = np.random.default_rng(0) + # Create random, orthogonal axes + r = Rotation.from_quat(xp.asarray(rng.normal(size=(4, 9, 1, 4)))) + axes = r.as_matrix() + angles = xp.asarray(rng.normal(size=(1, 4, 3))) + rot = Rotation.from_davenport(axes, 'e', angles) + # (4, 9, 1, 3) + (3,) axes, (1, 4, 3) angles -> (4, 9, 4) + (4,) for quaternion + assert rot.as_quat().shape == (4, 9, 4, 4) + + +@make_xp_test_case((Rotation, "from_davenport")) +def test_from_davenport_invalid_input(xp): + ez = [0, 0, 1] + ey = [0, 1, 0] + ezy = [0, 1, 1] + # We can only raise in non-lazy frameworks. + axes = xp.asarray([ez, ezy]) + if is_lazy_array(axes): + q = Rotation.from_davenport(axes, 'e', [0, 0]).as_quat() + assert xp.all(xp.isnan(q)) + else: + with pytest.raises(ValueError, match="must be orthogonal"): + Rotation.from_davenport(axes, 'e', [0, 0]) + axes = xp.asarray([ez, ey, ezy]) + if is_lazy_array(axes): + q = Rotation.from_davenport(axes, 'e', [0, 0, 0]).as_quat() + assert xp.all(xp.isnan(q)) + else: + with pytest.raises(ValueError, match="must be orthogonal"): + Rotation.from_davenport(axes, 'e', [0, 0, 0]) + with pytest.raises(ValueError, match="order should be"): + Rotation.from_davenport(xp.asarray([ez]), 'xyz', [0]) + with pytest.raises(ValueError, match="Expected `angles`"): + Rotation.from_davenport(xp.asarray([ez, ey, ez]), 'e', [0, 1, 2, 3]) + with pytest.raises(ValueError, match="Expected `angles`"): # Too many angles + Rotation.from_davenport(xp.asarray(ez), 'e', [0, 1]) + with pytest.raises(ValueError, match="Expected `angles`"): # Too few angles + Rotation.from_davenport(xp.asarray([ez, ey, ez]), 'e', [0, 1]) + + +def test_from_davenport_array_like(): + rng = np.random.default_rng(123) + # Single rotation + e1 = np.array([1, 0, 0]) + e2 = np.array([0, 1, 0]) + e3 = np.array([0, 0, 1]) + r_expected = Rotation.random(rng=rng) + angles = r_expected.as_davenport([e1, e2, e3], 'e') + r = Rotation.from_davenport([e1, e2, e3], 'e', angles.tolist()) + assert r_expected.approx_equal(r, atol=1e-12) + + # Multiple rotations + r_expected = Rotation.random(2, rng=rng) + angles = r_expected.as_davenport([e1, e2, e3], 'e') + r = Rotation.from_davenport([e1, e2, e3], 'e', angles.tolist()) + assert np.all(r_expected.approx_equal(r, atol=1e-12)) + + +@make_xp_test_case((Rotation, "from_davenport"), (Rotation, "as_davenport")) +def test_as_davenport(xp): + dtype = xpx.default_dtype(xp) + rnd = np.random.RandomState(0) + n = 100 + angles = np.empty((n, 3)) + angles[:, 0] = rnd.uniform(low=-np.pi, high=np.pi, size=(n,)) + angles_middle = rnd.uniform(low=0, high=np.pi, size=(n,)) + angles[:, 2] = rnd.uniform(low=-np.pi, high=np.pi, size=(n,)) + lambdas = rnd.uniform(low=0, high=np.pi, size=(20,)) + + e1 = xp.asarray([1.0, 0, 0]) + e2 = xp.asarray([0.0, 1, 0]) + + for lamb in lambdas: + e3 = xp.asarray(Rotation.from_rotvec(lamb*e2).apply(e1)) + ax_lamb = xp.stack([e1, e2, e3], axis=0) + angles[:, 1] = angles_middle - lamb + for order in ['extrinsic', 'intrinsic']: + ax = ax_lamb if order == "intrinsic" else xp.flip(ax_lamb, axis=0) + rot = Rotation.from_davenport(ax, order, xp.asarray(angles, dtype=dtype)) + angles_dav = rot.as_davenport(ax, order) + xp_assert_close(angles_dav, xp.asarray(angles, dtype=dtype)) + + +@make_xp_test_case((Rotation, "from_davenport"), (Rotation, "as_davenport")) +def test_as_davenport_nd(xp): + rng = np.random.default_rng(0) + r = Rotation.from_quat(xp.asarray(rng.normal(size=(4, 9, 1, 4)))) + axes = r.as_matrix() # Get orthogonal axes + angles = xp.asarray(rng.uniform(low=-np.pi, high=np.pi, size=(4, 9, 1, 3))) + angles = xpx.at(angles)[..., 1].set(angles[..., 1] / 2) + + for order in ['extrinsic', 'intrinsic']: + if order == "intrinsic": + axes = xp.flip(axes, axis=-2) + rot = Rotation.from_davenport(axes, order, angles) + angles_dav = rot.as_davenport(axes, order) + xp_assert_close(angles_dav, angles) + + +@make_xp_test_case( + (Rotation, "from_davenport"), (Rotation, "as_davenport"), (Rotation, "as_matrix") +) +@pytest.mark.parametrize("suppress_warnings", (False, True)) +def test_as_davenport_degenerate(xp, suppress_warnings): + dtype = xpx.default_dtype(xp) + atol = 1e-12 if dtype == xp.float64 else 1e-6 + # Since we cannot check for angle equality, we check for rotation matrix + # equality + rnd = np.random.RandomState(0) + n = 5 + angles = np.empty((n, 3)) + + # symmetric sequences + angles[:, 0] = rnd.uniform(low=-np.pi, high=np.pi, size=(n,)) + angles_middle = [rnd.choice([0, np.pi]) for i in range(n)] + angles[:, 2] = rnd.uniform(low=-np.pi, high=np.pi, size=(n,)) + lambdas = rnd.uniform(low=0, high=np.pi, size=(5,)) + + e1 = xp.asarray([1.0, 0, 0]) + e2 = xp.asarray([0.0, 1, 0]) + + for lamb in lambdas: + e3 = xp.asarray(Rotation.from_rotvec(lamb*e2).apply(e1)) + ax_lamb = xp.stack([e1, e2, e3], axis=0) + angles[:, 1] = angles_middle - lamb + for order in ['extrinsic', 'intrinsic']: + ax = ax_lamb if order == 'intrinsic' else xp.flip(ax_lamb, axis=0) + rot = Rotation.from_davenport(ax, order, xp.asarray(angles, dtype=dtype)) + with maybe_warn_gimbal_lock(not suppress_warnings, xp): + angles_dav = rot.as_davenport( + ax, + order, + suppress_warnings=suppress_warnings + ) + mat_expected = rot.as_matrix() + rot_estimated = Rotation.from_davenport(ax, order, angles_dav) + mat_estimated = rot_estimated.as_matrix() + xp_assert_close(mat_expected, mat_estimated, atol=atol) + + +@make_xp_test_case((Rotation, "from_euler"), (Rotation, "from_davenport")) +def test_compare_from_davenport_from_euler(xp): + dtype = xpx.default_dtype(xp) + rnd = np.random.RandomState(0) + n = 100 + angles = np.empty((n, 3)) + + # symmetric sequences + rtol = 1e-12 if dtype == xp.float64 else 1e-5 + angles[:, 0] = rnd.uniform(low=-np.pi, high=np.pi, size=(n,)) + angles[:, 1] = rnd.uniform(low=0, high=np.pi, size=(n,)) + angles[:, 2] = rnd.uniform(low=-np.pi, high=np.pi, size=(n,)) + angles = xp.asarray(angles, dtype=dtype) + for order in ['extrinsic', 'intrinsic']: + for seq_tuple in permutations('xyz'): + seq = ''.join([seq_tuple[0], seq_tuple[1], seq_tuple[0]]) + ax = xp.asarray([basis_vec(i) for i in seq], dtype=dtype) + if order == 'intrinsic': + seq = seq.upper() + eul = Rotation.from_euler(seq, angles) + dav = Rotation.from_davenport(ax, order, angles) + xp_assert_close(eul.as_quat(canonical=True), dav.as_quat(canonical=True), + rtol=rtol) + + # asymmetric sequences + angles = xpx.at(angles)[:, 1].subtract(np.pi / 2) + for order in ['extrinsic', 'intrinsic']: + for seq_tuple in permutations('xyz'): + seq = ''.join(seq_tuple) + ax = xp.asarray([basis_vec(i) for i in seq], dtype=dtype) + if order == 'intrinsic': + seq = seq.upper() + eul = Rotation.from_euler(seq, angles) + dav = Rotation.from_davenport(ax, order, angles) + xp_assert_close(eul.as_quat(), dav.as_quat(), rtol=rtol) + + +@make_xp_test_case( + (Rotation, "from_euler"), (Rotation, "as_euler"), (Rotation, "as_davenport") +) +def test_compare_as_davenport_as_euler(xp): + rnd = np.random.RandomState(0) + n = 100 + angles = np.empty((n, 3)) + + # symmetric sequences + angles[:, 0] = rnd.uniform(low=-np.pi, high=np.pi, size=(n,)) + angles[:, 1] = rnd.uniform(low=0, high=np.pi, size=(n,)) + angles[:, 2] = rnd.uniform(low=-np.pi, high=np.pi, size=(n,)) + for order in ['extrinsic', 'intrinsic']: + for seq_tuple in permutations('xyz'): + seq = ''.join([seq_tuple[0], seq_tuple[1], seq_tuple[0]]) + ax = [basis_vec(i) for i in seq] + if order == 'intrinsic': + seq = seq.upper() + rot = Rotation.from_euler(seq, xp.asarray(angles)) + eul = rot.as_euler(seq) + dav = rot.as_davenport(xp.asarray(ax), order) + xp_assert_close(eul, dav, rtol=1e-12) + + # asymmetric sequences + angles[:, 1] -= np.pi / 2 + for order in ['extrinsic', 'intrinsic']: + for seq_tuple in permutations('xyz'): + seq = ''.join(seq_tuple) + ax = [basis_vec(i) for i in seq] + if order == 'intrinsic': + seq = seq.upper() + rot = Rotation.from_euler(seq, xp.asarray(angles)) + eul = rot.as_euler(seq) + dav = rot.as_davenport(xp.asarray(ax), order) + xp_assert_close(eul, dav, rtol=1e-12) + + +@make_xp_test_case( + (Rotation, "from_matrix"), (Rotation, "from_euler"), (Rotation, "from_rotvec"), + (Rotation, "from_davenport"), (Rotation, "from_mrp") +) +def test_zero_rotation_construction(xp): + r = Rotation.random(num=0) + assert len(r) == 0 + + r_ide = Rotation.identity(num=0) + assert len(r_ide) == 0 + + r_get = Rotation.random(num=3)[[]] + assert len(r_get) == 0 + + r_quat = Rotation.from_quat(xp.zeros((0, 4))) + assert len(r_quat) == 0 + + r_matrix = Rotation.from_matrix(xp.zeros((0, 3, 3))) + assert len(r_matrix) == 0 + + r_euler = Rotation.from_euler("xyz", xp.zeros((0, 3))) + assert len(r_euler) == 0 + + r_vec = Rotation.from_rotvec(xp.zeros((0, 3))) + assert len(r_vec) == 0 + + r_dav = Rotation.from_davenport(xp.eye(3), "extrinsic", xp.zeros((0, 3))) + assert len(r_dav) == 0 + + r_mrp = Rotation.from_mrp(xp.zeros((0, 3))) + assert len(r_mrp) == 0 + + +@make_xp_test_case( + (Rotation, "as_matrix"), (Rotation, "as_euler"), (Rotation, "as_rotvec"), + (Rotation, "as_mrp"), (Rotation, "as_davenport") +) +def test_zero_rotation_representation(xp): + r = Rotation.from_quat(xp.zeros((0, 4))) + assert r.as_quat().shape == (0, 4) + assert r.as_matrix().shape == (0, 3, 3) + assert r.as_euler("xyz").shape == (0, 3) + assert r.as_rotvec().shape == (0, 3) + assert r.as_mrp().shape == (0, 3) + assert r.as_davenport(xp.eye(3), "extrinsic").shape == (0, 3) + + +@make_xp_test_case((Rotation, "apply")) +def test_zero_rotation_array_rotation(xp): + r = Rotation.from_quat(xp.zeros((0, 4))) + + v = xp.asarray([1, 2, 3]) + v_rotated = r.apply(v) + assert v_rotated.shape == (0, 3) + + v0 = xp.zeros((0, 3)) + v0_rot = r.apply(v0) + assert v0_rot.shape == (0, 3) + + v2 = xp.ones((2, 3)) + with pytest.raises( + ValueError, match="Cannot broadcast"): + r.apply(v2) + + +@make_xp_test_case((Rotation, "__mul__")) +def test_zero_rotation_multiplication(xp): + r = Rotation.from_quat(xp.zeros((0, 4))) + + r_single = Rotation.from_quat(xp.asarray([0.0, 0, 0, 1])) + r_mult_left = r * r_single + assert len(r_mult_left) == 0 + + r_mult_right = r_single * r + assert len(r_mult_right) == 0 + + r0 = Rotation.from_quat(xp.zeros((0, 4))) + r_mult = r * r0 + assert len(r_mult) == 0 + + r2 = rotation_to_xp(Rotation.random(2), xp) + with pytest.raises(ValueError, match="Cannot broadcast"): + r0 * r2 + + with pytest.raises(ValueError, match="Cannot broadcast"): + r2 * r0 + + +@make_xp_test_case((Rotation, "concatenate")) +def test_zero_rotation_concatentation(xp): + r = Rotation.from_quat(xp.zeros((0, 4))) + + r0 = Rotation.concatenate([r, r]) + assert len(r0) == 0 + + r1 = Rotation.from_quat(xp.asarray([0.0, 0, 0, 1])) + r1 = r.concatenate([r1, r]) + assert len(r1) == 1 + + r3 = rotation_to_xp(Rotation.random(3), xp) + r3 = r.concatenate([r3, r]) + assert len(r3) == 3 + + r4 = rotation_to_xp(Rotation.random(4), xp) + r4 = r.concatenate([r, r4]) + r4 = r.concatenate([r, r4]) + assert len(r4) == 4 + + +@make_xp_test_case((Rotation, "__pow__")) +def test_zero_rotation_power(xp): + r = Rotation.from_quat(xp.zeros((0, 4))) + for pp in [-1.5, -1, 0, 1, 1.5]: + pow0 = r**pp + assert len(pow0) == 0 + + +@make_xp_test_case((Rotation, "inv")) +def test_zero_rotation_inverse(xp): + r = Rotation.from_quat(xp.zeros((0, 4))) + r_inv = r.inv() + assert len(r_inv) == 0 + + +@make_xp_test_case((Rotation, "magnitude")) +def test_zero_rotation_magnitude(xp): + r = Rotation.from_quat(xp.zeros((0, 4))) + magnitude = r.magnitude() + assert magnitude.shape == (0,) + + +@make_xp_test_case((Rotation, "mean")) +def test_zero_rotation_mean(xp): + r = Rotation.from_quat(xp.zeros((0, 4))) + with pytest.raises(ValueError, match="Mean of an empty rotation set is undefined."): + r.mean() + + +@make_xp_test_case((Rotation, "approx_equal")) +def test_zero_rotation_approx_equal(xp): + r = Rotation.from_quat(xp.zeros((0, 4))) + r0 = Rotation.from_quat(xp.zeros((0, 4))) + assert r.approx_equal(r0).shape == (0,) + r1 = Rotation.from_quat(xp.asarray([0.0, 0, 0, 1])) + assert r.approx_equal(r1).shape == (0,) + r2 = rotation_to_xp(Rotation.random(), xp) + assert r2.approx_equal(r).shape == (0,) + + approx_msg = "Expected broadcastable shapes in both rotations" + r3 = rotation_to_xp(Rotation.random(2), xp) + with pytest.raises(ValueError, match=approx_msg): + r.approx_equal(r3) + + with pytest.raises(ValueError, match=approx_msg): + r3.approx_equal(r) + + +@pytest.mark.skip_xp_backends("jax.numpy", + reason="JAX out-of-bounds indexing deviates from numpy") +def test_zero_rotation_get_set(xp): + r = Rotation.from_quat(xp.zeros((0, 4))) + + r_get = r[xp.asarray([], dtype=xp.bool)] + assert len(r_get) == 0 + + r_slice = r[:0] + assert len(r_slice) == 0 + + with pytest.raises(IndexError): + r[xp.asarray([0])] + + with pytest.raises(IndexError): + r[xp.asarray([True])] + + with pytest.raises(IndexError): + r[0] = Rotation.from_quat(xp.asarray([0, 0, 0, 1])) + + +@make_xp_test_case((Rotation, "__getitem__")) +def test_boolean_indexes(xp): + r = rotation_to_xp(Rotation.random(3), xp) + + r0 = r[xp.asarray([False, False, False])] + assert len(r0) == 0 + + r1 = r[xp.asarray([False, True, False])] + assert len(r1) == 1 + + r3 = r[xp.asarray([True, True, True])] + assert len(r3) == 3 + + # Multiple dimensions + r = Rotation.from_quat(xp.ones((3, 2, 4))) + r4 = r[xp.asarray([True, False, False])] + assert len(r4) == 1 + assert r4.as_quat().shape == (1, 2, 4) + + with pytest.raises(IndexError): + r[xp.asarray([True, True])] + + +@make_xp_test_case((Rotation, "__iter__")) +def test_rotation_iter(xp): + r = rotation_to_xp(Rotation.random(3), xp) + for i, r_i in enumerate(r): + assert isinstance(r_i, Rotation) + xp_assert_equal(r_i.as_quat(), r[i].as_quat()) + if i > len(r): + raise RuntimeError("Iteration exceeded length of rotations") + + +@pytest.mark.parametrize("ndim", list(range(1, 5))) +def test_rotation_shape(xp, ndim: int): + shape = tuple(range(2, 2 + ndim)[:ndim - 1]) + quat = xp.ones(shape + (4,)) + r = Rotation.from_quat(quat) + assert r.shape == shape, f"Got {r.shape}, expected {shape}" + + +def test_non_writeable(): + q = np.array([0, 0, 0, 1.0]) + q.flags.writeable = False + Rotation.from_quat(q) # Regression test against gh-24354, should not raise + + +def test_gh_24751(): + points = np.array([[0.0, 0.0, 0.0]]) + points.flags.writeable = False + rot = Rotation.from_euler("z", 0.0) + rot.apply(points) + + +def test_rotvec_non_writeable(): + rotvec = np.array([0, 0, 1]) * np.pi / 2 + rotvec.flags.writeable = False + Rotation.from_rotvec(rotvec) diff --git a/v12x/pylib/scipy/spatial/transform/tests/test_rotation_groups.py b/v12x/pylib/scipy/spatial/transform/tests/test_rotation_groups.py new file mode 100644 index 0000000000000000000000000000000000000000..3f8a53b7abfede9b241fed2190e0f1244d76a9ee --- /dev/null +++ b/v12x/pylib/scipy/spatial/transform/tests/test_rotation_groups.py @@ -0,0 +1,169 @@ +import pytest + +import numpy as np +from numpy.testing import assert_array_almost_equal +from scipy.spatial.transform import Rotation +from scipy.optimize import linear_sum_assignment +from scipy.spatial.distance import cdist +from scipy.constants import golden as phi +from scipy.spatial import cKDTree + + +TOL = 1E-12 +NS = list(range(1, 13)) +NAMES = ["I", "O", "T"] + [f"C{n}" for n in NS] + [f"D{n}" for n in NS] +SIZES = [60, 24, 12] + list(NS) + [2 * n for n in NS] + + +def _calculate_rmsd(P, Q): + """Calculates the root-mean-square distance between the points of P and Q. + The distance is taken as the minimum over all possible matchings. It is + zero if P and Q are identical and non-zero if not. + """ + distance_matrix = cdist(P, Q, metric='sqeuclidean') + matching = linear_sum_assignment(distance_matrix) + return np.sqrt(distance_matrix[matching].sum()) + + +def _generate_pyramid(n, axis): + thetas = np.linspace(0, 2 * np.pi, n + 1)[:-1] + P = np.vstack([np.zeros(n), np.cos(thetas), np.sin(thetas)]).T + P = np.concatenate((P, [[1, 0, 0]])) + return np.roll(P, axis, axis=1) + + +def _generate_prism(n, axis): + thetas = np.linspace(0, 2 * np.pi, n + 1)[:-1] + bottom = np.vstack([-np.ones(n), np.cos(thetas), np.sin(thetas)]).T + top = np.vstack([+np.ones(n), np.cos(thetas), np.sin(thetas)]).T + P = np.concatenate((bottom, top)) + return np.roll(P, axis, axis=1) + + +def _generate_icosahedron(): + x = np.array([[0, -1, -phi], + [0, -1, +phi], + [0, +1, -phi], + [0, +1, +phi]]) + return np.concatenate([np.roll(x, i, axis=1) for i in range(3)]) + + +def _generate_octahedron(): + return np.array([[-1, 0, 0], [+1, 0, 0], [0, -1, 0], + [0, +1, 0], [0, 0, -1], [0, 0, +1]]) + + +def _generate_tetrahedron(): + return np.array([[1, 1, 1], [1, -1, -1], [-1, 1, -1], [-1, -1, 1]]) + + +@pytest.mark.parametrize("name", [-1, None, True, np.array(['C3'])]) +def test_group_type(name): + with pytest.raises(ValueError, + match="must be a string"): + Rotation.create_group(name) + + +@pytest.mark.parametrize("name", ["Q", " ", "CA", "C ", "DA", "D ", "I2", ""]) +def test_group_name(name): + with pytest.raises(ValueError, + match="must be one of 'I', 'O', 'T', 'Dn', 'Cn'"): + Rotation.create_group(name) + + +@pytest.mark.parametrize("name", ["C0", "D0"]) +def test_group_order_positive(name): + with pytest.raises(ValueError, + match="Group order must be positive"): + Rotation.create_group(name) + + +@pytest.mark.parametrize("axis", ['A', 'b', 0, 1, 2, 4, False, None]) +def test_axis_valid(axis): + with pytest.raises(ValueError, + match="`axis` must be one of"): + Rotation.create_group("C1", axis) + + +def test_icosahedral(): + """The icosahedral group fixes the rotations of an icosahedron. Here we + test that the icosahedron is invariant after application of the elements + of the rotation group.""" + P = _generate_icosahedron() + for g in Rotation.create_group("I"): + g = Rotation.from_quat(g.as_quat()) + assert _calculate_rmsd(P, g.apply(P)) < TOL + + +def test_octahedral(): + """Test that the octahedral group correctly fixes the rotations of an + octahedron.""" + P = _generate_octahedron() + for g in Rotation.create_group("O"): + assert _calculate_rmsd(P, g.apply(P)) < TOL + + +def test_tetrahedral(): + """Test that the tetrahedral group correctly fixes the rotations of a + tetrahedron.""" + P = _generate_tetrahedron() + for g in Rotation.create_group("T"): + assert _calculate_rmsd(P, g.apply(P)) < TOL + + +@pytest.mark.parametrize("n", NS) +@pytest.mark.parametrize("axis", 'XYZ') +def test_dicyclic(n, axis): + """Test that the dicyclic group correctly fixes the rotations of a + prism.""" + P = _generate_prism(n, axis='XYZ'.index(axis)) + for g in Rotation.create_group(f"D{n}", axis=axis): + assert _calculate_rmsd(P, g.apply(P)) < TOL + + +@pytest.mark.parametrize("n", NS) +@pytest.mark.parametrize("axis", 'XYZ') +def test_cyclic(n, axis): + """Test that the cyclic group correctly fixes the rotations of a + pyramid.""" + P = _generate_pyramid(n, axis='XYZ'.index(axis)) + for g in Rotation.create_group(f"C{n}", axis=axis): + assert _calculate_rmsd(P, g.apply(P)) < TOL + + +@pytest.mark.parametrize("name, size", list(zip(NAMES, SIZES))) +def test_group_sizes(name, size): + assert len(Rotation.create_group(name)) == size + + +@pytest.mark.parametrize("name, size", list(zip(NAMES, SIZES))) +def test_group_no_duplicates(name, size): + g = Rotation.create_group(name) + kdtree = cKDTree(g.as_quat()) + assert len(kdtree.query_pairs(1E-3)) == 0 + + +@pytest.mark.parametrize("name, size", list(zip(NAMES, SIZES))) +def test_group_symmetry(name, size): + g = Rotation.create_group(name) + q = np.concatenate((-g.as_quat(), g.as_quat())) + distance = np.sort(cdist(q, q)) + deltas = np.max(distance, axis=0) - np.min(distance, axis=0) + assert (deltas < TOL).all() + + +@pytest.mark.parametrize("name", NAMES) +def test_reduction(name): + """Test that the elements of the rotation group are correctly + mapped onto the identity rotation.""" + g = Rotation.create_group(name) + f = g.reduce(g) + assert_array_almost_equal(f.magnitude(), np.zeros(len(g))) + + +@pytest.mark.parametrize("name", NAMES) +def test_single_reduction(name): + g = Rotation.create_group(name) + f = g[-1].reduce(g) + assert_array_almost_equal(f.magnitude(), 0) + assert f.as_quat().shape == (4,) diff --git a/v12x/pylib/scipy/spatial/transform/tests/test_rotation_spline.py b/v12x/pylib/scipy/spatial/transform/tests/test_rotation_spline.py new file mode 100644 index 0000000000000000000000000000000000000000..d268659e06416506243593bc8a4c2e8febf3f98e --- /dev/null +++ b/v12x/pylib/scipy/spatial/transform/tests/test_rotation_spline.py @@ -0,0 +1,188 @@ +from itertools import product +import numpy as np +from numpy.testing import assert_allclose +import pytest +from scipy.spatial.transform import Rotation, RotationSpline +from scipy.spatial.transform._rotation_spline import ( + _angular_rate_to_rotvec_dot_matrix, + _rotvec_dot_to_angular_rate_matrix, + _matrix_vector_product_of_stacks, + _angular_acceleration_nonlinear_term, + _create_block_3_diagonal_matrix) + + +pytestmark = pytest.mark.skip_xp_backends(np_only=True) + + +def test_angular_rate_to_rotvec_conversions(): + np.random.seed(0) + rv = np.random.randn(4, 3) + A = _angular_rate_to_rotvec_dot_matrix(rv) + A_inv = _rotvec_dot_to_angular_rate_matrix(rv) + + # When the rotation vector is aligned with the angular rate, then + # the rotation vector rate and angular rate are the same. + assert_allclose(_matrix_vector_product_of_stacks(A, rv), rv) + assert_allclose(_matrix_vector_product_of_stacks(A_inv, rv), rv) + + # A and A_inv must be reciprocal to each other. + I_stack = np.empty((4, 3, 3)) + I_stack[:] = np.eye(3) + assert_allclose(np.matmul(A, A_inv), I_stack, atol=1e-15) + + +def test_angular_rate_nonlinear_term(): + # The only simple test is to check that the term is zero when + # the rotation vector + np.random.seed(0) + rv = np.random.rand(4, 3) + assert_allclose(_angular_acceleration_nonlinear_term(rv, rv), 0, + atol=1e-19) + + +def test_create_block_3_diagonal_matrix(): + np.random.seed(0) + A = np.empty((4, 3, 3)) + A[:] = np.arange(1, 5)[:, None, None] + + B = np.empty((4, 3, 3)) + B[:] = -np.arange(1, 5)[:, None, None] + d = 10 * np.arange(10, 15) + + banded = _create_block_3_diagonal_matrix(A, B, d) + + # Convert the banded matrix to the full matrix. + k, l = list(zip(*product(np.arange(banded.shape[0]), + np.arange(banded.shape[1])))) + k = np.asarray(k) + l = np.asarray(l) + + i = k - 5 + l + j = l + values = banded.ravel() + mask = (i >= 0) & (i < 15) + i = i[mask] + j = j[mask] + values = values[mask] + full = np.zeros((15, 15)) + full[i, j] = values + + zero = np.zeros((3, 3)) + eye = np.eye(3) + + # Create the reference full matrix in the most straightforward manner. + ref = np.block([ + [d[0] * eye, B[0], zero, zero, zero], + [A[0], d[1] * eye, B[1], zero, zero], + [zero, A[1], d[2] * eye, B[2], zero], + [zero, zero, A[2], d[3] * eye, B[3]], + [zero, zero, zero, A[3], d[4] * eye], + ]) + + assert_allclose(full, ref, atol=1e-19) + + +def test_spline_2_rotations(): + times = [0, 10] + rotations = Rotation.from_euler('xyz', [[0, 0, 0], [10, -20, 30]], + degrees=True) + spline = RotationSpline(times, rotations) + + rv = (rotations[0].inv() * rotations[1]).as_rotvec() + rate = rv / (times[1] - times[0]) + times_check = np.array([-1, 5, 12]) + dt = times_check - times[0] + rv_ref = rate * dt[:, None] + + assert_allclose(spline(times_check).as_rotvec(), rv_ref) + assert_allclose(spline(times_check, 1), np.resize(rate, (3, 3))) + assert_allclose(spline(times_check, 2), 0, atol=1e-16) + + +def test_constant_attitude(): + times = np.arange(10) + rotations = Rotation.from_rotvec(np.ones((10, 3))) + spline = RotationSpline(times, rotations) + + times_check = np.linspace(-1, 11) + assert_allclose(spline(times_check).as_rotvec(), 1, rtol=1e-15) + assert_allclose(spline(times_check, 1), 0, atol=1e-17) + assert_allclose(spline(times_check, 2), 0, atol=1e-17) + + assert_allclose(spline(5.5).as_rotvec(), 1, rtol=1e-15) + assert_allclose(spline(5.5, 1), 0, atol=1e-17) + assert_allclose(spline(5.5, 2), 0, atol=1e-17) + + +def test_spline_properties(): + times = np.array([0, 5, 15, 27]) + angles = [[-5, 10, 27], [3, 5, 38], [-12, 10, 25], [-15, 20, 11]] + + rotations = Rotation.from_euler('xyz', angles, degrees=True) + spline = RotationSpline(times, rotations) + + assert_allclose(spline(times).as_euler('xyz', degrees=True), angles) + assert_allclose(spline(0).as_euler('xyz', degrees=True), angles[0]) + + h = 1e-8 + rv0 = spline(times).as_rotvec() + rvm = spline(times - h).as_rotvec() + rvp = spline(times + h).as_rotvec() + # rtol bumped from 1e-15 to 1.5e-15 in gh18414 for linux 32 bit + assert_allclose(rv0, 0.5 * (rvp + rvm), rtol=1.5e-15) + + r0 = spline(times, 1) + rm = spline(times - h, 1) + rp = spline(times + h, 1) + assert_allclose(r0, 0.5 * (rm + rp), rtol=1e-14) + + a0 = spline(times, 2) + am = spline(times - h, 2) + ap = spline(times + h, 2) + assert_allclose(a0, am, rtol=1e-7) + assert_allclose(a0, ap, rtol=1e-7) + + +def test_error_handling(): + with pytest.raises(ValueError): + RotationSpline([1.0], Rotation.random()) + + r = Rotation.random(10) + t = np.arange(10).reshape(5, 2) + with pytest.raises(ValueError): + RotationSpline(t, r) + + t = np.arange(9) + with pytest.raises(ValueError): + RotationSpline(t, r) + + t = np.arange(10) + t[5] = 0 + with pytest.raises(ValueError): + RotationSpline(t, r) + + t = np.arange(10) + + s = RotationSpline(t, r) + with pytest.raises(ValueError): + s(10, -1) + + with pytest.raises(ValueError): + s(np.arange(10).reshape(5, 2)) + + r = Rotation.from_quat(np.array([[[0.0, 0, 0, 1], [1, 0, 0 ,0]]])) + t = np.arange(2) + with pytest.raises(ValueError): + RotationSpline(t, r) # More than 2 dimensions + + +@pytest.mark.skip_xp_backends("numpy") +def test_xp_errors(xp): + times = xp.asarray([0, 10]) + r = Rotation.random(2) + r = Rotation.from_quat(xp.asarray(r.as_quat())) + s = RotationSpline(times, r) + t = xp.asarray([0.5, 1.5]) + # RotationSpline does not have native Array API support, so we check that it + # converts any array to NumPy and outputs NumPy arrays. + assert isinstance(s(t).as_quat(), np.ndarray) diff --git a/v12x/pylib/scipy/special/__init__.py b/v12x/pylib/scipy/special/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9bcbfd132896a63ccf689c54ab9c36fae7db695a --- /dev/null +++ b/v12x/pylib/scipy/special/__init__.py @@ -0,0 +1,842 @@ +""" +======================================== +Special functions (:mod:`scipy.special`) +======================================== + +.. currentmodule:: scipy.special + +.. toctree:: + :hidden: + + special.cython_special + +Almost all of the functions below accept NumPy arrays as input +arguments as well as single numbers. This means they follow +broadcasting and automatic array-looping rules. Technically, +they are `NumPy universal functions +`_. +Functions which do not accept NumPy arrays are marked by a warning +in the section description. + +.. seealso:: + + `scipy.special.cython_special` -- Typed Cython versions of special functions + + +Error handling +============== + +Errors are handled by returning NaNs or other appropriate values. +Some of the special function routines can emit warnings or raise +exceptions when an error occurs. By default this is disabled, except +for memory allocation errors, which result in an exception being raised. +To query and control the current error handling state the following +functions are provided. + +.. autosummary:: + :toctree: generated/ + + geterr -- Get the current way of handling special-function errors. + seterr -- Set how special-function errors are handled. + errstate -- Context manager for special-function error handling. + SpecialFunctionWarning -- Warning that can be emitted by special functions. + SpecialFunctionError -- Exception that can be raised by special functions. + +Available functions +=================== + +Airy functions +-------------- + +.. autosummary:: + :toctree: generated/ + + airy -- Airy functions and their derivatives. + airye -- Exponentially scaled Airy functions and their derivatives. + ai_zeros -- Compute `nt` zeros and values of the Airy function Ai and its derivative. + bi_zeros -- Compute `nt` zeros and values of the Airy function Bi and its derivative. + itairy -- Integrals of Airy functions + + +Elliptic functions and integrals +-------------------------------- + +.. autosummary:: + :toctree: generated/ + + ellipj -- Jacobian elliptic functions. + ellipk -- Complete elliptic integral of the first kind. + ellipkm1 -- Complete elliptic integral of the first kind around `m` = 1. + ellipkinc -- Incomplete elliptic integral of the first kind. + ellipe -- Complete elliptic integral of the second kind. + ellipeinc -- Incomplete elliptic integral of the second kind. + elliprc -- Degenerate symmetric integral RC. + elliprd -- Symmetric elliptic integral of the second kind. + elliprf -- Completely-symmetric elliptic integral of the first kind. + elliprg -- Completely-symmetric elliptic integral of the second kind. + elliprj -- Symmetric elliptic integral of the third kind. + +Bessel functions +---------------- + +.. autosummary:: + :toctree: generated/ + + jv -- Bessel function of the first kind of real order and \ + complex argument. + jve -- Exponentially scaled Bessel function of order `v`. + yn -- Bessel function of the second kind of integer order and \ + real argument. + yv -- Bessel function of the second kind of real order and \ + complex argument. + yve -- Exponentially scaled Bessel function of the second kind \ + of real order. + iv -- Modified Bessel function of the first kind of real order. + ive -- Exponentially scaled modified Bessel function of the \ + first kind. + kn -- Modified Bessel function of the second kind of integer \ + order `n` + kv -- Modified Bessel function of the second kind of real order \ + `v` + kve -- Exponentially scaled modified Bessel function of the \ + second kind. + hankel1 -- Hankel function of the first kind. + hankel1e -- Exponentially scaled Hankel function of the first kind. + hankel2 -- Hankel function of the second kind. + hankel2e -- Exponentially scaled Hankel function of the second kind. + wright_bessel -- Wright's generalized Bessel function. + log_wright_bessel -- Logarithm of Wright's generalized Bessel function. + +The following function does not accept NumPy arrays (it is not a +universal function): + +.. autosummary:: + :toctree: generated/ + + lmbda -- Jahnke-Emden Lambda function, Lambdav(x). + +Zeros of Bessel functions +^^^^^^^^^^^^^^^^^^^^^^^^^ + +The following functions do not accept NumPy arrays (they are not +universal functions): + +.. autosummary:: + :toctree: generated/ + + jnjnp_zeros -- Compute zeros of integer-order Bessel functions Jn and Jn'. + jnyn_zeros -- Compute nt zeros of Bessel functions Jn(x), Jn'(x), Yn(x), and Yn'(x). + jn_zeros -- Compute zeros of integer-order Bessel function Jn(x). + jnp_zeros -- Compute zeros of integer-order Bessel function derivative Jn'(x). + yn_zeros -- Compute zeros of integer-order Bessel function Yn(x). + ynp_zeros -- Compute zeros of integer-order Bessel function derivative Yn'(x). + y0_zeros -- Compute nt zeros of Bessel function Y0(z), and derivative at each zero. + y1_zeros -- Compute nt zeros of Bessel function Y1(z), and derivative at each zero. + y1p_zeros -- Compute nt zeros of Bessel derivative Y1'(z), and value at each zero. + +Faster versions of common Bessel functions +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +.. autosummary:: + :toctree: generated/ + + j0 -- Bessel function of the first kind of order 0. + j1 -- Bessel function of the first kind of order 1. + y0 -- Bessel function of the second kind of order 0. + y1 -- Bessel function of the second kind of order 1. + i0 -- Modified Bessel function of order 0. + i0e -- Exponentially scaled modified Bessel function of order 0. + i1 -- Modified Bessel function of order 1. + i1e -- Exponentially scaled modified Bessel function of order 1. + k0 -- Modified Bessel function of the second kind of order 0, :math:`K_0`. + k0e -- Exponentially scaled modified Bessel function K of order 0 + k1 -- Modified Bessel function of the second kind of order 1, :math:`K_1(x)`. + k1e -- Exponentially scaled modified Bessel function K of order 1. + +Integrals of Bessel functions +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +.. autosummary:: + :toctree: generated/ + + itj0y0 -- Integrals of Bessel functions of order 0. + it2j0y0 -- Integrals related to Bessel functions of order 0. + iti0k0 -- Integrals of modified Bessel functions of order 0. + it2i0k0 -- Integrals related to modified Bessel functions of order 0. + besselpoly -- Weighted integral of a Bessel function. + +Derivatives of Bessel functions +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +.. autosummary:: + :toctree: generated/ + + jvp -- Compute nth derivative of Bessel function Jv(z) with respect to `z`. + yvp -- Compute nth derivative of Bessel function Yv(z) with respect to `z`. + ivp -- Compute nth derivative of modified Bessel function Iv(z) with respect to `z`. + kvp -- Compute nth derivative of real-order modified Bessel function Kv(z) + h1vp -- Compute nth derivative of Hankel function H1v(z) with respect to `z`. + h2vp -- Compute nth derivative of Hankel function H2v(z) with respect to `z`. + +Spherical Bessel functions +^^^^^^^^^^^^^^^^^^^^^^^^^^ + +.. autosummary:: + :toctree: generated/ + + spherical_jn -- Spherical Bessel function of the first kind or its derivative. + spherical_yn -- Spherical Bessel function of the second kind or its derivative. + spherical_in -- Modified spherical Bessel function of the first kind or its derivative. + spherical_kn -- Modified spherical Bessel function of the second kind or its derivative. + +Riccati-Bessel functions +^^^^^^^^^^^^^^^^^^^^^^^^ + +The following functions do not accept NumPy arrays (they are not +universal functions): + +.. autosummary:: + :toctree: generated/ + + riccati_jn -- Compute Riccati-Bessel function of the first kind and its derivative. + riccati_yn -- Compute Riccati-Bessel function of the second kind and its derivative. + +Struve functions +---------------- + +.. autosummary:: + :toctree: generated/ + + struve -- Struve function. + modstruve -- Modified Struve function. + itstruve0 -- Integral of the Struve function of order 0. + it2struve0 -- Integral related to the Struve function of order 0. + itmodstruve0 -- Integral of the modified Struve function of order 0. + + +Raw statistical functions +------------------------- + +.. seealso:: :mod:`scipy.stats`: Friendly versions of these functions. + +Binomial distribution +^^^^^^^^^^^^^^^^^^^^^ + +.. autosummary:: + :toctree: generated/ + + bdtr -- Binomial distribution cumulative distribution function. + bdtrc -- Binomial distribution survival function. + bdtri -- Inverse function to `bdtr` with respect to `p`. + bdtrik -- Inverse function to `bdtr` with respect to `k`. + bdtrin -- Inverse function to `bdtr` with respect to `n`. + +Beta distribution +^^^^^^^^^^^^^^^^^ + +.. autosummary:: + :toctree: generated/ + + btdtria -- Inverse of `betainc` with respect to `a`. + btdtrib -- Inverse of `betainc` with respect to `b`. + +F distribution +^^^^^^^^^^^^^^ + +.. autosummary:: + :toctree: generated/ + + fdtr -- F cumulative distribution function. + fdtrc -- F survival function. + fdtri -- The `p`-th quantile of the F-distribution. + fdtridfd -- Inverse to `fdtr` vs dfd. + +Gamma distribution +^^^^^^^^^^^^^^^^^^ + +.. autosummary:: + :toctree: generated/ + + gdtr -- Gamma distribution cumulative distribution function. + gdtrc -- Gamma distribution survival function. + gdtria -- Inverse of `gdtr` vs a. + gdtrib -- Inverse of `gdtr` vs b. + gdtrix -- Inverse of `gdtr` vs x. + +Negative binomial distribution +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +.. autosummary:: + :toctree: generated/ + + nbdtr -- Negative binomial cumulative distribution function. + nbdtrc -- Negative binomial survival function. + nbdtri -- Inverse of `nbdtr` vs `p`. + nbdtrik -- Inverse of `nbdtr` vs `k`. + nbdtrin -- Inverse of `nbdtr` vs `n`. + +Noncentral F distribution +^^^^^^^^^^^^^^^^^^^^^^^^^ + +.. autosummary:: + :toctree: generated/ + + ncfdtr -- Cumulative distribution function of the non-central F distribution. + ncfdtridfd -- Calculate degrees of freedom (denominator) for the noncentral F-distribution. + ncfdtridfn -- Calculate degrees of freedom (numerator) for the noncentral F-distribution. + ncfdtri -- Inverse cumulative distribution function of the non-central F distribution. + ncfdtrinc -- Calculate non-centrality parameter for non-central F distribution. + +Noncentral t distribution +^^^^^^^^^^^^^^^^^^^^^^^^^ + +.. autosummary:: + :toctree: generated/ + + nctdtr -- Cumulative distribution function of the non-central `t` distribution. + nctdtridf -- Calculate degrees of freedom for non-central t distribution. + nctdtrit -- Inverse cumulative distribution function of the non-central t distribution. + nctdtrinc -- Calculate non-centrality parameter for non-central t distribution. + +Normal distribution +^^^^^^^^^^^^^^^^^^^ + +.. autosummary:: + :toctree: generated/ + + nrdtrimn -- Calculate mean of normal distribution given other params. + nrdtrisd -- Calculate standard deviation of normal distribution given other params. + ndtr -- Normal cumulative distribution function. + log_ndtr -- Logarithm of normal cumulative distribution function. + ndtri -- Inverse of `ndtr` vs x. + ndtri_exp -- Inverse of `log_ndtr` vs x. + +Poisson distribution +^^^^^^^^^^^^^^^^^^^^ + +.. autosummary:: + :toctree: generated/ + + pdtr -- Poisson cumulative distribution function. + pdtrc -- Poisson survival function. + pdtri -- Inverse to `pdtr` vs m. + pdtrik -- Inverse to `pdtr` vs k. + +Student t distribution +^^^^^^^^^^^^^^^^^^^^^^ + +.. autosummary:: + :toctree: generated/ + + stdtr -- Student t distribution cumulative distribution function. + stdtridf -- Inverse of `stdtr` vs df. + stdtrit -- Inverse of `stdtr` vs `t`. + +Chi square distribution +^^^^^^^^^^^^^^^^^^^^^^^ + +.. autosummary:: + :toctree: generated/ + + chdtr -- Chi square cumulative distribution function. + chdtrc -- Chi square survival function. + chdtri -- Inverse to `chdtrc`. + chdtriv -- Inverse to `chdtr` vs `v`. + +Non-central chi square distribution +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +.. autosummary:: + :toctree: generated/ + + chndtr -- Non-central chi square cumulative distribution function. + chndtridf -- Inverse to `chndtr` vs `df`. + chndtrinc -- Inverse to `chndtr` vs `nc`. + chndtrix -- Inverse to `chndtr` vs `x`. + +Kolmogorov distribution +^^^^^^^^^^^^^^^^^^^^^^^ + +.. autosummary:: + :toctree: generated/ + + smirnov -- Kolmogorov-Smirnov complementary cumulative distribution function. + smirnovi -- Inverse to `smirnov`. + kolmogorov -- Complementary cumulative distribution function of Kolmogorov distribution. + kolmogi -- Inverse function to `kolmogorov`. + +Box-Cox transformation +^^^^^^^^^^^^^^^^^^^^^^ + +.. autosummary:: + :toctree: generated/ + + boxcox -- Compute the Box-Cox transformation. + boxcox1p -- Compute the Box-Cox transformation of 1 + `x`. + inv_boxcox -- Compute the inverse of the Box-Cox transformation. + inv_boxcox1p -- Compute the inverse of the Box-Cox transformation. + + +Sigmoidal functions +^^^^^^^^^^^^^^^^^^^ + +.. autosummary:: + :toctree: generated/ + + logit -- Logit ufunc for ndarrays. + expit -- Logistic sigmoid function. + log_expit -- Logarithm of the logistic sigmoid function. + +Miscellaneous +^^^^^^^^^^^^^ + +.. autosummary:: + :toctree: generated/ + + tklmbda -- Tukey-Lambda cumulative distribution function. + owens_t -- Owen's T Function. + + +Information Theory functions +---------------------------- + +.. autosummary:: + :toctree: generated/ + + entr -- Elementwise function for computing entropy. + rel_entr -- Elementwise function for computing relative entropy. + kl_div -- Elementwise function for computing Kullback-Leibler divergence. + huber -- Huber loss function. + pseudo_huber -- Pseudo-Huber loss function. + + +Gamma and related functions +--------------------------- + +.. autosummary:: + :toctree: generated/ + + gamma -- Gamma function. + gammaln -- Logarithm of the absolute value of the Gamma function for real inputs. + loggamma -- Principal branch of the logarithm of the Gamma function. + gammasgn -- Sign of the gamma function. + gammainc -- Regularized lower incomplete gamma function. + gammaincinv -- Inverse to `gammainc`. + gammaincc -- Regularized upper incomplete gamma function. + gammainccinv -- Inverse to `gammaincc`. + beta -- Beta function. + betaln -- Natural logarithm of absolute value of beta function. + betainc -- Incomplete beta integral. + betaincc -- Complemented incomplete beta integral. + betaincinv -- Inverse function to beta integral. + betainccinv -- Inverse of the complemented incomplete beta integral. + psi -- The digamma function. + rgamma -- Gamma function inverted. + polygamma -- Polygamma function n. + multigammaln -- Returns the log of multivariate gamma, also sometimes called the generalized gamma. + digamma -- psi(x[, out]). + poch -- Rising factorial (z)_m. + +Error function and Fresnel integrals +------------------------------------ + +.. autosummary:: + :toctree: generated/ + + erf -- Returns the error function of complex argument. + erfc -- Complementary error function, ``1 - erf(x)``. + erfcx -- Scaled complementary error function, ``exp(x**2) * erfc(x)``. + erfi -- Imaginary error function, ``-i erf(i z)``. + erfinv -- Inverse function for erf. + erfcinv -- Inverse function for erfc. + wofz -- Faddeeva function. + dawsn -- Dawson's integral. + fresnel -- Fresnel sin and cos integrals. + fresnel_zeros -- Compute nt complex zeros of sine and cosine Fresnel integrals S(z) and C(z). + modfresnelp -- Modified Fresnel positive integrals. + modfresnelm -- Modified Fresnel negative integrals. + voigt_profile -- Voigt profile. + +The following functions do not accept NumPy arrays (they are not +universal functions): + +.. autosummary:: + :toctree: generated/ + + erf_zeros -- Compute nt complex zeros of error function erf(z). + fresnelc_zeros -- Compute nt complex zeros of cosine Fresnel integral C(z). + fresnels_zeros -- Compute nt complex zeros of sine Fresnel integral S(z). + +Legendre functions +------------------ + +.. autosummary:: + :toctree: generated/ + + legendre_p -- Legendre polynomials of the first kind. + legendre_p_all -- All Legendre polynomials of the first kind up to a specified order. + assoc_legendre_p -- Associated Legendre polynomials of the first kind. + assoc_legendre_p_all -- All associated Legendre polynomials of the first kind up to a specified order and degree. + sph_legendre_p -- Spherical Legendre polynomials of the first kind. + sph_legendre_p_all -- All spherical Legendre polynomials of the first kind up to a specified order and degree. + sph_harm_y -- Spherical harmonics. + sph_harm_y_all -- All spherical harmonics up to a specified order and degree. + +The following functions are in the process of being deprecated in favor of the above, +which provide a more flexible and consistent interface. + +.. autosummary:: + :toctree: generated/ + + lpmv -- Associated Legendre function of integer order and real degree. + lqn -- Legendre function of the second kind. + lqmn -- Sequence of associated Legendre functions of the second kind. + +Ellipsoidal harmonics +--------------------- + +.. autosummary:: + :toctree: generated/ + + ellip_harm -- Ellipsoidal harmonic functions E^p_n(l). + ellip_harm_2 -- Ellipsoidal harmonic functions F^p_n(l). + ellip_normal -- Ellipsoidal harmonic normalization constants gamma^p_n. + +Orthogonal polynomials +---------------------- + +The following functions evaluate values of orthogonal polynomials: + +.. autosummary:: + :toctree: generated/ + + assoc_laguerre -- Compute the generalized (associated) Laguerre polynomial of degree n and order k. + eval_legendre -- Evaluate Legendre polynomial at a point. + eval_chebyt -- Evaluate Chebyshev polynomial of the first kind at a point. + eval_chebyu -- Evaluate Chebyshev polynomial of the second kind at a point. + eval_chebyc -- Evaluate Chebyshev polynomial of the first kind on [-2, 2] at a point. + eval_chebys -- Evaluate Chebyshev polynomial of the second kind on [-2, 2] at a point. + eval_jacobi -- Evaluate Jacobi polynomial at a point. + eval_laguerre -- Evaluate Laguerre polynomial at a point. + eval_genlaguerre -- Evaluate generalized Laguerre polynomial at a point. + eval_hermite -- Evaluate physicist's Hermite polynomial at a point. + eval_hermitenorm -- Evaluate probabilist's (normalized) Hermite polynomial at a point. + eval_gegenbauer -- Evaluate Gegenbauer polynomial at a point. + eval_sh_legendre -- Evaluate shifted Legendre polynomial at a point. + eval_sh_chebyt -- Evaluate shifted Chebyshev polynomial of the first kind at a point. + eval_sh_chebyu -- Evaluate shifted Chebyshev polynomial of the second kind at a point. + eval_sh_jacobi -- Evaluate shifted Jacobi polynomial at a point. + +The following functions compute roots and quadrature weights for +orthogonal polynomials: + +.. autosummary:: + :toctree: generated/ + + roots_legendre -- Gauss-Legendre quadrature. + roots_chebyt -- Gauss-Chebyshev (first kind) quadrature. + roots_chebyu -- Gauss-Chebyshev (second kind) quadrature. + roots_chebyc -- Gauss-Chebyshev (first kind) quadrature. + roots_chebys -- Gauss-Chebyshev (second kind) quadrature. + roots_jacobi -- Gauss-Jacobi quadrature. + roots_laguerre -- Gauss-Laguerre quadrature. + roots_genlaguerre -- Gauss-generalized Laguerre quadrature. + roots_hermite -- Gauss-Hermite (physicist's) quadrature. + roots_hermitenorm -- Gauss-Hermite (statistician's) quadrature. + roots_gegenbauer -- Gauss-Gegenbauer quadrature. + roots_sh_legendre -- Gauss-Legendre (shifted) quadrature. + roots_sh_chebyt -- Gauss-Chebyshev (first kind, shifted) quadrature. + roots_sh_chebyu -- Gauss-Chebyshev (second kind, shifted) quadrature. + roots_sh_jacobi -- Gauss-Jacobi (shifted) quadrature. + +The functions below, in turn, return the polynomial coefficients in +``orthopoly1d`` objects, which function similarly as `numpy.poly1d`. +The ``orthopoly1d`` class also has an attribute ``weights``, which returns +the roots, weights, and total weights for the appropriate form of Gaussian +quadrature. These are returned in an ``n x 3`` array with roots in the first +column, weights in the second column, and total weights in the final column. +Note that ``orthopoly1d`` objects are converted to `~numpy.poly1d` when doing +arithmetic, and lose information of the original orthogonal polynomial. + +.. autosummary:: + :toctree: generated/ + + legendre -- Legendre polynomial. + chebyt -- Chebyshev polynomial of the first kind. + chebyu -- Chebyshev polynomial of the second kind. + chebyc -- Chebyshev polynomial of the first kind on :math:`[-2, 2]`. + chebys -- Chebyshev polynomial of the second kind on :math:`[-2, 2]`. + jacobi -- Jacobi polynomial. + laguerre -- Laguerre polynomial. + genlaguerre -- Generalized (associated) Laguerre polynomial. + hermite -- Physicist's Hermite polynomial. + hermitenorm -- Normalized (probabilist's) Hermite polynomial. + gegenbauer -- Gegenbauer (ultraspherical) polynomial. + sh_legendre -- Shifted Legendre polynomial. + sh_chebyt -- Shifted Chebyshev polynomial of the first kind. + sh_chebyu -- Shifted Chebyshev polynomial of the second kind. + sh_jacobi -- Shifted Jacobi polynomial. + +.. warning:: + + Computing values of high-order polynomials (around ``order > 20``) using + polynomial coefficients is numerically unstable. To evaluate polynomial + values, the ``eval_*`` functions should be used instead. + + +Hypergeometric functions +------------------------ + +.. autosummary:: + :toctree: generated/ + + hyp2f1 -- Gauss hypergeometric function 2F1(a, b; c; z). + hyp1f1 -- Confluent hypergeometric function 1F1(a, b; x). + hyperu -- Confluent hypergeometric function U(a, b, x) of the second kind. + hyp0f1 -- Confluent hypergeometric limit function 0F1. + + +Parabolic cylinder functions +---------------------------- + +.. autosummary:: + :toctree: generated/ + + pbdv -- Parabolic cylinder function D. + pbvv -- Parabolic cylinder function V. + pbwa -- Parabolic cylinder function W. + +The following functions do not accept NumPy arrays (they are not +universal functions): + +.. autosummary:: + :toctree: generated/ + + pbdv_seq -- Parabolic cylinder functions Dv(x) and derivatives. + pbvv_seq -- Parabolic cylinder functions Vv(x) and derivatives. + pbdn_seq -- Parabolic cylinder functions Dn(z) and derivatives. + +Mathieu and related functions +----------------------------- + +.. autosummary:: + :toctree: generated/ + + mathieu_a -- Characteristic value of even Mathieu functions. + mathieu_b -- Characteristic value of odd Mathieu functions. + +The following functions do not accept NumPy arrays (they are not +universal functions): + +.. autosummary:: + :toctree: generated/ + + mathieu_even_coef -- Fourier coefficients for even Mathieu and modified Mathieu functions. + mathieu_odd_coef -- Fourier coefficients for even Mathieu and modified Mathieu functions. + +The following return both function and first derivative: + +.. autosummary:: + :toctree: generated/ + + mathieu_cem -- Even Mathieu function and its derivative. + mathieu_sem -- Odd Mathieu function and its derivative. + mathieu_modcem1 -- Even modified Mathieu function of the first kind and its derivative. + mathieu_modcem2 -- Even modified Mathieu function of the second kind and its derivative. + mathieu_modsem1 -- Odd modified Mathieu function of the first kind and its derivative. + mathieu_modsem2 -- Odd modified Mathieu function of the second kind and its derivative. + +Spheroidal wave functions +------------------------- + +.. autosummary:: + :toctree: generated/ + + pro_ang1 -- Prolate spheroidal angular function of the first kind and its derivative. + pro_rad1 -- Prolate spheroidal radial function of the first kind and its derivative. + pro_rad2 -- Prolate spheroidal radial function of the second kind and its derivative. + obl_ang1 -- Oblate spheroidal angular function of the first kind and its derivative. + obl_rad1 -- Oblate spheroidal radial function of the first kind and its derivative. + obl_rad2 -- Oblate spheroidal radial function of the second kind and its derivative. + pro_cv -- Characteristic value of prolate spheroidal function. + obl_cv -- Characteristic value of oblate spheroidal function. + pro_cv_seq -- Characteristic values for prolate spheroidal wave functions. + obl_cv_seq -- Characteristic values for oblate spheroidal wave functions. + +The following functions require pre-computed characteristic value: + +.. autosummary:: + :toctree: generated/ + + pro_ang1_cv -- Prolate spheroidal angular function pro_ang1 for precomputed characteristic value. + pro_rad1_cv -- Prolate spheroidal radial function pro_rad1 for precomputed characteristic value. + pro_rad2_cv -- Prolate spheroidal radial function pro_rad2 for precomputed characteristic value. + obl_ang1_cv -- Oblate spheroidal angular function obl_ang1 for precomputed characteristic value. + obl_rad1_cv -- Oblate spheroidal radial function obl_rad1 for precomputed characteristic value. + obl_rad2_cv -- Oblate spheroidal radial function obl_rad2 for precomputed characteristic value. + +Kelvin functions +---------------- + +.. autosummary:: + :toctree: generated/ + + kelvin -- Kelvin functions as complex numbers. + kelvin_zeros -- Compute nt zeros of all Kelvin functions. + ber -- Kelvin function ber. + bei -- Kelvin function bei + berp -- Derivative of the Kelvin function `ber`. + beip -- Derivative of the Kelvin function `bei`. + ker -- Kelvin function ker. + kei -- Kelvin function ker. + kerp -- Derivative of the Kelvin function ker. + keip -- Derivative of the Kelvin function kei. + +The following functions do not accept NumPy arrays (they are not +universal functions): + +.. autosummary:: + :toctree: generated/ + + ber_zeros -- Compute nt zeros of the Kelvin function ber(x). + bei_zeros -- Compute nt zeros of the Kelvin function bei(x). + berp_zeros -- Compute nt zeros of the Kelvin function ber'(x). + beip_zeros -- Compute nt zeros of the Kelvin function bei'(x). + ker_zeros -- Compute nt zeros of the Kelvin function ker(x). + kei_zeros -- Compute nt zeros of the Kelvin function kei(x). + kerp_zeros -- Compute nt zeros of the Kelvin function ker'(x). + keip_zeros -- Compute nt zeros of the Kelvin function kei'(x). + +Combinatorics +------------- + +.. autosummary:: + :toctree: generated/ + + comb -- The number of combinations of N things taken k at a time. + perm -- Permutations of N things taken k at a time, i.e., k-permutations of N. + stirling2 -- Stirling numbers of the second kind. + +Lambert W and related functions +------------------------------- + +.. autosummary:: + :toctree: generated/ + + lambertw -- Lambert W function. + wrightomega -- Wright Omega function. + +Other special functions +----------------------- + +.. autosummary:: + :toctree: generated/ + + agm -- Arithmetic, Geometric Mean. + bernoulli -- Bernoulli numbers B0..Bn (inclusive). + binom -- Binomial coefficient + diric -- Periodic sinc function, also called the Dirichlet kernel. + euler -- Euler numbers E0..En (inclusive). + expn -- Exponential integral E_n. + exp1 -- Exponential integral E_1 of complex argument z. + expi -- Exponential integral Ei. + factorial -- The factorial of a number or array of numbers. + factorial2 -- Double factorial. + factorialk -- Multifactorial of n of order k, n(!!...!). + shichi -- Hyperbolic sine and cosine integrals. + sici -- Sine and cosine integrals. + softmax -- Softmax function. + log_softmax -- Logarithm of softmax function. + spence -- Spence's function, also known as the dilogarithm. + zeta -- Riemann zeta function. + zetac -- Riemann zeta function minus 1. + softplus -- Softplus function. + +Convenience functions +--------------------- + +.. autosummary:: + :toctree: generated/ + + cbrt -- Cube root of `x`. + exp10 -- 10**x. + exp2 -- 2**x. + radian -- Convert from degrees to radians. + cosdg -- Cosine of the angle `x` given in degrees. + sindg -- Sine of angle given in degrees. + tandg -- Tangent of angle x given in degrees. + cotdg -- Cotangent of the angle `x` given in degrees. + log1p -- Calculates log(1+x) for use when `x` is near zero. + expm1 -- ``exp(x) - 1`` for use when `x` is near zero. + cosm1 -- ``cos(x) - 1`` for use when `x` is near zero. + powm1 -- ``x**y - 1`` for use when `y` is near zero or `x` is near 1. + round -- Round to nearest integer. + xlogy -- Compute ``x*log(y)`` so that the result is 0 if ``x = 0``. + xlog1py -- Compute ``x*log1p(y)`` so that the result is 0 if ``x = 0``. + logsumexp -- Compute the log of the sum of exponentials of input elements. + exprel -- Relative error exponential, (exp(x)-1)/x, for use when `x` is near zero. + sinc -- Return the sinc function. + +""" # noqa: E501 + + +from ._sf_error import SpecialFunctionWarning, SpecialFunctionError + +from . import _ufuncs +from ._ufuncs import * + +from . import _basic +from ._basic import * + +# Replace some function definitions from _ufuncs and _basic +# to add Array API support +from ._support_alternative_backends import * + +from ._logsumexp import logsumexp, softmax, log_softmax + +from . import _multiufuncs +from ._multiufuncs import * + +from . import _orthogonal +from ._orthogonal import * + +from ._ellip_harm import ( + ellip_harm, + ellip_harm_2, + ellip_normal +) +from ._lambertw import lambertw +from ._spherical_bessel import ( + spherical_jn, + spherical_yn, + spherical_in, + spherical_kn +) + +# Deprecated namespaces, to be removed in v2.0.0 +from . import add_newdocs, basic, orthogonal, specfun, sf_error, spfun_stats + +# We replace some function definitions from _ufuncs with those from +# _support_alternative_backends above, but those are all listed in _ufuncs.__all__, +# so there is no need to consider _support_alternative_backends.__all__ here. +__all__ = _ufuncs.__all__ + _basic.__all__ + _orthogonal.__all__ + _multiufuncs.__all__ +__all__ += [ + 'SpecialFunctionWarning', + 'SpecialFunctionError', + 'logsumexp', + 'softmax', + 'log_softmax', + 'multigammaln', # pyrefly:ignore[bad-dunder-all] + 'ellip_harm', + 'ellip_harm_2', + 'ellip_normal', + 'lambertw', + 'spherical_jn', + 'spherical_yn', + 'spherical_in', + 'spherical_kn', +] + +from scipy._lib._testutils import PytestTester +test = PytestTester(__name__) +del PytestTester diff --git a/v12x/pylib/scipy/special/_add_newdocs.py b/v12x/pylib/scipy/special/_add_newdocs.py new file mode 100644 index 0000000000000000000000000000000000000000..63d2581d620d4c20640f94573f36a25885f6a1a5 --- /dev/null +++ b/v12x/pylib/scipy/special/_add_newdocs.py @@ -0,0 +1,8913 @@ +# Docstrings for generated ufuncs +# +# The syntax is designed to look like the function add_newdoc is being +# called from numpy.lib, but in this file add_newdoc puts the +# docstrings in a dictionary. This dictionary is used in +# _generate_pyx.py to generate the docstrings for the ufuncs in +# scipy.special at the C level when the ufuncs are created at compile +# time. + +docdict: dict[str, str] = {} + + +def get(name): + return docdict.get(name) + + +def add_newdoc(name, doc): + docdict[name] = doc + + +add_newdoc("_sf_error_test_function", + """ + Private function; do not use. + """) + + +add_newdoc("_cosine_cdf", + """ + _cosine_cdf(x) + + Cumulative distribution function (CDF) of the cosine distribution:: + + { 0, x < -pi + cdf(x) = { (pi + x + sin(x))/(2*pi), -pi <= x <= pi + { 1, x > pi + + Parameters + ---------- + x : array_like + `x` must contain real numbers. + + Returns + ------- + scalar or ndarray + The cosine distribution CDF evaluated at `x`. + + """) + +add_newdoc("_cosine_invcdf", + """ + _cosine_invcdf(p) + + Inverse of the cumulative distribution function (CDF) of the cosine + distribution. + + The CDF of the cosine distribution is:: + + cdf(x) = (pi + x + sin(x))/(2*pi) + + This function computes the inverse of cdf(x). + + Parameters + ---------- + p : array_like + `p` must contain real numbers in the interval ``0 <= p <= 1``. + `nan` is returned for values of `p` outside the interval [0, 1]. + + Returns + ------- + scalar or ndarray + The inverse of the cosine distribution CDF evaluated at `p`. + + """) + +add_newdoc("_ellip_harm", + """ + Internal function, use `ellip_harm` instead. + """) + +add_newdoc("_ellip_norm", + """ + Internal function, use `ellip_norm` instead. + """) + +add_newdoc("wrightomega", + r""" + wrightomega(z, out=None) + + Wright Omega function. + + Defined as the solution to + + .. math:: + + \omega + \log(\omega) = z + + where :math:`\log` is the principal branch of the complex logarithm. + + Parameters + ---------- + z : array_like + Points at which to evaluate the Wright Omega function + out : ndarray, optional + Optional output array for the function values + + Returns + ------- + omega : scalar or ndarray + Values of the Wright Omega function + + See Also + -------- + lambertw : The Lambert W function + + Notes + ----- + .. versionadded:: 0.19.0 + + The function can also be defined as + + .. math:: + + \omega(z) = W_{K(z)}(e^z) + + where :math:`K(z) = \lceil (\Im(z) - \pi)/(2\pi) \rceil` is the + unwinding number and :math:`W` is the Lambert W function. + + The implementation here is taken from [1]_. + + References + ---------- + .. [1] Lawrence, Corless, and Jeffrey, "Algorithm 917: Complex + Double-Precision Evaluation of the Wright :math:`\omega` + Function." ACM Transactions on Mathematical Software, + 2012. :doi:`10.1145/2168773.2168779`. + + Examples + -------- + >>> import numpy as np + >>> from scipy.special import wrightomega, lambertw + + >>> wrightomega([-2, -1, 0, 1, 2]) + array([0.12002824, 0.27846454, 0.56714329, 1. , 1.5571456 ]) + + Complex input: + + >>> wrightomega(3 + 5j) + (1.5804428632097158+3.8213626783287937j) + + Verify that ``wrightomega(z)`` satisfies ``w + log(w) = z``: + + >>> w = -5 + 4j + >>> wrightomega(w + np.log(w)) + (-5+4j) + + Verify the connection to ``lambertw``: + + >>> z = 0.5 + 3j + >>> wrightomega(z) + (0.0966015889280649+1.4937828458191993j) + >>> lambertw(np.exp(z)) + (0.09660158892806493+1.4937828458191993j) + + >>> z = 0.5 + 4j + >>> wrightomega(z) + (-0.3362123489037213+2.282986001579032j) + >>> lambertw(np.exp(z), k=1) + (-0.33621234890372115+2.282986001579032j) + """) + + +add_newdoc("agm", + """ + agm(a, b, out=None) + + Compute the arithmetic-geometric mean of `a` and `b`. + + Start with a_0 = a and b_0 = b and iteratively compute:: + + a_{n+1} = (a_n + b_n)/2 + b_{n+1} = sqrt(a_n*b_n) + + a_n and b_n converge to the same limit as n increases; their common + limit is agm(a, b). + + Parameters + ---------- + a, b : array_like + Real values only. If the values are both negative, the result + is negative. If one value is negative and the other is positive, + `nan` is returned. + out : ndarray, optional + Optional output array for the function values + + Returns + ------- + scalar or ndarray + The arithmetic-geometric mean of `a` and `b`. + + Examples + -------- + >>> import numpy as np + >>> from scipy.special import agm + >>> a, b = 24.0, 6.0 + >>> agm(a, b) + 13.458171481725614 + + Compare that result to the iteration: + + >>> while a != b: + ... a, b = (a + b)/2, np.sqrt(a*b) + ... print("a = %19.16f b=%19.16f" % (a, b)) + ... + a = 15.0000000000000000 b=12.0000000000000000 + a = 13.5000000000000000 b=13.4164078649987388 + a = 13.4582039324993694 b=13.4581390309909850 + a = 13.4581714817451772 b=13.4581714817060547 + a = 13.4581714817256159 b=13.4581714817256159 + + When array-like arguments are given, broadcasting applies: + + >>> a = np.array([[1.5], [3], [6]]) # a has shape (3, 1). + >>> b = np.array([6, 12, 24, 48]) # b has shape (4,). + >>> agm(a, b) + array([[ 3.36454287, 5.42363427, 9.05798751, 15.53650756], + [ 4.37037309, 6.72908574, 10.84726853, 18.11597502], + [ 6. , 8.74074619, 13.45817148, 21.69453707]]) + """) + +add_newdoc("bdtr", + r""" + bdtr(k, n, p, out=None) + + Binomial distribution cumulative distribution function. + + Sum of the terms 0 through `floor(k)` of the Binomial probability density. + + .. math:: + \mathrm{bdtr}(k, n, p) = + \sum_{j=0}^{\lfloor k \rfloor} {{n}\choose{j}} p^j (1-p)^{n-j} + + Parameters + ---------- + k : array_like + Number of successes (double), rounded down to the nearest integer. + n : array_like + Number of events (int). + p : array_like + Probability of success in a single event (float). + out : ndarray, optional + Optional output array for the function values + + Returns + ------- + y : scalar or ndarray + Probability of `floor(k)` or fewer successes in `n` independent events with + success probabilities of `p`. + + Notes + ----- + The terms are not summed directly; instead the regularized incomplete beta + function is employed, according to the formula, + + .. math:: + \mathrm{bdtr}(k, n, p) = + I_{1 - p}(n - \lfloor k \rfloor, \lfloor k \rfloor + 1). + + Wrapper for the Cephes [1]_ routine `bdtr`. + + References + ---------- + .. [1] Cephes Mathematical Functions Library, + https://netlib.org/cephes/ + + """) + +add_newdoc("bdtrc", + r""" + bdtrc(k, n, p, out=None) + + Binomial distribution survival function. + + Sum of the terms `floor(k) + 1` through `n` of the binomial probability + density, + + .. math:: + \mathrm{bdtrc}(k, n, p) = + \sum_{j=\lfloor k \rfloor +1}^n {{n}\choose{j}} p^j (1-p)^{n-j} + + Parameters + ---------- + k : array_like + Number of successes (double), rounded down to nearest integer. + n : array_like + Number of events (int) + p : array_like + Probability of success in a single event. + out : ndarray, optional + Optional output array for the function values + + Returns + ------- + y : scalar or ndarray + Probability of `floor(k) + 1` or more successes in `n` independent + events with success probabilities of `p`. + + See Also + -------- + bdtr + betainc + + Notes + ----- + The terms are not summed directly; instead the regularized incomplete beta + function is employed, according to the formula, + + .. math:: + \mathrm{bdtrc}(k, n, p) = I_{p}(\lfloor k \rfloor + 1, n - \lfloor k \rfloor). + + Wrapper for the Cephes [1]_ routine `bdtrc`. + + References + ---------- + .. [1] Cephes Mathematical Functions Library, + https://netlib.org/cephes/ + + """) + +add_newdoc("bdtri", + r""" + bdtri(k, n, y, out=None) + + Inverse function to `bdtr` with respect to `p`. + + Finds the event probability `p` such that the sum of the terms 0 through + `k` of the binomial probability density is equal to the given cumulative + probability `y`. + + Parameters + ---------- + k : array_like + Number of successes (float), rounded down to the nearest integer. + n : array_like + Number of events (float) + y : array_like + Cumulative probability (probability of `k` or fewer successes in `n` + events). + out : ndarray, optional + Optional output array for the function values + + Returns + ------- + p : scalar or ndarray + The event probability such that `bdtr(\lfloor k \rfloor, n, p) = y`. + + See Also + -------- + bdtr + betaincinv + + Notes + ----- + The computation is carried out using the inverse beta integral function + and the relation,:: + + 1 - p = betaincinv(n - k, k + 1, y). + + Wrapper for the Cephes [1]_ routine `bdtri`. + + References + ---------- + .. [1] Cephes Mathematical Functions Library, + https://netlib.org/cephes/ + """) + +add_newdoc("bdtrik", + """ + bdtrik(y, n, p, out=None) + + Inverse function to `bdtr` with respect to `k`. + + Finds the number of successes `k` such that the sum of the terms 0 through + `k` of the Binomial probability density for `n` events with probability + `p` is equal to the given cumulative probability `y`. + + Parameters + ---------- + y : array_like + Cumulative probability (probability of `k` or fewer successes in `n` + events). + n : array_like + Number of events (float). + p : array_like + Success probability (float). + out : ndarray, optional + Optional output array for the function values + + Returns + ------- + k : scalar or ndarray + The number of successes `k` such that `bdtr(k, n, p) = y`. + + See Also + -------- + bdtr + + Notes + ----- + Formula 26.5.24 of [1]_ (or equivalently [2]_) is used to reduce the binomial + distribution to the cumulative incomplete beta distribution. + + Computation of `k` involves a search for a value that produces the desired + value of `y`. The search relies on the monotonicity of `y` with `k`. + + Wrapper for the CDFLIB [3]_ Fortran routine `cdfbin`. + + References + ---------- + .. [1] Milton Abramowitz and Irene A. Stegun, eds. + Handbook of Mathematical Functions with Formulas, + Graphs, and Mathematical Tables. New York: Dover, 1972. + .. [2] NIST Digital Library of Mathematical Functions + https://dlmf.nist.gov/8.17.5#E5 + .. [3] Barry Brown, James Lovato, and Kathy Russell, + CDFLIB: Library of Fortran Routines for Cumulative Distribution + Functions, Inverses, and Other Parameters. + + """) + +add_newdoc("bdtrin", + r""" + bdtrin(k, y, p, out=None) + + Inverse function to `bdtr` with respect to `n`. + + Finds the number of events `n` such that the sum of the terms 0 through + `k` of the Binomial probability density for events with probability `p` is + equal to the given cumulative probability `y`. + + Parameters + ---------- + k : array_like + Number of successes (float). + y : array_like + Cumulative probability (probability of `k` or fewer successes in `n` + events). + p : array_like + Success probability (float). + out : ndarray, optional + Optional output array for the function values + + Returns + ------- + n : scalar or ndarray + The number of events `n` such that `bdtr(k, n, p) = y`. + + See Also + -------- + bdtr + + Notes + ----- + This function uses the `find_minimum_number_of_trials` method of the + `binomial_distribution` class of the Boost.Math C++ library [1]_. + + References + ---------- + .. [1] The Boost Developers. "Boost C++ Libraries". https://www.boost.org/. + + Examples + -------- + How often do we have to flip a fair coin to have at least a 90% chance + of getting 10 heads? `bdtrin` answers this question: + + >>> import scipy.special as sc + >>> k = 10 # number of times we want heads + >>> p = 0.5 # probability of flipping heads + >>> y = 0.9 # cumulative probability + >>> result = sc.bdtrin(k, y, p) + >>> result + 15.90442928275109 + + To verify, compute the cumulative probability of getting 10 or fewer + successes in 16 trials with probability 0.5 using the binomial + distribution from `scipy.stats`. Since `bdtrin` returns a non-integer + number of trials, we round up to the next integer: + + >>> from scipy.stats import Binomial + >>> Binomial(n=16, p=p).cdf(k) + 0.8949432373046875 + + """) + +add_newdoc("btdtria", + r""" + btdtria(p, b, x, out=None) + + Inverse of `betainc` with respect to `a`. + + This is the inverse of the beta cumulative distribution function, `betainc`, + considered as a function of `a`, returning the value of `a` for which + `betainc(a, b, x) = p`, or + + .. math:: + p = \int_0^x \frac{\Gamma(a + b)}{\Gamma(a)\Gamma(b)} t^{a-1} (1-t)^{b-1}\,dt + + Parameters + ---------- + p : array_like + Cumulative probability, in [0, 1]. + b : array_like + Shape parameter (`b` > 0). + x : array_like + The quantile, in [0, 1]. + out : ndarray, optional + Optional output array for the function values + + Returns + ------- + a : scalar or ndarray + The value of the shape parameter `a` such that `betainc(a, b, x) = p`. + + See Also + -------- + betainc : Regularized incomplete beta function + betaincinv : Inverse of the regularized incomplete beta function + btdtrib : Inverse of the beta cumulative distribution function, with respect to `b`. + + Notes + ----- + This function wraps the ``ibeta_inva`` routine from the + Boost Math C++ library [1]_. + + References + ---------- + .. [1] The Boost Developers. "Boost C++ Libraries". https://www.boost.org/. + + Examples + -------- + >>> import scipy.special as sc + + This function is the inverse of `betainc` for fixed + values of :math:`b` and :math:`x`. + + >>> a, b, x = 1.2, 3.1, 0.2 + >>> y = sc.betainc(a, b, x) + >>> sc.btdtria(y, b, x) + 1.2 + + """) + +add_newdoc("btdtrib", + r""" + btdtria(a, p, x, out=None) + + Inverse of `betainc` with respect to `b`. + + This is the inverse of the beta cumulative distribution function, `betainc`, + considered as a function of `b`, returning the value of `b` for which + `betainc(a, b, x) = p`, or + + .. math:: + p = \int_0^x \frac{\Gamma(a + b)}{\Gamma(a)\Gamma(b)} t^{a-1} (1-t)^{b-1}\,dt + + Parameters + ---------- + a : array_like + Shape parameter (`a` > 0). + p : array_like + Cumulative probability, in [0, 1]. + x : array_like + The quantile, in [0, 1]. + out : ndarray, optional + Optional output array for the function values + + Returns + ------- + b : scalar or ndarray + The value of the shape parameter `b` such that `betainc(a, b, x) = p`. + + See Also + -------- + betainc : Regularized incomplete beta function + betaincinv : Inverse of the regularized incomplete beta function with + respect to `x`. + btdtria : Inverse of the beta cumulative distribution function, with respect to `a`. + + Notes + ----- + Wrapper for the `ibeta_invb` routine from the Boost Math C++ library [1]_. + + References + ---------- + .. [1] The Boost Developers. "Boost C++ Libraries". https://www.boost.org/. + + Examples + -------- + >>> import scipy.special as sc + >>> a, b, x = 1.2, 3.1, 0.2 + >>> y = sc.betainc(a, b, x) + + `btdtrib` is the inverse of `betainc` for fixed values of :math:`a` and + :math:`x`: + + >>> sc.btdtrib(a, y, x) + 3.1 + + """) + +add_newdoc( + "betainc", + r""" + betainc(a, b, x, out=None) + + Regularized incomplete beta function. + + Computes the regularized incomplete beta function, defined as [1]_: + + .. math:: + + I_x(a, b) = \frac{\Gamma(a+b)}{\Gamma(a)\Gamma(b)} \int_0^x + t^{a-1}(1-t)^{b-1}dt, + + for :math:`0 \leq x \leq 1`. + + This function is the cumulative distribution function for the beta + distribution; its range is [0, 1]. + + Parameters + ---------- + a, b : array_like + Positive, real-valued parameters. + x : array_like + Real-valued such that :math:`0 \leq x \leq 1`, + the upper limit of integration. + out : ndarray, optional + Optional output array for the function values. + + Returns + ------- + scalar or ndarray + Value of the regularized incomplete beta function. + + See Also + -------- + beta : beta function + betaincinv : inverse of the regularized incomplete beta function + betaincc : complement of the regularized incomplete beta function + scipy.stats.beta : beta distribution + + Notes + ----- + The term *regularized* in the name of this function refers to the + scaling of the function by the gamma function terms shown in the + formula. When not qualified as *regularized*, the name *incomplete + beta function* often refers to just the integral expression, + without the gamma terms. One can use the function `beta` from + `scipy.special` to get this "nonregularized" incomplete beta + function by multiplying the result of ``betainc(a, b, x)`` by + ``beta(a, b)``. + + ``betainc(a, b, x)`` is treated as a two parameter family of functions + of a single variable `x`, rather than as a function of three variables. + This impacts only the limiting cases ``a = 0``, ``b = 0``, ``a = inf``, + ``b = inf``. + + In general + + .. math:: + + \lim_{(a, b) \rightarrow (a_0, b_0)} \mathrm{betainc}(a, b, x) + + is treated as a pointwise limit in ``x``. Thus for example, + ``betainc(0, b, 0)`` equals ``0`` for ``b > 0``, although it would be + indeterminate when considering the simultaneous limit ``(a, x) -> (0+, 0+)``. + + This function wraps the ``ibeta`` routine from the + Boost Math C++ library [2]_. + + References + ---------- + .. [1] NIST Digital Library of Mathematical Functions + https://dlmf.nist.gov/8.17 + .. [2] The Boost Developers. "Boost C++ Libraries". https://www.boost.org/. + + Examples + -------- + + Let :math:`B(a, b)` be the `beta` function. + + >>> import scipy.special as sc + + The coefficient in terms of `gamma` is equal to + :math:`1/B(a, b)`. Also, when :math:`x=1` + the integral is equal to :math:`B(a, b)`. + Therefore, :math:`I_{x=1}(a, b) = 1` for any :math:`a, b`. + + >>> sc.betainc(0.2, 3.5, 1.0) + 1.0 + + It satisfies + :math:`I_x(a, b) = x^a F(a, 1-b, a+1, x)/ (aB(a, b))`, + where :math:`F` is the hypergeometric function `hyp2f1`: + + >>> a, b, x = 1.4, 3.1, 0.5 + >>> x**a * sc.hyp2f1(a, 1 - b, a + 1, x)/(a * sc.beta(a, b)) + 0.8148904036225295 + >>> sc.betainc(a, b, x) + 0.8148904036225296 + + This functions satisfies the relationship + :math:`I_x(a, b) = 1 - I_{1-x}(b, a)`: + + >>> sc.betainc(2.2, 3.1, 0.4) + 0.49339638807619446 + >>> 1 - sc.betainc(3.1, 2.2, 1 - 0.4) + 0.49339638807619446 + + """) + + +add_newdoc( + "betaincc", + r""" + betaincc(a, b, x, out=None) + + Complement of the regularized incomplete beta function. + + Computes the complement of the regularized incomplete beta function, + defined as [1]_: + + .. math:: + + \bar{I}_x(a, b) = 1 - I_x(a, b) + = 1 - \frac{\Gamma(a+b)}{\Gamma(a)\Gamma(b)} \int_0^x + t^{a-1}(1-t)^{b-1}dt, + + for :math:`0 \leq x \leq 1`. + + Parameters + ---------- + a, b : array_like + Positive, real-valued parameters. + x : array_like + Real-valued such that :math:`0 \leq x \leq 1`, + the upper limit of integration. + out : ndarray, optional + Optional output array for the function values. + + Returns + ------- + scalar or ndarray + Value of the complement of the regularized incomplete beta function. + + See Also + -------- + betainc : regularized incomplete beta function + betaincinv : inverse of the regularized incomplete beta function + betainccinv : + inverse of the complement of the regularized incomplete beta function + beta : beta function + scipy.stats.beta : beta distribution + + Notes + ----- + .. versionadded:: 1.11.0 + + Like `betainc`, ``betaincc(a, b, x)`` is treated as a two parameter + family of functions of a single variable `x`, rather than as a function of + three variables. See the `betainc` docstring for more info on how this + impacts limiting cases. + + This function wraps the ``ibetac`` routine from the + Boost Math C++ library [2]_. + + References + ---------- + .. [1] NIST Digital Library of Mathematical Functions + https://dlmf.nist.gov/8.17 + .. [2] The Boost Developers. "Boost C++ Libraries". https://www.boost.org/. + + Examples + -------- + >>> from scipy.special import betaincc, betainc + + The naive calculation ``1 - betainc(a, b, x)`` loses precision when + the values of ``betainc(a, b, x)`` are close to 1: + + >>> 1 - betainc(0.5, 8, [0.9, 0.99, 0.999]) + array([2.0574632e-09, 0.0000000e+00, 0.0000000e+00]) + + By using ``betaincc``, we get the correct values: + + >>> betaincc(0.5, 8, [0.9, 0.99, 0.999]) + array([2.05746321e-09, 1.97259354e-17, 1.96467954e-25]) + + """) + +add_newdoc( + "betaincinv", + r""" + betaincinv(a, b, y, out=None) + + Inverse of the regularized incomplete beta function. + + Computes :math:`x` such that: + + .. math:: + + y = I_x(a, b) = \frac{\Gamma(a+b)}{\Gamma(a)\Gamma(b)} + \int_0^x t^{a-1}(1-t)^{b-1}dt, + + where :math:`I_x` is the normalized incomplete beta function `betainc` + and :math:`\Gamma` is the `gamma` function [1]_. + + Parameters + ---------- + a, b : array_like + Positive, real-valued parameters + y : array_like + Real-valued input + out : ndarray, optional + Optional output array for function values + + Returns + ------- + scalar or ndarray + Value of the inverse of the regularized incomplete beta function + + See Also + -------- + betainc : regularized incomplete beta function + gamma : gamma function + + Notes + ----- + This function wraps the ``ibeta_inv`` routine from the + Boost Math C++ library [2]_. + + References + ---------- + .. [1] NIST Digital Library of Mathematical Functions + https://dlmf.nist.gov/8.17 + .. [2] The Boost Developers. "Boost C++ Libraries". https://www.boost.org/. + + Examples + -------- + >>> import scipy.special as sc + + This function is the inverse of `betainc` for fixed + values of :math:`a` and :math:`b`. + + >>> a, b = 1.2, 3.1 + >>> y = sc.betainc(a, b, 0.2) + >>> sc.betaincinv(a, b, y) + 0.2 + >>> + >>> a, b = 7.5, 0.4 + >>> x = sc.betaincinv(a, b, 0.5) + >>> sc.betainc(a, b, x) + 0.5 + + """) + + +add_newdoc( + "betainccinv", + r""" + betainccinv(a, b, y, out=None) + + Inverse of the complemented regularized incomplete beta function. + + Computes :math:`x` such that: + + .. math:: + + y = 1 - I_x(a, b) = 1 - \frac{\Gamma(a+b)}{\Gamma(a)\Gamma(b)} + \int_0^x t^{a-1}(1-t)^{b-1}dt, + + where :math:`I_x` is the normalized incomplete beta function `betainc` + and :math:`\Gamma` is the `gamma` function [1]_. + + Parameters + ---------- + a, b : array_like + Positive, real-valued parameters + y : array_like + Real-valued input + out : ndarray, optional + Optional output array for function values + + Returns + ------- + scalar or ndarray + Value of the inverse of the regularized incomplete beta function + + See Also + -------- + betainc : regularized incomplete beta function + betaincc : complement of the regularized incomplete beta function + + Notes + ----- + .. versionadded:: 1.11.0 + + This function wraps the ``ibetac_inv`` routine from the + Boost Math C++ library [2]_. + + References + ---------- + .. [1] NIST Digital Library of Mathematical Functions + https://dlmf.nist.gov/8.17 + .. [2] The Boost Developers. "Boost C++ Libraries". https://www.boost.org/. + + Examples + -------- + >>> from scipy.special import betainccinv, betaincc + + This function is the inverse of `betaincc` for fixed + values of :math:`a` and :math:`b`. + + >>> a, b = 1.2, 3.1 + >>> y = betaincc(a, b, 0.2) + >>> betainccinv(a, b, y) + 0.2 + + >>> a, b = 7, 2.5 + >>> x = betainccinv(a, b, 0.875) + >>> betaincc(a, b, x) + 0.875 + + """) + +add_newdoc("boxcox", + r""" + boxcox(x, lmbda, out=None) + + Compute the Box-Cox transformation. + + The Box-Cox transformation is + + .. math:: + + y = \begin{cases} + (x^\lambda - 1) / \lambda & \text{if } \lambda \neq 0 \\ + \log(x) & \text{if } \lambda = 0 + \end{cases} + + Returns ``nan`` if :math:`x < 0`. + Returns ``-inf`` if :math:`x = 0` and :math:`\lambda \leq 0`. + + Parameters + ---------- + x : array_like + Data to be transformed. + lmbda : array_like + Power parameter :math:`\lambda` of the Box-Cox transform. + out : ndarray, optional + Optional output array for the function values. + + Returns + ------- + y : scalar or ndarray + Transformed data. + + See Also + -------- + boxcox1p : Box-Cox transformation of ``1 + x``. + inv_boxcox : Inverse of the Box-Cox transformation. + + Notes + ----- + + .. versionadded:: 0.14.0 + + Examples + -------- + >>> from scipy.special import boxcox + >>> boxcox([1, 4, 10], 2.5) + array([ 0. , 12.4 , 126.09110641]) + >>> boxcox(2, [0, 1, 2]) + array([ 0.69314718, 1. , 1.5 ]) + """) + +add_newdoc("boxcox1p", + r""" + boxcox1p(x, lmbda, out=None) + + Compute the Box-Cox transformation of :math:`1 + x`. + + The Box-Cox transformation computed by `boxcox1p` is + + .. math:: + + y = \begin{cases} + ((1+x)^\lambda - 1) / \lambda & \text{if } \lambda \neq 0 \\ + \log(1+x) & \text{if } \lambda = 0 + \end{cases} + + Returns ``nan`` if :math:`x < -1`. + Returns ``-inf`` if :math:`x = -1` and :math:`\lambda \leq 0`. + + Parameters + ---------- + x : array_like + Data to be transformed. + lmbda : array_like + Power parameter :math:`\lambda` of the Box-Cox transform. + out : ndarray, optional + Optional output array for the function values. + + Returns + ------- + y : scalar or ndarray + Transformed data. + + See Also + -------- + boxcox : Box-Cox transformation. + inv_boxcox1p : Inverse of the Box-Cox transformation of ``1 + x``. + + Notes + ----- + + .. versionadded:: 0.14.0 + + Examples + -------- + >>> from scipy.special import boxcox1p + >>> boxcox1p(1e-4, [0, 0.5, 1]) + array([ 9.99950003e-05, 9.99975001e-05, 1.00000000e-04]) + >>> boxcox1p([0.01, 0.1], 0.25) + array([ 0.00996272, 0.09645476]) + """) + +add_newdoc("inv_boxcox", + r""" + inv_boxcox(y, lmbda, out=None) + + Compute the inverse of the Box-Cox transformation. + + Find :math:`x` such that + + .. math:: + + y = \begin{cases} + (x^\lambda - 1) / \lambda & \text{if } \lambda \neq 0 \\ + \log(x) & \text{if } \lambda = 0 + \end{cases} + + Parameters + ---------- + y : array_like + Transformed data (input to the inverse transform). + lmbda : array_like + Power parameter :math:`\lambda` of the Box-Cox transform. + out : ndarray, optional + Optional output array for the function values. + + Returns + ------- + x : scalar or ndarray + Original data (inverse Box-Cox transform of `y`). + + See Also + -------- + boxcox : Box-Cox transformation. + inv_boxcox1p : Inverse of the Box-Cox transformation of ``1 + x``. + + Notes + ----- + + .. versionadded:: 0.16.0 + + Examples + -------- + >>> from scipy.special import boxcox, inv_boxcox + >>> y = boxcox([1, 4, 10], 2.5) + >>> inv_boxcox(y, 2.5) + array([1., 4., 10.]) + """) + +add_newdoc("inv_boxcox1p", + r""" + inv_boxcox1p(y, lmbda, out=None) + + Compute the inverse of the Box-Cox transformation of :math:`1 + x`. + + Find :math:`x` such that + + .. math:: + + y = \begin{cases} + ((1+x)^\lambda - 1) / \lambda & \text{if } \lambda \neq 0 \\ + \log(1+x) & \text{if } \lambda = 0 + \end{cases} + + Parameters + ---------- + y : array_like + Transformed data (input to the inverse transform). + lmbda : array_like + Power parameter :math:`\lambda` of the Box-Cox transform. + out : ndarray, optional + Optional output array for the function values. + + Returns + ------- + x : scalar or ndarray + Original data (inverse Box-Cox transform of `y`). + + See Also + -------- + boxcox1p : Box-Cox transformation of ``1 + x``. + inv_boxcox : Inverse of the Box-Cox transformation. + + Notes + ----- + + .. versionadded:: 0.16.0 + + Examples + -------- + >>> from scipy.special import boxcox1p, inv_boxcox1p + >>> y = boxcox1p([1, 4, 10], 2.5) + >>> inv_boxcox1p(y, 2.5) + array([1., 4., 10.]) + """) + +add_newdoc("chdtr", + r""" + chdtr(v, x, out=None) + + Chi square cumulative distribution function. + + Returns the area under the left tail (from 0 to `x`) of the Chi + square probability density function with `v` degrees of freedom: + + .. math:: + + \frac{1}{2^{v/2} \Gamma(v/2)} \int_0^x t^{v/2 - 1} e^{-t/2} dt + + Here :math:`\Gamma` is the Gamma function; see `gamma`. This + integral can be expressed in terms of the regularized lower + incomplete gamma function `gammainc` as + ``gammainc(v / 2, x / 2)``. [1]_ + + Parameters + ---------- + v : array_like + Degrees of freedom. + x : array_like + Upper bound of the integral. + out : ndarray, optional + Optional output array for the function results. + + Returns + ------- + scalar or ndarray + Values of the cumulative distribution function. + + See Also + -------- + chdtrc, chdtri, chdtriv, gammainc + + References + ---------- + .. [1] Chi-Square distribution, + https://www.itl.nist.gov/div898/handbook/eda/section3/eda3666.htm + + Examples + -------- + >>> import numpy as np + >>> import scipy.special as sc + + It can be expressed in terms of the regularized lower incomplete + gamma function. + + >>> v = 1 + >>> x = np.arange(4) + >>> sc.chdtr(v, x) + array([0. , 0.68268949, 0.84270079, 0.91673548]) + >>> sc.gammainc(v / 2, x / 2) + array([0. , 0.68268949, 0.84270079, 0.91673548]) + + """) + +add_newdoc("chdtrc", + r""" + chdtrc(v, x, out=None) + + Chi square survival function. + + Returns the area under the right hand tail (from `x` to infinity) + of the Chi square probability density function with `v` degrees of + freedom: + + .. math:: + + \frac{1}{2^{v/2} \Gamma(v/2)} \int_x^\infty t^{v/2 - 1} e^{-t/2} dt + + Here :math:`\Gamma` is the Gamma function; see `gamma`. This + integral can be expressed in terms of the regularized upper + incomplete gamma function `gammaincc` as + ``gammaincc(v / 2, x / 2)``. [1]_ + + Parameters + ---------- + v : array_like + Degrees of freedom. + x : array_like + Lower bound of the integral. + out : ndarray, optional + Optional output array for the function results. + + Returns + ------- + scalar or ndarray + Values of the survival function. + + See Also + -------- + chdtr, chdtri, chdtriv, gammaincc + + References + ---------- + .. [1] Chi-Square distribution, + https://www.itl.nist.gov/div898/handbook/eda/section3/eda3666.htm + + Examples + -------- + >>> import numpy as np + >>> import scipy.special as sc + + It can be expressed in terms of the regularized upper incomplete + gamma function. + + >>> v = 1 + >>> x = np.arange(4) + >>> sc.chdtrc(v, x) + array([1. , 0.31731051, 0.15729921, 0.08326452]) + >>> sc.gammaincc(v / 2, x / 2) + array([1. , 0.31731051, 0.15729921, 0.08326452]) + + """) + +add_newdoc("chdtri", + """ + chdtri(v, p, out=None) + + Inverse to `chdtrc` with respect to `x`. + + Returns `x` such that ``chdtrc(v, x) == p``. + + Parameters + ---------- + v : array_like + Degrees of freedom. + p : array_like + Probability. + out : ndarray, optional + Optional output array for the function results. + + Returns + ------- + x : scalar or ndarray + Value so that the probability a Chi square random variable + with `v` degrees of freedom is greater than `x` equals `p`. + + See Also + -------- + chdtrc, chdtr, chdtriv + + References + ---------- + .. [1] Chi-Square distribution, + https://www.itl.nist.gov/div898/handbook/eda/section3/eda3666.htm + + Examples + -------- + >>> import scipy.special as sc + + It inverts `chdtrc`. + + >>> v, p = 1, 0.3 + >>> sc.chdtrc(v, sc.chdtri(v, p)) + 0.3 + >>> x = 1 + >>> sc.chdtri(v, sc.chdtrc(v, x)) + 1.0 + + """) + +add_newdoc("chdtriv", + """ + chdtriv(p, x, out=None) + + Inverse to `chdtr` with respect to `v`. + + Returns `v` such that ``chdtr(v, x) == p``. + + Parameters + ---------- + p : array_like + Probability that the Chi square random variable is less than + or equal to `x`. + x : array_like + Nonnegative input. + out : ndarray, optional + Optional output array for the function results. + + Returns + ------- + scalar or ndarray + Degrees of freedom. + + See Also + -------- + chdtr, chdtrc, chdtri + + Notes + ----- + This function wraps routines from the Boost Math C++ library [1]_. + + References + ---------- + .. [1] The Boost Developers. "Boost C++ Libraries". https://www.boost.org/. + .. [2] Chi-Square distribution, + https://www.itl.nist.gov/div898/handbook/eda/section3/eda3666.htm + + Examples + -------- + >>> import scipy.special as sc + + It inverts `chdtr`. + + >>> p, x = 0.5, 1 + >>> sc.chdtr(sc.chdtriv(p, x), x) + 0.5000000000000003 + >>> v = 1 + >>> sc.chdtriv(sc.chdtr(v, x), x) + 1.0 + + """) + +add_newdoc("chndtr", + r""" + chndtr(x, df, nc, out=None) + + Non-central chi-squared cumulative distribution function. + + The cumulative distribution function is given by + + .. math:: + + F_{\nu,\lambda}(x) + = \sum_{j=0}^{\infty} + e^{-\lambda / 2} + \frac{(\lambda / 2)^j}{j!} + F_{\chi^2_{\nu + 2j}}(x), + + where :math:`\nu > 0` is the degrees of freedom (``df``), :math:`\lambda \geq 0` + is the non-centrality parameter (``nc``), and :math:`F_{\chi^2_{\nu + 2j}}` is the + CDF of the central chi-squared distribution with :math:`\nu + 2j` degrees of + freedom. + + Parameters + ---------- + x : array_like + Upper bound of the integral; must satisfy ``x >= 0``. + df : array_like + Degrees of freedom; must satisfy ``df > 0``. + nc : array_like + Non-centrality parameter; must satisfy ``nc >= 0``. + out : ndarray, optional + Optional output array for the function results. + + Returns + ------- + cdf : scalar or ndarray + Value of the non-central chi-squared cumulative distribution function. + + See Also + -------- + chndtrix: Noncentral Chi Squared distribution quantile + chndtridf: Inverse of `chndtr` with respect to `df` + chndtrinc: Inverse of `chndtr` with respect to `nc` + scipy.stats.ncx2: Non-central chi-squared distribution + + Notes + ----- + The noncentral chi squared distribution is also available in + `scipy.stats.ncx2`. ``scipy.stats.ncx2.cdf`` is equivalent to `chndtr`. + + This function wraps routines from the Boost Math C++ library [1]_. + + References + ---------- + .. [1] The Boost Developers. "Boost C++ Libraries". https://www.boost.org/. + + Examples + -------- + >>> import numpy as np + >>> import scipy.special as sc + + Compute the noncentral chi squared distribution CDF at one point. + + >>> x = 4.0 + >>> df = 1.0 + >>> nc = 5.0 + >>> sc.chndtr(x, df, nc) + 0.40667858759710945 + + Plot the noncentral chi squared distribution CDF for different parameters. + + >>> import matplotlib.pyplot as plt + >>> x = np.linspace(0, 40, 1000) + >>> plt.plot(x, sc.chndtr(x, 1, 5), label=r"$df=1,\ nc=5$") + >>> plt.plot(x, sc.chndtr(x, 5, 10), label=r"$df=5,\ nc=10$") + >>> plt.legend() + >>> plt.show() + + """) + +add_newdoc("chndtrix", + """ + chndtrix(p, df, nc, out=None) + + Inverse to `chndtr` vs `x`. + + Calculated using a search to find a value for `x` that produces the + desired value of `p`. + + Parameters + ---------- + p : array_like + Probability; must satisfy ``0 <= p < 1`` + df : array_like + Degrees of freedom; must satisfy ``df > 0`` + nc : array_like + Non-centrality parameter; must satisfy ``nc >= 0`` + out : ndarray, optional + Optional output array for the function results + + Returns + ------- + x : scalar or ndarray + Value so that the probability a non-central Chi square random variable + with `df` degrees of freedom and non-centrality, `nc`, is greater than + `x` equals `p`. + + See Also + -------- + chndtr : Noncentral chi-squared distribution CDF + chndtridf : inverse of `chndtr` with respect to `cdf` + chndtrinc : inverse of `chndtr` with respect to `nc` + scipy.stats.ncx2 : Non-central chi-squared distribution + + Notes + ----- + The noncentral chi squared distribution is also available in + `scipy.stats.ncx2`. ``scipy.stats.ncx2.ppf`` is equivalent to `chndtrix`. + + This function wraps routines from the Boost Math C++ library [1]_. + + References + ---------- + .. [1] The Boost Developers. "Boost C++ Libraries". https://www.boost.org/. + + Examples + -------- + >>> from scipy.special import chndtrix, chndtr + + Compute the noncentral chi squared distribution CDF at one point. + >>> x, df, nc = 3, 5, 10 + >>> p = chndtr(x, df, nc) + + `chndtrix` is the inverse of `chndtr` with respect to `x`: + + >>> chndtrix(p, df, nc) + 3.0 + + """) + +add_newdoc("chndtridf", + """ + chndtridf(x, p, nc, out=None) + + Inverse to `chndtr` vs `df`. + + Calculated using a search to find a value for `df` that produces the + desired value of `p`. + + Parameters + ---------- + x : array_like + Upper bound of the integral; must satisfy ``x >= 0`` + p : array_like + Probability; must satisfy ``0 <= p < 1`` + nc : array_like + Non-centrality parameter; must satisfy ``nc >= 0`` + out : ndarray, optional + Optional output array for the function results + + Returns + ------- + df : scalar or ndarray + Degrees of freedom + + See Also + -------- + chndtr : Noncentral chi-squared distribution CDF + chndtrix : inverse of `chndtr` with respect to `x` + chndtrinc : inverse of `chndtr` with respect to `nc` + scipy.stats.ncx2 : Non-central chi-squared distribution + + Notes + ----- + The noncentral chi squared distribution is also available in + `scipy.stats.ncx2`. + + This function wraps routines from the Boost Math C++ library [1]_. + + References + ---------- + .. [1] The Boost Developers. "Boost C++ Libraries". https://www.boost.org/. + + Examples + -------- + >>> from scipy.special import chndtridf, chndtr + + Compute the noncentral chi squared distribution CDF at one point. + + >>> x, df, nc = 3, 5, 10 + >>> p = chndtr(x, df, nc) + + `chndtridf` is the inverse of `chndtr` with respect to `df`: + + >>> chndtridf(x, p, nc) + 5.0 + + """) + +add_newdoc("chndtrinc", + """ + chndtrinc(x, df, p, out=None) + + Inverse of `chndtr` with respect to `nc`. + + Finds the non-centrality parameter `nc` such that + + .. math:: + + \\operatorname{chndtr}(x, df, nc) = p. + + Parameters + ---------- + x : array_like + Upper bound of the integral; must satisfy ``x >= 0``. + df : array_like + Degrees of freedom; must satisfy ``df > 0``. + p : array_like + Probability; must satisfy ``0 <= p < 1``. + out : ndarray, optional + Optional output array for the function results. + + Returns + ------- + nc : scalar or ndarray + Non-centrality parameter. + + See Also + -------- + chndtr : Noncentral chi-squared distribution CDF + chndtridf : Inverse of `chndtr` with respect to `df` + chndtrinc : Inverse of `chndtr` with respect to `nc` + scipy.stats.ncx2 : Non-central chi-squared distribution + + Notes + ----- + The noncentral chi squared distribution is also available in + `scipy.stats.ncx2`. + + This function wraps routines from the Boost Math C++ library [1]_. + + References + ---------- + .. [1] The Boost Developers. "Boost C++ Libraries". https://www.boost.org/. + + Examples + -------- + >>> from scipy.special import chndtrinc, chndtr + + Compute the noncentral chi squared distribution CDF at one point. + + >>> x, df, nc = 3, 5, 10 + >>> p = chndtr(x, df, nc) + + `chndtrinc` is the inverse of `chndtr` with respect to `nc`: + + >>> chndtrinc(x, df, p) + 10.0 + + """) + +add_newdoc( + "elliprc", + r""" + elliprc(x, y, out=None) + + Degenerate symmetric elliptic integral. + + The function RC is defined as [1]_ + + .. math:: + + R_{\mathrm{C}}(x, y) = + \frac{1}{2} \int_0^{+\infty} (t + x)^{-1/2} (t + y)^{-1} dt + = R_{\mathrm{F}}(x, y, y) + + Parameters + ---------- + x, y : array_like + Real or complex input parameters. `x` can be any number in the + complex plane cut along the negative real axis. `y` must be non-zero. + out : ndarray, optional + Optional output array for the function values + + Returns + ------- + R : scalar or ndarray + Value of the integral. If `y` is real and negative, the Cauchy + principal value is returned. If both of `x` and `y` are real, the + return value is real. Otherwise, the return value is complex. + + See Also + -------- + elliprf : Completely-symmetric elliptic integral of the first kind. + elliprd : Symmetric elliptic integral of the second kind. + elliprg : Completely-symmetric elliptic integral of the second kind. + elliprj : Symmetric elliptic integral of the third kind. + + Notes + ----- + RC is a degenerate case of the symmetric integral RF: ``elliprc(x, y) == + elliprf(x, y, y)``. It is an elementary function rather than an elliptic + integral. + + The code implements Carlson's algorithm based on the duplication theorems + and series expansion up to the 7th order. [2]_ + + .. versionadded:: 1.8.0 + + References + ---------- + .. [1] B. C. Carlson, ed., Chapter 19 in "Digital Library of Mathematical + Functions," NIST, US Dept. of Commerce. + https://dlmf.nist.gov/19.16.E6 + .. [2] B. C. Carlson, "Numerical computation of real or complex elliptic + integrals," Numer. Algorithm, vol. 10, no. 1, pp. 13-26, 1995. + :doi:`10.1007/BF02198293`. https://arxiv.org/abs/math/9409227 + + Examples + -------- + Basic homogeneity property: + + >>> import numpy as np + >>> from scipy.special import elliprc + + >>> x = 1.2 + 3.4j + >>> y = 5. + >>> scale = 0.3 + 0.4j + >>> elliprc(scale*x, scale*y) + (0.5484493976710874-0.4169557678995833j) + + >>> elliprc(x, y)/np.sqrt(scale) + (0.5484493976710874-0.41695576789958333j) + + When the two arguments coincide, the integral is particularly + simple: + + >>> x = 1.2 + 3.4j + >>> elliprc(x, x) + (0.4299173120614631-0.3041729818745595j) + + >>> 1/np.sqrt(x) + (0.4299173120614631-0.30417298187455954j) + + Another simple case: the first argument vanishes: + + >>> y = 1.2 + 3.4j + >>> elliprc(0, y) + (0.6753125346116815-0.47779380263880866j) + + >>> np.pi/2/np.sqrt(y) + (0.6753125346116815-0.4777938026388088j) + + When `x` and `y` are both positive, we can express + :math:`R_C(x,y)` in terms of more elementary functions. For the + case :math:`0 \le x < y`, + + >>> x = 3.2 + >>> y = 6. + >>> elliprc(x, y) + 0.44942991498453444 + + >>> np.arctan(np.sqrt((y-x)/x))/np.sqrt(y-x) + 0.44942991498453433 + + And for the case :math:`0 \le y < x`, + + >>> x = 6. + >>> y = 3.2 + >>> elliprc(x,y) + 0.4989837501576147 + + >>> np.log((np.sqrt(x)+np.sqrt(x-y))/np.sqrt(y))/np.sqrt(x-y) + 0.49898375015761476 + + """) + +add_newdoc( + "elliprd", + r""" + elliprd(x, y, z, out=None) + + Symmetric elliptic integral of the second kind. + + The function RD is defined as [1]_ + + .. math:: + + R_{\mathrm{D}}(x, y, z) = + \frac{3}{2} \int_0^{+\infty} [(t + x) (t + y)]^{-1/2} (t + z)^{-3/2} + dt + + Parameters + ---------- + x, y, z : array_like + Real or complex input parameters. `x` or `y` can be any number in the + complex plane cut along the negative real axis, but at most one of them + can be zero, while `z` must be non-zero. + out : ndarray, optional + Optional output array for the function values + + Returns + ------- + R : scalar or ndarray + Value of the integral. If all of `x`, `y`, and `z` are real, the + return value is real. Otherwise, the return value is complex. + + See Also + -------- + elliprc : Degenerate symmetric elliptic integral. + elliprf : Completely-symmetric elliptic integral of the first kind. + elliprg : Completely-symmetric elliptic integral of the second kind. + elliprj : Symmetric elliptic integral of the third kind. + + Notes + ----- + RD is a degenerate case of the elliptic integral RJ: ``elliprd(x, y, z) == + elliprj(x, y, z, z)``. + + The code implements Carlson's algorithm based on the duplication theorems + and series expansion up to the 7th order. [2]_ + + .. versionadded:: 1.8.0 + + References + ---------- + .. [1] B. C. Carlson, ed., Chapter 19 in "Digital Library of Mathematical + Functions," NIST, US Dept. of Commerce. + https://dlmf.nist.gov/19.16.E5 + .. [2] B. C. Carlson, "Numerical computation of real or complex elliptic + integrals," Numer. Algorithm, vol. 10, no. 1, pp. 13-26, 1995. + :doi:`10.1007/BF02198293`. https://arxiv.org/abs/math/9409227 + + Examples + -------- + Basic homogeneity property: + + >>> import numpy as np + >>> from scipy.special import elliprd + + >>> x = 1.2 + 3.4j + >>> y = 5. + >>> z = 6. + >>> scale = 0.3 + 0.4j + >>> elliprd(scale*x, scale*y, scale*z) + (-0.03703043835680379-0.24500934665683802j) + + >>> elliprd(x, y, z)*np.power(scale, -1.5) + (-0.0370304383568038-0.24500934665683805j) + + All three arguments coincide: + + >>> x = 1.2 + 3.4j + >>> elliprd(x, x, x) + (-0.03986825876151896-0.14051741840449586j) + + >>> np.power(x, -1.5) + (-0.03986825876151894-0.14051741840449583j) + + The so-called "second lemniscate constant": + + >>> elliprd(0, 2, 1)/3 + 0.5990701173677961 + + >>> from scipy.special import gamma + >>> gamma(0.75)**2/np.sqrt(2*np.pi) + 0.5990701173677959 + + """) + +add_newdoc( + "elliprf", + r""" + elliprf(x, y, z, out=None) + + Completely-symmetric elliptic integral of the first kind. + + The function RF is defined as [1]_ + + .. math:: + + R_{\mathrm{F}}(x, y, z) = + \frac{1}{2} \int_0^{+\infty} [(t + x) (t + y) (t + z)]^{-1/2} dt + + Parameters + ---------- + x, y, z : array_like + Real or complex input parameters. `x`, `y`, or `z` can be any number in + the complex plane cut along the negative real axis, but at most one of + them can be zero. + out : ndarray, optional + Optional output array for the function values + + Returns + ------- + R : scalar or ndarray + Value of the integral. If all of `x`, `y`, and `z` are real, the return + value is real. Otherwise, the return value is complex. + + See Also + -------- + elliprc : Degenerate symmetric integral. + elliprd : Symmetric elliptic integral of the second kind. + elliprg : Completely-symmetric elliptic integral of the second kind. + elliprj : Symmetric elliptic integral of the third kind. + + Notes + ----- + The code implements Carlson's algorithm based on the duplication theorems + and series expansion up to the 7th order (cf.: + https://dlmf.nist.gov/19.36.i) and the AGM algorithm for the complete + integral. [2]_ + + .. versionadded:: 1.8.0 + + References + ---------- + .. [1] B. C. Carlson, ed., Chapter 19 in "Digital Library of Mathematical + Functions," NIST, US Dept. of Commerce. + https://dlmf.nist.gov/19.16.E1 + .. [2] B. C. Carlson, "Numerical computation of real or complex elliptic + integrals," Numer. Algorithm, vol. 10, no. 1, pp. 13-26, 1995. + :doi:`10.1007/BF02198293`. https://arxiv.org/abs/math/9409227 + + Examples + -------- + Basic homogeneity property: + + >>> import numpy as np + >>> from scipy.special import elliprf + + >>> x = 1.2 + 3.4j + >>> y = 5. + >>> z = 6. + >>> scale = 0.3 + 0.4j + >>> elliprf(scale*x, scale*y, scale*z) + (0.5328051227278146-0.4008623567957094j) + + >>> elliprf(x, y, z)/np.sqrt(scale) + (0.5328051227278147-0.4008623567957095j) + + All three arguments coincide: + + >>> x = 1.2 + 3.4j + >>> elliprf(x, x, x) + (0.42991731206146316-0.30417298187455954j) + + >>> 1/np.sqrt(x) + (0.4299173120614631-0.30417298187455954j) + + The so-called "first lemniscate constant": + + >>> elliprf(0, 1, 2) + 1.3110287771460598 + + >>> from scipy.special import gamma + >>> gamma(0.25)**2/(4*np.sqrt(2*np.pi)) + 1.3110287771460598 + + """) + +add_newdoc( + "elliprg", + r""" + elliprg(x, y, z, out=None) + + Completely-symmetric elliptic integral of the second kind. + + The function RG is defined as [1]_ + + .. math:: + + R_{\mathrm{G}}(x, y, z) = + \frac{1}{4} \int_0^{+\infty} [(t + x) (t + y) (t + z)]^{-1/2} + \left(\frac{x}{t + x} + \frac{y}{t + y} + \frac{z}{t + z}\right) t + dt + + Parameters + ---------- + x, y, z : array_like + Real or complex input parameters. `x`, `y`, or `z` can be any number in + the complex plane cut along the negative real axis. + out : ndarray, optional + Optional output array for the function values + + Returns + ------- + R : scalar or ndarray + Value of the integral. If all of `x`, `y`, and `z` are real, the return + value is real. Otherwise, the return value is complex. + + See Also + -------- + elliprc : Degenerate symmetric integral. + elliprd : Symmetric elliptic integral of the second kind. + elliprf : Completely-symmetric elliptic integral of the first kind. + elliprj : Symmetric elliptic integral of the third kind. + + Notes + ----- + The implementation uses the relation [1]_ + + .. math:: + + 2 R_{\mathrm{G}}(x, y, z) = + z R_{\mathrm{F}}(x, y, z) - + \frac{1}{3} (x - z) (y - z) R_{\mathrm{D}}(x, y, z) + + \sqrt{\frac{x y}{z}} + + and the symmetry of `x`, `y`, `z` when at least one non-zero parameter can + be chosen as the pivot. When one of the arguments is close to zero, the AGM + method is applied instead. Other special cases are computed following Ref. + [2]_ + + .. versionadded:: 1.8.0 + + References + ---------- + .. [1] B. C. Carlson, "Numerical computation of real or complex elliptic + integrals," Numer. Algorithm, vol. 10, no. 1, pp. 13-26, 1995. + :doi:`10.1007/BF02198293`. https://arxiv.org/abs/math/9409227 + .. [2] B. C. Carlson, ed., Chapter 19 in "Digital Library of Mathematical + Functions," NIST, US Dept. of Commerce. + https://dlmf.nist.gov/19.16.E1 + https://dlmf.nist.gov/19.20.ii + + Examples + -------- + Basic homogeneity property: + + >>> import numpy as np + >>> from scipy.special import elliprg + + >>> x = 1.2 + 3.4j + >>> y = 5. + >>> z = 6. + >>> scale = 0.3 + 0.4j + >>> elliprg(scale*x, scale*y, scale*z) + (1.195936862005246+0.8470988320464167j) + + >>> elliprg(x, y, z)*np.sqrt(scale) + (1.195936862005246+0.8470988320464165j) + + Simplifications: + + >>> elliprg(0, y, y) + 1.756203682760182 + + >>> 0.25*np.pi*np.sqrt(y) + 1.7562036827601817 + + >>> elliprg(0, 0, z) + 1.224744871391589 + + >>> 0.5*np.sqrt(z) + 1.224744871391589 + + The surface area of a triaxial ellipsoid with semiaxes ``a``, ``b``, and + ``c`` is given by + + .. math:: + + S = 4 \pi a b c R_{\mathrm{G}}(1 / a^2, 1 / b^2, 1 / c^2). + + >>> def ellipsoid_area(a, b, c): + ... r = 4.0 * np.pi * a * b * c + ... return r * elliprg(1.0 / (a * a), 1.0 / (b * b), 1.0 / (c * c)) + >>> print(ellipsoid_area(1, 3, 5)) + 108.62688289491807 + """) + +add_newdoc( + "elliprj", + r""" + elliprj(x, y, z, p, out=None) + + Symmetric elliptic integral of the third kind. + + The function RJ is defined as [1]_ + + .. math:: + + R_{\mathrm{J}}(x, y, z, p) = + \frac{3}{2} \int_0^{+\infty} [(t + x) (t + y) (t + z)]^{-1/2} + (t + p)^{-1} dt + + .. warning:: + This function should be considered experimental when the inputs are + unbalanced. Check correctness with another independent implementation. + + Parameters + ---------- + x, y, z, p : array_like + Real or complex input parameters. `x`, `y`, or `z` are numbers in + the complex plane cut along the negative real axis (subject to further + constraints, see Notes), and at most one of them can be zero. `p` must + be non-zero. + out : ndarray, optional + Optional output array for the function values + + Returns + ------- + R : scalar or ndarray + Value of the integral. If all of `x`, `y`, `z`, and `p` are real, the + return value is real. Otherwise, the return value is complex. + + If `p` is real and negative, while `x`, `y`, and `z` are real, + non-negative, and at most one of them is zero, the Cauchy principal + value is returned. [1]_ [2]_ + + See Also + -------- + elliprc : Degenerate symmetric integral. + elliprd : Symmetric elliptic integral of the second kind. + elliprf : Completely-symmetric elliptic integral of the first kind. + elliprg : Completely-symmetric elliptic integral of the second kind. + + Notes + ----- + The code implements Carlson's algorithm based on the duplication theorems + and series expansion up to the 7th order. [3]_ The algorithm is slightly + different from its earlier incarnation as it appears in [1]_, in that the + call to `elliprc` (or ``atan``/``atanh``, see [4]_) is no longer needed in + the inner loop. Asymptotic approximations are used where arguments differ + widely in the order of magnitude. [5]_ + + The input values are subject to certain sufficient but not necessary + constraints when input arguments are complex. Notably, ``x``, ``y``, and + ``z`` must have non-negative real parts, unless two of them are + non-negative and complex-conjugates to each other while the other is a real + non-negative number. [1]_ If the inputs do not satisfy the sufficient + condition described in Ref. [1]_ they are rejected outright with the output + set to NaN. + + In the case where one of ``x``, ``y``, and ``z`` is equal to ``p``, the + function ``elliprd`` should be preferred because of its less restrictive + domain. + + .. versionadded:: 1.8.0 + + References + ---------- + .. [1] B. C. Carlson, "Numerical computation of real or complex elliptic + integrals," Numer. Algorithm, vol. 10, no. 1, pp. 13-26, 1995. + :doi:`10.1007/BF02198293`. https://arxiv.org/abs/math/9409227 + .. [2] B. C. Carlson, ed., Chapter 19 in "Digital Library of Mathematical + Functions," NIST, US Dept. of Commerce. + https://dlmf.nist.gov/19.20.iii + .. [3] B. C. Carlson, J. FitzSimmons, "Reduction Theorems for Elliptic + Integrands with the Square Root of Two Quadratic Factors," J. + Comput. Appl. Math., vol. 118, nos. 1-2, pp. 71-85, 2000. + :doi:`10.1016/S0377-0427(00)00282-X`. + .. [4] F. Johansson, "Numerical Evaluation of Elliptic Functions, Elliptic + Integrals and Modular Forms," in J. Blumlein, C. Schneider, P. + Paule, eds., "Elliptic Integrals, Elliptic Functions and Modular + Forms in Quantum Field Theory," pp. 269-293, 2019 (Cham, + Switzerland: Springer Nature Switzerland). + :doi:`10.1007/978-3-030-04480-0`. https://arxiv.org/abs/1806.06725 + .. [5] B. C. Carlson, J. L. Gustafson, "Asymptotic Approximations for + Symmetric Elliptic Integrals," SIAM J. Math. Anls., vol. 25, no. 2, + pp. 288-303, 1994. :doi:`10.1137/S0036141092228477`. + https://arxiv.org/abs/math/9310223 + + Examples + -------- + Basic homogeneity property: + + >>> import numpy as np + >>> from scipy.special import elliprj + + >>> x = 1.2 + 3.4j + >>> y = 5. + >>> z = 6. + >>> p = 7. + >>> scale = 0.3 - 0.4j + >>> elliprj(scale*x, scale*y, scale*z, scale*p) + (0.10834905565679157+0.19694950747103812j) + + >>> elliprj(x, y, z, p)*np.power(scale, -1.5) + (0.10834905565679556+0.19694950747103854j) + + Reduction to simpler elliptic integral: + + >>> elliprj(x, y, z, z) + (0.08288462362195129-0.028376809745123258j) + + >>> from scipy.special import elliprd + >>> elliprd(x, y, z) + (0.08288462362195136-0.028376809745123296j) + + All arguments coincide: + + >>> elliprj(x, x, x, x) + (-0.03986825876151896-0.14051741840449586j) + + >>> np.power(x, -1.5) + (-0.03986825876151894-0.14051741840449583j) + + """) + +add_newdoc("entr", + r""" + entr(x, out=None) + + Elementwise function for computing entropy. + + .. math:: \text{entr}(x) = \begin{cases} - x \log(x) & x > 0 \\ 0 & x = 0 + \\ -\infty & \text{otherwise} \end{cases} + + Parameters + ---------- + x : ndarray + Input array. + out : ndarray, optional + Optional output array for the function values + + Returns + ------- + res : scalar or ndarray + The value of the elementwise entropy function at the given points `x`. + + See Also + -------- + kl_div, rel_entr, scipy.stats.entropy + + Notes + ----- + .. versionadded:: 0.15.0 + + This function is concave. + + The origin of this function is in convex programming; see [1]_. + Given a probability distribution :math:`p_1, \ldots, p_n`, + the definition of entropy in the context of *information theory* is + + .. math:: + + \sum_{i = 1}^n \mathrm{entr}(p_i). + + To compute the latter quantity, use `scipy.stats.entropy`. + + References + ---------- + .. [1] Boyd, Stephen and Lieven Vandenberghe. *Convex optimization*. + Cambridge University Press, 2004. + :doi:`10.1017/CBO9780511804441`. + + Examples + -------- + >>> import numpy as np + >>> from scipy.special import entr + + Calculate the entropy (in nats) of a 3-outcome probability distribution + + >>> p = np.array([0.2, 0.5, 0.3]) + >>> entr(p) + array([0.32188758, 0.34657359, 0.36119184]) + >>> entr(p).sum() + 1.0296530140645737 + + """) + +add_newdoc( + "erfinv", + """ + erfinv(y, out=None) + + Inverse of the error function. + + Computes the inverse of the error function. + + In the complex domain, there is no unique complex number w satisfying + erf(w)=z. This indicates a true inverse function would be multivalued. + When the domain restricts to the real, -1 < x < 1, there is a unique real + number satisfying erf(erfinv(x)) = x. + + Parameters + ---------- + y : ndarray + Argument at which to evaluate. Domain: [-1, 1] + out : ndarray, optional + Optional output array for the function values + + Returns + ------- + erfinv : scalar or ndarray + The inverse of erf of y, element-wise + + See Also + -------- + erf : Error function of a complex argument + erfc : Complementary error function, ``1 - erf(x)`` + erfcinv : Inverse of the complementary error function + + Notes + ----- + This function wraps the ``erf_inv`` routine from the + Boost Math C++ library [1]_. + + References + ---------- + .. [1] The Boost Developers. "Boost C++ Libraries". https://www.boost.org/. + + Examples + -------- + >>> import numpy as np + >>> import matplotlib.pyplot as plt + >>> from scipy.special import erfinv, erf + + >>> erfinv(0.5) + 0.4769362762044699 + + >>> y = np.linspace(-1.0, 1.0, num=9) + >>> x = erfinv(y) + >>> x + array([ -inf, -0.81341985, -0.47693628, -0.22531206, 0. , + 0.22531206, 0.47693628, 0.81341985, inf]) + + Verify that ``erf(erfinv(y))`` is ``y``. + + >>> erf(x) + array([-1. , -0.75, -0.5 , -0.25, 0. , 0.25, 0.5 , 0.75, 1. ]) + + Plot the function: + + >>> y = np.linspace(-1, 1, 200) + >>> fig, ax = plt.subplots() + >>> ax.plot(y, erfinv(y)) + >>> ax.grid(True) + >>> ax.set_xlabel('y') + >>> ax.set_title('erfinv(y)') + >>> plt.show() + + """) + +add_newdoc( + "erfcinv", + """ + erfcinv(y, out=None) + + Inverse of the complementary error function. + + Computes the inverse of the complementary error function. + + In the complex domain, there is no unique complex number :math:`w` satisfying + :math:`\\operatorname{erfc}(w) = z`. This indicates a true inverse function + would be multivalued. + When the domain restricts to the real interval :math:`0 < x < 2`, there is + a unique real number satisfying + + .. math:: + + \\operatorname{erfc}(\\operatorname{erfcinv}(x)) = x + + It is related to the inverse of the error function by + + .. math:: + + \\operatorname{erfcinv}(1 - x) = \\operatorname{erfinv}(x) + + Parameters + ---------- + y : ndarray + Argument at which to evaluate. Domain: :math:`[0, 2]` + out : ndarray, optional + Optional output array for the function values + + Returns + ------- + erfcinv : scalar or ndarray + The inverse of :math:`\\operatorname{erfc}` of :math:`y`, element-wise + + See Also + -------- + erf : Error function + erfc : Complementary error function + erfinv : Inverse of the error function + + Examples + -------- + >>> import numpy as np + >>> import matplotlib.pyplot as plt + >>> from scipy.special import erfcinv + + >>> erfcinv(0.5) + 0.4769362762044699 + + >>> y = np.linspace(0.0, 2.0, num=11) + >>> erfcinv(y) + array([ inf, 0.9061938 , 0.59511608, 0.37080716, 0.17914345, + -0. , -0.17914345, -0.37080716, -0.59511608, -0.9061938 , + -inf]) + + Plot the function: + + >>> y = np.linspace(0, 2, 200) + >>> fig, ax = plt.subplots() + >>> ax.plot(y, erfcinv(y)) + >>> ax.grid(True) + >>> ax.set_xlabel('y') + >>> ax.set_title('erfcinv(y)') + >>> plt.show() + + """) + +add_newdoc("eval_jacobi", + r""" + eval_jacobi(n, alpha, beta, x, out=None) + + Evaluate Jacobi polynomial at a point. + + The Jacobi polynomials can be defined via the Gauss hypergeometric + function :math:`{}_2F_1` as + + .. math:: + + P_n^{(\alpha, \beta)}(x) = \frac{(\alpha + 1)_n}{\Gamma(n + 1)} + {}_2F_1(-n, 1 + \alpha + \beta + n; \alpha + 1; (1 - x)/2) + + where :math:`(\cdot)_n` is the Pochhammer symbol; see `poch`. When + :math:`n` is an integer the result is a polynomial of degree + :math:`n`. See 22.5.42 in [AS]_ or [DLMF]_ for details. + + Parameters + ---------- + n : array_like + Degree of the polynomial. If not an integer the result is + determined via the relation to the Gauss hypergeometric + function. + alpha : array_like + Parameter. + beta : array_like + Parameter. + x : array_like + Points at which to evaluate the polynomial. + out : ndarray, optional + Optional output array for the function values. + + Returns + ------- + P : scalar or ndarray + Values of the Jacobi polynomial. + + See Also + -------- + roots_jacobi : roots and quadrature weights of Jacobi polynomials + jacobi : Jacobi polynomial object + hyp2f1 : Gauss hypergeometric function + + References + ---------- + .. [AS] Milton Abramowitz and Irene A. Stegun, eds. + Handbook of Mathematical Functions with Formulas, + Graphs, and Mathematical Tables. New York: Dover, 1972. + .. [DLMF] NIST Digital Library of Mathematical Functions, + https://dlmf.nist.gov/18.5.E7 + + """) + +add_newdoc("eval_sh_jacobi", + r""" + eval_sh_jacobi(n, p, q, x, out=None) + + Evaluate shifted Jacobi polynomial at a point. + + Defined by + + .. math:: + + G_n^{(p, q)}(x) + = \binom{2n + p - 1}{n}^{-1} P_n^{(p - q, q - 1)}(2x - 1), + + where :math:`P_n^{(\cdot, \cdot)}` is the n-th Jacobi polynomial. + See 22.5.2 in [AS]_ (or equivalently [DLMF]_) for details. + + Parameters + ---------- + n : int + Degree of the polynomial. If not an integer, the result is + determined via the relation to `binom` and `eval_jacobi`. + p : float + Parameter + q : float + Parameter + out : ndarray, optional + Optional output array for the function values + + Returns + ------- + G : scalar or ndarray + Values of the shifted Jacobi polynomial. + + See Also + -------- + roots_sh_jacobi : roots and quadrature weights of shifted Jacobi + polynomials + sh_jacobi : shifted Jacobi polynomial object + eval_jacobi : evaluate Jacobi polynomials + + References + ---------- + .. [AS] Milton Abramowitz and Irene A. Stegun, eds. + Handbook of Mathematical Functions with Formulas, + Graphs, and Mathematical Tables. New York: Dover, 1972. + .. [DLMF] NIST Digital Library of Mathematical Functions, + https://dlmf.nist.gov/18.1.E2 + + """) + +add_newdoc("eval_gegenbauer", + r""" + eval_gegenbauer(n, alpha, x, out=None) + + Evaluate Gegenbauer (ultraspherical) polynomial at a point. + + The Gegenbauer polynomials can be defined via the Gauss + hypergeometric function :math:`{}_2F_1` as + + .. math:: + + C_n^{(\alpha)}(x) = \frac{(2\alpha)_n}{\Gamma(n + 1)} + {}_2F_1(-n, 2\alpha + n; \alpha + 1/2; (1 - x)/2). + + where :math:`(\cdot)_n` is the Pochhammer symbol; see `poch`. When + :math:`n` is an integer the result is a polynomial of degree + :math:`n`. See 22.5.46 in [AS]_ (or equivalently [DLMF]_) for details. + + Parameters + ---------- + n : array_like + Degree of the polynomial. If not an integer, the result is + determined via the relation to the Gauss hypergeometric + function. + alpha : array_like + Parameter. + x : array_like + Points at which to evaluate the Gegenbauer polynomial. + out : ndarray, optional + Optional output array for the function values. + + Returns + ------- + C : scalar or ndarray + Values of the Gegenbauer polynomial. + + See Also + -------- + roots_gegenbauer : roots and quadrature weights of Gegenbauer + polynomials + gegenbauer : Gegenbauer polynomial object + hyp2f1 : Gauss hypergeometric function + + References + ---------- + .. [AS] Milton Abramowitz and Irene A. Stegun, eds. + Handbook of Mathematical Functions with Formulas, + Graphs, and Mathematical Tables. New York: Dover, 1972. + .. [DLMF] NIST Digital Library of Mathematical Functions, + https://dlmf.nist.gov/18.5.E9 + + """) + +add_newdoc("eval_chebyt", + r""" + eval_chebyt(n, x, out=None) + + Evaluate Chebyshev polynomial of the first kind at a point. + + The Chebyshev polynomials of the first kind can be defined via the + Gauss hypergeometric function :math:`{}_2F_1` as + + .. math:: + + T_n(x) = {}_2F_1(n, -n; 1/2; (1 - x)/2). + + When :math:`n` is an integer the result is a polynomial of degree + :math:`n`. See 22.5.47 in [AS]_ (or equivalently [DLMF]_) for details. + + Parameters + ---------- + n : array_like + Degree of the polynomial. If not an integer, the result is + determined via the relation to the Gauss hypergeometric + function. + x : array_like + Points at which to evaluate the Chebyshev polynomial + out : ndarray, optional + Optional output array for the function values + + Returns + ------- + T : scalar or ndarray + Values of the Chebyshev polynomial + + See Also + -------- + roots_chebyt : roots and quadrature weights of Chebyshev + polynomials of the first kind + chebyu : Chebychev polynomial object + eval_chebyu : evaluate Chebyshev polynomials of the second kind + hyp2f1 : Gauss hypergeometric function + numpy.polynomial.chebyshev.Chebyshev : Chebyshev series + + Notes + ----- + This routine is numerically stable for `x` in ``[-1, 1]`` at least + up to order ``10000``. + + References + ---------- + .. [AS] Milton Abramowitz and Irene A. Stegun, eds. + Handbook of Mathematical Functions with Formulas, + Graphs, and Mathematical Tables. New York: Dover, 1972. + .. [DLMF] NIST Digital Library of Mathematical Functions, + https://dlmf.nist.gov/18.5.E11_2 + + """) + +add_newdoc("eval_chebyu", + r""" + eval_chebyu(n, x, out=None) + + Evaluate Chebyshev polynomial of the second kind at a point. + + The Chebyshev polynomials of the second kind can be defined via + the Gauss hypergeometric function :math:`{}_2F_1` as + + .. math:: + + U_n(x) = (n + 1) {}_2F_1(-n, n + 2; 3/2; (1 - x)/2). + + When :math:`n` is an integer the result is a polynomial of degree + :math:`n`. See 22.5.48 in [AS]_ (or equivalently [DLMF]_) for details. + + Parameters + ---------- + n : array_like + Degree of the polynomial. If not an integer, the result is + determined via the relation to the Gauss hypergeometric + function. + x : array_like + Points at which to evaluate the Chebyshev polynomial + out : ndarray, optional + Optional output array for the function values + + Returns + ------- + U : scalar or ndarray + Values of the Chebyshev polynomial + + See Also + -------- + roots_chebyu : roots and quadrature weights of Chebyshev + polynomials of the second kind + chebyu : Chebyshev polynomial object + eval_chebyt : evaluate Chebyshev polynomials of the first kind + hyp2f1 : Gauss hypergeometric function + + References + ---------- + .. [AS] Milton Abramowitz and Irene A. Stegun, eds. + Handbook of Mathematical Functions with Formulas, + Graphs, and Mathematical Tables. New York: Dover, 1972. + .. [DLMF] NIST Digital Library of Mathematical Functions, + https://dlmf.nist.gov/18.5.E11_4 + + """) + +add_newdoc("eval_chebys", + r""" + eval_chebys(n, x, out=None) + + Evaluate Chebyshev polynomial of the second kind on [-2, 2] at a + point. + + These polynomials are defined as + + .. math:: + + S_n(x) = U_n(x/2) + + where :math:`U_n` is a Chebyshev polynomial of the second kind. + See 22.5.13 in [AS]_ (or equivalently [DLMF]_) for details. + + Parameters + ---------- + n : array_like + Degree of the polynomial. If not an integer, the result is + determined via the relation to `eval_chebyu`. + x : array_like + Points at which to evaluate the Chebyshev polynomial + out : ndarray, optional + Optional output array for the function values + + Returns + ------- + S : scalar or ndarray + Values of the Chebyshev polynomial + + See Also + -------- + roots_chebys : roots and quadrature weights of Chebyshev + polynomials of the second kind on [-2, 2] + chebys : Chebyshev polynomial object + eval_chebyu : evaluate Chebyshev polynomials of the second kind + + References + ---------- + .. [AS] Milton Abramowitz and Irene A. Stegun, eds. + Handbook of Mathematical Functions with Formulas, + Graphs, and Mathematical Tables. New York: Dover, 1972. + .. [DLMF] NIST Digital Library of Mathematical Functions, + https://dlmf.nist.gov/18.1.E3 + + Examples + -------- + >>> import numpy as np + >>> import scipy.special as sc + + They are a scaled version of the Chebyshev polynomials of the + second kind. + + >>> x = np.linspace(-2, 2, 6) + >>> sc.eval_chebys(3, x) + array([-4. , 0.672, 0.736, -0.736, -0.672, 4. ]) + >>> sc.eval_chebyu(3, x / 2) + array([-4. , 0.672, 0.736, -0.736, -0.672, 4. ]) + + """) + +add_newdoc("eval_chebyc", + r""" + eval_chebyc(n, x, out=None) + + Evaluate Chebyshev polynomial of the first kind on [-2, 2] at a + point. + + These polynomials are defined as + + .. math:: + + C_n(x) = 2 T_n(x/2) + + where :math:`T_n` is a Chebyshev polynomial of the first kind. See + 22.5.11 in [AS]_ (or equivalently [DLMF]_) for details. + + Parameters + ---------- + n : array_like + Degree of the polynomial. If not an integer, the result is + determined via the relation to `eval_chebyt`. + x : array_like + Points at which to evaluate the Chebyshev polynomial + out : ndarray, optional + Optional output array for the function values + + Returns + ------- + C : scalar or ndarray + Values of the Chebyshev polynomial + + See Also + -------- + roots_chebyc : roots and quadrature weights of Chebyshev + polynomials of the first kind on [-2, 2] + chebyc : Chebyshev polynomial object + numpy.polynomial.chebyshev.Chebyshev : Chebyshev series + eval_chebyt : evaluate Chebycshev polynomials of the first kind + + References + ---------- + .. [AS] Milton Abramowitz and Irene A. Stegun, eds. + Handbook of Mathematical Functions with Formulas, + Graphs, and Mathematical Tables. New York: Dover, 1972. + .. [DLMF] NIST Digital Library of Mathematical Functions, + https://dlmf.nist.gov/18.1.E3 + + Examples + -------- + >>> import numpy as np + >>> import scipy.special as sc + + They are a scaled version of the Chebyshev polynomials of the + first kind. + + >>> x = np.linspace(-2, 2, 6) + >>> sc.eval_chebyc(3, x) + array([-2. , 1.872, 1.136, -1.136, -1.872, 2. ]) + >>> 2 * sc.eval_chebyt(3, x / 2) + array([-2. , 1.872, 1.136, -1.136, -1.872, 2. ]) + + """) + +add_newdoc("eval_sh_chebyt", + r""" + eval_sh_chebyt(n, x, out=None) + + Evaluate shifted Chebyshev polynomial of the first kind at a + point. + + These polynomials are defined as + + .. math:: + + T_n^*(x) = T_n(2x - 1) + + where :math:`T_n` is a Chebyshev polynomial of the first kind. See + 22.5.14 in [AS]_ (or equivalently [DLMF]_) for details. + + Parameters + ---------- + n : array_like + Degree of the polynomial. If not an integer, the result is + determined via the relation to `eval_chebyt`. + x : array_like + Points at which to evaluate the shifted Chebyshev polynomial + out : ndarray, optional + Optional output array for the function values + + Returns + ------- + T : scalar or ndarray + Values of the shifted Chebyshev polynomial + + See Also + -------- + roots_sh_chebyt : roots and quadrature weights of shifted + Chebyshev polynomials of the first kind + sh_chebyt : shifted Chebyshev polynomial object + eval_chebyt : evaluate Chebyshev polynomials of the first kind + numpy.polynomial.chebyshev.Chebyshev : Chebyshev series + + References + ---------- + .. [AS] Milton Abramowitz and Irene A. Stegun, eds. + Handbook of Mathematical Functions with Formulas, + Graphs, and Mathematical Tables. New York: Dover, 1972. + .. [DLMF] NIST Digital Library of Mathematical Functions, + https://dlmf.nist.gov/18.7.E7 + + """) + +add_newdoc("eval_sh_chebyu", + r""" + eval_sh_chebyu(n, x, out=None) + + Evaluate shifted Chebyshev polynomial of the second kind at a + point. + + These polynomials are defined as + + .. math:: + + U_n^*(x) = U_n(2x - 1) + + where :math:`U_n` is a Chebyshev polynomial of the first kind. See + 22.5.15 in [AS]_ (or equivalently [DLMF]_) for details. + + Parameters + ---------- + n : array_like + Degree of the polynomial. If not an integer, the result is + determined via the relation to `eval_chebyu`. + x : array_like + Points at which to evaluate the shifted Chebyshev polynomial + out : ndarray, optional + Optional output array for the function values + + Returns + ------- + U : scalar or ndarray + Values of the shifted Chebyshev polynomial + + See Also + -------- + roots_sh_chebyu : roots and quadrature weights of shifted + Chebychev polynomials of the second kind + sh_chebyu : shifted Chebyshev polynomial object + eval_chebyu : evaluate Chebyshev polynomials of the second kind + + References + ---------- + .. [AS] Milton Abramowitz and Irene A. Stegun, eds. + Handbook of Mathematical Functions with Formulas, + Graphs, and Mathematical Tables. New York: Dover, 1972. + .. [DLMF] NIST Digital Library of Mathematical Functions, + https://dlmf.nist.gov/18.7.E8 + + """) + +add_newdoc("eval_legendre", + r""" + eval_legendre(n, x, out=None) + + Evaluate Legendre polynomial at a point. + + The Legendre polynomials can be defined via the Gauss + hypergeometric function :math:`{}_2F_1` as + + .. math:: + + P_n(x) = {}_2F_1(-n, n + 1; 1; (1 - x)/2). + + When :math:`n` is an integer the result is a polynomial of degree + :math:`n`. See 22.5.49 in [AS]_ (or equivalently [DLMF]_) for details. + + Parameters + ---------- + n : array_like + Degree of the polynomial. If not an integer, the result is + determined via the relation to the Gauss hypergeometric + function. + x : array_like + Points at which to evaluate the Legendre polynomial + out : ndarray, optional + Optional output array for the function values + + Returns + ------- + P : scalar or ndarray + Values of the Legendre polynomial + + See Also + -------- + roots_legendre : roots and quadrature weights of Legendre + polynomials + legendre : Legendre polynomial object + hyp2f1 : Gauss hypergeometric function + numpy.polynomial.legendre.Legendre : Legendre series + + References + ---------- + .. [AS] Milton Abramowitz and Irene A. Stegun, eds. + Handbook of Mathematical Functions with Formulas, + Graphs, and Mathematical Tables. New York: Dover, 1972. + .. [DLMF] NIST Digital Library of Mathematical Functions, + https://dlmf.nist.gov/15.9.E7 + + Examples + -------- + >>> import numpy as np + >>> from scipy.special import eval_legendre + + Evaluate the zero-order Legendre polynomial at x = 0 + + >>> eval_legendre(0, 0) + 1.0 + + Evaluate the first-order Legendre polynomial between -1 and 1 + + >>> X = np.linspace(-1, 1, 5) # Domain of Legendre polynomials + >>> eval_legendre(1, X) + array([-1. , -0.5, 0. , 0.5, 1. ]) + + Evaluate Legendre polynomials of order 0 through 4 at x = 0 + + >>> N = range(0, 5) + >>> eval_legendre(N, 0) + array([ 1. , 0. , -0.5 , 0. , 0.375]) + + Plot Legendre polynomials of order 0 through 4 + + >>> X = np.linspace(-1, 1) + + >>> import matplotlib.pyplot as plt + >>> for n in range(0, 5): + ... y = eval_legendre(n, X) + ... plt.plot(X, y, label=r'$P_{}(x)$'.format(n)) + + >>> plt.title("Legendre Polynomials") + >>> plt.xlabel("x") + >>> plt.ylabel(r'$P_n(x)$') + >>> plt.legend(loc='lower right') + >>> plt.show() + + """) + +add_newdoc("eval_sh_legendre", + r""" + eval_sh_legendre(n, x, out=None) + + Evaluate shifted Legendre polynomial at a point. + + These polynomials are defined as + + .. math:: + + P_n^*(x) = P_n(2x - 1) + + where :math:`P_n` is a Legendre polynomial. See 2.2.11 in [AS]_ + or [DLMF]_ for details. + + Parameters + ---------- + n : array_like + Degree of the polynomial. If not an integer, the value is + determined via the relation to `eval_legendre`. + x : array_like + Points at which to evaluate the shifted Legendre polynomial + out : ndarray, optional + Optional output array for the function values + + Returns + ------- + P : scalar or ndarray + Values of the shifted Legendre polynomial + + See Also + -------- + roots_sh_legendre : roots and quadrature weights of shifted + Legendre polynomials + sh_legendre : shifted Legendre polynomial object + eval_legendre : evaluate Legendre polynomials + numpy.polynomial.legendre.Legendre : Legendre series + + References + ---------- + .. [AS] Milton Abramowitz and Irene A. Stegun, eds. + Handbook of Mathematical Functions with Formulas, + Graphs, and Mathematical Tables. New York: Dover, 1972. + .. [DLMF] NIST Digital Library of Mathematical Functions, + https://dlmf.nist.gov/18.7.E10 + + """) + +add_newdoc("eval_genlaguerre", + r""" + eval_genlaguerre(n, alpha, x, out=None) + + Evaluate generalized Laguerre polynomial at a point. + + The generalized Laguerre polynomials can be defined via the + confluent hypergeometric function :math:`{}_1F_1` as + + .. math:: + + L_n^{(\alpha)}(x) = \binom{n + \alpha}{n} + {}_1F_1(-n, \alpha + 1, x). + + When :math:`n` is an integer the result is a polynomial of degree + :math:`n`. See 22.5.54 in [AS]_ or [DLMF]_ for details. The Laguerre + polynomials are the special case where :math:`\alpha = 0`. + + Parameters + ---------- + n : array_like + Degree of the polynomial. If not an integer, the result is + determined via the relation to the confluent hypergeometric + function. + alpha : array_like + Parameter; must have ``alpha > -1`` + x : array_like + Points at which to evaluate the generalized Laguerre + polynomial + out : ndarray, optional + Optional output array for the function values + + Returns + ------- + L : scalar or ndarray + Values of the generalized Laguerre polynomial + + See Also + -------- + roots_genlaguerre : roots and quadrature weights of generalized + Laguerre polynomials + genlaguerre : generalized Laguerre polynomial object + hyp1f1 : confluent hypergeometric function + eval_laguerre : evaluate Laguerre polynomials + + References + ---------- + .. [AS] Milton Abramowitz and Irene A. Stegun, eds. + Handbook of Mathematical Functions with Formulas, + Graphs, and Mathematical Tables. New York: Dover, 1972. + .. [DLMF] NIST Digital Library of Mathematical Functions, + https://dlmf.nist.gov/18.5.E12 + + """) + +add_newdoc("eval_laguerre", + r""" + eval_laguerre(n, x, out=None) + + Evaluate Laguerre polynomial at a point. + + The Laguerre polynomials can be defined via the confluent + hypergeometric function :math:`{}_1F_1` as + + .. math:: + + L_n(x) = {}_1F_1(-n, 1, x). + + See 22.5.16 and 22.5.54 in [AS]_ (or equivalently [DLMF1]_ and [DLMF2]_) + for details. When :math:`n` is an integer the result is a polynomial + of degree :math:`n`. + + Parameters + ---------- + n : array_like + Degree of the polynomial. If not an integer the result is + determined via the relation to the confluent hypergeometric + function. + x : array_like + Points at which to evaluate the Laguerre polynomial + out : ndarray, optional + Optional output array for the function values + + Returns + ------- + L : scalar or ndarray + Values of the Laguerre polynomial + + See Also + -------- + roots_laguerre : roots and quadrature weights of Laguerre + polynomials + laguerre : Laguerre polynomial object + numpy.polynomial.laguerre.Laguerre : Laguerre series + eval_genlaguerre : evaluate generalized Laguerre polynomials + + References + ---------- + .. [AS] Milton Abramowitz and Irene A. Stegun, eds. + Handbook of Mathematical Functions with Formulas, + Graphs, and Mathematical Tables. New York: Dover, 1972. + .. [DLMF1] NIST Digital Library of Mathematical Functions, + https://dlmf.nist.gov/18.1#I1.ix7.p1 + .. [DLMF2] NIST Digital Library of Mathematical Functions, + https://dlmf.nist.gov/18.5.E12 + + """) + +add_newdoc("eval_hermite", + r""" + eval_hermite(n, x, out=None) + + Evaluate physicist's Hermite polynomial at a point. + + Defined by + + .. math:: + + H_n(x) = (-1)^n e^{x^2} \frac{d^n}{dx^n} e^{-x^2}; + + :math:`H_n` is a polynomial of degree :math:`n`. See 22.11.7 in + [AS]_ or [DLMF]_ for details. + + Parameters + ---------- + n : array_like + Degree of the polynomial + x : array_like + Points at which to evaluate the Hermite polynomial + out : ndarray, optional + Optional output array for the function values + + Returns + ------- + H : scalar or ndarray + Values of the Hermite polynomial + + See Also + -------- + roots_hermite : roots and quadrature weights of physicist's + Hermite polynomials + hermite : physicist's Hermite polynomial object + numpy.polynomial.hermite.Hermite : Physicist's Hermite series + eval_hermitenorm : evaluate Probabilist's Hermite polynomials + + References + ---------- + .. [AS] Milton Abramowitz and Irene A. Stegun, eds. + Handbook of Mathematical Functions with Formulas, + Graphs, and Mathematical Tables. New York: Dover, 1972. + .. [DLMF] NIST Digital Library of Mathematical Functions, + https://dlmf.nist.gov/18.5.T1 + + """) + +add_newdoc("eval_hermitenorm", + r""" + eval_hermitenorm(n, x, out=None) + + Evaluate probabilist's (normalized) Hermite polynomial at a + point. + + Defined by + + .. math:: + + He_n(x) = (-1)^n e^{x^2/2} \frac{d^n}{dx^n} e^{-x^2/2}; + + :math:`He_n` is a polynomial of degree :math:`n`. See 22.11.8 in + [AS]_ or [DLMF]_ for details. + + Parameters + ---------- + n : array_like + Degree of the polynomial + x : array_like + Points at which to evaluate the Hermite polynomial + out : ndarray, optional + Optional output array for the function values + + Returns + ------- + He : scalar or ndarray + Values of the Hermite polynomial + + See Also + -------- + roots_hermitenorm : roots and quadrature weights of probabilist's + Hermite polynomials + hermitenorm : probabilist's Hermite polynomial object + numpy.polynomial.hermite_e.HermiteE : Probabilist's Hermite series + eval_hermite : evaluate physicist's Hermite polynomials + + References + ---------- + .. [AS] Milton Abramowitz and Irene A. Stegun, eds. + Handbook of Mathematical Functions with Formulas, + Graphs, and Mathematical Tables. New York: Dover, 1972. + .. [DLMF] NIST Digital Library of Mathematical Functions, + https://dlmf.nist.gov/18.5.T1 + + """) + +add_newdoc("expn", + r""" + expn(n, x, out=None) + + Generalized exponential integral En. + + For integer :math:`n \geq 0` and real :math:`x \geq 0` the + generalized exponential integral is defined as [DLMF]_ + + .. math:: + + E_n(x) = x^{n - 1} \int_x^\infty \frac{e^{-t}}{t^n} dt. + + Parameters + ---------- + n : array_like + Non-negative integers + x : array_like + Real argument + out : ndarray, optional + Optional output array for the function results + + Returns + ------- + scalar or ndarray + Values of the generalized exponential integral + + See Also + -------- + exp1 : special case of :math:`E_n` for :math:`n = 1` + expi : related to :math:`E_n` when :math:`n = 1` + + References + ---------- + .. [DLMF] Digital Library of Mathematical Functions, 8.19.2 + https://dlmf.nist.gov/8.19#E2 + + Examples + -------- + >>> import numpy as np + >>> import scipy.special as sc + + Its domain is nonnegative n and x. + + >>> sc.expn(-1, 1.0), sc.expn(1, -1.0) + (nan, nan) + + It has a pole at ``x = 0`` for ``n = 1, 2``; for larger ``n`` it + is equal to ``1 / (n - 1)``. + + >>> sc.expn([0, 1, 2, 3, 4], 0) + array([ inf, inf, 1. , 0.5 , 0.33333333]) + + For n equal to 0 it reduces to ``exp(-x) / x``. + + >>> x = np.array([1, 2, 3, 4]) + >>> sc.expn(0, x) + array([0.36787944, 0.06766764, 0.01659569, 0.00457891]) + >>> np.exp(-x) / x + array([0.36787944, 0.06766764, 0.01659569, 0.00457891]) + + For n equal to 1 it reduces to `exp1`. + + >>> sc.expn(1, x) + array([0.21938393, 0.04890051, 0.01304838, 0.00377935]) + >>> sc.exp1(x) + array([0.21938393, 0.04890051, 0.01304838, 0.00377935]) + + """) + +add_newdoc("fdtr", + r""" + fdtr(dfn, dfd, x, out=None) + + F cumulative distribution function. + + Returns the value of the cumulative distribution function of the + F-distribution, also known as Snedecor's F-distribution or the + Fisher-Snedecor distribution. + + The F-distribution with parameters :math:`d_n` and :math:`d_d` is the + distribution of the random variable, + + .. math:: + X = \frac{U_n/d_n}{U_d/d_d}, + + where :math:`U_n` and :math:`U_d` are random variables distributed + :math:`\chi^2`, with :math:`d_n` and :math:`d_d` degrees of freedom, + respectively. + + Parameters + ---------- + dfn : array_like + First parameter (positive float). + dfd : array_like + Second parameter (positive float). + x : array_like + Argument (nonnegative float). + out : ndarray, optional + Optional output array for the function values + + Returns + ------- + y : scalar or ndarray + The CDF of the F-distribution with parameters `dfn` and `dfd` at `x`. + + See Also + -------- + fdtrc : F distribution survival function + fdtri : F distribution inverse cumulative distribution + scipy.stats.f : F distribution + + Notes + ----- + The regularized incomplete beta function is used, according to the + formula, + + .. math:: + F(d_n, d_d; x) = I_{xd_n/(d_d + xd_n)}(d_n/2, d_d/2). + + Wrapper for a routine from the Boost Math C++ library [1]_. The + F distribution is also available as `scipy.stats.f`. Calling + `fdtr` directly can improve performance compared to the ``cdf`` + method of `scipy.stats.f` (see last example below). + + References + ---------- + .. [1] The Boost Developers. "Boost C++ Libraries". https://www.boost.org/. + + + Examples + -------- + Calculate the function for ``dfn=1`` and ``dfd=2`` at ``x=1``. + + >>> import numpy as np + >>> from scipy.special import fdtr + >>> fdtr(1, 2, 1) + 0.5773502691896258 + + Calculate the function at several points by providing a NumPy array for + `x`. + + >>> x = np.array([0.5, 2., 3.]) + >>> fdtr(1, 2, x) + array([0.4472136 , 0.70710678, 0.77459667]) + + Plot the function for several parameter sets. + + >>> import matplotlib.pyplot as plt + >>> dfn_parameters = [1, 5, 10, 50] + >>> dfd_parameters = [1, 1, 2, 3] + >>> linestyles = ['solid', 'dashed', 'dotted', 'dashdot'] + >>> parameters_list = list(zip(dfn_parameters, dfd_parameters, + ... linestyles)) + >>> x = np.linspace(0, 30, 1000) + >>> fig, ax = plt.subplots() + >>> for parameter_set in parameters_list: + ... dfn, dfd, style = parameter_set + ... fdtr_vals = fdtr(dfn, dfd, x) + ... ax.plot(x, fdtr_vals, label=rf"$d_n={dfn},\, d_d={dfd}$", + ... ls=style) + >>> ax.legend() + >>> ax.set_xlabel("$x$") + >>> ax.set_title("F distribution cumulative distribution function") + >>> plt.show() + + The F distribution is also available as `scipy.stats.f`. Using `fdtr` + directly can be much faster than calling the ``cdf`` method of + `scipy.stats.f`, especially for small arrays or individual values. + To get the same results one must use the following parametrization: + ``stats.f(dfn, dfd).cdf(x)=fdtr(dfn, dfd, x)``. + + >>> from scipy.stats import f + >>> dfn, dfd = 1, 2 + >>> x = 1 + >>> fdtr_res = fdtr(dfn, dfd, x) # this will often be faster than below + >>> f_dist_res = f(dfn, dfd).cdf(x) + >>> fdtr_res == f_dist_res # test that results are equal + True + """) + +add_newdoc("fdtrc", + r""" + fdtrc(dfn, dfd, x, out=None) + + F survival function. + + Returns the complemented F-distribution function (the integral of the + density from `x` to infinity). + + Parameters + ---------- + dfn : array_like + First parameter (positive float). + dfd : array_like + Second parameter (positive float). + x : array_like + Argument (nonnegative float). + out : ndarray, optional + Optional output array for the function values + + Returns + ------- + y : scalar or ndarray + The complemented F-distribution function with parameters `dfn` and + `dfd` at `x`. + + See Also + -------- + fdtr : F distribution cumulative distribution function + fdtri : F distribution inverse cumulative distribution function + scipy.stats.f : F distribution + + Notes + ----- + The regularized incomplete beta function is used, according to the + formula, + + .. math:: + F(d_n, d_d; x) = I_{d_d/(d_d + xd_n)}(d_d/2, d_n/2). + + Wrapper for a routine from the Boost Math C++ library [1]_. The + F distribution is also available as `scipy.stats.f`. Calling + `fdtrc` directly can improve performance compared to the ``sf`` + method of `scipy.stats.f` (see last example below). + + References + ---------- + .. [1] The Boost Developers. "Boost C++ Libraries". https://www.boost.org/. + + Examples + -------- + Calculate the function for ``dfn=1`` and ``dfd=2`` at ``x=1``. + + >>> import numpy as np + >>> from scipy.special import fdtrc + >>> fdtrc(1, 2, 1) + 0.42264973081037427 + + Calculate the function at several points by providing a NumPy array for + `x`. + + >>> x = np.array([0.5, 2., 3.]) + >>> fdtrc(1, 2, x) + array([0.5527864 , 0.29289322, 0.22540333]) + + Plot the function for several parameter sets. + + >>> import matplotlib.pyplot as plt + >>> dfn_parameters = [1, 5, 10, 50] + >>> dfd_parameters = [1, 1, 2, 3] + >>> linestyles = ['solid', 'dashed', 'dotted', 'dashdot'] + >>> parameters_list = list(zip(dfn_parameters, dfd_parameters, + ... linestyles)) + >>> x = np.linspace(0, 30, 1000) + >>> fig, ax = plt.subplots() + >>> for parameter_set in parameters_list: + ... dfn, dfd, style = parameter_set + ... fdtrc_vals = fdtrc(dfn, dfd, x) + ... ax.plot(x, fdtrc_vals, label=rf"$d_n={dfn},\, d_d={dfd}$", + ... ls=style) + >>> ax.legend() + >>> ax.set_xlabel("$x$") + >>> ax.set_title("F distribution survival function") + >>> plt.show() + + The F distribution is also available as `scipy.stats.f`. Using `fdtrc` + directly can be much faster than calling the ``sf`` method of + `scipy.stats.f`, especially for small arrays or individual values. + To get the same results one must use the following parametrization: + ``stats.f(dfn, dfd).sf(x)=fdtrc(dfn, dfd, x)``. + + >>> from scipy.stats import f + >>> dfn, dfd = 1, 2 + >>> x = 1 + >>> fdtrc_res = fdtrc(dfn, dfd, x) # this will often be faster than below + >>> f_dist_res = f(dfn, dfd).sf(x) + >>> f_dist_res == fdtrc_res # test that results are equal + True + """) + +add_newdoc("fdtri", + r""" + fdtri(dfn, dfd, p, out=None) + + The `p`-th quantile of the F-distribution. + + This function is the inverse of the F-distribution CDF, `fdtr`, returning + the `x` such that `fdtr(dfn, dfd, x) = p`. + + Parameters + ---------- + dfn : array_like + First parameter (positive float). + dfd : array_like + Second parameter (positive float). + p : array_like + Cumulative probability, in [0, 1]. + out : ndarray, optional + Optional output array for the function values + + Returns + ------- + x : scalar or ndarray + The quantile corresponding to `p`. + + See Also + -------- + fdtr : F distribution cumulative distribution function + fdtrc : F distribution survival function + scipy.stats.f : F distribution + + Notes + ----- + Wrapper for a routine from the Boost Math C++ library [1]_. The + F distribution is also available as `scipy.stats.f`. Calling + `fdtri` directly can improve performance compared to the ``ppf`` + method of `scipy.stats.f` (see last example below). + + References + ---------- + .. [1] The Boost Developers. "Boost C++ Libraries". https://www.boost.org/. + + Examples + -------- + `fdtri` represents the inverse of the F distribution CDF which is + available as `fdtr`. Here, we calculate the CDF for ``df1=1``, ``df2=2`` + at ``x=3``. `fdtri` then returns ``3`` given the same values for `df1`, + `df2` and the computed CDF value. + + >>> import numpy as np + >>> from scipy.special import fdtri, fdtr + >>> df1, df2 = 1, 2 + >>> x = 3 + >>> cdf_value = fdtr(df1, df2, x) + >>> fdtri(df1, df2, cdf_value) + 3.000000000000006 + + Calculate the function at several points by providing a NumPy array for + `x`. + + >>> x = np.array([0.1, 0.4, 0.7]) + >>> fdtri(1, 2, x) + array([0.02020202, 0.38095238, 1.92156863]) + + Plot the function for several parameter sets. + + >>> import matplotlib.pyplot as plt + >>> dfn_parameters = [50, 10, 1, 50] + >>> dfd_parameters = [0.5, 1, 1, 5] + >>> linestyles = ['solid', 'dashed', 'dotted', 'dashdot'] + >>> parameters_list = list(zip(dfn_parameters, dfd_parameters, + ... linestyles)) + >>> x = np.linspace(0, 1, 1000) + >>> fig, ax = plt.subplots() + >>> for parameter_set in parameters_list: + ... dfn, dfd, style = parameter_set + ... fdtri_vals = fdtri(dfn, dfd, x) + ... ax.plot(x, fdtri_vals, label=rf"$d_n={dfn},\, d_d={dfd}$", + ... ls=style) + >>> ax.legend() + >>> ax.set_xlabel("$x$") + >>> title = "F distribution inverse cumulative distribution function" + >>> ax.set_title(title) + >>> ax.set_ylim(0, 30) + >>> plt.show() + + The F distribution is also available as `scipy.stats.f`. Using `fdtri` + directly can be much faster than calling the ``ppf`` method of + `scipy.stats.f`, especially for small arrays or individual values. + To get the same results one must use the following parametrization: + ``stats.f(dfn, dfd).ppf(x)=fdtri(dfn, dfd, x)``. + + >>> from scipy.stats import f + >>> dfn, dfd = 1, 2 + >>> x = 0.7 + >>> fdtri_res = fdtri(dfn, dfd, x) # this will often be faster than below + >>> f_dist_res = f(dfn, dfd).ppf(x) + >>> f_dist_res == fdtri_res # test that results are equal + True + """) + +add_newdoc("fdtridfd", + """ + fdtridfd(dfn, p, x, out=None) + + Inverse to `fdtr` vs dfd. + + Finds the F density argument dfd such that ``fdtr(dfn, dfd, x) == p``. + + Parameters + ---------- + dfn : array_like + First parameter (positive float). + p : array_like + Cumulative probability, in [0, 1]. + x : array_like + Argument (nonnegative float). + out : ndarray, optional + Optional output array for the function values + + Returns + ------- + dfd : scalar or ndarray + `dfd` such that ``fdtr(dfn, dfd, x) == p``. + + See Also + -------- + fdtr : F distribution cumulative distribution function + fdtrc : F distribution survival function + fdtri : F distribution quantile function + scipy.stats.f : F distribution + + Examples + -------- + Compute the F distribution cumulative distribution function for one + parameter set. + + >>> from scipy.special import fdtridfd, fdtr + >>> dfn, dfd, x = 10, 5, 2 + >>> cdf_value = fdtr(dfn, dfd, x) + >>> cdf_value + 0.7700248806501017 + + Verify that `fdtridfd` recovers the original value for `dfd`: + + >>> fdtridfd(dfn, cdf_value, x) + 5.0 + """) + +''' +commented out as fdtridfn seems to have bugs and is not in functions.json +see: https://github.com/scipy/scipy/pull/15622#discussion_r811440983 + +add_newdoc( + "fdtridfn", + """ + fdtridfn(p, dfd, x, out=None) + + Inverse to `fdtr` vs dfn + + finds the F density argument dfn such that ``fdtr(dfn, dfd, x) == p``. + + + Parameters + ---------- + p : array_like + Cumulative probability, in [0, 1]. + dfd : array_like + Second parameter (positive float). + x : array_like + Argument (nonnegative float). + out : ndarray, optional + Optional output array for the function values + + Returns + ------- + dfn : scalar or ndarray + `dfn` such that ``fdtr(dfn, dfd, x) == p``. + + See Also + -------- + fdtr, fdtrc, fdtri, fdtridfd + + + """) +''' + +add_newdoc("gdtr", + r""" + gdtr(a, b, x, out=None) + + Gamma distribution cumulative distribution function. + + Returns the integral from zero to `x` of the gamma probability density + function, + + .. math:: + + F(x) = \int_0^x \frac{a^b}{\Gamma(b)} t^{b-1} e^{-at}\,dt, + + where :math:`\Gamma` is the gamma function. + + Parameters + ---------- + a : array_like + The rate parameter of the gamma distribution, sometimes denoted + :math:`\beta` (float). It is also the reciprocal of the scale + parameter :math:`\theta`. + b : array_like + The shape parameter of the gamma distribution, sometimes denoted + :math:`\alpha` (float). + x : array_like + The quantile (upper limit of integration; float). + out : ndarray, optional + Optional output array for the function values + + Returns + ------- + scalar or ndarray + The CDF of the gamma distribution with parameters `a` and `b` + evaluated at `x`. + + See Also + -------- + gdtrc : 1 - CDF of the gamma distribution. + scipy.stats.gamma: Gamma distribution + + Notes + ----- + The evaluation is carried out using the relation to the incomplete gamma + integral (regularized gamma function). + + Wrapper for the Cephes [1]_ routine `gdtr`. Calling `gdtr` directly can + improve performance compared to the ``cdf`` method of `scipy.stats.gamma` + (see last example below). + + References + ---------- + .. [1] Cephes Mathematical Functions Library, + https://netlib.org/cephes/ + + Examples + -------- + Compute the function for ``a=1``, ``b=2`` at ``x=5``. + + >>> import numpy as np + >>> from scipy.special import gdtr + >>> import matplotlib.pyplot as plt + >>> gdtr(1., 2., 5.) + 0.9595723180054873 + + Compute the function for ``a=1`` and ``b=2`` at several points by + providing a NumPy array for `x`. + + >>> xvalues = np.array([1., 2., 3., 4]) + >>> gdtr(1., 1., xvalues) + array([0.63212056, 0.86466472, 0.95021293, 0.98168436]) + + `gdtr` can evaluate different parameter sets by providing arrays with + broadcasting compatible shapes for `a`, `b` and `x`. Here we compute the + function for three different `a` at four positions `x` and ``b=3``, + resulting in a 3x4 array. + + >>> a = np.array([[0.5], [1.5], [2.5]]) + >>> x = np.array([1., 2., 3., 4]) + >>> a.shape, x.shape + ((3, 1), (4,)) + + >>> gdtr(a, 3., x) + array([[0.01438768, 0.0803014 , 0.19115317, 0.32332358], + [0.19115317, 0.57680992, 0.82642193, 0.9380312 ], + [0.45618688, 0.87534798, 0.97974328, 0.9972306 ]]) + + Plot the function for four different parameter sets. + + >>> a_parameters = [0.3, 1, 2, 6] + >>> b_parameters = [2, 10, 15, 20] + >>> linestyles = ['solid', 'dashed', 'dotted', 'dashdot'] + >>> parameters_list = list(zip(a_parameters, b_parameters, linestyles)) + >>> x = np.linspace(0, 30, 1000) + >>> fig, ax = plt.subplots() + >>> for parameter_set in parameters_list: + ... a, b, style = parameter_set + ... gdtr_vals = gdtr(a, b, x) + ... ax.plot(x, gdtr_vals, label=fr"$a= {a},\, b={b}$", ls=style) + >>> ax.legend() + >>> ax.set_xlabel("$x$") + >>> ax.set_title("Gamma distribution cumulative distribution function") + >>> plt.show() + + The gamma distribution is also available as `scipy.stats.gamma`. Using + `gdtr` directly can be much faster than calling the ``cdf`` method of + `scipy.stats.gamma`, especially for small arrays or individual values. + To get the same results one must use the following parametrization: + ``stats.gamma(b, scale=1/a).cdf(x)=gdtr(a, b, x)``. + + >>> from scipy.stats import gamma + >>> a = 2. + >>> b = 3 + >>> x = 1. + >>> gdtr_result = gdtr(a, b, x) # this will often be faster than below + >>> gamma_dist_result = gamma(b, scale=1/a).cdf(x) + >>> gdtr_result == gamma_dist_result # test that results are equal + True + """) + +add_newdoc("gdtrc", + r""" + gdtrc(a, b, x, out=None) + + Gamma distribution survival function. + + Integral from `x` to infinity of the gamma probability density function, + + .. math:: + + S(x) = \int_x^\infty \frac{a^b}{\Gamma(b)} t^{b-1} e^{-at}\,dt, + + where :math:`\Gamma` is the gamma function. + + Parameters + ---------- + a : array_like + The rate parameter of the gamma distribution, sometimes denoted + :math:`\beta` (float). It is also the reciprocal of the scale + parameter :math:`\theta`. + b : array_like + The shape parameter of the gamma distribution, sometimes denoted + :math:`\alpha` (float). + x : array_like + The quantile (lower limit of integration; float). + out : ndarray, optional + Optional output array for the function values + + Returns + ------- + scalar or ndarray + The survival function of the gamma distribution with parameters `a` + and `b` evaluated at `x`. + + See Also + -------- + gdtr: Gamma distribution cumulative distribution function + scipy.stats.gamma: Gamma distribution + gdtrix + + Notes + ----- + The evaluation is carried out using the relation to the incomplete gamma + integral (regularized gamma function). + + Wrapper for the Cephes [1]_ routine `gdtrc`. Calling `gdtrc` directly can + improve performance compared to the ``sf`` method of `scipy.stats.gamma` + (see last example below). + + References + ---------- + .. [1] Cephes Mathematical Functions Library, + https://netlib.org/cephes/ + + Examples + -------- + Compute the function for ``a=1`` and ``b=2`` at ``x=5``. + + >>> import numpy as np + >>> from scipy.special import gdtrc + >>> import matplotlib.pyplot as plt + >>> gdtrc(1., 2., 5.) + 0.04042768199451279 + + Compute the function for ``a=1``, ``b=2`` at several points by providing + a NumPy array for `x`. + + >>> xvalues = np.array([1., 2., 3., 4]) + >>> gdtrc(1., 1., xvalues) + array([0.36787944, 0.13533528, 0.04978707, 0.01831564]) + + `gdtrc` can evaluate different parameter sets by providing arrays with + broadcasting compatible shapes for `a`, `b` and `x`. Here we compute the + function for three different `a` at four positions `x` and ``b=3``, + resulting in a 3x4 array. + + >>> a = np.array([[0.5], [1.5], [2.5]]) + >>> x = np.array([1., 2., 3., 4]) + >>> a.shape, x.shape + ((3, 1), (4,)) + + >>> gdtrc(a, 3., x) + array([[0.98561232, 0.9196986 , 0.80884683, 0.67667642], + [0.80884683, 0.42319008, 0.17357807, 0.0619688 ], + [0.54381312, 0.12465202, 0.02025672, 0.0027694 ]]) + + Plot the function for four different parameter sets. + + >>> a_parameters = [0.3, 1, 2, 6] + >>> b_parameters = [2, 10, 15, 20] + >>> linestyles = ['solid', 'dashed', 'dotted', 'dashdot'] + >>> parameters_list = list(zip(a_parameters, b_parameters, linestyles)) + >>> x = np.linspace(0, 30, 1000) + >>> fig, ax = plt.subplots() + >>> for parameter_set in parameters_list: + ... a, b, style = parameter_set + ... gdtrc_vals = gdtrc(a, b, x) + ... ax.plot(x, gdtrc_vals, label=fr"$a= {a},\, b={b}$", ls=style) + >>> ax.legend() + >>> ax.set_xlabel("$x$") + >>> ax.set_title("Gamma distribution survival function") + >>> plt.show() + + The gamma distribution is also available as `scipy.stats.gamma`. + Using `gdtrc` directly can be much faster than calling the ``sf`` method + of `scipy.stats.gamma`, especially for small arrays or individual + values. To get the same results one must use the following parametrization: + ``stats.gamma(b, scale=1/a).sf(x)=gdtrc(a, b, x)``. + + >>> from scipy.stats import gamma + >>> a = 2 + >>> b = 3 + >>> x = 1. + >>> gdtrc_result = gdtrc(a, b, x) # this will often be faster than below + >>> gamma_dist_result = gamma(b, scale=1/a).sf(x) + >>> gdtrc_result == gamma_dist_result # test that results are equal + True + """) + +add_newdoc("gdtria", + """ + gdtria(p, b, x, out=None) + + Inverse of `gdtr` vs a. + + Returns the inverse with respect to the parameter `a` of ``p = + gdtr(a, b, x)``, the cumulative distribution function of the gamma + distribution. + + Parameters + ---------- + p : array_like + Probability values. + b : array_like + `b` parameter values of `gdtr(a, b, x)`. `b` is the "shape" parameter + of the gamma distribution. + x : array_like + Nonnegative real values, from the domain of the gamma distribution. + out : ndarray, optional + If a fourth argument is given, it must be a numpy.ndarray whose size + matches the broadcast result of `a`, `b` and `x`. `out` is then the + array returned by the function. + + Returns + ------- + a : scalar or ndarray + Values of the `a` parameter such that ``p = gdtr(a, b, x)`. ``1/a`` + is the "scale" parameter of the gamma distribution. + + See Also + -------- + gdtr : CDF of the gamma distribution. + gdtrib : Inverse with respect to `b` of `gdtr(a, b, x)`. + gdtrix : Inverse with respect to `x` of `gdtr(a, b, x)`. + gammaincinv : Inverse of the incomplete regularized gamma function. + + Notes + ----- + `gdtria` is implemented in terms of the incomplete gamma inverse as + ``gdtria(p, b, x) = gammaincinv(b, p)/x``. + + Examples + -------- + First evaluate `gdtr`. + + >>> from scipy.special import gdtr, gdtria + >>> p = gdtr(1.2, 3.4, 5.6) + >>> print(p) + 0.94378087442 + + Verify the inverse. + + >>> gdtria(p, 3.4, 5.6) + 1.2 + """) + +add_newdoc("gdtrib", + """ + gdtrib(a, p, x, out=None) + + Inverse of `gdtr` vs b. + + Returns the inverse with respect to the parameter `b` of ``p = + gdtr(a, b, x)``, the cumulative distribution function of the gamma + distribution. + + Parameters + ---------- + a : array_like + `a` parameter values of ``gdtr(a, b, x)`. ``1/a`` is the "scale" + parameter of the gamma distribution. + p : array_like + Probability values. + x : array_like + Nonnegative real values, from the domain of the gamma distribution. + out : ndarray, optional + If a fourth argument is given, it must be a numpy.ndarray whose size + matches the broadcast result of `a`, `b` and `x`. `out` is then the + array returned by the function. + + Returns + ------- + b : scalar or ndarray + Values of the `b` parameter such that `p = gdtr(a, b, x)`. `b` is + the "shape" parameter of the gamma distribution. + + See Also + -------- + gdtr : CDF of the gamma distribution. + gdtria : Inverse with respect to `a` of `gdtr(a, b, x)`. + gdtrix : Inverse with respect to `x` of `gdtr(a, b, x)`. + + Notes + ----- + + The cumulative distribution function `p` is computed using the Cephes [1]_ + routines `igam` and `igamc`. Computation of `b` involves a search for a value + that produces the desired value of `p` using Chandrupatla's bracketing + root finding algorithm [2]_. + + Note that there are some edge cases where `gdtrib` is extended by taking + limits where they are uniquely defined. In particular + ``x == 0`` with ``p > 0`` and ``p == 0`` with ``x > 0``. + For these edge cases, a numerical result will be returned for + ``gdtrib(a, p, x)`` even though ``gdtr(a, gdtrib(a, p, x), x)`` is + undefined. + + References + ---------- + .. [1] Cephes Mathematical Functions Library, + https://netlib.org/cephes/ + .. [2] Chandrupatla, Tirupathi R. + "A new hybrid quadratic/bisection algorithm for finding the zero of a + nonlinear function without using derivatives". + Advances in Engineering Software, 28(3), 145-149. + :doi:`10.1016/s0965-9978(96)00051-8`. + + Examples + -------- + First evaluate `gdtr`. + + >>> from scipy.special import gdtr, gdtrib + >>> p = gdtr(1.2, 3.4, 5.6) + >>> print(p) + 0.94378087442 + + Verify the inverse. + + >>> gdtrib(1.2, p, 5.6) + 3.3999999999999995 + """) + +add_newdoc("gdtrix", + """ + gdtrix(a, b, p, out=None) + + Inverse of `gdtr` vs x. + + Returns the inverse with respect to the parameter `x` of ``p = + gdtr(a, b, x)``, the cumulative distribution function of the gamma + distribution. This is also known as the pth quantile of the + distribution. + + Parameters + ---------- + a : array_like + `a` parameter values of ``gdtr(a, b, x)``. ``1/a`` is the "scale" + parameter of the gamma distribution. + b : array_like + `b` parameter values of ``gdtr(a, b, x)``. `b` is the "shape" parameter + of the gamma distribution. + p : array_like + Probability values. + out : ndarray, optional + If a fourth argument is given, it must be a numpy.ndarray whose size + matches the broadcast result of `a`, `b` and `x`. `out` is then the + array returned by the function. + + Returns + ------- + x : scalar or ndarray + Values of the `x` parameter such that `p = gdtr(a, b, x)`. + + See Also + -------- + gdtr : CDF of the gamma distribution. + gdtria : Inverse with respect to `a` of ``gdtr(a, b, x)``. + gdtrib : Inverse with respect to `b` of ``gdtr(a, b, x)``. + gammaincinv : Inverse of the incomplete regularized gamma function. + + Notes + ----- + `gdtrix` is implemented in terms of the incomplete gamma inverse as + ``gdtrix(a, b, p) = gammaincinv(b, p)/a``. + + Examples + -------- + First evaluate `gdtr`. + + >>> from scipy.special import gdtr, gdtrix + >>> p = gdtr(1.2, 3.4, 5.6) + >>> print(p) + 0.94378087442 + + Verify the inverse. + + >>> gdtrix(1.2, 3.4, p) + 5.6 + """) + + + +add_newdoc("huber", + r""" + huber(delta, r, out=None) + + Huber loss function. + + .. math:: \text{huber}(\delta, r) = \begin{cases} \infty & \delta < 0 \\ + \frac{1}{2}r^2 & 0 \le \delta, | r | \le \delta \\ + \delta ( |r| - \frac{1}{2}\delta ) & \text{otherwise} \end{cases} + + Parameters + ---------- + delta : ndarray + Input array, indicating the quadratic vs. linear loss changepoint. + r : ndarray + Input array, possibly representing residuals. + out : ndarray, optional + Optional output array for the function values + + Returns + ------- + scalar or ndarray + The computed Huber loss function values. + + See Also + -------- + pseudo_huber : smooth approximation of this function + + Notes + ----- + `huber` is useful as a loss function in robust statistics or machine + learning to reduce the influence of outliers as compared to the common + squared error loss, residuals with a magnitude higher than `delta` are + not squared [1]_. + + Typically, `r` represents residuals, the difference + between a model prediction and data. Then, for :math:`|r|\leq\delta`, + `huber` resembles the squared error and for :math:`|r|>\delta` the + absolute error. This way, the Huber loss often achieves + a fast convergence in model fitting for small residuals like the squared + error loss function and still reduces the influence of outliers + (:math:`|r|>\delta`) like the absolute error loss. As :math:`\delta` is + the cutoff between squared and absolute error regimes, it has + to be tuned carefully for each problem. `huber` is also + convex, making it suitable for gradient based optimization. + + .. versionadded:: 0.15.0 + + References + ---------- + .. [1] Peter Huber. "Robust Estimation of a Location Parameter", + 1964. Annals of Statistics. 53 (1): 73 - 101. + + Examples + -------- + Import all necessary modules. + + >>> import numpy as np + >>> from scipy.special import huber + >>> import matplotlib.pyplot as plt + + Compute the function for ``delta=1`` at ``r=2`` + + >>> huber(1., 2.) + 1.5 + + Compute the function for different `delta` by providing a NumPy array or + list for `delta`. + + >>> huber([1., 3., 5.], 4.) + array([3.5, 7.5, 8. ]) + + Compute the function at different points by providing a NumPy array or + list for `r`. + + >>> huber(2., np.array([1., 1.5, 3.])) + array([0.5 , 1.125, 4. ]) + + The function can be calculated for different `delta` and `r` by + providing arrays for both with compatible shapes for broadcasting. + + >>> r = np.array([1., 2.5, 8., 10.]) + >>> deltas = np.array([[1.], [5.], [9.]]) + >>> print(r.shape, deltas.shape) + (4,) (3, 1) + + >>> huber(deltas, r) + array([[ 0.5 , 2. , 7.5 , 9.5 ], + [ 0.5 , 3.125, 27.5 , 37.5 ], + [ 0.5 , 3.125, 32. , 49.5 ]]) + + Plot the function for different `delta`. + + >>> x = np.linspace(-4, 4, 500) + >>> deltas = [1, 2, 3] + >>> linestyles = ["dashed", "dotted", "dashdot"] + >>> fig, ax = plt.subplots() + >>> combined_plot_parameters = list(zip(deltas, linestyles)) + >>> for delta, style in combined_plot_parameters: + ... ax.plot(x, huber(delta, x), label=fr"$\delta={delta}$", ls=style) + >>> ax.legend(loc="upper center") + >>> ax.set_xlabel("$x$") + >>> ax.set_title(r"Huber loss function $h_{\delta}(x)$") + >>> ax.set_xlim(-4, 4) + >>> ax.set_ylim(0, 8) + >>> plt.show() + """) + +add_newdoc("hyp0f1", + r""" + hyp0f1(v, z, out=None) + + Confluent hypergeometric limit function 0F1. + + Parameters + ---------- + v : array_like + Real-valued parameter + z : array_like + Real- or complex-valued argument + out : ndarray, optional + Optional output array for the function results + + Returns + ------- + scalar or ndarray + The confluent hypergeometric limit function + + Notes + ----- + This function is defined as: + + .. math:: _0F_1(v, z) = \sum_{k=0}^{\infty}\frac{z^k}{(v)_k k!}. + + It's also the limit as :math:`q \to \infty` of :math:`_1F_1(q; v; z/q)`, + and satisfies the differential equation :math:`f''(z) + vf'(z) = + f(z)`. See [1]_ for more information. + + References + ---------- + .. [1] Wolfram MathWorld, "Confluent Hypergeometric Limit Function", + https://mathworld.wolfram.com/ConfluentHypergeometricLimitFunction.html + + Examples + -------- + >>> import numpy as np + >>> import scipy.special as sc + + It is one when `z` is zero. + + >>> sc.hyp0f1(1, 0) + 1.0 + + It is the limit of the confluent hypergeometric function as `q` + goes to infinity. + + >>> q = np.array([1, 10, 100, 1000]) + >>> v = 1 + >>> z = 1 + >>> sc.hyp1f1(q, v, z / q) + array([2.71828183, 2.31481985, 2.28303778, 2.27992985]) + >>> sc.hyp0f1(v, z) + 2.2795853023360673 + + It is related to Bessel functions. + + >>> n = 1 + >>> x = np.linspace(0, 1, 5) + >>> sc.jv(n, x) + array([0. , 0.12402598, 0.24226846, 0.3492436 , 0.44005059]) + >>> (0.5 * x)**n / sc.factorial(n) * sc.hyp0f1(n + 1, -0.25 * x**2) + array([0. , 0.12402598, 0.24226846, 0.3492436 , 0.44005059]) + + """) + +add_newdoc("hyp1f1", + r""" + hyp1f1(a, b, x, out=None) + + Confluent hypergeometric function 1F1. + + The confluent hypergeometric function is defined by the series + + .. math:: + + {}_1F_1(a; b; x) = \sum_{k = 0}^\infty \frac{(a)_k}{(b)_k k!} x^k. + + See [DLMF]_ for more details. Here :math:`(\cdot)_k` is the + Pochhammer symbol; see `poch`. + + Parameters + ---------- + a, b : array_like + Real parameters + x : array_like + Real or complex argument + out : ndarray, optional + Optional output array for the function results + + Returns + ------- + scalar or ndarray + Values of the confluent hypergeometric function + + See Also + -------- + hyperu : another confluent hypergeometric function + hyp0f1 : confluent hypergeometric limit function + hyp2f1 : Gaussian hypergeometric function + + Notes + ----- + For real values, this function uses the ``hyp1f1`` routine from the C++ Boost + library [2]_, for complex values a C translation of the specfun + Fortran library [3]_. + + References + ---------- + .. [DLMF] NIST Digital Library of Mathematical Functions + https://dlmf.nist.gov/13.2#E2 + .. [2] The Boost Developers. "Boost C++ Libraries". https://www.boost.org/. + .. [3] S. Zhang and J.M. Jin, "Computation of Special Functions", Wiley 1996. + + Examples + -------- + >>> import numpy as np + >>> import scipy.special as sc + + It is one when `x` is zero: + + >>> sc.hyp1f1(0.5, 0.5, 0) + 1.0 + + It is singular when `b` is a nonpositive integer. + + >>> sc.hyp1f1(0.5, -1, 0) + inf + + It is a polynomial when `a` is a nonpositive integer. + + >>> a, b, x = -1, 0.5, np.array([1.0, 2.0, 3.0, 4.0]) + >>> sc.hyp1f1(a, b, x) + array([-1., -3., -5., -7.]) + >>> 1 + (a / b) * x + array([-1., -3., -5., -7.]) + + It reduces to the exponential function when ``a = b``. + + >>> sc.hyp1f1(2, 2, [1, 2, 3, 4]) + array([ 2.71828183, 7.3890561 , 20.08553692, 54.59815003]) + >>> np.exp([1, 2, 3, 4]) + array([ 2.71828183, 7.3890561 , 20.08553692, 54.59815003]) + + """) + +add_newdoc("hyperu", + r""" + hyperu(a, b, x, out=None) + + Confluent hypergeometric function U. + + It is defined as the solution to the equation + + .. math:: + + x \frac{d^2w}{dx^2} + (b - x) \frac{dw}{dx} - aw = 0 + + which satisfies the property + + .. math:: + + U(a, b, x) \sim x^{-a} + + as :math:`x \to \infty`. See [DLMF]_ for more details. + + Parameters + ---------- + a, b : array_like + Real-valued parameters + x : array_like + Real-valued argument + out : ndarray, optional + Optional output array for the function values + + Returns + ------- + scalar or ndarray + Values of `U` + + References + ---------- + .. [DLMF] NIST Digital Library of Mathematics Functions + https://dlmf.nist.gov/13.2#E6 + + Examples + -------- + >>> import numpy as np + >>> import scipy.special as sc + + It has a branch cut along the negative `x` axis. + + >>> x = np.linspace(-0.1, -10, 5) + >>> sc.hyperu(1, 1, x) + array([nan, nan, nan, nan, nan]) + + It approaches zero as `x` goes to infinity. + + >>> x = np.array([1, 10, 100]) + >>> sc.hyperu(1, 1, x) + array([0.59634736, 0.09156333, 0.00990194]) + + It satisfies Kummer's transformation. + + >>> a, b, x = 2, 1, 1 + >>> sc.hyperu(a, b, x) + 0.1926947246463881 + >>> x**(1 - b) * sc.hyperu(a - b + 1, 2 - b, x) + 0.1926947246463881 + + """) + +add_newdoc("_igam_fac", + """ + Internal function, do not use. + """) + +add_newdoc("kl_div", + r""" + kl_div(x, y, out=None) + + Elementwise function for computing Kullback-Leibler divergence. + + .. math:: + + \mathrm{kl\_div}(x, y) = + \begin{cases} + x \log(x / y) - x + y & x > 0, y > 0 \\ + y & x = 0, y \ge 0 \\ + \infty & \text{otherwise} + \end{cases} + + Parameters + ---------- + x, y : array_like + Real arguments + out : ndarray, optional + Optional output array for the function results + + Returns + ------- + scalar or ndarray + Values of the Kullback-Liebler divergence. + + See Also + -------- + entr, rel_entr, scipy.stats.entropy + + Notes + ----- + .. versionadded:: 0.15.0 + + This function is non-negative and is jointly convex in `x` and `y`. + + The origin of this function is in convex programming; see [1]_ for + details. This is why the function contains the extra :math:`-x + + y` terms over what might be expected from the Kullback-Leibler + divergence. For a version of the function without the extra terms, + see `rel_entr`. + + References + ---------- + .. [1] Boyd, Stephen and Lieven Vandenberghe. *Convex optimization*. + Cambridge University Press, 2004. + :doi:`10.1017/CBO9780511804441`. + + """) + +add_newdoc("kn", + r""" + kn(n, x, out=None) + + Modified Bessel function of the second kind of integer order `n`. + + Returns the modified Bessel function of the second kind for integer order + `n` at real `z`. + + These are also sometimes called functions of the third kind, Basset + functions, or Macdonald functions. + + Parameters + ---------- + n : array_like of int + Order of Bessel functions (floats will truncate with a warning) + x : array_like of float + Argument at which to evaluate the Bessel functions + out : ndarray, optional + Optional output array for the function results. + + Returns + ------- + scalar or ndarray + Value of the Modified Bessel function of the second kind, + :math:`K_n(x)`. + + See Also + -------- + kv : Same function, but accepts real order and complex argument + kvp : Derivative of this function + + Notes + ----- + Wrapper for AMOS [1]_ routine `zbesk`. For a discussion of the + algorithm used, see [2]_ and the references therein. + + References + ---------- + .. [1] Donald E. Amos, "AMOS, A Portable Package for Bessel Functions + of a Complex Argument and Nonnegative Order", + https://netlib.org/amos/ + .. [2] Donald E. Amos, "Algorithm 644: A portable package for Bessel + functions of a complex argument and nonnegative order", ACM + TOMS Vol. 12 Issue 3, Sept. 1986, p. 265. + + Examples + -------- + Plot the function of several orders for real input: + + >>> import numpy as np + >>> from scipy.special import kn + >>> import matplotlib.pyplot as plt + >>> x = np.linspace(0, 5, 1000) + >>> for N in range(6): + ... plt.plot(x, kn(N, x), label='$K_{}(x)$'.format(N)) + >>> plt.ylim(0, 10) + >>> plt.legend() + >>> plt.title(r'Modified Bessel function of the second kind $K_n(x)$') + >>> plt.show() + + Calculate for a single value at multiple orders: + + >>> kn([4, 5, 6], 1) + array([ 44.23241585, 360.9605896 , 3653.83831186]) + """) + +add_newdoc("kolmogi", + """ + kolmogi(p, out=None) + + Inverse Survival Function of Kolmogorov distribution. + + It is the inverse function to `kolmogorov`. + Returns y such that ``kolmogorov(y) == p``. + + Parameters + ---------- + p : float array_like + Probability + out : ndarray, optional + Optional output array for the function results + + Returns + ------- + scalar or ndarray + The value(s) of kolmogi(p) + + See Also + -------- + kolmogorov : The Survival Function for the distribution + scipy.stats.kstwobign : Provides the functionality as a continuous distribution + smirnov, smirnovi : Functions for the one-sided distribution + + Notes + ----- + `kolmogorov` is used by `stats.kstest` in the application of the + Kolmogorov-Smirnov Goodness of Fit test. For historical reasons this + function is exposed in `scpy.special`, but the recommended way to achieve + the most accurate CDF/SF/PDF/PPF/ISF computations is to use the + `stats.kstwobign` distribution. + + Examples + -------- + >>> from scipy.special import kolmogi + >>> kolmogi([0, 0.1, 0.25, 0.5, 0.75, 0.9, 1.0]) + array([ inf, 1.22384787, 1.01918472, 0.82757356, 0.67644769, + 0.57117327, 0. ]) + + """) + +add_newdoc("kolmogorov", + r""" + kolmogorov(y, out=None) + + Complementary cumulative distribution (Survival Function) function of + Kolmogorov distribution. + + Returns the complementary cumulative distribution function of + Kolmogorov's limiting distribution (``D_n*\sqrt(n)`` as n goes to infinity) + of a two-sided test for equality between an empirical and a theoretical + distribution. It is equal to the (limit as n->infinity of the) + probability that ``sqrt(n) * max absolute deviation > y``. + + Parameters + ---------- + y : float array_like + Absolute deviation between the Empirical CDF (ECDF) and the target CDF, + multiplied by sqrt(n). + out : ndarray, optional + Optional output array for the function results + + Returns + ------- + scalar or ndarray + The value(s) of kolmogorov(y) + + See Also + -------- + kolmogi : The Inverse Survival Function for the distribution + scipy.stats.kstwobign : Provides the functionality as a continuous distribution + smirnov, smirnovi : Functions for the one-sided distribution + + Notes + ----- + `kolmogorov` is used by `stats.kstest` in the application of the + Kolmogorov-Smirnov Goodness of Fit test. For historical reasons this + function is exposed in `scpy.special`, but the recommended way to achieve + the most accurate CDF/SF/PDF/PPF/ISF computations is to use the + `stats.kstwobign` distribution. + + Examples + -------- + Show the probability of a gap at least as big as 0, 0.5 and 1.0. + + >>> import numpy as np + >>> from scipy.special import kolmogorov + >>> from scipy.stats import kstwobign + >>> kolmogorov([0, 0.5, 1.0]) + array([ 1. , 0.96394524, 0.26999967]) + + Compare a sample of size 1000 drawn from a Laplace(0, 1) distribution against + the target distribution, a Normal(0, 1) distribution. + + >>> from scipy.stats import norm, laplace + >>> rng = np.random.default_rng() + >>> n = 1000 + >>> lap01 = laplace(0, 1) + >>> x = np.sort(lap01.rvs(n, random_state=rng)) + >>> np.mean(x), np.std(x) + (-0.05841730131499543, 1.3968109101997568) + + Construct the Empirical CDF and the K-S statistic Dn. + + >>> target = norm(0,1) # Normal mean 0, stddev 1 + >>> cdfs = target.cdf(x) + >>> ecdfs = np.arange(n+1, dtype=float)/n + >>> gaps = np.column_stack([cdfs - ecdfs[:n], ecdfs[1:] - cdfs]) + >>> Dn = np.max(gaps) + >>> Kn = np.sqrt(n) * Dn + >>> print('Dn=%f, sqrt(n)*Dn=%f' % (Dn, Kn)) + Dn=0.043363, sqrt(n)*Dn=1.371265 + >>> print(chr(10).join(['For a sample of size n drawn from a N(0, 1) distribution:', + ... ' the approximate Kolmogorov probability that sqrt(n)*Dn>=%f is %f' % + ... (Kn, kolmogorov(Kn)), + ... ' the approximate Kolmogorov probability that sqrt(n)*Dn<=%f is %f' % + ... (Kn, kstwobign.cdf(Kn))])) + For a sample of size n drawn from a N(0, 1) distribution: + the approximate Kolmogorov probability that sqrt(n)*Dn>=1.371265 is 0.046533 + the approximate Kolmogorov probability that sqrt(n)*Dn<=1.371265 is 0.953467 + + Plot the Empirical CDF against the target N(0, 1) CDF. + + >>> import matplotlib.pyplot as plt + >>> plt.step(np.concatenate([[-3], x]), ecdfs, where='post', label='Empirical CDF') + >>> x3 = np.linspace(-3, 3, 100) + >>> plt.plot(x3, target.cdf(x3), label='CDF for N(0, 1)') + >>> plt.ylim([0, 1]); plt.grid(True); plt.legend(); + >>> # Add vertical lines marking Dn+ and Dn- + >>> iminus, iplus = np.argmax(gaps, axis=0) + >>> plt.vlines([x[iminus]], ecdfs[iminus], cdfs[iminus], + ... color='r', linestyle='dashed', lw=4) + >>> plt.vlines([x[iplus]], cdfs[iplus], ecdfs[iplus+1], + ... color='r', linestyle='dashed', lw=4) + >>> plt.show() + """) + +add_newdoc("_kolmogc", + r""" + Internal function, do not use. + """) + +add_newdoc("_kolmogci", + r""" + Internal function, do not use. + """) + +add_newdoc("_kolmogp", + r""" + Internal function, do not use. + """) + +add_newdoc("_lanczos_sum_expg_scaled", + """ + Internal function, do not use. + """) + +add_newdoc( + "_landau_pdf", + """ + _landau_pdf(x, loc, scale) + + Probability density function of the Landau distribution. + + Parameters + ---------- + x : array_like + Real-valued argument + loc : array_like + Real-valued distribution location + scale : array_like + Positive, real-valued distribution scale + + Returns + ------- + scalar or ndarray + """) + +add_newdoc( + "_landau_cdf", + """ + _landau_cdf(x, loc, scale) + + Cumulative distribution function of the Landau distribution. + + Parameters + ---------- + x : array_like + Real-valued argument + loc : array_like + Real-valued distribution location + scale : array_like + Positive, real-valued distribution scale + + Returns + ------- + scalar or ndarray + """) + +add_newdoc( + "_landau_sf", + """ + _landau_sf(x, loc, scale) + + Survival function of the Landau distribution. + + Parameters + ---------- + x : array_like + Real-valued argument + loc : array_like + Real-valued distribution location + scale : array_like + Positive, real-valued distribution scale + + Returns + ------- + scalar or ndarray + """) + +add_newdoc( + "_landau_ppf", + """ + _landau_ppf(p, loc, scale) + + Percent point function of the Landau distribution. + + Parameters + ---------- + p : array_like + Real-valued argument between 0 and 1 + loc : array_like + Real-valued distribution location + scale : array_like + Positive, real-valued distribution scale + + Returns + ------- + scalar or ndarray + """) + +add_newdoc( + "_landau_isf", + """ + _landau_isf(p, loc, scale) + + Inverse survival function of the Landau distribution. + + Parameters + ---------- + p : array_like + Real-valued argument between 0 and 1 + loc : array_like + Real-valued distribution location + scale : array_like + Positive, real-valued distribution scale + + Returns + ------- + scalar or ndarray + """) + +add_newdoc("_lgam1p", + """ + Internal function, do not use. + """) + +add_newdoc("lpmv", + r""" + lpmv(m, v, x, out=None) + + Associated Legendre function of integer order and real degree. + + Defined as + + .. math:: + + P_v^m(x) = (-1)^m (1 - x^2)^{m/2} \frac{d^m}{dx^m} (P_v(x)) + + where + + .. math:: + + P_v(x) = \sum_{k = 0}^\infty \frac{(-v)_k (v + 1)_k}{(k!)^2} + \left(\frac{1 - x}{2}\right)^k + + is the Legendre function of the first kind. Here :math:`(\cdot)_k` + is the Pochhammer symbol; see `poch`. + + Parameters + ---------- + m : array_like + Order (int or float). If passed a float not equal to an + integer the function returns NaN. + v : array_like + Degree (float). + x : array_like + Argument (float). Must have ``|x| <= 1``. + out : ndarray, optional + Optional output array for the function results. + + Returns + ------- + pmv : scalar or ndarray + Value of the associated Legendre function. + + Notes + ----- + Note that this implementation includes the Condon-Shortley phase. + + References + ---------- + .. [1] S. Zhang and J.M. Jin, "Computation of Special Functions", Wiley 1996. + + Examples + -------- + Compute a single associated Legendre function value: + + >>> from scipy.special import lpmv + >>> lpmv(1, 2.5, 0.3) + -0.1292299 + + Plot the associated Legendre functions for orders ``m = 1`` and + ``m = 2`` and degrees ``v = 0, ..., 5``: + + >>> import numpy as np + >>> import matplotlib.pyplot as plt + >>> nus = np.arange(0, 6) + >>> xs = np.linspace(-1+1e-5, 1-1e-5, 101) + + Order 1. + + >>> fig, ax = plt.subplots() + >>> ax.plot(xs, lpmv(1, nus, xs[:, np.newaxis]), "-") + >>> ax.legend([rf"$\nu={nu}$" for nu in nus]) + >>> ax.set_xlabel(r"$x$") + >>> ax.set_ylabel(r"$P_\nu^1(x)$") + >>> plt.show() + + Order 2. + + >>> fig, ax = plt.subplots() + >>> ax.plot(xs, lpmv(2, nus, xs[:, np.newaxis]), "-") + >>> ax.legend([rf"$\nu={nu}$" for nu in nus]) + >>> ax.set_xlabel(r"$x$") + >>> ax.set_ylabel(r"$P_\nu^2(x)$") + >>> plt.show() + """) + +add_newdoc("nbdtr", + r""" + nbdtr(k, n, p, out=None) + + Negative binomial cumulative distribution function. + + Returns the sum of the terms 0 through `k` of the negative binomial + distribution probability mass function, + + .. math:: + + F(k) = \sum_{j=0}^k {{n + j - 1}\choose{j}} p^n (1 - p)^j. + + In a sequence of Bernoulli trials with individual success probabilities + `p`, this is the probability that `k` or fewer failures precede the nth + success. + + Parameters + ---------- + k : array_like + The maximum number of allowed failures (nonnegative int). + n : array_like + The target number of successes (positive int). + p : array_like + Probability of success in a single event (float). + out : ndarray, optional + Optional output array for the function results + + Returns + ------- + scalar or ndarray + The probability of `k` or fewer failures before `n` successes in a + sequence of events with individual success probability `p`. + + See Also + -------- + nbdtrc : Negative binomial survival function + nbdtrik : Negative binomial quantile function + scipy.stats.nbinom : Negative binomial distribution + + Notes + ----- + If floating point values are passed for `k` or `n`, they will be truncated + to integers. + + The terms are not summed directly; instead the regularized incomplete beta + function is employed, according to the formula, + + .. math:: + \mathrm{nbdtr}(k, n, p) = I_{p}(n, k + 1). + + Wrapper for the Cephes [1]_ routine `nbdtr`. + + The negative binomial distribution is also available as + `scipy.stats.nbinom`. Using `nbdtr` directly can improve performance + compared to the ``cdf`` method of `scipy.stats.nbinom` (see last example). + + References + ---------- + .. [1] Cephes Mathematical Functions Library, + https://netlib.org/cephes/ + + Examples + -------- + Compute the function for ``k=10`` and ``n=5`` at ``p=0.5``. + + >>> import numpy as np + >>> from scipy.special import nbdtr + >>> nbdtr(10, 5, 0.5) + 0.940765380859375 + + Compute the function for ``n=10`` and ``p=0.5`` at several points by + providing a NumPy array or list for `k`. + + >>> nbdtr([5, 10, 15], 10, 0.5) + array([0.15087891, 0.58809853, 0.88523853]) + + Plot the function for four different parameter sets. + + >>> import matplotlib.pyplot as plt + >>> k = np.arange(130) + >>> n_parameters = [20, 20, 20, 80] + >>> p_parameters = [0.2, 0.5, 0.8, 0.5] + >>> linestyles = ['solid', 'dashed', 'dotted', 'dashdot'] + >>> parameters_list = list(zip(p_parameters, n_parameters, + ... linestyles)) + >>> fig, ax = plt.subplots(figsize=(8, 8)) + >>> for parameter_set in parameters_list: + ... p, n, style = parameter_set + ... nbdtr_vals = nbdtr(k, n, p) + ... ax.plot(k, nbdtr_vals, label=rf"$n={n},\, p={p}$", + ... ls=style) + >>> ax.legend() + >>> ax.set_xlabel("$k$") + >>> ax.set_title("Negative binomial cumulative distribution function") + >>> plt.show() + + The negative binomial distribution is also available as + `scipy.stats.nbinom`. Using `nbdtr` directly can be much faster than + calling the ``cdf`` method of `scipy.stats.nbinom`, especially for small + arrays or individual values. To get the same results one must use the + following parametrization: ``nbinom(n, p).cdf(k)=nbdtr(k, n, p)``. + + >>> from scipy.stats import nbinom + >>> k, n, p = 5, 3, 0.5 + >>> nbdtr_res = nbdtr(k, n, p) # this will often be faster than below + >>> stats_res = nbinom(n, p).cdf(k) + >>> stats_res, nbdtr_res # test that results are equal + (0.85546875, 0.85546875) + + `nbdtr` can evaluate different parameter sets by providing arrays with + shapes compatible for broadcasting for `k`, `n` and `p`. Here we compute + the function for three different `k` at four locations `p`, resulting in + a 3x4 array. + + >>> k = np.array([[5], [10], [15]]) + >>> p = np.array([0.3, 0.5, 0.7, 0.9]) + >>> k.shape, p.shape + ((3, 1), (4,)) + + >>> nbdtr(k, 5, p) + array([[0.15026833, 0.62304687, 0.95265101, 0.9998531 ], + [0.48450894, 0.94076538, 0.99932777, 0.99999999], + [0.76249222, 0.99409103, 0.99999445, 1. ]]) + """) + +add_newdoc("nbdtrc", + r""" + nbdtrc(k, n, p, out=None) + + Negative binomial survival function. + + Returns the sum of the terms `k + 1` to infinity of the negative binomial + distribution probability mass function, + + .. math:: + + S(k) = \sum_{j=k + 1}^\infty {{n + j - 1}\choose{j}} p^n (1 - p)^j. + + In a sequence of Bernoulli trials with individual success probabilities + `p`, this is the probability that more than `k` failures precede the nth + success. + + Parameters + ---------- + k : array_like + The maximum number of allowed failures (nonnegative int). + n : array_like + The target number of successes (positive int). + p : array_like + Probability of success in a single event (float). + out : ndarray, optional + Optional output array for the function results + + Returns + ------- + scalar or ndarray + The probability of `k + 1` or more failures before `n` successes in a + sequence of events with individual success probability `p`. + + See Also + -------- + nbdtr : Negative binomial cumulative distribution function + nbdtrik : Negative binomial percentile function + scipy.stats.nbinom : Negative binomial distribution + + Notes + ----- + If floating point values are passed for `k` or `n`, they will be truncated + to integers. + + The terms are not summed directly; instead the regularized incomplete beta + function is employed, according to the formula, + + .. math:: + \mathrm{nbdtrc}(k, n, p) = I_{1 - p}(k + 1, n). + + Wrapper for the Cephes [1]_ routine `nbdtrc`. + + The negative binomial distribution is also available as + `scipy.stats.nbinom`. Using `nbdtrc` directly can improve performance + compared to the ``sf`` method of `scipy.stats.nbinom` (see last example). + + References + ---------- + .. [1] Cephes Mathematical Functions Library, + https://netlib.org/cephes/ + + Examples + -------- + Compute the function for ``k=10`` and ``n=5`` at ``p=0.5``. + + >>> import numpy as np + >>> from scipy.special import nbdtrc + >>> nbdtrc(10, 5, 0.5) + 0.059234619140624986 + + Compute the function for ``n=10`` and ``p=0.5`` at several points by + providing a NumPy array or list for `k`. + + >>> nbdtrc([5, 10, 15], 10, 0.5) + array([0.84912109, 0.41190147, 0.11476147]) + + Plot the function for four different parameter sets. + + >>> import matplotlib.pyplot as plt + >>> k = np.arange(130) + >>> n_parameters = [20, 20, 20, 80] + >>> p_parameters = [0.2, 0.5, 0.8, 0.5] + >>> linestyles = ['solid', 'dashed', 'dotted', 'dashdot'] + >>> parameters_list = list(zip(p_parameters, n_parameters, + ... linestyles)) + >>> fig, ax = plt.subplots(figsize=(8, 8)) + >>> for parameter_set in parameters_list: + ... p, n, style = parameter_set + ... nbdtrc_vals = nbdtrc(k, n, p) + ... ax.plot(k, nbdtrc_vals, label=rf"$n={n},\, p={p}$", + ... ls=style) + >>> ax.legend() + >>> ax.set_xlabel("$k$") + >>> ax.set_title("Negative binomial distribution survival function") + >>> plt.show() + + The negative binomial distribution is also available as + `scipy.stats.nbinom`. Using `nbdtrc` directly can be much faster than + calling the ``sf`` method of `scipy.stats.nbinom`, especially for small + arrays or individual values. To get the same results one must use the + following parametrization: ``nbinom(n, p).sf(k)=nbdtrc(k, n, p)``. + + >>> from scipy.stats import nbinom + >>> k, n, p = 3, 5, 0.5 + >>> nbdtr_res = nbdtrc(k, n, p) # this will often be faster than below + >>> stats_res = nbinom(n, p).sf(k) + >>> stats_res, nbdtr_res # test that results are equal + (0.6367187499999999, 0.6367187499999999) + + `nbdtrc` can evaluate different parameter sets by providing arrays with + shapes compatible for broadcasting for `k`, `n` and `p`. Here we compute + the function for three different `k` at four locations `p`, resulting in + a 3x4 array. + + >>> k = np.array([[5], [10], [15]]) + >>> p = np.array([0.3, 0.5, 0.7, 0.9]) + >>> k.shape, p.shape + ((3, 1), (4,)) + + >>> nbdtrc(k, 5, p) + array([[8.49731667e-01, 3.76953125e-01, 4.73489874e-02, 1.46902600e-04], + [5.15491059e-01, 5.92346191e-02, 6.72234070e-04, 9.29610100e-09], + [2.37507779e-01, 5.90896606e-03, 5.55025308e-06, 3.26346760e-13]]) + """) + +add_newdoc( + "nbdtri", + r""" + nbdtri(k, n, y, out=None) + + Returns the inverse with respect to the parameter `p` of + ``y = nbdtr(k, n, p)``, the negative binomial cumulative distribution + function. + + Parameters + ---------- + k : array_like + The maximum number of allowed failures (nonnegative int). + n : array_like + The target number of successes (positive int). + y : array_like + The probability of `k` or fewer failures before `n` successes (float). + out : ndarray, optional + Optional output array for the function results + + Returns + ------- + p : scalar or ndarray + Probability of success in a single event (float) such that + `nbdtr(k, n, p) = y`. + + See Also + -------- + nbdtr : Cumulative distribution function of the negative binomial. + nbdtrc : Negative binomial survival function. + scipy.stats.nbinom : negative binomial distribution. + nbdtrik : Inverse with respect to `k` of `nbdtr(k, n, p)`. + nbdtrin : Inverse with respect to `n` of `nbdtr(k, n, p)`. + scipy.stats.nbinom : Negative binomial distribution + + Notes + ----- + Wrapper for the Cephes [1]_ routine `nbdtri`. + + The negative binomial distribution is also available as + `scipy.stats.nbinom`. Using `nbdtri` directly can improve performance + compared to the ``ppf`` method of `scipy.stats.nbinom`. + + References + ---------- + .. [1] Cephes Mathematical Functions Library, + https://netlib.org/cephes/ + + Examples + -------- + `nbdtri` is the inverse of `nbdtr` with respect to `p`. + Up to floating point errors the following holds: + ``nbdtri(k, n, nbdtr(k, n, p))=p``. + + >>> import numpy as np + >>> from scipy.special import nbdtri, nbdtr + >>> k, n, p = 5, 10, 0.2 + >>> cdf_val = nbdtr(k, n, p) + >>> nbdtri(k, n, cdf_val) + 0.20000000000000004 + + Compute the function for ``k=10`` and ``n=5`` at several points by + providing a NumPy array or list for `y`. + + >>> y = np.array([0.1, 0.4, 0.8]) + >>> nbdtri(3, 5, y) + array([0.34462319, 0.51653095, 0.69677416]) + + Plot the function for three different parameter sets. + + >>> import matplotlib.pyplot as plt + >>> n_parameters = [5, 20, 30, 30] + >>> k_parameters = [20, 20, 60, 80] + >>> linestyles = ['solid', 'dashed', 'dotted', 'dashdot'] + >>> parameters_list = list(zip(n_parameters, k_parameters, linestyles)) + >>> cdf_vals = np.linspace(0, 1, 1000) + >>> fig, ax = plt.subplots(figsize=(8, 8)) + >>> for parameter_set in parameters_list: + ... n, k, style = parameter_set + ... nbdtri_vals = nbdtri(k, n, cdf_vals) + ... ax.plot(cdf_vals, nbdtri_vals, label=rf"$k={k},\ n={n}$", + ... ls=style) + >>> ax.legend() + >>> ax.set_ylabel("$p$") + >>> ax.set_xlabel("$CDF$") + >>> title = "nbdtri: inverse of negative binomial CDF with respect to $p$" + >>> ax.set_title(title) + >>> plt.show() + + `nbdtri` can evaluate different parameter sets by providing arrays with + shapes compatible for broadcasting for `k`, `n` and `p`. Here we compute + the function for three different `k` at four locations `p`, resulting in + a 3x4 array. + + >>> k = np.array([[5], [10], [15]]) + >>> y = np.array([0.3, 0.5, 0.7, 0.9]) + >>> k.shape, y.shape + ((3, 1), (4,)) + + >>> nbdtri(k, 5, y) + array([[0.37258157, 0.45169416, 0.53249956, 0.64578407], + [0.24588501, 0.30451981, 0.36778453, 0.46397088], + [0.18362101, 0.22966758, 0.28054743, 0.36066188]]) + """) + +add_newdoc("nbdtrik", + r""" + nbdtrik(y, n, p, out=None) + + Negative binomial percentile function. + + Returns the inverse with respect to the parameter `k` of + ``y = nbdtr(k, n, p)``, the negative binomial cumulative distribution + function. + + Parameters + ---------- + y : array_like + The probability of `k` or fewer failures before `n` successes (float). + n : array_like + The target number of successes (positive int). + p : array_like + Probability of success in a single event (float). + out : ndarray, optional + Optional output array for the function results + + Returns + ------- + k : scalar or ndarray + The maximum number of allowed failures such that `nbdtr(k, n, p) = y`. + + See Also + -------- + nbdtr : Cumulative distribution function of the negative binomial. + nbdtrc : Survival function of the negative binomial. + nbdtri : Inverse with respect to `p` of `nbdtr(k, n, p)`. + nbdtrin : Inverse with respect to `n` of `nbdtr(k, n, p)`. + scipy.stats.nbinom : Negative binomial distribution + + Notes + ----- + This function wraps routines from the Boost Math C++ library [1]_. + + References + ---------- + .. [1] The Boost Developers. "Boost C++ Libraries". https://www.boost.org/. + + Examples + -------- + Compute the negative binomial cumulative distribution function for an + exemplary parameter set. + + >>> import numpy as np + >>> from scipy.special import nbdtr, nbdtrik + >>> k, n, p = 5, 2, 0.5 + >>> cdf_value = nbdtr(k, n, p) + >>> cdf_value + 0.9375 + + Verify that `nbdtrik` recovers the original value for `k`. + + >>> nbdtrik(cdf_value, n, p) + 5.0 + + Plot the function for different parameter sets. + + >>> import matplotlib.pyplot as plt + >>> p_parameters = [0.2, 0.5, 0.7, 0.5] + >>> n_parameters = [30, 30, 30, 80] + >>> linestyles = ['solid', 'dashed', 'dotted', 'dashdot'] + >>> parameters_list = list(zip(p_parameters, n_parameters, linestyles)) + >>> cdf_vals = np.linspace(0, 1, 1000) + >>> fig, ax = plt.subplots(figsize=(8, 8)) + >>> for parameter_set in parameters_list: + ... p, n, style = parameter_set + ... nbdtrik_vals = nbdtrik(cdf_vals, n, p) + ... ax.plot(cdf_vals, nbdtrik_vals, label=rf"$n={n},\ p={p}$", + ... ls=style) + >>> ax.legend() + >>> ax.set_ylabel("$k$") + >>> ax.set_xlabel("$CDF$") + >>> ax.set_title("Negative binomial percentile function") + >>> plt.show() + + The negative binomial distribution is also available as + `scipy.stats.nbinom`. The percentile function method ``ppf`` + returns the result of `nbdtrik` rounded up to integers: + + >>> from scipy.stats import nbinom + >>> q, n, p = 0.6, 5, 0.5 + >>> nbinom.ppf(q, n, p), nbdtrik(q, n, p) + (5.0, 4.800428460273882) + + """) + +add_newdoc("nbdtrin", + r""" + nbdtrin(k, y, p, out=None) + + Inverse of `nbdtr` vs `n`. + + Returns the inverse with respect to the parameter `n` of + ``y = nbdtr(k, n, p)``, the negative binomial cumulative distribution + function. + + Parameters + ---------- + k : array_like + The maximum number of allowed failures (nonnegative int). + y : array_like + The probability of `k` or fewer failures before `n` successes (float). + p : array_like + Probability of success in a single event (float). + out : ndarray, optional + Optional output array for the function results + + Returns + ------- + n : scalar or ndarray + The number of successes `n` such that `nbdtr(k, n, p) = y`. + + See Also + -------- + nbdtr : Cumulative distribution function of the negative binomial. + nbdtri : Inverse with respect to `p` of `nbdtr(k, n, p)`. + nbdtrik : Inverse with respect to `k` of `nbdtr(k, n, p)`. + + Notes + ----- + This function wraps routines from the Boost Math C++ library [1]_. + + Formula 26.5.26 of [2]_ or [3]_, + + .. math:: + \sum_{j=k + 1}^\infty {{n + j - 1} + \choose{j}} p^n (1 - p)^j = I_{1 - p}(k + 1, n), + + is used to reduce calculation of the cumulative distribution function to + that of a regularized incomplete beta :math:`I`. + + Computation of `n` involves a search for a value that produces the desired + value of `y`. The search relies on the monotonicity of `y` with `n`. + + References + ---------- + .. [1] The Boost Developers. "Boost C++ Libraries". https://www.boost.org/. + .. [2] Milton Abramowitz and Irene A. Stegun, eds. + Handbook of Mathematical Functions with Formulas, + Graphs, and Mathematical Tables. New York: Dover, 1972. + .. [3] NIST Digital Library of Mathematical Functions + https://dlmf.nist.gov/8.17.E24 + + Examples + -------- + Compute the negative binomial cumulative distribution function for an + exemplary parameter set. + + >>> from scipy.special import nbdtr, nbdtrin + >>> k, n, p = 5, 2, 0.5 + >>> cdf_value = nbdtr(k, n, p) + >>> cdf_value + 0.9375 + + Verify that `nbdtrin` recovers the original value for `n` up to floating + point accuracy. + + >>> nbdtrin(k, cdf_value, p) + 2.0 + """) + +add_newdoc("ncfdtr", + r""" + ncfdtr(dfn, dfd, nc, f, out=None) + + Cumulative distribution function of the non-central F distribution. + + The non-central F describes the distribution of, + + .. math:: + Z = \frac{X/d_n}{Y/d_d} + + where :math:`X` and :math:`Y` are independently distributed, with + :math:`X` distributed non-central :math:`\chi^2` with noncentrality + parameter `nc` and :math:`d_n` degrees of freedom, and :math:`Y` + distributed :math:`\chi^2` with :math:`d_d` degrees of freedom. + + Parameters + ---------- + dfn : array_like + Degrees of freedom of the numerator sum of squares. Range (0, inf). + dfd : array_like + Degrees of freedom of the denominator sum of squares. Range (0, inf). + nc : array_like + Noncentrality parameter. Range [0, inf). + f : array_like + Quantiles, i.e. the upper limit of integration. + out : ndarray, optional + Optional output array for the function results + + Returns + ------- + cdf : scalar or ndarray + The calculated CDF. If all inputs are scalar, the return will be a + float. Otherwise it will be an array. + + See Also + -------- + ncfdtri : Quantile function; inverse of `ncfdtr` with respect to `f`. + ncfdtridfd : Inverse of `ncfdtr` with respect to `dfd`. + ncfdtridfn : Inverse of `ncfdtr` with respect to `dfn`. + ncfdtrinc : Inverse of `ncfdtr` with respect to `nc`. + scipy.stats.ncf : Non-central F distribution. + + Notes + ----- + This function calculates the CDF of the non-central f distribution using + the Boost Math C++ library [1]_. + + The cumulative distribution function is computed using Formula 26.6.20 of + [2]_: + + .. math:: + F(d_n, d_d, n_c, f) = \sum_{j=0}^\infty e^{-n_c/2} + \frac{(n_c/2)^j}{j!} I_{x}(\frac{d_n}{2} + j, \frac{d_d}{2}), + + where :math:`I` is the regularized incomplete beta function, and + :math:`x = f d_n/(f d_n + d_d)`. + + Note that argument order of `ncfdtr` is different from that of the + similar ``cdf`` method of `scipy.stats.ncf`: `f` is the last + parameter of `ncfdtr` but the first parameter of ``scipy.stats.ncf.cdf``. + + References + ---------- + .. [1] The Boost Developers. "Boost C++ Libraries". https://www.boost.org/. + .. [2] Milton Abramowitz and Irene A. Stegun, eds. + Handbook of Mathematical Functions with Formulas, + Graphs, and Mathematical Tables. New York: Dover, 1972. + + Examples + -------- + >>> import numpy as np + >>> from scipy import special + >>> from scipy import stats + >>> import matplotlib.pyplot as plt + + Plot the CDF of the non-central F distribution, for nc=0. Compare with the + F-distribution from scipy.stats: + + >>> x = np.linspace(-1, 8, num=500) + >>> dfn = 3 + >>> dfd = 2 + >>> ncf_stats = stats.f.cdf(x, dfn, dfd) + >>> ncf_special = special.ncfdtr(dfn, dfd, 0, x) + + >>> fig = plt.figure() + >>> ax = fig.add_subplot(111) + >>> ax.plot(x, ncf_stats, 'b-', lw=3) + >>> ax.plot(x, ncf_special, 'r-') + >>> plt.show() + + """) + +add_newdoc("ncfdtri", + """ + ncfdtri(dfn, dfd, nc, p, out=None) + + Inverse with respect to `f` of the CDF of the non-central F distribution. + + See `ncfdtr` for more details. + + Parameters + ---------- + dfn : array_like + Degrees of freedom of the numerator sum of squares. Range (0, inf). + dfd : array_like + Degrees of freedom of the denominator sum of squares. Range (0, inf). + nc : array_like + Noncentrality parameter. Range [0, inf). + p : array_like + Value of the cumulative distribution function. Must be in the + range [0, 1]. + out : ndarray, optional + Optional output array for the function results + + Returns + ------- + f : scalar or ndarray + Quantiles, i.e., the upper limit of integration. + + See Also + -------- + ncfdtr : CDF of the non-central F distribution. + ncfdtridfd : Inverse of `ncfdtr` with respect to `dfd`. + ncfdtridfn : Inverse of `ncfdtr` with respect to `dfn`. + ncfdtrinc : Inverse of `ncfdtr` with respect to `nc`. + scipy.stats.ncf : Non-central F distribution. + + Notes + ----- + This function calculates the Quantile of the non-central f distribution + using the Boost Math C++ library [1]_. + + Note that argument order of `ncfdtri` is different from that of the + similar ``ppf`` method of `scipy.stats.ncf`. `p` is the last parameter + of `ncfdtri` but the first parameter of ``scipy.stats.ncf.ppf``. + + References + ---------- + .. [1] The Boost Developers. "Boost C++ Libraries". https://www.boost.org/. + + Examples + -------- + >>> from scipy.special import ncfdtr, ncfdtri + + Compute the CDF for several values of `f`: + + >>> f = [0.5, 1, 1.5] + >>> p = ncfdtr(2, 3, 1.5, f) + >>> p + array([ 0.20782291, 0.36107392, 0.47345752]) + + Compute the inverse. We recover the values of `f`, as expected: + + >>> ncfdtri(2, 3, 1.5, p) + array([ 0.5, 1. , 1.5]) + + """) + +add_newdoc("ncfdtridfd", + """ + ncfdtridfd(dfn, p, nc, f, out=None) + + Calculate degrees of freedom (denominator) for the noncentral F-distribution. + + This is the inverse with respect to `dfd` of `ncfdtr`. + See `ncfdtr` for more details. + + Parameters + ---------- + dfn : array_like + Degrees of freedom of the numerator sum of squares. Range (0, inf). + p : array_like + Value of the cumulative distribution function. Must be in the + range [0, 1]. + nc : array_like + Noncentrality parameter. Should be in range (0, 1e4). + f : array_like + Quantiles, i.e., the upper limit of integration. + out : ndarray, optional + Optional output array for the function results + + Returns + ------- + dfd : scalar or ndarray + Degrees of freedom of the denominator sum of squares. + + See Also + -------- + ncfdtr : CDF of the non-central F distribution. + ncfdtri : Quantile function; inverse of `ncfdtr` with respect to `f`. + ncfdtridfn : Inverse of `ncfdtr` with respect to `dfn`. + ncfdtrinc : Inverse of `ncfdtr` with respect to `nc`. + + Notes + ----- + The value of the cumulative noncentral F distribution is not necessarily + monotone in either degrees of freedom. There thus may be two values that + provide a given CDF value. This routine assumes monotonicity and will + find an arbitrary one of the two values. + + Examples + -------- + >>> from scipy.special import ncfdtr, ncfdtridfd + + Compute the CDF for several values of `dfd`: + + >>> dfd = [1, 2, 3] + >>> p = ncfdtr(2, dfd, 0.25, 15) + >>> p + array([ 0.8097138 , 0.93020416, 0.96787852]) + + Compute the inverse. We recover the values of `dfd`, as expected: + + >>> ncfdtridfd(2, p, 0.25, 15) + array([ 1., 2., 3.]) + + """) + +add_newdoc("ncfdtridfn", + """ + ncfdtridfn(p, dfd, nc, f, out=None) + + Calculate degrees of freedom (numerator) for the noncentral F-distribution. + + This is the inverse with respect to `dfn` of `ncfdtr`. + See `ncfdtr` for more details. + + Parameters + ---------- + p : array_like + Value of the cumulative distribution function. Must be in the + range [0, 1]. + dfd : array_like + Degrees of freedom of the denominator sum of squares. Range (0, inf). + nc : array_like + Noncentrality parameter. Should be in range (0, 1e4). + f : float + Quantiles, i.e., the upper limit of integration. + out : ndarray, optional + Optional output array for the function results + + Returns + ------- + dfn : scalar or ndarray + Degrees of freedom of the numerator sum of squares. + + See Also + -------- + ncfdtr : CDF of the non-central F distribution. + ncfdtri : Quantile function; inverse of `ncfdtr` with respect to `f`. + ncfdtridfd : Inverse of `ncfdtr` with respect to `dfd`. + ncfdtrinc : Inverse of `ncfdtr` with respect to `nc`. + + Notes + ----- + The value of the cumulative noncentral F distribution is not necessarily + monotone in either degrees of freedom. There thus may be two values that + provide a given CDF value. This routine assumes monotonicity and will + find an arbitrary one of the two values. + + Examples + -------- + >>> from scipy.special import ncfdtr, ncfdtridfn + + Compute the CDF for several values of `dfn`: + + >>> dfn = [1, 2, 3] + >>> p = ncfdtr(dfn, 2, 0.25, 15) + >>> p + array([ 0.92562363, 0.93020416, 0.93188394]) + + Compute the inverse. We recover the values of `dfn`, as expected: + + >>> ncfdtridfn(p, 2, 0.25, 15) + array([ 1., 2., 3.]) + + """) + +add_newdoc("ncfdtrinc", + """ + ncfdtrinc(dfn, dfd, p, f, out=None) + + Calculate non-centrality parameter for non-central F distribution. + + This is the inverse with respect to `nc` of `ncfdtr`. + See `ncfdtr` for more details. + + Parameters + ---------- + dfn : array_like + Degrees of freedom of the numerator sum of squares. Range (0, inf). + dfd : array_like + Degrees of freedom of the denominator sum of squares. Range (0, inf). + p : array_like + Value of the cumulative distribution function. Must be in the + range [0, 1]. + f : array_like + Quantiles, i.e., the upper limit of integration. + out : ndarray, optional + Optional output array for the function results + + Returns + ------- + nc : scalar or ndarray + Noncentrality parameter. + + See Also + -------- + ncfdtr : CDF of the non-central F distribution. + ncfdtri : Quantile function; inverse of `ncfdtr` with respect to `f`. + ncfdtridfd : Inverse of `ncfdtr` with respect to `dfd`. + ncfdtridfn : Inverse of `ncfdtr` with respect to `dfn`. + + Examples + -------- + >>> from scipy.special import ncfdtr, ncfdtrinc + + Compute the CDF for several values of `nc`: + + >>> nc = [0.5, 1.5, 2.0] + >>> p = ncfdtr(2, 3, nc, 15) + >>> p + array([ 0.96309246, 0.94327955, 0.93304098]) + + Compute the inverse. We recover the values of `nc`, as expected: + + >>> ncfdtrinc(2, 3, p, 15) + array([ 0.5, 1.5, 2. ]) + + """) + +add_newdoc("nctdtr", + """ + nctdtr(df, nc, t, out=None) + + Cumulative distribution function of the non-central `t` distribution. + + Parameters + ---------- + df : array_like + Degrees of freedom of the distribution. Should be in range (0, inf). + nc : array_like + Noncentrality parameter. + t : array_like + Quantiles, i.e., the upper limit of integration. + out : ndarray, optional + Optional output array for the function results + + Returns + ------- + cdf : scalar or ndarray + The calculated CDF. If all inputs are scalar, the return will be a + float. Otherwise, it will be an array. + + See Also + -------- + nctdtrit : Inverse CDF (iCDF) of the non-central t distribution. + nctdtridf : Calculate degrees of freedom, given CDF and iCDF values. + nctdtrinc : Calculate non-centrality parameter, given CDF iCDF values. + + Notes + ----- + This function calculates the CDF of the non-central t distribution using + the Boost Math C++ library [1]_. + + Note that the argument order of `nctdtr` is different from that of the + similar ``cdf`` method of `scipy.stats.nct`: `t` is the last + parameter of `nctdtr` but the first parameter of ``scipy.stats.nct.cdf``. + + References + ---------- + .. [1] The Boost Developers. "Boost C++ Libraries". https://www.boost.org/. + + Examples + -------- + >>> import numpy as np + >>> from scipy import special + >>> from scipy import stats + >>> import matplotlib.pyplot as plt + + Plot the CDF of the non-central t distribution, for nc=0. Compare with the + t-distribution from scipy.stats: + + >>> x = np.linspace(-5, 5, num=500) + >>> df = 3 + >>> nct_stats = stats.t.cdf(x, df) + >>> nct_special = special.nctdtr(df, 0, x) + + >>> fig = plt.figure() + >>> ax = fig.add_subplot(111) + >>> ax.plot(x, nct_stats, 'b-', lw=3) + >>> ax.plot(x, nct_special, 'r-') + >>> plt.show() + + """) + +add_newdoc("nctdtridf", + """ + nctdtridf(p, nc, t, out=None) + + Calculate degrees of freedom for non-central t distribution. + + See `nctdtr` for more details. + + Parameters + ---------- + p : array_like + CDF values, in range (0, 1]. + nc : array_like + Noncentrality parameter. Should be in range (-1e6, 1e6). + t : array_like + Quantiles, i.e., the upper limit of integration. + out : ndarray, optional + Optional output array for the function results + + Returns + ------- + df : scalar or ndarray + The degrees of freedom. If all inputs are scalar, the return will be a + float. Otherwise, it will be an array. + + See Also + -------- + nctdtr : CDF of the non-central `t` distribution. + nctdtrit : Inverse CDF (iCDF) of the non-central t distribution. + nctdtrinc : Calculate non-centrality parameter, given CDF iCDF values. + + Examples + -------- + >>> from scipy.special import nctdtr, nctdtridf + + Compute the CDF for several values of `df`: + + >>> df = [1, 2, 3] + >>> p = nctdtr(df, 0.25, 1) + >>> p + array([0.67491974, 0.716464 , 0.73349456]) + + Compute the inverse. We recover the values of `df`, as expected: + + >>> nctdtridf(p, 0.25, 1) + array([1., 2., 3.]) + + """) + +add_newdoc("nctdtrinc", + """ + nctdtrinc(df, p, t, out=None) + + Calculate non-centrality parameter for non-central t distribution. + + See `nctdtr` for more details. + + Parameters + ---------- + df : array_like + Degrees of freedom of the distribution. Should be in range (0, inf). + p : array_like + CDF values, in range (0, 1]. + t : array_like + Quantiles, i.e., the upper limit of integration. + out : ndarray, optional + Optional output array for the function results + + Returns + ------- + nc : scalar or ndarray + Noncentrality parameter + + See Also + -------- + nctdtr : CDF of the non-central `t` distribution. + nctdtrit : Inverse CDF (iCDF) of the non-central t distribution. + nctdtridf : Calculate degrees of freedom, given CDF and iCDF values. + + Examples + -------- + >>> from scipy.special import nctdtr, nctdtrinc + + Compute the CDF for several values of `nc`: + + >>> nc = [0.5, 1.5, 2.5] + >>> p = nctdtr(3, nc, 1.5) + >>> p + array([0.77569497, 0.45524533, 0.1668691 ]) + + Compute the inverse. We recover the values of `nc`, as expected: + + >>> nctdtrinc(3, p, 1.5) + array([0.5, 1.5, 2.5]) + + """) + +add_newdoc("nctdtrit", + """ + nctdtrit(df, nc, p, out=None) + + Inverse cumulative distribution function of the non-central t distribution. + + See `nctdtr` for more details. + + Parameters + ---------- + df : array_like + Degrees of freedom of the distribution. Should be in range (0, inf). + nc : array_like + Noncentrality parameter. + p : array_like + CDF values, in range (0, 1]. + out : ndarray, optional + Optional output array for the function results + + Returns + ------- + t : scalar or ndarray + Quantiles + + See Also + -------- + nctdtr : CDF of the non-central `t` distribution. + nctdtridf : Calculate degrees of freedom, given CDF and iCDF values. + nctdtrinc : Calculate non-centrality parameter, given CDF iCDF values. + + Notes + ----- + This function calculates the quantile of the non-central t distribution using + the Boost Math C++ library [1]_. + + Note that the argument order of `nctdtrit` is different from that of the + similar ``ppf`` method of `scipy.stats.nct`: `t` is the last + parameter of `nctdtrit` but the first parameter of ``scipy.stats.nct.ppf``. + + References + ---------- + .. [1] The Boost Developers. "Boost C++ Libraries". https://www.boost.org/. + + Examples + -------- + >>> from scipy.special import nctdtr, nctdtrit + + Compute the CDF for several values of `t`: + + >>> t = [0.5, 1, 1.5] + >>> p = nctdtr(3, 1, t) + >>> p + array([0.29811049, 0.46922687, 0.6257559 ]) + + Compute the inverse. We recover the values of `t`, as expected: + + >>> nctdtrit(3, 1, p) + array([0.5, 1. , 1.5]) + + """) + +add_newdoc("nrdtrimn", + """ + nrdtrimn(p, std, x, out=None) + + Calculate mean of normal distribution given other params. + + Parameters + ---------- + p : array_like + CDF values, in range (0, 1]. + std : array_like + Standard deviation. + x : array_like + Quantiles, i.e. the upper limit of integration. + out : ndarray, optional + Optional output array for the function results + + Returns + ------- + mn : scalar or ndarray + The mean of the normal distribution. + + See Also + -------- + scipy.stats.norm : Normal distribution + ndtr : Standard normal cumulative probability distribution + ndtri : Inverse of standard normal CDF with respect to quantile + nrdtrisd : Inverse of normal distribution CDF with respect to + standard deviation + + Examples + -------- + `nrdtrimn` can be used to recover the mean of a normal distribution + if we know the CDF value `p` for a given quantile `x` and the + standard deviation `std`. First, we calculate + the normal distribution CDF for an exemplary parameter set. + + >>> from scipy.stats import norm + >>> mean = 3. + >>> std = 2. + >>> x = 6. + >>> p = norm.cdf(x, loc=mean, scale=std) + >>> p + 0.9331927987311419 + + Verify that `nrdtrimn` returns the original value for `mean`. + + >>> from scipy.special import nrdtrimn + >>> nrdtrimn(p, std, x) + 3.0000000000000004 + + """) + +add_newdoc("nrdtrisd", + """ + nrdtrisd(mn, p, x, out=None) + + Calculate standard deviation of normal distribution given other params. + + Parameters + ---------- + mn : scalar or ndarray + The mean of the normal distribution. + p : array_like + CDF values, in range (0, 1]. + x : array_like + Quantiles, i.e. the upper limit of integration. + + out : ndarray, optional + Optional output array for the function results + + Returns + ------- + std : scalar or ndarray + Standard deviation. + + See Also + -------- + scipy.stats.norm : Normal distribution + ndtr : Standard normal cumulative probability distribution + ndtri : Inverse of standard normal CDF with respect to quantile + nrdtrimn : Inverse of normal distribution CDF with respect to + mean + + Examples + -------- + `nrdtrisd` can be used to recover the standard deviation of a normal + distribution if we know the CDF value `p` for a given quantile `x` and + the mean `mn`. First, we calculate the normal distribution CDF for an + exemplary parameter set. + + >>> from scipy.stats import norm + >>> mean = 3. + >>> std = 2. + >>> x = 6. + >>> p = norm.cdf(x, loc=mean, scale=std) + >>> p + 0.9331927987311419 + + Verify that `nrdtrisd` returns the original value for `std`. + + >>> from scipy.special import nrdtrisd + >>> nrdtrisd(mean, p, x) + 2.0000000000000004 + + """) + +add_newdoc("ndtri", + """ + ndtri(p, out=None) + + Inverse of `ndtr`. + + Returns the quantile `x` such that the cumulative distribution function of the + standard normal distribution evaluated at `x` equals `p`, that is, ``ndtr(x) == p``. + + Parameters + ---------- + p : array_like + Probability values. + out : ndarray, optional + Optional output array for the function results. + + Returns + ------- + x : scalar or ndarray + Quantile(s) corresponding to the probabilitie(s) in `p`. + + See Also + -------- + ndtr : Standard normal cumulative distribution function + ndtri_exp : Inverse of log_ndtr + + Examples + -------- + `ndtri` is the percentile (quantile) function of the standard normal distribution, + i.e., the inverse of the cumulative distribution function `ndtr`. + + First, compute a cumulative distribution value: + + >>> import numpy as np + >>> from scipy.special import ndtri, ndtr + >>> cdf_val = ndtr(2) + >>> cdf_val + 0.9772498680518208 + + Verify that `ndtri` yields the original value for `x` up to floating point errors. + + >>> ndtri(cdf_val) + 2.0000000000000004 + + Plot the percentile function over a range of probabilities. + + >>> import matplotlib.pyplot as plt + >>> p = np.linspace(1e-3, 1 - 1e-3, 201) + >>> fig, ax = plt.subplots() + >>> ax.plot(p, ndtri(p)) + >>> ax.set_title("Standard normal percentile function") + >>> plt.show() + """) + +add_newdoc("pdtr", + r""" + pdtr(k, m, out=None) + + Poisson cumulative distribution function. + + Defined as the probability that a Poisson-distributed random + variable with event rate :math:`m` is less than or equal to + :math:`k`. More concretely, this works out to be [1]_ + + .. math:: + + \exp(-m) \sum_{j = 0}^{\lfloor{k}\rfloor} \frac{m^j}{j!}. + + Parameters + ---------- + k : array_like + Number of occurrences (nonnegative, real) + m : array_like + Shape parameter (nonnegative, real) + out : ndarray, optional + Optional output array for the function results + + Returns + ------- + scalar or ndarray + Values of the Poisson cumulative distribution function + + See Also + -------- + pdtrc : Poisson survival function + pdtrik : inverse of `pdtr` with respect to `k` + pdtri : inverse of `pdtr` with respect to `m` + + References + ---------- + .. [1] https://en.wikipedia.org/wiki/Poisson_distribution + + Examples + -------- + >>> import numpy as np + >>> import scipy.special as sc + + It is a cumulative distribution function, so it converges to 1 + monotonically as `k` goes to infinity. + + >>> sc.pdtr([1, 10, 100, np.inf], 1) + array([0.73575888, 0.99999999, 1. , 1. ]) + + It is discontinuous at integers and constant between integers. + + >>> sc.pdtr([1, 1.5, 1.9, 2], 1) + array([0.73575888, 0.73575888, 0.73575888, 0.9196986 ]) + + """) + +add_newdoc("pdtrc", + """ + pdtrc(k, m, out=None) + + Poisson survival function. + + Returns the sum of the terms from k+1 to infinity of the Poisson + distribution: sum(exp(-m) * m**j / j!, j=k+1..inf) = gammainc( + k+1, m). Arguments must both be non-negative doubles. + + Parameters + ---------- + k : array_like + Number of occurrences (nonnegative, real) + m : array_like + Shape parameter (nonnegative, real) + out : ndarray, optional + Optional output array for the function results + + Returns + ------- + scalar or ndarray + Values of the Poisson survival function + + See Also + -------- + pdtr : Poisson cumulative distribution function + pdtrik : inverse of `pdtr` with respect to `k` + pdtri : inverse of `pdtr` with respect to `m` + + Examples + -------- + >>> import numpy as np + >>> import scipy.special as sc + + It is a survival function, so it decreases to 0 + monotonically as `k` goes to infinity. + + >>> k = np.array([1, 10, 100, np.inf]) + >>> sc.pdtrc(k, 1) + array([2.64241118e-001, 1.00477664e-008, 3.94147589e-161, 0.00000000e+000]) + + It can be expressed in terms of the lower incomplete gamma + function `gammainc`. + + >>> sc.gammainc(k + 1, 1) + array([2.64241118e-001, 1.00477664e-008, 3.94147589e-161, 0.00000000e+000]) + + """) + +add_newdoc("pdtri", + """ + pdtri(k, y, out=None) + + Inverse of `pdtr` with respect to `m`. + + Returns the Poisson variable `m` such that the sum from 0 to `k` of + the Poisson density is equal to the given probability `y`: + calculated by ``gammainccinv(k + 1, y)``. `k` must be a nonnegative + integer and `y` between 0 and 1. + + Parameters + ---------- + k : array_like + Number of occurrences (nonnegative, real). + y : array_like + Probability. + out : ndarray, optional + Optional output array for the function results. + + Returns + ------- + scalar or ndarray + Values of the shape parameter `m` such that ``pdtr(k, m) = y``. + + See Also + -------- + pdtr : Poisson cumulative distribution function + pdtrc : Poisson survival function + pdtrik : Inverse of `pdtr` with respect to `k` + + Examples + -------- + >>> import scipy.special as sc + + Compute the CDF for several values of `m`: + + >>> k = 1 + >>> m = [0.5, 1, 1.5] + >>> p = sc.pdtr(k, m) + >>> p + array([0.90979599, 0.73575888, 0.5578254 ]) + + Invert the CDF with respect to the Poisson mean. We recover the values + of `m`, as expected: + + >>> sc.pdtri(k, p) + array([0.5, 1. , 1.5]) + + Verify the relation with `gammainccinv`: + + >>> sc.gammainccinv(k + 1, p) + array([0.5, 1. , 1.5]) + + + """) + +add_newdoc("pdtrik", + """ + pdtrik(p, m, out=None) + + Inverse of `pdtr` with respect to `k`. + + Parameters + ---------- + p : array_like + Probability + m : array_like + Shape parameter (nonnegative, real) + out : ndarray, optional + Optional output array for the function results + + Returns + ------- + scalar or ndarray + The number of occurrences `k` such that ``pdtr(k, m) = p`` + + See Also + -------- + pdtr : Poisson cumulative distribution function + pdtrc : Poisson survival function + pdtri : inverse of `pdtr` with respect to `m` + + Notes + ----- + This function relies on the ``gamma_q_inva`` function from the Boost + Math C++ library [1]_. + + References + ---------- + .. [1] The Boost Developers. "Boost C++ Libraries". https://www.boost.org/. + + Examples + -------- + >>> import scipy.special as sc + + Compute the CDF for several values of `k`: + + >>> k = [1, 2, 3] + >>> p = sc.pdtr(k, 2) + >>> p + array([0.40600585, 0.67667642, 0.85712346]) + + Compute the inverse. We recover the values of `k`, as expected: + + >>> sc.pdtrik(p, 2) + array([1., 2., 3.]) + + """) + +add_newdoc("poch", + r""" + poch(z, m, out=None) + + Pochhammer symbol. + + The Pochhammer symbol (rising factorial) is defined as + + .. math:: + + (z)_m = \frac{\Gamma(z + m)}{\Gamma(z)} + + For positive integer `m` it reads + + .. math:: + + (z)_m = z (z + 1) ... (z + m - 1) + + See [DLMF]_ for more details. + + Parameters + ---------- + z, m : array_like + Real-valued arguments. + out : ndarray, optional + Optional output array for the function results + + Returns + ------- + scalar or ndarray + The value of the function. + + References + ---------- + .. [DLMF] Nist, Digital Library of Mathematical Functions + https://dlmf.nist.gov/5.2#iii + + Examples + -------- + >>> import scipy.special as sc + + It is 1 when m is 0. + + >>> sc.poch([1, 2, 3, 4], 0) + array([1., 1., 1., 1.]) + + For z equal to 1 it reduces to the factorial function. + + >>> sc.poch(1, 5) + 120.0 + >>> 1 * 2 * 3 * 4 * 5 + 120 + + It can be expressed in terms of the gamma function. + + >>> z, m = 3.7, 2.1 + >>> sc.poch(z, m) + 20.529581933776953 + >>> sc.gamma(z + m) / sc.gamma(z) + 20.52958193377696 + + """) + +add_newdoc("powm1", """ + powm1(x, y, out=None) + + Computes ``x**y - 1``. + + This function is useful when `y` is near 0, or when `x` is near 1. + + The function is implemented for real types only (unlike ``numpy.power``, + which accepts complex inputs). + + Parameters + ---------- + x : array_like + The base. Must be a real type (i.e. integer or float, not complex). + y : array_like + The exponent. Must be a real type (i.e. integer or float, not complex). + + Returns + ------- + array_like + Result of the calculation + + Notes + ----- + .. versionadded:: 1.10.0 + + The underlying code is implemented for single precision and double + precision floats only. Unlike `numpy.power`, integer inputs to + `powm1` are converted to floating point, and complex inputs are + not accepted. + + Note the following edge cases: + + * ``powm1(x, 0)`` returns 0 for any ``x``, including 0, ``inf`` + and ``nan``. + * ``powm1(1, y)`` returns 0 for any ``y``, including ``nan`` + and ``inf``. + + This function wraps the ``powm1`` routine from the + Boost Math C++ library [1]_. + + References + ---------- + .. [1] The Boost Developers. "Boost C++ Libraries". https://www.boost.org/. + + Examples + -------- + >>> import numpy as np + >>> from scipy.special import powm1 + + >>> x = np.array([1.2, 10.0, 0.9999999975]) + >>> y = np.array([1e-9, 1e-11, 0.1875]) + >>> powm1(x, y) + array([ 1.82321557e-10, 2.30258509e-11, -4.68749998e-10]) + + It can be verified that the relative errors in those results + are less than 2.5e-16. + + Compare that to the result of ``x**y - 1``, where the + relative errors are all larger than 8e-8: + + >>> x**y - 1 + array([ 1.82321491e-10, 2.30258035e-11, -4.68750039e-10]) + + """) + + +add_newdoc("pseudo_huber", + r""" + pseudo_huber(delta, r, out=None) + + Pseudo-Huber loss function. + + .. math:: \mathrm{pseudo\_huber}(\delta, r) = + \delta^2 \left( \sqrt{ 1 + \left( \frac{r}{\delta} \right)^2 } - 1 \right) + + Parameters + ---------- + delta : array_like + Input array, indicating the soft quadratic vs. linear loss changepoint. + r : array_like + Input array, possibly representing residuals. + out : ndarray, optional + Optional output array for the function results + + Returns + ------- + res : scalar or ndarray + The computed Pseudo-Huber loss function values. + + See Also + -------- + huber: Similar function which this function approximates + + Notes + ----- + Like `huber`, `pseudo_huber` often serves as a robust loss function + in statistics or machine learning to reduce the influence of outliers. + Unlike `huber`, `pseudo_huber` is smooth. + + Typically, `r` represents residuals, the difference + between a model prediction and data. Then, for :math:`|r|\leq\delta`, + `pseudo_huber` resembles the squared error and for :math:`|r|>\delta` the + absolute error. This way, the Pseudo-Huber loss often achieves + a fast convergence in model fitting for small residuals like the squared + error loss function and still reduces the influence of outliers + (:math:`|r|>\delta`) like the absolute error loss. As :math:`\delta` is + the cutoff between squared and absolute error regimes, it has + to be tuned carefully for each problem. `pseudo_huber` is also + convex, making it suitable for gradient based optimization. [1]_ [2]_ + + .. versionadded:: 0.15.0 + + References + ---------- + .. [1] Hartley, Zisserman, "Multiple View Geometry in Computer Vision". + 2003. Cambridge University Press. p. 619 + .. [2] Charbonnier et al. "Deterministic edge-preserving regularization + in computed imaging". 1997. IEEE Trans. Image Processing. + 6 (2): 298 - 311. + + Examples + -------- + Import all necessary modules. + + >>> import numpy as np + >>> from scipy.special import pseudo_huber, huber + >>> import matplotlib.pyplot as plt + + Calculate the function for ``delta=1`` at ``r=2``. + + >>> pseudo_huber(1., 2.) + 1.2360679774997898 + + Calculate the function at ``r=2`` for different `delta` by providing + a list or NumPy array for `delta`. + + >>> pseudo_huber([1., 2., 4.], 3.) + array([2.16227766, 3.21110255, 4. ]) + + Calculate the function for ``delta=1`` at several points by providing + a list or NumPy array for `r`. + + >>> pseudo_huber(2., np.array([1., 1.5, 3., 4.])) + array([0.47213595, 1. , 3.21110255, 4.94427191]) + + The function can be calculated for different `delta` and `r` by + providing arrays for both with compatible shapes for broadcasting. + + >>> r = np.array([1., 2.5, 8., 10.]) + >>> deltas = np.array([[1.], [5.], [9.]]) + >>> print(r.shape, deltas.shape) + (4,) (3, 1) + + >>> pseudo_huber(deltas, r) + array([[ 0.41421356, 1.6925824 , 7.06225775, 9.04987562], + [ 0.49509757, 2.95084972, 22.16990566, 30.90169944], + [ 0.49846624, 3.06693762, 27.37435121, 40.08261642]]) + + Plot the function for different `delta`. + + >>> x = np.linspace(-4, 4, 500) + >>> deltas = [1, 2, 3] + >>> linestyles = ["dashed", "dotted", "dashdot"] + >>> fig, ax = plt.subplots() + >>> combined_plot_parameters = list(zip(deltas, linestyles)) + >>> for delta, style in combined_plot_parameters: + ... ax.plot(x, pseudo_huber(delta, x), label=rf"$\delta={delta}$", + ... ls=style) + >>> ax.legend(loc="upper center") + >>> ax.set_xlabel("$x$") + >>> ax.set_title(r"Pseudo-Huber loss function $h_{\delta}(x)$") + >>> ax.set_xlim(-4, 4) + >>> ax.set_ylim(0, 8) + >>> plt.show() + + Finally, illustrate the difference between `huber` and `pseudo_huber` by + plotting them and their gradients with respect to `r`. The plot shows + that `pseudo_huber` is continuously differentiable while `huber` is not + at the points :math:`\pm\delta`. + + >>> def huber_grad(delta, x): + ... grad = np.copy(x) + ... linear_area = np.argwhere(np.abs(x) > delta) + ... grad[linear_area]=delta*np.sign(x[linear_area]) + ... return grad + >>> def pseudo_huber_grad(delta, x): + ... return x* (1+(x/delta)**2)**(-0.5) + >>> x=np.linspace(-3, 3, 500) + >>> delta = 1. + >>> fig, ax = plt.subplots(figsize=(7, 7)) + >>> ax.plot(x, huber(delta, x), label="Huber", ls="dashed") + >>> ax.plot(x, huber_grad(delta, x), label="Huber Gradient", ls="dashdot") + >>> ax.plot(x, pseudo_huber(delta, x), label="Pseudo-Huber", ls="dotted") + >>> ax.plot(x, pseudo_huber_grad(delta, x), label="Pseudo-Huber Gradient", + ... ls="solid") + >>> ax.legend(loc="upper center") + >>> plt.show() + """) + +add_newdoc("rel_entr", + r""" + rel_entr(x, y, out=None) + + Elementwise function for computing relative entropy. + + .. math:: + + \mathrm{rel\_entr}(x, y) = + \begin{cases} + x \log(x / y) & x > 0, y > 0 \\ + 0 & x = 0, y \ge 0 \\ + \infty & \text{otherwise} + \end{cases} + + Parameters + ---------- + x, y : array_like + Input arrays + out : ndarray, optional + Optional output array for the function results + + Returns + ------- + scalar or ndarray + Relative entropy of the inputs + + See Also + -------- + entr, kl_div, scipy.stats.entropy + + Notes + ----- + .. versionadded:: 0.15.0 + + This function is jointly convex in x and y. + + The origin of this function is in convex programming; see + [1]_. Given two discrete probability distributions :math:`p_1, + \ldots, p_n` and :math:`q_1, \ldots, q_n`, the definition of relative + entropy in the context of *information theory* is + + .. math:: + + \sum_{i = 1}^n \mathrm{rel\_entr}(p_i, q_i). + + To compute the latter quantity, use `scipy.stats.entropy`. + + See [2]_ for details. + + References + ---------- + .. [1] Boyd, Stephen and Lieven Vandenberghe. *Convex optimization*. + Cambridge University Press, 2004. + :doi:`10.1017/CBO9780511804441`. + .. [2] Kullback-Leibler divergence, + https://en.wikipedia.org/wiki/Kullback%E2%80%93Leibler_divergence + + """) + +add_newdoc("round", + """ + round(x, out=None) + + Round to the nearest integer. + + Returns the nearest integer to `x`. If `x` ends in 0.5 exactly, + the nearest even integer is chosen. + + Parameters + ---------- + x : array_like + Real valued input. + out : ndarray, optional + Optional output array for the function results. + + Returns + ------- + scalar or ndarray + The nearest integers to the elements of `x`. The result is of + floating type, not integer type. + + Examples + -------- + >>> import scipy.special as sc + + It rounds to even. + + >>> sc.round([0.5, 1.5]) + array([0., 2.]) + + """) + +add_newdoc("shichi", + r""" + shichi(x, out=None) + + Hyperbolic sine and cosine integrals. + + The hyperbolic sine integral is + + .. math:: + + \int_0^x \frac{\sinh{t}}{t}dt + + and the hyperbolic cosine integral is + + .. math:: + + \gamma + \log(x) + \int_0^x \frac{\cosh{t} - 1}{t} dt + + where :math:`\gamma` is Euler's constant and :math:`\log` is the + principal branch of the logarithm [1]_ (see also [2]_). + + Parameters + ---------- + x : array_like + Real or complex points at which to compute the hyperbolic sine + and cosine integrals. + out : tuple of ndarray, optional + Optional output arrays for the function results + + Returns + ------- + si : scalar or ndarray + Hyperbolic sine integral at ``x`` + ci : scalar or ndarray + Hyperbolic cosine integral at ``x`` + + See Also + -------- + sici : Sine and cosine integrals. + exp1 : Exponential integral E1. + expi : Exponential integral Ei. + + Notes + ----- + For real arguments with ``x < 0``, ``chi`` is the real part of the + hyperbolic cosine integral. For such points ``chi(x)`` and ``chi(x + + 0j)`` differ by a factor of ``1j*pi``. + + For real arguments the function is computed by calling Cephes' + [3]_ *shichi* routine. For complex arguments the algorithm is based + on Mpmath's [4]_ *shi* and *chi* routines. + + References + ---------- + .. [1] Milton Abramowitz and Irene A. Stegun, eds. + Handbook of Mathematical Functions with Formulas, + Graphs, and Mathematical Tables. New York: Dover, 1972. + (See Section 5.2.) + .. [2] NIST Digital Library of Mathematical Functions + https://dlmf.nist.gov/6.2.E15 and https://dlmf.nist.gov/6.2.E16 + .. [3] Cephes Mathematical Functions Library, + https://netlib.org/cephes/ + .. [4] Fredrik Johansson and others. + "mpmath: a Python library for arbitrary-precision floating-point + arithmetic" (Version 0.19) https://mpmath.org/ + + Examples + -------- + >>> import numpy as np + >>> import matplotlib.pyplot as plt + >>> from scipy.special import shichi, sici + + `shichi` accepts real or complex input: + + >>> shichi(0.5) + (0.5069967498196671, -0.05277684495649357) + >>> shichi(0.5 + 2.5j) + ((0.11772029666668238+1.831091777729851j), + (0.29912435887648825+1.7395351121166562j)) + + The hyperbolic sine and cosine integrals Shi(z) and Chi(z) are + related to the sine and cosine integrals Si(z) and Ci(z) by + + * Shi(z) = -i*Si(i*z) + * Chi(z) = Ci(-i*z) + i*pi/2 + + >>> z = 0.25 + 5j + >>> shi, chi = shichi(z) + >>> shi, -1j*sici(1j*z)[0] # Should be the same. + ((-0.04834719325101729+1.5469354086921228j), + (-0.04834719325101729+1.5469354086921228j)) + >>> chi, sici(-1j*z)[1] + 1j*np.pi/2 # Should be the same. + ((-0.19568708973868087+1.556276312103824j), + (-0.19568708973868087+1.556276312103824j)) + + Plot the functions evaluated on the real axis: + + >>> xp = np.geomspace(1e-8, 4.0, 250) + >>> x = np.concatenate((-xp[::-1], xp)) + >>> shi, chi = shichi(x) + + >>> fig, ax = plt.subplots() + >>> ax.plot(x, shi, label='Shi(x)') + >>> ax.plot(x, chi, '--', label='Chi(x)') + >>> ax.set_xlabel('x') + >>> ax.set_title('Hyperbolic Sine and Cosine Integrals') + >>> ax.legend(shadow=True, framealpha=1, loc='lower right') + >>> ax.grid(True) + >>> plt.show() + + """) + +add_newdoc("sici", + r""" + sici(x, out=None) + + Sine and cosine integrals. + + The sine integral is + + .. math:: + + \int_0^x \frac{\sin{t}}{t}dt + + and the cosine integral is + + .. math:: + + \gamma + \log(x) + \int_0^x \frac{\cos{t} - 1}{t}dt + + where :math:`\gamma` is Euler's constant and :math:`\log` is the + principal branch of the logarithm [1]_ (see also [2]_). + + Parameters + ---------- + x : array_like + Real or complex points at which to compute the sine and cosine + integrals. + out : tuple of ndarray, optional + Optional output arrays for the function results + + Returns + ------- + si : scalar or ndarray + Sine integral at ``x`` + ci : scalar or ndarray + Cosine integral at ``x`` + + See Also + -------- + shichi : Hyperbolic sine and cosine integrals. + exp1 : Exponential integral E1. + expi : Exponential integral Ei. + + Notes + ----- + For real arguments with ``x < 0``, ``ci`` is the real part of the + cosine integral. For such points ``ci(x)`` and ``ci(x + 0j)`` + differ by a factor of ``1j*pi``. + + For real arguments the function is computed by calling Cephes' + [3]_ *sici* routine. For complex arguments the algorithm is based + on Mpmath's [4]_ *si* and *ci* routines. + + References + ---------- + .. [1] Milton Abramowitz and Irene A. Stegun, eds. + Handbook of Mathematical Functions with Formulas, + Graphs, and Mathematical Tables. New York: Dover, 1972. + (See Section 5.2.) + .. [2] NIST Digital Library of Mathematical Functions + https://dlmf.nist.gov/6.2.E9, https://dlmf.nist.gov/6.2.E12, + and https://dlmf.nist.gov/6.2.E13 + .. [3] Cephes Mathematical Functions Library, + https://netlib.org/cephes/ + .. [4] Fredrik Johansson and others. + "mpmath: a Python library for arbitrary-precision floating-point + arithmetic" (Version 0.19) https://mpmath.org/ + + Examples + -------- + >>> import numpy as np + >>> import matplotlib.pyplot as plt + >>> from scipy.special import sici, exp1 + + `sici` accepts real or complex input: + + >>> sici(2.5) + (1.7785201734438267, 0.2858711963653835) + >>> sici(2.5 + 3j) + ((4.505735874563953+0.06863305018999577j), + (0.0793644206906966-2.935510262937543j)) + + For z in the right half plane, the sine and cosine integrals are + related to the exponential integral E1 (implemented in SciPy as + `scipy.special.exp1`) by + + * Si(z) = (E1(i*z) - E1(-i*z))/2i + pi/2 + * Ci(z) = -(E1(i*z) + E1(-i*z))/2 + + See [1]_ (equations 5.2.21 and 5.2.23). + + We can verify these relations: + + >>> z = 2 - 3j + >>> sici(z) + ((4.54751388956229-1.3991965806460565j), + (1.408292501520851+2.9836177420296055j)) + + >>> (exp1(1j*z) - exp1(-1j*z))/2j + np.pi/2 # Same as sine integral + (4.54751388956229-1.3991965806460565j) + + >>> -(exp1(1j*z) + exp1(-1j*z))/2 # Same as cosine integral + (1.408292501520851+2.9836177420296055j) + + Plot the functions evaluated on the real axis; the dotted horizontal + lines are at pi/2 and -pi/2: + + >>> x = np.linspace(-16, 16, 150) + >>> si, ci = sici(x) + + >>> fig, ax = plt.subplots() + >>> ax.plot(x, si, label='Si(x)') + >>> ax.plot(x, ci, '--', label='Ci(x)') + >>> ax.legend(shadow=True, framealpha=1, loc='upper left') + >>> ax.set_xlabel('x') + >>> ax.set_title('Sine and Cosine Integrals') + >>> ax.axhline(np.pi/2, linestyle=':', alpha=0.5, color='k') + >>> ax.axhline(-np.pi/2, linestyle=':', alpha=0.5, color='k') + >>> ax.grid(True) + >>> plt.show() + + """) + +add_newdoc("smirnov", + r""" + smirnov(n, d, out=None) + + Kolmogorov-Smirnov complementary cumulative distribution function. + + Returns the exact Kolmogorov-Smirnov complementary cumulative + distribution function,(aka the Survival Function) of Dn+ (or Dn-) + for a one-sided test of equality between an empirical and a + theoretical distribution. It is equal to the probability that the + maximum difference between a theoretical distribution and an empirical + one based on `n` samples is greater than d. + + Parameters + ---------- + n : int + Number of samples + d : float array_like + Deviation between the Empirical CDF (ECDF) and the target CDF. + out : ndarray, optional + Optional output array for the function results + + Returns + ------- + scalar or ndarray + The value(s) of smirnov(n, d), Prob(Dn+ >= d) (Also Prob(Dn- >= d)) + + See Also + -------- + smirnovi : The Inverse Survival Function for the distribution + scipy.stats.ksone : Provides the functionality as a continuous distribution + kolmogorov, kolmogi : Functions for the two-sided distribution + + Notes + ----- + `smirnov` is used by `stats.kstest` in the application of the + Kolmogorov-Smirnov Goodness of Fit test. For historical reasons this + function is exposed in `scpy.special`, but the recommended way to achieve + the most accurate CDF/SF/PDF/PPF/ISF computations is to use the + `stats.ksone` distribution. + + Examples + -------- + >>> import numpy as np + >>> from scipy.special import smirnov + >>> from scipy.stats import norm + + Show the probability of a gap at least as big as 0, 0.5 and 1.0 for a + sample of size 5. + + >>> smirnov(5, [0, 0.5, 1.0]) + array([ 1. , 0.056, 0. ]) + + Compare a sample of size 5 against N(0, 1), the standard normal + distribution with mean 0 and standard deviation 1. + + `x` is the sample. + + >>> x = np.array([-1.392, -0.135, 0.114, 0.190, 1.82]) + + >>> target = norm(0, 1) + >>> cdfs = target.cdf(x) + >>> cdfs + array([0.0819612 , 0.44630594, 0.5453811 , 0.57534543, 0.9656205 ]) + + Construct the empirical CDF and the K-S statistics (Dn+, Dn-, Dn). + + >>> n = len(x) + >>> ecdfs = np.arange(n+1, dtype=float)/n + >>> cols = np.column_stack([x, ecdfs[1:], cdfs, cdfs - ecdfs[:n], + ... ecdfs[1:] - cdfs]) + >>> with np.printoptions(precision=3): + ... print(cols) + [[-1.392 0.2 0.082 0.082 0.118] + [-0.135 0.4 0.446 0.246 -0.046] + [ 0.114 0.6 0.545 0.145 0.055] + [ 0.19 0.8 0.575 -0.025 0.225] + [ 1.82 1. 0.966 0.166 0.034]] + >>> gaps = cols[:, -2:] + >>> Dnpm = np.max(gaps, axis=0) + >>> print(f'Dn-={Dnpm[0]:f}, Dn+={Dnpm[1]:f}') + Dn-=0.246306, Dn+=0.224655 + >>> probs = smirnov(n, Dnpm) + >>> print(f'For a sample of size {n} drawn from N(0, 1):', + ... f' Smirnov n={n}: Prob(Dn- >= {Dnpm[0]:f}) = {probs[0]:.4f}', + ... f' Smirnov n={n}: Prob(Dn+ >= {Dnpm[1]:f}) = {probs[1]:.4f}', + ... sep='\n') + For a sample of size 5 drawn from N(0, 1): + Smirnov n=5: Prob(Dn- >= 0.246306) = 0.4711 + Smirnov n=5: Prob(Dn+ >= 0.224655) = 0.5245 + + Plot the empirical CDF and the standard normal CDF. + + >>> import matplotlib.pyplot as plt + >>> plt.step(np.concatenate(([-2.5], x, [2.5])), + ... np.concatenate((ecdfs, [1])), + ... where='post', label='Empirical CDF') + >>> xx = np.linspace(-2.5, 2.5, 100) + >>> plt.plot(xx, target.cdf(xx), '--', label='CDF for N(0, 1)') + + Add vertical lines marking Dn+ and Dn-. + + >>> iminus, iplus = np.argmax(gaps, axis=0) + >>> plt.vlines([x[iminus]], ecdfs[iminus], cdfs[iminus], color='r', + ... alpha=0.5, lw=4) + >>> plt.vlines([x[iplus]], cdfs[iplus], ecdfs[iplus+1], color='m', + ... alpha=0.5, lw=4) + + >>> plt.grid(True) + >>> plt.legend(framealpha=1, shadow=True) + >>> plt.show() + """) + +add_newdoc("smirnovi", + """ + smirnovi(n, p, out=None) + + Inverse to `smirnov`. + + Returns `d` such that ``smirnov(n, d) == p``, the critical value + corresponding to `p`. + + Parameters + ---------- + n : int + Number of samples + p : float array_like + Probability + out : ndarray, optional + Optional output array for the function results + + Returns + ------- + scalar or ndarray + The value(s) of smirnovi(n, p), the critical values. + + See Also + -------- + smirnov : The Survival Function (SF) for the distribution + scipy.stats.ksone : Provides the functionality as a continuous distribution + kolmogorov, kolmogi : Functions for the two-sided distribution + scipy.stats.kstwobign : Two-sided Kolmogorov-Smirnov distribution, large n + + Notes + ----- + `smirnov` is used by `stats.kstest` in the application of the + Kolmogorov-Smirnov Goodness of Fit test. For historical reasons this + function is exposed in `scpy.special`, but the recommended way to achieve + the most accurate CDF/SF/PDF/PPF/ISF computations is to use the + `stats.ksone` distribution. + + Examples + -------- + >>> from scipy.special import smirnovi, smirnov + + >>> n = 24 + >>> deviations = [0.1, 0.2, 0.3] + + Use `smirnov` to compute the complementary CDF of the Smirnov + distribution for the given number of samples and deviations. + + >>> p = smirnov(n, deviations) + >>> p + array([0.58105083, 0.12826832, 0.01032231]) + + The inverse function ``smirnovi(n, p)`` returns ``deviations``. + + >>> smirnovi(n, p) + array([0.1, 0.2, 0.3]) + + """) + +add_newdoc("_smirnovc", + """ + _smirnovc(n, d) + Internal function, do not use. + """) + +add_newdoc("_smirnovci", + """ + Internal function, do not use. + """) + +add_newdoc("_smirnovp", + """ + _smirnovp(n, p) + Internal function, do not use. + """) + +add_newdoc("spence", + r""" + spence(z, out=None) + + Spence's function, also known as the dilogarithm. + + It is defined to be + + .. math:: + \int_1^z \frac{\log(t)}{1 - t}dt + + for complex :math:`z`, where the contour of integration is taken + to avoid the branch cut of the logarithm. Spence's function is + analytic everywhere except the negative real axis where it has a + branch cut. + + Parameters + ---------- + z : array_like + Points at which to evaluate Spence's function + out : ndarray, optional + Optional output array for the function results + + Returns + ------- + s : scalar or ndarray + Computed values of Spence's function + + Notes + ----- + There is a different convention which defines Spence's function by + the integral + + .. math:: + -\int_0^z \frac{\log(1 - t)}{t}dt; + + this is our ``spence(1 - z)``. + + Examples + -------- + >>> import numpy as np + >>> from scipy.special import spence + >>> import matplotlib.pyplot as plt + + The function is defined for complex inputs: + + >>> spence([1-1j, 1.5+2j, 3j, -10-5j]) + array([-0.20561676+0.91596559j, -0.86766909-1.39560134j, + -0.59422064-2.49129918j, -1.14044398+6.80075924j]) + + For complex inputs on the branch cut, which is the negative real axis, + the function returns the limit for ``z`` with positive imaginary part. + For example, in the following, note the sign change of the imaginary + part of the output for ``z = -2`` and ``z = -2 - 1e-8j``: + + >>> spence([-2 + 1e-8j, -2, -2 - 1e-8j]) + array([2.32018041-3.45139229j, 2.32018042-3.4513923j , + 2.32018041+3.45139229j]) + + The function returns ``nan`` for real inputs on the branch cut: + + >>> spence(-1.5) + nan + + Verify some particular values: ``spence(0) = pi**2/6``, + ``spence(1) = 0`` and ``spence(2) = -pi**2/12``. + + >>> spence([0, 1, 2]) + array([ 1.64493407, 0. , -0.82246703]) + >>> np.pi**2/6, -np.pi**2/12 + (1.6449340668482264, -0.8224670334241132) + + Verify the identity:: + + spence(z) + spence(1 - z) = pi**2/6 - log(z)*log(1 - z) + + >>> z = 3 + 4j + >>> spence(z) + spence(1 - z) + (-2.6523186143876067+1.8853470951513935j) + >>> np.pi**2/6 - np.log(z)*np.log(1 - z) + (-2.652318614387606+1.885347095151394j) + + Plot the function for positive real input. + + >>> fig, ax = plt.subplots() + >>> x = np.linspace(0, 6, 400) + >>> ax.plot(x, spence(x)) + >>> ax.grid() + >>> ax.set_xlabel('x') + >>> ax.set_title('spence(x)') + >>> plt.show() + """) + +add_newdoc( + "stdtr", + r""" + stdtr(df, t, out=None) + + Student t distribution cumulative distribution function. + + Returns the integral: + + .. math:: + \frac{\Gamma((df+1)/2)}{\sqrt{\pi df} \Gamma(df/2)} + \int_{-\infty}^t (1+x^2/df)^{-(df+1)/2}\, dx + + Parameters + ---------- + df : array_like + Degrees of freedom + t : array_like + Upper bound of the integral + out : ndarray, optional + Optional output array for the function results + + Returns + ------- + scalar or ndarray + Value of the Student t CDF at t + + See Also + -------- + stdtridf : inverse of stdtr with respect to `df` + stdtrit : inverse of stdtr with respect to `t` + scipy.stats.t : student t distribution + + Notes + ----- + The student t distribution is also available as `scipy.stats.t`. + Calling `stdtr` directly can improve performance compared to the + ``cdf`` method of `scipy.stats.t` (see last example below). + + The function is computed using the Boost Math library [1]_, which + relies on the incomplete beta function. + + References + ---------- + .. [1] Boost C++ Libraries, https://www.boost.org/ + + Examples + -------- + Calculate the function for ``df=3`` at ``t=1``. + + >>> import numpy as np + >>> from scipy.special import stdtr + >>> import matplotlib.pyplot as plt + >>> stdtr(3, 1) + 0.8044988905221148 + + Plot the function for three different degrees of freedom. + + >>> x = np.linspace(-10, 10, 1000) + >>> fig, ax = plt.subplots() + >>> parameters = [(1, "solid"), (3, "dashed"), (10, "dotted")] + >>> for (df, linestyle) in parameters: + ... ax.plot(x, stdtr(df, x), ls=linestyle, label=f"$df={df}$") + >>> ax.legend() + >>> ax.set_title("Student t distribution cumulative distribution function") + >>> plt.show() + + The function can be computed for several degrees of freedom at the same + time by providing a NumPy array or list for `df`: + + >>> stdtr([1, 2, 3], 1) + array([0.75 , 0.78867513, 0.80449889]) + + It is possible to calculate the function at several points for several + different degrees of freedom simultaneously by providing arrays for `df` + and `t` with shapes compatible for broadcasting. Compute `stdtr` at + 4 points for 3 degrees of freedom resulting in an array of shape 3x4. + + >>> dfs = np.array([[1], [2], [3]]) + >>> t = np.array([2, 4, 6, 8]) + >>> dfs.shape, t.shape + ((3, 1), (4,)) + + >>> stdtr(dfs, t) + array([[0.85241638, 0.92202087, 0.94743154, 0.96041658], + [0.90824829, 0.97140452, 0.98666426, 0.99236596], + [0.93033702, 0.98599577, 0.99536364, 0.99796171]]) + + The t distribution is also available as `scipy.stats.t`. Calling `stdtr` + directly can be much faster than calling the ``cdf`` method of + `scipy.stats.t`. To get the same results, one must use the following + parametrization: ``scipy.stats.t(df).cdf(x) = stdtr(df, x)``. + + >>> from scipy.stats import t + >>> df, x = 3, 1 + >>> stdtr_result = stdtr(df, x) # this can be faster than below + >>> stats_result = t(df).cdf(x) + >>> stats_result == stdtr_result # test that results are equal + True + """) + +add_newdoc("stdtridf", + """ + stdtridf(p, t, out=None) + + Inverse of `stdtr` vs df. + + Returns the argument df such that stdtr(df, t) is equal to `p`. + + Parameters + ---------- + p : array_like + Probability + t : array_like + Upper bound of the integral + out : ndarray, optional + Optional output array for the function results + + Returns + ------- + df : scalar or ndarray + Value of `df` such that ``stdtr(df, t) == p`` + + See Also + -------- + stdtr : Student t CDF + stdtrit : inverse of stdtr with respect to `t` + scipy.stats.t : Student t distribution + + Examples + -------- + Compute the student t cumulative distribution function for one + parameter set. + + >>> from scipy.special import stdtr, stdtridf + >>> df, x = 5, 2 + >>> cdf_value = stdtr(df, x) + >>> cdf_value + 0.9490302605850709 + + Verify that `stdtridf` recovers the original value for `df` given + the CDF value and `x`. + + >>> stdtridf(cdf_value, x) + 5.0 + """) + +add_newdoc("stdtrit", + """ + stdtrit(df, p, out=None) + + The `p`-th quantile of the student t distribution. + + This function is the inverse of the student t distribution cumulative + distribution function (CDF), returning `t` such that `stdtr(df, t) = p`. + + Returns the argument `t` such that stdtr(df, t) is equal to `p`. + + Parameters + ---------- + df : array_like + Degrees of freedom + p : array_like + Probability + out : ndarray, optional + Optional output array for the function results + + Returns + ------- + t : scalar or ndarray + Value of `t` such that ``stdtr(df, t) == p`` + + See Also + -------- + stdtr : Student t CDF + stdtridf : inverse of stdtr with respect to `df` + scipy.stats.t : Student t distribution + + Notes + ----- + The student t distribution is also available as `scipy.stats.t`. Calling + `stdtrit` directly can improve performance compared to the ``ppf`` + method of `scipy.stats.t` (see last example below). + + The function is computed using the Boost Math library [1]_, which + relies on the incomplete beta function. + + References + ---------- + .. [1] Boost C++ Libraries, https://www.boost.org/ + + Examples + -------- + `stdtrit` represents the inverse of the student t distribution CDF which + is available as `stdtr`. Here, we calculate the CDF for ``df`` at + ``x=1``. `stdtrit` then returns ``1`` up to floating point errors + given the same value for `df` and the computed CDF value. + + >>> import numpy as np + >>> from scipy.special import stdtr, stdtrit + >>> import matplotlib.pyplot as plt + >>> df = 3 + >>> x = 1 + >>> cdf_value = stdtr(df, x) + >>> stdtrit(df, cdf_value) + 0.9999999994418539 + + Plot the function for three different degrees of freedom. + + >>> x = np.linspace(0, 1, 1000) + >>> parameters = [(1, "solid"), (2, "dashed"), (5, "dotted")] + >>> fig, ax = plt.subplots() + >>> for (df, linestyle) in parameters: + ... ax.plot(x, stdtrit(df, x), ls=linestyle, label=f"$df={df}$") + >>> ax.legend() + >>> ax.set_ylim(-10, 10) + >>> ax.set_title("Student t distribution quantile function") + >>> plt.show() + + The function can be computed for several degrees of freedom at the same + time by providing a NumPy array or list for `df`: + + >>> stdtrit([1, 2, 3], 0.7) + array([0.72654253, 0.6172134 , 0.58438973]) + + It is possible to calculate the function at several points for several + different degrees of freedom simultaneously by providing arrays for `df` + and `p` with shapes compatible for broadcasting. Compute `stdtrit` at + 4 points for 3 degrees of freedom resulting in an array of shape 3x4. + + >>> dfs = np.array([[1], [2], [3]]) + >>> p = np.array([0.2, 0.4, 0.7, 0.8]) + >>> dfs.shape, p.shape + ((3, 1), (4,)) + + >>> stdtrit(dfs, p) + array([[-1.37638192, -0.3249197 , 0.72654253, 1.37638192], + [-1.06066017, -0.28867513, 0.6172134 , 1.06066017], + [-0.97847231, -0.27667066, 0.58438973, 0.97847231]]) + + The t distribution is also available as `scipy.stats.t`. Calling `stdtrit` + directly can be much faster than calling the ``ppf`` method of + `scipy.stats.t`. To get the same results, one must use the following + parametrization: ``scipy.stats.t(df).ppf(x) = stdtrit(df, x)``. + + >>> from scipy.stats import t + >>> df, x = 3, 0.5 + >>> stdtrit_result = stdtrit(df, x) # this can be faster than below + >>> stats_result = t(df).ppf(x) + >>> stats_result == stdtrit_result # test that results are equal + True + """) + +add_newdoc( + "tklmbda", + r""" + tklmbda(x, lmbda, out=None) + + Cumulative distribution function of the Tukey lambda distribution. + + Parameters + ---------- + x, lmbda : array_like + Parameters + out : ndarray, optional + Optional output array for the function results + + Returns + ------- + cdf : scalar or ndarray + Value of the Tukey lambda CDF + + See Also + -------- + scipy.stats.tukeylambda : Tukey lambda distribution + + Examples + -------- + >>> import numpy as np + >>> import matplotlib.pyplot as plt + >>> from scipy.special import tklmbda, expit + + Compute the cumulative distribution function (CDF) of the Tukey lambda + distribution at several ``x`` values for `lmbda` = -1.5. + + >>> x = np.linspace(-2, 2, 9) + >>> x + array([-2. , -1.5, -1. , -0.5, 0. , 0.5, 1. , 1.5, 2. ]) + >>> tklmbda(x, -1.5) + array([0.34688734, 0.3786554 , 0.41528805, 0.45629737, 0.5 , + 0.54370263, 0.58471195, 0.6213446 , 0.65311266]) + + When `lmbda` is 0, the function is the logistic sigmoid function, + which is implemented in `scipy.special` as `expit`. + + >>> tklmbda(x, 0) + array([0.11920292, 0.18242552, 0.26894142, 0.37754067, 0.5 , + 0.62245933, 0.73105858, 0.81757448, 0.88079708]) + >>> expit(x) + array([0.11920292, 0.18242552, 0.26894142, 0.37754067, 0.5 , + 0.62245933, 0.73105858, 0.81757448, 0.88079708]) + + When `lmbda` is 1, the Tukey lambda distribution is uniform on the + interval [-1, 1], so the CDF increases linearly. + + >>> t = np.linspace(-1, 1, 9) + >>> tklmbda(t, 1) + array([0. , 0.125, 0.25 , 0.375, 0.5 , 0.625, 0.75 , 0.875, 1. ]) + + In the following, we generate plots for several values of `lmbda`. + + The first figure shows graphs for `lmbda` <= 0. + + >>> styles = ['-', '-.', '--', ':'] + >>> fig, ax = plt.subplots() + >>> x = np.linspace(-12, 12, 500) + >>> for k, lmbda in enumerate([-1.0, -0.5, 0.0]): + ... y = tklmbda(x, lmbda) + ... ax.plot(x, y, styles[k], label=rf'$\lambda$ = {lmbda:-4.1f}') + + >>> ax.set_title(r'tklmbda(x, $\lambda$)') + >>> ax.set_label('x') + >>> ax.legend(framealpha=1, shadow=True) + >>> ax.grid(True) + + The second figure shows graphs for `lmbda` > 0. The dots in the + graphs show the bounds of the support of the distribution. + + >>> fig, ax = plt.subplots() + >>> x = np.linspace(-4.2, 4.2, 500) + >>> lmbdas = [0.25, 0.5, 1.0, 1.5] + >>> for k, lmbda in enumerate(lmbdas): + ... y = tklmbda(x, lmbda) + ... ax.plot(x, y, styles[k], label=fr'$\lambda$ = {lmbda}') + + >>> ax.set_prop_cycle(None) + >>> for lmbda in lmbdas: + ... ax.plot([-1/lmbda, 1/lmbda], [0, 1], '.', ms=8) + + >>> ax.set_title(r'tklmbda(x, $\lambda$)') + >>> ax.set_xlabel('x') + >>> ax.legend(framealpha=1, shadow=True) + >>> ax.grid(True) + + >>> plt.tight_layout() + >>> plt.show() + + The CDF of the Tukey lambda distribution is also implemented as the + ``cdf`` method of `scipy.stats.tukeylambda`. In the following, + ``tukeylambda.cdf(x, -0.5)`` and ``tklmbda(x, -0.5)`` compute the + same values: + + >>> from scipy.stats import tukeylambda + >>> x = np.linspace(-2, 2, 9) + + >>> tukeylambda.cdf(x, -0.5) + array([0.21995157, 0.27093858, 0.33541677, 0.41328161, 0.5 , + 0.58671839, 0.66458323, 0.72906142, 0.78004843]) + + >>> tklmbda(x, -0.5) + array([0.21995157, 0.27093858, 0.33541677, 0.41328161, 0.5 , + 0.58671839, 0.66458323, 0.72906142, 0.78004843]) + + The implementation in ``tukeylambda`` also provides location and scale + parameters, and other methods such as ``pdf()`` (the probability + density function) and ``ppf()`` (the inverse of the CDF), so for + working with the Tukey lambda distribution, ``tukeylambda`` is more + generally useful. The primary advantage of ``tklmbda`` is that it is + significantly faster than ``tukeylambda.cdf``. + """) + +add_newdoc("yn", + r""" + yn(n, x, out=None) + + Bessel function of the second kind of integer order and real argument. + + Parameters + ---------- + n : array_like + Order (integer). + x : array_like + Argument (float). + out : ndarray, optional + Optional output array for the function results + + Returns + ------- + Y : scalar or ndarray + Value of the Bessel function, :math:`Y_n(x)`. + + See Also + -------- + yv : For real order and real or complex argument. + y0: faster implementation of this function for order 0 + y1: faster implementation of this function for order 1 + + Notes + ----- + Wrapper for the Cephes [1]_ routine `yn`. + + The function is evaluated by forward recurrence on `n`, starting with + values computed by the Cephes routines `y0` and `y1`. If ``n = 0`` or 1, + the routine for `y0` or `y1` is called directly. + + References + ---------- + .. [1] Cephes Mathematical Functions Library, + https://netlib.org/cephes/ + + Examples + -------- + Evaluate the function of order 0 at one point. + + >>> from scipy.special import yn + >>> yn(0, 1.) + 0.08825696421567697 + + Evaluate the function at one point for different orders. + + >>> yn(0, 1.), yn(1, 1.), yn(2, 1.) + (0.08825696421567697, -0.7812128213002888, -1.6506826068162546) + + The evaluation for different orders can be carried out in one call by + providing a list or NumPy array as argument for the `v` parameter: + + >>> yn([0, 1, 2], 1.) + array([ 0.08825696, -0.78121282, -1.65068261]) + + Evaluate the function at several points for order 0 by providing an + array for `z`. + + >>> import numpy as np + >>> points = np.array([0.5, 3., 8.]) + >>> yn(0, points) + array([-0.44451873, 0.37685001, 0.22352149]) + + If `z` is an array, the order parameter `v` must be broadcastable to + the correct shape if different orders shall be computed in one call. + To calculate the orders 0 and 1 for a 1D array: + + >>> orders = np.array([[0], [1]]) + >>> orders.shape + (2, 1) + + >>> yn(orders, points) + array([[-0.44451873, 0.37685001, 0.22352149], + [-1.47147239, 0.32467442, -0.15806046]]) + + Plot the functions of order 0 to 3 from 0 to 10. + + >>> import matplotlib.pyplot as plt + >>> fig, ax = plt.subplots() + >>> x = np.linspace(0., 10., 1000) + >>> for i in range(4): + ... ax.plot(x, yn(i, x), label=f'$Y_{i!r}$') + >>> ax.set_ylim(-3, 1) + >>> ax.legend() + >>> plt.show() + """) + + +add_newdoc("_struve_asymp_large_z", + """ + _struve_asymp_large_z(v, z, is_h) + + Internal function for testing `struve` & `modstruve` + + Evaluates using asymptotic expansion + + Returns + ------- + v, err + """) + +add_newdoc("_struve_power_series", + """ + _struve_power_series(v, z, is_h) + + Internal function for testing `struve` & `modstruve` + + Evaluates using power series + + Returns + ------- + v, err + """) + +add_newdoc("_struve_bessel_series", + """ + _struve_bessel_series(v, z, is_h) + + Internal function for testing `struve` & `modstruve` + + Evaluates using Bessel function series + + Returns + ------- + v, err + """) + +add_newdoc("_spherical_jn", + """ + Internal function, use `spherical_jn` instead. + """) + +add_newdoc("_spherical_jn_d", + """ + Internal function, use `spherical_jn` instead. + """) + +add_newdoc("_spherical_yn", + """ + Internal function, use `spherical_yn` instead. + """) + +add_newdoc("_spherical_yn_d", + """ + Internal function, use `spherical_yn` instead. + """) + +add_newdoc("_spherical_in", + """ + Internal function, use `spherical_in` instead. + """) + +add_newdoc("_spherical_in_d", + """ + Internal function, use `spherical_in` instead. + """) + +add_newdoc("_spherical_kn", + """ + Internal function, use `spherical_kn` instead. + """) + +add_newdoc("_spherical_kn_d", + """ + Internal function, use `spherical_kn` instead. + """) + +add_newdoc("owens_t", + """ + owens_t(h, a, out=None) + + Owen's T Function. + + The function T(h, a) gives the probability of the event + (X > h and 0 < Y < a * X) where X and Y are independent + standard normal random variables. + + Parameters + ---------- + h : array_like + Input value. + a : array_like + Input value. + out : ndarray, optional + Optional output array for the function results + + Returns + ------- + t: scalar or ndarray + Probability of the event (X > h and 0 < Y < a * X), + where X and Y are independent standard normal random variables. + + References + ---------- + .. [1] M. Patefield and D. Tandy, "Fast and accurate calculation of + Owen's T Function", Statistical Software vol. 5, pp. 1-25, 2000. + + Examples + -------- + >>> from scipy import special + >>> a = 3.5 + >>> h = 0.78 + >>> special.owens_t(h, a) + 0.10877216734852274 + """) + +add_newdoc("_factorial", + """ + Internal function, do not use. + """) + +add_newdoc("ndtri_exp", + r""" + ndtri_exp(y, out=None) + + Inverse of `log_ndtr` vs x. Allows for greater precision than + `ndtri` composed with `numpy.exp` for very small values of y and for + y close to 0. + + Parameters + ---------- + y : array_like of float + Function argument + out : ndarray, optional + Optional output array for the function results + + Returns + ------- + scalar or ndarray + Inverse of the log CDF of the standard normal distribution, evaluated + at y. + + See Also + -------- + log_ndtr : log of the standard normal cumulative distribution function + ndtr : standard normal cumulative distribution function + ndtri : standard normal percentile function + + Examples + -------- + >>> import numpy as np + >>> import scipy.special as sc + + `ndtri_exp` agrees with the naive implementation when the latter does + not suffer from underflow. + + >>> sc.ndtri_exp(-1) + -0.33747496376420244 + >>> sc.ndtri(np.exp(-1)) + -0.33747496376420244 + + For extreme values of y, the naive approach fails + + >>> sc.ndtri(np.exp(-800)) + -inf + >>> sc.ndtri(np.exp(-1e-20)) + inf + + whereas `ndtri_exp` is still able to compute the result to high precision. + + >>> sc.ndtri_exp(-800) + -39.88469483825668 + >>> sc.ndtri_exp(-1e-20) + 9.262340089798409 + """) + + +add_newdoc("_stirling2_inexact", + r""" + Internal function, do not use. + """) + +add_newdoc( + "_beta_pdf", + r""" + _beta_pdf(x, a, b) + + Probability density function of beta distribution. + + Parameters + ---------- + x : array_like + Real-valued such that :math:`0 \leq x \leq 1`, + the upper limit of integration + a, b : array_like + Positive, real-valued parameters + + Returns + ------- + scalar or ndarray + + """) + +add_newdoc( + "_beta_ppf", + r""" + _beta_ppf(x, a, b) + + Percent point function of beta distribution. + + Parameters + ---------- + x : array_like + Real-valued such that :math:`0 \leq x \leq 1`, + the upper limit of integration + a, b : array_like + Positive, real-valued parameters + + Returns + ------- + scalar or ndarray + + """) + +add_newdoc( + "_invgauss_ppf", + """ + _invgauss_ppf(x, mu) + + Percent point function of inverse gaussian distribution. + + Parameters + ---------- + x : array_like + Positive real-valued + mu : array_like + Positive, real-valued parameters + + Returns + ------- + scalar or ndarray + + """) + +add_newdoc( + "_invgauss_isf", + """ + _invgauss_isf(x, mu, s) + + Inverse survival function of inverse gaussian distribution. + + Parameters + ---------- + x : array_like + Positive real-valued + mu : array_like + Positive, real-valued parameters + s : array_like + Positive, real-valued parameters + + Returns + ------- + scalar or ndarray + + """) + +add_newdoc( + "_cauchy_ppf", + """ + _cauchy_ppf(p, loc, scale) + + Percent point function (i.e. quantile) of the Cauchy distribution. + + Parameters + ---------- + p : array_like + Probabilities + loc : array_like + Location parameter of the distribution. + scale : array_like + Scale parameter of the distribution. + + Returns + ------- + scalar or ndarray + + """) + +add_newdoc( + "_cauchy_isf", + """ + _cauchy_isf(p, loc, scale) + + Inverse survival function of the Cauchy distribution. + + Parameters + ---------- + p : array_like + Probabilities + loc : array_like + Location parameter of the distribution. + scale : array_like + Scale parameter of the distribution. + + Returns + ------- + scalar or ndarray + + """) + +add_newdoc( + "_ncx2_pdf", + """ + _ncx2_pdf(x, k, l) + + Probability density function of Non-central chi-squared distribution. + + Parameters + ---------- + x : array_like + Positive real-valued + k, l : array_like + Positive, real-valued parameters + + Returns + ------- + scalar or ndarray + + """) + +add_newdoc( + "_ncx2_cdf", + """ + _ncx2_cdf(x, k, l) + + Cumulative density function of Non-central chi-squared distribution. + + Parameters + ---------- + x : array_like + Positive real-valued + k, l : array_like + Positive, real-valued parameters + + Returns + ------- + scalar or ndarray + + """) + +add_newdoc( + "_ncx2_ppf", + """ + _ncx2_ppf(x, k, l) + + Percent point function of Non-central chi-squared distribution. + + Parameters + ---------- + x : array_like + Positive real-valued + k, l : array_like + Positive, real-valued parameters + + Returns + ------- + scalar or ndarray + + """) + +add_newdoc( + "_ncx2_sf", + """ + _ncx2_sf(x, k, l) + + Survival function of Non-central chi-squared distribution. + + Parameters + ---------- + x : array_like + Positive real-valued + k, l : array_like + Positive, real-valued parameters + + Returns + ------- + scalar or ndarray + + """) + +add_newdoc( + "_ncx2_isf", + """ + _ncx2_isf(x, k, l) + + Inverse survival function of Non-central chi-squared distribution. + + Parameters + ---------- + x : array_like + Positive real-valued + k, l : array_like + Positive, real-valued parameters + + Returns + ------- + scalar or ndarray + + """) + +add_newdoc( + "_ncf_pdf", + """ + _ncf_pdf(x, v1, v2, l) + + Probability density function of noncentral F-distribution. + + Parameters + ---------- + x : array_like + Positive real-valued + v1, v2, l : array_like + Positive, real-valued parameters + + Returns + ------- + scalar or ndarray + + """) + +add_newdoc( + "_ncf_cdf", + """ + _ncf_cdf(x, v1, v2, l) + + Cumulative density function of noncentral F-distribution. + + Parameters + ---------- + x : array_like + Positive real-valued + v1, v2, l : array_like + Positive, real-valued parameters + + Returns + ------- + scalar or ndarray + + """) + +add_newdoc( + "_ncf_ppf", + """ + _ncf_ppf(x, v1, v2, l) + + Percent point function of noncentral F-distribution. + + Parameters + ---------- + x : array_like + Positive real-valued + v1, v2, l : array_like + Positive, real-valued parameters + + Returns + ------- + scalar or ndarray + + """) + +add_newdoc( + "_ncf_sf", + """ + _ncf_sf(x, v1, v2, l) + + Survival function of noncentral F-distribution. + + Parameters + ---------- + x : array_like + Positive real-valued + v1, v2, l : array_like + Positive, real-valued parameters + + Returns + ------- + scalar or ndarray + + """) + +add_newdoc( + "_ncf_isf", + """ + _ncf_isf(x, v1, v2, l) + + Inverse survival function of noncentral F-distribution. + + Parameters + ---------- + x : array_like + Positive real-valued + v1, v2, l : array_like + Positive, real-valued parameters + + Returns + ------- + scalar or ndarray + + """) + +add_newdoc( + "_ncf_mean", + """ + _ncf_mean(v1, v2, l) + + Mean of noncentral F-distribution. + + Parameters + ---------- + v1, v2, l : array_like + Positive, real-valued parameters + + Returns + ------- + scalar or ndarray + + """) + +add_newdoc( + "_ncf_variance", + """ + _ncf_variance(v1, v2, l) + + Variance of noncentral F-distribution. + + Parameters + ---------- + v1, v2, l : array_like + Positive, real-valued parameters + + Returns + ------- + scalar or ndarray + + """) + +add_newdoc( + "_ncf_skewness", + """ + _ncf_skewness(v1, v2, l) + + Skewness of noncentral F-distribution. + + Parameters + ---------- + v1, v2, l : array_like + Positive, real-valued parameters + + Returns + ------- + scalar or ndarray + + """) + +add_newdoc( + "_ncf_kurtosis_excess", + """ + _ncf_kurtosis_excess(v1, v2, l) + + Kurtosis excess of noncentral F-distribution. + + Parameters + ---------- + v1, v2, l : array_like + Positive, real-valued parameters + + Returns + ------- + scalar or ndarray + + """) + +add_newdoc( + "_nct_cdf", + """ + _nct_cdf(x, v, l) + + Cumulative density function of noncentral t-distribution. + + Parameters + ---------- + x : array_like + Real-valued + v : array_like + Positive, real-valued parameters + l : array_like + Real-valued parameters + + Returns + ------- + scalar or ndarray + + """) + +add_newdoc( + "_nct_pdf", + """ + _nct_pdf(x, v, l) + + Probability density function of noncentral t-distribution. + + Parameters + ---------- + x : array_like + Real-valued + v : array_like + Positive, real-valued parameters + l : array_like + Real-valued parameters + + Returns + ------- + scalar or ndarray + + """) + + +add_newdoc( + "_nct_ppf", + """ + _nct_ppf(x, v, l) + + Percent point function of noncentral t-distribution. + + Parameters + ---------- + x : array_like + Real-valued + v : array_like + Positive, real-valued parameters + l : array_like + Real-valued parameters + + Returns + ------- + scalar or ndarray + + """) + +add_newdoc( + "_nct_sf", + """ + _nct_sf(x, v, l) + + Survival function of noncentral t-distribution. + + Parameters + ---------- + x : array_like + Real-valued + v : array_like + Positive, real-valued parameters + l : array_like + Real-valued parameters + + Returns + ------- + scalar or ndarray + + """) + +add_newdoc( + "_nct_isf", + """ + _nct_isf(x, v, l) + + Inverse survival function of noncentral t-distribution. + + Parameters + ---------- + x : array_like + Real-valued + v : array_like + Positive, real-valued parameters + l : array_like + Real-valued parameters + + Returns + ------- + scalar or ndarray + + """) + +add_newdoc( + "_nct_mean", + """ + _nct_mean(v, l) + + Mean of noncentral t-distribution. + + Parameters + ---------- + v : array_like + Positive, real-valued parameters + l : array_like + Real-valued parameters + + Returns + ------- + scalar or ndarray + + """) + +add_newdoc( + "_nct_variance", + """ + _nct_variance(v, l) + + Variance of noncentral t-distribution. + + Parameters + ---------- + v : array_like + Positive, real-valued parameters + l : array_like + Real-valued parameters + + Returns + ------- + scalar or ndarray + + """) + +add_newdoc( + "_nct_skewness", + """ + _nct_skewness(v, l) + + Skewness of noncentral t-distribution. + + Parameters + ---------- + v : array_like + Positive, real-valued parameters + l : array_like + Real-valued parameters + + Returns + ------- + scalar or ndarray + + """) + +add_newdoc( + "_nct_kurtosis_excess", + """ + _nct_kurtosis_excess(v, l) + + Kurtosis excess of noncentral t-distribution. + + Parameters + ---------- + v : array_like + Positive, real-valued parameters + l : array_like + Real-valued parameters + + Returns + ------- + scalar or ndarray + + """) + +add_newdoc( + "_skewnorm_cdf", + """ + _skewnorm_cdf(x, l, sc, sh) + + Cumulative density function of skewnorm distribution. + + Parameters + ---------- + x : array_like + Real-valued + l : array_like + Real-valued parameters + sc : array_like + Positive, Real-valued parameters + sh : array_like + Real-valued parameters + + Returns + ------- + scalar or ndarray + + """) + +add_newdoc( + "_skewnorm_ppf", + """ + _skewnorm_ppf(x, l, sc, sh) + + Percent point function of skewnorm distribution. + + Parameters + ---------- + x : array_like + Real-valued + l : array_like + Real-valued parameters + sc : array_like + Positive, Real-valued parameters + sh : array_like + Real-valued parameters + + Returns + ------- + scalar or ndarray + + """) + +add_newdoc( + "_skewnorm_isf", + """ + _skewnorm_isf(x, l, sc, sh) + + Inverse survival function of skewnorm distribution. + + Parameters + ---------- + x : array_like + Real-valued + l : array_like + Real-valued parameters + sc : array_like + Positive, Real-valued parameters + sh : array_like + Real-valued parameters + + Returns + ------- + scalar or ndarray + + """) + +add_newdoc( + "_binom_pmf", + """ + _binom_pmf(x, n, p) + + Probability mass function of binomial distribution. + + Parameters + ---------- + x : array_like + Real-valued + n : array_like + Positive, integer-valued parameter + p : array_like + Positive, real-valued parameter + + Returns + ------- + scalar or ndarray + + """) + +add_newdoc( + "_binom_cdf", + """ + _binom_cdf(x, n, p) + + Cumulative density function of binomial distribution. + + Parameters + ---------- + x : array_like + Real-valued + n : array_like + Positive, integer-valued parameter + p : array_like + Positive, real-valued parameter + + Returns + ------- + scalar or ndarray + + """) + +add_newdoc( + "_binom_ppf", + """ + _binom_ppf(x, n, p) + + Percent point function of binomial distribution. + + Parameters + ---------- + x : array_like + Real-valued + n : array_like + Positive, integer-valued parameter + p : array_like + Positive, real-valued parameter + + Returns + ------- + scalar or ndarray + + """) + +add_newdoc( + "_binom_sf", + """ + _binom_sf(x, n, p) + + Survival function of binomial distribution. + + Parameters + ---------- + x : array_like + Real-valued + n : array_like + Positive, integer-valued parameter + p : array_like + Positive, real-valued parameter + + Returns + ------- + scalar or ndarray + + """) + +add_newdoc( + "_binom_isf", + """ + _binom_isf(x, n, p) + + Inverse survival function of binomial distribution. + + Parameters + ---------- + x : array_like + Real-valued + n : array_like + Positive, integer-valued parameter + p : array_like + Positive, real-valued parameter + + Returns + ------- + scalar or ndarray + + """) + +add_newdoc( + "_nbinom_pmf", + """ + _nbinom_pmf(x, r, p) + + Probability mass function of negative binomial distribution. + + Parameters + ---------- + x : array_like + Real-valued + r : array_like + Positive, integer-valued parameter + p : array_like + Positive, real-valued parameter + + Returns + ------- + scalar or ndarray + + """) + +add_newdoc( + "_nbinom_cdf", + """ + _nbinom_cdf(x, r, p) + + Cumulative density function of negative binomial distribution. + + Parameters + ---------- + x : array_like + Real-valued + r : array_like + Positive, integer-valued parameter + p : array_like + Positive, real-valued parameter + + Returns + ------- + scalar or ndarray + + """) + +add_newdoc( + "_nbinom_ppf", + """ + _nbinom_ppf(x, r, p) + + Percent point function of negative binomial distribution. + + Parameters + ---------- + x : array_like + Real-valued + r : array_like + Positive, integer-valued parameter + p : array_like + Positive, real-valued parameter + + Returns + ------- + scalar or ndarray + + """) + +add_newdoc( + "_nbinom_sf", + """ + _nbinom_sf(x, r, p) + + Survival function of negative binomial distribution. + + Parameters + ---------- + x : array_like + Real-valued + r : array_like + Positive, integer-valued parameter + p : array_like + Positive, real-valued parameter + + Returns + ------- + scalar or ndarray + + """) + +add_newdoc( + "_nbinom_isf", + """ + _nbinom_isf(x, r, p) + + Inverse survival function of negative binomial distribution. + + Parameters + ---------- + x : array_like + Real-valued + r : array_like + Positive, integer-valued parameter + p : array_like + Positive, real-valued parameter + + Returns + ------- + scalar or ndarray + + """) + +add_newdoc( + "_nbinom_mean", + """ + _nbinom_mean(r, p) + + Mean of negative binomial distribution. + + Parameters + ---------- + r : array_like + Positive, integer-valued parameter + p : array_like + Positive, real-valued parameter + + Returns + ------- + scalar or ndarray + + """) + +add_newdoc( + "_nbinom_variance", + """ + _nbinom_variance(r, p) + + Variance of negative binomial distribution. + + Parameters + ---------- + r : array_like + Positive, integer-valued parameter + p : array_like + Positive, real-valued parameter + + Returns + ------- + scalar or ndarray + + """) + +add_newdoc( + "_nbinom_skewness", + """ + _nbinom_skewness(r, p) + + Skewness of negative binomial distribution. + + Parameters + ---------- + r : array_like + Positive, integer-valued parameter + p : array_like + Positive, real-valued parameter + + Returns + ------- + scalar or ndarray + + """) + +add_newdoc( + "_nbinom_kurtosis_excess", + """ + _nbinom_kurtosis_excess(r, p) + + Kurtosis excess of negative binomial distribution. + + Parameters + ---------- + r : array_like + Positive, integer-valued parameter + p : array_like + Positive, real-valued parameter + + Returns + ------- + scalar or ndarray + + """) + +add_newdoc( + "_hypergeom_pmf", + """ + _hypergeom_pmf(x, r, N, M) + + Probability mass function of hypergeometric distribution. + + Parameters + ---------- + x : array_like + Real-valued + r, N, M : array_like + Positive, integer-valued parameter + + Returns + ------- + scalar or ndarray + + """) + +add_newdoc( + "_hypergeom_cdf", + """ + _hypergeom_cdf(x, r, N, M) + + Cumulative density function of hypergeometric distribution. + + Parameters + ---------- + x : array_like + Real-valued + r, N, M : array_like + Positive, integer-valued parameter + + Returns + ------- + scalar or ndarray + """) + +add_newdoc( + "_hypergeom_sf", + """ + _hypergeom_sf(x, r, N, M) + + Survival function of hypergeometric distribution. + + Parameters + ---------- + x : array_like + Real-valued + r, N, M : array_like + Positive, integer-valued parameter + + Returns + ------- + scalar or ndarray + """) + +add_newdoc( + "_hypergeom_mean", + """ + _hypergeom_mean(r, N, M) + + Mean of hypergeometric distribution. + + Parameters + ---------- + r, N, M : array_like + Positive, integer-valued parameter + + Returns + ------- + scalar or ndarray + + """) + +add_newdoc( + "_hypergeom_variance", + """ + _hypergeom_variance(r, N, M) + + Variance of hypergeometric distribution. + + Parameters + ---------- + r, N, M : array_like + Positive, integer-valued parameter + + Returns + ------- + scalar or ndarray + + """) + +add_newdoc( + "_hypergeom_skewness", + """ + _hypergeom_skewness(r, N, M) + + Skewness of hypergeometric distribution. + + Parameters + ---------- + r, N, M : array_like + Positive, integer-valued parameter + + Returns + ------- + scalar or ndarray + + """) diff --git a/v12x/pylib/scipy/special/_basic.py b/v12x/pylib/scipy/special/_basic.py new file mode 100644 index 0000000000000000000000000000000000000000..8e017deb4db1eec2562a69f506fcdfed18d3de30 --- /dev/null +++ b/v12x/pylib/scipy/special/_basic.py @@ -0,0 +1,3504 @@ +# +# Author: Travis Oliphant, 2002 +# + +import numpy as np +import math +import warnings +from collections import defaultdict +from heapq import heapify, heappop +from numpy import (pi, asarray, floor, isscalar, sqrt, where, + sin, place, issubdtype, extract, inexact, nan, zeros, sinc) + +from . import _ufuncs +from ._ufuncs import (mathieu_a, mathieu_b, iv, jv, gamma, rgamma, + psi, hankel1, hankel2, yv, kv, poch, binom, + _stirling2_inexact) + +from ._gufuncs import _lqn, _lqmn, _rctj, _rcty +from ._input_validation import _nonneg_int_or_fail +from . import _specfun +from ._comb import _comb_int + + +__all__ = [ + 'ai_zeros', + 'assoc_laguerre', + 'bei_zeros', + 'beip_zeros', + 'ber_zeros', + 'bernoulli', + 'berp_zeros', + 'bi_zeros', + 'comb', + 'digamma', + 'diric', + 'erf_zeros', + 'euler', + 'factorial', + 'factorial2', + 'factorialk', + 'fresnel_zeros', + 'fresnelc_zeros', + 'fresnels_zeros', + 'h1vp', + 'h2vp', + 'ivp', + 'jn_zeros', + 'jnjnp_zeros', + 'jnp_zeros', + 'jnyn_zeros', + 'jvp', + 'kei_zeros', + 'keip_zeros', + 'kelvin_zeros', + 'ker_zeros', + 'kerp_zeros', + 'kvp', + 'lmbda', + 'lqmn', + 'lqn', + 'mathieu_even_coef', + 'mathieu_odd_coef', + 'obl_cv_seq', + 'pbdn_seq', + 'pbdv_seq', + 'pbvv_seq', + 'perm', + 'polygamma', + 'pro_cv_seq', + 'riccati_jn', + 'riccati_yn', + 'sinc', + 'softplus', + 'stirling2', + 'y0_zeros', + 'y1_zeros', + 'y1p_zeros', + 'yn_zeros', + 'ynp_zeros', + 'yvp', + 'zeta' +] + + +# mapping k to last n such that factorialk(n, k) < np.iinfo(np.int64).max +_FACTORIALK_LIMITS_64BITS = {1: 20, 2: 33, 3: 44, 4: 54, 5: 65, + 6: 74, 7: 84, 8: 93, 9: 101} +# mapping k to last n such that factorialk(n, k) < np.iinfo(np.int32).max +_FACTORIALK_LIMITS_32BITS = {1: 12, 2: 19, 3: 25, 4: 31, 5: 37, + 6: 43, 7: 47, 8: 51, 9: 56} + + +def diric(x, n): + """Periodic sinc function, also called the Dirichlet kernel. + + The Dirichlet kernel is defined as:: + + diric(x, n) = sin(x * n/2) / (n * sin(x / 2)), + + where `n` is a positive integer. + + Parameters + ---------- + x : array_like + Input data + n : int + Integer defining the periodicity. + + Returns + ------- + diric : ndarray + Value of periodic sinc function. + + Examples + -------- + >>> import numpy as np + >>> from scipy import special + >>> import matplotlib.pyplot as plt + + >>> x = np.linspace(-8*np.pi, 8*np.pi, num=201) + >>> plt.figure(figsize=(8, 8)); + >>> for idx, n in enumerate([2, 3, 4, 9]): + ... plt.subplot(2, 2, idx+1) + ... plt.plot(x, special.diric(x, n)) + ... plt.title('diric, n={}'.format(n)) + >>> plt.show() + + The following example demonstrates that `diric` gives the magnitudes + (modulo the sign and scaling) of the Fourier coefficients of a + rectangular pulse. + + Suppress output of values that are effectively 0: + + >>> np.set_printoptions(suppress=True) + + Create a signal `x` of length `m` with `k` ones: + + >>> m = 8 + >>> k = 3 + >>> x = np.zeros(m) + >>> x[:k] = 1 + + Use the FFT to compute the Fourier transform of `x`, and + inspect the magnitudes of the coefficients: + + >>> np.abs(np.fft.fft(x)) + array([ 3. , 2.41421356, 1. , 0.41421356, 1. , + 0.41421356, 1. , 2.41421356]) + + Now find the same values (up to sign) using `diric`. We multiply + by `k` to account for the different scaling conventions of + `numpy.fft.fft` and `diric`: + + >>> theta = np.linspace(0, 2*np.pi, m, endpoint=False) + >>> k * special.diric(theta, k) + array([ 3. , 2.41421356, 1. , -0.41421356, -1. , + -0.41421356, 1. , 2.41421356]) + """ + x, n = asarray(x), asarray(n) + n = asarray(n + (x-x)) + x = asarray(x + (n-n)) + if issubdtype(x.dtype, inexact): + ytype = x.dtype + else: + ytype = float + y = zeros(x.shape, ytype) + + # empirical minval for 32, 64 or 128 bit float computations + # where sin(x/2) < minval, result is fixed at +1 or -1 + if np.finfo(ytype).eps < 1e-18: + minval = 1e-11 + elif np.finfo(ytype).eps < 1e-15: + minval = 1e-7 + else: + minval = 1e-3 + + mask1 = (n <= 0) | (n != floor(n)) + place(y, mask1, nan) + + x = x / 2 + denom = sin(x) + mask2 = (1-mask1) & (abs(denom) < minval) + xsub = extract(mask2, x) + nsub = extract(mask2, n) + zsub = xsub / pi + place(y, mask2, pow(-1, np.round(zsub)*(nsub-1))) + + mask = (1-mask1) & (1-mask2) + xsub = extract(mask, x) + nsub = extract(mask, n) + dsub = extract(mask, denom) + place(y, mask, sin(nsub*xsub)/(nsub*dsub)) + return y + + +def jnjnp_zeros(nt): + """Compute zeros of integer-order Bessel functions Jn and Jn'. + + Results are arranged in order of the magnitudes of the zeros. + + Parameters + ---------- + nt : int + Number (<=1200) of zeros to compute + + Returns + ------- + zo[l-1] : ndarray + Value of the lth zero of Jn(x) and Jn'(x). Of length `nt`. + n[l-1] : ndarray + Order of the Jn(x) or Jn'(x) associated with lth zero. Of length `nt`. + m[l-1] : ndarray + Serial number of the zeros of Jn(x) or Jn'(x) associated + with lth zero. Of length `nt`. + t[l-1] : ndarray + 0 if lth zero in zo is zero of Jn(x), 1 if it is a zero of Jn'(x). Of + length `nt`. + + See Also + -------- + jn_zeros, jnp_zeros : to get separated arrays of zeros. + + References + ---------- + .. [1] Zhang, Shanjie and Jin, Jianming. "Computation of Special + Functions", John Wiley and Sons, 1996, chapter 5. + https://people.sc.fsu.edu/~jburkardt/f77_src/special_functions/special_functions.html + + """ + if not isscalar(nt) or (floor(nt) != nt) or (nt > 1200): + raise ValueError("Number must be integer <= 1200.") + nt = int(nt) + n, m, t, zo = _specfun.jdzo(nt) + return zo[1:nt+1], n[:nt], m[:nt], t[:nt] + + +def jnyn_zeros(n, nt): + """Compute nt zeros of Bessel functions Jn(x), Jn'(x), Yn(x), and Yn'(x). + + Returns 4 arrays of length `nt`, corresponding to the first `nt` + zeros of Jn(x), Jn'(x), Yn(x), and Yn'(x), respectively. The zeros + are returned in ascending order. + + Parameters + ---------- + n : int + Order of the Bessel functions + nt : int + Number (<=1200) of zeros to compute + + Returns + ------- + Jn : ndarray + First `nt` zeros of Jn + Jnp : ndarray + First `nt` zeros of Jn' + Yn : ndarray + First `nt` zeros of Yn + Ynp : ndarray + First `nt` zeros of Yn' + + See Also + -------- + jn_zeros, jnp_zeros, yn_zeros, ynp_zeros + + References + ---------- + .. [1] Zhang, Shanjie and Jin, Jianming. "Computation of Special + Functions", John Wiley and Sons, 1996, chapter 5. + https://people.sc.fsu.edu/~jburkardt/f77_src/special_functions/special_functions.html + + Examples + -------- + Compute the first three roots of :math:`J_1`, :math:`J_1'`, + :math:`Y_1` and :math:`Y_1'`. + + >>> from scipy.special import jnyn_zeros + >>> jn_roots, jnp_roots, yn_roots, ynp_roots = jnyn_zeros(1, 3) + >>> jn_roots, yn_roots + (array([ 3.83170597, 7.01558667, 10.17346814]), + array([2.19714133, 5.42968104, 8.59600587])) + + Plot :math:`J_1`, :math:`J_1'`, :math:`Y_1`, :math:`Y_1'` and their roots. + + >>> import numpy as np + >>> import matplotlib.pyplot as plt + >>> from scipy.special import jnyn_zeros, jvp, jn, yvp, yn + >>> jn_roots, jnp_roots, yn_roots, ynp_roots = jnyn_zeros(1, 3) + >>> fig, ax = plt.subplots() + >>> xmax= 11 + >>> x = np.linspace(0, xmax) + >>> x[0] += 1e-15 + >>> ax.plot(x, jn(1, x), label=r"$J_1$", c='r') + >>> ax.plot(x, jvp(1, x, 1), label=r"$J_1'$", c='b') + >>> ax.plot(x, yn(1, x), label=r"$Y_1$", c='y') + >>> ax.plot(x, yvp(1, x, 1), label=r"$Y_1'$", c='c') + >>> zeros = np.zeros((3, )) + >>> ax.scatter(jn_roots, zeros, s=30, c='r', zorder=5, + ... label=r"$J_1$ roots") + >>> ax.scatter(jnp_roots, zeros, s=30, c='b', zorder=5, + ... label=r"$J_1'$ roots") + >>> ax.scatter(yn_roots, zeros, s=30, c='y', zorder=5, + ... label=r"$Y_1$ roots") + >>> ax.scatter(ynp_roots, zeros, s=30, c='c', zorder=5, + ... label=r"$Y_1'$ roots") + >>> ax.hlines(0, 0, xmax, color='k') + >>> ax.set_ylim(-0.6, 0.6) + >>> ax.set_xlim(0, xmax) + >>> ax.legend(ncol=2, bbox_to_anchor=(1., 0.75)) + >>> plt.tight_layout() + >>> plt.show() + """ + if not (isscalar(nt) and isscalar(n)): + raise ValueError("Arguments must be scalars.") + if (floor(n) != n) or (floor(nt) != nt): + raise ValueError("Arguments must be integers.") + if (nt <= 0): + raise ValueError("nt > 0") + return _specfun.jyzo(abs(n), nt) + + +def jn_zeros(n, nt): + r"""Compute zeros of integer-order Bessel functions Jn. + + Compute `nt` zeros of the Bessel functions :math:`J_n(x)` on the + interval :math:`(0, \infty)`. The zeros are returned in ascending + order. Note that this interval excludes the zero at :math:`x = 0` + that exists for :math:`n > 0`. + + Parameters + ---------- + n : int + Order of Bessel function + nt : int + Number of zeros to return + + Returns + ------- + ndarray + First `nt` zeros of the Bessel function. + + See Also + -------- + jv: Real-order Bessel functions of the first kind + jnp_zeros: Zeros of :math:`Jn'` + + References + ---------- + .. [1] Zhang, Shanjie and Jin, Jianming. "Computation of Special + Functions", John Wiley and Sons, 1996, chapter 5. + https://people.sc.fsu.edu/~jburkardt/f77_src/special_functions/special_functions.html + + Examples + -------- + Compute the first four positive roots of :math:`J_3`. + + >>> from scipy.special import jn_zeros + >>> jn_zeros(3, 4) + array([ 6.3801619 , 9.76102313, 13.01520072, 16.22346616]) + + Plot :math:`J_3` and its first four positive roots. Note + that the root located at 0 is not returned by `jn_zeros`. + + >>> import numpy as np + >>> import matplotlib.pyplot as plt + >>> from scipy.special import jn, jn_zeros + >>> j3_roots = jn_zeros(3, 4) + >>> xmax = 18 + >>> xmin = -1 + >>> x = np.linspace(xmin, xmax, 500) + >>> fig, ax = plt.subplots() + >>> ax.plot(x, jn(3, x), label=r'$J_3$') + >>> ax.scatter(j3_roots, np.zeros((4, )), s=30, c='r', + ... label=r"$J_3$_Zeros", zorder=5) + >>> ax.scatter(0, 0, s=30, c='k', + ... label=r"Root at 0", zorder=5) + >>> ax.hlines(0, 0, xmax, color='k') + >>> ax.set_xlim(xmin, xmax) + >>> plt.legend() + >>> plt.show() + """ + return jnyn_zeros(n, nt)[0] + + +def jnp_zeros(n, nt): + r"""Compute zeros of integer-order Bessel function derivatives Jn'. + + Compute `nt` zeros of the functions :math:`J_n'(x)` on the + interval :math:`(0, \infty)`. The zeros are returned in ascending + order. Note that this interval excludes the zero at :math:`x = 0` + that exists for :math:`n > 1`. + + Parameters + ---------- + n : int + Order of Bessel function + nt : int + Number of zeros to return + + Returns + ------- + ndarray + First `nt` zeros of the Bessel function. + + See Also + -------- + jvp: Derivatives of integer-order Bessel functions of the first kind + jv: Float-order Bessel functions of the first kind + + References + ---------- + .. [1] Zhang, Shanjie and Jin, Jianming. "Computation of Special + Functions", John Wiley and Sons, 1996, chapter 5. + https://people.sc.fsu.edu/~jburkardt/f77_src/special_functions/special_functions.html + + Examples + -------- + Compute the first four roots of :math:`J_2'`. + + >>> from scipy.special import jnp_zeros + >>> jnp_zeros(2, 4) + array([ 3.05423693, 6.70613319, 9.96946782, 13.17037086]) + + As `jnp_zeros` yields the roots of :math:`J_n'`, it can be used to + compute the locations of the peaks of :math:`J_n`. Plot + :math:`J_2`, :math:`J_2'` and the locations of the roots of :math:`J_2'`. + + >>> import numpy as np + >>> import matplotlib.pyplot as plt + >>> from scipy.special import jn, jnp_zeros, jvp + >>> j2_roots = jnp_zeros(2, 4) + >>> xmax = 15 + >>> x = np.linspace(0, xmax, 500) + >>> fig, ax = plt.subplots() + >>> ax.plot(x, jn(2, x), label=r'$J_2$') + >>> ax.plot(x, jvp(2, x, 1), label=r"$J_2'$") + >>> ax.hlines(0, 0, xmax, color='k') + >>> ax.scatter(j2_roots, np.zeros((4, )), s=30, c='r', + ... label=r"Roots of $J_2'$", zorder=5) + >>> ax.set_ylim(-0.4, 0.8) + >>> ax.set_xlim(0, xmax) + >>> plt.legend() + >>> plt.show() + """ + return jnyn_zeros(n, nt)[1] + + +def yn_zeros(n, nt): + r"""Compute zeros of integer-order Bessel function Yn(x). + + Compute `nt` zeros of the functions :math:`Y_n(x)` on the interval + :math:`(0, \infty)`. The zeros are returned in ascending order. + + Parameters + ---------- + n : int + Order of Bessel function + nt : int + Number of zeros to return + + Returns + ------- + ndarray + First `nt` zeros of the Bessel function. + + See Also + -------- + yn: Bessel function of the second kind for integer order + yv: Bessel function of the second kind for real order + + References + ---------- + .. [1] Zhang, Shanjie and Jin, Jianming. "Computation of Special + Functions", John Wiley and Sons, 1996, chapter 5. + https://people.sc.fsu.edu/~jburkardt/f77_src/special_functions/special_functions.html + + Examples + -------- + Compute the first four roots of :math:`Y_2`. + + >>> from scipy.special import yn_zeros + >>> yn_zeros(2, 4) + array([ 3.38424177, 6.79380751, 10.02347798, 13.20998671]) + + Plot :math:`Y_2` and its first four roots. + + >>> import numpy as np + >>> import matplotlib.pyplot as plt + >>> from scipy.special import yn, yn_zeros + >>> xmin = 2 + >>> xmax = 15 + >>> x = np.linspace(xmin, xmax, 500) + >>> fig, ax = plt.subplots() + >>> ax.hlines(0, xmin, xmax, color='k') + >>> ax.plot(x, yn(2, x), label=r'$Y_2$') + >>> ax.scatter(yn_zeros(2, 4), np.zeros((4, )), s=30, c='r', + ... label='Roots', zorder=5) + >>> ax.set_ylim(-0.4, 0.4) + >>> ax.set_xlim(xmin, xmax) + >>> plt.legend() + >>> plt.show() + """ + return jnyn_zeros(n, nt)[2] + + +def ynp_zeros(n, nt): + r"""Compute zeros of integer-order Bessel function derivatives Yn'(x). + + Compute `nt` zeros of the functions :math:`Y_n'(x)` on the + interval :math:`(0, \infty)`. The zeros are returned in ascending + order. + + Parameters + ---------- + n : int + Order of Bessel function + nt : int + Number of zeros to return + + Returns + ------- + ndarray + First `nt` zeros of the Bessel derivative function. + + + See Also + -------- + yvp + + References + ---------- + .. [1] Zhang, Shanjie and Jin, Jianming. "Computation of Special + Functions", John Wiley and Sons, 1996, chapter 5. + https://people.sc.fsu.edu/~jburkardt/f77_src/special_functions/special_functions.html + + Examples + -------- + Compute the first four roots of the first derivative of the + Bessel function of second kind for order 0 :math:`Y_0'`. + + >>> from scipy.special import ynp_zeros + >>> ynp_zeros(0, 4) + array([ 2.19714133, 5.42968104, 8.59600587, 11.74915483]) + + Plot :math:`Y_0`, :math:`Y_0'` and confirm visually that the roots of + :math:`Y_0'` are located at local extrema of :math:`Y_0`. + + >>> import numpy as np + >>> import matplotlib.pyplot as plt + >>> from scipy.special import yn, ynp_zeros, yvp + >>> zeros = ynp_zeros(0, 4) + >>> xmax = 13 + >>> x = np.linspace(0, xmax, 500) + >>> fig, ax = plt.subplots() + >>> ax.plot(x, yn(0, x), label=r'$Y_0$') + >>> ax.plot(x, yvp(0, x, 1), label=r"$Y_0'$") + >>> ax.scatter(zeros, np.zeros((4, )), s=30, c='r', + ... label=r"Roots of $Y_0'$", zorder=5) + >>> for root in zeros: + ... y0_extremum = yn(0, root) + ... lower = min(0, y0_extremum) + ... upper = max(0, y0_extremum) + ... ax.vlines(root, lower, upper, color='r') + >>> ax.hlines(0, 0, xmax, color='k') + >>> ax.set_ylim(-0.6, 0.6) + >>> ax.set_xlim(0, xmax) + >>> plt.legend() + >>> plt.show() + """ + return jnyn_zeros(n, nt)[3] + + +def y0_zeros(nt, complex=False): + """Compute nt zeros of Bessel function Y0(z), and derivative at each zero. + + The derivatives are given by Y0'(z0) = -Y1(z0) at each zero z0. + + Parameters + ---------- + nt : int + Number of zeros to return + complex : bool, default False + Set to False to return only the real zeros; set to True to return only + the complex zeros with negative real part and positive imaginary part. + Note that the complex conjugates of the latter are also zeros of the + function, but are not returned by this routine. + + Returns + ------- + z0n : ndarray + Location of nth zero of Y0(z) + y0pz0n : ndarray + Value of derivative Y0'(z0) for nth zero + + References + ---------- + .. [1] Zhang, Shanjie and Jin, Jianming. "Computation of Special + Functions", John Wiley and Sons, 1996, chapter 5. + https://people.sc.fsu.edu/~jburkardt/f77_src/special_functions/special_functions.html + + Examples + -------- + Compute the first 4 real roots and the derivatives at the roots of + :math:`Y_0`: + + >>> import numpy as np + >>> from scipy.special import y0_zeros + >>> zeros, grads = y0_zeros(4) + >>> with np.printoptions(precision=5): + ... print(f"Roots: {zeros}") + ... print(f"Gradients: {grads}") + Roots: [ 0.89358+0.j 3.95768+0.j 7.08605+0.j 10.22235+0.j] + Gradients: [-0.87942+0.j 0.40254+0.j -0.3001 +0.j 0.2497 +0.j] + + Plot the real part of :math:`Y_0` and the first four computed roots. + + >>> import matplotlib.pyplot as plt + >>> from scipy.special import y0 + >>> xmin = 0 + >>> xmax = 11 + >>> x = np.linspace(xmin, xmax, 500) + >>> fig, ax = plt.subplots() + >>> ax.hlines(0, xmin, xmax, color='k') + >>> ax.plot(x, y0(x), label=r'$Y_0$') + >>> zeros, grads = y0_zeros(4) + >>> ax.scatter(zeros.real, np.zeros((4, )), s=30, c='r', + ... label=r'$Y_0$_zeros', zorder=5) + >>> ax.set_ylim(-0.5, 0.6) + >>> ax.set_xlim(xmin, xmax) + >>> plt.legend(ncol=2) + >>> plt.show() + + Compute the first 4 complex roots and the derivatives at the roots of + :math:`Y_0` by setting ``complex=True``: + + >>> y0_zeros(4, True) + (array([ -2.40301663+0.53988231j, -5.5198767 +0.54718001j, + -8.6536724 +0.54841207j, -11.79151203+0.54881912j]), + array([ 0.10074769-0.88196771j, -0.02924642+0.5871695j , + 0.01490806-0.46945875j, -0.00937368+0.40230454j])) + """ + if not isscalar(nt) or (floor(nt) != nt) or (nt <= 0): + raise ValueError("Arguments must be scalar positive integer.") + kf = 0 + kc = not complex + return _specfun.cyzo(nt, kf, kc) + + +def y1_zeros(nt, complex=False): + """Compute nt zeros of Bessel function Y1(z), and derivative at each zero. + + The derivatives are given by Y1'(z1) = Y0(z1) at each zero z1. + + Parameters + ---------- + nt : int + Number of zeros to return + complex : bool, default False + Set to False to return only the real zeros; set to True to return only + the complex zeros with negative real part and positive imaginary part. + Note that the complex conjugates of the latter are also zeros of the + function, but are not returned by this routine. + + Returns + ------- + z1n : ndarray + Location of nth zero of Y1(z) + y1pz1n : ndarray + Value of derivative Y1'(z1) for nth zero + + References + ---------- + .. [1] Zhang, Shanjie and Jin, Jianming. "Computation of Special + Functions", John Wiley and Sons, 1996, chapter 5. + https://people.sc.fsu.edu/~jburkardt/f77_src/special_functions/special_functions.html + + Examples + -------- + Compute the first 4 real roots and the derivatives at the roots of + :math:`Y_1`: + + >>> import numpy as np + >>> from scipy.special import y1_zeros + >>> zeros, grads = y1_zeros(4) + >>> with np.printoptions(precision=5): + ... print(f"Roots: {zeros}") + ... print(f"Gradients: {grads}") + Roots: [ 2.19714+0.j 5.42968+0.j 8.59601+0.j 11.74915+0.j] + Gradients: [ 0.52079+0.j -0.34032+0.j 0.27146+0.j -0.23246+0.j] + + Extract the real parts: + + >>> realzeros = zeros.real + >>> realzeros + array([ 2.19714133, 5.42968104, 8.59600587, 11.74915483]) + + Plot :math:`Y_1` and the first four computed roots. + + >>> import matplotlib.pyplot as plt + >>> from scipy.special import y1 + >>> xmin = 0 + >>> xmax = 13 + >>> x = np.linspace(xmin, xmax, 500) + >>> zeros, grads = y1_zeros(4) + >>> fig, ax = plt.subplots() + >>> ax.hlines(0, xmin, xmax, color='k') + >>> ax.plot(x, y1(x), label=r'$Y_1$') + >>> ax.scatter(zeros.real, np.zeros((4, )), s=30, c='r', + ... label=r'$Y_1$_zeros', zorder=5) + >>> ax.set_ylim(-0.5, 0.5) + >>> ax.set_xlim(xmin, xmax) + >>> plt.legend() + >>> plt.show() + + Compute the first 4 complex roots and the derivatives at the roots of + :math:`Y_1` by setting ``complex=True``: + + >>> y1_zeros(4, True) + (array([ -0.50274327+0.78624371j, -3.83353519+0.56235654j, + -7.01590368+0.55339305j, -10.17357383+0.55127339j]), + array([-0.45952768+1.31710194j, 0.04830191-0.69251288j, + -0.02012695+0.51864253j, 0.011614 -0.43203296j])) + """ + if not isscalar(nt) or (floor(nt) != nt) or (nt <= 0): + raise ValueError("Arguments must be scalar positive integer.") + kf = 1 + kc = not complex + return _specfun.cyzo(nt, kf, kc) + + +def y1p_zeros(nt, complex=False): + """Compute nt zeros of Bessel derivative Y1'(z), and value at each zero. + + The values are given by Y1(z1) at each z1 where Y1'(z1)=0. + + Parameters + ---------- + nt : int + Number of zeros to return + complex : bool, default False + Set to False to return only the real zeros; set to True to return only + the complex zeros with negative real part and positive imaginary part. + Note that the complex conjugates of the latter are also zeros of the + function, but are not returned by this routine. + + Returns + ------- + z1pn : ndarray + Location of nth zero of Y1'(z) + y1z1pn : ndarray + Value of derivative Y1(z1) for nth zero + + References + ---------- + .. [1] Zhang, Shanjie and Jin, Jianming. "Computation of Special + Functions", John Wiley and Sons, 1996, chapter 5. + https://people.sc.fsu.edu/~jburkardt/f77_src/special_functions/special_functions.html + + Examples + -------- + Compute the first four roots of :math:`Y_1'` and the values of + :math:`Y_1` at these roots. + + >>> import numpy as np + >>> from scipy.special import y1p_zeros + >>> y1grad_roots, y1_values = y1p_zeros(4) + >>> with np.printoptions(precision=5): + ... print(f"Y1' Roots: {y1grad_roots.real}") + ... print(f"Y1 values: {y1_values.real}") + Y1' Roots: [ 3.68302 6.9415 10.1234 13.28576] + Y1 values: [ 0.41673 -0.30317 0.25091 -0.21897] + + `y1p_zeros` can be used to calculate the extremal points of :math:`Y_1` + directly. Here we plot :math:`Y_1` and the first four extrema. + + >>> import matplotlib.pyplot as plt + >>> from scipy.special import y1, yvp + >>> y1_roots, y1_values_at_roots = y1p_zeros(4) + >>> real_roots = y1_roots.real + >>> xmax = 15 + >>> x = np.linspace(0, xmax, 500) + >>> x[0] += 1e-15 + >>> fig, ax = plt.subplots() + >>> ax.plot(x, y1(x), label=r'$Y_1$') + >>> ax.plot(x, yvp(1, x, 1), label=r"$Y_1'$") + >>> ax.scatter(real_roots, np.zeros((4, )), s=30, c='r', + ... label=r"Roots of $Y_1'$", zorder=5) + >>> ax.scatter(real_roots, y1_values_at_roots.real, s=30, c='k', + ... label=r"Extrema of $Y_1$", zorder=5) + >>> ax.hlines(0, 0, xmax, color='k') + >>> ax.set_ylim(-0.5, 0.5) + >>> ax.set_xlim(0, xmax) + >>> ax.legend(ncol=2, bbox_to_anchor=(1., 0.75)) + >>> plt.tight_layout() + >>> plt.show() + """ + if not isscalar(nt) or (floor(nt) != nt) or (nt <= 0): + raise ValueError("Arguments must be scalar positive integer.") + kf = 2 + kc = not complex + return _specfun.cyzo(nt, kf, kc) + + +def _bessel_diff_formula(v, z, n, L, phase): + # from AMS55. + # L(v, z) = J(v, z), Y(v, z), H1(v, z), H2(v, z), phase = -1 + # L(v, z) = I(v, z) or exp(v*pi*i)K(v, z), phase = 1 + # For K, you can pull out the exp((v-k)*pi*i) into the caller + v = asarray(v) + p = 1.0 + s = L(v-n, z) + for i in range(1, n+1): + p = phase * (p * (n-i+1)) / i # = choose(k, i) + s += p*L(v-n + i*2, z) + return s / (2.**n) + + +def jvp(v, z, n=1): + """Compute derivatives of Bessel functions of the first kind. + + Compute the nth derivative of the Bessel function `Jv` with + respect to `z`. + + Parameters + ---------- + v : array_like or float + Order of Bessel function + z : complex + Argument at which to evaluate the derivative; can be real or + complex. + n : int, default 1 + Order of derivative. For 0 returns the Bessel function `jv` itself. + + Returns + ------- + scalar or ndarray + Values of the derivative of the Bessel function. + + Notes + ----- + The derivative is computed using the relation DLFM 10.6.7 [2]_. + + References + ---------- + .. [1] Zhang, Shanjie and Jin, Jianming. "Computation of Special + Functions", John Wiley and Sons, 1996, chapter 5. + https://people.sc.fsu.edu/~jburkardt/f77_src/special_functions/special_functions.html + + .. [2] NIST Digital Library of Mathematical Functions. + https://dlmf.nist.gov/10.6.E7 + + Examples + -------- + + Compute the Bessel function of the first kind of order 0 and + its first two derivatives at 1. + + >>> from scipy.special import jvp + >>> jvp(0, 1, 0), jvp(0, 1, 1), jvp(0, 1, 2) + (0.7651976865579666, -0.44005058574493355, -0.3251471008130331) + + Compute the first derivative of the Bessel function of the first + kind for several orders at 1 by providing an array for `v`. + + >>> jvp([0, 1, 2], 1, 1) + array([-0.44005059, 0.3251471 , 0.21024362]) + + Compute the first derivative of the Bessel function of the first + kind of order 0 at several points by providing an array for `z`. + + >>> import numpy as np + >>> points = np.array([0., 1.5, 3.]) + >>> jvp(0, points, 1) + array([-0. , -0.55793651, -0.33905896]) + + Plot the Bessel function of the first kind of order 1 and its + first three derivatives. + + >>> import matplotlib.pyplot as plt + >>> x = np.linspace(-10, 10, 1000) + >>> fig, ax = plt.subplots() + >>> ax.plot(x, jvp(1, x, 0), label=r"$J_1$") + >>> ax.plot(x, jvp(1, x, 1), label=r"$J_1'$") + >>> ax.plot(x, jvp(1, x, 2), label=r"$J_1''$") + >>> ax.plot(x, jvp(1, x, 3), label=r"$J_1'''$") + >>> plt.legend() + >>> plt.show() + """ + n = _nonneg_int_or_fail(n, 'n') + if n == 0: + return jv(v, z) + else: + return _bessel_diff_formula(v, z, n, jv, -1) + + +def yvp(v, z, n=1): + """Compute derivatives of Bessel functions of the second kind. + + Compute the nth derivative of the Bessel function `Yv` with + respect to `z`. + + Parameters + ---------- + v : array_like of float + Order of Bessel function + z : complex + Argument at which to evaluate the derivative + n : int, default 1 + Order of derivative. For 0 returns the BEssel function `yv` + + Returns + ------- + scalar or ndarray + nth derivative of the Bessel function. + + See Also + -------- + yv : Bessel functions of the second kind + + Notes + ----- + The derivative is computed using the relation DLFM 10.6.7 [2]_. + + References + ---------- + .. [1] Zhang, Shanjie and Jin, Jianming. "Computation of Special + Functions", John Wiley and Sons, 1996, chapter 5. + https://people.sc.fsu.edu/~jburkardt/f77_src/special_functions/special_functions.html + + .. [2] NIST Digital Library of Mathematical Functions. + https://dlmf.nist.gov/10.6.E7 + + Examples + -------- + Compute the Bessel function of the second kind of order 0 and + its first two derivatives at 1. + + >>> from scipy.special import yvp + >>> yvp(0, 1, 0), yvp(0, 1, 1), yvp(0, 1, 2) + (0.088256964215677, 0.7812128213002889, -0.8694697855159659) + + Compute the first derivative of the Bessel function of the second + kind for several orders at 1 by providing an array for `v`. + + >>> yvp([0, 1, 2], 1, 1) + array([0.78121282, 0.86946979, 2.52015239]) + + Compute the first derivative of the Bessel function of the + second kind of order 0 at several points by providing an array for `z`. + + >>> import numpy as np + >>> points = np.array([0.5, 1.5, 3.]) + >>> yvp(0, points, 1) + array([ 1.47147239, 0.41230863, -0.32467442]) + + Plot the Bessel function of the second kind of order 1 and its + first three derivatives. + + >>> import matplotlib.pyplot as plt + >>> x = np.linspace(0, 5, 1000) + >>> x[0] += 1e-15 + >>> fig, ax = plt.subplots() + >>> ax.plot(x, yvp(1, x, 0), label=r"$Y_1$") + >>> ax.plot(x, yvp(1, x, 1), label=r"$Y_1'$") + >>> ax.plot(x, yvp(1, x, 2), label=r"$Y_1''$") + >>> ax.plot(x, yvp(1, x, 3), label=r"$Y_1'''$") + >>> ax.set_ylim(-10, 10) + >>> plt.legend() + >>> plt.show() + """ + n = _nonneg_int_or_fail(n, 'n') + if n == 0: + return yv(v, z) + else: + return _bessel_diff_formula(v, z, n, yv, -1) + + +def kvp(v, z, n=1): + """Compute derivatives of real-order modified Bessel function Kv(z). + + Kv(z) is the modified Bessel function of the second kind. + Derivative is calculated with respect to `z`. + + Parameters + ---------- + v : array_like of float + Order of Bessel function + z : array_like of complex + Argument at which to evaluate the derivative + n : int, default 1 + Order of derivative. For 0 returns the Bessel function `kv` itself. + + Returns + ------- + out : ndarray + The results + + See Also + -------- + kv + + Notes + ----- + The derivative is computed using the relation DLFM 10.29.5 [2]_. + + References + ---------- + .. [1] Zhang, Shanjie and Jin, Jianming. "Computation of Special + Functions", John Wiley and Sons, 1996, chapter 6. + https://people.sc.fsu.edu/~jburkardt/f77_src/special_functions/special_functions.html + + .. [2] NIST Digital Library of Mathematical Functions. + https://dlmf.nist.gov/10.29.E5 + + Examples + -------- + Compute the modified bessel function of the second kind of order 0 and + its first two derivatives at 1. + + >>> from scipy.special import kvp + >>> kvp(0, 1, 0), kvp(0, 1, 1), kvp(0, 1, 2) + (0.42102443824070834, -0.6019072301972346, 1.0229316684379428) + + Compute the first derivative of the modified Bessel function of the second + kind for several orders at 1 by providing an array for `v`. + + >>> kvp([0, 1, 2], 1, 1) + array([-0.60190723, -1.02293167, -3.85158503]) + + Compute the first derivative of the modified Bessel function of the + second kind of order 0 at several points by providing an array for `z`. + + >>> import numpy as np + >>> points = np.array([0.5, 1.5, 3.]) + >>> kvp(0, points, 1) + array([-1.65644112, -0.2773878 , -0.04015643]) + + Plot the modified bessel function of the second kind and its + first three derivatives. + + >>> import matplotlib.pyplot as plt + >>> x = np.linspace(0, 5, 1000) + >>> fig, ax = plt.subplots() + >>> ax.plot(x, kvp(1, x, 0), label=r"$K_1$") + >>> ax.plot(x, kvp(1, x, 1), label=r"$K_1'$") + >>> ax.plot(x, kvp(1, x, 2), label=r"$K_1''$") + >>> ax.plot(x, kvp(1, x, 3), label=r"$K_1'''$") + >>> ax.set_ylim(-2.5, 2.5) + >>> plt.legend() + >>> plt.show() + """ + n = _nonneg_int_or_fail(n, 'n') + if n == 0: + return kv(v, z) + else: + return (-1)**n * _bessel_diff_formula(v, z, n, kv, 1) + + +def ivp(v, z, n=1): + """Compute derivatives of modified Bessel functions of the first kind. + + Compute the nth derivative of the modified Bessel function `Iv` + with respect to `z`. + + Parameters + ---------- + v : array_like or float + Order of Bessel function + z : array_like + Argument at which to evaluate the derivative; can be real or + complex. + n : int, default 1 + Order of derivative. For 0, returns the Bessel function `iv` itself. + + Returns + ------- + scalar or ndarray + nth derivative of the modified Bessel function. + + See Also + -------- + iv + + Notes + ----- + The derivative is computed using the relation DLFM 10.29.5 [2]_. + + References + ---------- + .. [1] Zhang, Shanjie and Jin, Jianming. "Computation of Special + Functions", John Wiley and Sons, 1996, chapter 6. + https://people.sc.fsu.edu/~jburkardt/f77_src/special_functions/special_functions.html + + .. [2] NIST Digital Library of Mathematical Functions. + https://dlmf.nist.gov/10.29.E5 + + Examples + -------- + Compute the modified Bessel function of the first kind of order 0 and + its first two derivatives at 1. + + >>> from scipy.special import ivp + >>> ivp(0, 1, 0), ivp(0, 1, 1), ivp(0, 1, 2) + (1.2660658777520084, 0.565159103992485, 0.7009067737595233) + + Compute the first derivative of the modified Bessel function of the first + kind for several orders at 1 by providing an array for `v`. + + >>> ivp([0, 1, 2], 1, 1) + array([0.5651591 , 0.70090677, 0.29366376]) + + Compute the first derivative of the modified Bessel function of the + first kind of order 0 at several points by providing an array for `z`. + + >>> import numpy as np + >>> points = np.array([0., 1.5, 3.]) + >>> ivp(0, points, 1) + array([0. , 0.98166643, 3.95337022]) + + Plot the modified Bessel function of the first kind of order 1 and its + first three derivatives. + + >>> import matplotlib.pyplot as plt + >>> x = np.linspace(-5, 5, 1000) + >>> fig, ax = plt.subplots() + >>> ax.plot(x, ivp(1, x, 0), label=r"$I_1$") + >>> ax.plot(x, ivp(1, x, 1), label=r"$I_1'$") + >>> ax.plot(x, ivp(1, x, 2), label=r"$I_1''$") + >>> ax.plot(x, ivp(1, x, 3), label=r"$I_1'''$") + >>> plt.legend() + >>> plt.show() + """ + n = _nonneg_int_or_fail(n, 'n') + if n == 0: + return iv(v, z) + else: + return _bessel_diff_formula(v, z, n, iv, 1) + + +def h1vp(v, z, n=1): + """Compute derivatives of Hankel function H1v(z) with respect to `z`. + + Parameters + ---------- + v : array_like + Order of Hankel function + z : array_like + Argument at which to evaluate the derivative. Can be real or + complex. + n : int, default 1 + Order of derivative. For 0 returns the Hankel function `hankel1` itself. + + Returns + ------- + scalar or ndarray + Values of the derivative of the Hankel function. + + See Also + -------- + hankel1 + + Notes + ----- + The derivative is computed using the relation DLFM 10.6.7 [2]_. + + References + ---------- + .. [1] Zhang, Shanjie and Jin, Jianming. "Computation of Special + Functions", John Wiley and Sons, 1996, chapter 5. + https://people.sc.fsu.edu/~jburkardt/f77_src/special_functions/special_functions.html + + .. [2] NIST Digital Library of Mathematical Functions. + https://dlmf.nist.gov/10.6.E7 + + Examples + -------- + Compute the Hankel function of the first kind of order 0 and + its first two derivatives at 1. + + >>> from scipy.special import h1vp + >>> h1vp(0, 1, 0), h1vp(0, 1, 1), h1vp(0, 1, 2) + ((0.7651976865579664+0.088256964215677j), + (-0.44005058574493355+0.7812128213002889j), + (-0.3251471008130329-0.8694697855159659j)) + + Compute the first derivative of the Hankel function of the first kind + for several orders at 1 by providing an array for `v`. + + >>> h1vp([0, 1, 2], 1, 1) + array([-0.44005059+0.78121282j, 0.3251471 +0.86946979j, + 0.21024362+2.52015239j]) + + Compute the first derivative of the Hankel function of the first kind + of order 0 at several points by providing an array for `z`. + + >>> import numpy as np + >>> points = np.array([0.5, 1.5, 3.]) + >>> h1vp(0, points, 1) + array([-0.24226846+1.47147239j, -0.55793651+0.41230863j, + -0.33905896-0.32467442j]) + """ + n = _nonneg_int_or_fail(n, 'n') + if n == 0: + return hankel1(v, z) + else: + return _bessel_diff_formula(v, z, n, hankel1, -1) + + +def h2vp(v, z, n=1): + """Compute derivatives of Hankel function H2v(z) with respect to `z`. + + Parameters + ---------- + v : array_like + Order of Hankel function + z : array_like + Argument at which to evaluate the derivative. Can be real or + complex. + n : int, default 1 + Order of derivative. For 0 returns the Hankel function `hankel2` itself. + + Returns + ------- + scalar or ndarray + Values of the derivative of the Hankel function. + + See Also + -------- + hankel2 + + Notes + ----- + The derivative is computed using the relation DLFM 10.6.7 [2]_. + + References + ---------- + .. [1] Zhang, Shanjie and Jin, Jianming. "Computation of Special + Functions", John Wiley and Sons, 1996, chapter 5. + https://people.sc.fsu.edu/~jburkardt/f77_src/special_functions/special_functions.html + + .. [2] NIST Digital Library of Mathematical Functions. + https://dlmf.nist.gov/10.6.E7 + + Examples + -------- + Compute the Hankel function of the second kind of order 0 and + its first two derivatives at 1. + + >>> from scipy.special import h2vp + >>> h2vp(0, 1, 0), h2vp(0, 1, 1), h2vp(0, 1, 2) + ((0.7651976865579664-0.088256964215677j), + (-0.44005058574493355-0.7812128213002889j), + (-0.3251471008130329+0.8694697855159659j)) + + Compute the first derivative of the Hankel function of the second kind + for several orders at 1 by providing an array for `v`. + + >>> h2vp([0, 1, 2], 1, 1) + array([-0.44005059-0.78121282j, 0.3251471 -0.86946979j, + 0.21024362-2.52015239j]) + + Compute the first derivative of the Hankel function of the second kind + of order 0 at several points by providing an array for `z`. + + >>> import numpy as np + >>> points = np.array([0.5, 1.5, 3.]) + >>> h2vp(0, points, 1) + array([-0.24226846-1.47147239j, -0.55793651-0.41230863j, + -0.33905896+0.32467442j]) + """ + n = _nonneg_int_or_fail(n, 'n') + if n == 0: + return hankel2(v, z) + else: + return _bessel_diff_formula(v, z, n, hankel2, -1) + + +def riccati_jn(n, x): + r"""Compute Riccati-Bessel function of the first kind and its derivative. + + The Riccati-Bessel function of the first kind is defined as :math:`x + j_n(x)`, where :math:`j_n` is the spherical Bessel function of the first + kind of order :math:`n`. + + This function computes the value and first derivative of the + Riccati-Bessel function for all orders up to and including `n`. + + Parameters + ---------- + n : int + Maximum order of function to compute + x : float + Argument at which to evaluate + + Returns + ------- + jn : ndarray + Value of j0(x), ..., jn(x) + jnp : ndarray + First derivative j0'(x), ..., jn'(x) + + Notes + ----- + The computation is carried out via backward recurrence, using the + relation DLMF 10.51.1 [2]_. + + Wrapper for a Fortran routine created by Shanjie Zhang and Jianming + Jin [1]_. + + References + ---------- + .. [1] Zhang, Shanjie and Jin, Jianming. "Computation of Special + Functions", John Wiley and Sons, 1996. + https://people.sc.fsu.edu/~jburkardt/f77_src/special_functions/special_functions.html + .. [2] NIST Digital Library of Mathematical Functions. + https://dlmf.nist.gov/10.51.E1 + + """ + if not (isscalar(n) and isscalar(x)): + raise ValueError("arguments must be scalars.") + n = _nonneg_int_or_fail(n, 'n', strict=False) + if (n == 0): + n1 = 1 + else: + n1 = n + + jn = np.empty((n1 + 1,), dtype=np.float64) + jnp = np.empty_like(jn) + + _rctj(x, out=(jn, jnp)) + return jn[:(n+1)], jnp[:(n+1)] + + +def riccati_yn(n, x): + """Compute Riccati-Bessel function of the second kind and its derivative. + + The Riccati-Bessel function of the second kind is defined here as :math:`+x + y_n(x)`, where :math:`y_n` is the spherical Bessel function of the second + kind of order :math:`n`. *Note that this is in contrast to a common convention + that includes a minus sign in the definition.* + + This function computes the value and first derivative of the function for + all orders up to and including `n`. + + Parameters + ---------- + n : int + Maximum order of function to compute + x : float + Argument at which to evaluate + + Returns + ------- + yn : ndarray + Value of y0(x), ..., yn(x) + ynp : ndarray + First derivative y0'(x), ..., yn'(x) + + Notes + ----- + The computation is carried out via ascending recurrence, using the + relation DLMF 10.51.1 [2]_. + + Wrapper for a Fortran routine created by Shanjie Zhang and Jianming + Jin [1]_. + + References + ---------- + .. [1] Zhang, Shanjie and Jin, Jianming. "Computation of Special + Functions", John Wiley and Sons, 1996. + https://people.sc.fsu.edu/~jburkardt/f77_src/special_functions/special_functions.html + .. [2] NIST Digital Library of Mathematical Functions. + https://dlmf.nist.gov/10.51.E1 + + """ + if not (isscalar(n) and isscalar(x)): + raise ValueError("arguments must be scalars.") + n = _nonneg_int_or_fail(n, 'n', strict=False) + if (n == 0): + n1 = 1 + else: + n1 = n + + yn = np.empty((n1 + 1,), dtype=np.float64) + ynp = np.empty_like(yn) + _rcty(x, out=(yn, ynp)) + + return yn[:(n+1)], ynp[:(n+1)] + + +def erf_zeros(nt): + """Compute the first nt zero in the first quadrant, ordered by absolute value. + + Zeros in the other quadrants can be obtained by using the symmetries + erf(-z) = erf(z) and erf(conj(z)) = conj(erf(z)). + + + Parameters + ---------- + nt : int + The number of zeros to compute + + Returns + ------- + The locations of the zeros of erf : ndarray (complex) + Complex values at which zeros of erf(z) + + References + ---------- + .. [1] Zhang, Shanjie and Jin, Jianming. "Computation of Special + Functions", John Wiley and Sons, 1996. + https://people.sc.fsu.edu/~jburkardt/f77_src/special_functions/special_functions.html + + Examples + -------- + >>> from scipy import special + >>> special.erf_zeros(1) + array([1.45061616+1.880943j]) + + Check that erf is (close to) zero for the value returned by erf_zeros + + >>> special.erf(special.erf_zeros(1)) + array([4.95159469e-14-1.16407394e-16j]) + + """ + if (floor(nt) != nt) or (nt <= 0) or not isscalar(nt): + raise ValueError("Argument must be positive scalar integer.") + return _specfun.cerzo(nt) + + +def fresnelc_zeros(nt): + """Compute nt complex zeros of cosine Fresnel integral C(z). + + Parameters + ---------- + nt : int + Number of zeros to compute + + Returns + ------- + fresnelc_zeros: ndarray + Zeros of the cosine Fresnel integral + + References + ---------- + .. [1] Zhang, Shanjie and Jin, Jianming. "Computation of Special + Functions", John Wiley and Sons, 1996. + https://people.sc.fsu.edu/~jburkardt/f77_src/special_functions/special_functions.html + + """ + if (floor(nt) != nt) or (nt <= 0) or not isscalar(nt): + raise ValueError("Argument must be positive scalar integer.") + return _specfun.fcszo(1, nt) + + +def fresnels_zeros(nt): + """Compute nt complex zeros of sine Fresnel integral S(z). + + Parameters + ---------- + nt : int + Number of zeros to compute + + Returns + ------- + fresnels_zeros: ndarray + Zeros of the sine Fresnel integral + + References + ---------- + .. [1] Zhang, Shanjie and Jin, Jianming. "Computation of Special + Functions", John Wiley and Sons, 1996. + https://people.sc.fsu.edu/~jburkardt/f77_src/special_functions/special_functions.html + + """ + if (floor(nt) != nt) or (nt <= 0) or not isscalar(nt): + raise ValueError("Argument must be positive scalar integer.") + return _specfun.fcszo(2, nt) + + +def fresnel_zeros(nt): + """Compute nt complex zeros of sine and cosine Fresnel integrals S(z) and C(z). + + Parameters + ---------- + nt : int + Number of zeros to compute + + Returns + ------- + zeros_sine: ndarray + Zeros of the sine Fresnel integral + zeros_cosine : ndarray + Zeros of the cosine Fresnel integral + + References + ---------- + .. [1] Zhang, Shanjie and Jin, Jianming. "Computation of Special + Functions", John Wiley and Sons, 1996. + https://people.sc.fsu.edu/~jburkardt/f77_src/special_functions/special_functions.html + + """ + if (floor(nt) != nt) or (nt <= 0) or not isscalar(nt): + raise ValueError("Argument must be positive scalar integer.") + return _specfun.fcszo(2, nt), _specfun.fcszo(1, nt) + + +def assoc_laguerre(x, n, k=0.0): + r"""Compute the generalized (associated) Laguerre polynomial of degree :math:`n` + and order :math:`k`. + + This function evaluates the generalized Laguerre polynomial + + .. math:: + + L^{(k)}_n(x). + + The polynomial is orthogonal over :math:`[0, \infty)` with weight + function + + .. math:: + + e^{-x}x^k + + for :math:`k > -1`. + + Parameters + ---------- + x : float or ndarray + Points where to evaluate the Laguerre polynomial + n : int + Degree of the Laguerre polynomial + k : int + Order of the Laguerre polynomial + + Returns + ------- + assoc_laguerre: float or ndarray + Associated laguerre polynomial values + + Notes + ----- + `assoc_laguerre` is a simple wrapper around `eval_genlaguerre`, with + reversed argument order ``(x, n, k=0.0) --> (n, k, x)``. + + Examples + -------- + Evaluate the associated Laguerre polynomial :math:`L_3^{(2)}` at + :math:`x = 1`: + + >>> import numpy as np + >>> from scipy.special import assoc_laguerre, eval_genlaguerre + >>> np.isclose(assoc_laguerre(1, 3, 2), 7/3) + True + + `assoc_laguerre` is equivalent to `eval_genlaguerre` with reversed + argument order: + + >>> x = np.linspace(0, 5, 6) + >>> np.allclose(assoc_laguerre(x, 3, 2), eval_genlaguerre(3, 2, x)) + True + + Plot :math:`L_3^{(k)}` for several values of :math:`k`: + + >>> import matplotlib.pyplot as plt + >>> x = np.linspace(0, 8, 400) + >>> fig, ax = plt.subplots() + >>> for k in range(3): + ... ax.plot(x, assoc_laguerre(x, 3, k), label=rf"$k={k}$") + >>> ax.set_title(r"Associated Laguerre polynomials $L_3^{(k)}$") + >>> ax.set_xlabel("x") + >>> ax.legend(loc="best") + >>> plt.show() + + """ + return _ufuncs.eval_genlaguerre(n, k, x) + + +digamma = psi + + +def polygamma(n, x): + r"""Polygamma functions. + + Defined as :math:`\psi^{(n)}(x)` where :math:`\psi` is the + `digamma` function. See [dlmf]_ for details. + + Parameters + ---------- + n : array_like + The order of the derivative of the digamma function; must be + integral + x : array_like + Real valued input + + Returns + ------- + ndarray + Function results + + See Also + -------- + digamma + + References + ---------- + .. [dlmf] NIST, Digital Library of Mathematical Functions, + https://dlmf.nist.gov/5.15 + + Examples + -------- + >>> from scipy import special + >>> x = [2, 3, 25.5] + >>> special.polygamma(1, x) + array([ 0.64493407, 0.39493407, 0.03999467]) + >>> special.polygamma(0, x) == special.psi(x) + array([ True, True, True], dtype=bool) + + """ + n, x = asarray(n), asarray(x) + fac2 = (-1.0)**(n+1) * gamma(n+1.0) * zeta(n+1, x) + return where(n == 0, psi(x), fac2) + + +def mathieu_even_coef(m, q): + r"""Fourier coefficients for even Mathieu and modified Mathieu functions. + + The Fourier series of the even solutions of the Mathieu differential + equation are of the form + + .. math:: \mathrm{ce}_{2n}(z, q) = \sum_{k=0}^{\infty} A_{(2n)}^{(2k)} \cos 2kz + + .. math:: \mathrm{ce}_{2n+1}(z, q) = + \sum_{k=0}^{\infty} A_{(2n+1)}^{(2k+1)} \cos (2k+1)z + + This function returns the coefficients :math:`A_{(2n)}^{(2k)}` for even + input m=2n, and the coefficients :math:`A_{(2n+1)}^{(2k+1)}` for odd input + m=2n+1. + + Parameters + ---------- + m : int + Order of Mathieu functions. Must be non-negative. + q : float (>=0) + Parameter of Mathieu functions. Must be non-negative. + + Returns + ------- + Ak : ndarray + Even or odd Fourier coefficients, corresponding to even or odd m. + + References + ---------- + .. [1] Zhang, Shanjie and Jin, Jianming. "Computation of Special + Functions", John Wiley and Sons, 1996. + https://people.sc.fsu.edu/~jburkardt/f77_src/special_functions/special_functions.html + .. [2] NIST Digital Library of Mathematical Functions + https://dlmf.nist.gov/28.4#i + + """ + if not (isscalar(m) and isscalar(q)): + raise ValueError("m and q must be scalars.") + if (q < 0): + raise ValueError("q >=0") + if (m != floor(m)) or (m < 0): + raise ValueError("m must be an integer >=0.") + + if (q <= 1): + qm = 7.5 + 56.1*sqrt(q) - 134.7*q + 90.7*sqrt(q)*q + else: + qm = 17.0 + 3.1*sqrt(q) - .126*q + .0037*sqrt(q)*q + km = int(qm + 0.5*m) + if km > 251: + warnings.warn("Too many predicted coefficients.", RuntimeWarning, stacklevel=2) + kd = 1 + m = int(floor(m)) + if m % 2: + kd = 2 + + a = mathieu_a(m, q) + fc = _specfun.fcoef(kd, m, q, a) + return fc[:km] + + +def mathieu_odd_coef(m, q): + r"""Fourier coefficients for odd Mathieu and modified Mathieu functions. + + The Fourier series of the odd solutions of the Mathieu differential + equation are of the form + + .. math:: \mathrm{se}_{2n+1}(z, q) = + \sum_{k=0}^{\infty} B_{(2n+1)}^{(2k+1)} \sin (2k+1)z + + .. math:: \mathrm{se}_{2n+2}(z, q) = + \sum_{k=0}^{\infty} B_{(2n+2)}^{(2k+2)} \sin (2k+2)z + + This function returns the coefficients :math:`B_{(2n+2)}^{(2k+2)}` for even + input m=2n+2, and the coefficients :math:`B_{(2n+1)}^{(2k+1)}` for odd + input m=2n+1. + + Parameters + ---------- + m : int + Order of Mathieu functions. Must be non-negative. + q : float (>=0) + Parameter of Mathieu functions. Must be non-negative. + + Returns + ------- + Bk : ndarray + Even or odd Fourier coefficients, corresponding to even or odd m. + + References + ---------- + .. [1] Zhang, Shanjie and Jin, Jianming. "Computation of Special + Functions", John Wiley and Sons, 1996. + https://people.sc.fsu.edu/~jburkardt/f77_src/special_functions/special_functions.html + + """ + if not (isscalar(m) and isscalar(q)): + raise ValueError("m and q must be scalars.") + if (q < 0): + raise ValueError("q >=0") + if (m != floor(m)) or (m <= 0): + raise ValueError("m must be an integer > 0") + + if (q <= 1): + qm = 7.5 + 56.1*sqrt(q) - 134.7*q + 90.7*sqrt(q)*q + else: + qm = 17.0 + 3.1*sqrt(q) - .126*q + .0037*sqrt(q)*q + km = int(qm + 0.5*m) + if km > 251: + warnings.warn("Too many predicted coefficients.", RuntimeWarning, stacklevel=2) + kd = 4 + m = int(floor(m)) + if m % 2: + kd = 3 + + b = mathieu_b(m, q) + fc = _specfun.fcoef(kd, m, q, b) + return fc[:km] + + +def lqmn(m, n, z): + """Sequence of associated Legendre functions of the second kind. + + Computes the associated Legendre function of the second kind of order m and + degree n, ``Qmn(z)`` = :math:`Q_n^m(z)`, and its derivative, ``Qmn'(z)``. + Returns two arrays of size ``(m+1, n+1)`` containing ``Qmn(z)`` and + ``Qmn'(z)`` for all orders from ``0..m`` and degrees from ``0..n``. + + Parameters + ---------- + m : int + ``|m| <= n``; the order of the Legendre function. + n : int + where ``n >= 0``; the degree of the Legendre function. Often + called ``l`` (lower case L) in descriptions of the associated + Legendre function + z : array_like, complex + Input value. + + Returns + ------- + Qmn_z : (m+1, n+1) array + Values for all orders 0..m and degrees 0..n + Qmn_d_z : (m+1, n+1) array + Derivatives for all orders 0..m and degrees 0..n + + References + ---------- + .. [1] Zhang, Shanjie and Jin, Jianming. "Computation of Special + Functions", John Wiley and Sons, 1996. + https://people.sc.fsu.edu/~jburkardt/f77_src/special_functions/special_functions.html + + """ + if not isscalar(m) or (m < 0): + raise ValueError("m must be a non-negative integer.") + if not isscalar(n) or (n < 0): + raise ValueError("n must be a non-negative integer.") + + m, n = int(m), int(n) # Convert to int to maintain backwards compatibility. + # Ensure neither m nor n == 0 + mm = max(1, m) + nn = max(1, n) + + z = np.asarray(z) + if (not np.issubdtype(z.dtype, np.inexact)): + z = z.astype(np.float64) + + if np.iscomplexobj(z): + q = np.empty((mm + 1, nn + 1) + z.shape, dtype=np.complex128) + else: + q = np.empty((mm + 1, nn + 1) + z.shape, dtype=np.float64) + qd = np.empty_like(q) + if (z.ndim == 0): + _lqmn(z, out=(q, qd)) + else: + # new axes must be last for the ufunc + _lqmn(z, + out=(np.moveaxis(q, (0, 1), (-2, -1)), + np.moveaxis(qd, (0, 1), (-2, -1)))) + + return q[:(m+1), :(n+1)], qd[:(m+1), :(n+1)] + + +def bernoulli(n): + """Bernoulli numbers B0..Bn (inclusive). + + Parameters + ---------- + n : int + Indicated the number of terms in the Bernoulli series to generate. + + Returns + ------- + ndarray + The Bernoulli numbers ``[B(0), B(1), ..., B(n)]``. + + References + ---------- + .. [1] Zhang, Shanjie and Jin, Jianming. "Computation of Special + Functions", John Wiley and Sons, 1996. + https://people.sc.fsu.edu/~jburkardt/f77_src/special_functions/special_functions.html + .. [2] "Bernoulli number", Wikipedia, https://en.wikipedia.org/wiki/Bernoulli_number + + Examples + -------- + >>> import numpy as np + >>> from scipy.special import bernoulli, zeta + >>> bernoulli(4) + array([ 1. , -0.5 , 0.16666667, 0. , -0.03333333]) + + The Wikipedia article ([2]_) points out the relationship between the + Bernoulli numbers and the zeta function, ``B_n^+ = -n * zeta(1 - n)`` + for ``n > 0``: + + >>> n = np.arange(1, 5) + >>> -n * zeta(1 - n) + array([ 0.5 , 0.16666667, -0. , -0.03333333]) + + Note that, in the notation used in the wikipedia article, + `bernoulli` computes ``B_n^-`` (i.e. it used the convention that + ``B_1`` is -1/2). The relation given above is for ``B_n^+``, so the + sign of 0.5 does not match the output of ``bernoulli(4)``. + + """ + if not isscalar(n) or (n < 0): + raise ValueError("n must be a non-negative integer.") + n = int(n) + if (n < 2): + n1 = 2 + else: + n1 = n + return _specfun.bernob(int(n1))[:(n+1)] + + +def euler(n): + """Euler numbers E(0), E(1), ..., E(n). + + The Euler numbers [1]_ are also known as the secant numbers. + + Because ``euler(n)`` returns floating point values, it does not give + exact values for large `n`. The first inexact value is E(22). + + Parameters + ---------- + n : int + The highest index of the Euler number to be returned. + + Returns + ------- + ndarray + The Euler numbers [E(0), E(1), ..., E(n)]. + The odd Euler numbers, which are all zero, are included. + + References + ---------- + .. [1] Sequence A122045, The On-Line Encyclopedia of Integer Sequences, + https://oeis.org/A122045 + .. [2] Zhang, Shanjie and Jin, Jianming. "Computation of Special + Functions", John Wiley and Sons, 1996. + https://people.sc.fsu.edu/~jburkardt/f77_src/special_functions/special_functions.html + + Examples + -------- + >>> import numpy as np + >>> from scipy.special import euler + >>> euler(6) + array([ 1., 0., -1., 0., 5., 0., -61.]) + + >>> euler(13).astype(np.int64) + array([ 1, 0, -1, 0, 5, 0, -61, + 0, 1385, 0, -50521, 0, 2702765, 0]) + + >>> euler(22)[-1] # Exact value of E(22) is -69348874393137901. + -69348874393137976.0 + + """ + if not isscalar(n) or (n < 0): + raise ValueError("n must be a non-negative integer.") + n = int(n) + if (n < 2): + n1 = 2 + else: + n1 = n + return _specfun.eulerb(n1)[:(n+1)] + + +def lqn(n, z): + """Legendre functions of the second kind. + + Compute sequence of Legendre functions of the second kind, ``Qn(z)`` and + derivatives for all degrees from 0 to `n` (inclusive). + Returns two arrays of size ``(n+1,) + z.shape`` containing ``Qn(z)`` and + ``Qn'(z)``. + + Parameters + ---------- + n : int + Maximum degree of the Legendre functions. + z : array_like, complex + Real or complex input values. + + Returns + ------- + Qn_z : ndarray, shape (n+1,) + shape(z) + Values for all degrees ``0..n`` + Qn_d_z : ndarray, shape (n+1,) + shape(z) + Derivatives for all degrees ``0..n`` + + References + ---------- + .. [1] Zhang, Shanjie and Jin, Jianming. "Computation of Special + Functions", John Wiley and Sons, 1996. + https://people.sc.fsu.edu/~jburkardt/f77_src/special_functions/special_functions.html + + Examples + -------- + Compute :math:`Q_n(x)` and its derivatives on an interval. + + >>> import numpy as np + >>> from scipy.special import lqn + >>> import matplotlib.pyplot as plt + + >>> xs = np.linspace(-2, 2, 200) + >>> n_max = 3 + >>> Qn, dQn = lqn(n_max, xs) + + Plot the Legendre functions of the second kind :math:`Q_n(x)`. + + >>> fig, ax = plt.subplots() + >>> ax.plot(xs, Qn.T, "-") + >>> ax.set_xlabel(r"$x$") + >>> ax.set_ylabel(r"$Q_n(x)$") + >>> ax.legend([fr"$n={n}$" for n in range(n_max + 1)]) + >>> plt.show() + + Plot the derivatives :math:`Q_n'(x)`. + + >>> fig, ax = plt.subplots() + >>> ax.plot(xs, dQn.T, "-") + >>> ax.set_xlabel(r"$x$") + >>> ax.set_ylabel(r"$Q_n'(x)$") + >>> ax.legend([fr"$n={n}$" for n in range(n_max + 1)]) + >>> plt.show() + """ + n = _nonneg_int_or_fail(n, 'n', strict=False) + if (n < 1): + n1 = 1 + else: + n1 = n + + z = np.asarray(z) + if (not np.issubdtype(z.dtype, np.inexact)): + z = z.astype(float) + + if np.iscomplexobj(z): + qn = np.empty((n1 + 1,) + z.shape, dtype=np.complex128) + else: + qn = np.empty((n1 + 1,) + z.shape, dtype=np.float64) + qd = np.empty_like(qn) + if (z.ndim == 0): + _lqn(z, out=(qn, qd)) + else: + # new axes must be last for the ufunc + _lqn(z, + out=(np.moveaxis(qn, 0, -1), + np.moveaxis(qd, 0, -1))) + + return qn[:(n+1)], qd[:(n+1)] + + +def ai_zeros(nt): + """ + Compute `nt` zeros and values of the Airy function Ai and its derivative. + + Computes the first `nt` zeros, `a`, of the Airy function Ai(x); + first `nt` zeros, `ap`, of the derivative of the Airy function Ai'(x); + the corresponding values Ai(a'); + and the corresponding values Ai'(a). + + Parameters + ---------- + nt : int + Number of zeros to compute + + Returns + ------- + a : ndarray + First `nt` zeros of Ai(x) + ap : ndarray + First `nt` zeros of Ai'(x) + ai : ndarray + Values of Ai(x) evaluated at first `nt` zeros of Ai'(x) + aip : ndarray + Values of Ai'(x) evaluated at first `nt` zeros of Ai(x) + + References + ---------- + .. [1] Zhang, Shanjie and Jin, Jianming. "Computation of Special + Functions", John Wiley and Sons, 1996. + https://people.sc.fsu.edu/~jburkardt/f77_src/special_functions/special_functions.html + + Examples + -------- + >>> from scipy import special + >>> a, ap, ai, aip = special.ai_zeros(3) + >>> a + array([-2.33810741, -4.08794944, -5.52055983]) + >>> ap + array([-1.01879297, -3.24819758, -4.82009921]) + >>> ai + array([ 0.53565666, -0.41901548, 0.38040647]) + >>> aip + array([ 0.70121082, -0.80311137, 0.86520403]) + + """ + kf = 1 + if not isscalar(nt) or (floor(nt) != nt) or (nt <= 0): + raise ValueError("nt must be a positive integer scalar.") + return _specfun.airyzo(nt, kf) + + +def bi_zeros(nt): + """ + Compute `nt` zeros and values of the Airy function Bi and its derivative. + + Computes the first `nt` zeros, b, of the Airy function Bi(x); + first `nt` zeros, b', of the derivative of the Airy function Bi'(x); + the corresponding values Bi(b'); + and the corresponding values Bi'(b). + + Parameters + ---------- + nt : int + Number of zeros to compute + + Returns + ------- + b : ndarray + First `nt` zeros of Bi(x) + bp : ndarray + First `nt` zeros of Bi'(x) + bi : ndarray + Values of Bi(x) evaluated at first `nt` zeros of Bi'(x) + bip : ndarray + Values of Bi'(x) evaluated at first `nt` zeros of Bi(x) + + References + ---------- + .. [1] Zhang, Shanjie and Jin, Jianming. "Computation of Special + Functions", John Wiley and Sons, 1996. + https://people.sc.fsu.edu/~jburkardt/f77_src/special_functions/special_functions.html + + Examples + -------- + >>> from scipy import special + >>> b, bp, bi, bip = special.bi_zeros(3) + >>> b + array([-1.17371322, -3.2710933 , -4.83073784]) + >>> bp + array([-2.29443968, -4.07315509, -5.51239573]) + >>> bi + array([-0.45494438, 0.39652284, -0.36796916]) + >>> bip + array([ 0.60195789, -0.76031014, 0.83699101]) + + """ + kf = 2 + if not isscalar(nt) or (floor(nt) != nt) or (nt <= 0): + raise ValueError("nt must be a positive integer scalar.") + return _specfun.airyzo(nt, kf) + + +def lmbda(v, x): + r"""Jahnke-Emden Lambda function, Lambdav(x). + + This function is defined as [2]_, + + .. math:: \Lambda_v(x) = \Gamma(v+1) \frac{J_v(x)}{(x/2)^v}, + + where :math:`\Gamma` is the gamma function and :math:`J_v` is the + Bessel function of the first kind. + + Parameters + ---------- + v : float + Order of the Lambda function + x : float + Value at which to evaluate the function and derivatives + + Returns + ------- + vl : ndarray + Values of Lambda_vi(x), for vi=v-int(v), vi=1+v-int(v), ..., vi=v. + dl : ndarray + Derivatives Lambda_vi'(x), for vi=v-int(v), vi=1+v-int(v), ..., vi=v. + + References + ---------- + .. [1] Zhang, Shanjie and Jin, Jianming. "Computation of Special + Functions", John Wiley and Sons, 1996. + https://people.sc.fsu.edu/~jburkardt/f77_src/special_functions/special_functions.html + .. [2] Jahnke, E. and Emde, F. "Tables of Functions with Formulae and + Curves" (4th ed.), Dover, 1945 + """ + if not (isscalar(v) and isscalar(x)): + raise ValueError("arguments must be scalars.") + if (v < 0): + raise ValueError("argument must be > 0.") + n = int(v) + v0 = v - n + if (n < 1): + n1 = 1 + else: + n1 = n + v1 = n1 + v0 + if (v != floor(v)): + vm, vl, dl = _specfun.lamv(v1, x) + else: + vm, vl, dl = _specfun.lamn(v1, x) + return vl[:(n+1)], dl[:(n+1)] + + +def pbdv_seq(v, x): + """Parabolic cylinder functions Dv(x) and derivatives. + + Parameters + ---------- + v : float + Order of the parabolic cylinder function + x : float + Value at which to evaluate the function and derivatives + + Returns + ------- + dv : ndarray + Values of D_vi(x), for vi=v-int(v), vi=1+v-int(v), ..., vi=v. + dp : ndarray + Derivatives D_vi'(x), for vi=v-int(v), vi=1+v-int(v), ..., vi=v. + + References + ---------- + .. [1] Zhang, Shanjie and Jin, Jianming. "Computation of Special + Functions", John Wiley and Sons, 1996, chapter 13. + https://people.sc.fsu.edu/~jburkardt/f77_src/special_functions/special_functions.html + + """ + if not (isscalar(v) and isscalar(x)): + raise ValueError("arguments must be scalars.") + n = int(v) + v0 = v-n + if (n < 1): + n1 = 1 + else: + n1 = n + v1 = n1 + v0 + dv, dp, pdf, pdd = _specfun.pbdv(v1, x) + return dv[:n1+1], dp[:n1+1] + + +def pbvv_seq(v, x): + """Parabolic cylinder functions Vv(x) and derivatives. + + Parameters + ---------- + v : float + Order of the parabolic cylinder function + x : float + Value at which to evaluate the function and derivatives + + Returns + ------- + dv : ndarray + Values of V_vi(x), for vi=v-int(v), vi=1+v-int(v), ..., vi=v. + dp : ndarray + Derivatives V_vi'(x), for vi=v-int(v), vi=1+v-int(v), ..., vi=v. + + References + ---------- + .. [1] Zhang, Shanjie and Jin, Jianming. "Computation of Special + Functions", John Wiley and Sons, 1996, chapter 13. + https://people.sc.fsu.edu/~jburkardt/f77_src/special_functions/special_functions.html + + """ + if not (isscalar(v) and isscalar(x)): + raise ValueError("arguments must be scalars.") + n = int(v) + v0 = v-n + if (n <= 1): + n1 = 1 + else: + n1 = n + v1 = n1 + v0 + dv, dp, pdf, pdd = _specfun.pbvv(v1, x) + return dv[:n1+1], dp[:n1+1] + + +def pbdn_seq(n, z): + """Parabolic cylinder functions Dn(z) and derivatives. + + Parameters + ---------- + n : int + Order of the parabolic cylinder function + z : complex + Value at which to evaluate the function and derivatives + + Returns + ------- + dv : ndarray + Values of D_i(z), for i=0, ..., i=n. + dp : ndarray + Derivatives D_i'(z), for i=0, ..., i=n. + + References + ---------- + .. [1] Zhang, Shanjie and Jin, Jianming. "Computation of Special + Functions", John Wiley and Sons, 1996, chapter 13. + https://people.sc.fsu.edu/~jburkardt/f77_src/special_functions/special_functions.html + + """ + if not (isscalar(n) and isscalar(z)): + raise ValueError("arguments must be scalars.") + if (floor(n) != n): + raise ValueError("n must be an integer.") + if (abs(n) <= 1): + n1 = 1 + else: + n1 = n + cpb, cpd = _specfun.cpbdn(n1, z) + return cpb[:n1+1], cpd[:n1+1] + + +def ber_zeros(nt): + """Compute nt zeros of the Kelvin function ber. + + Parameters + ---------- + nt : int + Number of zeros to compute. Must be positive. + + Returns + ------- + ndarray + First `nt` zeros of the Kelvin function. + + See Also + -------- + ber + + References + ---------- + .. [1] Zhang, Shanjie and Jin, Jianming. "Computation of Special + Functions", John Wiley and Sons, 1996. + https://people.sc.fsu.edu/~jburkardt/f77_src/special_functions/special_functions.html + + """ + if not isscalar(nt) or (floor(nt) != nt) or (nt <= 0): + raise ValueError("nt must be positive integer scalar.") + return _specfun.klvnzo(nt, 1) + + +def bei_zeros(nt): + """Compute nt zeros of the Kelvin function bei. + + Parameters + ---------- + nt : int + Number of zeros to compute. Must be positive. + + Returns + ------- + ndarray + First `nt` zeros of the Kelvin function. + + See Also + -------- + bei + + References + ---------- + .. [1] Zhang, Shanjie and Jin, Jianming. "Computation of Special + Functions", John Wiley and Sons, 1996. + https://people.sc.fsu.edu/~jburkardt/f77_src/special_functions/special_functions.html + + """ + if not isscalar(nt) or (floor(nt) != nt) or (nt <= 0): + raise ValueError("nt must be positive integer scalar.") + return _specfun.klvnzo(nt, 2) + + +def ker_zeros(nt): + """Compute nt zeros of the Kelvin function ker. + + Parameters + ---------- + nt : int + Number of zeros to compute. Must be positive. + + Returns + ------- + ndarray + First `nt` zeros of the Kelvin function. + + See Also + -------- + ker + + References + ---------- + .. [1] Zhang, Shanjie and Jin, Jianming. "Computation of Special + Functions", John Wiley and Sons, 1996. + https://people.sc.fsu.edu/~jburkardt/f77_src/special_functions/special_functions.html + + """ + if not isscalar(nt) or (floor(nt) != nt) or (nt <= 0): + raise ValueError("nt must be positive integer scalar.") + return _specfun.klvnzo(nt, 3) + + +def kei_zeros(nt): + """Compute nt zeros of the Kelvin function kei. + + Parameters + ---------- + nt : int + Number of zeros to compute. Must be positive. + + Returns + ------- + ndarray + First `nt` zeros of the Kelvin function. + + See Also + -------- + kei + + References + ---------- + .. [1] Zhang, Shanjie and Jin, Jianming. "Computation of Special + Functions", John Wiley and Sons, 1996. + https://people.sc.fsu.edu/~jburkardt/f77_src/special_functions/special_functions.html + + """ + if not isscalar(nt) or (floor(nt) != nt) or (nt <= 0): + raise ValueError("nt must be positive integer scalar.") + return _specfun.klvnzo(nt, 4) + + +def berp_zeros(nt): + """Compute nt zeros of the derivative of the Kelvin function ber. + + Parameters + ---------- + nt : int + Number of zeros to compute. Must be positive. + + Returns + ------- + ndarray + First `nt` zeros of the derivative of the Kelvin function. + + See Also + -------- + ber, berp + + References + ---------- + .. [1] Zhang, Shanjie and Jin, Jianming. "Computation of Special + Functions", John Wiley and Sons, 1996. + https://people.sc.fsu.edu/~jburkardt/f77_src/special_functions/special_functions.html + + + Examples + -------- + Compute the first 5 zeros of the derivative of the Kelvin function. + + >>> from scipy.special import berp_zeros + >>> berp_zeros(5) + array([ 6.03871081, 10.51364251, 14.96844542, 19.41757493, 23.86430432]) + + """ + if not isscalar(nt) or (floor(nt) != nt) or (nt <= 0): + raise ValueError("nt must be positive integer scalar.") + return _specfun.klvnzo(nt, 5) + + +def beip_zeros(nt): + """Compute nt zeros of the derivative of the Kelvin function bei. + + Parameters + ---------- + nt : int + Number of zeros to compute. Must be positive. + + Returns + ------- + ndarray + First `nt` zeros of the derivative of the Kelvin function. + + See Also + -------- + bei, beip + + References + ---------- + .. [1] Zhang, Shanjie and Jin, Jianming. "Computation of Special + Functions", John Wiley and Sons, 1996. + https://people.sc.fsu.edu/~jburkardt/f77_src/special_functions/special_functions.html + + """ + if not isscalar(nt) or (floor(nt) != nt) or (nt <= 0): + raise ValueError("nt must be positive integer scalar.") + return _specfun.klvnzo(nt, 6) + + +def kerp_zeros(nt): + """Compute nt zeros of the derivative of the Kelvin function ker. + + Parameters + ---------- + nt : int + Number of zeros to compute. Must be positive. + + Returns + ------- + ndarray + First `nt` zeros of the derivative of the Kelvin function. + + See Also + -------- + ker, kerp + + References + ---------- + .. [1] Zhang, Shanjie and Jin, Jianming. "Computation of Special + Functions", John Wiley and Sons, 1996. + https://people.sc.fsu.edu/~jburkardt/f77_src/special_functions/special_functions.html + + """ + if not isscalar(nt) or (floor(nt) != nt) or (nt <= 0): + raise ValueError("nt must be positive integer scalar.") + return _specfun.klvnzo(nt, 7) + + +def keip_zeros(nt): + """Compute nt zeros of the derivative of the Kelvin function kei. + + Parameters + ---------- + nt : int + Number of zeros to compute. Must be positive. + + Returns + ------- + ndarray + First `nt` zeros of the derivative of the Kelvin function. + + See Also + -------- + kei, keip + + References + ---------- + .. [1] Zhang, Shanjie and Jin, Jianming. "Computation of Special + Functions", John Wiley and Sons, 1996. + https://people.sc.fsu.edu/~jburkardt/f77_src/special_functions/special_functions.html + + """ + if not isscalar(nt) or (floor(nt) != nt) or (nt <= 0): + raise ValueError("nt must be positive integer scalar.") + return _specfun.klvnzo(nt, 8) + + +def kelvin_zeros(nt): + """Compute `nt` zeros of all Kelvin functions. + + Parameters + ---------- + nt : int + Number of zeros to compute for each function. + + Returns + ------- + zeros : tuple of arrays + Length-8 tuple of arrays of length `nt`. The tuple contains the arrays of zeros + of ``(ber, bei, ker, kei, ber', bei', ker', kei')``. + + References + ---------- + .. [1] Zhang, Shanjie and Jin, Jianming. "Computation of Special + Functions", John Wiley and Sons, 1996. + https://people.sc.fsu.edu/~jburkardt/f77_src/special_functions/special_functions.html + + """ + if not isscalar(nt) or (floor(nt) != nt) or (nt <= 0): + raise ValueError("nt must be positive integer scalar.") + return (_specfun.klvnzo(nt, 1), + _specfun.klvnzo(nt, 2), + _specfun.klvnzo(nt, 3), + _specfun.klvnzo(nt, 4), + _specfun.klvnzo(nt, 5), + _specfun.klvnzo(nt, 6), + _specfun.klvnzo(nt, 7), + _specfun.klvnzo(nt, 8)) + + +def pro_cv_seq(m, n, c): + """Characteristic values for prolate spheroidal wave functions. + + Compute a sequence of characteristic values for the prolate + spheroidal wave functions for mode `m` and n'=m..n and spheroidal + parameter `c`. + + Parameters + ---------- + m, n : int + Non-negative mode parameters. + c : float + Spheroidal parameter. + + Returns + ------- + cv : array of floats + Characteristic values. + + References + ---------- + .. [1] Zhang, Shanjie and Jin, Jianming. "Computation of Special + Functions", John Wiley and Sons, 1996. + https://people.sc.fsu.edu/~jburkardt/f77_src/special_functions/special_functions.html + + """ + if not (isscalar(m) and isscalar(n) and isscalar(c)): + raise ValueError("Arguments must be scalars.") + if (n != floor(n)) or (m != floor(m)): + raise ValueError("Modes must be integers.") + if (n-m > 199): + raise ValueError("Difference between n and m is too large.") + maxL = n-m+1 + return _specfun.segv(m, n, c, 1)[1][:maxL] + + +def obl_cv_seq(m, n, c): + """Characteristic values for oblate spheroidal wave functions. + + Compute a sequence of characteristic values for the oblate + spheroidal wave functions for mode m and n'=m..n and spheroidal + parameter c. + + Parameters + ---------- + m, n : int + Non-negative mode parameters. + c : float + Spheroidal parameter. + + Returns + ------- + cv : array of floats + Characteristic values. + + References + ---------- + .. [1] Zhang, Shanjie and Jin, Jianming. "Computation of Special + Functions", John Wiley and Sons, 1996. + https://people.sc.fsu.edu/~jburkardt/f77_src/special_functions/special_functions.html + + """ + if not (isscalar(m) and isscalar(n) and isscalar(c)): + raise ValueError("Arguments must be scalars.") + if (n != floor(n)) or (m != floor(m)): + raise ValueError("Modes must be integers.") + if (n-m > 199): + raise ValueError("Difference between n and m is too large.") + maxL = n-m+1 + return _specfun.segv(m, n, c, -1)[1][:maxL] + + +def comb(N, k, *, exact=False, repetition=False): + """The number of combinations of N things taken k at a time. + + This is often expressed as "N choose k". + + Parameters + ---------- + N : int, ndarray + Number of things. + k : int, ndarray + Number of elements taken. + exact : bool, optional + For integers, if `exact` is False, then floating point precision is + used, otherwise the result is computed exactly. + repetition : bool, optional + If `repetition` is True, then the number of combinations with + repetition is computed. + + Returns + ------- + val : int, float, ndarray + The total number of combinations. + + See Also + -------- + binom : Binomial coefficient considered as a function of two real + variables. + + Notes + ----- + - Array arguments accepted only for exact=False case. + - If N < 0, or k < 0, then 0 is returned. + - If k > N and repetition=False, then 0 is returned. + + Examples + -------- + >>> import numpy as np + >>> from scipy.special import comb + >>> k = np.array([3, 4]) + >>> n = np.array([10, 10]) + >>> comb(n, k, exact=False) + array([ 120., 210.]) + >>> comb(10, 3, exact=True) + 120 + >>> comb(10, 3, exact=True, repetition=True) + 220 + + """ + if repetition: + # Special case: C(n, 0) with repetition = 1 for n >= 0 + # Without this check, comb(0, 0, repetition=True) would compute + # comb(-1, 0) which incorrectly returns 0 + if exact: + if k == 0 and int(N) == N and N >= 0: + return 1 + else: + k, N = asarray(k), asarray(N) + cond = (k == 0) & (N >= 0) + vals = binom(N + k - 1, k) + if isinstance(vals, np.ndarray): + vals[cond] = 1.0 + elif cond: + vals = np.float64(1.0) + return vals + return comb(N + k - 1, k, exact=exact) + if exact: + if int(N) == N and int(k) == k: + # _comb_int casts inputs to integers, which is safe & intended here + return _comb_int(N, k) + else: + raise ValueError("Non-integer `N` and `k` with `exact=True` is not " + "supported.") + else: + k, N = asarray(k), asarray(N) + cond = (k <= N) & (N >= 0) & (k >= 0) + vals = binom(N, k) + if isinstance(vals, np.ndarray): + vals[~cond] = 0 + elif not cond: + vals = np.float64(0) + return vals + + +def perm(N, k, exact=False): + """Permutations of N things taken k at a time, i.e., k-permutations of N. + + It's also known as "partial permutations". + + Parameters + ---------- + N : int, ndarray + Number of things. + k : int, ndarray + Number of elements taken. + exact : bool, optional + If ``True``, calculate the answer exactly using long integer arithmetic (`N` + and `k` must be scalar integers). If ``False``, a floating point approximation + is calculated (more rapidly) using `poch`. Default is ``False``. + + Returns + ------- + val : int, ndarray + The number of k-permutations of N. + + Notes + ----- + - Array arguments accepted only for exact=False case. + - If k > N, N < 0, or k < 0, then a 0 is returned. + + Examples + -------- + >>> import numpy as np + >>> from scipy.special import perm + >>> k = np.array([3, 4]) + >>> n = np.array([10, 10]) + >>> perm(n, k) + array([ 720., 5040.]) + >>> perm(10, 3, exact=True) + 720 + + """ + if exact: + N = np.squeeze(N)[()] # for backward compatibility (accepted size 1 arrays) + k = np.squeeze(k)[()] + if not (isscalar(N) and isscalar(k)): + raise ValueError("`N` and `k` must be scalar integers with `exact=True`.") + + floor_N, floor_k = int(N), int(k) + non_integral = not (floor_N == N and floor_k == k) + if non_integral: + raise ValueError("Non-integer `N` and `k` with `exact=True` is not " + "supported.") + + if (k > N) or (N < 0) or (k < 0): + return 0 + + val = 1 + for i in range(floor_N - floor_k + 1, floor_N + 1): + val *= i + return val + else: + k, N = asarray(k), asarray(N) + cond = (k <= N) & (N >= 0) & (k >= 0) + vals = poch(N - k + 1, k) + if isinstance(vals, np.ndarray): + vals[~cond] = 0 + elif not cond: + vals = np.float64(0) + return vals + + +# https://stackoverflow.com/a/16327037 +def _range_prod(lo, hi, k=1): + """ + Product of a range of numbers spaced k apart (from hi). + + For k=1, this returns the product of + lo * (lo+1) * (lo+2) * ... * (hi-2) * (hi-1) * hi + = hi! / (lo-1)! + + For k>1, it correspond to taking only every k'th number when + counting down from hi - e.g. 18!!!! = _range_prod(1, 18, 4). + + Breaks into smaller products first for speed: + _range_prod(2, 9) = ((2*3)*(4*5))*((6*7)*(8*9)) + """ + if lo == 1 and k == 1: + return math.factorial(hi) + + if lo + k < hi: + mid = (hi + lo) // 2 + if k > 1: + # make sure mid is a multiple of k away from hi + mid = mid - ((mid - hi) % k) + return _range_prod(lo, mid, k) * _range_prod(mid + k, hi, k) + elif lo + k == hi: + return lo * hi + else: + return hi + + +def _factorialx_array_exact(n, k=1): + """ + Exact computation of factorial for an array. + + The factorials are computed in incremental fashion, by taking + the sorted unique values of n and multiplying the intervening + numbers between the different unique values. + + In other words, the factorial for the largest input is only + computed once, with each other result computed in the process. + + k > 1 corresponds to the multifactorial. + """ + un = np.unique(n) + + # Convert to object array if np.int64 can't handle size + if k in _FACTORIALK_LIMITS_64BITS.keys(): + if un[-1] > _FACTORIALK_LIMITS_64BITS[k]: + # e.g. k=1: 21! > np.iinfo(np.int64).max + dt = object + elif un[-1] > _FACTORIALK_LIMITS_32BITS[k]: + # e.g. k=3: 26!!! > np.iinfo(np.int32).max + dt = np.int64 + else: + dt = np.dtype("long") + else: + # for k >= 10, we always use object + dt = object + + out = np.empty_like(n, dtype=dt) + + # Handle invalid/trivial values + un = un[un > 1] + out[n < 2] = 1 + out[n < 0] = 0 + + # Calculate products of each range of numbers + # we can only multiply incrementally if the values are k apart; + # therefore we partition `un` into "lanes", i.e. its residues modulo k + for lane in range(0, k): + ul = un[(un % k) == lane] if k > 1 else un + if ul.size: + # after np.unique, un resp. ul are sorted, ul[0] is the smallest; + # cast to python ints to avoid overflow with np.int-types + val = _range_prod(1, int(ul[0]), k=k) + out[n == ul[0]] = val + for i in range(len(ul) - 1): + # by the filtering above, we have ensured that prev & current + # are a multiple of k apart + prev = ul[i] + current = ul[i + 1] + # we already multiplied all factors until prev; continue + # building the full factorial from the following (`prev + 1`); + # use int() for the same reason as above + val *= _range_prod(int(prev + 1), int(current), k=k) + out[n == current] = val + + return out + + +def _factorialx_array_approx(n, k, extend): + """ + Calculate approximation to multifactorial for array n and integer k. + + Ensure that values aren't calculated unnecessarily. + """ + if extend == "complex": + return _factorialx_approx_core(n, k=k, extend=extend) + + # at this point we are guaranteed that extend='zero' and that k>0 is an integer + result = zeros(n.shape) + # keep nans as nans + place(result, np.isnan(n), np.nan) + # only compute where n >= 0 (excludes nans), everything else is 0 + cond = (n >= 0) + n_to_compute = extract(cond, n) + place(result, cond, _factorialx_approx_core(n_to_compute, k=k, extend=extend)) + return result + + +def _gamma1p(vals): + """ + returns gamma(n+1), though with NaN at -1 instead of inf, c.f. #21827 + """ + res = gamma(vals + 1) + # replace infinities at -1 (from gamma function at 0) with nan + # gamma only returns inf for real inputs; can ignore complex case + if isinstance(res, np.ndarray): + if not _is_subdtype(vals.dtype, "c"): + res[vals == -1] = np.nan + elif np.isinf(res) and vals == -1: + res = np.float64("nan") + return res + + +def _factorialx_approx_core(n, k, extend): + """ + Core approximation to multifactorial for array n and integer k. + """ + if k == 1: + # shortcut for k=1; same for both extensions, because we assume the + # handling of extend == 'zero' happens in _factorialx_array_approx + result = _gamma1p(n) + if isinstance(n, np.ndarray): + # gamma does not maintain 0-dim arrays; fix it + result = np.array(result) + return result + + if extend == "complex": + # see https://numpy.org/doc/stable/reference/generated/numpy.power.html + p_dtype = complex if (_is_subdtype(type(k), "c") or k < 0) else None + with warnings.catch_warnings(): + # do not warn about 0 * inf, nan / nan etc.; the results are correct + warnings.simplefilter("ignore", RuntimeWarning) + # don't use `(n-1)/k` in np.power; underflows if 0 is of a uintX type + result = np.power(k, n / k, dtype=p_dtype) * _gamma1p(n / k) + result *= rgamma(1 / k + 1) / np.power(k, 1 / k, dtype=p_dtype) + if isinstance(n, np.ndarray): + # ensure we keep array-ness for 0-dim inputs; already n/k above loses it + result = np.array(result) + return result + + # at this point we are guaranteed that extend='zero' and that k>0 is an integer + n_mod_k = n % k + # scalar case separately, unified handling would be inefficient for arrays; + # don't use isscalar due to numpy/numpy#23574; 0-dim arrays treated below + if not isinstance(n, np.ndarray): + with warnings.catch_warnings(): + # large n cause overflow warnings, but infinity is fine + warnings.simplefilter("ignore", RuntimeWarning) + return ( + np.power(k, (n - n_mod_k) / k) + * gamma(n / k + 1) / gamma(n_mod_k / k + 1) + * max(n_mod_k, 1) + ) + + # factor that's independent of the residue class (see factorialk docstring) + with warnings.catch_warnings(): + # large n cause overflow warnings, but infinity is fine + warnings.simplefilter("ignore", RuntimeWarning) + result = np.power(k, n / k) * gamma(n / k + 1) + # factor dependent on residue r (for `r=0` it's 1, so we skip `r=0` + # below and thus also avoid evaluating `max(r, 1)`) + def corr(k, r): return np.power(k, -r / k) / gamma(r / k + 1) * r + for r in np.unique(n_mod_k): + if r == 0: + continue + # cast to int because uint types break on `-r` + result[n_mod_k == r] *= corr(k, int(r)) + return result + + +def _is_subdtype(dtype, dtypes): + """ + Shorthand for calculating whether dtype is subtype of some dtypes. + + Also allows specifying a list instead of just a single dtype. + + Additionaly, the most important supertypes from + https://numpy.org/doc/stable/reference/arrays.scalars.html + can optionally be specified using abbreviations as follows: + "i": np.integer + "f": np.floating + "c": np.complexfloating + "n": np.number (contains the other three) + """ + dtypes = dtypes if isinstance(dtypes, list) else [dtypes] + # map single character abbreviations, if they are in dtypes + mapping = { + "i": np.integer, + "f": np.floating, + "c": np.complexfloating, + "n": np.number + } + dtypes = [mapping.get(x, x) for x in dtypes] + return any(np.issubdtype(dtype, dt) for dt in dtypes) + + +def _factorialx_wrapper(fname, n, k, exact, extend): + """ + Shared implementation for factorial, factorial2 & factorialk. + """ + if extend not in ("zero", "complex"): + raise ValueError( + f"argument `extend` must be either 'zero' or 'complex', received: {extend}" + ) + if exact and extend == "complex": + raise ValueError("Incompatible options: `exact=True` and `extend='complex'`") + + msg_unsup = ( + "Unsupported data type for {vname} in {fname}: {dtype}\n" + ) + if fname == "factorial": + msg_unsup += ( + "Permitted data types are integers and floating point numbers, " + "as well as complex numbers if `extend='complex' is passed." + ) + else: + msg_unsup += ( + "Permitted data types are integers, as well as floating point " + "numbers and complex numbers if `extend='complex' is passed." + ) + msg_exact_not_possible = ( + "`exact=True` only supports integers, cannot use data type {dtype}" + ) + msg_needs_complex = ( + "In order to use non-integer arguments, you must opt into this by passing " + "`extend='complex'`. Note that this changes the result for all negative " + "arguments (which by default return 0)." + ) + + if fname == "factorial2": + msg_needs_complex += (" Additionally, it will rescale the values of the double" + " factorial at even integers by a factor of sqrt(2/pi).") + elif fname == "factorialk": + msg_needs_complex += (" Additionally, it will perturb the values of the" + " multifactorial at most positive integers `n`.") + # check type of k + if not _is_subdtype(type(k), ["i", "f", "c"]): + raise ValueError(msg_unsup.format(vname="`k`", fname=fname, dtype=type(k))) + elif _is_subdtype(type(k), ["f", "c"]) and extend != "complex": + raise ValueError(msg_needs_complex) + # check value of k + if extend == "zero" and k < 1: + msg = f"For `extend='zero'`, k must be a positive integer, received: {k}" + raise ValueError(msg) + elif k == 0: + raise ValueError("Parameter k cannot be zero!") + + # factorial allows floats also for extend="zero" + types_requiring_complex = "c" if fname == "factorial" else ["f", "c"] + + # don't use isscalar due to numpy/numpy#23574; 0-dim arrays treated below + if np.ndim(n) == 0 and not isinstance(n, np.ndarray): + # scalar cases + if not _is_subdtype(type(n), ["i", "f", "c", type(None)]): + raise ValueError(msg_unsup.format(vname="`n`", fname=fname, dtype=type(n))) + elif _is_subdtype(type(n), types_requiring_complex) and extend != "complex": + raise ValueError(msg_needs_complex) + elif n is None or np.isnan(n): + complexify = (extend == "complex") and _is_subdtype(type(n), "c") + return np.complex128("nan+nanj") if complexify else np.float64("nan") + elif extend == "zero" and n < 0: + return 0 if exact else np.float64(0) + elif n in {0, 1}: + return 1 if exact else np.float64(1) + elif exact and _is_subdtype(type(n), "i"): + # calculate with integers; cast away other int types (like unsigned) + return _range_prod(1, int(n), k=k) + elif exact: + # only relevant for factorial + raise ValueError(msg_exact_not_possible.format(dtype=type(n))) + # approximation + return _factorialx_approx_core(n, k=k, extend=extend) + + # arrays & array-likes + n = asarray(n) + + if not _is_subdtype(n.dtype, ["i", "f", "c"]): + raise ValueError(msg_unsup.format(vname="`n`", fname=fname, dtype=n.dtype)) + elif _is_subdtype(n.dtype, types_requiring_complex) and extend != "complex": + raise ValueError(msg_needs_complex) + elif exact and _is_subdtype(n.dtype, ["f"]): + # only relevant for factorial + raise ValueError(msg_exact_not_possible.format(dtype=n.dtype)) + + if n.size == 0: + # return empty arrays unchanged + return n + elif exact: + # calculate with integers + return _factorialx_array_exact(n, k=k) + # approximation + return _factorialx_array_approx(n, k=k, extend=extend) + + +def factorial(n, exact=False, extend="zero"): + """ + The factorial of a number or array of numbers. + + The factorial of non-negative integer `n` is the product of all + positive integers less than or equal to `n`:: + + n! = n * (n - 1) * (n - 2) * ... * 1 + + Parameters + ---------- + n : int or float or complex (or array_like thereof) + Input values for ``n!``. Complex values require ``extend='complex'``. + By default, the return value for ``n < 0`` is 0. + exact : bool, optional + If ``exact`` is set to True, calculate the answer exactly using + integer arithmetic, otherwise approximate using the gamma function + (faster, but yields floats instead of integers). + Default is False. + extend : str, optional + One of ``'zero'`` or ``'complex'``; this determines how values ``n<0`` + are handled - by default they are 0, but it is possible to opt into the + complex extension of the factorial (see below). + + Returns + ------- + nf : int or float or complex or ndarray + Factorial of ``n``, as integer, float or complex (depending on ``exact`` + and ``extend``). Array inputs are returned as arrays. + + Notes + ----- + For arrays with ``exact=True``, the factorial is computed only once, for + the largest input, with each other result computed in the process. + The output dtype is increased to ``int64`` or ``object`` if necessary. + + With ``exact=False`` the factorial is approximated using the gamma + function (which is also the definition of the complex extension): + + .. math:: n! = \\Gamma(n+1) + + Examples + -------- + >>> import numpy as np + >>> from scipy.special import factorial + >>> arr = np.array([3, 4, 5]) + >>> factorial(arr, exact=False) + array([ 6., 24., 120.]) + >>> factorial(arr, exact=True) + array([ 6, 24, 120]) + >>> factorial(5, exact=True) + 120 + + """ + return _factorialx_wrapper("factorial", n, k=1, exact=exact, extend=extend) + + +def factorial2(n, exact=False, extend="zero"): + """ + Double factorial. + + This is the factorial with every second value skipped. E.g., ``7!! = 7 * 5 + * 3 * 1``. It can be approximated numerically as:: + + n!! = 2 ** (n / 2) * gamma(n / 2 + 1) * sqrt(2 / pi) n odd + = 2 ** (n / 2) * gamma(n / 2 + 1) n even + = 2 ** (n / 2) * (n / 2)! n even + + The formula for odd ``n`` is the basis for the complex extension. + + Parameters + ---------- + n : int or float or complex (or array_like thereof) + Input values for ``n!!``. Non-integer values require ``extend='complex'``. + By default, the return value for ``n < 0`` is 0. + exact : bool, optional + If ``exact`` is set to True, calculate the answer exactly using + integer arithmetic, otherwise use above approximation (faster, + but yields floats instead of integers). + Default is False. + extend : str, optional + One of ``'zero'`` or ``'complex'``; this determines how values ``n<0`` + are handled - by default they are 0, but it is possible to opt into the + complex extension of the double factorial. This also enables passing + complex values to ``n``. + + .. warning:: + + Using the ``'complex'`` extension also changes the values of the + double factorial for even integers, reducing them by a factor of + ``sqrt(2/pi) ~= 0.79``, see [1]. + + Returns + ------- + nf : int or float or complex or ndarray + Double factorial of ``n``, as integer, float or complex (depending on + ``exact`` and ``extend``). Array inputs are returned as arrays. + + References + ---------- + .. [1] Complex extension to double factorial + https://en.wikipedia.org/wiki/Double_factorial#Complex_arguments + + Examples + -------- + >>> from scipy.special import factorial2 + >>> factorial2(7, exact=False) + np.float64(105.00000000000001) + >>> factorial2(7, exact=True) + 105 + """ + return _factorialx_wrapper("factorial2", n, k=2, exact=exact, extend=extend) + + +def factorialk(n, k, exact=False, extend="zero"): + """ + Multifactorial of n of order k, n(!!...!). + + This is the multifactorial of n skipping k values. For example, + + factorialk(17, 4) = 17!!!! = 17 * 13 * 9 * 5 * 1 + + In particular, for any integer ``n``, we have + + factorialk(n, 1) = factorial(n) + + factorialk(n, 2) = factorial2(n) + + Parameters + ---------- + n : int or float or complex (or array_like thereof) + Input values for multifactorial. Non-integer values require + ``extend='complex'``. By default, the return value for ``n < 0`` is 0. + k : int or float or complex (or array_like thereof) + Order of multifactorial. Non-integer values require ``extend='complex'``. + exact : bool, optional + If ``exact`` is set to True, calculate the answer exactly using + integer arithmetic, otherwise use an approximation (faster, + but yields floats instead of integers) + Default is False. + extend : str, optional + One of ``'zero'`` or ``'complex'``; this determines how values ``n<0`` are + handled - by default they are 0, but it is possible to opt into the complex + extension of the multifactorial. This enables passing complex values, + not only to ``n`` but also to ``k``. + + .. warning:: + + Using the ``'complex'`` extension also changes the values of the + multifactorial at integers ``n != 1 (mod k)`` by a factor depending + on both ``k`` and ``n % k``, see below or [1]. + + Returns + ------- + nf : int or float or complex or ndarray + Multifactorial (order ``k``) of ``n``, as integer, float or complex (depending + on ``exact`` and ``extend``). Array inputs are returned as arrays. + + Notes + ----- + While less straight-forward than for the double-factorial, it's possible to + calculate a general approximation formula of n!(k) by studying ``n`` for a given + remainder ``r < k`` (thus ``n = m * k + r``, resp. ``r = n % k``), which can be + put together into something valid for all integer values ``n >= 0`` & ``k > 0``:: + + n!(k) = k ** ((n - r)/k) * gamma(n/k + 1) / gamma(r/k + 1) * max(r, 1) + + This is the basis of the approximation when ``exact=False``. + + In principle, any fixed choice of ``r`` (ignoring its relation ``r = n%k`` + to ``n``) would provide a suitable analytic continuation from integer ``n`` + to complex ``z`` (not only satisfying the functional equation but also + being logarithmically convex, c.f. Bohr-Mollerup theorem) -- in fact, the + choice of ``r`` above only changes the function by a constant factor. The + final constraint that determines the canonical continuation is ``f(1) = 1``, + which forces ``r = 1`` (see also [1]).:: + + z!(k) = k ** ((z - 1)/k) * gamma(z/k + 1) / gamma(1/k + 1) + + References + ---------- + .. [1] Complex extension to multifactorial + https://en.wikipedia.org/wiki/Double_factorial#Alternative_extension_of_the_multifactorial + + Examples + -------- + >>> from scipy.special import factorialk + >>> factorialk(5, k=1, exact=True) + 120 + >>> factorialk(5, k=3, exact=True) + 10 + >>> factorialk([5, 7, 9], k=3, exact=True) + array([ 10, 28, 162]) + >>> factorialk([5, 7, 9], k=3, exact=False) + array([ 10., 28., 162.]) + """ + return _factorialx_wrapper("factorialk", n, k=k, exact=exact, extend=extend) + + +def stirling2(N, K, *, exact=False): + r"""Generate Stirling number(s) of the second kind. + + Stirling numbers of the second kind count the number of ways to + partition a set with N elements into K non-empty subsets. + + The values this function returns are calculated using a dynamic + program which avoids redundant computation across the subproblems + in the solution. For array-like input, this implementation also + avoids redundant computation across the different Stirling number + calculations. + + The numbers are sometimes denoted + + .. math:: + + {N \brace{K}} + + see [1]_ for details. This is often expressed-verbally-as + "N subset K". + + Parameters + ---------- + N : int, ndarray + Number of things. + K : int, ndarray + Number of non-empty subsets taken. + exact : bool, optional + Uses dynamic programming (DP) with floating point + numbers for smaller arrays and uses a second order approximation due to + Temme for larger entries of `N` and `K` that allows trading speed for + accuracy. See [2]_ for a description. Temme approximation is used for + values ``n>50``. The max error from the DP has max relative error + ``4.5*10^-16`` for ``n<=50`` and the max error from the Temme approximation + has max relative error ``5*10^-5`` for ``51 <= n < 70`` and + ``9*10^-6`` for ``70 <= n < 101``. Note that these max relative errors will + decrease further as `n` increases. + + Returns + ------- + val : int, float, ndarray + The number of partitions. + + See Also + -------- + comb : The number of combinations of N things taken k at a time. + + Notes + ----- + - If N < 0, or K < 0, then 0 is returned. + - If K > N, then 0 is returned. + + The output type will always be `int` or ndarray of `object`. + The input must contain either numpy or python integers otherwise a + TypeError is raised. + + References + ---------- + .. [1] R. L. Graham, D. E. Knuth and O. Patashnik, "Concrete + Mathematics: A Foundation for Computer Science," Addison-Wesley + Publishing Company, Boston, 1989. Chapter 6, page 258. + + .. [2] Temme, Nico M. "Asymptotic estimates of Stirling numbers." + Studies in Applied Mathematics 89.3 (1993): 233-243. + + Examples + -------- + >>> import numpy as np + >>> from scipy.special import stirling2 + >>> k = np.array([3, -1, 3]) + >>> n = np.array([10, 10, 9]) + >>> stirling2(n, k) + array([9330.0, 0.0, 3025.0]) + + """ + output_is_scalar = np.isscalar(N) and np.isscalar(K) + # make a min-heap of unique (n,k) pairs + N, K = asarray(N), asarray(K) + if not np.issubdtype(N.dtype, np.integer): + raise TypeError("Argument `N` must contain only integers") + if not np.issubdtype(K.dtype, np.integer): + raise TypeError("Argument `K` must contain only integers") + if not exact: + # NOTE: here we allow np.uint via casting to double types prior to + # passing to private ufunc dispatcher. All dispatched functions + # take double type for (n,k) arguments and return double. + return _stirling2_inexact(N.astype(float), K.astype(float)) + nk_pairs = list( + set([(n.take(0), k.take(0)) + for n, k in np.nditer([N, K], ['refs_ok'])]) + ) + heapify(nk_pairs) + # base mapping for small values + snsk_vals = defaultdict(int) + for pair in [(0, 0), (1, 1), (2, 1), (2, 2)]: + snsk_vals[pair] = 1 + # for each pair in the min-heap, calculate the value, store for later + n_old, n_row = 2, [0, 1, 1] + while nk_pairs: + n, k = heappop(nk_pairs) + if n < 2 or k > n or k <= 0: + continue + elif k == n or k == 1: + snsk_vals[(n, k)] = 1 + continue + elif n != n_old: + num_iters = n - n_old + while num_iters > 0: + n_row.append(1) + # traverse from back to remove second row + for j in range(len(n_row)-2, 1, -1): + n_row[j] = n_row[j]*j + n_row[j-1] + num_iters -= 1 + snsk_vals[(n, k)] = n_row[k] + else: + snsk_vals[(n, k)] = n_row[k] + n_old, n_row = n, n_row + out_types = [object, object, object] if exact else [float, float, float] + # for each pair in the map, fetch the value, and populate the array + it = np.nditer( + [N, K, None], + ['buffered', 'refs_ok'], + [['readonly'], ['readonly'], ['writeonly', 'allocate']], + op_dtypes=out_types, + ) + with it: + while not it.finished: + it[2] = snsk_vals[(int(it[0]), int(it[1]))] + it.iternext() + output = it.operands[2] + # If N and K were both scalars, convert output to scalar. + if output_is_scalar: + output = output.take(0) + return output + + +def zeta(x, q=None, out=None): + r""" + Riemann or Hurwitz zeta function. + + Parameters + ---------- + x : array_like of float or complex + Input data + q : array_like of float, optional + Input data, must be real. Defaults to Riemann zeta. When `q` is + ``None``, complex inputs `x` are supported. If `q` is not ``None``, + then currently only real inputs `x` with ``x >= 1`` are supported, + even when ``q = 1.0`` (corresponding to the Riemann zeta function). + + out : ndarray, optional + Output array for the computed values. + + Returns + ------- + out : array_like + Values of zeta(x). + + See Also + -------- + zetac + + Notes + ----- + The two-argument version is the Hurwitz zeta function + + .. math:: + + \zeta(x, q) = \sum_{k=0}^{\infty} \frac{1}{(k + q)^x}; + + see [dlmf]_ for details. The Riemann zeta function corresponds to + the case when ``q = 1``. + + For complex inputs with ``q = None``, points with + ``abs(z.imag) > 1e9`` and ``0 <= abs(z.real) < 2.5`` are currently not + supported due to slow convergence causing excessive runtime. + + References + ---------- + .. [dlmf] NIST, Digital Library of Mathematical Functions, + https://dlmf.nist.gov/25.11#i + + Examples + -------- + >>> import numpy as np + >>> from scipy.special import zeta, polygamma, factorial + + Some specific values: + + >>> zeta(2), np.pi**2/6 + (1.6449340668482266, 1.6449340668482264) + + >>> zeta(4), np.pi**4/90 + (1.0823232337111381, 1.082323233711138) + + First nontrivial zero: + + >>> zeta(0.5 + 14.134725141734695j) + 0 + 0j + + Relation to the `polygamma` function: + + >>> m = 3 + >>> x = 1.25 + >>> polygamma(m, x) + array(2.782144009188397) + >>> (-1)**(m+1) * factorial(m) * zeta(m+1, x) + 2.7821440091883969 + + """ + if q is None: + return _ufuncs._riemann_zeta(x, out) + else: + return _ufuncs._zeta(x, q, out) + + +def softplus(x, **kwargs): + r""" + Compute the softplus function element-wise. + + The softplus function is defined as: ``softplus(x) = log(1 + exp(x))``. + It is a smooth approximation of the rectifier function (ReLU). + + Parameters + ---------- + x : array_like + Input value. + **kwargs + For other keyword-only arguments, see the + `ufunc docs `_. + + Returns + ------- + softplus : ndarray + Logarithm of ``exp(0) + exp(x)``. + + Examples + -------- + >>> from scipy import special + + >>> special.softplus(0) + 0.6931471805599453 + + >>> special.softplus([-1, 0, 1]) + array([0.31326169, 0.69314718, 1.31326169]) + """ + return np.logaddexp(0, x, **kwargs) diff --git a/v12x/pylib/scipy/special/_ellip_harm.py b/v12x/pylib/scipy/special/_ellip_harm.py new file mode 100644 index 0000000000000000000000000000000000000000..28ebd6cd4402674f4eed8bb9a790f733f3808927 --- /dev/null +++ b/v12x/pylib/scipy/special/_ellip_harm.py @@ -0,0 +1,214 @@ +import numpy as np + +from ._ufuncs import _ellip_harm +from ._ellip_harm_2 import _ellipsoid, _ellipsoid_norm + + +def ellip_harm(h2, k2, n, p, s, signm=1, signn=1): + r""" + Ellipsoidal harmonic functions E^p_n(l). + + These are also known as Lamé functions of the first kind, and are + solutions to the Lamé equation: + + .. math:: (s^2 - h^2)(s^2 - k^2)E''(s) + + s(2s^2 - h^2 - k^2)E'(s) + (a - q s^2)E(s) = 0 + + where :math:`q = (n+1)n` and :math:`a` is the eigenvalue (not + returned) corresponding to the solutions. + + Parameters + ---------- + h2 : float + ``h**2`` + k2 : float + ``k**2``; should be larger than ``h**2`` + n : int + Degree + p : int + Order, can range between ``[1, 2n+1]`` + s : float + Coordinate + signm : {1, -1}, optional + Sign of prefactor of functions. Can be +/-1. See Notes. + signn : {1, -1}, optional + Sign of prefactor of functions. Can be +/-1. See Notes. + + Returns + ------- + E : float + the harmonic :math:`E^p_n(s)` + + See Also + -------- + ellip_harm_2, ellip_normal + + Notes + ----- + The geometric interpretation of the ellipsoidal functions is + explained in [2]_, [3]_, [4]_. The `signm` and `signn` arguments control the + sign of prefactors for functions according to their type:: + + K : +1 + L : signm + M : signn + N : signm*signn + + .. versionadded:: 0.15.0 + + References + ---------- + .. [1] Digital Library of Mathematical Functions 29.12 + https://dlmf.nist.gov/29.12 + .. [2] Bardhan and Knepley, "Computational science and + re-discovery: open-source implementations of + ellipsoidal harmonics for problems in potential theory", + Comput. Sci. Disc. 5, 014006 (2012) + :doi:`10.1088/1749-4699/5/1/014006`. + .. [3] David J.and Dechambre P, "Computation of Ellipsoidal + Gravity Field Harmonics for small solar system bodies" + pp. 30-36, 2000 + .. [4] George Dassios, "Ellipsoidal Harmonics: Theory and Applications" + pp. 418, 2012 + + Examples + -------- + >>> from scipy.special import ellip_harm + >>> w = ellip_harm(5,8,1,1,2.5) + >>> w + 2.5 + + Check that the functions indeed are solutions to the Lamé equation: + + >>> import numpy as np + >>> from scipy.interpolate import UnivariateSpline + >>> def eigenvalue(f, df, ddf): + ... r = (((s**2 - h**2) * (s**2 - k**2) * ddf + ... + s * (2*s**2 - h**2 - k**2) * df + ... - n * (n + 1)*s**2*f) / f) + ... return -r.mean(), r.std() + >>> s = np.linspace(0.1, 10, 200) + >>> k, h, n, p = 8.0, 2.2, 3, 2 + >>> E = ellip_harm(h**2, k**2, n, p, s) + >>> E_spl = UnivariateSpline(s, E) + >>> a, a_err = eigenvalue(E_spl(s), E_spl(s,1), E_spl(s,2)) + >>> a, a_err + (583.44366156701483, 6.4580890640310646e-11) + + """ # noqa: E501 + return _ellip_harm(h2, k2, n, p, s, signm, signn) + + +_ellip_harm_2_vec = np.vectorize(_ellipsoid, otypes='d') + + +def ellip_harm_2(h2, k2, n, p, s): + r""" + Ellipsoidal harmonic functions F^p_n(l). + + These are also known as Lamé functions of the second kind, and are + solutions to the Lamé equation: + + .. math:: (s^2 - h^2)(s^2 - k^2)F''(s) + + s(2s^2 - h^2 - k^2)F'(s) + (a - q s^2)F(s) = 0 + + where :math:`q = (n+1)n` and :math:`a` is the eigenvalue (not + returned) corresponding to the solutions. + + Parameters + ---------- + h2 : float + ``h**2`` + k2 : float + ``k**2``; should be larger than ``h**2`` + n : int + Degree. + p : int + Order, can range between [1,2n+1]. + s : float + Coordinate + + Returns + ------- + F : float + The harmonic :math:`F^p_n(s)` + + See Also + -------- + ellip_harm, ellip_normal + + Notes + ----- + Lamé functions of the second kind are related to the functions of the first kind: + + .. math:: + + F^p_n(s)=(2n + 1)E^p_n(s)\int_{0}^{1/s} + \frac{du}{(E^p_n(1/u))^2\sqrt{(1-u^2k^2)(1-u^2h^2)}} + + .. versionadded:: 0.15.0 + + Examples + -------- + >>> from scipy.special import ellip_harm_2 + >>> w = ellip_harm_2(5,8,2,1,10) + >>> w + 0.00108056853382 + + """ + with np.errstate(all='ignore'): + return _ellip_harm_2_vec(h2, k2, n, p, s) + + +def _ellip_normal_vec(h2, k2, n, p): + return _ellipsoid_norm(h2, k2, n, p) + + +_ellip_normal_vec = np.vectorize(_ellip_normal_vec, otypes='d') + + +def ellip_normal(h2, k2, n, p): + r""" + Ellipsoidal harmonic normalization constants gamma^p_n. + + The normalization constant is defined as + + .. math:: + + \gamma^p_n=8\int_{0}^{h}dx\int_{h}^{k}dy + \frac{(y^2-x^2)(E^p_n(y)E^p_n(x))^2}{\sqrt((k^2-y^2)(y^2-h^2)(h^2-x^2)(k^2-x^2)} + + Parameters + ---------- + h2 : float + ``h**2`` + k2 : float + ``k**2``; should be larger than ``h**2`` + n : int + Degree. + p : int + Order, can range between [1,2n+1]. + + Returns + ------- + gamma : float + The normalization constant :math:`\gamma^p_n` + + See Also + -------- + ellip_harm, ellip_harm_2 + + Notes + ----- + .. versionadded:: 0.15.0 + + Examples + -------- + >>> from scipy.special import ellip_normal + >>> w = ellip_normal(5,8,3,7) + >>> w + 1723.38796997 + + """ + with np.errstate(all='ignore'): + return _ellip_normal_vec(h2, k2, n, p) diff --git a/v12x/pylib/scipy/special/_input_validation.py b/v12x/pylib/scipy/special/_input_validation.py new file mode 100644 index 0000000000000000000000000000000000000000..e5b7fe36df87617bf91655623ee0076c37a4d08a --- /dev/null +++ b/v12x/pylib/scipy/special/_input_validation.py @@ -0,0 +1,17 @@ +import math +import operator + +def _nonneg_int_or_fail(n, var_name, strict=True): + try: + if strict: + # Raises an exception if float + n = operator.index(n) + elif n == math.floor(n): + n = int(n) + else: + raise ValueError() + if n < 0: + raise ValueError() + except (ValueError, TypeError) as err: + raise err.__class__(f"{var_name} must be a non-negative integer") from err + return n diff --git a/v12x/pylib/scipy/special/_lambertw.py b/v12x/pylib/scipy/special/_lambertw.py new file mode 100644 index 0000000000000000000000000000000000000000..f758c7c21fdddc0ec1b84727d90c6de7f34a094e --- /dev/null +++ b/v12x/pylib/scipy/special/_lambertw.py @@ -0,0 +1,149 @@ +from ._ufuncs import _lambertw + +import numpy as np + + +def lambertw(z, k=0, tol=1e-8): + r""" + lambertw(z, k=0, tol=1e-8) + + Lambert W function. + + The Lambert W function `W(z)` is defined as the inverse function + of ``w * exp(w)``. In other words, the value of ``W(z)`` is + such that ``z = W(z) * exp(W(z))`` for any complex number + ``z``. + + The Lambert W function is a multivalued function with infinitely + many branches. Each branch gives a separate solution of the + equation ``z = w exp(w)``. Here, the branches are indexed by the + integer `k`. + + Parameters + ---------- + z : array_like + Input argument. + k : int, optional + Branch index. + tol : float, optional + Evaluation tolerance. + + Returns + ------- + w : array + `w` will have the same shape as `z`. + + See Also + -------- + wrightomega : the Wright Omega function + + Notes + ----- + All branches are supported by `lambertw`: + + * ``lambertw(z)`` gives the principal solution (branch 0) + * ``lambertw(z, k)`` gives the solution on branch `k` + + The Lambert W function has two partially real branches: the + principal branch (`k = 0`) is real for real ``z > -1/e``, and the + ``k = -1`` branch is real for ``-1/e < z < 0``. All branches except + ``k = 0`` have a logarithmic singularity at ``z = 0``. + + **Possible issues** + + The evaluation can become inaccurate very close to the branch point + at ``-1/e``. In some corner cases, `lambertw` might currently + fail to converge, or can end up on the wrong branch. + + **Algorithm** + + Halley's iteration is used to invert ``w * exp(w)``, using a first-order + asymptotic approximation (O(log(w)) or `O(w)`) as the initial estimate. + + The definition, implementation and choice of branches is based on [2]_. + + References + ---------- + .. [1] https://en.wikipedia.org/wiki/Lambert_W_function + .. [2] Corless et al, "On the Lambert W function", Adv. Comp. Math. 5 + (1996) 329-359. + https://cs.uwaterloo.ca/research/tr/1993/03/W.pdf + + Examples + -------- + The Lambert W function is the inverse of ``w exp(w)``: + + >>> import numpy as np + >>> from scipy.special import lambertw + >>> w = lambertw(1) + >>> w + (0.56714329040978384+0j) + >>> w * np.exp(w) + (1.0+0j) + + Any branch gives a valid inverse: + + >>> w = lambertw(1, k=3) + >>> w + (-2.8535817554090377+17.113535539412148j) + >>> w*np.exp(w) + (1.0000000000000002+1.609823385706477e-15j) + + **Applications to equation-solving** + + The Lambert W function may be used to solve various kinds of + equations. We give two examples here. + + First, the function can be used to solve implicit equations of the + form + + :math:`x = a + b e^{c x}` + + for :math:`x`. We assume :math:`c` is not zero. After a little + algebra, the equation may be written + + :math:`z e^z = -b c e^{a c}` + + where :math:`z = c (a - x)`. :math:`z` may then be expressed using + the Lambert W function + + :math:`z = W(-b c e^{a c})` + + giving + + :math:`x = a - W(-b c e^{a c})/c` + + For example, + + >>> a = 3 + >>> b = 2 + >>> c = -0.5 + + The solution to :math:`x = a + b e^{c x}` is: + + >>> x = a - lambertw(-b*c*np.exp(a*c))/c + >>> x + (3.3707498368978794+0j) + + Verify that it solves the equation: + + >>> a + b*np.exp(c*x) + (3.37074983689788+0j) + + The Lambert W function may also be used find the value of the infinite + power tower :math:`z^{z^{z^{\ldots}}}`: + + >>> def tower(z, n): + ... if n == 0: + ... return z + ... return z ** tower(z, n-1) + ... + >>> tower(0.5, 100) + 0.641185744504986 + >>> -lambertw(-np.log(0.5)) / np.log(0.5) + (0.64118574450498589+0j) + """ + # TODO: special expert should inspect this + # interception; better place to do it? + k = np.asarray(k, dtype=np.dtype("long")) + return _lambertw(z, k, tol) diff --git a/v12x/pylib/scipy/special/_logsumexp.py b/v12x/pylib/scipy/special/_logsumexp.py new file mode 100644 index 0000000000000000000000000000000000000000..fdfdcded06e22591ef48117d5ff8af83c339cf74 --- /dev/null +++ b/v12x/pylib/scipy/special/_logsumexp.py @@ -0,0 +1,420 @@ +import numpy as np +from scipy._lib._array_api import ( + array_namespace, + xp_capabilities, + xp_device, + xp_size, + xp_promote, + xp_float_to_complex, +) +from scipy._external import array_api_extra as xpx + +__all__ = ["logsumexp", "softmax", "log_softmax"] + + +@xp_capabilities() +def logsumexp(a, axis=None, b=None, keepdims=False, return_sign=False): + """Compute the log of the sum of exponentials of input elements. + + Parameters + ---------- + a : array_like + Input array. + axis : None or int or tuple of ints, optional + Axis or axes over which the sum is taken. By default `axis` is None, + and all elements are summed. + + .. versionadded:: 0.11.0 + b : array-like, optional + Scaling factor for exp(`a`) must be of the same shape as `a` or + broadcastable to `a`. These values may be negative in order to + implement subtraction. + + .. versionadded:: 0.12.0 + keepdims : bool, optional + If this is set to True, the axes which are reduced are left in the + result as dimensions with size one. With this option, the result + will broadcast correctly against the original array. + + .. versionadded:: 0.15.0 + return_sign : bool, optional + If this is set to True, the result will be a pair containing sign + information; if False, results that are negative will be returned + as NaN. Default is False (no sign information). + + .. versionadded:: 0.16.0 + + Returns + ------- + res : ndarray + The result, ``np.log(np.sum(np.exp(a)))`` calculated in a numerically + more stable way. If `b` is given then ``np.log(np.sum(b*np.exp(a)))`` + is returned. If ``return_sign`` is True, ``res`` contains the log of + the absolute value of the argument. + sgn : ndarray + If ``return_sign`` is True, this will be an array of floating-point + numbers matching res containing +1, 0, -1 (for real-valued inputs) + or a complex phase (for complex inputs). This gives the sign of the + argument of the logarithm in ``res``. + If ``return_sign`` is False, only one result is returned. + + See Also + -------- + :data:`numpy.logaddexp` + :data:`numpy.logaddexp2` + + Notes + ----- + NumPy has a logaddexp function which is very similar to `logsumexp`, but + only handles two arguments. `logaddexp.reduce` is similar to this + function, but may be less stable. + + The logarithm is a multivalued function: for each :math:`x` there is an + infinite number of :math:`z` such that :math:`exp(z) = x`. The convention + is to return the :math:`z` whose imaginary part lies in :math:`(-pi, pi]`. + + Examples + -------- + >>> import numpy as np + >>> from scipy.special import logsumexp + >>> a = np.arange(10) + >>> logsumexp(a) + 9.4586297444267107 + >>> np.log(np.sum(np.exp(a))) + 9.4586297444267107 + + With weights + + >>> a = np.arange(10) + >>> b = np.arange(10, 0, -1) + >>> logsumexp(a, b=b) + 9.9170178533034665 + >>> np.log(np.sum(b*np.exp(a))) + 9.9170178533034647 + + Returning a sign flag + + >>> logsumexp([1,2],b=[1,-1],return_sign=True) + (1.5413248546129181, -1.0) + + Notice that `logsumexp` does not directly support masked arrays. To use it + on a masked array, convert the mask into zero weights: + + >>> a = np.ma.array([np.log(2), 2, np.log(3)], + ... mask=[False, True, False]) + >>> b = (~a.mask).astype(int) + >>> logsumexp(a.data, b=b), np.log(5) + 1.6094379124341005, 1.6094379124341005 + + """ + xp = array_namespace(a, b) + a, b = xp_promote(a, b, broadcast=True, force_floating=True, xp=xp) + a = xpx.atleast_nd(a, ndim=1, xp=xp) + b = xpx.atleast_nd(b, ndim=1, xp=xp) if b is not None else b + axis = tuple(range(a.ndim)) if axis is None else axis + + if xp_size(a) != 0: + with np.errstate(divide='ignore', invalid='ignore', over='ignore'): + # Where result is infinite, we use the direct logsumexp calculation to + # delegate edge case handling to the behavior of `xp.log` and `xp.exp`, + # which should follow the C99 standard for complex values. + b_exp_a = xp.exp(a) if b is None else b * xp.exp(a) + sum_ = xp.sum(b_exp_a, axis=axis, keepdims=True) + sgn_inf = xp.sign(sum_) if return_sign else None + sum_ = xp.abs(sum_) if return_sign else sum_ + out_inf = xp.log(sum_) + + with np.errstate(divide='ignore', invalid='ignore'): # log of zero is OK + out, sgn = _logsumexp(a, b, axis=axis, return_sign=return_sign, xp=xp) + + # Replace infinite results. This probably could be done with an + # `apply_where`-like strategy to avoid redundant calculation, but currently + # `apply_where` itself is only for elementwise functions. + out_finite = xp.isfinite(out) + out = xp.where(out_finite, out, out_inf) + sgn = xp.where(out_finite, sgn, sgn_inf) if return_sign else sgn + else: + shape = np.asarray(a.shape) # NumPy is convenient for shape manipulation + shape[axis] = 1 + out = xp.full(tuple(shape), -xp.inf, dtype=a.dtype, device=xp_device(a)) + sgn = xp.sign(out) + + if xp.isdtype(out.dtype, 'complex floating'): + if return_sign: + real = xp.real(sgn) + imag = xp_float_to_complex(_wrap_radians(xp.imag(sgn), xp=xp), xp=xp) + sgn = real + imag*1j + else: + real = xp.real(out) + imag = xp_float_to_complex(_wrap_radians(xp.imag(out), xp=xp), xp=xp) + out = real + imag*1j + + # Deal with shape details - reducing dimensions and convert 0-D to scalar for NumPy + out = xp.squeeze(out, axis=axis) if not keepdims else out + sgn = xp.squeeze(sgn, axis=axis) if (sgn is not None and not keepdims) else sgn + out = out[()] if out.ndim == 0 else out + sgn = sgn[()] if (sgn is not None and sgn.ndim == 0) else sgn + + return (out, sgn) if return_sign else out + + +def _wrap_radians(x, *, xp): + # Wrap radians to (-pi, pi] interval + wrapped = -((-x + xp.pi) % (2 * xp.pi) - xp.pi) + # preserve relative precision + no_wrap = xp.abs(x) < xp.pi + return xp.where(no_wrap, x, wrapped) + + +def _elements_and_indices_with_max_real(a, *, axis=-1, xp): + # This is an array-API compatible `max` function that works something + # like `np.max` for complex input. The important part is that it finds + # the element with maximum real part. When there are multiple complex values + # with this real part, it doesn't matter which we choose. + # We could use `argmax` on real component, but array API doesn't yet have + # `take_along_axis`, and even if it did, we would have problems with axis tuples. + # Feel free to rewrite! It's ugly, but it's not the purpose of the PR, and + # it gets the job done. + + if xp.isdtype(a.dtype, "complex floating"): + # select all elements with max real part. + real_a = xp.real(a) + max_ = xp.max(real_a, axis=axis, keepdims=True) + mask = real_a == max_ + + # Of those, choose one arbitrarily. This is a reasonably + # simple, array-API compatible way of doing so that doesn't + # have a problem with `axis` being a tuple or None. + i = xp.reshape(xp.arange(xp_size(a), device=xp_device(a)), a.shape) + i = xpx.at(i, ~mask).set(-1) + max_i = xp.max(i, axis=axis, keepdims=True) + mask = i == max_i + a = xp.where(mask, a, 0.) + max_ = xp.sum(a, axis=axis, dtype=a.dtype, keepdims=True) + else: + max_ = xp.max(a, axis=axis, keepdims=True) + mask = a == max_ + + return max_, mask + + +def _logsumexp(a, b, *, axis, return_sign, xp): + # This has been around for about a decade, so let's consider it a feature: + # Even if element of `a` is infinite or NaN, it adds nothing to the sum if + # the corresponding weight is zero. + if b is not None: + a = xpx.at(a, b == 0).set(-xp.inf, copy=True) + + # Find element with maximum real part, since this is what affects the magnitude + # of the exponential. Possible enhancement: include log of `b` magnitude in `a`. + a_max, i_max = _elements_and_indices_with_max_real(a, axis=axis, xp=xp) + + # for precision, these terms are separated out of the main sum. + a = xpx.at(a, i_max).set(-xp.inf, copy=True if b is None else None) + i_max_dt = xp.astype(i_max, a.dtype) + # This is an inefficient way of getting `m` because it is the sum of a sparse + # array; however, this is the simplest way I can think of to get the right shape. + b_i_max = i_max_dt if b is None else b * i_max_dt + m = xp.sum(b_i_max, axis=axis, keepdims=True, dtype=a.dtype) + + # Shift, exponentiate, scale, and sum + exp = b * xp.exp(a - a_max) if b is not None else xp.exp(a - a_max) + s = xp.sum(exp, axis=axis, keepdims=True, dtype=exp.dtype) + s = xp.where(s == 0, s, s/m) + + # Separate sign/magnitude information + # Originally, this was only performed if `return_sign=True`. + # However, this is also needed if any elements of `m < 0` or `s < -1`. + # An improvement would be to perform the calculations only on these entries. + + sgn = xp.sign(s + 1) * xp.sign(m) + + if xp.isdtype(s.dtype, "real floating"): + # The log functions need positive arguments + s = xp.where(s < -1, -s - 2, s) + m = xp.abs(m) + else: + # `a_max` can have a sign component for complex input + sgn = sgn * xp.exp(xp.imag(a_max) * 1.0j) + + # Take log and undo shift + out = xp.log1p(s) + xp.log(m) + a_max + + if return_sign: + out = xp.real(out) + elif xp.isdtype(out.dtype, 'real floating'): + out = xpx.at(out)[sgn < 0].set(xp.nan) + + return out, sgn + + +@xp_capabilities() +def softmax(x, axis=None): + r"""Compute the softmax function. + + The softmax function transforms each element of a collection by + computing the exponential of each element divided by the sum of the + exponentials of all the elements. That is, if `x` is a one-dimensional + numpy array:: + + softmax(x) = np.exp(x)/sum(np.exp(x)) + + Parameters + ---------- + x : array_like + Input array. + axis : int or tuple of ints, optional + Axis to compute values along. Default is None and softmax will be + computed over the entire array `x`. + + Returns + ------- + s : ndarray + An array the same shape as `x`. The result will sum to 1 along the + specified axis. + + Notes + ----- + The formula for the softmax function :math:`\sigma(x)` for a vector + :math:`x = \{x_0, x_1, ..., x_{n-1}\}` is + + .. math:: \sigma(x)_j = \frac{e^{x_j}}{\sum_k e^{x_k}} + + The `softmax` function is the gradient of `logsumexp`. + + The implementation uses shifting to avoid overflow. See [1]_ for more + details. + + .. versionadded:: 1.2.0 + + References + ---------- + .. [1] P. Blanchard, D.J. Higham, N.J. Higham, "Accurately computing the + log-sum-exp and softmax functions", IMA Journal of Numerical Analysis, + Vol.41(4), :doi:`10.1093/imanum/draa038`. + + Examples + -------- + >>> import numpy as np + >>> from scipy.special import softmax + >>> np.set_printoptions(precision=5) + + >>> x = np.array([[1, 0.5, 0.2, 3], + ... [1, -1, 7, 3], + ... [2, 12, 13, 3]]) + ... + + Compute the softmax transformation over the entire array. + + >>> m = softmax(x) + >>> m + array([[ 4.48309e-06, 2.71913e-06, 2.01438e-06, 3.31258e-05], + [ 4.48309e-06, 6.06720e-07, 1.80861e-03, 3.31258e-05], + [ 1.21863e-05, 2.68421e-01, 7.29644e-01, 3.31258e-05]]) + + >>> m.sum() + 1.0 + + Compute the softmax transformation along the first axis (i.e., the + columns). + + >>> m = softmax(x, axis=0) + + >>> m + array([[ 2.11942e-01, 1.01300e-05, 2.75394e-06, 3.33333e-01], + [ 2.11942e-01, 2.26030e-06, 2.47262e-03, 3.33333e-01], + [ 5.76117e-01, 9.99988e-01, 9.97525e-01, 3.33333e-01]]) + + >>> m.sum(axis=0) + array([ 1., 1., 1., 1.]) + + Compute the softmax transformation along the second axis (i.e., the rows). + + >>> m = softmax(x, axis=1) + >>> m + array([[ 1.05877e-01, 6.42177e-02, 4.75736e-02, 7.82332e-01], + [ 2.42746e-03, 3.28521e-04, 9.79307e-01, 1.79366e-02], + [ 1.22094e-05, 2.68929e-01, 7.31025e-01, 3.31885e-05]]) + + >>> m.sum(axis=1) + array([ 1., 1., 1.]) + + """ + xp = array_namespace(x) + x = xp.asarray(x) + x_max = xp.max(x, axis=axis, keepdims=True) + exp_x_shifted = xp.exp(x - x_max) + return exp_x_shifted / xp.sum(exp_x_shifted, axis=axis, keepdims=True) + + +@xp_capabilities() +def log_softmax(x, axis=None): + r"""Compute the logarithm of the softmax function. + + In principle:: + + log_softmax(x) = log(softmax(x)) + + but using a more accurate implementation. + + Parameters + ---------- + x : array_like + Input array. + axis : int or tuple of ints, optional + Axis to compute values along. Default is None and softmax will be + computed over the entire array `x`. + + Returns + ------- + s : ndarray or scalar + An array with the same shape as `x`. Exponential of the result will + sum to 1 along the specified axis. If `x` is a scalar, a scalar is + returned. + + Notes + ----- + `log_softmax` is more accurate than ``np.log(softmax(x))`` with inputs that + make `softmax` saturate (see examples below). + + .. versionadded:: 1.5.0 + + Examples + -------- + >>> import numpy as np + >>> from scipy.special import log_softmax + >>> from scipy.special import softmax + >>> np.set_printoptions(precision=5) + + >>> x = np.array([1000.0, 1.0]) + + >>> y = log_softmax(x) + >>> y + array([ 0., -999.]) + + >>> with np.errstate(divide='ignore'): + ... y = np.log(softmax(x)) + ... + >>> y + array([ 0., -inf]) + + """ + xp = array_namespace(x) + x = xp.asarray(x) + + x_max = xp.max(x, axis=axis, keepdims=True) + + if x_max.ndim > 0: + x_max = xpx.at(x_max, ~xp.isfinite(x_max)).set(0) + elif not xp.isfinite(x_max): + x_max = 0 + + tmp = x - x_max + exp_tmp = xp.exp(tmp) + + # suppress warnings about log of zero + with np.errstate(divide='ignore'): + s = xp.sum(exp_tmp, axis=axis, keepdims=True) + out = xp.log(s) + + return tmp - out diff --git a/v12x/pylib/scipy/special/_mptestutils.py b/v12x/pylib/scipy/special/_mptestutils.py new file mode 100644 index 0000000000000000000000000000000000000000..9e519093dface79e21f16d7063541ad107f5ca96 --- /dev/null +++ b/v12x/pylib/scipy/special/_mptestutils.py @@ -0,0 +1,453 @@ +import os +import sys +import time +from itertools import zip_longest + +import numpy as np +from numpy.testing import assert_ +import pytest + +from scipy.special._testutils import assert_func_equal + +try: + import mpmath +except ImportError: + pass + + +# ------------------------------------------------------------------------------ +# Machinery for systematic tests with mpmath +# ------------------------------------------------------------------------------ + +class Arg: + """Generate a set of numbers on the real axis, concentrating on + 'interesting' regions and covering all orders of magnitude. + + """ + + def __init__(self, a=-np.inf, b=np.inf, inclusive_a=True, inclusive_b=True): + if a > b: + raise ValueError("a should be less than or equal to b") + if a == -np.inf: + a = -0.5*np.finfo(float).max + if b == np.inf: + b = 0.5*np.finfo(float).max + self.a, self.b = a, b + + self.inclusive_a, self.inclusive_b = inclusive_a, inclusive_b + + def _positive_values(self, a, b, n): + if a < 0: + raise ValueError("a should be positive") + + # Try to put half of the points into a linspace between a and + # 10 the other half in a logspace. + if n % 2 == 0: + nlogpts = n//2 + nlinpts = nlogpts + else: + nlogpts = n//2 + nlinpts = nlogpts + 1 + + if a >= 10: + # Outside of linspace range; just return a logspace. + pts = np.logspace(np.log10(a), np.log10(b), n) + elif a > 0 and b < 10: + # Outside of logspace range; just return a linspace + pts = np.linspace(a, b, n) + elif a > 0: + # Linspace between a and 10 and a logspace between 10 and + # b. + linpts = np.linspace(a, 10, nlinpts, endpoint=False) + logpts = np.logspace(1, np.log10(b), nlogpts) + pts = np.hstack((linpts, logpts)) + elif a == 0 and b <= 10: + # Linspace between 0 and b and a logspace between 0 and + # the smallest positive point of the linspace + linpts = np.linspace(0, b, nlinpts) + if linpts.size > 1: + right = np.log10(linpts[1]) + else: + right = -30 + logpts = np.logspace(-30, right, nlogpts, endpoint=False) + pts = np.hstack((logpts, linpts)) + else: + # Linspace between 0 and 10, logspace between 0 and the + # smallest positive point of the linspace, and a logspace + # between 10 and b. + if nlogpts % 2 == 0: + nlogpts1 = nlogpts//2 + nlogpts2 = nlogpts1 + else: + nlogpts1 = nlogpts//2 + nlogpts2 = nlogpts1 + 1 + linpts = np.linspace(0, 10, nlinpts, endpoint=False) + if linpts.size > 1: + right = np.log10(linpts[1]) + else: + right = -30 + logpts1 = np.logspace(-30, right, nlogpts1, endpoint=False) + logpts2 = np.logspace(1, np.log10(b), nlogpts2) + pts = np.hstack((logpts1, linpts, logpts2)) + + return np.sort(pts) + + def values(self, n): + """Return an array containing n numbers.""" + a, b = self.a, self.b + if a == b: + return np.zeros(n) + + if not self.inclusive_a: + n += 1 + if not self.inclusive_b: + n += 1 + + if n % 2 == 0: + n1 = n//2 + n2 = n1 + else: + n1 = n//2 + n2 = n1 + 1 + + if a >= 0: + pospts = self._positive_values(a, b, n) + negpts = [] + elif b <= 0: + pospts = [] + negpts = -self._positive_values(-b, -a, n) + else: + pospts = self._positive_values(0, b, n1) + negpts = -self._positive_values(0, -a, n2 + 1) + # Don't want to get zero twice + negpts = negpts[1:] + pts = np.hstack((negpts[::-1], pospts)) + + if not self.inclusive_a: + pts = pts[1:] + if not self.inclusive_b: + pts = pts[:-1] + return pts + + +class FixedArg: + def __init__(self, values): + self._values = np.asarray(values) + + def values(self, n): + return self._values + + +class ComplexArg: + def __init__(self, a=complex(-np.inf, -np.inf), b=complex(np.inf, np.inf)): + self.real = Arg(a.real, b.real) + self.imag = Arg(a.imag, b.imag) + + def values(self, n): + m = int(np.floor(np.sqrt(n))) + x = self.real.values(m) + y = self.imag.values(m + 1) + return (x[:,None] + 1j*y[None,:]).ravel() + + +class IntArg: + def __init__(self, a=-1000, b=1000): + self.a = a + self.b = b + + def values(self, n): + v1 = Arg(self.a, self.b).values(max(1 + n//2, n-5)).astype(int) + v2 = np.arange(-5, 5) + v = np.unique(np.r_[v1, v2]) + v = v[(v >= self.a) & (v < self.b)] + return v + + +def get_args(argspec, n): + if isinstance(argspec, np.ndarray): + args = argspec.copy() + else: + nargs = len(argspec) + ms = np.asarray( + [1.5 if isinstance(spec, ComplexArg) else 1.0 for spec in argspec] + ) + ms = (n**(ms/sum(ms))).astype(int) + 1 + + args = [spec.values(m) for spec, m in zip(argspec, ms)] + args = np.array(np.broadcast_arrays(*np.ix_(*args))).reshape(nargs, -1).T + + return args + + +class MpmathData: + def __init__(self, scipy_func, mpmath_func, arg_spec, name=None, + dps=None, prec=None, n=None, rtol=1e-7, atol=1e-300, + ignore_inf_sign=False, distinguish_nan_and_inf=True, + nan_ok=True, param_filter=None): + + # mpmath tests are really slow (see gh-6989). Use a small number of + # points by default, increase back to 5000 (old default) if XSLOW is + # set + if n is None: + try: + is_xslow = int(os.environ.get('SCIPY_XSLOW', '0')) + except ValueError: + is_xslow = False + + n = 5000 if is_xslow else 500 + + self.scipy_func = scipy_func + self.mpmath_func = mpmath_func + self.arg_spec = arg_spec + self.dps = dps + self.prec = prec + self.n = n + self.rtol = rtol + self.atol = atol + self.ignore_inf_sign = ignore_inf_sign + self.nan_ok = nan_ok + if isinstance(self.arg_spec, np.ndarray): + self.is_complex = np.issubdtype(self.arg_spec.dtype, np.complexfloating) + else: + self.is_complex = any( + [isinstance(arg, ComplexArg) for arg in self.arg_spec] + ) + self.ignore_inf_sign = ignore_inf_sign + self.distinguish_nan_and_inf = distinguish_nan_and_inf + if not name or name == '': + name = getattr(scipy_func, '__name__', None) + if not name or name == '': + name = getattr(mpmath_func, '__name__', None) + self.name = name + self.param_filter = param_filter + + def check(self): + np.random.seed(1234) + + # Generate values for the arguments + argarr = get_args(self.arg_spec, self.n) + + # Check + old_dps, old_prec = mpmath.mp.dps, mpmath.mp.prec + try: + if self.dps is not None: + dps_list = [self.dps] + else: + dps_list = [20] + if self.prec is not None: + mpmath.mp.prec = self.prec + + # Proper casting of mpmath input and output types. Using + # native mpmath types as inputs gives improved precision + # in some cases. + if np.issubdtype(argarr.dtype, np.complexfloating): + pytype = mpc2complex + + def mptype(x): + return mpmath.mpc(complex(x)) + else: + def mptype(x): + return mpmath.mpf(float(x)) + + def pytype(x): + if abs(x.imag) > 1e-16*(1 + abs(x.real)): + return np.nan + else: + return mpf2float(x.real) + + # Try out different dps until one (or none) works + for j, dps in enumerate(dps_list): + mpmath.mp.dps = dps + + try: + assert_func_equal( + self.scipy_func, + lambda *a: pytype(self.mpmath_func(*map(mptype, a))), + argarr, + vectorized=False, + rtol=self.rtol, + atol=self.atol, + ignore_inf_sign=self.ignore_inf_sign, + distinguish_nan_and_inf=self.distinguish_nan_and_inf, + nan_ok=self.nan_ok, + param_filter=self.param_filter + ) + break + except AssertionError: + if j >= len(dps_list)-1: + # reraise the Exception + tp, value, tb = sys.exc_info() + if value.__traceback__ is not tb: + raise value.with_traceback(tb) + raise value + finally: + mpmath.mp.dps, mpmath.mp.prec = old_dps, old_prec + + def __repr__(self): + if self.is_complex: + return f"" + else: + return f"" + + +def assert_mpmath_equal(*a, **kw): + d = MpmathData(*a, **kw) + d.check() + + +def nonfunctional_tooslow(func): + return pytest.mark.skip( + reason=" Test not yet functional (too slow), needs more work." + )(func) + + +# ------------------------------------------------------------------------------ +# Tools for dealing with mpmath quirks +# ------------------------------------------------------------------------------ + +def mpf2float(x): + """ + Convert an mpf to the nearest floating point number. Just using + float directly doesn't work because of results like this: + + with mp.workdps(50): + float(mpf("0.99999999999999999")) = 0.9999999999999999 + + """ + return float(mpmath.nstr(x, 17, min_fixed=0, max_fixed=0)) + + +def mpc2complex(x): + return complex(mpf2float(x.real), mpf2float(x.imag)) + + +def trace_args(func): + def tofloat(x): + if isinstance(x, mpmath.mpc): + return complex(x) + else: + return float(x) + + def wrap(*a, **kw): + sys.stderr.write(f"{tuple(map(tofloat, a))!r}: ") + sys.stderr.flush() + try: + r = func(*a, **kw) + sys.stderr.write(f"-> {r!r}") + finally: + sys.stderr.write("\n") + sys.stderr.flush() + return r + return wrap + + +try: + import signal + POSIX = ('setitimer' in dir(signal)) +except ImportError: + POSIX = False + + +class TimeoutError(Exception): + pass + + +def time_limited(timeout=0.5, return_val=np.nan, use_sigalrm=True): + """ + Decorator for setting a timeout for pure-Python functions. + + If the function does not return within `timeout` seconds, the + value `return_val` is returned instead. + + On POSIX this uses SIGALRM by default. On non-POSIX, settrace is + used. Do not use this with threads: the SIGALRM implementation + does probably not work well. The settrace implementation only + traces the current thread. + + The settrace implementation slows down execution speed. Slowdown + by a factor around 10 is probably typical. + """ + if POSIX and use_sigalrm: + def sigalrm_handler(signum, frame): + raise TimeoutError() + + def deco(func): + def wrap(*a, **kw): + old_handler = signal.signal(signal.SIGALRM, sigalrm_handler) + signal.setitimer(signal.ITIMER_REAL, timeout) + try: + return func(*a, **kw) + except TimeoutError: + return return_val + finally: + signal.setitimer(signal.ITIMER_REAL, 0) + signal.signal(signal.SIGALRM, old_handler) + return wrap + else: + def deco(func): + def wrap(*a, **kw): + start_time = time.time() + + def trace(frame, event, arg): + if time.time() - start_time > timeout: + raise TimeoutError() + return trace + sys.settrace(trace) + try: + return func(*a, **kw) + except TimeoutError: + sys.settrace(None) + return return_val + finally: + sys.settrace(None) + return wrap + return deco + + +def exception_to_nan(func): + """Decorate function to return nan if it raises an exception""" + def wrap(*a, **kw): + try: + return func(*a, **kw) + except Exception: + return np.nan + return wrap + + +def inf_to_nan(func): + """Decorate function to return nan if it returns inf""" + def wrap(*a, **kw): + v = func(*a, **kw) + if not np.isfinite(v): + return np.nan + return v + return wrap + + +def mp_assert_allclose(res, std, atol=0, rtol=1e-17): + """ + Compare lists of mpmath.mpf's or mpmath.mpc's directly so that it + can be done to higher precision than double. + """ + failures = [] + for k, (resval, stdval) in enumerate(zip_longest(res, std)): + if resval is None or stdval is None: + raise ValueError('Lengths of inputs res and std are not equal.') + if mpmath.fabs(resval - stdval) > atol + rtol*mpmath.fabs(stdval): + failures.append((k, resval, stdval)) + + nfail = len(failures) + if nfail > 0: + ndigits = int(abs(np.log10(rtol))) + msg = [""] + msg.append(f"Bad results ({nfail} out of {k + 1}) for the following points:") + for k, resval, stdval in failures: + resrep = mpmath.nstr(resval, ndigits, min_fixed=0, max_fixed=0) + stdrep = mpmath.nstr(stdval, ndigits, min_fixed=0, max_fixed=0) + if stdval == 0: + rdiff = "inf" + else: + rdiff = mpmath.fabs((resval - stdval)/stdval) + rdiff = mpmath.nstr(rdiff, 3) + msg.append(f"{k}: {resrep} != {stdrep} (rdiff {rdiff})") + assert_(False, "\n".join(msg)) diff --git a/v12x/pylib/scipy/special/_multiufuncs.py b/v12x/pylib/scipy/special/_multiufuncs.py new file mode 100644 index 0000000000000000000000000000000000000000..f62002ef701f0edffeb8ec348559f1436f80ab08 --- /dev/null +++ b/v12x/pylib/scipy/special/_multiufuncs.py @@ -0,0 +1,744 @@ +import collections +import numbers +import numpy as np + +from ._input_validation import _nonneg_int_or_fail + +from ._special_ufuncs import (legendre_p, assoc_legendre_p, + sph_legendre_p, sph_harm_y) +from ._gufuncs import (legendre_p_all, assoc_legendre_p_all, + sph_legendre_p_all, sph_harm_y_all) + +__all__ = [ + "assoc_legendre_p", + "assoc_legendre_p_all", + "legendre_p", + "legendre_p_all", + "sph_harm_y", + "sph_harm_y_all", + "sph_legendre_p", + "sph_legendre_p_all", +] + + +class MultiUFunc: + def __init__(self, ufunc_or_ufuncs, name=None, doc=None, *, + force_complex_output=False, **default_kwargs): + if not isinstance(ufunc_or_ufuncs, np.ufunc): + if isinstance(ufunc_or_ufuncs, collections.abc.Mapping): + ufuncs_iter = ufunc_or_ufuncs.values() + elif isinstance(ufunc_or_ufuncs, collections.abc.Iterable): + ufuncs_iter = ufunc_or_ufuncs + else: + raise ValueError("ufunc_or_ufuncs should be a ufunc or a" + " ufunc collection") + + # Perform input validation to ensure all ufuncs in ufuncs are + # actually ufuncs and all take the same input types. + seen_input_types = set() + for ufunc in ufuncs_iter: + if not isinstance(ufunc, np.ufunc): + raise ValueError("All ufuncs must have type `numpy.ufunc`." + f" Received {ufunc_or_ufuncs}") + seen_input_types.add(frozenset(x.split("->")[0] for x in ufunc.types)) + if len(seen_input_types) > 1: + raise ValueError("All ufuncs must take the same input types.") + + self.__name__ = name + self._ufunc_or_ufuncs = ufunc_or_ufuncs + self.__doc = doc + self.__force_complex_output = force_complex_output + self._default_kwargs = default_kwargs + self._resolve_out_shapes = None + self._finalize_out = None + self._key = None + self._ufunc_default_args = lambda *args, **kwargs: () + self._ufunc_default_kwargs = lambda *args, **kwargs: {} + self.__text_signature__ = doc.split("\n\n")[0] if doc else None + + @property + def __doc__(self): # pyrefly:ignore[bad-override] + return self.__doc + + def _override_key(self, func): + """Set `key` method by decorating a function. + """ + self._key = func + + def _override_ufunc_default_args(self, func): + self._ufunc_default_args = func + + def _override_ufunc_default_kwargs(self, func): + self._ufunc_default_kwargs = func + + def _override_resolve_out_shapes(self, func): + """Set `resolve_out_shapes` method by decorating a function.""" + if func.__doc__ is None: + func.__doc__ = \ + """Resolve to output shapes based on relevant inputs.""" + func.__name__ = "resolve_out_shapes" + self._resolve_out_shapes = func + + def _override_finalize_out(self, func): + self._finalize_out = func + + def _resolve_ufunc(self, **kwargs): + """Resolve to a ufunc based on keyword arguments.""" + + if isinstance(self._ufunc_or_ufuncs, np.ufunc): + return self._ufunc_or_ufuncs + + ufunc_key = self._key(**kwargs) + return self._ufunc_or_ufuncs[ufunc_key] + + def __call__(self, *args, **kwargs): + kwargs = self._default_kwargs | kwargs + + args += self._ufunc_default_args(**kwargs) + + ufunc = self._resolve_ufunc(**kwargs) + + # array arguments to be passed to the ufunc + ufunc_args = [np.asarray(arg) for arg in args[-ufunc.nin:]] + + ufunc_kwargs = self._ufunc_default_kwargs(**kwargs) + + if (self._resolve_out_shapes is not None): + ufunc_arg_shapes = tuple(np.shape(ufunc_arg) for ufunc_arg in ufunc_args) + ufunc_out_shapes = self._resolve_out_shapes(*args[:-ufunc.nin], + *ufunc_arg_shapes, ufunc.nout, + **kwargs) + + ufunc_arg_dtypes = tuple(ufunc_arg.dtype if hasattr(ufunc_arg, 'dtype') + else np.dtype(type(ufunc_arg)) + for ufunc_arg in ufunc_args) + + if hasattr(ufunc, 'resolve_dtypes'): + ufunc_dtypes = ufunc_arg_dtypes + ufunc.nout * (None,) + ufunc_dtypes = ufunc.resolve_dtypes(ufunc_dtypes) + ufunc_out_dtypes = ufunc_dtypes[-ufunc.nout:] + else: + ufunc_out_dtype = np.result_type(*ufunc_arg_dtypes) + if (not np.issubdtype(ufunc_out_dtype, np.inexact)): + ufunc_out_dtype = np.float64 + + ufunc_out_dtypes = ufunc.nout * (ufunc_out_dtype,) + + if self.__force_complex_output: + ufunc_out_dtypes = tuple(np.result_type(1j, ufunc_out_dtype) + for ufunc_out_dtype in ufunc_out_dtypes) + + out = tuple(np.empty(ufunc_out_shape, dtype=ufunc_out_dtype) + for ufunc_out_shape, ufunc_out_dtype + in zip(ufunc_out_shapes, ufunc_out_dtypes)) + + ufunc_kwargs['out'] = out + + out = ufunc(*ufunc_args, **ufunc_kwargs) + if (self._finalize_out is not None): + out = self._finalize_out(out) + + return out + + +sph_legendre_p = MultiUFunc( + sph_legendre_p, + "sph_legendre_p", + r"""sph_legendre_p(n, m, theta, *, diff_n=0) + + Spherical Legendre polynomial of the first kind. + + Parameters + ---------- + n : array_like of ints + Degree of the spherical Legendre polynomial. Must have ``n >= 0``. + m : array_like of ints + Order of the spherical Legendre polynomial. + theta : array_like + Input value. + diff_n : int, optional + A non-negative integer. Compute and return all derivatives up + to order `diff_n`. Default is 0. + + Returns + ------- + ndarray or tuple of ndarray + Spherical Legendre polynomial with ``diff_n`` derivatives. + + Notes + ----- + The spherical counterpart of an (unnormalized) associated Legendre polynomial has + the additional factor + + .. math:: + + \sqrt{\frac{(2 n + 1) (n - m)!}{4 \pi (n + m)!}} + + It is the same as the spherical harmonic :math:`Y_{n}^{m}(\theta, \phi)` + with :math:`\phi = 0`. + """, diff_n=0 +) + + +@sph_legendre_p._override_key +def _(diff_n): + diff_n = _nonneg_int_or_fail(diff_n, "diff_n", strict=False) + if not 0 <= diff_n <= 2: + raise ValueError( + "diff_n is currently only implemented for orders 0, 1, and 2," + f" received: {diff_n}." + ) + return diff_n + + +@sph_legendre_p._override_finalize_out +def _(out): + return np.moveaxis(out, -1, 0) + + +sph_legendre_p_all = MultiUFunc( + sph_legendre_p_all, + "sph_legendre_p_all", + """sph_legendre_p_all(n, m, theta, *, diff_n=0) + + All spherical Legendre polynomials of the first kind up to the + specified degree `n`, order `m`, and all derivatives up + to order `diff_n`. + + Parameters + ---------- + n : int + Degree of the spherical Legendre polynomials. Must have ``n >= 0``. + m : int + Order of the spherical Legendre polynomials. + theta : array_like + Input value. + diff_n : int, optional + A non-negative integer. Compute and return all derivatives up + to order `diff_n`. Default is 0. + + Returns + ------- + ndarray + Output shape is ``(diff_n + 1, n + 1, 2 * m + 1, ...)``. The entry at + ``(i, j, k)`` corresponds to the ``i``-th derivative, degree ``j``, and + order ``k`` for all ``0 <= i <= diff_n``, ``0 <= j <= n``, and + ``-m <= k <= m``. + + See Also + -------- + sph_legendre_p + """, diff_n=0 +) + + +@sph_legendre_p_all._override_key +def _(diff_n): + diff_n = _nonneg_int_or_fail(diff_n, "diff_n", strict=False) + if not 0 <= diff_n <= 2: + raise ValueError( + "diff_n is currently only implemented for orders 0, 1, and 2," + f" received: {diff_n}." + ) + return diff_n + + +@sph_legendre_p_all._override_ufunc_default_kwargs +def _(diff_n): + return {'axes': [()] + [(0, 1, -1)]} + + +@sph_legendre_p_all._override_resolve_out_shapes +def _(n, m, theta_shape, nout, diff_n): + if not isinstance(n, numbers.Integral) or (n < 0): + raise ValueError("n must be a non-negative integer.") + + return ((n + 1, 2 * abs(m) + 1) + theta_shape + (diff_n + 1,),) + + +@sph_legendre_p_all._override_finalize_out +def _(out): + return np.moveaxis(out, -1, 0) + + +assoc_legendre_p = MultiUFunc( + assoc_legendre_p, + "assoc_legendre_p", + r"""assoc_legendre_p(n, m, z, *, branch_cut=2, norm=False, diff_n=0) + + Associated Legendre polynomial of the first kind. + + Defined as + + .. math:: + + P_n^m(z) = (-1)^m (1 - z^2)^{m/2} + \frac{d^m}{dz^m} P_n(z) + + where :math:`P_n` is the Legendre polynomial. + + This definition includes the Condon-Shortley phase :math:`(-1)^m`. + + Parameters + ---------- + n : array_like of ints + Degree of the associated Legendre polynomial. Must have ``n >= 0``. + m : array_like of ints + Order of the associated Legendre polynomial. + z : array_like + Input value. + branch_cut : array_like of ints, optional + Selects branch cut. Must be 2 (default) or 3. + 2: cut on the real axis ``|z| > 1`` + 3: cut on the real axis ``-1 < z < 1`` + norm : bool, optional + If ``True``, compute the normalized associated Legendre polynomial. + Default is ``False``. + diff_n : int, optional + A non-negative integer. Compute and return all derivatives up + to order ``diff_n``. Default is 0. + + Returns + ------- + ndarray or tuple of ndarray + Associated Legendre polynomial with ``diff_n`` derivatives. + + Notes + ----- + The normalized counterpart of an (unnormalized) associated Legendre + polynomial has the additional factor + + .. math:: + + \sqrt{\frac{(2 n + 1) (n - m)!}{2 (n + m)!}} + + Examples + -------- + Evaluate :math:`P_2^1(z)` on :math:`[-1, 1]`: + + >>> import numpy as np + >>> from scipy.special import assoc_legendre_p + >>> z = np.linspace(-1, 1, 11) + >>> expected = -3 * z * np.sqrt(1 - z**2) + >>> np.allclose(assoc_legendre_p(2, 1, z), expected) + True + + Compute the normalized associated Legendre polynomial: + + >>> scale = np.sqrt(5/12) + >>> np.allclose(assoc_legendre_p(2, 1, z, norm=True), scale * expected) + True + """, branch_cut=2, norm=False, diff_n=0 +) + + +@assoc_legendre_p._override_key +def _(branch_cut, norm, diff_n): + diff_n = _nonneg_int_or_fail(diff_n, "diff_n", strict=False) + if not 0 <= diff_n <= 2: + raise ValueError( + "diff_n is currently only implemented for orders 0, 1, and 2," + f" received: {diff_n}." + ) + return norm, diff_n + + +@assoc_legendre_p._override_ufunc_default_args +def _(branch_cut, norm, diff_n): + return branch_cut, + + +@assoc_legendre_p._override_finalize_out +def _(out): + return np.moveaxis(out, -1, 0) + + +assoc_legendre_p_all = MultiUFunc( + assoc_legendre_p_all, + "assoc_legendre_p_all", + """assoc_legendre_p_all(n, m, z, *, branch_cut=2, norm=False, diff_n=0) + + All associated Legendre polynomials of the first kind up to the + specified degree `n`, order `m`, and all derivatives up + to order `diff_n`. + + Parameters + ---------- + n : int + Degree of the associated Legendre polynomials. Must have ``n >= 0``. + m : int + Order of the associated Legendre polynomials. + z : array_like + Input value. + branch_cut : array_like of ints, optional + Selects branch cut. Must be 2 (default) or 3. + 2: cut on the real axis ``|z| > 1`` + 3: cut on the real axis ``-1 < z < 1`` + norm : bool, optional + If ``True``, compute the normalized associated Legendre polynomials. + Default is ``False``. + diff_n : int, optional + A non-negative integer. Compute and return all derivatives up + to order ``diff_n``. Default is 0. + + Returns + ------- + ndarray + Output shape is ``(diff_n + 1, n + 1, 2 * m + 1, ...)``. The entry at + ``(i, j, k)`` corresponds to the ``i``-th derivative, degree ``j``, and + order ``k`` for all ``0 <= i <= diff_n``, ``0 <= j <= n``, and + ``-m <= k <= m``. + + See Also + -------- + assoc_legendre_p + """, branch_cut=2, norm=False, diff_n=0 +) + + +@assoc_legendre_p_all._override_key +def _(branch_cut, norm, diff_n): + if not ((isinstance(diff_n, numbers.Integral)) + and diff_n >= 0): + raise ValueError( + f"diff_n must be a non-negative integer, received: {diff_n}." + ) + if not 0 <= diff_n <= 2: + raise ValueError( + "diff_n is currently only implemented for orders 0, 1, and 2," + f" received: {diff_n}." + ) + return norm, diff_n + + +@assoc_legendre_p_all._override_ufunc_default_args +def _(branch_cut, norm, diff_n): + return branch_cut, + + +@assoc_legendre_p_all._override_ufunc_default_kwargs +def _(branch_cut, norm, diff_n): + return {'axes': [(), ()] + [(0, 1, -1)]} + + +@assoc_legendre_p_all._override_resolve_out_shapes +def _(n, m, z_shape, branch_cut_shape, nout, **kwargs): + diff_n = kwargs['diff_n'] + + if not isinstance(n, numbers.Integral) or (n < 0): + raise ValueError("n must be a non-negative integer.") + if not isinstance(m, numbers.Integral) or (m < 0): + raise ValueError("m must be a non-negative integer.") + + return ((n + 1, 2 * abs(m) + 1) + + np.broadcast_shapes(z_shape, branch_cut_shape) + (diff_n + 1,),) + + +@assoc_legendre_p_all._override_finalize_out +def _(out): + return np.moveaxis(out, -1, 0) + + +legendre_p = MultiUFunc( + legendre_p, + "legendre_p", + """legendre_p(n, z, *, diff_n=0) + + Legendre polynomial of the first kind. + + Parameters + ---------- + n : array_like of ints + Degree of the Legendre polynomial. Must have ``n >= 0``. + z : array_like + Input value. + diff_n : int, optional + A non-negative integer. Compute and return all derivatives up + to order ``diff_n``. Default is 0. + + Returns + ------- + p : ndarray or tuple of ndarray + Legendre polynomial with ``diff_n`` derivatives. + + See Also + -------- + legendre + + References + ---------- + .. [1] Zhang, Shanjie and Jin, Jianming. "Computation of Special + Functions", John Wiley and Sons, 1996. + https://people.sc.fsu.edu/~jburkardt/f77_src/special_functions/special_functions.html + + Examples + -------- + Evaluate the Legendre polynomial :math:`P_3` at :math:`z = 0.5`: + + >>> import numpy as np + >>> from scipy.special import legendre_p + >>> np.allclose(legendre_p(3, 0.5), -0.4375) + True + + Compute the value and first derivative with respect to ``z``: + + >>> p, dp = legendre_p(3, 0.5, diff_n=1) + >>> np.allclose([p, dp], [-0.4375, 0.375]) + True + """, diff_n=0 +) + + +@legendre_p._override_key +def _(diff_n): + if (not isinstance(diff_n, numbers.Integral)) or (diff_n < 0): + raise ValueError( + f"diff_n must be a non-negative integer, received: {diff_n}." + ) + if not 0 <= diff_n <= 2: + raise NotImplementedError( + "diff_n is currently only implemented for orders 0, 1, and 2," + f" received: {diff_n}." + ) + return diff_n + + +@legendre_p._override_finalize_out +def _(out): + return np.moveaxis(out, -1, 0) + + +legendre_p_all = MultiUFunc( + legendre_p_all, + "legendre_p_all", + """legendre_p_all(n, z, *, diff_n=0) + + All Legendre polynomials of the first kind up to the specified degree + `n` and all derivatives up to order `diff_n`. + + Parameters + ---------- + n : int + Degree of the Legendre polynomials. Must have ``n >= 0``. + z : array_like + Input value. + diff_n : int, optional + A non-negative integer. Compute and return all derivatives up + to order ``diff_n``. Default is 0. + + Returns + ------- + ndarray + Output shape is ``(diff_n + 1, n + 1, ...)``. The entry at ``(i, j)`` + corresponds to the ``i``-th derivative and degree ``j`` for all + ``0 <= i <= diff_n`` and ``0 <= j <= n``. + + See Also + -------- + legendre_p + """, diff_n=0 +) + + +@legendre_p_all._override_key +def _(diff_n): + diff_n = _nonneg_int_or_fail(diff_n, "diff_n", strict=False) + if not 0 <= diff_n <= 2: + raise ValueError( + "diff_n is currently only implemented for orders 0, 1, and 2," + f" received: {diff_n}." + ) + return diff_n + + +@legendre_p_all._override_ufunc_default_kwargs +def _(diff_n): + return {'axes': [(), (0, -1)]} + + +@legendre_p_all._override_resolve_out_shapes +def _(n, z_shape, nout, diff_n): + n = _nonneg_int_or_fail(n, 'n', strict=False) + + return nout * ((n + 1,) + z_shape + (diff_n + 1,),) + + +@legendre_p_all._override_finalize_out +def _(out): + return np.moveaxis(out, -1, 0) + + +sph_harm_y = MultiUFunc( + sph_harm_y, + "sph_harm_y", + r"""sph_harm_y(n, m, theta, phi, *, diff_n=0) + + Spherical harmonics. + + They are defined as + + .. math:: + + Y_n^m(\theta,\phi) = \sqrt{\frac{2 n + 1}{4 \pi} \frac{(n - m)!}{(n + m)!}} + P_n^m(\cos(\theta)) e^{i m \phi} + + where :math:`P_n^m` are the (unnormalized) associated Legendre polynomials. + + Parameters + ---------- + n : array_like of ints + Degree of the harmonic. Must have ``n >= 0``. This is + often denoted by ``l`` (lower case L) in descriptions of + spherical harmonics. + m : array_like of ints + Order of the harmonic. + theta : array_like + Polar (colatitudinal) coordinate; must be in ``[0, pi]``. + phi : array_like + Azimuthal (longitudinal) coordinate; must be in ``[0, 2*pi]``. + diff_n : int, optional + A non-negative integer. Compute and return all derivatives up + to order `diff_n`. Default is 0. + + Returns + ------- + ndarray or tuple of ndarray + Spherical harmonics with `diff_n` derivatives. + + Notes + ----- + There are different conventions for the meanings of the input + arguments ``theta`` and ``phi``. In SciPy ``theta`` is the + polar angle and ``phi`` is the azimuthal angle. It is common to + see the opposite convention, that is, ``theta`` as the azimuthal angle + and ``phi`` as the polar angle. + + Note that SciPy's spherical harmonics include the Condon-Shortley + phase [2]_ because it is part of `sph_legendre_p`. + + With SciPy's conventions, the first several spherical harmonics + are + + .. math:: + + Y_0^0(\theta, \phi) &= \frac{1}{2} \sqrt{\frac{1}{\pi}} \\ + Y_1^{-1}(\theta, \phi) &= \frac{1}{2} \sqrt{\frac{3}{2\pi}} + e^{-i\phi} \sin(\theta) \\ + Y_1^0(\theta, \phi) &= \frac{1}{2} \sqrt{\frac{3}{\pi}} + \cos(\theta) \\ + Y_1^1(\theta, \phi) &= -\frac{1}{2} \sqrt{\frac{3}{2\pi}} + e^{i\phi} \sin(\theta). + + References + ---------- + .. [1] Digital Library of Mathematical Functions, 14.30. + https://dlmf.nist.gov/14.30 + .. [2] https://en.wikipedia.org/wiki/Spherical_harmonics#Condon.E2.80.93Shortley_phase + """, force_complex_output=True, diff_n=0 +) + + +@sph_harm_y._override_key +def _(diff_n): + diff_n = _nonneg_int_or_fail(diff_n, "diff_n", strict=False) + if not 0 <= diff_n <= 2: + raise ValueError( + "diff_n is currently only implemented for orders 0, 1, and 2," + f" received: {diff_n}." + ) + return diff_n + + +@sph_harm_y._override_finalize_out +def _(out): + if (out.shape[-1] == 1): + return out[..., 0, 0] + + if (out.shape[-1] == 2): + return out[..., 0, 0], out[..., [1, 0], [0, 1]] + + if (out.shape[-1] == 3): + return (out[..., 0, 0], out[..., [1, 0], [0, 1]], + out[..., [[2, 1], [1, 0]], [[0, 1], [1, 2]]]) + + +sph_harm_y_all = MultiUFunc( + sph_harm_y_all, + "sph_harm_y_all", + """sph_harm_y_all(n, m, theta, phi, *, diff_n=0) + + All spherical harmonics up to the specified degree `n`, order `m`, + and all derivatives up to order `diff_n`. + + Parameters + ---------- + n : int + Degree of the harmonics. Must have ``n >= 0``. This is + often denoted by ``l`` (lower case L) in descriptions of + spherical harmonics. + m : int + Order of the harmonics. + theta : array_like + Polar (colatitudinal) coordinate; must be in ``[0, pi]``. + phi : array_like + Azimuthal (longitudinal) coordinate; must be in ``[0, 2*pi]``. + diff_n : int, optional + A non-negative integer. Compute and return all derivatives up + to order `diff_n`. Default is 0. + + Returns + ------- + ndarray or tuple of ndarray + Returns a tuple of length ``diff_n + 1`` (if ``diff_n > 0``). The first + entry corresponds to the spherical harmonics, the second entry + (if ``diff_n >= 1``) to the gradient, and the third entry + (if ``diff_n >= 2``) to the Hessian matrix. Each entry is an array of + shape ``(n + 1, 2 * m + 1, ...)``, where the entry at ``(i, j)`` + corresponds to degree ``i`` and order ``j`` for all ``0 <= i <= n`` + and ``-m <= j <= m``. + + See Also + -------- + sph_harm_y + """, force_complex_output=True, diff_n=0 +) + + +@sph_harm_y_all._override_key +def _(diff_n): + diff_n = _nonneg_int_or_fail(diff_n, "diff_n", strict=False) + if not 0 <= diff_n <= 2: + raise ValueError( + "diff_n is currently only implemented for orders 2," + f" received: {diff_n}." + ) + return diff_n + + +@sph_harm_y_all._override_ufunc_default_kwargs +def _(diff_n): + return {'axes': [(), ()] + [(0, 1, -2, -1)]} + + +@sph_harm_y_all._override_resolve_out_shapes +def _(n, m, theta_shape, phi_shape, nout, **kwargs): + diff_n = kwargs['diff_n'] + + if not isinstance(n, numbers.Integral) or (n < 0): + raise ValueError("n must be a non-negative integer.") + + return ((n + 1, 2 * abs(m) + 1) + np.broadcast_shapes(theta_shape, phi_shape) + + (diff_n + 1, diff_n + 1),) + + +@sph_harm_y_all._override_finalize_out +def _(out): + if (out.shape[-1] == 1): + return out[..., 0, 0] + + if (out.shape[-1] == 2): + return out[..., 0, 0], out[..., [1, 0], [0, 1]] + + if (out.shape[-1] == 3): + return (out[..., 0, 0], out[..., [1, 0], [0, 1]], + out[..., [[2, 1], [1, 0]], [[0, 1], [1, 2]]]) diff --git a/v12x/pylib/scipy/special/_orthogonal.py b/v12x/pylib/scipy/special/_orthogonal.py new file mode 100644 index 0000000000000000000000000000000000000000..de51d61f6912abd1d28686d2638d1b31527ab836 --- /dev/null +++ b/v12x/pylib/scipy/special/_orthogonal.py @@ -0,0 +1,2856 @@ +""" +A collection of functions to find the weights and abscissas for +Gaussian Quadrature. + +These calculations are done by finding the eigenvalues of a +tridiagonal matrix whose entries are dependent on the coefficients +in the recursion formula for the orthogonal polynomials with the +corresponding weighting function over the interval. + +Many recursion relations for orthogonal polynomials are given: + +.. math:: + + a1n f_{n+1} (x) = (a2n + a3n x ) f_n (x) - a4n f_{n-1} (x) + +The recursion relation of interest is + +.. math:: + + P_{n+1} (x) = (x - A_n) P_n (x) - B_n P_{n-1} (x) + +where :math:`P` has a different normalization than :math:`f`. + +The coefficients can be found as: + +.. math:: + + A_n = -a2n / a3n + \\qquad + B_n = ( a4n / a3n \\sqrt{h_n-1 / h_n})^2 + +where + +.. math:: + + h_n = \\int_a^b w(x) f_n(x)^2 + +assume: + +.. math:: + + P_0 (x) = 1 + \\qquad + P_{-1} (x) == 0 + +For the mathematical background, see [golub.welsch-1969-mathcomp]_ and +[abramowitz.stegun-1965]_. + +References +---------- +.. [golub.welsch-1969-mathcomp] + Golub, Gene H, and John H Welsch. 1969. Calculation of Gauss + Quadrature Rules. *Mathematics of Computation* 23, 221-230+s1--s10. + +.. [abramowitz.stegun-1965] + Abramowitz, Milton, and Irene A Stegun. (1965) *Handbook of + Mathematical Functions: with Formulas, Graphs, and Mathematical + Tables*. Gaithersburg, MD: National Bureau of Standards. + http://www.math.sfu.ca/~cbm/aands/ + +.. [townsend.trogdon.olver-2014] + Townsend, A. and Trogdon, T. and Olver, S. (2014) + *Fast computation of Gauss quadrature nodes and + weights on the whole real line*. :arXiv:`1410.5286`. + +.. [townsend.trogdon.olver-2015] + Townsend, A. and Trogdon, T. and Olver, S. (2015) + *Fast computation of Gauss quadrature nodes and + weights on the whole real line*. + IMA Journal of Numerical Analysis + :doi:`10.1093/imanum/drv002`. +""" +# +# Author: Travis Oliphant 2000 +# Updated Sep. 2003 (fixed bugs --- tested to be accurate) + +# SciPy imports. +import numpy as np +from numpy import (exp, inf, pi, sqrt, floor, sin, cos, around, + hstack, arccos, arange) +from scipy.special import airy + +# Local imports. +# There is no .pyi file for _specfun +from . import _specfun # type: ignore +from . import _ufuncs +_gam = _ufuncs.gamma + +_polyfuns = ['legendre', 'chebyt', 'chebyu', 'chebyc', 'chebys', + 'jacobi', 'laguerre', 'genlaguerre', 'hermite', + 'hermitenorm', 'gegenbauer', 'sh_legendre', 'sh_chebyt', + 'sh_chebyu', 'sh_jacobi'] + +# Correspondence between new and old names of root functions +_rootfuns_map = {'roots_legendre': 'p_roots', + 'roots_chebyt': 't_roots', + 'roots_chebyu': 'u_roots', + 'roots_chebyc': 'c_roots', + 'roots_chebys': 's_roots', + 'roots_jacobi': 'j_roots', + 'roots_laguerre': 'l_roots', + 'roots_genlaguerre': 'la_roots', + 'roots_hermite': 'h_roots', + 'roots_hermitenorm': 'he_roots', + 'roots_gegenbauer': 'cg_roots', + 'roots_sh_legendre': 'ps_roots', + 'roots_sh_chebyt': 'ts_roots', + 'roots_sh_chebyu': 'us_roots', + 'roots_sh_jacobi': 'js_roots'} + +__all__ = _polyfuns + list(_rootfuns_map.keys()) + + +class orthopoly1d(np.poly1d): + + def __init__(self, roots, weights=None, hn=1.0, kn=1.0, wfunc=None, + limits=None, monic=False, eval_func=None): + equiv_weights = [weights[k] / wfunc(roots[k]) for + k in range(len(roots))] + mu = sqrt(hn) + if monic: + evf = eval_func + if evf: + knn = kn + def eval_func(x): + return evf(x) / knn + mu = mu / abs(kn) + kn = 1.0 + + # compute coefficients from roots, then scale + poly = np.poly1d(roots, r=True) + np.poly1d.__init__(self, poly.coeffs * float(kn)) + + self.weights = np.array(list(zip(roots, weights, equiv_weights))) + self.weight_func = wfunc + self.limits = limits + self.normcoef = mu + + # Note: eval_func will be discarded on arithmetic + self._eval_func = eval_func + + def __call__(self, v): + if self._eval_func and not isinstance(v, np.poly1d): + return self._eval_func(v) + else: + return np.poly1d.__call__(self, v) + + def _scale(self, p): + if p == 1.0: + return + self._coeffs *= p + + evf = self._eval_func + if evf: + self._eval_func = lambda x: evf(x) * p + self.normcoef *= p + + +def _gen_roots_and_weights(n, mu0, an_func, bn_func, f, df, symmetrize, mu): + """[x,w] = gen_roots_and_weights(n,an_func,sqrt_bn_func,mu) + + Returns the roots (x) of an nth order orthogonal polynomial, + and weights (w) to use in appropriate Gaussian quadrature with that + orthogonal polynomial. + + The polynomials have the recurrence relation + P_n+1(x) = (x - A_n) P_n(x) - B_n P_n-1(x) + + an_func(n) should return A_n + sqrt_bn_func(n) should return sqrt(B_n) + mu ( = h_0 ) is the integral of the weight over the orthogonal + interval + """ + # lazy import to prevent to prevent linalg dependency for whole module (gh-23420) + from scipy import linalg + k = np.arange(n, dtype='d') + c = np.zeros((2, n)) + c[0,1:] = bn_func(k[1:]) + c[1,:] = an_func(k) + x = linalg.eigvals_banded(c, overwrite_a_band=True) + + # improve roots by one application of Newton's method + y = f(n, x) + dy = df(n, x) + x -= y/dy + + # fm and dy may contain very large/small values, so we + # log-normalize them to maintain precision in the product fm*dy + fm = f(n-1, x) + log_fm = np.log(np.abs(fm)) + log_dy = np.log(np.abs(dy)) + fm /= np.exp((log_fm.max() + log_fm.min()) / 2.) + dy /= np.exp((log_dy.max() + log_dy.min()) / 2.) + w = 1.0 / (fm * dy) + + if symmetrize: + w = (w + w[::-1]) / 2 + x = (x - x[::-1]) / 2 + + w *= mu0 / w.sum() + + if mu: + return x, w, mu0 + else: + return x, w + +# Jacobi Polynomials 1 P^(alpha,beta)_n(x) + + +def roots_jacobi(n, alpha, beta, mu=False): + r"""Gauss-Jacobi quadrature. + + Compute the sample points and weights for Gauss-Jacobi + quadrature. The sample points are the roots of the nth degree + Jacobi polynomial, :math:`P^{\alpha, \beta}_n(x)`. These sample + points and weights correctly integrate polynomials of degree + :math:`2n - 1` or less over the interval :math:`[-1, 1]` with + weight function :math:`w(x) = (1 - x)^{\alpha} (1 + + x)^{\beta}`. See 22.2.1 in [AS]_ for details. + + Parameters + ---------- + n : int + Quadrature order. + alpha : float + alpha must be > -1. + beta : float + beta must be > -1. + mu : bool, optional + If True, return the sum of the weights in addition to sample points and weights. + + Returns + ------- + x : ndarray + Sample points. + w : ndarray + Weights. + mu : float, optional + Sum of the weights, only returned if `mu=True`. + + See Also + -------- + scipy.integrate.fixed_quad + + References + ---------- + .. [AS] Milton Abramowitz and Irene A. Stegun, eds. + Handbook of Mathematical Functions with Formulas, + Graphs, and Mathematical Tables. New York: Dover, 1972. + + Examples + -------- + >>> from scipy.special import roots_jacobi + >>> x, w = roots_jacobi(3, 0.5, 0.5) + >>> x + array([-0.70710678, 0. , 0.70710678]) + >>> w + array([0.39269908, 0.78539816, 0.39269908]) + + >>> x, w, mu = roots_jacobi(3, 0.5, 0.5, mu=True) + >>> mu + 1.5707963267948966 # Sum of weights, equals pi/2 for alpha = beta = 0.5 + """ + + m = int(n) + if n < 1 or n != m: + raise ValueError("n must be a positive integer.") + if alpha <= -1 or beta <= -1: + raise ValueError("alpha and beta must be greater than -1.") + + if alpha == 0.0 and beta == 0.0: + return roots_legendre(m, mu) + if alpha == beta: + return roots_gegenbauer(m, alpha+0.5, mu) + + if (alpha + beta) <= 1000: + mu0 = 2.0**(alpha+beta+1) * _ufuncs.beta(alpha+1, beta+1) + else: + # Avoid overflows in pow and beta for very large parameters + mu0 = np.exp((alpha + beta + 1) * np.log(2.0) + + _ufuncs.betaln(alpha+1, beta+1)) + a = alpha + b = beta + if a + b == 0.0: + def an_func(k): + return np.where(k == 0, (b - a) / (2 + a + b), 0.0) + else: + def an_func(k): + return np.where( + k == 0, + (b - a) / (2 + a + b), + (b * b - a * a) / ((2.0 * k + a + b) * (2.0 * k + a + b + 2)) + ) + + def bn_func(k): + return ( + 2.0 / (2.0 * k + a + b) + * np.sqrt((k + a) * (k + b) / (2 * k + a + b + 1)) + * np.where(k == 1, 1.0, np.sqrt(k * (k + a + b) / (2.0 * k + a + b - 1))) + ) + + def f(n, x): + return _ufuncs.eval_jacobi(n, a, b, x) + def df(n, x): + return 0.5 * (n + a + b + 1) * _ufuncs.eval_jacobi(n - 1, a + 1, b + 1, x) + return _gen_roots_and_weights(m, mu0, an_func, bn_func, f, df, False, mu) + + +def jacobi(n, alpha, beta, monic=False): + r"""Jacobi polynomial. + + Defined to be the solution of + + .. math:: + + \begin{aligned} + (1 - x^2)\frac{d^2}{dx^2} P_n^{(\alpha, \beta)}(x) + &+ \left(\beta - \alpha - (\alpha + \beta + 2)x\right) + \frac{d}{dx} P_n^{(\alpha, \beta)}(x) \\ + &+ n(n + \alpha + \beta + 1) P_n^{(\alpha, \beta)}(x) = 0 + \end{aligned} + + for :math:`\alpha, \beta > -1`; :math:`P_n^{(\alpha, \beta)}` is a + polynomial of degree :math:`n`. + + Parameters + ---------- + n : int + Degree of the polynomial. + alpha : float + Parameter, must be greater than -1. + beta : float + Parameter, must be greater than -1. + monic : bool, optional + If `True`, scale the leading coefficient to be 1. Default is + `False`. + + Returns + ------- + P : orthopoly1d + Jacobi polynomial. + + Notes + ----- + For fixed :math:`\alpha, \beta`, the polynomials + :math:`P_n^{(\alpha, \beta)}` are orthogonal over :math:`[-1, 1]` + with weight function :math:`(1 - x)^\alpha(1 + x)^\beta`. + + References + ---------- + .. [AS] Milton Abramowitz and Irene A. Stegun, eds. + Handbook of Mathematical Functions with Formulas, + Graphs, and Mathematical Tables. New York: Dover, 1972. + + Examples + -------- + The Jacobi polynomials satisfy the recurrence relation: + + .. math:: + P_n^{(\alpha, \beta-1)}(x) - P_n^{(\alpha-1, \beta)}(x) + = P_{n-1}^{(\alpha, \beta)}(x) + + This can be verified, for example, for :math:`\alpha = \beta = 2` + and :math:`n = 1` over the interval :math:`[-1, 1]`: + + >>> import numpy as np + >>> from scipy.special import jacobi + >>> x = np.arange(-1.0, 1.0, 0.01) + >>> np.allclose(jacobi(0, 2, 2)(x), + ... jacobi(1, 2, 1)(x) - jacobi(1, 1, 2)(x)) + True + + Plot of the Jacobi polynomial :math:`P_5^{(\alpha, -0.5)}` for + different values of :math:`\alpha`: + + >>> import matplotlib.pyplot as plt + >>> x = np.arange(-1.0, 1.0, 0.01) + >>> fig, ax = plt.subplots() + >>> ax.set_ylim(-2.0, 2.0) + >>> ax.set_title(r'Jacobi polynomials $P_5^{(\alpha, -0.5)}$') + >>> for alpha in np.arange(0, 4, 1): + ... ax.plot(x, jacobi(5, alpha, -0.5)(x), label=rf'$\alpha={alpha}$') + >>> plt.legend(loc='best') + >>> plt.show() + + """ + if n < 0: + raise ValueError("n must be nonnegative.") + + def wfunc(x): + return (1 - x) ** alpha * (1 + x) ** beta + if n == 0: + return orthopoly1d([], [], 1.0, 1.0, wfunc, (-1, 1), monic, + eval_func=np.ones_like) + x, w, mu = roots_jacobi(n, alpha, beta, mu=True) + ab1 = alpha + beta + 1.0 + hn = 2**ab1 / (2 * n + ab1) * _gam(n + alpha + 1) + hn *= _gam(n + beta + 1.0) / _gam(n + 1) / _gam(n + ab1) + kn = _gam(2 * n + ab1) / 2.0**n / _gam(n + 1) / _gam(n + ab1) + # here kn = coefficient on x^n term + p = orthopoly1d(x, w, hn, kn, wfunc, (-1, 1), monic, + lambda x: _ufuncs.eval_jacobi(n, alpha, beta, x)) + return p + +# Jacobi Polynomials shifted G_n(p,q,x) + + +def roots_sh_jacobi(n, p1, q1, mu=False): + """Gauss-Jacobi (shifted) quadrature. + + Compute the sample points and weights for Gauss-Jacobi (shifted) + quadrature. The sample points are the roots of the nth degree + shifted Jacobi polynomial, :math:`G^{p,q}_n(x)`. These sample + points and weights correctly integrate polynomials of degree + :math:`2n - 1` or less over the interval :math:`[0, 1]` with + weight function :math:`w(x) = (1 - x)^{p-q} x^{q-1}`. See 22.2.2 + in [AS]_ for details. + + Parameters + ---------- + n : int + quadrature order. + p1 : float + (p1 - q1) must be > -1. + q1 : float + q1 must be > 0. + mu : bool, optional + If True, return the sum of the weights, optional. + + Returns + ------- + x : ndarray + Sample points + w : ndarray + Weights + mu : float + Sum of the weights + + See Also + -------- + scipy.integrate.fixed_quad + + References + ---------- + .. [AS] Milton Abramowitz and Irene A. Stegun, eds. + Handbook of Mathematical Functions with Formulas, + Graphs, and Mathematical Tables. New York: Dover, 1972. + + """ + if (p1-q1) <= -1 or q1 <= 0: + message = "(p - q) must be greater than -1, and q must be greater than 0." + raise ValueError(message) + x, w, m = roots_jacobi(n, p1-q1, q1-1, True) + x = (x + 1) / 2 + scale = 2.0**p1 + w /= scale + m /= scale + if mu: + return x, w, m + else: + return x, w + + +def sh_jacobi(n, p, q, monic=False): + r"""Shifted Jacobi polynomial. + + Defined by + + .. math:: + + G_n^{(p, q)}(x) + = \binom{2n + p - 1}{n}^{-1}P_n^{(p - q, q - 1)}(2x - 1), + + where :math:`P_n^{(\cdot, \cdot)}` is the nth Jacobi polynomial. + + Parameters + ---------- + n : int + Degree of the polynomial. + p : float + Parameter, must have :math:`p > q - 1`. + q : float + Parameter, must be greater than 0. + monic : bool, optional + If `True`, scale the leading coefficient to be 1. Default is + `False`. + + Returns + ------- + G : orthopoly1d + Shifted Jacobi polynomial. + + Notes + ----- + For fixed :math:`p, q`, the polynomials :math:`G_n^{(p, q)}` are + orthogonal over :math:`[0, 1]` with weight function :math:`(1 - + x)^{p - q}x^{q - 1}`. + + Examples + -------- + Evaluate the shifted Jacobi polynomial :math:`G_3^{(2, 1)}` at + :math:`x = 0.5`: + + >>> import numpy as np + >>> from scipy.special import binom, jacobi, sh_jacobi + >>> np.isclose(sh_jacobi(3, 2, 1)(0.5), -3/280) + True + + The polynomial is related to the Jacobi polynomial + :math:`P_n^{(p - q, q - 1)}`: + + >>> x = np.linspace(0, 1, 5) + >>> n, p, q = 3, 2, 1 + >>> scale = 1 / binom(2*n + p - 1, n) + >>> np.allclose(sh_jacobi(n, p, q)(x), + ... scale * jacobi(n, p - q, q - 1)(2*x - 1)) + True + + Plot :math:`G_3^{(p, 1)}` for several values of :math:`p`: + + >>> import matplotlib.pyplot as plt + >>> x = np.linspace(0, 1, 400) + >>> fig, ax = plt.subplots() + >>> for p in [1, 2, 3]: + ... ax.plot(x, sh_jacobi(3, p, 1)(x), label=rf"$p={p}$") + >>> ax.set_title(r"Shifted Jacobi polynomials $G_3^{(p, 1)}$") + >>> ax.set_xlabel("x") + >>> ax.legend(loc="best") + >>> plt.show() + + """ + if n < 0: + raise ValueError("n must be nonnegative.") + + def wfunc(x): + return (1.0 - x) ** (p - q) * x ** (q - 1.0) + if n == 0: + return orthopoly1d([], [], 1.0, 1.0, wfunc, (-1, 1), monic, + eval_func=np.ones_like) + n1 = n + x, w = roots_sh_jacobi(n1, p, q) + hn = _gam(n + 1) * _gam(n + q) * _gam(n + p) * _gam(n + p - q + 1) + hn /= (2 * n + p) * (_gam(2 * n + p)**2) + # kn = 1.0 in standard form so monic is redundant. Kept for compatibility. + kn = 1.0 + pp = orthopoly1d(x, w, hn, kn, wfunc=wfunc, limits=(0, 1), monic=monic, + eval_func=lambda x: _ufuncs.eval_sh_jacobi(n, p, q, x)) + return pp + +# Generalized Laguerre L^(alpha)_n(x) + + +def roots_genlaguerre(n, alpha, mu=False): + r"""Gauss-generalized Laguerre quadrature. + + Compute the sample points and weights for Gauss-generalized + Laguerre quadrature. The sample points are the roots of the nth + degree generalized Laguerre polynomial, :math:`L^{\alpha}_n(x)`. + These sample points and weights correctly integrate polynomials of + degree :math:`2n - 1` or less over the interval :math:`[0, + \infty]` with weight function :math:`w(x) = x^{\alpha} + e^{-x}`. See 22.3.9 in [AS]_ for details. + + Parameters + ---------- + n : int + quadrature order + alpha : float + alpha must be > -1 + mu : bool, optional + If True, return the sum of the weights, optional. + + Returns + ------- + x : ndarray + Sample points + w : ndarray + Weights + mu : float + Sum of the weights + + See Also + -------- + scipy.integrate.fixed_quad + + References + ---------- + .. [AS] Milton Abramowitz and Irene A. Stegun, eds. + Handbook of Mathematical Functions with Formulas, + Graphs, and Mathematical Tables. New York: Dover, 1972. + + """ + m = int(n) + if n < 1 or n != m: + raise ValueError("n must be a positive integer.") + if alpha < -1: + raise ValueError("alpha must be greater than -1.") + + mu0 = _ufuncs.gamma(alpha + 1) + + if m == 1: + x = np.array([alpha+1.0], 'd') + w = np.array([mu0], 'd') + if mu: + return x, w, mu0 + else: + return x, w + + def an_func(k): + return 2 * k + alpha + 1 + def bn_func(k): + return -np.sqrt(k * (k + alpha)) + def f(n, x): + return _ufuncs.eval_genlaguerre(n, alpha, x) + def df(n, x): + return (n * _ufuncs.eval_genlaguerre(n, alpha, x) + - (n + alpha) * _ufuncs.eval_genlaguerre(n - 1, alpha, x)) / x + return _gen_roots_and_weights(m, mu0, an_func, bn_func, f, df, False, mu) + + +def genlaguerre(n, alpha, monic=False): + r"""Generalized (associated) Laguerre polynomial. + + Defined to be the solution of + + .. math:: + x\frac{d^2}{dx^2}L_n^{(\alpha)}(x) + + (\alpha + 1 - x)\frac{d}{dx}L_n^{(\alpha)}(x) + + nL_n^{(\alpha)}(x) = 0, + + where :math:`\alpha > -1`; :math:`L_n^{(\alpha)}` is a polynomial + of degree :math:`n`. + + Parameters + ---------- + n : int + Degree of the polynomial. + alpha : float + Parameter, must be greater than -1. + monic : bool, optional + If `True`, scale the leading coefficient to be 1. Default is + `False`. + + Returns + ------- + L : orthopoly1d + Generalized Laguerre polynomial. + + See Also + -------- + laguerre : Laguerre polynomial. + hyp1f1 : confluent hypergeometric function + + Notes + ----- + For fixed :math:`\alpha`, the polynomials :math:`L_n^{(\alpha)}` + are orthogonal over :math:`[0, \infty)` with weight function + :math:`e^{-x}x^\alpha`. + + The Laguerre polynomials are the special case where :math:`\alpha + = 0`. + + References + ---------- + .. [AS] Milton Abramowitz and Irene A. Stegun, eds. + Handbook of Mathematical Functions with Formulas, + Graphs, and Mathematical Tables. New York: Dover, 1972. + + Examples + -------- + The generalized Laguerre polynomials are closely related to the confluent + hypergeometric function :math:`{}_1F_1`: + + .. math:: + L_n^{(\alpha)}(x) = \binom{n + \alpha}{n} {}_1F_1(-n, \alpha +1, x) + + This can be verified, for example, for :math:`n = \alpha = 3` over the + interval :math:`[-1, 1]`: + + >>> import numpy as np + >>> from scipy.special import binom, genlaguerre, hyp1f1 + >>> x = np.arange(-1.0, 1.0, 0.01) + >>> np.allclose(genlaguerre(3, 3)(x), binom(6, 3) * hyp1f1(-3, 4, x)) + True + + This is the plot of the generalized Laguerre polynomials + :math:`L_3^{(\alpha)}` for some values of :math:`\alpha`: + + >>> import matplotlib.pyplot as plt + >>> x = np.arange(-4.0, 12.0, 0.01) + >>> fig, ax = plt.subplots() + >>> ax.set_ylim(-5.0, 10.0) + >>> ax.set_title(r'Generalized Laguerre polynomials $L_3^{(\alpha)}$') + >>> for alpha in np.arange(0, 5): + ... ax.plot(x, genlaguerre(3, alpha)(x), label=rf"$L_3^{{({alpha})}}$") + >>> plt.legend(loc='best') + >>> plt.show() + + """ + if alpha <= -1: + raise ValueError("alpha must be > -1") + if n < 0: + raise ValueError("n must be nonnegative.") + + if n == 0: + n1 = n + 1 + else: + n1 = n + x, w = roots_genlaguerre(n1, alpha) + def wfunc(x): + return exp(-x) * x ** alpha + if n == 0: + x, w = [], [] + hn = _gam(n + alpha + 1) / _gam(n + 1) + kn = (-1)**n / _gam(n + 1) + p = orthopoly1d(x, w, hn, kn, wfunc, (0, inf), monic, + lambda x: _ufuncs.eval_genlaguerre(n, alpha, x)) + return p + +# Laguerre L_n(x) + + +def roots_laguerre(n, mu=False): + r"""Gauss-Laguerre quadrature. + + Compute the sample points and weights for Gauss-Laguerre + quadrature. The sample points are the roots of the nth degree + Laguerre polynomial, :math:`L_n(x)`. These sample points and + weights correctly integrate polynomials of degree :math:`2n - 1` + or less over the interval :math:`[0, \infty]` with weight function + :math:`w(x) = e^{-x}`. See 22.2.13 in [AS]_ for details. + + Parameters + ---------- + n : int + quadrature order + mu : bool, optional + If True, return the sum of the weights, optional. + + Returns + ------- + x : ndarray + Sample points + w : ndarray + Weights + mu : float + Sum of the weights + + See Also + -------- + scipy.integrate.fixed_quad + numpy.polynomial.laguerre.laggauss + + References + ---------- + .. [AS] Milton Abramowitz and Irene A. Stegun, eds. + Handbook of Mathematical Functions with Formulas, + Graphs, and Mathematical Tables. New York: Dover, 1972. + + """ + return roots_genlaguerre(n, 0.0, mu=mu) + + +def laguerre(n, monic=False): + r"""Laguerre polynomial. + + Defined to be the solution of + + .. math:: + x\frac{d^2}{dx^2}L_n(x) + (1 - x)\frac{d}{dx}L_n(x) + + nL_n(x) = 0; + + :math:`L_n` is a polynomial of degree :math:`n`. + + Parameters + ---------- + n : int + Degree of the polynomial. + monic : bool, optional + If `True`, scale the leading coefficient to be 1. Default is + `False`. + + Returns + ------- + L : orthopoly1d + Laguerre Polynomial. + + See Also + -------- + genlaguerre : Generalized (associated) Laguerre polynomial. + + Notes + ----- + The polynomials :math:`L_n` are orthogonal over :math:`[0, + \infty)` with weight function :math:`e^{-x}`. + + References + ---------- + .. [AS] Milton Abramowitz and Irene A. Stegun, eds. + Handbook of Mathematical Functions with Formulas, + Graphs, and Mathematical Tables. New York: Dover, 1972. + + Examples + -------- + The Laguerre polynomials :math:`L_n` are the special case + :math:`\alpha = 0` of the generalized Laguerre polynomials + :math:`L_n^{(\alpha)}`. + Let's verify it on the interval :math:`[-1, 1]`: + + >>> import numpy as np + >>> from scipy.special import genlaguerre + >>> from scipy.special import laguerre + >>> x = np.arange(-1.0, 1.0, 0.01) + >>> np.allclose(genlaguerre(3, 0)(x), laguerre(3)(x)) + True + + The polynomials :math:`L_n` also satisfy the recurrence relation: + + .. math:: + (n + 1)L_{n+1}(x) = (2n +1 -x)L_n(x) - nL_{n-1}(x) + + This can be easily checked on :math:`[0, 1]` for :math:`n = 3`: + + >>> x = np.arange(0.0, 1.0, 0.01) + >>> np.allclose(4 * laguerre(4)(x), + ... (7 - x) * laguerre(3)(x) - 3 * laguerre(2)(x)) + True + + This is the plot of the first few Laguerre polynomials :math:`L_n`: + + >>> import matplotlib.pyplot as plt + >>> x = np.arange(-1.0, 5.0, 0.01) + >>> fig, ax = plt.subplots() + >>> ax.set_ylim(-5.0, 5.0) + >>> ax.set_title(r'Laguerre polynomials $L_n$') + >>> for n in np.arange(0, 5): + ... ax.plot(x, laguerre(n)(x), label=rf'$L_{n}$') + >>> plt.legend(loc='best') + >>> plt.show() + + """ + if n < 0: + raise ValueError("n must be nonnegative.") + + if n == 0: + n1 = n + 1 + else: + n1 = n + x, w = roots_laguerre(n1) + if n == 0: + x, w = [], [] + hn = 1.0 + kn = (-1)**n / _gam(n + 1) + p = orthopoly1d(x, w, hn, kn, lambda x: exp(-x), (0, inf), monic, + lambda x: _ufuncs.eval_laguerre(n, x)) + return p + +# Hermite 1 H_n(x) + + +def roots_hermite(n, mu=False): + r"""Gauss-Hermite (physicist's) quadrature. + + Compute the sample points and weights for Gauss-Hermite + quadrature. The sample points are the roots of the nth degree + Hermite polynomial, :math:`H_n(x)`. These sample points and + weights correctly integrate polynomials of degree :math:`2n - 1` + or less over the interval :math:`[-\infty, \infty]` with weight + function :math:`w(x) = e^{-x^2}`. See 22.2.14 in [AS]_ for + details. + + Parameters + ---------- + n : int + quadrature order + mu : bool, optional + If True, return the sum of the weights, optional. + + Returns + ------- + x : ndarray + Sample points + w : ndarray + Weights + mu : float + Sum of the weights + + See Also + -------- + scipy.integrate.fixed_quad + numpy.polynomial.hermite.hermgauss + roots_hermitenorm + + Notes + ----- + For small n up to 150 a modified version of the Golub-Welsch + algorithm is used. Nodes are computed from the eigenvalue + problem and improved by one step of a Newton iteration. + The weights are computed from the well-known analytical formula. + + For n larger than 150 an optimal asymptotic algorithm is applied + which computes nodes and weights in a numerically stable manner. + The algorithm has linear runtime making computation for very + large n (several thousand or more) feasible. + + References + ---------- + .. [townsend.trogdon.olver-2014] + Townsend, A. and Trogdon, T. and Olver, S. (2014) + *Fast computation of Gauss quadrature nodes and + weights on the whole real line*. :arXiv:`1410.5286`. + .. [townsend.trogdon.olver-2015] + Townsend, A. and Trogdon, T. and Olver, S. (2015) + *Fast computation of Gauss quadrature nodes and + weights on the whole real line*. + IMA Journal of Numerical Analysis + :doi:`10.1093/imanum/drv002`. + .. [AS] Milton Abramowitz and Irene A. Stegun, eds. + Handbook of Mathematical Functions with Formulas, + Graphs, and Mathematical Tables. New York: Dover, 1972. + + """ + m = int(n) + if n < 1 or n != m: + raise ValueError("n must be a positive integer.") + + mu0 = np.sqrt(np.pi) + if n <= 150: + def an_func(k): + return 0.0 * k + def bn_func(k): + return np.sqrt(k / 2.0) + f = _ufuncs.eval_hermite + def df(n, x): + return 2.0 * n * _ufuncs.eval_hermite(n - 1, x) + return _gen_roots_and_weights(m, mu0, an_func, bn_func, f, df, True, mu) + else: + nodes, weights = _roots_hermite_asy(m) + if mu: + return nodes, weights, mu0 + else: + return nodes, weights + + +def _compute_tauk(n, k, maxit=5): + """Helper function for Tricomi initial guesses + + For details, see formula 3.1 in lemma 3.1 in the + original paper. + + Parameters + ---------- + n : int + Quadrature order + k : ndarray of type int + Index of roots :math:`\tau_k` to compute + maxit : int + Number of Newton maxit performed, the default + value of 5 is sufficient. + + Returns + ------- + tauk : ndarray + Roots of equation 3.1 + + See Also + -------- + initial_nodes_a + roots_hermite_asy + """ + a = n % 2 - 0.5 + c = (4.0*floor(n/2.0) - 4.0*k + 3.0)*pi / (4.0*floor(n/2.0) + 2.0*a + 2.0) + def f(x): + return x - sin(x) - c + def df(x): + return 1.0 - cos(x) + xi = 0.5*pi + for i in range(maxit): + xi = xi - f(xi)/df(xi) + return xi + + +def _initial_nodes_a(n, k): + r"""Tricomi initial guesses + + Computes an initial approximation to the square of the `k`-th + (positive) root :math:`x_k` of the Hermite polynomial :math:`H_n` + of order :math:`n`. The formula is the one from lemma 3.1 in the + original paper. The guesses are accurate except in the region + near :math:`\sqrt{2n + 1}`. + + Parameters + ---------- + n : int + Quadrature order + k : ndarray of type int + Index of roots to compute + + Returns + ------- + xksq : ndarray + Square of the approximate roots + + See Also + -------- + initial_nodes + roots_hermite_asy + """ + tauk = _compute_tauk(n, k) + sigk = cos(0.5*tauk)**2 + a = n % 2 - 0.5 + nu = 4.0*floor(n/2.0) + 2.0*a + 2.0 + # Initial approximation of Hermite roots (square) + xksq = nu*sigk - 1.0/(3.0*nu) * (5.0/(4.0*(1.0-sigk)**2) - 1.0/(1.0-sigk) - 0.25) + return xksq + + +def _initial_nodes_b(n, k): + r"""Gatteschi initial guesses + + Computes an initial approximation to the square of the kth + (positive) root :math:`x_k` of the Hermite polynomial :math:`H_n` + of order :math:`n`. The formula is the one from lemma 3.2 in the + original paper. The guesses are accurate in the region just + below :math:`\sqrt{2n + 1}`. + + Parameters + ---------- + n : int + Quadrature order + k : ndarray of type int + Index of roots to compute + + Returns + ------- + xksq : ndarray + Square of the approximate root + + See Also + -------- + initial_nodes + roots_hermite_asy + """ + a = n % 2 - 0.5 + nu = 4.0*floor(n/2.0) + 2.0*a + 2.0 + # Airy roots by approximation + ak = _specfun.airyzo(k.max(), 1)[0][::-1] + # Initial approximation of Hermite roots (square) + xksq = (nu + + 2.0**(2.0/3.0) * ak * nu**(1.0/3.0) + + 1.0/5.0 * 2.0**(4.0/3.0) * ak**2 * nu**(-1.0/3.0) + + (9.0/140.0 - 12.0/175.0 * ak**3) * nu**(-1.0) + + (16.0/1575.0 * ak + 92.0/7875.0 * ak**4) * 2.0**(2.0/3.0) * nu**(-5.0/3.0) + - (15152.0/3031875.0 * ak**5 + 1088.0/121275.0 * ak**2) + * 2.0**(1.0/3.0) * nu**(-7.0/3.0)) + return xksq + + +def _initial_nodes(n): + """Initial guesses for the Hermite roots + + Computes an initial approximation to the non-negative + roots :math:`x_k` of the Hermite polynomial :math:`H_n` + of order :math:`n`. The Tricomi and Gatteschi initial + guesses are used in the region where they are accurate. + + Parameters + ---------- + n : int + Quadrature order + + Returns + ------- + xk : ndarray + Approximate roots + + See Also + -------- + roots_hermite_asy + """ + # Turnover point + # linear polynomial fit to error of 10, 25, 40, ..., 1000 point rules + fit = 0.49082003*n - 4.37859653 + turnover = around(fit).astype(int) + # Compute all approximations + ia = arange(1, int(floor(n*0.5)+1)) + ib = ia[::-1] + xasq = _initial_nodes_a(n, ia[:turnover+1]) + xbsq = _initial_nodes_b(n, ib[turnover+1:]) + # Combine + iv = sqrt(hstack([xasq, xbsq])) + # Central node is always zero + if n % 2 == 1: + iv = hstack([0.0, iv]) + return iv + + +def _pbcf(n, theta): + r"""Asymptotic series expansion of parabolic cylinder function + + The implementation is based on sections 3.2 and 3.3 from the + original paper. Compared to the published version this code + adds one more term to the asymptotic series. The detailed + formulas can be found at [parabolic-asymptotics]_. The evaluation + is done in a transformed variable :math:`\theta := \arccos(t)` + where :math:`t := x / \mu` and :math:`\mu := \sqrt{2n + 1}`. + + Parameters + ---------- + n : int + Quadrature order + theta : ndarray + Transformed position variable + + Returns + ------- + U : ndarray + Value of the parabolic cylinder function :math:`U(a, \theta)`. + Ud : ndarray + Value of the derivative :math:`U^{\prime}(a, \theta)` of + the parabolic cylinder function. + + See Also + -------- + roots_hermite_asy + + References + ---------- + .. [parabolic-asymptotics] + https://dlmf.nist.gov/12.10#vii + """ + st = sin(theta) + ct = cos(theta) + # https://dlmf.nist.gov/12.10#vii + mu = 2.0*n + 1.0 + # https://dlmf.nist.gov/12.10#E23 + eta = 0.5*theta - 0.5*st*ct + # https://dlmf.nist.gov/12.10#E39 + zeta = -(3.0*eta/2.0) ** (2.0/3.0) + # https://dlmf.nist.gov/12.10#E40 + phi = (-zeta / st**2) ** (0.25) + # Coefficients + # https://dlmf.nist.gov/12.10#E43 + a0 = 1.0 + a1 = 0.10416666666666666667 + a2 = 0.08355034722222222222 + a3 = 0.12822657455632716049 + a4 = 0.29184902646414046425 + a5 = 0.88162726744375765242 + b0 = 1.0 + b1 = -0.14583333333333333333 + b2 = -0.09874131944444444444 + b3 = -0.14331205391589506173 + b4 = -0.31722720267841354810 + b5 = -0.94242914795712024914 + # Polynomials + # https://dlmf.nist.gov/12.10#E9 + # https://dlmf.nist.gov/12.10#E10 + ctp = ct ** arange(16).reshape((-1,1)) + u0 = 1.0 + u1 = (1.0*ctp[3,:] - 6.0*ct) / 24.0 + u2 = (-9.0*ctp[4,:] + 249.0*ctp[2,:] + 145.0) / 1152.0 + u3 = (-4042.0*ctp[9,:] + 18189.0*ctp[7,:] - 28287.0*ctp[5,:] + - 151995.0*ctp[3,:] - 259290.0*ct) / 414720.0 + u4 = (72756.0*ctp[10,:] - 321339.0*ctp[8,:] - 154982.0*ctp[6,:] + + 50938215.0*ctp[4,:] + 122602962.0*ctp[2,:] + 12773113.0) / 39813120.0 + u5 = (82393456.0*ctp[15,:] - 617950920.0*ctp[13,:] + 1994971575.0*ctp[11,:] + - 3630137104.0*ctp[9,:] + 4433574213.0*ctp[7,:] - 37370295816.0*ctp[5,:] + - 119582875013.0*ctp[3,:] - 34009066266.0*ct) / 6688604160.0 + v0 = 1.0 + v1 = (1.0*ctp[3,:] + 6.0*ct) / 24.0 + v2 = (15.0*ctp[4,:] - 327.0*ctp[2,:] - 143.0) / 1152.0 + v3 = (-4042.0*ctp[9,:] + 18189.0*ctp[7,:] - 36387.0*ctp[5,:] + + 238425.0*ctp[3,:] + 259290.0*ct) / 414720.0 + v4 = (-121260.0*ctp[10,:] + 551733.0*ctp[8,:] - 151958.0*ctp[6,:] + - 57484425.0*ctp[4,:] - 132752238.0*ctp[2,:] - 12118727) / 39813120.0 + v5 = (82393456.0*ctp[15,:] - 617950920.0*ctp[13,:] + 2025529095.0*ctp[11,:] + - 3750839308.0*ctp[9,:] + 3832454253.0*ctp[7,:] + 35213253348.0*ctp[5,:] + + 130919230435.0*ctp[3,:] + 34009066266*ct) / 6688604160.0 + # Airy Evaluation (Bi and Bip unused) + Ai, Aip, Bi, Bip = airy(mu**(4.0/6.0) * zeta) + # Prefactor for U + P = 2.0*sqrt(pi) * mu**(1.0/6.0) * phi + # Terms for U + # https://dlmf.nist.gov/12.10#E42 + phip = phi ** arange(6, 31, 6).reshape((-1,1)) + A0 = b0*u0 + A1 = (b2*u0 + phip[0,:]*b1*u1 + phip[1,:]*b0*u2) / zeta**3 + A2 = (b4*u0 + phip[0,:]*b3*u1 + phip[1,:]*b2*u2 + phip[2,:]*b1*u3 + + phip[3,:]*b0*u4) / zeta**6 + B0 = -(a1*u0 + phip[0,:]*a0*u1) / zeta**2 + B1 = -(a3*u0 + phip[0,:]*a2*u1 + phip[1,:]*a1*u2 + phip[2,:]*a0*u3) / zeta**5 + B2 = -(a5*u0 + phip[0,:]*a4*u1 + phip[1,:]*a3*u2 + phip[2,:]*a2*u3 + + phip[3,:]*a1*u4 + phip[4,:]*a0*u5) / zeta**8 + # U + # https://dlmf.nist.gov/12.10#E35 + U = P * (Ai * (A0 + A1/mu**2.0 + A2/mu**4.0) + + Aip * (B0 + B1/mu**2.0 + B2/mu**4.0) / mu**(8.0/6.0)) + # Prefactor for derivative of U + Pd = sqrt(2.0*pi) * mu**(2.0/6.0) / phi + # Terms for derivative of U + # https://dlmf.nist.gov/12.10#E46 + C0 = -(b1*v0 + phip[0,:]*b0*v1) / zeta + C1 = -(b3*v0 + phip[0,:]*b2*v1 + phip[1,:]*b1*v2 + phip[2,:]*b0*v3) / zeta**4 + C2 = -(b5*v0 + phip[0,:]*b4*v1 + phip[1,:]*b3*v2 + phip[2,:]*b2*v3 + + phip[3,:]*b1*v4 + phip[4,:]*b0*v5) / zeta**7 + D0 = a0*v0 + D1 = (a2*v0 + phip[0,:]*a1*v1 + phip[1,:]*a0*v2) / zeta**3 + D2 = (a4*v0 + phip[0,:]*a3*v1 + phip[1,:]*a2*v2 + phip[2,:]*a1*v3 + + phip[3,:]*a0*v4) / zeta**6 + # Derivative of U + # https://dlmf.nist.gov/12.10#E36 + Ud = Pd * (Ai * (C0 + C1/mu**2.0 + C2/mu**4.0) / mu**(4.0/6.0) + + Aip * (D0 + D1/mu**2.0 + D2/mu**4.0)) + return U, Ud + + +def _newton(n, x_initial, maxit=5): + """Newton iteration for polishing the asymptotic approximation + to the zeros of the Hermite polynomials. + + Parameters + ---------- + n : int + Quadrature order + x_initial : ndarray + Initial guesses for the roots + maxit : int + Maximal number of Newton iterations. + The default 5 is sufficient, usually + only one or two steps are needed. + + Returns + ------- + nodes : ndarray + Quadrature nodes + weights : ndarray + Quadrature weights + + See Also + -------- + roots_hermite_asy + """ + # Variable transformation + mu = sqrt(2.0*n + 1.0) + t = x_initial / mu + theta = arccos(t) + # Newton iteration + for i in range(maxit): + u, ud = _pbcf(n, theta) + dtheta = u / (sqrt(2.0) * mu * sin(theta) * ud) + theta = theta + dtheta + if max(abs(dtheta)) < 1e-14: + break + # Undo variable transformation + x = mu * cos(theta) + # Central node is always zero + if n % 2 == 1: + x[0] = 0.0 + # Compute weights + w = exp(-x**2) / (2.0*ud**2) + return x, w + + +def _roots_hermite_asy(n): + r"""Gauss-Hermite (physicist's) quadrature for large n. + + Computes the sample points and weights for Gauss-Hermite quadrature. + The sample points are the roots of the nth degree Hermite polynomial, + :math:`H_n(x)`. These sample points and weights correctly integrate + polynomials of degree :math:`2n - 1` or less over the interval + :math:`[-\infty, \infty]` with weight function :math:`f(x) = e^{-x^2}`. + + This method relies on asymptotic expansions which work best for n > 150. + The algorithm has linear runtime making computation for very large n + feasible. + + Parameters + ---------- + n : int + quadrature order + + Returns + ------- + nodes : ndarray + Quadrature nodes + weights : ndarray + Quadrature weights + + See Also + -------- + roots_hermite + + References + ---------- + .. [townsend.trogdon.olver-2014] + Townsend, A. and Trogdon, T. and Olver, S. (2014) + *Fast computation of Gauss quadrature nodes and + weights on the whole real line*. :arXiv:`1410.5286`. + + .. [townsend.trogdon.olver-2015] + Townsend, A. and Trogdon, T. and Olver, S. (2015) + *Fast computation of Gauss quadrature nodes and + weights on the whole real line*. + IMA Journal of Numerical Analysis + :doi:`10.1093/imanum/drv002`. + """ + iv = _initial_nodes(n) + nodes, weights = _newton(n, iv) + # Combine with negative parts + if n % 2 == 0: + nodes = hstack([-nodes[::-1], nodes]) + weights = hstack([weights[::-1], weights]) + else: + nodes = hstack([-nodes[-1:0:-1], nodes]) + weights = hstack([weights[-1:0:-1], weights]) + # Scale weights + weights *= sqrt(pi) / sum(weights) + return nodes, weights + + +def hermite(n, monic=False): + r"""Physicist's Hermite polynomial. + + Defined by + + .. math:: + + H_n(x) = (-1)^ne^{x^2}\frac{d^n}{dx^n}e^{-x^2}; + + :math:`H_n` is a polynomial of degree :math:`n`. + + Parameters + ---------- + n : int + Degree of the polynomial. + monic : bool, optional + If `True`, scale the leading coefficient to be 1. Default is + `False`. + + Returns + ------- + H : orthopoly1d + Physicist's Hermite polynomial. + + Notes + ----- + The polynomials :math:`H_n` are orthogonal over :math:`(-\infty, + \infty)` with weight function :math:`e^{-x^2}`. + + Examples + -------- + >>> from scipy import special + >>> import matplotlib.pyplot as plt + >>> import numpy as np + + >>> p_monic = special.hermite(3, monic=True) + >>> p_monic + poly1d([ 1. , 0. , -1.5, 0. ]) + >>> p_monic(1) + -0.49999999999999983 + >>> x = np.linspace(-3, 3, 400) + >>> y = p_monic(x) + >>> plt.plot(x, y) + >>> plt.title("Monic Hermite polynomial of degree 3") + >>> plt.xlabel("x") + >>> plt.ylabel("H_3(x)") + >>> plt.show() + + """ + if n < 0: + raise ValueError("n must be nonnegative.") + + if n == 0: + n1 = n + 1 + else: + n1 = n + x, w = roots_hermite(n1) + def wfunc(x): + return exp(-x * x) + if n == 0: + x, w = [], [] + hn = 2**n * _gam(n + 1) * sqrt(pi) + kn = 2**n + p = orthopoly1d(x, w, hn, kn, wfunc, (-inf, inf), monic, + lambda x: _ufuncs.eval_hermite(n, x)) + return p + +# Hermite 2 He_n(x) + + +def roots_hermitenorm(n, mu=False): + r"""Gauss-Hermite (statistician's) quadrature. + + Compute the sample points and weights for Gauss-Hermite + quadrature. The sample points are the roots of the nth degree + Hermite polynomial, :math:`He_n(x)`. These sample points and + weights correctly integrate polynomials of degree :math:`2n - 1` + or less over the interval :math:`[-\infty, \infty]` with weight + function :math:`w(x) = e^{-x^2/2}`. See 22.2.15 in [AS]_ for more + details. + + Parameters + ---------- + n : int + quadrature order. + mu : bool, optional + If True, return the sum of the weights, optional. + + Returns + ------- + x : ndarray + Sample points. + w : ndarray + Weights. + mu : float + Sum of the weights. + + See Also + -------- + scipy.integrate.fixed_quad + numpy.polynomial.hermite_e.hermegauss + + Notes + ----- + For small n up to 150 a modified version of the Golub-Welsch + algorithm is used. Nodes are computed from the eigenvalue + problem and improved by one step of a Newton iteration. + The weights are computed from the well-known analytical formula. + + For n larger than 150 an optimal asymptotic algorithm is used + which computes nodes and weights in a numerical stable manner. + The algorithm has linear runtime making computation for very + large n (several thousand or more) feasible. + + References + ---------- + .. [AS] Milton Abramowitz and Irene A. Stegun, eds. + Handbook of Mathematical Functions with Formulas, + Graphs, and Mathematical Tables. New York: Dover, 1972. + + """ + m = int(n) + if n < 1 or n != m: + raise ValueError("n must be a positive integer.") + + mu0 = np.sqrt(2.0*np.pi) + if n <= 150: + def an_func(k): + return 0.0 * k + def bn_func(k): + return np.sqrt(k) + f = _ufuncs.eval_hermitenorm + def df(n, x): + return n * _ufuncs.eval_hermitenorm(n - 1, x) + return _gen_roots_and_weights(m, mu0, an_func, bn_func, f, df, True, mu) + else: + nodes, weights = _roots_hermite_asy(m) + # Transform + nodes *= sqrt(2) + weights *= sqrt(2) + if mu: + return nodes, weights, mu0 + else: + return nodes, weights + + +def hermitenorm(n, monic=False): + r"""Probabilist's Hermite polynomial. + + Defined by + + .. math:: + + He_n(x) = (-1)^ne^{x^2/2}\frac{d^n}{dx^n}e^{-x^2/2}; + + :math:`He_n` is a polynomial of degree :math:`n`. + + Parameters + ---------- + n : int + Degree of the polynomial. + monic : bool, optional + If `True`, scale the leading coefficient to be 1. Default is + `False`. + + Returns + ------- + He : orthopoly1d + Probabilist's Hermite polynomial. + + Notes + ----- + + The polynomials :math:`He_n` are orthogonal over :math:`(-\infty, + \infty)` with weight function :math:`e^{-x^2/2}`. + + Examples + -------- + >>> import numpy as np + >>> import matplotlib.pyplot as plt + >>> from scipy.special import hermitenorm + + >>> p_monic = hermitenorm(3) + >>> p_monic + poly1d([ 1., 0., -3., 0.]) + + Evaluate the probabilist's Hermite polynomial of degree 3 at x = 1: + + >>> p_monic(1) + np.float64(-2.0) + + Plot probabilist's Hermite polynomials of degree 0 to 4: + + >>> x = np.linspace(-3, 3, 100) + >>> fig, ax = plt.subplots() + >>> for i in range(5): + ... ax.plot(x, hermitenorm(i)(x), label=f"n={i}") + >>> plt.title(f"Probabilist's Hermite polynomials $He_n$") + >>> plt.xlabel("x") + >>> plt.ylabel(rf"$He_n(x)$") + >>> plt.legend(loc="best") + >>> plt.show() + + """ + if n < 0: + raise ValueError("n must be nonnegative.") + + if n == 0: + n1 = n + 1 + else: + n1 = n + x, w = roots_hermitenorm(n1) + def wfunc(x): + return exp(-x * x / 2.0) + if n == 0: + x, w = [], [] + hn = sqrt(2 * pi) * _gam(n + 1) + kn = 1.0 + p = orthopoly1d(x, w, hn, kn, wfunc=wfunc, limits=(-inf, inf), monic=monic, + eval_func=lambda x: _ufuncs.eval_hermitenorm(n, x)) + return p + +# The remainder of the polynomials can be derived from the ones above. + +# Ultraspherical (Gegenbauer) C^(alpha)_n(x) + + +def roots_gegenbauer(n, alpha, mu=False): + r"""Gauss-Gegenbauer quadrature. + + Compute the sample points and weights for Gauss-Gegenbauer + quadrature. The sample points are the roots of the nth degree + Gegenbauer polynomial, :math:`C^{\alpha}_n(x)`. These sample + points and weights correctly integrate polynomials of degree + :math:`2n - 1` or less over the interval :math:`[-1, 1]` with + weight function :math:`w(x) = (1 - x^2)^{\alpha - 1/2}`. See + 22.2.3 in [AS]_ for more details. + + Parameters + ---------- + n : int + quadrature order. + alpha : float + alpha must be > -0.5. + mu : bool, optional + If True, return the sum of the weights, optional. + + Returns + ------- + x : ndarray + Sample points + w : ndarray + Weights + mu : float + Sum of the weights + + See Also + -------- + scipy.integrate.fixed_quad + + References + ---------- + .. [AS] Milton Abramowitz and Irene A. Stegun, eds. + Handbook of Mathematical Functions with Formulas, + Graphs, and Mathematical Tables. New York: Dover, 1972. + + """ + m = int(n) + if n < 1 or n != m: + raise ValueError("n must be a positive integer.") + if alpha < -0.5: + raise ValueError("alpha must be greater than -0.5.") + elif alpha == 0.0: + # C(n,0,x) == 0 uniformly, however, as alpha->0, C(n,alpha,x)->T(n,x) + # strictly, we should just error out here, since the roots are not + # really defined, but we used to return something useful, so let's + # keep doing so. + return roots_chebyt(n, mu) + + if alpha <= 170: + mu0 = (np.sqrt(np.pi) * _ufuncs.gamma(alpha + 0.5)) \ + / _ufuncs.gamma(alpha + 1) + else: + # For large alpha we use a Taylor series expansion around inf, + # expressed as a 6th order polynomial of a^-1 and using Horner's + # method to minimize computation and maximize precision + inv_alpha = 1. / alpha + coeffs = np.array([0.000207186, -0.00152206, -0.000640869, + 0.00488281, 0.0078125, -0.125, 1.]) + mu0 = coeffs[0] + for term in range(1, len(coeffs)): + mu0 = mu0 * inv_alpha + coeffs[term] + mu0 = mu0 * np.sqrt(np.pi / alpha) + def an_func(k): + return 0.0 * k + def bn_func(k): + return np.sqrt(k * (k + 2 * alpha - 1) / (4 * (k + alpha) * (k + alpha - 1))) + def f(n, x): + return _ufuncs.eval_gegenbauer(n, alpha, x) + def df(n, x): + return ( + -n * x * _ufuncs.eval_gegenbauer(n, alpha, x) + + (n + 2 * alpha - 1) * _ufuncs.eval_gegenbauer(n - 1, alpha, x) + ) / (1 - x ** 2) + return _gen_roots_and_weights(m, mu0, an_func, bn_func, f, df, True, mu) + + +def gegenbauer(n, alpha, monic=False): + r"""Gegenbauer (ultraspherical) polynomial. + + Defined to be the solution of the second-order linear ordinary differential equation + + .. math:: + (1 - x^2)\frac{d^2}{dx^2}C_n^{(\alpha)}(x) + - (2\alpha + 1)x\frac{d}{dx}C_n^{(\alpha)}(x) + + n(n + 2\alpha)C_n^{(\alpha)}(x) = 0 + + for :math:`\alpha > -1/2`; :math:`C_n^{(\alpha)}` is a polynomial + of degree :math:`n`. + + Parameters + ---------- + n : int + Degree of the polynomial. + alpha : float + Parameter, must be greater than -0.5. + monic : bool, optional + If `True`, scale the leading coefficient to be 1. Default is + `False`. + + Returns + ------- + C : orthopoly1d + Gegenbauer polynomial. + + Notes + ----- + The polynomials :math:`C_n^{(\alpha)}` are orthogonal on + :math:`[-1,1]` with respect to the weight function :math:`(1 - x^2)^{(\alpha - + 1/2)}`. + + Examples + -------- + >>> import numpy as np + >>> from scipy import special + >>> import matplotlib.pyplot as plt + + We can initialize a variable ``p`` as a Gegenbauer polynomial using the + `gegenbauer` function and evaluate at a point ``x = 1``. + + >>> p = special.gegenbauer(3, 0.5, monic=False) + >>> p + poly1d([ 2.5, 0. , -1.5, 0. ]) + >>> p(1) + 1.0 + + To evaluate ``p`` at various points ``x`` in the interval ``(-3, 3)``, + simply pass an array ``x`` to ``p`` as follows: + + >>> x = np.linspace(-3, 3, 400) + >>> y = p(x) + + We can then visualize ``x, y`` using `matplotlib.pyplot`. + + >>> fig, ax = plt.subplots() + >>> ax.plot(x, y) + >>> ax.set_title("Gegenbauer (ultraspherical) polynomial of degree 3") + >>> ax.set_xlabel("x") + >>> ax.set_ylabel(r"$C_3^{(0.5)}(x)$") + >>> plt.show() + + """ + if not np.isfinite(alpha) or alpha <= -0.5 : + raise ValueError("`alpha` must be a finite number greater than -1/2") + base = jacobi(n, alpha - 0.5, alpha - 0.5, monic=monic) + if monic or n == 0: + return base + # Abrahmowitz and Stegan 22.5.20 + factor = (_gam(2*alpha + n) * _gam(alpha + 0.5) / + _gam(2*alpha) / _gam(alpha + 0.5 + n)) + base._scale(factor) + base.__dict__['_eval_func'] = lambda x: _ufuncs.eval_gegenbauer(float(n), + alpha, x) + return base + +# Chebyshev of the first kind: T_n(x) = +# n! sqrt(pi) / _gam(n+1./2)* P^(-1/2,-1/2)_n(x) +# Computed anew. + + +def roots_chebyt(n, mu=False): + r"""Gauss-Chebyshev (first kind) quadrature. + + Computes the sample points and weights for Gauss-Chebyshev + quadrature. The sample points are the roots of the nth degree + Chebyshev polynomial of the first kind, :math:`T_n(x)`. These + sample points and weights correctly integrate polynomials of + degree :math:`2n - 1` or less over the interval :math:`[-1, 1]` + with weight function :math:`w(x) = 1/\sqrt{1 - x^2}`. See 22.2.4 + in [AS]_ for more details. + + Parameters + ---------- + n : int + quadrature order + mu : bool, optional + If True, return the sum of the weights, optional. + + Returns + ------- + x : ndarray + Sample points + w : ndarray + Weights + mu : float + Sum of the weights + + See Also + -------- + scipy.integrate.fixed_quad + numpy.polynomial.chebyshev.chebgauss + + References + ---------- + .. [AS] Milton Abramowitz and Irene A. Stegun, eds. + Handbook of Mathematical Functions with Formulas, + Graphs, and Mathematical Tables. New York: Dover, 1972. + + """ + m = int(n) + if n < 1 or n != m: + raise ValueError('n must be a positive integer.') + x = _ufuncs._sinpi(np.arange(-m + 1, m, 2) / (2*m)) + w = np.full_like(x, pi/m) + if mu: + return x, w, pi + else: + return x, w + + +def chebyt(n, monic=False): + r"""Chebyshev polynomial of the first kind. + + Defined to be the solution of + + .. math:: + (1 - x^2)\frac{d^2}{dx^2}T_n(x) - x\frac{d}{dx}T_n(x) + + n^2T_n(x) = 0; + + :math:`T_n` is a polynomial of degree :math:`n`. + + Parameters + ---------- + n : int + Degree of the polynomial. + monic : bool, optional + If `True`, scale the leading coefficient to be 1. Default is + `False`. + + Returns + ------- + T : orthopoly1d + Chebyshev polynomial of the first kind. + + See Also + -------- + chebyu : Chebyshev polynomial of the second kind. + + Notes + ----- + The polynomials :math:`T_n` are orthogonal over :math:`[-1, 1]` + with weight function :math:`(1 - x^2)^{-1/2}`. + + References + ---------- + .. [AS] Milton Abramowitz and Irene A. Stegun, eds. + Handbook of Mathematical Functions with Formulas, + Graphs, and Mathematical Tables. New York: Dover, 1972. + + Examples + -------- + Chebyshev polynomials of the first kind of order :math:`n` can + be obtained as the determinant of specific :math:`n \times n` + matrices. As an example we can check how the points obtained from + the determinant of the following :math:`3 \times 3` matrix + lay exactly on :math:`T_3`: + + >>> import numpy as np + >>> import matplotlib.pyplot as plt + >>> from scipy.linalg import det + >>> from scipy.special import chebyt + >>> x = np.arange(-1.0, 1.0, 0.01) + >>> fig, ax = plt.subplots() + >>> ax.set_ylim(-2.0, 2.0) + >>> ax.set_title(r'Chebyshev polynomial $T_3$') + >>> ax.plot(x, chebyt(3)(x), label=rf'$T_3$') + >>> for p in np.arange(-1.0, 1.0, 0.1): + ... ax.plot(p, + ... det(np.array([[p, 1, 0], [1, 2*p, 1], [0, 1, 2*p]])), + ... 'rx') + >>> plt.legend(loc='best') + >>> plt.show() + + They are also related to the Jacobi Polynomials + :math:`P_n^{(-0.5, -0.5)}` through the relation: + + .. math:: + P_n^{(-0.5, -0.5)}(x) = \frac{1}{4^n} \binom{2n}{n} T_n(x) + + Let's verify it for :math:`n = 3`: + + >>> from scipy.special import binom + >>> from scipy.special import jacobi + >>> x = np.arange(-1.0, 1.0, 0.01) + >>> np.allclose(jacobi(3, -0.5, -0.5)(x), + ... 1/64 * binom(6, 3) * chebyt(3)(x)) + True + + We can plot the Chebyshev polynomials :math:`T_n` for some values + of :math:`n`: + + >>> x = np.arange(-1.5, 1.5, 0.01) + >>> fig, ax = plt.subplots() + >>> ax.set_ylim(-4.0, 4.0) + >>> ax.set_title(r'Chebyshev polynomials $T_n$') + >>> for n in np.arange(2,5): + ... ax.plot(x, chebyt(n)(x), label=rf'$T_n={n}$') + >>> plt.legend(loc='best') + >>> plt.show() + + """ + if n < 0: + raise ValueError("n must be nonnegative.") + + def wfunc(x): + return 1.0 / sqrt(1 - x * x) + if n == 0: + return orthopoly1d([], [], pi, 1.0, wfunc, (-1, 1), monic, + lambda x: _ufuncs.eval_chebyt(n, x)) + n1 = n + x, w, mu = roots_chebyt(n1, mu=True) + hn = pi / 2 + kn = 2**(n - 1) + p = orthopoly1d(x, w, hn, kn, wfunc, (-1, 1), monic, + lambda x: _ufuncs.eval_chebyt(n, x)) + return p + +# Chebyshev of the second kind +# U_n(x) = (n+1)! sqrt(pi) / (2*_gam(n+3./2)) * P^(1/2,1/2)_n(x) + + +def roots_chebyu(n, mu=False): + r"""Gauss-Chebyshev (second kind) quadrature. + + Computes the sample points and weights for Gauss-Chebyshev + quadrature. The sample points are the roots of the nth degree + Chebyshev polynomial of the second kind, :math:`U_n(x)`. These + sample points and weights correctly integrate polynomials of + degree :math:`2n - 1` or less over the interval :math:`[-1, 1]` + with weight function :math:`w(x) = \sqrt{1 - x^2}`. See 22.2.5 in + [AS]_ for details. + + Parameters + ---------- + n : int + quadrature order + mu : bool, optional + If True, return the sum of the weights, optional. + + Returns + ------- + x : ndarray + Sample points + w : ndarray + Weights + mu : float + Sum of the weights + + See Also + -------- + scipy.integrate.fixed_quad + + References + ---------- + .. [AS] Milton Abramowitz and Irene A. Stegun, eds. + Handbook of Mathematical Functions with Formulas, + Graphs, and Mathematical Tables. New York: Dover, 1972. + + """ + m = int(n) + if n < 1 or n != m: + raise ValueError('n must be a positive integer.') + t = np.arange(m, 0, -1) * pi / (m + 1) + x = np.cos(t) + w = pi * np.sin(t)**2 / (m + 1) + if mu: + return x, w, pi / 2 + else: + return x, w + + +def chebyu(n, monic=False): + r"""Chebyshev polynomial of the second kind. + + Defined to be the solution of + + .. math:: + (1 - x^2)\frac{d^2}{dx^2}U_n(x) - 3x\frac{d}{dx}U_n(x) + + n(n + 2)U_n(x) = 0; + + :math:`U_n` is a polynomial of degree :math:`n`. + + Parameters + ---------- + n : int + Degree of the polynomial. + monic : bool, optional + If `True`, scale the leading coefficient to be 1. Default is + `False`. + + Returns + ------- + U : orthopoly1d + Chebyshev polynomial of the second kind. + + See Also + -------- + chebyt : Chebyshev polynomial of the first kind. + + Notes + ----- + The polynomials :math:`U_n` are orthogonal over :math:`[-1, 1]` + with weight function :math:`(1 - x^2)^{1/2}`. + + References + ---------- + .. [AS] Milton Abramowitz and Irene A. Stegun, eds. + Handbook of Mathematical Functions with Formulas, + Graphs, and Mathematical Tables. New York: Dover, 1972. + + Examples + -------- + Chebyshev polynomials of the second kind of order :math:`n` can + be obtained as the determinant of specific :math:`n \times n` + matrices. As an example we can check how the points obtained from + the determinant of the following :math:`3 \times 3` matrix + lay exactly on :math:`U_3`: + + >>> import numpy as np + >>> import matplotlib.pyplot as plt + >>> from scipy.linalg import det + >>> from scipy.special import chebyu + >>> x = np.arange(-1.0, 1.0, 0.01) + >>> fig, ax = plt.subplots() + >>> ax.set_ylim(-2.0, 2.0) + >>> ax.set_title(r'Chebyshev polynomial $U_3$') + >>> ax.plot(x, chebyu(3)(x), label=rf'$U_3$') + >>> for p in np.arange(-1.0, 1.0, 0.1): + ... ax.plot(p, + ... det(np.array([[2*p, 1, 0], [1, 2*p, 1], [0, 1, 2*p]])), + ... 'rx') + >>> plt.legend(loc='best') + >>> plt.show() + + They satisfy the recurrence relation: + + .. math:: + U_{2n-1}(x) = 2 T_n(x)U_{n-1}(x) + + where :math:`T_n` is the Chebyshev polynomial of the first kind. + Let's verify it for :math:`n = 2`: + + >>> from scipy.special import chebyt + >>> x = np.arange(-1.0, 1.0, 0.01) + >>> np.allclose(chebyu(3)(x), 2 * chebyt(2)(x) * chebyu(1)(x)) + True + + We can plot the Chebyshev polynomials :math:`U_n` for some values + of :math:`n`: + + >>> x = np.arange(-1.0, 1.0, 0.01) + >>> fig, ax = plt.subplots() + >>> ax.set_ylim(-1.5, 1.5) + >>> ax.set_title(r'Chebyshev polynomials $U_n$') + >>> for n in np.arange(1,5): + ... ax.plot(x, chebyu(n)(x), label=rf'$U_n={n}$') + >>> plt.legend(loc='best') + >>> plt.show() + + """ + base = jacobi(n, 0.5, 0.5, monic=monic) + if monic: + return base + factor = sqrt(pi) / 2.0 * _gam(n + 2) / _gam(n + 1.5) + base._scale(factor) + return base + +# Chebyshev of the first kind C_n(x) + + +def roots_chebyc(n, mu=False): + r"""Gauss-Chebyshev (first kind) quadrature. + + Compute the sample points and weights for Gauss-Chebyshev + quadrature. The sample points are the roots of the nth degree + Chebyshev polynomial of the first kind, :math:`C_n(x)`. These + sample points and weights correctly integrate polynomials of + degree :math:`2n - 1` or less over the interval :math:`[-2, 2]` + with weight function :math:`w(x) = 1 / \sqrt{1 - (x/2)^2}`. See + 22.2.6 in [AS]_ for more details. + + Parameters + ---------- + n : int + quadrature order + mu : bool, optional + If True, return the sum of the weights, optional. + + Returns + ------- + x : ndarray + Sample points + w : ndarray + Weights + mu : float + Sum of the weights + + See Also + -------- + scipy.integrate.fixed_quad + + References + ---------- + .. [AS] Milton Abramowitz and Irene A. Stegun, eds. + Handbook of Mathematical Functions with Formulas, + Graphs, and Mathematical Tables. New York: Dover, 1972. + + """ + x, w, m = roots_chebyt(n, True) + x *= 2 + w *= 2 + m *= 2 + if mu: + return x, w, m + else: + return x, w + + +def chebyc(n, monic=False): + r"""Chebyshev polynomial of the first kind on :math:`[-2, 2]`. + + Defined as :math:`C_n(x) = 2T_n(x/2)`, where :math:`T_n` is the + nth Chebyshev polynomial of the first kind. + + Parameters + ---------- + n : int + Degree of the polynomial. + monic : bool, optional + If `True`, scale the leading coefficient to be 1. Default is + `False`. + + Returns + ------- + C : orthopoly1d + Chebyshev polynomial of the first kind on :math:`[-2, 2]`. + + See Also + -------- + chebyt : Chebyshev polynomial of the first kind. + + Notes + ----- + The polynomials :math:`C_n(x)` are orthogonal over :math:`[-2, 2]` + with weight function :math:`1/\sqrt{1 - (x/2)^2}`. + + References + ---------- + .. [1] Abramowitz and Stegun, "Handbook of Mathematical Functions" + Section 22. National Bureau of Standards, 1972. + + Examples + -------- + Evaluate the Chebyshev polynomial :math:`C_3` at :math:`x = 1`: + + >>> import numpy as np + >>> from scipy.special import chebyc, chebyt + >>> np.isclose(chebyc(3)(1), -2.0) + True + + The polynomial :math:`C_n` is a scaled Chebyshev polynomial of the + first kind: + + >>> x = np.linspace(-2, 2, 5) + >>> np.allclose(chebyc(3)(x), 2 * chebyt(3)(x/2)) + True + + Plot :math:`C_n` for several values of :math:`n`: + + >>> import matplotlib.pyplot as plt + >>> x = np.linspace(-2, 2, 400) + >>> fig, ax = plt.subplots() + >>> for n in range(4): + ... ax.plot(x, chebyc(n)(x), label=rf"$C_{n}$") + >>> ax.set_title(r"Chebyshev polynomials $C_n$") + >>> ax.set_xlabel("x") + >>> ax.legend(loc="best") + >>> plt.show() + """ + if n < 0: + raise ValueError("n must be nonnegative.") + + if n == 0: + n1 = n + 1 + else: + n1 = n + x, w = roots_chebyc(n1) + if n == 0: + x, w = [], [] + hn = 4 * pi * ((n == 0) + 1) + kn = 1.0 + p = orthopoly1d(x, w, hn, kn, + wfunc=lambda x: 1.0 / sqrt(1 - x * x / 4.0), + limits=(-2, 2), monic=monic) + if not monic: + p._scale(2.0 / p(2)) + p.__dict__['_eval_func'] = lambda x: _ufuncs.eval_chebyc(n, x) + return p + +# Chebyshev of the second kind S_n(x) + + +def roots_chebys(n, mu=False): + r"""Gauss-Chebyshev (second kind) quadrature. + + Compute the sample points and weights for Gauss-Chebyshev + quadrature. The sample points are the roots of the nth degree + Chebyshev polynomial of the second kind, :math:`S_n(x)`. These + sample points and weights correctly integrate polynomials of + degree :math:`2n - 1` or less over the interval :math:`[-2, 2]` + with weight function :math:`w(x) = \sqrt{1 - (x/2)^2}`. See 22.2.7 + in [AS]_ for more details. + + Parameters + ---------- + n : int + quadrature order + mu : bool, optional + If True, return the sum of the weights, optional. + + Returns + ------- + x : ndarray + Sample points + w : ndarray + Weights + mu : float + Sum of the weights + + See Also + -------- + scipy.integrate.fixed_quad + + References + ---------- + .. [AS] Milton Abramowitz and Irene A. Stegun, eds. + Handbook of Mathematical Functions with Formulas, + Graphs, and Mathematical Tables. New York: Dover, 1972. + + """ + x, w, m = roots_chebyu(n, True) + x *= 2 + w *= 2 + m *= 2 + if mu: + return x, w, m + else: + return x, w + + +def chebys(n, monic=False): + r"""Chebyshev polynomial of the second kind on :math:`[-2, 2]`. + + Defined as :math:`S_n(x) = U_n(x/2)` where :math:`U_n` is the + nth Chebyshev polynomial of the second kind. + + Parameters + ---------- + n : int + Degree of the polynomial. + monic : bool, optional + If `True`, scale the leading coefficient to be 1. Default is + `False`. + + Returns + ------- + S : orthopoly1d + Chebyshev polynomial of the second kind on :math:`[-2, 2]`. + + See Also + -------- + chebyu : Chebyshev polynomial of the second kind + + Notes + ----- + The polynomials :math:`S_n(x)` are orthogonal over :math:`[-2, 2]` + with weight function :math:`\sqrt{1 - (x/2)^2}`. + + References + ---------- + .. [1] Abramowitz and Stegun, "Handbook of Mathematical Functions" + Section 22. National Bureau of Standards, 1972. + + Examples + -------- + Evaluate the Chebyshev polynomial of the second kind :math:`S_3` at :math:`x = 1`: + + >>> import numpy as np + >>> from scipy.special import chebys, chebyu + >>> np.isclose(chebys(3)(1), -1.0) + True + + The polynomial :math:`S_n` is a scaled Chebyshev polynomial of the + second kind: + + >>> x = np.linspace(-2, 2, 5) + >>> np.allclose(chebys(3)(x), chebyu(3)(x/2)) + True + + Plot :math:`S_n` for several values of :math:`n`: + + >>> import matplotlib.pyplot as plt + >>> x = np.linspace(-2, 2, 400) + >>> fig, ax = plt.subplots() + >>> for n in range(4): + ... ax.plot(x, chebys(n)(x), label=rf"$S_{n}$") + >>> ax.set_title(r"Chebyshev polynomials $S_n$") + >>> ax.set_xlabel("x") + >>> ax.legend(loc="best") + >>> plt.show() + """ + if n < 0: + raise ValueError("n must be nonnegative.") + + if n == 0: + n1 = n + 1 + else: + n1 = n + x, w = roots_chebys(n1) + if n == 0: + x, w = [], [] + hn = pi + kn = 1.0 + p = orthopoly1d(x, w, hn, kn, + wfunc=lambda x: sqrt(1 - x * x / 4.0), + limits=(-2, 2), monic=monic) + if not monic: + factor = (n + 1.0) / p(2) + p._scale(factor) + p.__dict__['_eval_func'] = lambda x: _ufuncs.eval_chebys(n, x) + return p + +# Shifted Chebyshev of the first kind T^*_n(x) + + +def roots_sh_chebyt(n, mu=False): + r"""Gauss-Chebyshev (first kind, shifted) quadrature. + + Compute the sample points and weights for Gauss-Chebyshev + quadrature. The sample points are the roots of the nth degree + shifted Chebyshev polynomial of the first kind, :math:`T_n(x)`. + These sample points and weights correctly integrate polynomials of + degree :math:`2n - 1` or less over the interval :math:`[0, 1]` + with weight function :math:`w(x) = 1/\sqrt{x - x^2}`. See 22.2.8 + in [AS]_ for more details. + + Parameters + ---------- + n : int + quadrature order. + mu : bool, optional + If True, return the sum of the weights, optional. + + Returns + ------- + x : ndarray + Sample points. + w : ndarray + Weights. + mu : float + Sum of the weights. + + See Also + -------- + scipy.integrate.fixed_quad + + References + ---------- + .. [AS] Milton Abramowitz and Irene A. Stegun, eds. + Handbook of Mathematical Functions with Formulas, + Graphs, and Mathematical Tables. New York: Dover, 1972. + + """ + xw = roots_chebyt(n, mu) + return ((xw[0] + 1) / 2,) + xw[1:] + + +def sh_chebyt(n, monic=False): + r"""Shifted Chebyshev polynomial of the first kind. + + Defined as :math:`T^*_n(x) = T_n(2x - 1)` for :math:`T_n` the nth + Chebyshev polynomial of the first kind. + + Parameters + ---------- + n : int + Degree of the polynomial. + monic : bool, optional + If `True`, scale the leading coefficient to be 1. Default is + `False`. + + Returns + ------- + T : orthopoly1d + Shifted Chebyshev polynomial of the first kind. + + Notes + ----- + The polynomials :math:`T^*_n` are orthogonal over :math:`[0, 1]` + with weight function :math:`(x - x^2)^{-1/2}`. + + Examples + -------- + Evaluate the shifted Chebyshev polynomial of the first kind :math:`T^*_3` at + :math:`x = 0.75`: + + >>> import numpy as np + >>> from scipy.special import chebyt, sh_chebyt + >>> np.isclose(sh_chebyt(3)(0.75), -1.0) + True + + The polynomial :math:`T^*_n` is the Chebyshev polynomial + :math:`T_n` shifted from :math:`[-1, 1]` to :math:`[0, 1]`: + + >>> x = np.linspace(0, 1, 5) + >>> np.allclose(sh_chebyt(3)(x), chebyt(3)(2*x - 1)) + True + + Plot :math:`T^*_n` for several values of :math:`n`: + + >>> import matplotlib.pyplot as plt + >>> x = np.linspace(0, 1, 400) + >>> fig, ax = plt.subplots() + >>> for n in range(4): + ... ax.plot(x, sh_chebyt(n)(x), label=rf"$T^*_{n}$") + >>> ax.set_title(r"Shifted Chebyshev polynomials $T^*_n$") + >>> ax.set_xlabel("x") + >>> ax.legend(loc="best") + >>> plt.show() + + """ + base = sh_jacobi(n, 0.0, 0.5, monic=monic) + if monic: + return base + if n > 0: + factor = 4**n / 2.0 + else: + factor = 1.0 + base._scale(factor) + return base + + +# Shifted Chebyshev of the second kind U^*_n(x) +def roots_sh_chebyu(n, mu=False): + r"""Gauss-Chebyshev (second kind, shifted) quadrature. + + Computes the sample points and weights for Gauss-Chebyshev + quadrature. The sample points are the roots of the nth degree + shifted Chebyshev polynomial of the second kind, :math:`U_n(x)`. + These sample points and weights correctly integrate polynomials of + degree :math:`2n - 1` or less over the interval :math:`[0, 1]` + with weight function :math:`w(x) = \sqrt{x - x^2}`. See 22.2.9 in + [AS]_ for more details. + + Parameters + ---------- + n : int + quadrature order + mu : bool, optional + If True, return the sum of the weights, optional. + + Returns + ------- + x : ndarray + Sample points + w : ndarray + Weights + mu : float + Sum of the weights + + See Also + -------- + scipy.integrate.fixed_quad + + References + ---------- + .. [AS] Milton Abramowitz and Irene A. Stegun, eds. + Handbook of Mathematical Functions with Formulas, + Graphs, and Mathematical Tables. New York: Dover, 1972. + + """ + x, w, m = roots_chebyu(n, True) + x = (x + 1) / 2 + m_us = _ufuncs.beta(1.5, 1.5) + w *= m_us / m + if mu: + return x, w, m_us + else: + return x, w + + +def sh_chebyu(n, monic=False): + r"""Shifted Chebyshev polynomial of the second kind. + + Defined as :math:`U^*_n(x) = U_n(2x - 1)` for :math:`U_n` the nth + Chebyshev polynomial of the second kind. + + Parameters + ---------- + n : int + Degree of the polynomial. + monic : bool, optional + If `True`, scale the leading coefficient to be 1. Default is + `False`. + + Returns + ------- + U : orthopoly1d + Shifted Chebyshev polynomial of the second kind. + + Notes + ----- + The polynomials :math:`U^*_n` are orthogonal over :math:`[0, 1]` + with weight function :math:`(x - x^2)^{1/2}`. + + Examples + -------- + Evaluate the shifted Chebyshev polynomial of the second kind :math:`U^*_3` at + :math:`x = 0.75`: + + >>> import numpy as np + >>> from scipy.special import chebyu, sh_chebyu + >>> np.isclose(sh_chebyu(3)(0.75), -1.0) + True + + The polynomial :math:`U^*_n` is the Chebyshev polynomial + :math:`U_n` shifted from :math:`[-1, 1]` to :math:`[0, 1]`: + + >>> x = np.linspace(0, 1, 5) + >>> np.allclose(sh_chebyu(3)(x), chebyu(3)(2*x - 1)) + True + + Plot :math:`U^*_n` for several values of :math:`n`: + + >>> import matplotlib.pyplot as plt + >>> x = np.linspace(0, 1, 400) + >>> fig, ax = plt.subplots() + >>> for n in range(4): + ... ax.plot(x, sh_chebyu(n)(x), label=rf"$U^*_{n}$") + >>> ax.set_title(r"Shifted Chebyshev polynomials $U^*_n$") + >>> ax.set_xlabel("x") + >>> ax.legend(loc="best") + >>> plt.show() + + """ + base = sh_jacobi(n, 2.0, 1.5, monic=monic) + if monic: + return base + factor = 4**n + base._scale(factor) + return base + +# Legendre + + +def roots_legendre(n, mu=False): + r"""Gauss-Legendre quadrature. + + Compute the sample points and weights for Gauss-Legendre + quadrature [GL]_. The sample points are the roots of the nth degree + Legendre polynomial :math:`P_n(x)`. These sample points and + weights correctly integrate polynomials of degree :math:`2n - 1` + or less over the interval :math:`[-1, 1]` with weight function + :math:`w(x) = 1`. See 2.2.10 in [AS]_ for more details. + + Parameters + ---------- + n : int + quadrature order. + mu : bool, optional + If True, return the sum of the weights, optional. + + Returns + ------- + x : ndarray + Sample points. + w : ndarray + Weights. + mu : float + Sum of the weights. + + See Also + -------- + scipy.integrate.fixed_quad + numpy.polynomial.legendre.leggauss + + References + ---------- + .. [AS] Milton Abramowitz and Irene A. Stegun, eds. + Handbook of Mathematical Functions with Formulas, + Graphs, and Mathematical Tables. New York: Dover, 1972. + .. [GL] Gauss-Legendre quadrature, Wikipedia, + https://en.wikipedia.org/wiki/Gauss%E2%80%93Legendre_quadrature + + Examples + -------- + >>> import numpy as np + >>> from scipy.special import roots_legendre, eval_legendre + >>> roots, weights = roots_legendre(9) + + ``roots`` holds the roots, and ``weights`` holds the weights for + Gauss-Legendre quadrature. + + >>> roots + array([-0.96816024, -0.83603111, -0.61337143, -0.32425342, 0. , + 0.32425342, 0.61337143, 0.83603111, 0.96816024]) + >>> weights + array([0.08127439, 0.18064816, 0.2606107 , 0.31234708, 0.33023936, + 0.31234708, 0.2606107 , 0.18064816, 0.08127439]) + + Verify that we have the roots by evaluating the degree 9 Legendre + polynomial at ``roots``. All the values are approximately zero: + + >>> eval_legendre(9, roots) + array([-8.88178420e-16, -2.22044605e-16, 1.11022302e-16, 1.11022302e-16, + 0.00000000e+00, -5.55111512e-17, -1.94289029e-16, 1.38777878e-16, + -8.32667268e-17]) + + Here we'll show how the above values can be used to estimate the + integral from 1 to 2 of f(t) = t + 1/t with Gauss-Legendre + quadrature [GL]_. First define the function and the integration + limits. + + >>> def f(t): + ... return t + 1/t + ... + >>> a = 1 + >>> b = 2 + + We'll use ``integral(f(t), t=a, t=b)`` to denote the definite integral + of f from t=a to t=b. The sample points in ``roots`` are from the + interval [-1, 1], so we'll rewrite the integral with the simple change + of variable:: + + x = 2/(b - a) * t - (a + b)/(b - a) + + with inverse:: + + t = (b - a)/2 * x + (a + b)/2 + + Then:: + + integral(f(t), a, b) = + (b - a)/2 * integral(f((b-a)/2*x + (a+b)/2), x=-1, x=1) + + We can approximate the latter integral with the values returned + by `roots_legendre`. + + Map the roots computed above from [-1, 1] to [a, b]. + + >>> t = (b - a)/2 * roots + (a + b)/2 + + Approximate the integral as the weighted sum of the function values. + + >>> (b - a)/2 * f(t).dot(weights) + 2.1931471805599276 + + Compare that to the exact result, which is 3/2 + log(2): + + >>> 1.5 + np.log(2) + 2.1931471805599454 + + """ + m = int(n) + if n < 1 or n != m: + raise ValueError("n must be a positive integer.") + + mu0 = 2.0 + def an_func(k): + return 0.0 * k + def bn_func(k): + return k * np.sqrt(1.0 / (4 * k * k - 1)) + f = _ufuncs.eval_legendre + def df(n, x): + return (-n * x * _ufuncs.eval_legendre(n, x) + + n * _ufuncs.eval_legendre(n - 1, x)) / (1 - x ** 2) + return _gen_roots_and_weights(m, mu0, an_func, bn_func, f, df, True, mu) + + +def legendre(n, monic=False): + r"""Legendre polynomial. + + Defined to be the solution of + + .. math:: + \frac{d}{dx}\left[(1 - x^2)\frac{d}{dx}P_n(x)\right] + + n(n + 1)P_n(x) = 0; + + :math:`P_n(x)` is a polynomial of degree :math:`n`. + + Parameters + ---------- + n : int + Degree of the polynomial. + monic : bool, optional + If `True`, scale the leading coefficient to be 1. Default is + `False`. + + Returns + ------- + P : orthopoly1d + Legendre polynomial. + + Notes + ----- + The polynomials :math:`P_n` are orthogonal over :math:`[-1, 1]` + with weight function 1. + + Examples + -------- + Generate the 3rd-order Legendre polynomial 1/2*(5x^3 + 0x^2 - 3x + 0): + + >>> from scipy.special import legendre + >>> legendre(3) + poly1d([ 2.5, 0. , -1.5, 0. ]) + + """ + if n < 0: + raise ValueError("n must be nonnegative.") + + if n == 0: + n1 = n + 1 + else: + n1 = n + x, w = roots_legendre(n1) + if n == 0: + x, w = [], [] + hn = 2.0 / (2 * n + 1) + kn = _gam(2 * n + 1) / _gam(n + 1)**2 / 2.0**n + p = orthopoly1d(x, w, hn, kn, wfunc=lambda x: 1.0, limits=(-1, 1), + monic=monic, + eval_func=lambda x: _ufuncs.eval_legendre(n, x)) + return p + +# Shifted Legendre P^*_n(x) + + +def roots_sh_legendre(n, mu=False): + r"""Gauss-Legendre (shifted) quadrature. + + Compute the sample points and weights for Gauss-Legendre + quadrature. The sample points are the roots of the nth degree + shifted Legendre polynomial :math:`P^*_n(x)`. These sample points + and weights correctly integrate polynomials of degree :math:`2n - + 1` or less over the interval :math:`[0, 1]` with weight function + :math:`w(x) = 1.0`. See 2.2.11 in [AS]_ for details. + + Parameters + ---------- + n : int + quadrature order + mu : bool, optional + If True, return the sum of the weights, optional. + + Returns + ------- + x : ndarray + Sample points + w : ndarray + Weights + mu : float + Sum of the weights + + See Also + -------- + scipy.integrate.fixed_quad + + References + ---------- + .. [AS] Milton Abramowitz and Irene A. Stegun, eds. + Handbook of Mathematical Functions with Formulas, + Graphs, and Mathematical Tables. New York: Dover, 1972. + + Examples + -------- + + Compute nodes and weights for a 7th-order shifted Gauss-Legendre quadrature: + + >>> from scipy.special import roots_sh_legendre, eval_sh_legendre + >>> x, w = roots_sh_legendre(7) + >>> x + array([0.02544604, 0.12923441, 0.29707742, 0.5, 0.70292258, 0.87076559, + 0.97455396]) + >>> w + array([0.06474248, 0.1398527, 0.19091503, 0.20897959, 0.19091503, 0.1398527, + 0.06474248]) + + Verify that ``x`` are the roots of the degree-7 shifted Legendre polynomial: + + >>> eval_sh_legendre(7, x) + array([5.55111512e-16, 1.11022302e-16, 3.33066907e-16, 0.00000000e+00, + -2.22044605e-16, -1.11022302e-16, -1.85962357e-15]) + + The sum of the weights for shifted Gauss-Legendre quadrature is always 1: + + >>> x, w, mu = roots_sh_legendre(10, mu=True) + >>> mu + 1.0 # Sum of weights of shifted Gauss-Legendre quadrature is always 1 + + """ + x, w = roots_legendre(n) + x = (x + 1) / 2 + w /= 2 + if mu: + return x, w, 1.0 + else: + return x, w + + +def sh_legendre(n, monic=False): + r""" + Shifted Legendre polynomial. + + Defined as :math:`P^*_n(x) = P_n(2x - 1)` for :math:`P_n` the nth + Legendre polynomial. + + Parameters + ---------- + n : int + Degree of the polynomial. + monic : bool, optional + If `True`, scale the leading coefficient to be 1. Default is + `False`. + + Returns + ------- + P : orthopoly1d + Shifted Legendre polynomial. + + See Also + -------- + scipy.special.legendre + scipy.special.roots_sh_legendre + + Notes + ----- + The polynomials :math:`P^*_n` are orthogonal over :math:`[0, 1]` + with weight function 1. + + Examples + -------- + The shifted Legendre polynomials :math:`P_n^*` are related + to the non-shifted polynomials :math:`P_n` by + :math:`P_n^*(x) = P_n(2x - 1)`. We can verify this on the + interval :math:`[0, 1]`: + + >>> import numpy as np + >>> from scipy.special import sh_legendre, legendre + >>> from scipy.integrate import trapezoid + >>> x = np.arange(0.0, 1.0, 0.01) + >>> n = 3 + >>> np.allclose(sh_legendre(n)(x), legendre(n)(2*x - 1)) + True + + The polynomials :math:`P_n^*` satisfy a recurrence + relation obtained by the change of variables + :math:`t = 2x - 1` in the standard Legendre recurrence: + + .. math:: + + (n+1) P_{n+1}^*(x) = (2n+1)(2x-1)\,P_n^*(x) - n\,P_{n-1}^*(x). + + This can be easily checked on :math:`[0, 1]` + for :math:`n = 3`: + + >>> n = 3 + >>> x = np.linspace(0.0, 1.0, 101) + >>> lhs = (n + 1) * sh_legendre(n + 1)(x) + >>> rhs = ( + ... (2*n + 1) * (2*x - 1) * sh_legendre(n)(x) + ... - n * sh_legendre(n - 1)(x) + ... ) + >>> np.allclose(lhs, rhs) + True + + Orthogonality over :math:`[0,1]` with weight 1 can be + checked numerically; for example, :math:`P_2^*` + is orthogonal to :math:`P_3^*`: + + >>> x = np.linspace(0.0, 1.0, 400) + >>> y = sh_legendre(2)(x) * sh_legendre(3)(x) + >>> np.isclose(trapezoid(y, x), 0.0, atol=1e-12) + True + """ + if n < 0: + raise ValueError("n must be nonnegative.") + + def wfunc(x): + return 0.0 * x + 1.0 + if n == 0: + return orthopoly1d([], [], 1.0, 1.0, wfunc, (0, 1), monic, + lambda x: _ufuncs.eval_sh_legendre(n, x)) + x, w = roots_sh_legendre(n) + hn = 1.0 / (2 * n + 1.0) + kn = _gam(2 * n + 1) / _gam(n + 1)**2 + p = orthopoly1d(x, w, hn, kn, wfunc, limits=(0, 1), monic=monic, + eval_func=lambda x: _ufuncs.eval_sh_legendre(n, x)) + return p + + +# Make the old root function names an alias for the new ones +_modattrs = globals() +for newfun, oldfun in _rootfuns_map.items(): + _modattrs[oldfun] = _modattrs[newfun] + __all__.append(oldfun) diff --git a/v12x/pylib/scipy/special/_precompute/__init__.py b/v12x/pylib/scipy/special/_precompute/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/v12x/pylib/scipy/special/_precompute/cosine_cdf.py b/v12x/pylib/scipy/special/_precompute/cosine_cdf.py new file mode 100644 index 0000000000000000000000000000000000000000..662c12bc74b31478c87471fbd1cce8bea285e765 --- /dev/null +++ b/v12x/pylib/scipy/special/_precompute/cosine_cdf.py @@ -0,0 +1,17 @@ +import mpmath + + +def f(x): + return (mpmath.pi + x + mpmath.sin(x)) / (2*mpmath.pi) + + +# Note: 40 digits might be overkill; a few more digits than the default +# might be sufficient. +mpmath.mp.dps = 40 +ts = mpmath.taylor(f, -mpmath.pi, 20) +p, q = mpmath.pade(ts, 9, 10) + +p = [float(c) for c in p] +q = [float(c) for c in q] +print('p =', p) +print('q =', q) diff --git a/v12x/pylib/scipy/special/_precompute/expn_asy.py b/v12x/pylib/scipy/special/_precompute/expn_asy.py new file mode 100644 index 0000000000000000000000000000000000000000..3491b8acd588a2cacfc48f0a3a60c6ae88c3e8c5 --- /dev/null +++ b/v12x/pylib/scipy/special/_precompute/expn_asy.py @@ -0,0 +1,54 @@ +"""Precompute the polynomials for the asymptotic expansion of the +generalized exponential integral. + +Sources +------- +[1] NIST, Digital Library of Mathematical Functions, + https://dlmf.nist.gov/8.20#ii + +""" +import os + +try: + import sympy + from sympy import Poly + x = sympy.symbols('x') +except ImportError: + pass + + +def generate_A(K): + A = [Poly(1, x)] + for k in range(K): + A.append(Poly(1 - 2*k*x, x)*A[k] + Poly(x*(x + 1))*A[k].diff()) + return A + + +WARNING = """\ +/* This file was automatically generated by _precompute/expn_asy.py. + * Do not edit it manually! + */ +""" + + +def main(): + print(__doc__) + fn = os.path.join('..', 'cephes', 'expn.h') + + K = 12 + A = generate_A(K) + with open(fn + '.new', 'w') as f: + f.write(WARNING) + f.write(f"#define nA {len(A)}\n") + for k, Ak in enumerate(A): + ', '.join([str(x.evalf(18)) for x in Ak.coeffs()]) + f.write(f"static const double A{k}[] = {{tmp}};\n") + ", ".join([f"A{k}" for k in range(K + 1)]) + f.write("static const double *A[] = {{tmp}};\n") + ", ".join([str(Ak.degree()) for Ak in A]) + f.write("static const int Adegs[] = {{tmp}};\n") + os.rename(fn + '.new', fn) + + +if __name__ == "__main__": + main() diff --git a/v12x/pylib/scipy/special/_precompute/gammainc_asy.py b/v12x/pylib/scipy/special/_precompute/gammainc_asy.py new file mode 100644 index 0000000000000000000000000000000000000000..98035457c78706ae01c02273ae1ab458b4ca140d --- /dev/null +++ b/v12x/pylib/scipy/special/_precompute/gammainc_asy.py @@ -0,0 +1,116 @@ +""" +Precompute coefficients of Temme's asymptotic expansion for gammainc. + +This takes about 8 hours to run on a 2.3 GHz Macbook Pro with 4GB ram. + +Sources: +[1] NIST, "Digital Library of Mathematical Functions", + https://dlmf.nist.gov/ + +""" +import os +from scipy.special._precompute.utils import lagrange_inversion + +try: + import mpmath as mp +except ImportError: + pass + + +def compute_a(n): + """a_k from DLMF 5.11.6""" + a = [mp.sqrt(2)/2] + for k in range(1, n): + ak = a[-1]/k + for j in range(1, len(a)): + ak -= a[j]*a[-j]/(j + 1) + ak /= a[0]*(1 + mp.mpf(1)/(k + 1)) + a.append(ak) + return a + + +def compute_g(n): + """g_k from DLMF 5.11.3/5.11.5""" + a = compute_a(2*n) + g = [mp.sqrt(2)*mp.rf(0.5, k)*a[2*k] for k in range(n)] + return g + + +def eta(lam): + """Function from DLMF 8.12.1 shifted to be centered at 0.""" + if lam > 0: + return mp.sqrt(2*(lam - mp.log(lam + 1))) + elif lam < 0: + return -mp.sqrt(2*(lam - mp.log(lam + 1))) + else: + return 0 + + +def compute_alpha(n): + """alpha_n from DLMF 8.12.13""" + coeffs = mp.taylor(eta, 0, n - 1) + return lagrange_inversion(coeffs) + + +def compute_d(K, N): + """d_{k, n} from DLMF 8.12.12""" + M = N + 2*K + d0 = [-mp.mpf(1)/3] + alpha = compute_alpha(M + 2) + for n in range(1, M): + d0.append((n + 2)*alpha[n+2]) + d = [d0] + g = compute_g(K) + for k in range(1, K): + dk = [] + for n in range(M - 2*k): + dk.append((-1)**k*g[k]*d[0][n] + (n + 2)*d[k-1][n+2]) + d.append(dk) + for k in range(K): + d[k] = d[k][:N] + return d + + +header = \ +r"""/* This file was automatically generated by _precomp/gammainc.py. + * Do not edit it manually! + */ + +#ifndef IGAM_H +#define IGAM_H + +#define K {} +#define N {} + +static const double d[K][N] = +{{""" + +footer = \ +r""" +#endif +""" + + +def main(): + print(__doc__) + K = 25 + N = 25 + with mp.workdps(50): + d = compute_d(K, N) + fn = os.path.join(os.path.dirname(__file__), '..', 'cephes', 'igam.h') + with open(fn + '.new', 'w') as f: + f.write(header.format(K, N)) + for k, row in enumerate(d): + row = [mp.nstr(x, 17, min_fixed=0, max_fixed=0) for x in row] + f.write('{') + f.write(", ".join(row)) + if k < K - 1: + f.write('},\n') + else: + f.write('}};\n') + f.write(footer) + os.rename(fn + '.new', fn) + + +if __name__ == "__main__": + main() diff --git a/v12x/pylib/scipy/special/_precompute/gammainc_data.py b/v12x/pylib/scipy/special/_precompute/gammainc_data.py new file mode 100644 index 0000000000000000000000000000000000000000..ebbe39f6159fb80c424dbb38eedeba46bd8cccf2 --- /dev/null +++ b/v12x/pylib/scipy/special/_precompute/gammainc_data.py @@ -0,0 +1,124 @@ +"""Compute gammainc and gammaincc for large arguments and parameters +and save the values to data files for use in tests. We can't just +compare to mpmath's gammainc in test_mpmath.TestSystematic because it +would take too long. + +Note that mpmath's gammainc is computed using hypercomb, but since it +doesn't allow the user to increase the maximum number of terms used in +the series it doesn't converge for many arguments. To get around this +we copy the mpmath implementation but use more terms. + +This takes about 17 minutes to run on a 2.3 GHz Macbook Pro with 4GB +ram. + +Sources: +[1] Fredrik Johansson and others. mpmath: a Python library for + arbitrary-precision floating-point arithmetic (version 0.19), + December 2013. http://mpmath.org/. + +""" +import os +from time import time +import numpy as np +from numpy import pi + +from scipy.special._mptestutils import mpf2float + +try: + import mpmath as mp +except ImportError: + pass + + +def gammainc(a, x, dps=50, maxterms=10**8): + """Compute gammainc exactly like mpmath does but allow for more + summands in hypercomb. See + + mpmath/functions/expintegrals.py#L134 + + in the mpmath GitHub repository. + + """ + with mp.workdps(dps): + z, a, b = mp.mpf(a), mp.mpf(x), mp.mpf(x) + G = [z] + negb = mp.fneg(b, exact=True) + + def h(z): + T1 = [mp.exp(negb), b, z], [1, z, -1], [], G, [1], [1+z], b + return (T1,) + + res = mp.hypercomb(h, [z], maxterms=maxterms) + return mpf2float(res) + + +def gammaincc(a, x, dps=50, maxterms=10**8): + """Compute gammaincc exactly like mpmath does but allow for more + terms in hypercomb. See + + mpmath/functions/expintegrals.py#L187 + + in the mpmath GitHub repository. + + """ + with mp.workdps(dps): + z, a = a, x + + if mp.isint(z): + try: + # mpmath has a fast integer path + return mpf2float(mp.gammainc(z, a=a, regularized=True)) + except mp.libmp.NoConvergence: + pass + nega = mp.fneg(a, exact=True) + G = [z] + # Use 2F0 series when possible; fall back to lower gamma representation + try: + def h(z): + r = z-1 + return [([mp.exp(nega), a], [1, r], [], G, [1, -r], [], 1/nega)] + return mpf2float(mp.hypercomb(h, [z], force_series=True)) + except mp.libmp.NoConvergence: + def h(z): + T1 = [], [1, z-1], [z], G, [], [], 0 + T2 = [-mp.exp(nega), a, z], [1, z, -1], [], G, [1], [1+z], a + return T1, T2 + return mpf2float(mp.hypercomb(h, [z], maxterms=maxterms)) + + +def main(): + t0 = time() + # It would be nice to have data for larger values, but either this + # requires prohibitively large precision (dps > 800) or mpmath has + # a bug. For example, gammainc(1e20, 1e20, dps=800) returns a + # value around 0.03, while the true value should be close to 0.5 + # (DLMF 8.12.15). + print(__doc__) + pwd = os.path.dirname(__file__) + r = np.logspace(4, 14, 30) + ltheta = np.logspace(np.log10(pi/4), np.log10(np.arctan(0.6)), 30) + utheta = np.logspace(np.log10(pi/4), np.log10(np.arctan(1.4)), 30) + + regimes = [(gammainc, ltheta), (gammaincc, utheta)] + for func, theta in regimes: + rg, thetag = np.meshgrid(r, theta) + a, x = rg*np.cos(thetag), rg*np.sin(thetag) + a, x = a.flatten(), x.flatten() + dataset = [] + for i, (a0, x0) in enumerate(zip(a, x)): + if func == gammaincc: + # Exploit the fast integer path in gammaincc whenever + # possible so that the computation doesn't take too + # long + a0, x0 = np.floor(a0), np.floor(x0) + dataset.append((a0, x0, func(a0, x0))) + dataset = np.array(dataset) + filename = os.path.join(pwd, '..', 'tests', 'data', 'local', + f'{func.__name__}.txt') + np.savetxt(filename, dataset) + + print(f"{(time() - t0)/60} minutes elapsed") + + +if __name__ == "__main__": + main() diff --git a/v12x/pylib/scipy/special/_precompute/hyp2f1_data.py b/v12x/pylib/scipy/special/_precompute/hyp2f1_data.py new file mode 100644 index 0000000000000000000000000000000000000000..c4adf14f49184bf75048a28a823909d24e778e04 --- /dev/null +++ b/v12x/pylib/scipy/special/_precompute/hyp2f1_data.py @@ -0,0 +1,484 @@ +"""This script evaluates scipy's implementation of hyp2f1 against mpmath's. + +Author: Albert Steppi + +This script is long running and generates a large output file. With default +arguments, the generated file is roughly 700MB in size and it takes around +40 minutes using an Intel(R) Core(TM) i5-8250U CPU with n_jobs set to 8 +(full utilization). There are optional arguments which can be used to restrict +(or enlarge) the computations performed. These are described below. +The output of this script can be analyzed to identify suitable test cases and +to find parameter and argument regions where hyp2f1 needs to be improved. + +The script has one mandatory positional argument for specifying the path to +the location where the output file is to be placed, and 4 optional arguments +--n_jobs, --grid_size, --regions, and --parameter_groups. --n_jobs specifies +the number of processes to use if running in parallel. The default value is 1. +The other optional arguments are explained below. + +Produces a tab separated values file with 11 columns. The first four columns +contain the parameters a, b, c and the argument z. The next two contain |z| and +a region code for which region of the complex plane belongs to. The regions are + + 0) z == 1 + 1) |z| < 0.9 and real(z) >= 0 + 2) |z| <= 1 and real(z) < 0 + 3) 0.9 <= |z| <= 1 and |1 - z| < 0.9: + 4) 0.9 <= |z| <= 1 and |1 - z| >= 0.9 and real(z) >= 0: + 5) 1 < |z| < 1.1 and |1 - z| >= 0.9 and real(z) >= 0 + 6) |z| > 1 and not in 5) + +The --regions optional argument allows the user to specify a list of regions +to which computation will be restricted. + +Parameters a, b, c are taken from a 10 * 10 * 10 grid with values at + + -16, -8, -4, -2, -1, 1, 2, 4, 8, 16 + +with random perturbations applied. + +There are 9 parameter groups handling the following cases. + + 1) A, B, C, B - A, C - A, C - B, C - A - B all non-integral. + 2) B - A integral + 3) C - A integral + 4) C - B integral + 5) C - A - B integral + 6) A integral + 7) B integral + 8) C integral + 9) Wider range with c - a - b > 0. + +The seventh column of the output file is an integer between 1 and 8 specifying +the parameter group as above. + +The --parameter_groups optional argument allows the user to specify a list of +parameter groups to which computation will be restricted. + +The argument z is taken from a grid in the box + -box_size <= real(z) <= box_size, -box_size <= imag(z) <= box_size. +with grid size specified using the optional command line argument --grid_size, +and box_size specified with the command line argument --box_size. +The default value of grid_size is 20 and the default value of box_size is 2.0, +yielding a 20 * 20 grid in the box with corners -2-2j, -2+2j, 2-2j, 2+2j. + +The final four columns have the expected value of hyp2f1 for the given +parameters and argument as calculated with mpmath, the observed value +calculated with scipy's hyp2f1, the relative error, and the absolute error. + +As special cases of hyp2f1 are moved from the original Fortran implementation +into Cython, this script can be used to ensure that no regressions occur and +to point out where improvements are needed. +""" + + +import os +import csv +import argparse +import numpy as np +from itertools import product +from multiprocessing import Pool + + +from scipy.special import hyp2f1 +from scipy.special.tests.test_hyp2f1 import mp_hyp2f1 + + +def get_region(z): + """Assign numbers for regions where hyp2f1 must be handled differently.""" + if z == 1 + 0j: + return 0 + elif abs(z) < 0.9 and z.real >= 0: + return 1 + elif abs(z) <= 1 and z.real < 0: + return 2 + elif 0.9 <= abs(z) <= 1 and abs(1 - z) < 0.9: + return 3 + elif 0.9 <= abs(z) <= 1 and abs(1 - z) >= 0.9: + return 4 + elif 1 < abs(z) < 1.1 and abs(1 - z) >= 0.9 and z.real >= 0: + return 5 + else: + return 6 + + +def get_result(a, b, c, z, group): + """Get results for given parameter and value combination.""" + expected, observed = mp_hyp2f1(a, b, c, z), hyp2f1(a, b, c, z) + if ( + np.isnan(observed) and np.isnan(expected) or + expected == observed + ): + relative_error = 0.0 + absolute_error = 0.0 + elif np.isnan(observed): + # Set error to infinity if result is nan when not expected to be. + # Makes results easier to interpret. + relative_error = float("inf") + absolute_error = float("inf") + else: + absolute_error = abs(expected - observed) + relative_error = absolute_error / abs(expected) + + return ( + a, + b, + c, + z, + abs(z), + get_region(z), + group, + expected, + observed, + relative_error, + absolute_error, + ) + + +def get_result_no_mp(a, b, c, z, group): + """Get results for given parameter and value combination.""" + expected, observed = complex('nan'), hyp2f1(a, b, c, z) + relative_error, absolute_error = float('nan'), float('nan') + return ( + a, + b, + c, + z, + abs(z), + get_region(z), + group, + expected, + observed, + relative_error, + absolute_error, + ) + + +def get_results(params, Z, n_jobs=1, compute_mp=True): + """Batch compute results for multiple parameter and argument values. + + Parameters + ---------- + params : iterable + iterable of tuples of floats (a, b, c) specifying parameter values + a, b, c for hyp2f1 + Z : iterable of complex + Arguments at which to evaluate hyp2f1 + n_jobs : Optional[int] + Number of jobs for parallel execution. + + Returns + ------- + list + List of tuples of results values. See return value in source code + of `get_result`. + """ + input_ = ( + (a, b, c, z, group) for (a, b, c, group), z in product(params, Z) + ) + + with Pool(n_jobs) as pool: + rows = pool.starmap( + get_result if compute_mp else get_result_no_mp, + input_ + ) + return rows + + +def _make_hyp2f1_test_case(a, b, c, z, rtol): + """Generate string for single test case as used in test_hyp2f1.py.""" + expected = mp_hyp2f1(a, b, c, z) + return ( + " pytest.param(\n" + " Hyp2f1TestCase(\n" + f" a={a},\n" + f" b={b},\n" + f" c={c},\n" + f" z={z},\n" + f" expected={expected},\n" + f" rtol={rtol},\n" + " ),\n" + " )," + ) + + +def make_hyp2f1_test_cases(rows): + """Generate string for a list of test cases for test_hyp2f1.py. + + Parameters + ---------- + rows : list + List of lists of the form [a, b, c, z, rtol] where a, b, c, z are + parameters and the argument for hyp2f1 and rtol is an expected + relative error for the associated test case. + + Returns + ------- + str + String for a list of test cases. The output string can be printed + or saved to a file and then copied into an argument for + `pytest.mark.parameterize` within `scipy.special.tests.test_hyp2f1.py`. + """ + result = "[\n" + result += '\n'.join( + _make_hyp2f1_test_case(a, b, c, z, rtol) + for a, b, c, z, rtol in rows + ) + result += "\n]" + return result + + +def main( + outpath, + n_jobs=1, + box_size=2.0, + grid_size=20, + regions=None, + parameter_groups=None, + compute_mp=True, +): + outpath = os.path.realpath(os.path.expanduser(outpath)) + + random_state = np.random.RandomState(1234) + # Parameters a, b, c selected near these values. + root_params = np.array( + [-16, -8, -4, -2, -1, 1, 2, 4, 8, 16] + ) + # Perturbations to apply to root values. + perturbations = 0.1 * random_state.random_sample( + size=(3, len(root_params)) + ) + + params = [] + # Parameter group 1 + # ----------------- + # No integer differences. This has been confirmed for the above seed. + A = root_params + perturbations[0, :] + B = root_params + perturbations[1, :] + C = root_params + perturbations[2, :] + params.extend( + sorted( + ((a, b, c, 1) for a, b, c in product(A, B, C)), + key=lambda x: max(abs(x[0]), abs(x[1])), + ) + ) + + # Parameter group 2 + # ----------------- + # B - A an integer + A = root_params + 0.5 + B = root_params + 0.5 + C = root_params + perturbations[1, :] + params.extend( + sorted( + ((a, b, c, 2) for a, b, c in product(A, B, C)), + key=lambda x: max(abs(x[0]), abs(x[1])), + ) + ) + + # Parameter group 3 + # ----------------- + # C - A an integer + A = root_params + 0.5 + B = root_params + perturbations[1, :] + C = root_params + 0.5 + params.extend( + sorted( + ((a, b, c, 3) for a, b, c in product(A, B, C)), + key=lambda x: max(abs(x[0]), abs(x[1])), + ) + ) + + # Parameter group 4 + # ----------------- + # C - B an integer + A = root_params + perturbations[0, :] + B = root_params + 0.5 + C = root_params + 0.5 + params.extend( + sorted( + ((a, b, c, 4) for a, b, c in product(A, B, C)), + key=lambda x: max(abs(x[0]), abs(x[1])), + ) + ) + + # Parameter group 5 + # ----------------- + # C - A - B an integer + A = root_params + 0.25 + B = root_params + 0.25 + C = root_params + 0.5 + params.extend( + sorted( + ((a, b, c, 5) for a, b, c in product(A, B, C)), + key=lambda x: max(abs(x[0]), abs(x[1])), + ) + ) + + # Parameter group 6 + # ----------------- + # A an integer + A = root_params + B = root_params + perturbations[0, :] + C = root_params + perturbations[1, :] + params.extend( + sorted( + ((a, b, c, 6) for a, b, c in product(A, B, C)), + key=lambda x: max(abs(x[0]), abs(x[1])), + ) + ) + + # Parameter group 7 + # ----------------- + # B an integer + A = root_params + perturbations[0, :] + B = root_params + C = root_params + perturbations[1, :] + params.extend( + sorted( + ((a, b, c, 7) for a, b, c in product(A, B, C)), + key=lambda x: max(abs(x[0]), abs(x[1])), + ) + ) + + # Parameter group 8 + # ----------------- + # C an integer + A = root_params + perturbations[0, :] + B = root_params + perturbations[1, :] + C = root_params + params.extend( + sorted( + ((a, b, c, 8) for a, b, c in product(A, B, C)), + key=lambda x: max(abs(x[0]), abs(x[1])), + ) + ) + + # Parameter group 9 + # ----------------- + # Wide range of magnitudes, c - a - b > 0. + phi = (1 + np.sqrt(5))/2 + P = phi**np.arange(16) + P = np.hstack([-P, P]) + group_9_params = sorted( + ( + (a, b, c, 9) for a, b, c in product(P, P, P) if c - a - b > 0 + ), + key=lambda x: max(abs(x[0]), abs(x[1])), + ) + + if parameter_groups is not None: + # Group 9 params only used if specified in arguments. + params.extend(group_9_params) + params = [ + (a, b, c, group) for a, b, c, group in params + if group in parameter_groups + ] + + # grid_size * grid_size grid in box with corners + # -2 - 2j, -2 + 2j, 2 - 2j, 2 + 2j + X, Y = np.meshgrid( + np.linspace(-box_size, box_size, grid_size), + np.linspace(-box_size, box_size, grid_size) + ) + Z = X + Y * 1j + Z = Z.flatten().tolist() + # Add z = 1 + 0j (region 0). + Z.append(1 + 0j) + if regions is not None: + Z = [z for z in Z if get_region(z) in regions] + + # Evaluate scipy and mpmath's hyp2f1 for all parameter combinations + # above against all arguments in the grid Z + rows = get_results(params, Z, n_jobs=n_jobs, compute_mp=compute_mp) + + with open(outpath, "w", newline="") as f: + writer = csv.writer(f, delimiter="\t") + writer.writerow( + [ + "a", + "b", + "c", + "z", + "|z|", + "region", + "parameter_group", + "expected", # mpmath's hyp2f1 + "observed", # scipy's hyp2f1 + "relative_error", + "absolute_error", + ] + ) + for row in rows: + writer.writerow(row) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Test scipy's hyp2f1 against mpmath's on a grid in the" + " complex plane over a grid of parameter values. Saves output to file" + " specified in positional argument \"outpath\"." + " Caution: With default arguments, the generated output file is" + " roughly 700MB in size. Script may take several hours to finish if" + " \"--n_jobs\" is set to 1." + ) + parser.add_argument( + "outpath", type=str, help="Path to output tsv file." + ) + parser.add_argument( + "--n_jobs", + type=int, + default=1, + help="Number of jobs for multiprocessing.", + ) + parser.add_argument( + "--box_size", + type=float, + default=2.0, + help="hyp2f1 is evaluated in box of side_length 2*box_size centered" + " at the origin." + ) + parser.add_argument( + "--grid_size", + type=int, + default=20, + help="hyp2f1 is evaluated on grid_size * grid_size grid in box of side" + " length 2*box_size centered at the origin." + ) + parser.add_argument( + "--parameter_groups", + type=int, + nargs='+', + default=None, + help="Restrict to supplied parameter groups. See the Docstring for" + " this module for more info on parameter groups. Calculate for all" + " parameter groups by default." + ) + parser.add_argument( + "--regions", + type=int, + nargs='+', + default=None, + help="Restrict to argument z only within the supplied regions. See" + " the Docstring for this module for more info on regions. Calculate" + " for all regions by default." + ) + parser.add_argument( + "--no_mp", + action='store_true', + help="If this flag is set, do not compute results with mpmath. Saves" + " time if results have already been computed elsewhere. Fills in" + " \"expected\" column with None values." + ) + args = parser.parse_args() + compute_mp = not args.no_mp + print(args.parameter_groups) + main( + args.outpath, + n_jobs=args.n_jobs, + box_size=args.box_size, + grid_size=args.grid_size, + parameter_groups=args.parameter_groups, + regions=args.regions, + compute_mp=compute_mp, + ) diff --git a/v12x/pylib/scipy/special/_precompute/lambertw.py b/v12x/pylib/scipy/special/_precompute/lambertw.py new file mode 100644 index 0000000000000000000000000000000000000000..1fdbf35b2cf85f1f7a6e73579546ed5cfe508fa6 --- /dev/null +++ b/v12x/pylib/scipy/special/_precompute/lambertw.py @@ -0,0 +1,68 @@ +"""Compute a Pade approximation for the principal branch of the +Lambert W function around 0 and compare it to various other +approximations. + +""" +import numpy as np + +try: + import mpmath + import matplotlib.pyplot as plt +except ImportError: + pass + + +def lambertw_pade(): + derivs = [mpmath.diff(mpmath.lambertw, 0, n=n) for n in range(6)] + p, q = mpmath.pade(derivs, 3, 2) + return p, q + + +def main(): + print(__doc__) + with mpmath.workdps(50): + p, q = lambertw_pade() + p, q = p[::-1], q[::-1] + print(f"p = {p}") + print(f"q = {q}") + + x, y = np.linspace(-1.5, 1.5, 75), np.linspace(-1.5, 1.5, 75) + x, y = np.meshgrid(x, y) + z = x + 1j*y + lambertw_std = [] + for z0 in z.flatten(): + lambertw_std.append(complex(mpmath.lambertw(z0))) + lambertw_std = np.array(lambertw_std).reshape(x.shape) + + fig, axes = plt.subplots(nrows=3, ncols=1) + # Compare Pade approximation to true result + p = np.array([float(p0) for p0 in p]) + q = np.array([float(q0) for q0 in q]) + pade_approx = np.polyval(p, z)/np.polyval(q, z) + pade_err = abs(pade_approx - lambertw_std) + axes[0].pcolormesh(x, y, pade_err) + # Compare two terms of asymptotic series to true result + asy_approx = np.log(z) - np.log(np.log(z)) + asy_err = abs(asy_approx - lambertw_std) + axes[1].pcolormesh(x, y, asy_err) + # Compare two terms of the series around the branch point to the + # true result + p = np.sqrt(2*(np.exp(1)*z + 1)) + series_approx = -1 + p - p**2/3 + series_err = abs(series_approx - lambertw_std) + im = axes[2].pcolormesh(x, y, series_err) + + fig.colorbar(im, ax=axes.ravel().tolist()) + plt.show() + + fig, ax = plt.subplots(nrows=1, ncols=1) + pade_better = pade_err < asy_err + im = ax.pcolormesh(x, y, pade_better) + t = np.linspace(-0.3, 0.3) + ax.plot(-2.5*abs(t) - 0.2, t, 'r') + fig.colorbar(im, ax=ax) + plt.show() + + +if __name__ == '__main__': + main() diff --git a/v12x/pylib/scipy/special/_precompute/loggamma.py b/v12x/pylib/scipy/special/_precompute/loggamma.py new file mode 100644 index 0000000000000000000000000000000000000000..74051ac7b46c70dc01919a362d05a8bbbe11333a --- /dev/null +++ b/v12x/pylib/scipy/special/_precompute/loggamma.py @@ -0,0 +1,43 @@ +"""Precompute series coefficients for log-Gamma.""" + +try: + import mpmath +except ImportError: + pass + + +def stirling_series(N): + with mpmath.workdps(100): + coeffs = [mpmath.bernoulli(2*n)/(2*n*(2*n - 1)) + for n in range(1, N + 1)] + return coeffs + + +def taylor_series_at_1(N): + coeffs = [] + with mpmath.workdps(100): + coeffs.append(-mpmath.euler) + for n in range(2, N + 1): + coeffs.append((-1)**n*mpmath.zeta(n)/n) + return coeffs + + +def main(): + print(__doc__) + print() + stirling_coeffs = [mpmath.nstr(x, 20, min_fixed=0, max_fixed=0) + for x in stirling_series(8)[::-1]] + taylor_coeffs = [mpmath.nstr(x, 20, min_fixed=0, max_fixed=0) + for x in taylor_series_at_1(23)[::-1]] + print("Stirling series coefficients") + print("----------------------------") + print("\n".join(stirling_coeffs)) + print() + print("Taylor series coefficients") + print("--------------------------") + print("\n".join(taylor_coeffs)) + print() + + +if __name__ == '__main__': + main() diff --git a/v12x/pylib/scipy/special/_precompute/struve_convergence.py b/v12x/pylib/scipy/special/_precompute/struve_convergence.py new file mode 100644 index 0000000000000000000000000000000000000000..dbf6009368540dbf603b61f5b72510f0acd1a65b --- /dev/null +++ b/v12x/pylib/scipy/special/_precompute/struve_convergence.py @@ -0,0 +1,131 @@ +""" +Convergence regions of the expansions used in ``struve.c`` + +Note that for v >> z both functions tend rapidly to 0, +and for v << -z, they tend to infinity. + +The floating-point functions over/underflow in the lower left and right +corners of the figure. + + +Figure legend +============= + +Red region + Power series is close (1e-12) to the mpmath result + +Blue region + Asymptotic series is close to the mpmath result + +Green region + Bessel series is close to the mpmath result + +Dotted colored lines + Boundaries of the regions + +Solid colored lines + Boundaries estimated by the routine itself. These will be used + for determining which of the results to use. + +Black dashed line + The line z = 0.7*|v| + 12 + +""" +import numpy as np +import matplotlib.pyplot as plt + +import mpmath + + +def err_metric(a, b, atol=1e-290): + m = abs(a - b) / (atol + abs(b)) + m[np.isinf(b) & (a == b)] = 0 + return m + + +def do_plot(is_h=True): + from scipy.special._ufuncs import (_struve_power_series, + _struve_asymp_large_z, + _struve_bessel_series) + + vs = np.linspace(-1000, 1000, 91) + zs = np.sort(np.r_[1e-5, 1.0, np.linspace(0, 700, 91)[1:]]) + + rp = _struve_power_series(vs[:,None], zs[None,:], is_h) + ra = _struve_asymp_large_z(vs[:,None], zs[None,:], is_h) + rb = _struve_bessel_series(vs[:,None], zs[None,:], is_h) + + mpmath.mp.dps = 50 + if is_h: + def sh(v, z): + return float(mpmath.struveh(mpmath.mpf(v), mpmath.mpf(z))) + else: + def sh(v, z): + return float(mpmath.struvel(mpmath.mpf(v), mpmath.mpf(z))) + ex = np.vectorize(sh, otypes='d')(vs[:,None], zs[None,:]) + + err_a = err_metric(ra[0], ex) + 1e-300 + err_p = err_metric(rp[0], ex) + 1e-300 + err_b = err_metric(rb[0], ex) + 1e-300 + + err_est_a = abs(ra[1]/ra[0]) + err_est_p = abs(rp[1]/rp[0]) + err_est_b = abs(rb[1]/rb[0]) + + z_cutoff = 0.7*abs(vs) + 12 + + levels = [-1000, -12] + + plt.cla() + + plt.hold(1) + plt.contourf(vs, zs, np.log10(err_p).T, + levels=levels, colors=['r', 'r'], alpha=0.1) + plt.contourf(vs, zs, np.log10(err_a).T, + levels=levels, colors=['b', 'b'], alpha=0.1) + plt.contourf(vs, zs, np.log10(err_b).T, + levels=levels, colors=['g', 'g'], alpha=0.1) + + plt.contour(vs, zs, np.log10(err_p).T, + levels=levels, colors=['r', 'r'], linestyles=[':', ':']) + plt.contour(vs, zs, np.log10(err_a).T, + levels=levels, colors=['b', 'b'], linestyles=[':', ':']) + plt.contour(vs, zs, np.log10(err_b).T, + levels=levels, colors=['g', 'g'], linestyles=[':', ':']) + + lp = plt.contour(vs, zs, np.log10(err_est_p).T, + levels=levels, colors=['r', 'r'], linestyles=['-', '-']) + la = plt.contour(vs, zs, np.log10(err_est_a).T, + levels=levels, colors=['b', 'b'], linestyles=['-', '-']) + lb = plt.contour(vs, zs, np.log10(err_est_b).T, + levels=levels, colors=['g', 'g'], linestyles=['-', '-']) + + plt.clabel(lp, fmt={-1000: 'P', -12: 'P'}) + plt.clabel(la, fmt={-1000: 'A', -12: 'A'}) + plt.clabel(lb, fmt={-1000: 'B', -12: 'B'}) + + plt.plot(vs, z_cutoff, 'k--') + + plt.xlim(vs.min(), vs.max()) + plt.ylim(zs.min(), zs.max()) + + plt.xlabel('v') + plt.ylabel('z') + + +def main(): + plt.clf() + plt.subplot(121) + do_plot(True) + plt.title('Struve H') + + plt.subplot(122) + do_plot(False) + plt.title('Struve L') + + plt.savefig('struve_convergence.png') + plt.show() + + +if __name__ == "__main__": + main() diff --git a/v12x/pylib/scipy/special/_precompute/utils.py b/v12x/pylib/scipy/special/_precompute/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..55cf4083ed5e5a6628fd3316c02ce1a5ce21a92c --- /dev/null +++ b/v12x/pylib/scipy/special/_precompute/utils.py @@ -0,0 +1,38 @@ +try: + import mpmath as mp +except ImportError: + pass + +try: + from sympy.abc import x +except ImportError: + pass + + +def lagrange_inversion(a): + """Given a series + + f(x) = a[1]*x + a[2]*x**2 + ... + a[n-1]*x**(n - 1), + + use the Lagrange inversion formula to compute a series + + g(x) = b[1]*x + b[2]*x**2 + ... + b[n-1]*x**(n - 1) + + so that f(g(x)) = g(f(x)) = x mod x**n. We must have a[0] = 0, so + necessarily b[0] = 0 too. + + The algorithm is naive and could be improved, but speed isn't an + issue here and it's easy to read. + + """ + n = len(a) + f = sum(a[i]*x**i for i in range(n)) + h = (x/f).series(x, 0, n).removeO() + hpower = [h**0] + for k in range(n): + hpower.append((hpower[-1]*h).expand()) + b = [mp.mpf(0)] + for k in range(1, n): + b.append(hpower[k].coeff(x, k - 1)/k) + b = [mp.mpf(x) for x in b] + return b diff --git a/v12x/pylib/scipy/special/_precompute/wright_bessel.py b/v12x/pylib/scipy/special/_precompute/wright_bessel.py new file mode 100644 index 0000000000000000000000000000000000000000..51d56b1cd5c47c7ef005d21aad9827a1e85ec0d9 --- /dev/null +++ b/v12x/pylib/scipy/special/_precompute/wright_bessel.py @@ -0,0 +1,342 @@ +"""Precompute coefficients of several series expansions +of Wright's generalized Bessel function Phi(a, b, x). + +See https://dlmf.nist.gov/10.46.E1 with rho=a, beta=b, z=x. +""" +from argparse import ArgumentParser, RawTextHelpFormatter +import numpy as np +from scipy.integrate import quad +from scipy.optimize import minimize_scalar, curve_fit +from time import time + +try: + import sympy + from sympy import EulerGamma, Rational, S, Sum, \ + factorial, gamma, gammasimp, pi, polygamma, symbols, zeta + from sympy.polys.polyfuncs import horner +except ImportError: + pass + + +def series_small_a(): + """Tylor series expansion of Phi(a, b, x) in a=0 up to order 5. + """ + order = 5 + a, b, x, k = symbols("a b x k") + A = [] # terms with a + X = [] # terms with x + B = [] # terms with b (polygammas) + # Phi(a, b, x) = exp(x)/gamma(b) * sum(A[i] * X[i] * B[i]) + expression = Sum(x**k/factorial(k)/gamma(a*k+b), (k, 0, S.Infinity)) + expression = gamma(b)/sympy.exp(x) * expression + + # nth term of taylor series in a=0: a^n/n! * (d^n Phi(a, b, x)/da^n at a=0) + for n in range(0, order+1): + term = expression.diff(a, n).subs(a, 0).simplify().doit() + # set the whole bracket involving polygammas to 1 + x_part = (term.subs(polygamma(0, b), 1) + .replace(polygamma, lambda *args: 0)) + # sign convention: x part always positive + x_part *= (-1)**n + + A.append(a**n/factorial(n)) + X.append(horner(x_part)) + B.append(horner((term/x_part).simplify())) + + s = "Tylor series expansion of Phi(a, b, x) in a=0 up to order 5.\n" + s += "Phi(a, b, x) = exp(x)/gamma(b) * sum(A[i] * X[i] * B[i], i=0..5)\n" + for name, c in zip(['A', 'X', 'B'], [A, X, B]): + for i in range(len(c)): + s += f"\n{name}[{i}] = " + str(c[i]) + return s + + +# expansion of digamma +def dg_series(z, n): + """Symbolic expansion of digamma(z) in z=0 to order n. + + See https://dlmf.nist.gov/5.7.E4 and with https://dlmf.nist.gov/5.5.E2 + """ + k = symbols("k") + return -1/z - EulerGamma + \ + sympy.summation((-1)**k * zeta(k) * z**(k-1), (k, 2, n+1)) + + +def pg_series(k, z, n): + """Symbolic expansion of polygamma(k, z) in z=0 to order n.""" + return sympy.diff(dg_series(z, n+k), z, k) + + +def series_small_a_small_b(): + """Tylor series expansion of Phi(a, b, x) in a=0 and b=0 up to order 5. + + Be aware of cancellation of poles in b=0 of digamma(b)/Gamma(b) and + polygamma functions. + + digamma(b)/Gamma(b) = -1 - 2*M_EG*b + O(b^2) + digamma(b)^2/Gamma(b) = 1/b + 3*M_EG + b*(-5/12*PI^2+7/2*M_EG^2) + O(b^2) + polygamma(1, b)/Gamma(b) = 1/b + M_EG + b*(1/12*PI^2 + 1/2*M_EG^2) + O(b^2) + and so on. + """ + order = 5 + a, b, x, k = symbols("a b x k") + M_PI, M_EG, M_Z3 = symbols("M_PI M_EG M_Z3") + c_subs = {pi: M_PI, EulerGamma: M_EG, zeta(3): M_Z3} + A = [] # terms with a + X = [] # terms with x + B = [] # terms with b (polygammas expanded) + C = [] # terms that generate B + # Phi(a, b, x) = exp(x) * sum(A[i] * X[i] * B[i]) + # B[0] = 1 + # B[k] = sum(C[k] * b**k/k!, k=0..) + # Note: C[k] can be obtained from a series expansion of 1/gamma(b). + expression = gamma(b)/sympy.exp(x) * \ + Sum(x**k/factorial(k)/gamma(a*k+b), (k, 0, S.Infinity)) + + # nth term of taylor series in a=0: a^n/n! * (d^n Phi(a, b, x)/da^n at a=0) + for n in range(0, order+1): + term = expression.diff(a, n).subs(a, 0).simplify().doit() + # set the whole bracket involving polygammas to 1 + x_part = (term.subs(polygamma(0, b), 1) + .replace(polygamma, lambda *args: 0)) + # sign convention: x part always positive + x_part *= (-1)**n + # expansion of polygamma part with 1/gamma(b) + pg_part = term/x_part/gamma(b) + if n >= 1: + # Note: highest term is digamma^n + pg_part = pg_part.replace(polygamma, + lambda k, x: pg_series(k, x, order+1+n)) + pg_part = (pg_part.series(b, 0, n=order+1-n) + .removeO() + .subs(polygamma(2, 1), -2*zeta(3)) + .simplify() + ) + + A.append(a**n/factorial(n)) + X.append(horner(x_part)) + B.append(pg_part) + + # Calculate C and put in the k! + C = sympy.Poly(B[1].subs(c_subs), b).coeffs() + C.reverse() + for i in range(len(C)): + C[i] = (C[i] * factorial(i)).simplify() + + s = "Tylor series expansion of Phi(a, b, x) in a=0 and b=0 up to order 5." + s += "\nPhi(a, b, x) = exp(x) * sum(A[i] * X[i] * B[i], i=0..5)\n" + s += "B[0] = 1\n" + s += "B[i] = sum(C[k+i-1] * b**k/k!, k=0..)\n" + s += "\nM_PI = pi" + s += "\nM_EG = EulerGamma" + s += "\nM_Z3 = zeta(3)" + for name, c in zip(['A', 'X'], [A, X]): + for i in range(len(c)): + s += f"\n{name}[{i}] = " + s += str(c[i]) + # For C, do also compute the values numerically + for i in range(len(C)): + s += f"\n# C[{i}] = " + s += str(C[i]) + s += f"\nC[{i}] = " + s += str(C[i].subs({M_EG: EulerGamma, M_PI: pi, M_Z3: zeta(3)}) + .evalf(17)) + + # Does B have the assumed structure? + s += "\n\nTest if B[i] does have the assumed structure." + s += "\nC[i] are derived from B[1] alone." + s += "\nTest B[2] == C[1] + b*C[2] + b^2/2*C[3] + b^3/6*C[4] + .." + test = sum([b**k/factorial(k) * C[k+1] for k in range(order-1)]) + test = (test - B[2].subs(c_subs)).simplify() + s += f"\ntest successful = {test==S(0)}" + s += "\nTest B[3] == C[2] + b*C[3] + b^2/2*C[4] + .." + test = sum([b**k/factorial(k) * C[k+2] for k in range(order-2)]) + test = (test - B[3].subs(c_subs)).simplify() + s += f"\ntest successful = {test==S(0)}" + return s + + +def asymptotic_series(): + """Asymptotic expansion for large x. + + Phi(a, b, x) ~ Z^(1/2-b) * exp((1+a)/a * Z) * sum_k (-1)^k * C_k / Z^k + Z = (a*x)^(1/(1+a)) + + Wright (1935) lists the coefficients C_0 and C_1 (he calls them a_0 and + a_1). With slightly different notation, Paris (2017) lists coefficients + c_k up to order k=3. + Paris (2017) uses ZP = (1+a)/a * Z (ZP = Z of Paris) and + C_k = C_0 * (-a/(1+a))^k * c_k + """ + order = 8 + + class g(sympy.Function): + """Helper function g according to Wright (1935) + + g(n, rho, v) = (1 + (rho+2)/3 * v + (rho+2)*(rho+3)/(2*3) * v^2 + ...) + + Note: Wright (1935) uses square root of above definition. + """ + nargs = 3 + + @classmethod + def eval(cls, n, rho, v): + if not n >= 0: + raise ValueError("must have n >= 0") + elif n == 0: + return 1 + else: + return g(n-1, rho, v) \ + + gammasimp(gamma(rho+2+n)/gamma(rho+2)) \ + / gammasimp(gamma(3+n)/gamma(3))*v**n + + class coef_C(sympy.Function): + """Calculate coefficients C_m for integer m. + + C_m is the coefficient of v^(2*m) in the Taylor expansion in v=0 of + Gamma(m+1/2)/(2*pi) * (2/(rho+1))^(m+1/2) * (1-v)^(-b) + * g(rho, v)^(-m-1/2) + """ + nargs = 3 + + @classmethod + def eval(cls, m, rho, beta): + if not m >= 0: + raise ValueError("must have m >= 0") + + v = symbols("v") + expression = (1-v)**(-beta) * g(2*m, rho, v)**(-m-Rational(1, 2)) + res = expression.diff(v, 2*m).subs(v, 0) / factorial(2*m) + res = res * (gamma(m + Rational(1, 2)) / (2*pi) + * (2/(rho+1))**(m + Rational(1, 2))) + return res + + # in order to have nice ordering/sorting of expressions, we set a = xa. + xa, b, xap1 = symbols("xa b xap1") + C0 = coef_C(0, xa, b) + # a1 = a(1, rho, beta) + s = "Asymptotic expansion for large x\n" + s += "Phi(a, b, x) = Z**(1/2-b) * exp((1+a)/a * Z) \n" + s += " * sum((-1)**k * C[k]/Z**k, k=0..6)\n\n" + s += "Z = pow(a * x, 1/(1+a))\n" + s += "A[k] = pow(a, k)\n" + s += "B[k] = pow(b, k)\n" + s += "Ap1[k] = pow(1+a, k)\n\n" + s += "C[0] = 1./sqrt(2. * M_PI * Ap1[1])\n" + for i in range(1, order+1): + expr = (coef_C(i, xa, b) / (C0/(1+xa)**i)).simplify() + factor = [x.denominator() for x in sympy.Poly(expr).coeffs()] + factor = sympy.lcm(factor) + expr = (expr * factor).simplify().collect(b, sympy.factor) + expr = expr.xreplace({xa+1: xap1}) + s += f"C[{i}] = C[0] / ({factor} * Ap1[{i}])\n" + s += f"C[{i}] *= {str(expr)}\n\n" + import re + re_a = re.compile(r'xa\*\*(\d+)') + s = re_a.sub(r'A[\1]', s) + re_b = re.compile(r'b\*\*(\d+)') + s = re_b.sub(r'B[\1]', s) + s = s.replace('xap1', 'Ap1[1]') + s = s.replace('xa', 'a') + # max integer = 2^31-1 = 2,147,483,647. Solution: Put a point after 10 + # or more digits. + re_digits = re.compile(r'(\d{10,})') + s = re_digits.sub(r'\1.', s) + return s + + +def optimal_epsilon_integral(): + """Fit optimal choice of epsilon for integral representation. + + The integrand of + int_0^pi P(eps, a, b, x, phi) * dphi + can exhibit oscillatory behaviour. It stems from the cosine of P and can be + minimized by minimizing the arc length of the argument + f(phi) = eps * sin(phi) - x * eps^(-a) * sin(a * phi) + (1 - b) * phi + of cos(f(phi)). + We minimize the arc length in eps for a grid of values (a, b, x) and fit a + parametric function to it. + """ + def fp(eps, a, b, x, phi): + """Derivative of f w.r.t. phi.""" + eps_a = np.power(1. * eps, -a) + return eps * np.cos(phi) - a * x * eps_a * np.cos(a * phi) + 1 - b + + def arclength(eps, a, b, x, epsrel=1e-2, limit=100): + """Compute Arc length of f. + + Note that the arc length of a function f from t0 to t1 is given by + int_t0^t1 sqrt(1 + f'(t)^2) dt + """ + return quad(lambda phi: np.sqrt(1 + fp(eps, a, b, x, phi)**2), + 0, np.pi, + epsrel=epsrel, limit=100)[0] + + # grid of minimal arc length values + data_a = [1e-3, 0.1, 0.5, 0.9, 1, 2, 4, 5, 6, 8] + data_b = [0, 1, 4, 7, 10] + data_x = [1, 1.5, 2, 4, 10, 20, 50, 100, 200, 500, 1e3, 5e3, 1e4] + data_a, data_b, data_x = np.meshgrid(data_a, data_b, data_x) + data_a, data_b, data_x = (data_a.flatten(), data_b.flatten(), + data_x.flatten()) + best_eps = [] + for i in range(data_x.size): + best_eps.append( + minimize_scalar(lambda eps: arclength(eps, data_a[i], data_b[i], + data_x[i]), + bounds=(1e-3, 1000), + method='Bounded', options={'xatol': 1e-3}).x + ) + best_eps = np.array(best_eps) + # pandas would be nice, but here a dictionary is enough + df = {'a': data_a, + 'b': data_b, + 'x': data_x, + 'eps': best_eps, + } + + def func(data, A0, A1, A2, A3, A4, A5): + """Compute parametric function to fit.""" + a = data['a'] + b = data['b'] + x = data['x'] + return (A0 * b * np.exp(-0.5 * a) + + np.exp(A1 + 1 / (1 + a) * np.log(x) - A2 * np.exp(-A3 * a) + + A4 / (1 + np.exp(A5 * a)))) + + func_params = list(curve_fit(func, df, df['eps'], method='trf')[0]) + + s = "Fit optimal eps for integrand P via minimal arc length\n" + s += "with parametric function:\n" + s += "optimal_eps = (A0 * b * exp(-a/2) + exp(A1 + 1 / (1 + a) * log(x)\n" + s += " - A2 * exp(-A3 * a) + A4 / (1 + exp(A5 * a)))\n\n" + s += "Fitted parameters A0 to A5 are:\n" + s += ', '.join([f'{x:.5g}' for x in func_params]) + return s + + +def main(): + t0 = time() + parser = ArgumentParser(description=__doc__, + formatter_class=RawTextHelpFormatter) + parser.add_argument('action', type=int, choices=[1, 2, 3, 4], + help='chose what expansion to precompute\n' + '1 : Series for small a\n' + '2 : Series for small a and small b\n' + '3 : Asymptotic series for large x\n' + ' This may take some time (>4h).\n' + '4 : Fit optimal eps for integral representation.' + ) + args = parser.parse_args() + + switch = {1: lambda: print(series_small_a()), + 2: lambda: print(series_small_a_small_b()), + 3: lambda: print(asymptotic_series()), + 4: lambda: print(optimal_epsilon_integral()) + } + switch.get(args.action, lambda: print("Invalid input."))() + print(f"\n{(time() - t0)/60:.1f} minutes elapsed.\n") + + +if __name__ == '__main__': + main() diff --git a/v12x/pylib/scipy/special/_precompute/wright_bessel_data.py b/v12x/pylib/scipy/special/_precompute/wright_bessel_data.py new file mode 100644 index 0000000000000000000000000000000000000000..1de9b4fe552ca9178c452194ab84af6ca5daac71 --- /dev/null +++ b/v12x/pylib/scipy/special/_precompute/wright_bessel_data.py @@ -0,0 +1,152 @@ +"""Compute a grid of values for Wright's generalized Bessel function +and save the values to data files for use in tests. Using mpmath directly in +tests would take too long. + +This takes about 10 minutes to run on a 2.7 GHz i7 Macbook Pro. +""" +from functools import lru_cache +import os +from time import time + +import numpy as np +from scipy.special._mptestutils import mpf2float + +try: + import mpmath as mp +except ImportError: + pass + +# exp_inf: smallest value x for which exp(x) == inf +exp_inf = 709.78271289338403 + + +# 64 Byte per value +@lru_cache(maxsize=100_000) +def rgamma_cached(x, dps): + with mp.workdps(dps): + return mp.rgamma(x) + + +def mp_wright_bessel(a, b, x, dps=50, maxterms=2000): + """Compute Wright's generalized Bessel function as Series with mpmath. + """ + with mp.workdps(dps): + a, b, x = mp.mpf(a), mp.mpf(b), mp.mpf(x) + res = mp.nsum(lambda k: x**k / mp.fac(k) + * rgamma_cached(a * k + b, dps=dps), + [0, mp.inf], + tol=dps, method='s', steps=[maxterms] + ) + return mpf2float(res) + + +def main(): + t0 = time() + print(__doc__) + pwd = os.path.dirname(__file__) + eps = np.finfo(float).eps * 100 + + a_range = np.array([eps, + 1e-4 * (1 - eps), 1e-4, 1e-4 * (1 + eps), + 1e-3 * (1 - eps), 1e-3, 1e-3 * (1 + eps), + 0.1, 0.5, + 1 * (1 - eps), 1, 1 * (1 + eps), + 1.5, 2, 4.999, 5, 10]) + b_range = np.array([0, eps, 1e-10, 1e-5, 0.1, 1, 2, 10, 20, 100]) + x_range = np.array([0, eps, 1 - eps, 1, 1 + eps, + 1.5, + 2 - eps, 2, 2 + eps, + 9 - eps, 9, 9 + eps, + 10 * (1 - eps), 10, 10 * (1 + eps), + 100 * (1 - eps), 100, 100 * (1 + eps), + 500, exp_inf, 1e3, 1e5, 1e10, 1e20]) + + a_range, b_range, x_range = np.meshgrid(a_range, b_range, x_range, + indexing='ij') + a_range = a_range.flatten() + b_range = b_range.flatten() + x_range = x_range.flatten() + + # filter out some values, especially too large x + bool_filter = ~((a_range < 5e-3) & (x_range >= exp_inf)) + bool_filter = bool_filter & ~((a_range < 0.2) & (x_range > exp_inf)) + bool_filter = bool_filter & ~((a_range < 0.5) & (x_range > 1e3)) + bool_filter = bool_filter & ~((a_range < 0.56) & (x_range > 5e3)) + bool_filter = bool_filter & ~((a_range < 1) & (x_range > 1e4)) + bool_filter = bool_filter & ~((a_range < 1.4) & (x_range > 1e5)) + bool_filter = bool_filter & ~((a_range < 1.8) & (x_range > 1e6)) + bool_filter = bool_filter & ~((a_range < 2.2) & (x_range > 1e7)) + bool_filter = bool_filter & ~((a_range < 2.5) & (x_range > 1e8)) + bool_filter = bool_filter & ~((a_range < 2.9) & (x_range > 1e9)) + bool_filter = bool_filter & ~((a_range < 3.3) & (x_range > 1e10)) + bool_filter = bool_filter & ~((a_range < 3.7) & (x_range > 1e11)) + bool_filter = bool_filter & ~((a_range < 4) & (x_range > 1e12)) + bool_filter = bool_filter & ~((a_range < 4.4) & (x_range > 1e13)) + bool_filter = bool_filter & ~((a_range < 4.7) & (x_range > 1e14)) + bool_filter = bool_filter & ~((a_range < 5.1) & (x_range > 1e15)) + bool_filter = bool_filter & ~((a_range < 5.4) & (x_range > 1e16)) + bool_filter = bool_filter & ~((a_range < 5.8) & (x_range > 1e17)) + bool_filter = bool_filter & ~((a_range < 6.2) & (x_range > 1e18)) + bool_filter = bool_filter & ~((a_range < 6.2) & (x_range > 1e18)) + bool_filter = bool_filter & ~((a_range < 6.5) & (x_range > 1e19)) + bool_filter = bool_filter & ~((a_range < 6.9) & (x_range > 1e20)) + + # filter out known values that do not meet the required numerical accuracy + # see test test_wright_data_grid_failures + failing = np.array([ + [0.1, 100, 709.7827128933841], + [0.5, 10, 709.7827128933841], + [0.5, 10, 1000], + [0.5, 100, 1000], + [1, 20, 100000], + [1, 100, 100000], + [1.0000000000000222, 20, 100000], + [1.0000000000000222, 100, 100000], + [1.5, 0, 500], + [1.5, 2.220446049250313e-14, 500], + [1.5, 1.e-10, 500], + [1.5, 1.e-05, 500], + [1.5, 0.1, 500], + [1.5, 20, 100000], + [1.5, 100, 100000], + ]).tolist() + + does_fail = np.full_like(a_range, False, dtype=bool) + for i in range(x_range.size): + if [a_range[i], b_range[i], x_range[i]] in failing: + does_fail[i] = True + + # filter and flatten + a_range = a_range[bool_filter] + b_range = b_range[bool_filter] + x_range = x_range[bool_filter] + does_fail = does_fail[bool_filter] + + dataset = [] + print(f"Computing {x_range.size} single points.") + print("Tests will fail for the following data points:") + for i in range(x_range.size): + a = a_range[i] + b = b_range[i] + x = x_range[i] + # take care of difficult corner cases + maxterms = 1000 + if a < 1e-6 and x >= exp_inf/10: + maxterms = 2000 + f = mp_wright_bessel(a, b, x, maxterms=maxterms) + if does_fail[i]: + print("failing data point a, b, x, value = " + f"[{a}, {b}, {x}, {f}]") + else: + dataset.append((a, b, x, f)) + dataset = np.array(dataset) + + filename = os.path.join(pwd, '..', 'tests', 'data', 'local', + 'wright_bessel.txt') + np.savetxt(filename, dataset) + + print(f"{(time() - t0)/60:.1f} minutes elapsed") + + +if __name__ == "__main__": + main() diff --git a/v12x/pylib/scipy/special/_precompute/wrightomega.py b/v12x/pylib/scipy/special/_precompute/wrightomega.py new file mode 100644 index 0000000000000000000000000000000000000000..0bcd0345a9c1b90c45b0e9e3340ab4da4ec5c6d7 --- /dev/null +++ b/v12x/pylib/scipy/special/_precompute/wrightomega.py @@ -0,0 +1,41 @@ +import numpy as np + +try: + import mpmath +except ImportError: + pass + + +def mpmath_wrightomega(x): + return mpmath.lambertw(mpmath.exp(x), mpmath.mpf('-0.5')) + + +def wrightomega_series_error(x): + series = x + desired = mpmath_wrightomega(x) + return abs(series - desired) / desired + + +def wrightomega_exp_error(x): + exponential_approx = mpmath.exp(x) + desired = mpmath_wrightomega(x) + return abs(exponential_approx - desired) / desired + + +def main(): + desired_error = 2 * np.finfo(float).eps + print('Series Error') + for x in [1e5, 1e10, 1e15, 1e20]: + with mpmath.workdps(100): + error = wrightomega_series_error(x) + print(x, error, error < desired_error) + + print('Exp error') + for x in [-10, -25, -50, -100, -200, -400, -700, -740]: + with mpmath.workdps(100): + error = wrightomega_exp_error(x) + print(x, error, error < desired_error) + + +if __name__ == '__main__': + main() diff --git a/v12x/pylib/scipy/special/_precompute/zetac.py b/v12x/pylib/scipy/special/_precompute/zetac.py new file mode 100644 index 0000000000000000000000000000000000000000..d408b1a2fffb6872287452923fcc9394adc13a7c --- /dev/null +++ b/v12x/pylib/scipy/special/_precompute/zetac.py @@ -0,0 +1,27 @@ +"""Compute the Taylor series for zeta(x) - 1 around x = 0.""" +try: + import mpmath +except ImportError: + pass + + +def zetac_series(N): + coeffs = [] + with mpmath.workdps(100): + coeffs.append(-1.5) + for n in range(1, N): + coeff = mpmath.diff(mpmath.zeta, 0, n)/mpmath.factorial(n) + coeffs.append(coeff) + return coeffs + + +def main(): + print(__doc__) + coeffs = zetac_series(10) + coeffs = [mpmath.nstr(x, 20, min_fixed=0, max_fixed=0) + for x in coeffs] + print("\n".join(coeffs[::-1])) + + +if __name__ == '__main__': + main() diff --git a/v12x/pylib/scipy/special/_sf_error.py b/v12x/pylib/scipy/special/_sf_error.py new file mode 100644 index 0000000000000000000000000000000000000000..e1edc9800759dfda9e49bde1becc775a64bce958 --- /dev/null +++ b/v12x/pylib/scipy/special/_sf_error.py @@ -0,0 +1,15 @@ +"""Warnings and Exceptions that can be raised by special functions.""" +import warnings + + +class SpecialFunctionWarning(Warning): + """Warning that can be emitted by special functions.""" + pass + + +warnings.simplefilter("always", category=SpecialFunctionWarning) + + +class SpecialFunctionError(Exception): + """Exception that can be raised by special functions.""" + pass diff --git a/v12x/pylib/scipy/special/_spfun_stats.py b/v12x/pylib/scipy/special/_spfun_stats.py new file mode 100644 index 0000000000000000000000000000000000000000..0b54f474ffd5fedeaa5cc44ce363e0761818f216 --- /dev/null +++ b/v12x/pylib/scipy/special/_spfun_stats.py @@ -0,0 +1,178 @@ +# Last Change: Sat Mar 21 02:00 PM 2009 J + +# Copyright (c) 2001, 2002 Enthought, Inc. +# +# All rights reserved. +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# a. Redistributions of source code must retain the above copyright notice, +# this list of conditions and the following disclaimer. +# b. Redistributions in binary form must reproduce the above copyright +# notice, this list of conditions and the following disclaimer in the +# documentation and/or other materials provided with the distribution. +# c. Neither the name of the Enthought nor the names of its contributors +# may be used to endorse or promote products derived from this software +# without specific prior written permission. +# +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +# ARE DISCLAIMED. IN NO EVENT SHALL THE REGENTS OR CONTRIBUTORS BE LIABLE FOR +# ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT +# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY +# OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH +# DAMAGE. + +"""Some more special functions which may be useful for multivariate statistical +analysis.""" + +import numpy as np + +import scipy.special._gufuncs as _gufuncs +from scipy.special import gammaln as loggam +from scipy.special._ufunc_tools import _with_cache_optimization + + +__all__ = ['multigammaln'] + + +def multigammaln(a, d): + r"""Returns the log of multivariate gamma, also sometimes called the + generalized gamma. + + Parameters + ---------- + a : ndarray + The multivariate gamma is computed for each item of `a`. + d : int + The dimension of the space of integration. + + Returns + ------- + res : ndarray + The values of the log multivariate gamma at the given points `a`. + + Notes + ----- + The formal definition of the multivariate gamma of dimension d for a real + `a` is + + .. math:: + + \Gamma_d(a) = \int_{A>0} e^{-tr(A)} |A|^{a - (d+1)/2} dA + + with the condition :math:`a > (d-1)/2`, and :math:`A > 0` being the set of + all the positive definite matrices of dimension `d`. Note that `a` is a + scalar: the integrand only is multivariate, the argument is not (the + function is defined over a subset of the real set). + + This can be proven to be equal to the much friendlier equation + + .. math:: + + \Gamma_d(a) = \pi^{d(d-1)/4} \prod_{i=1}^{d} \Gamma(a - (i-1)/2). + + References + ---------- + R. J. Muirhead, Aspects of multivariate statistical theory (Wiley Series in + probability and mathematical statistics). + + Examples + -------- + >>> import numpy as np + >>> from scipy.special import multigammaln, gammaln + >>> a = 23.5 + >>> d = 10 + >>> multigammaln(a, d) + 454.1488605074416 + + Verify that the result agrees with the logarithm of the equation + shown above: + + >>> d*(d-1)/4*np.log(np.pi) + gammaln(a - 0.5*np.arange(0, d)).sum() + 454.1488605074416 + """ + a = np.asarray(a) + # Support for 0d arrays is needed for array_api_strict and dask. + d = np.asarray(d)[()] + if not np.isscalar(d) or (np.floor(d) != d): + raise ValueError("d should be a positive integer (dimension)") + if np.any(a <= 0.5 * (d - 1)): + raise ValueError(f"condition a ({a}) > 0.5 * (d-1) ({0.5 * (d-1)}) not met") + + res = (d * (d-1) * 0.25) * np.log(np.pi) + res += np.sum(loggam([(a - (j - 1.)/2) for j in range(1, d+1)]), axis=0) + return res + + +_poisson_binom_pmf_doc = ( + """Returns pmf of Poisson Binomial distribution. + + Parameters + ---------- + k : array + Number of successes at which to evaluate pmf. + + p : array + Success probabilities of independent Bernoulli trials. + + Notes + ----- + This is equivalent to a gufunc with signature ``()(i)->()``. + The last dimension of `p` contains success probabilities and + the preceding dimensions are batch dimensions. The batch + dimensions are broadcast against ``k``. + + The output will be C contiguous regardless of the contiguity of + `k` and `p`. + + """ +) + + +_poisson_binom_pmf = _with_cache_optimization( + name="_poisson_binom_pmf", + arg_names=["k", "p"], + docstring=_poisson_binom_pmf_doc, + ufunc=_gufuncs._poisson_binom_pmf, + cache_arg_indices=[1], +) + + +_poisson_binom_cdf_doc = ( + """Returns cdf of Poisson Binomial distribution. + + Parameters + ---------- + k : array + Number of successes at which to evaluate cdf. + + p : array + Success probabilities of independent Bernoulli trials. + + Notes + ----- + This is equivalent to a gufunc with signature ``()(i)->()``. + The last dimension of `p` contains success probabilities and + the preceding dimensions are batch dimensions. The batch + dimensions are broadcast against ``k``. + + The output will be C contiguous regardless of the contiguity of + `k` and `p`. + + """ +) + +_poisson_binom_cdf = _with_cache_optimization( + name="_poisson_binom_cdf", + arg_names=["k", "p"], + docstring=_poisson_binom_cdf_doc, + ufunc=_gufuncs._poisson_binom_cdf, + cache_arg_indices=[1], +) diff --git a/v12x/pylib/scipy/special/_spherical_bessel.py b/v12x/pylib/scipy/special/_spherical_bessel.py new file mode 100644 index 0000000000000000000000000000000000000000..56b345b78ad41f8e5c3c0a303b287611f7be9f4c --- /dev/null +++ b/v12x/pylib/scipy/special/_spherical_bessel.py @@ -0,0 +1,401 @@ +from functools import wraps +import scipy._external.array_api_extra as xpx +import numpy as np +from ._ufuncs import (_spherical_jn, _spherical_yn, _spherical_in, + _spherical_kn, _spherical_jn_d, _spherical_yn_d, + _spherical_in_d, _spherical_kn_d) + + +def use_reflection(sign_n_even=None, reflection_fun=None): + # - If reflection_fun is not specified, reflects negative `z` and multiplies + # output by appropriate sign (indicated by `sign_n_even`). + # - If reflection_fun is specified, calls `reflection_fun` instead of `fun`. + # See DLMF 10.47(v) https://dlmf.nist.gov/10.47 + def decorator(fun): + def standard_reflection(n, z, derivative): + # sign_n_even indicates the sign when the order `n` is even + sign = np.where(n % 2 == 0, sign_n_even, -sign_n_even) + # By the chain rule, differentiation at `-z` adds a minus sign + sign = -sign if derivative else sign + # Evaluate at positive z (minus negative z) and adjust the sign + return fun(n, -z, derivative) * sign + + @wraps(fun) + def wrapper(n, z, derivative=False): + z = np.asarray(z) + + if np.issubdtype(z.dtype, np.complexfloating): + return fun(n, z, derivative) # complex dtype just works + + f2 = standard_reflection if reflection_fun is None else reflection_fun + return xpx.apply_where(z.real >= 0, (n, z), + lambda n, z: fun(n, z, derivative), + lambda n, z: f2(n, z, derivative))[()] + return wrapper + return decorator + + +@use_reflection(+1) # See DLMF 10.47(v) https://dlmf.nist.gov/10.47 +def spherical_jn(n, z, derivative=False): + r"""Spherical Bessel function of the first kind or its derivative. + + Defined as [1]_, + + .. math:: j_n(z) = \sqrt{\frac{\pi}{2z}} J_{n + 1/2}(z), + + where :math:`J_n` is the Bessel function of the first kind. + + Parameters + ---------- + n : int, array_like + Order of the Bessel function (n >= 0). + z : complex or float, array_like + Argument of the Bessel function. + derivative : bool, optional + If True, the value of the derivative (rather than the function + itself) is returned. + + Returns + ------- + jn : ndarray + Value or derivative of spherical Bessel function of the first kind. + + Notes + ----- + For real arguments greater than the order, the function is computed + using the ascending recurrence [2]_. For small real or complex + arguments, the definitional relation to the cylindrical Bessel function + of the first kind is used. + + The derivative is computed using the relations [3]_, + + .. math:: + j_n'(z) = j_{n-1}(z) - \frac{n + 1}{z} j_n(z). + + j_0'(z) = -j_1(z) + + + .. versionadded:: 0.18.0 + + References + ---------- + .. [1] https://dlmf.nist.gov/10.47.E3 + .. [2] https://dlmf.nist.gov/10.51.E1 + .. [3] https://dlmf.nist.gov/10.51.E2 + .. [AS] Milton Abramowitz and Irene A. Stegun, eds. + Handbook of Mathematical Functions with Formulas, + Graphs, and Mathematical Tables. New York: Dover, 1972. + + Examples + -------- + The spherical Bessel functions of the first kind :math:`j_n` accept + both real and complex second argument. They can return a complex type: + + >>> from scipy.special import spherical_jn + >>> spherical_jn(0, 3+5j) + (-9.878987731663194-8.021894345786002j) + >>> type(spherical_jn(0, 3+5j)) + + + We can verify the relation for the derivative from the Notes + for :math:`n=3` in the interval :math:`[1, 2]`: + + >>> import numpy as np + >>> x = np.arange(1.0, 2.0, 0.01) + >>> np.allclose(spherical_jn(3, x, True), + ... spherical_jn(2, x) - 4/x * spherical_jn(3, x)) + True + + The first few :math:`j_n` with real argument: + + >>> import matplotlib.pyplot as plt + >>> x = np.arange(0.0, 10.0, 0.01) + >>> fig, ax = plt.subplots() + >>> ax.set_ylim(-0.5, 1.5) + >>> ax.set_title(r'Spherical Bessel functions $j_n$') + >>> for n in np.arange(0, 4): + ... ax.plot(x, spherical_jn(n, x), label=rf'$j_{n}$') + >>> plt.legend(loc='best') + >>> plt.show() + + """ + n = np.asarray(n, dtype=np.dtype("long")) + if derivative: + return _spherical_jn_d(n, z) + else: + return _spherical_jn(n, z) + + +@use_reflection(-1) # See DLMF 10.47(v) https://dlmf.nist.gov/10.47 +def spherical_yn(n, z, derivative=False): + r"""Spherical Bessel function of the second kind or its derivative. + + Defined as [1]_, + + .. math:: y_n(z) = \sqrt{\frac{\pi}{2z}} Y_{n + 1/2}(z), + + where :math:`Y_n` is the Bessel function of the second kind. + + Parameters + ---------- + n : int, array_like + Order of the Bessel function (n >= 0). + z : complex or float, array_like + Argument of the Bessel function. + derivative : bool, optional + If True, the value of the derivative (rather than the function + itself) is returned. + + Returns + ------- + yn : ndarray + Value or derivative of spherical Bessel function of the second kind. + + Notes + ----- + For real arguments, the function is computed using the ascending + recurrence [2]_. For complex arguments, the definitional relation to + the cylindrical Bessel function of the second kind is used. + + The derivative is computed using the relations [3]_, + + .. math:: + y_n' = y_{n-1} - \frac{n + 1}{z} y_n. + + y_0' = -y_1 + + + .. versionadded:: 0.18.0 + + References + ---------- + .. [1] https://dlmf.nist.gov/10.47.E4 + .. [2] https://dlmf.nist.gov/10.51.E1 + .. [3] https://dlmf.nist.gov/10.51.E2 + .. [AS] Milton Abramowitz and Irene A. Stegun, eds. + Handbook of Mathematical Functions with Formulas, + Graphs, and Mathematical Tables. New York: Dover, 1972. + + Examples + -------- + The spherical Bessel functions of the second kind :math:`y_n` accept + both real and complex second argument. They can return a complex type: + + >>> from scipy.special import spherical_yn + >>> spherical_yn(0, 3+5j) + (8.022343088587197-9.880052589376795j) + >>> type(spherical_yn(0, 3+5j)) + + + We can verify the relation for the derivative from the Notes + for :math:`n=3` in the interval :math:`[1, 2]`: + + >>> import numpy as np + >>> x = np.arange(1.0, 2.0, 0.01) + >>> np.allclose(spherical_yn(3, x, True), + ... spherical_yn(2, x) - 4/x * spherical_yn(3, x)) + True + + The first few :math:`y_n` with real argument: + + >>> import matplotlib.pyplot as plt + >>> x = np.arange(0.0, 10.0, 0.01) + >>> fig, ax = plt.subplots() + >>> ax.set_ylim(-2.0, 1.0) + >>> ax.set_title(r'Spherical Bessel functions $y_n$') + >>> for n in np.arange(0, 4): + ... ax.plot(x, spherical_yn(n, x), label=rf'$y_{n}$') + >>> plt.legend(loc='best') + >>> plt.show() + + """ + n = np.asarray(n, dtype=np.dtype("long")) + if derivative: + return _spherical_yn_d(n, z) + else: + return _spherical_yn(n, z) + + +@use_reflection(+1) # See DLMF 10.47(v) https://dlmf.nist.gov/10.47 +def spherical_in(n, z, derivative=False): + r"""Modified spherical Bessel function of the first kind or its derivative. + + Defined as [1]_, + + .. math:: i_n(z) = \sqrt{\frac{\pi}{2z}} I_{n + 1/2}(z), + + where :math:`I_n` is the modified Bessel function of the first kind. + + Parameters + ---------- + n : int, array_like + Order of the Bessel function (n >= 0). + z : complex or float, array_like + Argument of the Bessel function. + derivative : bool, optional + If True, the value of the derivative (rather than the function + itself) is returned. + + Returns + ------- + in : ndarray + Value or derivative of modified spherical Bessel function of the first kind. + + Notes + ----- + The function is computed using its definitional relation to the + modified cylindrical Bessel function of the first kind. + + The derivative is computed using the relations [2]_, + + .. math:: + i_n' = i_{n-1} - \frac{n + 1}{z} i_n. + + i_1' = i_0 + + + .. versionadded:: 0.18.0 + + References + ---------- + .. [1] https://dlmf.nist.gov/10.47.E7 + .. [2] https://dlmf.nist.gov/10.51.E5 + .. [AS] Milton Abramowitz and Irene A. Stegun, eds. + Handbook of Mathematical Functions with Formulas, + Graphs, and Mathematical Tables. New York: Dover, 1972. + + Examples + -------- + The modified spherical Bessel functions of the first kind :math:`i_n` + accept both real and complex second argument. + They can return a complex type: + + >>> from scipy.special import spherical_in + >>> spherical_in(0, 3+5j) + (-1.1689867793369182-1.2697305267234222j) + >>> type(spherical_in(0, 3+5j)) + + + We can verify the relation for the derivative from the Notes + for :math:`n=3` in the interval :math:`[1, 2]`: + + >>> import numpy as np + >>> x = np.arange(1.0, 2.0, 0.01) + >>> np.allclose(spherical_in(3, x, True), + ... spherical_in(2, x) - 4/x * spherical_in(3, x)) + True + + The first few :math:`i_n` with real argument: + + >>> import matplotlib.pyplot as plt + >>> x = np.arange(0.0, 6.0, 0.01) + >>> fig, ax = plt.subplots() + >>> ax.set_ylim(-0.5, 5.0) + >>> ax.set_title(r'Modified spherical Bessel functions $i_n$') + >>> for n in np.arange(0, 4): + ... ax.plot(x, spherical_in(n, x), label=rf'$i_{n}$') + >>> plt.legend(loc='best') + >>> plt.show() + + """ + n = np.asarray(n, dtype=np.dtype("long")) + if derivative: + return _spherical_in_d(n, z) + else: + return _spherical_in(n, z) + + +def spherical_kn_reflection(n, z, derivative=False): + # More complex than the other cases, and this will likely be re-implemented + # in C++ anyway. Would require multiple function evaluations. Probably about + # as fast to just resort to complex math, and much simpler. + return spherical_kn(n, z + 0j, derivative=derivative).real + + +@use_reflection(reflection_fun=spherical_kn_reflection) +def spherical_kn(n, z, derivative=False): + r"""Modified spherical Bessel function of the second kind or its derivative. + + Defined as [1]_, + + .. math:: k_n(z) = \sqrt{\frac{\pi}{2z}} K_{n + 1/2}(z), + + where :math:`K_n` is the modified Bessel function of the second kind. + + Parameters + ---------- + n : int, array_like + Order of the Bessel function (n >= 0). + z : complex or float, array_like + Argument of the Bessel function. + derivative : bool, optional + If True, the value of the derivative (rather than the function + itself) is returned. + + Returns + ------- + kn : ndarray + Value or derivative of modified spherical Bessel function of the second kind. + + Notes + ----- + The function is computed using its definitional relation to the + modified cylindrical Bessel function of the second kind. + + The derivative is computed using the relations [2]_, + + .. math:: + k_n' = -k_{n-1} - \frac{n + 1}{z} k_n. + + k_0' = -k_1 + + + .. versionadded:: 0.18.0 + + References + ---------- + .. [1] https://dlmf.nist.gov/10.47.E9 + .. [2] https://dlmf.nist.gov/10.51.E5 + .. [AS] Milton Abramowitz and Irene A. Stegun, eds. + Handbook of Mathematical Functions with Formulas, + Graphs, and Mathematical Tables. New York: Dover, 1972. + + Examples + -------- + The modified spherical Bessel functions of the second kind :math:`k_n` + accept both real and complex second argument. + They can return a complex type: + + >>> from scipy.special import spherical_kn + >>> spherical_kn(0, 3+5j) + (0.012985785614001561+0.003354691603137546j) + >>> type(spherical_kn(0, 3+5j)) + + + We can verify the relation for the derivative from the Notes + for :math:`n=3` in the interval :math:`[1, 2]`: + + >>> import numpy as np + >>> x = np.arange(1.0, 2.0, 0.01) + >>> np.allclose(spherical_kn(3, x, True), + ... - 4/x * spherical_kn(3, x) - spherical_kn(2, x)) + True + + The first few :math:`k_n` with real argument: + + >>> import matplotlib.pyplot as plt + >>> x = np.arange(0.0, 4.0, 0.01) + >>> fig, ax = plt.subplots() + >>> ax.set_ylim(0.0, 5.0) + >>> ax.set_title(r'Modified spherical Bessel functions $k_n$') + >>> for n in np.arange(0, 4): + ... ax.plot(x, spherical_kn(n, x), label=rf'$k_{n}$') + >>> plt.legend(loc='best') + >>> plt.show() + + """ + n = np.asarray(n, dtype=np.dtype("long")) + if derivative: + return _spherical_kn_d(n, z) + else: + return _spherical_kn(n, z) diff --git a/v12x/pylib/scipy/special/_support_alternative_backends.py b/v12x/pylib/scipy/special/_support_alternative_backends.py new file mode 100644 index 0000000000000000000000000000000000000000..262b974ee26d3bc1c7f69a9275e7527faff78ffc --- /dev/null +++ b/v12x/pylib/scipy/special/_support_alternative_backends.py @@ -0,0 +1,933 @@ +import functools +import operator +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from types import ModuleType + +import numpy as np +from scipy._lib._array_api import ( + array_namespace, scipy_namespace_for, is_numpy, is_dask, is_marray, is_jax_array, + is_jax, xp_promote, xp_capabilities, SCIPY_ARRAY_API, get_native_namespace_name, + is_array_api_obj +) +import scipy._external.array_api_extra as xpx +from . import _basic +from . import _spfun_stats +from . import _ufuncs + +# mypy: disable-error-code=dict-item + +def _special_namespace_for(xp): + spx = scipy_namespace_for(xp) + return getattr(spx, "special", None) + + +@dataclass +class _FuncInfo: + # NumPy-only function. IT MUST BE ELEMENTWISE. + func: Callable + # Number of arguments, not counting out= + # This is for testing purposes only, due to the fact that + # inspect.signature() just returns *args for ufuncs. + n_args: int + # @xp_capabilities decorator, for the purpose of + # documentation and unit testing. Omit to indicate + # full support for all backends. + xp_capabilities: Callable[[Callable], Callable] | None = None + # Generic implementation to fall back on if there is no native dispatch + # available. This is a function that accepts (main namespace, scipy namespace) + # and returns the final callable, or None if not available. + generic_impl: Callable[ + [ModuleType, ModuleType | None], Callable | None + ] | None = None + # Handle case where a backend uses an alternative name for a function. + # Should map backend names to alternative function names. + alt_names_map: dict[str, str] | None = None + # Some functions only take integer arrays for some arguments. + int_only: tuple[bool, ...] | None = None + # For testing purposes, whether tests should only use positive values + # for some arguments. If bool and equal to True, restrict to positive + # values for all arguments. To restrict only some arguments to positive + # values, pass a tuple of bool of the same length as the number of + # arguments, the ith entry in the tuple controls positive_only for + # the ith argument. To make backend specific choices for positive_only, + # pass in a dict mapping backend names to bool or tuple[bool, ...]. + positive_only: ( + bool | tuple[bool, ...] | Mapping[str, bool | tuple[bool, ...]] + ) = False + # Some special functions are not ufuncs and ufunc-specific tests + # should not be applied to these. + is_ufunc: bool = True + # Some non-ufunc special functions take only Python ints for some arguments. + # If so, python_int_only should be a tuple of the same length as the number + # of arguments,with value True if the corresponding argument needs to be a + # Python int. + # Can also take a dict mapping backends to such tuples if an argument being + # Python int only is backend specific. + python_int_only: dict[str, tuple[bool, ...]] | tuple[bool, ...] | None = None + # Some functions which seem to be scalar also accept 0d arrays. + scalar_or_0d_only: dict[str, tuple[bool, ...]] | tuple[bool, ...] | None = None + # Some functions may not work well with very large integer valued arguments. + test_large_ints: bool = True + # Some non-ufunc special functions don't decay 0d arrays to scalar. + produces_0d: bool = False + # Whether or not uses native PyTorch or falls back to NumPy/SciPy. This + # is needed because in PyTorch, the default dtype affects promotion + # rules when mixing integer and floating dtypes, so relying on a + # NumPy/SciPy fallback when the default dtype is other than float64 can lead + # to float64 output when native PyTorch would have e.g. float32 output. This + # must be accounted for in tests. Not putting this in xp_capabilities for now, + # but in the future I think it's likely we may want to add a warning to + # xp_capabilities when not using native PyTorch on CPU. + torch_native: bool = True + # Place a backend in this tuple if `func` is available as `xp.func` but not + # available in the `scipy.special` namespace for this backend. + # One example is `jax.numpy.sinc` being available but not `jax.scipy.special.sinc`. + backends_with_func_in_xp: tuple[str, ...] = () + + @property + def name(self): + return self.func.__name__ + + # These are needed by @lru_cache below + def __hash__(self): + return hash(self.func) + + def __eq__(self, other): + return isinstance(other, _FuncInfo) and self.func == other.func + + @property + def wrapper(self): + if self.name in globals(): + # Already initialised. We are likely in a unit test. + # Return function potentially overridden by xpx.testing.lazy_xp_function. + import scipy.special + return getattr(scipy.special, self.name) + + if SCIPY_ARRAY_API: + @functools.wraps(self.func) + def wrapped(*args, **kwargs): + xp = array_namespace(*args) + return self._wrapper_for(xp)(*args, **kwargs) + + # Allow pickling the function. Normally this is done by @wraps, + # but in this case it doesn't work because self.func is a ufunc. + wrapped.__module__ = "scipy.special" + wrapped.__qualname__ = self.name + func = wrapped + else: + func = self.func + + capabilities = self.xp_capabilities or xp_capabilities() + # In order to retain a naked ufunc when SCIPY_ARRAY_API is + # disabled, xp_capabilities must apply its changes in place. + cap_func = capabilities(func) + assert cap_func is func + return func + + @functools.lru_cache(1000) + def _wrapper_for(self, xp): + if is_numpy(xp): + return self.func + + # If a native implementation is available, use that + in_xp = get_native_namespace_name(xp) in self.backends_with_func_in_xp + namespace = xp if in_xp else _special_namespace_for(xp) + f = _get_native_func( + xp, namespace, self.name, alt_names_map=self.alt_names_map + ) + if f is not None: + return f + + if in_xp: + # when namespace is passed to self.generic_impl below, we want to + # make sure that it is the special namespace for xp and not xp + # itself, so raise if xp was incorrectly placed in + # `backends_with_func_in_xp`. + raise RuntimeError( + f"func {self.func} is not available as {xp.__name__}.{self.func}" + f" but {xp.__name__} was passed in ``in_xp``." + ) + + # If generic Array API implementation is available, use that + if self.generic_impl is not None: + f = self.generic_impl(xp, namespace) + if f is not None: + return f + + if is_marray(xp): + # Unwrap the array, apply the function on the wrapped namespace, + # and then re-wrap it. + # IMPORTANT: this only works because all functions in this module + # are elementwise. Otherwise, we would not be able to define a + # general rule for mask propagation. + + _f = globals()[self.name] # Allow nested wrapping + def f(*args, _f=_f, xp=xp, **kwargs): + data_args = [getattr(arg, 'data', arg) for arg in args] + out = _f(*data_args, **kwargs) + mask = functools.reduce(operator.or_, + (getattr(arg, 'mask', False) for arg in args)) + return xp.asarray(out, mask=mask) + + return f + + if is_dask(xp): + # Apply the function to each block of the Dask array. + # IMPORTANT: map_blocks works only because all functions in this module + # are elementwise. It would be a grave mistake to apply this to gufuncs + # or any other function with reductions, as they would change their + # output depending on chunking! + + _f = globals()[self.name] # Allow nested wrapping + def f(*args, _f=_f, xp=xp, **kwargs): + # Hide dtype kwarg from map_blocks + return xp.map_blocks(functools.partial(_f, **kwargs), *args) + + return f + + # As a final resort, use the NumPy/SciPy implementation + _f = self.func + + if is_jax(xp) and self.is_ufunc: + # if this is a ufunc, we can use the resolve_dtypes method to figure + # out what the output dtype should be and use lazy_apply to make this + # work with the JAX JIT. Since pure_callback used in lazy_apply will + # convert all of the inputs to eager JAX arrays, we will also need to + # get the input dtypes inferred from resolve_dtypes so that Python + # scalar inputs can be cast to the correct dtype under NEP50 weak + # promotion rules rather than getting promoted to the default currently + # set in JAX. One cannot just use xp_promote for the input dtypes because + # some ufuncs have integer only args. + def f(*args, _f=_f, xp=xp, **kwargs): + dtypes = (arg.dtype if is_jax_array(arg) else type(arg) for arg in args) + # result_dtypes needs an arg for the dtype of the optional out param. + # Use `None` since `out` is incompatible with JAX's immutability. + dtypes = (*dtypes, None) + # JAX uses NumPy dtypes so we can just pass these directly to + # resolve_dtypes. TODO: generalize to other lazy backends. + dtypes = _f.resolve_dtypes(dtypes) + out_dtype = dtypes[-1] + args = [ + xp.asarray(arg, dtype=dtype) + for arg, dtype in zip(args, dtypes[:-1]) + ] + return xpx.lazy_apply( + _f, *args, xp=xp, as_numpy=True, dtype=out_dtype, **kwargs + ) + else: + def f(*args, _f=_f, xp=xp, **kwargs): + # Check with `is_array_api_obj` to keep Python scalars untouched so that + # NEP50 can be followed. + args = [ + np.asarray(arg) if is_array_api_obj(arg) else arg for arg in args + ] + out = _f(*args, **kwargs) + return xp.asarray(out) + + return f + + +def _get_native_func(xp, namespace, f_name, *, alt_names_map=None): + if alt_names_map is None: + alt_names_map = {} + f_name = alt_names_map.get(get_native_namespace_name(xp), f_name) + f = getattr(namespace, f_name, None) + if f is None and hasattr(xp, 'special'): + # Currently dead branch, in anticipation of 'special' Array API extension + # https://github.com/data-apis/array-api/issues/725 + f = getattr(xp.special, f_name, None) + return f + + +def _rel_entr(xp, spsx): + def __rel_entr(x, y, *, xp=xp): + # https://github.com/data-apis/array-api-extra/issues/160 + mxp = array_namespace(x._meta, y._meta) if is_dask(xp) else xp + x, y = xp_promote(x, y, broadcast=True, force_floating=True, xp=xp) + xy_pos = (x > 0) & (y > 0) + xy_inf = xp.isinf(x) & xp.isinf(y) + res = xpx.apply_where( + xy_pos & ~xy_inf, + (x, y), + # Note: for very large x, this can overflow. + lambda x, y: x * (mxp.log(x) - mxp.log(y)), + fill_value=xp.inf + ) + res = xpx.at(res)[(x == 0) & (y >= 0)].set(0) + res = xpx.at(res)[xp.isnan(x) | xp.isnan(y) | (xy_pos & xy_inf)].set(xp.nan) + return res + + return __rel_entr + + +def _xlogy(xp, spsx): + def __xlogy(x, y, *, xp=xp): + x, y = xp_promote(x, y, force_floating=True, xp=xp) + with np.errstate(divide='ignore', invalid='ignore'): + temp = x * xp.log(y) + return xp.where(x == 0., 0., temp) + return __xlogy + + + +def _chdtr(xp, spsx): + # The difference between this and just using `gammainc` + # defined by `get_array_special_func` is that if `gammainc` + # isn't found, we don't want to use the SciPy version; we'll + # return None here and use the SciPy version of `chdtr`. + gammainc = _get_native_func(xp, spsx, 'gammainc') + if gammainc is None: + return None + + def __chdtr(v, x): + res = gammainc(v / 2, x / 2) # this is almost all we need + # The rest can be removed when google/jax#20507 is resolved + mask = (v == 0) & (x > 0) # JAX returns NaN + res = xp.where(mask, 1., res) + mask = xp.isinf(v) & xp.isinf(x) # JAX returns 1.0 + return xp.where(mask, xp.nan, res) + return __chdtr + + +def _chdtrc(xp, spsx): + # The difference between this and just using `gammaincc` + # defined by `get_array_special_func` is that if `gammaincc` + # isn't found, we don't want to use the SciPy version; we'll + # return None here and use the SciPy version of `chdtrc`. + gammaincc = _get_native_func(xp, spsx, 'gammaincc') + if gammaincc is None: + return None + + def __chdtrc(v, x): + res = xp.where(x >= 0, gammaincc(v/2, x/2), 1) + i_nan = ((x == 0) & (v == 0)) | xp.isnan(x) | xp.isnan(v) | (v < 0) + res = xp.where(i_nan, xp.nan, res) + return res + return __chdtrc + + +def _betaincc(xp, spsx): + betainc = _get_native_func(xp, spsx, 'betainc') + if betainc is None: + return None + + def __betaincc(a, b, x): + # not perfect; might want to just rely on SciPy + return betainc(b, a, 1-x) + return __betaincc + + +def _stdtr(xp, spsx): + betainc = _get_native_func(xp, spsx, 'betainc') + if betainc is None: + return None + + def __stdtr(df, t): + x = df / (t ** 2 + df) + tail = betainc(df / 2, 0.5, x) / 2 + return xp.where(t < 0, tail, 1 - tail) + + return __stdtr + + +def _stdtrit(xp, spsx): + # Need either native stdtr or native betainc + stdtr = _get_native_func(xp, spsx, 'stdtr') or _stdtr(xp, spsx) + # If betainc is not defined, the root-finding would be done with `xp` + # despite `stdtr` being evaluated with SciPy/NumPy `stdtr`. Save the + # conversions: in this case, just evaluate `stdtrit` with SciPy/NumPy. + if stdtr is None: + return None + + from scipy.optimize.elementwise import bracket_root, find_root + + def __stdtrit(df, p): + def fun(t, df, p): return stdtr(df, t) - p + res_bracket = bracket_root(fun, xp.zeros_like(p), args=(df, p)) + res_root = find_root(fun, res_bracket.bracket, args=(df, p)) + return res_root.x + + return __stdtrit + + +# Inventory of automatically dispatched functions +# IMPORTANT: these must all be **elementwise** functions! + +# PyTorch doesn't implement `betainc`. +# On torch CPU we can fall back to NumPy, but on GPU it won't work. +_needs_betainc = xp_capabilities(cpu_only=True, exceptions=["jax.numpy", "cupy"]) + +_special_funcs = ( + _FuncInfo( + _ufuncs.bdtr, 3, + xp_capabilities( + cpu_only=True, exceptions=["cupy"], + jax_jit=True, + ), + int_only=(False, True, False), torch_native=False, + ), + _FuncInfo( + _ufuncs.bdtrc, 3, + xp_capabilities( + cpu_only=True, exceptions=["cupy"], + jax_jit=True, + ), + int_only=(False, True, False), torch_native=False, + ), + _FuncInfo( + _ufuncs.bdtri, 3, + xp_capabilities( + cpu_only=True, exceptions=["cupy"], + jax_jit=True, + ), + int_only=(False, True, False), torch_native=False, + ), + _FuncInfo(_ufuncs.betainc, 3, _needs_betainc, torch_native=False), + _FuncInfo(_ufuncs.betaincc, 3, _needs_betainc, generic_impl=_betaincc, + torch_native=False), + _FuncInfo( + _ufuncs.betaincinv, 3, + xp_capabilities( + cpu_only=True, exceptions=["cupy"], + jax_jit=True, + ), + test_large_ints=False, positive_only=True, torch_native=False, + ), + _FuncInfo( + _ufuncs.betaln, 2, + xp_capabilities(cpu_only=True, exceptions=["cupy", "jax.numpy"]), + # For betaln, nan mismatches can occur at negative integer a or b of + # sufficiently large magnitude. + positive_only={"jax.numpy": True}, torch_native=False, + ), + _FuncInfo( + _ufuncs.binom, 2, + xp_capabilities( + cpu_only=True, exceptions=["cupy"], + jax_jit=True, + ), + torch_native=False, + ), + _FuncInfo( + _ufuncs.boxcox, 2, + xp_capabilities( + cpu_only=True, exceptions=["cupy"], + jax_jit=True, + ), + torch_native=False, + ), + _FuncInfo( + _ufuncs.boxcox1p, 2, + xp_capabilities( + cpu_only=True, exceptions=["cupy"], + jax_jit=True, + ), + torch_native=False, + ), + _FuncInfo( + _ufuncs.cbrt, 1, + xp_capabilities( + cpu_only=True, exceptions=["cupy"], + jax_jit=True, + ), + torch_native=False, + ), + _FuncInfo(_ufuncs.chdtr, 2, generic_impl=_chdtr), + _FuncInfo(_ufuncs.chdtrc, 2, generic_impl=_chdtrc, + # scipy/scipy#20972 + positive_only={"cupy": True, "jax.numpy": True, "torch": True}), + _FuncInfo( + _ufuncs.chdtri, 2, + xp_capabilities( + cpu_only=True, exceptions=["cupy"], + jax_jit=True, + ), + torch_native=False, + ), + _FuncInfo( + _ufuncs.cosdg, 1, + xp_capabilities( + cpu_only=True, exceptions=["cupy"], + jax_jit=True, + ), + test_large_ints=False, torch_native=False, + ), + _FuncInfo( + _ufuncs.cosm1, 1, + xp_capabilities( + cpu_only=True, exceptions=["cupy"], + jax_jit=True, + ), + torch_native=False, + ), + _FuncInfo( + _ufuncs.cotdg, 1, + xp_capabilities( + cpu_only=True, exceptions=["cupy"], + jax_jit=True, + ), + torch_native=False, + ), + _FuncInfo( + _ufuncs.ellipk, 1, + xp_capabilities( + cpu_only=True, exceptions=["cupy"], + jax_jit=True, + ), + torch_native=False, + ), + _FuncInfo( + _ufuncs.ellipkm1, 1, + xp_capabilities( + cpu_only=True, exceptions=["cupy"], + jax_jit=True, + ), + torch_native=False, + ), + _FuncInfo(_ufuncs.entr, 1), + _FuncInfo(_ufuncs.erf, 1), + _FuncInfo(_ufuncs.erfc, 1), + _FuncInfo( + _ufuncs.erfcx, 1, + xp_capabilities( + cpu_only=True, exceptions=["cupy", "torch"], + jax_jit=True, + ), + torch_native=False, + ), + _FuncInfo(_ufuncs.erfinv, 1), + _FuncInfo( + _ufuncs.exp1, 1, + xp_capabilities( + cpu_only=True, exceptions=["cupy"], + jax_jit=True, + ), + torch_native=False, + ), + _FuncInfo( + _ufuncs.exp10, 1, + xp_capabilities( + cpu_only=True, exceptions=["cupy"], + jax_jit=True, + ), + torch_native=False, + ), + _FuncInfo( + _ufuncs.exp2, 1, + xp_capabilities( + cpu_only=True, exceptions=["cupy"], + jax_jit=True, + ), + torch_native=False, + ), + _FuncInfo( + _ufuncs.exprel, 1, + xp_capabilities( + cpu_only=True, exceptions=["cupy"], + jax_jit=True, + ), + torch_native=False, + ), + _FuncInfo( + _ufuncs.expi, 1, + xp_capabilities(cpu_only=True, exceptions=["cupy", "jax.numpy"]), + torch_native=False, + ), + _FuncInfo(_ufuncs.expit, 1), + _FuncInfo( + _ufuncs.expn, 2, + xp_capabilities(cpu_only=True, exceptions=["cupy", "jax.numpy"]), + # Inconsistent behavior for negative n. expn is not defined here without + # taking analytic continuation. + positive_only=True, + int_only=(True, False), test_large_ints=False, + torch_native=False, + ), + _FuncInfo( + _ufuncs.fdtr, 3, + xp_capabilities( + cpu_only=True, exceptions=["cupy"], + jax_jit=True, + ), + torch_native=False, + ), + _FuncInfo( + _ufuncs.fdtrc, 3, + xp_capabilities( + cpu_only=True, exceptions=["cupy"], + jax_jit=True, + ), + torch_native=False, + ), + _FuncInfo( + _ufuncs.fdtri, 3, + xp_capabilities( + cpu_only=True, exceptions=["cupy"], + jax_jit=True, + ), + torch_native=False, + ), + _FuncInfo( + _ufuncs.gamma, 1, + xp_capabilities(cpu_only=True, exceptions=["cupy", "jax.numpy"]), + torch_native=False, + ), + _FuncInfo(_ufuncs.gammainc, 2), + _FuncInfo( + _ufuncs.gammaincc, 2, + # google/jax#20699 + positive_only={"jax.numpy": True}, + ), + _FuncInfo( + _ufuncs.gammainccinv, 2, + xp_capabilities( + cpu_only=True, exceptions=["cupy"], + jax_jit=True, + ), + torch_native=False, + ), + _FuncInfo( + _ufuncs.gammaincinv, 2, + xp_capabilities( + cpu_only=True, exceptions=["cupy"], + jax_jit=True, + ), + torch_native=False, + ), + _FuncInfo(_ufuncs.gammaln, 1), + _FuncInfo( + _ufuncs.gammasgn, 1, + xp_capabilities(cpu_only=True, exceptions=["cupy", "jax.numpy"]), + torch_native=False, + ), + _FuncInfo( + _ufuncs.gdtr, 3, + xp_capabilities( + cpu_only=True, exceptions=["cupy"], + jax_jit=True, + ), + torch_native=False, + ), + _FuncInfo( + _ufuncs.gdtrc, 3, + xp_capabilities( + cpu_only=True, exceptions=["cupy"], + jax_jit=True, + ), + torch_native=False, + ), + _FuncInfo( + _ufuncs.huber, 2, + xp_capabilities( + cpu_only=True, exceptions=["cupy"], + jax_jit=True, + ), + torch_native=False, + ), + _FuncInfo( + _ufuncs.hyp1f1, 3, + xp_capabilities(cpu_only=True, exceptions=["jax.numpy"]), + positive_only={"jax.numpy": True}, test_large_ints=False, + torch_native=False, + ), + # Comment out when jax>=0.6.1 is available in Conda for CI. + # (or add version requirements to xp_capabilities). + # _FuncInfo( + # _ufuncs.hyp2f1, 4, + # xp_capabilities(cpu_only=True, exceptions=["jax.numpy"]), + # positive_only={"jax.numpy": True}, test_large_ints=False, + # torch_native=False, + # ), + _FuncInfo( + _ufuncs.inv_boxcox, 2, + xp_capabilities( + cpu_only=True, exceptions=["cupy"], + jax_jit=True, + ), + torch_native=False, + ), + _FuncInfo( + _ufuncs.inv_boxcox1p, 2, + xp_capabilities( + cpu_only=True, exceptions=["cupy"], + jax_jit=True, + ), + torch_native=False, + ), + _FuncInfo(_ufuncs.i0, 1), + _FuncInfo(_ufuncs.i0e, 1), + _FuncInfo(_ufuncs.i1, 1), + _FuncInfo(_ufuncs.i1e, 1), + _FuncInfo( + _ufuncs.j0, 1, + xp_capabilities( + cpu_only=True, exceptions=["cupy", "torch"], + jax_jit=True, + ), + alt_names_map={"torch": "bessel_j0"}, test_large_ints=False, + ), + _FuncInfo( + _ufuncs.j1, 1, + xp_capabilities( + cpu_only=True, exceptions=["cupy", "torch"], + jax_jit=True, + ), + alt_names_map={"torch": "bessel_j1"}, test_large_ints=False, + ), + _FuncInfo( + _ufuncs.k0, 1, + xp_capabilities( + cpu_only=True, exceptions=["cupy", "torch"], + jax_jit=True, + ), + alt_names_map={"torch": "modified_bessel_k0"}, + ), + _FuncInfo( + _ufuncs.k0e, 1, + xp_capabilities( + cpu_only=True, exceptions=["cupy", "torch"], + jax_jit=True, + ), + alt_names_map={"torch": "scaled_modified_bessel_k0"}, + test_large_ints=False, + ), + _FuncInfo( + _ufuncs.k1, 1, + xp_capabilities( + cpu_only=True, exceptions=["cupy", "torch"], + jax_jit=True, + ), + alt_names_map={"torch": "modified_bessel_k1"}, + ), + _FuncInfo( + _ufuncs.k1e, 1, + xp_capabilities( + cpu_only=True, exceptions=["cupy", "torch"], + jax_jit=True, + ), + alt_names_map={"torch": "scaled_modified_bessel_k1"}, + test_large_ints=False), + _FuncInfo( + _ufuncs.kl_div, 2, + xp_capabilities(cpu_only=True, exceptions=["cupy", "jax.numpy"]), + torch_native=False, + ), + _FuncInfo(_ufuncs.log_ndtr, 1), + _FuncInfo( + _ufuncs.loggamma, 1, + xp_capabilities( + cpu_only=True, exceptions=["cupy"], + jax_jit=True, + ), + torch_native=False, + ), + _FuncInfo(_ufuncs.logit, 1), + _FuncInfo( + _ufuncs.lpmv, 3, + xp_capabilities( + cpu_only=True, exceptions=["cupy"], + jax_jit=True, + ), + torch_native=False, + test_large_ints=False, + ), + _FuncInfo( + _spfun_stats.multigammaln, 2, + is_ufunc=False, + python_int_only={ + "cupy": (False, True), + "jax.numpy": (False, True), + "torch": (False, True), + }, + scalar_or_0d_only={ + "array_api_strict": (False, True), + "numpy": (False, True), + "dask.array": (False, True), + "marray": (False, True), + }, + int_only=(False, True), test_large_ints=False, + positive_only=True, torch_native=False, + ), + _FuncInfo( + _ufuncs.nbdtr, 3, + xp_capabilities( + cpu_only=True, exceptions=["cupy"], + jax_jit=True, + ), + int_only=(True, True, False), positive_only=True, + torch_native=False, + ), + _FuncInfo( + _ufuncs.nbdtrc, 3, + xp_capabilities( + cpu_only=True, exceptions=["cupy"], + jax_jit=True, + ), + int_only=(True, True, False), positive_only=True, + torch_native=False, + ), + _FuncInfo( + _ufuncs.nbdtri, 3, + xp_capabilities( + cpu_only=True, exceptions=["cupy"], + jax_jit=True, + ), + int_only=(True, True, False), positive_only=True, + torch_native=False, + ), + _FuncInfo(_ufuncs.ndtr, 1), + _FuncInfo(_ufuncs.ndtri, 1), + _FuncInfo( + _ufuncs.pdtr, 2, + xp_capabilities( + cpu_only=True, exceptions=["cupy"], + jax_jit=True, + ), + positive_only=True, torch_native=False, + ), + _FuncInfo( + _ufuncs.pdtrc, 2, + xp_capabilities( + cpu_only=True, exceptions=["cupy"], + jax_jit=True, + ), + positive_only=True, torch_native=False, + ), + _FuncInfo( + _ufuncs.pdtri, 2, + xp_capabilities( + cpu_only=True, exceptions=["cupy"], + jax_jit=True, + ), + int_only=(True, False), positive_only=True, + torch_native=False, + ), + _FuncInfo( + _ufuncs.poch, 2, + xp_capabilities(cpu_only=True, exceptions=["cupy", "jax.numpy"]), + test_large_ints=False, torch_native=False, + ), + _FuncInfo( + _ufuncs.pseudo_huber, 2, + xp_capabilities( + cpu_only=True, exceptions=["cupy"], + jax_jit=True, + ), + torch_native=False, + ), + _FuncInfo( + _basic.polygamma, 2, int_only=(True, False), is_ufunc=False, + scalar_or_0d_only={"torch": (True, False)}, produces_0d=True, + positive_only={"torch": (True, False), "jax.numpy": True}, + test_large_ints=False, + ), + _FuncInfo(_ufuncs.psi, 1, alt_names_map={"jax.numpy": "digamma"}), + _FuncInfo( + _ufuncs.radian, 3, + xp_capabilities( + cpu_only=True, exceptions=["cupy"], + jax_jit=True, + ), + torch_native=False, + ), + _FuncInfo(_ufuncs.rel_entr, 2, generic_impl=_rel_entr), + _FuncInfo( + _ufuncs.rgamma, 1, + xp_capabilities( + cpu_only=True, exceptions=["cupy"], + jax_jit=True, + ), + torch_native=False, + ), + _FuncInfo( + _basic.sinc, 1, + xp_capabilities( + cpu_only=True, exceptions=["cupy", "torch"], + jax_jit=True, + ), + is_ufunc=False, + backends_with_func_in_xp=("jax.numpy",), + ), + _FuncInfo( + _ufuncs.sindg, 1, + xp_capabilities( + cpu_only=True, exceptions=["cupy"], + jax_jit=True, + ), + test_large_ints=False, torch_native=False, + ), + _FuncInfo( + _ufuncs.spence, 1, + xp_capabilities(cpu_only=True, exceptions=["jax.numpy"]), + torch_native=False, + ), + _FuncInfo(_ufuncs.stdtr, 2, _needs_betainc, generic_impl=_stdtr, + torch_native=False), + _FuncInfo( + _ufuncs.stdtrit, 2, + xp_capabilities( + cpu_only=True, exceptions=["cupy"], # needs betainc + skip_backends=[("jax.numpy", "no scipy.optimize support")], + ), + generic_impl=_stdtrit, torch_native=False, + ), + _FuncInfo( + _ufuncs.tandg, 1, + xp_capabilities( + cpu_only=True, exceptions=["cupy"], + jax_jit=True, + ), + test_large_ints=False, torch_native=False, + ), + _FuncInfo(_ufuncs.xlog1py, 2), + _FuncInfo(_ufuncs.xlogy, 2, generic_impl=_xlogy), + _FuncInfo( + _ufuncs.y0, 1, + xp_capabilities( + cpu_only=True, exceptions=["cupy", "torch"], + jax_jit=True, + ), + alt_names_map={"torch": "bessel_y0"}, test_large_ints=False, + ), + _FuncInfo( + _ufuncs.y1, 1, + xp_capabilities( + cpu_only=True, exceptions=["cupy", "torch"], + jax_jit=True, + ), + alt_names_map={"torch": "bessel_y1"}, test_large_ints=False, + ), + _FuncInfo( + _ufuncs.yn, 2, + xp_capabilities( + cpu_only=True, exceptions=["cupy"], + jax_jit=True, + ), + positive_only={"cupy": (True, False)}, int_only=(True, False), + test_large_ints=False, torch_native=False, + ), + _FuncInfo( + _basic.zeta, 2, is_ufunc=False, + positive_only={"jax.numpy": True, "torch": (True, False)}, + test_large_ints=False, + ), + _FuncInfo( + _ufuncs.zetac, 1, + xp_capabilities( + cpu_only=True, exceptions=["cupy"], + jax_jit=True, + ), + torch_native=False, + ), +) + +# Override ufuncs. +# When SCIPY_ARRAY_API is disabled, this exclusively updates the docstrings in place +# and populates the xp_capabilities table, while retaining the original ufuncs. +globals().update({nfo.func.__name__: nfo.wrapper for nfo in _special_funcs}) +# digamma is an alias for psi. Define here so it also has alternative backend +# support. Add noqa because the linter gets confused by the sneaky way psi +# is inserted into globals above. +digamma = psi # type:ignore[name-defined] # noqa: F821 +__all__ = [nfo.func.__name__ for nfo in _special_funcs] + ["digamma"] diff --git a/v12x/pylib/scipy/special/_testutils.py b/v12x/pylib/scipy/special/_testutils.py new file mode 100644 index 0000000000000000000000000000000000000000..3ff4130dc3eb504a772f9b0ea2a2c74df714d999 --- /dev/null +++ b/v12x/pylib/scipy/special/_testutils.py @@ -0,0 +1,321 @@ +import os +import functools +import operator +from scipy._external.packaging_version import version + +import numpy as np +from numpy.testing import assert_ +import pytest + +import scipy.special as sc + +__all__ = ['with_special_errors', 'assert_func_equal', 'FuncData'] + + +#------------------------------------------------------------------------------ +# Check if a module is present to be used in tests +#------------------------------------------------------------------------------ + +class MissingModule: + def __init__(self, name): + self.name = name + + +def check_version(module, min_ver): + if type(module) is MissingModule: + return pytest.mark.skip(reason=f"{module.name} is not installed") + return pytest.mark.skipif( + version.parse(module.__version__) < version.Version(min_ver), + reason=f"{module.__name__} version >= {min_ver} required" + ) + + +#------------------------------------------------------------------------------ +# Enable convergence and loss of precision warnings -- turn off one by one +#------------------------------------------------------------------------------ + +def with_special_errors(func): + """ + Enable special function errors (such as underflow, overflow, + loss of precision, etc.) + """ + @functools.wraps(func) + def wrapper(*a, **kw): + with sc.errstate(all='raise'): + res = func(*a, **kw) + return res + return wrapper + + +#------------------------------------------------------------------------------ +# Comparing function values at many data points at once, with helpful +# error reports +#------------------------------------------------------------------------------ + +def assert_func_equal(func, results, points, rtol=None, atol=None, + param_filter=None, knownfailure=None, + vectorized=True, dtype=None, nan_ok=False, + ignore_inf_sign=False, distinguish_nan_and_inf=True): + if hasattr(points, 'next'): + # it's a generator + points = list(points) + + points = np.asarray(points) + if points.ndim == 1: + points = points[:,None] + nparams = points.shape[1] + + if hasattr(results, '__name__'): + # function + data = points + result_columns = None + result_func = results + else: + # dataset + data = np.c_[points, results] + result_columns = list(range(nparams, data.shape[1])) + result_func = None + + fdata = FuncData(func, data, list(range(nparams)), + result_columns=result_columns, result_func=result_func, + rtol=rtol, atol=atol, param_filter=param_filter, + knownfailure=knownfailure, nan_ok=nan_ok, vectorized=vectorized, + ignore_inf_sign=ignore_inf_sign, + distinguish_nan_and_inf=distinguish_nan_and_inf) + fdata.check() + + +class FuncData: + """ + Data set for checking a special function. + + Parameters + ---------- + func : function + Function to test + data : numpy array + columnar data to use for testing + param_columns : int or tuple of ints + Columns indices in which the parameters to `func` lie. + Can be imaginary integers to indicate that the parameter + should be cast to complex. + result_columns : int or tuple of ints, optional + Column indices for expected results from `func`. + result_func : callable, optional + Function to call to obtain results. + rtol : float, optional + Required relative tolerance. Default is 5*eps. + atol : float, optional + Required absolute tolerance. Default is 5*tiny. + param_filter : function, or tuple of functions/Nones, optional + Filter functions to exclude some parameter ranges. + If omitted, no filtering is done. + knownfailure : str, optional + Known failure error message to raise when the test is run. + If omitted, no exception is raised. + nan_ok : bool, optional + If nan is always an accepted result. + vectorized : bool, optional + Whether all functions passed in are vectorized. + ignore_inf_sign : bool, optional + Whether to ignore signs of infinities. + (Doesn't matter for complex-valued functions.) + distinguish_nan_and_inf : bool, optional + If True, treat numbers which contain nans or infs as + equal. Sets ignore_inf_sign to be True. + + """ + + def __init__(self, func, data, param_columns, result_columns=None, + result_func=None, rtol=None, atol=None, param_filter=None, + knownfailure=None, dataname=None, nan_ok=False, vectorized=True, + ignore_inf_sign=False, distinguish_nan_and_inf=True): + self.func = func + self.data = data + self.dataname = dataname + if not hasattr(param_columns, '__len__'): + param_columns = (param_columns,) + self.param_columns = tuple(param_columns) + if result_columns is not None: + if not hasattr(result_columns, '__len__'): + result_columns = (result_columns,) + self.result_columns = tuple(result_columns) + if result_func is not None: + message = "Only result_func or result_columns should be provided" + raise ValueError(message) + elif result_func is not None: + self.result_columns = None + else: + raise ValueError("Either result_func or result_columns should be provided") + self.result_func = result_func + self.rtol = rtol + self.atol = atol + if not hasattr(param_filter, '__len__'): + param_filter = (param_filter,) + self.param_filter = param_filter + self.knownfailure = knownfailure + self.nan_ok = nan_ok + self.vectorized = vectorized + self.ignore_inf_sign = ignore_inf_sign + self.distinguish_nan_and_inf = distinguish_nan_and_inf + if not self.distinguish_nan_and_inf: + self.ignore_inf_sign = True + + def get_tolerances(self, dtype): + if not np.issubdtype(dtype, np.inexact): + dtype = np.dtype(float) + info = np.finfo(dtype) + rtol, atol = self.rtol, self.atol + if rtol is None: + rtol = 5*info.eps + if atol is None: + atol = 5*info.tiny + return rtol, atol + + def check(self, data=None, dtype=None, dtypes=None): + """Check the special function against the data.""" + __tracebackhide__ = operator.methodcaller( + 'errisinstance', AssertionError + ) + + if self.knownfailure: + pytest.xfail(reason=self.knownfailure) + + if data is None: + data = self.data + + if dtype is None: + dtype = data.dtype + else: + data = data.astype(dtype) + + rtol, atol = self.get_tolerances(dtype) + + # Apply given filter functions + if self.param_filter: + param_mask = np.ones((data.shape[0],), np.bool_) + for j, filter in zip(self.param_columns, self.param_filter): + if filter: + param_mask &= list(filter(data[:,j])) + data = data[param_mask] + + # Pick parameters from the correct columns + params = [] + for idx, j in enumerate(self.param_columns): + if np.iscomplexobj(j): + j = int(j.imag) + params.append(data[:,j].astype(complex)) + elif dtypes and idx < len(dtypes): + params.append(data[:, j].astype(dtypes[idx])) + else: + params.append(data[:,j]) + + # Helper for evaluating results + def eval_func_at_params(func, skip_mask=None): + if self.vectorized: + got = func(*params) + else: + got = [] + for j in range(len(params[0])): + if skip_mask is not None and skip_mask[j]: + got.append(np.nan) + continue + got.append(func(*tuple([params[i][j] for i in range(len(params))]))) + got = np.asarray(got) + if not isinstance(got, tuple): + got = (got,) + return got + + # Evaluate function to be tested + got = eval_func_at_params(self.func) + + # Grab the correct results + if self.result_columns is not None: + # Correct results passed in with the data + wanted = tuple([data[:,icol] for icol in self.result_columns]) + else: + # Function producing correct results passed in + skip_mask = None + if self.nan_ok and len(got) == 1: + # Don't spend time evaluating what doesn't need to be evaluated + skip_mask = np.isnan(got[0]) + wanted = eval_func_at_params(self.result_func, skip_mask=skip_mask) + + # Check the validity of each output returned + assert_(len(got) == len(wanted)) + + for output_num, (x, y) in enumerate(zip(got, wanted)): + if np.issubdtype(x.dtype, np.complexfloating) or self.ignore_inf_sign: + pinf_x = np.isinf(x) + pinf_y = np.isinf(y) + minf_x = np.isinf(x) + minf_y = np.isinf(y) + else: + pinf_x = np.isposinf(x) + pinf_y = np.isposinf(y) + minf_x = np.isneginf(x) + minf_y = np.isneginf(y) + nan_x = np.isnan(x) + nan_y = np.isnan(y) + + with np.errstate(all='ignore'): + abs_y = np.absolute(y) + abs_y[~np.isfinite(abs_y)] = 0 + diff = np.absolute(x - y) + diff[~np.isfinite(diff)] = 0 + + rdiff = diff / np.absolute(y) + rdiff[~np.isfinite(rdiff)] = 0 + + tol_mask = (diff <= atol + rtol*abs_y) + pinf_mask = (pinf_x == pinf_y) + minf_mask = (minf_x == minf_y) + + nan_mask = (nan_x == nan_y) + + bad_j = ~(tol_mask & pinf_mask & minf_mask & nan_mask) + + point_count = bad_j.size + if self.nan_ok: + bad_j &= ~nan_x + bad_j &= ~nan_y + point_count -= (nan_x | nan_y).sum() + + if not self.distinguish_nan_and_inf and not self.nan_ok: + # If nan's are okay we've already covered all these cases + inf_x = np.isinf(x) + inf_y = np.isinf(y) + both_nonfinite = (inf_x & nan_y) | (nan_x & inf_y) + bad_j &= ~both_nonfinite + point_count -= both_nonfinite.sum() + + if np.any(bad_j): + # Some bad results: inform what, where, and how bad + msg = [""] + msg.append(f"Max |adiff|: {diff[bad_j].max():g}") + msg.append(f"Max |rdiff|: {rdiff[bad_j].max():g}") + msg.append(f"Bad results ({np.sum(bad_j)} out of " + f"{point_count}) for the following points " + f"(in output {output_num}):") + for j in np.nonzero(bad_j)[0]: + j = int(j) + def fmt(x): + return f'{np.array2string(x[j], precision=18):30s}' + a = " ".join(map(fmt, params)) + b = " ".join(map(fmt, got)) + c = " ".join(map(fmt, wanted)) + d = fmt(rdiff) + msg.append(f"{a} => {b} != {c} (rdiff {d})") + assert_(False, "\n".join(msg)) + + def __repr__(self): + """Pretty-printing""" + if np.any(list(map(np.iscomplexobj, self.param_columns))): + is_complex = " (complex)" + else: + is_complex = "" + if self.dataname: + return (f"") + else: + return f"" diff --git a/v12x/pylib/scipy/special/_ufunc_tools.py b/v12x/pylib/scipy/special/_ufunc_tools.py new file mode 100644 index 0000000000000000000000000000000000000000..d9f97ff78afa4857e18383c7e70c3fbd31da884f --- /dev/null +++ b/v12x/pylib/scipy/special/_ufunc_tools.py @@ -0,0 +1,162 @@ +"""Helpers for producing efficient wrappers of ufuncs. +""" + +import re +import numpy as np + + +def _parse_core_ndims(signature): + """Return tuple of num core dims per input from gufunc signature.""" + input_sig = signature.split('->')[0] + groups = re.findall(r"\((.*?)\)", input_sig) + return tuple(0 if not g.strip() else g.count(',') + 1 for g in groups) + + +def _with_cache_optimization( + *, + name, + arg_names, + docstring, + ufunc, + cache_arg_indices, +): + """Helper to ensure optimal iteration order for ufuncs that use caching. + + This concerns internal caches which are only live over the course of one + call to a ufunc to avoid repeated computations during the course of the + loops. See the notes below for more information. + + Parameters + ---------- + name : str + arg_names : list[str] + The function name and arg names are passed in so that the wrapper + can be generated with the create name and argument names, improving + documentation and autocomplete. + docstring : str + ufunc : numpy.ufunc + cache_arg_indices : list[int] + Arguments to ufunc which are used in the kernel to compute an output + which is being cached for reuse when iterating over other arguments. + + Returns + ------- + callable + A wrapper for ufunc which transposes the axes of the inputs to ensure + iteration precedes in such a way to allow the cache within the ufunc + kernel to eliminate redundant computation. + + Notes + ----- + There is a common pattern in ufunc kernels exemplified by the situation + where some of the arguments are used to compute coeffients of an expansion + that is taken over one or more of the other arguments. A classic example is + Mathieu functions, which compute coefficients corresponding to the shape + parameter q and order m which in principle could be reused for varying + values of the parameter x. + + It had long been the case that the expensive computation of coefficients is + repeated unnecessarily across values of x. It is possible to add a cache to + the ufunc kernel which stores the expansion coefficients and only updates + if the pointers into the q and m arrays advance during the course of the + ufunc loops. Such a cache is instantiated each time a ufunc is called and + only lives during the course of the loops that are carried out for that + particular call. + + Whether such a cache actually helps depends on the order in which iteration + occurs. Ideally, one would want q and m to advance most slowly and for the + iterations over x for fixed q and m to be pushed to the inner most loops. + This helper replaces each input array (and a pre-allocated output array) + with a view where the axes which should vary most slowly are transposed to + the ends and forces computation in C order. This ensures iteration proceeds + in the optimal order. + + Note that because the pre-allocated output array used internally is C + contiguous, the output will be C contiguous regardless of contiguity of + the inputs. + + """ + + # Need to keep track of the number of core dimensions for each input + # since core dimensions don't participate in broadcasting. + core_ndims = ( + _parse_core_ndims(ufunc.signature) + if ufunc.signature is not None + else (0,)*ufunc.nin + ) + + def _wrapper(*args): + args = [np.asarray(arg) for arg in args] + + # Fast path for when the arguments which are used in the cached + # computation don't have batches. + if all(args[i].ndim == core_ndims[i] for i in cache_arg_indices): + return ufunc(*args) + + # To get batch_shapes, need to exclude core dimensions. Again, the core + # dimensions won't participate in broadcasting. + batch_shapes = [ + arg.shape[:-core_ndims[i]] if core_ndims[i] > 0 else arg.shape + for i, arg in enumerate(args) + ] + batch_shape = np.broadcast_shapes(*batch_shapes) + batch_ndim = len(batch_shape) + + # Broadcast each arg so that the batch shapes all agree, but the + # core dimensions may still differ. + args_b = [ + np.broadcast_to(arg, batch_shape + arg.shape[-core_ndims[i]:]) + if core_ndims[i] > 0 else np.broadcast_to(arg, batch_shape) + for i, arg in enumerate(args) + ] + + # After broadcasting, determine which axes have stride-length + # zero for each of the args participating in the cache. The cached + # value will remain the same for iterations across these axes. + varying_axes = [] + constant_axes = [] + + for ax in range(batch_ndim): + is_constant = all( + args_b[i].strides[ax] == 0 for i in cache_arg_indices + ) + if is_constant: + constant_axes.append(ax) + else: + varying_axes.append(ax) + sorted_batch_axes = varying_axes + constant_axes + + # Push the non-varying axes to the end so that they will vary most + # slowly when iterating in C order. + args_t = [] + for i, arg_b in enumerate(args_b): + axes = sorted_batch_axes + list( + range(batch_ndim, batch_ndim + core_ndims[i]) + ) + args_t.append(np.transpose(arg_b, axes=axes)) + + # Premake the output array. + input_dtypes = tuple(arg.dtype for arg in args_t) + (None,)*ufunc.nout + out_dtype = ufunc.resolve_dtypes(input_dtypes)[-1] + out_final = np.empty(batch_shape, dtype=out_dtype, order='C') + # a view of the output array with axes sorted as needed. + out_t = np.transpose(out_final, axes=sorted_batch_axes) + + # Set out to the above view, but return the C contiguous output. + # This avoids having non-contiguous output. + ufunc(*args_t, out=out_t, order='C') + return out_final + + # Do some metaprogramming with exec so that the arg names and func + # name are as expected. + arg_str = ", ".join(arg_names) + code = ( + f"""def {name}({arg_str}): + return _wrapper({arg_str}) + """ + ) + namespace = {"_wrapper": _wrapper} + exec(code, namespace) + wrapper = namespace[name] + wrapper.__doc__ = docstring + return wrapper diff --git a/v12x/pylib/scipy/special/add_newdocs.py b/v12x/pylib/scipy/special/add_newdocs.py new file mode 100644 index 0000000000000000000000000000000000000000..5549717d35710d71655e42c836625cde9346bcc3 --- /dev/null +++ b/v12x/pylib/scipy/special/add_newdocs.py @@ -0,0 +1,15 @@ +# This file is not meant for public use and will be removed in SciPy v2.0.0. + +from scipy._lib.deprecation import _sub_module_deprecation + +__all__: list[str] = [] + + +def __dir__(): + return __all__ + + +def __getattr__(name): + return _sub_module_deprecation(sub_package="special", module="add_newdocs", + private_modules=["_add_newdocs"], all=__all__, + attribute=name) diff --git a/v12x/pylib/scipy/special/basic.py b/v12x/pylib/scipy/special/basic.py new file mode 100644 index 0000000000000000000000000000000000000000..46c179a893c4c9022a7cdd77c8923ff5a012b0eb --- /dev/null +++ b/v12x/pylib/scipy/special/basic.py @@ -0,0 +1,84 @@ +# This file is not meant for public use and will be removed in SciPy v2.0.0. +# Use the `scipy.special` namespace for importing the functions +# included below. + +from scipy._lib.deprecation import _sub_module_deprecation + + +__all__ = [ # noqa: F822 + 'ai_zeros', + 'assoc_laguerre', + 'bei_zeros', + 'beip_zeros', + 'ber_zeros', + 'bernoulli', + 'berp_zeros', + 'bi_zeros', + 'comb', + 'digamma', + 'diric', + 'erf_zeros', + 'euler', + 'factorial', + 'factorial2', + 'factorialk', + 'fresnel_zeros', + 'fresnelc_zeros', + 'fresnels_zeros', + 'gamma', + 'h1vp', + 'h2vp', + 'hankel1', + 'hankel2', + 'iv', + 'ivp', + 'jn_zeros', + 'jnjnp_zeros', + 'jnp_zeros', + 'jnyn_zeros', + 'jv', + 'jvp', + 'kei_zeros', + 'keip_zeros', + 'kelvin_zeros', + 'ker_zeros', + 'kerp_zeros', + 'kv', + 'kvp', + 'lmbda', + 'lqmn', + 'lqn', + 'mathieu_a', + 'mathieu_b', + 'mathieu_even_coef', + 'mathieu_odd_coef', + 'obl_cv_seq', + 'pbdn_seq', + 'pbdv_seq', + 'pbvv_seq', + 'perm', + 'polygamma', + 'pro_cv_seq', + 'psi', + 'riccati_jn', + 'riccati_yn', + 'sinc', + 'y0_zeros', + 'y1_zeros', + 'y1p_zeros', + 'yn_zeros', + 'ynp_zeros', + 'yv', + 'yvp', + 'zeta' +] + + +def __dir__(): + return __all__ + + +def __getattr__(name): + return _sub_module_deprecation(sub_package="special", module="basic", + private_modules=["_basic", "_ufuncs"], all=__all__, + attribute=name) diff --git a/v12x/pylib/scipy/special/orthogonal.py b/v12x/pylib/scipy/special/orthogonal.py new file mode 100644 index 0000000000000000000000000000000000000000..0b13a08a96cb683d72a4a00d6962446e1779c88a --- /dev/null +++ b/v12x/pylib/scipy/special/orthogonal.py @@ -0,0 +1,45 @@ +# This file is not meant for public use and will be removed in SciPy v2.0.0. +# Use the `scipy.special` namespace for importing the functions +# included below. + +from scipy._lib.deprecation import _sub_module_deprecation + + +_polyfuns = ['legendre', 'chebyt', 'chebyu', 'chebyc', 'chebys', + 'jacobi', 'laguerre', 'genlaguerre', 'hermite', + 'hermitenorm', 'gegenbauer', 'sh_legendre', 'sh_chebyt', + 'sh_chebyu', 'sh_jacobi'] + +# Correspondence between new and old names of root functions +_rootfuns_map = {'roots_legendre': 'p_roots', + 'roots_chebyt': 't_roots', + 'roots_chebyu': 'u_roots', + 'roots_chebyc': 'c_roots', + 'roots_chebys': 's_roots', + 'roots_jacobi': 'j_roots', + 'roots_laguerre': 'l_roots', + 'roots_genlaguerre': 'la_roots', + 'roots_hermite': 'h_roots', + 'roots_hermitenorm': 'he_roots', + 'roots_gegenbauer': 'cg_roots', + 'roots_sh_legendre': 'ps_roots', + 'roots_sh_chebyt': 'ts_roots', + 'roots_sh_chebyu': 'us_roots', + 'roots_sh_jacobi': 'js_roots'} + + +__all__ = _polyfuns + list(_rootfuns_map.keys()) + [ # noqa: F822 + 'airy', 'p_roots', 't_roots', 'u_roots', 'c_roots', 's_roots', + 'j_roots', 'l_roots', 'la_roots', 'h_roots', 'he_roots', 'cg_roots', + 'ps_roots', 'ts_roots', 'us_roots', 'js_roots' +] + + +def __dir__(): + return __all__ + + +def __getattr__(name): + return _sub_module_deprecation(sub_package="special", module="orthogonal", + private_modules=["_orthogonal"], all=__all__, + attribute=name) diff --git a/v12x/pylib/scipy/special/sf_error.py b/v12x/pylib/scipy/special/sf_error.py new file mode 100644 index 0000000000000000000000000000000000000000..00ff73756acd4219a4ba94eb089bce7d4c32266d --- /dev/null +++ b/v12x/pylib/scipy/special/sf_error.py @@ -0,0 +1,20 @@ +# This file is not meant for public use and will be removed in SciPy v2.0.0. +# Use the `scipy.special` namespace for importing the functions +# included below. + +from scipy._lib.deprecation import _sub_module_deprecation + +__all__ = [ # noqa: F822 + 'SpecialFunctionWarning', + 'SpecialFunctionError' +] + + +def __dir__(): + return __all__ + + +def __getattr__(name): + return _sub_module_deprecation(sub_package="special", module="sf_error", + private_modules=["_sf_error"], all=__all__, + attribute=name) diff --git a/v12x/pylib/scipy/special/specfun.py b/v12x/pylib/scipy/special/specfun.py new file mode 100644 index 0000000000000000000000000000000000000000..f23284e66f19b55ca1dee380bef02c9a26662ef8 --- /dev/null +++ b/v12x/pylib/scipy/special/specfun.py @@ -0,0 +1,21 @@ +# This file is not meant for public use and will be removed in SciPy v2.0.0. +# Use the `scipy.special` namespace for importing the functions +# included below. + +from scipy._lib.deprecation import _sub_module_deprecation + +# ruff: noqa: F822 +__all__ = [ + 'lqmn', + 'pbdv' +] + + +def __dir__(): + return __all__ + + +def __getattr__(name): + return _sub_module_deprecation(sub_package="special", module="specfun", + private_modules=["_basic", "_specfun"], all=__all__, + attribute=name) diff --git a/v12x/pylib/scipy/special/spfun_stats.py b/v12x/pylib/scipy/special/spfun_stats.py new file mode 100644 index 0000000000000000000000000000000000000000..a1e58487aaa547483c9f2531ac4efc2ad5e4795c --- /dev/null +++ b/v12x/pylib/scipy/special/spfun_stats.py @@ -0,0 +1,17 @@ +# This file is not meant for public use and will be removed in SciPy v2.0.0. +# Use the `scipy.special` namespace for importing the functions +# included below. + +from scipy._lib.deprecation import _sub_module_deprecation + +__all__ = ['multigammaln'] # noqa: F822 + + +def __dir__(): + return __all__ + + +def __getattr__(name): + return _sub_module_deprecation(sub_package="special", module="spfun_stats", + private_modules=["_spfun_stats"], all=__all__, + attribute=name) diff --git a/v12x/pylib/scipy/special/tests/__init__.py b/v12x/pylib/scipy/special/tests/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/v12x/pylib/scipy/special/tests/cython_abi_signatures.json b/v12x/pylib/scipy/special/tests/cython_abi_signatures.json new file mode 100644 index 0000000000000000000000000000000000000000..6c83a72d0d2760b350833e9b5fd947e3f59aa8fc --- /dev/null +++ b/v12x/pylib/scipy/special/tests/cython_abi_signatures.json @@ -0,0 +1,400 @@ +{ + "scipy.special.cython_special": { + "__pyx_fuse_0_0eval_chebyc": "__pyx_t_double_complex (double, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_0_0eval_chebys": "__pyx_t_double_complex (double, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_0_0eval_chebyt": "__pyx_t_double_complex (double, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_0_0eval_chebyu": "__pyx_t_double_complex (double, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_0_0eval_gegenbauer": "__pyx_t_double_complex (double, double, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_0_0eval_genlaguerre": "__pyx_t_double_complex (double, double, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_0_0eval_jacobi": "__pyx_t_double_complex (double, double, double, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_0_0eval_laguerre": "__pyx_t_double_complex (double, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_0_0eval_legendre": "__pyx_t_double_complex (double, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_0_0eval_sh_chebyt": "__pyx_t_double_complex (double, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_0_0eval_sh_chebyu": "__pyx_t_double_complex (double, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_0_0eval_sh_jacobi": "__pyx_t_double_complex (double, double, double, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_0_0eval_sh_legendre": "__pyx_t_double_complex (double, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_0_1eval_chebyc": "double (double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_0_1eval_chebys": "double (double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_0_1eval_chebyt": "double (double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_0_1eval_chebyu": "double (double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_0_1eval_gegenbauer": "double (double, double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_0_1eval_genlaguerre": "double (double, double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_0_1eval_jacobi": "double (double, double, double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_0_1eval_laguerre": "double (double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_0_1eval_legendre": "double (double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_0_1eval_sh_chebyt": "double (double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_0_1eval_sh_chebyu": "double (double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_0_1eval_sh_jacobi": "double (double, double, double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_0_1eval_sh_legendre": "double (double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_0airy": "void (__pyx_t_double_complex, __pyx_t_double_complex *, __pyx_t_double_complex *, __pyx_t_double_complex *, __pyx_t_double_complex *)", + "__pyx_fuse_0airye": "void (__pyx_t_double_complex, __pyx_t_double_complex *, __pyx_t_double_complex *, __pyx_t_double_complex *, __pyx_t_double_complex *)", + "__pyx_fuse_0bdtr": "double (double, double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_0bdtrc": "double (double, double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_0bdtri": "double (double, double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_0betainc": "double (double, double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_0betaincc": "double (double, double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_0betainccinv": "double (double, double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_0betaincinv": "double (double, double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_0btdtria": "double (double, double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_0btdtrib": "double (double, double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_0chdtriv": "double (double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_0chndtr": "double (double, double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_0chndtridf": "double (double, double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_0chndtrinc": "double (double, double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_0chndtrix": "double (double, double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_0dawsn": "__pyx_t_double_complex (__pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_0elliprc": "__pyx_t_double_complex (__pyx_t_double_complex, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_0elliprd": "__pyx_t_double_complex (__pyx_t_double_complex, __pyx_t_double_complex, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_0elliprf": "__pyx_t_double_complex (__pyx_t_double_complex, __pyx_t_double_complex, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_0elliprg": "__pyx_t_double_complex (__pyx_t_double_complex, __pyx_t_double_complex, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_0elliprj": "__pyx_t_double_complex (__pyx_t_double_complex, __pyx_t_double_complex, __pyx_t_double_complex, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_0erf": "__pyx_t_double_complex (__pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_0erfc": "__pyx_t_double_complex (__pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_0erfcx": "__pyx_t_double_complex (__pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_0erfi": "__pyx_t_double_complex (__pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_0erfinv": "double (double, int __pyx_skip_dispatch)", + "__pyx_fuse_0exp1": "__pyx_t_double_complex (__pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_0expi": "__pyx_t_double_complex (__pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_0expit": "double (double, int __pyx_skip_dispatch)", + "__pyx_fuse_0expm1": "__pyx_t_double_complex (__pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_0expn": "double (double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_0fdtr": "double (double, double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_0fdtrc": "double (double, double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_0fdtri": "double (double, double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_0fresnel": "void (__pyx_t_double_complex, __pyx_t_double_complex *, __pyx_t_double_complex *)", + "__pyx_fuse_0gamma": "__pyx_t_double_complex (__pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_0hyp0f1": "__pyx_t_double_complex (double, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_0hyp1f1": "__pyx_t_double_complex (double, double, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_0hyp2f1": "__pyx_t_double_complex (double, double, double, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_0iv": "__pyx_t_double_complex (double, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_0ive": "__pyx_t_double_complex (double, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_0jv": "__pyx_t_double_complex (double, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_0jve": "__pyx_t_double_complex (double, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_0kn": "double (double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_0kv": "__pyx_t_double_complex (double, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_0kve": "__pyx_t_double_complex (double, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_0log1p": "__pyx_t_double_complex (__pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_0log_expit": "double (double, int __pyx_skip_dispatch)", + "__pyx_fuse_0log_ndtr": "__pyx_t_double_complex (__pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_0loggamma": "__pyx_t_double_complex (__pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_0logit": "double (double, int __pyx_skip_dispatch)", + "__pyx_fuse_0nbdtr": "double (double, double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_0nbdtrc": "double (double, double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_0nbdtri": "double (double, double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_0ncfdtr": "double (double, double, double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_0ncfdtri": "double (double, double, double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_0nctdtr": "double (double, double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_0nctdtrit": "double (double, double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_0ndtr": "__pyx_t_double_complex (__pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_0pdtri": "double (double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_0pdtrik": "double (double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_0powm1": "double (double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_0psi": "__pyx_t_double_complex (__pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_0rgamma": "__pyx_t_double_complex (__pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_0shichi": "void (__pyx_t_double_complex, __pyx_t_double_complex *, __pyx_t_double_complex *)", + "__pyx_fuse_0sici": "void (__pyx_t_double_complex, __pyx_t_double_complex *, __pyx_t_double_complex *)", + "__pyx_fuse_0smirnov": "double (double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_0smirnovi": "double (double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_0spence": "__pyx_t_double_complex (__pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_0spherical_in": "__pyx_t_double_complex (Py_ssize_t, __pyx_t_double_complex, int __pyx_skip_dispatch, struct __pyx_fuse_0__pyx_opt_args_5scipy_7special_14cython_special_spherical_in *__pyx_optional_args)", + "__pyx_fuse_0spherical_jn": "__pyx_t_double_complex (Py_ssize_t, __pyx_t_double_complex, int __pyx_skip_dispatch, struct __pyx_fuse_0__pyx_opt_args_5scipy_7special_14cython_special_spherical_jn *__pyx_optional_args)", + "__pyx_fuse_0spherical_kn": "__pyx_t_double_complex (Py_ssize_t, __pyx_t_double_complex, int __pyx_skip_dispatch, struct __pyx_fuse_0__pyx_opt_args_5scipy_7special_14cython_special_spherical_kn *__pyx_optional_args)", + "__pyx_fuse_0spherical_yn": "__pyx_t_double_complex (Py_ssize_t, __pyx_t_double_complex, int __pyx_skip_dispatch, struct __pyx_fuse_0__pyx_opt_args_5scipy_7special_14cython_special_spherical_yn *__pyx_optional_args)", + "__pyx_fuse_0stdtr": "double (double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_0stdtrit": "double (double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_0wrightomega": "__pyx_t_double_complex (__pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_0xlog1py": "__pyx_t_double_complex (__pyx_t_double_complex, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_0xlogy": "__pyx_t_double_complex (__pyx_t_double_complex, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_0yn": "double (double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_0yv": "__pyx_t_double_complex (double, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_0yve": "__pyx_t_double_complex (double, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_1_0eval_chebyc": "__pyx_t_double_complex (long, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_1_0eval_chebys": "__pyx_t_double_complex (long, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_1_0eval_chebyt": "__pyx_t_double_complex (long, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_1_0eval_chebyu": "__pyx_t_double_complex (long, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_1_0eval_gegenbauer": "__pyx_t_double_complex (long, double, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_1_0eval_genlaguerre": "__pyx_t_double_complex (long, double, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_1_0eval_jacobi": "__pyx_t_double_complex (long, double, double, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_1_0eval_laguerre": "__pyx_t_double_complex (long, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_1_0eval_legendre": "__pyx_t_double_complex (long, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_1_0eval_sh_chebyt": "__pyx_t_double_complex (long, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_1_0eval_sh_chebyu": "__pyx_t_double_complex (long, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_1_0eval_sh_jacobi": "__pyx_t_double_complex (long, double, double, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_1_0eval_sh_legendre": "__pyx_t_double_complex (long, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_1_1eval_chebyc": "double (long, double, int __pyx_skip_dispatch)", + "__pyx_fuse_1_1eval_chebys": "double (long, double, int __pyx_skip_dispatch)", + "__pyx_fuse_1_1eval_chebyt": "double (long, double, int __pyx_skip_dispatch)", + "__pyx_fuse_1_1eval_chebyu": "double (long, double, int __pyx_skip_dispatch)", + "__pyx_fuse_1_1eval_gegenbauer": "double (long, double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_1_1eval_genlaguerre": "double (long, double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_1_1eval_jacobi": "double (long, double, double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_1_1eval_laguerre": "double (long, double, int __pyx_skip_dispatch)", + "__pyx_fuse_1_1eval_legendre": "double (long, double, int __pyx_skip_dispatch)", + "__pyx_fuse_1_1eval_sh_chebyt": "double (long, double, int __pyx_skip_dispatch)", + "__pyx_fuse_1_1eval_sh_chebyu": "double (long, double, int __pyx_skip_dispatch)", + "__pyx_fuse_1_1eval_sh_jacobi": "double (long, double, double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_1_1eval_sh_legendre": "double (long, double, int __pyx_skip_dispatch)", + "__pyx_fuse_1airy": "void (double, double *, double *, double *, double *)", + "__pyx_fuse_1airye": "void (double, double *, double *, double *, double *)", + "__pyx_fuse_1bdtr": "double (double, long, double, int __pyx_skip_dispatch)", + "__pyx_fuse_1bdtrc": "double (double, long, double, int __pyx_skip_dispatch)", + "__pyx_fuse_1bdtri": "double (double, long, double, int __pyx_skip_dispatch)", + "__pyx_fuse_1betainc": "float (float, float, float, int __pyx_skip_dispatch)", + "__pyx_fuse_1betaincc": "float (float, float, float, int __pyx_skip_dispatch)", + "__pyx_fuse_1betainccinv": "float (float, float, float, int __pyx_skip_dispatch)", + "__pyx_fuse_1betaincinv": "float (float, float, float, int __pyx_skip_dispatch)", + "__pyx_fuse_1btdtria": "float (float, float, float, int __pyx_skip_dispatch)", + "__pyx_fuse_1btdtrib": "float (float, float, float, int __pyx_skip_dispatch)", + "__pyx_fuse_1chdtriv": "float (float, float, int __pyx_skip_dispatch)", + "__pyx_fuse_1chndtr": "float (float, float, float, int __pyx_skip_dispatch)", + "__pyx_fuse_1chndtridf": "float (float, float, float, int __pyx_skip_dispatch)", + "__pyx_fuse_1chndtrinc": "float (float, float, float, int __pyx_skip_dispatch)", + "__pyx_fuse_1chndtrix": "float (float, float, float, int __pyx_skip_dispatch)", + "__pyx_fuse_1dawsn": "double (double, int __pyx_skip_dispatch)", + "__pyx_fuse_1elliprc": "double (double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_1elliprd": "double (double, double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_1elliprf": "double (double, double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_1elliprg": "double (double, double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_1elliprj": "double (double, double, double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_1erf": "double (double, int __pyx_skip_dispatch)", + "__pyx_fuse_1erfc": "double (double, int __pyx_skip_dispatch)", + "__pyx_fuse_1erfcx": "double (double, int __pyx_skip_dispatch)", + "__pyx_fuse_1erfi": "double (double, int __pyx_skip_dispatch)", + "__pyx_fuse_1erfinv": "float (float, int __pyx_skip_dispatch)", + "__pyx_fuse_1exp1": "double (double, int __pyx_skip_dispatch)", + "__pyx_fuse_1expi": "double (double, int __pyx_skip_dispatch)", + "__pyx_fuse_1expit": "float (float, int __pyx_skip_dispatch)", + "__pyx_fuse_1expm1": "double (double, int __pyx_skip_dispatch)", + "__pyx_fuse_1expn": "double (long, double, int __pyx_skip_dispatch)", + "__pyx_fuse_1fdtr": "float (float, float, float, int __pyx_skip_dispatch)", + "__pyx_fuse_1fdtrc": "float (float, float, float, int __pyx_skip_dispatch)", + "__pyx_fuse_1fdtri": "float (float, float, float, int __pyx_skip_dispatch)", + "__pyx_fuse_1fresnel": "void (double, double *, double *)", + "__pyx_fuse_1gamma": "double (double, int __pyx_skip_dispatch)", + "__pyx_fuse_1hyp0f1": "double (double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_1hyp1f1": "double (double, double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_1hyp2f1": "double (double, double, double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_1iv": "double (double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_1ive": "double (double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_1jv": "double (double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_1jve": "double (double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_1kn": "double (long, double, int __pyx_skip_dispatch)", + "__pyx_fuse_1kv": "double (double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_1kve": "double (double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_1log1p": "double (double, int __pyx_skip_dispatch)", + "__pyx_fuse_1log_expit": "float (float, int __pyx_skip_dispatch)", + "__pyx_fuse_1log_ndtr": "double (double, int __pyx_skip_dispatch)", + "__pyx_fuse_1loggamma": "double (double, int __pyx_skip_dispatch)", + "__pyx_fuse_1logit": "float (float, int __pyx_skip_dispatch)", + "__pyx_fuse_1nbdtr": "double (long, long, double, int __pyx_skip_dispatch)", + "__pyx_fuse_1nbdtrc": "double (long, long, double, int __pyx_skip_dispatch)", + "__pyx_fuse_1nbdtri": "double (long, long, double, int __pyx_skip_dispatch)", + "__pyx_fuse_1ncfdtr": "float (float, float, float, float, int __pyx_skip_dispatch)", + "__pyx_fuse_1ncfdtri": "float (float, float, float, float, int __pyx_skip_dispatch)", + "__pyx_fuse_1nctdtr": "float (float, float, float, int __pyx_skip_dispatch)", + "__pyx_fuse_1nctdtrit": "float (float, float, float, int __pyx_skip_dispatch)", + "__pyx_fuse_1ndtr": "double (double, int __pyx_skip_dispatch)", + "__pyx_fuse_1pdtri": "double (long, double, int __pyx_skip_dispatch)", + "__pyx_fuse_1pdtrik": "float (float, float, int __pyx_skip_dispatch)", + "__pyx_fuse_1powm1": "float (float, float, int __pyx_skip_dispatch)", + "__pyx_fuse_1psi": "double (double, int __pyx_skip_dispatch)", + "__pyx_fuse_1rgamma": "double (double, int __pyx_skip_dispatch)", + "__pyx_fuse_1shichi": "void (double, double *, double *)", + "__pyx_fuse_1sici": "void (double, double *, double *)", + "__pyx_fuse_1smirnov": "double (long, double, int __pyx_skip_dispatch)", + "__pyx_fuse_1smirnovi": "double (long, double, int __pyx_skip_dispatch)", + "__pyx_fuse_1spence": "double (double, int __pyx_skip_dispatch)", + "__pyx_fuse_1spherical_in": "double (Py_ssize_t, double, int __pyx_skip_dispatch, struct __pyx_fuse_1__pyx_opt_args_5scipy_7special_14cython_special_spherical_in *__pyx_optional_args)", + "__pyx_fuse_1spherical_jn": "double (Py_ssize_t, double, int __pyx_skip_dispatch, struct __pyx_fuse_1__pyx_opt_args_5scipy_7special_14cython_special_spherical_jn *__pyx_optional_args)", + "__pyx_fuse_1spherical_kn": "double (Py_ssize_t, double, int __pyx_skip_dispatch, struct __pyx_fuse_1__pyx_opt_args_5scipy_7special_14cython_special_spherical_kn *__pyx_optional_args)", + "__pyx_fuse_1spherical_yn": "double (Py_ssize_t, double, int __pyx_skip_dispatch, struct __pyx_fuse_1__pyx_opt_args_5scipy_7special_14cython_special_spherical_yn *__pyx_optional_args)", + "__pyx_fuse_1stdtr": "float (float, float, int __pyx_skip_dispatch)", + "__pyx_fuse_1stdtrit": "float (float, float, int __pyx_skip_dispatch)", + "__pyx_fuse_1wrightomega": "double (double, int __pyx_skip_dispatch)", + "__pyx_fuse_1xlog1py": "double (double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_1xlogy": "double (double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_1yn": "double (long, double, int __pyx_skip_dispatch)", + "__pyx_fuse_1yv": "double (double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_1yve": "double (double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_2_0eval_chebyc": "__pyx_t_double_complex (Py_ssize_t, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_2_0eval_chebys": "__pyx_t_double_complex (Py_ssize_t, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_2_0eval_chebyt": "__pyx_t_double_complex (Py_ssize_t, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_2_0eval_chebyu": "__pyx_t_double_complex (Py_ssize_t, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_2_0eval_gegenbauer": "__pyx_t_double_complex (Py_ssize_t, double, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_2_0eval_genlaguerre": "__pyx_t_double_complex (Py_ssize_t, double, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_2_0eval_jacobi": "__pyx_t_double_complex (Py_ssize_t, double, double, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_2_0eval_laguerre": "__pyx_t_double_complex (Py_ssize_t, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_2_0eval_legendre": "__pyx_t_double_complex (Py_ssize_t, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_2_0eval_sh_chebyt": "__pyx_t_double_complex (Py_ssize_t, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_2_0eval_sh_chebyu": "__pyx_t_double_complex (Py_ssize_t, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_2_0eval_sh_jacobi": "__pyx_t_double_complex (Py_ssize_t, double, double, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_2_0eval_sh_legendre": "__pyx_t_double_complex (Py_ssize_t, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "__pyx_fuse_2_1eval_chebyc": "double (Py_ssize_t, double, int __pyx_skip_dispatch)", + "__pyx_fuse_2_1eval_chebys": "double (Py_ssize_t, double, int __pyx_skip_dispatch)", + "__pyx_fuse_2_1eval_chebyt": "double (Py_ssize_t, double, int __pyx_skip_dispatch)", + "__pyx_fuse_2_1eval_chebyu": "double (Py_ssize_t, double, int __pyx_skip_dispatch)", + "__pyx_fuse_2_1eval_gegenbauer": "double (Py_ssize_t, double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_2_1eval_genlaguerre": "double (Py_ssize_t, double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_2_1eval_jacobi": "double (Py_ssize_t, double, double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_2_1eval_laguerre": "double (Py_ssize_t, double, int __pyx_skip_dispatch)", + "__pyx_fuse_2_1eval_legendre": "double (Py_ssize_t, double, int __pyx_skip_dispatch)", + "__pyx_fuse_2_1eval_sh_chebyt": "double (Py_ssize_t, double, int __pyx_skip_dispatch)", + "__pyx_fuse_2_1eval_sh_chebyu": "double (Py_ssize_t, double, int __pyx_skip_dispatch)", + "__pyx_fuse_2_1eval_sh_jacobi": "double (Py_ssize_t, double, double, double, int __pyx_skip_dispatch)", + "__pyx_fuse_2_1eval_sh_legendre": "double (Py_ssize_t, double, int __pyx_skip_dispatch)", + "__pyx_fuse_2bdtr": "double (double, Py_ssize_t, double, int __pyx_skip_dispatch)", + "__pyx_fuse_2bdtrc": "double (double, Py_ssize_t, double, int __pyx_skip_dispatch)", + "__pyx_fuse_2bdtri": "double (double, Py_ssize_t, double, int __pyx_skip_dispatch)", + "__pyx_fuse_2expit": "long double (long double, int __pyx_skip_dispatch)", + "__pyx_fuse_2expn": "double (Py_ssize_t, double, int __pyx_skip_dispatch)", + "__pyx_fuse_2kn": "double (Py_ssize_t, double, int __pyx_skip_dispatch)", + "__pyx_fuse_2log_expit": "long double (long double, int __pyx_skip_dispatch)", + "__pyx_fuse_2logit": "long double (long double, int __pyx_skip_dispatch)", + "__pyx_fuse_2nbdtr": "double (Py_ssize_t, Py_ssize_t, double, int __pyx_skip_dispatch)", + "__pyx_fuse_2nbdtrc": "double (Py_ssize_t, Py_ssize_t, double, int __pyx_skip_dispatch)", + "__pyx_fuse_2nbdtri": "double (Py_ssize_t, Py_ssize_t, double, int __pyx_skip_dispatch)", + "__pyx_fuse_2pdtri": "double (Py_ssize_t, double, int __pyx_skip_dispatch)", + "__pyx_fuse_2smirnov": "double (Py_ssize_t, double, int __pyx_skip_dispatch)", + "__pyx_fuse_2smirnovi": "double (Py_ssize_t, double, int __pyx_skip_dispatch)", + "__pyx_fuse_2yn": "double (Py_ssize_t, double, int __pyx_skip_dispatch)", + "agm": "double (double, double, int __pyx_skip_dispatch)", + "bdtrik": "double (double, double, double, int __pyx_skip_dispatch)", + "bdtrin": "double (double, double, double, int __pyx_skip_dispatch)", + "bei": "double (double, int __pyx_skip_dispatch)", + "beip": "double (double, int __pyx_skip_dispatch)", + "ber": "double (double, int __pyx_skip_dispatch)", + "berp": "double (double, int __pyx_skip_dispatch)", + "besselpoly": "double (double, double, double, int __pyx_skip_dispatch)", + "beta": "double (double, double, int __pyx_skip_dispatch)", + "betaln": "double (double, double, int __pyx_skip_dispatch)", + "binom": "double (double, double, int __pyx_skip_dispatch)", + "boxcox": "double (double, double, int __pyx_skip_dispatch)", + "boxcox1p": "double (double, double, int __pyx_skip_dispatch)", + "cbrt": "double (double, int __pyx_skip_dispatch)", + "chdtr": "double (double, double, int __pyx_skip_dispatch)", + "chdtrc": "double (double, double, int __pyx_skip_dispatch)", + "chdtri": "double (double, double, int __pyx_skip_dispatch)", + "cosdg": "double (double, int __pyx_skip_dispatch)", + "cosm1": "double (double, int __pyx_skip_dispatch)", + "cotdg": "double (double, int __pyx_skip_dispatch)", + "ellipe": "double (double, int __pyx_skip_dispatch)", + "ellipeinc": "double (double, double, int __pyx_skip_dispatch)", + "ellipj": "void (double, double, double *, double *, double *, double *)", + "ellipk": "double (double, int __pyx_skip_dispatch)", + "ellipkinc": "double (double, double, int __pyx_skip_dispatch)", + "ellipkm1": "double (double, int __pyx_skip_dispatch)", + "entr": "double (double, int __pyx_skip_dispatch)", + "erfcinv": "double (double, int __pyx_skip_dispatch)", + "eval_hermite": "double (Py_ssize_t, double, int __pyx_skip_dispatch)", + "eval_hermitenorm": "double (Py_ssize_t, double, int __pyx_skip_dispatch)", + "exp10": "double (double, int __pyx_skip_dispatch)", + "exp2": "double (double, int __pyx_skip_dispatch)", + "exprel": "double (double, int __pyx_skip_dispatch)", + "fdtridfd": "double (double, double, double, int __pyx_skip_dispatch)", + "gammainc": "double (double, double, int __pyx_skip_dispatch)", + "gammaincc": "double (double, double, int __pyx_skip_dispatch)", + "gammainccinv": "double (double, double, int __pyx_skip_dispatch)", + "gammaincinv": "double (double, double, int __pyx_skip_dispatch)", + "gammaln": "double (double, int __pyx_skip_dispatch)", + "gammasgn": "double (double, int __pyx_skip_dispatch)", + "gdtr": "double (double, double, double, int __pyx_skip_dispatch)", + "gdtrc": "double (double, double, double, int __pyx_skip_dispatch)", + "gdtria": "double (double, double, double, int __pyx_skip_dispatch)", + "gdtrib": "double (double, double, double, int __pyx_skip_dispatch)", + "gdtrix": "double (double, double, double, int __pyx_skip_dispatch)", + "hankel1": "__pyx_t_double_complex (double, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "hankel1e": "__pyx_t_double_complex (double, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "hankel2": "__pyx_t_double_complex (double, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "hankel2e": "__pyx_t_double_complex (double, __pyx_t_double_complex, int __pyx_skip_dispatch)", + "huber": "double (double, double, int __pyx_skip_dispatch)", + "hyperu": "double (double, double, double, int __pyx_skip_dispatch)", + "i0": "double (double, int __pyx_skip_dispatch)", + "i0e": "double (double, int __pyx_skip_dispatch)", + "i1": "double (double, int __pyx_skip_dispatch)", + "i1e": "double (double, int __pyx_skip_dispatch)", + "inv_boxcox": "double (double, double, int __pyx_skip_dispatch)", + "inv_boxcox1p": "double (double, double, int __pyx_skip_dispatch)", + "it2i0k0": "void (double, double *, double *)", + "it2j0y0": "void (double, double *, double *)", + "it2struve0": "double (double, int __pyx_skip_dispatch)", + "itairy": "void (double, double *, double *, double *, double *)", + "iti0k0": "void (double, double *, double *)", + "itj0y0": "void (double, double *, double *)", + "itmodstruve0": "double (double, int __pyx_skip_dispatch)", + "itstruve0": "double (double, int __pyx_skip_dispatch)", + "j0": "double (double, int __pyx_skip_dispatch)", + "j1": "double (double, int __pyx_skip_dispatch)", + "k0": "double (double, int __pyx_skip_dispatch)", + "k0e": "double (double, int __pyx_skip_dispatch)", + "k1": "double (double, int __pyx_skip_dispatch)", + "k1e": "double (double, int __pyx_skip_dispatch)", + "kei": "double (double, int __pyx_skip_dispatch)", + "keip": "double (double, int __pyx_skip_dispatch)", + "kelvin": "void (double, __pyx_t_double_complex *, __pyx_t_double_complex *, __pyx_t_double_complex *, __pyx_t_double_complex *)", + "ker": "double (double, int __pyx_skip_dispatch)", + "kerp": "double (double, int __pyx_skip_dispatch)", + "kl_div": "double (double, double, int __pyx_skip_dispatch)", + "kolmogi": "double (double, int __pyx_skip_dispatch)", + "kolmogorov": "double (double, int __pyx_skip_dispatch)", + "log_wright_bessel": "double (double, double, double, int __pyx_skip_dispatch)", + "lpmv": "double (double, double, double, int __pyx_skip_dispatch)", + "mathieu_a": "double (double, double, int __pyx_skip_dispatch)", + "mathieu_b": "double (double, double, int __pyx_skip_dispatch)", + "mathieu_cem": "void (double, double, double, double *, double *)", + "mathieu_modcem1": "void (double, double, double, double *, double *)", + "mathieu_modcem2": "void (double, double, double, double *, double *)", + "mathieu_modsem1": "void (double, double, double, double *, double *)", + "mathieu_modsem2": "void (double, double, double, double *, double *)", + "mathieu_sem": "void (double, double, double, double *, double *)", + "modfresnelm": "void (double, __pyx_t_double_complex *, __pyx_t_double_complex *)", + "modfresnelp": "void (double, __pyx_t_double_complex *, __pyx_t_double_complex *)", + "modstruve": "double (double, double, int __pyx_skip_dispatch)", + "nbdtrik": "double (double, double, double, int __pyx_skip_dispatch)", + "nbdtrin": "double (double, double, double, int __pyx_skip_dispatch)", + "ncfdtridfd": "double (double, double, double, double, int __pyx_skip_dispatch)", + "ncfdtridfn": "double (double, double, double, double, int __pyx_skip_dispatch)", + "ncfdtrinc": "double (double, double, double, double, int __pyx_skip_dispatch)", + "nctdtridf": "double (double, double, double, int __pyx_skip_dispatch)", + "nctdtrinc": "double (double, double, double, int __pyx_skip_dispatch)", + "ndtri": "double (double, int __pyx_skip_dispatch)", + "ndtri_exp": "double (double, int __pyx_skip_dispatch)", + "nrdtrimn": "double (double, double, double, int __pyx_skip_dispatch)", + "nrdtrisd": "double (double, double, double, int __pyx_skip_dispatch)", + "obl_ang1": "void (double, double, double, double, double *, double *)", + "obl_ang1_cv": "void (double, double, double, double, double, double *, double *)", + "obl_cv": "double (double, double, double, int __pyx_skip_dispatch)", + "obl_rad1": "void (double, double, double, double, double *, double *)", + "obl_rad1_cv": "void (double, double, double, double, double, double *, double *)", + "obl_rad2": "void (double, double, double, double, double *, double *)", + "obl_rad2_cv": "void (double, double, double, double, double, double *, double *)", + "owens_t": "double (double, double, int __pyx_skip_dispatch)", + "pbdv": "void (double, double, double *, double *)", + "pbvv": "void (double, double, double *, double *)", + "pbwa": "void (double, double, double *, double *)", + "pdtr": "double (double, double, int __pyx_skip_dispatch)", + "pdtrc": "double (double, double, int __pyx_skip_dispatch)", + "poch": "double (double, double, int __pyx_skip_dispatch)", + "pro_ang1": "void (double, double, double, double, double *, double *)", + "pro_ang1_cv": "void (double, double, double, double, double, double *, double *)", + "pro_cv": "double (double, double, double, int __pyx_skip_dispatch)", + "pro_rad1": "void (double, double, double, double, double *, double *)", + "pro_rad1_cv": "void (double, double, double, double, double, double *, double *)", + "pro_rad2": "void (double, double, double, double, double *, double *)", + "pro_rad2_cv": "void (double, double, double, double, double, double *, double *)", + "pseudo_huber": "double (double, double, int __pyx_skip_dispatch)", + "radian": "double (double, double, double, int __pyx_skip_dispatch)", + "rel_entr": "double (double, double, int __pyx_skip_dispatch)", + "round": "double (double, int __pyx_skip_dispatch)", + "sindg": "double (double, int __pyx_skip_dispatch)", + "stdtridf": "double (double, double, int __pyx_skip_dispatch)", + "struve": "double (double, double, int __pyx_skip_dispatch)", + "tandg": "double (double, int __pyx_skip_dispatch)", + "tklmbda": "double (double, double, int __pyx_skip_dispatch)", + "voigt_profile": "double (double, double, double, int __pyx_skip_dispatch)", + "wofz": "__pyx_t_double_complex (__pyx_t_double_complex, int __pyx_skip_dispatch)", + "wright_bessel": "double (double, double, double, int __pyx_skip_dispatch)", + "y0": "double (double, int __pyx_skip_dispatch)", + "y1": "double (double, int __pyx_skip_dispatch)", + "zetac": "double (double, int __pyx_skip_dispatch)" + } +} diff --git a/v12x/pylib/scipy/special/tests/test_basic.py b/v12x/pylib/scipy/special/tests/test_basic.py new file mode 100644 index 0000000000000000000000000000000000000000..cfa9587f9c19ce9d37b54e049c6ff9de89b80d3b --- /dev/null +++ b/v12x/pylib/scipy/special/tests/test_basic.py @@ -0,0 +1,4868 @@ +# this program corresponds to special.py + +### Means test is not done yet +# E Means test is giving error (E) +# F Means test is failing (F) +# EF Means test is giving error and Failing +#! Means test is segfaulting +# 8 Means test runs forever + +### test_besselpoly +### test_modfresnelp +### test_modfresnelm +# test_pbdv_seq +### test_pbvv_seq +### test_sph_harm + +import functools +import itertools +import operator +import platform +import sys +import warnings + +import numpy as np +from numpy import (array, isnan, r_, arange, finfo, pi, sin, cos, tan, exp, + log, zeros, sqrt, asarray, inf, nan_to_num, real, arctan, double, + array_equal) + +import pytest +from pytest import raises as assert_raises +from numpy.testing import (assert_equal, assert_array_equal, assert_, + assert_allclose, assert_array_almost_equal_nulp) + +from scipy import special +import scipy.special._ufuncs as cephes +from scipy.special import ellipe, ellipk, ellipkm1 +from scipy.special import elliprc, elliprd, elliprf, elliprg, elliprj +from scipy.special import softplus +from scipy.special import mathieu_odd_coef, mathieu_even_coef, stirling2 +from scipy._lib._array_api import xp_assert_close, xp_assert_equal, SCIPY_ARRAY_API + +from scipy.special._basic import ( + _FACTORIALK_LIMITS_64BITS, _FACTORIALK_LIMITS_32BITS, _is_subdtype +) +from scipy.special._testutils import with_special_errors, \ + assert_func_equal, FuncData +from scipy.integrate import quad + +import math + + +native_int = np.int32 if ( + sys.platform == 'win32' + or platform.architecture()[0] == "32bit" +) else np.int64 + + +class TestCephes: + def test_airy(self): + cephes.airy(0) + + def test_airye(self): + cephes.airye(0) + + def test_binom(self): + n = np.array([0.264, 4, 5.2, 17]) + k = np.array([2, 0.4, 7, 3.3]) + nk = np.array(np.broadcast_arrays(n[:,None], k[None,:]) + ).reshape(2, -1).T + rknown = np.array([[-0.097152, 0.9263051596159367, 0.01858423645695389, + -0.007581020651518199],[6, 2.0214389119675666, 0, 2.9827344527963846], + [10.92, 2.22993515861399, -0.00585728, 10.468891352063146], + [136, 3.5252179590758828, 19448, 1024.5526916174495]]) + assert_func_equal(cephes.binom, rknown.ravel(), nk, rtol=1e-13) + + # Test branches in implementation + rng = np.random.RandomState(1234) + n = np.r_[np.arange(-7, 30), 1000*rng.rand(30) - 500] + k = np.arange(0, 102) + nk = np.array(np.broadcast_arrays(n[:,None], k[None,:]) + ).reshape(2, -1).T + + assert_func_equal(cephes.binom, + cephes.binom(nk[:,0], nk[:,1] * (1 + 1e-15)), + nk, + atol=1e-10, rtol=1e-10) + + def test_binom_2(self): + # Test branches in implementation + np.random.seed(1234) + n = np.r_[np.logspace(1, 300, 20)] + k = np.arange(0, 102) + nk = np.array(np.broadcast_arrays(n[:,None], k[None,:]) + ).reshape(2, -1).T + + assert_func_equal(cephes.binom, + cephes.binom(nk[:,0], nk[:,1] * (1 + 1e-15)), + nk, + atol=1e-10, rtol=1e-10) + + def test_binom_exact(self): + @np.vectorize + def binom_int(n, k): + n = int(n) + k = int(k) + num = 1 + den = 1 + for i in range(1, k+1): + num *= i + n - k + den *= i + return float(num/den) + + np.random.seed(1234) + n = np.arange(1, 15) + k = np.arange(0, 15) + nk = np.array(np.broadcast_arrays(n[:,None], k[None,:]) + ).reshape(2, -1).T + nk = nk[nk[:,0] >= nk[:,1]] + assert_func_equal(cephes.binom, + binom_int(nk[:,0], nk[:,1]), + nk, + atol=0, rtol=0) + + def test_binom_nooverflow_8346(self): + # Test (binom(n, k) doesn't overflow prematurely */ + dataset = [ + (1000, 500, 2.70288240945436551e+299), + (1002, 501, 1.08007396880791225e+300), + (1004, 502, 4.31599279169058121e+300), + (1006, 503, 1.72468101616263781e+301), + (1008, 504, 6.89188009236419153e+301), + (1010, 505, 2.75402257948335448e+302), + (1012, 506, 1.10052048531923757e+303), + (1014, 507, 4.39774063758732849e+303), + (1016, 508, 1.75736486108312519e+304), + (1018, 509, 7.02255427788423734e+304), + (1020, 510, 2.80626776829962255e+305), + (1022, 511, 1.12140876377061240e+306), + (1024, 512, 4.48125455209897109e+306), + (1026, 513, 1.79075474304149900e+307), + (1028, 514, 7.15605105487789676e+307) + ] + dataset = np.asarray(dataset) + FuncData(cephes.binom, dataset, (0, 1), 2, rtol=1e-12).check() + + def test_bdtr(self): + assert_equal(cephes.bdtr(1,1,0.5),1.0) + + def test_bdtri(self): + assert_equal(cephes.bdtri(1,3,0.5),0.5) + + def test_bdtrc(self): + assert_equal(cephes.bdtrc(1,3,0.5),0.5) + + def test_bdtrik(self): + cephes.bdtrik(1,3,0.5) + + def test_bei(self): + assert_equal(cephes.bei(0),0.0) + + def test_beip(self): + assert_equal(cephes.beip(0),0.0) + + def test_ber(self): + assert_equal(cephes.ber(0),1.0) + + def test_berp(self): + assert_equal(cephes.berp(0),0.0) + + def test_besselpoly(self): + assert_equal(cephes.besselpoly(0,0,0),1.0) + + def test_cbrt(self): + assert_allclose(cephes.cbrt(1), 1.0, atol=1e-6, rtol=0) + + def test_chdtr(self): + assert_equal(cephes.chdtr(1,0),0.0) + + def test_chdtrc(self): + assert_equal(cephes.chdtrc(1,0),1.0) + + def test_chdtri(self): + assert_equal(cephes.chdtri(1,1),0.0) + + def test_chndtrix(self): + assert_equal(cephes.chndtrix(0,1,0),0.0) + + def test_cosdg(self): + assert_equal(cephes.cosdg(0),1.0) + + def test_cosm1(self): + assert_equal(cephes.cosm1(0),0.0) + + def test_cotdg(self): + assert_allclose(cephes.cotdg(45), 1.0, atol=1.5e-7, rtol=0) + + def test_dawsn(self): + assert_equal(cephes.dawsn(0),0.0) + assert_allclose(cephes.dawsn(1.23), 0.50053727749081767) + + def test_diric(self): + # Test behavior near multiples of 2pi. Regression test for issue + # described in gh-4001. + n_odd = [1, 5, 25] + x = np.array(2*np.pi + 5e-5).astype(np.float32) + assert_allclose(special.diric(x, n_odd), 1.0, atol=1.5e-7, rtol=0) + x = np.array(2*np.pi + 1e-9).astype(np.float64) + assert_allclose(special.diric(x, n_odd), 1.0, atol=1.5e-15, rtol=0) + x = np.array(2*np.pi + 1e-15).astype(np.float64) + assert_allclose(special.diric(x, n_odd), 1.0, atol=1.5e-15, rtol=0) + if hasattr(np, 'float128'): + # No float128 available in 32-bit numpy + x = np.array(2*np.pi + 1e-12).astype(np.float128) + assert_allclose(special.diric(x, n_odd), 1.0, atol=1.5e-19, rtol=0) + + n_even = [2, 4, 24] + x = np.array(2*np.pi + 1e-9).astype(np.float64) + assert_allclose(special.diric(x, n_even), -1.0, atol=1.5e-15, rtol=0) + + # Test at some values not near a multiple of pi + x = np.arange(0.2*np.pi, 1.0*np.pi, 0.2*np.pi) + octave_result = [0.872677996249965, 0.539344662916632, + 0.127322003750035, -0.206011329583298] + assert_allclose(special.diric(x, 3), octave_result, atol=1.5e-15, rtol=0) + + def test_diric_broadcasting(self): + x = np.arange(5) + n = np.array([1, 3, 7]) + assert_(special.diric(x[:, np.newaxis], n).shape == (x.size, n.size)) + + def test_ellipe(self): + assert_equal(cephes.ellipe(1),1.0) + + def test_ellipeinc(self): + assert_equal(cephes.ellipeinc(0,1),0.0) + + def test_ellipj(self): + cephes.ellipj(0,1) + + def test_ellipk(self): + assert_allclose(ellipk(0), pi/2) + + def test_ellipkinc(self): + assert_equal(cephes.ellipkinc(0,0),0.0) + + def test_erf(self): + assert_equal(cephes.erf(0), 0.0) + + def test_erf_symmetry(self): + x = 5.905732037710919 + assert_equal(cephes.erf(x) + cephes.erf(-x), 0.0) + + def test_erfc(self): + assert_equal(cephes.erfc(0), 1.0) + + def test_exp10(self): + assert_allclose(cephes.exp10(2), 100.0, atol=1e-6, rtol=0) + + def test_exp2(self): + assert_equal(cephes.exp2(2),4.0) + + def test_expm1(self): + assert_equal(cephes.expm1(0),0.0) + assert_equal(cephes.expm1(np.inf), np.inf) + assert_equal(cephes.expm1(-np.inf), -1) + assert_equal(cephes.expm1(np.nan), np.nan) + + def test_expm1_complex(self): + expm1 = cephes.expm1 + assert_equal(expm1(0 + 0j), 0 + 0j) + assert_equal(expm1(complex(np.inf, 0)), complex(np.inf, 0)) + assert_equal(expm1(complex(np.inf, 1)), complex(np.inf, np.inf)) + assert_equal(expm1(complex(np.inf, 2)), complex(-np.inf, np.inf)) + assert_equal(expm1(complex(np.inf, 4)), complex(-np.inf, -np.inf)) + assert_equal(expm1(complex(np.inf, 5)), complex(np.inf, -np.inf)) + assert_equal(expm1(complex(1, np.inf)), complex(np.nan, np.nan)) + assert_equal(expm1(complex(0, np.inf)), complex(np.nan, np.nan)) + assert_equal(expm1(complex(np.inf, np.inf)), complex(np.inf, np.nan)) + assert_equal(expm1(complex(-np.inf, np.inf)), complex(-1, 0)) + assert_equal(expm1(complex(-np.inf, np.nan)), complex(-1, 0)) + assert_equal(expm1(complex(np.inf, np.nan)), complex(np.inf, np.nan)) + assert_equal(expm1(complex(0, np.nan)), complex(np.nan, np.nan)) + assert_equal(expm1(complex(1, np.nan)), complex(np.nan, np.nan)) + assert_equal(expm1(complex(np.nan, 1)), complex(np.nan, np.nan)) + assert_equal(expm1(complex(np.nan, np.nan)), complex(np.nan, np.nan)) + + @pytest.mark.xfail(reason='The real part of expm1(z) bad at these points') + def test_expm1_complex_hard(self): + # The real part of this function is difficult to evaluate when + # z.real = -log(cos(z.imag)). + y = np.array([0.1, 0.2, 0.3, 5, 11, 20]) + x = -np.log(np.cos(y)) + z = x + 1j*y + + # evaluate using mpmath.expm1 with dps=1000 + expected = np.array([-5.5507901846769623e-17+0.10033467208545054j, + 2.4289354732893695e-18+0.20271003550867248j, + 4.5235500262585768e-17+0.30933624960962319j, + 7.8234305217489006e-17-3.3805150062465863j, + -1.3685191953697676e-16-225.95084645419513j, + 8.7175620481291045e-17+2.2371609442247422j]) + found = cephes.expm1(z) + # this passes. + assert_array_almost_equal_nulp(found.imag, expected.imag, 3) + # this fails. + assert_array_almost_equal_nulp(found.real, expected.real, 20) + + def test_fdtr(self): + assert_equal(cephes.fdtr(1, 1, 0), 0.0) + # Computed using Wolfram Alpha: CDF[FRatioDistribution[1e-6, 5], 10] + assert_allclose(cephes.fdtr(1e-6, 5, 10), 0.9999940790193488, + rtol=1e-12) + + def test_fdtrc(self): + assert_equal(cephes.fdtrc(1, 1, 0), 1.0) + # Computed using Wolfram Alpha: + # 1 - CDF[FRatioDistribution[2, 1/10], 1e10] + assert_allclose(cephes.fdtrc(2, 0.1, 1e10), 0.27223784621293512, + rtol=1e-12) + + def test_fdtri(self): + assert_allclose(cephes.fdtri(1, 1, [0.499, 0.501]), + array([0.9937365, 1.00630298]), rtol=1e-6) + # From Wolfram Alpha: + # CDF[FRatioDistribution[1/10, 1], 3] = 0.8756751669632105666874... + p = 0.8756751669632105666874 + assert_allclose(cephes.fdtri(0.1, 1, p), 3, rtol=1e-12) + + def test_gh20835(self): + # gh-20835 reported fdtri failing for extreme inputs + dfd, dfn, x = 1, 50000, 29.72591544307521 + assert_allclose(cephes.fdtri(dfd, dfn, cephes.fdtr(dfd, dfn, x)), x, rtol=1e-15) + + def test_fdtri_mysterious_failure(self): + assert_allclose(cephes.fdtri(1, 1, 0.5), 1) + + def test_fdtridfd(self): + assert_equal(cephes.fdtridfd(1,0,0),5.0) + + def test_fresnel(self): + assert_equal(cephes.fresnel(0),(0.0,0.0)) + + def test_gamma(self): + assert_equal(cephes.gamma(5),24.0) + + def test_gammainccinv(self): + assert_equal(cephes.gammainccinv(5,1),0.0) + + def test_gammaln(self): + cephes.gammaln(10) + + def test_gammasgn(self): + vals = np.array( + [-np.inf, -4, -3.5, -2.3, -0.0, 0.0, 1, 4.2, np.inf], np.float64 + ) + reference = np.array( + [np.nan, np.nan, 1.0, -1.0, -1.0, 1.0, 1.0, 1.0, 1.0], np.float64 + ) + assert_array_equal(cephes.gammasgn(vals), reference) + + def test_gdtr(self): + assert_equal(cephes.gdtr(1,1,0),0.0) + + def test_gdtr_inf(self): + assert_equal(cephes.gdtr(1,1,np.inf),1.0) + + def test_gdtrc(self): + assert_equal(cephes.gdtrc(1,1,0),1.0) + + def test_gdtria(self): + assert_equal(cephes.gdtria(0,1,1),0.0) + + def test_gdtrib(self): + cephes.gdtrib(1,0,1) + # assert_equal(cephes.gdtrib(1,0,1),5.0) + + def test_gdtrix(self): + cephes.gdtrix(1,1,.1) + + def test_hankel1(self): + cephes.hankel1(1,1) + + def test_hankel1e(self): + cephes.hankel1e(1,1) + + def test_hankel2(self): + cephes.hankel2(1,1) + + def test_hankel2e(self): + cephes.hankel2e(1,1) + + def test_hyp1f1(self): + assert_allclose(cephes.hyp1f1(1, 1, 1), exp(1.0), atol=1e-6, rtol=0) + assert_allclose(cephes.hyp1f1(3, 4, -6), 0.026056422099537251095, + atol=1e-6, rtol=0) + cephes.hyp1f1(1,1,1) + + def test_hyp2f1(self): + assert_equal(cephes.hyp2f1(1,1,1,0),1.0) + + def test_i0(self): + assert_equal(cephes.i0(0),1.0) + + def test_i0e(self): + assert_equal(cephes.i0e(0),1.0) + + def test_i1(self): + assert_equal(cephes.i1(0),0.0) + + def test_i1e(self): + assert_equal(cephes.i1e(0),0.0) + + def test_it2i0k0(self): + cephes.it2i0k0(1) + + def test_it2j0y0(self): + cephes.it2j0y0(1) + + def test_it2struve0(self): + cephes.it2struve0(1) + + def test_itairy(self): + cephes.itairy(1) + + def test_iti0k0(self): + assert_equal(cephes.iti0k0(0),(0.0,0.0)) + + def test_itj0y0(self): + assert_equal(cephes.itj0y0(0),(0.0,0.0)) + + def test_itmodstruve0(self): + assert_equal(cephes.itmodstruve0(0),0.0) + + def test_itstruve0(self): + assert_equal(cephes.itstruve0(0),0.0) + + def test_iv(self): + assert_equal(cephes.iv(1,0),0.0) + + def test_ive(self): + assert_equal(cephes.ive(1,0),0.0) + + def test_j0(self): + assert_equal(cephes.j0(0),1.0) + + def test_j1(self): + assert_equal(cephes.j1(0),0.0) + + def test_jn(self): + assert_equal(cephes.jn(0,0),1.0) + + def test_jv(self): + assert_equal(cephes.jv(0,0),1.0) + + def test_jve(self): + assert_equal(cephes.jve(0,0),1.0) + + def test_k0(self): + cephes.k0(2) + + def test_k0e(self): + cephes.k0e(2) + + def test_k1(self): + cephes.k1(2) + + def test_k1e(self): + cephes.k1e(2) + + def test_kei(self): + cephes.kei(2) + + def test_keip(self): + assert_equal(cephes.keip(0),0.0) + + def test_ker(self): + cephes.ker(2) + + def test_kerp(self): + cephes.kerp(2) + + def test_kelvin(self): + cephes.kelvin(2) + + def test_kn(self): + cephes.kn(1,1) + + def test_kolmogi(self): + assert_equal(cephes.kolmogi(1),0.0) + assert_(np.isnan(cephes.kolmogi(np.nan))) + + def test_kolmogorov(self): + assert_equal(cephes.kolmogorov(0), 1.0) + + def test_kolmogp(self): + assert_equal(cephes._kolmogp(0), -0.0) + + def test_kolmogc(self): + assert_equal(cephes._kolmogc(0), 0.0) + + def test_kolmogci(self): + assert_equal(cephes._kolmogci(0), 0.0) + assert_(np.isnan(cephes._kolmogci(np.nan))) + + def test_kv(self): + cephes.kv(1,1) + + def test_kve(self): + cephes.kve(1,1) + + def test_log1p(self): + log1p = cephes.log1p + assert_equal(log1p(0), 0.0) + assert_equal(log1p(-1), -np.inf) + assert_equal(log1p(-2), np.nan) + assert_equal(log1p(np.inf), np.inf) + + def test_log1p_complex(self): + log1p = cephes.log1p + c = complex + assert_equal(log1p(0 + 0j), 0 + 0j) + assert_equal(log1p(c(-1, 0)), c(-np.inf, 0)) + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", "invalid value encountered in multiply", RuntimeWarning) + assert_allclose(log1p(c(1, np.inf)), c(np.inf, np.pi/2)) + assert_equal(log1p(c(1, np.nan)), c(np.nan, np.nan)) + assert_allclose(log1p(c(-np.inf, 1)), c(np.inf, np.pi)) + assert_equal(log1p(c(np.inf, 1)), c(np.inf, 0)) + assert_allclose(log1p(c(-np.inf, np.inf)), c(np.inf, 3*np.pi/4)) + assert_allclose(log1p(c(np.inf, np.inf)), c(np.inf, np.pi/4)) + assert_equal(log1p(c(np.inf, np.nan)), c(np.inf, np.nan)) + assert_equal(log1p(c(-np.inf, np.nan)), c(np.inf, np.nan)) + assert_equal(log1p(c(np.nan, np.inf)), c(np.inf, np.nan)) + assert_equal(log1p(c(np.nan, 1)), c(np.nan, np.nan)) + assert_equal(log1p(c(np.nan, np.nan)), c(np.nan, np.nan)) + + def test_lpmv(self): + assert_equal(cephes.lpmv(0,0,1),1.0) + + def test_mathieu_a_q0(self): + # When q is 0, the exact result is m**2. + m = np.array([1, 2, 5]) + assert_equal(cephes.mathieu_a(m, 0), m**2) + + # Reference values were computed with Wolfram Alpha: + # MathieuCharacteristicA[m, q] + @pytest.mark.parametrize( + 'm, q, ref', + [(0, 8, -10.6067292355526479852024), + (3, 3/2, 9.19330104768060974047804), + (5, 1/4, 25.0013021454698022809572), + (8, -10, 64.8008910105046444848962)] + ) + def test_mathieu_a(self, m, q, ref): + y = cephes.mathieu_a(m, q) + assert_allclose(y, ref, rtol=1e-15) + + def test_mathieu_b_q0(self): + # When q is 0, the exact result is m**2. + m = np.array([1, 2, 5]) + assert_equal(cephes.mathieu_b(m, 0), m**2) + + # Reference values were computed with Wolfram Alpha: + # MathieuCharacteristicB[m, q] + @pytest.mark.parametrize( + 'm, q, ref', + [(1, 15, -22.5130034974234666335), + (5, 3, 25.1870798027185125480), + (9, 1/4, 81.00039062627570760760), + (10, -3, 100.0454683359769326164)] + ) + def test_mathieu_b(self, m, q, ref): + y = cephes.mathieu_b(m, q) + assert_allclose(y, ref, rtol=1e-15) + + def test_mathieu_cem(self): + assert_equal(cephes.mathieu_cem(1,0,0),(1.0,0.0)) + + # Test AMS 20.2.27 + @np.vectorize + def ce_smallq(m, q, z): + z *= np.pi/180 + if m == 0: + # + O(q^2) + return 2**(-0.5) * (1 - .5*q*cos(2*z)) + elif m == 1: + # + O(q^2) + return cos(z) - q/8 * cos(3*z) + elif m == 2: + # + O(q^2) + return cos(2*z) - q*(cos(4*z)/12 - 1/4) + else: + # + O(q^2) + return cos(m*z) - q*(cos((m+2)*z)/(4*(m+1)) - cos((m-2)*z)/(4*(m-1))) + m = np.arange(0, 100) + q = np.r_[0, np.logspace(-30, -9, 10)] + assert_allclose(cephes.mathieu_cem(m[:,None], q[None,:], 0.123)[0], + ce_smallq(m[:,None], q[None,:], 0.123), + rtol=1e-14, atol=0) + + def test_mathieu_sem(self): + assert_equal(cephes.mathieu_sem(1,0,0),(0.0,1.0)) + + # Test AMS 20.2.27 + @np.vectorize + def se_smallq(m, q, z): + z *= np.pi/180 + if m == 1: + # + O(q^2) + return sin(z) - q/8 * sin(3*z) + elif m == 2: + # + O(q^2) + return sin(2*z) - q*sin(4*z)/12 + else: + # + O(q^2) + return sin(m*z) - q*(sin((m+2)*z)/(4*(m+1)) - sin((m-2)*z)/(4*(m-1))) + m = np.arange(1, 100) + q = np.r_[0, np.logspace(-30, -9, 10)] + assert_allclose(cephes.mathieu_sem(m[:,None], q[None,:], 0.123)[0], + se_smallq(m[:,None], q[None,:], 0.123), + rtol=1e-14, atol=0) + + def test_mathieu_modcem1(self): + assert_equal(cephes.mathieu_modcem1(1,0,0),(0.0,0.0)) + + def test_mathieu_modcem2(self): + cephes.mathieu_modcem2(1,1,1) + + # Test reflection relation AMS 20.6.19 + m = np.arange(0, 4)[:,None,None] + q = np.r_[np.logspace(-2, 2, 10)][None,:,None] + z = np.linspace(0, 1, 7)[None,None,:] + + y1 = cephes.mathieu_modcem2(m, q, -z)[0] + + fr = -cephes.mathieu_modcem2(m, q, 0)[0] / cephes.mathieu_modcem1(m, q, 0)[0] + y2 = (-cephes.mathieu_modcem2(m, q, z)[0] + - 2*fr*cephes.mathieu_modcem1(m, q, z)[0]) + + assert_allclose(y1, y2, rtol=1e-10) + + def test_mathieu_modsem1(self): + assert_equal(cephes.mathieu_modsem1(1,0,0),(0.0,0.0)) + + def test_mathieu_modsem2(self): + cephes.mathieu_modsem2(1,1,1) + + # Test reflection relation AMS 20.6.20 + m = np.arange(1, 4)[:,None,None] + q = np.r_[np.logspace(-2, 2, 10)][None,:,None] + z = np.linspace(0, 1, 7)[None,None,:] + + y1 = cephes.mathieu_modsem2(m, q, -z)[0] + fr = cephes.mathieu_modsem2(m, q, 0)[1] / cephes.mathieu_modsem1(m, q, 0)[1] + y2 = (cephes.mathieu_modsem2(m, q, z)[0] + - 2*fr*cephes.mathieu_modsem1(m, q, z)[0]) + assert_allclose(y1, y2, rtol=1e-10) + + def test_mathieu_overflow(self): + # Check that these return NaNs instead of causing a SEGV + assert_equal(cephes.mathieu_cem(10000, 0, 1.3), (np.nan, np.nan)) + assert_equal(cephes.mathieu_sem(10000, 0, 1.3), (np.nan, np.nan)) + assert_equal(cephes.mathieu_cem(10000, 1.5, 1.3), (np.nan, np.nan)) + assert_equal(cephes.mathieu_sem(10000, 1.5, 1.3), (np.nan, np.nan)) + assert_equal(cephes.mathieu_modcem1(10000, 1.5, 1.3), (np.nan, np.nan)) + assert_equal(cephes.mathieu_modsem1(10000, 1.5, 1.3), (np.nan, np.nan)) + assert_equal(cephes.mathieu_modcem2(10000, 1.5, 1.3), (np.nan, np.nan)) + assert_equal(cephes.mathieu_modsem2(10000, 1.5, 1.3), (np.nan, np.nan)) + + def test_mathieu_ticket_1847(self): + # Regression test --- this call had some out-of-bounds access + # and could return nan occasionally + for k in range(60): + v = cephes.mathieu_modsem2(2, 100, -1) + # Values from ACM TOMS 804 (derivate by numerical differentiation) + assert_allclose(v[0], 0.1431742913063671074347, rtol=1e-10) + assert_allclose(v[1], 0.9017807375832909144719, rtol=1e-4) + + def test_modfresnelm(self): + cephes.modfresnelm(0) + + def test_modfresnelp(self): + cephes.modfresnelp(0) + + def test_modstruve(self): + assert_equal(cephes.modstruve(1,0),0.0) + + def test_nbdtr(self): + assert_equal(cephes.nbdtr(1,1,1),1.0) + + def test_nbdtrc(self): + assert_equal(cephes.nbdtrc(1,1,1),0.0) + + def test_nbdtri(self): + assert_equal(cephes.nbdtri(1,1,1),1.0) + + def test_nbdtrik(self): + cephes.nbdtrik(1,.4,.5) + + def test_ncfdtr(self): + assert_equal(cephes.ncfdtr(1,1,1,0),0.0) + + def test_ncfdtri(self): + assert_equal(cephes.ncfdtri(1, 1, 1, 0), 0.0) + f = [0.5, 1, 1.5] + p = cephes.ncfdtr(2, 3, 1.5, f) + assert_allclose(cephes.ncfdtri(2, 3, 1.5, p), f) + + @pytest.mark.xfail( + reason=( + "ncfdtr uses a Boost math implementation but ncfdtridfd" + "inverts the less accurate cdflib implementation of ncfdtr." + ) + ) + def test_ncfdtridfd(self): + dfd = [1, 2, 3] + p = cephes.ncfdtr(2, dfd, 0.25, 15) + assert_allclose(cephes.ncfdtridfd(2, p, 0.25, 15), dfd) + + @pytest.mark.xfail( + reason=( + "ncfdtr uses a Boost math implementation but ncfdtridfn" + "inverts the less accurate cdflib implementation of ncfdtr." + ) + ) + def test_ncfdtridfn(self): + dfn = [0.1, 1, 2, 3, 1e4] + p = cephes.ncfdtr(dfn, 2, 0.25, 15) + assert_allclose(cephes.ncfdtridfn(p, 2, 0.25, 15), dfn, rtol=1e-5) + + @pytest.mark.xfail( + reason=( + "ncfdtr uses a Boost math implementation but ncfdtrinc" + "inverts the less accurate cdflib implementation of ncfdtr." + ) + ) + def test_ncfdtrinc(self): + nc = [0.5, 1.5, 2.0] + p = cephes.ncfdtr(2, 3, nc, 15) + assert_allclose(cephes.ncfdtrinc(2, 3, p, 15), nc) + + def test_nctdtr(self): + assert_equal(cephes.nctdtr(1,0,0),0.5) + assert_equal(cephes.nctdtr(9, 65536, 45), 0.0) + + assert_allclose(cephes.nctdtr(np.inf, 1., 1.), 0.5, atol=1e-4, rtol=0) + assert_(np.isnan(cephes.nctdtr(2., np.inf, 10.))) + assert_allclose(cephes.nctdtr(2., 1., np.inf), 1., atol=1e-6, rtol=0) + + assert_(np.isnan(cephes.nctdtr(np.nan, 1., 1.))) + assert_(np.isnan(cephes.nctdtr(2., np.nan, 1.))) + assert_(np.isnan(cephes.nctdtr(2., 1., np.nan))) + + def test_nctdtridf(self): + cephes.nctdtridf(1,0.5,0) + + def test_nctdtrinc(self): + cephes.nctdtrinc(1,0,0) + + def test_nctdtrit(self): + cephes.nctdtrit(.1,0.2,.5) + + def test_nrdtrimn(self): + assert_allclose(cephes.nrdtrimn(0.5, 1, 1), 1.0, atol=1e-6, rtol=0) + + def test_obl_ang1(self): + cephes.obl_ang1(1,1,1,0) + + def test_obl_ang1_cv(self): + result = cephes.obl_ang1_cv(1,1,1,1,0) + assert_allclose(result[0], 1.0, atol=1.5e-7, rtol=0) + assert_allclose(result[1], 0.0, atol=1.5e-7, rtol=0) + + def test_obl_cv(self): + assert_equal(cephes.obl_cv(1,1,0),2.0) + + def test_obl_rad1(self): + cephes.obl_rad1(1,1,1,0) + + def test_obl_rad1_cv(self): + cephes.obl_rad1_cv(1,1,1,1,0) + + def test_obl_rad2(self): + cephes.obl_rad2(1,1,1,0) + + def test_obl_rad2_cv(self): + cephes.obl_rad2_cv(1,1,1,1,0) + + def test_pbdv(self): + assert_equal(cephes.pbdv(1,0),(0.0,1.0)) + + def test_pbvv(self): + cephes.pbvv(1,0) + + def test_pbwa(self): + cephes.pbwa(1,0) + + def test_pdtr(self): + val = cephes.pdtr(0, 1) + assert_allclose(val, np.exp(-1), atol=1.5e-7, rtol=0) + # Edge case: m = 0. + val = cephes.pdtr([0, 1, 2], 0) + assert_array_equal(val, [1, 1, 1]) + + def test_pdtrc(self): + val = cephes.pdtrc(0, 1) + assert_allclose(val, 1 - np.exp(-1), atol=1.5e-7, rtol=0) + # Edge case: m = 0. + val = cephes.pdtrc([0, 1, 2], 0.0) + assert_array_equal(val, [0, 0, 0]) + + def test_pdtri(self): + with warnings.catch_warnings(): + msg = "floating point number truncated to an integer" + warnings.filterwarnings("ignore", msg, RuntimeWarning) + cephes.pdtri(0.5,0.5) + + + def test_pro_ang1(self): + cephes.pro_ang1(1,1,1,0) + + def test_pro_ang1_cv(self): + assert_allclose(cephes.pro_ang1_cv(1, 1, 1, 1, 0), array((1.0, 0.0)), + atol=1.5e-7, rtol=0) + + def test_pro_cv(self): + assert_equal(cephes.pro_cv(1,1,0),2.0) + + def test_pro_rad1(self): + cephes.pro_rad1(1,1,1,0.1) + + def test_pro_rad1_cv(self): + cephes.pro_rad1_cv(1,1,1,1,0) + + def test_pro_rad2(self): + cephes.pro_rad2(1,1,1,0) + + def test_pro_rad2_cv(self): + cephes.pro_rad2_cv(1,1,1,1,0) + + def test_psi(self): + cephes.psi(1) + + def test_radian(self): + assert_equal(cephes.radian(0,0,0),0) + + def test_rgamma(self): + assert_equal(cephes.rgamma(1),1.0) + + def test_round(self): + assert_equal(cephes.round(3.4),3.0) + assert_equal(cephes.round(-3.4),-3.0) + assert_equal(cephes.round(3.6),4.0) + assert_equal(cephes.round(-3.6),-4.0) + assert_equal(cephes.round(3.5),4.0) + assert_equal(cephes.round(-3.5),-4.0) + + def test_shichi(self): + cephes.shichi(1) + + def test_sici(self): + cephes.sici(1) + + s, c = cephes.sici(np.inf) + assert_allclose(s, np.pi * 0.5, atol=1.5e-7, rtol=0) + assert_allclose(c, 0, atol=1.5e-7, rtol=0) + + s, c = cephes.sici(-np.inf) + assert_allclose(s, -np.pi * 0.5, atol=1.5e-7, rtol=0) + assert_(np.isnan(c), "cosine integral(-inf) is not nan") + + def test_sindg(self): + assert_equal(cephes.sindg(90),1.0) + + def test_smirnov(self): + assert_equal(cephes.smirnov(1,.1),0.9) + assert_(np.isnan(cephes.smirnov(1,np.nan))) + + def test_smirnovp(self): + assert_equal(cephes._smirnovp(1, .1), -1) + assert_equal(cephes._smirnovp(2, 0.75), -2*(0.25)**(2-1)) + assert_equal(cephes._smirnovp(3, 0.75), -3*(0.25)**(3-1)) + assert_(np.isnan(cephes._smirnovp(1, np.nan))) + + def test_smirnovc(self): + assert_equal(cephes._smirnovc(1,.1),0.1) + assert_(np.isnan(cephes._smirnovc(1,np.nan))) + x10 = np.linspace(0, 1, 11, endpoint=True) + assert_allclose(cephes._smirnovc(3, x10), 1 - cephes.smirnov(3, x10), + atol=1.5e-7, rtol=0) + x4 = np.linspace(0, 1, 5, endpoint=True) + assert_allclose(cephes._smirnovc(4, x4), 1 - cephes.smirnov(4, x4), + atol=1.5e-7, rtol=0) + + def test_smirnovi(self): + assert_allclose(cephes.smirnov(1, cephes.smirnovi(1, 0.4)), 0.4, + atol=1.5e-7, rtol=0) + assert_allclose(cephes.smirnov(1, cephes.smirnovi(1, 0.6)), 0.6, + atol=1.5e-7, rtol=0) + assert_(np.isnan(cephes.smirnovi(1,np.nan))) + + def test_smirnovci(self): + assert_allclose(cephes._smirnovc(1, cephes._smirnovci(1, 0.4)), 0.4, + atol=1.5e-7, rtol=0) + assert_allclose(cephes._smirnovc(1, cephes._smirnovci(1, 0.6)), 0.6, + atol=1.5e-7, rtol=0) + assert_(np.isnan(cephes._smirnovci(1,np.nan))) + + def test_spence(self): + assert_equal(cephes.spence(1),0.0) + + def test_stdtr(self): + assert_equal(cephes.stdtr(1,0),0.5) + assert_allclose(cephes.stdtr(1, 1), 0.75, atol=1.5e-7, rtol=0) + assert_allclose(cephes.stdtr(1, 2), 0.852416382349, atol=1.5e-7, rtol=0) + + def test_stdtridf(self): + cephes.stdtridf(0.7,1) + + def test_stdtrit(self): + cephes.stdtrit(1,0.7) + + def test_struve(self): + assert_equal(cephes.struve(0,0),0.0) + + def test_tandg(self): + assert_equal(cephes.tandg(45),1.0) + + def test_tklmbda(self): + assert_allclose(cephes.tklmbda(1, 1), 1.0, atol=1.5e-7, rtol=0) + + def test_y0(self): + cephes.y0(1) + + def test_y1(self): + cephes.y1(1) + + def test_yn(self): + cephes.yn(1,1) + + def test_yv(self): + cephes.yv(1,1) + + def test_yve(self): + cephes.yve(1,1) + + def test_wofz(self): + z = [complex(624.2,-0.26123), complex(-0.4,3.), complex(0.6,2.), + complex(-1.,1.), complex(-1.,-9.), complex(-1.,9.), + complex(-0.0000000234545,1.1234), complex(-3.,5.1), + complex(-53,30.1), complex(0.0,0.12345), + complex(11,1), complex(-22,-2), complex(9,-28), + complex(21,-33), complex(1e5,1e5), complex(1e14,1e14) + ] + w = [ + complex(-3.78270245518980507452677445620103199303131110e-7, + 0.000903861276433172057331093754199933411710053155), + complex(0.1764906227004816847297495349730234591778719532788, + -0.02146550539468457616788719893991501311573031095617), + complex(0.2410250715772692146133539023007113781272362309451, + 0.06087579663428089745895459735240964093522265589350), + complex(0.30474420525691259245713884106959496013413834051768, + -0.20821893820283162728743734725471561394145872072738), + complex(7.317131068972378096865595229600561710140617977e34, + 8.321873499714402777186848353320412813066170427e34), + complex(0.0615698507236323685519612934241429530190806818395, + -0.00676005783716575013073036218018565206070072304635), + complex(0.3960793007699874918961319170187598400134746631, + -5.593152259116644920546186222529802777409274656e-9), + complex(0.08217199226739447943295069917990417630675021771804, + -0.04701291087643609891018366143118110965272615832184), + complex(0.00457246000350281640952328010227885008541748668738, + -0.00804900791411691821818731763401840373998654987934), + complex(0.8746342859608052666092782112565360755791467973338452, + 0.), + complex(0.00468190164965444174367477874864366058339647648741, + 0.0510735563901306197993676329845149741675029197050), + complex(-0.0023193175200187620902125853834909543869428763219, + -0.025460054739731556004902057663500272721780776336), + complex(9.11463368405637174660562096516414499772662584e304, + 3.97101807145263333769664875189354358563218932e305), + complex(-4.4927207857715598976165541011143706155432296e281, + -2.8019591213423077494444700357168707775769028e281), + complex(2.820947917809305132678577516325951485807107151e-6, + 2.820947917668257736791638444590253942253354058e-6), + complex(2.82094791773878143474039725787438662716372268e-15, + 2.82094791773878143474039725773333923127678361e-15) + ] + assert_func_equal(cephes.wofz, w, z, rtol=1e-13) + + +class TestAiry: + def test_airy(self): + # This tests the airy function to ensure 8 place accuracy in computation + + x = special.airy(.99) + assert_allclose(x, array([0.13689066, -0.16050153, 1.19815925, 0.92046818]), + atol=1.5e-8, rtol=0) + x = special.airy(.41) + assert_allclose(x, array([0.25238916, -.23480512, 0.80686202, 0.51053919]), + atol=1.5e-8, rtol=0) + x = special.airy(-.36) + assert_allclose(x, array([0.44508477,-0.23186773,0.44939534,0.48105354]), + atol=1.5e-8, rtol=0) + + def test_airye(self): + a = special.airye(0.01) + b = special.airy(0.01) + b1 = [None]*4 + for n in range(2): + b1[n] = b[n]*exp(2.0/3.0*0.01*sqrt(0.01)) + for n in range(2,4): + b1[n] = b[n]*exp(-abs(real(2.0/3.0*0.01*sqrt(0.01)))) + assert_allclose(a, b1, atol=1.5e-6, rtol=0) + + def test_bi_zeros(self): + bi = special.bi_zeros(2) + bia = (array([-1.17371322, -3.2710930]), + array([-2.29443968, -4.07315509]), + array([-0.45494438, 0.39652284]), + array([0.60195789, -0.76031014])) + assert_allclose(bi, bia, atol=1.5e-4, rtol=0) + + bi = special.bi_zeros(5) + assert_allclose(bi[0], array([-1.173713222709127, + -3.271093302836352, + -4.830737841662016, + -6.169852128310251, + -7.376762079367764]), + atol=1.5e-11, rtol=0) + + assert_allclose(bi[1], array([-2.294439682614122, + -4.073155089071828, + -5.512395729663599, + -6.781294445990305, + -7.940178689168587]), + atol=1.5e-10, rtol=0) + + assert_allclose(bi[2], array([-0.454944383639657, + 0.396522836094465, + -0.367969161486959, + 0.349499116831805, + -0.336026240133662]), + atol=1.5e-11, rtol=0) + + assert_allclose(bi[3], array([0.601957887976239, + -0.760310141492801, + 0.836991012619261, + -0.88947990142654, + 0.929983638568022]), + atol=1.5e-10, rtol=0) + + def test_ai_zeros(self): + ai = special.ai_zeros(1) + assert_allclose(ai, (array([-2.33810741]), + array([-1.01879297]), + array([0.5357]), + array([0.7012])), + atol=1.5e-4, rtol=0) + + @pytest.mark.fail_slow(5) + def test_ai_zeros_big(self): + z, zp, ai_zpx, aip_zx = special.ai_zeros(50000) + ai_z, aip_z, _, _ = special.airy(z) + ai_zp, aip_zp, _, _ = special.airy(zp) + + ai_envelope = 1/abs(z)**(1./4) + aip_envelope = abs(zp)**(1./4) + + # Check values + assert_allclose(ai_zpx, ai_zp, rtol=1e-10) + assert_allclose(aip_zx, aip_z, rtol=1e-10) + + # Check they are zeros + assert_allclose(ai_z/ai_envelope, 0, atol=1e-10, rtol=0) + assert_allclose(aip_zp/aip_envelope, 0, atol=1e-10, rtol=0) + + # Check first zeros, DLMF 9.9.1 + assert_allclose(z[:6], + [-2.3381074105, -4.0879494441, -5.5205598281, + -6.7867080901, -7.9441335871, -9.0226508533], rtol=1e-10) + assert_allclose(zp[:6], + [-1.0187929716, -3.2481975822, -4.8200992112, + -6.1633073556, -7.3721772550, -8.4884867340], rtol=1e-10) + + @pytest.mark.fail_slow(5) + def test_bi_zeros_big(self): + z, zp, bi_zpx, bip_zx = special.bi_zeros(50000) + _, _, bi_z, bip_z = special.airy(z) + _, _, bi_zp, bip_zp = special.airy(zp) + + bi_envelope = 1/abs(z)**(1./4) + bip_envelope = abs(zp)**(1./4) + + # Check values + assert_allclose(bi_zpx, bi_zp, rtol=1e-10) + assert_allclose(bip_zx, bip_z, rtol=1e-10) + + # Check they are zeros + assert_allclose(bi_z/bi_envelope, 0, atol=1e-10, rtol=0) + assert_allclose(bip_zp/bip_envelope, 0, atol=1e-10, rtol=0) + + # Check first zeros, DLMF 9.9.2 + assert_allclose(z[:6], + [-1.1737132227, -3.2710933028, -4.8307378417, + -6.1698521283, -7.3767620794, -8.4919488465], rtol=1e-10) + assert_allclose(zp[:6], + [-2.2944396826, -4.0731550891, -5.5123957297, + -6.7812944460, -7.9401786892, -9.0195833588], rtol=1e-10) + + +class TestAssocLaguerre: + def test_assoc_laguerre(self): + a1 = special.genlaguerre(11,1) + a2 = special.assoc_laguerre(.2,11,1) + assert_allclose(a2, a1(.2), atol=1.5e-8, rtol=0) + a2 = special.assoc_laguerre(1,11,1) + assert_allclose(a2, a1(1), atol=1.5e-8, rtol=0) + + +class TestBesselpoly: + def test_besselpoly(self): + pass + + +class TestKelvin: + def test_bei(self): + mbei = special.bei(2) + # This may not be exact. + assert_allclose(mbei, 0.9722916273066613, atol=1.5e-5, rtol=0) + + def test_beip(self): + mbeip = special.beip(2) + # This may not be exact. + assert_allclose(mbeip, 0.91701361338403631, atol=1.5e-5, rtol=0) + + def test_ber(self): + mber = special.ber(2) + # This may not be exact. + assert_allclose(mber, 0.75173418271380821, atol=1.5e-5, rtol=0) + + def test_berp(self): + mberp = special.berp(2) + # This may not be exact. + assert_allclose(mberp, -0.49306712470943909, atol=1.5e-5, rtol=0) + + def test_bei_zeros(self): + # Abramowitz & Stegun, Table 9.12 + bi = special.bei_zeros(5) + assert_allclose(bi, array([5.02622, + 9.45541, + 13.89349, + 18.33398, + 22.77544]), + atol=1.5e-4, rtol=0) + + def test_beip_zeros(self): + bip = special.beip_zeros(5) + assert_allclose(bip, array([3.772673304934953, + 8.280987849760042, + 12.742147523633703, + 17.193431752512542, + 21.641143941167325]), + atol=1.5e-8, rtol=0) + + def test_ber_zeros(self): + ber = special.ber_zeros(5) + assert_allclose(ber, array([2.84892, + 7.23883, + 11.67396, + 16.11356, + 20.55463]), + atol=1.5e-4, rtol=0) + + def test_berp_zeros(self): + brp = special.berp_zeros(5) + assert_allclose(brp, array([6.03871, + 10.51364, + 14.96844, + 19.41758, + 23.86430]), + atol=1.5e-4, rtol=0) + + def test_kelvin(self): + mkelv = special.kelvin(2) + assert_allclose(mkelv, (special.ber(2) + special.bei(2)*1j, + special.ker(2) + special.kei(2)*1j, + special.berp(2) + special.beip(2)*1j, + special.kerp(2) + special.keip(2)*1j), + atol=1.5e-8, rtol=0) + + def test_kei(self): + mkei = special.kei(2) + assert_allclose(mkei, -0.20240006776470432, atol=1.5e-5, rtol=0) + + def test_keip(self): + mkeip = special.keip(2) + assert_allclose(mkeip, 0.21980790991960536, atol=1.5e-5, rtol=0) + + def test_ker(self): + mker = special.ker(2) + assert_allclose(mker, -0.041664513991509472, atol=1.5e-5, rtol=0) + + def test_kerp(self): + mkerp = special.kerp(2) + assert_allclose(mkerp, -0.10660096588105264, atol=1.5e-5, rtol=0) + + def test_kei_zeros(self): + kei = special.kei_zeros(5) + assert_allclose(kei, array([3.91467, + 8.34422, + 12.78256, + 17.22314, + 21.66464]), + atol=1.5e-4, rtol=0) + + def test_keip_zeros(self): + keip = special.keip_zeros(5) + assert_allclose(keip, array([4.93181, + 9.40405, + 13.85827, + 18.30717, + 22.75379]), + atol=1.5e-4, rtol=0) + + # numbers come from 9.9 of A&S pg. 381 + def test_kelvin_zeros(self): + tmp = special.kelvin_zeros(5) + berz, beiz, kerz, keiz, berpz, beipz, kerpz, keipz = tmp + assert_allclose(berz, array([2.84892, + 7.23883, + 11.67396, + 16.11356, + 20.55463]), + atol=1.5e-4, rtol=0) + assert_allclose(beiz, array([5.02622, + 9.45541, + 13.89349, + 18.33398, + 22.77544]), + atol=1.5e-4, rtol=0) + assert_allclose(kerz, array([1.71854, + 6.12728, + 10.56294, + 15.00269, + 19.44382]), + atol=1.5e-4, rtol=0) + assert_allclose(keiz, array([3.91467, + 8.34422, + 12.78256, + 17.22314, + 21.66464]), + atol=1.5e-4, rtol=0) + assert_allclose(berpz, array([6.03871, + 10.51364, + 14.96844, + 19.41758, + 23.86430]), + atol=1.5e-4, rtol=0) + assert_allclose(beipz, array([3.77267, + # table from 1927 had 3.77320 + # but this is more accurate + 8.28099, + 12.74215, + 17.19343, + 21.64114]), + atol=1.5e-4, rtol=0) + assert_allclose(kerpz, array([2.66584, + 7.17212, + 11.63218, + 16.08312, + 20.53068]), + atol=1.5e-4, rtol=0) + assert_allclose(keipz, array([4.93181, + 9.40405, + 13.85827, + 18.30717, + 22.75379]), + atol=1.5e-4, rtol=0) + + def test_ker_zeros(self): + ker = special.ker_zeros(5) + assert_allclose(ker, array([1.71854, + 6.12728, + 10.56294, + 15.00269, + 19.44381]), + atol=1.5e-4, rtol=0) + + def test_kerp_zeros(self): + kerp = special.kerp_zeros(5) + assert_allclose(kerp, array([2.66584, + 7.17212, + 11.63218, + 16.08312, + 20.53068]), + atol=1.5e-4, rtol=0) + + +class TestBernoulli: + def test_bernoulli(self): + brn = special.bernoulli(5) + assert_allclose(brn, array([1.0000, + -0.5000, + 0.1667, + 0.0000, + -0.0333, + 0.0000]), + atol=1.5e-4, rtol=0) + + +class TestBeta: + """ + Test beta and betaln. + """ + + def test_beta(self): + assert_equal(special.beta(1, 1), 1.0) + assert_allclose(special.beta(-100.3, 1e-200), special.gamma(1e-200)) + assert_allclose(special.beta(0.0342, 171), 24.070498359873497, + rtol=1e-13, atol=0) + + bet = special.beta(2, 4) + betg = (special.gamma(2)*special.gamma(4))/special.gamma(6) + assert_allclose(bet, betg, rtol=1e-13) + + def test_beta_inf(self): + assert_(np.isinf(special.beta(-1, 2))) + + def test_betaln(self): + assert_equal(special.betaln(1, 1), 0.0) + assert_allclose(special.betaln(-100.3, 1e-200), + special.gammaln(1e-200)) + assert_allclose(special.betaln(0.0342, 170), 3.1811881124242447, + rtol=1e-14, atol=0) + + betln = special.betaln(2, 4) + bet = log(abs(special.beta(2, 4))) + assert_allclose(betln, bet, rtol=1e-13) + + +class TestBetaInc: + """ + Tests for betainc, betaincinv, betaincc, betainccinv. + """ + + def test_a1_b1(self): + # betainc(1, 1, x) is x. + x = np.array([0, 0.25, 1]) + assert_equal(special.betainc(1, 1, x), x) + assert_equal(special.betaincinv(1, 1, x), x) + assert_equal(special.betaincc(1, 1, x), 1 - x) + assert_equal(special.betainccinv(1, 1, x), 1 - x) + + # Nontrivial expected values computed with mpmath: + # from mpmath import mp + # mp.dps = 100 + # p = mp.betainc(a, b, 0, x, regularized=True) + # + # or, e.g., + # + # p = 0.25 + # a, b = 0.0342, 171 + # x = mp.findroot( + # lambda t: mp.betainc(a, b, 0, t, regularized=True) - p, + # (8e-21, 9e-21), + # solver='anderson', + # ) + # + @pytest.mark.parametrize( + 'a, b, x, p', + [(2, 4, 0.3138101704556974, 0.5), + (0.0342, 171.0, 1e-10, 0.552699169018070910641), + # gh-3761: + (0.0342, 171, 8.42313169354797e-21, 0.25), + # gh-4244: + (0.0002742794749792665, 289206.03125, 1.639984034231756e-56, + 0.9688708782196045), + # gh-12796: + (4, 99997, 0.0001947841578892121, 0.999995)]) + def test_betainc_and_inverses(self, a, b, x, p): + p1 = special.betainc(a, b, x) + assert_allclose(p1, p, rtol=1e-15) + x1 = special.betaincinv(a, b, p) + assert_allclose(x1, x, rtol=5e-13) + a1 = special.btdtria(p, b, x) + assert_allclose(a1, a, rtol=1e-13) + b1 = special.btdtrib(a, p, x) + assert_allclose(b1, b, rtol=1e-13) + + # Expected values computed with mpmath: + # from mpmath import mp + # mp.dps = 100 + # p = mp.betainc(a, b, x, 1, regularized=True) + @pytest.mark.parametrize('a, b, x, p', + [(2.5, 3.0, 0.25, 0.833251953125), + (7.5, 13.25, 0.375, 0.43298734645560368593), + (0.125, 7.5, 0.425, 0.0006688257851314237), + (0.125, 18.0, 1e-6, 0.72982359145096327654), + (0.125, 18.0, 0.996, 7.2745875538380150586e-46), + (0.125, 24.0, 0.75, 3.70853404816862016966e-17), + (16.0, 0.75, 0.99999999975, + 5.4408759277418629909e-07), + # gh-4677 (numbers from stackoverflow question): + (0.4211959643503401, 16939.046996018118, + 0.000815296167195521, 1e-7)]) + def test_betaincc_betainccinv(self, a, b, x, p): + p1 = special.betaincc(a, b, x) + assert_allclose(p1, p, rtol=5e-15) + x1 = special.betainccinv(a, b, p) + assert_allclose(x1, x, rtol=8e-15) + + @pytest.mark.parametrize( + 'a, b, y, ref', + [(14.208308325339239, 14.208308325339239, 7.703145458496392e-307, + 8.566004561846704e-23), + (14.0, 14.5, 1e-280, 2.9343915006642424e-21), + (3.5, 15.0, 4e-95, 1.3290751429289227e-28), + (10.0, 1.25, 2e-234, 3.982659092143654e-24), + (4.0, 99997.0, 5e-88, 3.309800566862242e-27)] + ) + def test_betaincinv_tiny_y(self, a, b, y, ref): + # Test with extremely small y values. This test includes + # a regression test for an issue in the boost code; + # see https://github.com/boostorg/math/issues/961 + # + # The reference values were computed with mpmath. For example, + # + # from mpmath import mp + # mp.dps = 1000 + # a = 14.208308325339239 + # p = 7.703145458496392e-307 + # x = mp.findroot(lambda t: mp.betainc(a, a, 0, t, + # regularized=True) - p, + # x0=8.566e-23) + # print(float(x)) + # + x = special.betaincinv(a, b, y) + assert_allclose(x, ref, rtol=1e-14) + + @pytest.mark.parametrize('func', [special.betainc, special.betaincinv, + special.btdtria, special.btdtrib, + special.betaincc, special.betainccinv]) + @pytest.mark.parametrize('args', [(-1.0, 2, 0.5), (1.5, -2.0, 0.5), + (1.5, 2.0, -0.3), (1.5, 2.0, 1.1)]) + def test_betainc_domain_errors(self, func, args): + with special.errstate(domain='raise'): + with pytest.raises(special.SpecialFunctionError, match='domain'): + func(*args) + + @pytest.mark.parametrize( + "args,expected", + [ + ((0.0, 0.0, 0.0), np.nan), + ((0.0, 0.0, 0.5), np.nan), + ((0.0, 0.0, 1.0), np.nan), + ((np.inf, np.inf, 0.0), np.nan), + ((np.inf, np.inf, 0.5), np.nan), + ((np.inf, np.inf, 1.0), np.nan), + ((0.0, 1.0, 0.0), 0.0), + ((0.0, 1.0, 0.5), 1.0), + ((0.0, 1.0, 1.0), 1.0), + ((1.0, 0.0, 0.0), 0.0), + ((1.0, 0.0, 0.5), 0.0), + ((1.0, 0.0, 1.0), 1.0), + ((0.0, np.inf, 0.0), 0.0), + ((0.0, np.inf, 0.5), 1.0), + ((0.0, np.inf, 1.0), 1.0), + ((np.inf, 0.0, 0.0), 0.0), + ((np.inf, 0.0, 0.5), 0.0), + ((np.inf, 0.0, 1.0), 1.0), + ((1.0, np.inf, 0.0), 0.0), + ((1.0, np.inf, 0.5), 1.0), + ((1.0, np.inf, 1.0), 1.0), + ((np.inf, 1.0, 0.0), 0.0), + ((np.inf, 1.0, 0.5), 0.0), + ((np.inf, 1.0, 1.0), 1.0), + ] + ) + def test_betainc_edge_cases(self, args, expected): + observed = special.betainc(*args) + assert_equal(observed, expected) + + @pytest.mark.parametrize( + "args,expected", + [ + ((0.0, 0.0, 0.0), np.nan), + ((0.0, 0.0, 0.5), np.nan), + ((0.0, 0.0, 1.0), np.nan), + ((np.inf, np.inf, 0.0), np.nan), + ((np.inf, np.inf, 0.5), np.nan), + ((np.inf, np.inf, 1.0), np.nan), + ((0.0, 1.0, 0.0), 1.0), + ((0.0, 1.0, 0.5), 0.0), + ((0.0, 1.0, 1.0), 0.0), + ((1.0, 0.0, 0.0), 1.0), + ((1.0, 0.0, 0.5), 1.0), + ((1.0, 0.0, 1.0), 0.0), + ((0.0, np.inf, 0.0), 1.0), + ((0.0, np.inf, 0.5), 0.0), + ((0.0, np.inf, 1.0), 0.0), + ((np.inf, 0.0, 0.0), 1.0), + ((np.inf, 0.0, 0.5), 1.0), + ((np.inf, 0.0, 1.0), 0.0), + ((1.0, np.inf, 0.0), 1.0), + ((1.0, np.inf, 0.5), 0.0), + ((1.0, np.inf, 1.0), 0.0), + ((np.inf, 1.0, 0.0), 1.0), + ((np.inf, 1.0, 0.5), 1.0), + ((np.inf, 1.0, 1.0), 0.0), + ] + ) + def test_betaincc_edge_cases(self, args, expected): + observed = special.betaincc(*args) + assert_equal(observed, expected) + + @pytest.mark.parametrize('dtype', [np.float32, np.float64]) + def test_gh21426(self, dtype): + # Test for gh-21426: betaincinv must not return NaN + a = np.array([5.], dtype=dtype) + x = np.array([0.5], dtype=dtype) + result = special.betaincinv(a, a, x) + assert_allclose(result, x, rtol=10 * np.finfo(dtype).eps) + + @pytest.mark.parametrize("dtype, rtol", + [(np.float32, 1e-4), + (np.float64, 1e-15)]) + @pytest.mark.parametrize("a, b, x, reference", + [(1e-20, 1e-21, 0.5, 0.0909090909090909), + (1e-15, 1e-16, 0.5, 0.09090909090909091)]) + def test_gh22682(self, a, b, x, reference, dtype, rtol): + # gh-22682: betainc returned incorrect results for tiny + # single precision inputs. test that this is resolved + a = np.array(a, dtype=dtype) + b = np.array(b, dtype=dtype) + x = np.array(x, dtype=dtype) + res = special.betainc(a, b, x) + assert_allclose(res, reference, rtol=rtol) + + +class TestCombinatorics: + def test_comb(self): + assert_allclose(special.comb([10, 10], [3, 4]), [120., 210.]) + assert_allclose(special.comb(10, 3), 120.) + assert_equal(special.comb(10, 3, exact=True), 120) + assert_equal(special.comb(10, 3, exact=True, repetition=True), 220) + + assert_allclose([special.comb(20, k, exact=True) for k in range(21)], + special.comb(20, list(range(21))), atol=1e-15) + + ii = np.iinfo(int).max + 1 + assert_equal(special.comb(ii, ii-1, exact=True), ii) + + expected = 100891344545564193334812497256 + assert special.comb(100, 50, exact=True) == expected + + def test_comb_with_np_int64(self): + n = 70 + k = 30 + np_n = np.int64(n) + np_k = np.int64(k) + res_np = special.comb(np_n, np_k, exact=True) + res_py = special.comb(n, k, exact=True) + assert res_np == res_py + + def test_comb_zeros(self): + assert_equal(special.comb(2, 3, exact=True), 0) + assert_equal(special.comb(-1, 3, exact=True), 0) + assert_equal(special.comb(2, -1, exact=True), 0) + assert_equal(special.comb(2, -1, exact=False), 0) + assert_allclose(special.comb([2, -1, 2, 10], [3, 3, -1, 3]), [0., 0., 0., 120.]) + + def test_comb_exact_non_int_error(self): + msg = "`exact=True`" + with pytest.raises(ValueError, match=msg): + special.comb(3.4, 4, exact=True) + with pytest.raises(ValueError, match=msg): + special.comb(3, 4.4, exact=True) + + @pytest.mark.parametrize('N', [0, 5, 10]) + @pytest.mark.parametrize('exact', [True, False]) + def test_comb_repetition_k_zero(self, N, exact): + # Regression test for gh-23867 + # C(n, 0) should always be 1 for n >= 0, regardless of repetition + actual = special.comb(N, 0, exact=exact, repetition=True) + assert actual == 1 + assert type(actual) is int if exact else np.float64 + + def test_comb_repetition_k_zero_array(self): + # Test array-like input with exact=False for gh-23867 + N = np.array([0, 5, 10]) + result = special.comb(N, 0, exact=False, repetition=True) + expected = np.array([1.0, 1.0, 1.0]) + assert_equal(result, expected) + + def test_perm(self): + assert_allclose(special.perm([10, 10], [3, 4]), [720., 5040.]) + assert_allclose(special.perm(10, 3), 720., atol=1.5e-7, rtol=0) + assert_equal(special.perm(10, 3, exact=True), 720) + + def test_perm_zeros(self): + assert_equal(special.perm(2, 3, exact=True), 0) + assert_equal(special.perm(-1, 3, exact=True), 0) + assert_equal(special.perm(2, -1, exact=True), 0) + assert_equal(special.perm(2, -1, exact=False), 0) + assert_allclose(special.perm([2, -1, 2, 10], [3, 3, -1, 3]), [0., 0., 0., 720.]) + + def test_perm_iv(self): + # currently `exact=True` only support scalars + with pytest.raises(ValueError, match="scalar integers"): + special.perm([1, 2], [4, 5], exact=True) + + with pytest.raises(ValueError, match="Non-integer"): + special.perm(4.6, 6, exact=True) + with pytest.raises(ValueError, match="Non-integer"): + special.perm(-4.6, 3, exact=True) + with pytest.raises(ValueError, match="Non-integer"): + special.perm(4, -3.9, exact=True) + with pytest.raises(ValueError, match="Non-integer"): + special.perm(6.0, 4.6, exact=True) + + +class TestTrigonometric: + def test_cbrt(self): + cb = special.cbrt(27) + cbrl = 27**(1.0/3.0) + assert_allclose(cb, cbrl, atol=1.5e-7, rtol=0) + + def test_cbrtmore(self): + cb1 = special.cbrt(27.9) + cbrl1 = 27.9**(1.0/3.0) + assert_allclose(cb1, cbrl1, atol=1.5e-8, rtol=0) + + def test_cosdg(self): + cdg = special.cosdg(90) + cdgrl = cos(pi/2.0) + assert_allclose(cdg, cdgrl, atol=1.5e-8, rtol=0) + + def test_cosdgmore(self): + cdgm = special.cosdg(30) + cdgmrl = cos(pi/6.0) + assert_allclose(cdgm, cdgmrl, atol=1.5e-8, rtol=0) + + def test_cosm1(self): + cs = (special.cosm1(0),special.cosm1(.3),special.cosm1(pi/10)) + csrl = (cos(0)-1,cos(.3)-1,cos(pi/10)-1) + assert_allclose(cs, csrl, atol=1.5e-8, rtol=0) + + def test_cotdg(self): + ct = special.cotdg(30) + ctrl = tan(pi/6.0)**(-1) + assert_allclose(ct, ctrl, atol=1.5e-8, rtol=0) + + def test_cotdgmore(self): + ct1 = special.cotdg(45) + ctrl1 = tan(pi/4.0)**(-1) + assert_allclose(ct1, ctrl1, atol=1.5e-8, rtol=0) + + def test_specialpoints(self): + assert_allclose(special.cotdg(45), 1.0, atol=1.5e-14, rtol=0) + assert_allclose(special.cotdg(-45), -1.0, atol=1.5e-14, rtol=0) + assert_allclose(special.cotdg(90), 0.0, atol=1.5e-14, rtol=0) + assert_allclose(special.cotdg(-90), 0.0, atol=1.5e-14, rtol=0) + assert_allclose(special.cotdg(135), -1.0, atol=1.5e-14, rtol=0) + assert_allclose(special.cotdg(-135), 1.0, atol=1.5e-14, rtol=0) + assert_allclose(special.cotdg(225), 1.0, atol=1.5e-14, rtol=0) + assert_allclose(special.cotdg(-225), -1.0, atol=1.5e-14, rtol=0) + assert_allclose(special.cotdg(270), 0.0, atol=1.5e-14, rtol=0) + assert_allclose(special.cotdg(-270), 0.0, atol=1.5e-14, rtol=0) + assert_allclose(special.cotdg(315), -1.0, atol=1.5e-14, rtol=0) + assert_allclose(special.cotdg(-315), 1.0, atol=1.5e-14, rtol=0) + assert_allclose(special.cotdg(765), 1.0, atol=1.5e-14, rtol=0) + + def test_sinc(self): + # the sinc implementation and more extensive sinc tests are in numpy + assert_array_equal(special.sinc([0]), 1) + assert_equal(special.sinc(0.0), 1.0) + + def test_sindg(self): + sn = special.sindg(90) + assert_equal(sn,1.0) + + def test_sindgmore(self): + snm = special.sindg(30) + snmrl = sin(pi/6.0) + assert_allclose(snm, snmrl, atol=1.5e-8, rtol=0) + snm1 = special.sindg(45) + snmrl1 = sin(pi/4.0) + assert_allclose(snm1, snmrl1, atol=1.5e-8, rtol=0) + + +class TestTandg: + + def test_tandg(self): + tn = special.tandg(30) + tnrl = tan(pi/6.0) + assert_allclose(tn, tnrl, atol=1.5e-8, rtol=0) + + def test_tandgmore(self): + tnm = special.tandg(45) + tnmrl = tan(pi/4.0) + assert_allclose(tnm, tnmrl, atol=1.5e-8, rtol=0) + tnm1 = special.tandg(60) + tnmrl1 = tan(pi/3.0) + assert_allclose(tnm1, tnmrl1, atol=1.5e-8, rtol=0) + + def test_specialpoints(self): + assert_allclose(special.tandg(0), 0.0, atol=1.5e-14, rtol=0) + assert_allclose(special.tandg(45), 1.0, atol=1.5e-14, rtol=0) + assert_allclose(special.tandg(-45), -1.0, atol=1.5e-14, rtol=0) + assert_allclose(special.tandg(135), -1.0, atol=1.5e-14, rtol=0) + assert_allclose(special.tandg(-135), 1.0, atol=1.5e-14, rtol=0) + assert_allclose(special.tandg(180), 0.0, atol=1.5e-14, rtol=0) + assert_allclose(special.tandg(-180), 0.0, atol=1.5e-14, rtol=0) + assert_allclose(special.tandg(225), 1.0, atol=1.5e-14, rtol=0) + assert_allclose(special.tandg(-225), -1.0, atol=1.5e-14, rtol=0) + assert_allclose(special.tandg(315), -1.0, atol=1.5e-14, rtol=0) + assert_allclose(special.tandg(-315), 1.0, atol=1.5e-14, rtol=0) + + +class TestEllip: + def test_ellipj_nan(self): + """Regression test for #912.""" + special.ellipj(0.5, np.nan) + + def test_ellipj(self): + el = special.ellipj(0.2,0) + rel = [sin(0.2),cos(0.2),1.0,0.20] + assert_allclose(el, rel, atol=1.5e-13, rtol=0) + + def test_ellipk(self): + elk = special.ellipk(.2) + assert_allclose(elk, 1.659623598610528, atol=1.5e-11, rtol=0) + + assert_equal(special.ellipkm1(0.0), np.inf) + assert_equal(special.ellipkm1(1.0), pi/2) + assert_equal(special.ellipkm1(np.inf), 0.0) + assert_equal(special.ellipkm1(np.nan), np.nan) + assert_equal(special.ellipkm1(-1), np.nan) + assert_allclose(special.ellipk(-10), 0.7908718902387385) + + def test_ellipkinc(self): + elkinc = special.ellipkinc(pi/2,.2) + elk = special.ellipk(0.2) + assert_allclose(elkinc, elk, atol=1.5e-15, rtol=0) + alpha = 20*pi/180 + phi = 45*pi/180 + m = sin(alpha)**2 + elkinc = special.ellipkinc(phi,m) + assert_allclose(elkinc, 0.79398143, atol=1.5e-8, rtol=0) + # From pg. 614 of A & S + + assert_equal(special.ellipkinc(pi/2, 0.0), pi/2) + assert_equal(special.ellipkinc(pi/2, 1.0), np.inf) + assert_equal(special.ellipkinc(pi/2, -np.inf), 0.0) + assert_equal(special.ellipkinc(pi/2, np.nan), np.nan) + assert_equal(special.ellipkinc(pi/2, 2), np.nan) + assert_equal(special.ellipkinc(0, 0.5), 0.0) + assert_equal(special.ellipkinc(np.inf, 0.5), np.inf) + assert_equal(special.ellipkinc(-np.inf, 0.5), -np.inf) + assert_equal(special.ellipkinc(np.inf, np.inf), np.nan) + assert_equal(special.ellipkinc(np.inf, -np.inf), np.nan) + assert_equal(special.ellipkinc(-np.inf, -np.inf), np.nan) + assert_equal(special.ellipkinc(-np.inf, np.inf), np.nan) + assert_equal(special.ellipkinc(np.nan, 0.5), np.nan) + assert_equal(special.ellipkinc(np.nan, np.nan), np.nan) + + assert_allclose(special.ellipkinc(0.38974112035318718, 1), 0.4, rtol=1e-14) + assert_allclose(special.ellipkinc(1.5707, -10), 0.79084284661724946) + + def test_ellipkinc_2(self): + # Regression test for gh-3550 + # ellipkinc(phi, mbad) was NaN and mvals[2:6] were twice the correct value + mbad = 0.68359375000000011 + phi = 0.9272952180016123 + m = np.nextafter(mbad, 0) + mvals = [] + for j in range(10): + mvals.append(m) + m = np.nextafter(m, 1) + f = special.ellipkinc(phi, mvals) + assert_array_almost_equal_nulp(f, np.full_like(f, 1.0259330100195334), 1) + # this bug also appears at phi + n * pi for at least small n + f1 = special.ellipkinc(phi + pi, mvals) + assert_array_almost_equal_nulp(f1, np.full_like(f1, 5.1296650500976675), 2) + + def test_ellipkinc_singular(self): + # ellipkinc(phi, 1) has closed form and is finite only for phi in (-pi/2, pi/2) + xlog = np.logspace(-300, -17, 25) + xlin = np.linspace(1e-17, 0.1, 25) + xlin2 = np.linspace(0.1, pi/2, 25, endpoint=False) + + assert_allclose(special.ellipkinc(xlog, 1), np.arcsinh(np.tan(xlog)), + rtol=1e14) + assert_allclose(special.ellipkinc(xlin, 1), np.arcsinh(np.tan(xlin)), + rtol=1e14) + assert_allclose(special.ellipkinc(xlin2, 1), np.arcsinh(np.tan(xlin2)), + rtol=1e14) + assert_equal(special.ellipkinc(np.pi/2, 1), np.inf) + assert_allclose(special.ellipkinc(-xlog, 1), np.arcsinh(np.tan(-xlog)), + rtol=1e14) + assert_allclose(special.ellipkinc(-xlin, 1), np.arcsinh(np.tan(-xlin)), + rtol=1e14) + assert_allclose(special.ellipkinc(-xlin2, 1), np.arcsinh(np.tan(-xlin2)), + rtol=1e14) + assert_equal(special.ellipkinc(-np.pi/2, 1), np.inf) + + def test_ellipe(self): + ele = special.ellipe(.2) + assert_allclose(ele, 1.4890350580958529, atol=1.5e-8, rtol=0) + + assert_equal(special.ellipe(0.0), pi/2) + assert_equal(special.ellipe(1.0), 1.0) + assert_equal(special.ellipe(-np.inf), np.inf) + assert_equal(special.ellipe(np.nan), np.nan) + assert_equal(special.ellipe(2), np.nan) + assert_allclose(special.ellipe(-10), 3.6391380384177689) + + def test_ellipeinc(self): + eleinc = special.ellipeinc(pi/2,.2) + ele = special.ellipe(0.2) + assert_allclose(eleinc, ele, atol=1.5e-14, rtol=0) + # pg 617 of A & S + alpha, phi = 52*pi/180,35*pi/180 + m = sin(alpha)**2 + eleinc = special.ellipeinc(phi,m) + assert_allclose(eleinc, 0.58823065, atol=1.5e-8, rtol=0) + + assert_equal(special.ellipeinc(pi/2, 0.0), pi/2) + assert_equal(special.ellipeinc(pi/2, 1.0), 1.0) + assert_equal(special.ellipeinc(pi/2, -np.inf), np.inf) + assert_equal(special.ellipeinc(pi/2, np.nan), np.nan) + assert_equal(special.ellipeinc(pi/2, 2), np.nan) + assert_equal(special.ellipeinc(0, 0.5), 0.0) + assert_equal(special.ellipeinc(np.inf, 0.5), np.inf) + assert_equal(special.ellipeinc(-np.inf, 0.5), -np.inf) + assert_equal(special.ellipeinc(np.inf, -np.inf), np.inf) + assert_equal(special.ellipeinc(-np.inf, -np.inf), -np.inf) + assert_equal(special.ellipeinc(np.inf, np.inf), np.nan) + assert_equal(special.ellipeinc(-np.inf, np.inf), np.nan) + assert_equal(special.ellipeinc(np.nan, 0.5), np.nan) + assert_equal(special.ellipeinc(np.nan, np.nan), np.nan) + assert_allclose(special.ellipeinc(1.5707, -10), 3.6388185585822876) + + def test_ellipeinc_2(self): + # Regression test for gh-3550 + # ellipeinc(phi, mbad) was NaN and mvals[2:6] were twice the correct value + mbad = 0.68359375000000011 + phi = 0.9272952180016123 + m = np.nextafter(mbad, 0) + mvals = [] + for j in range(10): + mvals.append(m) + m = np.nextafter(m, 1) + f = special.ellipeinc(phi, mvals) + assert_array_almost_equal_nulp(f, np.full_like(f, 0.84442884574781019), 2) + # this bug also appears at phi + n * pi for at least small n + f1 = special.ellipeinc(phi + pi, mvals) + assert_array_almost_equal_nulp(f1, np.full_like(f1, 3.3471442287390509), 4) + + +class TestEllipCarlson: + """Test for Carlson elliptic integrals ellipr[cdfgj]. + The special values used in these tests can be found in Sec. 3 of Carlson + (1994), https://arxiv.org/abs/math/9409227 + """ + def test_elliprc(self): + assert_allclose(elliprc(1, 1), 1) + assert elliprc(1, inf) == 0.0 + assert isnan(elliprc(1, 0)) + assert elliprc(1, complex(1, inf)) == 0.0 + args = array([[0.0, 0.25], + [2.25, 2.0], + [0.0, 1.0j], + [-1.0j, 1.0j], + [0.25, -2.0], + [1.0j, -1.0]]) + expected_results = array([np.pi, + np.log(2.0), + 1.1107207345396 * (1.0-1.0j), + 1.2260849569072-0.34471136988768j, + np.log(2.0) / 3.0, + 0.77778596920447+0.19832484993429j]) + for i, arr in enumerate(args): + assert_allclose(elliprc(*arr), expected_results[i]) + + def test_elliprd(self): + assert_allclose(elliprd(1, 1, 1), 1) + assert_allclose(elliprd(0, 2, 1) / 3.0, 0.59907011736779610371) + assert elliprd(1, 1, inf) == 0.0 + assert np.isinf(elliprd(1, 1, 0)) + assert np.isinf(elliprd(1, 1, complex(0, 0))) + assert np.isinf(elliprd(0, 1, complex(0, 0))) + assert isnan(elliprd(1, 1, -np.finfo(np.float64).tiny / 2.0)) + assert isnan(elliprd(1, 1, complex(-1, 0))) + args = array([[0.0, 2.0, 1.0], + [2.0, 3.0, 4.0], + [1.0j, -1.0j, 2.0], + [0.0, 1.0j, -1.0j], + [0.0, -1.0+1.0j, 1.0j], + [-2.0-1.0j, -1.0j, -1.0+1.0j]]) + expected_results = array([1.7972103521034, + 0.16510527294261, + 0.65933854154220, + 1.2708196271910+2.7811120159521j, + -1.8577235439239-0.96193450888839j, + 1.8249027393704-1.2218475784827j]) + for i, arr in enumerate(args): + assert_allclose(elliprd(*arr), expected_results[i]) + + def test_elliprf(self): + assert_allclose(elliprf(1, 1, 1), 1) + assert_allclose(elliprf(0, 1, 2), 1.31102877714605990523) + assert elliprf(1, inf, 1) == 0.0 + assert np.isinf(elliprf(0, 1, 0)) + assert isnan(elliprf(1, 1, -1)) + assert elliprf(complex(inf), 0, 1) == 0.0 + assert isnan(elliprf(1, 1, complex(-inf, 1))) + args = array([[1.0, 2.0, 0.0], + [1.0j, -1.0j, 0.0], + [0.5, 1.0, 0.0], + [-1.0+1.0j, 1.0j, 0.0], + [2.0, 3.0, 4.0], + [1.0j, -1.0j, 2.0], + [-1.0+1.0j, 1.0j, 1.0-1.0j]]) + expected_results = array([1.3110287771461, + 1.8540746773014, + 1.8540746773014, + 0.79612586584234-1.2138566698365j, + 0.58408284167715, + 1.0441445654064, + 0.93912050218619-0.53296252018635j]) + for i, arr in enumerate(args): + assert_allclose(elliprf(*arr), expected_results[i]) + + def test_elliprg(self): + assert_allclose(elliprg(1, 1, 1), 1) + assert_allclose(elliprg(0, 0, 1), 0.5) + assert_allclose(elliprg(0, 0, 0), 0) + assert np.isinf(elliprg(1, inf, 1)) + assert np.isinf(elliprg(complex(inf), 1, 1)) + args = array([[0.0, 16.0, 16.0], + [2.0, 3.0, 4.0], + [0.0, 1.0j, -1.0j], + [-1.0+1.0j, 1.0j, 0.0], + [-1.0j, -1.0+1.0j, 1.0j], + [0.0, 0.0796, 4.0]]) + expected_results = array([np.pi, + 1.7255030280692, + 0.42360654239699, + 0.44660591677018+0.70768352357515j, + 0.36023392184473+0.40348623401722j, + 1.0284758090288]) + for i, arr in enumerate(args): + assert_allclose(elliprg(*arr), expected_results[i]) + + def test_elliprj(self): + assert_allclose(elliprj(1, 1, 1, 1), 1) + assert elliprj(1, 1, inf, 1) == 0.0 + assert isnan(elliprj(1, 0, 0, 0)) + assert isnan(elliprj(-1, 1, 1, 1)) + assert elliprj(1, 1, 1, inf) == 0.0 + args = array([[0.0, 1.0, 2.0, 3.0], + [2.0, 3.0, 4.0, 5.0], + [2.0, 3.0, 4.0, -1.0+1.0j], + [1.0j, -1.0j, 0.0, 2.0], + [-1.0+1.0j, -1.0-1.0j, 1.0, 2.0], + [1.0j, -1.0j, 0.0, 1.0-1.0j], + [-1.0+1.0j, -1.0-1.0j, 1.0, -3.0+1.0j], + [2.0, 3.0, 4.0, -0.5], # Cauchy principal value + [2.0, 3.0, 4.0, -5.0]]) # Cauchy principal value + expected_results = array([0.77688623778582, + 0.14297579667157, + 0.13613945827771-0.38207561624427j, + 1.6490011662711, + 0.94148358841220, + 1.8260115229009+1.2290661908643j, + -0.61127970812028-1.0684038390007j, + 0.24723819703052, # Cauchy principal value + -0.12711230042964]) # Caucny principal value + for i, arr in enumerate(args): + assert_allclose(elliprj(*arr), expected_results[i]) + + @pytest.mark.xfail(reason="Insufficient accuracy on 32-bit") + def test_elliprj_hard(self): + assert_allclose(elliprj(6.483625725195452e-08, + 1.1649136528196886e-27, + 3.6767340167168e+13, + 0.493704617023468), + 8.63426920644241857617477551054e-6, + rtol=5e-15, atol=1e-20) + assert_allclose(elliprj(14.375105857849121, + 9.993988969725365e-11, + 1.72844262269944e-26, + 5.898871222598245e-06), + 829774.1424801627252574054378691828, + rtol=5e-15, atol=1e-20) + + +class TestEllipLegendreCarlsonIdentities: + """Test identities expressing the Legendre elliptic integrals in terms + of Carlson's symmetric integrals. These identities can be found + in the DLMF https://dlmf.nist.gov/19.25#i . + """ + + def setup_class(self): + self.m_n1_1 = np.arange(-1., 1., 0.01) + # For double, this is -(2**1024) + self.max_neg = finfo(double).min + # Lots of very negative numbers + self.very_neg_m = -1. * 2.**arange(-1 + + np.log2(-self.max_neg), 0., + -1.) + self.ms_up_to_1 = np.concatenate(([self.max_neg], + self.very_neg_m, + self.m_n1_1)) + + def test_k(self): + """Test identity: + K(m) = R_F(0, 1-m, 1) + """ + m = self.ms_up_to_1 + assert_allclose(ellipk(m), elliprf(0., 1.-m, 1.)) + + def test_km1(self): + """Test identity: + K(m) = R_F(0, 1-m, 1) + But with the ellipkm1 function + """ + # For double, this is 2**-1022 + tiny = finfo(double).tiny + # All these small powers of 2, up to 2**-1 + m1 = tiny * 2.**arange(0., -np.log2(tiny)) + assert_allclose(ellipkm1(m1), elliprf(0., m1, 1.)) + + def test_e(self): + """Test identity: + E(m) = 2*R_G(0, 1-k^2, 1) + """ + m = self.ms_up_to_1 + assert_allclose(ellipe(m), 2.*elliprg(0., 1.-m, 1.)) + + +class TestErf: + + def test_erf(self): + er = special.erf(.25) + assert_allclose(er, 0.2763263902, atol=1.5e-8, rtol=0) + + def test_erf_zeros(self): + erz = special.erf_zeros(5) + erzr = array([1.45061616+1.88094300j, + 2.24465928+2.61657514j, + 2.83974105+3.17562810j, + 3.33546074+3.64617438j, + 3.76900557+4.06069723j]) + assert_allclose(erz, erzr, atol=1.5e-4, rtol=0) + + def _check_variant_func(self, func, other_func, rtol, atol=0): + rng = np.random.RandomState(1234) + n = 10000 + x = rng.pareto(0.02, n) * (2*rng.randint(0, 2, n) - 1) + y = rng.pareto(0.02, n) * (2*rng.randint(0, 2, n) - 1) + z = x + 1j*y + + with np.errstate(all='ignore'): + w = other_func(z) + w_real = other_func(x).real + + mask = np.isfinite(w) + w = w[mask] + z = z[mask] + + mask = np.isfinite(w_real) + w_real = w_real[mask] + x = x[mask] + + # test both real and complex variants + assert_func_equal(func, w, z, rtol=rtol, atol=atol) + assert_func_equal(func, w_real, x, rtol=rtol, atol=atol) + + def test_erfc_consistent(self): + self._check_variant_func( + cephes.erfc, + lambda z: 1 - cephes.erf(z), + rtol=1e-12, + atol=1e-14 # <- the test function loses precision + ) + + def test_erfcx_consistent(self): + self._check_variant_func( + cephes.erfcx, + lambda z: np.exp(z*z) * cephes.erfc(z), + rtol=1e-12 + ) + + def test_erfi_consistent(self): + self._check_variant_func( + cephes.erfi, + lambda z: -1j * cephes.erf(1j*z), + rtol=1e-12 + ) + + def test_dawsn_consistent(self): + self._check_variant_func( + cephes.dawsn, + lambda z: sqrt(pi)/2 * np.exp(-z*z) * cephes.erfi(z), + rtol=1e-12 + ) + + def test_erf_nan_inf(self): + vals = [np.nan, -np.inf, np.inf] + expected = [np.nan, -1, 1] + assert_allclose(special.erf(vals), expected, rtol=1e-15) + + def test_erfc_nan_inf(self): + vals = [np.nan, -np.inf, np.inf] + expected = [np.nan, 2, 0] + assert_allclose(special.erfc(vals), expected, rtol=1e-15) + + def test_erfcx_nan_inf(self): + vals = [np.nan, -np.inf, np.inf] + expected = [np.nan, np.inf, 0] + assert_allclose(special.erfcx(vals), expected, rtol=1e-15) + + def test_erfi_nan_inf(self): + vals = [np.nan, -np.inf, np.inf] + expected = [np.nan, -np.inf, np.inf] + assert_allclose(special.erfi(vals), expected, rtol=1e-15) + + def test_dawsn_nan_inf(self): + vals = [np.nan, -np.inf, np.inf] + expected = [np.nan, -0.0, 0.0] + assert_allclose(special.dawsn(vals), expected, rtol=1e-15) + + def test_wofz_nan_inf(self): + vals = [np.nan, -np.inf, np.inf] + expected = [np.nan + np.nan * 1.j, 0.-0.j, 0.+0.j] + assert_allclose(special.wofz(vals), expected, rtol=1e-15) + + +class TestEuler: + def test_euler(self): + eu0 = special.euler(0) + eu1 = special.euler(1) + eu2 = special.euler(2) # just checking segfaults + assert_allclose(eu0, [1], rtol=1e-15) + assert_allclose(eu1, [1, 0], rtol=1e-15) + assert_allclose(eu2, [1, 0, -1], rtol=1e-15) + eu24 = special.euler(24) + mathworld = [1,1,5,61,1385,50521,2702765,199360981, + 19391512145,2404879675441, + 370371188237525,69348874393137901, + 15514534163557086905] + correct = zeros((25,),'d') + for k in range(0,13): + if (k % 2): + correct[2*k] = -float(mathworld[k]) + else: + correct[2*k] = float(mathworld[k]) + with np.errstate(all='ignore'): + err = nan_to_num((eu24-correct)/correct) + errmax = max(err) + assert_allclose(errmax, 0.0, atol=1.5e-14, rtol=0) + + +class TestExp: + def test_exp2(self): + ex = special.exp2(2) + exrl = 2**2 + assert_equal(ex,exrl) + + def test_exp2more(self): + exm = special.exp2(2.5) + exmrl = 2**(2.5) + assert_allclose(exm, exmrl, atol=1.5e-8, rtol=0) + + def test_exp10(self): + ex = special.exp10(2) + exrl = 10**2 + assert_allclose(ex, exrl, atol=1e-6, rtol=0) + + def test_exp10more(self): + exm = special.exp10(2.5) + exmrl = 10**(2.5) + assert_allclose(exm, exmrl, atol=1.5e-8, rtol=0) + + def test_expm1(self): + ex = (special.expm1(2), special.expm1(3), special.expm1(4)) + exrl = (exp(2) - 1, exp(3) - 1, exp(4) - 1) + assert_allclose(ex, exrl, atol=1.5e-8, rtol=0) + + def test_expm1more(self): + ex1 = (special.expm1(2), special.expm1(2.1), special.expm1(2.2)) + exrl1 = (exp(2) - 1, exp(2.1) - 1, exp(2.2) - 1) + assert_allclose(ex1, exrl1, atol=1.5e-8, rtol=0) + + +def assert_really_equal(x, y, rtol=None): + """ + Sharper assertion function that is stricter about matching types, not just values + + This is useful/necessary in some cases: + * dtypes for arrays that have the same _values_ (e.g. element 1.0 vs 1) + * distinguishing complex from real NaN + * result types for scalars + + We still want to be able to allow a relative tolerance for the values though. + The main logic comparison logic is handled by the xp_assert_* functions. + """ + def assert_func(x, y): + xp_assert_equal(x, y) if rtol is None else xp_assert_close(x, y, rtol=rtol) + + def assert_complex_nan(x): + assert np.isnan(x.real) and np.isnan(x.imag) + + assert type(x) is type(y), f"types not equal: {type(x)}, {type(y)}" + + # ensure we also compare the values _within_ an array appropriately, + # e.g. assert_equal does not distinguish different complex nans in arrays + if isinstance(x, np.ndarray): + # assert_equal does not compare (all) types, only values + assert x.dtype == y.dtype + # for empty arrays resp. to ensure shapes match + assert_func(x, y) + for elem_x, elem_y in zip(x.ravel(), y.ravel()): + assert_really_equal(elem_x, elem_y, rtol=rtol) + elif np.isnan(x) and np.isnan(y) and _is_subdtype(type(x), "c"): + assert_complex_nan(x) and assert_complex_nan(y) + # no need to consider complex infinities due to numpy/numpy#25493 + else: + assert_func(x, y) + + +class TestFactorialFunctions: + def factorialk_ref(self, n, k, exact, extend): + if exact: + return special.factorialk(n, k=k, exact=True) + # for details / explanation see factorialk-docstring + r = np.mod(n, k) if extend == "zero" else 1 + vals = np.power(k, (n - r)/k) * special.gamma(n/k + 1) * special.rgamma(r/k + 1) + # np.maximum is element-wise, which is what we want + return vals * np.maximum(r, 1) + + @pytest.mark.parametrize("exact,extend", + [(True, "zero"), (False, "zero"), (False, "complex")]) + def test_factorialx_scalar_return_type(self, exact, extend): + kw = {"exact": exact, "extend": extend} + assert np.isscalar(special.factorial(1, **kw)) + assert np.isscalar(special.factorial2(1, **kw)) + assert np.isscalar(special.factorialk(1, k=3, **kw)) + + @pytest.mark.parametrize("n", [-1, -2, -3]) + @pytest.mark.parametrize("exact", [True, False]) + def test_factorialx_negative_extend_zero(self, exact, n): + kw = {"exact": exact} + assert_equal(special.factorial(n, **kw), 0) + assert_equal(special.factorial2(n, **kw), 0) + assert_equal(special.factorialk(n, k=3, **kw), 0) + + @pytest.mark.parametrize("exact", [True, False]) + def test_factorialx_negative_extend_zero_array(self, exact): + kw = {"exact": exact} + rtol = 1e-15 + n = [-5, -4, 0, 1] + # Consistent output for n < 0 + expected = np.array([0, 0, 1, 1], dtype=native_int if exact else np.float64) + assert_really_equal(special.factorial(n, **kw), expected, rtol=rtol) + assert_really_equal(special.factorial2(n, **kw), expected, rtol=rtol) + assert_really_equal(special.factorialk(n, k=3, **kw), expected, rtol=rtol) + + @pytest.mark.parametrize("n", [-1.1, -2.2, -3.3]) + def test_factorialx_negative_extend_complex(self, n): + kw = {"extend": "complex"} + exp_1 = {-1.1: -10.686287021193184771, + -2.2: 4.8509571405220931958, + -3.3: -1.4471073942559181166} + exp_2 = {-1.1: 1.0725776858167496309, + -2.2: -3.9777171783768419874, + -3.3: -0.99588841846200555977} + exp_k = {-1.1: 0.73565345382163025659, + -2.2: 1.1749163167190809498, + -3.3: -2.4780584257450583713} + rtol = 3e-15 + assert_allclose(special.factorial(n, **kw), exp_1[n], rtol=rtol) + assert_allclose(special.factorial2(n, **kw), exp_2[n], rtol=rtol) + assert_allclose(special.factorialk(n, k=3, **kw), exp_k[n], rtol=rtol) + assert_allclose(special.factorial([n], **kw)[0], exp_1[n], rtol=rtol) + assert_allclose(special.factorial2([n], **kw)[0], exp_2[n], rtol=rtol) + assert_allclose(special.factorialk([n], k=3, **kw)[0], exp_k[n], rtol=rtol) + + @pytest.mark.parametrize("imag", [0, 0j]) + @pytest.mark.parametrize("n_outer", [-1, -2, -3]) + def test_factorialx_negative_extend_complex_poles(self, n_outer, imag): + kw = {"extend": "complex"} + def _check(n): + complexify = _is_subdtype(type(n), "c") + # like for gamma, we expect complex nans for complex inputs + complex_nan = np.complex128("nan+nanj") + exp = np.complex128("nan+nanj") if complexify else np.float64("nan") + # poles are at negative integers that are multiples of k + assert_really_equal(special.factorial(n, **kw), exp) + assert_really_equal(special.factorial2(n * 2, **kw), exp) + assert_really_equal(special.factorialk(n * 3, k=3, **kw), exp) + # also test complex k for factorialk + c = 1.5 - 2j + assert_really_equal(special.factorialk(n * c, k=c, **kw), complex_nan) + # same for array case + assert_really_equal(special.factorial([n], **kw)[0], exp) + assert_really_equal(special.factorial2([n * 2], **kw)[0], exp) + assert_really_equal(special.factorialk([n * 3], k=3, **kw)[0], exp) + assert_really_equal(special.factorialk([n * c], k=c, **kw)[0], complex_nan) + # more specific tests in test_factorial{,2,k}_complex_reference + + # imag ensures we test both real and complex representations of the poles + _check(n_outer + imag) + # check for large multiple of period + _check(100_000 * n_outer + imag) + + @pytest.mark.parametrize("boxed", [True, False]) + @pytest.mark.parametrize("extend", ["zero", "complex"]) + @pytest.mark.parametrize( + "n", + [ + np.nan, np.float64("nan"), np.nan + np.nan*1j, np.complex128("nan+nanj"), + np.inf, np.inf + 0j, -np.inf, -np.inf + 0j, None, np.datetime64("nat", "s") + ], + ids=[ + "NaN", "np.float64('nan')", "NaN+i*NaN", "np.complex128('nan+nanj')", + "inf", "inf+0i", "-inf", "-inf+0i", "None", "NaT" + ] + ) + @pytest.mark.parametrize( + "factorialx", + [special.factorial, special.factorial2, special.factorialk] + ) + def test_factorialx_inf_nan(self, factorialx, n, extend, boxed): + # NaNs not allowed (by dtype) for exact=True + kw = {"exact": False, "extend": extend} + if factorialx == special.factorialk: + kw["k"] = 3 + + # None is allowed for scalars, but would cause object type in array case + permissible_types = ["i", "f", "c"] if boxed else ["i", "f", "c", type(None)] + # factorial allows floats also for extend="zero" + types_need_complex_ext = "c" if factorialx == special.factorial else ["f", "c"] + + if not _is_subdtype(type(n), permissible_types): + with pytest.raises(ValueError, match="Unsupported data type.*"): + factorialx([n] if boxed else n, **kw) + elif _is_subdtype(type(n), types_need_complex_ext) and extend != "complex": + with pytest.raises(ValueError, match="In order to use non-integer.*"): + factorialx([n] if boxed else n, **kw) + else: + # account for type and whether extend="complex" + complexify = (extend == "complex") and _is_subdtype(type(n), "c") + # note that the type of the naïve `np.nan + np.nan * 1j` is `complex` + # instead of `numpy.complex128`, which trips up assert_really_equal + expected = np.complex128("nan+nanj") if complexify else np.float64("nan") + # the only exception are real infinities + if _is_subdtype(type(n), "f") and np.isinf(n): + # unchanged for positive infinity; negative one depends on extension + neg_inf_result = np.float64(0 if (extend == "zero") else "nan") + expected = np.float64("inf") if (n > 0) else neg_inf_result + + result = factorialx([n], **kw)[0] if boxed else factorialx(n, **kw) + assert_really_equal(result, expected) + # also tested in test_factorial{,2,k}_{array,scalar}_corner_cases + + @pytest.mark.parametrize("extend", [0, 1.1, np.nan, "string"]) + def test_factorialx_raises_extend(self, extend): + with pytest.raises(ValueError, match="argument `extend` must be.*"): + special.factorial(1, extend=extend) + with pytest.raises(ValueError, match="argument `extend` must be.*"): + special.factorial2(1, extend=extend) + with pytest.raises(ValueError, match="argument `extend` must be.*"): + special.factorialk(1, k=3, exact=True, extend=extend) + + @pytest.mark.parametrize("levels", range(1, 5)) + @pytest.mark.parametrize("exact", [True, False]) + def test_factorialx_array_shape(self, levels, exact): + def _nest_me(x, k=1): + """ + Double x and nest it k times + + For example: + >>> _nest_me([3, 4], 2) + [[[3, 4], [3, 4]], [[3, 4], [3, 4]]] + """ + if k == 0: + return x + else: + return _nest_me([x, x], k-1) + + def _check(res, nucleus): + exp = np.array(_nest_me(nucleus, k=levels), dtype=object) + # test that ndarray shape is maintained + # need to cast to float due to numpy/numpy#21220 + assert_allclose(res.astype(np.float64), exp.astype(np.float64)) + + n = np.array(_nest_me([5, 25], k=levels)) + exp_nucleus = {1: [120, math.factorial(25)], + # correctness of factorial{2,k}() is tested elsewhere + 2: [15, special.factorial2(25, exact=True)], + 3: [10, special.factorialk(25, 3, exact=True)]} + + _check(special.factorial(n, exact=exact), exp_nucleus[1]) + _check(special.factorial2(n, exact=exact), exp_nucleus[2]) + _check(special.factorialk(n, 3, exact=exact), exp_nucleus[3]) + + @pytest.mark.parametrize("exact", [True, False]) + @pytest.mark.parametrize("dtype", [ + None, int, np.int8, np.int16, np.int32, np.int64, + np.uint8, np.uint16, np.uint32, np.uint64 + ]) + @pytest.mark.parametrize("dim", range(0, 5)) + def test_factorialx_array_dimension(self, dim, dtype, exact): + n = np.array(5, dtype=dtype, ndmin=dim) + exp = {1: 120, 2: 15, 3: 10} + assert_allclose(special.factorial(n, exact=exact), + np.array(exp[1], ndmin=dim)) + assert_allclose(special.factorial2(n, exact=exact), + np.array(exp[2], ndmin=dim)) + assert_allclose(special.factorialk(n, 3, exact=exact), + np.array(exp[3], ndmin=dim)) + + @pytest.mark.parametrize("exact", [True, False]) + @pytest.mark.parametrize("level", range(1, 5)) + def test_factorialx_array_like(self, level, exact): + def _nest_me(x, k=1): + if k == 0: + return x + else: + return _nest_me([x], k-1) + + n = _nest_me([5], k=level-1) # nested list + exp_nucleus = {1: 120, 2: 15, 3: 10} + assert_func = assert_array_equal if exact else assert_allclose + assert_func(special.factorial(n, exact=exact), + np.array(exp_nucleus[1], ndmin=level)) + assert_func(special.factorial2(n, exact=exact), + np.array(exp_nucleus[2], ndmin=level)) + assert_func(special.factorialk(n, 3, exact=exact), + np.array(exp_nucleus[3], ndmin=level)) + + @pytest.mark.fail_slow(5) + @pytest.mark.parametrize("dtype", [np.uint8, np.uint16, np.uint32, np.uint64]) + @pytest.mark.parametrize("exact,extend", + [(True, "zero"), (False, "zero"), (False, "complex")]) + def test_factorialx_uint(self, exact, extend, dtype): + # ensure that uint types work correctly as inputs + kw = {"exact": exact, "extend": extend} + assert_func = assert_array_equal if exact else assert_allclose + def _check(n): + n_ref = n.astype(np.int64) if isinstance(n, np.ndarray) else np.int64(n) + assert_func(special.factorial(n, **kw), special.factorial(n_ref, **kw)) + assert_func(special.factorial2(n, **kw), special.factorial2(n_ref, **kw)) + assert_func(special.factorialk(n, k=3, **kw), + special.factorialk(n_ref, k=3, **kw)) + def _check_inf(n): + # produce inf of same type/dimension + with warnings.catch_warnings(): + warnings.simplefilter("ignore", RuntimeWarning) + shaped_inf = n / 0 + assert_func(special.factorial(n, **kw), shaped_inf) + assert_func(special.factorial2(n, **kw), shaped_inf) + assert_func(special.factorialk(n, k=3, **kw), shaped_inf) + + _check(dtype(0)) + _check(dtype(1)) + _check(np.array(0, dtype=dtype)) + _check(np.array([0, 1], dtype=dtype)) + # test that maximal uint values work as well + N = dtype(np.iinfo(dtype).max) + # TODO: cannot use N itself yet; factorial uses `gamma(N+1)` resp. `(hi+lo)//2` + if dtype == np.uint64: + if exact: + # avoid attempting huge calculation + pass + else: + # N does not fit into int64 --> cannot use _check + _check_inf(dtype(N-1)) + _check_inf(np.array(N-1, dtype=dtype)) + _check_inf(np.array([N-1], dtype=dtype)) + elif dtype in [np.uint8, np.uint16] or not exact: + # factorial(65535, exact=True) has 287189 digits and is calculated almost + # instantaneously on modern hardware; however, dtypes bigger than uint16 + # would blow up runtime and memory consumption for exact=True + _check(N-1) + _check(np.array(N-1, dtype=dtype)) + _check(np.array([N-2, N-1], dtype=dtype)) + + # note that n=170 is the last integer such that factorial(n) fits float64 + @pytest.mark.parametrize('n', range(30, 180, 10)) + def test_factorial_accuracy(self, n): + # Compare exact=True vs False, i.e. that the accuracy of the + # approximation is better than the specified tolerance. + + rtol = 6e-14 if sys.platform == 'win32' else 1e-15 + # need to cast exact result to float due to numpy/numpy#21220 + assert_allclose(float(special.factorial(n, exact=True)), + special.factorial(n, exact=False), rtol=rtol) + assert_allclose(special.factorial([n], exact=True).astype(float), + special.factorial([n], exact=False), rtol=rtol) + + @pytest.mark.parametrize('n', + list(range(0, 22)) + list(range(30, 180, 10))) + def test_factorial_int_reference(self, n): + # Compare all with math.factorial + correct = math.factorial(n) + assert_array_equal(correct, special.factorial(n, exact=True)) + assert_array_equal(correct, special.factorial([n], exact=True)[0]) + + rtol = 8e-14 if sys.platform == 'win32' else 1e-15 + # need to cast exact result to float due to numpy/numpy#21220 + correct = float(correct) + assert_allclose(correct, special.factorial(n, exact=False), rtol=rtol) + assert_allclose(correct, special.factorial([n], exact=False)[0], rtol=rtol) + + # extend="complex" only works for exact=False + kw = {"exact": False, "extend": "complex"} + assert_allclose(correct, special.factorial(n, **kw), rtol=rtol) + assert_allclose(correct, special.factorial([n], **kw)[0], rtol=rtol) + + def test_factorial_float_reference(self): + def _check(n, expected): + rtol = 8e-14 if sys.platform == 'win32' else 1e-15 + assert_allclose(special.factorial(n), expected, rtol=rtol) + assert_allclose(special.factorial([n])[0], expected, rtol=rtol) + # using floats with `exact=True` raises an error for scalars and arrays + with pytest.raises(ValueError, match="`exact=True` only supports.*"): + special.factorial(n, exact=True) + with pytest.raises(ValueError, match="`exact=True` only supports.*"): + special.factorial([n], exact=True) + + # Reference values from mpmath for gamma(n+1) + _check(0.01, 0.994325851191506032181932988) + _check(1.11, 1.051609009483625091514147465) + _check(5.55, 314.9503192327208241614959052) + _check(11.1, 50983227.84411615655137170553) + _check(33.3, 2.493363339642036352229215273e+37) + _check(55.5, 9.479934358436729043289162027e+73) + _check(77.7, 3.060540559059579022358692625e+114) + _check(99.9, 5.885840419492871504575693337e+157) + # close to maximum for float64 + _check(170.6243, 1.79698185749571048960082e+308) + + def test_factorial_complex_reference(self): + def _check(n, expected): + rtol = 3e-15 if sys.platform == 'win32' else 2e-15 + kw = {"exact": False, "extend": "complex"} + assert_allclose(special.factorial(n, **kw), expected, rtol=rtol) + assert_allclose(special.factorial([n], **kw)[0], expected, rtol=rtol) + + # Reference values from mpmath.gamma(n+1) + # negative & complex values + _check(-0.5, expected=1.7724538509055160276) + _check(-0.5 + 0j, expected=1.7724538509055160276 + 0j) + _check(2 + 2j, expected=-0.42263728631120216694 + 0.87181425569650686062j) + # close to poles + _check(-0.9999, expected=9999.422883232725532) + _check(-1 + 0.0001j, expected=-0.57721565582674219 - 9999.9999010944009697j) + + @pytest.mark.parametrize("dtype", [np.int64, np.float64, + np.complex128, object]) + @pytest.mark.parametrize("extend", ["zero", "complex"]) + @pytest.mark.parametrize("exact", [True, False]) + @pytest.mark.parametrize("dim", range(0, 5)) + # test empty & non-empty arrays, with nans and mixed + @pytest.mark.parametrize( + "content", + [[], [1], [1.1], [np.nan], [np.nan + np.nan * 1j], [np.nan, 1]], + ids=["[]", "[1]", "[1.1]", "[NaN]", "[NaN+i*NaN]", "[NaN, 1]"], + ) + def test_factorial_array_corner_cases(self, content, dim, exact, extend, dtype): + if dtype is object and SCIPY_ARRAY_API: + pytest.skip("object arrays unsupported in array API mode") + # get dtype without calling array constructor (that might fail or mutate) + if dtype is np.int64 and any(np.isnan(x) or (x != int(x)) for x in content): + pytest.skip("impossible combination") + if dtype == np.float64 and any(_is_subdtype(type(x), "c") for x in content): + pytest.skip("impossible combination") + + kw = {"exact": exact, "extend": extend} + # np.array(x, ndim=0) will not be 0-dim. unless x is too + content = content if (dim > 0 or len(content) != 1) else content[0] + n = np.array(content, ndmin=dim, dtype=dtype) + + result = None + if extend == "complex" and exact: + with pytest.raises(ValueError, match="Incompatible options:.*"): + special.factorial(n, **kw) + elif not _is_subdtype(n.dtype, ["i", "f", "c"]): + with pytest.raises(ValueError, match="Unsupported data type.*"): + special.factorial(n, **kw) + elif _is_subdtype(n.dtype, "c") and extend != "complex": + with pytest.raises(ValueError, match="In order to use non-integer.*"): + special.factorial(n, **kw) + elif exact and not _is_subdtype(n.dtype, "i"): + with pytest.raises(ValueError, match="`exact=True` only supports.*"): + special.factorial(n, **kw) + else: + result = special.factorial(n, **kw) + + if result is not None: + # use scalar case as reference; tested separately in *_scalar_corner_cases + ref = [special.factorial(x, **kw) for x in n.ravel()] + # unpack length-1 lists so that np.array(x, ndim=0) works correctly + ref = ref[0] if len(ref) == 1 else ref + # result is empty if and only if n is empty, and has the same dimension + # as n; dtype stays the same, except when not empty and not exact: + if n.size: + cx = (extend == "complex") and _is_subdtype(n.dtype, "c") + dtype = np.complex128 if cx else (native_int if exact else np.float64) + expected = np.array(ref, ndmin=dim, dtype=dtype) + assert_really_equal(result, expected, rtol=1e-15) + + @pytest.mark.parametrize("extend", ["zero", "complex"]) + @pytest.mark.parametrize("exact", [True, False]) + @pytest.mark.parametrize("n", [1, 1.1, 2 + 2j, np.nan, np.nan + np.nan*1j, None], + ids=["1", "1.1", "2+2j", "NaN", "NaN+i*NaN", "None"]) + def test_factorial_scalar_corner_cases(self, n, exact, extend): + kw = {"exact": exact, "extend": extend} + if extend == "complex" and exact: + with pytest.raises(ValueError, match="Incompatible options:.*"): + special.factorial(n, **kw) + elif not _is_subdtype(type(n), ["i", "f", "c", type(None)]): + with pytest.raises(ValueError, match="Unsupported data type.*"): + special.factorial(n, **kw) + elif _is_subdtype(type(n), "c") and extend != "complex": + with pytest.raises(ValueError, match="In order to use non-integer.*"): + special.factorial(n, **kw) + elif n is None or np.isnan(n): + # account for dtype and whether extend="complex" + complexify = (extend == "complex") and _is_subdtype(type(n), "c") + expected = np.complex128("nan+nanj") if complexify else np.float64("nan") + assert_really_equal(special.factorial(n, **kw), expected) + elif exact and _is_subdtype(type(n), "f"): + with pytest.raises(ValueError, match="`exact=True` only supports.*"): + special.factorial(n, **kw) + else: + assert_equal(special.factorial(n, **kw), special.gamma(n + 1)) + + # use odd increment to make sure both odd & even numbers are tested! + @pytest.mark.parametrize('n', range(30, 180, 11)) + def test_factorial2_accuracy(self, n): + # Compare exact=True vs False, i.e. that the accuracy of the + # approximation is better than the specified tolerance. + + rtol = 2e-14 if sys.platform == 'win32' else 1e-15 + # need to cast exact result to float due to numpy/numpy#21220 + assert_allclose(float(special.factorial2(n, exact=True)), + special.factorial2(n, exact=False), rtol=rtol) + assert_allclose(special.factorial2([n], exact=True).astype(float), + special.factorial2([n], exact=False), rtol=rtol) + + @pytest.mark.parametrize('n', + list(range(0, 22)) + list(range(30, 180, 11))) + def test_factorial2_int_reference(self, n): + # Compare all with correct value + + # Cannot use np.product due to overflow + correct = functools.reduce(operator.mul, list(range(n, 0, -2)), 1) + + assert_array_equal(correct, special.factorial2(n, exact=True)) + assert_array_equal(correct, special.factorial2([n], exact=True)[0]) + + rtol = 2e-14 if sys.platform == 'win32' else 1e-15 + # need to cast exact result to float due to numpy/numpy#21220 + correct = float(correct) + assert_allclose(correct, special.factorial2(n, exact=False), rtol=rtol) + assert_allclose(correct, special.factorial2([n], exact=False)[0], rtol=rtol) + + # extend="complex" only works for exact=False + kw = {"exact": False, "extend": "complex"} + # approximation only matches exactly for `n == 1 (mod k)`, see docstring + if n % 2 == 1: + assert_allclose(correct, special.factorial2(n, **kw), rtol=rtol) + assert_allclose(correct, special.factorial2([n], **kw)[0], rtol=rtol) + + def test_factorial2_complex_reference(self): + # this tests for both floats and complex + def _check(n, expected): + rtol = 5e-15 + kw = {"exact": False, "extend": "complex"} + assert_allclose(special.factorial2(n, **kw), expected, rtol=rtol) + assert_allclose(special.factorial2([n], **kw)[0], expected, rtol=rtol) + + # Reference values from mpmath for: + # mpmath.power(2, n/2) * mpmath.gamma(n/2 + 1) * mpmath.sqrt(2 / mpmath.pi) + _check(3, expected=3) + _check(4, expected=special.factorial2(4) * math.sqrt(2 / math.pi)) + _check(20, expected=special.factorial2(20) * math.sqrt(2 / math.pi)) + # negative & complex values + _check(-0.5, expected=0.82217895866245855122) + _check(-0.5 + 0j, expected=0.82217895866245855122 + 0j) + _check(3 + 3j, expected=-1.0742236630142471526 + 1.4421398439387262897j) + # close to poles + _check(-1.9999, expected=7978.8918745523440682) + _check(-2 + 0.0001j, expected=0.0462499835314308444 - 7978.84559148876374493j) + + @pytest.mark.parametrize("dtype", [np.int64, np.float64, + np.complex128, object]) + @pytest.mark.parametrize("extend", ["zero", "complex"]) + @pytest.mark.parametrize("exact", [True, False]) + @pytest.mark.parametrize("dim", range(0, 5)) + # test empty & non-empty arrays, with nans and mixed + @pytest.mark.parametrize( + "content", + [[], [1], [1.1], [np.nan], [np.nan + np.nan * 1j], [np.nan, 1]], + ids=["[]", "[1]", "[1.1]", "[NaN]", "[NaN+i*NaN]", "[NaN, 1]"], + ) + def test_factorial2_array_corner_cases(self, content, dim, exact, extend, dtype): + # get dtype without calling array constructor (that might fail or mutate) + if dtype == np.int64 and any(np.isnan(x) or (x != int(x)) for x in content): + pytest.skip("impossible combination") + if dtype == np.float64 and any(_is_subdtype(type(x), "c") for x in content): + pytest.skip("impossible combination") + + kw = {"exact": exact, "extend": extend} + # np.array(x, ndim=0) will not be 0-dim. unless x is too + content = content if (dim > 0 or len(content) != 1) else content[0] + n = np.array(content, ndmin=dim, dtype=dtype) + + result = None + if extend == "complex" and exact: + with pytest.raises(ValueError, match="Incompatible options:.*"): + special.factorial2(n, **kw) + elif not _is_subdtype(n.dtype, ["i", "f", "c"]): + with pytest.raises(ValueError, match="Unsupported data type.*"): + special.factorial2(n, **kw) + elif _is_subdtype(n.dtype, ["f", "c"]) and extend != "complex": + with pytest.raises(ValueError, match="In order to use non-integer.*"): + special.factorial2(n, **kw) + else: + result = special.factorial2(n, **kw) + + if result is not None: + # use scalar case as reference; tested separately in *_scalar_corner_cases + ref = [special.factorial2(x, **kw) for x in n.ravel()] + # unpack length-1 lists so that np.array(x, ndim=0) works correctly + ref = ref[0] if len(ref) == 1 else ref + # result is empty if and only if n is empty, and has the same dimension + # as n; dtype stays the same, except when not empty and not exact: + if n.size: + cx = (extend == "complex") and _is_subdtype(n.dtype, "c") + dtype = np.complex128 if cx else (native_int if exact else np.float64) + expected = np.array(ref, ndmin=dim, dtype=dtype) + assert_really_equal(result, expected, rtol=2e-15) + + @pytest.mark.parametrize("extend", ["zero", "complex"]) + @pytest.mark.parametrize("exact", [True, False]) + @pytest.mark.parametrize("n", [1, 1.1, 2 + 2j, np.nan, np.nan + np.nan*1j, None], + ids=["1", "1.1", "2+2j", "NaN", "NaN+i*NaN", "None"]) + def test_factorial2_scalar_corner_cases(self, n, exact, extend): + kw = {"exact": exact, "extend": extend} + if extend == "complex" and exact: + with pytest.raises(ValueError, match="Incompatible options:.*"): + special.factorial2(n, **kw) + elif not _is_subdtype(type(n), ["i", "f", "c", type(None)]): + with pytest.raises(ValueError, match="Unsupported data type.*"): + special.factorial2(n, **kw) + elif _is_subdtype(type(n), ["f", "c"]) and extend != "complex": + with pytest.raises(ValueError, match="In order to use non-integer.*"): + special.factorial2(n, **kw) + elif n is None or np.isnan(n): + # account for dtype and whether extend="complex" + complexify = (extend == "complex") and _is_subdtype(type(n), "c") + expected = np.complex128("nan+nanj") if complexify else np.float64("nan") + assert_really_equal(special.factorial2(n, **kw), expected) + else: + expected = self.factorialk_ref(n, k=2, **kw) + assert_really_equal(special.factorial2(n, **kw), expected, rtol=1e-15) + + @pytest.mark.parametrize("k", range(1, 5)) + # note that n=170 is the last integer such that factorial(n) fits float64; + # use odd increment to make sure both odd & even numbers are tested + @pytest.mark.parametrize('n', range(170, 20, -29)) + def test_factorialk_accuracy(self, n, k): + # Compare exact=True vs False, i.e. that the accuracy of the + # approximation is better than the specified tolerance. + + rtol = 6e-14 if sys.platform == 'win32' else 2e-14 + # need to cast exact result to float due to numpy/numpy#21220 + assert_allclose(float(special.factorialk(n, k=k, exact=True)), + special.factorialk(n, k=k, exact=False), rtol=rtol) + assert_allclose(special.factorialk([n], k=k, exact=True).astype(float), + special.factorialk([n], k=k, exact=False), rtol=rtol) + + @pytest.mark.parametrize('k', list(range(1, 5)) + [10, 20]) + @pytest.mark.parametrize('n', + list(range(0, 22)) + list(range(22, 100, 11))) + def test_factorialk_int_reference(self, n, k): + # Compare all with correct value + + # Would be nice to use np.product here, but that's + # broken on windows, see numpy/numpy#21219 + correct = functools.reduce(operator.mul, list(range(n, 0, -k)), 1) + + assert_array_equal(correct, special.factorialk(n, k, exact=True)) + assert_array_equal(correct, special.factorialk([n], k, exact=True)[0]) + + rtol = 3e-14 if sys.platform == 'win32' else 1e-14 + # need to cast exact result to float due to numpy/numpy#21220 + correct = float(correct) + assert_allclose(correct, special.factorialk(n, k, exact=False), rtol=rtol) + assert_allclose(correct, special.factorialk([n], k, exact=False)[0], rtol=rtol) + + # extend="complex" only works for exact=False + kw = {"k": k, "exact": False, "extend": "complex"} + # approximation only matches exactly for `n == 1 (mod k)`, see docstring + if n % k == 1: + rtol = 2e-14 + assert_allclose(correct, special.factorialk(n, **kw), rtol=rtol) + assert_allclose(correct, special.factorialk([n], **kw)[0], rtol=rtol) + + def test_factorialk_complex_reference(self): + # this tests for both floats and complex + def _check(n, k, exp): + rtol = 1e-14 + kw = {"k": k, "exact": False, "extend": "complex"} + assert_allclose(special.factorialk(n, **kw), exp, rtol=rtol) + assert_allclose(special.factorialk([n], **kw)[0], exp, rtol=rtol) + + # Reference values from mpmath for: + # mpmath.power(k, (n-1)/k) * mpmath.gamma(n/k + 1) / mpmath.gamma(1/k + 1) + _check(n=4, k=3, exp=special.factorialk(4, k=3, exact=True)) + _check(n=5, k=3, exp=7.29011132947227083) + _check(n=6.5, k=3, exp=19.6805080113566010) + # non-integer k + _check(n=3, k=2.5, exp=2.58465740293218541) + _check(n=11, k=2.5, exp=1963.5) # ==11*8.5*6*3.5; c.f. n == 1 (mod k) + _check(n=-3 + 3j + 1, k=-3 + 3j, exp=-2 + 3j) + # complex values + _check(n=4 + 4j, k=4, exp=-0.67855904082768043854 + 2.1993925819930311497j) + _check(n=4, k=4 - 4j, exp=1.9775338957222718742 + 0.92607172675423901371j) + _check(n=4 + 4j, k=4 - 4j, exp=0.1868492880824934475 + 0.87660580316894290247j) + # negative values + _check(n=-0.5, k=3, exp=0.72981013240713739354) + _check(n=-0.5 + 0j, k=3, exp=0.72981013240713739354 + 0j) + _check(n=2.9, k=-0.7, exp=0.45396591474966867296 + 0.56925525174685228866j) + _check(n=-0.6, k=-0.7, exp=-0.07190820089634757334 - 0.090170031876701730081j) + # close to poles + _check(n=-2.9999, k=3, exp=7764.7170695908828364) + _check(n=-3 + 0.0001j, k=3, exp=0.1349475632879599864 - 7764.5821055158365027j) + + @pytest.mark.parametrize("dtype", [np.int64, np.float64, + np.complex128, object]) + @pytest.mark.parametrize("extend", ["zero", "complex"]) + @pytest.mark.parametrize("exact", [True, False]) + @pytest.mark.parametrize("dim", range(0, 5)) + # test empty & non-empty arrays, with nans and mixed + @pytest.mark.parametrize( + "content", + [[], [1], [1.1], [np.nan], [np.nan + np.nan * 1j], [np.nan, 1]], + ids=["[]", "[1]", "[1.1]", "[NaN]", "[NaN+i*NaN]", "[NaN, 1]"], + ) + def test_factorialk_array_corner_cases(self, content, dim, exact, extend, dtype): + # get dtype without calling array constructor (that might fail or mutate) + if dtype == np.int64 and any(np.isnan(x) or (x != int(x)) for x in content): + pytest.skip("impossible combination") + if dtype == np.float64 and any(_is_subdtype(type(x), "c") for x in content): + pytest.skip("impossible combination") + + kw = {"k": 3, "exact": exact, "extend": extend} + # np.array(x, ndim=0) will not be 0-dim. unless x is too + content = content if (dim > 0 or len(content) != 1) else content[0] + n = np.array(content, ndmin=dim, dtype=dtype) + + result = None + if extend == "complex" and exact: + with pytest.raises(ValueError, match="Incompatible options:.*"): + special.factorialk(n, **kw) + elif not _is_subdtype(n.dtype, ["i", "f", "c"]): + with pytest.raises(ValueError, match="Unsupported data type.*"): + special.factorialk(n, **kw) + elif _is_subdtype(n.dtype, ["f", "c"]) and extend != "complex": + with pytest.raises(ValueError, match="In order to use non-integer.*"): + special.factorialk(n, **kw) + else: + result = special.factorialk(n, **kw) + + if result is not None: + # use scalar case as reference; tested separately in *_scalar_corner_cases + ref = [special.factorialk(x, **kw) for x in n.ravel()] + # unpack length-1 lists so that np.array(x, ndim=0) works correctly + ref = ref[0] if len(ref) == 1 else ref + # result is empty if and only if n is empty, and has the same dimension + # as n; dtype stays the same, except when not empty and not exact: + if n.size: + cx = (extend == "complex") and _is_subdtype(n.dtype, "c") + dtype = np.complex128 if cx else (native_int if exact else np.float64) + expected = np.array(ref, ndmin=dim, dtype=dtype) + assert_really_equal(result, expected, rtol=2e-15) + + @pytest.mark.parametrize("extend", ["zero", "complex"]) + @pytest.mark.parametrize("exact", [True, False]) + @pytest.mark.parametrize("k", range(1, 5)) + @pytest.mark.parametrize("n", [1, 1.1, 2 + 2j, np.nan, np.nan + np.nan*1j, None], + ids=["1", "1.1", "2+2j", "NaN", "NaN+i*NaN", "None"]) + def test_factorialk_scalar_corner_cases(self, n, k, exact, extend): + kw = {"k": k, "exact": exact, "extend": extend} + if extend == "complex" and exact: + with pytest.raises(ValueError, match="Incompatible options:.*"): + special.factorialk(n, **kw) + elif not _is_subdtype(type(n), ["i", "f", "c", type(None)]): + with pytest.raises(ValueError, match="Unsupported data type.*"): + special.factorialk(n, **kw) + elif _is_subdtype(type(n), ["f", "c"]) and extend != "complex": + with pytest.raises(ValueError, match="In order to use non-integer.*"): + special.factorialk(n, **kw) + elif n is None or np.isnan(n): + # account for dtype and whether extend="complex" + complexify = (extend == "complex") and _is_subdtype(type(n), "c") + expected = np.complex128("nan+nanj") if complexify else np.float64("nan") + assert_really_equal(special.factorialk(n, **kw), expected) + else: + expected = self.factorialk_ref(n, **kw) + assert_really_equal(special.factorialk(n, **kw), expected, rtol=1e-15) + + @pytest.mark.parametrize("boxed", [True, False]) + @pytest.mark.parametrize("exact,extend", + [(True, "zero"), (False, "zero"), (False, "complex")]) + @pytest.mark.parametrize("k", [-1, -1.0, 0, 0.0, 0 + 1j, 1.1, np.nan]) + def test_factorialk_raises_k_complex(self, k, exact, extend, boxed): + n = [1] if boxed else 1 + kw = {"k": k, "exact": exact, "extend": extend} + if extend == "zero": + msg = "In order to use non-integer.*" + if _is_subdtype(type(k), "i") and (k < 1): + msg = "For `extend='zero'`.*" + with pytest.raises(ValueError, match=msg): + special.factorialk(n, **kw) + elif k == 0: + with pytest.raises(ValueError, match="Parameter k cannot be zero!"): + special.factorialk(n, **kw) + else: + # no error + special.factorialk(n, **kw) + + @pytest.mark.parametrize("boxed", [True, False]) + @pytest.mark.parametrize("exact,extend", + [(True, "zero"), (False, "zero"), (False, "complex")]) + # neither integer, float nor complex + @pytest.mark.parametrize("k", ["string", np.datetime64("nat", "s")], + ids=["string", "NaT"]) + def test_factorialk_raises_k_other(self, k, exact, extend, boxed): + n = [1] if boxed else 1 + kw = {"k": k, "exact": exact, "extend": extend} + with pytest.raises(ValueError, match="Unsupported data type.*"): + special.factorialk(n, **kw) + + @pytest.mark.parametrize("exact,extend", + [(True, "zero"), (False, "zero"), (False, "complex")]) + @pytest.mark.parametrize("k", range(1, 12)) + def test_factorialk_dtype(self, k, exact, extend): + kw = {"k": k, "exact": exact, "extend": extend} + if exact and k in _FACTORIALK_LIMITS_64BITS.keys(): + n = np.array([_FACTORIALK_LIMITS_32BITS[k]]) + assert_equal(special.factorialk(n, **kw).dtype, np.long) + assert_equal(special.factorialk(n + 1, **kw).dtype, np.int64) + # assert maximality of limits for given dtype + assert special.factorialk(n + 1, **kw) > np.iinfo(np.int32).max + + n = np.array([_FACTORIALK_LIMITS_64BITS[k]]) + assert_equal(special.factorialk(n, **kw).dtype, np.int64) + assert_equal(special.factorialk(n + 1, **kw).dtype, object) + assert special.factorialk(n + 1, **kw) > np.iinfo(np.int64).max + else: + n = np.array([_FACTORIALK_LIMITS_64BITS.get(k, 1)]) + # for exact=True and k >= 10, we always return object; + # for exact=False it's always float (unless input is complex) + dtype = object if exact else np.float64 + assert_equal(special.factorialk(n, **kw).dtype, dtype) + + def test_factorial_mixed_nan_inputs(self): + x = np.array([np.nan, 1, 2, 3, np.nan]) + expected = np.array([np.nan, 1, 2, 6, np.nan]) + assert_equal(special.factorial(x, exact=False), expected) + with pytest.raises(ValueError, match="`exact=True` only supports.*"): + special.factorial(x, exact=True) + + +class TestFresnel: + @pytest.mark.parametrize("z, s, c", [ + # some positive value + (.5, 0.064732432859999287, 0.49234422587144644), + (.5 + .0j, 0.064732432859999287, 0.49234422587144644), + # negative half annulus + # https://github.com/scipy/scipy/issues/12309 + # Reference values can be reproduced with + # https://www.wolframalpha.com/input/?i=FresnelS%5B-2.0+%2B+0.1i%5D + # https://www.wolframalpha.com/input/?i=FresnelC%5B-2.0+%2B+0.1i%5D + ( + -2.0 + 0.1j, + -0.3109538687728942-0.0005870728836383176j, + -0.4879956866358554+0.10670801832903172j + ), + ( + -0.1 - 1.5j, + -0.03918309471866977+0.7197508454568574j, + 0.09605692502968956-0.43625191013617465j + ), + # a different algorithm kicks in for "large" values, i.e., |z| >= 4.5, + # make sure to test both float and complex values; a different + # algorithm is used + (6.0, 0.44696076, 0.49953147), + (6.0 + 0.0j, 0.44696076, 0.49953147), + (6.0j, -0.44696076j, 0.49953147j), + (-6.0 + 0.0j, -0.44696076, -0.49953147), + (-6.0j, 0.44696076j, -0.49953147j), + # inf + (np.inf, 0.5, 0.5), + (-np.inf, -0.5, -0.5), + ]) + def test_fresnel_values(self, z, s, c): + frs = array(special.fresnel(z)) + assert_allclose(frs, array([s, c]), atol=1.5e-8, rtol=0) + + # values from pg 329 Table 7.11 of A & S + # slightly corrected in 4th decimal place + def test_fresnel_zeros(self): + szo, czo = special.fresnel_zeros(5) + assert_allclose(szo, array([2.0093+0.2885j, + 2.8335+0.2443j, + 3.4675+0.2185j, + 4.0026+0.2009j, + 4.4742+0.1877j]), + atol=1.5e-3, rtol=0) + assert_allclose(czo, array([1.7437+0.3057j, + 2.6515+0.2529j, + 3.3204+0.2240j, + 3.8757+0.2047j, + 4.3611+0.1907j]), + atol=1.5e-3, rtol=0) + vals1 = special.fresnel(szo)[0] + vals2 = special.fresnel(czo)[1] + assert_allclose(vals1, 0, atol=1.5e-14, rtol=0) + assert_allclose(vals2, 0, atol=1.5e-14, rtol=0) + + def test_fresnelc_zeros(self): + szo, czo = special.fresnel_zeros(6) + frc = special.fresnelc_zeros(6) + assert_allclose(frc, czo, atol=1.5e-12, rtol=0) + + def test_fresnels_zeros(self): + szo, czo = special.fresnel_zeros(5) + frs = special.fresnels_zeros(5) + assert_allclose(frs, szo, atol=1.5e-12, rtol=0) + + +class TestGamma: + def test_gamma(self): + gam = special.gamma(5) + assert_equal(gam,24.0) + + def test_gammaln(self): + gamln = special.gammaln(3) + lngam = log(special.gamma(3)) + assert_allclose(gamln, lngam, atol=1.5e-8, rtol=0) + + def test_gammainccinv(self): + gccinv = special.gammainccinv(.5,.5) + gcinv = special.gammaincinv(.5,.5) + assert_allclose(gccinv, gcinv, atol=1.5e-8, rtol=0) + + @with_special_errors + def test_gammaincinv(self): + y = special.gammaincinv(.4,.4) + x = special.gammainc(.4,y) + assert_allclose(x, 0.4, atol=1.5e-10, rtol=0) + y = special.gammainc(10, 0.05) + x = special.gammaincinv(10, 2.5715803516000736e-20) + assert_allclose(0.05, x, atol=1.5e-10, rtol=0) + assert_allclose(y, 2.5715803516000736e-20, atol=1.5e-10, rtol=0) + x = special.gammaincinv(50, 8.20754777388471303050299243573393e-18) + assert_allclose(11.0, x, atol=1.5e-10, rtol=0) + + @with_special_errors + def test_975(self): + # Regression test for ticket #975 -- switch point in algorithm + # check that things work OK at the point, immediately next floats + # around it, and a bit further away + pts = [0.25, + np.nextafter(0.25, 0), 0.25 - 1e-12, + np.nextafter(0.25, 1), 0.25 + 1e-12] + for pt in pts: + y = special.gammaincinv(.4, pt) + x = special.gammainc(0.4, y) + assert_allclose(x, pt, rtol=1e-12) + + def test_rgamma(self): + rgam = special.rgamma(8) + rlgam = 1/special.gamma(8) + assert_allclose(rgam, rlgam, atol=1.5e-8, rtol=0) + + def test_infinity(self): + assert_equal(special.rgamma(-1), 0) + + @pytest.mark.parametrize( + "x,expected", + [ + # infinities + ([-np.inf, np.inf], [np.nan, np.inf]), + # negative and positive zero + ([-0.0, 0.0], [-np.inf, np.inf]), + # small poles + (range(-32, 0), np.full(32, np.nan)), + # medium sized poles + (range(-1024, -32, 99), np.full(11, np.nan)), + # large pole + ([-4.141512231792294e+16], [np.nan]), + ] + ) + def test_poles(self, x, expected): + assert_array_equal(special.gamma(x), expected) + + +class TestHankel: + + def test_negv1(self): + assert_allclose(special.hankel1(-3, 2), -special.hankel1(3, 2), + atol=1.5e-14, rtol=0) + + def test_hankel1(self): + hank1 = special.hankel1(1,.1) + hankrl = (special.jv(1,.1) + special.yv(1,.1)*1j) + assert_allclose(hank1, hankrl, atol=1.5e-8, rtol=0) + + def test_negv1e(self): + assert_allclose(special.hankel1e(-3, 2), -special.hankel1e(3, 2), + atol=1.5e-14, rtol=0) + + def test_hankel1e(self): + hank1e = special.hankel1e(1,.1) + hankrle = special.hankel1(1,.1)*exp(-.1j) + assert_allclose(hank1e, hankrle, atol=1.5e-8, rtol=0) + + def test_negv2(self): + assert_allclose(special.hankel2(-3, 2), -special.hankel2(3, 2), + atol=1.5e-14, rtol=0) + + def test_hankel2(self): + hank2 = special.hankel2(1,.1) + hankrl2 = (special.jv(1,.1) - special.yv(1,.1)*1j) + assert_allclose(hank2, hankrl2, atol=1.5e-8, rtol=0) + + def test_neg2e(self): + assert_allclose(special.hankel2e(-3, 2), -special.hankel2e(3, 2), + atol=1.5e-14, rtol=0) + + def test_hankl2e(self): + hank2e = special.hankel2e(1,.1) + hankrl2e = special.hankel2e(1,.1) + assert_allclose(hank2e, hankrl2e, atol=1.5e-8, rtol=0) + + def test_hankel2_gh4517(self): + # Test edge case reported in https://github.com/scipy/scipy/issues/4517 + res = special.hankel2(0, 0) + assert np.isnan(res.real) + assert np.isposinf(res.imag) + + +class TestHyper: + def test_h1vp(self): + h1 = special.h1vp(1,.1) + h1real = (special.jvp(1,.1) + special.yvp(1,.1)*1j) + assert_allclose(h1, h1real, atol=1.5e-8, rtol=0) + + def test_h2vp(self): + h2 = special.h2vp(1,.1) + h2real = (special.jvp(1,.1) - special.yvp(1,.1)*1j) + assert_allclose(h2, h2real, atol=1.5e-8, rtol=0) + + def test_hyp0f1(self): + # scalar input + assert_allclose(special.hyp0f1(2.5, 0.5), 1.21482702689997, rtol=1e-12) + assert_allclose(special.hyp0f1(2.5, 0), 1.0, rtol=1e-15) + + # float input, expected values match mpmath + x = special.hyp0f1(3.0, [-1.5, -1, 0, 1, 1.5]) + expected = np.array([0.58493659229143, 0.70566805723127, 1.0, + 1.37789689539747, 1.60373685288480]) + assert_allclose(x, expected, rtol=1e-12) + + # complex input + x = special.hyp0f1(3.0, np.array([-1.5, -1, 0, 1, 1.5]) + 0.j) + assert_allclose(x, expected.astype(complex), rtol=1e-12) + + # test broadcasting + x1 = [0.5, 1.5, 2.5] + x2 = [0, 1, 0.5] + x = special.hyp0f1(x1, x2) + expected = [1.0, 1.8134302039235093, 1.21482702689997] + assert_allclose(x, expected, rtol=1e-12) + x = special.hyp0f1(np.vstack([x1] * 2), x2) + assert_allclose(x, np.vstack([expected] * 2), rtol=1e-12) + assert_raises(ValueError, special.hyp0f1, + np.vstack([x1] * 3), [0, 1]) + + def test_hyp0f1_gh5764(self): + # Just checks the point that failed; there's a more systematic + # test in test_mpmath + res = special.hyp0f1(0.8, 0.5 + 0.5*1J) + # The expected value was generated using mpmath + assert_allclose(res, 1.6139719776441115 + 1J*0.80893054061790665, + atol=1.5e-7, rtol=0) + + def test_hyp1f1(self): + hyp1 = special.hyp1f1(.1,.1,.3) + assert_allclose(hyp1, 1.3498588075760032, atol=1.5e-7, rtol=0) + + # test contributed by Moritz Deger (2008-05-29) + # https://github.com/scipy/scipy/issues/1186 (Trac #659) + + # reference data obtained from mathematica [ a, b, x, m(a,b,x)]: + # produced with test_hyp1f1.nb + ref_data = array([ + [-8.38132975e+00, -1.28436461e+01, -2.91081397e+01, 1.04178330e+04], + [2.91076882e+00, -6.35234333e+00, -1.27083993e+01, 6.68132725e+00], + [-1.42938258e+01, 1.80869131e-01, 1.90038728e+01, 1.01385897e+05], + [5.84069088e+00, 1.33187908e+01, 2.91290106e+01, 1.59469411e+08], + [-2.70433202e+01, -1.16274873e+01, -2.89582384e+01, 1.39900152e+24], + [4.26344966e+00, -2.32701773e+01, 1.91635759e+01, 6.13816915e+21], + [1.20514340e+01, -3.40260240e+00, 7.26832235e+00, 1.17696112e+13], + [2.77372955e+01, -1.99424687e+00, 3.61332246e+00, 3.07419615e+13], + [1.50310939e+01, -2.91198675e+01, -1.53581080e+01, -3.79166033e+02], + [1.43995827e+01, 9.84311196e+00, 1.93204553e+01, 2.55836264e+10], + [-4.08759686e+00, 1.34437025e+01, -1.42072843e+01, 1.70778449e+01], + [8.05595738e+00, -1.31019838e+01, 1.52180721e+01, 3.06233294e+21], + [1.81815804e+01, -1.42908793e+01, 9.57868793e+00, -2.84771348e+20], + [-2.49671396e+01, 1.25082843e+01, -1.71562286e+01, 2.36290426e+07], + [2.67277673e+01, 1.70315414e+01, 6.12701450e+00, 7.77917232e+03], + [2.49565476e+01, 2.91694684e+01, 6.29622660e+00, 2.35300027e+02], + [6.11924542e+00, -1.59943768e+00, 9.57009289e+00, 1.32906326e+11], + [-1.47863653e+01, 2.41691301e+01, -1.89981821e+01, 2.73064953e+03], + [2.24070483e+01, -2.93647433e+00, 8.19281432e+00, -6.42000372e+17], + [8.04042600e-01, 1.82710085e+01, -1.97814534e+01, 5.48372441e-01], + [1.39590390e+01, 1.97318686e+01, 2.37606635e+00, 5.51923681e+00], + [-4.66640483e+00, -2.00237930e+01, 7.40365095e+00, 4.50310752e+00], + [2.76821999e+01, -6.36563968e+00, 1.11533984e+01, -9.28725179e+23], + [-2.56764457e+01, 1.24544906e+00, 1.06407572e+01, 1.25922076e+01], + [3.20447808e+00, 1.30874383e+01, 2.26098014e+01, 2.03202059e+04], + [-1.24809647e+01, 4.15137113e+00, -2.92265700e+01, 2.39621411e+08], + [2.14778108e+01, -2.35162960e+00, -1.13758664e+01, 4.46882152e-01], + [-9.85469168e+00, -3.28157680e+00, 1.67447548e+01, -1.07342390e+07], + [1.08122310e+01, -2.47353236e+01, -1.15622349e+01, -2.91733796e+03], + [-2.67933347e+01, -3.39100709e+00, 2.56006986e+01, -5.29275382e+09], + [-8.60066776e+00, -8.02200924e+00, 1.07231926e+01, 1.33548320e+06], + [-1.01724238e-01, -1.18479709e+01, -2.55407104e+01, 1.55436570e+00], + [-3.93356771e+00, 2.11106818e+01, -2.57598485e+01, 2.13467840e+01], + [3.74750503e+00, 1.55687633e+01, -2.92841720e+01, 1.43873509e-02], + [6.99726781e+00, 2.69855571e+01, -1.63707771e+01, 3.08098673e-02], + [-2.31996011e+01, 3.47631054e+00, 9.75119815e-01, 1.79971073e-02], + [2.38951044e+01, -2.91460190e+01, -2.50774708e+00, 9.56934814e+00], + [1.52730825e+01, 5.77062507e+00, 1.21922003e+01, 1.32345307e+09], + [1.74673917e+01, 1.89723426e+01, 4.94903250e+00, 9.90859484e+01], + [1.88971241e+01, 2.86255413e+01, 5.52360109e-01, 1.44165360e+00], + [1.02002319e+01, -1.66855152e+01, -2.55426235e+01, 6.56481554e+02], + [-1.79474153e+01, 1.22210200e+01, -1.84058212e+01, 8.24041812e+05], + [-1.36147103e+01, 1.32365492e+00, -7.22375200e+00, 9.92446491e+05], + [7.57407832e+00, 2.59738234e+01, -1.34139168e+01, 3.64037761e-02], + [2.21110169e+00, 1.28012666e+01, 1.62529102e+01, 1.33433085e+02], + [-2.64297569e+01, -1.63176658e+01, -1.11642006e+01, -2.44797251e+13], + [-2.46622944e+01, -3.02147372e+00, 8.29159315e+00, -3.21799070e+05], + [-1.37215095e+01, -1.96680183e+01, 2.91940118e+01, 3.21457520e+12], + [-5.45566105e+00, 2.81292086e+01, 1.72548215e-01, 9.66973000e-01], + [-1.55751298e+00, -8.65703373e+00, 2.68622026e+01, -3.17190834e+16], + [2.45393609e+01, -2.70571903e+01, 1.96815505e+01, 1.80708004e+37], + [5.77482829e+00, 1.53203143e+01, 2.50534322e+01, 1.14304242e+06], + [-1.02626819e+01, 2.36887658e+01, -2.32152102e+01, 7.28965646e+02], + [-1.30833446e+00, -1.28310210e+01, 1.87275544e+01, -9.33487904e+12], + [5.83024676e+00, -1.49279672e+01, 2.44957538e+01, -7.61083070e+27], + [-2.03130747e+01, 2.59641715e+01, -2.06174328e+01, 4.54744859e+04], + [1.97684551e+01, -2.21410519e+01, -2.26728740e+01, 3.53113026e+06], + [2.73673444e+01, 2.64491725e+01, 1.57599882e+01, 1.07385118e+07], + [5.73287971e+00, 1.21111904e+01, 1.33080171e+01, 2.63220467e+03], + [-2.82751072e+01, 2.08605881e+01, 9.09838900e+00, -6.60957033e-07], + [1.87270691e+01, -1.74437016e+01, 1.52413599e+01, 6.59572851e+27], + [6.60681457e+00, -2.69449855e+00, 9.78972047e+00, -2.38587870e+12], + [1.20895561e+01, -2.51355765e+01, 2.30096101e+01, 7.58739886e+32], + [-2.44682278e+01, 2.10673441e+01, -1.36705538e+01, 4.54213550e+04], + [-4.50665152e+00, 3.72292059e+00, -4.83403707e+00, 2.68938214e+01], + [-7.46540049e+00, -1.08422222e+01, -1.72203805e+01, -2.09402162e+02], + [-2.00307551e+01, -7.50604431e+00, -2.78640020e+01, 4.15985444e+19], + [1.99890876e+01, 2.20677419e+01, -2.51301778e+01, 1.23840297e-09], + [2.03183823e+01, -7.66942559e+00, 2.10340070e+01, 1.46285095e+31], + [-2.90315825e+00, -2.55785967e+01, -9.58779316e+00, 2.65714264e-01], + [2.73960829e+01, -1.80097203e+01, -2.03070131e+00, 2.52908999e+02], + [-2.11708058e+01, -2.70304032e+01, 2.48257944e+01, 3.09027527e+08], + [2.21959758e+01, 4.00258675e+00, -1.62853977e+01, -9.16280090e-09], + [1.61661840e+01, -2.26845150e+01, 2.17226940e+01, -8.24774394e+33], + [-3.35030306e+00, 1.32670581e+00, 9.39711214e+00, -1.47303163e+01], + [7.23720726e+00, -2.29763909e+01, 2.34709682e+01, -9.20711735e+29], + [2.71013568e+01, 1.61951087e+01, -7.11388906e-01, 2.98750911e-01], + [8.40057933e+00, -7.49665220e+00, 2.95587388e+01, 6.59465635e+29], + [-1.51603423e+01, 1.94032322e+01, -7.60044357e+00, 1.05186941e+02], + [-8.83788031e+00, -2.72018313e+01, 1.88269907e+00, 1.81687019e+00], + [-1.87283712e+01, 5.87479570e+00, -1.91210203e+01, 2.52235612e+08], + [-5.61338513e-01, 2.69490237e+01, 1.16660111e-01, 9.97567783e-01], + [-5.44354025e+00, -1.26721408e+01, -4.66831036e+00, 1.06660735e-01], + [-2.18846497e+00, 2.33299566e+01, 9.62564397e+00, 3.03842061e-01], + [6.65661299e+00, -2.39048713e+01, 1.04191807e+01, 4.73700451e+13], + [-2.57298921e+01, -2.60811296e+01, 2.74398110e+01, -5.32566307e+11], + [-1.11431826e+01, -1.59420160e+01, -1.84880553e+01, -1.01514747e+02], + [6.50301931e+00, 2.59859051e+01, -2.33270137e+01, 1.22760500e-02], + [-1.94987891e+01, -2.62123262e+01, 3.90323225e+00, 1.71658894e+01], + [7.26164601e+00, -1.41469402e+01, 2.81499763e+01, -2.50068329e+31], + [-1.52424040e+01, 2.99719005e+01, -2.85753678e+01, 1.31906693e+04], + [5.24149291e+00, -1.72807223e+01, 2.22129493e+01, 2.50748475e+25], + [3.63207230e-01, -9.54120862e-02, -2.83874044e+01, 9.43854939e-01], + [-2.11326457e+00, -1.25707023e+01, 1.17172130e+00, 1.20812698e+00], + [2.48513582e+00, 1.03652647e+01, -1.84625148e+01, 6.47910997e-02], + [2.65395942e+01, 2.74794672e+01, 1.29413428e+01, 2.89306132e+05], + [-9.49445460e+00, 1.59930921e+01, -1.49596331e+01, 3.27574841e+02], + [-5.89173945e+00, 9.96742426e+00, 2.60318889e+01, -3.15842908e-01], + [-1.15387239e+01, -2.21433107e+01, -2.17686413e+01, 1.56724718e-01], + [-5.30592244e+00, -2.42752190e+01, 1.29734035e+00, 1.31985534e+00] + ]) + + for a,b,c,expected in ref_data: + result = special.hyp1f1(a,b,c) + assert_(abs(expected - result)/expected < 1e-4) + + def test_hyp1f1_gh2957(self): + hyp1 = special.hyp1f1(0.5, 1.5, -709.7827128933) + hyp2 = special.hyp1f1(0.5, 1.5, -709.7827128934) + assert_allclose(hyp1, hyp2, atol=1.5e-12, rtol=0) + + def test_hyp1f1_gh2282(self): + hyp = special.hyp1f1(0.5, 1.5, -1000) + assert_allclose(hyp, 0.028024956081989643, atol=1.5e-12, rtol=0) + + def test_hyp2f1(self): + # a collection of special cases taken from AMS 55 + values = [ + [0.5, 1, 1.5, 0.2**2, 0.5/0.2*log((1+0.2)/(1-0.2))], + [0.5, 1, 1.5, -0.2**2, 1./0.2*arctan(0.2)], + [1, 1, 2, 0.2, -1/0.2*log(1-0.2)], + [3, 3.5, 1.5, 0.2**2, 0.5/0.2/(-5)*((1+0.2)**(-5)-(1-0.2)**(-5))], + [-3, 3, 0.5, sin(0.2)**2, cos(2*3*0.2)], + [3, 4, 8, 1, + special.gamma(8) * special.gamma(8-4-3) + / special.gamma(8-3) / special.gamma(8-4)], + [3, 2, 3-2+1, -1, + 1./2**3*sqrt(pi) * special.gamma(1+3-2) + / special.gamma(1+0.5*3-2) / special.gamma(0.5+0.5*3)], + [5, 2, 5-2+1, -1, + 1./2**5*sqrt(pi) * special.gamma(1+5-2) + / special.gamma(1+0.5*5-2) / special.gamma(0.5+0.5*5)], + [4, 0.5+4, 1.5-2*4, -1./3, + (8./9)**(-2*4)*special.gamma(4./3) * special.gamma(1.5-2*4) + / special.gamma(3./2) / special.gamma(4./3-2*4)], + # and some others + # ticket #424 + [1.5, -0.5, 1.0, -10.0, 4.1300097765277476484], + # negative integer a or b, with c-a-b integer and x > 0.9 + [-2,3,1,0.95,0.715], + [2,-3,1,0.95,-0.007], + [-6,3,1,0.95,0.0000810625], + [2,-5,1,0.95,-0.000029375], + # huge negative integers + (10, -900, 10.5, 0.99, 1.91853705796607664803709475658e-24), + (10, -900, -10.5, 0.99, 3.54279200040355710199058559155e-18), + ] + for i, (a, b, c, x, v) in enumerate(values): + cv = special.hyp2f1(a, b, c, x) + assert_allclose(cv, v, atol=1.5e-8, rtol=0, err_msg=f'test #{i}') + + def test_hyperu(self): + val1 = special.hyperu(1,0.1,100) + assert_allclose(val1, 0.0098153, atol=1.5e-7, rtol=0) + a,b = [0.3,0.6,1.2,-2.7],[1.5,3.2,-0.4,-3.2] + a,b = asarray(a), asarray(b) + z = 0.5 + hypu = special.hyperu(a,b,z) + hprl = (pi/sin(pi*b))*(special.hyp1f1(a,b,z) / + (special.gamma(1+a-b)*special.gamma(b)) - + z**(1-b)*special.hyp1f1(1+a-b,2-b,z) + / (special.gamma(a)*special.gamma(2-b))) + assert_allclose(hypu, hprl, atol=1.5e-12, rtol=0) + + def test_hyperu_gh2287(self): + assert_allclose(special.hyperu(1, 1.5, 20.2), 0.048360918656699191, + atol=1.5e-12, rtol=0) + + +class TestBessel: + def test_itj0y0(self): + it0 = array(special.itj0y0(.2)) + assert_allclose(it0, array([0.19933433254006822, -0.34570883800412566]), + atol=1.5e-8, rtol=0) + + def test_it2j0y0(self): + it2 = array(special.it2j0y0(.2)) + assert_allclose(it2, array([0.0049937546274601858, -0.43423067011231614]), + atol=1.5e-8, rtol=0) + + def test_negv_iv(self): + assert_equal(special.iv(3,2), special.iv(-3,2)) + + def test_j0(self): + oz = special.j0(.1) + ozr = special.jn(0,.1) + assert_allclose(oz, ozr, atol=1.5e-8, rtol=0) + + def test_j1(self): + o1 = special.j1(.1) + o1r = special.jn(1,.1) + assert_allclose(o1, o1r, atol=1.5e-8, rtol=0) + + def test_jn(self): + jnnr = special.jn(1,.2) + assert_allclose(jnnr, 0.099500832639235995, atol=1.5e-8, rtol=0) + + def test_negv_jv(self): + assert_allclose(special.jv(-3, 2), -special.jv(3, 2), atol=1.5e-14, rtol=0) + + def test_jv(self): + values = [[0, 0.1, 0.99750156206604002], + [2./3, 1e-8, 0.3239028506761532e-5], + [2./3, 1e-10, 0.1503423854873779e-6], + [3.1, 1e-10, 0.1711956265409013e-32], + [2./3, 4.0, -0.2325440850267039], + ] + for i, (v, x, y) in enumerate(values): + yc = special.jv(v, x) + assert_allclose(yc, y, atol=1.5e-8, rtol=0, err_msg=f'test #{i}') + + def test_negv_jve(self): + assert_allclose(special.jve(-3, 2), -special.jve(3, 2), + atol=1.5e-14, rtol=0) + + def test_jve(self): + jvexp = special.jve(1,.2) + assert_allclose(jvexp, 0.099500832639235995, atol=1.5e-8, rtol=0) + jvexp1 = special.jve(1,.2+1j) + z = .2+1j + jvexpr = special.jv(1,z)*exp(-abs(z.imag)) + assert_allclose(jvexp1, jvexpr, atol=1.5e-8, rtol=0) + + def test_jn_zeros(self): + jn0 = special.jn_zeros(0,5) + jn1 = special.jn_zeros(1,5) + assert_allclose(jn0, array([2.4048255577, + 5.5200781103, + 8.6537279129, + 11.7915344391, + 14.9309177086]), + atol=1.5e-4, rtol=0) + assert_allclose(jn1, array([3.83171, + 7.01559, + 10.17347, + 13.32369, + 16.47063]), + atol=1.5e-4, rtol=0) + + jn102 = special.jn_zeros(102,5) + assert_allclose(jn102, array([110.89174935992040343, + 117.83464175788308398, + 123.70194191713507279, + 129.02417238949092824, + 134.00114761868422559]), rtol=1e-13) + + jn301 = special.jn_zeros(301,5) + assert_allclose(jn301, array([313.59097866698830153, + 323.21549776096288280, + 331.22338738656748796, + 338.39676338872084500, + 345.03284233056064157]), rtol=1e-13) + + def test_jn_zeros_slow(self): + jn0 = special.jn_zeros(0, 300) + assert_allclose(jn0[260-1], 816.02884495068867280, rtol=1e-13) + assert_allclose(jn0[280-1], 878.86068707124422606, rtol=1e-13) + assert_allclose(jn0[300-1], 941.69253065317954064, rtol=1e-13) + + jn10 = special.jn_zeros(10, 300) + assert_allclose(jn10[260-1], 831.67668514305631151, rtol=1e-13) + assert_allclose(jn10[280-1], 894.51275095371316931, rtol=1e-13) + assert_allclose(jn10[300-1], 957.34826370866539775, rtol=1e-13) + + jn3010 = special.jn_zeros(3010,5) + assert_allclose(jn3010, array([3036.86590780927, + 3057.06598526482, + 3073.66360690272, + 3088.37736494778, + 3101.86438139042]), rtol=1e-8) + + def test_jnjnp_zeros(self): + jn = special.jn + + def jnp(n, x): + return (jn(n-1,x) - jn(n+1,x))/2 + for nt in range(1, 30): + z, n, m, t = special.jnjnp_zeros(nt) + for zz, nn, tt in zip(z, n, t): + if tt == 0: + assert_allclose(jn(nn, zz), 0, atol=1e-6) + elif tt == 1: + assert_allclose(jnp(nn, zz), 0, atol=1e-6) + else: + raise AssertionError(f"Invalid t return for nt={nt}") + + def test_jnp_zeros(self): + jnp = special.jnp_zeros(1,5) + assert_allclose(jnp, array([1.84118, + 5.33144, + 8.53632, + 11.70600, + 14.86359]), + atol=1.5e-4, rtol=0) + jnp = special.jnp_zeros(443,5) + assert_allclose(special.jvp(443, jnp), 0, atol=1e-15) + + def test_jnyn_zeros(self): + jnz = special.jnyn_zeros(1, 5) + assert_allclose(jnz, (array([3.83171, + 7.01559, + 10.17347, + 13.32369, + 16.47063]), + array([1.84118, + 5.33144, + 8.53632, + 11.70600, + 14.86359]), + array([2.19714, + 5.42968, + 8.59601, + 11.74915, + 14.89744]), + array([3.68302, + 6.94150, + 10.12340, + 13.28576, + 16.44006])), + atol=1.5e-5, rtol=0) + + def test_jvp(self): + jvprim = special.jvp(2,2) + jv0 = (special.jv(1,2)-special.jv(3,2))/2 + assert_allclose(jvprim, jv0, atol=1.5e-10, rtol=0) + + def test_k0(self): + ozk = special.k0(.1) + ozkr = special.kv(0,.1) + assert_allclose(ozk,ozkr, atol=1.5e-8, rtol=0) + + def test_k0e(self): + ozke = special.k0e(.1) + ozker = special.kve(0,.1) + assert_allclose(ozke, ozker, atol=1.5e-8, rtol=0) + + def test_k1(self): + o1k = special.k1(.1) + o1kr = special.kv(1,.1) + assert_allclose(o1k,o1kr, atol=1.5e-8, rtol=0) + + def test_k1e(self): + o1ke = special.k1e(.1) + o1ker = special.kve(1,.1) + assert_allclose(o1ke, o1ker, atol=1.5e-8, rtol=0) + + def test_jacobi(self): + a = 5*np.random.random() - 1 + b = 5*np.random.random() - 1 + P0 = special.jacobi(0,a,b) + P1 = special.jacobi(1,a,b) + P2 = special.jacobi(2,a,b) + P3 = special.jacobi(3,a,b) + + assert_allclose(P0.c, [1], atol=1.5e-13, rtol=0) + assert_allclose(P1.c, array([a + b + 2, a - b]) / 2.0, + atol=1.5e-13, rtol=0) + cp = [(a+b+3)*(a+b+4), 4*(a+b+3)*(a+2), 4*(a+1)*(a+2)] + p2c = [cp[0],cp[1]-2*cp[0],cp[2]-cp[1]+cp[0]] + assert_allclose(P2.c, array(p2c) / 8.0, atol=1.5e-13, rtol=0) + cp = [(a+b+4)*(a+b+5)*(a+b+6),6*(a+b+4)*(a+b+5)*(a+3), + 12*(a+b+4)*(a+2)*(a+3),8*(a+1)*(a+2)*(a+3)] + p3c = [cp[0],cp[1]-3*cp[0],cp[2]-2*cp[1]+3*cp[0],cp[3]-cp[2]+cp[1]-cp[0]] + assert_allclose(P3.c, array(p3c) / 48.0, atol=1.5e-13, rtol=0) + + def test_kn(self): + kn1 = special.kn(0,.2) + assert_allclose(kn1, 1.7527038555281462, atol=1.5e-8, rtol=0) + + def test_negv_kv(self): + assert_equal(special.kv(3.0, 2.2), special.kv(-3.0, 2.2)) + + def test_kv0(self): + kv0 = special.kv(0,.2) + assert_allclose(kv0, 1.7527038555281462, atol=1.5e-10, rtol=0) + + def test_kv1(self): + kv1 = special.kv(1,0.2) + assert_allclose(kv1, 4.775972543220472, atol=1.5e-10, rtol=0) + + def test_kv2(self): + kv2 = special.kv(2,0.2) + assert_allclose(kv2, 49.51242928773287, atol=1.5e-10, rtol=0) + + def test_kn_largeorder(self): + assert_allclose(special.kn(32, 1), 1.7516596664574289e+43) + + def test_kv_largearg(self): + assert_equal(special.kv(0, 1e19), 0) + + def test_negv_kve(self): + assert_equal(special.kve(3.0, 2.2), special.kve(-3.0, 2.2)) + + def test_kve(self): + kve1 = special.kve(0,.2) + kv1 = special.kv(0,.2)*exp(.2) + assert_allclose(kve1, kv1, atol=1.5e-8, rtol=0) + z = .2+1j + kve2 = special.kve(0,z) + kv2 = special.kv(0,z)*exp(z) + assert_allclose(kve2, kv2, atol=1.5e-8, rtol=0) + + def test_kvp_v0n1(self): + z = 2.2 + assert_allclose(-special.kv(1, z), special.kvp(0, z, n=1), + atol=1.5e-10, rtol=0) + + def test_kvp_n1(self): + v = 3. + z = 2.2 + xc = -special.kv(v+1,z) + v/z*special.kv(v,z) + x = special.kvp(v,z, n=1) + # this function (kvp) is broken + assert_allclose(xc, x, atol=1.5e-10, rtol=0) + + def test_kvp_n2(self): + v = 3. + z = 2.2 + xc = (z**2+v**2-v)/z**2 * special.kv(v,z) + special.kv(v+1,z)/z + x = special.kvp(v, z, n=2) + assert_allclose(xc, x, atol=1.5e-10, rtol=0) + + def test_y0(self): + oz = special.y0(.1) + ozr = special.yn(0,.1) + assert_allclose(oz, ozr, atol=1.5e-8, rtol=0) + + def test_y1(self): + o1 = special.y1(.1) + o1r = special.yn(1,.1) + assert_allclose(o1,o1r, atol=1.5e-8, rtol=0) + + def test_y0_zeros(self): + yo,ypo = special.y0_zeros(2) + zo,zpo = special.y0_zeros(2,complex=1) + all = r_[yo,zo] + allval = r_[ypo,zpo] + assert_allclose(abs(special.yv(0.0, all)), 0.0, atol=1.5e-11, rtol=0) + assert_allclose(abs(special.yv(1, all) - allval), 0.0, atol=1.5e-11, rtol=0) + + def test_y1_zeros(self): + y1 = special.y1_zeros(1) + assert_allclose(y1, (array([2.19714]), array([0.52079])), + atol=1.5e-5, rtol=0) + + def test_y1p_zeros(self): + y1p = special.y1p_zeros(1,complex=1) + assert_allclose(y1p, (array([0.5768+0.904j]), array([-0.7635+0.5892j])), + atol=1.5e-3, rtol=0) + + def test_yn_zeros(self): + an = special.yn_zeros(4,2) + assert_allclose(an, array([5.64515, 9.36162]), atol=1.5e-5, rtol=0) + an = special.yn_zeros(443,5) + assert_allclose(an, [450.13573091578090314, + 463.05692376675001542, + 472.80651546418663566, + 481.27353184725625838, + 488.98055964441374646], + rtol=1e-15,) + + def test_ynp_zeros(self): + ao = special.ynp_zeros(0,2) + assert_allclose(ao, array([2.19714133, 5.42968104]), atol=1.5e-6, rtol=0) + ao = special.ynp_zeros(43,5) + assert_allclose(special.yvp(43, ao), 0, atol=1e-15) + ao = special.ynp_zeros(443,5) + assert_allclose(special.yvp(443, ao), 0, atol=1e-9) + + def test_ynp_zeros_large_order(self): + ao = special.ynp_zeros(443,5) + assert_allclose(special.yvp(443, ao), 0, atol=1e-14) + + def test_yn(self): + yn2n = special.yn(1,.2) + assert_allclose(yn2n, -3.3238249881118471, atol=1.5e-8, rtol=0) + + def test_yn_gh_20405(self): + # Enforce correct asymptotic behavior for large n. + observed = cephes.yn(500, 1) + assert observed == -np.inf + + def test_negv_yv(self): + assert_allclose(special.yv(-3, 2), -special.yv(3, 2), + atol=1.5e-14, rtol=0) + + def test_yv(self): + yv2 = special.yv(1,.2) + assert_allclose(yv2, -3.3238249881118471, atol=1.5e-8, rtol=0) + + def test_negv_yve(self): + assert_allclose(special.yve(-3, 2), -special.yve(3, 2), + atol=1.5e-14, rtol=0) + + def test_yve(self): + yve2 = special.yve(1,.2) + assert_allclose(yve2, -3.3238249881118471, atol=1.5e-8, rtol=0) + yve2r = special.yv(1,.2+1j)*exp(-1) + yve22 = special.yve(1,.2+1j) + assert_allclose(yve22, yve2r, atol=1.5e-8, rtol=0) + + def test_yvp(self): + yvpr = (special.yv(1,.2) - special.yv(3,.2))/2.0 + yvp1 = special.yvp(2,.2) + assert_allclose(yvp1, yvpr, atol=1.5e-10, rtol=0) + + def _cephes_vs_amos_points(self): + """Yield points at which to compare Cephes implementation to AMOS""" + # check several points, including large-amplitude ones + v = [-120, -100.3, -20., -10., -1., -.5, 0., 1., 12.49, 120., 301] + z = [-1300, -11, -10, -1, 1., 10., 200.5, 401., 600.5, 700.6, 1300, + 10003] + yield from itertools.product(v, z) + + # check half-integers; these are problematic points at least + # for cephes/iv + yield from itertools.product(0.5 + arange(-60, 60), [3.5]) + + def check_cephes_vs_amos(self, f1, f2, rtol=1e-11, atol=0, skip=None): + for v, z in self._cephes_vs_amos_points(): + if skip is not None and skip(v, z): + continue + c1, c2, c3 = f1(v, z), f1(v,z+0j), f2(int(v), z) + if np.isinf(c1): + assert_(np.abs(c2) >= 1e300, (v, z)) + elif np.isnan(c1): + assert_(c2.imag != 0, (v, z)) + else: + assert_allclose(c1, c2, err_msg=(v, z), rtol=rtol, atol=atol) + if v == int(v): + assert_allclose(c3, c2, err_msg=(v, z), + rtol=rtol, atol=atol) + + @pytest.mark.xfail(platform.machine() == 'ppc64le', + reason="fails on ppc64le") + def test_jv_cephes_vs_amos(self): + self.check_cephes_vs_amos(special.jv, special.jn, rtol=1e-10, atol=1e-305) + + @pytest.mark.xfail(platform.machine() == 'ppc64le', + reason="fails on ppc64le") + def test_yv_cephes_vs_amos(self): + self.check_cephes_vs_amos(special.yv, special.yn, rtol=1e-11, atol=1e-305) + + def test_yv_cephes_vs_amos_only_small_orders(self): + def skipper(v, z): + return abs(v) > 50 + self.check_cephes_vs_amos(special.yv, special.yn, rtol=1e-11, atol=1e-305, + skip=skipper) + + def test_iv_cephes_vs_amos(self): + with np.errstate(all='ignore'): + self.check_cephes_vs_amos(special.iv, special.iv, rtol=5e-9, atol=1e-305) + + @pytest.mark.slow + def test_iv_cephes_vs_amos_mass_test(self): + N = 1000000 + np.random.seed(1) + v = np.random.pareto(0.5, N) * (-1)**np.random.randint(2, size=N) + x = np.random.pareto(0.2, N) * (-1)**np.random.randint(2, size=N) + + imsk = (np.random.randint(8, size=N) == 0) + v[imsk] = v[imsk].astype(np.int64) + + with np.errstate(all='ignore'): + c1 = special.iv(v, x) + c2 = special.iv(v, x+0j) + + # deal with differences in the inf and zero cutoffs + c1[abs(c1) > 1e300] = np.inf + c2[abs(c2) > 1e300] = np.inf + c1[abs(c1) < 1e-300] = 0 + c2[abs(c2) < 1e-300] = 0 + + dc = abs(c1/c2 - 1) + dc[np.isnan(dc)] = 0 + + k = np.argmax(dc) + + # Most error apparently comes from AMOS and not our implementation; + # there are some problems near integer orders there + assert_( + dc[k] < 2e-7, + (v[k], x[k], special.iv(v[k], x[k]), special.iv(v[k], x[k]+0j)) + ) + + def test_kv_cephes_vs_amos(self): + self.check_cephes_vs_amos(special.kv, special.kn, rtol=1e-9, atol=1e-305) + self.check_cephes_vs_amos(special.kv, special.kv, rtol=1e-9, atol=1e-305) + + def test_ticket_623(self): + assert_allclose(special.jv(3, 4), 0.43017147387562193) + assert_allclose(special.jv(301, 1300), 0.0183487151115275) + assert_allclose(special.jv(301, 1296.0682), -0.0224174325312048) + + def test_ticket_853(self): + """Negative-order Bessels""" + # cephes + assert_allclose(special.jv(-1, 1), -0.4400505857449335) + assert_allclose(special.jv(-2, 1), 0.1149034849319005) + assert_allclose(special.yv(-1, 1), 0.7812128213002887) + assert_allclose(special.yv(-2, 1), -1.650682606816255) + assert_allclose(special.iv(-1, 1), 0.5651591039924851) + assert_allclose(special.iv(-2, 1), 0.1357476697670383) + assert_allclose(special.kv(-1, 1), 0.6019072301972347) + assert_allclose(special.kv(-2, 1), 1.624838898635178) + assert_allclose(special.jv(-0.5, 1), 0.43109886801837607952) + assert_allclose(special.yv(-0.5, 1), 0.6713967071418031) + assert_allclose(special.iv(-0.5, 1), 1.231200214592967) + assert_allclose(special.kv(-0.5, 1), 0.4610685044478945) + # amos + assert_allclose(special.jv(-1, 1+0j), -0.4400505857449335) + assert_allclose(special.jv(-2, 1+0j), 0.1149034849319005) + assert_allclose(special.yv(-1, 1+0j), 0.7812128213002887) + assert_allclose(special.yv(-2, 1+0j), -1.650682606816255) + + assert_allclose(special.iv(-1, 1+0j), 0.5651591039924851) + assert_allclose(special.iv(-2, 1+0j), 0.1357476697670383) + assert_allclose(special.kv(-1, 1+0j), 0.6019072301972347) + assert_allclose(special.kv(-2, 1+0j), 1.624838898635178) + + assert_allclose(special.jv(-0.5, 1+0j), 0.43109886801837607952) + assert_allclose(special.jv(-0.5, 1+1j), 0.2628946385649065-0.827050182040562j) + assert_allclose(special.yv(-0.5, 1+0j), 0.6713967071418031) + assert_allclose(special.yv(-0.5, 1+1j), 0.967901282890131+0.0602046062142816j) + + assert_allclose(special.iv(-0.5, 1+0j), 1.231200214592967) + assert_allclose(special.iv(-0.5, 1+1j), 0.77070737376928+0.39891821043561j) + assert_allclose(special.kv(-0.5, 1+0j), 0.4610685044478945) + assert_allclose(special.kv(-0.5, 1+1j), 0.06868578341999-0.38157825981268j) + + assert_allclose(special.jve(-0.5,1+0.3j), special.jv(-0.5, 1+0.3j)*exp(-0.3)) + assert_allclose(special.yve(-0.5,1+0.3j), special.yv(-0.5, 1+0.3j)*exp(-0.3)) + assert_allclose(special.ive(-0.5,0.3+1j), special.iv(-0.5, 0.3+1j)*exp(-0.3)) + assert_allclose(special.kve(-0.5,0.3+1j), special.kv(-0.5, 0.3+1j)*exp(0.3+1j)) + + assert_allclose( + special.hankel1(-0.5, 1+1j), + special.jv(-0.5, 1+1j) + 1j*special.yv(-0.5,1+1j) + ) + assert_allclose( + special.hankel2(-0.5, 1+1j), + special.jv(-0.5, 1+1j) - 1j*special.yv(-0.5,1+1j) + ) + + def test_ticket_854(self): + """Real-valued Bessel domains""" + assert_(isnan(special.jv(0.5, -1))) + assert_(isnan(special.iv(0.5, -1))) + assert_(isnan(special.yv(0.5, -1))) + assert_(isnan(special.yv(1, -1))) + assert_(isnan(special.kv(0.5, -1))) + assert_(isnan(special.kv(1, -1))) + assert_(isnan(special.jve(0.5, -1))) + assert_(isnan(special.ive(0.5, -1))) + assert_(isnan(special.yve(0.5, -1))) + assert_(isnan(special.yve(1, -1))) + assert_(isnan(special.kve(0.5, -1))) + assert_(isnan(special.kve(1, -1))) + assert_(isnan(special.airye(-1)[0:2]).all(), special.airye(-1)) + assert_(not isnan(special.airye(-1)[2:4]).any(), special.airye(-1)) + + def test_gh_7909(self): + assert_(special.kv(1.5, 0) == np.inf) + assert_(special.kve(1.5, 0) == np.inf) + + def test_ticket_503(self): + """Real-valued Bessel I overflow""" + assert_allclose(special.iv(1, 700), 1.528500390233901e302) + assert_allclose(special.iv(1000, 1120), 1.301564549405821e301) + + def test_iv_hyperg_poles(self): + assert_allclose(special.iv(-0.5, 1), 1.231200214592967) + + def iv_series(self, v, z, n=200): + k = arange(0, n).astype(double) + r = (v+2*k)*log(.5*z) - special.gammaln(k+1) - special.gammaln(v+k+1) + r[isnan(r)] = inf + r = exp(r) + err = abs(r).max() * finfo(double).eps * n + abs(r[-1])*10 + return r.sum(), err + + def test_i0_series(self): + for z in [1., 10., 200.5]: + value, err = self.iv_series(0, z) + assert_allclose(special.i0(z), value, atol=err, err_msg=z) + + def test_i1_series(self): + for z in [1., 10., 200.5]: + value, err = self.iv_series(1, z) + assert_allclose(special.i1(z), value, atol=err, err_msg=z) + + def test_iv_series(self): + for v in [-20., -10., -1., 0., 1., 12.49, 120.]: + for z in [1., 10., 200.5, -1+2j]: + value, err = self.iv_series(v, z) + assert_allclose(special.iv(v, z), value, atol=err, err_msg=(v, z)) + + def test_i0(self): + values = [[0.0, 1.0], + [1e-10, 1.0], + [0.1, 0.9071009258], + [0.5, 0.6450352706], + [1.0, 0.4657596077], + [2.5, 0.2700464416], + [5.0, 0.1835408126], + [20.0, 0.0897803119], + ] + for i, (x, v) in enumerate(values): + cv = special.i0(x) * exp(-x) + assert_allclose(cv, v, atol=1.5e-8, rtol=0, err_msg=f'test #{i}') + + def test_i0e(self): + oize = special.i0e(.1) + oizer = special.ive(0, .1) + assert_allclose(oize, oizer, atol=1.5e-8, rtol=0) + + def test_i1(self): + values = [[0.0, 0.0], + [1e-10, 0.4999999999500000e-10], + [0.1, 0.0452984468], + [0.5, 0.1564208032], + [1.0, 0.2079104154], + [5.0, 0.1639722669], + [20.0, 0.0875062222], + ] + for i, (x, v) in enumerate(values): + cv = special.i1(x) * exp(-x) + assert_allclose(cv, v, atol=1.5e-8, rtol=0, err_msg=f'test #{i}') + + def test_i1e(self): + oi1e = special.i1e(.1) + oi1er = special.ive(1, .1) + assert_allclose(oi1e, oi1er, atol=1.5e-8, rtol=0) + + def test_iti0k0(self): + iti0 = array(special.iti0k0(5)) + assert_allclose(iti0, array([31.848667776169801, 1.5673873907283657]), + atol=1.5e-5, rtol=0) + + def test_it2i0k0(self): + it2k = special.it2i0k0(.1) + assert_allclose(it2k, array([0.0012503906973464409, 3.3309450354686687]), + atol=1.5e-6, rtol=0) + + def test_iv(self): + iv1 = special.iv(0,.1)*exp(-.1) + assert_allclose(iv1, 0.90710092578230106, atol=1.5e-10, rtol=0) + + def test_negv_ive(self): + assert_equal(special.ive(3,2), special.ive(-3,2)) + + def test_ive(self): + ive1 = special.ive(0,.1) + iv1 = special.iv(0,.1)*exp(-.1) + assert_allclose(ive1, iv1, atol=1.5e-10, rtol=0) + + def test_ivp0(self): + assert_allclose(special.iv(1, 2), special.ivp(0, 2), atol=1.5e-10, rtol=0) + + def test_ivp(self): + y = (special.iv(0,2) + special.iv(2,2))/2 + x = special.ivp(1,2) + assert_allclose(x, y, atol=1.5e-10, rtol=0) + + @pytest.mark.parametrize("x, expected", + [(1e15, 6.156638646885021e-09), # 1/sqrt(eps) < x < 1/eps + (1e30, -5.589003016686147e-16) # x > 1/eps + ]) + def test_gh22705(self, x, expected): + # reference values computed with mpmath + # from mpmath import mp + # mp.dps = 1000 + # float(mp.besselj(0, mp.mpf('1e15'))) + assert_allclose(special.j0(x), expected, rtol=5e-15) + + +class TestLaguerre: + def test_laguerre(self): + lag0 = special.laguerre(0) + lag1 = special.laguerre(1) + lag2 = special.laguerre(2) + lag3 = special.laguerre(3) + lag4 = special.laguerre(4) + lag5 = special.laguerre(5) + assert_allclose(lag0.c, [1], atol=1.5e-13, rtol=0) + assert_allclose(lag1.c, [-1, 1], atol=1.5e-13, rtol=0) + assert_allclose(lag2.c, array([1, -4,2]) / 2.0, atol=1.5e-13, rtol=0) + assert_allclose(lag3.c, array([-1, 9,-18,6])/6.0, atol=1.5e-13, rtol=0) + assert_allclose(lag4.c, array([1, -16,72,-96,24])/24.0, + atol=1.5e-13, rtol=0) + assert_allclose(lag5.c, array([-1, 25, -200, 600, -600, 120]) / 120.0, + atol=1.5e-13, rtol=0) + + def test_genlaguerre(self): + k = 5*np.random.random() - 0.9 + lag0 = special.genlaguerre(0,k) + lag1 = special.genlaguerre(1,k) + lag2 = special.genlaguerre(2,k) + lag3 = special.genlaguerre(3,k) + assert_equal(lag0.c, [1]) + assert_equal(lag1.c, [-1, k + 1]) + assert_allclose(lag2.c, array([1, -2 * (k + 2), (k + 1.) * (k + 2.)]) / 2.0, + atol=1.5e-7, rtol=0) + expected = array([-1, + 3 * (k + 3), + -3 * (k + 2) * (k + 3), + (k + 1) * (k + 2) * (k + 3)]) / 6.0 + assert_allclose(lag3.c, expected, atol=1.5e-7, rtol=0) + + +class TestLambda: + def test_lmbda(self): + lam = special.lmbda(1,.1) + lamr = ( + array([special.jn(0,.1), 2*special.jn(1,.1)/.1]), + array([special.jvp(0,.1), -2*special.jv(1,.1)/.01 + 2*special.jvp(1,.1)/.1]) + ) + assert_allclose(lam, lamr, atol=1.5e-8, rtol=0) + + +class TestLog1p: + def test_log1p(self): + l1p = (special.log1p(10), special.log1p(11), special.log1p(12)) + l1prl = (log(11), log(12), log(13)) + assert_allclose(l1p, l1prl, atol=1.5e-8, rtol=0) + + def test_log1pmore(self): + l1pm = (special.log1p(1), special.log1p(1.1), special.log1p(1.2)) + l1pmrl = (log(2),log(2.1),log(2.2)) + assert_allclose(l1pm, l1pmrl, atol=1.5e-8, rtol=0) + + +def ce_fourier_coefficient_using_integral(k, n, q): + """ + Compute the Fourier coefficient of the even Mathieu function. + The integral definition of a Fourier coefficient is used. + This function is used as an alternative implementation of + mathieu_even_coef(). + """ + period = 180 if n % 2 == 0 else 360 + # For k = 0, the factor outside the integral is (1/period). + # For k = 1, 2, 3, ..., the factor is (2/period). + c = (1/period)*quad(lambda t: special.mathieu_cem(n, q, t)[0], + -period/2, period/2, + weight='cos', wvar=2*np.pi*k/period, epsrel=1e-14)[0] + if k > 0: + c *= 2 + return c + + +def se_fourier_coefficient_using_integral(k, n, q): + """ + Compute the Fourier coefficient of the odd Mathieu function. + The integral definition of a Fourier coefficient is used. + This function is used as an alternative implementation of + mathieu_odd_coef(). + """ + # For k == 0, the result is 0. (The test code won't call this + # function with k == 0, but we'll check anyway.) + if k == 0: + return 0.0 + period = 180 if n % 2 == 0 else 360 + c = (2/period)*quad(lambda t: special.mathieu_sem(n, q, t)[0], + -period/2, period/2, + weight='sin', wvar=2*np.pi*k/period, epsrel=1e-14)[0] + return c + + +class TestMathieu: + + @pytest.mark.parametrize('n, q', [(4, 3.5), (8, 4.25)]) + def test_mathieu_even_coef_against_integral_n_even(self, n, q): + # Get the nonzero Fourier coefficients. For the even Mathieu functions + # with even n, these are the coefficients of the cosine series. None of + # the coefficients are 0 for k = 0, 1, 2, 3, ... + A = special.mathieu_even_coef(n, q) + # Compare the first four nonzero Fourier coefficients to the coefficients + # computed using the integral definition. + c = [ce_fourier_coefficient_using_integral(k, n, q) for k in range(4)] + assert_allclose(c, A[:len(c)], rtol=1e-10) + + @pytest.mark.parametrize('n, q', [(3, 3.5), (7, 2)]) + def test_mathieu_even_coef_against_integral_n_odd(self, n, q): + # Get the nonzero Fourier coefficients. For the even Mathieu functions + # with odd n, these are the coefficients of the cosine series. Only the + # coefficients c[k] for k = 1, 3, 5, 7, ... are nonzero. These are the + # values returned by mathieu_even_coef(n, q). + A = special.mathieu_even_coef(n, q) + # Compare the first 4 nonzero Fourier coefficients to the coefficients + # computed using the integral definition. + c = [ce_fourier_coefficient_using_integral(k, n, q) for k in range(1, 9, 2)] + assert_allclose(c, A[:len(c)], rtol=1e-10) + + @pytest.mark.parametrize('n, q', [(2, 3.5), (10, 2)]) + def test_mathieu_odd_coef_against_integral_n_even(self, n, q): + # Get the nonzero Fourier coefficients. For the odd Mathieu functions + # with even n, these are the coefficients of the sine series. Only the + # coefficients c[k] for k = 1, 2, 3, 4, ... are nonzero. These are the + # values returned by mathieu_odd_coef(n, q). + B = special.mathieu_odd_coef(n, q) + # Compare the first 4 nonzero Fourier coefficients to the coefficients + # computed using the integral definition. + c = [se_fourier_coefficient_using_integral(k, n, q) for k in range(1, 5)] + assert_allclose(c, B[:len(c)], rtol=1e-10) + + @pytest.mark.parametrize('n, q', [(3, 3.5), (7, 2)]) + def test_mathieu_odd_coef_against_integral_n_odd(self, n, q): + # Get the nonzero Fourier coefficients. For the odd Mathieu functions + # with odd n, these are the coefficients of the sine series. Only the + # coefficients c[k] for k = 1, 3, 5, 7, ... are nonzero. These are the + # values returned by mathieu_odd_coef(n, q). + B = special.mathieu_odd_coef(n, q) + # Compare the first 4 nonzero Fourier coefficients to the coefficients + # computed using the integral definition. + c = [se_fourier_coefficient_using_integral(k, n, q) for k in range(1, 9, 2)] + assert_allclose(c, B[:len(c)], rtol=1e-10) + + +class TestFresnelIntegral: + + def test_modfresnelp(self): + pass + + def test_modfresnelm(self): + pass + + +class TestOblCvSeq: + def test_obl_cv_seq(self): + obl = special.obl_cv_seq(0,3,1) + assert_allclose(obl, array([-0.348602, + 1.393206, + 5.486800, + 11.492120]), + atol=1.5e-5, rtol=0) + + +class TestParabolicCylinder: + def test_pbdn_seq(self): + pb = special.pbdn_seq(1, .1) + assert_allclose(pb, (array([0.9975, + 0.0998]), + array([-0.0499, + 0.9925])), + atol=1.5e-4, rtol=0) + + def test_pbdv(self): + special.pbdv(1,.2) + 1/2*(.2)*special.pbdv(1,.2)[0] - special.pbdv(0,.2)[0] + + def test_pbdv_seq(self): + pbn = special.pbdn_seq(1,.1) + pbv = special.pbdv_seq(1,.1) + assert_allclose(pbv, (real(pbn[0]), real(pbn[1])), atol=1.5e-4, rtol=0) + + def test_pbdv_points(self): + # simple case + eta = np.linspace(-10, 10, 5) + z = 2**(eta/2)*np.sqrt(np.pi)*special.rgamma(.5-.5*eta) + assert_allclose(special.pbdv(eta, 0.)[0], z, rtol=1e-14, atol=1e-14) + + # some points + assert_allclose(special.pbdv(10.34, 20.44)[0], 1.3731383034455e-32, rtol=1e-12) + assert_allclose(special.pbdv(-9.53, 3.44)[0], 3.166735001119246e-8, rtol=1e-12) + + def test_pbdv_gradient(self): + x = np.linspace(-4, 4, 8)[:,None] + eta = np.linspace(-10, 10, 5)[None,:] + + p = special.pbdv(eta, x) + eps = 1e-7 + 1e-7*abs(x) + dp = (special.pbdv(eta, x + eps)[0] - special.pbdv(eta, x - eps)[0]) / eps / 2. + assert_allclose(p[1], dp, rtol=1e-6, atol=1e-6) + + def test_pbvv_gradient(self): + x = np.linspace(-4, 4, 8)[:,None] + eta = np.linspace(-10, 10, 5)[None,:] + + p = special.pbvv(eta, x) + eps = 1e-7 + 1e-7*abs(x) + dp = (special.pbvv(eta, x + eps)[0] - special.pbvv(eta, x - eps)[0]) / eps / 2. + assert_allclose(p[1], dp, rtol=1e-6, atol=1e-6) + + def test_pbvv_seq(self): + res1, res2 = special.pbvv_seq(2, 3) + assert_allclose(res1, np.array([2.976319645712036, + 1.358840996329579, + 0.5501016716383508])) + assert_allclose(res2, np.array([3.105638472238475, + 0.9380581512176672, + 0.533688488872053])) + + +class TestPolygamma: + # from Table 6.2 (pg. 271) of A&S + def test_polygamma(self): + poly2 = special.polygamma(2, 1) + poly3 = special.polygamma(3, 1) + assert_allclose(poly2, -2.4041138063, atol=1.5e-10, rtol=0) + assert_allclose(poly3, 6.4939394023, atol=1.5e-10, rtol=0) + + # Test polygamma(0, x) == psi(x) + x = [2, 3, 1.1e14] + assert_allclose(special.polygamma(0, x), special.psi(x), + atol=1.5e-7, rtol=0) + + # Test broadcasting + n = [0, 1, 2] + x = [0.5, 1.5, 2.5] + expected = [-1.9635100260214238, 0.93480220054467933, + -0.23620405164172739] + assert_allclose(special.polygamma(n, x), expected, atol=1.5e-7, rtol=0) + expected = np.vstack([expected]*2) + assert_allclose(special.polygamma(n, np.vstack([x]*2)), expected, + atol=1.5e-7, rtol=0) + assert_allclose(special.polygamma(np.vstack([n]*2), x), expected, + atol=1.5e-7, rtol=0) + + +class TestProCvSeq: + def test_pro_cv_seq(self): + prol = special.pro_cv_seq(0, 3, 1) + assert_allclose(prol, array([0.319000, + 2.593084, + 6.533471, + 12.514462]), + atol=1.5e-5, rtol=0) + + +class TestPsi: + def test_psi(self): + ps = special.psi(1) + assert_allclose(ps, -0.57721566490153287, atol=1.5e-8, rtol=0) + + +class TestRadian: + def test_radian(self): + rad = special.radian(90, 0, 0) + assert_allclose(rad, pi/2.0, atol=1.5e-5, rtol=0) + + def test_radianmore(self): + rad1 = special.radian(90, 1, 60) + assert_allclose(rad1, pi/2 + 0.0005816135199345904, atol=1.5e-5, rtol=0) + + +class TestRiccati: + def test_riccati_jn(self): + N, x = 2, 0.2 + S = np.empty((N, N)) + for n in range(N): + j = special.spherical_jn(n, x) + jp = special.spherical_jn(n, x, derivative=True) + S[0,n] = x*j + S[1,n] = x*jp + j + assert_allclose(S, special.riccati_jn(n, x), atol=1.5e-8, rtol=0) + + def test_riccati_yn(self): + N, x = 2, 0.2 + C = np.empty((N, N)) + for n in range(N): + y = special.spherical_yn(n, x) + yp = special.spherical_yn(n, x, derivative=True) + C[0,n] = x*y + C[1,n] = x*yp + y + assert_allclose(C, special.riccati_yn(n, x), atol=1.5e-8, rtol=0) + + +class TestSoftplus: + def test_softplus(self): + # Test cases for the softplus function. Selected based on Eq.(10) of: + # Mächler, M. (2012). log1mexp-note.pdf. Rmpfr: R MPFR - Multiple Precision + # Floating-Point Reliable. Retrieved from: + # https://cran.r-project.org/web/packages/Rmpfr/vignettes/log1mexp-note.pdf + # Reference values computed with `mpmath` + import numpy as np + rng = np.random.default_rng(3298432985245) + n = 3 + a1 = rng.uniform(-100, -37, size=n) + a2 = rng.uniform(-37, 18, size=n) + a3 = rng.uniform(18, 33.3, size=n) + a4 = rng.uniform(33.33, 100, size=n) + a = np.stack([a1, a2, a3, a4]) + + # from mpmath import mp + # mp.dps = 100 + # @np.vectorize + # def softplus(x): + # return float(mp.log(mp.one + mp.exp(x))) + # softplus(a).tolist() + ref = [[1.692721323272333e-42, 7.42673911145206e-41, 8.504608846033205e-35], + [1.8425343736349797, 9.488245799395577e-15, 7.225195764021444e-08], + [31.253760266045106, 27.758244090327832, 29.995959179643634], + [73.26040086468937, 76.24944728617226, 37.83955519155184]] + + res = softplus(a) + assert_allclose(res, ref, rtol=2e-15) + + def test_softplus_with_kwargs(self): + x = np.arange(5) - 2 + out = np.ones(5) + ref = out.copy() + where = x > 0 + + softplus(x, out=out, where=where) + ref[where] = softplus(x[where]) + assert_allclose(out, ref) + + +class TestRound: + def test_round(self): + rnd = list(map(int, (special.round(10.1), + special.round(10.4), + special.round(10.5), + special.round(10.6)))) + + # Note: According to the documentation, scipy.special.round is + # supposed to round to the nearest even number if the fractional + # part is exactly 0.5. On some platforms, this does not appear + # to work and thus this test may fail. However, this unit test is + # correctly written. + rndrl = (10,10,10,11) + assert_array_equal(rnd,rndrl) + +class TestStruve: + def _series(self, v, z, n=100): + """Compute Struve function & error estimate from its power series.""" + k = arange(0, n) + r = (-1)**k * (.5*z)**(2*k+v+1)/special.gamma(k+1.5)/special.gamma(k+v+1.5) + err = abs(r).max() * finfo(double).eps * n + return r.sum(), err + + def test_vs_series(self): + """Check Struve function versus its power series""" + for v in [-20, -10, -7.99, -3.4, -1, 0, 1, 3.4, 12.49, 16]: + for z in [1, 10, 19, 21, 30]: + value, err = self._series(v, z) + assert_allclose(special.struve(v, z), value, rtol=0, atol=err), (v, z) + + def test_some_values(self): + assert_allclose(special.struve(-7.99, 21), 0.0467547614113, rtol=1e-7) + assert_allclose(special.struve(-8.01, 21), 0.0398716951023, rtol=1e-8) + assert_allclose(special.struve(-3.0, 200), 0.0142134427432, rtol=1e-12) + assert_allclose(special.struve(-8.0, -41), 0.0192469727846, rtol=1e-11) + assert_equal(special.struve(-12, -41), -special.struve(-12, 41)) + assert_equal(special.struve(+12, -41), -special.struve(+12, 41)) + assert_equal(special.struve(-11, -41), +special.struve(-11, 41)) + assert_equal(special.struve(+11, -41), +special.struve(+11, 41)) + + assert_(isnan(special.struve(-7.1, -1))) + assert_(isnan(special.struve(-10.1, -1))) + + def test_regression_679(self): + """Regression test for #679""" + assert_allclose(special.struve(-1.0, 20 - 1e-8), + special.struve(-1.0, 20 + 1e-8)) + assert_allclose(special.struve(-2.0, 20 - 1e-8), + special.struve(-2.0, 20 + 1e-8)) + assert_allclose(special.struve(-4.3, 20 - 1e-8), + special.struve(-4.3, 20 + 1e-8)) + + +def test_chi2_smalldf(): + assert_allclose(special.chdtr(0.6, 3), 0.957890536704110, atol=1.5e-7, rtol=0) + + +def test_ch2_inf(): + assert_equal(special.chdtr(0.7,np.inf), 1.0) + + +@pytest.mark.parametrize("x", [-np.inf, -1.0, -0.0, 0.0, np.inf, np.nan]) +def test_chi2_v_nan(x): + assert np.isnan(special.chdtr(np.nan, x)) + + +@pytest.mark.parametrize("v", [-np.inf, -1.0, -0.0, 0.0, np.inf, np.nan]) +def test_chi2_x_nan(v): + assert np.isnan(special.chdtr(v, np.nan)) + + +@pytest.mark.parametrize("x", [-np.inf, -1.0, -0.0, 0.0, np.inf, np.nan]) +def test_chi2c_v_nan(x): + assert np.isnan(special.chdtrc(np.nan, x)) + + +@pytest.mark.parametrize("v", [-np.inf, -1.0, -0.0, 0.0, np.inf, np.nan]) +def test_chi2c_x_nan(v): + assert np.isnan(special.chdtrc(v, np.nan)) + + +def test_chi2_edgecases_gh20972(): + # Tests that a variety of edgecases for chi square distribution functions + # correctly return NaN when and only when they are supposed to, when + # computed through different related ufuncs. See gh-20972. + v = np.asarray([-0.01, 0, 0.01, 1, np.inf])[:, np.newaxis] + x = np.asarray([-np.inf, -0.01, 0, 0.01, np.inf]) + + # Check that `gammainc` is NaN when it should be and finite otherwise + ref = special.gammainc(v / 2, x / 2) + mask = (x < 0) | (v < 0) | (x == 0) & (v == 0) | np.isinf(v) & np.isinf(x) + assert np.all(np.isnan(ref[mask])) + assert np.all(np.isfinite(ref[~mask])) + + # Use `gammainc` as a reference for the rest + assert_allclose(special.chdtr(v, x), ref) + assert_allclose(special.gdtr(1, v / 2, x / 2), ref) + assert_allclose(1 - special.gammaincc(v / 2, x / 2), ref) + assert_allclose(1 - special.chdtrc(v, x), ref) + assert_allclose(1 - special.gdtrc(1, v / 2, x / 2), ref) + + +def test_chi2c_smalldf(): + assert_allclose(special.chdtrc(0.6, 3), 1 - 0.957890536704110, + atol=1.5e-7, rtol=0) + + +def test_chi2_inv_smalldf(): + assert_allclose(special.chdtri(0.6, 1 - 0.957890536704110), 3, + atol=1.5e-7, rtol=0) + + +def test_agm_simple(): + rtol = 1e-13 + + # Gauss's constant + assert_allclose(1/special.agm(1, np.sqrt(2)), 0.834626841674073186, + rtol=rtol) + + # These values were computed using Wolfram Alpha, with the + # function ArithmeticGeometricMean[a, b]. + agm13 = 1.863616783244897 + agm15 = 2.604008190530940 + agm35 = 3.936235503649555 + assert_allclose(special.agm([[1], [3]], [1, 3, 5]), + [[1, agm13, agm15], + [agm13, 3, agm35]], rtol=rtol) + + # Computed by the iteration formula using mpmath, + # with mpmath.mp.prec = 1000: + agm12 = 1.4567910310469068 + assert_allclose(special.agm(1, 2), agm12, rtol=rtol) + assert_allclose(special.agm(2, 1), agm12, rtol=rtol) + assert_allclose(special.agm(-1, -2), -agm12, rtol=rtol) + assert_allclose(special.agm(24, 6), 13.458171481725614, rtol=rtol) + assert_allclose(special.agm(13, 123456789.5), 11111458.498599306, + rtol=rtol) + assert_allclose(special.agm(1e30, 1), 2.229223055945383e+28, rtol=rtol) + assert_allclose(special.agm(1e-22, 1), 0.030182566420169886, rtol=rtol) + assert_allclose(special.agm(1e150, 1e180), 2.229223055945383e+178, + rtol=rtol) + assert_allclose(special.agm(1e180, 1e-150), 2.0634722510162677e+177, + rtol=rtol) + assert_allclose(special.agm(1e-150, 1e-170), 3.3112619670463756e-152, + rtol=rtol) + fi = np.finfo(1.0) + assert_allclose(special.agm(fi.tiny, fi.max), 1.9892072050015473e+305, + rtol=rtol) + assert_allclose(special.agm(0.75*fi.max, fi.max), 1.564904312298045e+308, + rtol=rtol) + assert_allclose(special.agm(fi.tiny, 3*fi.tiny), 4.1466849866735005e-308, + rtol=rtol) + + # zero, nan and inf cases. + assert_equal(special.agm(0, 0), 0) + assert_equal(special.agm(99, 0), 0) + + assert_equal(special.agm(-1, 10), np.nan) + assert_equal(special.agm(0, np.inf), np.nan) + assert_equal(special.agm(np.inf, 0), np.nan) + assert_equal(special.agm(0, -np.inf), np.nan) + assert_equal(special.agm(-np.inf, 0), np.nan) + assert_equal(special.agm(np.inf, -np.inf), np.nan) + assert_equal(special.agm(-np.inf, np.inf), np.nan) + assert_equal(special.agm(1, np.nan), np.nan) + assert_equal(special.agm(np.nan, -1), np.nan) + + assert_equal(special.agm(1, np.inf), np.inf) + assert_equal(special.agm(np.inf, 1), np.inf) + assert_equal(special.agm(-1, -np.inf), -np.inf) + assert_equal(special.agm(-np.inf, -1), -np.inf) + + +def test_legacy(): + # Legacy behavior: truncating arguments to integers + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", "floating point number truncated to an integer", RuntimeWarning) + assert_equal(special.expn(1, 0.3), special.expn(1.8, 0.3)) + assert_equal(special.nbdtrc(1, 2, 0.3), special.nbdtrc(1.8, 2.8, 0.3)) + assert_equal(special.nbdtr(1, 2, 0.3), special.nbdtr(1.8, 2.8, 0.3)) + assert_equal(special.nbdtri(1, 2, 0.3), special.nbdtri(1.8, 2.8, 0.3)) + assert_equal(special.pdtri(1, 0.3), special.pdtri(1.8, 0.3)) + assert_equal(special.kn(1, 0.3), special.kn(1.8, 0.3)) + assert_equal(special.yn(1, 0.3), special.yn(1.8, 0.3)) + assert_equal(special.smirnov(1, 0.3), special.smirnov(1.8, 0.3)) + assert_equal(special.smirnovi(1, 0.3), special.smirnovi(1.8, 0.3)) + + +# This lock can be removed once errstate is made thread-safe (see gh-21956) +@pytest.fixture +def errstate_lock(): + import threading + return threading.Lock() + + +@with_special_errors +def test_error_raising(errstate_lock): + with errstate_lock: + with special.errstate(all='raise'): + assert_raises(special.SpecialFunctionError, special.iv, 1, 1e99j) + + +def test_xlogy(): + def xfunc(x, y): + with np.errstate(invalid='ignore'): + if x == 0 and not np.isnan(y): + return x + else: + return x*np.log(y) + + z1 = np.asarray([(0,0), (0, np.nan), (0, np.inf), (1.0, 2.0)], dtype=float) + z2 = np.r_[z1, [(0, 1j), (1, 1j)]] + + w1 = np.vectorize(xfunc)(z1[:,0], z1[:,1]) + assert_func_equal(special.xlogy, w1, z1, rtol=1e-13, atol=1e-13) + w2 = np.vectorize(xfunc)(z2[:,0], z2[:,1]) + assert_func_equal(special.xlogy, w2, z2, rtol=1e-13, atol=1e-13) + + +def test_xlog1py(): + def xfunc(x, y): + with np.errstate(invalid='ignore'): + if x == 0 and not np.isnan(y): + return x + else: + return x * np.log1p(y) + + z1 = np.asarray([(0,0), (0, np.nan), (0, np.inf), (1.0, 2.0), + (1, 1e-30)], dtype=float) + w1 = np.vectorize(xfunc)(z1[:,0], z1[:,1]) + assert_func_equal(special.xlog1py, w1, z1, rtol=1e-13, atol=1e-13) + + +def test_entr(): + def xfunc(x): + if x < 0: + return -np.inf + else: + return -special.xlogy(x, x) + values = (0, 0.5, 1.0, np.inf) + signs = [-1, 1] + arr = [] + for sgn, v in itertools.product(signs, values): + arr.append(sgn * v) + z = np.array(arr, dtype=float) + w = np.vectorize(xfunc, otypes=[np.float64])(z) + assert_func_equal(special.entr, w, z, rtol=1e-13, atol=1e-13) + + +def test_kl_div(): + def xfunc(x, y): + if x < 0 or y < 0 or (y == 0 and x != 0): + # extension of natural domain to preserve convexity + return np.inf + elif np.isposinf(x) or np.isposinf(y): + # limits within the natural domain + return np.inf + elif x == 0: + return y + else: + return special.xlogy(x, x/y) - x + y + values = (0, 0.5, 1.0) + signs = [-1, 1] + arr = [] + for sgna, va, sgnb, vb in itertools.product(signs, values, signs, values): + arr.append((sgna*va, sgnb*vb)) + z = np.array(arr, dtype=float) + w = np.vectorize(xfunc, otypes=[np.float64])(z[:,0], z[:,1]) + assert_func_equal(special.kl_div, w, z, rtol=1e-13, atol=1e-13) + + +def test_rel_entr(): + def xfunc(x, y): + if x > 0 and y > 0: + return special.xlogy(x, x/y) + elif x == 0 and y >= 0: + return 0 + else: + return np.inf + values = (0, 0.5, 1.0) + signs = [-1, 1] + arr = [] + for sgna, va, sgnb, vb in itertools.product(signs, values, signs, values): + arr.append((sgna*va, sgnb*vb)) + z = np.array(arr, dtype=float) + w = np.vectorize(xfunc, otypes=[np.float64])(z[:,0], z[:,1]) + assert_func_equal(special.rel_entr, w, z, rtol=1e-13, atol=1e-13) + + +def test_rel_entr_gh_20710_near_zero(): + # Check accuracy of inputs which are very close + inputs = np.array([ + # x, y + (0.9456657713430001, 0.9456657713430094), + (0.48066098564791515, 0.48066098564794774), + (0.786048657854401, 0.7860486578542367), + ]) + # Known values produced using `x * mpmath.log(x / y)` with dps=30 + expected = [ + -9.325873406851269e-15, + -3.258504577274724e-14, + 1.6431300764454033e-13, + ] + x = inputs[:, 0] + y = inputs[:, 1] + assert_allclose(special.rel_entr(x, y), expected, rtol=1e-13, atol=0) + + +def test_rel_entr_gh_20710_overflow(): + special.seterr(all='ignore') + inputs = np.array([ + # x, y + # Overflow + (4, 2.22e-308), + # Underflow + (1e-200, 1e+200), + # Subnormal + (2.22e-308, 1e15), + ]) + # Known values produced using `x * mpmath.log(x / y)` with dps=30 + expected = [ + 2839.139983229607, + -9.210340371976183e-198, + -1.6493212008074475e-305, + ] + x = inputs[:, 0] + y = inputs[:, 1] + assert_allclose(special.rel_entr(x, y), expected, rtol=1e-13, atol=0) + + +def test_huber(): + assert_equal(special.huber(-1, 1.5), np.inf) + assert_allclose(special.huber(2, 1.5), 0.5 * np.square(1.5)) + assert_allclose(special.huber(2, 2.5), 2 * (2.5 - 0.5 * 2)) + + def xfunc(delta, r): + if delta < 0: + return np.inf + elif np.abs(r) < delta: + return 0.5 * np.square(r) + else: + return delta * (np.abs(r) - 0.5 * delta) + + z = np.random.randn(10, 2) + w = np.vectorize(xfunc, otypes=[np.float64])(z[:,0], z[:,1]) + assert_func_equal(special.huber, w, z, rtol=1e-13, atol=1e-13) + + +def test_pseudo_huber(): + def xfunc(delta, r): + if delta < 0: + return np.inf + elif (not delta) or (not r): + return 0 + else: + return delta**2 * (np.sqrt(1 + (r/delta)**2) - 1) + + z = np.array(np.random.randn(10, 2).tolist() + [[0, 0.5], [0.5, 0]]) + w = np.vectorize(xfunc, otypes=[np.float64])(z[:,0], z[:,1]) + assert_func_equal(special.pseudo_huber, w, z, rtol=1e-13, atol=1e-13) + + +def test_pseudo_huber_small_r(): + delta = 1.0 + r = 1e-18 + y = special.pseudo_huber(delta, r) + # expected computed with mpmath: + # import mpmath + # mpmath.mp.dps = 200 + # r = mpmath.mpf(1e-18) + # expected = float(mpmath.sqrt(1 + r**2) - 1) + expected = 5.0000000000000005e-37 + assert_allclose(y, expected, rtol=1e-13) + + +def test_runtime_warning(): + with pytest.warns(RuntimeWarning, + match=r'Too many predicted coefficients'): + mathieu_odd_coef(1000, 1000) + with pytest.warns(RuntimeWarning, + match=r'Too many predicted coefficients'): + mathieu_even_coef(1000, 1000) + + +class TestStirling2: + table = [ + [1], + [0, 1], + [0, 1, 1], + [0, 1, 3, 1], + [0, 1, 7, 6, 1], + [0, 1, 15, 25, 10, 1], + [0, 1, 31, 90, 65, 15, 1], + [0, 1, 63, 301, 350, 140, 21, 1], + [0, 1, 127, 966, 1701, 1050, 266, 28, 1], + [0, 1, 255, 3025, 7770, 6951, 2646, 462, 36, 1], + [0, 1, 511, 9330, 34105, 42525, 22827, 5880, 750, 45, 1], + ] + + @pytest.mark.parametrize("is_exact, comp, kwargs", [ + (True, assert_equal, {}), + (False, assert_allclose, {'rtol': 1e-12}) + ]) + def test_table_cases(self, is_exact, comp, kwargs): + for n in range(1, len(self.table)): + k_values = list(range(n+1)) + row = self.table[n] + comp(row, stirling2([n], k_values, exact=is_exact), **kwargs) + + @pytest.mark.parametrize("is_exact, comp, kwargs", [ + (True, assert_equal, {}), + (False, assert_allclose, {'rtol': 1e-12}) + ]) + def test_valid_single_integer(self, is_exact, comp, kwargs): + comp(stirling2(0, 0, exact=is_exact), self.table[0][0], **kwargs) + comp(stirling2(4, 2, exact=is_exact), self.table[4][2], **kwargs) + # a single 2-tuple of integers as arguments must return an int and not + # an array whereas arrays of single values should return array + comp(stirling2(5, 3, exact=is_exact), 25, **kwargs) + comp(stirling2([5], [3], exact=is_exact), [25], **kwargs) + + @pytest.mark.parametrize("is_exact, comp, kwargs", [ + (True, assert_equal, {}), + (False, assert_allclose, {'rtol': 1e-12}) + ]) + def test_negative_integer(self, is_exact, comp, kwargs): + # negative integers for n or k arguments return 0 + comp(stirling2(-1, -1, exact=is_exact), 0, **kwargs) + comp(stirling2(-1, 2, exact=is_exact), 0, **kwargs) + comp(stirling2(2, -1, exact=is_exact), 0, **kwargs) + + @pytest.mark.parametrize("is_exact, comp, kwargs", [ + (True, assert_equal, {}), + (False, assert_allclose, {'rtol': 1e-12}) + ]) + def test_array_inputs(self, is_exact, comp, kwargs): + ans = [self.table[10][3], self.table[10][4]] + comp(stirling2(asarray([10, 10]), + asarray([3, 4]), + exact=is_exact), + ans) + comp(stirling2([10, 10], + asarray([3, 4]), + exact=is_exact), + ans) + comp(stirling2(asarray([10, 10]), + [3, 4], + exact=is_exact), + ans) + + @pytest.mark.parametrize("is_exact, comp, kwargs", [ + (True, assert_equal, {}), + (False, assert_allclose, {'rtol': 1e-13}) + ]) + def test_mixed_values(self, is_exact, comp, kwargs): + # negative values-of either n or k-should return 0 for the entry + ans = [0, 1, 3, 25, 1050, 5880, 9330] + n = [-1, 0, 3, 5, 8, 10, 10] + k = [-2, 0, 2, 3, 5, 7, 3] + comp(stirling2(n, k, exact=is_exact), ans, **kwargs) + + def test_correct_parity(self): + """Test parity follows well known identity. + + en.wikipedia.org/wiki/Stirling_numbers_of_the_second_kind#Parity + """ + n, K = 100, np.arange(101) + assert_equal( + stirling2(n, K, exact=True) % 2, + [math.comb(n - (k // 2) - 1, n - k) % 2 for k in K], + ) + + def test_big_numbers(self): + # via mpmath (bigger than 32bit) + ans = asarray([48063331393110, 48004081105038305]) + n = [25, 30] + k = [17, 4] + assert array_equal(stirling2(n, k, exact=True), ans) + # bigger than 64 bit + ans = asarray([2801934359500572414253157841233849412, + 14245032222277144547280648984426251]) + n = [42, 43] + k = [17, 23] + assert array_equal(stirling2(n, k, exact=True), ans) + + @pytest.mark.parametrize("N", [4.5, 3., 4+1j, "12", np.nan]) + @pytest.mark.parametrize("K", [3.5, 3, "2", None]) + @pytest.mark.parametrize("is_exact", [True, False]) + def test_unsupported_input_types(self, N, K, is_exact): + # object, float, string, complex are not supported and raise TypeError + with pytest.raises(TypeError): + stirling2(N, K, exact=is_exact) + + @pytest.mark.parametrize("is_exact", [True, False]) + def test_numpy_array_int_object_dtype(self, is_exact): + # python integers with arbitrary precision are *not* allowed as + # object type in numpy arrays are inconsistent from api perspective + ans = asarray(self.table[4][1:]) + n = asarray([4, 4, 4, 4], dtype=object) + k = asarray([1, 2, 3, 4], dtype=object) + with pytest.raises(TypeError): + array_equal(stirling2(n, k, exact=is_exact), ans) + + @pytest.mark.parametrize("is_exact, comp, kwargs", [ + (True, assert_equal, {}), + (False, assert_allclose, {'rtol': 1e-13}) + ]) + def test_numpy_array_unsigned_int_dtype(self, is_exact, comp, kwargs): + # numpy unsigned integers are allowed as dtype in numpy arrays + ans = asarray(self.table[4][1:]) + n = asarray([4, 4, 4, 4], dtype=np.ulong) + k = asarray([1, 2, 3, 4], dtype=np.ulong) + comp(stirling2(n, k, exact=False), ans, **kwargs) + + @pytest.mark.parametrize("is_exact, comp, kwargs", [ + (True, assert_equal, {}), + (False, assert_allclose, {'rtol': 1e-13}) + ]) + def test_broadcasting_arrays_correctly(self, is_exact, comp, kwargs): + # broadcasting is handled by stirling2 + # test leading 1s are replicated + ans = asarray([[1, 15, 25, 10], [1, 7, 6, 1]]) # shape (2,4) + n = asarray([[5, 5, 5, 5], [4, 4, 4, 4]]) # shape (2,4) + k = asarray([1, 2, 3, 4]) # shape (4,) + comp(stirling2(n, k, exact=is_exact), ans, **kwargs) + # test that dims both mismatch broadcast correctly (5,1) & (6,) + n = asarray([[4], [4], [4], [4], [4]]) + k = asarray([0, 1, 2, 3, 4, 5]) + ans = asarray([[0, 1, 7, 6, 1, 0] for _ in range(5)]) + comp(stirling2(n, k, exact=False), ans, **kwargs) + + def test_temme_rel_max_error(self): + # python integers with arbitrary precision are *not* allowed as + # object type in numpy arrays are inconsistent from api perspective + x = list(range(51, 101, 5)) + for n in x: + k_entries = list(range(1, n+1)) + denom = stirling2([n], k_entries, exact=True) + num = denom - stirling2([n], k_entries, exact=False) + assert np.max(np.abs(num / denom)) < 2e-5 diff --git a/v12x/pylib/scipy/special/tests/test_bdtr.py b/v12x/pylib/scipy/special/tests/test_bdtr.py new file mode 100644 index 0000000000000000000000000000000000000000..4d59c1e8c594166a99b213a217b7851c292d9d64 --- /dev/null +++ b/v12x/pylib/scipy/special/tests/test_bdtr.py @@ -0,0 +1,114 @@ +import warnings + +import numpy as np +import scipy.special as sc +import pytest +from numpy.testing import assert_allclose, assert_array_equal + + +class TestBdtr: + def test(self): + val = sc.bdtr(0, 1, 0.5) + assert_allclose(val, 0.5) + + def test_sum_is_one(self): + val = sc.bdtr([0, 1, 2], 2, 0.5) + assert_array_equal(val, [0.25, 0.75, 1.0]) + + def test_rounding(self): + double_val = sc.bdtr([0.1, 1.1, 2.1], 2, 0.5) + int_val = sc.bdtr([0, 1, 2], 2, 0.5) + assert_array_equal(double_val, int_val) + + @pytest.mark.parametrize('k, n, p', [ + (np.inf, 2, 0.5), + (1.0, np.inf, 0.5), + (1.0, 2, np.inf) + ]) + def test_inf(self, k, n, p): + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + val = sc.bdtr(k, n, p) + assert np.isnan(val) + + def test_domain(self): + val = sc.bdtr(-1.1, 1, 0.5) + assert np.isnan(val) + + +class TestBdtrc: + def test_value(self): + val = sc.bdtrc(0, 1, 0.5) + assert_allclose(val, 0.5) + + def test_sum_is_one(self): + val = sc.bdtrc([0, 1, 2], 2, 0.5) + assert_array_equal(val, [0.75, 0.25, 0.0]) + + def test_rounding(self): + double_val = sc.bdtrc([0.1, 1.1, 2.1], 2, 0.5) + int_val = sc.bdtrc([0, 1, 2], 2, 0.5) + assert_array_equal(double_val, int_val) + + @pytest.mark.parametrize('k, n, p', [ + (np.inf, 2, 0.5), + (1.0, np.inf, 0.5), + (1.0, 2, np.inf) + ]) + def test_inf(self, k, n, p): + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + val = sc.bdtrc(k, n, p) + assert np.isnan(val) + + def test_domain(self): + val = sc.bdtrc(-1.1, 1, 0.5) + val2 = sc.bdtrc(2.1, 1, 0.5) + assert np.isnan(val2) + assert_allclose(val, 1.0) + + def test_bdtr_bdtrc_sum_to_one(self): + bdtr_vals = sc.bdtr([0, 1, 2], 2, 0.5) + bdtrc_vals = sc.bdtrc([0, 1, 2], 2, 0.5) + vals = bdtr_vals + bdtrc_vals + assert_allclose(vals, [1.0, 1.0, 1.0]) + + +class TestBdtri: + def test_value(self): + val = sc.bdtri(0, 1, 0.5) + assert_allclose(val, 0.5) + + def test_sum_is_one(self): + val = sc.bdtri([0, 1], 2, 0.5) + actual = np.asarray([1 - 1/np.sqrt(2), 1/np.sqrt(2)]) + assert_allclose(val, actual) + + def test_rounding(self): + double_val = sc.bdtri([0.1, 1.1], 2, 0.5) + int_val = sc.bdtri([0, 1], 2, 0.5) + assert_allclose(double_val, int_val) + + @pytest.mark.parametrize('k, n, p', [ + (np.inf, 2, 0.5), + (1.0, np.inf, 0.5), + (1.0, 2, np.inf) + ]) + def test_inf(self, k, n, p): + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + val = sc.bdtri(k, n, p) + assert np.isnan(val) + + @pytest.mark.parametrize('k, n, p', [ + (-1.1, 1, 0.5), + (2.1, 1, 0.5) + ]) + def test_domain(self, k, n, p): + val = sc.bdtri(k, n, p) + assert np.isnan(val) + + def test_bdtr_bdtri_roundtrip(self): + bdtr_vals = sc.bdtr([0, 1, 2], 2, 0.5) + roundtrip_vals = sc.bdtri([0, 1, 2], 2, bdtr_vals) + assert_allclose(roundtrip_vals, [0.5, 0.5, np.nan]) diff --git a/v12x/pylib/scipy/special/tests/test_boost_ufuncs.py b/v12x/pylib/scipy/special/tests/test_boost_ufuncs.py new file mode 100644 index 0000000000000000000000000000000000000000..e782765da6cf26d10317492c8da3e2f8b1b2b07f --- /dev/null +++ b/v12x/pylib/scipy/special/tests/test_boost_ufuncs.py @@ -0,0 +1,83 @@ +import pytest +import numpy as np +from numpy.testing import assert_allclose +import scipy.special._ufuncs as scu +from scipy.integrate import tanhsinh + + +type_char_to_type_tol = {'f': (np.float32, 32*np.finfo(np.float32).eps), + 'd': (np.float64, 32*np.finfo(np.float64).eps)} + + +# Each item in this list is +# (func, args, expected_value) +# All the values can be represented exactly, even with np.float32. +# +# This is not an exhaustive test data set of all the functions! +# It is a spot check of several functions, primarily for +# checking that the different data types are handled correctly. +test_data = [ + (scu._beta_pdf, (0.5, 2, 3), 1.5), + (scu._beta_pdf, (0, 1, 5), 5.0), + (scu._beta_pdf, (1, 5, 1), 5.0), + (scu._beta_ppf, (0.5, 5., 5.), 0.5), # gh-21303 + (scu._binom_cdf, (1, 3, 0.5), 0.5), + (scu._binom_pmf, (1, 4, 0.5), 0.25), + (scu._hypergeom_cdf, (2, 3, 5, 6), 0.5), + (scu._nbinom_cdf, (1, 4, 0.25), 0.015625), + (scu._ncf_mean, (10, 12, 2.5), 1.5), +] + + +@pytest.mark.parametrize('func, args, expected', test_data) +def test_stats_boost_ufunc(func, args, expected): + type_sigs = func.types + type_chars = [sig.split('->')[-1] for sig in type_sigs] + for type_char in type_chars: + typ, rtol = type_char_to_type_tol[type_char] + args = [typ(arg) for arg in args] + # Harmless overflow warnings are a "feature" of some wrappers on some + # platforms. This test is about dtype and accuracy, so let's avoid false + # test failures cause by these warnings. See gh-17432. + with np.errstate(over='ignore'): + value = func(*args) + assert isinstance(value, typ) + assert_allclose(value, expected, rtol=rtol) + + +def test_landau(): + # Test that Landau distribution ufuncs are wrapped as expected; + # accuracy is tested by Boost. + x = np.linspace(-3, 10, 10) + args = (0, 1) + res = tanhsinh(lambda x: scu._landau_pdf(x, *args), -np.inf, x) + cdf = scu._landau_cdf(x, *args) + assert_allclose(res.integral, cdf) + sf = scu._landau_sf(x, *args) + assert_allclose(sf, 1-cdf) + ppf = scu._landau_ppf(cdf, *args) + assert_allclose(ppf, x) + isf = scu._landau_isf(sf, *args) + assert_allclose(isf, x, rtol=1e-6) + +def test_gh22956(): + _ = scu._ncx2_pdf(30, 1e307, 16) + +@pytest.mark.parametrize("func", [scu._binom_cdf, scu._binom_sf]) +@pytest.mark.parametrize("dtype", [np.float32, np.float64]) +def test_extreme_inputs_for_binomial_probabilities(func, dtype): + # certain inputs caused C++ exceptions in boost + # resulting in Python interpreter crashes + k = 3e18 + n = 10e18 + p = 0.3 + func(dtype(k), dtype(n), dtype(p)) + +@pytest.mark.parametrize("func", [scu._binom_ppf, scu._binom_isf]) +@pytest.mark.parametrize("dtype", [np.float32, np.float64]) +def test_extreme_inputs_for_binomial_quantiles(func, dtype): + # certain inputs caused C++ exceptions in boost + # resulting in Python interpreter crashes + n = 10e18 + p = 0.5 + func(dtype(p), dtype(n), dtype(p)) diff --git a/v12x/pylib/scipy/special/tests/test_boxcox.py b/v12x/pylib/scipy/special/tests/test_boxcox.py new file mode 100644 index 0000000000000000000000000000000000000000..703a1fdee635a484fe53f14e98fb3ce08c811447 --- /dev/null +++ b/v12x/pylib/scipy/special/tests/test_boxcox.py @@ -0,0 +1,125 @@ +import numpy as np +from numpy.testing import assert_equal, assert_allclose +from scipy.special import boxcox, boxcox1p, inv_boxcox, inv_boxcox1p +import pytest + + +# There are more tests of boxcox and boxcox1p in test_mpmath.py. + +def test_boxcox_basic(): + x = np.array([0.5, 1, 2, 4]) + + # lambda = 0 => y = log(x) + y = boxcox(x, 0) + assert_allclose(y, np.log(x), atol=1.5e-7, rtol=0) + + # lambda = 1 => y = x - 1 + y = boxcox(x, 1) + assert_allclose(y, x - 1, atol=1.5e-7, rtol=0) + + # lambda = 2 => y = 0.5*(x**2 - 1) + y = boxcox(x, 2) + assert_allclose(y, 0.5*(x**2 - 1), atol=1.5e-7, rtol=0) + + # x = 0 and lambda > 0 => y = -1 / lambda + lam = np.array([0.5, 1, 2]) + y = boxcox(0, lam) + assert_allclose(y, -1.0 / lam, atol=1.5e-7, rtol=0) + +def test_boxcox_underflow(): + x = 1 + 1e-15 + lmbda = 1e-306 + y = boxcox(x, lmbda) + assert_allclose(y, np.log(x), rtol=1e-14) + + +def test_boxcox_nonfinite(): + # x < 0 => y = nan + x = np.array([-1, -1, -0.5]) + y = boxcox(x, [0.5, 2.0, -1.5]) + assert_equal(y, np.array([np.nan, np.nan, np.nan])) + + # x = 0 and lambda <= 0 => y = -inf + x = 0 + y = boxcox(x, [-2.5, 0]) + assert_equal(y, np.array([-np.inf, -np.inf])) + + +def test_boxcox1p_basic(): + x = np.array([-0.25, -1e-20, 0, 1e-20, 0.25, 1, 3]) + + # lambda = 0 => y = log(1+x) + y = boxcox1p(x, 0) + assert_allclose(y, np.log1p(x), atol=1.5e-7, rtol=0) + + # lambda = 1 => y = x + y = boxcox1p(x, 1) + assert_allclose(y, x, atol=1.5e-7, rtol=0) + + # lambda = 2 => y = 0.5*((1+x)**2 - 1) = 0.5*x*(2 + x) + y = boxcox1p(x, 2) + assert_allclose(y, 0.5*x*(2 + x), atol=1.5e-7, rtol=0) + + # x = -1 and lambda > 0 => y = -1 / lambda + lam = np.array([0.5, 1, 2]) + y = boxcox1p(-1, lam) + assert_allclose(y, -1.0 / lam, atol=1.5e-7, rtol=0) + + +def test_boxcox1p_underflow(): + x = np.array([1e-15, 1e-306]) + lmbda = np.array([1e-306, 1e-18]) + y = boxcox1p(x, lmbda) + assert_allclose(y, np.log1p(x), rtol=1e-14) + + +def test_boxcox1p_nonfinite(): + # x < -1 => y = nan + x = np.array([-2, -2, -1.5]) + y = boxcox1p(x, [0.5, 2.0, -1.5]) + assert_equal(y, np.array([np.nan, np.nan, np.nan])) + + # x = -1 and lambda <= 0 => y = -inf + x = -1 + y = boxcox1p(x, [-2.5, 0]) + assert_equal(y, np.array([-np.inf, -np.inf])) + + +def test_inv_boxcox(): + x = np.array([0., 1., 2.]) + lam = np.array([0., 1., 2.]) + y = boxcox(x, lam) + x2 = inv_boxcox(y, lam) + assert_allclose(x, x2, atol=1.5e-7, rtol=0) + + x = np.array([0., 1., 2.]) + lam = np.array([0., 1., 2.]) + y = boxcox1p(x, lam) + x2 = inv_boxcox1p(y, lam) + assert_allclose(x, x2, atol=1.5e-7, rtol=0) + + +def test_inv_boxcox1p_underflow(): + x = 1e-15 + lam = 1e-306 + y = inv_boxcox1p(x, lam) + assert_allclose(y, x, rtol=1e-14) + + +@pytest.mark.parametrize( + "x, lmb", + [[100, 155], + [0.01, -155]] +) +def test_boxcox_premature_overflow(x, lmb): + # test boxcox & inv_boxcox + y = boxcox(x, lmb) + assert np.isfinite(y) + x_inv = inv_boxcox(y, lmb) + assert_allclose(x, x_inv) + + # test boxcox1p & inv_boxcox1p + y1p = boxcox1p(x-1, lmb) + assert np.isfinite(y1p) + x1p_inv = inv_boxcox1p(y1p, lmb) + assert_allclose(x-1, x1p_inv) diff --git a/v12x/pylib/scipy/special/tests/test_cdflib.py b/v12x/pylib/scipy/special/tests/test_cdflib.py new file mode 100644 index 0000000000000000000000000000000000000000..48938605349829a024f1b9bff5b7af96e57cec6a --- /dev/null +++ b/v12x/pylib/scipy/special/tests/test_cdflib.py @@ -0,0 +1,980 @@ +""" +Test cdflib functions versus mpmath, if available. + +The following functions still need tests: + +- ncfdtridfn +- ncfdtridfd +- ncfdtrinc +- nctdtridf +- nctdtrinc + +""" +import itertools + +import numpy as np +from numpy.testing import assert_equal, assert_allclose, assert_array_equal +import pytest + +import scipy.special as sp +from scipy.special._testutils import ( + MissingModule, check_version, FuncData) +from scipy.special._mptestutils import ( + Arg, IntArg, get_args, mpf2float, assert_mpmath_equal) + +try: + import mpmath +except ImportError: + mpmath = MissingModule('mpmath') + + +class ProbArg: + """Generate a set of probabilities on [0, 1].""" + + def __init__(self): + # Include the endpoints for compatibility with Arg et. al. + self.a = 0 + self.b = 1 + + def values(self, n): + """Return an array containing approximately n numbers.""" + m = max(1, n//3) + v1 = np.logspace(-30, np.log10(0.3), m) + v2 = np.linspace(0.3, 0.7, m + 1, endpoint=False)[1:] + v3 = 1 - np.logspace(np.log10(0.3), -15, m) + v = np.r_[v1, v2, v3] + return np.unique(v) + + +class EndpointFilter: + def __init__(self, a, b, rtol, atol): + self.a = a + self.b = b + self.rtol = rtol + self.atol = atol + + def __call__(self, x): + mask1 = np.abs(x - self.a) < self.rtol*np.abs(self.a) + self.atol + mask2 = np.abs(x - self.b) < self.rtol*np.abs(self.b) + self.atol + return np.where(mask1 | mask2, False, True) + + +class _CDFData: + def __init__(self, spfunc, mpfunc, index, argspec, spfunc_first=True, + dps=20, n=5000, rtol=None, atol=None, + endpt_rtol=None, endpt_atol=None): + self.spfunc = spfunc + self.mpfunc = mpfunc + self.index = index + self.argspec = argspec + self.spfunc_first = spfunc_first + self.dps = dps + self.n = n + self.rtol = rtol + self.atol = atol + + if not isinstance(argspec, list): + self.endpt_rtol = None + self.endpt_atol = None + elif endpt_rtol is not None or endpt_atol is not None: + if isinstance(endpt_rtol, list): + self.endpt_rtol = endpt_rtol + else: + self.endpt_rtol = [endpt_rtol]*len(self.argspec) + if isinstance(endpt_atol, list): + self.endpt_atol = endpt_atol + else: + self.endpt_atol = [endpt_atol]*len(self.argspec) + else: + self.endpt_rtol = None + self.endpt_atol = None + + def idmap(self, *args): + if self.spfunc_first: + res = self.spfunc(*args) + if np.isnan(res): + return np.nan + args = list(args) + args[self.index] = res + with mpmath.workdps(self.dps): + res = self.mpfunc(*tuple(args)) + # Imaginary parts are spurious + res = mpf2float(res.real) + else: + with mpmath.workdps(self.dps): + res = self.mpfunc(*args) + res = mpf2float(res.real) + args = list(args) + args[self.index] = res + res = self.spfunc(*tuple(args)) + return res + + def get_param_filter(self): + if self.endpt_rtol is None and self.endpt_atol is None: + return None + + filters = [] + for rtol, atol, spec in zip(self.endpt_rtol, self.endpt_atol, self.argspec): + if rtol is None and atol is None: + filters.append(None) + continue + elif rtol is None: + rtol = 0.0 + elif atol is None: + atol = 0.0 + + filters.append(EndpointFilter(spec.a, spec.b, rtol, atol)) + return filters + + def check(self): + # Generate values for the arguments + args = get_args(self.argspec, self.n) + param_filter = self.get_param_filter() + param_columns = tuple(range(args.shape[1])) + result_columns = args.shape[1] + args = np.hstack((args, args[:, self.index].reshape(args.shape[0], 1))) + FuncData(self.idmap, args, + param_columns=param_columns, result_columns=result_columns, + rtol=self.rtol, atol=self.atol, vectorized=False, + param_filter=param_filter).check() + + +def _assert_inverts(*a, **kw): + d = _CDFData(*a, **kw) + d.check() + + +def _binomial_cdf(k, n, p): + k, n, p = mpmath.mpf(k), mpmath.mpf(n), mpmath.mpf(p) + if k <= 0: + return mpmath.mpf(0) + elif k >= n: + return mpmath.mpf(1) + + onemp = mpmath.fsub(1, p, exact=True) + return mpmath.betainc(n - k, k + 1, x2=onemp, regularized=True) + + +def _f_cdf(dfn, dfd, x): + if x < 0: + return mpmath.mpf(0) + dfn, dfd, x = mpmath.mpf(dfn), mpmath.mpf(dfd), mpmath.mpf(x) + ub = dfn*x/(dfn*x + dfd) + res = mpmath.betainc(dfn/2, dfd/2, x2=ub, regularized=True) + return res + + +def _student_t_cdf(df, t, dps=None): + if dps is None: + dps = mpmath.mp.dps + with mpmath.workdps(dps): + df, t = mpmath.mpf(df), mpmath.mpf(t) + fac = mpmath.hyp2f1(0.5, 0.5*(df + 1), 1.5, -t**2/df) + fac *= t*mpmath.gamma(0.5*(df + 1)) + fac /= mpmath.sqrt(mpmath.pi*df)*mpmath.gamma(0.5*df) + return 0.5 + fac + + +def _noncentral_chi_pdf(t, df, nc): + res = mpmath.besseli(df/2 - 1, mpmath.sqrt(nc*t)) + res *= mpmath.exp(-(t + nc)/2)*(t/nc)**(df/4 - 1/2)/2 + return res + + +def _noncentral_chi_cdf(x, df, nc, dps=None): + if dps is None: + dps = mpmath.mp.dps + x, df, nc = mpmath.mpf(x), mpmath.mpf(df), mpmath.mpf(nc) + with mpmath.workdps(dps): + res = mpmath.quad(lambda t: _noncentral_chi_pdf(t, df, nc), [0, x]) + return res + + +def _tukey_lmbda_quantile(p, lmbda): + # For lmbda != 0 + return (p**lmbda - (1 - p)**lmbda)/lmbda + + +@pytest.mark.slow +@check_version(mpmath, '0.19') +class TestCDFlib: + + def test_bdtrin(self): + _assert_inverts( + sp.bdtrin, + _binomial_cdf, + 1, [IntArg(1, 3000), ProbArg(), ProbArg()], + rtol=1e-12, endpt_atol=[None, None, 1e-6]) + + def test_btdtria(self): + _assert_inverts( + sp.btdtria, + lambda a, b, x: mpmath.betainc(a, b, x2=x, regularized=True), + 0, [ProbArg(), Arg(0, 3e2, inclusive_a=False), + Arg(0, 1, inclusive_a=False, inclusive_b=False)], + rtol=1e-12) + + def test_btdtrib(self): + # Use small values of a or mpmath doesn't converge + _assert_inverts( + sp.btdtrib, + lambda a, b, x: mpmath.betainc(a, b, x2=x, regularized=True), + 1, + [Arg(0, 1e2, inclusive_a=False), ProbArg(), + Arg(0, 1, inclusive_a=False, inclusive_b=False)], + rtol=1e-7, + endpt_atol=[None, 1e-18, 1e-15]) + + @pytest.mark.xfail(run=False) + def test_fdtridfd(self): + _assert_inverts( + sp.fdtridfd, + _f_cdf, + 1, + [IntArg(1, 100), ProbArg(), Arg(0, 100, inclusive_a=False)], + rtol=1e-7) + + def test_gdtria(self): + _assert_inverts( + sp.gdtria, + lambda a, b, x: mpmath.gammainc(b, b=a*x, regularized=True), + 0, + [ProbArg(), Arg(0, 1e3, inclusive_a=False), + Arg(0, 1e4, inclusive_a=False)], + rtol=1e-12, + endpt_atol=[None, 1e-7, 1e-10]) + + def test_gdtrib(self): + # Use small values of a and x or mpmath doesn't converge + _assert_inverts( + sp.gdtrib, + lambda a, b, x: mpmath.gammainc(b, b=a*x, regularized=True), + 1, + [Arg(0, 1e2, inclusive_a=False), ProbArg(), + Arg(0, 1e3, inclusive_a=False)], + rtol=1e-5) + + def test_gdtrix(self): + _assert_inverts( + sp.gdtrix, + lambda a, b, x: mpmath.gammainc(b, b=a*x, regularized=True), + 2, + [Arg(0, 1e3, inclusive_a=False), Arg(0, 1e3, inclusive_a=False), + ProbArg()], + rtol=1e-7, + endpt_atol=[None, 1e-7, 1e-10]) + + # Overall nrdtrimn and nrdtrisd are not performing well with infeasible/edge + # combinations of sigma and x, hence restricted the domains to still use the + # testing machinery, also see gh-20069 + + # nrdtrimn signature: p, sd, x + # nrdtrisd signature: mn, p, x + def test_nrdtrimn(self): + _assert_inverts( + sp.nrdtrimn, + lambda x, y, z: mpmath.ncdf(z, x, y), + 0, + [ProbArg(), # CDF value p + Arg(0.1, np.inf, inclusive_a=False, inclusive_b=False), # sigma + Arg(-1e10, 1e10)], # x + rtol=1e-5) + + def test_nrdtrisd(self): + _assert_inverts( + sp.nrdtrisd, + lambda x, y, z: mpmath.ncdf(z, x, y), + 1, + [Arg(-np.inf, 10, inclusive_a=False, inclusive_b=False), # mn + ProbArg(), # CDF value p + Arg(10, 1e100)], # x + rtol=1e-5) + + def test_stdtr(self): + # Ideally the left endpoint for Arg() should be 0. + assert_mpmath_equal( + sp.stdtr, + _student_t_cdf, + [IntArg(1, 100), Arg(1e-10, np.inf)], rtol=1e-7) + + @pytest.mark.xfail(run=False) + def test_stdtridf(self): + _assert_inverts( + sp.stdtridf, + _student_t_cdf, + 0, [ProbArg(), Arg()], rtol=1e-7) + + def test_stdtrit(self): + _assert_inverts( + sp.stdtrit, + _student_t_cdf, + 1, [IntArg(1, 100), ProbArg()], rtol=1e-7, + endpt_atol=[None, 1e-10]) + + def test_chdtriv(self): + _assert_inverts( + sp.chdtriv, + lambda v, x: mpmath.gammainc(v/2, b=x/2, regularized=True), + 0, [ProbArg(), IntArg(1, 100)], rtol=1e-4) + + + @pytest.mark.xfail(run=False) + def test_chndtridf(self): + # Use a larger atol since mpmath is doing numerical integration + _assert_inverts( + sp.chndtridf, + _noncentral_chi_cdf, + 1, [Arg(0, 100, inclusive_a=False), ProbArg(), + Arg(0, 100, inclusive_a=False)], + n=1000, rtol=1e-4, atol=1e-15) + + @pytest.mark.xfail(run=False) + def test_chndtrinc(self): + # Use a larger atol since mpmath is doing numerical integration + _assert_inverts( + sp.chndtrinc, + _noncentral_chi_cdf, + 2, [Arg(0, 100, inclusive_a=False), IntArg(1, 100), ProbArg()], + n=1000, rtol=1e-4, atol=1e-15) + + def test_chndtrix(self): + # Use a larger atol since mpmath is doing numerical integration + _assert_inverts( + sp.chndtrix, + _noncentral_chi_cdf, + 0, [ProbArg(), IntArg(1, 100), Arg(0, 100, inclusive_a=False)], + n=1000, rtol=1e-4, atol=1e-15, + endpt_atol=[1e-6, None, None]) + + def test_tklmbda_zero_shape(self): + # When lmbda = 0 the CDF has a simple closed form + one = mpmath.mpf(1) + assert_mpmath_equal( + lambda x: sp.tklmbda(x, 0), + lambda x: one/(mpmath.exp(-x) + one), + [Arg()], rtol=1e-7) + + def test_tklmbda_neg_shape(self): + _assert_inverts( + sp.tklmbda, + _tukey_lmbda_quantile, + 0, [ProbArg(), Arg(-25, 0, inclusive_b=False)], + spfunc_first=False, rtol=1e-5, + endpt_atol=[1e-9, 1e-5]) + + @pytest.mark.xfail(run=False) + def test_tklmbda_pos_shape(self): + _assert_inverts( + sp.tklmbda, + _tukey_lmbda_quantile, + 0, [ProbArg(), Arg(0, 100, inclusive_a=False)], + spfunc_first=False, rtol=1e-5) + + # The values of lmdba are chosen so that 1/lmbda is exact. + @pytest.mark.parametrize('lmbda', [0.5, 1.0, 8.0]) + def test_tklmbda_lmbda1(self, lmbda): + bound = 1/lmbda + assert_equal(sp.tklmbda([-bound, bound], lmbda), [0.0, 1.0]) + + +funcs = [ + ("btdtria", 3), + ("btdtrib", 3), + ("bdtrik", 3), + ("bdtrin", 3), + ("chdtriv", 2), + ("chndtr", 3), + ("chndtrix", 3), + ("chndtridf", 3), + ("chndtrinc", 3), + ("fdtridfd", 3), + ("ncfdtr", 4), + ("ncfdtri", 4), + ("ncfdtridfn", 4), + ("ncfdtridfd", 4), + ("ncfdtrinc", 4), + ("gdtrix", 3), + ("gdtrib", 3), + ("gdtria", 3), + ("nbdtrik", 3), + ("nbdtrin", 3), + ("nrdtrimn", 3), + ("nrdtrisd", 3), + ("pdtrik", 2), + ("stdtr", 2), + ("stdtrit", 2), + ("stdtridf", 2), + ("nctdtr", 3), + ("nctdtrit", 3), + ("nctdtridf", 3), + ("nctdtrinc", 3), + ("tklmbda", 2), +] + + +@pytest.mark.parametrize('func,numargs', funcs, ids=[x[0] for x in funcs]) +def test_nonfinite(func, numargs): + + rng = np.random.default_rng(1701299355559735) + func = getattr(sp, func) + args_choices = [(float(x), np.nan, np.inf, -np.inf) for x in rng.random(numargs)] + + for args in itertools.product(*args_choices): + res = func(*args) + + if any(np.isnan(x) for x in args): + # Nan inputs should result to nan output + assert_equal(res, np.nan) + else: + # All other inputs should return something (but not + # raise exceptions or cause hangs) + pass + + +def test_chndtrix_gh2158(): + # test that gh-2158 is resolved; previously this blew up + res = sp.chndtrix(0.999999, 2, np.arange(20.)+1e-6) + + # Generated in R + # options(digits=16) + # ncp <- seq(0, 19) + 1e-6 + # print(qchisq(0.999999, df = 2, ncp = ncp)) + res_exp = [27.63103493142305, 35.25728589950540, 39.97396073236288, + 43.88033702110538, 47.35206403482798, 50.54112500166103, + 53.52720257322766, 56.35830042867810, 59.06600769498512, + 61.67243118946381, 64.19376191277179, 66.64228141346548, + 69.02756927200180, 71.35726934749408, 73.63759723904816, + 75.87368842650227, 78.06984431185720, 80.22971052389806, + 82.35640899964173, 84.45263768373256] + assert_allclose(res, res_exp) + + +def test_nctdtrinc_gh19896(): + # test that gh-19896 is resolved. + # Compared to SciPy 1.11 results from Fortran code. + dfarr = [0.001, 0.98, 9.8, 98, 980, 10000, 98, 9.8, 0.98, 0.001] + parr = [0.001, 0.1, 0.3, 0.8, 0.999, 0.001, 0.1, 0.3, 0.8, 0.999] + tarr = [0.0015, 0.15, 1.5, 15, 300, 0.0015, 0.15, 1.5, 15, 300] + desired = [3.090232306168629, 1.406141304556198, 2.014225177124157, + 13.727067118283456, 278.9765683871208, 3.090232306168629, + 1.4312427877936222, 2.014225177124157, 3.712743137978295, + -3.086951096691082] + actual = sp.nctdtrinc(dfarr, parr, tarr) + assert_allclose(actual, desired, rtol=5e-12, atol=0.0) + + +def test_stdtr_stdtrit_neg_inf(): + # -inf was treated as +inf and values from the normal were returned + assert np.all(np.isnan(sp.stdtr(-np.inf, [-np.inf, -1.0, 0.0, 1.0, np.inf]))) + assert np.all(np.isnan(sp.stdtrit(-np.inf, [0.0, 0.25, 0.5, 0.75, 1.0]))) + + +def test_bdtrik_nbdtrik_inf(): + y = np.array( + [np.nan,-np.inf,-10.0, -1.0, 0.0, .00001, .5, 0.9999, 1.0, 10.0, np.inf]) + y = y[:,None] + p = np.atleast_2d( + [np.nan, -np.inf, -10.0, -1.0, 0.0, .00001, .5, 1.0, np.inf]) + assert np.all(np.isnan(sp.bdtrik(y, np.inf, p))) + assert np.all(np.isnan(sp.nbdtrik(y, np.inf, p))) + + +@pytest.mark.parametrize( + "dfn,dfd,nc,f,expected_cdf", + [[100.0, 0.1, 0.1, 100.0, 0.29787396410092676], + [100.0, 100.0, 0.01, 0.1, 4.4344737598690424e-26], + [100.0, 0.01, 0.1, 0.01, 0.002848616633080384], + [10.0, 0.01, 1.0, 0.1, 0.012339557729057956], + [100.0, 100.0, 0.01, 0.01, 1.8926477420964936e-72], + [1.0, 100.0, 100.0, 0.1, 1.7925940526821304e-22], + [1.0, 0.01, 100.0, 10.0, 0.012334711965024968], + [1.0, 0.01, 10.0, 0.01, 0.00021944525290299], + [10.0, 1.0, 0.1, 100.0, 0.9219345555070705], + [0.1, 0.1, 1.0, 1.0, 0.3136335813423239], + [100.0, 100.0, 0.1, 10.0, 1.0], + [1.0, 0.1, 100.0, 10.0, 0.02926064279680897], + [1e-100, 3, 1.5, 1e100, 0.611815287345399] + ] +) +def test_ncfdtr_ncfdtri(dfn, dfd, nc, f, expected_cdf): + # Reference values computed with mpmath with the following script + # + # import numpy as np + # + # from mpmath import mp + # from scipy.special import ncfdtr + # + # mp.dps = 100 + # + # def mp_ncfdtr(dfn, dfd, nc, f): + # # Uses formula 26.2.20 from Abramowitz and Stegun. + # dfn, dfd, nc, f = map(mp.mpf, (dfn, dfd, nc, f)) + # def term(j): + # result = mp.exp(-nc/2)*(nc/2)**j / mp.factorial(j) + # result *= mp.betainc( + # dfn/2 + j, dfd/2, 0, f*dfn/(f*dfn + dfd), regularized=True + # ) + # return result + # result = mp.nsum(term, [0, mp.inf]) + # return float(result) + # + # dfn = np.logspace(-2, 2, 5) + # dfd = np.logspace(-2, 2, 5) + # nc = np.logspace(-2, 2, 5) + # f = np.logspace(-2, 2, 5) + # + # dfn, dfd, nc, f = np.meshgrid(dfn, dfd, nc, f) + # dfn, dfd, nc, f = map(np.ravel, (dfn, dfd, nc, f)) + # + # cases = [] + # re = [] + # for x0, x1, x2, x3 in zip(*(dfn, dfd, nc, f)): + # observed = ncfdtr(x0, x1, x2, x3) + # expected = mp_ncfdtr(x0, x1, x2, x3) + # cases.append((x0, x1, x2, x3, expected)) + # re.append((abs(expected - observed)/abs(expected))) + # + # assert np.max(re) < 1e-13 + # + # rng = np.random.default_rng(1234) + # sample_idx = rng.choice(len(re), replace=False, size=12) + # cases = np.array(cases)[sample_idx].tolist() + assert_allclose(sp.ncfdtr(dfn, dfd, nc, f), expected_cdf, rtol=1e-13, atol=0) + # testing tails where the CDF reaches 0 or 1 does not make sense for inverses + # of a CDF as they are not bijective in these regions + if 0 < expected_cdf < 1: + assert_allclose(sp.ncfdtri(dfn, dfd, nc, expected_cdf), f, rtol=5e-11) + +@pytest.mark.parametrize( + "args", + [(-1.0, 0.1, 0.1, 0.5), + (1, -1.0, 0.1, 0.5), + (1, 1, -1.0, 0.5), + (1, 1, 1, 100), + (1, 1, 1, -1)] +) +def test_ncfdtri_domain_error(args): + with sp.errstate(domain="raise"): + with pytest.raises(sp.SpecialFunctionError, match="domain"): + sp.ncfdtri(*args) + +class TestNoncentralTFunctions: + + # Reference values computed with mpmath with the following script + # Formula from: + # Lenth, Russell V (1989). "Algorithm AS 243: Cumulative Distribution Function + # of the Non-central t Distribution". Journal of the Royal Statistical Society, + # Series C. 38 (1): 185-189 + # + # Warning: may take a long time to run + # + # from mpmath import mp + # mp.dps = 400 + + # def nct_cdf(df, nc, x): + # df, nc, x = map(mp.mpf, (df, nc, x)) + + # def f(df, nc, x): + # phi = mp.ncdf(-nc) + # y = x * x / (x * x + df) + # constant = mp.exp(-nc * nc / 2.) + # def term(j): + # intermediate = constant * (nc *nc / 2.)**j + # p = intermediate/mp.factorial(j) + # q = nc / (mp.sqrt(2.) * mp.gamma(j + 1.5)) * intermediate + # first_beta_term = mp.betainc(j + 0.5, df/2., x2=y, + # regularized=True) + # second_beta_term = mp.betainc(j + mp.one, df/2., x2=y, + # regularized=True) + # return p * first_beta_term + q * second_beta_term + + # sum_term = mp.nsum(term, [0, mp.inf]) + # f = phi + 0.5 * sum_term + # return f + + # if x >= 0: + # result = f(df, nc, x) + # else: + # result = mp.one - f(df, -nc, x) + # return float(result) + + @pytest.mark.parametrize("df, nc, x, expected_cdf", [ + (0.98, -3.8, 0.0015, 0.9999279987514815), + (0.98, -3.8, 0.15, 0.9999528361700505), + (0.98, -3.8, 1.5, 0.9999908823016942), + (0.98, -3.8, 15, 0.9999990264591945), + (0.98, 0.38, 0.0015, 0.35241533122693), + (0.98, 0.38, 0.15, 0.39749697267146983), + (0.98, 0.38, 1.5, 0.716862963488558), + (0.98, 0.38, 15, 0.9656246449257494), + (0.98, 3.8, 0.0015, 7.26973354942293e-05), + (0.98, 3.8, 0.15, 0.00012416481147589105), + (0.98, 3.8, 1.5, 0.035388035775454095), + (0.98, 3.8, 15, 0.7954826975430583), + (0.98, 38, 0.0015, 3.02106943e-316), + (0.98, 38, 0.15, 6.069970616996603e-309), + (0.98, 38, 1.5, 2.591995360483094e-97), + (0.98, 38, 15, 0.011927265886910935), + (9.8, -3.8, 0.0015, 0.9999280776192786), + (9.8, -3.8, 0.15, 0.9999599410685442), + (9.8, -3.8, 1.5, 0.9999997432394788), + (9.8, -3.8, 15, 0.9999999999999984), + (9.8, 0.38, 0.0015, 0.3525155979107491), + (9.8, 0.38, 0.15, 0.40763120140379194), + (9.8, 0.38, 1.5, 0.8476794017024651), + (9.8, 0.38, 15, 0.9999999297116268), + (9.8, 3.8, 0.0015, 7.277620328149153e-05), + (9.8, 3.8, 0.15, 0.00013024802220900652), + (9.8, 3.8, 1.5, 0.013477432800072933), + (9.8, 3.8, 15, 0.999850151230648), + (9.8, 38, 0.0015, 3.05066095e-316), + (9.8, 38, 0.15, 1.79065514676e-313), + (9.8, 38, 1.5, 2.0935940165900746e-249), + (9.8, 38, 15, 2.252076291604796e-09), + (98, -3.8, 0.0015, 0.9999280875149109), + (98, -3.8, 0.15, 0.9999608250170452), + (98, -3.8, 1.5, 0.9999999304757682), + (98, -3.8, 15, 1.0), + (98, 0.38, 0.0015, 0.35252817848596313), + (98, 0.38, 0.15, 0.40890253001794846), + (98, 0.38, 1.5, 0.8664672830006552), + (98, 0.38, 15, 1.0), + (98, 3.8, 0.0015, 7.278609891281275e-05), + (98, 3.8, 0.15, 0.0001310318674827004), + (98, 3.8, 1.5, 0.010990879189991727), + (98, 3.8, 15, 0.9999999999999989), + (98, 38, 0.0015, 3.05437385e-316), + (98, 38, 0.15, 9.1668336166e-314), + (98, 38, 1.5, 1.8085884236563926e-288), + (98, 38, 15, 2.7740532792035907e-50), + (980, -3.8, 0.0015, 0.9999280885188965), + (980, -3.8, 0.15, 0.9999609144559273), + (980, -3.8, 1.5, 0.9999999410050979), + (980, -3.8, 15, 1.0), + (980, 0.38, 0.0015, 0.3525294548792812), + (980, 0.38, 0.15, 0.4090315324657382), + (980, 0.38, 1.5, 0.8684247068517293), + (980, 0.38, 15, 1.0), + (980, 3.8, 0.0015, 7.278710289828983e-05), + (980, 3.8, 0.15, 0.00013111131667906573), + (980, 3.8, 1.5, 0.010750678886113882), + (980, 3.8, 15, 1.0), + (980, 38, 0.0015, 3.0547506e-316), + (980, 38, 0.15, 8.6191646313e-314), + (980, 38, 15, 5.407535300713606e-105) + ]) + def test_gh19896(self, df, nc, x, expected_cdf): + # test that gh-19896 is resolved. + # Originally this was a regression test that used the old Fortran results + # as a reference. The Fortran results were not accurate, so the reference + # values were recomputed with mpmath. + nctdtr_result = sp.nctdtr(df, nc, x) + assert_allclose(nctdtr_result, expected_cdf, rtol=1e-13, atol=1e-303) + + def test_gh19896_loose_tolerance(self): + # edge case with very small CDF value, requires a looser tolerance to pass + assert_allclose(sp.nctdtr(980, 38, 1.5), 1.1824454111413493e-291, rtol=1e-8) + + def test_nctdtr_gh8344(self): + # test that gh-8344 is resolved. + df, nc, x = 3000, 3, 0.1 + expected = 0.0018657780826323328 + assert_allclose(sp.nctdtr(df, nc, x), expected, rtol=1e-14) + + @pytest.mark.parametrize( + "df, nc, x, expected, rtol", + [[3., 5., -2., 1.5645373999149622e-09, 2e-8], + [1000., 10., 1., 1.1493552133826623e-19, 1e-13], + [1e-5, -6., 2., 0.9999999990135003, 1e-13], + [10., 20., 0.15, 6.426530505957303e-88, 1e-13], + [1., 1., np.inf, 1.0, 0.0], + [1., 1., -np.inf, 0.0, 0.0] + ] + ) + def test_nctdtr_accuracy(self, df, nc, x, expected, rtol): + assert_allclose(sp.nctdtr(df, nc, x), expected, rtol=rtol) + + @pytest.mark.parametrize("df, nc, x, expected_cdf", [ + (0.98, 38, 1.5, 2.591995360483094e-97), + (3000, 3, 0.1, 0.0018657780826323328), + (0.98, -3.8, 15, 0.9999990264591945), + (9.8, 38, 15, 2.252076291604796e-09), + + ]) + def test_nctdtrit(self, df, nc, x, expected_cdf): + assert_allclose(sp.nctdtrit(df, nc, expected_cdf), x, rtol=1e-10) + + +class TestNegativeBinomialFunctions: + + @pytest.mark.parametrize("function", [sp.nbdtrin, sp.nbdtrik]) + @pytest.mark.parametrize("args", + [(np.nan, 1, 1), (1, np.nan, 1), (1, 1, np.nan), + (np.nan, np.nan, 1), (np.nan, 1, np.nan), (1, np.nan, np.nan), + (np.nan, np.nan, np.nan)] + ) + def test_nan_propagation(self, function, args): + assert np.isnan(function(*args)) + + @pytest.mark.parametrize("args", + [(-1, 1, 1), (1, -1, 1), (1, 1, -1), (-1, -1, 1), + (-1, 1, -1), (1, -1, -1), (-1, -1, -1), + (1, 1.1, 0.9), (1, 0.9, 1.1)] + ) + def test_domain_error_nbdtrin(self, args): + with sp.errstate(domain="raise"): + with pytest.raises(sp.SpecialFunctionError, match="domain"): + sp.nbdtrin(*args) + + @pytest.mark.parametrize("args", + [(-1, 1, 1), (1, -1, 1), (1, 1, -1), (-1, -1, 1), + (-1, 1, -1), (1, -1, -1), (-1, -1, -1), + (1, 0.9, 1.1)] + ) + def test_domain_error_nbdtrik(self, args): + with sp.errstate(domain="raise"): + with pytest.raises(sp.SpecialFunctionError, match="domain"): + sp.nbdtrik(*args) + + @pytest.mark.parametrize("k, y, p, n, rtol", + [(5, 0.3214569091796875, 0.25, 3, 5e-16), + (10, 3.820379007878081e-70, 0.15, 100, 1e-15), + (100, 0.975574722927043, 0.15, 10, 1e-14), + (2, 3.503253250000002e-73, 0.001, 25, 5e-15), + (2, 0.9999993271029954, 0.999, 15, 1e-10), + (500, 0.9999999998603178, 0.1, 15, 1e-8), + (0, 1.0000000000000006e-10, 0.1, 10, 5e-16)]) + def test_inverses(self, k, y, p, n, rtol): + # The following code was used to generate the values. + # import mpmath + # import scipy + # mpmath.mp.dps = 1000 + + # def neg_binomial_cdf(k, n, p): + # return mpmath.betainc(n, k+1, x1=0, x2=p, regularized=True) + + assert_allclose(sp.nbdtrin(k, y, p), n, rtol) + assert_allclose(sp.nbdtrik(y, n, p), k, rtol) + + +class TestNoncentralChiSquaredFunctions: + + def test_chndtr_and_inverses_against_wolfram_alpha(self): + # Each row holds (x, nu, lam, expected_value) + # These values were computed using Wolfram Alpha with + # CDF[NoncentralChiSquareDistribution[nu, lam], x] + values = np.array([ + [25.00, 20.0, 400, 4.1210655112396197139e-57], + [25.00, 8.00, 250, 2.3988026526832425878e-29], + [0.001, 8.00, 40., 5.3761806201366039084e-24], + [0.010, 8.00, 40., 5.45396231055999457039e-20], + [20.00, 2.00, 107, 1.39390743555819597802e-9], + [22.50, 2.00, 107, 7.11803307138105870671e-9], + [25.00, 2.00, 107, 3.11041244829864897313e-8], + [3.000, 2.00, 1.0, 0.62064365321954362734], + [350.0, 300., 10., 0.93880128006276407710], + [100.0, 13.5, 10., 0.99999999650104210949], + [700.0, 20.0, 400, 0.99999999925680650105], + [150.0, 13.5, 10., 0.99999999999999983046], + [160.0, 13.5, 10., 0.99999999999999999518], # 1.0 + ]) + cdf = sp.chndtr(values[:, 0], values[:, 1], values[:, 2]) + assert_allclose(cdf, values[:, 3], rtol=1e-13) + # the last two values are very close to 1.0, so we do not + # test the inverses for them + x = sp.chndtrix(values[:, 3], values[:, 1], values[:, 2]) + assert_allclose(x[:-2], values[:-2, 0], rtol=1e-8) + df = sp.chndtridf(values[:, 0], values[:, 3], values[:, 2]) + assert_allclose(df[:-2], values[:-2, 1], rtol=1e-8) + nc = sp.chndtrinc(values[:, 0], values[:, 1], values[:, 3]) + assert_allclose(nc[:-2], values[:-2, 2], rtol=1e-8) + + # CDF Reference values computed with mpmath with the following script + # Formula from: + # https://www.boost.org/doc/libs/1_87_0/libs/math/doc/html/math_toolkit/dist_ref/dists/nc_chi_squared_dist.html # noqa: E501 + # which cites: + # Computing discrete mixtures of continuous distributions: noncentral chisquare, + # noncentral t and the distribution of the square of the + # sample multiple correlation coefficient", + # Denise Benton and K. Krishnamoorthy, + # Computational Statistics & Data Analysis, 43, (2003), 249-267 + + # Warning: may take a long time to run + # from mpmath import mp + # + # mp.dps = 400 + # + # def noncentral_chi_squared_cdf(x, df, nc): + # x, df, nc = map(mp.mpf, (x, df, nc)) + # def term(i): + # return mp.exp(-nc/2) * (nc/2)**i / mp.factorial(i) * + # mp.gammainc(df/2 + i, 0, x/2, regularized=True) + # return float(mp.nsum(term, [0, mp.inf])) + + @pytest.mark.parametrize( + "x, df, nc, expected_cdf, rtol_cdf, rtol_inv", + [(0.1, 200, 50, 1.1311224867205481e-299, 1e-13, 1e-13), + (1e-12, 20, 50, 3.737446313006551e-141, 1e-13, 1e-13), + (1, 200, 50, 8.09760974833666e-200, 1e-13, 1e-13), + (9e4, 1e5, 1e3, 1.6895533704217566e-141, 5e-12, 5e-12), + (30, 3, 1.5, 0.9999508759095675, 5e-13, 5e-13), + (500, 300, 1.5, 0.9999999999944232, 1e-13, 1e-5), + (1400, 30, 1e3, 0.9999999612246238, 1e-13, 1e-9), + (400, 50, 200, 0.9999948255072892, 1e-13, 1e-11)] + ) + def test_tails(self, x, df, nc, expected_cdf, rtol_cdf, rtol_inv): + chndtr_result = sp.chndtr(x, df, nc) + assert_allclose(chndtr_result, expected_cdf, rtol=rtol_cdf) + assert_allclose(sp.chndtrix(expected_cdf, df, nc), x, rtol=rtol_inv) + assert_allclose(sp.chndtridf(x, expected_cdf, nc), df, rtol=rtol_inv) + assert_allclose(sp.chndtrinc(x, df, expected_cdf), nc, rtol=rtol_inv) + + # test round-trip calculation + assert_allclose(sp.chndtrix(chndtr_result, df, nc), x, rtol=rtol_inv) + assert_allclose(sp.chndtridf(x, chndtr_result, nc), df, rtol=1e-7) + assert_allclose(sp.chndtrinc(x, df, chndtr_result), nc, rtol=1e-5) + + @pytest.mark.parametrize("x, df", + [(1, 3), (1, 0), (1, np.inf), (np.inf, 1)] + ) + def test_chndtr_with_nc_zero_equals_chdtr(self, x, df): + assert_allclose(sp.chndtr(x, df, 0), sp.chdtr(df, x), rtol=1e-15) + + @pytest.mark.parametrize("fun", + [sp.chndtr, sp.chndtrix, sp.chndtridf, sp.chndtrinc] + ) + @pytest.mark.parametrize("args", + [(-1, 1, 1), (1, -1, 1), (1, 1, -1), (-1, -1, 1), + (-1, 1, -1), (1, -1, -1), (-1, -1, -1)] + ) + def test_domain_error(self, args, fun): + with sp.errstate(domain="raise"): + with pytest.raises(sp.SpecialFunctionError, match="domain"): + fun(*args) + + @pytest.mark.parametrize("fun", + [sp.chndtr, sp.chndtrix, sp.chndtridf, sp.chndtrinc] + ) + @pytest.mark.parametrize("args", + [(np.nan, 1, 1), (1, np.nan, 1), (1, 1, np.nan), + (np.nan, np.nan, 1), (np.nan, 1, np.nan), (1, np.nan, np.nan), + (np.nan, np.nan, np.nan)] + ) + def test_nan_propagation(self, fun, args): + assert np.isnan(fun(*args)) + + @pytest.mark.parametrize( + "x, df, nc, expected", + [(1, 0, 1, np.nan), + (1, 0, np.inf, np.nan), + (1, np.inf, 1, np.nan), + (1, np.inf, np.inf, 0), + (1, 1, np.inf, 0), + (np.inf, 0, 1, np.nan), + (np.inf, 1, np.inf, np.nan), + (np.inf, 1, 1, 1), + (np.inf, np.inf, np.inf, np.nan)] + ) + def test_chndtr_edge_cases(self, x, df, nc, expected): + assert_allclose(sp.chndtr(x, df, nc), expected, rtol=1e-15) + + +@pytest.mark.parametrize("x", [0.1, 100]) +def test_chdtriv_p_equals_1_returns_0(x): + assert sp.chdtriv(1, x) == 0 + + +class TestPdtrik: + @pytest.mark.parametrize("p, m, expected", + [(0, 0.5, 0), + (0.5, 0, 0)]) + def test_edge_cases(self, p, m, expected): + assert sp.pdtrik(p, m) == expected + + @pytest.mark.parametrize("m", (0.1, 1, 10)) + def test_p_equals_1_returns_nan(self, m): + assert np.isnan(sp.pdtrik(1, m)) + + def test_small_probabilities(self): + # Edge case: m = 0 or very small. + k = sp.pdtrik([[0], [0.25], [0.95]], [0, 1e-20, 1e-6]) + assert_array_equal(k, np.zeros((3, 3))) + + def test_roundtrip_against_pdtr(self): + m = [10, 50, 500] + k = 5 + p = sp.pdtr(k, m) + assert_allclose(sp.pdtrik(p, m), k, rtol=1e-15) + + @pytest.mark.parametrize("p, m, k", + [(1.8976107553682285e-40, 100, 2), + (0.48670120172085135, 100, 99), + (8.30383406699052e-69, 1000, 500), + (2.252837804125894e-227, 100_000, 90_000)]) + def test_accuracy(self, p, m, k): + # Reference values for p were computed with mpmath using + # mp.gammainc(k+1, a=m, regularized=True) + assert_allclose(sp.pdtrik(p, m), k, rtol=1e-15) + +@pytest.mark.parametrize("a, b, p, ref", [ + (0, 0, 0, np.nan), + (0, 0, 1, np.nan), + (0, np.inf, 0, np.nan), + (0, np.inf, 1, np.nan), + (np.inf, 0, 0, np.nan), + (np.inf, 0, 1, np.nan), + (np.inf, np.inf, 0, np.nan), + (np.inf, np.inf, 1, np.nan) +]) +def test_gdtrix_edge_cases(a, b, p, ref): + assert_equal(sp.gdtrix(a, b, p), ref) + +@pytest.mark.parametrize("p, b, x, ref", [ + (0, 0, 0, np.nan), + (0, 0, np.inf, np.nan), + (0, 1, 0, np.nan), + (0, 1, np.inf, 0), + (np.inf, 0, 0, np.nan), + (np.inf, 0, np.inf, np.nan), + (np.inf, 1, 0, np.nan), + (np.inf, 1, np.inf, np.nan) +]) +def test_gdtria_edge_cases(p, b, x, ref): + assert_equal(sp.gdtria(p, b, x), ref) + +@pytest.mark.parametrize("y, n, p, k", [ + (6.110236824190975e-78, 100, 0.9, 10), + (4.0188255421199357e-224, 1000, 0.9, 500), + (0.9999998319058928, 1000, 0.5, 580), + (0.9999734386011124, 100, 0.9, 99) +]) +def test_bdtrik(y, n, p, k): + # Reference values for y were computed with mpmath using + # the _binomial_cdf function from above. + assert_allclose(sp.bdtrik(y, n, p), k, rtol=1e-11) + + +@pytest.mark.parametrize("p, std, x, ref", [ + (0, 0.1, 0, np.nan), + (0.1, np.inf, 1, np.inf), + (0.1, 1, np.inf, np.inf), + (0.1, 1, -np.inf, -np.inf), + (0.1, np.inf, np.inf, np.inf), + (0.1, 0.1, np.inf, np.inf), + (0.1, 1, -np.inf, -np.inf), + (0.3, -1, 1, np.nan) +]) +def test_nrdtrimn_edge_cases(p, std, x, ref): + assert_equal(sp.nrdtrimn(p, std, x), ref) + + +@pytest.mark.parametrize("mn, p, x, ref", [ + (0, 0, 0, np.nan), + (0, 1, 0, np.nan), + (0, 1, np.inf, np.nan), + (0.1, 1, 1, np.nan), + (0.1, 0.1, np.inf, -np.inf), + (0.1, 0.1, -np.inf, np.inf), +]) +def test_nrdtrisd_edge_cases(mn, p, x, ref): + assert_equal(sp.nrdtrisd(mn, p, x), ref) diff --git a/v12x/pylib/scipy/special/tests/test_cdft_asymptotic.py b/v12x/pylib/scipy/special/tests/test_cdft_asymptotic.py new file mode 100644 index 0000000000000000000000000000000000000000..2f21b197871a230e3170c21969a8edcc8ebbdfc2 --- /dev/null +++ b/v12x/pylib/scipy/special/tests/test_cdft_asymptotic.py @@ -0,0 +1,52 @@ +# gh-14777 regression tests +# Test stdtr and stdtrit with infinite df and large values of df + +import numpy as np +from numpy.testing import assert_allclose, assert_equal +from scipy.special import stdtr, stdtrit, ndtr, ndtri + + +def test_stdtr_vs_R_large_df(): + df = [1e10, 1e12, 1e120, np.inf] + t = 1. + res = stdtr(df, t) + # R Code: + # options(digits=20) + # pt(1., c(1e10, 1e12, 1e120, Inf)) + res_R = [0.84134474605644460343, + 0.84134474606842180044, + 0.84134474606854281475, + 0.84134474606854292578] + assert_allclose(res, res_R, rtol=1e-15) + # last value should also agree with ndtr + assert_equal(res[3], ndtr(1.)) + + +def test_stdtrit_vs_R_large_df(): + df = [1e10, 1e12, 1e120, np.inf] + p = 0.1 + res = stdtrit(df, p) + # R Code: + # options(digits=20) + # qt(0.1, c(1e10, 1e12, 1e120, Inf)) + res_R = [-1.2815515656292593150, + -1.2815515655454472466, + -1.2815515655446008125, + -1.2815515655446008125] + assert_allclose(res, res_R, rtol=1e-15, atol=1e-15) + # last value should also agree with ndtri + # actually the result from stdtrit is closer to R than ndtri, + # so we accept a deviation of one ULP + epsilon = np.finfo(np.float64).eps + assert_allclose(res[3], ndtri(0.1), rtol=epsilon, atol=epsilon) + + +def test_stdtr_stdtri_invalid(): + # a mix of large and inf df with t/p equal to nan + df = [1e10, 1e12, 1e120, np.inf] + x = np.nan + res1 = stdtr(df, x) + res2 = stdtrit(df, x) + res_ex = 4*[np.nan] + assert_equal(res1, res_ex) + assert_equal(res2, res_ex) diff --git a/v12x/pylib/scipy/special/tests/test_cephes_intp_cast.py b/v12x/pylib/scipy/special/tests/test_cephes_intp_cast.py new file mode 100644 index 0000000000000000000000000000000000000000..05f3d1ae5c101ff50c75d1065e5e234063d192e4 --- /dev/null +++ b/v12x/pylib/scipy/special/tests/test_cephes_intp_cast.py @@ -0,0 +1,29 @@ +import pytest +import numpy as np +from scipy.special._ufuncs import ( + _smirnovc, _smirnovci, _smirnovp, + _struve_asymp_large_z, _struve_bessel_series, _struve_power_series, + bdtr, bdtrc, bdtri, expn, kn, nbdtr, nbdtrc, nbdtri, pdtri, + smirnov, smirnovi, yn +) + + +# +# For each ufunc here, verify that the default integer type, np.intp, +# can be safely cast to the integer type found in the input type signatures. +# For this particular set of functions, the code expects to find just one +# integer type among the input signatures. +# +@pytest.mark.parametrize( + 'ufunc', + [_smirnovc, _smirnovci, _smirnovp, + _struve_asymp_large_z, _struve_bessel_series, _struve_power_series, + bdtr, bdtrc, bdtri, expn, kn, nbdtr, nbdtrc, nbdtri, pdtri, + smirnov, smirnovi, yn], +) +def test_intp_safe_cast(ufunc): + int_chars = {'i', 'l', 'q'} + int_input = [set(sig.split('->')[0]) & int_chars for sig in ufunc.types] + int_char = ''.join(s.pop() if s else '' for s in int_input) + assert len(int_char) == 1, "More integer types in the signatures than expected" + assert np.can_cast(np.intp, np.dtype(int_char)) diff --git a/v12x/pylib/scipy/special/tests/test_cosine_distr.py b/v12x/pylib/scipy/special/tests/test_cosine_distr.py new file mode 100644 index 0000000000000000000000000000000000000000..27e3ca2699d0d1d0b58665f125d99c166095696d --- /dev/null +++ b/v12x/pylib/scipy/special/tests/test_cosine_distr.py @@ -0,0 +1,83 @@ +import numpy as np +from numpy.testing import assert_allclose +import pytest +from scipy.special._ufuncs import _cosine_cdf, _cosine_invcdf + + +# These values are (x, p) where p is the expected exact value of +# _cosine_cdf(x). These values will be tested for exact agreement. +_coscdf_exact = [ + (-4.0, 0.0), + (0, 0.5), + (np.pi, 1.0), + (4.0, 1.0), +] + +@pytest.mark.parametrize("x, expected", _coscdf_exact) +def test_cosine_cdf_exact(x, expected): + assert _cosine_cdf(x) == expected + + +# These values are (x, p), where p is the expected value of +# _cosine_cdf(x). The expected values were computed with mpmath using +# 50 digits of precision. These values will be tested for agreement +# with the computed values using a very small relative tolerance. +# The value at -np.pi is not 0, because -np.pi does not equal -π. +_coscdf_close = [ + (3.1409, 0.999999999991185), + (2.25, 0.9819328173287907), + # -1.6 is the threshold below which the Pade approximant is used. + (-1.599, 0.08641959838382553), + (-1.601, 0.086110582992713), + (-2.0, 0.0369709335961611), + (-3.0, 7.522387241801384e-05), + (-3.1415, 2.109869685443648e-14), + (-3.14159, 4.956444476505336e-19), + (-np.pi, 4.871934450264861e-50), +] + +@pytest.mark.parametrize("x, expected", _coscdf_close) +def test_cosine_cdf(x, expected): + assert_allclose(_cosine_cdf(x), expected, rtol=5e-15) + + +# These values are (p, x) where x is the expected exact value of +# _cosine_invcdf(p). These values will be tested for exact agreement. +_cosinvcdf_exact = [ + (0.0, -np.pi), + (0.5, 0.0), + (1.0, np.pi), +] + +@pytest.mark.parametrize("p, expected", _cosinvcdf_exact) +def test_cosine_invcdf_exact(p, expected): + assert _cosine_invcdf(p) == expected + + +def test_cosine_invcdf_invalid_p(): + # Check that p values outside of [0, 1] return nan. + assert np.isnan(_cosine_invcdf([-0.1, 1.1])).all() + + +# These values are (p, x), where x is the expected value of _cosine_invcdf(p). +# The expected values were computed with mpmath using 50 digits of precision. +_cosinvcdf_close = [ + (1e-50, -np.pi), + (1e-14, -3.1415204137058454), + (1e-08, -3.1343686589124524), + (0.0018001, -2.732563923138336), + (0.010, -2.41276589008678), + (0.060, -1.7881244975330157), + (0.125, -1.3752523669869274), + (0.250, -0.831711193579736), + (0.400, -0.3167954512395289), + (0.419, -0.25586025626919906), + (0.421, -0.24947570750445663), + (0.750, 0.831711193579736), + (0.940, 1.7881244975330153), + (0.9999999996, 3.1391220839917167), +] + +@pytest.mark.parametrize("p, expected", _cosinvcdf_close) +def test_cosine_invcdf(p, expected): + assert_allclose(_cosine_invcdf(p), expected, rtol=1e-14) diff --git a/v12x/pylib/scipy/special/tests/test_cython_abi.py b/v12x/pylib/scipy/special/tests/test_cython_abi.py new file mode 100644 index 0000000000000000000000000000000000000000..56b61ca951d511bf2a73c8c87e3cc7028ed91b0e --- /dev/null +++ b/v12x/pylib/scipy/special/tests/test_cython_abi.py @@ -0,0 +1,87 @@ +# --------------------------------------------------------------------------- +# AUTO-GENERATED - do not edit by hand. +# Re-generate with: python tools/generate_cython_abi_tests.py special +# +# Generated against SciPy 1.17.0 on 2026-05-08 +# Python 3.14.4 +# --------------------------------------------------------------------------- + +""" +Regression test for the cython_special binary interface. + +The __pyx_capi__ dict of each module contains PyCapsule objects whose +*name* string encodes the C-level function signature that Cython checks +at import time when another module does ``cimport``. Any change to +those strings is a binary incompatibility that breaks every downstream +package (e.g., scikit-learn, statsmodels) compiled against the previous release, +without any compile-time warning. + +The expected signatures are stored in cython_abi_signatures.json. +The test fails if: + * an existing name disappears, or + * an existing name's signature string changes. +Adding new names is explicitly allowed (it is not a break). +""" + +import ctypes +import importlib +import importlib.resources +import json + +import pytest +import scipy + +CYTHON_BLAS_ILP64 = ( + scipy.show_config(mode='dicts')['Build Dependencies']['blas']['cython blas ilp64'] +) + + +# --------------------------------------------------------------------------- +# Helper +# --------------------------------------------------------------------------- + +_get_capsule_name = ctypes.pythonapi.PyCapsule_GetName +_get_capsule_name.restype = ctypes.c_char_p +_get_capsule_name.argtypes = [ctypes.py_object] + + +def _extract_capi(module): + """Return {name: signature} for every entry in module.__pyx_capi__.""" + result = {} + for name, capsule in module.__pyx_capi__.items(): + raw = _get_capsule_name(capsule) + if raw is not None: + result[name] = raw.decode('utf-8') + return result + + +_sig_file = importlib.resources.files('scipy.special.tests').joinpath( + 'cython_abi_signatures.json' +) +with importlib.resources.as_file(_sig_file) as _f: + _EXPECTED = json.loads(_f.read_text()) + + +def test_cython_special_abi_stability(): + """No existing cython_special signature may change or disappear.""" + expected = _EXPECTED['scipy.special.cython_special'] + mod = importlib.import_module('scipy.special.cython_special') + actual = _extract_capi(mod) + errors = [] + for name, expected_sig in expected.items(): + if name not in actual: + errors.append( + f"REMOVED {name!r} (was {expected_sig!r})" + ) + elif actual[name] != expected_sig: + errors.append( + f"CHANGED {name!r}\n" + f" expected: {expected_sig!r}\n" + f" actual: {actual[name]!r}" + ) + if errors: + joined = '\n'.join(errors) + pytest.fail( + f"scipy.special.cython_special.__pyx_capi__ has {len(errors)} " + f"ABI breakage(s):\n{joined}" + ) diff --git a/v12x/pylib/scipy/special/tests/test_cython_special.py b/v12x/pylib/scipy/special/tests/test_cython_special.py new file mode 100644 index 0000000000000000000000000000000000000000..a9307767ddaeb43dd8c319ac8197198f935f132a --- /dev/null +++ b/v12x/pylib/scipy/special/tests/test_cython_special.py @@ -0,0 +1,363 @@ +import warnings + +from collections.abc import Callable + +import pytest +from itertools import product +from numpy.testing import assert_allclose +from scipy import special +from scipy.special import cython_special + + +bint_points = [True, False] +int_points = [-10, -1, 1, 10] +real_points = [-10.0, -1.0, 1.0, 10.0] +complex_points = [complex(*tup) for tup in product(real_points, repeat=2)] + + +CYTHON_SIGNATURE_MAP = { + 'b': 'bint', + 'f': 'float', + 'd': 'double', + 'g': 'long double', + 'F': 'float complex', + 'D': 'double complex', + 'G': 'long double complex', + 'i': 'int', + 'l': 'long' +} + + +TEST_POINTS = { + 'b': bint_points, + 'f': real_points, + 'd': real_points, + 'g': real_points, + 'F': complex_points, + 'D': complex_points, + 'G': complex_points, + 'i': int_points, + 'l': int_points, +} + + +PARAMS: list[tuple[Callable, Callable, tuple[str, ...], str | None]] = [ + (special.agm, cython_special.agm, ('dd',), None), + (special.airy, cython_special._airy_pywrap, ('d', 'D'), None), + (special.airye, cython_special._airye_pywrap, ('d', 'D'), None), + (special.bdtr, cython_special.bdtr, ('dld', 'ddd'), None), + (special.bdtrc, cython_special.bdtrc, ('dld', 'ddd'), None), + (special.bdtri, cython_special.bdtri, ('dld', 'ddd'), None), + (special.bdtrik, cython_special.bdtrik, ('ddd',), None), + (special.bdtrin, cython_special.bdtrin, ('ddd',), None), + (special.bei, cython_special.bei, ('d',), None), + (special.beip, cython_special.beip, ('d',), None), + (special.ber, cython_special.ber, ('d',), None), + (special.berp, cython_special.berp, ('d',), None), + (special.besselpoly, cython_special.besselpoly, ('ddd',), None), + (special.beta, cython_special.beta, ('dd',), None), + (special.betainc, cython_special.betainc, ('ddd',), None), + (special.betaincc, cython_special.betaincc, ('ddd',), None), + (special.betaincinv, cython_special.betaincinv, ('ddd',), None), + (special.betainccinv, cython_special.betainccinv, ('ddd',), None), + (special.betaln, cython_special.betaln, ('dd',), None), + (special.binom, cython_special.binom, ('dd',), None), + (special.boxcox, cython_special.boxcox, ('dd',), None), + (special.boxcox1p, cython_special.boxcox1p, ('dd',), None), + (special.btdtria, cython_special.btdtria, ('ddd',), None), + (special.btdtrib, cython_special.btdtrib, ('ddd',), None), + (special.cbrt, cython_special.cbrt, ('d',), None), + (special.chdtr, cython_special.chdtr, ('dd',), None), + (special.chdtrc, cython_special.chdtrc, ('dd',), None), + (special.chdtri, cython_special.chdtri, ('dd',), None), + (special.chdtriv, cython_special.chdtriv, ('dd',), None), + (special.chndtr, cython_special.chndtr, ('ddd',), None), + (special.chndtridf, cython_special.chndtridf, ('ddd',), None), + (special.chndtrinc, cython_special.chndtrinc, ('ddd',), None), + (special.chndtrix, cython_special.chndtrix, ('ddd',), None), + (special.cosdg, cython_special.cosdg, ('d',), None), + (special.cosm1, cython_special.cosm1, ('d',), None), + (special.cotdg, cython_special.cotdg, ('d',), None), + (special.dawsn, cython_special.dawsn, ('d', 'D'), None), + (special.ellipe, cython_special.ellipe, ('d',), None), + (special.ellipeinc, cython_special.ellipeinc, ('dd',), None), + (special.ellipj, cython_special._ellipj_pywrap, ('dd',), None), + (special.ellipkinc, cython_special.ellipkinc, ('dd',), None), + (special.ellipkm1, cython_special.ellipkm1, ('d',), None), + (special.ellipk, cython_special.ellipk, ('d',), None), + (special.elliprc, cython_special.elliprc, ('dd', 'DD'), None), + (special.elliprd, cython_special.elliprd, ('ddd', 'DDD'), None), + (special.elliprf, cython_special.elliprf, ('ddd', 'DDD'), None), + (special.elliprg, cython_special.elliprg, ('ddd', 'DDD'), None), + (special.elliprj, cython_special.elliprj, ('dddd', 'DDDD'), None), + (special.entr, cython_special.entr, ('d',), None), + (special.erf, cython_special.erf, ('d', 'D'), None), + (special.erfc, cython_special.erfc, ('d', 'D'), None), + (special.erfcx, cython_special.erfcx, ('d', 'D'), None), + (special.erfi, cython_special.erfi, ('d', 'D'), None), + (special.erfinv, cython_special.erfinv, ('d',), None), + (special.erfcinv, cython_special.erfcinv, ('d',), None), + (special.eval_chebyc, cython_special.eval_chebyc, ('dd', 'dD', 'ld'), None), + (special.eval_chebys, cython_special.eval_chebys, ('dd', 'dD', 'ld'), + 'd and l differ for negative int'), + (special.eval_chebyt, cython_special.eval_chebyt, ('dd', 'dD', 'ld'), + 'd and l differ for negative int'), + (special.eval_chebyu, cython_special.eval_chebyu, ('dd', 'dD', 'ld'), + 'd and l differ for negative int'), + (special.eval_gegenbauer, cython_special.eval_gegenbauer, ('ddd', 'ddD', 'ldd'), + 'd and l differ for negative int'), + (special.eval_genlaguerre, cython_special.eval_genlaguerre, ('ddd', 'ddD', 'ldd'), + 'd and l differ for negative int'), + (special.eval_hermite, cython_special.eval_hermite, ('ld',), None), + (special.eval_hermitenorm, cython_special.eval_hermitenorm, ('ld',), None), + (special.eval_jacobi, cython_special.eval_jacobi, ('dddd', 'dddD', 'lddd'), + 'd and l differ for negative int'), + (special.eval_laguerre, cython_special.eval_laguerre, ('dd', 'dD', 'ld'), + 'd and l differ for negative int'), + (special.eval_legendre, cython_special.eval_legendre, ('dd', 'dD', 'ld'), None), + (special.eval_sh_chebyt, cython_special.eval_sh_chebyt, ('dd', 'dD', 'ld'), None), + (special.eval_sh_chebyu, cython_special.eval_sh_chebyu, ('dd', 'dD', 'ld'), + 'd and l differ for negative int'), + (special.eval_sh_jacobi, cython_special.eval_sh_jacobi, ('dddd', 'dddD', 'lddd'), + 'd and l differ for negative int'), + (special.eval_sh_legendre, cython_special.eval_sh_legendre, ('dd', 'dD', 'ld'), + None), + (special.exp1, cython_special.exp1, ('d', 'D'), None), + (special.exp10, cython_special.exp10, ('d',), None), + (special.exp2, cython_special.exp2, ('d',), None), + (special.expi, cython_special.expi, ('d', 'D'), None), + (special.expit, cython_special.expit, ('f', 'd', 'g'), None), + (special.expm1, cython_special.expm1, ('d', 'D'), None), + (special.expn, cython_special.expn, ('ld', 'dd'), None), + (special.exprel, cython_special.exprel, ('d',), None), + (special.fdtr, cython_special.fdtr, ('ddd',), None), + (special.fdtrc, cython_special.fdtrc, ('ddd',), None), + (special.fdtri, cython_special.fdtri, ('ddd',), None), + (special.fdtridfd, cython_special.fdtridfd, ('ddd',), None), + (special.fresnel, cython_special._fresnel_pywrap, ('d', 'D'), None), + (special.gamma, cython_special.gamma, ('d', 'D'), None), + (special.gammainc, cython_special.gammainc, ('dd',), None), + (special.gammaincc, cython_special.gammaincc, ('dd',), None), + (special.gammainccinv, cython_special.gammainccinv, ('dd',), None), + (special.gammaincinv, cython_special.gammaincinv, ('dd',), None), + (special.gammaln, cython_special.gammaln, ('d',), None), + (special.gammasgn, cython_special.gammasgn, ('d',), None), + (special.gdtr, cython_special.gdtr, ('ddd',), None), + (special.gdtrc, cython_special.gdtrc, ('ddd',), None), + (special.gdtria, cython_special.gdtria, ('ddd',), None), + (special.gdtrib, cython_special.gdtrib, ('ddd',), None), + (special.gdtrix, cython_special.gdtrix, ('ddd',), None), + (special.hankel1, cython_special.hankel1, ('dD',), None), + (special.hankel1e, cython_special.hankel1e, ('dD',), None), + (special.hankel2, cython_special.hankel2, ('dD',), None), + (special.hankel2e, cython_special.hankel2e, ('dD',), None), + (special.huber, cython_special.huber, ('dd',), None), + (special.hyp0f1, cython_special.hyp0f1, ('dd', 'dD'), None), + (special.hyp1f1, cython_special.hyp1f1, ('ddd', 'ddD'), None), + (special.hyp2f1, cython_special.hyp2f1, ('dddd', 'dddD'), None), + (special.hyperu, cython_special.hyperu, ('ddd',), None), + (special.i0, cython_special.i0, ('d',), None), + (special.i0e, cython_special.i0e, ('d',), None), + (special.i1, cython_special.i1, ('d',), None), + (special.i1e, cython_special.i1e, ('d',), None), + (special.inv_boxcox, cython_special.inv_boxcox, ('dd',), None), + (special.inv_boxcox1p, cython_special.inv_boxcox1p, ('dd',), None), + (special.it2i0k0, cython_special._it2i0k0_pywrap, ('d',), None), + (special.it2j0y0, cython_special._it2j0y0_pywrap, ('d',), None), + (special.it2struve0, cython_special.it2struve0, ('d',), None), + (special.itairy, cython_special._itairy_pywrap, ('d',), None), + (special.iti0k0, cython_special._iti0k0_pywrap, ('d',), None), + (special.itj0y0, cython_special._itj0y0_pywrap, ('d',), None), + (special.itmodstruve0, cython_special.itmodstruve0, ('d',), None), + (special.itstruve0, cython_special.itstruve0, ('d',), None), + (special.iv, cython_special.iv, ('dd', 'dD'), None), + (special.ive, cython_special.ive, ('dd', 'dD'), None), + (special.j0, cython_special.j0, ('d',), None), + (special.j1, cython_special.j1, ('d',), None), + (special.jv, cython_special.jv, ('dd', 'dD'), None), + (special.jve, cython_special.jve, ('dd', 'dD'), None), + (special.k0, cython_special.k0, ('d',), None), + (special.k0e, cython_special.k0e, ('d',), None), + (special.k1, cython_special.k1, ('d',), None), + (special.k1e, cython_special.k1e, ('d',), None), + (special.kei, cython_special.kei, ('d',), None), + (special.keip, cython_special.keip, ('d',), None), + (special.kelvin, cython_special._kelvin_pywrap, ('d',), None), + (special.ker, cython_special.ker, ('d',), None), + (special.kerp, cython_special.kerp, ('d',), None), + (special.kl_div, cython_special.kl_div, ('dd',), None), + (special.kn, cython_special.kn, ('ld', 'dd'), None), + (special.kolmogi, cython_special.kolmogi, ('d',), None), + (special.kolmogorov, cython_special.kolmogorov, ('d',), None), + (special.kv, cython_special.kv, ('dd', 'dD'), None), + (special.kve, cython_special.kve, ('dd', 'dD'), None), + (special.log1p, cython_special.log1p, ('d', 'D'), None), + (special.log_expit, cython_special.log_expit, ('f', 'd', 'g'), None), + (special.log_ndtr, cython_special.log_ndtr, ('d', 'D'), None), + (special.log_wright_bessel, cython_special.log_wright_bessel, ('ddd',), None), + (special.ndtri_exp, cython_special.ndtri_exp, ('d',), None), + (special.loggamma, cython_special.loggamma, ('D',), None), + (special.logit, cython_special.logit, ('f', 'd', 'g'), None), + (special.lpmv, cython_special.lpmv, ('ddd',), None), + (special.mathieu_a, cython_special.mathieu_a, ('dd',), None), + (special.mathieu_b, cython_special.mathieu_b, ('dd',), None), + (special.mathieu_cem, cython_special._mathieu_cem_pywrap, ('ddd',), None), + (special.mathieu_modcem1, cython_special._mathieu_modcem1_pywrap, ('ddd',), None), + (special.mathieu_modcem2, cython_special._mathieu_modcem2_pywrap, ('ddd',), None), + (special.mathieu_modsem1, cython_special._mathieu_modsem1_pywrap, ('ddd',), None), + (special.mathieu_modsem2, cython_special._mathieu_modsem2_pywrap, ('ddd',), None), + (special.mathieu_sem, cython_special._mathieu_sem_pywrap, ('ddd',), None), + (special.modfresnelm, cython_special._modfresnelm_pywrap, ('d',), None), + (special.modfresnelp, cython_special._modfresnelp_pywrap, ('d',), None), + (special.modstruve, cython_special.modstruve, ('dd',), None), + (special.nbdtr, cython_special.nbdtr, ('lld', 'ddd'), None), + (special.nbdtrc, cython_special.nbdtrc, ('lld', 'ddd'), None), + (special.nbdtri, cython_special.nbdtri, ('lld', 'ddd'), None), + (special.nbdtrik, cython_special.nbdtrik, ('ddd',), None), + (special.nbdtrin, cython_special.nbdtrin, ('ddd',), None), + (special.ncfdtr, cython_special.ncfdtr, ('dddd',), None), + (special.ncfdtri, cython_special.ncfdtri, ('dddd',), None), + (special.ncfdtridfd, cython_special.ncfdtridfd, ('dddd',), None), + (special.ncfdtridfn, cython_special.ncfdtridfn, ('dddd',), None), + (special.ncfdtrinc, cython_special.ncfdtrinc, ('dddd',), None), + (special.nctdtr, cython_special.nctdtr, ('ddd',), None), + (special.nctdtridf, cython_special.nctdtridf, ('ddd',), None), + (special.nctdtrinc, cython_special.nctdtrinc, ('ddd',), None), + (special.nctdtrit, cython_special.nctdtrit, ('ddd',), None), + (special.ndtr, cython_special.ndtr, ('d', 'D'), None), + (special.ndtri, cython_special.ndtri, ('d',), None), + (special.nrdtrimn, cython_special.nrdtrimn, ('ddd',), None), + (special.nrdtrisd, cython_special.nrdtrisd, ('ddd',), None), + (special.obl_ang1, cython_special._obl_ang1_pywrap, ('dddd',), None), + (special.obl_ang1_cv, cython_special._obl_ang1_cv_pywrap, ('ddddd',), None), + (special.obl_cv, cython_special.obl_cv, ('ddd',), None), + (special.obl_rad1, cython_special._obl_rad1_pywrap, ('dddd',), "see gh-6211"), + (special.obl_rad1_cv, cython_special._obl_rad1_cv_pywrap, ('ddddd',), + "see gh-6211"), + (special.obl_rad2, cython_special._obl_rad2_pywrap, ('dddd',), "see gh-6211"), + (special.obl_rad2_cv, cython_special._obl_rad2_cv_pywrap, ('ddddd',), + "see gh-6211"), + (special.pbdv, cython_special._pbdv_pywrap, ('dd',), None), + (special.pbvv, cython_special._pbvv_pywrap, ('dd',), None), + (special.pbwa, cython_special._pbwa_pywrap, ('dd',), None), + (special.pdtr, cython_special.pdtr, ('dd', 'dd'), None), + (special.pdtrc, cython_special.pdtrc, ('dd', 'dd'), None), + (special.pdtri, cython_special.pdtri, ('ld', 'dd'), None), + (special.pdtrik, cython_special.pdtrik, ('dd',), None), + (special.poch, cython_special.poch, ('dd',), None), + (special.powm1, cython_special.powm1, ('dd',), None), + (special.pro_ang1, cython_special._pro_ang1_pywrap, ('dddd',), None), + (special.pro_ang1_cv, cython_special._pro_ang1_cv_pywrap, ('ddddd',), None), + (special.pro_cv, cython_special.pro_cv, ('ddd',), None), + (special.pro_rad1, cython_special._pro_rad1_pywrap, ('dddd',), "see gh-6211"), + (special.pro_rad1_cv, cython_special._pro_rad1_cv_pywrap, ('ddddd',), + "see gh-6211"), + (special.pro_rad2, cython_special._pro_rad2_pywrap, ('dddd',), "see gh-6211"), + (special.pro_rad2_cv, cython_special._pro_rad2_cv_pywrap, ('ddddd',), + "see gh-6211"), + (special.pseudo_huber, cython_special.pseudo_huber, ('dd',), None), + (special.psi, cython_special.psi, ('d', 'D'), None), + (special.radian, cython_special.radian, ('ddd',), None), + (special.rel_entr, cython_special.rel_entr, ('dd',), None), + (special.rgamma, cython_special.rgamma, ('d', 'D'), None), + (special.round, cython_special.round, ('d',), None), + (special.spherical_jn, cython_special.spherical_jn, ('ld', 'ldb', 'lD', 'lDb'), + "Python version supports negative reals; Cython version doesn't - see gh-21629"), + (special.spherical_yn, cython_special.spherical_yn, ('ld', 'ldb', 'lD', 'lDb'), + "Python version supports negative reals; Cython version doesn't - see gh-21629"), + (special.spherical_in, cython_special.spherical_in, ('ld', 'ldb', 'lD', 'lDb'), + "Python version supports negative reals; Cython version doesn't - see gh-21629"), + (special.spherical_kn, cython_special.spherical_kn, ('ld', 'ldb', 'lD', 'lDb'), + "Python version supports negative reals; Cython version doesn't - see gh-21629"), + (special.shichi, cython_special._shichi_pywrap, ('d', 'D'), None), + (special.sici, cython_special._sici_pywrap, ('d', 'D'), None), + (special.sindg, cython_special.sindg, ('d',), None), + (special.smirnov, cython_special.smirnov, ('ld', 'dd'), None), + (special.smirnovi, cython_special.smirnovi, ('ld', 'dd'), None), + (special.spence, cython_special.spence, ('d', 'D'), None), + (special.stdtr, cython_special.stdtr, ('dd',), None), + (special.stdtridf, cython_special.stdtridf, ('dd',), None), + (special.stdtrit, cython_special.stdtrit, ('dd',), None), + (special.struve, cython_special.struve, ('dd',), None), + (special.tandg, cython_special.tandg, ('d',), None), + (special.tklmbda, cython_special.tklmbda, ('dd',), None), + (special.voigt_profile, cython_special.voigt_profile, ('ddd',), None), + (special.wofz, cython_special.wofz, ('D',), None), + (special.wright_bessel, cython_special.wright_bessel, ('ddd',), None), + (special.wrightomega, cython_special.wrightomega, ('D',), None), + (special.xlog1py, cython_special.xlog1py, ('dd', 'DD'), None), + (special.xlogy, cython_special.xlogy, ('dd', 'DD'), None), + (special.y0, cython_special.y0, ('d',), None), + (special.y1, cython_special.y1, ('d',), None), + (special.yn, cython_special.yn, ('ld', 'dd'), None), + (special.yv, cython_special.yv, ('dd', 'dD'), None), + (special.yve, cython_special.yve, ('dd', 'dD'), None), + (special.zetac, cython_special.zetac, ('d',), None), + (special.owens_t, cython_special.owens_t, ('dd',), None) +] + + +IDS = [x[0].__name__ for x in PARAMS] + + +def _generate_test_points(typecodes): + axes = tuple(TEST_POINTS[x] for x in typecodes) + pts = list(product(*axes)) + return pts + + +def test_cython_api_completeness(): + # Check that everything is tested + for name in dir(cython_special): + func = getattr(cython_special, name) + if callable(func) and not name.startswith('_'): + for _, cyfun, _, _ in PARAMS: + if cyfun is func: + break + else: + raise RuntimeError(f"{name} missing from tests!") + + +@pytest.mark.fail_slow(20) +@pytest.mark.parametrize("param", PARAMS, ids=IDS) +def test_cython_api(param): + pyfunc, cyfunc, specializations, knownfailure = param + if knownfailure: + pytest.xfail(reason=knownfailure) + + # Check which parameters are expected to be fused types + max_params = max(len(spec) for spec in specializations) + values = [set() for _ in range(max_params)] + for typecodes in specializations: + for j, v in enumerate(typecodes): + values[j].add(v) + seen = set() + is_fused_code = [False] * len(values) + for j, v in enumerate(values): + vv = tuple(sorted(v)) + if vv in seen: + continue + is_fused_code[j] = (len(v) > 1) + seen.add(vv) + + # Check results + for typecodes in specializations: + # Pick the correct specialized function + signature = [CYTHON_SIGNATURE_MAP[code] + for j, code in enumerate(typecodes) + if is_fused_code[j]] + + if signature: + cy_spec_func = cyfunc[tuple(signature)] + else: + signature = None + cy_spec_func = cyfunc + + # Test it + pts = _generate_test_points(typecodes) + for pt in pts: + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + pyval = pyfunc(*pt) + cyval = cy_spec_func(*pt) + assert_allclose(cyval, pyval, err_msg=f"{pt} {typecodes} {signature}") diff --git a/v12x/pylib/scipy/special/tests/test_data.py b/v12x/pylib/scipy/special/tests/test_data.py new file mode 100644 index 0000000000000000000000000000000000000000..3d44ff9bf96eb45dfe40c2b77b112c9a6e300243 --- /dev/null +++ b/v12x/pylib/scipy/special/tests/test_data.py @@ -0,0 +1,679 @@ +import importlib.resources +import warnings + +import numpy as np +import pytest + +from scipy.special import ( + lpmv, lqn, lqmn, eval_legendre, eval_hermite, + eval_laguerre, eval_genlaguerre, binom, cbrt, expm1, log1p, zeta, + jn, jv, jvp, yn, yv, yvp, iv, ivp, kn, kv, kvp, + gamma, gammaln, gammainc, gammaincc, gammaincinv, gammainccinv, digamma, + beta, betainc, betaincinv, poch, + ellipe, ellipeinc, ellipk, ellipkm1, ellipkinc, + elliprc, elliprd, elliprf, elliprg, elliprj, + erf, erfc, erfinv, erfcinv, exp1, expi, expn, + bdtrik, btdtria, btdtrib, chndtr, gdtr, gdtrc, gdtrix, gdtrib, + nbdtrik, pdtrik, owens_t, + mathieu_a, mathieu_b, mathieu_cem, mathieu_sem, mathieu_modcem1, + mathieu_modsem1, mathieu_modcem2, mathieu_modsem2, + ellip_harm, ellip_harm_2, spherical_jn, spherical_yn, wright_bessel +) +from scipy.integrate import IntegrationWarning + +from scipy.special._testutils import FuncData + + +# The npz files are generated, and hence may live in the build dir. We can only +# access them through `importlib.resources`, not an explicit path from `__file__` +_datadir = importlib.resources.files('scipy.special.tests.data') + +_boost_npz = _datadir.joinpath('boost.npz') +with importlib.resources.as_file(_boost_npz) as f: + DATASETS_BOOST = np.load(f) + +_gsl_npz = _datadir.joinpath('gsl.npz') +with importlib.resources.as_file(_gsl_npz) as f: + DATASETS_GSL = np.load(f) + +_local_npz = _datadir.joinpath('local.npz') +with importlib.resources.as_file(_local_npz) as f: + DATASETS_LOCAL = np.load(f) + + +def data(func, dataname, *a, **kw): + kw.setdefault('dataname', dataname) + return FuncData(func, DATASETS_BOOST[dataname], *a, **kw) + + +def data_gsl(func, dataname, *a, **kw): + kw.setdefault('dataname', dataname) + return FuncData(func, DATASETS_GSL[dataname], *a, **kw) + + +def data_local(func, dataname, *a, **kw): + kw.setdefault('dataname', dataname) + return FuncData(func, DATASETS_LOCAL[dataname], *a, **kw) + + +def ellipk_(k): + return ellipk(k*k) + + +def ellipkinc_(f, k): + return ellipkinc(f, k*k) + + +def ellipe_(k): + return ellipe(k*k) + + +def ellipeinc_(f, k): + return ellipeinc(f, k*k) + + +def zeta_(x): + return zeta(x, 1.) + + +def assoc_legendre_p_boost_(nu, mu, x): + # the boost test data is for integer orders only + return lpmv(mu, nu.astype(int), x) + +def legendre_p_via_assoc_(nu, x): + return lpmv(0, nu, x) + +def lqn_(n, x): + return lqn(n.astype('l'), x)[0][-1] + +def legendre_q_via_lqmn(n, x): + return lqmn(0, n, x)[0][0,-1] + +def mathieu_ce_rad(m, q, x): + return mathieu_cem(m, q, x*180/np.pi)[0] + + +def mathieu_se_rad(m, q, x): + return mathieu_sem(m, q, x*180/np.pi)[0] + + +def mathieu_mc1_scaled(m, q, x): + # GSL follows a different normalization. + # We follow Abramowitz & Stegun, they apparently something else. + return mathieu_modcem1(m, q, x)[0] * np.sqrt(np.pi/2) + + +def mathieu_ms1_scaled(m, q, x): + return mathieu_modsem1(m, q, x)[0] * np.sqrt(np.pi/2) + + +def mathieu_mc2_scaled(m, q, x): + return mathieu_modcem2(m, q, x)[0] * np.sqrt(np.pi/2) + + +def mathieu_ms2_scaled(m, q, x): + return mathieu_modsem2(m, q, x)[0] * np.sqrt(np.pi/2) + +def eval_legendre_ld(n, x): + return eval_legendre(n.astype('l'), x) + +def eval_legendre_dd(n, x): + return eval_legendre(n.astype('d'), x) + +def eval_hermite_ld(n, x): + return eval_hermite(n.astype('l'), x) + +def eval_laguerre_ld(n, x): + return eval_laguerre(n.astype('l'), x) + +def eval_laguerre_dd(n, x): + return eval_laguerre(n.astype('d'), x) + +def eval_genlaguerre_ldd(n, a, x): + return eval_genlaguerre(n.astype('l'), a, x) + +def eval_genlaguerre_ddd(n, a, x): + return eval_genlaguerre(n.astype('d'), a, x) + +def bdtrik_comp(y, n, p): + return bdtrik(1-y, n, p) + +def btdtria_comp(p, b, x): + return btdtria(1-p, b, x) + +def btdtrib_comp(a, p, x): + return btdtrib(a, 1-p, x) + +def gdtr_(p, x): + return gdtr(1.0, p, x) + +def gdtrc_(p, x): + return gdtrc(1.0, p, x) + +def gdtrix_(b, p): + return gdtrix(1.0, b, p) + +def gdtrix_comp(b, p): + return gdtrix(1.0, b, 1-p) + +def gdtrib_(p, x): + return gdtrib(1.0, p, x) + +def gdtrib_comp(p, x): + return gdtrib(1.0, 1-p, x) + +def nbdtrik_comp(y, n, p): + return nbdtrik(1-y, n, p) + +def pdtrik_comp(p, m): + return pdtrik(1-p, m) + +def poch_(z, m): + return 1.0 / poch(z, m) + +def poch_minus(z, m): + return 1.0 / poch(z, -m) + +def spherical_jn_(n, x): + return spherical_jn(n.astype('l'), x) + +def spherical_yn_(n, x): + return spherical_yn(n.astype('l'), x) + +def cexpm1(x, y): + z = expm1(x + 1j*y) + return z.real, z.imag + +def clog1p(x, y): + z = log1p(x + 1j*y) + return z.real, z.imag + + +BOOST_TESTS = [ + data(assoc_legendre_p_boost_, 'assoc_legendre_p_ipp-assoc_legendre_p', + (0,1,2), 3, rtol=1e-11), + + data(legendre_p_via_assoc_, 'legendre_p_ipp-legendre_p', + (0,1), 2, rtol=1e-11), + data(legendre_p_via_assoc_, 'legendre_p_large_ipp-legendre_p_large', + (0,1), 2, rtol=9.6e-14), + data(eval_legendre_ld, 'legendre_p_ipp-legendre_p', + (0,1), 2, rtol=6e-14), + data(eval_legendre_ld, 'legendre_p_large_ipp-legendre_p_large', + (0,1), 2, rtol=2e-13), + data(eval_legendre_dd, 'legendre_p_ipp-legendre_p', + (0,1), 2, rtol=2e-14), + data(eval_legendre_dd, 'legendre_p_large_ipp-legendre_p_large', + (0,1), 2, rtol=2e-13), + + data(lqn_, 'legendre_p_ipp-legendre_p', + (0,1), 3, rtol=2e-14, vectorized=False), + data(lqn_, 'legendre_p_large_ipp-legendre_p_large', + (0,1), 3, rtol=2e-12, vectorized=False), + data(legendre_q_via_lqmn, 'legendre_p_ipp-legendre_p', + (0,1), 3, rtol=2e-14, vectorized=False), + data(legendre_q_via_lqmn, 'legendre_p_large_ipp-legendre_p_large', + (0,1), 3, rtol=2e-12, vectorized=False), + + data(beta, 'beta_exp_data_ipp-beta_exp_data', + (0,1), 2, rtol=1e-13), + data(beta, 'beta_exp_data_ipp-beta_exp_data', + (0,1), 2, rtol=1e-13), + data(beta, 'beta_med_data_ipp-beta_med_data', + (0,1), 2, rtol=5e-13), + + data(betainc, 'ibeta_small_data_ipp-ibeta_small_data', + (0,1,2), 5, rtol=6e-15), + data(betainc, 'ibeta_data_ipp-ibeta_data', + (0,1,2), 5, rtol=5e-13), + data(betainc, 'ibeta_int_data_ipp-ibeta_int_data', + (0,1,2), 5, rtol=2e-14), + data(betainc, 'ibeta_large_data_ipp-ibeta_large_data', + (0,1,2), 5, rtol=4e-10), + + data(betaincinv, 'ibeta_inv_data_ipp-ibeta_inv_data', + (0,1,2), 3, rtol=1e-5), + + data(btdtria, 'ibeta_inva_data_ipp-ibeta_inva_data', + (2,0,1), 3, rtol=5e-9), + data(btdtria_comp, 'ibeta_inva_data_ipp-ibeta_inva_data', + (2,0,1), 4, rtol=5e-9), + + data(btdtrib, 'ibeta_inva_data_ipp-ibeta_inva_data', + (0,2,1), 5, rtol=5e-9), + data(btdtrib_comp, 'ibeta_inva_data_ipp-ibeta_inva_data', + (0,2,1), 6, rtol=5e-9), + + data(binom, 'binomial_data_ipp-binomial_data', + (0,1), 2, rtol=1e-13), + data(binom, 'binomial_large_data_ipp-binomial_large_data', + (0,1), 2, rtol=5e-13), + + data(bdtrik, 'binomial_quantile_ipp-binomial_quantile_data', + (2,0,1), 3, rtol=5e-9), + data(bdtrik_comp, 'binomial_quantile_ipp-binomial_quantile_data', + (2,0,1), 4, rtol=5e-9), + + data(nbdtrik, 'negative_binomial_quantile_ipp-negative_binomial_quantile_data', + (2,0,1), 3, rtol=4e-9), + data(nbdtrik_comp, + 'negative_binomial_quantile_ipp-negative_binomial_quantile_data', + (2,0,1), 4, rtol=4e-9), + + data(pdtrik, 'poisson_quantile_ipp-poisson_quantile_data', + (1,0), 2, rtol=3e-9), + data(pdtrik_comp, 'poisson_quantile_ipp-poisson_quantile_data', + (1,0), 3, rtol=4e-9), + + data(cbrt, 'cbrt_data_ipp-cbrt_data', 1, 0), + + data(digamma, 'digamma_data_ipp-digamma_data', 0, 1), + data(digamma, 'digamma_data_ipp-digamma_data', 0j, 1), + data(digamma, 'digamma_neg_data_ipp-digamma_neg_data', 0, 1, rtol=2e-13), + data(digamma, 'digamma_neg_data_ipp-digamma_neg_data', 0j, 1, rtol=1e-13), + data(digamma, 'digamma_root_data_ipp-digamma_root_data', 0, 1, rtol=1e-15), + data(digamma, 'digamma_root_data_ipp-digamma_root_data', 0j, 1, rtol=1e-15), + data(digamma, 'digamma_small_data_ipp-digamma_small_data', 0, 1, rtol=1e-15), + data(digamma, 'digamma_small_data_ipp-digamma_small_data', 0j, 1, rtol=1e-14), + + data(ellipk_, 'ellint_k_data_ipp-ellint_k_data', 0, 1), + data(ellipkinc_, 'ellint_f_data_ipp-ellint_f_data', (0,1), 2, rtol=1e-14), + data(ellipe_, 'ellint_e_data_ipp-ellint_e_data', 0, 1), + data(ellipeinc_, 'ellint_e2_data_ipp-ellint_e2_data', (0,1), 2, rtol=1e-14), + + data(erf, 'erf_data_ipp-erf_data', 0, 1), + data(erf, 'erf_data_ipp-erf_data', 0j, 1, rtol=1e-13), + data(erfc, 'erf_data_ipp-erf_data', 0, 2, rtol=6e-15), + data(erf, 'erf_large_data_ipp-erf_large_data', 0, 1), + data(erf, 'erf_large_data_ipp-erf_large_data', 0j, 1), + data(erfc, 'erf_large_data_ipp-erf_large_data', 0, 2, rtol=4e-14), + data(erf, 'erf_small_data_ipp-erf_small_data', 0, 1), + data(erf, 'erf_small_data_ipp-erf_small_data', 0j, 1, rtol=1e-13), + data(erfc, 'erf_small_data_ipp-erf_small_data', 0, 2), + + data(erfinv, 'erf_inv_data_ipp-erf_inv_data', 0, 1), + data(erfcinv, 'erfc_inv_data_ipp-erfc_inv_data', 0, 1), + data(erfcinv, 'erfc_inv_big_data_ipp-erfc_inv_big_data', 0, 1, + param_filter=(lambda s: s > 0)), + + data(exp1, 'expint_1_data_ipp-expint_1_data', 1, 2, rtol=1e-13), + data(exp1, 'expint_1_data_ipp-expint_1_data', 1j, 2, rtol=5e-9), + data(expi, 'expinti_data_ipp-expinti_data', 0, 1, rtol=1e-13), + data(expi, 'expinti_data_double_ipp-expinti_data_double', 0, 1, rtol=1e-13), + data(expi, 'expinti_data_long_ipp-expinti_data_long', 0, 1), + + data(expn, 'expint_small_data_ipp-expint_small_data', (0,1), 2), + data(expn, 'expint_data_ipp-expint_data', (0,1), 2, rtol=1e-14), + + data(gamma, 'test_gamma_data_ipp-near_0', 0, 1), + data(gamma, 'test_gamma_data_ipp-near_1', 0, 1), + data(gamma, 'test_gamma_data_ipp-near_2', 0, 1), + data(gamma, 'test_gamma_data_ipp-near_m10', 0, 1), + data(gamma, 'test_gamma_data_ipp-near_m55', 0, 1, rtol=7e-12), + data(gamma, 'test_gamma_data_ipp-factorials', 0, 1, rtol=4e-14), + data(gamma, 'test_gamma_data_ipp-near_0', 0j, 1, rtol=2e-9), + data(gamma, 'test_gamma_data_ipp-near_1', 0j, 1, rtol=2e-9), + data(gamma, 'test_gamma_data_ipp-near_2', 0j, 1, rtol=2e-9), + data(gamma, 'test_gamma_data_ipp-near_m10', 0j, 1, rtol=2e-9), + data(gamma, 'test_gamma_data_ipp-near_m55', 0j, 1, rtol=2e-9), + data(gamma, 'test_gamma_data_ipp-factorials', 0j, 1, rtol=2e-13), + data(gammaln, 'test_gamma_data_ipp-near_0', 0, 2, rtol=5e-11), + data(gammaln, 'test_gamma_data_ipp-near_1', 0, 2, rtol=5e-11), + data(gammaln, 'test_gamma_data_ipp-near_2', 0, 2, rtol=2e-10), + data(gammaln, 'test_gamma_data_ipp-near_m10', 0, 2, rtol=5e-11), + data(gammaln, 'test_gamma_data_ipp-near_m55', 0, 2, rtol=5e-11), + data(gammaln, 'test_gamma_data_ipp-factorials', 0, 2), + + data(gammainc, 'igamma_small_data_ipp-igamma_small_data', (0,1), 5, rtol=5e-15), + data(gammainc, 'igamma_med_data_ipp-igamma_med_data', (0,1), 5, rtol=2e-13), + data(gammainc, 'igamma_int_data_ipp-igamma_int_data', (0,1), 5, rtol=2e-13), + data(gammainc, 'igamma_big_data_ipp-igamma_big_data', (0,1), 5, rtol=1e-12), + + data(gdtr_, 'igamma_small_data_ipp-igamma_small_data', (0,1), 5, rtol=1e-13), + data(gdtr_, 'igamma_med_data_ipp-igamma_med_data', (0,1), 5, rtol=2e-13), + data(gdtr_, 'igamma_int_data_ipp-igamma_int_data', (0,1), 5, rtol=2e-13), + data(gdtr_, 'igamma_big_data_ipp-igamma_big_data', (0,1), 5, rtol=2e-9), + + data(gammaincc, 'igamma_small_data_ipp-igamma_small_data', + (0,1), 3, rtol=1e-13), + data(gammaincc, 'igamma_med_data_ipp-igamma_med_data', + (0,1), 3, rtol=2e-13), + data(gammaincc, 'igamma_int_data_ipp-igamma_int_data', + (0,1), 3, rtol=4e-14), + data(gammaincc, 'igamma_big_data_ipp-igamma_big_data', + (0,1), 3, rtol=1e-11), + + data(gdtrc_, 'igamma_small_data_ipp-igamma_small_data', (0,1), 3, rtol=1e-13), + data(gdtrc_, 'igamma_med_data_ipp-igamma_med_data', (0,1), 3, rtol=2e-13), + data(gdtrc_, 'igamma_int_data_ipp-igamma_int_data', (0,1), 3, rtol=4e-14), + data(gdtrc_, 'igamma_big_data_ipp-igamma_big_data', (0,1), 3, rtol=1e-11), + + data(gdtrib_, 'igamma_inva_data_ipp-igamma_inva_data', (1,0), 2, rtol=5e-9), + data(gdtrib_comp, 'igamma_inva_data_ipp-igamma_inva_data', (1,0), 3, rtol=5e-9), + + data(poch_, 'tgamma_delta_ratio_data_ipp-tgamma_delta_ratio_data', + (0,1), 2, rtol=2e-13), + data(poch_, 'tgamma_delta_ratio_int_ipp-tgamma_delta_ratio_int', + (0,1), 2,), + data(poch_, 'tgamma_delta_ratio_int2_ipp-tgamma_delta_ratio_int2', + (0,1), 2,), + data(poch_minus, 'tgamma_delta_ratio_data_ipp-tgamma_delta_ratio_data', + (0,1), 3, rtol=2e-13), + data(poch_minus, 'tgamma_delta_ratio_int_ipp-tgamma_delta_ratio_int', + (0,1), 3), + data(poch_minus, 'tgamma_delta_ratio_int2_ipp-tgamma_delta_ratio_int2', + (0,1), 3), + + data(eval_hermite_ld, 'hermite_ipp-hermite', + (0,1), 2, rtol=2e-14), + + data(eval_laguerre_ld, 'laguerre2_ipp-laguerre2', + (0,1), 2, rtol=7e-12), + data(eval_laguerre_dd, 'laguerre2_ipp-laguerre2', + (0,1), 2, knownfailure='hyp2f1 insufficiently accurate.'), + data(eval_genlaguerre_ldd, 'laguerre3_ipp-laguerre3', + (0,1,2), 3, rtol=2e-13), + data(eval_genlaguerre_ddd, 'laguerre3_ipp-laguerre3', + (0,1,2), 3, knownfailure='hyp2f1 insufficiently accurate.'), + + data(log1p, 'log1p_expm1_data_ipp-log1p_expm1_data', 0, 1), + data(expm1, 'log1p_expm1_data_ipp-log1p_expm1_data', 0, 2), + + data(iv, 'bessel_i_data_ipp-bessel_i_data', + (0,1), 2, rtol=1e-12), + data(iv, 'bessel_i_data_ipp-bessel_i_data', + (0,1j), 2, rtol=2e-10, atol=1e-306), + data(iv, 'bessel_i_int_data_ipp-bessel_i_int_data', + (0,1), 2, rtol=1e-9), + data(iv, 'bessel_i_int_data_ipp-bessel_i_int_data', + (0,1j), 2, rtol=2e-10), + + data(ivp, 'bessel_i_prime_int_data_ipp-bessel_i_prime_int_data', + (0,1), 2, rtol=1.2e-13), + data(ivp, 'bessel_i_prime_int_data_ipp-bessel_i_prime_int_data', + (0,1j), 2, rtol=1.2e-13, atol=1e-300), + + data(jn, 'bessel_j_int_data_ipp-bessel_j_int_data', (0,1), 2, rtol=1e-12), + data(jn, 'bessel_j_int_data_ipp-bessel_j_int_data', (0,1j), 2, rtol=1e-12), + data(jn, 'bessel_j_large_data_ipp-bessel_j_large_data', (0,1), 2, rtol=6e-11), + data(jn, 'bessel_j_large_data_ipp-bessel_j_large_data', (0,1j), 2, rtol=6e-11), + + data(jv, 'bessel_j_int_data_ipp-bessel_j_int_data', (0,1), 2, rtol=1e-12), + data(jv, 'bessel_j_int_data_ipp-bessel_j_int_data', (0,1j), 2, rtol=1e-12), + data(jv, 'bessel_j_data_ipp-bessel_j_data', (0,1), 2, rtol=1e-12), + data(jv, 'bessel_j_data_ipp-bessel_j_data', (0,1j), 2, rtol=1e-12), + + data(jvp, 'bessel_j_prime_int_data_ipp-bessel_j_prime_int_data', + (0,1), 2, rtol=1e-13), + data(jvp, 'bessel_j_prime_int_data_ipp-bessel_j_prime_int_data', + (0,1j), 2, rtol=1e-13), + data(jvp, 'bessel_j_prime_large_data_ipp-bessel_j_prime_large_data', + (0,1), 2, rtol=1e-11), + data(jvp, 'bessel_j_prime_large_data_ipp-bessel_j_prime_large_data', + (0,1j), 2, rtol=2e-11), + + data(kn, 'bessel_k_int_data_ipp-bessel_k_int_data', (0,1), 2, rtol=1e-12), + + data(kv, 'bessel_k_int_data_ipp-bessel_k_int_data', (0,1), 2, rtol=1e-12), + data(kv, 'bessel_k_int_data_ipp-bessel_k_int_data', (0,1j), 2, rtol=1e-12), + data(kv, 'bessel_k_data_ipp-bessel_k_data', (0,1), 2, rtol=1e-12), + data(kv, 'bessel_k_data_ipp-bessel_k_data', (0,1j), 2, rtol=1e-12), + + data(kvp, 'bessel_k_prime_int_data_ipp-bessel_k_prime_int_data', + (0,1), 2, rtol=3e-14), + data(kvp, 'bessel_k_prime_int_data_ipp-bessel_k_prime_int_data', + (0,1j), 2, rtol=3e-14), + data(kvp, 'bessel_k_prime_data_ipp-bessel_k_prime_data', (0,1), 2, rtol=7e-14), + data(kvp, 'bessel_k_prime_data_ipp-bessel_k_prime_data', (0,1j), 2, rtol=7e-14), + + data(yn, 'bessel_y01_data_ipp-bessel_y01_data', (0,1), 2, rtol=1e-12), + data(yn, 'bessel_yn_data_ipp-bessel_yn_data', (0,1), 2, rtol=1e-12), + + data(yv, 'bessel_yn_data_ipp-bessel_yn_data', (0,1), 2, rtol=1e-12), + data(yv, 'bessel_yn_data_ipp-bessel_yn_data', (0,1j), 2, rtol=1e-12), + data(yv, 'bessel_yv_data_ipp-bessel_yv_data', (0,1), 2, rtol=1e-10), + data(yv, 'bessel_yv_data_ipp-bessel_yv_data', (0,1j), 2, rtol=1e-10), + + data(yvp, 'bessel_yv_prime_data_ipp-bessel_yv_prime_data', + (0, 1), 2, rtol=4e-9), + data(yvp, 'bessel_yv_prime_data_ipp-bessel_yv_prime_data', + (0, 1j), 2, rtol=4e-9), + + data(zeta_, 'zeta_data_ipp-zeta_data', 0, 1, + param_filter=(lambda s: s > 1)), + data(zeta_, 'zeta_neg_data_ipp-zeta_neg_data', 0, 1, + param_filter=(lambda s: s > 1)), + data(zeta_, 'zeta_1_up_data_ipp-zeta_1_up_data', 0, 1, + param_filter=(lambda s: s > 1)), + data(zeta_, 'zeta_1_below_data_ipp-zeta_1_below_data', 0, 1, + param_filter=(lambda s: s > 1)), + + data(gammaincinv, 'gamma_inv_small_data_ipp-gamma_inv_small_data', + (0,1), 2, rtol=1e-11), + data(gammaincinv, 'gamma_inv_data_ipp-gamma_inv_data', + (0,1), 2, rtol=1e-14), + data(gammaincinv, 'gamma_inv_big_data_ipp-gamma_inv_big_data', + (0,1), 2, rtol=1e-11), + + data(gammainccinv, 'gamma_inv_small_data_ipp-gamma_inv_small_data', + (0,1), 3, rtol=1e-12), + data(gammainccinv, 'gamma_inv_data_ipp-gamma_inv_data', + (0,1), 3, rtol=1e-14), + data(gammainccinv, 'gamma_inv_big_data_ipp-gamma_inv_big_data', + (0,1), 3, rtol=1e-14), + + data(gdtrix_, 'gamma_inv_small_data_ipp-gamma_inv_small_data', + (0,1), 2, rtol=3e-13, knownfailure='gdtrix unflow some points'), + data(gdtrix_, 'gamma_inv_data_ipp-gamma_inv_data', + (0,1), 2, rtol=3e-15), + data(gdtrix_, 'gamma_inv_big_data_ipp-gamma_inv_big_data', + (0,1), 2), + data(gdtrix_comp, 'gamma_inv_small_data_ipp-gamma_inv_small_data', + (0,1), 2, knownfailure='gdtrix bad some points'), + data(gdtrix_comp, 'gamma_inv_data_ipp-gamma_inv_data', + (0,1), 3, rtol=6e-15), + data(gdtrix_comp, 'gamma_inv_big_data_ipp-gamma_inv_big_data', + (0,1), 3), + + data(chndtr, 'nccs_ipp-nccs', + (2,0,1), 3, rtol=3e-5), + data(chndtr, 'nccs_big_ipp-nccs_big', + (2,0,1), 3, rtol=5e-4, knownfailure='chndtr inaccurate some points'), + + data(spherical_jn_, 'sph_bessel_data_ipp-sph_bessel_data', + (0,1), 2, rtol=1e-13), + data(spherical_yn_, 'sph_neumann_data_ipp-sph_neumann_data', + (0,1), 2, rtol=8e-15), + + data(owens_t, 'owens_t_ipp-owens_t', + (0, 1), 2, rtol=5e-14), + data(owens_t, 'owens_t_large_data_ipp-owens_t_large_data', + (0, 1), 2, rtol=8e-12), + + # -- test data exists in boost but is not used in scipy -- + + # ibeta_derivative_data_ipp/ibeta_derivative_data.txt + # ibeta_derivative_int_data_ipp/ibeta_derivative_int_data.txt + # ibeta_derivative_large_data_ipp/ibeta_derivative_large_data.txt + # ibeta_derivative_small_data_ipp/ibeta_derivative_small_data.txt + + # bessel_y01_prime_data_ipp/bessel_y01_prime_data.txt + # bessel_yn_prime_data_ipp/bessel_yn_prime_data.txt + # sph_bessel_prime_data_ipp/sph_bessel_prime_data.txt + # sph_neumann_prime_data_ipp/sph_neumann_prime_data.txt + + # ellint_d2_data_ipp/ellint_d2_data.txt + # ellint_d_data_ipp/ellint_d_data.txt + # ellint_pi2_data_ipp/ellint_pi2_data.txt + # ellint_pi3_data_ipp/ellint_pi3_data.txt + # ellint_pi3_large_data_ipp/ellint_pi3_large_data.txt + data(elliprc, 'ellint_rc_data_ipp-ellint_rc_data', (0, 1), 2, + rtol=5e-16), + data(elliprd, 'ellint_rd_data_ipp-ellint_rd_data', (0, 1, 2), 3, + rtol=5e-16), + data(elliprd, 'ellint_rd_0xy_ipp-ellint_rd_0xy', (0, 1, 2), 3, + rtol=5e-16), + data(elliprd, 'ellint_rd_0yy_ipp-ellint_rd_0yy', (0, 1, 2), 3, + rtol=5e-16), + data(elliprd, 'ellint_rd_xxx_ipp-ellint_rd_xxx', (0, 1, 2), 3, + rtol=5e-16), + # Some of the following rtol for elliprd may be larger than 5e-16 to + # work around some hard cases in the Boost test where we get slightly + # larger error than the ideal bound when the x (==y) input is close to + # zero. + # Also the accuracy on 32-bit builds with g++ may suffer from excess + # loss of precision; see GCC bugzilla 323 + # https://gcc.gnu.org/bugzilla/show_bug.cgi?id=323 + data(elliprd, 'ellint_rd_xxz_ipp-ellint_rd_xxz', (0, 1, 2), 3, + rtol=6.5e-16), + data(elliprd, 'ellint_rd_xyy_ipp-ellint_rd_xyy', (0, 1, 2), 3, + rtol=6e-16), + data(elliprf, 'ellint_rf_data_ipp-ellint_rf_data', (0, 1, 2), 3, + rtol=5e-16), + data(elliprf, 'ellint_rf_xxx_ipp-ellint_rf_xxx', (0, 1, 2), 3, + rtol=5e-16), + data(elliprf, 'ellint_rf_xyy_ipp-ellint_rf_xyy', (0, 1, 2), 3, + rtol=5e-16), + data(elliprf, 'ellint_rf_xy0_ipp-ellint_rf_xy0', (0, 1, 2), 3, + rtol=5e-16), + data(elliprf, 'ellint_rf_0yy_ipp-ellint_rf_0yy', (0, 1, 2), 3, + rtol=5e-16), + # The accuracy of R_G is primarily limited by R_D that is used + # internally. It is generally worse than R_D. Notice that we increased + # the rtol for R_G here. The cases with duplicate arguments are + # slightly less likely to be unbalanced (at least two arguments are + # already balanced) so the error bound is slightly better. Again, + # precision with g++ 32-bit is even worse. + data(elliprg, 'ellint_rg_ipp-ellint_rg', (0, 1, 2), 3, + rtol=8.0e-16), + data(elliprg, 'ellint_rg_xxx_ipp-ellint_rg_xxx', (0, 1, 2), 3, + rtol=6e-16), + data(elliprg, 'ellint_rg_xyy_ipp-ellint_rg_xyy', (0, 1, 2), 3, + rtol=7.5e-16), + data(elliprg, 'ellint_rg_xy0_ipp-ellint_rg_xy0', (0, 1, 2), 3, + rtol=5.2e-16), + data(elliprg, 'ellint_rg_00x_ipp-ellint_rg_00x', (0, 1, 2), 3, + rtol=5e-16), + data(elliprj, 'ellint_rj_data_ipp-ellint_rj_data', (0, 1, 2, 3), 4, + rtol=5e-16, atol=1e-25, + param_filter=(lambda s: s <= 5e-26,)), + # ellint_rc_data_ipp/ellint_rc_data.txt + # ellint_rd_0xy_ipp/ellint_rd_0xy.txt + # ellint_rd_0yy_ipp/ellint_rd_0yy.txt + # ellint_rd_data_ipp/ellint_rd_data.txt + # ellint_rd_xxx_ipp/ellint_rd_xxx.txt + # ellint_rd_xxz_ipp/ellint_rd_xxz.txt + # ellint_rd_xyy_ipp/ellint_rd_xyy.txt + # ellint_rf_0yy_ipp/ellint_rf_0yy.txt + # ellint_rf_data_ipp/ellint_rf_data.txt + # ellint_rf_xxx_ipp/ellint_rf_xxx.txt + # ellint_rf_xy0_ipp/ellint_rf_xy0.txt + # ellint_rf_xyy_ipp/ellint_rf_xyy.txt + # ellint_rg_00x_ipp/ellint_rg_00x.txt + # ellint_rg_ipp/ellint_rg.txt + # ellint_rg_xxx_ipp/ellint_rg_xxx.txt + # ellint_rg_xy0_ipp/ellint_rg_xy0.txt + # ellint_rg_xyy_ipp/ellint_rg_xyy.txt + # ellint_rj_data_ipp/ellint_rj_data.txt + # ellint_rj_e2_ipp/ellint_rj_e2.txt + # ellint_rj_e3_ipp/ellint_rj_e3.txt + # ellint_rj_e4_ipp/ellint_rj_e4.txt + # ellint_rj_zp_ipp/ellint_rj_zp.txt + + # jacobi_elliptic_ipp/jacobi_elliptic.txt + # jacobi_elliptic_small_ipp/jacobi_elliptic_small.txt + # jacobi_large_phi_ipp/jacobi_large_phi.txt + # jacobi_near_1_ipp/jacobi_near_1.txt + # jacobi_zeta_big_phi_ipp/jacobi_zeta_big_phi.txt + # jacobi_zeta_data_ipp/jacobi_zeta_data.txt + + # heuman_lambda_data_ipp/heuman_lambda_data.txt + + # hypergeometric_0F2_ipp/hypergeometric_0F2.txt + # hypergeometric_1F1_big_ipp/hypergeometric_1F1_big.txt + # hypergeometric_1F1_ipp/hypergeometric_1F1.txt + # hypergeometric_1F1_small_random_ipp/hypergeometric_1F1_small_random.txt + # hypergeometric_1F2_ipp/hypergeometric_1F2.txt + # hypergeometric_1f1_large_regularized_ipp/hypergeometric_1f1_large_regularized.txt # noqa: E501 + # hypergeometric_1f1_log_large_unsolved_ipp/hypergeometric_1f1_log_large_unsolved.txt # noqa: E501 + # hypergeometric_2F0_half_ipp/hypergeometric_2F0_half.txt + # hypergeometric_2F0_integer_a2_ipp/hypergeometric_2F0_integer_a2.txt + # hypergeometric_2F0_ipp/hypergeometric_2F0.txt + # hypergeometric_2F0_large_z_ipp/hypergeometric_2F0_large_z.txt + # hypergeometric_2F1_ipp/hypergeometric_2F1.txt + # hypergeometric_2F2_ipp/hypergeometric_2F2.txt + + # ncbeta_big_ipp/ncbeta_big.txt + # nct_small_delta_ipp/nct_small_delta.txt + # nct_asym_ipp/nct_asym.txt + # ncbeta_ipp/ncbeta.txt + + # powm1_data_ipp/powm1_big_data.txt + # powm1_sqrtp1m1_test_hpp/sqrtp1m1_data.txt + + # sinc_data_ipp/sinc_data.txt + + # test_gamma_data_ipp/gammap1m1_data.txt + # tgamma_ratio_data_ipp/tgamma_ratio_data.txt + + # trig_data_ipp/trig_data.txt + # trig_data2_ipp/trig_data2.txt +] + + +@pytest.mark.parametrize('test', BOOST_TESTS, ids=repr) +def test_boost(test): + _test_factory(test) + + +GSL_TESTS = [ + data_gsl(mathieu_a, 'mathieu_ab', (0, 1), 2, rtol=1e-13, atol=1e-13), + data_gsl(mathieu_b, 'mathieu_ab', (0, 1), 3, rtol=1e-13, atol=1e-13), + + # Also the GSL output has limited accuracy... + data_gsl(mathieu_ce_rad, 'mathieu_ce_se', (0, 1, 2), 3, rtol=1e-7, atol=1e-13), + data_gsl(mathieu_se_rad, 'mathieu_ce_se', (0, 1, 2), 4, rtol=1e-7, atol=1e-13), + + data_gsl(mathieu_mc1_scaled, 'mathieu_mc_ms', + (0, 1, 2), 3, rtol=1e-7, atol=1e-13), + data_gsl(mathieu_ms1_scaled, 'mathieu_mc_ms', + (0, 1, 2), 4, rtol=1e-7, atol=1e-13), + + data_gsl(mathieu_mc2_scaled, 'mathieu_mc_ms', + (0, 1, 2), 5, rtol=1e-7, atol=1e-13), + data_gsl(mathieu_ms2_scaled, 'mathieu_mc_ms', + (0, 1, 2), 6, rtol=1e-7, atol=1e-13), +] + + +@pytest.mark.parametrize('test', GSL_TESTS, ids=repr) +def test_gsl(test): + _test_factory(test) + + +LOCAL_TESTS = [ + data_local(ellipkinc, 'ellipkinc_neg_m', (0, 1), 2), + data_local(ellipkm1, 'ellipkm1', 0, 1), + data_local(ellipeinc, 'ellipeinc_neg_m', (0, 1), 2), + data_local(clog1p, 'log1p_expm1_complex', (0,1), (2,3), rtol=1e-14), + data_local(cexpm1, 'log1p_expm1_complex', (0,1), (4,5), rtol=1e-14), + data_local(gammainc, 'gammainc', (0, 1), 2, rtol=1e-12), + data_local(gammaincc, 'gammaincc', (0, 1), 2, rtol=1e-11), + data_local(ellip_harm_2, 'ellip',(0, 1, 2, 3, 4), 6, rtol=1e-10, atol=1e-13), + data_local(ellip_harm, 'ellip',(0, 1, 2, 3, 4), 5, rtol=1e-10, atol=1e-13), + data_local(wright_bessel, 'wright_bessel', (0, 1, 2), 3, rtol=1e-11), +] + + +@pytest.mark.parametrize('test', LOCAL_TESTS, ids=repr) +def test_local(test): + _test_factory(test) + + +def _test_factory(test, dtype=np.float64): + """Boost test""" + with warnings.catch_warnings(): + msg = "The occurrence of roundoff error is detected" + warnings.filterwarnings("ignore", msg, IntegrationWarning) + with np.errstate(all='ignore'): + test.check(dtype=dtype) diff --git a/v12x/pylib/scipy/special/tests/test_dd.py b/v12x/pylib/scipy/special/tests/test_dd.py new file mode 100644 index 0000000000000000000000000000000000000000..6da2c8ddddd7341cbb777216e9f2f3cce536a51e --- /dev/null +++ b/v12x/pylib/scipy/special/tests/test_dd.py @@ -0,0 +1,42 @@ +# Tests for a few of the "double-double" C++ functions defined in +# special/cephes/dd_real.h. Prior to gh-20390 which translated these +# functions from C to C++, there were test cases for _dd_expm1. It +# was determined that this function is not used anywhere internally +# in SciPy, so this function was not translated. + + +import pytest +from numpy.testing import assert_allclose +from scipy.special._test_internal import _dd_exp, _dd_log + + +# Each tuple in test_data contains: +# (dd_func, xhi, xlo, expected_yhi, expected_ylo) +# The expected values were computed with mpmath, e.g. +# +# import mpmath +# mpmath.mp.dps = 100 +# xhi = 10.0 +# xlo = 0.0 +# x = mpmath.mpf(xhi) + mpmath.mpf(xlo) +# y = mpmath.log(x) +# expected_yhi = float(y) +# expected_ylo = float(y - expected_yhi) +# +test_data = [ + (_dd_exp, -0.3333333333333333, -1.850371707708594e-17, + 0.7165313105737893, -2.0286948382455594e-17), + (_dd_exp, 0.0, 0.0, 1.0, 0.0), + (_dd_exp, 10.0, 0.0, 22026.465794806718, -1.3780134700517372e-12), + (_dd_log, 0.03125, 0.0, -3.4657359027997265, -4.930038229799327e-18), + (_dd_log, 10.0, 0.0, 2.302585092994046, -2.1707562233822494e-16), +] + + +@pytest.mark.parametrize('dd_func, xhi, xlo, expected_yhi, expected_ylo', + test_data) +def test_dd(dd_func, xhi, xlo, expected_yhi, expected_ylo): + yhi, ylo = dd_func(xhi, xlo) + assert yhi == expected_yhi, (f"high double ({yhi}) does not equal the " + f"expected value {expected_yhi}") + assert_allclose(ylo, expected_ylo, rtol=5e-15) diff --git a/v12x/pylib/scipy/special/tests/test_digamma.py b/v12x/pylib/scipy/special/tests/test_digamma.py new file mode 100644 index 0000000000000000000000000000000000000000..d7f27dc7b71c1ae928b4bdd8bd987df9ca420bab --- /dev/null +++ b/v12x/pylib/scipy/special/tests/test_digamma.py @@ -0,0 +1,45 @@ +import numpy as np +from numpy import pi, log, sqrt +from numpy.testing import assert_, assert_equal + +from scipy.special._testutils import FuncData +import scipy.special as sc + +# Euler-Mascheroni constant +euler = 0.57721566490153286 + + +def test_consistency(): + # Make sure the implementation of digamma for real arguments + # agrees with the implementation of digamma for complex arguments. + + # It's all poles after -1e16 + x = np.r_[-np.logspace(15, -30, 200), np.logspace(-30, 300, 200)] + dataset = np.vstack((x + 0j, sc.digamma(x))).T + FuncData(sc.digamma, dataset, 0, 1, rtol=5e-14, nan_ok=True).check() + + +def test_special_values(): + # Test special values from Gauss's digamma theorem. See + # + # https://en.wikipedia.org/wiki/Digamma_function + + dataset = [ + (1, -euler), + (0.5, -2*log(2) - euler), + (1/3, -pi/(2*sqrt(3)) - 3*log(3)/2 - euler), + (1/4, -pi/2 - 3*log(2) - euler), + (1/6, -pi*sqrt(3)/2 - 2*log(2) - 3*log(3)/2 - euler), + (1/8, + -pi/2 - 4*log(2) - (pi + log(2 + sqrt(2)) - log(2 - sqrt(2)))/sqrt(2) - euler) + ] + + dataset = np.asarray(dataset) + FuncData(sc.digamma, dataset, 0, 1, rtol=1e-14).check() + + +def test_nonfinite(): + pts = [0.0, -0.0, np.inf] + std = [-np.inf, np.inf, np.inf] + assert_equal(sc.digamma(pts), std) + assert_(all(np.isnan(sc.digamma([-np.inf, -1])))) diff --git a/v12x/pylib/scipy/special/tests/test_ellip_harm.py b/v12x/pylib/scipy/special/tests/test_ellip_harm.py new file mode 100644 index 0000000000000000000000000000000000000000..68526bc31e7a6caedae7de1cb8f033f9b221062a --- /dev/null +++ b/v12x/pylib/scipy/special/tests/test_ellip_harm.py @@ -0,0 +1,288 @@ +# +# Tests for the Ellipsoidal Harmonic Function, +# Distributed under the same license as SciPy itself. +# +import warnings + +import numpy as np +from numpy.testing import assert_equal, assert_allclose, assert_ +from scipy.special._testutils import assert_func_equal +from scipy.special import ellip_harm, ellip_harm_2, ellip_normal +from scipy.integrate import IntegrationWarning +from numpy import sqrt, pi + + +def test_ellip_potential(): + def change_coefficient(lambda1, mu, nu, h2, k2): + x = sqrt(lambda1**2*mu**2*nu**2/(h2*k2)) + y = sqrt((lambda1**2 - h2)*(mu**2 - h2)*(h2 - nu**2)/(h2*(k2 - h2))) + z = sqrt((lambda1**2 - k2)*(k2 - mu**2)*(k2 - nu**2)/(k2*(k2 - h2))) + return x, y, z + + def solid_int_ellip(lambda1, mu, nu, n, p, h2, k2): + return (ellip_harm(h2, k2, n, p, lambda1)*ellip_harm(h2, k2, n, p, mu) + * ellip_harm(h2, k2, n, p, nu)) + + def solid_int_ellip2(lambda1, mu, nu, n, p, h2, k2): + return (ellip_harm_2(h2, k2, n, p, lambda1) + * ellip_harm(h2, k2, n, p, mu)*ellip_harm(h2, k2, n, p, nu)) + + def summation(lambda1, mu1, nu1, lambda2, mu2, nu2, h2, k2): + tol = 1e-8 + sum1 = 0 + for n in range(20): + xsum = 0 + for p in range(1, 2*n+2): + xsum += (4*pi*(solid_int_ellip(lambda2, mu2, nu2, n, p, h2, k2) + * solid_int_ellip2(lambda1, mu1, nu1, n, p, h2, k2)) / + (ellip_normal(h2, k2, n, p)*(2*n + 1))) + if abs(xsum) < 0.1*tol*abs(sum1): + break + sum1 += xsum + return sum1, xsum + + def potential(lambda1, mu1, nu1, lambda2, mu2, nu2, h2, k2): + x1, y1, z1 = change_coefficient(lambda1, mu1, nu1, h2, k2) + x2, y2, z2 = change_coefficient(lambda2, mu2, nu2, h2, k2) + res = sqrt((x2 - x1)**2 + (y2 - y1)**2 + (z2 - z1)**2) + return 1/res + + pts = [ + (120, sqrt(19), 2, 41, sqrt(17), 2, 15, 25), + (120, sqrt(16), 3.2, 21, sqrt(11), 2.9, 11, 20), + ] + + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", "The occurrence of roundoff error", IntegrationWarning) + warnings.filterwarnings( + "ignore", "The maximum number of subdivisions", IntegrationWarning) + + for p in pts: + err_msg = repr(p) + exact = potential(*p) + result, last_term = summation(*p) + assert_allclose(exact, result, atol=0, rtol=1e-8, err_msg=err_msg) + assert_(abs(result - exact) < 10*abs(last_term), err_msg) + + +def test_ellip_norm(): + + def G01(h2, k2): + return 4*pi + + def G11(h2, k2): + return 4*pi*h2*k2/3 + + def G12(h2, k2): + return 4*pi*h2*(k2 - h2)/3 + + def G13(h2, k2): + return 4*pi*k2*(k2 - h2)/3 + + def G22(h2, k2): + res = (2*(h2**4 + k2**4) - 4*h2*k2*(h2**2 + k2**2) + 6*h2**2*k2**2 + + sqrt(h2**2 + k2**2 - h2*k2)*(-2*(h2**3 + k2**3) + 3*h2*k2*(h2 + k2))) + return 16*pi/405*res + + def G21(h2, k2): + res = (2*(h2**4 + k2**4) - 4*h2*k2*(h2**2 + k2**2) + 6*h2**2*k2**2 + + sqrt(h2**2 + k2**2 - h2*k2)*(2*(h2**3 + k2**3) - 3*h2*k2*(h2 + k2))) + return 16*pi/405*res + + def G23(h2, k2): + return 4*pi*h2**2*k2*(k2 - h2)/15 + + def G24(h2, k2): + return 4*pi*h2*k2**2*(k2 - h2)/15 + + def G25(h2, k2): + return 4*pi*h2*k2*(k2 - h2)**2/15 + + def G32(h2, k2): + res = (16*(h2**4 + k2**4) - 36*h2*k2*(h2**2 + k2**2) + 46*h2**2*k2**2 + + sqrt(4*(h2**2 + k2**2) - 7*h2*k2)*(-8*(h2**3 + k2**3) + + 11*h2*k2*(h2 + k2))) + return 16*pi/13125*k2*h2*res + + def G31(h2, k2): + res = (16*(h2**4 + k2**4) - 36*h2*k2*(h2**2 + k2**2) + 46*h2**2*k2**2 + + sqrt(4*(h2**2 + k2**2) - 7*h2*k2)*(8*(h2**3 + k2**3) - + 11*h2*k2*(h2 + k2))) + return 16*pi/13125*h2*k2*res + + def G34(h2, k2): + res = (6*h2**4 + 16*k2**4 - 12*h2**3*k2 - 28*h2*k2**3 + 34*h2**2*k2**2 + + sqrt(h2**2 + 4*k2**2 - h2*k2)*(-6*h2**3 - 8*k2**3 + 9*h2**2*k2 + + 13*h2*k2**2)) + return 16*pi/13125*h2*(k2 - h2)*res + + def G33(h2, k2): + res = (6*h2**4 + 16*k2**4 - 12*h2**3*k2 - 28*h2*k2**3 + 34*h2**2*k2**2 + + sqrt(h2**2 + 4*k2**2 - h2*k2)*(6*h2**3 + 8*k2**3 - 9*h2**2*k2 - + 13*h2*k2**2)) + return 16*pi/13125*h2*(k2 - h2)*res + + def G36(h2, k2): + res = (16*h2**4 + 6*k2**4 - 28*h2**3*k2 - 12*h2*k2**3 + 34*h2**2*k2**2 + + sqrt(4*h2**2 + k2**2 - h2*k2)*(-8*h2**3 - 6*k2**3 + 13*h2**2*k2 + + 9*h2*k2**2)) + return 16*pi/13125*k2*(k2 - h2)*res + + def G35(h2, k2): + res = (16*h2**4 + 6*k2**4 - 28*h2**3*k2 - 12*h2*k2**3 + 34*h2**2*k2**2 + + sqrt(4*h2**2 + k2**2 - h2*k2)*(8*h2**3 + 6*k2**3 - 13*h2**2*k2 - + 9*h2*k2**2)) + return 16*pi/13125*k2*(k2 - h2)*res + + def G37(h2, k2): + return 4*pi*h2**2*k2**2*(k2 - h2)**2/105 + + known_funcs = {(0, 1): G01, (1, 1): G11, (1, 2): G12, (1, 3): G13, + (2, 1): G21, (2, 2): G22, (2, 3): G23, (2, 4): G24, + (2, 5): G25, (3, 1): G31, (3, 2): G32, (3, 3): G33, + (3, 4): G34, (3, 5): G35, (3, 6): G36, (3, 7): G37} + + def _ellip_norm(n, p, h2, k2): + func = known_funcs[n, p] + return func(h2, k2) + _ellip_norm = np.vectorize(_ellip_norm) + + def ellip_normal_known(h2, k2, n, p): + return _ellip_norm(n, p, h2, k2) + + # generate both large and small h2 < k2 pairs + np.random.seed(1234) + h2 = np.random.pareto(0.5, size=1) + k2 = h2 * (1 + np.random.pareto(0.5, size=h2.size)) + + points = [] + for n in range(4): + for p in range(1, 2*n+2): + points.append((h2, k2, np.full(h2.size, n), np.full(h2.size, p))) + points = np.array(points) + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", "The occurrence of roundoff error", IntegrationWarning) + assert_func_equal(ellip_normal, ellip_normal_known, points, rtol=1e-12) + + +def test_ellip_harm_2(): + + def I1(h2, k2, s): + res = (ellip_harm_2(h2, k2, 1, 1, s)/(3 * ellip_harm(h2, k2, 1, 1, s)) + + ellip_harm_2(h2, k2, 1, 2, s)/(3 * ellip_harm(h2, k2, 1, 2, s)) + + ellip_harm_2(h2, k2, 1, 3, s)/(3 * ellip_harm(h2, k2, 1, 3, s))) + return res + + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", "The occurrence of roundoff error", IntegrationWarning) + assert_allclose(I1(5, 8, 10), 1/(10*sqrt((100-5)*(100-8))), + atol=1.5e-7, rtol=0) + + # Values produced by code from arXiv:1204.0267 + assert_allclose(ellip_harm_2(5, 8, 2, 1, 10), 0.00108056853382, + atol=1.5e-7, rtol=0) + assert_allclose(ellip_harm_2(5, 8, 2, 2, 10), 0.00105820513809, + atol=1.5e-7, rtol=0) + assert_allclose(ellip_harm_2(5, 8, 2, 3, 10), 0.00106058384743, + atol=1.5e-7, rtol=0) + assert_allclose(ellip_harm_2(5, 8, 2, 4, 10), 0.00106774492306, + atol=1.5e-7, rtol=0) + assert_allclose(ellip_harm_2(5, 8, 2, 5, 10), 0.00107976356454, + atol=1.5e-7, rtol=0) + + +def test_ellip_harm(): + + def E01(h2, k2, s): + return 1 + + def E11(h2, k2, s): + return s + + def E12(h2, k2, s): + return sqrt(abs(s*s - h2)) + + def E13(h2, k2, s): + return sqrt(abs(s*s - k2)) + + def E21(h2, k2, s): + return s*s - 1/3*((h2 + k2) + sqrt(abs((h2 + k2)*(h2 + k2)-3*h2*k2))) + + def E22(h2, k2, s): + return s*s - 1/3*((h2 + k2) - sqrt(abs((h2 + k2)*(h2 + k2)-3*h2*k2))) + + def E23(h2, k2, s): + return s * sqrt(abs(s*s - h2)) + + def E24(h2, k2, s): + return s * sqrt(abs(s*s - k2)) + + def E25(h2, k2, s): + return sqrt(abs((s*s - h2)*(s*s - k2))) + + def E31(h2, k2, s): + return s*s*s - (s/5)*(2*(h2 + k2) + sqrt(4*(h2 + k2)*(h2 + k2) - + 15*h2*k2)) + + def E32(h2, k2, s): + return s*s*s - (s/5)*(2*(h2 + k2) - sqrt(4*(h2 + k2)*(h2 + k2) - + 15*h2*k2)) + + def E33(h2, k2, s): + return sqrt(abs(s*s - h2))*(s*s - 1/5*((h2 + 2*k2) + sqrt(abs((h2 + + 2*k2)*(h2 + 2*k2) - 5*h2*k2)))) + + def E34(h2, k2, s): + return sqrt(abs(s*s - h2))*(s*s - 1/5*((h2 + 2*k2) - sqrt(abs((h2 + + 2*k2)*(h2 + 2*k2) - 5*h2*k2)))) + + def E35(h2, k2, s): + return sqrt(abs(s*s - k2))*(s*s - 1/5*((2*h2 + k2) + sqrt(abs((2*h2 + + k2)*(2*h2 + k2) - 5*h2*k2)))) + + def E36(h2, k2, s): + return sqrt(abs(s*s - k2))*(s*s - 1/5*((2*h2 + k2) - sqrt(abs((2*h2 + + k2)*(2*h2 + k2) - 5*h2*k2)))) + + def E37(h2, k2, s): + return s * sqrt(abs((s*s - h2)*(s*s - k2))) + + assert_equal(ellip_harm(5, 8, 1, 2, 2.5, 1, 1), + ellip_harm(5, 8, 1, 2, 2.5)) + + known_funcs = {(0, 1): E01, (1, 1): E11, (1, 2): E12, (1, 3): E13, + (2, 1): E21, (2, 2): E22, (2, 3): E23, (2, 4): E24, + (2, 5): E25, (3, 1): E31, (3, 2): E32, (3, 3): E33, + (3, 4): E34, (3, 5): E35, (3, 6): E36, (3, 7): E37} + + point_ref = [] + + def ellip_harm_known(h2, k2, n, p, s): + for i in range(h2.size): + func = known_funcs[(int(n[i]), int(p[i]))] + point_ref.append(func(h2[i], k2[i], s[i])) + return point_ref + + rng = np.random.RandomState(1234) + h2 = rng.pareto(0.5, size=30) + k2 = h2*(1 + rng.pareto(0.5, size=h2.size)) + s = rng.pareto(0.5, size=h2.size) + points = [] + for i in range(h2.size): + for n in range(4): + for p in range(1, 2*n+2): + points.append((h2[i], k2[i], n, p, s[i])) + points = np.array(points) + assert_func_equal(ellip_harm, ellip_harm_known, points, rtol=1e-12) + + +def test_ellip_harm_invalid_p(): + # Regression test. This should return nan. + n = 4 + # Make p > 2*n + 1. + p = 2*n + 2 + result = ellip_harm(0.5, 2.0, n, p, 0.2) + assert np.isnan(result) diff --git a/v12x/pylib/scipy/special/tests/test_erfinv.py b/v12x/pylib/scipy/special/tests/test_erfinv.py new file mode 100644 index 0000000000000000000000000000000000000000..98739b93fc6ad75a41a7b80107ee696453b12a09 --- /dev/null +++ b/v12x/pylib/scipy/special/tests/test_erfinv.py @@ -0,0 +1,89 @@ +import numpy as np +from numpy.testing import assert_allclose, assert_equal +import pytest + +import scipy.special as sc + + +class TestInverseErrorFunction: + def test_compliment(self): + # Test erfcinv(1 - x) == erfinv(x) + x = np.linspace(-1, 1, 101) + assert_allclose(sc.erfcinv(1 - x), sc.erfinv(x), rtol=0, atol=1e-15) + + def test_literal_values(self): + # The expected values were calculated with mpmath: + # + # import mpmath + # mpmath.mp.dps = 200 + # for y in [0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]: + # x = mpmath.erfinv(y) + # print(x) + # + y = np.array([0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]) + actual = sc.erfinv(y) + expected = [ + 0.0, + 0.08885599049425769, + 0.1791434546212917, + 0.2724627147267543, + 0.37080715859355795, + 0.4769362762044699, + 0.5951160814499948, + 0.7328690779592167, + 0.9061938024368233, + 1.1630871536766743, + ] + assert_allclose(actual, expected, rtol=0, atol=1e-15) + + @pytest.mark.parametrize( + 'f, x, y', + [ + (sc.erfinv, -1, -np.inf), + (sc.erfinv, 0, 0), + (sc.erfinv, 1, np.inf), + (sc.erfinv, -100, np.nan), + (sc.erfinv, 100, np.nan), + (sc.erfcinv, 0, np.inf), + (sc.erfcinv, 1, -0.0), + (sc.erfcinv, 2, -np.inf), + (sc.erfcinv, -100, np.nan), + (sc.erfcinv, 100, np.nan), + ], + ids=[ + 'erfinv at lower bound', + 'erfinv at midpoint', + 'erfinv at upper bound', + 'erfinv below lower bound', + 'erfinv above upper bound', + 'erfcinv at lower bound', + 'erfcinv at midpoint', + 'erfcinv at upper bound', + 'erfcinv below lower bound', + 'erfcinv above upper bound', + ] + ) + def test_domain_bounds(self, f, x, y): + assert_equal(f(x), y) + + def test_erfinv_asympt(self): + # regression test for gh-12758: erfinv(x) loses precision at small x + # expected values precomputed with mpmath: + # >>> mpmath.mp.dps = 100 + # >>> expected = [float(mpmath.erfinv(t)) for t in x] + x = np.array([1e-20, 1e-15, 1e-14, 1e-10, 1e-8, 0.9e-7, 1.1e-7, 1e-6]) + expected = np.array([8.86226925452758e-21, + 8.862269254527581e-16, + 8.86226925452758e-15, + 8.862269254527581e-11, + 8.86226925452758e-09, + 7.97604232907484e-08, + 9.74849617998037e-08, + 8.8622692545299e-07]) + assert_allclose(sc.erfinv(x), expected, + rtol=1e-15) + + # also test the roundtrip consistency + assert_allclose(sc.erf(sc.erfinv(x)), + x, + rtol=5e-15) diff --git a/v12x/pylib/scipy/special/tests/test_exponential_integrals.py b/v12x/pylib/scipy/special/tests/test_exponential_integrals.py new file mode 100644 index 0000000000000000000000000000000000000000..d6081580b338fbf1c1eb39a846047ee0ea684e66 --- /dev/null +++ b/v12x/pylib/scipy/special/tests/test_exponential_integrals.py @@ -0,0 +1,123 @@ +import pytest + +import numpy as np +from numpy.testing import assert_allclose +import scipy.special as sc + + +class TestExp1: + + def test_branch_cut(self): + assert np.isnan(sc.exp1(-1)) + assert sc.exp1(complex(-1, 0)).imag == ( + -sc.exp1(complex(-1, -0.0)).imag + ) + + assert_allclose( + sc.exp1(complex(-1, 0)), + sc.exp1(-1 + 1e-20j), + atol=0, + rtol=1e-15 + ) + assert_allclose( + sc.exp1(complex(-1, -0.0)), + sc.exp1(-1 - 1e-20j), + atol=0, + rtol=1e-15 + ) + + def test_834(self): + # Regression test for #834 + a = sc.exp1(-complex(19.9999990)) + b = sc.exp1(-complex(19.9999991)) + assert_allclose(a.imag, b.imag, atol=0, rtol=1e-15) + + +class TestScaledExp1: + + @pytest.mark.parametrize('x, expected', [(0, 0), (np.inf, 1)]) + def test_limits(self, x, expected): + y = sc._ufuncs._scaled_exp1(x) + assert y == expected + + # The expected values were computed with mpmath, e.g.: + # + # from mpmath import mp + # mp.dps = 80 + # x = 1e-25 + # print(float(x*mp.exp(x)*np.expint(1, x))) + # + # prints 5.698741165994961e-24 + # + # The method used to compute _scaled_exp1 changes at x=1 + # and x=1250, so values at those inputs, and values just + # above and below them, are included in the test data. + @pytest.mark.parametrize('x, expected', + [(1e-25, 5.698741165994961e-24), + (0.1, 0.20146425447084518), + (0.9995, 0.5962509885831002), + (1.0, 0.5963473623231941), + (1.0005, 0.5964436833238044), + (2.5, 0.7588145912149602), + (10.0, 0.9156333393978808), + (100.0, 0.9901942286733019), + (500.0, 0.9980079523802055), + (1000.0, 0.9990019940238807), + (1249.5, 0.9992009578306811), + (1250.0, 0.9992012769377913), + (1250.25, 0.9992014363957858), + (2000.0, 0.9995004992514963), + (1e4, 0.9999000199940024), + (1e10, 0.9999999999), + (1e15, 0.999999999999999), + ]) + def test_scaled_exp1(self, x, expected): + y = sc._ufuncs._scaled_exp1(x) + assert_allclose(y, expected, rtol=2e-15) + + +class TestExpi: + + @pytest.mark.parametrize('result', [ + sc.expi(complex(-1, 0)), + sc.expi(complex(-1, -0.0)), + sc.expi(-1) + ]) + def test_branch_cut(self, result): + desired = -0.21938393439552027368 # Computed using Mpmath + assert_allclose(result, desired, atol=0, rtol=1e-14) + + def test_near_branch_cut(self): + lim_from_above = sc.expi(-1 + 1e-20j) + lim_from_below = sc.expi(-1 - 1e-20j) + assert_allclose( + lim_from_above.real, + lim_from_below.real, + atol=0, + rtol=1e-15 + ) + assert_allclose( + lim_from_above.imag, + -lim_from_below.imag, + atol=0, + rtol=1e-15 + ) + + def test_continuity_on_positive_real_axis(self): + assert_allclose( + sc.expi(complex(1, 0)), + sc.expi(complex(1, -0.0)), + atol=0, + rtol=1e-15 + ) + + @pytest.mark.parametrize('x, expected', [(0, -np.inf), (np.inf, np.inf)]) + def test_limits(self, x, expected): + y = sc.expi(x) + assert y == expected + + +class TestExpn: + + def test_out_of_domain(self): + assert all(np.isnan([sc.expn(-1, 1.0), sc.expn(1, -1.0)])) diff --git a/v12x/pylib/scipy/special/tests/test_extending.py b/v12x/pylib/scipy/special/tests/test_extending.py new file mode 100644 index 0000000000000000000000000000000000000000..3ecaf51545e006e37a451a08f356ce0392a3159c --- /dev/null +++ b/v12x/pylib/scipy/special/tests/test_extending.py @@ -0,0 +1,28 @@ +import os +import platform +import sysconfig + +import pytest + +from scipy._lib._testutils import IS_EDITABLE,_test_cython_extension, cython +from scipy.special import beta, gamma + + +@pytest.mark.fail_slow(40) +# essential per https://github.com/scipy/scipy/pull/20487#discussion_r1567057247 +@pytest.mark.skipif(IS_EDITABLE, + reason='Editable install cannot find .pxd headers.') +@pytest.mark.skipif((platform.system() == 'Windows' and + sysconfig.get_config_var('Py_GIL_DISABLED')), + reason='gh-22039') +@pytest.mark.skipif(platform.machine() in ["wasm32", "wasm64"], + reason="Can't start subprocess") +@pytest.mark.skipif(cython is None, reason="requires cython") +def test_cython(tmp_path): + srcdir = os.path.dirname(os.path.dirname(__file__)) + extensions, extensions_cpp = _test_cython_extension(tmp_path, srcdir) + # actually test the cython c-extensions + assert extensions.cy_beta(0.5, 0.1) == beta(0.5, 0.1) + assert extensions.cy_gamma(0.5 + 1.0j) == gamma(0.5 + 1.0j) + assert extensions_cpp.cy_beta(0.5, 0.1) == beta(0.5, 0.1) + assert extensions_cpp.cy_gamma(0.5 + 1.0j) == gamma(0.5 + 1.0j) diff --git a/v12x/pylib/scipy/special/tests/test_faddeeva.py b/v12x/pylib/scipy/special/tests/test_faddeeva.py new file mode 100644 index 0000000000000000000000000000000000000000..8868f66c47ce0d4bbb21c78435a6c89d44065252 --- /dev/null +++ b/v12x/pylib/scipy/special/tests/test_faddeeva.py @@ -0,0 +1,85 @@ +import pytest + +import numpy as np +from numpy.testing import assert_allclose +import scipy.special as sc +from scipy.special._testutils import FuncData + + +class TestVoigtProfile: + + @pytest.mark.parametrize('x, sigma, gamma', [ + (np.nan, 1, 1), + (0, np.nan, 1), + (0, 1, np.nan), + (1, np.nan, 0), + (np.nan, 1, 0), + (1, 0, np.nan), + (np.nan, 0, 1), + (np.nan, 0, 0) + ]) + def test_nan(self, x, sigma, gamma): + assert np.isnan(sc.voigt_profile(x, sigma, gamma)) + + @pytest.mark.parametrize('x, desired', [ + (-np.inf, 0), + (np.inf, 0) + ]) + def test_inf(self, x, desired): + assert sc.voigt_profile(x, 1, 1) == desired + + def test_against_mathematica(self): + # Results obtained from Mathematica by computing + # + # PDF[VoigtDistribution[gamma, sigma], x] + # + points = np.array([ + [-7.89, 45.06, 6.66, 0.0077921073660388806401], + [-0.05, 7.98, 24.13, 0.012068223646769913478], + [-13.98, 16.83, 42.37, 0.0062442236362132357833], + [-12.66, 0.21, 6.32, 0.010052516161087379402], + [11.34, 4.25, 21.96, 0.0113698923627278917805], + [-11.56, 20.40, 30.53, 0.0076332760432097464987], + [-9.17, 25.61, 8.32, 0.011646345779083005429], + [16.59, 18.05, 2.50, 0.013637768837526809181], + [9.11, 2.12, 39.33, 0.0076644040807277677585], + [-43.33, 0.30, 45.68, 0.0036680463875330150996] + ]) + FuncData( + sc.voigt_profile, + points, + (0, 1, 2), + 3, + atol=0, + rtol=1e-15 + ).check() + + def test_symmetry(self): + x = np.linspace(0, 10, 20) + assert_allclose( + sc.voigt_profile(x, 1, 1), + sc.voigt_profile(-x, 1, 1), + rtol=1e-15, + atol=0 + ) + + @pytest.mark.parametrize('x, sigma, gamma, desired', [ + (0, 0, 0, np.inf), + (1, 0, 0, 0) + ]) + def test_corner_cases(self, x, sigma, gamma, desired): + assert sc.voigt_profile(x, sigma, gamma) == desired + + @pytest.mark.parametrize('sigma1, gamma1, sigma2, gamma2', [ + (0, 1, 1e-16, 1), + (1, 0, 1, 1e-16), + (0, 0, 1e-16, 1e-16) + ]) + def test_continuity(self, sigma1, gamma1, sigma2, gamma2): + x = np.linspace(1, 10, 20) + assert_allclose( + sc.voigt_profile(x, sigma1, gamma1), + sc.voigt_profile(x, sigma2, gamma2), + rtol=1e-16, + atol=1e-16 + ) diff --git a/v12x/pylib/scipy/special/tests/test_gamma.py b/v12x/pylib/scipy/special/tests/test_gamma.py new file mode 100644 index 0000000000000000000000000000000000000000..2e3fbd17dddeed73d311566a930f52899e3b9db6 --- /dev/null +++ b/v12x/pylib/scipy/special/tests/test_gamma.py @@ -0,0 +1,12 @@ +import numpy as np +import scipy.special as sc + + +class TestRgamma: + + def test_gh_11315(self): + assert sc.rgamma(-35) == 0 + + def test_rgamma_zeros(self): + x = np.array([0, -10, -100, -1000, -10000]) + assert np.all(sc.rgamma(x) == 0) diff --git a/v12x/pylib/scipy/special/tests/test_gammainc.py b/v12x/pylib/scipy/special/tests/test_gammainc.py new file mode 100644 index 0000000000000000000000000000000000000000..b5a0ffb65d6deebf77fe82ebee2dca61f2067c59 --- /dev/null +++ b/v12x/pylib/scipy/special/tests/test_gammainc.py @@ -0,0 +1,152 @@ +import pytest + +import numpy as np +from numpy.testing import assert_allclose, assert_array_equal + +import scipy.special as sc +from scipy.special._testutils import FuncData + + +INVALID_POINTS = [ + (1, -1), + (0, 0), + (-1, 1), + (np.nan, 1), + (1, np.nan) +] + + +class TestGammainc: + + @pytest.mark.parametrize('a, x', INVALID_POINTS) + def test_domain(self, a, x): + assert np.isnan(sc.gammainc(a, x)) + + def test_a_eq_0_x_gt_0(self): + assert sc.gammainc(0, 1) == 1 + + @pytest.mark.parametrize('a, x, desired', [ + (np.inf, 1, 0), + (np.inf, 0, 0), + (np.inf, np.inf, np.nan), + (1, np.inf, 1) + ]) + def test_infinite_arguments(self, a, x, desired): + result = sc.gammainc(a, x) + if np.isnan(desired): + assert np.isnan(result) + else: + assert result == desired + + @pytest.mark.parametrize("x", [-np.inf, -1.0, -0.0, 0.0, np.inf, np.nan]) + def test_a_nan(self, x): + assert np.isnan(sc.gammainc(np.nan, x)) + + @pytest.mark.parametrize("a", [-np.inf, -1.0, -0.0, 0.0, np.inf, np.nan]) + def test_x_nan(self, a): + assert np.isnan(sc.gammainc(a, np.nan)) + + def test_infinite_limits(self): + # Test that large arguments converge to the hard-coded limits + # at infinity. + assert_allclose( + sc.gammainc(1000, 100), + sc.gammainc(np.inf, 100), + atol=1e-200, # Use `atol` since the function converges to 0. + rtol=0 + ) + assert sc.gammainc(100, 1000) == sc.gammainc(100, np.inf) + + def test_x_zero(self): + a = np.arange(1, 10) + assert_array_equal(sc.gammainc(a, 0), 0) + + def test_limit_check(self): + result = sc.gammainc(1e-10, 1) + limit = sc.gammainc(0, 1) + assert np.isclose(result, limit) + + def gammainc_line(self, x): + # The line a = x where a simpler asymptotic expansion (analog + # of DLMF 8.12.15) is available. + c = np.array([-1/3, -1/540, 25/6048, 101/155520, + -3184811/3695155200, -2745493/8151736420]) + res = 0 + xfac = 1 + for ck in c: + res -= ck*xfac + xfac /= x + res /= np.sqrt(2*np.pi*x) + res += 0.5 + return res + + def test_line(self): + x = np.logspace(np.log10(25), 300, 500) + a = x + dataset = np.vstack((a, x, self.gammainc_line(x))).T + FuncData(sc.gammainc, dataset, (0, 1), 2, rtol=1e-11).check() + + def test_roundtrip(self): + a = np.logspace(-5, 10, 100) + x = np.logspace(-5, 10, 100) + + y = sc.gammaincinv(a, sc.gammainc(a, x)) + assert_allclose(x, y, rtol=1e-10) + + +class TestGammaincc: + + @pytest.mark.parametrize('a, x', INVALID_POINTS) + def test_domain(self, a, x): + assert np.isnan(sc.gammaincc(a, x)) + + def test_a_eq_0_x_gt_0(self): + assert sc.gammaincc(0, 1) == 0 + + @pytest.mark.parametrize('a, x, desired', [ + (np.inf, 1, 1), + (np.inf, 0, 1), + (np.inf, np.inf, np.nan), + (1, np.inf, 0) + ]) + def test_infinite_arguments(self, a, x, desired): + result = sc.gammaincc(a, x) + if np.isnan(desired): + assert np.isnan(result) + else: + assert result == desired + + @pytest.mark.parametrize("x", [-np.inf, -1.0, -0.0, 0.0, np.inf, np.nan]) + def test_a_nan(self, x): + assert np.isnan(sc.gammaincc(np.nan, x)) + + @pytest.mark.parametrize("a", [-np.inf, -1.0, -0.0, 0.0, np.inf, np.nan]) + def test_x_nan(self, a): + assert np.isnan(sc.gammaincc(a, np.nan)) + + def test_infinite_limits(self): + # Test that large arguments converge to the hard-coded limits + # at infinity. + assert sc.gammaincc(1000, 100) == sc.gammaincc(np.inf, 100) + assert_allclose( + sc.gammaincc(100, 1000), + sc.gammaincc(100, np.inf), + atol=1e-200, # Use `atol` since the function converges to 0. + rtol=0 + ) + + def test_limit_check(self): + result = sc.gammaincc(1e-10,1) + limit = sc.gammaincc(0,1) + assert np.isclose(result, limit) + + def test_x_zero(self): + a = np.arange(1, 10) + assert_array_equal(sc.gammaincc(a, 0), 1) + + def test_roundtrip(self): + a = np.logspace(-5, 10, 100) + x = np.logspace(-5, 10, 100) + + y = sc.gammainccinv(a, sc.gammaincc(a, x)) + assert_allclose(x, y, rtol=1e-14) diff --git a/v12x/pylib/scipy/special/tests/test_gen_harmonic.py b/v12x/pylib/scipy/special/tests/test_gen_harmonic.py new file mode 100644 index 0000000000000000000000000000000000000000..0896142d571ae1d3815ab2eb99a775d52b7d9169 --- /dev/null +++ b/v12x/pylib/scipy/special/tests/test_gen_harmonic.py @@ -0,0 +1,99 @@ +import numpy as np +from numpy.testing import assert_allclose, assert_equal +import pytest +from scipy.special._ufuncs import _gen_harmonic, _normalized_gen_harmonic + + +# +# In the following tests, reference values were computed with mpmath. +# + +@pytest.mark.parametrize('typ', [np.int32, np.int64, np.float64]) +@pytest.mark.parametrize( + 'n, a, ref', + [(8, 9.0, 1.0020083884212339), + (1000, 2.5, 1.3414661912046497), + (10, 1.5, 1.9953364933456017), + (10000, 1.25, 4.1951168257387765), + (10000,1.00001, 9.787182620770265), + (80, 1.000002, 4.965460167788836), + (75, 1 + 1e-12, 4.901355630543771), + (100, 1 + 1e-14, 5.187377517639515), + (100, 1 + 8e-16, 5.187377517639611), + (100, 1.0, 5.187377517639621), + (7, 1.0, 2.592857142857143), + (8000, 1.0, 9.564474984261423), + (5, 1 - 1e-12, 2.2833333333347143), + (25000, 1 - 1e-12, 10.703866768669737), + (1000, 0.995, 7.6058022857089975), + (1000, 0.75, 19.055178975831392), + (10000, 0.25, 1332.5700547197382), + (5, 1e-8, 4.999999952125083), + (15, 1e-16, 14.999999999999996), + (100, 0.0, 100.0), + (4, -1.0, 10.0), + (75, -1.5, 19811.38815892374)] +) +def test_gen_harmonic(typ, n, a, ref): + h = _gen_harmonic(typ(n), a) + assert_allclose(h, ref, rtol=5e-15) + + +@pytest.mark.parametrize('typ', [np.int32, np.int64, np.float64]) +@pytest.mark.parametrize( + 'n, a, ref', + [(10, np.inf, 1.0), + (1, np.nan, 1.0), + (1, -np.inf, 1.0), + (3, np.nan, np.nan), + (-3, 1.0, np.nan)] +) +def test_gen_harmonic_exact_cases(typ, n, a, ref): + h = _gen_harmonic(typ(n), a) + assert_equal(h, ref) + + +def test_gen_harmonic_n_nan(): + h = _gen_harmonic(np.nan, 0.75) + assert_equal(h, np.nan) + + +@pytest.mark.parametrize('typ', [np.int32, np.int64, np.float64]) +@pytest.mark.parametrize( + 'j, k, n, a, ref', + [(400, 5000, 5000, 10.0, 4.2821759663214485e-25), + (400, 5000, 5000, 3.5, 1.11086549102426e-07), + (1, 2, 3, 1.5, 0.8755176866163012), + (300, 500, 500, 1 + 1e-14, 0.07559343891632035), + (1500, 2500, 3000, 1 - 1e-12, 0.05957291246371843), + (10, 12, 16, 0.5, 0.13601665344521513), + (16, 16, 20, 0.125, 0.04583107002260924), + (10, 12, 16, -0.5, 0.22359306724308234), + (1, 8000, 10000, -1.5, 0.5724512895513029)] +) +def test_normalized_gen_harmonic(typ, j, k, n, a, ref): + h = _normalized_gen_harmonic(typ(j), typ(k), typ(n), a) + assert_allclose(h, ref, 5e-15) + + +@pytest.mark.parametrize('typ', [np.int32, np.int64, np.float64]) +@pytest.mark.parametrize( + 'j, k, n, a, ref', + [(1, 1, 1, 0.5, 1.0), + (1, 1, 1, np.nan, 1.0), + (1, 2, 5, np.nan, np.nan), + (1, 2, 1, 1.25, np.nan), + (1, 2, 3, np.inf, 1.0), + (2, 3, 4, np.inf, 0.0), + (1, 1, 10, -np.inf, 0.0), + (2, 3, 4, -np.inf, np.nan), + (3, 6, 8, 0.0, 0.5)] +) +def test_normalized_gen_harmonic_exact_cases(typ, j, k, n, a, ref): + h = _normalized_gen_harmonic(typ(j), typ(k), typ(n), a) + assert_equal(h, ref) + + +def test_normalized_gen_harmonic_input_nan(): + h = _normalized_gen_harmonic(1.0, np.nan, 10.0, 1.05) + assert_equal(h, np.nan) diff --git a/v12x/pylib/scipy/special/tests/test_hyp2f1.py b/v12x/pylib/scipy/special/tests/test_hyp2f1.py new file mode 100644 index 0000000000000000000000000000000000000000..8e51a48f91440583ddaa88ffe09d48e4376055cf --- /dev/null +++ b/v12x/pylib/scipy/special/tests/test_hyp2f1.py @@ -0,0 +1,2566 @@ +"""Tests for hyp2f1 for complex values. + +Author: Albert Steppi, with credit to Adam Kullberg (FormerPhycisist) for +the implementation of mp_hyp2f1 below, which modifies mpmath's hyp2f1 to +return the same branch as scipy's on the standard branch cut. +""" + +import sys +import pytest +import numpy as np +from typing import NamedTuple +from numpy.testing import assert_allclose + +from scipy.special import hyp2f1 +from scipy.special._testutils import check_version, MissingModule + + +try: + import mpmath +except ImportError: + mpmath = MissingModule("mpmath") + + +def mp_hyp2f1(a, b, c, z): + """Return mpmath hyp2f1 calculated on same branch as scipy hyp2f1. + + For most values of a,b,c mpmath returns the x - 0j branch of hyp2f1 on the + branch cut x=(1,inf) whereas scipy's hyp2f1 calculates the x + 0j branch. + Thus, to generate the right comparison values on the branch cut, we + evaluate mpmath.hyp2f1 at x + 1e-15*j. + + The exception to this occurs when c-a=-m in which case both mpmath and + scipy calculate the x + 0j branch on the branch cut. When this happens + mpmath.hyp2f1 will be evaluated at the original z point. + """ + on_branch_cut = z.real > 1.0 and abs(z.imag) < 1.0e-15 + cond1 = abs(c - a - round(c - a)) < 1.0e-15 and round(c - a) <= 0 + cond2 = abs(c - b - round(c - b)) < 1.0e-15 and round(c - b) <= 0 + # Make sure imaginary part is *exactly* zero + if on_branch_cut: + z = z.real + 0.0j + if on_branch_cut and not (cond1 or cond2): + z_mpmath = z.real + 1.0e-15j + else: + z_mpmath = z + return complex(mpmath.hyp2f1(a, b, c, z_mpmath)) + + +class Hyp2f1TestCase(NamedTuple): + a: float + b: float + c: float + z: complex + expected: complex + rtol: float + + +class TestHyp2f1: + """Tests for hyp2f1 for complex values. + + Expected values for test cases were computed using mpmath. See + `scipy.special._precompute.hyp2f1_data`. The verbose style of specifying + test cases is used for readability and to make it easier to mark individual + cases as expected to fail. Expected failures are used to highlight cases + where improvements are needed. See + `scipy.special._precompute.hyp2f1_data.make_hyp2f1_test_cases` for a + function to generate the boilerplate for the test cases. + + Assertions have been added to each test to ensure that the test cases match + the situations that are intended. A final test `test_test_hyp2f1` checks + that the expected values in the test cases actually match what is computed + by mpmath. This test is marked slow even though it isn't particularly slow + so that it won't run by default on continuous integration builds. + """ + @pytest.mark.parametrize( + "hyp2f1_test_case", + [ + pytest.param( + Hyp2f1TestCase( + a=0.5, + b=0.2, + c=-10, + z=0.2 + 0.2j, + expected=np.inf + 0j, + rtol=0 + ) + ), + pytest.param( + Hyp2f1TestCase( + a=0.5, + b=0.2, + c=-10, + z=0 + 0j, + expected=1 + 0j, + rtol=0 + ), + ), + pytest.param( + Hyp2f1TestCase( + a=0.5, + b=0, + c=-10, + z=0.2 + 0.2j, + expected=1 + 0j, + rtol=0 + ), + ), + pytest.param( + Hyp2f1TestCase( + a=0.5, + b=0, + c=0, + z=0.2 + 0.2j, + expected=1 + 0j, + rtol=0, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=0.5, + b=0.2, + c=0, + z=0.2 + 0.2j, + expected=np.inf + 0j, + rtol=0, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=0.5, + b=0.2, + c=0, + z=0 + 0j, + expected=np.nan + 0j, + rtol=0, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=0.5, + b=-5, + c=-10, + z=0.2 + 0.2j, + expected=(1.0495404166666666+0.05708208333333334j), + rtol=1e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=0.5, + b=-10, + c=-10, + z=0.2 + 0.2j, + expected=(1.092966013125+0.13455014673750001j), + rtol=1e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-10, + b=-20, + c=-10, + z=0.2 + 0.2j, + expected=(-0.07712512000000005+0.12752814080000005j), + rtol=1e-13, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-1, + b=3.2, + c=-1, + z=0.2 + 0.2j, + expected=(1.6400000000000001+0.6400000000000001j), + rtol=1e-13, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-2, + b=1.2, + c=-4, + z=1 + 0j, + expected=1.8200000000000001 + 0j, + rtol=1e-15, + ), + ), + ] + ) + def test_c_non_positive_int(self, hyp2f1_test_case): + a, b, c, z, expected, rtol = hyp2f1_test_case + assert_allclose(hyp2f1(a, b, c, z), expected, rtol=rtol) + + @pytest.mark.parametrize( + "hyp2f1_test_case", + [ + pytest.param( + Hyp2f1TestCase( + a=0.5, + b=0.2, + c=1.5, + z=1 + 0j, + expected=1.1496439092239847 + 0j, + rtol=1e-15 + ), + ), + pytest.param( + Hyp2f1TestCase( + a=12.3, + b=8.0, + c=20.31, + z=1 + 0j, + expected=69280986.75273195 + 0j, + rtol=1e-15 + ), + ), + pytest.param( + Hyp2f1TestCase( + a=290.2, + b=321.5, + c=700.1, + z=1 + 0j, + expected=1.3396562400934e117 + 0j, + rtol=1e-12, + ), + ), + # Note that here even mpmath produces different results for + # results that should be equivalent. + pytest.param( + Hyp2f1TestCase( + a=9.2, + b=621.5, + c=700.1, + z=(1+0j), + expected=(952726652.4158565+0j), + rtol=5e-13, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=621.5, + b=9.2, + c=700.1, + z=(1+0j), + expected=(952726652.4160284+0j), + rtol=5e-12, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-101.2, + b=-400.4, + c=-172.1, + z=(1+0j), + expected=(2.2253618341394838e+37+0j), + rtol=1e-13, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-400.4, + b=-101.2, + c=-172.1, + z=(1+0j), + expected=(2.2253618341394838e+37+0j), + rtol=5e-13, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=172.5, + b=-201.3, + c=151.2, + z=(1+0j), + expected=(7.072266653650905e-135+0j), + rtol=5e-13, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-201.3, + b=172.5, + c=151.2, + z=(1+0j), + expected=(7.072266653650905e-135+0j), + rtol=5e-13, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-102.1, + b=-20.3, + c=1.3, + z=1 + 0j, + expected=2.7899070752746906e22 + 0j, + rtol=3e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-202.6, + b=60.3, + c=1.5, + z=1 + 0j, + expected=-1.3113641413099326e-56 + 0j, + rtol=1e-12, + ), + ), + ], + ) + def test_unital_argument(self, hyp2f1_test_case): + """Tests for case z = 1, c - a - b > 0. + + Expected answers computed using mpmath. + """ + a, b, c, z, expected, rtol = hyp2f1_test_case + assert z == 1 and c - a - b > 0 # Tests the test + assert_allclose(hyp2f1(a, b, c, z), expected, rtol=rtol) + + @pytest.mark.parametrize( + "hyp2f1_test_case", + [ + pytest.param( + Hyp2f1TestCase( + a=0.5, + b=0.2, + c=1.3, + z=-1 + 0j, + expected=0.9428846409614143 + 0j, + rtol=1e-15), + ), + pytest.param( + Hyp2f1TestCase( + a=12.3, + b=8.0, + c=5.300000000000001, + z=-1 + 0j, + expected=-4.845809986595704e-06 + 0j, + rtol=1e-15 + ), + ), + pytest.param( + Hyp2f1TestCase( + a=221.5, + b=90.2, + c=132.3, + z=-1 + 0j, + expected=2.0490488728377282e-42 + 0j, + rtol=1e-7, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-102.1, + b=-20.3, + c=-80.8, + z=-1 + 0j, + expected=45143784.46783885 + 0j, + rtol=1e-7, + ), + marks=pytest.mark.xfail( + condition=sys.maxsize < 2**32, + reason="Fails on 32 bit.", + ) + ), + ], + ) + def test_special_case_z_near_minus_1(self, hyp2f1_test_case): + """Tests for case z ~ -1, c ~ 1 + a - b + + Expected answers computed using mpmath. + """ + a, b, c, z, expected, rtol = hyp2f1_test_case + assert abs(1 + a - b - c) < 1e-15 and abs(z + 1) < 1e-15 + assert_allclose(hyp2f1(a, b, c, z), expected, rtol=rtol) + + @pytest.mark.parametrize( + "hyp2f1_test_case", + [ + pytest.param( + Hyp2f1TestCase( + a=-4, + b=2.02764642551431, + c=1.0561196186065624, + z=(0.9473684210526314-0.10526315789473695j), + expected=(0.0031961077109535375-0.0011313924606557173j), + rtol=1e-12, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-8, + b=-7.937789122896016, + c=-15.964218273004214, + z=(2-0.10526315789473695j), + expected=(0.005543763196412503-0.0025948879065698306j), + rtol=5e-13, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-8, + b=8.095813935368371, + c=4.0013768449590685, + z=(0.9473684210526314-0.10526315789473695j), + expected=(-0.0003054674127221263-9.261359291755414e-05j), + rtol=1e-10, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-4, + b=-3.956227226099288, + c=-3.9316537064827854, + z=(1.1578947368421053-0.3157894736842106j), + expected=(-0.0020809502580892937-0.0041877333232365095j), + rtol=5e-12, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=2.02764642551431, + b=-4, + c=2.050308316530781, + z=(0.9473684210526314-0.10526315789473695j), + expected=(0.0011282435590058734+0.0002027062303465851j), + rtol=5e-13, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-7.937789122896016, + b=-8, + c=-15.964218273004214, + z=(1.3684210526315788+0.10526315789473673j), + expected=(-9.134907719238265e-05-0.00040219233987390723j), + rtol=5e-12, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=4.080187217753502, + b=-4, + c=4.0013768449590685, + z=(0.9473684210526314-0.10526315789473695j), + expected=(-0.000519013062087489-0.0005855883076830948j), + rtol=5e-12, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-10000, + b=2.2, + c=93459345.3, + z=(2+2j), + expected=(0.9995292071559088-0.00047047067522659253j), + rtol=1e-12, + ), + ), + ] + ) + def test_a_b_negative_int(self, hyp2f1_test_case): + a, b, c, z, expected, rtol = hyp2f1_test_case + assert a == int(a) and a < 0 or b == int(b) and b < 0 # Tests the test + assert_allclose(hyp2f1(a, b, c, z), expected, rtol=rtol) + + @pytest.mark.parametrize( + "hyp2f1_test_case", + [ + pytest.param( + Hyp2f1TestCase( + a=-0.5, + b=-0.9629749245209605, + c=-15.5, + z=(1.1578947368421053-1.1578947368421053j), + expected=(0.9778506962676361+0.044083801141231616j), + rtol=3e-12, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=8.5, + b=-3.9316537064827854, + c=1.5, + z=(0.9473684210526314-0.10526315789473695j), + expected=(4.0793167523167675-10.11694246310966j), + rtol=6e-12, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=8.5, + b=-0.9629749245209605, + c=2.5, + z=(1.1578947368421053-0.10526315789473695j), + expected=(-2.9692999501916915+0.6394599899845594j), + rtol=1e-11, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-0.5, + b=-0.9629749245209605, + c=-15.5, + z=(1.5789473684210522-1.1578947368421053j), + expected=(0.9493076367106102-0.04316852977183447j), + rtol=1e-11, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-0.9220024191881196, + b=-0.5, + c=-15.5, + z=(0.5263157894736841+0.10526315789473673j), + expected=(0.9844377175631795-0.003120587561483841j), + rtol=1e-10, + ), + ), + ], + ) + def test_a_b_neg_int_after_euler_hypergeometric_transformation( + self, hyp2f1_test_case + ): + a, b, c, z, expected, rtol = hyp2f1_test_case + assert ( # Tests the test + (abs(c - a - int(c - a)) < 1e-15 and c - a < 0) or + (abs(c - b - int(c - b)) < 1e-15 and c - b < 0) + ) + assert_allclose(hyp2f1(a, b, c, z), expected, rtol=rtol) + + @pytest.mark.parametrize( + "hyp2f1_test_case", + [ + pytest.param( + Hyp2f1TestCase( + a=-0.9220024191881196, + b=-0.9629749245209605, + c=-15.963511401609862, + z=(0.10526315789473673-0.3157894736842106j), + expected=(0.9941449585778349+0.01756335047931358j), + rtol=1e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=1.0272592605282642, + b=-0.9629749245209605, + c=-15.963511401609862, + z=(0.5263157894736841+0.5263157894736841j), + expected=(1.0388722293372104-0.09549450380041416j), + rtol=5e-11, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=2.02764642551431, + b=1.0561196186065624, + c=-7.93846038215665, + z=(0.10526315789473673+0.7368421052631575j), + expected=(2.1948378809826434+24.934157235172222j), + rtol=5e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=2.02764642551431, + b=16.088264119063613, + c=8.031683612216888, + z=(0.3157894736842106-0.736842105263158j), + expected=(-0.4075277891264672-0.06819344579666956j), + rtol=2e-12, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=4.080187217753502, + b=2.050308316530781, + c=8.031683612216888, + z=(0.7368421052631575-0.10526315789473695j), + expected=(2.833535530740603-0.6925373701408158j), + rtol=5e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=2.02764642551431, + b=2.050308316530781, + c=4.078873014294075, + z=(0.10526315789473673-0.3157894736842106j), + expected=(1.005347176329683-0.3580736009337313j), + rtol=5e-16, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-0.9220024191881196, + b=-0.9629749245209605, + c=-15.963511401609862, + z=(0.3157894736842106-0.5263157894736843j), + expected=(0.9824353641135369+0.029271018868990268j), + rtol=5e-13, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-0.9220024191881196, + b=-0.9629749245209605, + c=-159.63511401609862, + z=(0.3157894736842106-0.5263157894736843j), + expected=(0.9982436200365834+0.002927268199671111j), + rtol=1e-7, + ), + marks=pytest.mark.xfail(reason="Poor convergence.") + ), + pytest.param( + Hyp2f1TestCase( + a=2.02764642551431, + b=16.088264119063613, + c=8.031683612216888, + z=(0.5263157894736841-0.5263157894736843j), + expected=(-0.6906825165778091+0.8176575137504892j), + rtol=5e-13, + ), + ), + ] + ) + def test_region1(self, hyp2f1_test_case): + """|z| < 0.9 and real(z) >= 0.""" + a, b, c, z, expected, rtol = hyp2f1_test_case + assert abs(z) < 0.9 and z.real >= 0 # Tests the test + assert_allclose(hyp2f1(a, b, c, z), expected, rtol=rtol) + + @pytest.mark.parametrize( + "hyp2f1_test_case", + [ + pytest.param( + Hyp2f1TestCase( + a=2.02764642551431, + b=1.0561196186065624, + c=4.078873014294075, + z=(-0.3157894736842106+0.7368421052631575j), + expected=(0.7751915029081136+0.24068493258607315j), + rtol=5e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=16.087593263474208, + b=16.088264119063613, + c=2.0397202577726152, + z=(-0.9473684210526316-0.3157894736842106j), + expected=(6.564549348474962e-07+1.6761570598334562e-06j), + rtol=5e-09, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=1.0272592605282642, + b=2.050308316530781, + c=16.056809865262608, + z=(-0.10526315789473695-0.10526315789473695j), + expected=(0.9862043298997204-0.013293151372712681j), + rtol=5e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=4.080187217753502, + b=8.077282662161238, + c=16.056809865262608, + z=(-0.3157894736842106-0.736842105263158j), + expected=(0.16163826638754716-0.41378530376373734j), + rtol=5e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=2.02764642551431, + b=2.050308316530781, + c=-0.906685989801748, + z=(-0.5263157894736843+0.3157894736842106j), + expected=(-6.256871535165936+0.13824973858225484j), + rtol=1e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=2.02764642551431, + b=8.077282662161238, + c=-3.9924618758357022, + z=(-0.9473684210526316-0.3157894736842106j), + expected=(75.54672526086316+50.56157041797548j), + rtol=5e-12, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=16.087593263474208, + b=8.077282662161238, + c=-1.9631175993998025, + z=(-0.5263157894736843+0.5263157894736841j), + expected=(282.0602536306534-82.31597306936214j), + rtol=5e-13, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=8.095813935368371, + b=-3.9316537064827854, + c=8.031683612216888, + z=(-0.5263157894736843-0.10526315789473695j), + expected=(5.179603735575851+1.4445374002099813j), + rtol=5e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=4.080187217753502, + b=-7.949900487447654, + c=1.0651378143226575, + z=(-0.3157894736842106-0.9473684210526316j), + expected=(2317.623517606141-269.51476321010324j), + rtol=5e-13, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=16.087593263474208, + b=-1.92872979730171, + c=2.0397202577726152, + z=(-0.736842105263158-0.3157894736842106j), + expected=(29.179154096175836+22.126690357535043j), + rtol=5e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=8.095813935368371, + b=-3.9316537064827854, + c=-15.963511401609862, + z=(-0.736842105263158-0.10526315789473695j), + expected=(0.20820247892032057-0.04763956711248794j), + rtol=5e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=1.0272592605282642, + b=-15.964218273004214, + c=-1.9631175993998025, + z=(-0.3157894736842106-0.5263157894736843j), + expected=(-157471.63920142158+991294.0587828817j), + rtol=5e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=8.095813935368371, + b=-7.949900487447654, + c=-7.93846038215665, + z=(-0.10526315789473695-0.10526315789473695j), + expected=(0.30765349653210194-0.2979706363594157j), + rtol=1e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-3.956227226099288, + b=1.0561196186065624, + c=8.031683612216888, + z=(-0.9473684210526316-0.10526315789473695j), + expected=(1.6787607400597109+0.10056620134616838j), + rtol=5e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-7.937789122896016, + b=16.088264119063613, + c=4.078873014294075, + z=(-0.5263157894736843-0.736842105263158j), + expected=(7062.07842506049-12768.77955655703j), + rtol=5e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-7.937789122896016, + b=16.088264119063613, + c=2.0397202577726152, + z=(-0.3157894736842106+0.7368421052631575j), + expected=(54749.216391029935-23078.144720887536j), + rtol=2e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-3.956227226099288, + b=1.0561196186065624, + c=-0.906685989801748, + z=(-0.10526315789473695-0.10526315789473695j), + expected=(1.21521766411428-4.449385173946672j), + rtol=5e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-15.980848054962111, + b=4.0013768449590685, + c=-1.9631175993998025, + z=(-0.736842105263158+0.5263157894736841j), + expected=(19234693144.196907+1617913967.7294445j), + rtol=5e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-1.9214641416286231, + b=1.0561196186065624, + c=-15.963511401609862, + z=(-0.5263157894736843+0.3157894736842106j), + expected=(0.9345201094534371+0.03745712558992195j), + rtol=5e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-7.937789122896016, + b=-0.9629749245209605, + c=2.0397202577726152, + z=(-0.10526315789473695+0.10526315789473673j), + expected=(0.605732446296829+0.398171533680972j), + rtol=5e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-1.9214641416286231, + b=-15.964218273004214, + c=2.0397202577726152, + z=(-0.10526315789473695-0.5263157894736843j), + expected=(-9.753761888305416-4.590126012666959j), + rtol=5e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-3.956227226099288, + b=-1.92872979730171, + c=2.0397202577726152, + z=(-0.10526315789473695+0.3157894736842106j), + expected=(0.45587226291120714+1.0694545265819797j), + rtol=5e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-0.9220024191881196, + b=-7.949900487447654, + c=-0.906685989801748, + z=(-0.736842105263158+0.3157894736842106j), + expected=(12.334808243233418-76.26089051819054j), + rtol=5e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-0.9220024191881196, + b=-7.949900487447654, + c=-15.963511401609862, + z=(-0.5263157894736843+0.10526315789473673j), + expected=(1.2396019687632678-0.047507973161146286j), + rtol=1e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-15.980848054962111, + b=-0.9629749245209605, + c=-0.906685989801748, + z=(-0.3157894736842106-0.5263157894736843j), + expected=(97.7889554372208-18.999754543400016j), + rtol=5e-13, + ), + ), + ] + ) + def test_region2(self, hyp2f1_test_case): + """|z| < 1 and real(z) < 0.""" + a, b, c, z, expected, rtol = hyp2f1_test_case + assert abs(z) < 1 and z.real < 0 # Tests the test + assert_allclose(hyp2f1(a, b, c, z), expected, rtol=rtol) + + @pytest.mark.parametrize( + "hyp2f1_test_case", + [ + pytest.param( + Hyp2f1TestCase( + a=16.25, + b=-3.75, + c=-3.5, + z=(0.5263157894736841+0.7368421052631575j), + expected=(-1279.4894322256655-2302.914821389276j), + rtol=5e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-15.75, + b=8.25, + c=-1.5, + z=(0.9473684210526314+0.3157894736842106j), + expected=(-8889.452798586273-11961.162305065242j), + rtol=5e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=2.25, + b=2.25, + c=-1.5, + z=(0.5263157894736841-0.736842105263158j), + expected=(-236.58971357952055-238.5228224781136j), + rtol=5e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-0.75, + b=-7.75, + c=-15.5, + z=(0.5263157894736841+0.7368421052631575j), + expected=(0.8116076584352279-0.29360565398246036j), + rtol=5e-16, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=1.25, + b=4.25, + c=-0.5, + z=(0.5263157894736841-0.736842105263158j), + expected=(-28.119407485189985+98.89858821348005j), + rtol=5e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-0.75, + b=2.25, + c=1.5, + z=(0.5263157894736841+0.7368421052631575j), + expected=(0.5311049067450484-0.9434347326448517j), + rtol=5e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=8.25, + b=-15.75, + c=-7.5, + z=(0.9473684210526314+0.10526315789473673j), + expected=(1262084.378141873+1775569.6338380123j), + rtol=5e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-15.75, + b=-7.75, + c=-15.5, + z=(0.5263157894736841-0.736842105263158j), + expected=(-0.009810480794804165+0.3648997569257999j), + rtol=1e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=8.25, + b=2.25, + c=-3.5, + z=(0.5263157894736841-0.736842105263158j), + expected=(585660.8815535795-33646.68398590896j), + rtol=5e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=4.080187217753502, + b=-3.9316537064827854, + c=-15.963511401609862, + z=(0.9473684210526314-0.10526315789473695j), + expected=(181899621848365.2-173207123998705.7j), + rtol=5e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-15.75, + b=8.25, + c=-0.5, + z=(0.9473684210526314-0.10526315789473695j), + expected=(0.04271686244952705-0.14087902824639406j), + rtol=5e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-7.937789122896016, + b=-1.92872979730171, + c=-0.906685989801748, + z=(0.9473684210526314-0.3157894736842106j), + expected=(-449.5119088817207+320.1423128036188j), + rtol=5e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-3.956227226099288, + b=1.0561196186065624, + c=8.031683612216888, + z=(0.9473684210526314-0.10526315789473695j), + expected=(0.6361479738012501+0.028575620091205088j), + rtol=5e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=4.25, + b=16.25, + c=16.5, + z=(0.5263157894736841+0.7368421052631575j), + expected=(-0.9038811840552261-1.5356250756164884j), + rtol=1e-8, + ), + marks=pytest.mark.xfail( + reason="Unhandled parameters." + ) + ), + pytest.param( + Hyp2f1TestCase( + a=8.25, + b=-1.75, + c=-1.5, + z=(0.9473684210526314+0.3157894736842106j), + expected=(653.0109150415394-4554.162605155542j), + rtol=1e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-15.75, + b=-3.75, + c=4.5, + z=(0.9473684210526314-0.10526315789473695j), + expected=(118.7009859241035-34.18713648654642j), + rtol=1e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=4.25, + b=-15.75, + c=-3.5, + z=(0.5263157894736841+0.7368421052631575j), + expected=(-540204.4774526551+4970059.109251281j), + rtol=1e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=8.25, + b=-15.75, + c=-0.5, + z=(0.5263157894736841-0.736842105263158j), + expected=(2253490.972258385+3318620.683390017j), + rtol=5e-13, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=16.25, + b=-7.75, + c=-7.5, + z=(0.9473684210526314+0.3157894736842106j), + expected=(-46159826.46716958-17880663.82218242j), + rtol=5e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-0.9220024191881196, + b=-7.949900487447654, + c=-7.93846038215665, + z=(0.9473684210526314-0.10526315789473695j), + expected=(0.07116833581404514+0.11823358038036977j), + rtol=5e-13, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=16.087593263474208, + b=4.0013768449590685, + c=-7.93846038215665, + z=(0.7368421052631575+0.5263157894736841j), + expected=(4.7724909620664006e+17-6.039064078946702e+16j), + rtol=5e-13, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=1.25, + b=-7.75, + c=1.5, + z=(0.9473684210526314-0.10526315789473695j), + expected=(0.0188022179759303+0.002921737281641378j), + rtol=5e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=1.0272592605282642, + b=1.0561196186065624, + c=-7.93846038215665, + z=(0.7368421052631575-0.5263157894736843j), + expected=(-9203.462928334846+12390.110518017136j), + rtol=5e-13, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-3.75, + b=-15.75, + c=8.5, + z=(0.7368421052631575+0.5263157894736841j), + expected=(6.468457061368628+24.190040684917374j), + rtol=6e-16, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=2.02764642551431, + b=16.088264119063613, + c=2.0397202577726152, + z=(0.7368421052631575+0.5263157894736841j), + expected=(2408.3451340186543-4275.257316636014j), + rtol=5e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-3.75, + b=-7.75, + c=8.5, + z=(0.7368421052631575-0.5263157894736843j), + expected=(4.1379984626381345-5.183654781039423j), + rtol=5e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=8.25, + b=-7.75, + c=-0.5, + z=(0.5263157894736841+0.7368421052631575j), + expected=(-81177.775295738+56079.73286548954j), + rtol=5e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=4.080187217753502, + b=2.050308316530781, + c=-0.906685989801748, + z=(0.9473684210526314+0.3157894736842106j), + expected=(1192868.5068926765+3624210.8182139914j), + rtol=5e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-7.937789122896016, + b=-1.92872979730171, + c=8.031683612216888, + z=(0.5263157894736841+0.7368421052631575j), + expected=(1.8286341846195202+1.9295255682312178j), + rtol=5e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=1.0272592605282642, + b=1.0561196186065624, + c=16.056809865262608, + z=(0.7368421052631575-0.5263157894736843j), + expected=(1.0514645669696452-0.0430834059440128j), + rtol=5e-10, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=16.087593263474208, + b=-15.964218273004214, + c=2.0397202577726152, + z=(0.5263157894736841+0.7368421052631575j), + expected=(541983.236432269+288200.2043029435j), + rtol=5e-13, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=2.25, + b=8.25, + c=1.5, + z=(0.5263157894736841-0.736842105263158j), + expected=(-10.931988086039945+1.9136272843579096j), + rtol=1e-15, + ), + ), + ] + ) + def test_region3(self, hyp2f1_test_case): + """0.9 <= |z| <= 1 and |1 - z| < 0.9.""" + a, b, c, z, expected, rtol = hyp2f1_test_case + assert 0.9 <= abs(z) <= 1 and abs(1 - z) < 0.9 # Tests the test + assert_allclose(hyp2f1(a, b, c, z), expected, rtol=rtol) + + @pytest.mark.parametrize( + "hyp2f1_test_case", + [ + pytest.param( + Hyp2f1TestCase( + a=16.25, + b=4.25, + c=2.5, + z=(0.4931034482758623-0.7965517241379311j), + expected=(38.41207903409937-30.510151276075792j), + rtol=5e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=2.0, + b=16.087593263474208, + c=16.088264119063613, + z=(0.5689655172413794-0.7965517241379311j), + expected=(-0.6667857912761286-1.0206224321443573j), + rtol=1e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=8.0, + b=1.0272592605282642, + c=-7.949900487447654, + z=(0.4931034482758623-0.7965517241379311j), + expected=(1679024.1647997478-2748129.775857212j), + rtol=5e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=4.080187217753502, + b=16.0, + c=-7.949900487447654, + z=(0.4931034482758623-0.7965517241379311j), + expected=(424747226301.16986-1245539049327.2856j), + rtol=1e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=2.02764642551431, + b=-15.964218273004214, + c=4.0, + z=(0.4931034482758623-0.7965517241379311j), + expected=(-0.0057826199201757595+0.026359861999025885j), + rtol=5e-06, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=2.02764642551431, + b=-0.9629749245209605, + c=2.0397202577726152, + z=(0.5689655172413794-0.7965517241379311j), + expected=(0.4671901063492606+0.7769632229834897j), + rtol=5e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=2.0, + b=-3.956227226099288, + c=-7.949900487447654, + z=(0.4931034482758623+0.7965517241379312j), + expected=(0.9422283708145973+1.3476905754773343j), + rtol=5e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=1.0, + b=-15.980848054962111, + c=-15.964218273004214, + z=(0.4931034482758623-0.7965517241379311j), + expected=(0.4168719497319604-0.9770953555235625j), + rtol=5e-10, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-0.5, + b=16.088264119063613, + c=2.5, + z=(0.5689655172413794+0.7965517241379312j), + expected=(1.279096377550619-2.173827694297929j), + rtol=5e-12, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-1.9214641416286231, + b=4.0013768449590685, + c=2.0397202577726152, + z=(0.4931034482758623+0.7965517241379312j), + expected=(-2.071520656161738-0.7846098268395909j), + rtol=5e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-0.9220024191881196, + b=8.0, + c=-0.9629749245209605, + z=(0.5689655172413794-0.7965517241379311j), + expected=(-7.740015495862889+3.386766435696699j), + rtol=5e-12, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-1.9214641416286231, + b=16.088264119063613, + c=-7.93846038215665, + z=(0.4931034482758623+0.7965517241379312j), + expected=(-6318.553685853241-7133.416085202879j), + rtol=5e-9, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-15.980848054962111, + b=-3.9316537064827854, + c=16.056809865262608, + z=(0.5689655172413794+0.7965517241379312j), + expected=(-0.8854577905547399+8.135089099967278j), + rtol=5e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-1.9214641416286231, + b=-0.9629749245209605, + c=4.078873014294075, + z=(0.4931034482758623+0.7965517241379312j), + expected=(1.224291301521487+0.36014711766402485j), + rtol=1e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-15.75, + b=-0.75, + c=-1.5, + z=(0.4931034482758623+0.7965517241379312j), + expected=(-1.5765685855028473-3.9399766961046323j), + rtol=1e-3, + ), + marks=pytest.mark.xfail( + reason="Unhandled parameters." + ) + ), + pytest.param( + Hyp2f1TestCase( + a=-15.980848054962111, + b=-1.92872979730171, + c=-7.93846038215665, + z=(0.5689655172413794-0.7965517241379311j), + expected=(56.794588688231194+4.556286783533971j), + rtol=5e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=4.5, + b=4.5, + c=2.050308316530781, + z=(0.5689655172413794+0.7965517241379312j), + expected=(-4.251456563455306+6.737837111569671j), + rtol=5e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=4.5, + b=8.5, + c=-1.92872979730171, + z=(0.4931034482758623-0.7965517241379311j), + expected=(2177143.9156599627-3313617.2748088865j), + rtol=5e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=2.5, + b=-1.5, + c=4.0013768449590685, + z=(0.4931034482758623-0.7965517241379311j), + expected=(0.45563554481603946+0.6212000158060831j), + rtol=5e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=8.5, + b=-7.5, + c=-15.964218273004214, + z=(0.4931034482758623+0.7965517241379312j), + expected=(61.03201617828073-37.185626416756214j), + rtol=5e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-15.5, + b=16.5, + c=4.0013768449590685, + z=(0.4931034482758623+0.7965517241379312j), + expected=(-33143.425963520735+20790.608514722644j), + rtol=1e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-0.5, + b=4.5, + c=-0.9629749245209605, + z=(0.5689655172413794+0.7965517241379312j), + expected=(30.778600270824423-26.65160354466787j), + rtol=5e-13, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-0.5, + b=-3.5, + c=16.088264119063613, + z=(0.5689655172413794-0.7965517241379311j), + expected=(1.0629792615560487-0.08308454486044772j), + rtol=1e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-3.5, + b=-7.5, + c=-0.9629749245209605, + z=(0.4931034482758623-0.7965517241379311j), + expected=(17431.571802591767+3553.7129767034507j), + rtol=5e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=2.25, + b=8.25, + c=16.5, + z=(0.11379310344827598+0.9482758620689657j), + expected=(0.4468600750211926+0.7313214934036885j), + rtol=1e-3, + ), + marks=pytest.mark.xfail( + reason="Unhandled parameters." + ) + ), + pytest.param( + Hyp2f1TestCase( + a=8.25, + b=16.25, + c=4.5, + z=(0.3413793103448277+0.8724137931034486j), + expected=(-3.905704438293991+3.693347860329299j), + rtol=5e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=4.25, + b=4.25, + c=-0.5, + z=(0.11379310344827598-0.9482758620689655j), + expected=(-40.31777941834244-89.89852492432011j), + rtol=5e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=1.0272592605282642, + b=8.0, + c=-15.964218273004214, + z=(0.11379310344827598-0.9482758620689655j), + expected=(52584.347773055284-109197.86244309516j), + rtol=5e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=8.095813935368371, + b=-15.964218273004214, + c=16.056809865262608, + z=(0.03793103448275881+0.9482758620689657j), + expected=(-1.187733570412592-1.5147865053584582j), + rtol=5e-10, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=4.080187217753502, + b=-3.9316537064827854, + c=1.0651378143226575, + z=(0.26551724137931054+0.9482758620689657j), + expected=(13.077494677898947+35.071599628224966j), + rtol=5e-13, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=4.080187217753502, + b=-3.5, + c=-3.5, + z=(0.26551724137931054+0.8724137931034486j), + expected=(-0.5359656237994614-0.2344483936591811j), + rtol=5e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=4.25, + b=-3.75, + c=-1.5, + z=(0.26551724137931054+0.9482758620689657j), + expected=(1204.8114871663133+64.41022826840198j), + rtol=5e-13, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-1.9214641416286231, + b=16.0, + c=4.0013768449590685, + z=(0.03793103448275881-0.9482758620689655j), + expected=(-9.85268872413994+7.011107558429154j), + rtol=5e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-7.937789122896016, + b=16.0, + c=4.0013768449590685, + z=(0.3413793103448277-0.8724137931034484j), + expected=(528.5522951158454-1412.21630264791j), + rtol=5e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-15.5, + b=1.0561196186065624, + c=-7.5, + z=(0.4172413793103451+0.8724137931034486j), + expected=(133306.45260685298+256510.7045225382j), + rtol=5e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-7.937789122896016, + b=8.077282662161238, + c=-15.963511401609862, + z=(0.3413793103448277-0.8724137931034484j), + expected=(-0.998555715276967+2.774198742229889j), + rtol=5e-11, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-7.75, + b=-0.75, + c=1.5, + z=(0.11379310344827598-0.9482758620689655j), + expected=(2.072445019723025-2.9793504811373515j), + rtol=5e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-15.5, + b=-1.92872979730171, + c=1.5, + z=(0.11379310344827598-0.9482758620689655j), + expected=(-41.87581944176649-32.52980303527139j), + rtol=5e-13, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-3.75, + b=-15.75, + c=-0.5, + z=(0.11379310344827598-0.9482758620689655j), + expected=(-3729.6214864209774-30627.510509112635j), + rtol=5e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-3.956227226099288, + b=-15.964218273004214, + c=-0.906685989801748, + z=(0.03793103448275881+0.9482758620689657j), + expected=(-131615.07820609974+145596.13384245415j), + rtol=5e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=1.5, + b=16.5, + c=16.088264119063613, + z=(0.26551724137931054+0.8724137931034486j), + expected=(0.18981844071070744+0.7855036242583742j), + rtol=1e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=16.5, + b=8.5, + c=-3.9316537064827854, + z=(0.11379310344827598-0.9482758620689655j), + expected=(110224529.2376068+128287212.04290268j), + rtol=5e-13, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=2.5, + b=-7.5, + c=4.0013768449590685, + z=(0.3413793103448277-0.8724137931034484j), + expected=(0.2722302180888523-0.21790187837266162j), + rtol=1.2e-12, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=8.5, + b=-7.5, + c=-15.964218273004214, + z=(0.11379310344827598-0.9482758620689655j), + expected=(-2.8252338010989035+2.430661949756161j), + rtol=5e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-3.5, + b=16.5, + c=4.0013768449590685, + z=(0.03793103448275881+0.9482758620689657j), + expected=(-20.604894257647945+74.5109432558078j), + rtol=5e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-7.5, + b=8.5, + c=-0.9629749245209605, + z=(0.3413793103448277+0.8724137931034486j), + expected=(-2764422.521269463-3965966.9965808876j), + rtol=5e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-1.5, + b=-0.5, + c=1.0561196186065624, + z=(0.26551724137931054+0.9482758620689657j), + expected=(1.2262338560994905+0.6545051266925549j), + rtol=1e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-0.5, + b=-15.5, + c=-7.949900487447654, + z=(0.4172413793103451-0.8724137931034484j), + expected=(-2258.1590330318213+8860.193389158803j), + rtol=1.4e-10, + ), + ), + ] + ) + def test_region4(self, hyp2f1_test_case): + """0.9 <= |z| <= 1 and |1 - z| >= 1. + + This region is unhandled by of the standard transformations and + needs special care. + """ + a, b, c, z, expected, rtol = hyp2f1_test_case + assert 0.9 <= abs(z) <= 1 and abs(1 - z) >= 0.9 # Tests the test + assert_allclose(hyp2f1(a, b, c, z), expected, rtol=rtol) + + @pytest.mark.parametrize( + "hyp2f1_test_case", + [ + pytest.param( + Hyp2f1TestCase( + a=4.5, + b=16.088264119063613, + c=8.5, + z=(0.6448275862068968+0.8724137931034486j), + expected=(0.018601324701770394-0.07618420586062377j), + rtol=5e-08, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=8.25, + b=4.25, + c=4.5, + z=(0.6448275862068968-0.8724137931034484j), + expected=(-1.391549471425551-0.118036604903893j), + rtol=5e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=2.02764642551431, + b=2.050308316530781, + c=-1.9631175993998025, + z=(0.6448275862068968+0.8724137931034486j), + expected=(-2309.178768155151-1932.7247727595172j), + rtol=5e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=16.087593263474208, + b=1.0, + c=-15.964218273004214, + z=(0.6448275862068968+0.8724137931034486j), + expected=(85592537010.05054-8061416766688.324j), + rtol=2e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=8.095813935368371, + b=-0.5, + c=1.5, + z=(0.6448275862068968+0.8724137931034486j), + expected=(1.2334498208515172-2.1639498536219732j), + rtol=5e-11, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=16.087593263474208, + b=-15.964218273004214, + c=4.0, + z=(0.6448275862068968+0.8724137931034486j), + expected=(102266.35398605966-44976.97828737755j), + rtol=1e-3, + ), + marks=pytest.mark.xfail( + reason="Unhandled parameters." + ) + ), + pytest.param( + Hyp2f1TestCase( + a=4.0, + b=-3.956227226099288, + c=-15.964218273004214, + z=(0.6448275862068968-0.8724137931034484j), + expected=(-2.9590030930007236-4.190770764773225j), + rtol=5e-13, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=4.080187217753502, + b=-15.5, + c=-7.5, + z=(0.5689655172413794-0.8724137931034484j), + expected=(-112554838.92074208+174941462.9202412j), + rtol=5e-05, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-15.980848054962111, + b=2.050308316530781, + c=1.0, + z=(0.6448275862068968-0.8724137931034484j), + expected=(3.7519882374080145+7.360753798667486j), + rtol=5e-13, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-7.937789122896016, + b=2.050308316530781, + c=4.0, + z=(0.6448275862068968-0.8724137931034484j), + expected=(0.000181132943964693+0.07742903103815582j), + rtol=5e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-7.937789122896016, + b=4.0013768449590685, + c=-1.9631175993998025, + z=(0.5689655172413794+0.8724137931034486j), + expected=(386338.760913596-386166.51762171905j), + rtol=5e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-15.980848054962111, + b=8.0, + c=-1.92872979730171, + z=(0.6448275862068968+0.8724137931034486j), + expected=(1348667126.3444858-2375132427.158893j), + rtol=5e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-3.5, + b=-0.9629749245209605, + c=4.5, + z=(0.5689655172413794+0.8724137931034486j), + expected=(1.428353429538678+0.6472718120804372j), + rtol=5e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-7.937789122896016, + b=-0.9629749245209605, + c=2.0397202577726152, + z=(0.5689655172413794-0.8724137931034484j), + expected=(3.1439267526119643-3.145305240375117j), + rtol=5e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-1.9214641416286231, + b=-15.964218273004214, + c=-7.93846038215665, + z=(0.6448275862068968-0.8724137931034484j), + expected=(75.27467675681773+144.0946946292215j), + rtol=1e-07, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-3.75, + b=-7.75, + c=-7.5, + z=(0.5689655172413794+0.8724137931034486j), + expected=(-0.3699450626264222+0.8732812475910993j), + rtol=5e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=1.5, + b=16.5, + c=1.0561196186065624, + z=(0.5689655172413794-0.8724137931034484j), + expected=(5.5361025821300665-2.4709693474656285j), + rtol=5e-09, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=1.5, + b=8.5, + c=-3.9316537064827854, + z=(0.6448275862068968-0.8724137931034484j), + expected=(-782805.6699207705-537192.581278909j), + rtol=5e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=2.5, + b=-15.5, + c=1.0561196186065624, + z=(0.6448275862068968+0.8724137931034486j), + expected=(12.345113400639693-14.993248992902007j), + rtol=0.0005, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=1.5, + b=-0.5, + c=-15.964218273004214, + z=(0.6448275862068968+0.8724137931034486j), + expected=(23.698109392667842+97.15002033534108j), + rtol=5e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-7.5, + b=16.5, + c=4.0013768449590685, + z=(0.6448275862068968-0.8724137931034484j), + expected=(1115.2978631811834+915.9212658718577j), + rtol=5e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-15.5, + b=16.5, + c=-0.9629749245209605, + z=(0.6448275862068968+0.8724137931034486j), + expected=(642077722221.6489+535274495398.21027j), + rtol=5e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-7.5, + b=-3.5, + c=4.0013768449590685, + z=(0.5689655172413794+0.8724137931034486j), + expected=(-5.689219222945697+16.877463062787143j), + rtol=5e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-15.5, + b=-1.5, + c=-0.9629749245209605, + z=(0.5689655172413794-0.8724137931034484j), + expected=(-44.32070290703576+1026.9127058617403j), + rtol=5e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=16.25, + b=2.25, + c=4.5, + z=(0.11379310344827598-1.024137931034483j), + expected=(-0.021965227124574663+0.009908300237809064j), + rtol=1e-3, + ), + marks=pytest.mark.xfail( + reason="Unhandled parameters." + ) + ), + pytest.param( + Hyp2f1TestCase( + a=2.02764642551431, + b=1.5, + c=16.5, + z=(0.26551724137931054+1.024137931034483j), + expected=(1.0046072901244183+0.19945500134119992j), + rtol=5e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=16.087593263474208, + b=1.0, + c=-3.9316537064827854, + z=(0.3413793103448277+0.9482758620689657j), + expected=(21022.30133421465+49175.98317370489j), + rtol=5e-13, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=4.080187217753502, + b=16.088264119063613, + c=-1.9631175993998025, + z=(0.4172413793103451-0.9482758620689655j), + expected=(-7024239.358547302+2481375.02681063j), + rtol=5e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=16.25, + b=-15.75, + c=1.5, + z=(0.18965517241379315+1.024137931034483j), + expected=(92371704.94848-403546832.548352j), + rtol=5e-06, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=8.5, + b=-7.949900487447654, + c=8.5, + z=(0.26551724137931054-1.024137931034483j), + expected=(1.9335109845308265+5.986542524829654j), + rtol=5e-10, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=8.095813935368371, + b=-1.92872979730171, + c=-7.93846038215665, + z=(0.4931034482758623+0.8724137931034486j), + expected=(-122.52639696039328-59.72428067512221j), + rtol=5e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=16.25, + b=-1.75, + c=-1.5, + z=(0.4931034482758623+0.9482758620689657j), + expected=(-90.40642053579428+50.50649180047921j), + rtol=5e-08, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-3.5, + b=8.077282662161238, + c=16.5, + z=(0.4931034482758623+0.9482758620689657j), + expected=(-0.2155745818150323-0.564628986876639j), + rtol=5e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-0.9220024191881196, + b=1.0561196186065624, + c=8.031683612216888, + z=(0.4172413793103451-0.9482758620689655j), + expected=(0.9503140488280465+0.11574960074292677j), + rtol=5e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-0.75, + b=2.25, + c=-15.5, + z=(0.4172413793103451+0.9482758620689657j), + expected=(0.9285862488442175+0.8203699266719692j), + rtol=5e-13, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-7.75, + b=4.25, + c=-15.5, + z=(0.3413793103448277-0.9482758620689655j), + expected=(-1.0509834850116921-1.1145522325486075j), + rtol=1.1e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-7.937789122896016, + b=-0.9629749245209605, + c=2.0397202577726152, + z=(0.4931034482758623-0.9482758620689655j), + expected=(2.88119116536769-3.4249933450696806j), + rtol=5e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-15.5, + b=-15.964218273004214, + c=16.5, + z=(0.18965517241379315+1.024137931034483j), + expected=(199.65868451496038+347.79384207302877j), + rtol=5e-12, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-15.75, + b=-15.75, + c=-3.5, + z=(0.4931034482758623-0.8724137931034484j), + expected=(-208138312553.07013+58631611809.026955j), + rtol=5e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-7.937789122896016, + b=-15.5, + c=-7.5, + z=(0.3413793103448277+0.9482758620689657j), + expected=(-23032.90519856288-18256.94050457296j), + rtol=5e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=4.5, + b=1.5, + c=1.0561196186065624, + z=(0.4931034482758623-0.8724137931034484j), + expected=(1.507342459587056+1.2332023580148403j), + rtol=1e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=2.5, + b=4.5, + c=-3.9316537064827854, + z=(0.4172413793103451+0.9482758620689657j), + expected=(7044.766127108853-40210.365567285575j), + rtol=5e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=1.5, + b=-1.5, + c=1.0561196186065624, + z=(0.03793103448275881+1.024137931034483j), + expected=(0.2725347741628333-2.247314875514784j), + rtol=1e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=4.5, + b=-1.5, + c=-7.949900487447654, + z=(0.26551724137931054+1.024137931034483j), + expected=(-11.250200011017546+12.597393659160472j), + rtol=5e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-7.5, + b=8.5, + c=16.088264119063613, + z=(0.26551724137931054+1.024137931034483j), + expected=(-0.18515160890991517+0.7959014164484782j), + rtol=2e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-7.5, + b=16.5, + c=-3.9316537064827854, + z=(0.3413793103448277-1.024137931034483j), + expected=(998246378.8556538+1112032928.103645j), + rtol=5e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-1.5, + b=-3.5, + c=2.050308316530781, + z=(0.03793103448275881+1.024137931034483j), + expected=(0.5527670397711952+2.697662715303637j), + rtol=1.2e-15, # rtol bumped from 1e-15 in gh18414 + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-15.5, + b=-1.5, + c=-0.9629749245209605, + z=(0.4931034482758623-0.8724137931034484j), + expected=(55.396931662136886+968.467463806326j), + rtol=5e-14, + ), + ), + ] + ) + def test_region5(self, hyp2f1_test_case): + """1 < |z| < 1.1 and |1 - z| >= 0.9 and real(z) >= 0""" + a, b, c, z, expected, rtol = hyp2f1_test_case + assert 1 < abs(z) < 1.1 and abs(1 - z) >= 0.9 and z.real >= 0 + assert_allclose(hyp2f1(a, b, c, z), expected, rtol=rtol) + + @pytest.mark.parametrize( + "hyp2f1_test_case", + [ + pytest.param( + Hyp2f1TestCase( + a=8.095813935368371, + b=4.0013768449590685, + c=4.078873014294075, + z=(-0.9473684210526316+0.5263157894736841j), + expected=(-0.0018093573941378783+0.003481887377423739j), + rtol=5e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=16.087593263474208, + b=2.050308316530781, + c=1.0651378143226575, + z=(-0.736842105263158-0.736842105263158j), + expected=(-0.00023401243818780545-1.7983496305603562e-05j), + rtol=1e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=1.0272592605282642, + b=8.077282662161238, + c=4.078873014294075, + z=(-0.5263157894736843-0.9473684210526316j), + expected=(0.22359773002226846-0.24092487123993353j), + rtol=1e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=1.0272592605282642, + b=2.050308316530781, + c=-15.963511401609862, + z=(-0.9473684210526316-0.5263157894736843j), + expected=(1.191573745740011+0.14347394589721466j), + rtol=5e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=4.080187217753502, + b=4.0013768449590685, + c=-15.963511401609862, + z=(-0.9473684210526316-0.5263157894736843j), + expected=(31.822620756901784-66.09094396747611j), + rtol=5e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=4.080187217753502, + b=8.077282662161238, + c=-7.93846038215665, + z=(-0.9473684210526316+0.5263157894736841j), + expected=(207.16750179245952+34.80478274924269j), + rtol=5e-12, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=8.095813935368371, + b=-7.949900487447654, + c=8.031683612216888, + z=(-0.736842105263158+0.7368421052631575j), + expected=(-159.62429364277145+9.154224290644898j), + rtol=5e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=1.0272592605282642, + b=-1.92872979730171, + c=16.056809865262608, + z=(-0.9473684210526316+0.5263157894736841j), + expected=(1.121122351247184-0.07170260470126685j), + rtol=5e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=16.087593263474208, + b=-0.9629749245209605, + c=16.056809865262608, + z=(-0.9473684210526316+0.5263157894736841j), + expected=(1.9040596681316053-0.4951799449960107j), + rtol=5e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=1.0272592605282642, + b=-1.92872979730171, + c=-0.906685989801748, + z=(-0.9473684210526316-0.5263157894736843j), + expected=(-14.496623497780739-21.897524523299875j), + rtol=5e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=4.080187217753502, + b=-3.9316537064827854, + c=-3.9924618758357022, + z=(-0.5263157894736843-0.9473684210526316j), + expected=(36.33473466026878+253.88728442029577j), + rtol=5e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=1.0272592605282642, + b=-15.964218273004214, + c=-0.906685989801748, + z=(-0.9473684210526316+0.5263157894736841j), + expected=(1505052.5653144997-50820766.81043443j), + rtol=1e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-3.956227226099288, + b=4.0013768449590685, + c=1.0651378143226575, + z=(-0.5263157894736843+0.9473684210526314j), + expected=(-127.79407519260877-28.69899444941112j), + rtol=5e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-1.9214641416286231, + b=8.077282662161238, + c=16.056809865262608, + z=(-0.9473684210526316-0.5263157894736843j), + expected=(2.0623331933754976+0.741234463565458j), + rtol=5e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-3.956227226099288, + b=8.077282662161238, + c=2.0397202577726152, + z=(-0.9473684210526316+0.5263157894736841j), + expected=(30.729193458862525-292.5700835046965j), + rtol=1e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-1.9214641416286231, + b=1.0561196186065624, + c=-1.9631175993998025, + z=(-0.5263157894736843-0.9473684210526316j), + expected=(1.1285917906203495-0.735264575450189j), + rtol=5e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-0.9220024191881196, + b=1.0561196186065624, + c=-3.9924618758357022, + z=(-0.736842105263158+0.7368421052631575j), + expected=(0.6356474446678052-0.02429663008952248j), + rtol=5e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-1.9214641416286231, + b=16.088264119063613, + c=-7.93846038215665, + z=(-0.736842105263158+0.7368421052631575j), + expected=(0.4718880510273174+0.655083067736377j), + rtol=1e-11, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-7.937789122896016, + b=-3.9316537064827854, + c=16.056809865262608, + z=(-0.9473684210526316+0.5263157894736841j), + expected=(-0.14681550942352714+0.16092206364265146j), + rtol=5e-11, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-0.9220024191881196, + b=-15.964218273004214, + c=1.0651378143226575, + z=(-0.5263157894736843+0.9473684210526314j), + expected=(-6.436835190526225+22.883156700606182j), + rtol=5e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-0.9220024191881196, + b=-7.949900487447654, + c=4.078873014294075, + z=(-0.9473684210526316-0.5263157894736843j), + expected=(-0.7505682955068583-1.1026583264249945j), + rtol=5e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-3.956227226099288, + b=-3.9316537064827854, + c=-7.93846038215665, + z=(-0.9473684210526316-0.5263157894736843j), + expected=(3.6247814989198166+2.596041360148318j), + rtol=5e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-3.956227226099288, + b=-15.964218273004214, + c=-1.9631175993998025, + z=(-0.5263157894736843-0.9473684210526316j), + expected=(-59537.65287927933-669074.4342539902j), + rtol=5e-15, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=-3.956227226099288, + b=-15.964218273004214, + c=-1.9631175993998025, + z=(-0.9473684210526316-0.5263157894736843j), + expected=(-433084.9970266166+431088.393918521j), + rtol=5e-14, + ), + ), + pytest.param( + Hyp2f1TestCase( + a=1, + b=1, + c=4, + z=(3 + 4j), + expected=(0.49234384000963544+0.6051340616612397j), + rtol=5e-14, + ), + ), + ] + ) + def test_region6(self, hyp2f1_test_case): + """|z| > 1 but not in region 5.""" + a, b, c, z, expected, rtol = hyp2f1_test_case + assert ( + abs(z) > 1 and + not (1 < abs(z) < 1.1 and abs(1 - z) >= 0.9 and z.real >= 0) + ) + assert_allclose(hyp2f1(a, b, c, z), expected, rtol=rtol) + + + @pytest.mark.parametrize( + "hyp2f1_test_case", + [ + # Broke when fixing gamma pole behavior in gh-21827 + pytest.param( + Hyp2f1TestCase( + a=1.3, + b=-0.2, + c=0.3, + z=-2.1, + expected=1.8202169687521206, + rtol=5e-15, + ), + ), + ] + ) + def test_miscellaneous(self, hyp2f1_test_case ): + a, b, c, z, expected, rtol = hyp2f1_test_case + assert_allclose(hyp2f1(a, b, c, z), expected, rtol=rtol) + + @pytest.mark.slow + @check_version(mpmath, "1.0.0") + def test_test_hyp2f1(self): + """Test that expected values match what is computed by mpmath. + + This gathers the parameters for the test cases out of the pytest marks. + The parameters are a, b, c, z, expected, rtol, where expected should + be the value of hyp2f1(a, b, c, z) computed with mpmath. The test + recomputes hyp2f1(a, b, c, z) using mpmath and verifies that expected + actually is the correct value. This allows the data for the tests to + live within the test code instead of an external datafile, while + avoiding having to compute the results with mpmath during the test, + except for when slow tests are being run. + """ + test_methods = [ + test_method for test_method in dir(self) + if test_method.startswith('test') and + # Filter properties and attributes (futureproofing). + callable(getattr(self, test_method)) and + # Filter out this test + test_method != 'test_test_hyp2f1' + ] + for test_method in test_methods: + params = self._get_test_parameters(getattr(self, test_method)) + for a, b, c, z, expected, _ in params: + assert_allclose(mp_hyp2f1(a, b, c, z), expected, rtol=2.25e-16) + + def _get_test_parameters(self, test_method): + """Get pytest.mark parameters for a test in this class.""" + return [ + case.values[0] for mark in test_method.pytestmark + if mark.name == 'parametrize' + for case in mark.args[1] + ] + +class TestHyp2f1ExtremeInputs: + + @pytest.mark.parametrize("a", [1.0, 2.0, 3.0, -np.inf, np.inf]) + @pytest.mark.parametrize("b", [3.0, 4.0, 5.0, -np.inf, np.inf]) + @pytest.mark.parametrize("c", [3.0, 5.0, 6.0, 7.0]) + @pytest.mark.parametrize("z", [4.0 + 1.0j]) + def test_inf_a_b(self, a, b, c, z): + if np.any(np.isinf(np.asarray([a, b]))): + assert(np.isnan(hyp2f1(a, b, c, z))) + + def test_large_a_b(self): + assert(np.isnan(hyp2f1(10**7, 1.0, 3.0, 4.0 + 1.0j))) + assert(np.isnan(hyp2f1(-10**7, 1.0, 3.0, 4.0 + 1.0j))) + + assert(np.isnan(hyp2f1(1.0, 10**7, 3.0, 4.0 + 1.0j))) + assert(np.isnan(hyp2f1(1.0, -10**7, 3.0, 4.0 + 1.0j))) + + # Already correct in main but testing for surety + assert(np.isnan(hyp2f1(np.inf, 1.0, 3.0, 4.0))) + assert(np.isnan(hyp2f1(1.0, np.inf, 3.0, 4.0))) diff --git a/v12x/pylib/scipy/special/tests/test_hypergeometric.py b/v12x/pylib/scipy/special/tests/test_hypergeometric.py new file mode 100644 index 0000000000000000000000000000000000000000..1ad092491905ae66c39cd881836054aef3dc6f6e --- /dev/null +++ b/v12x/pylib/scipy/special/tests/test_hypergeometric.py @@ -0,0 +1,234 @@ +import pytest +import numpy as np +from numpy.testing import assert_allclose, assert_equal +import scipy.special as sc + + +class TestHyperu: + + def test_negative_x(self): + a, b, x = np.meshgrid( + [-1, -0.5, 0, 0.5, 1], + [-1, -0.5, 0, 0.5, 1], + np.linspace(-100, -1, 10), + ) + assert np.all(np.isnan(sc.hyperu(a, b, x))) + + def test_special_cases(self): + assert sc.hyperu(0, 1, 1) == 1.0 + + @pytest.mark.parametrize('a', [0.5, 1, np.nan]) + @pytest.mark.parametrize('b', [1, 2, np.nan]) + @pytest.mark.parametrize('x', [0.25, 3, np.nan]) + def test_nan_inputs(self, a, b, x): + assert np.isnan(sc.hyperu(a, b, x)) == np.any(np.isnan([a, b, x])) + + @pytest.mark.parametrize( + 'a,b,x,expected', + [(0.21581740448533887, 1.0, 1e-05, 3.6030558839391325), + (0.21581740448533887, 1.0, 0.00021544346900318823, 2.8783254988948976), + (0.21581740448533887, 1.0, 0.004641588833612777, 2.154928216691109), + (0.21581740448533887, 1.0, 0.1, 1.446546638718792), + (0.0030949064301273865, 1.0, 1e-05, 1.0356696454116199), + (0.0030949064301273865, 1.0, 0.00021544346900318823, 1.0261510362481985), + (0.0030949064301273865, 1.0, 0.004641588833612777, 1.0166326903402296), + (0.0030949064301273865, 1.0, 0.1, 1.0071174207698674), + (0.1509924314279033, 1.0, 1e-05, 2.806173846998948), + (0.1509924314279033, 1.0, 0.00021544346900318823, 2.3092158526816124), + (0.1509924314279033, 1.0, 0.004641588833612777, 1.812905980588048), + (0.1509924314279033, 1.0, 0.1, 1.3239738117634872), + (-0.010678995342969011, 1.0, 1e-05, 0.8775194903781114), + (-0.010678995342969011, 1.0, 0.00021544346900318823, 0.9101008998540128), + (-0.010678995342969011, 1.0, 0.004641588833612777, 0.9426854294058609), + (-0.010678995342969011, 1.0, 0.1, 0.9753065150174902), + (-0.06556622211831487, 1.0, 1e-05, 0.26435429752668904), + (-0.06556622211831487, 1.0, 0.00021544346900318823, 0.4574756033875781), + (-0.06556622211831487, 1.0, 0.004641588833612777, 0.6507121093358457), + (-0.06556622211831487, 1.0, 0.1, 0.8453129788602187), + (-0.21628242470175185, 1.0, 1e-05, -1.2318314201114489), + (-0.21628242470175185, 1.0, 0.00021544346900318823, -0.6704694233529538), + (-0.21628242470175185, 1.0, 0.004641588833612777, -0.10795098653682857), + (-0.21628242470175185, 1.0, 0.1, 0.4687227684115524)] + ) + def test_gh_15650_mp(self, a, b, x, expected): + # See https://github.com/scipy/scipy/issues/15650 + # b == 1, |a| < 0.25, 0 < x < 1 + # + # This purpose of this test is to check the accuracy of results + # in the region that was impacted by gh-15650. + # + # Reference values computed with mpmath using the script: + # + # import itertools as it + # import numpy as np + # + # from mpmath import mp + # + # rng = np.random.default_rng(1234) + # + # cases = [] + # for a, x in it.product( + # np.random.uniform(-0.25, 0.25, size=6), + # np.logspace(-5, -1, 4), + # ): + # with mp.workdps(100): + # cases.append((float(a), 1.0, float(x), float(mp.hyperu(a, 1.0, x)))) + assert_allclose(sc.hyperu(a, b, x), expected, rtol=1e-13) + + def test_gh_15650_sanity(self): + # The purpose of this test is to sanity check hyperu in the region that + # was impacted by gh-15650 by making sure there are no excessively large + # results, as were reported there. + a = np.linspace(-0.5, 0.5, 500) + x = np.linspace(1e-6, 1e-1, 500) + a, x = np.meshgrid(a, x) + results = sc.hyperu(a, 1.0, x) + assert np.all(np.abs(results) < 1e3) + + +class TestHyp1f1: + + @pytest.mark.parametrize('a, b, x', [ + (np.nan, 1, 1), + (1, np.nan, 1), + (1, 1, np.nan) + ]) + def test_nan_inputs(self, a, b, x): + assert np.isnan(sc.hyp1f1(a, b, x)) + + def test_poles(self): + assert_equal(sc.hyp1f1(1, [0, -1, -2, -3, -4], 0.5), np.inf) + + @pytest.mark.parametrize('a, b, x, result', [ + (-1, 1, 0.5, 0.5), + (1, 1, 0.5, 1.6487212707001281468), + (2, 1, 0.5, 2.4730819060501922203), + (1, 2, 0.5, 1.2974425414002562937), + (-10, 1, 0.5, -0.38937441413785204475) + ]) + def test_special_cases(self, a, b, x, result): + # Hit all the special case branches at the beginning of the + # function. Desired answers computed using Mpmath. + assert_allclose(sc.hyp1f1(a, b, x), result, atol=0, rtol=1e-15) + + @pytest.mark.parametrize('a, b, x, result', [ + (1, 1, 0.44, 1.5527072185113360455), + (-1, 1, 0.44, 0.55999999999999999778), + (100, 100, 0.89, 2.4351296512898745592), + (-100, 100, 0.89, 0.40739062490768104667), + (1.5, 100, 59.99, 3.8073513625965598107), + (-1.5, 100, 59.99, 0.25099240047125826943) + ]) + def test_geometric_convergence(self, a, b, x, result): + # Test the region where we are relying on the ratio of + # + # (|a| + 1) * |x| / |b| + # + # being small. Desired answers computed using Mpmath + assert_allclose(sc.hyp1f1(a, b, x), result, atol=0, rtol=1e-15) + + @pytest.mark.parametrize('a, b, x, result', [ + (-1, 1, 1.5, -0.5), + (-10, 1, 1.5, 0.41801777430943080357), + (-25, 1, 1.5, 0.25114491646037839809), + (-50, 1, 1.5, -0.25683643975194756115), + (-80, 1, 1.5, -0.24554329325751503601), + (-150, 1, 1.5, -0.173364795515420454496), + ]) + def test_a_negative_integer(self, a, b, x, result): + # Desired answers computed using Mpmath. + assert_allclose(sc.hyp1f1(a, b, x), result, atol=0, rtol=2e-14) + + @pytest.mark.parametrize('a, b, x, expected', [ + (0.01, 150, -4, 0.99973683897677527773), # gh-3492 + (1, 5, 0.01, 1.0020033381011970966), # gh-3593 + (50, 100, 0.01, 1.0050126452421463411), # gh-3593 + (1, 0.3, -1e3, -7.011932249442947651455e-04), # gh-14149 + (1, 0.3, -1e4, -7.001190321418937164734e-05), # gh-14149 + (9, 8.5, -350, -5.224090831922378361082e-20), # gh-17120 + (9, 8.5, -355, -4.595407159813368193322e-20), # gh-17120 + (75, -123.5, 15, 3.425753920814889017493e+06), + ]) + def test_assorted_cases(self, a, b, x, expected): + # Expected values were computed with mpmath.hyp1f1(a, b, x). + assert_allclose(sc.hyp1f1(a, b, x), expected, atol=0, rtol=1e-14) + + def test_a_neg_int_and_b_equal_x(self): + # This is a case where the Boost wrapper will call hypergeometric_pFq + # instead of hypergeometric_1F1. When we use a version of Boost in + # which https://github.com/boostorg/math/issues/833 is fixed, this + # test case can probably be moved into test_assorted_cases. + # The expected value was computed with mpmath.hyp1f1(a, b, x). + a = -10.0 + b = 2.5 + x = 2.5 + expected = 0.0365323664364104338721 + computed = sc.hyp1f1(a, b, x) + assert_allclose(computed, expected, atol=0, rtol=1e-13) + + @pytest.mark.parametrize('a, b, x, desired', [ + (-1, -2, 2, 2), + (-1, -4, 10, 3.5), + (-2, -2, 1, 2.5) + ]) + def test_gh_11099(self, a, b, x, desired): + # All desired results computed using Mpmath + assert sc.hyp1f1(a, b, x) == desired + + @pytest.mark.parametrize('a', [-3, -2]) + def test_x_zero_a_and_b_neg_ints_and_a_ge_b(self, a): + assert sc.hyp1f1(a, -3, 0) == 1 + + # In the following tests with complex z, the reference values + # were computed with mpmath.hyp1f1(a, b, z), and verified with + # Wolfram Alpha Hypergeometric1F1(a, b, z), except for the + # case a=0.1, b=1, z=7-24j, where Wolfram Alpha reported + # "Standard computation time exceeded". That reference value + # was confirmed in an online Matlab session, with the commands + # + # > format long + # > hypergeom(0.1, 1, 7-24i) + # ans = + # -3.712349651834209 + 4.554636556672912i + # + @pytest.mark.parametrize( + 'a, b, z, ref', + [(-0.25, 0.5, 1+2j, 1.1814553180903435-1.2792130661292984j), + (0.25, 0.5, 1+2j, 0.24636797405707597+1.293434354945675j), + (25, 1.5, -2j, -516.1771262822523+407.04142751922024j), + (12, -1.5, -10+20j, -5098507.422706547-1341962.8043508842j), + pytest.param( + 10, 250, 10-15j, 1.1985998416598884-0.8613474402403436j, + marks=pytest.mark.xfail, + ), + pytest.param( + 0.1, 1, 7-24j, -3.712349651834209+4.554636556672913j, + marks=pytest.mark.xfail, + ) + ], + ) + def test_complex_z(self, a, b, z, ref): + h = sc.hyp1f1(a, b, z) + assert_allclose(h, ref, rtol=4e-15) + + # The "legacy edge cases" mentioned in the comments in the following + # tests refers to the behavior of hyp1f1(a, b, x) when b is a nonpositive + # integer. In some subcases, the behavior of SciPy does not match that + # of Boost (1.81+), mpmath and Mathematica (via Wolfram Alpha online). + # If the handling of these edges cases is changed to agree with those + # libraries, these test will have to be updated. + + @pytest.mark.parametrize('b', [0, -1, -5]) + def test_legacy_case1(self, b): + # Test results of hyp1f1(0, n, x) for n <= 0. + # This is a legacy edge case. + # Boost (versions greater than 1.80), Mathematica (via Wolfram Alpha + # online) and mpmath all return 1 in this case, but SciPy's hyp1f1 + # returns inf. + assert_equal(sc.hyp1f1(0, b, [-1.5, 0, 1.5]), [np.inf, np.inf, np.inf]) + + def test_legacy_case2(self): + # This is a legacy edge case. + # In software such as boost (1.81+), mpmath and Mathematica, + # the value is 1. + assert sc.hyp1f1(-4, -3, 0) == np.inf diff --git a/v12x/pylib/scipy/special/tests/test_iv_ratio.py b/v12x/pylib/scipy/special/tests/test_iv_ratio.py new file mode 100644 index 0000000000000000000000000000000000000000..dafc35982c4d4869ead2136e74466f1c44c64556 --- /dev/null +++ b/v12x/pylib/scipy/special/tests/test_iv_ratio.py @@ -0,0 +1,249 @@ +# This file contains unit tests for iv_ratio() and related functions. + +import pytest +import numpy as np +from numpy.testing import assert_equal, assert_allclose +from scipy.special._ufuncs import ( # type: ignore[attr-defined] + _iv_ratio as iv_ratio, + _iv_ratio_c as iv_ratio_c, +) + + +class TestIvRatio: + + @pytest.mark.parametrize('v,x,r', [ + (0.5, 0.16666666666666666, 0.16514041292462933), + (0.5, 0.3333333333333333, 0.32151273753163434), + (0.5, 0.5, 0.46211715726000974), + (0.5, 0.6666666666666666, 0.5827829453479101), + (0.5, 0.8333333333333335, 0.6822617902381698), + (1, 0.3380952380952381, 0.1666773049170313), + (1, 0.7083333333333333, 0.33366443586989925), + (1, 1.1666666666666667, 0.5023355231537423), + (1, 1.8666666666666665, 0.674616572252164), + (1, 3.560606060606061, 0.844207659503163), + (2.34, 0.7975238095238094, 0.16704903081553285), + (2.34, 1.7133333333333334, 0.3360215931268845), + (2.34, 2.953333333333333, 0.50681909317803), + (2.34, 5.0826666666666656, 0.6755252698800679), + (2.34, 10.869696969696973, 0.8379351104498762), + (56.789, 19.46575238095238, 0.1667020505391409), + (56.789, 42.55008333333333, 0.33353809996933026), + (56.789, 75.552, 0.5003932381177826), + (56.789, 135.76026666666667, 0.6670528221946127), + (56.789, 307.8642424242425, 0.8334999441460798), + ]) + def test_against_reference_values(self, v, x, r): + """The reference values are computed using mpmath as follows. + + from mpmath import mp + mp.dps = 100 + + def iv_ratio_mp(v, x): + return mp.besseli(v, x) / mp.besseli(v - 1, x) + + def _sample(n, *, v): + '''Return n positive real numbers x such that iv_ratio(v, x) are + roughly evenly spaced over (0, 1). The formula is taken from [1]. + + [1] Banerjee A., Dhillon, I. S., Ghosh, J., Sra, S. (2005). + "Clustering on the Unit Hypersphere using von Mises-Fisher + Distributions." Journal of Machine Learning Research, + 6(46):1345-1382. + ''' + r = np.arange(1, n+1) / (n+1) + return r * (2*v-r*r) / (1-r*r) + + for v in (0.5, 1, 2.34, 56.789): + xs = _sample(5, v=v) + for x in xs: + print(f"({v}, {x}, {float(iv_ratio_mp(v,x))}),") + """ + assert_allclose(iv_ratio(v, x), r, rtol=4e-16, atol=0) + + @pytest.mark.parametrize('v,x,r', [ + (1, np.inf, 1), + (np.inf, 1, 0), + ]) + def test_inf(self, v, x, r): + """If exactly one of v or x is inf and the other is within domain, + should return 0 or 1 accordingly.""" + assert_equal(iv_ratio(v, x), r) + + @pytest.mark.parametrize('v', [0.49, -np.inf, np.nan, np.inf]) + @pytest.mark.parametrize('x', [-np.finfo(float).smallest_normal, + -np.finfo(float).smallest_subnormal, + -np.inf, np.nan, np.inf]) + def test_nan(self, v, x): + """If at least one argument is out of domain, or if v = x = inf, + the function should return nan.""" + assert_equal(iv_ratio(v, x), np.nan) + + @pytest.mark.parametrize('v', [0.5, 1, np.finfo(float).max, np.inf]) + def test_zero_x(self, v): + """If x is +/-0.0, return x to ensure iv_ratio is an odd function.""" + assert_equal(iv_ratio(v, 0.0), 0.0) + assert_equal(iv_ratio(v, -0.0), -0.0) + + @pytest.mark.parametrize('v,x', [ + (1, np.finfo(float).smallest_normal), + (1, np.finfo(float).smallest_subnormal), + (1, np.finfo(float).smallest_subnormal*2), + (1e20, 123), + (np.finfo(float).max, 1), + (np.finfo(float).max, np.sqrt(np.finfo(float).max)), + ]) + def test_tiny_x(self, v, x): + """If x is much less than v, the bounds + + x x + --------------------------- <= R <= ----------------------- + v-0.5+sqrt(x**2+(v+0.5)**2) v-1+sqrt(x**2+(v+1)**2) + + collapses to R ~= x/2v. Test against this asymptotic expression. + """ + assert_equal(iv_ratio(v, x), (0.5*x)/v) + + @pytest.mark.parametrize('v,x', [ + (1, 1e16), + (1e20, 1e40), + (np.sqrt(np.finfo(float).max), np.finfo(float).max), + ]) + def test_huge_x(self, v, x): + """If x is much greater than v, the bounds + + x x + --------------------------- <= R <= --------------------------- + v-0.5+sqrt(x**2+(v+0.5)**2) v-0.5+sqrt(x**2+(v-0.5)**2) + + collapses to R ~= 1. Test against this asymptotic expression. + """ + assert_equal(iv_ratio(v, x), 1.0) + + @pytest.mark.parametrize('v,x', [ + (np.finfo(float).max, np.finfo(float).max), + (np.finfo(float).max / 3, np.finfo(float).max), + (np.finfo(float).max, np.finfo(float).max / 3), + ]) + def test_huge_v_x(self, v, x): + """If both x and v are very large, the bounds + + x x + --------------------------- <= R <= ----------------------- + v-0.5+sqrt(x**2+(v+0.5)**2) v-1+sqrt(x**2+(v+1)**2) + + collapses to R ~= x/(v+sqrt(x**2+v**2). Test against this asymptotic + expression, and in particular that no numerical overflow occurs during + intermediate calculations. + """ + t = x / v + expected = t / (1 + np.hypot(1, t)) + assert_allclose(iv_ratio(v, x), expected, rtol=4e-16, atol=0) + + +class TestIvRatioC: + + @pytest.mark.parametrize('v,x,r', [ + (0.5, 0.16666666666666666, 0.8348595870753707), + (0.5, 0.3333333333333333, 0.6784872624683657), + (0.5, 0.5, 0.5378828427399902), + (0.5, 0.6666666666666666, 0.4172170546520899), + (0.5, 0.8333333333333335, 0.3177382097618302), + (1, 0.3380952380952381, 0.8333226950829686), + (1, 0.7083333333333333, 0.6663355641301008), + (1, 1.1666666666666667, 0.4976644768462577), + (1, 1.8666666666666665, 0.325383427747836), + (1, 3.560606060606061, 0.155792340496837), + (2.34, 0.7975238095238094, 0.8329509691844672), + (2.34, 1.7133333333333334, 0.6639784068731155), + (2.34, 2.953333333333333, 0.49318090682197), + (2.34, 5.0826666666666656, 0.3244747301199321), + (2.34, 10.869696969696973, 0.16206488955012377), + (56.789, 19.46575238095238, 0.8332979494608591), + (56.789, 42.55008333333333, 0.6664619000306697), + (56.789, 75.552, 0.4996067618822174), + (56.789, 135.76026666666667, 0.3329471778053873), + (56.789, 307.8642424242425, 0.16650005585392025), + ]) + def test_against_reference_values(self, v, x, r): + """The reference values are one minus those of TestIvRatio.""" + assert_allclose(iv_ratio_c(v, x), r, rtol=1e-15, atol=0) + + @pytest.mark.parametrize('v,x,r', [ + (1, np.inf, 0), + (np.inf, 1, 1), + ]) + def test_inf(self, v, x, r): + """If exactly one of v or x is inf and the other is within domain, + should return 0 or 1 accordingly.""" + assert_equal(iv_ratio_c(v, x), r) + + @pytest.mark.parametrize('v', [0.49, -np.inf, np.nan, np.inf]) + @pytest.mark.parametrize('x', [-np.finfo(float).smallest_normal, + -np.finfo(float).smallest_subnormal, + -np.inf, np.nan, np.inf]) + def test_nan(self, v, x): + """If at least one argument is out of domain, or if v = x = inf, + the function should return nan.""" + assert_equal(iv_ratio_c(v, x), np.nan) + + @pytest.mark.parametrize('v', [0.5, 1, np.finfo(float).max, np.inf]) + def test_zero_x(self, v): + """If x is +/-0.0, return 1.""" + assert_equal(iv_ratio_c(v, 0.0), 1.0) + assert_equal(iv_ratio_c(v, -0.0), 1.0) + + @pytest.mark.parametrize('v,x', [ + (1, np.finfo(float).smallest_normal), + (1, np.finfo(float).smallest_subnormal), + (1, np.finfo(float).smallest_subnormal*2), + (1e20, 123), + (np.finfo(float).max, 1), + (np.finfo(float).max, np.sqrt(np.finfo(float).max)), + ]) + def test_tiny_x(self, v, x): + """If x is much less than v, the bounds + + x x + --------------------------- <= R <= ----------------------- + v-0.5+sqrt(x**2+(v+0.5)**2) v-1+sqrt(x**2+(v+1)**2) + + collapses to 1-R ~= 1-x/2v. Test against this asymptotic expression. + """ + assert_equal(iv_ratio_c(v, x), 1.0-(0.5*x)/v) + + @pytest.mark.parametrize('v,x', [ + (1, 1e16), + (1e20, 1e40), + (np.sqrt(np.finfo(float).max), np.finfo(float).max), + ]) + def test_huge_x(self, v, x): + """If x is much greater than v, the bounds + + x x + --------------------------- <= R <= --------------------------- + v-0.5+sqrt(x**2+(v+0.5)**2) v-0.5+sqrt(x**2+(v-0.5)**2) + + collapses to 1-R ~= (v-0.5)/x. Test against this asymptotic expression. + """ + assert_allclose(iv_ratio_c(v, x), (v-0.5)/x, rtol=1e-15, atol=0) + + @pytest.mark.parametrize('v,x', [ + (np.finfo(float).max, np.finfo(float).max), + (np.finfo(float).max / 3, np.finfo(float).max), + (np.finfo(float).max, np.finfo(float).max / 3), + ]) + def test_huge_v_x(self, v, x): + """If both x and v are very large, the bounds + + x x + --------------------------- <= R <= ----------------------- + v-0.5+sqrt(x**2+(v+0.5)**2) v-1+sqrt(x**2+(v+1)**2) + + collapses to 1 - R ~= 1 - x/(v+sqrt(x**2+v**2). Test against this + asymptotic expression, and in particular that no numerical overflow + occurs during intermediate calculations. + """ + t = x / v + expected = 1 - t / (1 + np.hypot(1, t)) + assert_allclose(iv_ratio_c(v, x), expected, rtol=4e-16, atol=0) diff --git a/v12x/pylib/scipy/special/tests/test_kolmogorov.py b/v12x/pylib/scipy/special/tests/test_kolmogorov.py new file mode 100644 index 0000000000000000000000000000000000000000..ade38115706808809890e1eba3b938562b0af6e6 --- /dev/null +++ b/v12x/pylib/scipy/special/tests/test_kolmogorov.py @@ -0,0 +1,491 @@ +import itertools + +import numpy as np +from numpy.testing import assert_ +from scipy.special._testutils import FuncData + +from scipy.special import kolmogorov, kolmogi, smirnov, smirnovi +from scipy.special._ufuncs import (_kolmogc, _kolmogci, _kolmogp, + _smirnovc, _smirnovci, _smirnovp) + +_rtol = 1e-10 + +class TestSmirnov: + def test_nan(self): + assert_(np.isnan(smirnov(1, np.nan))) + + def test_basic(self): + dataset = [(1, 0.1, 0.9), + (1, 0.875, 0.125), + (2, 0.875, 0.125 * 0.125), + (3, 0.875, 0.125 * 0.125 * 0.125)] + + dataset = np.asarray(dataset) + FuncData( + smirnov, dataset, (0, 1), 2, rtol=_rtol + ).check(dtypes=[int, float, float]) + dataset[:, -1] = 1 - dataset[:, -1] + FuncData( + _smirnovc, dataset, (0, 1), 2, rtol=_rtol + ).check(dtypes=[int, float, float]) + + def test_x_equals_0(self): + dataset = [(n, 0, 1) for n in itertools.chain(range(2, 20), range(1010, 1020))] + dataset = np.asarray(dataset) + FuncData( + smirnov, dataset, (0, 1), 2, rtol=_rtol + ).check(dtypes=[int, float, float]) + dataset[:, -1] = 1 - dataset[:, -1] + FuncData( + _smirnovc, dataset, (0, 1), 2, rtol=_rtol + ).check(dtypes=[int, float, float]) + + def test_x_equals_1(self): + dataset = [(n, 1, 0) for n in itertools.chain(range(2, 20), range(1010, 1020))] + dataset = np.asarray(dataset) + FuncData( + smirnov, dataset, (0, 1), 2, rtol=_rtol + ).check(dtypes=[int, float, float]) + dataset[:, -1] = 1 - dataset[:, -1] + FuncData( + _smirnovc, dataset, (0, 1), 2, rtol=_rtol + ).check(dtypes=[int, float, float]) + + def test_x_equals_0point5(self): + dataset = [(1, 0.5, 0.5), + (2, 0.5, 0.25), + (3, 0.5, 0.166666666667), + (4, 0.5, 0.09375), + (5, 0.5, 0.056), + (6, 0.5, 0.0327932098765), + (7, 0.5, 0.0191958707681), + (8, 0.5, 0.0112953186035), + (9, 0.5, 0.00661933257355), + (10, 0.5, 0.003888705)] + + dataset = np.asarray(dataset) + FuncData( + smirnov, dataset, (0, 1), 2, rtol=_rtol + ).check(dtypes=[int, float, float]) + dataset[:, -1] = 1 - dataset[:, -1] + FuncData( + _smirnovc, dataset, (0, 1), 2, rtol=_rtol + ).check(dtypes=[int, float, float]) + + def test_n_equals_1(self): + x = np.linspace(0, 1, 101, endpoint=True) + dataset = np.column_stack([[1]*len(x), x, 1-x]) + FuncData( + smirnov, dataset, (0, 1), 2, rtol=_rtol + ).check(dtypes=[int, float, float]) + dataset[:, -1] = 1 - dataset[:, -1] + FuncData( + _smirnovc, dataset, (0, 1), 2, rtol=_rtol + ).check(dtypes=[int, float, float]) + + def test_n_equals_2(self): + x = np.linspace(0.5, 1, 101, endpoint=True) + p = np.power(1-x, 2) + n = np.array([2] * len(x)) + dataset = np.column_stack([n, x, p]) + FuncData( + smirnov, dataset, (0, 1), 2, rtol=_rtol + ).check(dtypes=[int, float, float]) + dataset[:, -1] = 1 - dataset[:, -1] + FuncData( + _smirnovc, dataset, (0, 1), 2, rtol=_rtol + ).check(dtypes=[int, float, float]) + + def test_n_equals_3(self): + x = np.linspace(0.7, 1, 31, endpoint=True) + p = np.power(1-x, 3) + n = np.array([3] * len(x)) + dataset = np.column_stack([n, x, p]) + FuncData( + smirnov, dataset, (0, 1), 2, rtol=_rtol + ).check(dtypes=[int, float, float]) + dataset[:, -1] = 1 - dataset[:, -1] + FuncData( + _smirnovc, dataset, (0, 1), 2, rtol=_rtol + ).check(dtypes=[int, float, float]) + + def test_n_large(self): + # test for large values of n + # Probabilities should go down as n goes up + x = 0.4 + pvals = np.array([smirnov(n, x) for n in range(400, 1100, 20)]) + dfs = np.diff(pvals) + assert_(np.all(dfs <= 0), msg=f'Not all diffs negative {dfs}') + + +class TestSmirnovi: + def test_nan(self): + assert_(np.isnan(smirnovi(1, np.nan))) + + def test_basic(self): + dataset = [(1, 0.4, 0.6), + (1, 0.6, 0.4), + (1, 0.99, 0.01), + (1, 0.01, 0.99), + (2, 0.125 * 0.125, 0.875), + (3, 0.125 * 0.125 * 0.125, 0.875), + (10, 1.0 / 16 ** 10, 1 - 1.0 / 16)] + + dataset = np.asarray(dataset) + FuncData( + smirnovi, dataset, (0, 1), 2, rtol=_rtol + ).check(dtypes=[int, float, float]) + dataset[:, 1] = 1 - dataset[:, 1] + FuncData( + _smirnovci, dataset, (0, 1), 2, rtol=_rtol + ).check(dtypes=[int, float, float]) + + def test_x_equals_0(self): + dataset = [(n, 0, 1) for n in itertools.chain(range(2, 20), range(1010, 1020))] + dataset = np.asarray(dataset) + FuncData( + smirnovi, dataset, (0, 1), 2, rtol=_rtol + ).check(dtypes=[int, float, float]) + dataset[:, 1] = 1 - dataset[:, 1] + FuncData( + _smirnovci, dataset, (0, 1), 2, rtol=_rtol + ).check(dtypes=[int, float, float]) + + def test_x_equals_1(self): + dataset = [(n, 1, 0) for n in itertools.chain(range(2, 20), range(1010, 1020))] + dataset = np.asarray(dataset) + FuncData( + smirnovi, dataset, (0, 1), 2, rtol=_rtol + ).check(dtypes=[int, float, float]) + dataset[:, 1] = 1 - dataset[:, 1] + FuncData( + _smirnovci, dataset, (0, 1), 2, rtol=_rtol + ).check(dtypes=[int, float, float]) + + def test_n_equals_1(self): + pp = np.linspace(0, 1, 101, endpoint=True) + # dataset = np.array([(1, p, 1-p) for p in pp]) + dataset = np.column_stack([[1]*len(pp), pp, 1-pp]) + FuncData( + smirnovi, dataset, (0, 1), 2, rtol=_rtol + ).check(dtypes=[int, float, float]) + dataset[:, 1] = 1 - dataset[:, 1] + FuncData( + _smirnovci, dataset, (0, 1), 2, rtol=_rtol + ).check(dtypes=[int, float, float]) + + def test_n_equals_2(self): + x = np.linspace(0.5, 1, 101, endpoint=True) + p = np.power(1-x, 2) + n = np.array([2] * len(x)) + dataset = np.column_stack([n, p, x]) + FuncData( + smirnovi, dataset, (0, 1), 2, rtol=_rtol + ).check(dtypes=[int, float, float]) + dataset[:, 1] = 1 - dataset[:, 1] + FuncData( + _smirnovci, dataset, (0, 1), 2, rtol=_rtol + ).check(dtypes=[int, float, float]) + + def test_n_equals_3(self): + x = np.linspace(0.7, 1, 31, endpoint=True) + p = np.power(1-x, 3) + n = np.array([3] * len(x)) + dataset = np.column_stack([n, p, x]) + FuncData( + smirnovi, dataset, (0, 1), 2, rtol=_rtol + ).check(dtypes=[int, float, float]) + dataset[:, 1] = 1 - dataset[:, 1] + FuncData( + _smirnovci, dataset, (0, 1), 2, rtol=_rtol + ).check(dtypes=[int, float, float]) + + def test_round_trip(self): + def _sm_smi(n, p): + return smirnov(n, smirnovi(n, p)) + + def _smc_smci(n, p): + return _smirnovc(n, _smirnovci(n, p)) + + dataset = [(1, 0.4, 0.4), + (1, 0.6, 0.6), + (2, 0.875, 0.875), + (3, 0.875, 0.875), + (3, 0.125, 0.125), + (10, 0.999, 0.999), + (10, 0.0001, 0.0001)] + + dataset = np.asarray(dataset) + FuncData( + _sm_smi, dataset, (0, 1), 2, rtol=_rtol + ).check(dtypes=[int, float, float]) + FuncData( + _smc_smci, dataset, (0, 1), 2, rtol=_rtol + ).check(dtypes=[int, float, float]) + + def test_x_equals_0point5(self): + dataset = [(1, 0.5, 0.5), + (2, 0.5, 0.366025403784), + (2, 0.25, 0.5), + (3, 0.5, 0.297156508177), + (4, 0.5, 0.255520481121), + (5, 0.5, 0.234559536069), + (6, 0.5, 0.21715965898), + (7, 0.5, 0.202722580034), + (8, 0.5, 0.190621765256), + (9, 0.5, 0.180363501362), + (10, 0.5, 0.17157867006)] + + dataset = np.asarray(dataset) + FuncData( + smirnovi, dataset, (0, 1), 2, rtol=_rtol + ).check(dtypes=[int, float, float]) + dataset[:, 1] = 1 - dataset[:, 1] + FuncData( + _smirnovci, dataset, (0, 1), 2, rtol=_rtol + ).check(dtypes=[int, float, float]) + + +class TestSmirnovp: + def test_nan(self): + assert_(np.isnan(_smirnovp(1, np.nan))) + + def test_basic(self): + # Check derivative at endpoints + n1_10 = np.arange(1, 10) + dataset0 = np.column_stack([n1_10, + np.full_like(n1_10, 0), + np.full_like(n1_10, -1)]) + FuncData( + _smirnovp, dataset0, (0, 1), 2, rtol=_rtol + ).check(dtypes=[int, float, float]) + + n2_10 = np.arange(2, 10) + dataset1 = np.column_stack([n2_10, + np.full_like(n2_10, 1.0), + np.full_like(n2_10, 0)]) + FuncData( + _smirnovp, dataset1, (0, 1), 2, rtol=_rtol + ).check(dtypes=[int, float, float]) + + def test_oneminusoneovern(self): + # Check derivative at x=1-1/n + n = np.arange(1, 20) + x = 1.0/n + xm1 = 1-1.0/n + pp1 = -n * x**(n-1) + pp1 -= (1-np.sign(n-2)**2) * 0.5 # n=2, x=0.5, 1-1/n = 0.5, need to adjust + dataset1 = np.column_stack([n, xm1, pp1]) + FuncData( + _smirnovp, dataset1, (0, 1), 2, rtol=_rtol + ).check(dtypes=[int, float, float]) + + def test_oneovertwon(self): + # Check derivative at x=1/2n (Discontinuous at x=1/n, so check at x=1/2n) + n = np.arange(1, 20) + x = 1.0/2/n + pp = -(n*x+1) * (1+x)**(n-2) + dataset0 = np.column_stack([n, x, pp]) + FuncData( + _smirnovp, dataset0, (0, 1), 2, rtol=_rtol + ).check(dtypes=[int, float, float]) + + def test_oneovern(self): + # Check derivative at x=1/n + # (Discontinuous at x=1/n, hard to tell if x==1/n, only use n=power of 2) + n = 2**np.arange(1, 10) + x = 1.0/n + pp = -(n*x+1) * (1+x)**(n-2) + 0.5 + dataset0 = np.column_stack([n, x, pp]) + FuncData( + _smirnovp, dataset0, (0, 1), 2, rtol=_rtol + ).check(dtypes=[int, float, float]) + + def test_oneovernclose(self): + # Check derivative at x=1/n + # (Discontinuous at x=1/n, test on either side: x=1/n +/- 2epsilon) + n = np.arange(3, 20) + + x = 1.0/n - 2*np.finfo(float).eps + pp = -(n*x+1) * (1+x)**(n-2) + dataset0 = np.column_stack([n, x, pp]) + FuncData( + _smirnovp, dataset0, (0, 1), 2, rtol=_rtol + ).check(dtypes=[int, float, float]) + + x = 1.0/n + 2*np.finfo(float).eps + pp = -(n*x+1) * (1+x)**(n-2) + 1 + dataset1 = np.column_stack([n, x, pp]) + FuncData( + _smirnovp, dataset1, (0, 1), 2, rtol=_rtol + ).check(dtypes=[int, float, float]) + + +class TestKolmogorov: + def test_nan(self): + assert_(np.isnan(kolmogorov(np.nan))) + + def test_basic(self): + dataset = [(0, 1.0), + (0.5, 0.96394524366487511), + (0.8275735551899077, 0.5000000000000000), + (1, 0.26999967167735456), + (2, 0.00067092525577969533)] + + dataset = np.asarray(dataset) + FuncData(kolmogorov, dataset, (0,), 1, rtol=_rtol).check() + + def test_linspace(self): + x = np.linspace(0, 2.0, 21) + dataset = [1.0000000000000000, 1.0000000000000000, 0.9999999999994950, + 0.9999906941986655, 0.9971923267772983, 0.9639452436648751, + 0.8642827790506042, 0.7112351950296890, 0.5441424115741981, + 0.3927307079406543, 0.2699996716773546, 0.1777181926064012, + 0.1122496666707249, 0.0680922218447664, 0.0396818795381144, + 0.0222179626165251, 0.0119520432391966, 0.0061774306344441, + 0.0030676213475797, 0.0014636048371873, 0.0006709252557797] + + dataset_c = [0.0000000000000000, 6.609305242245699e-53, 5.050407338670114e-13, + 9.305801334566668e-06, 0.0028076732227017, 0.0360547563351249, + 0.1357172209493958, 0.2887648049703110, 0.4558575884258019, + 0.6072692920593457, 0.7300003283226455, 0.8222818073935988, + 0.8877503333292751, 0.9319077781552336, 0.9603181204618857, + 0.9777820373834749, 0.9880479567608034, 0.9938225693655559, + 0.9969323786524203, 0.9985363951628127, 0.9993290747442203] + + dataset = np.column_stack([x, dataset]) + FuncData(kolmogorov, dataset, (0,), 1, rtol=_rtol).check() + dataset_c = np.column_stack([x, dataset_c]) + FuncData(_kolmogc, dataset_c, (0,), 1, rtol=_rtol).check() + + def test_linspacei(self): + p = np.linspace(0, 1.0, 21, endpoint=True) + dataset = [np.inf, 1.3580986393225507, 1.2238478702170823, + 1.1379465424937751, 1.0727491749396481, 1.0191847202536859, + 0.9730633753323726, 0.9320695842357622, 0.8947644549851197, + 0.8601710725555463, 0.8275735551899077, 0.7964065373291559, + 0.7661855555617682, 0.7364542888171910, 0.7067326523068980, + 0.6764476915028201, 0.6448126061663567, 0.6105590999244391, + 0.5711732651063401, 0.5196103791686224, 0.0000000000000000] + + dataset_c = [0.0000000000000000, 0.5196103791686225, 0.5711732651063401, + 0.6105590999244391, 0.6448126061663567, 0.6764476915028201, + 0.7067326523068980, 0.7364542888171910, 0.7661855555617682, + 0.7964065373291559, 0.8275735551899077, 0.8601710725555463, + 0.8947644549851196, 0.9320695842357622, 0.9730633753323727, + 1.0191847202536859, 1.0727491749396481, 1.1379465424937754, + 1.2238478702170825, 1.3580986393225509, np.inf] + + dataset = np.column_stack([p[1:], dataset[1:]]) + FuncData(kolmogi, dataset, (0,), 1, rtol=_rtol).check() + dataset_c = np.column_stack([p[:-1], dataset_c[:-1]]) + FuncData(_kolmogci, dataset_c, (0,), 1, rtol=_rtol).check() + + def test_smallx(self): + epsilon = 0.1 ** np.arange(1, 14) + x = np.array([0.571173265106, 0.441027698518, 0.374219690278, 0.331392659217, + 0.300820537459, 0.277539353999, 0.259023494805, 0.243829561254, + 0.231063086389, 0.220135543236, 0.210641372041, 0.202290283658, + 0.19487060742]) + + dataset = np.column_stack([x, 1-epsilon]) + FuncData(kolmogorov, dataset, (0,), 1, rtol=_rtol).check() + + def test_round_trip(self): + def _ki_k(_x): + return kolmogi(kolmogorov(_x)) + + def _kci_kc(_x): + return _kolmogci(_kolmogc(_x)) + + x = np.linspace(0.0, 2.0, 21, endpoint=True) + # Exclude 0.1, 0.2. 0.2 almost makes succeeds, but 0.1 has no chance. + x02 = x[(x == 0) | (x > 0.21)] + dataset02 = np.column_stack([x02, x02]) + FuncData(_ki_k, dataset02, (0,), 1, rtol=_rtol).check() + + dataset = np.column_stack([x, x]) + FuncData(_kci_kc, dataset, (0,), 1, rtol=_rtol).check() + + +class TestKolmogi: + def test_nan(self): + assert_(np.isnan(kolmogi(np.nan))) + + def test_basic(self): + dataset = [(1.0, 0), + (0.96394524366487511, 0.5), + (0.9, 0.571173265106), + (0.5000000000000000, 0.8275735551899077), + (0.26999967167735456, 1), + (0.00067092525577969533, 2)] + + dataset = np.asarray(dataset) + FuncData(kolmogi, dataset, (0,), 1, rtol=_rtol).check() + + def test_smallpcdf(self): + epsilon = 0.5 ** np.arange(1, 55, 3) + # kolmogi(1-p) == _kolmogci(p) if 1-(1-p) == p, but not necessarily otherwise + # Use epsilon s.t. 1-(1-epsilon)) == epsilon, + # so can use same x-array for both results + + x = np.array([0.8275735551899077, 0.5345255069097583, 0.4320114038786941, + 0.3736868442620478, 0.3345161714909591, 0.3057833329315859, + 0.2835052890528936, 0.2655578150208676, 0.2506869966107999, + 0.2380971058736669, 0.2272549289962079, 0.2177876361600040, + 0.2094254686862041, 0.2019676748836232, 0.1952612948137504, + 0.1891874239646641, 0.1836520225050326, 0.1785795904846466]) + + dataset = np.column_stack([1-epsilon, x]) + FuncData(kolmogi, dataset, (0,), 1, rtol=_rtol).check() + + dataset = np.column_stack([epsilon, x]) + FuncData(_kolmogci, dataset, (0,), 1, rtol=_rtol).check() + + def test_smallpsf(self): + epsilon = 0.5 ** np.arange(1, 55, 3) + # kolmogi(p) == _kolmogci(1-p) if 1-(1-p) == p, but not necessarily otherwise + # Use epsilon s.t. 1-(1-epsilon)) == epsilon, + # so can use same x-array for both results + + x = np.array([0.8275735551899077, 1.3163786275161036, 1.6651092133663343, + 1.9525136345289607, 2.2027324540033235, 2.4272929437460848, + 2.6327688477341593, 2.8233300509220260, 3.0018183401530627, + 3.1702735084088891, 3.3302184446307912, 3.4828258153113318, + 3.6290214150152051, 3.7695513262825959, 3.9050272690877326, + 4.0359582187082550, 4.1627730557884890, 4.2858371743264527]) + + dataset = np.column_stack([epsilon, x]) + FuncData(kolmogi, dataset, (0,), 1, rtol=_rtol).check() + + dataset = np.column_stack([1-epsilon, x]) + FuncData(_kolmogci, dataset, (0,), 1, rtol=_rtol).check() + + def test_round_trip(self): + def _k_ki(_p): + return kolmogorov(kolmogi(_p)) + + p = np.linspace(0.1, 1.0, 10, endpoint=True) + dataset = np.column_stack([p, p]) + FuncData(_k_ki, dataset, (0,), 1, rtol=_rtol).check() + + +class TestKolmogp: + def test_nan(self): + assert_(np.isnan(_kolmogp(np.nan))) + + def test_basic(self): + dataset = [(0.000000, -0.0), + (0.200000, -1.532420541338916e-10), + (0.400000, -0.1012254419260496), + (0.600000, -1.324123244249925), + (0.800000, -1.627024345636592), + (1.000000, -1.071948558356941), + (1.200000, -0.538512430720529), + (1.400000, -0.2222133182429472), + (1.600000, -0.07649302775520538), + (1.800000, -0.02208687346347873), + (2.000000, -0.005367402045629683)] + + dataset = np.asarray(dataset) + FuncData(_kolmogp, dataset, (0,), 1, rtol=_rtol).check() diff --git a/v12x/pylib/scipy/special/tests/test_lambertw.py b/v12x/pylib/scipy/special/tests/test_lambertw.py new file mode 100644 index 0000000000000000000000000000000000000000..ac26dfb9ff65dbfd94740220a95a16c902b4be3d --- /dev/null +++ b/v12x/pylib/scipy/special/tests/test_lambertw.py @@ -0,0 +1,109 @@ +# +# Tests for the lambertw function, +# Adapted from the MPMath tests [1] by Yosef Meller, mellerf@netvision.net.il +# Distributed under the same license as SciPy itself. +# +# [1] mpmath source code, Subversion revision 992 +# http://code.google.com/p/mpmath/source/browse/trunk/mpmath/tests/test_functions2.py?spec=svn994&r=992 + +import pytest +import numpy as np +from numpy.testing import assert_, assert_equal, assert_allclose +from scipy.special import lambertw +from numpy import nan, inf, pi, e, isnan, log, r_, array, complex128 + +from scipy.special._testutils import FuncData + + +def test_values(): + assert_(isnan(lambertw(nan))) + assert_equal(lambertw(inf,1).real, inf) + assert_equal(lambertw(inf,1).imag, 2*pi) + assert_equal(lambertw(-inf,1).real, inf) + assert_equal(lambertw(-inf,1).imag, 3*pi) + + assert_equal(lambertw(1.), lambertw(1., 0)) + + data = [ + (0,0, 0), + (0+0j,0, 0), + (inf,0, inf), + (0,-1, -inf), + (0,1, -inf), + (0,3, -inf), + (e,0, 1), + (1,0, 0.567143290409783873), + (-pi/2,0, 1j*pi/2), + (-log(2)/2,0, -log(2)), + (0.25,0, 0.203888354702240164), + (-0.25,0, -0.357402956181388903), + (-1./10000,0, -0.000100010001500266719), + (-0.25,-1, -2.15329236411034965), + (0.25,-1, -3.00899800997004620-4.07652978899159763j), + (-0.25,-1, -2.15329236411034965), + (0.25,1, -3.00899800997004620+4.07652978899159763j), + (-0.25,1, -3.48973228422959210+7.41405453009603664j), + (-4,0, 0.67881197132094523+1.91195078174339937j), + (-4,1, -0.66743107129800988+7.76827456802783084j), + (-4,-1, 0.67881197132094523-1.91195078174339937j), + (1000,0, 5.24960285240159623), + (1000,1, 4.91492239981054535+5.44652615979447070j), + (1000,-1, 4.91492239981054535-5.44652615979447070j), + (1000,5, 3.5010625305312892+29.9614548941181328j), + (3+4j,0, 1.281561806123775878+0.533095222020971071j), + (-0.4+0.4j,0, -0.10396515323290657+0.61899273315171632j), + (3+4j,1, -0.11691092896595324+5.61888039871282334j), + (3+4j,-1, 0.25856740686699742-3.85211668616143559j), + (-0.5,-1, -0.794023632344689368-0.770111750510379110j), + (-1./10000,1, -11.82350837248724344+6.80546081842002101j), + (-1./10000,-1, -11.6671145325663544), + (-1./10000,-2, -11.82350837248724344-6.80546081842002101j), + (-1./100000,4, -14.9186890769540539+26.1856750178782046j), + (-1./100000,5, -15.0931437726379218666+32.5525721210262290086j), + ((2+1j)/10,0, 0.173704503762911669+0.071781336752835511j), + ((2+1j)/10,1, -3.21746028349820063+4.56175438896292539j), + ((2+1j)/10,-1, -3.03781405002993088-3.53946629633505737j), + ((2+1j)/10,4, -4.6878509692773249+23.8313630697683291j), + (-(2+1j)/10,0, -0.226933772515757933-0.164986470020154580j), + (-(2+1j)/10,1, -2.43569517046110001+0.76974067544756289j), + (-(2+1j)/10,-1, -3.54858738151989450-6.91627921869943589j), + (-(2+1j)/10,4, -4.5500846928118151+20.6672982215434637j), + (pi,0, 1.073658194796149172092178407024821347547745350410314531), + + # Former bug in generated branch, + (-0.5+0.002j,0, -0.78917138132659918344 + 0.76743539379990327749j), + (-0.5-0.002j,0, -0.78917138132659918344 - 0.76743539379990327749j), + (-0.448+0.4j,0, -0.11855133765652382241 + 0.66570534313583423116j), + (-0.448-0.4j,0, -0.11855133765652382241 - 0.66570534313583423116j), + ] + data = array(data, dtype=complex128) + + def w(x, y): + return lambertw(x, y.real.astype(int)) + with np.errstate(all='ignore'): + FuncData(w, data, (0,1), 2, rtol=1e-10, atol=1e-13).check() + + +def test_ufunc(): + assert_allclose(lambertw(r_[0., e, 1.]), r_[0., 1., 0.567143290409783873], + atol=1.5e-7, rtol=0) + + +def test_lambertw_ufunc_loop_selection(): + # see https://github.com/scipy/scipy/issues/4895 + dt = np.dtype(np.complex128) + assert_equal(lambertw(0, 0, 0).dtype, dt) + assert_equal(lambertw([0], 0, 0).dtype, dt) + assert_equal(lambertw(0, [0], 0).dtype, dt) + assert_equal(lambertw(0, 0, [0]).dtype, dt) + assert_equal(lambertw([0], [0], [0]).dtype, dt) + + +@pytest.mark.parametrize('z', [1e-316, -2e-320j, -5e-318+1e-320j]) +def test_lambertw_subnormal_k0(z): + # Verify that subnormal inputs are handled correctly on + # the branch k=0 (regression test for gh-16291). + w = lambertw(z) + # For values this small, we can be sure that numerically, + # lambertw(z) is z. + assert w == z diff --git a/v12x/pylib/scipy/special/tests/test_legendre.py b/v12x/pylib/scipy/special/tests/test_legendre.py new file mode 100644 index 0000000000000000000000000000000000000000..7625313ef258de23307c6c4a432069f5c4cd0cd1 --- /dev/null +++ b/v12x/pylib/scipy/special/tests/test_legendre.py @@ -0,0 +1,1421 @@ +import math +import warnings + +import numpy as np + +import pytest +from numpy.testing import assert_equal, assert_allclose + +from scipy import special +from scipy.special import (legendre_p, legendre_p_all, assoc_legendre_p, + assoc_legendre_p_all, sph_legendre_p, sph_legendre_p_all) + + +# Base polynomials come from Abrahmowitz and Stegan +class TestLegendre: + def test_legendre(self): + leg0 = special.legendre(0) + leg1 = special.legendre(1) + leg2 = special.legendre(2) + leg3 = special.legendre(3) + leg4 = special.legendre(4) + leg5 = special.legendre(5) + assert_equal(leg0.c, [1]) + assert_equal(leg1.c, [1,0]) + assert_allclose(leg2.c, np.array([3, 0, -1])/2.0, + atol=1.5e-13, rtol=0) + assert_allclose(leg3.c, np.array([5, 0, -3, 0])/2.0, + atol=1.5e-7, rtol=0) + assert_allclose(leg4.c, np.array([35, 0, -30, 0, 3])/8.0, + atol=1.5e-7, rtol=0) + assert_allclose(leg5.c, np.array([63, 0, -70, 0, 15, 0])/8.0, + atol=1.5e-7, rtol=0) + +class TestLegendreP: + @pytest.mark.parametrize("shape", [(10,), (4, 9), (3, 5, 7)]) + def test_ode(self, shape): + rng = np.random.default_rng(1234) + + n = rng.integers(0, 100, shape) + x = rng.uniform(-1, 1, shape) + + p, p_jac, p_hess = legendre_p(n, x, diff_n=2) + + assert p.shape == shape + assert p_jac.shape == p.shape + assert p_hess.shape == p_jac.shape + + err = (1 - x * x) * p_hess - 2 * x * p_jac + n * (n + 1) * p + np.testing.assert_allclose(err, 0, atol=1e-10) + + @pytest.mark.parametrize("n_max", [1, 2, 4, 8, 16, 32]) + @pytest.mark.parametrize("x_shape", [(10,), (4, 9), (3, 5, 7)]) + def test_all_ode(self, n_max, x_shape): + rng = np.random.default_rng(1234) + + x = rng.uniform(-1, 1, x_shape) + p, p_jac, p_hess = legendre_p_all(n_max, x, diff_n=2) + + n = np.arange(n_max + 1) + n = np.expand_dims(n, axis = tuple(range(1, x.ndim + 1))) + + assert p.shape == (len(n),) + x.shape + assert p_jac.shape == p.shape + assert p_hess.shape == p_jac.shape + + err = (1 - x * x) * p_hess - 2 * x * p_jac + n * (n + 1) * p + np.testing.assert_allclose(err, 0, atol=1e-10) + + +class TestAssocLegendreP: + def test_assoc_legendre_gh23101(self): + z = np.array([-1, -.5, 0, .5, 1]) + expected = assoc_legendre_p_1_0(z) + result = assoc_legendre_p(1, 0, z) + assert_allclose(np.squeeze(result), expected) + + expected = assoc_legendre_p_3_0(z) + result = assoc_legendre_p(3, 0, z) + assert_allclose(np.squeeze(result), expected) + + @pytest.mark.parametrize("shape", [(10,), (4, 9), (3, 5, 7, 10)]) + @pytest.mark.parametrize("m_max", [5, 4]) + @pytest.mark.parametrize("n_max", [7, 10]) + def test_lpmn(self, shape, n_max, m_max): + rng = np.random.default_rng(1234) + + x = rng.uniform(-0.99, 0.99, shape) + p_all, p_all_jac, p_all_hess = \ + assoc_legendre_p_all(n_max, m_max, x, diff_n=2) + + n = np.arange(n_max + 1) + n = np.expand_dims(n, axis = tuple(range(1, x.ndim + 2))) + + m = np.concatenate([np.arange(m_max + 1), np.arange(-m_max, 0)]) + m = np.expand_dims(m, axis = (0,) + tuple(range(2, x.ndim + 2))) + + x = np.expand_dims(x, axis = (0, 1)) + p, p_jac, p_hess = assoc_legendre_p(n, m, x, diff_n=2) + + np.testing.assert_allclose(p, p_all) + np.testing.assert_allclose(p_jac, p_all_jac) + np.testing.assert_allclose(p_hess, p_all_hess) + + @pytest.mark.parametrize("shape", [(10,), (4, 9), (3, 5, 7, 10)]) + @pytest.mark.parametrize("norm", [True, False]) + def test_ode(self, shape, norm): + rng = np.random.default_rng(1234) + + n = rng.integers(0, 10, shape) + m = rng.integers(-10, 10, shape) + x = rng.uniform(-1, 1, shape) + + p, p_jac, p_hess = assoc_legendre_p(n, m, x, norm=norm, diff_n=2) + + assert p.shape == shape + assert p_jac.shape == p.shape + assert p_hess.shape == p_jac.shape + + np.testing.assert_allclose((1 - x * x) * p_hess, + 2 * x * p_jac - (n * (n + 1) - m * m / (1 - x * x)) * p, + rtol=1e-05, atol=1e-08) + + @pytest.mark.parametrize("shape", [(10,), (4, 9), (3, 5, 7)]) + def test_all(self, shape): + rng = np.random.default_rng(1234) + + n_max = 20 + m_max = 20 + + x = rng.uniform(-0.99, 0.99, shape) + + p, p_jac, p_hess = assoc_legendre_p_all(n_max, m_max, x, diff_n=2) + + m = np.concatenate([np.arange(m_max + 1), np.arange(-m_max, 0)]) + n = np.arange(n_max + 1) + + n = np.expand_dims(n, axis = tuple(range(1, x.ndim + 2))) + m = np.expand_dims(m, axis = (0,) + tuple(range(2, x.ndim + 2))) + np.testing.assert_allclose((1 - x * x) * p_hess, + 2 * x * p_jac - (n * (n + 1) - m * m / (1 - x * x)) * p, + rtol=1e-05, atol=1e-08) + + @pytest.mark.parametrize("shape", [(10,), (4, 9), (3, 5, 7)]) + @pytest.mark.parametrize("norm", [True, False]) + def test_specific(self, shape, norm): + rng = np.random.default_rng(1234) + + x = rng.uniform(-0.99, 0.99, shape) + + p, p_jac = assoc_legendre_p_all(4, 4, x, norm=norm, diff_n=1) + + np.testing.assert_allclose(p[0, 0], + assoc_legendre_p_0_0(x, norm=norm)) + np.testing.assert_allclose(p[0, 1], 0) + np.testing.assert_allclose(p[0, 2], 0) + np.testing.assert_allclose(p[0, 3], 0) + np.testing.assert_allclose(p[0, 4], 0) + np.testing.assert_allclose(p[0, -3], 0) + np.testing.assert_allclose(p[0, -2], 0) + np.testing.assert_allclose(p[0, -1], 0) + + np.testing.assert_allclose(p[1, 0], + assoc_legendre_p_1_0(x, norm=norm)) + np.testing.assert_allclose(p[1, 1], + assoc_legendre_p_1_1(x, norm=norm)) + np.testing.assert_allclose(p[1, 2], 0) + np.testing.assert_allclose(p[1, 3], 0) + np.testing.assert_allclose(p[1, 4], 0) + np.testing.assert_allclose(p[1, -4], 0) + np.testing.assert_allclose(p[1, -3], 0) + np.testing.assert_allclose(p[1, -2], 0) + np.testing.assert_allclose(p[1, -1], + assoc_legendre_p_1_m1(x, norm=norm)) + + np.testing.assert_allclose(p[2, 0], + assoc_legendre_p_2_0(x, norm=norm)) + np.testing.assert_allclose(p[2, 1], + assoc_legendre_p_2_1(x, norm=norm)) + np.testing.assert_allclose(p[2, 2], + assoc_legendre_p_2_2(x, norm=norm)) + np.testing.assert_allclose(p[2, 3], 0) + np.testing.assert_allclose(p[2, 4], 0) + np.testing.assert_allclose(p[2, -4], 0) + np.testing.assert_allclose(p[2, -3], 0) + np.testing.assert_allclose(p[2, -2], + assoc_legendre_p_2_m2(x, norm=norm)) + np.testing.assert_allclose(p[2, -1], + assoc_legendre_p_2_m1(x, norm=norm)) + + np.testing.assert_allclose(p[3, 0], + assoc_legendre_p_3_0(x, norm=norm)) + np.testing.assert_allclose(p[3, 1], + assoc_legendre_p_3_1(x, norm=norm)) + np.testing.assert_allclose(p[3, 2], + assoc_legendre_p_3_2(x, norm=norm)) + np.testing.assert_allclose(p[3, 3], + assoc_legendre_p_3_3(x, norm=norm)) + np.testing.assert_allclose(p[3, 4], 0) + np.testing.assert_allclose(p[3, -4], 0) + np.testing.assert_allclose(p[3, -3], + assoc_legendre_p_3_m3(x, norm=norm)) + np.testing.assert_allclose(p[3, -2], + assoc_legendre_p_3_m2(x, norm=norm)) + np.testing.assert_allclose(p[3, -1], + assoc_legendre_p_3_m1(x, norm=norm)) + + np.testing.assert_allclose(p[4, 0], + assoc_legendre_p_4_0(x, norm=norm)) + np.testing.assert_allclose(p[4, 1], + assoc_legendre_p_4_1(x, norm=norm)) + np.testing.assert_allclose(p[4, 2], + assoc_legendre_p_4_2(x, norm=norm)) + np.testing.assert_allclose(p[4, 3], + assoc_legendre_p_4_3(x, norm=norm)) + np.testing.assert_allclose(p[4, 4], + assoc_legendre_p_4_4(x, norm=norm)) + np.testing.assert_allclose(p[4, -4], + assoc_legendre_p_4_m4(x, norm=norm)) + np.testing.assert_allclose(p[4, -3], + assoc_legendre_p_4_m3(x, norm=norm)) + np.testing.assert_allclose(p[4, -2], + assoc_legendre_p_4_m2(x, norm=norm)) + np.testing.assert_allclose(p[4, -1], + assoc_legendre_p_4_m1(x, norm=norm)) + + np.testing.assert_allclose(p_jac[0, 0], + assoc_legendre_p_0_0_jac(x, norm=norm)) + np.testing.assert_allclose(p_jac[0, 1], 0) + np.testing.assert_allclose(p_jac[0, 2], 0) + np.testing.assert_allclose(p_jac[0, 3], 0) + np.testing.assert_allclose(p_jac[0, 4], 0) + np.testing.assert_allclose(p_jac[0, -4], 0) + np.testing.assert_allclose(p_jac[0, -3], 0) + np.testing.assert_allclose(p_jac[0, -2], 0) + np.testing.assert_allclose(p_jac[0, -1], 0) + + np.testing.assert_allclose(p_jac[1, 0], + assoc_legendre_p_1_0_jac(x, norm=norm)) + np.testing.assert_allclose(p_jac[1, 1], + assoc_legendre_p_1_1_jac(x, norm=norm)) + np.testing.assert_allclose(p_jac[1, 2], 0) + np.testing.assert_allclose(p_jac[1, 3], 0) + np.testing.assert_allclose(p_jac[1, 4], 0) + np.testing.assert_allclose(p_jac[1, -4], 0) + np.testing.assert_allclose(p_jac[1, -3], 0) + np.testing.assert_allclose(p_jac[1, -2], 0) + np.testing.assert_allclose(p_jac[1, -1], + assoc_legendre_p_1_m1_jac(x, norm=norm)) + + np.testing.assert_allclose(p_jac[2, 0], + assoc_legendre_p_2_0_jac(x, norm=norm)) + np.testing.assert_allclose(p_jac[2, 1], + assoc_legendre_p_2_1_jac(x, norm=norm)) + np.testing.assert_allclose(p_jac[2, 2], + assoc_legendre_p_2_2_jac(x, norm=norm)) + np.testing.assert_allclose(p_jac[2, 3], 0) + np.testing.assert_allclose(p_jac[2, 4], 0) + np.testing.assert_allclose(p_jac[2, -4], 0) + np.testing.assert_allclose(p_jac[2, -3], 0) + np.testing.assert_allclose(p_jac[2, -2], + assoc_legendre_p_2_m2_jac(x, norm=norm)) + np.testing.assert_allclose(p_jac[2, -1], + assoc_legendre_p_2_m1_jac(x, norm=norm)) + + np.testing.assert_allclose(p_jac[3, 0], + assoc_legendre_p_3_0_jac(x, norm=norm)) + np.testing.assert_allclose(p_jac[3, 1], + assoc_legendre_p_3_1_jac(x, norm=norm)) + np.testing.assert_allclose(p_jac[3, 2], + assoc_legendre_p_3_2_jac(x, norm=norm)) + np.testing.assert_allclose(p_jac[3, 3], + assoc_legendre_p_3_3_jac(x, norm=norm)) + np.testing.assert_allclose(p_jac[3, 4], 0) + np.testing.assert_allclose(p_jac[3, -4], 0) + np.testing.assert_allclose(p_jac[3, -3], + assoc_legendre_p_3_m3_jac(x, norm=norm)) + np.testing.assert_allclose(p_jac[3, -2], + assoc_legendre_p_3_m2_jac(x, norm=norm)) + np.testing.assert_allclose(p_jac[3, -1], + assoc_legendre_p_3_m1_jac(x, norm=norm)) + + np.testing.assert_allclose(p_jac[4, 0], + assoc_legendre_p_4_0_jac(x, norm=norm)) + np.testing.assert_allclose(p_jac[4, 1], + assoc_legendre_p_4_1_jac(x, norm=norm)) + np.testing.assert_allclose(p_jac[4, 2], + assoc_legendre_p_4_2_jac(x, norm=norm)) + np.testing.assert_allclose(p_jac[4, 3], + assoc_legendre_p_4_3_jac(x, norm=norm)) + np.testing.assert_allclose(p_jac[4, 4], + assoc_legendre_p_4_4_jac(x, norm=norm)) + np.testing.assert_allclose(p_jac[4, -4], + assoc_legendre_p_4_m4_jac(x, norm=norm)) + np.testing.assert_allclose(p_jac[4, -3], + assoc_legendre_p_4_m3_jac(x, norm=norm)) + np.testing.assert_allclose(p_jac[4, -2], + assoc_legendre_p_4_m2_jac(x, norm=norm)) + np.testing.assert_allclose(p_jac[4, -1], + assoc_legendre_p_4_m1_jac(x, norm=norm)) + + @pytest.mark.parametrize("m_max", [7]) + @pytest.mark.parametrize("n_max", [10]) + @pytest.mark.parametrize("x", [1, -1]) + def test_all_limits(self, m_max, n_max, x): + p, p_jac = assoc_legendre_p_all(n_max, m_max, x, diff_n=1) + + n = np.arange(n_max + 1) + + np.testing.assert_allclose(p_jac[:, 0], + pow(x, n + 1) * n * (n + 1) / 2) + np.testing.assert_allclose(p_jac[:, 1], + np.where(n >= 1, pow(x, n) * np.inf, 0)) + np.testing.assert_allclose(p_jac[:, 2], + np.where(n >= 2, -pow(x, n + 1) * (n + 2) * (n + 1) * n * (n - 1) / 4, 0)) + np.testing.assert_allclose(p_jac[:, -2], + np.where(n >= 2, -pow(x, n + 1) / 4, 0)) + np.testing.assert_allclose(p_jac[:, -1], + np.where(n >= 1, -pow(x, n) * np.inf, 0)) + + for m in range(3, m_max + 1): + np.testing.assert_allclose(p_jac[:, m], 0) + np.testing.assert_allclose(p_jac[:, -m], 0) + + @pytest.mark.parametrize("n", np.arange(0, 20)) + def test_assoc_legendre_p_norm_m0(self, n): + # gh 24099 + m = 0 + z = np.linspace(-1, 1, 101) + p = assoc_legendre_p(n, m, z, norm=False) + p_norm = assoc_legendre_p(n, m, z, norm=True) + factor = np.sqrt((2 * n + 1) / 2) + np.testing.assert_allclose(p_norm, factor * p) + + @pytest.mark.parametrize("n", np.arange(0, 20)) + def test_assoc_legendre_p_all_norm_m0(self, n): + # gh 24099 + m = 0 + z = np.linspace(-1, 1, 101) + p_all = assoc_legendre_p_all(n, m, z, norm=False)[0, : n + 1, 0, :] + p_norm_all = assoc_legendre_p_all(n, m, z, norm=True)[0, : n + 1, 0, :] + factor = np.sqrt((2 * np.arange(n + 1) + 1) / 2) + factor = factor[:, np.newaxis] + np.testing.assert_allclose(p_norm_all, factor * p_all) + + +class TestMultiAssocLegendreP: + @pytest.mark.parametrize("shape", [(1000,), (4, 9), (3, 5, 7)]) + @pytest.mark.parametrize("branch_cut", [2, 3]) + @pytest.mark.parametrize("z_min, z_max", [(-10 - 10j, 10 + 10j), + (-1, 1), (-10j, 10j)]) + @pytest.mark.parametrize("norm", [True, False]) + def test_specific(self, shape, branch_cut, z_min, z_max, norm): + rng = np.random.default_rng(1234) + + z = rng.uniform(z_min.real, z_max.real, shape) + \ + 1j * rng.uniform(z_min.imag, z_max.imag, shape) + + p, p_jac = assoc_legendre_p_all(4, 4, + z, branch_cut=branch_cut, norm=norm, diff_n=1) + + np.testing.assert_allclose(p[0, 0], + assoc_legendre_p_0_0(z, branch_cut=branch_cut, norm=norm)) + np.testing.assert_allclose(p[0, 1], 0) + np.testing.assert_allclose(p[0, 2], 0) + np.testing.assert_allclose(p[0, 3], 0) + np.testing.assert_allclose(p[0, 4], 0) + np.testing.assert_allclose(p[0, -4], 0) + np.testing.assert_allclose(p[0, -3], 0) + np.testing.assert_allclose(p[0, -2], 0) + np.testing.assert_allclose(p[0, -1], 0) + + np.testing.assert_allclose(p[1, 0], + assoc_legendre_p_1_0(z, branch_cut=branch_cut, norm=norm)) + np.testing.assert_allclose(p[1, 1], + assoc_legendre_p_1_1(z, branch_cut=branch_cut, norm=norm)) + np.testing.assert_allclose(p[1, 2], 0) + np.testing.assert_allclose(p[1, 3], 0) + np.testing.assert_allclose(p[1, 4], 0) + np.testing.assert_allclose(p[1, -4], 0) + np.testing.assert_allclose(p[1, -3], 0) + np.testing.assert_allclose(p[1, -2], 0) + np.testing.assert_allclose(p[1, -1], + assoc_legendre_p_1_m1(z, branch_cut=branch_cut, norm=norm)) + + np.testing.assert_allclose(p[2, 0], + assoc_legendre_p_2_0(z, branch_cut=branch_cut, norm=norm)) + np.testing.assert_allclose(p[2, 1], + assoc_legendre_p_2_1(z, branch_cut=branch_cut, norm=norm)) + np.testing.assert_allclose(p[2, 2], + assoc_legendre_p_2_2(z, branch_cut=branch_cut, norm=norm)) + np.testing.assert_allclose(p[2, 3], 0) + np.testing.assert_allclose(p[2, 4], 0) + np.testing.assert_allclose(p[2, -4], 0) + np.testing.assert_allclose(p[2, -3], 0) + np.testing.assert_allclose(p[2, -2], + assoc_legendre_p_2_m2(z, branch_cut=branch_cut, norm=norm)) + np.testing.assert_allclose(p[2, -1], + assoc_legendre_p_2_m1(z, branch_cut=branch_cut, norm=norm)) + + np.testing.assert_allclose(p[3, 0], + assoc_legendre_p_3_0(z, branch_cut=branch_cut, norm=norm)) + np.testing.assert_allclose(p[3, 1], + assoc_legendre_p_3_1(z, branch_cut=branch_cut, norm=norm)) + np.testing.assert_allclose(p[3, 2], + assoc_legendre_p_3_2(z, branch_cut=branch_cut, norm=norm)) + np.testing.assert_allclose(p[3, 3], + assoc_legendre_p_3_3(z, branch_cut=branch_cut, norm=norm)) + np.testing.assert_allclose(p[3, 4], 0) + np.testing.assert_allclose(p[3, -4], 0) + np.testing.assert_allclose(p[3, -3], + assoc_legendre_p_3_m3(z, branch_cut=branch_cut, norm=norm)) + np.testing.assert_allclose(p[3, -2], + assoc_legendre_p_3_m2(z, branch_cut=branch_cut, norm=norm)) + np.testing.assert_allclose(p[3, -1], + assoc_legendre_p_3_m1(z, branch_cut=branch_cut, norm=norm)) + + np.testing.assert_allclose(p[4, 0], + assoc_legendre_p_4_0(z, branch_cut=branch_cut, norm=norm)) + np.testing.assert_allclose(p[4, 1], + assoc_legendre_p_4_1(z, branch_cut=branch_cut, norm=norm)) + np.testing.assert_allclose(p[4, 2], + assoc_legendre_p_4_2(z, branch_cut=branch_cut, norm=norm)) + np.testing.assert_allclose(p[4, 3], + assoc_legendre_p_4_3(z, branch_cut=branch_cut, norm=norm)) + np.testing.assert_allclose(p[4, 4], + assoc_legendre_p_4_4(z, branch_cut=branch_cut, norm=norm)) + np.testing.assert_allclose(p[4, -4], + assoc_legendre_p_4_m4(z, branch_cut=branch_cut, norm=norm)) + np.testing.assert_allclose(p[4, -3], + assoc_legendre_p_4_m3(z, branch_cut=branch_cut, norm=norm)) + np.testing.assert_allclose(p[4, -2], + assoc_legendre_p_4_m2(z, branch_cut=branch_cut, norm=norm)) + np.testing.assert_allclose(p[4, -1], + assoc_legendre_p_4_m1(z, branch_cut=branch_cut, norm=norm)) + + np.testing.assert_allclose(p_jac[0, 0], + assoc_legendre_p_0_0_jac(z, branch_cut=branch_cut, norm=norm)) + np.testing.assert_allclose(p_jac[0, 1], 0) + np.testing.assert_allclose(p_jac[0, 2], 0) + np.testing.assert_allclose(p_jac[0, 3], 0) + np.testing.assert_allclose(p_jac[0, 4], 0) + np.testing.assert_allclose(p_jac[0, -4], 0) + np.testing.assert_allclose(p_jac[0, -3], 0) + np.testing.assert_allclose(p_jac[0, -2], 0) + np.testing.assert_allclose(p_jac[0, -1], 0) + + np.testing.assert_allclose(p_jac[1, 0], + assoc_legendre_p_1_0_jac(z, branch_cut=branch_cut, norm=norm)) + np.testing.assert_allclose(p_jac[1, 1], + assoc_legendre_p_1_1_jac(z, branch_cut=branch_cut, norm=norm)) + np.testing.assert_allclose(p_jac[1, 2], 0) + np.testing.assert_allclose(p_jac[1, 3], 0) + np.testing.assert_allclose(p_jac[1, 4], 0) + np.testing.assert_allclose(p_jac[1, -4], 0) + np.testing.assert_allclose(p_jac[1, -3], 0) + np.testing.assert_allclose(p_jac[1, -2], 0) + np.testing.assert_allclose(p_jac[1, -1], + assoc_legendre_p_1_m1_jac(z, branch_cut=branch_cut, norm=norm)) + + np.testing.assert_allclose(p_jac[2, 0], + assoc_legendre_p_2_0_jac(z, branch_cut=branch_cut, norm=norm)) + np.testing.assert_allclose(p_jac[2, 1], + assoc_legendre_p_2_1_jac(z, branch_cut=branch_cut, norm=norm)) + np.testing.assert_allclose(p_jac[2, 2], + assoc_legendre_p_2_2_jac(z, branch_cut=branch_cut, norm=norm)) + np.testing.assert_allclose(p_jac[2, 3], 0) + np.testing.assert_allclose(p_jac[2, 4], 0) + np.testing.assert_allclose(p_jac[2, -4], 0) + np.testing.assert_allclose(p_jac[2, -3], 0) + np.testing.assert_allclose(p_jac[2, -2], + assoc_legendre_p_2_m2_jac(z, branch_cut=branch_cut, norm=norm)) + np.testing.assert_allclose(p_jac[2, -1], + assoc_legendre_p_2_m1_jac(z, branch_cut=branch_cut, norm=norm)) + + np.testing.assert_allclose(p_jac[3, 0], + assoc_legendre_p_3_0_jac(z, branch_cut=branch_cut, norm=norm)) + np.testing.assert_allclose(p_jac[3, 1], + assoc_legendre_p_3_1_jac(z, branch_cut=branch_cut, norm=norm)) + np.testing.assert_allclose(p_jac[3, 2], + assoc_legendre_p_3_2_jac(z, branch_cut=branch_cut, norm=norm)) + np.testing.assert_allclose(p_jac[3, 3], + assoc_legendre_p_3_3_jac(z, branch_cut=branch_cut, norm=norm)) + np.testing.assert_allclose(p_jac[3, 4], 0) + np.testing.assert_allclose(p_jac[3, -4], 0) + np.testing.assert_allclose(p_jac[3, -3], + assoc_legendre_p_3_m3_jac(z, branch_cut=branch_cut, norm=norm)) + np.testing.assert_allclose(p_jac[3, -2], + assoc_legendre_p_3_m2_jac(z, branch_cut=branch_cut, norm=norm)) + np.testing.assert_allclose(p_jac[3, -1], + assoc_legendre_p_3_m1_jac(z, branch_cut=branch_cut, norm=norm)) + + np.testing.assert_allclose(p_jac[4, 0], + assoc_legendre_p_4_0_jac(z, branch_cut=branch_cut, norm=norm)) + np.testing.assert_allclose(p_jac[4, 1], + assoc_legendre_p_4_1_jac(z, branch_cut=branch_cut, norm=norm)) + np.testing.assert_allclose(p_jac[4, 2], + assoc_legendre_p_4_2_jac(z, branch_cut=branch_cut, norm=norm)) + np.testing.assert_allclose(p_jac[4, 3], + assoc_legendre_p_4_3_jac(z, branch_cut=branch_cut, norm=norm)) + np.testing.assert_allclose(p_jac[4, 4], + assoc_legendre_p_4_4_jac(z, branch_cut=branch_cut, norm=norm)) + np.testing.assert_allclose(p_jac[4, -4], + assoc_legendre_p_4_m4_jac(z, branch_cut=branch_cut, norm=norm)) + np.testing.assert_allclose(p_jac[4, -3], + assoc_legendre_p_4_m3_jac(z, branch_cut=branch_cut, norm=norm)) + np.testing.assert_allclose(p_jac[4, -2], + assoc_legendre_p_4_m2_jac(z, branch_cut=branch_cut, norm=norm)) + np.testing.assert_allclose(p_jac[4, -1], + assoc_legendre_p_4_m1_jac(z, branch_cut=branch_cut, norm=norm)) + +class TestSphLegendreP: + @pytest.mark.parametrize("shape", [(10,), (4, 9), (3, 5, 7)]) + def test_specific(self, shape): + rng = np.random.default_rng(1234) + + theta = rng.uniform(-np.pi, np.pi, shape) + + p, p_jac = sph_legendre_p_all(4, 4, theta, diff_n=1) + + np.testing.assert_allclose(p[0, 0], + sph_legendre_p_0_0(theta)) + np.testing.assert_allclose(p[0, 1], 0) + np.testing.assert_allclose(p[0, 2], 0) + np.testing.assert_allclose(p[0, 3], 0) + np.testing.assert_allclose(p[0, 4], 0) + np.testing.assert_allclose(p[0, -3], 0) + np.testing.assert_allclose(p[0, -2], 0) + np.testing.assert_allclose(p[0, -1], 0) + + np.testing.assert_allclose(p[1, 0], + sph_legendre_p_1_0(theta)) + np.testing.assert_allclose(p[1, 1], + sph_legendre_p_1_1(theta)) + np.testing.assert_allclose(p[1, 2], 0) + np.testing.assert_allclose(p[1, 3], 0) + np.testing.assert_allclose(p[1, 4], 0) + np.testing.assert_allclose(p[1, -4], 0) + np.testing.assert_allclose(p[1, -3], 0) + np.testing.assert_allclose(p[1, -2], 0) + np.testing.assert_allclose(p[1, -1], + sph_legendre_p_1_m1(theta)) + + np.testing.assert_allclose(p[2, 0], + sph_legendre_p_2_0(theta)) + np.testing.assert_allclose(p[2, 1], + sph_legendre_p_2_1(theta)) + np.testing.assert_allclose(p[2, 2], + sph_legendre_p_2_2(theta)) + np.testing.assert_allclose(p[2, 3], 0) + np.testing.assert_allclose(p[2, 4], 0) + np.testing.assert_allclose(p[2, -4], 0) + np.testing.assert_allclose(p[2, -3], 0) + np.testing.assert_allclose(p[2, -2], + sph_legendre_p_2_m2(theta)) + np.testing.assert_allclose(p[2, -1], + sph_legendre_p_2_m1(theta)) + + np.testing.assert_allclose(p[3, 0], + sph_legendre_p_3_0(theta)) + np.testing.assert_allclose(p[3, 1], + sph_legendre_p_3_1(theta)) + np.testing.assert_allclose(p[3, 2], + sph_legendre_p_3_2(theta)) + np.testing.assert_allclose(p[3, 3], + sph_legendre_p_3_3(theta)) + np.testing.assert_allclose(p[3, 4], 0) + np.testing.assert_allclose(p[3, -4], 0) + np.testing.assert_allclose(p[3, -3], + sph_legendre_p_3_m3(theta)) + np.testing.assert_allclose(p[3, -2], + sph_legendre_p_3_m2(theta)) + np.testing.assert_allclose(p[3, -1], + sph_legendre_p_3_m1(theta)) + + np.testing.assert_allclose(p[4, 0], + sph_legendre_p_4_0(theta)) + np.testing.assert_allclose(p[4, 1], + sph_legendre_p_4_1(theta)) + np.testing.assert_allclose(p[4, 2], + sph_legendre_p_4_2(theta)) + np.testing.assert_allclose(p[4, 3], + sph_legendre_p_4_3(theta)) + np.testing.assert_allclose(p[4, 4], + sph_legendre_p_4_4(theta)) + np.testing.assert_allclose(p[4, -4], + sph_legendre_p_4_m4(theta)) + np.testing.assert_allclose(p[4, -3], + sph_legendre_p_4_m3(theta)) + np.testing.assert_allclose(p[4, -2], + sph_legendre_p_4_m2(theta)) + np.testing.assert_allclose(p[4, -1], + sph_legendre_p_4_m1(theta)) + + np.testing.assert_allclose(p_jac[0, 0], + sph_legendre_p_0_0_jac(theta)) + np.testing.assert_allclose(p_jac[0, 1], 0) + np.testing.assert_allclose(p_jac[0, 2], 0) + np.testing.assert_allclose(p_jac[0, 3], 0) + np.testing.assert_allclose(p_jac[0, 4], 0) + np.testing.assert_allclose(p_jac[0, -3], 0) + np.testing.assert_allclose(p_jac[0, -2], 0) + np.testing.assert_allclose(p_jac[0, -1], 0) + + np.testing.assert_allclose(p_jac[1, 0], + sph_legendre_p_1_0_jac(theta)) + np.testing.assert_allclose(p_jac[1, 1], + sph_legendre_p_1_1_jac(theta)) + np.testing.assert_allclose(p_jac[1, 2], 0) + np.testing.assert_allclose(p_jac[1, 3], 0) + np.testing.assert_allclose(p_jac[1, 4], 0) + np.testing.assert_allclose(p_jac[1, -4], 0) + np.testing.assert_allclose(p_jac[1, -3], 0) + np.testing.assert_allclose(p_jac[1, -2], 0) + np.testing.assert_allclose(p_jac[1, -1], + sph_legendre_p_1_m1_jac(theta)) + + np.testing.assert_allclose(p_jac[2, 0], + sph_legendre_p_2_0_jac(theta)) + np.testing.assert_allclose(p_jac[2, 1], + sph_legendre_p_2_1_jac(theta)) + np.testing.assert_allclose(p_jac[2, 2], + sph_legendre_p_2_2_jac(theta)) + np.testing.assert_allclose(p_jac[2, 3], 0) + np.testing.assert_allclose(p_jac[2, 4], 0) + np.testing.assert_allclose(p_jac[2, -4], 0) + np.testing.assert_allclose(p_jac[2, -3], 0) + np.testing.assert_allclose(p_jac[2, -2], + sph_legendre_p_2_m2_jac(theta)) + np.testing.assert_allclose(p_jac[2, -1], + sph_legendre_p_2_m1_jac(theta)) + + np.testing.assert_allclose(p_jac[3, 0], + sph_legendre_p_3_0_jac(theta)) + np.testing.assert_allclose(p_jac[3, 1], + sph_legendre_p_3_1_jac(theta)) + np.testing.assert_allclose(p_jac[3, 2], + sph_legendre_p_3_2_jac(theta)) + np.testing.assert_allclose(p_jac[3, 3], + sph_legendre_p_3_3_jac(theta)) + np.testing.assert_allclose(p_jac[3, 4], 0) + np.testing.assert_allclose(p_jac[3, -4], 0) + np.testing.assert_allclose(p_jac[3, -3], + sph_legendre_p_3_m3_jac(theta)) + np.testing.assert_allclose(p_jac[3, -2], + sph_legendre_p_3_m2_jac(theta)) + np.testing.assert_allclose(p_jac[3, -1], + sph_legendre_p_3_m1_jac(theta)) + + np.testing.assert_allclose(p_jac[4, 0], + sph_legendre_p_4_0_jac(theta)) + np.testing.assert_allclose(p_jac[4, 1], + sph_legendre_p_4_1_jac(theta)) + np.testing.assert_allclose(p_jac[4, 2], + sph_legendre_p_4_2_jac(theta)) + np.testing.assert_allclose(p_jac[4, 3], + sph_legendre_p_4_3_jac(theta)) + np.testing.assert_allclose(p_jac[4, 4], + sph_legendre_p_4_4_jac(theta)) + np.testing.assert_allclose(p_jac[4, -4], + sph_legendre_p_4_m4_jac(theta)) + np.testing.assert_allclose(p_jac[4, -3], + sph_legendre_p_4_m3_jac(theta)) + np.testing.assert_allclose(p_jac[4, -2], + sph_legendre_p_4_m2_jac(theta)) + np.testing.assert_allclose(p_jac[4, -1], + sph_legendre_p_4_m1_jac(theta)) + + @pytest.mark.parametrize("shape", [(10,), (4, 9), (3, 5, 7, 10)]) + def test_ode(self, shape): + rng = np.random.default_rng(1234) + + n = rng.integers(0, 10, shape) + m = rng.integers(-10, 10, shape) + theta = rng.uniform(-np.pi, np.pi, shape) + + p, p_jac, p_hess = sph_legendre_p(n, m, theta, diff_n=2) + + assert p.shape == shape + assert p_jac.shape == p.shape + assert p_hess.shape == p_jac.shape + + np.testing.assert_allclose(np.sin(theta) * p_hess, -np.cos(theta) * p_jac + - (n * (n + 1) * np.sin(theta) - m * m / np.sin(theta)) * p, + rtol=1e-05, atol=1e-08) + +class TestLegendreFunctions: + """ + @pytest.mark.parametrize("m_max", [3]) + @pytest.mark.parametrize("n_max", [5]) + @pytest.mark.parametrize("z", [-1]) + def test_clpmn_all_limits(self, m_max, n_max, z): + rng = np.random.default_rng(1234) + + type = 2 + + p, p_jac = special.clpmn_all(m_max, n_max, type, z, diff_n=1) + + n = np.arange(n_max + 1) + + np.testing.assert_allclose(p_jac[0], pow(z, n + 1) * n * (n + 1) / 2) + np.testing.assert_allclose(p_jac[1], np.where(n >= 1, pow(z, n) * np.inf, 0)) + np.testing.assert_allclose(p_jac[2], np.where(n >= 2, + -pow(z, n + 1) * (n + 2) * (n + 1) * n * (n - 1) / 4, 0)) + np.testing.assert_allclose(p_jac[-2], np.where(n >= 2, -pow(z, n + 1) / 4, 0)) + np.testing.assert_allclose(p_jac[-1], np.where(n >= 1, -pow(z, n) * np.inf, 0)) + + for m in range(3, m_max + 1): + np.testing.assert_allclose(p_jac[m], 0) + np.testing.assert_allclose(p_jac[-m], 0) + """ + + def test_lpmv(self): + lp = special.lpmv(0, 2, .5) + assert_allclose(lp, -0.125, atol=1.5e-7, rtol=0) + lp = special.lpmv(0, 40, .001) + assert_allclose(lp, 0.1252678976534484, atol=1.5e-7, rtol=0) + + # XXX: this is outside the domain of the current implementation, + # so ensure it returns a NaN rather than a wrong answer. + with np.errstate(all='ignore'): + lp = special.lpmv(-1,-1,.001) + assert lp != 0 or np.isnan(lp) + + def test_lqmn(self): + lqmnf = special.lqmn(0, 2, .5) + lqf = special.lqn(2, .5) + assert_allclose(lqmnf[0][0], lqf[0], atol=1.5e-4, rtol=0) + assert_allclose(lqmnf[1][0], lqf[1], atol=1.5e-4, rtol=0) + + def test_lqmn_gt1(self): + """algorithm for real arguments changes at 1.0001 + test against analytical result for m=2, n=1 + """ + x0 = 1.0001 + delta = 0.00002 + for x in (x0-delta, x0+delta): + lq = special.lqmn(2, 1, x)[0][-1, -1] + expected = 2/(x*x-1) + assert_allclose(lq, expected, atol=1.5e-7, rtol=0) + + def test_lqmn_shape(self): + a, b = special.lqmn(4, 4, 1.1) + assert_equal(a.shape, (5, 5)) + assert_equal(b.shape, (5, 5)) + + a, b = special.lqmn(4, 0, 1.1) + assert_equal(a.shape, (5, 1)) + assert_equal(b.shape, (5, 1)) + + def test_lqn(self): + lqf = special.lqn(2, .5) + assert_allclose(lqf, (np.array([0.5493, -0.7253, -0.8187]), + np.array([1.3333, 1.216, -0.8427])), + atol=1.5e-4, rtol=0) + + @pytest.mark.parametrize("function", [special.lqn]) + @pytest.mark.parametrize("n", [1, 2, 4, 8, 16, 32]) + @pytest.mark.parametrize("z_complex", [False, True]) + @pytest.mark.parametrize("z_inexact", [False, True]) + @pytest.mark.parametrize( + "input_shape", + [ + (), (1, ), (2, ), (2, 1), (1, 2), (2, 2), (2, 2, 1), (2, 2, 2) + ] + ) + def test_array_inputs_lxn(self, function, n, z_complex, z_inexact, input_shape): + """Tests for correct output shapes.""" + rng = np.random.default_rng(1234) + if z_inexact: + z = rng.integers(-3, 3, size=input_shape) + else: + z = rng.uniform(-1, 1, size=input_shape) + + if z_complex: + z = 1j * z + 0.5j * z + + with warnings.catch_warnings(): + warnings.simplefilter("ignore", category=DeprecationWarning) + P_z, P_d_z = function(n, z) + assert P_z.shape == (n + 1, ) + input_shape + assert P_d_z.shape == (n + 1, ) + input_shape + + @pytest.mark.parametrize("function", [special.lqmn]) + @pytest.mark.parametrize( + "m,n", + [(0, 1), (1, 2), (1, 4), (3, 8), (11, 16), (19, 32)] + ) + @pytest.mark.parametrize("z_inexact", [False, True]) + @pytest.mark.parametrize( + "input_shape", [ + (), (1, ), (2, ), (2, 1), (1, 2), (2, 2), (2, 2, 1) + ] + ) + def test_array_inputs_lxmn(self, function, m, n, z_inexact, input_shape): + """Tests for correct output shapes and dtypes.""" + rng = np.random.default_rng(1234) + if z_inexact: + z = rng.integers(-3, 3, size=input_shape) + else: + z = rng.uniform(-1, 1, size=input_shape) + + P_z, P_d_z = function(m, n, z) + assert P_z.shape == (m + 1, n + 1) + input_shape + assert P_d_z.shape == (m + 1, n + 1) + input_shape + + @pytest.mark.parametrize("function", [special.lqmn]) + @pytest.mark.parametrize( + "m,n", + [(0, 1), (1, 2), (1, 4), (3, 8), (11, 16), (19, 32)] + ) + @pytest.mark.parametrize( + "input_shape", [ + (), (1, ), (2, ), (2, 1), (1, 2), (2, 2), (2, 2, 1) + ] + ) + def test_array_inputs_clxmn(self, function, m, n, input_shape): + """Tests for correct output shapes and dtypes.""" + rng = np.random.default_rng(1234) + z = rng.uniform(-1, 1, size=input_shape) + z = 1j * z + 0.5j * z + + with warnings.catch_warnings(): + warnings.simplefilter("ignore", category=DeprecationWarning) + P_z, P_d_z = function(m, n, z) + + assert P_z.shape == (m + 1, n + 1) + input_shape + assert P_d_z.shape == (m + 1, n + 1) + input_shape + +def assoc_legendre_factor(n, m, norm): + if norm: + return (math.sqrt((2 * n + 1) * + math.factorial(n - m) / (2 * math.factorial(n + m)))) + + return 1 + +def assoc_legendre_p_0_0(z, *, branch_cut=2, norm=False): + fac = assoc_legendre_factor(0, 0, norm) + + return np.full_like(z, fac) + +def assoc_legendre_p_1_0(z, *, branch_cut=2, norm=False): + fac = assoc_legendre_factor(1, 0, norm) + + return fac * z + +def assoc_legendre_p_1_1(z, *, branch_cut=2, norm=False): + branch_sign = np.where(branch_cut == 3, np.where(np.signbit(np.real(z)), 1, -1), -1) + branch_cut_sign = np.where(branch_cut == 3, -1, 1) + fac = assoc_legendre_factor(1, 1, norm) + + w = np.sqrt(np.where(branch_cut == 3, z * z - 1, 1 - z * z)) + + return branch_cut_sign * branch_sign * fac * w + +def assoc_legendre_p_1_m1(z, *, branch_cut=2, norm=False): + branch_cut_sign = np.where(branch_cut == 3, -1, 1) + fac = assoc_legendre_factor(1, -1, norm) + + return (-branch_cut_sign * fac * + assoc_legendre_p_1_1(z, branch_cut=branch_cut) / 2) + +def assoc_legendre_p_2_0(z, *, branch_cut=2, norm=False): + fac = assoc_legendre_factor(2, 0, norm) + + return fac * (3 * z * z - 1) / 2 + +def assoc_legendre_p_2_1(z, *, branch_cut=2, norm=False): + fac = assoc_legendre_factor(2, 1, norm) + + return (3 * fac * z * + assoc_legendre_p_1_1(z, branch_cut=branch_cut)) + +def assoc_legendre_p_2_2(z, *, branch_cut=2, norm=False): + branch_cut_sign = np.where(branch_cut == 3, -1, 1) + fac = assoc_legendre_factor(2, 2, norm) + + return 3 * branch_cut_sign * fac * (1 - z * z) + +def assoc_legendre_p_2_m2(z, *, branch_cut=2, norm=False): + branch_cut_sign = np.where(branch_cut == 3, -1, 1) + fac = assoc_legendre_factor(2, -2, norm) + + return branch_cut_sign * fac * (1 - z * z) / 8 + +def assoc_legendre_p_2_m1(z, *, branch_cut=2, norm=False): + branch_cut_sign = np.where(branch_cut == 3, -1, 1) + fac = assoc_legendre_factor(2, -1, norm) + + return (-branch_cut_sign * fac * z * + assoc_legendre_p_1_1(z, branch_cut=branch_cut) / 2) + +def assoc_legendre_p_3_0(z, *, branch_cut=2, norm=False): + fac = assoc_legendre_factor(3, 0, norm) + + return fac * (5 * z * z - 3) * z / 2 + +def assoc_legendre_p_3_1(z, *, branch_cut=2, norm=False): + fac = assoc_legendre_factor(3, 1, norm) + + return (3 * fac * (5 * z * z - 1) * + assoc_legendre_p_1_1(z, branch_cut=branch_cut) / 2) + +def assoc_legendre_p_3_2(z, *, branch_cut=2, norm=False): + branch_cut_sign = np.where(branch_cut == 3, -1, 1) + fac = assoc_legendre_factor(3, 2, norm) + + return 15 * branch_cut_sign * fac * (1 - z * z) * z + +def assoc_legendre_p_3_3(z, *, branch_cut=2, norm=False): + branch_cut_sign = np.where(branch_cut == 3, -1, 1) + fac = assoc_legendre_factor(3, 3, norm) + + return (15 * branch_cut_sign * fac * (1 - z * z) * + assoc_legendre_p_1_1(z, branch_cut=branch_cut)) + +def assoc_legendre_p_3_m3(z, *, branch_cut=2, norm=False): + fac = assoc_legendre_factor(3, -3, norm) + + return (fac * (z * z - 1) * + assoc_legendre_p_1_1(z, branch_cut=branch_cut) / 48) + +def assoc_legendre_p_3_m2(z, *, branch_cut=2, norm=False): + branch_cut_sign = np.where(branch_cut == 3, -1, 1) + fac = assoc_legendre_factor(3, -2, norm) + + return branch_cut_sign * fac * (1 - z * z) * z / 8 + +def assoc_legendre_p_3_m1(z, *, branch_cut=2, norm=False): + branch_cut_sign = np.where(branch_cut == 3, -1, 1) + fac = assoc_legendre_factor(3, -1, norm) + + return (branch_cut_sign * fac * (1 - 5 * z * z) * + assoc_legendre_p_1_1(z, branch_cut=branch_cut) / 8) + +def assoc_legendre_p_4_0(z, *, branch_cut=2, norm=False): + fac = assoc_legendre_factor(4, 0, norm) + + return fac * ((35 * z * z - 30) * z * z + 3) / 8 + +def assoc_legendre_p_4_1(z, *, branch_cut=2, norm=False): + fac = assoc_legendre_factor(4, 1, norm) + + return (5 * fac * (7 * z * z - 3) * z * + assoc_legendre_p_1_1(z, branch_cut=branch_cut) / 2) + +def assoc_legendre_p_4_2(z, *, branch_cut=2, norm=False): + branch_cut_sign = np.where(branch_cut == 3, -1, 1) + fac = assoc_legendre_factor(4, 2, norm) + + return 15 * branch_cut_sign * fac * ((8 - 7 * z * z) * z * z - 1) / 2 + +def assoc_legendre_p_4_3(z, *, branch_cut=2, norm=False): + branch_cut_sign = np.where(branch_cut == 3, -1, 1) + fac = assoc_legendre_factor(4, 3, norm) + + return (105 * branch_cut_sign * fac * (1 - z * z) * z * + assoc_legendre_p_1_1(z, branch_cut=branch_cut)) + +def assoc_legendre_p_4_4(z, *, branch_cut=2, norm=False): + fac = assoc_legendre_factor(4, 4, norm) + + return 105 * fac * np.square(z * z - 1) + +def assoc_legendre_p_4_m4(z, *, branch_cut=2, norm=False): + fac = assoc_legendre_factor(4, -4, norm) + + return fac * np.square(z * z - 1) / 384 + +def assoc_legendre_p_4_m3(z, *, branch_cut=2, norm=False): + fac = assoc_legendre_factor(4, -3, norm) + + return (fac * (z * z - 1) * z * + assoc_legendre_p_1_1(z, branch_cut=branch_cut) / 48) + +def assoc_legendre_p_4_m2(z, *, branch_cut=2, norm=False): + branch_cut_sign = np.where(branch_cut == 3, -1, 1) + fac = assoc_legendre_factor(4, -2, norm) + + return branch_cut_sign * fac * ((8 - 7 * z * z) * z * z - 1) / 48 + +def assoc_legendre_p_4_m1(z, *, branch_cut=2, norm=False): + branch_cut_sign = np.where(branch_cut == 3, -1, 1) + fac = assoc_legendre_factor(4, -1, norm) + + return (branch_cut_sign * fac * (3 - 7 * z * z) * z * + assoc_legendre_p_1_1(z, branch_cut=branch_cut) / 8) + +def assoc_legendre_p_1_1_jac_div_z(z, branch_cut=2): + branch_sign = np.where(branch_cut == 3, np.where(np.signbit(np.real(z)), 1, -1), -1) + + out11_div_z = (-branch_sign / + np.sqrt(np.where(branch_cut == 3, z * z - 1, 1 - z * z))) + + return out11_div_z + +def assoc_legendre_p_0_0_jac(z, *, branch_cut=2, norm=False): + return np.zeros_like(z) + +def assoc_legendre_p_1_0_jac(z, *, branch_cut=2, norm=False): + fac = assoc_legendre_factor(1, 0, norm) + + return np.full_like(z, fac) + +def assoc_legendre_p_1_1_jac(z, *, branch_cut=2, norm=False): + fac = assoc_legendre_factor(1, 1, norm) + + return (fac * z * + assoc_legendre_p_1_1_jac_div_z(z, branch_cut=branch_cut)) + +def assoc_legendre_p_1_m1_jac(z, *, branch_cut=2, norm=False): + branch_cut_sign = np.where(branch_cut == 3, -1, 1) + fac = assoc_legendre_factor(1, -1, norm) + + return (-branch_cut_sign * fac * z * + assoc_legendre_p_1_1_jac_div_z(z, branch_cut=branch_cut) / 2) + +def assoc_legendre_p_2_0_jac(z, *, branch_cut=2, norm=False): + fac = assoc_legendre_factor(2, 0, norm) + + return 3 * fac * z + +def assoc_legendre_p_2_1_jac(z, *, branch_cut=2, norm=False): + fac = assoc_legendre_factor(2, 1, norm) + + return (3 * fac * (2 * z * z - 1) * + assoc_legendre_p_1_1_jac_div_z(z, branch_cut=branch_cut)) + +def assoc_legendre_p_2_2_jac(z, *, branch_cut=2, norm=False): + branch_cut_sign = np.where(branch_cut == 3, -1, 1) + fac = assoc_legendre_factor(2, 2, norm) + + return -6 * branch_cut_sign * fac * z + +def assoc_legendre_p_2_m1_jac(z, *, branch_cut=2, norm=False): + branch_cut_sign = np.where(branch_cut == 3, -1, 1) + fac = assoc_legendre_factor(2, -1, norm) + + return (branch_cut_sign * fac * (1 - 2 * z * z) * + assoc_legendre_p_1_1_jac_div_z(z, branch_cut=branch_cut) / 2) + +def assoc_legendre_p_2_m2_jac(z, *, branch_cut=2, norm=False): + branch_cut_sign = np.where(branch_cut == 3, -1, 1) + fac = assoc_legendre_factor(2, -2, norm) + + return -branch_cut_sign * fac * z / 4 + +def assoc_legendre_p_3_0_jac(z, *, branch_cut=2, norm=False): + fac = assoc_legendre_factor(3, 0, norm) + + return 3 * fac * (5 * z * z - 1) / 2 + +def assoc_legendre_p_3_1_jac(z, *, branch_cut=2, norm=False): + fac = assoc_legendre_factor(3, 1, norm) + + return (3 * fac * (15 * z * z - 11) * z * + assoc_legendre_p_1_1_jac_div_z(z, branch_cut=branch_cut) / 2) + +def assoc_legendre_p_3_2_jac(z, *, branch_cut=2, norm=False): + branch_cut_sign = np.where(branch_cut == 3, -1, 1) + fac = assoc_legendre_factor(3, 2, norm) + + return 15 * branch_cut_sign * fac * (1 - 3 * z * z) + +def assoc_legendre_p_3_3_jac(z, *, branch_cut=2, norm=False): + branch_cut_sign = np.where(branch_cut == 3, -1, 1) + fac = assoc_legendre_factor(3, 3, norm) + + return (45 * branch_cut_sign * fac * (1 - z * z) * z * + assoc_legendre_p_1_1_jac_div_z(z, branch_cut=branch_cut)) + +def assoc_legendre_p_3_m3_jac(z, *, branch_cut=2, norm=False): + fac = assoc_legendre_factor(3, -3, norm) + + return (fac * (z * z - 1) * z * + assoc_legendre_p_1_1_jac_div_z(z, branch_cut=branch_cut) / 16) + +def assoc_legendre_p_3_m2_jac(z, *, branch_cut=2, norm=False): + branch_cut_sign = np.where(branch_cut == 3, -1, 1) + fac = assoc_legendre_factor(3, -2, norm) + + return branch_cut_sign * fac * (1 - 3 * z * z) / 8 + +def assoc_legendre_p_3_m1_jac(z, *, branch_cut=2, norm=False): + branch_cut_sign = np.where(branch_cut == 3, -1, 1) + fac = assoc_legendre_factor(3, -1, norm) + + return (branch_cut_sign * fac * (11 - 15 * z * z) * z * + assoc_legendre_p_1_1_jac_div_z(z, branch_cut=branch_cut) / 8) + +def assoc_legendre_p_4_0_jac(z, *, branch_cut=2, norm=False): + fac = assoc_legendre_factor(4, 0, norm) + + return 5 * fac * (7 * z * z - 3) * z / 2 + +def assoc_legendre_p_4_1_jac(z, *, branch_cut=2, norm=False): + fac = assoc_legendre_factor(4, 1, norm) + + return (5 * fac * ((28 * z * z - 27) * z * z + 3) * + assoc_legendre_p_1_1_jac_div_z(z, branch_cut=branch_cut) / 2) + +def assoc_legendre_p_4_2_jac(z, *, branch_cut=2, norm=False): + branch_cut_sign = np.where(branch_cut == 3, -1, 1) + fac = assoc_legendre_factor(4, 2, norm) + + return 30 * branch_cut_sign * fac * (4 - 7 * z * z) * z + +def assoc_legendre_p_4_3_jac(z, *, branch_cut=2, norm=False): + branch_cut_sign = np.where(branch_cut == 3, -1, 1) + fac = assoc_legendre_factor(4, 3, norm) + + return (105 * branch_cut_sign * fac * ((5 - 4 * z * z) * z * z - 1) * + assoc_legendre_p_1_1_jac_div_z(z, branch_cut=branch_cut)) + +def assoc_legendre_p_4_4_jac(z, *, branch_cut=2, norm=False): + fac = assoc_legendre_factor(4, 4, norm) + + return 420 * fac * (z * z - 1) * z + +def assoc_legendre_p_4_m4_jac(z, *, branch_cut=2, norm=False): + fac = assoc_legendre_factor(4, -4, norm) + + return fac * (z * z - 1) * z / 96 + +def assoc_legendre_p_4_m3_jac(z, *, branch_cut=2, norm=False): + fac = assoc_legendre_factor(4, -3, norm) + + return (fac * ((4 * z * z - 5) * z * z + 1) * + assoc_legendre_p_1_1_jac_div_z(z, branch_cut=branch_cut) / 48) + +def assoc_legendre_p_4_m2_jac(z, *, branch_cut=2, norm=False): + branch_cut_sign = np.where(branch_cut == 3, -1, 1) + fac = assoc_legendre_factor(4, -2, norm) + + return branch_cut_sign * fac * (4 - 7 * z * z) * z / 12 + +def assoc_legendre_p_4_m1_jac(z, *, branch_cut=2, norm=False): + branch_cut_sign = np.where(branch_cut == 3, -1, 1) + fac = assoc_legendre_factor(4, -1, norm) + + return (branch_cut_sign * fac * ((27 - 28 * z * z) * z * z - 3) * + assoc_legendre_p_1_1_jac_div_z(z, branch_cut=branch_cut) / 8) + +def sph_legendre_factor(n, m): + return assoc_legendre_factor(n, m, norm=True) / np.sqrt(2 * np.pi) + +def sph_legendre_p_0_0(theta): + fac = sph_legendre_factor(0, 0) + + return np.full_like(theta, fac) + +def sph_legendre_p_1_0(theta): + fac = sph_legendre_factor(1, 0) + + return fac * np.cos(theta) + +def sph_legendre_p_1_1(theta): + fac = sph_legendre_factor(1, 1) + + return -fac * np.abs(np.sin(theta)) + +def sph_legendre_p_1_m1(theta): + fac = sph_legendre_factor(1, -1) + + return fac * np.abs(np.sin(theta)) / 2 + +def sph_legendre_p_2_0(theta): + fac = sph_legendre_factor(2, 0) + + return fac * (3 * np.square(np.cos(theta)) - 1) / 2 + +def sph_legendre_p_2_1(theta): + fac = sph_legendre_factor(2, 1) + + return -3 * fac * np.abs(np.sin(theta)) * np.cos(theta) + +def sph_legendre_p_2_2(theta): + fac = sph_legendre_factor(2, 2) + + return 3 * fac * (1 - np.square(np.cos(theta))) + +def sph_legendre_p_2_m2(theta): + fac = sph_legendre_factor(2, -2) + + return fac * (1 - np.square(np.cos(theta))) / 8 + +def sph_legendre_p_2_m1(theta): + fac = sph_legendre_factor(2, -1) + + return fac * np.cos(theta) * np.abs(np.sin(theta)) / 2 + +def sph_legendre_p_3_0(theta): + fac = sph_legendre_factor(3, 0) + + return (fac * (5 * np.square(np.cos(theta)) - 3) * + np.cos(theta) / 2) + +def sph_legendre_p_3_1(theta): + fac = sph_legendre_factor(3, 1) + + return (-3 * fac * (5 * np.square(np.cos(theta)) - 1) * + np.abs(np.sin(theta)) / 2) + +def sph_legendre_p_3_2(theta): + fac = sph_legendre_factor(3, 2) + + return (-15 * fac * (np.square(np.cos(theta)) - 1) * + np.cos(theta)) + +def sph_legendre_p_3_3(theta): + fac = sph_legendre_factor(3, 3) + + return -15 * fac * np.power(np.abs(np.sin(theta)), 3) + +def sph_legendre_p_3_m3(theta): + fac = sph_legendre_factor(3, -3) + + return fac * np.power(np.abs(np.sin(theta)), 3) / 48 + +def sph_legendre_p_3_m2(theta): + fac = sph_legendre_factor(3, -2) + + return (-fac * (np.square(np.cos(theta)) - 1) * + np.cos(theta) / 8) + +def sph_legendre_p_3_m1(theta): + fac = sph_legendre_factor(3, -1) + + return (fac * (5 * np.square(np.cos(theta)) - 1) * + np.abs(np.sin(theta)) / 8) + +def sph_legendre_p_4_0(theta): + fac = sph_legendre_factor(4, 0) + + return (fac * (35 * np.square(np.square(np.cos(theta))) - + 30 * np.square(np.cos(theta)) + 3) / 8) + +def sph_legendre_p_4_1(theta): + fac = sph_legendre_factor(4, 1) + + return (-5 * fac * (7 * np.square(np.cos(theta)) - 3) * + np.cos(theta) * np.abs(np.sin(theta)) / 2) + +def sph_legendre_p_4_2(theta): + fac = sph_legendre_factor(4, 2) + + return (-15 * fac * (7 * np.square(np.cos(theta)) - 1) * + (np.square(np.cos(theta)) - 1) / 2) + +def sph_legendre_p_4_3(theta): + fac = sph_legendre_factor(4, 3) + + return -105 * fac * np.power(np.abs(np.sin(theta)), 3) * np.cos(theta) + +def sph_legendre_p_4_4(theta): + fac = sph_legendre_factor(4, 4) + + return 105 * fac * np.square(np.square(np.cos(theta)) - 1) + +def sph_legendre_p_4_m4(theta): + fac = sph_legendre_factor(4, -4) + + return fac * np.square(np.square(np.cos(theta)) - 1) / 384 + +def sph_legendre_p_4_m3(theta): + fac = sph_legendre_factor(4, -3) + + return (fac * np.power(np.abs(np.sin(theta)), 3) * + np.cos(theta) / 48) + +def sph_legendre_p_4_m2(theta): + fac = sph_legendre_factor(4, -2) + + return (-fac * (7 * np.square(np.cos(theta)) - 1) * + (np.square(np.cos(theta)) - 1) / 48) + +def sph_legendre_p_4_m1(theta): + fac = sph_legendre_factor(4, -1) + + return (fac * (7 * np.square(np.cos(theta)) - 3) * + np.cos(theta) * np.abs(np.sin(theta)) / 8) + +def sph_legendre_p_0_0_jac(theta): + return np.zeros_like(theta) + +def sph_legendre_p_1_0_jac(theta): + fac = sph_legendre_factor(1, 0) + + return -fac * np.sin(theta) + +def sph_legendre_p_1_1_jac(theta): + fac = sph_legendre_factor(1, 1) + + return -fac * np.cos(theta) * (2 * np.heaviside(np.sin(theta), 1) - 1) + +def sph_legendre_p_1_m1_jac(theta): + fac = sph_legendre_factor(1, -1) + + return fac * np.cos(theta) * (2 * np.heaviside(np.sin(theta), 1) - 1) / 2 + +def sph_legendre_p_2_0_jac(theta): + fac = sph_legendre_factor(2, 0) + + return -3 * fac * np.cos(theta) * np.sin(theta) + +def sph_legendre_p_2_1_jac(theta): + fac = sph_legendre_factor(2, 1) + + return (3 * fac * (-np.square(np.cos(theta)) * + (2 * np.heaviside(np.sin(theta), 1) - 1) + + np.abs(np.sin(theta)) * np.sin(theta))) + +def sph_legendre_p_2_2_jac(theta): + fac = sph_legendre_factor(2, 2) + + return 6 * fac * np.sin(theta) * np.cos(theta) + +def sph_legendre_p_2_m2_jac(theta): + fac = sph_legendre_factor(2, -2) + + return fac * np.sin(theta) * np.cos(theta) / 4 + +def sph_legendre_p_2_m1_jac(theta): + fac = sph_legendre_factor(2, -1) + + return (-fac * (-np.square(np.cos(theta)) * + (2 * np.heaviside(np.sin(theta), 1) - 1) + + np.abs(np.sin(theta)) * np.sin(theta)) / 2) + +def sph_legendre_p_3_0_jac(theta): + fac = sph_legendre_factor(3, 0) + + return 3 * fac * (1 - 5 * np.square(np.cos(theta))) * np.sin(theta) / 2 + +def sph_legendre_p_3_1_jac(theta): + fac = sph_legendre_factor(3, 1) + + return (3 * fac * (11 - 15 * np.square(np.cos(theta))) * np.cos(theta) * + (2 * np.heaviside(np.sin(theta), 1) - 1) / 2) + +def sph_legendre_p_3_2_jac(theta): + fac = sph_legendre_factor(3, 2) + + return 15 * fac * (3 * np.square(np.cos(theta)) - 1) * np.sin(theta) + +def sph_legendre_p_3_3_jac(theta): + fac = sph_legendre_factor(3, 3) + + return -45 * fac * np.abs(np.sin(theta)) * np.sin(theta) * np.cos(theta) + +def sph_legendre_p_3_m3_jac(theta): + fac = sph_legendre_factor(3, -3) + + return fac * np.abs(np.sin(theta)) * np.sin(theta) * np.cos(theta) / 16 + +def sph_legendre_p_3_m2_jac(theta): + fac = sph_legendre_factor(3, -2) + + return fac * (3 * np.square(np.cos(theta)) - 1) * np.sin(theta) / 8 + +def sph_legendre_p_3_m1_jac(theta): + fac = sph_legendre_factor(3, -1) + + return (-fac * (11 - 15 * np.square(np.cos(theta))) * + np.cos(theta) * + (2 * np.heaviside(np.sin(theta), 1) - 1) / 8) + +def sph_legendre_p_4_0_jac(theta): + fac = sph_legendre_factor(4, 0) + + return (-5 * fac * (7 * np.square(np.cos(theta)) - 3) * + np.sin(theta) * np.cos(theta) / 2) + +def sph_legendre_p_4_1_jac(theta): + fac = sph_legendre_factor(4, 1) + + return (5 * fac * (-3 + 27 * np.square(np.cos(theta)) - + 28 * np.square(np.square(np.cos(theta)))) * + (2 * np.heaviside(np.sin(theta), 1) - 1) / 2) + +def sph_legendre_p_4_2_jac(theta): + fac = sph_legendre_factor(4, 2) + + return (30 * fac * (7 * np.square(np.cos(theta)) - 4) * + np.sin(theta) * np.cos(theta)) + +def sph_legendre_p_4_3_jac(theta): + fac = sph_legendre_factor(4, 3) + + return (-105 * fac * (4 * np.square(np.cos(theta)) - 1) * + np.abs(np.sin(theta)) * np.sin(theta)) + +def sph_legendre_p_4_4_jac(theta): + fac = sph_legendre_factor(4, 4) + + return (-420 * fac * (np.square(np.cos(theta)) - 1) * + np.sin(theta) * np.cos(theta)) + +def sph_legendre_p_4_m4_jac(theta): + fac = sph_legendre_factor(4, -4) + + return (-fac * (np.square(np.cos(theta)) - 1) * + np.sin(theta) * np.cos(theta) / 96) + +def sph_legendre_p_4_m3_jac(theta): + fac = sph_legendre_factor(4, -3) + + return (fac * (4 * np.square(np.cos(theta)) - 1) * + np.abs(np.sin(theta)) * np.sin(theta) / 48) + +def sph_legendre_p_4_m2_jac(theta): + fac = sph_legendre_factor(4, -2) + + return (fac * (7 * np.square(np.cos(theta)) - 4) * np.sin(theta) * + np.cos(theta) / 12) + +def sph_legendre_p_4_m1_jac(theta): + fac = sph_legendre_factor(4, -1) + + return (-fac * (-3 + 27 * np.square(np.cos(theta)) - + 28 * np.square(np.square(np.cos(theta)))) * + (2 * np.heaviside(np.sin(theta), 1) - 1) / 8) diff --git a/v12x/pylib/scipy/special/tests/test_log1mexp.py b/v12x/pylib/scipy/special/tests/test_log1mexp.py new file mode 100644 index 0000000000000000000000000000000000000000..ed8b1958eabd973704c32e9be157397c7b3472e3 --- /dev/null +++ b/v12x/pylib/scipy/special/tests/test_log1mexp.py @@ -0,0 +1,85 @@ +import numpy as np +import pytest + +from numpy.testing import assert_allclose, assert_equal + +from scipy.special._ufuncs import _log1mexp + +# # Test cases generated with the script +# +# import numpy as np + +# from mpmath import mp + + +# def mp_log1mexp(x): +# with mp.workdps(324): +# return float(mp.log(mp.one - mp.exp(x))) + +# X = np.concat([-np.logspace(-1, -300, 20), np.linspace(-745, -1, 20)]) + +# cases = [(float(x), mp_log1mexp(x)) for x in X] + +@pytest.mark.parametrize( + "x,expected", + [ + (-0.1, -2.3521684610440907), + (-1.8329807108324374e-17, -38.538003135374026), + (-3.359818286283788e-33, -74.773421177754), + (-6.1584821106602796e-49, -111.00883922013399), + (-1.1288378916846929e-64, -147.24425726251397), + (-2.0691380811148324e-80, -183.47967530489393), + (-3.792690190732269e-96, -219.71509334727392), + (-6.951927961775534e-112, -255.95051138965394), + (-1.2742749857031425e-127, -292.1859294320339), + (-2.3357214690901785e-143, -328.42134747441384), + (-4.281332398719571e-159, -364.6567655167938), + (-7.847599703514559e-175, -400.8921835591739), + (-1.4384498882876776e-190, -437.1276016015538), + (-2.6366508987304307e-206, -473.3630196439338), + (-4.832930238571653e-222, -509.59843768631384), + (-8.858667904100796e-238, -545.8338557286938), + (-1.623776739188744e-253, -582.0692737710738), + (-2.9763514416312156e-269, -618.3046918134538), + (-5.455594781168782e-285, -654.5401098558336), + (-1e-300, -690.7755278982137), + (-745.0, -5e-324), + (-705.8421052631579, -2.8619931451743316e-307), + (-666.6842105263158, -2.9021923726875757e-290), + (-627.5263157894738, -2.9429562339405562e-273), + (-588.3684210526316, -2.9842926597143714e-256), + (-549.2105263157895, -3.0262096921839423e-239), + (-510.0526315789474, -3.0687154864846747e-222), + (-470.89473684210526, -3.1118183122979086e-205), + (-431.7368421052632, -3.155526555459449e-188), + (-392.5789473684211, -3.1998487195921207e-171), + (-353.42105263157896, -3.2447934277596653e-154), + (-314.2631578947369, -3.2903694241438367e-137), + (-275.1052631578948, -3.3365855757467166e-120), + (-235.94736842105266, -3.3834508741152875e-103), + (-196.78947368421052, -3.4309744370903894e-86), + (-157.63157894736844, -3.4791655105810003e-69), + (-118.47368421052636, -3.528033470363468e-52), + (-79.31578947368428, -3.577587823905024e-35), + (-40.157894736842195, -3.627838212213697e-18), + (-1.0, -0.4586751453870819), + ] +) +def test_log1mexp(x, expected): + observed = _log1mexp(x) + assert_allclose(observed, expected, rtol=1e-15, atol=1e-300) + + +@pytest.mark.parametrize("x", [1.1, 1e10, np.inf]) +def test_log1mexp_out_of_domain(x): + observed = _log1mexp(x) + assert np.isnan(observed) + + +@pytest.mark.parametrize( + "x,expected", + [(-np.inf, -0.0), (0.0, -np.inf), (-0.0, -np.inf), (np.nan, np.nan)] +) +def test_log1mexp_extreme(x, expected): + observed = _log1mexp(x) + assert_equal(expected, observed) diff --git a/v12x/pylib/scipy/special/tests/test_loggamma.py b/v12x/pylib/scipy/special/tests/test_loggamma.py new file mode 100644 index 0000000000000000000000000000000000000000..2fcb5a20037de46df939895d38fbe5fe6b85c9aa --- /dev/null +++ b/v12x/pylib/scipy/special/tests/test_loggamma.py @@ -0,0 +1,70 @@ +import numpy as np +from numpy.testing import assert_allclose, assert_ + +from scipy.special._testutils import FuncData +from scipy.special import gamma, gammaln, loggamma + + +def test_identities1(): + # test the identity exp(loggamma(z)) = gamma(z) + x = np.array([-99.5, -9.5, -0.5, 0.5, 9.5, 99.5]) + y = x.copy() + x, y = np.meshgrid(x, y) + z = (x + 1J*y).flatten() + dataset = np.vstack((z, gamma(z))).T + + def f(z): + return np.exp(loggamma(z)) + + FuncData(f, dataset, 0, 1, rtol=1e-14, atol=1e-14).check() + + +def test_identities2(): + # test the identity loggamma(z + 1) = log(z) + loggamma(z) + x = np.array([-99.5, -9.5, -0.5, 0.5, 9.5, 99.5]) + y = x.copy() + x, y = np.meshgrid(x, y) + z = (x + 1J*y).flatten() + dataset = np.vstack((z, np.log(z) + loggamma(z))).T + + def f(z): + return loggamma(z + 1) + + FuncData(f, dataset, 0, 1, rtol=1e-14, atol=1e-14).check() + + +def test_complex_dispatch_realpart(): + # Test that the real parts of loggamma and gammaln agree on the + # real axis. + x = np.r_[-np.logspace(10, -10), np.logspace(-10, 10)] + 0.5 + + dataset = np.vstack((x, gammaln(x))).T + + def f(z): + z = np.array(z, dtype='complex128') + return loggamma(z).real + + FuncData(f, dataset, 0, 1, rtol=1e-14, atol=1e-14).check() + + +def test_real_dispatch(): + x = np.logspace(-10, 10) + 0.5 + dataset = np.vstack((x, gammaln(x))).T + + FuncData(loggamma, dataset, 0, 1, rtol=1e-14, atol=1e-14).check() + assert_(loggamma(0) == np.inf) + assert_(np.isnan(loggamma(-1))) + + +def test_gh_6536(): + z = loggamma(complex(-3.4, +0.0)) + zbar = loggamma(complex(-3.4, -0.0)) + assert_allclose(z, zbar.conjugate(), rtol=1e-15, atol=0) + + +def test_branch_cut(): + # Make sure negative zero is treated correctly + x = -np.logspace(300, -30, 100) + z = np.asarray([complex(x0, 0.0) for x0 in x]) + zbar = np.asarray([complex(x0, -0.0) for x0 in x]) + assert_allclose(z, zbar.conjugate(), rtol=1e-15, atol=0) diff --git a/v12x/pylib/scipy/special/tests/test_logit.py b/v12x/pylib/scipy/special/tests/test_logit.py new file mode 100644 index 0000000000000000000000000000000000000000..ae8dd8851d8932dcbbda68c7f11edbe2daafcf00 --- /dev/null +++ b/v12x/pylib/scipy/special/tests/test_logit.py @@ -0,0 +1,161 @@ +import numpy as np +from numpy.testing import assert_equal, assert_allclose +from scipy.special import logit, expit, log_expit + + +class TestLogit: + + def check_logit_out(self, a, expected): + actual = logit(a) + assert_equal(actual.dtype, a.dtype) + rtol = 16*np.finfo(a.dtype).eps + assert_allclose(actual, expected, rtol=rtol) + + def test_float32(self): + a = np.concatenate((np.linspace(0, 1, 10, dtype=np.float32), + [np.float32(0.0001), np.float32(0.49999), + np.float32(0.50001)])) + # Expected values computed with mpmath from float32 inputs, e.g. + # from mpmath import mp + # mp.dps = 200 + # a = np.float32(1/9) + # print(np.float32(mp.log(a) - mp.log1p(-a))) + # prints `-2.0794415`. + expected = np.array([-np.inf, -2.0794415, -1.2527629, -6.9314712e-01, + -2.2314353e-01, 2.2314365e-01, 6.9314724e-01, + 1.2527630, 2.0794415, np.inf, + -9.2102404, -4.0054321e-05, 4.0054321e-05], + dtype=np.float32) + self.check_logit_out(a, expected) + + def test_float64(self): + a = np.concatenate((np.linspace(0, 1, 10, dtype=np.float64), + [1e-8, 0.4999999999999, 0.50000000001])) + # Expected values computed with mpmath. + expected = np.array([-np.inf, + -2.079441541679836, + -1.252762968495368, + -0.6931471805599454, + -0.22314355131420985, + 0.22314355131420985, + 0.6931471805599452, + 1.2527629684953674, + 2.0794415416798353, + np.inf, + -18.420680733952366, + -3.999023334699814e-13, + 4.000000330961484e-11]) + self.check_logit_out(a, expected) + + def test_nan(self): + expected = np.array([np.nan]*4) + with np.errstate(invalid='ignore'): + actual = logit(np.array([-3., -2., 2., 3.])) + + assert_equal(expected, actual) + + +class TestExpit: + def check_expit_out(self, dtype, expected): + a = np.linspace(-4, 4, 10) + a = np.array(a, dtype=dtype) + actual = expit(a) + assert_allclose(actual, expected, atol=1.5e-7, rtol=0) + assert_equal(actual.dtype, np.dtype(dtype)) + + def test_float32(self): + expected = np.array([0.01798621, 0.04265125, + 0.09777259, 0.20860852, + 0.39068246, 0.60931754, + 0.79139149, 0.9022274, + 0.95734876, 0.98201376], dtype=np.float32) + self.check_expit_out('f4', expected) + + def test_float64(self): + expected = np.array([0.01798621, 0.04265125, + 0.0977726, 0.20860853, + 0.39068246, 0.60931754, + 0.79139147, 0.9022274, + 0.95734875, 0.98201379]) + self.check_expit_out('f8', expected) + + def test_large(self): + for dtype in (np.float32, np.float64, np.longdouble): + for n in (88, 89, 709, 710, 11356, 11357): + n = np.array(n, dtype=dtype) + assert_allclose(expit(n), 1.0, atol=1e-20) + assert_allclose(expit(-n), 0.0, atol=1e-20) + assert_equal(expit(n).dtype, dtype) + assert_equal(expit(-n).dtype, dtype) + + +class TestLogExpit: + + def test_large_negative(self): + x = np.array([-10000.0, -750.0, -500.0, -35.0]) + y = log_expit(x) + assert_equal(y, x) + + def test_large_positive(self): + x = np.array([750.0, 1000.0, 10000.0]) + y = log_expit(x) + # y will contain -0.0, and -0.0 is used in the expected value, + # but assert_equal does not check the sign of zeros, and I don't + # think the sign is an essential part of the test (i.e. it would + # probably be OK if log_expit(1000) returned 0.0 instead of -0.0). + assert_equal(y, np.array([-0.0, -0.0, -0.0])) + + def test_basic_float64(self): + x = np.array([-32, -20, -10, -3, -1, -0.1, -1e-9, + 0, 1e-9, 0.1, 1, 10, 100, 500, 710, 725, 735]) + y = log_expit(x) + # + # Expected values were computed with mpmath: + # + # import mpmath + # + # mpmath.mp.dps = 100 + # + # def mp_log_expit(x): + # return -mpmath.log1p(mpmath.exp(-x)) + # + # expected = [float(mp_log_expit(t)) for t in x] + # + expected = [-32.000000000000014, -20.000000002061153, + -10.000045398899218, -3.048587351573742, + -1.3132616875182228, -0.7443966600735709, + -0.6931471810599453, -0.6931471805599453, + -0.6931471800599454, -0.6443966600735709, + -0.3132616875182228, -4.539889921686465e-05, + -3.720075976020836e-44, -7.124576406741286e-218, + -4.47628622567513e-309, -1.36930634e-315, + -6.217e-320] + + # When tested locally, only one value in y was not exactly equal to + # expected. That was for x=1, and the y value differed from the + # expected by 1 ULP. For this test, however, I'll use rtol=1e-15. + assert_allclose(y, expected, rtol=1e-15) + + def test_basic_float32(self): + x = np.array([-32, -20, -10, -3, -1, -0.1, -1e-9, + 0, 1e-9, 0.1, 1, 10, 100], dtype=np.float32) + y = log_expit(x) + # + # Expected values were computed with mpmath: + # + # import mpmath + # + # mpmath.mp.dps = 100 + # + # def mp_log_expit(x): + # return -mpmath.log1p(mpmath.exp(-x)) + # + # expected = [np.float32(mp_log_expit(t)) for t in x] + # + expected = np.array([-32.0, -20.0, -10.000046, -3.0485873, + -1.3132616, -0.7443967, -0.6931472, + -0.6931472, -0.6931472, -0.64439666, + -0.3132617, -4.5398898e-05, -3.8e-44], + dtype=np.float32) + + assert_allclose(y, expected, rtol=5e-7) diff --git a/v12x/pylib/scipy/special/tests/test_logsumexp.py b/v12x/pylib/scipy/special/tests/test_logsumexp.py new file mode 100644 index 0000000000000000000000000000000000000000..51196a3d5e3ec3376b4311c374d13bec37f37f00 --- /dev/null +++ b/v12x/pylib/scipy/special/tests/test_logsumexp.py @@ -0,0 +1,481 @@ +import itertools as it +import math +import pytest + +import numpy as np + +from scipy._lib._array_api import (is_array_api_strict, make_xp_test_case, + xp_default_dtype, xp_device) +from scipy._lib._array_api_no_0d import (xp_assert_equal, xp_assert_close, + xp_assert_less) + +from scipy.special import log_softmax, logsumexp, softmax +from scipy.special._logsumexp import _wrap_radians + + +dtypes = ['float32', 'float64', 'int32', 'int64', 'complex64', 'complex128'] +integral_dtypes = ['int32', 'int64'] + + +@pytest.mark.uses_xp_capabilities(False, reason="private") +def test_wrap_radians(xp): + x = xp.asarray([-math.pi-1, -math.pi, -1, -1e-300, + 0, 1e-300, 1, math.pi, math.pi+1]) + ref = xp.asarray([math.pi-1, math.pi, -1, -1e-300, + 0, 1e-300, 1, math.pi, -math.pi+1]) + res = _wrap_radians(x, xp=xp) + xp_assert_close(res, ref, atol=0) + + +# numpy warning filters don't work for dask (dask/dask#3245) +# (also we should not expect the numpy warning filter to work for any Array API +# library) +@pytest.mark.filterwarnings("ignore:invalid value encountered:RuntimeWarning") +@pytest.mark.filterwarnings("ignore:divide by zero encountered:RuntimeWarning") +@pytest.mark.filterwarnings("ignore:overflow encountered:RuntimeWarning") +@make_xp_test_case(logsumexp) +class TestLogSumExp: + def test_logsumexp(self, xp): + # Test with zero-size array + a = xp.asarray([]) + desired = xp.asarray(-xp.inf) + xp_assert_equal(logsumexp(a), desired) + + # Test whether logsumexp() function correctly handles large inputs. + a = xp.arange(200., dtype=xp.float64) + desired = xp.log(xp.sum(xp.exp(a))) + xp_assert_close(logsumexp(a), desired) + + # Now test with large numbers + b = xp.asarray([1000., 1000.]) + desired = xp.asarray(1000.0 + math.log(2.0)) + xp_assert_close(logsumexp(b), desired) + + n = 1000 + b = xp.full((n,), 10000) + desired = xp.asarray(10000.0 + math.log(n)) + xp_assert_close(logsumexp(b), desired) + + x = xp.asarray([1e-40] * 1000000) + logx = xp.log(x) + X = xp.stack([x, x]) + logX = xp.stack([logx, logx]) + xp_assert_close(xp.exp(logsumexp(logX)), xp.sum(X)) + xp_assert_close(xp.exp(logsumexp(logX, axis=0)), xp.sum(X, axis=0)) + xp_assert_close(xp.exp(logsumexp(logX, axis=1)), xp.sum(X, axis=1)) + + # Handling special values properly + inf = xp.asarray([xp.inf]) + nan = xp.asarray([xp.nan]) + xp_assert_equal(logsumexp(inf), inf[0]) + xp_assert_equal(logsumexp(-inf), -inf[0]) + xp_assert_equal(logsumexp(nan), nan[0]) + xp_assert_equal(logsumexp(xp.asarray([-xp.inf, -xp.inf])), -inf[0]) + + # Handling an array with different magnitudes on the axes + a = xp.asarray([[1e10, 1e-10], + [-1e10, -np.inf]]) + ref = xp.asarray([1e10, -1e10]) + xp_assert_close(logsumexp(a, axis=-1), ref) + + # Test keeping dimensions + ref = xp.expand_dims(ref, axis=-1) + xp_assert_close(logsumexp(a, axis=-1, keepdims=True), ref) + + # Test multiple axes + xp_assert_close(logsumexp(a, axis=(-1, -2)), xp.asarray(1e10)) + + def test_logsumexp_b(self, xp): + a = xp.arange(200., dtype=xp.float64) + b = xp.arange(200., 0., -1.) + desired = xp.log(xp.sum(b*xp.exp(a))) + xp_assert_close(logsumexp(a, b=b), desired) + + a = xp.asarray([1000, 1000]) + b = xp.asarray([1.2, 1.2]) + desired = xp.asarray(1000 + math.log(2 * 1.2)) + xp_assert_close(logsumexp(a, b=b), desired) + + x = xp.asarray([1e-40] * 100000) + b = xp.linspace(1, 1000, 100000) + logx = xp.log(x) + X = xp.stack((x, x)) + logX = xp.stack((logx, logx)) + B = xp.stack((b, b)) + xp_assert_close(xp.exp(logsumexp(logX, b=B)), xp.sum(B * X)) + xp_assert_close(xp.exp(logsumexp(logX, b=B, axis=0)), xp.sum(B * X, axis=0)) + xp_assert_close(xp.exp(logsumexp(logX, b=B, axis=1)), xp.sum(B * X, axis=1)) + + def test_logsumexp_sign(self, xp): + a = xp.asarray([1, 1, 1]) + b = xp.asarray([1, -1, -1]) + + r, s = logsumexp(a, b=b, return_sign=True) + xp_assert_close(r, xp.asarray(1.)) + xp_assert_equal(s, xp.asarray(-1.)) + + def test_logsumexp_sign_zero(self, xp): + a = xp.asarray([1, 1]) + b = xp.asarray([1, -1]) + + r, s = logsumexp(a, b=b, return_sign=True) + assert not xp.isfinite(r) + assert not xp.isnan(r) + assert r < 0 + assert s == 0 + + def test_logsumexp_sign_shape(self, xp): + a = xp.ones((1, 2, 3, 4)) + b = xp.ones_like(a) + + r, s = logsumexp(a, axis=2, b=b, return_sign=True) + assert r.shape == s.shape == (1, 2, 4) + + r, s = logsumexp(a, axis=(1, 3), b=b, return_sign=True) + assert r.shape == s.shape == (1,3) + + def test_logsumexp_complex_sign(self, xp): + a = xp.asarray([1 + 1j, 2 - 1j, -2 + 3j]) + + r, s = logsumexp(a, return_sign=True) + + expected_sumexp = xp.sum(xp.exp(a)) + # This is the numpy>=2.0 convention for np.sign + expected_sign = expected_sumexp / xp.abs(expected_sumexp) + + xp_assert_close(s, expected_sign) + xp_assert_close(s * xp.exp(r), expected_sumexp) + + def test_logsumexp_shape(self, xp): + a = xp.ones((1, 2, 3, 4)) + b = xp.ones_like(a) + + r = logsumexp(a, axis=2, b=b) + assert r.shape == (1, 2, 4) + + r = logsumexp(a, axis=(1, 3), b=b) + assert r.shape == (1, 3) + + def test_logsumexp_b_zero(self, xp): + a = xp.asarray([1, 10000]) + b = xp.asarray([1, 0]) + + xp_assert_close(logsumexp(a, b=b), xp.asarray(1.)) + + def test_logsumexp_b_shape(self, xp): + a = xp.zeros((4, 1, 2, 1)) + b = xp.ones((3, 1, 5)) + + logsumexp(a, b=b) + + @pytest.mark.parametrize('arg', (1, [1, 2, 3])) + def test_xp_invalid_input(self, arg): + assert logsumexp(arg) == logsumexp(np.asarray(np.atleast_1d(arg))) + + def test_array_like(self): + a = [1000, 1000] + desired = np.asarray(1000.0 + math.log(2.0)) + xp_assert_close(logsumexp(a), desired) + + @pytest.mark.parametrize('dtype', dtypes) + def test_dtypes_a(self, dtype, xp): + dtype = getattr(xp, dtype) + a = xp.asarray([1000., 1000.], dtype=dtype) + desired_dtype = (xp.asarray(1.).dtype if xp.isdtype(dtype, 'integral') + else dtype) # true for all libraries tested + desired = xp.asarray(1000.0 + math.log(2.0), dtype=desired_dtype) + xp_assert_close(logsumexp(a), desired) + + @pytest.mark.parametrize('dtype_a', dtypes) + @pytest.mark.parametrize('dtype_b', dtypes) + def test_dtypes_ab(self, dtype_a, dtype_b, xp): + xp_dtype_a = getattr(xp, dtype_a) + xp_dtype_b = getattr(xp, dtype_b) + a = xp.asarray([2, 1], dtype=xp_dtype_a) + b = xp.asarray([1, -1], dtype=xp_dtype_b) + if is_array_api_strict(xp): + # special-case for `TypeError: array_api_strict.float32 and + # and array_api_strict.int64 cannot be type promoted together` + xp_float_dtypes = [dtype for dtype in [xp_dtype_a, xp_dtype_b] + if not xp.isdtype(dtype, 'integral')] + if len(xp_float_dtypes) < 2: # at least one is integral + xp_float_dtypes.append(xp.asarray(1.).dtype) + desired_dtype = xp.result_type(*xp_float_dtypes) + else: + desired_dtype = xp.result_type(xp_dtype_a, xp_dtype_b) + if xp.isdtype(desired_dtype, 'integral'): + desired_dtype = xp_default_dtype(xp) + desired = xp.asarray(math.log(math.exp(2) - math.exp(1)), dtype=desired_dtype) + xp_assert_close(logsumexp(a, b=b), desired) + + def test_gh18295(self, xp): + # gh-18295 noted loss of precision when real part of one element is much + # larger than the rest. Check that this is resolved. + a = xp.asarray([0.0, -40.0]) + res = logsumexp(a) + ref = xp.logaddexp(a[0], a[1]) + xp_assert_close(res, ref) + + @pytest.mark.parametrize('dtype', ['complex64', 'complex128']) + def test_gh21610(self, xp, dtype): + # gh-21610 noted that `logsumexp` could return imaginary components + # outside the range (-pi, pi]. Check that this is resolved. + # While working on this, I noticed that all other tests passed even + # when the imaginary component of the result was zero. This suggested + # the need of a stronger test with imaginary dtype. + rng = np.random.default_rng(324984329582349862) + dtype = getattr(xp, dtype) + shape = (10, 100) + x = rng.uniform(1, 40, shape) + 1.j * rng.uniform(1, 40, shape) + x = xp.asarray(x, dtype=dtype) + + res = logsumexp(x, axis=1) + ref = xp.log(xp.sum(xp.exp(x), axis=1)) + max = xp.full_like(xp.imag(res), xp.pi) + xp_assert_less(xp.abs(xp.imag(res)), max) + xp_assert_close(res, ref) + + out, sgn = logsumexp(x, return_sign=True, axis=1) + ref = xp.sum(xp.exp(x), axis=1) + xp_assert_less(xp.abs(xp.imag(sgn)), max) + xp_assert_close(out, xp.real(xp.log(ref))) + xp_assert_close(sgn, ref/xp.abs(ref)) + + def test_gh21709_small_imaginary(self, xp): + # Test that `logsumexp` does not lose relative precision of + # small imaginary components + x = xp.asarray([0, 0.+2.2204460492503132e-17j]) + res = logsumexp(x) + # from mpmath import mp + # mp.dps = 100 + # x, y = mp.mpc(0), mp.mpc('0', '2.2204460492503132e-17') + # ref = complex(mp.log(mp.exp(x) + mp.exp(y))) + ref = xp.asarray(0.6931471805599453+1.1102230246251566e-17j) + xp_assert_close(xp.real(res), xp.real(ref)) + xp_assert_close(xp.imag(res), xp.imag(ref), atol=0, rtol=1e-15) + + + @pytest.mark.parametrize('x,y', list(it.product( + [ + -np.inf, + np.inf, + complex(-np.inf, 0.), + complex(-np.inf, -0.), + complex(-np.inf, np.inf), + complex(-np.inf, -np.inf), + complex(np.inf, 0.), + complex(np.inf, -0.), + complex(np.inf, np.inf), + complex(np.inf, -np.inf), + # Phase in each quadrant. + complex(-np.inf, 0.7533), + complex(-np.inf, 2.3562), + complex(-np.inf, 3.9270), + complex(-np.inf, 5.4978), + complex(np.inf, 0.7533), + complex(np.inf, 2.3562), + complex(np.inf, 3.9270), + complex(np.inf, 5.4978), + ], repeat=2) + )) + def test_gh22601_infinite_elements(self, x, y, xp): + # Test that `logsumexp` does reasonable things in the presence of + # real and complex infinities. + res = logsumexp(xp.asarray([x, y])) + ref = xp.log(xp.sum(xp.exp(xp.asarray([x, y])))) + xp_assert_equal(res, ref) + + def test_no_writeback(self, xp): + """Test that logsumexp doesn't accidentally write back to its parameters.""" + a = xp.asarray([5., 4.]) + b = xp.asarray([3., 2.]) + logsumexp(a) + logsumexp(a, b=b) + xp_assert_equal(a, xp.asarray([5., 4.])) + xp_assert_equal(b, xp.asarray([3., 2.])) + + @pytest.mark.parametrize("x_raw", [1.0, 1.0j, []]) + def test_device(self, x_raw, xp, devices): + """Test input device propagation to output.""" + for d in devices: + x = xp.asarray(x_raw, device=d) + assert xp_device(logsumexp(x)) == xp_device(x) + assert xp_device(logsumexp(x, b=x)) == xp_device(x) + + def test_gh22903(self, xp): + # gh-22903 reported that `logsumexp` produced NaN where the weight associated + # with the max magnitude element was negative and `return_sign=False`, even if + # the net result should be the log of a positive number. + + # result is log of positive number + a = xp.asarray([3.06409428, 0.37251854, 3.87471931]) + b = xp.asarray([1.88190708, 2.84174795, -0.85016884]) + xp_assert_close(logsumexp(a, b=b), logsumexp(a, b=b, return_sign=True)[0]) + + # result is log of negative number + b = xp.asarray([1.88190708, 2.84174795, -3.85016884]) + xp_assert_close(logsumexp(a, b=b), xp.asarray(xp.nan)) + + @pytest.mark.parametrize("a, b, sign_ref", + [([np.inf], None, 1.), + ([np.inf], [-1.], -1.)]) + def test_gh23548(self, xp, a, b, sign_ref): + # gh-23548 reported that `logsumexp` with `return_sign=True` returned a sign + # of NaN with infinite reals + a, b = xp.asarray(a), xp.asarray(b) if b is not None else None + val, sign = logsumexp(a, b=b, return_sign=True) + assert xp.isinf(val) + xp_assert_equal(sign, xp.asarray(sign_ref)) + + +@make_xp_test_case(softmax) +class TestSoftmax: + def test_softmax_fixtures(self, xp): + xp_assert_close(softmax(xp.asarray([1000., 0., 0., 0.])), + xp.asarray([1., 0., 0., 0.]), rtol=1e-13) + xp_assert_close(softmax(xp.asarray([1., 1.])), + xp.asarray([.5, .5]), rtol=1e-13) + xp_assert_close(softmax(xp.asarray([0., 1.])), + xp.asarray([1., np.e])/(1 + np.e), + rtol=1e-13) + + # Expected value computed using mpmath (with mpmath.mp.dps = 200) and then + # converted to float. + x = xp.arange(4, dtype=xp.float64) + expected = xp.asarray([0.03205860328008499, + 0.08714431874203256, + 0.23688281808991013, + 0.6439142598879722], dtype=xp.float64) + + xp_assert_close(softmax(x), expected, rtol=1e-13) + + # Translation property. If all the values are changed by the same amount, + # the softmax result does not change. + xp_assert_close(softmax(x + 100), expected, rtol=1e-13) + + # When axis=None, softmax operates on the entire array, and preserves + # the shape. + xp_assert_close(softmax(xp.reshape(x, (2, 2))), + xp.reshape(expected, (2, 2)), rtol=1e-13) + + def test_softmax_multi_axes(self, xp): + xp_assert_close(softmax(xp.asarray([[1000., 0.], [1000., 0.]]), axis=0), + xp.asarray([[.5, .5], [.5, .5]]), rtol=1e-13) + xp_assert_close(softmax(xp.asarray([[1000., 0.], [1000., 0.]]), axis=1), + xp.asarray([[1., 0.], [1., 0.]]), rtol=1e-13) + + # Expected value computed using mpmath (with mpmath.mp.dps = 200) and then + # converted to float. + x = xp.asarray([[-25., 0., 25., 50.], + [ 1., 325., 749., 750.]]) + expected = xp.asarray([[2.678636961770877e-33, + 1.9287498479371314e-22, + 1.3887943864771144e-11, + 0.999999999986112], + [0.0, + 1.9444526359919372e-185, + 0.2689414213699951, + 0.7310585786300048]]) + xp_assert_close(softmax(x, axis=1), expected, rtol=1e-13) + xp_assert_close(softmax(x.T, axis=0), expected.T, rtol=1e-13) + + # 3-d input, with a tuple for the axis. + x3d = xp.reshape(x, (2, 2, 2)) + xp_assert_close(softmax(x3d, axis=(1, 2)), + xp.reshape(expected, (2, 2, 2)), rtol=1e-13) + + @pytest.mark.xfail_xp_backends("array_api_strict", reason="int->float promotion") + def test_softmax_int_array(self, xp): + xp_assert_close(softmax(xp.asarray([1000, 0, 0, 0])), + xp.asarray([1., 0., 0., 0.]), rtol=1e-13) + + def test_softmax_scalar(self): + xp_assert_close(softmax(1000), np.asarray(1.), rtol=1e-13) + + def test_softmax_array_like(self): + xp_assert_close(softmax([1000, 0, 0, 0]), + np.asarray([1., 0., 0., 0.]), rtol=1e-13) + + +@make_xp_test_case(log_softmax) +class TestLogSoftmax: + def test_log_softmax_basic(self, xp): + xp_assert_close(log_softmax(xp.asarray([1000., 1.])), + xp.asarray([0., -999.]), rtol=1e-13) + + @pytest.mark.xfail_xp_backends("array_api_strict", reason="int->float promotion") + def test_log_softmax_int_array(self, xp): + xp_assert_close(log_softmax(xp.asarray([1000, 1])), + xp.asarray([0., -999.]), rtol=1e-13) + + def test_log_softmax_scalar(self): + xp_assert_close(log_softmax(1.0), 0.0, rtol=1e-13) + + def test_log_softmax_array_like(self): + xp_assert_close(log_softmax([1000, 1]), + np.asarray([0., -999.]), rtol=1e-13) + + @staticmethod + def data_1d(xp): + x = xp.arange(4, dtype=xp.float64) + # Expected value computed using mpmath (with mpmath.mp.dps = 200) + expect = [-3.4401896985611953, + -2.4401896985611953, + -1.4401896985611953, + -0.44018969856119533] + return x, xp.asarray(expect, dtype=xp.float64) + + @staticmethod + def data_2d(xp): + x = xp.reshape(xp.arange(8, dtype=xp.float64), (2, 4)) + + # Expected value computed using mpmath (with mpmath.mp.dps = 200) + expect = [[-3.4401896985611953, + -2.4401896985611953, + -1.4401896985611953, + -0.44018969856119533], + [-3.4401896985611953, + -2.4401896985611953, + -1.4401896985611953, + -0.44018969856119533]] + return x, xp.asarray(expect, dtype=xp.float64) + + @pytest.mark.parametrize("offset", [0, 100]) + def test_log_softmax_translation(self, offset, xp): + # Translation property. If all the values are changed by the same amount, + # the softmax result does not change. + x, expect = self.data_1d(xp) + x += offset + xp_assert_close(log_softmax(x), expect, rtol=1e-13) + + def test_log_softmax_noneaxis(self, xp): + # When axis=None, softmax operates on the entire array, and preserves + # the shape. + x, expect = self.data_1d(xp) + x = xp.reshape(x, (2, 2)) + expect = xp.reshape(expect, (2, 2)) + xp_assert_close(log_softmax(x), expect, rtol=1e-13) + + @pytest.mark.parametrize('axis_2d, expected_2d', [ + (0, np.log(0.5) * np.ones((2, 2))), + (1, [[0., -999.], [0., -999.]]), + ]) + def test_axes(self, axis_2d, expected_2d, xp): + x = xp.asarray([[1000., 1.], [1000., 1.]]) + xp_assert_close(log_softmax(x, axis=axis_2d), + xp.asarray(expected_2d, dtype=x.dtype), rtol=1e-13) + + def test_log_softmax_2d_axis1(self, xp): + x, expect = self.data_2d(xp) + xp_assert_close(log_softmax(x, axis=1), expect, rtol=1e-13) + + def test_log_softmax_2d_axis0(self, xp): + x, expect = self.data_2d(xp) + xp_assert_close(log_softmax(x.T, axis=0), expect.T, rtol=1e-13) + + def test_log_softmax_3d(self, xp): + # 3D input, with a tuple for the axis. + x, expect = self.data_2d(xp) + x = xp.reshape(x, (2, 2, 2)) + expect = xp.reshape(expect, (2, 2, 2)) + xp_assert_close(log_softmax(x, axis=(1, 2)), expect, rtol=1e-13) diff --git a/v12x/pylib/scipy/special/tests/test_mpmath.py b/v12x/pylib/scipy/special/tests/test_mpmath.py new file mode 100644 index 0000000000000000000000000000000000000000..8a2c5d9874eb1f80c211a2bfa1bcbf06502ac84b --- /dev/null +++ b/v12x/pylib/scipy/special/tests/test_mpmath.py @@ -0,0 +1,2160 @@ +""" +Test SciPy functions versus mpmath, if available. + +""" +import numpy as np +from numpy.testing import assert_, assert_allclose +from numpy import pi +import pytest +import itertools + +from scipy._external.packaging_version import version + +import scipy.special as sc +from scipy.special._testutils import ( + MissingModule, check_version, FuncData, + assert_func_equal) +from scipy.special._mptestutils import ( + Arg, FixedArg, ComplexArg, IntArg, assert_mpmath_equal, + nonfunctional_tooslow, trace_args, time_limited, exception_to_nan, + inf_to_nan) +from scipy.special._ufuncs import ( + _sinpi, _cospi, _lgam1p, _lanczos_sum_expg_scaled, _log1pmx, + _igam_fac) + +try: + import mpmath +except ImportError: + mpmath = MissingModule('mpmath') + +pytestmark = pytest.mark.thread_unsafe( + reason=("mpmath gmpy2 backend is not thread-safe, " + "see https://github.com/mpmath/mpmath/issues/974")) + +# ------------------------------------------------------------------------------ +# expi +# ------------------------------------------------------------------------------ + +@check_version(mpmath, '0.10') +def test_expi_complex(): + dataset = [] + for r in np.logspace(-99, 2, 10): + for p in np.linspace(0, 2*np.pi, 30): + z = r*np.exp(1j*p) + dataset.append((z, complex(mpmath.ei(z)))) + dataset = np.array(dataset, dtype=np.cdouble) + + FuncData(sc.expi, dataset, 0, 1).check() + + +# ------------------------------------------------------------------------------ +# expn +# ------------------------------------------------------------------------------ + +@check_version(mpmath, '0.19') +def test_expn_large_n(): + # Test the transition to the asymptotic regime of n. + dataset = [] + for n in [50, 51]: + for x in np.logspace(0, 4, 200): + with mpmath.workdps(100): + dataset.append((n, x, float(mpmath.expint(n, x)))) + dataset = np.asarray(dataset) + + FuncData(sc.expn, dataset, (0, 1), 2, rtol=1e-13).check() + +# ------------------------------------------------------------------------------ +# hyp0f1 +# ------------------------------------------------------------------------------ + + +@check_version(mpmath, '0.19') +def test_hyp0f1_gh5764(): + # Do a small and somewhat systematic test that runs quickly + dataset = [] + axis = [-99.5, -9.5, -0.5, 0.5, 9.5, 99.5] + for v in axis: + for x in axis: + for y in axis: + z = x + 1j*y + # mpmath computes the answer correctly at dps ~ 17 but + # fails for 20 < dps < 120 (uses a different method); + # set the dps high enough that this isn't an issue + with mpmath.workdps(120): + res = complex(mpmath.hyp0f1(v, z)) + dataset.append((v, z, res)) + dataset = np.array(dataset) + + FuncData(lambda v, z: sc.hyp0f1(v.real, z), dataset, (0, 1), 2, + rtol=1e-13).check() + + +@check_version(mpmath, '0.19') +def test_hyp0f1_gh_1609(): + # this is a regression test for gh-1609 + vv = np.linspace(150, 180, 21) + af = sc.hyp0f1(vv, 0.5) + mf = np.array([mpmath.hyp0f1(v, 0.5) for v in vv]) + assert_allclose(af, mf.astype(float), rtol=1e-12) + + +# ------------------------------------------------------------------------------ +# hyperu +# ------------------------------------------------------------------------------ + +@check_version(mpmath, '1.1.0') +def test_hyperu_around_0(): + dataset = [] + # DLMF 13.2.14-15 test points. + for n in np.arange(-5, 5): + for b in np.linspace(-5, 5, 20): + a = -n + dataset.append((a, b, 0, float(mpmath.hyperu(a, b, 0)))) + a = -n + b - 1 + dataset.append((a, b, 0, float(mpmath.hyperu(a, b, 0)))) + # DLMF 13.2.16-22 test points. + for a in [-10.5, -1.5, -0.5, 0, 0.5, 1, 10]: + for b in [-1.0, -0.5, 0, 0.5, 1, 1.5, 2, 2.5]: + dataset.append((a, b, 0, float(mpmath.hyperu(a, b, 0)))) + dataset = np.array(dataset) + + FuncData(sc.hyperu, dataset, (0, 1, 2), 3, rtol=1e-15, atol=5e-13).check() + + +# ------------------------------------------------------------------------------ +# hyp2f1 +# ------------------------------------------------------------------------------ + +@check_version(mpmath, '1.0.0') +def test_hyp2f1_strange_points(): + pts = [ + (2, -1, -1, 0.7), # expected: 2.4 + (2, -2, -2, 0.7), # expected: 3.87 + ] + pts += list(itertools.product([2, 1, -0.7, -1000], repeat=4)) + pts = [ + (a, b, c, x) for a, b, c, x in pts + if b == c and round(b) == b and b < 0 and b != -1000 + ] + kw = dict(eliminate=True) + dataset = [p + (float(mpmath.hyp2f1(*p, **kw)),) for p in pts] + dataset = np.array(dataset, dtype=np.float64) + + FuncData(sc.hyp2f1, dataset, (0,1,2,3), 4, rtol=1e-10).check() + + +@check_version(mpmath, '0.13') +def test_hyp2f1_real_some_points(): + pts = [ + (1, 2, 3, 0), + (1./3, 2./3, 5./6, 27./32), + (1./4, 1./2, 3./4, 80./81), + (2,-2, -3, 3), + (2, -3, -2, 3), + (2, -1.5, -1.5, 3), + (1, 2, 3, 0), + (0.7235, -1, -5, 0.3), + (0.25, 1./3, 2, 0.999), + (0.25, 1./3, 2, -1), + (2, 3, 5, 0.99), + (3./2, -0.5, 3, 0.99), + (2, 2.5, -3.25, 0.999), + (-8, 18.016500331508873, 10.805295997850628, 0.90875647507000001), + (-10, 900, -10.5, 0.99), + (-10, 900, 10.5, 0.99), + (-1, 2, 1, 1.0), + (-1, 2, 1, -1.0), + (-3, 13, 5, 1.0), + (-3, 13, 5, -1.0), + (0.5, 1 - 270.5, 1.5, 0.999**2), # from issue 1561 + ] + dataset = [p + (float(mpmath.hyp2f1(*p)),) for p in pts] + dataset = np.array(dataset, dtype=np.float64) + + with np.errstate(invalid='ignore'): + FuncData(sc.hyp2f1, dataset, (0,1,2,3), 4, rtol=1e-10).check() + + +@check_version(mpmath, '0.14') +def test_hyp2f1_some_points_2(): + # Taken from mpmath unit tests -- this point failed for mpmath 0.13 but + # was fixed in their SVN since then + pts = [ + (112, (51,10), (-9,10), -0.99999), + (10,-900,10.5,0.99), + (10,-900,-10.5,0.99), + ] + + def fev(x): + if isinstance(x, tuple): + return float(x[0]) / x[1] + else: + return x + + dataset = [tuple(map(fev, p)) + (float(mpmath.hyp2f1(*p)),) for p in pts] + dataset = np.array(dataset, dtype=np.float64) + + FuncData(sc.hyp2f1, dataset, (0,1,2,3), 4, rtol=1e-10).check() + + +@check_version(mpmath, '0.13') +def test_hyp2f1_real_some(): + dataset = [] + for a in [-10, -5, -1.8, 1.8, 5, 10]: + for b in [-2.5, -1, 1, 7.4]: + for c in [-9, -1.8, 5, 20.4]: + for z in [-10, -1.01, -0.99, 0, 0.6, 0.95, 1.5, 10]: + try: + v = float(mpmath.hyp2f1(a, b, c, z)) + except Exception: + continue + dataset.append((a, b, c, z, v)) + dataset = np.array(dataset, dtype=np.float64) + + with np.errstate(invalid='ignore'): + FuncData(sc.hyp2f1, dataset, (0,1,2,3), 4, rtol=1e-9, + ignore_inf_sign=True).check() + + +@check_version(mpmath, '0.12') +@pytest.mark.slow +def test_hyp2f1_real_random(): + npoints = 500 + dataset = np.zeros((npoints, 5), np.float64) + + np.random.seed(1234) + dataset[:, 0] = np.random.pareto(1.5, npoints) + dataset[:, 1] = np.random.pareto(1.5, npoints) + dataset[:, 2] = np.random.pareto(1.5, npoints) + dataset[:, 3] = 2*np.random.rand(npoints) - 1 + + dataset[:, 0] *= (-1)**np.random.randint(2, npoints) + dataset[:, 1] *= (-1)**np.random.randint(2, npoints) + dataset[:, 2] *= (-1)**np.random.randint(2, npoints) + + for ds in dataset: + if mpmath.__version__ < '0.14': + # mpmath < 0.14 fails for c too much smaller than a, b + if abs(ds[:2]).max() > abs(ds[2]): + ds[2] = abs(ds[:2]).max() + ds[4] = float(mpmath.hyp2f1(*tuple(ds[:4]))) + + FuncData(sc.hyp2f1, dataset, (0, 1, 2, 3), 4, rtol=1e-9).check() + + +# ------------------------------------------------------------------------------ +# erf (complex) +# ------------------------------------------------------------------------------ + +@pytest.mark.slow() +@check_version(mpmath, '0.14') +def test_erf_complex(): + # need to increase mpmath precision for this test + old_dps, old_prec = mpmath.mp.dps, mpmath.mp.prec + try: + mpmath.mp.dps = 70 + x1, y1 = np.meshgrid(np.linspace(-10, 1, 31), np.linspace(-10, 1, 11)) + x2, y2 = np.meshgrid(np.logspace(-80, .8, 31), np.logspace(-80, .8, 11)) + points = np.r_[x1.ravel(),x2.ravel()] + 1j*np.r_[y1.ravel(), y2.ravel()] + + assert_func_equal(sc.erf, lambda x: complex(mpmath.erf(x)), points, + vectorized=False, rtol=1e-13) + assert_func_equal(sc.erfc, lambda x: complex(mpmath.erfc(x)), points, + vectorized=False, rtol=1e-13) + finally: + mpmath.mp.dps, mpmath.mp.prec = old_dps, old_prec + + +# ------------------------------------------------------------------------------ +# lpmv +# ------------------------------------------------------------------------------ + +@check_version(mpmath, '0.15') +def test_lpmv(): + pts = [] + for x in [-0.99, -0.557, 1e-6, 0.132, 1]: + pts.extend([ + (1, 1, x), + (1, -1, x), + (-1, 1, x), + (-1, -2, x), + (1, 1.7, x), + (1, -1.7, x), + (-1, 1.7, x), + (-1, -2.7, x), + (1, 10, x), + (1, 11, x), + (3, 8, x), + (5, 11, x), + (-3, 8, x), + (-5, 11, x), + (3, -8, x), + (5, -11, x), + (-3, -8, x), + (-5, -11, x), + (3, 8.3, x), + (5, 11.3, x), + (-3, 8.3, x), + (-5, 11.3, x), + (3, -8.3, x), + (5, -11.3, x), + (-3, -8.3, x), + (-5, -11.3, x), + ]) + + def mplegenp(nu, mu, x): + if mu == int(mu) and x == 1: + # mpmath 0.17 gets this wrong + if mu == 0: + return 1 + else: + return 0 + return mpmath.legenp(nu, mu, x) + + dataset = [p + (mplegenp(p[1], p[0], p[2]),) for p in pts] + dataset = np.array(dataset, dtype=np.float64) + + def evf(mu, nu, x): + return sc.lpmv(mu.astype(int), nu, x) + + with np.errstate(invalid='ignore'): + FuncData(evf, dataset, (0,1,2), 3, rtol=1e-10, atol=1e-14).check() + + +# ------------------------------------------------------------------------------ +# beta +# ------------------------------------------------------------------------------ + +@pytest.mark.slow +@check_version(mpmath, '0.15') +def test_beta(): + np.random.seed(1234) + + b = np.r_[np.logspace(-200, 200, 4), + np.logspace(-10, 10, 4), + np.logspace(-1, 1, 4), + np.arange(-10, 11, 1), + np.arange(-10, 11, 1) + 0.5, + -1, -2.3, -3, -100.3, -10003.4] + a = b + + ab = np.array(np.broadcast_arrays(a[:,None], b[None,:])).reshape(2, -1).T + + old_dps, old_prec = mpmath.mp.dps, mpmath.mp.prec + try: + mpmath.mp.dps = 400 + + assert_func_equal(sc.beta, + lambda a, b: float(mpmath.beta(a, b)), + ab, + vectorized=False, + rtol=1e-10, + ignore_inf_sign=True) + + assert_func_equal( + sc.betaln, + lambda a, b: float(mpmath.log(abs(mpmath.beta(a, b)))), + ab, + vectorized=False, + rtol=1e-10) + finally: + mpmath.mp.dps, mpmath.mp.prec = old_dps, old_prec + + +# ------------------------------------------------------------------------------ +# loggamma +# ------------------------------------------------------------------------------ + +LOGGAMMA_TAYLOR_RADIUS = 0.2 + + +@check_version(mpmath, '0.19') +def test_loggamma_taylor_transition(): + # Make sure there isn't a big jump in accuracy when we move from + # using the Taylor series to using the recurrence relation. + + r = LOGGAMMA_TAYLOR_RADIUS + np.array([-0.1, -0.01, 0, 0.01, 0.1]) + theta = np.linspace(0, 2*np.pi, 20) + r, theta = np.meshgrid(r, theta) + dz = r*np.exp(1j*theta) + z = np.r_[1 + dz, 2 + dz].flatten() + + dataset = [(z0, complex(mpmath.loggamma(z0))) for z0 in z] + dataset = np.array(dataset) + + FuncData(sc.loggamma, dataset, 0, 1, rtol=5e-14).check() + + +@check_version(mpmath, '0.19') +def test_loggamma_taylor(): + # Test around the zeros at z = 1, 2. + + r = np.logspace(-16, np.log10(LOGGAMMA_TAYLOR_RADIUS), 10) + theta = np.linspace(0, 2*np.pi, 20) + r, theta = np.meshgrid(r, theta) + dz = r*np.exp(1j*theta) + z = np.r_[1 + dz, 2 + dz].flatten() + + dataset = [(z0, complex(mpmath.loggamma(z0))) for z0 in z] + dataset = np.array(dataset) + + FuncData(sc.loggamma, dataset, 0, 1, rtol=5e-14).check() + + +# ------------------------------------------------------------------------------ +# rgamma +# ------------------------------------------------------------------------------ + +@check_version(mpmath, '0.19') +@pytest.mark.slow +def test_rgamma_zeros(): + # Test around the zeros at z = 0, -1, -2, ..., -169. (After -169 we + # get values that are out of floating point range even when we're + # within 0.1 of the zero.) + + # Can't use too many points here or the test takes forever. + dx = np.r_[-np.logspace(-1, -13, 3), 0, np.logspace(-13, -1, 3)] + dy = dx.copy() + dx, dy = np.meshgrid(dx, dy) + dz = dx + 1j*dy + zeros = np.arange(0, -170, -1).reshape(1, 1, -1) + z = (zeros + np.dstack((dz,)*zeros.size)).flatten() + with mpmath.workdps(100): + dataset = [(z0, complex(mpmath.rgamma(z0))) for z0 in z] + + dataset = np.array(dataset) + FuncData(sc.rgamma, dataset, 0, 1, rtol=1e-12).check() + + +# ------------------------------------------------------------------------------ +# digamma +# ------------------------------------------------------------------------------ + +@check_version(mpmath, '0.19') +@pytest.mark.slow +def test_digamma_roots(): + # Test the special-cased roots for digamma. + root = mpmath.findroot(mpmath.digamma, 1.5) + roots = [float(root)] + root = mpmath.findroot(mpmath.digamma, -0.5) + roots.append(float(root)) + roots = np.array(roots) + + # If we test beyond a radius of 0.24 mpmath will take forever. + dx = np.r_[-0.24, -np.logspace(-1, -15, 10), 0, np.logspace(-15, -1, 10), 0.24] + dy = dx.copy() + dx, dy = np.meshgrid(dx, dy) + dz = dx + 1j*dy + z = (roots + np.dstack((dz,)*roots.size)).flatten() + with mpmath.workdps(30): + dataset = [(z0, complex(mpmath.digamma(z0))) for z0 in z] + + dataset = np.array(dataset) + FuncData(sc.digamma, dataset, 0, 1, rtol=1e-14).check() + + +@check_version(mpmath, '0.19') +def test_digamma_negreal(): + # Test digamma around the negative real axis. Don't do this in + # TestSystematic because the points need some jiggering so that + # mpmath doesn't take forever. + + digamma = exception_to_nan(mpmath.digamma) + + x = -np.logspace(300, -30, 100) + y = np.r_[-np.logspace(0, -3, 5), 0, np.logspace(-3, 0, 5)] + x, y = np.meshgrid(x, y) + z = (x + 1j*y).flatten() + + with mpmath.workdps(40): + dataset = [(z0, complex(digamma(z0))) for z0 in z] + dataset = np.asarray(dataset) + + FuncData(sc.digamma, dataset, 0, 1, rtol=1e-13).check() + + +@check_version(mpmath, '0.19') +def test_digamma_boundary(): + # Check that there isn't a jump in accuracy when we switch from + # using the asymptotic series to the reflection formula. + + x = -np.logspace(300, -30, 100) + y = np.array([-6.1, -5.9, 5.9, 6.1]) + x, y = np.meshgrid(x, y) + z = (x + 1j*y).flatten() + + with mpmath.workdps(30): + dataset = [(z0, complex(mpmath.digamma(z0))) for z0 in z] + dataset = np.asarray(dataset) + + FuncData(sc.digamma, dataset, 0, 1, rtol=1e-13).check() + + +# ------------------------------------------------------------------------------ +# gammainc +# ------------------------------------------------------------------------------ + +@check_version(mpmath, '0.19') +@pytest.mark.slow +def test_gammainc_boundary(): + # Test the transition to the asymptotic series. + small = 20 + a = np.linspace(0.5*small, 2*small, 50) + x = a.copy() + a, x = np.meshgrid(a, x) + a, x = a.flatten(), x.flatten() + with mpmath.workdps(100): + dataset = [(a0, x0, float(mpmath.gammainc(a0, b=x0, regularized=True))) + for a0, x0 in zip(a, x)] + dataset = np.array(dataset) + + FuncData(sc.gammainc, dataset, (0, 1), 2, rtol=1e-12).check() + + +# ------------------------------------------------------------------------------ +# spence +# ------------------------------------------------------------------------------ + +@check_version(mpmath, '0.19') +@pytest.mark.slow +def test_spence_circle(): + # The trickiest region for spence is around the circle |z - 1| = 1, + # so test that region carefully. + + def spence(z): + return complex(mpmath.polylog(2, 1 - z)) + + r = np.linspace(0.5, 1.5) + theta = np.linspace(0, 2*pi) + z = (1 + np.outer(r, np.exp(1j*theta))).flatten() + dataset = np.asarray([(z0, spence(z0)) for z0 in z]) + + FuncData(sc.spence, dataset, 0, 1, rtol=1e-14).check() + + +# ------------------------------------------------------------------------------ +# sinpi and cospi +# ------------------------------------------------------------------------------ + +@check_version(mpmath, '0.19') +def test_sinpi_zeros(): + eps = np.finfo(float).eps + dx = np.r_[-np.logspace(0, -13, 3), 0, np.logspace(-13, 0, 3)] + dy = dx.copy() + dx, dy = np.meshgrid(dx, dy) + dz = dx + 1j*dy + zeros = np.arange(-100, 100, 1).reshape(1, 1, -1) + z = (zeros + np.dstack((dz,)*zeros.size)).flatten() + dataset = np.asarray([(z0, complex(mpmath.sinpi(z0))) + for z0 in z]) + FuncData(_sinpi, dataset, 0, 1, rtol=2*eps).check() + + +@check_version(mpmath, '0.19') +def test_cospi_zeros(): + eps = np.finfo(float).eps + dx = np.r_[-np.logspace(0, -13, 3), 0, np.logspace(-13, 0, 3)] + dy = dx.copy() + dx, dy = np.meshgrid(dx, dy) + dz = dx + 1j*dy + zeros = (np.arange(-100, 100, 1) + 0.5).reshape(1, 1, -1) + z = (zeros + np.dstack((dz,)*zeros.size)).flatten() + dataset = np.asarray([(z0, complex(mpmath.cospi(z0))) + for z0 in z]) + + FuncData(_cospi, dataset, 0, 1, rtol=2*eps).check() + + +# ------------------------------------------------------------------------------ +# ellipj +# ------------------------------------------------------------------------------ + +@check_version(mpmath, '0.19') +def test_dn_quarter_period(): + def dn(u, m): + return sc.ellipj(u, m)[2] + + def mpmath_dn(u, m): + return float(mpmath.ellipfun("dn", u=u, m=m)) + + m = np.linspace(0, 1, 20) + du = np.r_[-np.logspace(-1, -15, 10), 0, np.logspace(-15, -1, 10)] + dataset = [] + for m0 in m: + u0 = float(mpmath.ellipk(m0)) + for du0 in du: + p = u0 + du0 + dataset.append((p, m0, mpmath_dn(p, m0))) + dataset = np.asarray(dataset) + + FuncData(dn, dataset, (0, 1), 2, rtol=1e-10).check() + + +# ------------------------------------------------------------------------------ +# Wright Omega +# ------------------------------------------------------------------------------ + +def _mpmath_wrightomega(z, dps): + with mpmath.workdps(dps): + z = mpmath.mpc(z) + unwind = mpmath.ceil((z.imag - mpmath.pi)/(2*mpmath.pi)) + res = mpmath.lambertw(mpmath.exp(z), unwind) + return res + + +@pytest.mark.slow +@check_version(mpmath, '0.19') +def test_wrightomega_branch(): + x = -np.logspace(10, 0, 25) + picut_above = [np.nextafter(np.pi, np.inf)] + picut_below = [np.nextafter(np.pi, -np.inf)] + npicut_above = [np.nextafter(-np.pi, np.inf)] + npicut_below = [np.nextafter(-np.pi, -np.inf)] + for i in range(50): + picut_above.append(np.nextafter(picut_above[-1], np.inf)) + picut_below.append(np.nextafter(picut_below[-1], -np.inf)) + npicut_above.append(np.nextafter(npicut_above[-1], np.inf)) + npicut_below.append(np.nextafter(npicut_below[-1], -np.inf)) + y = np.hstack((picut_above, picut_below, npicut_above, npicut_below)) + x, y = np.meshgrid(x, y) + z = (x + 1j*y).flatten() + + dataset = np.asarray([(z0, complex(_mpmath_wrightomega(z0, 25))) + for z0 in z]) + + FuncData(sc.wrightomega, dataset, 0, 1, rtol=1e-8).check() + + +@pytest.mark.slow +@check_version(mpmath, '0.19') +def test_wrightomega_region1(): + # This region gets less coverage in the TestSystematic test + x = np.linspace(-2, 1) + y = np.linspace(1, 2*np.pi) + x, y = np.meshgrid(x, y) + z = (x + 1j*y).flatten() + + dataset = np.asarray([(z0, complex(_mpmath_wrightomega(z0, 25))) + for z0 in z]) + + FuncData(sc.wrightomega, dataset, 0, 1, rtol=1e-15).check() + + +@pytest.mark.slow +@check_version(mpmath, '0.19') +def test_wrightomega_region2(): + # This region gets less coverage in the TestSystematic test + x = np.linspace(-2, 1) + y = np.linspace(-2*np.pi, -1) + x, y = np.meshgrid(x, y) + z = (x + 1j*y).flatten() + + dataset = np.asarray([(z0, complex(_mpmath_wrightomega(z0, 25))) + for z0 in z]) + + FuncData(sc.wrightomega, dataset, 0, 1, rtol=1e-15).check() + + +# ------------------------------------------------------------------------------ +# lambertw +# ------------------------------------------------------------------------------ + +@pytest.mark.slow +@check_version(mpmath, '0.19') +def test_lambertw_smallz(): + x, y = np.linspace(-1, 1, 25), np.linspace(-1, 1, 25) + x, y = np.meshgrid(x, y) + z = (x + 1j*y).flatten() + + dataset = np.asarray([(z0, complex(mpmath.lambertw(z0))) + for z0 in z]) + + FuncData(sc.lambertw, dataset, 0, 1, rtol=1e-13).check() + + +# ------------------------------------------------------------------------------ +# Systematic tests +# ------------------------------------------------------------------------------ + +HYPERKW = dict(maxprec=200, maxterms=200) + + +@pytest.mark.slow +@check_version(mpmath, '0.17') +class TestSystematic: + + def test_airyai(self): + # oscillating function, limit range + assert_mpmath_equal(lambda z: sc.airy(z)[0], + mpmath.airyai, + [Arg(-1e8, 1e8)], + rtol=1e-5) + assert_mpmath_equal(lambda z: sc.airy(z)[0], + mpmath.airyai, + [Arg(-1e3, 1e3)]) + + def test_airyai_complex(self): + assert_mpmath_equal(lambda z: sc.airy(z)[0], + mpmath.airyai, + [ComplexArg()]) + + def test_airyai_prime(self): + # oscillating function, limit range + assert_mpmath_equal(lambda z: sc.airy(z)[1], lambda z: + mpmath.airyai(z, derivative=1), + [Arg(-1e8, 1e8)], + rtol=1e-5) + assert_mpmath_equal(lambda z: sc.airy(z)[1], lambda z: + mpmath.airyai(z, derivative=1), + [Arg(-1e3, 1e3)]) + + def test_airyai_prime_complex(self): + assert_mpmath_equal(lambda z: sc.airy(z)[1], lambda z: + mpmath.airyai(z, derivative=1), + [ComplexArg()]) + + def test_airybi(self): + # oscillating function, limit range + assert_mpmath_equal(lambda z: sc.airy(z)[2], lambda z: + mpmath.airybi(z), + [Arg(-1e8, 1e8)], + rtol=1e-5) + assert_mpmath_equal(lambda z: sc.airy(z)[2], lambda z: + mpmath.airybi(z), + [Arg(-1e3, 1e3)]) + + def test_airybi_complex(self): + assert_mpmath_equal(lambda z: sc.airy(z)[2], lambda z: + mpmath.airybi(z), + [ComplexArg()]) + + def test_airybi_prime(self): + # oscillating function, limit range + assert_mpmath_equal(lambda z: sc.airy(z)[3], lambda z: + mpmath.airybi(z, derivative=1), + [Arg(-1e8, 1e8)], + rtol=1e-5) + assert_mpmath_equal(lambda z: sc.airy(z)[3], lambda z: + mpmath.airybi(z, derivative=1), + [Arg(-1e3, 1e3)]) + + def test_airybi_prime_complex(self): + assert_mpmath_equal(lambda z: sc.airy(z)[3], lambda z: + mpmath.airybi(z, derivative=1), + [ComplexArg()]) + + def test_bei(self): + assert_mpmath_equal(sc.bei, + exception_to_nan(lambda z: mpmath.bei(0, z, **HYPERKW)), + [Arg(-1e3, 1e3)]) + + def test_ber(self): + assert_mpmath_equal(sc.ber, + exception_to_nan(lambda z: mpmath.ber(0, z, **HYPERKW)), + [Arg(-1e3, 1e3)]) + + def test_bernoulli(self): + assert_mpmath_equal(lambda n: sc.bernoulli(int(n))[int(n)], + lambda n: float(mpmath.bernoulli(int(n))), + [IntArg(0, 13000)], + rtol=1e-9, n=13000) + + def test_besseli(self): + assert_mpmath_equal( + sc.iv, + exception_to_nan(lambda v, z: mpmath.besseli(v, z, **HYPERKW)), + [Arg(-1e100, 1e100), Arg()], + atol=1e-270, + ) + + def test_besseli_complex(self): + assert_mpmath_equal( + lambda v, z: sc.iv(v.real, z), + exception_to_nan(lambda v, z: mpmath.besseli(v, z, **HYPERKW)), + [Arg(-1e100, 1e100), ComplexArg()], + ) + + def test_besselj(self): + assert_mpmath_equal( + sc.jv, + exception_to_nan(lambda v, z: mpmath.besselj(v, z, **HYPERKW)), + [Arg(-1e100, 1e100), Arg(-1e3, 1e3)], + ignore_inf_sign=True, + ) + + # loss of precision at large arguments due to oscillation + assert_mpmath_equal( + sc.jv, + exception_to_nan(lambda v, z: mpmath.besselj(v, z, **HYPERKW)), + [Arg(-1e100, 1e100), Arg(-1e8, 1e8)], + ignore_inf_sign=True, + rtol=1e-5, + ) + + def test_besselj_complex(self): + assert_mpmath_equal( + lambda v, z: sc.jv(v.real, z), + exception_to_nan(lambda v, z: mpmath.besselj(v, z, **HYPERKW)), + [Arg(), ComplexArg()] + ) + + def test_besselk(self): + assert_mpmath_equal( + sc.kv, + mpmath.besselk, + [Arg(-200, 200), Arg(0, np.inf)], + nan_ok=False, + rtol=1e-12, + ) + + def test_besselk_int(self): + assert_mpmath_equal( + sc.kn, + mpmath.besselk, + [IntArg(-200, 200), Arg(0, np.inf)], + nan_ok=False, + rtol=1e-12, + ) + + def test_besselk_complex(self): + assert_mpmath_equal( + lambda v, z: sc.kv(v.real, z), + exception_to_nan(lambda v, z: mpmath.besselk(v, z, **HYPERKW)), + [Arg(-1e100, 1e100), ComplexArg()], + ) + + def test_bessely(self): + def mpbessely(v, x): + r = float(mpmath.bessely(v, x, **HYPERKW)) + if abs(r) > 1e305: + # overflowing to inf a bit earlier is OK + r = np.inf * np.sign(r) + if abs(r) == 0 and x == 0: + # invalid result from mpmath, point x=0 is a divergence + return np.nan + return r + assert_mpmath_equal( + sc.yv, + exception_to_nan(mpbessely), + [Arg(-1e100, 1e100), Arg(-1e8, 1e8)], + n=5000, + ) + + def test_bessely_complex(self): + def mpbessely(v, x): + r = complex(mpmath.bessely(v, x, **HYPERKW)) + if abs(r) > 1e305: + # overflowing to inf a bit earlier is OK + with np.errstate(invalid='ignore'): + r = np.inf * np.sign(r) + return r + assert_mpmath_equal( + lambda v, z: sc.yv(v.real, z), + exception_to_nan(mpbessely), + [Arg(), ComplexArg()], + n=15000, + ) + + def test_bessely_int(self): + def mpbessely(v, x): + r = float(mpmath.bessely(v, x)) + if abs(r) == 0 and x == 0: + # invalid result from mpmath, point x=0 is a divergence + return np.nan + return r + assert_mpmath_equal( + lambda v, z: sc.yn(int(v), z), + exception_to_nan(mpbessely), + [IntArg(-1000, 1000), Arg(-1e8, 1e8)], + ) + + def test_beta(self): + bad_points = [] + + def beta(a, b, nonzero=False): + if a < -1e12 or b < -1e12: + # Function is defined here only at integers, but due + # to loss of precision this is numerically + # ill-defined. Don't compare values here. + return np.nan + if (a < 0 or b < 0) and (abs(float(a + b)) % 1) == 0: + # close to a zero of the function: mpmath and scipy + # will not round here the same, so the test needs to be + # run with an absolute tolerance + if nonzero: + bad_points.append((float(a), float(b))) + return np.nan + return mpmath.beta(a, b) + + assert_mpmath_equal( + sc.beta, + lambda a, b: beta(a, b, nonzero=True), + [Arg(), Arg()], + dps=400, + ignore_inf_sign=True, + ) + + assert_mpmath_equal( + sc.beta, + beta, + np.array(bad_points), + dps=400, + ignore_inf_sign=True, + atol=1e-11, + ) + + def test_betainc(self): + assert_mpmath_equal( + sc.betainc, + time_limited()( + exception_to_nan( + lambda a, b, x: mpmath.betainc(a, b, 0, x, regularized=True) + ) + ), + [Arg(), Arg(), Arg()], + ) + + def test_betaincc(self): + assert_mpmath_equal( + sc.betaincc, + time_limited()( + exception_to_nan( + lambda a, b, x: mpmath.betainc(a, b, x, 1, regularized=True) + ) + ), + [Arg(), Arg(), Arg()], + dps=400, + ) + + def test_binom(self): + bad_points = [] + + def binomial(n, k, nonzero=False): + if abs(k) > 1e8*(abs(n) + 1): + # The binomial is rapidly oscillating in this region, + # and the function is numerically ill-defined. Don't + # compare values here. + return np.nan + if n < k and abs(float(n-k) - np.round(float(n-k))) < 1e-15: + # close to a zero of the function: mpmath and scipy + # will not round here the same, so the test needs to be + # run with an absolute tolerance + if nonzero: + bad_points.append((float(n), float(k))) + return np.nan + return mpmath.binomial(n, k) + + assert_mpmath_equal( + sc.binom, + lambda n, k: binomial(n, k, nonzero=True), + [Arg(), Arg()], + dps=400, + ) + + assert_mpmath_equal( + sc.binom, + binomial, + np.array(bad_points), + dps=400, + atol=1e-14, + ) + + def test_chebyt_int(self): + assert_mpmath_equal( + lambda n, x: sc.eval_chebyt(int(n), x), + exception_to_nan(lambda n, x: mpmath.chebyt(n, x, **HYPERKW)), + [IntArg(), Arg()], + dps=50, + ) + + @pytest.mark.xfail(run=False, reason="some cases in hyp2f1 not fully accurate") + def test_chebyt(self): + assert_mpmath_equal( + sc.eval_chebyt, + lambda n, x: time_limited()( + exception_to_nan(mpmath.chebyt) + )(n, x, **HYPERKW), + [Arg(-101, 101), Arg()], + n=10000, + ) + + def test_chebyu_int(self): + assert_mpmath_equal( + lambda n, x: sc.eval_chebyu(int(n), x), + exception_to_nan(lambda n, x: mpmath.chebyu(n, x, **HYPERKW)), + [IntArg(), Arg()], + dps=50, + ) + + @pytest.mark.xfail(run=False, reason="some cases in hyp2f1 not fully accurate") + def test_chebyu(self): + assert_mpmath_equal( + sc.eval_chebyu, + lambda n, x: time_limited()( + exception_to_nan(mpmath.chebyu) + )(n, x, **HYPERKW), + [Arg(-101, 101), Arg()], + ) + + def test_chi(self): + def chi(x): + return sc.shichi(x)[1] + assert_mpmath_equal(chi, mpmath.chi, [Arg()]) + # check asymptotic series cross-over + assert_mpmath_equal(chi, mpmath.chi, [FixedArg([88 - 1e-9, 88, 88 + 1e-9])]) + + def test_chi_complex(self): + def chi(z): + return sc.shichi(z)[1] + # chi oscillates as Im[z] -> +- inf, so limit range + assert_mpmath_equal( + chi, + mpmath.chi, + [ComplexArg(complex(-np.inf, -1e8), complex(np.inf, 1e8))], + rtol=1e-12, + ) + + def test_ci(self): + def ci(x): + return sc.sici(x)[1] + # oscillating function: limit range + assert_mpmath_equal(ci, mpmath.ci, [Arg(-1e8, 1e8)]) + + def test_ci_complex(self): + def ci(z): + return sc.sici(z)[1] + # ci oscillates as Re[z] -> +- inf, so limit range + assert_mpmath_equal( + ci, + mpmath.ci, + [ComplexArg(complex(-1e8, -np.inf), complex(1e8, np.inf))], + rtol=1e-8, + ) + + def test_cospi(self): + eps = np.finfo(float).eps + assert_mpmath_equal(_cospi, mpmath.cospi, [Arg()], nan_ok=False, rtol=2*eps) + + def test_cospi_complex(self): + assert_mpmath_equal( + _cospi, + mpmath.cospi, + [ComplexArg()], + nan_ok=False, + rtol=1e-13, + ) + + def test_digamma(self): + assert_mpmath_equal( + sc.digamma, + exception_to_nan(mpmath.digamma), + [Arg()], + rtol=1e-12, + dps=50, + ) + + def test_digamma_complex(self): + # Test on a cut plane because mpmath will hang. See + # test_digamma_negreal for tests on the negative real axis. + def param_filter(z): + return np.where((z.real < 0) & (np.abs(z.imag) < 1.12), False, True) + + assert_mpmath_equal( + sc.digamma, + exception_to_nan(mpmath.digamma), + [ComplexArg()], + rtol=1e-13, + dps=40, + param_filter=param_filter + ) + + def test_e1(self): + assert_mpmath_equal( + sc.exp1, + mpmath.e1, + [Arg()], + rtol=1e-14, + ) + + def test_e1_complex(self): + # E_1 oscillates as Im[z] -> +- inf, so limit range + assert_mpmath_equal( + sc.exp1, + mpmath.e1, + [ComplexArg(complex(-np.inf, -1e8), complex(np.inf, 1e8))], + rtol=1e-11, + ) + + # Check cross-over region + assert_mpmath_equal( + sc.exp1, + mpmath.e1, + (np.linspace(-50, 50, 171)[:, None] + + np.r_[0, np.logspace(-3, 2, 61), -np.logspace(-3, 2, 11)]*1j).ravel(), + rtol=1e-11, + ) + assert_mpmath_equal( + sc.exp1, + mpmath.e1, + (np.linspace(-50, -35, 10000) + 0j), + rtol=1e-11, + ) + + def test_exprel(self): + assert_mpmath_equal( + sc.exprel, + lambda x: mpmath.expm1(x)/x if x != 0 else mpmath.mpf('1.0'), + [Arg(a=-np.log(np.finfo(np.float64).max), + b=np.log(np.finfo(np.float64).max))], + ) + assert_mpmath_equal( + sc.exprel, + lambda x: mpmath.expm1(x)/x if x != 0 else mpmath.mpf('1.0'), + np.array([1e-12, 1e-24, 0, 1e12, 1e24, np.inf]), + rtol=1e-11, + ) + assert_(np.isinf(sc.exprel(np.inf))) + assert_(sc.exprel(-np.inf) == 0) + + def test_expm1_complex(self): + # Oscillates as a function of Im[z], so limit range to avoid loss of precision + assert_mpmath_equal( + sc.expm1, + mpmath.expm1, + [ComplexArg(complex(-np.inf, -1e7), complex(np.inf, 1e7))], + ) + + def test_log1p_complex(self): + assert_mpmath_equal( + sc.log1p, + lambda x: mpmath.log(x+1), + [ComplexArg()], + dps=60, + ) + + def test_log1pmx(self): + assert_mpmath_equal( + _log1pmx, + lambda x: mpmath.log(x + 1) - x, + [Arg()], + dps=60, + rtol=1e-14, + ) + + def test_ei(self): + assert_mpmath_equal(sc.expi, mpmath.ei, [Arg()], rtol=1e-11) + + def test_ei_complex(self): + # Ei oscillates as Im[z] -> +- inf, so limit range + assert_mpmath_equal( + sc.expi, + mpmath.ei, + [ComplexArg(complex(-np.inf, -1e8), complex(np.inf, 1e8))], + rtol=1e-9, + ) + + def test_ellipe(self): + assert_mpmath_equal(sc.ellipe, mpmath.ellipe, [Arg(b=1.0)]) + + def test_ellipeinc(self): + assert_mpmath_equal(sc.ellipeinc, mpmath.ellipe, [Arg(-1e3, 1e3), Arg(b=1.0)]) + + def test_ellipeinc_largephi(self): + assert_mpmath_equal(sc.ellipeinc, mpmath.ellipe, [Arg(), Arg()]) + + def test_ellipf(self): + assert_mpmath_equal(sc.ellipkinc, mpmath.ellipf, [Arg(-1e3, 1e3), Arg()]) + + def test_ellipf_largephi(self): + assert_mpmath_equal(sc.ellipkinc, mpmath.ellipf, [Arg(), Arg()]) + + def test_ellipk(self): + assert_mpmath_equal(sc.ellipk, mpmath.ellipk, [Arg(b=1.0)]) + assert_mpmath_equal( + sc.ellipkm1, + lambda m: mpmath.ellipk(1 - m), + [Arg(a=0.0)], + dps=400, + ) + + def test_ellipkinc(self): + def ellipkinc(phi, m): + return mpmath.ellippi(0, phi, m) + assert_mpmath_equal( + sc.ellipkinc, + ellipkinc, + [Arg(-1e3, 1e3), Arg(b=1.0)], + ignore_inf_sign=True, + ) + + def test_ellipkinc_largephi(self): + def ellipkinc(phi, m): + return mpmath.ellippi(0, phi, m) + assert_mpmath_equal( + sc.ellipkinc, + ellipkinc, + [Arg(), Arg(b=1.0)], + ignore_inf_sign=True, + ) + + def test_ellipfun_sn(self): + def sn(u, m): + # mpmath doesn't get the zero at u = 0--fix that + if u == 0: + return 0 + else: + return mpmath.ellipfun("sn", u=u, m=m) + + # Oscillating function --- limit range of first argument; the + # loss of precision there is an expected numerical feature + # rather than an actual bug + assert_mpmath_equal( + lambda u, m: sc.ellipj(u, m)[0], + sn, + [Arg(-1e6, 1e6), Arg(a=0, b=1)], + rtol=1e-8, + ) + + def test_ellipfun_cn(self): + # see comment in ellipfun_sn + assert_mpmath_equal( + lambda u, m: sc.ellipj(u, m)[1], + lambda u, m: mpmath.ellipfun("cn", u=u, m=m), + [Arg(-1e6, 1e6), Arg(a=0, b=1)], + rtol=1e-8, + ) + + def test_ellipfun_dn(self): + # see comment in ellipfun_sn + assert_mpmath_equal( + lambda u, m: sc.ellipj(u, m)[2], + lambda u, m: mpmath.ellipfun("dn", u=u, m=m), + [Arg(-1e6, 1e6), Arg(a=0, b=1)], + rtol=1e-8, + ) + + def test_erf(self): + assert_mpmath_equal(sc.erf, lambda z: mpmath.erf(z), [Arg()]) + + def test_erf_complex(self): + assert_mpmath_equal(sc.erf, lambda z: mpmath.erf(z), [ComplexArg()], n=200) + + def test_erfc(self): + assert_mpmath_equal( + sc.erfc, + exception_to_nan(lambda z: mpmath.erfc(z)), + [Arg()], + rtol=1e-13, + ) + + def test_erfc_complex(self): + assert_mpmath_equal( + sc.erfc, + exception_to_nan(lambda z: mpmath.erfc(z)), + [ComplexArg()], + n=200, + ) + + def test_erfi(self): + assert_mpmath_equal(sc.erfi, mpmath.erfi, [Arg()], n=200) + + def test_erfi_complex(self): + assert_mpmath_equal(sc.erfi, mpmath.erfi, [ComplexArg()], n=200) + + def test_ndtr(self): + assert_mpmath_equal( + sc.ndtr, + exception_to_nan(lambda z: mpmath.ncdf(z)), + [Arg()], + n=200, + ) + + def test_ndtr_complex(self): + assert_mpmath_equal( + sc.ndtr, + lambda z: mpmath.erfc(-z/np.sqrt(2.))/2., + [ComplexArg(a=complex(-10000, -10000), b=complex(10000, 10000))], + n=400, + ) + + def test_log_ndtr(self): + assert_mpmath_equal( + sc.log_ndtr, + exception_to_nan(lambda z: mpmath.log(mpmath.ncdf(z))), + [Arg()], n=600, dps=300, rtol=1e-13, + ) + + def test_log_ndtr_complex(self): + assert_mpmath_equal( + sc.log_ndtr, + exception_to_nan(lambda z: mpmath.log(mpmath.erfc(-z/np.sqrt(2.))/2.)), + [ComplexArg(a=complex(-10000, -100), b=complex(10000, 100))], + n=200, dps=300, + ) + + def test_eulernum(self): + assert_mpmath_equal( + lambda n: sc.euler(n)[-1], + mpmath.eulernum, + [IntArg(1, 10000)], + n=10000, + ) + + def test_expint(self): + assert_mpmath_equal( + sc.expn, + mpmath.expint, + [IntArg(0, 200), Arg(0, np.inf)], + rtol=1e-13, + dps=160, + ) + + def test_fresnels(self): + def fresnels(x): + return sc.fresnel(x)[0] + assert_mpmath_equal(fresnels, mpmath.fresnels, [Arg()]) + + def test_fresnelc(self): + def fresnelc(x): + return sc.fresnel(x)[1] + assert_mpmath_equal(fresnelc, mpmath.fresnelc, [Arg()]) + + def test_gamma(self): + assert_mpmath_equal(sc.gamma, exception_to_nan(mpmath.gamma), [Arg()]) + + def test_gamma_complex(self): + assert_mpmath_equal( + sc.gamma, + exception_to_nan(mpmath.gamma), + [ComplexArg()], + rtol=5e-13, + ) + + def test_gammainc(self): + # Larger arguments are tested in test_data.py:test_local + assert_mpmath_equal( + sc.gammainc, + lambda z, b: mpmath.gammainc(z, b=b, regularized=True), + [Arg(0, 1e4, inclusive_a=False), Arg(0, 1e4)], + nan_ok=False, + rtol=1e-11, + ) + + def test_gammaincc(self): + # Larger arguments are tested in test_data.py:test_local + assert_mpmath_equal( + sc.gammaincc, + lambda z, a: mpmath.gammainc(z, a=a, regularized=True), + [Arg(0, 1e4, inclusive_a=False), Arg(0, 1e4)], + nan_ok=False, + rtol=1e-11, + ) + + def test_gammaln(self): + # The real part of loggamma is log(|gamma(z)|). + def f(z): + return mpmath.loggamma(z).real + + assert_mpmath_equal(sc.gammaln, exception_to_nan(f), [Arg()]) + + @pytest.mark.xfail(run=False) + def test_gegenbauer(self): + assert_mpmath_equal( + sc.eval_gegenbauer, + exception_to_nan(mpmath.gegenbauer), + [Arg(-1e3, 1e3), Arg(), Arg()], + ) + + def test_gegenbauer_int(self): + # Redefine functions to deal with numerical + mpmath issues + def gegenbauer(n, a, x): + # Avoid overflow at large `a` (mpmath would need an even larger + # dps to handle this correctly, so just skip this region) + if abs(a) > 1e100: + return np.nan + + # Deal with n=0, n=1 correctly; mpmath 0.17 doesn't do these + # always correctly + if n == 0: + r = 0.0 if a == 0.0 else 1.0 + elif n == 1: + r = 2*a*x + else: + r = mpmath.gegenbauer(n, a, x) + + # Mpmath 0.17 gives wrong results (spurious zero) in some cases, so + # compute the value by perturbing the result + if float(r) == 0 and a < -1 and float(a) == int(float(a)): + r = mpmath.gegenbauer(n, a + mpmath.mpf('1e-50'), x) + if abs(r) < mpmath.mpf('1e-50'): + r = mpmath.mpf('0.0') + + # Differing overflow thresholds in scipy vs. mpmath + if abs(r) > 1e270: + return np.inf + return r + + def sc_gegenbauer(n, a, x): + r = sc.eval_gegenbauer(int(n), a, x) + # Differing overflow thresholds in scipy vs. mpmath + if abs(r) > 1e270: + return np.inf + return r + assert_mpmath_equal( + sc_gegenbauer, + exception_to_nan(gegenbauer), + [IntArg(0, 100), Arg(-1e9, 1e9), Arg()], + dps=100, ignore_inf_sign=True, rtol=1e-6, + ) + + # Check the small-x expansion + assert_mpmath_equal( + sc_gegenbauer, + exception_to_nan(gegenbauer), + [IntArg(0, 100), Arg(), FixedArg(np.logspace(-30, -4, 30))], + dps=100, ignore_inf_sign=True, + ) + + @pytest.mark.xfail(run=False) + def test_gegenbauer_complex(self): + assert_mpmath_equal( + lambda n, a, x: sc.eval_gegenbauer(int(n), a.real, x), + exception_to_nan(mpmath.gegenbauer), + [IntArg(0, 100), Arg(), ComplexArg()], + ) + + @nonfunctional_tooslow + def test_gegenbauer_complex_general(self): + assert_mpmath_equal( + lambda n, a, x: sc.eval_gegenbauer(n.real, a.real, x), + exception_to_nan(mpmath.gegenbauer), + [Arg(-1e3, 1e3), Arg(), ComplexArg()], + ) + + def test_hankel1(self): + assert_mpmath_equal( + sc.hankel1, + exception_to_nan(lambda v, x: mpmath.hankel1(v, x, **HYPERKW)), + [Arg(-1e20, 1e20), Arg()], + ) + + def test_hankel2(self): + assert_mpmath_equal( + sc.hankel2, + exception_to_nan(lambda v, x: mpmath.hankel2(v, x, **HYPERKW)), + [Arg(-1e20, 1e20), Arg()], + ) + + @pytest.mark.xfail(run=False, reason="issues at intermediately large orders") + def test_hermite(self): + assert_mpmath_equal( + lambda n, x: sc.eval_hermite(int(n), x), + exception_to_nan(mpmath.hermite), + [IntArg(0, 10000), Arg()], + ) + + # hurwitz: same as zeta + + def test_hyp0f1(self): + # mpmath reports no convergence unless maxterms is large enough + KW = dict(maxprec=400, maxterms=1500) + # n=500 (non-xslow default) fails for one bad point + assert_mpmath_equal( + sc.hyp0f1, + lambda a, x: mpmath.hyp0f1(a, x, **KW), + [Arg(-1e7, 1e7), Arg(0, 1e5)], + n=5000, + ) + # NB: The range of the second parameter ("z") is limited from below + # because of an overflow in the intermediate calculations. The way + # for fix it is to implement an asymptotic expansion for Bessel J + # (similar to what is implemented for Bessel I here). + + def test_hyp0f1_complex(self): + assert_mpmath_equal( + lambda a, z: sc.hyp0f1(a.real, z), + exception_to_nan(lambda a, x: mpmath.hyp0f1(a, x, **HYPERKW)), + [Arg(-10, 10), ComplexArg(complex(-120, -120), complex(120, 120))], + ) + # NB: The range of the first parameter ("v") are limited by an overflow + # in the intermediate calculations. Can be fixed by implementing an + # asymptotic expansion for Bessel functions for large order. + + def test_hyp1f1(self): + def mpmath_hyp1f1(a, b, x): + try: + return mpmath.hyp1f1(a, b, x) + except ZeroDivisionError: + return np.inf + + assert_mpmath_equal( + sc.hyp1f1, + mpmath_hyp1f1, + [Arg(-50, 50), Arg(1, 50, inclusive_a=False), Arg(-50, 50)], + n=500, + nan_ok=False, + ) + + @pytest.mark.xfail(run=False) + def test_hyp1f1_complex(self): + assert_mpmath_equal( + inf_to_nan(lambda a, b, x: sc.hyp1f1(a.real, b.real, x)), + exception_to_nan(lambda a, b, x: mpmath.hyp1f1(a, b, x, **HYPERKW)), + [Arg(-1e3, 1e3), Arg(-1e3, 1e3), ComplexArg()], + n=2000, + ) + + @nonfunctional_tooslow + def test_hyp2f1_complex(self): + # SciPy's hyp2f1 seems to have performance and accuracy problems + assert_mpmath_equal( + lambda a, b, c, x: sc.hyp2f1(a.real, b.real, c.real, x), + exception_to_nan(lambda a, b, c, x: mpmath.hyp2f1(a, b, c, x, **HYPERKW)), + [Arg(-1e2, 1e2), Arg(-1e2, 1e2), Arg(-1e2, 1e2), ComplexArg()], + n=10, + ) + + @pytest.mark.xfail(run=False) + def test_hyperu(self): + assert_mpmath_equal( + sc.hyperu, + exception_to_nan(lambda a, b, x: mpmath.hyperu(a, b, x, **HYPERKW)), + [Arg(), Arg(), Arg()], + ) + + @pytest.mark.xfail_on_32bit("mpmath issue gh-342: " + "unsupported operand mpz, long for pow") + def test_igam_fac(self): + def mp_igam_fac(a, x): + return mpmath.power(x, a)*mpmath.exp(-x)/mpmath.gamma(a) + + assert_mpmath_equal( + _igam_fac, + mp_igam_fac, + [Arg(0, 1e14, inclusive_a=False), Arg(0, 1e14)], + rtol=1e-10, + dps=29, + ) + + def test_j0(self): + # The Bessel function at large arguments is j0(x) ~ cos(x + phi)/sqrt(x) + # and at large arguments the phase of the cosine loses precision. + # + # This is numerically expected behavior, so we compare only up to + # 1e8 = 1e15 * 1e-7 + assert_mpmath_equal(sc.j0, mpmath.j0, [Arg(-1e3, 1e3)]) + assert_mpmath_equal(sc.j0, mpmath.j0, [Arg(-1e8, 1e8)], rtol=1e-5) + + def test_j1(self): + # See comment in test_j0 + assert_mpmath_equal(sc.j1, mpmath.j1, [Arg(-1e3, 1e3)]) + assert_mpmath_equal(sc.j1, mpmath.j1, [Arg(-1e8, 1e8)], rtol=1e-5) + + @pytest.mark.xfail(run=False) + def test_jacobi(self): + assert_mpmath_equal( + sc.eval_jacobi, + exception_to_nan(lambda a, b, c, x: mpmath.jacobi(a, b, c, x, **HYPERKW)), + [Arg(), Arg(), Arg(), Arg()], + ) + assert_mpmath_equal( + lambda n, b, c, x: sc.eval_jacobi(int(n), b, c, x), + exception_to_nan(lambda a, b, c, x: mpmath.jacobi(a, b, c, x, **HYPERKW)), + [IntArg(), Arg(), Arg(), Arg()], + ) + + def test_jacobi_int(self): + # Redefine functions to deal with numerical + mpmath issues + def jacobi(n, a, b, x): + # Mpmath does not handle n=0 case always correctly + if n == 0: + return 1.0 + return mpmath.jacobi(n, a, b, x) + assert_mpmath_equal( + lambda n, a, b, x: sc.eval_jacobi(int(n), a, b, x), + lambda n, a, b, x: exception_to_nan(jacobi)(n, a, b, x, **HYPERKW), + [IntArg(), Arg(), Arg(), Arg()], + n=4095, + dps=50, + ) + + def test_kei(self): + def kei(x): + if x == 0: + # work around mpmath issue at x=0 + return -pi/4 + return exception_to_nan(mpmath.kei)(0, x, **HYPERKW) + assert_mpmath_equal(sc.kei, kei, [Arg(-1e30, 1e30)], n=1000) + + def test_ker(self): + assert_mpmath_equal( + sc.ker, + exception_to_nan(lambda x: mpmath.ker(0, x, **HYPERKW)), + [Arg(-1e30, 1e30)], + n=1000, + ) + + @nonfunctional_tooslow + def test_laguerre(self): + assert_mpmath_equal( + trace_args(sc.eval_laguerre), + lambda n, x: exception_to_nan(mpmath.laguerre)(n, x, **HYPERKW), + [Arg(), Arg()], + ) + + def test_laguerre_int(self): + assert_mpmath_equal( + lambda n, x: sc.eval_laguerre(int(n), x), + lambda n, x: exception_to_nan(mpmath.laguerre)(n, x, **HYPERKW), + [IntArg(), Arg()], + n=20000, + ) + + @pytest.mark.xfail_on_32bit("see gh-3551 for bad points") + def test_lambertw_real(self): + assert_mpmath_equal( + lambda x, k: sc.lambertw(x, int(k.real)), + lambda x, k: mpmath.lambertw(x, int(k.real)), + [ComplexArg(-np.inf, np.inf), IntArg(0, 10)], + rtol=1e-13, nan_ok=False, + ) + + def test_lanczos_sum_expg_scaled(self): + maxgamma = 171.624376956302725 + e = np.exp(1) + g = 6.024680040776729583740234375 + + def gamma(x): + with np.errstate(over='ignore'): + fac = ((x + g - 0.5)/e)**(x - 0.5) + if fac != np.inf: + res = fac*_lanczos_sum_expg_scaled(x) + else: + fac = ((x + g - 0.5)/e)**(0.5*(x - 0.5)) + res = fac*_lanczos_sum_expg_scaled(x) + res *= fac + return res + + assert_mpmath_equal( + gamma, + mpmath.gamma, + [Arg(0, maxgamma, inclusive_a=False)], + rtol=1e-13, + ) + + @nonfunctional_tooslow + def test_legendre(self): + assert_mpmath_equal(sc.eval_legendre, mpmath.legendre, [Arg(), Arg()]) + + def test_legendre_int(self): + assert_mpmath_equal( + lambda n, x: sc.eval_legendre(int(n), x), + lambda n, x: exception_to_nan(mpmath.legendre)(n, x, **HYPERKW), + [IntArg(), Arg()], + n=20000, + ) + + # Check the small-x expansion + assert_mpmath_equal( + lambda n, x: sc.eval_legendre(int(n), x), + lambda n, x: exception_to_nan(mpmath.legendre)(n, x, **HYPERKW), + [IntArg(), FixedArg(np.logspace(-30, -4, 20))], + ) + + @pytest.mark.xfail(run=False, reason="apparently picks wrong function at |z| > 1") + def test_legenq(self): + def lqnm(n, m, z): + return sc.lqmn(m, n, z)[0][-1,-1] + + def legenq(n, m, z): + if abs(z) < 1e-15: + # mpmath has bad performance here + return np.nan + return exception_to_nan(mpmath.legenq)(n, m, z, type=2) + + assert_mpmath_equal( + lqnm, + legenq, + [IntArg(0, 100), IntArg(0, 100), Arg()], + ) + + @nonfunctional_tooslow + def test_legenq_complex(self): + def lqnm(n, m, z): + return sc.lqmn(int(m.real), int(n.real), z)[0][-1,-1] + + def legenq(n, m, z): + if abs(z) < 1e-15: + # mpmath has bad performance here + return np.nan + return exception_to_nan(mpmath.legenq)(int(n.real), int(m.real), z, type=2) + + assert_mpmath_equal( + lqnm, + legenq, + [IntArg(0, 100), IntArg(0, 100), ComplexArg()], + n=100, + ) + + def test_lgam1p(self): + def param_filter(x): + # Filter the poles + return np.where((np.floor(x) == x) & (x <= 0), False, True) + + def mp_lgam1p(z): + # The real part of loggamma is log(|gamma(z)|) + return mpmath.loggamma(1 + z).real + + assert_mpmath_equal( + _lgam1p, + mp_lgam1p, + [Arg()], + rtol=1e-13, + dps=100, + param_filter=param_filter, + ) + + def test_loggamma(self): + def mpmath_loggamma(z): + try: + res = mpmath.loggamma(z) + except ValueError: + res = complex(np.nan, np.nan) + return res + + assert_mpmath_equal( + sc.loggamma, + mpmath_loggamma, + [ComplexArg()], + nan_ok=False, + distinguish_nan_and_inf=False, + rtol=5e-14, + ) + + @pytest.mark.xfail(run=False) + def test_pcfd(self): + def pcfd(v, x): + return sc.pbdv(v, x)[0] + assert_mpmath_equal( + pcfd, + exception_to_nan(lambda v, x: mpmath.pcfd(v, x, **HYPERKW)), + [Arg(), Arg()], + ) + + @pytest.mark.xfail(run=False, reason="it's not the same as the mpmath function --- " + "maybe different definition?") + def test_pcfv(self): + def pcfv(v, x): + return sc.pbvv(v, x)[0] + assert_mpmath_equal( + pcfv, + lambda v, x: time_limited()(exception_to_nan(mpmath.pcfv))(v, x, **HYPERKW), + [Arg(), Arg()], + n=1000, + ) + + def test_pcfw(self): + def pcfw(a, x): + return sc.pbwa(a, x)[0] + + def dpcfw(a, x): + return sc.pbwa(a, x)[1] + + def mpmath_dpcfw(a, x): + return mpmath.diff(mpmath.pcfw, (a, x), (0, 1)) + + # The Zhang and Jin implementation only uses Taylor series and + # is thus accurate in only a very small range. + assert_mpmath_equal( + pcfw, + mpmath.pcfw, + [Arg(-5, 5), Arg(-5, 5)], + rtol=2e-8, + n=100, + ) + + assert_mpmath_equal( + dpcfw, + mpmath_dpcfw, + [Arg(-5, 5), Arg(-5, 5)], + rtol=2e-9, + n=100, + ) + + @pytest.mark.xfail(run=False, + reason="issues at large arguments (atol OK, rtol not) " + "and = version.Version("1.0.0"): + # no workarounds needed + mppoch = mpmath.rf + else: + def mppoch(a, m): + # deal with cases where the result in double precision + # hits exactly a non-positive integer, but the + # corresponding extended-precision mpf floats don't + if float(a + m) == int(a + m) and float(a + m) <= 0: + a = mpmath.mpf(a) + m = int(a + m) - a + return mpmath.rf(a, m) + + assert_mpmath_equal(sc.poch, mppoch, [Arg(), Arg()], dps=400) + + def test_sinpi(self): + eps = np.finfo(float).eps + assert_mpmath_equal( + _sinpi, + mpmath.sinpi, + [Arg()], + nan_ok=False, + rtol=2*eps, + ) + + def test_sinpi_complex(self): + assert_mpmath_equal( + _sinpi, + mpmath.sinpi, + [ComplexArg()], + nan_ok=False, + rtol=2e-14, + ) + + def test_shi(self): + def shi(x): + return sc.shichi(x)[0] + assert_mpmath_equal(shi, mpmath.shi, [Arg()]) + # check asymptotic series cross-over + assert_mpmath_equal(shi, mpmath.shi, [FixedArg([88 - 1e-9, 88, 88 + 1e-9])]) + + def test_shi_complex(self): + def shi(z): + return sc.shichi(z)[0] + # shi oscillates as Im[z] -> +- inf, so limit range + assert_mpmath_equal( + shi, + mpmath.shi, + [ComplexArg(complex(-np.inf, -1e8), complex(np.inf, 1e8))], + rtol=1e-12, + ) + + def test_si(self): + def si(x): + return sc.sici(x)[0] + assert_mpmath_equal(si, mpmath.si, [Arg()]) + + def test_si_complex(self): + def si(z): + return sc.sici(z)[0] + # si oscillates as Re[z] -> +- inf, so limit range + assert_mpmath_equal( + si, + mpmath.si, + [ComplexArg(complex(-1e8, -np.inf), complex(1e8, np.inf))], + rtol=1e-12, + ) + + def test_spence(self): + # mpmath uses a different convention for the dilogarithm + def dilog(x): + return mpmath.polylog(2, 1 - x) + # Spence has a branch cut on the negative real axis + assert_mpmath_equal( + sc.spence, + exception_to_nan(dilog), + [Arg(0, np.inf)], + rtol=1e-14, + ) + + def test_spence_complex(self): + def dilog(z): + return mpmath.polylog(2, 1 - z) + assert_mpmath_equal( + sc.spence, + exception_to_nan(dilog), + [ComplexArg()], + rtol=1e-14, + ) + + def test_struveh(self): + assert_mpmath_equal( + sc.struve, + exception_to_nan(mpmath.struveh), + [Arg(-1e4, 1e4), Arg(0, 1e4)], + rtol=5e-10, + ) + + def test_struvel(self): + def mp_struvel(v, z): + if v < 0 and z < -v and abs(v) > 1000: + # larger DPS needed for correct results + old_dps = mpmath.mp.dps + try: + mpmath.mp.dps = 500 + return mpmath.struvel(v, z) + finally: + mpmath.mp.dps = old_dps + return mpmath.struvel(v, z) + + assert_mpmath_equal( + sc.modstruve, + exception_to_nan(mp_struvel), + [Arg(-1e4, 1e4), Arg(0, 1e4)], + rtol=5e-10, + ignore_inf_sign=True, + ) + + def test_wrightomega_real(self): + def mpmath_wrightomega_real(x): + return mpmath.lambertw(mpmath.exp(x), mpmath.mpf('-0.5')) + + # For x < -1000 the Wright Omega function is just 0 to double + # precision, and for x > 1e21 it is just x to double + # precision. + assert_mpmath_equal( + sc.wrightomega, + mpmath_wrightomega_real, + [Arg(-1000, 1e21)], + rtol=5e-15, + atol=0, + nan_ok=False, + ) + + def test_wrightomega(self): + assert_mpmath_equal( + sc.wrightomega, + lambda z: _mpmath_wrightomega(z, 25), + [ComplexArg()], + rtol=1e-14, + nan_ok=False, + ) + + def test_hurwitz_zeta(self): + assert_mpmath_equal( + sc.zeta, + exception_to_nan(mpmath.zeta), + [Arg(a=1, b=1e10, inclusive_a=False), Arg(a=0, inclusive_a=False)], + ) + + def test_riemann_zeta(self): + assert_mpmath_equal( + sc.zeta, + lambda x: mpmath.zeta(x) if x != 1 else mpmath.inf, + [Arg(-100, 100)], + nan_ok=False, + rtol=5e-13, + ) + + def test_zetac(self): + assert_mpmath_equal( + sc.zetac, + lambda x: mpmath.zeta(x) - 1 if x != 1 else mpmath.inf, + [Arg(-100, 100)], + nan_ok=False, + dps=45, + rtol=5e-13, + ) + + def test_boxcox(self): + + def mp_boxcox(x, lmbda): + x = mpmath.mp.mpf(x) + lmbda = mpmath.mp.mpf(lmbda) + if lmbda == 0: + return mpmath.mp.log(x) + else: + return mpmath.mp.powm1(x, lmbda) / lmbda + + assert_mpmath_equal( + sc.boxcox, + exception_to_nan(mp_boxcox), + [Arg(a=0, inclusive_a=False), Arg()], + n=200, + dps=60, + rtol=1e-13, + ) + + def test_boxcox1p(self): + + def mp_boxcox1p(x, lmbda): + x = mpmath.mp.mpf(x) + lmbda = mpmath.mp.mpf(lmbda) + one = mpmath.mp.mpf(1) + if lmbda == 0: + return mpmath.mp.log(one + x) + else: + return mpmath.mp.powm1(one + x, lmbda) / lmbda + + assert_mpmath_equal( + sc.boxcox1p, + exception_to_nan(mp_boxcox1p), + [Arg(a=-1, inclusive_a=False), Arg()], + n=200, + dps=60, + rtol=1e-13, + ) + + def test_spherical_jn(self): + def mp_spherical_jn(n, z): + arg = mpmath.mpmathify(z) + out = (mpmath.besselj(n + mpmath.mpf(1)/2, arg) / + mpmath.sqrt(2*arg/mpmath.pi)) + if arg.imag == 0: + return out.real + else: + return out + + assert_mpmath_equal( + lambda n, z: sc.spherical_jn(int(n), z), + exception_to_nan(mp_spherical_jn), + [IntArg(0, 200), Arg(-1e8, 1e8)], + dps=300, + # underflow of `spherical_jn` is a bit premature; see gh-21629 + param_filter=(None, lambda z: np.abs(z) > 1e-20), + ) + + def test_spherical_jn_complex(self): + def mp_spherical_jn(n, z): + arg = mpmath.mpmathify(z) + out = (mpmath.besselj(n + mpmath.mpf(1)/2, arg) / + mpmath.sqrt(2*arg/mpmath.pi)) + if arg.imag == 0: + return out.real + else: + return out + + assert_mpmath_equal( + lambda n, z: sc.spherical_jn(int(n.real), z), + exception_to_nan(mp_spherical_jn), + [IntArg(0, 200), ComplexArg()] + ) + + def test_spherical_yn(self): + def mp_spherical_yn(n, z): + arg = mpmath.mpmathify(z) + out = (mpmath.bessely(n + mpmath.mpf(1)/2, arg) / + mpmath.sqrt(2*arg/mpmath.pi)) + if arg.imag == 0: + return out.real + else: + return out + + assert_mpmath_equal( + lambda n, z: sc.spherical_yn(int(n), z), + exception_to_nan(mp_spherical_yn), + [IntArg(0, 200), Arg(-1e10, 1e10)], + dps=100, + ) + + def test_spherical_yn_complex(self): + def mp_spherical_yn(n, z): + arg = mpmath.mpmathify(z) + out = (mpmath.bessely(n + mpmath.mpf(1)/2, arg) / + mpmath.sqrt(2*arg/mpmath.pi)) + if arg.imag == 0: + return out.real + else: + return out + + assert_mpmath_equal( + lambda n, z: sc.spherical_yn(int(n.real), z), + exception_to_nan(mp_spherical_yn), + [IntArg(0, 200), ComplexArg()], + ) + + def test_spherical_in(self): + def mp_spherical_in(n, z): + arg = mpmath.mpmathify(z) + out = (mpmath.besseli(n + mpmath.mpf(1)/2, arg) / + mpmath.sqrt(2*arg/mpmath.pi)) + if arg.imag == 0: + return out.real + else: + return out + + assert_mpmath_equal( + lambda n, z: sc.spherical_in(int(n), z), + exception_to_nan(mp_spherical_in), + [IntArg(0, 200), Arg()], + dps=200, + atol=10**(-278), + ) + + def test_spherical_in_complex(self): + def mp_spherical_in(n, z): + arg = mpmath.mpmathify(z) + out = (mpmath.besseli(n + mpmath.mpf(1)/2, arg) / + mpmath.sqrt(2*arg/mpmath.pi)) + if arg.imag == 0: + return out.real + else: + return out + + assert_mpmath_equal( + lambda n, z: sc.spherical_in(int(n.real), z), + exception_to_nan(mp_spherical_in), + [IntArg(0, 200), ComplexArg()], + ) + + def test_spherical_kn(self): + def mp_spherical_kn(n, z): + arg = mpmath.mpmathify(z) + out = (mpmath.besselk(n + mpmath.mpf(1)/2, arg) / + mpmath.sqrt(2*arg/mpmath.pi)) + if mpmath.mpmathify(z).imag == 0: + return out.real + else: + return out + + assert_mpmath_equal( + lambda n, z: sc.spherical_kn(int(n), z), + exception_to_nan(mp_spherical_kn), + [IntArg(0, 150), Arg()], + dps=100, + ) + + @pytest.mark.xfail(run=False, + reason="Accuracy issues near z = -1 inherited from kv.") + def test_spherical_kn_complex(self): + def mp_spherical_kn(n, z): + arg = mpmath.mpmathify(z) + out = (mpmath.besselk(n + mpmath.mpf(1)/2, arg) / + mpmath.sqrt(2*arg/mpmath.pi)) + if arg.imag == 0: + return out.real + else: + return out + + assert_mpmath_equal( + lambda n, z: sc.spherical_kn(int(n.real), z), + exception_to_nan(mp_spherical_kn), + [IntArg(0, 200), ComplexArg()], + dps=200, + ) diff --git a/v12x/pylib/scipy/special/tests/test_nan_inputs.py b/v12x/pylib/scipy/special/tests/test_nan_inputs.py new file mode 100644 index 0000000000000000000000000000000000000000..76af3032b62484b630713631e23d95110161b759 --- /dev/null +++ b/v12x/pylib/scipy/special/tests/test_nan_inputs.py @@ -0,0 +1,72 @@ +"""Test how the ufuncs in special handle nan inputs. + +""" +import warnings + +from collections.abc import Callable + +import numpy as np +from numpy.testing import assert_array_equal, assert_ +import pytest +import scipy.special as sc + + +KNOWNFAILURES: dict[str, Callable] = {} + +POSTPROCESSING: dict[str, Callable] = {} + + +def _get_ufuncs(): + ufuncs = [] + ufunc_names = [] + for name in sorted(sc.__dict__): + obj = sc.__dict__[name] + if not isinstance(obj, np.ufunc): + continue + msg = KNOWNFAILURES.get(obj) + if msg is None: + ufuncs.append(obj) + ufunc_names.append(name) + else: + fail = pytest.mark.xfail(run=False, reason=msg) + ufuncs.append(pytest.param(obj, marks=fail)) + ufunc_names.append(name) + return ufuncs, ufunc_names + + +UFUNCS, UFUNC_NAMES = _get_ufuncs() + + +@pytest.mark.parametrize("func", UFUNCS, ids=UFUNC_NAMES) +def test_nan_inputs(func): + args = (np.nan,)*func.nin + with warnings.catch_warnings(): + # Ignore warnings about unsafe casts from legacy wrappers + warnings.filterwarnings( + "ignore", + "floating point number truncated to an integer", + RuntimeWarning + ) + try: + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + res = func(*args) + except TypeError: + # One of the arguments doesn't take real inputs + return + if func in POSTPROCESSING: + res = POSTPROCESSING[func](*res) + + msg = f"got {res} instead of nan" + assert_array_equal(np.isnan(res), True, err_msg=msg) + + +def test_legacy_cast(): + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", + "floating point number truncated to an integer", + RuntimeWarning + ) + res = sc.bdtrc(np.nan, 1, 0.5) + assert_(np.isnan(res)) diff --git a/v12x/pylib/scipy/special/tests/test_ndtr.py b/v12x/pylib/scipy/special/tests/test_ndtr.py new file mode 100644 index 0000000000000000000000000000000000000000..ba9b689b34384585cc65204000febcb99c910d55 --- /dev/null +++ b/v12x/pylib/scipy/special/tests/test_ndtr.py @@ -0,0 +1,77 @@ +import numpy as np +from numpy.testing import assert_equal, assert_allclose +import scipy.special as sc + + +def test_ndtr(): + assert_equal(sc.ndtr(0), 0.5) + assert_allclose(sc.ndtr(1), 0.8413447460685429) + + +class TestNdtri: + + def test_zero(self): + assert sc.ndtri(0.5) == 0.0 + + def test_asymptotes(self): + assert_equal(sc.ndtri([0.0, 1.0]), [-np.inf, np.inf]) + + def test_outside_of_domain(self): + assert all(np.isnan(sc.ndtri([-1.5, 1.5]))) + + +class TestLogNdtr: + + # The expected values in these tests were computed with mpmath: + # + # def log_ndtr_mp(x): + # return mpmath.log(mpmath.ncdf(x)) + # + + def test_log_ndtr_moderate_le8(self): + x = np.array([-0.75, -0.25, 0, 0.5, 1.5, 2.5, 3, 4, 5, 7, 8]) + expected = np.array([-1.4844482299196562, + -0.9130617648111351, + -0.6931471805599453, + -0.3689464152886564, + -0.06914345561223398, + -0.006229025485860002, + -0.0013508099647481938, + -3.167174337748927e-05, + -2.866516129637636e-07, + -1.279812543886654e-12, + -6.220960574271786e-16]) + y = sc.log_ndtr(x) + assert_allclose(y, expected, rtol=1e-14) + + def test_log_ndtr_values_8_16(self): + x = np.array([8.001, 8.06, 8.15, 8.5, 10, 12, 14, 16]) + expected = [-6.170639424817055e-16, + -3.814722443652823e-16, + -1.819621363526629e-16, + -9.479534822203318e-18, + -7.619853024160525e-24, + -1.776482112077679e-33, + -7.7935368191928e-45, + -6.388754400538087e-58] + y = sc.log_ndtr(x) + assert_allclose(y, expected, rtol=5e-14) + + def test_log_ndtr_values_16_31(self): + x = np.array([16.15, 20.3, 21.4, 26.2, 30.9]) + expected = [-5.678084565148492e-59, + -6.429244467698346e-92, + -6.680402412553295e-102, + -1.328698078458869e-151, + -5.972288641838264e-210] + y = sc.log_ndtr(x) + assert_allclose(y, expected, rtol=2e-13) + + def test_log_ndtr_values_gt31(self): + x = np.array([31.6, 32.8, 34.9, 37.1]) + expected = [-1.846036234858162e-219, + -2.9440539964066835e-236, + -3.71721649450857e-267, + -1.4047119663106221e-301] + y = sc.log_ndtr(x) + assert_allclose(y, expected, rtol=3e-13) diff --git a/v12x/pylib/scipy/special/tests/test_ndtri_exp.py b/v12x/pylib/scipy/special/tests/test_ndtri_exp.py new file mode 100644 index 0000000000000000000000000000000000000000..82a9fbd3bcda117770e00018facda3f56630a6bc --- /dev/null +++ b/v12x/pylib/scipy/special/tests/test_ndtri_exp.py @@ -0,0 +1,94 @@ +import pytest +import numpy as np +from numpy.testing import assert_equal, assert_allclose +from scipy.special import log_ndtr, ndtri_exp +from scipy.special._testutils import assert_func_equal + + +def log_ndtr_ndtri_exp(y): + return log_ndtr(ndtri_exp(y)) + + +@pytest.fixture(scope="class") +def uniform_random_points(): + random_state = np.random.RandomState(1234) + points = random_state.random_sample(1000) + return points + + +class TestNdtriExp: + """Tests that ndtri_exp is sufficiently close to an inverse of log_ndtr. + + We have separate tests for the five intervals (-inf, -10), + [-10, -2), [-2, -0.14542), [-0.14542, -1e-6), and [-1e-6, 0). + ndtri_exp(y) is computed in three different ways depending on if y + is in (-inf, -2), [-2, log(1 - exp(-2))], or [log(1 - exp(-2), 0). + Each of these intervals is given its own test with two additional tests + for handling very small values and values very close to zero. + """ + + @pytest.mark.parametrize( + "test_input", [-1e1, -1e2, -1e10, -1e20, -np.finfo(float).max] + ) + def test_very_small_arg(self, test_input, uniform_random_points): + scale = test_input + points = scale * (0.5 * uniform_random_points + 0.5) + assert_func_equal( + log_ndtr_ndtri_exp, + lambda y: y, points, + rtol=1e-14, + nan_ok=True + ) + + @pytest.mark.parametrize( + "interval,expected_rtol", + [ + ((-10, -2), 1e-14), + ((-2, -0.14542), 1e-12), + ((-0.14542, -1e-6), 1e-10), + ((-1e-6, 0), 1e-6), + ], + ) + def test_in_interval(self, interval, expected_rtol, uniform_random_points): + left, right = interval + points = (right - left) * uniform_random_points + left + assert_func_equal( + log_ndtr_ndtri_exp, + lambda y: y, points, + rtol=expected_rtol, + nan_ok=True + ) + + def test_extreme(self): + # bigneg is not quite the largest negative double precision value. + # Here's why: + # The round-trip calculation + # y = ndtri_exp(bigneg) + # bigneg2 = log_ndtr(y) + # where bigneg is a very large negative value, would--with infinite + # precision--result in bigneg2 == bigneg. When bigneg is large enough, + # y is effectively equal to -sqrt(2)*sqrt(-bigneg), and log_ndtr(y) is + # effectively -(y/sqrt(2))**2. If we use bigneg = np.finfo(float).min, + # then by construction, the theoretical value is the most negative + # finite value that can be represented with 64 bit float point. This + # means tiny changes in how the computation proceeds can result in the + # return value being -inf. (E.g. changing the constant representation + # of 1/sqrt(2) from 0.7071067811865475--which is the value returned by + # 1/np.sqrt(2)--to 0.7071067811865476--which is the most accurate 64 + # bit floating point representation of 1/sqrt(2)--results in the + # round-trip that starts with np.finfo(float).min returning -inf. So + # we'll move the bigneg value a few ULPs towards 0 to avoid this + # sensitivity. + # Use the reduce method to apply nextafter four times. + bigneg = np.nextafter.reduce([np.finfo(float).min, 0, 0, 0, 0]) + # tinyneg is approx. -2.225e-308. + tinyneg = -np.finfo(float).tiny + x = np.array([tinyneg, bigneg]) + result = log_ndtr_ndtri_exp(x) + assert_allclose(result, x, rtol=1e-12) + + def test_asymptotes(self): + assert_equal(ndtri_exp([-np.inf, 0.0]), [-np.inf, np.inf]) + + def test_outside_domain(self): + assert np.isnan(ndtri_exp(1.0)) diff --git a/v12x/pylib/scipy/special/tests/test_orthogonal.py b/v12x/pylib/scipy/special/tests/test_orthogonal.py new file mode 100644 index 0000000000000000000000000000000000000000..130d52073b3da0a27fc7a15ca18506ca9094942f --- /dev/null +++ b/v12x/pylib/scipy/special/tests/test_orthogonal.py @@ -0,0 +1,823 @@ +import pytest +from pytest import raises as assert_raises + +import numpy as np +from numpy import array, sqrt +from numpy.testing import assert_equal, assert_allclose + +from scipy import integrate +import scipy.special as sc +from scipy.special import gamma +import scipy.special._orthogonal as orth + + +class TestCheby: + def test_chebyc(self): + C0 = orth.chebyc(0) + C1 = orth.chebyc(1) + with np.errstate(all='ignore'): + C2 = orth.chebyc(2) + C3 = orth.chebyc(3) + C4 = orth.chebyc(4) + C5 = orth.chebyc(5) + + assert_allclose(C0.c, [2], atol=1.5e-13, rtol=0) + assert_allclose(C1.c, [1, 0], atol=1.5e-13, rtol=0) + assert_allclose(C2.c, [1, 0, -2], atol=1.5e-13, rtol=0) + assert_allclose(C3.c, [1, 0, -3, 0], atol=1.5e-13, rtol=0) + assert_allclose(C4.c, [1, 0, -4, 0, 2], atol=1.5e-13, rtol=0) + assert_allclose(C5.c, [1, 0, -5, 0, 5, 0],atol=1.5e-13, rtol=0) + + def test_chebys(self): + S0 = orth.chebys(0) + S1 = orth.chebys(1) + S2 = orth.chebys(2) + S3 = orth.chebys(3) + S4 = orth.chebys(4) + S5 = orth.chebys(5) + assert_allclose(S0.c, [1], atol=1.5e-13, rtol=0) + assert_allclose(S1.c, [1, 0], atol=1.5e-13, rtol=0) + assert_allclose(S2.c, [1, 0, -1], atol=1.5e-13, rtol=0) + assert_allclose(S3.c, [1, 0, -2, 0], atol=1.5e-13, rtol=0) + assert_allclose(S4.c, [1, 0, -3, 0, 1], atol=1.5e-13, rtol=0) + assert_allclose(S5.c, [1, 0, -4, 0, 3, 0], atol=1.5e-13, rtol=0) + + def test_chebyt(self): + T0 = orth.chebyt(0) + T1 = orth.chebyt(1) + T2 = orth.chebyt(2) + T3 = orth.chebyt(3) + T4 = orth.chebyt(4) + T5 = orth.chebyt(5) + assert_allclose(T0.c, [1], atol=1.5e-13, rtol=0) + assert_allclose(T1.c, [1, 0], atol=1.5e-13, rtol=0) + assert_allclose(T2.c, [2, 0, -1], atol=1.5e-13, rtol=0) + assert_allclose(T3.c, [4, 0, -3, 0], atol=1.5e-13, rtol=0) + assert_allclose(T4.c, [8, 0, -8, 0, 1], atol=1.5e-13, rtol=0) + assert_allclose(T5.c, [16, 0, -20, 0, 5, 0], atol=1.5e-13, rtol=0) + + def test_chebyu(self): + U0 = orth.chebyu(0) + U1 = orth.chebyu(1) + U2 = orth.chebyu(2) + U3 = orth.chebyu(3) + U4 = orth.chebyu(4) + U5 = orth.chebyu(5) + assert_allclose(U0.c, [1], atol=1.5e-13, rtol=0) + assert_allclose(U1.c, [2, 0], atol=1.5e-13, rtol=0) + assert_allclose(U2.c, [4, 0, -1], atol=1.5e-13, rtol=0) + assert_allclose(U3.c, [8, 0, -4, 0], atol=1.5e-13, rtol=0) + assert_allclose(U4.c, [16, 0, -12, 0, 1], atol=1.5e-13, rtol=0) + assert_allclose(U5.c, [32, 0, -32, 0, 6, 0], atol=1.5e-13, rtol=0) + + +class TestGegenbauer: + + def test_gegenbauer(self): + a = 5*np.random.random() - 0.5 + if np.any(a == 0): + a = -0.2 + Ca0 = orth.gegenbauer(0,a) + Ca1 = orth.gegenbauer(1,a) + Ca2 = orth.gegenbauer(2,a) + Ca3 = orth.gegenbauer(3,a) + Ca4 = orth.gegenbauer(4,a) + Ca5 = orth.gegenbauer(5,a) + + assert_allclose(Ca0.c, array([1]), atol=1.5e-13, rtol=0) + assert_allclose(Ca1.c, array([2*a, 0]), atol=1.5e-13, rtol=0) + assert_allclose(Ca2.c, array([2*a*(a + 1), 0, -a]), + atol=1.5e-13, rtol=0) + assert_allclose(Ca3.c, array([4*sc.poch(a, 3), 0,-6*a*(a + 1), 0])/3.0, + atol=1.5e-11, rtol=0) + assert_allclose(Ca4.c, array([4*sc.poch(a, 4), 0, -12*sc.poch(a, 3), + 0, 3*a*(a + 1)])/6.0, + atol=1.5e-11, rtol=0) + assert_allclose(Ca5.c, array([4*sc.poch(a, 5), 0, -20*sc.poch(a, 4), + 0, 15*sc.poch(a, 3), 0])/15.0, + atol=1.5e-11, rtol=0) + + @pytest.mark.parametrize('a', [0, 1]) + def test_n_zero_gh8888(self, a): + # gh-8888 reported that gegenbauer(0, 0) returns NaN polynomial + Cn0 = orth.gegenbauer(0, a) + assert_equal(Cn0.c, np.asarray([1.])) + + def test_valid_alpha(self): + # Check input validation of `alpha` + message = '`alpha` must be a finite number greater...' + with pytest.raises(ValueError, match=message): + orth.gegenbauer(0, np.nan) + with pytest.raises(ValueError, match=message): + orth.gegenbauer(1, -0.5) + with pytest.raises(ValueError, match=message): + orth.gegenbauer(2, -np.inf) + + +class TestHermite: + def test_hermite(self): + H0 = orth.hermite(0) + H1 = orth.hermite(1) + H2 = orth.hermite(2) + H3 = orth.hermite(3) + H4 = orth.hermite(4) + H5 = orth.hermite(5) + assert_allclose(H0.c, [1], atol=1.5e-13, rtol=0) + assert_allclose(H1.c, [2, 0], atol=1.5e-13, rtol=0) + assert_allclose(H2.c, [4, 0, -2], atol=1.5e-13, rtol=0) + assert_allclose(H3.c, [8, 0, -12, 0], atol=1.5e-13, rtol=0) + assert_allclose(H4.c, [16, 0, -48, 0, 12], atol=1.5e-12, rtol=0) + assert_allclose(H5.c, [32, 0, -160, 0, 120, 0], atol=1.5e-12, rtol=0) + + def test_hermitenorm(self): + # He_n(x) = 2**(-n/2) H_n(x/sqrt(2)) + psub = np.poly1d([1.0/sqrt(2),0]) + H0 = orth.hermitenorm(0) + H1 = orth.hermitenorm(1) + H2 = orth.hermitenorm(2) + H3 = orth.hermitenorm(3) + H4 = orth.hermitenorm(4) + H5 = orth.hermitenorm(5) + he0 = orth.hermite(0)(psub) + he1 = orth.hermite(1)(psub) / sqrt(2) + he2 = orth.hermite(2)(psub) / 2.0 + he3 = orth.hermite(3)(psub) / (2*sqrt(2)) + he4 = orth.hermite(4)(psub) / 4.0 + he5 = orth.hermite(5)(psub) / (4.0*sqrt(2)) + + assert_allclose(H0.c, he0.c, atol=1.5e-13, rtol=0) + assert_allclose(H1.c, he1.c, atol=1.5e-13, rtol=0) + assert_allclose(H2.c, he2.c, atol=1.5e-13, rtol=0) + assert_allclose(H3.c, he3.c, atol=1.5e-13, rtol=0) + assert_allclose(H4.c, he4.c, atol=1.5e-13, rtol=0) + assert_allclose(H5.c, he5.c, atol=1.5e-13, rtol=0) + + +class TestShLegendre: + def test_sh_legendre(self): + # P*_n(x) = P_n(2x-1) + psub = np.poly1d([2,-1]) + Ps0 = orth.sh_legendre(0) + Ps1 = orth.sh_legendre(1) + Ps2 = orth.sh_legendre(2) + Ps3 = orth.sh_legendre(3) + Ps4 = orth.sh_legendre(4) + Ps5 = orth.sh_legendre(5) + pse0 = orth.legendre(0)(psub) + pse1 = orth.legendre(1)(psub) + pse2 = orth.legendre(2)(psub) + pse3 = orth.legendre(3)(psub) + pse4 = orth.legendre(4)(psub) + pse5 = orth.legendre(5)(psub) + assert_allclose(Ps0.c, pse0.c, atol=1.5e-13, rtol=0) + assert_allclose(Ps1.c, pse1.c, atol=1.5e-13, rtol=0) + assert_allclose(Ps2.c, pse2.c, atol=1.5e-13, rtol=0) + assert_allclose(Ps3.c, pse3.c, atol=1.5e-13, rtol=0) + assert_allclose(Ps4.c, pse4.c, atol=1.5e-12, rtol=0) + assert_allclose(Ps5.c, pse5.c, atol=1.5e-12, rtol=0) + + +class TestShChebyt: + def test_sh_chebyt(self): + # T*_n(x) = T_n(2x-1) + psub = np.poly1d([2,-1]) + Ts0 = orth.sh_chebyt(0) + Ts1 = orth.sh_chebyt(1) + Ts2 = orth.sh_chebyt(2) + Ts3 = orth.sh_chebyt(3) + Ts4 = orth.sh_chebyt(4) + Ts5 = orth.sh_chebyt(5) + tse0 = orth.chebyt(0)(psub) + tse1 = orth.chebyt(1)(psub) + tse2 = orth.chebyt(2)(psub) + tse3 = orth.chebyt(3)(psub) + tse4 = orth.chebyt(4)(psub) + tse5 = orth.chebyt(5)(psub) + assert_allclose(Ts0.c, tse0.c, atol=1.5e-13, rtol=0) + assert_allclose(Ts1.c, tse1.c, atol=1.5e-13, rtol=0) + assert_allclose(Ts2.c, tse2.c, atol=1.5e-13, rtol=0) + assert_allclose(Ts3.c, tse3.c, atol=1.5e-13, rtol=0) + assert_allclose(Ts4.c, tse4.c, atol=1.5e-12, rtol=0) + assert_allclose(Ts5.c, tse5.c, atol=1.5e-12, rtol=0) + + +class TestShChebyu: + def test_sh_chebyu(self): + # U*_n(x) = U_n(2x-1) + psub = np.poly1d([2,-1]) + Us0 = orth.sh_chebyu(0) + Us1 = orth.sh_chebyu(1) + Us2 = orth.sh_chebyu(2) + Us3 = orth.sh_chebyu(3) + Us4 = orth.sh_chebyu(4) + Us5 = orth.sh_chebyu(5) + use0 = orth.chebyu(0)(psub) + use1 = orth.chebyu(1)(psub) + use2 = orth.chebyu(2)(psub) + use3 = orth.chebyu(3)(psub) + use4 = orth.chebyu(4)(psub) + use5 = orth.chebyu(5)(psub) + assert_allclose(Us0.c, use0.c, atol=1.5e-13, rtol=0) + assert_allclose(Us1.c, use1.c, atol=1.5e-13, rtol=0) + assert_allclose(Us2.c, use2.c, atol=1.5e-13, rtol=0) + assert_allclose(Us3.c, use3.c, atol=1.5e-13, rtol=0) + assert_allclose(Us4.c, use4.c, atol=1.5e-12, rtol=0) + assert_allclose(Us5.c, use5.c, atol=1.5e-11, rtol=0) + + +class TestShJacobi: + def test_sh_jacobi(self): + # G^(p,q)_n(x) = n! gamma(n+p)/gamma(2*n+p) * P^(p-q,q-1)_n(2*x-1) + def conv(n, p): + return gamma(n + 1) * gamma(n + p) / gamma(2 * n + p) + psub = np.poly1d([2,-1]) + q = 4 * np.random.random() + p = q-1 + 2*np.random.random() + # print("shifted jacobi p,q = ", p, q) + G0 = orth.sh_jacobi(0,p,q) + G1 = orth.sh_jacobi(1,p,q) + G2 = orth.sh_jacobi(2,p,q) + G3 = orth.sh_jacobi(3,p,q) + G4 = orth.sh_jacobi(4,p,q) + G5 = orth.sh_jacobi(5,p,q) + ge0 = orth.jacobi(0,p-q,q-1)(psub) * conv(0,p) + ge1 = orth.jacobi(1,p-q,q-1)(psub) * conv(1,p) + ge2 = orth.jacobi(2,p-q,q-1)(psub) * conv(2,p) + ge3 = orth.jacobi(3,p-q,q-1)(psub) * conv(3,p) + ge4 = orth.jacobi(4,p-q,q-1)(psub) * conv(4,p) + ge5 = orth.jacobi(5,p-q,q-1)(psub) * conv(5,p) + + assert_allclose(G0.c, ge0.c, atol=1.5e-13, rtol=0) + assert_allclose(G1.c, ge1.c, atol=1.5e-13, rtol=0) + assert_allclose(G2.c, ge2.c, atol=1.5e-13, rtol=0) + assert_allclose(G3.c, ge3.c, atol=1.5e-13, rtol=0) + assert_allclose(G4.c, ge4.c, atol=1.5e-13, rtol=0) + assert_allclose(G5.c, ge5.c, atol=1.5e-13, rtol=0) + + +class TestCall: + def test_call(self): + poly = [] + for n in range(5): + poly.extend([x.strip() for x in + (f""" + orth.jacobi({n},0.3,0.9) + orth.sh_jacobi({n},0.3,0.9) + orth.genlaguerre({n},0.3) + orth.laguerre({n}) + orth.hermite({n}) + orth.hermitenorm({n}) + orth.gegenbauer({n},0.3) + orth.chebyt({n}) + orth.chebyu({n}) + orth.chebyc({n}) + orth.chebys({n}) + orth.sh_chebyt({n}) + orth.sh_chebyu({n}) + orth.legendre({n}) + orth.sh_legendre({n}) + """).split()]) + with np.errstate(all='ignore'): + for pstr in poly: + p = eval(pstr) + assert_allclose(p(0.315), np.poly1d(p.coef)(0.315), + atol=1.5e-7, rtol=0, err_msg=pstr) + + +class TestGenlaguerre: + def test_regression(self): + assert_equal(orth.genlaguerre(1, 1, monic=False)(0), 2.) + assert_equal(orth.genlaguerre(1, 1, monic=True)(0), -2.) + assert_equal(orth.genlaguerre(1, 1, monic=False), np.poly1d([-1, 2])) + assert_equal(orth.genlaguerre(1, 1, monic=True), np.poly1d([1, -2])) + + +def verify_gauss_quad(root_func, eval_func, weight_func, a, b, N, + rtol=1e-15, atol=5e-14): + # this test is copied from numpy's TestGauss in test_hermite.py + x, w, mu = root_func(N, True) + + n = np.arange(N, dtype=np.dtype("long")) + v = eval_func(n[:,np.newaxis], x) + vv = np.dot(v*w, v.T) + vd = 1 / np.sqrt(vv.diagonal()) + vv = vd[:, np.newaxis] * vv * vd + assert_allclose(vv, np.eye(N), rtol, atol) + + # check that the integral of 1 is correct + assert_allclose(w.sum(), mu, rtol, atol) + + # compare the results of integrating a function with quad. + def f(x): + return x ** 3 - 3 * x ** 2 + x - 2 + resI = integrate.quad(lambda x: f(x)*weight_func(x), a, b) + resG = np.vdot(f(x), w) + rtol = 1e-6 if 1e-6 < resI[1] else resI[1] * 10 + assert_allclose(resI[0], resG, rtol=rtol) + +def test_roots_jacobi(): + def rf(a, b): + return lambda n, mu: sc.roots_jacobi(n, a, b, mu) + def ef(a, b): + return lambda n, x: sc.eval_jacobi(n, a, b, x) + def wf(a, b): + return lambda x: (1 - x) ** a * (1 + x) ** b + + vgq = verify_gauss_quad + vgq(rf(-0.5, -0.75), ef(-0.5, -0.75), wf(-0.5, -0.75), -1., 1., 5) + vgq(rf(-0.5, -0.75), ef(-0.5, -0.75), wf(-0.5, -0.75), -1., 1., + 25, atol=1e-12) + vgq(rf(-0.5, -0.75), ef(-0.5, -0.75), wf(-0.5, -0.75), -1., 1., + 100, atol=1e-11) + + vgq(rf(0.5, -0.5), ef(0.5, -0.5), wf(0.5, -0.5), -1., 1., 5) + vgq(rf(0.5, -0.5), ef(0.5, -0.5), wf(0.5, -0.5), -1., 1., 25, atol=1.5e-13) + vgq(rf(0.5, -0.5), ef(0.5, -0.5), wf(0.5, -0.5), -1., 1., 100, atol=2e-12) + + vgq(rf(1, 0.5), ef(1, 0.5), wf(1, 0.5), -1., 1., 5, atol=2e-13) + vgq(rf(1, 0.5), ef(1, 0.5), wf(1, 0.5), -1., 1., 25, atol=2e-13) + vgq(rf(1, 0.5), ef(1, 0.5), wf(1, 0.5), -1., 1., 100, atol=1e-12) + + vgq(rf(0.9, 2), ef(0.9, 2), wf(0.9, 2), -1., 1., 5) + vgq(rf(0.9, 2), ef(0.9, 2), wf(0.9, 2), -1., 1., 25, atol=1e-13) + vgq(rf(0.9, 2), ef(0.9, 2), wf(0.9, 2), -1., 1., 100, atol=3e-13) + + vgq(rf(18.24, 27.3), ef(18.24, 27.3), wf(18.24, 27.3), -1., 1., 5) + vgq(rf(18.24, 27.3), ef(18.24, 27.3), wf(18.24, 27.3), -1., 1., 25, + atol=1.1e-14) + vgq(rf(18.24, 27.3), ef(18.24, 27.3), wf(18.24, 27.3), -1., 1., + 100, atol=1e-13) + + vgq(rf(47.1, -0.2), ef(47.1, -0.2), wf(47.1, -0.2), -1., 1., 5, atol=1e-13) + vgq(rf(47.1, -0.2), ef(47.1, -0.2), wf(47.1, -0.2), -1., 1., 25, atol=2e-13) + vgq(rf(47.1, -0.2), ef(47.1, -0.2), wf(47.1, -0.2), -1., 1., + 100, atol=1e-11) + + vgq(rf(1., 658.), ef(1., 658.), wf(1., 658.), -1., 1., 5, atol=2e-13) + vgq(rf(1., 658.), ef(1., 658.), wf(1., 658.), -1., 1., 25, atol=1e-12) + vgq(rf(1., 658.), ef(1., 658.), wf(1., 658.), -1., 1., 100, atol=1e-11) + vgq(rf(1., 658.), ef(1., 658.), wf(1., 658.), -1., 1., 250, atol=1e-11) + + vgq(rf(511., 511.), ef(511., 511.), wf(511., 511.), -1., 1., 5, + atol=1e-12) + vgq(rf(511., 511.), ef(511., 511.), wf(511., 511.), -1., 1., 25, + atol=1e-11) + vgq(rf(511., 511.), ef(511., 511.), wf(511., 511.), -1., 1., 100, + atol=1e-10) + + vgq(rf(511., 512.), ef(511., 512.), wf(511., 512.), -1., 1., 5, + atol=1e-12) + vgq(rf(511., 512.), ef(511., 512.), wf(511., 512.), -1., 1., 25, + atol=1e-11) + vgq(rf(511., 512.), ef(511., 512.), wf(511., 512.), -1., 1., 100, + atol=1e-10) + + vgq(rf(1000., 500.), ef(1000., 500.), wf(1000., 500.), -1., 1., 5, + atol=1e-12) + vgq(rf(1000., 500.), ef(1000., 500.), wf(1000., 500.), -1., 1., 25, + atol=1e-11) + vgq(rf(1000., 500.), ef(1000., 500.), wf(1000., 500.), -1., 1., 100, + atol=1e-10) + + vgq(rf(2.25, 68.9), ef(2.25, 68.9), wf(2.25, 68.9), -1., 1., 5) + vgq(rf(2.25, 68.9), ef(2.25, 68.9), wf(2.25, 68.9), -1., 1., 25, + atol=1e-13) + vgq(rf(2.25, 68.9), ef(2.25, 68.9), wf(2.25, 68.9), -1., 1., 100, + atol=1e-13) + + # when alpha == beta == 0, P_n^{a,b}(x) == P_n(x) + xj, wj = sc.roots_jacobi(6, 0.0, 0.0) + xl, wl = sc.roots_legendre(6) + assert_allclose(xj, xl, 1e-14, 1e-14) + assert_allclose(wj, wl, 1e-14, 1e-14) + + # when alpha == beta != 0, P_n^{a,b}(x) == C_n^{alpha+0.5}(x) + xj, wj = sc.roots_jacobi(6, 4.0, 4.0) + xc, wc = sc.roots_gegenbauer(6, 4.5) + assert_allclose(xj, xc, 1e-14, 1e-14) + assert_allclose(wj, wc, 1e-14, 1e-14) + + x, w = sc.roots_jacobi(5, 2, 3, False) + y, v, m = sc.roots_jacobi(5, 2, 3, True) + assert_allclose(x, y, 1e-14, 1e-14) + assert_allclose(w, v, 1e-14, 1e-14) + + muI, muI_err = integrate.quad(wf(2,3), -1, 1) + assert_allclose(m, muI, rtol=muI_err) + + assert_raises(ValueError, sc.roots_jacobi, 0, 1, 1) + assert_raises(ValueError, sc.roots_jacobi, 3.3, 1, 1) # type: ignore[call-overload] + assert_raises(ValueError, sc.roots_jacobi, 3, -2, 1) + assert_raises(ValueError, sc.roots_jacobi, 3, 1, -2) + assert_raises(ValueError, sc.roots_jacobi, 3, -2, -2) + +def test_roots_sh_jacobi(): + def rf(a, b): + return lambda n, mu: sc.roots_sh_jacobi(n, a, b, mu) + def ef(a, b): + return lambda n, x: sc.eval_sh_jacobi(n, a, b, x) + def wf(a, b): + return lambda x: (1.0 - x) ** (a - b) * x ** (b - 1.0) + + vgq = verify_gauss_quad + vgq(rf(-0.5, 0.25), ef(-0.5, 0.25), wf(-0.5, 0.25), 0., 1., 5) + vgq(rf(-0.5, 0.25), ef(-0.5, 0.25), wf(-0.5, 0.25), 0., 1., + 25, atol=1e-12) + vgq(rf(-0.5, 0.25), ef(-0.5, 0.25), wf(-0.5, 0.25), 0., 1., + 100, atol=1e-11) + + vgq(rf(0.5, 0.5), ef(0.5, 0.5), wf(0.5, 0.5), 0., 1., 5) + vgq(rf(0.5, 0.5), ef(0.5, 0.5), wf(0.5, 0.5), 0., 1., 25, atol=1e-13) + vgq(rf(0.5, 0.5), ef(0.5, 0.5), wf(0.5, 0.5), 0., 1., 100, atol=1e-12) + + vgq(rf(1, 0.5), ef(1, 0.5), wf(1, 0.5), 0., 1., 5) + vgq(rf(1, 0.5), ef(1, 0.5), wf(1, 0.5), 0., 1., 25, atol=1.5e-13) + vgq(rf(1, 0.5), ef(1, 0.5), wf(1, 0.5), 0., 1., 100, atol=2e-12) + + vgq(rf(2, 0.9), ef(2, 0.9), wf(2, 0.9), 0., 1., 5) + vgq(rf(2, 0.9), ef(2, 0.9), wf(2, 0.9), 0., 1., 25, atol=1e-13) + vgq(rf(2, 0.9), ef(2, 0.9), wf(2, 0.9), 0., 1., 100, atol=1e-12) + + vgq(rf(27.3, 18.24), ef(27.3, 18.24), wf(27.3, 18.24), 0., 1., 5) + vgq(rf(27.3, 18.24), ef(27.3, 18.24), wf(27.3, 18.24), 0., 1., 25) + vgq(rf(27.3, 18.24), ef(27.3, 18.24), wf(27.3, 18.24), 0., 1., + 100, atol=1e-13) + + vgq(rf(47.1, 0.2), ef(47.1, 0.2), wf(47.1, 0.2), 0., 1., 5, atol=1e-12) + vgq(rf(47.1, 0.2), ef(47.1, 0.2), wf(47.1, 0.2), 0., 1., 25, atol=1e-11) + vgq(rf(47.1, 0.2), ef(47.1, 0.2), wf(47.1, 0.2), 0., 1., 100, atol=1e-10) + + vgq(rf(68.9, 2.25), ef(68.9, 2.25), wf(68.9, 2.25), 0., 1., 5, atol=3.5e-14) + vgq(rf(68.9, 2.25), ef(68.9, 2.25), wf(68.9, 2.25), 0., 1., 25, atol=2e-13) + vgq(rf(68.9, 2.25), ef(68.9, 2.25), wf(68.9, 2.25), 0., 1., + 100, atol=1e-12) + + x, w = sc.roots_sh_jacobi(5, 3, 2, False) + y, v, m = sc.roots_sh_jacobi(5, 3, 2, True) + assert_allclose(x, y, 1e-14, 1e-14) + assert_allclose(w, v, 1e-14, 1e-14) + + muI, muI_err = integrate.quad(wf(3,2), 0, 1) + assert_allclose(m, muI, rtol=muI_err) + + assert_raises(ValueError, sc.roots_sh_jacobi, 0, 1, 1) + assert_raises(ValueError, sc.roots_sh_jacobi, 3.3, 1, 1) # type: ignore[call-overload] + assert_raises(ValueError, sc.roots_sh_jacobi, 3, 1, 2) # p - q <= -1 + assert_raises(ValueError, sc.roots_sh_jacobi, 3, 2, -1) # q <= 0 + assert_raises(ValueError, sc.roots_sh_jacobi, 3, -2, -1) # both + +def test_roots_hermite(): + rootf = sc.roots_hermite + evalf = sc.eval_hermite + weightf = orth.hermite(5).weight_func + + verify_gauss_quad(rootf, evalf, weightf, -np.inf, np.inf, 5) + verify_gauss_quad(rootf, evalf, weightf, -np.inf, np.inf, 25, atol=1e-13) + verify_gauss_quad(rootf, evalf, weightf, -np.inf, np.inf, 100, atol=1e-12) + + # Golub-Welsch branch + x, w = sc.roots_hermite(5, False) + y, v, m = sc.roots_hermite(5, True) + assert_allclose(x, y, 1e-14, 1e-14) + assert_allclose(w, v, 1e-14, 1e-14) + + muI, muI_err = integrate.quad(weightf, -np.inf, np.inf) + assert_allclose(m, muI, rtol=muI_err) + + # Asymptotic branch (switch over at n >= 150) + x, w = sc.roots_hermite(200, False) + y, v, m = sc.roots_hermite(200, True) + assert_allclose(x, y, 1e-14, 1e-14) + assert_allclose(w, v, 1e-14, 1e-14) + assert_allclose(sum(v), m, 1e-14, 1e-14) + + assert_raises(ValueError, sc.roots_hermite, 0) + assert_raises(ValueError, sc.roots_hermite, 3.3) # type: ignore[call-overload] + +def test_roots_hermite_asy(): + # Recursion for Hermite functions + def hermite_recursion(n, nodes): + H = np.zeros((n, nodes.size)) + H[0,:] = np.pi**(-0.25) * np.exp(-0.5*nodes**2) + if n > 1: + H[1,:] = sqrt(2.0) * nodes * H[0,:] + for k in range(2, n): + H[k,:] = sqrt(2.0/k) * nodes * H[k-1,:] - sqrt((k-1.0)/k) * H[k-2,:] + return H + + # This tests only the nodes + def test(N, rtol=1e-15, atol=1e-14): + x, w = orth._roots_hermite_asy(N) + H = hermite_recursion(N+1, x) + assert_allclose(H[-1,:], np.zeros(N), rtol, atol) + assert_allclose(sum(w), sqrt(np.pi), rtol, atol) + + test(150, atol=1e-12) + test(151, atol=1e-12) + test(300, atol=1e-12) + test(301, atol=1e-12) + test(500, atol=1e-12) + test(501, atol=1e-12) + test(999, atol=1e-12) + test(1000, atol=1e-12) + test(2000, atol=1e-12) + test(5000, atol=1e-12) + +def test_roots_hermitenorm(): + rootf = sc.roots_hermitenorm + evalf = sc.eval_hermitenorm + weightf = orth.hermitenorm(5).weight_func + + verify_gauss_quad(rootf, evalf, weightf, -np.inf, np.inf, 5) + verify_gauss_quad(rootf, evalf, weightf, -np.inf, np.inf, 25, atol=1e-13) + verify_gauss_quad(rootf, evalf, weightf, -np.inf, np.inf, 100, atol=1e-12) + + x, w = sc.roots_hermitenorm(5, False) + y, v, m = sc.roots_hermitenorm(5, True) + assert_allclose(x, y, 1e-14, 1e-14) + assert_allclose(w, v, 1e-14, 1e-14) + + muI, muI_err = integrate.quad(weightf, -np.inf, np.inf) + assert_allclose(m, muI, rtol=muI_err) + + assert_raises(ValueError, sc.roots_hermitenorm, 0) + assert_raises(ValueError, sc.roots_hermitenorm, 3.3) # type: ignore[call-overload] + +def test_roots_gegenbauer(): + def rootf(a): + return lambda n, mu: sc.roots_gegenbauer(n, a, mu) + def evalf(a): + return lambda n, x: sc.eval_gegenbauer(n, a, x) + def weightf(a): + return lambda x: (1 - x ** 2) ** (a - 0.5) + + vgq = verify_gauss_quad + vgq(rootf(-0.25), evalf(-0.25), weightf(-0.25), -1., 1., 5) + vgq(rootf(-0.25), evalf(-0.25), weightf(-0.25), -1., 1., 25, atol=1e-12) + vgq(rootf(-0.25), evalf(-0.25), weightf(-0.25), -1., 1., 100, atol=1e-11) + + vgq(rootf(0.1), evalf(0.1), weightf(0.1), -1., 1., 5) + vgq(rootf(0.1), evalf(0.1), weightf(0.1), -1., 1., 25, atol=1e-13) + vgq(rootf(0.1), evalf(0.1), weightf(0.1), -1., 1., 100, atol=1e-12) + + vgq(rootf(1), evalf(1), weightf(1), -1., 1., 5) + vgq(rootf(1), evalf(1), weightf(1), -1., 1., 25, atol=1e-13) + vgq(rootf(1), evalf(1), weightf(1), -1., 1., 100, atol=1e-12) + + vgq(rootf(10), evalf(10), weightf(10), -1., 1., 5) + vgq(rootf(10), evalf(10), weightf(10), -1., 1., 25, atol=1e-13) + vgq(rootf(10), evalf(10), weightf(10), -1., 1., 100, atol=1e-12) + + vgq(rootf(50), evalf(50), weightf(50), -1., 1., 5, atol=1e-13) + vgq(rootf(50), evalf(50), weightf(50), -1., 1., 25, atol=1e-12) + vgq(rootf(50), evalf(50), weightf(50), -1., 1., 100, atol=1e-11) + + # Alpha=170 is where the approximation used in roots_gegenbauer changes + vgq(rootf(170), evalf(170), weightf(170), -1., 1., 5, atol=1e-13) + vgq(rootf(170), evalf(170), weightf(170), -1., 1., 25, atol=1e-12) + vgq(rootf(170), evalf(170), weightf(170), -1., 1., 100, atol=1e-11) + vgq(rootf(170.5), evalf(170.5), weightf(170.5), -1., 1., 5, atol=1.25e-13) + vgq(rootf(170.5), evalf(170.5), weightf(170.5), -1., 1., 25, atol=1e-12) + vgq(rootf(170.5), evalf(170.5), weightf(170.5), -1., 1., 100, atol=1e-11) + + # Test for failures, e.g. overflows, resulting from large alphas + vgq(rootf(238), evalf(238), weightf(238), -1., 1., 5, atol=1e-13) + vgq(rootf(238), evalf(238), weightf(238), -1., 1., 25, atol=1e-12) + vgq(rootf(238), evalf(238), weightf(238), -1., 1., 100, atol=1e-11) + vgq(rootf(512.5), evalf(512.5), weightf(512.5), -1., 1., 5, atol=1e-12) + vgq(rootf(512.5), evalf(512.5), weightf(512.5), -1., 1., 25, atol=1e-11) + vgq(rootf(512.5), evalf(512.5), weightf(512.5), -1., 1., 100, atol=1e-10) + + # this is a special case that the old code supported. + # when alpha = 0, the gegenbauer polynomial is uniformly 0. but it goes + # to a scaled down copy of T_n(x) there. + vgq(rootf(0), sc.eval_chebyt, weightf(0), -1., 1., 5) + vgq(rootf(0), sc.eval_chebyt, weightf(0), -1., 1., 25) + vgq(rootf(0), sc.eval_chebyt, weightf(0), -1., 1., 100, atol=1e-12) + + x, w = sc.roots_gegenbauer(5, 2, False) + y, v, m = sc.roots_gegenbauer(5, 2, True) + assert_allclose(x, y, 1e-14, 1e-14) + assert_allclose(w, v, 1e-14, 1e-14) + + muI, muI_err = integrate.quad(weightf(2), -1, 1) + assert_allclose(m, muI, rtol=muI_err) + + assert_raises(ValueError, sc.roots_gegenbauer, 0, 2) + assert_raises(ValueError, sc.roots_gegenbauer, 3.3, 2) # type: ignore[call-overload] + assert_raises(ValueError, sc.roots_gegenbauer, 3, -.75) + +def test_roots_chebyt(): + weightf = orth.chebyt(5).weight_func + verify_gauss_quad(sc.roots_chebyt, sc.eval_chebyt, weightf, -1., 1., 5) + verify_gauss_quad(sc.roots_chebyt, sc.eval_chebyt, weightf, -1., 1., 25) + verify_gauss_quad(sc.roots_chebyt, sc.eval_chebyt, weightf, -1., 1., 100, + atol=1e-12) + + x, w = sc.roots_chebyt(5, False) + y, v, m = sc.roots_chebyt(5, True) + assert_allclose(x, y, 1e-14, 1e-14) + assert_allclose(w, v, 1e-14, 1e-14) + + muI, muI_err = integrate.quad(weightf, -1, 1) + assert_allclose(m, muI, rtol=muI_err) + + assert_raises(ValueError, sc.roots_chebyt, 0) + assert_raises(ValueError, sc.roots_chebyt, 3.3) # type: ignore[call-overload] + +def test_chebyt_symmetry(): + x, w = sc.roots_chebyt(21) + pos, neg = x[:10], x[11:] + assert_equal(neg, -pos[::-1]) + assert_equal(x[10], 0) + +def test_roots_chebyu(): + weightf = orth.chebyu(5).weight_func + verify_gauss_quad(sc.roots_chebyu, sc.eval_chebyu, weightf, -1., 1., 5) + verify_gauss_quad(sc.roots_chebyu, sc.eval_chebyu, weightf, -1., 1., 25) + verify_gauss_quad(sc.roots_chebyu, sc.eval_chebyu, weightf, -1., 1., 100) + + x, w = sc.roots_chebyu(5, False) + y, v, m = sc.roots_chebyu(5, True) + assert_allclose(x, y, 1e-14, 1e-14) + assert_allclose(w, v, 1e-14, 1e-14) + + muI, muI_err = integrate.quad(weightf, -1, 1) + assert_allclose(m, muI, rtol=muI_err) + + assert_raises(ValueError, sc.roots_chebyu, 0) + assert_raises(ValueError, sc.roots_chebyu, 3.3) # type: ignore[call-overload] + +def test_roots_chebyc(): + weightf = orth.chebyc(5).weight_func + verify_gauss_quad(sc.roots_chebyc, sc.eval_chebyc, weightf, -2., 2., 5) + verify_gauss_quad(sc.roots_chebyc, sc.eval_chebyc, weightf, -2., 2., 25) + verify_gauss_quad(sc.roots_chebyc, sc.eval_chebyc, weightf, -2., 2., 100, + atol=1e-12) + + x, w = sc.roots_chebyc(5, False) + y, v, m = sc.roots_chebyc(5, True) + assert_allclose(x, y, 1e-14, 1e-14) + assert_allclose(w, v, 1e-14, 1e-14) + + muI, muI_err = integrate.quad(weightf, -2, 2) + assert_allclose(m, muI, rtol=muI_err) + + assert_raises(ValueError, sc.roots_chebyc, 0) + assert_raises(ValueError, sc.roots_chebyc, 3.3) # type: ignore[call-overload] + +def test_roots_chebys(): + weightf = orth.chebys(5).weight_func + verify_gauss_quad(sc.roots_chebys, sc.eval_chebys, weightf, -2., 2., 5) + verify_gauss_quad(sc.roots_chebys, sc.eval_chebys, weightf, -2., 2., 25) + verify_gauss_quad(sc.roots_chebys, sc.eval_chebys, weightf, -2., 2., 100) + + x, w = sc.roots_chebys(5, False) + y, v, m = sc.roots_chebys(5, True) + assert_allclose(x, y, 1e-14, 1e-14) + assert_allclose(w, v, 1e-14, 1e-14) + + muI, muI_err = integrate.quad(weightf, -2, 2) + assert_allclose(m, muI, rtol=muI_err) + + assert_raises(ValueError, sc.roots_chebys, 0) + assert_raises(ValueError, sc.roots_chebys, 3.3) # type: ignore[call-overload] + +def test_roots_sh_chebyt(): + weightf = orth.sh_chebyt(5).weight_func + verify_gauss_quad(sc.roots_sh_chebyt, sc.eval_sh_chebyt, weightf, 0., 1., 5) + verify_gauss_quad(sc.roots_sh_chebyt, sc.eval_sh_chebyt, weightf, 0., 1., 25) + verify_gauss_quad(sc.roots_sh_chebyt, sc.eval_sh_chebyt, weightf, 0., 1., + 100, atol=1e-13) + + x, w = sc.roots_sh_chebyt(5, False) + y, v, m = sc.roots_sh_chebyt(5, True) + assert_allclose(x, y, 1e-14, 1e-14) + assert_allclose(w, v, 1e-14, 1e-14) + + muI, muI_err = integrate.quad(weightf, 0, 1) + assert_allclose(m, muI, rtol=muI_err) + + assert_raises(ValueError, sc.roots_sh_chebyt, 0) + assert_raises(ValueError, sc.roots_sh_chebyt, 3.3) # type: ignore[call-overload] + +def test_roots_sh_chebyu(): + weightf = orth.sh_chebyu(5).weight_func + verify_gauss_quad(sc.roots_sh_chebyu, sc.eval_sh_chebyu, weightf, 0., 1., 5) + verify_gauss_quad(sc.roots_sh_chebyu, sc.eval_sh_chebyu, weightf, 0., 1., 25) + verify_gauss_quad(sc.roots_sh_chebyu, sc.eval_sh_chebyu, weightf, 0., 1., + 100, atol=1e-13) + + x, w = sc.roots_sh_chebyu(5, False) + y, v, m = sc.roots_sh_chebyu(5, True) + assert_allclose(x, y, 1e-14, 1e-14) + assert_allclose(w, v, 1e-14, 1e-14) + + muI, muI_err = integrate.quad(weightf, 0, 1) + assert_allclose(m, muI, rtol=muI_err) + + assert_raises(ValueError, sc.roots_sh_chebyu, 0) + assert_raises(ValueError, sc.roots_sh_chebyu, 3.3) # type: ignore[call-overload] + +def test_roots_legendre(): + weightf = orth.legendre(5).weight_func + verify_gauss_quad(sc.roots_legendre, sc.eval_legendre, weightf, -1., 1., 5) + verify_gauss_quad(sc.roots_legendre, sc.eval_legendre, weightf, -1., 1., + 25, atol=1e-13) + verify_gauss_quad(sc.roots_legendre, sc.eval_legendre, weightf, -1., 1., + 100, atol=1e-12) + + x, w = sc.roots_legendre(5, False) + y, v, m = sc.roots_legendre(5, True) + assert_allclose(x, y, 1e-14, 1e-14) + assert_allclose(w, v, 1e-14, 1e-14) + + muI, muI_err = integrate.quad(weightf, -1, 1) + assert_allclose(m, muI, rtol=muI_err) + + assert_raises(ValueError, sc.roots_legendre, 0) + assert_raises(ValueError, sc.roots_legendre, 3.3) # type: ignore[call-overload] + +def test_roots_sh_legendre(): + weightf = orth.sh_legendre(5).weight_func + verify_gauss_quad(sc.roots_sh_legendre, sc.eval_sh_legendre, weightf, 0., 1., 5) + verify_gauss_quad(sc.roots_sh_legendre, sc.eval_sh_legendre, weightf, 0., 1., + 25, atol=1e-13) + verify_gauss_quad(sc.roots_sh_legendre, sc.eval_sh_legendre, weightf, 0., 1., + 100, atol=1e-12) + + x, w = sc.roots_sh_legendre(5, False) + y, v, m = sc.roots_sh_legendre(5, True) + assert_allclose(x, y, 1e-14, 1e-14) + assert_allclose(w, v, 1e-14, 1e-14) + + muI, muI_err = integrate.quad(weightf, 0, 1) + assert_allclose(m, muI, rtol=muI_err) + + assert_raises(ValueError, sc.roots_sh_legendre, 0) + assert_raises(ValueError, sc.roots_sh_legendre, 3.3) # type: ignore[call-overload] + +def test_roots_laguerre(): + weightf = orth.laguerre(5).weight_func + verify_gauss_quad(sc.roots_laguerre, sc.eval_laguerre, weightf, 0., np.inf, 5) + verify_gauss_quad(sc.roots_laguerre, sc.eval_laguerre, weightf, 0., np.inf, + 25, atol=1e-13) + verify_gauss_quad(sc.roots_laguerre, sc.eval_laguerre, weightf, 0., np.inf, + 100, atol=1e-12) + + x, w = sc.roots_laguerre(5, False) + y, v, m = sc.roots_laguerre(5, True) + assert_allclose(x, y, 1e-14, 1e-14) + assert_allclose(w, v, 1e-14, 1e-14) + + muI, muI_err = integrate.quad(weightf, 0, np.inf) + assert_allclose(m, muI, rtol=muI_err) + + assert_raises(ValueError, sc.roots_laguerre, 0) + assert_raises(ValueError, sc.roots_laguerre, 3.3) # type: ignore[call-overload] + +def test_roots_genlaguerre(): + def rootf(a): + return lambda n, mu: sc.roots_genlaguerre(n, a, mu) + def evalf(a): + return lambda n, x: sc.eval_genlaguerre(n, a, x) + def weightf(a): + return lambda x: x ** a * np.exp(-x) + + vgq = verify_gauss_quad + vgq(rootf(-0.5), evalf(-0.5), weightf(-0.5), 0., np.inf, 5) + vgq(rootf(-0.5), evalf(-0.5), weightf(-0.5), 0., np.inf, 25, atol=1e-13) + vgq(rootf(-0.5), evalf(-0.5), weightf(-0.5), 0., np.inf, 100, atol=1e-12) + + vgq(rootf(0.1), evalf(0.1), weightf(0.1), 0., np.inf, 5) + vgq(rootf(0.1), evalf(0.1), weightf(0.1), 0., np.inf, 25, atol=1e-13) + vgq(rootf(0.1), evalf(0.1), weightf(0.1), 0., np.inf, 100, atol=1.6e-13) + + vgq(rootf(1), evalf(1), weightf(1), 0., np.inf, 5) + vgq(rootf(1), evalf(1), weightf(1), 0., np.inf, 25, atol=1e-13) + vgq(rootf(1), evalf(1), weightf(1), 0., np.inf, 100, atol=1.03e-13) + + vgq(rootf(10), evalf(10), weightf(10), 0., np.inf, 5) + vgq(rootf(10), evalf(10), weightf(10), 0., np.inf, 25, atol=1e-13) + vgq(rootf(10), evalf(10), weightf(10), 0., np.inf, 100, atol=1e-12) + + vgq(rootf(50), evalf(50), weightf(50), 0., np.inf, 5) + vgq(rootf(50), evalf(50), weightf(50), 0., np.inf, 25, atol=1e-13) + vgq(rootf(50), evalf(50), weightf(50), 0., np.inf, 100, rtol=1e-14, atol=2e-13) + + x, w = sc.roots_genlaguerre(5, 2, False) + y, v, m = sc.roots_genlaguerre(5, 2, True) + assert_allclose(x, y, 1e-14, 1e-14) + assert_allclose(w, v, 1e-14, 1e-14) + + muI, muI_err = integrate.quad(weightf(2.), 0., np.inf) + assert_allclose(m, muI, rtol=muI_err) + + assert_raises(ValueError, sc.roots_genlaguerre, 0, 2) + assert_raises(ValueError, sc.roots_genlaguerre, 3.3, 2) # type: ignore[call-overload] + assert_raises(ValueError, sc.roots_genlaguerre, 3, -1.1) + + +def test_gh_6721(): + # Regression test for gh_6721. This should not raise. + sc.chebyt(65)(0.2) diff --git a/v12x/pylib/scipy/special/tests/test_orthogonal_eval.py b/v12x/pylib/scipy/special/tests/test_orthogonal_eval.py new file mode 100644 index 0000000000000000000000000000000000000000..b88027d4cf531c4dc0c5ad9c80e2a5c264f06844 --- /dev/null +++ b/v12x/pylib/scipy/special/tests/test_orthogonal_eval.py @@ -0,0 +1,383 @@ +import numpy as np +from numpy.testing import assert_, assert_allclose +import pytest + +from scipy.special import _ufuncs, gammasgn +import scipy.special._orthogonal as orth +from scipy.special._testutils import FuncData + + +def test_eval_chebyt(): + n = np.arange(0, 10000, 7, dtype=np.dtype("long")) + x = 2*np.random.rand() - 1 + v1 = np.cos(n*np.arccos(x)) + v2 = _ufuncs.eval_chebyt(n, x) + assert_(np.allclose(v1, v2, rtol=1e-15)) + + +def test_eval_chebyt_gh20129(): + # https://github.com/scipy/scipy/issues/20129 + assert _ufuncs.eval_chebyt(7, 2 + 0j) == 5042.0 + + +def test_eval_genlaguerre_restriction(): + # check it returns nan for alpha <= -1 + assert_(np.isnan(_ufuncs.eval_genlaguerre(0, -1, 0))) + assert_(np.isnan(_ufuncs.eval_genlaguerre(0.1, -1, 0))) + + +def test_warnings(): + # ticket 1334 + with np.errstate(all='raise'): + # these should raise no fp warnings + _ufuncs.eval_legendre(1, 0) + _ufuncs.eval_laguerre(1, 1) + _ufuncs.eval_gegenbauer(1, 1, 0) + + +class TestPolys: + """ + Check that the eval_* functions agree with the constructed polynomials + + """ + + def check_poly(self, func, cls, param_ranges=(), x_range=(), nn=10, + nparam=10, nx=10, rtol=1e-8): + rng = np.random.RandomState(1234) + + dataset = [] + for n in np.arange(nn): + params = [a + (b-a)*rng.rand(nparam) for a,b in param_ranges] + params = np.asarray(params).T + if not param_ranges: + params = [0] + for p in params: + if param_ranges: + p = (n,) + tuple(p) + else: + p = (n,) + x = x_range[0] + (x_range[1] - x_range[0])*rng.rand(nx) + x[0] = x_range[0] # always include domain start point + x[1] = x_range[1] # always include domain end point + poly = np.poly1d(cls(*p).coef) + z = np.c_[np.tile(p, (nx,1)), x, poly(x)] + dataset.append(z) + + dataset = np.concatenate(dataset, axis=0) + + def polyfunc(*p): + p = (p[0].astype(np.dtype("long")),) + p[1:] + return func(*p) + + with np.errstate(all='raise'): + ds = FuncData(polyfunc, dataset, list(range(len(param_ranges)+2)), -1, + rtol=rtol) + ds.check() + + def test_jacobi(self): + self.check_poly(_ufuncs.eval_jacobi, orth.jacobi, + param_ranges=[(-0.99, 10), (-0.99, 10)], + x_range=[-1, 1], rtol=1e-5) + + def test_sh_jacobi(self): + self.check_poly(_ufuncs.eval_sh_jacobi, orth.sh_jacobi, + param_ranges=[(1, 10), (0, 1)], x_range=[0, 1], + rtol=1e-5) + + def test_gegenbauer(self): + self.check_poly(_ufuncs.eval_gegenbauer, orth.gegenbauer, + param_ranges=[(-0.499, 10)], x_range=[-1, 1], + rtol=1e-7) + + def test_chebyt(self): + self.check_poly(_ufuncs.eval_chebyt, orth.chebyt, + param_ranges=[], x_range=[-1, 1]) + + def test_chebyu(self): + self.check_poly(_ufuncs.eval_chebyu, orth.chebyu, + param_ranges=[], x_range=[-1, 1]) + + def test_chebys(self): + self.check_poly(_ufuncs.eval_chebys, orth.chebys, + param_ranges=[], x_range=[-2, 2]) + + def test_chebyc(self): + self.check_poly(_ufuncs.eval_chebyc, orth.chebyc, + param_ranges=[], x_range=[-2, 2]) + + def test_sh_chebyt(self): + with np.errstate(all='ignore'): + self.check_poly(_ufuncs.eval_sh_chebyt, orth.sh_chebyt, + param_ranges=[], x_range=[0, 1]) + + def test_sh_chebyu(self): + self.check_poly(_ufuncs.eval_sh_chebyu, orth.sh_chebyu, + param_ranges=[], x_range=[0, 1]) + + def test_legendre(self): + self.check_poly(_ufuncs.eval_legendre, orth.legendre, + param_ranges=[], x_range=[-1, 1]) + + def test_sh_legendre(self): + with np.errstate(all='ignore'): + self.check_poly(_ufuncs.eval_sh_legendre, orth.sh_legendre, + param_ranges=[], x_range=[0, 1]) + + def test_genlaguerre(self): + self.check_poly(_ufuncs.eval_genlaguerre, orth.genlaguerre, + param_ranges=[(-0.99, 10)], x_range=[0, 100]) + + def test_laguerre(self): + self.check_poly(_ufuncs.eval_laguerre, orth.laguerre, + param_ranges=[], x_range=[0, 100]) + + def test_hermite(self): + self.check_poly(_ufuncs.eval_hermite, orth.hermite, + param_ranges=[], x_range=[-100, 100]) + + def test_hermitenorm(self): + self.check_poly(_ufuncs.eval_hermitenorm, orth.hermitenorm, + param_ranges=[], x_range=[-100, 100]) + + +class TestRecurrence: + """ + Check that the eval_* functions sig='ld->d' and 'dd->d' agree. + + """ + + def check_poly(self, func, param_ranges=(), x_range=(), nn=10, + nparam=10, nx=10, rtol=1e-8): + rng = np.random.default_rng(1234) + + dataset = [] + for n in np.arange(nn): + params = [a + (b-a)*rng.random(nparam) for a,b in param_ranges] + params = np.asarray(params).T + if not param_ranges: + params = [0] + for p in params: + if param_ranges: + p = (n,) + tuple(p) + else: + p = (n,) + x = x_range[0] + (x_range[1] - x_range[0])*rng.random(nx) + x[0] = x_range[0] # always include domain start point + x[1] = x_range[1] # always include domain end point + kw = dict(sig=(len(p)+1)*'d'+'->d') + z = np.c_[np.tile(p, (nx,1)), x, func(*(p + (x,)), **kw)] + dataset.append(z) + + dataset = np.concatenate(dataset, axis=0) + + def polyfunc(*p): + p0 = p[0].astype(np.intp) + p = (p0,) + p[1:] + p0_type_char = p0.dtype.char + kw = dict(sig=p0_type_char + (len(p)-1)*'d' + '->d') + return func(*p, **kw) + + with np.errstate(all='raise'): + ds = FuncData(polyfunc, dataset, list(range(len(param_ranges)+2)), -1, + rtol=rtol) + ds.check() + + def test_jacobi(self): + self.check_poly(_ufuncs.eval_jacobi, + param_ranges=[(-0.99, 10), (-0.99, 10)], + x_range=[-1, 1]) + + def test_sh_jacobi(self): + self.check_poly(_ufuncs.eval_sh_jacobi, + param_ranges=[(1, 10), (0, 1)], x_range=[0, 1]) + + def test_gegenbauer(self): + self.check_poly(_ufuncs.eval_gegenbauer, + param_ranges=[(-0.499, 10)], x_range=[-1, 1]) + + def test_chebyt(self): + self.check_poly(_ufuncs.eval_chebyt, + param_ranges=[], x_range=[-1, 1]) + + def test_chebyu(self): + self.check_poly(_ufuncs.eval_chebyu, + param_ranges=[], x_range=[-1, 1]) + + def test_chebys(self): + self.check_poly(_ufuncs.eval_chebys, + param_ranges=[], x_range=[-2, 2]) + + def test_chebyc(self): + self.check_poly(_ufuncs.eval_chebyc, + param_ranges=[], x_range=[-2, 2]) + + def test_sh_chebyt(self): + self.check_poly(_ufuncs.eval_sh_chebyt, + param_ranges=[], x_range=[0, 1]) + + def test_sh_chebyu(self): + self.check_poly(_ufuncs.eval_sh_chebyu, + param_ranges=[], x_range=[0, 1]) + + def test_legendre(self): + self.check_poly(_ufuncs.eval_legendre, + param_ranges=[], x_range=[-1, 1]) + + def test_sh_legendre(self): + self.check_poly(_ufuncs.eval_sh_legendre, + param_ranges=[], x_range=[0, 1]) + + def test_genlaguerre(self): + self.check_poly(_ufuncs.eval_genlaguerre, + param_ranges=[(-0.99, 10)], x_range=[0, 100]) + + def test_laguerre(self): + self.check_poly(_ufuncs.eval_laguerre, + param_ranges=[], x_range=[0, 100]) + + def test_hermite(self): + v = _ufuncs.eval_hermite(70, 1.0) + a = -1.457076485701412e60 + assert_allclose(v, a) + + +def test_hermite_domain(): + # Regression test for gh-11091. + assert np.isnan(_ufuncs.eval_hermite(-1, 1.0)) + assert np.isnan(_ufuncs.eval_hermitenorm(-1, 1.0)) + + +@pytest.mark.parametrize("n", [0, 1, 2]) +@pytest.mark.parametrize("x", [0, 1, np.nan]) +def test_hermite_nan(n, x): + # Regression test for gh-11369. + assert np.isnan(_ufuncs.eval_hermite(n, x)) == np.any(np.isnan([n, x])) + assert np.isnan(_ufuncs.eval_hermitenorm(n, x)) == np.any(np.isnan([n, x])) + + +@pytest.mark.parametrize('n', [0, 1, 2, 3.2]) +@pytest.mark.parametrize('alpha', [1, np.nan]) +@pytest.mark.parametrize('x', [2, np.nan]) +def test_genlaguerre_nan(n, alpha, x): + # Regression test for gh-11361. + nan_laguerre = np.isnan(_ufuncs.eval_genlaguerre(n, alpha, x)) + nan_arg = np.any(np.isnan([n, alpha, x])) + assert nan_laguerre == nan_arg + + +@pytest.mark.parametrize('n', [0, 1, 2, 3.2, np.nan]) +@pytest.mark.parametrize('alpha', [0.0, 1, np.nan]) +@pytest.mark.parametrize('x', [1e-6, 2, np.nan]) +def test_gegenbauer_nan(n, alpha, x): + # Regression test for gh-11370. + nan_gegenbauer = np.isnan(_ufuncs.eval_gegenbauer(n, alpha, x)) + nan_arg = np.any(np.isnan([n, alpha, x])) + assert nan_gegenbauer == nan_arg + +@pytest.mark.parametrize("n", np.arange(10)) +@pytest.mark.parametrize("alpha", [-0.45, -0.25, -0.125, 0.0, 0.25, 1.0, 2.0]) +@pytest.mark.parametrize("x", [-np.inf, np.inf]) +def test_gegenbauer_infinity(n, alpha, x): + # gh-11713 - check correct handling of x = +inf and x = -inf + if alpha == 0.0: + expected = 0.0 + elif n == 0.0: + expected = 1.0 + else: + # sign of leading coefficient: 2^n * Gamma(n+alpha) / (n! Gamma(alpha)) + lead_sign = gammasgn(n + alpha) * gammasgn(alpha) + expected = lead_sign * (np.sign(x) ** n) * np.inf + assert_allclose(_ufuncs.eval_gegenbauer(int(n), alpha, x), expected, rtol=1e-10) + assert_allclose(_ufuncs.eval_gegenbauer(float(n), alpha, x), expected, rtol=1e-10) + +@pytest.mark.parametrize( + "n, expected", + [ + (0, [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]), + (1, [-2.0, -1.8, -1.6, -1.4, -1.2, -1.0, -0.8, -0.6, -0.4, -0.2, 0.0]), + (2, [3.0, 2.16, 1.44, 0.84, 0.36, 0.0, -0.24, -0.36, -0.36, -0.24, 0.0]), + (3, [-4.0, -1.98, -0.64, 0.14, 0.48, 0.5, 0.32, 0.06, -0.16, -0.22, 0.0]), + ( + 4, + [5.0, 1.332, -0.288, -0.658, -0.408, 0.0, 0.272, 0.282, 0.072, -0.148, 0.0], + ), + ( + 5, + [ + -6.0, + -0.43308, + 0.79104, + 0.36876, + -0.21312, + -0.375, + -0.14208, + 0.15804, + 0.19776, + -0.04812, + 0.0, + ], + ), + ], +) +def test_jacobi_alpha_minus_one_beta_plus_one(n, expected): + # gh-7001 - expected values were computed with mathematica + x = np.linspace(-1.0, 1.0, 11) + a, b = -1, 1 # alpha, beta + assert_allclose(_ufuncs.eval_jacobi(n, a, b, x), expected, rtol=1e-10, atol=1e-14) + assert_allclose( + _ufuncs.eval_jacobi(float(n), a, b, x), expected, rtol=1e-10, atol=1e-14 + ) + + +@pytest.mark.parametrize( + "n, expected", + [ + (0, [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]), + (1, [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]), + (2, [0.0, -0.09, -0.16, -0.21, -0.24, -0.25, -0.24, -0.21, -0.16, -0.09, 0.0]), + ( + 3, + [0.0, 0.144, 0.192, 0.168, 0.096, 0.0, -0.096, -0.168, -0.192, -0.144, 0.0], + ), + ( + 4, + [ + 0.0, + -0.1485, + -0.096, + 0.0315, + 0.144, + 0.1875, + 0.144, + 0.0315, + -0.096, + -0.1485, + 0.0, + ], + ), + ( + 5, + [ + 0.0, + 0.10656, + -0.04608, + -0.15792, + -0.13056, + 0.0, + 0.13056, + 0.15792, + 0.04608, + -0.10656, + 0.0, + ], + ), + ], +) +def test_jacobi_alpha_minus_one_beta_minus_one(n, expected): + # gh-7001 - expected values were computed with mathematica + x = np.linspace(-1.0, 1.0, 11) + a, b = -1, -1 # alpha, beta + assert_allclose(_ufuncs.eval_jacobi(n, a, b, x), expected, rtol=1e-10, atol=1e-14) + assert_allclose( + _ufuncs.eval_jacobi(float(n), a, b, x), expected, rtol=1e-10, atol=1e-14 + ) diff --git a/v12x/pylib/scipy/special/tests/test_owens_t.py b/v12x/pylib/scipy/special/tests/test_owens_t.py new file mode 100644 index 0000000000000000000000000000000000000000..8d15aead25302023c5f07d8392c0931995764ced --- /dev/null +++ b/v12x/pylib/scipy/special/tests/test_owens_t.py @@ -0,0 +1,53 @@ +import numpy as np +from numpy.testing import assert_equal, assert_allclose + +import scipy.special as sc + + +def test_symmetries(): + np.random.seed(1234) + a, h = np.random.rand(100), np.random.rand(100) + assert_equal(sc.owens_t(h, a), sc.owens_t(-h, a)) + assert_equal(sc.owens_t(h, a), -sc.owens_t(h, -a)) + + +def test_special_cases(): + assert_equal(sc.owens_t(5, 0), 0) + assert_allclose(sc.owens_t(0, 5), 0.5*np.arctan(5)/np.pi, + rtol=5e-14) + # Target value is 0.5*Phi(5)*(1 - Phi(5)) for Phi the CDF of the + # standard normal distribution + assert_allclose(sc.owens_t(5, 1), 1.4332574485503512543e-07, + rtol=5e-14) + + +def test_nans(): + assert_equal(sc.owens_t(20, np.nan), np.nan) + assert_equal(sc.owens_t(np.nan, 20), np.nan) + assert_equal(sc.owens_t(np.nan, np.nan), np.nan) + + +def test_infs(): + h, a = 0, np.inf + # T(0, a) = 1/2π * arctan(a) + res = 1/(2*np.pi) * np.arctan(a) + assert_allclose(sc.owens_t(h, a), res, rtol=5e-14) + assert_allclose(sc.owens_t(h, -a), -res, rtol=5e-14) + + h = 1 + # Refer Owens T function definition in Wikipedia + # https://en.wikipedia.org/wiki/Owen%27s_T_function + # Value approximated through Numerical Integration + # using scipy.integrate.quad + # quad(lambda x: 1/(2*pi)*(exp(-0.5*(1*1)*(1+x*x))/(1+x*x)), 0, inf) + res = 0.07932762696572854 + assert_allclose(sc.owens_t(h, np.inf), res, rtol=5e-14) + assert_allclose(sc.owens_t(h, -np.inf), -res, rtol=5e-14) + + assert_equal(sc.owens_t(np.inf, 1), 0) + assert_equal(sc.owens_t(-np.inf, 1), 0) + + assert_equal(sc.owens_t(np.inf, np.inf), 0) + assert_equal(sc.owens_t(-np.inf, np.inf), 0) + assert_equal(sc.owens_t(np.inf, -np.inf), -0.0) + assert_equal(sc.owens_t(-np.inf, -np.inf), -0.0) diff --git a/v12x/pylib/scipy/special/tests/test_pcf.py b/v12x/pylib/scipy/special/tests/test_pcf.py new file mode 100644 index 0000000000000000000000000000000000000000..a8c42aa688081fb58f79ad2c8ea932d03b33523b --- /dev/null +++ b/v12x/pylib/scipy/special/tests/test_pcf.py @@ -0,0 +1,24 @@ +"""Tests for parabolic cylinder functions. + +""" +import numpy as np +from numpy.testing import assert_allclose, assert_equal +import scipy.special as sc + + +def test_pbwa_segfault(): + # Regression test for https://github.com/scipy/scipy/issues/6208. + # + # Data generated by mpmath. + # + w = 1.02276567211316867161 + wp = -0.48887053372346189882 + assert_allclose(sc.pbwa(0, 0), (w, wp), rtol=1e-13, atol=0) + + +def test_pbwa_nan(): + # Check that NaN's are returned outside of the range in which the + # implementation is accurate. + pts = [(-6, -6), (-6, 6), (6, -6), (6, 6)] + for p in pts: + assert_equal(sc.pbwa(*p), (np.nan, np.nan)) diff --git a/v12x/pylib/scipy/special/tests/test_pdtr.py b/v12x/pylib/scipy/special/tests/test_pdtr.py new file mode 100644 index 0000000000000000000000000000000000000000..f2159e6bd3acd60f4bdc98fc9ce0f6d44ac76485 --- /dev/null +++ b/v12x/pylib/scipy/special/tests/test_pdtr.py @@ -0,0 +1,48 @@ +import numpy as np +import scipy.special as sc +from numpy.testing import assert_allclose, assert_array_equal + + +class TestPdtr: + def test(self): + val = sc.pdtr(0, 1) + assert_allclose(val, np.exp(-1), atol=1.5e-7, rtol=0) + + def test_m_zero(self): + val = sc.pdtr([0, 1, 2], 0) + assert_allclose(val, [1, 1, 1], atol=1.5e-7, rtol=0) + + def test_rounding(self): + double_val = sc.pdtr([0.1, 1.1, 2.1], 1.0) + int_val = sc.pdtr([0, 1, 2], 1.0) + assert_array_equal(double_val, int_val) + + def test_inf(self): + val = sc.pdtr(np.inf, 1.0) + assert_allclose(val, 1.0, atol=1.5e-7, rtol=0) + + def test_domain(self): + val = sc.pdtr(-1.1, 1.0) + assert np.isnan(val) + +class TestPdtrc: + def test_value(self): + val = sc.pdtrc(0, 1) + assert_allclose(val, 1 - np.exp(-1), atol=1.5e-7, rtol=0) + + def test_m_zero(self): + val = sc.pdtrc([0, 1, 2], 0.0) + assert_array_equal(val, [0, 0, 0]) + + def test_rounding(self): + double_val = sc.pdtrc([0.1, 1.1, 2.1], 1.0) + int_val = sc.pdtrc([0, 1, 2], 1.0) + assert_array_equal(double_val, int_val) + + def test_inf(self): + val = sc.pdtrc(np.inf, 1.0) + assert_allclose(val, 0.0, atol=1.5e-7, rtol=0) + + def test_domain(self): + val = sc.pdtrc(-1.1, 1.0) + assert np.isnan(val) diff --git a/v12x/pylib/scipy/special/tests/test_powm1.py b/v12x/pylib/scipy/special/tests/test_powm1.py new file mode 100644 index 0000000000000000000000000000000000000000..3d809963f64ddaedf6b59de80dcd5f7ca8fa18a9 --- /dev/null +++ b/v12x/pylib/scipy/special/tests/test_powm1.py @@ -0,0 +1,65 @@ +import pytest +import numpy as np +from numpy.testing import assert_allclose +from scipy.special import powm1 + + +# Expected values were computed with mpmath, e.g. +# +# >>> import mpmath +# >>> mpmath.np.dps = 200 +# >>> print(float(mpmath.powm1(2.0, 1e-7)) +# 6.931472045825965e-08 +# +powm1_test_cases = [ + (1.25, 0.75, 0.18217701125396976, 1e-15), + (2.0, 1e-7, 6.931472045825965e-08, 1e-15), + (25.0, 5e-11, 1.6094379125636148e-10, 1e-15), + (0.99996, 0.75, -3.0000150002530058e-05, 1e-15), + (0.9999999999990905, 20, -1.81898940353014e-11, 1e-15), + (-1.25, 751.0, -6.017550852453444e+72, 2e-15) +] + + +@pytest.mark.parametrize('x, y, expected, rtol', powm1_test_cases) +def test_powm1(x, y, expected, rtol): + p = powm1(x, y) + assert_allclose(p, expected, rtol=rtol) + + +@pytest.mark.parametrize('x, y, expected', + [(0.0, 0.0, 0.0), + (0.0, -1.5, np.inf), + (0.0, 1.75, -1.0), + (-1.5, 2.0, 1.25), + (-1.5, 3.0, -4.375), + (np.nan, 0.0, 0.0), + (1.0, np.nan, 0.0), + (1.0, np.inf, 0.0), + (1.0, -np.inf, 0.0), + (np.inf, 7.5, np.inf), + (np.inf, -7.5, -1.0), + (3.25, np.inf, np.inf), + (np.inf, np.inf, np.inf), + (np.inf, -np.inf, -1.0), + (np.inf, 0.0, 0.0), + (-np.inf, 0.0, 0.0), + (-np.inf, 2.0, np.inf), + (-np.inf, 3.0, -np.inf), + (-1.0, float(2**53 - 1), -2.0)]) +def test_powm1_exact_cases(x, y, expected): + # Test cases where we have an exact expected value. + p = powm1(x, y) + assert p == expected + + +@pytest.mark.parametrize('x, y', + [(-1.25, 751.03), + (-1.25, np.inf), + (np.nan, np.nan), + (-np.inf, -np.inf), + (-np.inf, 2.5)]) +def test_powm1_return_nan(x, y): + # Test cases where the expected return value is nan. + p = powm1(x, y) + assert np.isnan(p) diff --git a/v12x/pylib/scipy/special/tests/test_precompute_expn_asy.py b/v12x/pylib/scipy/special/tests/test_precompute_expn_asy.py new file mode 100644 index 0000000000000000000000000000000000000000..7b6c6cba21d5c1d5fdfab879153ba8f125e98d5f --- /dev/null +++ b/v12x/pylib/scipy/special/tests/test_precompute_expn_asy.py @@ -0,0 +1,24 @@ +from numpy.testing import assert_equal + +from scipy.special._testutils import check_version, MissingModule +from scipy.special._precompute.expn_asy import generate_A + +try: + import sympy + from sympy import Poly +except ImportError: + sympy = MissingModule("sympy") + + +@check_version(sympy, "1.0") +def test_generate_A(): + # Data from DLMF 8.20.5 + x = sympy.symbols('x') + Astd = [Poly(1, x), + Poly(1, x), + Poly(1 - 2*x), + Poly(1 - 8*x + 6*x**2)] + Ares = generate_A(len(Astd)) + + for p, q in zip(Astd, Ares): + assert_equal(p, q) diff --git a/v12x/pylib/scipy/special/tests/test_precompute_gammainc.py b/v12x/pylib/scipy/special/tests/test_precompute_gammainc.py new file mode 100644 index 0000000000000000000000000000000000000000..86278b065a58bdde3fcde32a2f7ec0d53d17077d --- /dev/null +++ b/v12x/pylib/scipy/special/tests/test_precompute_gammainc.py @@ -0,0 +1,112 @@ +import pytest + +from scipy.special._testutils import MissingModule, check_version +from scipy.special._mptestutils import ( + Arg, IntArg, mp_assert_allclose, assert_mpmath_equal) +from scipy.special._precompute.gammainc_asy import ( + compute_g, compute_alpha, compute_d) +from scipy.special._precompute.gammainc_data import gammainc, gammaincc + +try: + import sympy +except ImportError: + sympy = MissingModule('sympy') + +try: + import mpmath as mp +except ImportError: + mp = MissingModule('mpmath') + + +pytestmark = pytest.mark.thread_unsafe( + reason=("mpmath gmpy2 backend is not thread-safe, " + "see https://github.com/mpmath/mpmath/issues/974")) + +@check_version(mp, '0.19') +def test_g(): + # Test data for the g_k. See DLMF 5.11.4. + with mp.workdps(30): + g = [mp.mpf(1), mp.mpf(1)/12, mp.mpf(1)/288, + -mp.mpf(139)/51840, -mp.mpf(571)/2488320, + mp.mpf(163879)/209018880, mp.mpf(5246819)/75246796800] + mp_assert_allclose(compute_g(7), g) + + +@pytest.mark.slow +@check_version(mp, '0.19') +@check_version(sympy, '0.7') +@pytest.mark.xfail_on_32bit("rtol only 2e-11, see gh-6938") +def test_alpha(): + # Test data for the alpha_k. See DLMF 8.12.14. + with mp.workdps(30): + alpha = [mp.mpf(0), mp.mpf(1), mp.mpf(1)/3, mp.mpf(1)/36, + -mp.mpf(1)/270, mp.mpf(1)/4320, mp.mpf(1)/17010, + -mp.mpf(139)/5443200, mp.mpf(1)/204120] + mp_assert_allclose(compute_alpha(9), alpha) + + +@pytest.mark.xslow +@check_version(mp, '0.19') +@check_version(sympy, '0.7') +def test_d(): + # Compare the d_{k, n} to the results in appendix F of [1]. + # + # Sources + # ------- + # [1] DiDonato and Morris, Computation of the Incomplete Gamma + # Function Ratios and their Inverse, ACM Transactions on + # Mathematical Software, 1986. + + with mp.workdps(50): + dataset = [(0, 0, -mp.mpf('0.333333333333333333333333333333')), + (0, 12, mp.mpf('0.102618097842403080425739573227e-7')), + (1, 0, -mp.mpf('0.185185185185185185185185185185e-2')), + (1, 12, mp.mpf('0.119516285997781473243076536700e-7')), + (2, 0, mp.mpf('0.413359788359788359788359788360e-2')), + (2, 12, -mp.mpf('0.140925299108675210532930244154e-7')), + (3, 0, mp.mpf('0.649434156378600823045267489712e-3')), + (3, 12, -mp.mpf('0.191111684859736540606728140873e-7')), + (4, 0, -mp.mpf('0.861888290916711698604702719929e-3')), + (4, 12, mp.mpf('0.288658297427087836297341274604e-7')), + (5, 0, -mp.mpf('0.336798553366358150308767592718e-3')), + (5, 12, mp.mpf('0.482409670378941807563762631739e-7')), + (6, 0, mp.mpf('0.531307936463992223165748542978e-3')), + (6, 12, -mp.mpf('0.882860074633048352505085243179e-7')), + (7, 0, mp.mpf('0.344367606892377671254279625109e-3')), + (7, 12, -mp.mpf('0.175629733590604619378669693914e-6')), + (8, 0, -mp.mpf('0.652623918595309418922034919727e-3')), + (8, 12, mp.mpf('0.377358774161109793380344937299e-6')), + (9, 0, -mp.mpf('0.596761290192746250124390067179e-3')), + (9, 12, mp.mpf('0.870823417786464116761231237189e-6'))] + d = compute_d(10, 13) + res = [d[k][n] for k, n, std in dataset] + std = [x[2] for x in dataset] + mp_assert_allclose(res, std) + + +@check_version(mp, '0.19') +def test_gammainc(): + # Quick check that the gammainc in + # special._precompute.gammainc_data agrees with mpmath's + # gammainc. + assert_mpmath_equal(gammainc, + lambda a, x: mp.gammainc(a, b=x, regularized=True), + [Arg(0, 100, inclusive_a=False), Arg(0, 100)], + nan_ok=False, rtol=1e-17, n=50, dps=50) + + +@pytest.mark.xslow +@check_version(mp, '0.19') +def test_gammaincc(): + # Check that the gammaincc in special._precompute.gammainc_data + # agrees with mpmath's gammainc. + assert_mpmath_equal(lambda a, x: gammaincc(a, x, dps=1000), + lambda a, x: mp.gammainc(a, a=x, regularized=True), + [Arg(20, 100), Arg(20, 100)], + nan_ok=False, rtol=1e-17, n=50, dps=1000) + + # Test the fast integer path + assert_mpmath_equal(gammaincc, + lambda a, x: mp.gammainc(a, a=x, regularized=True), + [IntArg(1, 100), Arg(0, 100)], + nan_ok=False, rtol=1e-17, n=50, dps=50) diff --git a/v12x/pylib/scipy/special/tests/test_precompute_utils.py b/v12x/pylib/scipy/special/tests/test_precompute_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..89616b92329691ca76039fe11a7e08f7f3db1150 --- /dev/null +++ b/v12x/pylib/scipy/special/tests/test_precompute_utils.py @@ -0,0 +1,36 @@ +import pytest + +from scipy.special._testutils import MissingModule, check_version +from scipy.special._mptestutils import mp_assert_allclose +from scipy.special._precompute.utils import lagrange_inversion + +try: + import sympy +except ImportError: + sympy = MissingModule('sympy') + +try: + import mpmath as mp +except ImportError: + mp = MissingModule('mpmath') + + +@pytest.mark.slow +@check_version(sympy, '0.7') +@check_version(mp, '0.19') +class TestInversion: + @pytest.mark.xfail_on_32bit("rtol only 2e-9, see gh-6938") + def test_log(self): + with mp.workdps(30): + logcoeffs = mp.taylor(lambda x: mp.log(1 + x), 0, 10) + expcoeffs = mp.taylor(lambda x: mp.exp(x) - 1, 0, 10) + invlogcoeffs = lagrange_inversion(logcoeffs) + mp_assert_allclose(invlogcoeffs, expcoeffs) + + @pytest.mark.xfail_on_32bit("rtol only 1e-15, see gh-6938") + def test_sin(self): + with mp.workdps(30): + sincoeffs = mp.taylor(mp.sin, 0, 10) + asincoeffs = mp.taylor(mp.asin, 0, 10) + invsincoeffs = lagrange_inversion(sincoeffs) + mp_assert_allclose(invsincoeffs, asincoeffs, atol=1e-30) diff --git a/v12x/pylib/scipy/special/tests/test_round.py b/v12x/pylib/scipy/special/tests/test_round.py new file mode 100644 index 0000000000000000000000000000000000000000..ba28b7c1cff95327df0ed69086dd92b93399c3c4 --- /dev/null +++ b/v12x/pylib/scipy/special/tests/test_round.py @@ -0,0 +1,18 @@ +import numpy as np +import pytest + +from scipy.special import _test_internal + + +@pytest.mark.fail_slow(20) +@pytest.mark.skipif(not _test_internal.have_fenv(), reason="no fenv()") +def test_add_round_up(): + rng = np.random.RandomState(1234) + _test_internal.test_add_round(10**5, 'up', rng) + + +@pytest.mark.fail_slow(20) +@pytest.mark.skipif(not _test_internal.have_fenv(), reason="no fenv()") +def test_add_round_down(): + rng = np.random.RandomState(1234) + _test_internal.test_add_round(10**5, 'down', rng) diff --git a/v12x/pylib/scipy/special/tests/test_sf_error.py b/v12x/pylib/scipy/special/tests/test_sf_error.py new file mode 100644 index 0000000000000000000000000000000000000000..337bd0c65ebd712bfda1385072c15a47e8a495c2 --- /dev/null +++ b/v12x/pylib/scipy/special/tests/test_sf_error.py @@ -0,0 +1,153 @@ +import sys +import warnings + +import numpy as np +from numpy.testing import assert_, assert_equal, HAS_REFCOUNT +import pytest +from pytest import raises as assert_raises + +import scipy.special as sc +from scipy.special._ufuncs import _sf_error_test_function + +_sf_error_code_map = { + # skip 'ok' + 'singular': 1, + 'underflow': 2, + 'overflow': 3, + 'slow': 4, + 'loss': 5, + 'no_result': 6, + 'domain': 7, + 'arg': 8, + 'other': 9, + 'memory': 10, +} + +_sf_error_actions = [ + 'ignore', + 'warn', + 'raise' +] + + +def _check_action(fun, args, action): + # TODO: special expert should correct + # the coercion at the true location? + args = np.asarray(args, dtype=np.dtype("long")) + if action == 'warn': + with pytest.warns(sc.SpecialFunctionWarning): + fun(*args) + elif action == 'raise': + with assert_raises(sc.SpecialFunctionError): + fun(*args) + else: + # action == 'ignore', make sure there are no warnings/exceptions + with warnings.catch_warnings(): + warnings.simplefilter("error") + fun(*args) + + +def test_geterr(): + err = sc.geterr() + for key, value in err.items(): + assert_(key in _sf_error_code_map) + assert_(value in _sf_error_actions) + + +def test_seterr(): + entry_err = sc.geterr() + try: + for category, error_code in _sf_error_code_map.items(): + for action in _sf_error_actions: + geterr_olderr = sc.geterr() + seterr_olderr = sc.seterr(**{category: action}) + assert_(geterr_olderr == seterr_olderr) + newerr = sc.geterr() + assert_(newerr[category] == action) + geterr_olderr.pop(category) + newerr.pop(category) + assert_(geterr_olderr == newerr) + _check_action(_sf_error_test_function, (error_code,), action) + finally: + sc.seterr(**entry_err) + + +@pytest.mark.thread_unsafe(reason="module refcounts are not stable in multiple threads") +@pytest.mark.skipif(not HAS_REFCOUNT, reason="Python lacks refcounts") +def test_sf_error_special_refcount(): + # Regression test for gh-16233. + # Check that the reference count of scipy.special is not increased + # when a SpecialFunctionError is raised. + refcount_before = sys.getrefcount(sc) + with sc.errstate(all='raise'): + with pytest.raises(sc.SpecialFunctionError, match='domain error'): + sc.ndtri(2.0) + refcount_after = sys.getrefcount(sc) + assert refcount_after == refcount_before + + +def test_errstate_pyx_basic(): + olderr = sc.geterr() + with sc.errstate(singular='raise'): + with assert_raises(sc.SpecialFunctionError): + sc.loggamma(0) + assert_equal(olderr, sc.geterr()) + + +def test_errstate_c_basic(): + olderr = sc.geterr() + with sc.errstate(domain='raise'): + with assert_raises(sc.SpecialFunctionError): + sc.spence(-1) + assert_equal(olderr, sc.geterr()) + + +def test_errstate_cpp_basic(): + olderr = sc.geterr() + with sc.errstate(underflow='raise'): + with assert_raises(sc.SpecialFunctionError): + sc.wrightomega(-1000) + assert_equal(olderr, sc.geterr()) + + +def test_errstate_cpp_scipy_special(): + olderr = sc.geterr() + with sc.errstate(singular='raise'): + with assert_raises(sc.SpecialFunctionError): + sc.lambertw(0, 1) + assert_equal(olderr, sc.geterr()) + + +def test_errstate_cpp_alt_ufunc_machinery(): + olderr = sc.geterr() + with sc.errstate(singular='raise'): + with assert_raises(sc.SpecialFunctionError): + sc.gammaln(0) + assert_equal(olderr, sc.geterr()) + + +def test_errstate(): + for category, error_code in _sf_error_code_map.items(): + for action in _sf_error_actions: + olderr = sc.geterr() + with sc.errstate(**{category: action}): + _check_action(_sf_error_test_function, (error_code,), action) + assert_equal(olderr, sc.geterr()) + + +def test_errstate_all_but_one(): + olderr = sc.geterr() + with sc.errstate(all='raise', singular='ignore'): + sc.gammaln(0) + with assert_raises(sc.SpecialFunctionError): + sc.spence(-1.0) + assert_equal(olderr, sc.geterr()) + + +def test_check_overflow_message(): + # Regression test for a bug where the overflow and underflow + # messages were switched. + with pytest.raises(sc.SpecialFunctionError, match="overflow"): + with sc.errstate(all='raise'): + # This should trigger an overflow: + sc.yn(3, 1e-105) diff --git a/v12x/pylib/scipy/special/tests/test_sici.py b/v12x/pylib/scipy/special/tests/test_sici.py new file mode 100644 index 0000000000000000000000000000000000000000..d33c1795641ba74f777fcfdcffe80d86463477e3 --- /dev/null +++ b/v12x/pylib/scipy/special/tests/test_sici.py @@ -0,0 +1,36 @@ +import numpy as np + +import scipy.special as sc +from scipy.special._testutils import FuncData + + +def test_sici_consistency(): + # Make sure the implementation of sici for real arguments agrees + # with the implementation of sici for complex arguments. + + # On the negative real axis Cephes drops the imaginary part in ci + def sici(x): + si, ci = sc.sici(x + 0j) + return si.real, ci.real + + x = np.r_[-np.logspace(8, -30, 200), 0, np.logspace(-30, 8, 200)] + si, ci = sc.sici(x) + dataset = np.column_stack((x, si, ci)) + FuncData(sici, dataset, 0, (1, 2), rtol=1e-12).check() + + +def test_shichi_consistency(): + # Make sure the implementation of shichi for real arguments agrees + # with the implementation of shichi for complex arguments. + + # On the negative real axis Cephes drops the imaginary part in chi + def shichi(x): + shi, chi = sc.shichi(x + 0j) + return shi.real, chi.real + + # Overflow happens quickly, so limit range + x = np.r_[-np.logspace(np.log10(700), -30, 200), 0, + np.logspace(-30, np.log10(700), 200)] + shi, chi = sc.shichi(x) + dataset = np.column_stack((x, shi, chi)) + FuncData(shichi, dataset, 0, (1, 2), rtol=1e-14).check() diff --git a/v12x/pylib/scipy/special/tests/test_specfun.py b/v12x/pylib/scipy/special/tests/test_specfun.py new file mode 100644 index 0000000000000000000000000000000000000000..f36fd2915be8c8d14138493c34dfd5733a2a1e6d --- /dev/null +++ b/v12x/pylib/scipy/special/tests/test_specfun.py @@ -0,0 +1,48 @@ +""" +Various made-up tests to hit different branches of the code in specfun.c +""" + +import numpy as np +from numpy.testing import assert_allclose +from scipy import special + + +def test_cva2_cv0_branches(): + res, resp = special.mathieu_cem([40, 129], [13, 14], [30, 45]) + assert_allclose(res, np.array([-0.3741211, 0.74441928])) + assert_allclose(resp, np.array([-37.02872758, -86.13549877])) + + res, resp = special.mathieu_sem([40, 129], [13, 14], [30, 45]) + assert_allclose(res, np.array([0.92955551, 0.66771207])) + assert_allclose(resp, np.array([-14.91073448, 96.02954185])) + + +def test_chgm_branches(): + res = special.eval_genlaguerre(-3.2, 3, 2.5) + assert_allclose(res, -0.7077721935779854) + + +def test_hygfz_branches(): + """(z == 1.0) && (c-a-b > 0.0)""" + res = special.hyp2f1(1.5, 2.5, 4.5, 1.+0.j) + assert_allclose(res, 10.30835089459151+0j) + """(cabs(z+1) < eps) && (fabs(c-a+b - 1.0) < eps)""" + res = special.hyp2f1(5+5e-16, 2, 2, -1.0 + 5e-16j) + assert_allclose(res, 0.031249999999999986+3.9062499999999994e-17j) + + +def test_pro_rad1(): + # https://github.com/scipy/scipy/issues/21058 + # Reference values taken from WolframAlpha + # SpheroidalS1(1, 1, 30, 1.1) + # SpheroidalS1Prime(1, 1, 30, 1.1) + res = special.pro_rad1(1, 1, 30, 1.1) + assert_allclose(res, (0.009657872296166435, 3.253369651472877), rtol=2e-5) + +def test_pro_rad2(): + # https://github.com/scipy/scipy/issues/21461 + # Reference values taken from WolframAlpha + # SpheroidalS2(0, 0, 3, 1.02) + # SpheroidalS2Prime(0, 0, 3, 1.02) + res = special.pro_rad2(0, 0, 3, 1.02) + assert_allclose(res, (-0.35089596858528077, 13.652764213480872), rtol=10e-10) diff --git a/v12x/pylib/scipy/special/tests/test_spence.py b/v12x/pylib/scipy/special/tests/test_spence.py new file mode 100644 index 0000000000000000000000000000000000000000..fbb26ac281dff81ea71b30318731065fe5a78f94 --- /dev/null +++ b/v12x/pylib/scipy/special/tests/test_spence.py @@ -0,0 +1,32 @@ +import numpy as np +from numpy import sqrt, log, pi +from scipy.special._testutils import FuncData +from scipy.special import spence + + +def test_consistency(): + # Make sure the implementation of spence for real arguments + # agrees with the implementation of spence for imaginary arguments. + + x = np.logspace(-30, 300, 200) + dataset = np.vstack((x + 0j, spence(x))).T + FuncData(spence, dataset, 0, 1, rtol=1e-14).check() + + +def test_special_points(): + # Check against known values of Spence's function. + + phi = (1 + sqrt(5))/2 + dataset = [(1, 0), + (2, -pi**2/12), + (0.5, pi**2/12 - log(2)**2/2), + (0, pi**2/6), + (-1, pi**2/4 - 1j*pi*log(2)), + ((-1 + sqrt(5))/2, pi**2/15 - log(phi)**2), + ((3 - sqrt(5))/2, pi**2/10 - log(phi)**2), + (phi, -pi**2/15 + log(phi)**2/2), + # Corrected from Zagier, "The Dilogarithm Function" + ((3 + sqrt(5))/2, -pi**2/10 - log(phi)**2)] + + dataset = np.asarray(dataset) + FuncData(spence, dataset, 0, 1, rtol=1e-14).check() diff --git a/v12x/pylib/scipy/special/tests/test_spfun_stats.py b/v12x/pylib/scipy/special/tests/test_spfun_stats.py new file mode 100644 index 0000000000000000000000000000000000000000..7dd4164bbab3ba88765617148dc7ea488bf81c65 --- /dev/null +++ b/v12x/pylib/scipy/special/tests/test_spfun_stats.py @@ -0,0 +1,61 @@ +import numpy as np +from numpy.testing import (assert_array_equal, assert_array_almost_equal_nulp, + assert_allclose) +from pytest import raises as assert_raises + +from scipy.special import gammaln, multigammaln # type:ignore[attr-defined] + + +class TestMultiGammaLn: + + def test1(self): + # A test of the identity + # Gamma_1(a) = Gamma(a) + np.random.seed(1234) + a = np.abs(np.random.randn()) + assert_array_equal(multigammaln(a, 1), gammaln(a)) + + def test2(self): + # A test of the identity + # Gamma_2(a) = sqrt(pi) * Gamma(a) * Gamma(a - 0.5) + a = np.array([2.5, 10.0]) + result = multigammaln(a, 2) + expected = np.log(np.sqrt(np.pi)) + gammaln(a) + gammaln(a - 0.5) + assert_allclose(result, expected, atol=1.5e-7, rtol=0) + + def test_bararg(self): + assert_raises(ValueError, multigammaln, 0.5, 1.2) + + +def _check_multigammaln_array_result(a, d): + # Test that the shape of the array returned by multigammaln + # matches the input shape, and that all the values match + # the value computed when multigammaln is called with a scalar. + result = multigammaln(a, d) + assert_array_equal(a.shape, result.shape) + a1 = a.ravel() + result1 = result.ravel() + for i in range(a.size): + assert_array_almost_equal_nulp(result1[i], multigammaln(a1[i], d)) + + +def test_multigammaln_array_arg(): + # Check that the array returned by multigammaln has the correct + # shape and contains the correct values. The cases have arrays + # with several different shapes. + # The cases include a regression test for ticket #1849 + # (a = np.array([2.0]), an array with a single element). + np.random.seed(1234) + + cases = [ + # a, d + (np.abs(np.random.randn(3, 2)) + 5, 5), + (np.abs(np.random.randn(1, 2)) + 5, 5), + (np.arange(10.0, 18.0).reshape(2, 2, 2), 3), + (np.array([2.0]), 3), + (np.float64(2.0), 3), + ] + + for a, d in cases: + _check_multigammaln_array_result(a, d) + diff --git a/v12x/pylib/scipy/special/tests/test_sph_harm.py b/v12x/pylib/scipy/special/tests/test_sph_harm.py new file mode 100644 index 0000000000000000000000000000000000000000..117f0938379082125f97d738d27364f7821eaefc --- /dev/null +++ b/v12x/pylib/scipy/special/tests/test_sph_harm.py @@ -0,0 +1,48 @@ +import numpy as np +import pytest + +from numpy.testing import assert_allclose +import scipy.special as sc + +class TestSphHarm: + @pytest.mark.slow + def test_p(self): + m_max = 20 + n_max = 10 + + theta = np.linspace(0, np.pi) + phi = np.linspace(0, 2*np.pi) + theta, phi = np.meshgrid(theta, phi) + + y, y_jac, y_hess = sc.sph_harm_y_all(n_max, m_max, theta, phi, diff_n=2) + p, p_jac, p_hess = sc.sph_legendre_p_all(n_max, m_max, theta, diff_n=2) + + m = np.concatenate([np.arange(m_max + 1), np.arange(-m_max, 0)]) + m = np.expand_dims(m, axis=(0,)+tuple(range(2,theta.ndim+2))) + + assert_allclose(y, p * np.exp(1j * m * phi)) + + assert_allclose(y_jac[..., 0], p_jac * np.exp(1j * m * phi)) + assert_allclose(y_jac[..., 1], 1j * m * p * np.exp(1j * m * phi)) + + assert_allclose(y_hess[..., 0, 0], p_hess * np.exp(1j * m * phi)) + assert_allclose(y_hess[..., 0, 1], 1j * m * p_jac * np.exp(1j * m * phi)) + assert_allclose(y_hess[..., 1, 0], y_hess[..., 0, 1]) + assert_allclose(y_hess[..., 1, 1], -m * m * p * np.exp(1j * m * phi)) + + @pytest.mark.parametrize("n_max", [7, 10, 50]) + @pytest.mark.parametrize("m_max", [1, 4, 5, 9, 14]) + def test_all(self, n_max, m_max): + theta = np.linspace(0, np.pi) + phi = np.linspace(0, 2 * np.pi) + + n = np.arange(n_max + 1) + n = np.expand_dims(n, axis=tuple(range(1,theta.ndim+2))) + + m = np.concatenate([np.arange(m_max + 1), np.arange(-m_max, 0)]) + m = np.expand_dims(m, axis=(0,)+tuple(range(2,theta.ndim+2))) + + y_actual = sc.sph_harm_y_all(n_max, m_max, theta, phi) + y_desired = sc.sph_harm_y(n, m, theta, phi) + + np.testing.assert_allclose(y_actual, y_desired, rtol=1e-05) diff --git a/v12x/pylib/scipy/special/tests/test_spherical_bessel.py b/v12x/pylib/scipy/special/tests/test_spherical_bessel.py new file mode 100644 index 0000000000000000000000000000000000000000..8444a4c16aa9e3b9ffc204f824a3f136a030f2be --- /dev/null +++ b/v12x/pylib/scipy/special/tests/test_spherical_bessel.py @@ -0,0 +1,407 @@ +# +# Tests of spherical Bessel functions. +# +import warnings + +import numpy as np +from numpy.testing import assert_allclose +import pytest +from numpy import sin, cos, sinh, cosh, exp, inf, nan, r_, pi + +from scipy.special import spherical_jn, spherical_yn, spherical_in, spherical_kn +from scipy.integrate import quad + + +class TestSphericalJn: + def test_spherical_jn_exact(self): + # https://dlmf.nist.gov/10.49.E3 + # Note: exact expression is numerically stable only for small + # n or z >> n. + x = np.array([0.12, 1.23, 12.34, 123.45, 1234.5]) + assert_allclose(spherical_jn(2, x), + (-1/x + 3/x**3)*sin(x) - 3/x**2*cos(x)) + + def test_spherical_jn_recurrence_complex(self): + # https://dlmf.nist.gov/10.51.E1 + n = np.array([1, 2, 3, 7, 12]) + x = 1.1 + 1.5j + assert_allclose(spherical_jn(n - 1, x) + spherical_jn(n + 1, x), + (2*n + 1)/x*spherical_jn(n, x)) + + def test_spherical_jn_recurrence_real(self): + # https://dlmf.nist.gov/10.51.E1 + n = np.array([1, 2, 3, 7, 12]) + x = 0.12 + assert_allclose(spherical_jn(n - 1, x) + spherical_jn(n + 1,x), + (2*n + 1)/x*spherical_jn(n, x)) + + def test_spherical_jn_inf_real(self): + # https://dlmf.nist.gov/10.52.E3 + n = 6 + x = np.array([-inf, inf]) + assert_allclose(spherical_jn(n, x), np.array([0, 0])) + + def test_spherical_jn_inf_complex(self): + # https://dlmf.nist.gov/10.52.E3 + n = 7 + x = np.array([-inf + 0j, inf + 0j, inf*(1+1j)]) + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", "invalid value encountered in multiply", RuntimeWarning) + assert_allclose(spherical_jn(n, x), np.array([0, 0, inf*(1+1j)])) + + def test_spherical_jn_large_arg_1(self): + # https://github.com/scipy/scipy/issues/2165 + # Reference value computed using mpmath, via + # besselj(n + mpf(1)/2, z)*sqrt(pi/(2*z)) + assert_allclose(spherical_jn(2, 3350.507), -0.00029846226538040747) + + def test_spherical_jn_large_arg_2(self): + # https://github.com/scipy/scipy/issues/1641 + # Reference value computed using mpmath, via + # besselj(n + mpf(1)/2, z)*sqrt(pi/(2*z)) + assert_allclose(spherical_jn(2, 10000), 3.0590002633029811e-05) + + def test_spherical_jn_at_zero(self): + # https://dlmf.nist.gov/10.52.E1 + # But note that n = 0 is a special case: j0 = sin(x)/x -> 1 + n = np.array([0, 1, 2, 5, 10, 100]) + x = 0 + assert_allclose(spherical_jn(n, x), np.array([1, 0, 0, 0, 0, 0])) + + +class TestSphericalYn: + def test_spherical_yn_exact(self): + # https://dlmf.nist.gov/10.49.E5 + # Note: exact expression is numerically stable only for small + # n or z >> n. + x = np.array([0.12, 1.23, 12.34, 123.45, 1234.5]) + assert_allclose(spherical_yn(2, x), + (1/x - 3/x**3)*cos(x) - 3/x**2*sin(x)) + + def test_spherical_yn_recurrence_real(self): + # https://dlmf.nist.gov/10.51.E1 + n = np.array([1, 2, 3, 7, 12]) + x = 0.12 + assert_allclose(spherical_yn(n - 1, x) + spherical_yn(n + 1,x), + (2*n + 1)/x*spherical_yn(n, x)) + + def test_spherical_yn_recurrence_complex(self): + # https://dlmf.nist.gov/10.51.E1 + n = np.array([1, 2, 3, 7, 12]) + x = 1.1 + 1.5j + assert_allclose(spherical_yn(n - 1, x) + spherical_yn(n + 1, x), + (2*n + 1)/x*spherical_yn(n, x)) + + def test_spherical_yn_inf_real(self): + # https://dlmf.nist.gov/10.52.E3 + n = 6 + x = np.array([-inf, inf]) + assert_allclose(spherical_yn(n, x), np.array([0, 0])) + + def test_spherical_yn_inf_complex(self): + # https://dlmf.nist.gov/10.52.E3 + n = 7 + x = np.array([-inf + 0j, inf + 0j, inf*(1+1j)]) + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", "invalid value encountered in multiply", RuntimeWarning) + assert_allclose(spherical_yn(n, x), np.array([0, 0, inf*(1+1j)])) + + def test_spherical_yn_at_zero(self): + # https://dlmf.nist.gov/10.52.E2 + n = np.array([0, 1, 2, 5, 10, 100]) + x = 0 + assert_allclose(spherical_yn(n, x), np.full(n.shape, -inf)) + + def test_spherical_yn_at_zero_complex(self): + # Consistently with numpy: + # >>> -np.cos(0)/0 + # -inf + # >>> -np.cos(0+0j)/(0+0j) + # (-inf + nan*j) + n = np.array([0, 1, 2, 5, 10, 100]) + x = 0 + 0j + assert_allclose(spherical_yn(n, x), np.full(n.shape, nan)) + + +class TestSphericalJnYnCrossProduct: + def test_spherical_jn_yn_cross_product_1(self): + # https://dlmf.nist.gov/10.50.E3 + n = np.array([1, 5, 8]) + x = np.array([0.1, 1, 10]) + left = (spherical_jn(n + 1, x) * spherical_yn(n, x) - + spherical_jn(n, x) * spherical_yn(n + 1, x)) + right = 1/x**2 + assert_allclose(left, right) + + def test_spherical_jn_yn_cross_product_2(self): + # https://dlmf.nist.gov/10.50.E3 + n = np.array([1, 5, 8]) + x = np.array([0.1, 1, 10]) + left = (spherical_jn(n + 2, x) * spherical_yn(n, x) - + spherical_jn(n, x) * spherical_yn(n + 2, x)) + right = (2*n + 3)/x**3 + assert_allclose(left, right) + + +class TestSphericalIn: + def test_spherical_in_exact(self): + # https://dlmf.nist.gov/10.49.E9 + x = np.array([0.12, 1.23, 12.34, 123.45]) + assert_allclose(spherical_in(2, x), + (1/x + 3/x**3)*sinh(x) - 3/x**2*cosh(x)) + + def test_spherical_in_recurrence_real(self): + # https://dlmf.nist.gov/10.51.E4 + n = np.array([1, 2, 3, 7, 12]) + x = 0.12 + assert_allclose(spherical_in(n - 1, x) - spherical_in(n + 1,x), + (2*n + 1)/x*spherical_in(n, x)) + + def test_spherical_in_recurrence_complex(self): + # https://dlmf.nist.gov/10.51.E1 + n = np.array([1, 2, 3, 7, 12]) + x = 1.1 + 1.5j + assert_allclose(spherical_in(n - 1, x) - spherical_in(n + 1,x), + (2*n + 1)/x*spherical_in(n, x)) + + def test_spherical_in_inf_real(self): + # https://dlmf.nist.gov/10.52.E3 + n = 5 + x = np.array([-inf, inf]) + assert_allclose(spherical_in(n, x), np.array([-inf, inf])) + + def test_spherical_in_inf_complex(self): + # https://dlmf.nist.gov/10.52.E5 + # Ideally, i1n(n, 1j*inf) = 0 and i1n(n, (1+1j)*inf) = (1+1j)*inf, but + # this appears impossible to achieve because C99 regards any complex + # value with at least one infinite part as a complex infinity, so + # 1j*inf cannot be distinguished from (1+1j)*inf. Therefore, nan is + # the correct return value. + n = 7 + x = np.array([-inf + 0j, inf + 0j, inf*(1+1j)]) + assert_allclose(spherical_in(n, x), np.array([-inf, inf, nan])) + + def test_spherical_in_at_zero(self): + # https://dlmf.nist.gov/10.52.E1 + # But note that n = 0 is a special case: i0 = sinh(x)/x -> 1 + n = np.array([0, 1, 2, 5, 10, 100]) + x = 0 + assert_allclose(spherical_in(n, x), np.array([1, 0, 0, 0, 0, 0])) + + +class TestSphericalKn: + def test_spherical_kn_exact(self): + # https://dlmf.nist.gov/10.49.E13 + x = np.array([0.12, 1.23, 12.34, 123.45]) + assert_allclose(spherical_kn(2, x), + pi/2*exp(-x)*(1/x + 3/x**2 + 3/x**3)) + + def test_spherical_kn_recurrence_real(self): + # https://dlmf.nist.gov/10.51.E4 + n = np.array([1, 2, 3, 7, 12]) + x = 0.12 + assert_allclose( + (-1)**(n - 1)*spherical_kn(n - 1, x) - (-1)**(n + 1)*spherical_kn(n + 1,x), + (-1)**n*(2*n + 1)/x*spherical_kn(n, x) + ) + + def test_spherical_kn_recurrence_complex(self): + # https://dlmf.nist.gov/10.51.E4 + n = np.array([1, 2, 3, 7, 12]) + x = 1.1 + 1.5j + assert_allclose( + (-1)**(n - 1)*spherical_kn(n - 1, x) - (-1)**(n + 1)*spherical_kn(n + 1,x), + (-1)**n*(2*n + 1)/x*spherical_kn(n, x) + ) + + def test_spherical_kn_inf_real(self): + # https://dlmf.nist.gov/10.52.E6 + n = 5 + x = np.array([-inf, inf]) + assert_allclose(spherical_kn(n, x), np.array([-inf, 0])) + + def test_spherical_kn_inf_complex(self): + # https://dlmf.nist.gov/10.52.E6 + # The behavior at complex infinity depends on the sign of the real + # part: if Re(z) >= 0, then the limit is 0; if Re(z) < 0, then it's + # z*inf. This distinction cannot be captured, so we return nan. + n = 7 + x = np.array([-inf + 0j, inf + 0j, inf*(1+1j)]) + assert_allclose(spherical_kn(n, x), np.array([-inf, 0, nan])) + + def test_spherical_kn_at_zero(self): + # https://dlmf.nist.gov/10.52.E2 + n = np.array([0, 1, 2, 5, 10, 100]) + x = 0 + assert_allclose(spherical_kn(n, x), np.full(n.shape, inf)) + + def test_spherical_kn_at_zero_complex(self): + # https://dlmf.nist.gov/10.52.E2 + n = np.array([0, 1, 2, 5, 10, 100]) + x = 0 + 0j + assert_allclose(spherical_kn(n, x), np.full(n.shape, nan)) + + +class SphericalDerivativesTestCase: + def fundamental_theorem(self, n, a, b): + integral, tolerance = quad(lambda z: self.df(n, z), a, b) + assert_allclose(integral, + self.f(n, b) - self.f(n, a), + atol=tolerance) + + @pytest.mark.slow + def test_fundamental_theorem_0(self): + self.fundamental_theorem(0, 3.0, 15.0) + + @pytest.mark.slow + def test_fundamental_theorem_7(self): + self.fundamental_theorem(7, 0.5, 1.2) + + +class TestSphericalJnDerivatives(SphericalDerivativesTestCase): + def f(self, n, z): + return spherical_jn(n, z) + + def df(self, n, z): + return spherical_jn(n, z, derivative=True) + + def test_spherical_jn_d_zero(self): + n = np.array([0, 1, 2, 3, 7, 15]) + assert_allclose(spherical_jn(n, 0, derivative=True), + np.array([0, 1/3, 0, 0, 0, 0])) + + +class TestSphericalYnDerivatives(SphericalDerivativesTestCase): + def f(self, n, z): + return spherical_yn(n, z) + + def df(self, n, z): + return spherical_yn(n, z, derivative=True) + + +class TestSphericalInDerivatives(SphericalDerivativesTestCase): + def f(self, n, z): + return spherical_in(n, z) + + def df(self, n, z): + return spherical_in(n, z, derivative=True) + + def test_spherical_in_d_zero(self): + n = np.array([0, 1, 2, 3, 7, 15]) + spherical_in(n, 0, derivative=False) + assert_allclose(spherical_in(n, 0, derivative=True), + np.array([0, 1/3, 0, 0, 0, 0])) + + +class TestSphericalKnDerivatives(SphericalDerivativesTestCase): + def f(self, n, z): + return spherical_kn(n, z) + + def df(self, n, z): + return spherical_kn(n, z, derivative=True) + + +class TestSphericalOld: + # These are tests from the TestSpherical class of test_basic.py, + # rewritten to use spherical_* instead of sph_* but otherwise unchanged. + + def test_sph_in(self): + # This test reproduces test_basic.TestSpherical.test_sph_in. + i1n = np.empty((2,2)) + x = 0.2 + + i1n[0][0] = spherical_in(0, x) + i1n[0][1] = spherical_in(1, x) + i1n[1][0] = spherical_in(0, x, derivative=True) + i1n[1][1] = spherical_in(1, x, derivative=True) + + inp0 = (i1n[0][1]) + inp1 = (i1n[0][0] - 2.0/0.2 * i1n[0][1]) + assert_allclose(i1n[0], np.array([1.0066800127054699381, + 0.066933714568029540839]), + atol=1.5e-12, rtol=0.0) + assert_allclose(i1n[1], [inp0, inp1], atol=1.5e-12, rtol=0) + + def test_sph_in_kn_order0(self): + x = 1. + sph_i0 = np.empty((2,)) + sph_i0[0] = spherical_in(0, x) + sph_i0[1] = spherical_in(0, x, derivative=True) + sph_i0_expected = np.array([np.sinh(x)/x, + np.cosh(x)/x-np.sinh(x)/x**2]) + assert_allclose(r_[sph_i0], sph_i0_expected, atol=1.5e-7, rtol=0) + + sph_k0 = np.empty((2,)) + sph_k0[0] = spherical_kn(0, x) + sph_k0[1] = spherical_kn(0, x, derivative=True) + sph_k0_expected = np.array([0.5*pi*exp(-x)/x, + -0.5*pi*exp(-x)*(1/x+1/x**2)]) + assert_allclose(r_[sph_k0], sph_k0_expected, atol=1.5e-7, rtol=0) + + def test_sph_jn(self): + s1 = np.empty((2,3)) + x = 0.2 + + s1[0][0] = spherical_jn(0, x) + s1[0][1] = spherical_jn(1, x) + s1[0][2] = spherical_jn(2, x) + s1[1][0] = spherical_jn(0, x, derivative=True) + s1[1][1] = spherical_jn(1, x, derivative=True) + s1[1][2] = spherical_jn(2, x, derivative=True) + + s10 = -s1[0][1] + s11 = s1[0][0]-2.0/0.2*s1[0][1] + s12 = s1[0][1]-3.0/0.2*s1[0][2] + assert_allclose(s1[0], [0.99334665397530607731, + 0.066400380670322230863, + 0.0026590560795273856680], + atol=1.5e-12, rtol=0) + assert_allclose(s1[1], [s10, s11, s12], atol=1.5e-12, rtol=0) + + def test_sph_kn(self): + kn = np.empty((2,3)) + x = 0.2 + + kn[0][0] = spherical_kn(0, x) + kn[0][1] = spherical_kn(1, x) + kn[0][2] = spherical_kn(2, x) + kn[1][0] = spherical_kn(0, x, derivative=True) + kn[1][1] = spherical_kn(1, x, derivative=True) + kn[1][2] = spherical_kn(2, x, derivative=True) + + kn0 = -kn[0][1] + kn1 = -kn[0][0]-2.0/0.2*kn[0][1] + kn2 = -kn[0][1]-3.0/0.2*kn[0][2] + assert_allclose(kn[0], [6.4302962978445670140, + 38.581777787067402086, + 585.15696310385559829], + atol=1.5e-12, rtol=0) + assert_allclose(kn[1], [kn0, kn1, kn2], atol=1.5e-9, rtol=0) + + def test_sph_yn(self): + sy1 = spherical_yn(2, 0.2) + sy2 = spherical_yn(0, 0.2) + # previous values in the system + assert_allclose(sy1, -377.52483, atol=1.5e-5, rtol=0) + assert_allclose(sy2, -4.9003329, atol=1.5e-5, rtol=0) + sphpy = (spherical_yn(0, 0.2) - 2*spherical_yn(2, 0.2))/3 + sy3 = spherical_yn(1, 0.2, derivative=True) + # compare correct derivative val. (correct =-system val). + assert_allclose(sy3, sphpy, atol=1.5e-4, rtol=0) + + +@pytest.mark.parametrize('derivative', [False, True]) +@pytest.mark.parametrize('fun', [spherical_jn, spherical_in, + spherical_yn, spherical_kn]) +def test_negative_real_gh14582(derivative, fun): + # gh-14582 reported that the spherical Bessel functions did not work + # with negative real argument `z`. Check that this is resolved. + rng = np.random.default_rng(3598435982345987234) + size = 25 + n = rng.integers(0, 10, size=size) + z = rng.standard_normal(size=size) + res = fun(n, z, derivative=derivative) + ref = fun(n, z+0j, derivative=derivative) + assert_allclose(res, ref.real) diff --git a/v12x/pylib/scipy/special/tests/test_support_alternative_backends.py b/v12x/pylib/scipy/special/tests/test_support_alternative_backends.py new file mode 100644 index 0000000000000000000000000000000000000000..34d239e4e8f8a619b59ef79f5ce2e4842e5c41a0 --- /dev/null +++ b/v12x/pylib/scipy/special/tests/test_support_alternative_backends.py @@ -0,0 +1,404 @@ +from functools import partial +import pickle + +import pytest +from hypothesis import given, strategies +import hypothesis.extra.numpy as npst +from packaging import version + +from scipy import special +from scipy.special._support_alternative_backends import _special_funcs +from scipy._lib._array_api_no_0d import xp_assert_close +from scipy._lib._array_api import (is_cupy, is_dask, is_jax, is_torch, + make_xp_pytest_param, make_xp_test_case, + get_native_namespace_name) +from scipy._external.array_api_compat import numpy as np +import scipy._external.array_api_extra as xpx + + +# Run all tests in this module in the Array API CI, including those without +# the xp fixture +pytestmark = pytest.mark.array_api_backends + +lazy_xp_modules = [special] + + +def _skip_or_tweak_alternative_backends(xp, nfo, dtypes, int_only): + """Skip tests for specific intersections of scipy.special functions + vs. backends vs. dtypes vs. devices. + Also suggest bespoke tweaks. + + Returns + ------- + positive_only : list[bool] + Whether you should exclusively test positive inputs. + dtypes : list[str] + dtype strings 'float64', 'int32', 'int64', etc. with integer types + mapped to the type of the NumPy default int. + """ + f_name = nfo.name + if isinstance(nfo.positive_only, dict): + positive_only = nfo.positive_only.get(get_native_namespace_name(xp), False) + else: + positive_only = nfo.positive_only + if isinstance(positive_only, bool): + positive_only = [positive_only]*nfo.n_args + + dtypes = [np.intp.__name__ if dtype == "intp" else dtype for dtype in dtypes] + + if f_name in {'betaincinv'} and is_cupy(xp): + pytest.xfail("CuPy uses different convention for out of domain input.") + + if not any('int' in dtype for dtype in dtypes): + return positive_only, dtypes + + # Integer-specific issues from this point onwards + + if f_name in {'gamma', 'gammasgn'} and is_cupy(xp): + # CuPy has not yet updated gamma pole behavior to match + # https://github.com/scipy/scipy/pull/21827. + positive_only = [True] + + if f_name in {'poch'} and is_jax(xp): + # Jax uses a different convention at gamma poles. + positive_only = [True, True] + + if f_name == 'multigammaln': + pytest.skip("multigammaln raises for out of domain inputs.") + + if ((is_torch(xp) and f_name in {'gammainc', 'gammaincc'}) + or (is_cupy(xp) and f_name in {'stdtr', 'i0e', 'i1e'}) + or (is_jax(xp) and f_name in {'stdtr', 'ndtr', 'ndtri', 'log_ndtr', 'hyp1f1', + 'hyp2f1', 'spence', 'kl_div'}) + ): + pytest.skip(f"`{f_name}` does not support integer types") + + # int/float mismatched args support is sketchy + if (any('float' in dtype for dtype in dtypes) + and ((is_torch(xp) and f_name in ('rel_entr', 'xlogy', 'polygamma', + 'zeta', 'xlog1py')) + or (is_jax(xp) and f_name in ('gammainc', 'gammaincc', 'expn', + 'rel_entr', 'xlogy', 'betaln', + 'polygamma', 'zeta', 'poch', + 'xlog1py'))) + ): + pytest.xfail("dtypes do not match") + + if (is_torch(xp) and xpx.default_dtype(xp) == xp.float32): + # On PyTorch with float32 default dtype, all ints are promoted + # to float32, but when falling back to NumPy/SciPy int64 is promoted + # instead to float64. Integer only parameters essentially do not + # participate in determination of the result type in PyTorch with + # float32 default dtype, but will impact the output dtype as if + # they were float64 when falling back to NumPy/SciPy. + if not nfo.torch_native: + pytest.xfail("dtypes do not match") + + return positive_only, dtypes + + +@pytest.mark.filterwarnings("ignore:invalid value encountered:RuntimeWarning:dask") +@pytest.mark.filterwarnings("ignore:overflow encountered:RuntimeWarning") +@pytest.mark.parametrize('shapes', [[(0,)]*4, [tuple()]*4, [(10,)]*4, + [(10,), (11, 1), (12, 1, 1), (13, 1, 1, 1)]]) +@pytest.mark.parametrize('base_dtype', ['float32', 'float64', 'intp']) +@pytest.mark.parametrize( + 'func,nfo', [make_xp_pytest_param(i.wrapper, i) for i in _special_funcs]) +def test_support_alternative_backends(xp, func, nfo, base_dtype, shapes): + int_only = nfo.int_only + if int_only is None: + int_only = (False, ) * nfo.n_args + dtypes = (base_dtype, ) * nfo.n_args + else: + dtypes = tuple( + 'intp' if needs_int else base_dtype for needs_int in int_only + ) + + positive_only, dtypes = _skip_or_tweak_alternative_backends( + xp, nfo, dtypes, int_only + ) + + dtypes_np = [getattr(np, dtype) for dtype in dtypes] + dtypes_xp = [getattr(xp, dtype) for dtype in dtypes] + + shapes = shapes[:nfo.n_args] + rng = np.random.default_rng(984254252920492019) + args_np = [] + + # Handle cases where there's an argument which only takes scalar values. + python_int_only = nfo.python_int_only + if isinstance(python_int_only, dict): + python_int_only = python_int_only.get(get_native_namespace_name(xp)) + scalar_or_0d_only = nfo.scalar_or_0d_only + if isinstance(scalar_or_0d_only, dict): + scalar_or_0d_only = scalar_or_0d_only.get(get_native_namespace_name(xp)) + + test_large_ints = nfo.test_large_ints + if isinstance(nfo.test_large_ints, dict): + test_large_ints = test_large_ints.get(get_native_namespace_name(xp), False) + + if python_int_only is None: + python_int_only = [False] * nfo.n_args + if scalar_or_0d_only is None: + scalar_or_0d_only = [False] * nfo.n_args + + no_shape = [ + cond1 or cond2 for cond1, cond2 in zip(python_int_only, scalar_or_0d_only) + ] + + shapes = [shape if not cond else None for shape, cond in zip(shapes, no_shape)] + + for dtype, dtype_np, shape, needs_python_int in zip( + dtypes, dtypes_np, shapes, python_int_only + ): + if 'int' in dtype and test_large_ints: + iinfo = np.iinfo(dtype_np) + rand = partial(rng.integers, iinfo.min, iinfo.max + 1) + elif 'int' in dtype: + rand = partial(rng.integers, -20, 21) + else: + rand = rng.standard_normal + val = rand(size=shape, dtype=dtype_np) + if needs_python_int: + # The logic above for determining shapes guarantees that + # shape will be None in the above line when a Python int is required, + # so this can safely be converted to an int. + val = int(val) + args_np.append(val) + + args_np = [ + np.abs(arg) if cond else arg for arg, cond in zip(args_np, positive_only) + ] + + args_xp = [ + xp.asarray(arg, dtype=dtype_xp) if not needs_python_int + else arg + for arg, dtype_xp, needs_python_int + in zip(args_np, dtypes_xp, python_int_only) + ] + + args_np = [ + np.asarray(arg, dtype=dtype_np) if not needs_python_int + else arg + for arg, dtype_np, needs_python_int + in zip(args_np, dtypes_np, python_int_only) + ] + + if is_dask(xp): + # We're using map_blocks to dispatch the function to Dask. + # This is the correct thing to do IF all tested functions are elementwise; + # otherwise the output would change depending on chunking. + # Try to trigger bugs related to having multiple chunks. + args_xp = [arg.rechunk(5) for arg in args_xp] + + res = nfo.wrapper(*args_xp) # Also wrapped by lazy_xp_function + ref = nfo.func(*args_np) # Unwrapped ufunc + if ( + is_torch(xp) + and xpx.default_dtype(xp) == xp.float32 + and "float64" not in dtypes + ): + # int64 promotes like float32 on torch with default dtype = float32 + # cast reference if needed + ref = np.float32(ref) + # When dtype_np is integer, the output dtype can be float + atol = 0 if ref.dtype.kind in 'iu' else 10 * np.finfo(ref.dtype).eps + rtol = None + if is_torch(xp) and func.__name__ == 'j1': + # If we end up needing more function/backend specific tolerance + # adjustments, this should be factored out properly. + atol = 1e-7 + rtol = 1e-5 + xp_assert_close( + res, xp.asarray(ref), rtol=rtol, atol=atol, check_0d=nfo.produces_0d + ) + + +@pytest.mark.parametrize( + 'func, nfo', + [make_xp_pytest_param(i.wrapper, i) for i in _special_funcs if i.n_args >= 2]) +@pytest.mark.filterwarnings("ignore:invalid value encountered:RuntimeWarning:dask") +def test_support_alternative_backends_mismatched_dtypes(xp, func, nfo): + """Test mix-n-match of int and float arguments""" + if func.__name__ in {'expn', 'polygamma', 'multigammaln', 'bdtr', 'bdtrc', 'bdtri', + 'nbdtr', 'nbdtrc', 'nbdtri', 'pdtri'}: + pytest.skip(f"dtypes for {func.__name__} make it a bad fit for this test.") + dtypes = ['intp', 'float32', 'float64', 'float64'][:nfo.n_args] + + positive_only, dtypes = _skip_or_tweak_alternative_backends( + xp, nfo, dtypes, (False,)*nfo.n_args + ) + dtypes_np = [getattr(np, dtype) for dtype in dtypes] + dtypes_xp = [getattr(xp, dtype) for dtype in dtypes] + + rng = np.random.default_rng(984254252920492019) + iinfo = np.iinfo(np.intp) + if nfo.test_large_ints: + randint = partial(rng.integers, iinfo.min, iinfo.max + 1) + else: + randint = partial(rng.integers, -20, 21) + args_np = [ + randint(size=1, dtype=np.int64), + rng.standard_normal(size=1, dtype=np.float32), + rng.standard_normal(size=1, dtype=np.float64), + rng.standard_normal(size=1, dtype=np.float64), + ][:nfo.n_args] + args_np = [ + np.abs(arg) if cond else arg for arg, cond in zip(args_np, positive_only) + ] + + args_xp = [xp.asarray(arg, dtype=dtype_xp) + for arg, dtype_xp in zip(args_np, dtypes_xp)] + args_np = [np.asarray(arg, dtype=dtype_np) + for arg, dtype_np in zip(args_np, dtypes_np)] + + res = nfo.wrapper(*args_xp) # Also wrapped by lazy_xp_function + ref = nfo.func(*args_np) # Unwrapped ufunc + if ( + is_torch(xp) + and xpx.default_dtype(xp) == xp.float32 + and "float64" not in dtypes + ): + # int64 promotes like float32 on torch with default dtype = float32 + # cast reference if needed + ref = np.float32(ref) + + atol = 10 * np.finfo(ref.dtype).eps + xp_assert_close(res, xp.asarray(ref), atol=atol) + + +@pytest.mark.xslow +@given(data=strategies.data()) +@pytest.mark.fail_slow(5) +@pytest.mark.parametrize( + 'func,nfo', [make_xp_pytest_param(nfo.wrapper, nfo) for nfo in _special_funcs]) +@pytest.mark.filterwarnings("ignore:invalid value encountered:RuntimeWarning:dask") +@pytest.mark.filterwarnings("ignore:divide by zero encountered:RuntimeWarning:dask") +@pytest.mark.filterwarnings("ignore:overflow encountered:RuntimeWarning:dask") +@pytest.mark.filterwarnings( + "ignore:overflow encountered:RuntimeWarning:array_api_strict" +) +def test_support_alternative_backends_hypothesis(xp, func, nfo, data): + if func.__name__ in {'expn', 'polygamma', 'multigammaln', 'bdtr', 'bdtrc', 'bdtri', + 'nbdtr', 'nbdtrc', 'nbdtri', 'pdtri'}: + pytest.skip(f"dtypes for {func.__name__} make it a bad fit for this test.") + dtype = data.draw(strategies.sampled_from(['float32', 'float64', 'intp'])) + positive_only, dtypes = _skip_or_tweak_alternative_backends( + xp, nfo, [dtype], (False,)*nfo.n_args + ) + dtype_np = getattr(np, dtypes[0]) + dtype_xp = getattr(xp, dtypes[0]) + + elements = {'allow_subnormal': False} + # Most failures are due to NaN or infinity; uncomment to suppress them + # elements['allow_infinity'] = False + # elements['allow_nan'] = False + if any(positive_only): + elements['min_value'] = 0 + + shapes, _ = data.draw( + npst.mutually_broadcastable_shapes(num_shapes=nfo.n_args)) + args_np = [data.draw(npst.arrays(dtype_np, shape, elements=elements)) + for shape in shapes] + + args_xp = [xp.asarray(arg, dtype=dtype_xp) for arg in args_np] + args_np = [np.asarray(arg, dtype=dtype_np) for arg in args_np] + + res = nfo.wrapper(*args_xp) # Also wrapped by lazy_xp_function + ref = nfo.func(*args_np) # Unwrapped ufunc + if ( + is_torch(xp) + and xpx.default_dtype(xp) == xp.float32 + and dtype != "float64" + ): + # int64 promotes like float32 on torch with default dtype = float32 + # cast reference if needed + ref = np.float32(ref) + + # When dtype_np is integer, the output dtype can be float + atol = 0 if ref.dtype.kind in 'iu' else 10 * np.finfo(ref.dtype).eps + xp_assert_close(res, xp.asarray(ref), atol=atol) + + +@pytest.mark.filterwarnings("ignore:numpy.core is deprecated:DeprecationWarning") +@pytest.mark.parametrize("func", [nfo.wrapper for nfo in _special_funcs]) +def test_pickle(func): + roundtrip = pickle.loads(pickle.dumps(func)) + assert roundtrip is func + + +@pytest.mark.parametrize("func", [nfo.wrapper for nfo in _special_funcs]) +def test_repr(func): + assert func.__name__ in repr(func) + assert "locals" not in repr(func) + + +@pytest.mark.skipif( + version.parse(np.__version__) < version.parse("2.2"), # type:ignore[attr-defined] + reason="Can't update ufunc __doc__ when SciPy is compiled vs. NumPy < 2.2") +@pytest.mark.parametrize('func', [nfo.wrapper for nfo in _special_funcs]) +def test_doc(func): + """xp_capabilities updates the docstring in place. + Make sure it does so exactly once, including when SCIPY_ARRAY_API is not set. + """ + match = "has experimental support for Python Array API" + assert func.__doc__.count(match) == 1 + + +@pytest.mark.parametrize( + 'func,n_args,int_only,is_ufunc', + [(nfo.wrapper, nfo.n_args, nfo.int_only, nfo.is_ufunc) + for nfo in _special_funcs] +) +def test_ufunc_kwargs(func, n_args, int_only, is_ufunc): + """Test that numpy-specific out= and dtype= keyword arguments + of ufuncs still work when SCIPY_ARRAY_API is set. + """ + if not is_ufunc: + pytest.skip(f"{func.__name__} is not a ufunc.") + if int_only is None: + int_only = (False, ) * n_args + # out= + args = [ + np.asarray([.1, .2]) if not needs_int + else np.asarray([1, 2]) + for needs_int in int_only + ] + out = np.empty(2) + y = func(*args, out=out) + xp_assert_close(y, out) + + # out= with out.dtype != args.dtype + out = np.empty(2, dtype=np.float32) + y = func(*args, out=out) + xp_assert_close(y, out) + + if func.__name__ in {"bdtr", "bdtrc", "bdtri"}: + # The below function evaluation will trigger a deprecation warning + # with dtype=np.float32. This will go away if the trigger is actually + # pulled on the deprecation. + return + + # dtype= + y = func(*args, dtype=np.float32) + assert y.dtype == np.float32 + + +@make_xp_test_case(special.chdtr) +def test_chdtr_gh21311(xp): + # the edge case behavior of generic chdtr was not right; see gh-21311 + # be sure to test at least these cases + # should add `np.nan` into the mix when gh-21317 is resolved + x = np.asarray([-np.inf, -1., 0., 1., np.inf]) + v = x.reshape(-1, 1) + ref = special.chdtr(v, x) + res = special.chdtr(xp.asarray(v), xp.asarray(x)) + xp_assert_close(res, xp.asarray(ref)) + + +@make_xp_test_case(special.fdtrc) +def test_mixed_arrays_and_python_scalars(xp): + # Tests that the delegation infrastructure respects NEP50. + res = special.fdtrc(1.1, 2., xp.asarray(1., dtype=xp.float32)) + ref = xp.asarray(0.4349004, dtype=xp.float32) + xp_assert_close(res, ref) diff --git a/v12x/pylib/scipy/special/tests/test_trig.py b/v12x/pylib/scipy/special/tests/test_trig.py new file mode 100644 index 0000000000000000000000000000000000000000..e5f8f79f8d2cac51bf9476f04f13df2bfac587a5 --- /dev/null +++ b/v12x/pylib/scipy/special/tests/test_trig.py @@ -0,0 +1,76 @@ +import warnings + +import pytest +import numpy as np +from numpy.testing import assert_equal, assert_allclose + +from scipy.special._ufuncs import _sinpi as sinpi +from scipy.special._ufuncs import _cospi as cospi + + +def test_integer_real_part(): + x = np.arange(-100, 101) + y = np.hstack((-np.linspace(310, -30, 10), np.linspace(-30, 310, 10))) + x, y = np.meshgrid(x, y) + z = x + 1j*y + # In the following we should be *exactly* right + res = sinpi(z) + assert_equal(res.real, 0.0) + res = cospi(z) + assert_equal(res.imag, 0.0) + + +def test_half_integer_real_part(): + x = np.arange(-100, 101) + 0.5 + y = np.hstack((-np.linspace(310, -30, 10), np.linspace(-30, 310, 10))) + x, y = np.meshgrid(x, y) + z = x + 1j*y + # In the following we should be *exactly* right + res = sinpi(z) + assert_equal(res.imag, 0.0) + res = cospi(z) + assert_equal(res.real, 0.0) + + +@pytest.mark.skip("Temporary skip while gh-19526 is being resolved") +def test_intermediate_overlow(): + # Make sure we avoid overflow in situations where cosh/sinh would + # overflow but the product with sin/cos would not + sinpi_pts = [complex(1 + 1e-14, 227), + complex(1e-35, 250), + complex(1e-301, 445)] + # Data generated with mpmath + sinpi_std = [complex(-8.113438309924894e+295, -np.inf), + complex(1.9507801934611995e+306, np.inf), + complex(2.205958493464539e+306, np.inf)] + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", "invalid value encountered in multiply", RuntimeWarning) + for p, std in zip(sinpi_pts, sinpi_std): + res = sinpi(p) + assert_allclose(res.real, std.real) + assert_allclose(res.imag, std.imag) + + # Test for cosine, less interesting because cos(0) = 1. + p = complex(0.5 + 1e-14, 227) + std = complex(-8.113438309924894e+295, -np.inf) + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", "invalid value encountered in multiply", RuntimeWarning) + res = cospi(p) + assert_allclose(res.real, std.real) + assert_allclose(res.imag, std.imag) + + +def test_zero_sign(): + y = sinpi(-0.0) + assert y == 0.0 + assert np.signbit(y) + + y = sinpi(0.0) + assert y == 0.0 + assert not np.signbit(y) + + y = cospi(0.5) + assert y == 0.0 + assert not np.signbit(y) diff --git a/v12x/pylib/scipy/special/tests/test_ufunc_infra.py b/v12x/pylib/scipy/special/tests/test_ufunc_infra.py new file mode 100644 index 0000000000000000000000000000000000000000..30b08cde365de9c115935d1100001f8ba1349276 --- /dev/null +++ b/v12x/pylib/scipy/special/tests/test_ufunc_infra.py @@ -0,0 +1,34 @@ +import numpy as np +import pytest + +from scipy.special import mathieu_sem +# from scipy.special._ufuncs import _mathieu_sem + + +@pytest.mark.parametrize("m_shape,q_shape,x_shape,where_shape", [ + ((5, 1), (1, 1), (1, 10), ()), + ((2, 3), (2, 3), (2, 3), ()), + ((1, 5), (1, 5), (10, 5), ()), + ((5, 1), (1, 1), (1, 10), (1, 10)), + ((2, 3), (2, 3), (2, 3), (2, 3)), + ((1, 5), (1, 5), (10, 5), (10, 5)), + ((5, 1), (1, 1), (1, 10), (5, 1)), +]) +def test_out(m_shape, q_shape, x_shape, where_shape): + rng = np.random.default_rng(1234) + + m = rng.integers(1, 20, m_shape) + q = rng.uniform(0, 10, q_shape) + x = rng.uniform(0, 90, x_shape) + + if where_shape == (): + batch_shape = np.broadcast_shapes(m_shape, q_shape, x_shape) + where = True + else: + where = rng.choice([True, False], size=where_shape) + batch_shape = np.broadcast_shapes(m_shape, q_shape, x_shape, where_shape) + + out0 = np.empty(batch_shape) + out1 = np.empty(batch_shape) + res0, res1 = mathieu_sem(m, q, x, out=(out0, out1), where=where) + assert res0 is out0 and res1 is out1 diff --git a/v12x/pylib/scipy/special/tests/test_ufunc_signatures.py b/v12x/pylib/scipy/special/tests/test_ufunc_signatures.py new file mode 100644 index 0000000000000000000000000000000000000000..e4fdb02a855797923a792dfce7747a10f6b407db --- /dev/null +++ b/v12x/pylib/scipy/special/tests/test_ufunc_signatures.py @@ -0,0 +1,119 @@ +"""Test that all ufuncs have float32-preserving signatures. + +This was once guaranteed through the code generation script for +generating ufuncs, `scipy/special/_generate_pyx.py`. Starting with +gh-20260, SciPy developers have begun moving to generate ufuncs +through direct use of the NumPy C API (through C++). Existence of +float32 preserving signatures must now be tested since it is no +longer guaranteed. +""" + +import numpy as np +import pytest +import scipy.special._ufuncs +import scipy.special._gufuncs +import scipy.special as special + +from numpy.testing import assert_equal + + +# Single precision is not implemented for these ufuncs; +# floating point inputs must be float64. +exceptions = ['_gen_harmonic', '_normalized_gen_harmonic'] + +_ufuncs = [] +for funcname in dir(scipy.special._ufuncs): + if funcname not in exceptions: + _ufuncs.append(getattr(scipy.special._ufuncs, funcname)) +for funcname in dir(scipy.special._gufuncs): + _ufuncs.append(getattr(scipy.special._gufuncs, funcname)) + +# Not all module members are actually ufuncs +_ufuncs = [func for func in _ufuncs if isinstance(func, np.ufunc)] + +@pytest.mark.parametrize("ufunc", _ufuncs) +def test_ufunc_signatures(ufunc): + + # From _generate_pyx.py + # "Don't add float32 versions of ufuncs with integer arguments, as this + # can lead to incorrect dtype selection if the integer arguments are + # arrays, but float arguments are scalars. + # This may be a NumPy bug, but we need to work around it. + # cf. gh-4895, https://github.com/numpy/numpy/issues/5895" + types = set(sig for sig in ufunc.types + if not ("l" in sig or "i" in sig or "q" in sig or "p" in sig)) + + # Generate the full expanded set of signatures which should exist. There + # should be matching float and double versions of any existing signature. + expanded_types = set() + for sig in types: + expanded_types.update( + [sig.replace("d", "f").replace("D", "F"), + sig.replace("f", "d").replace("F", "D")] + ) + assert types == expanded_types + + +def _get_nan_val(typecode): + return { + "f": np.asarray(np.nan, dtype=np.float32), + "d": np.asarray(np.nan, dtype=np.float64), + "g": np.asarray(np.nan, dtype=np.longdouble), + "F": np.asarray(np.nan + np.nan*1j, dtype=np.complex64), + "D": np.asarray(np.nan + np.nan*1j, dtype=np.complex128), + "G": np.asarray(np.nan + np.nan*1j, dtype=np.clongdouble), + }[typecode] + +skips = { + # These hit https://github.com/scipy/xsf/issues/94 + "eval_chebyc", + "eval_chebys", + "eval_chebyt", + "eval_chebyu", + "eval_gegenbauer", + "eval_genlaguerre", + "eval_jacobi", + "eval_laguerre", + "eval_legendre", + "eval_sh_chebyt", + "eval_sh_chebyu", + "eval_sh_jacobi", + "eval_sh_legendre", + "hyp1f1", + "hyp2f1", + # warns if called with noninteger values of some args + "bdtr", + "bdtrc", + "bdtri", +} +_multi_arg_public_ufuncs = [ + ufunc for ufunc in _ufuncs if ufunc.nargs - ufunc.nout > 1 + and ufunc.__name__ not in skips and ufunc.__name__ in special.__all__ +] + + +@pytest.mark.parametrize("ufunc", _multi_arg_public_ufuncs) +def test_nep50(ufunc): + # Test that functions with multiple arguments respect nep50 promotion rules. + rng = np.random.default_rng(1234) + # As in test_ufunc_signatures, filter out signatures involving integers. + types = set(sig for sig in ufunc.types + if not ("l" in sig or "i" in sig or "q" in sig or "p" in sig)) + for sig in types: + input_types, output_types = sig.split("->") + # since we only care about dtypes and not values here, just use an appropriately + # typed nan for each argument. + args = [_get_nan_val(typecode) for typecode in input_types] + # swap out a random one of the nans with an appropriately typed numpy scalar. + idx = rng.choice(len(args)) + args[idx] = float("nan") if np.isrealobj(args[idx]) else complex("nan") + result = ufunc(*args) + result = [result] if len(output_types) == 1 else result + result = np.asarray(result) + + # Test that the output is an appropriately typed nan. This also implicitly + # tests that ufuncs propagate nans correctly. + assert_equal( + result, np.asarray([_get_nan_val(typecode) for typecode in output_types]), + strict=True + ) diff --git a/v12x/pylib/scipy/special/tests/test_wright_bessel.py b/v12x/pylib/scipy/special/tests/test_wright_bessel.py new file mode 100644 index 0000000000000000000000000000000000000000..ef30163f55c90850fe32ca600eb47547e4fa5e03 --- /dev/null +++ b/v12x/pylib/scipy/special/tests/test_wright_bessel.py @@ -0,0 +1,205 @@ +# Reference MPMATH implementation: +# +# import mpmath +# from mpmath import nsum +# +# def Wright_Series_MPMATH(a, b, z, dps=50, method='r+s+e', steps=[1000]): +# """Compute Wright' generalized Bessel function as Series. +# +# This uses mpmath for arbitrary precision. +# """ +# with mpmath.workdps(dps): +# res = nsum(lambda k: z**k/mpmath.fac(k) * mpmath.rgamma(a*k+b), +# [0, mpmath.inf], +# tol=dps, method=method, steps=steps +# ) +# +# return res + +from itertools import product + +import pytest +import numpy as np +from numpy.testing import assert_equal, assert_allclose + +import scipy.special as sc +from scipy.special import log_wright_bessel, loggamma, rgamma, wright_bessel + + +@pytest.mark.parametrize('a', [0, 1e-6, 0.1, 0.5, 1, 10]) +@pytest.mark.parametrize('b', [0, 1e-6, 0.1, 0.5, 1, 10]) +def test_wright_bessel_zero(a, b): + """Test at x = 0.""" + assert_equal(wright_bessel(a, b, 0.), rgamma(b)) + assert_allclose(log_wright_bessel(a, b, 0.), -loggamma(b)) + + +@pytest.mark.parametrize('b', [0, 1e-6, 0.1, 0.5, 1, 10]) +@pytest.mark.parametrize('x', [0, 1e-6, 0.1, 0.5, 1]) +def test_wright_bessel_iv(b, x): + """Test relation of wright_bessel and modified bessel function iv. + + iv(z) = (1/2*z)**v * Phi(1, v+1; 1/4*z**2). + See https://dlmf.nist.gov/10.46.E2 + """ + if x != 0: + v = b - 1 + wb = wright_bessel(1, v + 1, x**2 / 4.) + # Note: iv(v, x) has precision of less than 1e-12 for some cases + # e.g v=1-1e-6 and x=1e-06) + assert_allclose(np.power(x / 2., v) * wb, + sc.iv(v, x), + rtol=1e-11, atol=1e-11) + + +@pytest.mark.parametrize('a', [0, 1e-6, 0.1, 0.5, 1, 10]) +@pytest.mark.parametrize('b', [1, 1 + 1e-3, 2, 5, 10]) +@pytest.mark.parametrize('x', [0, 1e-6, 0.1, 0.5, 1, 5, 10, 100]) +def test_wright_functional(a, b, x): + """Test functional relation of wright_bessel. + + Phi(a, b-1, z) = a*z*Phi(a, b+a, z) + (b-1)*Phi(a, b, z) + + Note that d/dx Phi(a, b, x) = Phi(a, b-1, x) + See Eq. (22) of + B. Stankovic, On the Function of E. M. Wright, + Publ. de l' Institut Mathematique, Beograd, + Nouvelle S`er. 10 (1970), 113-124. + """ + assert_allclose(wright_bessel(a, b - 1, x), + a * x * wright_bessel(a, b + a, x) + + (b - 1) * wright_bessel(a, b, x), + rtol=1e-8, atol=1e-8) + + +# grid of rows [a, b, x, value, accuracy] that do not reach 1e-11 accuracy +# see output of: +# cd scipy/scipy/_precompute +# python wright_bessel_data.py +grid_a_b_x_value_acc = np.array([ + [0.1, 100.0, 709.7827128933841, 8.026353022981087e+34, 2e-8], + [0.5, 10.0, 709.7827128933841, 2.680788404494657e+48, 9e-8], + [0.5, 10.0, 1000.0, 2.005901980702872e+64, 1e-8], + [0.5, 100.0, 1000.0, 3.4112367580445246e-117, 6e-8], + [1.0, 20.0, 100000.0, 1.7717158630699857e+225, 3e-11], + [1.0, 100.0, 100000.0, 1.0269334596230763e+22, np.nan], + [1.0000000000000222, 20.0, 100000.0, 1.7717158630001672e+225, 3e-11], + [1.0000000000000222, 100.0, 100000.0, 1.0269334595866202e+22, np.nan], + [1.5, 0.0, 500.0, 15648961196.432373, 3e-11], + [1.5, 2.220446049250313e-14, 500.0, 15648961196.431465, 3e-11], + [1.5, 1e-10, 500.0, 15648961192.344728, 3e-11], + [1.5, 1e-05, 500.0, 15648552437.334162, 3e-11], + [1.5, 0.1, 500.0, 12049870581.10317, 2e-11], + [1.5, 20.0, 100000.0, 7.81930438331405e+43, 3e-9], + [1.5, 100.0, 100000.0, 9.653370857459075e-130, np.nan], + ]) + + +@pytest.mark.xfail +@pytest.mark.parametrize( + 'a, b, x, phi', + grid_a_b_x_value_acc[:, :4].tolist()) +def test_wright_data_grid_failures(a, b, x, phi): + """Test cases of test_data that do not reach relative accuracy of 1e-11""" + assert_allclose(wright_bessel(a, b, x), phi, rtol=1e-11) + + +@pytest.mark.parametrize( + 'a, b, x, phi, accuracy', + grid_a_b_x_value_acc.tolist()) +def test_wright_data_grid_less_accurate(a, b, x, phi, accuracy): + """Test cases of test_data that do not reach relative accuracy of 1e-11 + + Here we test for reduced accuracy or even nan. + """ + if np.isnan(accuracy): + assert np.isnan(wright_bessel(a, b, x)) + else: + assert_allclose(wright_bessel(a, b, x), phi, rtol=accuracy) + + +@pytest.mark.parametrize( + 'a, b, x', + list( + product([0, 0.1, 0.5, 1.5, 5, 10], [1, 2], [1e-3, 1, 1.5, 5, 10]) + ) +) +def test_log_wright_bessel_same_as_wright_bessel(a, b, x): + """Test that log_wright_bessel equals log of wright_bessel.""" + assert_allclose( + log_wright_bessel(a, b, x), + np.log(wright_bessel(a, b, x)), + rtol=1e-8, + ) + + +# Computed with, see also mp_wright_bessel from wright_bessel_data.py: +# +# from functools import lru_cache +# import mpmath as mp +# +# @lru_cache(maxsize=1_000_000) +# def rgamma_cached(x, dps): +# with mp.workdps(dps): +# return mp.rgamma(x) +# +# def mp_log_wright_bessel(a, b, x, dps=100, maxterms=10_000, method="d"): +# """Compute log of Wright's generalized Bessel function as Series with mpmath.""" +# with mp.workdps(dps): +# a, b, x = mp.mpf(a), mp.mpf(b), mp.mpf(x) +# res = mp.nsum(lambda k: x**k / mp.fac(k) +# * rgamma_cached(a * k + b, dps=dps), +# [0, mp.inf], +# tol=dps, method=method, steps=[maxterms] +# ) +# return mp.log(res) +# +# Sometimes, one needs to set maxterms as high as 1_00_000 to get accurate results for +# phi. +# At the end of the day, we can only hope that results are correct for very large x, +# e.g. by the asymptotic series, as there is no way to produce those in "exact" +# arithmetic. +# Note: accuracy = np.nan means log_wright_bessel returns nan. +@pytest.mark.parametrize( + 'a, b, x, phi, accuracy', + [ + (0, 0, 0, -np.inf, 1e-11), + (0, 0, 1, -np.inf, 1e-11), + (0, 1, 1.23, 1.23, 1e-11), + (0, 1, 1e50, 1e50, 1e-11), + (1e-5, 0, 700, 695.0421608273609, 1e-11), + (1e-5, 0, 1e3, 995.40052566540066, 1e-11), + (1e-5, 100, 1e3, 640.8197935670078, 1e-11), + (1e-3, 0, 1e4, 9987.2229532297262, 1e-11), + (1e-3, 0, 1e5, 99641.920687169507, 1e-11), + (1e-3, 0, 1e6, 994118.55560054416, 1e-11), # maxterms=1_000_000 + (1e-3, 10, 1e5, 99595.47710802537, 1e-11), + (1e-3, 50, 1e5, 99401.240922855647, 1e-3), + (1e-3, 100, 1e5, 99143.465191656527, np.nan), + (0.5, 0, 1e5, 4074.1112442197941, 1e-11), + (0.5, 0, 1e7, 87724.552120038896, 1e-11), + (0.5, 100, 1e5, 3350.3928746306163, np.nan), + (0.5, 100, 1e7, 86696.109975301719, 1e-11), + (1, 0, 1e5, 634.06765787997266, 1e-11), + (1, 0, 1e8, 20003.339639312035, 1e-11), + (1.5, 0, 1e5, 197.01777556071194, 1e-11), + (1.5, 0, 1e8, 3108.987414395706, 1e-11), + (1.5, 100, 1e8, 2354.8915946283275, np.nan), + (5, 0, 1e5, 9.8980480013203547, 1e-11), + (5, 0, 1e8, 33.642337258687465, 1e-11), + (5, 0, 1e12, 157.53704288117429, 1e-11), + (5, 100, 1e5, -359.13419630792148, 1e-11), + (5, 100, 1e12, -337.07722086995229, 1e-4), + (5, 100, 1e20, 2588.2471229986845, 2e-6), + (100, 0, 1e5, -347.62127990460517, 1e-11), + (100, 0, 1e20, -313.08250350969449, 1e-11), + (100, 100, 1e5, -359.1342053695754, 1e-11), + (100, 100, 1e20, -359.1342053695754, 1e-11), + ] +) +def test_log_wright_bessel(a, b, x, phi, accuracy): + """Test for log_wright_bessel, in particular for large x.""" + if np.isnan(accuracy): + assert np.isnan(log_wright_bessel(a, b, x)) + else: + assert_allclose(log_wright_bessel(a, b, x), phi, rtol=accuracy) diff --git a/v12x/pylib/scipy/special/tests/test_wrightomega.py b/v12x/pylib/scipy/special/tests/test_wrightomega.py new file mode 100644 index 0000000000000000000000000000000000000000..c1a93ca007e42fea3a0dec634c51b37f03effa9e --- /dev/null +++ b/v12x/pylib/scipy/special/tests/test_wrightomega.py @@ -0,0 +1,117 @@ +import pytest +import numpy as np +from numpy.testing import assert_, assert_equal, assert_allclose + +import scipy.special as sc +from scipy.special._testutils import assert_func_equal + + +def test_wrightomega_nan(): + pts = [complex(np.nan, 0), + complex(0, np.nan), + complex(np.nan, np.nan), + complex(np.nan, 1), + complex(1, np.nan)] + for p in pts: + res = sc.wrightomega(p) + assert_(np.isnan(res.real)) + assert_(np.isnan(res.imag)) + + +def test_wrightomega_inf_branch(): + pts = [complex(-np.inf, np.pi/4), + complex(-np.inf, -np.pi/4), + complex(-np.inf, 3*np.pi/4), + complex(-np.inf, -3*np.pi/4)] + expected_results = [complex(0.0, 0.0), + complex(0.0, -0.0), + complex(-0.0, 0.0), + complex(-0.0, -0.0)] + for p, expected in zip(pts, expected_results): + res = sc.wrightomega(p) + # We can't use assert_equal(res, expected) because in older versions of + # numpy, assert_equal doesn't check the sign of the real and imaginary + # parts when comparing complex zeros. It does check the sign when the + # arguments are *real* scalars. + assert_equal(res.real, expected.real) + assert_equal(res.imag, expected.imag) + + +def test_wrightomega_inf(): + pts = [complex(np.inf, 10), + complex(-np.inf, 10), + complex(10, np.inf), + complex(10, -np.inf)] + for p in pts: + assert_equal(sc.wrightomega(p), p) + + +def test_wrightomega_singular(): + pts = [complex(-1.0, np.pi), + complex(-1.0, -np.pi)] + for p in pts: + res = sc.wrightomega(p) + assert_equal(res, -1.0) + assert_(np.signbit(res.imag) == np.bool_(False)) + + +@pytest.mark.parametrize('x, desired', [ + (-np.inf, 0), + (np.inf, np.inf), +]) +def test_wrightomega_real_infinities(x, desired): + assert sc.wrightomega(x) == desired + + +def test_wrightomega_real_nan(): + assert np.isnan(sc.wrightomega(np.nan)) + + +def test_wrightomega_real_series_crossover(): + desired_error = 2 * np.finfo(float).eps + crossover = 1e20 + x_before_crossover = np.nextafter(crossover, -np.inf) + x_after_crossover = np.nextafter(crossover, np.inf) + # Computed using Mpmath + desired_before_crossover = 99999999999999983569.948 + desired_after_crossover = 100000000000000016337.948 + assert_allclose( + sc.wrightomega(x_before_crossover), + desired_before_crossover, + atol=0, + rtol=desired_error, + ) + assert_allclose( + sc.wrightomega(x_after_crossover), + desired_after_crossover, + atol=0, + rtol=desired_error, + ) + + +def test_wrightomega_exp_approximation_crossover(): + desired_error = 2 * np.finfo(float).eps + crossover = -50 + x_before_crossover = np.nextafter(crossover, np.inf) + x_after_crossover = np.nextafter(crossover, -np.inf) + # Computed using Mpmath + desired_before_crossover = 1.9287498479639314876e-22 + desired_after_crossover = 1.9287498479639040784e-22 + assert_allclose( + sc.wrightomega(x_before_crossover), + desired_before_crossover, + atol=0, + rtol=desired_error, + ) + assert_allclose( + sc.wrightomega(x_after_crossover), + desired_after_crossover, + atol=0, + rtol=desired_error, + ) + + +def test_wrightomega_real_versus_complex(): + x = np.linspace(-500, 500, 1001) + results = sc.wrightomega(x + 0j).real + assert_func_equal(sc.wrightomega, results, x, atol=0, rtol=1e-14) diff --git a/v12x/pylib/scipy/special/tests/test_zeta.py b/v12x/pylib/scipy/special/tests/test_zeta.py new file mode 100644 index 0000000000000000000000000000000000000000..988f9a716196fb8f65a92a0a2e894038c2474cfa --- /dev/null +++ b/v12x/pylib/scipy/special/tests/test_zeta.py @@ -0,0 +1,301 @@ +import scipy +import scipy.special as sc +import sys +import numpy as np +import pytest + +from numpy.testing import assert_equal, assert_allclose + + +def test_zeta(): + assert_allclose(sc.zeta(2,2), np.pi**2/6 - 1, rtol=1e-12) + + +def test_zetac(): + # Expected values in the following were computed using Wolfram + # Alpha's `Zeta[x] - 1` + x = [-2.1, 0.8, 0.9999, 9, 50, 75] + desired = [ + -0.9972705002153750, + -5.437538415895550, + -10000.42279161673, + 0.002008392826082214, + 8.881784210930816e-16, + 2.646977960169853e-23, + ] + assert_allclose(sc.zetac(x), desired, rtol=1e-12) + + +def test_zetac_special_cases(): + assert sc.zetac(np.inf) == 0 + assert np.isnan(sc.zetac(-np.inf)) + assert sc.zetac(0) == -1.5 + assert sc.zetac(1.0) == np.inf + + assert_equal(sc.zetac([-2, -50, -100]), -1) + + +def test_riemann_zeta_special_cases(): + assert np.isnan(sc.zeta(np.nan)) + assert sc.zeta(np.inf) == 1 + assert sc.zeta(0) == -0.5 + + # Riemann zeta is zero add negative even integers. + assert_equal(sc.zeta([-2, -4, -6, -8, -10]), 0) + + assert_allclose(sc.zeta(2), np.pi**2/6, rtol=1e-12) + assert_allclose(sc.zeta(4), np.pi**4/90, rtol=1e-12) + + +def test_riemann_zeta_avoid_overflow(): + s = -260.00000000001 + desired = -5.6966307844402683127e+297 # Computed with Mpmath + assert_allclose(sc.zeta(s), desired, atol=0, rtol=5e-14) + + +@pytest.mark.parametrize( + "z, desired, rtol", + [ + ## Test cases taken from mpmath with the script: + + # import numpy as np + # import scipy.stats as stats + + # from mpmath import mp + + # # seed = np.random.SeedSequence().entropy + # seed = 154689806791763421822480125722191067828 + # rng = np.random.default_rng(seed) + # default_rtol = 1e-13 + + # # A small point in each quadrant outside of the critical strip + # cases = [] + # for x_sign, y_sign in [1, 1], [1, -1], [-1, 1], [-1, -1]: + # x = x_sign * rng.uniform(2, 8) + # y = y_sign * rng.uniform(2, 8) + # z = x + y*1j + # reference = complex(mp.zeta(z)) + # cases.append((z, reference, default_rtol)) + + # # Moderately large imaginary part in each quadrant outside of critical strip + # for x_sign, y_sign in [1, 1], [1, -1], [-1, 1], [-1, -1]: + # x = x_sign * rng.uniform(2, 8) + # y = y_sign * rng.uniform(50, 80) + # z = x + y*1j + # reference = complex(mp.zeta(z)) + # cases.append((z, reference, default_rtol)) + + # # points in critical strip + # x = rng.uniform(0.0, 1.0, size=5) + # y = np.exp(rng.uniform(0, 5, size=5)) + # z = x + y*1j + # for t in z: + # reference = complex(mp.zeta(t)) + # cases.append((complex(t), reference, default_rtol)) + # z = x - y*1j + # for t in z: + # reference = complex(mp.zeta(t)) + # cases.append((complex(t), reference, default_rtol)) + + # # Near small trivial zeros + # x = np.array([-2, -4, -6, -8]) + # y = np.array([1e-15, -1e-15]) + # x, y = np.meshgrid(x, y) + # x, y = x.ravel(), y.ravel() + # z = x + y*1j + # for t in z: + # reference = complex(mp.zeta(t)) + # cases.append((complex(t), reference, 1e-7)) + + # # Some other points near real axis + # x = np.array([-0.5, 0, 0.2, 0.75]) + # y = np.array([1e-15, -1e-15]) + # x, y = np.meshgrid(x, y) + # x, y = x.ravel(), y.ravel() + # z = x + y*1j + # for t in z: + # reference = complex(mp.zeta(t)) + # cases.append((complex(t), reference, 1e-7)) + + # # Moderately large real part + # x = np.array([49.33915930750887, 50.55805244181687]) + # y = rng.uniform(20, 100, size=3) + # x, y = np.meshgrid(x, y) + # x, y = x.ravel(), y.ravel() + # z = x + y*1j + # for t in z: + # reference = complex(mp.zeta(t)) + # cases.append((complex(t), reference, default_rtol)) + + # # Very large imaginary part + # x = np.array([0.5, 34.812847097948854, 50.55805244181687]) + # y = np.array([1e6, -1e6]) + # x, y = np.meshgrid(x, y) + # x, y = x.ravel(), y.ravel() + # z = x + y*1j + # for t in z: + # reference = complex(mp.zeta(t)) + # rtol = 1e-7 if t.real == 0.5 else default_rtol + # cases.append((complex(t), reference, rtol)) + # + # # Naive implementation of reflection formula suffers internal overflow + # x = -rng.uniform(200, 300, 3) + # y = np.array([rng.uniform(10, 30), -rng.uniform(10, 30)]) + # x, y = np.meshgrid(x, y) + # x, y = x.ravel(), y.ravel() + # z = x + y*1j + # for t in z: + # reference = complex(mp.zeta(t)) + # cases.append((complex(t), reference, default_rtol)) + # + # A small point in each quadrant outside of the critical strip + ((3.12838509346655+7.111085974836645j), + (1.0192654793474945+0.08795174413289127j), + 1e-13), + ((7.06791362314716-7.219497492626728j), + (1.0020740683598117-0.006752725913243711j), + 1e-13), + ((-6.806227077655519+2.724411451005281j), + (0.06312488213559667-0.061641496333765956j), + 1e-13), + ((-3.0170751511621026-6.3686522550665945j), + (-0.10330747857150148-1.214541994832571j), + 1e-13), + # Moderately large imaginary part in each quadrant outside of critical strip + ((6.133994402212294+60.03091448000761j), + (0.9885701843417336+0.009636925981078128j), + 1e-13), + ((6.17268142822657-64.74883149743795j), + (1.0080474225840865+0.012032804974965354j), + 1e-13), + ((-3.462191939791879+76.16258975567534j), + (18672.072070850158+2908.5104826247184j), + 1e-13), + ((-6.955735216531752-74.75791554155748j), + (-77672258.72276545+71625206.0401107j), + 1e-13), + # Points in critical strip + ((0.4088038289823922+1.4596830498094384j), + (0.3032837969400845-0.47272237994110344j), + 1e-13), + ((0.9673493951209633+4.918968547259143j), + (0.7488756907431944+0.17281553371482428j), + 1e-13), + ((0.8692482679977754+66.6142398421354j), + (0.5831942469552066-0.26848904799062334j), + 1e-13), + ((0.42771847720003764+21.783747851715468j), + (0.4767032638444329+0.6898148744603123j), + 1e-13), + ((0.20479494678428956+33.17656449538932j), + (-0.6983038977487848+0.18060923618150224j), + 1e-13), + ((0.4088038289823922-1.4596830498094384j), + (0.3032837969400845+0.47272237994110344j), + 1e-13), + ((0.9673493951209633-4.918968547259143j), + (0.7488756907431944-0.17281553371482428j), + 1e-13), + ((0.8692482679977754-66.6142398421354j), + (0.5831942469552066+0.26848904799062334j), + 1e-13), + ((0.42771847720003764-21.783747851715468j), + (0.4767032638444329-0.6898148744603123j), + 1e-13), + ((0.20479494678428956-33.17656449538932j), + (-0.6983038977487848-0.18060923618150224j), + 1e-13), + # Near small trivial zeros + ((-2+1e-15j), (3.288175809370978e-32-3.0448457058393275e-17j), 1e-07), + ((-4+1e-15j), (-2.868707923051182e-33+7.983811450268625e-18j), 1e-07), + ((-6+1e-15j), (-1.7064292323640116e-34-5.8997591435159376e-18j), 1e-07), + ((-8+1e-15j), (2.5060859548261706e-33+8.316161985602247e-18j), 1e-07), + ((-2-1e-15j), (3.288175809371319e-32+3.0448457058393275e-17j), 1e-07), + ((-4-1e-15j), (-2.8687079230520114e-33-7.983811450268625e-18j), 1e-07), + ((-6-1e-15j), (-1.70642923235801e-34+5.8997591435159376e-18j), 1e-07), + ((-8-1e-15j), (2.5060859548253293e-33-8.316161985602247e-18j), 1e-07), + # Some other points near real axis + ((-0.5+1e-15j), (-0.20788622497735457-3.608543395999408e-16j), 1e-07), + (1e-15j, (-0.5-9.189384239689193e-16j), 1e-07), + ((0.2+1e-15j), (-0.7339209248963406-1.4828001150329085e-15j), 1e-07), + ((0.75+1e-15j), (-3.4412853869452227-1.5924832114302393e-14j), 1e-13), + ((-0.5-1e-15j), (-0.20788622497735457+3.608543395999408e-16j), 1e-07), + (-1e-15j, (-0.5+9.189387416062746e-16j), 1e-07), + ((0.2-1e-15j), (-0.7339209248963406+1.4828004007675122e-15j), 1e-07), + ((0.75-1e-15j), (-3.4412853869452227+1.5924831974403957e-14j), 1e-13), + # Moderately large real part + ((49.33915930750887+53.213478698903955j), + (1.0000000000000009+1.0212452494616078e-15j), + 1e-13), + ((50.55805244181687+53.213478698903955j), + (1.0000000000000004+4.387394180390787e-16j), + 1e-13), + ((49.33915930750887+40.6366015728302j), + (0.9999999999999986-1.502268709924849e-16j), + 1e-13), + ((50.55805244181687+40.6366015728302j), + (0.9999999999999994-6.453929613571651e-17j), + 1e-13), + ((49.33915930750887+85.83555435273925j), + (0.9999999999999987-2.7014400611995846e-16j), + 1e-13), + ((50.55805244181687+85.83555435273925j), + (0.9999999999999994-1.160571605555322e-16j), + 1e-13), + # Very large imaginary part + ((0.5+1e6j), (0.0760890697382271+2.805102101019299j), 1e-07), + ((34.812847097948854+1e6j), + (1.0000000000102545+3.150848654056419e-11j), + 1e-13), + ((50.55805244181687+1e6j), + (1.0000000000000002+5.736517078070873e-16j), + 1e-13), + ((0.5-1e6j), (0.0760890697382271-2.805102101019299j), 1e-07), + ((34.812847097948854-1e6j), + (1.0000000000102545-3.150848654056419e-11j), + 1e-13), + ((50.55805244181687-1e6j), + (1.0000000000000002-5.736517078070873e-16j), + 1e-13), + ((-294.86605461349745+13.992648136816397j), (-np.inf+np.inf*1j), 1e-13), + ((-294.86605461349745-16.147667799398363j), (np.inf-np.inf*1j), 1e-13), + ] +) +def test_riemann_zeta_complex(z, desired, rtol): + assert_allclose(sc.zeta(z), desired, rtol=rtol) + + +# Some of the test cases below fail for intel compilers +cpp_compiler = scipy.__config__.CONFIG["Compilers"]["c++"]["name"] +gcc_linux = cpp_compiler == "gcc" and sys.platform == "linux" +clang_macOS = cpp_compiler == "clang" and sys.platform == "darwin" + + +@pytest.mark.skipif( + not (gcc_linux or clang_macOS), + reason="Underflow may not be avoided on other platforms", +) +@pytest.mark.parametrize( + "z, desired, rtol", + [ + # Test cases generated as part of same script for + # test_riemann_zeta_complex. These cases are split off because + # they fail on some platforms. + # + # Naive implementation of reflection formula suffers internal overflow + ((-217.40285743524163+13.992648136816397j), + (-6.012818500554211e+249-1.926943776932387e+250j), + 5e-13,), + ((-237.71710702931668+13.992648136816397j), + (-8.823803086106129e+281-5.009074181335139e+281j), + 1e-13,), + ((-217.40285743524163-16.147667799398363j), + (-5.111612904844256e+251-4.907132127666742e+250j), + 5e-13,), + ((-237.71710702931668-16.147667799398363j), + (-1.3256112779883167e+283-2.253002003455494e+283j), + 5e-13,), + ], +) +def test_riemann_zeta_complex_avoid_underflow(z, desired, rtol): + assert_allclose(sc.zeta(z), desired, rtol=rtol) diff --git a/v12x/pylib/scipy/stats/__init__.py b/v12x/pylib/scipy/stats/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e67ecbe12091add03d50815fe8307e746af3b85b --- /dev/null +++ b/v12x/pylib/scipy/stats/__init__.py @@ -0,0 +1,644 @@ +""" +.. _statsrefmanual: + +========================================== +Statistical functions (:mod:`scipy.stats`) +========================================== + +.. currentmodule:: scipy.stats + +This module contains a large number of probability distributions, +summary and frequency statistics, correlation functions and statistical +tests, masked statistics, kernel density estimation, quasi-Monte Carlo +functionality, and more. + +Statistics is a very large area, and there are topics that are out of scope +for SciPy and are covered by other packages. Some of the most important ones +are: + +- `statsmodels `__: + regression, linear models, time series analysis, extensions to topics + also covered by ``scipy.stats``. +- `Pandas `__: tabular data, time series + functionality, interfaces to other statistical languages. +- `PyMC `__: Bayesian statistical + modeling, probabilistic machine learning. +- `scikit-learn `__: classification, regression, + model selection. +- `Seaborn `__: statistical data visualization. +- `rpy2 `__: Python to R bridge. + + +Probability distributions +========================= + +Each univariate distribution is an instance of a subclass of `rv_continuous` +(`rv_discrete` for discrete distributions): + +.. autosummary:: + :toctree: generated/ + + rv_continuous + rv_discrete + rv_histogram + +Continuous distributions +------------------------ + +.. autosummary:: + :toctree: generated/ + + alpha -- Alpha + anglit -- Anglit + arcsine -- Arcsine + argus -- Argus + beta -- Beta + betaprime -- Beta Prime + bradford -- Bradford + burr -- Burr (Type III) + burr12 -- Burr (Type XII) + cauchy -- Cauchy + chi -- Chi + chi2 -- Chi-squared + cosine -- Cosine + crystalball -- Crystalball + dgamma -- Double Gamma + dpareto_lognorm -- Double Pareto Lognormal + dweibull -- Double Weibull + erlang -- Erlang + expon -- Exponential + exponnorm -- Exponentially Modified Normal + exponweib -- Exponentiated Weibull + exponpow -- Exponential Power + f -- F (Snecdor F) + fatiguelife -- Fatigue Life (Birnbaum-Saunders) + fisk -- Fisk + foldcauchy -- Folded Cauchy + foldnorm -- Folded Normal + genlogistic -- Generalized Logistic + gennorm -- Generalized normal + genpareto -- Generalized Pareto + genexpon -- Generalized Exponential + genextreme -- Generalized Extreme Value + gausshyper -- Gauss Hypergeometric + gamma -- Gamma + gengamma -- Generalized gamma + genhalflogistic -- Generalized Half Logistic + genhyperbolic -- Generalized Hyperbolic + geninvgauss -- Generalized Inverse Gaussian + gibrat -- Gibrat + gompertz -- Gompertz (Truncated Gumbel) + gumbel_r -- Right Sided Gumbel, Log-Weibull, Fisher-Tippett, Extreme Value Type I + gumbel_l -- Left Sided Gumbel, etc. + halfcauchy -- Half Cauchy + halflogistic -- Half Logistic + halfnorm -- Half Normal + halfgennorm -- Generalized Half Normal + hypsecant -- Hyperbolic Secant + invgamma -- Inverse Gamma + invgauss -- Inverse Gaussian + invweibull -- Inverse Weibull + irwinhall -- Irwin-Hall + jf_skew_t -- Jones and Faddy Skew-T + johnsonsb -- Johnson SB + johnsonsu -- Johnson SU + kappa4 -- Kappa 4 parameter + kappa3 -- Kappa 3 parameter + ksone -- Distribution of Kolmogorov-Smirnov one-sided test statistic + kstwo -- Distribution of Kolmogorov-Smirnov two-sided test statistic + kstwobign -- Limiting Distribution of scaled Kolmogorov-Smirnov two-sided test statistic. + landau -- Landau + laplace -- Laplace + laplace_asymmetric -- Asymmetric Laplace + levy -- Levy + levy_l + levy_stable + logistic -- Logistic + loggamma -- Log-Gamma + loglaplace -- Log-Laplace (Log Double Exponential) + lognorm -- Log-Normal + loguniform -- Log-Uniform + lomax -- Lomax (Pareto of the second kind) + maxwell -- Maxwell + mielke -- Mielke's Beta-Kappa + moyal -- Moyal + nakagami -- Nakagami + ncx2 -- Non-central chi-squared + ncf -- Non-central F + nct -- Non-central Student's T + norm -- Normal (Gaussian) + norminvgauss -- Normal Inverse Gaussian + pareto -- Pareto + pearson3 -- Pearson type III + powerlaw -- Power-function + powerlognorm -- Power log normal + powernorm -- Power normal + rdist -- R-distribution + rayleigh -- Rayleigh + rel_breitwigner -- Relativistic Breit-Wigner + rice -- Rice + recipinvgauss -- Reciprocal Inverse Gaussian + semicircular -- Semicircular + skewcauchy -- Skew Cauchy + skewnorm -- Skew normal + studentized_range -- Studentized Range + t -- Student's T + trapezoid -- Trapezoidal + triang -- Triangular + truncexpon -- Truncated Exponential + truncnorm -- Truncated Normal + truncpareto -- Truncated Pareto + truncweibull_min -- Truncated minimum Weibull distribution + tukeylambda -- Tukey-Lambda + uniform -- Uniform + vonmises -- Von-Mises (Circular) + vonmises_line -- Von-Mises (Line) + wald -- Wald + weibull_min -- Minimum Weibull (see Frechet) + weibull_max -- Maximum Weibull (see Frechet) + wrapcauchy -- Wrapped Cauchy + +The ``fit`` method of the univariate continuous distributions uses +maximum likelihood estimation to fit the distribution to a data set. +The ``fit`` method can accept regular data or *censored data*. +Censored data is represented with instances of the `CensoredData` +class. + +.. autosummary:: + :toctree: generated/ + + CensoredData + + +Multivariate distributions +-------------------------- + +.. autosummary:: + :toctree: generated/ + + multivariate_normal -- Multivariate normal distribution + matrix_normal -- Matrix normal distribution + dirichlet -- Dirichlet + dirichlet_multinomial -- Dirichlet multinomial distribution + wishart -- Wishart + invwishart -- Inverse Wishart + multinomial -- Multinomial distribution + special_ortho_group -- SO(N) group + ortho_group -- O(N) group + unitary_group -- U(N) group + random_correlation -- random correlation matrices + multivariate_t -- Multivariate t-distribution + multivariate_hypergeom -- Multivariate hypergeometric distribution + normal_inverse_gamma -- Normal-inverse-gamma distribution + random_table -- Distribution of random tables with given marginals + uniform_direction -- Uniform distribution on S(N-1) + vonmises_fisher -- Von Mises-Fisher distribution + matrix_t -- Matrix variate t distribution + +`scipy.stats.multivariate_normal` methods accept instances +of the following class to represent the covariance. + +.. autosummary:: + :toctree: generated/ + + Covariance -- Representation of a covariance matrix + + +Discrete distributions +---------------------- + +.. autosummary:: + :toctree: generated/ + + bernoulli -- Bernoulli + betabinom -- Beta-Binomial + betanbinom -- Beta-Negative Binomial + binom -- Binomial + boltzmann -- Boltzmann (Truncated Discrete Exponential) + dlaplace -- Discrete Laplacian + geom -- Geometric + hypergeom -- Hypergeometric + logser -- Logarithmic (Log-Series, Series) + nbinom -- Negative Binomial + nchypergeom_fisher -- Fisher's Noncentral Hypergeometric + nchypergeom_wallenius -- Wallenius's Noncentral Hypergeometric + nhypergeom -- Negative Hypergeometric + planck -- Planck (Discrete Exponential) + poisson -- Poisson + poisson_binom -- Poisson Binomial + randint -- Discrete Uniform + skellam -- Skellam + yulesimon -- Yule-Simon + zipf -- Zipf (Zeta) + zipfian -- Zipfian + + +An overview of statistical functions is given below. Many of these functions +have a similar version in `scipy.stats.mstats` which work for masked arrays. + +Summary statistics +================== + +.. autosummary:: + :toctree: generated/ + + describe -- Descriptive statistics + gmean -- Geometric mean + hmean -- Harmonic mean + pmean -- Power mean + kurtosis -- Fisher or Pearson kurtosis + mode -- Modal value + moment -- Central moment + lmoment + expectile -- Expectile + skew -- Skewness + kstat -- + kstatvar -- + tmean -- Truncated arithmetic mean + tvar -- Truncated variance + tmin -- + tmax -- + tstd -- + tsem -- + variation -- Coefficient of variation + rankdata + tiecorrect + trim_mean + gstd -- Geometric Standard Deviation + iqr + sem + bayes_mvs + mvsdist + entropy + differential_entropy + median_abs_deviation + +Frequency statistics +==================== + +.. autosummary:: + :toctree: generated/ + + quantile + estimated_cdf + cumfreq + relfreq + percentileofscore + scoreatpercentile + +.. autosummary:: + :toctree: generated/ + + binned_statistic -- Compute a binned statistic for a set of data. + binned_statistic_2d -- Compute a 2-D binned statistic for a set of data. + binned_statistic_dd -- Compute a d-D binned statistic for a set of data. + +.. _hypotests: + +Hypothesis Tests and related functions +====================================== +SciPy has many functions for performing hypothesis tests that return a +test statistic and a p-value, and several of them return confidence intervals +and/or other related information. + +The headings below are based on common uses of the functions within, but due to +the wide variety of statistical procedures, any attempt at coarse-grained +categorization will be imperfect. Also, note that tests within the same heading +are not interchangeable in general (e.g. many have different distributional +assumptions). + +One Sample Tests / Paired Sample Tests +-------------------------------------- +One sample tests are typically used to assess whether a single sample was +drawn from a specified distribution or a distribution with specified properties +(e.g. zero mean). + +.. autosummary:: + :toctree: generated/ + + ttest_1samp + binomtest + quantile_test + skewtest + kurtosistest + normaltest + jarque_bera + shapiro + anderson + cramervonmises + ks_1samp + goodness_of_fit + chisquare + power_divergence + +Paired sample tests are often used to assess whether two samples were drawn +from the same distribution; they differ from the independent sample tests below +in that each observation in one sample is treated as paired with a +closely-related observation in the other sample (e.g. when environmental +factors are controlled between observations within a pair but not among pairs). +They can also be interpreted or used as one-sample tests (e.g. tests on the +mean or median of *differences* between paired observations). + +.. autosummary:: + :toctree: generated/ + + ttest_rel + wilcoxon + +Association/Correlation Tests +----------------------------- + +These tests are often used to assess whether there is a relationship (e.g. +linear) between paired observations in multiple samples or among the +coordinates of multivariate observations. + +.. autosummary:: + :toctree: generated/ + + linregress + pearsonr + spearmanrho + pointbiserialr + kendalltau + chatterjeexi + weightedtau + somersd + siegelslopes + theilslopes + page_trend_test + multiscale_graphcorr + spearmanr + +These association tests and are to work with samples in the form of contingency +tables. Supporting functions are available in `scipy.stats.contingency`. + +.. autosummary:: + :toctree: generated/ + + chi2_contingency + fisher_exact + barnard_exact + boschloo_exact + +Independent Sample Tests +------------------------ +Independent sample tests are typically used to assess whether multiple samples +were independently drawn from the same distribution or different distributions +with a shared property (e.g. equal means). + +Some tests are specifically for comparing two samples. + +.. autosummary:: + :toctree: generated/ + + ttest_ind_from_stats + poisson_means_test + ttest_ind + mannwhitneyu + bws_test + ranksums + brunnermunzel + mood + ansari + cramervonmises_2samp + epps_singleton_2samp + ks_2samp + kstest + +Others are generalized to multiple samples. + +.. autosummary:: + :toctree: generated/ + + f_oneway + tukey_hsd + dunnett + kruskal + alexandergovern + fligner + levene + bartlett + median_test + friedmanchisquare + anderson_ksamp + +Resampling and Monte Carlo Methods +---------------------------------- +The following functions can reproduce the p-value and confidence interval +results of most of the functions above, and often produce accurate results in a +wider variety of conditions. They can also be used to perform hypothesis tests +and generate confidence intervals for custom statistics. This flexibility comes +at the cost of greater computational requirements and stochastic results. + +.. autosummary:: + :toctree: generated/ + + monte_carlo_test + permutation_test + bootstrap + power + +Instances of the following object can be passed into some hypothesis test +functions to perform a resampling or Monte Carlo version of the hypothesis +test. + +.. autosummary:: + :toctree: generated/ + + MonteCarloMethod + PermutationMethod + BootstrapMethod + +Multiple Hypothesis Testing and Meta-Analysis +--------------------------------------------- +These functions are for assessing the results of individual tests as a whole. +Functions for performing specific multiple hypothesis tests (e.g. post hoc +tests) are listed above. + +.. autosummary:: + :toctree: generated/ + + combine_pvalues + false_discovery_control + +Random Variables +================ + +.. autosummary:: + :toctree: generated/ + + make_distribution + Normal + Logistic + Uniform + Binomial + Mixture + order_statistic + truncate + abs + exp + log + +Other statistical functionality +=============================== +.. toctree:: + :maxdepth: 1 + + stats.qmc + stats.contingency + stats.mstats + stats.sampling + +Transformations +--------------- + +.. autosummary:: + :toctree: generated/ + + boxcox + boxcox_normmax + boxcox_llf + yeojohnson + yeojohnson_normmax + yeojohnson_llf + obrientransform + sigmaclip + trimboth + trim1 + zmap + zscore + gzscore + +Statistical distances +--------------------- + +.. autosummary:: + :toctree: generated/ + + wasserstein_distance + wasserstein_distance_nd + energy_distance + +Fitting / Survival Analysis +--------------------------- + +.. autosummary:: + :toctree: generated/ + + fit + ecdf + logrank + +Directional statistical functions +--------------------------------- + +.. autosummary:: + :toctree: generated/ + + directional_stats + circmean + circvar + circstd + +Sensitivity Analysis +-------------------- + +.. autosummary:: + :toctree: generated/ + + sobol_indices + +Plot-tests +---------- + +.. autosummary:: + :toctree: generated/ + + ppcc_max + ppcc_plot + probplot + boxcox_normplot + yeojohnson_normplot + +Univariate and multivariate kernel density estimation +----------------------------------------------------- + +.. autosummary:: + :toctree: generated/ + + gaussian_kde + +Warnings / Errors used in :mod:`scipy.stats` +-------------------------------------------- + +.. autosummary:: + :toctree: generated/ + + DegenerateDataWarning + ConstantInputWarning + NearConstantInputWarning + FitError + +Result classes used in :mod:`scipy.stats` +----------------------------------------- + +.. warning:: + + These classes are private, but they are included here because instances + of them are returned by other statistical functions. User import and + instantiation is not supported. + +.. toctree:: + :maxdepth: 2 + + stats._result_classes + +""" # noqa: E501 + +from ._warnings_errors import (ConstantInputWarning, NearConstantInputWarning, + DegenerateDataWarning, FitError) +from ._stats_py import * +from ._variation import variation +from .distributions import * +from ._morestats import * +from ._multicomp import * +from ._binomtest import binomtest +from ._binned_statistic import * +from ._kde import gaussian_kde +from . import mstats +from . import qmc +from ._multivariate import * +from . import contingency +from .contingency import chi2_contingency +from ._censored_data import CensoredData +from ._resampling import (bootstrap, monte_carlo_test, permutation_test, power, + MonteCarloMethod, PermutationMethod, BootstrapMethod) +from ._entropy import * +from ._hypotests import * +from ._page_trend_test import page_trend_test +from ._mannwhitneyu import mannwhitneyu +from ._bws_test import bws_test +from ._fit import fit, goodness_of_fit +from ._covariance import Covariance +from ._sensitivity_analysis import * +from ._survival import * +from ._distribution_infrastructure import ( + make_distribution, Mixture, order_statistic, truncate, exp, log, abs +) +from ._new_distributions import Normal, Logistic, Uniform, Binomial +from ._mgc import multiscale_graphcorr +from ._correlation import chatterjeexi, spearmanrho, theilslopes, siegelslopes +from ._quantile import quantile, estimated_cdf + + +# Deprecated namespaces, to be removed in v2.0.0 +from . import ( + biasedurn, kde, morestats, mstats_basic, mstats_extras, mvn, stats +) + + +__all__ = [s for s in dir() if not s.startswith("_")] # Remove dunders. + +from scipy._lib._testutils import PytestTester +test = PytestTester(__name__) +del PytestTester diff --git a/v12x/pylib/scipy/stats/_axis_nan_policy.py b/v12x/pylib/scipy/stats/_axis_nan_policy.py new file mode 100644 index 0000000000000000000000000000000000000000..5b47171d2667e066af9a6e74c6a0fb739f8d9480 --- /dev/null +++ b/v12x/pylib/scipy/stats/_axis_nan_policy.py @@ -0,0 +1,725 @@ +# Many scipy.stats functions support `axis` and `nan_policy` parameters. +# When the two are combined, it can be tricky to get all the behavior just +# right. This file contains utility functions useful for scipy.stats functions +# that support `axis` and `nan_policy`, including a decorator that +# automatically adds `axis` and `nan_policy` arguments to a function. + +import math +import warnings +import numpy as np +from functools import wraps +from scipy._lib._array_api import xp_ravel +from scipy._lib._docscrape import FunctionDoc, Parameter +from scipy._lib._util import _contains_nan, AxisError, _get_nan +from scipy._lib._array_api import (array_namespace, is_numpy, xp_size, xp_copy, + xp_promote, is_dask, is_jax, xp_capabilities) +import scipy._external.array_api_extra as xpx + +import inspect + +too_small_1d_not_omit = ( + "One or more sample arguments is too small; all " + "returned values will be NaN. " + "See documentation for sample size requirements.") + +too_small_1d_omit = ( + "After omitting NaNs, one or more sample arguments " + "is too small; all returned values will be NaN. " + "See documentation for sample size requirements.") + +too_small_nd_not_omit = ( + "All axis-slices of one or more sample arguments are " + "too small; all elements of returned arrays will be NaN. " + "See documentation for sample size requirements.") + +too_small_nd_omit = ( + "After omitting NaNs, one or more axis-slices of one " + "or more sample arguments is too small; corresponding " + "elements of returned arrays will be NaN. " + "See documentation for sample size requirements.") + +class SmallSampleWarning(RuntimeWarning): + pass + + +def _broadcast_arrays(arrays, axis=None, xp=None): + """ + Broadcast shapes of arrays, ignoring incompatibility of specified axes + """ + arrays = tuple(arrays) + if not arrays: + return arrays + xp = array_namespace(*arrays) if xp is None else xp + arrays = [xp.asarray(arr) for arr in arrays] + shapes = [arr.shape for arr in arrays] + new_shapes = _broadcast_shapes(shapes, axis) + if axis is None: + new_shapes = [new_shapes]*len(arrays) + return [xp.broadcast_to(array, new_shape) + for array, new_shape in zip(arrays, new_shapes)] + + +def _broadcast_shapes(shapes, axis=None): + """ + Broadcast shapes, ignoring incompatibility of specified axes + """ + if not shapes: + return shapes + + # input validation + if axis is not None: + axis = np.atleast_1d(axis) + message = '`axis` must be an integer, a tuple of integers, or `None`.' + try: + with np.errstate(invalid='ignore'): + axis_int = axis.astype(int) + except ValueError as e: + raise AxisError(message) from e + if not np.array_equal(axis_int, axis): + raise AxisError(message) + axis = axis_int + + # First, ensure all shapes have same number of dimensions by prepending 1s. + n_dims = max([len(shape) for shape in shapes]) + new_shapes = np.ones((len(shapes), n_dims), dtype=int) + for row, shape in zip(new_shapes, shapes): + row[len(row)-len(shape):] = shape # can't use negative indices (-0:) + + # Remove the shape elements of the axes to be ignored, but remember them. + if axis is not None: + axis[axis < 0] = n_dims + axis[axis < 0] + axis = np.sort(axis) + if axis[-1] >= n_dims or axis[0] < 0: + message = (f"`axis` is out of bounds " + f"for array of dimension {n_dims}") + raise AxisError(message) + + if len(np.unique(axis)) != len(axis): + raise AxisError("`axis` must contain only distinct elements") + + removed_shapes = new_shapes[:, axis] + new_shapes = np.delete(new_shapes, axis, axis=1) + + # If arrays are broadcastable, shape elements that are 1 may be replaced + # with a corresponding non-1 shape element. Assuming arrays are + # broadcastable, that final shape element can be found with: + new_shape = np.max(new_shapes, axis=0) + # except in case of an empty array: + new_shape *= new_shapes.all(axis=0) + + # Among all arrays, there can only be one unique non-1 shape element. + # Therefore, if any non-1 shape element does not match what we found + # above, the arrays must not be broadcastable after all. + if np.any(~((new_shapes == 1) | (new_shapes == new_shape))): + raise ValueError("Array shapes are incompatible for broadcasting.") + + if axis is not None: + # Add back the shape elements that were ignored + new_axis = axis - np.arange(len(axis)) + new_shapes = [tuple(np.insert(new_shape, new_axis, removed_shape)) + for removed_shape in removed_shapes] + return new_shapes + else: + return tuple(new_shape) + + +def _broadcast_array_shapes_remove_axis(arrays, axis=None): + """ + Broadcast shapes of arrays, dropping specified axes + + Given a sequence of arrays `arrays` and an integer or tuple `axis`, find + the shape of the broadcast result after consuming/dropping `axis`. + In other words, return output shape of a typical hypothesis test on + `arrays` vectorized along `axis`. + + Examples + -------- + >>> import numpy as np + >>> from scipy.stats._axis_nan_policy import _broadcast_array_shapes_remove_axis + >>> a = np.zeros((5, 2, 1)) + >>> b = np.zeros((9, 3)) + >>> _broadcast_array_shapes_remove_axis((a, b), 1) + (5, 3) + """ + # Note that here, `axis=None` means do not consume/drop any axes - _not_ + # ravel arrays before broadcasting. + shapes = [arr.shape for arr in arrays] + return _broadcast_shapes_remove_axis(shapes, axis) + + +def _broadcast_shapes_remove_axis(shapes, axis=None): + """ + Broadcast shapes, dropping specified axes + + Same as _broadcast_array_shapes_remove_axis, but given a sequence + of array shapes `shapes` instead of the arrays themselves. + """ + shapes = _broadcast_shapes(shapes, axis) + shape = shapes[0] + if axis is not None: + shape = np.delete(shape, axis) + return tuple(shape) + + +@xp_capabilities() +def _broadcast_concatenate(arrays, axis, paired=False, xp=None): + """Concatenate arrays along an axis with broadcasting.""" + xp = array_namespace(*arrays) if xp is None else xp + arrays = _broadcast_arrays(arrays, axis if not paired else None, xp=xp) + res = xp.concat(arrays, axis=axis) + return res + + +def _remove_nans(samples, paired, xp=None): + "Remove nans from paired or unpaired 1D samples" + # potential optimization: don't copy arrays that don't contain nans + xp = array_namespace(*samples) + if not paired: + return [sample[~xp.isnan(sample)] for sample in samples] + + # for paired samples, we need to remove the whole pair when any part + # has a nan + nans = xp.isnan(samples[0]) + for sample in samples[1:]: + nans = nans | xp.isnan(sample) + not_nans = ~nans + return [sample[not_nans] for sample in samples] + + +def _remove_sentinel(samples, paired, sentinel): + "Remove sentinel values from paired or unpaired 1D samples" + # could consolidate with `_remove_nans`, but it's not quite as simple as + # passing `sentinel=np.nan` because `(np.nan == np.nan) is False` + + # potential optimization: don't copy arrays that don't contain sentinel + if not paired: + return [sample[sample != sentinel] for sample in samples] + + # for paired samples, we need to remove the whole pair when any part + # has a nan + sentinels = (samples[0] == sentinel) + for sample in samples[1:]: + sentinels = sentinels | (sample == sentinel) + not_sentinels = ~sentinels + return [sample[not_sentinels] for sample in samples] + + +def _masked_arrays_2_sentinel_arrays(samples): + # masked arrays in `samples` are converted to regular arrays, and values + # corresponding with masked elements are replaced with a sentinel value + + # return without modifying arrays if none have a mask + has_mask = False + for sample in samples: + mask = getattr(sample, 'mask', False) + has_mask = has_mask or np.any(mask) + if not has_mask: + return samples, None # None means there is no sentinel value + + # Choose a sentinel value. We can't use `np.nan`, because sentinel (masked) + # values are always omitted, but there are different nan policies. + dtype = np.result_type(*samples) + dtype = dtype if np.issubdtype(dtype, np.number) else np.float64 + for i in range(len(samples)): + # Things get more complicated if the arrays are of different types. + # We could have different sentinel values for each array, but + # the purpose of this code is convenience, not efficiency. + samples[i] = samples[i].astype(dtype, copy=False) + + inexact = np.issubdtype(dtype, np.inexact) + info = np.finfo if inexact else np.iinfo + max_possible, min_possible = info(dtype).max, info(dtype).min + nextafter = np.nextafter if inexact else (lambda x, _: x - 1) + + sentinel = max_possible + # For simplicity, min_possible/np.infs are not candidate sentinel values + while sentinel > min_possible: + for sample in samples: + if np.any(sample == sentinel): # choose a new sentinel value + sentinel = nextafter(sentinel, -np.inf) + break + else: # when sentinel value is OK, break the while loop + break + else: + message = ("This function replaces masked elements with sentinel " + "values, but the data contains all distinct values of this " + "data type. Consider promoting the dtype to `np.float64`.") + raise ValueError(message) + + # replace masked elements with sentinel value + out_samples = [] + for sample in samples: + mask = getattr(sample, 'mask', None) + if mask is not None: # turn all masked arrays into sentinel arrays + mask = np.broadcast_to(mask, sample.shape) + sample = sample.data.copy() if np.any(mask) else sample.data + sample = np.asarray(sample) # `sample.data` could be a memoryview? + sample[mask] = sentinel + out_samples.append(sample) + + return out_samples, sentinel + + +def _check_empty_inputs(samples, axis, xp=None): + """ + Check for empty sample; return appropriate output for a vectorized hypotest + """ + xp = array_namespace(*samples) if xp is None else xp + # if none of the samples are empty, we need to perform the test + if not any(xp_size(sample) == 0 for sample in samples): + return None + # otherwise, the statistic and p-value will be either empty arrays or + # arrays with NaNs. Produce the appropriate array and return it. + output_shape = _broadcast_array_shapes_remove_axis(samples, axis) + NaN = _get_nan(*samples) + output = xp.full(output_shape, xp.nan, dtype=NaN.dtype) + return output + + +def _add_reduced_axes(res, reduced_axes, keepdims, xp=np): + """ + Add reduced axes back to all the arrays in the result object + if keepdims = True. + """ + return ([xpx.expand_dims(output, axis=reduced_axes) + if not isinstance(output, int) else output for output in res] + if keepdims else res) + + +# Standard docstring / signature entries for `axis`, `nan_policy`, `keepdims` +_name = 'axis' +_desc = ( + """If an int, the axis of the input along which to compute the statistic. +The statistic of each axis-slice (e.g. row) of the input will appear in a +corresponding element of the output. +If ``None``, the input will be raveled before computing the statistic.""" + .split('\n')) + + +def _get_axis_params(default_axis=0, _name=_name, _desc=_desc): # bind NOW + _type = f"int or None, default: {default_axis}" + _axis_parameter_doc = Parameter(_name, _type, _desc) + _axis_parameter = inspect.Parameter(_name, + inspect.Parameter.KEYWORD_ONLY, + default=default_axis) + return _axis_parameter_doc, _axis_parameter + + +_name = 'nan_policy' +_type = "{'propagate', 'omit', 'raise'}" +_desc = ( + """Defines how to handle input NaNs. + +- ``propagate``: if a NaN is present in the axis slice (e.g. row) along + which the statistic is computed, the corresponding entry of the output + will be NaN. +- ``omit``: NaNs will be omitted when performing the calculation. + If insufficient data remains in the axis slice along which the + statistic is computed, the corresponding entry of the output will be + NaN. +- ``raise``: if a NaN is present, a ``ValueError`` will be raised.""" + .split('\n')) +_nan_policy_parameter_doc = Parameter(_name, _type, _desc) +_nan_policy_parameter = inspect.Parameter(_name, + inspect.Parameter.KEYWORD_ONLY, + default='propagate') + +_name = 'keepdims' +_type = "bool, default: False" +_desc = ( + """If this is set to True, the axes which are reduced are left +in the result as dimensions with size one. With this option, +the result will broadcast correctly against the input array.""" + .split('\n')) +_keepdims_parameter_doc = Parameter(_name, _type, _desc) +_keepdims_parameter = inspect.Parameter(_name, + inspect.Parameter.KEYWORD_ONLY, + default=False) + +_standard_note_addition = ( + """\nBeginning in SciPy 1.9, ``np.matrix`` inputs (not recommended for new +code) are converted to ``np.ndarray`` before the calculation is performed. In +this case, the output will be a scalar or ``np.ndarray`` of appropriate shape +rather than a 2D ``np.matrix``. Similarly, while masked elements of masked +arrays are ignored, the output will be a scalar or ``np.ndarray`` rather than a +masked array with ``mask=False``.""").split('\n') + + +def _axis_nan_policy_factory(tuple_to_result, default_axis=0, + n_samples=1, paired=False, + result_to_tuple=None, too_small=0, + n_outputs=2, kwd_samples=(), override=None): + """Factory for a wrapper that adds axis/nan_policy params to a function. + + Parameters + ---------- + tuple_to_result : callable + Callable that returns an object of the type returned by the function + being wrapped (e.g. the namedtuple or dataclass returned by a + statistical test) provided the separate components (e.g. statistic, + pvalue). + default_axis : int, default: 0 + The default value of the axis argument. Standard is 0 except when + backwards compatibility demands otherwise (e.g. `None`). + n_samples : int or callable, default: 1 + The number of data samples accepted by the function + (e.g. `mannwhitneyu`), a callable that accepts a dictionary of + parameters passed into the function and returns the number of data + samples (e.g. `wilcoxon`), or `None` to indicate an arbitrary number + of samples (e.g. `kruskal`). + paired : {False, True} + Whether the function being wrapped treats the samples as paired (i.e. + corresponding elements of each sample should be considered as different + components of the same sample.) + result_to_tuple : callable, optional + Function that unpacks the results of the function being wrapped into + a tuple. This is essentially the inverse of `tuple_to_result`. Default + is `None`, which is appropriate for statistical tests that return a + statistic, pvalue tuple (rather than, e.g., a non-iterable dataclass). + too_small : int or callable, default: 0 + The largest unacceptably small sample for the function being wrapped. + For example, some functions require samples of size two or more or they + raise an error. This argument prevents the error from being raised when + input is not 1D and instead places a NaN in the corresponding element + of the result. If callable, it must accept a list of samples, axis, + and a dictionary of keyword arguments passed to the wrapper function as + arguments and return a bool indicating whether the samples passed are + too small. + n_outputs : int or callable, default: 2 + The number of outputs produced by the function given 1d sample(s). For + example, hypothesis tests that return a namedtuple or result object + with attributes ``statistic`` and ``pvalue`` use the default + ``n_outputs=2``; summary statistics with scalar output use + ``n_outputs=1``. Alternatively, may be a callable that accepts a + dictionary of arguments passed into the wrapped function and returns + the number of outputs corresponding with those arguments. + kwd_samples : sequence, default: () + The names of keyword parameters that should be treated as samples. For + example, `gmean` accepts as its first argument a sample `a` but + also `weights` as a fourth, optional keyword argument. In this case, we + use `n_samples=1` and kwd_samples=['weights']. + override : dict, default: {'vectorization': False, 'nan_propagation': True} + Pass a dictionary with ``'vectorization': True`` to ensure that the + decorator overrides the function's behavior for multimensional input. + Use ``'nan_propagation': False`` to ensure that the decorator does not + override the function's behavior for ``nan_policy='propagate'``. + """ + # Specify which existing behaviors the decorator must override + temp = override or {} + override = {'vectorization': False, + 'nan_propagation': True} + override.update(temp) + + if result_to_tuple is None: + def result_to_tuple(res, _): + return res + + if not callable(too_small): + def is_too_small(samples, *ts_args, axis=-1, **ts_kwargs): + for sample in samples: + if sample.shape[axis] <= too_small: + return True + return False + else: + is_too_small = too_small + + def axis_nan_policy_decorator(hypotest_fun_in): + @wraps(hypotest_fun_in) + def axis_nan_policy_wrapper(*args, _no_deco=False, **kwds): + + if _no_deco: # for testing, decorator does nothing + return hypotest_fun_in(*args, **kwds) + + # For now, skip the decorator entirely if using array API. In the future, + # we'll probably want to use it for `keepdims`, `axis` tuples, etc. + if len(args) == 0: # extract sample from `kwds` if there are no `args` + used_kwd_samples = list(set(kwds).intersection(set(kwd_samples))) + temp = used_kwd_samples[:1] + else: + temp = args[0] + + if is_dask(array_namespace(temp)): + msg = ("Use of `nan_policy` and `keepdims` " + "is incompatible with lazy arrays.") + if 'nan_policy' in kwds or 'keepdims' in kwds: + raise NotImplementedError(msg) + return hypotest_fun_in(*args, **kwds) + + # We need to be flexible about whether position or keyword + # arguments are used, but we need to make sure users don't pass + # both for the same parameter. To complicate matters, some + # functions accept samples with *args, and some functions already + # accept `axis` and `nan_policy` as positional arguments. + # The strategy is to make sure that there is no duplication + # between `args` and `kwds`, combine the two into `kwds`, then + # the samples, `nan_policy`, and `axis` from `kwds`, as they are + # dealt with separately. + + # Check for intersection between positional and keyword args + params = list(inspect.signature(hypotest_fun_in).parameters) + if n_samples is None: + # Give unique names to each positional sample argument + # Note that *args can't be provided as a keyword argument + params = [f"arg{i}" for i in range(len(args))] + params[1:] + + # raise if there are too many positional args + maxarg = (np.inf if inspect.getfullargspec(hypotest_fun_in).varargs + else len(inspect.getfullargspec(hypotest_fun_in).args)) + if len(args) > maxarg: # let the function raise the right error + hypotest_fun_in(*args, **kwds) + + # raise if multiple values passed for same parameter + d_args = dict(zip(params, args)) + intersection = set(d_args) & set(kwds) + if intersection: # let the function raise the right error + hypotest_fun_in(*args, **kwds) + + # Consolidate other positional and keyword args into `kwds` + kwds.update(d_args) + + # rename avoids UnboundLocalError + if callable(n_samples): + # Future refactoring idea: no need for callable n_samples. + # Just replace `n_samples` and `kwd_samples` with a single + # list of the names of all samples, and treat all of them + # as `kwd_samples` are treated below. + n_samp = n_samples(kwds) + else: + n_samp = n_samples or len(args) + + # get the number of outputs + n_out = n_outputs # rename to avoid UnboundLocalError + if callable(n_out): + n_out = n_out(kwds) + + # If necessary, rearrange function signature: accept other samples + # as positional args right after the first n_samp args + kwd_samp = [name for name in kwd_samples + if kwds.get(name, None) is not None] + n_kwd_samp = len(kwd_samp) + if not kwd_samp: + hypotest_fun_out = hypotest_fun_in + else: + def hypotest_fun_out(*samples, **kwds): + new_kwds = dict(zip(kwd_samp, samples[n_samp:])) + kwds.update(new_kwds) + return hypotest_fun_in(*samples[:n_samp], **kwds) + + # Extract the things we need here + try: # if something is missing + samples = [kwds.pop(param) for param in (params[:n_samp] + kwd_samp)] + xp = array_namespace(*samples) + samples = xp_promote(*samples, xp=xp) + samples = (samples,) if not isinstance(samples, tuple) else samples + samples = [xpx.atleast_nd(sample, ndim=1) for sample in samples] + except KeyError: # let the function raise the right error + # might need to revisit this if required arg is not a "sample" + hypotest_fun_in(*args, **kwds) + vectorized = True if 'axis' in params else False + vectorized = vectorized and not override['vectorization'] + axis = kwds.pop('axis', default_axis) + nan_policy = kwds.pop('nan_policy', 'propagate') + keepdims = kwds.pop("keepdims", False) + del args # avoid the possibility of passing both `args` and `kwds` + + # convert masked arrays to regular arrays with sentinel values + sentinel = None + if is_numpy(xp): + samples, sentinel = _masked_arrays_2_sentinel_arrays(samples) + + # standardize to always work along last axis + reduced_axes = axis + if axis is None: + if samples: + # when axis=None, take the maximum of all dimensions since + # all the dimensions are reduced. + n_dims = max([xp.asarray(sample).ndim for sample in samples]) + reduced_axes = tuple(range(n_dims)) + samples = [xp_ravel(sample) for sample in samples] + else: + # don't ignore any axes when broadcasting if paired + samples = _broadcast_arrays(samples, axis=axis if not paired else None) + axis = (axis,) if np.isscalar(axis) else axis + n_axes = len(axis) + # move all axes in `axis` to the end to be raveled + samples = [xp.moveaxis(sample, axis, tuple(range(-len(axis), 0))) + for sample in samples] + shapes = [sample.shape for sample in samples] + # New shape is unchanged for all axes _not_ in `axis` + # At the end, we append the product of the shapes of the axes + # in `axis`. Appending -1 doesn't work for zero-size arrays! + new_shapes = [shape[:-n_axes] + (math.prod(shape[-n_axes:]),) + for shape in shapes] + samples = [xp.reshape(sample, new_shape) + for sample, new_shape in zip(samples, new_shapes)] + axis = -1 # work over the last axis + + NaN = _get_nan(*samples) if samples else xp.nan + + # if axis is not needed, just handle nan_policy and return + ndims = np.array([sample.ndim for sample in samples]) # NumPy OK for ndims + if np.all(ndims <= 1): + # Addresses nan_policy == "raise" + if nan_policy != 'propagate' or override['nan_propagation']: + contains_nan = [_contains_nan(sample, nan_policy) + for sample in samples] + else: + # Behave as though there are no NaNs (even if there are) + contains_nan = [False] * len(samples) + + # To give JAX most benefits of the `_axis_nan_policy` decorator without + # value-dependent branching, the decorator will always treat JAX arrays + # as if there are no NaNs. Typically, this means that NaNs will + # multidimensional naturally. However, there may be cases in which this + # doesn't produce the desired result, so for now, JAX users should + # treat all functions as though there is no explicit`nan_policy`. + # Future work tracked in gh-14651. + any_contains_nan = not is_jax(xp) and any(contains_nan) + # Addresses nan_policy == "propagate" + if any_contains_nan and (nan_policy == 'propagate' + and override['nan_propagation']): + res = xp.full(n_out, xp.nan, dtype=NaN.dtype) + res = _add_reduced_axes(res, reduced_axes, keepdims) + return tuple_to_result(*res) + + # Addresses nan_policy == "omit" + too_small_msg = too_small_1d_not_omit + if any_contains_nan and nan_policy == 'omit': + # consider passing in contains_nan + samples = _remove_nans(samples, paired) + too_small_msg = too_small_1d_omit + + if sentinel: + samples = _remove_sentinel(samples, paired, sentinel) + + if is_too_small(samples, kwds): + warnings.warn(too_small_msg, SmallSampleWarning, stacklevel=2) + res = xp.full(n_out, xp.nan, dtype=NaN.dtype) + res = _add_reduced_axes(res, reduced_axes, keepdims) + return tuple_to_result(*res) + + res = hypotest_fun_out(*samples, **kwds) + res = result_to_tuple(res, n_out) + res = _add_reduced_axes(res, reduced_axes, keepdims) + return tuple_to_result(*res) + + # check for empty input + empty_output = _check_empty_inputs(samples, axis, xp=xp) + # only return empty output if zero sized input is too small. + if ( + empty_output is not None + and (is_too_small(samples, kwds) or xp_size(empty_output) == 0) + ): + if is_too_small(samples, kwds) and xp_size(empty_output) != 0: + warnings.warn(too_small_nd_not_omit, SmallSampleWarning, + stacklevel=2) + res = [xp_copy(empty_output) for i in range(n_out)] + res = _add_reduced_axes(res, reduced_axes, keepdims) + return tuple_to_result(*res) + + if not is_numpy(xp) and 'nan_policy' in kwds: + msg = ("Use of `nan_policy` is incompatible with multidimensional " + "non-NumPy arrays.") + raise NotImplementedError(msg) + + if not is_numpy(xp): + res = hypotest_fun_out(*samples, axis=axis, **kwds) + res = result_to_tuple(res, n_out) + res = _add_reduced_axes(res, reduced_axes, keepdims, xp=xp) + return tuple_to_result(*res) + + # otherwise, concatenate all samples along axis, remembering where + # each separate sample begins + lengths = np.array([sample.shape[axis] for sample in samples]) + split_indices = np.cumsum(lengths) + x = _broadcast_concatenate(samples, axis, paired=paired) + + # Addresses nan_policy == "raise" + if nan_policy != 'propagate' or override['nan_propagation']: + contains_nan = _contains_nan(x, nan_policy) + else: + contains_nan = False # behave like there are no NaNs + + if vectorized and not contains_nan and not sentinel: + res = hypotest_fun_out(*samples, axis=axis, **kwds) + res = result_to_tuple(res, n_out) + res = _add_reduced_axes(res, reduced_axes, keepdims) + return tuple_to_result(*res) + + # Addresses nan_policy == "omit" + if contains_nan and nan_policy == 'omit': + def hypotest_fun(x): + samples = np.split(x, split_indices)[:n_samp+n_kwd_samp] + samples = _remove_nans(samples, paired) + if sentinel: + samples = _remove_sentinel(samples, paired, sentinel) + if is_too_small(samples, kwds): + warnings.warn(too_small_nd_omit, SmallSampleWarning, + stacklevel=4) + return np.full(n_out, NaN) + return result_to_tuple(hypotest_fun_out(*samples, **kwds), n_out) + + # Addresses nan_policy == "propagate" + elif (contains_nan and nan_policy == 'propagate' + and override['nan_propagation']): + def hypotest_fun(x): + if np.isnan(x).any(): + return np.full(n_out, NaN) + + samples = np.split(x, split_indices)[:n_samp+n_kwd_samp] + if sentinel: + samples = _remove_sentinel(samples, paired, sentinel) + if is_too_small(samples, kwds): + return np.full(n_out, NaN) + return result_to_tuple(hypotest_fun_out(*samples, **kwds), n_out) + + else: + def hypotest_fun(x): + samples = np.split(x, split_indices)[:n_samp+n_kwd_samp] + if sentinel: + samples = _remove_sentinel(samples, paired, sentinel) + if is_too_small(samples, kwds): + return np.full(n_out, NaN) + return result_to_tuple(hypotest_fun_out(*samples, **kwds), n_out) + + x = np.moveaxis(x, axis, 0) + res = np.apply_along_axis(hypotest_fun, axis=0, arr=x) + res = _add_reduced_axes(res, reduced_axes, keepdims) + return tuple_to_result(*res) + + _axis_parameter_doc, _axis_parameter = _get_axis_params(default_axis) + doc = FunctionDoc(axis_nan_policy_wrapper) + parameter_names = [param.name for param in doc['Parameters']] + if 'axis' in parameter_names: + doc['Parameters'][parameter_names.index('axis')] = ( + _axis_parameter_doc) + else: + doc['Parameters'].append(_axis_parameter_doc) + if 'nan_policy' in parameter_names: + doc['Parameters'][parameter_names.index('nan_policy')] = ( + _nan_policy_parameter_doc) + else: + doc['Parameters'].append(_nan_policy_parameter_doc) + if 'keepdims' in parameter_names: + doc['Parameters'][parameter_names.index('keepdims')] = ( + _keepdims_parameter_doc) + else: + doc['Parameters'].append(_keepdims_parameter_doc) + doc['Notes'] += _standard_note_addition + doc = str(doc).split("\n", 1)[1].lstrip(" \n") # remove signature + axis_nan_policy_wrapper.__doc__ = str(doc) + + sig = inspect.signature(axis_nan_policy_wrapper) + parameters = sig.parameters + parameter_list = list(parameters.values()) + if 'axis' not in parameters: + parameter_list.append(_axis_parameter) + if 'nan_policy' not in parameters: + parameter_list.append(_nan_policy_parameter) + if 'keepdims' not in parameters: + parameter_list.append(_keepdims_parameter) + sig = sig.replace(parameters=parameter_list) + axis_nan_policy_wrapper.__signature__ = sig + + return axis_nan_policy_wrapper + return axis_nan_policy_decorator diff --git a/v12x/pylib/scipy/stats/_binned_statistic.py b/v12x/pylib/scipy/stats/_binned_statistic.py new file mode 100644 index 0000000000000000000000000000000000000000..5a11a1caa4bc48d2e6a0a42793f79635a3e58d36 --- /dev/null +++ b/v12x/pylib/scipy/stats/_binned_statistic.py @@ -0,0 +1,806 @@ +import builtins +from warnings import catch_warnings, simplefilter +import numpy as np +from operator import index +from collections import namedtuple + +from scipy._lib._array_api import xp_capabilities + +__all__ = ['binned_statistic', + 'binned_statistic_2d', + 'binned_statistic_dd'] + + +BinnedStatisticResult = namedtuple('BinnedStatisticResult', + ('statistic', 'bin_edges', 'binnumber')) + + +@xp_capabilities(np_only=True) +def binned_statistic(x, values, statistic='mean', + bins=10, range=None): + """ + Compute a binned statistic for one or more sets of data. + + This is a generalization of a histogram function. A histogram divides + the space into bins, and returns the count of the number of points in + each bin. This function allows the computation of the sum, mean, median, + or other statistic of the values (or set of values) within each bin. + + Parameters + ---------- + x : (N,) array_like + A sequence of values to be binned. + values : (N,) array_like or list of (N,) array_like + The data on which the statistic will be computed. This must be + the same shape as `x`, or a set of sequences - each the same shape as + `x`. If `values` is a set of sequences, the statistic will be computed + on each independently. + statistic : str or callable, optional + The statistic to compute (default is 'mean'). + The following statistics are available: + + * 'mean' : compute the mean of values for points within each bin. + Empty bins will be represented by NaN. + * 'std' : compute the standard deviation within each bin. This + is implicitly calculated with ddof=0. + * 'median' : compute the median of values for points within each + bin. Empty bins will be represented by NaN. + * 'count' : compute the count of points within each bin. This is + identical to an unweighted histogram. `values` array is not + referenced. + * 'sum' : compute the sum of values for points within each bin. + This is identical to a weighted histogram. + * 'min' : compute the minimum of values for points within each bin. + Empty bins will be represented by NaN. + * 'max' : compute the maximum of values for point within each bin. + Empty bins will be represented by NaN. + * function : a user-defined function which takes a 1D array of + values, and outputs a single numerical statistic. This function + will be called on the values in each bin. Empty bins will be + represented by function([]), or NaN if this returns an error. + + bins : int or sequence of scalars, optional + If `bins` is an int, it defines the number of equal-width bins in the + given range (10 by default). If `bins` is a sequence, it defines the + bin edges, including the rightmost edge, allowing for non-uniform bin + widths. Values in `x` that are smaller than lowest bin edge are + assigned to bin number 0, values beyond the highest bin are assigned to + ``bins[-1]``. If the bin edges are specified, the number of bins will + be, (nx = len(bins)-1). + range : (float, float) or [(float, float)], optional + The lower and upper range of the bins. If not provided, range + is simply ``(x.min(), x.max())``. Values outside the range are + ignored. + + Returns + ------- + statistic : array + The values of the selected statistic in each bin. + bin_edges : array of dtype float + Return the bin edges ``(length(statistic)+1)``. + binnumber: 1-D ndarray of ints + Indices of the bins (corresponding to `bin_edges`) in which each value + of `x` belongs. Same length as `values`. A binnumber of `i` means the + corresponding value is between (bin_edges[i-1], bin_edges[i]). + + See Also + -------- + numpy.digitize, numpy.histogram, binned_statistic_2d, binned_statistic_dd + + Notes + ----- + All but the last (righthand-most) bin is half-open. In other words, if + `bins` is ``[1, 2, 3, 4]``, then the first bin is ``[1, 2)`` (including 1, + but excluding 2) and the second ``[2, 3)``. The last bin, however, is + ``[3, 4]``, which *includes* 4. + + .. versionadded:: 0.11.0 + + Examples + -------- + >>> import numpy as np + >>> from scipy import stats + >>> import matplotlib.pyplot as plt + + First some basic examples: + + Create two evenly spaced bins in the range of the given sample, and sum the + corresponding values in each of those bins: + + >>> values = [1.0, 1.0, 2.0, 1.5, 3.0] + >>> stats.binned_statistic([1, 1, 2, 5, 7], values, 'sum', bins=2) + BinnedStatisticResult(statistic=array([4. , 4.5]), + bin_edges=array([1., 4., 7.]), binnumber=array([1, 1, 1, 2, 2])) + + Multiple arrays of values can also be passed. The statistic is calculated + on each set independently: + + >>> values = [[1.0, 1.0, 2.0, 1.5, 3.0], [2.0, 2.0, 4.0, 3.0, 6.0]] + >>> stats.binned_statistic([1, 1, 2, 5, 7], values, 'sum', bins=2) + BinnedStatisticResult(statistic=array([[4. , 4.5], + [8. , 9. ]]), bin_edges=array([1., 4., 7.]), + binnumber=array([1, 1, 1, 2, 2])) + + >>> stats.binned_statistic([1, 2, 1, 2, 4], np.arange(5), statistic='mean', + ... bins=3) + BinnedStatisticResult(statistic=array([1., 2., 4.]), + bin_edges=array([1., 2., 3., 4.]), + binnumber=array([1, 2, 1, 2, 3])) + + As a second example, we now generate some random data of sailing boat speed + as a function of wind speed, and then determine how fast our boat is for + certain wind speeds: + + >>> rng = np.random.default_rng() + >>> windspeed = 8 * rng.random(500) + >>> boatspeed = .3 * windspeed**.5 + .2 * rng.random(500) + >>> bin_means, bin_edges, binnumber = stats.binned_statistic(windspeed, + ... boatspeed, statistic='median', bins=[1,2,3,4,5,6,7]) + >>> plt.figure() + >>> plt.plot(windspeed, boatspeed, 'b.', label='raw data') + >>> plt.hlines(bin_means, bin_edges[:-1], bin_edges[1:], colors='g', lw=5, + ... label='binned statistic of data') + >>> plt.legend() + + Now we can use ``binnumber`` to select all datapoints with a windspeed + below 1: + + >>> low_boatspeed = boatspeed[binnumber == 0] + + As a final example, we will use ``bin_edges`` and ``binnumber`` to make a + plot of a distribution that shows the mean and distribution around that + mean per bin, on top of a regular histogram and the probability + distribution function: + + >>> x = np.linspace(0, 5, num=500) + >>> x_pdf = stats.maxwell.pdf(x) + >>> samples = stats.maxwell.rvs(size=10000) + + >>> bin_means, bin_edges, binnumber = stats.binned_statistic(x, x_pdf, + ... statistic='mean', bins=25) + >>> bin_width = (bin_edges[1] - bin_edges[0]) + >>> bin_centers = bin_edges[1:] - bin_width/2 + + >>> plt.figure() + >>> plt.hist(samples, bins=50, density=True, histtype='stepfilled', + ... alpha=0.2, label='histogram of data') + >>> plt.plot(x, x_pdf, 'r-', label='analytical pdf') + >>> plt.hlines(bin_means, bin_edges[:-1], bin_edges[1:], colors='g', lw=2, + ... label='binned statistic of data') + >>> plt.plot((binnumber - 0.5) * bin_width, x_pdf, 'g.', alpha=0.5) + >>> plt.legend(fontsize=10) + >>> plt.show() + + """ + try: + N = len(bins) + except TypeError: + N = 1 + + if N != 1: + bins = [np.asarray(bins, float)] + + if range is not None: + if len(range) == 2: + range = [range] + + medians, edges, binnumbers = binned_statistic_dd( + [x], values, statistic, bins, range) + + return BinnedStatisticResult(medians, edges[0], binnumbers) + + +BinnedStatistic2dResult = namedtuple('BinnedStatistic2dResult', + ('statistic', 'x_edge', 'y_edge', + 'binnumber')) + + +@xp_capabilities(np_only=True) +def binned_statistic_2d(x, y, values, statistic='mean', + bins=10, range=None, expand_binnumbers=False): + """ + Compute a bidimensional binned statistic for one or more sets of data. + + This is a generalization of a histogram2d function. A histogram divides + the space into bins, and returns the count of the number of points in + each bin. This function allows the computation of the sum, mean, median, + or other statistic of the values (or set of values) within each bin. + + Parameters + ---------- + x : (N,) array_like + A sequence of values to be binned along the first dimension. + y : (N,) array_like + A sequence of values to be binned along the second dimension. + values : (N,) array_like or list of (N,) array_like + The data on which the statistic will be computed. This must be + the same shape as `x`, or a list of sequences - each with the same + shape as `x`. If `values` is such a list, the statistic will be + computed on each independently. + statistic : str or callable, optional + The statistic to compute (default is 'mean'). + The following statistics are available: + + * 'mean' : compute the mean of values for points within each bin. + Empty bins will be represented by NaN. + * 'std' : compute the standard deviation within each bin. This + is implicitly calculated with ddof=0. + * 'median' : compute the median of values for points within each + bin. Empty bins will be represented by NaN. + * 'count' : compute the count of points within each bin. This is + identical to an unweighted histogram. `values` array is not + referenced. + * 'sum' : compute the sum of values for points within each bin. + This is identical to a weighted histogram. + * 'min' : compute the minimum of values for points within each bin. + Empty bins will be represented by NaN. + * 'max' : compute the maximum of values for point within each bin. + Empty bins will be represented by NaN. + * function : a user-defined function which takes a 1D array of + values, and outputs a single numerical statistic. This function + will be called on the values in each bin. Empty bins will be + represented by function([]), or NaN if this returns an error. + + bins : int or [int, int] or array_like or [array, array], optional + The bin specification: + + * the number of bins for the two dimensions (nx = ny = bins), + * the number of bins in each dimension (nx, ny = bins), + * the bin edges for the two dimensions (x_edge = y_edge = bins), + * the bin edges in each dimension (x_edge, y_edge = bins). + + If the bin edges are specified, the number of bins will be, + (nx = len(x_edge)-1, ny = len(y_edge)-1). + + range : (2,2) array_like, optional + The leftmost and rightmost edges of the bins along each dimension + (if not specified explicitly in the `bins` parameters): + [[xmin, xmax], [ymin, ymax]]. All values outside of this range will be + considered outliers and not tallied in the histogram. + expand_binnumbers : bool, optional + 'False' (default): the returned `binnumber` is a shape (N,) array of + linearized bin indices. + 'True': the returned `binnumber` is 'unraveled' into a shape (2,N) + ndarray, where each row gives the bin numbers in the corresponding + dimension. + See the `binnumber` returned value, and the `Examples` section. + + .. versionadded:: 0.17.0 + + Returns + ------- + statistic : (nx, ny) ndarray + The values of the selected statistic in each two-dimensional bin. + x_edge : (nx + 1) ndarray + The bin edges along the first dimension. + y_edge : (ny + 1) ndarray + The bin edges along the second dimension. + binnumber : (N,) array of ints or (2,N) ndarray of ints + This assigns to each element of `sample` an integer that represents the + bin in which this observation falls. The representation depends on the + `expand_binnumbers` argument. See `Notes` for details. + + + See Also + -------- + numpy.digitize, numpy.histogram2d, binned_statistic, binned_statistic_dd + + Notes + ----- + Binedges: + All but the last (righthand-most) bin is half-open. In other words, if + `bins` is ``[1, 2, 3, 4]``, then the first bin is ``[1, 2)`` (including 1, + but excluding 2) and the second ``[2, 3)``. The last bin, however, is + ``[3, 4]``, which *includes* 4. + + `binnumber`: + This returned argument assigns to each element of `sample` an integer that + represents the bin in which it belongs. The representation depends on the + `expand_binnumbers` argument. If 'False' (default): The returned + `binnumber` is a shape (N,) array of linearized indices mapping each + element of `sample` to its corresponding bin (using row-major ordering). + Note that the returned linearized bin indices are used for an array with + extra bins on the outer binedges to capture values outside of the defined + bin bounds. + If 'True': The returned `binnumber` is a shape (2,N) ndarray where + each row indicates bin placements for each dimension respectively. In each + dimension, a binnumber of `i` means the corresponding value is between + (D_edge[i-1], D_edge[i]), where 'D' is either 'x' or 'y'. + + .. versionadded:: 0.11.0 + + Examples + -------- + >>> from scipy import stats + + Calculate the counts with explicit bin-edges: + + >>> x = [0.1, 0.1, 0.1, 0.6] + >>> y = [2.1, 2.6, 2.1, 2.1] + >>> binx = [0.0, 0.5, 1.0] + >>> biny = [2.0, 2.5, 3.0] + >>> ret = stats.binned_statistic_2d(x, y, None, 'count', bins=[binx, biny]) + >>> ret.statistic + array([[2., 1.], + [1., 0.]]) + + The bin in which each sample is placed is given by the `binnumber` + returned parameter. By default, these are the linearized bin indices: + + >>> ret.binnumber + array([5, 6, 5, 9]) + + The bin indices can also be expanded into separate entries for each + dimension using the `expand_binnumbers` parameter: + + >>> ret = stats.binned_statistic_2d(x, y, None, 'count', bins=[binx, biny], + ... expand_binnumbers=True) + >>> ret.binnumber + array([[1, 1, 1, 2], + [1, 2, 1, 1]]) + + Which shows that the first three elements belong in the xbin 1, and the + fourth into xbin 2; and so on for y. + + """ + + # This code is based on np.histogram2d + try: + N = len(bins) + except TypeError: + N = 1 + + if N != 1 and N != 2: + xedges = yedges = np.asarray(bins, float) + bins = [xedges, yedges] + + medians, edges, binnumbers = binned_statistic_dd( + [x, y], values, statistic, bins, range, + expand_binnumbers=expand_binnumbers) + + return BinnedStatistic2dResult(medians, edges[0], edges[1], binnumbers) + + +BinnedStatisticddResult = namedtuple('BinnedStatisticddResult', + ('statistic', 'bin_edges', + 'binnumber')) + + +def _bincount(x, weights): + if np.iscomplexobj(weights): + a = np.bincount(x, np.real(weights)) + b = np.bincount(x, np.imag(weights)) + z = a + b*1j + + else: + z = np.bincount(x, weights) + return z + + +@xp_capabilities(np_only=True) +def binned_statistic_dd(sample, values, statistic='mean', + bins=10, range=None, expand_binnumbers=False, + binned_statistic_result=None): + """ + Compute a multidimensional binned statistic for a set of data. + + This is a generalization of a histogramdd function. A histogram divides + the space into bins, and returns the count of the number of points in + each bin. This function allows the computation of the sum, mean, median, + or other statistic of the values within each bin. + + Parameters + ---------- + sample : array_like + Data to histogram passed as a sequence of N arrays of length D, or + as an (N,D) array. + values : (N,) array_like or list of (N,) array_like + The data on which the statistic will be computed. This must be + the same shape as `sample`, or a list of sequences - each with the + same shape as `sample`. If `values` is such a list, the statistic + will be computed on each independently. + statistic : str or callable, optional + The statistic to compute (default is 'mean'). + The following statistics are available: + + * 'mean' : compute the mean of values for points within each bin. + Empty bins will be represented by NaN. + * 'median' : compute the median of values for points within each + bin. Empty bins will be represented by NaN. + * 'count' : compute the count of points within each bin. This is + identical to an unweighted histogram. `values` array is not + referenced. + * 'sum' : compute the sum of values for points within each bin. + This is identical to a weighted histogram. + * 'std' : compute the standard deviation within each bin. This + is implicitly calculated with ddof=0. If the number of values + within a given bin is 0 or 1, the computed standard deviation value + will be 0 for the bin. + * 'min' : compute the minimum of values for points within each bin. + Empty bins will be represented by NaN. + * 'max' : compute the maximum of values for point within each bin. + Empty bins will be represented by NaN. + * function : a user-defined function which takes a 1D array of + values, and outputs a single numerical statistic. This function + will be called on the values in each bin. Empty bins will be + represented by function([]), or NaN if this returns an error. + + bins : sequence or positive int, optional + The bin specification must be in one of the following forms: + + * A sequence of arrays describing the bin edges along each dimension. + * The number of bins for each dimension (nx, ny, ... = bins). + * The number of bins for all dimensions (nx = ny = ... = bins). + + range : sequence, optional + A sequence of lower and upper bin edges to be used if the edges are + not given explicitly in `bins`. Defaults to the minimum and maximum + values along each dimension. + expand_binnumbers : bool, optional + 'False' (default): the returned `binnumber` is a shape (N,) array of + linearized bin indices. + 'True': the returned `binnumber` is 'unraveled' into a shape (D,N) + ndarray, where each row gives the bin numbers in the corresponding + dimension. + See the `binnumber` returned value, and the `Examples` section of + `binned_statistic_2d`. + binned_statistic_result : binnedStatisticddResult + Result of a previous call to the function in order to reuse bin edges + and bin numbers with new values and/or a different statistic. + To reuse bin numbers, `expand_binnumbers` must have been set to False + (the default) + + .. versionadded:: 0.17.0 + + Returns + ------- + statistic : ndarray, shape(nx1, nx2, nx3,...) + The values of the selected statistic in each two-dimensional bin. + bin_edges : list of ndarrays + A list of D arrays describing the (nxi + 1) bin edges for each + dimension. + binnumber : (N,) array of ints or (D,N) ndarray of ints + This assigns to each element of `sample` an integer that represents the + bin in which this observation falls. The representation depends on the + `expand_binnumbers` argument. See `Notes` for details. + + + See Also + -------- + numpy.digitize, numpy.histogramdd, binned_statistic, binned_statistic_2d + + Notes + ----- + Binedges: + All but the last (righthand-most) bin is half-open in each dimension. In + other words, if `bins` is ``[1, 2, 3, 4]``, then the first bin is + ``[1, 2)`` (including 1, but excluding 2) and the second ``[2, 3)``. The + last bin, however, is ``[3, 4]``, which *includes* 4. + + `binnumber`: + This returned argument assigns to each element of `sample` an integer that + represents the bin in which it belongs. The representation depends on the + `expand_binnumbers` argument. If 'False' (default): The returned + `binnumber` is a shape (N,) array of linearized indices mapping each + element of `sample` to its corresponding bin (using row-major ordering). + If 'True': The returned `binnumber` is a shape (D,N) ndarray where + each row indicates bin placements for each dimension respectively. In each + dimension, a binnumber of `i` means the corresponding value is between + (bin_edges[D][i-1], bin_edges[D][i]), for each dimension 'D'. + + .. versionadded:: 0.11.0 + + Examples + -------- + >>> import numpy as np + >>> from scipy import stats + >>> import matplotlib.pyplot as plt + >>> from mpl_toolkits.mplot3d import Axes3D + + Take an array of 600 (x, y) coordinates as an example. + `binned_statistic_dd` can handle arrays of higher dimension `D`. But a plot + of dimension `D+1` is required. + + >>> mu = np.array([0., 1.]) + >>> sigma = np.array([[1., -0.5],[-0.5, 1.5]]) + >>> multinormal = stats.multivariate_normal(mu, sigma) + >>> data = multinormal.rvs(size=600, random_state=235412) + >>> data.shape + (600, 2) + + Create bins and count how many arrays fall in each bin: + + >>> N = 60 + >>> x = np.linspace(-3, 3, N) + >>> y = np.linspace(-3, 4, N) + >>> ret = stats.binned_statistic_dd(data, np.arange(600), bins=[x, y], + ... statistic='count') + >>> bincounts = ret.statistic + + Set the volume and the location of bars: + + >>> dx = x[1] - x[0] + >>> dy = y[1] - y[0] + >>> x, y = np.meshgrid(x[:-1]+dx/2, y[:-1]+dy/2) + >>> z = 0 + + >>> bincounts = bincounts.ravel() + >>> x = x.ravel() + >>> y = y.ravel() + + >>> fig = plt.figure() + >>> ax = fig.add_subplot(111, projection='3d') + >>> with np.errstate(divide='ignore'): # silence random axes3d warning + ... ax.bar3d(x, y, z, dx, dy, bincounts) + + Reuse bin numbers and bin edges with new values: + + >>> ret2 = stats.binned_statistic_dd(data, -np.arange(600), + ... binned_statistic_result=ret, + ... statistic='mean') + """ + known_stats = ['mean', 'median', 'count', 'sum', 'std', 'min', 'max'] + if not callable(statistic) and statistic not in known_stats: + raise ValueError(f'invalid statistic {statistic!r}') + + try: + bins = index(bins) + except TypeError: + # bins is not an integer + pass + # If bins was an integer-like object, now it is an actual Python int. + + # For context: see gh-23751 and linked issues/PRs + if np.isnan(sample).any(): + message = ("The sample on which `values` is to be binned contains at least one " + "NaN, and the meaning is ambiguous. Consider removing or replacing " + "with numerical outliers.") + raise ValueError(message) + + # `Ndim` is the number of dimensions (e.g. `2` for `binned_statistic_2d`) + # `Dlen` is the length of elements along each dimension. + # This code is based on np.histogramdd + try: + # `sample` is an ND-array. + Dlen, Ndim = sample.shape + except (AttributeError, ValueError): + # `sample` is a sequence of 1D arrays. + sample = np.atleast_2d(sample).T + Dlen, Ndim = sample.shape + + # Store initial shape of `values` to preserve it in the output + values = np.asarray(values) + input_shape = list(values.shape) + # Make sure that `values` is 2D to iterate over rows + values = np.atleast_2d(values) + Vdim, Vlen = values.shape + + # Make sure `values` match `sample` + if statistic != 'count' and Vlen != Dlen: + raise AttributeError('The number of `values` elements must match the ' + 'length of each `sample` dimension.') + + try: + M = len(bins) + if M != Ndim: + raise AttributeError('The dimension of bins must be equal ' + 'to the dimension of the sample x.') + except TypeError: + bins = Ndim * [bins] + + if binned_statistic_result is None: + nbin, edges, dedges = _bin_edges(sample, bins, range) + binnumbers = _bin_numbers(sample, nbin, edges, dedges) + else: + edges = binned_statistic_result.bin_edges + nbin = np.array([len(edges[i]) + 1 for i in builtins.range(Ndim)]) + # +1 for outlier bins + dedges = [np.diff(edges[i]) for i in builtins.range(Ndim)] + binnumbers = binned_statistic_result.binnumber + + # Avoid overflow with double precision. Complex `values` -> `complex128`. + result_type = np.result_type(values, np.float64) + result = np.empty([Vdim, nbin.prod()], dtype=result_type) + + if statistic in {'mean', np.mean}: + result.fill(np.nan) + flatcount = _bincount(binnumbers, None) + a = flatcount.nonzero() + for vv in builtins.range(Vdim): + flatsum = _bincount(binnumbers, values[vv]) + result[vv, a] = flatsum[a] / flatcount[a] + elif statistic in {'std', np.std}: + result.fill(np.nan) + flatcount = _bincount(binnumbers, None) + a = flatcount.nonzero() + for vv in builtins.range(Vdim): + flatsum = _bincount(binnumbers, values[vv]) + delta = values[vv] - flatsum[binnumbers] / flatcount[binnumbers] + std = np.sqrt( + _bincount(binnumbers, delta*np.conj(delta))[a] / flatcount[a] + ) + result[vv, a] = std + result = np.real(result) + elif statistic == 'count': + result = np.empty([Vdim, nbin.prod()], dtype=np.float64) + result.fill(0) + flatcount = _bincount(binnumbers, None) + a = np.arange(len(flatcount)) + result[:, a] = flatcount[np.newaxis, :] + elif statistic in {'sum', np.sum}: + result.fill(0) + for vv in builtins.range(Vdim): + flatsum = _bincount(binnumbers, values[vv]) + a = np.arange(len(flatsum)) + result[vv, a] = flatsum + elif statistic in {'median', np.median}: + result.fill(np.nan) + for vv in builtins.range(Vdim): + i = np.lexsort((values[vv], binnumbers)) + _, j, counts = np.unique(binnumbers[i], + return_index=True, return_counts=True) + mid = j + (counts - 1) / 2 + mid_a = values[vv, i][np.floor(mid).astype(int)] + mid_b = values[vv, i][np.ceil(mid).astype(int)] + medians = (mid_a + mid_b) / 2 + result[vv, binnumbers[i][j]] = medians + elif statistic in {'min', np.min}: + result.fill(np.nan) + for vv in builtins.range(Vdim): + i = np.argsort(values[vv])[::-1] # Reversed so the min is last + result[vv, binnumbers[i]] = values[vv, i] + elif statistic in {'max', np.max}: + result.fill(np.nan) + for vv in builtins.range(Vdim): + i = np.argsort(values[vv]) + result[vv, binnumbers[i]] = values[vv, i] + elif callable(statistic): + with np.errstate(invalid='ignore'), catch_warnings(): + simplefilter("ignore", RuntimeWarning) + try: + null = statistic([]) + except Exception: + null = np.nan + if np.iscomplexobj(null): + result = result.astype(np.complex128) + result.fill(null) + try: + _calc_binned_statistic( + Vdim, binnumbers, result, values, statistic + ) + except ValueError: + result = result.astype(np.complex128) + _calc_binned_statistic( + Vdim, binnumbers, result, values, statistic + ) + + # Shape into a proper matrix + result = result.reshape(np.append(Vdim, nbin)) + + # Remove outliers (indices 0 and -1 for each bin-dimension). + core = tuple([slice(None)] + Ndim * [slice(1, -1)]) + result = result[core] + + # Unravel binnumbers into an ndarray, each row the bins for each dimension + if expand_binnumbers and Ndim > 1: + binnumbers = np.asarray(np.unravel_index(binnumbers, nbin)) + + if np.any(result.shape[1:] != nbin - 2): + raise RuntimeError('Internal Shape Error') + + # Reshape to have output (`result`) match input (`values`) shape + result = result.reshape(input_shape[:-1] + list(nbin-2)) + + return BinnedStatisticddResult(result, edges, binnumbers) + + +def _calc_binned_statistic(Vdim, bin_numbers, result, values, stat_func): + unique_bin_numbers = np.unique(bin_numbers) + for vv in builtins.range(Vdim): + bin_map = _create_binned_data(bin_numbers, unique_bin_numbers, + values, vv) + for i in unique_bin_numbers: + stat = stat_func(np.array(bin_map[i])) + if np.iscomplexobj(stat) and not np.iscomplexobj(result): + raise ValueError("The statistic function returns complex ") + result[vv, i] = stat + + +def _create_binned_data(bin_numbers, unique_bin_numbers, values, vv): + """ Create hashmap of bin ids to values in bins + key: bin number + value: list of binned data + """ + bin_map = dict() + for i in unique_bin_numbers: + bin_map[i] = [] + for i in builtins.range(len(bin_numbers)): + bin_map[bin_numbers[i]].append(values[vv, i]) + return bin_map + + +def _bin_edges(sample, bins=None, range=None): + """ Create edge arrays + """ + Dlen, Ndim = sample.shape + + nbin = np.empty(Ndim, int) # Number of bins in each dimension + edges = Ndim * [None] # Bin edges for each dim (will be 2D array) + dedges = Ndim * [None] # Spacing between edges (will be 2D array) + + # Preserve sample floating point precision in bin edges + edges_dtype = (sample.dtype if np.issubdtype(sample.dtype, np.floating) + else float) + + # Select range for each dimension + # Used only if number of bins is given. + if range is None: + sample = sample.astype(float, copy=True) + sample[np.isinf(sample)] = np.nan + smin = np.atleast_1d(np.nanmin(sample, axis=0)) + smax = np.atleast_1d(np.nanmax(sample, axis=0)) + else: + if len(range) != Ndim: + raise ValueError( + f"range given for {len(range)} dimensions; {Ndim} required") + smin = np.empty(Ndim) + smax = np.empty(Ndim) + for i in builtins.range(Ndim): + if range[i][1] < range[i][0]: + raise ValueError( + f"In {f'dimension {i + 1} of ' if Ndim > 1 else ''}range," + " start must be <= stop") + smin[i], smax[i] = range[i] + + # Make sure the bins have a finite width. + for i in builtins.range(len(smin)): + if smin[i] == smax[i]: + smin[i] = smin[i] - .5 + smax[i] = smax[i] + .5 + + # Create edge arrays + for i in builtins.range(Ndim): + if np.isscalar(bins[i]): + nbin[i] = bins[i] + 2 # +2 for outlier bins + edges[i] = np.linspace(smin[i], smax[i], nbin[i] - 1, + dtype=edges_dtype) + else: + edges[i] = np.asarray(bins[i], edges_dtype) + nbin[i] = len(edges[i]) + 1 # +1 for outlier bins + dedges[i] = np.diff(edges[i]) + + nbin = np.asarray(nbin) + + return nbin, edges, dedges + + +def _bin_numbers(sample, nbin, edges, dedges): + """Compute the bin number each sample falls into, in each dimension + """ + Dlen, Ndim = sample.shape + + sampBin = [ + np.digitize(sample[:, i], edges[i]) + for i in range(Ndim) + ] + + # Using `digitize`, values that fall on an edge are put in the right bin. + # For the rightmost bin, we want values equal to the right + # edge to be counted in the last bin, and not as an outlier. + for i in range(Ndim): + # Find the rounding precision + dedges_min = dedges[i].min() + if dedges_min == 0: + raise ValueError('The smallest edge difference is numerically 0.') + decimal = int(-np.log10(dedges_min)) + 6 + # Find which points are on the rightmost edge. + on_edge = np.where((sample[:, i] >= edges[i][-1]) & + (np.around(sample[:, i], decimal) == + np.around(edges[i][-1], decimal)))[0] + # Shift these points one bin to the left. + sampBin[i][on_edge] -= 1 + + # Compute the sample indices in the flattened statistic matrix. + binnumbers = np.ravel_multi_index(sampBin, nbin) + + return binnumbers diff --git a/v12x/pylib/scipy/stats/_binomtest.py b/v12x/pylib/scipy/stats/_binomtest.py new file mode 100644 index 0000000000000000000000000000000000000000..746be06a07c949902d9173bc8fe8eb6cb1cbeece --- /dev/null +++ b/v12x/pylib/scipy/stats/_binomtest.py @@ -0,0 +1,417 @@ +import numpy as np +import scipy._external.array_api_extra as xpx +from scipy._lib._array_api import (xp_capabilities, array_namespace, xp_promote, is_jax, + is_lazy_array) +from scipy.optimize.elementwise import find_root +from scipy.special import ndtri +from scipy.special import _ufuncs as scu +from ._common import ConfidenceInterval + + +class BinomTestResult: + """ + Result of `scipy.stats.binomtest`. + + Attributes + ---------- + k : int + The number of successes (copied from `binomtest` input). + n : int + The number of trials (copied from `binomtest` input). + alternative : str + Indicates the alternative hypothesis specified in the input + to `binomtest`. It will be one of ``'two-sided'``, ``'greater'``, + or ``'less'``. + statistic: float + The estimate of the proportion of successes. + pvalue : float + The p-value of the hypothesis test. + + """ + def __init__(self, k, n, alternative, statistic, pvalue, xp): + lazy = is_lazy_array(k) + self.k = int(k) if not lazy and statistic.ndim == 0 and not xp.isnan(k) else k + self.n = int(n) if not lazy and statistic.ndim == 0 and not xp.isnan(n) else n + self.alternative = alternative + self.statistic = statistic + self.pvalue = pvalue + self._xp = xp + + # add alias for backward compatibility + self.proportion_estimate = statistic + + def __repr__(self): + s = ("BinomTestResult(" + f"k={self.k}, " + f"n={self.n}, " + f"alternative={self.alternative!r}, " + f"statistic={self.statistic}, " + f"pvalue={self.pvalue})") + return s + + def proportion_ci(self, confidence_level=0.95, method='exact'): + """ + Compute the confidence interval for ``statistic``. + + Parameters + ---------- + confidence_level : float, optional + Confidence level for the computed confidence interval + of the estimated proportion. Default is 0.95. + method : {'exact', 'wilson', 'wilsoncc'}, optional + Selects the method used to compute the confidence interval + for the estimate of the proportion: + + 'exact' : + Use the Clopper-Pearson exact method [1]_. + 'wilson' : + Wilson's method, without continuity correction ([2]_, [3]_). + 'wilsoncc' : + Wilson's method, with continuity correction ([2]_, [3]_). + + Default is ``'exact'``. + + Returns + ------- + ci : ``ConfidenceInterval`` object + The object has attributes ``low`` and ``high`` that hold the + lower and upper bounds of the confidence interval. + + References + ---------- + .. [1] C. J. Clopper and E. S. Pearson, The use of confidence or + fiducial limits illustrated in the case of the binomial, + Biometrika, Vol. 26, No. 4, pp 404-413 (Dec. 1934). + .. [2] E. B. Wilson, Probable inference, the law of succession, and + statistical inference, J. Amer. Stat. Assoc., 22, pp 209-212 + (1927). + .. [3] Robert G. Newcombe, Two-sided confidence intervals for the + single proportion: comparison of seven methods, Statistics + in Medicine, 17, pp 857-872 (1998). + + Examples + -------- + >>> from scipy.stats import binomtest + >>> result = binomtest(k=7, n=50, p=0.1) + >>> result.statistic + 0.14 + >>> result.proportion_ci() + ConfidenceInterval(low=0.05819170033997342, high=0.26739600249700846) + """ + if method not in ('exact', 'wilson', 'wilsoncc'): + raise ValueError(f"method ('{method}') must be one of 'exact', " + "'wilson' or 'wilsoncc'.") + if not (0 <= confidence_level <= 1): + raise ValueError(f'confidence_level ({confidence_level}) must be in ' + 'the interval [0, 1].') + + xp = self._xp + k, n, confidence_level = xp_promote( + self.k, self.n, confidence_level, + xp=xp, force_floating=True, + ) + if method == 'exact': + low, high = _binom_exact_conf_int( + k, n, confidence_level, self.alternative, xp=xp) + else: + # method is 'wilson' or 'wilsoncc' + correction = method == 'wilsoncc' + low, high = _binom_wilson_conf_int( + k, n, confidence_level, self.alternative, correction, xp=xp) + low = low[()] if low.ndim == 0 else low + high = high[()] if high.ndim == 0 else high + return ConfidenceInterval(low=low, high=high) + + +def _binom_exact_conf_int(k, n, confidence_level, alternative, *, xp): + """ + Compute the estimate and confidence interval for the binomial test. + + Returns proportion, prop_low, prop_high + """ + init = (xp.zeros_like(k), xp.ones_like(k)) + args = (k, n) + alpha = ((1 - confidence_level) / 2 if alternative == 'two-sided' + else 1 - confidence_level) + + # I think using the private methods here is fine, since we will only evaluate with + # valid `p`, `k`, and `n` (or all NaNs). One exception is when `k=0` and + # `binom._sf(k-1, n, p)`: evaluates to NaN, but that's not a problem because + # `plow` has a special case for `k=0` below. + plow = (xp.zeros_like(k) if alternative == 'less' else + find_root(lambda p, k, n: _SimpleBinomial(n, p).sf(k-1) - alpha, + init, args=args).x) + phigh = (xp.ones_like(k) if alternative == 'greater' else + find_root(lambda p, k, n: _SimpleBinomial(n, p).cdf(k) - alpha, + init, args=args).x) + + plow = xp.where(k == 0, 0.0, plow) + phigh = xp.where(k == n, 1.0, phigh) + return plow, phigh + + +def _binom_wilson_conf_int(k, n, confidence_level, alternative, correction, *, xp): + # This function assumes that the arguments have already been validated. + # In particular, `alternative` must be one of 'two-sided', 'less' or + # 'greater'. + p = k / n + if alternative == 'two-sided': + z = ndtri(0.5 + 0.5*confidence_level) + else: + z = ndtri(confidence_level) + + # For reference, the formulas implemented here are from + # Newcombe (1998) (ref. [3] in the proportion_ci docstring). + denom = 2*(n + z**2) + center = (2*n*p + z**2)/denom + q = 1 - p + + if correction: + with np.errstate(divide='ignore', invalid='ignore'): + dlo = (1 + z*xp.sqrt(z**2 - 2 - 1/n + 4*p*(n*q + 1)))/denom + dhi = (1 + z*xp.sqrt(z**2 + 2 - 1/n + 4*p*(n*q - 1)))/denom + else: + delta = z / denom * xp.sqrt(4*n*p*q + z**2) + dlo, dhi = delta, delta + + lo = xp.where((k == 0) | (alternative == 'less'), 0.0, center - dlo) + hi = xp.where((k == n) | (alternative == 'greater'), 1.0, center + dhi) + return lo, hi + + +@xp_capabilities( + skip_backends=[('dask.array', "")], cpu_only=True, + reason="binomial distribution ufuncs only available for NumPy", + extra_note="`alternative='two-sided'` is incompatible with JAX arrays.", +) +def binomtest(k, n, p=0.5, alternative='two-sided'): + """ + Perform a test that the probability of success is p. + + The binomial test [1]_ is a test of the null hypothesis that the + probability of success in a Bernoulli experiment is `p`. + + Details of the test can be found in many texts on statistics, such + as section 24.5 of [2]_. + + The documentation is written as though the function accepts and returns Python + scalars, but the function is vectorized to work elementwise with NumPy arrays. + + Parameters + ---------- + k : int + The number of successes. + n : int + The number of trials. + p : float, optional + The hypothesized probability of success, i.e. the expected + proportion of successes. The value must be in the interval + ``0 <= p <= 1``. The default value is ``p = 0.5``. + alternative : {'two-sided', 'greater', 'less'}, optional + Indicates the alternative hypothesis. The default value is + 'two-sided'. + + Returns + ------- + result : `~scipy.stats._result_classes.BinomTestResult` instance + The return value is an object with the following attributes: + + k : int + The number of successes (copied from `binomtest` input). + n : int + The number of trials (copied from `binomtest` input). + alternative : str + Indicates the alternative hypothesis specified in the input + to `binomtest`. It will be one of ``'two-sided'``, ``'greater'``, + or ``'less'``. + statistic : float + The estimate of the proportion of successes. + pvalue : float + The p-value of the hypothesis test. + + The object has the following methods: + + proportion_ci(confidence_level=0.95, method='exact') : + Compute the confidence interval for ``statistic``. + + Notes + ----- + .. versionadded:: 1.7.0 + + References + ---------- + .. [1] Binomial test, https://en.wikipedia.org/wiki/Binomial_test + .. [2] Jerrold H. Zar, Biostatistical Analysis (fifth edition), + Prentice Hall, Upper Saddle River, New Jersey USA (2010) + + Examples + -------- + >>> from scipy.stats import binomtest + + A car manufacturer claims that no more than 10% of their cars are unsafe. + 15 cars are inspected for safety, 3 were found to be unsafe. Test the + manufacturer's claim: + + >>> result = binomtest(3, n=15, p=0.1, alternative='greater') + >>> result.pvalue + 0.18406106910639114 + + The null hypothesis cannot be rejected at the 5% level of significance + because the returned p-value is greater than the critical value of 5%. + + The test statistic is equal to the estimated proportion, which is simply + ``3/15``: + + >>> result.statistic + 0.2 + + We can use the `proportion_ci()` method of the result to compute the + confidence interval of the estimate: + + >>> result.proportion_ci(confidence_level=0.95) + ConfidenceInterval(low=0.05684686759024681, high=1.0) + + """ + xp = array_namespace(k, n, p) + k, n, p = xp_promote(k, n, p, force_floating=True, broadcast=True, xp=xp) + k_valid = (k >= 0) & (k <= n) & (k == xp.floor(k)) + n_valid = (n >= 1) & (n == xp.floor(n)) + p_valid = (p >= 0) & (p <= 1) + valid = k_valid & n_valid & p_valid + k = xp.where(valid, k, xp.nan) + n = xp.where(valid, n, xp.nan) + p = xp.where(valid, p, xp.nan) + + if alternative not in ('two-sided', 'less', 'greater'): + raise ValueError(f"alternative ('{alternative}') not recognized; \n" + "must be 'two-sided', 'less' or 'greater'") + + B = _SimpleBinomial(n, p) + if alternative == 'less': + pval = B.cdf(k) + elif alternative == 'greater': + pval = B.sf(k - 1) + else: + if is_jax(xp): + message = "`alternative='two-sided'` is incompatible with JAX arrays." + raise ValueError(message) + + # alternative is 'two-sided' + d = B.pmf(k) + rerr = 1 + 1e-7 + + def k_lt_pn(d, k, p, n): + B = _SimpleBinomial(n, p) + ix = _binary_search_for_binom_tst(lambda x1: -B.pmf(x1), -d*rerr, + xp.ceil(p * n), n, xp=xp) + # y is the number of terms between mode and n that are <= d*rerr. + # ix gave us the first term where a(ix) <= d*rerr < a(ix-1) + # if the first equality doesn't hold, y=n-ix. Otherwise, we + # need to include ix as well as the equality holds. Note that + # the equality will hold in very very rare situations due to rerr. + y = n - ix + xp.asarray(d*rerr == B.pmf(ix), dtype=ix.dtype) + pval = B.cdf(k) + B.sf(n - y) + return pval + + def k_gte_pn(d, k, p, n): + B = _SimpleBinomial(n, p) + ix = _binary_search_for_binom_tst(B.pmf, d*rerr, + xp.zeros_like(n), xp.floor(p * n), xp=xp) + # y is the number of terms between 0 and mode that are <= d*rerr. + # we need to add a 1 to account for the 0 index. + # For comparing this with old behavior, see + # tst_binary_srch_for_binom_tst method in test_morestats. + y = ix + 1 + pval = B.cdf(y-1) + B.sf(k-1) + return pval + + pval = xpx.apply_where(k < p*n, (d, k, p, n), k_lt_pn, k_gte_pn) + # xp.minimum(1.0, pval) but for data-apis/array-api-compat#271 + pval = xp.minimum(xp.asarray(1.0, dtype=pval.dtype), pval) + + statistic = xp.where(valid, k/n, xp.nan) + pval = xp.where(valid, pval, xp.nan) + if statistic.ndim == 0: + k, n, statistic, pval = k[()], n[()], statistic[()], pval[()] + + result = BinomTestResult(k=k, n=n, alternative=alternative, + statistic=statistic, pvalue=pval, xp=xp) + return result + + +def _binary_search_for_binom_tst(a, d, lo, hi, *, xp): + """ + Conducts an implicit binary search on a function specified by `a`. + + Meant to be used on the binomial PMF for the case of two-sided tests + to obtain the value on the other side of the mode where the tail + probability should be computed. The values on either side of + the mode are always in order, meaning binary search is applicable. + + Parameters + ---------- + a : callable + The function over which to perform binary search. Its values + for inputs lo and hi should be in ascending order. + d : float + The value to search. + lo : int + The lower end of range to search. + hi : int + The higher end of the range to search. + + Returns + ------- + int + The index, i between lo and hi + such that a(i)<=d d + hi = xpx.at(hi)[i_gt].set(mid[i_gt] - 1) + + i_eq = (midval == d) + mid_i_eq = mid[i_eq] + lo = xpx.at(lo)[i_eq].set(mid_i_eq) + hi = xpx.at(hi)[i_eq].set(mid_i_eq) + + return xp.where(a(lo) <= d, lo, lo-1) + + +class _SimpleBinomial: + # A very simple, array-API compatible binomial distribution for use in + # hypothesis tests. May be replaced by new infrastructure Binomial + # distribution in due time. + def __init__(self, n, p, *, xp=None): + xp = array_namespace(n, p) if xp is None else xp + self.n = n + self.p = p + self.xp = xp + + def f(self, x, fun): + xp = self.xp + return xpx.lazy_apply(fun, x, self.n, self.p, as_numpy=True, xp=xp) + + def cdf(self, x): + return self.xp.where(x >= 0, self.f(self.xp.floor(x), scu._binom_cdf), 0.0) + + def sf(self, x): + return self.xp.where(x >= 0, self.f(self.xp.floor(x), scu._binom_sf), 1.0) + + def ppf(self, x): + return self.f(x, scu._binom_ppf) + + def isf(self, x): + return self.f(x, scu._binom_isf) + + def pmf(self, x): + return self.f(self.xp.floor(x), scu._binom_pmf) diff --git a/v12x/pylib/scipy/stats/_bws_test.py b/v12x/pylib/scipy/stats/_bws_test.py new file mode 100644 index 0000000000000000000000000000000000000000..786b5ebe57ada5aa0bda7dc583494004769daf3f --- /dev/null +++ b/v12x/pylib/scipy/stats/_bws_test.py @@ -0,0 +1,204 @@ +import numpy as np +from functools import partial + +from scipy import stats +from scipy.stats._axis_nan_policy import _broadcast_arrays +from scipy._external import array_api_extra as xpx +from scipy._lib._array_api import (xp_capabilities, array_namespace, xp_ravel, + xp_result_type) + + +def _bws_input_validation(x, y, alternative, axis, method): + ''' Input validation and standardization for bws test''' + xp = array_namespace(x, y) + x, y = xp.asarray(x), xp.asarray(y) + x, y = xpx.atleast_nd(x, ndim=1), xpx.atleast_nd(y, ndim=1) + + axis_none = axis is None + ndim = max(x.ndim, y.ndim) + if axis_none: + x = xp_ravel(x, xp=xp) + y = xp_ravel(y, xp=xp) + axis = 0 + elif np.iterable(axis) or int(axis) != axis: + message = "`axis` must be an integer or None." + raise ValueError(message) + elif (axis >= ndim) or (axis < -ndim): + message = "`axis` is not compatible with the shapes of the inputs." + raise ValueError(message) + axis = int(axis) + + x, y = _broadcast_arrays((x, y), axis=axis, xp=xp) + nx, ny = x.shape[-1], y.shape[-1] + if nx < 2 or ny < 2: + raise ValueError('`x` and `y` must contain at least two observations each.') + + x, y = xp.moveaxis(x, axis, -1), xp.moveaxis(y, axis, -1) + z = stats.rankdata(xp.concat((x, y), axis=-1), axis=-1) + z = xp.astype(z, xp_result_type(x, y, force_floating=True, xp=xp)) + x, y = z[..., :x.shape[-1]], z[..., x.shape[-1]:] + x, y = xp.moveaxis(x, -1, axis), xp.moveaxis(y, -1, axis) + + alternatives = {'two-sided', 'less', 'greater'} + alternative = alternative.lower() + if alternative not in alternatives: + raise ValueError(f'`alternative` must be one of {alternatives}.') + + method = stats.PermutationMethod() if method is None else method + if not isinstance(method, stats.PermutationMethod): + raise ValueError('`method` must be an instance of ' + '`scipy.stats.PermutationMethod`') + + return x, y, alternative, axis, method, xp + + +def _bws_statistic(x, y, alternative, axis, xp): + '''Compute the BWS test statistic for two independent samples''' + # Public function currently does not accept `axis`, but `permutation_test` + # uses `axis` to make vectorized call. + + Ri, Hj = xp.sort(x, axis=axis), xp.sort(y, axis=axis) + n, m = Ri.shape[axis], Hj.shape[axis] + i, j = xp.arange(1, n+1, dtype=Ri.dtype), xp.arange(1, m+1, dtype=Hj.dtype) + + Bx_num = Ri - (m + n)/n * i + By_num = Hj - (m + n)/m * j + + if alternative == 'two-sided': + Bx_num = xpx.at(Bx_num)[...].multiply(Bx_num) + By_num = xpx.at(By_num)[...].multiply(By_num) + else: + Bx_num = xpx.at(Bx_num)[...].multiply(xp.abs(Bx_num)) + By_num = xpx.at(By_num)[...].multiply(xp.abs(By_num)) + + Bx_den = i/(n+1) * (1 - i/(n+1)) * m*(m+n)/n + By_den = j/(m+1) * (1 - j/(m+1)) * n*(m+n)/m + + Bx = 1/n * xp.sum(Bx_num/Bx_den, axis=axis) + By = 1/m * xp.sum(By_num/By_den, axis=axis) + + B = (Bx + By) / 2 if alternative == 'two-sided' else (Bx - By) / 2 + + return B + + +@xp_capabilities(skip_backends=[('dask.array', 'no rankdata')]) +def bws_test(x, y, *, alternative="two-sided", axis=0, method=None): + r'''Perform the Baumgartner-Weiss-Schindler test on two independent samples. + + The Baumgartner-Weiss-Schindler (BWS) test is a nonparametric test of + the null hypothesis that the distribution underlying sample `x` + is the same as the distribution underlying sample `y`. Unlike + the Kolmogorov-Smirnov, Wilcoxon, and Cramer-Von Mises tests, + the BWS test weights the integral by the variance of the difference + in cumulative distribution functions (CDFs), emphasizing the tails of the + distributions, which increases the power of the test in many applications. + + Parameters + ---------- + x, y : array-like + Arrays of samples. Shapes must be broadcastable except along `axis`. + alternative : {'two-sided', 'less', 'greater'}, optional + Defines the alternative hypothesis. Default is 'two-sided'. + Let *F(u)* and *G(u)* be the cumulative distribution functions of the + distributions underlying `x` and `y`, respectively. Then the following + alternative hypotheses are available: + + * 'two-sided': the distributions are not equal, i.e. *F(u) ≠ G(u)* for + at least one *u*. + * 'less': the distribution underlying `x` is stochastically less than + the distribution underlying `y`, i.e. *F(u) >= G(u)* for all *u*. + * 'greater': the distribution underlying `x` is stochastically greater + than the distribution underlying `y`, i.e. *F(u) <= G(u)* for all + *u*. + + Under a more restrictive set of assumptions, the alternative hypotheses + can be expressed in terms of the locations of the distributions; + see [2] section 5.1. + axis : int or None, default: 0 + Axis along which the quantiles are computed. + ``None`` ravels both `x` and `y` before performing the calculation. + method : PermutationMethod, optional + Configures the method used to compute the p-value. The default is + the default `PermutationMethod` object. + + Returns + ------- + res : PermutationTestResult + An object with attributes: + + statistic : float + The observed test statistic of the data. + pvalue : float + The p-value for the given alternative. + null_distribution : ndarray + The values of the test statistic generated under the null hypothesis. + + See Also + -------- + scipy.stats.wilcoxon, scipy.stats.mannwhitneyu, scipy.stats.ttest_ind + + Notes + ----- + When ``alternative=='two-sided'``, the statistic is defined by the + equations given in [1]_ Section 2. This statistic is not appropriate for + one-sided alternatives; in that case, the statistic is the *negative* of + that given by the equations in [1]_ Section 2. Consequently, when the + distribution of the first sample is stochastically greater than that of the + second sample, the statistic will tend to be positive. + + References + ---------- + .. [1] Neuhäuser, M. (2005). Exact Tests Based on the + Baumgartner-Weiss-Schindler Statistic: A Survey. Statistical Papers, + 46(1), 1-29. + .. [2] Fay, M. P., & Proschan, M. A. (2010). Wilcoxon-Mann-Whitney or t-test? + On assumptions for hypothesis tests and multiple interpretations of + decision rules. Statistics surveys, 4, 1. + + Examples + -------- + We follow the example of table 3 in [1]_: Fourteen children were divided + randomly into two groups. Their ranks at performing a specific tests are + as follows. + + >>> import numpy as np + >>> x = [1, 2, 3, 4, 6, 7, 8] + >>> y = [5, 9, 10, 11, 12, 13, 14] + + We use the BWS test to assess whether there is a statistically significant + difference between the two groups. + The null hypothesis is that there is no difference in the distributions of + performance between the two groups. We decide that a significance level of + 1% is required to reject the null hypothesis in favor of the alternative + that the distributions are different. + Since the number of samples is very small, we can compare the observed test + statistic against the *exact* distribution of the test statistic under the + null hypothesis. + + >>> from scipy.stats import bws_test + >>> res = bws_test(x, y) + >>> print(res.statistic) + 5.132167152575315 + + This agrees with :math:`B = 5.132` reported in [1]_. The *p*-value produced + by `bws_test` also agrees with :math:`p = 0.0029` reported in [1]_. + + >>> print(res.pvalue) + 0.002913752913752914 + + Because the p-value is below our threshold of 1%, we take this as evidence + against the null hypothesis in favor of the alternative that there is a + difference in performance between the two groups. + ''' + + x, y, alternative, axis, method, xp = _bws_input_validation(x, y, alternative, + axis, method) + bws_statistic = partial(_bws_statistic, alternative=alternative, xp=xp) + + permutation_alternative = 'less' if alternative == 'less' else 'greater' + res = stats.permutation_test((x, y), bws_statistic, + alternative=permutation_alternative, axis=axis, + **method._asdict()) + + return res diff --git a/v12x/pylib/scipy/stats/_censored_data.py b/v12x/pylib/scipy/stats/_censored_data.py new file mode 100644 index 0000000000000000000000000000000000000000..14c90aa3edb8c5303642a9b36a3323fe6fbcac2c --- /dev/null +++ b/v12x/pylib/scipy/stats/_censored_data.py @@ -0,0 +1,468 @@ +from types import GenericAlias +import numpy as np + + +def _validate_1d(a, name, allow_inf=False): + if np.ndim(a) != 1: + raise ValueError(f'`{name}` must be a one-dimensional sequence.') + if np.isnan(a).any(): + raise ValueError(f'`{name}` must not contain nan.') + if not allow_inf and np.isinf(a).any(): + raise ValueError(f'`{name}` must contain only finite values.') + + +def _validate_interval(interval): + interval = np.asarray(interval) + if interval.shape == (0,): + # The input was a sequence with length 0. + interval = interval.reshape((0, 2)) + if interval.ndim != 2 or interval.shape[-1] != 2: + raise ValueError('`interval` must be a two-dimensional array with ' + 'shape (m, 2), where m is the number of ' + 'interval-censored values, but got shape ' + f'{interval.shape}') + + if np.isnan(interval).any(): + raise ValueError('`interval` must not contain nan.') + if np.isinf(interval).all(axis=1).any(): + raise ValueError('In each row in `interval`, both values must not' + ' be infinite.') + if (interval[:, 0] > interval[:, 1]).any(): + raise ValueError('In each row of `interval`, the left value must not' + ' exceed the right value.') + + uncensored_mask = interval[:, 0] == interval[:, 1] + left_mask = np.isinf(interval[:, 0]) + right_mask = np.isinf(interval[:, 1]) + interval_mask = np.isfinite(interval).all(axis=1) & ~uncensored_mask + + uncensored2 = interval[uncensored_mask, 0] + left2 = interval[left_mask, 1] + right2 = interval[right_mask, 0] + interval2 = interval[interval_mask] + + return uncensored2, left2, right2, interval2 + + +def _validate_x_censored(x, censored): + x = np.asarray(x) + if x.ndim != 1: + raise ValueError('`x` must be one-dimensional.') + censored = np.asarray(censored) + if censored.ndim != 1: + raise ValueError('`censored` must be one-dimensional.') + if (~np.isfinite(x)).any(): + raise ValueError('`x` must not contain nan or inf.') + if censored.size != x.size: + raise ValueError('`x` and `censored` must have the same length.') + return x, censored.astype(bool) + + +class CensoredData: + """ + Instances of this class represent censored data. + + Instances may be passed to the ``fit`` method of continuous + univariate SciPy distributions for maximum likelihood estimation. + The *only* method of the univariate continuous distributions that + understands `CensoredData` is the ``fit`` method. An instance of + `CensoredData` can not be passed to methods such as ``pdf`` and + ``cdf``. + + An observation is said to be *censored* when the precise value is unknown, + but it has a known upper and/or lower bound. The conventional terminology + is: + + * left-censored: an observation is below a certain value but it is + unknown by how much. + * right-censored: an observation is above a certain value but it is + unknown by how much. + * interval-censored: an observation lies somewhere on an interval between + two values. + + Left-, right-, and interval-censored data can be represented by + `CensoredData`. + + For convenience, the class methods ``left_censored`` and + ``right_censored`` are provided to create a `CensoredData` + instance from a single one-dimensional array of measurements + and a corresponding boolean array to indicate which measurements + are censored. The class method ``interval_censored`` accepts two + one-dimensional arrays that hold the lower and upper bounds of the + intervals. + + Parameters + ---------- + uncensored : array_like, 1D + Uncensored observations. + left : array_like, 1D + Left-censored observations. + right : array_like, 1D + Right-censored observations. + interval : array_like, 2D, with shape (m, 2) + Interval-censored observations. Each row ``interval[k, :]`` + represents the interval for the kth interval-censored observation. + + Notes + ----- + In the input array `interval`, the lower bound of the interval may + be ``-inf``, and the upper bound may be ``inf``, but at least one must be + finite. When the lower bound is ``-inf``, the row represents a left- + censored observation, and when the upper bound is ``inf``, the row + represents a right-censored observation. If the length of an interval + is 0 (i.e. ``interval[k, 0] == interval[k, 1]``, the observation is + treated as uncensored. So one can represent all the types of censored + and uncensored data in ``interval``, but it is generally more convenient + to use `uncensored`, `left` and `right` for uncensored, left-censored and + right-censored observations, respectively. + + Examples + -------- + In the most general case, a censored data set may contain values that + are left-censored, right-censored, interval-censored, and uncensored. + For example, here we create a data set with five observations. Two + are uncensored (values 1 and 1.5), one is a left-censored observation + of 0, one is a right-censored observation of 10 and one is + interval-censored in the interval [2, 3]. + + >>> import numpy as np + >>> from scipy.stats import CensoredData + >>> data = CensoredData(uncensored=[1, 1.5], left=[0], right=[10], + ... interval=[[2, 3]]) + >>> print(data) + CensoredData(5 values: 2 not censored, 1 left-censored, + 1 right-censored, 1 interval-censored) + + Equivalently, + + >>> data = CensoredData(interval=[[1, 1], + ... [1.5, 1.5], + ... [-np.inf, 0], + ... [10, np.inf], + ... [2, 3]]) + >>> print(data) + CensoredData(5 values: 2 not censored, 1 left-censored, + 1 right-censored, 1 interval-censored) + + A common case is to have a mix of uncensored observations and censored + observations that are all right-censored (or all left-censored). For + example, consider an experiment in which six devices are started at + various times and left running until they fail. Assume that time is + measured in hours, and the experiment is stopped after 30 hours, even + if all the devices have not failed by that time. We might end up with + data such as this:: + + Device Start-time Fail-time Time-to-failure + 1 0 13 13 + 2 2 24 22 + 3 5 22 17 + 4 8 23 15 + 5 10 *** >20 + 6 12 *** >18 + + Two of the devices had not failed when the experiment was stopped; + the observations of the time-to-failure for these two devices are + right-censored. We can represent this data with + + >>> data = CensoredData(uncensored=[13, 22, 17, 15], right=[20, 18]) + >>> print(data) + CensoredData(6 values: 4 not censored, 2 right-censored) + + Alternatively, we can use the method `CensoredData.right_censored` to + create a representation of this data. The time-to-failure observations + are put the list ``ttf``. The ``censored`` list indicates which values + in ``ttf`` are censored. + + >>> ttf = [13, 22, 17, 15, 20, 18] + >>> censored = [False, False, False, False, True, True] + + Pass these lists to `CensoredData.right_censored` to create an + instance of `CensoredData`. + + >>> data = CensoredData.right_censored(ttf, censored) + >>> print(data) + CensoredData(6 values: 4 not censored, 2 right-censored) + + If the input data is interval censored and already stored in two + arrays, one holding the low end of the intervals and another + holding the high ends, the class method ``interval_censored`` can + be used to create the `CensoredData` instance. + + This example creates an instance with four interval-censored values. + The intervals are [10, 11], [0.5, 1], [2, 3], and [12.5, 13.5]. + + >>> a = [10, 0.5, 2, 12.5] # Low ends of the intervals + >>> b = [11, 1.0, 3, 13.5] # High ends of the intervals + >>> data = CensoredData.interval_censored(low=a, high=b) + >>> print(data) + CensoredData(4 values: 0 not censored, 4 interval-censored) + + Finally, we create and censor some data from the `weibull_min` + distribution, and then fit `weibull_min` to that data. We'll assume + that the location parameter is known to be 0. + + >>> from scipy.stats import weibull_min + >>> rng = np.random.default_rng() + + Create the random data set. + + >>> x = weibull_min.rvs(2.5, loc=0, scale=30, size=250, random_state=rng) + >>> x[x > 40] = 40 # Right-censor values greater or equal to 40. + + Create the `CensoredData` instance with the `right_censored` method. + The censored values are those where the value is 40. + + >>> data = CensoredData.right_censored(x, x == 40) + >>> print(data) + CensoredData(250 values: 215 not censored, 35 right-censored) + + 35 values have been right-censored. + + Fit `weibull_min` to the censored data. We expect to shape and scale + to be approximately 2.5 and 30, respectively. + + >>> weibull_min.fit(data, floc=0) + (2.3575922823897315, 0, 30.40650074451254) + + """ + + # generic type compatibility with scipy-stubs + __class_getitem__: classmethod = classmethod(GenericAlias) + + def __init__(self, uncensored=None, *, left=None, right=None, + interval=None): + if uncensored is None: + uncensored = [] + if left is None: + left = [] + if right is None: + right = [] + if interval is None: + interval = np.empty((0, 2)) + + _validate_1d(uncensored, 'uncensored') + _validate_1d(left, 'left') + _validate_1d(right, 'right') + uncensored2, left2, right2, interval2 = _validate_interval(interval) + + self._uncensored = np.concatenate((uncensored, uncensored2)) + self._left = np.concatenate((left, left2)) + self._right = np.concatenate((right, right2)) + # Note that by construction, the private attribute _interval + # will be a 2D array that contains only finite values representing + # intervals with nonzero but finite length. + self._interval = interval2 + + def __repr__(self): + uncensored_str = " ".join(np.array_repr(self._uncensored).split()) + left_str = " ".join(np.array_repr(self._left).split()) + right_str = " ".join(np.array_repr(self._right).split()) + interval_str = " ".join(np.array_repr(self._interval).split()) + return (f"CensoredData(uncensored={uncensored_str}, left={left_str}, " + f"right={right_str}, interval={interval_str})") + + def __str__(self): + num_nc = len(self._uncensored) + num_lc = len(self._left) + num_rc = len(self._right) + num_ic = len(self._interval) + n = num_nc + num_lc + num_rc + num_ic + parts = [f'{num_nc} not censored'] + if num_lc > 0: + parts.append(f'{num_lc} left-censored') + if num_rc > 0: + parts.append(f'{num_rc} right-censored') + if num_ic > 0: + parts.append(f'{num_ic} interval-censored') + return f'CensoredData({n} values: ' + ', '.join(parts) + ')' + + # This is not a complete implementation of the arithmetic operators. + # All we need is subtracting a scalar and dividing by a scalar. + + def __sub__(self, other): + return CensoredData(uncensored=self._uncensored - other, + left=self._left - other, + right=self._right - other, + interval=self._interval - other) + + def __truediv__(self, other): + return CensoredData(uncensored=self._uncensored / other, + left=self._left / other, + right=self._right / other, + interval=self._interval / other) + + def __len__(self): + """ + The number of values (censored and not censored). + """ + return (len(self._uncensored) + len(self._left) + len(self._right) + + len(self._interval)) + + def num_censored(self): + """ + Number of censored values. + + Returns + ------- + int + The number of censored values. + """ + return len(self._left) + len(self._right) + len(self._interval) + + @classmethod + def right_censored(cls, x, censored): + """ + Create a `CensoredData` instance of right-censored data. + + Parameters + ---------- + x : array_like + `x` is the array of observed data or measurements. + `x` must be a one-dimensional sequence of finite numbers. + censored : array_like of bool + `censored` must be a one-dimensional sequence of boolean + values. If ``censored[k]`` is True, the corresponding value + in `x` is right-censored. That is, the value ``x[k]`` + is the lower bound of the true (but unknown) value. + + Returns + ------- + data : `CensoredData` + An instance of `CensoredData` that represents the + collection of uncensored and right-censored values. + + Examples + -------- + >>> from scipy.stats import CensoredData + + Two uncensored values (4 and 10) and two right-censored values + (24 and 25). + + >>> data = CensoredData.right_censored([4, 10, 24, 25], + ... [False, False, True, True]) + >>> data + CensoredData(uncensored=array([ 4., 10.]), + left=array([], dtype=float64), right=array([24., 25.]), + interval=array([], shape=(0, 2), dtype=float64)) + >>> print(data) + CensoredData(4 values: 2 not censored, 2 right-censored) + """ + x, censored = _validate_x_censored(x, censored) + return cls(uncensored=x[~censored], right=x[censored]) + + @classmethod + def left_censored(cls, x, censored): + """ + Create a `CensoredData` instance of left-censored data. + + Parameters + ---------- + x : array_like + `x` is the array of observed data or measurements. + `x` must be a one-dimensional sequence of finite numbers. + censored : array_like of bool + `censored` must be a one-dimensional sequence of boolean + values. If ``censored[k]`` is True, the corresponding value + in `x` is left-censored. That is, the value ``x[k]`` + is the upper bound of the true (but unknown) value. + + Returns + ------- + data : `CensoredData` + An instance of `CensoredData` that represents the + collection of uncensored and left-censored values. + + Examples + -------- + >>> from scipy.stats import CensoredData + + Two uncensored values (0.12 and 0.033) and two left-censored values + (both 1e-3). + + >>> data = CensoredData.left_censored([0.12, 0.033, 1e-3, 1e-3], + ... [False, False, True, True]) + >>> data + CensoredData(uncensored=array([0.12 , 0.033]), + left=array([0.001, 0.001]), right=array([], dtype=float64), + interval=array([], shape=(0, 2), dtype=float64)) + >>> print(data) + CensoredData(4 values: 2 not censored, 2 left-censored) + """ + x, censored = _validate_x_censored(x, censored) + return cls(uncensored=x[~censored], left=x[censored]) + + @classmethod + def interval_censored(cls, low, high): + """ + Create a `CensoredData` instance of interval-censored data. + + This method is useful when all the data is interval-censored, and + the low and high ends of the intervals are already stored in + separate one-dimensional arrays. + + Parameters + ---------- + low : array_like + The one-dimensional array containing the low ends of the + intervals. + high : array_like + The one-dimensional array containing the high ends of the + intervals. + + Returns + ------- + data : `CensoredData` + An instance of `CensoredData` that represents the + collection of censored values. + + Examples + -------- + >>> import numpy as np + >>> from scipy.stats import CensoredData + + ``a`` and ``b`` are the low and high ends of a collection of + interval-censored values. + + >>> a = [0.5, 2.0, 3.0, 5.5] + >>> b = [1.0, 2.5, 3.5, 7.0] + >>> data = CensoredData.interval_censored(low=a, high=b) + >>> print(data) + CensoredData(4 values: 0 not censored, 4 interval-censored) + """ + _validate_1d(low, 'low', allow_inf=True) + _validate_1d(high, 'high', allow_inf=True) + if len(low) != len(high): + raise ValueError('`low` and `high` must have the same length.') + interval = np.column_stack((low, high)) + uncensored, left, right, interval = _validate_interval(interval) + return cls(uncensored=uncensored, left=left, right=right, + interval=interval) + + def _uncensor(self): + """ + This function is used when a non-censored version of the data + is needed to create a rough estimate of the parameters of a + distribution via the method of moments or some similar method. + The data is "uncensored" by taking the given endpoints as the + data for the left- or right-censored data, and the mean for the + interval-censored data. + """ + data = np.concatenate((self._uncensored, self._left, self._right, + self._interval.mean(axis=1))) + return data + + def _supported(self, a, b): + """ + Return a subset of self containing the values that are in + (or overlap with) the interval (a, b). + """ + uncensored = self._uncensored + uncensored = uncensored[(a < uncensored) & (uncensored < b)] + left = self._left + left = left[a < left] + right = self._right + right = right[right < b] + interval = self._interval + interval = interval[(a < interval[:, 1]) & (interval[:, 0] < b)] + return CensoredData(uncensored, left=left, right=right, + interval=interval) diff --git a/v12x/pylib/scipy/stats/_common.py b/v12x/pylib/scipy/stats/_common.py new file mode 100644 index 0000000000000000000000000000000000000000..4011d425cc4afea3c7ee8937526b13f1f92b0850 --- /dev/null +++ b/v12x/pylib/scipy/stats/_common.py @@ -0,0 +1,5 @@ +from collections import namedtuple + + +ConfidenceInterval = namedtuple("ConfidenceInterval", ["low", "high"]) +ConfidenceInterval. __doc__ = "Class for confidence intervals." diff --git a/v12x/pylib/scipy/stats/_constants.py b/v12x/pylib/scipy/stats/_constants.py new file mode 100644 index 0000000000000000000000000000000000000000..b539ce8146ebdbc8e08c66143461b04d742804f2 --- /dev/null +++ b/v12x/pylib/scipy/stats/_constants.py @@ -0,0 +1,42 @@ +""" +Statistics-related constants. + +""" +import numpy as np + + +# The smallest representable positive number such that 1.0 + _EPS != 1.0. +_EPS = np.finfo(float).eps + +# The largest [in magnitude] usable floating value. +_XMAX = np.finfo(float).max + +# The log of the largest usable floating value; useful for knowing +# when exp(something) will overflow +_LOGXMAX = np.log(_XMAX) + +# The smallest [in magnitude] usable (i.e. not subnormal) double precision +# floating value. +_XMIN = np.finfo(float).tiny + +# The log of the smallest [in magnitude] usable (i.e not subnormal) +# double precision floating value. +_LOGXMIN = np.log(_XMIN) + +# -special.psi(1) +_EULER = 0.577215664901532860606512090082402431042 + +# special.zeta(3, 1) Apery's constant +_ZETA3 = 1.202056903159594285399738161511449990765 + +# sqrt(pi) +_SQRT_PI = 1.772453850905516027298167483341145182798 + +# sqrt(2/pi) +_SQRT_2_OVER_PI = 0.7978845608028654 + +# log(pi) +_LOG_PI = 1.1447298858494002 + +# log(sqrt(2/pi)) +_LOG_SQRT_2_OVER_PI = -0.22579135264472744 diff --git a/v12x/pylib/scipy/stats/_continued_fraction.py b/v12x/pylib/scipy/stats/_continued_fraction.py new file mode 100644 index 0000000000000000000000000000000000000000..4966b1c99ed092ba3578aa7ca15b5ba8ea4bc6a9 --- /dev/null +++ b/v12x/pylib/scipy/stats/_continued_fraction.py @@ -0,0 +1,387 @@ +import numpy as np + +from scipy._lib._array_api import ( + array_namespace, xp_ravel, xp_copy, xp_promote +) +import scipy._lib._elementwise_iterative_method as eim +from scipy._lib._util import _RichResult +from scipy import special + +# Todo: +# Avoid special-casing key 'n' in _lib._elementwise_iterative_method::_check_termination +# Rearrange termination condition to allow absolute and relative tolerances? +# Interpret/return |f_n - f_{n-1}| as an error estimate? +# Return gracefully for size=0 arrays + +def _logaddexp(x, y, xp=None): + # logaddexp that supports complex numbers + xp = array_namespace(x, y) if xp is None else xp + x, y = xp.broadcast_arrays(x, y) + xy = xp.stack((x, y), axis=0) + return special.logsumexp(xy, axis=0) + + +def _continued_fraction_iv(a, b, args, tolerances, maxiter, log): + # Input validation for `_continued_fraction` + + if not callable(a) or not callable(b): + raise ValueError('`a` and `b` must be callable.') + + if not np.iterable(args): + args = (args,) + + # Call each callable once to determine namespace and dtypes + a0, b0 = a(0, *args), b(0, *args) + xp = array_namespace(a0, b0, *args) + a0, b0, *args = xp_promote(a0, b0, *args, force_floating=True, broadcast=True, + xp=xp) + shape, dtype = a0.shape, a0.dtype + a0, b0, *args = (xp_ravel(arg) for arg in (a0, b0) + tuple(args)) + + tolerances = {} if tolerances is None else tolerances + eps = tolerances.get('eps', None) + tiny = tolerances.get('tiny', None) + + # tolerances are floats, not arrays, so it's OK to use NumPy + message = ('`eps` and `tiny` must be (or represent the logarithm of) ' + 'finite, positive, real scalars.') + tols = np.asarray([eps if eps is not None else 1, + tiny if tiny is not None else 1]) + not_real = (not np.issubdtype(tols.dtype, np.number) + or np.issubdtype(tols.dtype, np.complexfloating)) + not_positive = np.any(tols <= 0) if not log else False + not_finite = not np.all(np.isfinite(tols)) + not_scalar = tols.shape != (2,) + if not_real or not_positive or not_finite or not_scalar: + raise ValueError(message) + + maxiter_int = int(maxiter) + if maxiter != maxiter_int or maxiter < 0: + raise ValueError('`maxiter` must be a non-negative integer.') + + if not isinstance(log, bool): + raise ValueError('`log` must be boolean.') + + return a, b, args, eps, tiny, maxiter, log, a0, b0, shape, dtype, xp + + +def _continued_fraction(a, b, *, args=(), tolerances=None, maxiter=100, log=False): + r"""Evaluate a generalized continued fraction numerically. + + `_continued_fraction` iteratively evaluates convergents of a continued fraction + given coefficients returned by callables `a` and `b`. Iteration terminates when + `maxiter` terms have been evaluated or a termination criterion controlled by + `tolerances` is satisfied, and the final convergent is returned as the ``f`` + attribute of the result object. + + This function works elementwise when `args` contains (broadcastable) arrays. + + Parameters + ---------- + a, b: callable + Functions that provide the *numerator* and *denominator* coefficients of + the continued fraction, respectively. + + The signature of each must be:: + + a(n: int, *argsj) -> ndarray + + where ``n`` is the coefficient number and ``argsj`` is a tuple, which may + contain an arbitrary number of arrays of any shape. `a` and `b` must be + elementwise functions: each scalar element ``a(n, *argsj)[i]`` must equal + ``a(n, *[argj[i] for argj in argsj])`` for valid indices ``i``. + `a` and `b` must not mutate the arrays in ``argsj``. + + The result shape is the broadcasted shape of ``a(0, *args)`` and + ``b(0, *args)``. The dtype used throughout computation is the result dtype + of these terms if it is a float, and the default float of the array library + otherwise. The numerical value of ``a(0, *args)`` is ignored, and + the value of the leading term ``b(0, *args)`` is the so-called "integer" + part of the continued fraction (although it need not be integral). + + args : tuple of array_like, optional + Additional positional *array* arguments to be passed to `a` and `b`. Arrays + must be broadcastable with one another. If the coefficient callables + require additional arguments that are not broadcastable with one + another, wrap them with callables `a` and `b` such that `a` and `b` accept + only ``n`` and broadcastable array arguments. + tolerances : dictionary of floats, optional + Tolerances and numerical thresholds used by the algorithm. Currently, + valid keys of the dictionary are: + + - ``eps`` - the convergence threshold of Lentz' algorithm + - ``tiny`` - not strictly a "tolerance", but a very small positive number + used to avoid division by zero + + The default `eps` is the precision of the appropriate dtype, and the default + `tiny` is the precision squared. [1]_ advises that ``eps`` is "as small as + you like", but for most purposes, it should not be set smaller than the default + because it may prevent convergence of the algorithm. [1]_ also advises that + ``tiny`` should be less than typical values of ``eps * b(n)``, so the default + is a good choice unless the :math:`b_n` are very small. See [1]_ for details. + maxiter : int, default: 100 + The maximum number of iterations of the algorithm to perform. + log : bool, default: False + If True, `a` and `b` return the (natural) logarithm of the terms, `tolerances` + contains the logarithm of the tolerances, and the result object reports the + logarithm of the convergent. + + Returns + ------- + res : _RichResult + An object similar to an instance of `scipy.optimize.OptimizeResult` with the + following attributes. The descriptions are written as though the values will + be scalars; however, if `f` returns an array, the outputs will be + arrays of the same shape. + + success : bool array + ``True`` where the algorithm terminated successfully (status ``0``); + ``False`` otherwise. + status : int array + An integer representing the exit status of the algorithm. + + - ``0`` : The algorithm converged to the specified tolerances. + - ``-2`` : The maximum number of iterations was reached. + - ``-3`` : A non-finite value was encountered. + + f : float array + The convergent which satisfied a termination criterion. + nit : int array + The number of iterations of the algorithm that were performed. + nfev : int array + The number of terms that were evaluated. + + Notes + ----- + A generalized continued fraction is an expression of the form + + .. math:: + + b_0 + \frac{a_1}{b_1 + \frac{a_2}{b_2 + \frac{a_3}{b_3 + \cdots}}} + + Successive "convergents" approximate the infinitely recursive continued fraction + with a finite number of terms :math:`a_n` and :math:`b_n`, which are provided + by callables `a` and `b`, respectively. This implementation follows the modified + Lentz algorithm ([1]_, [2]_) to evaluate successive convergents until a + termination condition is satisfied. + + References + ---------- + .. [1] Press, William H., and Saul A. Teukolsky. "Evaluating continued fractions + and computing exponential integrals." Computers in Physics 2.5 (1988): 88-89. + .. [2] Lentz's algorithm. Wikipedia. + https://en.wikipedia.org/wiki/Lentz%27s_algorithm + .. [3] Continued fraction. Wikipedia. + https://en.wikipedia.org/wiki/Continued_fraction + .. [4] Generalized continued fraction. Wikipedia. + https://en.wikipedia.org/wiki/Generalized_continued_fraction + + Examples + -------- + The "simple continued fraction" of :math:`\pi` is given in [3]_ as + + .. math:: + + 3 + \frac{1}{7 + \frac{1}{15 + \frac{1}{1 + \cdots}}} + + where the :math:`b_n` terms follow no obvious pattern: + + >>> b = [3, 7, 15, 1, 292, 1, 1, 1, 2, 1, 3, 1] + + and the :math:`a_n` terms are all :math:`1`. + In this case, all the terms have been precomputed, so we call `_continued_fraction` + with simple callables which simply return the precomputed coefficients: + + >>> import numpy as np + >>> from scipy.special._continued_fraction import _continued_fraction + >>> res = _continued_fraction(a=lambda n: 1, b=lambda n: b[n], maxiter=len(b) - 1) + >>> (res.f - np.pi) / np.pi + np.float64(7.067899292141148e-15) + + A generalized continued fraction for :math:`\pi` is given by: + + .. math:: + + 3 + \frac{1^2}{6 + \frac{3^2}{6 + \frac{5^2}{6 + \cdots}}} + + We define the coefficient callables as: + + >>> def a(n): + ... return (2*n - 1)**2 + >>> + >>> def b(n): + ... if n == 0: + ... return 3 + ... else: + ... return 6 + + Then the continued fraction can be evaluated as: + + >>> res = _continued_fraction(a, b) + >>> res + success: False + status: -2 + f: 3.1415924109719846 + nit: 100 + nfev: 101 + + Note that the requested tolerance was not reached within the (default) + maximum number of iterations because it converges very slowly. + An expression that converges more rapidly is expressed as the difference + between two continued fractions. We will compute both of them in one + vectorized call to `_continued_fraction`. + + >>> u, v = 5, 239 + >>> + >>> def a(n, a1, _): + ... # The shape of the output must be the shape of the arguments + ... shape = a1.shape + ... if n == 0: + ... return np.zeros(shape) + ... elif n == 1: + ... return a1 + ... else: + ... return np.full(shape, (n-1)**2) + >>> + >>> def b(n, _, uv): + ... shape = uv.shape + ... if n == 0: + ... return np.zeros(shape) + ... return np.full(shape, (2*n - 1)*uv) + >>> + >>> res = _continued_fraction(a, b, args=([16, 4], [u, v])) + >>> res + success: [ True True] + status: [0 0] + f: [ 3.158e+00 1.674e-02] + nit: [10 4] + nfev: [11 5] + + Note that the second term converged in fewer than half the number of iterations + as the first. The approximation of :math:`\pi` is the difference between the two: + + >>> pi = res.f[0] - res.f[1] + >>> (pi - np.pi) / np.pi + np.float64(2.8271597168564594e-16) + + If it is more efficient to compute the :math:`a_n` and :math:`b_n` terms together, + consider instantiating a class with a method that computes both terms and stores + the results in an attribute. Separate methods of the class retrieve the + coefficients, and these methods are passed to `_continued_fraction` as arguments + `a` and `b`. Similarly,if the coefficients can be computed recursively in terms of + previous coefficients, use a class to maintain state between callable evaluations. + + """ + + res = _continued_fraction_iv(a, b, args, tolerances, maxiter, log) + a, b, args, eps, tiny, maxiter, log, a0, b0, shape, dtype, xp = res + callback = None # don't want to test it, but easy to add later + + # The EIM framework was designed for the case in where there would + # be only one callable, and all arguments of the callable would be + # arrays. We're going a bit beyond that here, since we have two callables, + # and the first argument is an integer (the number of the term). Rather + # than complicate the framework, we wrap the user-provided callables to + # make this problem fit within the existing framework. + + def a(n, *args, a=a): + n = int(xp.real(xp_ravel(n))[0]) + return a(n, *args) + + def b(n, *args, b=b): + n = int(xp.real(xp_ravel(n))[0]) + return b(n, *args) + + def func(n, *args): + return xp.stack((a(n, *args), b(n, *args)), axis=-1) + + status = xp.full_like(a0, eim._EINPROGRESS, dtype=xp.int32) # in progress + nit, nfev = 0, 1 # one function evaluation (per function) performed above + maxiter = 100 if maxiter is None else maxiter + + # Quotations describing the algorithm are from [1]_ + # "... as small as you like, say eps" + if eps is None: + eps = xp.finfo(dtype).eps if not log else np.log(xp.finfo(dtype).eps) + + # "The parameter tiny should be less than typical values of eps |b_n|" + if tiny is None: + tiny = xp.finfo(dtype).eps**2 if not log else 2*np.log(xp.finfo(dtype).eps) + + # "Set f0 and C0 to the value b0 or to tiny if b0=0. Set D0 = 0. + zero = -xp.inf if log else 0 + fn = xp.where(b0 == zero, tiny, b0) + Cnm1 = xp_copy(fn) + Dnm1 = xp.full_like(fn, zero) + + CnDn = xp.full_like(fn, xp.inf) + + work = _RichResult(n=0, fn=fn, Cnm1=Cnm1, Dnm1=Dnm1, CnDn=CnDn, + eps=eps, tiny=tiny, + nit=nit, nfev=nfev, status=status) + res_work_pairs = [('status', 'status'), ('f', 'fn'), + ('nit', 'nit'), ('nfev', 'nfev')] + + def pre_func_eval(work): + work.n = xp.reshape(xp.asarray(work.n + 1), (-1,)) + return work.n + + def post_func_eval(n, ab, work): + an, bn = ab[..., 0], ab[..., 1] + + zero = 0 if not log else -xp.inf + + # "Set D_n = 1/(b_n + a_n D_{n-1}) or 1/tiny, if the denominator vanishes" + denominator = (bn + an*work.Dnm1 if not log + else _logaddexp(bn, an + work.Dnm1, xp=xp)) + denominator[denominator == zero] = tiny + Dn = (1/denominator if not log + else -denominator) + + # "Set C_n = b_n + a_n / C_{n-1} (or =tiny, if the expression vanishes)" + Cn = (bn + an / work.Cnm1 if not log + else _logaddexp(bn, an - work.Cnm1, xp=xp)) + Cn[Cn == zero] = tiny + + # "and set f_n = f_{n-1} C_n D_n" + work.CnDn = (Cn * Dn if not log + else Cn + Dn) + work.fn = (work.fn * work.CnDn if not log + else work.fn + work.CnDn) + + + work.Cnm1, work.Dnm1 = Cn, Dn + + def check_termination(work): + # Check for all terminal conditions and record statuses. + stop = xp.zeros_like(work.CnDn, dtype=xp.bool) + + # "You quit when |D_n C_n - 1| is as small as you like, say eps" + pij = xp.full_like(work.CnDn, xp.pi*1j) if log else None + residual = (xp.abs(work.CnDn - 1) if not log + else xp.real(_logaddexp(work.CnDn, pij, xp=xp))) + i = residual < work.eps + work.status[i] = eim._ECONVERGED + stop[i] = True + + # If function value is NaN, report failure. + i = (~xp.isfinite(work.fn) if not log + else ~(xp.isfinite(work.fn) | (work.fn == -xp.inf))) + work.status[i] = eim._EVALUEERR + stop[i] = True + + return stop + + def post_termination_check(work): + pass + + def customize_result(res, shape): + # Only needed pre-NEP 50 + res['f'] = xp.asarray(res['f'], dtype=dtype) + res['f'] = res['f'][()] if res['f'].ndim == 0 else res['f'] + return shape + + return eim._loop(work, callback, shape, maxiter, func, args, dtype, + pre_func_eval, post_func_eval, check_termination, + post_termination_check, customize_result, res_work_pairs, + xp=xp) diff --git a/v12x/pylib/scipy/stats/_continuous_distns.py b/v12x/pylib/scipy/stats/_continuous_distns.py new file mode 100644 index 0000000000000000000000000000000000000000..98029b3ec326e29bb33ce9832232b89ab31e2641 --- /dev/null +++ b/v12x/pylib/scipy/stats/_continuous_distns.py @@ -0,0 +1,12628 @@ +# +# Author: Travis Oliphant 2002-2011 with contributions from +# SciPy Developers 2004-2011 +# +import warnings +from collections.abc import Iterable +from functools import wraps, cached_property +import ctypes +import operator + +import numpy as np +from numpy.polynomial import Polynomial +from scipy.interpolate import BSpline +from scipy._lib.doccer import (extend_notes_in_docstring, + replace_notes_in_docstring, + inherit_docstring_from) +from scipy._lib._ccallback import LowLevelCallable +from scipy import optimize +from scipy import integrate +import scipy.special as sc + +import scipy.special._ufuncs as scu +from scipy._lib._util import _lazyselect +import scipy._external.array_api_extra as xpx +from scipy._lib._array_api import xp_promote + +from . import _stats +from ._tukeylambda_stats import (tukeylambda_variance as _tlvar, + tukeylambda_kurtosis as _tlkurt) +from ._distn_infrastructure import (_vectorize_rvs_over_shapes, + get_distribution_names, _kurtosis, _isintegral, + rv_continuous, _skew, _get_fixed_fit_value, _check_shape, _ShapeInfo) +from ._ksstats import kolmogn, kolmognp, kolmogni +from ._constants import (_XMIN, _LOGXMIN, _EULER, _ZETA3, _SQRT_PI, + _SQRT_2_OVER_PI, _LOG_PI, _LOG_SQRT_2_OVER_PI) +from ._censored_data import CensoredData +from scipy.optimize import root_scalar +from scipy.stats._warnings_errors import FitError +import scipy.stats as stats + +def _remove_optimizer_parameters(kwds): + """ + Remove the optimizer-related keyword arguments 'loc', 'scale' and + 'optimizer' from `kwds`. Then check that `kwds` is empty, and + raise `TypeError("Unknown arguments: %s." % kwds)` if it is not. + + This function is used in the fit method of distributions that override + the default method and do not use the default optimization code. + + `kwds` is modified in-place. + """ + kwds.pop('loc', None) + kwds.pop('scale', None) + kwds.pop('optimizer', None) + kwds.pop('method', None) + if kwds: + raise TypeError(f"Unknown arguments: {kwds}.") + + +def _call_super_mom(fun): + # If fit method is overridden only for MLE and doesn't specify what to do + # if method == 'mm' or with censored data, this decorator calls the generic + # implementation. + @wraps(fun) + def wrapper(self, data, *args, **kwds): + method = kwds.get('method', 'mle').lower() + censored = isinstance(data, CensoredData) + if method == 'mm' or (censored and data.num_censored() > 0): + return super(type(self), self).fit(data, *args, **kwds) + else: + if censored: + # data is an instance of CensoredData, but actually holds + # no censored values, so replace it with the array of + # uncensored values. + data = data._uncensored + return fun(self, data, *args, **kwds) + + return wrapper + + +def _get_left_bracket(fun, rbrack, lbrack=None): + # find left bracket for `root_scalar`. A guess for lbrack may be provided. + lbrack = lbrack or rbrack - 1 + diff = rbrack - lbrack + + # if there is no sign change in `fun` between the brackets, expand + # rbrack - lbrack until a sign change occurs + def interval_contains_root(lbrack, rbrack): + # return true if the signs disagree. + return np.sign(fun(lbrack)) != np.sign(fun(rbrack)) + + while not interval_contains_root(lbrack, rbrack): + diff *= 2 + lbrack = rbrack - diff + + msg = ("The solver could not find a bracket containing a " + "root to an MLE first order condition.") + if np.isinf(lbrack): + raise FitSolverError(msg) + + return lbrack + + +class ksone_gen(rv_continuous): + r"""Kolmogorov-Smirnov one-sided test statistic distribution. + + This is the distribution of the one-sided Kolmogorov-Smirnov (KS) + statistics :math:`D_n^+` and :math:`D_n^-` + for a finite sample size ``n >= 1`` (the shape parameter). + + %(before_notes)s + + See Also + -------- + kstwobign, kstwo, kstest + + Notes + ----- + :math:`D_n^+` and :math:`D_n^-` are given by + + .. math:: + + D_n^+ &= \text{sup}_x (F_n(x) - F(x)),\\ + D_n^- &= \text{sup}_x (F(x) - F_n(x)),\\ + + where :math:`F` is a continuous CDF and :math:`F_n` is an empirical CDF. + `ksone` describes the distribution under the null hypothesis of the KS test + that the empirical CDF corresponds to :math:`n` i.i.d. random variates + with CDF :math:`F`. + + %(after_notes)s + + References + ---------- + .. [1] Birnbaum, Z. W. and Tingey, F.H. "One-sided confidence contours + for probability distribution functions", The Annals of Mathematical + Statistics, 22(4), pp 592-596 (1951). + + Examples + -------- + >>> import numpy as np + >>> from scipy.stats import ksone + >>> import matplotlib.pyplot as plt + >>> fig, ax = plt.subplots(1, 1) + + Display the probability density function (``pdf``): + + >>> n = 1e+03 + >>> x = np.linspace(ksone.ppf(0.01, n), + ... ksone.ppf(0.99, n), 100) + >>> ax.plot(x, ksone.pdf(x, n), + ... 'r-', lw=5, alpha=0.6, label='ksone pdf') + + Alternatively, the distribution object can be called (as a function) + to fix the shape, location and scale parameters. This returns a "frozen" + RV object holding the given parameters fixed. + + Freeze the distribution and display the frozen ``pdf``: + + >>> rv = ksone(n) + >>> ax.plot(x, rv.pdf(x), 'k-', lw=2, label='frozen pdf') + >>> ax.legend(loc='best', frameon=False) + >>> plt.show() + + Check accuracy of ``cdf`` and ``ppf``: + + >>> vals = ksone.ppf([0.001, 0.5, 0.999], n) + >>> np.allclose([0.001, 0.5, 0.999], ksone.cdf(vals, n)) + True + + """ + def _argcheck(self, n): + return (n >= 1) & (n == np.round(n)) + + def _shape_info(self): + return [_ShapeInfo("n", True, (1, np.inf), (True, False))] + + def _pdf(self, x, n): + return -scu._smirnovp(n, x) + + def _cdf(self, x, n): + return scu._smirnovc(n, x) + + def _sf(self, x, n): + return sc.smirnov(n, x) + + def _ppf(self, q, n): + return scu._smirnovci(n, q) + + def _isf(self, q, n): + return sc.smirnovi(n, q) + + +ksone = ksone_gen(a=0.0, b=1.0, name='ksone') + + +class kstwo_gen(rv_continuous): + r"""Kolmogorov-Smirnov two-sided test statistic distribution. + + This is the distribution of the two-sided Kolmogorov-Smirnov (KS) + statistic :math:`D_n` for a finite sample size ``n >= 1`` + (the shape parameter). + + %(before_notes)s + + See Also + -------- + kstwobign, ksone, kstest + + Notes + ----- + :math:`D_n` is given by + + .. math:: + + D_n = \text{sup}_x |F_n(x) - F(x)| + + where :math:`F` is a (continuous) CDF and :math:`F_n` is an empirical CDF. + `kstwo` describes the distribution under the null hypothesis of the KS test + that the empirical CDF corresponds to :math:`n` i.i.d. random variates + with CDF :math:`F`. + + %(after_notes)s + + References + ---------- + .. [1] Simard, R., L'Ecuyer, P. "Computing the Two-Sided + Kolmogorov-Smirnov Distribution", Journal of Statistical Software, + Vol 39, 11, 1-18 (2011). + + Examples + -------- + >>> import numpy as np + >>> from scipy.stats import kstwo + >>> import matplotlib.pyplot as plt + >>> fig, ax = plt.subplots(1, 1) + + Display the probability density function (``pdf``): + + >>> n = 10 + >>> x = np.linspace(kstwo.ppf(0.01, n), + ... kstwo.ppf(0.99, n), 100) + >>> ax.plot(x, kstwo.pdf(x, n), + ... 'r-', lw=5, alpha=0.6, label='kstwo pdf') + + Alternatively, the distribution object can be called (as a function) + to fix the shape, location and scale parameters. This returns a "frozen" + RV object holding the given parameters fixed. + + Freeze the distribution and display the frozen ``pdf``: + + >>> rv = kstwo(n) + >>> ax.plot(x, rv.pdf(x), 'k-', lw=2, label='frozen pdf') + >>> ax.legend(loc='best', frameon=False) + >>> plt.show() + + Check accuracy of ``cdf`` and ``ppf``: + + >>> vals = kstwo.ppf([0.001, 0.5, 0.999], n) + >>> np.allclose([0.001, 0.5, 0.999], kstwo.cdf(vals, n)) + True + + """ + def _argcheck(self, n): + return (n >= 1) & (n == np.round(n)) + + def _shape_info(self): + return [_ShapeInfo("n", True, (1, np.inf), (True, False))] + + def _get_support(self, n): + return (0.5/(n if not isinstance(n, Iterable) else np.asanyarray(n)), + 1.0) + + def _pdf(self, x, n): + return kolmognp(n, x) + + def _cdf(self, x, n): + return kolmogn(n, x) + + def _sf(self, x, n): + return kolmogn(n, x, cdf=False) + + def _ppf(self, q, n): + return kolmogni(n, q, cdf=True) + + def _isf(self, q, n): + return kolmogni(n, q, cdf=False) + + +# Use the pdf, (not the ppf) to compute moments +kstwo = kstwo_gen(momtype=0, a=0.0, b=1.0, name='kstwo') + + +class kstwobign_gen(rv_continuous): + r"""Limiting distribution of scaled Kolmogorov-Smirnov two-sided test statistic. + + This is the asymptotic distribution of the two-sided Kolmogorov-Smirnov + statistic :math:`\sqrt{n} D_n` that measures the maximum absolute + distance of the theoretical (continuous) CDF from the empirical CDF. + (see `kstest`). + + %(before_notes)s + + See Also + -------- + ksone, kstwo, kstest + + Notes + ----- + :math:`\sqrt{n} D_n` is given by + + .. math:: + + D_n = \text{sup}_x |F_n(x) - F(x)| + + where :math:`F` is a continuous CDF and :math:`F_n` is an empirical CDF. + `kstwobign` describes the asymptotic distribution (i.e. the limit of + :math:`\sqrt{n} D_n`) under the null hypothesis of the KS test that the + empirical CDF corresponds to i.i.d. random variates with CDF :math:`F`. + + %(after_notes)s + + References + ---------- + .. [1] Feller, W. "On the Kolmogorov-Smirnov Limit Theorems for Empirical + Distributions", Ann. Math. Statist. Vol 19, 177-189 (1948). + + %(example)s + + """ + def _shape_info(self): + return [] + + def _pdf(self, x): + return -scu._kolmogp(x) + + def _cdf(self, x): + return scu._kolmogc(x) + + def _sf(self, x): + return sc.kolmogorov(x) + + def _ppf(self, q): + return scu._kolmogci(q) + + def _isf(self, q): + return sc.kolmogi(q) + + +kstwobign = kstwobign_gen(a=0.0, name='kstwobign') + + +## Normal distribution + +# loc = mu, scale = std +# Keep these implementations out of the class definition so they can be reused +# by other distributions. +_norm_pdf_C = np.sqrt(2*np.pi) +_norm_pdf_logC = np.log(_norm_pdf_C) + + +def _norm_pdf(x): + return np.exp(-x**2/2.0) / _norm_pdf_C + + +def _norm_logpdf(x): + return -x**2 / 2.0 - _norm_pdf_logC + + +def _norm_cdf(x): + return sc.ndtr(x) + + +def _norm_logcdf(x): + return sc.log_ndtr(x) + + +def _norm_ppf(q): + return sc.ndtri(q) + + +def _norm_sf(x): + return _norm_cdf(-x) + + +def _norm_logsf(x): + return _norm_logcdf(-x) + + +def _norm_isf(q): + return -_norm_ppf(q) + + +class norm_gen(rv_continuous): + r"""A normal continuous random variable. + + The location (``loc``) keyword specifies the mean. + The scale (``scale``) keyword specifies the standard deviation. + + %(before_notes)s + + Notes + ----- + The probability density function for `norm` is: + + .. math:: + + f(x) = \frac{\exp(-x^2/2)}{\sqrt{2\pi}} + + for a real number :math:`x`. + + %(after_notes)s + + %(example)s + + """ + def _shape_info(self): + return [] + + def _rvs(self, size=None, random_state=None): + return random_state.standard_normal(size) + + def _pdf(self, x): + # norm.pdf(x) = exp(-x**2/2)/sqrt(2*pi) + return _norm_pdf(x) + + def _logpdf(self, x): + return _norm_logpdf(x) + + def _cdf(self, x): + return _norm_cdf(x) + + def _logcdf(self, x): + return _norm_logcdf(x) + + def _sf(self, x): + return _norm_sf(x) + + def _logsf(self, x): + return _norm_logsf(x) + + def _ppf(self, q): + return _norm_ppf(q) + + def _isf(self, q): + return _norm_isf(q) + + def _stats(self): + return 0.0, 1.0, 0.0, 0.0 + + def _entropy(self): + return 0.5*(np.log(2*np.pi)+1) + + @_call_super_mom + @replace_notes_in_docstring(rv_continuous, notes="""\ + For the normal distribution, method of moments and maximum likelihood + estimation give identical fits, and explicit formulas for the estimates + are available. + This function uses these explicit formulas for the maximum likelihood + estimation of the normal distribution parameters, so the + `optimizer` and `method` arguments are ignored.\n\n""") + def fit(self, data, **kwds): + floc = kwds.pop('floc', None) + fscale = kwds.pop('fscale', None) + + _remove_optimizer_parameters(kwds) + + if floc is not None and fscale is not None: + # This check is for consistency with `rv_continuous.fit`. + # Without this check, this function would just return the + # parameters that were given. + raise ValueError("All parameters fixed. There is nothing to " + "optimize.") + + data = np.asarray(data) + + if not np.isfinite(data).all(): + raise ValueError("The data contains non-finite values.") + + if floc is None: + loc = data.mean() + else: + loc = floc + + if fscale is None: + scale = np.sqrt(((data - loc)**2).mean()) + else: + scale = fscale + + return loc, scale + + def _munp(self, n): + """ + @returns Moments of standard normal distribution for integer n >= 0 + + See eq. 16 of https://arxiv.org/abs/1209.4340v2 + """ + if n == 0: + return 1. + if n % 2 == 0: + return sc.factorial2(int(n) - 1) + else: + return 0. + + +norm = norm_gen(name='norm') + + +class alpha_gen(rv_continuous): + r"""An alpha continuous random variable. + + %(before_notes)s + + Notes + ----- + The probability density function for `alpha` ([1]_, [2]_) is: + + .. math:: + + f(x, a) = \frac{1}{x^2 \Phi(a) \sqrt{2\pi}} * + \exp(-\frac{1}{2} (a-1/x)^2) + + where :math:`\Phi` is the normal CDF, :math:`x > 0`, and :math:`a > 0`. + + `alpha` takes ``a`` as a shape parameter. + + %(after_notes)s + + References + ---------- + .. [1] Johnson, Kotz, and Balakrishnan, "Continuous Univariate + Distributions, Volume 1", Second Edition, John Wiley and Sons, + p. 173 (1994). + .. [2] Anthony A. Salvia, "Reliability applications of the Alpha + Distribution", IEEE Transactions on Reliability, Vol. R-34, + No. 3, pp. 251-252 (1985). + + %(example)s + + """ + _support_mask = rv_continuous._open_support_mask + + def _shape_info(self): + return [_ShapeInfo("a", False, (0, np.inf), (False, False))] + + def _pdf(self, x, a): + # alpha.pdf(x, a) = 1/(x**2*Phi(a)*sqrt(2*pi)) * exp(-1/2 * (a-1/x)**2) + return 1.0/(x**2)/_norm_cdf(a)*_norm_pdf(a-1.0/x) + + def _logpdf(self, x, a): + return -2*np.log(x) + _norm_logpdf(a-1.0/x) - np.log(_norm_cdf(a)) + + def _cdf(self, x, a): + return _norm_cdf(a-1.0/x) / _norm_cdf(a) + + def _ppf(self, q, a): + return 1.0/np.asarray(a - _norm_ppf(q*_norm_cdf(a))) + + def _stats(self, a): + return [np.inf]*2 + [np.nan]*2 + + +alpha = alpha_gen(a=0.0, name='alpha') + + +class anglit_gen(rv_continuous): + r"""An anglit continuous random variable. + + %(before_notes)s + + Notes + ----- + The probability density function for `anglit` is: + + .. math:: + + f(x) = \sin(2x + \pi/2) = \cos(2x) + + for :math:`-\pi/4 \le x \le \pi/4`. + + %(after_notes)s + + %(example)s + + """ + def _shape_info(self): + return [] + + def _pdf(self, x): + # anglit.pdf(x) = sin(2*x + \pi/2) = cos(2*x) + return np.cos(2*x) + + def _cdf(self, x): + return np.sin(x+np.pi/4)**2.0 + + def _sf(self, x): + return np.cos(x + np.pi / 4) ** 2.0 + + def _ppf(self, q): + return np.arcsin(np.sqrt(q))-np.pi/4 + + def _stats(self): + return 0.0, np.pi*np.pi/16-0.5, 0.0, -2*(np.pi**4 - 96)/(np.pi*np.pi-8)**2 + + def _entropy(self): + return 1-np.log(2) + + +anglit = anglit_gen(a=-np.pi/4, b=np.pi/4, name='anglit') + + +class arcsine_gen(rv_continuous): + r"""An arcsine continuous random variable. + + %(before_notes)s + + Notes + ----- + The probability density function for `arcsine` is: + + .. math:: + + f(x) = \frac{1}{\pi \sqrt{x (1-x)}} + + for :math:`0 < x < 1`. + + %(after_notes)s + + %(example)s + + """ + def _shape_info(self): + return [] + + def _pdf(self, x): + # arcsine.pdf(x) = 1/(pi*sqrt(x*(1-x))) + with np.errstate(divide='ignore'): + return 1.0/np.pi/np.sqrt(x*(1-x)) + + def _cdf(self, x): + return 2.0/np.pi*np.arcsin(np.sqrt(x)) + + def _ppf(self, q): + return np.sin(np.pi/2.0*q)**2.0 + + def _stats(self): + mu = 0.5 + mu2 = 1.0/8 + g1 = 0 + g2 = -3.0/2.0 + return mu, mu2, g1, g2 + + def _entropy(self): + return -0.24156447527049044468 + + +arcsine = arcsine_gen(a=0.0, b=1.0, name='arcsine') + + +class FitDataError(ValueError): + """Raised when input data is inconsistent with fixed parameters.""" + # This exception is raised by, for example, beta_gen.fit when both floc + # and fscale are fixed and there are values in the data not in the open + # interval (floc, floc+fscale). + def __init__(self, distr, lower, upper): + self.args = ( + "Invalid values in `data`. Maximum likelihood " + f"estimation with {distr!r} requires that {lower!r} < " + f"(x - loc)/scale < {upper!r} for each x in `data`.", + ) + + +class FitSolverError(FitError): + """ + Raised when a solver fails to converge while fitting a distribution. + """ + # This exception is raised by, for example, beta_gen.fit when + # optimize.fsolve returns with ier != 1. + def __init__(self, mesg): + emsg = "Solver for the MLE equations failed to converge: " + emsg += mesg.replace('\n', '') + self.args = (emsg,) + + +def _beta_mle_a(a, b, n, s1): + # The zeros of this function give the MLE for `a`, with + # `b`, `n` and `s1` given. `s1` is the sum of the logs of + # the data. `n` is the number of data points. + psiab = sc.psi(a + b) + func = s1 - n * (-psiab + sc.psi(a)) + return func + + +def _beta_mle_ab(theta, n, s1, s2): + # Zeros of this function are critical points of + # the maximum likelihood function. Solving this system + # for theta (which contains a and b) gives the MLE for a and b + # given `n`, `s1` and `s2`. `s1` is the sum of the logs of the data, + # and `s2` is the sum of the logs of 1 - data. `n` is the number + # of data points. + a, b = theta + psiab = sc.psi(a + b) + func = [s1 - n * (-psiab + sc.psi(a)), + s2 - n * (-psiab + sc.psi(b))] + return func + + +class beta_gen(rv_continuous): + r"""A beta continuous random variable. + + %(before_notes)s + + Notes + ----- + The probability density function for `beta` is: + + .. math:: + + f(x, a, b) = \frac{\Gamma(a+b) x^{a-1} (1-x)^{b-1}} + {\Gamma(a) \Gamma(b)} + + for :math:`0 <= x <= 1`, :math:`a > 0`, :math:`b > 0`, where + :math:`\Gamma` is the gamma function (`scipy.special.gamma`). + + `beta` takes :math:`a` and :math:`b` as shape parameters. + + This distribution uses routines from the Boost Math C++ library for + the computation of the ``pdf``, ``cdf``, ``ppf``, ``sf`` and ``isf`` + methods. [1]_ + + Maximum likelihood estimates of parameters are only available when the location and + scale are fixed. When either of these parameters is free, ``beta.fit`` resorts to + numerical optimization, but this problem is unbounded: the location and scale may be + chosen to make the minimum and maximum elements of the data coincide with the + endpoints of the support, and the shape parameters may be chosen to make the PDF at + these points infinite. For best results, pass ``floc`` and ``fscale`` keyword + arguments to fix the location and scale, or use `scipy.stats.fit` with + ``method='mse'``. + + %(after_notes)s + + References + ---------- + .. [1] The Boost Developers. "Boost C++ Libraries". https://www.boost.org/. + + %(example)s + + """ + def _shape_info(self): + ia = _ShapeInfo("a", False, (0, np.inf), (False, False)) + ib = _ShapeInfo("b", False, (0, np.inf), (False, False)) + return [ia, ib] + + def _rvs(self, a, b, size=None, random_state=None): + return random_state.beta(a, b, size) + + def _pdf(self, x, a, b): + # gamma(a+b) * x**(a-1) * (1-x)**(b-1) + # beta.pdf(x, a, b) = ------------------------------------ + # gamma(a)*gamma(b) + with np.errstate(over='ignore'): + return scu._beta_pdf(x, a, b) + + def _logpdf(self, x, a, b): + lPx = sc.xlog1py(b - 1.0, -x) + sc.xlogy(a - 1.0, x) + lPx -= sc.betaln(a, b) + return lPx + + def _cdf(self, x, a, b): + return sc.betainc(a, b, x) + + def _sf(self, x, a, b): + return sc.betaincc(a, b, x) + + def _isf(self, x, a, b): + return sc.betainccinv(a, b, x) + + def _ppf(self, q, a, b): + return scu._beta_ppf(q, a, b) + + def _stats(self, a, b): + a_plus_b = a + b + _beta_mean = a/a_plus_b + _beta_variance = a*b / (a_plus_b**2 * (a_plus_b + 1)) + _beta_skewness = ((2 * (b - a) * np.sqrt(a_plus_b + 1)) / + ((a_plus_b + 2) * np.sqrt(a * b))) + _beta_kurtosis_excess_n = 6 * ((a - b)**2 * (a_plus_b + 1) - + a * b * (a_plus_b + 2)) + _beta_kurtosis_excess_d = a * b * (a_plus_b + 2) * (a_plus_b + 3) + _beta_kurtosis_excess = _beta_kurtosis_excess_n / _beta_kurtosis_excess_d + return ( + _beta_mean, + _beta_variance, + _beta_skewness, + _beta_kurtosis_excess) + + def _fitstart(self, data): + if isinstance(data, CensoredData): + data = data._uncensor() + + g1 = _skew(data) + g2 = _kurtosis(data) + + def func(x): + a, b = x + sk = 2*(b-a)*np.sqrt(a + b + 1) / (a + b + 2) / np.sqrt(a*b) + ku = a**3 - a**2*(2*b-1) + b**2*(b+1) - 2*a*b*(b+2) + ku /= a*b*(a+b+2)*(a+b+3) + ku *= 6 + return [sk-g1, ku-g2] + a, b = optimize.fsolve(func, (1.0, 1.0)) + return super()._fitstart(data, args=(a, b)) + + @_call_super_mom + @extend_notes_in_docstring(rv_continuous, notes="""\ + In the special case where `method="MLE"` and + both `floc` and `fscale` are given, a + `ValueError` is raised if any value `x` in `data` does not satisfy + `floc < x < floc + fscale`.\n\n""") + def fit(self, data, *args, **kwds): + # Override rv_continuous.fit, so we can more efficiently handle the + # case where floc and fscale are given. + + floc = kwds.get('floc', None) + fscale = kwds.get('fscale', None) + + if floc is None or fscale is None: + # do general fit + return super().fit(data, *args, **kwds) + + # We already got these from kwds, so just pop them. + kwds.pop('floc', None) + kwds.pop('fscale', None) + + f0 = _get_fixed_fit_value(kwds, ['f0', 'fa', 'fix_a']) + f1 = _get_fixed_fit_value(kwds, ['f1', 'fb', 'fix_b']) + + _remove_optimizer_parameters(kwds) + + if f0 is not None and f1 is not None: + # This check is for consistency with `rv_continuous.fit`. + raise ValueError("All parameters fixed. There is nothing to " + "optimize.") + + # Special case: loc and scale are constrained, so we are fitting + # just the shape parameters. This can be done much more efficiently + # than the method used in `rv_continuous.fit`. (See the subsection + # "Two unknown parameters" in the section "Maximum likelihood" of + # the Wikipedia article on the Beta distribution for the formulas.) + + if not np.isfinite(data).all(): + raise ValueError("The data contains non-finite values.") + + # Normalize the data to the interval [0, 1]. + data = (np.ravel(data) - floc) / fscale + if np.any(data <= 0) or np.any(data >= 1): + raise FitDataError("beta", lower=floc, upper=floc + fscale) + + xbar = data.mean() + + if f0 is not None or f1 is not None: + # One of the shape parameters is fixed. + + if f0 is not None: + # The shape parameter a is fixed, so swap the parameters + # and flip the data. We always solve for `a`. The result + # will be swapped back before returning. + b = f0 + data = 1 - data + xbar = 1 - xbar + else: + b = f1 + + # Initial guess for a. Use the formula for the mean of the beta + # distribution, E[x] = a / (a + b), to generate a reasonable + # starting point based on the mean of the data and the given + # value of b. + a = b * xbar / (1 - xbar) + + # Compute the MLE for `a` by solving _beta_mle_a. + theta, info, ier, mesg = optimize.fsolve( + _beta_mle_a, a, + args=(b, len(data), np.log(data).sum()), + full_output=True + ) + if ier != 1: + raise FitSolverError(mesg=mesg) + a = theta[0] + + if f0 is not None: + # The shape parameter a was fixed, so swap back the + # parameters. + a, b = b, a + + else: + # Neither of the shape parameters is fixed. + + # s1 and s2 are used in the extra arguments passed to _beta_mle_ab + # by optimize.fsolve. + s1 = np.log(data).sum() + s2 = sc.log1p(-data).sum() + + # Use the "method of moments" to estimate the initial + # guess for a and b. + fac = xbar * (1 - xbar) / data.var(ddof=0) - 1 + a = xbar * fac + b = (1 - xbar) * fac + + # Compute the MLE for a and b by solving _beta_mle_ab. + theta, info, ier, mesg = optimize.fsolve( + _beta_mle_ab, [a, b], + args=(len(data), s1, s2), + full_output=True + ) + if ier != 1: + raise FitSolverError(mesg=mesg) + a, b = theta + + return a, b, floc, fscale + + def _entropy(self, a, b): + def regular(a, b): + return (sc.betaln(a, b) - (a - 1) * sc.psi(a) - + (b - 1) * sc.psi(b) + (a + b - 2) * sc.psi(a + b)) + + def asymptotic_ab_large(a, b): + sum_ab = a + b + log_term = 0.5 * ( + np.log(2*np.pi) + np.log(a) + np.log(b) - 3*np.log(sum_ab) + 1 + ) + t1 = 110/sum_ab + 20*sum_ab**-2.0 + sum_ab**-3.0 - 2*sum_ab**-4.0 + t2 = -50/a - 10*a**-2.0 - a**-3.0 + a**-4.0 + t3 = -50/b - 10*b**-2.0 - b**-3.0 + b**-4.0 + return log_term + (t1 + t2 + t3) / 120 + + def asymptotic_b_large(a, b): + sum_ab = a + b + t1 = sc.gammaln(a) - (a - 1) * sc.psi(a) + t2 = ( + - 1/(2*b) + 1/(12*b) - b**-2.0/12 - b**-3.0/120 + b**-4.0/120 + + b**-5.0/252 - b**-6.0/252 + 1/sum_ab - 1/(12*sum_ab) + + sum_ab**-2.0/6 + sum_ab**-3.0/120 - sum_ab**-4.0/60 + - sum_ab**-5.0/252 + sum_ab**-6.0/126 + ) + log_term = sum_ab*np.log1p(a/b) + np.log(b) - 2*np.log(sum_ab) + return t1 + t2 + log_term + + def asymptotic_a_large(a, b): + return asymptotic_b_large(b, a) + + def threshold_large(v): + j = np.floor(np.log10(v)) + d = np.floor(v / 10 ** j) + 2 + return xpx.apply_where(v != 1.0, (d, j), lambda d_, j_: d_ * 10**(7 + j_), + fill_value=1000) + + threshold_a = threshold_large(a) + threshold_b = threshold_large(b) + return _lazyselect([(a >= 4.96e6) & (b >= 4.96e6), + (a <= 4.9e6) & (b - a >= 1e6) & (b >= threshold_a), + (b <= 4.9e6) & (a - b >= 1e6) & (a >= threshold_b), + (a < 4.9e6) & (b < 4.9e6) + ], + [asymptotic_ab_large, asymptotic_b_large, + asymptotic_a_large, regular], + [a, b] + ) + + +beta = beta_gen(a=0.0, b=1.0, name='beta') + + +class betaprime_gen(rv_continuous): + r"""A beta prime continuous random variable. + + %(before_notes)s + + Notes + ----- + The probability density function for `betaprime` is: + + .. math:: + + f(x, a, b) = \frac{x^{a-1} (1+x)^{-a-b}}{\beta(a, b)} + + for :math:`x >= 0`, :math:`a > 0`, :math:`b > 0`, where + :math:`\beta(a, b)` is the beta function (see `scipy.special.beta`). + + `betaprime` takes ``a`` and ``b`` as shape parameters. + + The distribution is related to the `beta` distribution as follows: + If :math:`X` follows a beta distribution with parameters :math:`a, b`, + then :math:`Y = X/(1-X)` has a beta prime distribution with + parameters :math:`a, b` ([1]_). + + The beta prime distribution is a reparametrized version of the + F distribution. The beta prime distribution with shape parameters + ``a`` and ``b`` and ``scale = s`` is equivalent to the F distribution + with parameters ``d1 = 2*a``, ``d2 = 2*b`` and ``scale = (a/b)*s``. + For example, + + >>> from scipy.stats import betaprime, f + >>> x = [1, 2, 5, 10] + >>> a = 12 + >>> b = 5 + >>> betaprime.pdf(x, a, b, scale=2) + array([0.00541179, 0.08331299, 0.14669185, 0.03150079]) + >>> f.pdf(x, 2*a, 2*b, scale=(a/b)*2) + array([0.00541179, 0.08331299, 0.14669185, 0.03150079]) + + %(after_notes)s + + References + ---------- + .. [1] Beta prime distribution, Wikipedia, + https://en.wikipedia.org/wiki/Beta_prime_distribution + + %(example)s + + """ + _support_mask = rv_continuous._open_support_mask + + def _shape_info(self): + ia = _ShapeInfo("a", False, (0, np.inf), (False, False)) + ib = _ShapeInfo("b", False, (0, np.inf), (False, False)) + return [ia, ib] + + def _rvs(self, a, b, size=None, random_state=None): + u1 = gamma.rvs(a, size=size, random_state=random_state) + u2 = gamma.rvs(b, size=size, random_state=random_state) + return u1 / u2 + + def _pdf(self, x, a, b): + # betaprime.pdf(x, a, b) = x**(a-1) * (1+x)**(-a-b) / beta(a, b) + return np.exp(self._logpdf(x, a, b)) + + def _logpdf(self, x, a, b): + return sc.xlogy(a - 1.0, x) - sc.xlog1py(a + b, x) - sc.betaln(a, b) + + def _cdf(self, x, a, b): + # note: f2 is the direct way to compute the cdf if the relationship + # to the beta distribution is used. + # however, for very large x, x/(1+x) == 1. since the distribution + # has very fat tails if b is small, this can cause inaccurate results + # use the following relationship of the incomplete beta function: + # betainc(x, a, b) = 1 - betainc(1-x, b, a) + # see gh-17631 + return xpx.apply_where( + x > 1, (x, a, b), + lambda x_, a_, b_: beta._sf(1 / (1 + x_), b_, a_), + lambda x_, a_, b_: beta._cdf(x_ / (1 + x_), a_, b_)) + + def _sf(self, x, a, b): + return xpx.apply_where( + x > 1, (x, a, b), + lambda x_, a_, b_: beta._cdf(1 / (1 + x_), b_, a_), + lambda x_, a_, b_: beta._sf(x_ / (1 + x_), a_, b_)) + + def _ppf(self, p, a, b): + p, a, b = np.broadcast_arrays(p, a, b) + # By default, compute the ppf by solving the following: + # p = beta._cdf(x/(1+x), a, b). This implies x = r/(1-r) with + # r = beta._ppf(p, a, b). This can cause numerical issues if r is + # very close to 1. In that case, invert the alternative expression of + # the cdf: p = beta._sf(1/(1+x), b, a). + r = stats.beta._ppf(p, a, b) + with np.errstate(divide='ignore'): + out = r / (1 - r) + rnear1 = r > 0.9999 + if np.isscalar(r): + if rnear1: + out = 1/stats.beta._isf(p, b, a) - 1 + else: + out[rnear1] = 1/stats.beta._isf(p[rnear1], b[rnear1], a[rnear1]) - 1 + return out + + def _munp(self, n, a, b): + return xpx.apply_where( + b > n, (a, b), + lambda a, b: np.prod([(a+i-1)/(b-i) for i in range(1, int(n)+1)], axis=0), + fill_value=np.inf) + + +betaprime = betaprime_gen(a=0.0, name='betaprime') + + +class bradford_gen(rv_continuous): + r"""A Bradford continuous random variable. + + %(before_notes)s + + Notes + ----- + The probability density function for `bradford` is: + + .. math:: + + f(x, c) = \frac{c}{\log(1+c) (1+cx)} + + for :math:`0 <= x <= 1` and :math:`c > 0`. + + `bradford` takes ``c`` as a shape parameter for :math:`c`. + + %(after_notes)s + + %(example)s + + """ + def _shape_info(self): + return [_ShapeInfo("c", False, (0, np.inf), (False, False))] + + def _pdf(self, x, c): + # bradford.pdf(x, c) = c / (k * (1+c*x)) + return c / (c*x + 1.0) / sc.log1p(c) + + def _cdf(self, x, c): + return sc.log1p(c*x) / sc.log1p(c) + + def _ppf(self, q, c): + return sc.expm1(q * sc.log1p(c)) / c + + def _stats(self, c, moments='mv'): + k = np.log(1.0+c) + mu = (c-k)/(c*k) + mu2 = ((c+2.0)*k-2.0*c)/(2*c*k*k) + g1 = None + g2 = None + if 's' in moments: + g1 = np.sqrt(2)*(12*c*c-9*c*k*(c+2)+2*k*k*(c*(c+3)+3)) + g1 /= np.sqrt(c*(c*(k-2)+2*k))*(3*c*(k-2)+6*k) + if 'k' in moments: + g2 = (c**3*(k-3)*(k*(3*k-16)+24)+12*k*c*c*(k-4)*(k-3) + + 6*c*k*k*(3*k-14) + 12*k**3) + g2 /= 3*c*(c*(k-2)+2*k)**2 + return mu, mu2, g1, g2 + + def _entropy(self, c): + k = np.log(1+c) + return k/2.0 - np.log(c/k) + + +bradford = bradford_gen(a=0.0, b=1.0, name='bradford') + + +class burr_gen(rv_continuous): + r"""A Burr (Type III) continuous random variable. + + %(before_notes)s + + See Also + -------- + fisk : a special case of either `burr` or `burr12` with ``d=1`` + burr12 : Burr Type XII distribution + mielke : Mielke Beta-Kappa / Dagum distribution + + Notes + ----- + The probability density function for `burr` is: + + .. math:: + + f(x; c, d) = c d \frac{x^{-c - 1}} + {{(1 + x^{-c})}^{d + 1}} + + for :math:`x >= 0` and :math:`c, d > 0`. + + `burr` takes ``c`` and ``d`` as shape parameters for :math:`c` and + :math:`d`. + + This is the PDF corresponding to the third CDF given in Burr's list; + specifically, it is equation (11) in Burr's paper [1]_. The distribution + is also commonly referred to as the Dagum distribution [2]_. If the + parameter :math:`c < 1` then the mean of the distribution does not + exist and if :math:`c < 2` the variance does not exist [2]_. + The PDF is finite at the left endpoint :math:`x = 0` if :math:`c * d >= 1`. + + %(after_notes)s + + References + ---------- + .. [1] Burr, I. W. "Cumulative frequency functions", Annals of + Mathematical Statistics, 13(2), pp 215-232 (1942). + .. [2] https://en.wikipedia.org/wiki/Dagum_distribution + .. [3] Kleiber, Christian. "A guide to the Dagum distributions." + Modeling Income Distributions and Lorenz Curves pp 97-117 (2008). + + %(example)s + + """ + # Do not set _support_mask to rv_continuous._open_support_mask + # Whether the left-hand endpoint is suitable for pdf evaluation is dependent + # on the values of c and d: if c*d >= 1, the pdf is finite, otherwise infinite. + + def _shape_info(self): + ic = _ShapeInfo("c", False, (0, np.inf), (False, False)) + id = _ShapeInfo("d", False, (0, np.inf), (False, False)) + return [ic, id] + + def _pdf(self, x, c, d): + # burr.pdf(x, c, d) = c * d * x**(-c-1) * (1+x**(-c))**(-d-1) + output = xpx.apply_where( + x == 0, (x, c, d), + lambda x_, c_, d_: c_ * d_ * (x_**(c_*d_-1)) / (1 + x_**c_), + lambda x_, c_, d_: (c_ * d_ * (x_ ** (-c_ - 1.0)) / + ((1 + x_ ** (-c_)) ** (d_ + 1.0)))) + return output[()] if output.ndim == 0 else output + + def _logpdf(self, x, c, d): + output = xpx.apply_where( + x == 0, (x, c, d), + lambda x_, c_, d_: (np.log(c_) + np.log(d_) + sc.xlogy(c_*d_ - 1, x_) + - (d_+1) * sc.log1p(x_**(c_))), + lambda x_, c_, d_: (np.log(c_) + np.log(d_) + + sc.xlogy(-c_ - 1, x_) + - sc.xlog1py(d_+1, x_**(-c_)))) + return output[()] if output.ndim == 0 else output + + def _cdf(self, x, c, d): + return (1 + x**(-c))**(-d) + + def _logcdf(self, x, c, d): + return sc.log1p(x**(-c)) * (-d) + + def _sf(self, x, c, d): + return np.exp(self._logsf(x, c, d)) + + def _logsf(self, x, c, d): + return np.log1p(- (1 + x**(-c))**(-d)) + + def _ppf(self, q, c, d): + return (q**(-1.0/d) - 1)**(-1.0/c) + + def _isf(self, q, c, d): + _q = sc.xlog1py(-1.0 / d, -q) + return sc.expm1(_q) ** (-1.0 / c) + + def _stats(self, c, d): + nc = np.arange(1, 5).reshape(4,1) / c + # ek is the kth raw moment, e1 is the mean e2-e1**2 variance etc. + e1, e2, e3, e4 = sc.beta(d + nc, 1. - nc) * d + mu = np.where(c > 1.0, e1, np.nan) + mu2_if_c = e2 - mu**2 + mu2 = np.where(c > 2.0, mu2_if_c, np.nan) + g1 = xpx.apply_where( + c > 3.0, (e1, e2, e3, mu2_if_c), + lambda e1, e2, e3, mu2_if_c: ((e3 - 3*e2*e1 + 2*e1**3) + / np.sqrt((mu2_if_c)**3)), + fill_value=np.nan) + g2 = xpx.apply_where( + c > 4.0, (e1, e2, e3, e4, mu2_if_c), + lambda e1, e2, e3, e4, mu2_if_c: ( + ((e4 - 4*e3*e1 + 6*e2*e1**2 - 3*e1**4) / mu2_if_c**2) - 3), + fill_value=np.nan) + if np.ndim(c) == 0: + return mu.item(), mu2.item(), g1.item(), g2.item() + return mu, mu2, g1, g2 + + def _munp(self, n, c, d): + def __munp(n, c, d): + nc = 1. * n / c + return d * sc.beta(1.0 - nc, d + nc) + n, c, d = np.asarray(n), np.asarray(c), np.asarray(d) + return xpx.apply_where((c > n) & (n == n) & (d == d), + (n, c, d), __munp, fill_value=np.nan) + + +burr = burr_gen(a=0.0, name='burr') + + +class burr12_gen(rv_continuous): + r"""A Burr (Type XII) continuous random variable. + + %(before_notes)s + + See Also + -------- + fisk : a special case of either `burr` or `burr12` with ``d=1`` + burr : Burr Type III distribution + + Notes + ----- + The probability density function for `burr12` is: + + .. math:: + + f(x; c, d) = c d \frac{x^{c-1}} + {(1 + x^c)^{d + 1}} + + for :math:`x >= 0` and :math:`c, d > 0`. + + `burr12` takes ``c`` and ``d`` as shape parameters for :math:`c` + and :math:`d`. + + This is the PDF corresponding to the twelfth CDF given in Burr's list; + specifically, it is equation (20) in Burr's paper [1]_. + + %(after_notes)s + + The Burr type 12 distribution is also sometimes referred to as + the Singh-Maddala distribution from NIST [2]_. + + References + ---------- + .. [1] Burr, I. W. "Cumulative frequency functions", Annals of + Mathematical Statistics, 13(2), pp 215-232 (1942). + + .. [2] https://www.itl.nist.gov/div898/software/dataplot/refman2/auxillar/b12pdf.htm + + .. [3] "Burr distribution", + https://en.wikipedia.org/wiki/Burr_distribution + + %(example)s + + """ + def _shape_info(self): + ic = _ShapeInfo("c", False, (0, np.inf), (False, False)) + id = _ShapeInfo("d", False, (0, np.inf), (False, False)) + return [ic, id] + + def _pdf(self, x, c, d): + # burr12.pdf(x, c, d) = c * d * x**(c-1) * (1+x**(c))**(-d-1) + return np.exp(self._logpdf(x, c, d)) + + def _logpdf(self, x, c, d): + return np.log(c) + np.log(d) + sc.xlogy(c - 1, x) + sc.xlog1py(-d-1, x**c) + + def _cdf(self, x, c, d): + return -sc.expm1(self._logsf(x, c, d)) + + def _logcdf(self, x, c, d): + return sc.log1p(-(1 + x**c)**(-d)) + + def _sf(self, x, c, d): + return np.exp(self._logsf(x, c, d)) + + def _logsf(self, x, c, d): + return sc.xlog1py(-d, x**c) + + def _ppf(self, q, c, d): + # The following is an implementation of + # ((1 - q)**(-1.0/d) - 1)**(1.0/c) + # that does a better job handling small values of q. + return sc.expm1(-1/d * sc.log1p(-q))**(1/c) + + def _isf(self, p, c, d): + return sc.expm1(-1/d * np.log(p))**(1/c) + + def _munp(self, n, c, d): + def moment_if_exists(n, c, d): + nc = 1. * n / c + return d * sc.beta(1.0 + nc, d - nc) + + return xpx.apply_where(c * d > n, (n, c, d), moment_if_exists, + fill_value=np.nan) + + +burr12 = burr12_gen(a=0.0, name='burr12') + + +class fisk_gen(burr_gen): + r"""A Fisk continuous random variable. + + The Fisk distribution is also known as the log-logistic distribution. + + %(before_notes)s + + See Also + -------- + burr + + Notes + ----- + The probability density function for `fisk` is: + + .. math:: + + f(x, c) = \frac{c x^{c-1}} + {(1 + x^c)^2} + + for :math:`x >= 0` and :math:`c > 0`. + + Please note that the above expression can be transformed into the following + one, which is also commonly used: + + .. math:: + + f(x, c) = \frac{c x^{-c-1}} + {(1 + x^{-c})^2} + + `fisk` takes ``c`` as a shape parameter for :math:`c`. + + `fisk` is a special case of `burr` or `burr12` with ``d=1``. + + Suppose ``X`` is a logistic random variable with location ``l`` + and scale ``s``. Then ``Y = exp(X)`` is a Fisk (log-logistic) + random variable with ``scale = exp(l)`` and shape ``c = 1/s``. + + %(after_notes)s + + %(example)s + + """ + def _shape_info(self): + return [_ShapeInfo("c", False, (0, np.inf), (False, False))] + + def _pdf(self, x, c): + # fisk.pdf(x, c) = c * x**(-c-1) * (1 + x**(-c))**(-2) + return burr._pdf(x, c, 1.0) + + def _cdf(self, x, c): + return burr._cdf(x, c, 1.0) + + def _sf(self, x, c): + return burr._sf(x, c, 1.0) + + def _logpdf(self, x, c): + # fisk.pdf(x, c) = c * x**(-c-1) * (1 + x**(-c))**(-2) + return burr._logpdf(x, c, 1.0) + + def _logcdf(self, x, c): + return burr._logcdf(x, c, 1.0) + + def _logsf(self, x, c): + return burr._logsf(x, c, 1.0) + + def _ppf(self, x, c): + return burr._ppf(x, c, 1.0) + + def _isf(self, q, c): + return burr._isf(q, c, 1.0) + + def _munp(self, n, c): + return burr._munp(n, c, 1.0) + + def _stats(self, c): + return burr._stats(c, 1.0) + + def _entropy(self, c): + return 2 - np.log(c) + + +fisk = fisk_gen(a=0.0, name='fisk') + + +class cauchy_gen(rv_continuous): + r"""A Cauchy continuous random variable. + + %(before_notes)s + + Notes + ----- + The probability density function for `cauchy` is + + .. math:: + + f(x) = \frac{1}{\pi (1 + x^2)} + + for a real number :math:`x`. + + This distribution uses routines from the Boost Math C++ library for + the computation of the ``ppf`` and ``isf`` methods. [1]_ + + %(after_notes)s + + References + ---------- + .. [1] The Boost Developers. "Boost C++ Libraries". https://www.boost.org/. + + %(example)s + + """ + def _shape_info(self): + return [] + + def _pdf(self, x): + # cauchy.pdf(x) = 1 / (pi * (1 + x**2)) + with np.errstate(over='ignore'): + return 1.0/np.pi/(1.0+x*x) + + def _logpdf(self, x): + # The formulas + # log(1/(pi*(1 + x**2))) = -log(pi) - log(1 + x**2) + # = -log(pi) - log(x**2*(1 + 1/x**2)) + # = -log(pi) - (2log(|x|) + log1p(1/x**2)) + # are used here. + absx = np.abs(x) + # In the following apply_where, `f1` provides better precision than `f2` + # for small and moderate x, while `f2` avoids the overflow that can + # occur with absx**2. + return xpx.apply_where( + absx < 1, absx, + lambda absx: -_LOG_PI - np.log1p(absx**2), + lambda absx: (-_LOG_PI - (2*np.log(absx) + np.log1p((1/absx)**2)))) + + def _cdf(self, x): + return np.arctan2(1, -x)/np.pi + + def _ppf(self, q): + return scu._cauchy_ppf(q, 0, 1) + + def _sf(self, x): + return np.arctan2(1, x)/np.pi + + def _isf(self, q): + return scu._cauchy_isf(q, 0, 1) + + def _stats(self): + return np.nan, np.nan, np.nan, np.nan + + def _entropy(self): + return np.log(4*np.pi) + + def _fitstart(self, data, args=None): + # Initialize ML guesses using quartiles instead of moments. + if isinstance(data, CensoredData): + data = data._uncensor() + p25, p50, p75 = np.percentile(data, [25, 50, 75]) + return p50, (p75 - p25)/2 + + +cauchy = cauchy_gen(name='cauchy') + + +class chi_gen(rv_continuous): + r"""A chi continuous random variable. + + %(before_notes)s + + Notes + ----- + The probability density function for `chi` is: + + .. math:: + + f(x, k) = \frac{1}{2^{k/2-1} \Gamma \left( k/2 \right)} + x^{k-1} \exp \left( -x^2/2 \right) + + for :math:`x >= 0` and :math:`k > 0` (degrees of freedom, denoted ``df`` + in the implementation). :math:`\Gamma` is the gamma function + (`scipy.special.gamma`). + + Special cases of `chi` are: + + - ``chi(1, loc, scale)`` is equivalent to `halfnorm` + - ``chi(2, 0, scale)`` is equivalent to `rayleigh` + - ``chi(3, 0, scale)`` is equivalent to `maxwell` + + `chi` takes ``df`` as a shape parameter. + + %(after_notes)s + + %(example)s + + """ + def _shape_info(self): + return [_ShapeInfo("df", False, (0, np.inf), (False, False))] + + def _rvs(self, df, size=None, random_state=None): + return np.sqrt(chi2.rvs(df, size=size, random_state=random_state)) + + def _pdf(self, x, df): + # x**(df-1) * exp(-x**2/2) + # chi.pdf(x, df) = ------------------------- + # 2**(df/2-1) * gamma(df/2) + return np.exp(self._logpdf(x, df)) + + def _logpdf(self, x, df): + l = np.log(2) - .5*np.log(2)*df - sc.gammaln(.5*df) + return l + sc.xlogy(df - 1., x) - .5*x**2 + + def _cdf(self, x, df): + return sc.gammainc(.5*df, .5*x**2) + + def _sf(self, x, df): + return sc.gammaincc(.5*df, .5*x**2) + + def _ppf(self, q, df): + return np.sqrt(2*sc.gammaincinv(.5*df, q)) + + def _isf(self, q, df): + return np.sqrt(2*sc.gammainccinv(.5*df, q)) + + def _stats(self, df): + # poch(df/2, 1/2) = gamma(df/2 + 1/2) / gamma(df/2) + mu = np.sqrt(2) * sc.poch(0.5 * df, 0.5) + mu2 = df - mu*mu + g1 = (2*mu**3.0 + mu*(1-2*df))/np.asarray(np.power(mu2, 1.5)) + g2 = 2*df*(1.0-df)-6*mu**4 + 4*mu**2 * (2*df-1) + g2 /= np.asarray(mu2**2.0) + return mu, mu2, g1, g2 + + def _entropy(self, df): + + def regular_formula(df): + return (sc.gammaln(.5 * df) + + 0.5 * (df - np.log(2) - (df - 1) * sc.digamma(0.5 * df))) + + def asymptotic_formula(df): + return (0.5 + np.log(np.pi)/2 - (df**-1)/6 - (df**-2)/6 + - 4/45*(df**-3) + (df**-4)/15) + + return xpx.apply_where(df < 300, df, regular_formula, asymptotic_formula) + + +chi = chi_gen(a=0.0, name='chi') + + +class chi2_gen(rv_continuous): + r"""A chi-squared continuous random variable. + + For the noncentral chi-square distribution, see `ncx2`. + + %(before_notes)s + + See Also + -------- + ncx2 + + Notes + ----- + The probability density function for `chi2` is: + + .. math:: + + f(x, k) = \frac{1}{2^{k/2} \Gamma \left( k/2 \right)} + x^{k/2-1} \exp \left( -x/2 \right) + + for :math:`x > 0` and :math:`k > 0` (degrees of freedom, denoted ``df`` + in the implementation). + + `chi2` takes ``df`` as a shape parameter. + + The chi-squared distribution is a special case of the gamma + distribution, with gamma parameters ``a = df/2``, ``loc = 0`` and + ``scale = 2``. + + %(after_notes)s + + %(example)s + + """ + def _shape_info(self): + return [_ShapeInfo("df", False, (0, np.inf), (False, False))] + + def _rvs(self, df, size=None, random_state=None): + return random_state.chisquare(df, size) + + def _pdf(self, x, df): + # chi2.pdf(x, df) = 1 / (2*gamma(df/2)) * (x/2)**(df/2-1) * exp(-x/2) + return np.exp(self._logpdf(x, df)) + + def _logpdf(self, x, df): + return sc.xlogy(df/2.-1, x) - x/2. - sc.gammaln(df/2.) - (np.log(2)*df)/2. + + def _cdf(self, x, df): + return sc.chdtr(df, x) + + def _sf(self, x, df): + return sc.chdtrc(df, x) + + def _isf(self, p, df): + return sc.chdtri(df, p) + + def _ppf(self, p, df): + return 2*sc.gammaincinv(df/2, p) + + def _stats(self, df): + mu = df + mu2 = 2*df + g1 = 2*np.sqrt(2.0/df) + g2 = 12.0/df + return mu, mu2, g1, g2 + + def _entropy(self, df): + half_df = 0.5 * df + + def regular_formula(half_df): + return (half_df + np.log(2) + sc.gammaln(half_df) + + (1 - half_df) * sc.psi(half_df)) + + def asymptotic_formula(half_df): + # plug in the above formula the following asymptotic + # expansions: + # ln(gamma(a)) ~ (a - 0.5) * ln(a) - a + 0.5 * ln(2 * pi) + + # 1/(12 * a) - 1/(360 * a**3) + # psi(a) ~ ln(a) - 1/(2 * a) - 1/(3 * a**2) + 1/120 * a**4) + c = np.log(2) + 0.5*(1 + np.log(2*np.pi)) + h = 0.5/half_df + return (h*(-2/3 + h*(-1/3 + h*(-4/45 + h/7.5))) + + 0.5*np.log(half_df) + c) + + return xpx.apply_where(half_df < 125, half_df, + regular_formula, asymptotic_formula) + + +chi2 = chi2_gen(a=0.0, name='chi2') + + +class cosine_gen(rv_continuous): + r"""A cosine continuous random variable. + + %(before_notes)s + + Notes + ----- + The cosine distribution is an approximation to the normal distribution. + The probability density function for `cosine` is: + + .. math:: + + f(x) = \frac{1}{2\pi} (1+\cos(x)) + + for :math:`-\pi \le x \le \pi`. + + %(after_notes)s + + %(example)s + + """ + def _shape_info(self): + return [] + + def _pdf(self, x): + # cosine.pdf(x) = 1/(2*pi) * (1+cos(x)) + return 1.0/2/np.pi*(1+np.cos(x)) + + def _logpdf(self, x): + c = np.cos(x) + return xpx.apply_where(c != -1, c, + lambda c: np.log1p(c) - np.log(2*np.pi), + fill_value=-np.inf) + + def _cdf(self, x): + return scu._cosine_cdf(x) + + def _sf(self, x): + return scu._cosine_cdf(-x) + + def _ppf(self, p): + return scu._cosine_invcdf(p) + + def _isf(self, p): + return -scu._cosine_invcdf(p) + + def _stats(self): + v = (np.pi * np.pi / 3.0) - 2.0 + k = -6.0 * (np.pi**4 - 90) / (5.0 * (np.pi * np.pi - 6)**2) + return 0.0, v, 0.0, k + + def _entropy(self): + return np.log(4*np.pi)-1.0 + + +cosine = cosine_gen(a=-np.pi, b=np.pi, name='cosine') + + +class dgamma_gen(rv_continuous): + r"""A double gamma continuous random variable. + + The double gamma distribution is also known as the reflected gamma + distribution [1]_. + + %(before_notes)s + + Notes + ----- + The probability density function for `dgamma` is: + + .. math:: + + f(x, a) = \frac{1}{2\Gamma(a)} |x|^{a-1} \exp(-|x|) + + for a real number :math:`x` and :math:`a > 0`. :math:`\Gamma` is the + gamma function (`scipy.special.gamma`). + + `dgamma` takes ``a`` as a shape parameter for :math:`a`. + + %(after_notes)s + + References + ---------- + .. [1] Johnson, Kotz, and Balakrishnan, "Continuous Univariate + Distributions, Volume 1", Second Edition, John Wiley and Sons + (1994). + + %(example)s + + """ + def _shape_info(self): + return [_ShapeInfo("a", False, (0, np.inf), (False, False))] + + def _rvs(self, a, size=None, random_state=None): + u = random_state.uniform(size=size) + gm = gamma.rvs(a, size=size, random_state=random_state) + return gm * np.where(u >= 0.5, 1, -1) + + def _pdf(self, x, a): + # dgamma.pdf(x, a) = 1 / (2*gamma(a)) * abs(x)**(a-1) * exp(-abs(x)) + ax = abs(x) + return 1.0/(2*sc.gamma(a))*ax**(a-1.0) * np.exp(-ax) + + def _logpdf(self, x, a): + ax = abs(x) + return sc.xlogy(a - 1.0, ax) - ax - np.log(2) - sc.gammaln(a) + + def _cdf(self, x, a): + return np.where(x > 0, + 0.5 + 0.5*sc.gammainc(a, x), + 0.5*sc.gammaincc(a, -x)) + + def _sf(self, x, a): + return np.where(x > 0, + 0.5*sc.gammaincc(a, x), + 0.5 + 0.5*sc.gammainc(a, -x)) + + def _entropy(self, a): + return stats.gamma._entropy(a) - np.log(0.5) + + def _ppf(self, q, a): + return np.where(q > 0.5, + sc.gammaincinv(a, 2*q - 1), + -sc.gammainccinv(a, 2*q)) + + def _isf(self, q, a): + return np.where(q > 0.5, + -sc.gammaincinv(a, 2*q - 1), + sc.gammainccinv(a, 2*q)) + + def _stats(self, a): + mu2 = a*(a+1.0) + return 0.0, mu2, 0.0, (a+2.0)*(a+3.0)/mu2-3.0 + + +dgamma = dgamma_gen(name='dgamma') + + +class dpareto_lognorm_gen(rv_continuous): + r"""A double Pareto lognormal continuous random variable. + + %(before_notes)s + + Notes + ----- + The probability density function for `dpareto_lognorm` is: + + .. math:: + + f(x, \mu, \sigma, \alpha, \beta) = + \frac{\alpha \beta}{(\alpha + \beta) x} + \phi\left( \frac{\log x - \mu}{\sigma} \right) + \left( R(y_1) + R(y_2) \right) + + where :math:`R(t) = \frac{1 - \Phi(t)}{\phi(t)}`, + :math:`\phi` and :math:`\Phi` are the normal PDF and CDF, respectively, + :math:`y_1 = \alpha \sigma - \frac{\log x - \mu}{\sigma}`, + and :math:`y_2 = \beta \sigma + \frac{\log x - \mu}{\sigma}` + for real numbers :math:`x` and :math:`\mu`, :math:`\sigma > 0`, + :math:`\alpha > 0`, and :math:`\beta > 0` [1]_. + + `dpareto_lognorm` takes + ``u`` as a shape parameter for :math:`\mu`, + ``s`` as a shape parameter for :math:`\sigma`, + ``a`` as a shape parameter for :math:`\alpha`, and + ``b`` as a shape parameter for :math:`\beta`. + + A random variable :math:`X` distributed according to the PDF above + can be represented as :math:`X = U \frac{V_1}{V_2}` where :math:`U`, + :math:`V_1`, and :math:`V_2` are independent, :math:`U` is lognormally + distributed such that :math:`\log U \sim N(\mu, \sigma^2)`, and + :math:`V_1` and :math:`V_2` follow Pareto distributions with parameters + :math:`\alpha` and :math:`\beta`, respectively [2]_. + + %(after_notes)s + + References + ---------- + .. [1] Hajargasht, Gholamreza, and William E. Griffiths. "Pareto-lognormal + distributions: Inequality, poverty, and estimation from grouped income + data." Economic Modelling 33 (2013): 593-604. + .. [2] Reed, William J., and Murray Jorgensen. "The double Pareto-lognormal + distribution - a new parametric model for size distributions." + Communications in Statistics - Theory and Methods 33.8 (2004): 1733-1753. + + %(example)s + + """ + _logphi = norm._logpdf + _logPhi = norm._logcdf + _logPhic = norm._logsf + _phi = norm._pdf + _Phi = norm._cdf + _Phic = norm._sf + + def _R(self, z): + return self._Phic(z) / self._phi(z) + + def _logR(self, z): + return self._logPhic(z) - self._logphi(z) + + def _shape_info(self): + return [_ShapeInfo("u", False, (-np.inf, np.inf), (False, False)), + _ShapeInfo("s", False, (0, np.inf), (False, False)), + _ShapeInfo("a", False, (0, np.inf), (False, False)), + _ShapeInfo("b", False, (0, np.inf), (False, False))] + + def _argcheck(self, u, s, a, b): + return (s > 0) & (a > 0) & (b > 0) + + def _rvs(self, u, s, a, b, size=None, random_state=None): + # From [1] after Equation (12): "To generate pseudo-random + # deviates from the dPlN distribution, one can exponentiate + # pseudo-random deviates from NL generated using (6)." + Z = random_state.normal(u, s, size=size) + E1 = random_state.standard_exponential(size=size) + E2 = random_state.standard_exponential(size=size) + return np.exp(Z + E1 / a - E2 / b) + + def _logpdf(self, x, u, s, a, b): + with np.errstate(invalid='ignore', divide='ignore'): + log_y, m = np.log(x), u # compare against [1] Eq. 1 + z = (log_y - m) / s + x1 = a * s - z + x2 = b * s + z + out = np.asarray(np.log(a) + np.log(b) - np.log(a + b) - log_y) + out += self._logphi(z) + out += np.logaddexp(self._logR(x1), self._logR(x2)) + out[(x == 0) | np.isinf(x)] = -np.inf + return out[()] + + def _logcdf(self, x, u, s, a, b): + with np.errstate(invalid='ignore', divide='ignore'): + log_y, m = np.log(x), u # compare against [1] Eq. 2 + z = (log_y - m) / s + x1 = a * s - z + x2 = b * s + z + t1 = self._logPhi(z) + t2 = self._logphi(z) + t3 = (np.log(b) + self._logR(x1)) + t4 = (np.log(a) + self._logR(x2)) + t1, t2, t3, t4, one = np.broadcast_arrays(t1, t2, t3, t4, 1) + # t3 can be smaller than t4, so we have to consider log of negative number + # This would be much simpler, but `return_sign` is available, so use it? + # t5 = sc.logsumexp([t3, t4 + np.pi*1j]) + t5, sign = sc.logsumexp([t3, t4], b=[one, -one], axis=0, return_sign=True) + temp = [t1, t2 + t5 - np.log(a + b)] + out = np.asarray(sc.logsumexp(temp, b=[one, -one*sign], axis=0)) + out[x == 0] = -np.inf + return out[()] + + def _logsf(self, x, u, s, a, b): + return scu._log1mexp(self._logcdf(x, u, s, a, b)) + + # Infrastructure doesn't seem to do this, so... + + def _pdf(self, x, u, s, a, b): + return np.exp(self._logpdf(x, u, s, a, b)) + + def _cdf(self, x, u, s, a, b): + return np.exp(self._logcdf(x, u, s, a, b)) + + def _sf(self, x, u, s, a, b): + return np.exp(self._logsf(x, u, s, a, b)) + + def _munp(self, n, u, s, a, b): + m, k = u, float(n) # compare against [1] Eq. 6 + out = (a * b) / ((a - k) * (b + k)) * np.exp(k * m + k ** 2 * s ** 2 / 2) + out = np.asarray(out) + out[a <= k] = np.nan + return out + + +dpareto_lognorm = dpareto_lognorm_gen(a=0, name='dpareto_lognorm') + + +class dweibull_gen(rv_continuous): + r"""A double Weibull continuous random variable. + + %(before_notes)s + + Notes + ----- + The probability density function for `dweibull` is given by + + .. math:: + + f(x, c) = c / 2 |x|^{c-1} \exp(-|x|^c) + + for a real number :math:`x` and :math:`c > 0`. + + `dweibull` takes ``c`` as a shape parameter for :math:`c`. + + %(after_notes)s + + %(example)s + + """ + def _shape_info(self): + return [_ShapeInfo("c", False, (0, np.inf), (False, False))] + + def _rvs(self, c, size=None, random_state=None): + u = random_state.uniform(size=size) + w = weibull_min.rvs(c, size=size, random_state=random_state) + return w * (np.where(u >= 0.5, 1, -1)) + + def _pdf(self, x, c): + # dweibull.pdf(x, c) = c / 2 * abs(x)**(c-1) * exp(-abs(x)**c) + ax = abs(x) + Px = c / 2.0 * ax**(c-1.0) * np.exp(-ax**c) + return Px + + def _logpdf(self, x, c): + ax = abs(x) + return np.log(c) - np.log(2.0) + sc.xlogy(c - 1.0, ax) - ax**c + + def _cdf(self, x, c): + Cx1 = 0.5 * np.exp(-abs(x)**c) + return np.where(x > 0, 1 - Cx1, Cx1) + + def _ppf(self, q, c): + fac = 2. * np.where(q <= 0.5, q, 1. - q) + fac = np.power(-np.log(fac), 1.0 / c) + return np.where(q > 0.5, fac, -fac) + + def _sf(self, x, c): + half_weibull_min_sf = 0.5 * stats.weibull_min._sf(np.abs(x), c) + return np.where(x > 0, half_weibull_min_sf, 1 - half_weibull_min_sf) + + def _isf(self, q, c): + double_q = 2. * np.where(q <= 0.5, q, 1. - q) + weibull_min_isf = stats.weibull_min._isf(double_q, c) + return np.where(q > 0.5, -weibull_min_isf, weibull_min_isf) + + def _munp(self, n, c): + return (1 - (n % 2)) * sc.gamma(1.0 + 1.0 * n / c) + + # since we know that all odd moments are zeros, return them at once. + # returning Nones from _stats makes the public stats call _munp + # so overall we're saving one or two gamma function evaluations here. + def _stats(self, c): + return 0, None, 0, None + + def _entropy(self, c): + h = stats.weibull_min._entropy(c) - np.log(0.5) + return h + + +dweibull = dweibull_gen(name='dweibull') + + +class expon_gen(rv_continuous): + r"""An exponential continuous random variable. + + %(before_notes)s + + Notes + ----- + The probability density function for `expon` is: + + .. math:: + + f(x) = \exp(-x) + + for :math:`x \ge 0`. + + %(after_notes)s + + A common parameterization for `expon` is in terms of the rate parameter + ``lambda``, such that ``pdf = lambda * exp(-lambda * x)``. This + parameterization corresponds to using ``scale = 1 / lambda``. + + The exponential distribution is a special case of the gamma + distributions, with gamma shape parameter ``a = 1``. + + %(example)s + + """ + def _shape_info(self): + return [] + + def _rvs(self, size=None, random_state=None): + return random_state.standard_exponential(size) + + def _pdf(self, x): + # expon.pdf(x) = exp(-x) + return np.exp(-x) + + def _logpdf(self, x): + return -x + + def _cdf(self, x): + return -sc.expm1(-x) + + def _ppf(self, q): + return -sc.log1p(-q) + + def _sf(self, x): + return np.exp(-x) + + def _logsf(self, x): + return -x + + def _isf(self, q): + return -np.log(q) + + def _stats(self): + return 1.0, 1.0, 2.0, 6.0 + + def _entropy(self): + return 1.0 + + @_call_super_mom + @replace_notes_in_docstring(rv_continuous, notes="""\ + When `method='MLE'`, + this function uses explicit formulas for the maximum likelihood + estimation of the exponential distribution parameters, so the + `optimizer`, `loc` and `scale` keyword arguments are + ignored.\n\n""") + def fit(self, data, *args, **kwds): + if len(args) > 0: + raise TypeError("Too many arguments.") + + floc = kwds.pop('floc', None) + fscale = kwds.pop('fscale', None) + + _remove_optimizer_parameters(kwds) + + if floc is not None and fscale is not None: + # This check is for consistency with `rv_continuous.fit`. + raise ValueError("All parameters fixed. There is nothing to " + "optimize.") + + data = np.asarray(data) + + if not np.isfinite(data).all(): + raise ValueError("The data contains non-finite values.") + + data_min = data.min() + + if floc is None: + # ML estimate of the location is the minimum of the data. + loc = data_min + else: + loc = floc + if data_min < loc: + # There are values that are less than the specified loc. + raise FitDataError("expon", lower=floc, upper=np.inf) + + if fscale is None: + # ML estimate of the scale is the shifted mean. + scale = data.mean() - loc + else: + scale = fscale + + # We expect the return values to be floating point, so ensure it + # by explicitly converting to float. + return float(loc), float(scale) + + +expon = expon_gen(a=0.0, name='expon') + + +class exponnorm_gen(rv_continuous): + r"""An exponentially modified Normal continuous random variable. + + Also known as the exponentially modified Gaussian distribution [1]_. + + %(before_notes)s + + Notes + ----- + The probability density function for `exponnorm` is: + + .. math:: + + f(x, K) = \frac{1}{2K} \exp\left(\frac{1}{2 K^2} - x / K \right) + \text{erfc}\left(-\frac{x - 1/K}{\sqrt{2}}\right) + + where :math:`x` is a real number and :math:`K > 0`. + + It can be thought of as the sum of a standard normal random variable + and an independent exponentially distributed random variable with rate + ``1/K``. + + %(after_notes)s + + An alternative parameterization of this distribution (for example, in + the Wikipedia article [1]_) involves three parameters, :math:`\mu`, + :math:`\lambda` and :math:`\sigma`. + + In the present parameterization this corresponds to having ``loc`` and + ``scale`` equal to :math:`\mu` and :math:`\sigma`, respectively, and + shape parameter :math:`K = 1/(\sigma\lambda)`. + + .. versionadded:: 0.16.0 + + References + ---------- + .. [1] Exponentially modified Gaussian distribution, Wikipedia, + https://en.wikipedia.org/wiki/Exponentially_modified_Gaussian_distribution + + %(example)s + + """ + def _shape_info(self): + return [_ShapeInfo("K", False, (0, np.inf), (False, False))] + + def _rvs(self, K, size=None, random_state=None): + expval = random_state.standard_exponential(size) * K + gval = random_state.standard_normal(size) + return expval + gval + + def _pdf(self, x, K): + return np.exp(self._logpdf(x, K)) + + def _logpdf(self, x, K): + u = (-x + 1.0 / K) / np.sqrt(2) + def logpdf_erfcx(x, K, u): + erfcx_term = np.log(sc.erfcx(u)) + return -np.log(2) - 0.5 * (x ** 2) - np.log(K) + erfcx_term + def logpdf_default(x, K, u): + invK = 1.0 / K + exparg = invK * (-x + 0.5 * invK) + return exparg + _norm_logcdf(x - invK) - np.log(K) + use_erfcx = np.logical_and(u >= 0, K < 1) + return xpx.apply_where( + use_erfcx, (x, K, u), logpdf_erfcx, logpdf_default) + + def _cdf(self, x, K): + u = (-x + 1.0 / K) / np.sqrt(2) + def cdf_erfcx(x, K, u): + return _norm_cdf(x) - 0.5 * np.exp(-0.5 * x ** 2) * sc.erfcx(u) + def cdf_default(x, K, u): + invK = 1.0 / K + expval = invK * (0.5 * invK - x) + logprod = expval + _norm_logcdf(x - invK) + return _norm_cdf(x) - np.exp(logprod) + use_erfcx = np.logical_and(u >= 0, K < 1) + return xpx.apply_where( + use_erfcx, (x, K, u), cdf_erfcx, cdf_default) + + def _sf(self, x, K): + u = (-x + 1.0 / K) / np.sqrt(2) + def sf_erfcx(x, K, u): + return _norm_cdf(-x) + 0.5 * np.exp(-0.5 * x ** 2) * sc.erfcx(u) + def sf_default(x, K, u): + invK = 1.0 / K + expval = invK * (0.5 * invK - x) + logprod = expval + _norm_logcdf(x - invK) + return _norm_cdf(-x) + np.exp(logprod) + use_erfcx = np.logical_and(u >= 0, K < 1) + return xpx.apply_where( + use_erfcx, (x, K, u), sf_erfcx, sf_default) + + def _stats(self, K): + K2 = K * K + opK2 = 1.0 + K2 + skw = 2 * K**3 * opK2**(-1.5) + krt = 6.0 * K2 * K2 * opK2**(-2) + return K, opK2, skw, krt + +exponnorm = exponnorm_gen(name='exponnorm') + + +def _pow1pm1(x, y): + """ + Compute (1 + x)**y - 1. + + Uses expm1 and xlog1py to avoid loss of precision when + (1 + x)**y is close to 1. + + Note that the inverse of this function with respect to x is + ``_pow1pm1(x, 1/y)``. That is, if + + t = _pow1pm1(x, y) + + then + + x = _pow1pm1(t, 1/y) + """ + return np.expm1(sc.xlog1py(y, x)) + + +class exponweib_gen(rv_continuous): + r"""An exponentiated Weibull continuous random variable. + + %(before_notes)s + + See Also + -------- + weibull_min, numpy.random.Generator.weibull + + Notes + ----- + The probability density function for `exponweib` is: + + .. math:: + + f(x, a, c) = a c [1-\exp(-x^c)]^{a-1} \exp(-x^c) x^{c-1} + + and its cumulative distribution function is: + + .. math:: + + F(x, a, c) = [1-\exp(-x^c)]^a + + for :math:`x > 0`, :math:`a > 0`, :math:`c > 0`. + + `exponweib` takes :math:`a` and :math:`c` as shape parameters: + + * :math:`a` is the exponentiation parameter, + with the special case :math:`a=1` corresponding to the + (non-exponentiated) Weibull distribution `weibull_min`. + * :math:`c` is the shape parameter of the non-exponentiated Weibull law. + + %(after_notes)s + + References + ---------- + https://en.wikipedia.org/wiki/Exponentiated_Weibull_distribution + + %(example)s + + """ + def _shape_info(self): + ia = _ShapeInfo("a", False, (0, np.inf), (False, False)) + ic = _ShapeInfo("c", False, (0, np.inf), (False, False)) + return [ia, ic] + + def _pdf(self, x, a, c): + # exponweib.pdf(x, a, c) = + # a * c * (1-exp(-x**c))**(a-1) * exp(-x**c)*x**(c-1) + return np.exp(self._logpdf(x, a, c)) + + def _logpdf(self, x, a, c): + negxc = -x**c + exm1c = -sc.expm1(negxc) + logp = (np.log(a) + np.log(c) + sc.xlogy(a - 1.0, exm1c) + + negxc + sc.xlogy(c - 1.0, x)) + return logp + + def _cdf(self, x, a, c): + exm1c = -sc.expm1(-x**c) + return exm1c**a + + def _ppf(self, q, a, c): + return (-sc.log1p(-q**(1.0/a)))**np.asarray(1.0/c) + + def _sf(self, x, a, c): + return -_pow1pm1(-np.exp(-x**c), a) + + def _isf(self, p, a, c): + return (-np.log(-_pow1pm1(-p, 1/a)))**(1/c) + + +exponweib = exponweib_gen(a=0.0, name='exponweib') + + +class exponpow_gen(rv_continuous): + r"""An exponential power continuous random variable. + + %(before_notes)s + + Notes + ----- + The probability density function for `exponpow` is: + + .. math:: + + f(x, b) = b x^{b-1} \exp(1 + x^b - \exp(x^b)) + + for :math:`x \ge 0`, :math:`b > 0`. Note that this is a different + distribution from the exponential power distribution that is also known + under the names "generalized normal" or "generalized Gaussian". + + `exponpow` takes ``b`` as a shape parameter for :math:`b`. + + %(after_notes)s + + References + ---------- + http://www.math.wm.edu/~leemis/chart/UDR/PDFs/Exponentialpower.pdf + + %(example)s + + """ + def _shape_info(self): + return [_ShapeInfo("b", False, (0, np.inf), (False, False))] + + def _pdf(self, x, b): + # exponpow.pdf(x, b) = b * x**(b-1) * exp(1 + x**b - exp(x**b)) + return np.exp(self._logpdf(x, b)) + + def _logpdf(self, x, b): + xb = x**b + f = 1 + np.log(b) + sc.xlogy(b - 1.0, x) + xb - np.exp(xb) + return f + + def _cdf(self, x, b): + return -sc.expm1(-sc.expm1(x**b)) + + def _sf(self, x, b): + return np.exp(-sc.expm1(x**b)) + + def _isf(self, x, b): + return (sc.log1p(-np.log(x)))**(1./b) + + def _ppf(self, q, b): + return pow(sc.log1p(-sc.log1p(-q)), 1.0/b) + + +exponpow = exponpow_gen(a=0.0, name='exponpow') + + +class fatiguelife_gen(rv_continuous): + r"""A fatigue-life (Birnbaum-Saunders) continuous random variable. + + %(before_notes)s + + Notes + ----- + The probability density function for `fatiguelife` is: + + .. math:: + + f(x, c) = \frac{x+1}{2c\sqrt{2\pi x^3}} \exp(-\frac{(x-1)^2}{2x c^2}) + + for :math:`x >= 0` and :math:`c > 0`. + + `fatiguelife` takes ``c`` as a shape parameter for :math:`c`. + + %(after_notes)s + + References + ---------- + .. [1] "Birnbaum-Saunders distribution", + https://en.wikipedia.org/wiki/Birnbaum-Saunders_distribution + + %(example)s + + """ + _support_mask = rv_continuous._open_support_mask + + def _shape_info(self): + return [_ShapeInfo("c", False, (0, np.inf), (False, False))] + + def _rvs(self, c, size=None, random_state=None): + z = random_state.standard_normal(size) + x = 0.5*c*z + x2 = x*x + t = 1.0 + 2*x2 + 2*x*np.sqrt(1 + x2) + return t + + def _pdf(self, x, c): + # fatiguelife.pdf(x, c) = + # (x+1) / (2*c*sqrt(2*pi*x**3)) * exp(-(x-1)**2/(2*x*c**2)) + return np.exp(self._logpdf(x, c)) + + def _logpdf(self, x, c): + return (np.log(x+1) - (x-1)**2 / (2.0*x*c**2) - np.log(2*c) - + 0.5*(np.log(2*np.pi) + 3*np.log(x))) + + def _cdf(self, x, c): + return _norm_cdf(1.0 / c * (np.sqrt(x) - 1.0/np.sqrt(x))) + + def _ppf(self, q, c): + tmp = c * _norm_ppf(q) + return 0.25 * (tmp + np.sqrt(tmp**2 + 4))**2 + + def _sf(self, x, c): + return _norm_sf(1.0 / c * (np.sqrt(x) - 1.0/np.sqrt(x))) + + def _isf(self, q, c): + tmp = -c * _norm_ppf(q) + return 0.25 * (tmp + np.sqrt(tmp**2 + 4))**2 + + def _stats(self, c): + # NB: the formula for kurtosis in wikipedia seems to have an error: + # it's 40, not 41. At least it disagrees with the one from Wolfram + # Alpha. And the latter one, below, passes the tests, while the wiki + # one doesn't So far I didn't have the guts to actually check the + # coefficients from the expressions for the raw moments. + c2 = c*c + mu = c2 / 2.0 + 1.0 + den = 5.0 * c2 + 4.0 + mu2 = c2*den / 4.0 + g1 = 4 * c * (11*c2 + 6.0) / np.power(den, 1.5) + g2 = 6 * c2 * (93*c2 + 40.0) / den**2.0 + return mu, mu2, g1, g2 + + +fatiguelife = fatiguelife_gen(a=0.0, name='fatiguelife') + + +class foldcauchy_gen(rv_continuous): + r"""A folded Cauchy continuous random variable. + + %(before_notes)s + + Notes + ----- + The probability density function for `foldcauchy` is: + + .. math:: + + f(x, c) = \frac{1}{\pi (1+(x-c)^2)} + \frac{1}{\pi (1+(x+c)^2)} + + for :math:`x \ge 0` and :math:`c \ge 0`. + + `foldcauchy` takes ``c`` as a shape parameter for :math:`c`. + + %(example)s + + """ + def _argcheck(self, c): + return c >= 0 + + def _shape_info(self): + return [_ShapeInfo("c", False, (0, np.inf), (True, False))] + + def _rvs(self, c, size=None, random_state=None): + return abs(cauchy.rvs(loc=c, size=size, + random_state=random_state)) + + def _pdf(self, x, c): + # foldcauchy.pdf(x, c) = 1/(pi*(1+(x-c)**2)) + 1/(pi*(1+(x+c)**2)) + return 1.0/np.pi*(1.0/(1+(x-c)**2) + 1.0/(1+(x+c)**2)) + + def _cdf(self, x, c): + return 1.0/np.pi*(np.arctan(x-c) + np.arctan(x+c)) + + def _sf(self, x, c): + # 1 - CDF(x, c) = 1 - (atan(x - c) + atan(x + c))/pi + # = ((pi/2 - atan(x - c)) + (pi/2 - atan(x + c)))/pi + # = (acot(x - c) + acot(x + c))/pi + # = (atan2(1, x - c) + atan2(1, x + c))/pi + return (np.arctan2(1, x - c) + np.arctan2(1, x + c))/np.pi + + def _stats(self, c): + return np.inf, np.inf, np.nan, np.nan + + +foldcauchy = foldcauchy_gen(a=0.0, name='foldcauchy') + + +class f_gen(rv_continuous): + r"""An F continuous random variable. + + For the noncentral F distribution, see `ncf`. + + %(before_notes)s + + See Also + -------- + ncf + + Notes + ----- + The F distribution with :math:`df_1 > 0` and :math:`df_2 > 0` degrees of freedom is + the distribution of the ratio of two independent chi-squared distributions with + :math:`df_1` and :math:`df_2` degrees of freedom, after rescaling by + :math:`df_2 / df_1`. + + The probability density function for `f` is: + + .. math:: + + f(x, df_1, df_2) = \frac{df_2^{df_2/2} df_1^{df_1/2} x^{df_1 / 2-1}} + {(df_2+df_1 x)^{(df_1+df_2)/2} + B(df_1/2, df_2/2)} + + for :math:`x > 0`. + + `f` accepts shape parameters ``dfn`` and ``dfd`` for :math:`df_1`, the degrees of + freedom of the chi-squared distribution in the numerator, and :math:`df_2`, the + degrees of freedom of the chi-squared distribution in the denominator, respectively. + + %(after_notes)s + + %(example)s + + """ + def _shape_info(self): + idfn = _ShapeInfo("dfn", False, (0, np.inf), (False, False)) + idfd = _ShapeInfo("dfd", False, (0, np.inf), (False, False)) + return [idfn, idfd] + + def _rvs(self, dfn, dfd, size=None, random_state=None): + return random_state.f(dfn, dfd, size) + + def _pdf(self, x, dfn, dfd): + # df2**(df2/2) * df1**(df1/2) * x**(df1/2-1) + # F.pdf(x, df1, df2) = -------------------------------------------- + # (df2+df1*x)**((df1+df2)/2) * B(df1/2, df2/2) + return np.exp(self._logpdf(x, dfn, dfd)) + + def _logpdf(self, x, dfn, dfd): + n = 1.0 * dfn + m = 1.0 * dfd + lPx = (m/2 * np.log(m) + n/2 * np.log(n) + sc.xlogy(n/2 - 1, x) + - (((n+m)/2) * np.log(m + n*x) + sc.betaln(n/2, m/2))) + return lPx + + def _cdf(self, x, dfn, dfd): + return sc.fdtr(dfn, dfd, x) + + def _sf(self, x, dfn, dfd): + return sc.fdtrc(dfn, dfd, x) + + def _ppf(self, q, dfn, dfd): + return sc.fdtri(dfn, dfd, q) + + def _stats(self, dfn, dfd): + v1, v2 = 1. * dfn, 1. * dfd + v2_2, v2_4, v2_6, v2_8 = v2 - 2., v2 - 4., v2 - 6., v2 - 8. + + mu = xpx.apply_where( + v2 > 2, (v2, v2_2), + lambda v2, v2_2: v2 / v2_2, + fill_value=np.inf) + + mu2 = xpx.apply_where( + v2 > 4, (v1, v2, v2_2, v2_4), + lambda v1, v2, v2_2, v2_4: + 2 * v2 * v2 * (v1 + v2_2) / (v1 * v2_2**2 * v2_4), + fill_value=np.inf) + + g1 = xpx.apply_where( + v2 > 6, (v1, v2_2, v2_4, v2_6), + lambda v1, v2_2, v2_4, v2_6: + (2 * v1 + v2_2) / v2_6 * np.sqrt(v2_4 / (v1 * (v1 + v2_2))), + fill_value=np.nan) + g1 *= np.sqrt(8.) + + g2 = xpx.apply_where( + v2 > 8, (g1, v2_6, v2_8), + lambda g1, v2_6, v2_8: (8 + g1 * g1 * v2_6) / v2_8, + fill_value=np.nan) + g2 *= 3. / 2. + + return mu, mu2, g1, g2 + + def _entropy(self, dfn, dfd): + # the formula found in literature is incorrect. This one yields the + # same result as numerical integration using the generic entropy + # definition. This is also tested in tests/test_conntinous_basic + half_dfn = 0.5 * dfn + half_dfd = 0.5 * dfd + half_sum = 0.5 * (dfn + dfd) + + return (np.log(dfd) - np.log(dfn) + sc.betaln(half_dfn, half_dfd) + + (1 - half_dfn) * sc.psi(half_dfn) - (1 + half_dfd) * + sc.psi(half_dfd) + half_sum * sc.psi(half_sum)) + + +f = f_gen(a=0.0, name='f') + + +## Folded Normal +## abs(Z) where (Z is normal with mu=L and std=S so that c=abs(L)/S) +## +## note: regress docs have scale parameter correct, but first parameter +## he gives is a shape parameter A = c * scale + +## Half-normal is folded normal with shape-parameter c=0. + +class foldnorm_gen(rv_continuous): + r"""A folded normal continuous random variable. + + %(before_notes)s + + Notes + ----- + The probability density function for `foldnorm` is: + + .. math:: + + f(x, c) = \sqrt{2/\pi} cosh(c x) \exp(-\frac{x^2+c^2}{2}) + + for :math:`x \ge 0` and :math:`c \ge 0`. + + `foldnorm` takes ``c`` as a shape parameter for :math:`c`. + + %(after_notes)s + + %(example)s + + """ + def _argcheck(self, c): + return c >= 0 + + def _shape_info(self): + return [_ShapeInfo("c", False, (0, np.inf), (True, False))] + + def _rvs(self, c, size=None, random_state=None): + return abs(random_state.standard_normal(size) + c) + + def _pdf(self, x, c): + # foldnormal.pdf(x, c) = sqrt(2/pi) * cosh(c*x) * exp(-(x**2+c**2)/2) + return _norm_pdf(x + c) + _norm_pdf(x-c) + + def _cdf(self, x, c): + sqrt_two = np.sqrt(2) + return 0.5 * (sc.erf((x - c)/sqrt_two) + sc.erf((x + c)/sqrt_two)) + + def _sf(self, x, c): + return _norm_sf(x - c) + _norm_sf(x + c) + + def _stats(self, c): + # Regina C. Elandt, Technometrics 3, 551 (1961) + # https://www.jstor.org/stable/1266561 + # + c2 = c*c + expfac = np.exp(-0.5*c2) / np.sqrt(2.*np.pi) + + mu = 2.*expfac + c * sc.erf(c/np.sqrt(2)) + mu2 = c2 + 1 - mu*mu + + g1 = 2. * (mu*mu*mu - c2*mu - expfac) + g1 /= np.power(mu2, 1.5) + + g2 = c2 * (c2 + 6.) + 3 + 8.*expfac*mu + g2 += (2. * (c2 - 3.) - 3. * mu**2) * mu**2 + g2 = g2 / mu2**2.0 - 3. + + return mu, mu2, g1, g2 + + +foldnorm = foldnorm_gen(a=0.0, name='foldnorm') + + +class weibull_min_gen(rv_continuous): + r"""Weibull minimum continuous random variable. + + The Weibull Minimum Extreme Value distribution, from extreme value theory + (Fisher-Gnedenko theorem), is also often simply called the Weibull + distribution. It arises as the limiting distribution of the rescaled + minimum of iid random variables. + + %(before_notes)s + + See Also + -------- + weibull_max, numpy.random.Generator.weibull, exponweib + + Notes + ----- + The probability density function for `weibull_min` is: + + .. math:: + + f(x, c) = c x^{c-1} \exp(-x^c) + + for :math:`x > 0`, :math:`c > 0`. + + `weibull_min` takes ``c`` as a shape parameter for :math:`c`. + (named :math:`k` in Wikipedia article and :math:`a` in + ``numpy.random.weibull``). Special shape values are :math:`c=1` and + :math:`c=2` where Weibull distribution reduces to the `expon` and + `rayleigh` distributions respectively. + + Suppose ``X`` is an exponentially distributed random variable with + scale ``s``. Then ``Y = X**k`` is `weibull_min` distributed with shape + ``c = 1/k`` and scale ``s**k``. + + %(after_notes)s + + References + ---------- + https://en.wikipedia.org/wiki/Weibull_distribution + + https://en.wikipedia.org/wiki/Fisher-Tippett-Gnedenko_theorem + + %(example)s + + """ + def _shape_info(self): + return [_ShapeInfo("c", False, (0, np.inf), (False, False))] + + def _pdf(self, x, c): + # weibull_min.pdf(x, c) = c * x**(c-1) * exp(-x**c) + return c*pow(x, c-1)*np.exp(-pow(x, c)) + + def _logpdf(self, x, c): + return np.log(c) + sc.xlogy(c - 1, x) - pow(x, c) + + def _cdf(self, x, c): + return -sc.expm1(-pow(x, c)) + + def _ppf(self, q, c): + return pow(-sc.log1p(-q), 1.0/c) + + def _sf(self, x, c): + return np.exp(self._logsf(x, c)) + + def _logsf(self, x, c): + return -pow(x, c) + + def _isf(self, q, c): + return (-np.log(q))**(1/c) + + def _munp(self, n, c): + return sc.gamma(1.0+n*1.0/c) + + def _entropy(self, c): + return -_EULER / c - np.log(c) + _EULER + 1 + + @extend_notes_in_docstring(rv_continuous, notes="""\ + If ``method='mm'``, parameters fixed by the user are respected, and the + remaining parameters are used to match distribution and sample moments + where possible. For example, if the user fixes the location with + ``floc``, the parameters will only match the distribution skewness and + variance to the sample skewness and variance; no attempt will be made + to match the means or minimize a norm of the errors. + \n\n""") + def fit(self, data, *args, **kwds): + + if isinstance(data, CensoredData): + if data.num_censored() == 0: + data = data._uncensor() + else: + return super().fit(data, *args, **kwds) + + if kwds.pop('superfit', False): + return super().fit(data, *args, **kwds) + + # this extracts fixed shape, location, and scale however they + # are specified, and also leaves them in `kwds` + data, fc, floc, fscale = _check_fit_input_parameters(self, data, + args, kwds) + method = kwds.get("method", "mle").lower() + + # See https://en.wikipedia.org/wiki/Weibull_distribution#Moments for + # moment formulas. + def skew(c): + gamma1 = sc.gamma(1+1/c) + gamma2 = sc.gamma(1+2/c) + gamma3 = sc.gamma(1+3/c) + num = 2 * gamma1**3 - 3*gamma1*gamma2 + gamma3 + den = (gamma2 - gamma1**2)**(3/2) + return num/den + + # For c in [1e2, 3e4], population skewness appears to approach + # asymptote near -1.139, but past c > 3e4, skewness begins to vary + # wildly, and MoM won't provide a good guess. Get out early. + s = stats.skew(data) + max_c = 1e4 + s_min = skew(max_c) + if s < s_min and method != "mm" and fc is None and not args: + return super().fit(data, *args, **kwds) + + # If method is method of moments, we don't need the user's guesses. + # Otherwise, extract the guesses from args and kwds. + if method == "mm": + c, loc, scale = None, None, None + else: + c = args[0] if len(args) else None + loc = kwds.pop('loc', None) + scale = kwds.pop('scale', None) + + if fc is None and c is None: # not fixed and no guess: use MoM + # Solve for c that matches sample distribution skewness to sample + # skewness. + # we start having numerical issues with `weibull_min` with + # parameters outside this range - and not just in this method. + # We could probably improve the situation by doing everything + # in the log space, but that is for another time. + c = root_scalar(lambda c: skew(c) - s, bracket=[0.02, max_c], + method='bisect').root + elif fc is not None: # fixed: use it + c = fc + + if fscale is None and scale is None: + v = np.var(data) + scale = np.sqrt(v / (sc.gamma(1+2/c) - sc.gamma(1+1/c)**2)) + elif fscale is not None: + scale = fscale + + if floc is None and loc is None: + m = np.mean(data) + loc = m - scale*sc.gamma(1 + 1/c) + elif floc is not None: + loc = floc + + if method == 'mm': + return c, loc, scale + else: + # At this point, parameter "guesses" may equal the fixed parameters + # in kwds. No harm in passing them as guesses, too. + return super().fit(data, c, loc=loc, scale=scale, **kwds) + + +weibull_min = weibull_min_gen(a=0.0, name='weibull_min') + + +class truncweibull_min_gen(rv_continuous): + r"""A doubly truncated Weibull minimum continuous random variable. + + %(before_notes)s + + See Also + -------- + weibull_min, truncexpon + + Notes + ----- + The probability density function for `truncweibull_min` is: + + .. math:: + + f(x, a, b, c) = \frac{c x^{c-1} \exp(-x^c)}{\exp(-a^c) - \exp(-b^c)} + + for :math:`a < x <= b`, :math:`0 \le a < b` and :math:`c > 0`. + + `truncweibull_min` takes :math:`a`, :math:`b`, and :math:`c` as shape + parameters. + + Notice that the truncation values, :math:`a` and :math:`b`, are defined in + standardized form: + + .. math:: + + a = (u_l - loc)/scale + b = (u_r - loc)/scale + + where :math:`u_l` and :math:`u_r` are the specific left and right + truncation values, respectively. In other words, the support of the + distribution becomes :math:`(a*scale + loc) < x <= (b*scale + loc)` when + :math:`loc` and/or :math:`scale` are provided. + + %(after_notes)s + + References + ---------- + + .. [1] Rinne, H. "The Weibull Distribution: A Handbook". CRC Press (2009). + + %(example)s + + """ + def _argcheck(self, c, a, b): + return (a >= 0.) & (b > a) & (c > 0.) + + def _shape_info(self): + ic = _ShapeInfo("c", False, (0, np.inf), (False, False)) + ia = _ShapeInfo("a", False, (0, np.inf), (True, False)) + ib = _ShapeInfo("b", False, (0, np.inf), (False, False)) + return [ic, ia, ib] + + def _fitstart(self, data): + # Arbitrary, but default a=b=c=1 is not valid + return super()._fitstart(data, args=(1, 0, 1)) + + def _get_support(self, c, a, b): + return a, b + + def _pdf(self, x, c, a, b): + denum = (np.exp(-pow(a, c)) - np.exp(-pow(b, c))) + return (c * pow(x, c-1) * np.exp(-pow(x, c))) / denum + + def _logpdf(self, x, c, a, b): + logdenum = np.log(np.exp(-pow(a, c)) - np.exp(-pow(b, c))) + return np.log(c) + sc.xlogy(c - 1, x) - pow(x, c) - logdenum + + def _cdf(self, x, c, a, b): + num = (np.exp(-pow(a, c)) - np.exp(-pow(x, c))) + denum = (np.exp(-pow(a, c)) - np.exp(-pow(b, c))) + return num / denum + + def _logcdf(self, x, c, a, b): + lognum = np.log(np.exp(-pow(a, c)) - np.exp(-pow(x, c))) + logdenum = np.log(np.exp(-pow(a, c)) - np.exp(-pow(b, c))) + return lognum - logdenum + + def _sf(self, x, c, a, b): + num = (np.exp(-pow(x, c)) - np.exp(-pow(b, c))) + denum = (np.exp(-pow(a, c)) - np.exp(-pow(b, c))) + return num / denum + + def _logsf(self, x, c, a, b): + lognum = np.log(np.exp(-pow(x, c)) - np.exp(-pow(b, c))) + logdenum = np.log(np.exp(-pow(a, c)) - np.exp(-pow(b, c))) + return lognum - logdenum + + def _isf(self, q, c, a, b): + return pow( + -np.log((1 - q) * np.exp(-pow(b, c)) + q * np.exp(-pow(a, c))), 1/c + ) + + def _ppf(self, q, c, a, b): + return pow( + -np.log((1 - q) * np.exp(-pow(a, c)) + q * np.exp(-pow(b, c))), 1/c + ) + + def _munp(self, n, c, a, b): + gamma_fun = sc.gamma(n/c + 1.) * ( + sc.gammainc(n/c + 1., pow(b, c)) - sc.gammainc(n/c + 1., pow(a, c)) + ) + denum = (np.exp(-pow(a, c)) - np.exp(-pow(b, c))) + return gamma_fun / denum + + +truncweibull_min = truncweibull_min_gen(name='truncweibull_min') +truncweibull_min._support = ('a', 'b') + + +class weibull_max_gen(rv_continuous): + r"""Weibull maximum continuous random variable. + + The Weibull Maximum Extreme Value distribution, from extreme value theory + (Fisher-Gnedenko theorem), is the limiting distribution of rescaled + maximum of iid random variables. This is the distribution of -X + if X is from the `weibull_min` function. + + %(before_notes)s + + See Also + -------- + weibull_min + + Notes + ----- + The probability density function for `weibull_max` is: + + .. math:: + + f(x, c) = c (-x)^{c-1} \exp(-(-x)^c) + + for :math:`x < 0`, :math:`c > 0`. + + `weibull_max` takes ``c`` as a shape parameter for :math:`c`. + + %(after_notes)s + + References + ---------- + https://en.wikipedia.org/wiki/Weibull_distribution + + https://en.wikipedia.org/wiki/Fisher-Tippett-Gnedenko_theorem + + %(example)s + + """ + def _shape_info(self): + return [_ShapeInfo("c", False, (0, np.inf), (False, False))] + + def _pdf(self, x, c): + # weibull_max.pdf(x, c) = c * (-x)**(c-1) * exp(-(-x)**c) + return c*pow(-x, c-1)*np.exp(-pow(-x, c)) + + def _logpdf(self, x, c): + return np.log(c) + sc.xlogy(c-1, -x) - pow(-x, c) + + def _cdf(self, x, c): + return np.exp(-pow(-x, c)) + + def _logcdf(self, x, c): + return -pow(-x, c) + + def _sf(self, x, c): + return -sc.expm1(-pow(-x, c)) + + def _ppf(self, q, c): + return -pow(-np.log(q), 1.0/c) + + def _munp(self, n, c): + val = sc.gamma(1.0+n*1.0/c) + if int(n) % 2: + sgn = -1 + else: + sgn = 1 + return sgn * val + + def _entropy(self, c): + return -_EULER / c - np.log(c) + _EULER + 1 + + +weibull_max = weibull_max_gen(b=0.0, name='weibull_max') + + +class genlogistic_gen(rv_continuous): + r"""A generalized logistic continuous random variable. + + %(before_notes)s + + Notes + ----- + The probability density function for `genlogistic` is: + + .. math:: + + f(x, c) = c \frac{\exp(-x)} + {(1 + \exp(-x))^{c+1}} + + for real :math:`x` and :math:`c > 0`. In literature, different + generalizations of the logistic distribution can be found. This is the type 1 + generalized logistic distribution according to [1]_. It is also referred to + as the skew-logistic distribution [2]_. + + `genlogistic` takes ``c`` as a shape parameter for :math:`c`. + + %(after_notes)s + + References + ---------- + .. [1] Johnson et al. "Continuous Univariate Distributions", Volume 2, + Wiley. 1995. + .. [2] "Generalized Logistic Distribution", Wikipedia, + https://en.wikipedia.org/wiki/Generalized_logistic_distribution + + %(example)s + + """ + def _shape_info(self): + return [_ShapeInfo("c", False, (0, np.inf), (False, False))] + + def _pdf(self, x, c): + # genlogistic.pdf(x, c) = c * exp(-x) / (1 + exp(-x))**(c+1) + return np.exp(self._logpdf(x, c)) + + def _logpdf(self, x, c): + # Two mathematically equivalent expressions for log(pdf(x, c)): + # log(pdf(x, c)) = log(c) - x - (c + 1)*log(1 + exp(-x)) + # = log(c) + c*x - (c + 1)*log(1 + exp(x)) + mult = -(c - 1) * (x < 0) - 1 + absx = np.abs(x) + return np.log(c) + mult*absx - (c+1) * sc.log1p(np.exp(-absx)) + + def _cdf(self, x, c): + Cx = (1+np.exp(-x))**(-c) + return Cx + + def _logcdf(self, x, c): + return -c * np.log1p(np.exp(-x)) + + def _ppf(self, q, c): + return -np.log(sc.powm1(q, -1.0/c)) + + def _sf(self, x, c): + return -sc.expm1(self._logcdf(x, c)) + + def _isf(self, q, c): + return self._ppf(1 - q, c) + + def _stats(self, c): + mu = _EULER + sc.psi(c) + mu2 = np.pi*np.pi/6.0 + sc.zeta(2, c) + g1 = -2*sc.zeta(3, c) + 2*_ZETA3 + g1 /= np.power(mu2, 1.5) + g2 = np.pi**4/15.0 + 6*sc.zeta(4, c) + g2 /= mu2**2.0 + return mu, mu2, g1, g2 + + def _entropy(self, c): + return xpx.apply_where( + c < 8e6, c, + lambda c: -np.log(c) + sc.psi(c + 1) + _EULER + 1, + # asymptotic expansion: psi(c) ~ log(c) - 1 / (2 * c) + # a = -log(c) + psi(c + 1) + # = -log(c) + psi(c) + 1 / c + # ~ -log(c) + log(c) - 1 / (2 * c) + 1 / c + # = 1 / (2 * c) + lambda c: 1 / (2 * c) + _EULER + 1) + + +genlogistic = genlogistic_gen(name='genlogistic') + + +class genpareto_gen(rv_continuous): + r"""A generalized Pareto continuous random variable. + + %(before_notes)s + + Notes + ----- + The probability density function for `genpareto` is: + + .. math:: + + f(x, c) = (1 + c x)^{-1 - 1/c} + + defined for :math:`x \ge 0` if :math:`c \ge 0`, and for + :math:`0 \le x \le -1/c` if :math:`c < 0`. + + `genpareto` takes ``c`` as a shape parameter for :math:`c`. + + For :math:`c=0`, `genpareto` reduces to the exponential + distribution, `expon`: + + .. math:: + + f(x, 0) = \exp(-x) + + For :math:`c=-1`, `genpareto` is uniform on ``[0, 1]``: + + .. math:: + + f(x, -1) = 1 + + %(after_notes)s + + %(example)s + + """ + def _argcheck(self, c): + return np.isfinite(c) + + def _shape_info(self): + return [_ShapeInfo("c", False, (-np.inf, np.inf), (False, False))] + + def _get_support(self, c): + c = np.asarray(c) + a = np.broadcast_arrays(self.a, c)[0].copy() + b = xpx.apply_where(c < 0, c, lambda c: -1. / c, + fill_value=np.inf) + return a, b + + def _pdf(self, x, c): + # genpareto.pdf(x, c) = (1 + c * x)**(-1 - 1/c) + return np.exp(self._logpdf(x, c)) + + def _logpdf(self, x, c): + return xpx.apply_where((x == x) & (c != 0), (x, c), + lambda x, c: -sc.xlog1py(c + 1., c*x) / c, + fill_value=-x) + + def _cdf(self, x, c): + return -sc.inv_boxcox1p(-x, -c) + + def _sf(self, x, c): + return sc.inv_boxcox(-x, -c) + + def _logsf(self, x, c): + return xpx.apply_where((x == x) & (c != 0), (x, c), + lambda x, c: -sc.log1p(c*x) / c, + fill_value=-x) + + def _ppf(self, q, c): + return -sc.boxcox1p(-q, -c) + + def _isf(self, q, c): + return -sc.boxcox(q, -c) + + def _stats(self, c, moments='mv'): + m, v, s, k = None, None, None, None + + if 'm' in moments: + m = xpx.apply_where(c < 1, c, + lambda xi: 1 / (1 - xi), + fill_value=np.inf) + + if 'v' in moments: + v = xpx.apply_where(c < 1/2, c, + lambda xi: 1 / (1 - xi)**2 / (1 - 2 * xi), + fill_value=np.nan) + + if 's' in moments: + s = xpx.apply_where( + c < 1/3, c, + lambda xi: 2 * (1 + xi) * np.sqrt(1 - 2*xi) / (1 - 3*xi), + fill_value=np.nan) + + if 'k' in moments: + k = xpx.apply_where( + c < 1/4, c, + lambda xi: 3 * (1 - 2*xi) * (2*xi**2 + xi + 3) + / (1 - 3*xi) / (1 - 4*xi) - 3, + fill_value=np.nan) + + return m, v, s, k + + def _munp(self, n, c): + def __munp(c): + val = 0.0 + k = np.arange(0, n + 1) + for ki, cnk in zip(k, sc.comb(n, k)): + val = val + cnk * (-1) ** ki / (1.0 - c * ki) + return np.where(c * n < 1, val * (-1.0 / c) ** n, np.inf) + + return xpx.apply_where(c != 0, c, __munp, fill_value=sc.gamma(n + 1)) + + def _entropy(self, c): + return 1. + c + + +genpareto = genpareto_gen(a=0.0, name='genpareto') + + +class genexpon_gen(rv_continuous): + r"""A generalized exponential continuous random variable. + + %(before_notes)s + + Notes + ----- + The probability density function for `genexpon` is: + + .. math:: + + f(x, a, b, c) = (a + b (1 - \exp(-c x))) + \exp(-a x - b x + \frac{b}{c} (1-\exp(-c x))) + + for :math:`x \ge 0`, :math:`a, b, c > 0`. + + `genexpon` takes :math:`a`, :math:`b` and :math:`c` as shape parameters. + + %(after_notes)s + + References + ---------- + H.K. Ryu, "An Extension of Marshall and Olkin's Bivariate Exponential + Distribution", Journal of the American Statistical Association, 1993. + + N. Balakrishnan, Asit P. Basu (editors), *The Exponential Distribution: + Theory, Methods and Applications*, Gordon and Breach, 1995. + ISBN 10: 2884491929 + + %(example)s + + """ + def _shape_info(self): + ia = _ShapeInfo("a", False, (0, np.inf), (False, False)) + ib = _ShapeInfo("b", False, (0, np.inf), (False, False)) + ic = _ShapeInfo("c", False, (0, np.inf), (False, False)) + return [ia, ib, ic] + + def _pdf(self, x, a, b, c): + # genexpon.pdf(x, a, b, c) = (a + b * (1 - exp(-c*x))) * \ + # exp(-a*x - b*x + b/c * (1-exp(-c*x))) + return (a + b*(-sc.expm1(-c*x)))*np.exp((-a-b)*x + + b*(-sc.expm1(-c*x))/c) + + def _logpdf(self, x, a, b, c): + return np.log(a+b*(-sc.expm1(-c*x))) + (-a-b)*x+b*(-sc.expm1(-c*x))/c + + def _cdf(self, x, a, b, c): + return -sc.expm1((-a-b)*x + b*(-sc.expm1(-c*x))/c) + + def _ppf(self, p, a, b, c): + s = a + b + t = (b - c*np.log1p(-p))/s + return (t + sc.lambertw(-b/s * np.exp(-t)).real)/c + + def _sf(self, x, a, b, c): + return np.exp((-a-b)*x + b*(-sc.expm1(-c*x))/c) + + def _isf(self, p, a, b, c): + s = a + b + t = (b - c*np.log(p))/s + return (t + sc.lambertw(-b/s * np.exp(-t)).real)/c + + +genexpon = genexpon_gen(a=0.0, name='genexpon') + + +class genextreme_gen(rv_continuous): + r"""A generalized extreme value continuous random variable. + + %(before_notes)s + + See Also + -------- + gumbel_r + + Notes + ----- + For :math:`c=0`, `genextreme` is equal to `gumbel_r` with + probability density function + + .. math:: + + f(x) = \exp(-\exp(-x)) \exp(-x), + + where :math:`-\infty < x < \infty`. + + For :math:`c \ne 0`, the probability density function for `genextreme` is: + + .. math:: + + f(x, c) = \exp(-(1-c x)^{1/c}) (1-c x)^{1/c-1}, + + where :math:`-\infty < x \le 1/c` if :math:`c > 0` and + :math:`1/c \le x < \infty` if :math:`c < 0`. + + Note that several sources and software packages use the opposite + convention for the sign of the shape parameter :math:`c`. + + `genextreme` takes ``c`` as a shape parameter for :math:`c`. + + %(after_notes)s + + %(example)s + + """ + def _argcheck(self, c): + return np.isfinite(c) + + def _shape_info(self): + return [_ShapeInfo("c", False, (-np.inf, np.inf), (False, False))] + + def _get_support(self, c): + c = np.asarray(c) + _b = np.where(c > 0, 1.0 / np.maximum(c, _XMIN), np.inf) + _a = np.where(c < 0, 1.0 / np.minimum(c, -_XMIN), -np.inf) + return _a, _b + + def _loglogcdf(self, x, c): + # Returns log(-log(cdf(x, c))) + return xpx.apply_where( + (x == x) & (c != 0), (x, c), + lambda x, c: sc.log1p(-c*x)/c, + fill_value=-x) + + def _pdf(self, x, c): + # genextreme.pdf(x, c) = + # exp(-exp(-x))*exp(-x), for c==0 + # exp(-(1-c*x)**(1/c))*(1-c*x)**(1/c-1), for x \le 1/c, c > 0 + return np.exp(self._logpdf(x, c)) + + def _logpdf(self, x, c): + # Suppress warnings 0 * inf + cx = xpx.apply_where((x == x) & (c != 0), (c, x), + operator.mul, fill_value=0.0) + logex2 = sc.log1p(-cx) + logpex2 = self._loglogcdf(x, c) + pex2 = np.exp(logpex2) + # Handle special cases + np.putmask(logpex2, (c == 0) & (x == -np.inf), 0.0) + logpdf = xpx.apply_where( + ~((cx == 1) | (cx == -np.inf)), + (pex2, logpex2, logex2), + lambda pex2, lpex2, lex2: -pex2 + lpex2 - lex2, + fill_value=-np.inf) + np.putmask(logpdf, (c == 1) & (x == 1), 0.0) + return logpdf + + def _logcdf(self, x, c): + return -np.exp(self._loglogcdf(x, c)) + + def _cdf(self, x, c): + return np.exp(self._logcdf(x, c)) + + def _sf(self, x, c): + return -sc.expm1(self._logcdf(x, c)) + + def _ppf(self, q, c): + x = -np.log(-np.log(q)) + return xpx.apply_where( + (x == x) & (c != 0), (x, c), + lambda x, c: -sc.expm1(-c * x) / c, + fill_value=x) + + def _isf(self, q, c): + x = -np.log(-sc.log1p(-q)) + return xpx.apply_where( + (x == x) & (c != 0), (x, c), + lambda x, c: -sc.expm1(-c * x) / c, + fill_value=x) + + def _stats(self, c): + def g(n): + return sc.gamma(n * c + 1) + g1 = g(1) + g2 = g(2) + g3 = g(3) + g4 = g(4) + g2mg12 = np.where(abs(c) < 1e-7, (c*np.pi)**2.0/6.0, g2-g1**2.0) + def gam2k_f(c): + return sc.expm1(sc.gammaln(2.0*c+1.0)-2*sc.gammaln(c + 1.0))/c**2.0 + gam2k = xpx.apply_where(abs(c) >= 1e-7, c, gam2k_f, fill_value=np.pi**2.0/6.0) + eps = 1e-14 + def gamk_f(c): + return sc.expm1(sc.gammaln(c + 1))/c + gamk = xpx.apply_where(abs(c) >= eps, c, gamk_f, fill_value=-_EULER) + + # mean + m = np.where(c < -1.0, np.nan, -gamk) + + # variance + v = np.where(c < -0.5, np.nan, g1**2.0*gam2k) + + # skewness + def sk1_eval(c, *args): + def sk1_eval_f(c, g1, g2, g3, g2mg12): + return np.sign(c)*(-g3 + (g2 + 2*g2mg12)*g1)/g2mg12**1.5 + return xpx.apply_where(c >= -1./3, (c, *args), + sk1_eval_f, fill_value=np.nan) + + sk_fill = 12*np.sqrt(6)*_ZETA3/np.pi**3 + args = (g1, g2, g3, g2mg12) + sk = xpx.apply_where(abs(c) > eps**0.29, (c, *args), + sk1_eval, fill_value=sk_fill) + + # kurtosis + def ku1_eval(c, *args): + def ku1_eval_f(g1, g2, g3, g4, g2mg12): + return (g4 + (-4*g3 + 3*(g2 + g2mg12)*g1)*g1)/g2mg12**2 - 3 + return xpx.apply_where(c >= -1./4, args, ku1_eval_f, fill_value=np.nan) + + args = (g1, g2, g3, g4, g2mg12) + ku = xpx.apply_where(abs(c) > eps**0.23, (c, *args), + ku1_eval, fill_value=12.0/5.0) + + return m, v, sk, ku + + def _fitstart(self, data): + if isinstance(data, CensoredData): + data = data._uncensor() + # This is better than the default shape of (1,). + g = _skew(data) + if g < 0: + a = 0.5 + else: + a = -0.5 + return super()._fitstart(data, args=(a,)) + + def _munp(self, n, c): + k = np.arange(0, n+1) + k = np.reshape(k, (-1,) + (1,)*c.ndim) + vals = 1.0/c**n * np.sum( + sc.comb(n, k) * (-1)**k * sc.gamma(c*k + 1), + axis=0) + return np.where(c*n > -1, vals, np.inf) + + def _entropy(self, c): + return _EULER*(1 - c) + 1 + + +genextreme = genextreme_gen(name='genextreme') + + +def _digammainv(y): + """Inverse of the digamma function (real positive arguments only). + + This function is used in the `fit` method of `gamma_gen`. + The function uses either optimize.fsolve or optimize.newton + to solve `sc.digamma(x) - y = 0`. There is probably room for + improvement, but currently it works over a wide range of y: + + >>> import numpy as np + >>> rng = np.random.default_rng() + >>> y = 64*rng.standard_normal(1000000) + >>> y.min(), y.max() + (-311.43592651416662, 351.77388222276869) + >>> x = [_digammainv(t) for t in y] + >>> np.abs(sc.digamma(x) - y).max() + 1.1368683772161603e-13 + + """ + _em = 0.5772156649015328606065120 + + def func(x): + return sc.digamma(x) - y + + if y > -0.125: + x0 = np.exp(y) + 0.5 + if y < 10: + # Some experimentation shows that newton reliably converges + # must faster than fsolve in this y range. For larger y, + # newton sometimes fails to converge. + value = optimize.newton(func, x0, tol=1e-10) + return value + elif y > -3: + x0 = np.exp(y/2.332) + 0.08661 + else: + x0 = 1.0 / (-y - _em) + + value, info, ier, mesg = optimize.fsolve(func, x0, xtol=1e-11, + full_output=True) + if ier != 1: + raise RuntimeError(f"_digammainv: fsolve failed, y = {y!r}") + + return value[0] + + +## Gamma (Use MATLAB and MATHEMATICA (b=theta=scale, a=alpha=shape) definition) + +## gamma(a, loc, scale) with a an integer is the Erlang distribution +## gamma(1, loc, scale) is the Exponential distribution +## gamma(df/2, 0, 2) is the chi2 distribution with df degrees of freedom. + +class gamma_gen(rv_continuous): + r"""A gamma continuous random variable. + + %(before_notes)s + + See Also + -------- + erlang, expon + + Notes + ----- + The probability density function for `gamma` is: + + .. math:: + + f(x, a) = \frac{x^{a-1} e^{-x}}{\Gamma(a)} + + for :math:`x \ge 0`, :math:`a > 0`. Here :math:`\Gamma(a)` refers to the + gamma function. + + `gamma` takes ``a`` as a shape parameter for :math:`a`. + + When :math:`a` is an integer, `gamma` reduces to the Erlang + distribution, and when :math:`a=1` to the exponential distribution. + + Gamma distributions are sometimes parameterized with two variables, + with a probability density function of: + + .. math:: + + f(x, \alpha, \beta) = + \frac{\beta^\alpha x^{\alpha - 1} e^{-\beta x }}{\Gamma(\alpha)} + + Note that this parameterization is equivalent to the above, with + ``scale = 1 / beta``. + + %(after_notes)s + + %(example)s + + """ + def _shape_info(self): + return [_ShapeInfo("a", False, (0, np.inf), (False, False))] + + def _rvs(self, a, size=None, random_state=None): + return random_state.standard_gamma(a, size) + + def _pdf(self, x, a): + # gamma.pdf(x, a) = x**(a-1) * exp(-x) / gamma(a) + return np.exp(self._logpdf(x, a)) + + def _logpdf(self, x, a): + return sc.xlogy(a-1.0, x) - x - sc.gammaln(a) + + def _cdf(self, x, a): + return sc.gammainc(a, x) + + def _sf(self, x, a): + return sc.gammaincc(a, x) + + def _ppf(self, q, a): + return sc.gammaincinv(a, q) + + def _isf(self, q, a): + return sc.gammainccinv(a, q) + + def _stats(self, a): + return a, a, 2.0/np.sqrt(a), 6.0/a + + def _munp(self, n, a): + return sc.poch(a, n) + + def _entropy(self, a): + + def regular_formula(a): + return sc.psi(a) * (1-a) + a + sc.gammaln(a) + + def asymptotic_formula(a): + # plug in above formula the expansions: + # psi(a) ~ ln(a) - 1/2a - 1/12a^2 + 1/120a^4 + # gammaln(a) ~ a * ln(a) - a - 1/2 * ln(a) + 1/2 ln(2 * pi) + + # 1/12a - 1/360a^3 + return (0.5 * (1. + np.log(2*np.pi) + np.log(a)) - 1/(3 * a) + - (a**-2.)/12 - (a**-3.)/90 + (a**-4.)/120) + + return xpx.apply_where(a < 250, a, regular_formula, asymptotic_formula) + + def _fitstart(self, data): + # The skewness of the gamma distribution is `2 / np.sqrt(a)`. + # We invert that to estimate the shape `a` using the skewness + # of the data. The formula is regularized with 1e-8 in the + # denominator to allow for degenerate data where the skewness + # is close to 0. + if isinstance(data, CensoredData): + data = data._uncensor() + sk = _skew(data) + a = 4 / (1e-8 + sk**2) + return super()._fitstart(data, args=(a,)) + + @extend_notes_in_docstring(rv_continuous, notes="""\ + When the location is fixed by using the argument `floc` + and `method='MLE'`, this + function uses explicit formulas or solves a simpler numerical + problem than the full ML optimization problem. So in that case, + the `optimizer`, `loc` and `scale` arguments are ignored. + \n\n""") + def fit(self, data, *args, **kwds): + floc = kwds.get('floc', None) + method = kwds.get('method', 'mle') + + if (isinstance(data, CensoredData) or + floc is None and method.lower() != 'mm'): + # loc is not fixed or we're not doing standard MLE. + # Use the default fit method. + return super().fit(data, *args, **kwds) + + # We already have this value, so just pop it from kwds. + kwds.pop('floc', None) + + f0 = _get_fixed_fit_value(kwds, ['f0', 'fa', 'fix_a']) + fscale = kwds.pop('fscale', None) + + _remove_optimizer_parameters(kwds) + + if f0 is not None and floc is not None and fscale is not None: + # This check is for consistency with `rv_continuous.fit`. + # Without this check, this function would just return the + # parameters that were given. + raise ValueError("All parameters fixed. There is nothing to " + "optimize.") + + # Fixed location is handled by shifting the data. + data = np.asarray(data) + + if not np.isfinite(data).all(): + raise ValueError("The data contains non-finite values.") + + # Use explicit formulas for mm (gh-19884) + if method.lower() == 'mm': + m1 = np.mean(data) + m2 = np.var(data) + m3 = np.mean((data - m1) ** 3) + a, loc, scale = f0, floc, fscale + # Three unknowns + if a is None and loc is None and scale is None: + scale = m3 / (2 * m2) + # Two unknowns + if loc is None and scale is None: + scale = np.sqrt(m2 / a) + if a is None and scale is None: + scale = m2 / (m1 - loc) + if a is None and loc is None: + a = m2 / (scale ** 2) + # One unknown + if a is None: + a = (m1 - loc) / scale + if loc is None: + loc = m1 - a * scale + if scale is None: + scale = (m1 - loc) / a + return a, loc, scale + + # Special case: loc is fixed. + + # NB: data == loc is ok if a >= 1; the below check is more strict. + if np.any(data <= floc): + raise FitDataError("gamma", lower=floc, upper=np.inf) + + if floc != 0: + # Don't do the subtraction in-place, because `data` might be a + # view of the input array. + data = data - floc + xbar = data.mean() + + # Three cases to handle: + # * shape and scale both free + # * shape fixed, scale free + # * shape free, scale fixed + + if fscale is None: + # scale is free + if f0 is not None: + # shape is fixed + a = f0 + else: + # shape and scale are both free. + # The MLE for the shape parameter `a` is the solution to: + # np.log(a) - sc.digamma(a) - np.log(xbar) + + # np.log(data).mean() = 0 + s = np.log(xbar) - np.log(data).mean() + aest = (3-s + np.sqrt((s-3)**2 + 24*s)) / (12*s) + xa = aest*(1-0.4) + xb = aest*(1+0.4) + a = optimize.brentq(lambda a: np.log(a) - sc.digamma(a) - s, + xa, xb, disp=0) + + # The MLE for the scale parameter is just the data mean + # divided by the shape parameter. + scale = xbar / a + else: + # scale is fixed, shape is free + # The MLE for the shape parameter `a` is the solution to: + # sc.digamma(a) - np.log(data).mean() + np.log(fscale) = 0 + c = np.log(data).mean() - np.log(fscale) + a = _digammainv(c) + scale = fscale + + return a, floc, scale + + +gamma = gamma_gen(a=0.0, name='gamma') + + +class erlang_gen(gamma_gen): + """An Erlang continuous random variable. + + %(before_notes)s + + See Also + -------- + gamma + + Notes + ----- + The Erlang distribution is a special case of the Gamma distribution, with + the shape parameter `a` an integer. Note that this restriction is not + enforced by `erlang`. It will, however, generate a warning the first time + a non-integer value is used for the shape parameter. + + Refer to `gamma` for examples. + + """ + + def _argcheck(self, a): + allint = np.all(np.floor(a) == a) + if not allint: + # An Erlang distribution shouldn't really have a non-integer + # shape parameter, so warn the user. + message = ('The shape parameter of the erlang distribution ' + f'has been given a non-integer value {a!r}.') + warnings.warn(message, RuntimeWarning, stacklevel=3) + return a > 0 + + def _shape_info(self): + return [_ShapeInfo("a", True, (1, np.inf), (True, False))] + + def _fitstart(self, data): + # Override gamma_gen_fitstart so that an integer initial value is + # used. (Also regularize the division, to avoid issues when + # _skew(data) is 0 or close to 0.) + if isinstance(data, CensoredData): + data = data._uncensor() + a = int(4.0 / (1e-8 + _skew(data)**2)) + return super(gamma_gen, self)._fitstart(data, args=(a,)) + + # Trivial override of the fit method, so we can monkey-patch its + # docstring. + @extend_notes_in_docstring(rv_continuous, notes="""\ + The Erlang distribution is generally defined to have integer values + for the shape parameter. This is not enforced by the `erlang` class. + When fitting the distribution, it will generally return a non-integer + value for the shape parameter. By using the keyword argument + `f0=`, the fit method can be constrained to fit the data to + a specific integer shape parameter.""") + def fit(self, data, *args, **kwds): + return super().fit(data, *args, **kwds) + + +erlang = erlang_gen(a=0.0, name='erlang') + + +class gengamma_gen(rv_continuous): + r"""A generalized gamma continuous random variable. + + %(before_notes)s + + See Also + -------- + gamma, halfgennorm, invgamma, weibull_min + + Notes + ----- + The probability density function for `gengamma` is ([1]_, [2]_): + + .. math:: + + f(x, a, c) = \frac{|c| x^{c a-1} \exp(-x^c)}{\Gamma(a)} + + for :math:`x \ge 0`, :math:`a > 0`, and :math:`c \ne 0`. + :math:`\Gamma` is the gamma function (`scipy.special.gamma`). + + `gengamma` takes :math:`a` and :math:`c` as shape parameters. + + The SciPy distribution `halfgennorm` is a special case of + `gengamma`: ``halfgennorm(beta) = gengamma(a=1/beta, c=beta)``. + + %(after_notes)s + + References + ---------- + .. [1] E.W. Stacy, "A Generalization of the Gamma Distribution", + Annals of Mathematical Statistics, Vol 33(3), pp. 1187--1192. + .. [2] "Generalized gamma distribution", Wikipedia, + https://en.wikipedia.org/wiki/Generalized_gamma_distribution + + %(example)s + + """ + def _argcheck(self, a, c): + return (a > 0) & (c != 0) + + def _shape_info(self): + ia = _ShapeInfo("a", False, (0, np.inf), (False, False)) + ic = _ShapeInfo("c", False, (-np.inf, np.inf), (False, False)) + return [ia, ic] + + def _pdf(self, x, a, c): + return np.exp(self._logpdf(x, a, c)) + + def _logpdf(self, x, a, c): + return xpx.apply_where( + (x != 0) | (c > 0), (x, c, a), + lambda x, c, a: (np.log(abs(c)) + sc.xlogy(c*a - 1, x) + - x**c - sc.gammaln(a)), + fill_value=-np.inf) + + def _cdf(self, x, a, c): + xc = x**c + val1 = sc.gammainc(a, xc) + val2 = sc.gammaincc(a, xc) + return np.where(c > 0, val1, val2) + + def _rvs(self, a, c, size=None, random_state=None): + r = random_state.standard_gamma(a, size=size) + return r**(1./c) + + def _sf(self, x, a, c): + xc = x**c + val1 = sc.gammainc(a, xc) + val2 = sc.gammaincc(a, xc) + return np.where(c > 0, val2, val1) + + def _ppf(self, q, a, c): + val1 = sc.gammaincinv(a, q) + val2 = sc.gammainccinv(a, q) + return np.where(c > 0, val1, val2)**(1.0/c) + + def _isf(self, q, a, c): + val1 = sc.gammaincinv(a, q) + val2 = sc.gammainccinv(a, q) + return np.where(c > 0, val2, val1)**(1.0/c) + + def _munp(self, n, a, c): + # Pochhammer symbol: sc.pocha,n) = gamma(a+n)/gamma(a) + return sc.poch(a, n*1.0/c) + + def _entropy(self, a, c): + def regular(a, c): + val = sc.psi(a) + A = a * (1 - val) + val / c + B = sc.gammaln(a) - np.log(abs(c)) + h = A + B + return h + + def asymptotic(a, c): + # using asymptotic expansions for gammaln and psi (see gh-18093) + return (norm._entropy() - np.log(a)/2 + - np.log(np.abs(c)) + (a**-1.)/6 - (a**-3.)/90 + + (np.log(a) - (a**-1.)/2 - (a**-2.)/12 + (a**-4.)/120)/c) + + return xpx.apply_where(a >= 200, (a, c), asymptotic, regular) + + +gengamma = gengamma_gen(a=0.0, name='gengamma') + + +class genhalflogistic_gen(rv_continuous): + r"""A generalized half-logistic continuous random variable. + + %(before_notes)s + + Notes + ----- + The probability density function for `genhalflogistic` is: + + .. math:: + + f(x, c) = \frac{2 (1 - c x)^{1/(c-1)}}{[1 + (1 - c x)^{1/c}]^2} + + for :math:`0 \le x \le 1/c`, and :math:`c > 0`. + + `genhalflogistic` takes ``c`` as a shape parameter for :math:`c`. + + %(after_notes)s + + %(example)s + + """ + def _shape_info(self): + return [_ShapeInfo("c", False, (0, np.inf), (False, False))] + + def _get_support(self, c): + return self.a, 1.0/c + + def _pdf(self, x, c): + # genhalflogistic.pdf(x, c) = + # 2 * (1-c*x)**(1/c-1) / (1+(1-c*x)**(1/c))**2 + limit = 1.0/c + tmp = np.asarray(1-c*x) + tmp0 = tmp**(limit-1) + tmp2 = tmp0*tmp + return 2*tmp0 / (1+tmp2)**2 + + def _cdf(self, x, c): + limit = 1.0/c + tmp = np.asarray(1-c*x) + tmp2 = tmp**(limit) + return (1.0-tmp2) / (1+tmp2) + + def _ppf(self, q, c): + return 1.0/c*(1-((1.0-q)/(1.0+q))**c) + + def _entropy(self, c): + return 2 - (2*c+1)*np.log(2) + + +genhalflogistic = genhalflogistic_gen(a=0.0, name='genhalflogistic') + + +class genhyperbolic_gen(rv_continuous): + r"""A generalized hyperbolic continuous random variable. + + %(before_notes)s + + See Also + -------- + t, norminvgauss, geninvgauss, laplace, cauchy + + Notes + ----- + The probability density function for `genhyperbolic` is: + + .. math:: + + f(x, p, a, b) = + \frac{(a^2 - b^2)^{p/2}} + {\sqrt{2\pi}a^{p-1/2} + K_p\Big(\sqrt{a^2 - b^2}\Big)} + e^{bx} \times \frac{K_{p - 1/2} + (a \sqrt{1 + x^2})} + {(\sqrt{1 + x^2})^{1/2 - p}} + + for :math:`x, p \in ( - \infty; \infty)`, + :math:`|b| < a` if :math:`p \ge 0`, + :math:`|b| \le a` if :math:`p < 0`. + :math:`K_{p}(.)` denotes the modified Bessel function of the second + kind and order :math:`p` (`scipy.special.kv`) + + `genhyperbolic` takes ``p`` as a tail parameter, + ``a`` as a shape parameter, + ``b`` as a skewness parameter. + + %(after_notes)s + + The original parameterization of the Generalized Hyperbolic Distribution + is found in [1]_ as follows + + .. math:: + + f(x, \lambda, \alpha, \beta, \delta, \mu) = + \frac{(\gamma/\delta)^\lambda}{\sqrt{2\pi}K_\lambda(\delta \gamma)} + e^{\beta (x - \mu)} \times \frac{K_{\lambda - 1/2} + (\alpha \sqrt{\delta^2 + (x - \mu)^2})} + {(\sqrt{\delta^2 + (x - \mu)^2} / \alpha)^{1/2 - \lambda}} + + for :math:`x \in ( - \infty; \infty)`, + :math:`\gamma := \sqrt{\alpha^2 - \beta^2}`, + :math:`\lambda, \mu \in ( - \infty; \infty)`, + :math:`\delta \ge 0, |\beta| < \alpha` if :math:`\lambda \ge 0`, + :math:`\delta > 0, |\beta| \le \alpha` if :math:`\lambda < 0`. + + The location-scale-based parameterization implemented in + SciPy is based on [2]_, where :math:`a = \alpha\delta`, + :math:`b = \beta\delta`, :math:`p = \lambda`, + :math:`scale=\delta` and :math:`loc=\mu` + + Moments are implemented based on [3]_ and [4]_. + + For the distributions that are a special case such as Student's t, + it is not recommended to rely on the implementation of genhyperbolic. + To avoid potential numerical problems and for performance reasons, + the methods of the specific distributions should be used. + + References + ---------- + .. [1] O. Barndorff-Nielsen, "Hyperbolic Distributions and Distributions + on Hyperbolae", Scandinavian Journal of Statistics, Vol. 5(3), + pp. 151-157, 1978. https://www.jstor.org/stable/4615705 + + .. [2] Eberlein E., Prause K. (2002) The Generalized Hyperbolic Model: + Financial Derivatives and Risk Measures. In: Geman H., Madan D., + Pliska S.R., Vorst T. (eds) Mathematical Finance - Bachelier + Congress 2000. Springer Finance. Springer, Berlin, Heidelberg. + :doi:`10.1007/978-3-662-12429-1_12` + + .. [3] Scott, David J, Würtz, Diethelm, Dong, Christine and Tran, + Thanh Tam, (2009), Moments of the generalized hyperbolic + distribution, MPRA Paper, University Library of Munich, Germany, + https://EconPapers.repec.org/RePEc:pra:mprapa:19081. + + .. [4] E. Eberlein and E. A. von Hammerstein. Generalized hyperbolic + and inverse Gaussian distributions: Limiting cases and approximation + of processes. FDM Preprint 80, April 2003. University of Freiburg. + https://freidok.uni-freiburg.de/fedora/objects/freidok:7974/datastreams/FILE1/content + + %(example)s + + """ + + def _argcheck(self, p, a, b): + return (np.logical_and(np.abs(b) < a, p >= 0) + | np.logical_and(np.abs(b) <= a, p < 0)) + + def _shape_info(self): + ip = _ShapeInfo("p", False, (-np.inf, np.inf), (False, False)) + ia = _ShapeInfo("a", False, (0, np.inf), (True, False)) + ib = _ShapeInfo("b", False, (-np.inf, np.inf), (False, False)) + return [ip, ia, ib] + + def _fitstart(self, data): + # Arbitrary, but the default p = a = b = 1 is not valid; the + # distribution requires |b| < a if p >= 0. + return super()._fitstart(data, args=(1, 1, 0.5)) + + def _logpdf(self, x, p, a, b): + # kve instead of kv works better for large values of p + # and smaller values of sqrt(a^2 - b^2) + @np.vectorize + def _logpdf_single(x, p, a, b): + return _stats.genhyperbolic_logpdf(x, p, a, b) + + return _logpdf_single(x, p, a, b) + + def _pdf(self, x, p, a, b): + # kve instead of kv works better for large values of p + # and smaller values of sqrt(a^2 - b^2) + @np.vectorize + def _pdf_single(x, p, a, b): + return _stats.genhyperbolic_pdf(x, p, a, b) + + return _pdf_single(x, p, a, b) + + # np.vectorize isn't currently designed to be used as a decorator, + # so use a lambda instead. This allows us to decorate the function + # with `np.vectorize` and still provide the `otypes` parameter. + @lambda func: np.vectorize(func, otypes=[np.float64]) + @staticmethod + def _integrate_pdf(x0, x1, p, a, b): + """ + Integrate the pdf of the genhyberbolic distribution from x0 to x1. + This is a private function used by _cdf() and _sf() only; either x0 + will be -inf or x1 will be inf. + """ + user_data = np.array([p, a, b], float).ctypes.data_as(ctypes.c_void_p) + llc = LowLevelCallable.from_cython(_stats, '_genhyperbolic_pdf', + user_data) + d = np.sqrt((a + b)*(a - b)) + mean = b/d * sc.kv(p + 1, d) / sc.kv(p, d) + epsrel = 1e-10 + epsabs = 0 + if x0 < mean < x1: + # If the interval includes the mean, integrate over the two + # intervals [x0, mean] and [mean, x1] and add. If we try to do + # the integral in one call of quad and the non-infinite endpoint + # is far in the tail, quad might return an incorrect result + # because it does not "see" the peak of the PDF. + intgrl = (integrate.quad(llc, x0, mean, + epsrel=epsrel, epsabs=epsabs)[0] + + integrate.quad(llc, mean, x1, + epsrel=epsrel, epsabs=epsabs)[0]) + else: + intgrl = integrate.quad(llc, x0, x1, + epsrel=epsrel, epsabs=epsabs)[0] + if np.isnan(intgrl): + msg = ("Infinite values encountered in scipy.special.kve. " + "Values replaced by NaN to avoid incorrect results.") + warnings.warn(msg, RuntimeWarning, stacklevel=3) + return max(0.0, min(1.0, intgrl)) + + def _cdf(self, x, p, a, b): + return self._integrate_pdf(-np.inf, x, p, a, b) + + def _sf(self, x, p, a, b): + return self._integrate_pdf(x, np.inf, p, a, b) + + def _rvs(self, p, a, b, size=None, random_state=None): + # note: X = b * V + sqrt(V) * X has a + # generalized hyperbolic distribution + # if X is standard normal and V is + # geninvgauss(p = p, b = t2, loc = loc, scale = t3) + t1 = np.float_power(a, 2) - np.float_power(b, 2) + # b in the GIG + t2 = np.float_power(t1, 0.5) + # scale in the GIG + t3 = np.float_power(t1, - 0.5) + gig = geninvgauss.rvs( + p=p, + b=t2, + scale=t3, + size=size, + random_state=random_state + ) + normst = norm.rvs(size=size, random_state=random_state) + + return b * gig + np.sqrt(gig) * normst + + def _stats(self, p, a, b): + # https://mpra.ub.uni-muenchen.de/19081/1/MPRA_paper_19081.pdf + # https://freidok.uni-freiburg.de/fedora/objects/freidok:7974/datastreams/FILE1/content + # standardized moments + p, a, b = np.broadcast_arrays(p, a, b) + t1 = np.float_power(a, 2) - np.float_power(b, 2) + t1 = np.float_power(t1, 0.5) + t2 = np.float_power(1, 2) * np.float_power(t1, - 1) + integers = np.linspace(0, 4, 5) + # make integers perpendicular to existing dimensions + integers = integers.reshape(integers.shape + (1,) * p.ndim) + b0, b1, b2, b3, b4 = sc.kv(p + integers, t1) + r1, r2, r3, r4 = (b / b0 for b in (b1, b2, b3, b4)) + + m = b * t2 * r1 + v = ( + t2 * r1 + np.float_power(b, 2) * np.float_power(t2, 2) * + (r2 - np.float_power(r1, 2)) + ) + m3e = ( + np.float_power(b, 3) * np.float_power(t2, 3) * + (r3 - 3 * b2 * b1 * np.float_power(b0, -2) + + 2 * np.float_power(r1, 3)) + + 3 * b * np.float_power(t2, 2) * + (r2 - np.float_power(r1, 2)) + ) + s = m3e * np.float_power(v, - 3 / 2) + m4e = ( + np.float_power(b, 4) * np.float_power(t2, 4) * + (r4 - 4 * b3 * b1 * np.float_power(b0, - 2) + + 6 * b2 * np.float_power(b1, 2) * np.float_power(b0, - 3) - + 3 * np.float_power(r1, 4)) + + np.float_power(b, 2) * np.float_power(t2, 3) * + (6 * r3 - 12 * b2 * b1 * np.float_power(b0, - 2) + + 6 * np.float_power(r1, 3)) + + 3 * np.float_power(t2, 2) * r2 + ) + k = m4e * np.float_power(v, -2) - 3 + + return m, v, s, k + + +genhyperbolic = genhyperbolic_gen(name='genhyperbolic') + + +class gompertz_gen(rv_continuous): + r"""A Gompertz (or truncated Gumbel) continuous random variable. + + %(before_notes)s + + Notes + ----- + The probability density function for `gompertz` is: + + .. math:: + + f(x, c) = c \exp(x) \exp(-c (e^x-1)) + + for :math:`x \ge 0`, :math:`c > 0`. + + `gompertz` takes ``c`` as a shape parameter for :math:`c`. + + %(after_notes)s + + %(example)s + + """ + def _shape_info(self): + return [_ShapeInfo("c", False, (0, np.inf), (False, False))] + + def _pdf(self, x, c): + # gompertz.pdf(x, c) = c * exp(x) * exp(-c*(exp(x)-1)) + return np.exp(self._logpdf(x, c)) + + def _logpdf(self, x, c): + return np.log(c) + x - c * sc.expm1(x) + + def _cdf(self, x, c): + return -sc.expm1(-c * sc.expm1(x)) + + def _ppf(self, q, c): + return sc.log1p(-1.0 / c * sc.log1p(-q)) + + def _sf(self, x, c): + return np.exp(-c * sc.expm1(x)) + + def _isf(self, p, c): + return sc.log1p(-np.log(p)/c) + + def _entropy(self, c): + return 1.0 - np.log(c) - sc._ufuncs._scaled_exp1(c)/c + + +gompertz = gompertz_gen(a=0.0, name='gompertz') + + +def _average_with_log_weights(x, logweights): + x = np.asarray(x) + logweights = np.asarray(logweights) + maxlogw = logweights.max() + weights = np.exp(logweights - maxlogw) + return np.average(x, weights=weights) + + +class gumbel_r_gen(rv_continuous): + r"""A right-skewed Gumbel continuous random variable. + + %(before_notes)s + + See Also + -------- + gumbel_l, gompertz, genextreme + + Notes + ----- + The probability density function for `gumbel_r` is: + + .. math:: + + f(x) = \exp(-(x + e^{-x})) + + for real :math:`x`. + + The Gumbel distribution is sometimes referred to as a type I Fisher-Tippett + distribution. It is also related to the extreme value distribution, + log-Weibull and Gompertz distributions. + + %(after_notes)s + + %(example)s + + """ + def _shape_info(self): + return [] + + def _pdf(self, x): + # gumbel_r.pdf(x) = exp(-(x + exp(-x))) + return np.exp(self._logpdf(x)) + + def _logpdf(self, x): + return -x - np.exp(-x) + + def _cdf(self, x): + return np.exp(-np.exp(-x)) + + def _logcdf(self, x): + return -np.exp(-x) + + def _ppf(self, q): + return -np.log(-np.log(q)) + + def _sf(self, x): + return -sc.expm1(-np.exp(-x)) + + def _isf(self, p): + return -np.log(-np.log1p(-p)) + + def _stats(self): + return _EULER, np.pi*np.pi/6.0, 12*np.sqrt(6)/np.pi**3 * _ZETA3, 12.0/5 + + def _entropy(self): + # https://en.wikipedia.org/wiki/Gumbel_distribution + return _EULER + 1. + + @_call_super_mom + @inherit_docstring_from(rv_continuous) + def fit(self, data, *args, **kwds): + data, floc, fscale = _check_fit_input_parameters(self, data, + args, kwds) + + # By the method of maximum likelihood, the estimators of the + # location and scale are the roots of the equations defined in + # `func` and the value of the expression for `loc` that follows. + # The first `func` is a first order derivative of the log-likelihood + # equation and the second is from Source: Statistical Distributions, + # 3rd Edition. Evans, Hastings, and Peacock (2000), Page 101. + + def get_loc_from_scale(scale): + return -scale * (sc.logsumexp(-data / scale) - np.log(len(data))) + + if fscale is not None: + # if the scale is fixed, the location can be analytically + # determined. + scale = fscale + loc = get_loc_from_scale(scale) + else: + # A different function is solved depending on whether the location + # is fixed. + if floc is not None: + loc = floc + + # equation to use if the location is fixed. + # note that one cannot use the equation in Evans, Hastings, + # and Peacock (2000) (since it assumes that the derivative + # w.r.t. the log-likelihood is zero). however, it is easy to + # derive the MLE condition directly if loc is fixed + def func(scale): + term1 = (loc - data) * np.exp((loc - data) / scale) + data + term2 = len(data) * (loc + scale) + return term1.sum() - term2 + else: + + # equation to use if both location and scale are free + def func(scale): + sdata = -data / scale + wavg = _average_with_log_weights(data, logweights=sdata) + return data.mean() - wavg - scale + + # set brackets for `root_scalar` to use when optimizing over the + # scale such that a root is likely between them. Use user supplied + # guess or default 1. + brack_start = kwds.get('scale', 1) + lbrack, rbrack = brack_start / 2, brack_start * 2 + + # if a root is not between the brackets, iteratively expand them + # until they include a sign change, checking after each bracket is + # modified. + def interval_contains_root(lbrack, rbrack): + # return true if the signs disagree. + return (np.sign(func(lbrack)) != + np.sign(func(rbrack))) + while (not interval_contains_root(lbrack, rbrack) + and (lbrack > 0 or rbrack < np.inf)): + lbrack /= 2 + rbrack *= 2 + + res = optimize.root_scalar(func, bracket=(lbrack, rbrack), + rtol=1e-14, xtol=1e-14) + scale = res.root + loc = floc if floc is not None else get_loc_from_scale(scale) + return loc, scale + + +gumbel_r = gumbel_r_gen(name='gumbel_r') + + +class gumbel_l_gen(rv_continuous): + r"""A left-skewed Gumbel continuous random variable. + + %(before_notes)s + + See Also + -------- + gumbel_r, gompertz, genextreme + + Notes + ----- + The probability density function for `gumbel_l` is: + + .. math:: + + f(x) = \exp(x - e^x) + + for real :math:`x`. + + The Gumbel distribution is sometimes referred to as a type I Fisher-Tippett + distribution. It is also related to the extreme value distribution, + log-Weibull and Gompertz distributions. + + %(after_notes)s + + %(example)s + + """ + + def _shape_info(self): + return [] + + def _pdf(self, x): + # gumbel_l.pdf(x) = exp(x - exp(x)) + return np.exp(self._logpdf(x)) + + def _logpdf(self, x): + return x - np.exp(x) + + def _cdf(self, x): + return -sc.expm1(-np.exp(x)) + + def _ppf(self, q): + return np.log(-sc.log1p(-q)) + + def _logsf(self, x): + return -np.exp(x) + + def _sf(self, x): + return np.exp(-np.exp(x)) + + def _isf(self, x): + return np.log(-np.log(x)) + + def _stats(self): + return -_EULER, np.pi*np.pi/6.0, \ + -12*np.sqrt(6)/np.pi**3 * _ZETA3, 12.0/5 + + def _entropy(self): + return _EULER + 1. + + @_call_super_mom + @inherit_docstring_from(rv_continuous) + def fit(self, data, *args, **kwds): + # The fit method of `gumbel_r` can be used for this distribution with + # small modifications. The process to do this is + # 1. pass the sign negated data into `gumbel_r.fit` + # - if the location is fixed, it should also be negated. + # 2. negate the sign of the resulting location, leaving the scale + # unmodified. + # `gumbel_r.fit` holds necessary input checks. + + if kwds.get('floc') is not None: + kwds['floc'] = -kwds['floc'] + loc_r, scale_r, = gumbel_r.fit(-np.asarray(data), *args, **kwds) + return -loc_r, scale_r + + +gumbel_l = gumbel_l_gen(name='gumbel_l') + + +class halfcauchy_gen(rv_continuous): + r"""A Half-Cauchy continuous random variable. + + %(before_notes)s + + Notes + ----- + The probability density function for `halfcauchy` is: + + .. math:: + + f(x) = \frac{2}{\pi (1 + x^2)} + + for :math:`x \ge 0`. + + %(after_notes)s + + %(example)s + + """ + def _shape_info(self): + return [] + + def _pdf(self, x): + # halfcauchy.pdf(x) = 2 / (pi * (1 + x**2)) + return 2.0/np.pi/(1.0+x*x) + + def _logpdf(self, x): + return np.log(2.0/np.pi) - sc.log1p(x*x) + + def _cdf(self, x): + return 2.0/np.pi*np.arctan(x) + + def _ppf(self, q): + return np.tan(np.pi/2*q) + + def _sf(self, x): + return 2.0/np.pi * np.arctan2(1, x) + + def _isf(self, p): + return 1.0/np.tan(np.pi*p/2) + + def _stats(self): + return np.inf, np.inf, np.nan, np.nan + + def _entropy(self): + return np.log(2*np.pi) + + @_call_super_mom + @inherit_docstring_from(rv_continuous) + def fit(self, data, *args, **kwds): + if kwds.pop('superfit', False): + return super().fit(data, *args, **kwds) + + data, floc, fscale = _check_fit_input_parameters(self, data, + args, kwds) + + # location is independent from the scale + data_min = np.min(data) + if floc is not None: + if data_min < floc: + # There are values that are less than the specified loc. + raise FitDataError("halfcauchy", lower=floc, upper=np.inf) + loc = floc + else: + # if not provided, location MLE is the minimal data point + loc = data_min + + # find scale + def find_scale(loc, data): + shifted_data = data - loc + n = data.size + shifted_data_squared = np.square(shifted_data) + + def fun_to_solve(scale): + denominator = scale**2 + shifted_data_squared + return 2 * np.sum(shifted_data_squared/denominator) - n + + small = np.finfo(1.0).tiny**0.5 # avoid underflow + res = root_scalar(fun_to_solve, bracket=(small, np.max(shifted_data))) + return res.root + + if fscale is not None: + scale = fscale + else: + scale = find_scale(loc, data) + + return loc, scale + + +halfcauchy = halfcauchy_gen(a=0.0, name='halfcauchy') + + +class halflogistic_gen(rv_continuous): + r"""A half-logistic continuous random variable. + + %(before_notes)s + + Notes + ----- + The probability density function for `halflogistic` is: + + .. math:: + + f(x) = \frac{ 2 e^{-x} }{ (1+e^{-x})^2 } + = \frac{1}{2} \text{sech}(x/2)^2 + + for :math:`x \ge 0`. + + %(after_notes)s + + References + ---------- + .. [1] Asgharzadeh et al (2011). "Comparisons of Methods of Estimation for the + Half-Logistic Distribution". Selcuk J. Appl. Math. 93-108. + + %(example)s + + """ + def _shape_info(self): + return [] + + def _pdf(self, x): + # halflogistic.pdf(x) = 2 * exp(-x) / (1+exp(-x))**2 + # = 1/2 * sech(x/2)**2 + return np.exp(self._logpdf(x)) + + def _logpdf(self, x): + return np.log(2) - x - 2. * sc.log1p(np.exp(-x)) + + def _cdf(self, x): + return np.tanh(x/2.0) + + def _ppf(self, q): + return 2*np.arctanh(q) + + def _sf(self, x): + return 2 * sc.expit(-x) + + def _isf(self, q): + return xpx.apply_where(q < 0.5, q, + lambda q: -sc.logit(0.5 * q), + lambda q: 2*np.arctanh(1 - q)) + + def _munp(self, n): + if n == 0: + return 1 # otherwise returns NaN + if n == 1: + return 2*np.log(2) + if n == 2: + return np.pi*np.pi/3.0 + if n == 3: + return 9*_ZETA3 + if n == 4: + return 7*np.pi**4 / 15.0 + return 2*(1-pow(2.0, 1-n))*sc.gamma(n+1)*sc.zeta(n, 1) + + def _entropy(self): + return 2-np.log(2) + + @_call_super_mom + @inherit_docstring_from(rv_continuous) + def fit(self, data, *args, **kwds): + if kwds.pop('superfit', False): + return super().fit(data, *args, **kwds) + + data, floc, fscale = _check_fit_input_parameters(self, data, + args, kwds) + + def find_scale(data, loc): + # scale is solution to a fix point problem ([1] 2.6) + # use approximate MLE as starting point ([1] 3.1) + n_observations = data.shape[0] + sorted_data = np.sort(data, axis=0) + p = np.arange(1, n_observations + 1)/(n_observations + 1) + q = 1 - p + pp1 = 1 + p + alpha = p - 0.5 * q * pp1 * np.log(pp1 / q) + beta = 0.5 * q * pp1 + sorted_data = sorted_data - loc + B = 2 * np.sum(alpha[1:] * sorted_data[1:]) + C = 2 * np.sum(beta[1:] * sorted_data[1:]**2) + # starting guess + scale = ((B + np.sqrt(B**2 + 8 * n_observations * C)) + /(4 * n_observations)) + + # relative tolerance of fix point iterator + rtol = 1e-8 + relative_residual = 1 + shifted_mean = sorted_data.mean() # y_mean - y_min + + # find fix point by repeated application of eq. (2.6) + # simplify as + # exp(-x) / (1 + exp(-x)) = 1 / (1 + exp(x)) + # = expit(-x)) + while relative_residual > rtol: + sum_term = sorted_data * sc.expit(-sorted_data/scale) + scale_new = shifted_mean - 2/n_observations * sum_term.sum() + relative_residual = abs((scale - scale_new)/scale) + scale = scale_new + return scale + + # location is independent from the scale + data_min = np.min(data) + if floc is not None: + if data_min < floc: + # There are values that are less than the specified loc. + raise FitDataError("halflogistic", lower=floc, upper=np.inf) + loc = floc + else: + # if not provided, location MLE is the minimal data point + loc = data_min + + # scale depends on location + scale = fscale if fscale is not None else find_scale(data, loc) + + return loc, scale + + +halflogistic = halflogistic_gen(a=0.0, name='halflogistic') + + +class halfnorm_gen(rv_continuous): + r"""A half-normal continuous random variable. + + %(before_notes)s + + Notes + ----- + The probability density function for `halfnorm` is: + + .. math:: + + f(x) = \sqrt{2/\pi} \exp(-x^2 / 2) + + for :math:`x >= 0`. + + `halfnorm` is a special case of `chi` with ``df=1``. + + %(after_notes)s + + %(example)s + + """ + def _shape_info(self): + return [] + + def _rvs(self, size=None, random_state=None): + return abs(random_state.standard_normal(size=size)) + + def _pdf(self, x): + # halfnorm.pdf(x) = sqrt(2/pi) * exp(-x**2/2) + return np.sqrt(2.0/np.pi)*np.exp(-x*x/2.0) + + def _logpdf(self, x): + return 0.5 * np.log(2.0/np.pi) - x*x/2.0 + + def _cdf(self, x): + return sc.erf(x / np.sqrt(2)) + + def _ppf(self, q): + return _norm_ppf((1+q)/2.0) + + def _sf(self, x): + return 2 * _norm_sf(x) + + def _isf(self, p): + return _norm_isf(p/2) + + def _stats(self): + return (np.sqrt(2.0/np.pi), + 1-2.0/np.pi, + np.sqrt(2)*(4-np.pi)/(np.pi-2)**1.5, + 8*(np.pi-3)/(np.pi-2)**2) + + def _entropy(self): + return 0.5*np.log(np.pi/2.0)+0.5 + + @_call_super_mom + @inherit_docstring_from(rv_continuous) + def fit(self, data, *args, **kwds): + if kwds.pop('superfit', False): + return super().fit(data, *args, **kwds) + + data, floc, fscale = _check_fit_input_parameters(self, data, + args, kwds) + + data_min = np.min(data) + + if floc is not None: + if data_min < floc: + # There are values that are less than the specified loc. + raise FitDataError("halfnorm", lower=floc, upper=np.inf) + loc = floc + else: + loc = data_min + + if fscale is not None: + scale = fscale + else: + scale = stats.moment(data, order=2, center=loc)**0.5 + + return loc, scale + + +halfnorm = halfnorm_gen(a=0.0, name='halfnorm') + + +class hypsecant_gen(rv_continuous): + r"""A hyperbolic secant continuous random variable. + + %(before_notes)s + + Notes + ----- + The probability density function for `hypsecant` is: + + .. math:: + + f(x) = \frac{1}{\pi} \text{sech}(x) + + for a real number :math:`x`. + + %(after_notes)s + + %(example)s + + """ + def _shape_info(self): + return [] + + def _pdf(self, x): + # hypsecant.pdf(x) = 1/pi * sech(x) + return 1.0/(np.pi*np.cosh(x)) + + def _cdf(self, x): + return 2.0/np.pi*np.arctan(np.exp(x)) + + def _ppf(self, q): + return np.log(np.tan(np.pi*q/2.0)) + + def _sf(self, x): + return 2.0/np.pi*np.arctan(np.exp(-x)) + + def _isf(self, q): + return -np.log(np.tan(np.pi*q/2.0)) + + def _stats(self): + return 0, np.pi*np.pi/4, 0, 2 + + def _entropy(self): + return np.log(2*np.pi) + + +hypsecant = hypsecant_gen(name='hypsecant') + + +class gausshyper_gen(rv_continuous): + r"""A Gauss hypergeometric continuous random variable. + + %(before_notes)s + + Notes + ----- + The probability density function for `gausshyper` is: + + .. math:: + + f(x, a, b, c, z) = C x^{a-1} (1-x)^{b-1} (1+zx)^{-c} + + for :math:`0 \le x \le 1`, :math:`a,b > 0`, :math:`c` a real number, + :math:`z > -1`, and :math:`C = \frac{1}{B(a, b) F[2, 1](c, a; a+b; -z)}`. + :math:`F[2, 1]` is the Gauss hypergeometric function + `scipy.special.hyp2f1`. + + `gausshyper` takes :math:`a`, :math:`b`, :math:`c` and :math:`z` as shape + parameters. + + %(after_notes)s + + References + ---------- + .. [1] Armero, C., and M. J. Bayarri. "Prior Assessments for Prediction in + Queues." *Journal of the Royal Statistical Society*. Series D (The + Statistician) 43, no. 1 (1994): 139-53. :doi:`10.2307/2348939`. + + %(example)s + + """ + + def _argcheck(self, a, b, c, z): + # z > -1 per gh-10134 + return (a > 0) & (b > 0) & (c == c) & (z > -1) + + def _shape_info(self): + ia = _ShapeInfo("a", False, (0, np.inf), (False, False)) + ib = _ShapeInfo("b", False, (0, np.inf), (False, False)) + ic = _ShapeInfo("c", False, (-np.inf, np.inf), (False, False)) + iz = _ShapeInfo("z", False, (-1, np.inf), (False, False)) + return [ia, ib, ic, iz] + + def _pdf(self, x, a, b, c, z): + normalization_constant = sc.beta(a, b) * sc.hyp2f1(c, a, a + b, -z) + return (1./normalization_constant * x**(a - 1.) * (1. - x)**(b - 1.0) + / (1.0 + z*x)**c) + + def _munp(self, n, a, b, c, z): + fac = sc.beta(n+a, b) / sc.beta(a, b) + num = sc.hyp2f1(c, a+n, a+b+n, -z) + den = sc.hyp2f1(c, a, a+b, -z) + return fac*num / den + + +gausshyper = gausshyper_gen(a=0.0, b=1.0, name='gausshyper') + + +class invgamma_gen(rv_continuous): + r"""An inverted gamma continuous random variable. + + %(before_notes)s + + Notes + ----- + The probability density function for `invgamma` is: + + .. math:: + + f(x, a) = \frac{x^{-a-1}}{\Gamma(a)} \exp(-\frac{1}{x}) + + for :math:`x >= 0`, :math:`a > 0`. :math:`\Gamma` is the gamma function + (`scipy.special.gamma`). + + `invgamma` takes ``a`` as a shape parameter for :math:`a`. + + `invgamma` is a special case of `gengamma` with ``c=-1``, and it is a + different parameterization of the scaled inverse chi-squared distribution. + Specifically, if the scaled inverse chi-squared distribution is + parameterized with degrees of freedom :math:`\nu` and scaling parameter + :math:`\tau^2`, then it can be modeled using `invgamma` with + ``a=`` :math:`\nu/2` and ``scale=`` :math:`\nu \tau^2/2`. + + %(after_notes)s + + %(example)s + + """ + _support_mask = rv_continuous._open_support_mask + + def _shape_info(self): + return [_ShapeInfo("a", False, (0, np.inf), (False, False))] + + def _pdf(self, x, a): + # invgamma.pdf(x, a) = x**(-a-1) / gamma(a) * exp(-1/x) + return np.exp(self._logpdf(x, a)) + + def _logpdf(self, x, a): + return -(a+1) * np.log(x) - sc.gammaln(a) - 1.0/x + + def _cdf(self, x, a): + return sc.gammaincc(a, 1.0 / x) + + def _ppf(self, q, a): + return 1.0 / sc.gammainccinv(a, q) + + def _sf(self, x, a): + return sc.gammainc(a, 1.0 / x) + + def _isf(self, q, a): + return 1.0 / sc.gammaincinv(a, q) + + def _stats(self, a, moments='mvsk'): + m1 = xpx.apply_where(a > 1, a, + lambda x: 1. / (x - 1.), + fill_value=np.inf) + m2 = xpx.apply_where(a > 2, a, + lambda x: 1. / (x - 1.)**2 / (x - 2.), + fill_value=np.inf) + + g1, g2 = None, None + if 's' in moments: + g1 = xpx.apply_where(a > 3, a, + lambda x: 4. * np.sqrt(x - 2.) / (x - 3.), + fill_value=np.nan) + if 'k' in moments: + g2 = xpx.apply_where(a > 4, a, + lambda x: 6. * (5. * x - 11.) / (x - 3.) / (x - 4.), + fill_value=np.nan) + + return m1, m2, g1, g2 + + def _entropy(self, a): + def regular(a): + h = a - (a + 1.0) * sc.psi(a) + sc.gammaln(a) + return h + + def asymptotic(a): + # gammaln(a) ~ a * ln(a) - a - 0.5 * ln(a) + 0.5 * ln(2 * pi) + # psi(a) ~ ln(a) - 1 / (2 * a) + h = ((1 - 3*np.log(a) + np.log(2) + np.log(np.pi))/2 + + 2/3*a**-1. + a**-2./12 - a**-3./90 - a**-4./120) + return h + + h = xpx.apply_where(a >= 200, a, asymptotic, regular) + return h + + +invgamma = invgamma_gen(a=0.0, name='invgamma') + + +class invgauss_gen(rv_continuous): + r"""An inverse Gaussian continuous random variable. + + %(before_notes)s + + Notes + ----- + The probability density function for `invgauss` is: + + .. math:: + + f(x; \mu) = \frac{1}{\sqrt{2 \pi x^3}} + \exp\left(-\frac{(x-\mu)^2}{2 \mu^2 x}\right) + + for :math:`x \ge 0` and :math:`\mu > 0`. + + `invgauss` takes ``mu`` as a shape parameter for :math:`\mu`. + + %(after_notes)s + + A common shape-scale parameterization of the inverse Gaussian distribution + has density + + .. math:: + + f(x; \nu, \lambda) = \sqrt{\frac{\lambda}{2 \pi x^3}} + \exp\left( -\frac{\lambda(x-\nu)^2}{2 \nu^2 x}\right) + + Using ``nu`` for :math:`\nu` and ``lam`` for :math:`\lambda`, this + parameterization is equivalent to the one above with ``mu = nu/lam``, + ``loc = 0``, and ``scale = lam``. + + This distribution uses routines from the Boost Math C++ library for + the computation of the ``ppf`` and ``isf`` methods. [1]_ + + References + ---------- + .. [1] The Boost Developers. "Boost C++ Libraries". https://www.boost.org/. + + %(example)s + + """ + _support_mask = rv_continuous._open_support_mask + + def _shape_info(self): + return [_ShapeInfo("mu", False, (0, np.inf), (False, False))] + + def _rvs(self, mu, size=None, random_state=None): + return random_state.wald(mu, 1.0, size=size) + + def _pdf(self, x, mu): + # invgauss.pdf(x, mu) = + # 1 / sqrt(2*pi*x**3) * exp(-(x-mu)**2/(2*x*mu**2)) + return 1.0/np.sqrt(2*np.pi*x**3.0)*np.exp(-1.0/(2*x)*(x/mu - 1)**2) + + def _logpdf(self, x, mu): + return -0.5*np.log(2*np.pi) - 1.5*np.log(x) - (x/mu - 1)**2/(2*x) + + # approach adapted from equations in + # https://journal.r-project.org/archive/2016-1/giner-smyth.pdf, + # not R code. see gh-13616 + + def _logcdf(self, x, mu): + fac = 1 / np.sqrt(x) + a = _norm_logcdf(fac * (x/mu - 1)) + b = 2 / mu + _norm_logcdf(-fac * (x/mu + 1)) + return a + np.log1p(np.exp(b - a)) + + def _logsf(self, x, mu): + fac = 1 / np.sqrt(x) + a = _norm_logsf(fac * (x/mu - 1)) + b = 2 / mu + _norm_logcdf(-fac * (x/mu + 1)) + return a + np.log1p(-np.exp(b - a)) + + def _sf(self, x, mu): + return np.exp(self._logsf(x, mu)) + + def _cdf(self, x, mu): + return np.exp(self._logcdf(x, mu)) + + def _ppf(self, x, mu): + with np.errstate(divide='ignore', over='ignore', invalid='ignore'): + x, mu = np.broadcast_arrays(x, mu) + ppf = np.asarray(scu._invgauss_ppf(x, mu, 1)) + i_wt = x > 0.5 # "wrong tail" - sometimes too inaccurate + ppf[i_wt] = scu._invgauss_isf(1-x[i_wt], mu[i_wt], 1) + i_nan = np.isnan(ppf) + ppf[i_nan] = super()._ppf(x[i_nan], mu[i_nan]) + return ppf + + def _isf(self, x, mu): + with np.errstate(divide='ignore', over='ignore', invalid='ignore'): + x, mu = np.broadcast_arrays(x, mu) + isf = scu._invgauss_isf(x, mu, 1) + i_wt = x > 0.5 # "wrong tail" - sometimes too inaccurate + isf[i_wt] = scu._invgauss_ppf(1-x[i_wt], mu[i_wt], 1) + i_nan = np.isnan(isf) + isf[i_nan] = super()._isf(x[i_nan], mu[i_nan]) + return isf + + def _stats(self, mu): + return mu, mu**3.0, 3*np.sqrt(mu), 15*mu + + @inherit_docstring_from(rv_continuous) + def fit(self, data, *args, **kwds): + method = kwds.get('method', 'mle') + + if (isinstance(data, CensoredData) or isinstance(self, wald_gen) + or method.lower() == 'mm'): + return super().fit(data, *args, **kwds) + + data, fshape_s, floc, fscale = _check_fit_input_parameters(self, data, + args, kwds) + ''' + Source: Statistical Distributions, 3rd Edition. Evans, Hastings, + and Peacock (2000), Page 121. Their shape parameter is equivalent to + SciPy's with the conversion `fshape_s = fshape / scale`. + + MLE formulas are not used in 3 conditions: + + - `loc` is not fixed + - `mu` is fixed + + These cases fall back on the superclass fit method. + + - `loc` is fixed but translation results in negative data raises + a `FitDataError`. + ''' + if floc is None or fshape_s is not None: + return super().fit(data, *args, **kwds) + elif np.any(data - floc < 0): + raise FitDataError("invgauss", lower=0, upper=np.inf) + else: + data = data - floc + fshape_n = np.mean(data) + if fscale is None: + fscale = len(data) / (np.sum(data ** -1 - fshape_n ** -1)) + fshape_s = fshape_n / fscale + return fshape_s, floc, fscale + + def _entropy(self, mu): + """ + Ref.: https://moser-isi.ethz.ch/docs/papers/smos-2012-10.pdf (eq. 9) + """ + # a = log(2*pi*e*mu**3) + # = 1 + log(2*pi) + 3 * log(mu) + a = 1. + np.log(2 * np.pi) + 3 * np.log(mu) + # b = exp(2/mu) * exp1(2/mu) + # = _scaled_exp1(2/mu) / (2/mu) + r = 2/mu + b = sc._ufuncs._scaled_exp1(r)/r + return 0.5 * a - 1.5 * b + + +invgauss = invgauss_gen(a=0.0, name='invgauss') + + +class geninvgauss_gen(rv_continuous): + r"""A Generalized Inverse Gaussian continuous random variable. + + %(before_notes)s + + Notes + ----- + The probability density function for `geninvgauss` is: + + .. math:: + + f(x, p, b) = x^{p-1} \exp(-b (x + 1/x) / 2) / (2 K_p(b)) + + where ``x > 0``, `p` is a real number and ``b > 0``\([1]_). + :math:`K_p` is the modified Bessel function of second kind of order `p` + (`scipy.special.kv`). + + %(after_notes)s + + The inverse Gaussian distribution `stats.invgauss(mu)` is a special case of + `geninvgauss` with ``p = -1/2``, ``b = 1 / mu`` and ``scale = mu``. + + Generating random variates is challenging for this distribution. The + implementation is based on [2]_. + + References + ---------- + .. [1] O. Barndorff-Nielsen, P. Blaesild, C. Halgreen, "First hitting time + models for the generalized inverse gaussian distribution", + Stochastic Processes and their Applications 7, pp. 49--54, 1978. + + .. [2] W. Hoermann and J. Leydold, "Generating generalized inverse Gaussian + random variates", Statistics and Computing, 24(4), p. 547--557, 2014. + + %(example)s + + """ + def _argcheck(self, p, b): + return (p == p) & (b > 0) + + def _shape_info(self): + ip = _ShapeInfo("p", False, (-np.inf, np.inf), (False, False)) + ib = _ShapeInfo("b", False, (0, np.inf), (False, False)) + return [ip, ib] + + def _logpdf(self, x, p, b): + # kve instead of kv works better for large values of b + # warn if kve produces infinite values and replace by nan + # otherwise c = -inf and the results are often incorrect + def logpdf_single(x, p, b): + return _stats.geninvgauss_logpdf(x, p, b) + + logpdf_single = np.vectorize(logpdf_single, otypes=[np.float64]) + + z = logpdf_single(x, p, b) + if np.isnan(z).any(): + msg = ("Infinite values encountered in scipy.special.kve(p, b). " + "Values replaced by NaN to avoid incorrect results.") + warnings.warn(msg, RuntimeWarning, stacklevel=3) + return z + + def _pdf(self, x, p, b): + # relying on logpdf avoids overflow of x**(p-1) for large x and p + return np.exp(self._logpdf(x, p, b)) + + def _cdf(self, x, p, b): + _a, _b = self._get_support(p, b) + + def _cdf_single(x, p, b): + user_data = np.array([p, b], float).ctypes.data_as(ctypes.c_void_p) + llc = LowLevelCallable.from_cython(_stats, '_geninvgauss_pdf', + user_data) + + return integrate.quad(llc, _a, x)[0] + + _cdf_single = np.vectorize(_cdf_single, otypes=[np.float64]) + + return _cdf_single(x, p, b) + + def _logquasipdf(self, x, p, b): + # log of the quasi-density (w/o normalizing constant) used in _rvs + return xpx.apply_where(x > 0, (x, p, b), + lambda x, p, b: (p - 1)*np.log(x) - b*(x + 1/x)/2, + fill_value=-np.inf) + + def _rvs(self, p, b, size=None, random_state=None): + # if p and b are scalar, use _rvs_scalar, otherwise need to create + # output by iterating over parameters + if np.isscalar(p) and np.isscalar(b): + out = self._rvs_scalar(p, b, size, random_state) + elif p.size == 1 and b.size == 1: + out = self._rvs_scalar(p.item(), b.item(), size, random_state) + else: + # When this method is called, size will be a (possibly empty) + # tuple of integers. It will not be None; if `size=None` is passed + # to `rvs()`, size will be the empty tuple (). + + p, b = np.broadcast_arrays(p, b) + # p and b now have the same shape. + + # `shp` is the shape of the blocks of random variates that are + # generated for each combination of parameters associated with + # broadcasting p and b. + # bc is a tuple the same length as size. The values + # in bc are bools. If bc[j] is True, it means that + # entire axis is filled in for a given combination of the + # broadcast arguments. + shp, bc = _check_shape(p.shape, size) + + # `numsamples` is the total number of variates to be generated + # for each combination of the input arguments. + numsamples = int(np.prod(shp)) + + # `out` is the array to be returned. It is filled in the + # loop below. + out = np.empty(size) + + it = np.nditer([p, b], + flags=['multi_index'], + op_flags=[['readonly'], ['readonly']]) + while not it.finished: + # Convert the iterator's multi_index into an index into the + # `out` array where the call to _rvs_scalar() will be stored. + # Where bc is True, we use a full slice; otherwise we use the + # index value from it.multi_index. len(it.multi_index) might + # be less than len(bc), and in that case we want to align these + # two sequences to the right, so the loop variable j runs from + # -len(size) to 0. This doesn't cause an IndexError, as + # bc[j] will be True in those cases where it.multi_index[j] + # would cause an IndexError. + idx = tuple((it.multi_index[j] if not bc[j] else slice(None)) + for j in range(-len(size), 0)) + out[idx] = self._rvs_scalar(it[0], it[1], numsamples, + random_state).reshape(shp) + it.iternext() + + if size == (): + out = out.item() + return out + + def _rvs_scalar(self, p, b, numsamples, random_state): + # following [2], the quasi-pdf is used instead of the pdf for the + # generation of rvs + invert_res = False + if not numsamples: + numsamples = 1 + if p < 0: + # note: if X is geninvgauss(p, b), then 1/X is geninvgauss(-p, b) + p = -p + invert_res = True + m = self._mode(p, b) + + # determine method to be used following [2] + ratio_unif = True + if p >= 1 or b > 1: + # ratio of uniforms with mode shift below + mode_shift = True + elif b >= min(0.5, 2 * np.sqrt(1 - p) / 3): + # ratio of uniforms without mode shift below + mode_shift = False + else: + # new algorithm in [2] + ratio_unif = False + + # prepare sampling of rvs + size1d = tuple(np.atleast_1d(numsamples)) + N = np.prod(size1d) # number of rvs needed, reshape upon return + x = np.zeros(N) + simulated = 0 + + if ratio_unif: + # use ratio of uniforms method + if mode_shift: + a2 = -2 * (p + 1) / b - m + a1 = 2 * m * (p - 1) / b - 1 + # find roots of x**3 + a2*x**2 + a1*x + m (Cardano's formula) + p1 = a1 - a2**2 / 3 + q1 = 2 * a2**3 / 27 - a2 * a1 / 3 + m + phi = np.arccos(-q1 * np.sqrt(-27 / p1**3) / 2) + s1 = -np.sqrt(-4 * p1 / 3) + root1 = s1 * np.cos(phi / 3 + np.pi / 3) - a2 / 3 + root2 = -s1 * np.cos(phi / 3) - a2 / 3 + # root3 = s1 * np.cos(phi / 3 - np.pi / 3) - a2 / 3 + + # if g is the quasipdf, rescale: g(x) / g(m) which we can write + # as exp(log(g(x)) - log(g(m))). This is important + # since for large values of p and b, g cannot be evaluated. + # denote the rescaled quasipdf by h + lm = self._logquasipdf(m, p, b) + d1 = self._logquasipdf(root1, p, b) - lm + d2 = self._logquasipdf(root2, p, b) - lm + # compute the bounding rectangle w.r.t. h. Note that + # np.exp(0.5*d1) = np.sqrt(g(root1)/g(m)) = np.sqrt(h(root1)) + vmin = (root1 - m) * np.exp(0.5 * d1) + vmax = (root2 - m) * np.exp(0.5 * d2) + umax = 1 # umax = sqrt(h(m)) = 1 + + def logqpdf(x): + return self._logquasipdf(x, p, b) - lm + + c = m + else: + # ratio of uniforms without mode shift + # compute np.sqrt(quasipdf(m)) + umax = np.exp(0.5*self._logquasipdf(m, p, b)) + xplus = ((1 + p) + np.sqrt((1 + p)**2 + b**2))/b + vmin = 0 + # compute xplus * np.sqrt(quasipdf(xplus)) + vmax = xplus * np.exp(0.5 * self._logquasipdf(xplus, p, b)) + c = 0 + + def logqpdf(x): + return self._logquasipdf(x, p, b) + + if vmin >= vmax: + raise ValueError("vmin must be smaller than vmax.") + if umax <= 0: + raise ValueError("umax must be positive.") + + i = 1 + while simulated < N: + k = N - simulated + # simulate uniform rvs on [0, umax] and [vmin, vmax] + u = umax * random_state.uniform(size=k) + v = random_state.uniform(size=k) + v = vmin + (vmax - vmin) * v + rvs = v / u + c + # rewrite acceptance condition u**2 <= pdf(rvs) by taking logs + accept = (2*np.log(u) <= logqpdf(rvs)) + num_accept = np.sum(accept) + if num_accept > 0: + x[simulated:(simulated + num_accept)] = rvs[accept] + simulated += num_accept + + if (simulated == 0) and (i*N >= 50000): + msg = ("Not a single random variate could be generated " + f"in {i*N} attempts. Sampling does not appear to " + "work for the provided parameters.") + raise RuntimeError(msg) + i += 1 + else: + # use new algorithm in [2] + x0 = b / (1 - p) + xs = np.max((x0, 2 / b)) + k1 = np.exp(self._logquasipdf(m, p, b)) + A1 = k1 * x0 + if x0 < 2 / b: + k2 = np.exp(-b) + if p > 0: + A2 = k2 * ((2 / b)**p - x0**p) / p + else: + A2 = k2 * np.log(2 / b**2) + else: + k2, A2 = 0, 0 + k3 = xs**(p - 1) + A3 = 2 * k3 * np.exp(-xs * b / 2) / b + A = A1 + A2 + A3 + + # [2]: rejection constant is < 2.73; so expected runtime is finite + while simulated < N: + k = N - simulated + h, rvs = np.zeros(k), np.zeros(k) + # simulate uniform rvs on [x1, x2] and [0, y2] + u = random_state.uniform(size=k) + v = A * random_state.uniform(size=k) + cond1 = v <= A1 + cond2 = np.logical_not(cond1) & (v <= A1 + A2) + cond3 = np.logical_not(cond1 | cond2) + # subdomain (0, x0) + rvs[cond1] = x0 * v[cond1] / A1 + h[cond1] = k1 + # subdomain (x0, 2 / b) + if p > 0: + rvs[cond2] = (x0**p + (v[cond2] - A1) * p / k2)**(1 / p) + else: + rvs[cond2] = b * np.exp((v[cond2] - A1) * np.exp(b)) + h[cond2] = k2 * rvs[cond2]**(p - 1) + # subdomain (xs, infinity) + z = np.exp(-xs * b / 2) - b * (v[cond3] - A1 - A2) / (2 * k3) + rvs[cond3] = -2 / b * np.log(z) + h[cond3] = k3 * np.exp(-rvs[cond3] * b / 2) + # apply rejection method + accept = (np.log(u * h) <= self._logquasipdf(rvs, p, b)) + num_accept = sum(accept) + if num_accept > 0: + x[simulated:(simulated + num_accept)] = rvs[accept] + simulated += num_accept + + rvs = np.reshape(x, size1d) + if invert_res: + rvs = 1 / rvs + return rvs + + def _mode(self, p, b): + # distinguish cases to avoid catastrophic cancellation (see [2]) + if p < 1: + return b / (np.sqrt((p - 1)**2 + b**2) + 1 - p) + else: + return (np.sqrt((1 - p)**2 + b**2) - (1 - p)) / b + + def _munp(self, n, p, b): + num = sc.kve(p + n, b) + denom = sc.kve(p, b) + inf_vals = np.isinf(num) | np.isinf(denom) + if inf_vals.any(): + msg = ("Infinite values encountered in the moment calculation " + "involving scipy.special.kve. Values replaced by NaN to " + "avoid incorrect results.") + warnings.warn(msg, RuntimeWarning, stacklevel=3) + m = np.full_like(num, np.nan, dtype=np.float64) + m[~inf_vals] = num[~inf_vals] / denom[~inf_vals] + else: + m = num / denom + return m + + +geninvgauss = geninvgauss_gen(a=0.0, name="geninvgauss") + + +@np.vectorize(otypes=[np.float64]) +def _norminvgauss_quadrature(x, y, a, b): + # gh-23196 reported that the norminvgauss CDF would drop to zero in the far right + # tail. The SF had a similar problem, dropping to zero at the far left. + # This fixes the bug by using 1 - SF to compute the CDF in the right tail and + # 1 - CDF to compute the SF in the left tail. The mean is guaranteed to be beyond + # the median, so there is no loss of precision due to subtractive cancellation. + mean = b / np.sqrt((a + b) * (a - b)) + if np.isneginf(x) and y > abs(mean): + return 1 - integrate.quad(norminvgauss._pdf, y, np.inf, args=(a, b))[0] + if np.isposinf(y) and x < -abs(mean): + return 1 - integrate.quad(norminvgauss._pdf, -np.inf, x, args=(a, b))[0] + else: + res = integrate.quad(norminvgauss._pdf, x, y, args=(a, b))[0] + return np.clip(res, 0, 1) + + +class norminvgauss_gen(rv_continuous): + r"""A Normal Inverse Gaussian continuous random variable. + + %(before_notes)s + + Notes + ----- + The probability density function for `norminvgauss` is: + + .. math:: + + f(x, a, b) = \frac{a \, K_1(a \sqrt{1 + x^2})}{\pi \sqrt{1 + x^2}} \, + \exp(\sqrt{a^2 - b^2} + b x) + + where :math:`x` is a real number, the parameter :math:`a` is the tail + heaviness and :math:`b` is the asymmetry parameter satisfying + :math:`a > 0` and :math:`|b| <= a`. + :math:`K_1` is the modified Bessel function of second kind + (`scipy.special.k1`). + + %(after_notes)s + + A normal inverse Gaussian random variable `Y` with parameters `a` and `b` + can be expressed as a normal mean-variance mixture: + ``Y = b * V + sqrt(V) * X`` where `X` is ``norm(0,1)`` and `V` is + ``invgauss(mu=1/sqrt(a**2 - b**2))``. This representation is used + to generate random variates. + + Another common parametrization of the distribution (see Equation 2.1 in + [2]_) is given by the following expression of the pdf: + + .. math:: + + g(x, \alpha, \beta, \delta, \mu) = + \frac{\alpha\delta K_1\left(\alpha\sqrt{\delta^2 + (x - \mu)^2}\right)} + {\pi \sqrt{\delta^2 + (x - \mu)^2}} \, + e^{\delta \sqrt{\alpha^2 - \beta^2} + \beta (x - \mu)} + + In SciPy, this corresponds to + :math:`a=\alpha \delta, b=\beta \delta, \text{loc}=\mu, \text{scale}=\delta`. + + References + ---------- + .. [1] O. Barndorff-Nielsen, "Hyperbolic Distributions and Distributions on + Hyperbolae", Scandinavian Journal of Statistics, Vol. 5(3), + pp. 151-157, 1978. + + .. [2] O. Barndorff-Nielsen, "Normal Inverse Gaussian Distributions and + Stochastic Volatility Modelling", Scandinavian Journal of + Statistics, Vol. 24, pp. 1-13, 1997. + + %(example)s + + """ + _support_mask = rv_continuous._open_support_mask + + def _argcheck(self, a, b): + return (a > 0) & (np.absolute(b) < a) + + def _shape_info(self): + ia = _ShapeInfo("a", False, (0, np.inf), (False, False)) + ib = _ShapeInfo("b", False, (-np.inf, np.inf), (False, False)) + return [ia, ib] + + def _fitstart(self, data): + # Arbitrary, but the default a = b = 1 is not valid; the distribution + # requires |b| < a. + return super()._fitstart(data, args=(1, 0.5)) + + def _pdf(self, x, a, b): + gamma = np.sqrt((a + b) * (a - b)) + fac1 = a / np.pi + sq = np.hypot(1, x) # reduce overflows + return fac1 * sc.k1e(a * sq) * np.exp(b*x - a*sq + gamma) / sq + + def _cdf(self, x, a, b): + return _norminvgauss_quadrature(-np.inf, x, a, b) + + def _sf(self, x, a, b): + return _norminvgauss_quadrature(x, np.inf, a, b) + + def _isf(self, q, a, b): + def _isf_scalar(q, a, b): + + def eq(x, a, b, q): + # Solve eq(x, a, b, q) = 0 to obtain isf(x, a, b) = q. + return self._sf(x, a, b) - q + + # Find a bracketing interval for the root. + # Start at the mean, and grow the length of the interval + # by 2 each iteration until there is a sign change in eq. + xm = self.mean(a, b) + em = eq(xm, a, b, q) + if em == 0: + # Unlikely, but might as well check. + return xm + if em > 0: + delta = 1 + left = xm + right = xm + delta + while eq(right, a, b, q) > 0: + delta = 2*delta + right = xm + delta + else: + # em < 0 + delta = 1 + right = xm + left = xm - delta + while eq(left, a, b, q) < 0: + delta = 2*delta + left = xm - delta + result = optimize.brentq(eq, left, right, args=(a, b, q), + xtol=self.xtol) + return result + + if np.isscalar(q): + return _isf_scalar(q, a, b) + else: + result = [] + for (q0, a0, b0) in zip(q, a, b): + result.append(_isf_scalar(q0, a0, b0)) + return np.array(result) + + def _rvs(self, a, b, size=None, random_state=None): + # note: X = b * V + sqrt(V) * X is norminvgaus(a,b) if X is standard + # normal and V is invgauss(mu=1/sqrt(a**2 - b**2)) + gamma = np.sqrt((a + b) * (a - b)) + ig = invgauss.rvs(mu=1/gamma, size=size, random_state=random_state) + return b * ig + np.sqrt(ig) * norm.rvs(size=size, + random_state=random_state) + + def _stats(self, a, b): + gamma = np.sqrt((a + b) * (a - b)) + mean = b / gamma + variance = a**2 / gamma**3 + skewness = 3.0 * b / (a * np.sqrt(gamma)) + kurtosis = 3.0 * (1 + 4 * b**2 / a**2) / gamma + return mean, variance, skewness, kurtosis + + +norminvgauss = norminvgauss_gen(name="norminvgauss") + + +class invweibull_gen(rv_continuous): + """An inverted Weibull continuous random variable. + + This distribution is also known as the Fréchet distribution or the + type II extreme value distribution. + + %(before_notes)s + + Notes + ----- + The probability density function for `invweibull` is: + + .. math:: + + f(x, c) = c x^{-c-1} \\exp(-x^{-c}) + + for :math:`x > 0`, :math:`c > 0`. + + `invweibull` takes ``c`` as a shape parameter for :math:`c`. + + %(after_notes)s + + References + ---------- + F.R.S. de Gusmao, E.M.M Ortega and G.M. Cordeiro, "The generalized inverse + Weibull distribution", Stat. Papers, vol. 52, pp. 591-619, 2011. + + %(example)s + + """ + _support_mask = rv_continuous._open_support_mask + + def _shape_info(self): + return [_ShapeInfo("c", False, (0, np.inf), (False, False))] + + def _pdf(self, x, c): + # invweibull.pdf(x, c) = c * x**(-c-1) * exp(-x**(-c)) + xc1 = np.power(x, -c - 1.0) + xc2 = np.power(x, -c) + xc2 = np.exp(-xc2) + return c * xc1 * xc2 + + def _cdf(self, x, c): + xc1 = np.power(x, -c) + return np.exp(-xc1) + + def _sf(self, x, c): + return -np.expm1(-x**-c) + + def _ppf(self, q, c): + return np.power(-np.log(q), -1.0/c) + + def _isf(self, p, c): + return (-np.log1p(-p))**(-1/c) + + def _munp(self, n, c): + return sc.gamma(1 - n / c) + + def _entropy(self, c): + return 1+_EULER + _EULER / c - np.log(c) + + def _fitstart(self, data, args=None): + # invweibull requires c > 1 for the first moment to exist, so use 2.0 + args = (2.0,) if args is None else args + return super()._fitstart(data, args=args) + + +invweibull = invweibull_gen(a=0, name='invweibull') + + +class jf_skew_t_gen(rv_continuous): + r"""Jones and Faddy skew-t distribution. + + %(before_notes)s + + Notes + ----- + The probability density function for `jf_skew_t` is: + + .. math:: + + f(x; a, b) = C_{a,b}^{-1} + \left(1+\frac{x}{\left(a+b+x^2\right)^{1/2}}\right)^{a+1/2} + \left(1-\frac{x}{\left(a+b+x^2\right)^{1/2}}\right)^{b+1/2} + + for real numbers :math:`a>0` and :math:`b>0`, where + :math:`C_{a,b} = 2^{a+b-1}B(a,b)(a+b)^{1/2}`, and :math:`B` denotes the + beta function (`scipy.special.beta`). + + When :math:`ab`, the distribution is positively skewed. If :math:`a=b`, then + we recover the `t` distribution with :math:`2a` degrees of freedom. + + `jf_skew_t` takes :math:`a` and :math:`b` as shape parameters. + + %(after_notes)s + + References + ---------- + .. [1] M.C. Jones and M.J. Faddy. "A skew extension of the t distribution, + with applications" *Journal of the Royal Statistical Society*. + Series B (Statistical Methodology) 65, no. 1 (2003): 159-174. + :doi:`10.1111/1467-9868.00378` + + %(example)s + + """ + def _shape_info(self): + ia = _ShapeInfo("a", False, (0, np.inf), (False, False)) + ib = _ShapeInfo("b", False, (0, np.inf), (False, False)) + return [ia, ib] + + def _pdf(self, x, a, b): + c = 2 ** (a + b - 1) * sc.beta(a, b) * np.sqrt(a + b) + d1 = (1 + x / np.sqrt(a + b + x ** 2)) ** (a + 0.5) + d2 = (1 - x / np.sqrt(a + b + x ** 2)) ** (b + 0.5) + return d1 * d2 / c + + def _rvs(self, a, b, size=None, random_state=None): + d1 = random_state.beta(a, b, size) + d2 = (2 * d1 - 1) * np.sqrt(a + b) + d3 = 2 * np.sqrt(d1 * (1 - d1)) + return d2 / d3 + + def _cdf(self, x, a, b): + y = (1 + x / np.sqrt(a + b + x ** 2)) * 0.5 + return sc.betainc(a, b, y) + + def _sf(self, x, a, b): + y = (1 + x / np.sqrt(a + b + x ** 2)) * 0.5 + return sc.betaincc(a, b, y) + + def _ppf(self, q, a, b): + d1 = beta.ppf(q, a, b) + d2 = (2 * d1 - 1) * np.sqrt(a + b) + d3 = 2 * np.sqrt(d1 * (1 - d1)) + return d2 / d3 + + def _munp(self, n, a, b): + """Returns the n-th moment(s) where all the following hold: + + - n >= 0 + - a > n / 2 + - b > n / 2 + + The result is np.nan in all other cases. + """ + def nth_moment(n_k, a_k, b_k): + """Computes E[T^(n_k)] where T is skew-t distributed with + parameters a_k and b_k. + """ + num = (a_k + b_k) ** (0.5 * n_k) + denom = 2 ** n_k * sc.beta(a_k, b_k) + + indices = np.arange(n_k + 1) + sgn = np.where(indices % 2 > 0, -1, 1) + d = sc.beta(a_k + 0.5 * n_k - indices, b_k - 0.5 * n_k + indices) + sum_terms = sc.comb(n_k, indices) * sgn * d + + return num / denom * sum_terms.sum() + + nth_moment_valid = (a > 0.5 * n) & (b > 0.5 * n) & (n >= 0) + return xpx.apply_where( + nth_moment_valid, + (n, a, b), + np.vectorize(nth_moment, otypes=[np.float64]), + fill_value=np.nan, + ) + + +jf_skew_t = jf_skew_t_gen(name='jf_skew_t') + + +class johnsonsb_gen(rv_continuous): + r"""A Johnson SB continuous random variable. + + %(before_notes)s + + See Also + -------- + johnsonsu + + Notes + ----- + The probability density function for `johnsonsb` is: + + .. math:: + + f(x, a, b) = \frac{b}{x(1-x)} \phi(a + b \log \frac{x}{1-x} ) + + where :math:`x`, :math:`a`, and :math:`b` are real scalars; :math:`b > 0` + and :math:`x \in [0,1]`. :math:`\phi` is the pdf of the normal + distribution. + + `johnsonsb` takes :math:`a` and :math:`b` as shape parameters. + + %(after_notes)s + + %(example)s + + """ + _support_mask = rv_continuous._open_support_mask + + def _argcheck(self, a, b): + return (b > 0) & (a == a) + + def _shape_info(self): + ia = _ShapeInfo("a", False, (-np.inf, np.inf), (False, False)) + ib = _ShapeInfo("b", False, (0, np.inf), (False, False)) + return [ia, ib] + + def _pdf(self, x, a, b): + # johnsonsb.pdf(x, a, b) = b / (x*(1-x)) * phi(a + b * log(x/(1-x))) + trm = _norm_pdf(a + b*sc.logit(x)) + return b*1.0/(x*(1-x))*trm + + def _cdf(self, x, a, b): + return _norm_cdf(a + b*sc.logit(x)) + + def _ppf(self, q, a, b): + return sc.expit(1.0 / b * (_norm_ppf(q) - a)) + + def _sf(self, x, a, b): + return _norm_sf(a + b*sc.logit(x)) + + def _isf(self, q, a, b): + return sc.expit(1.0 / b * (_norm_isf(q) - a)) + + +johnsonsb = johnsonsb_gen(a=0.0, b=1.0, name='johnsonsb') + + +class johnsonsu_gen(rv_continuous): + r"""A Johnson SU continuous random variable. + + %(before_notes)s + + See Also + -------- + johnsonsb + + Notes + ----- + The probability density function for `johnsonsu` is: + + .. math:: + + f(x, a, b) = \frac{b}{\sqrt{x^2 + 1}} + \phi(a + b \log(x + \sqrt{x^2 + 1})) + + where :math:`x`, :math:`a`, and :math:`b` are real scalars; :math:`b > 0`. + :math:`\phi` is the pdf of the normal distribution. + + `johnsonsu` takes :math:`a` and :math:`b` as shape parameters. + + The first four central moments are calculated according to the formulas + in [1]_. + + %(after_notes)s + + References + ---------- + .. [1] Taylor Enterprises. "Johnson Family of Distributions". + https://variation.com/wp-content/distribution_analyzer_help/hs126.htm + + %(example)s + + """ + def _argcheck(self, a, b): + return (b > 0) & (a == a) + + def _shape_info(self): + ia = _ShapeInfo("a", False, (-np.inf, np.inf), (False, False)) + ib = _ShapeInfo("b", False, (0, np.inf), (False, False)) + return [ia, ib] + + def _pdf(self, x, a, b): + # johnsonsu.pdf(x, a, b) = b / sqrt(x**2 + 1) * + # phi(a + b * log(x + sqrt(x**2 + 1))) + x2 = x*x + trm = _norm_pdf(a + b * np.arcsinh(x)) + return b*1.0/np.sqrt(x2+1.0)*trm + + def _cdf(self, x, a, b): + return _norm_cdf(a + b * np.arcsinh(x)) + + def _ppf(self, q, a, b): + return np.sinh((_norm_ppf(q) - a) / b) + + def _sf(self, x, a, b): + return _norm_sf(a + b * np.arcsinh(x)) + + def _isf(self, x, a, b): + return np.sinh((_norm_isf(x) - a) / b) + + def _stats(self, a, b, moments='mv'): + # Naive implementation of first and second moment to address gh-18071. + # https://variation.com/wp-content/distribution_analyzer_help/hs126.htm + # Numerical improvements left to future enhancements. + mu, mu2, g1, g2 = None, None, None, None + + bn2 = b**-2. + expbn2 = np.exp(bn2) + a_b = a / b + + if 'm' in moments: + mu = -expbn2**0.5 * np.sinh(a_b) + if 'v' in moments: + mu2 = 0.5*sc.expm1(bn2)*(expbn2*np.cosh(2*a_b) + 1) + if 's' in moments: + t1 = expbn2**.5 * sc.expm1(bn2)**0.5 + t2 = 3*np.sinh(a_b) + t3 = expbn2 * (expbn2 + 2) * np.sinh(3*a_b) + denom = np.sqrt(2) * (1 + expbn2 * np.cosh(2*a_b))**(3/2) + g1 = -t1 * (t2 + t3) / denom + if 'k' in moments: + t1 = 3 + 6*expbn2 + t2 = 4*expbn2**2 * (expbn2 + 2) * np.cosh(2*a_b) + t3 = expbn2**2 * np.cosh(4*a_b) + t4 = -3 + 3*expbn2**2 + 2*expbn2**3 + expbn2**4 + denom = 2*(1 + expbn2*np.cosh(2*a_b))**2 + g2 = (t1 + t2 + t3*t4) / denom - 3 + return mu, mu2, g1, g2 + + +johnsonsu = johnsonsu_gen(name='johnsonsu') + + +class landau_gen(rv_continuous): + r"""A Landau continuous random variable. + + %(before_notes)s + + See Also + -------- + :ref:`landau_energy_loss` : Extended example in a particle physics context. + + Notes + ----- + The probability density function for `landau` ([1]_, [2]_) is: + + .. math:: + + f(x) = \frac{1}{\pi}\int_0^\infty \exp(-t \log t - xt)\sin(\pi t) dt + + for a real number :math:`x`. + + %(after_notes)s + + Often (e.g. [2]_), the Landau distribution is parameterized in terms of a + location parameter :math:`\mu` and scale parameter :math:`c`, the latter of + which *also* introduces a location shift. If ``mu`` and ``c`` are used to + represent these parameters, this corresponds with SciPy's parameterization + with ``loc = mu + 2*c / np.pi * np.log(c)`` and ``scale = c``. + + This distribution uses routines from the Boost Math C++ library for + the computation of the ``pdf``, ``cdf``, ``ppf``, ``sf`` and ``isf`` + methods. [1]_ + + References + ---------- + .. [1] Landau, L. (1944). "On the energy loss of fast particles by + ionization". J. Phys. (USSR). 8: 201. + .. [2] "Landau Distribution", Wikipedia, + https://en.wikipedia.org/wiki/Landau_distribution + .. [3] Chambers, J. M., Mallows, C. L., & Stuck, B. (1976). + "A method for simulating stable random variables." + Journal of the American Statistical Association, 71(354), 340-344. + .. [4] The Boost Developers. "Boost C++ Libraries". https://www.boost.org/. + .. [5] Yoshimura, T. "Numerical Evaluation and High Precision Approximation + Formula for Landau Distribution". + :doi:`10.36227/techrxiv.171822215.53612870/v2` + + %(example)s + + """ + def _shape_info(self): + return [] + + def _entropy(self): + # Computed with mpmath - see gh-19145 + return 2.37263644000448182 + + def _pdf(self, x): + return scu._landau_pdf(x, 0, 1) + + def _cdf(self, x): + return scu._landau_cdf(x, 0, 1) + + def _sf(self, x): + return scu._landau_sf(x, 0, 1) + + def _ppf(self, p): + return scu._landau_ppf(p, 0, 1) + + def _isf(self, p): + return scu._landau_isf(p, 0, 1) + + def _stats(self): + return np.nan, np.nan, np.nan, np.nan + + def _munp(self, n): + return np.nan if n > 0 else 1 + + def _fitstart(self, data, args=None): + # Initialize ML guesses using quartiles instead of moments. + if isinstance(data, CensoredData): + data = data._uncensor() + p25, p50, p75 = np.percentile(data, [25, 50, 75]) + return p50, (p75 - p25)/2 + + def _rvs(self, size=None, random_state=None): + # Method from https://www.jstor.org/stable/2285309 Eq. 2.4 + pi_2 = np.pi / 2 + U = random_state.uniform(-np.pi / 2, np.pi / 2, size=size) + W = random_state.standard_exponential(size=size) + S = 2 / np.pi * ((pi_2 + U) * np.tan(U) + - np.log((pi_2 * W * np.cos(U)) / (pi_2 + U))) + return S + + +landau = landau_gen(name='landau') + + +class laplace_gen(rv_continuous): + r"""A Laplace continuous random variable. + + %(before_notes)s + + Notes + ----- + The probability density function for `laplace` is + + .. math:: + + f(x) = \frac{1}{2} \exp(-|x|) + + for a real number :math:`x`. + + %(after_notes)s + + %(example)s + + """ + def _shape_info(self): + return [] + + def _rvs(self, size=None, random_state=None): + return random_state.laplace(0, 1, size=size) + + def _pdf(self, x): + # laplace.pdf(x) = 1/2 * exp(-abs(x)) + return 0.5*np.exp(-abs(x)) + + def _cdf(self, x): + with np.errstate(over='ignore'): + return np.where(x > 0, 1.0 - 0.5*np.exp(-x), 0.5*np.exp(x)) + + def _sf(self, x): + # By symmetry... + return self._cdf(-x) + + def _ppf(self, q): + return np.where(q > 0.5, -np.log(2*(1-q)), np.log(2*q)) + + def _isf(self, q): + # By symmetry... + return -self._ppf(q) + + def _stats(self): + return 0, 2, 0, 3 + + def _entropy(self): + return np.log(2)+1 + + @_call_super_mom + @replace_notes_in_docstring(rv_continuous, notes="""\ + This function uses explicit formulas for the maximum likelihood + estimation of the Laplace distribution parameters, so the keyword + arguments `loc`, `scale`, and `optimizer` are ignored.\n\n""") + def fit(self, data, *args, **kwds): + data, floc, fscale = _check_fit_input_parameters(self, data, + args, kwds) + + # Source: Statistical Distributions, 3rd Edition. Evans, Hastings, + # and Peacock (2000), Page 124 + + if floc is None: + floc = np.median(data) + + if fscale is None: + fscale = (np.sum(np.abs(data - floc))) / len(data) + + return floc, fscale + + +laplace = laplace_gen(name='laplace') + + +class laplace_asymmetric_gen(rv_continuous): + r"""An asymmetric Laplace continuous random variable. + + %(before_notes)s + + See Also + -------- + laplace : Laplace distribution + + Notes + ----- + The probability density function for `laplace_asymmetric` is + + .. math:: + + f(x, \kappa) &= \frac{1}{\kappa+\kappa^{-1}}\exp(-x\kappa),\quad x\ge0\\ + &= \frac{1}{\kappa+\kappa^{-1}}\exp(x/\kappa),\quad x<0\\ + + for :math:`-\infty < x < \infty`, :math:`\kappa > 0`. + + `laplace_asymmetric` takes ``kappa`` as a shape parameter for + :math:`\kappa`. For :math:`\kappa = 1`, it is identical to a + Laplace distribution. + + %(after_notes)s + + Note that the scale parameter of some references is the reciprocal of + SciPy's ``scale``. For example, :math:`\lambda = 1/2` in the + parameterization of [1]_ is equivalent to ``scale = 2`` with + `laplace_asymmetric`. + + References + ---------- + .. [1] "Asymmetric Laplace distribution", Wikipedia + https://en.wikipedia.org/wiki/Asymmetric_Laplace_distribution + + .. [2] Kozubowski TJ and Podgórski K. A Multivariate and + Asymmetric Generalization of Laplace Distribution, + Computational Statistics 15, 531--540 (2000). + :doi:`10.1007/PL00022717` + + %(example)s + + """ + def _shape_info(self): + return [_ShapeInfo("kappa", False, (0, np.inf), (False, False))] + + def _pdf(self, x, kappa): + return np.exp(self._logpdf(x, kappa)) + + def _logpdf(self, x, kappa): + kapinv = 1/kappa + lPx = x * np.where(x >= 0, -kappa, kapinv) + lPx -= np.log(kappa+kapinv) + return lPx + + def _cdf(self, x, kappa): + kapinv = 1/kappa + kappkapinv = kappa+kapinv + return np.where(x >= 0, + 1 - np.exp(-x*kappa)*(kapinv/kappkapinv), + np.exp(x*kapinv)*(kappa/kappkapinv)) + + def _sf(self, x, kappa): + kapinv = 1/kappa + kappkapinv = kappa+kapinv + return np.where(x >= 0, + np.exp(-x*kappa)*(kapinv/kappkapinv), + 1 - np.exp(x*kapinv)*(kappa/kappkapinv)) + + def _ppf(self, q, kappa): + kapinv = 1/kappa + kappkapinv = kappa+kapinv + return np.where(q >= kappa/kappkapinv, + -np.log((1 - q)*kappkapinv*kappa)*kapinv, + np.log(q*kappkapinv/kappa)*kappa) + + def _isf(self, q, kappa): + kapinv = 1/kappa + kappkapinv = kappa+kapinv + return np.where(q <= kapinv/kappkapinv, + -np.log(q*kappkapinv*kappa)*kapinv, + np.log((1 - q)*kappkapinv/kappa)*kappa) + + def _stats(self, kappa): + kapinv = 1/kappa + mn = kapinv - kappa + var = kapinv*kapinv + kappa*kappa + g1 = 2.0*(1-np.power(kappa, 6))/np.power(1+np.power(kappa, 4), 1.5) + g2 = 6.0*(1+np.power(kappa, 8))/np.power(1+np.power(kappa, 4), 2) + return mn, var, g1, g2 + + def _entropy(self, kappa): + return 1 + np.log(kappa+1/kappa) + + +laplace_asymmetric = laplace_asymmetric_gen(name='laplace_asymmetric') + + +def _check_fit_input_parameters(dist, data, args, kwds): + if not isinstance(data, CensoredData): + data = np.asarray(data) + + floc = kwds.get('floc', None) + fscale = kwds.get('fscale', None) + + num_shapes = len(dist.shapes.split(",")) if dist.shapes else 0 + fshape_keys = [] + fshapes = [] + + # user has many options for fixing the shape, so here we standardize it + # into 'f' + the number of the shape. + # Adapted from `_reduce_func` in `_distn_infrastructure.py`: + if dist.shapes: + shapes = dist.shapes.replace(',', ' ').split() + for j, s in enumerate(shapes): + key = 'f' + str(j) + names = [key, 'f' + s, 'fix_' + s] + val = _get_fixed_fit_value(kwds, names) + fshape_keys.append(key) + fshapes.append(val) + if val is not None: + kwds[key] = val + + # determine if there are any unknown arguments in kwds + known_keys = {'loc', 'scale', 'optimizer', 'method', + 'floc', 'fscale', *fshape_keys} + unknown_keys = set(kwds).difference(known_keys) + if unknown_keys: + raise TypeError(f"Unknown keyword arguments: {unknown_keys}.") + + if len(args) > num_shapes: + raise TypeError("Too many positional arguments.") + + if None not in {floc, fscale, *fshapes}: + # This check is for consistency with `rv_continuous.fit`. + # Without this check, this function would just return the + # parameters that were given. + raise RuntimeError("All parameters fixed. There is nothing to " + "optimize.") + + uncensored = data._uncensor() if isinstance(data, CensoredData) else data + if not np.isfinite(uncensored).all(): + raise ValueError("The data contains non-finite values.") + + return (data, *fshapes, floc, fscale) + + +class levy_gen(rv_continuous): + r"""A Levy continuous random variable. + + %(before_notes)s + + See Also + -------- + levy_stable, levy_l + + Notes + ----- + The probability density function for `levy` is: + + .. math:: + + f(x) = \frac{1}{\sqrt{2\pi x^3}} \exp\left(-\frac{1}{2x}\right) + + for :math:`x > 0`. + + This is the same as the Levy-stable distribution with :math:`a=1/2` and + :math:`b=1`. + + %(after_notes)s + + Examples + -------- + >>> import numpy as np + >>> from scipy.stats import levy + >>> import matplotlib.pyplot as plt + >>> fig, ax = plt.subplots(1, 1) + + Calculate the first four moments: + + >>> mean, var, skew, kurt = levy.stats(moments='mvsk') + + Display the probability density function (``pdf``): + + >>> # `levy` is very heavy-tailed. + >>> # To show a nice plot, let's cut off the upper 40 percent. + >>> a, b = levy.ppf(0), levy.ppf(0.6) + >>> x = np.linspace(a, b, 100) + >>> ax.plot(x, levy.pdf(x), + ... 'r-', lw=5, alpha=0.6, label='levy pdf') + + Alternatively, the distribution object can be called (as a function) + to fix the shape, location and scale parameters. This returns a "frozen" + RV object holding the given parameters fixed. + + Freeze the distribution and display the frozen ``pdf``: + + >>> rv = levy() + >>> ax.plot(x, rv.pdf(x), 'k-', lw=2, label='frozen pdf') + + Check accuracy of ``cdf`` and ``ppf``: + + >>> vals = levy.ppf([0.001, 0.5, 0.999]) + >>> np.allclose([0.001, 0.5, 0.999], levy.cdf(vals)) + True + + Generate random numbers: + + >>> r = levy.rvs(size=1000) + + And compare the histogram: + + >>> # manual binning to ignore the tail + >>> bins = np.concatenate((np.linspace(a, b, 20), [np.max(r)])) + >>> ax.hist(r, bins=bins, density=True, histtype='stepfilled', alpha=0.2) + >>> ax.set_xlim([x[0], x[-1]]) + >>> ax.legend(loc='best', frameon=False) + >>> plt.show() + + """ + _support_mask = rv_continuous._open_support_mask + + def _shape_info(self): + return [] + + def _pdf(self, x): + # levy.pdf(x) = 1 / (x * sqrt(2*pi*x)) * exp(-1/(2*x)) + return 1 / np.sqrt(2*np.pi*x) / x * np.exp(-1/(2*x)) + + def _cdf(self, x): + # Equivalent to 2*norm.sf(np.sqrt(1/x)) + return sc.erfc(np.sqrt(0.5 / x)) + + def _sf(self, x): + return sc.erf(np.sqrt(0.5 / x)) + + def _ppf(self, q): + # Equivalent to 1.0/(norm.isf(q/2)**2) or 0.5/(erfcinv(q)**2) + val = _norm_isf(q/2) + return 1.0 / (val * val) + + def _isf(self, p): + return 1/(2*sc.erfinv(p)**2) + + def _stats(self): + return np.inf, np.inf, np.nan, np.nan + + +levy = levy_gen(a=0.0, name="levy") + + +class levy_l_gen(rv_continuous): + r"""A left-skewed Levy continuous random variable. + + %(before_notes)s + + See Also + -------- + levy, levy_stable + + Notes + ----- + The probability density function for `levy_l` is: + + .. math:: + f(x) = \frac{1}{|x| \sqrt{2\pi |x|}} \exp{ \left(-\frac{1}{2|x|} \right)} + + for :math:`x < 0`. + + This is the same as the Levy-stable distribution with :math:`a=1/2` and + :math:`b=-1`. + + %(after_notes)s + + Examples + -------- + >>> import numpy as np + >>> from scipy.stats import levy_l + >>> import matplotlib.pyplot as plt + >>> fig, ax = plt.subplots(1, 1) + + Calculate the first four moments: + + >>> mean, var, skew, kurt = levy_l.stats(moments='mvsk') + + Display the probability density function (``pdf``): + + >>> # `levy_l` is very heavy-tailed. + >>> # To show a nice plot, let's cut off the lower 40 percent. + >>> a, b = levy_l.ppf(0.4), levy_l.ppf(1) + >>> x = np.linspace(a, b, 100) + >>> ax.plot(x, levy_l.pdf(x), + ... 'r-', lw=5, alpha=0.6, label='levy_l pdf') + + Alternatively, the distribution object can be called (as a function) + to fix the shape, location and scale parameters. This returns a "frozen" + RV object holding the given parameters fixed. + + Freeze the distribution and display the frozen ``pdf``: + + >>> rv = levy_l() + >>> ax.plot(x, rv.pdf(x), 'k-', lw=2, label='frozen pdf') + + Check accuracy of ``cdf`` and ``ppf``: + + >>> vals = levy_l.ppf([0.001, 0.5, 0.999]) + >>> np.allclose([0.001, 0.5, 0.999], levy_l.cdf(vals)) + True + + Generate random numbers: + + >>> r = levy_l.rvs(size=1000) + + And compare the histogram: + + >>> # manual binning to ignore the tail + >>> bins = np.concatenate(([np.min(r)], np.linspace(a, b, 20))) + >>> ax.hist(r, bins=bins, density=True, histtype='stepfilled', alpha=0.2) + >>> ax.set_xlim([x[0], x[-1]]) + >>> ax.legend(loc='best', frameon=False) + >>> plt.show() + + """ + _support_mask = rv_continuous._open_support_mask + + def _shape_info(self): + return [] + + def _pdf(self, x): + # levy_l.pdf(x) = 1 / (abs(x) * sqrt(2*pi*abs(x))) * exp(-1/(2*abs(x))) + ax = abs(x) + return 1/np.sqrt(2*np.pi*ax)/ax*np.exp(-1/(2*ax)) + + def _cdf(self, x): + ax = abs(x) + return 2 * _norm_cdf(1 / np.sqrt(ax)) - 1 + + def _sf(self, x): + ax = abs(x) + return 2 * _norm_sf(1 / np.sqrt(ax)) + + def _ppf(self, q): + val = _norm_ppf((q + 1.0) / 2) + return -1.0 / (val * val) + + def _isf(self, p): + return -1/_norm_isf(p/2)**2 + + def _stats(self): + return np.inf, np.inf, np.nan, np.nan + + +levy_l = levy_l_gen(b=0.0, name="levy_l") + + +class logistic_gen(rv_continuous): + r"""A logistic (or Sech-squared) continuous random variable. + + %(before_notes)s + + Notes + ----- + The probability density function for `logistic` is: + + .. math:: + + f(x) = \frac{\exp(-x)} + {(1+\exp(-x))^2} + + `logistic` is a special case of `genlogistic` with ``c=1``. + + Remark that the survival function (``logistic.sf``) is equal to the + Fermi-Dirac distribution describing fermionic statistics. + + %(after_notes)s + + %(example)s + + """ + def _shape_info(self): + return [] + + def _rvs(self, size=None, random_state=None): + return random_state.logistic(size=size) + + def _pdf(self, x): + # logistic.pdf(x) = exp(-x) / (1+exp(-x))**2 + return np.exp(self._logpdf(x)) + + def _logpdf(self, x): + y = -np.abs(x) + return y - 2. * sc.log1p(np.exp(y)) + + def _cdf(self, x): + return sc.expit(x) + + def _logcdf(self, x): + return sc.log_expit(x) + + def _ppf(self, q): + return sc.logit(q) + + def _sf(self, x): + return sc.expit(-x) + + def _logsf(self, x): + return sc.log_expit(-x) + + def _isf(self, q): + return -sc.logit(q) + + def _stats(self): + return 0, np.pi*np.pi/3.0, 0, 6.0/5.0 + + def _entropy(self): + # https://en.wikipedia.org/wiki/Logistic_distribution + return 2.0 + + @_call_super_mom + @inherit_docstring_from(rv_continuous) + def fit(self, data, *args, **kwds): + if kwds.pop('superfit', False): + return super().fit(data, *args, **kwds) + + data, floc, fscale = _check_fit_input_parameters(self, data, + args, kwds) + n = len(data) + + # rv_continuous provided guesses + loc, scale = self._fitstart(data) + # these are trumped by user-provided guesses + loc, scale = kwds.get('loc', loc), kwds.get('scale', scale) + + # the maximum likelihood estimators `a` and `b` of the location and + # scale parameters are roots of the two equations described in `func`. + # Source: Statistical Distributions, 3rd Edition. Evans, Hastings, and + # Peacock (2000), Page 130 + + def dl_dloc(loc, scale=fscale): + c = (data - loc) / scale + return np.sum(sc.expit(c)) - n/2 + + def dl_dscale(scale, loc=floc): + c = (data - loc) / scale + return np.sum(c*np.tanh(c/2)) - n + + def func(params): + loc, scale = params + return dl_dloc(loc, scale), dl_dscale(scale, loc) + + if fscale is not None and floc is None: + res = optimize.root(dl_dloc, (loc,)) + loc = res.x[0] + scale = fscale + elif floc is not None and fscale is None: + res = optimize.root(dl_dscale, (scale,)) + scale = res.x[0] + loc = floc + else: + res = optimize.root(func, (loc, scale)) + loc, scale = res.x + + # Note: gh-18176 reported data for which the reported MLE had + # `scale < 0`. To fix the bug, we return abs(scale). This is OK because + # `dl_dscale` and `dl_dloc` are even and odd functions of `scale`, + # respectively, so if `-scale` is a solution, so is `scale`. + scale = abs(scale) + return ((loc, scale) if res.success + else super().fit(data, *args, **kwds)) + + +logistic = logistic_gen(name='logistic') + + +class loggamma_gen(rv_continuous): + r"""A log gamma continuous random variable. + + %(before_notes)s + + Notes + ----- + The probability density function for `loggamma` is: + + .. math:: + + f(x, c) = \frac{\exp(c x - \exp(x))} + {\Gamma(c)} + + for all :math:`x, c > 0`. Here, :math:`\Gamma` is the + gamma function (`scipy.special.gamma`). + + `loggamma` takes ``c`` as a shape parameter for :math:`c`. + + %(after_notes)s + + %(example)s + + """ + + def _shape_info(self): + return [_ShapeInfo("c", False, (0, np.inf), (False, False))] + + def _rvs(self, c, size=None, random_state=None): + # Use the property of the gamma distribution Gamma(c) + # Gamma(c) ~ Gamma(c + 1)*U**(1/c), + # where U is uniform on [0, 1]. (See, e.g., + # G. Marsaglia and W.W. Tsang, "A simple method for generating gamma + # variables", https://doi.org/10.1145/358407.358414) + # So + # log(Gamma(c)) ~ log(Gamma(c + 1)) + log(U)/c + # Generating a sample with this formulation is a bit slower + # than the more obvious log(Gamma(c)), but it avoids loss + # of precision when c << 1. + return (np.log(random_state.gamma(c + 1, size=size)) + + np.log(random_state.uniform(size=size))/c) + + def _pdf(self, x, c): + # loggamma.pdf(x, c) = exp(c*x-exp(x)) / gamma(c) + return np.exp(c*x-np.exp(x)-sc.gammaln(c)) + + def _logpdf(self, x, c): + return c*x - np.exp(x) - sc.gammaln(c) + + def _cdf(self, x, c): + # This function is gammainc(c, exp(x)), where gammainc(c, z) is + # the regularized incomplete gamma function. + # The first term in a series expansion of gamminc(c, z) is + # z**c/Gamma(c+1); see 6.5.29 of Abramowitz & Stegun (and refer + # back to 6.5.1, 6.5.2 and 6.5.4 for the relevant notation). + # This can also be found in the wikipedia article + # https://en.wikipedia.org/wiki/Incomplete_gamma_function. + # Here we use that formula when x is sufficiently negative that + # exp(x) will result in subnormal numbers and lose precision. + # We evaluate the log of the expression first to allow the possible + # cancellation of the terms in the division, and then exponentiate. + # That is, + # exp(x)**c/Gamma(c+1) = exp(log(exp(x)**c/Gamma(c+1))) + # = exp(c*x - gammaln(c+1)) + return xpx.apply_where( + x < _LOGXMIN, (x, c), + lambda x, c: np.exp(c*x - sc.gammaln(c+1)), + lambda x, c: sc.gammainc(c, np.exp(x))) + + def _ppf(self, q, c): + # The expression used when g < _XMIN inverts the one term expansion + # given in the comments of _cdf(). + g = sc.gammaincinv(c, q) + return xpx.apply_where( + g < _XMIN, (g, q, c), + lambda g, q, c: (np.log(q) + sc.gammaln(c+1))/c, + lambda g, q, c: np.log(g)) + + def _sf(self, x, c): + # See the comments for _cdf() for how x < _LOGXMIN is handled. + return xpx.apply_where( + x < _LOGXMIN, (x, c), + lambda x, c: -np.expm1(c*x - sc.gammaln(c+1)), + lambda x, c: sc.gammaincc(c, np.exp(x))) + + def _isf(self, q, c): + # The expression used when g < _XMIN inverts the complement of + # the one term expansion given in the comments of _cdf(). + g = sc.gammainccinv(c, q) + return xpx.apply_where( + g < _XMIN, (g, q, c), + lambda g, q, c: (np.log1p(-q) + sc.gammaln(c+1))/c, + lambda g, q, c: np.log(g)) + + def _stats(self, c): + # See, for example, "A Statistical Study of Log-Gamma Distribution", by + # Ping Shing Chan (thesis, McMaster University, 1993). + mean = sc.digamma(c) + var = sc.polygamma(1, c) + skewness = sc.polygamma(2, c) / np.power(var, 1.5) + excess_kurtosis = sc.polygamma(3, c) / (var*var) + return mean, var, skewness, excess_kurtosis + + def _entropy(self, c): + def regular(c): + h = sc.gammaln(c) - c * sc.digamma(c) + c + return h + + def asymptotic(c): + # using asymptotic expansions for gammaln and psi (see gh-18093) + term = -0.5*np.log(c) + c**-1./6 - c**-3./90 + c**-5./210 + h = norm._entropy() + term + return h + + return xpx.apply_where(c >= 45, c, asymptotic, regular) + + +loggamma = loggamma_gen(name='loggamma') + + +class loglaplace_gen(rv_continuous): + r"""A log-Laplace continuous random variable. + + %(before_notes)s + + Notes + ----- + The probability density function for `loglaplace` is: + + .. math:: + + f(x, c) = \begin{cases}\frac{c}{2} x^{ c-1} &\text{for } 0 < x < 1\\ + \frac{c}{2} x^{-c-1} &\text{for } x \ge 1 + \end{cases} + + for :math:`c > 0`. + + `loglaplace` takes ``c`` as a shape parameter for :math:`c`. + + %(after_notes)s + + Suppose a random variable ``X`` follows the Laplace distribution with + location ``a`` and scale ``b``. Then ``Y = exp(X)`` follows the + log-Laplace distribution with ``c = 1 / b`` and ``scale = exp(a)``. + + References + ---------- + T.J. Kozubowski and K. Podgorski, "A log-Laplace growth rate model", + The Mathematical Scientist, vol. 28, pp. 49-60, 2003. + + %(example)s + + """ + def _shape_info(self): + return [_ShapeInfo("c", False, (0, np.inf), (False, False))] + + def _pdf(self, x, c): + # loglaplace.pdf(x, c) = c / 2 * x**(c-1), for 0 < x < 1 + # = c / 2 * x**(-c-1), for x >= 1 + cd2 = c/2.0 + c = np.where(x < 1, c, -c) + return cd2*x**(c-1) + + def _cdf(self, x, c): + return np.where(x < 1, 0.5*x**c, 1-0.5*x**(-c)) + + def _sf(self, x, c): + return np.where(x < 1, 1 - 0.5*x**c, 0.5*x**(-c)) + + def _ppf(self, q, c): + return np.where(q < 0.5, (2.0*q)**(1.0/c), (2*(1.0-q))**(-1.0/c)) + + def _isf(self, q, c): + return np.where(q > 0.5, (2.0*(1.0 - q))**(1.0/c), (2*q)**(-1.0/c)) + + def _munp(self, n, c): + with np.errstate(divide='ignore'): + c2, n2 = c**2, n**2 + return np.where(n2 < c2, c2 / (c2 - n2), np.inf) + + def _entropy(self, c): + return np.log(2.0/c) + 1.0 + + @_call_super_mom + @inherit_docstring_from(rv_continuous) + def fit(self, data, *args, **kwds): + data, fc, floc, fscale = _check_fit_input_parameters(self, data, + args, kwds) + + # Specialize MLE only when location is known. + if floc is None: + return super(type(self), self).fit(data, *args, **kwds) + + # Raise an error if any observation has zero likelihood. + if np.any(data <= floc): + raise FitDataError("loglaplace", lower=floc, upper=np.inf) + + # Remove location from data. + if floc != 0: + data = data - floc + + # When location is zero, the log-Laplace distribution is related to + # the Laplace distribution in that if X ~ Laplace(loc=a, scale=b), + # then Y = exp(X) ~ LogLaplace(c=1/b, loc=0, scale=exp(a)). It can + # be shown that the MLE for Y is the same as the MLE for X = ln(Y). + # Therefore, we reuse the formulas from laplace.fit() and transform + # the result back into log-laplace's parameter space. + a, b = laplace.fit(np.log(data), + floc=np.log(fscale) if fscale is not None else None, + fscale=1/fc if fc is not None else None, + method='mle') + loc = floc + scale = np.exp(a) if fscale is None else fscale + c = 1 / b if fc is None else fc + return c, loc, scale + +loglaplace = loglaplace_gen(a=0.0, name='loglaplace') + + +def _lognorm_logpdf(x, s): + return xpx.apply_where( + x != 0, (x, s), + lambda x, s: (-np.log(x)**2 / (2 * s**2) + - np.log(s * x * np.sqrt(2 * np.pi))), + fill_value=-np.inf) + + +class lognorm_gen(rv_continuous): + r"""A lognormal continuous random variable. + + %(before_notes)s + + Notes + ----- + The probability density function for `lognorm` is: + + .. math:: + + f(x, s) = \frac{1}{s x \sqrt{2\pi}} + \exp\left(-\frac{\log^2(x)}{2s^2}\right) + + for :math:`x > 0`, :math:`s > 0`. + + `lognorm` takes ``s`` as a shape parameter for :math:`s`. + + %(after_notes)s + + Suppose a normally distributed random variable ``X`` has mean ``mu`` and + standard deviation ``sigma``. Then ``Y = exp(X)`` is lognormally + distributed with ``s = sigma`` and ``scale = exp(mu)``. + + %(example)s + + The logarithm of a log-normally distributed random variable is + normally distributed: + + >>> import numpy as np + >>> import matplotlib.pyplot as plt + >>> from scipy import stats + >>> fig, ax = plt.subplots(1, 1) + >>> mu, sigma = 2, 0.5 + >>> X = stats.norm(loc=mu, scale=sigma) + >>> Y = stats.lognorm(s=sigma, scale=np.exp(mu)) + >>> x = np.linspace(*X.interval(0.999)) + >>> y = Y.rvs(size=10000) + >>> ax.plot(x, X.pdf(x), label='X (pdf)') + >>> ax.hist(np.log(y), density=True, bins=x, label='log(Y) (histogram)') + >>> ax.legend() + >>> plt.show() + + """ + _support_mask = rv_continuous._open_support_mask + + def _shape_info(self): + return [_ShapeInfo("s", False, (0, np.inf), (False, False))] + + def _rvs(self, s, size=None, random_state=None): + return np.exp(s * random_state.standard_normal(size)) + + def _pdf(self, x, s): + # lognorm.pdf(x, s) = 1 / (s*x*sqrt(2*pi)) * exp(-1/2*(log(x)/s)**2) + return np.exp(self._logpdf(x, s)) + + def _logpdf(self, x, s): + return _lognorm_logpdf(x, s) + + def _cdf(self, x, s): + return _norm_cdf(np.log(x) / s) + + def _logcdf(self, x, s): + return _norm_logcdf(np.log(x) / s) + + def _ppf(self, q, s): + return np.exp(s * _norm_ppf(q)) + + def _sf(self, x, s): + return _norm_sf(np.log(x) / s) + + def _logsf(self, x, s): + return _norm_logsf(np.log(x) / s) + + def _isf(self, q, s): + return np.exp(s * _norm_isf(q)) + + def _stats(self, s): + p = np.exp(s*s) + mu = np.sqrt(p) + mu2 = p*(p-1) + g1 = np.sqrt(p-1)*(2+p) + g2 = np.polyval([1, 2, 3, 0, -6.0], p) + return mu, mu2, g1, g2 + + def _entropy(self, s): + return 0.5 * (1 + np.log(2*np.pi) + 2 * np.log(s)) + + @_call_super_mom + @extend_notes_in_docstring(rv_continuous, notes="""\ + When `method='MLE'` and + the location parameter is fixed by using the `floc` argument, + this function uses explicit formulas for the maximum likelihood + estimation of the log-normal shape and scale parameters, so the + `optimizer`, `loc` and `scale` keyword arguments are ignored. + If the location is free, a likelihood maximum is found by + setting its partial derivative w.r.t. location to 0, and + solving by substituting the analytical expressions of shape + and scale (or provided parameters). + See, e.g., equation 3.1 in + A. Clifford Cohen & Betty Jones Whitten (1980) + Estimation in the Three-Parameter Lognormal Distribution, + Journal of the American Statistical Association, 75:370, 399-404 + https://doi.org/10.2307/2287466 + \n\n""") + def fit(self, data, *args, **kwds): + if kwds.pop('superfit', False): + return super().fit(data, *args, **kwds) + + parameters = _check_fit_input_parameters(self, data, args, kwds) + data, fshape, floc, fscale = parameters + data_min = np.min(data) + + def get_shape_scale(loc): + # Calculate maximum likelihood scale and shape with analytical + # formulas unless provided by the user + if fshape is None or fscale is None: + lndata = np.log(data - loc) + scale = fscale or np.exp(lndata.mean()) + shape = fshape or np.sqrt(np.mean((lndata - np.log(scale))**2)) + return shape, scale + + def dL_dLoc(loc): + # Derivative of (positive) LL w.r.t. loc + shape, scale = get_shape_scale(loc) + shifted = data - loc + return np.sum((1 + np.log(shifted/scale)/shape**2)/shifted) + + def ll(loc): + # (Positive) log-likelihood + shape, scale = get_shape_scale(loc) + return -self.nnlf((shape, loc, scale), data) + + if floc is None: + # The location must be less than the minimum of the data. + # Back off a bit to avoid numerical issues. + spacing = np.spacing(data_min) + rbrack = data_min - spacing + + # Find the right end of the bracket by successive doubling of the + # distance to data_min. We're interested in a maximum LL, so the + # slope dL_dLoc_rbrack should be negative at the right end. + # optimization for later: share shape, scale + dL_dLoc_rbrack = dL_dLoc(rbrack) + ll_rbrack = ll(rbrack) + delta = 2 * spacing # 2 * (data_min - rbrack) + while dL_dLoc_rbrack >= -1e-6: + rbrack = data_min - delta + dL_dLoc_rbrack = dL_dLoc(rbrack) + delta *= 2 + + if not np.isfinite(rbrack) or not np.isfinite(dL_dLoc_rbrack): + # If we never find a negative slope, either we missed it or the + # slope is always positive. It's usually the latter, + # which means + # loc = data_min - spacing + # But sometimes when shape and/or scale are fixed there are + # other issues, so be cautious. + return super().fit(data, *args, **kwds) + + # Now find the left end of the bracket. Guess is `rbrack-1` + # unless that is too small of a difference to resolve. Double + # the size of the interval until the left end is found. + lbrack = np.minimum(np.nextafter(rbrack, -np.inf), rbrack-1) + dL_dLoc_lbrack = dL_dLoc(lbrack) + delta = 2 * (rbrack - lbrack) + while (np.isfinite(lbrack) and np.isfinite(dL_dLoc_lbrack) + and np.sign(dL_dLoc_lbrack) == np.sign(dL_dLoc_rbrack)): + lbrack = rbrack - delta + dL_dLoc_lbrack = dL_dLoc(lbrack) + delta *= 2 + + # I don't recall observing this, but just in case... + if not np.isfinite(lbrack) or not np.isfinite(dL_dLoc_lbrack): + return super().fit(data, *args, **kwds) + + # If we have a valid bracket, find the root + res = root_scalar(dL_dLoc, bracket=(lbrack, rbrack)) + if not res.converged: + return super().fit(data, *args, **kwds) + + # If the slope was positive near the minimum of the data, + # the maximum LL could be there instead of at the root. Compare + # the LL of the two points to decide. + ll_root = ll(res.root) + loc = res.root if ll_root > ll_rbrack else data_min-spacing + + else: + if floc >= data_min: + raise FitDataError("lognorm", lower=0., upper=np.inf) + loc = floc + + shape, scale = get_shape_scale(loc) + if not (self._argcheck(shape) and scale > 0): + return super().fit(data, *args, **kwds) + return shape, loc, scale + + +lognorm = lognorm_gen(a=0.0, name='lognorm') + + +class gibrat_gen(rv_continuous): + r"""A Gibrat continuous random variable. + + %(before_notes)s + + Notes + ----- + The probability density function for `gibrat` is: + + .. math:: + + f(x) = \frac{1}{x \sqrt{2\pi}} \exp(-\frac{1}{2} (\log(x))^2) + + for :math:`x >= 0`. + + `gibrat` is a special case of `lognorm` with ``s=1``. + + %(after_notes)s + + %(example)s + + """ + _support_mask = rv_continuous._open_support_mask + + def _shape_info(self): + return [] + + def _rvs(self, size=None, random_state=None): + return np.exp(random_state.standard_normal(size)) + + def _pdf(self, x): + # gibrat.pdf(x) = 1/(x*sqrt(2*pi)) * exp(-1/2*(log(x))**2) + return np.exp(self._logpdf(x)) + + def _logpdf(self, x): + return _lognorm_logpdf(x, 1.0) + + def _cdf(self, x): + return _norm_cdf(np.log(x)) + + def _ppf(self, q): + return np.exp(_norm_ppf(q)) + + def _sf(self, x): + return _norm_sf(np.log(x)) + + def _isf(self, p): + return np.exp(_norm_isf(p)) + + def _stats(self): + p = np.e + mu = np.sqrt(p) + mu2 = p * (p - 1) + g1 = np.sqrt(p - 1) * (2 + p) + g2 = np.polyval([1, 2, 3, 0, -6.0], p) + return mu, mu2, g1, g2 + + def _entropy(self): + return 0.5 * np.log(2 * np.pi) + 0.5 + + +gibrat = gibrat_gen(a=0.0, name='gibrat') + + +class maxwell_gen(rv_continuous): + r"""A Maxwell continuous random variable. + + %(before_notes)s + + Notes + ----- + A special case of a `chi` distribution, with ``df=3``, ``loc=0.0``, + and given ``scale = a``, where ``a`` is the parameter used in the + Mathworld description [1]_. + + The probability density function for `maxwell` is: + + .. math:: + + f(x) = \sqrt{2/\pi}x^2 \exp(-x^2/2) + + for :math:`x >= 0`. + + %(after_notes)s + + References + ---------- + .. [1] http://mathworld.wolfram.com/MaxwellDistribution.html + + %(example)s + """ + def _shape_info(self): + return [] + + def _rvs(self, size=None, random_state=None): + return chi.rvs(3.0, size=size, random_state=random_state) + + def _pdf(self, x): + # maxwell.pdf(x) = sqrt(2/pi)x**2 * exp(-x**2/2) + return _SQRT_2_OVER_PI*x*x*np.exp(-x*x/2.0) + + def _logpdf(self, x): + # Allow x=0 without 'divide by zero' warnings + with np.errstate(divide='ignore'): + return _LOG_SQRT_2_OVER_PI + 2*np.log(x) - 0.5*x*x + + def _cdf(self, x): + return sc.gammainc(1.5, x*x/2.0) + + def _ppf(self, q): + return np.sqrt(2*sc.gammaincinv(1.5, q)) + + def _sf(self, x): + return sc.gammaincc(1.5, x*x/2.0) + + def _isf(self, q): + return np.sqrt(2*sc.gammainccinv(1.5, q)) + + def _stats(self): + val = 3*np.pi-8 + return (2*np.sqrt(2.0/np.pi), + 3-8/np.pi, + np.sqrt(2)*(32-10*np.pi)/val**1.5, + (-12*np.pi*np.pi + 160*np.pi - 384) / val**2.0) + + def _entropy(self): + return _EULER + 0.5*np.log(2*np.pi)-0.5 + + +maxwell = maxwell_gen(a=0.0, name='maxwell') + + +class mielke_gen(rv_continuous): + r"""A Mielke Beta-Kappa / Dagum continuous random variable. + + %(before_notes)s + + Notes + ----- + The probability density function for `mielke` is: + + .. math:: + + f(x, k, s) = \frac{k x^{k-1}}{(1+x^s)^{1+k/s}} + + for :math:`x > 0` and :math:`k, s > 0`. The distribution is sometimes + called Dagum distribution ([2]_). It was already defined in [3]_, called + a Burr Type III distribution (`burr` with parameters ``c=s`` and + ``d=k/s``). + + `mielke` takes ``k`` and ``s`` as shape parameters. + + %(after_notes)s + + References + ---------- + .. [1] Mielke, P.W., 1973 "Another Family of Distributions for Describing + and Analyzing Precipitation Data." J. Appl. Meteor., 12, 275-280 + .. [2] Dagum, C., 1977 "A new model for personal income distribution." + Economie Appliquee, 33, 327-367. + .. [3] Burr, I. W. "Cumulative frequency functions", Annals of + Mathematical Statistics, 13(2), pp 215-232 (1942). + + %(example)s + + """ + def _shape_info(self): + ik = _ShapeInfo("k", False, (0, np.inf), (False, False)) + i_s = _ShapeInfo("s", False, (0, np.inf), (False, False)) + return [ik, i_s] + + def _pdf(self, x, k, s): + return k*x**(k-1.0) / (1.0+x**s)**(1.0+k*1.0/s) + + def _logpdf(self, x, k, s): + # Allow x=0 without 'divide by zero' warnings. + with np.errstate(divide='ignore'): + return np.log(k) + np.log(x)*(k - 1) - np.log1p(x**s)*(1 + k/s) + + def _cdf(self, x, k, s): + return x**k / (1.0+x**s)**(k*1.0/s) + + def _ppf(self, q, k, s): + qsk = pow(q, s*1.0/k) + return pow(qsk/(1.0-qsk), 1.0/s) + + def _munp(self, n, k, s): + def nth_moment(n, k, s): + # n-th moment is defined for -k < n < s + return sc.gamma((k+n)/s)*sc.gamma(1-n/s)/sc.gamma(k/s) + + return xpx.apply_where(n < s, (n, k, s), nth_moment, fill_value=np.inf) + + +mielke = mielke_gen(a=0.0, name='mielke') + + +class kappa4_gen(rv_continuous): + r"""Kappa 4 parameter distribution. + + %(before_notes)s + + Notes + ----- + The probability density function for kappa4 is: + + .. math:: + + f(x, h, k) = (1 - k x)^{1/k - 1} (1 - h (1 - k x)^{1/k})^{1/h-1} + + if :math:`h` and :math:`k` are not equal to 0. + + If :math:`h` or :math:`k` are zero then the pdf can be simplified: + + :math:`h = 0` and :math:`k \neq 0`:: + + kappa4.pdf(x, h, k) = (1.0 - k*x)**(1.0/k - 1.0)* + exp(-(1.0 - k*x)**(1.0/k)) + + :math:`h \neq 0` and :math:`k = 0`:: + + kappa4.pdf(x, h, k) = exp(-x)*(1.0 - h*exp(-x))**(1.0/h - 1.0) + + :math:`h = 0` and :math:`k = 0`:: + + kappa4.pdf(x, h, k) = exp(-x)*exp(-exp(-x)) + + kappa4 takes :math:`h` and :math:`k` as shape parameters. + + The kappa4 distribution returns other distributions when certain + :math:`h` and :math:`k` values are used. + + +------+-------------+----------------+------------------+ + | h | k=0.0 | k=1.0 | -inf<=k<=inf | + +======+=============+================+==================+ + | -1.0 | Logistic | | Generalized | + | | | | Logistic(1) | + | | | | | + | | logistic(x) | | | + +------+-------------+----------------+------------------+ + | 0.0 | Gumbel | Reverse | Generalized | + | | | Exponential(2) | Extreme Value | + | | | | | + | | gumbel_r(x) | | genextreme(x, k) | + +------+-------------+----------------+------------------+ + | 1.0 | Exponential | Uniform | Generalized | + | | | | Pareto | + | | | | | + | | expon(x) | uniform(x) | genpareto(x, -k) | + +------+-------------+----------------+------------------+ + + (1) There are at least five generalized logistic distributions. + Four are described here: + https://en.wikipedia.org/wiki/Generalized_logistic_distribution + The "fifth" one is the one kappa4 should match which currently + isn't implemented in scipy: + https://en.wikipedia.org/wiki/Talk:Generalized_logistic_distribution + https://www.mathwave.com/help/easyfit/html/analyses/distributions/gen_logistic.html + (2) This distribution is currently not in scipy. + + References + ---------- + J.C. Finney, "Optimization of a Skewed Logistic Distribution With Respect + to the Kolmogorov-Smirnov Test", A Dissertation Submitted to the Graduate + Faculty of the Louisiana State University and Agricultural and Mechanical + College, (August, 2004), + https://digitalcommons.lsu.edu/gradschool_dissertations/3672 + + J.R.M. Hosking, "The four-parameter kappa distribution". IBM J. Res. + Develop. 38 (3), 25 1-258 (1994). + + B. Kumphon, A. Kaew-Man, P. Seenoi, "A Rainfall Distribution for the Lampao + Site in the Chi River Basin, Thailand", Journal of Water Resource and + Protection, vol. 4, 866-869, (2012). + :doi:`10.4236/jwarp.2012.410101` + + C. Winchester, "On Estimation of the Four-Parameter Kappa Distribution", A + Thesis Submitted to Dalhousie University, Halifax, Nova Scotia, (March + 2000). + http://www.nlc-bnc.ca/obj/s4/f2/dsk2/ftp01/MQ57336.pdf + + %(after_notes)s + + %(example)s + + """ + def _argcheck(self, h, k): + shape = np.broadcast_arrays(h, k)[0].shape + return np.full(shape, fill_value=True) + + def _shape_info(self): + ih = _ShapeInfo("h", False, (-np.inf, np.inf), (False, False)) + ik = _ShapeInfo("k", False, (-np.inf, np.inf), (False, False)) + return [ih, ik] + + def _get_support(self, h, k): + condlist = [np.logical_and(h > 0, k > 0), + np.logical_and(h > 0, k == 0), + np.logical_and(h > 0, k < 0), + np.logical_and(h <= 0, k > 0), + np.logical_and(h <= 0, k == 0), + np.logical_and(h <= 0, k < 0)] + + def f0(h, k): + return (1.0 - np.float_power(h, -k))/k + + def f1(h, k): + return np.log(h) + + def f3(h, k): + a = np.empty(np.shape(h)) + a[:] = -np.inf + return a + + def f5(h, k): + return 1.0/k + + _a = _lazyselect(condlist, + [f0, f1, f0, f3, f3, f5], + [h, k], + default=np.nan) + + def f0(h, k): + return 1.0/k + + def f1(h, k): + a = np.empty(np.shape(h)) + a[:] = np.inf + return a + + _b = _lazyselect(condlist, + [f0, f1, f1, f0, f1, f1], + [h, k], + default=np.nan) + return _a, _b + + def _pdf(self, x, h, k): + # kappa4.pdf(x, h, k) = (1.0 - k*x)**(1.0/k - 1.0)* + # (1.0 - h*(1.0 - k*x)**(1.0/k))**(1.0/h-1) + return np.exp(self._logpdf(x, h, k)) + + def _logpdf(self, x, h, k): + condlist = [np.logical_and(h != 0, k != 0), + np.logical_and(h == 0, k != 0), + np.logical_and(h != 0, k == 0), + np.logical_and(h == 0, k == 0)] + + def f0(x, h, k): + '''pdf = (1.0 - k*x)**(1.0/k - 1.0)*( + 1.0 - h*(1.0 - k*x)**(1.0/k))**(1.0/h-1.0) + logpdf = ... + ''' + return (sc.xlog1py(1.0/k - 1.0, -k*x) + + sc.xlog1py(1.0/h - 1.0, -h*(1.0 - k*x)**(1.0/k))) + + def f1(x, h, k): + '''pdf = (1.0 - k*x)**(1.0/k - 1.0)*np.exp(-( + 1.0 - k*x)**(1.0/k)) + logpdf = ... + ''' + return sc.xlog1py(1.0/k - 1.0, -k*x) - (1.0 - k*x)**(1.0/k) + + def f2(x, h, k): + '''pdf = np.exp(-x)*(1.0 - h*np.exp(-x))**(1.0/h - 1.0) + logpdf = ... + ''' + return -x + sc.xlog1py(1.0/h - 1.0, -h*np.exp(-x)) + + def f3(x, h, k): + '''pdf = np.exp(-x-np.exp(-x)) + logpdf = ... + ''' + return -x - np.exp(-x) + + return _lazyselect(condlist, + [f0, f1, f2, f3], + [x, h, k], + default=np.nan) + + def _cdf(self, x, h, k): + return np.exp(self._logcdf(x, h, k)) + + def _logcdf(self, x, h, k): + condlist = [np.logical_and(h != 0, k != 0), + np.logical_and(h == 0, k != 0), + np.logical_and(h != 0, k == 0), + np.logical_and(h == 0, k == 0)] + + def f0(x, h, k): + '''cdf = (1.0 - h*(1.0 - k*x)**(1.0/k))**(1.0/h) + logcdf = ... + ''' + return (1.0/h)*sc.log1p(-h*(1.0 - k*x)**(1.0/k)) + + def f1(x, h, k): + '''cdf = np.exp(-(1.0 - k*x)**(1.0/k)) + logcdf = ... + ''' + return -(1.0 - k*x)**(1.0/k) + + def f2(x, h, k): + '''cdf = (1.0 - h*np.exp(-x))**(1.0/h) + logcdf = ... + ''' + return (1.0/h)*sc.log1p(-h*np.exp(-x)) + + def f3(x, h, k): + '''cdf = np.exp(-np.exp(-x)) + logcdf = ... + ''' + return -np.exp(-x) + + return _lazyselect(condlist, + [f0, f1, f2, f3], + [x, h, k], + default=np.nan) + + def _ppf(self, q, h, k): + condlist = [np.logical_and(h != 0, k != 0), + np.logical_and(h == 0, k != 0), + np.logical_and(h != 0, k == 0), + np.logical_and(h == 0, k == 0)] + + def f0(q, h, k): + return 1.0/k*(1.0 - ((1.0 - (q**h))/h)**k) + + def f1(q, h, k): + return 1.0/k*(1.0 - (-np.log(q))**k) + + def f2(q, h, k): + '''ppf = -np.log((1.0 - (q**h))/h) + ''' + return -sc.log1p(-(q**h)) + np.log(h) + + def f3(q, h, k): + return -np.log(-np.log(q)) + + return _lazyselect(condlist, + [f0, f1, f2, f3], + [q, h, k], + default=np.nan) + + def _get_stats_info(self, h, k): + condlist = [ + np.logical_and(h < 0, k >= 0), + k < 0, + ] + + def f0(h, k): + return (-1.0/h*k).astype(int) + + def f1(h, k): + return (-1.0/k).astype(int) + + return _lazyselect(condlist, [f0, f1], [h, k], default=5) + + def _stats(self, h, k): + maxr = self._get_stats_info(h, k) + outputs = [None if np.any(r < maxr) else np.nan for r in range(1, 5)] + return outputs[:] + + def _mom1_sc(self, m, *args): + maxr = self._get_stats_info(args[0], args[1]) + if m >= maxr: + return np.nan + return integrate.quad(self._mom_integ1, 0, 1, args=(m,)+args)[0] + + +kappa4 = kappa4_gen(name='kappa4') + + +class kappa3_gen(rv_continuous): + r"""Kappa 3 parameter distribution. + + %(before_notes)s + + Notes + ----- + The probability density function for `kappa3` is: + + .. math:: + + f(x, a) = a (a + x^a)^{-(a + 1)/a} + + for :math:`x > 0` and :math:`a > 0`. + + `kappa3` takes ``a`` as a shape parameter for :math:`a`. + + References + ---------- + P.W. Mielke and E.S. Johnson, "Three-Parameter Kappa Distribution Maximum + Likelihood and Likelihood Ratio Tests", Methods in Weather Research, + 701-707, (September, 1973), + :doi:`10.1175/1520-0493(1973)101<0701:TKDMLE>2.3.CO;2` + + B. Kumphon, "Maximum Entropy and Maximum Likelihood Estimation for the + Three-Parameter Kappa Distribution", Open Journal of Statistics, vol 2, + 415-419 (2012), :doi:`10.4236/ojs.2012.24050` + + %(after_notes)s + + %(example)s + + """ + def _shape_info(self): + return [_ShapeInfo("a", False, (0, np.inf), (False, False))] + + def _pdf(self, x, a): + # kappa3.pdf(x, a) = a*(a + x**a)**(-(a + 1)/a), for x > 0 + return a*(a + x**a)**(-1.0/a-1) + + def _cdf(self, x, a): + return x*(a + x**a)**(-1.0/a) + + def _sf(self, x, a): + x, a = np.broadcast_arrays(x, a) # some code paths pass scalars + sf = super()._sf(x, a) + + # When the SF is small, another formulation is typically more accurate. + # However, it blows up for large `a`, so use it only if it also returns + # a small value of the SF. + cutoff = 0.01 + i = sf < cutoff + sf2 = -sc.expm1(sc.xlog1py(-1.0 / a[i], a[i] * x[i]**-a[i])) + i2 = sf2 > cutoff + sf2[i2] = sf[i][i2] # replace bad values with original values + + sf[i] = sf2 + return sf + + def _ppf(self, q, a): + return (a/(q**-a - 1.0))**(1.0/a) + + def _isf(self, q, a): + lg = sc.xlog1py(-a, -q) + denom = sc.expm1(lg) + return (a / denom)**(1.0 / a) + + def _stats(self, a): + outputs = [None if np.any(i < a) else np.nan for i in range(1, 5)] + return outputs[:] + + def _mom1_sc(self, m, *args): + if np.any(m >= args[0]): + return np.nan + return integrate.quad(self._mom_integ1, 0, 1, args=(m,)+args)[0] + + +kappa3 = kappa3_gen(a=0.0, name='kappa3') + + +class moyal_gen(rv_continuous): + r"""A Moyal continuous random variable. + + %(before_notes)s + + Notes + ----- + The probability density function for `moyal` is: + + .. math:: + + f(x) = \exp(-(x + \exp(-x))/2) / \sqrt{2\pi} + + for a real number :math:`x`. + + %(after_notes)s + + This distribution has utility in high-energy physics and radiation + detection. It describes the energy loss of a charged relativistic + particle due to ionization of the medium [1]_. It also provides an + approximation for the Landau distribution. For an in depth description + see [2]_. For additional description, see [3]_. + + References + ---------- + .. [1] J.E. Moyal, "XXX. Theory of ionization fluctuations", + The London, Edinburgh, and Dublin Philosophical Magazine + and Journal of Science, vol 46, 263-280, (1955). + :doi:`10.1080/14786440308521076` (gated) + .. [2] G. Cordeiro et al., "The beta Moyal: A useful skew distribution", + International Journal of Research and Reviews in Applied Sciences, + vol 10, 171-192, (2012). + https://www.arpapress.com/files/volumes/vol10issue2/ijrras_10_2_02.pdf + .. [3] C. Walck, "Handbook on Statistical Distributions for + Experimentalists; International Report SUF-PFY/96-01", Chapter 26, + University of Stockholm: Stockholm, Sweden, (2007). + http://www.stat.rice.edu/~dobelman/textfiles/DistributionsHandbook.pdf + + .. versionadded:: 1.1.0 + + %(example)s + + """ + def _shape_info(self): + return [] + + def _rvs(self, size=None, random_state=None): + u1 = gamma.rvs(a=0.5, scale=2, size=size, + random_state=random_state) + return -np.log(u1) + + def _pdf(self, x): + return np.exp(-0.5 * (x + np.exp(-x))) / np.sqrt(2*np.pi) + + def _cdf(self, x): + return sc.erfc(np.exp(-0.5 * x) / np.sqrt(2)) + + def _sf(self, x): + return sc.erf(np.exp(-0.5 * x) / np.sqrt(2)) + + def _ppf(self, x): + return -np.log(2 * sc.erfcinv(x)**2) + + def _stats(self): + mu = np.log(2) + np.euler_gamma + mu2 = np.pi**2 / 2 + g1 = 28 * np.sqrt(2) * sc.zeta(3) / np.pi**3 + g2 = 4. + return mu, mu2, g1, g2 + + def _munp(self, n): + if n == 1.0: + return np.log(2) + np.euler_gamma + elif n == 2.0: + return np.pi**2 / 2 + (np.log(2) + np.euler_gamma)**2 + elif n == 3.0: + tmp1 = 1.5 * np.pi**2 * (np.log(2)+np.euler_gamma) + tmp2 = (np.log(2)+np.euler_gamma)**3 + tmp3 = 14 * sc.zeta(3) + return tmp1 + tmp2 + tmp3 + elif n == 4.0: + tmp1 = 4 * 14 * sc.zeta(3) * (np.log(2) + np.euler_gamma) + tmp2 = 3 * np.pi**2 * (np.log(2) + np.euler_gamma)**2 + tmp3 = (np.log(2) + np.euler_gamma)**4 + tmp4 = 7 * np.pi**4 / 4 + return tmp1 + tmp2 + tmp3 + tmp4 + else: + # return generic for higher moments + # return rv_continuous._mom1_sc(self, n, b) + return self._mom1_sc(n) + + +moyal = moyal_gen(name="moyal") + + +class nakagami_gen(rv_continuous): + r"""A Nakagami continuous random variable. + + %(before_notes)s + + Notes + ----- + The probability density function for `nakagami` is: + + .. math:: + + f(x, \nu) = \frac{2 \nu^\nu}{\Gamma(\nu)} x^{2\nu-1} \exp(-\nu x^2) + + for :math:`x >= 0`, :math:`\nu > 0`. The distribution was introduced in + [2]_, see also [1]_ for further information. + + `nakagami` takes ``nu`` as a shape parameter for :math:`\nu`. + + %(after_notes)s + + References + ---------- + .. [1] "Nakagami distribution", Wikipedia + https://en.wikipedia.org/wiki/Nakagami_distribution + .. [2] M. Nakagami, "The m-distribution - A general formula of intensity + distribution of rapid fading", Statistical methods in radio wave + propagation, Pergamon Press, 1960, 3-36. + :doi:`10.1016/B978-0-08-009306-2.50005-4` + + %(example)s + + """ + def _argcheck(self, nu): + return nu > 0 + + def _shape_info(self): + return [_ShapeInfo("nu", False, (0, np.inf), (False, False))] + + def _pdf(self, x, nu): + return np.exp(self._logpdf(x, nu)) + + def _logpdf(self, x, nu): + # nakagami.pdf(x, nu) = 2 * nu**nu / gamma(nu) * + # x**(2*nu-1) * exp(-nu*x**2) + return (np.log(2) + sc.xlogy(nu, nu) - sc.gammaln(nu) + + sc.xlogy(2*nu - 1, x) - nu*x**2) + + def _cdf(self, x, nu): + return sc.gammainc(nu, nu*x*x) + + def _ppf(self, q, nu): + return np.sqrt(1.0/nu*sc.gammaincinv(nu, q)) + + def _sf(self, x, nu): + return sc.gammaincc(nu, nu*x*x) + + def _isf(self, p, nu): + return np.sqrt(1/nu * sc.gammainccinv(nu, p)) + + def _stats(self, nu): + mu = sc.poch(nu, 0.5)/np.sqrt(nu) + mu2 = 1.0-mu*mu + g1 = mu * (1 - 4*nu*mu2) / 2.0 / nu / np.power(mu2, 1.5) + g2 = -6*mu**4*nu + (8*nu-2)*mu**2-2*nu + 1 + g2 /= nu*mu2**2.0 + return mu, mu2, g1, g2 + + def _entropy(self, nu): + shape = np.shape(nu) + # because somehow this isn't taken care of by the infrastructure... + nu = np.atleast_1d(nu) + A = sc.gammaln(nu) + B = nu - (nu - 0.5) * sc.digamma(nu) + C = -0.5 * np.log(nu) - np.log(2) + h = A + B + C + # This is the asymptotic sum of A and B (see gh-17868) + norm_entropy = stats.norm._entropy() + # Above, this is lost to rounding error for large nu, so use the + # asymptotic sum when the approximation becomes accurate + i = nu > 5e4 # roundoff error ~ approximation error + # -1 / (12 * nu) is the O(1/nu) term; see gh-17929 + h[i] = C[i] + norm_entropy - 1/(12*nu[i]) + return h.reshape(shape)[()] + + def _rvs(self, nu, size=None, random_state=None): + # this relationship can be found in [1] or by a direct calculation + return np.sqrt(random_state.standard_gamma(nu, size=size) / nu) + + def _fitstart(self, data, args=None): + if isinstance(data, CensoredData): + data = data._uncensor() + if args is None: + args = (1.0,) * self.numargs + # Analytical justified estimates + # see: https://docs.scipy.org/doc/scipy/reference/tutorial/stats/continuous_nakagami.html + loc = np.min(data) + scale = np.sqrt(np.sum((data - loc)**2) / len(data)) + return args + (loc, scale) + + +nakagami = nakagami_gen(a=0.0, name="nakagami") + + +# The function name ncx2 is an abbreviation for noncentral chi squared. +def _ncx2_log_pdf(x, df, nc): + # We use (xs**2 + ns**2)/2 = (xs - ns)**2/2 + xs*ns, and include the + # factor of exp(-xs*ns) into the ive function to improve numerical + # stability at large values of xs. See also `rice.pdf`. + df2 = df/2.0 - 1.0 + xs, ns = np.sqrt(x), np.sqrt(nc) + res = sc.xlogy(df2/2.0, x/nc) - 0.5*(xs - ns)**2 + corr = sc.ive(df2, xs*ns) / 2.0 + # Return res + np.log(corr) avoiding np.log(0) + return xpx.apply_where( + corr > 0, + (res, corr), + lambda r, c: r + np.log(c), + fill_value=-np.inf) + + +class ncx2_gen(rv_continuous): + r"""A non-central chi-squared continuous random variable. + + %(before_notes)s + + Notes + ----- + The probability density function for `ncx2` is: + + .. math:: + + f(x, k, \lambda) = \frac{1}{2} \exp(-(\lambda+x)/2) + (x/\lambda)^{(k-2)/4} I_{(k-2)/2}(\sqrt{\lambda x}) + + for :math:`x >= 0`, :math:`k > 0` and :math:`\lambda \ge 0`. + :math:`k` specifies the degrees of freedom (denoted ``df`` in the + implementation) and :math:`\lambda` is the non-centrality parameter + (denoted ``nc`` in the implementation). :math:`I_\nu` denotes the + modified Bessel function of first order of degree :math:`\nu` + (`scipy.special.iv`). + + `ncx2` takes ``df`` and ``nc`` as shape parameters. + + This distribution uses routines from the Boost Math C++ library for + the computation of the ``pdf``, ``cdf``, ``ppf``, ``sf`` and ``isf`` + methods. [1]_ + + %(after_notes)s + + References + ---------- + .. [1] The Boost Developers. "Boost C++ Libraries". https://www.boost.org/. + + %(example)s + + """ + def _argcheck(self, df, nc): + return (df > 0) & np.isfinite(df) & (nc >= 0) + + def _shape_info(self): + idf = _ShapeInfo("df", False, (0, np.inf), (False, False)) + inc = _ShapeInfo("nc", False, (0, np.inf), (True, False)) + return [idf, inc] + + def _rvs(self, df, nc, size=None, random_state=None): + return random_state.noncentral_chisquare(df, nc, size) + + def _logpdf(self, x, df, nc): + return xpx.apply_where(nc != 0, (x, df, nc), _ncx2_log_pdf, + lambda x, df, _: chi2._logpdf(x, df)) + + def _pdf(self, x, df, nc): + with np.errstate(over='ignore'): # see gh-17432 + return xpx.apply_where(nc != 0, (x, df, nc), scu._ncx2_pdf, + lambda x, df, _: chi2._pdf(x, df)) + + def _cdf(self, x, df, nc): + with np.errstate(over='ignore'): # see gh-17432 + return xpx.apply_where(nc != 0, (x, df, nc), sc.chndtr, + lambda x, df, _: chi2._cdf(x, df)) + + def _ppf(self, q, df, nc): + with np.errstate(over='ignore'): # see gh-17432 + return xpx.apply_where(nc != 0, (q, df, nc), sc.chndtrix, + lambda x, df, _: chi2._ppf(x, df)) + + def _sf(self, x, df, nc): + with np.errstate(over='ignore'): # see gh-17432 + return xpx.apply_where(nc != 0, (x, df, nc), scu._ncx2_sf, + lambda x, df, _: chi2._sf(x, df)) + + def _isf(self, x, df, nc): + with np.errstate(over='ignore'): # see gh-17432 + return xpx.apply_where(nc != 0, (x, df, nc), scu._ncx2_isf, + lambda x, df, _: chi2._isf(x, df)) + + def _stats(self, df, nc): + _ncx2_mean = df + nc + def k_plus_cl(k, l, c): + return k + c*l + _ncx2_variance = 2.0 * k_plus_cl(df, nc, 2.0) + _ncx2_skewness = (np.sqrt(8.0) * k_plus_cl(df, nc, 3) / + np.sqrt(k_plus_cl(df, nc, 2.0)**3)) + _ncx2_kurtosis_excess = (12.0 * k_plus_cl(df, nc, 4.0) / + k_plus_cl(df, nc, 2.0)**2) + return ( + _ncx2_mean, + _ncx2_variance, + _ncx2_skewness, + _ncx2_kurtosis_excess, + ) + + +ncx2 = ncx2_gen(a=0.0, name='ncx2') + + +class ncf_gen(rv_continuous): + r"""A non-central F distribution continuous random variable. + + %(before_notes)s + + See Also + -------- + scipy.stats.f : Fisher distribution + + Notes + ----- + The probability density function for `ncf` is: + + .. math:: + + f(x, n_1, n_2, \lambda) = + \exp\left(\frac{\lambda}{2} + + \lambda n_1 \frac{x}{2(n_1 x + n_2)} + \right) + n_1^{n_1/2} n_2^{n_2/2} x^{n_1/2 - 1} \\ + (n_2 + n_1 x)^{-(n_1 + n_2)/2} + \gamma(n_1/2) \gamma(1 + n_2/2) \\ + \frac{L^{\frac{n_1}{2}-1}_{n_2/2} + \left(-\lambda n_1 \frac{x}{2(n_1 x + n_2)}\right)} + {B(n_1/2, n_2/2) + \gamma\left(\frac{n_1 + n_2}{2}\right)} + + for :math:`n_1, n_2 > 0`, :math:`\lambda \ge 0`. Here :math:`n_1` is the + degrees of freedom in the numerator, :math:`n_2` the degrees of freedom in + the denominator, :math:`\lambda` the non-centrality parameter, + :math:`\gamma` is the logarithm of the Gamma function, :math:`L_n^k` is a + generalized Laguerre polynomial and :math:`B` is the beta function. + + `ncf` takes ``dfn``, ``dfd`` and ``nc`` as shape parameters. If ``nc=0``, + the distribution becomes equivalent to the Fisher distribution. + + This distribution uses routines from the Boost Math C++ library for + the computation of the ``pdf``, ``cdf``, ``ppf``, ``stats``, ``sf`` and + ``isf`` methods. [1]_ + + %(after_notes)s + + References + ---------- + .. [1] The Boost Developers. "Boost C++ Libraries". https://www.boost.org/. + + %(example)s + + """ + def _argcheck(self, dfn, dfd, nc): + return (dfn > 0) & (dfd > 0) & (nc >= 0) + + def _shape_info(self): + idf1 = _ShapeInfo("dfn", False, (0, np.inf), (False, False)) + idf2 = _ShapeInfo("dfd", False, (0, np.inf), (False, False)) + inc = _ShapeInfo("nc", False, (0, np.inf), (True, False)) + return [idf1, idf2, inc] + + def _rvs(self, dfn, dfd, nc, size=None, random_state=None): + return random_state.noncentral_f(dfn, dfd, nc, size) + + def _pdf(self, x, dfn, dfd, nc): + return scu._ncf_pdf(x, dfn, dfd, nc) + + def _cdf(self, x, dfn, dfd, nc): + return sc.ncfdtr(dfn, dfd, nc, x) + + def _ppf(self, q, dfn, dfd, nc): + with np.errstate(over='ignore'): # see gh-17432 + return sc.ncfdtri(dfn, dfd, nc, q) + + def _sf(self, x, dfn, dfd, nc): + return scu._ncf_sf(x, dfn, dfd, nc) + + def _isf(self, x, dfn, dfd, nc): + with np.errstate(over='ignore'): # see gh-17432 + return scu._ncf_isf(x, dfn, dfd, nc) + + # # Produces bogus values as written - maybe it's close, though? + # def _munp(self, n, dfn, dfd, nc): + # val = (dfn * 1.0/dfd)**n + # term = sc.gammaln(n+0.5*dfn) + sc.gammaln(0.5*dfd-n) - sc.gammaln(dfd*0.5) + # val *= np.exp(-nc / 2.0+term) + # val *= sc.hyp1f1(n+0.5*dfn, 0.5*dfn, 0.5*nc) + # return val + + def _stats(self, dfn, dfd, nc, moments='mv'): + mu = scu._ncf_mean(dfn, dfd, nc) + mu2 = scu._ncf_variance(dfn, dfd, nc) + g1 = scu._ncf_skewness(dfn, dfd, nc) if 's' in moments else None + g2 = scu._ncf_kurtosis_excess( # isn't really excess kurtosis! + dfn, dfd, nc) - 3 if 'k' in moments else None + # Mathematica: Kurtosis[NoncentralFRatioDistribution[27, 27, 0.415784417992261]] + return mu, mu2, g1, g2 + + +ncf = ncf_gen(a=0.0, name='ncf') + + +class t_gen(rv_continuous): + r"""A Student's t continuous random variable. + + For the noncentral t distribution, see `nct`. + + %(before_notes)s + + See Also + -------- + nct + + Notes + ----- + The probability density function for `t` is: + + .. math:: + + f(x, \nu) = \frac{\Gamma((\nu+1)/2)} + {\sqrt{\pi \nu} \Gamma(\nu/2)} + (1+x^2/\nu)^{-(\nu+1)/2} + + where :math:`x` is a real number and the degrees of freedom parameter + :math:`\nu` (denoted ``df`` in the implementation) satisfies + :math:`\nu > 0`. :math:`\Gamma` is the gamma function + (`scipy.special.gamma`). + + %(after_notes)s + + %(example)s + + """ + def _shape_info(self): + return [_ShapeInfo("df", False, (0, np.inf), (False, False))] + + def _rvs(self, df, size=None, random_state=None): + return random_state.standard_t(df, size=size) + + def _pdf(self, x, df): + return xpx.apply_where( + df == np.inf, (x, df), + lambda x, df: norm._pdf(x), + lambda x, df: np.exp(self._logpdf(x, df))) + + def _logpdf(self, x, df): + + def t_logpdf(x, df): + return (np.log(sc.poch(0.5 * df, 0.5)) + - 0.5 * (np.log(df) + np.log(np.pi)) + - (df + 1)/2*np.log1p(x * x/df)) + + def norm_logpdf(x, df): + return norm._logpdf(x) + + return xpx.apply_where(df == np.inf, (x, df), norm_logpdf, t_logpdf) + + def _cdf(self, x, df): + return sc.stdtr(df, x) + + def _sf(self, x, df): + return sc.stdtr(df, -x) + + def _ppf(self, q, df): + return sc.stdtrit(df, q) + + def _isf(self, q, df): + return -sc.stdtrit(df, q) + + def _stats(self, df): + # infinite df -> normal distribution (0.0, 1.0, 0.0, 0.0) + infinite_df = np.isposinf(df) + + mu = np.where(df > 1, 0.0, np.inf) + + condlist = ((df > 1) & (df <= 2), + (df > 2) & np.isfinite(df), + infinite_df) + choicelist = (lambda df: np.broadcast_to(np.inf, df.shape), + lambda df: df / (df-2.0), + lambda df: np.broadcast_to(1, df.shape)) + mu2 = _lazyselect(condlist, choicelist, (df,), np.nan) + + g1 = np.where(df > 3, 0.0, np.nan) + + condlist = ((df > 2) & (df <= 4), + (df > 4) & np.isfinite(df), + infinite_df) + choicelist = (lambda df: np.broadcast_to(np.inf, df.shape), + lambda df: 6.0 / (df-4.0), + lambda df: np.broadcast_to(0, df.shape)) + g2 = _lazyselect(condlist, choicelist, (df,), np.nan) + + return mu, mu2, g1, g2 + + def _entropy(self, df): + + def regular(df): + half = df/2 + half1 = (df + 1)/2 + return (half1*(sc.digamma(half1) - sc.digamma(half)) + + np.log(np.sqrt(df)*sc.beta(half, 0.5))) + + def asymptotic(df): + # Formula from Wolfram Alpha: + # "asymptotic expansion (d+1)/2 * (digamma((d+1)/2) - digamma(d/2)) + # + log(sqrt(d) * beta(d/2, 1/2))" + h = (norm._entropy() + 1/df + (df**-2.)/4 - (df**-3.)/6 + - (df**-4.)/8 + 3/10*(df**-5.) + (df**-6.)/4) + return h + + return xpx.apply_where(df >= 100, df, asymptotic, regular) + + +t = t_gen(name='t') + + +class nct_gen(rv_continuous): + r"""A non-central Student's t continuous random variable. + + %(before_notes)s + + Notes + ----- + If :math:`Y` is a standard normal random variable and :math:`V` is + an independent chi-square random variable (`chi2`) with :math:`k` degrees + of freedom, then + + .. math:: + + X = \frac{Y + c}{\sqrt{V/k}} + + has a non-central Student's t distribution on the real line. + The degrees of freedom parameter :math:`k` (denoted ``df`` in the + implementation) satisfies :math:`k > 0` and the noncentrality parameter + :math:`c` (denoted ``nc`` in the implementation) is a real number. + + This distribution uses routines from the Boost Math C++ library for + the computation of the ``pdf``, ``cdf``, ``ppf``, ``sf`` and ``isf`` + methods. [1]_ + + %(after_notes)s + + References + ---------- + .. [1] The Boost Developers. "Boost C++ Libraries". https://www.boost.org/. + + %(example)s + + """ + def _argcheck(self, df, nc): + return (df > 0) & (nc == nc) + + def _shape_info(self): + idf = _ShapeInfo("df", False, (0, np.inf), (False, False)) + inc = _ShapeInfo("nc", False, (-np.inf, np.inf), (False, False)) + return [idf, inc] + + def _rvs(self, df, nc, size=None, random_state=None): + n = norm.rvs(loc=nc, size=size, random_state=random_state) + c2 = chi2.rvs(df, size=size, random_state=random_state) + return n * np.sqrt(df) / np.sqrt(c2) + + def _pdf(self, x, df, nc): + return scu._nct_pdf(x, df, nc) + + def _cdf(self, x, df, nc): + return sc.nctdtr(df, nc, x) + + def _ppf(self, q, df, nc): + return sc.nctdtrit(df, nc, q) + + def _sf(self, x, df, nc): + with np.errstate(over='ignore'): # see gh-17432 + return np.clip(scu._nct_sf(x, df, nc), 0, 1) + + def _isf(self, x, df, nc): + with np.errstate(over='ignore'): # see gh-17432 + return scu._nct_isf(x, df, nc) + + def _stats(self, df, nc, moments='mv'): + mu = scu._nct_mean(df, nc) + mu2 = scu._nct_variance(df, nc) + g1 = scu._nct_skewness(df, nc) if 's' in moments else None + g2 = scu._nct_kurtosis_excess(df, nc) if 'k' in moments else None + return mu, mu2, g1, g2 + + +nct = nct_gen(name="nct") + + +class pareto_gen(rv_continuous): + r"""A Pareto continuous random variable. + + %(before_notes)s + + Notes + ----- + The probability density function for `pareto` is: + + .. math:: + + f(x, b) = \frac{b}{x^{b+1}} + + for :math:`x \ge 1`, :math:`b > 0`. + + `pareto` takes ``b`` as a shape parameter for :math:`b`. + + %(after_notes)s + + %(example)s + + """ + def _shape_info(self): + return [_ShapeInfo("b", False, (0, np.inf), (False, False))] + + def _pdf(self, x, b): + # pareto.pdf(x, b) = b / x**(b+1) + return b * x**(-b-1) + + def _cdf(self, x, b): + return 1 - x**(-b) + + def _ppf(self, q, b): + return pow(1-q, -1.0/b) + + def _sf(self, x, b): + return x**(-b) + + def _isf(self, q, b): + return np.power(q, -1.0 / b) + + def _stats(self, b, moments='mv'): + mu, mu2, g1, g2 = None, None, None, None + if 'm' in moments: + mask = b > 1 + bt = np.extract(mask, b) + mu = np.full(np.shape(b), fill_value=np.inf) + np.place(mu, mask, bt / (bt-1.0)) + if 'v' in moments: + mask = b > 2 + bt = np.extract(mask, b) + mu2 = np.full(np.shape(b), fill_value=np.inf) + np.place(mu2, mask, bt / (bt-2.0) / (bt-1.0)**2) + if 's' in moments: + mask = b > 3 + bt = np.extract(mask, b) + g1 = np.full(np.shape(b), fill_value=np.nan) + vals = 2 * (bt + 1.0) * np.sqrt(bt - 2.0) / ((bt - 3.0) * np.sqrt(bt)) + np.place(g1, mask, vals) + if 'k' in moments: + mask = b > 4 + bt = np.extract(mask, b) + g2 = np.full(np.shape(b), fill_value=np.nan) + vals = (6.0*np.polyval([1.0, 1.0, -6, -2], bt) / + np.polyval([1.0, -7.0, 12.0, 0.0], bt)) + np.place(g2, mask, vals) + return mu, mu2, g1, g2 + + def _entropy(self, b): + return 1 + 1.0/b - np.log(b) + + @_call_super_mom + @inherit_docstring_from(rv_continuous) + def fit(self, data, *args, **kwds): + parameters = _check_fit_input_parameters(self, data, args, kwds) + data, fshape, floc, fscale = parameters + + # ensure that any fixed parameters don't violate constraints of the + # distribution before continuing. + if floc is not None and np.min(data) - floc < (fscale or 0): + raise FitDataError("pareto", lower=1, upper=np.inf) + + ndata = data.shape[0] + + def get_shape(scale, location): + # The first-order necessary condition on `shape` can be solved in + # closed form + return ndata / np.sum(np.log((data - location) / scale)) + + if floc is fscale is None: + # The support of the distribution is `(x - loc)/scale > 0`. + # The method of Lagrange multipliers turns this constraint + # into an equation that can be solved numerically. + # See gh-12545 for details. + + def dL_dScale(shape, scale): + # The partial derivative of the log-likelihood function w.r.t. + # the scale. + return ndata * shape / scale + + def dL_dLocation(shape, location): + # The partial derivative of the log-likelihood function w.r.t. + # the location. + return (shape + 1) * np.sum(1 / (data - location)) + + def fun_to_solve(scale): + # optimize the scale by setting the partial derivatives + # w.r.t. location and scale equal and solving. + location = np.min(data) - scale + shape = fshape or get_shape(scale, location) + return dL_dLocation(shape, location) - dL_dScale(shape, scale) + + def interval_contains_root(lbrack, rbrack): + # return true if the signs disagree. + return (np.sign(fun_to_solve(lbrack)) != + np.sign(fun_to_solve(rbrack))) + + # set brackets for `root_scalar` to use when optimizing over the + # scale such that a root is likely between them. Use user supplied + # guess or default 1. + brack_start = float(kwds.get('scale', 1)) + lbrack, rbrack = brack_start / 2, brack_start * 2 + # if a root is not between the brackets, iteratively expand them + # until they include a sign change, checking after each bracket is + # modified. + while (not interval_contains_root(lbrack, rbrack) + and (lbrack > 0 or rbrack < np.inf)): + lbrack /= 2 + rbrack *= 2 + res = root_scalar(fun_to_solve, bracket=[lbrack, rbrack]) + if res.converged: + scale = res.root + loc = np.min(data) - scale + shape = fshape or get_shape(scale, loc) + + # The Pareto distribution requires that its parameters satisfy + # the condition `fscale + floc <= min(data)`. However, to + # avoid numerical issues, we require that `fscale + floc` + # is strictly less than `min(data)`. If this condition + # is not satisfied, reduce the scale with `np.nextafter` to + # ensure that data does not fall outside of the support. + if not (scale + loc) < np.min(data): + scale = np.min(data) - loc + scale = np.nextafter(scale, 0) + return shape, loc, scale + else: + return super().fit(data, **kwds) + elif floc is None: + loc = np.min(data) - fscale + else: + loc = floc + # Source: Evans, Hastings, and Peacock (2000), Statistical + # Distributions, 3rd. Ed., John Wiley and Sons. Page 149. + scale = fscale or np.min(data) - loc + shape = fshape or get_shape(scale, loc) + return shape, loc, scale + + +pareto = pareto_gen(a=1.0, name="pareto") + + +class lomax_gen(rv_continuous): + r"""A Lomax (Pareto of the second kind) continuous random variable. + + %(before_notes)s + + Notes + ----- + The probability density function for `lomax` is: + + .. math:: + + f(x, c) = \frac{c}{(1+x)^{c+1}} + + for :math:`x \ge 0`, :math:`c > 0`. + + `lomax` takes ``c`` as a shape parameter for :math:`c`. + + `lomax` is a special case of `pareto` with ``loc=-1.0``. + + %(after_notes)s + + %(example)s + + """ + def _shape_info(self): + return [_ShapeInfo("c", False, (0, np.inf), (False, False))] + + def _pdf(self, x, c): + # lomax.pdf(x, c) = c / (1+x)**(c+1) + return c*1.0/(1.0+x)**(c+1.0) + + def _logpdf(self, x, c): + return np.log(c) - (c+1)*sc.log1p(x) + + def _cdf(self, x, c): + return -sc.expm1(-c*sc.log1p(x)) + + def _sf(self, x, c): + return np.exp(-c*sc.log1p(x)) + + def _logsf(self, x, c): + return -c*sc.log1p(x) + + def _ppf(self, q, c): + return sc.expm1(-sc.log1p(-q)/c) + + def _isf(self, q, c): + return q**(-1.0 / c) - 1 + + def _stats(self, c): + mu, mu2, g1, g2 = pareto.stats(c, loc=-1.0, moments='mvsk') + return mu, mu2, g1, g2 + + def _entropy(self, c): + return 1+1.0/c-np.log(c) + + +lomax = lomax_gen(a=0.0, name="lomax") + + +class pearson3_gen(rv_continuous): + r""" + A pearson type III continuous random variable. + + %(before_notes)s + + Notes + ----- + The probability density function for `pearson3` is: + + .. math:: + + f(x, \kappa) = \frac{|\beta|}{\Gamma(\alpha)} + (\beta (x - \zeta))^{\alpha - 1} + \exp(-\beta (x - \zeta)) + + where: + + .. math:: + + \beta = \frac{2}{\kappa} + + \alpha = \beta^2 = \frac{4}{\kappa^2} + + \zeta = -\frac{\alpha}{\beta} = -\beta + + :math:`\Gamma` is the gamma function (`scipy.special.gamma`). + Pass the skew :math:`\kappa` into `pearson3` as the shape parameter + ``skew``. + + %(after_notes)s + + References + ---------- + R.W. Vogel and D.E. McMartin, "Probability Plot Goodness-of-Fit and + Skewness Estimation Procedures for the Pearson Type 3 Distribution", Water + Resources Research, Vol.27, 3149-3158 (1991). + + L.R. Salvosa, "Tables of Pearson's Type III Function", Ann. Math. Statist., + Vol.1, 191-198 (1930). + + "Using Modern Computing Tools to Fit the Pearson Type III Distribution to + Aviation Loads Data", Office of Aviation Research (2003). + + %(example)s + """ + def _preprocess(self, x, skew): + # The real 'loc' and 'scale' are handled in the calling pdf(...). The + # local variables 'loc' and 'scale' within pearson3._pdf are set to + # the defaults just to keep them as part of the equations for + # documentation. + loc = 0.0 + scale = 1.0 + + # If skew is small, return _norm_pdf. The divide between pearson3 + # and norm was found by brute force and is approximately a skew of + # 0.000016. No one, I hope, would actually use a skew value even + # close to this small. + norm2pearson_transition = 0.000016 + + ans, x, skew = np.broadcast_arrays(1.0, x, skew) + ans = ans.copy() + + # mask is True where skew is small enough to use the normal approx. + mask = np.absolute(skew) < norm2pearson_transition + invmask = ~mask + + beta = 2.0 / (skew[invmask] * scale) + alpha = (scale * beta)**2 + zeta = loc - alpha / beta + + transx = beta * (x[invmask] - zeta) + return ans, x, transx, mask, invmask, beta, alpha, zeta + + def _argcheck(self, skew): + # The _argcheck function in rv_continuous only allows positive + # arguments. The skew argument for pearson3 can be zero (which I want + # to handle inside pearson3._pdf) or negative. So just return True + # for all skew args. + return np.isfinite(skew) + + def _shape_info(self): + return [_ShapeInfo("skew", False, (-np.inf, np.inf), (False, False))] + + def _stats(self, skew): + m = 0.0 + v = 1.0 + s = skew + k = 1.5*skew**2 + return m, v, s, k + + def _pdf(self, x, skew): + # pearson3.pdf(x, skew) = abs(beta) / gamma(alpha) * + # (beta * (x - zeta))**(alpha - 1) * exp(-beta*(x - zeta)) + # Do the calculation in _logpdf since helps to limit + # overflow/underflow problems + ans = np.exp(self._logpdf(x, skew)) + if ans.ndim == 0: + if np.isnan(ans): + return 0.0 + return ans + ans[np.isnan(ans)] = 0.0 + return ans + + def _logpdf(self, x, skew): + # PEARSON3 logpdf GAMMA logpdf + # np.log(abs(beta)) + # + (alpha - 1)*np.log(beta*(x - zeta)) + (a - 1)*np.log(x) + # - beta*(x - zeta) - x + # - sc.gammalnalpha) - sc.gammalna) + ans, x, transx, mask, invmask, beta, alpha, _ = ( + self._preprocess(x, skew)) + + ans[mask] = np.log(_norm_pdf(x[mask])) + # use logpdf instead of _logpdf to fix issue mentioned in gh-12640 + # (_logpdf does not return correct result for alpha = 1) + ans[invmask] = np.log(abs(beta)) + gamma.logpdf(transx, alpha) + return ans + + def _cdf(self, x, skew): + ans, x, transx, mask, invmask, _, alpha, _ = ( + self._preprocess(x, skew)) + + ans[mask] = _norm_cdf(x[mask]) + + skew = np.broadcast_to(skew, invmask.shape) + invmask1a = np.logical_and(invmask, skew > 0) + invmask1b = skew[invmask] > 0 + # use cdf instead of _cdf to fix issue mentioned in gh-12640 + # (_cdf produces NaNs for inputs outside support) + ans[invmask1a] = gamma.cdf(transx[invmask1b], alpha[invmask1b]) + + # The gamma._cdf approach wasn't working with negative skew. + # Note that multiplying the skew by -1 reflects about x=0. + # So instead of evaluating the CDF with negative skew at x, + # evaluate the SF with positive skew at -x. + invmask2a = np.logical_and(invmask, skew < 0) + invmask2b = skew[invmask] < 0 + # gamma._sf produces NaNs when transx < 0, so use gamma.sf + ans[invmask2a] = gamma.sf(transx[invmask2b], alpha[invmask2b]) + + return ans + + def _sf(self, x, skew): + ans, x, transx, mask, invmask, _, alpha, _ = ( + self._preprocess(x, skew)) + + ans[mask] = _norm_sf(x[mask]) + + skew = np.broadcast_to(skew, invmask.shape) + invmask1a = np.logical_and(invmask, skew > 0) + invmask1b = skew[invmask] > 0 + ans[invmask1a] = gamma.sf(transx[invmask1b], alpha[invmask1b]) + + invmask2a = np.logical_and(invmask, skew < 0) + invmask2b = skew[invmask] < 0 + ans[invmask2a] = gamma.cdf(transx[invmask2b], alpha[invmask2b]) + + return ans + + def _rvs(self, skew, size=None, random_state=None): + skew = np.broadcast_to(skew, size) + ans, _, _, mask, invmask, beta, alpha, zeta = ( + self._preprocess([0], skew)) + + nsmall = mask.sum() + nbig = mask.size - nsmall + ans[mask] = random_state.standard_normal(nsmall) + ans[invmask] = random_state.standard_gamma(alpha, nbig)/beta + zeta + + if size == (): + ans = ans[0] + return ans + + def _ppf(self, q, skew): + ans, q, _, mask, invmask, beta, alpha, zeta = ( + self._preprocess(q, skew)) + ans[mask] = _norm_ppf(q[mask]) + q = q[invmask] + q[beta < 0] = 1 - q[beta < 0] # for negative skew; see gh-17050 + ans[invmask] = sc.gammaincinv(alpha, q)/beta + zeta + return ans + + @_call_super_mom + @extend_notes_in_docstring(rv_continuous, notes="""\ + Note that method of moments (`method='MM'`) is not + available for this distribution.\n\n""") + def fit(self, data, *args, **kwds): + if kwds.get("method", None) == 'MM': + raise NotImplementedError("Fit `method='MM'` is not available for " + "the Pearson3 distribution. Please try " + "the default `method='MLE'`.") + else: + return super(type(self), self).fit(data, *args, **kwds) + + +pearson3 = pearson3_gen(name="pearson3") + + +class powerlaw_gen(rv_continuous): + r"""A power-function continuous random variable. + + %(before_notes)s + + See Also + -------- + pareto + + Notes + ----- + The probability density function for `powerlaw` is: + + .. math:: + + f(x, a) = a x^{a-1} + + for :math:`0 \le x \le 1`, :math:`a > 0`. + + `powerlaw` takes ``a`` as a shape parameter for :math:`a`. + + %(after_notes)s + + For example, the support of `powerlaw` can be adjusted from the default + interval ``[0, 1]`` to the interval ``[c, c+d]`` by setting ``loc=c`` and + ``scale=d``. For a power-law distribution with infinite support, see + `pareto`. For a power-law distribution described by PDF: + + .. math:: + + f(x; a, l, h) = \frac{a}{h^a - l^2} x^{a-1} + + with :math:`a \neq 0` and :math:`0 < l < x < h`, see `truncpareto`. + + `powerlaw` is a special case of `beta` with ``b=1``. + + %(example)s + + """ + def _shape_info(self): + return [_ShapeInfo("a", False, (0, np.inf), (False, False))] + + def _pdf(self, x, a): + # powerlaw.pdf(x, a) = a * x**(a-1) + return a*x**(a-1.0) + + def _logpdf(self, x, a): + return np.log(a) + sc.xlogy(a - 1, x) + + def _cdf(self, x, a): + return x**(a*1.0) + + def _logcdf(self, x, a): + return a*np.log(x) + + def _ppf(self, q, a): + return pow(q, 1.0/a) + + def _sf(self, p, a): + return -sc.powm1(p, a) + + def _munp(self, n, a): + # The following expression is correct for all real n (provided a > 0). + return a / (a + n) + + def _stats(self, a): + return (a / (a + 1.0), + a / (a + 2.0) / (a + 1.0) ** 2, + -2.0 * ((a - 1.0) / (a + 3.0)) * np.sqrt((a + 2.0) / a), + 6 * np.polyval([1, -1, -6, 2], a) / (a * (a + 3.0) * (a + 4))) + + def _entropy(self, a): + return 1 - 1.0/a - np.log(a) + + def _support_mask(self, x, a): + return (super()._support_mask(x, a) + & ((x != 0) | (a >= 1))) + + @_call_super_mom + @extend_notes_in_docstring(rv_continuous, notes="""\ + Notes specifically for ``powerlaw.fit``: If the location is a free + parameter and the value returned for the shape parameter is less than + one, the true maximum likelihood approaches infinity. This causes + numerical difficulties, and the resulting estimates are approximate. + \n\n""") + def fit(self, data, *args, **kwds): + # Summary of the strategy: + # + # 1) If the scale and location are fixed, return the shape according + # to a formula. + # + # 2) If the scale is fixed, there are two possibilities for the other + # parameters - one corresponding with shape less than one, and + # another with shape greater than one. Calculate both, and return + # whichever has the better log-likelihood. + # + # At this point, the scale is known to be free. + # + # 3) If the location is fixed, return the scale and shape according to + # formulas (or, if the shape is fixed, the fixed shape). + # + # At this point, the location and scale are both free. There are + # separate equations depending on whether the shape is less than one or + # greater than one. + # + # 4a) If the shape is less than one, there are formulas for shape, + # location, and scale. + # 4b) If the shape is greater than one, there are formulas for shape + # and scale, but there is a condition for location to be solved + # numerically. + # + # If the shape is fixed and less than one, we use 4a. + # If the shape is fixed and greater than one, we use 4b. + # If the shape is also free, we calculate fits using both 4a and 4b + # and choose the one that results a better log-likelihood. + # + # In many cases, the use of `np.nextafter` is used to avoid numerical + # issues. + if kwds.pop('superfit', False): + return super().fit(data, *args, **kwds) + + if len(np.unique(data)) == 1: + return super().fit(data, *args, **kwds) + + data, fshape, floc, fscale = _check_fit_input_parameters(self, data, + args, kwds) + penalized_nllf_args = [data, (self._fitstart(data),)] + penalized_nllf = self._reduce_func(penalized_nllf_args, {})[1] + + # ensure that any fixed parameters don't violate constraints of the + # distribution before continuing. The support of the distribution + # is `0 < (x - loc)/scale < 1`. + if floc is not None: + if not data.min() > floc: + raise FitDataError('powerlaw', 0, 1) + if fscale is not None and not data.max() <= floc + fscale: + raise FitDataError('powerlaw', 0, 1) + + if fscale is not None: + if fscale <= 0: + raise ValueError("Negative or zero `fscale` is outside the " + "range allowed by the distribution.") + if fscale <= np.ptp(data): + msg = "`fscale` must be greater than the range of data." + raise ValueError(msg) + + def get_shape(data, loc, scale): + # The first-order necessary condition on `shape` can be solved in + # closed form. It can be used no matter the assumption of the + # value of the shape. + N = len(data) + return - N / (np.sum(np.log(data - loc)) - N*np.log(scale)) + + def get_scale(data, loc): + # analytical solution for `scale` based on the location. + # It can be used no matter the assumption of the value of the + # shape. + return data.max() - loc + + # 1) The location and scale are both fixed. Analytically determine the + # shape. + if fscale is not None and floc is not None: + return get_shape(data, floc, fscale), floc, fscale + + # 2) The scale is fixed. There are two possibilities for the other + # parameters. Choose the option with better log-likelihood. + if fscale is not None: + # using `data.min()` as the optimal location + loc_lt1 = np.nextafter(data.min(), -np.inf) + shape_lt1 = fshape or get_shape(data, loc_lt1, fscale) + ll_lt1 = penalized_nllf((shape_lt1, loc_lt1, fscale), data) + + # using `data.max() - scale` as the optimal location + loc_gt1 = np.nextafter(data.max() - fscale, np.inf) + shape_gt1 = fshape or get_shape(data, loc_gt1, fscale) + ll_gt1 = penalized_nllf((shape_gt1, loc_gt1, fscale), data) + + if ll_lt1 < ll_gt1: + return shape_lt1, loc_lt1, fscale + else: + return shape_gt1, loc_gt1, fscale + + # 3) The location is fixed. Return the analytical scale and the + # analytical (or fixed) shape. + if floc is not None: + scale = get_scale(data, floc) + shape = fshape or get_shape(data, floc, scale) + return shape, floc, scale + + # 4) Location and scale are both free + # 4a) Use formulas that assume `shape <= 1`. + + def fit_loc_scale_w_shape_lt_1(): + loc = np.nextafter(data.min(), -np.inf) + if np.abs(loc) < np.finfo(loc.dtype).tiny: + loc = np.sign(loc) * np.finfo(loc.dtype).tiny + scale = np.nextafter(get_scale(data, loc), np.inf) + shape = fshape or get_shape(data, loc, scale) + return shape, loc, scale + + # 4b) Fit under the assumption that `shape > 1`. The support + # of the distribution is `(x - loc)/scale <= 1`. The method of Lagrange + # multipliers turns this constraint into the condition that + # dL_dScale - dL_dLocation must be zero, which is solved numerically. + # (Alternatively, substitute the constraint into the objective + # function before deriving the likelihood equation for location.) + + def dL_dScale(data, shape, scale): + # The partial derivative of the log-likelihood function w.r.t. + # the scale. + return -data.shape[0] * shape / scale + + def dL_dLocation(data, shape, loc): + # The partial derivative of the log-likelihood function w.r.t. + # the location. + return (shape - 1) * np.sum(1 / (loc - data)) # -1/(data-loc) + + def dL_dLocation_star(loc): + # The derivative of the log-likelihood function w.r.t. + # the location, given optimal shape and scale + scale = np.nextafter(get_scale(data, loc), -np.inf) + shape = fshape or get_shape(data, loc, scale) + return dL_dLocation(data, shape, loc) + + def fun_to_solve(loc): + # optimize the location by setting the partial derivatives + # w.r.t. location and scale equal and solving. + scale = np.nextafter(get_scale(data, loc), -np.inf) + shape = fshape or get_shape(data, loc, scale) + return (dL_dScale(data, shape, scale) + - dL_dLocation(data, shape, loc)) + + def fit_loc_scale_w_shape_gt_1(): + # set brackets for `root_scalar` to use when optimizing over the + # location such that a root is likely between them. + rbrack = np.nextafter(data.min(), -np.inf) + + # if the sign of `dL_dLocation_star` is positive at rbrack, + # we're not going to find the root we're looking for + delta = (data.min() - rbrack) + while dL_dLocation_star(rbrack) > 0: + rbrack = data.min() - delta + delta *= 2 + + def interval_contains_root(lbrack, rbrack): + # Check if the interval (lbrack, rbrack) contains the root. + return (np.sign(fun_to_solve(lbrack)) + != np.sign(fun_to_solve(rbrack))) + + lbrack = rbrack - 1 + + # if the sign doesn't change between the brackets, move the left + # bracket until it does. (The right bracket remains fixed at the + # maximum permissible value.) + i = 1.0 + while (not interval_contains_root(lbrack, rbrack) + and lbrack != -np.inf): + lbrack = (data.min() - i) + i *= 2 + + root = optimize.root_scalar(fun_to_solve, bracket=(lbrack, rbrack)) + + loc = np.nextafter(root.root, -np.inf) + scale = np.nextafter(get_scale(data, loc), np.inf) + shape = fshape or get_shape(data, loc, scale) + return shape, loc, scale + + # Shape is fixed - choose 4a or 4b accordingly. + if fshape is not None and fshape <= 1: + return fit_loc_scale_w_shape_lt_1() + elif fshape is not None and fshape > 1: + return fit_loc_scale_w_shape_gt_1() + + # Shape is free + fit_shape_lt1 = fit_loc_scale_w_shape_lt_1() + ll_lt1 = self.nnlf(fit_shape_lt1, data) + + fit_shape_gt1 = fit_loc_scale_w_shape_gt_1() + ll_gt1 = self.nnlf(fit_shape_gt1, data) + + if ll_lt1 <= ll_gt1 and fit_shape_lt1[0] <= 1: + return fit_shape_lt1 + elif ll_lt1 > ll_gt1 and fit_shape_gt1[0] > 1: + return fit_shape_gt1 + else: + return super().fit(data, *args, **kwds) + + +powerlaw = powerlaw_gen(a=0.0, b=1.0, name="powerlaw") + + +class powerlognorm_gen(rv_continuous): + r"""A power log-normal continuous random variable. + + %(before_notes)s + + Notes + ----- + The probability density function for `powerlognorm` is: + + .. math:: + + f(x, c, s) = \frac{c}{x s} \phi(\log(x)/s) + (\Phi(-\log(x)/s))^{c-1} + + where :math:`\phi` is the normal pdf, and :math:`\Phi` is the normal cdf, + and :math:`x > 0`, :math:`s, c > 0`. + + `powerlognorm` takes :math:`c` and :math:`s` as shape parameters. + + %(after_notes)s + + %(example)s + + """ + _support_mask = rv_continuous._open_support_mask + + def _shape_info(self): + ic = _ShapeInfo("c", False, (0, np.inf), (False, False)) + i_s = _ShapeInfo("s", False, (0, np.inf), (False, False)) + return [ic, i_s] + + def _pdf(self, x, c, s): + return np.exp(self._logpdf(x, c, s)) + + def _logpdf(self, x, c, s): + return (np.log(c) - np.log(x) - np.log(s) + + _norm_logpdf(np.log(x) / s) + + _norm_logcdf(-np.log(x) / s) * (c - 1.)) + + def _cdf(self, x, c, s): + return -sc.expm1(self._logsf(x, c, s)) + + def _ppf(self, q, c, s): + return self._isf(1 - q, c, s) + + def _sf(self, x, c, s): + return np.exp(self._logsf(x, c, s)) + + def _logsf(self, x, c, s): + return _norm_logcdf(-np.log(x) / s) * c + + def _isf(self, q, c, s): + return np.exp(-_norm_ppf(q**(1/c)) * s) + + +powerlognorm = powerlognorm_gen(a=0.0, name="powerlognorm") + + +class powernorm_gen(rv_continuous): + r"""A power normal continuous random variable. + + %(before_notes)s + + Notes + ----- + The probability density function for `powernorm` is: + + .. math:: + + f(x, c) = c \phi(x) (\Phi(-x))^{c-1} + + where :math:`\phi` is the normal pdf, :math:`\Phi` is the normal cdf, + :math:`x` is any real, and :math:`c > 0` [1]_. + + `powernorm` takes ``c`` as a shape parameter for :math:`c`. + + %(after_notes)s + + References + ---------- + .. [1] NIST Engineering Statistics Handbook, Section 1.3.6.6.13, + https://www.itl.nist.gov/div898/handbook//eda/section3/eda366d.htm + + %(example)s + + """ + def _shape_info(self): + return [_ShapeInfo("c", False, (0, np.inf), (False, False))] + + def _pdf(self, x, c): + # powernorm.pdf(x, c) = c * phi(x) * (Phi(-x))**(c-1) + return c*_norm_pdf(x) * (_norm_cdf(-x)**(c-1.0)) + + def _logpdf(self, x, c): + return np.log(c) + _norm_logpdf(x) + (c-1)*_norm_logcdf(-x) + + def _cdf(self, x, c): + return -sc.expm1(self._logsf(x, c)) + + def _ppf(self, q, c): + return -_norm_ppf(pow(1.0 - q, 1.0 / c)) + + def _sf(self, x, c): + return np.exp(self._logsf(x, c)) + + def _logsf(self, x, c): + return c * _norm_logcdf(-x) + + def _isf(self, q, c): + return -_norm_ppf(np.exp(np.log(q) / c)) + + +powernorm = powernorm_gen(name='powernorm') + + +class rdist_gen(rv_continuous): + r"""An R-distributed (symmetric beta) continuous random variable. + + %(before_notes)s + + Notes + ----- + The probability density function for `rdist` is: + + .. math:: + + f(x, c) = \frac{(1-x^2)^{c/2-1}}{B(1/2, c/2)} + + for :math:`-1 \le x \le 1`, :math:`c > 0`. `rdist` is also called the + symmetric beta distribution: if B has a `beta` distribution with + parameters (c/2, c/2), then X = 2*B - 1 follows a R-distribution with + parameter c. + + `rdist` takes ``c`` as a shape parameter for :math:`c`. + + This distribution includes the following distribution kernels as + special cases:: + + c = 2: uniform + c = 3: `semicircular` + c = 4: Epanechnikov (parabolic) + c = 6: quartic (biweight) + c = 8: triweight + + %(after_notes)s + + %(example)s + + """ + def _shape_info(self): + return [_ShapeInfo("c", False, (0, np.inf), (False, False))] + + # use relation to the beta distribution for pdf, cdf, etc + def _pdf(self, x, c): + return np.exp(self._logpdf(x, c)) + + def _logpdf(self, x, c): + return -np.log(2) + beta._logpdf((x + 1)/2, c/2, c/2) + + def _cdf(self, x, c): + return beta._cdf((x + 1)/2, c/2, c/2) + + def _sf(self, x, c): + return beta._sf((x + 1)/2, c/2, c/2) + + def _ppf(self, q, c): + return 2*beta._ppf(q, c/2, c/2) - 1 + + def _rvs(self, c, size=None, random_state=None): + return 2 * random_state.beta(c/2, c/2, size) - 1 + + def _munp(self, n, c): + numerator = (1 - (n % 2)) * sc.beta((n + 1.0) / 2, c / 2.0) + return numerator / sc.beta(1. / 2, c / 2.) + + +rdist = rdist_gen(a=-1.0, b=1.0, name="rdist") + + +class rayleigh_gen(rv_continuous): + r"""A Rayleigh continuous random variable. + + %(before_notes)s + + Notes + ----- + The probability density function for `rayleigh` is: + + .. math:: + + f(x) = x \exp(-x^2/2) + + for :math:`x \ge 0`. + + `rayleigh` is a special case of `chi` with ``df=2``. + + %(after_notes)s + + %(example)s + + """ + _support_mask = rv_continuous._open_support_mask + + def _shape_info(self): + return [] + + def _rvs(self, size=None, random_state=None): + return chi.rvs(2, size=size, random_state=random_state) + + def _pdf(self, r): + # rayleigh.pdf(r) = r * exp(-r**2/2) + return np.exp(self._logpdf(r)) + + def _logpdf(self, r): + return np.log(r) - 0.5 * r * r + + def _cdf(self, r): + return -sc.expm1(-0.5 * r**2) + + def _ppf(self, q): + return np.sqrt(-2 * sc.log1p(-q)) + + def _sf(self, r): + return np.exp(self._logsf(r)) + + def _logsf(self, r): + return -0.5 * r * r + + def _isf(self, q): + return np.sqrt(-2 * np.log(q)) + + def _stats(self): + val = 4 - np.pi + return (np.sqrt(np.pi/2), + val/2, + 2*(np.pi-3)*np.sqrt(np.pi)/val**1.5, + 6*np.pi/val-16/val**2) + + def _entropy(self): + return _EULER/2.0 + 1 - 0.5*np.log(2) + + @_call_super_mom + @extend_notes_in_docstring(rv_continuous, notes="""\ + Notes specifically for ``rayleigh.fit``: If the location is fixed with + the `floc` parameter, this method uses an analytical formula to find + the scale. Otherwise, this function uses a numerical root finder on + the first order conditions of the log-likelihood function to find the + MLE. Only the (optional) `loc` parameter is used as the initial guess + for the root finder; the `scale` parameter and any other parameters + for the optimizer are ignored.\n\n""") + def fit(self, data, *args, **kwds): + if kwds.pop('superfit', False): + return super().fit(data, *args, **kwds) + data, floc, fscale = _check_fit_input_parameters(self, data, + args, kwds) + + def scale_mle(loc): + # Source: Statistical Distributions, 3rd Edition. Evans, Hastings, + # and Peacock (2000), Page 175 + return (np.sum((data - loc) ** 2) / (2 * len(data))) ** .5 + + def loc_mle(loc): + # This implicit equation for `loc` is used when + # both `loc` and `scale` are free. + xm = data - loc + s1 = xm.sum() + s2 = (xm**2).sum() + s3 = (1/xm).sum() + return s1 - s2/(2*len(data))*s3 + + def loc_mle_scale_fixed(loc, scale=fscale): + # This implicit equation for `loc` is used when + # `scale` is fixed but `loc` is not. + xm = data - loc + return xm.sum() - scale**2 * (1/xm).sum() + + if floc is not None: + # `loc` is fixed, analytically determine `scale`. + if np.any(data - floc <= 0): + raise FitDataError("rayleigh", lower=1, upper=np.inf) + else: + return floc, scale_mle(floc) + + # Account for user provided guess of `loc`. + loc0 = kwds.get('loc') + if loc0 is None: + # Use _fitstart to estimate loc; ignore the returned scale. + loc0 = self._fitstart(data)[0] + + fun = loc_mle if fscale is None else loc_mle_scale_fixed + rbrack = np.nextafter(np.min(data), -np.inf) + lbrack = _get_left_bracket(fun, rbrack) + res = optimize.root_scalar(fun, bracket=(lbrack, rbrack)) + if not res.converged: + raise FitSolverError(res.flag) + loc = res.root + scale = fscale or scale_mle(loc) + return loc, scale + + +rayleigh = rayleigh_gen(a=0.0, name="rayleigh") + + +class reciprocal_gen(rv_continuous): + r"""A loguniform or reciprocal continuous random variable. + + %(before_notes)s + + Notes + ----- + The probability density function for this class is: + + .. math:: + + f(x, a, b) = \frac{1}{x \log(b/a)} + + for :math:`a \le x \le b`, :math:`b > a > 0`. This class takes + :math:`a` and :math:`b` as shape parameters. + + %(after_notes)s + + %(example)s + + This doesn't show the equal probability of ``0.01``, ``0.1`` and + ``1``. This is best when the x-axis is log-scaled: + + >>> import numpy as np + >>> import matplotlib.pyplot as plt + >>> fig, ax = plt.subplots(1, 1) + >>> ax.hist(np.log10(r)) + >>> ax.set_ylabel("Frequency") + >>> ax.set_xlabel("Value of random variable") + >>> ax.xaxis.set_major_locator(plt.FixedLocator([-2, -1, 0])) + >>> ticks = ["$10^{{ {} }}$".format(i) for i in [-2, -1, 0]] + >>> ax.set_xticklabels(ticks) # doctest: +SKIP + >>> plt.show() + + This random variable will be log-uniform regardless of the base chosen for + ``a`` and ``b``. Let's specify with base ``2`` instead: + + >>> rvs = %(name)s(2**-2, 2**0).rvs(size=1000) + + Values of ``1/4``, ``1/2`` and ``1`` are equally likely with this random + variable. Here's the histogram: + + >>> fig, ax = plt.subplots(1, 1) + >>> ax.hist(np.log2(rvs)) + >>> ax.set_ylabel("Frequency") + >>> ax.set_xlabel("Value of random variable") + >>> ax.xaxis.set_major_locator(plt.FixedLocator([-2, -1, 0])) + >>> ticks = ["$2^{{ {} }}$".format(i) for i in [-2, -1, 0]] + >>> ax.set_xticklabels(ticks) # doctest: +SKIP + >>> plt.show() + + """ + def _argcheck(self, a, b): + return (a > 0) & (b > a) + + def _shape_info(self): + ia = _ShapeInfo("a", False, (0, np.inf), (False, False)) + ib = _ShapeInfo("b", False, (0, np.inf), (False, False)) + return [ia, ib] + + def _fitstart(self, data): + if isinstance(data, CensoredData): + data = data._uncensor() + # Reasonable, since support is [a, b] + return super()._fitstart(data, args=(np.min(data), np.max(data))) + + def _get_support(self, a, b): + return a, b + + def _pdf(self, x, a, b): + # reciprocal.pdf(x, a, b) = 1 / (x*(log(b) - log(a))) + return np.exp(self._logpdf(x, a, b)) + + def _logpdf(self, x, a, b): + return -np.log(x) - np.log(np.log(b) - np.log(a)) + + def _cdf(self, x, a, b): + return (np.log(x)-np.log(a)) / (np.log(b) - np.log(a)) + + def _ppf(self, q, a, b): + return np.exp(np.log(a) + q*(np.log(b) - np.log(a))) + + def _munp(self, n, a, b): + if n == 0: + return 1.0 + t1 = 1 / (np.log(b) - np.log(a)) / n + t2 = np.real(np.exp(_log_diff(n * np.log(b), n*np.log(a)))) + return t1 * t2 + + def _entropy(self, a, b): + return 0.5*(np.log(a) + np.log(b)) + np.log(np.log(b) - np.log(a)) + + fit_note = """\ + `loguniform`/`reciprocal` is over-parameterized. `fit` automatically + fixes `scale` to 1 unless `fscale` is provided by the user.\n\n""" + + @extend_notes_in_docstring(rv_continuous, notes=fit_note) + def fit(self, data, *args, **kwds): + fscale = kwds.pop('fscale', 1) + return super().fit(data, *args, fscale=fscale, **kwds) + + # Details related to the decision of not defining + # the survival function for this distribution can be + # found in the PR: https://github.com/scipy/scipy/pull/18614 + + +loguniform = reciprocal_gen(name="loguniform") +reciprocal = reciprocal_gen(name="reciprocal") +loguniform._support = ('a', 'b') +reciprocal._support = ('a', 'b') + + +class rice_gen(rv_continuous): + r"""A Rice continuous random variable. + + %(before_notes)s + + Notes + ----- + The probability density function for `rice` is: + + .. math:: + + f(x, b) = x \exp(- \frac{x^2 + b^2}{2}) I_0(x b) + + for :math:`x >= 0`, :math:`b > 0`. :math:`I_0` is the modified Bessel + function of order zero (`scipy.special.i0`). + + `rice` takes ``b`` as a shape parameter for :math:`b`. + + %(after_notes)s + + The Rice distribution describes the length, :math:`r`, of a 2-D vector with + components :math:`(U+u, V+v)`, where :math:`U, V` are constant, :math:`u, + v` are independent Gaussian random variables with standard deviation + :math:`s`. Let :math:`R = \sqrt{U^2 + V^2}`. Then the pdf of :math:`r` is + ``rice.pdf(x, R/s, scale=s)``. + + %(example)s + + """ + def _argcheck(self, b): + return b >= 0 + + def _shape_info(self): + return [_ShapeInfo("b", False, (0, np.inf), (True, False))] + + def _rvs(self, b, size=None, random_state=None): + # https://en.wikipedia.org/wiki/Rice_distribution + t = b/np.sqrt(2) + random_state.standard_normal(size=(2,) + size) + return np.sqrt((t*t).sum(axis=0)) + + def _cdf(self, x, b): + return sc.chndtr(np.square(x), 2, np.square(b)) + + def _ppf(self, q, b): + return np.sqrt(sc.chndtrix(q, 2, np.square(b))) + + def _pdf(self, x, b): + # rice.pdf(x, b) = x * exp(-(x**2+b**2)/2) * I[0](x*b) + # + # We use (x**2 + b**2)/2 = ((x-b)**2)/2 + xb. + # The factor of np.exp(-xb) is then included in the i0e function + # in place of the modified Bessel function, i0, improving + # numerical stability for large values of xb. + return x * np.exp(-(x-b)*(x-b)/2.0) * sc.i0e(x*b) + + def _munp(self, n, b): + nd2 = n/2.0 + n1 = 1 + nd2 + b2 = b*b/2.0 + return (2.0**(nd2) * np.exp(-b2) * sc.gamma(n1) * + sc.hyp1f1(n1, 1, b2)) + + +rice = rice_gen(a=0.0, name="rice") + +class irwinhall_gen(rv_continuous): + r"""An Irwin-Hall (Uniform Sum) continuous random variable. + + An `Irwin-Hall `_ + continuous random variable is the sum of :math:`n` independent + standard uniform random variables [1]_ [2]_. + + %(before_notes)s + + Notes + ----- + Applications include `Rao's Spacing Test + `_, + a more powerful alternative to the Rayleigh test + when the data are not unimodal, and radar [3]_. + + Conveniently, the pdf and cdf are the :math:`n`-fold convolution of + the ones for the standard uniform distribution, which is also the + definition of the cardinal B-splines of degree :math:`n-1` + having knots evenly spaced from :math:`1` to :math:`n` [4]_ [5]_. + + The Bates distribution, which represents the *mean* of statistically + independent, uniformly distributed random variables, is simply the + Irwin-Hall distribution scaled by :math:`1/n`. For example, the frozen + distribution ``bates = irwinhall(10, scale=1/10)`` represents the + distribution of the mean of 10 uniformly distributed random variables. + + %(after_notes)s + + References + ---------- + .. [1] P. Hall, "The distribution of means for samples of size N drawn + from a population in which the variate takes values between 0 and 1, + all such values being equally probable", + Biometrika, Volume 19, Issue 3-4, December 1927, Pages 240-244, + :doi:`10.1093/biomet/19.3-4.240`. + .. [2] J. O. Irwin, "On the frequency distribution of the means of samples + from a population having any law of frequency with finite moments, + with special reference to Pearson's Type II, + Biometrika, Volume 19, Issue 3-4, December 1927, Pages 225-239, + :doi:`0.1093/biomet/19.3-4.225`. + .. [3] K. Buchanan, T. Adeyemi, C. Flores-Molina, S. Wheeland and D. Overturf, + "Sidelobe behavior and bandwidth characteristics + of distributed antenna arrays," + 2018 United States National Committee of + URSI National Radio Science Meeting (USNC-URSI NRSM), + Boulder, CO, USA, 2018, pp. 1-2. + https://www.usnc-ursi-archive.org/nrsm/2018/papers/B15-9.pdf. + .. [4] Amos Ron, "Lecture 1: Cardinal B-splines and convolution operators", p. 1 + https://pages.cs.wisc.edu/~deboor/887/lec1new.pdf. + .. [5] Trefethen, N. (2012, July). B-splines and convolution. Chebfun. + Retrieved April 30, 2024, from http://www.chebfun.org/examples/approx/BSplineConv.html. + + %(example)s + """ # noqa: E501 + + @replace_notes_in_docstring(rv_continuous, notes="""\ + Raises a ``NotImplementedError`` for the Irwin-Hall distribution because + the generic `fit` implementation is unreliable and no custom implementation + is available. Consider using `scipy.stats.fit`.\n\n""") + def fit(self, data, *args, **kwds): + fit_notes = ("The generic `fit` implementation is unreliable for this " + "distribution, and no custom implementation is available. " + "Consider using `scipy.stats.fit`.") + raise NotImplementedError(fit_notes) + + def _argcheck(self, n): + return (n > 0) & _isintegral(n) & np.isrealobj(n) + + def _get_support(self, n): + return 0, n + + def _shape_info(self): + return [_ShapeInfo("n", True, (1, np.inf), (True, False))] + + def _munp(self, order, n): + # see https://link.springer.com/content/pdf/10.1007/s10959-020-01050-9.pdf + # page 640, with m=n, j=n+order + def vmunp(order, n): + n = np.asarray(n, dtype=np.int64) + return (sc.stirling2(n+order, n, exact=True) + / sc.comb(n+order, n, exact=True)) + + # exact rationals, but we convert to float anyway + return np.vectorize(vmunp, otypes=[np.float64])(order, n) + + @staticmethod + def _cardbspl(n): + t = np.arange(n+1) + return BSpline.basis_element(t) + + def _pdf(self, x, n): + def vpdf(x, n): + return self._cardbspl(n)(x) + return np.vectorize(vpdf, otypes=[np.float64])(x, n) + + def _cdf(self, x, n): + def vcdf(x, n): + return self._cardbspl(n).antiderivative()(x) + return np.vectorize(vcdf, otypes=[np.float64])(x, n) + + def _sf(self, x, n): + def vsf(x, n): + return self._cardbspl(n).antiderivative()(n-x) + return np.vectorize(vsf, otypes=[np.float64])(x, n) + + def _rvs(self, n, size=None, random_state=None, *args): + @_vectorize_rvs_over_shapes + def _rvs1(n, size=None, random_state=None): + n = np.floor(n).astype(int) + usize = (n,) if size is None else (n, *size) + return random_state.uniform(size=usize).sum(axis=0) + return _rvs1(n, size=size, random_state=random_state) + + def _stats(self, n): + # mgf = ((exp(t) - 1)/t)**n + # m'th derivative follows from the generalized Leibniz rule + # Moments follow directly from the definition as the sum of n iid unif(0,1) + # and the summation rules for moments of a sum of iid random variables + # E(IH((n))) = n*E(U(0,1)) = n/2 + # Var(IH((n))) = n*Var(U(0,1)) = n/12 + # Skew(IH((n))) = Skew(U(0,1))/sqrt(n) = 0 + # Kurt(IH((n))) = Kurt(U(0,1))/n = -6/(5*n) -- Fisher's excess kurtosis + # See e.g. https://en.wikipedia.org/wiki/Irwin%E2%80%93Hall_distribution + + return n/2, n/12, 0, -6/(5*n) + + +irwinhall = irwinhall_gen(name="irwinhall") +irwinhall._support = (0.0, 'n') + + +class recipinvgauss_gen(rv_continuous): + r"""A reciprocal inverse Gaussian continuous random variable. + + %(before_notes)s + + Notes + ----- + The probability density function for `recipinvgauss` is: + + .. math:: + + f(x, \mu) = \frac{1}{\sqrt{2\pi x}} + \exp\left(\frac{-(1-\mu x)^2}{2\mu^2x}\right) + + for :math:`x \ge 0`. + + `recipinvgauss` takes ``mu`` as a shape parameter for :math:`\mu`. + + %(after_notes)s + + %(example)s + + """ + def _shape_info(self): + return [_ShapeInfo("mu", False, (0, np.inf), (False, False))] + + def _pdf(self, x, mu): + # recipinvgauss.pdf(x, mu) = + # 1/sqrt(2*pi*x) * exp(-(1-mu*x)**2/(2*x*mu**2)) + return np.exp(self._logpdf(x, mu)) + + def _logpdf(self, x, mu): + return xpx.apply_where( + x > 0, (x, mu), + lambda x, mu: (-(1 - mu*x)**2.0 / (2*x*mu**2.0) + - 0.5*np.log(2*np.pi*x)), + fill_value=-np.inf) + + def _cdf(self, x, mu): + trm1 = 1.0/mu - x + trm2 = 1.0/mu + x + isqx = 1.0/np.sqrt(x) + return _norm_cdf(-isqx*trm1) - np.exp(2.0/mu)*_norm_cdf(-isqx*trm2) + + def _sf(self, x, mu): + trm1 = 1.0/mu - x + trm2 = 1.0/mu + x + isqx = 1.0/np.sqrt(x) + return _norm_cdf(isqx*trm1) + np.exp(2.0/mu)*_norm_cdf(-isqx*trm2) + + def _rvs(self, mu, size=None, random_state=None): + return 1.0/random_state.wald(mu, 1.0, size=size) + + +recipinvgauss = recipinvgauss_gen(a=0.0, name='recipinvgauss') + + +class semicircular_gen(rv_continuous): + r"""A semicircular continuous random variable. + + %(before_notes)s + + See Also + -------- + rdist + + Notes + ----- + The probability density function for `semicircular` is: + + .. math:: + + f(x) = \frac{2}{\pi} \sqrt{1-x^2} + + for :math:`-1 \le x \le 1`. + + The distribution is a special case of `rdist` with ``c = 3``. + + %(after_notes)s + + References + ---------- + .. [1] "Wigner semicircle distribution", + https://en.wikipedia.org/wiki/Wigner_semicircle_distribution + + %(example)s + + """ + def _shape_info(self): + return [] + + def _pdf(self, x): + return 2.0/np.pi*np.sqrt(1-x*x) + + def _logpdf(self, x): + return np.log(2/np.pi) + 0.5*sc.log1p(-x*x) + + def _cdf(self, x): + return 0.5+1.0/np.pi*(x*np.sqrt(1-x*x) + np.arcsin(x)) + + def _ppf(self, q): + return rdist._ppf(q, 3) + + def _rvs(self, size=None, random_state=None): + # generate values uniformly distributed on the area under the pdf + # (semi-circle) by randomly generating the radius and angle + r = np.sqrt(random_state.uniform(size=size)) + a = np.cos(np.pi * random_state.uniform(size=size)) + return r * a + + def _stats(self): + return 0, 0.25, 0, -1.0 + + def _entropy(self): + return 0.64472988584940017414 + + +semicircular = semicircular_gen(a=-1.0, b=1.0, name="semicircular") + + +class skewcauchy_gen(rv_continuous): + r"""A skewed Cauchy random variable. + + %(before_notes)s + + See Also + -------- + cauchy : Cauchy distribution + + Notes + ----- + + The probability density function for `skewcauchy` is: + + .. math:: + + f(x) = \frac{1}{\pi \left(\frac{x^2}{\left(a\, \text{sign}(x) + 1 + \right)^2} + 1 \right)} + + for a real number :math:`x` and skewness parameter :math:`-1 < a < 1`. + + When :math:`a=0`, the distribution reduces to the usual Cauchy + distribution. + + %(after_notes)s + + References + ---------- + .. [1] "Skewed generalized *t* distribution", Wikipedia + https://en.wikipedia.org/wiki/Skewed_generalized_t_distribution#Skewed_Cauchy_distribution + + %(example)s + + """ + def _argcheck(self, a): + return np.abs(a) < 1 + + def _shape_info(self): + return [_ShapeInfo("a", False, (-1.0, 1.0), (False, False))] + + def _pdf(self, x, a): + return 1 / (np.pi * (x**2 / (a * np.sign(x) + 1)**2 + 1)) + + def _cdf(self, x, a): + return np.where(x <= 0, + (1 - a) / 2 + (1 - a) / np.pi * np.arctan(x / (1 - a)), + (1 - a) / 2 + (1 + a) / np.pi * np.arctan(x / (1 + a))) + + def _ppf(self, x, a): + i = x < self._cdf(0, a) + return np.where(i, + np.tan(np.pi / (1 - a) * (x - (1 - a) / 2)) * (1 - a), + np.tan(np.pi / (1 + a) * (x - (1 - a) / 2)) * (1 + a)) + + def _stats(self, a, moments='mvsk'): + return np.nan, np.nan, np.nan, np.nan + + def _fitstart(self, data): + # Use 0 as the initial guess of the skewness shape parameter. + # For the location and scale, estimate using the median and + # quartiles. + if isinstance(data, CensoredData): + data = data._uncensor() + p25, p50, p75 = np.percentile(data, [25, 50, 75]) + return 0.0, p50, (p75 - p25)/2 + + +skewcauchy = skewcauchy_gen(name='skewcauchy') + + +class skewnorm_gen(rv_continuous): + r"""A skew-normal random variable. + + %(before_notes)s + + Notes + ----- + The pdf is:: + + skewnorm.pdf(x, a) = 2 * norm.pdf(x) * norm.cdf(a*x) + + `skewnorm` takes a real number :math:`a` as a skewness parameter + When ``a = 0`` the distribution is identical to a normal distribution + (`norm`). `rvs` implements the method of [1]_. + + This distribution uses routines from the Boost Math C++ library for + the computation of ``cdf``, ``ppf`` and ``isf`` methods. [2]_ + + %(after_notes)s + + References + ---------- + .. [1] A. Azzalini and A. Capitanio (1999). Statistical applications of + the multivariate skew-normal distribution. J. Roy. Statist. Soc., + B 61, 579-602. :arxiv:`0911.2093` + .. [2] The Boost Developers. "Boost C++ Libraries". https://www.boost.org/. + + %(example)s + + """ + def _argcheck(self, a): + return np.isfinite(a) + + def _shape_info(self): + return [_ShapeInfo("a", False, (-np.inf, np.inf), (False, False))] + + def _pdf(self, x, a): + return xpx.apply_where( + a == 0, (x, a), + lambda x, a: _norm_pdf(x), + lambda x, a: 2.*_norm_pdf(x)*_norm_cdf(a*x)) + + def _logpdf(self, x, a): + return xpx.apply_where( + a == 0, (x, a), + lambda x, a: _norm_logpdf(x), + lambda x, a: np.log(2)+_norm_logpdf(x)+_norm_logcdf(a*x)) + + def _cdf(self, x, a): + a = np.atleast_1d(a) + cdf = scu._skewnorm_cdf(x, 0.0, 1.0, a) + # for some reason, a isn't broadcasted if some of x are invalid + a = np.broadcast_to(a, cdf.shape) + # Boost is not accurate in left tail when a > 0 + i_small_cdf = (cdf < 1e-6) & (a > 0) + cdf[i_small_cdf] = super()._cdf(x[i_small_cdf], a[i_small_cdf]) + return np.clip(cdf, 0, 1) + + def _ppf(self, x, a): + return scu._skewnorm_ppf(x, 0.0, 1.0, a) + + def _sf(self, x, a): + # Boost's SF is implemented this way. Use whatever customizations + # we made in the _cdf. + return self._cdf(-x, -a) + + def _isf(self, x, a): + return scu._skewnorm_isf(x, 0.0, 1.0, a) + + def _rvs(self, a, size=None, random_state=None): + u0 = random_state.normal(size=size) + v = random_state.normal(size=size) + d = a/np.sqrt(1 + a**2) + u1 = d*u0 + v*np.sqrt(1 - d**2) + return np.where(u0 >= 0, u1, -u1) + + def _stats(self, a, moments='mvsk'): + output = [None, None, None, None] + const = np.sqrt(2/np.pi) * a/np.sqrt(1 + a**2) + + if 'm' in moments: + output[0] = const + if 'v' in moments: + output[1] = 1 - const**2 + if 's' in moments: + output[2] = ((4 - np.pi)/2) * (const/np.sqrt(1 - const**2))**3 + if 'k' in moments: + output[3] = (2*(np.pi - 3)) * (const**4/(1 - const**2)**2) + + return output + + # For odd order, the each noncentral moment of the skew-normal distribution + # with location 0 and scale 1 can be expressed as a polynomial in delta, + # where delta = a/sqrt(1 + a**2) and `a` is the skew-normal shape + # parameter. The dictionary _skewnorm_odd_moments defines those + # polynomials for orders up to 19. The dict is implemented as a cached + # property to reduce the impact of the creation of the dict on import time. + @cached_property + def _skewnorm_odd_moments(self): + skewnorm_odd_moments = { + 1: Polynomial([1]), + 3: Polynomial([3, -1]), + 5: Polynomial([15, -10, 3]), + 7: Polynomial([105, -105, 63, -15]), + 9: Polynomial([945, -1260, 1134, -540, 105]), + 11: Polynomial([10395, -17325, 20790, -14850, 5775, -945]), + 13: Polynomial([135135, -270270, 405405, -386100, 225225, -73710, + 10395]), + 15: Polynomial([2027025, -4729725, 8513505, -10135125, 7882875, + -3869775, 1091475, -135135]), + 17: Polynomial([34459425, -91891800, 192972780, -275675400, + 268017750, -175429800, 74220300, -18378360, + 2027025]), + 19: Polynomial([654729075, -1964187225, 4714049340, -7856748900, + 9166207050, -7499623950, 4230557100, -1571349780, + 346621275, -34459425]), + } + return skewnorm_odd_moments + + def _munp(self, order, a): + if order % 2: + if order > 19: + raise NotImplementedError("skewnorm noncentral moments not " + "implemented for odd orders greater " + "than 19.") + # Use the precomputed polynomials that were derived from the + # moment generating function. + delta = a/np.sqrt(1 + a**2) + return (delta * self._skewnorm_odd_moments[order](delta**2) + * _SQRT_2_OVER_PI) + else: + # For even order, the moment is just (order-1)!!, where !! is the + # notation for the double factorial; for an odd integer m, m!! is + # m*(m-2)*...*3*1. + # We could use special.factorial2, but we know the argument is odd, + # so avoid the overhead of that function and compute the result + # directly here. + return sc.gamma((order + 1)/2) * 2**(order/2) / _SQRT_PI + + @extend_notes_in_docstring(rv_continuous, notes="""\ + If ``method='mm'``, parameters fixed by the user are respected, and the + remaining parameters are used to match distribution and sample moments + where possible. For example, if the user fixes the location with + ``floc``, the parameters will only match the distribution skewness and + variance to the sample skewness and variance; no attempt will be made + to match the means or minimize a norm of the errors. + Note that the maximum possible skewness magnitude of a + `scipy.stats.skewnorm` distribution is approximately 0.9952717; if the + magnitude of the data's sample skewness exceeds this, the returned + shape parameter ``a`` will be infinite. + \n\n""") + def fit(self, data, *args, **kwds): + if kwds.pop("superfit", False): + return super().fit(data, *args, **kwds) + if isinstance(data, CensoredData): + if data.num_censored() == 0: + data = data._uncensor() + else: + return super().fit(data, *args, **kwds) + + # this extracts fixed shape, location, and scale however they + # are specified, and also leaves them in `kwds` + data, fa, floc, fscale = _check_fit_input_parameters(self, data, + args, kwds) + method = kwds.get("method", "mle").lower() + + # See https://en.wikipedia.org/wiki/Skew_normal_distribution for + # moment formulas. + def skew_d(d): # skewness in terms of delta + return (4-np.pi)/2 * ((d * np.sqrt(2 / np.pi))**3 + / (1 - 2*d**2 / np.pi)**(3/2)) + + def d_skew(skew): # delta in terms of skewness + s_23 = np.abs(skew)**(2/3) + return np.sign(skew) * np.sqrt( + np.pi/2 * s_23 / (s_23 + ((4 - np.pi)/2)**(2/3)) + ) + + # If method is method of moments, we don't need the user's guesses. + # Otherwise, extract the guesses from args and kwds. + if method == "mm": + a, loc, scale = None, None, None + else: + a = args[0] if len(args) else None + loc = kwds.pop('loc', None) + scale = kwds.pop('scale', None) + + if fa is None and a is None: # not fixed and no guess: use MoM + # Solve for a that matches sample distribution skewness to sample + # skewness. + s = stats.skew(data) + if method == 'mle': + # For MLE initial conditions, clip skewness to a large but + # reasonable value in case the data skewness is out-of-range. + s = np.clip(s, -0.99, 0.99) + else: + s_max = skew_d(1) + s = np.clip(s, -s_max, s_max) + d = d_skew(s) + with np.errstate(divide='ignore'): + a = np.sqrt(np.divide(d**2, (1-d**2)))*np.sign(s) + else: + a = fa if fa is not None else a + d = a / np.sqrt(1 + a**2) + + if fscale is None and scale is None: + v = np.var(data) + scale = np.sqrt(v / (1 - 2*d**2/np.pi)) + elif fscale is not None: + scale = fscale + + if floc is None and loc is None: + m = np.mean(data) + loc = m - scale*d*np.sqrt(2/np.pi) + elif floc is not None: + loc = floc + + if method == 'mm': + return a, loc, scale + else: + # At this point, parameter "guesses" may equal the fixed parameters + # in kwds. No harm in passing them as guesses, too. + return super().fit(data, a, loc=loc, scale=scale, **kwds) + + +skewnorm = skewnorm_gen(name='skewnorm') + + +class trapezoid_gen(rv_continuous): + r""" + A trapezoidal continuous random variable. + + %(before_notes)s + + Notes + ----- + The trapezoidal distribution can be represented with an up-sloping line + from ``loc`` to ``(loc + c*scale)``, then constant to ``(loc + d*scale)`` + and then downsloping from ``(loc + d*scale)`` to ``(loc+scale)``. This + defines the trapezoid base from ``loc`` to ``(loc+scale)`` and the flat + top from ``c`` to ``d`` proportional to the position along the base + with ``0 <= c <= d <= 1``. When ``c=d``, this is equivalent to `triang` + with the same values for `loc`, `scale` and `c`. + The method of [1]_ is used for computing moments. + + `trapezoid` takes :math:`c` and :math:`d` as shape parameters. + + %(after_notes)s + + The standard form is in the range [0, 1] with c the mode. + The location parameter shifts the start to `loc`. + The scale parameter changes the width from 1 to `scale`. + + References + ---------- + .. [1] Kacker, R.N. and Lawrence, J.F. (2007). Trapezoidal and triangular + distributions for Type B evaluation of standard uncertainty. + Metrologia 44, 117-127. :doi:`10.1088/0026-1394/44/2/003` + + %(example)s + """ + def _argcheck(self, c, d): + return (c >= 0) & (c <= 1) & (d >= 0) & (d <= 1) & (d >= c) + + def _shape_info(self): + ic = _ShapeInfo("c", False, (0, 1.0), (True, True)) + id = _ShapeInfo("d", False, (0, 1.0), (True, True)) + return [ic, id] + + def _pdf(self, x, c, d): + u = 2 / (d-c+1) + + return _lazyselect([x < c, + (c <= x) & (x <= d), + x > d], + [lambda x, c, d, u: u * x / c, + lambda x, c, d, u: u, + lambda x, c, d, u: u * (1-x) / (1-d)], + (x, c, d, u)) + + def _cdf(self, x, c, d): + return _lazyselect([x < c, + (c <= x) & (x <= d), + x > d], + [lambda x, c, d: x**2 / c / (d-c+1), + lambda x, c, d: (c + 2 * (x-c)) / (d-c+1), + lambda x, c, d: 1-((1-x) ** 2 + / (d-c+1) / (1-d))], + (x, c, d)) + + def _ppf(self, q, c, d): + qc, qd = self._cdf(c, c, d), self._cdf(d, c, d) + condlist = [q < qc, q <= qd, q > qd] + choicelist = [np.sqrt(q * c * (1 + d - c)), + 0.5 * q * (1 + d - c) + 0.5 * c, + 1 - np.sqrt((1 - q) * (d - c + 1) * (1 - d))] + return np.select(condlist, choicelist) + + def _munp(self, n, c, d): + # Using the parameterization from Kacker, 2007, with + # a=bottom left, c=top left, d=top right, b=bottom right, then + # E[X^n] = h/(n+1)/(n+2) [(b^{n+2}-d^{n+2})/(b-d) + # - ((c^{n+2} - a^{n+2})/(c-a)] + # with h = 2/((b-a) - (d-c)). The corresponding parameterization + # in scipy, has a'=loc, c'=loc+c*scale, d'=loc+d*scale, b'=loc+scale, + # which for standard form reduces to a'=0, b'=1, c'=c, d'=d. + # Substituting into E[X^n] gives the bd' term as (1 - d^{n+2})/(1 - d) + # and the ac' term as c^{n-1} for the standard form. The bd' term has + # numerical difficulties near d=1, so replace (1 - d^{n+2})/(1-d) + # with expm1((n+2)*log(d))/(d-1). + # Testing with n=18 for c=(1e-30,1-eps) shows that this is stable. + # We still require an explicit test for d=1 to prevent divide by zero, + # and now a test for d=0 to prevent log(0). + ab_term = c**(n+1) + dc_term = _lazyselect( + [d == 0.0, (0.0 < d) & (d < 1.0), d == 1.0], + [lambda d: 1.0, + lambda d: np.expm1((n+2) * np.log(d)) / (d-1.0), + lambda d: n+2], + [d]) + val = 2.0 / (1.0+d-c) * (dc_term - ab_term) / ((n+1) * (n+2)) + return val + + def _entropy(self, c, d): + # Using the parameterization from Wikipedia (van Dorp, 2003) + # with a=bottom left, c=top left, d=top right, b=bottom right + # gives a'=loc, b'=loc+c*scale, c'=loc+d*scale, d'=loc+scale, + # which for loc=0, scale=1 is a'=0, b'=c, c'=d, d'=1. + # Substituting into the entropy formula from Wikipedia gives + # the following result. + return 0.5 * (1.0-d+c) / (1.0+d-c) + np.log(0.5 * (1.0+d-c)) + + def _fitstart(self, data, args=None): + # Arbitrary, but c=d=1 fails due to being on edge of bounds + if args is None: + args = (0.33, 0.66) + return super()._fitstart(data, args=args) + + +trapezoid = trapezoid_gen(a=0.0, b=1.0, name="trapezoid") + + +class triang_gen(rv_continuous): + r"""A triangular continuous random variable. + + %(before_notes)s + + Notes + ----- + The triangular distribution can be represented with an up-sloping line from + ``loc`` to ``(loc + c*scale)`` and then downsloping for ``(loc + c*scale)`` + to ``(loc + scale)``. + + `triang` takes ``c`` as a shape parameter for :math:`0 \le c \le 1`. + + %(after_notes)s + + The standard form is in the range [0, 1] with c the mode. + The location parameter shifts the start to `loc`. + The scale parameter changes the width from 1 to `scale`. + + %(example)s + + """ + def _rvs(self, c, size=None, random_state=None): + return random_state.triangular(0, c, 1, size) + + def _argcheck(self, c): + return (c >= 0) & (c <= 1) + + def _shape_info(self): + return [_ShapeInfo("c", False, (0, 1.0), (True, True))] + + def _pdf(self, x, c): + # 0: edge case where c=0 + # 1: generalised case for x < c, don't use x <= c, as it doesn't cope + # with c = 0. + # 2: generalised case for x >= c, but doesn't cope with c = 1 + # 3: edge case where c=1 + r = _lazyselect([c == 0, + x < c, + (x >= c) & (c != 1), + c == 1], + [lambda x, c: 2 - 2 * x, + lambda x, c: 2 * x / c, + lambda x, c: 2 * (1 - x) / (1 - c), + lambda x, c: 2 * x], + (x, c)) + return r + + def _cdf(self, x, c): + r = _lazyselect([c == 0, + x < c, + (x >= c) & (c != 1), + c == 1], + [lambda x, c: 2*x - x*x, + lambda x, c: x * x / c, + lambda x, c: (x*x - 2*x + c) / (c-1), + lambda x, c: x * x], + (x, c)) + return r + + def _ppf(self, q, c): + return np.where(q < c, np.sqrt(c * q), 1-np.sqrt((1-c) * (1-q))) + + def _stats(self, c): + return ((c+1.0)/3.0, + (1.0-c+c*c)/18, + np.sqrt(2)*(2*c-1)*(c+1)*(c-2) / (5*np.power((1.0-c+c*c), 1.5)), + -3.0/5.0) + + def _entropy(self, c): + return 0.5-np.log(2) + + +triang = triang_gen(a=0.0, b=1.0, name="triang") + + +class truncexpon_gen(rv_continuous): + r"""A truncated exponential continuous random variable. + + %(before_notes)s + + Notes + ----- + The probability density function for `truncexpon` is: + + .. math:: + + f(x, b) = \frac{\exp(-x)}{1 - \exp(-b)} + + for :math:`0 <= x <= b`. + + `truncexpon` takes ``b`` as a shape parameter for :math:`b`. + + %(after_notes)s + + %(example)s + + """ + def _shape_info(self): + return [_ShapeInfo("b", False, (0, np.inf), (False, False))] + + def _get_support(self, b): + return self.a, b + + def _pdf(self, x, b): + # truncexpon.pdf(x, b) = exp(-x) / (1-exp(-b)) + return np.exp(-x)/(-sc.expm1(-b)) + + def _logpdf(self, x, b): + return -x - np.log(-sc.expm1(-b)) + + def _cdf(self, x, b): + return sc.expm1(-x)/sc.expm1(-b) + + def _ppf(self, q, b): + return -sc.log1p(q*sc.expm1(-b)) + + def _sf(self, x, b): + return (np.exp(-b) - np.exp(-x))/sc.expm1(-b) + + def _isf(self, q, b): + return -np.log(np.exp(-b) - q * sc.expm1(-b)) + + def _munp(self, n, b): + # wrong answer with formula, same as in continuous.pdf + # return sc.gamman+1)-sc.gammainc1+n, b) + if n == 1: + return (1-(b+1)*np.exp(-b))/(-sc.expm1(-b)) + elif n == 2: + return 2*(1-0.5*(b*b+2*b+2)*np.exp(-b))/(-sc.expm1(-b)) + else: + # return generic for higher moments + return super()._munp(n, b) + + def _entropy(self, b): + eB = np.exp(b) + return np.log(eB-1)+(1+eB*(b-1.0))/(1.0-eB) + + +truncexpon = truncexpon_gen(a=0.0, name='truncexpon') +truncexpon._support = (0.0, 'b') + + +# logsumexp trick for log(p + q) with only log(p) and log(q) +def _log_sum(log_p, log_q): + return sc.logsumexp([log_p, log_q], axis=0) + + +# same as above, but using -exp(x) = exp(x + πi) +def _log_diff(log_p, log_q): + return sc.logsumexp([log_p, log_q+np.pi*1j], axis=0) + + +def _log_gauss_mass(a, b): + """Log of Gaussian probability mass within an interval""" + a, b = np.broadcast_arrays(a, b) + + # Calculations in right tail are inaccurate, so we'll exploit the + # symmetry and work only in the left tail + case_left = b <= 0 + case_right = a > 0 + case_central = ~(case_left | case_right) + + def mass_case_left(a, b): + return _log_diff(_norm_logcdf(b), _norm_logcdf(a)) + + def mass_case_right(a, b): + return mass_case_left(-b, -a) + + def mass_case_central(a, b): + # Previously, this was implemented as: + # left_mass = mass_case_left(a, 0) + # right_mass = mass_case_right(0, b) + # return _log_sum(left_mass, right_mass) + # Catastrophic cancellation occurs as np.exp(log_mass) approaches 1. + # Correct for this with an alternative formulation. + # We're not concerned with underflow here: if only one term + # underflows, it was insignificant; if both terms underflow, + # the result can't accurately be represented in logspace anyway + # because sc.log1p(x) ~ x for small x. + return sc.log1p(-_norm_cdf(a) - _norm_cdf(-b)) + + # _lazyselect not working; don't care to debug it + out = np.full_like(a, fill_value=np.nan, dtype=np.complex128) + if a[case_left].size: + out[case_left] = mass_case_left(a[case_left], b[case_left]) + if a[case_right].size: + out[case_right] = mass_case_right(a[case_right], b[case_right]) + if a[case_central].size: + out[case_central] = mass_case_central(a[case_central], b[case_central]) + return np.real(out) # discard ~0j + + +class truncnorm_gen(rv_continuous): + r"""A truncated normal continuous random variable. + + %(before_notes)s + + Notes + ----- + This distribution is the normal distribution centered on ``loc`` (default + 0), with standard deviation ``scale`` (default 1), and truncated at ``a`` + and ``b`` *standard deviations* from ``loc``. For arbitrary ``loc`` and + ``scale``, ``a`` and ``b`` are *not* the abscissae at which the shifted + and scaled distribution is truncated. + + .. note:: + If ``a_trunc`` and ``b_trunc`` are the abscissae at which we wish + to truncate the distribution (as opposed to the number of standard + deviations from ``loc``), then we can calculate the distribution + parameters ``a`` and ``b`` as follows:: + + a, b = (a_trunc - loc) / scale, (b_trunc - loc) / scale + + This is a common point of confusion. For additional clarification, + please see the example below. + + %(example)s + + In the examples above, ``loc=0`` and ``scale=1``, so the plot is truncated + at ``a`` on the left and ``b`` on the right. However, suppose we were to + produce the same histogram with ``loc = 1`` and ``scale=0.5``. + + >>> loc, scale = 1, 0.5 + >>> rv = truncnorm(a, b, loc=loc, scale=scale) + >>> x = np.linspace(truncnorm.ppf(0.01, a, b), + ... truncnorm.ppf(0.99, a, b), 100) + >>> r = rv.rvs(size=1000) + + >>> fig, ax = plt.subplots(1, 1) + >>> ax.plot(x, rv.pdf(x), 'k-', lw=2, label='frozen pdf') + >>> ax.hist(r, density=True, bins='auto', histtype='stepfilled', alpha=0.2) + >>> ax.set_xlim(a, b) + >>> ax.legend(loc='best', frameon=False) + >>> plt.show() + + Note that the distribution is no longer appears to be truncated at + abscissae ``a`` and ``b``. That is because the *standard* normal + distribution is first truncated at ``a`` and ``b``, *then* the resulting + distribution is scaled by ``scale`` and shifted by ``loc``. If we instead + want the shifted and scaled distribution to be truncated at ``a`` and + ``b``, we need to transform these values before passing them as the + distribution parameters. + + >>> a_transformed, b_transformed = (a - loc) / scale, (b - loc) / scale + >>> rv = truncnorm(a_transformed, b_transformed, loc=loc, scale=scale) + >>> x = np.linspace(truncnorm.ppf(0.01, a, b), + ... truncnorm.ppf(0.99, a, b), 100) + >>> r = rv.rvs(size=10000) + + >>> fig, ax = plt.subplots(1, 1) + >>> ax.plot(x, rv.pdf(x), 'k-', lw=2, label='frozen pdf') + >>> ax.hist(r, density=True, bins='auto', histtype='stepfilled', alpha=0.2) + >>> ax.set_xlim(a-0.1, b+0.1) + >>> ax.legend(loc='best', frameon=False) + >>> plt.show() + """ + + def _argcheck(self, a, b): + return a < b + + def _shape_info(self): + ia = _ShapeInfo("a", False, (-np.inf, np.inf), (True, False)) + ib = _ShapeInfo("b", False, (-np.inf, np.inf), (False, True)) + return [ia, ib] + + def _fitstart(self, data): + # Reasonable, since support is [a, b] + if isinstance(data, CensoredData): + data = data._uncensor() + return super()._fitstart(data, args=(np.min(data), np.max(data))) + + def _get_support(self, a, b): + return a, b + + def _pdf(self, x, a, b): + return np.exp(self._logpdf(x, a, b)) + + def _logpdf(self, x, a, b): + return _norm_logpdf(x) - _log_gauss_mass(a, b) + + def _cdf(self, x, a, b): + return np.exp(self._logcdf(x, a, b)) + + def _logcdf(self, x, a, b): + x, a, b = np.broadcast_arrays(x, a, b) + logcdf = np.asarray(_log_gauss_mass(a, x) - _log_gauss_mass(a, b)) + i = logcdf > -0.1 # avoid catastrophic cancellation + if np.any(i): + logcdf[i] = np.log1p(-np.exp(self._logsf(x[i], a[i], b[i]))) + return logcdf + + def _sf(self, x, a, b): + return np.exp(self._logsf(x, a, b)) + + def _logsf(self, x, a, b): + x, a, b = np.broadcast_arrays(x, a, b) + logsf = np.asarray(_log_gauss_mass(x, b) - _log_gauss_mass(a, b)) + i = logsf > -0.1 # avoid catastrophic cancellation + if np.any(i): + logsf[i] = np.log1p(-np.exp(self._logcdf(x[i], a[i], b[i]))) + return logsf + + def _entropy(self, a, b): + A = _norm_cdf(a) + B = _norm_cdf(b) + Z = B - A + C = np.log(np.sqrt(2 * np.pi * np.e) * Z) + D = (a * _norm_pdf(a) - b * _norm_pdf(b)) / (2 * Z) + h = C + D + return h + + def _ppf(self, q, a, b): + q, a, b = np.broadcast_arrays(q, a, b) + + case_left = a < 0 + case_right = ~case_left + + def ppf_left(q, a, b): + log_Phi_x = _log_sum(_norm_logcdf(a), + np.log(q) + _log_gauss_mass(a, b)) + return sc.ndtri_exp(log_Phi_x) + + def ppf_right(q, a, b): + log_Phi_x = _log_sum(_norm_logcdf(-b), + np.log1p(-q) + _log_gauss_mass(a, b)) + return -sc.ndtri_exp(log_Phi_x) + + out = np.empty_like(q) + + q_left = q[case_left] + q_right = q[case_right] + + if q_left.size: + out[case_left] = ppf_left(q_left, a[case_left], b[case_left]) + if q_right.size: + out[case_right] = ppf_right(q_right, a[case_right], b[case_right]) + + return out + + def _isf(self, q, a, b): + # Mostly copy-paste of _ppf, but I think this is simpler than combining + q, a, b = np.broadcast_arrays(q, a, b) + + case_left = b < 0 + case_right = ~case_left + + def isf_left(q, a, b): + log_Phi_x = _log_diff(_norm_logcdf(b), + np.log(q) + _log_gauss_mass(a, b)) + return sc.ndtri_exp(np.real(log_Phi_x)) + + def isf_right(q, a, b): + log_Phi_x = _log_diff(_norm_logcdf(-a), + np.log1p(-q) + _log_gauss_mass(a, b)) + return -sc.ndtri_exp(np.real(log_Phi_x)) + + out = np.empty_like(q) + + q_left = q[case_left] + q_right = q[case_right] + + if q_left.size: + out[case_left] = isf_left(q_left, a[case_left], b[case_left]) + if q_right.size: + out[case_right] = isf_right(q_right, a[case_right], b[case_right]) + + return out + + def _munp(self, n, a, b): + def n_th_moment(n, a, b): + """ + Returns n-th moment. Defined only if n >= 0. + Function cannot broadcast due to the loop over n + """ + ab = np.asarray([a, b]) + pA, pB = self._pdf(ab, a, b) + probs = np.asarray([pA, -pB]) + cond = probs != 0 + moments = [0, 1] + for k in range(1, n+1): + # a or b might be infinite, and the corresponding pdf value + # is 0 in that case, but nan is returned for the + # multiplication. However, as b->infinity, pdf(b)*b**k -> 0. + # So it is safe to use xpx.apply_where to avoid the nan. + vals = xpx.apply_where(cond, (probs, ab), + lambda x, y: x * y**(k-1), + fill_value=0) + mk = np.sum(vals) + (k-1) * moments[-2] + moments.append(mk) + return moments[-1] + + return xpx.apply_where((n >= 0) & (a == a) & (b == b), (n, a, b), + np.vectorize(n_th_moment, otypes=[np.float64]), + fill_value=np.nan) + + def _stats(self, a, b, moments='mv'): + pA, pB = self.pdf(np.array([a, b]), a, b) + + def _truncnorm_stats_scalar(a, b, pA, pB): + ab = np.asarray([a, b]) + m1 = pA - pB + mu = m1 + # use xpx.apply_where to avoid nan (See detailed comment in _munp) + probs = np.asarray([pA, -pB]) + cond = probs != 0 + vals = xpx.apply_where(cond, (probs, ab), lambda x, y: x*y, + fill_value=0) + m2 = 1 + np.sum(vals) + vals = xpx.apply_where(cond, (probs, ab - mu), lambda x, y: x*y, + fill_value=0) + # mu2 = m2 - mu**2, but not as numerically stable as: + # mu2 = (a-mu)*pA - (b-mu)*pB + 1 + mu2 = 1 + np.sum(vals) + vals = xpx.apply_where(cond, (probs, ab), lambda x, y: x*y**2, + fill_value=0) + m3 = 2*m1 + np.sum(vals) + vals = xpx.apply_where(cond, (probs, ab), lambda x, y: x*y**3, + fill_value=0) + m4 = 3*m2 + np.sum(vals) + + mu3 = m3 + m1 * (-3*m2 + 2*m1**2) + g1 = mu3 / np.power(mu2, 1.5) + mu4 = m4 + m1*(-4*m3 + 3*m1*(2*m2 - m1**2)) + g2 = mu4 / mu2**2 - 3 + return mu, mu2, g1, g2 + + _truncnorm_stats = np.vectorize(_truncnorm_stats_scalar) + return _truncnorm_stats(a, b, pA, pB) + + +truncnorm = truncnorm_gen(name='truncnorm', momtype=1) +truncnorm._support = ('a', 'b') + + +class truncpareto_gen(rv_continuous): + r"""An upper truncated Pareto continuous random variable. + + %(before_notes)s + + See Also + -------- + pareto : Pareto distribution + + Notes + ----- + The probability density function for `truncpareto` is: + + .. math:: + + f(x, b, c) = \frac{b}{1 - c^{-b}} \frac{1}{x^{b+1}} + + for :math:`b \neq 0`, :math:`c > 1` and :math:`1 \le x \le c`. + + `truncpareto` takes `b` and `c` as shape parameters for :math:`b` and + :math:`c`. + + Notice that the upper truncation value :math:`c` is defined in + standardized form so that random values of an unscaled, unshifted variable + are within the range ``[1, c]``. + If ``u_r`` is the upper bound to a scaled and/or shifted variable, + then ``c = (u_r - loc) / scale``. In other words, the support of the + distribution becomes ``(scale + loc) <= x <= (c*scale + loc)`` when + `scale` and/or `loc` are provided. + + The ``fit`` method assumes that :math:`b` is positive; it does not produce + good results when the data is more consistent with negative :math:`b`. + + `truncpareto` can also be used to model a general power law distribution + with PDF: + + .. math:: + + f(x; a, l, h) = \frac{a}{h^a - l^a} x^{a-1} + + for :math:`a \neq 0` and :math:`0 < l < x < h`. Suppose :math:`a`, + :math:`l`, and :math:`h` are represented in code as ``a``, ``l``, and + ``h``, respectively. In this case, use `truncpareto` with parameters + ``b = -a``, ``c = h / l``, ``scale = l``, and ``loc = 0``. + + %(after_notes)s + + References + ---------- + .. [1] Burroughs, S. M., and Tebbens S. F. + "Upper-truncated power laws in natural systems." + Pure and Applied Geophysics 158.4 (2001): 741-757. + + %(example)s + + """ + + def _shape_info(self): + ib = _ShapeInfo("b", False, (-np.inf, np.inf), (False, False)) + ic = _ShapeInfo("c", False, (1.0, np.inf), (False, False)) + return [ib, ic] + + def _argcheck(self, b, c): + return (b != 0.) & (c > 1.) + + def _get_support(self, b, c): + return self.a, c + + def _pdf(self, x, b, c): + # here and below, avoid int to negative int power + x, b, c = xp_promote(x, b, c, force_floating=True, xp=np) + return b * x**-(b+1) / (1 - 1/c**b) + + def _logpdf(self, x, b, c): + x, b, c = xp_promote(x, b, c, force_floating=True, xp=np) + return xpx.apply_where(b > 0, (x, b, c), self._logpdf_pos_b, super()._logpdf) + + def _logpdf_pos_b(self, x, b, c): + return np.log(b) - np.log(-np.expm1(-b*np.log(c))) - (b+1)*np.log(x) + + def _cdf(self, x, b, c): + x, b, c = xp_promote(x, b, c, force_floating=True, xp=np) + return (1 - x**-b) / (1 - 1/c**b) + + def _logcdf(self, x, b, c): + x, b, c = xp_promote(x, b, c, force_floating=True, xp=np) + return xpx.apply_where(b > 0, (x, b, c), self._logcdf_pos_b, super()._logcdf) + + def _logcdf_pos_b(self, x, b, c): + return np.log1p(-x**-b) - np.log1p(-1/c**b) + + def _ppf(self, q, b, c): + q, b, c = xp_promote(q, b, c, force_floating=True, xp=np) + return pow(1 - (1 - 1/c**b)*q, -1/b) + + def _sf(self, x, b, c): + x, b, c = xp_promote(x, b, c, force_floating=True, xp=np) + return (x**-b - 1/c**b) / (1 - 1/c**b) + + def _logsf(self, x, b, c): + x, b, c = xp_promote(x, b, c, force_floating=True, xp=np) + return xpx.apply_where(b > 0, (x, b, c), self._logsf_pos_b, super()._logsf) + + def _logsf_pos_b(self, x, b, c): + return np.log(x**-b - 1/c**b) - np.log1p(-1/c**b) + + def _isf(self, q, b, c): + q, b, c = xp_promote(q, b, c, force_floating=True, xp=np) + return pow(1/c**b + (1 - 1/c**b)*q, -1/b) + + def _entropy(self, b, c): + return -(np.log(b/(1 - 1/c**b)) + + (b+1)*(np.log(c)/(c**b - 1) - 1/b)) + + def _munp(self, n, b, c): + n, b, c = xp_promote(n, b, c, force_floating=True, xp=np) + if (n == b).all(): + return b*np.log(c) / (1 - 1/c**b) + else: + return b / (b-n) * (c**b - c**n) / (c**b - 1) + + def _fitstart(self, data): + if isinstance(data, CensoredData): + data = data._uncensor() + b, loc, scale = pareto.fit(data) + c = (max(data) - loc)/scale + return b, c, loc, scale + + @_call_super_mom + @inherit_docstring_from(rv_continuous) + def fit(self, data, *args, **kwds): + if kwds.pop("superfit", False): + return super().fit(data, *args, **kwds) + + def log_mean(x): + return np.mean(np.log(x)) + + def harm_mean(x): + return 1/np.mean(1/x) + + def get_b(c, loc, scale): + u = (data-loc)/scale + harm_m = harm_mean(u) + log_m = log_mean(u) + quot = (harm_m-1)/log_m + return (1 - (quot-1) / (quot - (1 - 1/c)*harm_m/np.log(c)))/log_m + + def get_c(loc, scale): + return (mx - loc)/scale + + def get_loc(fc, fscale): + if fscale: # (fscale and fc) or (fscale and not fc) + loc = mn - fscale + return loc + if fc: + loc = (fc*mn - mx)/(fc - 1) + return loc + + def get_scale(loc): + return mn - loc + + # Functions used for optimisation; partial derivatives of + # the Lagrangian, set to equal 0. + + def dL_dLoc(loc, b_=None): + # Partial derivative wrt location. + # Optimised upon when no parameters, or only b, are fixed. + scale = get_scale(loc) + c = get_c(loc, scale) + b = get_b(c, loc, scale) if b_ is None else b_ + harm_m = harm_mean((data - loc)/scale) + return 1 - (1 + (c - 1)/(c**(b+1) - c)) * (1 - 1/(b+1)) * harm_m + + def dL_dB(b, logc, logm): + # Partial derivative wrt b. + # Optimised upon whenever at least one parameter but b is fixed, + # and b is free. + return b - np.log1p(b*logc / (1 - b*logm)) / logc + + def fallback(data, *args, **kwargs): + # Should any issue arise, default to the general fit method. + return super(truncpareto_gen, self).fit(data, *args, **kwargs) + + parameters = _check_fit_input_parameters(self, data, args, kwds) + data, fb, fc, floc, fscale = parameters + mn, mx = data.min(), data.max() + mn_inf = np.nextafter(mn, -np.inf) + + if (fb is not None + and fc is not None + and floc is not None + and fscale is not None): + raise ValueError("All parameters fixed." + "There is nothing to optimize.") + elif fc is None and floc is None and fscale is None: + if fb is None: + def cond_b(loc): + # b is positive only if this function is positive + scale = get_scale(loc) + c = get_c(loc, scale) + harm_m = harm_mean((data - loc)/scale) + return (1 + 1/(c-1)) * np.log(c) / harm_m - 1 + + # This gives an upper bound on loc allowing for a positive b. + # Iteratively look for a bracket for root_scalar. + mn_inf = np.nextafter(mn, -np.inf) + rbrack = mn_inf + i = 0 + lbrack = rbrack - 1 + while ((lbrack > -np.inf) + and (cond_b(lbrack)*cond_b(rbrack) >= 0)): + i += 1 + lbrack = rbrack - np.power(2., i) + if not lbrack > -np.inf: + return fallback(data, *args, **kwds) + res = root_scalar(cond_b, bracket=(lbrack, rbrack)) + if not res.converged: + return fallback(data, *args, **kwds) + + # Determine the MLE for loc. + # Iteratively look for a bracket for root_scalar. + rbrack = res.root - 1e-3 # grad_loc is numerically ill-behaved + lbrack = rbrack - 1 + i = 0 + while ((lbrack > -np.inf) + and (dL_dLoc(lbrack)*dL_dLoc(rbrack) >= 0)): + i += 1 + lbrack = rbrack - np.power(2., i) + if not lbrack > -np.inf: + return fallback(data, *args, **kwds) + res = root_scalar(dL_dLoc, bracket=(lbrack, rbrack)) + if not res.converged: + return fallback(data, *args, **kwds) + loc = res.root + scale = get_scale(loc) + c = get_c(loc, scale) + b = get_b(c, loc, scale) + + std_data = (data - loc)/scale + # The expression of b relies on b being bounded above. + up_bound_b = min(1/log_mean(std_data), + 1/(harm_mean(std_data)-1)) + if not (b < up_bound_b): + return fallback(data, *args, **kwds) + else: + # We know b is positive (or a FitError will be triggered) + # so we let loc get close to min(data). + rbrack = mn_inf + lbrack = mn_inf - 1 + i = 0 + # Iteratively look for a bracket for root_scalar. + while (lbrack > -np.inf + and (dL_dLoc(lbrack, fb) + * dL_dLoc(rbrack, fb) >= 0)): + i += 1 + lbrack = rbrack - 2**i + if not lbrack > -np.inf: + return fallback(data, *args, **kwds) + res = root_scalar(dL_dLoc, (fb,), + bracket=(lbrack, rbrack)) + if not res.converged: + return fallback(data, *args, **kwds) + loc = res.root + scale = get_scale(loc) + c = get_c(loc, scale) + b = fb + else: + # At least one of the parameters determining the support is fixed; + # the others then have analytical expressions from the constraints. + # The completely determined case (fixed c, loc and scale) + # has to be checked for not overflowing the support. + # If not fixed, b has to be determined numerically. + loc = floc if floc is not None else get_loc(fc, fscale) + scale = fscale or get_scale(loc) + c = fc or get_c(loc, scale) + + # Unscaled, translated values should be positive when the location + # is fixed. If it is not the case, we end up with negative `scale` + # and `c`, which would trigger a FitError before exiting the + # method. + if floc is not None and data.min() - floc < 0: + raise FitDataError("truncpareto", lower=1, upper=c) + + # Standardised values should be within the distribution support + # when all parameters controlling it are fixed. If it not the case, + # `fc` is overridden by `c` determined from `floc` and `fscale` when + # raising the exception. + if fc and (floc is not None) and fscale: + if data.max() > fc*fscale + floc: + raise FitDataError("truncpareto", lower=1, + upper=get_c(loc, scale)) + + # The other constraints should be automatically satisfied + # from the analytical expressions of the parameters. + # If fc or fscale are respectively less than one or less than 0, + # a FitError is triggered before exiting the method. + + if fb is None: + std_data = (data - loc)/scale + logm = log_mean(std_data) + logc = np.log(c) + # Condition for a positive root to exist. + if not (2*logm < logc): + return fallback(data, *args, **kwds) + + lbrack = 1/logm + 1/(logm - logc) + rbrack = np.nextafter(1/logm, 0) + try: + res = root_scalar(dL_dB, (logc, logm), + bracket=(lbrack, rbrack)) + # we should then never get there + if not res.converged: + return fallback(data, *args, **kwds) + b = res.root + except ValueError: + b = rbrack + else: + b = fb + + # The distribution requires that `scale+loc <= data <= c*scale+loc`. + # To avoid numerical issues, some tuning may be necessary. + # We adjust `scale` to satisfy the lower bound, and we adjust + # `c` to satisfy the upper bound. + if not (scale+loc) < mn: + if fscale: + loc = np.nextafter(loc, -np.inf) + else: + scale = get_scale(loc) + scale = np.nextafter(scale, 0) + if not (c*scale+loc) > mx: + c = get_c(loc, scale) + c = np.nextafter(c, np.inf) + + if not (np.all(self._argcheck(b, c)) and (scale > 0)): + return fallback(data, *args, **kwds) + + params_override = b, c, loc, scale + if floc is None and fscale is None: + # Based on testing in gh-16782, the following methods are only + # reliable if either `floc` or `fscale` are provided. They are + # fast, though, so might as well see if they are better than the + # generic method. + params_super = fallback(data, *args, **kwds) + nllf_override = self.nnlf(params_override, data) + nllf_super = self.nnlf(params_super, data) + if nllf_super < nllf_override: + return params_super + + return params_override + + +truncpareto = truncpareto_gen(a=1.0, name='truncpareto') +truncpareto._support = (1.0, 'c') + + +class tukeylambda_gen(rv_continuous): + r"""A Tukey-Lamdba continuous random variable. + + %(before_notes)s + + Notes + ----- + A flexible distribution, able to represent and interpolate between the + following distributions: + + - Cauchy (:math:`lambda = -1`) + - logistic (:math:`lambda = 0`) + - approx Normal (:math:`lambda = 0.14`) + - uniform from -1 to 1 (:math:`lambda = 1`) + + `tukeylambda` takes a real number :math:`lambda` (denoted ``lam`` + in the implementation) as a shape parameter. + + %(after_notes)s + + %(example)s + + """ + _support_mask = rv_continuous._open_support_mask + + def _argcheck(self, lam): + return np.isfinite(lam) + + def _shape_info(self): + return [_ShapeInfo("lam", False, (-np.inf, np.inf), (False, False))] + + def _get_support(self, lam): + b = xpx.apply_where(lam > 0, lam, + lambda lam: 1/lam, + fill_value=np.inf) + return -b, b + + def _pdf(self, x, lam): + Fx = np.asarray(sc.tklmbda(x, lam)) + Px = Fx**(lam-1.0) + (np.asarray(1-Fx))**(lam-1.0) + with np.errstate(divide='ignore'): + Px = 1.0/np.asarray(Px) + return np.where((lam <= 0) | (abs(x) < 1.0/np.asarray(lam)), Px, 0.0) + + def _cdf(self, x, lam): + return sc.tklmbda(x, lam) + + def _ppf(self, q, lam): + return sc.boxcox(q, lam) - sc.boxcox1p(-q, lam) + + def _stats(self, lam): + return 0, _tlvar(lam), 0, _tlkurt(lam) + + def _entropy(self, lam): + @np.vectorize + def entropy_1d(lam): + def integ(p): + return np.log(pow(p, lam - 1) + pow(1 - p, lam - 1)) + return integrate.quad(integ, 0, 1)[0] + return entropy_1d(lam) + + +tukeylambda = tukeylambda_gen(name='tukeylambda') + + +class FitUniformFixedScaleDataError(FitDataError): + def __init__(self, ptp, fscale): + self.args = ( + "Invalid values in `data`. Maximum likelihood estimation with " + "the uniform distribution and fixed scale requires that " + f"np.ptp(data) <= fscale, but np.ptp(data) = {ptp} and " + f"fscale = {fscale}." + ) + + +class uniform_gen(rv_continuous): + r"""A uniform continuous random variable. + + %(before_notes)s + + Notes + ----- + The probability density function for `uniform` is: + + .. math:: + + f(x) = \begin{cases} + 1 & \text{for } 0 \le x \le 1 \\ + 0 & \text{otherwise} + \end{cases} + + In the standard form, the distribution is uniform on ``[0, 1]``. Using + the parameters ``loc`` and ``scale``, one obtains the uniform distribution + on ``[loc, loc + scale]``. + + References + ---------- + .. [1] "Continuous uniform distribution", Wikipedia, + https://en.wikipedia.org/wiki/Continuous_uniform_distribution + + %(example)s + + """ + def _shape_info(self): + return [] + + def _rvs(self, size=None, random_state=None): + return random_state.uniform(0.0, 1.0, size) + + def _pdf(self, x): + return 1.0*(x == x) + + def _cdf(self, x): + return x + + def _ppf(self, q): + return q + + def _stats(self): + return 0.5, 1.0/12, 0, -1.2 + + def _entropy(self): + return 0.0 + + @_call_super_mom + def fit(self, data, *args, **kwds): + """ + Maximum likelihood estimate for the location and scale parameters. + + `uniform.fit` uses only the following parameters. Because exact + formulas are used, the parameters related to optimization that are + available in the `fit` method of other distributions are ignored + here. The only positional argument accepted is `data`. + + Parameters + ---------- + data : array_like + Data to use in calculating the maximum likelihood estimate. + floc : float, optional + Hold the location parameter fixed to the specified value. + fscale : float, optional + Hold the scale parameter fixed to the specified value. + + Returns + ------- + loc, scale : float + Maximum likelihood estimates for the location and scale. + + Notes + ----- + An error is raised if `floc` is given and any values in `data` are + less than `floc`, or if `fscale` is given and `fscale` is less + than ``data.max() - data.min()``. An error is also raised if both + `floc` and `fscale` are given. + + Examples + -------- + >>> import numpy as np + >>> from scipy.stats import uniform + + We'll fit the uniform distribution to `x`: + + >>> x = np.array([2, 2.5, 3.1, 9.5, 13.0]) + + For a uniform distribution MLE, the location is the minimum of the + data, and the scale is the maximum minus the minimum. + + >>> loc, scale = uniform.fit(x) + >>> loc + 2.0 + >>> scale + 11.0 + + If we know the data comes from a uniform distribution where the support + starts at 0, we can use ``floc=0``: + + >>> loc, scale = uniform.fit(x, floc=0) + >>> loc + 0.0 + >>> scale + 13.0 + + Alternatively, if we know the length of the support is 12, we can use + ``fscale=12``: + + >>> loc, scale = uniform.fit(x, fscale=12) + >>> loc + 1.5 + >>> scale + 12.0 + + In that last example, the support interval is [1.5, 13.5]. This + solution is not unique. For example, the distribution with ``loc=2`` + and ``scale=12`` has the same likelihood as the one above. When + `fscale` is given and it is larger than ``data.max() - data.min()``, + the parameters returned by the `fit` method center the support over + the interval ``[data.min(), data.max()]``. + + """ + if len(args) > 0: + raise TypeError("Too many arguments.") + + floc = kwds.pop('floc', None) + fscale = kwds.pop('fscale', None) + + _remove_optimizer_parameters(kwds) + + if floc is not None and fscale is not None: + # This check is for consistency with `rv_continuous.fit`. + raise ValueError("All parameters fixed. There is nothing to " + "optimize.") + + data = np.asarray(data) + + if not np.isfinite(data).all(): + raise ValueError("The data contains non-finite values.") + + # MLE for the uniform distribution + # -------------------------------- + # The PDF is + # + # f(x, loc, scale) = {1/scale for loc <= x <= loc + scale + # {0 otherwise} + # + # The likelihood function is + # L(x, loc, scale) = (1/scale)**n + # where n is len(x), assuming loc <= x <= loc + scale for all x. + # The log-likelihood is + # l(x, loc, scale) = -n*log(scale) + # The log-likelihood is maximized by making scale as small as possible, + # while keeping loc <= x <= loc + scale. So if neither loc nor scale + # are fixed, the log-likelihood is maximized by choosing + # loc = x.min() + # scale = np.ptp(x) + # If loc is fixed, it must be less than or equal to x.min(), and then + # the scale is + # scale = x.max() - loc + # If scale is fixed, it must not be less than np.ptp(x). If scale is + # greater than np.ptp(x), the solution is not unique. Note that the + # likelihood does not depend on loc, except for the requirement that + # loc <= x <= loc + scale. All choices of loc for which + # x.max() - scale <= loc <= x.min() + # have the same log-likelihood. In this case, we choose loc such that + # the support is centered over the interval [data.min(), data.max()]: + # loc = x.min() = 0.5*(scale - np.ptp(x)) + + if fscale is None: + # scale is not fixed. + if floc is None: + # loc is not fixed, scale is not fixed. + loc = data.min() + scale = np.ptp(data) + else: + # loc is fixed, scale is not fixed. + loc = floc + scale = data.max() - loc + if data.min() < loc: + raise FitDataError("uniform", lower=loc, upper=loc + scale) + else: + # loc is not fixed, scale is fixed. + ptp = np.ptp(data) + if ptp > fscale: + raise FitUniformFixedScaleDataError(ptp=ptp, fscale=fscale) + # If ptp < fscale, the ML estimate is not unique; see the comments + # above. We choose the distribution for which the support is + # centered over the interval [data.min(), data.max()]. + loc = data.min() - 0.5*(fscale - ptp) + scale = fscale + + # We expect the return values to be floating point, so ensure it + # by explicitly converting to float. + return float(loc), float(scale) + + +uniform = uniform_gen(a=0.0, b=1.0, name='uniform') + + +class vonmises_gen(rv_continuous): + r"""A Von Mises continuous random variable. + + %(before_notes)s + + See Also + -------- + scipy.stats.vonmises_fisher : Von-Mises Fisher distribution on a + hypersphere + + Notes + ----- + The probability density function for `vonmises` and `vonmises_line` is: + + .. math:: + + f(x, \kappa) = \frac{ \exp(\kappa \cos(x)) }{ 2 \pi I_0(\kappa) } + + for :math:`-\pi \le x \le \pi`, :math:`\kappa \ge 0`. :math:`I_0` is the + modified Bessel function of order zero (`scipy.special.i0`). + + `vonmises` is a circular distribution which does not restrict the + distribution to a fixed interval. Currently, there is no circular + distribution framework in SciPy. The ``cdf`` is implemented such that + ``cdf(x + 2*np.pi) == cdf(x) + 1``. + + `vonmises_line` is the same distribution, defined on :math:`[-\pi, \pi]` + on the real line. This is a regular (i.e. non-circular) distribution. + + Note about distribution parameters: `vonmises` and `vonmises_line` take + ``kappa`` as a shape parameter (concentration) and ``loc`` as the location + (circular mean). A ``scale`` parameter is accepted but does not have any + effect. + + References + ---------- + .. [1] Mardia, K. V. and Jupp, P. E. *Directional Statistics*. + John Wiley & Sons, 1999, p. 36. + .. [2] "von Mises distribution", Wikipedia, + https://en.wikipedia.org/wiki/Von_Mises_distribution + + Examples + -------- + Import the necessary modules. + + >>> import numpy as np + >>> import matplotlib.pyplot as plt + >>> from scipy.stats import vonmises + + Define distribution parameters. + + >>> loc = 0.5 * np.pi # circular mean + >>> kappa = 1 # concentration + + Compute the probability density at ``x=0`` via the ``pdf`` method. + + >>> vonmises.pdf(0, loc=loc, kappa=kappa) + 0.12570826359722018 + + Verify that the percentile function ``ppf`` inverts the cumulative + distribution function ``cdf`` up to floating point accuracy. + + >>> x = 1 + >>> cdf_value = vonmises.cdf(x, loc=loc, kappa=kappa) + >>> ppf_value = vonmises.ppf(cdf_value, loc=loc, kappa=kappa) + >>> x, cdf_value, ppf_value + (1, 0.31489339900904967, 1.0000000000000004) + + Draw 1000 random variates by calling the ``rvs`` method. + + >>> sample_size = 1000 + >>> sample = vonmises(loc=loc, kappa=kappa).rvs(sample_size) + + Plot the von Mises density on a Cartesian and polar grid to emphasize + that it is a circular distribution. + + >>> fig = plt.figure(figsize=(12, 6)) + >>> left = plt.subplot(121) + >>> right = plt.subplot(122, projection='polar') + >>> x = np.linspace(-np.pi, np.pi, 500) + >>> vonmises_pdf = vonmises.pdf(x, loc=loc, kappa=kappa) + >>> ticks = [0, 0.15, 0.3] + + The left image contains the Cartesian plot. + + >>> left.plot(x, vonmises_pdf) + >>> left.set_yticks(ticks) + >>> number_of_bins = int(np.sqrt(sample_size)) + >>> left.hist(sample, density=True, bins=number_of_bins) + >>> left.set_title("Cartesian plot") + >>> left.set_xlim(-np.pi, np.pi) + >>> left.grid(True) + + The right image contains the polar plot. + + >>> right.plot(x, vonmises_pdf, label="PDF") + >>> right.set_yticks(ticks) + >>> right.hist(sample, density=True, bins=number_of_bins, + ... label="Histogram") + >>> right.set_title("Polar plot") + >>> right.legend(bbox_to_anchor=(0.15, 1.06)) + + """ + def _shape_info(self): + return [_ShapeInfo("kappa", False, (0, np.inf), (True, False))] + + def _argcheck(self, kappa): + return kappa >= 0 + + def _rvs(self, kappa, size=None, random_state=None): + return random_state.vonmises(0.0, kappa, size=size) + + @inherit_docstring_from(rv_continuous) + def rvs(self, *args, **kwds): + rvs = super().rvs(*args, **kwds) + return np.mod(rvs + np.pi, 2*np.pi) - np.pi + + def _pdf(self, x, kappa): + # vonmises.pdf(x, kappa) = exp(kappa * cos(x)) / (2*pi*I[0](kappa)) + # = exp(kappa * (cos(x) - 1)) / + # (2*pi*exp(-kappa)*I[0](kappa)) + # = exp(kappa * cosm1(x)) / (2*pi*i0e(kappa)) + return np.exp(kappa*sc.cosm1(x)) / (2*np.pi*sc.i0e(kappa)) + + def _logpdf(self, x, kappa): + # vonmises.pdf(x, kappa) = exp(kappa * cosm1(x)) / (2*pi*i0e(kappa)) + return kappa * sc.cosm1(x) - np.log(2*np.pi) - np.log(sc.i0e(kappa)) + + def _cdf(self, x, kappa): + return _stats.von_mises_cdf(kappa, x) + + def _stats_skip(self, kappa): + return 0, None, 0, None + + def _entropy(self, kappa): + # vonmises.entropy(kappa) = -kappa * I[1](kappa) / I[0](kappa) + + # log(2 * np.pi * I[0](kappa)) + # = -kappa * I[1](kappa) * exp(-kappa) / + # (I[0](kappa) * exp(-kappa)) + + # log(2 * np.pi * + # I[0](kappa) * exp(-kappa) / exp(-kappa)) + # = -kappa * sc.i1e(kappa) / sc.i0e(kappa) + + # log(2 * np.pi * i0e(kappa)) + kappa + return (-kappa * sc.i1e(kappa) / sc.i0e(kappa) + + np.log(2 * np.pi * sc.i0e(kappa)) + kappa) + + @extend_notes_in_docstring(rv_continuous, notes="""\ + The default limits of integration are endpoints of the interval + of width ``2*pi`` centered at `loc` (e.g. ``[-pi, pi]`` when + ``loc=0``).\n\n""") + def expect(self, func=None, args=(), loc=0, scale=1, lb=None, ub=None, + conditional=False, **kwds): + _a, _b = -np.pi, np.pi + + if lb is None: + lb = loc + _a + if ub is None: + ub = loc + _b + + return super().expect(func, args, loc, + scale, lb, ub, conditional, **kwds) + + @_call_super_mom + @extend_notes_in_docstring(rv_continuous, notes="""\ + Fit data is assumed to represent angles and will be wrapped onto the + unit circle. `f0` and `fscale` are ignored; the returned shape is + always the maximum likelihood estimate and the scale is always + 1. Initial guesses are ignored.\n\n""") + def fit(self, data, *args, **kwds): + if kwds.pop('superfit', False): + return super().fit(data, *args, **kwds) + + data, fshape, floc, fscale = _check_fit_input_parameters(self, data, + args, kwds) + if self.a == -np.pi: + # vonmises line case, here the default fit method will be used + return super().fit(data, *args, **kwds) + + # wrap data to interval [0, 2*pi] + data = np.mod(data, 2 * np.pi) + + def find_mu(data): + return stats.circmean(data) + + def find_kappa(data, loc): + # Usually, sources list the following as the equation to solve for + # the MLE of the shape parameter: + # r = I[1](kappa)/I[0](kappa), where r = mean resultant length + # This is valid when the location is the MLE of location. + # More generally, when the location may be fixed at an arbitrary + # value, r should be defined as follows: + r = np.sum(np.cos(loc - data))/len(data) + # See gh-18128 for more information. + + # The function r[0](kappa) := I[1](kappa)/I[0](kappa) is monotonic + # increasing from r[0](0) = 0 to r[0](+inf) = 1. The partial + # derivative of the log likelihood function with respect to kappa + # is monotonic decreasing in kappa. + if r == 1: + # All observations are (almost) equal to the mean. Return + # some large kappa such that r[0](kappa) = 1.0 numerically. + return 1e16 + elif r > 0: + def solve_for_kappa(kappa): + return sc.i1e(kappa)/sc.i0e(kappa) - r + + # The bounds of the root of r[0](kappa) = r are derived from + # selected bounds of r[0](x) given in [1, Eq. 11 & 16]. See + # gh-20102 for details. + # + # [1] Amos, D. E. (1973). Computation of Modified Bessel + # Functions and Their Ratios. Mathematics of Computation, + # 28(125): 239-251. + lower_bound = r/(1-r)/(1+r) + upper_bound = 2*lower_bound + + # The bounds are violated numerically for certain values of r, + # where solve_for_kappa evaluated at the bounds have the same + # sign. This indicates numerical imprecision of i1e()/i0e(). + # Return the violated bound in this case as it's more accurate. + if solve_for_kappa(lower_bound) >= 0: + return lower_bound + elif solve_for_kappa(upper_bound) <= 0: + return upper_bound + else: + root_res = root_scalar(solve_for_kappa, method="brentq", + bracket=(lower_bound, upper_bound)) + return root_res.root + else: + # if the provided floc is very far from the circular mean, + # the mean resultant length r can become negative. + # In that case, the equation + # I[1](kappa)/I[0](kappa) = r does not have a solution. + # The maximum likelihood kappa is then 0 which practically + # results in the uniform distribution on the circle. As + # vonmises is defined for kappa > 0, return instead the + # smallest floating point value. + # See gh-18190 for more information + return np.finfo(float).tiny + + # location likelihood equation has a solution independent of kappa + loc = floc if floc is not None else find_mu(data) + # shape likelihood equation depends on location + shape = fshape if fshape is not None else find_kappa(data, loc) + + loc = np.mod(loc + np.pi, 2 * np.pi) - np.pi # ensure in [-pi, pi] + return shape, loc, 1 # scale is not handled + + +vonmises = vonmises_gen(name='vonmises') +vonmises_line = vonmises_gen(a=-np.pi, b=np.pi, name='vonmises_line') + + +class wald_gen(invgauss_gen): + r"""A Wald continuous random variable. + + %(before_notes)s + + Notes + ----- + The probability density function for `wald` is: + + .. math:: + + f(x) = \frac{1}{\sqrt{2\pi x^3}} \exp(- \frac{ (x-1)^2 }{ 2x }) + + for :math:`x >= 0`. + + `wald` is a special case of `invgauss` with ``mu=1``. + + %(after_notes)s + + %(example)s + """ + _support_mask = rv_continuous._open_support_mask + + def _shape_info(self): + return [] + + def _rvs(self, size=None, random_state=None): + return random_state.wald(1.0, 1.0, size=size) + + def _pdf(self, x): + # wald.pdf(x) = 1/sqrt(2*pi*x**3) * exp(-(x-1)**2/(2*x)) + return invgauss._pdf(x, 1.0) + + def _cdf(self, x): + return invgauss._cdf(x, 1.0) + + def _sf(self, x): + return invgauss._sf(x, 1.0) + + def _ppf(self, x): + return invgauss._ppf(x, 1.0) + + def _isf(self, x): + return invgauss._isf(x, 1.0) + + def _logpdf(self, x): + return invgauss._logpdf(x, 1.0) + + def _logcdf(self, x): + return invgauss._logcdf(x, 1.0) + + def _logsf(self, x): + return invgauss._logsf(x, 1.0) + + def _stats(self): + return 1.0, 1.0, 3.0, 15.0 + + def _entropy(self): + return invgauss._entropy(1.0) + + +wald = wald_gen(a=0.0, name="wald") + + +class wrapcauchy_gen(rv_continuous): + r"""A wrapped Cauchy continuous random variable. + + %(before_notes)s + + Notes + ----- + The probability density function for `wrapcauchy` is: + + .. math:: + + f(x, c) = \frac{1-c^2}{2\pi (1+c^2 - 2c \cos(x))} + + for :math:`0 \le x \le 2\pi`, :math:`0 < c < 1`. + + `wrapcauchy` takes ``c`` as a shape parameter for :math:`c`. + + %(after_notes)s + + %(example)s + + """ + def _argcheck(self, c): + return (c > 0) & (c < 1) + + def _shape_info(self): + return [_ShapeInfo("c", False, (0, 1), (False, False))] + + def _pdf(self, x, c): + # wrapcauchy.pdf(x, c) = (1-c**2) / (2*pi*(1+c**2-2*c*cos(x))) + return (1.0-c*c)/(2*np.pi*(1+c*c-2*c*np.cos(x))) + + def _cdf(self, x, c): + + def f1(x, cr): + # CDF for 0 <= x < pi + return 1/np.pi * np.arctan(cr*np.tan(x/2)) + + def f2(x, cr): + # CDF for pi <= x <= 2*pi + return 1 - 1/np.pi * np.arctan(cr*np.tan((2*np.pi - x)/2)) + + cr = (1 + c)/(1 - c) + return xpx.apply_where(x < np.pi, (x, cr), f1, f2) + + def _ppf(self, q, c): + val = (1.0-c)/(1.0+c) + rcq = 2*np.arctan(val*np.tan(np.pi*q)) + rcmq = 2*np.pi-2*np.arctan(val*np.tan(np.pi*(1-q))) + return np.where(q < 1.0/2, rcq, rcmq) + + def _entropy(self, c): + return np.log(2*np.pi*(1-c*c)) + + def _fitstart(self, data): + # Use 0.5 as the initial guess of the shape parameter. + # For the location and scale, use the minimum and + # peak-to-peak/(2*pi), respectively. + if isinstance(data, CensoredData): + data = data._uncensor() + return 0.5, np.min(data), np.ptp(data)/(2*np.pi) + + @inherit_docstring_from(rv_continuous) + def rvs(self, *args, **kwds): + rvs = super().rvs(*args, **kwds) + return np.mod(rvs, 2*np.pi) + +wrapcauchy = wrapcauchy_gen(a=0.0, b=2*np.pi, name='wrapcauchy') + + +class gennorm_gen(rv_continuous): + r"""A (symmetric) generalized normal continuous random variable. + + %(before_notes)s + + See Also + -------- + laplace : Laplace distribution + norm : normal distribution + + Notes + ----- + The (symmetric) generalized normal distribution is also known as the + Subbotin distribution, exponential power distribution, and generalized + error distribution [1]_. + + The probability density function for `gennorm` is [1]_: + + .. math:: + + f(x, \beta) = \frac{\beta}{2 \Gamma(1/\beta)} \exp(-|x|^\beta), + + where :math:`x` is a real number, :math:`\beta > 0` and + :math:`\Gamma` is the gamma function (`scipy.special.gamma`). + + `gennorm` takes ``beta`` as a shape parameter for :math:`\beta`. + For :math:`\beta = 1`, it is identical to a Laplace distribution. + For :math:`\beta = 2`, it is identical to a normal distribution + (with ``scale=1/sqrt(2)``). + + References + ---------- + + .. [1] "Generalized normal distribution, Version 1", + https://en.wikipedia.org/wiki/Generalized_normal_distribution#Version_1 + + .. [2] Nardon, Martina, and Paolo Pianca. "Simulation techniques for + generalized Gaussian densities." Journal of Statistical + Computation and Simulation 79.11 (2009): 1317-1329 + + .. [3] Wicklin, Rick. "Simulate data from a generalized Gaussian + distribution" in The DO Loop blog, September 21, 2016, + https://blogs.sas.com/content/iml/2016/09/21/simulate-generalized-gaussian-sas.html + + %(example)s + + """ + def _shape_info(self): + return [_ShapeInfo("beta", False, (0, np.inf), (False, False))] + + def _pdf(self, x, beta): + return np.exp(self._logpdf(x, beta)) + + def _logpdf(self, x, beta): + return np.log(0.5*beta) - sc.gammaln(1.0/beta) - abs(x)**beta + + def _cdf(self, x, beta): + c = 0.5 * np.sign(x) + # evaluating (.5 + c) first prevents numerical cancellation + return (0.5 + c) - c * sc.gammaincc(1.0/beta, abs(x)**beta) + + def _ppf(self, x, beta): + c = np.sign(x - 0.5) + # evaluating (1. + c) first prevents numerical cancellation + return c * sc.gammainccinv(1.0/beta, (1.0 + c) - 2.0*c*x)**(1.0/beta) + + def _sf(self, x, beta): + return self._cdf(-x, beta) + + def _isf(self, x, beta): + return -self._ppf(x, beta) + + def _munp(self, n, beta): + if n == 0: + return 1. + if n % 2 == 0: + c1, cn = sc.gammaln([1.0/beta, (n + 1.0)/beta]) + return np.exp(cn - c1) + else: + return 0. + + def _stats(self, beta): + c1, c3, c5 = sc.gammaln([1.0/beta, 3.0/beta, 5.0/beta]) + return 0., np.exp(c3 - c1), 0., np.exp(c5 + c1 - 2.0*c3) - 3. + + def _entropy(self, beta): + return 1. / beta - np.log(.5 * beta) + sc.gammaln(1. / beta) + + def _rvs(self, beta, size=None, random_state=None): + # see [2]_ for the algorithm + # see [3]_ for reference implementation in SAS + z = random_state.gamma(1/beta, size=size) + y = z ** (1/beta) + # convert y to array to ensure masking support + y = np.asarray(y) + mask = random_state.random(size=y.shape) < 0.5 + y[mask] = -y[mask] + return y + + +gennorm = gennorm_gen(name='gennorm') + + +class halfgennorm_gen(rv_continuous): + r"""The upper half of a generalized normal continuous random variable. + + %(before_notes)s + + See Also + -------- + gengamma : generalized gamma distribution + gennorm : generalized normal distribution + expon : exponential distribution + halfnorm : half normal distribution + + Notes + ----- + The probability density function for `halfgennorm` is: + + .. math:: + + f(x, \beta) = \frac{\beta}{\Gamma(1/\beta)} \exp(-|x|^\beta) + + for :math:`x, \beta > 0`. :math:`\Gamma` is the gamma function + (`scipy.special.gamma`). + + `halfgennorm` takes ``beta`` as a shape parameter for :math:`\beta`. + For :math:`\beta = 1`, it is identical to an exponential distribution. + For :math:`\beta = 2`, it is identical to a half normal distribution + (with ``scale=1/sqrt(2)``). + + `halfgennorm` is the upper half of a generalized normal continuous + random variable [1]_. + + `halfgennorm` is a special case of the generalized gamma distribution [2]_, + which is implemented in SciPy as `gengamma`: + ``halfgennorm(beta) = gengamma(a=1/beta, c=beta)``. + + References + ---------- + + .. [1] "Generalized normal distribution, Version 1", Wikipedia, + https://en.wikipedia.org/wiki/Generalized_normal_distribution#Version_1 + .. [2] "Generalized gamma distribution", Wikipedia, + https://en.wikipedia.org/wiki/Generalized_gamma_distribution + + %(example)s + + """ + def _shape_info(self): + return [_ShapeInfo("beta", False, (0, np.inf), (False, False))] + + def _pdf(self, x, beta): + # beta + # halfgennorm.pdf(x, beta) = ------------- exp(-|x|**beta) + # gamma(1/beta) + return np.exp(self._logpdf(x, beta)) + + def _logpdf(self, x, beta): + return np.log(beta) - sc.gammaln(1.0/beta) - x**beta + + def _cdf(self, x, beta): + return sc.gammainc(1.0/beta, x**beta) + + def _ppf(self, x, beta): + return sc.gammaincinv(1.0/beta, x)**(1.0/beta) + + def _sf(self, x, beta): + return sc.gammaincc(1.0/beta, x**beta) + + def _isf(self, x, beta): + return sc.gammainccinv(1.0/beta, x)**(1.0/beta) + + def _munp(self, n, beta): + return sc.poch(1/beta, n/beta) + + def _entropy(self, beta): + return 1.0/beta - np.log(beta) + sc.gammaln(1.0/beta) + + +halfgennorm = halfgennorm_gen(a=0, name='halfgennorm') + + +class crystalball_gen(rv_continuous): + r""" + Crystalball distribution. + + %(before_notes)s + + Notes + ----- + The probability density function for `crystalball` is: + + .. math:: + + f(x, \beta, m) = \begin{cases} + N \exp(-x^2 / 2), &\text{for } x > -\beta\\ + N A (B - x)^{-m} &\text{for } x \le -\beta + \end{cases} + + where :math:`A = (m / |\beta|)^m \exp(-\beta^2 / 2)`, + :math:`B = m/|\beta| - |\beta|` and :math:`N` is a normalisation constant. + + `crystalball` takes :math:`\beta > 0` and :math:`m > 1` as shape + parameters. :math:`\beta` defines the point where the pdf changes + from a power-law to a Gaussian distribution. :math:`m` is the power + of the power-law tail. + + %(after_notes)s + + .. versionadded:: 0.19.0 + + References + ---------- + .. [1] "Crystal Ball Function", + https://en.wikipedia.org/wiki/Crystal_Ball_function + + %(example)s + """ + def _argcheck(self, beta, m): + """ + Shape parameter bounds are m > 1 and beta > 0. + """ + return (m > 1) & (beta > 0) + + def _shape_info(self): + ibeta = _ShapeInfo("beta", False, (0, np.inf), (False, False)) + im = _ShapeInfo("m", False, (1, np.inf), (False, False)) + return [ibeta, im] + + def _fitstart(self, data): + # Arbitrary, but the default m=1 is not valid + return super()._fitstart(data, args=(1, 1.5)) + + def _pdf(self, x, beta, m): + """ + Return PDF of the crystalball function. + + -- + | exp(-x**2 / 2), for x > -beta + crystalball.pdf(x, beta, m) = N * | + | A * (B - x)**(-m), for x <= -beta + -- + """ + N = 1.0 / (m/beta / (m-1) * np.exp(-beta**2 / 2.0) + + _norm_pdf_C * _norm_cdf(beta)) + + def rhs(x, beta, m): + return np.exp(-x**2 / 2) + + def lhs(x, beta, m): + return ((m/beta)**m * np.exp(-beta**2 / 2.0) * + (m/beta - beta - x)**(-m)) + + return N * xpx.apply_where(x > -beta, (x, beta, m), rhs, lhs) + + def _logpdf(self, x, beta, m): + """ + Return the log of the PDF of the crystalball function. + """ + N = 1.0 / (m/beta / (m-1) * np.exp(-beta**2 / 2.0) + + _norm_pdf_C * _norm_cdf(beta)) + + def rhs(x, beta, m): + return -x**2/2 + + def lhs(x, beta, m): + return m*np.log(m/beta) - beta**2/2 - m*np.log(m/beta - beta - x) + + return np.log(N) + xpx.apply_where(x > -beta, (x, beta, m), rhs, lhs) + + def _cdf(self, x, beta, m): + """ + Return CDF of the crystalball function + """ + N = 1.0 / (m/beta / (m-1) * np.exp(-beta**2 / 2.0) + + _norm_pdf_C * _norm_cdf(beta)) + + def rhs(x, beta, m): + return ((m/beta) * np.exp(-beta**2 / 2.0) / (m-1) + + _norm_pdf_C * (_norm_cdf(x) - _norm_cdf(-beta))) + + def lhs(x, beta, m): + return ((m/beta)**m * np.exp(-beta**2 / 2.0) * + (m/beta - beta - x)**(-m+1) / (m-1)) + + return N * xpx.apply_where(x > -beta, (x, beta, m), rhs, lhs) + + def _sf(self, x, beta, m): + """ + Survival function of the crystalball distribution. + """ + + def rhs(x, beta, m): + # M is the same as 1/N used elsewhere. + M = m/beta/(m - 1)*np.exp(-beta**2/2) + _norm_pdf_C*_norm_cdf(beta) + return _norm_pdf_C*_norm_sf(x)/M + + def lhs(x, beta, m): + # Default behavior is OK in the left tail of the SF. + return 1 - self._cdf(x, beta, m) + + return xpx.apply_where(x > -beta, (x, beta, m), rhs, lhs) + + def _ppf(self, p, beta, m): + N = 1.0 / (m/beta / (m-1) * np.exp(-beta**2 / 2.0) + + _norm_pdf_C * _norm_cdf(beta)) + pbeta = N * (m/beta) * np.exp(-beta**2/2) / (m - 1) + + def ppf_less(p, beta, m): + eb2 = np.exp(-beta**2/2) + C = (m/beta) * eb2 / (m-1) + N = 1/(C + _norm_pdf_C * _norm_cdf(beta)) + return (m/beta - beta - + ((m - 1)*(m/beta)**(-m)/eb2*p/N)**(1/(1-m))) + + def ppf_greater(p, beta, m): + eb2 = np.exp(-beta**2/2) + C = (m/beta) * eb2 / (m-1) + N = 1/(C + _norm_pdf_C * _norm_cdf(beta)) + return _norm_ppf(_norm_cdf(-beta) + (1/_norm_pdf_C)*(p/N - C)) + + return xpx.apply_where(p < pbeta, (p, beta, m), ppf_less, ppf_greater) + + def _munp(self, n, beta, m): + """ + Returns the n-th non-central moment of the crystalball function. + """ + N = 1.0 / (m/beta / (m-1) * np.exp(-beta**2 / 2.0) + + _norm_pdf_C * _norm_cdf(beta)) + + def n_th_moment(n, beta, m): + """ + Returns n-th moment. Defined only if n+1 < m + Function cannot broadcast due to the loop over n + """ + A = (m/beta)**m * np.exp(-beta**2 / 2.0) + B = m/beta - beta + rhs = (2**((n-1)/2.0) * sc.gamma((n+1)/2) * + (1.0 + (-1)**n * sc.gammainc((n+1)/2, beta**2 / 2))) + lhs = np.zeros(rhs.shape) + for k in range(int(n) + 1): + lhs += (sc.binom(n, k) * B**(n-k) * (-1)**k / (m - k - 1) * + (m/beta)**(-m + k + 1)) + return A * lhs + rhs + + return N * xpx.apply_where(n + 1 < m, (n, beta, m), + np.vectorize(n_th_moment, otypes=[np.float64]), + fill_value=np.inf) + + +crystalball = crystalball_gen(name='crystalball', longname="A Crystalball Function") + + +def _argus_phi(chi): + """ + Utility function for the argus distribution used in the pdf, sf and + moment calculation. + Note that for all x > 0: + gammainc(1.5, x**2/2) = 2 * (_norm_cdf(x) - x * _norm_pdf(x) - 0.5). + This can be verified directly by noting that the cdf of Gamma(1.5) can + be written as erf(sqrt(x)) - 2*sqrt(x)*exp(-x)/sqrt(Pi). + We use gammainc instead of the usual definition because it is more precise + for small chi. + """ + return sc.gammainc(1.5, chi**2/2) / 2 + + +class argus_gen(rv_continuous): + r""" + Argus distribution. + + %(before_notes)s + + Notes + ----- + The probability density function for `argus` is: + + .. math:: + + f(x, \chi) = \frac{\chi^3}{\sqrt{2\pi} \Psi(\chi)} x \sqrt{1-x^2} + \exp(-\chi^2 (1 - x^2)/2) + + for :math:`0 < x < 1` and :math:`\chi > 0`, where + + .. math:: + + \Psi(\chi) = \Phi(\chi) - \chi \phi(\chi) - 1/2 + + with :math:`\Phi` and :math:`\phi` being the CDF and PDF of a standard + normal distribution, respectively. + + `argus` takes :math:`\chi` as shape a parameter. Details about sampling + from the ARGUS distribution can be found in [2]_. + + %(after_notes)s + + References + ---------- + .. [1] "ARGUS distribution", + https://en.wikipedia.org/wiki/ARGUS_distribution + .. [2] Christoph Baumgarten "Random variate generation by fast numerical + inversion in the varying parameter case." Research in Statistics, + vol. 1, 2023. :doi:`10.1080/27684520.2023.2279060` + + .. versionadded:: 0.19.0 + + %(example)s + """ + def _shape_info(self): + return [_ShapeInfo("chi", False, (0, np.inf), (False, False))] + + def _logpdf(self, x, chi): + # for x = 0 or 1, logpdf returns -np.inf + with np.errstate(divide='ignore'): + y = 1.0 - x*x + A = 3*np.log(chi) - _norm_pdf_logC - np.log(_argus_phi(chi)) + return A + np.log(x) + 0.5*np.log1p(-x*x) - chi**2 * y / 2 + + def _pdf(self, x, chi): + return np.exp(self._logpdf(x, chi)) + + def _cdf(self, x, chi): + return 1.0 - self._sf(x, chi) + + def _sf(self, x, chi): + return _argus_phi(chi * np.sqrt((1 - x)*(1 + x))) / _argus_phi(chi) + + def _rvs(self, chi, size=None, random_state=None): + chi = np.asarray(chi) + if chi.size == 1: + out = self._rvs_scalar(chi, numsamples=size, + random_state=random_state) + else: + shp, bc = _check_shape(chi.shape, size) + numsamples = int(np.prod(shp)) + out = np.empty(size) + it = np.nditer([chi], + flags=['multi_index'], + op_flags=[['readonly']]) + while not it.finished: + idx = tuple((it.multi_index[j] if not bc[j] else slice(None)) + for j in range(-len(size), 0)) + r = self._rvs_scalar(it[0], numsamples=numsamples, + random_state=random_state) + out[idx] = r.reshape(shp) + it.iternext() + + if size == (): + out = out[()] + return out + + def _rvs_scalar(self, chi, numsamples=None, random_state=None): + # if chi <= 1.8: + # use rejection method, see Devroye: + # Non-Uniform Random Variate Generation, 1986, section II.3.2. + # write: PDF f(x) = c * g(x) * h(x), where + # h is [0,1]-valued and g is a density + # we use two ways to write f + # + # Case 1: + # write g(x) = 3*x*sqrt(1-x**2), h(x) = exp(-chi**2 (1-x**2) / 2) + # If X has a distribution with density g its ppf G_inv is given by: + # G_inv(u) = np.sqrt(1 - u**(2/3)) + # + # Case 2: + # g(x) = chi**2 * x * exp(-chi**2 * (1-x**2)/2) / (1 - exp(-chi**2 /2)) + # h(x) = sqrt(1 - x**2), 0 <= x <= 1 + # one can show that + # G_inv(u) = np.sqrt(2*np.log(u*(np.exp(chi**2/2)-1)+1))/chi + # = np.sqrt(1 + 2*np.log(np.exp(-chi**2/2)*(1-u)+u)/chi**2) + # the latter expression is used for precision with small chi + # + # In both cases, the inverse cdf of g can be written analytically, and + # we can apply the rejection method: + # + # REPEAT + # Generate U uniformly distributed on [0, 1] + # Generate X with density g (e.g. via inverse transform sampling: + # X = G_inv(V) with V uniformly distributed on [0, 1]) + # UNTIL X <= h(X) + # RETURN X + # + # We use case 1 for chi <= 0.5 as it maintains precision for small chi + # and case 2 for 0.5 < chi <= 1.8 due to its speed for moderate chi. + # + # if chi > 1.8: + # use relation to the Gamma distribution: if X is ARGUS with parameter + # chi), then Y = chi**2 * (1 - X**2) / 2 has density proportional to + # sqrt(u) * exp(-u) on [0, chi**2 / 2], i.e. a Gamma(3/2) distribution + # conditioned on [0, chi**2 / 2]). Therefore, to sample X from the + # ARGUS distribution, we sample Y from the gamma distribution, keeping + # only samples on [0, chi**2 / 2], and apply the inverse + # transformation X = (1 - 2*Y/chi**2)**(1/2). Since we only + # look at chi > 1.8, gamma(1.5).cdf(chi**2/2) is large enough such + # Y falls in the interval [0, chi**2 / 2] with a high probability: + # stats.gamma(1.5).cdf(1.8**2/2) = 0.644... + # + # The points to switch between the different methods are determined + # by a comparison of the runtime of the different methods. However, + # the runtime is platform-dependent. The implemented values should + # ensure a good overall performance and are supported by an analysis + # of the rejection constants of different methods. + + size1d = tuple(np.atleast_1d(numsamples)) + N = int(np.prod(size1d)) + x = np.zeros(N) + simulated = 0 + chi2 = chi * chi + if chi <= 0.5: + d = -chi2 / 2 + while simulated < N: + k = N - simulated + u = random_state.uniform(size=k) + v = random_state.uniform(size=k) + z = v**(2/3) + # acceptance condition: u <= h(G_inv(v)). This simplifies to + accept = (np.log(u) <= d * z) + num_accept = np.sum(accept) + if num_accept > 0: + # we still need to transform z=v**(2/3) to X = G_inv(v) + rvs = np.sqrt(1 - z[accept]) + x[simulated:(simulated + num_accept)] = rvs + simulated += num_accept + elif chi <= 1.8: + echi = np.exp(-chi2 / 2) + while simulated < N: + k = N - simulated + u = random_state.uniform(size=k) + v = random_state.uniform(size=k) + z = 2 * np.log(echi * (1 - v) + v) / chi2 + # as in case one, simplify u <= h(G_inv(v)) and then transform + # z to the target distribution X = G_inv(v) + accept = (u**2 + z <= 0) + num_accept = np.sum(accept) + if num_accept > 0: + rvs = np.sqrt(1 + z[accept]) + x[simulated:(simulated + num_accept)] = rvs + simulated += num_accept + else: + # conditional Gamma for chi > 1.8 + while simulated < N: + k = N - simulated + g = random_state.standard_gamma(1.5, size=k) + accept = (g <= chi2 / 2) + num_accept = np.sum(accept) + if num_accept > 0: + x[simulated:(simulated + num_accept)] = g[accept] + simulated += num_accept + x = np.sqrt(1 - 2 * x / chi2) + + return np.reshape(x, size1d) + + def _stats(self, chi): + # need to ensure that dtype is float + # otherwise the mask below does not work for integers + chi = np.asarray(chi, dtype=float) + phi = _argus_phi(chi) + m = np.sqrt(np.pi/8) * chi * sc.ive(1, chi**2/4) / phi + # compute second moment, use Taylor expansion for small chi (<= 0.1) + mu2 = np.empty_like(chi) + mask = chi > 0.1 + c = chi[mask] + mu2[mask] = 1 - 3 / c**2 + c * _norm_pdf(c) / phi[mask] + c = chi[~mask] + coef = [-358/65690625, 0, -94/1010625, 0, 2/2625, 0, 6/175, 0, 0.4] + mu2[~mask] = np.polyval(coef, c) + return m, mu2 - m**2, None, None + + +argus = argus_gen(name='argus', longname="An Argus Function", a=0.0, b=1.0) + + +class rv_histogram(rv_continuous): + """ + Generates a distribution given by a histogram. + This is useful to generate a template distribution from a binned + datasample. + + As a subclass of the `rv_continuous` class, `rv_histogram` inherits from it + a collection of generic methods (see `rv_continuous` for the full list), + and implements them based on the properties of the provided binned + datasample. + + Parameters + ---------- + histogram : tuple of array_like + Tuple containing two array_like objects. + The first containing the content of n bins, + the second containing the (n+1) bin boundaries. + In particular, the return value of `numpy.histogram` is accepted. + + density : bool, optional + If False, assumes the histogram is proportional to counts per bin; + otherwise, assumes it is proportional to a density. + For constant bin widths, these are equivalent, but the distinction + is important when bin widths vary (see Notes). + If None (default), sets ``density=True`` for backwards compatibility, + but warns if the bin widths are variable. Set `density` explicitly + to silence the warning. + + .. versionadded:: 1.10.0 + + Notes + ----- + When a histogram has unequal bin widths, there is a distinction between + histograms that are proportional to counts per bin and histograms that are + proportional to probability density over a bin. If `numpy.histogram` is + called with its default ``density=False``, the resulting histogram is the + number of counts per bin, so ``density=False`` should be passed to + `rv_histogram`. If `numpy.histogram` is called with ``density=True``, the + resulting histogram is in terms of probability density, so ``density=True`` + should be passed to `rv_histogram`. To avoid warnings, always pass + ``density`` explicitly when the input histogram has unequal bin widths. + + There are no additional shape parameters except for the loc and scale. + The pdf is defined as a stepwise function from the provided histogram. + The cdf is a linear interpolation of the pdf. + + .. versionadded:: 0.19.0 + + Examples + -------- + + Create a scipy.stats distribution from a numpy histogram + + >>> import scipy.stats + >>> import numpy as np + >>> data = scipy.stats.norm.rvs(size=100000, loc=0, scale=1.5, + ... random_state=123) + >>> hist = np.histogram(data, bins=100) + >>> hist_dist = scipy.stats.rv_histogram(hist, density=False) + + Behaves like an ordinary scipy rv_continuous distribution + + >>> hist_dist.pdf(1.0) + 0.20538577847618705 + >>> hist_dist.cdf(2.0) + 0.90818568543056499 + + PDF is zero above (below) the highest (lowest) bin of the histogram, + defined by the max (min) of the original dataset + + >>> hist_dist.pdf(np.max(data)) + 0.0 + >>> hist_dist.cdf(np.max(data)) + 1.0 + >>> hist_dist.pdf(np.min(data)) + 7.7591907244498314e-05 + >>> hist_dist.cdf(np.min(data)) + 0.0 + + PDF and CDF follow the histogram + + >>> import matplotlib.pyplot as plt + >>> X = np.linspace(-5.0, 5.0, 100) + >>> fig, ax = plt.subplots() + >>> ax.set_title("PDF from Template") + >>> ax.hist(data, density=True, bins=100) + >>> ax.plot(X, hist_dist.pdf(X), label='PDF') + >>> ax.plot(X, hist_dist.cdf(X), label='CDF') + >>> ax.legend() + >>> fig.show() + + """ + _support_mask = rv_continuous._support_mask + + def __init__(self, histogram, *args, density=None, **kwargs): + """ + Create a new distribution using the given histogram + + Parameters + ---------- + histogram : tuple of array_like + Tuple containing two array_like objects. + The first containing the content of n bins, + the second containing the (n+1) bin boundaries. + In particular, the return value of np.histogram is accepted. + density : bool, optional + If False, assumes the histogram is proportional to counts per bin; + otherwise, assumes it is proportional to a density. + For constant bin widths, these are equivalent. + If None (default), sets ``density=True`` for backward + compatibility, but warns if the bin widths are variable. Set + `density` explicitly to silence the warning. + """ + self._histogram = histogram + self._density = density + if len(histogram) != 2: + raise ValueError("Expected length 2 for parameter histogram") + self._hpdf = np.asarray(histogram[0]) + self._hbins = np.asarray(histogram[1]) + if len(self._hpdf) + 1 != len(self._hbins): + raise ValueError("Number of elements in histogram content " + "and histogram boundaries do not match, " + "expected n and n+1.") + self._hbin_widths = self._hbins[1:] - self._hbins[:-1] + bins_vary = not np.allclose(self._hbin_widths, self._hbin_widths[0]) + if density is None and bins_vary: + message = ("Bin widths are not constant. Assuming `density=True`." + "Specify `density` explicitly to silence this warning.") + warnings.warn(message, RuntimeWarning, stacklevel=2) + density = True + elif not density: + self._hpdf = self._hpdf / self._hbin_widths + + self._hpdf = self._hpdf / float(np.sum(self._hpdf * self._hbin_widths)) + self._hcdf = np.cumsum(self._hpdf * self._hbin_widths) + self._hpdf = np.hstack([0.0, self._hpdf, 0.0]) + self._hcdf = np.hstack([0.0, self._hcdf]) + # Set support + kwargs['a'] = self.a = self._hbins[0] + kwargs['b'] = self.b = self._hbins[-1] + super().__init__(*args, **kwargs) + + def _pdf(self, x): + """ + PDF of the histogram + """ + return self._hpdf[np.searchsorted(self._hbins, x, side='right')] + + def _cdf(self, x): + """ + CDF calculated from the histogram + """ + return np.interp(x, self._hbins, self._hcdf) + + def _ppf(self, x): + """ + Percentile function calculated from the histogram + """ + return np.interp(x, self._hcdf, self._hbins) + + def _munp(self, n): + """Compute the n-th non-central moment.""" + integrals = (self._hbins[1:]**(n+1) - self._hbins[:-1]**(n+1)) / (n+1) + return np.sum(self._hpdf[1:-1] * integrals) + + def _entropy(self): + """Compute entropy of distribution""" + hpdf = self._hpdf[1:-1] + res = xpx.apply_where(hpdf > 0.0, hpdf, np.log, fill_value=0.0) + return -np.sum(hpdf * res * self._hbin_widths) + + def _updated_ctor_param(self): + """ + Set the histogram as additional constructor argument + """ + dct = super()._updated_ctor_param() + dct['histogram'] = self._histogram + dct['density'] = self._density + return dct + + +class studentized_range_gen(rv_continuous): + r"""A studentized range continuous random variable. + + %(before_notes)s + + See Also + -------- + t: Student's t distribution + + Notes + ----- + The probability density function for `studentized_range` is: + + .. math:: + + f(x; k, \nu) = \frac{k(k-1)\nu^{\nu/2}}{\Gamma(\nu/2) + 2^{\nu/2-1}} \int_{0}^{\infty} \int_{-\infty}^{\infty} + s^{\nu} e^{-\nu s^2/2} \phi(z) \phi(sx + z) + [\Phi(sx + z) - \Phi(z)]^{k-2} \,dz \,ds + + for :math:`x ≥ 0`, :math:`k > 1`, and :math:`\nu > 0`. + + `studentized_range` takes ``k`` for :math:`k` and ``df`` for :math:`\nu` + as shape parameters. + + When :math:`\nu` exceeds 100,000, an asymptotic approximation (infinite + degrees of freedom) is used to compute the cumulative distribution + function [4]_ and probability distribution function. + + %(after_notes)s + + References + ---------- + + .. [1] "Studentized range distribution", + https://en.wikipedia.org/wiki/Studentized_range_distribution + .. [2] Batista, Ben Dêivide, et al. "Externally Studentized Normal Midrange + Distribution." Ciência e Agrotecnologia, vol. 41, no. 4, 2017, pp. + 378-389., :doi:`10.1590/1413-70542017414047716`. + .. [3] Harter, H. Leon. "Tables of Range and Studentized Range." The Annals + of Mathematical Statistics, vol. 31, no. 4, 1960, pp. 1122-1147. + JSTOR, www.jstor.org/stable/2237810. Accessed 18 Feb. 2021. + .. [4] Lund, R. E., and J. R. Lund. "Algorithm AS 190: Probabilities and + Upper Quantiles for the Studentized Range." Journal of the Royal + Statistical Society. Series C (Applied Statistics), vol. 32, no. 2, + 1983, pp. 204-210. JSTOR, www.jstor.org/stable/2347300. Accessed 18 + Feb. 2021. + + Examples + -------- + >>> import numpy as np + >>> from scipy.stats import studentized_range + >>> import matplotlib.pyplot as plt + >>> fig, ax = plt.subplots(1, 1) + + Display the probability density function (``pdf``): + + >>> k, df = 3, 10 + >>> x = np.linspace(studentized_range.ppf(0.01, k, df), + ... studentized_range.ppf(0.99, k, df), 100) + >>> ax.plot(x, studentized_range.pdf(x, k, df), + ... 'r-', lw=5, alpha=0.6, label='studentized_range pdf') + + Alternatively, the distribution object can be called (as a function) + to fix the shape, location and scale parameters. This returns a "frozen" + RV object holding the given parameters fixed. + + Freeze the distribution and display the frozen ``pdf``: + + >>> rv = studentized_range(k, df) + >>> ax.plot(x, rv.pdf(x), 'k-', lw=2, label='frozen pdf') + + Check accuracy of ``cdf`` and ``ppf``: + + >>> vals = studentized_range.ppf([0.001, 0.5, 0.999], k, df) + >>> np.allclose([0.001, 0.5, 0.999], studentized_range.cdf(vals, k, df)) + True + + Rather than using (``studentized_range.rvs``) to generate random variates, + which is very slow for this distribution, we can approximate the inverse + CDF using an interpolator, and then perform inverse transform sampling + with this approximate inverse CDF. + + This distribution has an infinite but thin right tail, so we focus our + attention on the leftmost 99.9 percent. + + >>> a, b = studentized_range.ppf([0, .999], k, df) + >>> a, b + 0, 7.41058083802274 + + >>> from scipy.interpolate import interp1d + >>> rng = np.random.default_rng() + >>> xs = np.linspace(a, b, 50) + >>> cdf = studentized_range.cdf(xs, k, df) + # Create an interpolant of the inverse CDF + >>> ppf = interp1d(cdf, xs, fill_value='extrapolate') + # Perform inverse transform sampling using the interpolant + >>> r = ppf(rng.uniform(size=1000)) + + And compare the histogram: + + >>> ax.hist(r, density=True, histtype='stepfilled', alpha=0.2) + >>> ax.legend(loc='best', frameon=False) + >>> plt.show() + + """ + + def _argcheck(self, k, df): + return (k > 1) & (df > 0) + + def _shape_info(self): + ik = _ShapeInfo("k", False, (1, np.inf), (False, False)) + idf = _ShapeInfo("df", False, (0, np.inf), (False, False)) + return [ik, idf] + + def _fitstart(self, data): + # Default is k=1, but that is not a valid value of the parameter. + return super()._fitstart(data, args=(2, 1)) + + def _munp(self, K, k, df): + cython_symbol = '_studentized_range_moment' + _a, _b = self._get_support() + # all three of these are used to create a numpy array so they must + # be the same shape. + + def _single_moment(K, k, df): + log_const = _stats._studentized_range_pdf_logconst(k, df) + arg = [K, k, df, log_const] + usr_data = np.array(arg, float).ctypes.data_as(ctypes.c_void_p) + + llc = LowLevelCallable.from_cython(_stats, cython_symbol, usr_data) + + ranges = [(-np.inf, np.inf), (0, np.inf), (_a, _b)] + opts = dict(epsabs=1e-11, epsrel=1e-12) + + return integrate.nquad(llc, ranges=ranges, opts=opts)[0] + + ufunc = np.frompyfunc(_single_moment, 3, 1) + return np.asarray(ufunc(K, k, df), dtype=np.float64)[()] + + def _pdf(self, x, k, df): + + def _single_pdf(q, k, df): + # The infinite form of the PDF is derived from the infinite + # CDF. + if df < 100000: + cython_symbol = '_studentized_range_pdf' + log_const = _stats._studentized_range_pdf_logconst(k, df) + arg = [q, k, df, log_const] + usr_data = np.array(arg, float).ctypes.data_as(ctypes.c_void_p) + ranges = [(-np.inf, np.inf), (0, np.inf)] + + else: + cython_symbol = '_studentized_range_pdf_asymptotic' + arg = [q, k] + usr_data = np.array(arg, float).ctypes.data_as(ctypes.c_void_p) + ranges = [(-np.inf, np.inf)] + + llc = LowLevelCallable.from_cython(_stats, cython_symbol, usr_data) + opts = dict(epsabs=1e-11, epsrel=1e-12) + return integrate.nquad(llc, ranges=ranges, opts=opts)[0] + + ufunc = np.frompyfunc(_single_pdf, 3, 1) + return np.asarray(ufunc(x, k, df), dtype=np.float64)[()] + + def _cdf(self, x, k, df): + + def _single_cdf(q, k, df): + # "When the degrees of freedom V are infinite the probability + # integral takes [on a] simpler form," and a single asymptotic + # integral is evaluated rather than the standard double integral. + # (Lund, Lund, page 205) + if df < 100000: + cython_symbol = '_studentized_range_cdf' + log_const = _stats._studentized_range_cdf_logconst(k, df) + arg = [q, k, df, log_const] + usr_data = np.array(arg, float).ctypes.data_as(ctypes.c_void_p) + ranges = [(-np.inf, np.inf), (0, np.inf)] + + else: + cython_symbol = '_studentized_range_cdf_asymptotic' + arg = [q, k] + usr_data = np.array(arg, float).ctypes.data_as(ctypes.c_void_p) + ranges = [(-np.inf, np.inf)] + + llc = LowLevelCallable.from_cython(_stats, cython_symbol, usr_data) + opts = dict(epsabs=1e-11, epsrel=1e-12) + return integrate.nquad(llc, ranges=ranges, opts=opts)[0] + + ufunc = np.frompyfunc(_single_cdf, 3, 1) + + # clip p-values to ensure they are in [0, 1]. + return np.clip(np.asarray(ufunc(x, k, df), dtype=np.float64)[()], 0, 1) + + +studentized_range = studentized_range_gen(name='studentized_range', a=0, + b=np.inf) + + +class rel_breitwigner_gen(rv_continuous): + r"""A relativistic Breit-Wigner random variable. + + %(before_notes)s + + See Also + -------- + cauchy: Cauchy distribution, also known as the Breit-Wigner distribution. + + Notes + ----- + + The probability density function for `rel_breitwigner` is + + .. math:: + + f(x, \rho) = \frac{k}{(x^2 - \rho^2)^2 + \rho^2} + + where + + .. math:: + k = \frac{2\sqrt{2}\rho^2\sqrt{\rho^2 + 1}} + {\pi\sqrt{\rho^2 + \rho\sqrt{\rho^2 + 1}}} + + The relativistic Breit-Wigner distribution is used in high energy physics + to model resonances [1]_. It gives the uncertainty in the invariant mass, + :math:`M` [2]_, of a resonance with characteristic mass :math:`M_0` and + decay-width :math:`\Gamma`, where :math:`M`, :math:`M_0` and :math:`\Gamma` + are expressed in natural units. In SciPy's parametrization, the shape + parameter :math:`\rho` is equal to :math:`M_0/\Gamma` and takes values in + :math:`(0, \infty)`. + + Equivalently, the relativistic Breit-Wigner distribution is said to give + the uncertainty in the center-of-mass energy :math:`E_{\text{cm}}`. In + natural units, the speed of light :math:`c` is equal to 1 and the invariant + mass :math:`M` is equal to the rest energy :math:`Mc^2`. In the + center-of-mass frame, the rest energy is equal to the total energy [3]_. + + %(after_notes)s + + :math:`\rho = M/\Gamma` and :math:`\Gamma` is the scale parameter. For + example, if one seeks to model the :math:`Z^0` boson with :math:`M_0 + \approx 91.1876 \text{ GeV}` and :math:`\Gamma \approx 2.4952\text{ GeV}` + [4]_ one can set ``rho=91.1876/2.4952`` and ``scale=2.4952``. + + To ensure a physically meaningful result when using the `fit` method, one + should set ``floc=0`` to fix the location parameter to 0. + + References + ---------- + .. [1] Relativistic Breit-Wigner distribution, Wikipedia, + https://en.wikipedia.org/wiki/Relativistic_Breit-Wigner_distribution + .. [2] Invariant mass, Wikipedia, + https://en.wikipedia.org/wiki/Invariant_mass + .. [3] Center-of-momentum frame, Wikipedia, + https://en.wikipedia.org/wiki/Center-of-momentum_frame + .. [4] M. Tanabashi et al. (Particle Data Group) Phys. Rev. D 98, 030001 - + Published 17 August 2018 + + %(example)s + + """ + def _argcheck(self, rho): + return rho > 0 + + def _shape_info(self): + return [_ShapeInfo("rho", False, (0, np.inf), (False, False))] + + def _pdf(self, x, rho): + # C = k / rho**2 + C = np.sqrt( + 2 * (1 + 1/rho**2) / (1 + np.sqrt(1 + 1/rho**2)) + ) * 2 / np.pi + with np.errstate(over='ignore'): + return C / (((x - rho)*(x + rho)/rho)**2 + 1) + + def _cdf(self, x, rho): + # C = k / (2 * rho**2) / np.sqrt(1 + 1/rho**2) + C = np.sqrt(2/(1 + np.sqrt(1 + 1/rho**2)))/np.pi + result = ( + np.sqrt(-1 + 1j/rho) + * np.arctan(x/np.sqrt(-rho*(rho + 1j))) + ) + result = C * 2 * np.imag(result) + # Sometimes above formula produces values greater than 1. + return np.clip(result, None, 1) + + def _munp(self, n, rho): + if n == 0: + return 1. + if n == 1: + # C = k / (2 * rho) + C = np.sqrt( + 2 * (1 + 1/rho**2) / (1 + np.sqrt(1 + 1/rho**2)) + ) / np.pi * rho + return C * (np.pi/2 + np.arctan(rho)) + if n == 2: + # C = pi * k / (4 * rho) + C = np.sqrt( + (1 + 1/rho**2) / (2 * (1 + np.sqrt(1 + 1/rho**2))) + ) * rho + result = (1 - rho * 1j) / np.sqrt(-1 - 1j/rho) + return 2 * C * np.real(result) + else: + return np.inf + + def _stats(self, rho): + # Returning None from stats makes public stats use _munp. + # nan values will be omitted from public stats. Skew and + # kurtosis are actually infinite. + return None, None, np.nan, np.nan + + @inherit_docstring_from(rv_continuous) + def fit(self, data, *args, **kwds): + # Override rv_continuous.fit to better handle case where floc is set. + data, _, floc, fscale = _check_fit_input_parameters( + self, data, args, kwds + ) + + censored = isinstance(data, CensoredData) + if censored: + if data.num_censored() == 0: + # There are no censored values in data, so replace the + # CensoredData instance with a regular array. + data = data._uncensored + censored = False + + if floc is None or censored: + return super().fit(data, *args, **kwds) + + if fscale is None: + # The interquartile range approximates the scale parameter gamma. + # The median approximates rho * gamma. + p25, p50, p75 = np.quantile(data - floc, [0.25, 0.5, 0.75]) + scale_0 = p75 - p25 + rho_0 = p50 / scale_0 + if not args: + args = [rho_0] + if "scale" not in kwds: + kwds["scale"] = scale_0 + else: + M_0 = np.median(data - floc) + rho_0 = M_0 / fscale + if not args: + args = [rho_0] + return super().fit(data, *args, **kwds) + + +rel_breitwigner = rel_breitwigner_gen(a=0.0, name="rel_breitwigner") + + +# Collect names of classes and objects in this module. +pairs = list(globals().copy().items()) +_distn_names, _distn_gen_names = get_distribution_names(pairs, rv_continuous) + +__all__ = _distn_names + _distn_gen_names + ['rv_histogram'] diff --git a/v12x/pylib/scipy/stats/_correlation.py b/v12x/pylib/scipy/stats/_correlation.py new file mode 100644 index 0000000000000000000000000000000000000000..5c72e5eac9300911c8b6281eceb9ca91d3b03a4c --- /dev/null +++ b/v12x/pylib/scipy/stats/_correlation.py @@ -0,0 +1,731 @@ +import numpy as np +import math +from scipy import stats, special +from scipy._external import array_api_extra as xpx +from scipy._lib._array_api import (xp_capabilities, array_namespace, xp_promote, + is_numpy, _share_masks, _count_nonmasked, is_marray) +from scipy.stats._stats_py import (_SimpleNormal, SignificanceResult, _get_pvalue, + _rankdata) +from scipy.stats._axis_nan_policy import _axis_nan_policy_factory +from scipy.stats._stats_mstats_common import (TheilslopesResult, _n_samples_optional_x, + SiegelslopesResult) + + +__all__ = ['chatterjeexi', 'spearmanrho', 'theilslopes', 'siegelslopes'] + + +def _xi_statistic(x, y, y_continuous, xp): + # Compute xi correlation statistic + + # `axis=-1` is guaranteed by _axis_nan_policy decorator + n = x.shape[-1] + + # "Rearrange the data as (X(1), Y(1)), . . . ,(X(n), Y(n)) + # such that X(1) ≤ ··· ≤ X(n)" + j = xp.argsort(x, axis=-1) + j, y = xp.broadcast_arrays(j, y) + y = xp.take_along_axis(y, j, axis=-1) + + # "Let ri be the rank of Y(i), that is, the number of j such that Y(j) ≤ Y(i)" + r = stats.rankdata(y, method='max', axis=-1) + # " additionally define li to be the number of j such that Y(j) ≥ Y(i)" + # Could probably compute this from r, but that can be an enhancement + l = stats.rankdata(-y, method='max', axis=-1) + + num = xp.sum(xp.abs(xp.diff(r, axis=-1)), axis=-1) + if y_continuous: # [1] Eq. 1.1 + statistic = 1 - 3 * num / (n ** 2 - 1) + else: # [1] Eq. 1.2 + den = 2 * xp.sum((n - l) * l, axis=-1) + statistic = 1 - n * num / den + + return statistic, r, l + + +def _xi_std(r, l, y_continuous, xp): + # Compute asymptotic standard deviation of xi under null hypothesis of independence + + # `axis=-1` is guaranteed by _axis_nan_policy decorator + n = r.shape[-1] + + # "Suppose that X and Y are independent and Y is continuous. Then + # √n·ξn(X, Y) → N(0, 2/5) in distribution as n → ∞" + if y_continuous: # [1] Theorem 2.1 + return xp.asarray(math.sqrt(2 / 5) / math.sqrt(n), dtype=r.dtype) + + # "Suppose that X and Y are independent. Then √n·ξn(X, Y) + # converges to N(0, τ²) in distribution as n → ∞ + # [1] Eq. 2.2 and surrounding math + i = xp.arange(1, n + 1, dtype=r.dtype) + u = xp.sort(r, axis=-1) + v = xp.cumulative_sum(u, axis=-1) + an = 1 / n**4 * xp.sum((2*n - 2*i + 1) * u**2, axis=-1) + bn = 1 / n**5 * xp.sum((v + (n - i)*u)**2, axis=-1) + cn = 1 / n**3 * xp.sum((2*n - 2*i + 1) * u, axis=-1) + dn = 1 / n**3 * xp.sum((l * (n - l)), axis=-1) + tau2 = (an - 2*bn + cn**2) / dn**2 + + return xp.sqrt(tau2) / math.sqrt(n) + + +def _chatterjeexi_iv(y_continuous, method): + # Input validation for `chatterjeexi` + # x, y, `axis` input validation taken care of by decorator + + if y_continuous not in {True, False}: + raise ValueError('`y_continuous` must be boolean.') + + if not isinstance(method, stats.PermutationMethod): + method = method.lower() + message = "`method` must be 'asymptotic' or a `PermutationMethod` instance." + if method != 'asymptotic': + raise ValueError(message) + + return y_continuous, method + + +def _unpack(res, _): + return res.statistic, res.pvalue + + +@xp_capabilities(skip_backends=[('dask.array', 'no take_along_axis')]) +@_axis_nan_policy_factory(SignificanceResult, paired=True, n_samples=2, + result_to_tuple=_unpack, n_outputs=2, too_small=1) +def chatterjeexi(x, y, *, axis=0, y_continuous=False, method='asymptotic'): + r"""Compute the xi correlation and perform a test of independence. + + The xi correlation coefficient is a measure of association between two + variables; the value tends to be close to zero when the variables are + independent and close to 1 when there is a strong association. Unlike + other correlation coefficients, the xi correlation is effective even + when the association is not monotonic. + + Parameters + ---------- + x, y : array-like + The samples: corresponding observations of the independent and + dependent variable. The (N-d) arrays must be broadcastable. + axis : int, default: 0 + Axis along which to perform the test. + y_continuous : bool, default: False + Whether `y` is assumed to be drawn from a continuous distribution. + If `y` is drawn from a continuous distribution, results are valid + whether this is assumed or not, but enabling this assumption will + result in faster computation and typically produce similar results. + method : 'asymptotic' or `PermutationMethod` instance, optional + Selects the method used to calculate the *p*-value. + Default is 'asymptotic'. The following options are available. + + * ``'asymptotic'``: compares the standardized test statistic + against the normal distribution. + * `PermutationMethod` instance. In this case, the p-value + is computed using `permutation_test` with the provided + configuration options and other appropriate settings. + + Returns + ------- + res : SignificanceResult + An object containing attributes: + + statistic : float + The xi correlation statistic. + pvalue : float + The associated *p*-value: the probability of a statistic at least as + high as the observed value under the null hypothesis of independence. + + See Also + -------- + scipy.stats.pearsonr, scipy.stats.spearmanr, scipy.stats.kendalltau + + Notes + ----- + There is currently no special handling of ties in `x`; they are broken arbitrarily + by the implementation. [1]_ recommends: "if there are ties among the Xi's, then + choose an increasing rearrangement as above by breaking ties uniformly at random." + This is easily accomplished by adding a small amount of random noise to `x`; see + examples. + + [1]_ notes that the statistic is not symmetric in `x` and `y` *by design*: + "...we may want to understand if :math:`Y` is a function :math:`X`, and not just + if one of the variables is a function of the other." See [1]_ Remark 1. + + References + ---------- + .. [1] Chatterjee, Sourav. "A new coefficient of correlation." Journal of + the American Statistical Association 116.536 (2021): 2009-2022. + :doi:`10.1080/01621459.2020.1758115`. + + Examples + -------- + Generate perfectly correlated data, and observe that the xi correlation is + nearly 1.0. + + >>> import numpy as np + >>> from scipy import stats + >>> rng = np.random.default_rng(348932549825235) + >>> x = rng.uniform(0, 10, size=100) + >>> y = np.sin(x) + >>> res = stats.chatterjeexi(x, y) + >>> res.statistic + np.float64(0.9012901290129013) + + The probability of observing such a high value of the statistic under the + null hypothesis of independence is very low. + + >>> res.pvalue + np.float64(2.2206974648177804e-46) + + As noise is introduced, the correlation coefficient decreases. + + >>> noise = rng.normal(scale=[[0.1], [0.5], [1]], size=(3, 100)) + >>> res = stats.chatterjeexi(x, y + noise, axis=-1) + >>> res.statistic + array([0.79507951, 0.41824182, 0.16651665]) + + Because the distribution of `y` is continuous, it is valid to pass + ``y_continuous=True``. The statistic is identical, and the p-value + (not shown) is only slightly different. + + >>> stats.chatterjeexi(x, y + noise, y_continuous=True, axis=-1).statistic + array([0.79507951, 0.41824182, 0.16651665]) + + Consider a case in which there are ties in `x`. + + >>> x = rng.integers(10, size=1000) + >>> y = rng.integers(10, size=1000) + + [1]_ recommends breaking the ties uniformly at random. + + >>> d = rng.uniform(1e-5, size=x.size) + >>> res = stats.chatterjeexi(x + d, y) + >>> res.statistic + -0.029919991638798438 + + Since this gives a randomized estimate of the statistic, [1]_ also suggests + considering the average over all possibilities of breaking ties. This is + computationally infeasible when there are many ties, but a randomized estimate of + *this* quantity can be obtained by considering many random possibilities of breaking + ties. + + >>> d = rng.uniform(1e-5, size=(9999, x.size)) + >>> res = stats.chatterjeexi(x + d, y, axis=1) + >>> np.mean(res.statistic) + 0.001186895213756626 + + """ + xp = array_namespace(x, y) + + # x, y, `axis` input validation taken care of by decorator + # In fact, `axis` is guaranteed to be -1 + y_continuous, method = _chatterjeexi_iv(y_continuous, method) + x, y = xp_promote(x, y, force_floating=True, xp=xp) + + # A highly negative statistic is possible, e.g. + # x = np.arange(100.), y = (x % 2 == 0) + # Unclear whether we should expose `alternative`, though. + alternative = 'greater' + + if method == 'asymptotic': + xi, r, l = _xi_statistic(x, y, y_continuous, xp=xp) + std = _xi_std(r, l, y_continuous, xp=xp) + norm = _SimpleNormal() + pvalue = _get_pvalue(xi / std, norm, alternative=alternative, xp=xp) + elif isinstance(method, stats.PermutationMethod): + res = stats.permutation_test( + # Could be faster if we just permuted the ranks; for now, keep it simple. + data=(y,), + statistic=lambda y, axis: _xi_statistic(x, y, y_continuous, xp=xp)[0], + alternative=alternative, permutation_type='pairings', **method._asdict(), + axis=-1) # `axis=-1` is guaranteed by _axis_nan_policy decorator + + xi, pvalue = res.statistic, res.pvalue + + xi = xi[()] if xi.ndim == 0 else xi + pvalue = pvalue[()] if pvalue.ndim == 0 else pvalue + return SignificanceResult(xi, pvalue) + + +@xp_capabilities(cpu_only=True, exceptions=['jax.numpy'], marray=True, + skip_backends=[('dask.array', 'not supported by rankdata (take_along_axis)')], + extra_note='Only the default `method` is compatible with MArray input.' +) +@_axis_nan_policy_factory(SignificanceResult, paired=True, n_samples=2, + result_to_tuple=_unpack, n_outputs=2, too_small=1) +def spearmanrho(x, y, /, *, alternative='two-sided', method=None, axis=0): + r"""Calculate a Spearman rho correlation coefficient with associated p-value. + + The Spearman rank-order correlation coefficient is a nonparametric measure + of the monotonicity of the relationship between two datasets. + Like other correlation coefficients, it varies between -1 and +1 with 0 + implying no correlation. Coefficients of -1 or +1 are associated with an exact + monotonic relationship. Positive correlations indicate that as `x` increases, + so does `y`; negative correlations indicate that as `x` increases, `y` decreases. + The p-value is the probability of an uncorrelated system producing datasets + with a Spearman correlation at least as extreme as the one computed from the + observed dataset. + + Parameters + ---------- + x, y : array-like + The samples: corresponding observations of the independent and + dependent variable. The (N-d) arrays must be broadcastable. + alternative : {'two-sided', 'less', 'greater'}, optional + Defines the alternative hypothesis. Default is 'two-sided'. + The following options are available: + + * 'two-sided': the correlation is nonzero + * 'less': the correlation is negative (less than zero) + * 'greater': the correlation is positive (greater than zero) + + method : ResamplingMethod, optional + Defines the method used to compute the p-value. If `method` is an + instance of `PermutationMethod`/`MonteCarloMethod`, the p-value is + computed using + `scipy.stats.permutation_test`/`scipy.stats.monte_carlo_test` with the + provided configuration options and other appropriate settings. + Otherwise, the p-value is computed using an asymptotic approximation of + the null distribution. + axis : int or None, optional + If axis=0 (default), then each column represents a variable, with + observations in the rows. If axis=1, the relationship is transposed: + each row represents a variable, while the columns contain observations. + If axis=None, then both arrays will be raveled. + Like other `scipy.stats` functions, `axis` is interpreted after the + arrays are broadcasted. + + Returns + ------- + res : SignificanceResult + An object containing attributes: + + statistic : floating point array or NumPy scalar + Spearman correlation coefficient + pvalue : floating point array NumPy scalar + The p-value - the probabilitiy of realizing such an extreme statistic + value under the null hypothesis that two samples have no ordinal + correlation. See `alternative` above for alternative hypotheses. + + Warns + ----- + `~scipy.stats.ConstantInputWarning` + Raised if an input is a constant array. The correlation coefficient + is not defined in this case, so ``np.nan`` is returned. + + Notes + ----- + `spearmanrho` was created to make improvements to SciPy's implementation of + the Spearman correlation test without making backward-incompatible changes + to `spearmanr`. Advantages of `spearmanrho` over `spearmanr` include: + + - `spearmanrho` follows standard array broadcasting rules. + - `spearmanrho` is compatible with some non-NumPy arrays. + - `spearmanrho` can compute exact p-values, even in the presence of ties, + when an appropriate instance of `PermutationMethod` is provided via the + `method` argument. + + References + ---------- + .. [1] Zwillinger, D. and Kokoska, S. (2000). CRC Standard + Probability and Statistics Tables and Formulae. Chapman & Hall: New + York. 2000. + Section 14.7 + .. [2] Kendall, M. G. and Stuart, A. (1973). + The Advanced Theory of Statistics, Volume 2: Inference and Relationship. + Griffin. 1973. + Section 31.18 + + Examples + -------- + Univariate samples, approximate p-value. + + >>> import numpy as np + >>> from scipy import stats + >>> x = [1, 2, 3, 4, 5] + >>> y = [5, 6, 7, 8, 7] + >>> res = stats.spearmanrho(x, y) + >>> res.statistic + np.float64(0.8207826816681233) + >>> res.pvalue + np.float64(0.08858700531354405) + + Univariate samples, exact p-value. + + >>> res = stats.spearmanrho(x, y, method=stats.PermutationMethod()) + >>> res.statistic + np.float64(0.8207826816681233) + >>> res.pvalue + np.float64(0.13333333333333333) + + Batch of univariate samples, one vectorized call. + + >>> rng = np.random.default_rng(98145152315484) + >>> x2 = rng.standard_normal((2, 100)) + >>> y2 = rng.standard_normal((2, 100)) + >>> res = stats.spearmanrho(x2, y2, axis=-1) + >>> res.statistic + array([ 0.16585659, -0.12151215]) + >>> res.pvalue + array([0.0991155 , 0.22846869]) + + Bivariate samples using standard broadcasting rules. + + >>> res = stats.spearmanrho(x2[np.newaxis, :], x2[:, np.newaxis], axis=-1) + >>> res.statistic + array([[ 1. , -0.14670267], + [-0.14670267, 1. ]]) + >>> res.pvalue + array([[0. , 0.14526128], + [0.14526128, 0. ]]) + + """ + xp = array_namespace(x, y) + x, y = _share_masks(x, y, xp=xp) + rx = stats.rankdata(x, axis=axis) + ry = stats.rankdata(y, axis=axis) + res = stats.pearsonr(rx, ry, method=method, alternative=alternative, axis=axis) + return SignificanceResult(res.statistic, res.pvalue) + + +@xp_capabilities(skip_backends=[("dask.array", "no take_along_axis"), + ("jax.numpy", "non-concrete boolean indexing")], + marray=True) +@_axis_nan_policy_factory(TheilslopesResult, default_axis=None, n_outputs=4, + n_samples=_n_samples_optional_x, + result_to_tuple=lambda x, _: tuple(x), paired=True, + too_small=1) +def theilslopes(y, x=None, alpha=0.95, method='separate', *, axis=None): + r""" + Computes the Theil-Sen estimator for a set of points (x, y). + + `theilslopes` implements a method for robust linear regression. It + computes the slope as the median of all slopes between paired values. + + Parameters + ---------- + y : array_like + Dependent variable. + x : array_like or None, optional + Independent variable. If None, use ``arange(len(y))`` instead. + alpha : float, optional + Confidence degree between 0 and 1. Default is 95% confidence. + Note that `alpha` is symmetric around 0.5, i.e. both 0.1 and 0.9 are + interpreted as "find the 90% confidence interval". + method : {'joint', 'separate'}, optional + Method to be used for computing estimate for intercept. + Following methods are supported, + + * 'joint': Uses ``np.median(y - slope * x)`` as intercept. + * 'separate': Uses ``np.median(y) - slope * np.median(x)`` + as intercept. + + The default is 'separate'. + + .. versionadded:: 1.8.0 + + axis : int or tuple of ints, default: None + If an int or tuple of ints, the axis or axes of the input along which + to compute the statistic. The statistic of each axis-slice (e.g. row) + of the input will appear in a corresponding element of the output. + If ``None``, the input will be raveled before computing the statistic. + + Returns + ------- + result : ``TheilslopesResult`` instance + The return value is an object with the following attributes: + + slope : float + Theil slope. + intercept : float + Intercept of the Theil line. + low_slope : float + Lower bound of the confidence interval on `slope`. + high_slope : float + Upper bound of the confidence interval on `slope`. + + See Also + -------- + siegelslopes : a similar technique using repeated medians + + Notes + ----- + The implementation of `theilslopes` follows [1]_. The intercept is + not defined in [1]_, and here it is defined as ``median(y) - + slope*median(x)``, which is given in [3]_. Other definitions of + the intercept exist in the literature such as ``median(y - slope*x)`` + in [4]_. The approach to compute the intercept can be determined by the + parameter ``method``. A confidence interval for the intercept is not + given as this question is not addressed in [1]_. + + For compatibility with older versions of SciPy, the return value acts + like a ``namedtuple`` of length 4, with fields ``slope``, ``intercept``, + ``low_slope``, and ``high_slope``, so one can continue to write:: + + slope, intercept, low_slope, high_slope = theilslopes(y, x) + + References + ---------- + .. [1] P.K. Sen, "Estimates of the regression coefficient based on + Kendall's tau", J. Am. Stat. Assoc., Vol. 63, pp. 1379-1389, 1968. + .. [2] H. Theil, "A rank-invariant method of linear and polynomial + regression analysis I, II and III", Nederl. Akad. Wetensch., Proc. + 53:, pp. 386-392, pp. 521-525, pp. 1397-1412, 1950. + .. [3] W.L. Conover, "Practical nonparametric statistics", 2nd ed., + John Wiley and Sons, New York, pp. 493. + .. [4] https://en.wikipedia.org/wiki/Theil%E2%80%93Sen_estimator + + Examples + -------- + >>> import numpy as np + >>> from scipy import stats + >>> import matplotlib.pyplot as plt + + >>> x = np.linspace(-5, 5, num=150) + >>> y = x + np.random.normal(size=x.size) + >>> y[11:15] += 10 # add outliers + >>> y[-5:] -= 7 + + Compute the slope, intercept and 90% confidence interval. For comparison, + also compute the least-squares fit with `linregress`: + + >>> res = stats.theilslopes(y, x, 0.90, method='separate') + >>> lsq_res = stats.linregress(x, y) + + Plot the results. The Theil-Sen regression line is shown in red, with the + dashed red lines illustrating the confidence interval of the slope (note + that the dashed red lines are not the confidence interval of the regression + as the confidence interval of the intercept is not included). The green + line shows the least-squares fit for comparison. + + >>> fig = plt.figure() + >>> ax = fig.add_subplot(111) + >>> ax.plot(x, y, 'b.') + >>> ax.plot(x, res[1] + res[0] * x, 'r-') + >>> ax.plot(x, res[1] + res[2] * x, 'r--') + >>> ax.plot(x, res[1] + res[3] * x, 'r--') + >>> ax.plot(x, lsq_res[1] + lsq_res[0] * x, 'g-') + >>> plt.show() + + """ + return _robust_slopes(y, x=x, alpha=alpha, method=method, pfun='theilslopes') + + +@xp_capabilities(skip_backends=[("dask.array", "no take_along_axis"), + ("jax.numpy", "quantile needs lazy nan_policy")], + marray=True) +@_axis_nan_policy_factory(SiegelslopesResult, default_axis=None, n_outputs=2, + n_samples=_n_samples_optional_x, + result_to_tuple=lambda x, _: tuple(x), paired=True, + too_small=1) +def siegelslopes(y, x=None, method='hierarchical', *, axis=None): + r""" + Computes the Siegel estimator for a set of points (x, y). + + `siegelslopes` implements a method for robust linear regression + using repeated medians (see [1]_) to fit a line to the points (x, y). + The method is robust to outliers with an asymptotic breakdown point + of 50%. + + Parameters + ---------- + y : array_like + Dependent variable. + x : array_like or None, optional + Independent variable. If None, use ``arange(len(y))`` instead. + method : {'hierarchical', 'separate'} + If 'hierarchical', estimate the intercept using the estimated + slope ``slope`` (default option). + If 'separate', estimate the intercept independent of the estimated + slope. See Notes for details. + axis : int or tuple of ints, default: None + If an int or tuple of ints, the axis or axes of the input along which + to compute the statistic. The statistic of each axis-slice (e.g. row) + of the input will appear in a corresponding element of the output. + If ``None``, the input will be raveled before computing the statistic. + + Returns + ------- + result : ``SiegelslopesResult`` instance + The return value is an object with the following attributes: + + slope : float + Estimate of the slope of the regression line. + intercept : float + Estimate of the intercept of the regression line. + + See Also + -------- + theilslopes : a similar technique without repeated medians + + Notes + ----- + With ``n = len(y)``, compute ``m_j`` as the median of + the slopes of the lines from the point ``(x[j], y[j])`` to all other ``n-1`` points. + ``slope`` is then the median of all slopes ``m_j``. + Two ways are given to estimate the intercept in [1]_ which can be chosen + via the parameter ``method``. + The hierarchical approach uses the estimated slope ``slope`` + and computes ``intercept`` as the median of ``y - slope*x``. + The other approach estimates the intercept separately as follows: for + each point ``(x[j], y[j])``, compute the intercepts of all the ``n-1`` + lines through the remaining points and take the median ``i_j``. + ``intercept`` is the median of the ``i_j``. + + The implementation computes `n` times the median of a vector of size `n` + which can be slow for large vectors. There are more efficient algorithms + (see [2]_) which are not implemented here. + + For compatibility with older versions of SciPy, the return value acts + like a ``namedtuple`` of length 2, with fields ``slope`` and + ``intercept``, so one can continue to write:: + + slope, intercept = siegelslopes(y, x) + + References + ---------- + .. [1] A. Siegel, "Robust Regression Using Repeated Medians", + Biometrika, Vol. 69, pp. 242-244, 1982. + + .. [2] A. Stein and M. Werman, "Finding the repeated median regression + line", Proceedings of the Third Annual ACM-SIAM Symposium on + Discrete Algorithms, pp. 409-413, 1992. + + Examples + -------- + >>> import numpy as np + >>> from scipy import stats + >>> import matplotlib.pyplot as plt + + >>> x = np.linspace(-5, 5, num=150) + >>> y = x + np.random.normal(size=x.size) + >>> y[11:15] += 10 # add outliers + >>> y[-5:] -= 7 + + Compute the slope and intercept. For comparison, also compute the + least-squares fit with `linregress`: + + >>> res = stats.siegelslopes(y, x) + >>> lsq_res = stats.linregress(x, y) + + Plot the results. The Siegel regression line is shown in red. The green + line shows the least-squares fit for comparison. + + >>> fig = plt.figure() + >>> ax = fig.add_subplot(111) + >>> ax.plot(x, y, 'b.') + >>> ax.plot(x, res[1] + res[0] * x, 'r-') + >>> ax.plot(x, lsq_res[1] + lsq_res[0] * x, 'g-') + >>> plt.show() + + """ + return _robust_slopes(y, x=x, method=method, pfun='siegelslopes') + + +def _robust_slopes(y, *, x, alpha=None, method, pfun): + other_method = 'joint' if pfun == 'theilslopes' else 'hierarchical' + if method not in {other_method, 'separate'}: + raise ValueError(f"method must be either '{other_method}' or 'separate'. " + f"'{method}' is invalid.") + + xp = array_namespace(y, x) + y, x = xp_promote(y, x, force_floating=True, xp=xp) + x = xp.arange(y.shape[-1], dtype=y.dtype) if x is None else x + y, x = xp.broadcast_arrays(y, x) + x, y = _share_masks(x, y, xp=xp) + + if x.shape[-1] < 2 or y.shape[-1] < 2: # only needed by test_axis_nan_policy + raise ValueError("`x` and `y` must have length at least 2.") + + # Compute sorted slopes only when deltax > 0 + deltax = x[..., :, xp.newaxis] - x[..., xp.newaxis, :] + deltay = y[..., :, xp.newaxis] - y[..., xp.newaxis, :] + + if pfun == 'theilslopes': + i = xp.astype(xp.triu(xp.ones(deltax.shape[-2:]), k=1), xp.bool) + if is_numpy(xp): + deltax, deltay = deltax[..., i], deltay[..., i] + else: + # With array API, mask must be sole index, so we need to broadcast it. + # Indexing ravels the array, so we need to reshape the results. + i = xp.broadcast_to(i, deltax.shape) + deltax, deltay = deltax[i], deltay[i] + deltax = xp.reshape(deltax, (*i.shape[:-2], -1)) + deltay = xp.reshape(deltay, (*i.shape[:-2], -1)) + + # Use `.data` to avoid indexing with masked indices. If masked elements of deltax=0, + # they can safely be set to xp.nan, too - they're masked, after all. + deltax0 = (deltax.data == 0) if is_marray(xp) else (deltax == 0) + deltax = xpx.at(deltax)[deltax0].set(xp.nan) + slopes = deltay / deltax + + if is_marray(xp): + def nanmedian(x, axis): # use mask to ignore NaNs + x_nans_masked = xp.asarray(x.data, mask=x.mask | x._xp.isnan(x.data)) + return stats.quantile(x_nans_masked, 0.5, axis=axis) + else: + def nanmedian(x, axis): # use nan_policy to ignore NaNs + return stats.quantile(x, 0.5, axis=axis, nan_policy='omit') + + def median(x, axis): return stats.quantile(x, 0.5, axis=axis) + + # `theilslopes` is median of all slopes. Indexing with `i` above has already raveled + # all the slopes, so we only need to take the median along the last axis. + # `siegelslope` is a median of medians: we take the median of slopes from point i + # to all other points, then the median of those medians. + # The slope is NaN wherever the two points are the same, and those don't contribute, + # hence the first median omits NaNs. NaNs in the input are propagated by the + # `_axis_nan_policy` decorator. + medslope = (median(nanmedian(slopes, axis=-1), axis=-1) if pfun == 'siegelslopes' + else nanmedian(slopes, axis=-1)) + + if method in {'joint', 'hierarchical'}: + medinter = median(y - medslope[..., np.newaxis] * x, axis=-1) + elif pfun == 'theilslopes': + medinter = median(y, axis=-1) - medslope * median(x, axis=-1) + else: + # Calculate pairwise intercepts given each point (row i) and the slope to each + # other point (column j). Then calculate the median of (row) medians. + intercepts = y[..., :, xp.newaxis] - slopes*x[..., :, xp.newaxis] + medinter = median(nanmedian(intercepts, axis=-1), axis=-1) + + if pfun == 'siegelslopes': + return SiegelslopesResult(slope=medslope[()], intercept=medinter[()]) + + # Now compute confidence intervals + if alpha > 0.5: + alpha = 1. - alpha + + z = float(special.ndtri(alpha / 2.)) + # This implements (2.6) from Sen (1968) + # we don't actually need ranks, so an enhancement could be to have + # `rankdata` return only the third output. In the meantime, use the + # least expensive `method`. + _, _, nxreps = _rankdata(x, method='min', return_ties=True) + _, _, nyreps = _rankdata(y, method='min', return_ties=True) + nt = xp.count_nonzero(xp.isfinite(slopes), axis=-1, keepdims=True) # N in Sen 1968 + nt = xp.asarray(nt, dtype=y.dtype) + ny = _count_nonmasked(y, keepdims=True, axis=-1) # n in Sen 1968 + # Equation 2.6 in Sen (1968): + sigsq = 1/18. * ( + ny * (ny-1) * (2*ny+5) + - xp.sum(nxreps * (nxreps-1) * (2*nxreps + 5), axis=-1, keepdims=True) + - xp.sum(nyreps * (nyreps-1) * (2*nyreps + 5), axis=-1, keepdims=True)) + # Find the confidence interval indices in `slopes` + sigma = xp.sqrt(xp.maximum(sigsq, xp.zeros_like(sigsq))) + Ru = xp.minimum(xp.astype(xp.round((nt - z*sigma)/2.), xp.int64), + xp.astype(nt, xp.int64)-1) + Rl = xp.maximum(xp.astype(xp.round((nt + z*sigma)/2.), xp.int64) - 1, + xp.asarray(0, dtype=xp.int64)) + R = xp.concat((xpx.atleast_nd(Rl, ndim=1), xpx.atleast_nd(Ru, ndim=1)), axis=-1) + slopes = xp.sort(slopes, axis=-1) + delta = xp.take_along_axis(slopes, R, axis=-1) + i_nan = xp.broadcast_to(sigsq < 0, delta.shape) + delta = xpx.at(delta)[i_nan].set(xp.nan) + + slope = medslope[()] if medslope.ndim == 0 else medslope + intercept = medinter[()] if medinter.ndim == 0 else medinter + low_slope = delta[..., 0] + high_slope = delta[..., 1] + low_slope = low_slope[()] if low_slope.ndim == 0 else low_slope + high_slope = high_slope[()] if high_slope.ndim == 0 else high_slope + return TheilslopesResult(slope=slope, intercept=intercept, + low_slope=low_slope, high_slope=high_slope) diff --git a/v12x/pylib/scipy/stats/_covariance.py b/v12x/pylib/scipy/stats/_covariance.py new file mode 100644 index 0000000000000000000000000000000000000000..f2133b9f3278d19fbb94fd6b8ba74876078dc6dd --- /dev/null +++ b/v12x/pylib/scipy/stats/_covariance.py @@ -0,0 +1,672 @@ +from functools import cached_property +from types import GenericAlias + +import numpy as np +from scipy import linalg +from scipy.stats import _multivariate + + +__all__ = ["Covariance"] + + +class Covariance: + """ + Representation of a covariance matrix. + + Calculations involving covariance matrices (e.g. data whitening, + multivariate normal function evaluation) are often performed more + efficiently using a decomposition of the covariance matrix instead of the + covariance matrix itself. This class allows the user to construct an + object representing a covariance matrix using any of several + decompositions and perform calculations using a common interface. + + .. note:: + + The `Covariance` class cannot be instantiated directly. Instead, use + one of the factory methods (e.g. `Covariance.from_diagonal`). + + Examples + -------- + The `Covariance` class is used by calling one of its + factory methods to create a `Covariance` object, then pass that + representation of the `Covariance` matrix as a shape parameter of a + multivariate distribution. + + For instance, the multivariate normal distribution can accept an array + representing a covariance matrix: + + >>> from scipy import stats + >>> import numpy as np + >>> d = [1, 2, 3] + >>> A = np.diag(d) # a diagonal covariance matrix + >>> x = [4, -2, 5] # a point of interest + >>> dist = stats.multivariate_normal(mean=[0, 0, 0], cov=A) + >>> dist.pdf(x) + 4.9595685102808205e-08 + + but the calculations are performed in a very generic way that does not + take advantage of any special properties of the covariance matrix. Because + our covariance matrix is diagonal, we can use ``Covariance.from_diagonal`` + to create an object representing the covariance matrix, and + `multivariate_normal` can use this to compute the probability density + function more efficiently. + + >>> cov = stats.Covariance.from_diagonal(d) + >>> dist = stats.multivariate_normal(mean=[0, 0, 0], cov=cov) + >>> dist.pdf(x) + 4.9595685102808205e-08 + + """ + + # generic type compatibility with scipy-stubs + __class_getitem__: classmethod = classmethod(GenericAlias) + + def __init__(self): + message = ("The `Covariance` class cannot be instantiated directly. " + "Please use one of the factory methods " + "(e.g. `Covariance.from_diagonal`).") + raise NotImplementedError(message) + + @staticmethod + def from_diagonal(diagonal): + r""" + Return a representation of a covariance matrix from its diagonal. + + Parameters + ---------- + diagonal : array_like + The diagonal elements of a diagonal matrix. + + Returns + ------- + Covariance + A `Covariance` object representing `diagonal`. + + Notes + ----- + Let the diagonal elements of a diagonal covariance matrix :math:`D` be + stored in the vector :math:`d`. + + When all elements of :math:`d` are strictly positive, whitening of a + data point :math:`x` is performed by computing + :math:`x \cdot d^{-1/2}`, where the inverse square root can be taken + element-wise. + :math:`\log\det{D}` is calculated as :math:`-2 \sum(\log{d})`, + where the :math:`\log` operation is performed element-wise. + + This `Covariance` class supports singular covariance matrices. When + computing ``_log_pdet``, non-positive elements of :math:`d` are + ignored. Whitening is not well defined when the point to be whitened + does not lie in the span of the columns of the covariance matrix. The + convention taken here is to treat the inverse square root of + non-positive elements of :math:`d` as zeros. + + Examples + -------- + Prepare a symmetric positive definite covariance matrix ``A`` and a + data point ``x``. + + >>> import numpy as np + >>> from scipy import stats + >>> rng = np.random.default_rng() + >>> n = 5 + >>> A = np.diag(rng.random(n)) + >>> x = rng.random(size=n) + + Extract the diagonal from ``A`` and create the `Covariance` object. + + >>> d = np.diag(A) + >>> cov = stats.Covariance.from_diagonal(d) + + Compare the functionality of the `Covariance` object against a + reference implementations. + + >>> res = cov.whiten(x) + >>> ref = np.diag(d**-0.5) @ x + >>> np.allclose(res, ref) + True + >>> res = cov.log_pdet + >>> ref = np.linalg.slogdet(A)[-1] + >>> np.allclose(res, ref) + True + + """ + return CovViaDiagonal(diagonal) + + @staticmethod + def from_precision(precision, covariance=None): + r""" + Return a representation of a covariance from its precision matrix. + + Parameters + ---------- + precision : array_like + The precision matrix; that is, the inverse of a square, symmetric, + positive definite covariance matrix. + covariance : array_like, optional + The square, symmetric, positive definite covariance matrix. If not + provided, this may need to be calculated (e.g. to evaluate the + cumulative distribution function of + `scipy.stats.multivariate_normal`) by inverting `precision`. + + Returns + ------- + Covariance + A `Covariance` object representing the covariance matrix + corresponding to `precision`. + + Notes + ----- + Let the covariance matrix be :math:`A`, its precision matrix be + :math:`P = A^{-1}`, and :math:`L` be the lower Cholesky factor such + that :math:`L L^T = P`. + Whitening of a data point :math:`x` is performed by computing + :math:`x^T L`. :math:`\log\det{A}` is calculated as + :math:`-2tr(\log{L})`, where the :math:`\log` operation is performed + element-wise. + + This `Covariance` class does not support singular covariance matrices + because the precision matrix does not exist for a singular covariance + matrix. + + Examples + -------- + Prepare a symmetric positive definite precision matrix ``P`` and a + data point ``x``. (If the precision matrix is not already available, + consider the other factory methods of the `Covariance` class.) + + >>> import numpy as np + >>> from scipy import stats + >>> rng = np.random.default_rng() + >>> n = 5 + >>> P = rng.random(size=(n, n)) + >>> P = P @ P.T # a precision matrix must be positive definite + >>> x = rng.random(size=n) + + Create the `Covariance` object. + + >>> cov = stats.Covariance.from_precision(P) + + Compare the functionality of the `Covariance` object against + reference implementations. + + >>> res = cov.whiten(x) + >>> ref = x @ np.linalg.cholesky(P) + >>> np.allclose(res, ref) + True + >>> res = cov.log_pdet + >>> ref = -np.linalg.slogdet(P)[-1] + >>> np.allclose(res, ref) + True + + """ + return CovViaPrecision(precision, covariance) + + @staticmethod + def from_cholesky(cholesky): + r""" + Representation of a covariance provided via the (lower) Cholesky factor. + + Parameters + ---------- + cholesky : array_like + The lower triangular Cholesky factor of the covariance matrix. + + Returns + ------- + Covariance + A `Covariance` object representing the covariance matrix + corresponding to `cholesky`. + + Notes + ----- + Let the covariance matrix be :math:`A` and :math:`L` be the lower + Cholesky factor such that :math:`L L^T = A`. + Whitening of a data point :math:`x` is performed by computing + :math:`L^{-1} x`. :math:`\log\det{A}` is calculated as + :math:`2tr(\log{L})`, where the :math:`\log` operation is performed + element-wise. + + This `Covariance` class does not support singular covariance matrices + because the Cholesky decomposition does not exist for a singular + covariance matrix. + + Examples + -------- + Prepare a symmetric positive definite covariance matrix ``A`` and a + data point ``x``. + + >>> import numpy as np + >>> from scipy import stats + >>> rng = np.random.default_rng() + >>> n = 5 + >>> A = rng.random(size=(n, n)) + >>> A = A @ A.T # make the covariance symmetric positive definite + >>> x = rng.random(size=n) + + Perform the Cholesky decomposition of ``A`` and create the + `Covariance` object. + + >>> L = np.linalg.cholesky(A) + >>> cov = stats.Covariance.from_cholesky(L) + + Compare the functionality of the `Covariance` object against + reference implementation. + + >>> from scipy.linalg import solve_triangular + >>> res = cov.whiten(x) + >>> ref = solve_triangular(L, x, lower=True) + >>> np.allclose(res, ref) + True + >>> res = cov.log_pdet + >>> ref = np.linalg.slogdet(A)[-1] + >>> np.allclose(res, ref) + True + + """ + return CovViaCholesky(cholesky) + + @staticmethod + def from_eigendecomposition(eigendecomposition): + r""" + Representation of a covariance provided via eigendecomposition. + + Parameters + ---------- + eigendecomposition : sequence + A sequence (nominally a tuple) containing the eigenvalue and + eigenvector arrays as computed by `scipy.linalg.eigh` or + `numpy.linalg.eigh`. + + Returns + ------- + Covariance + A `Covariance` object representing the covariance matrix + corresponding to `eigendecomposition`. + + Notes + ----- + Let the covariance matrix be :math:`A`, let :math:`V` be matrix of + eigenvectors, and let :math:`W` be the diagonal matrix of eigenvalues + such that `V W V^T = A`. + + When all of the eigenvalues are strictly positive, whitening of a + data point :math:`x` is performed by computing + :math:`x^T (V W^{-1/2})`, where the inverse square root can be taken + element-wise. + :math:`\log\det{A}` is calculated as :math:`tr(\log{W})`, + where the :math:`\log` operation is performed element-wise. + + This `Covariance` class supports singular covariance matrices. When + computing ``_log_pdet``, non-positive eigenvalues are ignored. + Whitening is not well defined when the point to be whitened + does not lie in the span of the columns of the covariance matrix. The + convention taken here is to treat the inverse square root of + non-positive eigenvalues as zeros. + + Examples + -------- + Prepare a symmetric positive definite covariance matrix ``A`` and a + data point ``x``. + + >>> import numpy as np + >>> from scipy import stats + >>> rng = np.random.default_rng() + >>> n = 5 + >>> A = rng.random(size=(n, n)) + >>> A = A @ A.T # make the covariance symmetric positive definite + >>> x = rng.random(size=n) + + Perform the eigendecomposition of ``A`` and create the `Covariance` + object. + + >>> w, v = np.linalg.eigh(A) + >>> cov = stats.Covariance.from_eigendecomposition((w, v)) + + Compare the functionality of the `Covariance` object against + reference implementations. + + >>> res = cov.whiten(x) + >>> ref = x @ (v @ np.diag(w**-0.5)) + >>> np.allclose(res, ref) + True + >>> res = cov.log_pdet + >>> ref = np.linalg.slogdet(A)[-1] + >>> np.allclose(res, ref) + True + + """ + return CovViaEigendecomposition(eigendecomposition) + + def whiten(self, x): + """ + Perform a whitening transformation on data. + + "Whitening" ("white" as in "white noise", in which each frequency has + equal magnitude) transforms a set of random variables into a new set of + random variables with unit-diagonal covariance. When a whitening + transform is applied to a sample of points distributed according to + a multivariate normal distribution with zero mean, the covariance of + the transformed sample is approximately the identity matrix. + + Parameters + ---------- + x : array_like + An array of points. The last dimension must correspond with the + dimensionality of the space, i.e., the number of columns in the + covariance matrix. + + Returns + ------- + x_ : array_like + The transformed array of points. + + References + ---------- + .. [1] "Whitening Transformation". Wikipedia. + https://en.wikipedia.org/wiki/Whitening_transformation + .. [2] Novak, Lukas, and Miroslav Vorechovsky. "Generalization of + coloring linear transformation". Transactions of VSB 18.2 + (2018): 31-35. :doi:`10.31490/tces-2018-0013` + + Examples + -------- + >>> import numpy as np + >>> from scipy import stats + >>> rng = np.random.default_rng() + >>> n = 3 + >>> A = rng.random(size=(n, n)) + >>> cov_array = A @ A.T # make matrix symmetric positive definite + >>> precision = np.linalg.inv(cov_array) + >>> cov_object = stats.Covariance.from_precision(precision) + >>> x = rng.multivariate_normal(np.zeros(n), cov_array, size=(10000)) + >>> x_ = cov_object.whiten(x) + >>> np.cov(x_, rowvar=False) # near-identity covariance + array([[0.97862122, 0.00893147, 0.02430451], + [0.00893147, 0.96719062, 0.02201312], + [0.02430451, 0.02201312, 0.99206881]]) + + """ + return self._whiten(np.asarray(x)) + + def colorize(self, x): + """ + Perform a colorizing transformation on data. + + "Colorizing" ("color" as in "colored noise", in which different + frequencies may have different magnitudes) transforms a set of + uncorrelated random variables into a new set of random variables with + the desired covariance. When a coloring transform is applied to a + sample of points distributed according to a multivariate normal + distribution with identity covariance and zero mean, the covariance of + the transformed sample is approximately the covariance matrix used + in the coloring transform. + + Parameters + ---------- + x : array_like + An array of points. The last dimension must correspond with the + dimensionality of the space, i.e., the number of columns in the + covariance matrix. + + Returns + ------- + x_ : array_like + The transformed array of points. + + References + ---------- + .. [1] "Whitening Transformation". Wikipedia. + https://en.wikipedia.org/wiki/Whitening_transformation + .. [2] Novak, Lukas, and Miroslav Vorechovsky. "Generalization of + coloring linear transformation". Transactions of VSB 18.2 + (2018): 31-35. :doi:`10.31490/tces-2018-0013` + + Examples + -------- + >>> import numpy as np + >>> from scipy import stats + >>> rng = np.random.default_rng(1638083107694713882823079058616272161) + >>> n = 3 + >>> A = rng.random(size=(n, n)) + >>> cov_array = A @ A.T # make matrix symmetric positive definite + >>> cholesky = np.linalg.cholesky(cov_array) + >>> cov_object = stats.Covariance.from_cholesky(cholesky) + >>> x = rng.multivariate_normal(np.zeros(n), np.eye(n), size=(10000)) + >>> x_ = cov_object.colorize(x) + >>> cov_data = np.cov(x_, rowvar=False) + >>> np.allclose(cov_data, cov_array, rtol=3e-2) + True + """ + return self._colorize(np.asarray(x)) + + @property + def log_pdet(self): + """ + Log of the pseudo-determinant of the covariance matrix + """ + return np.array(self._log_pdet, dtype=float)[()] + + @property + def rank(self): + """ + Rank of the covariance matrix + """ + return np.array(self._rank, dtype=int)[()] + + @property + def covariance(self): + """ + Explicit representation of the covariance matrix + """ + return self._covariance + + @property + def shape(self): + """ + Shape of the covariance array + """ + return self._shape + + def _validate_matrix(self, A, name): + A = np.atleast_2d(A) + m, n = A.shape[-2:] + if m != n or A.ndim != 2 or not (np.issubdtype(A.dtype, np.integer) or + np.issubdtype(A.dtype, np.floating)): + message = (f"The input `{name}` must be a square, " + "two-dimensional array of real numbers.") + raise ValueError(message) + return A + + def _validate_vector(self, A, name): + A = np.atleast_1d(A) + if A.ndim != 1 or not (np.issubdtype(A.dtype, np.integer) or + np.issubdtype(A.dtype, np.floating)): + message = (f"The input `{name}` must be a one-dimensional array " + "of real numbers.") + raise ValueError(message) + return A + + +class CovViaPrecision(Covariance): + + __class_getitem__ = None # type:ignore[assignment] + + def __init__(self, precision, covariance=None): + precision = self._validate_matrix(precision, 'precision') + if covariance is not None: + covariance = self._validate_matrix(covariance, 'covariance') + message = "`precision.shape` must equal `covariance.shape`." + if precision.shape != covariance.shape: + raise ValueError(message) + + self._chol_P = np.linalg.cholesky(precision) + self._log_pdet = -2*np.log(np.diag(self._chol_P)).sum(axis=-1) + self._rank = precision.shape[-1] # must be full rank if invertible + self._precision = precision + self._cov_matrix = covariance + self._shape = precision.shape + self._allow_singular = False + + def _whiten(self, x): + return x @ self._chol_P + + @cached_property + def _covariance(self): + n = self._shape[-1] + return (linalg.cho_solve((self._chol_P, True), np.eye(n)) + if self._cov_matrix is None else self._cov_matrix) + + def _colorize(self, x): + m = x.T.shape[0] + res = linalg.solve_triangular(self._chol_P.T, x.T.reshape(m, -1), lower=False) + return res.reshape(x.T.shape).T + + +def _dot_diag(x, d): + # If d were a full diagonal matrix, x @ d would always do what we want. + # Special treatment is needed for n-dimensional `d` in which each row + # includes only the diagonal elements of a covariance matrix. + return x * d if x.ndim < 2 else x * np.expand_dims(d, -2) + + +class CovViaDiagonal(Covariance): + + def __init__(self, diagonal): + diagonal = self._validate_vector(diagonal, 'diagonal') + + i_zero = diagonal <= 0 + positive_diagonal = np.array(diagonal, dtype=np.float64) + + positive_diagonal[i_zero] = 1 # ones don't affect determinant + self._log_pdet = np.sum(np.log(positive_diagonal), axis=-1) + + psuedo_reciprocals = 1 / np.sqrt(positive_diagonal) + psuedo_reciprocals[i_zero] = 0 + + self._sqrt_diagonal = np.sqrt(diagonal) + self._LP = psuedo_reciprocals + self._rank = positive_diagonal.shape[-1] - i_zero.sum(axis=-1) + self._covariance = np.apply_along_axis(np.diag, -1, diagonal) + self._i_zero = i_zero + self._shape = self._covariance.shape + self._allow_singular = True + + def _whiten(self, x): + return _dot_diag(x, self._LP) + + def _colorize(self, x): + return _dot_diag(x, self._sqrt_diagonal) + + def _support_mask(self, x): + """ + Check whether x lies in the support of the distribution. + """ + return ~np.any(_dot_diag(x, self._i_zero), axis=-1) + + +class CovViaCholesky(Covariance): + + __class_getitem__ = None # type:ignore[assignment] + + def __init__(self, cholesky): + L = self._validate_matrix(cholesky, 'cholesky') + + self._factor = L + self._log_pdet = 2*np.log(np.diag(self._factor)).sum(axis=-1) + self._rank = L.shape[-1] # must be full rank for cholesky + self._shape = L.shape + self._allow_singular = False + + @cached_property + def _covariance(self): + return self._factor @ self._factor.T + + def _whiten(self, x): + m = x.T.shape[0] + res = linalg.solve_triangular(self._factor, x.T.reshape(m, -1), lower=True) + return res.reshape(x.T.shape).T + + def _colorize(self, x): + return x @ self._factor.T + + +class CovViaEigendecomposition(Covariance): + + __class_getitem__ = None # type:ignore[assignment] + + def __init__(self, eigendecomposition): + eigenvalues, eigenvectors = eigendecomposition + eigenvalues = self._validate_vector(eigenvalues, 'eigenvalues') + eigenvectors = self._validate_matrix(eigenvectors, 'eigenvectors') + message = ("The shapes of `eigenvalues` and `eigenvectors` " + "must be compatible.") + try: + eigenvalues = np.expand_dims(eigenvalues, -2) + eigenvectors, eigenvalues = np.broadcast_arrays(eigenvectors, + eigenvalues) + eigenvalues = eigenvalues[..., 0, :] + except ValueError: + raise ValueError(message) + + i_zero = eigenvalues <= 0 + positive_eigenvalues = np.array(eigenvalues, dtype=np.float64) + + positive_eigenvalues[i_zero] = 1 # ones don't affect determinant + self._log_pdet = np.sum(np.log(positive_eigenvalues), axis=-1) + + psuedo_reciprocals = 1 / np.sqrt(positive_eigenvalues) + psuedo_reciprocals[i_zero] = 0 + + self._LP = eigenvectors * psuedo_reciprocals + self._LA = eigenvectors * np.sqrt(eigenvalues) + self._rank = positive_eigenvalues.shape[-1] - i_zero.sum(axis=-1) + self._w = eigenvalues + self._v = eigenvectors + self._shape = eigenvectors.shape + self._null_basis = eigenvectors * i_zero + # This is only used for `_support_mask`, not to decide whether + # the covariance is singular or not. + self._eps = _multivariate._eigvalsh_to_eps(eigenvalues) * 10**3 + self._allow_singular = True + + def _whiten(self, x): + return x @ self._LP + + def _colorize(self, x): + return x @ self._LA.T + + @cached_property + def _covariance(self): + return (self._v * self._w) @ self._v.T + + def _support_mask(self, x): + """ + Check whether x lies in the support of the distribution. + """ + residual = np.linalg.norm(x @ self._null_basis, axis=-1) + in_support = residual < self._eps + return in_support + + +class CovViaPSD(Covariance): + """ + Representation of a covariance provided via an instance of _PSD + """ + + __class_getitem__ = None # type:ignore[assignment] + + def __init__(self, psd): + self._LP = psd.U + self._log_pdet = psd.log_pdet + self._rank = psd.rank + self._covariance = psd._M + self._shape = psd._M.shape + self._psd = psd + self._allow_singular = False # by default + + def _whiten(self, x): + return x @ self._LP + + def _support_mask(self, x): + return self._psd._support_mask(x) diff --git a/v12x/pylib/scipy/stats/_crosstab.py b/v12x/pylib/scipy/stats/_crosstab.py new file mode 100644 index 0000000000000000000000000000000000000000..dd016a11627968c5fa8361c439c9c55a60a592c2 --- /dev/null +++ b/v12x/pylib/scipy/stats/_crosstab.py @@ -0,0 +1,198 @@ +import numpy as np +from scipy.sparse import coo_array +from scipy._lib._bunch import _make_tuple_bunch + + +CrosstabResult = _make_tuple_bunch( + "CrosstabResult", ["elements", "count"] +) + + +def crosstab(*args, levels=None, sparse=False): + """ + Return table of counts for each possible unique combination in ``*args``. + + When ``len(args) > 1``, the array computed by this function is + often referred to as a *contingency table* [1]_. + + The arguments must be sequences with the same length. The second return + value, `count`, is an integer array with ``len(args)`` dimensions. If + `levels` is None, the shape of `count` is ``(n0, n1, ...)``, where ``nk`` + is the number of unique elements in ``args[k]``. + + Parameters + ---------- + *args : sequences + A sequence of sequences whose unique aligned elements are to be + counted. The sequences in args must all be the same length. + levels : sequence, optional + If `levels` is given, it must be a sequence that is the same length as + `args`. Each element in `levels` is either a sequence or None. If it + is a sequence, it gives the values in the corresponding sequence in + `args` that are to be counted. If any value in the sequences in `args` + does not occur in the corresponding sequence in `levels`, that value + is ignored and not counted in the returned array `count`. The default + value of `levels` for ``args[i]`` is ``np.unique(args[i])`` + sparse : bool, optional + If True, return a sparse array. The array will be an instance of + the `scipy.sparse.coo_array` class. + Default is False. + + Returns + ------- + res : CrosstabResult + An object containing the following attributes: + + elements : tuple of numpy.ndarrays. + Tuple of length ``len(args)`` containing the arrays of elements + that are counted in `count`. These can be interpreted as the + labels of the corresponding dimensions of `count`. If `levels` was + given, then if ``levels[i]`` is not None, ``elements[i]`` will + hold the values given in ``levels[i]``. + count : numpy.ndarray or scipy.sparse.coo_array + Counts of the unique elements in ``zip(*args)``, stored in an + array. Also known as a *contingency table* when ``len(args) > 1``. + + See Also + -------- + numpy.unique + + Notes + ----- + .. versionadded:: 1.7.0 + + References + ---------- + .. [1] "Contingency table", http://en.wikipedia.org/wiki/Contingency_table + + Examples + -------- + >>> from scipy.stats.contingency import crosstab + + Given the lists `a` and `x`, create a contingency table that counts the + frequencies of the corresponding pairs. + + >>> a = ['A', 'B', 'A', 'A', 'B', 'B', 'A', 'A', 'B', 'B'] + >>> x = ['X', 'X', 'X', 'Y', 'Z', 'Z', 'Y', 'Y', 'Z', 'Z'] + >>> res = crosstab(a, x) + >>> avals, xvals = res.elements + >>> avals + array(['A', 'B'], dtype='>> xvals + array(['X', 'Y', 'Z'], dtype='>> res.count + array([[2, 3, 0], + [1, 0, 4]]) + + So ``('A', 'X')`` occurs twice, ``('A', 'Y')`` occurs three times, etc. + + Higher dimensional contingency tables can be created. + + >>> p = [0, 0, 0, 0, 1, 1, 1, 0, 0, 1] + >>> res = crosstab(a, x, p) + >>> res.count + array([[[2, 0], + [2, 1], + [0, 0]], + [[1, 0], + [0, 0], + [1, 3]]]) + >>> res.count.shape + (2, 3, 2) + + The values to be counted can be set by using the `levels` argument. + It allows the elements of interest in each input sequence to be + given explicitly instead finding the unique elements of the sequence. + + For example, suppose one of the arguments is an array containing the + answers to a survey question, with integer values 1 to 4. Even if the + value 1 does not occur in the data, we want an entry for it in the table. + + >>> q1 = [2, 3, 3, 2, 4, 4, 2, 3, 4, 4, 4, 3, 3, 3, 4] # 1 does not occur. + >>> q2 = [4, 4, 2, 2, 2, 4, 1, 1, 2, 2, 4, 2, 2, 2, 4] # 3 does not occur. + >>> options = [1, 2, 3, 4] + >>> res = crosstab(q1, q2, levels=(options, options)) + >>> res.count + array([[0, 0, 0, 0], + [1, 1, 0, 1], + [1, 4, 0, 1], + [0, 3, 0, 3]]) + + If `levels` is given, but an element of `levels` is None, the unique values + of the corresponding argument are used. For example, + + >>> res = crosstab(q1, q2, levels=(None, options)) + >>> res.elements + [array([2, 3, 4]), [1, 2, 3, 4]] + >>> res.count + array([[1, 1, 0, 1], + [1, 4, 0, 1], + [0, 3, 0, 3]]) + + If we want to ignore the pairs where 4 occurs in ``q2``, we can + give just the values [1, 2] to `levels`, and the 4 will be ignored: + + >>> res = crosstab(q1, q2, levels=(None, [1, 2])) + >>> res.elements + [array([2, 3, 4]), [1, 2]] + >>> res.count + array([[1, 1], + [1, 4], + [0, 3]]) + + Finally, let's repeat the first example, but return a sparse array: + + >>> res = crosstab(a, x, sparse=True) + >>> res.count + + >>> res.count.toarray() + array([[2, 3, 0], + [1, 0, 4]]) + + """ + nargs = len(args) + if nargs == 0: + raise TypeError("At least one input sequence is required.") + + len0 = len(args[0]) + if not all(len(a) == len0 for a in args[1:]): + raise ValueError("All input sequences must have the same length.") + + if levels is None: + # Call np.unique with return_inverse=True on each argument. + actual_levels, indices = zip(*[np.unique(a, return_inverse=True) + for a in args]) + else: + # `levels` is not None... + if len(levels) != nargs: + raise ValueError('len(levels) must equal the number of input ' + 'sequences') + + args = [np.asarray(arg) for arg in args] + mask = np.zeros((nargs, len0), dtype=np.bool_) + inv = np.zeros((nargs, len0), dtype=np.intp) + actual_levels = [] + for k, (levels_list, arg) in enumerate(zip(levels, args)): + if levels_list is None: + levels_list, inv[k, :] = np.unique(arg, return_inverse=True) + mask[k, :] = True + else: + q = arg == np.asarray(levels_list).reshape(-1, 1) + mask[k, :] = np.any(q, axis=0) + qnz = q.T.nonzero() + inv[k, qnz[0]] = qnz[1] + actual_levels.append(levels_list) + + mask_all = mask.all(axis=0) + indices = tuple(inv[:, mask_all]) + + if sparse: + count = coo_array((np.ones(len(indices[0]), dtype=int), indices)) + count.sum_duplicates() + else: + shape = [len(u) for u in actual_levels] + count = np.zeros(shape, dtype=int) + np.add.at(count, indices, 1) + + return CrosstabResult(actual_levels, count) diff --git a/v12x/pylib/scipy/stats/_discrete_distns.py b/v12x/pylib/scipy/stats/_discrete_distns.py new file mode 100644 index 0000000000000000000000000000000000000000..07b02cf7283c33fa1e9d133d7500dc99bff6df27 --- /dev/null +++ b/v12x/pylib/scipy/stats/_discrete_distns.py @@ -0,0 +1,2105 @@ +# +# Author: Travis Oliphant 2002-2011 with contributions from +# SciPy Developers 2004-2011 +# +from functools import partial +from types import MethodType + +from scipy import special +from scipy.special import entr, logsumexp, betaln, gammaln as gamln +import scipy.special._ufuncs as scu +from scipy.special._spfun_stats import _poisson_binom_pmf, _poisson_binom_cdf +from scipy._lib._util import rng_integers +import scipy._external.array_api_extra as xpx +from scipy.interpolate import interp1d + +from numpy import floor, ceil, log, exp, sqrt, log1p, expm1, tanh, cosh, sinh + +import numpy as np + +from ._distn_infrastructure import (rv_discrete, get_distribution_names, + _vectorize_rvs_over_shapes, + _ShapeInfo, _isintegral, + rv_discrete_frozen) +from ._biasedurn import (_PyFishersNCHypergeometric, + _PyWalleniusNCHypergeometric, + _PyStochasticLib3) + + +class binom_gen(rv_discrete): + r""" + A binomial discrete random variable. + + %(before_notes)s + + See Also + -------- + hypergeom, nbinom, nhypergeom + + Notes + ----- + The probability mass function for `binom` is: + + .. math:: + + f(k) = \binom{n}{k} p^k (1-p)^{n-k} + + for :math:`k \in \{0, 1, \dots, n\}`, :math:`0 \leq p \leq 1` + + `binom` takes :math:`n` and :math:`p` as shape parameters, + where :math:`p` is the probability of a single success + and :math:`1-p` is the probability of a single failure. + + This distribution uses routines from the Boost Math C++ library for + the computation of the ``pmf``, ``cdf``, ``sf``, ``ppf`` and ``isf`` + methods. [1]_ + + %(after_notes)s + + References + ---------- + .. [1] The Boost Developers. "Boost C++ Libraries". https://www.boost.org/. + + %(example)s + """ + def _shape_info(self): + return [_ShapeInfo("n", True, (0, np.inf), (True, False)), + _ShapeInfo("p", False, (0, 1), (True, True))] + + def _rvs(self, n, p, size=None, random_state=None): + if not np.all(n == np.floor(n)): + raise ValueError("`n` must be integral.") + return random_state.binomial(np.asarray(n, dtype=int), p, size) + + def _argcheck(self, n, p): + return (n >= 0) & _isintegral(n) & (p >= 0) & (p <= 1) + + def _get_support(self, n, p): + return self.a, n + + def _logpmf(self, x, n, p): + k = floor(x) + combiln = (gamln(n+1) - (gamln(k+1) + gamln(n-k+1))) + return combiln + special.xlogy(k, p) + special.xlog1py(n-k, -p) + + def _pmf(self, x, n, p): + # binom.pmf(k) = choose(n, k) * p**k * (1-p)**(n-k) + return scu._binom_pmf(x, n, p) + + def _cdf(self, x, n, p): + k = floor(x) + return scu._binom_cdf(k, n, p) + + def _sf(self, x, n, p): + k = floor(x) + return scu._binom_sf(k, n, p) + + def _isf(self, x, n, p): + return scu._binom_isf(x, n, p) + + def _ppf(self, q, n, p): + return scu._binom_ppf(q, n, p) + + def _stats(self, n, p, moments='mv'): + mu = n * p + var = mu - n * np.square(p) + g1, g2 = None, None + if 's' in moments: + pq = p - np.square(p) + npq_sqrt = np.sqrt(n * pq) + t1 = np.reciprocal(npq_sqrt) + t2 = (2.0 * p) / npq_sqrt + g1 = t1 - t2 + if 'k' in moments: + pq = p - np.square(p) + npq = n * pq + t1 = np.reciprocal(npq) + t2 = 6.0/n + g2 = t1 - t2 + return mu, var, g1, g2 + + +binom = binom_gen(name='binom') + + +class bernoulli_gen(binom_gen): + r"""A Bernoulli discrete random variable. + + %(before_notes)s + + Notes + ----- + The probability mass function for `bernoulli` is: + + .. math:: + + f(k) = \begin{cases}1-p &\text{if } k = 0\\ + p &\text{if } k = 1\end{cases} + + for :math:`k` in :math:`\{0, 1\}`, :math:`0 \leq p \leq 1` + + `bernoulli` takes :math:`p` as shape parameter, + where :math:`p` is the probability of a single success + and :math:`1-p` is the probability of a single failure. + + %(after_notes)s + + %(example)s + + """ + def _shape_info(self): + return [_ShapeInfo("p", False, (0, 1), (True, True))] + + def _rvs(self, p, size=None, random_state=None): + return binom_gen._rvs(self, 1, p, size=size, random_state=random_state) + + def _argcheck(self, p): + return (p >= 0) & (p <= 1) + + def _get_support(self, p): + # Overrides binom_gen._get_support!x + return self.a, self.b + + def _logpmf(self, x, p): + return binom._logpmf(x, 1, p) + + def _pmf(self, x, p): + # bernoulli.pmf(k) = 1-p if k = 0 + # = p if k = 1 + return binom._pmf(x, 1, p) + + def _cdf(self, x, p): + return binom._cdf(x, 1, p) + + def _sf(self, x, p): + return binom._sf(x, 1, p) + + def _isf(self, x, p): + return binom._isf(x, 1, p) + + def _ppf(self, q, p): + return binom._ppf(q, 1, p) + + def _stats(self, p): + return binom._stats(1, p) + + def _entropy(self, p): + return entr(p) + entr(1-p) + + +bernoulli = bernoulli_gen(b=1, name='bernoulli') + + +class betabinom_gen(rv_discrete): + r""" + A beta-binomial discrete random variable. + + %(before_notes)s + + See Also + -------- + beta, binom + + Notes + ----- + The beta-binomial distribution is a binomial distribution with a + probability of success `p` that follows a beta distribution. + + The probability mass function for `betabinom` is: + + .. math:: + + f(k) = \binom{n}{k} \frac{B(k + a, n - k + b)}{B(a, b)} + + for :math:`k \in \{0, 1, \dots, n\}`, :math:`n \geq 0`, :math:`a > 0`, + :math:`b > 0`, where :math:`B(a, b)` is the beta function. + + `betabinom` takes :math:`n`, :math:`a`, and :math:`b` as shape parameters. + + %(after_notes)s + + References + ---------- + .. [1] https://en.wikipedia.org/wiki/Beta-binomial_distribution + + .. versionadded:: 1.4.0 + + %(example)s + """ + def _shape_info(self): + return [_ShapeInfo("n", True, (0, np.inf), (True, False)), + _ShapeInfo("a", False, (0, np.inf), (False, False)), + _ShapeInfo("b", False, (0, np.inf), (False, False))] + + def _rvs(self, n, a, b, size=None, random_state=None): + p = random_state.beta(a, b, size) + if not np.all(n == np.floor(n)): + raise ValueError("`n` must be integral.") + return random_state.binomial(np.asarray(n, dtype=int), p, size) + + def _get_support(self, n, a, b): + return 0, n + + def _argcheck(self, n, a, b): + return (n >= 0) & _isintegral(n) & (a > 0) & (b > 0) + + def _logpmf(self, x, n, a, b): + k = floor(x) + combiln = -log(n + 1) - betaln(n - k + 1, k + 1) + return combiln + betaln(k + a, n - k + b) - betaln(a, b) + + def _pmf(self, x, n, a, b): + return exp(self._logpmf(x, n, a, b)) + + def _stats(self, n, a, b, moments='mv'): + e_p = a / (a + b) + e_q = 1 - e_p + mu = n * e_p + var = n * (a + b + n) * e_p * e_q / (a + b + 1) + g1, g2 = None, None + if 's' in moments: + g1 = 1.0 / sqrt(var) + g1 *= (a + b + 2 * n) * (b - a) + g1 /= (a + b + 2) * (a + b) + if 'k' in moments: + g2 = (a + b).astype(e_p.dtype) + g2 *= (a + b - 1 + 6 * n) + g2 += 3 * a * b * (n - 2) + g2 += 6 * n ** 2 + g2 -= 3 * e_p * b * n * (6 - n) + g2 -= 18 * e_p * e_q * n ** 2 + g2 *= (a + b) ** 2 * (1 + a + b) + g2 /= (n * a * b * (a + b + 2) * (a + b + 3) * (a + b + n)) + g2 -= 3 + return mu, var, g1, g2 + + +betabinom = betabinom_gen(name='betabinom') + + +class nbinom_gen(rv_discrete): + r""" + A negative binomial discrete random variable. + + %(before_notes)s + + See Also + -------- + hypergeom, binom, nhypergeom + + Notes + ----- + Negative binomial distribution describes a sequence of i.i.d. Bernoulli + trials, repeated until a predefined, non-random number of successes occurs. + + The probability mass function of the number of failures for `nbinom` is: + + .. math:: + + f(k) = \binom{k+n-1}{n-1} p^n (1-p)^k + + for :math:`k \ge 0`, :math:`0 < p \leq 1` + + `nbinom` takes :math:`n` and :math:`p` as shape parameters where :math:`n` + is the number of successes, :math:`p` is the probability of a single + success, and :math:`1-p` is the probability of a single failure. + + Another common parameterization of the negative binomial distribution is + in terms of the mean number of failures :math:`\mu` to achieve :math:`n` + successes. The mean :math:`\mu` is related to the probability of success + as + + .. math:: + + p = \frac{n}{n + \mu} + + The number of successes :math:`n` may also be specified in terms of a + "dispersion", "heterogeneity", or "aggregation" parameter :math:`\alpha`, + which relates the mean :math:`\mu` to the variance :math:`\sigma^2`, + e.g. :math:`\sigma^2 = \mu + \alpha \mu^2`. Regardless of the convention + used for :math:`\alpha`, + + .. math:: + + p &= \frac{\mu}{\sigma^2} \\ + n &= \frac{\mu^2}{\sigma^2 - \mu} + + This distribution uses routines from the Boost Math C++ library for + the computation of the ``pmf``, ``cdf``, ``sf``, ``ppf``, ``isf`` + and ``stats`` methods. [1]_ + + %(after_notes)s + + References + ---------- + .. [1] The Boost Developers. "Boost C++ Libraries". https://www.boost.org/. + + %(example)s + """ + def _shape_info(self): + return [_ShapeInfo("n", True, (0, np.inf), (True, False)), + _ShapeInfo("p", False, (0, 1), (True, True))] + + def _rvs(self, n, p, size=None, random_state=None): + return random_state.negative_binomial(n, p, size) + + def _argcheck(self, n, p): + return (n > 0) & (p > 0) & (p <= 1) + + def _pmf(self, x, n, p): + # nbinom.pmf(k) = choose(k+n-1, n-1) * p**n * (1-p)**k + return scu._nbinom_pmf(x, n, p) + + def _logpmf(self, x, n, p): + coeff = gamln(n+x) - gamln(x+1) - gamln(n) + return coeff + n*log(p) + special.xlog1py(x, -p) + + def _cdf(self, x, n, p): + k = floor(x) + return scu._nbinom_cdf(k, n, p) + + def _logcdf(self, x, n, p): + k = floor(x) + k, n, p = np.broadcast_arrays(k, n, p) + cdf = self._cdf(k, n, p) + cond = cdf > 0.5 + def f1(k, n, p): + return np.log1p(-special.betainc(k + 1, n, 1 - p)) + + # do calc in place + logcdf = cdf + with np.errstate(divide='ignore'): + logcdf[cond] = f1(k[cond], n[cond], p[cond]) + logcdf[~cond] = np.log(cdf[~cond]) + return logcdf + + def _sf(self, x, n, p): + k = floor(x) + return scu._nbinom_sf(k, n, p) + + def _isf(self, x, n, p): + with np.errstate(over='ignore'): # see gh-17432 + return scu._nbinom_isf(x, n, p) + + def _ppf(self, q, n, p): + with np.errstate(over='ignore'): # see gh-17432 + return scu._nbinom_ppf(q, n, p) + + def _stats(self, n, p): + return ( + scu._nbinom_mean(n, p), + scu._nbinom_variance(n, p), + scu._nbinom_skewness(n, p), + scu._nbinom_kurtosis_excess(n, p), + ) + + +nbinom = nbinom_gen(name='nbinom') + + +class betanbinom_gen(rv_discrete): + r""" + A beta-negative-binomial discrete random variable. + + %(before_notes)s + + See Also + -------- + betabinom : Beta binomial distribution + + Notes + ----- + The beta-negative-binomial distribution is a negative binomial + distribution with a probability of success `p` that follows a + beta distribution. + + The probability mass function for `betanbinom` is: + + .. math:: + + f(k) = \binom{n + k - 1}{k} \frac{B(a + n, b + k)}{B(a, b)} + + for :math:`k \ge 0`, :math:`n \geq 0`, :math:`a > 0`, + :math:`b > 0`, where :math:`B(a, b)` is the beta function. + + `betanbinom` takes :math:`n`, :math:`a`, and :math:`b` as shape parameters. + + %(after_notes)s + + References + ---------- + .. [1] https://en.wikipedia.org/wiki/Beta_negative_binomial_distribution + + .. versionadded:: 1.12.0 + + %(example)s + """ + def _shape_info(self): + return [_ShapeInfo("n", True, (0, np.inf), (True, False)), + _ShapeInfo("a", False, (0, np.inf), (False, False)), + _ShapeInfo("b", False, (0, np.inf), (False, False))] + + def _rvs(self, n, a, b, size=None, random_state=None): + p = random_state.beta(a, b, size) + return random_state.negative_binomial(n, p, size) + + def _argcheck(self, n, a, b): + return (n >= 0) & _isintegral(n) & (a > 0) & (b > 0) + + def _logpmf(self, x, n, a, b): + k = floor(x) + combiln = -np.log(n + k) - betaln(n, k + 1) + return combiln + betaln(a + n, b + k) - betaln(a, b) + + def _pmf(self, x, n, a, b): + return exp(self._logpmf(x, n, a, b)) + + def _stats(self, n, a, b, moments='mv'): + # reference: Wolfram Alpha input + # BetaNegativeBinomialDistribution[a, b, n] + def mean(n, a, b): + return n * b / (a - 1.) + mu = xpx.apply_where(a > 1, (n, a, b), mean, fill_value=np.inf) + def var(n, a, b): + return (n * b * (n + a - 1.) * (a + b - 1.) + / ((a - 2.) * (a - 1.)**2.)) + var = xpx.apply_where(a > 2, (n, a, b), var, fill_value=np.inf) + g1, g2 = None, None + def skew(n, a, b): + return ((2 * n + a - 1.) * (2 * b + a - 1.) + / (a - 3.) / sqrt(n * b * (n + a - 1.) * (b + a - 1.) + / (a - 2.))) + if 's' in moments: + g1 = xpx.apply_where(a > 3, (n, a, b), skew, fill_value=np.inf) + def kurtosis(n, a, b): + term = (a - 2.) + term_2 = ((a - 1.)**2. * (a**2. + a * (6 * b - 1.) + + 6. * (b - 1.) * b) + + 3. * n**2. * ((a + 5.) * b**2. + (a + 5.) + * (a - 1.) * b + 2. * (a - 1.)**2) + + 3 * (a - 1.) * n + * ((a + 5.) * b**2. + (a + 5.) * (a - 1.) * b + + 2. * (a - 1.)**2.)) + denominator = ((a - 4.) * (a - 3.) * b * n + * (a + b - 1.) * (a + n - 1.)) + # Wolfram Alpha uses Pearson kurtosis, so we subtract 3 to get + # scipy's Fisher kurtosis + return term * term_2 / denominator - 3. + if 'k' in moments: + g2 = xpx.apply_where(a > 4, (n, a, b), kurtosis, fill_value=np.inf) + return mu, var, g1, g2 + + +betanbinom = betanbinom_gen(name='betanbinom') + + +class geom_gen(rv_discrete): + r"""A geometric discrete random variable. + + %(before_notes)s + + See Also + -------- + planck + + Notes + ----- + The probability mass function for `geom` is: + + .. math:: + + f(k) = (1-p)^{k-1} p + + for :math:`k \ge 1`, :math:`0 < p \leq 1` + + `geom` takes :math:`p` as shape parameter, + where :math:`p` is the probability of a single success + and :math:`1-p` is the probability of a single failure. + + Note that when drawing random samples, the probability of observations that exceed + ``np.iinfo(np.int64).max`` increases rapidly as $p$ decreases below $10^{-17}$. For + $p < 10^{-20}$, almost all observations would exceed the maximum ``int64``; however, + the output dtype is always ``int64``, so these values are clipped to the maximum. + + %(after_notes)s + + %(example)s + """ + + def _shape_info(self): + return [_ShapeInfo("p", False, (0, 1), (True, True))] + + def _rvs(self, p, size=None, random_state=None): + res = random_state.geometric(p, size=size) + # RandomState.geometric can wrap around to negative values; make behavior + # consistent with Generator.geometric by replacing with maximum integer. + max_int = np.iinfo(res.dtype).max + return np.where(res < 0, max_int, res) + + def _argcheck(self, p): + return (p <= 1) & (p > 0) + + def _pmf(self, k, p): + return np.power(1-p, k-1) * p + + def _logpmf(self, k, p): + return special.xlog1py(k - 1, -p) + log(p) + + def _cdf(self, x, p): + k = floor(x) + return -expm1(log1p(-p)*k) + + def _sf(self, x, p): + return np.exp(self._logsf(x, p)) + + def _logsf(self, x, p): + k = floor(x) + return k*log1p(-p) + + def _ppf(self, q, p): + vals = ceil(log1p(-q) / log1p(-p)) + temp = self._cdf(vals-1, p) + return np.where((temp >= q) & (vals > 0), vals-1, vals) + + def _stats(self, p): + mu = 1.0/p + qr = 1.0-p + var = qr / p / p + g1 = (2.0-p) / sqrt(qr) + g2 = np.polyval([1, -6, 6], p)/(1.0-p) + return mu, var, g1, g2 + + def _entropy(self, p): + return -np.log(p) - np.log1p(-p) * (1.0-p) / p + + +geom = geom_gen(a=1, name='geom', longname="A geometric") + + +class hypergeom_gen(rv_discrete): + r"""A hypergeometric discrete random variable. + + The hypergeometric distribution models drawing objects from a bin. + `M` is the total number of objects, `n` is total number of Type I objects. + The random variate represents the number of Type I objects in `N` drawn + without replacement from the total population. + + %(before_notes)s + + See Also + -------- + nhypergeom, binom, nbinom + + Notes + ----- + The symbols used to denote the shape parameters (`M`, `n`, and `N`) are not + universally accepted. See the Examples for a clarification of the + definitions used here. + + The probability mass function is defined as, + + .. math:: p(k, M, n, N) = \frac{\binom{n}{k} \binom{M - n}{N - k}} + {\binom{M}{N}} + + for :math:`k \in [\max(0, N - M + n), \min(n, N)]`, where the binomial + coefficients are defined as, + + .. math:: \binom{n}{k} \equiv \frac{n!}{k! (n - k)!}. + + This distribution uses routines from the Boost Math C++ library for + the computation of the ``pmf``, ``cdf``, ``sf`` and ``stats`` methods. [1]_ + + %(after_notes)s + + References + ---------- + .. [1] The Boost Developers. "Boost C++ Libraries". https://www.boost.org/. + + Examples + -------- + >>> import numpy as np + >>> from scipy.stats import hypergeom + >>> import matplotlib.pyplot as plt + + Suppose we have a collection of 20 animals, of which 7 are dogs. Then if + we want to know the probability of finding a given number of dogs if we + choose at random 12 of the 20 animals, we can initialize a frozen + distribution and plot the probability mass function: + + >>> [M, n, N] = [20, 7, 12] + >>> rv = hypergeom(M, n, N) + >>> x = np.arange(0, n+1) + >>> pmf_dogs = rv.pmf(x) + + >>> fig = plt.figure() + >>> ax = fig.add_subplot(111) + >>> ax.plot(x, pmf_dogs, 'bo') + >>> ax.vlines(x, 0, pmf_dogs, lw=2) + >>> ax.set_xlabel('# of dogs in our group of chosen animals') + >>> ax.set_ylabel('hypergeom PMF') + >>> plt.show() + + Instead of using a frozen distribution we can also use `hypergeom` + methods directly. To for example obtain the cumulative distribution + function, use: + + >>> prb = hypergeom.cdf(x, M, n, N) + + And to generate random numbers: + + >>> R = hypergeom.rvs(M, n, N, size=10) + """ + def _shape_info(self): + return [_ShapeInfo("M", True, (0, np.inf), (True, False)), + _ShapeInfo("n", True, (0, np.inf), (True, False)), + _ShapeInfo("N", True, (0, np.inf), (True, False))] + + def _rvs(self, M, n, N, size=None, random_state=None): + return random_state.hypergeometric(n, M-n, N, size=size) + + def _get_support(self, M, n, N): + return np.maximum(N-(M-n), 0), np.minimum(n, N) + + def _argcheck(self, M, n, N): + cond = (M > 0) & (n >= 0) & (N >= 0) + cond &= (n <= M) & (N <= M) + cond &= _isintegral(M) & _isintegral(n) & _isintegral(N) + return cond + + def _logpmf(self, k, M, n, N): + tot, good = M, n + bad = tot - good + result = (betaln(good+1, 1) + betaln(bad+1, 1) + betaln(tot-N+1, N+1) - + betaln(k+1, good-k+1) - betaln(N-k+1, bad-N+k+1) - + betaln(tot+1, 1)) + return result + + def _pmf(self, k, M, n, N): + return scu._hypergeom_pmf(k, n, N, M) + + def _cdf(self, k, M, n, N): + return scu._hypergeom_cdf(k, n, N, M) + + def _stats(self, M, n, N): + M, n, N = 1. * M, 1. * n, 1. * N + m = M - n + + # Boost kurtosis_excess doesn't return the same as the value + # computed here. + g2 = M * (M + 1) - 6. * N * (M - N) - 6. * n * m + g2 *= (M - 1) * M * M + g2 += 6. * n * N * (M - N) * m * (5. * M - 6) + g2 /= n * N * (M - N) * m * (M - 2.) * (M - 3.) + return ( + scu._hypergeom_mean(n, N, M), + scu._hypergeom_variance(n, N, M), + scu._hypergeom_skewness(n, N, M), + g2, + ) + + def _entropy(self, M, n, N): + k = np.r_[N - (M - n):min(n, N) + 1] + vals = self.pmf(k, M, n, N) + return np.sum(entr(vals), axis=0) + + def _sf(self, k, M, n, N): + return scu._hypergeom_sf(k, n, N, M) + + def _logsf(self, k, M, n, N): + res = [] + for quant, tot, good, draw in zip(*np.broadcast_arrays(k, M, n, N)): + if (quant + 0.5) * (tot + 0.5) < (good - 0.5) * (draw - 0.5): + # Less terms to sum if we calculate log(1-cdf) + res.append(log1p(-exp(self.logcdf(quant, tot, good, draw)))) + else: + # Integration over probability mass function using logsumexp + k2 = np.arange(quant + 1, draw + 1) + res.append(logsumexp(self._logpmf(k2, tot, good, draw))) + return np.asarray(res) + + def _logcdf(self, k, M, n, N): + res = [] + for quant, tot, good, draw in zip(*np.broadcast_arrays(k, M, n, N)): + if (quant + 0.5) * (tot + 0.5) > (good - 0.5) * (draw - 0.5): + # Less terms to sum if we calculate log(1-sf) + res.append(log1p(-exp(self.logsf(quant, tot, good, draw)))) + else: + # Integration over probability mass function using logsumexp + k2 = np.arange(0, quant + 1) + res.append(logsumexp(self._logpmf(k2, tot, good, draw))) + return np.asarray(res) + + +hypergeom = hypergeom_gen(name='hypergeom') + + +class nhypergeom_gen(rv_discrete): + r"""A negative hypergeometric discrete random variable. + + Consider a box containing :math:`M` balls:, :math:`n` red and + :math:`M-n` blue. We randomly sample balls from the box, one + at a time and *without* replacement, until we have picked :math:`r` + blue balls. `nhypergeom` is the distribution of the number of + red balls :math:`k` we have picked. + + %(before_notes)s + + See Also + -------- + hypergeom, binom, nbinom + + Notes + ----- + The symbols used to denote the shape parameters (`M`, `n`, and `r`) are not + universally accepted. See the Examples for a clarification of the + definitions used here. + + The probability mass function is defined as, + + .. math:: f(k; M, n, r) = \frac{{{k+r-1}\choose{k}}{{M-r-k}\choose{n-k}}} + {{M \choose n}} + + for :math:`k \in [0, n]`, :math:`n \in [0, M]`, :math:`r \in [0, M-n]`, + and the binomial coefficient is: + + .. math:: \binom{n}{k} \equiv \frac{n!}{k! (n - k)!}. + + It is equivalent to observing :math:`k` successes in :math:`k+r-1` + samples with :math:`k+r`'th sample being a failure. The former + can be modelled as a hypergeometric distribution. The probability + of the latter is simply the number of failures remaining + :math:`M-n-(r-1)` divided by the size of the remaining population + :math:`M-(k+r-1)`. This relationship can be shown as: + + .. math:: NHG(k;M,n,r) = HG(k;M,n,k+r-1)\frac{(M-n-(r-1))}{(M-(k+r-1))} + + where :math:`NHG` is probability mass function (PMF) of the + negative hypergeometric distribution and :math:`HG` is the + PMF of the hypergeometric distribution. + + %(after_notes)s + + References + ---------- + .. [1] Negative Hypergeometric Distribution on Wikipedia + https://en.wikipedia.org/wiki/Negative_hypergeometric_distribution + + .. [2] Negative Hypergeometric Distribution from + http://www.math.wm.edu/~leemis/chart/UDR/PDFs/Negativehypergeometric.pdf + + Examples + -------- + >>> import numpy as np + >>> from scipy.stats import nhypergeom + >>> import matplotlib.pyplot as plt + + Suppose we have a collection of 20 animals, of which 7 are dogs. + Then if we want to know the probability of finding a given number + of dogs (successes) in a sample with exactly 12 animals that + aren't dogs (failures), we can initialize a frozen distribution + and plot the probability mass function: + + >>> M, n, r = [20, 7, 12] + >>> rv = nhypergeom(M, n, r) + >>> x = np.arange(0, n+2) + >>> pmf_dogs = rv.pmf(x) + + >>> fig = plt.figure() + >>> ax = fig.add_subplot(111) + >>> ax.plot(x, pmf_dogs, 'bo') + >>> ax.vlines(x, 0, pmf_dogs, lw=2) + >>> ax.set_xlabel('# of dogs in our group with given 12 failures') + >>> ax.set_ylabel('nhypergeom PMF') + >>> plt.show() + + Instead of using a frozen distribution we can also use `nhypergeom` + methods directly. To for example obtain the probability mass + function, use: + + >>> prb = nhypergeom.pmf(x, M, n, r) + + And to generate random numbers: + + >>> R = nhypergeom.rvs(M, n, r, size=10) + + To verify the relationship between `hypergeom` and `nhypergeom`, use: + + >>> from scipy.stats import hypergeom, nhypergeom + >>> M, n, r = 45, 13, 8 + >>> k = 6 + >>> nhypergeom.pmf(k, M, n, r) + 0.06180776620271643 + >>> hypergeom.pmf(k, M, n, k+r-1) * (M - n - (r-1)) / (M - (k+r-1)) + 0.06180776620271644 + """ + + def _shape_info(self): + return [_ShapeInfo("M", True, (0, np.inf), (True, False)), + _ShapeInfo("n", True, (0, np.inf), (True, False)), + _ShapeInfo("r", True, (0, np.inf), (True, False))] + + def _get_support(self, M, n, r): + return 0, n + + def _argcheck(self, M, n, r): + cond = (n >= 0) & (n <= M) & (r >= 0) & (r <= M-n) + cond &= _isintegral(M) & _isintegral(n) & _isintegral(r) + return cond + + def _rvs(self, M, n, r, size=None, random_state=None): + + @_vectorize_rvs_over_shapes + def _rvs1(M, n, r, size, random_state): + # invert cdf by calculating all values in support, scalar M, n, r + a, b = self.support(M, n, r) + ks = np.arange(a, b+1) + cdf = self.cdf(ks, M, n, r) + ppf = interp1d(cdf, ks, kind='next', fill_value='extrapolate') + rvs = ppf(random_state.uniform(size=size)).astype(int) + if size is None: + return rvs.item() + return rvs + + return _rvs1(M, n, r, size=size, random_state=random_state) + + def _logpmf(self, k, M, n, r): + return xpx.apply_where( + (r != 0) | (k != 0), (k, M, n, r), + lambda k, M, n, r: + (-betaln(k+1, r) + betaln(k+r, 1) + - betaln(n-k+1, M-r-n+1) + betaln(M-r-k+1, 1) + + betaln(n+1, M-n+1) - betaln(M+1, 1)), + fill_value=0.0) + + def _pmf(self, k, M, n, r): + # same as the following but numerically more precise + # return comb(k+r-1, k) * comb(M-r-k, n-k) / comb(M, n) + return exp(self._logpmf(k, M, n, r)) + + def _stats(self, M, n, r): + # Promote the datatype to at least float + # mu = rn / (M-n+1) + M, n, r = 1.*M, 1.*n, 1.*r + mu = r*n / (M-n+1) + + var = r*(M+1)*n / ((M-n+1)*(M-n+2)) * (1 - r / (M-n+1)) + + # The skew and kurtosis are mathematically + # intractable so return `None`. See [2]_. + g1, g2 = None, None + return mu, var, g1, g2 + + +nhypergeom = nhypergeom_gen(name='nhypergeom') + + +# FIXME: Fails _cdfvec +class logser_gen(rv_discrete): + r"""A Logarithmic (Log-Series, Series) discrete random variable. + + %(before_notes)s + + Notes + ----- + The probability mass function for `logser` is: + + .. math:: + + f(k) = - \frac{p^k}{k \log(1-p)} + + for :math:`k \ge 1`, :math:`0 < p < 1` + + `logser` takes :math:`p` as shape parameter, + where :math:`p` is the probability of a single success + and :math:`1-p` is the probability of a single failure. + + %(after_notes)s + + %(example)s + + """ + + def _shape_info(self): + return [_ShapeInfo("p", False, (0, 1), (True, True))] + + def _rvs(self, p, size=None, random_state=None): + # looks wrong for p>0.5, too few k=1 + # trying to use generic is worse, no k=1 at all + return random_state.logseries(p, size=size) + + def _argcheck(self, p): + return (p > 0) & (p < 1) + + def _pmf(self, k, p): + # logser.pmf(k) = - p**k / (k*log(1-p)) + return -np.power(p, k) * 1.0 / k / special.log1p(-p) + + def _sf(self, k, p): + tiny = 1e-100 + # Ideally, this is the unregularized beta function with `b=0`. We don't have + # an unregularized beta function (yet, although we could get it from Boost), + # and neither technically support `b=0` - despite the function being accurate + # for `b` super close to zero. See https://github.com/scipy/scipy/issues/3890. + return -special.betainc(k+1, tiny, p) * special.beta(k+1, tiny) / np.log1p(-p) + + def _stats(self, p): + r = special.log1p(-p) + mu = p / (p - 1.0) / r + mu2p = -p / r / (p - 1.0)**2 + var = mu2p - mu*mu + mu3p = -p / r * (1.0+p) / (1.0 - p)**3 + mu3 = mu3p - 3*mu*mu2p + 2*mu**3 + g1 = mu3 / np.power(var, 1.5) + + mu4p = -p / r * ( + 1.0 / (p-1)**2 - 6*p / (p - 1)**3 + 6*p*p / (p-1)**4) + mu4 = mu4p - 4*mu3p*mu + 6*mu2p*mu*mu - 3*mu**4 + g2 = mu4 / var**2 - 3.0 + return mu, var, g1, g2 + + +logser = logser_gen(a=1, name='logser', longname='A logarithmic') + + +class poisson_gen(rv_discrete): + r"""A Poisson discrete random variable. + + %(before_notes)s + + Notes + ----- + The probability mass function for `poisson` is: + + .. math:: + + f(k) = \exp(-\mu) \frac{\mu^k}{k!} + + for :math:`k \ge 0`. + + `poisson` takes :math:`\mu \geq 0` as shape parameter. + When :math:`\mu = 0`, the ``pmf`` method + returns ``1.0`` at quantile :math:`k = 0`. + + %(after_notes)s + + %(example)s + + """ + + def _shape_info(self): + return [_ShapeInfo("mu", False, (0, np.inf), (True, False))] + + # Override rv_discrete._argcheck to allow mu=0. + def _argcheck(self, mu): + return mu >= 0 + + def _rvs(self, mu, size=None, random_state=None): + return random_state.poisson(mu, size) + + def _logpmf(self, k, mu): + Pk = special.xlogy(k, mu) - gamln(k + 1) - mu + return Pk + + def _pmf(self, k, mu): + # poisson.pmf(k) = exp(-mu) * mu**k / k! + return exp(self._logpmf(k, mu)) + + def _cdf(self, x, mu): + k = floor(x) + return special.pdtr(k, mu) + + def _sf(self, x, mu): + k = floor(x) + return special.pdtrc(k, mu) + + def _ppf(self, q, mu): + vals = ceil(special.pdtrik(q, mu)) + vals1 = np.maximum(vals - 1, 0) + temp = special.pdtr(vals1, mu) + return np.where(temp >= q, vals1, vals) + + def _stats(self, mu): + var = mu + tmp = np.asarray(mu) + mu_nonzero = tmp > 0 + g1 = xpx.apply_where(mu_nonzero, tmp, lambda x: sqrt(1.0/x), fill_value=np.inf) + g2 = xpx.apply_where(mu_nonzero, tmp, lambda x: 1.0/x, fill_value=np.inf) + return mu, var, g1, g2 + + +poisson = poisson_gen(name="poisson", longname='A Poisson') + + +class planck_gen(rv_discrete): + r"""A Planck discrete exponential random variable. + + %(before_notes)s + + See Also + -------- + geom + + Notes + ----- + The probability mass function for `planck` is: + + .. math:: + + f(k) = (1-\exp(-\lambda)) \exp(-\lambda k) + + for :math:`k \ge 0` and :math:`\lambda > 0`. + + `planck` takes :math:`\lambda` as shape parameter. The Planck distribution + can be written as a geometric distribution (`geom`) with + :math:`p = 1 - \exp(-\lambda)` shifted by ``loc = -1``. + + %(after_notes)s + + %(example)s + """ + def _shape_info(self): + return [_ShapeInfo("lambda_", False, (0, np.inf), (False, False))] + + def _argcheck(self, lambda_): + return lambda_ > 0 + + def _pmf(self, k, lambda_): + return -expm1(-lambda_)*exp(-lambda_*k) + + def _cdf(self, x, lambda_): + k = floor(x) + return -expm1(-lambda_*(k+1)) + + def _sf(self, x, lambda_): + return exp(self._logsf(x, lambda_)) + + def _logsf(self, x, lambda_): + k = floor(x) + return -lambda_*(k+1) + + def _ppf(self, q, lambda_): + vals = ceil(-1.0/lambda_ * log1p(-q)-1) + vals1 = (vals-1).clip(*(self._get_support(lambda_))) + temp = self._cdf(vals1, lambda_) + return np.where(temp >= q, vals1, vals) + + def _rvs(self, lambda_, size=None, random_state=None): + # use relation to geometric distribution for sampling + p = -expm1(-lambda_) + return random_state.geometric(p, size=size) - 1.0 + + def _stats(self, lambda_): + mu = 1/expm1(lambda_) + var = exp(-lambda_)/(expm1(-lambda_))**2 + g1 = 2*cosh(lambda_/2.0) + g2 = 4+2*cosh(lambda_) + return mu, var, g1, g2 + + def _entropy(self, lambda_): + C = -expm1(-lambda_) + return lambda_*exp(-lambda_)/C - log(C) + + +planck = planck_gen(a=0, name='planck', longname='A discrete exponential ') + + +class boltzmann_gen(rv_discrete): + r"""A Boltzmann (Truncated Discrete Exponential) random variable. + + %(before_notes)s + + Notes + ----- + The probability mass function for `boltzmann` is: + + .. math:: + + f(k) = (1-\exp(-\lambda)) \exp(-\lambda k) / (1-\exp(-\lambda N)) + + for :math:`k = 0,..., N-1`. + + `boltzmann` takes :math:`\lambda > 0` and :math:`N > 0` as shape parameters. + + %(after_notes)s + + %(example)s + + """ + def _shape_info(self): + return [_ShapeInfo("lambda_", False, (0, np.inf), (False, False)), + _ShapeInfo("N", True, (0, np.inf), (False, False))] + + def _argcheck(self, lambda_, N): + return (lambda_ > 0) & (N > 0) & _isintegral(N) + + def _get_support(self, lambda_, N): + return self.a, N - 1 + + def _pmf(self, k, lambda_, N): + # boltzmann.pmf(k) = + # (1-exp(-lambda_)*exp(-lambda_*k)/(1-exp(-lambda_*N)) + fact = (1-exp(-lambda_))/(1-exp(-lambda_*N)) + return fact*exp(-lambda_*k) + + def _cdf(self, x, lambda_, N): + k = floor(x) + return (1-exp(-lambda_*(k+1)))/(1-exp(-lambda_*N)) + + def _ppf(self, q, lambda_, N): + qnew = q*(1-exp(-lambda_*N)) + vals = ceil(-1.0/lambda_ * log(1-qnew)-1) + vals1 = (vals-1).clip(0.0, np.inf) + temp = self._cdf(vals1, lambda_, N) + return np.where(temp >= q, vals1, vals) + + def _stats(self, lambda_, N): + z = exp(-lambda_) + zN = exp(-lambda_*N) + mu = z/(1.0-z)-N*zN/(1-zN) + var = z/(1.0-z)**2 - N*N*zN/(1-zN)**2 + trm = (1-zN)/(1-z) + trm2 = (z*trm**2 - N*N*zN) + g1 = z*(1+z)*trm**3 - N**3*zN*(1+zN) + g1 = g1 / trm2**(1.5) + g2 = z*(1+4*z+z*z)*trm**4 - N**4 * zN*(1+4*zN+zN*zN) + g2 = g2 / trm2 / trm2 + return mu, var, g1, g2 + + +boltzmann = boltzmann_gen(name='boltzmann', a=0, + longname='A truncated discrete exponential ') + + +class randint_gen(rv_discrete): + r"""A uniform discrete random variable. + + %(before_notes)s + + Notes + ----- + The probability mass function for `randint` is: + + .. math:: + + f(k) = \frac{1}{\texttt{high} - \texttt{low}} + + for :math:`k \in \{\texttt{low}, \dots, \texttt{high} - 1\}`. + + `randint` takes :math:`\texttt{low}` and :math:`\texttt{high}` as shape + parameters. + + %(after_notes)s + + Examples + -------- + >>> import numpy as np + >>> from scipy.stats import randint + >>> import matplotlib.pyplot as plt + >>> fig, ax = plt.subplots(1, 1) + + Calculate the first four moments: + + >>> low, high = 7, 31 + >>> mean, var, skew, kurt = randint.stats(low, high, moments='mvsk') + + Display the probability mass function (``pmf``): + + >>> x = np.arange(low - 5, high + 5) + >>> ax.plot(x, randint.pmf(x, low, high), 'bo', ms=8, label='randint pmf') + >>> ax.vlines(x, 0, randint.pmf(x, low, high), colors='b', lw=5, alpha=0.5) + + Alternatively, the distribution object can be called (as a function) to + fix the shape and location. This returns a "frozen" RV object holding the + given parameters fixed. + + Freeze the distribution and display the frozen ``pmf``: + + >>> rv = randint(low, high) + >>> ax.vlines(x, 0, rv.pmf(x), colors='k', linestyles='-', + ... lw=1, label='frozen pmf') + >>> ax.legend(loc='lower center') + >>> plt.show() + + Check the relationship between the cumulative distribution function + (``cdf``) and its inverse, the percent point function (``ppf``): + + >>> q = np.arange(low, high) + >>> p = randint.cdf(q, low, high) + >>> np.allclose(q, randint.ppf(p, low, high)) + True + + Generate random numbers: + + >>> r = randint.rvs(low, high, size=1000) + + """ + + def _shape_info(self): + return [_ShapeInfo("low", True, (-np.inf, np.inf), (False, False)), + _ShapeInfo("high", True, (-np.inf, np.inf), (False, False))] + + def _argcheck(self, low, high): + return (high > low) & _isintegral(low) & _isintegral(high) + + def _get_support(self, low, high): + return low, high-1 + + def _pmf(self, k, low, high): + # randint.pmf(k) = 1./(high - low) + p = np.ones_like(k) / (np.asarray(high, dtype=np.int64) - low) + return np.where((k >= low) & (k < high), p, 0.) + + def _cdf(self, x, low, high): + k = floor(x) + return (k - low + 1.) / (high - low) + + def _ppf(self, q, low, high): + vals = ceil(q * (high - low) + low) - 1 + vals1 = (vals - 1).clip(low, high) + temp = self._cdf(vals1, low, high) + return np.where(temp >= q, vals1, vals) + + def _stats(self, low, high): + m2, m1 = np.asarray(high), np.asarray(low) + mu = (m2 + m1 - 1.0) / 2 + d = m2 - m1 + var = (d*d - 1) / 12.0 + g1 = 0.0 + g2 = -6.0/5.0 * (d*d + 1.0) / (d*d - 1.0) + return mu, var, g1, g2 + + def _rvs(self, low, high, size=None, random_state=None): + """An array of *size* random integers >= ``low`` and < ``high``.""" + if np.asarray(low).size == 1 and np.asarray(high).size == 1: + # no need to vectorize in that case + return rng_integers(random_state, low, high, size=size) + + if size is not None: + # NumPy's RandomState.randint() doesn't broadcast its arguments. + # Use `broadcast_to()` to extend the shapes of low and high + # up to size. Then we can use the numpy.vectorize'd + # randint without needing to pass it a `size` argument. + low = np.broadcast_to(low, size) + high = np.broadcast_to(high, size) + randint = np.vectorize(partial(rng_integers, random_state), + otypes=[np.dtype(int)]) + return randint(low, high) + + def _entropy(self, low, high): + return log(high - low) + + +randint = randint_gen(name='randint', longname='A discrete uniform ' + '(random integer)') + + +# FIXME: problems sampling. +class zipf_gen(rv_discrete): + r"""A Zipf (Zeta) discrete random variable. + + %(before_notes)s + + See Also + -------- + zipfian + + Notes + ----- + The probability mass function for `zipf` is: + + .. math:: + + f(k, a) = \frac{1}{\zeta(a) k^a} + + for :math:`k \ge 1`, :math:`a > 1`. + + `zipf` takes :math:`a > 1` as shape parameter. :math:`\zeta` is the + Riemann zeta function (`scipy.special.zeta`) + + The Zipf distribution is also known as the zeta distribution, which is + a special case of the Zipfian distribution (`zipfian`). + + %(after_notes)s + + References + ---------- + .. [1] "Zeta Distribution", Wikipedia, + https://en.wikipedia.org/wiki/Zeta_distribution + + %(example)s + + Confirm that `zipf` is the large `n` limit of `zipfian`. + + >>> import numpy as np + >>> from scipy.stats import zipf, zipfian + >>> k = np.arange(11) + >>> np.allclose(zipf.pmf(k, a), zipfian.pmf(k, a, n=10000000)) + True + + """ + + def _shape_info(self): + return [_ShapeInfo("a", False, (1, np.inf), (False, False))] + + def _rvs(self, a, size=None, random_state=None): + return random_state.zipf(a, size=size) + + def _argcheck(self, a): + return a > 1 + + def _pmf(self, k, a): + k = k.astype(np.float64) + # zipf.pmf(k, a) = 1/(zeta(a) * k**a) + Pk = 1.0 / special.zeta(a, 1) * k**-a + return Pk + + def _munp(self, n, a): + return xpx.apply_where( + a > n + 1, (a, n), + lambda a, n: special.zeta(a - n, 1) / special.zeta(a, 1), + fill_value=np.inf) + + +zipf = zipf_gen(a=1, name='zipf', longname='A Zipf') + + +class zipfian_gen(rv_discrete): + r"""A Zipfian discrete random variable. + + %(before_notes)s + + See Also + -------- + zipf + + Notes + ----- + The probability mass function for `zipfian` is: + + .. math:: + + f(k, a, n) = \frac{1}{H_{n,a} k^a} + + for :math:`k \in \{1, 2, \dots, n-1, n\}`, :math:`a \ge 0`, + :math:`n \in \{1, 2, 3, \dots\}`. + + `zipfian` takes :math:`a` and :math:`n` as shape parameters. + :math:`H_{n,a}` is the :math:`n`:sup:`th` generalized harmonic + number of order :math:`a`. + + The SciPy implementation of this distribution requires :math:`1 \le n \le 2^{53}`. + For larger values of :math:`n`, the `zipfian` methods (`pmf`, `cdf`, `mean`, etc.) + will return `nan`. + + When :math:`a > 1`, the Zipfian distribution reduces to the Zipf (zeta) + distribution as :math:`n \rightarrow \infty`. + + %(after_notes)s + + References + ---------- + .. [1] "Zipf's Law", Wikipedia, https://en.wikipedia.org/wiki/Zipf's_law + .. [2] Larry Leemis, "Zipf Distribution", Univariate Distribution + Relationships. http://www.math.wm.edu/~leemis/chart/UDR/PDFs/Zipf.pdf + + %(example)s + + Confirm that `zipfian` reduces to `zipf` for large `n`, ``a > 1``. + + >>> import numpy as np + >>> from scipy.stats import zipf, zipfian + >>> k = np.arange(11) + >>> np.allclose(zipfian.pmf(k, a=3.5, n=10000000), zipf.pmf(k, a=3.5)) + True + + """ + + def _shape_info(self): + return [_ShapeInfo("a", False, (0, np.inf), (True, False)), + _ShapeInfo("n", True, (0, np.inf), (False, False))] + + def _argcheck(self, a, n): + # The upper bound on n is for practical numerical reasons. The numerical + # methods for computing the PMF, CDF and SF involve sums over range(1, n+1) + # when a <= 1, so there is no way they can be computed for extremely large + # n--even 2**53 is ridiculously large for those calculations. The bound is + # also required to ensure that the loops compute the sums correctly when + # the inputs are double precision instead of integers. + # + # n may be an integer or a float, but the value must be an integer in the + # range 1 <= n <= 2**53. The expression below clips n to the accepted range + # before attempting to cast to integer to avoid warnings generated by an + # input such as n=1e100. The extra `np.asarray()` wrapper avoids the error + # that arises with an input such as `n=2**100`. + # upper limit changed in gh-24119 because of overflow on windows, linux32 bit + return ((a >= 0) & + (n == np.asarray(np.clip(n, 1, 2**31 - 1)).astype(dtype=np.int64))) + + def _get_support(self, a, n): + return 1, np.floor(n) + + def _pmf(self, k, a, n): + k = np.floor(k) + n = np.floor(n) + return scu._normalized_gen_harmonic(k, k, n, a) + + def _cdf(self, k, a, n): + k = np.floor(k) + n = np.floor(n) + return scu._normalized_gen_harmonic(1, k, n, a) + + def _sf(self, k, a, n): + k = np.floor(k) + n = np.floor(n) + return scu._normalized_gen_harmonic(k + 1, n, n, a) + + def _stats(self, a, n): + n = np.floor(n) + # see http://www.math.wm.edu/~leemis/chart/UDR/PDFs/Zipf.pdf + Hna = scu._gen_harmonic(n, a) + Hna1 = scu._gen_harmonic(n, a-1) + Hna2 = scu._gen_harmonic(n, a-2) + Hna3 = scu._gen_harmonic(n, a-3) + Hna4 = scu._gen_harmonic(n, a-4) + mu1 = Hna1/Hna + mu2n = (Hna2*Hna - Hna1**2) + mu2d = Hna**2 + mu2 = mu2n / mu2d + g1 = (Hna3/Hna - 3*Hna1*Hna2/Hna**2 + 2*Hna1**3/Hna**3)/mu2**(3/2) + g2 = (Hna**3*Hna4 - 4*Hna**2*Hna1*Hna3 + 6*Hna*Hna1**2*Hna2 + - 3*Hna1**4) / mu2n**2 + g2 -= 3 + return mu1, mu2, g1, g2 + + +zipfian = zipfian_gen(a=1, name='zipfian', longname='A Zipfian') + + +class dlaplace_gen(rv_discrete): + r"""A Laplacian discrete random variable. + + %(before_notes)s + + Notes + ----- + The probability mass function for `dlaplace` is: + + .. math:: + + f(k) = \tanh(a/2) \exp(-a |k|) + + for integers :math:`k` and :math:`a > 0`. + + `dlaplace` takes :math:`a` as shape parameter. + + %(after_notes)s + + %(example)s + + """ + + def _shape_info(self): + return [_ShapeInfo("a", False, (0, np.inf), (False, False))] + + def _pmf(self, k, a): + # dlaplace.pmf(k) = tanh(a/2) * exp(-a*abs(k)) + return tanh(a/2.0) * exp(-a * abs(k)) + + def _cdf(self, x, a): + k = floor(x) + + def f1(k, a): + return 1.0 - exp(-a * k) / (exp(a) + 1) + + def f2(k, a): + return exp(a * (k + 1)) / (exp(a) + 1) + + return xpx.apply_where(k >= 0, (k, a), f1, f2) + + def _ppf(self, q, a): + const = 1 + exp(a) + vals = ceil(np.where(q < 1.0 / (1 + exp(-a)), + log(q*const) / a - 1, + -log((1-q) * const) / a)) + vals1 = vals - 1 + return np.where(self._cdf(vals1, a) >= q, vals1, vals) + + def _stats(self, a): + ea = exp(a) + mu2 = 2.*ea/(ea-1.)**2 + mu4 = 2.*ea*(ea**2+10.*ea+1.) / (ea-1.)**4 + return 0., mu2, 0., mu4/mu2**2 - 3. + + def _entropy(self, a): + return a / sinh(a) - log(tanh(a/2.0)) + + def _rvs(self, a, size=None, random_state=None): + # The discrete Laplace is equivalent to the two-sided geometric + # distribution with PMF: + # f(k) = (1 - alpha)/(1 + alpha) * alpha^abs(k) + # Reference: + # https://www.sciencedirect.com/science/ + # article/abs/pii/S0378375804003519 + # Furthermore, the two-sided geometric distribution is + # equivalent to the difference between two iid geometric + # distributions. + # Reference (page 179): + # https://pdfs.semanticscholar.org/61b3/ + # b99f466815808fd0d03f5d2791eea8b541a1.pdf + # Thus, we can leverage the following: + # 1) alpha = e^-a + # 2) probability_of_success = 1 - alpha (Bernoulli trial) + probOfSuccess = -np.expm1(-np.asarray(a)) + x = random_state.geometric(probOfSuccess, size=size) + y = random_state.geometric(probOfSuccess, size=size) + return x - y + + +dlaplace = dlaplace_gen(a=-np.inf, + name='dlaplace', longname='A discrete Laplacian') + + +class poisson_binom_gen(rv_discrete): + r"""A Poisson Binomial discrete random variable. + + %(before_notes)s + + See Also + -------- + binom + + Notes + ----- + The probability mass function for `poisson_binom` is: + + .. math:: + + f(k; p_1, p_2, ..., p_n) = \sum_{A \in F_k} \prod_{i \in A} p_i \prod_{j \in A^C} 1 - p_j + + where :math:`k \in \{0, 1, \dots, n-1, n\}`, :math:`F_k` is the set of all + subsets of :math:`k` integers that can be selected :math:`\{0, 1, \dots, n-1, n\}`, + and :math:`A^C` is the complement of a set :math:`A`. + + `poisson_binom` accepts a single array argument ``p`` for shape parameters + :math:`0 ≤ p_i ≤ 1`, where the last axis corresponds with the index :math:`i` and + any others are for batch dimensions. Broadcasting behaves according to the usual + rules except that the last axis of ``p`` is ignored. Instances of this class do + not support serialization/unserialization. + + %(after_notes)s + + References + ---------- + .. [1] "Poisson binomial distribution", Wikipedia, + https://en.wikipedia.org/wiki/Poisson_binomial_distribution + .. [2] Biscarri, William, Sihai Dave Zhao, and Robert J. Brunner. "A simple and + fast method for computing the Poisson binomial distribution function". + Computational Statistics & Data Analysis 122 (2018) 92-100. + :doi:`10.1016/j.csda.2018.01.007` + + %(example)s + + """ # noqa: E501 + + _parse_args_rvs: MethodType + _parse_args_stats: MethodType + _parse_args: MethodType + + def _shape_info(self): + # message = 'Fitting is not implemented for this distribution." + # raise NotImplementedError(message) + return [] + + def _argcheck(self, *args): + p = np.stack(args, axis=0) + conds = (0 <= p) & (p <= 1) + return np.all(conds, axis=0) + + def _rvs(self, *args, size=None, random_state=None): + # convenient to work along the last axis here to avoid interference with `size` + p = np.stack(args, axis=-1) + # Size passed by the user is the *shape of the returned array*, so it won't + # contain the length of the last axis of p. + size = (p.shape if size is None else + (size, 1) if np.isscalar(size) else tuple(size) + (1,)) + size = np.broadcast_shapes(p.shape, size) + return bernoulli._rvs(p, size=size, random_state=random_state).sum(axis=-1) + + def _get_support(self, *args): + return 0, len(args) + + def _pmf(self, k, *args): + k = np.atleast_1d(k).astype(np.int64) + k, *args = np.broadcast_arrays(k, *args) + p = np.stack(args, dtype=np.float64, axis=-1) + return _poisson_binom_pmf(k, p) + + def _cdf(self, k, *args): + k = np.atleast_1d(k).astype(np.int64) + k, *args = np.broadcast_arrays(k, *args) + p = np.stack(args, dtype=np.float64, axis=-1) + return _poisson_binom_cdf(k, p) + + def _stats(self, *args, **kwds): + p = np.stack(args, axis=0) + mean = np.sum(p, axis=0) + var = np.sum(p * (1-p), axis=0) + return (mean, var, None, None) + + def __call__(self, *args, **kwds): + return poisson_binomial_frozen(self, *args, **kwds) + + +poisson_binom = poisson_binom_gen(name='poisson_binom', longname='A Poisson binomial', + shapes='p') + +# The _parse_args methods don't work with vector-valued shape parameters, so we rewrite +# them. Note that `p` is accepted as an array with the index `i` of `p_i` corresponding +# with the last axis; we return it as a tuple (p_1, p_2, ..., p_n) so that it looks +# like `n` scalar (or arrays of scalar-valued) shape parameters to the infrastructure. + +def _parse_args_rvs(self, p, loc=0, size=None): + return tuple(np.moveaxis(p, -1, 0)), loc, 1.0, size + +def _parse_args_stats(self, p, loc=0, moments='mv'): + return tuple(np.moveaxis(p, -1, 0)), loc, 1.0, moments + +def _parse_args(self, p, loc=0): + return tuple(np.moveaxis(p, -1, 0)), loc, 1.0 + +# The infrastructure manually binds these methods to the instance, so +# we can only override them by manually binding them, too. +_pb_obj, _pb_cls = poisson_binom, poisson_binom_gen # shorter names (for PEP8) +poisson_binom._parse_args_rvs = _parse_args_rvs.__get__(_pb_obj, _pb_cls) +poisson_binom._parse_args_stats = _parse_args_stats.__get__(_pb_obj, _pb_cls) +poisson_binom._parse_args = _parse_args.__get__(_pb_obj, _pb_cls) + +class poisson_binomial_frozen(rv_discrete_frozen): + # copied from rv_frozen; we just need to bind the `_parse_args` methods + def __init__(self, dist, *args, **kwds): # verbatim + self.args = args # verbatim + self.kwds = kwds # verbatim + + # create a new instance # verbatim + self.dist = dist.__class__(**dist._updated_ctor_param()) # verbatim + + # Here is the only modification + self.dist._parse_args_rvs = _parse_args_rvs.__get__(_pb_obj, _pb_cls) + self.dist._parse_args_stats = _parse_args_stats.__get__(_pb_obj, _pb_cls) + self.dist._parse_args = _parse_args.__get__(_pb_obj, _pb_cls) + + shapes, _, _ = self.dist._parse_args(*args, **kwds) # verbatim + self.a, self.b = self.dist._get_support(*shapes) # verbatim + + def expect(self, func=None, lb=None, ub=None, conditional=False, **kwds): + a, loc, scale = self.dist._parse_args(*self.args, **self.kwds) + # Here's the modification: we pass all args (including `loc`) into the `args` + # parameter of `expect` so the shape only goes through `_parse_args` once. + return self.dist.expect(func, self.args, loc, lb, ub, conditional, **kwds) + + +class skellam_gen(rv_discrete): + r"""A Skellam discrete random variable. + + %(before_notes)s + + Notes + ----- + Probability distribution of the difference of two correlated or + uncorrelated Poisson random variables. + + Let :math:`k_1` and :math:`k_2` be two Poisson-distributed r.v. with + expected values :math:`\lambda_1` and :math:`\lambda_2`. Then, + :math:`k_1 - k_2` follows a Skellam distribution with parameters + :math:`\mu_1 = \lambda_1 - \rho \sqrt{\lambda_1 \lambda_2}` and + :math:`\mu_2 = \lambda_2 - \rho \sqrt{\lambda_1 \lambda_2}`, where + :math:`\rho` is the correlation coefficient between :math:`k_1` and + :math:`k_2`. If the two Poisson-distributed r.v. are independent then + :math:`\rho = 0`. + + Parameters :math:`\mu_1` and :math:`\mu_2` must be strictly positive. + + `skellam` takes :math:`\mu_1` and :math:`\mu_2` as shape parameters. + + %(after_notes)s + + References + ---------- + .. [1] Skellam, J. G. "The Frequency Distribution of the Difference + Between Two Poisson Variates Belonging to Different Populations." + *Journal of the Royal Statistical Society* 109, no. 3 (1946): 296-296. + :doi:`10.2307/2981372` + .. [2] "Skellam distribution", Wikipedia, + https://en.wikipedia.org/wiki/Skellam_distribution + + %(example)s + + """ + def _shape_info(self): + return [_ShapeInfo("mu1", False, (0, np.inf), (False, False)), + _ShapeInfo("mu2", False, (0, np.inf), (False, False))] + + def _rvs(self, mu1, mu2, size=None, random_state=None): + n = size + return (random_state.poisson(mu1, n) - + random_state.poisson(mu2, n)) + + def _pmf(self, x, mu1, mu2): + with np.errstate(over='ignore'): # see gh-17432 + px = np.where(x < 0, + scu._ncx2_pdf(2*mu2, 2*(1-x), 2*mu1)*2, + scu._ncx2_pdf(2*mu1, 2*(1+x), 2*mu2)*2) + # ncx2.pdf() returns nan's for extremely low probabilities + return px + + def _cdf(self, x, mu1, mu2): + x = floor(x) + with np.errstate(over='ignore'): # see gh-17432 + px = np.where(x < 0, + special.chndtr(2*mu2, -2*x, 2*mu1), + scu._ncx2_sf(2*mu1, 2*(x+1), 2*mu2)) + return px + + def _stats(self, mu1, mu2): + mean = mu1 - mu2 + var = mu1 + mu2 + g1 = mean / sqrt((var)**3) + g2 = 1 / var + return mean, var, g1, g2 + + +skellam = skellam_gen(a=-np.inf, name="skellam", longname='A Skellam') + + +class yulesimon_gen(rv_discrete): + r"""A Yule-Simon discrete random variable. + + %(before_notes)s + + Notes + ----- + + The probability mass function for the `yulesimon` is: + + .. math:: + + f(k) = \alpha B(k, \alpha+1) + + for :math:`k=1,2,3,...`, where :math:`\alpha>0`. + Here :math:`B` refers to the `scipy.special.beta` function. + + The sampling of random variates is based on pg 553, Section 6.3 of [1]_. + Our notation maps to the referenced logic via :math:`\alpha=a-1`. + + For details see the wikipedia entry [2]_. + + References + ---------- + .. [1] Devroye, Luc. "Non-uniform Random Variate Generation", + (1986) Springer, New York. + + .. [2] https://en.wikipedia.org/wiki/Yule-Simon_distribution + + %(after_notes)s + + %(example)s + + """ + def _shape_info(self): + return [_ShapeInfo("alpha", False, (0, np.inf), (False, False))] + + def _rvs(self, alpha, size=None, random_state=None): + E1 = random_state.standard_exponential(size) + E2 = random_state.standard_exponential(size) + ans = ceil(-E1 / log1p(-exp(-E2 / alpha))) + return ans + + def _pmf(self, x, alpha): + return alpha * special.beta(x, alpha + 1) + + def _argcheck(self, alpha): + return (alpha > 0) + + def _logpmf(self, x, alpha): + return log(alpha) + special.betaln(x, alpha + 1) + + def _cdf(self, x, alpha): + return 1 - x * special.beta(x, alpha + 1) + + def _sf(self, x, alpha): + return x * special.beta(x, alpha + 1) + + def _logsf(self, x, alpha): + return log(x) + special.betaln(x, alpha + 1) + + def _stats(self, alpha): + mu = np.where(alpha <= 1, np.inf, alpha / (alpha - 1)) + mu2 = np.where(alpha > 2, + alpha**2 / ((alpha - 2.0) * (alpha - 1)**2), + np.inf) + mu2 = np.where(alpha <= 1, np.nan, mu2) + g1 = np.where(alpha > 3, + sqrt(alpha - 2) * (alpha + 1)**2 / (alpha * (alpha - 3)), + np.inf) + g1 = np.where(alpha <= 2, np.nan, g1) + g2 = np.where(alpha > 4, + alpha + 3 + ((11 * alpha**3 - 49 * alpha - 22) / + (alpha * (alpha - 4) * (alpha - 3))), + np.inf) + g2 = np.where(alpha <= 2, np.nan, g2) + return mu, mu2, g1, g2 + + +yulesimon = yulesimon_gen(name='yulesimon', a=1) + + +class _nchypergeom_gen(rv_discrete): + r"""A noncentral hypergeometric discrete random variable. + + For subclassing by nchypergeom_fisher_gen and nchypergeom_wallenius_gen. + + """ + + rvs_name = None + dist = None + + def _shape_info(self): + return [_ShapeInfo("M", True, (0, np.inf), (True, False)), + _ShapeInfo("n", True, (0, np.inf), (True, False)), + _ShapeInfo("N", True, (0, np.inf), (True, False)), + _ShapeInfo("odds", False, (0, np.inf), (False, False))] + + def _get_support(self, M, n, N, odds): + N, m1, n = M, n, N # follow Wikipedia notation + m2 = N - m1 + x_min = np.maximum(0, n - m2) + x_max = np.minimum(n, m1) + return x_min, x_max + + def _argcheck(self, M, n, N, odds): + M, n = np.asarray(M), np.asarray(n), + N, odds = np.asarray(N), np.asarray(odds) + cond1 = (~np.isnan(M)) & (M.astype(int) == M) & (M >= 0) + cond2 = (~np.isnan(n)) & (n.astype(int) == n) & (n >= 0) + cond3 = (~np.isnan(N)) & (N.astype(int) == N) & (N >= 0) + cond4 = odds > 0 + cond5 = N <= M + cond6 = n <= M + return cond1 & cond2 & cond3 & cond4 & cond5 & cond6 + + def _rvs(self, M, n, N, odds, size=None, random_state=None): + + @_vectorize_rvs_over_shapes + def _rvs1(M, n, N, odds, size, random_state): + if np.isnan(M) | np.isnan(n) | np.isnan(N): + return np.full(size, np.nan) + length = np.prod(size) + urn = _PyStochasticLib3() + rv_gen = getattr(urn, self.rvs_name) + rvs = rv_gen(N, n, M, odds, length, random_state) + rvs = rvs.reshape(size) + return rvs + + return _rvs1(M, n, N, odds, size=size, random_state=random_state) + + def _pmf(self, x, M, n, N, odds): + + x, M, n, N, odds = np.broadcast_arrays(x, M, n, N, odds) + if x.size == 0: # np.vectorize doesn't work with zero size input + return np.empty_like(x) + + @np.vectorize + def _pmf1(x, M, n, N, odds): + if np.isnan(x) | np.isnan(M) | np.isnan(n) | np.isnan(N): + return np.nan + urn = self.dist(N, n, M, odds, 1e-12) + return urn.probability(x) + + return _pmf1(x, M, n, N, odds) + + def _stats(self, M, n, N, odds, moments='mv'): + + @np.vectorize + def _moments1(M, n, N, odds): + if np.isnan(M) | np.isnan(n) | np.isnan(N): + return np.nan, np.nan + urn = self.dist(N, n, M, odds, 1e-12) + return urn.moments() + + m, v = (_moments1(M, n, N, odds) if ("m" in moments or "v" in moments) + else (None, None)) + s, k = None, None + return m, v, s, k + + +class nchypergeom_fisher_gen(_nchypergeom_gen): + r"""A Fisher's noncentral hypergeometric discrete random variable. + + Fisher's noncentral hypergeometric distribution models drawing objects of + two types from a bin. `M` is the total number of objects, `n` is the + number of Type I objects, and `odds` is the odds ratio: the odds of + selecting a Type I object rather than a Type II object when there is only + one object of each type. + The random variate represents the number of Type I objects drawn if we + take a handful of objects from the bin at once and find out afterwards + that we took `N` objects. + + %(before_notes)s + + See Also + -------- + nchypergeom_wallenius, hypergeom, nhypergeom + + Notes + ----- + Let mathematical symbols :math:`N`, :math:`n`, and :math:`M` correspond + with parameters `N`, `n`, and `M` (respectively) as defined above. + + The probability mass function is defined as + + .. math:: + + p(x; M, n, N, \omega) = + \frac{\binom{n}{x}\binom{M - n}{N-x}\omega^x}{P_0}, + + for + :math:`x \in [x_l, x_u]`, + :math:`M \in {\mathbb N}`, + :math:`n \in [0, M]`, + :math:`N \in [0, M]`, + :math:`\omega > 0`, + where + :math:`x_l = \max(0, N - (M - n))`, + :math:`x_u = \min(N, n)`, + + .. math:: + + P_0 = \sum_{y=x_l}^{x_u} \binom{n}{y}\binom{M - n}{N-y}\omega^y, + + and the binomial coefficients are defined as + + .. math:: \binom{n}{k} \equiv \frac{n!}{k! (n - k)!}. + + `nchypergeom_fisher` uses the BiasedUrn package by Agner Fog with + permission for it to be distributed under SciPy's license. + + The symbols used to denote the shape parameters (`N`, `n`, and `M`) are not + universally accepted; they are chosen for consistency with `hypergeom`. + + Note that Fisher's noncentral hypergeometric distribution is distinct + from Wallenius' noncentral hypergeometric distribution, which models + drawing a pre-determined `N` objects from a bin one by one. + When the odds ratio is unity, however, both distributions reduce to the + ordinary hypergeometric distribution. + + %(after_notes)s + + References + ---------- + .. [1] Agner Fog, "Biased Urn Theory". + https://cran.r-project.org/web/packages/BiasedUrn/vignettes/UrnTheory.pdf + + .. [2] "Fisher's noncentral hypergeometric distribution", Wikipedia, + https://en.wikipedia.org/wiki/Fisher's_noncentral_hypergeometric_distribution + + %(example)s + + """ + + rvs_name = "rvs_fisher" + dist = _PyFishersNCHypergeometric + + +nchypergeom_fisher = nchypergeom_fisher_gen( + name='nchypergeom_fisher', + longname="A Fisher's noncentral hypergeometric") + + +class nchypergeom_wallenius_gen(_nchypergeom_gen): + r"""A Wallenius' noncentral hypergeometric discrete random variable. + + Wallenius' noncentral hypergeometric distribution models drawing objects of + two types from a bin. `M` is the total number of objects, `n` is the + number of Type I objects, and `odds` is the odds ratio: the odds of + selecting a Type I object rather than a Type II object when there is only + one object of each type. + The random variate represents the number of Type I objects drawn if we + draw a pre-determined `N` objects from a bin one by one. + + %(before_notes)s + + See Also + -------- + nchypergeom_fisher, hypergeom, nhypergeom + + Notes + ----- + Let mathematical symbols :math:`N`, :math:`n`, and :math:`M` correspond + with parameters `N`, `n`, and `M` (respectively) as defined above. + + The probability mass function is defined as + + .. math:: + + p(x; N, n, M) = \binom{n}{x} \binom{M - n}{N-x} + \int_0^1 \left(1-t^{\omega/D}\right)^x\left(1-t^{1/D}\right)^{N-x} dt + + for + :math:`x \in [x_l, x_u]`, + :math:`M \in {\mathbb N}`, + :math:`n \in [0, M]`, + :math:`N \in [0, M]`, + :math:`\omega > 0`, + where + :math:`x_l = \max(0, N - (M - n))`, + :math:`x_u = \min(N, n)`, + + .. math:: + + D = \omega(n - x) + ((M - n)-(N-x)), + + and the binomial coefficients are defined as + + .. math:: \binom{n}{k} \equiv \frac{n!}{k! (n - k)!}. + + `nchypergeom_wallenius` uses the BiasedUrn package by Agner Fog with + permission for it to be distributed under SciPy's license. + + The symbols used to denote the shape parameters (`N`, `n`, and `M`) are not + universally accepted; they are chosen for consistency with `hypergeom`. + + Note that Wallenius' noncentral hypergeometric distribution is distinct + from Fisher's noncentral hypergeometric distribution, which models + take a handful of objects from the bin at once, finding out afterwards + that `N` objects were taken. + When the odds ratio is unity, however, both distributions reduce to the + ordinary hypergeometric distribution. + + %(after_notes)s + + References + ---------- + .. [1] Agner Fog, "Biased Urn Theory". + https://cran.r-project.org/web/packages/BiasedUrn/vignettes/UrnTheory.pdf + + .. [2] "Wallenius' noncentral hypergeometric distribution", Wikipedia, + https://en.wikipedia.org/wiki/Wallenius'_noncentral_hypergeometric_distribution + + %(example)s + + """ + + rvs_name = "rvs_wallenius" + dist = _PyWalleniusNCHypergeometric + + +nchypergeom_wallenius = nchypergeom_wallenius_gen( + name='nchypergeom_wallenius', + longname="A Wallenius' noncentral hypergeometric") + + +# Collect names of classes and objects in this module. +pairs = list(globals().copy().items()) +_distn_names, _distn_gen_names = get_distribution_names(pairs, rv_discrete) + +__all__ = _distn_names + _distn_gen_names diff --git a/v12x/pylib/scipy/stats/_distn_infrastructure.py b/v12x/pylib/scipy/stats/_distn_infrastructure.py new file mode 100644 index 0000000000000000000000000000000000000000..457b9ddb511d4dcba37e375a77c1a8a1dedd6554 --- /dev/null +++ b/v12x/pylib/scipy/stats/_distn_infrastructure.py @@ -0,0 +1,4274 @@ +# +# Author: Travis Oliphant 2002-2011 with contributions from +# SciPy Developers 2004-2011 +# +from scipy._lib._util import getfullargspec_no_self as _getfullargspec + +import sys +import keyword +import re +import types +import warnings +from itertools import zip_longest + +from scipy._lib import doccer +from scipy._lib._docscrape import FunctionDoc +from ._distr_params import distcont, distdiscrete +from scipy._lib._util import check_random_state +import scipy._external.array_api_extra as xpx + +from scipy.special import comb, entr + + +# for root finding for continuous distribution ppf, and maximum likelihood +# estimation +from scipy import optimize + +# for functions of continuous distributions (e.g. moments, entropy, cdf) +from scipy import integrate + +# to approximate the pdf of a continuous distribution given its cdf +from scipy.stats._finite_differences import _derivative + +# for scipy.stats.entropy. Attempts to import just that function or file +# have cause import problems +from scipy import stats + +from numpy import (arange, putmask, ones, shape, ndarray, zeros, floor, + logical_and, log, sqrt, place, argmax, vectorize, asarray, + nan, inf, isinf, empty) + +import numpy as np +from ._constants import _XMAX, _LOGXMAX +from ._censored_data import CensoredData +from scipy.stats._warnings_errors import FitError + +# These are the docstring parts used for substitution in specific +# distribution docstrings + +docheaders = {'methods': """\nMethods\n-------\n""", + 'notes': """\nNotes\n-----\n""", + 'examples': """\nExamples\n--------\n"""} + +_doc_rvs = """\ +rvs(%(shapes)s, loc=0, scale=1, size=1, random_state=None) + Random variates. +""" +_doc_pdf = """\ +pdf(x, %(shapes)s, loc=0, scale=1) + Probability density function. +""" +_doc_logpdf = """\ +logpdf(x, %(shapes)s, loc=0, scale=1) + Log of the probability density function. +""" +_doc_pmf = """\ +pmf(k, %(shapes)s, loc=0, scale=1) + Probability mass function. +""" +_doc_logpmf = """\ +logpmf(k, %(shapes)s, loc=0, scale=1) + Log of the probability mass function. +""" +_doc_cdf = """\ +cdf(x, %(shapes)s, loc=0, scale=1) + Cumulative distribution function. +""" +_doc_logcdf = """\ +logcdf(x, %(shapes)s, loc=0, scale=1) + Log of the cumulative distribution function. +""" +_doc_sf = """\ +sf(x, %(shapes)s, loc=0, scale=1) + Survival function (also defined as ``1 - cdf``, but `sf` is sometimes more accurate). +""" # noqa: E501 +_doc_logsf = """\ +logsf(x, %(shapes)s, loc=0, scale=1) + Log of the survival function. +""" +_doc_ppf = """\ +ppf(q, %(shapes)s, loc=0, scale=1) + Percent point function (inverse of ``cdf`` --- percentiles). +""" +_doc_isf = """\ +isf(q, %(shapes)s, loc=0, scale=1) + Inverse survival function (inverse of ``sf``). +""" +_doc_moment = """\ +moment(order, %(shapes)s, loc=0, scale=1) + Non-central moment of the specified order. +""" +_doc_stats = """\ +stats(%(shapes)s, loc=0, scale=1, moments='mv') + Mean('m'), variance('v'), skew('s'), and/or kurtosis('k'). +""" +_doc_entropy = """\ +entropy(%(shapes)s, loc=0, scale=1) + (Differential) entropy of the RV. +""" +_doc_fit = """\ +fit(data) + Parameter estimates for generic data. + See `scipy.stats.rv_continuous.fit `__ for detailed documentation of the + keyword arguments. +""" # noqa: E501 +_doc_expect = """\ +expect(func, args=(%(shapes_)s), loc=0, scale=1, lb=None, ub=None, conditional=False, **kwds) + Expected value of a function (of one argument) with respect to the distribution. +""" # noqa: E501 +_doc_expect_discrete = """\ +expect(func, args=(%(shapes_)s), loc=0, lb=None, ub=None, conditional=False) + Expected value of a function (of one argument) with respect to the distribution. +""" +_doc_median = """\ +median(%(shapes)s, loc=0, scale=1) + Median of the distribution. +""" +_doc_mean = """\ +mean(%(shapes)s, loc=0, scale=1) + Mean of the distribution. +""" +_doc_var = """\ +var(%(shapes)s, loc=0, scale=1) + Variance of the distribution. +""" +_doc_std = """\ +std(%(shapes)s, loc=0, scale=1) + Standard deviation of the distribution. +""" +_doc_interval = """\ +interval(confidence, %(shapes)s, loc=0, scale=1) + Confidence interval with equal areas around the median. +""" +_doc_allmethods = ''.join([docheaders['methods'], _doc_rvs, _doc_pdf, + _doc_logpdf, _doc_cdf, _doc_logcdf, _doc_sf, + _doc_logsf, _doc_ppf, _doc_isf, _doc_moment, + _doc_stats, _doc_entropy, _doc_fit, + _doc_expect, _doc_median, + _doc_mean, _doc_var, _doc_std, _doc_interval]) + +_doc_default_longsummary = """\ +As an instance of the `rv_continuous` class, `%(name)s` object inherits from it +a collection of generic methods (see below for the full list), +and completes them with details specific for this particular distribution. +""" + +_doc_default_frozen_note = """ +Alternatively, the object may be called (as a function) to fix the shape, +location, and scale parameters returning a "frozen" continuous RV object: + +rv = %(name)s(%(shapes)s, loc=0, scale=1) + - Frozen RV object with the same methods but holding the given shape, + location, and scale fixed. +""" +_doc_default_example = """\ +Examples +-------- +>>> import numpy as np +>>> from scipy.stats import %(name)s +>>> import matplotlib.pyplot as plt +>>> fig, ax = plt.subplots(1, 1) + +Get the support: + +%(set_vals_stmt)s +>>> lb, ub = %(name)s.support(%(shapes)s) + +Calculate the first four moments: + +>>> mean, var, skew, kurt = %(name)s.stats(%(shapes)s, moments='mvsk') + +Display the probability density function (``pdf``): + +>>> x = np.linspace(%(name)s.ppf(0.01, %(shapes)s), +... %(name)s.ppf(0.99, %(shapes)s), 100) +>>> ax.plot(x, %(name)s.pdf(x, %(shapes)s), +... 'r-', lw=5, alpha=0.6, label='%(name)s pdf') + +Alternatively, the distribution object can be called (as a function) +to fix the shape, location and scale parameters. This returns a "frozen" +RV object holding the given parameters fixed. + +Freeze the distribution and display the frozen ``pdf``: + +>>> rv = %(name)s(%(shapes)s) +>>> ax.plot(x, rv.pdf(x), 'k-', lw=2, label='frozen pdf') + +Check accuracy of ``cdf`` and ``ppf``: + +>>> vals = %(name)s.ppf([0.001, 0.5, 0.999], %(shapes)s) +>>> np.allclose([0.001, 0.5, 0.999], %(name)s.cdf(vals, %(shapes)s)) +True + +Generate random numbers: + +>>> r = %(name)s.rvs(%(shapes)s, size=1000) + +And compare the histogram: + +>>> ax.hist(r, density=True, bins='auto', histtype='stepfilled', alpha=0.2) +>>> ax.set_xlim([x[0], x[-1]]) +>>> ax.legend(loc='best', frameon=False) +>>> plt.show() + +""" + +_doc_default_locscale = """\ +The probability density above is defined in the "standardized" form. To shift +and/or scale the distribution use the ``loc`` and ``scale`` parameters. +Specifically, ``%(name)s.pdf(x, %(shapes)s, loc, scale)`` is identically +equivalent to ``%(name)s.pdf(y, %(shapes)s) / scale`` with +``y = (x - loc) / scale``. Note that shifting the location of a distribution +does not make it a "noncentral" distribution; noncentral generalizations of +some distributions are available in separate classes. +""" + +_doc_default = ''.join([_doc_default_longsummary, + _doc_allmethods, + '\n', + _doc_default_example]) + +_doc_default_before_notes = ''.join([_doc_default_longsummary, + _doc_allmethods]) + +docdict = { + 'rvs': _doc_rvs, + 'pdf': _doc_pdf, + 'logpdf': _doc_logpdf, + 'cdf': _doc_cdf, + 'logcdf': _doc_logcdf, + 'sf': _doc_sf, + 'logsf': _doc_logsf, + 'ppf': _doc_ppf, + 'isf': _doc_isf, + 'stats': _doc_stats, + 'entropy': _doc_entropy, + 'fit': _doc_fit, + 'moment': _doc_moment, + 'expect': _doc_expect, + 'interval': _doc_interval, + 'mean': _doc_mean, + 'std': _doc_std, + 'var': _doc_var, + 'median': _doc_median, + 'allmethods': _doc_allmethods, + 'longsummary': _doc_default_longsummary, + 'frozennote': _doc_default_frozen_note, + 'example': _doc_default_example, + 'default': _doc_default, + 'before_notes': _doc_default_before_notes, + 'after_notes': _doc_default_locscale +} + +# Reuse common content between continuous and discrete docs, change some +# minor bits. +docdict_discrete = docdict.copy() + +docdict_discrete['pmf'] = _doc_pmf +docdict_discrete['logpmf'] = _doc_logpmf +docdict_discrete['expect'] = _doc_expect_discrete +_doc_disc_methods = ['rvs', 'pmf', 'logpmf', 'cdf', 'logcdf', 'sf', 'logsf', + 'ppf', 'isf', 'stats', 'entropy', 'expect', 'median', + 'mean', 'var', 'std', 'interval'] +for obj in _doc_disc_methods: + docdict_discrete[obj] = docdict_discrete[obj].replace(', scale=1', '') + +_doc_disc_methods_err_varname = ['cdf', 'logcdf', 'sf', 'logsf'] +for obj in _doc_disc_methods_err_varname: + docdict_discrete[obj] = docdict_discrete[obj].replace('(x, ', '(k, ') + +docdict_discrete.pop('pdf') +docdict_discrete.pop('logpdf') + +_doc_allmethods = ''.join([docdict_discrete[obj] for obj in _doc_disc_methods]) +docdict_discrete['allmethods'] = docheaders['methods'] + _doc_allmethods + +docdict_discrete['longsummary'] = _doc_default_longsummary.replace( + 'rv_continuous', 'rv_discrete') + +_doc_default_frozen_note = """ +Alternatively, the object may be called (as a function) to fix the shape and +location parameters returning a "frozen" discrete RV object: + +rv = %(name)s(%(shapes)s, loc=0) + - Frozen RV object with the same methods but holding the given shape and + location fixed. +""" +docdict_discrete['frozennote'] = _doc_default_frozen_note + +_doc_default_discrete_example = """\ +Examples +-------- +>>> import numpy as np +>>> from scipy.stats import %(name)s +>>> import matplotlib.pyplot as plt +>>> fig, ax = plt.subplots(1, 1) + +Get the support: + +%(set_vals_stmt)s +>>> lb, ub = %(name)s.support(%(shapes)s) + +Calculate the first four moments: + +>>> mean, var, skew, kurt = %(name)s.stats(%(shapes)s, moments='mvsk') + +Display the probability mass function (``pmf``): + +>>> x = np.arange(%(name)s.ppf(0.01, %(shapes)s), +... %(name)s.ppf(0.99, %(shapes)s)) +>>> ax.plot(x, %(name)s.pmf(x, %(shapes)s), 'bo', ms=8, label='%(name)s pmf') +>>> ax.vlines(x, 0, %(name)s.pmf(x, %(shapes)s), colors='b', lw=5, alpha=0.5) + +Alternatively, the distribution object can be called (as a function) +to fix the shape and location. This returns a "frozen" RV object holding +the given parameters fixed. + +Freeze the distribution and display the frozen ``pmf``: + +>>> rv = %(name)s(%(shapes)s) +>>> ax.vlines(x, 0, rv.pmf(x), colors='k', linestyles='-', lw=1, +... label='frozen pmf') +>>> ax.legend(loc='best', frameon=False) +>>> plt.show() + +Check accuracy of ``cdf`` and ``ppf``: + +>>> prob = %(name)s.cdf(x, %(shapes)s) +>>> np.allclose(x, %(name)s.ppf(prob, %(shapes)s)) +True + +Generate random numbers: + +>>> r = %(name)s.rvs(%(shapes)s, size=1000) +""" + + +_doc_default_discrete_locscale = """\ +The probability mass function above is defined in the "standardized" form. +To shift distribution use the ``loc`` parameter. +Specifically, ``%(name)s.pmf(k, %(shapes)s, loc)`` is identically +equivalent to ``%(name)s.pmf(k - loc, %(shapes)s)``. +""" + +docdict_discrete['example'] = _doc_default_discrete_example +docdict_discrete['after_notes'] = _doc_default_discrete_locscale + +_doc_default_before_notes = ''.join([docdict_discrete['longsummary'], + docdict_discrete['allmethods']]) +docdict_discrete['before_notes'] = _doc_default_before_notes + +_doc_default_disc = ''.join([docdict_discrete['longsummary'], + docdict_discrete['allmethods'], + docdict_discrete['frozennote'], + docdict_discrete['example']]) +docdict_discrete['default'] = _doc_default_disc + +# clean up all the separate docstring elements, we do not need them anymore +for obj in [s for s in dir() if s.startswith('_doc_')]: + exec('del ' + obj) +del obj + + +def _moment(data, n, mu=None): + if mu is None: + mu = data.mean() + return ((data - mu)**n).mean() + + +def _moment_from_stats(n, mu, mu2, g1, g2, moment_func, args): + if (n == 0): + return 1.0 + elif (n == 1): + if mu is None: + val = moment_func(1, *args) + else: + val = mu + elif (n == 2): + if mu2 is None or mu is None: + val = moment_func(2, *args) + else: + val = mu2 + mu*mu + elif (n == 3): + if g1 is None or mu2 is None or mu is None: + val = moment_func(3, *args) + else: + mu3 = g1 * np.power(mu2, 1.5) # 3rd central moment + val = mu3+3*mu*mu2+mu*mu*mu # 3rd non-central moment + elif (n == 4): + if g1 is None or g2 is None or mu2 is None or mu is None: + val = moment_func(4, *args) + else: + mu4 = (g2+3.0)*(mu2**2.0) # 4th central moment + mu3 = g1*np.power(mu2, 1.5) # 3rd central moment + val = mu4+4*mu*mu3+6*mu*mu*mu2+mu*mu*mu*mu + else: + val = moment_func(n, *args) + + return val + + +def _skew(data): + """ + skew is third central moment / variance**(1.5) + """ + data = np.ravel(data) + mu = data.mean() + m2 = ((data - mu)**2).mean() + m3 = ((data - mu)**3).mean() + return m3 / np.power(m2, 1.5) + + +def _kurtosis(data): + """Fisher's excess kurtosis is fourth central moment / variance**2 - 3.""" + data = np.ravel(data) + mu = data.mean() + m2 = ((data - mu)**2).mean() + m4 = ((data - mu)**4).mean() + return m4 / m2**2 - 3 + +def _vectorize_rvs_over_shapes(_rvs1): + """Decorator that vectorizes _rvs method to work on ndarray shapes""" + # _rvs1 must be a _function_ that accepts _scalar_ args as positional + # arguments, `size` and `random_state` as keyword arguments. + # _rvs1 must return a random variate array with shape `size`. If `size` is + # None, _rvs1 must return a scalar. + # When applied to _rvs1, this decorator broadcasts ndarray args + # and loops over them, calling _rvs1 for each set of scalar args. + # For usage example, see _nchypergeom_gen + def _rvs(*args, size, random_state): + _rvs1_size, _rvs1_indices = _check_shape(args[0].shape, size) + + size = np.array(size) + _rvs1_size = np.array(_rvs1_size) + _rvs1_indices = np.array(_rvs1_indices) + + if np.all(_rvs1_indices): # all args are scalars + return _rvs1(*args, size, random_state) + + out = np.empty(size) + + # out.shape can mix dimensions associated with arg_shape and _rvs1_size + # Sort them to arg_shape + _rvs1_size for easy indexing of dimensions + # corresponding with the different sets of scalar args + j0 = np.arange(out.ndim) + j1 = np.hstack((j0[~_rvs1_indices], j0[_rvs1_indices])) + out = np.moveaxis(out, j1, j0) + + for i in np.ndindex(*size[~_rvs1_indices]): + # arg can be squeezed because singleton dimensions will be + # associated with _rvs1_size, not arg_shape per _check_shape + out[i] = _rvs1(*[np.squeeze(arg)[i] for arg in args], + _rvs1_size, random_state) + + return np.moveaxis(out, j0, j1) # move axes back before returning + return _rvs + + +def _fit_determine_optimizer(optimizer): + if not callable(optimizer) and isinstance(optimizer, str): + if not optimizer.startswith('fmin_'): + optimizer = "fmin_"+optimizer + if optimizer == 'fmin_': + optimizer = 'fmin' + try: + optimizer = getattr(optimize, optimizer) + except AttributeError as e: + raise ValueError(f"{optimizer} is not a valid optimizer") from e + return optimizer + +def _isintegral(x): + return x == np.round(x) + +def _sum_finite(x): + """ + For a 1D array x, return a tuple containing the sum of the + finite values of x and the number of nonfinite values. + + This is a utility function used when evaluating the negative + loglikelihood for a distribution and an array of samples. + + Examples + -------- + >>> import numpy as np + >>> from scipy.stats._distn_infrastructure import _sum_finite + >>> tot, nbad = _sum_finite(np.array([-2, -np.inf, 5, 1])) + >>> tot + 4.0 + >>> nbad + 1 + """ + finite_x = np.isfinite(x) + bad_count = finite_x.size - np.count_nonzero(finite_x) + return np.sum(x[finite_x]), bad_count + + +# Frozen RV class +class rv_frozen: + + # generic type compatibility with scipy-stubs + __class_getitem__: classmethod = classmethod(types.GenericAlias) + + def __init__(self, dist, *args, **kwds): + self.args = args + self.kwds = kwds + + # create a new instance + self.dist = dist.__class__(**dist._updated_ctor_param()) + + shapes, _, _ = self.dist._parse_args(*args, **kwds) + self.a, self.b = self.dist._get_support(*shapes) + + @property + def random_state(self): + return self.dist._random_state + + @random_state.setter + def random_state(self, seed): + self.dist._random_state = check_random_state(seed) + + def cdf(self, x): + return self.dist.cdf(x, *self.args, **self.kwds) + + def logcdf(self, x): + return self.dist.logcdf(x, *self.args, **self.kwds) + + def ppf(self, q): + return self.dist.ppf(q, *self.args, **self.kwds) + + def isf(self, q): + return self.dist.isf(q, *self.args, **self.kwds) + + def rvs(self, size=None, random_state=None): + kwds = self.kwds.copy() + kwds.update({'size': size, 'random_state': random_state}) + return self.dist.rvs(*self.args, **kwds) + + def sf(self, x): + return self.dist.sf(x, *self.args, **self.kwds) + + def logsf(self, x): + return self.dist.logsf(x, *self.args, **self.kwds) + + def stats(self, moments='mv'): + kwds = self.kwds.copy() + kwds.update({'moments': moments}) + return self.dist.stats(*self.args, **kwds) + + def median(self): + return self.dist.median(*self.args, **self.kwds) + + def mean(self): + return self.dist.mean(*self.args, **self.kwds) + + def var(self): + return self.dist.var(*self.args, **self.kwds) + + def std(self): + return self.dist.std(*self.args, **self.kwds) + + def moment(self, order=None): + return self.dist.moment(order, *self.args, **self.kwds) + + def entropy(self): + return self.dist.entropy(*self.args, **self.kwds) + + def interval(self, confidence=None): + return self.dist.interval(confidence, *self.args, **self.kwds) + + def expect(self, func=None, lb=None, ub=None, conditional=False, **kwds): + # expect method only accepts shape parameters as positional args + # hence convert self.args, self.kwds, also loc/scale + # See the .expect method docstrings for the meaning of + # other parameters. + a, loc, scale = self.dist._parse_args(*self.args, **self.kwds) + if isinstance(self.dist, rv_discrete): + return self.dist.expect(func, a, loc, lb, ub, conditional, **kwds) + else: + return self.dist.expect(func, a, loc, scale, lb, ub, + conditional, **kwds) + + def support(self): + return self.dist.support(*self.args, **self.kwds) + + +class rv_discrete_frozen(rv_frozen): + + def pmf(self, k): + return self.dist.pmf(k, *self.args, **self.kwds) + + def logpmf(self, k): # No error + return self.dist.logpmf(k, *self.args, **self.kwds) + + +class rv_continuous_frozen(rv_frozen): + + def pdf(self, x): + return self.dist.pdf(x, *self.args, **self.kwds) + + def logpdf(self, x): + return self.dist.logpdf(x, *self.args, **self.kwds) + + +def argsreduce(cond, *args): + """Clean arguments to: + + 1. Ensure all arguments are iterable (arrays of dimension at least one + 2. If cond != True and size > 1, ravel(args[i]) where ravel(condition) is + True, in 1D. + + Return list of processed arguments. + + Examples + -------- + >>> import numpy as np + >>> from scipy.stats._distn_infrastructure import argsreduce + >>> rng = np.random.default_rng() + >>> A = rng.random((4, 5)) + >>> B = 2 + >>> C = rng.random((1, 5)) + >>> cond = np.ones(A.shape) + >>> [A1, B1, C1] = argsreduce(cond, A, B, C) + >>> A1.shape + (4, 5) + >>> B1.shape + (1,) + >>> C1.shape + (1, 5) + >>> cond[2,:] = 0 + >>> [A1, B1, C1] = argsreduce(cond, A, B, C) + >>> A1.shape + (15,) + >>> B1.shape + (1,) + >>> C1.shape + (15,) + + """ + # some distributions assume arguments are iterable. + newargs = np.atleast_1d(*args) + + # np.atleast_1d returns an array if only one argument, or a list of arrays + # if more than one argument. + if not isinstance(newargs, (list | tuple)): + newargs = (newargs,) + + if np.all(cond): + # broadcast arrays with cond + *newargs, cond = np.broadcast_arrays(*newargs, cond) + return [arg.ravel() for arg in newargs] + + s = cond.shape + # np.extract returns flattened arrays, which are not broadcastable together + # unless they are either the same size or size == 1. + return [(arg if np.size(arg) == 1 + else np.extract(cond, np.broadcast_to(arg, s))) + for arg in newargs] + + +parse_arg_template = """ +def _parse_args(self, %(shape_arg_str)s %(locscale_in)s): + return (%(shape_arg_str)s), %(locscale_out)s + +def _parse_args_rvs(self, %(shape_arg_str)s %(locscale_in)s, size=None): + return self._argcheck_rvs(%(shape_arg_str)s %(locscale_out)s, size=size) + +def _parse_args_stats(self, %(shape_arg_str)s %(locscale_in)s, moments='mv'): + return (%(shape_arg_str)s), %(locscale_out)s, moments +""" + + +class rv_generic: + """Class which encapsulates common functionality between rv_discrete + and rv_continuous. + + """ + + def __init__(self, seed=None): + super().__init__() + + # figure out if _stats signature has 'moments' keyword + sig = _getfullargspec(self._stats) + self._stats_has_moments = ((sig.varkw is not None) or + ('moments' in sig.args) or + ('moments' in sig.kwonlyargs)) + self._random_state = check_random_state(seed) + + @property + def random_state(self): + """Get or set the generator object for generating random variates. + + If `random_state` is None (or `np.random`), the + `numpy.random.RandomState` singleton is used. + If `random_state` is an int, a new ``RandomState`` instance is used, + seeded with `random_state`. + If `random_state` is already a ``Generator`` or ``RandomState`` + instance, that instance is used. + + """ + return self._random_state + + @random_state.setter + def random_state(self, seed): + self._random_state = check_random_state(seed) + + def __setstate__(self, state): + try: + self.__dict__.update(state) + # attaches the dynamically created methods on each instance. + # if a subclass overrides rv_generic.__setstate__, or implements + # it's own _attach_methods, then it must make sure that + # _attach_argparser_methods is called. + self._attach_methods() + except ValueError: + # reconstitute an old pickle scipy<1.6, that contains + # (_ctor_param, random_state) as state + self._ctor_param = state[0] + self._random_state = state[1] + self.__init__() + + def _attach_methods(self): + """Attaches dynamically created methods to the rv_* instance. + + This method must be overridden by subclasses, and must itself call + _attach_argparser_methods. This method is called in __init__ in + subclasses, and in __setstate__ + """ + raise NotImplementedError + + def _attach_argparser_methods(self): + """ + Generates the argument-parsing functions dynamically and attaches + them to the instance. + + Should be called from `_attach_methods`, typically in __init__ and + during unpickling (__setstate__) + """ + ns = {} + exec(self._parse_arg_template, ns) + # NB: attach to the instance, not class + for name in ['_parse_args', '_parse_args_stats', '_parse_args_rvs']: + setattr(self, name, types.MethodType(ns[name], self)) + + def _construct_argparser( + self, meths_to_inspect, locscale_in, locscale_out): + """Construct the parser string for the shape arguments. + + This method should be called in __init__ of a class for each + distribution. It creates the `_parse_arg_template` attribute that is + then used by `_attach_argparser_methods` to dynamically create and + attach the `_parse_args`, `_parse_args_stats`, `_parse_args_rvs` + methods to the instance. + + If self.shapes is a non-empty string, interprets it as a + comma-separated list of shape parameters. + + Otherwise inspects the call signatures of `meths_to_inspect` + and constructs the argument-parsing functions from these. + In this case also sets `shapes` and `numargs`. + """ + + if self.shapes: + # sanitize the user-supplied shapes + if not isinstance(self.shapes, str): + raise TypeError('shapes must be a string.') + + shapes = self.shapes.replace(',', ' ').split() + + for field in shapes: + if keyword.iskeyword(field): + raise SyntaxError('keywords cannot be used as shapes.') + if not re.match('^[_a-zA-Z][_a-zA-Z0-9]*$', field): + raise SyntaxError( + 'shapes must be valid python identifiers') + else: + # find out the call signatures (_pdf, _cdf etc), deduce shape + # arguments. Generic methods only have 'self, x', any further args + # are shapes. + shapes_list = [] + for meth in meths_to_inspect: + shapes_args = _getfullargspec(meth) # NB does not contain self + args = shapes_args.args[1:] # peel off 'x', too + + if args: + shapes_list.append(args) + + # *args or **kwargs are not allowed w/automatic shapes + if shapes_args.varargs is not None: + raise TypeError( + '*args are not allowed w/out explicit shapes') + if shapes_args.varkw is not None: + raise TypeError( + '**kwds are not allowed w/out explicit shapes') + if shapes_args.kwonlyargs: + raise TypeError( + 'kwonly args are not allowed w/out explicit shapes') + if shapes_args.defaults is not None: + raise TypeError('defaults are not allowed for shapes') + + if shapes_list: + shapes = shapes_list[0] + + # make sure the signatures are consistent + for item in shapes_list: + if item != shapes: + raise TypeError('Shape arguments are inconsistent.') + else: + shapes = [] + + # have the arguments, construct the method from template + shapes_str = ', '.join(shapes) + ', ' if shapes else '' # NB: not None + dct = dict(shape_arg_str=shapes_str, + locscale_in=locscale_in, + locscale_out=locscale_out, + ) + + # this string is used by _attach_argparser_methods + self._parse_arg_template = parse_arg_template % dct + + self.shapes = ', '.join(shapes) if shapes else None + if not hasattr(self, 'numargs'): + # allows more general subclassing with *args + self.numargs = len(shapes) + + def _construct_doc(self, docdict, shapes_vals=None): + """Construct the instance docstring with string substitutions.""" + if sys.flags.optimize > 1: + # if run with -OO, docstrings are stripped + # see https://docs.python.org/3/using/cmdline.html#cmdoption-OO + return + tempdict = docdict.copy() + tempdict['name'] = self.name or 'distname' + tempdict['shapes'] = self.shapes or '' + + if shapes_vals is None: + shapes_vals = () + try: + vals = ', '.join(f'{val:.3g}' for val in shapes_vals) + except TypeError: + vals = ', '.join(f'{val}' for val in shapes_vals) + tempdict['vals'] = vals + + tempdict['shapes_'] = self.shapes or '' + if self.shapes and self.numargs == 1: + tempdict['shapes_'] += ',' + + if self.shapes: + tempdict['set_vals_stmt'] = f'>>> {self.shapes} = {vals}' + else: + tempdict['set_vals_stmt'] = '' + + if self.shapes is None: + # remove shapes from call parameters if there are none + for item in ['default', 'before_notes']: + tempdict[item] = tempdict[item].replace( + "\n%(shapes)s : array_like\n shape parameters", "") + for i in range(2): + if self.shapes is None: + # necessary because we use %(shapes)s in two forms (w w/o ", ") + self.__doc__ = self.__doc__.replace("%(shapes)s, ", "") + try: + self.__doc__ = doccer.docformat(self.__doc__, tempdict) + except TypeError as e: + raise Exception("Unable to construct docstring for " + f"distribution \"{self.name}\": {repr(e)}") from e + + # correct for empty shapes + self.__doc__ = self.__doc__.replace('(, ', '(').replace(', )', ')') + + def _construct_default_doc(self, longname=None, + docdict=None, discrete='continuous'): + """Construct instance docstring from the default template.""" + if sys.flags.optimize > 1: + # if run with -OO, docstrings are stripped + # see https://docs.python.org/3/using/cmdline.html#cmdoption-OO + return + if longname is None: + longname = 'A' + self.__doc__ = ''.join([f'{longname} {discrete} random variable.', + '\n\n%(before_notes)s\n', docheaders['notes'], + '\n%(example)s']) + self._construct_doc(docdict) + + def freeze(self, *args, **kwds): + """Freeze the distribution for the given arguments. + + Parameters + ---------- + arg1, arg2, arg3,... : array_like + The shape parameter(s) for the distribution. Should include all + the non-optional arguments, may include ``loc`` and ``scale``. + + Returns + ------- + rv_frozen : rv_frozen instance + The frozen distribution. + + """ + if isinstance(self, rv_continuous): + return rv_continuous_frozen(self, *args, **kwds) + else: + return rv_discrete_frozen(self, *args, **kwds) + + def __call__(self, *args, **kwds): + return self.freeze(*args, **kwds) + __call__.__doc__ = freeze.__doc__ + + # The actual calculation functions (no basic checking need be done) + # If these are defined, the others won't be looked at. + # Otherwise, the other set can be defined. + def _stats(self, *args, **kwds): + return None, None, None, None + + # Noncentral moments (also known as the moment about the origin). + # Expressed in LaTeX, munp would be $\mu'_{n}$, i.e. "mu-sub-n-prime". + # The primed mu is a widely used notation for the noncentral moment. + def _munp(self, n, *args): + # Silence floating point warnings from integration. + with np.errstate(all='ignore'): + vals = self.generic_moment(n, *args) + return vals + + def _argcheck_rvs(self, *args, **kwargs): + # Handle broadcasting and size validation of the rvs method. + # Subclasses should not have to override this method. + # The rule is that if `size` is not None, then `size` gives the + # shape of the result (integer values of `size` are treated as + # tuples with length 1; i.e. `size=3` is the same as `size=(3,)`.) + # + # `args` is expected to contain the shape parameters (if any), the + # location and the scale in a flat tuple (e.g. if there are two + # shape parameters `a` and `b`, `args` will be `(a, b, loc, scale)`). + # The only keyword argument expected is 'size'. + size = kwargs.get('size', None) + all_bcast = np.broadcast_arrays(*args) + + def squeeze_left(a): + while a.ndim > 0 and a.shape[0] == 1: + a = a[0] + return a + + # Eliminate trivial leading dimensions. In the convention + # used by numpy's random variate generators, trivial leading + # dimensions are effectively ignored. In other words, when `size` + # is given, trivial leading dimensions of the broadcast parameters + # in excess of the number of dimensions in size are ignored, e.g. + # >>> np.random.normal([[1, 3, 5]], [[[[0.01]]]], size=3) + # array([ 1.00104267, 3.00422496, 4.99799278]) + # If `size` is not given, the exact broadcast shape is preserved: + # >>> np.random.normal([[1, 3, 5]], [[[[0.01]]]]) + # array([[[[ 1.00862899, 3.00061431, 4.99867122]]]]) + # + all_bcast = [squeeze_left(a) for a in all_bcast] + bcast_shape = all_bcast[0].shape + bcast_ndim = all_bcast[0].ndim + + if size is None: + size_ = bcast_shape + else: + size_ = tuple(np.atleast_1d(size)) + + # Check compatibility of size_ with the broadcast shape of all + # the parameters. This check is intended to be consistent with + # how the numpy random variate generators (e.g. np.random.normal, + # np.random.beta) handle their arguments. The rule is that, if size + # is given, it determines the shape of the output. Broadcasting + # can't change the output size. + + # This is the standard broadcasting convention of extending the + # shape with fewer dimensions with enough dimensions of length 1 + # so that the two shapes have the same number of dimensions. + ndiff = bcast_ndim - len(size_) + if ndiff < 0: + bcast_shape = (1,)*(-ndiff) + bcast_shape + elif ndiff > 0: + size_ = (1,)*ndiff + size_ + + # This compatibility test is not standard. In "regular" broadcasting, + # two shapes are compatible if for each dimension, the lengths are the + # same or one of the lengths is 1. Here, the length of a dimension in + # size_ must not be less than the corresponding length in bcast_shape. + ok = all([bcdim == 1 or bcdim == szdim + for (bcdim, szdim) in zip(bcast_shape, size_)]) + if not ok: + raise ValueError("size does not match the broadcast shape of " + f"the parameters. {size}, {size_}, {bcast_shape}") + + param_bcast = all_bcast[:-2] + loc_bcast = all_bcast[-2] + scale_bcast = all_bcast[-1] + + return param_bcast, loc_bcast, scale_bcast, size_ + + # These are the methods you must define (standard form functions) + # NB: generic _pdf, _logpdf, _cdf are different for + # rv_continuous and rv_discrete hence are defined in there + def _argcheck(self, *args): + """Default check for correct values on args and keywords. + + Returns condition array of 1's where arguments are correct and + 0's where they are not. + + """ + cond = 1 + for arg in args: + cond = logical_and(cond, (asarray(arg) > 0)) + return cond + + def _get_support(self, *args, **kwargs): + """Return the support of the (unscaled, unshifted) distribution. + + *Must* be overridden by distributions which have support dependent + upon the shape parameters of the distribution. Any such override + *must not* set or change any of the class members, as these members + are shared amongst all instances of the distribution. + + Parameters + ---------- + arg1, arg2, ... : array_like + The shape parameter(s) for the distribution (see docstring of the + instance object for more information). + + Returns + ------- + a, b : numeric (float, or int or +/-np.inf) + end-points of the distribution's support for the specified + shape parameters. + """ + return self.a, self.b + + def _support_mask(self, x, *args): + a, b = self._get_support(*args) + with np.errstate(invalid='ignore'): + return (a <= x) & (x <= b) + + def _open_support_mask(self, x, *args): + a, b = self._get_support(*args) + with np.errstate(invalid='ignore'): + return (a < x) & (x < b) + + def _rvs(self, *args, size=None, random_state=None): + # This method must handle size being a tuple, and it must + # properly broadcast *args and size. size might be + # an empty tuple, which means a scalar random variate is to be + # generated. + + # Use basic inverse cdf algorithm for RV generation as default. + U = random_state.uniform(size=size) + Y = self._ppf(U, *args) + return Y + + def _logcdf(self, x, *args): + with np.errstate(divide='ignore'): + return log(self._cdf(x, *args)) + + def _sf(self, x, *args): + return 1.0-self._cdf(x, *args) + + def _logsf(self, x, *args): + with np.errstate(divide='ignore'): + return log(self._sf(x, *args)) + + def _ppf(self, q, *args): + return self._ppfvec(q, *args) + + def _isf(self, q, *args): + return self._ppf(1.0-q, *args) # use correct _ppf for subclasses + + # These are actually called, and should not be overwritten if you + # want to keep error checking. + def rvs(self, *args, **kwds): + """Random variates of given type. + + Parameters + ---------- + arg1, arg2, arg3,... : array_like + The shape parameter(s) for the distribution (see docstring of the + instance object for more information). + loc : array_like, optional + Location parameter (default=0). + scale : array_like, optional + Scale parameter (default=1). + size : int or tuple of ints, optional + Defining number of random variates (default is 1). + random_state : {None, int, `numpy.random.Generator`, + `numpy.random.RandomState`}, optional + + If `random_state` is None (or `np.random`), the + `numpy.random.RandomState` singleton is used. + If `random_state` is an int, a new ``RandomState`` instance is + used, seeded with `random_state`. + If `random_state` is already a ``Generator`` or ``RandomState`` + instance, that instance is used. + + Returns + ------- + rvs : ndarray or scalar + Random variates of given `size`. + + """ + discrete = kwds.pop('discrete', None) + rndm = kwds.pop('random_state', None) + args, loc, scale, size = self._parse_args_rvs(*args, **kwds) + cond = logical_and(self._argcheck(*args), (scale >= 0)) + if not np.all(cond): + message = ("Domain error in arguments. The `scale` parameter must " + "be positive for all distributions, and many " + "distributions have restrictions on shape parameters. " + f"Please see the `scipy.stats.{self.name}` " + "documentation for details.") + raise ValueError(message) + + if np.all(scale == 0): + return loc*ones(size, 'd') + + # extra gymnastics needed for a custom random_state + if rndm is not None: + random_state_saved = self._random_state + random_state = check_random_state(rndm) + else: + random_state = self._random_state + + vals = self._rvs(*args, size=size, random_state=random_state) + + vals = vals * scale + loc + + # do not forget to restore the _random_state + if rndm is not None: + self._random_state = random_state_saved + + # Cast to int if discrete + if discrete and not isinstance(self, rv_sample): + if size == (): + vals = int(vals) + else: + vals = vals.astype(np.int64) + + return vals + + def stats(self, *args, **kwds): + """Some statistics of the given RV. + + Parameters + ---------- + arg1, arg2, arg3,... : array_like + The shape parameter(s) for the distribution (see docstring of the + instance object for more information) + loc : array_like, optional + location parameter (default=0) + scale : array_like, optional + scale parameter (default=1) + moments : str, optional + composed of letters ['mvsk'] defining which moments to compute: + 'm' = mean, + 'v' = variance, + 's' = (Fisher's) skew, + 'k' = (Fisher's) kurtosis. + (default is 'mv') + + Returns + ------- + stats : sequence + of requested moments. + + """ + args, loc, scale, moments = self._parse_args_stats(*args, **kwds) + # scale = 1 by construction for discrete RVs + loc, scale = map(asarray, (loc, scale)) + args = tuple(map(asarray, args)) + cond = self._argcheck(*args) & (scale > 0) & (loc == loc) + output = [] + default = np.full(shape(cond), fill_value=self.badvalue) + + # Use only entries that are valid in calculation + if np.any(cond): + goodargs = argsreduce(cond, *(args+(scale, loc))) + scale, loc, goodargs = goodargs[-2], goodargs[-1], goodargs[:-2] + + if self._stats_has_moments: + mu, mu2, g1, g2 = self._stats(*goodargs, + **{'moments': moments}) + else: + mu, mu2, g1, g2 = self._stats(*goodargs) + + if 'm' in moments: + if mu is None: + mu = self._munp(1, *goodargs) + out0 = default.copy() + place(out0, cond, mu * scale + loc) + output.append(out0) + + if 'v' in moments: + if mu2 is None: + mu2p = self._munp(2, *goodargs) + if mu is None: + mu = self._munp(1, *goodargs) + # if mean is inf then var is also inf + with np.errstate(invalid='ignore'): + mu2 = np.where(~np.isinf(mu), mu2p - mu**2, np.inf) + out0 = default.copy() + place(out0, cond, mu2 * scale * scale) + output.append(out0) + + if 's' in moments: + if g1 is None: + mu3p = self._munp(3, *goodargs) + if mu is None: + mu = self._munp(1, *goodargs) + if mu2 is None: + mu2p = self._munp(2, *goodargs) + with np.errstate(invalid='ignore'): + mu2 = mu2p - mu * mu + with np.errstate(invalid='ignore'): + mu3 = (-mu*mu - 3*mu2)*mu + mu3p + g1 = mu3 / np.power(mu2, 1.5) + out0 = default.copy() + place(out0, cond, g1) + output.append(out0) + + if 'k' in moments: + if g2 is None: + mu4p = self._munp(4, *goodargs) + if mu is None: + mu = self._munp(1, *goodargs) + if mu2 is None: + mu2p = self._munp(2, *goodargs) + with np.errstate(invalid='ignore'): + mu2 = mu2p - mu * mu + if g1 is None: + mu3 = None + else: + # (mu2**1.5) breaks down for nan and inf + mu3 = g1 * np.power(mu2, 1.5) + if mu3 is None: + mu3p = self._munp(3, *goodargs) + with np.errstate(invalid='ignore'): + mu3 = (-mu * mu - 3 * mu2) * mu + mu3p + with np.errstate(invalid='ignore'): + mu4 = ((-mu**2 - 6*mu2) * mu - 4*mu3)*mu + mu4p + g2 = mu4 / mu2**2.0 - 3.0 + out0 = default.copy() + place(out0, cond, g2) + output.append(out0) + else: # no valid args + output = [default.copy() for _ in moments] + + output = [out[()] for out in output] + if len(output) == 1: + return output[0] + else: + return tuple(output) + + def entropy(self, *args, **kwds): + """Differential entropy of the RV. + + Parameters + ---------- + arg1, arg2, arg3,... : array_like + The shape parameter(s) for the distribution (see docstring of the + instance object for more information). + loc : array_like, optional + Location parameter (default=0). + scale : array_like, optional + Scale parameter (default=1). + + Notes + ----- + Entropy is defined base `e`: + + >>> import numpy as np + >>> from scipy.stats._distn_infrastructure import rv_discrete + >>> drv = rv_discrete(values=((0, 1), (0.5, 0.5))) + >>> np.allclose(drv.entropy(), np.log(2.0)) + True + + """ + args, loc, scale = self._parse_args(*args, **kwds) + # NB: for discrete distributions scale=1 by construction in _parse_args + loc, scale = map(asarray, (loc, scale)) + args = tuple(map(asarray, args)) + cond0 = self._argcheck(*args) & (scale > 0) & (loc == loc) + output = zeros(shape(cond0), 'd') + place(output, (1-cond0), self.badvalue) + goodargs = argsreduce(cond0, scale, *args) + goodscale = goodargs[0] + goodargs = goodargs[1:] + place(output, cond0, self.vecentropy(*goodargs) + log(goodscale)) + return output[()] + + def moment(self, order, *args, **kwds): + """non-central moment of distribution of specified order. + + Parameters + ---------- + order : int, order >= 1 + Order of moment. + arg1, arg2, arg3,... : float + The shape parameter(s) for the distribution (see docstring of the + instance object for more information). + loc : array_like, optional + location parameter (default=0) + scale : array_like, optional + scale parameter (default=1) + + """ + n = order + shapes, loc, scale = self._parse_args(*args, **kwds) + args = np.broadcast_arrays(*(*shapes, loc, scale)) + *shapes, loc, scale = args + + i0 = np.logical_and(self._argcheck(*shapes), scale > 0) + i1 = np.logical_and(i0, loc == 0) + i2 = np.logical_and(i0, loc != 0) + + args = argsreduce(i0, *shapes, loc, scale) + *shapes, loc, scale = args + + if (floor(n) != n): + raise ValueError("Moment must be an integer.") + if (n < 0): + raise ValueError("Moment must be positive.") + mu, mu2, g1, g2 = None, None, None, None + if (n > 0) and (n < 5): + if self._stats_has_moments: + mdict = {'moments': {1: 'm', 2: 'v', 3: 'vs', 4: 'mvsk'}[n]} + else: + mdict = {} + mu, mu2, g1, g2 = self._stats(*shapes, **mdict) + val = np.empty(loc.shape) # val needs to be indexed by loc + val[...] = _moment_from_stats(n, mu, mu2, g1, g2, self._munp, shapes) + + # Convert to transformed X = L + S*Y + # E[X^n] = E[(L+S*Y)^n] = L^n sum(comb(n, k)*(S/L)^k E[Y^k], k=0...n) + result = zeros(i0.shape) + place(result, ~i0, self.badvalue) + + if i1.any(): + res1 = scale[loc == 0]**n * val[loc == 0] + place(result, i1, res1) + + if i2.any(): + mom = [mu, mu2, g1, g2] + arrs = [i for i in mom if i is not None] + idx = [i for i in range(4) if mom[i] is not None] + if any(idx): + arrs = argsreduce(loc != 0, *arrs) + j = 0 + for i in idx: + mom[i] = arrs[j] + j += 1 + mu, mu2, g1, g2 = mom + args = argsreduce(loc != 0, *shapes, loc, scale, val) + *shapes, loc, scale, val = args + + res2 = zeros(loc.shape, dtype='d') + fac = scale / loc + for k in range(n): + valk = _moment_from_stats(k, mu, mu2, g1, g2, self._munp, + shapes) + res2 += comb(n, k, exact=True)*fac**k * valk + res2 += fac**n * val + res2 *= loc**n + place(result, i2, res2) + + return result[()] + + def median(self, *args, **kwds): + """Median of the distribution. + + Parameters + ---------- + arg1, arg2, arg3,... : array_like + The shape parameter(s) for the distribution (see docstring of the + instance object for more information) + loc : array_like, optional + Location parameter, Default is 0. + scale : array_like, optional + Scale parameter, Default is 1. + + Returns + ------- + median : float + The median of the distribution. + + See Also + -------- + rv_discrete.ppf + Inverse of the CDF + + """ + return self.ppf(0.5, *args, **kwds) + + def mean(self, *args, **kwds): + """Mean of the distribution. + + Parameters + ---------- + arg1, arg2, arg3,... : array_like + The shape parameter(s) for the distribution (see docstring of the + instance object for more information) + loc : array_like, optional + location parameter (default=0) + scale : array_like, optional + scale parameter (default=1) + + Returns + ------- + mean : float + the mean of the distribution + + """ + kwds['moments'] = 'm' + res = self.stats(*args, **kwds) + if isinstance(res, ndarray) and res.ndim == 0: + return res[()] + return res + + def var(self, *args, **kwds): + """Variance of the distribution. + + Parameters + ---------- + arg1, arg2, arg3,... : array_like + The shape parameter(s) for the distribution (see docstring of the + instance object for more information) + loc : array_like, optional + location parameter (default=0) + scale : array_like, optional + scale parameter (default=1) + + Returns + ------- + var : float + the variance of the distribution + + """ + kwds['moments'] = 'v' + res = self.stats(*args, **kwds) + if isinstance(res, ndarray) and res.ndim == 0: + return res[()] + return res + + def std(self, *args, **kwds): + """Standard deviation of the distribution. + + Parameters + ---------- + arg1, arg2, arg3,... : array_like + The shape parameter(s) for the distribution (see docstring of the + instance object for more information) + loc : array_like, optional + location parameter (default=0) + scale : array_like, optional + scale parameter (default=1) + + Returns + ------- + std : float + standard deviation of the distribution + + """ + kwds['moments'] = 'v' + res = sqrt(self.stats(*args, **kwds)) + return res + + def interval(self, confidence, *args, **kwds): + """Confidence interval with equal areas around the median. + + Parameters + ---------- + confidence : array_like of float + Probability that an rv will be drawn from the returned range. + Each value should be in the range [0, 1]. + arg1, arg2, ... : array_like + The shape parameter(s) for the distribution (see docstring of the + instance object for more information). + loc : array_like, optional + location parameter, Default is 0. + scale : array_like, optional + scale parameter, Default is 1. + + Returns + ------- + a, b : ndarray of float + end-points of range that contain ``100 * alpha %`` of the rv's + possible values. + + Notes + ----- + This is implemented as ``ppf([p_tail, 1-p_tail])``, where + ``ppf`` is the inverse cumulative distribution function and + ``p_tail = (1-confidence)/2``. Suppose ``[c, d]`` is the support of a + discrete distribution; then ``ppf([0, 1]) == (c-1, d)``. Therefore, + when ``confidence=1`` and the distribution is discrete, the left end + of the interval will be beyond the support of the distribution. + For discrete distributions, the interval will limit the probability + in each tail to be less than or equal to ``p_tail`` (usually + strictly less). + + """ + alpha = confidence + + alpha = asarray(alpha) + if np.any((alpha > 1) | (alpha < 0)): + raise ValueError("alpha must be between 0 and 1 inclusive") + q1 = (1.0-alpha)/2 + q2 = (1.0+alpha)/2 + a = self.ppf(q1, *args, **kwds) + b = self.ppf(q2, *args, **kwds) + return a, b + + def support(self, *args, **kwargs): + """Support of the distribution. + + Parameters + ---------- + arg1, arg2, ... : array_like + The shape parameter(s) for the distribution (see docstring of the + instance object for more information). + loc : array_like, optional + location parameter, Default is 0. + scale : array_like, optional + scale parameter, Default is 1. + + Returns + ------- + a, b : array_like + end-points of the distribution's support. + + """ + args, loc, scale = self._parse_args(*args, **kwargs) + arrs = np.broadcast_arrays(*args, loc, scale) + args, loc, scale = arrs[:-2], arrs[-2], arrs[-1] + cond = self._argcheck(*args) & (scale > 0) + _a, _b = self._get_support(*args) + if cond.all(): + return _a * scale + loc, _b * scale + loc + elif cond.ndim == 0: + return self.badvalue, self.badvalue + # promote bounds to at least float to fill in the badvalue + _a, _b = np.asarray(_a).astype('d'), np.asarray(_b).astype('d') + out_a, out_b = _a * scale + loc, _b * scale + loc + place(out_a, 1-cond, self.badvalue) + place(out_b, 1-cond, self.badvalue) + return out_a, out_b + + def nnlf(self, theta, x): + """Negative loglikelihood function. + Notes + ----- + This is ``-sum(log pdf(x, theta), axis=0)`` where `theta` are the + parameters (including loc and scale). + """ + loc, scale, args = self._unpack_loc_scale(theta) + if not self._argcheck(*args) or scale <= 0: + return inf + x = (asarray(x)-loc) / scale + n_log_scale = len(x) * log(scale) + if np.any(~self._support_mask(x, *args)): + return inf + return self._nnlf(x, *args) + n_log_scale + + def _nnlf(self, x, *args): + return -np.sum(self._logpxf(x, *args), axis=0) + + def _nlff_and_penalty(self, x, args, log_fitfun): + # negative log fit function + cond0 = ~self._support_mask(x, *args) + n_bad = np.count_nonzero(cond0, axis=0) + if n_bad > 0: + x = argsreduce(~cond0, x)[0] + logff = log_fitfun(x, *args) + finite_logff = np.isfinite(logff) + n_bad += np.sum(~finite_logff, axis=0) + if n_bad > 0: + penalty = n_bad * log(_XMAX) * 100 + return -np.sum(logff[finite_logff], axis=0) + penalty + return -np.sum(logff, axis=0) + + def _penalized_nnlf(self, theta, x): + """Penalized negative loglikelihood function. + i.e., - sum (log pdf(x, theta), axis=0) + penalty + where theta are the parameters (including loc and scale) + """ + loc, scale, args = self._unpack_loc_scale(theta) + if not self._argcheck(*args) or scale <= 0: + return inf + x = asarray((x-loc) / scale) + n_log_scale = len(x) * log(scale) + return self._nlff_and_penalty(x, args, self._logpxf) + n_log_scale + + def _penalized_nlpsf(self, theta, x): + """Penalized negative log product spacing function. + i.e., - sum (log (diff (cdf (x, theta))), axis=0) + penalty + where theta are the parameters (including loc and scale) + Follows reference [1] of scipy.stats.fit + """ + loc, scale, args = self._unpack_loc_scale(theta) + if not self._argcheck(*args) or scale <= 0: + return inf + x = (np.sort(x) - loc)/scale + + def log_psf(x, *args): + x, lj = np.unique(x, return_counts=True) # fast for sorted x + cdf_data = self._cdf(x, *args) if x.size else [] + if not (x.size and 1 - cdf_data[-1] <= 0): + cdf = np.concatenate(([0], cdf_data, [1])) + lj = np.concatenate((lj, [1])) + else: + cdf = np.concatenate(([0], cdf_data)) + # here we could use logcdf w/ logsumexp trick to take differences, + # but in the context of the method, it seems unlikely to matter + return lj * np.log(np.diff(cdf) / lj) + + return self._nlff_and_penalty(x, args, log_psf) + + +class _ShapeInfo: + def __init__(self, name, integrality=False, domain=(-np.inf, np.inf), + inclusive=(True, True)): + self.name = name + self.integrality = integrality + self.endpoints = domain + self.inclusive = inclusive + + domain = list(domain) + if np.isfinite(domain[0]) and not inclusive[0]: + domain[0] = np.nextafter(domain[0], np.inf) + if np.isfinite(domain[1]) and not inclusive[1]: + domain[1] = np.nextafter(domain[1], -np.inf) + self.domain = domain + + +def _get_fixed_fit_value(kwds, names): + """ + Given names such as ``['f0', 'fa', 'fix_a']``, check that there is + at most one non-None value in `kwds` associated with those names. + Return that value, or None if none of the names occur in `kwds`. + As a side effect, all occurrences of those names in `kwds` are + removed. + """ + vals = [(name, kwds.pop(name)) for name in names if name in kwds] + if len(vals) > 1: + repeated = [name for name, val in vals] + raise ValueError("fit method got multiple keyword arguments to " + "specify the same fixed parameter: " + + ', '.join(repeated)) + return vals[0][1] if vals else None + + +# continuous random variables: implement maybe later +# +# hf --- Hazard Function (PDF / SF) +# chf --- Cumulative hazard function (-log(SF)) +# psf --- Probability sparsity function (reciprocal of the pdf) in +# units of percent-point-function (as a function of q). +# Also, the derivative of the percent-point function. + + +class rv_continuous(rv_generic): + """A generic continuous random variable class meant for subclassing. + + `rv_continuous` is a base class to construct specific distribution classes + and instances for continuous random variables. It cannot be used + directly as a distribution. + + Parameters + ---------- + momtype : int, optional + The type of generic moment calculation to use: 0 for pdf, 1 (default) + for ppf. + a : float, optional + Lower bound of the support of the distribution, default is minus + infinity. + b : float, optional + Upper bound of the support of the distribution, default is plus + infinity. + xtol : float, optional + The tolerance for fixed point calculation for generic ppf. + badvalue : float, optional + The value in a result arrays that indicates a value that for which + some argument restriction is violated, default is np.nan. + name : str, optional + The name of the instance. This string is used to construct the default + example for distributions. + longname : str, optional + This string is used as part of the first line of the docstring returned + when a subclass has no docstring of its own. Note: `longname` exists + for backwards compatibility, do not use for new subclasses. + shapes : str, optional + The shape of the distribution. For example ``"m, n"`` for a + distribution that takes two integers as the two shape arguments for all + its methods. If not provided, shape parameters will be inferred from + the signature of the private methods, ``_pdf`` and ``_cdf`` of the + instance. + seed : {None, int, `numpy.random.Generator`, `numpy.random.RandomState`}, optional + If `seed` is None (or `np.random`), the `numpy.random.RandomState` + singleton is used. + If `seed` is an int, a new ``RandomState`` instance is used, + seeded with `seed`. + If `seed` is already a ``Generator`` or ``RandomState`` instance then + that instance is used. + + Attributes + ---------- + a, b : float, optional + Lower/upper bound of the support of the unshifted/unscaled distribution. + This value is unaffected by the `loc` and `scale` parameters. + To calculate the support of the shifted/scaled distribution, + use the `support` method. + + Methods + ------- + rvs + pdf + logpdf + cdf + logcdf + sf + logsf + ppf + isf + moment + stats + entropy + expect + median + mean + std + var + interval + __call__ + fit + fit_loc_scale + nnlf + support + + Notes + ----- + Public methods of an instance of a distribution class (e.g., ``pdf``, + ``cdf``) check their arguments and pass valid arguments to private, + computational methods (``_pdf``, ``_cdf``). For ``pdf(x)``, ``x`` is valid + if it is within the support of the distribution. + Whether a shape parameter is valid is decided by an ``_argcheck`` method + (which defaults to checking that its arguments are strictly positive.) + + **Subclassing** + + New random variables can be defined by subclassing the `rv_continuous` class + and re-defining at least the ``_pdf`` or the ``_cdf`` method (normalized + to location 0 and scale 1). + + If positive argument checking is not correct for your RV + then you will also need to re-define the ``_argcheck`` method. + + For most of the scipy.stats distributions, the support interval doesn't + depend on the shape parameters. ``x`` being in the support interval is + equivalent to ``self.a <= x <= self.b``. If either of the endpoints of + the support do depend on the shape parameters, then + i) the distribution must implement the ``_get_support`` method; and + ii) those dependent endpoints must be omitted from the distribution's + call to the ``rv_continuous`` initializer. + + Correct, but potentially slow defaults exist for the remaining + methods but for speed and/or accuracy you can over-ride:: + + _logpdf, _cdf, _logcdf, _ppf, _rvs, _isf, _sf, _logsf + + The default method ``_rvs`` relies on the inverse of the cdf, ``_ppf``, + applied to a uniform random variate. In order to generate random variates + efficiently, either the default ``_ppf`` needs to be overwritten (e.g. + if the inverse cdf can expressed in an explicit form) or a sampling + method needs to be implemented in a custom ``_rvs`` method. + + If possible, you should override ``_isf``, ``_sf`` or ``_logsf``. + The main reason would be to improve numerical accuracy: for example, + the survival function ``_sf`` is computed as ``1 - _cdf`` which can + result in loss of precision if ``_cdf(x)`` is close to one. + + **Methods that can be overwritten by subclasses** + :: + + _rvs + _pdf + _cdf + _sf + _ppf + _isf + _stats + _munp + _entropy + _argcheck + _get_support + + There are additional (internal and private) generic methods that can + be useful for cross-checking and for debugging, but might work in all + cases when directly called. + + A note on ``shapes``: subclasses need not specify them explicitly. In this + case, `shapes` will be automatically deduced from the signatures of the + overridden methods (`pdf`, `cdf` etc). + If, for some reason, you prefer to avoid relying on introspection, you can + specify ``shapes`` explicitly as an argument to the instance constructor. + + + **Frozen Distributions** + + Normally, you must provide shape parameters (and, optionally, location and + scale parameters to each call of a method of a distribution. + + Alternatively, the object may be called (as a function) to fix the shape, + location, and scale parameters returning a "frozen" continuous RV object: + + rv = generic(, loc=0, scale=1) + `rv_frozen` object with the same methods but holding the given shape, + location, and scale fixed + + **Statistics** + + Statistics are computed using numerical integration by default. + For speed you can redefine this using ``_stats``: + + - take shape parameters and return mu, mu2, g1, g2 + - If you can't compute one of these, return it as None + - Can also be defined with a keyword argument ``moments``, which is a + string composed of "m", "v", "s", and/or "k". + Only the components appearing in string should be computed and + returned in the order "m", "v", "s", or "k" with missing values + returned as None. + + Alternatively, you can override ``_munp``, which takes ``n`` and shape + parameters and returns the n-th non-central moment of the distribution. + + **Deepcopying / Pickling** + + If a distribution or frozen distribution is deepcopied (pickled/unpickled, + etc.), any underlying random number generator is deepcopied with it. An + implication is that if a distribution relies on the singleton RandomState + before copying, it will rely on a copy of that random state after copying, + and ``np.random.seed`` will no longer control the state. + + Examples + -------- + To create a new Gaussian distribution, we would do the following: + + >>> from scipy.stats import rv_continuous + >>> class gaussian_gen(rv_continuous): + ... "Gaussian distribution" + ... def _pdf(self, x): + ... return np.exp(-x**2 / 2.) / np.sqrt(2.0 * np.pi) + >>> gaussian = gaussian_gen(name='gaussian') + + ``scipy.stats`` distributions are *instances*, so here we subclass + `rv_continuous` and create an instance. With this, we now have + a fully functional distribution with all relevant methods automagically + generated by the framework. + + Note that above we defined a standard normal distribution, with zero mean + and unit variance. Shifting and scaling of the distribution can be done + by using ``loc`` and ``scale`` parameters: ``gaussian.pdf(x, loc, scale)`` + essentially computes ``y = (x - loc) / scale`` and + ``gaussian._pdf(y) / scale``. + + """ + + _support: tuple[float | str, float | str] # not required + + def __init__(self, momtype=1, a=None, b=None, xtol=1e-14, + badvalue=None, name=None, longname=None, + shapes=None, seed=None): + + super().__init__(seed) + + # save the ctor parameters, cf generic freeze + self._ctor_param = dict( + momtype=momtype, a=a, b=b, xtol=xtol, + badvalue=badvalue, name=name, longname=longname, + shapes=shapes, seed=seed) + + if badvalue is None: + badvalue = nan + if name is None: + name = 'Distribution' + self.badvalue = badvalue + self.name = name + self.a = a + self.b = b + if a is None: + self.a = -inf + if b is None: + self.b = inf + self.xtol = xtol + self.moment_type = momtype + self.shapes = shapes + + self._construct_argparser(meths_to_inspect=[self._pdf, self._cdf], + locscale_in='loc=0, scale=1', + locscale_out='loc, scale') + self._attach_methods() + + if longname is None: + if name[0] in ['aeiouAEIOU']: + hstr = "An " + else: + hstr = "A " + longname = hstr + name + + if sys.flags.optimize < 2: + # Skip adding docstrings if interpreter is run with -OO + if self.__doc__ is None: + self._construct_default_doc(longname=longname, + docdict=docdict, + discrete='continuous') + else: + dct = dict(distcont) + self._construct_doc(docdict, dct.get(self.name)) + + def __getstate__(self): + dct = self.__dict__.copy() + + # these methods will be remade in __setstate__ + # _random_state attribute is taken care of by rv_generic + attrs = ["_parse_args", "_parse_args_stats", "_parse_args_rvs", + "_cdfvec", "_ppfvec", "vecentropy", "generic_moment"] + [dct.pop(attr, None) for attr in attrs] + return dct + + def _attach_methods(self): + """ + Attaches dynamically created methods to the rv_continuous instance. + """ + # _attach_methods is responsible for calling _attach_argparser_methods + self._attach_argparser_methods() + + # nin correction + self._ppfvec = vectorize(self._ppf_single, otypes='d') + self._ppfvec.nin = self.numargs + 1 + self.vecentropy = vectorize(self._entropy, otypes='d') + self._cdfvec = vectorize(self._cdf_single, otypes='d') + self._cdfvec.nin = self.numargs + 1 + + if self.moment_type == 0: + self.generic_moment = vectorize(self._mom0_sc, otypes='d') + else: + self.generic_moment = vectorize(self._mom1_sc, otypes='d') + # Because of the *args argument of _mom0_sc, vectorize cannot count the + # number of arguments correctly. + self.generic_moment.nin = self.numargs + 1 + + def _updated_ctor_param(self): + """Return the current version of _ctor_param, possibly updated by user. + + Used by freezing. + Keep this in sync with the signature of __init__. + """ + dct = self._ctor_param.copy() + dct['a'] = self.a + dct['b'] = self.b + dct['xtol'] = self.xtol + dct['badvalue'] = self.badvalue + dct['name'] = self.name + dct['shapes'] = self.shapes + return dct + + def _ppf_to_solve(self, x, q, *args): + return self.cdf(*(x, )+args)-q + + def _ppf_single(self, q, *args): + factor = 10. + left, right = self._get_support(*args) + + if np.isinf(left): + left = min(-factor, right) + while self._ppf_to_solve(left, q, *args) > 0.: + left, right = left * factor, left + # left is now such that cdf(left) <= q + # if right has changed, then cdf(right) > q + + if np.isinf(right): + right = max(factor, left) + while self._ppf_to_solve(right, q, *args) < 0.: + left, right = right, right * factor + # right is now such that cdf(right) >= q + + return optimize.brentq(self._ppf_to_solve, + left, right, args=(q,)+args, xtol=self.xtol) + + # moment from definition + def _mom_integ0(self, x, m, *args): + return x**m * self.pdf(x, *args) + + def _mom0_sc(self, m, *args): + _a, _b = self._get_support(*args) + return integrate.quad(self._mom_integ0, _a, _b, + args=(m,)+args)[0] + + # moment calculated using ppf + def _mom_integ1(self, q, m, *args): + return (self.ppf(q, *args))**m + + def _mom1_sc(self, m, *args): + return integrate.quad(self._mom_integ1, 0, 1, args=(m,)+args)[0] + + def _pdf(self, x, *args): + return _derivative(self._cdf, x, dx=1e-5, args=args, order=5) + + # Could also define any of these + def _logpdf(self, x, *args): + p = self._pdf(x, *args) + with np.errstate(divide='ignore'): + return log(p) + + def _logpxf(self, x, *args): + # continuous distributions have PDF, discrete have PMF, but sometimes + # the distinction doesn't matter. This lets us use `_logpxf` for both + # discrete and continuous distributions. + return self._logpdf(x, *args) + + def _cdf_single(self, x, *args): + _a, _b = self._get_support(*args) + return integrate.quad(self._pdf, _a, x, args=args)[0] + + def _cdf(self, x, *args): + return self._cdfvec(x, *args) + + def _logcdf(self, x, *args): + median = self._ppf(0.5, *args) + with np.errstate(divide='ignore'): + return xpx.apply_where( + x < median, (x,) + args, + lambda x, *args: np.log(self._cdf(x, *args)), + lambda x, *args: np.log1p(-self._sf(x, *args))) + + def _logsf(self, x, *args): + median = self._ppf(0.5, *args) + with np.errstate(divide='ignore'): + return xpx.apply_where( + x > median, (x,) + args, + lambda x, *args: np.log(self._sf(x, *args)), + lambda x, *args: np.log1p(-self._cdf(x, *args))) + + # generic _argcheck, _sf, _ppf, _isf, _rvs are defined + # in rv_generic + + def pdf(self, x, *args, **kwds): + """Probability density function at x of the given RV. + + Parameters + ---------- + x : array_like + quantiles + arg1, arg2, arg3,... : array_like + The shape parameter(s) for the distribution (see docstring of the + instance object for more information) + loc : array_like, optional + location parameter (default=0) + scale : array_like, optional + scale parameter (default=1) + + Returns + ------- + pdf : ndarray + Probability density function evaluated at x + + """ + args, loc, scale = self._parse_args(*args, **kwds) + x, loc, scale = map(asarray, (x, loc, scale)) + args = tuple(map(asarray, args)) + dtyp = np.promote_types(x.dtype, np.float64) + x = np.asarray((x - loc)/scale, dtype=dtyp) + cond0 = self._argcheck(*args) & (scale > 0) + cond1 = self._support_mask(x, *args) & (scale > 0) + cond = cond0 & cond1 + output = zeros(shape(cond), dtyp) + putmask(output, (1-cond0)+np.isnan(x), self.badvalue) + if np.any(cond): + goodargs = argsreduce(cond, *((x,)+args+(scale,))) + scale, goodargs = goodargs[-1], goodargs[:-1] + place(output, cond, self._pdf(*goodargs) / scale) + if output.ndim == 0: + return output[()] + return output + + def logpdf(self, x, *args, **kwds): + """Log of the probability density function at x of the given RV. + + This uses a more numerically accurate calculation if available. + + Parameters + ---------- + x : array_like + quantiles + arg1, arg2, arg3,... : array_like + The shape parameter(s) for the distribution (see docstring of the + instance object for more information) + loc : array_like, optional + location parameter (default=0) + scale : array_like, optional + scale parameter (default=1) + + Returns + ------- + logpdf : array_like + Log of the probability density function evaluated at x + + """ + args, loc, scale = self._parse_args(*args, **kwds) + x, loc, scale = map(asarray, (x, loc, scale)) + args = tuple(map(asarray, args)) + dtyp = np.promote_types(x.dtype, np.float64) + x = np.asarray((x - loc)/scale, dtype=dtyp) + cond0 = self._argcheck(*args) & (scale > 0) + cond1 = self._support_mask(x, *args) & (scale > 0) + cond = cond0 & cond1 + output = empty(shape(cond), dtyp) + output.fill(-inf) + putmask(output, (1-cond0)+np.isnan(x), self.badvalue) + if np.any(cond): + goodargs = argsreduce(cond, *((x,)+args+(scale,))) + scale, goodargs = goodargs[-1], goodargs[:-1] + place(output, cond, self._logpdf(*goodargs) - log(scale)) + if output.ndim == 0: + return output[()] + return output + + def cdf(self, x, *args, **kwds): + """ + Cumulative distribution function of the given RV. + + Parameters + ---------- + x : array_like + quantiles + arg1, arg2, arg3,... : array_like + The shape parameter(s) for the distribution (see docstring of the + instance object for more information) + loc : array_like, optional + location parameter (default=0) + scale : array_like, optional + scale parameter (default=1) + + Returns + ------- + cdf : ndarray + Cumulative distribution function evaluated at `x` + + """ + args, loc, scale = self._parse_args(*args, **kwds) + x, loc, scale = map(asarray, (x, loc, scale)) + args = tuple(map(asarray, args)) + _a, _b = self._get_support(*args) + dtyp = np.promote_types(x.dtype, np.float64) + x = np.asarray((x - loc)/scale, dtype=dtyp) + cond0 = self._argcheck(*args) & (scale > 0) + cond1 = self._open_support_mask(x, *args) & (scale > 0) + cond2 = (x >= np.asarray(_b)) & cond0 + cond = cond0 & cond1 + output = zeros(shape(cond), dtyp) + place(output, (1-cond0)+np.isnan(x), self.badvalue) + place(output, cond2, 1.0) + if np.any(cond): # call only if at least 1 entry + goodargs = argsreduce(cond, *((x,)+args)) + place(output, cond, self._cdf(*goodargs)) + if output.ndim == 0: + return output[()] + return output + + def logcdf(self, x, *args, **kwds): + """Log of the cumulative distribution function at x of the given RV. + + Parameters + ---------- + x : array_like + quantiles + arg1, arg2, arg3,... : array_like + The shape parameter(s) for the distribution (see docstring of the + instance object for more information) + loc : array_like, optional + location parameter (default=0) + scale : array_like, optional + scale parameter (default=1) + + Returns + ------- + logcdf : array_like + Log of the cumulative distribution function evaluated at x + + """ + args, loc, scale = self._parse_args(*args, **kwds) + x, loc, scale = map(asarray, (x, loc, scale)) + args = tuple(map(asarray, args)) + _a, _b = self._get_support(*args) + dtyp = np.promote_types(x.dtype, np.float64) + x = np.asarray((x - loc)/scale, dtype=dtyp) + cond0 = self._argcheck(*args) & (scale > 0) + cond1 = self._open_support_mask(x, *args) & (scale > 0) + cond2 = (x >= _b) & cond0 + cond = cond0 & cond1 + output = empty(shape(cond), dtyp) + output.fill(-inf) + place(output, (1-cond0)*(cond1 == cond1)+np.isnan(x), self.badvalue) + place(output, cond2, 0.0) + if np.any(cond): # call only if at least 1 entry + goodargs = argsreduce(cond, *((x,)+args)) + place(output, cond, self._logcdf(*goodargs)) + if output.ndim == 0: + return output[()] + return output + + def sf(self, x, *args, **kwds): + """Survival function (1 - `cdf`) at x of the given RV. + + Parameters + ---------- + x : array_like + quantiles + arg1, arg2, arg3,... : array_like + The shape parameter(s) for the distribution (see docstring of the + instance object for more information) + loc : array_like, optional + location parameter (default=0) + scale : array_like, optional + scale parameter (default=1) + + Returns + ------- + sf : array_like + Survival function evaluated at x + + """ + args, loc, scale = self._parse_args(*args, **kwds) + x, loc, scale = map(asarray, (x, loc, scale)) + args = tuple(map(asarray, args)) + _a, _b = self._get_support(*args) + dtyp = np.promote_types(x.dtype, np.float64) + x = np.asarray((x - loc)/scale, dtype=dtyp) + cond0 = self._argcheck(*args) & (scale > 0) + cond1 = self._open_support_mask(x, *args) & (scale > 0) + cond2 = cond0 & (x <= _a) + cond = cond0 & cond1 + output = zeros(shape(cond), dtyp) + place(output, (1-cond0)+np.isnan(x), self.badvalue) + place(output, cond2, 1.0) + if np.any(cond): + goodargs = argsreduce(cond, *((x,)+args)) + place(output, cond, self._sf(*goodargs)) + if output.ndim == 0: + return output[()] + return output + + def logsf(self, x, *args, **kwds): + """Log of the survival function of the given RV. + + Returns the log of the "survival function," defined as (1 - `cdf`), + evaluated at `x`. + + Parameters + ---------- + x : array_like + quantiles + arg1, arg2, arg3,... : array_like + The shape parameter(s) for the distribution (see docstring of the + instance object for more information) + loc : array_like, optional + location parameter (default=0) + scale : array_like, optional + scale parameter (default=1) + + Returns + ------- + logsf : ndarray + Log of the survival function evaluated at `x`. + + """ + args, loc, scale = self._parse_args(*args, **kwds) + x, loc, scale = map(asarray, (x, loc, scale)) + args = tuple(map(asarray, args)) + _a, _b = self._get_support(*args) + dtyp = np.promote_types(x.dtype, np.float64) + x = np.asarray((x - loc)/scale, dtype=dtyp) + cond0 = self._argcheck(*args) & (scale > 0) + cond1 = self._open_support_mask(x, *args) & (scale > 0) + cond2 = cond0 & (x <= _a) + cond = cond0 & cond1 + output = empty(shape(cond), dtyp) + output.fill(-inf) + place(output, (1-cond0)+np.isnan(x), self.badvalue) + place(output, cond2, 0.0) + if np.any(cond): + goodargs = argsreduce(cond, *((x,)+args)) + place(output, cond, self._logsf(*goodargs)) + if output.ndim == 0: + return output[()] + return output + + def ppf(self, q, *args, **kwds): + """Percent point function (inverse of `cdf`) at q of the given RV. + + Parameters + ---------- + q : array_like + lower tail probability + arg1, arg2, arg3,... : array_like + The shape parameter(s) for the distribution (see docstring of the + instance object for more information) + loc : array_like, optional + location parameter (default=0) + scale : array_like, optional + scale parameter (default=1) + + Returns + ------- + x : array_like + quantile corresponding to the lower tail probability q. + + """ + args, loc, scale = self._parse_args(*args, **kwds) + q, loc, scale = map(asarray, (q, loc, scale)) + args = tuple(map(asarray, args)) + _a, _b = self._get_support(*args) + cond0 = self._argcheck(*args) & (scale > 0) & (loc == loc) + cond1 = (0 < q) & (q < 1) + cond2 = cond0 & (q == 0) + cond3 = cond0 & (q == 1) + cond = cond0 & cond1 + output = np.full(shape(cond), fill_value=self.badvalue) + + lower_bound = _a * scale + loc + upper_bound = _b * scale + loc + place(output, cond2, argsreduce(cond2, lower_bound)[0]) + place(output, cond3, argsreduce(cond3, upper_bound)[0]) + + if np.any(cond): # call only if at least 1 entry + goodargs = argsreduce(cond, *((q,)+args+(scale, loc))) + scale, loc, goodargs = goodargs[-2], goodargs[-1], goodargs[:-2] + place(output, cond, self._ppf(*goodargs) * scale + loc) + if output.ndim == 0: + return output[()] + return output + + def isf(self, q, *args, **kwds): + """Inverse survival function (inverse of `sf`) at q of the given RV. + + Parameters + ---------- + q : array_like + upper tail probability + arg1, arg2, arg3,... : array_like + The shape parameter(s) for the distribution (see docstring of the + instance object for more information) + loc : array_like, optional + location parameter (default=0) + scale : array_like, optional + scale parameter (default=1) + + Returns + ------- + x : ndarray or scalar + Quantile corresponding to the upper tail probability q. + + """ + args, loc, scale = self._parse_args(*args, **kwds) + q, loc, scale = map(asarray, (q, loc, scale)) + args = tuple(map(asarray, args)) + _a, _b = self._get_support(*args) + cond0 = self._argcheck(*args) & (scale > 0) & (loc == loc) + cond1 = (0 < q) & (q < 1) + cond2 = cond0 & (q == 1) + cond3 = cond0 & (q == 0) + cond = cond0 & cond1 + output = np.full(shape(cond), fill_value=self.badvalue) + + lower_bound = _a * scale + loc + upper_bound = _b * scale + loc + place(output, cond2, argsreduce(cond2, lower_bound)[0]) + place(output, cond3, argsreduce(cond3, upper_bound)[0]) + + if np.any(cond): + goodargs = argsreduce(cond, *((q,)+args+(scale, loc))) + scale, loc, goodargs = goodargs[-2], goodargs[-1], goodargs[:-2] + place(output, cond, self._isf(*goodargs) * scale + loc) + if output.ndim == 0: + return output[()] + return output + + def _unpack_loc_scale(self, theta): + try: + loc = theta[-2] + scale = theta[-1] + args = tuple(theta[:-2]) + except IndexError as e: + raise ValueError("Not enough input arguments.") from e + return loc, scale, args + + def _nnlf_and_penalty(self, x, args): + """ + Compute the penalized negative log-likelihood for the + "standardized" data (i.e. already shifted by loc and + scaled by scale) for the shape parameters in `args`. + + `x` can be a 1D numpy array or a CensoredData instance. + """ + if isinstance(x, CensoredData): + # Filter out the data that is not in the support. + xs = x._supported(*self._get_support(*args)) + n_bad = len(x) - len(xs) + i1, i2 = xs._interval.T + terms = [ + # logpdf of the noncensored data. + self._logpdf(xs._uncensored, *args), + # logcdf of the left-censored data. + self._logcdf(xs._left, *args), + # logsf of the right-censored data. + self._logsf(xs._right, *args), + # log of probability of the interval-censored data. + np.log(self._delta_cdf(i1, i2, *args)), + ] + else: + cond0 = ~self._support_mask(x, *args) + n_bad = np.count_nonzero(cond0) + if n_bad > 0: + x = argsreduce(~cond0, x)[0] + terms = [self._logpdf(x, *args)] + + totals, bad_counts = zip(*[_sum_finite(term) for term in terms]) + total = sum(totals) + n_bad += sum(bad_counts) + + return -total + n_bad * _LOGXMAX * 100 + + def _penalized_nnlf(self, theta, x): + """Penalized negative loglikelihood function. + + i.e., - sum (log pdf(x, theta), axis=0) + penalty + where theta are the parameters (including loc and scale) + """ + loc, scale, args = self._unpack_loc_scale(theta) + if not self._argcheck(*args) or scale <= 0: + return inf + if isinstance(x, CensoredData): + x = (x - loc) / scale + n_log_scale = (len(x) - x.num_censored()) * log(scale) + else: + x = (x - loc) / scale + n_log_scale = len(x) * log(scale) + + return self._nnlf_and_penalty(x, args) + n_log_scale + + def _fitstart(self, data, args=None): + """Starting point for fit (shape arguments + loc + scale).""" + if args is None: + args = (1.0,)*self.numargs + loc, scale = self._fit_loc_scale_support(data, *args) + return args + (loc, scale) + + def _reduce_func(self, args, kwds, data=None): + """ + Return the (possibly reduced) function to optimize in order to find MLE + estimates for the .fit method. + """ + # Convert fixed shape parameters to the standard numeric form: e.g. for + # stats.beta, shapes='a, b'. To fix `a`, the caller can give a value + # for `f0`, `fa` or 'fix_a'. The following converts the latter two + # into the first (numeric) form. + shapes = [] + if self.shapes: + shapes = self.shapes.replace(',', ' ').split() + for j, s in enumerate(shapes): + key = 'f' + str(j) + names = [key, 'f' + s, 'fix_' + s] + val = _get_fixed_fit_value(kwds, names) + if val is not None: + kwds[key] = val + + args = list(args) + Nargs = len(args) + fixedn = [] + names = [f'f{n}' for n in range(Nargs - 2)] + ['floc', 'fscale'] + x0 = [] + for n, key in enumerate(names): + if key in kwds: + fixedn.append(n) + args[n] = kwds.pop(key) + else: + x0.append(args[n]) + + methods = {"mle", "mm"} + method = kwds.pop('method', "mle").lower() + if method == "mm": + n_params = len(shapes) + 2 - len(fixedn) + exponents = (np.arange(1, n_params+1))[:, np.newaxis] + data_moments = np.sum(data[None, :]**exponents/len(data), axis=1) + + def objective(theta, x): + return self._moment_error(theta, x, data_moments) + + elif method == "mle": + objective = self._penalized_nnlf + else: + raise ValueError(f"Method '{method}' not available; " + f"must be one of {methods}") + + if len(fixedn) == 0: + func = objective + restore = None + else: + if len(fixedn) == Nargs: + raise ValueError( + "All parameters fixed. There is nothing to optimize.") + + def restore(args, theta): + # Replace with theta for all numbers not in fixedn + # This allows the non-fixed values to vary, but + # we still call self.nnlf with all parameters. + i = 0 + for n in range(Nargs): + if n not in fixedn: + args[n] = theta[i] + i += 1 + return args + + def func(theta, x): + newtheta = restore(args[:], theta) + return objective(newtheta, x) + + return x0, func, restore, args + + def _moment_error(self, theta, x, data_moments): + loc, scale, args = self._unpack_loc_scale(theta) + if not self._argcheck(*args) or scale <= 0: + return inf + + dist_moments = np.array([self.moment(i+1, *args, loc=loc, scale=scale) + for i in range(len(data_moments))]) + if np.any(np.isnan(dist_moments)): + raise ValueError("Method of moments encountered a non-finite " + "distribution moment and cannot continue. " + "Consider trying method='MLE'.") + + return (((data_moments - dist_moments) / + np.maximum(np.abs(data_moments), 1e-8))**2).sum() + + def fit(self, data, *args, **kwds): + r""" + Return estimates of shape (if applicable), location, and scale + parameters from data. The default estimation method is Maximum + Likelihood Estimation (MLE), but Method of Moments (MM) + is also available. + + Starting estimates for the fit are given by input arguments; + for any arguments not provided with starting estimates, + ``self._fitstart(data)`` is called to generate such. + + One can hold some parameters fixed to specific values by passing in + keyword arguments ``f0``, ``f1``, ..., ``fn`` (for shape parameters) + and ``floc`` and ``fscale`` (for location and scale parameters, + respectively). + + Parameters + ---------- + data : array_like or `CensoredData` instance + Data to use in estimating the distribution parameters. + arg1, arg2, arg3,... : floats, optional + Starting value(s) for any shape-characterizing arguments (those not + provided will be determined by a call to ``_fitstart(data)``). + No default value. + **kwds : floats, optional + - `loc`: initial guess of the distribution's location parameter. + - `scale`: initial guess of the distribution's scale parameter. + + Special keyword arguments are recognized as holding certain + parameters fixed: + + - f0, ..., fn : hold respective shape parameters fixed. + Alternatively, shape parameters to fix can be specified by name. + For example, if ``self.shapes == "a, b"``, ``fa`` and ``fix_a`` + are equivalent to ``f0``, and ``fb`` and ``fix_b`` are + equivalent to ``f1``. + + - floc : hold location parameter fixed to specified value. + + - fscale : hold scale parameter fixed to specified value. + + - optimizer : The optimizer to use. The optimizer must take + ``func`` and starting position as the first two arguments, + plus ``args`` (for extra arguments to pass to the + function to be optimized) and ``disp``. + The ``fit`` method calls the optimizer with ``disp=0`` to suppress output. + The optimizer must return the estimated parameters. + + - method : The method to use. The default is "MLE" (Maximum + Likelihood Estimate); "MM" (Method of Moments) + is also available. + + Raises + ------ + TypeError, ValueError + If an input is invalid + `~scipy.stats.FitError` + If fitting fails or the fit produced would be invalid + + Returns + ------- + parameter_tuple : tuple of floats + Estimates for any shape parameters (if applicable), followed by + those for location and scale. For most random variables, shape + statistics will be returned, but there are exceptions (e.g. + ``norm``). + + Notes + ----- + With ``method="MLE"`` (default), the fit is computed by minimizing + the negative log-likelihood function. A large, finite penalty + (rather than infinite negative log-likelihood) is applied for + observations beyond the support of the distribution. + + With ``method="MM"``, the fit is computed by minimizing the L2 norm + of the relative errors between the first *k* raw (about zero) data + moments and the corresponding distribution moments, where *k* is the + number of non-fixed parameters. + More precisely, the objective function is:: + + (((data_moments - dist_moments) + / np.maximum(np.abs(data_moments), 1e-8))**2).sum() + + where the constant ``1e-8`` avoids division by zero in case of + vanishing data moments. Typically, this error norm can be reduced to + zero. + Note that the standard method of moments can produce parameters for + which some data are outside the support of the fitted distribution; + this implementation does nothing to prevent this. + + For either method, + the returned answer is not guaranteed to be globally optimal; it + may only be locally optimal, or the optimization may fail altogether. + If the data contain any of ``np.nan``, ``np.inf``, or ``-np.inf``, + the `fit` method will raise a ``RuntimeError``. + + When passing a ``CensoredData`` instance to ``data``, the log-likelihood + function is defined as: + + .. math:: + + l(\pmb{\theta}; k) & = \sum + \log(f(k_u; \pmb{\theta})) + + \sum + \log(F(k_l; \pmb{\theta})) \\ + & + \sum + \log(1 - F(k_r; \pmb{\theta})) \\ + & + \sum + \log(F(k_{\text{high}, i}; \pmb{\theta}) + - F(k_{\text{low}, i}; \pmb{\theta})) + + where :math:`f` and :math:`F` are the pdf and cdf, respectively, of the + function being fitted, :math:`\pmb{\theta}` is the parameter vector, + :math:`u` are the indices of uncensored observations, + :math:`l` are the indices of left-censored observations, + :math:`r` are the indices of right-censored observations, + subscripts "low"/"high" denote endpoints of interval-censored observations, and + :math:`i` are the indices of interval-censored observations. + + Examples + -------- + + Generate some data to fit: draw random variates from the `beta` + distribution + + >>> import numpy as np + >>> from scipy.stats import beta + >>> a, b = 1., 2. + >>> rng = np.random.default_rng(172786373191770012695001057628748821561) + >>> x = beta.rvs(a, b, size=1000, random_state=rng) + + Now we can fit all four parameters (``a``, ``b``, ``loc`` and + ``scale``): + + >>> a1, b1, loc1, scale1 = beta.fit(x) + >>> a1, b1, loc1, scale1 + (1.0198945204435628, 1.9484708982737828, 4.372241314917588e-05, 0.9979078845964814) + + The fit can be done also using a custom optimizer: + + >>> from scipy.optimize import minimize + >>> def custom_optimizer(func, x0, args=(), disp=0): + ... res = minimize(func, x0, args, method="slsqp", options={"disp": disp}) + ... if res.success: + ... return res.x + ... raise RuntimeError('optimization routine failed') + >>> a1, b1, loc1, scale1 = beta.fit(x, method="MLE", optimizer=custom_optimizer) + >>> a1, b1, loc1, scale1 + (1.0198821087258905, 1.948484145914738, 4.3705304486881485e-05, 0.9979104663953395) + + We can also use some prior knowledge about the dataset: let's keep + ``loc`` and ``scale`` fixed: + + >>> a1, b1, loc1, scale1 = beta.fit(x, floc=0, fscale=1) + >>> loc1, scale1 + (0, 1) + + We can also keep shape parameters fixed by using ``f``-keywords. To + keep the zero-th shape parameter ``a`` equal 1, use ``f0=1`` or, + equivalently, ``fa=1``: + + >>> a1, b1, loc1, scale1 = beta.fit(x, fa=1, floc=0, fscale=1) + >>> a1 + 1 + + Not all distributions return estimates for the shape parameters. + ``norm`` for example just returns estimates for location and scale: + + >>> from scipy.stats import norm + >>> x = norm.rvs(a, b, size=1000, random_state=123) + >>> loc1, scale1 = norm.fit(x) + >>> loc1, scale1 + (0.92087172783841631, 2.0015750750324668) + """ # noqa: E501 + method = kwds.get('method', "mle").lower() + + censored = isinstance(data, CensoredData) + if censored: + if method != 'mle': + raise ValueError('For censored data, the method must' + ' be "MLE".') + if data.num_censored() == 0: + # There are no censored values in data, so replace the + # CensoredData instance with a regular array. + data = data._uncensored + censored = False + + Narg = len(args) + if Narg > self.numargs: + raise TypeError("Too many input arguments.") + + # Check the finiteness of data only if data is not an instance of + # CensoredData. The arrays in a CensoredData instance have already + # been validated. + if not censored: + # Note: `ravel()` is called for backwards compatibility. + data = np.asarray(data).ravel() + if not np.isfinite(data).all(): + raise ValueError("The data contains non-finite values.") + + start = [None]*2 + if (Narg < self.numargs) or not ('loc' in kwds and + 'scale' in kwds): + # get distribution specific starting locations + start = self._fitstart(data) + args += start[Narg:-2] + loc = kwds.pop('loc', start[-2]) + scale = kwds.pop('scale', start[-1]) + args += (loc, scale) + x0, func, restore, args = self._reduce_func(args, kwds, data=data) + optimizer = kwds.pop('optimizer', optimize.fmin) + # convert string to function in scipy.optimize + optimizer = _fit_determine_optimizer(optimizer) + # by now kwds must be empty, since everybody took what they needed + if kwds: + raise TypeError(f"Unknown arguments: {kwds}.") + + # In some cases, method of moments can be done with fsolve/root + # instead of an optimizer, but sometimes no solution exists, + # especially when the user fixes parameters. Minimizing the sum + # of squares of the error generalizes to these cases. + vals = optimizer(func, x0, args=(data,), disp=0) + obj = func(vals, data) + + if restore is not None: + vals = restore(args, vals) + vals = tuple(vals) + + loc, scale, shapes = self._unpack_loc_scale(vals) + if not (np.all(self._argcheck(*shapes)) and scale > 0): + raise FitError("Optimization converged to parameters that are " + "outside the range allowed by the distribution.") + + if method == 'mm': + if not np.isfinite(obj): + raise FitError("Optimization failed: either a data moment " + "or fitted distribution moment is " + "non-finite.") + + return vals + + def _fit_loc_scale_support(self, data, *args): + """Estimate loc and scale parameters from data accounting for support. + + Parameters + ---------- + data : array_like + Data to fit. + arg1, arg2, arg3,... : array_like + The shape parameter(s) for the distribution (see docstring of the + instance object for more information). + + Returns + ------- + Lhat : float + Estimated location parameter for the data. + Shat : float + Estimated scale parameter for the data. + + """ + if isinstance(data, CensoredData): + # For this estimate, "uncensor" the data by taking the + # given endpoints as the data for the left- or right-censored + # data, and the mean for the interval-censored data. + data = data._uncensor() + else: + data = np.asarray(data) + + # Estimate location and scale according to the method of moments. + loc_hat, scale_hat = self.fit_loc_scale(data, *args) + + # Compute the support according to the shape parameters. + self._argcheck(*args) + _a, _b = self._get_support(*args) + a, b = _a, _b + support_width = b - a + + # If the support is empty then return the moment-based estimates. + if support_width <= 0: + return loc_hat, scale_hat + + # Compute the proposed support according to the loc and scale + # estimates. + a_hat = loc_hat + a * scale_hat + b_hat = loc_hat + b * scale_hat + + # Use the moment-based estimates if they are compatible with the data. + data_a = np.min(data) + data_b = np.max(data) + if a_hat < data_a and data_b < b_hat: + return loc_hat, scale_hat + + # Otherwise find other estimates that are compatible with the data. + data_width = data_b - data_a + rel_margin = 0.1 + margin = data_width * rel_margin + + # For a finite interval, both the location and scale + # should have interesting values. + if support_width < np.inf: + loc_hat = (data_a - a) - margin + scale_hat = (data_width + 2 * margin) / support_width + return loc_hat, scale_hat + + # For a one-sided interval, use only an interesting location parameter. + if a > -np.inf: + return (data_a - a) - margin, 1 + elif b < np.inf: + return (data_b - b) + margin, 1 + else: + raise RuntimeError + + def fit_loc_scale(self, data, *args): + """ + Estimate loc and scale parameters from data using 1st and 2nd moments. + + Parameters + ---------- + data : array_like + Data to fit. + arg1, arg2, arg3,... : array_like + The shape parameter(s) for the distribution (see docstring of the + instance object for more information). + + Returns + ------- + Lhat : float + Estimated location parameter for the data. + Shat : float + Estimated scale parameter for the data. + + """ + mu, mu2 = self.stats(*args, **{'moments': 'mv'}) + tmp = asarray(data) + muhat = tmp.mean() + mu2hat = tmp.var() + Shat = sqrt(mu2hat / mu2) + with np.errstate(invalid='ignore'): + Lhat = muhat - Shat*mu + if not np.isfinite(Lhat): + Lhat = 0 + if not (np.isfinite(Shat) and (0 < Shat)): + Shat = 1 + return Lhat, Shat + + def _entropy(self, *args): + def integ(x): + val = self._pdf(x, *args) + return entr(val) + + # upper limit is often inf, so suppress warnings when integrating + _a, _b = self._get_support(*args) + with np.errstate(over='ignore'): + h = integrate.quad(integ, _a, _b)[0] + + if not np.isnan(h): + return h + else: + # try with different limits if integration problems + low, upp = self.ppf([1e-10, 1. - 1e-10], *args) + if np.isinf(_b): + upper = upp + else: + upper = _b + if np.isinf(_a): + lower = low + else: + lower = _a + return integrate.quad(integ, lower, upper)[0] + + def expect(self, func=None, args=(), loc=0, scale=1, lb=None, ub=None, + conditional=False, **kwds): + """Calculate expected value of a function with respect to the + distribution by numerical integration. + + The expected value of a function ``f(x)`` with respect to a + distribution ``dist`` is defined as:: + + ub + E[f(x)] = Integral(f(x) * dist.pdf(x)), + lb + + where ``ub`` and ``lb`` are arguments and ``x`` has the ``dist.pdf(x)`` + distribution. If the bounds ``lb`` and ``ub`` correspond to the + support of the distribution, e.g. ``[-inf, inf]`` in the default + case, then the integral is the unrestricted expectation of ``f(x)``. + Also, the function ``f(x)`` may be defined such that ``f(x)`` is ``0`` + outside a finite interval in which case the expectation is + calculated within the finite range ``[lb, ub]``. + + Parameters + ---------- + func : callable, optional + Function for which integral is calculated. Takes only one argument. + The default is the identity mapping f(x) = x. + args : tuple, optional + Shape parameters of the distribution. + loc : float, optional + Location parameter (default=0). + scale : float, optional + Scale parameter (default=1). + lb, ub : scalar, optional + Lower and upper bound for integration. Default is set to the + support of the distribution. + conditional : bool, optional + If True, the integral is corrected by the conditional probability + of the integration interval. The return value is the expectation + of the function, conditional on being in the given interval. + Default is False. + + Additional keyword arguments are passed to the integration routine. + + Returns + ------- + expect : float + The calculated expected value. + + Notes + ----- + The integration behavior of this function is inherited from + `scipy.integrate.quad`. Neither this function nor + `scipy.integrate.quad` can verify whether the integral exists or is + finite. For example ``cauchy(0).mean()`` returns ``np.nan`` and + ``cauchy(0).expect()`` returns ``0.0``. + + Likewise, the accuracy of results is not verified by the function. + `scipy.integrate.quad` is typically reliable for integrals that are + numerically favorable, but it is not guaranteed to converge + to a correct value for all possible intervals and integrands. This + function is provided for convenience; for critical applications, + check results against other integration methods. + + The function is not vectorized. + + Examples + -------- + + To understand the effect of the bounds of integration consider + + >>> from scipy.stats import expon + >>> expon(1).expect(lambda x: 1, lb=0.0, ub=2.0) + 0.6321205588285578 + + This is close to + + >>> expon(1).cdf(2.0) - expon(1).cdf(0.0) + 0.6321205588285577 + + If ``conditional=True`` + + >>> expon(1).expect(lambda x: 1, lb=0.0, ub=2.0, conditional=True) + 1.0000000000000002 + + The slight deviation from 1 is due to numerical integration. + + The integrand can be treated as a complex-valued function + by passing ``complex_func=True`` to `scipy.integrate.quad` . + + >>> import numpy as np + >>> from scipy.stats import vonmises + >>> res = vonmises(loc=2, kappa=1).expect(lambda x: np.exp(1j*x), + ... complex_func=True) + >>> res + (-0.18576377217422957+0.40590124735052263j) + + >>> np.angle(res) # location of the (circular) distribution + 2.0 + + """ + lockwds = {'loc': loc, + 'scale': scale} + self._argcheck(*args) + _a, _b = self._get_support(*args) + if func is None: + def fun(x, *args): + return x * self.pdf(x, *args, **lockwds) + else: + def fun(x, *args): + return func(x) * self.pdf(x, *args, **lockwds) + if lb is None: + lb = loc + _a * scale + if ub is None: + ub = loc + _b * scale + + cdf_bounds = self.cdf([lb, ub], *args, **lockwds) + invfac = cdf_bounds[1] - cdf_bounds[0] + + kwds['args'] = args + + # split interval to help integrator w/ infinite support; see gh-8928 + alpha = 0.05 # split body from tails at probability mass `alpha` + inner_bounds = np.array([alpha, 1-alpha]) + cdf_inner_bounds = cdf_bounds[0] + invfac * inner_bounds + c, d = loc + self._ppf(cdf_inner_bounds, *args) * scale + + # Do not silence warnings from integration. + lbc = integrate.quad(fun, lb, c, **kwds)[0] + cd = integrate.quad(fun, c, d, **kwds)[0] + dub = integrate.quad(fun, d, ub, **kwds)[0] + vals = (lbc + cd + dub) + + if conditional: + vals /= invfac + return np.array(vals)[()] # make it a numpy scalar like other methods + + def _param_info(self): + shape_info = self._shape_info() + loc_info = _ShapeInfo("loc", False, (-np.inf, np.inf), (False, False)) + scale_info = _ShapeInfo("scale", False, (0, np.inf), (False, False)) + param_info = shape_info + [loc_info, scale_info] + return param_info + + # For now, _delta_cdf is a private method. + def _delta_cdf(self, x1, x2, *args, loc=0, scale=1): + """ + Compute CDF(x2) - CDF(x1). + + Where x1 is greater than the median, compute SF(x1) - SF(x2), + otherwise compute CDF(x2) - CDF(x1). + + This function is only useful if `dist.sf(x, ...)` has an implementation + that is numerically more accurate than `1 - dist.cdf(x, ...)`. + """ + cdf1 = self.cdf(x1, *args, loc=loc, scale=scale) + # Possible optimizations (needs investigation-these might not be + # better): + # * Use xpx.apply_where instead of np.where + # * Instead of cdf1 > 0.5, compare x1 to the median. + result = np.where(cdf1 > 0.5, + (self.sf(x1, *args, loc=loc, scale=scale) + - self.sf(x2, *args, loc=loc, scale=scale)), + self.cdf(x2, *args, loc=loc, scale=scale) - cdf1) + if result.ndim == 0: + result = result[()] + return result + + +# Helpers for the discrete distributions +def _drv2_moment(self, n, *args): + """Non-central moment of discrete distribution.""" + def fun(x): + return np.power(x, n) * self._pmf(x, *args) + + _a, _b = self._get_support(*args) + return _expect(fun, _a, _b, self._ppf(0.5, *args), self.inc) + + +def _drv2_ppfsingle(self, q, *args): # Use basic bisection algorithm + _a, _b = self._get_support(*args) + b = _b + a = _a + + step = 10 + if isinf(b): # Be sure ending point is > q + b = float(max(100*q, 10)) + while 1: + if b >= _b: + qb = 1.0 + break + qb = self._cdf(b, *args) + if (qb < q): + b += step + step *= 2 + else: + break + else: + qb = 1.0 + + step = 10 + if isinf(a): # be sure starting point < q + a = float(min(-100*q, -10)) + while 1: + if a <= _a: + qb = 0.0 + break + qa = self._cdf(a, *args) + if (qa > q): + a -= step + step *= 2 + else: + break + else: + qa = self._cdf(a, *args) + + if np.isinf(a) or np.isinf(b): + message = "Arguments that bracket the requested quantile could not be found." + raise RuntimeError(message) + + # maximum number of bisections within the normal float64s + # maxiter = int(np.log2(finfo.max) - np.log2(finfo.smallest_normal)) + maxiter = 2046 + for i in range(maxiter): + if (qa == q): + return a + if (qb == q): + return b + if b <= a+1: + if qa > q: + return a + else: + return b + c = int((a+b)/2.0) + qc = self._cdf(c, *args) + if (qc < q): + if a != c: + a = c + else: + raise RuntimeError('updating stopped, endless loop') + qa = qc + elif (qc > q): + if b != c: + b = c + else: + raise RuntimeError('updating stopped, endless loop') + qb = qc + else: + return c + + +# Must over-ride one of _pmf or _cdf or pass in +# x_k, p(x_k) lists in initialization + + +class rv_discrete(rv_generic): + """A generic discrete random variable class meant for subclassing. + + `rv_discrete` is a base class to construct specific distribution classes + and instances for discrete random variables. It can also be used + to construct an arbitrary distribution defined by a list of support + points and corresponding probabilities. + + Parameters + ---------- + a : float, optional + Lower bound of the support of the distribution, default: 0 + b : float, optional + Upper bound of the support of the distribution, default: plus infinity + moment_tol : float, optional + The tolerance for the generic calculation of moments. + values : tuple of two array_like, optional + ``(xk, pk)`` where ``xk`` are integers and ``pk`` are the non-zero + probabilities between 0 and 1 with ``sum(pk) = 1``. ``xk`` + and ``pk`` must have the same shape, and ``xk`` must be unique. + inc : int, optional + Increment for the support of the distribution. + Default is 1. (other values have not been tested) + badvalue : float, optional + The value in a result arrays that indicates a value that for which + some argument restriction is violated, default is np.nan. + name : str, optional + The name of the instance. This string is used to construct the default + example for distributions. + longname : str, optional + This string is used as part of the first line of the docstring returned + when a subclass has no docstring of its own. Note: `longname` exists + for backwards compatibility, do not use for new subclasses. + shapes : str, optional + The shape of the distribution. For example "m, n" for a distribution + that takes two integers as the two shape arguments for all its methods + If not provided, shape parameters will be inferred from + the signatures of the private methods, ``_pmf`` and ``_cdf`` of + the instance. + seed : {None, int, `numpy.random.Generator`, `numpy.random.RandomState`}, optional + If `seed` is None (or `np.random`), the `numpy.random.RandomState` + singleton is used. + If `seed` is an int, a new ``RandomState`` instance is used, + seeded with `seed`. + If `seed` is already a ``Generator`` or ``RandomState`` instance then + that instance is used. + + Attributes + ---------- + a, b : float, optional + Lower/upper bound of the support of the unshifted/unscaled distribution. + This value is unaffected by the `loc` and `scale` parameters. + To calculate the support of the shifted/scaled distribution, + use the `support` method. + + Methods + ------- + rvs + pmf + logpmf + cdf + logcdf + sf + logsf + ppf + isf + moment + stats + entropy + expect + median + mean + std + var + interval + __call__ + support + + Notes + ----- + This class is similar to `rv_continuous`. Whether a shape parameter is + valid is decided by an ``_argcheck`` method (which defaults to checking + that its arguments are strictly positive.) + The main differences are as follows. + + - The support of the distribution is a set of integers. + - Instead of the probability density function, ``pdf`` (and the + corresponding private ``_pdf``), this class defines the + *probability mass function*, `pmf` (and the corresponding + private ``_pmf``.) + - There is no ``scale`` parameter. + - The default implementations of methods (e.g. ``_cdf``) are not designed + for distributions with support that is unbounded below (i.e. + ``a=-np.inf``), so they must be overridden. + + To create a new discrete distribution, we would do the following: + + >>> from scipy.stats import rv_discrete + >>> class poisson_gen(rv_discrete): + ... "Poisson distribution" + ... def _pmf(self, k, mu): + ... return exp(-mu) * mu**k / factorial(k) + + and create an instance:: + + >>> poisson = poisson_gen(name="poisson") + + Note that above we defined the Poisson distribution in the standard form. + Shifting the distribution can be done by providing the ``loc`` parameter + to the methods of the instance. For example, ``poisson.pmf(x, mu, loc)`` + delegates the work to ``poisson._pmf(x-loc, mu)``. + + **Discrete distributions from a list of probabilities** + + Alternatively, you can construct an arbitrary discrete rv defined + on a finite set of values ``xk`` with ``Prob{X=xk} = pk`` by using the + ``values`` keyword argument to the `rv_discrete` constructor. + + **Deepcopying / Pickling** + + If a distribution or frozen distribution is deepcopied (pickled/unpickled, + etc.), any underlying random number generator is deepcopied with it. An + implication is that if a distribution relies on the singleton RandomState + before copying, it will rely on a copy of that random state after copying, + and ``np.random.seed`` will no longer control the state. + + Examples + -------- + Custom made discrete distribution: + + >>> import numpy as np + >>> from scipy import stats + >>> xk = np.arange(7) + >>> pk = (0.1, 0.2, 0.3, 0.1, 0.1, 0.0, 0.2) + >>> custm = stats.rv_discrete(name='custm', values=(xk, pk)) + >>> + >>> import matplotlib.pyplot as plt + >>> fig, ax = plt.subplots(1, 1) + >>> ax.plot(xk, custm.pmf(xk), 'ro', ms=12, mec='r') + >>> ax.vlines(xk, 0, custm.pmf(xk), colors='r', lw=4) + >>> plt.show() + + Random number generation: + + >>> R = custm.rvs(size=100) + + """ + def __new__(cls, a=0, b=inf, name=None, badvalue=None, + moment_tol=1e-8, values=None, inc=1, longname=None, + shapes=None, seed=None): + + if values is not None: + # dispatch to a subclass + return super().__new__(rv_sample) + else: + # business as usual + return super().__new__(cls) + + def __init__(self, a=0, b=inf, name=None, badvalue=None, + moment_tol=1e-8, values=None, inc=1, longname=None, + shapes=None, seed=None): + + super().__init__(seed) + + # cf generic freeze + self._ctor_param = dict( + a=a, b=b, name=name, badvalue=badvalue, + moment_tol=moment_tol, values=values, inc=inc, + longname=longname, shapes=shapes, seed=seed) + + if badvalue is None: + badvalue = nan + self.badvalue = badvalue + self.a = a + self.b = b + self.moment_tol = moment_tol + self.inc = inc + self.shapes = shapes + + if values is not None: + raise ValueError("rv_discrete.__init__(..., values != None, ...)") + + self._construct_argparser(meths_to_inspect=[self._pmf, self._cdf], + locscale_in='loc=0', + # scale=1 for discrete RVs + locscale_out='loc, 1') + self._attach_methods() + self._construct_docstrings(name, longname) + + def __getstate__(self): + dct = self.__dict__.copy() + # these methods will be remade in __setstate__ + attrs = ["_parse_args", "_parse_args_stats", "_parse_args_rvs", + "_cdfvec", "_ppfvec", "generic_moment"] + [dct.pop(attr, None) for attr in attrs] + return dct + + def _attach_methods(self): + """Attaches dynamically created methods to the rv_discrete instance.""" + self._cdfvec = vectorize(self._cdf_single, otypes='d') + self.vecentropy = vectorize(self._entropy) + + # _attach_methods is responsible for calling _attach_argparser_methods + self._attach_argparser_methods() + + # nin correction needs to be after we know numargs + # correct nin for generic moment vectorization + _vec_generic_moment = vectorize(_drv2_moment, otypes='d') + _vec_generic_moment.nin = self.numargs + 2 + self.generic_moment = types.MethodType(_vec_generic_moment, self) + + # correct nin for ppf vectorization + _vppf = vectorize(_drv2_ppfsingle, otypes='d') + _vppf.nin = self.numargs + 2 + self._ppfvec = types.MethodType(_vppf, self) + + # now that self.numargs is defined, we can adjust nin + self._cdfvec.nin = self.numargs + 1 + + def _construct_docstrings(self, name, longname): + if name is None: + name = 'Distribution' + self.name = name + + # generate docstring for subclass instances + if longname is None: + if name[0] in ['aeiouAEIOU']: + hstr = "An " + else: + hstr = "A " + longname = hstr + name + + if sys.flags.optimize < 2: + # Skip adding docstrings if interpreter is run with -OO + if self.__doc__ is None: + self._construct_default_doc(longname=longname, + docdict=docdict_discrete, + discrete='discrete') + else: + dct = dict(distdiscrete) + self._construct_doc(docdict_discrete, dct.get(self.name)) + + # discrete RV do not have the scale parameter, remove it + self.__doc__ = self.__doc__.replace( + '\n scale : array_like, ' + 'optional\n scale parameter (default=1)', '') + + def _updated_ctor_param(self): + """Return the current version of _ctor_param, possibly updated by user. + + Used by freezing. + Keep this in sync with the signature of __init__. + """ + dct = self._ctor_param.copy() + dct['a'] = self.a + dct['b'] = self.b + dct['badvalue'] = self.badvalue + dct['moment_tol'] = self.moment_tol + dct['inc'] = self.inc + dct['name'] = self.name + dct['shapes'] = self.shapes + return dct + + def _nonzero(self, k, *args): + return floor(k) == k + + def _pmf(self, k, *args): + return self._cdf(k, *args) - self._cdf(k-1, *args) + + def _logpmf(self, k, *args): + with np.errstate(divide='ignore'): + return log(self._pmf(k, *args)) + + def _logpxf(self, k, *args): + # continuous distributions have PDF, discrete have PMF, but sometimes + # the distinction doesn't matter. This lets us use `_logpxf` for both + # discrete and continuous distributions. + return self._logpmf(k, *args) + + def _unpack_loc_scale(self, theta): + try: + loc = theta[-1] + scale = 1 + args = tuple(theta[:-1]) + except IndexError as e: + raise ValueError("Not enough input arguments.") from e + return loc, scale, args + + def _cdf_single(self, k, *args): + _a, _b = self._get_support(*args) + m = arange(int(_a), k+1) + return np.sum(self._pmf(m, *args), axis=0) + + def _cdf(self, x, *args): + k = floor(x).astype(np.float64) + return self._cdfvec(k, *args) + + # generic _logcdf, _sf, _logsf, _ppf, _isf, _rvs defined in rv_generic + + def rvs(self, *args, **kwargs): + """Random variates of given type. + + Parameters + ---------- + arg1, arg2, arg3,... : array_like + The shape parameter(s) for the distribution (see docstring of the + instance object for more information). + loc : array_like, optional + Location parameter (default=0). + size : int or tuple of ints, optional + Defining number of random variates (Default is 1). Note that `size` + has to be given as keyword, not as positional argument. + random_state : {None, int, `numpy.random.Generator`, + `numpy.random.RandomState`}, optional + + If `random_state` is None (or `np.random`), the + `numpy.random.RandomState` singleton is used. + If `random_state` is an int, a new ``RandomState`` instance is + used, seeded with `random_state`. + If `random_state` is already a ``Generator`` or ``RandomState`` + instance, that instance is used. + + Returns + ------- + rvs : ndarray or scalar + Random variates of given `size`. + + """ + kwargs['discrete'] = True + return super().rvs(*args, **kwargs) + + def pmf(self, k, *args, **kwds): + """Probability mass function at k of the given RV. + + Parameters + ---------- + k : array_like + Quantiles. + arg1, arg2, arg3,... : array_like + The shape parameter(s) for the distribution (see docstring of the + instance object for more information) + loc : array_like, optional + Location parameter (default=0). + + Returns + ------- + pmf : array_like + Probability mass function evaluated at k + + """ + args, loc, _ = self._parse_args(*args, **kwds) + k, loc = map(asarray, (k, loc)) + args = tuple(map(asarray, args)) + _a, _b = self._get_support(*args) + k = asarray(k-loc) + cond0 = self._argcheck(*args) + cond1 = (k >= _a) & (k <= _b) + if not isinstance(self, rv_sample): + cond1 = cond1 & self._nonzero(k, *args) + cond = cond0 & cond1 + output = zeros(shape(cond), 'd') + place(output, (1-cond0) + np.isnan(k), self.badvalue) + if np.any(cond): + goodargs = argsreduce(cond, *((k,)+args)) + place(output, cond, np.clip(self._pmf(*goodargs), 0, 1)) + if output.ndim == 0: + return output[()] + return output + + def logpmf(self, k, *args, **kwds): + """Log of the probability mass function at k of the given RV. + + Parameters + ---------- + k : array_like + Quantiles. + arg1, arg2, arg3,... : array_like + The shape parameter(s) for the distribution (see docstring of the + instance object for more information). + loc : array_like, optional + Location parameter. Default is 0. + + Returns + ------- + logpmf : array_like + Log of the probability mass function evaluated at k. + + """ + args, loc, _ = self._parse_args(*args, **kwds) + k, loc = map(asarray, (k, loc)) + args = tuple(map(asarray, args)) + _a, _b = self._get_support(*args) + k = asarray(k-loc) + cond0 = self._argcheck(*args) + cond1 = (k >= _a) & (k <= _b) + if not isinstance(self, rv_sample): + cond1 = cond1 & self._nonzero(k, *args) + cond = cond0 & cond1 + output = empty(shape(cond), 'd') + output.fill(-inf) + place(output, (1-cond0) + np.isnan(k), self.badvalue) + if np.any(cond): + goodargs = argsreduce(cond, *((k,)+args)) + place(output, cond, self._logpmf(*goodargs)) + if output.ndim == 0: + return output[()] + return output + + def cdf(self, k, *args, **kwds): + """Cumulative distribution function of the given RV. + + Parameters + ---------- + k : array_like, int + Quantiles. + arg1, arg2, arg3,... : array_like + The shape parameter(s) for the distribution (see docstring of the + instance object for more information). + loc : array_like, optional + Location parameter (default=0). + + Returns + ------- + cdf : ndarray + Cumulative distribution function evaluated at `k`. + + """ + args, loc, _ = self._parse_args(*args, **kwds) + k, loc = map(asarray, (k, loc)) + args = tuple(map(asarray, args)) + _a, _b = self._get_support(*args) + k = asarray(k-loc) + cond0 = self._argcheck(*args) + cond1 = (k >= _a) & (k < _b) + cond2 = (k >= _b) + cond3 = np.isneginf(k) + cond = cond0 & cond1 & np.isfinite(k) + + output = zeros(shape(cond), 'd') + place(output, cond2*(cond0 == cond0), 1.0) + place(output, cond3*(cond0 == cond0), 0.0) + place(output, (1-cond0) + np.isnan(k), self.badvalue) + + if np.any(cond): + goodargs = argsreduce(cond, *((k,)+args)) + place(output, cond, np.clip(self._cdf(*goodargs), 0, 1)) + if output.ndim == 0: + return output[()] + return output + + def logcdf(self, k, *args, **kwds): + """Log of the cumulative distribution function at k of the given RV. + + Parameters + ---------- + k : array_like, int + Quantiles. + arg1, arg2, arg3,... : array_like + The shape parameter(s) for the distribution (see docstring of the + instance object for more information). + loc : array_like, optional + Location parameter (default=0). + + Returns + ------- + logcdf : array_like + Log of the cumulative distribution function evaluated at k. + + """ + args, loc, _ = self._parse_args(*args, **kwds) + k, loc = map(asarray, (k, loc)) + args = tuple(map(asarray, args)) + _a, _b = self._get_support(*args) + k = asarray(k-loc) + cond0 = self._argcheck(*args) + cond1 = (k >= _a) & (k < _b) + cond2 = (k >= _b) + cond = cond0 & cond1 + output = empty(shape(cond), 'd') + output.fill(-inf) + place(output, (1-cond0) + np.isnan(k), self.badvalue) + place(output, cond2*(cond0 == cond0), 0.0) + + if np.any(cond): + goodargs = argsreduce(cond, *((k,)+args)) + place(output, cond, self._logcdf(*goodargs)) + if output.ndim == 0: + return output[()] + return output + + def sf(self, k, *args, **kwds): + """Survival function (1 - `cdf`) at k of the given RV. + + Parameters + ---------- + k : array_like + Quantiles. + arg1, arg2, arg3,... : array_like + The shape parameter(s) for the distribution (see docstring of the + instance object for more information). + loc : array_like, optional + Location parameter (default=0). + + Returns + ------- + sf : array_like + Survival function evaluated at k. + + """ + args, loc, _ = self._parse_args(*args, **kwds) + k, loc = map(asarray, (k, loc)) + args = tuple(map(asarray, args)) + _a, _b = self._get_support(*args) + k = asarray(k-loc) + cond0 = self._argcheck(*args) + cond1 = (k >= _a) & (k < _b) + cond2 = ((k < _a) | np.isneginf(k)) & cond0 + cond = cond0 & cond1 & np.isfinite(k) + output = zeros(shape(cond), 'd') + place(output, (1-cond0) + np.isnan(k), self.badvalue) + place(output, cond2, 1.0) + if np.any(cond): + goodargs = argsreduce(cond, *((k,)+args)) + place(output, cond, np.clip(self._sf(*goodargs), 0, 1)) + if output.ndim == 0: + return output[()] + return output + + def logsf(self, k, *args, **kwds): + """Log of the survival function of the given RV. + + Returns the log of the "survival function," defined as 1 - `cdf`, + evaluated at `k`. + + Parameters + ---------- + k : array_like + Quantiles. + arg1, arg2, arg3,... : array_like + The shape parameter(s) for the distribution (see docstring of the + instance object for more information). + loc : array_like, optional + Location parameter (default=0). + + Returns + ------- + logsf : ndarray + Log of the survival function evaluated at `k`. + + """ + args, loc, _ = self._parse_args(*args, **kwds) + k, loc = map(asarray, (k, loc)) + args = tuple(map(asarray, args)) + _a, _b = self._get_support(*args) + k = asarray(k-loc) + cond0 = self._argcheck(*args) + cond1 = (k >= _a) & (k < _b) + cond2 = (k < _a) & cond0 + cond = cond0 & cond1 + output = empty(shape(cond), 'd') + output.fill(-inf) + place(output, (1-cond0) + np.isnan(k), self.badvalue) + place(output, cond2, 0.0) + if np.any(cond): + goodargs = argsreduce(cond, *((k,)+args)) + place(output, cond, self._logsf(*goodargs)) + if output.ndim == 0: + return output[()] + return output + + def ppf(self, q, *args, **kwds): + """Percent point function (inverse of `cdf`) at q of the given RV. + + Parameters + ---------- + q : array_like + Lower tail probability. + arg1, arg2, arg3,... : array_like + The shape parameter(s) for the distribution (see docstring of the + instance object for more information). + loc : array_like, optional + Location parameter (default=0). + + Returns + ------- + k : array_like + Quantile corresponding to the lower tail probability, q. + + Notes + ----- + For discrete distributions, the `cdf` is not strictly invertible. By convention, + this method returns the minimum value `k` for which the `cdf` at `k` is at + least `q`. There is one exception: the `ppf` of ``0`` is ``a-1``, + where ``a`` is the left endpoint of the support. + + """ + args, loc, _ = self._parse_args(*args, **kwds) + q, loc = map(asarray, (q, loc)) + args = tuple(map(asarray, args)) + _a, _b = self._get_support(*args) + cond0 = self._argcheck(*args) & (loc == loc) + cond1 = (q > 0) & (q < 1) + cond2 = (q == 0) & cond0 + cond3 = (q == 1) & cond0 + cond = cond0 & cond1 + output = np.full(shape(cond), fill_value=self.badvalue, dtype='d') + # output type 'd' to handle nin and inf + + place(output, cond2, argsreduce(cond2, _a-1 + loc)[0]) + place(output, cond3, argsreduce(cond3, _b + loc)[0]) + if np.any(cond): + goodargs = argsreduce(cond, *((q,)+args+(loc,))) + loc, goodargs = goodargs[-1], goodargs[:-1] + place(output, cond, self._ppf(*goodargs) + loc) + + if output.ndim == 0: + return output[()] + return output + + def isf(self, q, *args, **kwds): + """Inverse survival function (inverse of `sf`) at q of the given RV. + + Parameters + ---------- + q : array_like + Upper tail probability. + arg1, arg2, arg3,... : array_like + The shape parameter(s) for the distribution (see docstring of the + instance object for more information). + loc : array_like, optional + Location parameter (default=0). + + Returns + ------- + k : ndarray or scalar + Quantile corresponding to the upper tail probability, q. + + Notes + ----- + For discrete distributions, the `sf` is not strictly invertible. By convention, + this method returns the minimum value `k` for which the `sf` at `k` is + no greater than `q`. There is one exception: the `isf` of ``1`` is ``a-1``, + where ``a`` is the left endpoint of the support. + + """ + args, loc, _ = self._parse_args(*args, **kwds) + q, loc = map(asarray, (q, loc)) + args = tuple(map(asarray, args)) + _a, _b = self._get_support(*args) + cond0 = self._argcheck(*args) & (loc == loc) + cond1 = (q > 0) & (q < 1) + cond2 = (q == 1) & cond0 + cond3 = (q == 0) & cond0 + cond = cond0 & cond1 + + # same problem as with ppf; copied from ppf and changed + output = np.full(shape(cond), fill_value=self.badvalue, dtype='d') + # output type 'd' to handle nin and inf + lower_bound = _a - 1 + loc + upper_bound = _b + loc + place(output, cond2, argsreduce(cond2, lower_bound)[0]) + place(output, cond3, argsreduce(cond3, upper_bound)[0]) + + # call place only if at least 1 valid argument + if np.any(cond): + goodargs = argsreduce(cond, *((q,)+args+(loc,))) + loc, goodargs = goodargs[-1], goodargs[:-1] + # PB same as ticket 766 + place(output, cond, self._isf(*goodargs) + loc) + + if output.ndim == 0: + return output[()] + return output + + def _entropy(self, *args): + if hasattr(self, 'pk'): + return stats.entropy(self.pk) + else: + _a, _b = self._get_support(*args) + return _expect(lambda x: entr(self._pmf(x, *args)), + _a, _b, self._ppf(0.5, *args), self.inc) + + def expect(self, func=None, args=(), loc=0, lb=None, ub=None, + conditional=False, maxcount=1000, tolerance=1e-10, chunksize=32): + """ + Calculate expected value of a function with respect to the distribution + for discrete distribution by numerical summation. + + Parameters + ---------- + func : callable, optional + Function for which the expectation value is calculated. + Takes only one argument. + The default is the identity mapping f(k) = k. + args : tuple, optional + Shape parameters of the distribution. + loc : float, optional + Location parameter. + Default is 0. + lb, ub : int, optional + Lower and upper bound for the summation, default is set to the + support of the distribution, inclusive (``lb <= k <= ub``). + conditional : bool, optional + If true then the expectation is corrected by the conditional + probability of the summation interval. The return value is the + expectation of the function, `func`, conditional on being in + the given interval (k such that ``lb <= k <= ub``). + Default is False. + maxcount : int, optional + Maximal number of terms to evaluate (to avoid an endless loop for + an infinite sum). Default is 1000. + tolerance : float, optional + Absolute tolerance for the summation. Default is 1e-10. + chunksize : int, optional + Iterate over the support of a distributions in chunks of this size. + Default is 32. + + Returns + ------- + expect : float + Expected value. + + Notes + ----- + For heavy-tailed distributions, the expected value may or + may not exist, + depending on the function, `func`. If it does exist, but the + sum converges + slowly, the accuracy of the result may be rather low. For instance, for + ``zipf(4)``, accuracy for mean, variance in example is only 1e-5. + increasing `maxcount` and/or `chunksize` may improve the result, + but may also make zipf very slow. + + The function is not vectorized. + + """ + # Although `args` is just the shape parameters, `poisson_binom` needs this + # to split the vector-valued shape into a tuple of separate shapes + args, _, _ = self._parse_args(*args) + + if func is None: + def fun(x): + # loc and args from outer scope + return (x+loc)*self._pmf(x, *args) + else: + def fun(x): + # loc and args from outer scope + return func(x+loc)*self._pmf(x, *args) + # used pmf because _pmf does not check support in randint and there + # might be problems(?) with correct self.a, self.b at this stage maybe + # not anymore, seems to work now with _pmf + + _a, _b = self._get_support(*args) + if lb is None: + lb = _a + else: + lb = lb - loc # convert bound for standardized distribution + if ub is None: + ub = _b + else: + ub = ub - loc # convert bound for standardized distribution + if conditional: + invfac = self.sf(lb-1, *args) - self.sf(ub, *args) + else: + invfac = 1.0 + + if isinstance(self, rv_sample): + res = self._expect(fun, lb, ub) + return res / invfac + + # iterate over the support, starting from the median + x0 = self._ppf(0.5, *args) + res = _expect(fun, lb, ub, x0, self.inc, maxcount, tolerance, chunksize) + return res / invfac + + def _param_info(self): + shape_info = self._shape_info() + loc_info = _ShapeInfo("loc", True, (-np.inf, np.inf), (False, False)) + param_info = shape_info + [loc_info] + return param_info + + # these methods are identical to those of `rv_generic` except that + # they don't accept the `scale` parameter. They need their own copy + # of the methods so we can adjust the documentation. + + def entropy(self, *args, **kwargs): + return super().entropy(*args, **kwargs) + + def interval(self, confidence, *args, **kwargs): + return super().interval(confidence, *args, **kwargs) + + def mean(self, *args, **kwargs): + return super().mean(*args, **kwargs) + + def median(self, *args, **kwargs): + return super().median(*args, **kwargs) + + def moment(self, order, *args, **kwargs): + return super().moment(order, *args, **kwargs) + + def stats(self, *args, **kwargs): + return super().stats(*args, **kwargs) + + def std(self, *args, **kwargs): + return super().std(*args, **kwargs) + + def support(self, *args, **kwargs): + return super().support(*args, **kwargs) + + def var(self, *args, **kwargs): + return super().var(*args, **kwargs) + + +_remove_scale_methods = ['entropy', 'interval', 'mean', 'median', + 'moment', 'stats', 'std', 'support', 'var'] +for method_name in _remove_scale_methods: + # remove `scale` parameter from method documentation + method = getattr(rv_discrete, method_name) + doc = FunctionDoc(method) + doc['Parameters'] = [p for p in doc['Parameters'] if p.name != 'scale'] + doc = str(doc).split("\n", 1)[1].lstrip(" \n") # remove signature + method.__doc__ = doc +del method_name, method, doc # pyrefly:ignore[unbound-name] + + +def _expect(fun, lb, ub, x0, inc, maxcount=1000, tolerance=1e-10, + chunksize=32): + """Helper for computing the expectation value of `fun`.""" + # short-circuit if the support size is small enough + if (ub - lb) <= chunksize: + supp = np.arange(lb, ub+1, inc) + vals = fun(supp) + return np.sum(vals) + + # otherwise, iterate starting from x0 + if x0 < lb: + x0 = lb + if x0 > ub: + x0 = ub + + count, tot = 0, 0. + # iterate over [x0, ub] inclusive + for x in _iter_chunked(x0, ub+1, chunksize=chunksize, inc=inc): + count += x.size + delta = np.sum(fun(x)) + tot += delta + if abs(delta) < tolerance * x.size: + break + if count > maxcount: + warnings.warn('expect(): sum did not converge', + RuntimeWarning, stacklevel=3) + return tot + + # iterate over [lb, x0) + for x in _iter_chunked(x0-1, lb-1, chunksize=chunksize, inc=-inc): + count += x.size + delta = np.sum(fun(x)) + tot += delta + if abs(delta) < tolerance * x.size: + break + if count > maxcount: + warnings.warn('expect(): sum did not converge', + RuntimeWarning, stacklevel=3) + break + + return tot + + +def _iter_chunked(x0, x1, chunksize=4, inc=1): + """Iterate from x0 to x1 in chunks of chunksize and steps inc. + + x0 must be finite, x1 need not be. In the latter case, the iterator is + infinite. + Handles both x0 < x1 and x0 > x1. In the latter case, iterates downwards + (make sure to set inc < 0.) + + >>> from scipy.stats._distn_infrastructure import _iter_chunked + >>> [x for x in _iter_chunked(2, 5, inc=2)] + [array([2, 4])] + >>> [x for x in _iter_chunked(2, 11, inc=2)] + [array([2, 4, 6, 8]), array([10])] + >>> [x for x in _iter_chunked(2, -5, inc=-2)] + [array([ 2, 0, -2, -4])] + >>> [x for x in _iter_chunked(2, -9, inc=-2)] + [array([ 2, 0, -2, -4]), array([-6, -8])] + + """ + if inc == 0: + raise ValueError('Cannot increment by zero.') + if chunksize <= 0: + raise ValueError(f'Chunk size must be positive; got {chunksize}.') + + s = 1 if inc > 0 else -1 + stepsize = abs(chunksize * inc) + + x = np.copy(x0) + while (x - x1) * inc < 0: + delta = min(stepsize, abs(x - x1)) + step = delta * s + supp = np.arange(x, x + step, inc) + x += step + yield supp + + +class rv_sample(rv_discrete): + """A 'sample' discrete distribution defined by the support and values. + + The ctor ignores most of the arguments, only needs the `values` argument. + """ + + def __init__(self, a=0, b=inf, name=None, badvalue=None, + moment_tol=1e-8, values=None, inc=1, longname=None, + shapes=None, seed=None): + + super(rv_discrete, self).__init__(seed) + + if values is None: + raise ValueError("rv_sample.__init__(..., values=None,...)") + + # cf generic freeze + self._ctor_param = dict( + a=a, b=b, name=name, badvalue=badvalue, + moment_tol=moment_tol, values=values, inc=inc, + longname=longname, shapes=shapes, seed=seed) + + if badvalue is None: + badvalue = nan + self.badvalue = badvalue + self.moment_tol = moment_tol + self.inc = inc + self.shapes = shapes + self.vecentropy = self._entropy + + xk, pk = values + + if np.shape(xk) != np.shape(pk): + raise ValueError("xk and pk must have the same shape.") + if np.less(pk, 0.0).any(): + raise ValueError("All elements of pk must be non-negative.") + if not np.allclose(np.sum(pk), 1): + raise ValueError("The sum of provided pk is not 1.") + if not len(set(np.ravel(xk))) == np.size(xk): + raise ValueError("xk may not contain duplicate values.") + + indx = np.argsort(np.ravel(xk)) + self.xk = np.take(np.ravel(xk), indx, 0) + self.pk = np.take(np.ravel(pk), indx, 0) + self.a = self.xk[0] + self.b = self.xk[-1] + + self.qvals = np.cumsum(self.pk, axis=0) + + self.shapes = ' ' # bypass inspection + + self._construct_argparser(meths_to_inspect=[self._pmf], + locscale_in='loc=0', + # scale=1 for discrete RVs + locscale_out='loc, 1') + + self._attach_methods() + + self._construct_docstrings(name, longname) + + def __getstate__(self): + dct = self.__dict__.copy() + + # these methods will be remade in rv_generic.__setstate__, + # which calls rv_generic._attach_methods + attrs = ["_parse_args", "_parse_args_stats", "_parse_args_rvs"] + [dct.pop(attr, None) for attr in attrs] + + return dct + + def _attach_methods(self): + """Attaches dynamically created argparser methods.""" + self._attach_argparser_methods() + + def _get_support(self, *args): + """Return the support of the (unscaled, unshifted) distribution. + + Parameters + ---------- + arg1, arg2, ... : array_like + The shape parameter(s) for the distribution (see docstring of the + instance object for more information). + + Returns + ------- + a, b : numeric (float, or int or +/-np.inf) + end-points of the distribution's support. + """ + return self.a, self.b + + def _pmf(self, x): + return np.select([x == k for k in self.xk], + [np.broadcast_arrays(p, x)[0] for p in self.pk], 0) + + def _cdf(self, x): + xx, xxk = np.broadcast_arrays(x[:, None], self.xk) + indx = np.argmax(xxk > xx, axis=-1) - 1 + return self.qvals[indx] + + def _ppf(self, q): + qq, sqq = np.broadcast_arrays(q[..., None], self.qvals) + indx = argmax(sqq >= qq, axis=-1) + return self.xk[indx] + + def _rvs(self, size=None, random_state=None): + # Need to define it explicitly, otherwise .rvs() with size=None + # fails due to explicit broadcasting in _ppf + U = random_state.uniform(size=size) + if size is None: + U = np.array(U, ndmin=1) + Y = self._ppf(U)[0] + else: + Y = self._ppf(U) + return Y + + def _entropy(self): + return stats.entropy(self.pk) + + def generic_moment(self, n): + n = asarray(n) + return np.sum(self.xk**n[np.newaxis, ...] * self.pk, axis=0) + + def _expect(self, fun, lb, ub, *args, **kwds): + # ignore all args, just do a brute force summation + supp = self.xk[(lb <= self.xk) & (self.xk <= ub)] + vals = fun(supp) + return np.sum(vals) + + +def _check_shape(argshape, size): + """ + This is a utility function used by `_rvs()` in the class geninvgauss_gen. + It compares the tuple argshape to the tuple size. + + Parameters + ---------- + argshape : tuple of integers + Shape of the arguments. + size : tuple of integers or integer + Size argument of rvs(). + + Returns + ------- + The function returns two tuples, scalar_shape and bc. + + scalar_shape : tuple + Shape to which the 1-d array of random variates returned by + _rvs_scalar() is converted when it is copied into the + output array of _rvs(). + + bc : tuple of booleans + bc is a tuple the same length as size. bc[j] is True if the data + associated with that index is generated in one call of _rvs_scalar(). + + """ + scalar_shape = [] + bc = [] + for argdim, sizedim in zip_longest(argshape[::-1], size[::-1], + fillvalue=1): + if sizedim > argdim or (argdim == sizedim == 1): + scalar_shape.append(sizedim) + bc.append(True) + else: + bc.append(False) + return tuple(scalar_shape[::-1]), tuple(bc[::-1]) + + +def get_distribution_names(namespace_pairs, rv_base_class): + """Collect names of statistical distributions and their generators. + + Parameters + ---------- + namespace_pairs : sequence + A snapshot of (name, value) pairs in the namespace of a module. + rv_base_class : class + The base class of random variable generator classes in a module. + + Returns + ------- + distn_names : list of strings + Names of the statistical distributions. + distn_gen_names : list of strings + Names of the generators of the statistical distributions. + Note that these are not simply the names of the statistical + distributions, with a _gen suffix added. + + """ + distn_names = [] + distn_gen_names = [] + for name, value in namespace_pairs: + if name.startswith('_'): + continue + if name.endswith('_gen') and issubclass(value, rv_base_class): + distn_gen_names.append(name) + if isinstance(value, rv_base_class): + distn_names.append(name) + return distn_names, distn_gen_names diff --git a/v12x/pylib/scipy/stats/_distr_params.py b/v12x/pylib/scipy/stats/_distr_params.py new file mode 100644 index 0000000000000000000000000000000000000000..3234bae358fffd18b70a008a00110c6f3e348530 --- /dev/null +++ b/v12x/pylib/scipy/stats/_distr_params.py @@ -0,0 +1,300 @@ +""" +Sane parameters for stats.distributions. +""" +import numpy as np + +distcont = [ + ['alpha', (3.5704770516650459,)], + ['anglit', ()], + ['arcsine', ()], + ['argus', (1.0,)], + ['beta', (2.3098496451481823, 0.62687954300963677)], + ['betaprime', (5, 6)], + ['bradford', (0.29891359763170633,)], + ['burr', (10.5, 4.3)], + ['burr12', (10, 4)], + ['cauchy', ()], + ['chi', (78,)], + ['chi2', (55,)], + ['cosine', ()], + ['crystalball', (2.0, 3.0)], + ['dgamma', (1.1023326088288166,)], + ['dpareto_lognorm', (3, 1.2, 1.5, 2)], + ['dweibull', (2.0685080649914673,)], + ['erlang', (10,)], + ['expon', ()], + ['exponnorm', (1.5,)], + ['exponpow', (2.697119160358469,)], + ['exponweib', (2.8923945291034436, 1.9505288745913174)], + ['f', (29, 18)], + ['fatiguelife', (29,)], # correction numargs = 1 + ['fisk', (3.0857548622253179,)], + ['foldcauchy', (4.7164673455831894,)], + ['foldnorm', (1.9521253373555869,)], + ['gamma', (1.9932305483800778,)], + ['gausshyper', (13.763771604130699, 3.1189636648681431, + 2.5145980350183019, 5.1811649903971615)], # veryslow + ['genexpon', (9.1325976465418908, 16.231956600590632, 3.2819552690843983)], + ['genextreme', (-0.1,)], + ['gengamma', (4.4162385429431925, 3.1193091679242761)], + ['gengamma', (4.4162385429431925, -3.1193091679242761)], + ['genhalflogistic', (0.77274727809929322,)], + ['genhyperbolic', (0.5, 1.5, -0.5,)], + ['geninvgauss', (2.3, 1.5)], + ['genlogistic', (0.41192440799679475,)], + ['gennorm', (1.2988442399460265,)], + ['halfgennorm', (0.6748054997000371,)], + ['genpareto', (0.1,)], # use case with finite moments + ['gibrat', ()], + ['gompertz', (0.94743713075105251,)], + ['gumbel_l', ()], + ['gumbel_r', ()], + ['halfcauchy', ()], + ['halflogistic', ()], + ['halfnorm', ()], + ['hypsecant', ()], + ['invgamma', (4.0668996136993067,)], + ['invgauss', (0.14546264555347513,)], + ['invweibull', (10.58,)], + ['irwinhall', (10,)], + ['jf_skew_t', (8, 4)], + ['johnsonsb', (4.3172675099141058, 3.1837781130785063)], + ['johnsonsu', (2.554395574161155, 2.2482281679651965)], + ['kappa4', (0.0, 0.0)], + ['kappa4', (-0.1, 0.1)], + ['kappa4', (0.0, 0.1)], + ['kappa4', (0.1, 0.0)], + ['kappa3', (1.0,)], + ['ksone', (1000,)], # replace 22 by 100 to avoid failing range, ticket 956 + ['kstwo', (10,)], + ['kstwobign', ()], + ['landau', ()], + ['laplace', ()], + ['laplace_asymmetric', (2,)], + ['levy', ()], + ['levy_l', ()], + ['levy_stable', (1.8, -0.5)], + ['loggamma', (0.41411931826052117,)], + ['logistic', ()], + ['loglaplace', (3.2505926592051435,)], + ['lognorm', (0.95368226960575331,)], + ['loguniform', (0.01, 1.25)], + ['lomax', (1.8771398388773268,)], + ['maxwell', ()], + ['mielke', (10.4, 4.6)], + ['moyal', ()], + ['nakagami', (4.9673794866666237,)], + ['ncf', (27, 27, 0.41578441799226107)], + ['nct', (14, 0.24045031331198066)], + ['ncx2', (21, 1.0560465975116415)], + ['norm', ()], + ['norminvgauss', (1.25, 0.5)], + ['pareto', (2.621716532144454,)], + ['pearson3', (0.1,)], + ['pearson3', (-2,)], + ['powerlaw', (1.6591133289905851,)], + ['powerlaw', (0.6591133289905851,)], + ['powerlognorm', (2.1413923530064087, 0.44639540782048337)], + ['powernorm', (4.4453652254590779,)], + ['rayleigh', ()], + ['rdist', (1.6,)], + ['recipinvgauss', (0.63004267809369119,)], + ['reciprocal', (0.01, 1.25)], + ['rel_breitwigner', (36.545206797050334, )], + ['rice', (0.7749725210111873,)], + ['semicircular', ()], + ['skewcauchy', (0.5,)], + ['skewnorm', (4.0,)], + ['studentized_range', (3.0, 10.0)], + ['t', (2.7433514990818093,)], + ['trapezoid', (0.2, 0.8)], + ['triang', (0.15785029824528218,)], + ['truncexpon', (4.6907725456810478,)], + ['truncnorm', (-1.0978730080013919, 2.7306754109031979)], + ['truncnorm', (0.1, 2.)], + ['truncpareto', (1.8, 5.3)], + ['truncpareto', (2, 5)], + ['truncpareto', (-2, 5)], + ['truncweibull_min', (2.5, 0.25, 1.75)], + ['tukeylambda', (3.1321477856738267,)], + ['uniform', ()], + ['vonmises', (3.9939042581071398,)], + ['vonmises_line', (3.9939042581071398,)], + ['wald', ()], + ['weibull_max', (2.8687961709100187,)], + ['weibull_min', (1.7866166930421596,)], + ['wrapcauchy', (0.031071279018614728,)] +] + + +distdiscrete = [ + ['bernoulli',(0.3,)], + ['betabinom', (5, 2.3, 0.63)], + ['betanbinom', (5, 9.3, 1)], + ['binom', (5, 0.4)], + ['boltzmann',(1.4, 19)], + ['dlaplace', (0.8,)], # 0.5 + ['geom', (0.5,)], + ['hypergeom',(30, 12, 6)], + ['hypergeom',(21,3,12)], # numpy.random (3,18,12) numpy ticket:921 + ['hypergeom',(21,18,11)], # numpy.random (18,3,11) numpy ticket:921 + ['nchypergeom_fisher', (140, 80, 60, 0.5)], + ['nchypergeom_wallenius', (140, 80, 60, 0.5)], + ['logser', (0.6,)], # re-enabled, numpy ticket:921 + ['nbinom', (0.4, 0.4)], # from tickets: 583 + ['nbinom', (5, 0.5)], + ['planck', (0.51,)], # 4.1 + ['poisson', (0.6,)], + ['poisson_binom', ([0.1, 0.6, 0.7, 0.8],)], + ['randint', (7, 31)], + ['skellam', (15, 8)], + ['zipf', (6.6,)], + ['zipfian', (0.75, 15)], + ['zipfian', (1.25, 10)], + ['yulesimon', (11.0,)], + ['nhypergeom', (20, 7, 1)] +] + + +invdistdiscrete = [ + # In each of the following, at least one shape parameter is invalid + ['hypergeom', (3, 3, 4)], + ['nhypergeom', (5, 2, 8)], + ['nchypergeom_fisher', (3, 3, 4, 1)], + ['nchypergeom_wallenius', (3, 3, 4, 1)], + ['bernoulli', (1.5, )], + ['binom', (10, 1.5)], + ['betabinom', (10, -0.4, -0.5)], + ['betanbinom', (10, -0.4, -0.5)], + ['boltzmann', (-1, 4)], + ['dlaplace', (-0.5, )], + ['geom', (1.5, )], + ['logser', (1.5, )], + ['nbinom', (10, 1.5)], + ['planck', (-0.5, )], + ['poisson', (-0.5, )], + ['poisson_binom', ([-1, 2, 0.5],)], + ['randint', (5, 2)], + ['skellam', (-5, -2)], + ['zipf', (-2, )], + ['yulesimon', (-2, )], + ['zipfian', (-0.75, 15)] +] + + +invdistcont = [ + # In each of the following, at least one shape parameter is invalid + ['alpha', (-1, )], + ['anglit', ()], + ['arcsine', ()], + ['argus', (-1, )], + ['beta', (-2, 2)], + ['betaprime', (-2, 2)], + ['bradford', (-1, )], + ['burr', (-1, 1)], + ['burr12', (-1, 1)], + ['cauchy', ()], + ['chi', (-1, )], + ['chi2', (-1, )], + ['cosine', ()], + ['crystalball', (-1, 2)], + ['dgamma', (-1, )], + ['dpareto_lognorm', (3, -1.2, 1.5, 2)], + ['dweibull', (-1, )], + ['erlang', (-1, )], + ['expon', ()], + ['exponnorm', (-1, )], + ['exponweib', (1, -1)], + ['exponpow', (-1, )], + ['f', (10, -10)], + ['fatiguelife', (-1, )], + ['fisk', (-1, )], + ['foldcauchy', (-1, )], + ['foldnorm', (-1, )], + ['genlogistic', (-1, )], + ['gennorm', (-1, )], + ['genpareto', (np.inf, )], + ['genexpon', (1, 2, -3)], + ['genextreme', (np.inf, )], + ['genhyperbolic', (0.5, -0.5, -1.5,)], + ['gausshyper', (1, 2, 3, -4)], + ['gamma', (-1, )], + ['gengamma', (-1, 0)], + ['genhalflogistic', (-1, )], + ['geninvgauss', (1, 0)], + ['gibrat', ()], + ['gompertz', (-1, )], + ['gumbel_r', ()], + ['gumbel_l', ()], + ['halfcauchy', ()], + ['halflogistic', ()], + ['halfnorm', ()], + ['halfgennorm', (-1, )], + ['hypsecant', ()], + ['invgamma', (-1, )], + ['invgauss', (-1, )], + ['invweibull', (-1, )], + ['irwinhall', (-1,)], + ['irwinhall', (0,)], + ['irwinhall', (2.5,)], + ['jf_skew_t', (-1, 0)], + ['johnsonsb', (1, -2)], + ['johnsonsu', (1, -2)], + ['kappa4', (np.nan, 0)], + ['kappa3', (-1, )], + ['ksone', (-1, )], + ['kstwo', (-1, )], + ['kstwobign', ()], + ['landau', ()], + ['laplace', ()], + ['laplace_asymmetric', (-1, )], + ['levy', ()], + ['levy_l', ()], + ['levy_stable', (-1, 1)], + ['logistic', ()], + ['loggamma', (-1, )], + ['loglaplace', (-1, )], + ['lognorm', (-1, )], + ['loguniform', (10, 5)], + ['lomax', (-1, )], + ['maxwell', ()], + ['mielke', (1, -2)], + ['moyal', ()], + ['nakagami', (-1, )], + ['ncx2', (-1, 2)], + ['ncf', (10, 20, -1)], + ['nct', (-1, 2)], + ['norm', ()], + ['norminvgauss', (5, -10)], + ['pareto', (-1, )], + ['pearson3', (np.nan, )], + ['powerlaw', (-1, )], + ['powerlognorm', (1, -2)], + ['powernorm', (-1, )], + ['rdist', (-1, )], + ['rayleigh', ()], + ['rice', (-1, )], + ['recipinvgauss', (-1, )], + ['semicircular', ()], + ['skewnorm', (np.inf, )], + ['studentized_range', (-1, 1)], + ['rel_breitwigner', (-2, )], + ['t', (-1, )], + ['trapezoid', (0, 2)], + ['triang', (2, )], + ['truncexpon', (-1, )], + ['truncnorm', (10, 5)], + ['truncpareto', (-1, -5)], + ['truncpareto', (1.8, .5)], + ['truncweibull_min', (-2.5, 0.25, 1.75)], + ['tukeylambda', (np.nan, )], + ['uniform', ()], + ['vonmises', (-1, )], + ['vonmises_line', (-1, )], + ['wald', ()], + ['weibull_min', (-1, )], + ['weibull_max', (-1, )], + ['wrapcauchy', (2, )], + ['reciprocal', (15, 10)], + ['skewcauchy', (2, )] +] diff --git a/v12x/pylib/scipy/stats/_distribution_infrastructure.py b/v12x/pylib/scipy/stats/_distribution_infrastructure.py new file mode 100644 index 0000000000000000000000000000000000000000..ab05372648e813dc962c4ddcd3fc5171dcab07c0 --- /dev/null +++ b/v12x/pylib/scipy/stats/_distribution_infrastructure.py @@ -0,0 +1,5764 @@ +import functools +from abc import ABC, abstractmethod +from functools import cached_property +from types import GenericAlias +import inspect +import math + +import numpy as np +from numpy import inf + +from scipy._lib._array_api import xp_capabilities, xp_promote +from scipy._lib._util import _rng_spawn +from scipy._lib._docscrape import ClassDoc, NumpyDocString +from scipy._external import array_api_extra as xpx +from scipy import special, stats +from scipy.special._ufuncs import _log1mexp +from scipy.integrate import tanhsinh as _tanhsinh, nsum +from scipy.optimize import elementwise +from scipy.stats._probability_distribution import _ProbabilityDistribution +from scipy.stats import qmc + +# in case we need to distinguish between None and not specified +# Typically this is used to determine whether the tolerance has been set by the +# user and make a decision about which method to use to evaluate a distribution +# function. Sometimes, the logic does not consider the value of the tolerance, +# only whether this has been defined or not. This is not intended to be the +# best possible logic; the intent is to establish the structure, which can +# be refined in follow-up work. +# See https://github.com/scipy/scipy/pull/21050#discussion_r1714195433. +_null = object() +def _isnull(x): + return type(x) is object or x is None + +__all__ = ['make_distribution', 'Mixture', 'order_statistic', + 'truncate', 'abs', 'exp', 'log'] + +# Could add other policies for broadcasting and edge/out-of-bounds case handling +# For instance, when edge case handling is known not to be needed, it's much +# faster to turn it off, but it might still be nice to have array conversion +# and shaping done so the user doesn't need to be so careful. +_SKIP_ALL = "skip_all" +# Other cache policies would be useful, too. +_NO_CACHE = "no_cache" + +# TODO: +# Test sample dtypes +# Add dtype kwarg (especially for distributions with no parameters) +# Rewrite `test_continuous` without Hypothesis, more modular tests that are easier to +# run individually. +# Add distribution-specific tests - just enough to confirm that distribution is the +# one it is supposed to be; maybe check finite vs undefined vs infinite moments. +# `_tanhsinh` special case when there are no abscissae between the limits +# example: cdf of uniform betweeen 1.0 and np.nextafter(1.0, np.inf) +# implement symmetric distribution (take advantage of symmetry) +# implement composite distribution (for using different implementations depending on +# shape parameters - e.g. t distribution with df = np.inf delegates to normal) +# Add array API support +# Investigate use `median` information throughout, e.g. to improve integration of +# CDF or to provide initial bracket for ICDF. (Only if `_median_formula` is +# provided.) +# Document user tips for faster execution: +# - pass NumPy arrays +# - pass inputs of floating point type (not integers) +# - prefer NumPy scalars or 0d arrays over other size 1 arrays +# - pass no invalid parameters and disable invalid parameter checks (validation_policy) +# - provide a Generator if you're going to do sampling +# Add options for drawing parameters: log-spacing +# Accuracy benchmark suite. Ideally, when all code is array API compatible and mparray +# provides an array-API compatible arbitrary precision array, we can just compare +# results when using NumPy backend against mparray backend. +# Should caches be attributes so we can more easily ensure that they are not +# modified when caching is turned off? +# Make ShiftedScaledDistribution more efficient - only process underlying +# distribution parameters as necessary. +# Reconsider `all_inclusive` - see comment in `contains` method of `_interval` +# `validation_policy` needs testing +# `tol` does not offer very fine-grained control; consider improving +# report accuracy estimates? How? +# reconsider distribution mutability + +# Originally, I planned to filter out invalid distribution parameters; +# distribution implementation functions would always work with "compressed", +# 1D arrays containing only valid distribution parameters. There are two +# problems with this: +# - This essentially requires copying all arrays, even if there is only a +# single invalid parameter combination. This is expensive. Then, to output +# the original size data to the user, we need to "decompress" the arrays +# and fill in the NaNs, so more copying. Unless we branch the code when +# there are no invalid data, these copies happen even in the normal case, +# where there are no invalid parameter combinations. We should not incur +# all this overhead in the normal case. +# - For methods that accept arguments other than distribution parameters, the +# user will pass in arrays that are broadcastable with the original arrays, +# not the compressed arrays. This means that this same sort of invalid +# value detection needs to be repeated every time one of these methods is +# called. +# The much simpler solution is to keep the data uncompressed but to replace +# the invalid parameters and arguments with NaNs (and only if some are +# invalid). With this approach, the copying happens only if/when it is +# needed. Most functions involved in stats distribution calculations don't +# mind NaNs; they just return NaN. The behavior "If x_i is NaN, the result +# is NaN" is explicit in the array API. So this should be fine. +# +# Currently, I am still leaving the parameters and function arguments +# in their broadcasted shapes rather than, say, raveling. The intent +# is to avoid back and forth reshaping. If authors of distributions have +# trouble dealing with N-D arrays, we can reconsider this. +# +# Another important decision is that the *private* methods must accept +# the distribution parameters as inputs rather than relying on these +# cached properties directly (although the public methods typically pass +# the cached values to the private methods). This is because the elementwise +# algorithms for quadrature, differentiation, root-finding, and minimization +# prefer that the input functions are strictly elementwise in the sense +# that the value output for a given input element does not depend on the +# shape of the input or that element's location within the input array. +# When the computation has converged for an element, it is removed from +# the computation entirely. As a result, the shape of the arrays passed to +# the function will almost never be broadcastable with the shape of the +# cached parameter arrays. +# +# I've sprinkled in some optimizations for scalars and same-shape/type arrays +# throughout. The biggest time sinks before were: +# - broadcast_arrays +# - result_dtype +# - is_subdtype +# It is much faster to check whether these are necessary than to do them. + + +class _Domain(ABC): + r""" Representation of the applicable domain of a parameter or variable. + + A `_Domain` object is responsible for storing information about the + domain of a parameter or variable, determining whether a value is within + the domain (`contains`), and providing a text/mathematical representation + of itself (`__str__`). Because the domain of a parameter/variable can have + a complicated relationship with other parameters and variables of a + distribution, `_Domain` itself does not try to represent all possibilities; + in fact, it has no implementation and is meant for subclassing. + + Attributes + ---------- + symbols : dict + A map from special numerical values to symbols for use in `__str__` + + Methods + ------- + contains(x) + Determine whether the argument is contained within the domain (True) + or not (False). Used for input validation. + get_numerical_endpoints() + Gets the numerical values of the domain endpoints, which may have been + defined symbolically or through a callable. + __str__() + Returns a text representation of the domain (e.g. ``[0, b)``). + Used for generating documentation. + + """ + symbols = {np.inf: r"\infty", -np.inf: r"-\infty", np.pi: r"\pi", -np.pi: r"-\pi"} + + # generic type compatibility with scipy-stubs + __class_getitem__: classmethod = classmethod(GenericAlias) + + @abstractmethod + def contains(self, x): + raise NotImplementedError() + + @abstractmethod + def draw(self, n): + raise NotImplementedError() + + @abstractmethod + def get_numerical_endpoints(self, x): + raise NotImplementedError() + + @abstractmethod + def __str__(self): + raise NotImplementedError() + + +class _Interval(_Domain): + r""" Representation of an interval defined by two endpoints. + + Each endpoint may be a finite scalar, positive or negative infinity, or + be given by a single parameter. The domain may include the endpoints or + not. + + This class still does not provide an implementation of the __str__ method, + so it is meant for subclassing (e.g. a subclass for domains on the real + line). + + Attributes + ---------- + symbols : dict + Inherited. A map from special values to symbols for use in `__str__`. + endpoints : 2-tuple of float(s) and/or str(s) and/or callable(s). + A tuple with two values. Each may be either a float (the numerical + value of the endpoints of the domain), a string (the name of the + parameters that will define the endpoint), or a callable taking the + parameters used to define the endpoints of the domain as keyword only + arguments and returning a numerical value for the endpoint. + inclusive : 2-tuple of bools + A tuple with two boolean values; each indicates whether the + corresponding endpoint is included within the domain or not. + + Methods + ------- + define_parameters(*parameters) + Records any parameters used to define the endpoints of the domain + get_numerical_endpoints(parameter_values) + Gets the numerical values of the domain endpoints, which may have been + defined symbolically or through a callable. + contains(item, parameter_values) + Determines whether the argument is contained within the domain + draw(size, rng, proportions, parameter_values) + Draws random values based on the domain. + + """ + def __init__(self, endpoints=(-inf, inf), inclusive=(False, False)): + self.symbols = super().symbols.copy() + a, b = endpoints + self.endpoints = np.asarray(a)[()], np.asarray(b)[()] + self.inclusive = inclusive + + def define_parameters(self, *parameters): + r""" Records any parameters used to define the endpoints of the domain. + + Adds the keyword name of each parameter and its text representation + to the `symbols` attribute as key:value pairs. + For instance, a parameter may be passed into to a distribution's + initializer using the keyword `log_a`, and the corresponding + string representation may be '\log(a)'. To form the text + representation of the domain for use in documentation, the + _Domain object needs to map from the keyword name used in the code + to the string representation. + + Returns None, but updates the `symbols` attribute. + + Parameters + ---------- + *parameters : _Parameter objects + Parameters that may define the endpoints of the domain. + + """ + new_symbols = {param.name: param.symbol for param in parameters} + self.symbols.update(new_symbols) + + def get_numerical_endpoints(self, parameter_values): + r""" Get the numerical values of the domain endpoints. + + Domain endpoints may be defined symbolically or through a callable. + This returns numerical values of the endpoints given numerical values for + any variables. + + Parameters + ---------- + parameter_values : dict + A dictionary that maps between string variable names and numerical + values of parameters, which may define the endpoints. + + Returns + ------- + a, b : ndarray + Numerical values of the endpoints + + """ + a, b = self.endpoints + # If `a` (`b`) is a string - the name of the parameter that defines + # the endpoint of the domain - then corresponding numerical values + # will be found in the `parameter_values` dictionary. + # If a callable, it will be executed with `parameter_values` passed as + # keyword arguments, and it will return the numerical values. + # Otherwise, it is itself the array of numerical values of the endpoint. + try: + if callable(a): + a = a(**parameter_values) + else: + a = np.asarray(parameter_values.get(a, a)) + if callable(b): + b = b(**parameter_values) + else: + b = np.asarray(parameter_values.get(b, b)) + except TypeError as e: + message = ("The endpoints of the distribution are defined by " + "parameters, but their values were not provided. When " + f"using a private method of {self.__class__}, pass " + "all required distribution parameters as keyword " + "arguments.") + raise TypeError(message) from e + # Floating point types are used for even integer parameters. + # Convert to float here to ensure consistency throughout framework. + a, b = xp_promote(a, b, force_floating=True, xp=np) + return a, b + + def contains(self, item, parameter_values=None): + r"""Determine whether the argument is contained within the domain. + + Parameters + ---------- + item : ndarray + The argument + parameter_values : dict + A dictionary that maps between string variable names and numerical + values of parameters, which may define the endpoints. + + Returns + ------- + out : bool + True if `item` is within the domain; False otherwise. + + """ + parameter_values = parameter_values or {} + # if self.all_inclusive: + # # Returning a 0d value here makes things much faster. + # # I'm not sure if it's safe, though. If it causes a bug someday, + # # I guess it wasn't. + # # Even if there is no bug because of the shape, it is incorrect for + # # `contains` to return True when there are invalid (e.g. NaN) + # # parameters. + # return np.asarray(True) + + a, b = self.get_numerical_endpoints(parameter_values) + left_inclusive, right_inclusive = self.inclusive + + in_left = item >= a if left_inclusive else item > a + in_right = item <= b if right_inclusive else item < b + return in_left & in_right + + def draw(self, n, type_, min, max, squeezed_base_shape, rng=None): + r""" Draw random values from the domain. + + Parameters + ---------- + n : int + The number of values to be drawn from the domain. + type_ : str + A string indicating whether the values are + + - strictly within the domain ('in'), + - at one of the two endpoints ('on'), + - strictly outside the domain ('out'), or + - NaN ('nan'). + min, max : ndarray + The endpoints of the domain. + squeezed_based_shape : tuple of ints + See _RealParameter.draw. + rng : np.Generator + The Generator used for drawing random values. + + """ + rng = np.random.default_rng(rng) + + def ints(*args, **kwargs): return rng.integers(*args, **kwargs, endpoint=True) + uniform = rng.uniform if isinstance(self, _RealInterval) else ints + + # get copies of min and max with no nans so that uniform doesn't fail + min_nn, max_nn = min.copy(), max.copy() + i = np.isnan(min_nn) | np.isnan(max_nn) + min_nn[i] = 0 + max_nn[i] = 1 + + shape = (n,) + squeezed_base_shape + + if type_ == 'in': + z = uniform(min_nn, max_nn, size=shape) + + elif type_ == 'on': + z_on_shape = shape + z = np.ones(z_on_shape) + i = rng.random(size=n) < 0.5 + z[i] = min + z[~i] = max + + elif type_ == 'out': + z = min_nn - uniform(1, 5, size=shape) # 1, 5 is arbitary; we just want + zr = max_nn + uniform(1, 5, size=shape) # some numbers outside domain + i = rng.random(size=n) < 0.5 + z[i] = zr[i] + + elif type_ == 'nan': + z = np.full(shape, np.nan) + + return z + + +class _RealInterval(_Interval): + r""" Represents a simply-connected subset of the real line; i.e., an interval + + Completes the implementation of the `_Interval` class for intervals + on the real line. + + Methods + ------- + define_parameters(*parameters) + (Inherited) Records any parameters used to define the endpoints of the + domain. + get_numerical_endpoints(parameter_values) + (Inherited) Gets the numerical values of the domain endpoints, which + may have been defined symbolically. + contains(item, parameter_values) + (Inherited) Determines whether the argument is contained within the + domain + __str__() + Returns a string representation of the domain, e.g. "[a, b)". + """ + + def __str__(self): + a, b = self.endpoints + a, b = self._get_endpoint_str(a, "f1"), self._get_endpoint_str(b, "f2") + left_inclusive, right_inclusive = self.inclusive + left = "[" if left_inclusive else "(" + right = "]" if right_inclusive else ")" + + return f"{left}{a}, {b}{right}" + + def _get_endpoint_str(self, endpoint, funcname): + if callable(endpoint): + if endpoint.__doc__ is not None: + return endpoint.__doc__ + params = inspect.signature(endpoint).parameters.values() + params = [ + p.name for p in params if p.kind == inspect.Parameter.KEYWORD_ONLY + ] + return f"{funcname}({','.join(params)})" + return self.symbols.get(endpoint, f"{endpoint}") + + +class _IntegerInterval(_Interval): + r""" Represents an interval of integers + + Completes the implementation of the `_Interval` class for simple + domains on the integers. + + Methods + ------- + define_parameters(*parameters) + (Inherited) Records any parameters used to define the endpoints of the + domain. + get_numerical_endpoints(parameter_values) + (Inherited) Gets the numerical values of the domain endpoints, which + may have been defined symbolically. + contains(item, parameter_values) + (Overridden) Determines whether the argument is contained within the + domain + draw(n, type_, min, max, squeezed_base_shape, rng=None) + (Inherited) Draws random values based on the domain. + __str__() + Returns a string representation of the domain, e.g. "{a, a+1, ..., b-1, b}". + + """ + def contains(self, item, parameter_values=None): + super_contains = super().contains(item, parameter_values) + integral = (item == np.round(item)) + return super_contains & integral + + def __str__(self): + a, b = self.endpoints + a = self.symbols.get(a, a) + b = self.symbols.get(b, b) + + a_str, b_str = isinstance(a, str), isinstance(b, str) + a_inf = a == r"-\infty" if a_str else np.isinf(a) + b_inf = b == r"\infty" if b_str else np.isinf(b) + + # This doesn't work well for cases where ``a`` is floating point + # number large enough that ``nextafter(a, inf) > a + 1``, and + # similarly for ``b`` and nextafter(b, -inf). There may not be any + # distributions fit for SciPy where we would actually need to handle these + # cases though. + ap1 = f"{a} + 1" if a_str else f"{a + 1}" + bm1 = f"{b} - 1" if b_str else f"{b - 1}" + + if not a_str and not b_str: + gap = b - a + if gap == 3: + return f"\\{{{a}, {ap1}, {bm1}, {b}\\}}" + if gap == 2: + return f"\\{{{a}, {ap1}, {b}\\}}" + if gap == 1: + return f"\\{{{a}, {b}\\}}" + if gap == 0: + return f"\\{{{a}\\}}" + + if not a_inf and b_inf: + ap2 = f"{a} + 2" if a_str else f"{a + 2}" + return f"\\{{{a}, {ap1}, {ap2}, ...\\}}" + if a_inf and not b_inf: + bm2 = f"{b} - 2" if b_str else f"{b - 2}" + return f"\\{{{b}, {bm1}, {bm2}, ...\\}}" + if a_inf and b_inf: + return "\\{..., -2, -1, 0, 1, 2, ...\\}" + + return f"\\{{{a}, {ap1}, ..., {bm1}, {b}\\}}" + + +class _Parameter(ABC): + r""" Representation of a distribution parameter or variable. + + A `_Parameter` object is responsible for storing information about a + parameter or variable, providing input validation/standardization of + values passed for that parameter, providing a text/mathematical + representation of the parameter for the documentation (`__str__`), and + drawing random values of itself for testing and benchmarking. It does + not provide a complete implementation of this functionality and is meant + for subclassing. + + Attributes + ---------- + name : str + The keyword used to pass numerical values of the parameter into the + initializer of the distribution + symbol : str + The text representation of the variable in the documentation. May + include LaTeX. + domain : _Domain + The domain of the parameter for which the distribution is valid. + typical : 2-tuple of floats or strings (consider making a _Domain) + Defines the endpoints of a typical range of values of the parameter. + Used for sampling. + + Methods + ------- + __str__(): + Returns a string description of the variable for use in documentation, + including the keyword used to represent it in code, the symbol used to + represent it mathemtatically, and a description of the valid domain. + draw(size, *, rng, domain, proportions) + Draws random values of the parameter. Proportions of values within + the valid domain, on the endpoints of the domain, outside the domain, + and having value NaN are specified by `proportions`. + validate(x): + Validates and standardizes the argument for use as numerical values + of the parameter. + + """ + + # generic type compatibility with scipy-stubs + __class_getitem__: classmethod = classmethod(GenericAlias) + + def __init__(self, name, *, domain, symbol=None, typical=None): + self.name = name + self.symbol = symbol or name + self.domain = domain + if typical is not None and not isinstance(typical, _Domain): + typical = domain.__class__(typical) + self.typical = typical or domain + + def __str__(self): + r""" String representation of the parameter for use in documentation.""" + return f"`{self.name}` for :math:`{self.symbol} \\in {str(self.domain)}`" + + def draw(self, size=None, *, rng=None, region='domain', proportions=None, + parameter_values=None): + r""" Draw random values of the parameter for use in testing. + + Parameters + ---------- + size : tuple of ints + The shape of the array of valid values to be drawn. + rng : np.Generator + The Generator used for drawing random values. + region : str + The region of the `_Parameter` from which to draw. Default is + "domain" (the *full* domain); alternative is "typical". An + enhancement would give a way to interpolate between the two. + proportions : tuple of numbers + A tuple of four non-negative numbers that indicate the expected + relative proportion of elements that: + + - are strictly within the domain, + - are at one of the two endpoints, + - are strictly outside the domain, and + - are NaN, + + respectively. Default is (1, 0, 0, 0). The number of elements in + each category is drawn from the multinomial distribution with + `np.prod(size)` as the number of trials and `proportions` as the + event probabilities. The values in `proportions` are automatically + normalized to sum to 1. + parameter_values : dict + Map between the names of parameters (that define the endpoints of + `typical`) and numerical values (arrays). + + """ + parameter_values = parameter_values or {} + domain = self.domain + proportions = (1, 0, 0, 0) if proportions is None else proportions + + pvals = proportions / np.sum(proportions) + + a, b = domain.get_numerical_endpoints(parameter_values) + a, b = np.broadcast_arrays(a, b) + + base_shape = a.shape + extended_shape = np.broadcast_shapes(size, base_shape) + n_extended = np.prod(extended_shape) + n_base = np.prod(base_shape) + n = int(n_extended / n_base) if n_extended else 0 + + rng = np.random.default_rng(rng) + n_in, n_on, n_out, n_nan = rng.multinomial(n, pvals) + + # `min` and `max` can have singleton dimensions that correspond with + # non-singleton dimensions in `size`. We need to be careful to avoid + # shuffling results (e.g. a value that was generated for the domain + # [min[i], max[i]] ends up at index j). To avoid this: + # - Squeeze the singleton dimensions out of `min`/`max`. Squeezing is + # often not the right thing to do, but here is equivalent to moving + # all the dimensions that are singleton in `min`/`max` (which may be + # non-singleton in the result) to the left. This is what we want. + # - Now all the non-singleton dimensions of the result are on the left. + # Ravel them to a single dimension of length `n`, which is now along + # the 0th axis. + # - Reshape the 0th axis back to the required dimensions, and move + # these axes back to their original places. + base_shape_padded = ((1,)*(len(extended_shape) - len(base_shape)) + + base_shape) + base_singletons = np.where(np.asarray(base_shape_padded)==1)[0] + new_base_singletons = tuple(range(len(base_singletons))) + # Base singleton dimensions are going to get expanded to these lengths + shape_expansion = np.asarray(extended_shape)[base_singletons] + + # assert(np.prod(shape_expansion) == n) # check understanding + # min = np.reshape(min, base_shape_padded) + # max = np.reshape(max, base_shape_padded) + # min = np.moveaxis(min, base_singletons, new_base_singletons) + # max = np.moveaxis(max, base_singletons, new_base_singletons) + # squeezed_base_shape = max.shape[len(base_singletons):] + # assert np.all(min.reshape(squeezed_base_shape) == min.squeeze()) + # assert np.all(max.reshape(squeezed_base_shape) == max.squeeze()) + + # min = np.maximum(a, _fiinfo(a).min/10) if np.any(np.isinf(a)) else a + # max = np.minimum(b, _fiinfo(b).max/10) if np.any(np.isinf(b)) else b + min = np.asarray(a.squeeze()) + max = np.asarray(b.squeeze()) + squeezed_base_shape = max.shape + + if region == 'typical': + typical = self.typical + a, b = typical.get_numerical_endpoints(parameter_values) + a, b = np.broadcast_arrays(a, b) + min_here = np.asarray(a.squeeze()) + max_here = np.asarray(b.squeeze()) + z_in = typical.draw(n_in, 'in', min_here, max_here, squeezed_base_shape, + rng=rng) + else: + z_in = domain.draw(n_in, 'in', min, max, squeezed_base_shape, rng=rng) + z_on = domain.draw(n_on, 'on', min, max, squeezed_base_shape, rng=rng) + z_out = domain.draw(n_out, 'out', min, max, squeezed_base_shape, rng=rng) + z_nan= domain.draw(n_nan, 'nan', min, max, squeezed_base_shape, rng=rng) + + z = np.concatenate((z_in, z_on, z_out, z_nan), axis=0) + z = rng.permuted(z, axis=0) + + z = np.reshape(z, tuple(shape_expansion) + squeezed_base_shape) + z = np.moveaxis(z, new_base_singletons, base_singletons) + return z + + @abstractmethod + def validate(self, arr): + raise NotImplementedError() + + +class _RealParameter(_Parameter): + r""" Represents a real-valued parameter. + + Implements the remaining methods of _Parameter for real parameters. + All attributes are inherited. + + """ + def validate(self, arr, parameter_values): + r""" Input validation/standardization of numerical values of a parameter. + + Checks whether elements of the argument `arr` are reals, ensuring that + the dtype reflects this. Also produces a logical array that indicates + which elements meet the requirements. + + Parameters + ---------- + arr : ndarray + The argument array to be validated and standardized. + parameter_values : dict + Map of parameter names to parameter value arrays. + + Returns + ------- + arr : ndarray + The argument array that has been validated and standardized + (converted to an appropriate dtype, if necessary). + dtype : NumPy dtype + The appropriate floating point dtype of the parameter. + valid : bool ndarray + Logical array indicating which elements are valid (True) and + which are not (False). The arrays of all distribution parameters + will be broadcasted, and elements for which any parameter value + does not meet the requirements will be replaced with NaN. + + """ + arr = np.asarray(arr) + + valid_dtype = None + # minor optimization - fast track the most common types to avoid + # overhead of np.issubdtype. Checking for `in {...}` doesn't work : / + if arr.dtype == np.float64 or arr.dtype == np.float32: + pass + elif arr.dtype == np.int32 or arr.dtype == np.int64: + arr = np.asarray(arr, dtype=np.float64) + elif np.issubdtype(arr.dtype, np.floating): + pass + elif np.issubdtype(arr.dtype, np.integer): + arr = np.asarray(arr, dtype=np.float64) + else: + message = f"Parameter `{self.name}` must be of real dtype." + raise TypeError(message) + + valid = self.domain.contains(arr, parameter_values) + valid = valid & valid_dtype if valid_dtype is not None else valid + + return arr[()], arr.dtype, valid + + +class _Parameterization: + r""" Represents a parameterization of a distribution. + + Distributions can have multiple parameterizations. A `_Parameterization` + object is responsible for recording the parameters used by the + parameterization, checking whether keyword arguments passed to the + distribution match the parameterization, and performing input validation + of the numerical values of these parameters. + + Attributes + ---------- + parameters : dict + String names (of keyword arguments) and the corresponding _Parameters. + + Methods + ------- + __len__() + Returns the number of parameters in the parameterization. + __str__() + Returns a string representation of the parameterization. + copy + Returns a copy of the parameterization. This is needed for transformed + distributions that add parameters to the parameterization. + matches(parameters) + Checks whether the keyword arguments match the parameterization. + validation(parameter_values) + Input validation / standardization of parameterization. Validates the + numerical values of all parameters. + draw(sizes, rng, proportions) + Draw random values of all parameters of the parameterization for use + in testing. + """ + def __init__(self, *parameters): + self.parameters = {param.name: param for param in parameters} + + def __len__(self): + return len(self.parameters) + + def copy(self): + return _Parameterization(*self.parameters.values()) + + def matches(self, parameters): + r""" Checks whether the keyword arguments match the parameterization. + + Parameters + ---------- + parameters : set + Set of names of parameters passed into the distribution as keyword + arguments. + + Returns + ------- + out : bool + True if the keyword arguments names match the names of the + parameters of this parameterization. + """ + return parameters == set(self.parameters.keys()) + + def validation(self, parameter_values): + r""" Input validation / standardization of parameterization. + + Parameters + ---------- + parameter_values : dict + The keyword arguments passed as parameter values to the + distribution. + + Returns + ------- + all_valid : ndarray + Logical array indicating the elements of the broadcasted arrays + for which all parameter values are valid. + dtype : dtype + The common dtype of the parameter arrays. This will determine + the dtype of the output of distribution methods. + """ + all_valid = True + dtypes = set() # avoid np.result_type if there's only one type + for name, arr in parameter_values.items(): + parameter = self.parameters[name] + arr, dtype, valid = parameter.validate(arr, parameter_values) + dtypes.add(dtype) + all_valid = all_valid & valid + parameter_values[name] = arr + dtype = arr.dtype if len(dtypes)==1 else np.result_type(*list(dtypes)) + + return all_valid, dtype + + def __str__(self): + r"""Returns a string representation of the parameterization.""" + messages = [str(param) for name, param in self.parameters.items()] + return ", ".join(messages) + + def draw(self, sizes=None, rng=None, proportions=None, region='domain'): + r"""Draw random values of all parameters for use in testing. + + Parameters + ---------- + sizes : iterable of shape tuples + The size of the array to be generated for each parameter in the + parameterization. Note that the order of sizes is arbitary; the + size of the array generated for a specific parameter is not + controlled individually as written. + rng : NumPy Generator + The generator used to draw random values. + proportions : tuple + A tuple of four non-negative numbers that indicate the expected + relative proportion of elements that are within the parameter's + domain, are on the boundary of the parameter's domain, are outside + the parameter's domain, and have value NaN. For more information, + see the `draw` method of the _Parameter subclasses. + domain : str + The domain of the `_Parameter` from which to draw. Default is + "domain" (the *full* domain); alternative is "typical". + + Returns + ------- + parameter_values : dict (string: array) + A dictionary of parameter name/value pairs. + """ + # ENH: be smart about the order. The domains of some parameters + # depend on others. If the relationshp is simple (e.g. a < b < c), + # we can draw values in order a, b, c. + parameter_values = {} + + if sizes is None or not len(sizes) or not np.iterable(sizes[0]): + sizes = [sizes]*len(self.parameters) + + for size, param in zip(sizes, self.parameters.values()): + parameter_values[param.name] = param.draw( + size, rng=rng, proportions=proportions, + parameter_values=parameter_values, + region=region + ) + + return parameter_values + + +def _set_invalid_nan(f): + # Wrapper for input / output validation and standardization of distribution + # functions that accept either the quantile or percentile as an argument: + # logpdf, pdf + # logpmf, pmf + # logcdf, cdf + # logccdf, ccdf + # ilogcdf, icdf + # ilogccdf, iccdf + # Arguments that are outside the required range are replaced by NaN before + # passing them into the underlying function. The corresponding outputs + # are replaced by the appropriate value before being returned to the user. + # For example, when the argument of `cdf` exceeds the right end of the + # distribution's support, the wrapper replaces the argument with NaN, + # ignores the output of the underlying function, and returns 1.0. It also + # ensures that output is of the appropriate shape and dtype. + + endpoints = {'icdf': (0, 1), 'iccdf': (0, 1), + 'ilogcdf': (-np.inf, 0), 'ilogccdf': (-np.inf, 0)} + replacements = {'logpdf': (-inf, -inf), 'pdf': (0, 0), + 'logpmf': (-inf, -inf), 'pmf': (0, 0), + '_logcdf1': (-inf, 0), '_logccdf1': (0, -inf), + '_cdf1': (0, 1), '_ccdf1': (1, 0)} + replace_strict = {'pdf', 'logpdf', 'pmf', 'logpmf'} + replace_exact = {'icdf', 'iccdf', 'ilogcdf', 'ilogccdf'} + clip = {'_cdf1', '_ccdf1'} + clip_log = {'_logcdf1', '_logccdf1'} + # relevant to discrete distributions only + replace_non_integral = {'pmf', 'logpmf', 'pdf', 'logpdf'} + + @functools.wraps(f) + def filtered(self, x, *args, **kwargs): + if self.validation_policy == _SKIP_ALL: + return f(self, x, *args, **kwargs) + + method_name = f.__name__ + x = np.asarray(x) + dtype = self._dtype + shape = self._shape + discrete = isinstance(self, DiscreteDistribution) + keep_low_endpoint = discrete and method_name in {'_cdf1', '_logcdf1', + '_ccdf1', '_logccdf1'} + + # Ensure that argument is at least as precise as distribution + # parameters, which are already at least floats. This will avoid issues + # with raising integers to negative integer powers and failure to replace + # invalid integers with NaNs. + if x.dtype != dtype: + dtype = np.result_type(x.dtype, dtype) + x = np.asarray(x, dtype=dtype) + + # Broadcasting is slow. Do it only if necessary. + if not x.shape == shape: + try: + shape = np.broadcast_shapes(x.shape, shape) + x = np.broadcast_to(x, shape) + # Should we broadcast the distribution parameters to this shape, too? + except ValueError as e: + message = ( + f"The argument provided to `{self.__class__.__name__}" + f".{method_name}` cannot be be broadcast to the same " + "shape as the distribution parameters.") + raise ValueError(message) from e + + low, high = endpoints.get(method_name, self.support()) + + # Check for arguments outside of domain. They'll be replaced with NaNs, + # and the result will be set to the appropriate value. + left_inc, right_inc = self._variable.domain.inclusive + mask_low = (x < low if (method_name in replace_strict and left_inc) + or keep_low_endpoint else x <= low) + mask_high = (x > high if (method_name in replace_strict and right_inc) + else x >= high) + mask_invalid = (mask_low | mask_high) + any_invalid = (mask_invalid if mask_invalid.shape == () + else np.any(mask_invalid)) + + # Check for arguments at domain endpoints, whether they + # are part of the domain or not. + any_endpoint = False + if method_name in replace_exact: + mask_low_endpoint = (x == low) + mask_high_endpoint = (x == high) + mask_endpoint = (mask_low_endpoint | mask_high_endpoint) + any_endpoint = (mask_endpoint if mask_endpoint.shape == () + else np.any(mask_endpoint)) + + # Check for non-integral arguments to PMF method + # or PDF of a discrete distribution. + any_non_integral = False + if discrete and method_name in replace_non_integral: + mask_non_integral = (x != np.floor(x)) + any_non_integral = (mask_non_integral if mask_non_integral.shape == () + else np.any(mask_non_integral)) + + # Set out-of-domain arguments to NaN. The result will be set to the + # appropriate value later. + if any_invalid: + x = np.array(x, dtype=dtype, copy=True) + x[mask_invalid] = np.nan + + res = np.asarray(f(self, x, *args, **kwargs)) + + # Ensure that the result is the correct dtype and shape, + # copying (only once) if necessary. + res_needs_copy = False + if res.dtype != dtype: + dtype = np.result_type(dtype, self._dtype) + res_needs_copy = True + + if res.shape != shape: # faster to check first + res = np.broadcast_to(res, self._shape) + res_needs_copy = (res_needs_copy or any_invalid + or any_endpoint or any_non_integral) + + if res_needs_copy: + res = np.array(res, dtype=dtype, copy=True) + + # For non-integral arguments to PMF (and PDF of discrete distribution) + # replace with zero. + if any_non_integral: + zero = -np.inf if method_name in {'logpmf', 'logpdf'} else 0 + res[mask_non_integral & ~np.isnan(res)] = zero + + # For arguments outside the function domain, replace results + if any_invalid: + replace_low, replace_high = ( + replacements.get(method_name, (np.nan, np.nan))) + res[mask_low] = replace_low + res[mask_high] = replace_high + + # For arguments at the endpoints of the domain, replace results + if any_endpoint: + a, b = self.support() + if a.shape != shape: + a = np.array(np.broadcast_to(a, shape), copy=True) + b = np.array(np.broadcast_to(b, shape), copy=True) + + replace_low_endpoint = ( + b[mask_low_endpoint] if method_name.endswith('ccdf') + else a[mask_low_endpoint]) + replace_high_endpoint = ( + a[mask_high_endpoint] if method_name.endswith('ccdf') + else b[mask_high_endpoint]) + + if not keep_low_endpoint: + res[mask_low_endpoint] = replace_low_endpoint + res[mask_high_endpoint] = replace_high_endpoint + + # Clip probabilities to [0, 1] + if method_name in clip: + res = np.clip(res, 0., 1.) + elif method_name in clip_log: + res = res.real # exp(res) > 0 + res = np.clip(res, None, 0.) # exp(res) < 1 + + return res[()] + + return filtered + + +def _set_invalid_nan_property(f): + # Wrapper for input / output validation and standardization of distribution + # functions that represent properties of the distribution itself: + # logentropy, entropy + # median, mode + # moment + # It ensures that the output is of the correct shape and dtype and that + # there are NaNs wherever the distribution parameters were invalid. + + @functools.wraps(f) + def filtered(self, *args, **kwargs): + if self.validation_policy == _SKIP_ALL: + return f(self, *args, **kwargs) + + res = f(self, *args, **kwargs) + if res is None: + # message could be more appropriate + raise NotImplementedError(self._not_implemented) + + res = np.asarray(res) + needs_copy = False + dtype = res.dtype + + if dtype != self._dtype: # this won't work for logmoments (complex) + dtype = np.result_type(dtype, self._dtype) + needs_copy = True + + if res.shape != self._shape: # faster to check first + res = np.broadcast_to(res, self._shape) + needs_copy = needs_copy or self._any_invalid + + if needs_copy: + res = res.astype(dtype=dtype, copy=True) + + if self._any_invalid: + # may be redundant when quadrature is used, but not necessarily + # when formulas are used. + res[self._invalid] = np.nan + + return res[()] + + return filtered + + +def _dispatch(f): + # For each public method (instance function) of a distribution (e.g. ccdf), + # there may be several ways ("method"s) that it can be computed (e.g. a + # formula, as the complement of the CDF, or via numerical integration). + # Each "method" is implemented by a different private method (instance + # function). + # This wrapper calls the appropriate private method based on the public + # method and any specified `method` keyword option. + # - If `method` is specified as a string (by the user), the appropriate + # private method is called. + # - If `method` is None: + # - The appropriate private method for the public method is looked up + # in a cache. + # - If the cache does not have an entry for the public method, the + # appropriate "dispatch " function is called to determine which method + # is most appropriate given the available private methods and + # settings (e.g. tolerance). + + @functools.wraps(f) + def wrapped(self, *args, method=None, **kwargs): + func_name = f.__name__ + method = method or self._method_cache.get(func_name, None) + if callable(method): + pass + elif method is not None: + method = 'logexp' if method == 'log/exp' else method + method_name = func_name.replace('dispatch', method) + method = getattr(self, method_name) + else: + method = f(self, *args, method=method, **kwargs) + if func_name != '_sample_dispatch' and self.cache_policy != _NO_CACHE: + self._method_cache[func_name] = method + + try: + return method(*args, **kwargs) + except KeyError as e: + raise NotImplementedError(self._not_implemented) from e + + return wrapped + + +def _cdf2_input_validation(f): + # Wrapper that does the job of `_set_invalid_nan` when `cdf` or `logcdf` + # is called with two quantile arguments. + # Let's keep it simple; no special cases for speed right now. + # The strategy is a bit different than for 1-arg `cdf` (and other methods + # covered by `_set_invalid_nan`). For 1-arg `cdf`, elements of `x` that + # are outside (or at the edge of) the support get replaced by `nan`, + # and then the results get replaced by the appropriate value (0 or 1). + # We *could* do something similar, dispatching to `_cdf1` in these + # cases. That would be a bit more robust, but it would also be quite + # a bit more complex, since we'd have to do different things when + # `x` and `y` are both out of bounds, when just `x` is out of bounds, + # when just `y` is out of bounds, and when both are out of bounds. + # I'm not going to do that right now. Instead, simply replace values + # outside the support by those at the edge of the support. Here, we also + # omit some of the optimizations that make `_set_invalid_nan` faster for + # simple arguments (e.g. float64 scalars). + + @functools.wraps(f) + def wrapped(self, x, y, *args, **kwargs): + func_name = f.__name__ + + low, high = self.support() + x, y, low, high = np.broadcast_arrays(x, y, low, high) + dtype = np.result_type(x.dtype, y.dtype, self._dtype) + # yes, copy to avoid modifying input arrays + x, y = x.astype(dtype, copy=True), y.astype(dtype, copy=True) + + # Swap arguments to ensure that x < y, and replace + # out-of domain arguments with domain endpoints. We'll + # transform the result later. + i_swap = y < x + x[i_swap], y[i_swap] = y[i_swap], x[i_swap] + i = x < low + x[i] = low[i] + i = y < low + y[i] = low[i] + i = x > high + x[i] = high[i] + i = y > high + y[i] = high[i] + + res = f(self, x, y, *args, **kwargs) + + # Clipping probability to [0, 1] + if func_name in {'_cdf2', '_ccdf2'}: + res = np.clip(res, 0., 1.) + else: + res = np.clip(res, None, 0.) # exp(res) < 1 + + # Transform the result to account for swapped argument order + res = np.asarray(res) + if func_name == '_cdf2': + res[i_swap] *= -1. + elif func_name == '_ccdf2': + res[i_swap] *= -1 + res[i_swap] += 2. + elif func_name == '_logcdf2': + res = np.asarray(res + 0j) if np.any(i_swap) else res + res[i_swap] = res[i_swap] + np.pi*1j + else: + # res[i_swap] is always positive and less than 1, so it's + # safe to ensure that the result is real + res[i_swap] = _logexpxmexpy(np.log(2), res[i_swap]).real + return res[()] + + return wrapped + + +def _fiinfo(x): + if np.issubdtype(x.dtype, np.inexact): + return np.finfo(x.dtype) + else: + return np.iinfo(x) + + +def _logexpxmexpy(x, y): + """ Compute the log of the difference of the exponentials of two arguments. + + Avoids over/underflow, but does not prevent loss of precision otherwise. + """ + return xpx.apply_where(x != y, (x, y), + lambda x, y: special.logsumexp([x, y+np.pi*1j], axis=0), + fill_value=-np.inf) + + +def _guess_bracket(xmin, xmax): + a = np.full_like(xmin, -1.0) + b = np.ones_like(xmax) + + i = np.isfinite(xmin) & np.isfinite(xmax) + a[i] = xmin[i] + b[i] = xmax[i] + + i = np.isfinite(xmin) & ~np.isfinite(xmax) + a[i] = xmin[i] + b[i] = xmin[i] + 1 + + i = np.isfinite(xmax) & ~np.isfinite(xmin) + a[i] = xmax[i] - 1 + b[i] = xmax[i] + + return a, b + + +def _log_real_standardize(x): + """Standardizes the (complex) logarithm of a real number. + + The logarithm of a real number may be represented by a complex number with + imaginary part that is a multiple of pi*1j. Even multiples correspond with + a positive real and odd multiples correspond with a negative real. + + Given a logarithm of a real number `x`, this function returns an equivalent + representation in a standard form: the log of a positive real has imaginary + part `0` and the log of a negative real has imaginary part `pi`. + + """ + shape = x.shape + x = np.atleast_1d(x) + real = np.real(x).astype(x.dtype) + complex = np.imag(x) + y = real + negative = np.exp(complex*1j) < 0.5 + y[negative] = y[negative] + np.pi * 1j + return y.reshape(shape)[()] + + +def _combine_docs(dist_family, *, include_examples=True): + fields = set(NumpyDocString.sections) + fields.remove('index') + if not include_examples: + fields.remove('Examples') + + doc = ClassDoc(dist_family) + superdoc = ClassDoc(UnivariateDistribution) + for field in fields: + if field in {"Methods", "Attributes"}: + doc[field] = superdoc[field] + elif field in {"Summary"}: + pass + elif field == "Extended Summary": + doc[field].append(_generate_domain_support(dist_family)) + elif field == 'Examples': + doc[field] = [_generate_example(dist_family)] + else: + doc[field] += superdoc[field] + return str(doc) + + +def _generate_domain_support(dist_family): + n_parameterizations = len(dist_family._parameterizations) + + domain = f"\nfor :math:`x \\in {dist_family._variable.domain}`.\n" + + if n_parameterizations == 0: + support = """ + This class accepts no distribution parameters. + """ + elif n_parameterizations == 1: + support = f""" + This class accepts one parameterization: + {str(dist_family._parameterizations[0])}. + """ + else: + number = {2: 'two', 3: 'three', 4: 'four', 5: 'five'}[ + n_parameterizations] + parameterizations = [f"- {str(p)}" for p in + dist_family._parameterizations] + parameterizations = "\n".join(parameterizations) + support = f""" + This class accepts {number} parameterizations: + + {parameterizations} + """ + support = "\n".join([line.lstrip() for line in support.split("\n")][1:]) + return domain + support + + +def _generate_example(dist_family): + n_parameters = dist_family._num_parameters(0) + shapes = [()] * n_parameters + rng = np.random.default_rng(615681484984984) + i = 0 + dist = dist_family._draw(shapes, rng=rng, i_parameterization=i) + + rng = np.random.default_rng(2354873452) + name = dist_family.__name__ + if n_parameters: + parameter_names = list(dist._parameterizations[i].parameters) + parameter_values = [round(getattr(dist, name), 2) for name in + parameter_names] + name_values = [f"{name}={value}" for name, value in + zip(parameter_names, parameter_values)] + instantiation = f"{name}({', '.join(name_values)})" + attributes = ", ".join([f"X.{param}" for param in dist._parameters]) + X = dist_family(**dict(zip(parameter_names, parameter_values))) + else: + instantiation = f"{name}()" + X = dist + + p = 0.32 + x = round(X.icdf(p), 2) + y = round(X.icdf(2 * p), 2) # noqa: F841 + + example = f""" + To use the distribution class, it must be instantiated using keyword + parameters corresponding with one of the accepted parameterizations. + + >>> import numpy as np + >>> import matplotlib.pyplot as plt + >>> from scipy import stats + >>> from scipy.stats import {name} + >>> X = {instantiation} + + For convenience, the ``plot`` method can be used to visualize the density + and other functions of the distribution. + + >>> X.plot() + >>> plt.show() + + The support of the underlying distribution is available using the ``support`` + method. + + >>> X.support() + {X.support()} + """ + + if n_parameters: + example += f""" + The numerical values of parameters associated with all parameterizations + are available as attributes. + + >>> {attributes} + {tuple(X._parameters.values())} + """ + + example += f""" + To evaluate the probability density/mass function of the underlying distribution + at argument ``x={x}``: + + >>> x = {x} + >>> X.pdf(x), X.pmf(x) + {X.pdf(x), X.pmf(x)} + + The cumulative distribution function, its complement, and the logarithm + of these functions are evaluated similarly. + + >>> np.allclose(np.exp(X.logccdf(x)), 1 - X.cdf(x)) + True + """ + + # When two-arg CDF is implemented for DiscreteDistribution, consider removing + # the special-casing here. + if issubclass(dist_family, ContinuousDistribution): + example_continuous = f""" + The inverse of these functions with respect to the argument ``x`` is also + available. + + >>> logp = np.log(1 - X.ccdf(x)) + >>> np.allclose(X.ilogcdf(logp), x) + True + + Note that distribution functions and their logarithms also have two-argument + versions for working with the probability mass between two arguments. The + result tends to be more accurate than the naive implementation because it avoids + subtractive cancellation. + + >>> y = {y} + >>> np.allclose(X.ccdf(x, y), 1 - (X.cdf(y) - X.cdf(x))) + True + """ + example += example_continuous + + example += f""" + There are methods for computing measures of central tendency, + dispersion, higher moments, and entropy. + + >>> X.mean(), X.median(), X.mode() + {X.mean(), X.median(), X.mode()} + + >>> X.variance(), X.standard_deviation() + {X.variance(), X.standard_deviation()} + + >>> X.skewness(), X.kurtosis() + {X.skewness(), X.kurtosis()} + + >>> np.allclose(X.moment(order=6, kind='standardized'), + ... X.moment(order=6, kind='central') / X.variance()**3) + True + """ + + # When logentropy is implemented for DiscreteDistribution, remove special-casing + if issubclass(dist_family, ContinuousDistribution): + example += """ + >>> np.allclose(np.exp(X.logentropy()), X.entropy()) + True + """ + else: + example += f""" + >>> X.entropy() + {X.entropy()} + """ + + example += f""" + Pseudo-random samples can be drawn from + the underlying distribution using ``sample``. + + >>> X.sample(shape=(4,)) + {repr(X.sample(shape=(4,)))} # may vary + """ + # remove the indentation due to use of block quote within function; + # eliminate blank first line + example = "\n".join([line.lstrip() for line in example.split("\n")][1:]) + return example + + +class UnivariateDistribution(_ProbabilityDistribution): + r""" Class that represents a continuous statistical distribution. + + Parameters + ---------- + tol : positive float, optional + The desired relative tolerance of calculations. Left unspecified, + calculations may be faster; when provided, calculations may be + more likely to meet the desired accuracy. + validation_policy : {None, "skip_all"} + Specifies the level of input validation to perform. Left unspecified, + input validation is performed to ensure appropriate behavior in edge + case (e.g. parameters out of domain, argument outside of distribution + support, etc.) and improve consistency of output dtype, shape, etc. + Pass ``'skip_all'`` to avoid the computational overhead of these + checks when rough edges are acceptable. + cache_policy : {None, "no_cache"} + Specifies the extent to which intermediate results are cached. Left + unspecified, intermediate results of some calculations (e.g. distribution + support, moments, etc.) are cached to improve performance of future + calculations. Pass ``'no_cache'`` to reduce memory reserved by the class + instance. + + Attributes + ---------- + All parameters are available as attributes. + + Methods + ------- + support + + plot + + sample + + moment + lmoment + + mean + median + mode + + variance + standard_deviation + + skewness + kurtosis + + pdf + logpdf + + cdf + icdf + ccdf + iccdf + + logcdf + ilogcdf + logccdf + ilogccdf + + entropy + logentropy + + See Also + -------- + :ref:`rv_infrastructure` : Tutorial + + Notes + ----- + The following abbreviations are used throughout the documentation. + + - PDF: probability density function + - CDF: cumulative distribution function + - CCDF: complementary CDF + - entropy: differential entropy + - log-*F*: logarithm of *F* (e.g. log-CDF) + - inverse *F*: inverse function of *F* (e.g. inverse CDF) + + The API documentation is written to describe the API, not to serve as + a statistical reference. Effort is made to be correct at the level + required to use the functionality, not to be mathematically rigorous. + For example, continuity and differentiability may be implicitly assumed. + For precise mathematical definitions, consult your preferred mathematical + text. + + """ + __array_priority__ = 1 + _parameterizations = [] # type: ignore[var-annotated] + + ### Initialization + + def __init__(self, *, tol=_null, validation_policy=None, cache_policy=None, + **parameters): + self.tol = tol + self.validation_policy = validation_policy + self.cache_policy = cache_policy + self._not_implemented = ( + f"`{self.__class__.__name__}` does not provide an accurate " + "implementation of the required method. Consider leaving " + "`method` and `tol` unspecified to use another implementation." + ) + self._original_parameters = {} + # We may want to override the `__init__` method with parameters so + # IDEs can suggest parameter names. If there are multiple parameterizations, + # we'll need the default values of parameters to be None; this will + # filter out the parameters that were not actually specified by the user. + parameters = {key: val for key, val in + sorted(parameters.items()) if val is not None} + self._update_parameters(**parameters) + + def _update_parameters(self, *, validation_policy=None, **params): + r""" Update the numerical values of distribution parameters. + + Parameters + ---------- + validation_policy : {None, "skip_all"} + Specifies the level of input validation to perform. Left unspecified, + input validation is performed to ensure appropriate behavior in edge + case (e.g. parameters out of domain, argument outside of distribution + support, etc.) and improve consistency of output dtype, shape, etc. + Pass ``'skip_all'`` to avoid the computational overhead of these + checks when rough edges are acceptable. + + **params : array_like + Desired numerical values of the distribution parameters. Any or all + of the parameters initially used to instantiate the distribution + may be modified. Parameters used in alternative parameterizations + are not accepted. + """ + + parameters = original_parameters = self._original_parameters.copy() + parameters.update(**params) + parameterization = None + self._invalid = np.asarray(False) + self._any_invalid = False + self._shape = tuple() + self._ndim = 0 + self._size = 1 + self._dtype = np.float64 + + if (validation_policy or self.validation_policy) == _SKIP_ALL: + parameters = self._process_parameters(**parameters) + elif not len(self._parameterizations): + if parameters: + message = (f"The `{self.__class__.__name__}` distribution " + "family does not accept parameters, but parameters " + f"`{set(parameters)}` were provided.") + raise ValueError(message) + else: + # This is default behavior, which re-runs all parameter validations + # even when only a single parameter is modified. For many + # distributions, the domain of a parameter doesn't depend on other + # parameters, so parameters could safely be modified without + # re-validating all other parameters. To handle these cases more + # efficiently, we could allow the developer to override this + # behavior. + + # Currently the user can only update the original parameterization. + # Even though that parameterization is already known, + # `_identify_parameterization` is called to produce a nice error + # message if the user passes other values. To be a little more + # efficient, we could detect whether the values passed are + # consistent with the original parameterization rather than finding + # it from scratch. However, we might want other parameterizations + # to be accepted, which would require other changes, so I didn't + # optimize this. + + parameterization = self._identify_parameterization(parameters) + parameters, shape, size, ndim = self._broadcast(parameters) + parameters, invalid, any_invalid, dtype = ( + self._validate(parameterization, parameters)) + parameters = self._process_parameters(**parameters) + + self._invalid = invalid + self._any_invalid = any_invalid + self._shape = shape + self._size = size + self._ndim = ndim + self._dtype = dtype + + self.reset_cache() + self._parameters = parameters + self._parameterization = parameterization + self._original_parameters = original_parameters + for name in self._parameters.keys(): + # Make parameters properties of the class; return values from the instance + if hasattr(self.__class__, name): + continue + setattr(self.__class__, name, property(lambda self_, name_=name: + self_._parameters[name_].copy()[()])) + + def reset_cache(self): + r""" Clear all cached values. + + To improve the speed of some calculations, the distribution's support + and moments are cached. + + This function is called automatically whenever the distribution + parameters are updated. + + """ + # We could offer finer control over what is cleared. + # For simplicity, these will still exist even if cache_policy is + # NO_CACHE; they just won't be populated. This allows caching to be + # turned on and off easily. + self._moment_raw_cache = {} + self._moment_central_cache = {} + self._moment_standardized_cache = {} + self._lmoment_cache = {} + self._support_cache = None + self._method_cache = {} + self._constant_cache = None + + def _identify_parameterization(self, parameters): + # Determine whether a `parameters` dictionary matches is consistent + # with one of the parameterizations of the distribution. If so, + # return that parameterization object; if not, raise an error. + # + # I've come back to this a few times wanting to avoid this explicit + # loop. I've considered several possibilities, but they've all been a + # little unusual. For example, we could override `_eq_` so we can + # use _parameterizations.index() to retrieve the parameterization, + # or the user could put the parameterizations in a dictionary so we + # could look them up with a key (e.g. frozenset of parameter names). + # I haven't been sure enough of these approaches to implement them. + parameter_names_set = set(parameters) + + for parameterization in self._parameterizations: + if parameterization.matches(parameter_names_set): + break + else: + if not parameter_names_set: + message = (f"The `{self.__class__.__name__}` distribution " + "family requires parameters, but none were " + "provided.") + else: + parameter_names = self._get_parameter_str(parameters) + message = (f"The provided parameters `{parameter_names}` " + "do not match a supported parameterization of the " + f"`{self.__class__.__name__}` distribution family.") + raise ValueError(message) + + return parameterization + + def _broadcast(self, parameters): + # Broadcast the distribution parameters to the same shape. If the + # arrays are not broadcastable, raise a meaningful error. + # + # We always make sure that the parameters *are* the same shape + # and not just broadcastable. Users can access parameters as + # attributes, and I think they should see the arrays as the same shape. + # More importantly, arrays should be the same shape before logical + # indexing operations, which are needed in infrastructure code when + # there are invalid parameters, and may be needed in + # distribution-specific code. We don't want developers to need to + # broadcast in implementation functions. + + # It's much faster to check whether broadcasting is necessary than to + # broadcast when it's not necessary. + parameter_vals = [np.asarray(parameter) + for parameter in parameters.values()] + parameter_shapes = set(parameter.shape for parameter in parameter_vals) + if len(parameter_shapes) == 1: + return (parameters, parameter_vals[0].shape, + parameter_vals[0].size, parameter_vals[0].ndim) + + try: + parameter_vals = np.broadcast_arrays(*parameter_vals) + except ValueError as e: + parameter_names = self._get_parameter_str(parameters) + message = (f"The parameters `{parameter_names}` provided to the " + f"`{self.__class__.__name__}` distribution family " + "cannot be broadcast to the same shape.") + raise ValueError(message) from e + return (dict(zip(parameters.keys(), parameter_vals)), + parameter_vals[0].shape, + parameter_vals[0].size, + parameter_vals[0].ndim) + + def _validate(self, parameterization, parameters): + # Broadcasts distribution parameter arrays and converts them to a + # consistent dtype. Replaces invalid parameters with `np.nan`. + # Returns the validated parameters, a boolean mask indicated *which* + # elements are invalid, a boolean scalar indicating whether *any* + # are invalid (to skip special treatments if none are invalid), and + # the common dtype. + valid, dtype = parameterization.validation(parameters) + invalid = ~valid + any_invalid = invalid if invalid.shape == () else np.any(invalid) + # If necessary, make the arrays contiguous and replace invalid with NaN + if any_invalid: + for parameter_name in parameters: + parameters[parameter_name] = np.copy( + parameters[parameter_name]) + parameters[parameter_name][invalid] = np.nan + + return parameters, invalid, any_invalid, dtype + + def _process_parameters(self, **params): + r""" Process and cache distribution parameters for reuse. + + This is intended to be overridden by subclasses. It allows distribution + authors to pre-process parameters for re-use. For instance, when a user + parameterizes a LogUniform distribution with `a` and `b`, it makes + sense to calculate `log(a)` and `log(b)` because these values will be + used in almost all distribution methods. The dictionary returned by + this method is passed to all private methods that calculate functions + of the distribution. + """ + return params + + def _get_parameter_str(self, parameters): + # Get a string representation of the parameters like "{a, b, c}". + return f"{{{', '.join(parameters.keys())}}}" + + def _copy_parameterization(self): + self._parameterizations = self._parameterizations.copy() + for i in range(len(self._parameterizations)): + self._parameterizations[i] = self._parameterizations[i].copy() + + ### Attributes + + @property + def tol(self): + r"""positive float: + The desired relative tolerance of calculations. Left unspecified, + calculations may be faster; when provided, calculations may be + more likely to meet the desired accuracy. + """ + return self._tol + + @tol.setter + def tol(self, tol): + if _isnull(tol): + self._tol = tol + return + + tol = np.asarray(tol) + if (tol.shape != () or not tol > 0 or # catches NaNs + not np.issubdtype(tol.dtype, np.floating)): + message = (f"Attribute `tol` of `{self.__class__.__name__}` must " + "be a positive float, if specified.") + raise ValueError(message) + self._tol = tol[()] + + @property + def cache_policy(self): + r"""{None, "no_cache"}: + Specifies the extent to which intermediate results are cached. Left + unspecified, intermediate results of some calculations (e.g. distribution + support, moments, etc.) are cached to improve performance of future + calculations. Pass ``'no_cache'`` to reduce memory reserved by the class + instance. + """ + return self._cache_policy + + @cache_policy.setter + def cache_policy(self, cache_policy): + cache_policy = str(cache_policy).lower() if cache_policy is not None else None + cache_policies = {None, 'no_cache'} + if cache_policy not in cache_policies: + message = (f"Attribute `cache_policy` of `{self.__class__.__name__}` " + f"must be one of {cache_policies}, if specified.") + raise ValueError(message) + self._cache_policy = cache_policy + + @property + def validation_policy(self): + r"""{None, "skip_all"}: + Specifies the level of input validation to perform. Left unspecified, + input validation is performed to ensure appropriate behavior in edge + case (e.g. parameters out of domain, argument outside of distribution + support, etc.) and improve consistency of output dtype, shape, etc. + Use ``'skip_all'`` to avoid the computational overhead of these + checks when rough edges are acceptable. + """ + return self._validation_policy + + @validation_policy.setter + def validation_policy(self, validation_policy): + validation_policy = (str(validation_policy).lower() + if validation_policy is not None else None) + iv_policies = {None, 'skip_all'} + if validation_policy not in iv_policies: + message = (f"Attribute `validation_policy` of `{self.__class__.__name__}` " + f"must be one of {iv_policies}, if specified.") + raise ValueError(message) + self._validation_policy = validation_policy + + ### Other magic methods + + def __repr__(self): + r""" Returns a string representation of the distribution. + + Includes the name of the distribution family, the names of the + parameters and the `repr` of each of their values. + + + """ + class_name = self.__class__.__name__ + parameters = list(self._original_parameters.items()) + info = [] + with np.printoptions(threshold=10): + str_parameters = [f"{symbol}={repr(value)}" for symbol, value in parameters] + str_parameters = f"{', '.join(str_parameters)}" + info.append(str_parameters) + return f"{class_name}({', '.join(info)})" + + def __str__(self): + class_name = self.__class__.__name__ + parameters = list(self._original_parameters.items()) + info = [] + with np.printoptions(threshold=10): + str_parameters = [f"{symbol}={str(value)}" for symbol, value in parameters] + str_parameters = f"{', '.join(str_parameters)}" + info.append(str_parameters) + return f"{class_name}({', '.join(info)})" + + def __add__(self, loc): + return ShiftedScaledDistribution(self, loc=loc) + + def __sub__(self, loc): + return ShiftedScaledDistribution(self, loc=-loc) + + def __mul__(self, scale): + return ShiftedScaledDistribution(self, scale=scale) + + def __truediv__(self, scale): + return ShiftedScaledDistribution(self, scale=1/scale) + + def __pow__(self, other): + if not np.isscalar(other) or other <= 0 or other != int(other): + message = ("Raising a random variable to the power of an argument is only " + "implemented when the argument is a positive integer.") + raise NotImplementedError(message) + + # Fill in repr_pattern with the repr of self before taking abs. + # Avoids having unnecessary abs in the repr. + with np.printoptions(threshold=10): + repr_pattern = f"({repr(self)})**{repr(other)}" + str_pattern = f"({str(self)})**{str(other)}" + X = abs(self) if other % 2 == 0 else self + + funcs = dict(g=lambda u: u**other, repr_pattern=repr_pattern, + str_pattern=str_pattern, + h=lambda u: np.sign(u) * np.abs(u)**(1 / other), + dh=lambda u: 1/other * np.abs(u)**(1/other - 1)) + + return MonotonicTransformedDistribution(X, **funcs, increasing=True) + + def __radd__(self, other): + return self.__add__(other) + + def __rsub__(self, other): + return self.__neg__().__add__(other) + + def __rmul__(self, other): + return self.__mul__(other) + + def __rtruediv__(self, other): + a, b = self.support() + with np.printoptions(threshold=10): + funcs = dict(g=lambda u: 1 / u, + repr_pattern=f"{repr(other)}/({repr(self)})", + str_pattern=f"{str(other)}/({str(self)})", + h=lambda u: 1 / u, dh=lambda u: 1 / u ** 2) + if np.all(a >= 0) or np.all(b <= 0): + out = MonotonicTransformedDistribution(self, **funcs, increasing=False) + else: + message = ("Division by a random variable is only implemented " + "when the support is either non-negative or non-positive.") + raise NotImplementedError(message) + if np.all(other == 1): + return out + else: + return out * other + + def __rpow__(self, other): + with np.printoptions(threshold=10): + funcs = dict(g=lambda u: other**u, + h=lambda u: np.log(u) / np.log(other), + dh=lambda u: 1 / np.abs(u * np.log(other)), + repr_pattern=f"{repr(other)}**({repr(self)})", + str_pattern=f"{str(other)}**({str(self)})",) + + if not np.isscalar(other) or other <= 0 or other == 1: + message = ("Raising an argument to the power of a random variable is only " + "implemented when the argument is a positive scalar other than " + "1.") + raise NotImplementedError(message) + + if other > 1: + return MonotonicTransformedDistribution(self, **funcs, increasing=True) + else: + return MonotonicTransformedDistribution(self, **funcs, increasing=False) + + def __neg__(self): + return self * -1 + + def __abs__(self): + return FoldedDistribution(self) + + ### Utilities + + ## Input validation + + def _validate_order(self, order, fname='moment', min_order=0): + # Yet another integer validating function. Unlike others in SciPy, it + # Is quite flexible about what is allowed as an integer, and it + # raises a distribution-specific error message to facilitate + # identification of the source of the error. + if self.validation_policy == _SKIP_ALL: + return order + + order = np.asarray(order, dtype=self._dtype)[()] + message = (f"Argument `order` of `{self.__class__.__name__}.{fname}` " + f"must be a finite integer greater than or equal to {min_order}.") + try: + order_int = round(order.item()) + # If this fails for any reason (e.g. it's an array, it's infinite) + # it's not a valid `order`. + except Exception as e: + raise ValueError(message) from e + + if order_int < min_order or order_int != order: + raise ValueError(message) + + return order + + def _validate_kind(self, kind, kinds): + message = (f"Argument `kind` of `{self.__class__.__name__}.moment` " + f"must be one of {set(kinds)}.") + kind = kind.lower() + if kind not in kinds: + raise ValueError(message) + return kind + + def _preserve_type(self, x): + x = np.asarray(x) + if x.dtype != self._dtype: + x = x.astype(self._dtype) + return x[()] + + ## Testing + + @classmethod + def _draw(cls, sizes=None, rng=None, i_parameterization=None, + proportions=None): + r""" Draw a specific (fully-defined) distribution from the family. + + See _Parameterization.draw for documentation details. + """ + rng = np.random.default_rng(rng) + if len(cls._parameterizations) == 0: + return cls() + if i_parameterization is None: + n = cls._num_parameterizations() + i_parameterization = rng.integers(0, max(0, n - 1), endpoint=True) + + parameterization = cls._parameterizations[i_parameterization] + parameters = parameterization.draw(sizes, rng, proportions=proportions, + region='typical') + return cls(**parameters) + + @classmethod + def _num_parameterizations(cls): + # Returns the number of parameterizations accepted by the family. + return len(cls._parameterizations) + + @classmethod + def _num_parameters(cls, i_parameterization=0): + # Returns the number of parameters used in the specified + # parameterization. + return (0 if not cls._num_parameterizations() + else len(cls._parameterizations[i_parameterization])) + + ## Algorithms + + def _quadrature(self, integrand, limits=None, args=(), params=None, log=False): + # Performs numerical integration/summation between limits or over support. + a, b = self._support(**params) if limits is None else limits + rtol = None if _isnull(self.tol) else self.tol + + # For now, we ignore the status, but I want to return the error estimate + if isinstance(self, ContinuousDistribution): + return _tanhsinh(integrand, a, b, args=args, kwargs=params, + log=log, rtol=rtol).integral + else: + res = nsum(integrand, a, b, args=args, kwargs=params, + log=log, tolerances=dict(rtol=rtol)).sum + + # b > a -> sum is zero... unless the parameters are nan. + # (Resolving gh-22321 should fix this in nsum.) + cond = np.isnan(params.popitem()[1]) if params else np.True_ + return np.where(cond, np.nan, res)[()] + + def _solve_bounded(self, f, p, *, bounds=None, params=None, xatol=None): + # Finds the argument of a function that produces the desired output. + xmin, xmax = self._support(**params) if bounds is None else bounds + + def f2(x, _p, **kwargs): # named `_p` to avoid conflict with shape `p` + return f(x, **kwargs) - _p + + xl0, xr0 = _guess_bracket(xmin, xmax) + + res = elementwise.bracket_root(f2, xl0=xl0, xr0=xr0, xmin=xmin, xmax=xmax, + args=(p,), kwargs=params) + # For now, we ignore the status, but I want to use the bracket width + # as an error estimate. + + xrtol = None if _isnull(self.tol) else self.tol + xatol = None if xatol is None else xatol + tolerances = dict(xrtol=xrtol, xatol=xatol, fatol=0, frtol=0) + return elementwise.find_root(f2, res.bracket, args=(p,), + kwargs=params, tolerances=tolerances) + + ## Other + + def _overrides(self, method_name): + # Determines whether a class overrides a specified method. + # Returns True if the method implementation exists and is the same as + # that of the `ContinuousDistribution` class; otherwise returns False. + + # Sometimes we use `_overrides` to check whether a certain method is overridden + # and if so, call it. This begs the questions of why we don't do the more + # obvious thing: restructure so that if the private method is overridden, + # Python will call it instead of the inherited version automatically. The short + # answer is that there are multiple ways a use might wish to evaluate a method, + # and simply overriding the method with a formula is not always the best option. + # For more complete discussion of the considerations, see: + # https://github.com/scipy/scipy/pull/21050#discussion_r1707798901 + method = getattr(self.__class__, method_name, None) + super_method = getattr(UnivariateDistribution, method_name, None) + return method is not super_method + + ### Distribution properties + # The following "distribution properties" are exposed via a public method + # that accepts only options (not distribution parameters or quantile/ + # percentile argument). + # support + # logentropy, entropy, + # median, mode, mean, + # variance, standard_deviation + # skewness, kurtosis + # Common options are: + # method - a string that indicates which method should be used to compute + # the quantity (e.g. a formula or numerical integration). + # Input/output validation is provided by the `_set_invalid_nan_property` + # decorator. These are the methods meant to be called by users. + # + # Each public method calls a private "dispatch" method that + # determines which "method" (strategy for calculating the desired quantity) + # to use by default and, via the `@_dispatch` decorator, calls the + # method and computes the result. + # Dispatch methods always accept: + # method - as passed from the public method + # params - a dictionary of distribution shape parameters passed by + # the public method. + # Dispatch methods accept `params` rather than relying on the state of the + # object because iterative algorithms like `tanhsinh` and `find_root` + # need their callable to follow a strict elementwise protocol: each element + # of the output is determined solely by the values of the inputs at the + # corresponding location. The public methods do not satisfy this protocol + # because they do not accept the parameters as arguments, producing an + # output that generally has a different shape than that of the input. Also, + # by calling "dispatch" methods rather than the public methods, the + # iterative algorithms avoid the overhead of input validation. + # + # Each dispatch method can designate the responsibility of computing + # the required value to any of several "implementation" methods. These + # methods accept only `**params`, the parameter dictionary passed from + # the public method via the dispatch method. We separate the implementation + # methods from the dispatch methods for the sake of simplicity (via + # compartmentalization) and to allow subclasses to override certain + # implementation methods (typically only the "formula" methods). The names + # of implementation methods are combinations of the public method name and + # the name of the "method" (strategy for calculating the desired quantity) + # string. (In fact, the name of the implementation method is calculated + # from these two strings in the `_dispatch` decorator.) Common method + # strings are: + # formula - distribution-specific analytical expressions to be implemented + # by subclasses. + # log/exp - Compute the log of a number and then exponentiate it or vice + # versa. + # quadrature - Compute the value via numerical integration. + # + # The default method (strategy) is determined based on what implementation + # methods are available and the error tolerance of the user. Typically, + # a formula is always used if available. We fall back to "log/exp" if a + # formula for the logarithm or exponential of the quantity is available, + # and we use quadrature otherwise. + + def support(self): + # If this were a `cached_property`, we couldn't update the value + # when the distribution parameters change. + # Caching is important, though, because calls to _support take a few + # microseconds even when `a` and `b` are already the same shape. + if self._support_cache is not None: + return self._support_cache + + a, b = self._support(**self._parameters) + if a.shape != self._shape: + a = np.broadcast_to(a, self._shape) + if b.shape != self._shape: + b = np.broadcast_to(b, self._shape) + + if self._any_invalid: + a, b = np.asarray(a).copy(), np.asarray(b).copy() + a[self._invalid], b[self._invalid] = np.nan, np.nan + a, b = a[()], b[()] + + support = (a, b) + + if self.cache_policy != _NO_CACHE: + self._support_cache = support + + return support + + def _support(self, **params): + # Computes the support given distribution parameters + a, b = self._variable.domain.get_numerical_endpoints(params) + if len(params): + # the parameters should all be of the same dtype and shape at this point + vals = list(params.values()) + shape = vals[0].shape + a = np.broadcast_to(a, shape) if a.shape != shape else a + b = np.broadcast_to(b, shape) if b.shape != shape else b + return self._preserve_type(a), self._preserve_type(b) + + @_set_invalid_nan_property + def logentropy(self, *, method=None): + return self._logentropy_dispatch(method=method, **self._parameters) + 0j + + @_dispatch + def _logentropy_dispatch(self, method=None, **params): + if self._overrides('_logentropy_formula'): + method = self._logentropy_formula + elif self._overrides('_entropy_formula'): + method = self._logentropy_logexp_safe + else: + method = self._logentropy_quadrature + return method + + def _logentropy_formula(self, **params): + raise NotImplementedError(self._not_implemented) + + def _logentropy_logexp(self, **params): + res = np.log(self._entropy_dispatch(**params)+0j) + return _log_real_standardize(res) + + def _logentropy_logexp_safe(self, **params): + out = self._logentropy_logexp(**params) + mask = np.isinf(out.real) + if np.any(mask): + params_mask = {key:val[mask] for key, val in params.items()} + out = np.asarray(out) + out[mask] = self._logentropy_quadrature(**params_mask) + return out[()] + + def _logentropy_quadrature(self, **params): + def logintegrand(x, **params): + logpxf = self._logpxf_dispatch(x, **params) + return logpxf + np.log(0j+logpxf) + res = self._quadrature(logintegrand, params=params, log=True) + return _log_real_standardize(res + np.pi*1j) + + @_set_invalid_nan_property + def entropy(self, *, method=None): + return self._entropy_dispatch(method=method, **self._parameters) + + @_dispatch + def _entropy_dispatch(self, method=None, **params): + if self._overrides('_entropy_formula'): + method = self._entropy_formula + elif self._overrides('_logentropy_formula'): + method = self._entropy_logexp + else: + method = self._entropy_quadrature + return method + + def _entropy_formula(self, **params): + raise NotImplementedError(self._not_implemented) + + def _entropy_logexp(self, **params): + return np.real(np.exp(self._logentropy_dispatch(**params))) + + def _entropy_quadrature(self, **params): + def integrand(x, **params): + pxf = self._pxf_dispatch(x, **params) + logpxf = self._logpxf_dispatch(x, **params) + temp = np.asarray(pxf) + i = (pxf != 0) # 0 * inf -> nan; should be 0 + temp[i] = -pxf[i]*logpxf[i] + return temp + return self._quadrature(integrand, params=params) + + @_set_invalid_nan_property + def median(self, *, method=None): + return self._median_dispatch(method=method, **self._parameters) + + @_dispatch + def _median_dispatch(self, method=None, **params): + if self._overrides('_median_formula'): + method = self._median_formula + else: + method = self._median_icdf + return method + + def _median_formula(self, **params): + raise NotImplementedError(self._not_implemented) + + def _median_icdf(self, **params): + return self._icdf_dispatch(np.asarray(0.5, dtype=self._dtype), **params) + + @_set_invalid_nan_property + def mode(self, *, method=None): + return self._mode_dispatch(method=method, **self._parameters) + + @_dispatch + def _mode_dispatch(self, method=None, **params): + # We could add a method that looks for a critical point with + # differentiation and the root finder + if self._overrides('_mode_formula'): + method = self._mode_formula + else: + method = self._mode_optimization + return method + + def _mode_formula(self, **params): + raise NotImplementedError(self._not_implemented) + + def _mode_optimization(self, xatol=None, **params): + a, b = self._support(**params) + m = self._median_dispatch(**params) + + def f(x, **params): + return -self._pxf_dispatch(x, **params) + + res_b = elementwise.bracket_minimum(f, m, xmin=a, xmax=b, kwargs=params) + res = elementwise.find_minimum(f, res_b.bracket, kwargs=params, + tolerances=dict(xatol=xatol)) + mode = np.asarray(res.x) + + # bracket_minimum fails with status == -1 if the mode is at a boundary. + # We have to consider this as a special case; there's no way to include + # this logic in find_minimum; it has to treat the bracket as invalid. + mode_at_boundary = res.status == -1 + fl, fm, fr = res_b.f_bracket + mode_at_left = mode_at_boundary & (fl <= fm) + mode_at_right = mode_at_boundary & (fr < fm) + mode[mode_at_left] = a[mode_at_left] + mode[mode_at_right] = b[mode_at_right] + return mode[()] + + def mean(self, *, method=None): + return self.moment(1, kind='raw', method=method) + + def variance(self, *, method=None): + return self.moment(2, kind='central', method=method) + + def standard_deviation(self, *, method=None): + return np.sqrt(self.variance(method=method)) + + def skewness(self, *, method=None): + return self.moment(3, kind='standardized', method=method) + + def kurtosis(self, *, method=None, convention='non-excess'): + conventions = {'non-excess', 'excess'} + message = (f'Parameter `convention` of `{self.__class__.__name__}.kurtosis` ' + f"must be one of {conventions}.") + convention = convention.lower() + if convention not in conventions: + raise ValueError(message) + k = self.moment(4, kind='standardized', method=method) + return k - 3 if convention == 'excess' else k + + ### Distribution functions + # The following functions related to the distribution PDF and CDF are + # exposed via a public method that accepts one positional argument - the + # quantile - and keyword options (but not distribution parameters). + # logpdf, pdf + # logcdf, cdf + # logccdf, ccdf + # The `logcdf` and `cdf` functions can also be called with two positional + # arguments - lower and upper quantiles - and they return the probability + # mass (integral of the PDF) between them. The 2-arg versions of `logccdf` + # and `ccdf` return the complement of this quantity. + # All the (1-arg) cumulative distribution functions have inverse + # functions, which accept one positional argument - the percentile. + # ilogcdf, icdf + # ilogccdf, iccdf + # Common keyword options include: + # method - a string that indicates which method should be used to compute + # the quantity (e.g. a formula or numerical integration). + # Tolerance options should be added. + # Input/output validation is provided by the `_set_invalid_nan` + # decorator. These are the methods meant to be called by users. + # + # Each public method calls a private "dispatch" method that + # determines which "method" (strategy for calculating the desired quantity) + # to use by default and, via the `@_dispatch` decorator, calls the + # method and computes the result. + # Each dispatch method can designate the responsibility of computing + # the required value to any of several "implementation" methods. These + # methods accept only `**params`, the parameter dictionary passed from + # the public method via the dispatch method. + # See the note corresponding with the "Distribution Parameters" for more + # information. + + ## Probability Density/Mass Functions + + @_set_invalid_nan + def logpdf(self, x, /, *, method=None): + return self._logpdf_dispatch(x, method=method, **self._parameters) + + @_dispatch + def _logpdf_dispatch(self, x, *, method=None, **params): + if self._overrides('_logpdf_formula'): + method = self._logpdf_formula + elif _isnull(self.tol): # ensure that developers override _logpdf + method = self._logpdf_logexp + return method + + def _logpdf_formula(self, x, **params): + raise NotImplementedError(self._not_implemented) + + def _logpdf_logexp(self, x, **params): + return np.log(self._pdf_dispatch(x, **params)) + + @_set_invalid_nan + def pdf(self, x, /, *, method=None): + return self._pdf_dispatch(x, method=method, **self._parameters) + + @_dispatch + def _pdf_dispatch(self, x, *, method=None, **params): + if self._overrides('_pdf_formula'): + method = self._pdf_formula + else: + method = self._pdf_logexp + return method + + def _pdf_formula(self, x, **params): + raise NotImplementedError(self._not_implemented) + + def _pdf_logexp(self, x, **params): + return np.exp(self._logpdf_dispatch(x, **params)) + + @_set_invalid_nan + def logpmf(self, x, /, *, method=None): + return self._logpmf_dispatch(x, method=method, **self._parameters) + + @_dispatch + def _logpmf_dispatch(self, x, *, method=None, **params): + if self._overrides('_logpmf_formula'): + method = self._logpmf_formula + elif _isnull(self.tol): # ensure that developers override _logpmf + method = self._logpmf_logexp + return method + + def _logpmf_formula(self, x, **params): + raise NotImplementedError(self._not_implemented) + + def _logpmf_logexp(self, x, **params): + with np.errstate(divide='ignore'): + return np.log(self._pmf_dispatch(x, **params)) + + @_set_invalid_nan + def pmf(self, x, /, *, method=None): + return self._pmf_dispatch(x, method=method, **self._parameters) + + @_dispatch + def _pmf_dispatch(self, x, *, method=None, **params): + if self._overrides('_pmf_formula'): + method = self._pmf_formula + else: + method = self._pmf_logexp + return method + + def _pmf_formula(self, x, **params): + raise NotImplementedError(self._not_implemented) + + def _pmf_logexp(self, x, **params): + return np.exp(self._logpmf_dispatch(x, **params)) + + ## Cumulative Distribution Functions + + def logcdf(self, x, y=None, /, *, method=None): + if y is None: + return self._logcdf1(x, method=method) + else: + return self._logcdf2(x, y, method=method) + + @_cdf2_input_validation + def _logcdf2(self, x, y, *, method): + out = self._logcdf2_dispatch(x, y, method=method, **self._parameters) + return (out + 0j) if not np.issubdtype(out.dtype, np.complexfloating) else out + + @_dispatch + def _logcdf2_dispatch(self, x, y, *, method=None, **params): + if self._overrides('_logcdf2_formula'): + method = self._logcdf2_formula + elif (self._overrides('_logcdf_formula') + or self._overrides('_logccdf_formula')): + method = self._logcdf2_subtraction + elif (self._overrides('_cdf_formula') + or self._overrides('_ccdf_formula')): + method = self._logcdf2_logexp_safe + else: + method = self._logcdf2_quadrature + return method + + def _logcdf2_formula(self, x, y, **params): + raise NotImplementedError(self._not_implemented) + + def _logcdf2_subtraction(self, x, y, **params): + flip_sign = x > y # some results will be negative + x, y = np.minimum(x, y), np.maximum(x, y) + logcdf_x = self._logcdf_dispatch(x, **params) + logcdf_y = self._logcdf_dispatch(y, **params) + logccdf_x = self._logccdf_dispatch(x, **params) + logccdf_y = self._logccdf_dispatch(y, **params) + case_left = (logcdf_x < -1) & (logcdf_y < -1) + case_right = (logccdf_x < -1) & (logccdf_y < -1) + case_central = ~(case_left | case_right) + log_mass = _logexpxmexpy(logcdf_y, logcdf_x) + log_mass[case_right] = _logexpxmexpy(logccdf_x, logccdf_y)[case_right] + log_tail = np.logaddexp(logcdf_x, logccdf_y)[case_central] + log_mass[case_central] = _log1mexp(log_tail) + log_mass[flip_sign] += np.pi * 1j + return log_mass[()] if np.any(flip_sign) else log_mass.real[()] + + def _logcdf2_logexp(self, x, y, **params): + expres = self._cdf2_dispatch(x, y, **params) + expres = expres + 0j if np.any(x > y) else expres + return np.log(expres) + + def _logcdf2_logexp_safe(self, x, y, **params): + out = self._logcdf2_logexp(x, y, **params) + mask = np.isinf(out.real) + if np.any(mask): + params_mask = {key: np.broadcast_to(val, mask.shape)[mask] + for key, val in params.items()} + out = np.asarray(out) + out[mask] = self._logcdf2_quadrature(x[mask], y[mask], **params_mask) + return out[()] + + def _logcdf2_quadrature(self, x, y, **params): + logres = self._quadrature(self._logpxf_dispatch, limits=(x, y), + log=True, params=params) + return logres + + @_set_invalid_nan + def _logcdf1(self, x, *, method=None): + return self._logcdf_dispatch(x, method=method, **self._parameters) + + @_dispatch + def _logcdf_dispatch(self, x, *, method=None, **params): + if self._overrides('_logcdf_formula'): + method = self._logcdf_formula + elif self._overrides('_logccdf_formula'): + method = self._logcdf_complement + elif self._overrides('_cdf_formula'): + method = self._logcdf_logexp_safe + else: + method = self._logcdf_quadrature + return method + + def _logcdf_formula(self, x, **params): + raise NotImplementedError(self._not_implemented) + + def _logcdf_complement(self, x, **params): + return _log1mexp(self._logccdf_dispatch(x, **params)) + + def _logcdf_logexp(self, x, **params): + return np.log(self._cdf_dispatch(x, **params)) + + def _logcdf_logexp_safe(self, x, **params): + out = self._logcdf_logexp(x, **params) + mask = np.isinf(out) + if np.any(mask): + params_mask = {key:np.broadcast_to(val, mask.shape)[mask] + for key, val in params.items()} + out = np.asarray(out) + out[mask] = self._logcdf_quadrature(x[mask], **params_mask) + return out[()] + + def _logcdf_quadrature(self, x, **params): + a, _ = self._support(**params) + return self._quadrature(self._logpxf_dispatch, limits=(a, x), + params=params, log=True) + + def cdf(self, x, y=None, /, *, method=None): + if y is None: + return self._cdf1(x, method=method) + else: + return self._cdf2(x, y, method=method) + + @_cdf2_input_validation + def _cdf2(self, x, y, *, method): + return self._cdf2_dispatch(x, y, method=method, **self._parameters) + + @_dispatch + def _cdf2_dispatch(self, x, y, *, method=None, **params): + if self._overrides('_cdf2_formula'): + method = self._cdf2_formula + elif (self._overrides('_logcdf_formula') + or self._overrides('_logccdf_formula')): + method = self._cdf2_logexp + elif self._overrides('_cdf_formula') or self._overrides('_ccdf_formula'): + method = self._cdf2_subtraction_safe + else: + method = self._cdf2_quadrature + return method + + def _cdf2_formula(self, x, y, **params): + raise NotImplementedError(self._not_implemented) + + def _cdf2_logexp(self, x, y, **params): + return np.real(np.exp(self._logcdf2_dispatch(x, y, **params))) + + def _cdf2_subtraction(self, x, y, **params): + # Improvements: + # Lazy evaluation of cdf/ccdf only where needed + # Stack x and y to reduce function calls? + cdf_x = self._cdf_dispatch(x, **params) + cdf_y = self._cdf_dispatch(y, **params) + ccdf_x = self._ccdf_dispatch(x, **params) + ccdf_y = self._ccdf_dispatch(y, **params) + i = (ccdf_x < 0.5) & (ccdf_y < 0.5) + return np.where(i, ccdf_x-ccdf_y, cdf_y-cdf_x) + + def _cdf2_subtraction_safe(self, x, y, **params): + cdf_x = self._cdf_dispatch(x, **params) + cdf_y = self._cdf_dispatch(y, **params) + ccdf_x = self._ccdf_dispatch(x, **params) + ccdf_y = self._ccdf_dispatch(y, **params) + i = (ccdf_x < 0.5) & (ccdf_y < 0.5) + out = np.where(i, ccdf_x-ccdf_y, cdf_y-cdf_x) + + eps = np.finfo(self._dtype).eps + tol = self.tol if not _isnull(self.tol) else np.sqrt(eps) + + cdf_max = np.maximum(cdf_x, cdf_y) + ccdf_max = np.maximum(ccdf_x, ccdf_y) + spacing = np.spacing(np.where(i, ccdf_max, cdf_max)) + mask = np.abs(tol * out) < spacing + + if np.any(mask): + params_mask = {key: np.broadcast_to(val, mask.shape)[mask] + for key, val in params.items()} + out = np.asarray(out) + out[mask] = self._cdf2_quadrature(x[mask], y[mask], **params_mask) + return out[()] + + def _cdf2_quadrature(self, x, y, **params): + return self._quadrature(self._pxf_dispatch, limits=(x, y), params=params) + + @_set_invalid_nan + def _cdf1(self, x, *, method): + return self._cdf_dispatch(x, method=method, **self._parameters) + + @_dispatch + def _cdf_dispatch(self, x, *, method=None, **params): + if self._overrides('_cdf_formula'): + method = self._cdf_formula + elif self._overrides('_logcdf_formula'): + method = self._cdf_logexp + elif self._overrides('_ccdf_formula'): + method = self._cdf_complement_safe + else: + method = self._cdf_quadrature + return method + + def _cdf_formula(self, x, **params): + raise NotImplementedError(self._not_implemented) + + def _cdf_logexp(self, x, **params): + return np.exp(self._logcdf_dispatch(x, **params)) + + def _cdf_complement(self, x, **params): + return 1 - self._ccdf_dispatch(x, **params) + + def _cdf_complement_safe(self, x, **params): + ccdf = self._ccdf_dispatch(x, **params) + out = 1 - ccdf + eps = np.finfo(self._dtype).eps + tol = self.tol if not _isnull(self.tol) else np.sqrt(eps) + mask = tol * out < np.spacing(ccdf) + if np.any(mask): + params_mask = {key: np.broadcast_to(val, mask.shape)[mask] + for key, val in params.items()} + out = np.asarray(out) + out[mask] = self._cdf_quadrature(x[mask], **params_mask) + return out[()] + + def _cdf_quadrature(self, x, **params): + a, _ = self._support(**params) + return self._quadrature(self._pxf_dispatch, limits=(a, x), + params=params) + + def logccdf(self, x, y=None, /, *, method=None): + if y is None: + return self._logccdf1(x, method=method) + else: + return self._logccdf2(x, y, method=method) + + @_cdf2_input_validation + def _logccdf2(self, x, y, *, method): + return self._logccdf2_dispatch(x, y, method=method, **self._parameters) + + @_dispatch + def _logccdf2_dispatch(self, x, y, *, method=None, **params): + # if _logccdf2_formula exists, we could use the complement + # if _ccdf2_formula exists, we could use log/exp + if self._overrides('_logccdf2_formula'): + method = self._logccdf2_formula + else: + method = self._logccdf2_addition + return method + + def _logccdf2_formula(self, x, y, **params): + raise NotImplementedError(self._not_implemented) + + def _logccdf2_addition(self, x, y, **params): + logcdf_x = self._logcdf_dispatch(x, **params) + logccdf_y = self._logccdf_dispatch(y, **params) + return special.logsumexp([logcdf_x, logccdf_y], axis=0) + + @_set_invalid_nan + def _logccdf1(self, x, *, method=None): + return self._logccdf_dispatch(x, method=method, **self._parameters) + + @_dispatch + def _logccdf_dispatch(self, x, method=None, **params): + if self._overrides('_logccdf_formula'): + method = self._logccdf_formula + elif self._overrides('_logcdf_formula'): + method = self._logccdf_complement + elif self._overrides('_ccdf_formula'): + method = self._logccdf_logexp_safe + else: + method = self._logccdf_quadrature + return method + + def _logccdf_formula(self, x, **params): + raise NotImplementedError(self._not_implemented) + + def _logccdf_complement(self, x, **params): + return _log1mexp(self._logcdf_dispatch(x, **params)) + + def _logccdf_logexp(self, x, **params): + return np.log(self._ccdf_dispatch(x, **params)) + + def _logccdf_logexp_safe(self, x, **params): + out = self._logccdf_logexp(x, **params) + mask = np.isinf(out) + if np.any(mask): + params_mask = {key: np.broadcast_to(val, mask.shape)[mask] + for key, val in params.items()} + out = np.asarray(out) + out[mask] = self._logccdf_quadrature(x[mask], **params_mask) + return out[()] + + def _logccdf_quadrature(self, x, **params): + _, b = self._support(**params) + return self._quadrature(self._logpxf_dispatch, limits=(x, b), + params=params, log=True) + + def ccdf(self, x, y=None, /, *, method=None): + if y is None: + return self._ccdf1(x, method=method) + else: + return self._ccdf2(x, y, method=method) + + @_cdf2_input_validation + def _ccdf2(self, x, y, *, method): + return self._ccdf2_dispatch(x, y, method=method, **self._parameters) + + @_dispatch + def _ccdf2_dispatch(self, x, y, *, method=None, **params): + if self._overrides('_ccdf2_formula'): + method = self._ccdf2_formula + else: + method = self._ccdf2_addition + return method + + def _ccdf2_formula(self, x, y, **params): + raise NotImplementedError(self._not_implemented) + + def _ccdf2_addition(self, x, y, **params): + cdf_x = self._cdf_dispatch(x, **params) + ccdf_y = self._ccdf_dispatch(y, **params) + # even if x > y, cdf(x, y) + ccdf(x,y) sums to 1 + return cdf_x + ccdf_y + + @_set_invalid_nan + def _ccdf1(self, x, *, method): + return self._ccdf_dispatch(x, method=method, **self._parameters) + + @_dispatch + def _ccdf_dispatch(self, x, method=None, **params): + if self._overrides('_ccdf_formula'): + method = self._ccdf_formula + elif self._overrides('_logccdf_formula'): + method = self._ccdf_logexp + elif self._overrides('_cdf_formula'): + method = self._ccdf_complement_safe + else: + method = self._ccdf_quadrature + return method + + def _ccdf_formula(self, x, **params): + raise NotImplementedError(self._not_implemented) + + def _ccdf_logexp(self, x, **params): + return np.exp(self._logccdf_dispatch(x, **params)) + + def _ccdf_complement(self, x, **params): + return 1 - self._cdf_dispatch(x, **params) + + def _ccdf_complement_safe(self, x, **params): + cdf = self._cdf_dispatch(x, **params) + out = 1 - cdf + eps = np.finfo(self._dtype).eps + tol = self.tol if not _isnull(self.tol) else np.sqrt(eps) + mask = tol * out < np.spacing(cdf) + if np.any(mask): + params_mask = {key: np.broadcast_to(val, mask.shape)[mask] + for key, val in params.items()} + out = np.asarray(out) + out[mask] = self._ccdf_quadrature(x[mask], **params_mask) + return out[()] + + def _ccdf_quadrature(self, x, **params): + _, b = self._support(**params) + return self._quadrature(self._pxf_dispatch, limits=(x, b), + params=params) + + ## Inverse cumulative distribution functions + + @_set_invalid_nan + def ilogcdf(self, logp, /, *, method=None): + return self._ilogcdf_dispatch(logp, method=method, **self._parameters) + + @_dispatch + def _ilogcdf_dispatch(self, x, method=None, **params): + if self._overrides('_ilogcdf_formula'): + method = self._ilogcdf_formula + elif self._overrides('_ilogccdf_formula'): + method = self._ilogcdf_complement + else: + method = self._ilogcdf_inversion + return method + + def _ilogcdf_formula(self, x, **params): + raise NotImplementedError(self._not_implemented) + + def _ilogcdf_complement(self, x, **params): + return self._ilogccdf_dispatch(_log1mexp(x), **params) + + def _ilogcdf_inversion(self, x, **params): + return self._solve_bounded_continuous(self._logcdf_dispatch, x, params=params) + + @_set_invalid_nan + def icdf(self, p, /, *, method=None): + return self._icdf_dispatch(p, method=method, **self._parameters) + + @_dispatch + def _icdf_dispatch(self, x, method=None, **params): + if self._overrides('_icdf_formula'): + method = self._icdf_formula + elif self._overrides('_iccdf_formula'): + method = self._icdf_complement_safe + else: + method = self._icdf_inversion + return method + + def _icdf_formula(self, x, **params): + raise NotImplementedError(self._not_implemented) + + def _icdf_complement(self, x, **params): + return self._iccdf_dispatch(1 - x, **params) + + def _icdf_complement_safe(self, x, **params): + out = self._icdf_complement(x, **params) + eps = np.finfo(self._dtype).eps + tol = self.tol if not _isnull(self.tol) else np.sqrt(eps) + mask = tol * x < np.spacing(1 - x) + if np.any(mask): + params_mask = {key: np.broadcast_to(val, mask.shape)[mask] + for key, val in params.items()} + out = np.asarray(out) + out[mask] = self._icdf_inversion(x[mask], **params_mask) + return out[()] + + def _icdf_inversion(self, x, **params): + return self._solve_bounded_continuous(self._cdf_dispatch, x, params=params) + + @_set_invalid_nan + def ilogccdf(self, logp, /, *, method=None): + return self._ilogccdf_dispatch(logp, method=method, **self._parameters) + + @_dispatch + def _ilogccdf_dispatch(self, x, method=None, **params): + if self._overrides('_ilogccdf_formula'): + method = self._ilogccdf_formula + elif self._overrides('_ilogcdf_formula'): + method = self._ilogccdf_complement + else: + method = self._ilogccdf_inversion + return method + + def _ilogccdf_formula(self, x, **params): + raise NotImplementedError(self._not_implemented) + + def _ilogccdf_complement(self, x, **params): + return self._ilogcdf_dispatch(_log1mexp(x), **params) + + def _ilogccdf_inversion(self, x, **params): + return self._solve_bounded_continuous(self._logccdf_dispatch, x, params=params) + + @_set_invalid_nan + def iccdf(self, p, /, *, method=None): + return self._iccdf_dispatch(p, method=method, **self._parameters) + + @_dispatch + def _iccdf_dispatch(self, x, method=None, **params): + if self._overrides('_iccdf_formula'): + method = self._iccdf_formula + elif self._overrides('_icdf_formula'): + method = self._iccdf_complement_safe + else: + method = self._iccdf_inversion + return method + + def _iccdf_formula(self, x, **params): + raise NotImplementedError(self._not_implemented) + + def _iccdf_complement(self, x, **params): + return self._icdf_dispatch(1 - x, **params) + + def _iccdf_complement_safe(self, x, **params): + out = self._iccdf_complement(x, **params) + eps = np.finfo(self._dtype).eps + tol = self.tol if not _isnull(self.tol) else np.sqrt(eps) + mask = tol * x < np.spacing(1 - x) + if np.any(mask): + params_mask = {key: np.broadcast_to(val, mask.shape)[mask] + for key, val in params.items()} + out = np.asarray(out) + out[mask] = self._iccdf_inversion(x[mask], **params_mask) + return out[()] + + def _iccdf_inversion(self, x, **params): + return self._solve_bounded_continuous(self._ccdf_dispatch, x, params=params) + + ### Sampling Functions + # The following functions for drawing samples from the distribution are + # exposed via a public method that accepts one positional argument - the + # shape of the sample - and keyword options (but not distribution + # parameters). + # sample + # ~~qmc_sample~~ built into sample now + # + # Common keyword options include: + # method - a string that indicates which method should be used to compute + # the quantity (e.g. a formula or numerical integration). + # rng - the NumPy Generator/SciPy QMCEngine object to used for drawing numbers. + # + # Input/output validation is included in each function, since there is + # little code to be shared. + # These are the methods meant to be called by users. + # + # Each public method calls a private "dispatch" method that + # determines which "method" (strategy for calculating the desired quantity) + # to use by default and, via the `@_dispatch` decorator, calls the + # method and computes the result. + # Each dispatch method can designate the responsibility of sampling to any + # of several "implementation" methods. These methods accept only + # `**params`, the parameter dictionary passed from the public method via + # the "dispatch" method. + # See the note corresponding with the "Distribution Parameters" for more + # information. + + def sample(self, shape=(), *, method=None, rng=None): + # needs output validation to ensure that developer returns correct + # dtype and shape + sample_shape = (shape,) if not np.iterable(shape) else tuple(shape) + full_shape = sample_shape + self._shape + rng = np.random.default_rng(rng) if not isinstance(rng, qmc.QMCEngine) else rng + res = self._sample_dispatch(full_shape, method=method, rng=rng, + **self._parameters) + + return res.astype(self._dtype, copy=False) + + @_dispatch + def _sample_dispatch(self, full_shape, *, method, rng, **params): + # make sure that tests catch if sample is 0d array + if self._overrides('_sample_formula') and not isinstance(rng, qmc.QMCEngine): + method = self._sample_formula + else: + method = self._sample_inverse_transform + return method + + def _sample_formula(self, full_shape, *, rng, **params): + raise NotImplementedError(self._not_implemented) + + def _sample_inverse_transform(self, full_shape, *, rng, **params): + if isinstance(rng, qmc.QMCEngine): + uniform = self._qmc_uniform(full_shape, qrng=rng, **params) + else: + uniform = rng.random(size=full_shape, dtype=self._dtype) + return self._icdf_dispatch(uniform, **params) + + def _qmc_uniform(self, full_shape, *, qrng, **params): + # Generate QMC uniform sample(s) on unit interval with specified shape; + # if `sample_shape != ()`, then each slice along axis 0 is independent. + + sample_shape = full_shape[:len(full_shape)-len(self._shape)] + # Determine the number of independent sequences and the length of each. + n_low_discrepancy = sample_shape[0] if sample_shape else 1 + n_independent = math.prod(full_shape[1:] if sample_shape else full_shape) + + # For each independent sequence, we'll need a new QRNG of the appropriate class + # with its own RNG. (If scramble=False, we don't really need all the separate + # rngs, but I'm not going to add a special code path right now.) + rngs = _rng_spawn(qrng.rng, n_independent) + qrng_class = qrng.__class__ + kwargs = dict(d=1, scramble=qrng.scramble, optimization=qrng._optimization) + if isinstance(qrng, qmc.Sobol): + kwargs['bits'] = qrng.bits + + # Draw uniform low-discrepancy sequences scrambled with each RNG + uniforms = [] + for rng in rngs: + qrng = qrng_class(seed=rng, **kwargs) + uniform = qrng.random(n_low_discrepancy) + uniform = uniform.reshape(n_low_discrepancy if sample_shape else ())[()] + uniforms.append(uniform) + + # Reorder the axes and ensure that the shape is correct + uniform = np.moveaxis(np.stack(uniforms), -1, 0) if uniforms else np.asarray([]) + return uniform.reshape(full_shape) + + ### Moments + # The `moment` method accepts two positional arguments - the order and kind + # (raw, central, or standard) of the moment - and a keyword option: + # method - a string that indicates which method should be used to compute + # the quantity (e.g. a formula or numerical integration). + # Like the distribution properties, input/output validation is provided by + # the `_set_invalid_nan_property` decorator. + # + # Unlike most public methods above, `moment` dispatches to one of three + # private methods - one for each 'kind'. Like most *public* methods above, + # each of these private methods calls a private "dispatch" method that + # determines which "method" (strategy for calculating the desired quantity) + # to use. Also, each dispatch method can designate the responsibility + # computing the moment to one of several "implementation" methods. + # Unlike the dispatch methods above, however, the `@_dispatch` decorator + # is not used, and both logic and method calls are included in the function + # itself. + # Instead of determining which method will be used based solely on the + # implementation methods available and calling only the corresponding + # implementation method, *all* the implementation methods are called + # in sequence until one returns the desired information. When an + # implementation methods cannot provide the requested information, it + # returns the object None (which is distinct from arrays with NaNs or infs, + # which are valid values of moments). + # The reason for this approach is that although formulae for the first + # few moments of a distribution may be found, general formulae that work + # for all orders are not always easy to find. This approach allows the + # developer to write "formula" implementation functions that return the + # desired moment when it is available and None otherwise. + # + # Note that the first implementation method called is a cache. This is + # important because lower-order moments are often needed to compute + # higher moments from formulae, so we eliminate redundant calculations + # when moments of several orders are needed. + + @cached_property + def _moment_methods(self): + return {'cache', 'formula', 'transform', + 'normalize', 'general', 'quadrature'} + + @cached_property + def _lmoment_methods(self): + return {'cache', 'formula', 'general', 'order_statistics', 'quadrature_icdf'} + + @property + def _zero(self): + return self._constants()[0] + + @property + def _one(self): + return self._constants()[1] + + def _constants(self): + if self._constant_cache is not None: + return self._constant_cache + + constants = self._preserve_type([0, 1]) + + if self.cache_policy != _NO_CACHE: + self._constant_cache = constants + + return constants + + @_set_invalid_nan_property + def moment(self, order=1, kind='raw', *, method=None): + kinds = {'raw': self._moment_raw, + 'central': self._moment_central, + 'standardized': self._moment_standardized} + order = self._validate_order(order) + kind = self._validate_kind(kind, kinds) + moment_kind = kinds[kind] + return moment_kind(order, method=method) + + def _moment_raw(self, order=1, *, method=None): + """Raw distribution moment about the origin.""" + # Consider exposing the point about which moments are taken as an + # option. This is easy to support, since `_moment_transform_center` + # does all the work. + methods = self._moment_methods if method is None else {method} + return self._moment_raw_dispatch(order, methods=methods, **self._parameters) + + def _moment_raw_dispatch(self, order, *, methods, **params): + moment = None + + if 'cache' in methods: + moment = self._moment_raw_cache.get(order, None) + + if moment is None and 'formula' in methods: + moment = self._moment_raw_formula(order, **params) + + if moment is None and 'transform' in methods and order > 1: + moment = self._moment_raw_transform(order, **params) + + if moment is None and 'general' in methods: + moment = self._moment_raw_general(order, **params) + + if moment is None and 'quadrature' in methods: + moment = self._moment_from_pxf(order, center=self._zero, **params) + + if moment is None and 'quadrature_icdf' in methods: + moment = self._moment_integrate_icdf(order, center=self._zero, **params) + + if moment is not None and self.cache_policy != _NO_CACHE: + self._moment_raw_cache[order] = moment + + return moment + + def _moment_raw_formula(self, order, **params): + return None + + def _moment_raw_transform(self, order, **params): + central_moments = [] + for i in range(int(order) + 1): + methods = {'cache', 'formula', 'normalize', 'general'} + moment_i = self._moment_central_dispatch(order=i, methods=methods, **params) + if moment_i is None: + return None + central_moments.append(moment_i) + + # Doesn't make sense to get the mean by "transform", since that's + # how we got here. Questionable whether 'quadrature' should be here. + mean_methods = {'cache', 'formula', 'quadrature'} + mean = self._moment_raw_dispatch(self._one, methods=mean_methods, **params) + if mean is None: + return None + + moment = self._moment_transform_center(order, central_moments, mean, self._zero) + return moment + + def _moment_raw_general(self, order, **params): + # This is the only general formula for a raw moment of a probability + # distribution + return self._one if order == 0 else None + + def _moment_central(self, order=1, *, method=None): + """Distribution moment about the mean.""" + methods = self._moment_methods if method is None else {method} + return self._moment_central_dispatch(order, methods=methods, **self._parameters) + + def _moment_central_dispatch(self, order, *, methods, **params): + moment = None + + if 'cache' in methods: + moment = self._moment_central_cache.get(order, None) + + if moment is None and 'formula' in methods: + moment = self._moment_central_formula(order, **params) + + if moment is None and 'transform' in methods: + moment = self._moment_central_transform(order, **params) + + if moment is None and 'normalize' in methods and order > 2: + moment = self._moment_central_normalize(order, **params) + + if moment is None and 'general' in methods: + moment = self._moment_central_general(order, **params) + + if moment is None and 'quadrature' in methods: + mean = self._moment_raw_dispatch(self._one, **params, + methods=self._moment_methods) + moment = self._moment_from_pxf(order, center=mean, **params) + + if moment is None and 'quadrature_icdf' in methods: + mean = self._moment_raw_dispatch(self._one, **params, + methods=self._moment_methods) + moment = self._moment_integrate_icdf(order, center=mean, **params) + + if moment is not None and self.cache_policy != _NO_CACHE: + self._moment_central_cache[order] = moment + + return moment + + def _moment_central_formula(self, order, **params): + return None + + def _moment_central_transform(self, order, **params): + + raw_moments = [] + for i in range(int(order) + 1): + methods = {'cache', 'formula', 'general'} + moment_i = self._moment_raw_dispatch(order=i, methods=methods, **params) + if moment_i is None: + return None + raw_moments.append(moment_i) + + mean_methods = self._moment_methods + mean = self._moment_raw_dispatch(self._one, methods=mean_methods, **params) + + moment = self._moment_transform_center(order, raw_moments, self._zero, mean) + return moment + + def _moment_central_normalize(self, order, **params): + methods = {'cache', 'formula', 'general'} + standard_moment = self._moment_standardized_dispatch(order, **params, + methods=methods) + if standard_moment is None: + return None + var = self._moment_central_dispatch(2, methods=self._moment_methods, **params) + return standard_moment*var**(order/2) + + def _moment_central_general(self, order, **params): + general_central_moments = {0: self._one, 1: self._zero} + return general_central_moments.get(order, None) + + def _moment_standardized(self, order=1, *, method=None): + """Standardized distribution moment.""" + methods = self._moment_methods if method is None else {method} + return self._moment_standardized_dispatch(order, methods=methods, + **self._parameters) + + def _moment_standardized_dispatch(self, order, *, methods, **params): + moment = None + + if 'cache' in methods: + moment = self._moment_standardized_cache.get(order, None) + + if moment is None and 'formula' in methods: + moment = self._moment_standardized_formula(order, **params) + + if moment is None and 'normalize' in methods: + moment = self._moment_standardized_normalize(order, False, **params) + + if moment is None and 'general' in methods: + moment = self._moment_standardized_general(order, **params) + + if moment is None and 'normalize' in methods: + moment = self._moment_standardized_normalize(order, True, **params) + + if moment is not None and self.cache_policy != _NO_CACHE: + self._moment_standardized_cache[order] = moment + + return moment + + def _moment_standardized_formula(self, order, **params): + return None + + def _moment_standardized_normalize(self, order, use_quadrature, **params): + methods = ({'quadrature'} if use_quadrature + else {'cache', 'formula', 'transform'}) + central_moment = self._moment_central_dispatch(order, **params, + methods=methods) + if central_moment is None: + return None + var = self._moment_central_dispatch(2, methods=self._moment_methods, + **params) + return central_moment/var**(order/2) + + def _moment_standardized_general(self, order, **params): + general_standard_moments = {0: self._one, 1: self._zero, 2: self._one} + return general_standard_moments.get(order, None) + + def _moment_from_pxf(self, order, center, **params): + def integrand(x, order, center, **params): + pxf = self._pxf_dispatch(x, **params) + return pxf*(x-center)**order + return self._quadrature(integrand, args=(order, center), params=params) + + def _moment_integrate_icdf(self, order, center, **params): + def integrand(x, order, center, **params): + x = self._icdf_dispatch(x, **params) + return (x-center)**order + return self._quadrature(integrand, limits=(0., 1.), + args=(order, center), params=params) + + def _moment_transform_center(self, order, moment_as, a, b): + a, b, *moment_as = np.broadcast_arrays(a, b, *moment_as) + n = order + i = np.arange(n+1).reshape([-1]+[1]*a.ndim) # orthogonal to other axes + i = self._preserve_type(i) + n_choose_i = special.binom(n, i) + with np.errstate(invalid='ignore'): # can happen with infinite moment + moment_b = np.sum(n_choose_i*moment_as*(a-b)**(n-i), axis=0) + return moment_b + + ### L-Moments + + @_set_invalid_nan_property + def lmoment(self, order=1, *, standardize=True, method=None): + order = self._validate_order(order, fname='lmoment', min_order=1) + return self._lmoment(order, standardize=standardize, method=method) + + def _lmoment(self, order, *, standardize, method): + methods = self._lmoment_methods if method is None else {method} + + lmoment = self._lmoment_dispatch(order, methods=methods, **self._parameters) + if lmoment is None: + return None + + if standardize and order >= 3 and lmoment is not None: + lscale = self._lmoment_dispatch(2, methods=methods, **self._parameters) + if lscale is None: + return None + lmoment = lmoment / lscale + + return lmoment + + def _lmoment_dispatch(self, order, *, methods, **params): + lmoment = None + + if 'cache' in methods: + lmoment = self._lmoment_cache.get(order, None) + + if lmoment is None and 'formula' in methods: + lmoment = self._lmoment_formula(order, **params) + + if lmoment is None and 'general' in methods: + lmoment = self._lmoment_general(order, **params) + + if lmoment is None and 'quadrature_icdf' in methods and ( + self._overrides('_icdf_formula') or self._overrides('_iccdf_formula')): + lmoment = self._lmoment_integrate_icdf(order, **params) + + if lmoment is None and 'order_statistics' in methods: + lmoment = self._lmoment_from_order_statistics(order, **params) + + if lmoment is None and 'quadrature_icdf' in methods: + lmoment = self._lmoment_integrate_icdf(order, **params) + + if lmoment is not None and self.cache_policy != _NO_CACHE: + self._lmoment_cache[order] = lmoment + + return lmoment + + def _lmoment_formula(self, order, **params): + return None + + def _lmoment_general(self, order, **params): + return self.mean() if order == 1 else None + + def _lmoment_from_order_statistics(self, order, **params): + k = np.arange(order) + k = xpx.atleast_nd(k, ndim=self._ndim + 1).T + E = order_statistic(self, r=order-k, n=order).mean() + bc = special.binom(order-1, k) + return np.sum((-1)**k * bc * E, axis=0) / order + + def _lmoment_integrate_icdf(self, order, **params): + def integrand(p, **params): + x = self._icdf_dispatch(p, **params) + P = special.eval_sh_legendre(order - 1, p) + return x * P + return self._quadrature(integrand, limits=(0., 1.), params=params) + + ### Convenience + + def plot(self, x='x', y=None, *, t=None, ax=None): + r"""Plot a function of the distribution. + + Convenience function for quick visualization of the distribution + underlying the random variable. + + Parameters + ---------- + x, y : str, optional + String indicating the quantities to be used as the abscissa and + ordinate (horizontal and vertical coordinates), respectively. + Defaults are ``'x'`` (the domain of the random variable) and either + ``'pdf'`` (the probability density function) (continuous) or + ``'pdf'`` (the probability density function) (discrete). + Valid values are: + 'x', 'pdf', 'pmf', 'cdf', 'ccdf', 'icdf', 'iccdf', 'logpdf', 'logpmf', + 'logcdf', 'logccdf', 'ilogcdf', 'ilogccdf'. + t : 3-tuple of (str, float, float), optional + Tuple indicating the limits within which the quantities are plotted. + The default is ``('cdf', 0.0005, 0.9995)`` if the domain is infinite, + indicating that the central 99.9% of the distribution is to be shown; + otherwise, endpoints of the support are used where they are finite. + Valid values are: + 'x', 'cdf', 'ccdf', 'icdf', 'iccdf', 'logcdf', 'logccdf', + 'ilogcdf', 'ilogccdf'. + ax : `matplotlib.axes`, optional + Axes on which to generate the plot. If not provided, use the + current axes. + + Returns + ------- + ax : `matplotlib.axes` + Axes on which the plot was generated. + The plot can be customized by manipulating this object. + + Examples + -------- + Instantiate a distribution with the desired parameters: + + >>> import numpy as np + >>> import matplotlib.pyplot as plt + >>> from scipy import stats + >>> X = stats.Normal(mu=1., sigma=2.) + + Plot the PDF over the central 99.9% of the distribution. + Compare against a histogram of a random sample. + + >>> ax = X.plot() + >>> sample = X.sample(10000) + >>> ax.hist(sample, density=True, bins=50, alpha=0.5) + >>> plt.show() + + Plot ``logpdf(x)`` as a function of ``x`` in the left tail, + where the log of the CDF is between -10 and ``np.log(0.5)``. + + >>> X.plot('x', 'logpdf', t=('logcdf', -10, np.log(0.5))) + >>> plt.show() + + Plot the PDF of the normal distribution as a function of the + CDF for various values of the scale parameter. + + >>> X = stats.Normal(mu=0., sigma=[0.5, 1., 2]) + >>> X.plot('cdf', 'pdf') + >>> plt.show() + + """ + + # Strategy: given t limits, get quantile limits. Form grid of + # quantiles, compute requested x and y at quantiles, and plot. + # Currently, the grid of quantiles is always linearly spaced. + # Instead of always computing linearly-spaced quantiles, it + # would be better to choose: + # a) quantiles or probabilities + # b) linearly or logarithmically spaced + # based on the specified `t`. + # TODO: + # - smart spacing of points + # - when the parameters of the distribution are an array, + # use the full range of abscissae for all curves + + discrete = isinstance(self, DiscreteDistribution) + t_is_quantile = {'x', 'icdf', 'iccdf', 'ilogcdf', 'ilogccdf'} + t_is_probability = {'cdf', 'ccdf', 'logcdf', 'logccdf'} + valid_t = t_is_quantile.union(t_is_probability) + valid_xy = valid_t.union({'pdf', 'logpdf', 'pmf', 'logpmf'}) + y_default = 'pmf' if discrete else 'pdf' + y = y_default if y is None else y + + ndim = self._ndim + x_name, y_name = x, y + t_name = 'cdf' if t is None else t[0] + + a, b = self.support() + tliml_default = 0 if np.all(np.isfinite(a)) else 0.0005 + tliml = tliml_default if t is None else t[1] + tlimr_default = 1 if np.all(np.isfinite(b)) else 0.9995 + tlimr = tlimr_default if t is None else t[2] + tlim = np.asarray([tliml, tlimr]) + tlim = tlim[:, np.newaxis] if ndim else tlim + + # pdf/logpdf are not valid for `t` because we can't easily invert them + message = (f'Argument `t` of `{self.__class__.__name__}.plot` "' + f'must be one of {valid_t}') + if t_name not in valid_t: + raise ValueError(message) + + message = (f'Argument `x` of `{self.__class__.__name__}.plot` "' + f'must be one of {valid_xy}') + if x_name not in valid_xy: + raise ValueError(message) + + message = (f'Argument `y` of `{self.__class__.__name__}.plot` "' + f'must be one of {valid_xy}') + if y_name not in valid_xy: + raise ValueError(message) + + # This could just be a warning + message = (f'`{self.__class__.__name__}.plot` was called on a random ' + 'variable with at least one invalid shape parameters. When ' + 'a parameter is invalid, no plot can be shown.') + if self._any_invalid: + raise ValueError(message) + + # We could automatically ravel, but do we want to? For now, raise. + message = ("To use `plot`, distribution parameters must be " + "scalars or arrays with one or fewer dimensions.") + if ndim > 1: + raise ValueError(message) + + try: + import matplotlib.pyplot as plt # noqa: F401, E402 + except ModuleNotFoundError as exc: + message = ("`matplotlib` must be installed to use " + f"`{self.__class__.__name__}.plot`.") + raise ModuleNotFoundError(message) from exc + ax = plt.gca() if ax is None else ax + + # get quantile limits given t limits + qlim = tlim if t_name in t_is_quantile else getattr(self, 'i'+t_name)(tlim) + + message = (f"`{self.__class__.__name__}.plot` received invalid input for `t`: " + f"calling {'i'+t_name}({tlim}) produced {qlim}.") + if not np.all(np.isfinite(qlim)): + raise ValueError(message) + + # form quantile grid + if discrete and x_name in t_is_quantile: + # should probably aggregate for large ranges + q = np.arange(np.min(qlim[0]), np.max(qlim[1]) + 1) + q = q[:, np.newaxis] if ndim else q + else: + grid = np.linspace(0, 1, 300) + grid = grid[:, np.newaxis] if ndim else grid + q = qlim[0] + (qlim[1] - qlim[0]) * grid + q = np.round(q) if discrete else q + + # compute requested x and y at quantile grid + x = q if x_name in t_is_quantile else getattr(self, x_name)(q) + y = q if y_name in t_is_quantile else getattr(self, y_name)(q) + + # make plot + x, y = np.broadcast_arrays(x.T, np.atleast_2d(y.T)) + for xi, yi in zip(x, y): # plot is vectorized, but bar/step don't seem to be + if discrete and x_name in t_is_quantile and y_name == 'pmf': + # should this just be a step plot, too? + ax.bar(xi, yi, alpha=np.sqrt(1/y.shape[0])) # alpha heuristic + elif discrete and x_name in t_is_quantile: + values = yi + edges = np.concatenate((xi, [xi[-1]+1])) + ax.stairs(values, edges, baseline=None) + else: + ax.plot(xi, yi) + ax.set_xlabel(f"${x_name}$") + ax.set_ylabel(f"${y_name}$") + ax.set_title(str(self)) + + # only need a legend if distribution has parameters + if len(self._parameters): + label = [] + parameters = self._parameterization.parameters + param_names = list(parameters) + param_arrays = [np.atleast_1d(self._parameters[pname]) + for pname in param_names] + for param_vals in zip(*param_arrays): + assignments = [f"${parameters[name].symbol}$ = {val:.4g}" + for name, val in zip(param_names, param_vals)] + label.append(", ".join(assignments)) + ax.legend(label) + + return ax + + + ### Fitting + # All methods above treat the distribution parameters as fixed, and the + # variable argument may be a quantile or probability. The fitting functions + # are fundamentally different because the quantiles (often observations) + # are considered to be fixed, and the distribution parameters are the + # variables. In a sense, they are like an inverse of the sampling + # functions. + # + # At first glance, it would seem ideal for `fit` to be a classmethod, + # called like `LogUniform.fit(sample=sample)`. + # I tried this. I insisted on it for a while. But if `fit` is a + # classmethod, it cannot call instance methods. If we want to support MLE, + # MPS, MoM, MoLM, then we end up with most of the distribution functions + # above needing to be classmethods, too. All state information, such as + # tolerances and the underlying distribution of `ShiftedScaledDistribution` + # and `OrderStatisticDistribution`, would need to be passed into all + # methods. And I'm not really sure how we would call `fit` as a + # classmethod of a transformed distribution - maybe + # ShiftedScaledDistribution.fit would accept the class of the + # shifted/scaled distribution as an argument? + # + # In any case, it was a conscious decision for the infrastructure to + # treat the parameters as "fixed" and the quantile/percentile arguments + # as "variable". There are a lot of advantages to this structure, and I + # don't think the fact that a few methods reverse the fixed and variable + # quantities should make us question that choice. It can still accommodate + # these methods reasonably efficiently. + + +class ContinuousDistribution(UnivariateDistribution): + def _overrides(self, method_name): + if method_name in {'_logpmf_formula', '_pmf_formula'}: + return True + return super()._overrides(method_name) + + def _pmf_formula(self, x, **params): + return np.zeros_like(x) + + def _logpmf_formula(self, x, **params): + return np.full_like(x, -np.inf) + + def _pxf_dispatch(self, x, *, method=None, **params): + return self._pdf_dispatch(x, method=method, **params) + + def _logpxf_dispatch(self, x, *, method=None, **params): + return self._logpdf_dispatch(x, method=method, **params) + + def _solve_bounded_continuous(self, func, p, params, xatol=None): + return self._solve_bounded(func, p, params=params, xatol=xatol).x + + +class DiscreteDistribution(UnivariateDistribution): + def _overrides(self, method_name): + if method_name in {'_logpdf_formula', '_pdf_formula'}: + return True + return super()._overrides(method_name) + + def _logpdf_formula(self, x, **params): + if params: + p = next(iter(params.values())) + nan_result = np.isnan(x) | np.isnan(p) + else: + nan_result = np.isnan(x) + return np.where(nan_result, np.nan, np.inf) + + def _pdf_formula(self, x, **params): + if params: + p = next(iter(params.values())) + nan_result = np.isnan(x) | np.isnan(p) + else: + nan_result = np.isnan(x) + return np.where(nan_result, np.nan, np.inf) + + def _pxf_dispatch(self, x, *, method=None, **params): + return self._pmf_dispatch(x, method=method, **params) + + def _logpxf_dispatch(self, x, *, method=None, **params): + return self._logpmf_dispatch(x, method=method, **params) + + def _cdf_quadrature(self, x, **params): + return super()._cdf_quadrature(np.floor(x), **params) + + def _logcdf_quadrature(self, x, **params): + return super()._logcdf_quadrature(np.floor(x), **params) + + def _ccdf_quadrature(self, x, **params): + return super()._ccdf_quadrature(np.floor(x + 1), **params) + + def _logccdf_quadrature(self, x, **params): + return super()._logccdf_quadrature(np.floor(x + 1), **params) + + def _cdf2(self, x, y, *, method): + raise NotImplementedError( + "Two argument cdf functions are currently only supported for " + "continuous distributions.") + + def _ccdf2(self, x, y, *, method): + raise NotImplementedError( + "Two argument cdf functions are currently only supported for " + "continuous distributions.") + + def _logcdf2(self, x, y, *, method): + raise NotImplementedError( + "Two argument cdf functions are currently only supported for " + "continuous distributions.") + + def _logccdf2(self, x, y, *, method): + raise NotImplementedError( + "Two argument cdf functions are currently only supported for " + "continuous distributions.") + + def _lmoment(self, order, *, standardize, method): + raise NotImplementedError( + "L-moments are currently available only " + "for continuous distributions.") + + def _solve_bounded_discrete(self, func, p, params, comp): + # We're trying to solve one of these two problems: + # a) find the smallest integer x* within the support s.t. F(x*) >= p + # b) find the smallest integer x* within the support s.t. G(x*) = 1 - F(x*) <= p + # Our approach is to solve a continuous version of the problem that narrows the + # solution down to an integer x s.t. either x* = x or x* = x + 1. At the end, + # we'll choose between them. + + # First, solve func(x) == p where func is a continuous, monotone interpolant + # of either the monotone increasing F or monotone decreasing G. + res = self._solve_bounded(func, p, params=params, xatol=0.9) + # Here, `_solve_bounded` can terminate for one of three reasons: + # 1. `func(res.x) == p` (`fatol = 0` is satisfied), + # 2. `res.xl` and `res.xr` bracket the root and `|res.xr - res.xl| <= xatol`, or + # 3. There is no solution within the support. + # There are several possible strategies for using `res.xl`, `res.x`, and/or + # `res.xr` to find a solution to the original, discrete problem. Here is ours. + + # Consider case 2a. Because F is an increasing function, we know + # that F(xr) >= p (and F(xl) <= p), so F(floor(xr) + 1) >= p. + # F(floor(xr)) *may* be >= p, but we can't know until we evaluate it. + # F(floor(xr) - 1) < p (strictly) because floor(xr) - 1 < xl and F decreases + # monotonically as the argument decreases. So we choose x = floor(xr), and + # later we'll choose between x* = x and x* = x + 1. + x = np.asarray(np.floor(res.bracket[-1])) + # This is also suitable for case 2b. Because G is a *decreasing* function, we + # know that G(xr) <= p (and G(xl) >= p), so G(floor(xr) + 1) <= p. + # G(floor(xr)) *may* be <= p, but we can't know until we evaluate it. + # G(floor(xr) - 1) > p (strictly) because floor(xr) - 1 < xl and G increases + # as the argument decreases. So we would still want to choose x = floor(xr), and + # later we'll choose between x* = x and x* = x + 1. + + # Now we consider case 1a/b. In this case, `res.x` solved the equation + # *exactly*, so the algorithm may have terminated before the bracket is tight + # enough to rely on `res.xr`. If `res.x` happens to be integral, `res.x` is + # the solution to the discrete problem, and floor(res.x) == res.x, so + # floor(res.x) is the solution to the discrete problem. If not: + # a) F(floor(res.x)) < p (strictly) and F(floor(res.x) + 1) > p (strictly). So + # floor(res.x) + 1 is the solution to the discrete problem. + # b) G(floor(res.x)) > p (strictly) and G(floor(res.x) + 1) < p (strictly). So + # floor(res.x) + 1 is again the solution to the discrete problem. + # Either way, we can choose x = res.x, and at the end we'll choose between + # x* = x and x* = x + 1. + mask = res.f_x == 0 + x[mask] = np.floor(res.x[mask]) + + # For case 3, let xmin be the left endpoint of the support, and note that in + # general, F(xmin) > 0 and G(xmin) < 1. Therefore it is possible that: + # a) F(x) > p for all x in the support (e.g. because p ~ 0) + # a) G(x) < p for all x in the support (e.g. because p ~ 1) + # In these cases, `_solve_bounded` would fail to find a root of the continuous + # equation above, but the solution to the original, discrete problem is the left + # endpoint of the support. + # This case is handled before we get to this function; otherwise, + # `_solve_bounded` may spin its wheels for a long time in vain. + + # Now, we choose between x* = x and x* = x + 1: if func(x) satisfies the + # comparison `comp` (>= for cdf, <= for ccdf), the solution is x* = x; + # otherwise the solution must be x* = x + 1. + f = func(x, **params) + x = np.where(comp(f, p), x, x + 1.0) + x[np.isnan(f)] = np.nan # needed? why would func(x) be NaN within support? + + return x + + def _base_discrete_inversion(self, p, func, comp, /, **params): + # For discrete distributions, icdf(p) is defined as the minimum integer x* + # within the support such that F(x*) >= p; iccdf(p) is the minimum integer x* + # within the support such that G(x*) <= p. + + # Identify where the solution is xmin. + # (See rationale in `_solve_bounded_discrete`.) + xmin, xmax = self._support(**params) + p, xmin, _ = np.broadcast_arrays(p, xmin, xmax) + mask = comp(func(xmin, **params), p) + + # Use `apply_where` to perform the inversion only when necessary. + def f1(p, *args): + return self._solve_bounded_discrete( + func, p, params=dict(zip(params.keys(), args)), comp=comp) + + x = xpx.apply_where(~mask, (p, *params.values()), f1, fill_value=xmin) + + # x above may be a finite value even when p is NaN, so the returned value + # should be NaN. We need to handle this as a special case. + x[np.isnan(p)] = np.nan + return x[()] + + def _icdf_inversion(self, x, **params): + return self._base_discrete_inversion(x, self._cdf_dispatch, + np.greater_equal, **params) + + def _ilogcdf_inversion(self, x, **params): + return self._base_discrete_inversion(x, self._logcdf_dispatch, + np.greater_equal, **params) + + def _iccdf_inversion(self, x, **params): + return self._base_discrete_inversion(x, self._ccdf_dispatch, + np.less_equal, **params) + + def _ilogccdf_inversion(self, x, **params): + return self._base_discrete_inversion(x, self._logccdf_dispatch, + np.less_equal, **params) + + def _mode_optimization(self, **params): + # If `x` is the true mode of a unimodal continuous function, we can find + # the mode among integers by rounding in each direction and checking + # which is better. If the difference between `x` and the nearest integer + # is less than `xatol`, the computed value of `x` may end up on the wrong + # side of the nearest integer. Setting `xatol=0.5` guarantees that at most + # three integers need to be checked, the two nearest integers, ``floor(x)`` + # and ``round(x)`` and the nearest integer other than these. + x = super()._mode_optimization(xatol=0.5, **params) + low, high = self.support() + xl, xr = np.floor(x), np.ceil(x) + nearest = np.round(x) + # Clip to stay within support. There will be redundant calculation + # when clipping since `xo` will be one of `xl` or `xr`, but let's + # keep the implementation simple for now. + xo = np.clip(nearest + np.copysign(1, nearest - x), low, high) + x = np.stack([xl, xo, xr]) + idx = np.argmax(self._pmf_dispatch(x, **params), axis=0) + return np.choose(idx, [xl, xo, xr]) + + def _logentropy_quadrature(self, **params): + def logintegrand(x, **params): + logpmf = self._logpmf_dispatch(x, **params) + # Entropy summand is -pmf*log(pmf), so log-entropy summand is + # logpmf + log(logpmf) + pi*j. But pmf is always between 0 and 1, + # so logpmf is always negative, and so log(logpmf) = log(-logpmf) + pi*j. + # The two imaginary components "cancel" each other out (which we would + # expect because each term of the entropy summand is positive). + return np.where(np.isfinite(logpmf), logpmf + np.log(-logpmf), -np.inf) + return self._quadrature(logintegrand, params=params, log=True) + + +# Special case the names of some new-style distributions in `make_distribution` +_distribution_names = { + # Continuous + 'argus': 'ARGUS', + 'betaprime': 'BetaPrime', + 'chi2': 'ChiSquared', + 'crystalball': 'CrystalBall', + 'dgamma': 'DoubleGamma', + 'dweibull': 'DoubleWeibull', + 'expon': 'Exponential', + 'exponnorm': 'ExponentiallyModifiedNormal', + 'exponweib': 'ExponentialWeibull', + 'exponpow': 'ExponentialPower', + 'fatiguelife': 'FatigueLife', + 'foldcauchy': 'FoldedCauchy', + 'foldnorm': 'FoldedNormal', + 'genlogistic': 'GeneralizedLogistic', + 'gennorm': 'GeneralizedNormal', + 'genpareto': 'GeneralizedPareto', + 'genexpon': 'GeneralizedExponential', + 'genextreme': 'GeneralizedExtremeValue', + 'gausshyper': 'GaussHypergeometric', + 'gengamma': 'GeneralizedGamma', + 'genhalflogistic': 'GeneralizedHalfLogistic', + 'geninvgauss': 'GeneralizedInverseGaussian', + 'gumbel_r': 'Gumbel', + 'gumbel_l': 'ReflectedGumbel', + 'halfcauchy': 'HalfCauchy', + 'halflogistic': 'HalfLogistic', + 'halfnorm': 'HalfNormal', + 'halfgennorm': 'HalfGeneralizedNormal', + 'hypsecant': 'HyperbolicSecant', + 'invgamma': 'InverseGammma', + 'invgauss': 'InverseGaussian', + 'invweibull': 'InverseWeibull', + 'irwinhall': 'IrwinHall', + 'jf_skew_t': 'JonesFaddySkewT', + 'johnsonsb': 'JohnsonSB', + 'johnsonsu': 'JohnsonSU', + 'ksone': 'KSOneSided', + 'kstwo': 'KSTwoSided', + 'kstwobign': 'KSTwoSidedAsymptotic', + 'laplace_asymmetric': 'LaplaceAsymmetric', + 'levy_l': 'LevyLeft', + 'levy_stable': 'LevyStable', + 'loggamma': 'ExpGamma', # really the Exponential Gamma Distribution + 'loglaplace': 'LogLaplace', + 'lognorm': 'LogNormal', + 'loguniform': 'LogUniform', + 'ncx2': 'NoncentralChiSquared', + 'nct': 'NoncentralT', + 'norm': 'Normal', + 'norminvgauss': 'NormalInverseGaussian', + 'powerlaw': 'PowerLaw', + 'powernorm': 'PowerNormal', + 'rdist': 'R', + 'rel_breitwigner': 'RelativisticBreitWigner', + 'recipinvgauss': 'ReciprocalInverseGaussian', + 'reciprocal': 'LogUniform', + 'semicircular': 'SemiCircular', + 'skewcauchy': 'SkewCauchy', + 'skewnorm': 'SkewNormal', + 'studentized_range': 'StudentizedRange', + 't': 'StudentT', + 'trapezoid': 'Trapezoidal', + 'triang': 'Triangular', + 'truncexpon': 'TruncatedExponential', + 'truncnorm': 'TruncatedNormal', + 'truncpareto': 'TruncatedPareto', + 'truncweibull_min': 'TruncatedWeibull', + 'tukeylambda': 'TukeyLambda', + 'vonmises_line': 'VonMisesLine', + 'weibull_min': 'Weibull', + 'weibull_max': 'ReflectedWeibull', + 'wrapcauchy': 'WrappedCauchyLine', + # Discrete + 'betabinom': 'BetaBinomial', + 'betanbinom': 'BetaNegativeBinomial', + 'dlaplace': 'LaplaceDiscrete', + 'geom': 'Geometric', + 'hypergeom': 'Hypergeometric', + 'logser': 'LogarithmicSeries', + 'nbinom': 'NegativeBinomial', + 'nchypergeom_fisher': 'NoncentralHypergeometricFisher', + 'nchypergeom_wallenius': 'NoncentralHypergeometricWallenius', + 'nhypergeom': 'NegativeHypergeometric', + 'poisson_binom': 'PoissonBinomial', + 'randint': 'UniformDiscrete', + 'yulesimon': 'YuleSimon', + 'zipf': 'Zeta', +} + + +@xp_capabilities(np_only=True) +def make_distribution(dist): + """Generate a `UnivariateDistribution` class from a compatible object. + + The argument may be an instance of `rv_continuous` or an instance of + another class that satisfies the interface described below. + + The returned value is a `ContinuousDistribution` subclass if the input is an + instance of `rv_continuous` or a `DiscreteDistribution` subclass if the input + is an instance of `rv_discrete`. Like any subclass of `UnivariateDistribution`, + it must be instantiated (i.e. by passing all shape parameters as keyword + arguments) before use. Once instantiated, the resulting object will have the + same interface as any other instance of `UnivariateDistribution`; e.g., + `scipy.stats.Normal`, `scipy.stats.Binomial`. + + .. note:: + + `make_distribution` does not work perfectly with all instances of + `rv_continuous`. Known failures include `levy_stable`, `vonmises`, + `hypergeom`, 'nchypergeom_fisher', 'nchypergeom_wallenius', and + `poisson_binom`. Some methods of some distributions will not support + array shape parameters. + + Parameters + ---------- + dist : `rv_continuous` + Instance of `rv_continuous`, `rv_discrete`, or an instance of any class with + the following attributes: + + __make_distribution_version__ : str + A string containing the version number of SciPy in which this interface + is defined. The preferred interface may change in future SciPy versions, + in which case support for an old interface version may be deprecated + and eventually removed. + parameters : dict or tuple + If a dictionary, each key is the name of a parameter, + and the corresponding value is either a dictionary or tuple. + If the value is a dictionary, it may have the following items, with default + values used for entries which aren't present. + + endpoints : tuple, default: (-inf, inf) + A tuple defining the lower and upper endpoints of the domain of the + parameter; allowable values are floats, the name (string) of another + parameter, or a callable taking parameters as keyword only + arguments and returning the numerical value of an endpoint for + given parameter values. + + inclusive : tuple of bool, default: (False, False) + A tuple specifying whether the endpoints are included within the domain + of the parameter. + + typical : tuple, default: ``endpoints`` + Defining endpoints of a typical range of values of a parameter. Can be + used for sampling parameter values for testing. Behaves like the + ``endpoints`` tuple above, and should define a subinterval of the + domain given by ``endpoints``. + + A tuple value ``(a, b)`` associated to a key in the ``parameters`` + dictionary is equivalent to ``{endpoints: (a, b)}``. + + Custom distributions with multiple parameterizations can be defined by + having the ``parameters`` attribute be a tuple of dictionaries with + the structure described above. In this case, ``dist``\'s class must also + define a method ``process_parameters`` to map between the different + parameterizations. It must take all parameters from all parameterizations + as optional keyword arguments and return a dictionary mapping parameters to + values, filling in values from other parameterizations using values from + the supplied parameterization. See example. + + support : dict or tuple + A dictionary describing the support of the distribution or a tuple + describing the endpoints of the support. This behaves identically to + the values of the parameters dict described above. + + The class **must** also define a ``pdf`` method and **may** define methods + ``logentropy``, ``entropy``, ``median``, ``mode``, ``logpdf``, + ``logcdf``, ``cdf``, ``logccdf``, ``ccdf``, + ``ilogcdf``, ``icdf``, ``ilogccdf``, ``iccdf``, + ``moment``, ``lmoment``, and ``sample``. + If defined, these methods must accept the parameters of the distribution as + keyword arguments and also accept any positional-only arguments accepted by + the corresponding method of `ContinuousDistribution`. + When multiple parameterizations are defined, these methods must accept + all parameters from all parameterizations. The ``moment`` method + must accept the ``order`` and ``kind`` arguments by position or keyword, but + may return ``None`` if a formula is not available for the arguments; in this + case, the infrastructure will fall back to a default implementation. + The ``lmoment`` method must accept the ``order`` argument by position or + keyword and return the specified L-moment (not the L-moment ratio), but may + return ``None`` if a formula is not available for the arguments. The + ``sample`` method must accept ``shape`` by position or keyword, but contrary + to the public method of the same name, the argument it receives will be the + *full* shape of the output array - that is, the shape passed to the public + method prepended to the broadcasted shape of random variable parameters. + + Returns + ------- + CustomDistribution : `UnivariateDistribution` + A subclass of `UnivariateDistribution` corresponding with `dist`. The + initializer requires all shape parameters to be passed as keyword arguments + (using the same names as the instance of `rv_continuous`/`rv_discrete`). + + Notes + ----- + The documentation of `UnivariateDistribution` is not rendered. See below for + an example of how to instantiate the class (i.e. pass all shape parameters of + `dist` to the initializer as keyword arguments). Documentation of all methods + is identical to that of `scipy.stats.Normal`. Use ``help`` on the returned + class or its methods for more information. + + Examples + -------- + >>> import numpy as np + >>> import matplotlib.pyplot as plt + >>> from scipy import stats + >>> from scipy import special + + Create a `ContinuousDistribution` from `scipy.stats.loguniform`. + + >>> LogUniform = stats.make_distribution(stats.loguniform) + >>> X = LogUniform(a=1.0, b=3.0) + >>> np.isclose((X + 0.25).median(), stats.loguniform.ppf(0.5, 1, 3, loc=0.25)) + np.True_ + >>> X.plot() + >>> sample = X.sample(10000, rng=np.random.default_rng()) + >>> plt.hist(sample, density=True, bins=30) + >>> plt.legend(('pdf', 'histogram')) + >>> plt.show() + + Create a custom distribution. + + >>> class MyLogUniform: + ... @property + ... def __make_distribution_version__(self): + ... return "1.16.0" + ... + ... @property + ... def parameters(self): + ... return {'a': {'endpoints': (0, np.inf), + ... 'inclusive': (False, False)}, + ... 'b': {'endpoints': ('a', np.inf), + ... 'inclusive': (False, False)}} + ... + ... @property + ... def support(self): + ... return {'endpoints': ('a', 'b'), 'inclusive': (True, True)} + ... + ... def pdf(self, x, a, b): + ... return 1 / (x * (np.log(b)- np.log(a))) + >>> + >>> MyLogUniform = stats.make_distribution(MyLogUniform()) + >>> Y = MyLogUniform(a=1.0, b=3.0) + >>> np.isclose(Y.cdf(2.), X.cdf(2.)) + np.True_ + + Create a custom distribution with variable support. + + >>> class MyUniformCube: + ... @property + ... def __make_distribution_version__(self): + ... return "1.16.0" + ... + ... @property + ... def parameters(self): + ... return {"a": (-np.inf, np.inf), + ... "b": {'endpoints':('a', np.inf), 'inclusive':(True, False)}} + ... + ... @property + ... def support(self): + ... def left(*, a, b): + ... return a**3 + ... + ... def right(*, a, b): + ... return b**3 + ... return (left, right) + ... + ... def pdf(self, x, *, a, b): + ... return 1 / (3*(b - a)*np.cbrt(x)**2) + ... + ... def cdf(self, x, *, a, b): + ... return (np.cbrt(x) - a) / (b - a) + >>> + >>> MyUniformCube = stats.make_distribution(MyUniformCube()) + >>> X = MyUniformCube(a=-2, b=2) + >>> Y = stats.Uniform(a=-2, b=2)**3 + >>> X.support() + (-8.0, 8.0) + >>> np.isclose(X.cdf(2.1), Y.cdf(2.1)) + np.True_ + + Create a custom distribution with multiple parameterizations. Here we create a + custom version of the beta distribution that has an alternative parameterization + in terms of the mean ``mu`` and a dispersion parameter ``nu``. + + >>> class MyBeta: + ... @property + ... def __make_distribution_version__(self): + ... return "1.16.0" + ... + ... @property + ... def parameters(self): + ... return ({"a": (0, np.inf), "b": (0, np.inf)}, + ... {"mu": (0, 1), "nu": (0, np.inf)}) + ... + ... def process_parameters(self, a=None, b=None, mu=None, nu=None): + ... if a is not None and b is not None: + ... nu = a + b + ... mu = a / nu + ... else: + ... a = mu * nu + ... b = nu - a + ... return dict(a=a, b=b, mu=mu, nu=nu) + ... + ... @property + ... def support(self): + ... return {'endpoints': (0, 1)} + ... + ... def pdf(self, x, a, b, mu, nu): + ... return special._ufuncs._beta_pdf(x, a, b) + ... + ... def cdf(self, x, a, b, mu, nu): + ... return special.betainc(a, b, x) + >>> + >>> MyBeta = stats.make_distribution(MyBeta()) + >>> X = MyBeta(a=2.0, b=2.0) + >>> Y = MyBeta(mu=0.5, nu=4.0) + >>> np.isclose(X.pdf(0.3), Y.pdf(0.3)) + np.True_ + + """ + if dist in {stats.levy_stable, stats.vonmises, stats.hypergeom, + stats.nchypergeom_fisher, stats.nchypergeom_wallenius, + stats.poisson_binom}: + raise NotImplementedError(f"`{dist.name}` is not supported.") + + if isinstance(dist, stats.rv_continuous | stats.rv_discrete): + return _make_distribution_rv_generic(dist) + elif getattr(dist, "__make_distribution_version__", "0.0.0") >= "1.16.0": + return _make_distribution_custom(dist) + else: + message = ("The argument must be an instance of `rv_continuous`, " + "`rv_discrete`, or an instance of a class with attribute " + "`__make_distribution_version__ >= 1.16`.") + raise ValueError(message) + +def _make_distribution_rv_generic(dist): + parameters = [] + names = [] + support = getattr(dist, '_support', (dist.a, dist.b)) + for shape_info in dist._shape_info(): + domain = _RealInterval(endpoints=shape_info.endpoints, + inclusive=shape_info.inclusive) + param = _RealParameter(shape_info.name, domain=domain) + parameters.append(param) + names.append(shape_info.name) + + repr_str = _distribution_names.get(dist.name, dist.name.capitalize()) + if isinstance(dist, stats.rv_continuous): + old_class, new_class = stats.rv_continuous, ContinuousDistribution + else: + old_class, new_class = stats.rv_discrete, DiscreteDistribution + + def _overrides(method_name): + return (getattr(dist.__class__, method_name, None) + is not getattr(old_class, method_name, None)) + + if _overrides("_get_support"): + def left(**parameter_values): + a, _ = dist._get_support(**parameter_values) + return np.asarray(a)[()] + + def right(**parameter_values): + _, b = dist._get_support(**parameter_values) + return np.asarray(b)[()] + + endpoints = (left, right) + else: + endpoints = support + + _x_support = _RealInterval(endpoints=endpoints, inclusive=(True, True)) + _x_param = _RealParameter('x', domain=_x_support, typical=(-1, 1)) + + class CustomDistribution(new_class): + _parameterizations = ([_Parameterization(*parameters)] if parameters + else []) + _variable = _x_param + + __class_getitem__ = None + + def __repr__(self): + s = super().__repr__() + return s.replace('CustomDistribution', repr_str) + + def __str__(self): + s = super().__str__() + return s.replace('CustomDistribution', repr_str) + + def _sample_formula(self, full_shape=(), *, rng=None, **kwargs): + return dist._rvs(size=full_shape, random_state=rng, **kwargs) + + def _moment_raw_formula(self, order, **kwargs): + return dist._munp(int(order), **kwargs) + + def _moment_raw_formula_1(self, order, **kwargs): + if order != 1: + return None + return dist._stats(**kwargs)[0] + + def _moment_central_formula(self, order, **kwargs): + if order != 2: + return None + return dist._stats(**kwargs)[1] + + def _moment_standard_formula(self, order, **kwargs): + if order == 3: + if dist._stats_has_moments: + kwargs['moments'] = 's' + return dist._stats(**kwargs)[int(order - 1)] + elif order == 4: + if dist._stats_has_moments: + kwargs['moments'] = 'k' + k = dist._stats(**kwargs)[int(order - 1)] + return k if k is None else k + 3 + else: + return None + + methods = {'_logpdf': '_logpdf_formula', + '_pdf': '_pdf_formula', + '_logpmf': '_logpmf_formula', + '_pmf': '_pmf_formula', + '_logcdf': '_logcdf_formula', + '_cdf': '_cdf_formula', + '_logsf': '_logccdf_formula', + '_sf': '_ccdf_formula', + '_ppf': '_icdf_formula', + '_isf': '_iccdf_formula', + '_entropy': '_entropy_formula', + '_median': '_median_formula'} + + # These are not desirable overrides for the new infrastructure + skip_override = {'norminvgauss': {'_sf', '_isf'}} + + for old_method, new_method in methods.items(): + if dist.name in skip_override and old_method in skip_override[dist.name]: + continue + # If method of old distribution overrides generic implementation... + method = getattr(dist.__class__, old_method, None) + super_method = getattr(old_class, old_method, None) + if method is not super_method: + # Make it an attribute of the new object with the new name + setattr(CustomDistribution, new_method, getattr(dist, old_method)) + + if _overrides('_munp'): + CustomDistribution._moment_raw_formula = _moment_raw_formula + + if _overrides('_rvs'): + CustomDistribution._sample_formula = _sample_formula + + if _overrides('_stats'): + CustomDistribution._moment_standardized_formula = _moment_standard_formula + if not _overrides('_munp'): + CustomDistribution._moment_raw_formula = _moment_raw_formula_1 + CustomDistribution._moment_central_formula = _moment_central_formula + + support_etc = _combine_docs(CustomDistribution, include_examples=False).lstrip() + docs = [ + f"This class represents `scipy.stats.{dist.name}` as a subclass of " + f"`{new_class}`.", + f"The `repr`/`str` of class instances is `{repr_str}`.", + f"The PDF of the distribution is defined {support_etc}" + ] + CustomDistribution.__doc__ = ("\n".join(docs)) + + return CustomDistribution + + +def _get_domain_info(info): + domain_info = {"endpoints": info} if isinstance(info, tuple) else info + typical = domain_info.pop("typical", None) + return domain_info, typical + + +def _make_distribution_custom(dist): + dist_parameters = ( + dist.parameters if isinstance(dist.parameters, tuple) else (dist.parameters, ) + ) + parameterizations = [] + for parameterization in dist_parameters: + # The attribute name ``parameters`` appears reasonable from a user facing + # perspective, but there is a little tension here with the internal. It's + # important to keep in mind that the ``parameters`` attribute in a + # user-created custom distribution specifies ``_parameterizations`` within + # the infrastructure. + parameters = [] + + for name, info in parameterization.items(): + domain_info, typical = _get_domain_info(info) + domain = _RealInterval(**domain_info) + param = _RealParameter(name, domain=domain, typical=typical) + parameters.append(param) + + if parameters: + parameterizations.append(_Parameterization(*parameters)) + + domain_info, typical = _get_domain_info(dist.support) + _x_support = _RealInterval(**domain_info) + _x_param = _RealParameter('x', domain=_x_support, typical=typical) + repr_str = dist.__class__.__name__ + + class CustomDistribution(ContinuousDistribution): + _parameterizations = parameterizations + _variable = _x_param + + def __repr__(self): + s = super().__repr__() + return s.replace('CustomDistribution', repr_str) + + def __str__(self): + s = super().__str__() + return s.replace('CustomDistribution', repr_str) + + methods = {'sample', 'logentropy', 'entropy', + 'median', 'mode', 'logpdf', 'pdf', + 'logcdf2', 'logcdf', 'cdf2', 'cdf', + 'logccdf2', 'logccdf', 'ccdf2', 'ccdf', + 'ilogcdf', 'icdf', 'ilogccdf', 'iccdf', + 'lmoment'} + + for method in methods: + if hasattr(dist, method): + # Make it an attribute of the new object with the new name + new_method = f"_{method}_formula" + setattr(CustomDistribution, new_method, getattr(dist, method)) + + if hasattr(dist, 'moment'): + def _moment_raw_formula(self, order, **kwargs): + return dist.moment(order, kind='raw', **kwargs) + + def _moment_central_formula(self, order, **kwargs): + return dist.moment(order, kind='central', **kwargs) + + def _moment_standardized_formula(self, order, **kwargs): + return dist.moment(order, kind='standardized', **kwargs) + + CustomDistribution._moment_raw_formula = _moment_raw_formula + CustomDistribution._moment_central_formula = _moment_central_formula + CustomDistribution._moment_standardized_formula = _moment_standardized_formula + + if hasattr(dist, 'process_parameters'): + setattr( + CustomDistribution, + "_process_parameters", + getattr(dist, "process_parameters") + ) + + support_etc = _combine_docs(CustomDistribution, include_examples=False).lstrip() + docs = [ + f"This class represents `{repr_str}` as a subclass of " + "`ContinuousDistribution`.", + f"The PDF of the distribution is defined {support_etc}" + ] + CustomDistribution.__doc__ = ("\n".join(docs)) + + return CustomDistribution + + +# The purpose of making ShiftedScaledDistribution a separate class is for +# a) simplicity of the ContinuousDistribution class and +# b) avoiding the requirement that every distribution accept loc/scale. +# The simplicity of ContinuousDistribution is important, because there are +# several other distribution transformations to be supported; e.g., truncation, +# wrapping, folding, and doubling. We wouldn't want to cram all of this +# into the `ContinuousDistribution` class. Also, the order of the composition +# matters (e.g. truncate then shift/scale or vice versa). It's easier to +# accommodate different orders if the transformation is built up from +# components rather than all built into `ContinuousDistribution`. + +def _shift_scale_distribution_function_2arg(func): + def wrapped(self, x, y, *args, loc, scale, sign, **kwargs): + item = func.__name__ + + f = getattr(self._dist, item) + + # Obviously it's possible to get away with half of the work here. + # Let's focus on correct results first and optimize later. + xt = self._transform(x, loc, scale) + yt = self._transform(y, loc, scale) + fxy = f(xt, yt, *args, **kwargs) + fyx = f(yt, xt, *args, **kwargs) + return np.real_if_close(np.where(sign, fxy, fyx))[()] + + return wrapped + +def _shift_scale_distribution_function(func): + # c is for complementary + citem = {'_logcdf_dispatch': '_logccdf_dispatch', + '_cdf_dispatch': '_ccdf_dispatch', + '_logccdf_dispatch': '_logcdf_dispatch', + '_ccdf_dispatch': '_cdf_dispatch'} + def wrapped(self, x, *args, loc, scale, sign, **kwargs): + item = func.__name__ + + f = getattr(self._dist, item) + cf = getattr(self._dist, citem[item]) + + # Obviously it's possible to get away with half of the work here. + # Let's focus on correct results first and optimize later. + xt = self._transform(x, loc, scale) + fx = f(xt, *args, **kwargs) + cfx = cf(xt, *args, **kwargs) + return np.where(sign, fx, cfx)[()] + + return wrapped + +def _shift_scale_inverse_function(func): + citem = {'_ilogcdf_dispatch': '_ilogccdf_dispatch', + '_icdf_dispatch': '_iccdf_dispatch', + '_ilogccdf_dispatch': '_ilogcdf_dispatch', + '_iccdf_dispatch': '_icdf_dispatch'} + def wrapped(self, p, *args, loc, scale, sign, **kwargs): + item = func.__name__ + + f = getattr(self._dist, item) + cf = getattr(self._dist, citem[item]) + + # Obviously it's possible to get away with half of the work here. + # Let's focus on correct results first and optimize later. + fx = self._itransform(f(p, *args, **kwargs), loc, scale) + cfx = self._itransform(cf(p, *args, **kwargs), loc, scale) + return np.where(sign, fx, cfx)[()] + + return wrapped + + +class TransformedDistribution(ContinuousDistribution): + def __init__(self, X, /, *args, **kwargs): + if not isinstance(X, ContinuousDistribution): + message = "Transformations are currently only supported for continuous RVs." + raise NotImplementedError(message) + self._copy_parameterization() + self._variable = X._variable + self._dist = X + if X._parameterization: + # Add standard distribution parameters to our parameterization + dist_parameters = X._parameterization.parameters + set_params = set(dist_parameters) + if not self._parameterizations: + self._parameterizations.append(_Parameterization()) + for parameterization in self._parameterizations: + if set_params.intersection(parameterization.parameters): + message = (f"One or more of the parameters of {X} has " + "the same name as a parameter of " + f"{self.__class__.__name__}. Name collisions " + "create ambiguities and are not supported.") + raise ValueError(message) + parameterization.parameters.update(dist_parameters) + super().__init__(*args, **kwargs) + + def _overrides(self, method_name): + return (self._dist._overrides(method_name) + or super()._overrides(method_name)) + + def reset_cache(self): + super().reset_cache() + + def _update_parameters(self, *, validation_policy=None, **params): + # maybe broadcast everything before processing? + parameters = {} + # There may be some issues with _original_parameters + # We only want to update with _dist._original_parameters during + # initialization. Afterward that, we want to start with + # self._original_parameters. + parameters.update(self._dist._original_parameters) + parameters.update(params) + super()._update_parameters(validation_policy=validation_policy, **parameters) + + def _process_parameters(self, **params): + return self._dist._process_parameters(**params) + + def __repr__(self): + raise NotImplementedError() + + def __str__(self): + raise NotImplementedError() + + +class TruncatedDistribution(TransformedDistribution): + """Truncated distribution.""" + # TODO: + # - consider avoiding catastropic cancellation by using appropriate tail + # - if the mode of `_dist` is within the support, it's still the mode + # - rejection sampling might be more efficient than inverse transform + + _lb_domain = _RealInterval(endpoints=(-inf, 'ub'), inclusive=(True, False)) + _lb_param = _RealParameter('lb', symbol=r'b_l', + domain=_lb_domain, typical=(0.1, 0.2)) + + _ub_domain = _RealInterval(endpoints=('lb', inf), inclusive=(False, True)) + _ub_param = _RealParameter('ub', symbol=r'b_u', + domain=_ub_domain, typical=(0.8, 0.9)) + + _parameterizations = [_Parameterization(_lb_param, _ub_param), + _Parameterization(_lb_param), + _Parameterization(_ub_param)] + + def __init__(self, X, /, *args, lb=-np.inf, ub=np.inf, **kwargs): + return super().__init__(X, *args, lb=lb, ub=ub, **kwargs) + + def _process_parameters(self, lb=None, ub=None, **params): + lb = lb if lb is not None else np.full_like(lb, -np.inf)[()] + ub = ub if ub is not None else np.full_like(ub, np.inf)[()] + parameters = self._dist._process_parameters(**params) + a, b = self._support(lb=lb, ub=ub, **parameters) + logmass = self._dist._logcdf2_dispatch(a, b, **parameters) + parameters.update(dict(lb=lb, ub=ub, _a=a, _b=b, logmass=logmass)) + return parameters + + def _support(self, lb, ub, **params): + a, b = self._dist._support(**params) + return np.maximum(a, lb), np.minimum(b, ub) + + def _overrides(self, method_name): + return False + + def _logpdf_dispatch(self, x, *args, lb, ub, _a, _b, logmass, **params): + logpdf = self._dist._logpdf_dispatch(x, *args, **params) + return logpdf - logmass + + def _logcdf_dispatch(self, x, *args, lb, ub, _a, _b, logmass, **params): + logcdf = self._dist._logcdf2_dispatch(_a, x, *args, **params) + # of course, if this result is small we could compute with the other tail + return logcdf - logmass + + def _logccdf_dispatch(self, x, *args, lb, ub, _a, _b, logmass, **params): + logccdf = self._dist._logcdf2_dispatch(x, _b, *args, **params) + return logccdf - logmass + + def _logcdf2_dispatch(self, x, y, *args, lb, ub, _a, _b, logmass, **params): + logcdf2 = self._dist._logcdf2_dispatch(x, y, *args, **params) + return logcdf2 - logmass + + def _ilogcdf_dispatch(self, logp, *args, lb, ub, _a, _b, logmass, **params): + log_Fa = self._dist._logcdf_dispatch(_a, *args, **params) + logp_adjusted = np.logaddexp(log_Fa, logp + logmass) + return self._dist._ilogcdf_dispatch(logp_adjusted, *args, **params) + + def _ilogccdf_dispatch(self, logp, *args, lb, ub, _a, _b, logmass, **params): + log_cFb = self._dist._logccdf_dispatch(_b, *args, **params) + logp_adjusted = np.logaddexp(log_cFb, logp + logmass) + return self._dist._ilogccdf_dispatch(logp_adjusted, *args, **params) + + def _icdf_dispatch(self, p, *args, lb, ub, _a, _b, logmass, **params): + Fa = self._dist._cdf_dispatch(_a, *args, **params) + p_adjusted = Fa + p*np.exp(logmass) + return self._dist._icdf_dispatch(p_adjusted, *args, **params) + + def _iccdf_dispatch(self, p, *args, lb, ub, _a, _b, logmass, **params): + cFb = self._dist._ccdf_dispatch(_b, *args, **params) + p_adjusted = cFb + p*np.exp(logmass) + return self._dist._iccdf_dispatch(p_adjusted, *args, **params) + + def __repr__(self): + with np.printoptions(threshold=10): + return (f"truncate({repr(self._dist)}, " + f"lb={repr(self.lb)}, ub={repr(self.ub)})") + + def __str__(self): + with np.printoptions(threshold=10): + return (f"truncate({str(self._dist)}, " + f"lb={str(self.lb)}, ub={str(self.ub)})") + + +@xp_capabilities(np_only=True) +def truncate(X, lb=-np.inf, ub=np.inf): + """Truncate the support of a random variable. + + Given a random variable `X`, `truncate` returns a random variable with + support truncated to the interval between `lb` and `ub`. The underlying + probability density function is normalized accordingly. + + Parameters + ---------- + X : `ContinuousDistribution` + The random variable to be truncated. + lb, ub : float array-like + The lower and upper truncation points, respectively. Must be + broadcastable with one another and the shape of `X`. + + Returns + ------- + X : `ContinuousDistribution` + The truncated random variable. + + References + ---------- + .. [1] "Truncated Distribution". *Wikipedia*. + https://en.wikipedia.org/wiki/Truncated_distribution + + Examples + -------- + Compare against `scipy.stats.truncnorm`, which truncates a standard normal, + *then* shifts and scales it. + + >>> import numpy as np + >>> import matplotlib.pyplot as plt + >>> from scipy import stats + >>> loc, scale, lb, ub = 1, 2, -2, 2 + >>> X = stats.truncnorm(lb, ub, loc, scale) + >>> Y = scale * stats.truncate(stats.Normal(), lb, ub) + loc + >>> x = np.linspace(-3, 5, 300) + >>> plt.plot(x, X.pdf(x), '-', label='X') + >>> plt.plot(x, Y.pdf(x), '--', label='Y') + >>> plt.xlabel('x') + >>> plt.ylabel('PDF') + >>> plt.title('Truncated, then Shifted/Scaled Normal') + >>> plt.legend() + >>> plt.show() + + However, suppose we wish to shift and scale a normal random variable, + then truncate its support to given values. This is straightforward with + `truncate`. + + >>> Z = stats.truncate(scale * stats.Normal() + loc, lb, ub) + >>> Z.plot() + >>> plt.show() + + Furthermore, `truncate` can be applied to any random variable: + + >>> Rayleigh = stats.make_distribution(stats.rayleigh) + >>> W = stats.truncate(Rayleigh(), lb=0.5, ub=3) + >>> W.plot() + >>> plt.show() + + """ + return TruncatedDistribution(X, lb=lb, ub=ub) + + +class ShiftedScaledDistribution(TransformedDistribution): + """Distribution with a standard shift/scale transformation.""" + # Unclear whether infinite loc/scale will work reasonably in all cases + _loc_domain = _RealInterval(endpoints=(-inf, inf), inclusive=(True, True)) + _loc_param = _RealParameter('loc', symbol=r'\mu', + domain=_loc_domain, typical=(1, 2)) + + _scale_domain = _RealInterval(endpoints=(-inf, inf), inclusive=(True, True)) + _scale_param = _RealParameter('scale', symbol=r'\sigma', + domain=_scale_domain, typical=(0.1, 10)) + + _parameterizations = [_Parameterization(_loc_param, _scale_param), + _Parameterization(_loc_param), + _Parameterization(_scale_param)] + + def _process_parameters(self, loc=None, scale=None, **params): + loc = loc if loc is not None else np.zeros_like(scale)[()] + scale = scale if scale is not None else np.ones_like(loc)[()] + sign = scale > 0 + parameters = self._dist._process_parameters(**params) + parameters.update(dict(loc=loc, scale=scale, sign=sign)) + return parameters + + def _transform(self, x, loc, scale, **kwargs): + return (x - loc)/scale + + def _itransform(self, x, loc, scale, **kwargs): + return x * scale + loc + + def _support(self, loc, scale, sign, **params): + # Add shortcut for infinite support? + a, b = self._dist._support(**params) + a, b = self._itransform(a, loc, scale), self._itransform(b, loc, scale) + return np.where(sign, a, b)[()], np.where(sign, b, a)[()] + + def __repr__(self): + with np.printoptions(threshold=10): + result = f"{repr(self.scale)}*{repr(self._dist)}" + if not self.loc.ndim and self.loc < 0: + result += f" - {repr(-self.loc)}" + elif (np.any(self.loc != 0) + or not np.can_cast(self.loc.dtype, self.scale.dtype)): + # We don't want to hide a zero array loc if it can cause + # a type promotion. + result += f" + {repr(self.loc)}" + return result + + def __str__(self): + with np.printoptions(threshold=10): + result = f"{str(self.scale)}*{str(self._dist)}" + if not self.loc.ndim and self.loc < 0: + result += f" - {str(-self.loc)}" + elif (np.any(self.loc != 0) + or not np.can_cast(self.loc.dtype, self.scale.dtype)): + # We don't want to hide a zero array loc if it can cause + # a type promotion. + result += f" + {str(self.loc)}" + return result + + # Here, we override all the `_dispatch` methods rather than the public + # methods or _formula methods. Why not the public methods? + # If we were to override the public methods, then other + # TransformedDistribution classes (which could transform a + # ShiftedScaledDistribution) would need to call the public methods of + # ShiftedScaledDistribution, which would run the input validation again. + # Why not the _formula methods? For distributions that rely on the + # default implementation of methods (e.g. `quadrature`, `inversion`), + # the implementation would "see" the location and scale like other + # distribution parameters, so they could affect the accuracy of the + # calculations. I think it is cleaner if `loc` and `scale` do not affect + # the underlying calculations at all. + + def _entropy_dispatch(self, *args, loc, scale, sign, **params): + return (self._dist._entropy_dispatch(*args, **params) + + np.log(np.abs(scale))) + + def _logentropy_dispatch(self, *args, loc, scale, sign, **params): + lH0 = self._dist._logentropy_dispatch(*args, **params) + lls = np.log(np.log(np.abs(scale))+0j) + return special.logsumexp(np.broadcast_arrays(lH0, lls), axis=0) + + def _median_dispatch(self, *, method, loc, scale, sign, **params): + raw = self._dist._median_dispatch(method=method, **params) + return self._itransform(raw, loc, scale) + + def _mode_dispatch(self, *, method, loc, scale, sign, **params): + raw = self._dist._mode_dispatch(method=method, **params) + return self._itransform(raw, loc, scale) + + def _logpdf_dispatch(self, x, *args, loc, scale, sign, **params): + x = self._transform(x, loc, scale) + logpdf = self._dist._logpdf_dispatch(x, *args, **params) + return logpdf - np.log(np.abs(scale)) + + def _pdf_dispatch(self, x, *args, loc, scale, sign, **params): + x = self._transform(x, loc, scale) + pdf = self._dist._pdf_dispatch(x, *args, **params) + return pdf / np.abs(scale) + + def _logpmf_dispatch(self, x, *args, loc, scale, sign, **params): + x = self._transform(x, loc, scale) + logpmf = self._dist._logpmf_dispatch(x, *args, **params) + return logpmf - np.log(np.abs(scale)) + + def _pmf_dispatch(self, x, *args, loc, scale, sign, **params): + x = self._transform(x, loc, scale) + pmf = self._dist._pmf_dispatch(x, *args, **params) + return pmf / np.abs(scale) + + def _logpxf_dispatch(self, x, *args, loc, scale, sign, **params): + x = self._transform(x, loc, scale) + logpxf = self._dist._logpxf_dispatch(x, *args, **params) + return logpxf - np.log(np.abs(scale)) + + def _pxf_dispatch(self, x, *args, loc, scale, sign, **params): + x = self._transform(x, loc, scale) + pxf = self._dist._pxf_dispatch(x, *args, **params) + return pxf / np.abs(scale) + + @_shift_scale_distribution_function + def _logcdf_dispatch(self, x, *, method=None, **params): + pass + + @_shift_scale_distribution_function + def _cdf_dispatch(self, x, *, method=None, **params): + pass + + @_shift_scale_distribution_function + def _logccdf_dispatch(self, x, *, method=None, **params): + pass + + @_shift_scale_distribution_function + def _ccdf_dispatch(self, x, *, method=None, **params): + pass + + @_shift_scale_distribution_function_2arg + def _logcdf2_dispatch(self, x, y, *, method=None, **params): + pass + + @_shift_scale_distribution_function_2arg + def _cdf2_dispatch(self, x, y, *, method=None, **params): + pass + + @_shift_scale_distribution_function_2arg + def _logccdf2_dispatch(self, x, y, *, method=None, **params): + pass + + @_shift_scale_distribution_function_2arg + def _ccdf2_dispatch(self, x, y, *, method=None, **params): + pass + + @_shift_scale_inverse_function + def _ilogcdf_dispatch(self, x, *, method=None, **params): + pass + + @_shift_scale_inverse_function + def _icdf_dispatch(self, x, *, method=None, **params): + pass + + @_shift_scale_inverse_function + def _ilogccdf_dispatch(self, x, *, method=None, **params): + pass + + @_shift_scale_inverse_function + def _iccdf_dispatch(self, x, *, method=None, **params): + pass + + def _moment_standardized_dispatch(self, order, *, loc, scale, sign, methods, + **params): + res = (self._dist._moment_standardized_dispatch( + order, methods=methods, **params)) + return None if res is None else res * np.sign(scale)**order + + def _moment_central_dispatch(self, order, *, loc, scale, sign, methods, + **params): + res = (self._dist._moment_central_dispatch( + order, methods=methods, **params)) + return None if res is None else res * scale**order + + def _moment_raw_dispatch(self, order, *, loc, scale, sign, methods, + **params): + raw_moments = [] + methods_highest_order = methods + for i in range(int(order) + 1): + methods = (self._moment_methods if i < order + else methods_highest_order) + raw = self._dist._moment_raw_dispatch(i, methods=methods, **params) + if raw is None: + return None + moment_i = raw * scale**i + raw_moments.append(moment_i) + + return self._moment_transform_center( + order, raw_moments, loc, self._zero) + + def _lmoment_dispatch(self, order, *, loc, scale, sign, methods, **params): + res = self._dist._lmoment_dispatch(order, methods=methods, **params) + if res is None: # if a specific method is requested but not available + return None + res = res * np.abs(scale) * np.sign(scale)**order + return res + loc if order == 1 else res + + def _sample_dispatch(self, full_shape, *, + rng, loc, scale, sign, method, **params): + rvs = self._dist._sample_dispatch(full_shape, method=method, rng=rng, **params) + return self._itransform(rvs, loc=loc, scale=scale, sign=sign, **params) + + def __add__(self, loc): + return ShiftedScaledDistribution(self._dist, loc=self.loc + loc, + scale=self.scale) + + def __sub__(self, loc): + return ShiftedScaledDistribution(self._dist, loc=self.loc - loc, + scale=self.scale) + + def __mul__(self, scale): + return ShiftedScaledDistribution(self._dist, + loc=self.loc * scale, + scale=self.scale * scale) + + def __truediv__(self, scale): + return ShiftedScaledDistribution(self._dist, + loc=self.loc / scale, + scale=self.scale / scale) + + +class OrderStatisticDistribution(TransformedDistribution): + r"""Probability distribution of an order statistic. + + An instance of this class represents a random variable that follows the + distribution of the :math:`r^{\text{th}}` order statistic of a + sample of :math:`n` observations of a random variable :math:`X`. + + Parameters + ---------- + dist : `ContinuousDistribution` + The random variable :math:`X`. + n : array_like + The (positive integer) sample size :math:`n`. + r : array_like + The (positive integer) rank of the order statistic :math:`r`, + satisfying :math:`1 \leq r \leq n`. + + + Notes + ----- + If we make :math:`n` observations of a continuous random variable + :math:`X` and sort them in increasing order + :math:`X_{(1)}, \dots, X_{(r)}, \dots, X_{(n)}`, then :math:`X_{(r)}` + is known as the :math:`r^{\text{th}}` order statistic. + + If the PDF, CDF, and CCDF of :math:`X` are denoted by :math:`f`, + :math:`F`, and :math:`G = 1 - F`, respectively, then the PDF of + :math:`X_{(r)}` is given by: + + .. math:: + + f_r(x) = \frac{n!}{(r-1)! (n-r)!} f(x) F(x)^{r-1} G(x)^{n - r} + + The CDF and other methods of the distribution of :math:`X_{(r)}` + are calculated using the fact that :math:`X = F^{-1}(U)`, where :math:`U` is a + standard uniform random variable, together with the fact that the order statistics + of i.i.d. uniform random variables follow a beta distribution + :math:`B(r, n - r + 1)`. + + References + ---------- + .. [1] Order statistic. *Wikipedia*. https://en.wikipedia.org/wiki/Order_statistic + + Examples + -------- + Suppose we are interested in order statistics of samples of size five drawn + from the standard normal distribution. Plot the PDF of the fourth + order statistic and compare with a normalized histogram from simulation. + + >>> import numpy as np + >>> import matplotlib.pyplot as plt + >>> from scipy import stats + >>> from scipy.stats._distribution_infrastructure import OrderStatisticDistribution + >>> + >>> X = stats.Normal() + >>> data = X.sample(shape=(10000, 5)) + >>> sorted_data = np.sort(data, axis=1) + >>> Y = OrderStatisticDistribution(X, r=4, n=5) + >>> + >>> ax = plt.gca() + >>> Y.plot(ax=ax) + >>> ax.hist(sorted_data[:, 3], density=True, bins=30) + >>> plt.show() + + """ + + _r_domain = _IntegerInterval(endpoints=(1, 'n'), inclusive=(True, True)) + _r_param = _RealParameter('r', domain=_r_domain, typical=(1, 2)) + + _n_domain = _IntegerInterval(endpoints=(1, np.inf), inclusive=(True, True)) + _n_param = _RealParameter('n', domain=_n_domain, typical=(1, 4)) + + _r_domain.define_parameters(_n_param) + + _parameterizations = [_Parameterization(_r_param, _n_param)] + + def __init__(self, dist, /, *args, r, n, **kwargs): + super().__init__(dist, *args, r=r, n=n, **kwargs) + + def _support(self, *args, r, n, **kwargs): + return self._dist._support(*args, **kwargs) + + def _process_parameters(self, r=None, n=None, **params): + parameters = self._dist._process_parameters(**params) + parameters.update(dict(r=r, n=n)) + return parameters + + def _overrides(self, method_name): + return method_name in {'_logpdf_formula', '_pdf_formula', + '_cdf_formula', '_ccdf_formula', + '_icdf_formula', '_iccdf_formula'} + + def _logpdf_formula(self, x, r, n, **kwargs): + log_factor = special.betaln(r, n - r + 1) + log_fX = self._dist._logpdf_dispatch(x, **kwargs) + # log-methods sometimes use complex dtype with 0 imaginary component, + # but `_tanhsinh` doesn't accept complex limits of integration; take `real`. + log_FX = self._dist._logcdf_dispatch(x.real, **kwargs) + log_cFX = self._dist._logccdf_dispatch(x.real, **kwargs) + # This can be problematic when (r - 1)|(n-r) = 0 and `log_FX`|log_cFX = -inf + # The PDF in these cases is 0^0, so these should be replaced with log(1)=0 + # return log_fX + (r-1)*log_FX + (n-r)*log_cFX - log_factor + rm1_log_FX = np.where((r - 1 == 0) & np.isneginf(log_FX), 0, (r-1)*log_FX) + nmr_log_cFX = np.where((n - r == 0) & np.isneginf(log_cFX), 0, (n-r)*log_cFX) + return log_fX + rm1_log_FX + nmr_log_cFX - log_factor + + def _pdf_formula(self, x, r, n, **kwargs): + # 1 / factor = factorial(n) / (factorial(r-1) * factorial(n-r)) + factor = special.beta(r, n - r + 1) + fX = self._dist._pdf_dispatch(x, **kwargs) + FX = self._dist._cdf_dispatch(x, **kwargs) + cFX = self._dist._ccdf_dispatch(x, **kwargs) + return fX * FX**(r-1) * cFX**(n-r) / factor + + def _cdf_formula(self, x, r, n, **kwargs): + x_ = self._dist._cdf_dispatch(x, **kwargs) + return special.betainc(r, n-r+1, x_) + + def _ccdf_formula(self, x, r, n, **kwargs): + x_ = self._dist._cdf_dispatch(x, **kwargs) + return special.betaincc(r, n-r+1, x_) + + def _icdf_formula(self, p, r, n, **kwargs): + p_ = special.betaincinv(r, n-r+1, p) + return self._dist._icdf_dispatch(p_, **kwargs) + + def _iccdf_formula(self, p, r, n, **kwargs): + p_ = special.betainccinv(r, n-r+1, p) + return self._dist._icdf_dispatch(p_, **kwargs) + + def __repr__(self): + with np.printoptions(threshold=10): + return (f"order_statistic({repr(self._dist)}, r={repr(self.r)}, " + f"n={repr(self.n)})") + + def __str__(self): + with np.printoptions(threshold=10): + return (f"order_statistic({str(self._dist)}, r={str(self.r)}, " + f"n={str(self.n)})") + + +@xp_capabilities(np_only=True) +def order_statistic(X, /, *, r, n): + r"""Probability distribution of an order statistic. + + Returns a random variable that follows the distribution of the + :math:`r^{\text{th}}` order statistic of a sample of :math:`n` + observations of a random variable :math:`X`. + + Parameters + ---------- + X : `ContinuousDistribution` + The random variable :math:`X`. + r : array_like + The (positive integer) rank of the order statistic :math:`r`, + satisfying ``1 <= r <= n``. + n : array_like + The (positive integer) sample size :math:`n`. + + Returns + ------- + Y : `ContinuousDistribution` + A random variable that follows the distribution of the prescribed + order statistic. + + Notes + ----- + If we make :math:`n` observations of a continuous random variable + :math:`X` and sort them in increasing order + :math:`X_{(1)}, \dots, X_{(r)}, \dots, X_{(n)}`, then :math:`X_{(r)}` + is known as the :math:`r^{\text{th}}` order statistic. + + If the PDF, CDF, and CCDF of :math:`X` are denoted by :math:`f`, + :math:`F`, and :math:`G = 1 - F`, respectively, then the PDF of + :math:`X_{(r)}` is given by: + + .. math:: + + f_r(x) = \frac{n!}{(r-1)! (n-r)!} f(x) F(x)^{r-1} G(x)^{n - r} + + The CDF and other methods of the distribution of :math:`X_{(r)}` + are calculated using the fact that :math:`X = F^{-1}(U)`, where :math:`U` is a + standard uniform random variable, together with the fact that the order statistics + of i.i.d. uniform random variables follow a beta distribution + :math:`B(r, n - r + 1)`. + + References + ---------- + .. [1] Order statistic. *Wikipedia*. https://en.wikipedia.org/wiki/Order_statistic + + Examples + -------- + Suppose we are interested in order statistics of samples of size five drawn + from the standard normal distribution. Plot the PDF of each + order statistic and compare with a normalized histogram from simulation. + + >>> import numpy as np + >>> import matplotlib.pyplot as plt + >>> from scipy import stats + >>> + >>> X = stats.Normal() + >>> data = X.sample(shape=(10000, 5)) + >>> sorted_data = np.sort(data, axis=1) + >>> Y = stats.order_statistic(X, r=[1, 2, 3, 4, 5], n=5) + >>> + >>> ax = plt.gca() + >>> colors = plt.rcParams['axes.prop_cycle'].by_key()['color'] + >>> for i in range(5): + ... y = sorted_data[:, i] + ... ax.hist(y, density=True, bins=30, alpha=0.1, color=colors[i]) + >>> Y.plot(ax=ax) + >>> plt.show() + + """ + r, n = np.asarray(r), np.asarray(n) + if np.any((r != np.floor(r)) | (r < 0)) or np.any((n != np.floor(n)) | (n < 0)): + message = "`r` and `n` must contain only positive integers." + raise ValueError(message) + return OrderStatisticDistribution(X, r=r, n=n) + + +class Mixture(_ProbabilityDistribution): + r"""Representation of a mixture distribution. + + A mixture distribution is the distribution of a random variable + defined in the following way: first, a random variable is selected + from `components` according to the probabilities given by `weights`, then + the selected random variable is realized. + + Parameters + ---------- + components : sequence of `ContinuousDistribution` + The underlying instances of `ContinuousDistribution`. + All must have scalar shape parameters (if any); e.g., the `pdf` evaluated + at a scalar argument must return a scalar. + weights : sequence of floats, optional + The corresponding probabilities of selecting each random variable. + Must be non-negative and sum to one. The default behavior is to weight + all components equally. + + Attributes + ---------- + components : sequence of `ContinuousDistribution` + The underlying instances of `ContinuousDistribution`. + weights : ndarray + The corresponding probabilities of selecting each random variable. + + Methods + ------- + support + + sample + + moment + + mean + median + mode + + variance + standard_deviation + + skewness + kurtosis + + pdf + logpdf + + cdf + icdf + ccdf + iccdf + + logcdf + ilogcdf + logccdf + ilogccdf + + entropy + + Notes + ----- + The following abbreviations are used throughout the documentation. + + - PDF: probability density function + - CDF: cumulative distribution function + - CCDF: complementary CDF + - entropy: differential entropy + - log-*F*: logarithm of *F* (e.g. log-CDF) + - inverse *F*: inverse function of *F* (e.g. inverse CDF) + + References + ---------- + .. [1] Mixture distribution, *Wikipedia*, + https://en.wikipedia.org/wiki/Mixture_distribution + + + Examples + -------- + A mixture of normal distributions: + + >>> import numpy as np + >>> from scipy import stats + >>> import matplotlib.pyplot as plt + >>> X1 = stats.Normal(mu=-2, sigma=1) + >>> X2 = stats.Normal(mu=2, sigma=1) + >>> mixture = stats.Mixture([X1, X2], weights=[0.4, 0.6]) + >>> print(f'mean: {mixture.mean():.2f}, ' + ... f'median: {mixture.median():.2f}, ' + ... f'mode: {mixture.mode():.2f}') + mean: 0.40, median: 1.04, mode: 2.00 + >>> x = np.linspace(-10, 10, 300) + >>> plt.plot(x, mixture.pdf(x)) + >>> plt.title('PDF of normal distribution mixture') + >>> plt.show() + + """ + # Todo: + # Add support for array shapes, weights + + def _input_validation(self, components, weights): + if len(components) == 0: + message = ("`components` must contain at least one random variable.") + raise ValueError(message) + + for var in components: + # will generalize to other kinds of distributions when there + # *are* other kinds of distributions + if not isinstance(var, ContinuousDistribution): + message = ("Each element of `components` must be an instance of " + "`ContinuousDistribution`.") + raise ValueError(message) + if not var._shape == (): + message = "All elements of `components` must have scalar shapes." + raise ValueError(message) + + if weights is None: + return components, weights + + weights = np.asarray(weights) + if weights.shape != (len(components),): + message = "`components` and `weights` must have the same length." + raise ValueError(message) + + if not np.issubdtype(weights.dtype, np.inexact): + message = "`weights` must have floating point dtype." + raise ValueError(message) + + if not np.isclose(np.sum(weights), 1.0): + message = "`weights` must sum to 1.0." + raise ValueError(message) + + if not np.all(weights >= 0): + message = "All `weights` must be non-negative." + raise ValueError(message) + + return components, weights + + def __init__(self, components, *, weights=None): + components, weights = self._input_validation(components, weights) + n = len(components) + dtype = np.result_type(*(var._dtype for var in components)) + self._shape = np.broadcast_shapes(*(var._shape for var in components)) + self._dtype, self._components = dtype, components + self._weights = np.full(n, 1/n, dtype=dtype) if weights is None else weights + self.validation_policy = None + + @property + def components(self): + return list(self._components) + + @property + def weights(self): + return self._weights.copy() + + def _full(self, val, *args): + args = [np.asarray(arg) for arg in args] + dtype = np.result_type(self._dtype, *(arg.dtype for arg in args)) + shape = np.broadcast_shapes(self._shape, *(arg.shape for arg in args)) + return np.full(shape, val, dtype=dtype) + + def _sum(self, fun, *args): + out = self._full(0, *args) + for var, weight in zip(self._components, self._weights): + out += getattr(var, fun)(*args) * weight + return out[()] + + def _logsum(self, fun, *args): + out = self._full(-np.inf, *args) + for var, log_weight in zip(self._components, np.log(self._weights)): + np.logaddexp(out, getattr(var, fun)(*args) + log_weight, out=out) + return out[()] + + def support(self): + a = self._full(np.inf) + b = self._full(-np.inf) + for var in self._components: + a = np.minimum(a, var.support()[0]) + b = np.maximum(b, var.support()[1]) + return a, b + + def _raise_if_method(self, method): + if method is not None: + raise NotImplementedError("`method` not implemented for this distribution.") + + def logentropy(self, *, method=None): + self._raise_if_method(method) + def log_integrand(x): + # `x` passed by `_tanhsinh` will be of complex dtype because + # `log_integrand` returns complex values, but the imaginary + # component is always zero. Extract the real part because + # `logpdf` uses `logaddexp`, which fails for complex input. + return self.logpdf(x.real) + np.log(self.logpdf(x.real) + 0j) + + res = _tanhsinh(log_integrand, *self.support(), log=True).integral + return _log_real_standardize(res + np.pi*1j) + + def entropy(self, *, method=None): + self._raise_if_method(method) + return _tanhsinh(lambda x: -self.pdf(x) * self.logpdf(x), + *self.support()).integral + + def mode(self, *, method=None): + self._raise_if_method(method) + a, b = self.support() + def f(x): return -self.pdf(x) + res = elementwise.bracket_minimum(f, 1., xmin=a, xmax=b) + res = elementwise.find_minimum(f, res.bracket) + return res.x + + def median(self, *, method=None): + self._raise_if_method(method) + return self.icdf(0.5) + + def mean(self, *, method=None): + self._raise_if_method(method) + return self._sum('mean') + + def variance(self, *, method=None): + self._raise_if_method(method) + return self._moment_central(2) + + def standard_deviation(self, *, method=None): + self._raise_if_method(method) + return self.variance()**0.5 + + def skewness(self, *, method=None): + self._raise_if_method(method) + return self._moment_standardized(3) + + def kurtosis(self, *, method=None): + self._raise_if_method(method) + return self._moment_standardized(4) + + def moment(self, order=1, kind='raw', *, method=None): + self._raise_if_method(method) + kinds = {'raw': self._moment_raw, + 'central': self._moment_central, + 'standardized': self._moment_standardized} + order = ContinuousDistribution._validate_order(self, order) + kind = ContinuousDistribution._validate_kind(self, kind, kinds) + moment_kind = kinds[kind] + return moment_kind(order) + + def _moment_raw(self, order): + out = self._full(0) + for var, weight in zip(self._components, self._weights): + out += var.moment(order, kind='raw') * weight + return out[()] + + def _moment_central(self, order): + order = int(order) + out = self._full(0) + for var, weight in zip(self._components, self._weights): + moment_as = [var.moment(order, kind='central') + for order in range(order + 1)] + a, b = var.mean(), self.mean() + moment = var._moment_transform_center(order, moment_as, a, b) + out += moment * weight + return out[()] + + def lmoment(self, order=1, *, standardize=True, method=None): + message = "L-moments are not currently available for mixture distributions." + raise NotImplementedError(message) + + def _moment_standardized(self, order): + return self._moment_central(order) / self.standard_deviation()**order + + def pdf(self, x, /, *, method=None): + self._raise_if_method(method) + return self._sum('pdf', x) + + def logpdf(self, x, /, *, method=None): + self._raise_if_method(method) + return self._logsum('logpdf', x) + + def pmf(self, x, /, *, method=None): + self._raise_if_method(method) + return self._sum('pmf', x) + + def logpmf(self, x, /, *, method=None): + self._raise_if_method(method) + return self._logsum('logpmf', x) + + def cdf(self, x, y=None, /, *, method=None): + self._raise_if_method(method) + args = (x,) if y is None else (x, y) + return self._sum('cdf', *args) + + def logcdf(self, x, y=None, /, *, method=None): + self._raise_if_method(method) + args = (x,) if y is None else (x, y) + return self._logsum('logcdf', *args) + + def ccdf(self, x, y=None, /, *, method=None): + self._raise_if_method(method) + args = (x,) if y is None else (x, y) + return self._sum('ccdf', *args) + + def logccdf(self, x, y=None, /, *, method=None): + self._raise_if_method(method) + args = (x,) if y is None else (x, y) + return self._logsum('logccdf', *args) + + def _invert(self, fun, p): + xmin, xmax = self.support() + fun = getattr(self, fun) + f = lambda x, p: fun(x) - p # noqa: E731 is silly + xl0, xr0 = _guess_bracket(xmin, xmax) + res = elementwise.bracket_root(f, xl0=xl0, xr0=xr0, + xmin=xmin, xmax=xmax, args=(p,)) + return elementwise.find_root(f, res.bracket, args=(p,)).x + + def icdf(self, p, /, *, method=None): + self._raise_if_method(method) + return self._invert('cdf', p) + + def iccdf(self, p, /, *, method=None): + self._raise_if_method(method) + return self._invert('ccdf', p) + + def ilogcdf(self, p, /, *, method=None): + self._raise_if_method(method) + return self._invert('logcdf', p) + + def ilogccdf(self, p, /, *, method=None): + self._raise_if_method(method) + return self._invert('logccdf', p) + + def sample(self, shape=(), *, rng=None, method=None): + self._raise_if_method(method) + rng = np.random.default_rng(rng) + size = np.prod(np.atleast_1d(shape)) + ns = rng.multinomial(size, self._weights) + x = [var.sample(shape=n, rng=rng) for n, var in zip(ns, self._components)] + x = np.reshape(rng.permuted(np.concatenate(x)), shape) + return x[()] + + def __repr__(self): + result = "Mixture(\n" + result += " [\n" + with np.printoptions(threshold=10): + for component in self.components: + result += f" {repr(component)},\n" + result += " ],\n" + result += f" weights={repr(self.weights)},\n" + result += ")" + return result + + def __str__(self): + result = "Mixture(\n" + result += " [\n" + with np.printoptions(threshold=10): + for component in self.components: + result += f" {str(component)},\n" + result += " ],\n" + result += f" weights={str(self.weights)},\n" + result += ")" + return result + + +class MonotonicTransformedDistribution(TransformedDistribution): + r"""Distribution underlying a strictly monotonic function of a random variable + + Given a random variable :math:`X`; a strictly monotonic function + :math:`g(u)`, its inverse :math:`h(u) = g^{-1}(u)`, and the derivative magnitude + :math: `|h'(u)| = \left| \frac{dh(u)}{du} \right|`, define the distribution + underlying the random variable :math:`Y = g(X)`. + + Parameters + ---------- + X : `ContinuousDistribution` + The random variable :math:`X`. + g, h, dh : callable + Elementwise functions representing the mathematical functions + :math:`g(u)`, :math:`h(u)`, and :math:`|h'(u)|` + logdh : callable, optional + Elementwise function representing :math:`\log(h'(u))`. + The default is ``lambda u: np.log(dh(u))``, but providing + a custom implementation may avoid over/underflow. + increasing : bool, optional + Whether the function is strictly increasing (True, default) + or strictly decreasing (False). + repr_pattern : str, optional + A string pattern for determining the __repr__. The __repr__ + for X will be substituted into the position where `***` appears. + For example: + ``"exp(***)"`` for the repr of an exponentially transformed + distribution + The default is ``f"{g.__name__}(***)"``. + str_pattern : str, optional + A string pattern for determining `__str__`. The `__str__` + for X will be substituted into the position where `***` appears. + For example: + ``"exp(***)"`` for the repr of an exponentially transformed + distribution + The default is the value `repr_pattern` takes. + """ + + def __init__(self, X, /, *args, g, h, dh, logdh=None, + increasing=True, repr_pattern=None, + str_pattern=None, **kwargs): + super().__init__(X, *args, **kwargs) + self._g = g + self._h = h + self._dh = dh + self._logdh = (logdh if logdh is not None + else lambda u: np.log(dh(u))) + if increasing: + self._xdf = self._dist._cdf_dispatch + self._cxdf = self._dist._ccdf_dispatch + self._ixdf = self._dist._icdf_dispatch + self._icxdf = self._dist._iccdf_dispatch + self._logxdf = self._dist._logcdf_dispatch + self._logcxdf = self._dist._logccdf_dispatch + self._ilogxdf = self._dist._ilogcdf_dispatch + self._ilogcxdf = self._dist._ilogccdf_dispatch + else: + self._xdf = self._dist._ccdf_dispatch + self._cxdf = self._dist._cdf_dispatch + self._ixdf = self._dist._iccdf_dispatch + self._icxdf = self._dist._icdf_dispatch + self._logxdf = self._dist._logccdf_dispatch + self._logcxdf = self._dist._logcdf_dispatch + self._ilogxdf = self._dist._ilogccdf_dispatch + self._ilogcxdf = self._dist._ilogcdf_dispatch + self._increasing = increasing + self._repr_pattern = repr_pattern or f"{g.__name__}(***)" + self._str_pattern = str_pattern or self._repr_pattern + + def __repr__(self): + with np.printoptions(threshold=10): + return self._repr_pattern.replace("***", repr(self._dist)) + + def __str__(self): + with np.printoptions(threshold=10): + return self._str_pattern.replace("***", str(self._dist)) + + def _overrides(self, method_name): + # Do not use the generic overrides of TransformedDistribution + return False + + def _support(self, **params): + a, b = self._dist._support(**params) + # For reciprocal transformation, we want this zero to become -inf + b = np.where(b==0, np.asarray("-0", dtype=b.dtype), b) + with np.errstate(divide='ignore'): + if self._increasing: + return self._g(a), self._g(b) + else: + return self._g(b), self._g(a) + + def _logpdf_dispatch(self, x, *args, **params): + return self._dist._logpdf_dispatch(self._h(x), *args, **params) + self._logdh(x) + + def _pdf_dispatch(self, x, *args, **params): + return self._dist._pdf_dispatch(self._h(x), *args, **params) * self._dh(x) + + def _logcdf_dispatch(self, x, *args, **params): + return self._logxdf(self._h(x), *args, **params) + + def _cdf_dispatch(self, x, *args, **params): + return self._xdf(self._h(x), *args, **params) + + def _logccdf_dispatch(self, x, *args, **params): + return self._logcxdf(self._h(x), *args, **params) + + def _ccdf_dispatch(self, x, *args, **params): + return self._cxdf(self._h(x), *args, **params) + + def _ilogcdf_dispatch(self, p, *args, **params): + return self._g(self._ilogxdf(p, *args, **params)) + + def _icdf_dispatch(self, p, *args, **params): + return self._g(self._ixdf(p, *args, **params)) + + def _ilogccdf_dispatch(self, p, *args, **params): + return self._g(self._ilogcxdf(p, *args, **params)) + + def _iccdf_dispatch(self, p, *args, **params): + return self._g(self._icxdf(p, *args, **params)) + + def _sample_dispatch(self, full_shape, *, method, rng, **params): + rvs = self._dist._sample_dispatch(full_shape, method=method, rng=rng, **params) + return self._g(rvs) + + +class FoldedDistribution(TransformedDistribution): + r"""Distribution underlying the absolute value of a random variable + + Given a random variable :math:`X`; define the distribution + underlying the random variable :math:`Y = |X|`. + + Parameters + ---------- + X : `ContinuousDistribution` + The random variable :math:`X`. + + Returns + ------- + Y : `ContinuousDistribution` + The random variable :math:`Y = |X|` + + """ + # Many enhancements are possible if distribution is symmetric. Start + # with the general case; enhance later. + + def __init__(self, X, /, *args, **kwargs): + super().__init__(X, *args, **kwargs) + # I think we need to allow `_support` to define whether the endpoints + # are inclusive or not. In the meantime, it's best to ensure that the lower + # endpoint (typically 0 for folded distribution) is inclusive so PDF evaluates + # correctly at that point. + self._variable.domain.inclusive = (True, self._variable.domain.inclusive[1]) + + def _overrides(self, method_name): + # Do not use the generic overrides of TransformedDistribution + return False + + def _support(self, **params): + a, b = self._dist._support(**params) + a_, b_ = np.abs(a), np.abs(b) + a_, b_ = np.minimum(a_, b_), np.maximum(a_, b_) + i = (a < 0) & (b > 0) + a_ = np.asarray(a_) + a_[i] = 0 + return a_[()], b_[()] + + def _logpdf_dispatch(self, x, *args, method=None, **params): + x = np.abs(x) + right = self._dist._logpdf_dispatch(x, *args, method=method, **params) + left = self._dist._logpdf_dispatch(-x, *args, method=method, **params) + left = np.asarray(left) + right = np.asarray(right) + a, b = self._dist._support(**params) + left[-x < a] = -np.inf + right[x > b] = -np.inf + logpdfs = np.stack([left, right]) + return special.logsumexp(logpdfs, axis=0) + + def _pdf_dispatch(self, x, *args, method=None, **params): + x = np.abs(x) + right = self._dist._pdf_dispatch(x, *args, method=method, **params) + left = self._dist._pdf_dispatch(-x, *args, method=method, **params) + left = np.asarray(left) + right = np.asarray(right) + a, b = self._dist._support(**params) + left[-x < a] = 0 + right[x > b] = 0 + return left + right + + def _logcdf_dispatch(self, x, *args, method=None, **params): + x = np.abs(x) + a, b = self._dist._support(**params) + xl = np.maximum(-x, a) + xr = np.minimum(x, b) + return self._dist._logcdf2_dispatch(xl, xr, *args, method=method, **params).real + + def _cdf_dispatch(self, x, *args, method=None, **params): + x = np.abs(x) + a, b = self._dist._support(**params) + xl = np.maximum(-x, a) + xr = np.minimum(x, b) + return self._dist._cdf2_dispatch(xl, xr, *args, **params) + + def _logccdf_dispatch(self, x, *args, method=None, **params): + x = np.abs(x) + a, b = self._dist._support(**params) + xl = np.maximum(-x, a) + xr = np.minimum(x, b) + return self._dist._logccdf2_dispatch(xl, xr, *args, method=method, + **params).real + + def _ccdf_dispatch(self, x, *args, method=None, **params): + x = np.abs(x) + a, b = self._dist._support(**params) + xl = np.maximum(-x, a) + xr = np.minimum(x, b) + return self._dist._ccdf2_dispatch(xl, xr, *args, method=method, **params) + + def _sample_dispatch(self, full_shape, *, method, rng, **params): + rvs = self._dist._sample_dispatch(full_shape, method=method, rng=rng, **params) + return np.abs(rvs) + + def __repr__(self): + with np.printoptions(threshold=10): + return f"abs({repr(self._dist)})" + + def __str__(self): + with np.printoptions(threshold=10): + return f"abs({str(self._dist)})" + + +@xp_capabilities(np_only=True) +def abs(X, /): + r"""Absolute value of a random variable. + + Parameters + ---------- + X : `ContinuousDistribution` + The random variable :math:`X`. + + Returns + ------- + Y : `ContinuousDistribution` + A random variable :math:`Y = |X|`. + + Examples + -------- + Suppose we have a normally distributed random variable :math:`X`: + + >>> import numpy as np + >>> from scipy import stats + >>> X = stats.Normal() + + We wish to have a random variable :math:`Y` distributed according to + the folded normal distribution; that is, a random variable :math:`|X|`. + + >>> Y = stats.abs(X) + + The PDF of the distribution in the left half plane is "folded" over to + the right half plane. Because the normal PDF is symmetric, the resulting + PDF is zero for negative arguments and doubled for positive arguments. + + >>> import matplotlib.pyplot as plt + >>> x = np.linspace(0, 5, 300) + >>> ax = plt.gca() + >>> Y.plot(x='x', y='pdf', t=('x', -1, 5), ax=ax) + >>> plt.plot(x, 2 * X.pdf(x), '--') + >>> plt.legend(('PDF of `Y`', 'Doubled PDF of `X`')) + >>> plt.show() + + """ + return FoldedDistribution(X) + + +@xp_capabilities(np_only=True) +def exp(X, /): + r"""Natural exponential of a random variable. + + Parameters + ---------- + X : `ContinuousDistribution` + The random variable :math:`X`. + + Returns + ------- + Y : `ContinuousDistribution` + A random variable :math:`Y = \exp(X)`. + + Examples + -------- + Suppose we have a normally distributed random variable :math:`X`: + + >>> import numpy as np + >>> from scipy import stats + >>> X = stats.Normal() + + We wish to have a lognormally distributed random variable :math:`Y`, + a random variable whose natural logarithm is :math:`X`. + If :math:`X` is to be the natural logarithm of :math:`Y`, then we + must take :math:`Y` to be the natural exponential of :math:`X`. + + >>> Y = stats.exp(X) + + To demonstrate that ``X`` represents the logarithm of ``Y``, + we plot a normalized histogram of the logarithm of observations of + ``Y`` against the PDF underlying ``X``. + + >>> import matplotlib.pyplot as plt + >>> rng = np.random.default_rng(435383595582522) + >>> y = Y.sample(shape=10000, rng=rng) + >>> ax = plt.gca() + >>> ax.hist(np.log(y), bins=50, density=True) + >>> X.plot(ax=ax) + >>> plt.legend(('PDF of `X`', 'histogram of `log(y)`')) + >>> plt.show() + + """ + return MonotonicTransformedDistribution(X, g=np.exp, h=np.log, dh=lambda u: 1 / u, + logdh=lambda u: -np.log(u)) + + +@xp_capabilities(np_only=True) +def log(X, /): + r"""Natural logarithm of a non-negative random variable. + + Parameters + ---------- + X : `ContinuousDistribution` + The random variable :math:`X` with positive support. + + Returns + ------- + Y : `ContinuousDistribution` + A random variable :math:`Y = \log(X)`. + + Examples + -------- + Suppose we have a gamma distributed random variable :math:`X`: + + >>> import numpy as np + >>> from scipy import stats + >>> Gamma = stats.make_distribution(stats.gamma) + >>> X = Gamma(a=1.0) + + We wish to have an exp-gamma distributed random variable :math:`Y`, + a random variable whose natural exponential is :math:`X`. + If :math:`X` is to be the natural exponential of :math:`Y`, then we + must take :math:`Y` to be the natural logarithm of :math:`X`. + + >>> Y = stats.log(X) + + To demonstrate that ``X`` represents the exponential of ``Y``, + we plot a normalized histogram of the exponential of observations of + ``Y`` against the PDF underlying ``X``. + + >>> import matplotlib.pyplot as plt + >>> rng = np.random.default_rng(435383595582522) + >>> y = Y.sample(shape=10000, rng=rng) + >>> ax = plt.gca() + >>> ax.hist(np.exp(y), bins=50, density=True) + >>> X.plot(ax=ax) + >>> plt.legend(('PDF of `X`', 'histogram of `exp(y)`')) + >>> plt.show() + + """ + if np.any(X.support()[0] < 0): + message = ("The logarithm of a random variable is only implemented when the " + "support is non-negative.") + raise NotImplementedError(message) + return MonotonicTransformedDistribution(X, g=np.log, h=np.exp, dh=np.exp, + logdh=lambda u: u) diff --git a/v12x/pylib/scipy/stats/_entropy.py b/v12x/pylib/scipy/stats/_entropy.py new file mode 100644 index 0000000000000000000000000000000000000000..ab5a874cf4aeaa0bb12e63292cd3f5122658f509 --- /dev/null +++ b/v12x/pylib/scipy/stats/_entropy.py @@ -0,0 +1,431 @@ +""" +Created on Fri Apr 2 09:06:05 2021 + +@author: matth +""" + +import math +import numpy as np +from scipy import special +from ._axis_nan_policy import _axis_nan_policy_factory +from scipy._lib._array_api import (array_namespace, xp_promote, xp_device, + _masked_apply, _share_masks, xp_capabilities) + +__all__ = ['entropy', 'differential_entropy'] + + +@xp_capabilities(marray=True) +@_axis_nan_policy_factory( + lambda x: x, + n_samples=lambda kwgs: ( + 2 if ("qk" in kwgs and kwgs["qk"] is not None) + else 1 + ), + n_outputs=1, result_to_tuple=lambda x, _: (x,), paired=True, + too_small=-1 # entropy doesn't have too small inputs +) +def entropy(pk: np.typing.ArrayLike, + qk: np.typing.ArrayLike | None = None, + base: float | None = None, + axis: int = 0 + ) -> np.number | np.ndarray: + """ + Calculate the Shannon entropy/relative entropy of given distribution(s). + + If only probabilities `pk` are given, the Shannon entropy is calculated as + ``H = -sum(pk * log(pk))``. + + If `qk` is not None, then compute the relative entropy + ``D = sum(pk * log(pk / qk))``. This quantity is also known + as the Kullback-Leibler divergence. + + This routine will normalize `pk` and `qk` if they don't sum to 1. + + Parameters + ---------- + pk : array_like + Defines the (discrete) distribution. Along each axis-slice of ``pk``, + element ``i`` is the (possibly unnormalized) probability of event + ``i``. + qk : array_like, optional + Sequence against which the relative entropy is computed. Should be in + the same format as `pk`. + base : float, optional + The logarithmic base to use, defaults to ``e`` (natural logarithm). + axis : int, optional + The axis along which the entropy is calculated. Default is 0. + + Returns + ------- + S : {float, array_like} + The calculated entropy. + + Notes + ----- + Informally, the Shannon entropy quantifies the expected uncertainty + inherent in the possible outcomes of a discrete random variable. + For example, + if messages consisting of sequences of symbols from a set are to be + encoded and transmitted over a noiseless channel, then the Shannon entropy + ``H(pk)`` gives a tight lower bound for the average number of units of + information needed per symbol if the symbols occur with frequencies + governed by the discrete distribution `pk` [1]_. The choice of base + determines the choice of units; e.g., ``e`` for nats, ``2`` for bits, etc. + + The relative entropy, ``D(pk|qk)``, quantifies the increase in the average + number of units of information needed per symbol if the encoding is + optimized for the probability distribution `qk` instead of the true + distribution `pk`. Informally, the relative entropy quantifies the expected + excess in surprise experienced if one believes the true distribution is + `qk` when it is actually `pk`. + + A related quantity, the cross entropy ``CE(pk, qk)``, satisfies the + equation ``CE(pk, qk) = H(pk) + D(pk|qk)`` and can also be calculated with + the formula ``CE = -sum(pk * log(qk))``. It gives the average + number of units of information needed per symbol if an encoding is + optimized for the probability distribution `qk` when the true distribution + is `pk`. It is not computed directly by `entropy`, but it can be computed + using two calls to the function (see Examples). + + See [2]_ for more information. + + References + ---------- + .. [1] Shannon, C.E. (1948), A Mathematical Theory of Communication. + Bell System Technical Journal, 27: 379-423. + :doi:`10.1002/j.1538-7305.1948.tb01338.x`. + .. [2] Thomas M. Cover and Joy A. Thomas. 2006. Elements of Information + Theory (Wiley Series in Telecommunications and Signal Processing). + Wiley-Interscience, USA. + + + Examples + -------- + The outcome of a fair coin is the most uncertain: + + >>> import numpy as np + >>> from scipy.stats import entropy + >>> base = 2 # work in units of bits + >>> pk = np.array([1/2, 1/2]) # fair coin + >>> H = entropy(pk, base=base) + >>> H + 1.0 + >>> H == -np.sum(pk * np.log(pk)) / np.log(base) + True + + The outcome of a biased coin is less uncertain: + + >>> qk = np.array([9/10, 1/10]) # biased coin + >>> entropy(qk, base=base) + 0.46899559358928117 + + The relative entropy between the fair coin and biased coin is calculated + as: + + >>> D = entropy(pk, qk, base=base) + >>> D + 0.7369655941662062 + >>> np.isclose(D, np.sum(pk * np.log(pk/qk)) / np.log(base), rtol=4e-16, atol=0) + True + + The cross entropy can be calculated as the sum of the entropy and + relative entropy`: + + >>> CE = entropy(pk, base=base) + entropy(pk, qk, base=base) + >>> CE + 1.736965594166206 + >>> CE == -np.sum(pk * np.log(qk)) / np.log(base) + True + + """ + if base is not None and base <= 0: + raise ValueError("`base` must be a positive number or `None`.") + + xp = array_namespace(pk, qk) + pk, qk = xp_promote(pk, qk, broadcast=True, xp=xp) + + with np.errstate(invalid='ignore'): + if qk is not None: + pk, qk = _share_masks(pk, qk, xp=xp) + qk = qk / xp.sum(qk, axis=axis, keepdims=True) + pk = pk / xp.sum(pk, axis=axis, keepdims=True) + + if qk is None: + vec = special.entr(pk) + else: + vec = _masked_apply(special.rel_entr, args=(pk, qk), xp=xp) + + S = xp.sum(vec, axis=axis) + if base is not None: + S /= math.log(base) + return S + + +def _differential_entropy_is_too_small(samples, kwargs, axis=-1): + values = samples[0] + n = values.shape[axis] + window_length = kwargs.get("window_length") + if window_length is None: + window_length = math.floor(math.sqrt(n) + 0.5) + if not 2 <= 2 * window_length < n: + return True + return False + + +@xp_capabilities() +@_axis_nan_policy_factory( + lambda x: x, n_outputs=1, result_to_tuple=lambda x, _: (x,), + too_small=_differential_entropy_is_too_small +) +def differential_entropy( + values: np.typing.ArrayLike, + *, + window_length: int | None = None, + base: float | None = None, + axis: int = 0, + method: str = "auto", +) -> np.number | np.ndarray: + r"""Given a sample of a distribution, estimate the differential entropy. + + Several estimation methods are available using the `method` parameter. By + default, a method is selected based the size of the sample. + + Parameters + ---------- + values : sequence + Sample from a continuous distribution. + window_length : int, optional + Window length for computing Vasicek estimate. Must be an integer + between 1 and half of the sample size. If ``None`` (the default), it + uses the heuristic value + + .. math:: + \left \lfloor \sqrt{n} + 0.5 \right \rfloor + + where :math:`n` is the sample size. This heuristic was originally + proposed in [2]_ and has become common in the literature. + base : float, optional + The logarithmic base to use, defaults to ``e`` (natural logarithm). + axis : int, optional + The axis along which the differential entropy is calculated. + Default is 0. + method : {'vasicek', 'van es', 'ebrahimi', 'correa', 'auto'}, optional + The method used to estimate the differential entropy from the sample. + Default is ``'auto'``. See Notes for more information. + + Returns + ------- + entropy : float + The calculated differential entropy. + + Notes + ----- + This function will converge to the true differential entropy in the limit + + .. math:: + n \to \infty, \quad m \to \infty, \quad \frac{m}{n} \to 0 + + The optimal choice of ``window_length`` for a given sample size depends on + the (unknown) distribution. Typically, the smoother the density of the + distribution, the larger the optimal value of ``window_length`` [1]_. + + The following options are available for the `method` parameter. + + * ``'vasicek'`` uses the estimator presented in [1]_. This is + one of the first and most influential estimators of differential entropy. + * ``'van es'`` uses the bias-corrected estimator presented in [3]_, which + is not only consistent but, under some conditions, asymptotically normal. + * ``'ebrahimi'`` uses an estimator presented in [4]_, which was shown + in simulation to have smaller bias and mean squared error than + the Vasicek estimator. + * ``'correa'`` uses the estimator presented in [5]_ based on local linear + regression. In a simulation study, it had consistently smaller mean + square error than the Vasiceck estimator, but it is more expensive to + compute. + * ``'auto'`` selects the method automatically (default). Currently, + this selects ``'van es'`` for very small samples (<10), ``'ebrahimi'`` + for moderate sample sizes (11-1000), and ``'vasicek'`` for larger + samples, but this behavior is subject to change in future versions. + + All estimators are implemented as described in [6]_. + + References + ---------- + .. [1] Vasicek, O. (1976). A test for normality based on sample entropy. + Journal of the Royal Statistical Society: + Series B (Methodological), 38(1), 54-59. + .. [2] Crzcgorzewski, P., & Wirczorkowski, R. (1999). Entropy-based + goodness-of-fit test for exponentiality. Communications in + Statistics-Theory and Methods, 28(5), 1183-1202. + .. [3] Van Es, B. (1992). Estimating functionals related to a density by a + class of statistics based on spacings. Scandinavian Journal of + Statistics, 61-72. + .. [4] Ebrahimi, N., Pflughoeft, K., & Soofi, E. S. (1994). Two measures + of sample entropy. Statistics & Probability Letters, 20(3), 225-234. + .. [5] Correa, J. C. (1995). A new estimator of entropy. Communications + in Statistics-Theory and Methods, 24(10), 2439-2449. + .. [6] Noughabi, H. A. (2015). Entropy Estimation Using Numerical Methods. + Annals of Data Science, 2(2), 231-241. + https://link.springer.com/article/10.1007/s40745-015-0045-9 + + Examples + -------- + >>> import numpy as np + >>> from scipy.stats import differential_entropy, norm + + Entropy of a standard normal distribution: + + >>> rng = np.random.default_rng() + >>> values = rng.standard_normal(100) + >>> differential_entropy(values) + 1.3407817436640392 + + Compare with the true entropy: + + >>> float(norm.entropy()) + 1.4189385332046727 + + For several sample sizes between 5 and 1000, compare the accuracy of + the ``'vasicek'``, ``'van es'``, and ``'ebrahimi'`` methods. Specifically, + compare the root mean squared error (over 1000 trials) between the estimate + and the true differential entropy of the distribution. + + >>> from scipy import stats + >>> import matplotlib.pyplot as plt + >>> + >>> + >>> def rmse(res, expected): + ... '''Root mean squared error''' + ... return np.sqrt(np.mean((res - expected)**2)) + >>> + >>> + >>> a, b = np.log10(5), np.log10(1000) + >>> ns = np.round(np.logspace(a, b, 10)).astype(int) + >>> reps = 1000 # number of repetitions for each sample size + >>> expected = stats.expon.entropy() + >>> + >>> method_errors = {'vasicek': [], 'van es': [], 'ebrahimi': []} + >>> for method in method_errors: + ... for n in ns: + ... rvs = stats.expon.rvs(size=(reps, n), random_state=rng) + ... res = stats.differential_entropy(rvs, method=method, axis=-1) + ... error = rmse(res, expected) + ... method_errors[method].append(error) + >>> + >>> for method, errors in method_errors.items(): + ... plt.loglog(ns, errors, label=method) + >>> + >>> plt.legend() + >>> plt.xlabel('sample size') + >>> plt.ylabel('RMSE (1000 trials)') + >>> plt.title('Entropy Estimator Error (Exponential Distribution)') + + """ + xp = array_namespace(values) + values = xp_promote(values, force_floating=True, xp=xp) + values = xp.moveaxis(values, axis, -1) + n = values.shape[-1] # type: ignore[union-attr] + + if window_length is None: + window_length = math.floor(math.sqrt(n) + 0.5) + + if not 2 <= 2 * window_length < n: + raise ValueError( + f"Window length ({window_length}) must be positive and less " + f"than half the sample size ({n}).", + ) + + if base is not None and base <= 0: + raise ValueError("`base` must be a positive number or `None`.") + + sorted_data = xp.sort(values, axis=-1) + + methods = {"vasicek": _vasicek_entropy, + "van es": _van_es_entropy, + "correa": _correa_entropy, + "ebrahimi": _ebrahimi_entropy, + "auto": _vasicek_entropy} + method = method.lower() + if method not in methods: + message = f"`method` must be one of {set(methods)}" + raise ValueError(message) + + if method == "auto": + if n <= 10: + method = 'van es' + elif n <= 1000: + method = 'ebrahimi' + else: + method = 'vasicek' + + res = methods[method](sorted_data, window_length, xp=xp) + + if base is not None: + res /= math.log(base) + + # avoid dtype changes due to data-apis/array-api-compat#152 + # can be removed when data-apis/array-api-compat#152 is resolved + return xp.astype(res, values.dtype) # type: ignore[union-attr] + + +def _pad_along_last_axis(X, m, *, xp): + """Pad the data for computing the rolling window difference.""" + # scales a bit better than method in _vasicek_like_entropy + shape = X.shape[:-1] + (m,) + Xl = xp.broadcast_to(X[..., :1], shape) # :1 vs 0 to maintain shape + Xr = xp.broadcast_to(X[..., -1:], shape) + return xp.concat((Xl, X, Xr), axis=-1) + + +def _vasicek_entropy(X, m, *, xp): + """Compute the Vasicek estimator as described in [6] Eq. 1.3.""" + n = X.shape[-1] + X = _pad_along_last_axis(X, m, xp=xp) + differences = X[..., 2 * m:] - X[..., : -2 * m:] + logs = xp.log(n/(2*m) * differences) + return xp.mean(logs, axis=-1) + + +def _van_es_entropy(X, m, *, xp): + """Compute the van Es estimator as described in [6].""" + # No equation number, but referred to as HVE_mn. + # Typo: there should be a log within the summation. + n = X.shape[-1] + difference = X[..., m:] - X[..., :-m] + term1 = 1/(n-m) * xp.sum(xp.log((n+1)/m * difference), axis=-1) + k = xp.arange(m, n+1, dtype=term1.dtype, device=xp_device(X)) + return term1 + xp.sum(1/k) + math.log(m) - math.log(n+1) + + +def _ebrahimi_entropy(X, m, *, xp): + """Compute the Ebrahimi estimator as described in [6].""" + # No equation number, but referred to as HE_mn + n = X.shape[-1] + X = _pad_along_last_axis(X, m, xp=xp) + + differences = X[..., 2 * m:] - X[..., : -2 * m:] + + i = xp.arange(1, n+1, dtype=X.dtype, device=xp_device(X)) + ci = xp.where(i <= m, 1 + (i - 1)/m, 2.) + ci = xp.where(i >= n - m + 1, 1 + (n - i)/m, ci) + + logs = xp.log(n * differences / (ci * m)) + return xp.mean(logs, axis=-1) + + +def _correa_entropy(X, m, *, xp): + """Compute the Correa estimator as described in [6].""" + # No equation number, but referred to as HC_mn + n = X.shape[-1] + X = _pad_along_last_axis(X, m, xp=xp) + + i = xp.arange(1, n+1, device=xp_device(X)) + dj = xp.arange(-m, m+1, device=xp_device(X))[:, None] + j = i + dj + j0 = j + m - 1 # 0-indexed version of j + + Xibar = xp.mean(X[..., j0], axis=-2, keepdims=True) + difference = X[..., j0] - Xibar + num = xp.sum(difference*dj, axis=-2) # dj is d-i + den = n*xp.sum(difference**2, axis=-2) + return -xp.mean(xp.log(num/den), axis=-1) diff --git a/v12x/pylib/scipy/stats/_finite_differences.py b/v12x/pylib/scipy/stats/_finite_differences.py new file mode 100644 index 0000000000000000000000000000000000000000..6ec91654c4e95f73f32edc5aa357a4e3f80fff1f --- /dev/null +++ b/v12x/pylib/scipy/stats/_finite_differences.py @@ -0,0 +1,145 @@ +from numpy import arange, newaxis, hstack, prod, array +from scipy import linalg + + +def _central_diff_weights(Np, ndiv=1): + """ + Return weights for an Np-point central derivative. + + Assumes equally-spaced function points. + + If weights are in the vector w, then + derivative is w[0] * f(x-ho*dx) + ... + w[-1] * f(x+h0*dx) + + Parameters + ---------- + Np : int + Number of points for the central derivative. + ndiv : int, optional + Number of divisions. Default is 1. + + Returns + ------- + w : ndarray + Weights for an Np-point central derivative. Its size is `Np`. + + Notes + ----- + Can be inaccurate for a large number of points. + + Examples + -------- + We can calculate a derivative value of a function. + + >>> def f(x): + ... return 2 * x**2 + 3 + >>> x = 3.0 # derivative point + >>> h = 0.1 # differential step + >>> Np = 3 # point number for central derivative + >>> weights = _central_diff_weights(Np) # weights for first derivative + >>> vals = [f(x + (i - Np/2) * h) for i in range(Np)] + >>> sum(w * v for (w, v) in zip(weights, vals))/h + 11.79999999999998 + + This value is close to the analytical solution: + f'(x) = 4x, so f'(3) = 12 + + References + ---------- + .. [1] https://en.wikipedia.org/wiki/Finite_difference + + """ + if Np < ndiv + 1: + raise ValueError( + "Number of points must be at least the derivative order + 1." + ) + if Np % 2 == 0: + raise ValueError("The number of points must be odd.") + + ho = Np >> 1 + x = arange(-ho, ho + 1.0) + x = x[:, newaxis] + X = x**0.0 + for k in range(1, Np): + X = hstack([X, x**k]) + w = prod(arange(1, ndiv + 1), axis=0) * linalg.inv(X)[ndiv] + return w + + +def _derivative(func, x0, dx=1.0, n=1, args=(), order=3): + """ + Find the nth derivative of a function at a point. + + Given a function, use a central difference formula with spacing `dx` to + compute the nth derivative at `x0`. + + Parameters + ---------- + func : function + Input function. + x0 : float + The point at which the nth derivative is found. + dx : float, optional + Spacing. + n : int, optional + Order of the derivative. Default is 1. + args : tuple, optional + Arguments + order : int, optional + Number of points to use, must be odd. + + Notes + ----- + Decreasing the step size too small can result in round-off error. + + Examples + -------- + >>> def f(x): + ... return x**3 + x**2 + >>> _derivative(f, 1.0, dx=1e-6) + 4.9999999999217337 + + """ + if order < n + 1: + raise ValueError( + "'order' (the number of points used to compute the derivative), " + "must be at least the derivative order 'n' + 1." + ) + if order % 2 == 0: + raise ValueError( + "'order' (the number of points used to compute the derivative) " + "must be odd." + ) + # pre-computed for n=1 and 2 and low-order for speed. + if n == 1: + if order == 3: + weights = array([-1, 0, 1]) / 2.0 + elif order == 5: + weights = array([1, -8, 0, 8, -1]) / 12.0 + elif order == 7: + weights = array([-1, 9, -45, 0, 45, -9, 1]) / 60.0 + elif order == 9: + weights = array([3, -32, 168, -672, 0, 672, -168, 32, -3]) / 840.0 + else: + weights = _central_diff_weights(order, 1) + elif n == 2: + if order == 3: + weights = array([1, -2.0, 1]) + elif order == 5: + weights = array([-1, 16, -30, 16, -1]) / 12.0 + elif order == 7: + weights = array([2, -27, 270, -490, 270, -27, 2]) / 180.0 + elif order == 9: + weights = ( + array([-9, 128, -1008, 8064, -14350, 8064, -1008, 128, -9]) + / 5040.0 + ) + else: + weights = _central_diff_weights(order, 2) + else: + weights = _central_diff_weights(order, n) + val = 0.0 + ho = order >> 1 + for k in range(order): + val += weights[k] * func(x0 + (k - ho) * dx, *args) + return val / prod((dx,) * n, axis=0) diff --git a/v12x/pylib/scipy/stats/_fit.py b/v12x/pylib/scipy/stats/_fit.py new file mode 100644 index 0000000000000000000000000000000000000000..a73e410753dcfc9992e31f9395bbe72e45077749 --- /dev/null +++ b/v12x/pylib/scipy/stats/_fit.py @@ -0,0 +1,1349 @@ +import warnings +from collections import namedtuple +import numpy as np +from scipy import optimize, stats +from scipy._lib._array_api import xp_capabilities +from scipy._lib._util import check_random_state, _transition_to_rng + + +def _combine_bounds(name, user_bounds, shape_domain, integral): + """Intersection of user-defined bounds and distribution PDF/PMF domain""" + + user_bounds = np.atleast_1d(user_bounds) + + if user_bounds[0] > user_bounds[1]: + message = (f"There are no values for `{name}` on the interval " + f"{list(user_bounds)}.") + raise ValueError(message) + + bounds = (max(user_bounds[0], shape_domain[0]), + min(user_bounds[1], shape_domain[1])) + + if integral and (np.ceil(bounds[0]) > np.floor(bounds[1])): + message = (f"There are no integer values for `{name}` on the interval " + f"defined by the user-provided bounds and the domain " + "of the distribution.") + raise ValueError(message) + elif not integral and (bounds[0] > bounds[1]): + message = (f"There are no values for `{name}` on the interval " + f"defined by the user-provided bounds and the domain " + "of the distribution.") + raise ValueError(message) + + if not np.all(np.isfinite(bounds)): + message = (f"The intersection of user-provided bounds for `{name}` " + f"and the domain of the distribution is not finite. Please " + f"provide finite bounds for shape `{name}` in `bounds`.") + raise ValueError(message) + + return bounds + + +class FitResult: + r"""Result of fitting a discrete or continuous distribution to data + + Attributes + ---------- + params : namedtuple + A namedtuple containing the maximum likelihood estimates of the + shape parameters, location, and (if applicable) scale of the + distribution. + success : bool or None + Whether the optimizer considered the optimization to terminate + successfully or not. + message : str or None + Any status message provided by the optimizer. + + """ + + def __init__(self, dist, data, discrete, res): + self._dist = dist + self._data = data + self.discrete = discrete + self.pxf = getattr(dist, "pmf", None) or getattr(dist, "pdf", None) + + shape_names = [] if dist.shapes is None else dist.shapes.split(", ") + if not discrete: + FitParams = namedtuple('FitParams', shape_names + ['loc', 'scale']) + else: + FitParams = namedtuple('FitParams', shape_names + ['loc']) + + self.params = FitParams(*res.x) + + # Optimizer can report success even when nllf is infinite + if res.success and not np.isfinite(self.nllf()): + res.success = False + res.message = ("Optimization converged to parameter values that " + "are inconsistent with the data.") + self.success = getattr(res, "success", None) + self.message = getattr(res, "message", None) + + def __repr__(self): + keys = ["params", "success", "message"] + m = max(map(len, keys)) + 1 + return '\n'.join([key.rjust(m) + ': ' + repr(getattr(self, key)) + for key in keys if getattr(self, key) is not None]) + + def nllf(self, params=None, data=None): + """Negative log-likelihood function + + Evaluates the negative of the log-likelihood function of the provided + data at the provided parameters. + + Parameters + ---------- + params : tuple, optional + The shape parameters, location, and (if applicable) scale of the + distribution as a single tuple. Default is the maximum likelihood + estimates (``self.params``). + data : array_like, optional + The data for which the log-likelihood function is to be evaluated. + Default is the data to which the distribution was fit. + + Returns + ------- + nllf : float + The negative of the log-likelihood function. + + """ + params = params if params is not None else self.params + data = data if data is not None else self._data + return self._dist.nnlf(theta=params, x=data) + + def plot(self, ax=None, *, plot_type="hist"): + """Visually compare the data against the fitted distribution. + + Available only if `matplotlib` is installed. + + Parameters + ---------- + ax : `matplotlib.axes.Axes` + Axes object to draw the plot onto, otherwise uses the current Axes. + plot_type : {"hist", "qq", "pp", "cdf"} + Type of plot to draw. Options include: + + - "hist": Superposes the PDF/PMF of the fitted distribution + over a normalized histogram of the data. + - "qq": Scatter plot of theoretical quantiles against the + empirical quantiles. Specifically, the x-coordinates are the + values of the fitted distribution PPF evaluated at the + percentiles ``(np.arange(1, n) - 0.5)/n``, where ``n`` is the + number of data points, and the y-coordinates are the sorted + data points. + - "pp": Scatter plot of theoretical percentiles against the + observed percentiles. Specifically, the x-coordinates are the + percentiles ``(np.arange(1, n) - 0.5)/n``, where ``n`` is + the number of data points, and the y-coordinates are the values + of the fitted distribution CDF evaluated at the sorted + data points. + - "cdf": Superposes the CDF of the fitted distribution over the + empirical CDF. Specifically, the x-coordinates of the empirical + CDF are the sorted data points, and the y-coordinates are the + percentiles ``(np.arange(1, n) - 0.5)/n``, where ``n`` is + the number of data points. + + Returns + ------- + ax : `matplotlib.axes.Axes` + The matplotlib Axes object on which the plot was drawn. + + Examples + -------- + >>> import numpy as np + >>> from scipy import stats + >>> import matplotlib.pyplot as plt # matplotlib must be installed + >>> rng = np.random.default_rng() + >>> data = stats.nbinom(5, 0.5).rvs(size=1000, random_state=rng) + >>> bounds = [(0, 30), (0, 1)] + >>> res = stats.fit(stats.nbinom, data, bounds) + >>> ax = res.plot() # save matplotlib Axes object + + The `matplotlib.axes.Axes` object can be used to customize the plot. + See `matplotlib.axes.Axes` documentation for details. + + >>> ax.set_xlabel('number of trials') # customize axis label + >>> ax.get_children()[0].set_linewidth(5) # customize line widths + >>> ax.legend() + >>> plt.show() + """ + try: + import matplotlib # noqa: F401 + except ModuleNotFoundError as exc: + message = "matplotlib must be installed to use method `plot`." + raise ModuleNotFoundError(message) from exc + + plots = {'histogram': self._hist_plot, 'qq': self._qq_plot, + 'pp': self._pp_plot, 'cdf': self._cdf_plot, + 'hist': self._hist_plot} + if plot_type.lower() not in plots: + message = f"`plot_type` must be one of {set(plots.keys())}" + raise ValueError(message) + plot = plots[plot_type.lower()] + + if ax is None: + import matplotlib.pyplot as plt + ax = plt.gca() + + fit_params = np.atleast_1d(self.params) + + return plot(ax=ax, fit_params=fit_params) + + def _hist_plot(self, ax, fit_params): + from matplotlib.ticker import MaxNLocator + + support = self._dist.support(*fit_params) + lb = support[0] if np.isfinite(support[0]) else min(self._data) + ub = support[1] if np.isfinite(support[1]) else max(self._data) + pxf = "PMF" if self.discrete else "PDF" + + if self.discrete: + x = np.arange(lb, ub + 2) + y = self.pxf(x, *fit_params) + ax.vlines(x[:-1], 0, y[:-1], label='Fitted Distribution PMF', + color='C0') + options = dict(density=True, bins=x, align='left', color='C1') + ax.xaxis.set_major_locator(MaxNLocator(integer=True)) + ax.set_xlabel('k') + ax.set_ylabel('PMF') + else: + x = np.linspace(lb, ub, 200) + y = self.pxf(x, *fit_params) + ax.plot(x, y, '--', label='Fitted Distribution PDF', color='C0') + options = dict(density=True, bins=50, align='mid', color='C1') + ax.set_xlabel('x') + ax.set_ylabel('PDF') + + if len(self._data) > 50 or self.discrete: + ax.hist(self._data, label="Histogram of Data", **options) + else: + ax.plot(self._data, np.zeros_like(self._data), "*", + label='Data', color='C1') + + ax.set_title(rf"Fitted $\tt {self._dist.name}$ {pxf} and Histogram") + ax.legend(*ax.get_legend_handles_labels()) + return ax + + def _qp_plot(self, ax, fit_params, qq): + data = np.sort(self._data) + ps = self._plotting_positions(len(self._data)) + + if qq: + qp = "Quantiles" + plot_type = 'Q-Q' + x = self._dist.ppf(ps, *fit_params) + y = data + else: + qp = "Percentiles" + plot_type = 'P-P' + x = ps + y = self._dist.cdf(data, *fit_params) + + ax.plot(x, y, '.', label=f'Fitted Distribution {plot_type}', + color='C0', zorder=1) + xlim = ax.get_xlim() + ylim = ax.get_ylim() + lim = [min(xlim[0], ylim[0]), max(xlim[1], ylim[1])] + if not qq: + lim = max(lim[0], 0), min(lim[1], 1) + + if self.discrete and qq: + q_min, q_max = int(lim[0]), int(lim[1]+1) + q_ideal = np.arange(q_min, q_max) + # q_ideal = np.unique(self._dist.ppf(ps, *fit_params)) + ax.plot(q_ideal, q_ideal, 'o', label='Reference', color='k', + alpha=0.25, markerfacecolor='none', clip_on=True) + elif self.discrete and not qq: + # The intent of this is to match the plot that would be produced + # if x were continuous on [0, 1] and y were cdf(ppf(x)). + # It can be approximated by letting x = np.linspace(0, 1, 1000), + # but this might not look great when zooming in. The vertical + # portions are included to indicate where the transition occurs + # where the data completely obscures the horizontal portions. + p_min, p_max = lim + a, b = self._dist.support(*fit_params) + p_min = max(p_min, 0 if np.isfinite(a) else 1e-3) + p_max = min(p_max, 1 if np.isfinite(b) else 1-1e-3) + q_min, q_max = self._dist.ppf([p_min, p_max], *fit_params) + qs = np.arange(q_min-1, q_max+1) + ps = self._dist.cdf(qs, *fit_params) + ax.step(ps, ps, '-', label='Reference', color='k', alpha=0.25, + clip_on=True) + else: + ax.plot(lim, lim, '-', label='Reference', color='k', alpha=0.25, + clip_on=True) + + ax.set_xlim(lim) + ax.set_ylim(lim) + ax.set_xlabel(rf"Fitted $\tt {self._dist.name}$ Theoretical {qp}") + ax.set_ylabel(f"Data {qp}") + ax.set_title(rf"Fitted $\tt {self._dist.name}$ {plot_type} Plot") + ax.legend(*ax.get_legend_handles_labels()) + ax.set_aspect('equal') + return ax + + def _qq_plot(self, **kwargs): + return self._qp_plot(qq=True, **kwargs) + + def _pp_plot(self, **kwargs): + return self._qp_plot(qq=False, **kwargs) + + def _plotting_positions(self, n, a=.5): + # See https://en.wikipedia.org/wiki/Q%E2%80%93Q_plot#Plotting_positions + k = np.arange(1, n+1) + return (k-a) / (n + 1 - 2*a) + + def _cdf_plot(self, ax, fit_params): + data = np.sort(self._data) + ecdf = self._plotting_positions(len(self._data)) + ls = '--' if len(np.unique(data)) < 30 else '.' + xlabel = 'k' if self.discrete else 'x' + ax.step(data, ecdf, ls, label='Empirical CDF', color='C1', zorder=0) + + xlim = ax.get_xlim() + q = np.linspace(*xlim, 300) + tcdf = self._dist.cdf(q, *fit_params) + + ax.plot(q, tcdf, label='Fitted Distribution CDF', color='C0', zorder=1) + ax.set_xlim(xlim) + ax.set_ylim(0, 1) + ax.set_xlabel(xlabel) + ax.set_ylabel("CDF") + ax.set_title(rf"Fitted $\tt {self._dist.name}$ and Empirical CDF") + handles, labels = ax.get_legend_handles_labels() + ax.legend(handles[::-1], labels[::-1]) + return ax + + +@xp_capabilities(out_of_scope=True) +def fit(dist, data, bounds=None, *, guess=None, method='mle', + optimizer=optimize.differential_evolution): + r"""Fit a discrete or continuous distribution to data. + + Given a distribution, data, and bounds on the parameters of the + distribution, return maximum likelihood estimates of the parameters. + + Parameters + ---------- + dist : `scipy.stats.rv_continuous` or `scipy.stats.rv_discrete` + The object representing the distribution to be fit to the data. + data : 1D array_like + The data to which the distribution is to be fit. If the data contain + any of ``np.nan``, ``np.inf``, or -``np.inf``, the fit method will + raise a ``ValueError``. + bounds : dict or sequence of tuples, optional + If a dictionary, each key is the name of a parameter of the + distribution, and the corresponding value is a tuple containing the + lower and upper bound on that parameter. If the distribution is + defined only for a finite range of values of that parameter, no entry + for that parameter is required; e.g., some distributions have + parameters which must be on the interval [0, 1]. Bounds for parameters + location (``loc``) and scale (``scale``) are optional; by default, + they are fixed to 0 and 1, respectively. + + If a sequence, element *i* is a tuple containing the lower and upper + bound on the *i*\ th parameter of the distribution. In this case, + bounds for *all* distribution shape parameters must be provided. + Optionally, bounds for location and scale may follow the + distribution shape parameters. + + If a shape is to be held fixed (e.g. if it is known), the + lower and upper bounds may be equal. If a user-provided lower or upper + bound is beyond a bound of the domain for which the distribution is + defined, the bound of the distribution's domain will replace the + user-provided value. Similarly, parameters which must be integral + will be constrained to integral values within the user-provided bounds. + guess : dict or array_like, optional + If a dictionary, each key is the name of a parameter of the + distribution, and the corresponding value is a guess for the value + of the parameter. + + If a sequence, element *i* is a guess for the *i*\ th parameter of the + distribution. In this case, guesses for *all* distribution shape + parameters must be provided. + + If `guess` is not provided, guesses for the decision variables will + not be passed to the optimizer. If `guess` is provided, guesses for + any missing parameters will be set at the mean of the lower and + upper bounds. Guesses for parameters which must be integral will be + rounded to integral values, and guesses that lie outside the + intersection of the user-provided bounds and the domain of the + distribution will be clipped. + method : {'mle', 'mse'} + With ``method="mle"`` (default), the fit is computed by minimizing + the negative log-likelihood function. A large, finite penalty + (rather than infinite negative log-likelihood) is applied for + observations beyond the support of the distribution. + With ``method="mse"``, the fit is computed by minimizing + the negative log-product spacing function. The same penalty is applied + for observations beyond the support. We follow the approach of [1]_, + which is generalized for samples with repeated observations. + optimizer : callable, optional + `optimizer` is a callable that accepts the following positional + argument. + + fun : callable + The objective function to be optimized. `fun` accepts one argument + ``x``, candidate shape parameters of the distribution, and returns + the objective function value given ``x``, `dist`, and the provided + `data`. + The job of `optimizer` is to find values of the decision variables + that minimizes `fun`. + + `optimizer` must also accept the following keyword argument. + + bounds : sequence of tuples + The bounds on values of the decision variables; each element will + be a tuple containing the lower and upper bound on a decision + variable. + + If `guess` is provided, `optimizer` must also accept the following + keyword argument. + + x0 : array_like + The guesses for each decision variable. + + If the distribution has any shape parameters that must be integral or + if the distribution is discrete and the location parameter is not + fixed, `optimizer` must also accept the following keyword argument. + + integrality : array_like of bools + For each decision variable, True if the decision variable + must be constrained to integer values and False if the decision + variable is continuous. + + `optimizer` must return an object, such as an instance of + `scipy.optimize.OptimizeResult`, which holds the optimal values of + the decision variables in an attribute ``x``. If attributes + ``fun``, ``status``, or ``message`` are provided, they will be + included in the result object returned by `fit`. + + Returns + ------- + result : `~scipy.stats._result_classes.FitResult` + An object with the following fields. + + params : namedtuple + A namedtuple containing the maximum likelihood estimates of the + shape parameters, location, and (if applicable) scale of the + distribution. + success : bool or None + Whether the optimizer considered the optimization to terminate + successfully or not. + message : str or None + Any status message provided by the optimizer. + + The object has the following method: + + nllf(params=None, data=None) + By default, the negative log-likelihood function at the fitted + `params` for the given `data`. Accepts a tuple containing + alternative shapes, location, and scale of the distribution and + an array of alternative data. + + plot(ax=None) + Superposes the PDF/PMF of the fitted distribution over a normalized + histogram of the data. + + See Also + -------- + rv_continuous, rv_discrete + + Notes + ----- + Optimization is more likely to converge to the maximum likelihood estimate + when the user provides tight bounds containing the maximum likelihood + estimate. For example, when fitting a binomial distribution to data, the + number of experiments underlying each sample may be known, in which case + the corresponding shape parameter ``n`` can be fixed. + + References + ---------- + .. [1] Shao, Yongzhao, and Marjorie G. Hahn. "Maximum product of spacings + method: a unified formulation with illustration of strong + consistency." Illinois Journal of Mathematics 43.3 (1999): 489-499. + + Examples + -------- + Suppose we wish to fit a distribution to the following data. + + >>> import numpy as np + >>> from scipy import stats + >>> rng = np.random.default_rng() + >>> dist = stats.nbinom + >>> shapes = (5, 0.5) + >>> data = dist.rvs(*shapes, size=1000, random_state=rng) + + Suppose we do not know how the data were generated, but we suspect that + it follows a negative binomial distribution with parameters *n* and *p*\. + (See `scipy.stats.nbinom`.) We believe that the parameter *n* was fewer + than 30, and we know that the parameter *p* must lie on the interval + [0, 1]. We record this information in a variable `bounds` and pass + this information to `fit`. + + >>> bounds = [(0, 30), (0, 1)] + >>> res = stats.fit(dist, data, bounds) + + `fit` searches within the user-specified `bounds` for the + values that best match the data (in the sense of maximum likelihood + estimation). In this case, it found shape values similar to those + from which the data were actually generated. + + >>> res.params + FitParams(n=5.0, p=0.5028157644634368, loc=0.0) # may vary + + We can visualize the results by superposing the probability mass function + of the distribution (with the shapes fit to the data) over a normalized + histogram of the data. + + >>> import matplotlib.pyplot as plt # matplotlib must be installed to plot + >>> res.plot() + >>> plt.show() + + Note that the estimate for *n* was exactly integral; this is because + the domain of the `nbinom` PMF includes only integral *n*, and the `nbinom` + object "knows" that. `nbinom` also knows that the shape *p* must be a + value between 0 and 1. In such a case - when the domain of the distribution + with respect to a parameter is finite - we are not required to specify + bounds for the parameter. + + >>> bounds = {'n': (0, 30)} # omit parameter p using a `dict` + >>> res2 = stats.fit(dist, data, bounds) + >>> res2.params + FitParams(n=5.0, p=0.5016492009232932, loc=0.0) # may vary + + If we wish to force the distribution to be fit with *n* fixed at 6, we can + set both the lower and upper bounds on *n* to 6. Note, however, that the + value of the objective function being optimized is typically worse (higher) + in this case. + + >>> bounds = {'n': (6, 6)} # fix parameter `n` + >>> res3 = stats.fit(dist, data, bounds) + >>> res3.params + FitParams(n=6.0, p=0.5486556076755706, loc=0.0) # may vary + >>> res3.nllf() > res.nllf() + True # may vary + + Note that the numerical results of the previous examples are typical, but + they may vary because the default optimizer used by `fit`, + `scipy.optimize.differential_evolution`, is stochastic. However, we can + customize the settings used by the optimizer to ensure reproducibility - + or even use a different optimizer entirely - using the `optimizer` + parameter. + + >>> from scipy.optimize import differential_evolution + >>> rng = np.random.default_rng(767585560716548) + >>> def optimizer(fun, bounds, *, integrality): + ... return differential_evolution(fun, bounds, strategy='best2bin', + ... rng=rng, integrality=integrality) + >>> bounds = [(0, 30), (0, 1)] + >>> res4 = stats.fit(dist, data, bounds, optimizer=optimizer) + >>> res4.params + FitParams(n=5.0, p=0.5015183149259951, loc=0.0) + + """ + # --- Input Validation / Standardization --- # + user_bounds = bounds + user_guess = guess + + # distribution input validation and information collection + if hasattr(dist, "pdf"): # can't use isinstance for types + default_bounds = {'loc': (0, 0), 'scale': (1, 1)} + discrete = False + elif hasattr(dist, "pmf"): + default_bounds = {'loc': (0, 0)} + discrete = True + else: + message = ("`dist` must be an instance of `rv_continuous` " + "or `rv_discrete.`") + raise ValueError(message) + + try: + param_info = dist._param_info() + except AttributeError as e: + message = (f"Distribution `{dist.name}` is not yet supported by " + "`scipy.stats.fit` because shape information has " + "not been defined.") + raise ValueError(message) from e + + # data input validation + data = np.asarray(data) + if data.ndim != 1: + message = "`data` must be exactly one-dimensional." + raise ValueError(message) + if not (np.issubdtype(data.dtype, np.number) + and np.all(np.isfinite(data))): + message = "All elements of `data` must be finite numbers." + raise ValueError(message) + + # bounds input validation and information collection + n_params = len(param_info) + n_shapes = n_params - (1 if discrete else 2) + param_list = [param.name for param in param_info] + param_names = ", ".join(param_list) + shape_names = ", ".join(param_list[:n_shapes]) + + if user_bounds is None: + user_bounds = {} + + if isinstance(user_bounds, dict): + default_bounds.update(user_bounds) + user_bounds = default_bounds + user_bounds_array = np.empty((n_params, 2)) + for i in range(n_params): + param_name = param_info[i].name + user_bound = user_bounds.pop(param_name, None) + if user_bound is None: + user_bound = param_info[i].domain + user_bounds_array[i] = user_bound + if user_bounds: + message = ("Bounds provided for the following unrecognized " + f"parameters will be ignored: {set(user_bounds)}") + warnings.warn(message, RuntimeWarning, stacklevel=2) + + else: + try: + user_bounds = np.asarray(user_bounds, dtype=float) + if user_bounds.size == 0: + user_bounds = np.empty((0, 2)) + except ValueError as e: + message = ("Each element of a `bounds` sequence must be a tuple " + "containing two elements: the lower and upper bound of " + "a distribution parameter.") + raise ValueError(message) from e + if (user_bounds.ndim != 2 or user_bounds.shape[1] != 2): + message = ("Each element of `bounds` must be a tuple specifying " + "the lower and upper bounds of a shape parameter") + raise ValueError(message) + if user_bounds.shape[0] < n_shapes: + message = (f"A `bounds` sequence must contain at least {n_shapes} " + "elements: tuples specifying the lower and upper " + f"bounds of all shape parameters {shape_names}.") + raise ValueError(message) + if user_bounds.shape[0] > n_params: + message = ("A `bounds` sequence may not contain more than " + f"{n_params} elements: tuples specifying the lower and " + "upper bounds of distribution parameters " + f"{param_names}.") + raise ValueError(message) + + user_bounds_array = np.empty((n_params, 2)) + user_bounds_array[n_shapes:] = list(default_bounds.values()) + user_bounds_array[:len(user_bounds)] = user_bounds + + user_bounds = user_bounds_array + validated_bounds = [] + for i in range(n_params): + name = param_info[i].name + user_bound = user_bounds_array[i] + param_domain = param_info[i].domain + integral = param_info[i].integrality + combined = _combine_bounds(name, user_bound, param_domain, integral) + validated_bounds.append(combined) + + bounds = np.asarray(validated_bounds) + integrality = [param.integrality for param in param_info] + + # guess input validation + + if user_guess is None: + guess_array = None + elif isinstance(user_guess, dict): + default_guess = {param.name: np.mean(bound) + for param, bound in zip(param_info, bounds)} + unrecognized = set(user_guess) - set(default_guess) + if unrecognized: + message = ("Guesses provided for the following unrecognized " + f"parameters will be ignored: {unrecognized}") + warnings.warn(message, RuntimeWarning, stacklevel=2) + default_guess.update(user_guess) + + message = ("Each element of `guess` must be a scalar " + "guess for a distribution parameter.") + try: + guess_array = np.asarray([default_guess[param.name] + for param in param_info], dtype=float) + except ValueError as e: + raise ValueError(message) from e + + else: + message = ("Each element of `guess` must be a scalar " + "guess for a distribution parameter.") + try: + user_guess = np.asarray(user_guess, dtype=float) + except ValueError as e: + raise ValueError(message) from e + if user_guess.ndim != 1: + raise ValueError(message) + if user_guess.shape[0] < n_shapes: + message = (f"A `guess` sequence must contain at least {n_shapes} " + "elements: scalar guesses for the distribution shape " + f"parameters {shape_names}.") + raise ValueError(message) + if user_guess.shape[0] > n_params: + message = ("A `guess` sequence may not contain more than " + f"{n_params} elements: scalar guesses for the " + f"distribution parameters {param_names}.") + raise ValueError(message) + + guess_array = np.mean(bounds, axis=1) + guess_array[:len(user_guess)] = user_guess + + if guess_array is not None: + guess_rounded = guess_array.copy() + + guess_rounded[integrality] = np.round(guess_rounded[integrality]) + rounded = np.where(guess_rounded != guess_array)[0] + for i in rounded: + message = (f"Guess for parameter `{param_info[i].name}` " + f"rounded from {guess_array[i]} to {guess_rounded[i]}.") + warnings.warn(message, RuntimeWarning, stacklevel=2) + + guess_clipped = np.clip(guess_rounded, bounds[:, 0], bounds[:, 1]) + clipped = np.where(guess_clipped != guess_rounded)[0] + for i in clipped: + message = (f"Guess for parameter `{param_info[i].name}` " + f"clipped from {guess_rounded[i]} to " + f"{guess_clipped[i]}.") + warnings.warn(message, RuntimeWarning, stacklevel=2) + + guess = guess_clipped + else: + guess = None + + # --- Fitting --- # + def nllf(free_params, data=data): # bind data NOW + with np.errstate(invalid='ignore', divide='ignore'): + return dist._penalized_nnlf(free_params, data) + + def nlpsf(free_params, data=data): # bind data NOW + with np.errstate(invalid='ignore', divide='ignore'): + return dist._penalized_nlpsf(free_params, data) + + methods = {'mle': nllf, 'mse': nlpsf} + objective = methods[method.lower()] + + with np.errstate(invalid='ignore', divide='ignore'): + kwds = {} + if bounds is not None: + kwds['bounds'] = bounds + if np.any(integrality): + kwds['integrality'] = integrality + if guess is not None: + kwds['x0'] = guess + res = optimizer(objective, **kwds) + + return FitResult(dist, data, discrete, res) + + +GoodnessOfFitResult = namedtuple('GoodnessOfFitResult', + ('fit_result', 'statistic', 'pvalue', + 'null_distribution')) + + +@xp_capabilities(out_of_scope=True) +@_transition_to_rng('random_state') +def goodness_of_fit(dist, data, *, known_params=None, fit_params=None, + guessed_params=None, statistic='ad', n_mc_samples=9999, + rng=None): + r""" + Perform a goodness of fit test comparing data to a distribution family. + + Given a distribution family and data, perform a test of the null hypothesis + that the data were drawn from a distribution in that family. Any known + parameters of the distribution may be specified. Remaining parameters of + the distribution will be fit to the data, and the p-value of the test + is computed accordingly. Several statistics for comparing the distribution + to data are available. + + Parameters + ---------- + dist : `scipy.stats.rv_continuous` + The object representing the distribution family under the null + hypothesis. + data : 1D array_like + Finite, uncensored data to be tested. + known_params : dict, optional + A dictionary containing name-value pairs of known distribution + parameters. Monte Carlo samples are randomly drawn from the + null-hypothesized distribution with these values of the parameters. + Before the statistic is evaluated for the observed `data` and each + Monte Carlo sample, only remaining unknown parameters of the + null-hypothesized distribution family are fit to the samples; the + known parameters are held fixed. If all parameters of the distribution + family are known, then the step of fitting the distribution family to + each sample is omitted. + fit_params : dict, optional + A dictionary containing name-value pairs of distribution parameters + that have already been fit to the data, e.g. using `scipy.stats.fit` + or the ``fit`` method of `dist`. Monte Carlo samples are drawn from the + null-hypothesized distribution with these specified values of the + parameter. However, these and all other unknown parameters of the + null-hypothesized distribution family are always fit to the sample, + whether that is the observed `data` or a Monte Carlo sample, before + the statistic is evaluated. + guessed_params : dict, optional + A dictionary containing name-value pairs of distribution parameters + which have been guessed. These parameters are always considered as + free parameters and are fit both to the provided `data` as well as + to the Monte Carlo samples drawn from the null-hypothesized + distribution. The purpose of these `guessed_params` is to be used as + initial values for the numerical fitting procedure. + statistic : {"ad", "ks", "cvm", "filliben"} or callable, optional + The statistic used to compare data to a distribution after fitting + unknown parameters of the distribution family to the data. The + Anderson-Darling ("ad") [1]_, Kolmogorov-Smirnov ("ks") [1]_, + Cramer-von Mises ("cvm") [1]_, and Filliben ("filliben") [7]_ + statistics are available. Alternatively, a callable with signature + ``(dist, data, axis)`` may be supplied to compute the statistic. Here + ``dist`` is a frozen distribution object (potentially with array + parameters), ``data`` is an array of Monte Carlo samples (of + compatible shape), and ``axis`` is the axis of ``data`` along which + the statistic must be computed. + n_mc_samples : int, default: 9999 + The number of Monte Carlo samples drawn from the null hypothesized + distribution to form the null distribution of the statistic. The + sample size of each is the same as the given `data`. + rng : `numpy.random.Generator`, optional + Pseudorandom number generator state. When `rng` is None, a new + `numpy.random.Generator` is created using entropy from the + operating system. Types other than `numpy.random.Generator` are + passed to `numpy.random.default_rng` to instantiate a ``Generator``. + + Returns + ------- + res : GoodnessOfFitResult + An object with the following attributes. + + fit_result : `~scipy.stats._result_classes.FitResult` + An object representing the fit of the provided `dist` to `data`. + This object includes the values of distribution family parameters + that fully define the null-hypothesized distribution, that is, + the distribution from which Monte Carlo samples are drawn. + statistic : float + The value of the statistic comparing provided `data` to the + null-hypothesized distribution. + pvalue : float + The proportion of elements in the null distribution with + statistic values at least as extreme as the statistic value of the + provided `data`. + null_distribution : ndarray + The value of the statistic for each Monte Carlo sample + drawn from the null-hypothesized distribution. + + Notes + ----- + This is a generalized Monte Carlo goodness-of-fit procedure, special cases + of which correspond with various Anderson-Darling tests, Lilliefors' test, + etc. The test is described in [2]_, [3]_, and [4]_ as a parametric + bootstrap test. This is a Monte Carlo test in which parameters that + specify the distribution from which samples are drawn have been estimated + from the data. We describe the test using "Monte Carlo" rather than + "parametric bootstrap" throughout to avoid confusion with the more familiar + nonparametric bootstrap, and describe how the test is performed below. + + *Traditional goodness of fit tests* + + Traditionally, critical values corresponding with a fixed set of + significance levels are pre-calculated using Monte Carlo methods. Users + perform the test by calculating the value of the test statistic only for + their observed `data` and comparing this value to tabulated critical + values. This practice is not very flexible, as tables are not available for + all distributions and combinations of known and unknown parameter values. + Also, results can be inaccurate when critical values are interpolated from + limited tabulated data to correspond with the user's sample size and + fitted parameter values. To overcome these shortcomings, this function + allows the user to perform the Monte Carlo trials adapted to their + particular data. + + *Algorithmic overview* + + In brief, this routine executes the following steps: + + 1. Fit unknown parameters to the given `data`, thereby forming the + "null-hypothesized" distribution, and compute the statistic of + this pair of data and distribution. + 2. Draw random samples from this null-hypothesized distribution. + 3. Fit the unknown parameters to each random sample. + 4. Calculate the statistic between each sample and the distribution that + has been fit to the sample. + 5. Compare the value of the statistic corresponding with `data` from (1) + against the values of the statistic corresponding with the random + samples from (4). The p-value is the proportion of samples with a + statistic value greater than or equal to the statistic of the observed + data. + + In more detail, the steps are as follows. + + First, any unknown parameters of the distribution family specified by + `dist` are fit to the provided `data` using maximum likelihood estimation. + (One exception is the normal distribution with unknown location and scale: + we use the bias-corrected standard deviation ``np.std(data, ddof=1)`` for + the scale as recommended in [1]_.) + These values of the parameters specify a particular member of the + distribution family referred to as the "null-hypothesized distribution", + that is, the distribution from which the data were sampled under the null + hypothesis. The `statistic`, which compares data to a distribution, is + computed between `data` and the null-hypothesized distribution. + + Next, many (specifically `n_mc_samples`) new samples, each containing the + same number of observations as `data`, are drawn from the + null-hypothesized distribution. All unknown parameters of the distribution + family `dist` are fit to *each resample*, and the `statistic` is computed + between each sample and its corresponding fitted distribution. These + values of the statistic form the Monte Carlo null distribution (not to be + confused with the "null-hypothesized distribution" above). + + The p-value of the test is the proportion of statistic values in the Monte + Carlo null distribution that are at least as extreme as the statistic value + of the provided `data`. More precisely, the p-value is given by + + .. math:: + + p = \frac{b + 1} + {m + 1} + + where :math:`b` is the number of statistic values in the Monte Carlo null + distribution that are greater than or equal to the statistic value + calculated for `data`, and :math:`m` is the number of elements in the + Monte Carlo null distribution (`n_mc_samples`). The addition of :math:`1` + to the numerator and denominator can be thought of as including the + value of the statistic corresponding with `data` in the null distribution, + but a more formal explanation is given in [5]_. + + *Limitations* + + The test can be very slow for some distribution families because unknown + parameters of the distribution family must be fit to each of the Monte + Carlo samples, and for most distributions in SciPy, distribution fitting + performed via numerical optimization. + + *Anti-Pattern* + + For this reason, it may be tempting + to treat parameters of the distribution pre-fit to `data` (by the user) + as though they were `known_params`, as specification of all parameters of + the distribution precludes the need to fit the distribution to each Monte + Carlo sample. (This is essentially how the original Kilmogorov-Smirnov + test is performed.) Although such a test can provide evidence against the + null hypothesis, the test is conservative in the sense that small p-values + will tend to (greatly) *overestimate* the probability of making a type I + error (that is, rejecting the null hypothesis although it is true), and the + power of the test is low (that is, it is less likely to reject the null + hypothesis even when the null hypothesis is false). + This is because the Monte Carlo samples are less likely to agree with the + null-hypothesized distribution as well as `data`. This tends to increase + the values of the statistic recorded in the null distribution, so that a + larger number of them exceed the value of statistic for `data`, thereby + inflating the p-value. + + References + ---------- + .. [1] M. A. Stephens (1974). "EDF Statistics for Goodness of Fit and + Some Comparisons." Journal of the American Statistical Association, + Vol. 69, pp. 730-737. + .. [2] W. Stute, W. G. Manteiga, and M. P. Quindimil (1993). + "Bootstrap based goodness-of-fit-tests." Metrika 40.1: 243-256. + .. [3] C. Genest, & B Rémillard. (2008). "Validity of the parametric + bootstrap for goodness-of-fit testing in semiparametric models." + Annales de l'IHP Probabilités et statistiques. Vol. 44. No. 6. + .. [4] I. Kojadinovic and J. Yan (2012). "Goodness-of-fit testing based on + a weighted bootstrap: A fast large-sample alternative to the + parametric bootstrap." Canadian Journal of Statistics 40.3: 480-500. + .. [5] B. Phipson and G. K. Smyth (2010). "Permutation P-values Should + Never Be Zero: Calculating Exact P-values When Permutations Are + Randomly Drawn." Statistical Applications in Genetics and Molecular + Biology 9.1. + .. [6] H. W. Lilliefors (1967). "On the Kolmogorov-Smirnov test for + normality with mean and variance unknown." Journal of the American + statistical Association 62.318: 399-402. + .. [7] Filliben, James J. "The probability plot correlation coefficient + test for normality." Technometrics 17.1 (1975): 111-117. + + Examples + -------- + A well-known test of the null hypothesis that data were drawn from a + given distribution is the Kolmogorov-Smirnov (KS) test, available in SciPy + as `scipy.stats.ks_1samp`. Suppose we wish to test whether the following + data: + + >>> import numpy as np + >>> from scipy import stats + >>> rng = np.random.default_rng(1638083107694713882823079058616272161) + >>> x = stats.uniform.rvs(size=75, random_state=rng) + + were sampled from a normal distribution. To perform a KS test, the + empirical distribution function of the observed data will be compared + against the (theoretical) cumulative distribution function of a normal + distribution. Of course, to do this, the normal distribution under the null + hypothesis must be fully specified. This is commonly done by first fitting + the ``loc`` and ``scale`` parameters of the distribution to the observed + data, then performing the test. + + >>> loc, scale = np.mean(x), np.std(x, ddof=1) + >>> cdf = stats.norm(loc, scale).cdf + >>> stats.ks_1samp(x, cdf) + KstestResult(statistic=0.1119257570456813, + pvalue=0.2827756409939257, + statistic_location=0.7751845155861765, + statistic_sign=-1) + + An advantage of the KS-test is that the p-value - the probability of + obtaining a value of the test statistic under the null hypothesis as + extreme as the value obtained from the observed data - can be calculated + exactly and efficiently. `goodness_of_fit` can only approximate these + results. + + >>> known_params = {'loc': loc, 'scale': scale} + >>> res = stats.goodness_of_fit(stats.norm, x, known_params=known_params, + ... statistic='ks', rng=rng) + >>> res.statistic, res.pvalue + (0.1119257570456813, 0.2788) + + The statistic matches exactly, but the p-value is estimated by forming + a "Monte Carlo null distribution", that is, by explicitly drawing random + samples from `scipy.stats.norm` with the provided parameters and + calculating the stastic for each. The fraction of these statistic values + at least as extreme as ``res.statistic`` approximates the exact p-value + calculated by `scipy.stats.ks_1samp`. + + However, in many cases, we would prefer to test only that the data were + sampled from one of *any* member of the normal distribution family, not + specifically from the normal distribution with the location and scale + fitted to the observed sample. In this case, Lilliefors [6]_ argued that + the KS test is far too conservative (that is, the p-value overstates + the actual probability of rejecting a true null hypothesis) and thus lacks + power - the ability to reject the null hypothesis when the null hypothesis + is actually false. + Indeed, our p-value above is approximately 0.28, which is far too large + to reject the null hypothesis at any common significance level. + + Consider why this might be. Note that in the KS test above, the statistic + always compares data against the CDF of a normal distribution fitted to the + *observed data*. This tends to reduce the value of the statistic for the + observed data, but it is "unfair" when computing the statistic for other + samples, such as those we randomly draw to form the Monte Carlo null + distribution. It is easy to correct for this: whenever we compute the KS + statistic of a sample, we use the CDF of a normal distribution fitted + to *that sample*. The null distribution in this case has not been + calculated exactly and is tyically approximated using Monte Carlo methods + as described above. This is where `goodness_of_fit` excels. + + >>> res = stats.goodness_of_fit(stats.norm, x, statistic='ks', rng=rng) + >>> res.statistic, res.pvalue + (0.1119257570456813, 0.0196) + + Indeed, this p-value is much smaller, and small enough to (correctly) + reject the null hypothesis at common significance levels, including 5% and + 2.5%. + + However, the KS statistic is not very sensitive to all deviations from + normality. The original advantage of the KS statistic was the ability + to compute the null distribution theoretically, but a more sensitive + statistic - resulting in a higher test power - can be used now that we can + approximate the null distribution + computationally. The Anderson-Darling statistic [1]_ tends to be more + sensitive, and critical values of this statistic have been tabulated + for various significance levels and sample sizes using Monte Carlo methods. + + >>> res = stats.anderson(x, 'norm', method='interpolate') + >>> print(res.statistic) + 1.2139573337497467 + >>> print(res.pvalue) + 0.01 + + Here, the observed value of the statistic exceeds the critical value + corresponding with a 1% significance level. This tells us that the p-value + of the observed data is 1% or less, but what is it? `goodness_of_fit` can + estimate it directly. + + >>> res = stats.goodness_of_fit(stats.norm, x, statistic='ad', rng=rng) + >>> res.statistic, res.pvalue + (1.2139573337497467, 0.0034) + + A further advantage is that use of `goodness_of_fit` is not limited to + a particular set of distributions or conditions on which parameters + are known versus which must be estimated from data. Instead, + `goodness_of_fit` can estimate p-values relatively quickly for any + distribution with a sufficiently fast and reliable ``fit`` method. For + instance, here we perform a goodness of fit test using the Cramer-von Mises + statistic against the Rayleigh distribution with known location and unknown + scale. + + >>> rng = np.random.default_rng() + >>> x = stats.chi(df=2.2, loc=0, scale=2).rvs(size=1000, random_state=rng) + >>> res = stats.goodness_of_fit(stats.rayleigh, x, statistic='cvm', + ... known_params={'loc': 0}, rng=rng) + + This executes fairly quickly, but to check the reliability of the ``fit`` + method, we should inspect the fit result. + + >>> res.fit_result # location is as specified, and scale is reasonable + params: FitParams(loc=0.0, scale=2.1026719844231243) + success: True + message: 'The fit was performed successfully.' + >>> import matplotlib.pyplot as plt # matplotlib must be installed to plot + >>> res.fit_result.plot() + >>> plt.show() + + If the distribution is not fit to the observed data as well as possible, + the test may not control the type I error rate, that is, the chance of + rejecting the null hypothesis even when it is true. + + We should also look for extreme outliers in the null distribution that + may be caused by unreliable fitting. These do not necessarily invalidate + the result, but they tend to reduce the test's power. + + >>> _, ax = plt.subplots() + >>> ax.hist(np.log10(res.null_distribution)) + >>> ax.set_xlabel("log10 of CVM statistic under the null hypothesis") + >>> ax.set_ylabel("Frequency") + >>> ax.set_title("Histogram of the Monte Carlo null distribution") + >>> plt.show() + + This plot seems reassuring. + + If ``fit`` method is working reliably, and if the distribution of the test + statistic is not particularly sensitive to the values of the fitted + parameters, then the p-value provided by `goodness_of_fit` is expected to + be a good approximation. + + >>> res.statistic, res.pvalue + (0.2231991510248692, 0.0525) + + """ + args = _gof_iv(dist, data, known_params, fit_params, guessed_params, + statistic, n_mc_samples, rng) + (dist, data, fixed_nhd_params, fixed_rfd_params, guessed_nhd_params, + guessed_rfd_params, statistic, n_mc_samples_int, rng) = args + + # Fit null hypothesis distribution to data + nhd_fit_fun = _get_fit_fun(dist, data, guessed_nhd_params, + fixed_nhd_params) + nhd_vals = nhd_fit_fun(data) + nhd_dist = dist(*nhd_vals) + + def rvs(size): + return nhd_dist.rvs(size=size, random_state=rng) + + # Define statistic + fit_fun = _get_fit_fun(dist, data, guessed_rfd_params, fixed_rfd_params) + if callable(statistic): + compare_fun = statistic + else: + compare_fun = _compare_dict[statistic] + alternative = getattr(compare_fun, 'alternative', 'greater') + + def statistic_fun(data, axis): + # Make things simple by always working along the last axis. + data = np.moveaxis(data, axis, -1) + rfd_vals = fit_fun(data) + rfd_dist = dist(*rfd_vals) + return compare_fun(rfd_dist, data, axis=-1) + + res = stats.monte_carlo_test(data, rvs, statistic_fun, vectorized=True, + n_resamples=n_mc_samples, axis=-1, + alternative=alternative) + opt_res = optimize.OptimizeResult() + opt_res.success = True + opt_res.message = "The fit was performed successfully." + opt_res.x = nhd_vals + # Only continuous distributions for now, hence discrete=False + # There's no fundamental limitation; it's just that we're not using + # stats.fit, discrete distributions don't have `fit` method, and + # we haven't written any vectorized fit functions for a discrete + # distribution yet. + return GoodnessOfFitResult(FitResult(dist, data, False, opt_res), + res.statistic, res.pvalue, + res.null_distribution) + + +def _get_fit_fun(dist, data, guessed_params, fixed_params): + + shape_names = [] if dist.shapes is None else dist.shapes.split(", ") + param_names = shape_names + ['loc', 'scale'] + fparam_names = ['f'+name for name in param_names] + all_fixed = not set(fparam_names).difference(fixed_params) + guessed_shapes = [guessed_params.pop(x, None) + for x in shape_names if x in guessed_params] + + if all_fixed: + def fit_fun(data): + return [fixed_params[name] for name in fparam_names] + # Define statistic, including fitting distribution to data + elif dist in _fit_funs: + def fit_fun(data): + params = _fit_funs[dist](data, **fixed_params) + params = np.asarray(np.broadcast_arrays(*params)) + if params.ndim > 1: + params = params[..., np.newaxis] + return params + else: + def fit_fun_1d(data): + return dist.fit(data, *guessed_shapes, **guessed_params, + **fixed_params) + + def fit_fun(data): + params = np.apply_along_axis(fit_fun_1d, axis=-1, arr=data) + if params.ndim > 1: + params = params.T[..., np.newaxis] + return params + + return fit_fun + + +# Vectorized fitting functions. These are to accept ND `data` in which each +# row (slice along last axis) is a sample to fit and scalar fixed parameters. +# They return a tuple of shape parameter arrays, each of shape data.shape[:-1]. +def _fit_norm(data, floc=None, fscale=None): + loc = floc + scale = fscale + if loc is None and scale is None: + loc = np.mean(data, axis=-1) + scale = np.std(data, ddof=1, axis=-1) + elif loc is None: + loc = np.mean(data, axis=-1) + elif scale is None: + scale = np.sqrt(((data - loc)**2).mean(axis=-1)) + return loc, scale + + +_fit_funs = {stats.norm: _fit_norm} # type: ignore[attr-defined] + + +# Vectorized goodness of fit statistic functions. These accept a frozen +# distribution object and `data` in which each row (slice along last axis) is +# a sample. + + +def _anderson_darling(dist, data, axis): + x = np.sort(data, axis=-1) + n = data.shape[-1] + i = np.arange(1, n+1) + Si = (2*i - 1)/n * (dist.logcdf(x) + dist.logsf(x[..., ::-1])) + S = np.sum(Si, axis=-1) + return -n - S + + +def _compute_dplus(cdfvals): # adapted from _stats_py before gh-17062 + n = cdfvals.shape[-1] + return (np.arange(1.0, n + 1) / n - cdfvals).max(axis=-1) + + +def _compute_dminus(cdfvals): + n = cdfvals.shape[-1] + return (cdfvals - np.arange(0.0, n)/n).max(axis=-1) + + +def _kolmogorov_smirnov(dist, data, axis=-1): + x = np.sort(data, axis=axis) + cdfvals = dist.cdf(x) + cdfvals = np.moveaxis(cdfvals, axis, -1) + Dplus = _compute_dplus(cdfvals) # always works along last axis + Dminus = _compute_dminus(cdfvals) + return np.maximum(Dplus, Dminus) + + +def _corr(X, M): + # Correlation coefficient r, simplified and vectorized as we need it. + # See [7] Equation (2). Lemma 1/2 are only for distributions symmetric + # about 0. + Xm = X.mean(axis=-1, keepdims=True) + Mm = M.mean(axis=-1, keepdims=True) + num = np.sum((X - Xm) * (M - Mm), axis=-1) + den = np.sqrt(np.sum((X - Xm)**2, axis=-1) * np.sum((M - Mm)**2, axis=-1)) + return num/den + + +def _filliben(dist, data, axis): + # [7] Section 8 # 1 + X = np.sort(data, axis=-1) + + # [7] Section 8 # 2 + n = data.shape[-1] + k = np.arange(1, n+1) + # Filliben used an approximation for the uniform distribution order + # statistic medians. + # m = (k - .3175)/(n + 0.365) + # m[-1] = 0.5**(1/n) + # m[0] = 1 - m[-1] + # We can just as easily use the (theoretically) exact values. See e.g. + # https://en.wikipedia.org/wiki/Order_statistic + # "Order statistics sampled from a uniform distribution" + m = stats.beta(k, n + 1 - k).median() + + # [7] Section 8 # 3 + M = dist.ppf(m) + + # [7] Section 8 # 4 + return _corr(X, M) +_filliben.alternative = 'less' # type: ignore[attr-defined] + + +def _cramer_von_mises(dist, data, axis): + x = np.sort(data, axis=-1) + n = data.shape[-1] + cdfvals = dist.cdf(x) + u = (2*np.arange(1, n+1) - 1)/(2*n) + w = 1 / (12*n) + np.sum((u - cdfvals)**2, axis=-1) + return w + + +_compare_dict = {"ad": _anderson_darling, "ks": _kolmogorov_smirnov, + "cvm": _cramer_von_mises, "filliben": _filliben} + + +def _gof_iv(dist, data, known_params, fit_params, guessed_params, statistic, + n_mc_samples, rng): + + if not isinstance(dist, stats.rv_continuous): + message = ("`dist` must be a (non-frozen) instance of " + "`stats.rv_continuous`.") + raise TypeError(message) + + data = np.asarray(data, dtype=float) + if not data.ndim == 1: + message = "`data` must be a one-dimensional array of numbers." + raise ValueError(message) + + # Leave validation of these key/value pairs to the `fit` method, + # but collect these into dictionaries that will be used + known_params = known_params or dict() + fit_params = fit_params or dict() + guessed_params = guessed_params or dict() + + known_params_f = {("f"+key): val for key, val in known_params.items()} + fit_params_f = {("f"+key): val for key, val in fit_params.items()} + + # These are the values of parameters of the null distribution family + # with which resamples are drawn + fixed_nhd_params = known_params_f.copy() + fixed_nhd_params.update(fit_params_f) + + # These are fixed when fitting the distribution family to resamples + fixed_rfd_params = known_params_f.copy() + + # These are used as guesses when fitting the distribution family to + # the original data + guessed_nhd_params = guessed_params.copy() + + # These are used as guesses when fitting the distribution family to + # resamples + guessed_rfd_params = fit_params.copy() + guessed_rfd_params.update(guessed_params) + + if not callable(statistic): + statistic = statistic.lower() + statistics = {'ad', 'ks', 'cvm', 'filliben'} + if statistic not in statistics: + message = f"`statistic` must be one of {statistics}." + raise ValueError(message) + + n_mc_samples_int = int(n_mc_samples) + if n_mc_samples_int != n_mc_samples: + message = "`n_mc_samples` must be an integer." + raise TypeError(message) + + rng = check_random_state(rng) + + return (dist, data, fixed_nhd_params, fixed_rfd_params, guessed_nhd_params, + guessed_rfd_params, statistic, n_mc_samples_int, rng) diff --git a/v12x/pylib/scipy/stats/_hypotests.py b/v12x/pylib/scipy/stats/_hypotests.py new file mode 100644 index 0000000000000000000000000000000000000000..b905654a99a08b2a856923279aa3a112363cb6cf --- /dev/null +++ b/v12x/pylib/scipy/stats/_hypotests.py @@ -0,0 +1,2201 @@ +from collections import namedtuple +from dataclasses import dataclass +import math +import numpy as np +import warnings +from itertools import combinations +import scipy.stats +from scipy.optimize import shgo +from . import distributions +from ._common import ConfidenceInterval +from ._continuous_distns import norm +from scipy._lib._array_api import (xp_capabilities, array_namespace, xp_size, + xp_promote, xp_result_type, xp_copy, is_numpy, + is_lazy_array, _count_nonmasked, is_marray, + _masked_apply) +import scipy._external.array_api_extra as xpx +from scipy.special import gamma, kv, gammaln +from scipy.fft import ifft +from ._stats_pythran import _a_ij_Aij_Dij2 +from ._stats_pythran import ( + _concordant_pairs as _P, _discordant_pairs as _Q +) +from ._axis_nan_policy import _axis_nan_policy_factory +from scipy.stats import _stats_py + +__all__ = ['epps_singleton_2samp', 'cramervonmises', 'somersd', + 'barnard_exact', 'boschloo_exact', 'cramervonmises_2samp', + 'tukey_hsd', 'poisson_means_test'] + +Epps_Singleton_2sampResult = namedtuple('Epps_Singleton_2sampResult', + ('statistic', 'pvalue')) + + +@xp_capabilities(skip_backends=[("dask.array", "lazy -> no _axis_nan_policy")]) +@_axis_nan_policy_factory(Epps_Singleton_2sampResult, n_samples=2, too_small=4) +def epps_singleton_2samp(x, y, t=(0.4, 0.8), *, axis=0): + """Compute the Epps-Singleton (ES) test statistic. + + Test the null hypothesis that two samples have the same underlying + probability distribution. + + Parameters + ---------- + x, y : array-like + The two samples of observations to be tested. Input must not have more + than one dimension. Samples can have different lengths, but both + must have at least five observations. + t : array-like, optional + The points (t1, ..., tn) where the empirical characteristic function is + to be evaluated. It should be positive distinct numbers. The default + value (0.4, 0.8) is proposed in [1]_. Input must not have more than + one dimension. + axis : int or tuple of ints, default: 0 + If an int or tuple of ints, the axis or axes of the input along which + to compute the statistic. The statistic of each axis-slice (e.g. row) + of the input will appear in a corresponding element of the output. + If ``None``, the input will be raveled before computing the statistic. + + Returns + ------- + statistic : float + The test statistic. + pvalue : float + The associated p-value based on the asymptotic chi2-distribution. + + See Also + -------- + ks_2samp, anderson_ksamp + + Notes + ----- + Testing whether two samples are generated by the same underlying + distribution is a classical question in statistics. A widely used test is + the Kolmogorov-Smirnov (KS) test which relies on the empirical + distribution function. Epps and Singleton introduce a test based on the + empirical characteristic function in [1]_. + + One advantage of the ES test compared to the KS test is that is does + not assume a continuous distribution. In [1]_, the authors conclude + that the test also has a higher power than the KS test in many + examples. They recommend the use of the ES test for discrete samples as + well as continuous samples with at least 25 observations each, whereas + `anderson_ksamp` is recommended for smaller sample sizes in the + continuous case. + + The p-value is computed from the asymptotic distribution of the test + statistic which follows a `chi2` distribution. If the sample size of both + `x` and `y` is below 25, the small sample correction proposed in [1]_ is + applied to the test statistic. + + The default values of `t` are determined in [1]_ by considering + various distributions and finding good values that lead to a high power + of the test in general. Table III in [1]_ gives the optimal values for + the distributions tested in that study. The values of `t` are scaled by + the semi-interquartile range in the implementation, see [1]_. + + References + ---------- + .. [1] T. W. Epps and K. J. Singleton, "An omnibus test for the two-sample + problem using the empirical characteristic function", Journal of + Statistical Computation and Simulation 26, p. 177--203, 1986. + + .. [2] S. J. Goerg and J. Kaiser, "Nonparametric testing of distributions + - the Epps-Singleton two-sample test using the empirical characteristic + function", The Stata Journal 9(3), p. 454--465, 2009. + + Examples + -------- + Generate random samples ``x`` and ``y``. + + >>> import numpy as np + >>> from scipy import stats + >>> rng = np.random.default_rng(66326777) + >>> x = rng.normal(0, 1, size=100) + >>> y = rng.normal(5, 3, size=100) + + Test the null hypothesis that ``x`` and ``y`` are sampled from the same + distribution. + + >>> res = stats.epps_singleton_2samp(x, y) + >>> res.statistic + np.float64(316.6441136688085) + >>> res.pvalue + np.float64(2.778946959815902e-67) + + The large value of the statistic and small p-value may be taken as evidence that + ``x`` and ``y`` were not sampled from the same distribution. + """ + xp = array_namespace(x, y) + # x and y are converted to arrays by the decorator + # and `axis` is guaranteed to be -1. + x, y = xp_promote(x, y, force_floating=True, xp=xp) + t = xp.asarray(t, dtype=x.dtype) + # check if x and y are valid inputs + nx, ny = x.shape[-1], y.shape[-1] + if (nx < 5) or (ny < 5): # only used by test_axis_nan_policy + raise ValueError('x and y should have at least 5 elements, but len(x) ' + f'= {nx} and len(y) = {ny}.') + n = nx + ny + + # check if t is valid + if t.ndim > 1: + raise ValueError(f't must be 1d, but t.ndim equals {t.ndim}.') + if is_lazy_array(t): + t = xpx.at(t)[t <= 0].set(xp.nan) + elif xp.any(t <= 0): + raise ValueError('t must contain positive elements only.') + + # Previously, non-finite input caused an error in linalg functions. + # To prevent an issue in one slice from halting the calculation, replace non-finite + # values with a harmless one, and replace results with NaN at the end. + i_x = ~xp.isfinite(x) + i_y = ~xp.isfinite(y) + # Ideally we would avoid copying all data here; see + # discussion in data-apis/array-api-extra#506. + x = xp.where(i_x, 1., x) + y = xp.where(i_y, 1., y) + invalid_result = xp.any(i_x, axis=-1) | xp.any(i_y, axis=-1) + + # rescale t with semi-iqr as proposed in [1]; import iqr here to avoid + # circular import + from scipy.stats import iqr + sigma = iqr(xp.concat((x, y), axis=-1), axis=-1, keepdims=True) / 2 + ts = xp.reshape(t, (-1,) + (1,)*x.ndim) / sigma + + # covariance estimation of ES test + gx = xp.concat((xp.cos(ts*x), xp.sin(ts*x)), axis=0) + gy = xp.concat((xp.cos(ts*y), xp.sin(ts*y)), axis=0) + gx, gy = xp.moveaxis(gx, 0, -2), xp.moveaxis(gy, 0, -2) + cov_x = xpx.cov(gx) * (nx-1)/nx # the test uses biased cov-estimate + cov_y = xpx.cov(gy) * (ny-1)/ny + cov_x, cov_y = xp.astype(cov_x, x.dtype), xp.astype(cov_y, y.dtype) + est_cov = (n/nx)*cov_x + (n/ny)*cov_y + est_cov_inv = xp.linalg.pinv(est_cov) + r = xp.asarray(xp.linalg.matrix_rank(est_cov_inv), dtype=est_cov_inv.dtype) + if not is_lazy_array(r) and xp.any(r < 2*xp_size(t)): + warnings.warn('Estimated covariance matrix does not have full rank. ' + 'This indicates a bad choice of the input t and the ' + 'test might not be consistent.', # see p. 183 in [1]_ + stacklevel=2) + + # compute test statistic w distributed asympt. as chisquare with df=r + g_diff = xp.mean(gx, axis=-1, keepdims=True) - xp.mean(gy, axis=-1, keepdims=True) + w = n*xp.matmul(xp.matrix_transpose(g_diff), xp.matmul(est_cov_inv, g_diff)) + w = w[..., 0, 0] + + # apply small-sample correction + if (max(nx, ny) < 25): + corr = 1.0/(1.0 + n**(-0.45) + 10.1*(nx**(-1.7) + ny**(-1.7))) + w *= corr + + chi2 = _stats_py._SimpleChi2(r) + p = _stats_py._get_pvalue(w, chi2, alternative='greater', symmetric=False, xp=xp) + + w = xpx.at(w)[invalid_result].set(xp.nan) + p = xpx.at(p)[invalid_result].set(xp.nan) + w = w[()] if w.ndim == 0 else w + p = p[()] if p.ndim == 0 else p + return Epps_Singleton_2sampResult(w, p) + + +@xp_capabilities(np_only=True) +def poisson_means_test(k1, n1, k2, n2, *, diff=0, alternative='two-sided'): + r""" + Performs the Poisson means test, AKA the "E-test". + + This is a test of the null hypothesis that the difference between means of + two Poisson distributions is `diff`. The samples are provided as the + number of events `k1` and `k2` observed within measurement intervals + (e.g. of time, space, number of observations) of sizes `n1` and `n2`. + + Parameters + ---------- + k1 : int + Number of events observed from distribution 1. + n1 : float + Size of sample from distribution 1. + k2 : int + Number of events observed from distribution 2. + n2 : float + Size of sample from distribution 2. + diff : float, default=0 + The hypothesized difference in means between the distributions + underlying the samples. + alternative : {'two-sided', 'less', 'greater'}, optional + Defines the alternative hypothesis. + The following options are available (default is 'two-sided'): + + * 'two-sided': the difference between distribution means is not + equal to `diff` + * 'less': the difference between distribution means is less than + `diff` + * 'greater': the difference between distribution means is greater + than `diff` + + Returns + ------- + statistic : float + The test statistic (see [1]_ equation 3.3). + pvalue : float + The probability of achieving such an extreme value of the test + statistic under the null hypothesis. + + Notes + ----- + + Let: + + .. math:: X_1 \sim \mbox{Poisson}(\mathtt{n1}\lambda_1) + + be a random variable independent of + + .. math:: X_2 \sim \mbox{Poisson}(\mathtt{n2}\lambda_2) + + and let ``k1`` and ``k2`` be the observed values of :math:`X_1` + and :math:`X_2`, respectively. Then `poisson_means_test` uses the number + of observed events ``k1`` and ``k2`` from samples of size ``n1`` and + ``n2``, respectively, to test the null hypothesis that + + .. math:: + H_0: \lambda_1 - \lambda_2 = \mathtt{diff} + + A benefit of the E-test is that it has good power for small sample sizes, + which can reduce sampling costs [1]_. It has been evaluated and determined + to be more powerful than the comparable C-test, sometimes referred to as + the Poisson exact test. + + References + ---------- + .. [1] Krishnamoorthy, K., & Thomson, J. (2004). A more powerful test for + comparing two Poisson means. Journal of Statistical Planning and + Inference, 119(1), 23-35. + + .. [2] Przyborowski, J., & Wilenski, H. (1940). Homogeneity of results in + testing samples from Poisson series: With an application to testing + clover seed for dodder. Biometrika, 31(3/4), 313-323. + + Examples + -------- + + Suppose that a gardener wishes to test the number of dodder (weed) seeds + in a sack of clover seeds that they buy from a seed company. It has + previously been established that the number of dodder seeds in clover + follows the Poisson distribution. + + A 100 gram sample is drawn from the sack before being shipped to the + gardener. The sample is analyzed, and it is found to contain no dodder + seeds; that is, `k1` is 0. However, upon arrival, the gardener draws + another 100 gram sample from the sack. This time, three dodder seeds are + found in the sample; that is, `k2` is 3. The gardener would like to + know if the difference is significant and not due to chance. The + null hypothesis is that the difference between the two samples is merely + due to chance, or that :math:`\lambda_1 - \lambda_2 = \mathtt{diff}` + where :math:`\mathtt{diff} = 0`. The alternative hypothesis is that the + difference is not due to chance, or :math:`\lambda_1 - \lambda_2 \ne 0`. + The gardener selects a significance level of 5% to reject the null + hypothesis in favor of the alternative [2]_. + + >>> import scipy.stats as stats + >>> res = stats.poisson_means_test(0, 100, 3, 100) + >>> res.statistic, res.pvalue + (-1.7320508075688772, 0.08837900929018157) + + The p-value is .088, indicating a near 9% chance of observing a value of + the test statistic under the null hypothesis. This exceeds 5%, so the + gardener does not reject the null hypothesis as the difference cannot be + regarded as significant at this level. + """ + + _poisson_means_test_iv(k1, n1, k2, n2, diff, alternative) + + # "for a given k_1 and k_2, an estimate of \lambda_2 is given by" [1] (3.4) + lmbd_hat2 = ((k1 + k2) / (n1 + n2) - diff * n1 / (n1 + n2)) + + # "\hat{\lambda_{2k}} may be less than or equal to zero ... and in this + # case the null hypothesis cannot be rejected ... [and] it is not necessary + # to compute the p-value". [1] page 26 below eq. (3.6). + if lmbd_hat2 <= 0: + return _stats_py.SignificanceResult(0, 1) + + # The unbiased variance estimate [1] (3.2) + var = k1 / (n1 ** 2) + k2 / (n2 ** 2) + + # The _observed_ pivot statistic from the input. It follows the + # unnumbered equation following equation (3.3) This is used later in + # comparison with the computed pivot statistics in an indicator function. + t_k1k2 = (k1 / n1 - k2 / n2 - diff) / np.sqrt(var) + + # Equation (3.5) of [1] is lengthy, so it is broken into several parts, + # beginning here. Note that the probability mass function of poisson is + # exp^(-\mu)*\mu^k/k!, so and this is called with shape \mu, here noted + # here as nlmbd_hat*. The strategy for evaluating the double summation in + # (3.5) is to create two arrays of the values of the two products inside + # the summation and then broadcast them together into a matrix, and then + # sum across the entire matrix. + + # Compute constants (as seen in the first and second separated products in + # (3.5).). (This is the shape (\mu) parameter of the poisson distribution.) + nlmbd_hat1 = n1 * (lmbd_hat2 + diff) + nlmbd_hat2 = n2 * lmbd_hat2 + + # Determine summation bounds for tail ends of distribution rather than + # summing to infinity. `x1*` is for the outer sum and `x2*` is the inner + # sum. + x1_lb, x1_ub = distributions.poisson.ppf([1e-10, 1 - 1e-16], nlmbd_hat1) + x2_lb, x2_ub = distributions.poisson.ppf([1e-10, 1 - 1e-16], nlmbd_hat2) + + # Construct arrays to function as the x_1 and x_2 counters on the summation + # in (3.5). `x1` is in columns and `x2` is in rows to allow for + # broadcasting. + x1 = np.arange(x1_lb, x1_ub + 1) + x2 = np.arange(x2_lb, x2_ub + 1)[:, None] + + # These are the two products in equation (3.5) with `prob_x1` being the + # first (left side) and `prob_x2` being the second (right side). (To + # make as clear as possible: the 1st contains a "+ d" term, the 2nd does + # not.) + prob_x1 = distributions.poisson.pmf(x1, nlmbd_hat1) + prob_x2 = distributions.poisson.pmf(x2, nlmbd_hat2) + + # compute constants for use in the "pivot statistic" per the + # unnumbered equation following (3.3). + lmbd_x1 = x1 / n1 + lmbd_x2 = x2 / n2 + lmbds_diff = lmbd_x1 - lmbd_x2 - diff + var_x1x2 = lmbd_x1 / n1 + lmbd_x2 / n2 + + # This is the 'pivot statistic' for use in the indicator of the summation + # (left side of "I[.]"). + with np.errstate(invalid='ignore', divide='ignore'): + t_x1x2 = lmbds_diff / np.sqrt(var_x1x2) + + # `[indicator]` implements the "I[.] ... the indicator function" per + # the paragraph following equation (3.5). + if alternative == 'two-sided': + indicator = np.abs(t_x1x2) >= np.abs(t_k1k2) + elif alternative == 'less': + indicator = t_x1x2 <= t_k1k2 + else: + indicator = t_x1x2 >= t_k1k2 + + # Multiply all combinations of the products together, exclude terms + # based on the `indicator` and then sum. (3.5) + pvalue = np.sum((prob_x1 * prob_x2)[indicator]) + return _stats_py.SignificanceResult(t_k1k2, pvalue) + + +def _poisson_means_test_iv(k1, n1, k2, n2, diff, alternative): + # """check for valid types and values of input to `poisson_mean_test`.""" + if k1 != int(k1) or k2 != int(k2): + raise TypeError('`k1` and `k2` must be integers.') + + count_err = '`k1` and `k2` must be greater than or equal to 0.' + if k1 < 0 or k2 < 0: + raise ValueError(count_err) + + if n1 <= 0 or n2 <= 0: + raise ValueError('`n1` and `n2` must be greater than 0.') + + if diff < 0: + raise ValueError('diff must be greater than or equal to 0.') + + alternatives = {'two-sided', 'less', 'greater'} + if alternative.lower() not in alternatives: + raise ValueError(f"Alternative must be one of '{alternatives}'.") + + +class CramerVonMisesResult: + def __init__(self, statistic, pvalue): + self.statistic = statistic + self.pvalue = pvalue + + def __repr__(self): + return (f"{self.__class__.__name__}(statistic={self.statistic}, " + f"pvalue={self.pvalue})") + + +def _psi1_mod(x, *, xp=None): + """ + psi1 is defined in equation 1.10 in Csörgő, S. and Faraway, J. (1996). + This implements a modified version by excluding the term V(x) / 12 + (here: _cdf_cvm_inf(x) / 12) to avoid evaluating _cdf_cvm_inf(x) + twice in _cdf_cvm. + + Implementation based on MAPLE code of Julian Faraway and R code of the + function pCvM in the package goftest (v1.1.1), permission granted + by Adrian Baddeley. Main difference in the implementation: the code + here keeps adding terms of the series until the terms are small enough. + """ + xp = array_namespace(x) if xp is None else xp + + def _ed2(y): + z = y**2 / 4 + z_ = np.asarray(z) + b = xp.asarray(kv(1/4, z_) + kv(3/4, z_)) + return xp.exp(-z) * (y/2)**(3/2) * b / math.sqrt(np.pi) + + def _ed3(y): + z = y**2 / 4 + z_ = np.asarray(z) + c = xp.exp(-z) / math.sqrt(np.pi) + kv_terms = xp.asarray(2*kv(1/4, z_) + + 3*kv(3/4, z_) - kv(5/4, z_)) + return c * (y/2)**(5/2) * kv_terms + + def _Ak(k, x): + m = 2*k + 1 + sx = 2 * xp.sqrt(x) + y1 = x**(3/4) + y2 = x**(5/4) + + gamma_kp1_2 = float(gamma(k + 1 / 2)) + gamma_kp3_2 = float(gamma(k + 3 / 2)) + + e1 = m * gamma_kp1_2 * _ed2((4 * k + 3)/sx) / (9 * y1) + e2 = gamma_kp1_2 * _ed3((4 * k + 1) / sx) / (72 * y2) + e3 = 2 * (m + 2) * gamma_kp3_2 * _ed3((4 * k + 5) / sx) / (12 * y2) + e4 = 7 * m * gamma_kp1_2 * _ed2((4 * k + 1) / sx) / (144 * y1) + e5 = 7 * m * gamma_kp1_2 * _ed2((4 * k + 5) / sx) / (144 * y1) + + return e1 + e2 + e3 + e4 + e5 + + x = xp.asarray(x) + tot = xp.zeros_like(x) + cond = xp.ones_like(x, dtype=xp.bool) + k = 0 + while xp.any(cond): + gamma_kp1 = float(gamma(k + 1)) + z = -_Ak(k, x[cond]) / (xp.pi * gamma_kp1) + tot = xpx.at(tot)[cond].set(tot[cond] + z) + # For float32 arithmetic, the tolerance may need to be adjusted or the + # algorithm may prove to be unsuitable. + cond = xpx.at(cond)[xp_copy(cond)].set(xp.abs(z) >= 1e-7) + k += 1 + + return tot + + +def _cdf_cvm_inf(x, *, xp=None): + """ + Calculate the cdf of the Cramér-von Mises statistic (infinite sample size). + + See equation 1.2 in Csörgő, S. and Faraway, J. (1996). + + Implementation based on MAPLE code of Julian Faraway and R code of the + function pCvM in the package goftest (v1.1.1), permission granted + by Adrian Baddeley. Main difference in the implementation: the code + here keeps adding terms of the series until the terms are small enough. + + The function is not expected to be accurate for large values of x, say + x > 4, when the cdf is very close to 1. + """ + xp = array_namespace(x) if xp is None else xp + x = xp.asarray(x) + + def term(x, k): + # this expression can be found in [2], second line of (1.3) + u = math.exp(gammaln(k + 0.5) - gammaln(k+1)) / (xp.pi**1.5 * xp.sqrt(x)) + y = 4*k + 1 + q = y**2 / (16*x) + b = xp.asarray(kv(0.25, np.asarray(q)), dtype=u.dtype) # not automatic? + return u * math.sqrt(y) * xp.exp(-q) * b + + tot = xp.zeros_like(x, dtype=x.dtype) + cond = xp.ones_like(x, dtype=xp.bool) + k = 0 + while xp.any(cond): + z = term(x[cond], k) + # tot[cond] = tot[cond] + z + tot = xpx.at(tot)[cond].add(z) + # cond[cond] = np.abs(z) >= 1e-7 + cond = xpx.at(cond)[xp_copy(cond)].set(xp.abs(z) >= 1e-7) # torch needs copy + k += 1 + + return tot + + +def _cdf_cvm(x, n=None, *, xp=None): + """ + Calculate the cdf of the Cramér-von Mises statistic for a finite sample + size n. If N is None, use the asymptotic cdf (n=inf). + + See equation 1.8 in Csörgő, S. and Faraway, J. (1996) for finite samples, + 1.2 for the asymptotic cdf. + + The function is not expected to be accurate for large values of x, say + x > 2, when the cdf is very close to 1 and it might return values > 1 + in that case, e.g. _cdf_cvm(2.0, 12) = 1.0000027556716846. Moreover, it + is not accurate for small values of n, especially close to the bounds of + the distribution's domain, [1/(12*n), n/3], where the value jumps to 0 + and 1, respectively. These are limitations of the approximation by Csörgő + and Faraway (1996) implemented in this function. + """ + xp = array_namespace(x) if xp is None else xp + x = xp.asarray(x) + + if n is None: + y = _cdf_cvm_inf(x, xp=xp) + else: + # support of the test statistic is [12/n, n/3], see 1.1 in [2] + y = xp.zeros_like(x, dtype=x.dtype) + sup = (1./(12*n) < x) & (x < n/3.) + # note: _psi1_mod does not include the term _cdf_cvm_inf(x) / 12 + # therefore, we need to add it here + y = xpx.at(y)[sup].set(_cdf_cvm_inf(x[sup], xp=xp) * (1 + 1./(12*n)) + + _psi1_mod(x[sup], xp=xp) / n) + y = xpx.at(y)[x >= n/3].set(1.) + + return y[()] if y.ndim == 0 else y + + +def _cvm_result_to_tuple(res, _): + return res.statistic, res.pvalue + + +@xp_capabilities(cpu_only=True, # needs special function `kv` + skip_backends=[('dask.array', 'typical dask issues')], + jax_jit=False, # array boolean indices must be concrete + marray=True) +@_axis_nan_policy_factory(CramerVonMisesResult, n_samples=1, too_small=1, + result_to_tuple=_cvm_result_to_tuple) +def cramervonmises(rvs, cdf, args=(), *, axis=0): + r"""Perform the one-sample Cramér-von Mises test for goodness of fit. + + This performs a test of the goodness of fit of a cumulative distribution + function (cdf) :math:`F` compared to the empirical distribution function + :math:`F_n` of observed random variates :math:`X_1, ..., X_n` that are + assumed to be independent and identically distributed ([1]_). + The null hypothesis is that the :math:`X_i` have cumulative distribution + :math:`F`. + + The test statistic :math:`T` is defined as in [1]_, where :math:`\omega^2` + is the Cramér-von Mises criterion and :math:`x_i` are the observed values. + + .. math:: + T = n\omega^2 = + \frac{1}{12n} + \sum_{i=1}^n \left[ \frac{2i-1}{2n} - F(x_i) \right]^2 + + Parameters + ---------- + rvs : array_like + A 1-D array of observed values of the random variables :math:`X_i`. + The sample must contain at least two observations. + cdf : str or callable + The cumulative distribution function :math:`F` to test the + observations against. If a string, it should be the name of a + distribution in `scipy.stats`. If a callable, that callable is used + to calculate the cdf: ``cdf(x, *args) -> float``. + args : tuple, optional + Distribution parameters. These are assumed to be known; see Notes. + axis : int or tuple of ints, default: 0 + If an int or tuple of ints, the axis or axes of the input along which + to compute the statistic. The statistic of each axis-slice (e.g. row) + of the input will appear in a corresponding element of the output. + If ``None``, the input will be raveled before computing the statistic. + + Returns + ------- + res : object with attributes + statistic : float + Cramér-von Mises statistic :math:`T`. + pvalue : float + The p-value. + + See Also + -------- + kstest, cramervonmises_2samp + + Notes + ----- + .. versionadded:: 1.6.0 + + The p-value relies on the approximation given by equation 1.8 in [2]_. + It is important to keep in mind that the p-value is only accurate if + one tests a simple hypothesis, i.e. the parameters of the reference + distribution are known. If the parameters are estimated from the data + (composite hypothesis), the computed p-value is not reliable. + + References + ---------- + .. [1] Cramér-von Mises criterion, Wikipedia, + https://en.wikipedia.org/wiki/Cram%C3%A9r%E2%80%93von_Mises_criterion + .. [2] Csörgő, S. and Faraway, J. (1996). The Exact and Asymptotic + Distribution of Cramér-von Mises Statistics. Journal of the + Royal Statistical Society, pp. 221-234. + + Examples + -------- + + Suppose we wish to test whether data generated by ``scipy.stats.norm.rvs`` + were, in fact, drawn from the standard normal distribution. We choose a + significance level of ``alpha=0.05``. + + >>> import numpy as np + >>> from scipy import stats + >>> rng = np.random.default_rng(165417232101553420507139617764912913465) + >>> x = stats.norm.rvs(size=500, random_state=rng) + >>> res = stats.cramervonmises(x, 'norm') + >>> res.statistic, res.pvalue + (0.1072085112565724, 0.5508482238203407) + + The p-value exceeds our chosen significance level, so we do not + reject the null hypothesis that the observed sample is drawn from the + standard normal distribution. + + Now suppose we wish to check whether the same samples shifted by 2.1 is + consistent with being drawn from a normal distribution with a mean of 2. + + >>> y = x + 2.1 + >>> res = stats.cramervonmises(y, 'norm', args=(2,)) + >>> res.statistic, res.pvalue + (0.8364446265294695, 0.00596286797008283) + + Here we have used the `args` keyword to specify the mean (``loc``) + of the normal distribution to test the data against. This is equivalent + to the following, in which we create a frozen normal distribution with + mean 2.1, then pass its ``cdf`` method as an argument. + + >>> frozen_dist = stats.norm(loc=2) + >>> res = stats.cramervonmises(y, frozen_dist.cdf) + >>> res.statistic, res.pvalue + (0.8364446265294695, 0.00596286797008283) + + In either case, we would reject the null hypothesis that the observed + sample is drawn from a normal distribution with a mean of 2 (and default + variance of 1) because the p-value is less than our chosen + significance level. + + """ + # `_axis_nan_policy` decorator ensures `axis=-1` + xp = array_namespace(rvs) + + if isinstance(cdf, str) and is_numpy(xp): + cdf = getattr(distributions, cdf).cdf + elif isinstance(cdf, str): + message = "`cdf` must be a callable if `rvs` is a non-NumPy array." + raise ValueError(message) + + n_length = rvs.shape[-1] + if n_length <= 1: # only needed for `test_axis_nan_policy.py`; not user-facing + raise ValueError('The sample must contain at least two observations.') + + rvs = xp_promote(rvs, force_floating=True, xp=xp) + vals = xp.sort(rvs, axis=-1) + cdfvals = _masked_apply(cdf, args=(vals, *args), xp=xp) + n_count = xp.asarray(_count_nonmasked(rvs, axis=-1, xp=xp), dtype=rvs.dtype) + + u = (2*xp.arange(1, n_length+1, dtype=rvs.dtype) - 1)/(2*n_count[..., xp.newaxis]) + w = 1/(12*n_count) + xp.sum((u - cdfvals)**2, axis=-1) + + # avoid small negative values that can occur due to the approximation + p = xp.clip(1. - _masked_apply(_cdf_cvm, args=(w, n_count), xp=xp), 0., None) + + return CramerVonMisesResult(statistic=w, pvalue=p) + + +def _get_wilcoxon_distr(n): + """ + Distribution of probability of the Wilcoxon ranksum statistic r_plus (sum + of ranks of positive differences). + Returns an array with the probabilities of all the possible ranks + r = 0, ..., n*(n+1)/2 + """ + c = np.ones(1, dtype=np.float64) + for k in range(1, n + 1): + prev_c = c + c = np.zeros(k * (k + 1) // 2 + 1, dtype=np.float64) + m = len(prev_c) + c[:m] = prev_c * 0.5 + c[-m:] += prev_c * 0.5 + return c + + +def _get_wilcoxon_distr2(n): + """ + Distribution of probability of the Wilcoxon ranksum statistic r_plus (sum + of ranks of positive differences). + Returns an array with the probabilities of all the possible ranks + r = 0, ..., n*(n+1)/2 + This is a slower reference function + References + ---------- + .. [1] 1. Harris T, Hardin JW. Exact Wilcoxon Signed-Rank and Wilcoxon + Mann-Whitney Ranksum Tests. The Stata Journal. 2013;13(2):337-343. + """ + ai = np.arange(1, n+1)[:, None] + t = n*(n+1)/2 + q = 2*t + j = np.arange(q) + theta = 2*np.pi/q*j + phi_sp = np.prod(np.cos(theta*ai), axis=0) + phi_s = np.exp(1j*theta*t) * phi_sp + p = np.real(ifft(phi_s)) + res = np.zeros(int(t)+1) + res[:-1:] = p[::2] + res[0] /= 2 + res[-1] = res[0] + return res + + +def _tau_b(A): + """Calculate Kendall's tau-b and p-value from contingency table.""" + # See [2] 2.2 and 4.2 + + # contingency table must be truly 2D + if A.shape[0] == 1 or A.shape[1] == 1: + return np.nan, np.nan + + NA = A.sum() + PA = _P(A) + QA = _Q(A) + Sri2 = (A.sum(axis=1)**2).sum() + Scj2 = (A.sum(axis=0)**2).sum() + denominator = (NA**2 - Sri2)*(NA**2 - Scj2) + + tau = (PA-QA)/(denominator)**0.5 + + numerator = 4*(_a_ij_Aij_Dij2(A) - (PA - QA)**2 / NA) + s02_tau_b = numerator/denominator + if s02_tau_b == 0: # Avoid divide by zero + return tau, 0 + Z = tau/s02_tau_b**0.5 + p = 2*norm.sf(abs(Z)) # 2-sided p-value + + return tau, p + + +def _somers_d(A, alternative='two-sided'): + """Calculate Somers' D and p-value from contingency table.""" + # See [3] page 1740 + + # contingency table must be truly 2D + if A.shape[0] <= 1 or A.shape[1] <= 1: + return np.nan, np.nan + + NA = A.sum() + NA2 = NA**2 + PA = _P(A) + QA = _Q(A) + Sri2 = (A.sum(axis=1)**2).sum() + + d = (PA - QA)/(NA2 - Sri2) + + S = _a_ij_Aij_Dij2(A) - (PA-QA)**2/NA + + with np.errstate(divide='ignore'): + Z = (PA - QA)/(4*(S))**0.5 + + norm = _stats_py._SimpleNormal() + p = _stats_py._get_pvalue(Z, norm, alternative, xp=np) + + return d, p + + +@dataclass +class SomersDResult: + statistic: float + pvalue: float + table: np.ndarray + + +@xp_capabilities(np_only=True) +def somersd(x, y=None, alternative='two-sided'): + r"""Calculates Somers' D, an asymmetric measure of ordinal association. + + Like Kendall's :math:`\tau`, Somers' :math:`D` is a measure of the + correspondence between two rankings. Both statistics consider the + difference between the number of concordant and discordant pairs in two + rankings :math:`X` and :math:`Y`, and both are normalized such that values + close to 1 indicate strong agreement and values close to -1 indicate + strong disagreement. They differ in how they are normalized. To show the + relationship, Somers' :math:`D` can be defined in terms of Kendall's + :math:`\tau_a`: + + .. math:: + D(Y|X) = \frac{\tau_a(X, Y)}{\tau_a(X, X)} + + Suppose the first ranking :math:`X` has :math:`r` distinct ranks and the + second ranking :math:`Y` has :math:`s` distinct ranks. These two lists of + :math:`n` rankings can also be viewed as an :math:`r \times s` contingency + table in which element :math:`i, j` is the number of rank pairs with rank + :math:`i` in ranking :math:`X` and rank :math:`j` in ranking :math:`Y`. + Accordingly, `somersd` also allows the input data to be supplied as a + single, 2D contingency table instead of as two separate, 1D rankings. + + Note that the definition of Somers' :math:`D` is asymmetric: in general, + :math:`D(Y|X) \neq D(X|Y)`. ``somersd(x, y)`` calculates Somers' + :math:`D(Y|X)`: the "row" variable :math:`X` is treated as an independent + variable, and the "column" variable :math:`Y` is dependent. For Somers' + :math:`D(X|Y)`, swap the input lists or transpose the input table. + + Parameters + ---------- + x : array_like + 1D array of rankings, treated as the (row) independent variable. + Alternatively, a 2D contingency table. + y : array_like, optional + If `x` is a 1D array of rankings, `y` is a 1D array of rankings of the + same length, treated as the (column) dependent variable. + If `x` is 2D, `y` is ignored. + alternative : {'two-sided', 'less', 'greater'}, optional + Defines the alternative hypothesis. Default is 'two-sided'. + The following options are available: + + * 'two-sided': the rank correlation is nonzero + * 'less': the rank correlation is negative (less than zero) + * 'greater': the rank correlation is positive (greater than zero) + + Returns + ------- + res : SomersDResult + A `SomersDResult` object with the following fields: + + statistic : float + The Somers' :math:`D` statistic. + pvalue : float + The p-value for a hypothesis test whose null + hypothesis is an absence of association, :math:`D=0`. + See notes for more information. + table : 2D array + The contingency table formed from rankings `x` and `y` (or the + provided contingency table, if `x` is a 2D array) + + See Also + -------- + kendalltau : Calculates Kendall's tau, another correlation measure. + weightedtau : Computes a weighted version of Kendall's tau. + spearmanr : Calculates a Spearman rank-order correlation coefficient. + pearsonr : Calculates a Pearson correlation coefficient. + + Notes + ----- + This function follows the contingency table approach of [2]_ and + [3]_. *p*-values are computed based on an asymptotic approximation of + the test statistic distribution under the null hypothesis :math:`D=0`. + + Theoretically, hypothesis tests based on Kendall's :math:`tau` and Somers' + :math:`D` should be identical. + However, the *p*-values returned by `kendalltau` are based + on the null hypothesis of *independence* between :math:`X` and :math:`Y` + (i.e. the population from which pairs in :math:`X` and :math:`Y` are + sampled contains equal numbers of all possible pairs), which is more + specific than the null hypothesis :math:`D=0` used here. If the null + hypothesis of independence is desired, it is acceptable to use the + *p*-value returned by `kendalltau` with the statistic returned by + `somersd` and vice versa. For more information, see [2]_. + + Contingency tables are formatted according to the convention used by + SAS and R: the first ranking supplied (``x``) is the "row" variable, and + the second ranking supplied (``y``) is the "column" variable. This is + opposite the convention of Somers' original paper [1]_. + + References + ---------- + .. [1] Robert H. Somers, "A New Asymmetric Measure of Association for + Ordinal Variables", *American Sociological Review*, Vol. 27, No. 6, + pp. 799--811, 1962. + + .. [2] Morton B. Brown and Jacqueline K. Benedetti, "Sampling Behavior of + Tests for Correlation in Two-Way Contingency Tables", *Journal of + the American Statistical Association* Vol. 72, No. 358, pp. + 309--315, 1977. + + .. [3] SAS Institute, Inc., "The FREQ Procedure (Book Excerpt)", + *SAS/STAT 9.2 User's Guide, Second Edition*, SAS Publishing, 2009. + + .. [4] Laerd Statistics, "Somers' d using SPSS Statistics", *SPSS + Statistics Tutorials and Statistical Guides*, + https://statistics.laerd.com/spss-tutorials/somers-d-using-spss-statistics.php, + Accessed July 31, 2020. + + Examples + -------- + We calculate Somers' D for the example given in [4]_, in which a hotel + chain owner seeks to determine the association between hotel room + cleanliness and customer satisfaction. The independent variable, hotel + room cleanliness, is ranked on an ordinal scale: "below average (1)", + "average (2)", or "above average (3)". The dependent variable, customer + satisfaction, is ranked on a second scale: "very dissatisfied (1)", + "moderately dissatisfied (2)", "neither dissatisfied nor satisfied (3)", + "moderately satisfied (4)", or "very satisfied (5)". 189 customers + respond to the survey, and the results are cast into a contingency table + with the hotel room cleanliness as the "row" variable and customer + satisfaction as the "column" variable. + + +-----+-----+-----+-----+-----+-----+ + | | (1) | (2) | (3) | (4) | (5) | + +=====+=====+=====+=====+=====+=====+ + | (1) | 27 | 25 | 14 | 7 | 0 | + +-----+-----+-----+-----+-----+-----+ + | (2) | 7 | 14 | 18 | 35 | 12 | + +-----+-----+-----+-----+-----+-----+ + | (3) | 1 | 3 | 2 | 7 | 17 | + +-----+-----+-----+-----+-----+-----+ + + For example, 27 customers assigned their room a cleanliness ranking of + "below average (1)" and a corresponding satisfaction of "very + dissatisfied (1)". We perform the analysis as follows. + + >>> from scipy.stats import somersd + >>> table = [[27, 25, 14, 7, 0], [7, 14, 18, 35, 12], [1, 3, 2, 7, 17]] + >>> res = somersd(table) + >>> res.statistic + 0.6032766111513396 + >>> res.pvalue + 1.0007091191074533e-27 + + The value of the Somers' D statistic is approximately 0.6, indicating + a positive correlation between room cleanliness and customer satisfaction + in the sample. + The *p*-value is very small, indicating a very small probability of + observing such an extreme value of the statistic under the null + hypothesis that the statistic of the entire population (from which + our sample of 189 customers is drawn) is zero. This supports the + alternative hypothesis that the true value of Somers' D for the population + is nonzero. + + """ + x, y = np.array(x), np.array(y) + if x.ndim == 1: + if x.size != y.size: + raise ValueError("Rankings must be of equal length.") + table = scipy.stats.contingency.crosstab(x, y)[1] + elif x.ndim == 2: + if np.any(x < 0): + raise ValueError("All elements of the contingency table must be " + "non-negative.") + if np.any(x != x.astype(int)): + raise ValueError("All elements of the contingency table must be " + "integer.") + if x.nonzero()[0].size < 2: + raise ValueError("At least two elements of the contingency table " + "must be nonzero.") + table = x + else: + raise ValueError("x must be either a 1D or 2D array") + # The table type is converted to a float to avoid an integer overflow + d, p = _somers_d(table.astype(float), alternative) + + # add alias for consistency with other correlation functions + res = SomersDResult(d, p, table) + res.correlation = d + return res + + +# This could be combined with `_all_partitions` in `_resampling.py` +def _all_partitions(nx, ny): + """ + Partition a set of indices into two fixed-length sets in all possible ways + + Partition a set of indices 0 ... nx + ny - 1 into two sets of length nx and + ny in all possible ways (ignoring order of elements). + """ + z = np.arange(nx+ny) + for c in combinations(z, nx): + x = np.array(c) + mask = np.ones(nx+ny, bool) + mask[x] = False + y = z[mask] + yield x, y + + +def _compute_log_combinations(n): + """Compute all log combination of C(n, k).""" + gammaln_arr = gammaln(np.arange(n + 1) + 1) + return gammaln(n + 1) - gammaln_arr - gammaln_arr[::-1] + + +@dataclass +class BarnardExactResult: + statistic: float + pvalue: float + + +@xp_capabilities(np_only=True) +def barnard_exact(table, alternative="two-sided", pooled=True, n=32): + r"""Perform a Barnard exact test on a 2x2 contingency table. + + Parameters + ---------- + table : array_like of ints + A 2x2 contingency table. Elements should be non-negative integers. + + alternative : {'two-sided', 'less', 'greater'}, optional + Defines the null and alternative hypotheses. Default is 'two-sided'. + Please see explanations in the Notes section below. + + pooled : bool, optional + Whether to compute score statistic with pooled variance (as in + Student's t-test, for example) or unpooled variance (as in Welch's + t-test). Default is ``True``. + + n : int, optional + Number of sampling points used in the construction of the sampling + method. Note that this argument will automatically be converted to + the next higher power of 2 since `scipy.stats.qmc.Sobol` is used to + select sample points. Default is 32. Must be positive. In most cases, + 32 points is enough to reach good precision. More points comes at + performance cost. + + Returns + ------- + ber : BarnardExactResult + A result object with the following attributes. + + statistic : float + The Wald statistic with pooled or unpooled variance, depending + on the user choice of `pooled`. + + pvalue : float + P-value, the probability of obtaining a distribution at least as + extreme as the one that was actually observed, assuming that the + null hypothesis is true. + + See Also + -------- + chi2_contingency : Chi-square test of independence of variables in a + contingency table. + fisher_exact : Fisher exact test on a 2x2 contingency table. + boschloo_exact : Boschloo's exact test on a 2x2 contingency table, + which is a uniformly more powerful alternative to Fisher's exact test. + + Notes + ----- + Barnard's test is an exact test used in the analysis of contingency + tables. It examines the association of two categorical variables, and + is a more powerful alternative than Fisher's exact test + for 2x2 contingency tables. + + Let's define :math:`X_0` a 2x2 matrix representing the observed sample, + where each column stores the binomial experiment, as in the example + below. Let's also define :math:`p_1, p_2` the theoretical binomial + probabilities for :math:`x_{11}` and :math:`x_{12}`. When using + Barnard exact test, we can assert three different null hypotheses : + + - :math:`H_0 : p_1 \geq p_2` versus :math:`H_1 : p_1 < p_2`, + with `alternative` = "less" + + - :math:`H_0 : p_1 \leq p_2` versus :math:`H_1 : p_1 > p_2`, + with `alternative` = "greater" + + - :math:`H_0 : p_1 = p_2` versus :math:`H_1 : p_1 \neq p_2`, + with `alternative` = "two-sided" (default one) + + In order to compute Barnard's exact test, we are using the Wald + statistic [3]_ with pooled or unpooled variance. + Under the default assumption that both variances are equal + (``pooled = True``), the statistic is computed as: + + .. math:: + + T(X) = \frac{ + \hat{p}_1 - \hat{p}_2 + }{ + \sqrt{ + \hat{p}(1 - \hat{p}) + (\frac{1}{c_1} + + \frac{1}{c_2}) + } + } + + with :math:`\hat{p}_1, \hat{p}_2` and :math:`\hat{p}` the estimator of + :math:`p_1, p_2` and :math:`p`, the latter being the combined probability, + given the assumption that :math:`p_1 = p_2`. + + If this assumption is invalid (``pooled = False``), the statistic is: + + .. math:: + + T(X) = \frac{ + \hat{p}_1 - \hat{p}_2 + }{ + \sqrt{ + \frac{\hat{p}_1 (1 - \hat{p}_1)}{c_1} + + \frac{\hat{p}_2 (1 - \hat{p}_2)}{c_2} + } + } + + The p-value is then computed as: + + .. math:: + + \sum + \binom{c_1}{x_{11}} + \binom{c_2}{x_{12}} + \pi^{x_{11} + x_{12}} + (1 - \pi)^{t - x_{11} - x_{12}} + + where the sum is over all 2x2 contingency tables :math:`X` such that: + + * :math:`T(X) \leq T(X_0)` when `alternative` = "less", + * :math:`T(X) \geq T(X_0)` when `alternative` = "greater", or + * :math:`T(X) \geq |T(X_0)|` when `alternative` = "two-sided". + + Above, :math:`c_1, c_2` are the sum of the columns 1 and 2, + and :math:`t` the total (sum of the 4 sample's element). + + The returned p-value is the maximum p-value taken over the nuisance + parameter :math:`\pi`, where :math:`0 \leq \pi \leq 1`. + + This function's complexity is :math:`O(n c_1 c_2)`, where `n` is the + number of sample points. + + References + ---------- + .. [1] Barnard, G. A. "Significance Tests for 2x2 Tables". *Biometrika*. + 34.1/2 (1947): 123-138. :doi:`10.2307/2332517`. + + .. [2] Mehta, Cyrus R., and Pralay Senchaudhuri. "Conditional versus + unconditional exact tests for comparing two binomials." + *Cytel Software Corporation* 675 (2003): 1-5. + + .. [3] "Wald Test". *Wikipedia*. https://en.wikipedia.org/wiki/Wald_test + + Examples + -------- + An example use of Barnard's test is presented in [2]_. + + Consider the following example of a vaccine efficacy study + (Chan, 1998). In a randomized clinical trial of 30 subjects, 15 were + inoculated with a recombinant DNA influenza vaccine and the 15 were + inoculated with a placebo. Twelve of the 15 subjects in the placebo + group (80%) eventually became infected with influenza whereas for the + vaccine group, only 7 of the 15 subjects (47%) became infected. The + data are tabulated as a 2 x 2 table:: + + Vaccine Placebo + Yes 7 12 + No 8 3 + + When working with statistical hypothesis testing, we usually use a + threshold probability or significance level upon which we decide + to reject the null hypothesis :math:`H_0`. Suppose we choose the common + significance level of 5%. + + Our alternative hypothesis is that the vaccine will lower the chance of + becoming infected with the virus; that is, the probability :math:`p_1` of + catching the virus with the vaccine will be *less than* the probability + :math:`p_2` of catching the virus without the vaccine. Therefore, we call + `barnard_exact` with the ``alternative="less"`` option: + + >>> import scipy.stats as stats + >>> res = stats.barnard_exact([[7, 12], [8, 3]], alternative="less") + >>> res.statistic + -1.894 + >>> res.pvalue + 0.03407 + + Under the null hypothesis that the vaccine will not lower the chance of + becoming infected, the probability of obtaining test results at least as + extreme as the observed data is approximately 3.4%. Since this p-value is + less than our chosen significance level, we have evidence to reject + :math:`H_0` in favor of the alternative. + + Suppose we had used Fisher's exact test instead: + + >>> _, pvalue = stats.fisher_exact([[7, 12], [8, 3]], alternative="less") + >>> pvalue + 0.0640 + + With the same threshold significance of 5%, we would not have been able + to reject the null hypothesis in favor of the alternative. As stated in + [2]_, Barnard's test is uniformly more powerful than Fisher's exact test + because Barnard's test does not condition on any margin. Fisher's test + should only be used when both sets of marginals are fixed. + + """ + if n <= 0: + raise ValueError( + "Number of points `n` must be strictly positive, " + f"found {n!r}" + ) + + table = np.asarray(table, dtype=np.int64) + + if not table.shape == (2, 2): + raise ValueError("The input `table` must be of shape (2, 2).") + + if np.any(table < 0): + raise ValueError("All values in `table` must be nonnegative.") + + if 0 in table.sum(axis=0): + # If both values in column are zero, the p-value is 1 and + # the score's statistic is NaN. + return BarnardExactResult(np.nan, 1.0) + + total_col_1, total_col_2 = table.sum(axis=0) + + x1 = np.arange(total_col_1 + 1, dtype=np.int64).reshape(-1, 1) + x2 = np.arange(total_col_2 + 1, dtype=np.int64).reshape(1, -1) + + # We need to calculate the wald statistics for each combination of x1 and + # x2. + p1, p2 = x1 / total_col_1, x2 / total_col_2 + + if pooled: + p = (x1 + x2) / (total_col_1 + total_col_2) + variances = p * (1 - p) * (1 / total_col_1 + 1 / total_col_2) + else: + variances = p1 * (1 - p1) / total_col_1 + p2 * (1 - p2) / total_col_2 + + # To avoid warning when dividing by 0 + with np.errstate(divide="ignore", invalid="ignore"): + wald_statistic = np.divide((p1 - p2), np.sqrt(variances)) + + wald_statistic[p1 == p2] = 0 # Removing NaN values + + wald_stat_obs = wald_statistic[table[0, 0], table[0, 1]] + + if alternative == "two-sided": + index_arr = np.abs(wald_statistic) >= abs(wald_stat_obs) + elif alternative == "less": + index_arr = wald_statistic <= wald_stat_obs + elif alternative == "greater": + index_arr = wald_statistic >= wald_stat_obs + else: + msg = ( + "`alternative` should be one of {'two-sided', 'less', 'greater'}," + f" found {alternative!r}" + ) + raise ValueError(msg) + + x1_sum_x2 = x1 + x2 + + x1_log_comb = _compute_log_combinations(total_col_1) + x2_log_comb = _compute_log_combinations(total_col_2) + x1_sum_x2_log_comb = x1_log_comb[x1] + x2_log_comb[x2] + + result = shgo( + _get_binomial_log_p_value_with_nuisance_param, + args=(x1_sum_x2, x1_sum_x2_log_comb, index_arr), + bounds=((0, 1),), + n=n, + sampling_method="sobol", + ) + + # result.fun is the negative log pvalue and therefore needs to be + # changed before return + p_value = np.clip(np.exp(-result.fun), a_min=0, a_max=1) + return BarnardExactResult(wald_stat_obs, p_value) + + +@dataclass +class BoschlooExactResult: + statistic: float + pvalue: float + + +@xp_capabilities(np_only=True) +def boschloo_exact(table, alternative="two-sided", n=32): + r"""Perform Boschloo's exact test on a 2x2 contingency table. + + Parameters + ---------- + table : array_like of ints + A 2x2 contingency table. Elements should be non-negative integers. + + alternative : {'two-sided', 'less', 'greater'}, optional + Defines the null and alternative hypotheses. Default is 'two-sided'. + Please see explanations in the Notes section below. + + n : int, optional + Number of sampling points used in the construction of the sampling + method. Note that this argument will automatically be converted to + the next higher power of 2 since `scipy.stats.qmc.Sobol` is used to + select sample points. Default is 32. Must be positive. In most cases, + 32 points is enough to reach good precision. More points comes at + performance cost. + + Returns + ------- + ber : BoschlooExactResult + A result object with the following attributes. + + statistic : float + The statistic used in Boschloo's test; that is, the p-value + from Fisher's exact test. + + pvalue : float + P-value, the probability of obtaining a distribution at least as + extreme as the one that was actually observed, assuming that the + null hypothesis is true. + + See Also + -------- + chi2_contingency : Chi-square test of independence of variables in a + contingency table. + fisher_exact : Fisher exact test on a 2x2 contingency table. + barnard_exact : Barnard's exact test, which is a more powerful alternative + than Fisher's exact test for 2x2 contingency tables. + + Notes + ----- + Boschloo's test is an exact test used in the analysis of contingency + tables. It examines the association of two categorical variables, and + is a uniformly more powerful alternative to Fisher's exact test + for 2x2 contingency tables. + + Boschloo's exact test uses the p-value of Fisher's exact test as a + statistic, and Boschloo's p-value is the probability under the null + hypothesis of observing such an extreme value of this statistic. + + Let's define :math:`X_0` a 2x2 matrix representing the observed sample, + where each column stores the binomial experiment, as in the example + below. Let's also define :math:`p_1, p_2` the theoretical binomial + probabilities for :math:`x_{11}` and :math:`x_{12}`. When using + Boschloo exact test, we can assert three different alternative hypotheses: + + - :math:`H_0 : p_1=p_2` versus :math:`H_1 : p_1 < p_2`, + with `alternative` = "less" + + - :math:`H_0 : p_1=p_2` versus :math:`H_1 : p_1 > p_2`, + with `alternative` = "greater" + + - :math:`H_0 : p_1=p_2` versus :math:`H_1 : p_1 \neq p_2`, + with `alternative` = "two-sided" (default) + + There are multiple conventions for computing a two-sided p-value when the + null distribution is asymmetric. Here, we apply the convention that the + p-value of a two-sided test is twice the minimum of the p-values of the + one-sided tests (clipped to 1.0). Note that `fisher_exact` follows a + different convention, so for a given `table`, the statistic reported by + `boschloo_exact` may differ from the p-value reported by `fisher_exact` + when ``alternative='two-sided'``. + + .. versionadded:: 1.7.0 + + References + ---------- + .. [1] R.D. Boschloo. "Raised conditional level of significance for the + 2 x 2-table when testing the equality of two probabilities", + Statistica Neerlandica, 24(1), 1970 + + .. [2] "Boschloo's test", Wikipedia, + https://en.wikipedia.org/wiki/Boschloo%27s_test + + .. [3] Lise M. Saari et al. "Employee attitudes and job satisfaction", + Human Resource Management, 43(4), 395-407, 2004, + :doi:`10.1002/hrm.20032`. + + Examples + -------- + In the following example, we consider the article "Employee + attitudes and job satisfaction" [3]_ + which reports the results of a survey from 63 scientists and 117 college + professors. Of the 63 scientists, 31 said they were very satisfied with + their jobs, whereas 74 of the college professors were very satisfied + with their work. Is this significant evidence that college + professors are happier with their work than scientists? + The following table summarizes the data mentioned above:: + + college professors scientists + Very Satisfied 74 31 + Dissatisfied 43 32 + + When working with statistical hypothesis testing, we usually use a + threshold probability or significance level upon which we decide + to reject the null hypothesis :math:`H_0`. Suppose we choose the common + significance level of 5%. + + Our alternative hypothesis is that college professors are truly more + satisfied with their work than scientists. Therefore, we expect + :math:`p_1` the proportion of very satisfied college professors to be + greater than :math:`p_2`, the proportion of very satisfied scientists. + We thus call `boschloo_exact` with the ``alternative="greater"`` option: + + >>> import scipy.stats as stats + >>> res = stats.boschloo_exact([[74, 31], [43, 32]], alternative="greater") + >>> res.statistic + 0.0483 + >>> res.pvalue + 0.0355 + + Under the null hypothesis that scientists are happier in their work than + college professors, the probability of obtaining test + results at least as extreme as the observed data is approximately 3.55%. + Since this p-value is less than our chosen significance level, we have + evidence to reject :math:`H_0` in favor of the alternative hypothesis. + + """ + hypergeom = distributions.hypergeom + + if n <= 0: + raise ValueError( + "Number of points `n` must be strictly positive," + f" found {n!r}" + ) + + table = np.asarray(table, dtype=np.int64) + + if not table.shape == (2, 2): + raise ValueError("The input `table` must be of shape (2, 2).") + + if np.any(table < 0): + raise ValueError("All values in `table` must be nonnegative.") + + if 0 in table.sum(axis=0): + # If both values in column are zero, the p-value is 1 and + # the score's statistic is NaN. + return BoschlooExactResult(np.nan, np.nan) + + total_col_1, total_col_2 = table.sum(axis=0) + total = total_col_1 + total_col_2 + x1 = np.arange(total_col_1 + 1, dtype=np.int64).reshape(1, -1) + x2 = np.arange(total_col_2 + 1, dtype=np.int64).reshape(-1, 1) + x1_sum_x2 = x1 + x2 + + if alternative == 'less': + pvalues = hypergeom.cdf(x1, total, x1_sum_x2, total_col_1).T + elif alternative == 'greater': + # Same formula as the 'less' case, but with the second column. + pvalues = hypergeom.cdf(x2, total, x1_sum_x2, total_col_2).T + elif alternative == 'two-sided': + boschloo_less = boschloo_exact(table, alternative="less", n=n) + boschloo_greater = boschloo_exact(table, alternative="greater", n=n) + + res = ( + boschloo_less if boschloo_less.pvalue < boschloo_greater.pvalue + else boschloo_greater + ) + + # Two-sided p-value is defined as twice the minimum of the one-sided + # p-values + pvalue = np.clip(2 * res.pvalue, a_min=0, a_max=1) + return BoschlooExactResult(res.statistic, pvalue) + else: + msg = ( + f"`alternative` should be one of {'two-sided', 'less', 'greater'}," + f" found {alternative!r}" + ) + raise ValueError(msg) + + fisher_stat = pvalues[table[0, 0], table[0, 1]] + + # fisher_stat * (1+1e-13) guards us from small numerical error. It is + # equivalent to np.isclose with relative tol of 1e-13 and absolute tol of 0 + # For more throughout explanations, see gh-14178 + index_arr = pvalues <= fisher_stat * (1+1e-13) + + x1, x2, x1_sum_x2 = x1.T, x2.T, x1_sum_x2.T + x1_log_comb = _compute_log_combinations(total_col_1) + x2_log_comb = _compute_log_combinations(total_col_2) + x1_sum_x2_log_comb = x1_log_comb[x1] + x2_log_comb[x2] + + result = shgo( + _get_binomial_log_p_value_with_nuisance_param, + args=(x1_sum_x2, x1_sum_x2_log_comb, index_arr), + bounds=((0, 1),), + n=n, + sampling_method="sobol", + ) + + # result.fun is the negative log pvalue and therefore needs to be + # changed before return + p_value = np.clip(np.exp(-result.fun), a_min=0, a_max=1) + return BoschlooExactResult(fisher_stat, p_value) + + +def _get_binomial_log_p_value_with_nuisance_param( + nuisance_param, x1_sum_x2, x1_sum_x2_log_comb, index_arr +): + r""" + Compute the log pvalue in respect of a nuisance parameter considering + a 2x2 sample space. + + Parameters + ---------- + nuisance_param : float + nuisance parameter used in the computation of the maximisation of + the p-value. Must be between 0 and 1 + + x1_sum_x2 : ndarray + Sum of x1 and x2 inside barnard_exact + + x1_sum_x2_log_comb : ndarray + sum of the log combination of x1 and x2 + + index_arr : ndarray of boolean + + Returns + ------- + p_value : float + Return the maximum p-value considering every nuisance parameter + between 0 and 1 + + Notes + ----- + + Both Barnard's test and Boschloo's test iterate over a nuisance parameter + :math:`\pi \in [0, 1]` to find the maximum p-value. To search this + maxima, this function return the negative log pvalue with respect to the + nuisance parameter passed in params. This negative log p-value is then + used in `shgo` to find the minimum negative pvalue which is our maximum + pvalue. + + Also, to compute the different combination used in the + p-values' computation formula, this function uses `gammaln` which is + more tolerant for large value than `scipy.special.comb`. `gammaln` gives + a log combination. For the little precision loss, performances are + improved a lot. + """ + t1, t2 = x1_sum_x2.shape + n = t1 + t2 - 2 + with np.errstate(divide="ignore", invalid="ignore"): + log_nuisance = np.log( + nuisance_param, + out=np.zeros_like(nuisance_param), + where=nuisance_param >= 0, + ) + log_1_minus_nuisance = np.log( + 1 - nuisance_param, + out=np.zeros_like(nuisance_param), + where=1 - nuisance_param >= 0, + ) + + nuisance_power_x1_x2 = log_nuisance * x1_sum_x2 + nuisance_power_x1_x2[(x1_sum_x2 == 0)[:, :]] = 0 + + nuisance_power_n_minus_x1_x2 = log_1_minus_nuisance * (n - x1_sum_x2) + nuisance_power_n_minus_x1_x2[(x1_sum_x2 == n)[:, :]] = 0 + + tmp_log_values_arr = ( + x1_sum_x2_log_comb + + nuisance_power_x1_x2 + + nuisance_power_n_minus_x1_x2 + ) + + tmp_values_from_index = tmp_log_values_arr[index_arr] + + # To avoid dividing by zero in log function and getting inf value, + # values are centered according to the max + max_value = tmp_values_from_index.max() + + # To have better result's precision, the log pvalue is taken here. + # Indeed, pvalue is included inside [0, 1] interval. Passing the + # pvalue to log makes the interval a lot bigger ([-inf, 0]), and thus + # help us to achieve better precision + with np.errstate(divide="ignore", invalid="ignore"): + log_probs = np.exp(tmp_values_from_index - max_value).sum() + log_pvalue = max_value + np.log( + log_probs, + out=np.full_like(log_probs, -np.inf), + where=log_probs > 0, + ) + + # Since shgo find the minima, minus log pvalue is returned + return -log_pvalue + + +@np.vectorize(otypes=[np.float64]) +def _pval_cvm_2samp_exact(s, m, n): + """ + Compute the exact p-value of the Cramer-von Mises two-sample test + for a given value s of the test statistic. + m and n are the sizes of the samples. + + [1] Y. Xiao, A. Gordon, and A. Yakovlev, "A C++ Program for + the Cramér-Von Mises Two-Sample Test", J. Stat. Soft., + vol. 17, no. 8, pp. 1-15, Dec. 2006. + [2] T. W. Anderson "On the Distribution of the Two-Sample Cramer-von Mises + Criterion," The Annals of Mathematical Statistics, Ann. Math. Statist. + 33(3), 1148-1159, (September, 1962) + """ + + # [1, p. 3] + m, n = int(m), int(n) + lcm = np.lcm(m, n) + # [1, p. 4], below eq. 3 + a = lcm // m + b = lcm // n + # Combine Eq. 9 in [2] with Eq. 2 in [1] and solve for $\zeta$ + # Hint: `s` is $U$ in [2], and $T_2$ in [1] is $T$ in [2] + mn = m * n + zeta = lcm ** 2 * (m + n) * (6 * s - mn * (4 * mn - 1)) // (6 * mn ** 2) + + # bound maximum value that may appear in `gs` (remember both rows!) + zeta_bound = lcm**2 * (m + n) # bound elements in row 1 + combinations = math.comb(m + n, m) # sum of row 2 + max_gs = max(zeta_bound, combinations) + dtype = np.min_scalar_type(max_gs) + + # the frequency table of $g_{u, v}^+$ defined in [1, p. 6] + gs = ([np.array([[0], [1]], dtype=dtype)] + + [np.empty((2, 0), dtype=dtype) for _ in range(m)]) + for u in range(n + 1): + next_gs = [] + tmp = np.empty((2, 0), dtype=dtype) + for v, g in enumerate(gs): + # Calculate g recursively with eq. 11 in [1]. Even though it + # doesn't look like it, this also does 12/13 (all of Algorithm 1). + vi, i0, i1 = np.intersect1d(tmp[0], g[0], return_indices=True) + tmp = np.concatenate([ + np.stack([vi, tmp[1, i0] + g[1, i1]]), + np.delete(tmp, i0, 1), + np.delete(g, i1, 1) + ], 1) + res = (a * v - b * u) ** 2 + tmp[0] += res.astype(dtype) + next_gs.append(tmp) + gs = next_gs + value, freq = gs[m] + return np.float64(np.sum(freq[value >= zeta]) / combinations) + + +def _pval_cvm_2samp_asymptotic(t, N, nx, ny, k, *, xp): + # compute expected value and variance of T (eq. 11 and 14 in [2]) + nx, ny = xp.asarray(nx, dtype=t.dtype), xp.asarray(ny, dtype=t.dtype) + et = (1 + 1 / N) / 6 + vt = (N + 1) * (4 * k * N - 3 * (nx ** 2 + ny ** 2) - 2 * k) + vt = vt / (45 * N ** 2 * 4 * k) + + # computed the normalized statistic (eq. 15 in [2]) + tn = 1 / 6 + (t - et) / xp.sqrt(45 * vt) + + # approximate distribution of tn with limiting distribution + # of the one-sample test statistic + # if tn < 0.003, the _cdf_cvm_inf(tn) < 1.28*1e-18, return 1.0 directly + p = xpx.apply_where(tn >= 0.003, + (tn,), + lambda tn: xp.clip(1. - _cdf_cvm_inf(tn, xp=xp), 0.), + fill_value = 1.) + return p + + +@xp_capabilities(skip_backends=[('dask.array', 'needs rankdata')], + cpu_only=True, jax_jit=False, # due to p-value calculation + marray=True, + extra_note="Only `method='exact'` is compatible with MArray input.") +@_axis_nan_policy_factory(CramerVonMisesResult, n_samples=2, too_small=1, + result_to_tuple=_cvm_result_to_tuple) +def cramervonmises_2samp(x, y, method='auto', *, axis=0): + r"""Perform the two-sample Cramér-von Mises test for goodness of fit. + + This is the two-sample version of the Cramér-von Mises test ([1]_): + for two independent samples :math:`X_1, ..., X_n` and + :math:`Y_1, ..., Y_m`, the null hypothesis is that the samples + come from the same (unspecified) continuous distribution. + + The test statistic :math:`T` is defined as in [1]_: + + .. math:: + T = \frac{nm}{n+m}\omega^2 = + \frac{U}{n m (n+m)} - \frac{4 m n - 1}{6(m+n)} + + where :math:`U` is defined as below, and :math:`\omega^2` is the Cramér-von + Mises criterion. The function :math:`r(\cdot)` here denotes the rank of the + observed values :math:`x_i` and :math:`y_j` within the pooled sample of size + :math:`n + m`, with ties assigned mid-rank values: + + .. math:: + U = n \sum_{i=1}^n (r(x_i)-i)^2 + m \sum_{j=1}^m (r(y_j)-j)^2 + + Parameters + ---------- + x : array_like + A 1-D array of observed values of the random variables :math:`X_i`. + Must contain at least two observations. + y : array_like + A 1-D array of observed values of the random variables :math:`Y_i`. + Must contain at least two observations. + method : {'auto', 'asymptotic', 'exact'}, optional + The method used to compute the p-value, see Notes for details. + The default is 'auto'. + axis : int or tuple of ints, default: 0 + If an int or tuple of ints, the axis or axes of the input along which + to compute the statistic. The statistic of each axis-slice (e.g. row) + of the input will appear in a corresponding element of the output. + If ``None``, the input will be raveled before computing the statistic. + + Returns + ------- + res : object with attributes + statistic : float + Cramér-von Mises statistic :math:`T`. + pvalue : float + The p-value. + + See Also + -------- + cramervonmises, anderson_ksamp, epps_singleton_2samp, ks_2samp + + Notes + ----- + .. versionadded:: 1.7.0 + + The statistic is computed according to equation 9 in [2]_. The + calculation of the p-value depends on the keyword `method`: + + - ``asymptotic``: The p-value is approximated by using the limiting + distribution of the test statistic. + - ``exact``: The exact p-value is computed by enumerating all + possible combinations of the test statistic, see [2]_. + + If ``method='auto'``, the exact approach is used + if both samples contain equal to or less than 20 observations, + otherwise the asymptotic distribution is used. + + If the underlying distribution is not continuous, the p-value is likely to + be conservative (Section 6.2 in [3]_). When ranking the data to compute + the test statistic, midranks are used if there are ties. + + References + ---------- + .. [1] https://en.wikipedia.org/wiki/Cramer-von_Mises_criterion + .. [2] Anderson, T.W. (1962). On the distribution of the two-sample + Cramer-von-Mises criterion. The Annals of Mathematical + Statistics, pp. 1148-1159. + .. [3] Conover, W.J., Practical Nonparametric Statistics, 1971. + + Examples + -------- + + Suppose we wish to test whether two samples generated by + ``scipy.stats.norm.rvs`` have the same distribution. We choose a + significance level of alpha=0.05. + + >>> import numpy as np + >>> from scipy import stats + >>> rng = np.random.default_rng() + >>> x = stats.norm.rvs(size=100, random_state=rng) + >>> y = stats.norm.rvs(size=70, random_state=rng) + >>> res = stats.cramervonmises_2samp(x, y) + >>> res.statistic, res.pvalue + (0.29376470588235293, 0.1412873014573014) + + The p-value exceeds our chosen significance level, so we do not + reject the null hypothesis that the observed samples are drawn from the + same distribution. + + For small sample sizes, one can compute the exact p-values: + + >>> x = stats.norm.rvs(size=7, random_state=rng) + >>> y = stats.t.rvs(df=2, size=6, random_state=rng) + >>> res = stats.cramervonmises_2samp(x, y, method='exact') + >>> res.statistic, res.pvalue + (0.197802197802198, 0.31643356643356646) + + The p-value based on the asymptotic distribution is a good approximation + even though the sample size is small. + + >>> res = stats.cramervonmises_2samp(x, y, method='asymptotic') + >>> res.statistic, res.pvalue + (0.197802197802198, 0.2966041181527128) + + Independent of the method, one would not reject the null hypothesis at the + chosen significance level in this example. + + """ + xp = array_namespace(x, y) + length_x = x.shape[-1] + length_y = y.shape[-1] + count_x = _count_nonmasked(x, axis=-1, xp=xp) + count_y = _count_nonmasked(y, axis=-1, xp=xp) + + if length_x <= 1 or length_y <= 1: # only needed for `test_axis_nan_policy` + raise ValueError('x and y must contain at least two observations.') + if method not in ['auto', 'exact', 'asymptotic']: + raise ValueError('method must be either auto, exact or asymptotic.') + + if is_marray(xp) and method in {'auto', 'asymptotic'}: + message = "Only `method='exact'` is compatible with MArray." + raise ValueError(message) + + if method == 'auto': + if max(count_x, count_y) > 20: + method = 'asymptotic' + else: + method = 'exact' + + # axis=-1 is guaranteed by _axis_nan_policy decorator + xa = xp.sort(x, axis=-1) + ya = xp.sort(y, axis=-1) + + # get ranks of x and y in the pooled sample + z = xp.concat([xa, ya], axis=-1) + # in case of ties, use midrank (see [1]) + r = scipy.stats.rankdata(z, method='average', axis=-1) + dtype = xp_result_type(x, y, force_floating=True, xp=xp) + rx = r[..., :length_x] + ry = r[..., length_x:] + + # compute U (eq. 10 in [2]) + u = (count_x * xp.sum((rx - xp.arange(1, length_x+1, dtype=dtype))**2, axis=-1) + + count_y * xp.sum((ry - xp.arange(1, length_y+1, dtype=dtype))**2, axis=-1)) + + # compute T (eq. 9 in [2]) + k, N = count_x*count_y, count_x + count_y + t = u / (k*N) - (4*k - 1)/(6*N) + + if method == 'exact': + if is_marray(xp): + u, count_x, count_y = u.data, count_x.data, count_y.data + p = _pval_cvm_2samp_exact(np.asarray(u), count_x, count_y) + p = xp.asarray(p, dtype=dtype) + else: + p = _pval_cvm_2samp_asymptotic(t, N, count_x, count_y, k, xp=xp) + + t = t[()] if t.ndim == 0 else t + p = p[()] if p.ndim == 0 else p + return CramerVonMisesResult(statistic=t, pvalue=p) + + +class TukeyHSDResult: + """Result of `scipy.stats.tukey_hsd`. + + Attributes + ---------- + statistic : float ndarray + The computed statistic of the test for each comparison. The element + at index ``(i, j)`` is the statistic for the comparison between groups + ``i`` and ``j``. + pvalue : float ndarray + The associated p-value from the studentized range distribution. The + element at index ``(i, j)`` is the p-value for the comparison + between groups ``i`` and ``j``. + + Notes + ----- + The string representation of this object displays the most recently + calculated confidence interval, and if none have been previously + calculated, it will evaluate ``confidence_interval()``. + + References + ---------- + .. [1] NIST/SEMATECH e-Handbook of Statistical Methods, "7.4.7.1. Tukey's + Method." + https://www.itl.nist.gov/div898/handbook/prc/section4/prc471.htm, + 28 November 2020. + .. [2] P. A. Games and J. F. Howell, "Pairwise Multiple Comparison Procedures + with Unequal N's and/or Variances: A Monte Carlo Study," Journal of + Educational Statistics, vol. 1, no. 2, pp. 113-125, Jun. 1976, + :doi:`10.3102/10769986001002113`. + """ + + def __init__(self, statistic, pvalue, _ntreatments, _df, _stand_err): + self.statistic = statistic + self.pvalue = pvalue + self._ntreatments = _ntreatments + self._df = _df + self._stand_err = _stand_err + self._ci = None + self._ci_cl = None + + def __str__(self): + # Note: `__str__` prints the confidence intervals from the most + # recent call to `confidence_interval`. If it has not been called, + # it will be called with the default CL of .95. + if self._ci is None: + self.confidence_interval(confidence_level=.95) + s = ("Pairwise Group Comparisons" + f" ({self._ci_cl*100:.1f}% Confidence Interval)\n") + s += "Comparison Statistic p-value Lower CI Upper CI\n" + for i, j in np.ndindex(self.pvalue.shape): + if i != j: + s += (f" ({i} - {j}) {self.statistic[i, j]:>10.3f}" + f"{self.pvalue[i, j]:>10.3f}" + f"{self._ci.low[i, j]:>10.3f}" + f"{self._ci.high[i, j]:>10.3f}\n") + return s + + def confidence_interval(self, confidence_level=.95): + """Compute the confidence interval for the specified confidence level. + + Parameters + ---------- + confidence_level : float, optional + Confidence level for the computed confidence interval + of the estimated proportion. Default is .95. + + Returns + ------- + ci : ``ConfidenceInterval`` object + The object has attributes ``low`` and ``high`` that hold the + lower and upper bounds of the confidence intervals for each + comparison. The high and low values are accessible for each + comparison at index ``(i, j)`` between groups ``i`` and ``j``. + + References + ---------- + .. [1] NIST/SEMATECH e-Handbook of Statistical Methods, "7.4.7.1. + Tukey's Method." + https://www.itl.nist.gov/div898/handbook/prc/section4/prc471.htm, + 28 November 2020. + .. [2] P. A. Games and J. F. Howell, "Pairwise Multiple Comparison Procedures + with Unequal N's and/or Variances: A Monte Carlo Study," Journal of + Educational Statistics, vol. 1, no. 2, pp. 113-125, Jun. 1976, + :doi:`10.3102/10769986001002113`. + + Examples + -------- + >>> from scipy.stats import tukey_hsd + >>> group0 = [24.5, 23.5, 26.4, 27.1, 29.9] + >>> group1 = [28.4, 34.2, 29.5, 32.2, 30.1] + >>> group2 = [26.1, 28.3, 24.3, 26.2, 27.8] + >>> result = tukey_hsd(group0, group1, group2) + >>> ci = result.confidence_interval() + >>> ci.low + array([[-3.649159, -8.249159, -3.909159], + [ 0.950841, -3.649159, 0.690841], + [-3.389159, -7.989159, -3.649159]]) + >>> ci.high + array([[ 3.649159, -0.950841, 3.389159], + [ 8.249159, 3.649159, 7.989159], + [ 3.909159, -0.690841, 3.649159]]) + """ + # check to see if the supplied confidence level matches that of the + # previously computed CI. + if (self._ci is not None and self._ci_cl is not None and + confidence_level == self._ci_cl): + return self._ci + + if not 0 < confidence_level < 1: + raise ValueError("Confidence level must be between 0 and 1.") + # determine the critical value of the studentized range using the + # appropriate confidence level, number of treatments, and degrees + # of freedom. See [1] "Confidence limits for Tukey's method" / [2] p.117 + # "H0 was rejected if...". Note that in the cases of unequal sample sizes, + # there will be a criterion for each group comparison. + params = (confidence_level, self._ntreatments, self._df) + srd = distributions.studentized_range.ppf(*params) + # also called maximum critical value, the confidence_radius is the + # studentized range critical value * the square root of mean square + # error over the sample size. + confidence_radius = srd * self._stand_err + # the confidence levels are determined by the + # `mean_differences` +- `confidence_radius` + upper_conf = self.statistic + confidence_radius + lower_conf = self.statistic - confidence_radius + self._ci = ConfidenceInterval(low=lower_conf, high=upper_conf) + self._ci_cl = confidence_level + return self._ci + + +def _tukey_hsd_iv(args, equal_var): + if (len(args)) < 2: + raise ValueError("There must be more than 1 treatment.") + if not isinstance(equal_var, bool): + raise TypeError("Expected a boolean value for 'equal_var'") + args = [np.asarray(arg) for arg in args] + for arg in args: + if arg.ndim != 1: + raise ValueError("Input samples must be one-dimensional.") + if arg.size <= 1: + raise ValueError("Input sample size must be greater than one.") + if np.isinf(arg).any(): + raise ValueError("Input samples must be finite.") + return args + + +@xp_capabilities(np_only=True) +def tukey_hsd(*samples, equal_var=True): + """Perform Tukey's HSD test for equality of means over multiple treatments. + + Tukey's honestly significant difference (HSD) test performs pairwise + comparison of means for a set of samples. Whereas ANOVA (e.g. `f_oneway`) + assesses whether the true means underlying each sample are identical, + Tukey's HSD is a post hoc test used to compare the mean of each sample + to the mean of each other sample. + + The null hypothesis is that the distributions underlying the samples all + have the same mean. The test statistic, which is computed for every + possible pairing of samples, is simply the difference between the sample + means. For each pair, the p-value is the probability under the null + hypothesis (and other assumptions; see notes) of observing such an extreme + value of the statistic, considering that many pairwise comparisons are + being performed. Confidence intervals for the difference between each pair + of means are also available. + + Parameters + ---------- + *samples : array_like + The sample measurements for each group. There must be at least + two arguments. + equal_var : bool, optional + If True (default) and equal sample size, perform Tukey-HSD test [6]. + If True and unequal sample size, perform Tukey-Kramer test [4]_. + If False, perform Games-Howell test [7]_, which does not assume equal variances. + + Returns + ------- + result : `~scipy.stats._result_classes.TukeyHSDResult` instance + The return value is an object with the following attributes: + + statistic : float ndarray + The computed statistic of the test for each comparison. The element + at index ``(i, j)`` is the statistic for the comparison between + groups ``i`` and ``j``. + pvalue : float ndarray + The computed p-value of the test for each comparison. The element + at index ``(i, j)`` is the p-value for the comparison between + groups ``i`` and ``j``. + + The object has the following methods: + + confidence_interval(confidence_level=0.95): + Compute the confidence interval for the specified confidence level. + + See Also + -------- + dunnett : performs comparison of means against a control group. + + Notes + ----- + The use of this test relies on several assumptions. + + 1. The observations are independent within and among groups. + 2. The observations within each group are normally distributed. + 3. The distributions from which the samples are drawn have the same finite + variance. + + The original formulation of the test was for samples of equal size drawn from + populations assumed to have equal variances [6]_. In case of unequal sample sizes, + the test uses the Tukey-Kramer method [4]_. When equal variances are not assumed + (``equal_var=False``), the test uses the Games-Howell method [7]_. + + References + ---------- + .. [1] NIST/SEMATECH e-Handbook of Statistical Methods, "7.4.7.1. Tukey's + Method." + https://www.itl.nist.gov/div898/handbook/prc/section4/prc471.htm, + 28 November 2020. + .. [2] Abdi, Herve & Williams, Lynne. (2021). "Tukey's Honestly Significant + Difference (HSD) Test." + https://personal.utdallas.edu/~herve/abdi-HSD2010-pretty.pdf + .. [3] "One-Way ANOVA Using SAS PROC ANOVA & PROC GLM." SAS + Tutorials, 2007. + https://www.stattutorials.com/SAS/TUTORIAL-PROC-GLM.htm + .. [4] Kramer, Clyde Young. "Extension of Multiple Range Tests to Group + Means with Unequal Numbers of Replications." Biometrics, vol. 12, + no. 3, 1956, pp. 307-310. https://www.jstor.org/stable/3001469 + .. [5] NIST/SEMATECH e-Handbook of Statistical Methods, "7.4.3.3. + The ANOVA table and tests of hypotheses about means" + https://www.itl.nist.gov/div898/handbook/prc/section4/prc433.htm, + 2 June 2021. + .. [6] Tukey, John W. "Comparing Individual Means in the Analysis of + Variance." Biometrics, vol. 5, no. 2, 1949, pp. 99-114. + https://www.jstor.org/stable/3001913 + .. [7] P. A. Games and J. F. Howell, "Pairwise Multiple Comparison Procedures + with Unequal N's and/or Variances: A Monte Carlo Study," Journal of + Educational Statistics, vol. 1, no. 2, pp. 113-125, Jun. 1976. + :doi:`10.3102/10769986001002113`. + + + Examples + -------- + Here are some data comparing the time to relief of three brands of + headache medicine, reported in minutes. Data adapted from [3]_. + + >>> import numpy as np + >>> from scipy.stats import tukey_hsd + >>> group0 = [24.5, 23.5, 26.4, 27.1, 29.9] + >>> group1 = [28.4, 34.2, 29.5, 32.2, 30.1] + >>> group2 = [26.1, 28.3, 24.3, 26.2, 27.8] + + We would like to see if the means between any of the groups are + significantly different. First, visually examine a box and whisker plot. + + >>> import matplotlib.pyplot as plt + >>> fig, ax = plt.subplots(1, 1) + >>> ax.boxplot([group0, group1, group2]) + >>> ax.set_xticklabels(["group0", "group1", "group2"]) # doctest: +SKIP + >>> ax.set_ylabel("mean") # doctest: +SKIP + >>> plt.show() + + From the box and whisker plot, we can see overlap in the interquartile + ranges group 1 to group 2 and group 3, but we can apply the ``tukey_hsd`` + test to determine if the difference between means is significant. We + set a significance level of .05 to reject the null hypothesis. + + >>> res = tukey_hsd(group0, group1, group2) + >>> print(res) + Pairwise Group Comparisons (95.0% Confidence Interval) + Comparison Statistic p-value Lower CI Upper CI + (0 - 1) -4.600 0.014 -8.249 -0.951 + (0 - 2) -0.260 0.980 -3.909 3.389 + (1 - 0) 4.600 0.014 0.951 8.249 + (1 - 2) 4.340 0.020 0.691 7.989 + (2 - 0) 0.260 0.980 -3.389 3.909 + (2 - 1) -4.340 0.020 -7.989 -0.691 + + The null hypothesis is that each group has the same mean. The p-value for + comparisons between ``group0`` and ``group1`` as well as ``group1`` and + ``group2`` do not exceed .05, so we reject the null hypothesis that they + have the same means. The p-value of the comparison between ``group0`` + and ``group2`` exceeds .05, so we accept the null hypothesis that there + is not a significant difference between their means. + + We can also compute the confidence interval associated with our chosen + confidence level. + + >>> group0 = [24.5, 23.5, 26.4, 27.1, 29.9] + >>> group1 = [28.4, 34.2, 29.5, 32.2, 30.1] + >>> group2 = [26.1, 28.3, 24.3, 26.2, 27.8] + >>> result = tukey_hsd(group0, group1, group2) + >>> conf = res.confidence_interval(confidence_level=.99) + >>> for ((i, j), l) in np.ndenumerate(conf.low): + ... # filter out self comparisons + ... if i != j: + ... h = conf.high[i,j] + ... print(f"({i} - {j}) {l:>6.3f} {h:>6.3f}") + (0 - 1) -9.480 0.280 + (0 - 2) -5.140 4.620 + (1 - 0) -0.280 9.480 + (1 - 2) -0.540 9.220 + (2 - 0) -4.620 5.140 + (2 - 1) -9.220 0.540 + """ + args = _tukey_hsd_iv(samples, equal_var) + ntreatments = len(args) + means = np.asarray([np.mean(arg) for arg in args]) + nsamples_treatments = np.asarray([a.size for a in args]) + nobs = np.sum(nsamples_treatments) + vars_ = np.asarray([np.var(arg, ddof=1) for arg in args]) + + if equal_var: + # determine mean square error [5]. Note that this is sometimes called + # mean square error within. + mse = (np.sum(vars_ * (nsamples_treatments - 1)) / (nobs - ntreatments)) + + # The calculation of the standard error differs when treatments differ in + # size. See ("Unequal sample sizes")[1]. + if np.unique(nsamples_treatments).size == 1: + # all input groups are the same length, so only one value needs to be + # calculated [1]. + normalize = 2 / nsamples_treatments[0] + else: + # to compare groups of differing sizes, we must compute a variance + # value for each individual comparison. Use broadcasting to get the + # resulting matrix. [3], verified against [4] (page 308). + normalize = 1 / nsamples_treatments + 1 / nsamples_treatments[None].T + + # the standard error is used in the computation of the tukey criterion and + # finding the p-values. + stand_err = np.sqrt(normalize * mse / 2) + df = nobs - ntreatments + else: + # `stand_err` is the denominator of the Behrens-Fisher statistic ($v$) + # with a factor of $\sqrt{2}$. Compare [7] p.116 "t-solution rejects H0 if...", + # [7] p. 117 "H0 was rejected", and definition of `t_stat` below. + sj2_nj = vars_ / nsamples_treatments + si2_ni = sj2_nj[:, np.newaxis] + stand_err = np.sqrt(si2_ni + sj2_nj) / 2**0.5 + + # `df` is the Welch degree of freedom $\nu$. + # See [7] p. 116 "and the degrees of freedom, $\nu$, are given by...". + njm1 = nsamples_treatments - 1 + nim1 = njm1[:, np.newaxis] + df = (si2_ni + sj2_nj)**2 / (si2_ni**2 / nim1 + sj2_nj**2 / njm1) + + # the mean difference is the test statistic. + mean_differences = means[None].T - means + + # Calculate the t-statistic to use within the survival function of the + # studentized range to get the p-value. + t_stat = np.abs(mean_differences) / stand_err + + params = t_stat, ntreatments, df + pvalues = distributions.studentized_range.sf(*params) + + return TukeyHSDResult(mean_differences, pvalues, ntreatments, + df, stand_err) diff --git a/v12x/pylib/scipy/stats/_kde.py b/v12x/pylib/scipy/stats/_kde.py new file mode 100644 index 0000000000000000000000000000000000000000..a1be45bfa59f7816235886763761dc3a2107c9cd --- /dev/null +++ b/v12x/pylib/scipy/stats/_kde.py @@ -0,0 +1,765 @@ +#------------------------------------------------------------------------------- +# +# Define classes for (uni/multi)-variate kernel density estimation. +# +# Currently, only Gaussian kernels are implemented. +# +# Written by: Robert Kern +# +# Date: 2004-08-09 +# +# Modified: 2005-02-10 by Robert Kern. +# Contributed to SciPy +# 2005-10-07 by Robert Kern. +# Some fixes to match the new scipy_core +# +# Copyright 2004-2005 by Enthought, Inc. +# +#------------------------------------------------------------------------------- + +from types import GenericAlias + +# SciPy imports. +from scipy import linalg, special +from scipy._lib._util import check_random_state + +from numpy import (asarray, atleast_2d, reshape, zeros, newaxis, exp, pi, + sqrt, ravel, power, atleast_1d, squeeze, sum, transpose, + ones, cov) +import numpy as np + +# Local imports. +from ._stats import gaussian_kernel_estimate, gaussian_kernel_estimate_log +from ._multivariate import multivariate_normal + +__all__ = ['gaussian_kde'] + + +class gaussian_kde: + """Representation of a kernel-density estimate using Gaussian kernels. + + Kernel density estimation is a way to estimate the probability density + function (PDF) of a random variable in a non-parametric way. + `gaussian_kde` works for both uni-variate and multi-variate data. It + includes automatic bandwidth determination. The estimation works best for + a unimodal distribution; bimodal or multi-modal distributions tend to be + oversmoothed. + + Parameters + ---------- + dataset : array_like + Datapoints to estimate from. In case of univariate data this is a 1-D + array, otherwise a 2-D array with shape (# of dims, # of data). + bw_method : str, scalar or callable, optional + The method used to calculate the bandwidth factor. This can be + 'scott', 'silverman', a scalar constant or a callable. If a scalar, + this will be used directly as `factor`. If a callable, it should + take a `gaussian_kde` instance as only parameter and return a scalar. + If None (default), 'scott' is used. See Notes for more details. + weights : array_like, optional + weights of datapoints. This must be the same shape as dataset. + If None (default), the samples are assumed to be equally weighted + + Attributes + ---------- + dataset : ndarray + The dataset with which `gaussian_kde` was initialized. + d : int + Number of dimensions. + n : int + Number of datapoints. + neff : int + Effective number of datapoints. + + .. versionadded:: 1.2.0 + factor : float + The bandwidth factor obtained from `covariance_factor`. + covariance : ndarray + The kernel covariance matrix; this is the data covariance matrix + multiplied by the square of the bandwidth factor, e.g. + ``np.cov(dataset) * factor**2``. + inv_cov : ndarray + The inverse of `covariance`. + + Methods + ------- + evaluate + __call__ + integrate_gaussian + integrate_box_1d + integrate_box + integrate_kde + pdf + logpdf + resample + set_bandwidth + covariance_factor + marginal + + Notes + ----- + Bandwidth selection strongly influences the estimate obtained from the KDE + (much more so than the actual shape of the kernel). Bandwidth selection + can be done by a "rule of thumb", by cross-validation, by "plug-in + methods" or by other means; see [3]_, [4]_ for reviews. `gaussian_kde` + uses a rule of thumb, the default is Scott's Rule. + + Scott's Rule [1]_, implemented as `scotts_factor`, is:: + + n**(-1./(d+4)), + + with ``n`` the number of data points and ``d`` the number of dimensions. + In the case of unequally weighted points, `scotts_factor` becomes:: + + neff**(-1./(d+4)), + + with ``neff`` the effective number of datapoints. + Silverman's suggestion for *multivariate* data [2]_, implemented as + `silverman_factor`, is:: + + (n * (d + 2) / 4.)**(-1. / (d + 4)). + + or in the case of unequally weighted points:: + + (neff * (d + 2) / 4.)**(-1. / (d + 4)). + + Note that this is not the same as "Silverman's rule of thumb" [6]_, which + may be more robust in the univariate case; see documentation of the + ``set_bandwidth`` method for implementing a custom bandwidth rule. + + Good general descriptions of kernel density estimation can be found in [1]_ + and [2]_, the mathematics for this multi-dimensional implementation can be + found in [1]_. + + With a set of weighted samples, the effective number of datapoints ``neff`` + is defined by:: + + neff = sum(weights)^2 / sum(weights^2) + + as detailed in [5]_. + + `gaussian_kde` does not currently support data that lies in a + lower-dimensional subspace of the space in which it is expressed. For such + data, consider performing principal component analysis / dimensionality + reduction and using `gaussian_kde` with the transformed data. + + References + ---------- + .. [1] D.W. Scott, "Multivariate Density Estimation: Theory, Practice, and + Visualization", John Wiley & Sons, New York, Chicester, 1992. + .. [2] B.W. Silverman, "Density Estimation for Statistics and Data + Analysis", Vol. 26, Monographs on Statistics and Applied Probability, + Chapman and Hall, London, 1986. + .. [3] B.A. Turlach, "Bandwidth Selection in Kernel Density Estimation: A + Review", CORE and Institut de Statistique, Vol. 19, pp. 1-33, 1993. + .. [4] D.M. Bashtannyk and R.J. Hyndman, "Bandwidth selection for kernel + conditional density estimation", Computational Statistics & Data + Analysis, Vol. 36, pp. 279-298, 2001. + .. [5] Gray P. G., 1969, Journal of the Royal Statistical Society. + Series A (General), 132, 272 + .. [6] Kernel density estimation. *Wikipedia.* + https://en.wikipedia.org/wiki/Kernel_density_estimation + + Examples + -------- + Generate some random two-dimensional data: + + >>> import numpy as np + >>> from scipy import stats + >>> def measure(n): + ... "Measurement model, return two coupled measurements." + ... m1 = np.random.normal(size=n) + ... m2 = np.random.normal(scale=0.5, size=n) + ... return m1+m2, m1-m2 + + >>> m1, m2 = measure(2000) + >>> xmin = m1.min() + >>> xmax = m1.max() + >>> ymin = m2.min() + >>> ymax = m2.max() + + Perform a kernel density estimate on the data: + + >>> X, Y = np.mgrid[xmin:xmax:100j, ymin:ymax:100j] + >>> positions = np.vstack([X.ravel(), Y.ravel()]) + >>> values = np.vstack([m1, m2]) + >>> kernel = stats.gaussian_kde(values) + >>> Z = np.reshape(kernel(positions).T, X.shape) + + Plot the results: + + >>> import matplotlib.pyplot as plt + >>> fig, ax = plt.subplots() + >>> ax.imshow(np.rot90(Z), cmap=plt.cm.gist_earth_r, + ... extent=[xmin, xmax, ymin, ymax]) + >>> ax.plot(m1, m2, 'k.', markersize=2) + >>> ax.set_xlim([xmin, xmax]) + >>> ax.set_ylim([ymin, ymax]) + >>> plt.show() + + Compare against manual KDE at a point: + + >>> point = [1, 2] + >>> mean = values.T + >>> cov = kernel.factor**2 * np.cov(values) + >>> X = stats.multivariate_normal(cov=cov) + >>> res = kernel.pdf(point) + >>> ref = X.pdf(point - mean).sum() / len(mean) + >>> np.allclose(res, ref) + True + """ + + # generic type compatibility with scipy-stubs + __class_getitem__: classmethod = classmethod(GenericAlias) + + def __init__(self, dataset, bw_method=None, weights=None): + self.dataset = atleast_2d(asarray(dataset)) + if not self.dataset.size > 1: + raise ValueError("`dataset` input should have multiple elements.") + + self.d, self.n = self.dataset.shape + + if weights is not None: + self._weights = atleast_1d(weights).astype(float) + self._weights /= sum(self._weights) + if self.weights.ndim != 1: + raise ValueError("`weights` input should be one-dimensional.") + if len(self._weights) != self.n: + raise ValueError("`weights` input should be of length n") + self._neff = 1/np.vecdot(self._weights, self._weights) + + # This can be converted to a warning once gh-10205 is resolved + if self.d > self.n: + msg = ("Number of dimensions is greater than number of samples. " + "This results in a singular data covariance matrix, which " + "cannot be treated using the algorithms implemented in " + "`gaussian_kde`. Note that `gaussian_kde` interprets each " + "*column* of `dataset` to be a point; consider transposing " + "the input to `dataset`.") + raise ValueError(msg) + + try: + self.set_bandwidth(bw_method=bw_method) + except linalg.LinAlgError as e: + msg = ("The data appears to lie in a lower-dimensional subspace " + "of the space in which it is expressed. This has resulted " + "in a singular data covariance matrix, which cannot be " + "treated using the algorithms implemented in " + "`gaussian_kde`. Consider performing principal component " + "analysis / dimensionality reduction and using " + "`gaussian_kde` with the transformed data.") + raise linalg.LinAlgError(msg) from e + + def evaluate(self, points): + """Evaluate the estimated pdf on a set of points. + + Parameters + ---------- + points : (# of dimensions, # of points)-array + Alternatively, a (# of dimensions,) vector can be passed in and + treated as a single point. + + Returns + ------- + values : (# of points,)-array + The values at each point. + + Raises + ------ + ValueError : if the dimensionality of the input points is different than + the dimensionality of the KDE. + + """ + points = atleast_2d(asarray(points)) + + d, m = points.shape + if d != self.d: + if d == 1 and m == self.d: + # points was passed in as a row vector + points = reshape(points, (self.d, 1)) + m = 1 + else: + msg = (f"points have dimension {d}, " + f"dataset has dimension {self.d}") + raise ValueError(msg) + + output_dtype, spec = _get_output_dtype(self.covariance, points) + result = gaussian_kernel_estimate[spec]( + self.dataset.T, self.weights[:, None], + points.T, self.cho_cov, output_dtype) + + return result[:, 0] + + __call__ = evaluate + + def integrate_gaussian(self, mean, cov): + """ + Multiply estimated density by a multivariate Gaussian and integrate + over the whole space. + + Parameters + ---------- + mean : aray_like + A 1-D array, specifying the mean of the Gaussian. + cov : array_like + A 2-D array, specifying the covariance matrix of the Gaussian. + + Returns + ------- + result : scalar + The value of the integral. + + Raises + ------ + ValueError + If the mean or covariance of the input Gaussian differs from + the KDE's dimensionality. + + """ + mean = atleast_1d(squeeze(mean)) + cov = atleast_2d(cov) + + if mean.shape != (self.d,): + raise ValueError(f"mean does not have dimension {self.d}") + if cov.shape != (self.d, self.d): + raise ValueError(f"covariance does not have dimension {self.d}") + + # make mean a column vector + mean = mean[:, newaxis] + + sum_cov = self.covariance + cov + + # This will raise LinAlgError if the new cov matrix is not s.p.d + # cho_factor returns (ndarray, bool) where bool is a flag for whether + # or not ndarray is upper or lower triangular + sum_cov_chol = linalg.cho_factor(sum_cov) + + diff = self.dataset - mean + tdiff = linalg.cho_solve(sum_cov_chol, diff) + + sqrt_det = np.prod(np.diagonal(sum_cov_chol[0])) + norm_const = power(2 * pi, sum_cov.shape[0] / 2.0) * sqrt_det + + energies = np.vecdot(diff, tdiff, axis=0) / 2.0 + result = np.vecdot(exp(-energies), self.weights, axis=0) / norm_const + + return result + + def integrate_box_1d(self, low, high): + """ + Computes the integral of a 1D pdf between two bounds. + + Parameters + ---------- + low : scalar + Lower bound of integration. + high : scalar + Upper bound of integration. + + Returns + ------- + value : scalar + The result of the integral. + + Raises + ------ + ValueError + If the KDE is over more than one dimension. + + """ + if self.d != 1: + raise ValueError("integrate_box_1d() only handles 1D pdfs") + + stdev = ravel(sqrt(self.covariance))[0] + + normalized_low = ravel((low - self.dataset) / stdev) + normalized_high = ravel((high - self.dataset) / stdev) + + delta = special.ndtr(normalized_high) - special.ndtr(normalized_low) + value = np.vecdot(self.weights, delta) + return value + + def integrate_box(self, low_bounds, high_bounds, maxpts=None, *, rng=None): + """Computes the integral of a pdf over a rectangular interval. + + Parameters + ---------- + low_bounds : array_like + A 1-D array containing the lower bounds of integration. + high_bounds : array_like + A 1-D array containing the upper bounds of integration. + maxpts : int, optional + The maximum number of points to use for integration. + rng : `numpy.random.Generator`, optional + Pseudorandom number generator state. When `rng` is None, a new + generator is created using entropy from the operating system. Types + other than `numpy.random.Generator` are passed to + `numpy.random.default_rng` to instantiate a ``Generator``. + + Returns + ------- + value : scalar + The result of the integral. + + """ + low, high = low_bounds - self.dataset.T, high_bounds - self.dataset.T + values = multivariate_normal.cdf( + high, lower_limit=low, cov=self.covariance, maxpts=maxpts, + rng=rng + ) + return np.vecdot(values, self.weights, axis=-1) + + def integrate_kde(self, other): + """ + Computes the integral of the product of this kernel density estimate + with another. + + Parameters + ---------- + other : gaussian_kde instance + The other kde. + + Returns + ------- + value : scalar + The result of the integral. + + Raises + ------ + ValueError + If the KDEs have different dimensionality. + + """ + if other.d != self.d: + raise ValueError("KDEs are not the same dimensionality") + + # we want to iterate over the smallest number of points + if other.n < self.n: + small = other + large = self + else: + small = self + large = other + + sum_cov = small.covariance + large.covariance + sum_cov_chol = linalg.cho_factor(sum_cov) + result = 0.0 + for i in range(small.n): + mean = small.dataset[:, i, newaxis] + diff = large.dataset - mean + tdiff = linalg.cho_solve(sum_cov_chol, diff) + + energies = np.vecdot(diff, tdiff, axis=0) / 2.0 + result += np.vecdot(exp(-energies), large.weights, axis=0)*small.weights[i] + + sqrt_det = np.prod(np.diagonal(sum_cov_chol[0])) + norm_const = power(2 * pi, sum_cov.shape[0] / 2.0) * sqrt_det + + result /= norm_const + + return result + + def resample(self, size=None, seed=None): + """Randomly sample a dataset from the estimated pdf. + + Parameters + ---------- + size : int, optional + The number of samples to draw. If not provided, then the size is + the same as the effective number of samples in the underlying + dataset. + seed : {None, int, `numpy.random.Generator`, `numpy.random.RandomState`}, optional + If `seed` is None (or `np.random`), the `numpy.random.RandomState` + singleton is used. + If `seed` is an int, a new ``RandomState`` instance is used, + seeded with `seed`. + If `seed` is already a ``Generator`` or ``RandomState`` instance then + that instance is used. + + Returns + ------- + resample : (self.d, `size`) ndarray + The sampled dataset. + + """ # numpy/numpydoc#87 # noqa: E501 + if size is None: + size = int(self.neff) + + random_state = check_random_state(seed) + norm = transpose(random_state.multivariate_normal( + zeros((self.d,), float), self.covariance, size=size + )) + indices = random_state.choice(self.n, size=size, p=self.weights) + means = self.dataset[:, indices] + + return means + norm + + def scotts_factor(self): + """Compute Scott's factor. + + Returns + ------- + s : float + Scott's factor. + """ + return power(self.neff, -1./(self.d+4)) + + def silverman_factor(self): + """Compute the Silverman factor. + + Returns + ------- + s : float + The silverman factor. + """ + return power(self.neff*(self.d+2.0)/4.0, -1./(self.d+4)) + + # Default method to calculate bandwidth, can be overwritten by subclass + covariance_factor = scotts_factor + covariance_factor.__doc__ = """Computes the bandwidth factor `factor`. + The default is `scotts_factor`. A subclass can overwrite this + method to provide a different method, or set it through a call to + `set_bandwidth`. + + Returns + ------- + factor : float + The bandwidth factor. + """ + + def set_bandwidth(self, bw_method=None): + """Compute the bandwidth factor with given method. + + The new bandwidth calculated after a call to `set_bandwidth` is used + for subsequent evaluations of the estimated density. + + Parameters + ---------- + bw_method : str, scalar or callable, optional + The method used to calculate the bandwidth factor. This can be + 'scott', 'silverman', a scalar constant or a callable. If a + scalar, this will be used directly as `factor`. If a callable, + it should take a `gaussian_kde` instance as only parameter and + return a scalar. If None (default), nothing happens; the current + `covariance_factor` method is kept. + + Notes + ----- + .. versionadded:: 0.11 + + Examples + -------- + >>> import numpy as np + >>> import scipy.stats as stats + >>> x1 = np.array([-7, -5, 1, 4, 5.]) + >>> kde = stats.gaussian_kde(x1) + >>> xs = np.linspace(-10, 10, num=50) + >>> y1 = kde(xs) + >>> kde.set_bandwidth(bw_method='silverman') + >>> y2 = kde(xs) + >>> kde.set_bandwidth(bw_method=kde.factor / 3.) + >>> y3 = kde(xs) + + >>> import matplotlib.pyplot as plt + >>> fig, ax = plt.subplots() + >>> ax.plot(x1, np.full(x1.shape, 1 / (4. * x1.size)), 'bo', + ... label='Data points (rescaled)') + >>> ax.plot(xs, y1, label='Scott (default)') + >>> ax.plot(xs, y2, label='Silverman') + >>> ax.plot(xs, y3, label='Const (1/3 * Silverman)') + >>> ax.legend() + >>> plt.show() + + """ + if bw_method is None: + pass + elif bw_method == 'scott': + self.covariance_factor = self.scotts_factor + elif bw_method == 'silverman': + self.covariance_factor = self.silverman_factor + elif np.isscalar(bw_method) and not isinstance(bw_method, str): + self._bw_method = 'use constant' + self.covariance_factor = lambda: bw_method + elif callable(bw_method): + self._bw_method = bw_method + self.covariance_factor = lambda: self._bw_method(self) + else: + msg = "`bw_method` should be 'scott', 'silverman', a scalar " \ + "or a callable." + raise ValueError(msg) + + self._compute_covariance() + + def _compute_covariance(self): + """Computes the covariance matrix for each Gaussian kernel using + covariance_factor(). + """ + self.factor = self.covariance_factor() + # Cache covariance and Cholesky decomp of covariance + if not hasattr(self, '_data_cho_cov'): + self._data_covariance = atleast_2d(cov(self.dataset, rowvar=1, + bias=False, + aweights=self.weights)) + self._data_cho_cov = linalg.cholesky(self._data_covariance, + lower=True) + + self.covariance = self._data_covariance * self.factor**2 + self.cho_cov = (self._data_cho_cov * self.factor).astype(np.float64) + self.log_det = 2*np.log(np.diag(self.cho_cov + * np.sqrt(2*pi))).sum() + + @property + def inv_cov(self): + # Re-compute from scratch each time because I'm not sure how this is + # used in the wild. (Perhaps users change the `dataset`, since it's + # not a private attribute?) `_compute_covariance` used to recalculate + # all these, so we'll recalculate everything now that this is a + # a property. + self.factor = self.covariance_factor() + self._data_covariance = atleast_2d(cov(self.dataset, rowvar=1, + bias=False, aweights=self.weights)) + return linalg.inv(self._data_covariance) / self.factor**2 + + def pdf(self, x): + """ + Evaluate the estimated pdf on a provided set of points. + + Parameters + ---------- + x : array_like + Points at which to evaluate the pdf. + + Returns + ------- + pdf : ndarray + The pdf evaluated at `x`. + + Notes + ----- + This is an alias for `gaussian_kde.evaluate`. See the ``evaluate`` + docstring for more details. + + """ + return self.evaluate(x) + + def logpdf(self, x): + """ + Evaluate the log of the estimated pdf on a provided set of points. + + Parameters + ---------- + x : array_like + Points at which to evaluate the log-pdf. + + Returns + ------- + logpdf : ndarray + The log-pdf evaluated at `x`. + """ + points = atleast_2d(x) + + d, m = points.shape + if d != self.d: + if d == 1 and m == self.d: + # points was passed in as a row vector + points = reshape(points, (self.d, 1)) + m = 1 + else: + msg = (f"points have dimension {d}, " + f"dataset has dimension {self.d}") + raise ValueError(msg) + + output_dtype, spec = _get_output_dtype(self.covariance, points) + result = gaussian_kernel_estimate_log[spec]( + self.dataset.T, self.weights[:, None], + points.T, self.cho_cov, output_dtype) + + return result[:, 0] + + def marginal(self, dimensions): + """Return a marginal KDE distribution. + + Parameters + ---------- + dimensions : int or 1-d array_like + The dimensions of the multivariate distribution corresponding + with the marginal variables, that is, the indices of the dimensions + that are being retained. The other dimensions are marginalized out. + + Returns + ------- + marginal_kde : gaussian_kde + An object representing the marginal distribution. + + Notes + ----- + .. versionadded:: 1.10.0 + + """ + + dims = np.atleast_1d(dimensions) + + if not np.issubdtype(dims.dtype, np.integer): + msg = ("Elements of `dimensions` must be integers - the indices " + "of the marginal variables being retained.") + raise ValueError(msg) + + n = len(self.dataset) # number of dimensions + original_dims = dims.copy() + + dims[dims < 0] = n + dims[dims < 0] + + if len(np.unique(dims)) != len(dims): + msg = ("All elements of `dimensions` must be unique.") + raise ValueError(msg) + + i_invalid = (dims < 0) | (dims >= n) + if np.any(i_invalid): + msg = (f"Dimensions {original_dims[i_invalid]} are invalid " + f"for a distribution in {n} dimensions.") + raise ValueError(msg) + + dataset = self.dataset[dims] + weights = self.weights + + return gaussian_kde(dataset, bw_method=self.covariance_factor(), + weights=weights) + + @property + def weights(self): + try: + return self._weights + except AttributeError: + self._weights = ones(self.n)/self.n + return self._weights + + @property + def neff(self): + try: + return self._neff + except AttributeError: + self._neff = 1/np.vecdot(self.weights, self.weights) + return self._neff + + +def _get_output_dtype(covariance, points): + """ + Calculates the output dtype and the "spec" (=C type name). + + This was necessary in order to deal with the fused types in the Cython + routine `gaussian_kernel_estimate`. See gh-10824 for details. + """ + output_dtype = np.common_type(covariance, points) + itemsize = np.dtype(output_dtype).itemsize + if itemsize == 4: + spec = 'float' + elif itemsize == 8: + spec = 'double' + elif itemsize in (12, 16): + spec = 'long double' + else: + raise ValueError( + f"{output_dtype} has unexpected item size: {itemsize}" + ) + + return output_dtype, spec diff --git a/v12x/pylib/scipy/stats/_ksstats.py b/v12x/pylib/scipy/stats/_ksstats.py new file mode 100644 index 0000000000000000000000000000000000000000..f6884bb834bbb060fdec5ed49869517e5b7c941b --- /dev/null +++ b/v12x/pylib/scipy/stats/_ksstats.py @@ -0,0 +1,600 @@ +# Compute the two-sided one-sample Kolmogorov-Smirnov Prob(Dn <= d) where: +# D_n = sup_x{|F_n(x) - F(x)|}, +# F_n(x) is the empirical CDF for a sample of size n {x_i: i=1,...,n}, +# F(x) is the CDF of a probability distribution. +# +# Exact methods: +# Prob(D_n >= d) can be computed via a matrix algorithm of Durbin[1] +# or a recursion algorithm due to Pomeranz[2]. +# Marsaglia, Tsang & Wang[3] gave a computation-efficient way to perform +# the Durbin algorithm. +# D_n >= d <==> D_n+ >= d or D_n- >= d (the one-sided K-S statistics), hence +# Prob(D_n >= d) = 2*Prob(D_n+ >= d) - Prob(D_n+ >= d and D_n- >= d). +# For d > 0.5, the latter intersection probability is 0. +# +# Approximate methods: +# For d close to 0.5, ignoring that intersection term may still give a +# reasonable approximation. +# Li-Chien[4] and Korolyuk[5] gave an asymptotic formula extending +# Kolmogorov's initial asymptotic, suitable for large d. (See +# scipy.special.kolmogorov for that asymptotic) +# Pelz-Good[6] used the functional equation for Jacobi theta functions to +# transform the Li-Chien/Korolyuk formula produce a computational formula +# suitable for small d. +# +# Simard and L'Ecuyer[7] provided an algorithm to decide when to use each of +# the above approaches and it is that which is used here. +# +# Other approaches: +# Carvalho[8] optimizes Durbin's matrix algorithm for large values of d. +# Moscovich and Nadler[9] use FFTs to compute the convolutions. + +# References: +# [1] Durbin J (1968). +# "The Probability that the Sample Distribution Function Lies Between Two +# Parallel Straight Lines." +# Annals of Mathematical Statistics, 39, 398-411. +# [2] Pomeranz J (1974). +# "Exact Cumulative Distribution of the Kolmogorov-Smirnov Statistic for +# Small Samples (Algorithm 487)." +# Communications of the ACM, 17(12), 703-704. +# [3] Marsaglia G, Tsang WW, Wang J (2003). +# "Evaluating Kolmogorov's Distribution." +# Journal of Statistical Software, 8(18), 1-4. +# [4] LI-CHIEN, C. (1956). +# "On the exact distribution of the statistics of A. N. Kolmogorov and +# their asymptotic expansion." +# Acta Matematica Sinica, 6, 55-81. +# [5] KOROLYUK, V. S. (1960). +# "Asymptotic analysis of the distribution of the maximum deviation in +# the Bernoulli scheme." +# Theor. Probability Appl., 4, 339-366. +# [6] Pelz W, Good IJ (1976). +# "Approximating the Lower Tail-areas of the Kolmogorov-Smirnov One-sample +# Statistic." +# Journal of the Royal Statistical Society, Series B, 38(2), 152-156. +# [7] Simard, R., L'Ecuyer, P. (2011) +# "Computing the Two-Sided Kolmogorov-Smirnov Distribution", +# Journal of Statistical Software, Vol 39, 11, 1-18. +# [8] Carvalho, Luis (2015) +# "An Improved Evaluation of Kolmogorov's Distribution" +# Journal of Statistical Software, Code Snippets; Vol 65(3), 1-8. +# [9] Amit Moscovich, Boaz Nadler (2017) +# "Fast calculation of boundary crossing probabilities for Poisson +# processes", +# Statistics & Probability Letters, Vol 123, 177-182. + + +import numpy as np +import scipy.special +import scipy.special._ufuncs as scu +from scipy.stats._finite_differences import _derivative + +_E128 = 128 +_EP128 = np.ldexp(np.longdouble(1), _E128) +_EM128 = np.ldexp(np.longdouble(1), -_E128) + +_SQRT2PI = np.sqrt(2 * np.pi) +_LOG_2PI = np.log(2 * np.pi) +_MIN_LOG = -708 +_SQRT3 = np.sqrt(3) +_PI_SQUARED = np.pi ** 2 +_PI_FOUR = np.pi ** 4 +_PI_SIX = np.pi ** 6 + +# [Lifted from _loggamma.pxd.] If B_m are the Bernoulli numbers, +# then Stirling coeffs are B_{2j}/(2j)/(2j-1) for j=8,...1. +_STIRLING_COEFFS = [-2.955065359477124183e-2, 6.4102564102564102564e-3, + -1.9175269175269175269e-3, 8.4175084175084175084e-4, + -5.952380952380952381e-4, 7.9365079365079365079e-4, + -2.7777777777777777778e-3, 8.3333333333333333333e-2] + + +def _log_nfactorial_div_n_pow_n(n): + # Computes n! / n**n + # = (n-1)! / n**(n-1) + # Uses Stirling's approximation, but removes n*log(n) up-front to + # avoid subtractive cancellation. + # = log(n)/2 - n + log(sqrt(2pi)) + sum B_{2j}/(2j)/(2j-1)/n**(2j-1) + rn = 1.0/n + return np.log(n)/2 - n + _LOG_2PI/2 + rn * np.polyval(_STIRLING_COEFFS, rn/n) + + +def _clip_prob(p): + """clips a probability to range 0<=p<=1.""" + return np.clip(p, 0.0, 1.0) + + +def _select_and_clip_prob(cdfprob, sfprob, cdf=True): + """Selects either the CDF or SF, and then clips to range 0<=p<=1.""" + p = np.where(cdf, cdfprob, sfprob) + return _clip_prob(p) + + +def _kolmogn_DMTW(n, d, cdf=True): + r"""Computes the Kolmogorov CDF: Pr(D_n <= d) using the MTW approach to + the Durbin matrix algorithm. + + Durbin (1968); Marsaglia, Tsang, Wang (2003). [1], [3]. + """ + # Write d = (k-h)/n, where k is positive integer and 0 <= h < 1 + # Generate initial matrix H of size m*m where m=(2k-1) + # Compute k-th row of (n!/n^n) * H^n, scaling intermediate results. + # Requires memory O(m^2) and computation O(m^2 log(n)). + # Most suitable for small m. + + if d >= 1.0: + return _select_and_clip_prob(1.0, 0.0, cdf) + nd = n * d + if nd <= 0.5: + return _select_and_clip_prob(0.0, 1.0, cdf) + k = int(np.ceil(nd)) + h = k - nd + m = 2 * k - 1 + + H = np.zeros([m, m]) + + # Initialize: v is first column (and last row) of H + # v[j] = (1-h^(j+1)/(j+1)! (except for v[-1]) + # w[j] = 1/(j)! + # q = k-th row of H (actually i!/n^i*H^i) + intm = np.arange(1, m + 1) + v = 1.0 - h ** intm + w = np.empty(m) + fac = 1.0 + for j in intm: + w[j - 1] = fac + fac /= j # This might underflow. Isn't a problem. + v[j - 1] *= fac + tt = max(2 * h - 1.0, 0)**m - 2*h**m + v[-1] = (1.0 + tt) * fac + + for i in range(1, m): + H[i - 1:, i] = w[:m - i + 1] + H[:, 0] = v + H[-1, :] = np.flip(v, axis=0) + + Hpwr = np.eye(np.shape(H)[0]) # Holds intermediate powers of H + nn = n + expnt = 0 # Scaling of Hpwr + Hexpnt = 0 # Scaling of H + while nn > 0: + if nn % 2: + Hpwr = np.matmul(Hpwr, H) + expnt += Hexpnt + H = np.matmul(H, H) + Hexpnt *= 2 + # Scale as needed. + if np.abs(H[k - 1, k - 1]) > _EP128: + H /= _EP128 + Hexpnt += _E128 + nn = nn // 2 + + p = Hpwr[k - 1, k - 1] + + # Multiply by n!/n^n + for i in range(1, n + 1): + p = i * p / n + if np.abs(p) < _EM128: + p *= _EP128 + expnt -= _E128 + + # unscale + if expnt != 0: + p = np.ldexp(p, expnt) + + return _select_and_clip_prob(p, 1.0-p, cdf) + + +def _pomeranz_compute_j1j2(i, n, ll, ceilf, roundf): + """Compute the endpoints of the interval for row i.""" + if i == 0: + j1, j2 = -ll - ceilf - 1, ll + ceilf - 1 + else: + # i + 1 = 2*ip1div2 + ip1mod2 + ip1div2, ip1mod2 = divmod(i + 1, 2) + if ip1mod2 == 0: # i is odd + if ip1div2 == n + 1: + j1, j2 = n - ll - ceilf - 1, n + ll + ceilf - 1 + else: + j1, j2 = ip1div2 - 1 - ll - roundf - 1, ip1div2 + ll - 1 + ceilf - 1 + else: + j1, j2 = ip1div2 - 1 - ll - 1, ip1div2 + ll + roundf - 1 + + return max(j1 + 2, 0), min(j2, n) + + +def _kolmogn_Pomeranz(n, x, cdf=True): + r"""Computes Pr(D_n <= d) using the Pomeranz recursion algorithm. + + Pomeranz (1974) [2] + """ + + # V is n*(2n+2) matrix. + # Each row is convolution of the previous row and probabilities from a + # Poisson distribution. + # Desired CDF probability is n! V[n-1, 2n+1] (final entry in final row). + # Only two rows are needed at any given stage: + # - Call them V0 and V1. + # - Swap each iteration + # Only a few (contiguous) entries in each row can be non-zero. + # - Keep track of start and end (j1 and j2 below) + # - V0s and V1s track the start in the two rows + # Scale intermediate results as needed. + # Only a few different Poisson distributions can occur + t = n * x + ll = int(np.floor(t)) + f = 1.0 * (t - ll) # fractional part of t + g = min(f, 1.0 - f) + ceilf = (1 if f > 0 else 0) + roundf = (1 if f > 0.5 else 0) + npwrs = 2 * (ll + 1) # Maximum number of powers needed in convolutions + gpower = np.empty(npwrs) # gpower = (g/n)^m/m! + twogpower = np.empty(npwrs) # twogpower = (2g/n)^m/m! + onem2gpower = np.empty(npwrs) # onem2gpower = ((1-2g)/n)^m/m! + # gpower etc are *almost* Poisson probs, just missing normalizing factor. + + gpower[0] = 1.0 + twogpower[0] = 1.0 + onem2gpower[0] = 1.0 + expnt = 0 + g_over_n, two_g_over_n, one_minus_two_g_over_n = g/n, 2*g/n, (1 - 2*g)/n + for m in range(1, npwrs): + gpower[m] = gpower[m - 1] * g_over_n / m + twogpower[m] = twogpower[m - 1] * two_g_over_n / m + onem2gpower[m] = onem2gpower[m - 1] * one_minus_two_g_over_n / m + + V0 = np.zeros([npwrs]) + V1 = np.zeros([npwrs]) + V1[0] = 1 # first row + V0s, V1s = 0, 0 # start indices of the two rows + + j1, j2 = _pomeranz_compute_j1j2(0, n, ll, ceilf, roundf) + for i in range(1, 2 * n + 2): + # Preserve j1, V1, V1s, V0s from last iteration + k1 = j1 + V0, V1 = V1, V0 + V0s, V1s = V1s, V0s + V1.fill(0.0) + j1, j2 = _pomeranz_compute_j1j2(i, n, ll, ceilf, roundf) + if i == 1 or i == 2 * n + 1: + pwrs = gpower + else: + pwrs = (twogpower if i % 2 else onem2gpower) + ln2 = j2 - k1 + 1 + if ln2 > 0: + conv = np.convolve(V0[k1 - V0s:k1 - V0s + ln2], pwrs[:ln2]) + conv_start = j1 - k1 # First index to use from conv + conv_len = j2 - j1 + 1 # Number of entries to use from conv + V1[:conv_len] = conv[conv_start:conv_start + conv_len] + # Scale to avoid underflow. + if 0 < np.max(V1) < _EM128: + V1 *= _EP128 + expnt -= _E128 + V1s = V0s + j1 - k1 + + # multiply by n! + ans = V1[n - V1s] + for m in range(1, n + 1): + if np.abs(ans) > _EP128: + ans *= _EM128 + expnt += _E128 + ans *= m + + # Undo any intermediate scaling + if expnt != 0: + ans = np.ldexp(ans, expnt) + ans = _select_and_clip_prob(ans, 1.0 - ans, cdf) + return ans + + +def _kolmogn_PelzGood(n, x, cdf=True): + """Computes the Pelz-Good approximation to Prob(Dn <= x) with 0<=x<=1. + + Start with Li-Chien, Korolyuk approximation: + Prob(Dn <= x) ~ K0(z) + K1(z)/sqrt(n) + K2(z)/n + K3(z)/n**1.5 + where z = x*sqrt(n). + Transform each K_(z) using Jacobi theta functions into a form suitable + for small z. + Pelz-Good (1976). [6] + """ + if x <= 0.0: + return _select_and_clip_prob(0.0, 1.0, cdf=cdf) + if x >= 1.0: + return _select_and_clip_prob(1.0, 0.0, cdf=cdf) + + z = np.sqrt(n) * x + zsquared, zthree, zfour, zsix = z**2, z**3, z**4, z**6 + + qlog = -_PI_SQUARED / 8 / zsquared + if qlog < _MIN_LOG: # z ~ 0.041743441416853426 + return _select_and_clip_prob(0.0, 1.0, cdf=cdf) + + q = np.exp(qlog) + + # Coefficients of terms in the sums for K1, K2 and K3 + k1a = -zsquared + k1b = _PI_SQUARED / 4 + + k2a = 6 * zsix + 2 * zfour + k2b = (2 * zfour - 5 * zsquared) * _PI_SQUARED / 4 + k2c = _PI_FOUR * (1 - 2 * zsquared) / 16 + + k3d = _PI_SIX * (5 - 30 * zsquared) / 64 + k3c = _PI_FOUR * (-60 * zsquared + 212 * zfour) / 16 + k3b = _PI_SQUARED * (135 * zfour - 96 * zsix) / 4 + k3a = -30 * zsix - 90 * z**8 + + K0to3 = np.zeros(4) + # Use a Horner scheme to evaluate sum c_i q^(i^2) + # Reduces to a sum over odd integers. + maxk = int(np.ceil(16 * z / np.pi)) + for k in range(maxk, 0, -1): + m = 2 * k - 1 + msquared, mfour, msix = m**2, m**4, m**6 + qpower = np.power(q, 8 * k) + coeffs = np.array([1.0, + k1a + k1b*msquared, + k2a + k2b*msquared + k2c*mfour, + k3a + k3b*msquared + k3c*mfour + k3d*msix]) + K0to3 *= qpower + K0to3 += coeffs + K0to3 *= q + K0to3 *= _SQRT2PI + # z**10 > 0 as z > 0.04 + K0to3 /= np.array([z, 6 * zfour, 72 * z**7, 6480 * z**10]) + + # Now do the other sum over the other terms, all integers k + # K_2: (pi^2 k^2) q^(k^2), + # K_3: (3pi^2 k^2 z^2 - pi^4 k^4)*q^(k^2) + # Don't expect much subtractive cancellation so use direct calculation + q = np.exp(-_PI_SQUARED / 2 / zsquared) + ks = np.arange(maxk, 0, -1) + ksquared = ks ** 2 + sqrt3z = _SQRT3 * z + kspi = np.pi * ks + qpwers = q ** ksquared + k2extra = np.sum(ksquared * qpwers) + k2extra *= _PI_SQUARED * _SQRT2PI/(-36 * zthree) + K0to3[2] += k2extra + k3extra = np.sum((sqrt3z + kspi) * (sqrt3z - kspi) * ksquared * qpwers) + k3extra *= _PI_SQUARED * _SQRT2PI/(216 * zsix) + K0to3[3] += k3extra + powers_of_n = np.power(n * 1.0, np.arange(len(K0to3)) / 2.0) + K0to3 /= powers_of_n + + if not cdf: + K0to3 *= -1 + K0to3[0] += 1 + + Ksum = sum(K0to3) + return Ksum + + +def _kolmogn(n, x, cdf=True): + """Computes the CDF(or SF) for the two-sided Kolmogorov-Smirnov statistic. + + x must be of type float, n of type integer. + + Simard & L'Ecuyer (2011) [7]. + """ + if np.isnan(n): + return n # Keep the same type of nan + if int(n) != n or n <= 0: + return np.nan + if x >= 1.0: + return _select_and_clip_prob(1.0, 0.0, cdf=cdf) + if x <= 0.0: + return _select_and_clip_prob(0.0, 1.0, cdf=cdf) + t = n * x + if t <= 1.0: # Ruben-Gambino: 1/2n <= x <= 1/n + if t <= 0.5: + return _select_and_clip_prob(0.0, 1.0, cdf=cdf) + if n <= 140: + prob = np.prod(np.arange(1, n+1) * (1.0/n) * (2*t - 1)) + else: + prob = np.exp(_log_nfactorial_div_n_pow_n(n) + n * np.log(2*t-1)) + return _select_and_clip_prob(prob, 1.0 - prob, cdf=cdf) + if t >= n - 1: # Ruben-Gambino + prob = 2 * (1.0 - x)**n + return _select_and_clip_prob(1 - prob, prob, cdf=cdf) + if x >= 0.5: # Exact: 2 * smirnov + prob = 2 * scipy.special.smirnov(n, x) + return _select_and_clip_prob(1.0 - prob, prob, cdf=cdf) + + nxsquared = t * x + if n <= 140: + if nxsquared <= 0.754693: + prob = _kolmogn_DMTW(n, x, cdf=True) + return _select_and_clip_prob(prob, 1.0 - prob, cdf=cdf) + if nxsquared <= 4: + prob = _kolmogn_Pomeranz(n, x, cdf=True) + return _select_and_clip_prob(prob, 1.0 - prob, cdf=cdf) + # Now use Miller approximation of 2*smirnov + prob = 2 * scipy.special.smirnov(n, x) + return _select_and_clip_prob(1.0 - prob, prob, cdf=cdf) + + # Split CDF and SF as they have different cutoffs on nxsquared. + if not cdf: + if nxsquared >= 370.0: + return 0.0 + if nxsquared >= 2.2: + prob = 2 * scipy.special.smirnov(n, x) + return _clip_prob(prob) + # Fall through and compute the SF as 1.0-CDF + if nxsquared >= 18.0: + cdfprob = 1.0 + elif n <= 100000 and n * x**1.5 <= 1.4: + cdfprob = _kolmogn_DMTW(n, x, cdf=True) + else: + cdfprob = _kolmogn_PelzGood(n, x, cdf=True) + return _select_and_clip_prob(cdfprob, 1.0 - cdfprob, cdf=cdf) + + +def _kolmogn_p(n, x): + """Computes the PDF for the two-sided Kolmogorov-Smirnov statistic. + + x must be of type float, n of type integer. + """ + if np.isnan(n): + return n # Keep the same type of nan + if int(n) != n or n <= 0: + return np.nan + if x >= 1.0 or x <= 0: + return 0 + t = n * x + if t <= 1.0: + # Ruben-Gambino: n!/n^n * (2t-1)^n -> 2 n!/n^n * n^2 * (2t-1)^(n-1) + if t <= 0.5: + return 0.0 + if n <= 140: + prd = np.prod(np.arange(1, n) * (1.0 / n) * (2 * t - 1)) + else: + prd = np.exp(_log_nfactorial_div_n_pow_n(n) + (n-1) * np.log(2 * t - 1)) + return prd * 2 * n**2 + if t >= n - 1: + # Ruben-Gambino : 1-2(1-x)**n -> 2n*(1-x)**(n-1) + return 2 * (1.0 - x) ** (n-1) * n + if x >= 0.5: + return 2 * scipy.stats.ksone.pdf(x, n) + + # Just take a small delta. + # Ideally x +/- delta would stay within [i/n, (i+1)/n] for some integer a. + # as the CDF is a piecewise degree n polynomial. + # It has knots at 1/n, 2/n, ... (n-1)/n + # and is not a C-infinity function at the knots + delta = x / 2.0**16 + delta = min(delta, x - 1.0/n) + delta = min(delta, 0.5 - x) + + def _kk(_x): + return kolmogn(n, _x) + + return _derivative(_kk, x, dx=delta, order=5) + + +def _kolmogni(n, p, q): + """Computes the PPF/ISF of kolmogn. + + n of type integer, n>= 1 + p is the CDF, q the SF, p+q=1 + """ + if np.isnan(n): + return n # Keep the same type of nan + if int(n) != n or n <= 0: + return np.nan + if p <= 0: + return 1.0/n + if q <= 0: + return 1.0 + delta = np.exp((np.log(p) - scipy.special.loggamma(n+1))/n) + if delta <= 1.0/n: + return (delta + 1.0 / n) / 2 + x = -np.expm1(np.log(q/2.0)/n) + if x >= 1 - 1.0/n: + return x + x1 = scu._kolmogci(p)/np.sqrt(n) + x1 = min(x1, 1.0 - 1.0/n) + + def _f(x): + return _kolmogn(n, x) - p + + return scipy.optimize.brentq(_f, 1.0/n, x1, xtol=1e-14) + + +def kolmogn(n, x, cdf=True): + """Computes the CDF for the two-sided Kolmogorov-Smirnov distribution. + + The two-sided Kolmogorov-Smirnov distribution has as its CDF Pr(D_n <= x), + for a sample of size n drawn from a distribution with CDF F(t), where + :math:`D_n &= sup_t |F_n(t) - F(t)|`, and + :math:`F_n(t)` is the Empirical Cumulative Distribution Function of the sample. + + Parameters + ---------- + n : int, array_like + the number of samples + x : float, array_like + The K-S statistic, float between 0 and 1 + cdf : bool, optional + whether to compute the CDF(default=true) or the SF. + + Returns + ------- + cdf : ndarray + CDF (or SF it cdf is False) at the specified locations. + + The return value has shape the result of numpy broadcasting n and x. + """ + it = np.nditer([n, x, cdf, None], flags=['zerosize_ok'], + op_dtypes=[None, np.float64, np.bool_, np.float64]) + for _n, _x, _cdf, z in it: + if np.isnan(_n): + z[...] = _n + continue + if int(_n) != _n: + raise ValueError(f'n is not integral: {_n}') + z[...] = _kolmogn(int(_n), _x, cdf=_cdf) + result = it.operands[-1] + return result + + +def kolmognp(n, x): + """Computes the PDF for the two-sided Kolmogorov-Smirnov distribution. + + Parameters + ---------- + n : int, array_like + the number of samples + x : float, array_like + The K-S statistic, float between 0 and 1 + + Returns + ------- + pdf : ndarray + The PDF at the specified locations + + The return value has shape the result of numpy broadcasting n and x. + """ + it = np.nditer([n, x, None]) + for _n, _x, z in it: + if np.isnan(_n): + z[...] = _n + continue + if int(_n) != _n: + raise ValueError(f'n is not integral: {_n}') + z[...] = _kolmogn_p(int(_n), _x) + result = it.operands[-1] + return result + + +def kolmogni(n, q, cdf=True): + """Computes the PPF(or ISF) for the two-sided Kolmogorov-Smirnov distribution. + + Parameters + ---------- + n : int, array_like + the number of samples + q : float, array_like + Probabilities, float between 0 and 1 + cdf : bool, optional + whether to compute the PPF(default=true) or the ISF. + + Returns + ------- + ppf : ndarray + PPF (or ISF if cdf is False) at the specified locations + + The return value has shape the result of numpy broadcasting n and x. + """ + it = np.nditer([n, q, cdf, None]) + for _n, _q, _cdf, z in it: + if np.isnan(_n): + z[...] = _n + continue + if int(_n) != _n: + raise ValueError(f'n is not integral: {_n}') + _pcdf, _psf = (_q, 1-_q) if _cdf else (1-_q, _q) + z[...] = _kolmogni(int(_n), _pcdf, _psf) + result = it.operands[-1] + return result diff --git a/v12x/pylib/scipy/stats/_levy_stable/__init__.py b/v12x/pylib/scipy/stats/_levy_stable/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..127c2464cebd4c168e5f01f0c88cfcdd1c9da8ad --- /dev/null +++ b/v12x/pylib/scipy/stats/_levy_stable/__init__.py @@ -0,0 +1,1247 @@ +# + +import warnings +from functools import partial + +import numpy as np + +from scipy import optimize +from scipy import integrate +from scipy.integrate._quadrature import _builtincoeffs +from scipy import interpolate +from scipy.interpolate import RectBivariateSpline +import scipy._external.array_api_extra as xpx +import scipy.special as sc +from .._distn_infrastructure import rv_continuous, _ShapeInfo, rv_continuous_frozen +from .._continuous_distns import uniform, expon, _norm_pdf, _norm_cdf +from .levyst import Nolan +from scipy._lib.doccer import inherit_docstring_from + + +__all__ = ["levy_stable", "levy_stable_gen", "pdf_from_cf_with_fft"] + +# Stable distributions are known for various parameterisations +# some being advantageous for numerical considerations and others +# useful due to their location/scale awareness. +# +# Here we follow [NO] convention (see the references in the docstring +# for levy_stable_gen below). +# +# S0 / Z0 / x0 (aka Zoleterav's M) +# S1 / Z1 / x1 +# +# Where S* denotes parameterisation, Z* denotes standardized +# version where gamma = 1, delta = 0 and x* denotes variable. +# +# Scipy's original Stable was a random variate generator. It +# uses S1 and unfortunately is not a location/scale aware. + + +# default numerical integration tolerance +# used for epsrel in piecewise and both epsrel and epsabs in dni +# (epsabs needed in dni since weighted quad requires epsabs > 0) +_QUAD_EPS = 1.2e-14 + + +def _Phi_Z0(alpha, t): + return ( + -np.tan(np.pi * alpha / 2) * (np.abs(t) ** (1 - alpha) - 1) + if alpha != 1 + else -2.0 * np.log(np.abs(t)) / np.pi + ) + + +def _Phi_Z1(alpha, t): + return ( + np.tan(np.pi * alpha / 2) + if alpha != 1 + else -2.0 * np.log(np.abs(t)) / np.pi + ) + + +def _cf(Phi, t, alpha, beta): + """Characteristic function.""" + return np.exp( + -(np.abs(t) ** alpha) * (1 - 1j * beta * np.sign(t) * Phi(alpha, t)) + ) + + +_cf_Z0 = partial(_cf, _Phi_Z0) +_cf_Z1 = partial(_cf, _Phi_Z1) + + +def _pdf_single_value_cf_integrate(Phi, x, alpha, beta, **kwds): + """To improve DNI accuracy convert characteristic function in to real + valued integral using Euler's formula, then exploit cosine symmetry to + change limits to [0, inf). Finally use cosine addition formula to split + into two parts that can be handled by weighted quad pack. + """ + quad_eps = kwds.get("quad_eps", _QUAD_EPS) + + def integrand1(t): + if t == 0: + return 0 + return np.exp(-(t ** alpha)) * ( + np.cos(beta * (t ** alpha) * Phi(alpha, t)) + ) + + def integrand2(t): + if t == 0: + return 0 + return np.exp(-(t ** alpha)) * ( + np.sin(beta * (t ** alpha) * Phi(alpha, t)) + ) + + with np.errstate(invalid="ignore"): + int1, *ret1 = integrate.quad( + integrand1, + 0, + np.inf, + weight="cos", + wvar=x, + limit=1000, + epsabs=quad_eps, + epsrel=quad_eps, + full_output=1, + ) + + int2, *ret2 = integrate.quad( + integrand2, + 0, + np.inf, + weight="sin", + wvar=x, + limit=1000, + epsabs=quad_eps, + epsrel=quad_eps, + full_output=1, + ) + + return (int1 + int2) / np.pi + + +_pdf_single_value_cf_integrate_Z0 = partial( + _pdf_single_value_cf_integrate, _Phi_Z0 +) +_pdf_single_value_cf_integrate_Z1 = partial( + _pdf_single_value_cf_integrate, _Phi_Z1 +) + + +def _nolan_round_x_near_zeta(x0, alpha, zeta, x_tol_near_zeta): + """Round x close to zeta for Nolan's method in [NO].""" + # "8. When |x0-beta*tan(pi*alpha/2)| is small, the + # computations of the density and cumulative have numerical problems. + # The program works around this by setting + # z = beta*tan(pi*alpha/2) when + # |z-beta*tan(pi*alpha/2)| < tol(5)*alpha**(1/alpha). + # (The bound on the right is ad hoc, to get reasonable behavior + # when alpha is small)." + # where tol(5) = 0.5e-2 by default. + # + # We seem to have partially addressed this through re-expression of + # g(theta) here, but it still needs to be used in some extreme cases. + # Perhaps tol(5) = 0.5e-2 could be reduced for our implementation. + if np.abs(x0 - zeta) < x_tol_near_zeta * alpha ** (1 / alpha): + x0 = zeta + return x0 + + +def _nolan_round_difficult_input( + x0, alpha, beta, zeta, x_tol_near_zeta, alpha_tol_near_one +): + """Round difficult input values for Nolan's method in [NO].""" + + # following Nolan's STABLE, + # "1. When 0 < |alpha-1| < 0.005, the program has numerical problems + # evaluating the pdf and cdf. The current version of the program sets + # alpha=1 in these cases. This approximation is not bad in the S0 + # parameterization." + if np.abs(alpha - 1) < alpha_tol_near_one: + alpha = 1.0 + + # "2. When alpha=1 and |beta| < 0.005, the program has numerical + # problems. The current version sets beta=0." + # We seem to have addressed this through re-expression of g(theta) here + + x0 = _nolan_round_x_near_zeta(x0, alpha, zeta, x_tol_near_zeta) + return x0, alpha, beta + + +def _pdf_single_value_piecewise_Z1(x, alpha, beta, **kwds): + # convert from Nolan's S_1 (aka S) to S_0 (aka Zolaterev M) + # parameterization + + zeta = -beta * np.tan(np.pi * alpha / 2.0) + x0 = x + zeta if alpha != 1 else x + + return _pdf_single_value_piecewise_Z0(x0, alpha, beta, **kwds) + + +def _pdf_single_value_piecewise_Z0(x0, alpha, beta, **kwds): + + quad_eps = kwds.get("quad_eps", _QUAD_EPS) + x_tol_near_zeta = kwds.get("piecewise_x_tol_near_zeta", 0.005) + alpha_tol_near_one = kwds.get("piecewise_alpha_tol_near_one", 0.005) + + zeta = -beta * np.tan(np.pi * alpha / 2.0) + x0, alpha, beta = _nolan_round_difficult_input( + x0, alpha, beta, zeta, x_tol_near_zeta, alpha_tol_near_one + ) + + # some other known distribution pdfs / analytical cases + # TODO: add more where possible with test coverage, + # eg https://en.wikipedia.org/wiki/Stable_distribution#Other_analytic_cases + if alpha == 2.0: + # normal + return _norm_pdf(x0 / np.sqrt(2)) / np.sqrt(2) + elif alpha == 0.5 and beta == 1.0: + # levy + # since S(1/2, 1, gamma, delta; ) == + # S(1/2, 1, gamma, gamma + delta; ). + _x = x0 + 1 + if _x <= 0: + return 0 + + return 1 / np.sqrt(2 * np.pi * _x) / _x * np.exp(-1 / (2 * _x)) + elif alpha == 0.5 and beta == 0.0 and x0 != 0: + # analytical solution [HO] + S, C = sc.fresnel([1 / np.sqrt(2 * np.pi * np.abs(x0))]) + arg = 1 / (4 * np.abs(x0)) + return ( + np.sin(arg) * (0.5 - S[0]) + np.cos(arg) * (0.5 - C[0]) + ) / np.sqrt(2 * np.pi * np.abs(x0) ** 3) + elif alpha == 1.0 and beta == 0.0: + # cauchy + return 1 / (1 + x0 ** 2) / np.pi + + return _pdf_single_value_piecewise_post_rounding_Z0( + x0, alpha, beta, quad_eps, x_tol_near_zeta + ) + + +def _pdf_single_value_piecewise_post_rounding_Z0(x0, alpha, beta, quad_eps, + x_tol_near_zeta): + """Calculate pdf using Nolan's methods as detailed in [NO].""" + + _nolan = Nolan(alpha, beta, x0) + zeta = _nolan.zeta + xi = _nolan.xi + c2 = _nolan.c2 + g = _nolan.g + + # round x0 to zeta again if needed. zeta was recomputed and may have + # changed due to floating point differences. + # See https://github.com/scipy/scipy/pull/18133 + x0 = _nolan_round_x_near_zeta(x0, alpha, zeta, x_tol_near_zeta) + # handle Nolan's initial case logic + if x0 == zeta: + return ( + sc.gamma(1 + 1 / alpha) + * np.cos(xi) + / np.pi + / ((1 + zeta ** 2) ** (1 / alpha / 2)) + ) + elif x0 < zeta: + return _pdf_single_value_piecewise_post_rounding_Z0( + -x0, alpha, -beta, quad_eps, x_tol_near_zeta + ) + + # following Nolan, we may now assume + # x0 > zeta when alpha != 1 + # beta != 0 when alpha == 1 + + # spare calculating integral on null set + # use isclose as macos has fp differences + if np.isclose(-xi, np.pi / 2, rtol=1e-014, atol=1e-014): + return 0.0 + + def integrand(theta): + # limit any numerical issues leading to g_1 < 0 near theta limits + g_1 = g(theta) + if not np.isfinite(g_1) or g_1 < 0: + g_1 = 0 + return g_1 * np.exp(-g_1) + + with np.errstate(all="ignore"): + peak = optimize.bisect( + lambda t: g(t) - 1, -xi, np.pi / 2, xtol=quad_eps + ) + + # this integrand can be very peaked, so we need to force + # QUADPACK to evaluate the function inside its support + # + + # lastly, we add additional samples at + # ~exp(-100), ~exp(-10), ~exp(-5), ~exp(-1) + # to improve QUADPACK's detection of rapidly descending tail behavior + # (this choice is fairly ad hoc) + tail_points = [ + optimize.bisect(lambda t: g(t) - exp_height, -xi, np.pi / 2) + for exp_height in [100, 10, 5] + # exp_height = 1 is handled by peak + ] + # Symmetry is used to ensure x > zeta as above. Integration + # becomes numerically unstable so let `quad` handle + # integration points + if beta != -1: + intg_points = [0, peak] + tail_points + else: + intg_points = None + + intg, *ret = integrate.quad( + integrand, + -xi, + np.pi / 2, + points=intg_points, + limit=100, + epsrel=quad_eps, + epsabs=0, + full_output=1, + ) + + return c2 * intg + + +def _cdf_single_value_piecewise_Z1(x, alpha, beta, **kwds): + # convert from Nolan's S_1 (aka S) to S_0 (aka Zolaterev M) + # parameterization + + zeta = -beta * np.tan(np.pi * alpha / 2.0) + x0 = x + zeta if alpha != 1 else x + + return _cdf_single_value_piecewise_Z0(x0, alpha, beta, **kwds) + + +def _cdf_single_value_piecewise_Z0(x0, alpha, beta, **kwds): + + quad_eps = kwds.get("quad_eps", _QUAD_EPS) + x_tol_near_zeta = kwds.get("piecewise_x_tol_near_zeta", 0.005) + alpha_tol_near_one = kwds.get("piecewise_alpha_tol_near_one", 0.005) + + zeta = -beta * np.tan(np.pi * alpha / 2.0) + x0, alpha, beta = _nolan_round_difficult_input( + x0, alpha, beta, zeta, x_tol_near_zeta, alpha_tol_near_one + ) + + # some other known distribution cdfs / analytical cases + # TODO: add more where possible with test coverage, + # eg https://en.wikipedia.org/wiki/Stable_distribution#Other_analytic_cases + if alpha == 2.0: + # normal + return _norm_cdf(x0 / np.sqrt(2)) + elif alpha == 0.5 and beta == 1.0: + # levy + # since S(1/2, 1, gamma, delta; ) == + # S(1/2, 1, gamma, gamma + delta; ). + _x = x0 + 1 + if _x <= 0: + return 0 + + return sc.erfc(np.sqrt(0.5 / _x)) + elif alpha == 1.0 and beta == 0.0: + # cauchy + return 0.5 + np.arctan(x0) / np.pi + + return _cdf_single_value_piecewise_post_rounding_Z0( + x0, alpha, beta, quad_eps, x_tol_near_zeta + ) + + +def _cdf_single_value_piecewise_post_rounding_Z0(x0, alpha, beta, quad_eps, + x_tol_near_zeta): + """Calculate cdf using Nolan's methods as detailed in [NO].""" + _nolan = Nolan(alpha, beta, x0) + zeta = _nolan.zeta + xi = _nolan.xi + c1 = _nolan.c1 + # c2 = _nolan.c2 + c3 = _nolan.c3 + g = _nolan.g + # round x0 to zeta again if needed. zeta was recomputed and may have + # changed due to floating point differences. + # See https://github.com/scipy/scipy/pull/18133 + x0 = _nolan_round_x_near_zeta(x0, alpha, zeta, x_tol_near_zeta) + # handle Nolan's initial case logic + if (alpha == 1 and beta < 0) or x0 < zeta: + # NOTE: Nolan's paper has a typo here! + # He states F(x) = 1 - F(x, alpha, -beta), but this is clearly + # incorrect since F(-infty) would be 1.0 in this case + # Indeed, the alpha != 1, x0 < zeta case is correct here. + return 1 - _cdf_single_value_piecewise_post_rounding_Z0( + -x0, alpha, -beta, quad_eps, x_tol_near_zeta + ) + elif x0 == zeta: + return 0.5 - xi / np.pi + + # following Nolan, we may now assume + # x0 > zeta when alpha != 1 + # beta > 0 when alpha == 1 + + # spare calculating integral on null set + # use isclose as macos has fp differences + if np.isclose(-xi, np.pi / 2, rtol=1e-014, atol=1e-014): + return c1 + + def integrand(theta): + g_1 = g(theta) + return np.exp(-g_1) + + with np.errstate(all="ignore"): + # shrink supports where required + left_support = -xi + right_support = np.pi / 2 + if alpha > 1: + # integrand(t) monotonic 0 to 1 + if integrand(-xi) != 0.0: + res = optimize.minimize( + integrand, + (-xi,), + method="L-BFGS-B", + bounds=[(-xi, np.pi / 2)], + ) + left_support = res.x[0] + else: + # integrand(t) monotonic 1 to 0 + if integrand(np.pi / 2) != 0.0: + res = optimize.minimize( + integrand, + (np.pi / 2,), + method="L-BFGS-B", + bounds=[(-xi, np.pi / 2)], + ) + right_support = res.x[0] + + intg, *ret = integrate.quad( + integrand, + left_support, + right_support, + points=[left_support, right_support], + limit=100, + epsrel=quad_eps, + epsabs=0, + full_output=1, + ) + + return c1 + c3 * intg + + +def _rvs_Z1(alpha, beta, size=None, random_state=None): + """Simulate random variables using Nolan's methods as detailed in [NO]. + """ + + def alpha1func(alpha, beta, TH, aTH, bTH, cosTH, tanTH, W): + return ( + 2 + / np.pi + * ( + (np.pi / 2 + bTH) * tanTH + - beta * np.log((np.pi / 2 * W * cosTH) / (np.pi / 2 + bTH)) + ) + ) + + def beta0func(alpha, beta, TH, aTH, bTH, cosTH, tanTH, W): + return ( + W + / (cosTH / np.tan(aTH) + np.sin(TH)) + * ((np.cos(aTH) + np.sin(aTH) * tanTH) / W) ** (1.0 / alpha) + ) + + def otherwise(alpha, beta, TH, aTH, bTH, cosTH, tanTH, W): + # alpha is not 1 and beta is not 0 + val0 = beta * np.tan(np.pi * alpha / 2) + th0 = np.arctan(val0) / alpha + val3 = W / (cosTH / np.tan(alpha * (th0 + TH)) + np.sin(TH)) + res3 = val3 * ( + ( + np.cos(aTH) + + np.sin(aTH) * tanTH + - val0 * (np.sin(aTH) - np.cos(aTH) * tanTH) + ) + / W + ) ** (1.0 / alpha) + return res3 + + def alphanot1func(alpha, beta, TH, aTH, bTH, cosTH, tanTH, W): + return xpx.apply_where( + beta == 0, (alpha, beta, TH, aTH, bTH, cosTH, tanTH, W), + beta0func, otherwise) + + alpha = np.broadcast_to(alpha, size) + beta = np.broadcast_to(beta, size) + TH = uniform.rvs( + loc=-np.pi / 2.0, scale=np.pi, size=size, random_state=random_state + ) + W = expon.rvs(size=size, random_state=random_state) + aTH = alpha * TH + bTH = beta * TH + cosTH = np.cos(TH) + tanTH = np.tan(TH) + return xpx.apply_where( + alpha == 1, (alpha, beta, TH, aTH, bTH, cosTH, tanTH, W), + alpha1func, alphanot1func) + + +def _fitstart_S0(data): + alpha, beta, delta1, gamma = _fitstart_S1(data) + + # Formulas for mapping parameters in S1 parameterization to + # those in S0 parameterization can be found in [NO]. Note that + # only delta changes. + if alpha != 1: + delta0 = delta1 + beta * gamma * np.tan(np.pi * alpha / 2.0) + else: + delta0 = delta1 + 2 * beta * gamma * np.log(gamma) / np.pi + + return alpha, beta, delta0, gamma + + +def _fitstart_S1(data): + # We follow McCullock 1986 method - Simple Consistent Estimators + # of Stable Distribution Parameters + + # fmt: off + # Table III and IV + nu_alpha_range = [2.439, 2.5, 2.6, 2.7, 2.8, 3, 3.2, 3.5, 4, + 5, 6, 8, 10, 15, 25] + nu_beta_range = [0, 0.1, 0.2, 0.3, 0.5, 0.7, 1] + + # table III - alpha = psi_1(nu_alpha, nu_beta) + alpha_table = np.array([ + [2.000, 2.000, 2.000, 2.000, 2.000, 2.000, 2.000], + [1.916, 1.924, 1.924, 1.924, 1.924, 1.924, 1.924], + [1.808, 1.813, 1.829, 1.829, 1.829, 1.829, 1.829], + [1.729, 1.730, 1.737, 1.745, 1.745, 1.745, 1.745], + [1.664, 1.663, 1.663, 1.668, 1.676, 1.676, 1.676], + [1.563, 1.560, 1.553, 1.548, 1.547, 1.547, 1.547], + [1.484, 1.480, 1.471, 1.460, 1.448, 1.438, 1.438], + [1.391, 1.386, 1.378, 1.364, 1.337, 1.318, 1.318], + [1.279, 1.273, 1.266, 1.250, 1.210, 1.184, 1.150], + [1.128, 1.121, 1.114, 1.101, 1.067, 1.027, 0.973], + [1.029, 1.021, 1.014, 1.004, 0.974, 0.935, 0.874], + [0.896, 0.892, 0.884, 0.883, 0.855, 0.823, 0.769], + [0.818, 0.812, 0.806, 0.801, 0.780, 0.756, 0.691], + [0.698, 0.695, 0.692, 0.689, 0.676, 0.656, 0.597], + [0.593, 0.590, 0.588, 0.586, 0.579, 0.563, 0.513]]).T + # transpose because interpolation with `RectBivariateSpline` is with + # `nu_beta` as `x` and `nu_alpha` as `y` + + # table IV - beta = psi_2(nu_alpha, nu_beta) + beta_table = np.array([ + [0, 2.160, 1.000, 1.000, 1.000, 1.000, 1.000], + [0, 1.592, 3.390, 1.000, 1.000, 1.000, 1.000], + [0, 0.759, 1.800, 1.000, 1.000, 1.000, 1.000], + [0, 0.482, 1.048, 1.694, 1.000, 1.000, 1.000], + [0, 0.360, 0.760, 1.232, 2.229, 1.000, 1.000], + [0, 0.253, 0.518, 0.823, 1.575, 1.000, 1.000], + [0, 0.203, 0.410, 0.632, 1.244, 1.906, 1.000], + [0, 0.165, 0.332, 0.499, 0.943, 1.560, 1.000], + [0, 0.136, 0.271, 0.404, 0.689, 1.230, 2.195], + [0, 0.109, 0.216, 0.323, 0.539, 0.827, 1.917], + [0, 0.096, 0.190, 0.284, 0.472, 0.693, 1.759], + [0, 0.082, 0.163, 0.243, 0.412, 0.601, 1.596], + [0, 0.074, 0.147, 0.220, 0.377, 0.546, 1.482], + [0, 0.064, 0.128, 0.191, 0.330, 0.478, 1.362], + [0, 0.056, 0.112, 0.167, 0.285, 0.428, 1.274]]).T + + # Table V and VII + # These are ordered with decreasing `alpha_range`; so we will need to + # reverse them as required by RectBivariateSpline. + alpha_range = [2, 1.9, 1.8, 1.7, 1.6, 1.5, 1.4, 1.3, 1.2, 1.1, + 1, 0.9, 0.8, 0.7, 0.6, 0.5][::-1] + beta_range = [0, 0.25, 0.5, 0.75, 1] + + # Table V - nu_c = psi_3(alpha, beta) + nu_c_table = np.array([ + [1.908, 1.908, 1.908, 1.908, 1.908], + [1.914, 1.915, 1.916, 1.918, 1.921], + [1.921, 1.922, 1.927, 1.936, 1.947], + [1.927, 1.930, 1.943, 1.961, 1.987], + [1.933, 1.940, 1.962, 1.997, 2.043], + [1.939, 1.952, 1.988, 2.045, 2.116], + [1.946, 1.967, 2.022, 2.106, 2.211], + [1.955, 1.984, 2.067, 2.188, 2.333], + [1.965, 2.007, 2.125, 2.294, 2.491], + [1.980, 2.040, 2.205, 2.435, 2.696], + [2.000, 2.085, 2.311, 2.624, 2.973], + [2.040, 2.149, 2.461, 2.886, 3.356], + [2.098, 2.244, 2.676, 3.265, 3.912], + [2.189, 2.392, 3.004, 3.844, 4.775], + [2.337, 2.634, 3.542, 4.808, 6.247], + [2.588, 3.073, 4.534, 6.636, 9.144]])[::-1].T + # transpose because interpolation with `RectBivariateSpline` is with + # `beta` as `x` and `alpha` as `y` + + # Table VII - nu_zeta = psi_5(alpha, beta) + nu_zeta_table = np.array([ + [0, 0.000, 0.000, 0.000, 0.000], + [0, -0.017, -0.032, -0.049, -0.064], + [0, -0.030, -0.061, -0.092, -0.123], + [0, -0.043, -0.088, -0.132, -0.179], + [0, -0.056, -0.111, -0.170, -0.232], + [0, -0.066, -0.134, -0.206, -0.283], + [0, -0.075, -0.154, -0.241, -0.335], + [0, -0.084, -0.173, -0.276, -0.390], + [0, -0.090, -0.192, -0.310, -0.447], + [0, -0.095, -0.208, -0.346, -0.508], + [0, -0.098, -0.223, -0.380, -0.576], + [0, -0.099, -0.237, -0.424, -0.652], + [0, -0.096, -0.250, -0.469, -0.742], + [0, -0.089, -0.262, -0.520, -0.853], + [0, -0.078, -0.272, -0.581, -0.997], + [0, -0.061, -0.279, -0.659, -1.198]])[::-1].T + # fmt: on + + psi_1 = RectBivariateSpline(nu_beta_range, nu_alpha_range, + alpha_table, kx=1, ky=1, s=0) + + def psi_1_1(nu_beta, nu_alpha): + return psi_1(nu_beta, nu_alpha) \ + if nu_beta > 0 else psi_1(-nu_beta, nu_alpha) + + psi_2 = RectBivariateSpline(nu_beta_range, nu_alpha_range, + beta_table, kx=1, ky=1, s=0) + + def psi_2_1(nu_beta, nu_alpha): + return psi_2(nu_beta, nu_alpha) \ + if nu_beta > 0 else -psi_2(-nu_beta, nu_alpha) + + phi_3 = RectBivariateSpline(beta_range, alpha_range, nu_c_table, + kx=1, ky=1, s=0) + + def phi_3_1(beta, alpha): + return phi_3(beta, alpha) if beta > 0 else phi_3(-beta, alpha) + + phi_5 = RectBivariateSpline(beta_range, alpha_range, nu_zeta_table, + kx=1, ky=1, s=0) + + def phi_5_1(beta, alpha): + return phi_5(beta, alpha) if beta > 0 else -phi_5(-beta, alpha) + + # quantiles + p05 = np.percentile(data, 5) + p50 = np.percentile(data, 50) + p95 = np.percentile(data, 95) + p25 = np.percentile(data, 25) + p75 = np.percentile(data, 75) + + nu_alpha = (p95 - p05) / (p75 - p25) + nu_beta = (p95 + p05 - 2 * p50) / (p95 - p05) + + if nu_alpha >= 2.439: + eps = np.finfo(float).eps + alpha = np.clip(psi_1_1(nu_beta, nu_alpha)[0, 0], eps, 2.) + beta = np.clip(psi_2_1(nu_beta, nu_alpha)[0, 0], -1.0, 1.0) + else: + alpha = 2.0 + beta = np.sign(nu_beta) + c = (p75 - p25) / phi_3_1(beta, alpha)[0, 0] + zeta = p50 + c * phi_5_1(beta, alpha)[0, 0] + delta = zeta-beta*c*np.tan(np.pi*alpha/2.) if alpha != 1. else zeta + + return (alpha, beta, delta, c) + + +class levy_stable_gen(rv_continuous): + r"""A Levy-stable continuous random variable. + + %(before_notes)s + + See Also + -------- + levy, levy_l, cauchy, norm + + Notes + ----- + The distribution for `levy_stable` has characteristic function: + + .. math:: + + \varphi(t, \alpha, \beta, c, \mu) = + e^{it\mu -|ct|^{\alpha}(1-i\beta\operatorname{sign}(t)\Phi)} + + where + + - :math:`\alpha` is the stability parameter (:math:`0 < \alpha \le 2`), + - :math:`\beta` is the skewness parameter (:math:`-1 \le \beta \le 1`), + - :math:`c` is the scale parameter (:math:`c > 0`), + - :math:`\mu` is the location parameter (:math:`-\infty < \mu < \infty`). + + Two parameterizations of :math:`\Phi` are supported. + + The :math:`S_1` parameterization (default): + + .. math:: + + \Phi = \begin{cases} + \tan \left({\frac {\pi \alpha }{2}}\right)&\alpha \neq 1\\ + -{\frac {2}{\pi }}\log |t|&\alpha =1 + \end{cases} + + The :math:`S_0` parameterization: + + .. math:: + + \Phi = \begin{cases} + -\tan \left({\frac {\pi \alpha }{2}}\right)(|ct|^{1-\alpha}-1) + &\alpha \neq 1\\ + -{\frac {2}{\pi }}\log |ct|&\alpha =1 + \end{cases} + + + The probability density function for `levy_stable` is: + + .. math:: + + f(x) = \frac{1}{2\pi}\int_{-\infty}^\infty \varphi(t)e^{-ixt}\,dt + + where :math:`-\infty < x < \infty`. This integral does not have a known + closed form. + + `levy_stable` generalizes several distributions. Where possible, they + should be used instead. Specifically, when the shape parameters + assume the values in the table below, the corresponding equivalent + distribution should be used. + + ========= ======== =========== + ``alpha`` ``beta`` Equivalent + ========= ======== =========== + 1/2 -1 `levy_l` + 1/2 1 `levy` + 1 0 `cauchy` + 2 any `norm` (with ``scale=sqrt(2)``) + ========= ======== =========== + + Evaluation of the pdf uses Nolan's piecewise integration approach with the + Zolotarev :math:`M` parameterization by default. There is also the option + to use direct numerical integration of the standard parameterization of the + characteristic function or to evaluate by taking the FFT of the + characteristic function. + + The default method can changed by setting the class variable + ``levy_stable.pdf_default_method`` to one of 'piecewise' for Nolan's + approach, 'dni' for direct numerical integration, or 'fft-simpson' for the + FFT based approach. For the sake of backwards compatibility, the methods + 'best' and 'zolotarev' are equivalent to 'piecewise' and the method + 'quadrature' is equivalent to 'dni'. + + The parameterization can be changed by setting the class variable + ``levy_stable.parameterization`` to either 'S0' or 'S1'. + The default is 'S1'. + + To improve performance of piecewise and direct numerical integration one + can specify ``levy_stable.quad_eps`` (defaults to 1.2e-14). This is used + as both the absolute and relative quadrature tolerance for direct numerical + integration and as the relative quadrature tolerance for the piecewise + method. One can also specify ``levy_stable.piecewise_x_tol_near_zeta`` + (defaults to 0.005) for how close x is to zeta before it is considered the + same as x [NO]. The exact check is + ``abs(x0 - zeta) < piecewise_x_tol_near_zeta*alpha**(1/alpha)``. One can + also specify ``levy_stable.piecewise_alpha_tol_near_one`` (defaults to + 0.005) for how close alpha is to 1 before being considered equal to 1. + + To increase accuracy of FFT calculation one can specify + ``levy_stable.pdf_fft_grid_spacing`` (defaults to 0.001) and + ``pdf_fft_n_points_two_power`` (defaults to None which means a value is + calculated that sufficiently covers the input range). + + Further control over FFT calculation is available by setting + ``pdf_fft_interpolation_degree`` (defaults to 3) for spline order and + ``pdf_fft_interpolation_level`` for determining the number of points to use + in the Newton-Cotes formula when approximating the characteristic function + (considered experimental). + + Evaluation of the cdf uses Nolan's piecewise integration approach with the + Zolatarev :math:`S_0` parameterization by default. There is also the option + to evaluate through integration of an interpolated spline of the pdf + calculated by means of the FFT method. The settings affecting FFT + calculation are the same as for pdf calculation. The default cdf method can + be changed by setting ``levy_stable.cdf_default_method`` to either + 'piecewise' or 'fft-simpson'. For cdf calculations the Zolatarev method is + superior in accuracy, so FFT is disabled by default. + + Fitting estimate uses quantile estimation method in [MC]. MLE estimation of + parameters in fit method uses this quantile estimate initially. Note that + MLE doesn't always converge if using FFT for pdf calculations; this will be + the case if alpha <= 1 where the FFT approach doesn't give good + approximations. + + Any non-missing value for the attribute + ``levy_stable.pdf_fft_min_points_threshold`` will set + ``levy_stable.pdf_default_method`` to 'fft-simpson' if a valid + default method is not otherwise set. + + + + .. warning:: + + For pdf calculations FFT calculation is considered experimental. + + For cdf calculations FFT calculation is considered experimental. Use + Zolatarev's method instead (default). + + The probability density above is defined in the "standardized" form. To + shift and/or scale the distribution use the ``loc`` and ``scale`` + parameters. + Generally ``%(name)s.pdf(x, %(shapes)s, loc, scale)`` is identically + equivalent to ``%(name)s.pdf(y, %(shapes)s) / scale`` with + ``y = (x - loc) / scale``, except in the ``S1`` parameterization if + ``alpha == 1``. In that case ``%(name)s.pdf(x, %(shapes)s, loc, scale)`` + is identically equivalent to ``%(name)s.pdf(y, %(shapes)s) / scale`` with + ``y = (x - loc - 2 * beta * scale * np.log(scale) / np.pi) / scale``. + See [NO2]_ Definition 1.8 for more information. + Note that shifting the location of a distribution + does not make it a "noncentral" distribution. + + References + ---------- + .. [MC] McCulloch, J., 1986. Simple consistent estimators of stable + distribution parameters. Communications in Statistics - Simulation and + Computation 15, 11091136. + .. [WZ] Wang, Li and Zhang, Ji-Hong, 2008. Simpson's rule based FFT method + to compute densities of stable distribution. + .. [NO] Nolan, J., 1997. Numerical Calculation of Stable Densities and + distributions Functions. + .. [NO2] Nolan, J., 2018. Stable Distributions: Models for Heavy Tailed + Data. + .. [HO] Hopcraft, K. I., Jakeman, E., Tanner, R. M. J., 1999. Lévy random + walks with fluctuating step number and multiscale behavior. + + %(example)s + + """ + # Configurable options as class variables + # (accessible from self by attribute lookup). + parameterization = "S1" + pdf_default_method = "piecewise" + cdf_default_method = "piecewise" + quad_eps = _QUAD_EPS + piecewise_x_tol_near_zeta = 0.005 + piecewise_alpha_tol_near_one = 0.005 + pdf_fft_min_points_threshold = None + pdf_fft_grid_spacing = 0.001 + pdf_fft_n_points_two_power = None + pdf_fft_interpolation_level = 3 + pdf_fft_interpolation_degree = 3 + + def __call__(self, *args, **params): + dist = levy_stable_frozen(self, *args, **params) + dist.parameterization = self.parameterization + return dist + + def _argcheck(self, alpha, beta): + return (alpha > 0) & (alpha <= 2) & (beta <= 1) & (beta >= -1) + + def _shape_info(self): + ialpha = _ShapeInfo("alpha", False, (0, 2), (False, True)) + ibeta = _ShapeInfo("beta", False, (-1, 1), (True, True)) + return [ialpha, ibeta] + + def _parameterization(self): + allowed = ("S0", "S1") + pz = self.parameterization + if pz not in allowed: + raise RuntimeError( + f"Parameterization '{pz}' in supported list: {allowed}" + ) + return pz + + @inherit_docstring_from(rv_continuous) + def rvs(self, *args, **kwds): + X1 = super().rvs(*args, **kwds) + + kwds.pop("discrete", None) + kwds.pop("random_state", None) + (alpha, beta), delta, gamma, size = self._parse_args_rvs(*args, **kwds) + + # shift location for this parameterisation (S1) + X1 = np.where( + alpha == 1.0, X1 + 2 * beta * gamma * np.log(gamma) / np.pi, X1 + ) + + if self._parameterization() == "S0": + return np.where( + alpha == 1.0, + X1 - (beta * 2 * gamma * np.log(gamma) / np.pi), + X1 - gamma * beta * np.tan(np.pi * alpha / 2.0), + ) + elif self._parameterization() == "S1": + return X1 + + def _rvs(self, alpha, beta, size=None, random_state=None): + return _rvs_Z1(alpha, beta, size, random_state) + + @inherit_docstring_from(rv_continuous) + def pdf(self, x, *args, **kwds): + # override base class version to correct + # location for S1 parameterization + if self._parameterization() == "S0": + return super().pdf(x, *args, **kwds) + elif self._parameterization() == "S1": + (alpha, beta), delta, gamma = self._parse_args(*args, **kwds) + if np.all(np.reshape(alpha, (1, -1))[0, :] != 1): + return super().pdf(x, *args, **kwds) + else: + # correct location for this parameterisation + x = np.reshape(x, (1, -1))[0, :] + x, alpha, beta = np.broadcast_arrays(x, alpha, beta) + + data_in = np.dstack((x, alpha, beta))[0] + data_out = np.empty(shape=(len(data_in), 1)) + # group data in unique arrays of alpha, beta pairs + uniq_param_pairs = np.unique(data_in[:, 1:], axis=0) + for pair in uniq_param_pairs: + _alpha, _beta = pair + _delta = ( + delta + 2 * _beta * gamma * np.log(gamma) / np.pi + if _alpha == 1.0 + else delta + ) + data_mask = np.all(data_in[:, 1:] == pair, axis=-1) + _x = data_in[data_mask, 0] + data_out[data_mask] = ( + super() + .pdf(_x, _alpha, _beta, loc=_delta, scale=gamma) + .reshape(len(_x), 1) + ) + output = data_out.T[0] + if output.shape == (1,): + return output[0] + return output + + def _pdf(self, x, alpha, beta): + if self._parameterization() == "S0": + _pdf_single_value_piecewise = _pdf_single_value_piecewise_Z0 + _pdf_single_value_cf_integrate = _pdf_single_value_cf_integrate_Z0 + _cf = _cf_Z0 + elif self._parameterization() == "S1": + _pdf_single_value_piecewise = _pdf_single_value_piecewise_Z1 + _pdf_single_value_cf_integrate = _pdf_single_value_cf_integrate_Z1 + _cf = _cf_Z1 + + x = np.asarray(x).reshape(1, -1)[0, :] + + x, alpha, beta = np.broadcast_arrays(x, alpha, beta) + + data_in = np.dstack((x, alpha, beta))[0] + data_out = np.empty(shape=(len(data_in), 1)) + + pdf_default_method_name = self.pdf_default_method + if pdf_default_method_name in ("piecewise", "best", "zolotarev"): + pdf_single_value_method = _pdf_single_value_piecewise + elif pdf_default_method_name in ("dni", "quadrature"): + pdf_single_value_method = _pdf_single_value_cf_integrate + elif ( + pdf_default_method_name == "fft-simpson" + or self.pdf_fft_min_points_threshold is not None + ): + pdf_single_value_method = None + + pdf_single_value_kwds = { + "quad_eps": self.quad_eps, + "piecewise_x_tol_near_zeta": self.piecewise_x_tol_near_zeta, + "piecewise_alpha_tol_near_one": self.piecewise_alpha_tol_near_one, + } + + fft_grid_spacing = self.pdf_fft_grid_spacing + fft_n_points_two_power = self.pdf_fft_n_points_two_power + fft_interpolation_level = self.pdf_fft_interpolation_level + fft_interpolation_degree = self.pdf_fft_interpolation_degree + + # group data in unique arrays of alpha, beta pairs + uniq_param_pairs = np.unique(data_in[:, 1:], axis=0) + for pair in uniq_param_pairs: + data_mask = np.all(data_in[:, 1:] == pair, axis=-1) + data_subset = data_in[data_mask] + if pdf_single_value_method is not None: + data_out[data_mask] = np.array( + [ + pdf_single_value_method( + _x, _alpha, _beta, **pdf_single_value_kwds + ) + for _x, _alpha, _beta in data_subset + ] + ).reshape(len(data_subset), 1) + else: + warnings.warn( + "Density calculations experimental for FFT method." + + " Use combination of piecewise and dni methods instead.", + RuntimeWarning, stacklevel=3, + ) + _alpha, _beta = pair + _x = data_subset[:, (0,)] + + if _alpha < 1.0: + raise RuntimeError( + "FFT method does not work well for alpha less than 1." + ) + + # need enough points to "cover" _x for interpolation + if fft_grid_spacing is None and fft_n_points_two_power is None: + raise ValueError( + "One of fft_grid_spacing or fft_n_points_two_power " + + "needs to be set." + ) + max_abs_x = np.max(np.abs(_x)) + h = ( + 2 ** (3 - fft_n_points_two_power) * max_abs_x + if fft_grid_spacing is None + else fft_grid_spacing + ) + q = ( + np.ceil(np.log(2 * max_abs_x / h) / np.log(2)) + 2 + if fft_n_points_two_power is None + else int(fft_n_points_two_power) + ) + + # for some parameters, the range of x can be quite + # large, let's choose an arbitrary cut off (8GB) to save on + # computer memory. + MAX_Q = 30 + if q > MAX_Q: + raise RuntimeError( + "fft_n_points_two_power has a maximum " + + f"value of {MAX_Q}" + ) + + density_x, density = pdf_from_cf_with_fft( + lambda t: _cf(t, _alpha, _beta), + h=h, + q=q, + level=fft_interpolation_level, + ) + f = interpolate.InterpolatedUnivariateSpline( + density_x, np.real(density), k=fft_interpolation_degree + ) # patch FFT to use cubic + data_out[data_mask] = f(_x) + + return data_out.T[0] + + @inherit_docstring_from(rv_continuous) + def cdf(self, x, *args, **kwds): + # override base class version to correct + # location for S1 parameterization + # NOTE: this is near identical to pdf() above + if self._parameterization() == "S0": + return super().cdf(x, *args, **kwds) + elif self._parameterization() == "S1": + (alpha, beta), delta, gamma = self._parse_args(*args, **kwds) + if np.all(np.reshape(alpha, (1, -1))[0, :] != 1): + return super().cdf(x, *args, **kwds) + else: + # correct location for this parameterisation + x = np.reshape(x, (1, -1))[0, :] + x, alpha, beta = np.broadcast_arrays(x, alpha, beta) + + data_in = np.dstack((x, alpha, beta))[0] + data_out = np.empty(shape=(len(data_in), 1)) + # group data in unique arrays of alpha, beta pairs + uniq_param_pairs = np.unique(data_in[:, 1:], axis=0) + for pair in uniq_param_pairs: + _alpha, _beta = pair + _delta = ( + delta + 2 * _beta * gamma * np.log(gamma) / np.pi + if _alpha == 1.0 + else delta + ) + data_mask = np.all(data_in[:, 1:] == pair, axis=-1) + _x = data_in[data_mask, 0] + data_out[data_mask] = ( + super() + .cdf(_x, _alpha, _beta, loc=_delta, scale=gamma) + .reshape(len(_x), 1) + ) + output = data_out.T[0] + if output.shape == (1,): + return output[0] + return output + + def _cdf(self, x, alpha, beta): + if self._parameterization() == "S0": + _cdf_single_value_piecewise = _cdf_single_value_piecewise_Z0 + _cf = _cf_Z0 + elif self._parameterization() == "S1": + _cdf_single_value_piecewise = _cdf_single_value_piecewise_Z1 + _cf = _cf_Z1 + + x = np.asarray(x).reshape(1, -1)[0, :] + + x, alpha, beta = np.broadcast_arrays(x, alpha, beta) + + data_in = np.dstack((x, alpha, beta))[0] + data_out = np.empty(shape=(len(data_in), 1)) + + cdf_default_method_name = self.cdf_default_method + if cdf_default_method_name == "piecewise": + cdf_single_value_method = _cdf_single_value_piecewise + elif cdf_default_method_name == "fft-simpson": + cdf_single_value_method = None + + cdf_single_value_kwds = { + "quad_eps": self.quad_eps, + "piecewise_x_tol_near_zeta": self.piecewise_x_tol_near_zeta, + "piecewise_alpha_tol_near_one": self.piecewise_alpha_tol_near_one, + } + + fft_grid_spacing = self.pdf_fft_grid_spacing + fft_n_points_two_power = self.pdf_fft_n_points_two_power + fft_interpolation_level = self.pdf_fft_interpolation_level + fft_interpolation_degree = self.pdf_fft_interpolation_degree + + # group data in unique arrays of alpha, beta pairs + uniq_param_pairs = np.unique(data_in[:, 1:], axis=0) + for pair in uniq_param_pairs: + data_mask = np.all(data_in[:, 1:] == pair, axis=-1) + data_subset = data_in[data_mask] + if cdf_single_value_method is not None: + data_out[data_mask] = np.array( + [ + cdf_single_value_method( + _x, _alpha, _beta, **cdf_single_value_kwds + ) + for _x, _alpha, _beta in data_subset + ] + ).reshape(len(data_subset), 1) + else: + warnings.warn( + "Cumulative density calculations experimental for FFT" + + " method. Use piecewise method instead.", + RuntimeWarning, stacklevel=3, + ) + _alpha, _beta = pair + _x = data_subset[:, (0,)] + + # need enough points to "cover" _x for interpolation + if fft_grid_spacing is None and fft_n_points_two_power is None: + raise ValueError( + "One of fft_grid_spacing or fft_n_points_two_power " + + "needs to be set." + ) + max_abs_x = np.max(np.abs(_x)) + h = ( + 2 ** (3 - fft_n_points_two_power) * max_abs_x + if fft_grid_spacing is None + else fft_grid_spacing + ) + q = ( + np.ceil(np.log(2 * max_abs_x / h) / np.log(2)) + 2 + if fft_n_points_two_power is None + else int(fft_n_points_two_power) + ) + + density_x, density = pdf_from_cf_with_fft( + lambda t: _cf(t, _alpha, _beta), + h=h, + q=q, + level=fft_interpolation_level, + ) + f = interpolate.InterpolatedUnivariateSpline( + density_x, np.real(density), k=fft_interpolation_degree + ) + data_out[data_mask] = np.array( + [f.integral(self.a, float(x_1.squeeze())) for x_1 in _x] + ).reshape(data_out[data_mask].shape) + + return data_out.T[0] + + def _fitstart(self, data): + if self._parameterization() == "S0": + _fitstart = _fitstart_S0 + elif self._parameterization() == "S1": + _fitstart = _fitstart_S1 + return _fitstart(data) + + def _stats(self, alpha, beta): + mu = 0 if alpha > 1 else np.nan + mu2 = 2 if alpha == 2 else np.inf + g1 = 0.0 if alpha == 2.0 else np.nan + g2 = 0.0 if alpha == 2.0 else np.nan + return mu, mu2, g1, g2 + + +# cotes numbers - see sequence from http://oeis.org/A100642 +Cotes_table = np.array( + [[], [1]] + [v[2] for v in _builtincoeffs.values()], dtype=object +) +Cotes = np.array( + [ + np.pad(r, (0, len(Cotes_table) - 1 - len(r)), mode='constant') + for r in Cotes_table + ] +) + + +def pdf_from_cf_with_fft(cf, h=0.01, q=9, level=3): + """Calculates pdf from characteristic function. + + Uses fast Fourier transform with Newton-Cotes integration following [WZ]. + Defaults to using Simpson's method (3-point Newton-Cotes integration). + + Parameters + ---------- + cf : callable + Single argument function from float -> complex expressing a + characteristic function for some distribution. + h : Optional[float] + Step size for Newton-Cotes integration. Default: 0.01 + q : Optional[int] + Use 2**q steps when performing Newton-Cotes integration. + The infinite integral in the inverse Fourier transform will then + be restricted to the interval [-2**q * h / 2, 2**q * h / 2]. Setting + the number of steps equal to a power of 2 allows the fft to be + calculated in O(n*log(n)) time rather than O(n**2). + Default: 9 + level : Optional[int] + Calculate integral using n-point Newton-Cotes integration for + n = level. The 3-point Newton-Cotes formula corresponds to Simpson's + rule. Default: 3 + + Returns + ------- + x_l : ndarray + Array of points x at which pdf is estimated. 2**q equally spaced + points from -pi/h up to but not including pi/h. + density : ndarray + Estimated values of pdf corresponding to cf at points in x_l. + + References + ---------- + .. [WZ] Wang, Li and Zhang, Ji-Hong, 2008. Simpson's rule based FFT method + to compute densities of stable distribution. + """ + n = level + N = 2**q + steps = np.arange(0, N) + L = N * h / 2 + x_l = np.pi * (steps - N / 2) / L + if level > 1: + indices = np.arange(n).reshape(n, 1) + s1 = np.sum( + (-1) ** steps * Cotes[n, indices] * np.fft.fft( + (-1)**steps * cf(-L + h * steps + h * indices / (n - 1)) + ) * np.exp( + 1j * np.pi * indices / (n - 1) + - 2 * 1j * np.pi * indices * steps / + (N * (n - 1)) + ), + axis=0 + ) + else: + s1 = (-1) ** steps * Cotes[n, 0] * np.fft.fft( + (-1) ** steps * cf(-L + h * steps) + ) + density = h * s1 / (2 * np.pi * np.sum(Cotes[n])) + return (x_l, density) + + +levy_stable = levy_stable_gen(name="levy_stable") + + +class levy_stable_frozen(rv_continuous_frozen): + @property + def parameterization(self): + return self.dist.parameterization + + @parameterization.setter + def parameterization(self, value): + self.dist.parameterization = value diff --git a/v12x/pylib/scipy/stats/_mannwhitneyu.py b/v12x/pylib/scipy/stats/_mannwhitneyu.py new file mode 100644 index 0000000000000000000000000000000000000000..def3e264ba8a623b190c787e29bd6aa72ebd1a8b --- /dev/null +++ b/v12x/pylib/scipy/stats/_mannwhitneyu.py @@ -0,0 +1,537 @@ +import threading +import numpy as np + +from scipy._lib._array_api import (array_namespace, xp_capabilities, xp_size, + xp_promote, is_lazy_array, is_jax, is_marray, + _count_nonmasked) +from scipy._lib._bunch import _make_tuple_bunch +import scipy._external.array_api_extra as xpx +from scipy import special +from scipy import stats +from scipy.stats._stats_py import _rankdata, _get_pvalue, _SimpleNormal +from scipy.stats._morestats import wilcoxon_result_unpacker, wilcoxon_outputs +from scipy.stats._wilcoxon import _correction_sign +from ._axis_nan_policy import _axis_nan_policy_factory + +class _MWU: + '''Distribution of MWU statistic under the null hypothesis''' + + def __init__(self, n1, n2): + self._reset(n1, n2) + + def set_shapes(self, n1, n2): + n1, n2 = min(n1, n2), max(n1, n2) + if (n1, n2) == (self.n1, self.n2): + return + + self.n1 = n1 + self.n2 = n2 + self.s_array = np.zeros(0, dtype=int) + self.configurations = np.zeros(0, dtype=np.uint64) + + def reset(self): + self._reset(self.n1, self.n2) + + def _reset(self, n1, n2): + self.n1 = None + self.n2 = None + self.set_shapes(n1, n2) + + def pmf(self, k): + + # In practice, `pmf` is never called with k > m*n/2. + # If it were, we'd exploit symmetry here: + # k = np.array(k, copy=True) + # k2 = m*n - k + # i = k2 < k + # k[i] = k2[i] + + pmfs = self.build_u_freqs_array(np.max(k)) + return pmfs[k] + + def cdf(self, k): + '''Cumulative distribution function''' + + # In practice, `cdf` is never called with k > m*n/2. + # If it were, we'd exploit symmetry here rather than in `sf` + pmfs = self.build_u_freqs_array(np.max(k)) + cdfs = np.cumsum(pmfs) + return cdfs[k] + + def sf(self, k): + '''Survival function''' + # Note that both CDF and SF include the PMF at k. The p-value is + # calculated from the SF and should include the mass at k, so this + # is desirable + + # Use the fact that the distribution is symmetric and sum from the left + kc = np.asarray(self.n1*self.n2 - k) # complement of k + i = k < kc + if np.any(i): + kc[i] = k[i] + cdfs = np.asarray(self.cdf(kc)) + cdfs[i] = 1. - cdfs[i] + self.pmf(kc[i]) + else: + cdfs = np.asarray(self.cdf(kc)) + return cdfs[()] + + # build_sigma_array and build_u_freqs_array adapted from code + # by @toobaz with permission. Thanks to @andreasloe for the suggestion. + # See https://github.com/scipy/scipy/pull/4933#issuecomment-1898082691 + def build_sigma_array(self, a): + n1, n2 = self.n1, self.n2 + if a + 1 <= self.s_array.size: + return self.s_array[1:a+1] + + s_array = np.zeros(a + 1, dtype=int) + + for d in np.arange(1, n1 + 1): + # All multiples of d, except 0: + indices = np.arange(d, a + 1, d) + # \epsilon_d = 1: + s_array[indices] += d + + for d in np.arange(n2 + 1, n2 + n1 + 1): + # All multiples of d, except 0: + indices = np.arange(d, a + 1, d) + # \epsilon_d = -1: + s_array[indices] -= d + + # We don't need 0: + self.s_array = s_array + return s_array[1:] + + def build_u_freqs_array(self, maxu): + """ + Build all the array of frequencies for u from 0 to maxu. + Assumptions: + n1 <= n2 + maxu <= n1 * n2 / 2 + """ + n1, n2 = self.n1, self.n2 + total = special.binom(n1 + n2, n1) + + if maxu + 1 <= self.configurations.size: + return self.configurations[:maxu + 1] / total + + s_array = self.build_sigma_array(maxu) + + # Start working with ints, for maximum precision and efficiency: + configurations = np.zeros(maxu + 1, dtype=np.uint64) + configurations_is_uint = True + uint_max = np.iinfo(np.uint64).max + # How many ways to have U=0? 1 + configurations[0] = 1 + + for u in np.arange(1, maxu + 1): + coeffs = s_array[u - 1::-1] + new_val = np.dot(configurations[:u], coeffs) / u + if new_val > uint_max and configurations_is_uint: + # OK, we got into numbers too big for uint64. + # So now we start working with floats. + # By doing this since the beginning, we would have lost precision. + # (And working on python long ints would be unbearably slow) + configurations = configurations.astype(float) + configurations_is_uint = False + configurations[u] = new_val + + self.configurations = configurations + return configurations / total + + +# Maintain state for faster repeat calls to `mannwhitneyu`. +# _MWU() is calculated once per thread and stored as an attribute on +# this thread-local variable inside mannwhitneyu(). +_mwu_state = threading.local() + + +def _get_mwu_z(U, n1, n2, t, continuity=True, *, alternative, xp): + '''Standardized MWU statistic''' + # Follows mannwhitneyu [2] + mu = n1 * n2 / 2 + n = n1 + n2 + + # Tie correction according to [2], "Normal approximation and tie correction" + # "A more computationally-efficient form..." + tie_term = xp.sum(t**3 - t, axis=-1) + s = xp.sqrt(n1*n2/12 * ((n + 1) - tie_term/(n*(n-1)))) + + numerator = U - mu + + # Continuity correction. + # Sign is chosen to always increase the p-value to account for the rest of the + # probability mass _at_ q = U. + if continuity: + sign = _correction_sign(numerator, alternative, xp=xp) + numerator -= 0.5 * sign + + # no problem evaluating the norm SF at an infinity + with np.errstate(divide='ignore', invalid='ignore'): + z = numerator / s + return z + + +def _mwu_input_validation(x, y, use_continuity, alternative, axis, method): + ''' Input validation and standardization for mannwhitneyu ''' + xp = array_namespace(x, y) + + if not is_lazy_array(x) and (xp.any(xp.isnan(x)) or xp.any(xp.isnan(y))): + raise ValueError('`x` and `y` must not contain NaNs.') + if xp_size(x) == 0 or xp_size(y) == 0: + raise ValueError('`x` and `y` must be of nonzero size.') + x, y = xp_promote(x, y, force_floating=True, xp=xp) + + bools = {True, False} + if use_continuity not in bools: + raise ValueError(f'`use_continuity` must be one of {bools}.') + + alternatives = {"two-sided", "less", "greater"} + alternative = alternative.lower() + if alternative not in alternatives: + raise ValueError(f'`alternative` must be one of {alternatives}.') + + axis_int = int(axis) + if axis != axis_int: + raise ValueError('`axis` must be an integer.') + + if not isinstance(method, stats.PermutationMethod): + methods = {"asymptotic", "exact", "auto"} + method = method.lower() + if method not in methods: + raise ValueError(f'`method` must be one of {methods}.') + + if is_marray(xp) and method != "asymptotic": + message = "Only `method='asymptotic'` is compatible with MArrays." + raise ValueError(message) + + output_z = True if method == 'asymptotic' else False + + return x, y, use_continuity, alternative, axis_int, method, output_z, xp + + +def _mwu_choose_method(n1, n2, ties): + """Choose method 'asymptotic' or 'exact' depending on input size, ties""" + + # if both inputs are large, asymptotic is OK + if n1 > 8 and n2 > 8: + return "asymptotic" + + # if there are any ties, asymptotic is preferred + if ties: + return "asymptotic" + + return "exact" + + +MannwhitneyuResult = _make_tuple_bunch('MannwhitneyuResult', ['statistic', 'pvalue']) + + +def mwu_result_object(statistic, pvalue, zstatistic=None): + res = MannwhitneyuResult(statistic, pvalue) + if zstatistic is not None: + res.zstatistic = zstatistic + return res + + +@xp_capabilities( + cpu_only=True, # exact calculation only implemented in NumPy + skip_backends=[('dask.array', 'needs rankdata')], + jax_jit=False, # the exact null distribution is NumPy-only + marray=True, + extra_note=("Only ``method='asymptotic'`` is compatible with MArray input." + "``method='auto'`` is incompatible with JAX arrays.")) +@_axis_nan_policy_factory(mwu_result_object, n_samples=2, + result_to_tuple=wilcoxon_result_unpacker, + n_outputs=wilcoxon_outputs) +def mannwhitneyu(x, y, use_continuity=True, alternative="two-sided", + axis=0, method="auto"): + r'''Perform the Mann-Whitney U rank test on two independent samples. + + The Mann-Whitney U test is a nonparametric test of the null hypothesis + that the distribution underlying sample `x` is the same as the + distribution underlying sample `y`. It is often used as a test of + difference in location between distributions. + + Parameters + ---------- + x, y : array-like + N-d arrays of samples. The arrays must be broadcastable except along + the dimension given by `axis`. + use_continuity : bool, optional + Whether a continuity correction (1/2) should be applied. + Default is True when `method` is ``'asymptotic'``; has no effect + otherwise. + alternative : {'two-sided', 'less', 'greater'}, optional + Defines the alternative hypothesis. Default is 'two-sided'. + Let *SX(u)* and *SY(u)* be the survival functions of the + distributions underlying `x` and `y`, respectively. Then the following + alternative hypotheses are available: + + * 'two-sided': the distributions are not equal, i.e. *SX(u) ≠ SY(u)* for + at least one *u*. + * 'less': the distribution underlying `x` is stochastically less + than the distribution underlying `y`, i.e. *SX(u) < SY(u)* for all *u*. + * 'greater': the distribution underlying `x` is stochastically greater + than the distribution underlying `y`, i.e. *SX(u) > SY(u)* for all *u*. + + Under a more restrictive set of assumptions, the alternative hypotheses + can be expressed in terms of the locations of the distributions; + see [5]_ section 5.1. + axis : int, optional + Axis along which to perform the test. Default is 0. + method : {'auto', 'asymptotic', 'exact'} or `PermutationMethod` instance, optional + Selects the method used to calculate the *p*-value. + Default is 'auto'. The following options are available. + + * ``'asymptotic'``: compares the standardized test statistic + against the normal distribution, correcting for ties. + * ``'exact'``: computes the exact *p*-value by comparing the observed + :math:`U` statistic against the exact distribution of the :math:`U` + statistic under the null hypothesis. No correction is made for ties. + * ``'auto'``: chooses ``'exact'`` when the size of one of the samples + is less than or equal to 8 and there are no ties; + chooses ``'asymptotic'`` otherwise. + * `PermutationMethod` instance. In this case, the p-value + is computed using `permutation_test` with the provided + configuration options and other appropriate settings. + + Returns + ------- + res : MannwhitneyuResult + An object containing attributes: + + statistic : float + The Mann-Whitney U statistic corresponding with sample `x`. See + Notes for the test statistic corresponding with sample `y`. + pvalue : float + The associated *p*-value for the chosen `alternative`. + zstatistic : array_like + When ``method = 'approx'``, this is the normalized z-statistic:: + + z = (U - mn - d) / se + + where ``U`` is `statistic` as defined above, ``mn`` is the mean of the + distribution under the null hypothesis, ``d`` is a continuity + correction, and ``se`` is the standard error. + When ``method != 'asymptotic'``, this attribute is not available. + + Notes + ----- + If ``U1`` is the statistic corresponding with sample `x`, then the + statistic corresponding with sample `y` is + ``U2 = x.shape[axis] * y.shape[axis] - U1``. + + The p-value and standardized z-statistic produced by the Mann-Whitney U test is + identical to that produced by a Wilcoxon *rank sum* test, but different raw + statistics are associated with the tests. Specifically, the rank sum statistic + corresponding with the first sample is + ``R1 = U1 + x.shape[axis] * (x.shape[axis] + 1) / 2``. + + `mannwhitneyu` is for independent samples. For related / paired samples, + consider `scipy.stats.wilcoxon`. + + `method` ``'exact'`` is recommended when there are no ties and when either + sample size is less than 8 [1]_. The implementation follows the algorithm + reported in [3]_. + Note that the exact method is *not* corrected for ties, but + `mannwhitneyu` will not raise errors or warnings if there are ties in the + data. If there are ties and either samples is small (fewer than ~10 + observations), consider passing an instance of `PermutationMethod` + as the `method` to perform a permutation test. + + The Mann-Whitney U test is a non-parametric version of the t-test for + independent samples. When the means of samples from the populations + are normally distributed, consider `scipy.stats.ttest_ind`. + + See Also + -------- + scipy.stats.wilcoxon, scipy.stats.ttest_ind + + References + ---------- + .. [1] H.B. Mann and D.R. Whitney, "On a test of whether one of two random + variables is stochastically larger than the other", The Annals of + Mathematical Statistics, Vol. 18, pp. 50-60, 1947. + .. [2] Mann-Whitney U Test, Wikipedia, + http://en.wikipedia.org/wiki/Mann-Whitney_U_test + .. [3] Andreas Löffler, + "Über eine Partition der nat. Zahlen und ihr Anwendung beim U-Test", + Wiss. Z. Univ. Halle, XXXII'83 pp. 87-89. + .. [4] Rosie Shier, "Statistics: 2.3 The Mann-Whitney U Test", Mathematics + Learning Support Centre, 2004. + .. [5] Michael P. Fay and Michael A. Proschan. "Wilcoxon-Mann-Whitney + or t-test? On assumptions for hypothesis tests and multiple \ + interpretations of decision rules." Statistics surveys, Vol. 4, pp. + 1-39, 2010. https://www.ncbi.nlm.nih.gov/pmc/articles/PMC2857732/ + .. [6] MathWorks. "ranksum - Wilcoxon rank sum test - MATLAB". + https://www.mathworks.com/help/stats/ranksum.html. + + Examples + -------- + We follow the example from [4]_: nine randomly sampled young adults were + diagnosed with type II diabetes at the ages below. + + >>> males = [19, 22, 16, 29, 24] + >>> females = [20, 11, 17, 12] + + We use the Mann-Whitney U test to assess whether there is a statistically + significant difference in the diagnosis age of males and females. + The null hypothesis is that the distribution of male diagnosis ages is + the same as the distribution of female diagnosis ages. We decide + that a confidence level of 95% is required to reject the null hypothesis + in favor of the alternative that the distributions are different. + Since the number of samples is very small and there are no ties in the + data, we can compare the observed test statistic against the *exact* + distribution of the test statistic under the null hypothesis. + + >>> from scipy.stats import mannwhitneyu + >>> U1, p = mannwhitneyu(males, females, method="exact") + >>> print(U1) + 17.0 + + `mannwhitneyu` always reports the statistic associated with the first + sample, which, in this case, is males. This agrees with :math:`U_M = 17` + reported in [4]_. The statistic associated with the second statistic + can be calculated: + + >>> nx, ny = len(males), len(females) + >>> U2 = nx*ny - U1 + >>> print(U2) + 3.0 + + This agrees with :math:`U_F = 3` reported in [4]_. The two-sided + *p*-value can be calculated from either statistic, and the value produced + by `mannwhitneyu` agrees with :math:`p = 0.11` reported in [4]_. + + >>> print(p) + 0.1111111111111111 + + The exact distribution of the test statistic is asymptotically normal, so + the example continues by comparing the exact *p*-value against the + *p*-value produced using the normal approximation. + + >>> _, pnorm = mannwhitneyu(males, females, method="asymptotic") + >>> print(pnorm) + 0.11134688653314041 + + Here `mannwhitneyu`'s reported *p*-value appears to conflict with the + value :math:`p = 0.09` given in [4]_. The reason is that [4]_ + does not apply the continuity correction performed by `mannwhitneyu`; + `mannwhitneyu` reduces the distance between the test statistic and the + mean :math:`\mu = n_x n_y / 2` by 0.5 to correct for the fact that the + discrete statistic is being compared against a continuous distribution. + Here, the :math:`U` statistic used is less than the mean, so we reduce + the distance by adding 0.5 in the numerator. + + >>> import numpy as np + >>> from scipy.stats import norm + >>> U = min(U1, U2) + >>> N = nx + ny + >>> z = (U - nx*ny/2 + 0.5) / np.sqrt(nx*ny * (N + 1)/ 12) + >>> p = 2 * norm.cdf(z) # use CDF to get p-value from smaller statistic + >>> print(p) + 0.11134688653314041 + + If desired, we can disable the continuity correction to get a result + that agrees with that reported in [4]_. + + >>> _, pnorm = mannwhitneyu(males, females, use_continuity=False, + ... method="asymptotic") + >>> print(pnorm) + 0.0864107329737 + + Regardless of whether we perform an exact or asymptotic test, the + probability of the test statistic being as extreme or more extreme by + chance exceeds 5%, so we do not consider the results statistically + significant. + + Suppose that, before seeing the data, we had hypothesized that females + would tend to be diagnosed at a younger age than males. + In that case, it would be natural to provide the female ages as the + first input, and we would have performed a one-sided test using + ``alternative = 'less'``: females are diagnosed at an age that is + stochastically less than that of males. + + >>> res = mannwhitneyu(females, males, alternative="less", method="exact") + >>> print(res) + MannwhitneyuResult(statistic=3.0, pvalue=0.05555555555555555) + + Again, the probability of getting a sufficiently low value of the + test statistic by chance under the null hypothesis is greater than 5%, + so we do not reject the null hypothesis in favor of our alternative. + + If it is reasonable to assume that the means of samples from the + populations are normally distributed, we could have used a t-test to + perform the analysis. + + >>> from scipy.stats import ttest_ind + >>> res = ttest_ind(females, males, alternative="less") + >>> print(res) + TtestResult(statistic=-2.239334696520584, + pvalue=0.030068441095757924, + df=7.0) + + Under this assumption, the *p*-value would be low enough to reject the + null hypothesis in favor of the alternative. + + ''' + + x, y, use_continuity, alternative, axis_int, method, output_z, xp = ( + _mwu_input_validation(x, y, use_continuity, alternative, axis, method)) + + # axis=-1 and arrays broadcasted due to _axis_nan_policy decorator + xy = xp.concat((x, y), axis=-1) + + n1 = _count_nonmasked(x, axis=-1, xp=xp) + n2 = _count_nonmasked(y, axis=-1, xp=xp) + + # Follows [2] + ranks, _, t = _rankdata(xy, 'average', return_ties=True) # method 2, step 1 + ranks = xp.astype(ranks, x.dtype, copy=False) + t = xp.astype(t, x.dtype, copy=False) + R1 = xp.sum(ranks[..., :x.shape[-1]], axis=-1) # method 2, step 2 + U1 = R1 - n1*(n1+1)/2 # method 2, step 3 + U2 = n1 * n2 - U1 # as U1 + U2 = n1 * n2 + + if alternative == "greater": + U, f = U1, 1 # U is the statistic to use for exact p-value, f is a factor + elif alternative == "less": + U, f = U2, 1 # Due to symmetry, it uses SF of U2 rather than CDF of U1 + else: + U, f = xp.maximum(U1, U2), 2 # and will multiply SF by two for two-sided test + + if method == "auto": + if is_jax(xp): + message = "`method='auto'` is incompatible with JAX arrays." + raise ValueError(message) + method = _mwu_choose_method(n1, n2, xp.any(t > 1)) + + if method == "exact": + if not hasattr(_mwu_state, 's'): + _mwu_state.s = _MWU(0, 0) + _mwu_state.s.set_shapes(n1, n2) + p = xpx.lazy_apply(_mwu_state.s.sf, xp.astype(U, xp.int64), + as_numpy=True, dtype=xp.float64) + p = xp.astype(p, x.dtype, copy=False) * f + elif method == "asymptotic": + z = _get_mwu_z(U1, n1, n2, t, continuity=use_continuity, + alternative=alternative, xp=xp) + p = _get_pvalue(z, _SimpleNormal(), alternative, xp=xp) + else: # `PermutationMethod` instance (already validated) + def statistic(x, y, axis): + return mannwhitneyu(x, y, use_continuity=use_continuity, + alternative=alternative, axis=axis, + method="asymptotic").statistic + + res = stats.permutation_test((x, y), statistic, axis=axis, + **method._asdict(), alternative=alternative) + p = res.pvalue + + # Ensure that test statistic is not greater than 1 + # This could happen for exact test when U = m*n/2 + p = xp.clip(p, 0., 1.) + + res = MannwhitneyuResult(U1, p) + if output_z: + res.zstatistic = z[()] if z.ndim == 0 else z + return res diff --git a/v12x/pylib/scipy/stats/_mgc.py b/v12x/pylib/scipy/stats/_mgc.py new file mode 100644 index 0000000000000000000000000000000000000000..f2ecbe9747acdcb4731a2755c88352896be10d0a --- /dev/null +++ b/v12x/pylib/scipy/stats/_mgc.py @@ -0,0 +1,552 @@ +import warnings +import numpy as np + +from scipy._lib._array_api import xp_capabilities +from scipy._lib._util import check_random_state, MapWrapper, rng_integers, _contains_nan +from scipy._lib._bunch import _make_tuple_bunch +from scipy.spatial.distance import cdist +from scipy.ndimage import _measurements + +from ._stats import _local_correlations # type: ignore[import-not-found] +from . import distributions + +__all__ = ['multiscale_graphcorr'] + +# FROM MGCPY: https://github.com/neurodata/mgcpy + + +class _ParallelP: + """Helper function to calculate parallel p-value.""" + + def __init__(self, x, y, random_states): + self.x = x + self.y = y + self.random_states = random_states + + def __call__(self, index): + order = self.random_states[index].permutation(self.y.shape[0]) + permy = self.y[order][:, order] + + # calculate permuted stats, store in null distribution + perm_stat = _mgc_stat(self.x, permy)[0] + + return perm_stat + + +def _perm_test(x, y, stat, reps=1000, workers=-1, random_state=None): + r"""Helper function that calculates the p-value. See below for uses. + + Parameters + ---------- + x, y : ndarray + `x` and `y` have shapes ``(n, p)`` and ``(n, q)``. + stat : float + The sample test statistic. + reps : int, optional + The number of replications used to estimate the null when using the + permutation test. The default is 1000 replications. + workers : int or map-like callable, optional + If `workers` is an int the population is subdivided into `workers` + sections and evaluated in parallel (uses + `multiprocessing.Pool `). Supply `-1` to use all cores + available to the Process. Alternatively supply a map-like callable, + such as `multiprocessing.Pool.map` for evaluating the population in + parallel. This evaluation is carried out as `workers(func, iterable)`. + Requires that `func` be pickleable. + random_state : {None, int, `numpy.random.Generator`, + `numpy.random.RandomState`}, optional + + If `seed` is None (or `np.random`), the `numpy.random.RandomState` + singleton is used. + If `seed` is an int, a new ``RandomState`` instance is used, + seeded with `seed`. + If `seed` is already a ``Generator`` or ``RandomState`` instance then + that instance is used. + + Returns + ------- + pvalue : float + The sample test p-value. + null_dist : list + The approximated null distribution. + + """ + # generate seeds for each rep (change to new parallel random number + # capabilities in numpy >= 1.17+) + random_state = check_random_state(random_state) + random_states = [np.random.RandomState(rng_integers(random_state, 1 << 32, + size=4, dtype=np.uint32)) for _ in range(reps)] + + # parallelizes with specified workers over number of reps and set seeds + parallelp = _ParallelP(x=x, y=y, random_states=random_states) + with MapWrapper(workers) as mapwrapper: + null_dist = np.array(list(mapwrapper(parallelp, range(reps)))) + + # calculate p-value and significant permutation map through list + pvalue = (1 + (null_dist >= stat).sum()) / (1 + reps) + + return pvalue, null_dist + + +def _euclidean_dist(x): + return cdist(x, x) + + +MGCResult = _make_tuple_bunch('MGCResult', + ['statistic', 'pvalue', 'mgc_dict'], []) + + +@xp_capabilities(np_only=True) +def multiscale_graphcorr(x, y, compute_distance=_euclidean_dist, reps=1000, + workers=1, is_twosamp=False, random_state=None): + r"""Computes the Multiscale Graph Correlation (MGC) test statistic. + + Specifically, for each point, MGC finds the :math:`k`-nearest neighbors for + one property (e.g. cloud density), and the :math:`l`-nearest neighbors for + the other property (e.g. grass wetness) [1]_. This pair :math:`(k, l)` is + called the "scale". A priori, however, it is not know which scales will be + most informative. So, MGC computes all distance pairs, and then efficiently + computes the distance correlations for all scales. The local correlations + illustrate which scales are relatively informative about the relationship. + The key, therefore, to successfully discover and decipher relationships + between disparate data modalities is to adaptively determine which scales + are the most informative, and the geometric implication for the most + informative scales. Doing so not only provides an estimate of whether the + modalities are related, but also provides insight into how the + determination was made. This is especially important in high-dimensional + data, where simple visualizations do not reveal relationships to the + unaided human eye. Characterizations of this implementation in particular + have been derived from and benchmarked within in [2]_. + + Parameters + ---------- + x, y : ndarray + If ``x`` and ``y`` have shapes ``(n, p)`` and ``(n, q)`` where `n` is + the number of samples and `p` and `q` are the number of dimensions, + then the MGC independence test will be run. Alternatively, ``x`` and + ``y`` can have shapes ``(n, n)`` if they are distance or similarity + matrices, and ``compute_distance`` must be sent to ``None``. If ``x`` + and ``y`` have shapes ``(n, p)`` and ``(m, p)``, an unpaired + two-sample MGC test will be run. + compute_distance : callable, optional + A function that computes the distance or similarity among the samples + within each data matrix. Set to ``None`` if ``x`` and ``y`` are + already distance matrices. The default uses the euclidean norm metric. + If you are calling a custom function, either create the distance + matrix before-hand or create a function of the form + ``compute_distance(x)`` where `x` is the data matrix for which + pairwise distances are calculated. + reps : int, optional + The number of replications used to estimate the null when using the + permutation test. The default is ``1000``. + workers : int or map-like callable, optional + If ``workers`` is an int the population is subdivided into ``workers`` + sections and evaluated in parallel (uses ``multiprocessing.Pool + ``). Supply ``-1`` to use all cores available to the + Process. Alternatively supply a map-like callable, such as + ``multiprocessing.Pool.map`` for evaluating the p-value in parallel. + This evaluation is carried out as ``workers(func, iterable)``. + Requires that `func` be pickleable. The default is ``1``. + is_twosamp : bool, optional + If `True`, a two sample test will be run. If ``x`` and ``y`` have + shapes ``(n, p)`` and ``(m, p)``, this optional will be overridden and + set to ``True``. Set to ``True`` if ``x`` and ``y`` both have shapes + ``(n, p)`` and a two sample test is desired. The default is ``False``. + Note that this will not run if inputs are distance matrices. + random_state : {None, int, `numpy.random.Generator`, + `numpy.random.RandomState`}, optional + + If `seed` is None (or `np.random`), the `numpy.random.RandomState` + singleton is used. + If `seed` is an int, a new ``RandomState`` instance is used, + seeded with `seed`. + If `seed` is already a ``Generator`` or ``RandomState`` instance then + that instance is used. + + Returns + ------- + res : MGCResult + An object containing attributes: + + statistic : float + The sample MGC test statistic within ``[-1, 1]``. + pvalue : float + The p-value obtained via permutation. + mgc_dict : dict + Contains additional useful results: + + - mgc_map : ndarray + A 2D representation of the latent geometry of the + relationship. + - opt_scale : (int, int) + The estimated optimal scale as a ``(x, y)`` pair. + - null_dist : list + The null distribution derived from the permuted matrices. + + See Also + -------- + pearsonr : Pearson correlation coefficient and p-value for testing + non-correlation. + kendalltau : Calculates Kendall's tau. + spearmanr : Calculates a Spearman rank-order correlation coefficient. + + Notes + ----- + A description of the process of MGC and applications on neuroscience data + can be found in [1]_. It is performed using the following steps: + + #. Two distance matrices :math:`D^X` and :math:`D^Y` are computed and + modified to be mean zero columnwise. This results in two + :math:`n \times n` distance matrices :math:`A` and :math:`B` (the + centering and unbiased modification) [3]_. + + #. For all values :math:`k` and :math:`l` from :math:`1, ..., n`, + + * The :math:`k`-nearest neighbor and :math:`l`-nearest neighbor graphs + are calculated for each property. Here, :math:`G_k (i, j)` indicates + the :math:`k`-smallest values of the :math:`i`-th row of :math:`A` + and :math:`H_l (i, j)` indicates the :math:`l` smallested values of + the :math:`i`-th row of :math:`B` + + * Let :math:`\circ` denotes the entry-wise matrix product, then local + correlations are summed and normalized using the following statistic: + + .. math:: + + c^{kl} = \frac{\sum_{ij} A G_k B H_l} + {\sqrt{\sum_{ij} A^2 G_k \times \sum_{ij} B^2 H_l}} + + #. The MGC test statistic is the smoothed optimal local correlation of + :math:`\{ c^{kl} \}`. Denote the smoothing operation as :math:`R(\cdot)` + (which essentially set all isolated large correlations) as 0 and + connected large correlations the same as before, see [3]_.) MGC is, + + .. math:: + + MGC_n (x, y) = \max_{(k, l)} R \left(c^{kl} \left( x_n, y_n \right) + \right) + + The test statistic returns a value between :math:`(-1, 1)` since it is + normalized. + + The p-value returned is calculated using a permutation test. This process + is completed by first randomly permuting :math:`y` to estimate the null + distribution and then calculating the probability of observing a test + statistic, under the null, at least as extreme as the observed test + statistic. + + MGC requires at least 5 samples to run with reliable results. It can also + handle high-dimensional data sets. + In addition, by manipulating the input data matrices, the two-sample + testing problem can be reduced to the independence testing problem [4]_. + Given sample data :math:`U` and :math:`V` of sizes :math:`p \times n` + :math:`p \times m`, data matrix :math:`X` and :math:`Y` can be created as + follows: + + .. math:: + + X = [U | V] \in \mathcal{R}^{p \times (n + m)} + Y = [0_{1 \times n} | 1_{1 \times m}] \in \mathcal{R}^{(n + m)} + + Then, the MGC statistic can be calculated as normal. This methodology can + be extended to similar tests such as distance correlation [4]_. + + .. versionadded:: 1.4.0 + + References + ---------- + .. [1] Vogelstein, J. T., Bridgeford, E. W., Wang, Q., Priebe, C. E., + Maggioni, M., & Shen, C. (2019). Discovering and deciphering + relationships across disparate data modalities. ELife. + .. [2] Panda, S., Palaniappan, S., Xiong, J., Swaminathan, A., + Ramachandran, S., Bridgeford, E. W., ... Vogelstein, J. T. (2019). + mgcpy: A Comprehensive High Dimensional Independence Testing Python + Package. :arXiv:`1907.02088` + .. [3] Shen, C., Priebe, C.E., & Vogelstein, J. T. (2019). From distance + correlation to multiscale graph correlation. Journal of the American + Statistical Association. + .. [4] Shen, C. & Vogelstein, J. T. (2018). The Exact Equivalence of + Distance and Kernel Methods for Hypothesis Testing. + :arXiv:`1806.05514` + + Examples + -------- + >>> import numpy as np + >>> from scipy.stats import multiscale_graphcorr + >>> x = np.arange(100) + >>> y = x + >>> res = multiscale_graphcorr(x, y) + >>> res.statistic, res.pvalue + (1.0, 0.001) + + To run an unpaired two-sample test, + + >>> x = np.arange(100) + >>> y = np.arange(79) + >>> res = multiscale_graphcorr(x, y) + >>> res.statistic, res.pvalue # doctest: +SKIP + (0.033258146255703246, 0.023) + + or, if shape of the inputs are the same, + + >>> x = np.arange(100) + >>> y = x + >>> res = multiscale_graphcorr(x, y, is_twosamp=True) + >>> res.statistic, res.pvalue # doctest: +SKIP + (-0.008021809890200488, 1.0) + + """ + if not isinstance(x, np.ndarray) or not isinstance(y, np.ndarray): + raise ValueError("x and y must be ndarrays") + + # convert arrays of type (n,) to (n, 1) + if x.ndim == 1: + x = x[:, np.newaxis] + elif x.ndim != 2: + raise ValueError(f"Expected a 2-D array `x`, found shape {x.shape}") + if y.ndim == 1: + y = y[:, np.newaxis] + elif y.ndim != 2: + raise ValueError(f"Expected a 2-D array `y`, found shape {y.shape}") + + nx, px = x.shape + ny, py = y.shape + + # check for NaNs + _contains_nan(x, nan_policy='raise') + _contains_nan(y, nan_policy='raise') + + # check for positive or negative infinity and raise error + if np.sum(np.isinf(x)) > 0 or np.sum(np.isinf(y)) > 0: + raise ValueError("Inputs contain infinities") + + if nx != ny: + if px == py: + # reshape x and y for two sample testing + is_twosamp = True + else: + raise ValueError("Shape mismatch, x and y must have shape [n, p] " + "and [n, q] or have shape [n, p] and [m, p].") + + if nx < 5 or ny < 5: + raise ValueError("MGC requires at least 5 samples to give reasonable " + "results.") + + # convert x and y to float + x = x.astype(np.float64) + y = y.astype(np.float64) + + # check if compute_distance_matrix if a callable() + if not callable(compute_distance) and compute_distance is not None: + raise ValueError("Compute_distance must be a function.") + + # check if number of reps exists, integer, or > 0 (if under 1000 raises + # warning) + if not isinstance(reps, int) or reps < 0: + raise ValueError("Number of reps must be an integer greater than 0.") + elif reps < 1000: + msg = ("The number of replications is low (under 1000), and p-value " + "calculations may be unreliable. Use the p-value result, with " + "caution!") + warnings.warn(msg, RuntimeWarning, stacklevel=2) + + if is_twosamp: + if compute_distance is None: + raise ValueError("Cannot run if inputs are distance matrices") + x, y = _two_sample_transform(x, y) + + if compute_distance is not None: + # compute distance matrices for x and y + x = compute_distance(x) + y = compute_distance(y) + + # calculate MGC stat + stat, stat_dict = _mgc_stat(x, y) + stat_mgc_map = stat_dict["stat_mgc_map"] + opt_scale = stat_dict["opt_scale"] + + # calculate permutation MGC p-value + pvalue, null_dist = _perm_test(x, y, stat, reps=reps, workers=workers, + random_state=random_state) + + # save all stats (other than stat/p-value) in dictionary + mgc_dict = {"mgc_map": stat_mgc_map, + "opt_scale": opt_scale, + "null_dist": null_dist} + + # create result object with alias for backward compatibility + res = MGCResult(stat, pvalue, mgc_dict) + res.stat = stat + return res + + +def _mgc_stat(distx, disty): + r"""Helper function that calculates the MGC stat. See above for use. + + Parameters + ---------- + distx, disty : ndarray + `distx` and `disty` have shapes ``(n, p)`` and ``(n, q)`` or + ``(n, n)`` and ``(n, n)`` + if distance matrices. + + Returns + ------- + stat : float + The sample MGC test statistic within ``[-1, 1]``. + stat_dict : dict + Contains additional useful additional returns containing the following + keys: + + - stat_mgc_map : ndarray + MGC-map of the statistics. + - opt_scale : (float, float) + The estimated optimal scale as a ``(x, y)`` pair. + + """ + # calculate MGC map and optimal scale + stat_mgc_map = _local_correlations(distx, disty, global_corr='mgc') + + m, n = stat_mgc_map.shape + if m == 1 or n == 1: + # the global scale at is the statistic calculated at maximal nearest + # neighbors. There is not enough local scale to search over, so + # default to global scale + stat = stat_mgc_map[m - 1][n - 1] + opt_scale = m * n + else: + samp_size = len(distx) - 1 + + # threshold to find connected region of significant local correlations + sig_connect = _threshold_mgc_map(stat_mgc_map, samp_size) + + # maximum within the significant region + stat, opt_scale = _smooth_mgc_map(sig_connect, stat_mgc_map) + + stat_dict = {"stat_mgc_map": stat_mgc_map, + "opt_scale": opt_scale} + + return stat, stat_dict + + +def _threshold_mgc_map(stat_mgc_map, samp_size): + r""" + Finds a connected region of significance in the MGC-map by thresholding. + + Parameters + ---------- + stat_mgc_map : ndarray + All local correlations within ``[-1,1]``. + samp_size : int + The sample size of original data. + + Returns + ------- + sig_connect : ndarray + A binary matrix with 1's indicating the significant region. + + """ + m, n = stat_mgc_map.shape + + # 0.02 is simply an empirical threshold, this can be set to 0.01 or 0.05 + # with varying levels of performance. Threshold is based on a beta + # approximation. + per_sig = 1 - (0.02 / samp_size) # Percentile to consider as significant + threshold = samp_size * (samp_size - 3)/4 - 1/2 # Beta approximation + threshold = distributions.beta.ppf(per_sig, threshold, threshold) * 2 - 1 + + # the global scale at is the statistic calculated at maximal nearest + # neighbors. Threshold is the maximum on the global and local scales + threshold = max(threshold, stat_mgc_map[m - 1][n - 1]) + + # find the largest connected component of significant correlations + sig_connect = stat_mgc_map > threshold + if np.sum(sig_connect) > 0: + sig_connect, _ = _measurements.label(sig_connect) + _, label_counts = np.unique(sig_connect, return_counts=True) + + # skip the first element in label_counts, as it is count(zeros) + max_label = np.argmax(label_counts[1:]) + 1 + sig_connect = sig_connect == max_label + else: + sig_connect = np.array([[False]]) + + return sig_connect + + +def _smooth_mgc_map(sig_connect, stat_mgc_map): + """Finds the smoothed maximal within the significant region R. + + If area of R is too small it returns the last local correlation. Otherwise, + returns the maximum within significant_connected_region. + + Parameters + ---------- + sig_connect : ndarray + A binary matrix with 1's indicating the significant region. + stat_mgc_map : ndarray + All local correlations within ``[-1, 1]``. + + Returns + ------- + stat : float + The sample MGC statistic within ``[-1, 1]``. + opt_scale: (float, float) + The estimated optimal scale as an ``(x, y)`` pair. + + """ + m, n = stat_mgc_map.shape + + # the global scale at is the statistic calculated at maximal nearest + # neighbors. By default, statistic and optimal scale are global. + stat = stat_mgc_map[m - 1][n - 1] + opt_scale = [m, n] + + if np.linalg.norm(sig_connect) != 0: + # proceed only when the connected region's area is sufficiently large + # 0.02 is simply an empirical threshold, this can be set to 0.01 or 0.05 + # with varying levels of performance + if np.sum(sig_connect) >= np.ceil(0.02 * max(m, n)) * min(m, n): + max_corr = max(stat_mgc_map[sig_connect]) + + # find all scales within significant_connected_region that maximize + # the local correlation + max_corr_index = np.where((stat_mgc_map >= max_corr) & sig_connect) + + if max_corr >= stat: + stat = max_corr + + k, l = max_corr_index + one_d_indices = k * n + l # 2D to 1D indexing + k = np.max(one_d_indices) // n + l = np.max(one_d_indices) % n + opt_scale = [k+1, l+1] # adding 1s to match R indexing + + return stat, opt_scale + + +def _two_sample_transform(u, v): + """Helper function that concatenates x and y for two sample MGC stat. + + See above for use. + + Parameters + ---------- + u, v : ndarray + `u` and `v` have shapes ``(n, p)`` and ``(m, p)``. + + Returns + ------- + x : ndarray + Concatenate `u` and `v` along the ``axis = 0``. `x` thus has shape + ``(2n, p)``. + y : ndarray + Label matrix for `x` where 0 refers to samples that comes from `u` and + 1 refers to samples that come from `v`. `y` thus has shape ``(2n, 1)``. + + """ + nx = u.shape[0] + ny = v.shape[0] + x = np.concatenate([u, v], axis=0) + y = np.concatenate([np.zeros(nx), np.ones(ny)], axis=0).reshape(-1, 1) + return x, y diff --git a/v12x/pylib/scipy/stats/_morestats.py b/v12x/pylib/scipy/stats/_morestats.py new file mode 100644 index 0000000000000000000000000000000000000000..3ead2e29323cd51229e4e5bb018d8bf52bd87fb6 --- /dev/null +++ b/v12x/pylib/scipy/stats/_morestats.py @@ -0,0 +1,5086 @@ +import itertools +import math +import warnings +import threading +from collections import namedtuple + +import numpy as np +from numpy import (isscalar, log, around, arange, sort, amin, amax, sqrt, array, + pi, exp, ravel, count_nonzero) + +from scipy import optimize, special, interpolate, stats +from scipy._lib._bunch import _make_tuple_bunch +from scipy._lib._util import _rename_parameter, _contains_nan, _get_nan +from scipy._lib.deprecation import _NoValue +import scipy._external.array_api_extra as xpx + +from scipy._lib._array_api import ( + array_namespace, + is_marray, + xp_capabilities, + is_numpy, + is_jax, + is_dask, + xp_size, + xp_vector_norm, + xp_promote, + xp_result_type, + xp_device, + xp_ravel, + _count_nonmasked, + is_lazy_array, +) + +from ._ansari_swilk_statistics import gscale +from . import _stats_py, _wilcoxon +from ._fit import FitResult +from ._stats_py import (_get_pvalue, SignificanceResult, # noqa:F401 + _SimpleNormal, _SimpleChi2, _SimpleF, _demean) +from .contingency import chi2_contingency +from . import distributions +from ._distn_infrastructure import rv_generic +from ._axis_nan_policy import _axis_nan_policy_factory, _broadcast_arrays + + +__all__ = ['mvsdist', + 'bayes_mvs', 'kstat', 'kstatvar', 'probplot', 'ppcc_max', 'ppcc_plot', + 'boxcox_llf', 'boxcox', 'boxcox_normmax', 'boxcox_normplot', + 'shapiro', 'anderson', 'ansari', 'bartlett', 'levene', + 'fligner', 'mood', 'wilcoxon', 'median_test', + 'circmean', 'circvar', 'circstd', 'anderson_ksamp', + 'yeojohnson_llf', 'yeojohnson', 'yeojohnson_normmax', + 'yeojohnson_normplot', 'directional_stats', + 'false_discovery_control' + ] + + +Mean = namedtuple('Mean', ('statistic', 'minmax')) +Variance = namedtuple('Variance', ('statistic', 'minmax')) +Std_dev = namedtuple('Std_dev', ('statistic', 'minmax')) + + +@xp_capabilities(np_only=True) +def bayes_mvs(data, alpha=0.90): + r""" + Bayesian confidence intervals for the mean, var, and std. + + Parameters + ---------- + data : array_like + Input data, if multi-dimensional it is flattened to 1-D by `bayes_mvs`. + Requires 2 or more data points. + alpha : float, optional + Probability that the returned confidence interval contains + the true parameter. + + Returns + ------- + mean_cntr, var_cntr, std_cntr : tuple + The three results are for the mean, variance and standard deviation, + respectively. Each result is a tuple of the form:: + + (center, (lower, upper)) + + with ``center`` the mean of the conditional pdf of the value given the + data, and ``(lower, upper)`` a confidence interval, centered on the + median, containing the estimate to a probability ``alpha``. + + See Also + -------- + mvsdist + + Notes + ----- + Each tuple of mean, variance, and standard deviation estimates represent + the (center, (lower, upper)) with center the mean of the conditional pdf + of the value given the data and (lower, upper) is a confidence interval + centered on the median, containing the estimate to a probability + ``alpha``. + + Converts data to 1-D and assumes all data has the same mean and variance. + Uses Jeffrey's prior for variance and std. + + Equivalent to ``tuple((x.mean(), x.interval(alpha)) for x in mvsdist(dat))`` + + References + ---------- + T.E. Oliphant, "A Bayesian perspective on estimating mean, variance, and + standard-deviation from data", https://scholarsarchive.byu.edu/facpub/278, + 2006. + + Examples + -------- + First a basic example to demonstrate the outputs: + + >>> from scipy import stats + >>> data = [6, 9, 12, 7, 8, 8, 13] + >>> mean, var, std = stats.bayes_mvs(data) + >>> mean + Mean(statistic=9.0, minmax=(7.103650222612533, 10.896349777387467)) + >>> var + Variance(statistic=10.0, minmax=(3.176724206, 24.45910382)) + >>> std + Std_dev(statistic=2.9724954732045084, + minmax=(1.7823367265645143, 4.945614605014631)) + + Now we generate some normally distributed random data, and get estimates of + mean and standard deviation with 95% confidence intervals for those + estimates: + + >>> n_samples = 100000 + >>> data = stats.norm.rvs(size=n_samples) + >>> res_mean, res_var, res_std = stats.bayes_mvs(data, alpha=0.95) + + >>> import matplotlib.pyplot as plt + >>> fig = plt.figure() + >>> ax = fig.add_subplot(111) + >>> ax.hist(data, bins=100, density=True, label='Histogram of data') + >>> ax.vlines(res_mean.statistic, 0, 0.5, colors='r', label='Estimated mean') + >>> ax.axvspan(res_mean.minmax[0],res_mean.minmax[1], facecolor='r', + ... alpha=0.2, label=r'Estimated mean (95% limits)') + >>> ax.vlines(res_std.statistic, 0, 0.5, colors='g', label='Estimated scale') + >>> ax.axvspan(res_std.minmax[0],res_std.minmax[1], facecolor='g', alpha=0.2, + ... label=r'Estimated scale (95% limits)') + + >>> ax.legend(fontsize=10) + >>> ax.set_xlim([-4, 4]) + >>> ax.set_ylim([0, 0.5]) + >>> plt.show() + + """ + m, v, s = mvsdist(data) + if alpha >= 1 or alpha <= 0: + raise ValueError(f"0 < alpha < 1 is required, but {alpha=} was given.") + + m_res = Mean(m.mean(), m.interval(alpha)) + v_res = Variance(v.mean(), v.interval(alpha)) + s_res = Std_dev(s.mean(), s.interval(alpha)) + + return m_res, v_res, s_res + + +@xp_capabilities(np_only=True) +def mvsdist(data): + """ + 'Frozen' distributions for mean, variance, and standard deviation of data. + + Parameters + ---------- + data : array_like + Input array. Converted to 1-D using ravel. + Requires 2 or more data-points. + + Returns + ------- + mdist : "frozen" distribution object + Distribution object representing the mean of the data. + vdist : "frozen" distribution object + Distribution object representing the variance of the data. + sdist : "frozen" distribution object + Distribution object representing the standard deviation of the data. + + See Also + -------- + bayes_mvs + + Notes + ----- + The return values from ``bayes_mvs(data)`` is equivalent to + ``tuple((x.mean(), x.interval(0.90)) for x in mvsdist(data))``. + + In other words, calling ``.mean()`` and ``.interval(0.90)`` + on the three distribution objects returned from this function will give + the same results that are returned from `bayes_mvs`. + + References + ---------- + T.E. Oliphant, "A Bayesian perspective on estimating mean, variance, and + standard-deviation from data", https://scholarsarchive.byu.edu/facpub/278, + 2006. + + Examples + -------- + >>> from scipy import stats + >>> data = [6, 9, 12, 7, 8, 8, 13] + >>> mean, var, std = stats.mvsdist(data) + + We now have frozen distribution objects "mean", "var" and "std" that we can + examine: + + >>> mean.mean() + 9.0 + >>> mean.interval(0.95) + (6.6120585482655692, 11.387941451734431) + >>> mean.std() + 1.1952286093343936 + + """ + x = ravel(data) + n = len(x) + if n < 2: + raise ValueError("Need at least 2 data-points.") + xbar = x.mean() + C = x.var() + if n > 1000: # gaussian approximations for large n + mdist = distributions.norm(loc=xbar, scale=math.sqrt(C / n)) + sdist = distributions.norm(loc=math.sqrt(C), scale=math.sqrt(C / (2. * n))) + vdist = distributions.norm(loc=C, scale=math.sqrt(2.0 / n) * C) + else: + nm1 = n - 1 + fac = n * C / 2. + val = nm1 / 2. + mdist = distributions.t(nm1, loc=xbar, scale=math.sqrt(C / nm1)) + sdist = distributions.gengamma(val, -2, scale=math.sqrt(fac)) + vdist = distributions.invgamma(val, scale=fac) + return mdist, vdist, sdist + + +@xp_capabilities(marray=True) +@_axis_nan_policy_factory( + lambda x: x, result_to_tuple=lambda x, _: (x,), n_outputs=1, default_axis=None +) +def kstat(data, n=2, *, axis=None): + r""" + Return the `n` th k-statistic ( ``1<=n<=4`` so far). + + The `n` th k-statistic ``k_n`` is the unique symmetric unbiased estimator of the + `n` th cumulant :math:`\kappa_n` [1]_ [2]_. + + Parameters + ---------- + data : array_like + Input array. + n : int, {1, 2, 3, 4}, optional + Default is equal to 2. + axis : int or None, default: None + If an int, the axis of the input along which to compute the statistic. + The statistic of each axis-slice (e.g. row) of the input will appear + in a corresponding element of the output. If ``None``, the input will + be raveled before computing the statistic. + + Returns + ------- + kstat : float + The `n` th k-statistic. + + See Also + -------- + kstatvar : Returns an unbiased estimator of the variance of the k-statistic + moment : Returns the n-th central moment about the mean for a sample. + + Notes + ----- + For a sample size :math:`n`, the first few k-statistics are given by + + .. math:: + + k_1 &= \frac{S_1}{n}, \\ + k_2 &= \frac{nS_2 - S_1^2}{n(n-1)}, \\ + k_3 &= \frac{2S_1^3 - 3nS_1S_2 + n^2S_3}{n(n-1)(n-2)}, \\ + k_4 &= \frac{-6S_1^4 + 12nS_1^2S_2 - 3n(n-1)S_2^2 - 4n(n+1)S_1S_3 + + n^2(n+1)S_4}{n (n-1)(n-2)(n-3)}, + + where + + .. math:: + + S_r \equiv \sum_{i=1}^n X_i^r, + + and :math:`X_i` is the :math:`i` th data point. + + References + ---------- + .. [1] http://mathworld.wolfram.com/k-Statistic.html + + .. [2] http://mathworld.wolfram.com/Cumulant.html + + Examples + -------- + >>> from scipy import stats + >>> from numpy.random import default_rng + >>> rng = default_rng() + + As sample size increases, `n`-th moment and `n`-th k-statistic converge to the + same number (although they aren't identical). In the case of the normal + distribution, they converge to zero. + + >>> for i in range(2,8): + ... x = rng.normal(size=10**i) + ... m, k = stats.moment(x, 3), stats.kstat(x, 3) + ... print(f"{i=}: {m=:.3g}, {k=:.3g}, {(m-k)=:.3g}") + i=2: m=-0.631, k=-0.651, (m-k)=0.0194 # random + i=3: m=0.0282, k=0.0283, (m-k)=-8.49e-05 + i=4: m=-0.0454, k=-0.0454, (m-k)=1.36e-05 + i=6: m=7.53e-05, k=7.53e-05, (m-k)=-2.26e-09 + i=7: m=0.00166, k=0.00166, (m-k)=-4.99e-09 + i=8: m=-2.88e-06 k=-2.88e-06, (m-k)=8.63e-13 + """ + xp = array_namespace(data) + data = xp.asarray(data) + if n > 4 or n < 1: + raise ValueError("k-statistics only supported for 1<=n<=4") + n = int(n) + if axis is None: + data = xp.reshape(data, (-1,)) + axis = 0 + + N = _count_nonmasked(data, axis, xp=xp) + + S = [None] + [xp.sum(data**k, axis=axis) for k in range(1, n + 1)] + if n == 1: + return S[1] * 1.0/N + elif n == 2: + return (N*S[2] - S[1]**2.0) / (N*(N - 1.0)) + elif n == 3: + return (2*S[1]**3 - 3*N*S[1]*S[2] + N*N*S[3]) / (N*(N - 1.0)*(N - 2.0)) + elif n == 4: + return ((-6*S[1]**4 + 12*N*S[1]**2 * S[2] - 3*N*(N-1.0)*S[2]**2 - + 4*N*(N+1)*S[1]*S[3] + N*N*(N+1)*S[4]) / + (N*(N-1.0)*(N-2.0)*(N-3.0))) + else: + raise ValueError("Should not be here.") + + +@xp_capabilities(marray=True) +@_axis_nan_policy_factory( + lambda x: x, result_to_tuple=lambda x, _: (x,), n_outputs=1, default_axis=None +) +def kstatvar(data, n=2, *, axis=None): + r"""Return an unbiased estimator of the variance of the k-statistic. + + See `kstat` and [1]_ for more details about the k-statistic. + + Parameters + ---------- + data : array_like + Input array. + n : int, {1, 2}, optional + Default is equal to 2. + axis : int or None, default: None + If an int, the axis of the input along which to compute the statistic. + The statistic of each axis-slice (e.g. row) of the input will appear + in a corresponding element of the output. If ``None``, the input will + be raveled before computing the statistic. + + Returns + ------- + kstatvar : float + The `n` th k-statistic variance. + + See Also + -------- + kstat : Returns the n-th k-statistic. + moment : Returns the n-th central moment about the mean for a sample. + + Notes + ----- + Unbiased estimators of the variances of the first two k-statistics are given by + + .. math:: + + \mathrm{var}(k_1) &= \frac{k_2}{n}, \\ + \mathrm{var}(k_2) &= \frac{2k_2^2n + (n-1)k_4}{n(n + 1)}. + + References + ---------- + .. [1] http://mathworld.wolfram.com/k-Statistic.html + + Examples + -------- + >>> import numpy as np + >>> from scipy import stats + >>> rng = np.random.default_rng(92366746) + + As the sample size increases, the estimated variance of the k-statistic converges + to zero. + + >>> for n in np.astype(np.logspace(1, 6, 6), int): + ... x = rng.normal(size=n) + ... kvar = stats.kstatvar(x, 1) + ... print(f"{n=:<8}: {kvar=:.3g}") + n=10 : kvar=0.0954 + n=100 : kvar=0.00974 + n=1000 : kvar=0.000962 + n=10000 : kvar=0.0001 + n=100000 : kvar=9.94e-06 + n=1000000 : kvar=9.99e-07 + """ # noqa: E501 + xp = array_namespace(data) + data = xp.asarray(data) + if axis is None: + data = xp.reshape(data, (-1,)) + axis = 0 + N = _count_nonmasked(data, axis, xp=xp) + + if n == 1: + return kstat(data, n=2, axis=axis, _no_deco=True) * 1.0/N + elif n == 2: + k2 = kstat(data, n=2, axis=axis, _no_deco=True) + k4 = kstat(data, n=4, axis=axis, _no_deco=True) + return (2*N*k2**2 + (N-1)*k4) / (N*(N+1)) + else: + raise ValueError("Only n=1 or n=2 supported.") + + +def _calc_uniform_order_statistic_medians(n): + """Approximations of uniform order statistic medians. + + Parameters + ---------- + n : int + Sample size. + + Returns + ------- + v : 1d float array + Approximations of the order statistic medians. + + References + ---------- + .. [1] James J. Filliben, "The Probability Plot Correlation Coefficient + Test for Normality", Technometrics, Vol. 17, pp. 111-117, 1975. + + Examples + -------- + Order statistics of the uniform distribution on the unit interval + are marginally distributed according to beta distributions. + The expectations of these order statistic are evenly spaced across + the interval, but the distributions are skewed in a way that + pushes the medians slightly towards the endpoints of the unit interval: + + >>> import numpy as np + >>> n = 4 + >>> k = np.arange(1, n+1) + >>> from scipy.stats import beta + >>> a = k + >>> b = n-k+1 + >>> beta.mean(a, b) + array([0.2, 0.4, 0.6, 0.8]) + >>> beta.median(a, b) + array([0.15910358, 0.38572757, 0.61427243, 0.84089642]) + + The Filliben approximation uses the exact medians of the smallest + and greatest order statistics, and the remaining medians are approximated + by points spread evenly across a sub-interval of the unit interval: + + >>> from scipy.stats._morestats import _calc_uniform_order_statistic_medians + >>> _calc_uniform_order_statistic_medians(n) + array([0.15910358, 0.38545246, 0.61454754, 0.84089642]) + + This plot shows the skewed distributions of the order statistics + of a sample of size four from a uniform distribution on the unit interval: + + >>> import matplotlib.pyplot as plt + >>> x = np.linspace(0.0, 1.0, num=50, endpoint=True) + >>> pdfs = [beta.pdf(x, a[i], b[i]) for i in range(n)] + >>> plt.figure() + >>> plt.plot(x, pdfs[0], x, pdfs[1], x, pdfs[2], x, pdfs[3]) + + """ + v = np.empty(n, dtype=np.float64) + v[-1] = 0.5**(1.0 / n) + v[0] = 1 - v[-1] + i = np.arange(2, n) + v[1:-1] = (i - 0.3175) / (n + 0.365) + return v + + +def _parse_dist_kw(dist, enforce_subclass=True): + """Parse `dist` keyword. + + Parameters + ---------- + dist : str or stats.distributions instance. + Several functions take `dist` as a keyword, hence this utility + function. + enforce_subclass : bool, optional + If True (default), `dist` needs to be a + `_distn_infrastructure.rv_generic` instance. + It can sometimes be useful to set this keyword to False, if a function + wants to accept objects that just look somewhat like such an instance + (for example, they have a ``ppf`` method). + + """ + if isinstance(dist, rv_generic): + pass + elif isinstance(dist, str): + try: + dist = getattr(distributions, dist) + except AttributeError as e: + raise ValueError(f"{dist} is not a valid distribution name") from e + elif enforce_subclass: + msg = ("`dist` should be a stats.distributions instance or a string " + "with the name of such a distribution.") + raise ValueError(msg) + + return dist + + +def _add_axis_labels_title(plot, xlabel, ylabel, title): + """Helper function to add axes labels and a title to stats plots.""" + try: + if hasattr(plot, 'set_title'): + # Matplotlib Axes instance or something that looks like it + plot.set_title(title) + plot.set_xlabel(xlabel) + plot.set_ylabel(ylabel) + else: + # matplotlib.pyplot module + plot.title(title) + plot.xlabel(xlabel) + plot.ylabel(ylabel) + except Exception: + # Not an MPL object or something that looks (enough) like it. + # Don't crash on adding labels or title + pass + + +@xp_capabilities(np_only=True) +def probplot(x, sparams=(), dist='norm', fit=True, plot=None, rvalue=False): + """ + Calculate quantiles for a probability plot, and optionally show the plot. + + Generates a probability plot of sample data against the quantiles of a + specified theoretical distribution (the normal distribution by default). + `probplot` optionally calculates a best-fit line for the data and plots the + results using Matplotlib or a given plot function. + + Parameters + ---------- + x : array_like + Sample/response data from which `probplot` creates the plot. + sparams : tuple, optional + Distribution-specific shape parameters (shape parameters plus location + and scale). + dist : str or stats.distributions instance, optional + Distribution or distribution function name. The default is 'norm' for a + normal probability plot. Objects that look enough like a + stats.distributions instance (i.e. they have a ``ppf`` method) are also + accepted. + fit : bool, optional + Fit a least-squares regression (best-fit) line to the sample data if + True (default). + plot : object, optional + If given, plots the quantiles. + If given and `fit` is True, also plots the least squares fit. + `plot` is an object that has to have methods "plot" and "text". + The `matplotlib.pyplot` module or a Matplotlib Axes object can be used, + or a custom object with the same methods. + Default is None, which means that no plot is created. + rvalue : bool, optional + If `plot` is provided and `fit` is True, setting `rvalue` to True + includes the coefficient of determination on the plot. + Default is False. + + Returns + ------- + (osm, osr) : tuple of ndarrays + Tuple of theoretical quantiles (osm, or order statistic medians) and + ordered responses (osr). `osr` is simply sorted input `x`. + For details on how `osm` is calculated see the Notes section. + (slope, intercept, r) : tuple of floats, optional + Tuple containing the result of the least-squares fit, if that is + performed by `probplot`. `r` is the square root of the coefficient of + determination. If ``fit=False`` and ``plot=None``, this tuple is not + returned. + + Notes + ----- + Even if `plot` is given, the figure is not shown or saved by `probplot`; + ``plt.show()`` or ``plt.savefig('figname.png')`` should be used after + calling `probplot`. + + `probplot` generates a probability plot, which should not be confused with + a Q-Q or a P-P plot. Statsmodels has more extensive functionality of this + type, see ``statsmodels.api.ProbPlot``. + + The formula used for the theoretical quantiles (horizontal axis of the + probability plot) is Filliben's estimate:: + + quantiles = dist.ppf(val), for + + 0.5**(1/n), for i = n + val = (i - 0.3175) / (n + 0.365), for i = 2, ..., n-1 + 1 - 0.5**(1/n), for i = 1 + + where ``i`` indicates the i-th ordered value and ``n`` is the total number + of values. + + Examples + -------- + >>> import numpy as np + >>> from scipy import stats + >>> import matplotlib.pyplot as plt + >>> nsample = 100 + >>> rng = np.random.default_rng() + + A t distribution with small degrees of freedom: + + >>> ax1 = plt.subplot(221) + >>> x = stats.t.rvs(3, size=nsample, random_state=rng) + >>> res = stats.probplot(x, plot=plt) + + A t distribution with larger degrees of freedom: + + >>> ax2 = plt.subplot(222) + >>> x = stats.t.rvs(25, size=nsample, random_state=rng) + >>> res = stats.probplot(x, plot=plt) + + A mixture of two normal distributions with broadcasting: + + >>> ax3 = plt.subplot(223) + >>> x = stats.norm.rvs(loc=[0,5], scale=[1,1.5], + ... size=(nsample//2,2), random_state=rng).ravel() + >>> res = stats.probplot(x, plot=plt) + + A standard normal distribution: + + >>> ax4 = plt.subplot(224) + >>> x = stats.norm.rvs(loc=0, scale=1, size=nsample, random_state=rng) + >>> res = stats.probplot(x, plot=plt) + + Produce a new figure with a loggamma distribution, using the ``dist`` and + ``sparams`` keywords: + + >>> fig = plt.figure() + >>> ax = fig.add_subplot(111) + >>> x = stats.loggamma.rvs(c=2.5, size=500, random_state=rng) + >>> res = stats.probplot(x, dist=stats.loggamma, sparams=(2.5,), plot=ax) + >>> ax.set_title("Probplot for loggamma dist with shape parameter 2.5") + + Show the results with Matplotlib: + + >>> plt.show() + + """ + x = np.asarray(x) + if x.size == 0: + if fit: + return (x, x), (np.nan, np.nan, 0.0) + else: + return x, x + + osm_uniform = _calc_uniform_order_statistic_medians(len(x)) + dist = _parse_dist_kw(dist, enforce_subclass=False) + if sparams is None: + sparams = () + if isscalar(sparams): + sparams = (sparams,) + if not isinstance(sparams, tuple): + sparams = tuple(sparams) + + osm = dist.ppf(osm_uniform, *sparams) + osr = sort(x) + if fit: + # perform a linear least squares fit. + slope, intercept, r, prob, _ = _stats_py.linregress(osm, osr) + + if plot is not None: + plot.plot(osm, osr, 'bo') + if fit: + plot.plot(osm, slope*osm + intercept, 'r-') + _add_axis_labels_title(plot, xlabel='Theoretical quantiles', + ylabel='Ordered Values', + title='Probability Plot') + + # Add R^2 value to the plot as text + if fit and rvalue: + xmin = amin(osm) + xmax = amax(osm) + ymin = amin(x) + ymax = amax(x) + posx = xmin + 0.70 * (xmax - xmin) + posy = ymin + 0.01 * (ymax - ymin) + plot.text(posx, posy, f"$R^2={r ** 2:1.4f}$") + + if fit: + return (osm, osr), (slope, intercept, r) + else: + return osm, osr + + +@xp_capabilities(np_only=True) +def ppcc_max(x, brack=(0.0, 1.0), dist='tukeylambda'): + """Calculate the shape parameter that maximizes the PPCC. + + The probability plot correlation coefficient (PPCC) plot can be used + to determine the optimal shape parameter for a one-parameter family + of distributions. ``ppcc_max`` returns the shape parameter that would + maximize the probability plot correlation coefficient for the given + data to a one-parameter family of distributions. + + Parameters + ---------- + x : array_like + Input array. + brack : tuple, optional + Triple (a,b,c) where (a>> import numpy as np + >>> from scipy import stats + >>> import matplotlib.pyplot as plt + >>> rng = np.random.default_rng() + >>> c = 2.5 + >>> x = stats.weibull_min.rvs(c, scale=4, size=2000, random_state=rng) + + Generate the PPCC plot for this data with the Weibull distribution. + + >>> fig, ax = plt.subplots(figsize=(8, 6)) + >>> res = stats.ppcc_plot(x, c/2, 2*c, dist='weibull_min', plot=ax) + + We calculate the value where the shape should reach its maximum and a + red line is drawn there. The line should coincide with the highest + point in the PPCC graph. + + >>> cmax = stats.ppcc_max(x, brack=(c/2, 2*c), dist='weibull_min') + >>> ax.axvline(cmax, color='r') + >>> plt.show() + + """ + dist = _parse_dist_kw(dist) + osm_uniform = _calc_uniform_order_statistic_medians(len(x)) + osr = sort(x) + + # this function computes the x-axis values of the probability plot + # and computes a linear regression (including the correlation) + # and returns 1-r so that a minimization function maximizes the + # correlation + def tempfunc(shape, mi, yvals, func): + xvals = func(mi, shape) + r, prob = _stats_py.pearsonr(xvals, yvals) + return 1 - r + + return optimize.brent(tempfunc, brack=brack, + args=(osm_uniform, osr, dist.ppf)) + + +@xp_capabilities(np_only=True) +def ppcc_plot(x, a, b, dist='tukeylambda', plot=None, N=80): + """Calculate and optionally plot probability plot correlation coefficient. + + The probability plot correlation coefficient (PPCC) plot can be used to + determine the optimal shape parameter for a one-parameter family of + distributions. It cannot be used for distributions without shape + parameters + (like the normal distribution) or with multiple shape parameters. + + By default a Tukey-Lambda distribution (`stats.tukeylambda`) is used. A + Tukey-Lambda PPCC plot interpolates from long-tailed to short-tailed + distributions via an approximately normal one, and is therefore + particularly useful in practice. + + Parameters + ---------- + x : array_like + Input array. + a, b : scalar + Lower and upper bounds of the shape parameter to use. + dist : str or stats.distributions instance, optional + Distribution or distribution function name. Objects that look enough + like a stats.distributions instance (i.e. they have a ``ppf`` method) + are also accepted. The default is ``'tukeylambda'``. + plot : object, optional + If given, plots PPCC against the shape parameter. + `plot` is an object that has to have methods "plot" and "text". + The `matplotlib.pyplot` module or a Matplotlib Axes object can be used, + or a custom object with the same methods. + Default is None, which means that no plot is created. + N : int, optional + Number of points on the horizontal axis (equally distributed from + `a` to `b`). + + Returns + ------- + svals : ndarray + The shape values for which `ppcc` was calculated. + ppcc : ndarray + The calculated probability plot correlation coefficient values. + + See Also + -------- + ppcc_max, probplot, boxcox_normplot, tukeylambda + + References + ---------- + J.J. Filliben, "The Probability Plot Correlation Coefficient Test for + Normality", Technometrics, Vol. 17, pp. 111-117, 1975. + + Examples + -------- + First we generate some random data from a Weibull distribution + with shape parameter 2.5, and plot the histogram of the data: + + >>> import numpy as np + >>> from scipy import stats + >>> import matplotlib.pyplot as plt + >>> rng = np.random.default_rng() + >>> c = 2.5 + >>> x = stats.weibull_min.rvs(c, scale=4, size=2000, random_state=rng) + + Take a look at the histogram of the data. + + >>> fig1, ax = plt.subplots(figsize=(9, 4)) + >>> ax.hist(x, bins=50) + >>> ax.set_title('Histogram of x') + >>> plt.show() + + Now we explore this data with a PPCC plot as well as the related + probability plot and Box-Cox normplot. A red line is drawn where we + expect the PPCC value to be maximal (at the shape parameter ``c`` + used above): + + >>> fig2 = plt.figure(figsize=(12, 4)) + >>> ax1 = fig2.add_subplot(1, 3, 1) + >>> ax2 = fig2.add_subplot(1, 3, 2) + >>> ax3 = fig2.add_subplot(1, 3, 3) + >>> res = stats.probplot(x, plot=ax1) + >>> res = stats.boxcox_normplot(x, -4, 4, plot=ax2) + >>> res = stats.ppcc_plot(x, c/2, 2*c, dist='weibull_min', plot=ax3) + >>> ax3.axvline(c, color='r') + >>> plt.show() + + """ + if b <= a: + raise ValueError("`b` has to be larger than `a`.") + + svals = np.linspace(a, b, num=N) + ppcc = np.empty_like(svals) + for k, sval in enumerate(svals): + _, r2 = probplot(x, sval, dist=dist, fit=True) + ppcc[k] = r2[-1] + + if plot is not None: + plot.plot(svals, ppcc, 'x') + _add_axis_labels_title(plot, xlabel='Shape Values', + ylabel='Prob Plot Corr. Coef.', + title=f'({dist}) PPCC Plot') + + return svals, ppcc + + +def _log_mean(logx, axis): + # compute log of mean of x from log(x) + return ( + special.logsumexp(logx, axis=axis, keepdims=True) + - math.log(logx.shape[axis]) + ) + + +def _log_var(logx, xp, axis): + # compute log of variance of x from log(x) + logmean = xp.broadcast_to(_log_mean(logx, axis=axis), logx.shape) + ones = xp.ones_like(logx) + logxmu, _ = special.logsumexp(xp.stack((logx, logmean), axis=0), axis=0, + b=xp.stack((ones, -ones), axis=0), return_sign=True) + return special.logsumexp(2 * logxmu, axis=axis) - math.log(logx.shape[axis]) + + +@xp_capabilities() +def boxcox_llf(lmb, data, *, axis=0, keepdims=False, nan_policy='propagate'): + r"""The boxcox log-likelihood function. + + Parameters + ---------- + lmb : scalar + Parameter for Box-Cox transformation. See `boxcox` for details. + data : array_like + Data to calculate Box-Cox log-likelihood for. If `data` is + multi-dimensional, the log-likelihood is calculated along the first + axis. + axis : int, default: 0 + If an int, the axis of the input along which to compute the statistic. + The statistic of each axis-slice (e.g. row) of the input will appear in a + corresponding element of the output. + If ``None``, the input will be raveled before computing the statistic. + keepdims : bool, default: False + If this is set to True, the axes which are reduced are left + in the result as dimensions with size one. With this option, + the result will broadcast correctly against the input array. + nan_policy : {'propagate', 'omit', 'raise'} + Defines how to handle input NaNs. + + - ``propagate``: if a NaN is present in the axis slice (e.g. row) along + which the statistic is computed, the corresponding entry of the output + will be NaN. + - ``omit``: NaNs will be omitted when performing the calculation. + If insufficient data remains in the axis slice along which the + statistic is computed, the corresponding entry of the output will be + NaN. + - ``raise``: if a NaN is present, a ``ValueError`` will be raised. + + Returns + ------- + llf : float or ndarray + Box-Cox log-likelihood of `data` given `lmb`. A float for 1-D `data`, + an array otherwise. + + See Also + -------- + boxcox, probplot, boxcox_normplot, boxcox_normmax + + Notes + ----- + The Box-Cox log-likelihood function :math:`l` is defined here as + + .. math:: + + l = (\lambda - 1) \sum_i^N \log(x_i) - + \frac{N}{2} \log\left(\sum_i^N (y_i - \bar{y})^2 / N\right), + + where :math:`N` is the number of data points ``data`` and :math:`y` is the Box-Cox + transformed input data. + This corresponds to the *profile log-likelihood* of the original data :math:`x` + with some constant terms dropped. + + Examples + -------- + >>> import numpy as np + >>> from scipy import stats + >>> import matplotlib.pyplot as plt + >>> from mpl_toolkits.axes_grid1.inset_locator import inset_axes + + Generate some random variates and calculate Box-Cox log-likelihood values + for them for a range of ``lmbda`` values: + + >>> rng = np.random.default_rng() + >>> x = stats.loggamma.rvs(5, loc=10, size=1000, random_state=rng) + >>> lmbdas = np.linspace(-2, 10) + >>> llf = np.zeros(lmbdas.shape, dtype=float) + >>> for ii, lmbda in enumerate(lmbdas): + ... llf[ii] = stats.boxcox_llf(lmbda, x) + + Also find the optimal lmbda value with `boxcox`: + + >>> x_most_normal, lmbda_optimal = stats.boxcox(x) + + Plot the log-likelihood as function of lmbda. Add the optimal lmbda as a + horizontal line to check that that's really the optimum: + + >>> fig = plt.figure() + >>> ax = fig.add_subplot(111) + >>> ax.plot(lmbdas, llf, 'b.-') + >>> ax.axhline(stats.boxcox_llf(lmbda_optimal, x), color='r') + >>> ax.set_xlabel('lmbda parameter') + >>> ax.set_ylabel('Box-Cox log-likelihood') + + Now add some probability plots to show that where the log-likelihood is + maximized the data transformed with `boxcox` looks closest to normal: + + >>> locs = [3, 10, 4] # 'lower left', 'center', 'lower right' + >>> for lmbda, loc in zip([-1, lmbda_optimal, 9], locs): + ... xt = stats.boxcox(x, lmbda=lmbda) + ... (osm, osr), (slope, intercept, r_sq) = stats.probplot(xt) + ... ax_inset = inset_axes(ax, width="20%", height="20%", loc=loc) + ... ax_inset.plot(osm, osr, 'c.', osm, slope*osm + intercept, 'k-') + ... ax_inset.set_xticklabels([]) + ... ax_inset.set_yticklabels([]) + ... ax_inset.set_title(r'$\lambda=%1.2f$' % lmbda) + + >>> plt.show() + + """ + # _axis_nan_policy decorator does not currently support these for lazy arrays. + # We want to run tests with lazy backends, so don't pass the arguments explicitly + # unless necessary. + kwargs = {} + if keepdims is not False: + kwargs['keepdims'] = keepdims + if nan_policy != 'propagate': + kwargs['nan_policy'] = nan_policy + return _boxcox_llf(data, lmb=lmb, axis=axis, **kwargs) + + +@_axis_nan_policy_factory(lambda x: x, n_outputs=1, default_axis=0, + result_to_tuple=lambda x, _: (x,)) +def _boxcox_llf(data, axis=0, *, lmb): + xp = array_namespace(data) + dtype = xp_result_type(lmb, data, force_floating=True, xp=xp) + data = xp.asarray(data, dtype=dtype) + N = data.shape[axis] + if N == 0: + return _get_nan(data, xp=xp) + + logdata = xp.log(data) + + # Compute the variance of the transformed data. + if lmb == 0: + logvar = xp.log(xp.var(logdata, axis=axis)) + else: + # Transform without the constant offset 1/lmb. The offset does + # not affect the variance, and the subtraction of the offset can + # lead to loss of precision. + # Division by lmb can be factored out to enhance numerical stability. + logx = lmb * logdata + logvar = _log_var(logx, xp, axis) - 2 * math.log(abs(lmb)) + + res = (lmb - 1) * xp.sum(logdata, axis=axis) - N/2 * logvar + res = xp.astype(res, data.dtype, copy=False) # compensate for NumPy <2.0 + res = res[()] if res.ndim == 0 else res + return res + + +def _boxcox_conf_interval(x, lmax, alpha): + # Need to find the lambda for which + # f(x,lmbda) >= f(x,lmax) - 0.5*chi^2_alpha;1 + fac = 0.5 * distributions.chi2.ppf(1 - alpha, 1) + target = boxcox_llf(lmax, x) - fac + + def rootfunc(lmbda, data, target): + return boxcox_llf(lmbda, data) - target + + # Find positive endpoint of interval in which answer is to be found + newlm = lmax + 0.5 + N = 0 + while (rootfunc(newlm, x, target) > 0.0) and (N < 500): + newlm += 0.1 + N += 1 + + if N == 500: + raise RuntimeError("Could not find endpoint.") + + lmplus = optimize.brentq(rootfunc, lmax, newlm, args=(x, target)) + + # Now find negative interval in the same way + newlm = lmax - 0.5 + N = 0 + while (rootfunc(newlm, x, target) > 0.0) and (N < 500): + newlm -= 0.1 + N += 1 + + if N == 500: + raise RuntimeError("Could not find endpoint.") + + lmminus = optimize.brentq(rootfunc, newlm, lmax, args=(x, target)) + return lmminus, lmplus + + +@xp_capabilities(np_only=True) +def boxcox(x, lmbda=None, alpha=None, optimizer=None, *, nan_policy='propagate'): + r"""Return a dataset transformed by a Box-Cox power transformation. + + Parameters + ---------- + x : ndarray + Input array to be transformed. + + If `lmbda` is not None, this is an alias of + `scipy.special.boxcox`. + Returns nan if ``x < 0``; returns -inf if ``x == 0 and lmbda < 0``. + + If `lmbda` is None, array must be positive, 1-dimensional, and + non-constant. + + lmbda : scalar, optional + If `lmbda` is None (default), find the value of `lmbda` that maximizes + the log-likelihood function and return it as the second output + argument. + + If `lmbda` is not None, do the transformation for that value. + + alpha : float, optional + If `lmbda` is None and `alpha` is not None (default), return the + ``100 * (1-alpha)%`` confidence interval for `lmbda` as the third + output argument. Must be between 0.0 and 1.0. + + If `lmbda` is not None, `alpha` is ignored. + optimizer : callable, optional + If `lmbda` is None, `optimizer` is the scalar optimizer used to find + the value of `lmbda` that minimizes the negative log-likelihood + function. `optimizer` is a callable that accepts one argument: + + fun : callable + The objective function, which evaluates the negative + log-likelihood function at a provided value of `lmbda` + + and returns an object, such as an instance of + `scipy.optimize.OptimizeResult`, which holds the optimal value of + `lmbda` in an attribute `x`. + + See the example in `boxcox_normmax` or the documentation of + `scipy.optimize.minimize_scalar` for more information. + + If `lmbda` is not None, `optimizer` is ignored. + nan_policy : {'propagate', 'omit', 'raise'} + Defines how to handle NaNs in `x`. + + - ``propagate``: if a NaN is present, all outputs will contain NaNs. + - ``omit``: NaNs will be omitted when calculating the optimal `maxlog` and + confidence interval; NaNs in `x` will remain NaNs in the transformed data. + - ``raise``: if a NaN is present, a ``ValueError`` will be raised. + + Returns + ------- + boxcox : ndarray + Box-Cox power transformed array. + maxlog : float, optional + If the `lmbda` parameter is None, the second returned argument is + the `lmbda` that maximizes the log-likelihood function. + (min_ci, max_ci) : tuple of float, optional + If `lmbda` parameter is None and `alpha` is not None, this returned + tuple of floats represents the minimum and maximum confidence limits + given `alpha`. + + See Also + -------- + probplot, boxcox_normplot, boxcox_normmax, boxcox_llf + + Notes + ----- + The Box-Cox transform is given by: + + .. math:: + + y = + \begin{cases} + \frac{x^\lambda - 1}{\lambda}, &\text{for } \lambda \neq 0 \\ + \log(x), &\text{for } \lambda = 0 + \end{cases} + + `boxcox` requires the input data to be positive. Sometimes a Box-Cox + transformation provides a shift parameter to achieve this; `boxcox` does + not. Such a shift parameter is equivalent to adding a positive constant to + `x` before calling `boxcox`. + + The confidence limits returned when `alpha` is provided give the interval + where: + + .. math:: + + l(\hat{\lambda}) - l(\lambda) < \frac{1}{2}\chi^2(1 - \alpha, 1), + + with :math:`l` the log-likelihood function and :math:`\chi^2` the chi-squared + function. + + References + ---------- + G.E.P. Box and D.R. Cox, "An Analysis of Transformations", Journal of the + Royal Statistical Society B, 26, 211-252 (1964). + + Examples + -------- + >>> from scipy import stats + >>> import matplotlib.pyplot as plt + + We generate some random variates from a non-normal distribution and make a + probability plot for it, to show it is non-normal in the tails: + + >>> fig = plt.figure() + >>> ax1 = fig.add_subplot(211) + >>> x = stats.loggamma.rvs(5, size=500) + 5 + >>> prob = stats.probplot(x, dist=stats.norm, plot=ax1) + >>> ax1.set_xlabel('') + >>> ax1.set_title('Probplot against normal distribution') + + We now use `boxcox` to transform the data so it's closest to normal: + + >>> ax2 = fig.add_subplot(212) + >>> xt, _ = stats.boxcox(x) + >>> prob = stats.probplot(xt, dist=stats.norm, plot=ax2) + >>> ax2.set_title('Probplot after Box-Cox transformation') + + >>> plt.show() + + """ + x = np.asarray(x) + + if lmbda is not None: # single transformation + _contains_nan(x, nan_policy, xp_omit_okay=True) # handle nan_policy='raise' + return special.boxcox(x, lmbda) + + if x.ndim != 1: + raise ValueError("Data must be 1-dimensional.") + + if x.size == 0: + return x + + if np.all(x == x[0]): + raise ValueError("Data must not be constant.") + + if np.any(x <= 0): + raise ValueError("Data must be positive.") + + # If lmbda=None, find the lmbda that maximizes the log-likelihood function. + lmax = boxcox_normmax(x, method='mle', optimizer=optimizer, nan_policy=nan_policy) + y = special.boxcox(x, lmax) + + if alpha is None: + return y, lmax + else: + # Find confidence interval + interval = ((lmax, lmax) if np.isnan(lmax) else + _boxcox_conf_interval(x[~np.isnan(x)], lmax, alpha)) + return y, lmax, interval + + +def _boxcox_inv_lmbda(x, y): + # compute lmbda given x and y for Box-Cox transformation + num = special.lambertw(-(x ** (-1 / y)) * np.log(x) / y, k=-1) + return np.real(-num / np.log(x) - 1 / y) + + +class _BigFloat: + def __repr__(self): + return "BIG_FLOAT" + + +_BigFloat_singleton = _BigFloat() + + +@xp_capabilities(np_only=True) +def boxcox_normmax( + x, brack=None, method='pearsonr', optimizer=None, *, ymax=_BigFloat_singleton, + nan_policy='propagate' +): + """Compute optimal Box-Cox transform parameter for input data. + + Parameters + ---------- + x : array_like + Input array. All entries must be positive, finite, real numbers. + brack : 2-tuple, optional, default (-2.0, 2.0) + The starting interval for a downhill bracket search for the default + `optimize.brent` solver. Note that this is in most cases not + critical; the final result is allowed to be outside this bracket. + If `optimizer` is passed, `brack` must be None. + method : str, optional + The method to determine the optimal transform parameter (`boxcox` + ``lmbda`` parameter). Options are: + + 'pearsonr' (default) + Maximizes the Pearson correlation coefficient between + ``y = boxcox(x)`` and the expected values for ``y`` if `x` would be + normally-distributed. + + 'mle' + Maximizes the log-likelihood `boxcox_llf`. This is the method used + in `boxcox`. + + 'all' + Use all optimization methods available, and return all results. + Useful to compare different methods. + optimizer : callable, optional + `optimizer` is a callable that accepts one argument: + + fun : callable + The objective function to be minimized. `fun` accepts one argument, + the Box-Cox transform parameter `lmbda`, and returns the value of + the function (e.g., the negative log-likelihood) at the provided + argument. The job of `optimizer` is to find the value of `lmbda` + that *minimizes* `fun`. + + and returns an object, such as an instance of + `scipy.optimize.OptimizeResult`, which holds the optimal value of + `lmbda` in an attribute `x`. + + See the example below or the documentation of + `scipy.optimize.minimize_scalar` for more information. + ymax : float, optional + The unconstrained optimal transform parameter may cause Box-Cox + transformed data to have extreme magnitude or even overflow. + This parameter constrains MLE optimization such that the magnitude + of the transformed `x` does not exceed `ymax`. The default is + the maximum value of the input dtype. If set to infinity, + `boxcox_normmax` returns the unconstrained optimal lambda. + Ignored when ``method='pearsonr'``. + nan_policy : {'propagate', 'omit', 'raise'} + Defines how to handle input NaNs. + + - ``propagate``: if a NaN is present in the input, the output will be NaN. + - ``omit``: NaNs will be omitted when performing the calculation. + - ``raise``: if a NaN is present, a ``ValueError`` will be raised. + + Returns + ------- + maxlog : float or ndarray + The optimal transform parameter found. An array instead of a scalar + for ``method='all'``. + + See Also + -------- + boxcox, boxcox_llf, boxcox_normplot, scipy.optimize.minimize_scalar + + Examples + -------- + >>> import numpy as np + >>> from scipy import stats + >>> import matplotlib.pyplot as plt + + We can generate some data and determine the optimal ``lmbda`` in various + ways: + + >>> rng = np.random.default_rng() + >>> x = stats.loggamma.rvs(5, size=30, random_state=rng) + 5 + >>> y, lmax_mle = stats.boxcox(x) + >>> lmax_pearsonr = stats.boxcox_normmax(x) + + >>> lmax_mle + 2.217563431465757 + >>> lmax_pearsonr + 2.238318660200961 + >>> stats.boxcox_normmax(x, method='all') + array([2.23831866, 2.21756343]) + + >>> fig = plt.figure() + >>> ax = fig.add_subplot(111) + >>> prob = stats.boxcox_normplot(x, -10, 10, plot=ax) + >>> ax.axvline(lmax_mle, color='r') + >>> ax.axvline(lmax_pearsonr, color='g', ls='--') + + >>> plt.show() + + Alternatively, we can define our own `optimizer` function. Suppose we + are only interested in values of `lmbda` on the interval [6, 7], we + want to use `scipy.optimize.minimize_scalar` with ``method='bounded'``, + and we want to use tighter tolerances when optimizing the log-likelihood + function. To do this, we define a function that accepts positional argument + `fun` and uses `scipy.optimize.minimize_scalar` to minimize `fun` subject + to the provided bounds and tolerances: + + >>> from scipy import optimize + >>> options = {'xatol': 1e-12} # absolute tolerance on `x` + >>> def optimizer(fun): + ... return optimize.minimize_scalar(fun, bounds=(6, 7), + ... method="bounded", options=options) + >>> stats.boxcox_normmax(x, optimizer=optimizer) + 6.000000000 + """ + x = np.asarray(x) + + contains_nan = _contains_nan(x, nan_policy, xp_omit_okay=True) + if contains_nan: + if nan_policy == 'propagate': + NaN = _get_nan(x, xp=np)[()] + return np.asarray([NaN, NaN]) if method == 'all' else NaN + x = x[~np.isnan(x)] # `nan_policy='omit + + if not np.all(np.isfinite(x) & (x >= 0)): + message = ("The `x` argument of `boxcox_normmax` must contain " + "only positive, finite, real numbers.") + raise ValueError(message) + + end_msg = "exceed specified `ymax`." + if ymax is _BigFloat_singleton: + dtype = x.dtype if np.issubdtype(x.dtype, np.floating) else np.float64 + # 10000 is a safety factor because `special.boxcox` overflows prematurely. + ymax = np.finfo(dtype).max / 10000 + end_msg = f"overflow in {dtype}." + elif ymax <= 0: + raise ValueError("`ymax` must be strictly positive") + + # If optimizer is not given, define default 'brent' optimizer. + if optimizer is None: + + # Set default value for `brack`. + if brack is None: + brack = (-2.0, 2.0) + + def _optimizer(func, args): + return optimize.brent(func, args=args, brack=brack) + + # Otherwise check optimizer. + else: + if not callable(optimizer): + raise ValueError("`optimizer` must be a callable") + + if brack is not None: + raise ValueError("`brack` must be None if `optimizer` is given") + + # `optimizer` is expected to return a `OptimizeResult` object, we here + # get the solution to the optimization problem. + def _optimizer(func, args): + def func_wrapped(x): + return func(x, *args) + return getattr(optimizer(func_wrapped), 'x', None) + + def _pearsonr(x): + osm_uniform = _calc_uniform_order_statistic_medians(len(x)) + xvals = distributions.norm.ppf(osm_uniform) + + def _eval_pearsonr(lmbda, xvals, samps): + # This function computes the x-axis values of the probability plot + # and computes a linear regression (including the correlation) and + # returns ``1 - r`` so that a minimization function maximizes the + # correlation. + y = boxcox(samps, lmbda) + yvals = np.sort(y) + r, prob = _stats_py.pearsonr(xvals, yvals) + return 1 - r + + return _optimizer(_eval_pearsonr, args=(xvals, x)) + + def _mle(x): + def _eval_mle(lmb, data): + # function to minimize + return -boxcox_llf(lmb, data) + + return _optimizer(_eval_mle, args=(x,)) + + def _all(x): + maxlog = np.empty(2, dtype=float) + maxlog[0] = _pearsonr(x) + maxlog[1] = _mle(x) + return maxlog + + methods = {'pearsonr': _pearsonr, + 'mle': _mle, + 'all': _all} + if method not in methods.keys(): + raise ValueError(f"Method {method} not recognized.") + + optimfunc = methods[method] + + res = optimfunc(x) + + if res is None: + message = ("The `optimizer` argument of `boxcox_normmax` must return " + "an object containing the optimal `lmbda` in attribute `x`.") + raise ValueError(message) + elif not np.isinf(ymax): # adjust the final lambda + # x > 1, boxcox(x) > 0; x < 1, boxcox(x) < 0 + xmax, xmin = np.max(x), np.min(x) + if xmin >= 1: + x_treme = xmax + elif xmax <= 1: + x_treme = xmin + else: # xmin < 1 < xmax + indicator = special.boxcox(xmax, res) > abs(special.boxcox(xmin, res)) + if isinstance(res, np.ndarray): + indicator = indicator[1] # select corresponds with 'mle' + x_treme = xmax if indicator else xmin + + mask = abs(special.boxcox(x_treme, res)) > ymax + if np.any(mask): + message = ( + f"The optimal lambda is {res}, but the returned lambda is the " + f"constrained optimum to ensure that the maximum or the minimum " + f"of the transformed data does not " + end_msg + ) + warnings.warn(message, stacklevel=2) + + # Return the constrained lambda to ensure the transformation + # does not cause overflow or exceed specified `ymax` + constrained_res = _boxcox_inv_lmbda(x_treme, ymax * np.sign(x_treme - 1)) + + if isinstance(res, np.ndarray): + res[mask] = constrained_res + else: + res = constrained_res + return res + + +def _normplot(method, x, la, lb, plot=None, N=80): + """Compute parameters for a Box-Cox or Yeo-Johnson normality plot, + optionally show it. + + See `boxcox_normplot` or `yeojohnson_normplot` for details. + """ + + if method == 'boxcox': + title = 'Box-Cox Normality Plot' + transform_func = boxcox + else: + title = 'Yeo-Johnson Normality Plot' + transform_func = yeojohnson + + x = np.asarray(x) + if x.size == 0: + return x + + if lb <= la: + raise ValueError("`lb` has to be larger than `la`.") + + if method == 'boxcox' and np.any(x <= 0): + raise ValueError("Data must be positive.") + + lmbdas = np.linspace(la, lb, num=N) + ppcc = lmbdas * 0.0 + for i, val in enumerate(lmbdas): + # Determine for each lmbda the square root of correlation coefficient + # of transformed x + z = transform_func(x, lmbda=val) + _, (_, _, r) = probplot(z, dist='norm', fit=True) + ppcc[i] = r + + if plot is not None: + plot.plot(lmbdas, ppcc, 'x') + _add_axis_labels_title(plot, xlabel='$\\lambda$', + ylabel='Prob Plot Corr. Coef.', + title=title) + + return lmbdas, ppcc + + +@xp_capabilities(np_only=True) +def boxcox_normplot(x, la, lb, plot=None, N=80): + """Compute parameters for a Box-Cox normality plot, optionally show it. + + A Box-Cox normality plot shows graphically what the best transformation + parameter is to use in `boxcox` to obtain a distribution that is close + to normal. + + Parameters + ---------- + x : array_like + Input array. + la, lb : scalar + The lower and upper bounds for the ``lmbda`` values to pass to `boxcox` + for Box-Cox transformations. These are also the limits of the + horizontal axis of the plot if that is generated. + plot : object, optional + If given, plots the quantiles and least squares fit. + `plot` is an object that has to have methods "plot" and "text". + The `matplotlib.pyplot` module or a Matplotlib Axes object can be used, + or a custom object with the same methods. + Default is None, which means that no plot is created. + N : int, optional + Number of points on the horizontal axis (equally distributed from + `la` to `lb`). + + Returns + ------- + lmbdas : ndarray + The ``lmbda`` values for which a Box-Cox transform was done. + ppcc : ndarray + Probability Plot Correlation Coefficient, as obtained from `probplot` + when fitting the Box-Cox transformed input `x` against a normal + distribution. + + See Also + -------- + probplot, boxcox, boxcox_normmax, boxcox_llf, ppcc_max + + Notes + ----- + Even if `plot` is given, the figure is not shown or saved by + `boxcox_normplot`; ``plt.show()`` or ``plt.savefig('figname.png')`` + should be used after calling `probplot`. + + Examples + -------- + >>> from scipy import stats + >>> import matplotlib.pyplot as plt + + Generate some non-normally distributed data, and create a Box-Cox plot: + + >>> x = stats.loggamma.rvs(5, size=500) + 5 + >>> fig = plt.figure() + >>> ax = fig.add_subplot(111) + >>> prob = stats.boxcox_normplot(x, -20, 20, plot=ax) + + Determine and plot the optimal ``lmbda`` to transform ``x`` and plot it in + the same plot: + + >>> _, maxlog = stats.boxcox(x) + >>> ax.axvline(maxlog, color='r') + + >>> plt.show() + + """ + return _normplot('boxcox', x, la, lb, plot, N) + + +@xp_capabilities(np_only=True) +def yeojohnson(x, lmbda=None, *, nan_policy='propagate'): + r"""Return a dataset transformed by a Yeo-Johnson power transformation. + + Parameters + ---------- + x : ndarray + Input array. Should be 1-dimensional. + lmbda : float, optional + If ``lmbda`` is ``None``, find the lambda that maximizes the + log-likelihood function and return it as the second output argument. + Otherwise the transformation is done for the given value. + nan_policy : {'propagate', 'omit', 'raise'} + Defines how to handle NaNs in `x`. + + - ``propagate``: if a NaN is present, all outputs will contain NaNs. + - ``omit``: NaNs will be omitted when calculating the optimal `maxlog`; + NaNs in `x` will remain NaNs in the transformed data. + - ``raise``: if a NaN is present, a ``ValueError`` will be raised. + + Returns + ------- + yeojohnson: ndarray + Yeo-Johnson power transformed array. + maxlog : float, optional + If the `lmbda` parameter is None, the second returned argument is + the lambda that maximizes the log-likelihood function. + + See Also + -------- + probplot, yeojohnson_normplot, yeojohnson_normmax, yeojohnson_llf, boxcox + + Notes + ----- + The Yeo-Johnson transform is given by: + + .. math:: + + y = + \begin{cases} + \frac{(x + 1)^\lambda - 1}{\lambda}, + &\text{for } x \geq 0, \lambda \neq 0 + \\ + \log(x + 1), + &\text{for } x \geq 0, \lambda = 0 + \\ + -\frac{(-x + 1)^{2 - \lambda} - 1}{2 - \lambda}, + &\text{for } x < 0, \lambda \neq 2 + \\ + -\log(-x + 1), + &\text{for } x < 0, \lambda = 2 + \end{cases} + + Unlike `boxcox`, `yeojohnson` does not require the input data to be + positive. + + .. versionadded:: 1.2.0 + + + References + ---------- + I. Yeo and R.A. Johnson, "A New Family of Power Transformations to + Improve Normality or Symmetry", Biometrika 87.4 (2000): + + + Examples + -------- + >>> from scipy import stats + >>> import matplotlib.pyplot as plt + + We generate some random variates from a non-normal distribution and make a + probability plot for it, to show it is non-normal in the tails: + + >>> fig = plt.figure() + >>> ax1 = fig.add_subplot(211) + >>> x = stats.loggamma.rvs(5, size=500) + 5 + >>> prob = stats.probplot(x, dist=stats.norm, plot=ax1) + >>> ax1.set_xlabel('') + >>> ax1.set_title('Probplot against normal distribution') + + We now use `yeojohnson` to transform the data so it's closest to normal: + + >>> ax2 = fig.add_subplot(212) + >>> xt, lmbda = stats.yeojohnson(x) + >>> prob = stats.probplot(xt, dist=stats.norm, plot=ax2) + >>> ax2.set_title('Probplot after Yeo-Johnson transformation') + + >>> plt.show() + + """ + x = np.asarray(x) + if x.size == 0: + return x + + if np.issubdtype(x.dtype, np.complexfloating): + raise ValueError('Yeo-Johnson transformation is not defined for ' + 'complex numbers.') + + if np.issubdtype(x.dtype, np.integer): + x = x.astype(np.float64, copy=False) + + if lmbda is not None: + _contains_nan(x, nan_policy, xp_omit_okay=True) # handle nan_policy='raise' + return _yeojohnson_transform(x, lmbda) + + # if lmbda=None, find the lmbda that maximizes the log-likelihood function. + lmax = yeojohnson_normmax(x, nan_policy=nan_policy) + y = _yeojohnson_transform(x, lmax) + + return y, lmax + + +def _yeojohnson_transform(x, lmbda, xp=None): + """Returns `x` transformed by the Yeo-Johnson power transform with given + parameter `lmbda`. + """ + xp = array_namespace(x) if xp is None else xp + dtype = xp_result_type(x, lmbda, force_floating=True, xp=xp) + eps = xp.finfo(dtype).eps + out = xp.zeros_like(x, dtype=dtype) + pos = x >= 0 # binary mask + + if is_jax(xp): + return xp.select( + [(abs(lmbda) < eps) & pos, (abs(lmbda - 2) < eps) & ~pos, pos], + [xp.log1p(x), -xp.log1p(-x), xp.expm1(lmbda * xp.log1p(x)) / lmbda], + -xp.expm1((2 - lmbda) * xp.log1p(-x)) / (2 - lmbda), + ) + + # when x >= 0 + if abs(lmbda) < eps: + out = xpx.at(out)[pos].set(xp.log1p(x[pos])) + else: # lmbda != 0 + # more stable version of: ((x + 1) ** lmbda - 1) / lmbda + out = xpx.at(out)[pos].set(xp.expm1(lmbda * xp.log1p(x[pos])) / lmbda) + + # when x < 0 + if abs(lmbda - 2) > eps: + out = xpx.at(out)[~pos].set( + -xp.expm1((2 - lmbda) * xp.log1p(-x[~pos])) / (2 - lmbda)) + else: # lmbda == 2 + out = xpx.at(out)[~pos].set(-xp.log1p(-x[~pos])) + + return out + + +@xp_capabilities(skip_backends=[("dask.array", "Dask can't broadcast nan shapes")]) +def yeojohnson_llf(lmb, data, *, axis=0, nan_policy='propagate', keepdims=False): + r"""The Yeo-Johnson log-likelihood function. + + Parameters + ---------- + lmb : scalar + Parameter for Yeo-Johnson transformation. See `yeojohnson` for + details. + data : array_like + Data to calculate Yeo-Johnson log-likelihood for. + axis : int, default: 0 + If an int, the axis of the input along which to compute the statistic. + The statistic of each axis-slice (e.g. row) of the input will appear in a + corresponding element of the output. + If ``None``, the input will be raveled before computing the statistic. + nan_policy : {'propagate', 'omit', 'raise' + Defines how to handle input NaNs. + + - ``propagate``: if a NaN is present in the axis slice (e.g. row) along + which the statistic is computed, the corresponding entry of the output + will be NaN. + - ``omit``: NaNs will be omitted when performing the calculation. + If insufficient data remains in the axis slice along which the + statistic is computed, the corresponding entry of the output will be + NaN. + - ``raise``: if a NaN is present, a ``ValueError`` will be raised. + + keepdims : bool, default: False + If this is set to True, the axes which are reduced are left + in the result as dimensions with size one. With this option, + the result will broadcast correctly against the input array. + + Returns + ------- + llf : float + Yeo-Johnson log-likelihood of `data` given `lmb`. + + See Also + -------- + yeojohnson, probplot, yeojohnson_normplot, yeojohnson_normmax + + Notes + ----- + The Yeo-Johnson log-likelihood function :math:`l` is defined here as + + .. math:: + + l = -\frac{N}{2} \log(\hat{\sigma}^2) + (\lambda - 1) + \sum_i^N \text{sign}(x_i) \log(|x_i| + 1) + + where :math:`N` is the number of data points :math:`x` = ``data`` and + :math:`\hat{\sigma}^2` is the estimated variance of the Yeo-Johnson transformed + input data :math:`x`. + This corresponds to the *profile log-likelihood* of the original data :math:`x` + with some constant terms dropped. + + .. versionadded:: 1.2.0 + + Examples + -------- + >>> import numpy as np + >>> from scipy import stats + >>> import matplotlib.pyplot as plt + >>> from mpl_toolkits.axes_grid1.inset_locator import inset_axes + + Generate some random variates and calculate Yeo-Johnson log-likelihood + values for them for a range of ``lmbda`` values: + + >>> x = stats.loggamma.rvs(5, loc=10, size=1000) + >>> lmbdas = np.linspace(-2, 10) + >>> llf = np.zeros(lmbdas.shape, dtype=float) + >>> for ii, lmbda in enumerate(lmbdas): + ... llf[ii] = stats.yeojohnson_llf(lmbda, x) + + Also find the optimal lmbda value with `yeojohnson`: + + >>> x_most_normal, lmbda_optimal = stats.yeojohnson(x) + + Plot the log-likelihood as function of lmbda. Add the optimal lmbda as a + horizontal line to check that that's really the optimum: + + >>> fig = plt.figure() + >>> ax = fig.add_subplot(111) + >>> ax.plot(lmbdas, llf, 'b.-') + >>> ax.axhline(stats.yeojohnson_llf(lmbda_optimal, x), color='r') + >>> ax.set_xlabel('lmbda parameter') + >>> ax.set_ylabel('Yeo-Johnson log-likelihood') + + Now add some probability plots to show that where the log-likelihood is + maximized the data transformed with `yeojohnson` looks closest to normal: + + >>> locs = [3, 10, 4] # 'lower left', 'center', 'lower right' + >>> for lmbda, loc in zip([-1, lmbda_optimal, 9], locs): + ... xt = stats.yeojohnson(x, lmbda=lmbda) + ... (osm, osr), (slope, intercept, r_sq) = stats.probplot(xt) + ... ax_inset = inset_axes(ax, width="20%", height="20%", loc=loc) + ... ax_inset.plot(osm, osr, 'c.', osm, slope*osm + intercept, 'k-') + ... ax_inset.set_xticklabels([]) + ... ax_inset.set_yticklabels([]) + ... ax_inset.set_title(r'$\lambda=%1.2f$' % lmbda) + + >>> plt.show() + + """ + # _axis_nan_policy decorator does not currently support these for lazy arrays. + # We want to run tests with lazy backends, so don't pass the arguments explicitly + # unless necessary. + kwargs = {} + if keepdims is not False: + kwargs['keepdims'] = keepdims + if nan_policy != 'propagate': + kwargs['nan_policy'] = nan_policy + res = _yeojohnson_llf(data, lmb=lmb, axis=axis, **kwargs) + return res[()] if res.ndim == 0 else res + + +@_axis_nan_policy_factory(lambda x: x, n_outputs=1, default_axis=0, + result_to_tuple=lambda x, _: (x,)) +def _yeojohnson_llf(data, *, lmb, axis=0): + xp = array_namespace(data) + dtype = xp_result_type(lmb, data, force_floating=True, xp=xp) + data = xp.asarray(data, dtype=dtype) + + n = data.shape[axis] + if n == 0: + return _get_nan(data, xp=xp) + eps = xp.finfo(dtype).eps + # Special case all-positive/negative data to avoid overflow and precision loss + pos = data >= 0 # binary mask + + # There exists numerical instability when abs(lmb) or abs(lmb - 2) is very small + if not is_lazy_array(pos) and xp.all(pos): + if abs(lmb) < eps: + logvar = xp.log(xp.var(xp.log1p(data), axis=axis)) + else: + logvar = _log_var(lmb * xp.log1p(data), xp, axis) - 2 * math.log(abs(lmb)) + + elif not is_lazy_array(pos) and xp.all(~pos): + if abs(lmb - 2) < eps: + logvar = xp.log(xp.var(xp.log1p(-data), axis=axis)) + else: + logvar = _log_var((2 - lmb) * xp.log1p(-data), xp, axis) - 2 * math.log( + abs(2 - lmb) + ) + + # overflow/precision loss not reported for mixed data; calculate `logvar` directly + else: # mixed positive and negative data + y = _yeojohnson_transform(data, lmb, xp=xp) + sigma = xp.var(y, axis=axis) + + # Suppress RuntimeWarning raised by np.log when the variance is too low + finite_variance = sigma >= xp.finfo(sigma.dtype).smallest_normal + logvar = xpx.apply_where(finite_variance, (sigma,), xp.log, fill_value=-xp.inf) + + loglike = (-n / 2 * logvar + + (lmb - 1) * xp.sum(xp.sign(data) * xp.log1p(xp.abs(data)), axis=axis)) + + return loglike + + +@xp_capabilities(np_only=True) +def yeojohnson_normmax(x, brack=None, *, nan_policy='propagate'): + """Compute optimal Yeo-Johnson transform parameter. + + Compute optimal Yeo-Johnson transform parameter for input data, using + maximum likelihood estimation. + + Parameters + ---------- + x : array_like + Input array. + brack : 2-tuple, optional + The starting interval for a downhill bracket search with + `optimize.brent`. Note that this is in most cases not critical; the + final result is allowed to be outside this bracket. If None, + `optimize.fminbound` is used with bounds that avoid overflow. + nan_policy : {'propagate', 'omit', 'raise'} + Defines how to handle input NaNs. + + - ``propagate``: if a NaN is present in the input, the output will be NaN. + - ``omit``: NaNs will be omitted when performing the calculation. + - ``raise``: if a NaN is present, a ``ValueError`` will be raised. + + Returns + ------- + maxlog : float + The optimal transform parameter found. + + See Also + -------- + yeojohnson, yeojohnson_llf, yeojohnson_normplot + + Notes + ----- + .. versionadded:: 1.2.0 + + Examples + -------- + >>> import numpy as np + >>> from scipy import stats + >>> import matplotlib.pyplot as plt + + Generate some data and determine optimal ``lmbda`` + + >>> rng = np.random.default_rng() + >>> x = stats.loggamma.rvs(5, size=30, random_state=rng) + 5 + >>> lmax = stats.yeojohnson_normmax(x) + + >>> fig = plt.figure() + >>> ax = fig.add_subplot(111) + >>> prob = stats.yeojohnson_normplot(x, -10, 10, plot=ax) + >>> ax.axvline(lmax, color='r') + + >>> plt.show() + + """ + def _neg_llf(lmbda, data): + llf = np.asarray(yeojohnson_llf(lmbda, data)) + # reject likelihoods that are inf which are likely due to small + # variance in the transformed space + llf[np.isinf(llf)] = -np.inf + return -llf + + with np.errstate(invalid='ignore'): + x = np.asarray(x) + + contains_nan = _contains_nan(x, nan_policy, xp_omit_okay=True) + if contains_nan: + if nan_policy == 'propagate': + return _get_nan(x, xp=np)[()] + x = x[~np.isnan(x)] # `nan_policy='omit + + if not np.all(np.isfinite(x)): + raise ValueError('Yeo-Johnson input must be finite.') + if np.all(x == 0): + return 1.0 + if brack is not None: + return optimize.brent(_neg_llf, brack=brack, args=(x,)) + + dtype = x.dtype if np.issubdtype(x.dtype, np.floating) else np.float64 + # Allow values up to 20 times the maximum observed value to be safely + # transformed without over- or underflow. + log1p_max_x = np.log1p(20 * np.max(np.abs(x))) + # Use half of floating point's exponent range to allow safe computation + # of the variance of the transformed data. + log_eps = np.log(np.finfo(dtype).eps) + log_tiny_float = (np.log(np.finfo(dtype).tiny) - log_eps) / 2 + log_max_float = (np.log(np.finfo(dtype).max) + log_eps) / 2 + # Compute the bounds by approximating the inverse of the Yeo-Johnson + # transform on the smallest and largest floating point exponents, given + # the largest data we expect to observe. See [1] for further details. + # [1] https://github.com/scipy/scipy/pull/18852#issuecomment-1630286174 + lb = log_tiny_float / log1p_max_x + ub = log_max_float / log1p_max_x + # Convert the bounds if all or some of the data is negative. + if np.all(x < 0): + lb, ub = 2 - ub, 2 - lb + elif np.any(x < 0): + lb, ub = max(2 - ub, lb), min(2 - lb, ub) + # Match `optimize.brent`'s tolerance. + tol_brent = 1.48e-08 + return optimize.fminbound(_neg_llf, lb, ub, args=(x,), xtol=tol_brent) + + +@xp_capabilities(np_only=True) +def yeojohnson_normplot(x, la, lb, plot=None, N=80): + """Compute parameters for a Yeo-Johnson normality plot, optionally show it. + + A Yeo-Johnson normality plot shows graphically what the best + transformation parameter is to use in `yeojohnson` to obtain a + distribution that is close to normal. + + Parameters + ---------- + x : array_like + Input array. + la, lb : scalar + The lower and upper bounds for the ``lmbda`` values to pass to + `yeojohnson` for Yeo-Johnson transformations. These are also the + limits of the horizontal axis of the plot if that is generated. + plot : object, optional + If given, plots the quantiles and least squares fit. + `plot` is an object that has to have methods "plot" and "text". + The `matplotlib.pyplot` module or a Matplotlib Axes object can be used, + or a custom object with the same methods. + Default is None, which means that no plot is created. + N : int, optional + Number of points on the horizontal axis (equally distributed from + `la` to `lb`). + + Returns + ------- + lmbdas : ndarray + The ``lmbda`` values for which a Yeo-Johnson transform was done. + ppcc : ndarray + Probability Plot Correlation Coefficient, as obtained from `probplot` + when fitting the Box-Cox transformed input `x` against a normal + distribution. + + See Also + -------- + probplot, yeojohnson, yeojohnson_normmax, yeojohnson_llf, ppcc_max + + Notes + ----- + Even if `plot` is given, the figure is not shown or saved by + `boxcox_normplot`; ``plt.show()`` or ``plt.savefig('figname.png')`` + should be used after calling `probplot`. + + .. versionadded:: 1.2.0 + + Examples + -------- + >>> from scipy import stats + >>> import matplotlib.pyplot as plt + + Generate some non-normally distributed data, and create a Yeo-Johnson plot: + + >>> x = stats.loggamma.rvs(5, size=500) + 5 + >>> fig = plt.figure() + >>> ax = fig.add_subplot(111) + >>> prob = stats.yeojohnson_normplot(x, -20, 20, plot=ax) + + Determine and plot the optimal ``lmbda`` to transform ``x`` and plot it in + the same plot: + + >>> _, maxlog = stats.yeojohnson(x) + >>> ax.axvline(maxlog, color='r') + + >>> plt.show() + + """ + return _normplot('yeojohnson', x, la, lb, plot, N) + + +ShapiroResult = namedtuple('ShapiroResult', ('statistic', 'pvalue')) + + +@xp_capabilities() +@_axis_nan_policy_factory(ShapiroResult, n_samples=1, too_small=2, default_axis=None) +def shapiro(x, *, axis=None): + r"""Perform the Shapiro-Wilk test for normality. + + The Shapiro-Wilk test tests the null hypothesis that the + data was drawn from a normal distribution. + + Parameters + ---------- + x : array_like + Array of sample data. Must contain at least three observations. + axis : int or tuple of ints, default: 0 + If an int or tuple of ints, the axis or axes of the input along which + to compute the statistic. The statistic of each axis-slice (e.g. row) + of the input will appear in a corresponding element of the output. + If ``None``, the input will be raveled before computing the statistic. + + Returns + ------- + statistic : float + The test statistic. + p-value : float + The p-value for the hypothesis test. + + See Also + -------- + anderson : The Anderson-Darling test for normality + kstest : The Kolmogorov-Smirnov test for goodness of fit. + :ref:`hypothesis_shapiro` : Extended example + + Notes + ----- + The algorithm used is described in [4]_, but censoring parameters as + described are not implemented. For N > 5000 the W test statistic is + accurate, but the p-value may not be. + + References + ---------- + .. [1] https://www.itl.nist.gov/div898/handbook/prc/section2/prc213.htm + :doi:`10.18434/M32189` + .. [2] Shapiro, S. S. & Wilk, M.B, "An analysis of variance test for + normality (complete samples)", Biometrika, 1965, Vol. 52, + pp. 591-611, :doi:`10.2307/2333709` + .. [3] Razali, N. M. & Wah, Y. B., "Power comparisons of Shapiro-Wilk, + Kolmogorov-Smirnov, Lilliefors and Anderson-Darling tests", Journal + of Statistical Modeling and Analytics, 2011, Vol. 2, pp. 21-33. + .. [4] Royston, P. "A toolkit for testing for non-normality in complete and + censored samples." Journal of the Royal Statistical Society: Series D + (The Statistician) 42.1 (1993): 37-43. + + Examples + -------- + + >>> import numpy as np + >>> from scipy import stats + >>> rng = np.random.default_rng() + >>> x = stats.norm.rvs(loc=5, scale=3, size=100, random_state=rng) + >>> shapiro_test = stats.shapiro(x) + >>> shapiro_test + ShapiroResult(statistic=0.9813305735588074, pvalue=0.16855233907699585) + >>> shapiro_test.statistic + 0.9813305735588074 + >>> shapiro_test.pvalue + 0.16855233907699585 + + For a more detailed example, see :ref:`hypothesis_shapiro`. + """ + # `x` is an array and axis=-1 due to _axis_nan_policy decorator + xp = array_namespace(x) + + N = x.shape[-1] + if N < 3: + raise ValueError("Data must be at least length 3.") + + y = xp.sort(x, axis=-1) + y -= x[..., N//2:N//2+1] # subtract the median (or a nearby value); see gh-15777 + + y = xp_promote(y, force_floating=True, xp=xp) + w, pw = _swilk(y, xp=xp) + if N > 5000: + warnings.warn("scipy.stats.shapiro: For N > 5000, computed p-value " + f"may not be accurate. Current N is {N}.", + stacklevel=2) + + # `w` and `pw` are always Python floats, which are double precision. + # We want to ensure that they are NumPy floats, so until dtypes are + # respected, we can explicitly convert each to float64 (faster than + # `np.array([w, pw])`). + return ShapiroResult(w[()], pw[()]) + + +def _swilk_w(y, a, *, xp): + # calculate Shapiro-Wilk statistic given sorted sample and weights + # Follows [4] Section 2.1 + num = xp.vecdot(a, y, axis=-1) ** 2 + y_ = _demean(y, mean=xp.mean(y, axis=-1, keepdims=True), axis=-1, xp=xp) + den = xp.vecdot(y_, y_, axis=-1) + return num / den + + +def _swilk(y, *, xp): + # calculate Shapiro-Wilk statistic and p-value from sorted sample + n = y.shape[-1] + + if n == 3: + # [2] Table 5 gives the first four digits + c = math.sqrt(2) / 2 + a = xp.asarray([-c, 0, c]) + # [2] Corollary 4; discussed in https://github.com/scipy/scipy/issues/18322 + W = xp.clip(_swilk_w(y, a, xp=xp), 0.75, 1.) + pvalue = xp.clip(1. - 6/np.pi * xp.acos(xp.sqrt(W)), 0., 1.) + return W, pvalue + + # Follows [4] section 2.2 + # could calculate half the coefficients and get the rest by antisymmetry + i = xp.arange(1, n + 1, dtype=y.dtype) + m = special.ndtri((i - 3 / 8) / (n + 1 / 4)) + u = n**(-0.5) + mTm = xp.vecdot(m, m) + c = mTm**(-0.5) * m + mnm1, mn = m[..., -2], m[..., -1] + cnm1, cn = c[..., -2], c[..., -1] + an = (cn + 0.221157*u - 0.147981*u**2 + - 2.071190*u**3 + 4.434685*u**4 - 2.706056*u**5) + anm1 = (cnm1 + 0.042981*u - 0.293762*u**2 + - 1.752461*u**3 + 5.682633*u**4 - 3.582633*u**5) + phi = ((mTm - 2*mn**2) / (1 - 2*an**2) if n <= 5 + else (mTm - 2*mn**2 - 2*mnm1**2) / (1 - 2*an**2 - 2*anm1**2)) + a = phi**(-0.5) * m + if n > 5: + a = xpx.at(a)[..., -2].set(anm1) + a = xpx.at(a)[..., 1].set(-anm1) + a = xpx.at(a)[..., -1].set(an) + a = xpx.at(a)[..., -0].set(-an) + W = _swilk_w(y, a, xp=xp) + + # Follows [4] Table 1 + if n <= 11: + u = n + gamma = -2.273 + 0.459*u + mu = 0.5440 - 0.39978*u + 0.025054*u**2 - 0.0006714*u**3 + sigma = math.exp(1.3822 - 0.77857*u + 0.062767*u**2 - 0.0020322*u**3) + gW = -xp.log(gamma - xp.log(1 - W)) + else: + u = math.log(n) + mu = -1.5861 - 0.31082*u - 0.083751*u**2 + 0.0038915*u**3 + sigma = math.exp(-0.4803 - 0.082676*u + 0.0030302*u**2) + gW = xp.log(1 - W) + + z = (gW - mu) / sigma + pvalue = special.ndtr(-z) + return W, pvalue + + +# Values from [8] +_Avals_norm = array([0.561, 0.631, 0.752, 0.873, 1.035]) +_Avals_expon = array([0.916, 1.062, 1.321, 1.591, 1.959]) +# From Stephens, M A, "Goodness of Fit for the Extreme Value Distribution", +# Biometrika, Vol. 64, Issue 3, Dec. 1977, pp 583-588. +_Avals_gumbel = array([0.474, 0.637, 0.757, 0.877, 1.038]) +# From Stephens, M A, "Tests of Fit for the Logistic Distribution Based +# on the Empirical Distribution Function.", Biometrika, +# Vol. 66, Issue 3, Dec. 1979, pp 591-595. +_Avals_logistic = array([0.426, 0.563, 0.660, 0.769, 0.906, 1.010]) +# From Richard A. Lockhart and Michael A. Stephens "Estimation and Tests of +# Fit for the Three-Parameter Weibull Distribution" +# Journal of the Royal Statistical Society.Series B(Methodological) +# Vol. 56, No. 3 (1994), pp. 491-500, table 1. Keys are c*100 +_Avals_weibull = [[0.292, 0.395, 0.467, 0.522, 0.617, 0.711, 0.836, 0.931], + [0.295, 0.399, 0.471, 0.527, 0.623, 0.719, 0.845, 0.941], + [0.298, 0.403, 0.476, 0.534, 0.631, 0.728, 0.856, 0.954], + [0.301, 0.408, 0.483, 0.541, 0.640, 0.738, 0.869, 0.969], + [0.305, 0.414, 0.490, 0.549, 0.650, 0.751, 0.885, 0.986], + [0.309, 0.421, 0.498, 0.559, 0.662, 0.765, 0.902, 1.007], + [0.314, 0.429, 0.508, 0.570, 0.676, 0.782, 0.923, 1.030], + [0.320, 0.438, 0.519, 0.583, 0.692, 0.802, 0.947, 1.057], + [0.327, 0.448, 0.532, 0.598, 0.711, 0.824, 0.974, 1.089], + [0.334, 0.469, 0.547, 0.615, 0.732, 0.850, 1.006, 1.125], + [0.342, 0.472, 0.563, 0.636, 0.757, 0.879, 1.043, 1.167]] +_Avals_weibull = np.array(_Avals_weibull) +_cvals_weibull = np.linspace(0, 0.5, 11) +_get_As_weibull = interpolate.interp1d(_cvals_weibull, _Avals_weibull.T, + kind='linear', bounds_error=False, + fill_value=_Avals_weibull[-1]) + + +def _weibull_fit_check(params, x): + # Refine the fit returned by `weibull_min.fit` to ensure that the first + # order necessary conditions are satisfied. If not, raise an error. + # Here, use `m` for the shape parameter to be consistent with [7] + # and avoid confusion with `c` as defined in [7]. + n = len(x) + m, u, s = params + + def dnllf_dm(m, u): + # Partial w.r.t. shape w/ optimal scale. See [7] Equation 5. + xu = x-u + return (1/m - (xu**m*np.log(xu)).sum()/(xu**m).sum() + + np.log(xu).sum()/n) + + def dnllf_du(m, u): + # Partial w.r.t. loc w/ optimal scale. See [7] Equation 6. + xu = x-u + return (m-1)/m*(xu**-1).sum() - n*(xu**(m-1)).sum()/(xu**m).sum() + + def get_scale(m, u): + # Partial w.r.t. scale solved in terms of shape and location. + # See [7] Equation 7. + return ((x-u)**m/n).sum()**(1/m) + + def dnllf(params): + # Partial derivatives of the NLLF w.r.t. parameters, i.e. + # first order necessary conditions for MLE fit. + return [dnllf_dm(*params), dnllf_du(*params)] + + suggestion = ("Maximum likelihood estimation is known to be challenging " + "for the three-parameter Weibull distribution. Consider " + "performing a custom goodness-of-fit test using " + "`scipy.stats.monte_carlo_test`.") + + if np.allclose(u, np.min(x)) or m < 1: + # The critical values provided by [7] don't seem to control the + # Type I error rate in this case. Error out. + message = ("Maximum likelihood estimation has converged to " + "a solution in which the location is equal to the minimum " + "of the data, the shape parameter is less than 2, or both. " + "The table of critical values in [7] does not " + "include this case. " + suggestion) + raise ValueError(message) + + try: + # Refine the MLE / verify that first-order necessary conditions are + # satisfied. If so, the critical values provided in [7] seem reliable. + with np.errstate(over='raise', invalid='raise'): + res = optimize.root(dnllf, params[:-1]) + + message = ("Solution of MLE first-order conditions failed: " + f"{res.message}. `anderson` cannot continue. " + suggestion) + if not res.success: + raise ValueError(message) + + except (FloatingPointError, ValueError) as e: + message = ("An error occurred while fitting the Weibull distribution " + "to the data, so `anderson` cannot continue. " + suggestion) + raise ValueError(message) from e + + m, u = res.x + s = get_scale(m, u) + return m, u, s + + +AndersonResult = _make_tuple_bunch('AndersonResult', + ['statistic', 'critical_values', + 'significance_level'], ['fit_result']) + + +_anderson_warning_message = ( +"""As of SciPy 1.17, users must choose a p-value calculation method by providing the +`method` parameter. `method='interpolate'` interpolates the p-value from pre-calculated +tables; `method` may also be an instance of `MonteCarloMethod` to approximate the +p-value via Monte Carlo simulation. When `method` is specified, the result object will +include a `pvalue` attribute and not attributes `critical_value`, `significance_level`, +or `fit_result`. Beginning in 1.19.0, these other attributes will no longer be +available, and a p-value will always be computed according to one of the available +`method` options.""".replace('\n', ' ')) + + +@xp_capabilities(np_only=True) +def anderson(x, dist='norm', *, method=None): + """Anderson-Darling test for data coming from a particular distribution. + + The Anderson-Darling test tests the null hypothesis that a sample is + drawn from a population that follows a particular distribution. + For the Anderson-Darling test, the critical values depend on + which distribution is being tested against. This function works + for normal, exponential, logistic, weibull_min, or Gumbel (Extreme Value + Type I) distributions. + + Parameters + ---------- + x : array_like + Array of sample data. + dist : {'norm', 'expon', 'logistic', 'gumbel', 'gumbel_l', 'gumbel_r', 'extreme1', 'weibull_min'}, optional + The type of distribution to test against. The default is 'norm'. + The names 'extreme1', 'gumbel_l' and 'gumbel' are synonyms for the + same distribution. + method : str or instance of `MonteCarloMethod` + Defines the method used to compute the p-value. + If `method` is ``"interpolated"``, the p-value is interpolated from + pre-calculated tables. + If `method` is an instance of `MonteCarloMethod`, the p-value is computed using + `scipy.stats.monte_carlo_test` with the provided configuration options and other + appropriate settings. + + .. versionadded:: 1.17.0 + If `method` is not specified, `anderson` will emit a ``FutureWarning`` + specifying that the user must opt into a p-value calculation method. + When `method` is specified, the object returned will include a ``pvalue`` + attribute, but no ``critical_value``, ``significance_level``, or + ``fit_result`` attributes. Beginning in 1.19.0, these other attributes will + no longer be available, and a p-value will always be computed according to + one of the available `method` options. + + Returns + ------- + result : AndersonResult + If `method` is provided, this is an object with the following attributes: + + statistic : float + The Anderson-Darling test statistic. + pvalue: float + The p-value corresponding with the test statistic, calculated according to + the specified `method`. + + If `method` is unspecified, this is an object with the following attributes: + + statistic : float + The Anderson-Darling test statistic. + critical_values : list + The critical values for this distribution. + significance_level : list + The significance levels for the corresponding critical values + in percents. The function returns critical values for a + differing set of significance levels depending on the + distribution that is being tested against. + fit_result : `~scipy.stats._result_classes.FitResult` + An object containing the results of fitting the distribution to + the data. + + .. deprecated:: 1.17.0 + The tuple-unpacking behavior of the return object and attributes + ``critical_values``, ``significance_level``, and ``fit_result`` are + deprecated. Beginning in SciPy 1.19.0, these features will no longer be + available, and the object returned will have attributes ``statistic`` and + ``pvalue``. + + See Also + -------- + kstest : The Kolmogorov-Smirnov test for goodness-of-fit. + + Notes + ----- + Critical values provided when `method` is unspecified are for the following + significance levels: + + normal/exponential + 15%, 10%, 5%, 2.5%, 1% + logistic + 25%, 10%, 5%, 2.5%, 1%, 0.5% + gumbel_l / gumbel_r + 25%, 10%, 5%, 2.5%, 1% + weibull_min + 50%, 25%, 15%, 10%, 5%, 2.5%, 1%, 0.5% + + If the returned statistic is larger than these critical values then + for the corresponding significance level, the null hypothesis that + the data come from the chosen distribution can be rejected. + The returned statistic is referred to as 'A2' in the references. + + For `weibull_min`, maximum likelihood estimation is known to be + challenging. If the test returns successfully, then the first order + conditions for a maximum likelihood estimate have been verified and + the critical values correspond relatively well to the significance levels, + provided that the sample is sufficiently large (>10 observations [7]). + However, for some data - especially data with no left tail - `anderson` + is likely to result in an error message. In this case, consider + performing a custom goodness of fit test using + `scipy.stats.monte_carlo_test`. + + References + ---------- + .. [1] https://www.itl.nist.gov/div898/handbook/prc/section2/prc213.htm + .. [2] Stephens, M. A. (1974). EDF Statistics for Goodness of Fit and + Some Comparisons, Journal of the American Statistical Association, + Vol. 69, pp. 730-737. + .. [3] Stephens, M. A. (1976). Asymptotic Results for Goodness-of-Fit + Statistics with Unknown Parameters, Annals of Statistics, Vol. 4, + pp. 357-369. + .. [4] Stephens, M. A. (1977). Goodness of Fit for the Extreme Value + Distribution, Biometrika, Vol. 64, pp. 583-588. + .. [5] Stephens, M. A. (1977). Goodness of Fit with Special Reference + to Tests for Exponentiality , Technical Report No. 262, + Department of Statistics, Stanford University, Stanford, CA. + .. [6] Stephens, M. A. (1979). Tests of Fit for the Logistic Distribution + Based on the Empirical Distribution Function, Biometrika, Vol. 66, + pp. 591-595. + .. [7] Richard A. Lockhart and Michael A. Stephens "Estimation and Tests of + Fit for the Three-Parameter Weibull Distribution" + Journal of the Royal Statistical Society.Series B(Methodological) + Vol. 56, No. 3 (1994), pp. 491-500, Table 0. + .. [8] D'Agostino, Ralph B. (1986). "Tests for the Normal Distribution". + In: Goodness-of-Fit Techniques. Ed. by Ralph B. D'Agostino and + Michael A. Stephens. New York: Marcel Dekker, pp. 122-141. ISBN: + 0-8247-7487-6. + + Examples + -------- + Test the null hypothesis that a random sample was drawn from a normal + distribution (with unspecified mean and standard deviation). + + >>> import numpy as np + >>> from scipy.stats import anderson + >>> rng = np.random.default_rng(9781234521) + >>> data = rng.random(size=35) + >>> res = anderson(data, dist='norm', method='interpolate') + >>> res.statistic + np.float64(0.9887620209957291) + >>> res.pvalue + np.float64(0.012111200538380142) + + The p-value is approximately 0.012,, so the null hypothesis may be rejected + at a significance level of 2.5%, but not at a significance level of 1%. + + """ # numpy/numpydoc#87 # noqa: E501 + dist = dist.lower() + if dist in {'extreme1', 'gumbel'}: + dist = 'gumbel_l' + dists = {'norm', 'expon', 'gumbel_l', + 'gumbel_r', 'logistic', 'weibull_min'} + + if dist not in dists: + raise ValueError(f"Invalid distribution; dist must be in {dists}.") + y = sort(x) + xbar = np.mean(x, axis=0) + N = len(y) + if dist == 'norm': + s = np.std(x, ddof=1, axis=0) + w = (y - xbar) / s + fit_params = xbar, s + logcdf = distributions.norm.logcdf(w) + logsf = distributions.norm.logsf(w) + sig = array([15, 10, 5, 2.5, 1]) + critical = around(_Avals_norm / (1.0 + 0.75/N + 2.25/N/N), 3) + elif dist == 'expon': + w = y / xbar + fit_params = 0, xbar + logcdf = distributions.expon.logcdf(w) + logsf = distributions.expon.logsf(w) + sig = array([15, 10, 5, 2.5, 1]) + critical = around(_Avals_expon / (1.0 + 0.6/N), 3) + elif dist == 'logistic': + def rootfunc(ab, xj, N): + a, b = ab + tmp = (xj - a) / b + tmp2 = exp(tmp) + val = [np.sum(1.0/(1+tmp2), axis=0) - 0.5*N, + np.sum(tmp*(1.0-tmp2)/(1+tmp2), axis=0) + N] + return array(val) + + sol0 = array([xbar, np.std(x, ddof=1, axis=0)]) + sol = optimize.fsolve(rootfunc, sol0, args=(x, N), xtol=1e-5) + w = (y - sol[0]) / sol[1] + fit_params = sol + logcdf = distributions.logistic.logcdf(w) + logsf = distributions.logistic.logsf(w) + sig = array([25, 10, 5, 2.5, 1, 0.5]) + critical = around(_Avals_logistic / (1.0 + 0.25/N), 3) + elif dist == 'gumbel_r': + xbar, s = distributions.gumbel_r.fit(x) + w = (y - xbar) / s + fit_params = xbar, s + logcdf = distributions.gumbel_r.logcdf(w) + logsf = distributions.gumbel_r.logsf(w) + sig = array([25, 10, 5, 2.5, 1]) + critical = around(_Avals_gumbel / (1.0 + 0.2/sqrt(N)), 3) + elif dist == 'gumbel_l': + xbar, s = distributions.gumbel_l.fit(x) + w = (y - xbar) / s + fit_params = xbar, s + logcdf = distributions.gumbel_l.logcdf(w) + logsf = distributions.gumbel_l.logsf(w) + sig = array([25, 10, 5, 2.5, 1]) + critical = around(_Avals_gumbel / (1.0 + 0.2/sqrt(N)), 3) + elif dist == 'weibull_min': + message = ("Critical values of the test statistic are given for the " + "asymptotic distribution. These may not be accurate for " + "samples with fewer than 10 observations. Consider using " + "`scipy.stats.monte_carlo_test`.") + if N < 10: + warnings.warn(message, stacklevel=2) + # [7] writes our 'c' as 'm', and they write `c = 1/m`. Use their names. + m, loc, scale = distributions.weibull_min.fit(y) + m, loc, scale = _weibull_fit_check((m, loc, scale), y) + fit_params = m, loc, scale + logcdf = stats.weibull_min(*fit_params).logcdf(y) + logsf = stats.weibull_min(*fit_params).logsf(y) + c = 1 / m # m and c are as used in [7] + sig = array([0.5, 0.75, 0.85, 0.9, 0.95, 0.975, 0.99, 0.995]) + critical = _get_As_weibull(c) + # Goodness-of-fit tests should only be used to provide evidence + # _against_ the null hypothesis. Be conservative and round up. + critical = np.round(critical + 0.0005, decimals=3) + + i = arange(1, N + 1) + A2 = -N - np.sum((2*i - 1.0) / N * (logcdf + logsf[::-1]), axis=0) + + # FitResult initializer expects an optimize result, so let's work with it + message = '`anderson` successfully fit the distribution to the data.' + res = optimize.OptimizeResult(success=True, message=message) + res.x = np.array(fit_params) + fit_result = FitResult(getattr(distributions, dist), y, + discrete=False, res=res) + + if method is None: + warnings.warn(_anderson_warning_message, FutureWarning, stacklevel=2) + return AndersonResult(A2, critical, sig, fit_result=fit_result) + + if method == 'interpolate': + sig = 1 - sig if dist == 'weibull_min' else sig / 100 + pvalue = np.interp(A2, critical, sig) + elif isinstance(method, stats.MonteCarloMethod): + pvalue = _anderson_simulate_pvalue(x, dist, method) + else: + message = ("`method` must be either 'interpolate' or " + "an instance of `MonteCarloMethod`.") + raise ValueError(message) + return SignificanceResult(statistic=A2, pvalue=pvalue) + + +def _anderson_simulate_pvalue(x, dist, method): + message = ("The `___` attribute of a `MonteCarloMethod` object passed as the " + "`method` parameter of `scipy.stats.anderson` is ignored.") + + method = method._asdict() + if method.pop('rvs', False): + warnings.warn(message.replace('___', 'rvs'), UserWarning, stacklevel=3) + if method.pop('batch', False): + warnings.warn(message.replace('___', 'batch'), UserWarning, stacklevel=3) + method['n_mc_samples'] = method.pop('n_resamples') + + kwargs= {'known_params': {'loc': 0}} if dist == 'expon' else {} + dist = getattr(stats, dist) + res = stats.goodness_of_fit(dist, x, statistic='ad', **kwargs, **method) + return res.pvalue + + +def _anderson_ksamp_continuous(samples, Z, Zstar, k, n, N): + """Compute A2akN equation 3 of Scholz & Stephens. + + Parameters + ---------- + samples : sequence of 1-D array_like + Array of sample arrays. + Z : array_like + Sorted array of all observations. + Zstar : array_like + Sorted array of unique observations. Unused. + k : int + Number of samples. + n : array_like + Number of observations in each sample. + N : int + Total number of observations. + + Returns + ------- + A2KN : float + The A2KN statistics of Scholz and Stephens 1987. + + """ + A2kN = 0. + + j = np.arange(1, N) + for i in arange(0, k): + s = np.sort(samples[i]) + Mij = s.searchsorted(Z[:-1], side='right') + inner = (N*Mij - j*n[i])**2 / (j * (N - j)) + A2kN += inner.sum() / n[i] + return A2kN / N + + +def _anderson_ksamp_midrank(samples, Z, Zstar, k, n, N): + """Compute A2akN equation 7 of Scholz and Stephens. + + Parameters + ---------- + samples : sequence of 1-D array_like + Array of sample arrays. + Z : array_like + Sorted array of all observations. + Zstar : array_like + Sorted array of unique observations. + k : int + Number of samples. + n : array_like + Number of observations in each sample. + N : int + Total number of observations. + + Returns + ------- + A2aKN : float + The A2aKN statistics of Scholz and Stephens 1987. + + """ + A2akN = 0. + Z_ssorted_left = Z.searchsorted(Zstar, 'left') + if N == Zstar.size: + lj = 1. + else: + lj = Z.searchsorted(Zstar, 'right') - Z_ssorted_left + Bj = Z_ssorted_left + lj / 2. + for i in arange(0, k): + s = np.sort(samples[i]) + s_ssorted_right = s.searchsorted(Zstar, side='right') + Mij = s_ssorted_right.astype(float) + fij = s_ssorted_right - s.searchsorted(Zstar, 'left') + Mij -= fij / 2. + inner = lj / float(N) * (N*Mij - Bj*n[i])**2 / (Bj*(N - Bj) - N*lj/4.) + A2akN += inner.sum() / n[i] + A2akN *= (N - 1.) / N + return A2akN + + +def _anderson_ksamp_right(samples, Z, Zstar, k, n, N): + """Compute A2akN equation 6 of Scholz & Stephens. + + Parameters + ---------- + samples : sequence of 1-D array_like + Array of sample arrays. + Z : array_like + Sorted array of all observations. + Zstar : array_like + Sorted array of unique observations. + k : int + Number of samples. + n : array_like + Number of observations in each sample. + N : int + Total number of observations. + + Returns + ------- + A2KN : float + The A2KN statistics of Scholz and Stephens 1987. + + """ + A2kN = 0. + lj = Z.searchsorted(Zstar[:-1], 'right') - Z.searchsorted(Zstar[:-1], + 'left') + Bj = lj.cumsum() + for i in arange(0, k): + s = np.sort(samples[i]) + Mij = s.searchsorted(Zstar[:-1], side='right') + inner = lj / float(N) * (N * Mij - Bj * n[i])**2 / (Bj * (N - Bj)) + A2kN += inner.sum() / n[i] + return A2kN + + +Anderson_ksampResult = _make_tuple_bunch( + 'Anderson_ksampResult', + ['statistic', 'critical_values', 'pvalue'], [] +) + + +@xp_capabilities(np_only=True) +def anderson_ksamp(samples, midrank=_NoValue, *, variant=_NoValue, method=None): + """The Anderson-Darling test for k-samples. + + The k-sample Anderson-Darling test is a modification of the + one-sample Anderson-Darling test. It tests the null hypothesis + that k-samples are drawn from the same population without having + to specify the distribution function of that population. The + critical values depend on the number of samples. + + Parameters + ---------- + samples : sequence of 1-D array_like + Array of sample data in arrays. + midrank : bool, optional + Variant of Anderson-Darling test which is computed. Default + (True) is the midrank test applicable to continuous and + discrete populations. If False, the right side empirical + distribution is used. + + .. deprecated:: 1.17.0 + Use parameter `variant` instead. + variant : {'midrank', 'right', 'continuous'} + Variant of Anderson-Darling test to be computed. ``'midrank'`` is applicable + to both continuous and discrete populations. ``'discrete'`` and ``'continuous'`` + perform alternative versions of the test for discrete and continuous + populations, respectively. + When `variant` is specified, the return object will not be unpackable as a + tuple, and only attributes ``statistic`` and ``pvalue`` will be present. + method : PermutationMethod, optional + Defines the method used to compute the p-value. If `method` is an + instance of `PermutationMethod`, the p-value is computed using + `scipy.stats.permutation_test` with the provided configuration options + and other appropriate settings. Otherwise, the p-value is interpolated + from tabulated values. + + Returns + ------- + res : Anderson_ksampResult + An object containing attributes: + + statistic : float + Normalized k-sample Anderson-Darling test statistic. + critical_values : array + The critical values for significance levels 25%, 10%, 5%, 2.5%, 1%, + 0.5%, 0.1%. + + .. deprecated:: 1.17.0 + Present only when `variant` is unspecified. + + pvalue : float + The approximate p-value of the test. If `method` is not + provided, the value is floored / capped at 0.1% / 25%. + + Raises + ------ + ValueError + If fewer than 2 samples are provided, a sample is empty, or no + distinct observations are in the samples. + + See Also + -------- + ks_2samp : 2 sample Kolmogorov-Smirnov test + anderson : 1 sample Anderson-Darling test + + Notes + ----- + [1]_ defines three versions of the k-sample Anderson-Darling test: + one for continuous distributions and two for discrete + distributions, in which ties between samples may occur. The + default of this routine is to compute the version based on the + midrank empirical distribution function. This test is applicable + to continuous and discrete data. If `variant` is set to ``'discrete'``, the + right side empirical distribution is used for a test for discrete + data; if `variant` is ``'continuous'``, the same test statistic and p-value are + computed for data with no ties, but with less computation. According to [1]_, + the two discrete test statistics differ only slightly if a few collisions due + to round-off errors occur in the test not adjusted for ties between samples. + + The critical values corresponding to the significance levels from 0.01 + to 0.25 are taken from [1]_. p-values are floored / capped + at 0.1% / 25%. Since the range of critical values might be extended in + future releases, it is recommended not to test ``p == 0.25``, but rather + ``p >= 0.25`` (analogously for the lower bound). + + .. versionadded:: 0.14.0 + + References + ---------- + .. [1] Scholz, F. W and Stephens, M. A. (1987), K-Sample + Anderson-Darling Tests, Journal of the American Statistical + Association, Vol. 82, pp. 918-924. + + Examples + -------- + >>> import numpy as np + >>> from scipy import stats + >>> rng = np.random.default_rng(44925884305279435) + >>> res = stats.anderson_ksamp([rng.normal(size=50), rng.normal(loc=0.5, size=30)], + ... variant='midrank') + >>> res.statistic, res.pvalue + (3.4444310693448936, 0.013106682406720973) + + The null hypothesis that the two random samples come from the same + distribution can be rejected at the 5% level because the returned + p-value is less than 0.05, but not at the 1% level. + + >>> samples = [rng.normal(size=50), rng.normal(size=30), + ... rng.normal(size=20)] + >>> res = stats.anderson_ksamp(samples, variant='continuous') + >>> res.statistic, res.pvalue + (-0.6309662273193832, 0.25) + + As we might expect, the null hypothesis cannot be rejected here for three samples + from an identical distribution. The reported p-value (25%) has been capped at the + maximum value for which pre-computed p-values are available. + + In such cases where the p-value is capped or when sample sizes are + small, a permutation test may be more accurate. + + >>> method = stats.PermutationMethod(n_resamples=9999, random_state=rng) + >>> res = stats.anderson_ksamp(samples, variant='continuous', method=method) + >>> res.pvalue + 0.699 + + """ + k = len(samples) + if (k < 2): + raise ValueError("anderson_ksamp needs at least two samples") + + samples = list(map(np.asarray, samples)) + Z = np.sort(np.hstack(samples)) + N = Z.size + Zstar = np.unique(Z) + if Zstar.size < 2: + raise ValueError("anderson_ksamp needs more than one distinct " + "observation") + + n = np.array([sample.size for sample in samples]) + if np.any(n == 0): + raise ValueError("anderson_ksamp encountered sample without " + "observations") + + if variant == _NoValue or midrank != _NoValue: + message = ("Parameter `variant` has been introduced to replace `midrank`; " + "`midrank` will be removed in SciPy 1.19.0. Specify `variant` to " + "silence this warning. Note that the returned object will no longer " + "be unpackable as a tuple, and `critical_values` will be omitted.") + warnings.warn(message, category=UserWarning, stacklevel=2) + + return_critical_values = False + if variant == _NoValue: + return_critical_values = True + variant = 'midrank' if midrank else 'right' + + if variant == 'midrank': + A2kN_fun = _anderson_ksamp_midrank + elif variant == 'right': + A2kN_fun = _anderson_ksamp_right + elif variant == 'continuous': + A2kN_fun = _anderson_ksamp_continuous + else: + message = "`variant` must be one of 'midrank', 'right', or 'continuous'." + raise ValueError(message) + + A2kN = A2kN_fun(samples, Z, Zstar, k, n, N) + + def statistic(*samples): + return A2kN_fun(samples, Z, Zstar, k, n, N) + + if method is not None: + res = stats.permutation_test(samples, statistic, **method._asdict(), + alternative='greater') + + H = (1. / n).sum() + hs_cs = (1. / arange(N - 1, 1, -1)).cumsum() + h = hs_cs[-1] + 1 + g = (hs_cs / arange(2, N)).sum() + + a = (4*g - 6) * (k - 1) + (10 - 6*g)*H + b = (2*g - 4)*k**2 + 8*h*k + (2*g - 14*h - 4)*H - 8*h + 4*g - 6 + c = (6*h + 2*g - 2)*k**2 + (4*h - 4*g + 6)*k + (2*h - 6)*H + 4*h + d = (2*h + 6)*k**2 - 4*h*k + sigmasq = (a*N**3 + b*N**2 + c*N + d) / ((N - 1.) * (N - 2.) * (N - 3.)) + m = k - 1 + A2 = (A2kN - m) / math.sqrt(sigmasq) + + # The b_i values are the interpolation coefficients from Table 2 + # of Scholz and Stephens 1987 + b0 = np.array([0.675, 1.281, 1.645, 1.96, 2.326, 2.573, 3.085]) + b1 = np.array([-0.245, 0.25, 0.678, 1.149, 1.822, 2.364, 3.615]) + b2 = np.array([-0.105, -0.305, -0.362, -0.391, -0.396, -0.345, -0.154]) + critical = b0 + b1 / math.sqrt(m) + b2 / m + + sig = np.array([0.25, 0.1, 0.05, 0.025, 0.01, 0.005, 0.001]) + + if A2 < critical.min() and method is None: + p = sig.max() + msg = (f"p-value capped: true value larger than {p}. Consider " + "specifying `method` " + "(e.g. `method=stats.PermutationMethod()`.)") + warnings.warn(msg, stacklevel=2) + elif A2 > critical.max() and method is None: + p = sig.min() + msg = (f"p-value floored: true value smaller than {p}. Consider " + "specifying `method` " + "(e.g. `method=stats.PermutationMethod()`.)") + warnings.warn(msg, stacklevel=2) + elif method is None: + # interpolation of probit of significance level + pf = np.polyfit(critical, log(sig), 2) + p = math.exp(np.polyval(pf, A2)) + else: + p = res.pvalue if method is not None else p + + if return_critical_values: + # create result object with alias for backward compatibility + res = Anderson_ksampResult(A2, critical, p) + res.significance_level = p + else: + res = SignificanceResult(statistic=A2, pvalue=p) + + return res + + + +AnsariResult = namedtuple('AnsariResult', ('statistic', 'pvalue')) + + +class _ABW: + """Distribution of Ansari-Bradley W-statistic under the null hypothesis.""" + # TODO: calculate exact distribution considering ties + # We could avoid summing over more than half the frequencies, + # but initially it doesn't seem worth the extra complexity + + def __init__(self): + """Minimal initializer.""" + self.m = None + self.n = None + self.astart = None + self.total = None + self.freqs = None + + def _recalc(self, n, m): + """When necessary, recalculate exact distribution.""" + if n != self.n or m != self.m: + self.n, self.m = n, m + # distribution is NOT symmetric when m + n is odd + # n is len(x), m is len(y), and ratio of scales is defined x/y + astart, a1, _ = gscale(n, m) + self.astart = astart # minimum value of statistic + # Exact distribution of test statistic under null hypothesis + # expressed as frequencies/counts/integers to maintain precision. + # Stored as floats to avoid overflow of sums. + self.freqs = a1.astype(np.float64) + self.total = self.freqs.sum() # could calculate from m and n + # probability mass is self.freqs / self.total; + + def pmf(self, k, n, m): + """Probability mass function.""" + self._recalc(n, m) + # The convention here is that PMF at k = 12.5 is the same as at k = 12, + # -> use `floor` in case of ties. + ind = np.floor(k - self.astart).astype(int) + return self.freqs[ind] / self.total + + def cdf(self, k, n, m): + """Cumulative distribution function.""" + self._recalc(n, m) + # Null distribution derived without considering ties is + # approximate. Round down to avoid Type I error. + ind = np.ceil(k - self.astart).astype(int) + return self.freqs[:ind+1].sum() / self.total + + def sf(self, k, n, m): + """Survival function.""" + self._recalc(n, m) + # Null distribution derived without considering ties is + # approximate. Round down to avoid Type I error. + ind = np.floor(k - self.astart).astype(int) + return self.freqs[ind:].sum() / self.total + + +# Maintain state for faster repeat calls to ansari w/ method='exact' +# _ABW() is calculated once per thread and stored as an attribute on +# this thread-local variable inside ansari(). +_abw_state = threading.local() + + +@xp_capabilities(cpu_only=True, skip_backends=[('dask.array', 'no rankdata')]) +@_axis_nan_policy_factory(AnsariResult, n_samples=2) +def ansari(x, y, alternative='two-sided', *, axis=0, method='auto'): + """Perform the Ansari-Bradley test for equal scale parameters. + + The Ansari-Bradley test ([1]_, [2]_) is a non-parametric test + for the equality of the scale parameter of the distributions + from which two samples were drawn. The null hypothesis states that + the ratio of the scale of the distribution underlying `x` to the scale + of the distribution underlying `y` is 1. + + Parameters + ---------- + x, y : array_like + Arrays of sample data. + alternative : {'two-sided', 'less', 'greater'}, optional + Defines the alternative hypothesis. Default is 'two-sided'. + The following options are available: + + * 'two-sided': the ratio of scales is not equal to 1. + * 'less': the ratio of scales is less than 1. + * 'greater': the ratio of scales is greater than 1. + + .. versionadded:: 1.7.0 + axis : int or tuple of ints, default: 0 + If an int or tuple of ints, the axis or axes of the input along which + to compute the statistic. The statistic of each axis-slice (e.g. row) + of the input will appear in a corresponding element of the output. + If ``None``, the input will be raveled before computing the statistic. + method : {'auto', 'asymptotic', 'exact'} or `PermutationMethod` instance, optional + Selects the method used to calculate the *p*-value. + Default is 'auto'. The following options are available. + + * ``'asymptotic'``: compares the standardized test statistic + against the normal distribution, correcting for ties. + * ``'exact'``: computes the exact *p*-value by comparing the observed + statistic against the exact distribution of the statistic under the + null hypothesis. No correction is made for ties. + * ``'auto'``: chooses ``'exact'`` when the size of both + samples is less than or equal to 55 and there are no ties; + chooses ``'asymptotic'`` otherwise. + * `PermutationMethod` instance. In this case, the p-value + is computed using `permutation_test` with the provided + configuration options and other appropriate settings. + + Returns + ------- + statistic : float + The Ansari-Bradley test statistic. + pvalue : float + The p-value of the hypothesis test. + + See Also + -------- + fligner : A non-parametric test for the equality of k variances + mood : A non-parametric test for the equality of two scale parameters + + Notes + ----- + The p-value given is exact when the sample sizes are both less than + 55 and there are no ties, otherwise a normal approximation for the + p-value is used. + + References + ---------- + .. [1] Ansari, A. R. and Bradley, R. A. (1960) Rank-sum tests for + dispersions, Annals of Mathematical Statistics, 31, 1174-1189. + .. [2] Sprent, Peter and N.C. Smeeton. Applied nonparametric + statistical methods. 3rd ed. Chapman and Hall/CRC. 2001. + Section 5.8.2. + .. [3] Nathaniel E. Helwig "Nonparametric Dispersion and Equality + Tests" at http://users.stat.umn.edu/~helwig/notes/npde-Notes.pdf + + Examples + -------- + >>> import numpy as np + >>> from scipy.stats import ansari + >>> rng = np.random.default_rng() + + For these examples, we'll create three random data sets. The first + two, with sizes 35 and 25, are drawn from a normal distribution with + mean 0 and standard deviation 2. The third data set has size 25 and + is drawn from a normal distribution with standard deviation 1.25. + + >>> x1 = rng.normal(loc=0, scale=2, size=35) + >>> x2 = rng.normal(loc=0, scale=2, size=25) + >>> x3 = rng.normal(loc=0, scale=1.25, size=25) + + First we apply `ansari` to `x1` and `x2`. These samples are drawn + from the same distribution, so we expect the Ansari-Bradley test + should not lead us to conclude that the scales of the distributions + are different. + + >>> ansari(x1, x2) + AnsariResult(statistic=541.0, pvalue=0.9762532927399098) + + With a p-value close to 1, we cannot conclude that there is a + significant difference in the scales (as expected). + + Now apply the test to `x1` and `x3`: + + >>> ansari(x1, x3) + AnsariResult(statistic=425.0, pvalue=0.0003087020407974518) + + The probability of observing such an extreme value of the statistic + under the null hypothesis of equal scales is only 0.03087%. We take this + as evidence against the null hypothesis in favor of the alternative: + the scales of the distributions from which the samples were drawn + are not equal. + + We can use the `alternative` parameter to perform a one-tailed test. + In the above example, the scale of `x1` is greater than `x3` and so + the ratio of scales of `x1` and `x3` is greater than 1. This means + that the p-value when ``alternative='greater'`` should be near 0 and + hence we should be able to reject the null hypothesis: + + >>> ansari(x1, x3, alternative='greater') + AnsariResult(statistic=425.0, pvalue=0.0001543510203987259) + + As we can see, the p-value is indeed quite low. Use of + ``alternative='less'`` should thus yield a large p-value: + + >>> ansari(x1, x3, alternative='less') + AnsariResult(statistic=425.0, pvalue=0.9998643258449039) + + """ + xp = array_namespace(x, y) + dtype = xp_result_type(x, y, force_floating=True, xp=xp) + + if alternative not in {'two-sided', 'greater', 'less'}: + raise ValueError("'alternative' must be 'two-sided'," + " 'greater', or 'less'.") + methods = {'auto', 'asymptotic', 'exact'} + if str(method) not in methods and not isinstance(method, stats.PermutationMethod): + message = (f"`method` must be one of {methods} or " + "an instance of `PermutationMethod`.") + raise ValueError(message) + if is_lazy_array(x) and method == 'auto': + message = ("`method` must be 'exact', 'asymptotic', or an instance of " + "`PermutationMethod` when used with lazy arrays.") + raise ValueError(message) + + if not hasattr(_abw_state, 'a'): + _abw_state.a = _ABW() + + # _axis_nan_policy decorator guarantees that axis=-1 + n = x.shape[-1] + m = y.shape[-1] + if m < 1: # needed by test_axis_nan_policy; not user-facing + raise ValueError("Not enough other observations.") + if n < 1: + raise ValueError("Not enough test observations.") + + N = m + n + xy = xp.concat([x, y], axis=-1) # combine + rank, _, t = _stats_py._rankdata(xy, method='average', return_ties=True) + symrank = xp.minimum(rank, N - rank + 1) + AB = xp.sum(symrank[..., :n], axis=-1) + repeats = xp.any(t > 1) # in theory we could branch for each slice separately + + if method == "auto": + method = 'exact' if ((m < 55) and (n < 55) and not repeats) else 'asymptotic' + + if method == 'exact': + # np.vectorize converts to NumPy here, and we convert back to the result + # type before returning + cdf = np.vectorize(_abw_state.a.cdf, otypes=[np.float64]) + sf = np.vectorize(_abw_state.a.sf, otypes=[np.float64]) + def get_ansari_pvalue(AB): + if alternative == 'two-sided': + pval = 2.0 * np.minimum(cdf(AB, n, m), sf(AB, n, m)) + elif alternative == 'greater': + # AB statistic is _smaller_ when ratio of scales is larger, + # so this is the opposite of the usual calculation + pval = cdf(AB, n, m) + else: + pval = sf(AB, n, m) + return pval + + pval = xpx.lazy_apply(get_ansari_pvalue, AB, shape=AB.shape) + pval = xp.clip(xp.asarray(pval, dtype=dtype), max=1.0) + AB = AB[()] if AB.ndim == 0 else AB + pval = pval[()] if pval.ndim == 0 else pval + return AnsariResult(AB, pval) + + elif isinstance(method, stats.PermutationMethod): + def statistic(x, y, axis): + return ansari(x, y, axis=axis, method="asymptotic").statistic + alternative = dict(less='greater', greater='less').get(alternative, 'two-sided') + res = stats.permutation_test((x, y), statistic, axis=axis, + **method._asdict(), alternative=alternative) + return AnsariResult(AB, res.pvalue) + + mnAB = (n * (N + 1.0) ** 2 / 4.0 / N) if N % 2 else (n * (N + 2.0) / 4.0) + + if is_lazy_array(repeats) or repeats: # adjust variance estimates + # compute np.sum(tj * rj**2,axis=0) + fac = xp.sum(symrank**2, axis=-1) + if N % 2: # N odd + varAB = m * n * (16*N*fac - (N+1)**4) / (16.0 * N**2 * (N-1)) + else: # N even + varAB = m * n * (16*fac - N*(N+2)**2) / (16.0 * N * (N-1)) + else: + # otherwise compute normal approximation + if N % 2: # N odd + varAB = n * m * (N + 1.0) * (3 + N ** 2) / (48.0 * N ** 2) + else: + varAB = m * n * (N + 2) * (N - 2.0) / 48 / (N - 1.0) + varAB = xp.asarray(varAB, dtype=dtype) + + # Small values of AB indicate larger dispersion for the x sample. + # Large values of AB indicate larger dispersion for the y sample. + # This is opposite to the way we define the ratio of scales. see [1]_. + z = (mnAB - AB) / xp.sqrt(varAB) + pvalue = _get_pvalue(z, _SimpleNormal(), alternative, xp=xp) + AB = AB[()] if AB.ndim == 0 else AB + pvalue = pvalue[()] if pvalue.ndim == 0 else pvalue + return AnsariResult(AB, pvalue) + + +BartlettResult = namedtuple('BartlettResult', ('statistic', 'pvalue')) + +@xp_capabilities(marray=True) +@_axis_nan_policy_factory(BartlettResult, n_samples=None) +def bartlett(*samples, axis=0): + r"""Perform Bartlett's test for equal variances. + + Bartlett's test tests the null hypothesis that all input samples + are from populations with equal variances. For samples + from significantly non-normal populations, Levene's test + `levene` is more robust. + + Parameters + ---------- + *samples : array_like + arrays of sample data. Only 1d arrays are accepted, they may have + different lengths. + + Returns + ------- + statistic : float + The test statistic. + pvalue : float + The p-value of the test. + + See Also + -------- + fligner : A non-parametric test for the equality of k variances + levene : A robust parametric test for equality of k variances + :ref:`hypothesis_bartlett` : Extended example + + Notes + ----- + Conover et al. (1981) examine many of the existing parametric and + nonparametric tests by extensive simulations and they conclude that the + tests proposed by Fligner and Killeen (1976) and Levene (1960) appear to be + superior in terms of robustness of departures from normality and power + ([3]_). + + References + ---------- + .. [1] https://www.itl.nist.gov/div898/handbook/eda/section3/eda357.htm + .. [2] Snedecor, George W. and Cochran, William G. (1989), Statistical + Methods, Eighth Edition, Iowa State University Press. + .. [3] Park, C. and Lindsay, B. G. (1999). Robust Scale Estimation and + Hypothesis Testing based on Quadratic Inference Function. Technical + Report #99-03, Center for Likelihood Studies, Pennsylvania State + University. + .. [4] Bartlett, M. S. (1937). Properties of Sufficiency and Statistical + Tests. Proceedings of the Royal Society of London. Series A, + Mathematical and Physical Sciences, Vol. 160, No.901, pp. 268-282. + + Examples + -------- + + Test whether the lists `a`, `b` and `c` come from populations + with equal variances. + + >>> import numpy as np + >>> from scipy import stats + >>> a = [8.88, 9.12, 9.04, 8.98, 9.00, 9.08, 9.01, 8.85, 9.06, 8.99] + >>> b = [8.88, 8.95, 9.29, 9.44, 9.15, 9.58, 8.36, 9.18, 8.67, 9.05] + >>> c = [8.95, 9.12, 8.95, 8.85, 9.03, 8.84, 9.07, 8.98, 8.86, 8.98] + >>> stat, p = stats.bartlett(a, b, c) + >>> p + 1.1254782518834628e-05 + + The very small p-value suggests that the populations do not have equal + variances. + + This is not surprising, given that the sample variance of `b` is much + larger than that of `a` and `c`: + + >>> [np.var(x, ddof=1) for x in [a, b, c]] + [0.007054444444444413, 0.13073888888888888, 0.008890000000000002] + + For a more detailed example, see :ref:`hypothesis_bartlett`. + """ + xp = array_namespace(*samples) + + k = len(samples) + if k < 2: + raise ValueError("Must enter at least two input sample vectors.") + + if axis is None: + samples = [xp_ravel(sample) for sample in samples] + else: + samples = _broadcast_arrays(samples, axis=axis, xp=xp) + samples = [xp.moveaxis(sample, axis, -1) for sample in samples] + + Ni = [xp.asarray(_count_nonmasked(sample, axis=-1, xp=xp), + dtype=sample.dtype, device=xp_device(sample)) + for sample in samples] + Ni = [xp.broadcast_to(N, samples[0].shape[:-1]) for N in Ni] + ssq = [xp.var(sample, correction=1, axis=-1) for sample in samples] + Ni = [arr[xp.newaxis, ...] for arr in Ni] + ssq = [arr[xp.newaxis, ...] for arr in ssq] + Ni = xp.concat(Ni, axis=0) + Ni = xpx.at(Ni)[Ni == 0].set(xp.nan) + ssq = xp.concat(ssq, axis=0) + dtype = Ni.dtype + Ntot = xp.sum(Ni, axis=0) + spsq = xp.sum((Ni - 1)*ssq, axis=0, dtype=dtype) / (Ntot - k) + numer = ((Ntot - k) * xp.log(spsq) + - xp.sum((Ni - 1)*xp.log(ssq), axis=0, dtype=dtype)) + denom = (1 + 1/(3*(k - 1)) + * ((xp.sum(1/(Ni - 1), axis=0)) - 1/(Ntot - k))) + T = numer / denom + + chi2 = _SimpleChi2(xp.asarray(k-1, dtype=dtype, device=xp_device(T))) + pvalue = _get_pvalue(T, chi2, alternative='greater', symmetric=False, xp=xp) + + T = xp.clip(T, min=0., max=xp.inf) + T = T[()] if T.ndim == 0 else T + pvalue = pvalue[()] if pvalue.ndim == 0 else pvalue + + return BartlettResult(T, pvalue) + + +LeveneResult = namedtuple('LeveneResult', ('statistic', 'pvalue')) + + +@xp_capabilities(cpu_only=True, exceptions=['cupy'], marray=True) +@_axis_nan_policy_factory(LeveneResult, n_samples=None) +def levene(*samples, center='median', proportiontocut=0.05, axis=0): + r"""Perform Levene test for equal variances. + + The Levene test tests the null hypothesis that all input samples + are from populations with equal variances. Levene's test is an + alternative to Bartlett's test `bartlett` in the case where + there are significant deviations from normality. + + Parameters + ---------- + *samples : array_like + The sample data, possibly with different lengths. + center : {'mean', 'median', 'trimmed'}, optional + Which statistics to use to center data points within each sample. Default + is 'median'. + proportiontocut : float, optional + When `center` is 'trimmed', this gives the proportion of data points + to cut from each end. (See `scipy.stats.trim_mean`.) + Default is 0.05. + axis : int or tuple of ints, default: 0 + If an int or tuple of ints, the axis or axes of the input along which + to compute the statistic. The statistic of each axis-slice (e.g. row) + of the input will appear in a corresponding element of the output. + If ``None``, the input will be raveled before computing the statistic. + + Returns + ------- + statistic : float + The test statistic. + pvalue : float + The p-value for the test. + + See Also + -------- + fligner : A non-parametric test for the equality of k variances + bartlett : A parametric test for equality of k variances in normal samples + :ref:`hypothesis_levene` : Extended example + + Notes + ----- + Three variations of Levene's test are possible. The possibilities + and their recommended usages are: + + * 'median' : Recommended for skewed (non-normal) distributions> + * 'mean' : Recommended for symmetric, moderate-tailed distributions. + * 'trimmed' : Recommended for heavy-tailed distributions. + + The test version using the mean was proposed in the original article + of Levene ([2]_) while the median and trimmed mean have been studied by + Brown and Forsythe ([3]_), sometimes also referred to as Brown-Forsythe + test. + + References + ---------- + .. [1] https://www.itl.nist.gov/div898/handbook/eda/section3/eda35a.htm + .. [2] Levene, H. (1960). In Contributions to Probability and Statistics: + Essays in Honor of Harold Hotelling, I. Olkin et al. eds., + Stanford University Press, pp. 278-292. + .. [3] Brown, M. B. and Forsythe, A. B. (1974), Journal of the American + Statistical Association, 69, 364-367 + + Examples + -------- + + Test whether the lists `a`, `b` and `c` come from populations + with equal variances. + + >>> import numpy as np + >>> from scipy import stats + >>> a = [8.88, 9.12, 9.04, 8.98, 9.00, 9.08, 9.01, 8.85, 9.06, 8.99] + >>> b = [8.88, 8.95, 9.29, 9.44, 9.15, 9.58, 8.36, 9.18, 8.67, 9.05] + >>> c = [8.95, 9.12, 8.95, 8.85, 9.03, 8.84, 9.07, 8.98, 8.86, 8.98] + >>> stat, p = stats.levene(a, b, c) + >>> p + 0.002431505967249681 + + The small p-value suggests that the populations do not have equal + variances. + + This is not surprising, given that the sample variance of `b` is much + larger than that of `a` and `c`: + + >>> [np.var(x, ddof=1) for x in [a, b, c]] + [0.007054444444444413, 0.13073888888888888, 0.008890000000000002] + + For a more detailed example, see :ref:`hypothesis_levene`. + """ + xp = array_namespace(*samples) + + if center not in ['mean', 'median', 'trimmed']: + raise ValueError("center must be 'mean', 'median' or 'trimmed'.") + + k = len(samples) + if k < 2: + raise ValueError("Must provide at least two samples.") + + if center == 'median': + + def func(x): + return (xp.median(x, axis=-1, keepdims=True) + if (is_numpy(xp) or is_dask(xp)) + else stats.quantile(x, 0.5, axis=-1, keepdims=True)) + + elif center == 'mean': + + def func(x): + return xp.mean(x, axis=-1, keepdims=True) + + else: # center == 'trimmed' + + def func(x): + # keepdims=True doesn't currently work for Dask + return _stats_py.trim_mean(x, proportiontocut, axis=-1)[..., xp.newaxis] + + Nis = [_count_nonmasked(sample, axis=-1, keepdims=True, xp=xp) + for sample in samples] + Ycis = [func(sample) for sample in samples] + Ntot = sum(Nis) + + # compute Zij's + Zijs = [xp.abs(sample - Yc) for sample, Yc in zip(samples, Ycis)] + + # compute Zbari + Zbaris = [xp.mean(Zij, axis=-1, keepdims=True) for Zij in Zijs] + Zbar = sum(Ni*Zbari for Ni, Zbari in zip(Nis, Zbaris)) / Ntot + + # compute numerator and denominator + dfd = (Ntot - k) + numer = dfd * sum(Ni * (Zbari - Zbar)**2 + for Ni, Zbari in zip(Nis, Zbaris)) + dfn = (k - 1.0) + denom = dfn * sum(xp.sum((Zij - Zbari)**2, axis=-1, keepdims=True) + for Zij, Zbari in zip(Zijs, Zbaris)) + + W = numer / denom + W = xp.squeeze(W, axis=-1) + dfd = xp.squeeze(dfd, axis=-1) if is_marray(xp) else dfd + dfn, dfd = xp.asarray(dfn, dtype=W.dtype), xp.asarray(dfd, dtype=W.dtype) + pval = _get_pvalue(W, _SimpleF(dfn, dfd), 'greater', xp=xp) + W = W[()] if W.ndim == 0 else W + pval = pval[()] if pval.ndim == 0 else pval + return LeveneResult(W, pval) + + +FlignerResult = namedtuple('FlignerResult', ('statistic', 'pvalue')) + + +@xp_capabilities(skip_backends=[('dask.array', 'no rankdata')], + marray=True) +@_axis_nan_policy_factory(FlignerResult, n_samples=None) +def fligner(*samples, center='median', proportiontocut=0.05, axis=0): + r"""Perform Fligner-Killeen test for equality of variance. + + Fligner's test tests the null hypothesis that all input samples + are from populations with equal variances. Fligner-Killeen's test is + distribution free when populations are identical [2]_. + + Parameters + ---------- + *samples : array_like + Arrays of sample data. Need not be the same length. + center : {'mean', 'median', 'trimmed'}, optional + Which statistics to use to center data points within each sample. Default + is 'median'. + proportiontocut : float, optional + When `center` is 'trimmed', this gives the proportion of data points + to cut from each end. (See `scipy.stats.trim_mean`.) + Default is 0.05. + axis : int or tuple of ints, default: 0 + If an int or tuple of ints, the axis or axes of the input along which + to compute the statistic. The statistic of each axis-slice (e.g. row) + of the input will appear in a corresponding element of the output. + If ``None``, the input will be raveled before computing the statistic. + + Returns + ------- + statistic : float + The test statistic. + pvalue : float + The p-value for the hypothesis test. + + See Also + -------- + bartlett : A parametric test for equality of k variances in normal samples + levene : A robust parametric test for equality of k variances + :ref:`hypothesis_fligner` : Extended example + + Notes + ----- + As with Levene's test there are three variants of Fligner's test that + differ by the measure of central tendency used in the test. See `levene` + for more information. + + Conover et al. (1981) examine many of the existing parametric and + nonparametric tests by extensive simulations and they conclude that the + tests proposed by Fligner and Killeen (1976) and Levene (1960) appear to be + superior in terms of robustness of departures from normality and power + [3]_. + + References + ---------- + .. [1] Qu, A., Lindsay, B. G., and Li, B. (2000). Improving generalized + estimating equations using quadratic inference functions. + Biometrika, 87(4), 823-836. + :doi:`10.1093/biomet/87.4.823` + .. [2] Fligner, M.A. and Killeen, T.J. (1976). Distribution-free two-sample + tests for scale. Journal of the American Statistical Association. + 71(353), 210-213. + .. [3] Conover, W. J., Johnson, M. E. and Johnson M. M. (1981). A + comparative study of tests for homogeneity of variances, with + applications to the outer continental shelf bidding data. + Technometrics, 23(4), 351-361. + + Examples + -------- + + >>> import numpy as np + >>> from scipy import stats + + Test whether the lists `a`, `b` and `c` come from populations + with equal variances. + + >>> a = [8.88, 9.12, 9.04, 8.98, 9.00, 9.08, 9.01, 8.85, 9.06, 8.99] + >>> b = [8.88, 8.95, 9.29, 9.44, 9.15, 9.58, 8.36, 9.18, 8.67, 9.05] + >>> c = [8.95, 9.12, 8.95, 8.85, 9.03, 8.84, 9.07, 8.98, 8.86, 8.98] + >>> stat, p = stats.fligner(a, b, c) + >>> p + 0.00450826080004775 + + The small p-value suggests that the populations do not have equal + variances. + + This is not surprising, given that the sample variance of `b` is much + larger than that of `a` and `c`: + + >>> [np.var(x, ddof=1) for x in [a, b, c]] + [0.007054444444444413, 0.13073888888888888, 0.008890000000000002] + + For a more detailed example, see :ref:`hypothesis_fligner`. + """ + xp = array_namespace(*samples) + + if center not in ['mean', 'median', 'trimmed']: + raise ValueError("center must be 'mean', 'median' or 'trimmed'.") + + k = len(samples) + if k < 2: + raise ValueError("Must provide at least two samples.") + + samples = xp_promote(*samples, force_floating=True, xp=xp) + dtype = samples[0].dtype + + # Handle empty input + for sample in samples: + if sample.size == 0: + NaN = _get_nan(*samples, xp=xp) + return FlignerResult(NaN, NaN) + + if center == 'median': + + def func(x): + return (xp.median(x, axis=-1, keepdims=True) + if (is_numpy(xp) or is_dask(xp)) + else stats.quantile(x, 0.5, axis=-1, keepdims=True)) + + elif center == 'mean': + + def func(x): + return xp.mean(x, axis=-1, keepdims=True) + + else: # center == 'trimmed' + + def func(x): + # keepdims=True doesn't currently work for lazy arrays + return _stats_py.trim_mean(x, proportiontocut, axis=-1)[..., xp.newaxis] + + lengths = [sample.shape[-1] for sample in samples] + ni = [_count_nonmasked(sample, axis=-1, keepdims=True, xp=xp) + for sample in samples] + N = sum(ni) + ni = [xp.squeeze(ni_, axis=-1) for ni_ in ni] if is_marray(xp) else ni + + # Implementation follows [3] pg 355 F-K. + Xibar = [func(sample) for sample in samples] + Xij_Xibar = [xp.abs(sample - Xibar_) for sample, Xibar_ in zip(samples, Xibar)] + Xij_Xibar = xp.concat(Xij_Xibar, axis=-1) + ranks = stats.rankdata(Xij_Xibar, method='average', axis=-1) + a_Ni = special.ndtri(ranks / (2*(N + 1.0)) + 0.5) + + # [3] Equation 2.1 + splits = list(itertools.accumulate(lengths, initial=0)) + Ai = [a_Ni[..., i:j] for i, j in zip(splits[:-1], splits[1:])] + Aibar = [xp.mean(Ai_, axis=-1) for Ai_ in Ai] + abar = xp.mean(a_Ni, axis=-1) + V2 = xp.var(a_Ni, axis=-1, correction=1) + statistic = sum(ni_ * (Aibar_ - abar)**2 for ni_, Aibar_ in zip(ni, Aibar)) / V2 + + chi2 = _SimpleChi2(xp.asarray(k-1, dtype=dtype)) + pval = _get_pvalue(statistic, chi2, alternative='greater', symmetric=False, xp=xp) + return FlignerResult(statistic, pval) + + +def _mood_statistic_with_ties(x, y, t, m, n, N, xp): + # First equation of "Mood's Squared Rank Test", Mielke pg 313 + E_0_T = m * (N * N - 1) / 12 + + # m, n, N, t, and S are defined in the second paragraph of Mielke pg 312 + # The only difference is that our `t` has zeros interspersed with the relevant + # numbers to keep the array rectangular, but these terms add nothing to the sum. + S = xp.cumulative_sum(t, include_initial=True, axis=-1) + S_i, S_i_m1 = S[..., 1:], S[..., :-1] + # Second equation of "Mood's Squared Rank Test", Mielke pg 313 + varM = (m * n * (N + 1.0) * (N**2 - 4) / 180 + - m * n / (180 * N * (N - 1)) + * xp.sum(t * (t ** 2 - 1) * (t ** 2 - 4 + (15 * (N - S_i - S_i_m1) ** 2)), + axis=-1)) + + # There is a formula for Phi (`phi` in code) in terms of t, S, and Psi(I) at the + # bottom of Mielke pg 312. Psi(I) = [I - (N+1)/2]^2 is defined (with a mistake in + # the location of the ^2) at the beginning of "Mood's Squared Rank Test" (pg 313). + # To vectorize this calculation, let c = (N + 1) / 2, so Psi(I) = I^2 - 2*c*I + c^2. + # We sum each of these three parts of Psi separately using formulas for sums from a + # to b (inclusive) of terms I^2, I, and 1 where I takes on successive integers. + def sum_I2(a, b=None): + return (a * (a + 1) * (2 * a + 1) / 6 if b is None + else sum_I2(b) - sum_I2(a) + a**2) + + def sum_I(a, b=None): + return (a * (a + 1) / 2 if b is None + else sum_I(b) - sum_I(a) + a) + + def sum_1(a, b): + return (b - a) + 1 + + with np.errstate(invalid='ignore', divide='ignore'): + sum_I2 = sum_I2(S_i_m1 + 1, S_i) + sum_I = sum_I(S_i_m1 + 1, S_i) + sum_1 = sum_1(S_i_m1 + 1, S_i) + c = (N + 1) / 2 + phi = (sum_I2 - 2*c*sum_I + sum_1*c**2) / t + + phi = xpx.at(phi)[t == 0].set(0.) # where t = 0 we get NaNs; eliminate them + + # Mielke pg 312 defines `a` as the count of elements in sample `x` for each of the + # unique values in the combined sample. The tricky thing is getting these to line + # up with the locations of nonzero elements in `t`/`phi`. + x = xp.sort(x, axis=-1) + xy = xp.concat((x, y), axis=-1) + i = xp.argsort(xy, stable=True, axis=-1) + _, _, a = _stats_py._rankdata(x, method='average', return_ties=True) + + zeros = xp.zeros(a.shape[:-1] + (n,), dtype=a.dtype) + a = xp.concat((a, zeros), axis=-1) + a = xp.take_along_axis(a, i, axis=-1) + + # Mielke pg 312 defines test statistic `T` as the inner product `a` and `phi` + T = xp.vecdot(a, phi, axis=-1) + + return (T - E_0_T) / xp.sqrt(varM) + + +def _mood_statistic_no_ties(r, m, n, N, xp): + rx = r[..., :m] + M = xp.sum((rx - (N + 1.0) / 2) ** 2, axis=-1) + E_0_T = m * (N * N - 1.0) / 12 + varM = m * n * (N + 1.0) * (N + 2) * (N - 2) / 180 + return (M - E_0_T) / math.sqrt(varM) + + +def _mood_too_small(samples, kwargs, axis=-1): + x, y = samples + m = x.shape[axis] + n = y.shape[axis] + N = m + n + return N < 3 + + +@xp_capabilities(skip_backends=[('dask.array', 'no rankdata')]) +@_axis_nan_policy_factory(SignificanceResult, n_samples=2, too_small=_mood_too_small) +def mood(x, y, axis=0, alternative="two-sided"): + """Perform Mood's test for equal scale parameters. + + Mood's two-sample test for scale parameters is a non-parametric + test for the null hypothesis that two samples are drawn from the + same distribution with the same scale parameter. + + Parameters + ---------- + x, y : array_like + Arrays of sample data. There must be at least three observations + total. + axis : int, optional + The axis along which the samples are tested. `x` and `y` can be of + different length along `axis`. + If `axis` is None, `x` and `y` are flattened and the test is done on + all values in the flattened arrays. + alternative : {'two-sided', 'less', 'greater'}, optional + Defines the alternative hypothesis. Default is 'two-sided'. + The following options are available: + + * 'two-sided': the scales of the distributions underlying `x` and `y` + are different. + * 'less': the scale of the distribution underlying `x` is less than + the scale of the distribution underlying `y`. + * 'greater': the scale of the distribution underlying `x` is greater + than the scale of the distribution underlying `y`. + + .. versionadded:: 1.7.0 + + Returns + ------- + res : SignificanceResult + An object containing attributes: + + statistic : scalar or ndarray + The z-score for the hypothesis test. For 1-D inputs a scalar is + returned. + pvalue : scalar ndarray + The p-value for the hypothesis test. + + See Also + -------- + fligner : A non-parametric test for the equality of k variances + ansari : A non-parametric test for the equality of 2 variances + bartlett : A parametric test for equality of k variances in normal samples + levene : A parametric test for equality of k variances + + Notes + ----- + The data are assumed to be drawn from probability distributions ``f(x)`` + and ``f(x/s) / s`` respectively, for some probability density function f. + The null hypothesis is that ``s == 1``. + + For multi-dimensional arrays, if the inputs are of shapes + ``(n0, n1, n2, n3)`` and ``(n0, m1, n2, n3)``, then if ``axis=1``, the + resulting z and p values will have shape ``(n0, n2, n3)``. Note that + ``n1`` and ``m1`` don't have to be equal, but the other dimensions do. + + References + ---------- + [1] Mielke, Paul W. "Note on Some Squared Rank Tests with Existing Ties." + Technometrics, vol. 9, no. 2, 1967, pp. 312-14. JSTOR, + :doi:`10.2307/1266427`. Accessed 18 May 2022. + + Examples + -------- + >>> import numpy as np + >>> from scipy import stats + >>> rng = np.random.default_rng() + >>> x2 = rng.standard_normal((2, 45, 6, 7)) + >>> x1 = rng.standard_normal((2, 30, 6, 7)) + >>> res = stats.mood(x1, x2, axis=1) + >>> res.pvalue.shape + (2, 6, 7) + + Find the number of points where the difference in scale is not significant: + + >>> (res.pvalue > 0.1).sum() + 78 + + Perform the test with different scales: + + >>> x1 = rng.standard_normal((2, 30)) + >>> x2 = rng.standard_normal((2, 35)) * 10.0 + >>> stats.mood(x1, x2, axis=1) + SignificanceResult(statistic=array([-5.76174136, -6.12650783]), + pvalue=array([8.32505043e-09, 8.98287869e-10])) + + """ + xp = array_namespace(x, y) + x, y = xp_promote(x, y, force_floating=True, xp=xp) + + # _axis_nan_policy decorator ensures axis=-1 + xy = xp.concat((x, y), axis=-1) + + m = x.shape[-1] + n = y.shape[-1] + N = m + n + + if m == 0 or n == 0 or N < 3: # only needed for test_axis_nan_policy + NaN = _get_nan(x, y, xp=xp) + return SignificanceResult(NaN, NaN) + + # determine if any of the samples contain ties + # `a` represents ties within `x`; `t` represents ties within `xy` + r, _, t = _stats_py._rankdata(xy, method='average', return_ties=True) + + if is_lazy_array(t) or xp.any(t > 1): + z = _mood_statistic_with_ties(x, y, t, m, n, N, xp=xp) + else: + z = _mood_statistic_no_ties(r, m, n, N, xp=xp) + + pval = _get_pvalue(z, _SimpleNormal(), alternative, xp=xp) + + z = z[()] if z.ndim == 0 else z + pval = pval[()] if pval.ndim == 0 else pval + return SignificanceResult(z, pval) + + +WilcoxonResult = _make_tuple_bunch('WilcoxonResult', ['statistic', 'pvalue']) + + +def wilcoxon_result_unpacker(res, _): + if hasattr(res, 'zstatistic'): + return res.statistic, res.pvalue, res.zstatistic + else: + return res.statistic, res.pvalue + + +def wilcoxon_result_object(statistic, pvalue, zstatistic=None): + res = WilcoxonResult(statistic, pvalue) + if zstatistic is not None: + res.zstatistic = zstatistic + return res + + +def wilcoxon_outputs(kwds): + method = kwds.get('method', 'auto') + if method == 'asymptotic': + return 3 + return 2 + + +@xp_capabilities(skip_backends=[("dask.array", "no rankdata")], + cpu_only=True, # null distribution is CPU only + marray=True, + extra_note=("Only ``method='asymptotic'``/``zero_method='zsplit'`` is " + "compatible with MArray input. " + "``method='auto'`` is incompatible with JAX arrays.")) +@_rename_parameter("mode", "method") +@_axis_nan_policy_factory( + wilcoxon_result_object, paired=True, + n_samples=lambda kwds: 2 if kwds.get('y', None) is not None else 1, + result_to_tuple=wilcoxon_result_unpacker, n_outputs=wilcoxon_outputs, +) +def wilcoxon(x, y=None, zero_method="wilcox", correction=False, + alternative="two-sided", method='auto', *, axis=0): + """Calculate the Wilcoxon signed-rank test. + + The Wilcoxon signed-rank test tests the null hypothesis that two + related paired samples come from the same distribution. In particular, + it tests whether the distribution of the differences ``x - y`` is symmetric + about zero. It is a non-parametric version of the paired T-test. + + Parameters + ---------- + x : array_like + Either the first set of measurements (in which case ``y`` is the second + set of measurements), or the differences between two sets of + measurements (in which case ``y`` is not to be specified.) Must be + one-dimensional. + y : array_like, optional + Either the second set of measurements (if ``x`` is the first set of + measurements), or not specified (if ``x`` is the differences between + two sets of measurements.) Must be one-dimensional. + + .. warning:: + When `y` is provided, `wilcoxon` calculates the test statistic + based on the ranks of the absolute values of ``d = x - y``. + Roundoff error in the subtraction can result in elements of ``d`` + being assigned different ranks even when they would be tied with + exact arithmetic. Rather than passing `x` and `y` separately, + consider computing the difference ``x - y``, rounding as needed to + ensure that only truly unique elements are numerically distinct, + and passing the result as `x`, leaving `y` at the default (None). + + zero_method : {"wilcox", "pratt", "zsplit"}, optional + There are different conventions for handling pairs of observations + with equal values ("zero-differences", or "zeros"). + + * "wilcox": Discards all zero-differences (default); see [4]_. + * "pratt": Includes zero-differences in the ranking process, + but drops the ranks of the zeros (more conservative); see [3]_. + In this case, the normal approximation is adjusted as in [5]_. + * "zsplit": Includes zero-differences in the ranking process and + splits the zero rank between positive and negative ones. + + correction : bool, optional + If True, apply continuity correction by adjusting the Wilcoxon rank + statistic by 0.5 towards the mean value when computing the + z-statistic if a normal approximation is used. Default is False. + alternative : {"two-sided", "greater", "less"}, optional + Defines the alternative hypothesis. Default is 'two-sided'. + In the following, let ``d`` represent the difference between the paired + samples: ``d = x - y`` if both ``x`` and ``y`` are provided, or + ``d = x`` otherwise. + + * 'two-sided': the distribution underlying ``d`` is not symmetric + about zero. + * 'less': the distribution underlying ``d`` is stochastically less + than a distribution symmetric about zero. + * 'greater': the distribution underlying ``d`` is stochastically + greater than a distribution symmetric about zero. + + method : {"auto", "exact", "asymptotic"} or `PermutationMethod` instance, optional + Method to calculate the p-value, see Notes. Default is "auto". + + axis : int or None, default: 0 + If an int, the axis of the input along which to compute the statistic. + The statistic of each axis-slice (e.g. row) of the input will appear + in a corresponding element of the output. If ``None``, the input will + be raveled before computing the statistic. + + Returns + ------- + result : WilcoxonResult + An object with the following attributes. + + statistic : array_like + If `alternative` is "two-sided", the sum of the ranks of the + differences above or below zero, whichever is smaller. + Otherwise the sum of the ranks of the differences above zero. + pvalue : array_like + The p-value for the test depending on `alternative` and `method`. + zstatistic : array_like + When ``method = 'asymptotic'``, this is the normalized z-statistic:: + + z = (T - mn - d) / se + + where ``T`` is `statistic` as defined above, ``mn`` is the mean of the + distribution under the null hypothesis, ``d`` is a continuity + correction, and ``se`` is the standard error. + When ``method != 'asymptotic'``, this attribute is not available. + + See Also + -------- + kruskal, mannwhitneyu + + Notes + ----- + In the following, let ``d`` represent the difference between the paired + samples: ``d = x - y`` if both ``x`` and ``y`` are provided, or ``d = x`` + otherwise. Assume that all elements of ``d`` are independent and + identically distributed observations, and all are distinct and nonzero. + + - When ``len(d)`` is sufficiently large, the null distribution of the + normalized test statistic (`zstatistic` above) is approximately normal, + and ``method = 'asymptotic'`` can be used to compute the p-value. + - When ``len(d)`` is small, the normal approximation may not be accurate, + and ``method='exact'`` is preferred (at the cost of additional + execution time). + - The default, ``method='auto'``, selects between the two: + ``method='exact'`` is used when ``len(d) <= 50``, and + ``method='asymptotic'`` is used otherwise. + + The presence of "ties" (i.e. not all elements of ``d`` are unique) or + "zeros" (i.e. elements of ``d`` are zero) changes the null distribution + of the test statistic, and ``method='exact'`` no longer calculates + the exact p-value. If ``method='asymptotic'``, the z-statistic is adjusted + for more accurate comparison against the standard normal, but still, + for finite sample sizes, the standard normal is only an approximation of + the true null distribution of the z-statistic. For such situations, the + `method` parameter also accepts instances of `PermutationMethod`. In this + case, the p-value is computed using `permutation_test` with the provided + configuration options and other appropriate settings. + + The presence of ties and zeros affects the resolution of ``method='auto'`` + accordingly: exhasutive permutations are performed when ``len(d) <= 13``, + and the asymptotic method is used otherwise. Note that they asymptotic + method may not be very accurate even for ``len(d) > 14``; the threshold + was chosen as a compromise between execution time and accuracy under the + constraint that the results must be deterministic. Consider providing an + instance of `PermutationMethod` method manually, choosing the + ``n_resamples`` parameter to balance time constraints and accuracy + requirements. + + Please also note that in the edge case that all elements of ``d`` are zero, + the p-value relying on the normal approximaton cannot be computed (NaN) + if ``zero_method='wilcox'`` or ``zero_method='pratt'``. + + References + ---------- + .. [1] https://en.wikipedia.org/wiki/Wilcoxon_signed-rank_test + .. [2] Conover, W.J., Practical Nonparametric Statistics, 1971. + .. [3] Pratt, J.W., Remarks on Zeros and Ties in the Wilcoxon Signed + Rank Procedures, Journal of the American Statistical Association, + Vol. 54, 1959, pp. 655-667. :doi:`10.1080/01621459.1959.10501526` + .. [4] Wilcoxon, F., Individual Comparisons by Ranking Methods, + Biometrics Bulletin, Vol. 1, 1945, pp. 80-83. :doi:`10.2307/3001968` + .. [5] Cureton, E.E., The Normal Approximation to the Signed-Rank + Sampling Distribution When Zero Differences are Present, + Journal of the American Statistical Association, Vol. 62, 1967, + pp. 1068-1069. :doi:`10.1080/01621459.1967.10500917` + + Examples + -------- + In [4]_, the differences in height between cross- and self-fertilized + corn plants is given as follows: + + >>> d = [6, 8, 14, 16, 23, 24, 28, 29, 41, -48, 49, 56, 60, -67, 75] + + Cross-fertilized plants appear to be higher. To test the null + hypothesis that there is no height difference, we can apply the + two-sided test: + + >>> from scipy.stats import wilcoxon + >>> res = wilcoxon(d) + >>> res.statistic, res.pvalue + (24.0, 0.041259765625) + + Hence, we would reject the null hypothesis at a confidence level of 5%, + concluding that there is a difference in height between the groups. + To confirm that the median of the differences can be assumed to be + positive, we use: + + >>> res = wilcoxon(d, alternative='greater') + >>> res.statistic, res.pvalue + (96.0, 0.0206298828125) + + This shows that the null hypothesis that the median is negative can be + rejected at a confidence level of 5% in favor of the alternative that + the median is greater than zero. The p-values above are exact. Using the + normal approximation gives very similar values: + + >>> res = wilcoxon(d, method='asymptotic') + >>> res.statistic, res.pvalue + (24.0, 0.04088813291185591) + + Note that the statistic changed to 96 in the one-sided case (the sum + of ranks of positive differences) whereas it is 24 in the two-sided + case (the minimum of sum of ranks above and below zero). + + In the example above, the differences in height between paired plants are + provided to `wilcoxon` directly. Alternatively, `wilcoxon` accepts two + samples of equal length, calculates the differences between paired + elements, then performs the test. Consider the samples ``x`` and ``y``: + + >>> import numpy as np + >>> x = np.array([0.5, 0.825, 0.375, 0.5]) + >>> y = np.array([0.525, 0.775, 0.325, 0.55]) + >>> res = wilcoxon(x, y, alternative='greater') + >>> res + WilcoxonResult(statistic=5.0, pvalue=0.5625) + + Note that had we calculated the differences by hand, the test would have + produced different results: + + >>> d = [-0.025, 0.05, 0.05, -0.05] + >>> ref = wilcoxon(d, alternative='greater') + >>> ref + WilcoxonResult(statistic=6.0, pvalue=0.5) + + The substantial difference is due to roundoff error in the results of + ``x-y``: + + >>> d - (x-y) + array([2.08166817e-17, 6.93889390e-17, 1.38777878e-17, 4.16333634e-17]) + + Even though we expected all the elements of ``(x-y)[1:]`` to have the same + magnitude ``0.05``, they have slightly different magnitudes in practice, + and therefore are assigned different ranks in the test. Before performing + the test, consider calculating ``d`` and adjusting it as necessary to + ensure that theoretically identically values are not numerically distinct. + For example: + + >>> d2 = np.around(x - y, decimals=3) + >>> wilcoxon(d2, alternative='greater') + WilcoxonResult(statistic=6.0, pvalue=0.5) + + """ + # replace approx by asymptotic to ensure backwards compatability + if method == "approx": + method = "asymptotic" + return _wilcoxon._wilcoxon_nd(x, y, zero_method, correction, alternative, + method, axis) + + +MedianTestResult = _make_tuple_bunch( + 'MedianTestResult', + ['statistic', 'pvalue', 'median', 'table'], [] +) + + +@xp_capabilities(np_only=True) +def median_test(*samples, ties='below', correction=True, lambda_=1, + nan_policy='propagate'): + """Perform a Mood's median test. + + Test that two or more samples come from populations with the same median. + + Let ``n = len(samples)`` be the number of samples. The "grand median" of + all the data is computed, and a contingency table is formed by + classifying the values in each sample as being above or below the grand + median. The contingency table, along with `correction` and `lambda_`, + are passed to `scipy.stats.chi2_contingency` to compute the test statistic + and p-value. + + Parameters + ---------- + *samples : array_like + The set of samples. There must be at least two samples. + Each sample must be a one-dimensional sequence containing at least + one value. The samples are not required to have the same length. + ties : str, optional + Determines how values equal to the grand median are classified in + the contingency table. The string must be one of:: + + "below": + Values equal to the grand median are counted as "below". + "above": + Values equal to the grand median are counted as "above". + "ignore": + Values equal to the grand median are not counted. + + The default is "below". + correction : bool, optional + If True, *and* there are just two samples, apply Yates' correction + for continuity when computing the test statistic associated with + the contingency table. Default is True. + lambda_ : float or str, optional + By default, the statistic computed in this test is Pearson's + chi-squared statistic. `lambda_` allows a statistic from the + Cressie-Read power divergence family to be used instead. See + `power_divergence` for details. + Default is 1 (Pearson's chi-squared statistic). + nan_policy : {'propagate', 'raise', 'omit'}, optional + Defines how to handle when input contains nan. 'propagate' returns nan, + 'raise' throws an error, 'omit' performs the calculations ignoring nan + values. Default is 'propagate'. + + Returns + ------- + res : MedianTestResult + An object containing attributes: + + statistic : float + The test statistic. The statistic that is returned is determined + by `lambda_`. The default is Pearson's chi-squared statistic. + pvalue : float + The p-value of the test. + median : float + The grand median. + table : ndarray + The contingency table. The shape of the table is (2, n), where + n is the number of samples. The first row holds the counts of the + values above the grand median, and the second row holds the counts + of the values below the grand median. The table allows further + analysis with, for example, `scipy.stats.chi2_contingency`, or with + `scipy.stats.fisher_exact` if there are two samples, without having + to recompute the table. If ``nan_policy`` is "propagate" and there + are nans in the input, the return value for ``table`` is ``None``. + + See Also + -------- + kruskal : Compute the Kruskal-Wallis H-test for independent samples. + mannwhitneyu : Computes the Mann-Whitney rank test on samples x and y. + + Notes + ----- + .. versionadded:: 0.15.0 + + References + ---------- + .. [1] Mood, A. M., Introduction to the Theory of Statistics. McGraw-Hill + (1950), pp. 394-399. + .. [2] Zar, J. H., Biostatistical Analysis, 5th ed. Prentice Hall (2010). + See Sections 8.12 and 10.15. + + Examples + -------- + A biologist runs an experiment in which there are three groups of plants. + Group 1 has 16 plants, group 2 has 15 plants, and group 3 has 17 plants. + Each plant produces a number of seeds. The seed counts for each group + are:: + + Group 1: 10 14 14 18 20 22 24 25 31 31 32 39 43 43 48 49 + Group 2: 28 30 31 33 34 35 36 40 44 55 57 61 91 92 99 + Group 3: 0 3 9 22 23 25 25 33 34 34 40 45 46 48 62 67 84 + + The following code applies Mood's median test to these samples. + + >>> g1 = [10, 14, 14, 18, 20, 22, 24, 25, 31, 31, 32, 39, 43, 43, 48, 49] + >>> g2 = [28, 30, 31, 33, 34, 35, 36, 40, 44, 55, 57, 61, 91, 92, 99] + >>> g3 = [0, 3, 9, 22, 23, 25, 25, 33, 34, 34, 40, 45, 46, 48, 62, 67, 84] + >>> from scipy.stats import median_test + >>> res = median_test(g1, g2, g3) + + The median is + + >>> res.median + 34.0 + + and the contingency table is + + >>> res.table + array([[ 5, 10, 7], + [11, 5, 10]]) + + `p` is too large to conclude that the medians are not the same: + + >>> res.pvalue + 0.12609082774093244 + + The "G-test" can be performed by passing ``lambda_="log-likelihood"`` to + `median_test`. + + >>> res = median_test(g1, g2, g3, lambda_="log-likelihood") + >>> res.pvalue + 0.12224779737117837 + + The median occurs several times in the data, so we'll get a different + result if, for example, ``ties="above"`` is used: + + >>> res = median_test(g1, g2, g3, ties="above") + >>> res.pvalue + 0.063873276069553273 + + >>> res.table + array([[ 5, 11, 9], + [11, 4, 8]]) + + This example demonstrates that if the data set is not large and there + are values equal to the median, the p-value can be sensitive to the + choice of `ties`. + + """ + if len(samples) < 2: + raise ValueError('median_test requires two or more samples.') + + ties_options = ['below', 'above', 'ignore'] + if ties not in ties_options: + raise ValueError(f"invalid 'ties' option '{ties}'; 'ties' must be one " + f"of: {str(ties_options)[1:-1]}") + + data = [np.asarray(sample) for sample in samples] + + # Validate the sizes and shapes of the arguments. + for k, d in enumerate(data): + if d.size == 0: + raise ValueError(f"Sample {k + 1} is empty. All samples must " + f"contain at least one value.") + if d.ndim != 1: + raise ValueError(f"Sample {k + 1} has {d.ndim} dimensions. " + f"All samples must be one-dimensional sequences.") + + cdata = np.concatenate(data) + contains_nan = _contains_nan(cdata, nan_policy) + if nan_policy == 'propagate' and contains_nan: + return MedianTestResult(np.nan, np.nan, np.nan, None) + + if contains_nan: + grand_median = np.median(cdata[~np.isnan(cdata)]) + else: + grand_median = np.median(cdata) + # When the minimum version of numpy supported by scipy is 1.9.0, + # the above if/else statement can be replaced by the single line: + # grand_median = np.nanmedian(cdata) + + # Create the contingency table. + table = np.zeros((2, len(data)), dtype=np.int64) + for k, sample in enumerate(data): + sample = sample[~np.isnan(sample)] + + nabove = count_nonzero(sample > grand_median) + nbelow = count_nonzero(sample < grand_median) + nequal = sample.size - (nabove + nbelow) + table[0, k] += nabove + table[1, k] += nbelow + if ties == "below": + table[1, k] += nequal + elif ties == "above": + table[0, k] += nequal + + # Check that no row or column of the table is all zero. + # Such a table can not be given to chi2_contingency, because it would have + # a zero in the table of expected frequencies. + rowsums = table.sum(axis=1) + if rowsums[0] == 0: + raise ValueError(f"All values are below the grand median ({grand_median}).") + if rowsums[1] == 0: + raise ValueError(f"All values are above the grand median ({grand_median}).") + if ties == "ignore": + # We already checked that each sample has at least one value, but it + # is possible that all those values equal the grand median. If `ties` + # is "ignore", that would result in a column of zeros in `table`. We + # check for that case here. + zero_cols = np.nonzero((table == 0).all(axis=0))[0] + if len(zero_cols) > 0: + raise ValueError( + f"All values in sample {zero_cols[0] + 1} are equal to the grand " + f"median ({grand_median!r}), so they are ignored, resulting in an " + f"empty sample." + ) + + stat, p, dof, expected = chi2_contingency(table, lambda_=lambda_, + correction=correction) + return MedianTestResult(stat, p, grand_median, table) + + +def _circfuncs_common(samples, period, xp=None): + xp = array_namespace(samples) if xp is None else xp + + samples = xp_promote(samples, force_floating=True, xp=xp) + + # Recast samples as radians that range between 0 and 2 pi and calculate + # the sine and cosine + scaled_samples = samples * ((2.0 * pi) / period) + sin_samp = xp.sin(scaled_samples) + cos_samp = xp.cos(scaled_samples) + + return samples, sin_samp, cos_samp + + +@xp_capabilities(marray=True) +@_axis_nan_policy_factory( + lambda x: x, n_outputs=1, default_axis=None, + result_to_tuple=lambda x, _: (x,) +) +def circmean(samples, high=2*pi, low=0, axis=None, nan_policy='propagate'): + r"""Compute the circular mean of a sample of angle observations. + + Given :math:`n` angle observations :math:`x_1, \cdots, x_n` measured in + radians, their *circular mean* is defined by ([1]_, Eq. 2.2.4) + + .. math:: + + \mathrm{Arg} \left( \frac{1}{n} \sum_{k=1}^n e^{i x_k} \right) + + where :math:`i` is the imaginary unit and :math:`\mathop{\mathrm{Arg}} z` + gives the principal value of the argument of complex number :math:`z`, + restricted to the range :math:`[0,2\pi]` by default. :math:`z` in the + above expression is known as the `mean resultant vector`. + + Parameters + ---------- + samples : array_like + Input array of angle observations. The value of a full angle is + equal to ``(high - low)``. + high : float, optional + Upper boundary of the principal value of an angle. Default is ``2*pi``. + low : float, optional + Lower boundary of the principal value of an angle. Default is ``0``. + + Returns + ------- + circmean : float + Circular mean, restricted to the range ``[low, high]``. + + If the mean resultant vector is zero, an input-dependent, + implementation-defined number between ``[low, high]`` is returned. + If the input array is empty, ``np.nan`` is returned. + + See Also + -------- + circstd : Circular standard deviation. + circvar : Circular variance. + + References + ---------- + .. [1] Mardia, K. V. and Jupp, P. E. *Directional Statistics*. + John Wiley & Sons, 1999. + + Examples + -------- + For readability, all angles are printed out in degrees. + + >>> import numpy as np + >>> from scipy.stats import circmean + >>> import matplotlib.pyplot as plt + >>> angles = np.deg2rad(np.array([20, 30, 330])) + >>> circmean = circmean(angles) + >>> np.rad2deg(circmean) + 7.294976657784009 + + >>> mean = angles.mean() + >>> np.rad2deg(mean) + 126.66666666666666 + + Plot and compare the circular mean against the arithmetic mean. + + >>> plt.plot(np.cos(np.linspace(0, 2*np.pi, 500)), + ... np.sin(np.linspace(0, 2*np.pi, 500)), + ... c='k') + >>> plt.scatter(np.cos(angles), np.sin(angles), c='k') + >>> plt.scatter(np.cos(circmean), np.sin(circmean), c='b', + ... label='circmean') + >>> plt.scatter(np.cos(mean), np.sin(mean), c='r', label='mean') + >>> plt.legend() + >>> plt.axis('equal') + >>> plt.show() + + """ + xp = array_namespace(samples) + # Needed for non-NumPy arrays to get appropriate NaN result + # Apparently atan2(0, 0) is 0, even though it is mathematically undefined + if xp_size(samples) == 0: + return xp.mean(samples, axis=axis) + period = high - low + samples, sin_samp, cos_samp = _circfuncs_common(samples, period, xp=xp) + sin_sum = xp.sum(sin_samp, axis=axis) + cos_sum = xp.sum(cos_samp, axis=axis) + res = xp.atan2(sin_sum, cos_sum) + + res = res[()] if res.ndim == 0 else res + return (res * (period / (2.0 * pi)) - low) % period + low + + +@xp_capabilities(marray=True) +@_axis_nan_policy_factory( + lambda x: x, n_outputs=1, default_axis=None, + result_to_tuple=lambda x, _: (x,) +) +def circvar(samples, high=2*pi, low=0, axis=None, nan_policy='propagate'): + r"""Compute the circular variance of a sample of angle observations. + + Given :math:`n` angle observations :math:`x_1, \cdots, x_n` measured in + radians, their *circular variance* is defined by ([2]_, Eq. 2.3.3) + + .. math:: + + 1 - \left| \frac{1}{n} \sum_{k=1}^n e^{i x_k} \right| + + where :math:`i` is the imaginary unit and :math:`|z|` gives the length + of the complex number :math:`z`. :math:`|z|` in the above expression + is known as the `mean resultant length`. + + Parameters + ---------- + samples : array_like + Input array of angle observations. The value of a full angle is + equal to ``(high - low)``. + high : float, optional + Upper boundary of the principal value of an angle. Default is ``2*pi``. + low : float, optional + Lower boundary of the principal value of an angle. Default is ``0``. + + Returns + ------- + circvar : float + Circular variance. The returned value is in the range ``[0, 1]``, + where ``0`` indicates no variance and ``1`` indicates large variance. + + If the input array is empty, ``np.nan`` is returned. + + See Also + -------- + circmean : Circular mean. + circstd : Circular standard deviation. + + Notes + ----- + In the limit of small angles, the circular variance is close to + half the 'linear' variance if measured in radians. + + References + ---------- + .. [1] Fisher, N.I. *Statistical analysis of circular data*. Cambridge + University Press, 1993. + .. [2] Mardia, K. V. and Jupp, P. E. *Directional Statistics*. + John Wiley & Sons, 1999. + + Examples + -------- + >>> import numpy as np + >>> from scipy.stats import circvar + >>> import matplotlib.pyplot as plt + >>> samples_1 = np.array([0.072, -0.158, 0.077, 0.108, 0.286, + ... 0.133, -0.473, -0.001, -0.348, 0.131]) + >>> samples_2 = np.array([0.111, -0.879, 0.078, 0.733, 0.421, + ... 0.104, -0.136, -0.867, 0.012, 0.105]) + >>> circvar_1 = circvar(samples_1) + >>> circvar_2 = circvar(samples_2) + + Plot the samples. + + >>> fig, (left, right) = plt.subplots(ncols=2) + >>> for image in (left, right): + ... image.plot(np.cos(np.linspace(0, 2*np.pi, 500)), + ... np.sin(np.linspace(0, 2*np.pi, 500)), + ... c='k') + ... image.axis('equal') + ... image.axis('off') + >>> left.scatter(np.cos(samples_1), np.sin(samples_1), c='k', s=15) + >>> left.set_title(f"circular variance: {np.round(circvar_1, 2)!r}") + >>> right.scatter(np.cos(samples_2), np.sin(samples_2), c='k', s=15) + >>> right.set_title(f"circular variance: {np.round(circvar_2, 2)!r}") + >>> plt.show() + + """ + xp = array_namespace(samples) + period = high - low + samples, sin_samp, cos_samp = _circfuncs_common(samples, period, xp=xp) + sin_mean = xp.mean(sin_samp, axis=axis) + cos_mean = xp.mean(cos_samp, axis=axis) + hypotenuse = (sin_mean**2. + cos_mean**2.)**0.5 + # hypotenuse can go slightly above 1 due to rounding errors + R = xp.clip(hypotenuse, max=1.) + + res = 1. - R + return res + + +@xp_capabilities(marray=True) +@_axis_nan_policy_factory( + lambda x: x, n_outputs=1, default_axis=None, + result_to_tuple=lambda x, _: (x,) +) +def circstd(samples, high=2*pi, low=0, axis=None, nan_policy='propagate', *, + normalize=False): + r""" + Compute the circular standard deviation of a sample of angle observations. + + Given :math:`n` angle observations :math:`x_1, \cdots, x_n` measured in + radians, their `circular standard deviation` is defined by + ([2]_, Eq. 2.3.11) + + .. math:: + + \sqrt{ -2 \log \left| \frac{1}{n} \sum_{k=1}^n e^{i x_k} \right| } + + where :math:`i` is the imaginary unit and :math:`|z|` gives the length + of the complex number :math:`z`. :math:`|z|` in the above expression + is known as the `mean resultant length`. + + Parameters + ---------- + samples : array_like + Input array of angle observations. The value of a full angle is + equal to ``(high - low)``. + high : float, optional + Upper boundary of the principal value of an angle. Default is ``2*pi``. + low : float, optional + Lower boundary of the principal value of an angle. Default is ``0``. + axis : int or None, default: None + If an int, the axis of the input along which to compute the statistic. + The statistic of each axis-slice (e.g. row) of the input will appear in a + corresponding element of the output. + If ``None``, the input will be raveled before computing the statistic. + nan_policy : {'propagate', 'omit', 'raise'} + Defines how to handle input NaNs. + + - ``propagate``: if a NaN is present in the axis slice (e.g. row) along + which the statistic is computed, the corresponding entry of the output + will be NaN. + - ``omit``: NaNs will be omitted when performing the calculation. + If insufficient data remains in the axis slice along which the + statistic is computed, the corresponding entry of the output will be + NaN. + - ``raise``: if a NaN is present, a ``ValueError`` will be raised. + + normalize : bool, optional + If ``False`` (the default), the return value is computed from the + above formula with the input scaled by ``(2*pi)/(high-low)`` and + the output scaled (back) by ``(high-low)/(2*pi)``. If ``True``, + the output is not scaled and is returned directly. + + Returns + ------- + circstd : float + Circular standard deviation, optionally normalized. + + If the input array is empty, ``np.nan`` is returned. + + See Also + -------- + circmean : Circular mean. + circvar : Circular variance. + + Notes + ----- + In the limit of small angles, the circular standard deviation is close + to the 'linear' standard deviation if ``normalize`` is ``False``. + + References + ---------- + .. [1] Mardia, K. V. (1972). 2. In *Statistics of Directional Data* + (pp. 18-24). Academic Press. :doi:`10.1016/C2013-0-07425-7`. + .. [2] Mardia, K. V. and Jupp, P. E. *Directional Statistics*. + John Wiley & Sons, 1999. + + Examples + -------- + >>> import numpy as np + >>> from scipy.stats import circstd + >>> import matplotlib.pyplot as plt + >>> samples_1 = np.array([0.072, -0.158, 0.077, 0.108, 0.286, + ... 0.133, -0.473, -0.001, -0.348, 0.131]) + >>> samples_2 = np.array([0.111, -0.879, 0.078, 0.733, 0.421, + ... 0.104, -0.136, -0.867, 0.012, 0.105]) + >>> circstd_1 = circstd(samples_1) + >>> circstd_2 = circstd(samples_2) + + Plot the samples. + + >>> fig, (left, right) = plt.subplots(ncols=2) + >>> for image in (left, right): + ... image.plot(np.cos(np.linspace(0, 2*np.pi, 500)), + ... np.sin(np.linspace(0, 2*np.pi, 500)), + ... c='k') + ... image.axis('equal') + ... image.axis('off') + >>> left.scatter(np.cos(samples_1), np.sin(samples_1), c='k', s=15) + >>> left.set_title(f"circular std: {np.round(circstd_1, 2)!r}") + >>> right.plot(np.cos(np.linspace(0, 2*np.pi, 500)), + ... np.sin(np.linspace(0, 2*np.pi, 500)), + ... c='k') + >>> right.scatter(np.cos(samples_2), np.sin(samples_2), c='k', s=15) + >>> right.set_title(f"circular std: {np.round(circstd_2, 2)!r}") + >>> plt.show() + + """ + xp = array_namespace(samples) + period = high - low + samples, sin_samp, cos_samp = _circfuncs_common(samples, period, xp=xp) + sin_mean = xp.mean(sin_samp, axis=axis) # [1] (2.2.3) + cos_mean = xp.mean(cos_samp, axis=axis) # [1] (2.2.3) + hypotenuse = (sin_mean**2. + cos_mean**2.)**0.5 + # hypotenuse can go slightly above 1 due to rounding errors + R = xp.clip(hypotenuse, max=1.) # [1] (2.2.4) + + res = (-2*xp.log(R))**0.5+0.0 # torch.pow returns -0.0 if R==1 + if not normalize: + res *= (high-low)/(2.*pi) # [1] (2.3.14) w/ (2.3.7) + return res + + +class DirectionalStats: + def __init__(self, mean_direction, mean_resultant_length): + self.mean_direction = mean_direction + self.mean_resultant_length = mean_resultant_length + + def __repr__(self): + return (f"DirectionalStats(mean_direction={self.mean_direction}," + f" mean_resultant_length={self.mean_resultant_length})") + + +@xp_capabilities(marray=True) +def directional_stats(samples, *, axis=0, normalize=True): + """ + Computes sample statistics for directional data. + + Computes the directional mean (also called the mean direction vector) and + mean resultant length of a sample of vectors. + + The directional mean is a measure of "preferred direction" of vector data. + It is analogous to the sample mean, but it is for use when the length of + the data is irrelevant (e.g. unit vectors). + + The mean resultant length is a value between 0 and 1 used to quantify the + dispersion of directional data: the smaller the mean resultant length, the + greater the dispersion. Several definitions of directional variance + involving the mean resultant length are given in [1]_ and [2]_. + + Parameters + ---------- + samples : array_like + Input array. Must be at least two-dimensional, and the last axis of the + input must correspond with the dimensionality of the vector space. + When the input is exactly two dimensional, this means that each row + of the data is a vector observation. + axis : int, default: 0 + Axis along which the directional mean is computed. + normalize : bool, default: True + If True, normalize the input to ensure that each observation is a + unit vector. It the observations are already unit vectors, consider + setting this to False to avoid unnecessary computation. + + Returns + ------- + res : DirectionalStats + An object containing attributes: + + mean_direction : ndarray + Directional mean. + mean_resultant_length : ndarray + The mean resultant length [1]_. + + See Also + -------- + circmean: circular mean; i.e. directional mean for 2D *angles* + circvar: circular variance; i.e. directional variance for 2D *angles* + + Notes + ----- + This uses a definition of directional mean from [1]_. + Assuming the observations are unit vectors, the calculation is as follows. + + .. code-block:: python + + mean = samples.mean(axis=0) + mean_resultant_length = np.linalg.norm(mean) + mean_direction = mean / mean_resultant_length + + This definition is appropriate for *directional* data (i.e. vector data + for which the magnitude of each observation is irrelevant) but not + for *axial* data (i.e. vector data for which the magnitude and *sign* of + each observation is irrelevant). + + Several definitions of directional variance involving the mean resultant + length ``R`` have been proposed, including ``1 - R`` [1]_, ``1 - R**2`` + [2]_, and ``2 * (1 - R)`` [2]_. Rather than choosing one, this function + returns ``R`` as attribute `mean_resultant_length` so the user can compute + their preferred measure of dispersion. + + References + ---------- + .. [1] Mardia, Jupp. (2000). *Directional Statistics* + (p. 163). Wiley. + + .. [2] https://en.wikipedia.org/wiki/Directional_statistics + + Examples + -------- + >>> import numpy as np + >>> from scipy.stats import directional_stats + >>> data = np.array([[3, 4], # first observation, 2D vector space + ... [6, -8]]) # second observation + >>> dirstats = directional_stats(data) + >>> dirstats.mean_direction + array([1., 0.]) + + In contrast, the regular sample mean of the vectors would be influenced + by the magnitude of each observation. Furthermore, the result would not be + a unit vector. + + >>> data.mean(axis=0) + array([4.5, -2.]) + + An exemplary use case for `directional_stats` is to find a *meaningful* + center for a set of observations on a sphere, e.g. geographical locations. + + >>> data = np.array([[0.8660254, 0.5, 0.], + ... [0.8660254, -0.5, 0.]]) + >>> dirstats = directional_stats(data) + >>> dirstats.mean_direction + array([1., 0., 0.]) + + The regular sample mean on the other hand yields a result which does not + lie on the surface of the sphere. + + >>> data.mean(axis=0) + array([0.8660254, 0., 0.]) + + The function also returns the mean resultant length, which + can be used to calculate a directional variance. For example, using the + definition ``Var(z) = 1 - R`` from [2]_ where ``R`` is the + mean resultant length, we can calculate the directional variance of the + vectors in the above example as: + + >>> 1 - dirstats.mean_resultant_length + 0.13397459716167093 + """ + xp = array_namespace(samples) + samples = xp.asarray(samples) + + if samples.ndim < 2: + raise ValueError("samples must at least be two-dimensional. " + f"Instead samples has shape: {tuple(samples.shape)}") + samples = xp.moveaxis(samples, axis, 0) + + if is_marray(xp): + _xp = array_namespace(samples.mask) + mask = _xp.any(samples.mask, axis=-1, keepdims=True) + samples = xp.asarray(samples.data, mask=mask) + + if normalize: + vectornorms = xp_vector_norm(samples, axis=-1, keepdims=True, xp=xp) + samples = samples/vectornorms + mean = xp.mean(samples, axis=0) + mean_resultant_length = xp_vector_norm(mean, axis=-1, keepdims=True, xp=xp) + mean_direction = mean / mean_resultant_length + mrl = xp.squeeze(mean_resultant_length, axis=-1) + mean_resultant_length = mrl[()] if mrl.ndim == 0 else mrl + return DirectionalStats(mean_direction, mean_resultant_length) + + +@xp_capabilities(skip_backends=[('dask.array', "no take_along_axis")]) +def false_discovery_control(ps, *, axis=0, method='bh'): + """Adjust p-values to control the false discovery rate. + + The false discovery rate (FDR) is the expected proportion of rejected null + hypotheses that are actually true. + If the null hypothesis is rejected when the *adjusted* p-value falls below + a specified level, the false discovery rate is controlled at that level. + + Parameters + ---------- + ps : 1D array_like + The p-values to adjust. Elements must be real numbers between 0 and 1. + axis : int + The axis along which to perform the adjustment. The adjustment is + performed independently along each axis-slice. If `axis` is None, `ps` + is raveled before performing the adjustment. + method : {'bh', 'by'} + The false discovery rate control procedure to apply: ``'bh'`` is for + Benjamini-Hochberg [1]_ (Eq. 1), ``'by'`` is for Benjaminini-Yekutieli + [2]_ (Theorem 1.3). The latter is more conservative, but it is + guaranteed to control the FDR even when the p-values are not from + independent tests. + + Returns + ------- + ps_adusted : array_like + The adjusted p-values. If the null hypothesis is rejected where these + fall below a specified level, the false discovery rate is controlled + at that level. + + See Also + -------- + combine_pvalues + statsmodels.stats.multitest.multipletests + + Notes + ----- + In multiple hypothesis testing, false discovery control procedures tend to + offer higher power than familywise error rate control procedures (e.g. + Bonferroni correction [1]_). + + If the p-values correspond with independent tests (or tests with + "positive regression dependencies" [2]_), rejecting null hypotheses + corresponding with Benjamini-Hochberg-adjusted p-values below :math:`q` + controls the false discovery rate at a level less than or equal to + :math:`q m_0 / m`, where :math:`m_0` is the number of true null hypotheses + and :math:`m` is the total number of null hypotheses tested. The same is + true even for dependent tests when the p-values are adjusted accorded to + the more conservative Benjaminini-Yekutieli procedure. + + The adjusted p-values produced by this function are comparable to those + produced by the R function ``p.adjust`` and the statsmodels function + `statsmodels.stats.multitest.multipletests`. Please consider the latter + for more advanced methods of multiple comparison correction. + + References + ---------- + .. [1] Benjamini, Yoav, and Yosef Hochberg. "Controlling the false + discovery rate: a practical and powerful approach to multiple + testing." Journal of the Royal statistical society: series B + (Methodological) 57.1 (1995): 289-300. + + .. [2] Benjamini, Yoav, and Daniel Yekutieli. "The control of the false + discovery rate in multiple testing under dependency." Annals of + statistics (2001): 1165-1188. + + .. [3] TileStats. FDR - Benjamini-Hochberg explained - Youtube. + https://www.youtube.com/watch?v=rZKa4tW2NKs. + + .. [4] Neuhaus, Karl-Ludwig, et al. "Improved thrombolysis in acute + myocardial infarction with front-loaded administration of alteplase: + results of the rt-PA-APSAC patency study (TAPS)." Journal of the + American College of Cardiology 19.5 (1992): 885-891. + + Examples + -------- + We follow the example from [1]_. + + Thrombolysis with recombinant tissue-type plasminogen activator (rt-PA) + and anisoylated plasminogen streptokinase activator (APSAC) in + myocardial infarction has been proved to reduce mortality. [4]_ + investigated the effects of a new front-loaded administration of rt-PA + versus those obtained with a standard regimen of APSAC, in a randomized + multicentre trial in 421 patients with acute myocardial infarction. + + There were four families of hypotheses tested in the study, the last of + which was "cardiac and other events after the start of thrombolitic + treatment". FDR control may be desired in this family of hypotheses + because it would not be appropriate to conclude that the front-loaded + treatment is better if it is merely equivalent to the previous treatment. + + The p-values corresponding with the 15 hypotheses in this family were + + >>> ps = [0.0001, 0.0004, 0.0019, 0.0095, 0.0201, 0.0278, 0.0298, 0.0344, + ... 0.0459, 0.3240, 0.4262, 0.5719, 0.6528, 0.7590, 1.000] + + If the chosen significance level is 0.05, we may be tempted to reject the + null hypotheses for the tests corresponding with the first nine p-values, + as the first nine p-values fall below the chosen significance level. + However, this would ignore the problem of "multiplicity": if we fail to + correct for the fact that multiple comparisons are being performed, we + are more likely to incorrectly reject true null hypotheses. + + One approach to the multiplicity problem is to control the family-wise + error rate (FWER), that is, the rate at which the null hypothesis is + rejected when it is actually true. A common procedure of this kind is the + Bonferroni correction [1]_. We begin by multiplying the p-values by the + number of hypotheses tested. + + >>> import numpy as np + >>> np.array(ps) * len(ps) + array([1.5000e-03, 6.0000e-03, 2.8500e-02, 1.4250e-01, 3.0150e-01, + 4.1700e-01, 4.4700e-01, 5.1600e-01, 6.8850e-01, 4.8600e+00, + 6.3930e+00, 8.5785e+00, 9.7920e+00, 1.1385e+01, 1.5000e+01]) + + To control the FWER at 5%, we reject only the hypotheses corresponding + with adjusted p-values less than 0.05. In this case, only the hypotheses + corresponding with the first three p-values can be rejected. According to + [1]_, these three hypotheses concerned "allergic reaction" and "two + different aspects of bleeding." + + An alternative approach is to control the false discovery rate: the + expected fraction of rejected null hypotheses that are actually true. The + advantage of this approach is that it typically affords greater power: an + increased rate of rejecting the null hypothesis when it is indeed false. To + control the false discovery rate at 5%, we apply the Benjamini-Hochberg + p-value adjustment. + + >>> from scipy import stats + >>> stats.false_discovery_control(ps) + array([0.0015 , 0.003 , 0.0095 , 0.035625 , 0.0603 , + 0.06385714, 0.06385714, 0.0645 , 0.0765 , 0.486 , + 0.58118182, 0.714875 , 0.75323077, 0.81321429, 1. ]) + + Now, the first *four* adjusted p-values fall below 0.05, so we would reject + the null hypotheses corresponding with these *four* p-values. Rejection + of the fourth null hypothesis was particularly important to the original + study as it led to the conclusion that the new treatment had a + "substantially lower in-hospital mortality rate." + + For simplicity of exposition, the p-values in the example above were given in + sorted order, but this is not required; `false_discovery_control` returns + adjusted p-values in order corresponding with the input `ps`. + + >>> stats.false_discovery_control([0.5, 0.6, 0.1, 0.001]) + array([0.6 , 0.6 , 0.2 , 0.004]) + + """ + xp = array_namespace(ps) + + # Input Validation and Special Cases + ps = xp.asarray(ps) + + if not xp.isdtype(ps.dtype, ("integral", "real floating")): + raise ValueError("`ps` must contain only real numbers.") + + if is_lazy_array(ps): + ps = xp.where((ps < 0.) | (ps > 1.), xp.nan, ps) + else: + if not xp.all(ps == xp.clip(ps, 0., 1.)): + raise ValueError("All values in `ps` must lie between 0. and 1.") + + methods = {'bh', 'by'} + if method.lower() not in methods: + raise ValueError(f"Unrecognized `method` '{method}'." + f"Method must be one of {methods}.") + method = method.lower() + + if axis is None: + axis = 0 + ps = xp_ravel(ps) + + axis = np.asarray(axis)[()] # use of NumPy for input validation is OK + if not np.issubdtype(axis.dtype, np.integer) or axis.size != 1: + raise ValueError("`axis` must be an integer or `None`") + axis = int(axis) + + if xp_size(ps) <= 1 or ps.shape[axis] <= 1: + return ps[()] if ps.ndim == 0 else ps + + ps = xp.moveaxis(ps, axis, -1) + m = ps.shape[-1] + + # Main Algorithm + # Equivalent to the ideas of [1] and [2], except that this adjusts the + # p-values as described in [3]. The results are similar to those produced + # by R's p.adjust. + + # "Let [ps] be the ordered observed p-values..." + order = xp.argsort(ps, axis=-1) + ps = xp.take_along_axis(ps, order, axis=-1) # this copies ps + + # Equation 1 of [1] rearranged to reject when p is less than specified q + i = xp.arange(1, m+1, dtype=ps.dtype, device=xp_device(ps)) + # ps *= m / i + ps = xpx.at(ps)[...].multiply(m / i) + + # Theorem 1.3 of [2] + if method == 'by': + # ps *= np.sum(1 / i) + ps = xpx.at(ps)[...].multiply(xp.sum(1 / i)) + + # accounts for rejecting all null hypotheses i for i < k, where k is + # defined in Eq. 1 of either [1] or [2]. See [3]. Starting with the index j + # of the second to last element, we replace element j with element j+1 if + # the latter is smaller. + if is_numpy(xp): + np.minimum.accumulate(ps[..., ::-1], out=ps[..., ::-1], axis=-1) + else: + n = ps.shape[-1] + for j in range(n-2, -1, -1): + # ps[..., j] = xp.minimum(ps[..., j], ps[..., j+1]) + ps = xpx.at(ps)[..., j].set(xp.minimum(ps[..., j], ps[..., j+1])) + + # Restore original order of axes and data + ps = _reorder_along_axis(ps, order, axis=-1, xp=xp) + ps = xp.moveaxis(ps, -1, axis) + + return xp.clip(ps, 0., 1.) + + +def _reorder_along_axis(x, i, *, axis, xp): + if is_jax(xp): + return xp.put_along_axis(x, i, values=x, axis=axis, inplace=False) + if hasattr(xp, 'put_along_axis'): + xp.put_along_axis(x, i, values=x.copy(), axis=axis) + return x + else: + return xp.take_along_axis(x, xp.argsort(i, axis=-1), axis=-1) diff --git a/v12x/pylib/scipy/stats/_mstats_basic.py b/v12x/pylib/scipy/stats/_mstats_basic.py new file mode 100644 index 0000000000000000000000000000000000000000..96517b32ae2a8865403d2a02ac9a64fbd3ea5805 --- /dev/null +++ b/v12x/pylib/scipy/stats/_mstats_basic.py @@ -0,0 +1,3833 @@ +""" +An extension of scipy.stats._stats_py to support masked arrays + +""" +# Original author (2007): Pierre GF Gerard-Marchant + + +__all__ = ['argstoarray', + 'count_tied_groups', + 'describe', + 'f_oneway', 'find_repeats','friedmanchisquare', + 'kendalltau','kendalltau_seasonal','kruskal','kruskalwallis', + 'ks_twosamp', 'ks_2samp', 'kurtosis', 'kurtosistest', + 'ks_1samp', 'kstest', + 'linregress', + 'mannwhitneyu', 'meppf','mode','moment','mquantiles','msign', + 'normaltest', + 'obrientransform', + 'pearsonr','plotting_positions','pointbiserialr', + 'rankdata', + 'scoreatpercentile','sem', + 'sen_seasonal_slopes','skew','skewtest','spearmanr', + 'siegelslopes', 'theilslopes', + 'tmax','tmean','tmin','trim','trimboth', + 'trimtail','trima','trimr','trimmed_mean','trimmed_std', + 'trimmed_stde','trimmed_var','tsem','ttest_1samp','ttest_onesamp', + 'ttest_ind','ttest_rel','tvar', + 'variation', + 'winsorize', + 'brunnermunzel', + ] + +import numpy as np +from numpy import ndarray +import numpy.ma as ma +from numpy.ma import masked, nomask +import math + +import itertools +import warnings +from collections import namedtuple + +from . import distributions +from scipy._lib._util import _rename_parameter, _contains_nan, _dedent_for_py313 +from scipy._lib._bunch import _make_tuple_bunch +import scipy.special as special +import scipy.stats._stats_py +import scipy.stats._stats_py as _stats_py +from scipy.stats._stats_pythran import siegelslopes as siegelslopes_pythran +from ._stats_mstats_common import _find_repeats, TheilslopesResult, SiegelslopesResult + + +def _chk_asarray(a, axis): + # Always returns a masked array, raveled for axis=None + a = ma.asanyarray(a) + if axis is None: + a = ma.ravel(a) + outaxis = 0 + else: + outaxis = axis + return a, outaxis + + +def _chk2_asarray(a, b, axis): + a = ma.asanyarray(a) + b = ma.asanyarray(b) + if axis is None: + a = ma.ravel(a) + b = ma.ravel(b) + outaxis = 0 + else: + outaxis = axis + return a, b, outaxis + + +def _chk_size(a, b): + a = ma.asanyarray(a) + b = ma.asanyarray(b) + (na, nb) = (a.size, b.size) + if na != nb: + raise ValueError("The size of the input array should match!" + f" ({na} <> {nb})") + return (a, b, na) + + +def _ttest_finish(df, t, alternative): + """Common code between all 3 t-test functions.""" + # We use ``stdtr`` directly here to preserve masked arrays + + if alternative == 'less': + pval = special._ufuncs.stdtr(df, t) + elif alternative == 'greater': + pval = special._ufuncs.stdtr(df, -t) + elif alternative == 'two-sided': + pval = special._ufuncs.stdtr(df, -np.abs(t))*2 + else: + raise ValueError("alternative must be " + "'less', 'greater' or 'two-sided'") + + if t.ndim == 0: + t = t[()] + if pval.ndim == 0: + pval = pval[()] + + return t, pval + + +def argstoarray(*args): + """ + Constructs a 2D array from a group of sequences. + + Sequences are filled with missing values to match the length of the longest + sequence. + + Parameters + ---------- + *args : sequences + Group of sequences. + + Returns + ------- + argstoarray : MaskedArray + A ( `m` x `n` ) masked array, where `m` is the number of arguments and + `n` the length of the longest argument. + + Notes + ----- + `numpy.ma.vstack` has identical behavior, but is called with a sequence + of sequences. + + Examples + -------- + A 2D masked array constructed from a group of sequences is returned. + + >>> from scipy.stats.mstats import argstoarray + >>> argstoarray([1, 2, 3], [4, 5, 6]) + masked_array( + data=[[1.0, 2.0, 3.0], + [4.0, 5.0, 6.0]], + mask=[[False, False, False], + [False, False, False]], + fill_value=1e+20) + + The returned masked array filled with missing values when the lengths of + sequences are different. + + >>> argstoarray([1, 3], [4, 5, 6]) + masked_array( + data=[[1.0, 3.0, --], + [4.0, 5.0, 6.0]], + mask=[[False, False, True], + [False, False, False]], + fill_value=1e+20) + + """ + if len(args) == 1 and not isinstance(args[0], ndarray): + output = ma.asarray(args[0]) + if output.ndim != 2: + raise ValueError("The input should be 2D") + else: + n = len(args) + m = max([len(k) for k in args]) + output = ma.array(np.empty((n,m), dtype=float), mask=True) + for (k,v) in enumerate(args): + output[k,:len(v)] = v + + output[np.logical_not(np.isfinite(output._data))] = masked + return output + + +def find_repeats(arr): + """Find repeats in arr and return a tuple (repeats, repeat_count). + + The input is cast to float64. Masked values are discarded. + + Parameters + ---------- + arr : sequence + Input array. The array is flattened if it is not 1D. + + Returns + ------- + repeats : ndarray + Array of repeated values. + counts : ndarray + Array of counts. + + Examples + -------- + >>> from scipy.stats import mstats + >>> mstats.find_repeats([2, 1, 2, 3, 2, 2, 5]) + (array([2.]), array([4])) + + In the above example, 2 repeats 4 times. + + >>> mstats.find_repeats([[10, 20, 1, 2], [5, 5, 4, 4]]) + (array([4., 5.]), array([2, 2])) + + In the above example, both 4 and 5 repeat 2 times. + + """ + # Make sure we get a copy. ma.compressed promises a "new array", but can + # actually return a reference. + compr = np.asarray(ma.compressed(arr), dtype=np.float64) + need_copy = np.may_share_memory(compr, arr) + if need_copy: + compr = compr.copy() + return _find_repeats(compr) + + +def count_tied_groups(x, use_missing=False): + """ + Counts the number of tied values. + + Parameters + ---------- + x : sequence + Sequence of data on which to counts the ties + use_missing : bool, optional + Whether to consider missing values as tied. + + Returns + ------- + count_tied_groups : dict + Returns a dictionary (nb of ties: nb of groups). + + Examples + -------- + >>> from scipy.stats import mstats + >>> import numpy as np + >>> z = [0, 0, 0, 2, 2, 2, 3, 3, 4, 5, 6] + >>> mstats.count_tied_groups(z) + {2: 1, 3: 2} + + In the above example, the ties were 0 (3x), 2 (3x) and 3 (2x). + + >>> z = np.ma.array([0, 0, 1, 2, 2, 2, 3, 3, 4, 5, 6]) + >>> mstats.count_tied_groups(z) + {2: 2, 3: 1} + >>> z[[1,-1]] = np.ma.masked + >>> mstats.count_tied_groups(z, use_missing=True) + {2: 2, 3: 1} + + """ + nmasked = ma.getmask(x).sum() + # We need the copy as find_repeats will overwrite the initial data + data = ma.compressed(x).copy() + (ties, counts) = find_repeats(data) + nties = {} + if len(ties): + nties = dict(zip(np.unique(counts), itertools.repeat(1))) + nties.update(dict(zip(*find_repeats(counts)))) + + if nmasked and use_missing: + try: + nties[nmasked] += 1 + except KeyError: + nties[nmasked] = 1 + + return nties + + +def rankdata(data, axis=None, use_missing=False): + """Returns the rank (also known as order statistics) of each data point + along the given axis. + + If some values are tied, their rank is averaged. + If some values are masked, their rank is set to 0 if use_missing is False, + or set to the average rank of the unmasked values if use_missing is True. + + Parameters + ---------- + data : sequence + Input data. The data is transformed to a masked array + axis : {None,int}, optional + Axis along which to perform the ranking. + If None, the array is first flattened. An exception is raised if + the axis is specified for arrays with a dimension larger than 2 + use_missing : bool, optional + Whether the masked values have a rank of 0 (False) or equal to the + average rank of the unmasked values (True). + + """ # numpydoc ignore=RT01 + def _rank1d(data, use_missing=False): + n = data.count() + rk = np.empty(data.size, dtype=float) + idx = data.argsort() + rk[idx[:n]] = np.arange(1,n+1) + + if use_missing: + rk[idx[n:]] = (n+1)/2. + else: + rk[idx[n:]] = 0 + + repeats = find_repeats(data.copy()) + for r in repeats[0]: + condition = (data == r).filled(False) + rk[condition] = rk[condition].mean() + return rk + + data = ma.array(data, copy=False) + if axis is None: + if data.ndim > 1: + return _rank1d(data.ravel(), use_missing).reshape(data.shape) + else: + return _rank1d(data, use_missing) + else: + return ma.apply_along_axis(_rank1d,axis,data,use_missing).view(ndarray) + + +ModeResult = namedtuple('ModeResult', ('mode', 'count')) + + +def mode(a, axis=0): + """ + Returns an array of the modal (most common) value in the passed array. + + Parameters + ---------- + a : array_like + n-dimensional array of which to find mode(s). + axis : int or None, optional + Axis along which to operate. Default is 0. If None, compute over + the whole array `a`. + + Returns + ------- + mode : ndarray + Array of modal values. + count : ndarray + Array of counts for each mode. + + Notes + ----- + For more details, see `scipy.stats.mode`. + + Examples + -------- + >>> import numpy as np + >>> from scipy import stats + >>> from scipy.stats import mstats + >>> m_arr = np.ma.array([1, 1, 0, 0, 0, 0], mask=[0, 0, 1, 1, 1, 0]) + >>> mstats.mode(m_arr) # note that most zeros are masked + ModeResult(mode=array([1.]), count=array([2.])) + + """ + return _mode(a, axis=axis, keepdims=True) + + +def _mode(a, axis=0, keepdims=True): + # Don't want to expose `keepdims` from the public `mstats.mode` + a, axis = _chk_asarray(a, axis) + + def _mode1D(a): + (rep,cnt) = find_repeats(a) + if not cnt.ndim: + return (0, 0) + elif cnt.size: + return (rep[cnt.argmax()], cnt.max()) + else: + return (a.min(), 1) + + if axis is None: + output = _mode1D(ma.ravel(a)) + output = (ma.array(output[0]), ma.array(output[1])) + else: + output = ma.apply_along_axis(_mode1D, axis, a) + if keepdims is None or keepdims: + newshape = list(a.shape) + newshape[axis] = 1 + slices = [slice(None)] * output.ndim + slices[axis] = 0 + modes = output[tuple(slices)].reshape(newshape) + slices[axis] = 1 + counts = output[tuple(slices)].reshape(newshape) + output = (modes, counts) + else: + output = np.moveaxis(output, axis, 0) + + return ModeResult(*output) + + +def _betai(a, b, x): + x = np.asanyarray(x) + x = ma.where(x < 1.0, x, 1.0) # if x > 1 then return 1.0 + return special.betainc(a, b, x) + + +def msign(x): + """Returns the sign of x, or 0 if x is masked.""" # numpydoc ignore=RT01 + return ma.filled(np.sign(x), 0) + + +def pearsonr(x, y): + r""" + Pearson correlation coefficient and p-value for testing non-correlation. + + The Pearson correlation coefficient [1]_ measures the linear relationship + between two datasets. The calculation of the p-value relies on the + assumption that each dataset is normally distributed. (See Kowalski [3]_ + for a discussion of the effects of non-normality of the input on the + distribution of the correlation coefficient.) Like other correlation + coefficients, this one varies between -1 and +1 with 0 implying no + correlation. Correlations of -1 or +1 imply an exact linear relationship. + + Parameters + ---------- + x : (N,) array_like + Input array. + y : (N,) array_like + Input array. + + Returns + ------- + r : float + Pearson's correlation coefficient. + p-value : float + Two-tailed p-value. + + Warns + ----- + `~scipy.stats.ConstantInputWarning` + Raised if an input is a constant array. The correlation coefficient + is not defined in this case, so ``np.nan`` is returned. + + `~scipy.stats.NearConstantInputWarning` + Raised if an input is "nearly" constant. The array ``x`` is considered + nearly constant if ``norm(x - mean(x)) < 1e-13 * abs(mean(x))``. + Numerical errors in the calculation ``x - mean(x)`` in this case might + result in an inaccurate calculation of r. + + See Also + -------- + spearmanr : Spearman rank-order correlation coefficient. + kendalltau : Kendall's tau, a correlation measure for ordinal data. + + Notes + ----- + The correlation coefficient is calculated as follows: + + .. math:: + + r = \frac{\sum (x - m_x) (y - m_y)} + {\sqrt{\sum (x - m_x)^2 \sum (y - m_y)^2}} + + where :math:`m_x` is the mean of the vector x and :math:`m_y` is + the mean of the vector y. + + Under the assumption that x and y are drawn from + independent normal distributions (so the population correlation coefficient + is 0), the probability density function of the sample correlation + coefficient r is ([1]_, [2]_): + + .. math:: + + f(r) = \frac{{(1-r^2)}^{n/2-2}}{\mathrm{B}(\frac{1}{2},\frac{n}{2}-1)} + + where n is the number of samples, and B is the beta function. This + is sometimes referred to as the exact distribution of r. This is + the distribution that is used in `pearsonr` to compute the p-value. + The distribution is a beta distribution on the interval [-1, 1], + with equal shape parameters a = b = n/2 - 1. In terms of SciPy's + implementation of the beta distribution, the distribution of r is:: + + dist = scipy.stats.beta(n/2 - 1, n/2 - 1, loc=-1, scale=2) + + The p-value returned by `pearsonr` is a two-sided p-value. The p-value + roughly indicates the probability of an uncorrelated system + producing datasets that have a Pearson correlation at least as extreme + as the one computed from these datasets. More precisely, for a + given sample with correlation coefficient r, the p-value is + the probability that abs(r') of a random sample x' and y' drawn from + the population with zero correlation would be greater than or equal + to abs(r). In terms of the object ``dist`` shown above, the p-value + for a given r and length n can be computed as:: + + p = 2*dist.cdf(-abs(r)) + + When n is 2, the above continuous distribution is not well-defined. + One can interpret the limit of the beta distribution as the shape + parameters a and b approach a = b = 0 as a discrete distribution with + equal probability masses at r = 1 and r = -1. More directly, one + can observe that, given the data x = [x1, x2] and y = [y1, y2], and + assuming x1 != x2 and y1 != y2, the only possible values for r are 1 + and -1. Because abs(r') for any sample x' and y' with length 2 will + be 1, the two-sided p-value for a sample of length 2 is always 1. + + References + ---------- + .. [1] "Pearson correlation coefficient", Wikipedia, + https://en.wikipedia.org/wiki/Pearson_correlation_coefficient + .. [2] Student, "Probable error of a correlation coefficient", + Biometrika, Volume 6, Issue 2-3, 1 September 1908, pp. 302-310. + .. [3] C. J. Kowalski, "On the Effects of Non-Normality on the Distribution + of the Sample Product-Moment Correlation Coefficient" + Journal of the Royal Statistical Society. Series C (Applied + Statistics), Vol. 21, No. 1 (1972), pp. 1-12. + + Examples + -------- + >>> import numpy as np + >>> from scipy import stats + >>> from scipy.stats import mstats + >>> mstats.pearsonr([1, 2, 3, 4, 5], [10, 9, 2.5, 6, 4]) + (-0.7426106572325057, 0.1505558088534455) + + There is a linear dependence between x and y if y = a + b*x + e, where + a,b are constants and e is a random error term, assumed to be independent + of x. For simplicity, assume that x is standard normal, a=0, b=1 and let + e follow a normal distribution with mean zero and standard deviation s>0. + + >>> s = 0.5 + >>> x = stats.norm.rvs(size=500) + >>> e = stats.norm.rvs(scale=s, size=500) + >>> y = x + e + >>> mstats.pearsonr(x, y) + (0.9029601878969703, 8.428978827629898e-185) # may vary + + This should be close to the exact value given by + + >>> 1/np.sqrt(1 + s**2) + 0.8944271909999159 + + For s=0.5, we observe a high level of correlation. In general, a large + variance of the noise reduces the correlation, while the correlation + approaches one as the variance of the error goes to zero. + + It is important to keep in mind that no correlation does not imply + independence unless (x, y) is jointly normal. Correlation can even be zero + when there is a very simple dependence structure: if X follows a + standard normal distribution, let y = abs(x). Note that the correlation + between x and y is zero. Indeed, since the expectation of x is zero, + cov(x, y) = E[x*y]. By definition, this equals E[x*abs(x)] which is zero + by symmetry. The following lines of code illustrate this observation: + + >>> y = np.abs(x) + >>> mstats.pearsonr(x, y) + (-0.016172891856853524, 0.7182823678751942) # may vary + + A non-zero correlation coefficient can be misleading. For example, if X has + a standard normal distribution, define y = x if x < 0 and y = 0 otherwise. + A simple calculation shows that corr(x, y) = sqrt(2/Pi) = 0.797..., + implying a high level of correlation: + + >>> y = np.where(x < 0, x, 0) + >>> mstats.pearsonr(x, y) + (0.8537091583771509, 3.183461621422181e-143) # may vary + + This is unintuitive since there is no dependence of x and y if x is larger + than zero which happens in about half of the cases if we sample x and y. + """ + (x, y, n) = _chk_size(x, y) + (x, y) = (x.ravel(), y.ravel()) + # Get the common mask and the total nb of unmasked elements + m = ma.mask_or(ma.getmask(x), ma.getmask(y)) + n -= m.sum() + df = n-2 + if df < 0: + return (masked, masked) + + return scipy.stats._stats_py.pearsonr( + ma.masked_array(x, mask=m).compressed(), + ma.masked_array(y, mask=m).compressed()) + + +def spearmanr(x, y=None, use_ties=True, axis=None, nan_policy='propagate', + alternative='two-sided'): + """ + Calculates a Spearman rank-order correlation coefficient and the p-value + to test for non-correlation. + + The Spearman correlation is a nonparametric measure of the linear + relationship between two datasets. Unlike the Pearson correlation, the + Spearman correlation does not assume that both datasets are normally + distributed. Like other correlation coefficients, this one varies + between -1 and +1 with 0 implying no correlation. Correlations of -1 or + +1 imply a monotonic relationship. Positive correlations imply that + as `x` increases, so does `y`. Negative correlations imply that as `x` + increases, `y` decreases. + + Missing values are discarded pair-wise: if a value is missing in `x`, the + corresponding value in `y` is masked. + + The p-value roughly indicates the probability of an uncorrelated system + producing datasets that have a Spearman correlation at least as extreme + as the one computed from these datasets. The p-values are not entirely + reliable but are probably reasonable for datasets larger than 500 or so. + + Parameters + ---------- + x, y : 1D or 2D array_like, y is optional + One or two 1-D or 2-D arrays containing multiple variables and + observations. When these are 1-D, each represents a vector of + observations of a single variable. For the behavior in the 2-D case, + see under ``axis``, below. + use_ties : bool, optional + DO NOT USE. Does not do anything, keyword is only left in place for + backwards compatibility reasons. + axis : int or None, optional + If axis=0 (default), then each column represents a variable, with + observations in the rows. If axis=1, the relationship is transposed: + each row represents a variable, while the columns contain observations. + If axis=None, then both arrays will be raveled. + nan_policy : {'propagate', 'raise', 'omit'}, optional + Defines how to handle when input contains nan. 'propagate' returns nan, + 'raise' throws an error, 'omit' performs the calculations ignoring nan + values. Default is 'propagate'. + alternative : {'two-sided', 'less', 'greater'}, optional + Defines the alternative hypothesis. Default is 'two-sided'. + The following options are available: + + * 'two-sided': the correlation is nonzero + * 'less': the correlation is negative (less than zero) + * 'greater': the correlation is positive (greater than zero) + + .. versionadded:: 1.7.0 + + Returns + ------- + res : SignificanceResult + An object containing attributes: + + statistic : float or ndarray (2-D square) + Spearman correlation matrix or correlation coefficient (if only 2 + variables are given as parameters). Correlation matrix is square + with length equal to total number of variables (columns or rows) in + ``a`` and ``b`` combined. + pvalue : float + The p-value for a hypothesis test whose null hypothesis + is that two sets of data are linearly uncorrelated. See + `alternative` above for alternative hypotheses. `pvalue` has the + same shape as `statistic`. + + References + ---------- + [CRCProbStat2000] section 14.7 + + """ + if not use_ties: + raise ValueError("`use_ties=False` is not supported in SciPy >= 1.2.0") + + # Always returns a masked array, raveled if axis=None + x, axisout = _chk_asarray(x, axis) + if y is not None: + # Deal only with 2-D `x` case. + y, _ = _chk_asarray(y, axis) + if axisout == 0: + x = ma.column_stack((x, y)) + else: + x = ma.vstack((x, y)) + + if axisout == 1: + # To simplify the code that follow (always use `n_obs, n_vars` shape) + x = x.T + + if nan_policy == 'omit': + x = ma.masked_invalid(x) + + def _spearmanr_2cols(x): + # Mask the same observations for all variables, and then drop those + # observations (can't leave them masked, rankdata is weird). + x = ma.mask_rowcols(x, axis=0) + x = x[~x.mask.any(axis=1), :] + + # If either column is entirely NaN or Inf + if not np.any(x.data): + res = scipy.stats._stats_py.SignificanceResult(np.nan, np.nan) + res.correlation = np.nan + return res + + m = ma.getmask(x) + n_obs = x.shape[0] + dof = n_obs - 2 - int(m.sum(axis=0)[0]) + if dof < 0: + raise ValueError("The input must have at least 3 entries!") + + # Gets the ranks and rank differences + x_ranked = rankdata(x, axis=0) + rs = ma.corrcoef(x_ranked, rowvar=False).data + + # rs can have elements equal to 1, so avoid zero division warnings + with np.errstate(divide='ignore'): + # clip the small negative values possibly caused by rounding + # errors before taking the square root + t = rs * np.sqrt((dof / ((rs+1.0) * (1.0-rs))).clip(0)) + + t, prob = _ttest_finish(dof, t, alternative) + + # For backwards compatibility, return scalars when comparing 2 columns + if rs.shape == (2, 2): + res = scipy.stats._stats_py.SignificanceResult(rs[1, 0], + prob[1, 0]) + res.correlation = rs[1, 0] + return res + else: + res = scipy.stats._stats_py.SignificanceResult(rs, prob) + res.correlation = rs + return res + + # Need to do this per pair of variables, otherwise the dropped observations + # in a third column mess up the result for a pair. + n_vars = x.shape[1] + if n_vars == 2: + return _spearmanr_2cols(x) + else: + rs = np.ones((n_vars, n_vars), dtype=float) + prob = np.zeros((n_vars, n_vars), dtype=float) + for var1 in range(n_vars - 1): + for var2 in range(var1+1, n_vars): + result = _spearmanr_2cols(x[:, [var1, var2]]) + rs[var1, var2] = result.correlation + rs[var2, var1] = result.correlation + prob[var1, var2] = result.pvalue + prob[var2, var1] = result.pvalue + + res = scipy.stats._stats_py.SignificanceResult(rs, prob) + res.correlation = rs + return res + + +def _kendall_p_exact(n, c, alternative='two-sided'): + + # Use the fact that distribution is symmetric: always calculate a CDF in + # the left tail. + # This will be the one-sided p-value if `c` is on the side of + # the null distribution predicted by the alternative hypothesis. + # The two-sided p-value will be twice this value. + # If `c` is on the other side of the null distribution, we'll need to + # take the complement and add back the probability mass at `c`. + in_right_tail = (c >= (n*(n-1))//2 - c) + alternative_greater = (alternative == 'greater') + c = int(min(c, (n*(n-1))//2 - c)) + + # Exact p-value, see Maurice G. Kendall, "Rank Correlation Methods" + # (4th Edition), Charles Griffin & Co., 1970. + if n <= 0: + raise ValueError(f'n ({n}) must be positive') + elif c < 0 or 4*c > n*(n-1): + raise ValueError(f'c ({c}) must satisfy 0 <= 4c <= n(n-1) = {n*(n-1)}.') + elif n == 1: + prob = 1.0 + p_mass_at_c = 1 + elif n == 2: + prob = 1.0 + p_mass_at_c = 0.5 + elif c == 0: + prob = 2.0/math.factorial(n) if n < 171 else 0.0 + p_mass_at_c = prob/2 + elif c == 1: + prob = 2.0/math.factorial(n-1) if n < 172 else 0.0 + p_mass_at_c = (n-1)/math.factorial(n) + elif 4*c == n*(n-1) and alternative == 'two-sided': + # I'm sure there's a simple formula for p_mass_at_c in this + # case, but I don't know it. Use generic formula for one-sided p-value. + prob = 1.0 + elif n < 171: + new = np.zeros(c+1) + new[0:2] = 1.0 + for j in range(3,n+1): + new = np.cumsum(new) + if j <= c: + new[j:] -= new[:c+1-j] + prob = 2.0*np.sum(new)/math.factorial(n) + p_mass_at_c = new[-1]/math.factorial(n) + else: + new = np.zeros(c+1) + new[0:2] = 1.0 + for j in range(3, n+1): + new = np.cumsum(new)/j + if j <= c: + new[j:] -= new[:c+1-j] + prob = np.sum(new) + p_mass_at_c = new[-1]/2 + + if alternative != 'two-sided': + # if the alternative hypothesis and alternative agree, + # one-sided p-value is half the two-sided p-value + if in_right_tail == alternative_greater: + prob /= 2 + else: + prob = 1 - prob/2 + p_mass_at_c + + prob = np.clip(prob, 0, 1) + + return prob + + +def kendalltau(x, y, use_ties=True, use_missing=False, method='auto', + alternative='two-sided'): + """ + Computes Kendall's rank correlation tau on two variables *x* and *y*. + + Parameters + ---------- + x : sequence + First data list (for example, time). + y : sequence + Second data list. + use_ties : {True, False}, optional + Whether ties correction should be performed. + use_missing : {False, True}, optional + Whether missing data should be allocated a rank of 0 (False) or the + average rank (True) + method : {'auto', 'asymptotic', 'exact'}, optional + Defines which method is used to calculate the p-value [1]_. + 'asymptotic' uses a normal approximation valid for large samples. + 'exact' computes the exact p-value, but can only be used if no ties + are present. As the sample size increases, the 'exact' computation + time may grow and the result may lose some precision. + 'auto' is the default and selects the appropriate + method based on a trade-off between speed and accuracy. + alternative : {'two-sided', 'less', 'greater'}, optional + Defines the alternative hypothesis. Default is 'two-sided'. + The following options are available: + + * 'two-sided': the rank correlation is nonzero + * 'less': the rank correlation is negative (less than zero) + * 'greater': the rank correlation is positive (greater than zero) + + Returns + ------- + res : SignificanceResult + An object containing attributes: + + statistic : float + The tau statistic. + pvalue : float + The p-value for a hypothesis test whose null hypothesis is + an absence of association, tau = 0. + + References + ---------- + .. [1] Maurice G. Kendall, "Rank Correlation Methods" (4th Edition), + Charles Griffin & Co., 1970. + + """ + (x, y, n) = _chk_size(x, y) + (x, y) = (x.flatten(), y.flatten()) + m = ma.mask_or(ma.getmask(x), ma.getmask(y)) + if m is not nomask: + x = ma.array(x, mask=m, copy=True) + y = ma.array(y, mask=m, copy=True) + # need int() here, otherwise numpy defaults to 32 bit + # integer on all Windows architectures, causing overflow. + # int() will keep it infinite precision. + n -= int(m.sum()) + + if n < 2: + res = scipy.stats._stats_py.SignificanceResult(np.nan, np.nan) + res.correlation = np.nan + return res + + rx = ma.masked_equal(rankdata(x, use_missing=use_missing), 0) + ry = ma.masked_equal(rankdata(y, use_missing=use_missing), 0) + idx = rx.argsort() + (rx, ry) = (rx[idx], ry[idx]) + C = np.sum([((ry[i+1:] > ry[i]) * (rx[i+1:] > rx[i])).filled(0).sum() + for i in range(len(ry)-1)], dtype=float) + D = np.sum([((ry[i+1:] < ry[i])*(rx[i+1:] > rx[i])).filled(0).sum() + for i in range(len(ry)-1)], dtype=float) + xties = count_tied_groups(x) + yties = count_tied_groups(y) + if use_ties: + corr_x = np.sum([v*k*(k-1) for (k,v) in xties.items()], dtype=float) + corr_y = np.sum([v*k*(k-1) for (k,v) in yties.items()], dtype=float) + denom = ma.sqrt((n*(n-1)-corr_x)/2. * (n*(n-1)-corr_y)/2.) + else: + denom = n*(n-1)/2. + tau = (C-D) / denom + + if method == 'exact' and (xties or yties): + raise ValueError("Ties found, exact method cannot be used.") + + if method == 'auto': + if (not xties and not yties) and (n <= 33 or min(C, n*(n-1)/2.0-C) <= 1): + method = 'exact' + else: + method = 'asymptotic' + + if not xties and not yties and method == 'exact': + prob = _kendall_p_exact(n, C, alternative) + + elif method == 'asymptotic': + var_s = n*(n-1)*(2*n+5) + if use_ties: + var_s -= np.sum([v*k*(k-1)*(2*k+5)*1. for (k,v) in xties.items()]) + var_s -= np.sum([v*k*(k-1)*(2*k+5)*1. for (k,v) in yties.items()]) + v1 = (np.sum([v*k*(k-1) for (k, v) in xties.items()], dtype=float) * + np.sum([v*k*(k-1) for (k, v) in yties.items()], dtype=float)) + v1 /= 2.*n*(n-1) + if n > 2: + v2 = np.sum([v*k*(k-1)*(k-2) for (k,v) in xties.items()], + dtype=float) * \ + np.sum([v*k*(k-1)*(k-2) for (k,v) in yties.items()], + dtype=float) + v2 /= 9.*n*(n-1)*(n-2) + else: + v2 = 0 + else: + v1 = v2 = 0 + + var_s /= 18. + var_s += (v1 + v2) + z = (C-D)/np.sqrt(var_s) + prob = scipy.stats._stats_py._get_pvalue(z, distributions.norm, alternative) + else: + raise ValueError("Unknown method "+str(method)+" specified, please " + "use auto, exact or asymptotic.") + + res = scipy.stats._stats_py.SignificanceResult(tau[()], prob[()]) + res.correlation = tau + return res + + +def kendalltau_seasonal(x): + """ + Computes a multivariate Kendall's rank correlation tau, for seasonal data. + + Parameters + ---------- + x : 2-D ndarray + Array of seasonal data, with seasons in columns. + + """ # numpydoc ignore=RT01 + x = ma.array(x, subok=True, copy=False, ndmin=2) + (n,m) = x.shape + n_p = x.count(0) + + S_szn = sum(msign(x[i:]-x[i]).sum(0) for i in range(n)) + S_tot = S_szn.sum() + + n_tot = x.count() + ties = count_tied_groups(x.compressed()) + corr_ties = sum(v*k*(k-1) for (k,v) in ties.items()) + denom_tot = ma.sqrt(1.*n_tot*(n_tot-1)*(n_tot*(n_tot-1)-corr_ties))/2. + + R = rankdata(x, axis=0, use_missing=True) + K = ma.empty((m,m), dtype=int) + covmat = ma.empty((m,m), dtype=float) + denom_szn = ma.empty(m, dtype=float) + for j in range(m): + ties_j = count_tied_groups(x[:,j].compressed()) + corr_j = sum(v*k*(k-1) for (k,v) in ties_j.items()) + cmb = n_p[j]*(n_p[j]-1) + for k in range(j,m,1): + K[j,k] = sum(msign((x[i:,j]-x[i,j])*(x[i:,k]-x[i,k])).sum() + for i in range(n)) + covmat[j,k] = (K[j,k] + 4*(R[:,j]*R[:,k]).sum() - + n*(n_p[j]+1)*(n_p[k]+1))/3. + K[k,j] = K[j,k] + covmat[k,j] = covmat[j,k] + + denom_szn[j] = ma.sqrt(cmb*(cmb-corr_j)) / 2. + + var_szn = covmat.diagonal() + + z_szn = msign(S_szn) * (abs(S_szn)-1) / ma.sqrt(var_szn) + z_tot_ind = msign(S_tot) * (abs(S_tot)-1) / ma.sqrt(var_szn.sum()) + z_tot_dep = msign(S_tot) * (abs(S_tot)-1) / ma.sqrt(covmat.sum()) + + prob_szn = special.erfc(abs(z_szn.data)/np.sqrt(2)) + prob_tot_ind = special.erfc(abs(z_tot_ind)/np.sqrt(2)) + prob_tot_dep = special.erfc(abs(z_tot_dep)/np.sqrt(2)) + + chi2_tot = (z_szn*z_szn).sum() + chi2_trd = m * z_szn.mean()**2 + output = {'seasonal tau': S_szn/denom_szn, + 'global tau': S_tot/denom_tot, + 'global tau (alt)': S_tot/denom_szn.sum(), + 'seasonal p-value': prob_szn, + 'global p-value (indep)': prob_tot_ind, + 'global p-value (dep)': prob_tot_dep, + 'chi2 total': chi2_tot, + 'chi2 trend': chi2_trd, + } + return output + + +PointbiserialrResult = namedtuple('PointbiserialrResult', ('correlation', + 'pvalue')) + + +def pointbiserialr(x, y): + """Calculates a point biserial correlation coefficient and its p-value. + + Parameters + ---------- + x : array_like of bools + Input array. + y : array_like + Input array. + + Returns + ------- + correlation : float + R value + pvalue : float + 2-tailed p-value + + Notes + ----- + Missing values are considered pair-wise: if a value is missing in x, + the corresponding value in y is masked. + + For more details on `pointbiserialr`, see `scipy.stats.pointbiserialr`. + + """ + x = ma.fix_invalid(x, copy=True).astype(bool) + y = ma.fix_invalid(y, copy=True).astype(float) + # Get rid of the missing data + m = ma.mask_or(ma.getmask(x), ma.getmask(y)) + if m is not nomask: + unmask = np.logical_not(m) + x = x[unmask] + y = y[unmask] + + n = len(x) + # phat is the fraction of x values that are True + phat = x.sum() / float(n) + y0 = y[~x] # y-values where x is False + y1 = y[x] # y-values where x is True + y0m = y0.mean() + y1m = y1.mean() + + rpb = (y1m - y0m)*np.sqrt(phat * (1-phat)) / y.std() + + df = n-2 + t = rpb*ma.sqrt(df/(1.0-rpb**2)) + prob = _betai(0.5*df, 0.5, df/(df+t*t)) + + return PointbiserialrResult(rpb, prob) + + +def linregress(x, y=None): + r""" + Calculate a linear least-squares regression for two sets of measurements. + + Parameters + ---------- + x, y : array_like + Two sets of measurements. Both arrays should have the same length N. If + only `x` is given (and ``y=None``), then it must be a two-dimensional + array where one dimension has length 2. The two sets of measurements + are then found by splitting the array along the length-2 dimension. In + the case where ``y=None`` and `x` is a 2xN array, ``linregress(x)`` is + equivalent to ``linregress(x[0], x[1])``. + + Returns + ------- + result : ``LinregressResult`` instance + The return value is an object with the following attributes: + + slope : float + Slope of the regression line. + intercept : float + Intercept of the regression line. + rvalue : float + The Pearson correlation coefficient. The square of ``rvalue`` + is equal to the coefficient of determination. + pvalue : float + The p-value for a hypothesis test whose null hypothesis is + that the slope is zero, using Wald Test with t-distribution of + the test statistic. See `alternative` above for alternative + hypotheses. + stderr : float + Standard error of the estimated slope (gradient), under the + assumption of residual normality. + intercept_stderr : float + Standard error of the estimated intercept, under the assumption + of residual normality. + + See Also + -------- + scipy.optimize.curve_fit : + Use non-linear least squares to fit a function to data. + scipy.optimize.leastsq : + Minimize the sum of squares of a set of equations. + + Notes + ----- + Missing values are considered pair-wise: if a value is missing in `x`, + the corresponding value in `y` is masked. + + For compatibility with older versions of SciPy, the return value acts + like a ``namedtuple`` of length 5, with fields ``slope``, ``intercept``, + ``rvalue``, ``pvalue`` and ``stderr``, so one can continue to write:: + + slope, intercept, r, p, se = linregress(x, y) + + With that style, however, the standard error of the intercept is not + available. To have access to all the computed values, including the + standard error of the intercept, use the return value as an object + with attributes, e.g.:: + + result = linregress(x, y) + print(result.intercept, result.intercept_stderr) + + Examples + -------- + >>> import numpy as np + >>> import matplotlib.pyplot as plt + >>> from scipy import stats + >>> rng = np.random.default_rng() + + Generate some data: + + >>> x = rng.random(10) + >>> y = 1.6*x + rng.random(10) + + Perform the linear regression: + + >>> res = stats.mstats.linregress(x, y) + + Coefficient of determination (R-squared): + + >>> print(f"R-squared: {res.rvalue**2:.6f}") + R-squared: 0.717533 + + Plot the data along with the fitted line: + + >>> plt.plot(x, y, 'o', label='original data') + >>> plt.plot(x, res.intercept + res.slope*x, 'r', label='fitted line') + >>> plt.legend() + >>> plt.show() + + Calculate 95% confidence interval on slope and intercept: + + >>> # Two-sided inverse Students t-distribution + >>> # p - probability, df - degrees of freedom + >>> from scipy.stats import t + >>> tinv = lambda p, df: abs(t.ppf(p/2, df)) + + >>> ts = tinv(0.05, len(x)-2) + >>> print(f"slope (95%): {res.slope:.6f} +/- {ts*res.stderr:.6f}") + slope (95%): 1.453392 +/- 0.743465 + >>> print(f"intercept (95%): {res.intercept:.6f}" + ... f" +/- {ts*res.intercept_stderr:.6f}") + intercept (95%): 0.616950 +/- 0.544475 + + """ + if y is None: + x = ma.array(x) + if x.shape[0] == 2: + x, y = x + elif x.shape[1] == 2: + x, y = x.T + else: + raise ValueError("If only `x` is given as input, " + "it has to be of shape (2, N) or (N, 2), " + f"provided shape was {x.shape}") + else: + x = ma.array(x) + y = ma.array(y) + + x = x.flatten() + y = y.flatten() + + if np.amax(x) == np.amin(x) and len(x) > 1: + raise ValueError("Cannot calculate a linear regression " + "if all x values are identical") + + m = ma.mask_or(ma.getmask(x), ma.getmask(y), shrink=False) + if m is not nomask: + x = ma.array(x, mask=m) + y = ma.array(y, mask=m) + if np.any(~m): + result = _stats_py.linregress(x.data[~m], y.data[~m]) + else: + # All data is masked + result = _stats_py.LinregressResult(slope=None, intercept=None, + rvalue=None, pvalue=None, + stderr=None, + intercept_stderr=None) + else: + result = _stats_py.linregress(x.data, y.data) + + return result + + +def _theilslopes(y, x=None, alpha=0.95, method='separate'): + r""" + Computes the Theil-Sen estimator for a set of points (x, y). + + `theilslopes` implements a method for robust linear regression. It + computes the slope as the median of all slopes between paired values. + + Parameters + ---------- + y : array_like + Dependent variable. + x : array_like or None, optional + Independent variable. If None, use ``arange(len(y))`` instead. + alpha : float, optional + Confidence degree between 0 and 1. Default is 95% confidence. + Note that `alpha` is symmetric around 0.5, i.e. both 0.1 and 0.9 are + interpreted as "find the 90% confidence interval". + method : {'joint', 'separate'}, optional + Method to be used for computing estimate for intercept. + Following methods are supported, + + * 'joint': Uses np.median(y - slope * x) as intercept. + * 'separate': Uses np.median(y) - slope * np.median(x) + as intercept. + + The default is 'separate'. + + .. versionadded:: 1.8.0 + + Returns + ------- + result : ``TheilslopesResult`` instance + The return value is an object with the following attributes: + + slope : float + Theil slope. + intercept : float + Intercept of the Theil line. + low_slope : float + Lower bound of the confidence interval on `slope`. + high_slope : float + Upper bound of the confidence interval on `slope`. + + See Also + -------- + siegelslopes : a similar technique using repeated medians + + Notes + ----- + The implementation of `theilslopes` follows [1]_. The intercept is + not defined in [1]_, and here it is defined as ``median(y) - + slope*median(x)``, which is given in [3]_. Other definitions of + the intercept exist in the literature such as ``median(y - slope*x)`` + in [4]_. The approach to compute the intercept can be determined by the + parameter ``method``. A confidence interval for the intercept is not + given as this question is not addressed in [1]_. + + For compatibility with older versions of SciPy, the return value acts + like a ``namedtuple`` of length 4, with fields ``slope``, ``intercept``, + ``low_slope``, and ``high_slope``, so one can continue to write:: + + slope, intercept, low_slope, high_slope = theilslopes(y, x) + + References + ---------- + .. [1] P.K. Sen, "Estimates of the regression coefficient based on + Kendall's tau", J. Am. Stat. Assoc., Vol. 63, pp. 1379-1389, 1968. + .. [2] H. Theil, "A rank-invariant method of linear and polynomial + regression analysis I, II and III", Nederl. Akad. Wetensch., Proc. + 53:, pp. 386-392, pp. 521-525, pp. 1397-1412, 1950. + .. [3] W.L. Conover, "Practical nonparametric statistics", 2nd ed., + John Wiley and Sons, New York, pp. 493. + .. [4] https://en.wikipedia.org/wiki/Theil%E2%80%93Sen_estimator + + Examples + -------- + >>> import numpy as np + >>> from scipy import stats + >>> import matplotlib.pyplot as plt + + >>> x = np.linspace(-5, 5, num=150) + >>> y = x + np.random.normal(size=x.size) + >>> y[11:15] += 10 # add outliers + >>> y[-5:] -= 7 + + Compute the slope, intercept and 90% confidence interval. For comparison, + also compute the least-squares fit with `linregress`: + + >>> res = stats.theilslopes(y, x, 0.90, method='separate') + >>> lsq_res = stats.linregress(x, y) + + Plot the results. The Theil-Sen regression line is shown in red, with the + dashed red lines illustrating the confidence interval of the slope (note + that the dashed red lines are not the confidence interval of the regression + as the confidence interval of the intercept is not included). The green + line shows the least-squares fit for comparison. + + >>> fig = plt.figure() + >>> ax = fig.add_subplot(111) + >>> ax.plot(x, y, 'b.') + >>> ax.plot(x, res[1] + res[0] * x, 'r-') + >>> ax.plot(x, res[1] + res[2] * x, 'r--') + >>> ax.plot(x, res[1] + res[3] * x, 'r--') + >>> ax.plot(x, lsq_res[1] + lsq_res[0] * x, 'g-') + >>> plt.show() + + """ + if method not in ['joint', 'separate']: + raise ValueError("method must be either 'joint' or 'separate'." + f"'{method}' is invalid.") + # We copy both x and y so we can use _find_repeats. + y = np.array(y, dtype=float, copy=True).ravel() + if x is None: + x = np.arange(len(y), dtype=float) + else: + x = np.array(x, dtype=float, copy=True).ravel() + if len(x) != len(y): + raise ValueError("Array shapes are incompatible for broadcasting.") + if len(x) < 2: + raise ValueError("`x` and `y` must have length at least 2.") + + # Compute sorted slopes only when deltax > 0 + deltax = x[:, np.newaxis] - x + deltay = y[:, np.newaxis] - y + slopes = deltay[deltax > 0] / deltax[deltax > 0] + if not slopes.size: + msg = "All `x` coordinates are identical." + warnings.warn(msg, RuntimeWarning, stacklevel=2) + slopes.sort() + medslope = np.median(slopes) + if method == 'joint': + medinter = np.median(y - medslope * x) + else: + medinter = np.median(y) - medslope * np.median(x) + # Now compute confidence intervals + if alpha > 0.5: + alpha = 1. - alpha + + z = distributions.norm.ppf(alpha / 2.) + # This implements (2.6) from Sen (1968) + _, nxreps = _find_repeats(x) + _, nyreps = _find_repeats(y) + nt = len(slopes) # N in Sen (1968) + ny = len(y) # n in Sen (1968) + # Equation 2.6 in Sen (1968): + sigsq = 1/18. * (ny * (ny-1) * (2*ny+5) - + sum(k * (k-1) * (2*k + 5) for k in nxreps) - + sum(k * (k-1) * (2*k + 5) for k in nyreps)) + # Find the confidence interval indices in `slopes` + try: + sigma = np.sqrt(sigsq) + Ru = min(int(np.round((nt - z*sigma)/2.)), len(slopes)-1) + Rl = max(int(np.round((nt + z*sigma)/2.)) - 1, 0) + delta = slopes[[Rl, Ru]] + except (ValueError, IndexError): + delta = (np.nan, np.nan) + + return TheilslopesResult(slope=medslope, intercept=medinter, + low_slope=delta[0], high_slope=delta[1]) + + +def theilslopes(y, x=None, alpha=0.95, method='separate'): + r""" + Computes the Theil-Sen estimator for a set of points (x, y). + + `theilslopes` implements a method for robust linear regression. It + computes the slope as the median of all slopes between paired values. + + Parameters + ---------- + y : array_like + Dependent variable. + x : array_like or None, optional + Independent variable. If None, use ``arange(len(y))`` instead. + alpha : float, optional + Confidence degree between 0 and 1. Default is 95% confidence. + Note that `alpha` is symmetric around 0.5, i.e. both 0.1 and 0.9 are + interpreted as "find the 90% confidence interval". + method : {'joint', 'separate'}, optional + Method to be used for computing estimate for intercept. + Following methods are supported, + + * 'joint': Uses np.median(y - slope * x) as intercept. + * 'separate': Uses np.median(y) - slope * np.median(x) + as intercept. + + The default is 'separate'. + + .. versionadded:: 1.8.0 + + Returns + ------- + result : ``TheilslopesResult`` instance + The return value is an object with the following attributes: + + slope : float + Theil slope. + intercept : float + Intercept of the Theil line. + low_slope : float + Lower bound of the confidence interval on `slope`. + high_slope : float + Upper bound of the confidence interval on `slope`. + + See Also + -------- + siegelslopes : a similar technique using repeated medians + + + Notes + ----- + For more details on `theilslopes`, see `scipy.stats.theilslopes`. + + """ + y = ma.asarray(y).flatten() + if x is None: + x = ma.arange(len(y), dtype=float) + else: + x = ma.asarray(x).flatten() + if len(x) != len(y): + raise ValueError(f"Incompatible lengths ! ({len(y)}<>{len(x)})") + + m = ma.mask_or(ma.getmask(x), ma.getmask(y)) + y._mask = x._mask = m + # Disregard any masked elements of x or y + y = y.compressed() + x = x.compressed().astype(float) + # We now have unmasked arrays so can use `scipy.stats.theilslopes` + return _theilslopes(y, x, alpha=alpha, method=method) + + +def _siegelslopes(y, x=None, method="hierarchical"): + if method not in ['hierarchical', 'separate']: + raise ValueError("method can only be 'hierarchical' or 'separate'") + y = np.asarray(y).ravel() + if x is None: + x = np.arange(len(y), dtype=float) + else: + x = np.asarray(x, dtype=float).ravel() + if len(x) != len(y): + raise ValueError("Array shapes are incompatible for broadcasting.") + if len(x) < 2: + raise ValueError("`x` and `y` must have length at least 2.") + + dtype = np.result_type(x, y, np.float32) # use at least float32 + y, x = y.astype(dtype), x.astype(dtype) + medslope, medinter = siegelslopes_pythran(y, x, method) + medslope, medinter = np.asarray(medslope)[()], np.asarray(medinter)[()] + return SiegelslopesResult(slope=medslope, intercept=medinter) + + +def siegelslopes(y, x=None, method="hierarchical"): + r""" + Computes the Siegel estimator for a set of points (x, y). + + `siegelslopes` implements a method for robust linear regression + using repeated medians to fit a line to the points (x, y). + The method is robust to outliers with an asymptotic breakdown point + of 50%. + + Parameters + ---------- + y : array_like + Dependent variable. + x : array_like or None, optional + Independent variable. If None, use ``arange(len(y))`` instead. + method : {'hierarchical', 'separate'} + If 'hierarchical', estimate the intercept using the estimated + slope ``slope`` (default option). + If 'separate', estimate the intercept independent of the estimated + slope. See Notes for details. + + Returns + ------- + result : ``SiegelslopesResult`` instance + The return value is an object with the following attributes: + + slope : float + Estimate of the slope of the regression line. + intercept : float + Estimate of the intercept of the regression line. + + See Also + -------- + theilslopes : a similar technique without repeated medians + + Notes + ----- + For more details on `siegelslopes`, see `scipy.stats.siegelslopes`. + + """ + y = ma.asarray(y).ravel() + if x is None: + x = ma.arange(len(y), dtype=float) + else: + x = ma.asarray(x).ravel() + if len(x) != len(y): + raise ValueError(f"Incompatible lengths ! ({len(y)}<>{len(x)})") + + m = ma.mask_or(ma.getmask(x), ma.getmask(y)) + y._mask = x._mask = m + # Disregard any masked elements of x or y + y = y.compressed() + x = x.compressed().astype(float) + # We now have unmasked arrays so can use `scipy.stats.siegelslopes` + return _siegelslopes(y, x, method=method) + + +SenSeasonalSlopesResult = _make_tuple_bunch('SenSeasonalSlopesResult', + ['intra_slope', 'inter_slope']) + + +def sen_seasonal_slopes(x): + r""" + Computes seasonal Theil-Sen and Kendall slope estimators. + + The seasonal generalization of Sen's slope computes the slopes between all + pairs of values within a "season" (column) of a 2D array. It returns an + array containing the median of these "within-season" slopes for each + season (the Theil-Sen slope estimator of each season), and it returns the + median of the within-season slopes across all seasons (the seasonal Kendall + slope estimator). + + Parameters + ---------- + x : 2D array_like + Each column of `x` contains measurements of the dependent variable + within a season. The independent variable (usually time) of each season + is assumed to be ``np.arange(x.shape[0])``. + + Returns + ------- + result : ``SenSeasonalSlopesResult`` instance + The return value is an object with the following attributes: + + intra_slope : ndarray + For each season, the Theil-Sen slope estimator: the median of + within-season slopes. + inter_slope : float + The seasonal Kendall slope estimator: the median of within-season + slopes *across all* seasons. + + See Also + -------- + theilslopes : the analogous function for non-seasonal data + scipy.stats.theilslopes : non-seasonal slopes for non-masked arrays + + Notes + ----- + The slopes :math:`d_{ijk}` within season :math:`i` are: + + .. math:: + + d_{ijk} = \frac{x_{ij} - x_{ik}} + {j - k} + + for pairs of distinct integer indices :math:`j, k` of :math:`x`. + + Element :math:`i` of the returned `intra_slope` array is the median of the + :math:`d_{ijk}` over all :math:`j < k`; this is the Theil-Sen slope + estimator of season :math:`i`. The returned `inter_slope` value, better + known as the seasonal Kendall slope estimator, is the median of the + :math:`d_{ijk}` over all :math:`i, j, k`. + + References + ---------- + .. [1] Hirsch, Robert M., James R. Slack, and Richard A. Smith. + "Techniques of trend analysis for monthly water quality data." + *Water Resources Research* 18.1 (1982): 107-121. + + Examples + -------- + Suppose we have 100 observations of a dependent variable for each of four + seasons: + + >>> import numpy as np + >>> rng = np.random.default_rng() + >>> x = rng.random(size=(100, 4)) + + We compute the seasonal slopes as: + + >>> from scipy import stats + >>> intra_slope, inter_slope = stats.mstats.sen_seasonal_slopes(x) + + If we define a function to compute all slopes between observations within + a season: + + >>> def dijk(yi): + ... n = len(yi) + ... x = np.arange(n) + ... dy = yi - yi[:, np.newaxis] + ... dx = x - x[:, np.newaxis] + ... # we only want unique pairs of distinct indices + ... mask = np.triu(np.ones((n, n), dtype=bool), k=1) + ... return dy[mask]/dx[mask] + + then element ``i`` of ``intra_slope`` is the median of ``dijk[x[:, i]]``: + + >>> i = 2 + >>> np.allclose(np.median(dijk(x[:, i])), intra_slope[i]) + True + + and ``inter_slope`` is the median of the values returned by ``dijk`` for + all seasons: + + >>> all_slopes = np.concatenate([dijk(x[:, i]) for i in range(x.shape[1])]) + >>> np.allclose(np.median(all_slopes), inter_slope) + True + + Because the data are randomly generated, we would expect the median slopes + to be nearly zero both within and across all seasons, and indeed they are: + + >>> intra_slope.data + array([ 0.00124504, -0.00277761, -0.00221245, -0.00036338]) + >>> inter_slope + -0.0010511779872922058 + + """ + x = ma.array(x, subok=True, copy=False, ndmin=2) + (n,_) = x.shape + # Get list of slopes per season + szn_slopes = ma.vstack([(x[i+1:]-x[i])/np.arange(1,n-i)[:,None] + for i in range(n)]) + szn_medslopes = ma.median(szn_slopes, axis=0) + medslope = ma.median(szn_slopes, axis=None) + return SenSeasonalSlopesResult(szn_medslopes, medslope) + + +Ttest_1sampResult = namedtuple('Ttest_1sampResult', ('statistic', 'pvalue')) + + +def ttest_1samp(a, popmean, axis=0, alternative='two-sided'): + """ + Calculates the T-test for the mean of ONE group of scores. + + Parameters + ---------- + a : array_like + sample observation + popmean : float or array_like + expected value in null hypothesis, if array_like than it must have the + same shape as `a` excluding the axis dimension + axis : int or None, optional + Axis along which to compute test. If None, compute over the whole + array `a`. + alternative : {'two-sided', 'less', 'greater'}, optional + Defines the alternative hypothesis. + The following options are available (default is 'two-sided'): + + * 'two-sided': the mean of the underlying distribution of the sample + is different than the given population mean (`popmean`) + * 'less': the mean of the underlying distribution of the sample is + less than the given population mean (`popmean`) + * 'greater': the mean of the underlying distribution of the sample is + greater than the given population mean (`popmean`) + + .. versionadded:: 1.7.0 + + Returns + ------- + statistic : float or array + t-statistic + pvalue : float or array + The p-value + + Notes + ----- + For more details on `ttest_1samp`, see `scipy.stats.ttest_1samp`. + + """ + a, axis = _chk_asarray(a, axis) + if a.size == 0: + return (np.nan, np.nan) + + x = a.mean(axis=axis) + v = a.var(axis=axis, ddof=1) + n = a.count(axis=axis) + # force df to be an array for masked division not to throw a warning + df = ma.asanyarray(n - 1.0) + svar = ((n - 1.0) * v) / df + with np.errstate(divide='ignore', invalid='ignore'): + t = (x - popmean) / ma.sqrt(svar / n) + + t, prob = _ttest_finish(df, t, alternative) + return Ttest_1sampResult(t, prob) + + +ttest_onesamp = ttest_1samp + + +Ttest_indResult = namedtuple('Ttest_indResult', ('statistic', 'pvalue')) + + +def ttest_ind(a, b, axis=0, equal_var=True, alternative='two-sided'): + """ + Calculates the T-test for the means of TWO INDEPENDENT samples of scores. + + Parameters + ---------- + a, b : array_like + The arrays must have the same shape, except in the dimension + corresponding to `axis` (the first, by default). + axis : int or None, optional + Axis along which to compute test. If None, compute over the whole + arrays, `a`, and `b`. + equal_var : bool, optional + If True, perform a standard independent 2 sample test that assumes equal + population variances. + If False, perform Welch's t-test, which does not assume equal population + variance. + + .. versionadded:: 0.17.0 + alternative : {'two-sided', 'less', 'greater'}, optional + Defines the alternative hypothesis. + The following options are available (default is 'two-sided'): + + * 'two-sided': the means of the distributions underlying the samples + are unequal. + * 'less': the mean of the distribution underlying the first sample + is less than the mean of the distribution underlying the second + sample. + * 'greater': the mean of the distribution underlying the first + sample is greater than the mean of the distribution underlying + the second sample. + + .. versionadded:: 1.7.0 + + Returns + ------- + statistic : float or array + The calculated t-statistic. + pvalue : float or array + The p-value. + + Notes + ----- + For more details on `ttest_ind`, see `scipy.stats.ttest_ind`. + + """ + a, b, axis = _chk2_asarray(a, b, axis) + + if a.size == 0 or b.size == 0: + return Ttest_indResult(np.nan, np.nan) + + (x1, x2) = (a.mean(axis), b.mean(axis)) + (v1, v2) = (a.var(axis=axis, ddof=1), b.var(axis=axis, ddof=1)) + (n1, n2) = (a.count(axis), b.count(axis)) + + if equal_var: + # force df to be an array for masked division not to throw a warning + df = ma.asanyarray(n1 + n2 - 2.0) + svar = ((n1-1)*v1+(n2-1)*v2) / df + denom = ma.sqrt(svar*(1.0/n1 + 1.0/n2)) # n-D computation here! + else: + vn1 = v1/n1 + vn2 = v2/n2 + with np.errstate(divide='ignore', invalid='ignore'): + df = (vn1 + vn2)**2 / (vn1**2 / (n1 - 1) + vn2**2 / (n2 - 1)) + + # If df is undefined, variances are zero. + # It doesn't matter what df is as long as it is not NaN. + df = np.where(np.isnan(df), 1, df) + denom = ma.sqrt(vn1 + vn2) + + with np.errstate(divide='ignore', invalid='ignore'): + t = (x1-x2) / denom + + t, prob = _ttest_finish(df, t, alternative) + return Ttest_indResult(t, prob) + + +Ttest_relResult = namedtuple('Ttest_relResult', ('statistic', 'pvalue')) + + +def ttest_rel(a, b, axis=0, alternative='two-sided'): + """ + Calculates the T-test on TWO RELATED samples of scores, a and b. + + Parameters + ---------- + a, b : array_like + The arrays must have the same shape. + axis : int or None, optional + Axis along which to compute test. If None, compute over the whole + arrays, `a`, and `b`. + alternative : {'two-sided', 'less', 'greater'}, optional + Defines the alternative hypothesis. + The following options are available (default is 'two-sided'): + + * 'two-sided': the means of the distributions underlying the samples + are unequal. + * 'less': the mean of the distribution underlying the first sample + is less than the mean of the distribution underlying the second + sample. + * 'greater': the mean of the distribution underlying the first + sample is greater than the mean of the distribution underlying + the second sample. + + .. versionadded:: 1.7.0 + + Returns + ------- + statistic : float or array + t-statistic + pvalue : float or array + two-tailed p-value + + Notes + ----- + For more details on `ttest_rel`, see `scipy.stats.ttest_rel`. + + """ + a, b, axis = _chk2_asarray(a, b, axis) + if len(a) != len(b): + raise ValueError('unequal length arrays') + + if a.size == 0 or b.size == 0: + return Ttest_relResult(np.nan, np.nan) + + n = a.count(axis) + df = ma.asanyarray(n-1.0) + d = (a-b).astype('d') + dm = d.mean(axis) + v = d.var(axis=axis, ddof=1) + denom = ma.sqrt(v / n) + with np.errstate(divide='ignore', invalid='ignore'): + t = dm / denom + + t, prob = _ttest_finish(df, t, alternative) + return Ttest_relResult(t, prob) + + +MannwhitneyuResult = namedtuple('MannwhitneyuResult', ('statistic', + 'pvalue')) + + +def mannwhitneyu(x,y, use_continuity=True): + """ + Computes the Mann-Whitney statistic + + Missing values in `x` and/or `y` are discarded. + + Parameters + ---------- + x : sequence + Input + y : sequence + Input + use_continuity : {True, False}, optional + Whether a continuity correction (1/2.) should be taken into account. + + Returns + ------- + statistic : float + The minimum of the Mann-Whitney statistics + pvalue : float + Approximate two-sided p-value assuming a normal distribution. + + """ + x = ma.asarray(x).compressed().view(ndarray) + y = ma.asarray(y).compressed().view(ndarray) + ranks = rankdata(np.concatenate([x,y])) + (nx, ny) = (len(x), len(y)) + nt = nx + ny + U = ranks[:nx].sum() - nx*(nx+1)/2. + U = max(U, nx*ny - U) + u = nx*ny - U + + mu = (nx*ny)/2. + sigsq = (nt**3 - nt)/12. + ties = count_tied_groups(ranks) + sigsq -= sum(v*(k**3-k) for (k,v) in ties.items())/12. + sigsq *= nx*ny/float(nt*(nt-1)) + + if use_continuity: + z = (U - 1/2. - mu) / ma.sqrt(sigsq) + else: + z = (U - mu) / ma.sqrt(sigsq) + + prob = special.erfc(abs(z)/np.sqrt(2)) + return MannwhitneyuResult(u, prob) + + +KruskalResult = namedtuple('KruskalResult', ('statistic', 'pvalue')) + + +def kruskal(*samples): + """ + Compute the Kruskal-Wallis H-test for independent samples + + Parameters + ---------- + *samples : array_like + Two or more arrays with the sample measurements can be given as + arguments. + + Returns + ------- + statistic : float + The Kruskal-Wallis H statistic, corrected for ties + pvalue : float + The p-value for the test using the assumption that H has a chi + square distribution + + Notes + ----- + For more details on `kruskal`, see `scipy.stats.kruskal`. + + Examples + -------- + >>> from scipy.stats.mstats import kruskal + + Random samples from three different brands of batteries were tested + to see how long the charge lasted. Results were as follows: + + >>> a = [6.3, 5.4, 5.7, 5.2, 5.0] + >>> b = [6.9, 7.0, 6.1, 7.9] + >>> c = [7.2, 6.9, 6.1, 6.5] + + Test the hypothesis that the distribution functions for all of the brands' + durations are identical. Use 5% level of significance. + + >>> kruskal(a, b, c) + KruskalResult(statistic=7.113812154696133, pvalue=0.028526948491942164) + + The null hypothesis is rejected at the 5% level of significance + because the returned p-value is less than the critical value of 5%. + + """ + output = argstoarray(*samples) + ranks = ma.masked_equal(rankdata(output, use_missing=False), 0) + sumrk = ranks.sum(-1) + ngrp = ranks.count(-1) + ntot = ranks.count() + H = 12./(ntot*(ntot+1)) * (sumrk**2/ngrp).sum() - 3*(ntot+1) + # Tie correction + ties = count_tied_groups(ranks) + T = 1. - sum(v*(k**3-k) for (k,v) in ties.items())/float(ntot**3-ntot) + if T == 0: + raise ValueError('All numbers are identical in kruskal') + + H /= T + df = len(output) - 1 + prob = distributions.chi2.sf(H, df) + return KruskalResult(H, prob) + + +kruskalwallis = kruskal + + +@_rename_parameter("mode", "method") +def ks_1samp(x, cdf, args=(), alternative="two-sided", method='auto'): + """ + Computes the Kolmogorov-Smirnov test on one sample of masked values. + + Missing values in `x` are discarded. + + Parameters + ---------- + x : array_like + a 1-D array of observations of random variables. + cdf : str or callable + If a string, it should be the name of a distribution in `scipy.stats`. + If a callable, that callable is used to calculate the cdf. + args : tuple, sequence, optional + Distribution parameters, used if `cdf` is a string. + alternative : {'two-sided', 'less', 'greater'}, optional + Indicates the alternative hypothesis. Default is 'two-sided'. + method : {'auto', 'exact', 'asymp'}, optional + Defines the method used for calculating the p-value. + The following options are available (default is 'auto'): + + * 'auto' : use 'exact' for small size arrays, 'asymp' for large + * 'exact' : use approximation to exact distribution of test statistic + * 'asymp' : use asymptotic distribution of test statistic + + Returns + ------- + d : float + Value of the Kolmogorov Smirnov test + p : float + Corresponding p-value. + + """ + alternative = {'t': 'two-sided', 'g': 'greater', 'l': 'less'}.get( + alternative.lower()[0], alternative) + return scipy.stats._stats_py.ks_1samp( + x, cdf, args=args, alternative=alternative, method=method) + + +@_rename_parameter("mode", "method") +def ks_2samp(data1, data2, alternative="two-sided", method='auto'): + """ + Computes the Kolmogorov-Smirnov test on two samples. + + Missing values in `x` and/or `y` are discarded. + + Parameters + ---------- + data1 : array_like + First data set + data2 : array_like + Second data set + alternative : {'two-sided', 'less', 'greater'}, optional + Indicates the alternative hypothesis. Default is 'two-sided'. + method : {'auto', 'exact', 'asymp'}, optional + Defines the method used for calculating the p-value. + The following options are available (default is 'auto'): + + * 'auto' : use 'exact' for small size arrays, 'asymp' for large + * 'exact' : use approximation to exact distribution of test statistic + * 'asymp' : use asymptotic distribution of test statistic + + Returns + ------- + d : float + Value of the Kolmogorov Smirnov test + p : float + Corresponding p-value. + + """ + # Ideally this would be accomplished by + # ks_2samp = scipy.stats._stats_py.ks_2samp + # but the circular dependencies between _mstats_basic and stats prevent that. + alternative = {'t': 'two-sided', 'g': 'greater', 'l': 'less'}.get( + alternative.lower()[0], alternative) + return scipy.stats._stats_py.ks_2samp(data1, data2, + alternative=alternative, + method=method) + + +ks_twosamp = ks_2samp + + +@_rename_parameter("mode", "method") +def kstest(data1, data2, args=(), alternative='two-sided', method='auto'): + """ + Performs the Kolmogorov-Smirnov test for goodness of fit. + + Parameters + ---------- + data1 : array_like + data2 : str, callable or array_like + args : tuple, sequence, optional + Distribution parameters, used if `data1` or `data2` are strings. + alternative : str, as documented in stats.kstest + method : str, as documented in stats.kstest + + Returns + ------- + tuple of (K-S statistic, probability) + + """ # numpydoc ignore=RT03 + return scipy.stats._stats_py.kstest(data1, data2, args, + alternative=alternative, method=method) + + +def trima(a, limits=None, inclusive=(True,True)): + """ + Trims an array by masking the data outside some given limits. + + Returns a masked version of the input array. + + Parameters + ---------- + a : array_like + Input array. + limits : {None, tuple}, optional + Tuple of (lower limit, upper limit) in absolute values. + Values of the input array lower (greater) than the lower (upper) limit + will be masked. A limit is None indicates an open interval. + inclusive : (bool, bool) tuple, optional + Tuple of (lower flag, upper flag), indicating whether values exactly + equal to the lower (upper) limit are allowed. + + Examples + -------- + >>> from scipy.stats.mstats import trima + >>> import numpy as np + + >>> a = np.arange(10) + + The interval is left-closed and right-open, i.e., `[2, 8)`. + Trim the array by keeping only values in the interval. + + >>> trima(a, limits=(2, 8), inclusive=(True, False)) + masked_array(data=[--, --, 2, 3, 4, 5, 6, 7, --, --], + mask=[ True, True, False, False, False, False, False, False, + True, True], + fill_value=999999) + + """ # numpydoc ignore=RT01 + a = ma.asarray(a) + a.unshare_mask() + if (limits is None) or (limits == (None, None)): + return a + + (lower_lim, upper_lim) = limits + (lower_in, upper_in) = inclusive + condition = False + if lower_lim is not None: + if lower_in: + condition |= (a < lower_lim) + else: + condition |= (a <= lower_lim) + + if upper_lim is not None: + if upper_in: + condition |= (a > upper_lim) + else: + condition |= (a >= upper_lim) + + a[condition.filled(True)] = masked + return a + + +def trimr(a, limits=None, inclusive=(True, True), axis=None): + """ + Trims an array by masking some proportion of the data on each end. + Returns a masked version of the input array. + + Parameters + ---------- + a : sequence + Input array. + limits : {None, tuple}, optional + Tuple of the percentages to cut on each side of the array, with respect + to the number of unmasked data, as floats between 0. and 1. + Noting n the number of unmasked data before trimming, the + (n*limits[0])th smallest data and the (n*limits[1])th largest data are + masked, and the total number of unmasked data after trimming is + n*(1.-sum(limits)). The value of one limit can be set to None to + indicate an open interval. + inclusive : {(True,True) tuple}, optional + Tuple of flags indicating whether the number of data being masked on + the left (right) end should be truncated (True) or rounded (False) to + integers. + axis : {None,int}, optional + Axis along which to trim. If None, the whole array is trimmed, but its + shape is maintained. + + """ # numpydoc ignore=RT01 + def _trimr1D(a, low_limit, up_limit, low_inclusive, up_inclusive): + n = a.count() + idx = a.argsort() + if low_limit: + if low_inclusive: + lowidx = int(low_limit*n) + else: + lowidx = int(np.round(low_limit*n)) + a[idx[:lowidx]] = masked + if up_limit is not None: + if up_inclusive: + upidx = n - int(n*up_limit) + else: + upidx = n - int(np.round(n*up_limit)) + a[idx[upidx:]] = masked + return a + + a = ma.asarray(a) + a.unshare_mask() + if limits is None: + return a + + # Check the limits + (lolim, uplim) = limits + errmsg = "The proportion to cut from the %s should be between 0. and 1." + if lolim is not None: + if lolim > 1. or lolim < 0: + raise ValueError(errmsg % 'beginning' + f"(got {lolim})") + if uplim is not None: + if uplim > 1. or uplim < 0: + raise ValueError(errmsg % 'end' + f"(got {uplim})") + + (loinc, upinc) = inclusive + + if axis is None: + shp = a.shape + return _trimr1D(a.ravel(),lolim,uplim,loinc,upinc).reshape(shp) + else: + return ma.apply_along_axis(_trimr1D, axis, a, lolim,uplim,loinc,upinc) + + +trimdoc = _dedent_for_py313(""" + Parameters + ---------- + a : sequence + Input array + limits : {None, tuple}, optional + If `relative` is False, tuple (lower limit, upper limit) in absolute values. + Values of the input array lower (greater) than the lower (upper) limit are + masked. + + If `relative` is True, tuple (lower percentage, upper percentage) to cut + on each side of the array, with respect to the number of unmasked data. + + Noting n the number of unmasked data before trimming, the (n*limits[0])th + smallest data and the (n*limits[1])th largest data are masked, and the + total number of unmasked data after trimming is n*(1.-sum(limits)) + In each case, the value of one limit can be set to None to indicate an + open interval. + + If limits is None, no trimming is performed + inclusive : {(bool, bool) tuple}, optional + If `relative` is False, tuple indicating whether values exactly equal + to the absolute limits are allowed. + If `relative` is True, tuple indicating whether the number of data + being masked on each side should be rounded (True) or truncated + (False). + relative : bool, optional + Whether to consider the limits as absolute values (False) or proportions + to cut (True). + axis : int, optional + Axis along which to trim.""") + + +def trim(a, limits=None, inclusive=(True,True), relative=False, axis=None): + """ + Trims an array by masking the data outside some given limits. + + Returns a masked version of the input array. + + %s + + Examples + -------- + >>> from scipy.stats.mstats import trim + >>> z = [ 1, 2, 3, 4, 5, 6, 7, 8, 9,10] + >>> print(trim(z,(3,8))) + [-- -- 3 4 5 6 7 8 -- --] + >>> print(trim(z,(0.1,0.2),relative=True)) + [-- 2 3 4 5 6 7 8 -- --] + + """ # numpydoc ignore=RT01 + if relative: + return trimr(a, limits=limits, inclusive=inclusive, axis=axis) + else: + return trima(a, limits=limits, inclusive=inclusive) + + +if trim.__doc__: + trim.__doc__ = trim.__doc__ % trimdoc + + +def trimboth(data, proportiontocut=0.2, inclusive=(True,True), axis=None): + """ + Trims the smallest and largest data values. + + Trims the `data` by masking the ``int(proportiontocut * n)`` smallest and + ``int(proportiontocut * n)`` largest values of data along the given axis, + where n is the number of unmasked values before trimming. + + Parameters + ---------- + data : ndarray + Data to trim. + proportiontocut : float, optional + Percentage of trimming (as a float between 0 and 1). + If n is the number of unmasked values before trimming, the number of + values after trimming is ``(1 - 2*proportiontocut) * n``. + Default is 0.2. + inclusive : {(bool, bool) tuple}, optional + Tuple indicating whether the number of data being masked on each side + should be rounded (True) or truncated (False). + axis : int, optional + Axis along which to perform the trimming. + If None, the input array is first flattened. + + """ # numpydoc ignore=RT01 + return trimr(data, limits=(proportiontocut,proportiontocut), + inclusive=inclusive, axis=axis) + + +def trimtail(data, proportiontocut=0.2, tail='left', inclusive=(True,True), + axis=None): + """ + Trims the data by masking values from one tail. + + Parameters + ---------- + data : array_like + Data to trim. + proportiontocut : float, optional + Percentage of trimming. If n is the number of unmasked values + before trimming, the number of values after trimming is + ``(1 - proportiontocut) * n``. Default is 0.2. + tail : {'left','right'}, optional + If 'left' the `proportiontocut` lowest values will be masked. + If 'right' the `proportiontocut` highest values will be masked. + Default is 'left'. + inclusive : {(bool, bool) tuple}, optional + Tuple indicating whether the number of data being masked on each side + should be rounded (True) or truncated (False). Default is + (True, True). + axis : int, optional + Axis along which to perform the trimming. + If None, the input array is first flattened. Default is None. + + Returns + ------- + trimtail : ndarray + Returned array of same shape as `data` with masked tail values. + + """ + tail = str(tail).lower()[0] + if tail == 'l': + limits = (proportiontocut,None) + elif tail == 'r': + limits = (None, proportiontocut) + else: + raise TypeError("The tail argument should be in ('left','right')") + + return trimr(data, limits=limits, axis=axis, inclusive=inclusive) + + +trim1 = trimtail + + +def trimmed_mean(a, limits=(0.1,0.1), inclusive=(1,1), relative=True, + axis=None): + """Returns the trimmed mean of the data along the given axis. + + %s + + """ # numpydoc ignore=RT01 + if (not isinstance(limits,tuple)) and isinstance(limits,float): + limits = (limits, limits) + if relative: + return trimr(a,limits=limits,inclusive=inclusive,axis=axis).mean(axis=axis) + else: + return trima(a,limits=limits,inclusive=inclusive).mean(axis=axis) + + +if trimmed_mean.__doc__: + trimmed_mean.__doc__ = trimmed_mean.__doc__ % trimdoc + + +def trimmed_var(a, limits=(0.1,0.1), inclusive=(1,1), relative=True, + axis=None, ddof=0): + """Returns the trimmed variance of the data along the given axis. + + %s + ddof : {0,integer}, optional + Means Delta Degrees of Freedom. The denominator used during computations + is (n-ddof). DDOF=0 corresponds to a biased estimate, DDOF=1 to an un- + biased estimate of the variance. + + """ # numpydoc ignore=RT01 + if (not isinstance(limits,tuple)) and isinstance(limits,float): + limits = (limits, limits) + if relative: + out = trimr(a,limits=limits, inclusive=inclusive,axis=axis) + else: + out = trima(a,limits=limits,inclusive=inclusive) + + return out.var(axis=axis, ddof=ddof) + + +if trimmed_var.__doc__: + trimmed_var.__doc__ = trimmed_var.__doc__ % trimdoc + + +def trimmed_std(a, limits=(0.1,0.1), inclusive=(1,1), relative=True, + axis=None, ddof=0): + """Returns the trimmed standard deviation of the data along the given axis. + + %s + ddof : {0,integer}, optional + Means Delta Degrees of Freedom. The denominator used during computations + is (n-ddof). DDOF=0 corresponds to a biased estimate, DDOF=1 to an un- + biased estimate of the variance. + + """ # numpydoc ignore=RT01 + if (not isinstance(limits,tuple)) and isinstance(limits,float): + limits = (limits, limits) + if relative: + out = trimr(a,limits=limits,inclusive=inclusive,axis=axis) + else: + out = trima(a,limits=limits,inclusive=inclusive) + return out.std(axis=axis,ddof=ddof) + + +if trimmed_std.__doc__: + trimmed_std.__doc__ = trimmed_std.__doc__ % trimdoc + + +def trimmed_stde(a, limits=(0.1,0.1), inclusive=(1,1), axis=None): + """ + Returns the standard error of the trimmed mean along the given axis. + + Parameters + ---------- + a : sequence + Input array + limits : {(0.1,0.1), tuple of float}, optional + tuple (lower percentage, upper percentage) to cut on each side of the + array, with respect to the number of unmasked data. + + If n is the number of unmasked data before trimming, the values + smaller than ``n * limits[0]`` and the values larger than + ``n * `limits[1]`` are masked, and the total number of unmasked + data after trimming is ``n * (1.-sum(limits))``. In each case, + the value of one limit can be set to None to indicate an open interval. + If `limits` is None, no trimming is performed. + inclusive : {(bool, bool) tuple} optional + Tuple indicating whether the number of data being masked on each side + should be rounded (True) or truncated (False). + axis : int, optional + Axis along which to trim. + + Returns + ------- + trimmed_stde : scalar or ndarray + + """ # numpydoc ignore=RT03 + def _trimmed_stde_1D(a, low_limit, up_limit, low_inclusive, up_inclusive): + "Returns the standard error of the trimmed mean for a 1D input data." + n = a.count() + idx = a.argsort() + if low_limit: + if low_inclusive: + lowidx = int(low_limit*n) + else: + lowidx = np.round(low_limit*n) + a[idx[:lowidx]] = masked + if up_limit is not None: + if up_inclusive: + upidx = n - int(n*up_limit) + else: + upidx = n - np.round(n*up_limit) + a[idx[upidx:]] = masked + a[idx[:lowidx]] = a[idx[lowidx]] + a[idx[upidx:]] = a[idx[upidx-1]] + winstd = a.std(ddof=1) + return winstd / ((1-low_limit-up_limit)*np.sqrt(len(a))) + + a = ma.array(a, copy=True, subok=True) + a.unshare_mask() + if limits is None: + return a.std(axis=axis,ddof=1)/ma.sqrt(a.count(axis)) + if (not isinstance(limits,tuple)) and isinstance(limits,float): + limits = (limits, limits) + + # Check the limits + (lolim, uplim) = limits + errmsg = "The proportion to cut from the %s should be between 0. and 1." + if lolim is not None: + if lolim > 1. or lolim < 0: + raise ValueError(errmsg % 'beginning' + f"(got {lolim})") + if uplim is not None: + if uplim > 1. or uplim < 0: + raise ValueError(errmsg % 'end' + f"(got {uplim})") + + (loinc, upinc) = inclusive + if (axis is None): + return _trimmed_stde_1D(a.ravel(),lolim,uplim,loinc,upinc) + else: + if a.ndim > 2: + raise ValueError(f"Array 'a' must be at most two dimensional, " + f"but got a.ndim = {a.ndim}") + return ma.apply_along_axis(_trimmed_stde_1D, axis, a, + lolim,uplim,loinc,upinc) + + +def _mask_to_limits(a, limits, inclusive): + """Mask an array for values outside of given limits. + + This is primarily a utility function. + + Parameters + ---------- + a : array + limits : (float or None, float or None) + A tuple consisting of the (lower limit, upper limit). Values in the + input array less than the lower limit or greater than the upper limit + will be masked out. None implies no limit. + inclusive : (bool, bool) + A tuple consisting of the (lower flag, upper flag). These flags + determine whether values exactly equal to lower or upper are allowed. + + Returns + ------- + A MaskedArray. + + Raises + ------ + A ValueError if there are no values within the given limits. + """ + lower_limit, upper_limit = limits + lower_include, upper_include = inclusive + am = ma.MaskedArray(a) + if lower_limit is not None: + if lower_include: + am = ma.masked_less(am, lower_limit) + else: + am = ma.masked_less_equal(am, lower_limit) + + if upper_limit is not None: + if upper_include: + am = ma.masked_greater(am, upper_limit) + else: + am = ma.masked_greater_equal(am, upper_limit) + + if am.count() == 0: + raise ValueError("No array values within given limits") + + return am + + +def tmean(a, limits=None, inclusive=(True, True), axis=None): + """ + Compute the trimmed mean. + + Parameters + ---------- + a : array_like + Array of values. + limits : None or (lower limit, upper limit), optional + Values in the input array less than the lower limit or greater than the + upper limit will be ignored. When limits is None (default), then all + values are used. Either of the limit values in the tuple can also be + None representing a half-open interval. + inclusive : (bool, bool), optional + A tuple consisting of the (lower flag, upper flag). These flags + determine whether values exactly equal to the lower or upper limits + are included. The default value is (True, True). + axis : int or None, optional + Axis along which to operate. If None, compute over the + whole array. Default is None. + + Returns + ------- + tmean : float + + Notes + ----- + For more details on `tmean`, see `scipy.stats.tmean`. + + Examples + -------- + >>> import numpy as np + >>> from scipy.stats import mstats + >>> a = np.array([[6, 8, 3, 0], + ... [3, 9, 1, 2], + ... [8, 7, 8, 2], + ... [5, 6, 0, 2], + ... [4, 5, 5, 2]]) + ... + ... + >>> mstats.tmean(a, (2,5)) + 3.3 + >>> mstats.tmean(a, (2,5), axis=0) + masked_array(data=[4.0, 5.0, 4.0, 2.0], + mask=[False, False, False, False], + fill_value=1e+20) + + """ # numpydoc ignore=RT03 + return trima(a, limits=limits, inclusive=inclusive).mean(axis=axis) + + +def tvar(a, limits=None, inclusive=(True, True), axis=0, ddof=1): + """ + Compute the trimmed variance + + This function computes the sample variance of an array of values, + while ignoring values which are outside of given `limits`. + + Parameters + ---------- + a : array_like + Array of values. + limits : None or (lower limit, upper limit), optional + Values in the input array less than the lower limit or greater than the + upper limit will be ignored. When limits is None, then all values are + used. Either of the limit values in the tuple can also be None + representing a half-open interval. The default value is None. + inclusive : (bool, bool), optional + A tuple consisting of the (lower flag, upper flag). These flags + determine whether values exactly equal to the lower or upper limits + are included. The default value is (True, True). + axis : int or None, optional + Axis along which to operate. If None, compute over the + whole array. Default is zero. + ddof : int, optional + Delta degrees of freedom. Default is 1. + + Returns + ------- + tvar : float + Trimmed variance. + + Notes + ----- + For more details on `tvar`, see `scipy.stats.tvar`. + + """ + a = a.astype(float).ravel() + if limits is None: + n = (~a.mask).sum() # todo: better way to do that? + return np.ma.var(a) * n/(n-1.) + am = _mask_to_limits(a, limits=limits, inclusive=inclusive) + + return np.ma.var(am, axis=axis, ddof=ddof) + + +def tmin(a, lowerlimit=None, axis=0, inclusive=True): + """ + Compute the trimmed minimum + + Parameters + ---------- + a : array_like + array of values + lowerlimit : None or float, optional + Values in the input array less than the given limit will be ignored. + When lowerlimit is None, then all values are used. The default value + is None. + axis : int or None, optional + Axis along which to operate. Default is 0. If None, compute over the + whole array `a`. + inclusive : {True, False}, optional + This flag determines whether values exactly equal to the lower limit + are included. The default value is True. + + Returns + ------- + tmin : float, int or ndarray + + Notes + ----- + For more details on `tmin`, see `scipy.stats.tmin`. + + Examples + -------- + >>> import numpy as np + >>> from scipy.stats import mstats + >>> a = np.array([[6, 8, 3, 0], + ... [3, 2, 1, 2], + ... [8, 1, 8, 2], + ... [5, 3, 0, 2], + ... [4, 7, 5, 2]]) + ... + >>> mstats.tmin(a, 5) + masked_array(data=[5, 7, 5, --], + mask=[False, False, False, True], + fill_value=999999) + + """ # numpydoc ignore=RT03 + a, axis = _chk_asarray(a, axis) + am = trima(a, (lowerlimit, None), (inclusive, False)) + return ma.minimum.reduce(am, axis) + + +def tmax(a, upperlimit=None, axis=0, inclusive=True): + """ + Compute the trimmed maximum + + This function computes the maximum value of an array along a given axis, + while ignoring values larger than a specified upper limit. + + Parameters + ---------- + a : array_like + array of values + upperlimit : None or float, optional + Values in the input array greater than the given limit will be ignored. + When upperlimit is None, then all values are used. The default value + is None. + axis : int or None, optional + Axis along which to operate. Default is 0. If None, compute over the + whole array `a`. + inclusive : {True, False}, optional + This flag determines whether values exactly equal to the upper limit + are included. The default value is True. + + Returns + ------- + tmax : float, int or ndarray + + Notes + ----- + For more details on `tmax`, see `scipy.stats.tmax`. + + Examples + -------- + >>> import numpy as np + >>> from scipy.stats import mstats + >>> a = np.array([[6, 8, 3, 0], + ... [3, 9, 1, 2], + ... [8, 7, 8, 2], + ... [5, 6, 0, 2], + ... [4, 5, 5, 2]]) + ... + ... + >>> mstats.tmax(a, 4) + masked_array(data=[4, --, 3, 2], + mask=[False, True, False, False], + fill_value=999999) + + """ # numpydoc ignore=RT03 + a, axis = _chk_asarray(a, axis) + am = trima(a, (None, upperlimit), (False, inclusive)) + return ma.maximum.reduce(am, axis) + + +def tsem(a, limits=None, inclusive=(True, True), axis=0, ddof=1): + """ + Compute the trimmed standard error of the mean. + + This function finds the standard error of the mean for given + values, ignoring values outside the given `limits`. + + Parameters + ---------- + a : array_like + array of values + limits : None or (lower limit, upper limit), optional + Values in the input array less than the lower limit or greater than the + upper limit will be ignored. When limits is None, then all values are + used. Either of the limit values in the tuple can also be None + representing a half-open interval. The default value is None. + inclusive : (bool, bool), optional + A tuple consisting of the (lower flag, upper flag). These flags + determine whether values exactly equal to the lower or upper limits + are included. The default value is (True, True). + axis : int or None, optional + Axis along which to operate. If None, compute over the + whole array. Default is zero. + ddof : int, optional + Delta degrees of freedom. Default is 1. + + Returns + ------- + tsem : float + + Notes + ----- + For more details on `tsem`, see `scipy.stats.tsem`. + + """ # numpydoc ignore=RT03 + a = ma.asarray(a).ravel() + if limits is None: + n = float(a.count()) + return a.std(axis=axis, ddof=ddof)/ma.sqrt(n) + + am = trima(a.ravel(), limits, inclusive) + sd = np.sqrt(am.var(axis=axis, ddof=ddof)) + return sd / np.sqrt(am.count()) + + +def winsorize(a, limits=None, inclusive=(True, True), inplace=False, + axis=None, nan_policy='propagate'): + """Returns a Winsorized version of the input array. + + The (limits[0])th lowest values are set to the (limits[0])th percentile, + and the (limits[1])th highest values are set to the (1 - limits[1])th + percentile. + Masked values are skipped. + + + Parameters + ---------- + a : sequence + Input array. + limits : {None, tuple of float}, optional + Tuple of the percentages to cut on each side of the array, with respect + to the number of unmasked data, as floats between 0. and 1. + Noting n the number of unmasked data before trimming, the + (n*limits[0])th smallest data and the (n*limits[1])th largest data are + masked, and the total number of unmasked data after trimming + is n*(1.-sum(limits)) The value of one limit can be set to None to + indicate an open interval. + inclusive : {(True, True) tuple}, optional + Tuple indicating whether the number of data being masked on each side + should be truncated (True) or rounded (False). + inplace : {False, True}, optional + Whether to winsorize in place (True) or to use a copy (False) + axis : {None, int}, optional + Axis along which to trim. If None, the whole array is trimmed, but its + shape is maintained. + nan_policy : {'propagate', 'raise', 'omit'}, optional + Defines how to handle when input contains nan. + The following options are available (default is 'propagate'): + + * 'propagate': allows nan values and may overwrite or propagate them + * 'raise': throws an error + * 'omit': performs the calculations ignoring nan values + + Notes + ----- + This function is applied to reduce the effect of possibly spurious outliers + by limiting the extreme values. + + Examples + -------- + >>> import numpy as np + >>> from scipy.stats.mstats import winsorize + + A shuffled array contains integers from 1 to 10. + + >>> a = np.array([10, 4, 9, 8, 5, 3, 7, 2, 1, 6]) + + The 10% of the lowest value (i.e., ``1``) and the 20% of the highest + values (i.e., ``9`` and ``10``) are replaced. + + >>> winsorize(a, limits=[0.1, 0.2]) + masked_array(data=[8, 4, 8, 8, 5, 3, 7, 2, 2, 6], + mask=False, + fill_value=999999) + + """ # numpydoc ignore=RT01 + def _winsorize1D(a, low_limit, up_limit, low_include, up_include, + contains_nan, nan_policy): + n = a.count() + idx = a.argsort() + if contains_nan: + nan_count = np.count_nonzero(np.isnan(a)) + if low_limit: + if low_include: + lowidx = int(low_limit * n) + else: + lowidx = np.round(low_limit * n).astype(int) + if contains_nan and nan_policy == 'omit': + lowidx = min(lowidx, n-nan_count-1) + a[idx[:lowidx]] = a[idx[lowidx]] + if up_limit is not None: + if up_include: + upidx = n - int(n * up_limit) + else: + upidx = n - np.round(n * up_limit).astype(int) + if contains_nan and nan_policy == 'omit': + a[idx[upidx:-nan_count]] = a[idx[upidx - 1]] + else: + a[idx[upidx:]] = a[idx[upidx - 1]] + return a + + contains_nan = _contains_nan(a, nan_policy) + # We are going to modify a: better make a copy + a = ma.array(a, copy=np.logical_not(inplace)) + + if limits is None: + return a + if (not isinstance(limits, tuple)) and isinstance(limits, float): + limits = (limits, limits) + + # Check the limits + (lolim, uplim) = limits + errmsg = "The proportion to cut from the %s should be between 0. and 1." + if lolim is not None: + if lolim > 1. or lolim < 0: + raise ValueError(errmsg % 'beginning' + f"(got {lolim})") + if uplim is not None: + if uplim > 1. or uplim < 0: + raise ValueError(errmsg % 'end' + f"(got {uplim})") + + (loinc, upinc) = inclusive + + if axis is None: + shp = a.shape + return _winsorize1D(a.ravel(), lolim, uplim, loinc, upinc, + contains_nan, nan_policy).reshape(shp) + else: + return ma.apply_along_axis(_winsorize1D, axis, a, lolim, uplim, loinc, + upinc, contains_nan, nan_policy) + + +def moment(a, moment=1, axis=0): + """ + Calculates the nth moment about the mean for a sample. + + Parameters + ---------- + a : array_like + data + moment : int, optional + order of central moment that is returned + axis : int or None, optional + Axis along which the central moment is computed. Default is 0. + If None, compute over the whole array `a`. + + Returns + ------- + n-th central moment : ndarray or float + The appropriate moment along the given axis or over all values if axis + is None. The denominator for the moment calculation is the number of + observations, no degrees of freedom correction is done. + + Notes + ----- + For more details about `moment`, see `scipy.stats.moment`. + + """ + a, axis = _chk_asarray(a, axis) + if a.size == 0: + moment_shape = list(a.shape) + del moment_shape[axis] + dtype = a.dtype.type if a.dtype.kind in 'fc' else np.float64 + # empty array, return nan(s) with shape matching `moment` + out_shape = (moment_shape if np.isscalar(moment) + else [len(moment)] + moment_shape) + if len(out_shape) == 0: + return dtype(np.nan) + else: + return ma.array(np.full(out_shape, np.nan, dtype=dtype)) + + # for array_like moment input, return a value for each. + if not np.isscalar(moment): + mean = a.mean(axis, keepdims=True) + mmnt = [_moment(a, i, axis, mean=mean) for i in moment] + return ma.array(mmnt) + else: + return _moment(a, moment, axis) + + +# Moment with optional pre-computed mean, equal to a.mean(axis, keepdims=True) +def _moment(a, moment, axis, *, mean=None): + if np.abs(moment - np.round(moment)) > 0: + raise ValueError("All moment parameters must be integers") + + if moment == 0 or moment == 1: + # By definition the zeroth moment about the mean is 1, and the first + # moment is 0. + shape = list(a.shape) + del shape[axis] + dtype = a.dtype.type if a.dtype.kind in 'fc' else np.float64 + + if len(shape) == 0: + return dtype(1.0 if moment == 0 else 0.0) + else: + return (ma.ones(shape, dtype=dtype) if moment == 0 + else ma.zeros(shape, dtype=dtype)) + else: + # Exponentiation by squares: form exponent sequence + n_list = [moment] + current_n = moment + while current_n > 2: + if current_n % 2: + current_n = (current_n-1)/2 + else: + current_n /= 2 + n_list.append(current_n) + + # Starting point for exponentiation by squares + mean = a.mean(axis, keepdims=True) if mean is None else mean + a_zero_mean = a - mean + if n_list[-1] == 1: + s = a_zero_mean.copy() + else: + s = a_zero_mean**2 + + # Perform multiplications + for n in n_list[-2::-1]: + s = s**2 + if n % 2: + s *= a_zero_mean + return s.mean(axis) + + +def variation(a, axis=0, ddof=0): + """ + Compute the coefficient of variation. + + The coefficient of variation is the standard deviation divided by the + mean. This function is equivalent to:: + + np.std(x, axis=axis, ddof=ddof) / np.mean(x) + + The default for ``ddof`` is 0, but many definitions of the coefficient + of variation use the square root of the unbiased sample variance + for the sample standard deviation, which corresponds to ``ddof=1``. + + Parameters + ---------- + a : array_like + Input array. + axis : int or None, optional + Axis along which to calculate the coefficient of variation. Default + is 0. If None, compute over the whole array `a`. + ddof : int, optional + Delta degrees of freedom. Default is 0. + + Returns + ------- + variation : ndarray + The calculated variation along the requested axis. + + Notes + ----- + For more details about `variation`, see `scipy.stats.variation`. + + Examples + -------- + >>> import numpy as np + >>> from scipy.stats.mstats import variation + >>> a = np.array([2,8,4]) + >>> variation(a) + 0.5345224838248487 + >>> b = np.array([2,8,3,4]) + >>> c = np.ma.masked_array(b, mask=[0,0,1,0]) + >>> variation(c) + 0.5345224838248487 + + In the example above, it can be seen that this works the same as + `scipy.stats.variation` except 'stats.mstats.variation' ignores masked + array elements. + + """ + a, axis = _chk_asarray(a, axis) + return a.std(axis, ddof=ddof)/a.mean(axis) + + +def skew(a, axis=0, bias=True): + """ + Computes the skewness of a data set. + + Parameters + ---------- + a : ndarray + data + axis : int or None, optional + Axis along which skewness is calculated. Default is 0. + If None, compute over the whole array `a`. + bias : bool, optional + If False, then the calculations are corrected for statistical bias. + + Returns + ------- + skewness : ndarray + The skewness of values along an axis, returning 0 where all values are + equal. + + Notes + ----- + For more details about `skew`, see `scipy.stats.skew`. + + """ + a, axis = _chk_asarray(a,axis) + mean = a.mean(axis, keepdims=True) + m2 = _moment(a, 2, axis, mean=mean) + m3 = _moment(a, 3, axis, mean=mean) + zero = (m2 <= (np.finfo(m2.dtype).resolution * mean.squeeze(axis))**2) + with np.errstate(all='ignore'): + vals = ma.where(zero, 0, m3 / m2**1.5) + + if not bias and zero is not ma.masked and m2 is not ma.masked: + n = a.count(axis) + can_correct = ~zero & (n > 2) + if can_correct.any(): + n = np.extract(can_correct, n) + m2 = np.extract(can_correct, m2) + m3 = np.extract(can_correct, m3) + nval = ma.sqrt((n-1.0)*n)/(n-2.0)*m3/m2**1.5 + np.place(vals, can_correct, nval) + return vals + + +def kurtosis(a, axis=0, fisher=True, bias=True): + """ + Computes the kurtosis (Fisher or Pearson) of a dataset. + + Kurtosis is the fourth central moment divided by the square of the + variance. If Fisher's definition is used, then 3.0 is subtracted from + the result to give 0.0 for a normal distribution. + + If bias is False then the kurtosis is calculated using k statistics to + eliminate bias coming from biased moment estimators + + Use `kurtosistest` to see if result is close enough to normal. + + Parameters + ---------- + a : array + data for which the kurtosis is calculated + axis : int or None, optional + Axis along which the kurtosis is calculated. Default is 0. + If None, compute over the whole array `a`. + fisher : bool, optional + If True, Fisher's definition is used (normal ==> 0.0). If False, + Pearson's definition is used (normal ==> 3.0). + bias : bool, optional + If False, then the calculations are corrected for statistical bias. + + Returns + ------- + kurtosis : array + The kurtosis of values along an axis. If all values are equal, + return -3 for Fisher's definition and 0 for Pearson's definition. + + Notes + ----- + For more details about `kurtosis`, see `scipy.stats.kurtosis`. + + """ + a, axis = _chk_asarray(a, axis) + mean = a.mean(axis, keepdims=True) + m2 = _moment(a, 2, axis, mean=mean) + m4 = _moment(a, 4, axis, mean=mean) + zero = (m2 <= (np.finfo(m2.dtype).resolution * mean.squeeze(axis))**2) + with np.errstate(all='ignore'): + vals = ma.where(zero, 0, m4 / m2**2.0) + + if not bias and zero is not ma.masked and m2 is not ma.masked: + n = a.count(axis) + can_correct = ~zero & (n > 3) + if can_correct.any(): + n = np.extract(can_correct, n) + m2 = np.extract(can_correct, m2) + m4 = np.extract(can_correct, m4) + nval = 1.0/(n-2)/(n-3)*((n*n-1.0)*m4/m2**2.0-3*(n-1)**2.0) + np.place(vals, can_correct, nval+3.0) + if fisher: + return vals - 3 + else: + return vals + + +DescribeResult = namedtuple('DescribeResult', ('nobs', 'minmax', 'mean', + 'variance', 'skewness', + 'kurtosis')) + + +def describe(a, axis=0, ddof=0, bias=True): + """ + Computes several descriptive statistics of the passed array. + + Parameters + ---------- + a : array_like + Data array + axis : int or None, optional + Axis along which to calculate statistics. Default 0. If None, + compute over the whole array `a`. + ddof : int, optional + degree of freedom (default 0); note that default ddof is different + from the same routine in stats.describe + bias : bool, optional + If False, then the skewness and kurtosis calculations are corrected for + statistical bias. + + Returns + ------- + nobs : int + (size of the data (discarding missing values) + + minmax : (int, int) + min, max + + mean : float + arithmetic mean + + variance : float + unbiased variance + + skewness : float + biased skewness + + kurtosis : float + biased kurtosis + + Examples + -------- + >>> import numpy as np + >>> from scipy.stats.mstats import describe + >>> ma = np.ma.array(range(6), mask=[0, 0, 0, 1, 1, 1]) + >>> describe(ma) + DescribeResult(nobs=np.int64(3), minmax=(masked_array(data=0, + mask=False, + fill_value=999999), masked_array(data=2, + mask=False, + fill_value=999999)), mean=np.float64(1.0), + variance=np.float64(0.6666666666666666), + skewness=masked_array(data=0., mask=False, fill_value=1e+20), + kurtosis=np.float64(-1.5)) + + """ + a, axis = _chk_asarray(a, axis) + n = a.count(axis) + mm = (ma.minimum.reduce(a, axis=axis), ma.maximum.reduce(a, axis=axis)) + m = a.mean(axis) + v = a.var(axis, ddof=ddof) + sk = skew(a, axis, bias=bias) + kurt = kurtosis(a, axis, bias=bias) + + return DescribeResult(n, mm, m, v, sk, kurt) + + +def stde_median(data, axis=None): + """Returns the McKean-Schrader estimate of the standard error of the sample + median along the given axis. masked values are discarded. + + Parameters + ---------- + data : ndarray + Data to trim. + axis : {None,int}, optional + Axis along which to perform the trimming. + If None, the input array is first flattened. + + """ + def _stdemed_1D(data): + data = np.sort(data.compressed()) + n = len(data) + z = 2.5758293035489004 + k = int(np.round((n+1)/2. - z * np.sqrt(n/4.),0)) + return ((data[n-k] - data[k-1])/(2.*z)) + + data = ma.array(data, copy=False, subok=True) + if (axis is None): + return _stdemed_1D(data) + else: + if data.ndim > 2: + raise ValueError(f"Array 'data' must be at most two dimensional, " + f"but got data.ndim = {data.ndim}") + return ma.apply_along_axis(_stdemed_1D, axis, data) + + +SkewtestResult = namedtuple('SkewtestResult', ('statistic', 'pvalue')) + + +def skewtest(a, axis=0, alternative='two-sided'): + """ + Tests whether the skew is different from the normal distribution. + + Parameters + ---------- + a : array_like + The data to be tested + axis : int or None, optional + Axis along which statistics are calculated. Default is 0. + If None, compute over the whole array `a`. + alternative : {'two-sided', 'less', 'greater'}, optional + Defines the alternative hypothesis. Default is 'two-sided'. + The following options are available: + + * 'two-sided': the skewness of the distribution underlying the sample + is different from that of the normal distribution (i.e. 0) + * 'less': the skewness of the distribution underlying the sample + is less than that of the normal distribution + * 'greater': the skewness of the distribution underlying the sample + is greater than that of the normal distribution + + .. versionadded:: 1.7.0 + + Returns + ------- + statistic : array_like + The computed z-score for this test. + pvalue : array_like + A p-value for the hypothesis test + + Notes + ----- + For more details about `skewtest`, see `scipy.stats.skewtest`. + + """ + a, axis = _chk_asarray(a, axis) + if axis is None: + a = a.ravel() + axis = 0 + b2 = skew(a,axis) + n = a.count(axis) + if np.min(n) < 8: + raise ValueError(f"skewtest is not valid with less than 8 samples; " + f"{np.min(n)} samples were given.") + + y = b2 * ma.sqrt(((n+1)*(n+3)) / (6.0*(n-2))) + beta2 = (3.0*(n*n+27*n-70)*(n+1)*(n+3)) / ((n-2.0)*(n+5)*(n+7)*(n+9)) + W2 = -1 + ma.sqrt(2*(beta2-1)) + delta = 1/ma.sqrt(0.5*ma.log(W2)) + alpha = ma.sqrt(2.0/(W2-1)) + y = ma.where(y == 0, 1, y) + Z = delta*ma.log(y/alpha + ma.sqrt((y/alpha)**2+1)) + pvalue = scipy.stats._stats_py._get_pvalue(Z, distributions.norm, alternative) + + return SkewtestResult(Z[()], pvalue[()]) + + +KurtosistestResult = namedtuple('KurtosistestResult', ('statistic', 'pvalue')) + + +def kurtosistest(a, axis=0, alternative='two-sided'): + """ + Tests whether a dataset has normal kurtosis + + Parameters + ---------- + a : array_like + array of the sample data + axis : int or None, optional + Axis along which to compute test. Default is 0. If None, + compute over the whole array `a`. + alternative : {'two-sided', 'less', 'greater'}, optional + Defines the alternative hypothesis. + The following options are available (default is 'two-sided'): + + * 'two-sided': the kurtosis of the distribution underlying the sample + is different from that of the normal distribution + * 'less': the kurtosis of the distribution underlying the sample + is less than that of the normal distribution + * 'greater': the kurtosis of the distribution underlying the sample + is greater than that of the normal distribution + + .. versionadded:: 1.7.0 + + Returns + ------- + statistic : array_like + The computed z-score for this test. + pvalue : array_like + The p-value for the hypothesis test + + Notes + ----- + For more details about `kurtosistest`, see `scipy.stats.kurtosistest`. + + """ + a, axis = _chk_asarray(a, axis) + n = a.count(axis=axis) + if np.min(n) < 5: + raise ValueError(f"kurtosistest requires at least 5 observations; " + f"{np.min(n)} observations were given.") + if np.min(n) < 20: + warnings.warn(f"kurtosistest only valid for n>=20 ... continuing " + f"anyway, n={np.min(n)}", stacklevel=2) + + b2 = kurtosis(a, axis, fisher=False) + E = 3.0*(n-1) / (n+1) + varb2 = 24.0*n*(n-2.)*(n-3) / ((n+1)*(n+1.)*(n+3)*(n+5)) + x = (b2-E)/ma.sqrt(varb2) + sqrtbeta1 = 6.0*(n*n-5*n+2)/((n+7)*(n+9)) * np.sqrt((6.0*(n+3)*(n+5)) / + (n*(n-2)*(n-3))) + A = 6.0 + 8.0/sqrtbeta1 * (2.0/sqrtbeta1 + np.sqrt(1+4.0/(sqrtbeta1**2))) + term1 = 1 - 2./(9.0*A) + denom = 1 + x*ma.sqrt(2/(A-4.0)) + if np.ma.isMaskedArray(denom): + # For multi-dimensional array input + denom[denom == 0.0] = masked + elif denom == 0.0: + denom = masked + + term2 = np.ma.where(denom > 0, ma.power((1-2.0/A)/denom, 1/3.0), + -ma.power(-(1-2.0/A)/denom, 1/3.0)) + Z = (term1 - term2) / np.sqrt(2/(9.0*A)) + pvalue = scipy.stats._stats_py._get_pvalue(Z, distributions.norm, alternative) + + return KurtosistestResult(Z[()], pvalue[()]) + + +NormaltestResult = namedtuple('NormaltestResult', ('statistic', 'pvalue')) + + +def normaltest(a, axis=0): + """ + Tests whether a sample differs from a normal distribution. + + Parameters + ---------- + a : array_like + The array containing the data to be tested. + axis : int or None, optional + Axis along which to compute test. Default is 0. If None, + compute over the whole array `a`. + + Returns + ------- + statistic : float or array + ``s^2 + k^2``, where ``s`` is the z-score returned by `skewtest` and + ``k`` is the z-score returned by `kurtosistest`. + pvalue : float or array + A 2-sided chi squared probability for the hypothesis test. + + Notes + ----- + For more details about `normaltest`, see `scipy.stats.normaltest`. + + """ + a, axis = _chk_asarray(a, axis) + s, _ = skewtest(a, axis) + k, _ = kurtosistest(a, axis) + k2 = s*s + k*k + + return NormaltestResult(k2, distributions.chi2.sf(k2, 2)) + + +def mquantiles(a, prob=(.25, .5, .75), alphap=.4, betap=.4, axis=None, + limit=()): + """ + Computes empirical quantiles for a data array. + + Samples quantile are defined by ``Q(p) = (1-gamma)*x[j] + gamma*x[j+1]``, + where ``x[j]`` is the j-th order statistic, and gamma is a function of + ``j = floor(n*p + m)``, ``m = alphap + p*(1 - alphap - betap)`` and + ``g = n*p + m - j``. + + Reinterpreting the above equations to compare to **R** lead to the + equation: ``p(k) = (k - alphap)/(n + 1 - alphap - betap)`` + + Typical values of (alphap,betap) are: + + - (0,1) : ``p(k) = k/n`` : linear interpolation of cdf + (**R** type 4) + - (.5,.5) : ``p(k) = (k - 1/2.)/n`` : piecewise linear function + (**R** type 5) + - (0,0) : ``p(k) = k/(n+1)`` : + (**R** type 6) + - (1,1) : ``p(k) = (k-1)/(n-1)``: p(k) = mode[F(x[k])]. + (**R** type 7, **R** default) + - (1/3,1/3): ``p(k) = (k-1/3)/(n+1/3)``: Then p(k) ~ median[F(x[k])]. + The resulting quantile estimates are approximately median-unbiased + regardless of the distribution of x. + (**R** type 8) + - (3/8,3/8): ``p(k) = (k-3/8)/(n+1/4)``: Blom. + The resulting quantile estimates are approximately unbiased + if x is normally distributed + (**R** type 9) + - (.4,.4) : approximately quantile unbiased (Cunnane) + - (.35,.35): APL, used with PWM + + Parameters + ---------- + a : array_like + Input data, as a sequence or array of dimension at most 2. + prob : array_like, optional + List of quantiles to compute. + alphap : float, optional + Plotting positions parameter, default is 0.4. + betap : float, optional + Plotting positions parameter, default is 0.4. + axis : int, optional + Axis along which to perform the trimming. + If None (default), the input array is first flattened. + limit : tuple, optional + Tuple of (lower, upper) values. + Values of `a` outside this open interval are ignored. + + Returns + ------- + mquantiles : MaskedArray + An array containing the calculated quantiles. + + Notes + ----- + This formulation is very similar to **R** except the calculation of + ``m`` from ``alphap`` and ``betap``, where in **R** ``m`` is defined + with each type. + + References + ---------- + .. [1] *R* statistical software: https://www.r-project.org/ + .. [2] *R* ``quantile`` function: + http://stat.ethz.ch/R-manual/R-devel/library/stats/html/quantile.html + + Examples + -------- + >>> import numpy as np + >>> from scipy.stats.mstats import mquantiles + >>> a = np.array([6., 47., 49., 15., 42., 41., 7., 39., 43., 40., 36.]) + >>> mquantiles(a) + array([ 19.2, 40. , 42.8]) + + Using a 2D array, specifying axis and limit. + + >>> data = np.array([[ 6., 7., 1.], + ... [ 47., 15., 2.], + ... [ 49., 36., 3.], + ... [ 15., 39., 4.], + ... [ 42., 40., -999.], + ... [ 41., 41., -999.], + ... [ 7., -999., -999.], + ... [ 39., -999., -999.], + ... [ 43., -999., -999.], + ... [ 40., -999., -999.], + ... [ 36., -999., -999.]]) + >>> print(mquantiles(data, axis=0, limit=(0, 50))) + [[19.2 14.6 1.45] + [40. 37.5 2.5 ] + [42.8 40.05 3.55]] + + >>> data[:, 2] = -999. + >>> print(mquantiles(data, axis=0, limit=(0, 50))) + [[19.200000000000003 14.6 --] + [40.0 37.5 --] + [42.800000000000004 40.05 --]] + + """ + def _quantiles1D(data,m,p): + x = np.sort(data.compressed()) + n = len(x) + if n == 0: + return ma.array(np.empty(len(p), dtype=float), mask=True) + elif n == 1: + return ma.array(np.resize(x, p.shape), mask=nomask) + aleph = (n*p + m) + k = np.floor(aleph.clip(1, n-1)).astype(int) + gamma = (aleph-k).clip(0,1) + return (1.-gamma)*x[(k-1).tolist()] + gamma*x[k.tolist()] + + data = ma.array(a, copy=False) + if data.ndim > 2: + raise TypeError("Array should be 2D at most !") + + if limit: + condition = (limit[0] < data) & (data < limit[1]) + data[~condition.filled(True)] = masked + + p = np.atleast_1d(np.asarray(prob)) + m = alphap + p*(1.-alphap-betap) + # Computes quantiles along axis (or globally) + if (axis is None): + return _quantiles1D(data, m, p) + + return ma.apply_along_axis(_quantiles1D, axis, data, m, p) + + +def scoreatpercentile(data, per, limit=(), alphap=.4, betap=.4): + """Calculate the score at the given 'per' percentile of the + sequence a. For example, the score at per=50 is the median. + + This function is a shortcut to mquantile + + """ # numpydoc ignore=RT01 + if (per < 0) or (per > 100.): + raise ValueError(f"The percentile should be between 0. and 100. ! (got {per})") + + return mquantiles(data, prob=[per/100.], alphap=alphap, betap=betap, + limit=limit, axis=0).squeeze() + + +def plotting_positions(data, alpha=0.4, beta=0.4): + """ + Returns plotting positions (or empirical percentile points) for the data. + + Plotting positions are defined as ``(i-alpha)/(n+1-alpha-beta)``, where: + + - i is the rank order statistics + - n is the number of unmasked values along the given axis + - `alpha` and `beta` are two parameters. + + Typical values for `alpha` and `beta` are: + + - (0,1) : ``p(k) = k/n``, linear interpolation of cdf (R, type 4) + - (.5,.5) : ``p(k) = (k-1/2.)/n``, piecewise linear function + (R, type 5) + - (0,0) : ``p(k) = k/(n+1)``, Weibull (R type 6) + - (1,1) : ``p(k) = (k-1)/(n-1)``, in this case, + ``p(k) = mode[F(x[k])]``. That's R default (R type 7) + - (1/3,1/3): ``p(k) = (k-1/3)/(n+1/3)``, then + ``p(k) ~ median[F(x[k])]``. + The resulting quantile estimates are approximately median-unbiased + regardless of the distribution of x. (R type 8) + - (3/8,3/8): ``p(k) = (k-3/8)/(n+1/4)``, Blom. + The resulting quantile estimates are approximately unbiased + if x is normally distributed (R type 9) + - (.4,.4) : approximately quantile unbiased (Cunnane) + - (.35,.35): APL, used with PWM + - (.3175, .3175): used in scipy.stats.probplot + + Parameters + ---------- + data : array_like + Input data, as a sequence or array of dimension at most 2. + alpha : float, optional + Plotting positions parameter. Default is 0.4. + beta : float, optional + Plotting positions parameter. Default is 0.4. + + Returns + ------- + positions : MaskedArray + The calculated plotting positions. + + """ + data = ma.array(data, copy=False).reshape(1,-1) + n = data.count() + plpos = np.empty(data.size, dtype=float) + plpos[n:] = 0 + plpos[data.argsort(axis=None)[:n]] = ((np.arange(1, n+1) - alpha) / + (n + 1.0 - alpha - beta)) + return ma.array(plpos, mask=data._mask) + + +meppf = plotting_positions + + +def obrientransform(*args): + """ + Computes a transform on input data (any number of columns). Used to + test for homogeneity of variance prior to running one-way stats. Each + array in ``*args`` is one level of a factor. If an `f_oneway()` run on + the transformed data and found significant, variances are unequal. From + Maxwell and Delaney, p.112. + + Returns: transformed data for use in an ANOVA + """ # numpydoc ignore=RT01 + data = argstoarray(*args).T + v = data.var(axis=0,ddof=1) + m = data.mean(0) + n = data.count(0).astype(float) + # result = ((N-1.5)*N*(a-m)**2 - 0.5*v*(n-1))/((n-1)*(n-2)) + data -= m + data **= 2 + data *= (n-1.5)*n + data -= 0.5*v*(n-1) + data /= (n-1.)*(n-2.) + if not ma.allclose(v,data.mean(0)): + raise ValueError("Lack of convergence in obrientransform.") + + return data + + +def sem(a, axis=0, ddof=1): + """ + Calculates the standard error of the mean of the input array. + + Also sometimes called standard error of measurement. + + Parameters + ---------- + a : array_like + An array containing the values for which the standard error is + returned. + axis : int or None, optional + If axis is None, ravel `a` first. If axis is an integer, this will be + the axis over which to operate. Defaults to 0. + ddof : int, optional + Delta degrees-of-freedom. How many degrees of freedom to adjust + for bias in limited samples relative to the population estimate + of variance. Defaults to 1. + + Returns + ------- + s : ndarray or float + The standard error of the mean in the sample(s), along the input axis. + + Notes + ----- + The default value for `ddof` changed in scipy 0.15.0 to be consistent with + `scipy.stats.sem` as well as with the most common definition used (like in + the R documentation). + + Examples + -------- + Find standard error along the first axis: + + >>> import numpy as np + >>> from scipy import stats + >>> a = np.arange(20).reshape(5,4) + >>> print(stats.mstats.sem(a)) + [2.8284271247461903 2.8284271247461903 2.8284271247461903 + 2.8284271247461903] + + Find standard error across the whole array, using n degrees of freedom: + + >>> print(stats.mstats.sem(a, axis=None, ddof=0)) + 1.2893796958227628 + + """ + a, axis = _chk_asarray(a, axis) + n = a.count(axis=axis) + s = a.std(axis=axis, ddof=ddof) / ma.sqrt(n) + return s + + +F_onewayResult = namedtuple('F_onewayResult', ('statistic', 'pvalue')) + + +def f_oneway(*args): + """ + Performs a 1-way ANOVA, returning an F-value and probability given + any number of groups. From Heiman, pp.394-7. + + Usage: ``f_oneway(*args)``, where ``*args`` is 2 or more arrays, + one per treatment group. + + Returns + ------- + statistic : float + The computed F-value of the test. + pvalue : float + The associated p-value from the F-distribution. + + """ + # Construct a single array of arguments: each row is a group + data = argstoarray(*args) + ngroups = len(data) + ntot = data.count() + sstot = (data**2).sum() - (data.sum())**2/float(ntot) + ssbg = (data.count(-1) * (data.mean(-1)-data.mean())**2).sum() + sswg = sstot-ssbg + dfbg = ngroups-1 + dfwg = ntot - ngroups + msb = ssbg/float(dfbg) + msw = sswg/float(dfwg) + f = msb/msw + prob = special.fdtrc(dfbg, dfwg, f) # equivalent to stats.f.sf + + return F_onewayResult(f, prob) + + +FriedmanchisquareResult = namedtuple('FriedmanchisquareResult', + ('statistic', 'pvalue')) + + +def friedmanchisquare(*args): + """Friedman Chi-Square is a non-parametric, one-way within-subjects ANOVA. + This function calculates the Friedman Chi-square test for repeated measures + and returns the result, along with the associated probability value. + + Each input is considered a given group. Ideally, the number of treatments + among each group should be equal. If this is not the case, only the first + n treatments are taken into account, where n is the number of treatments + of the smallest group. + If a group has some missing values, the corresponding treatments are masked + in the other groups. + The test statistic is corrected for ties. + + Masked values in one group are propagated to the other groups. + + Returns + ------- + statistic : float + the test statistic. + pvalue : float + the associated p-value. + + """ + data = argstoarray(*args).astype(float) + k = len(data) + if k < 3: + raise ValueError(f"Less than 3 groups ({k}): the Friedman test " + f"is NOT appropriate.") + + ranked = ma.masked_values(rankdata(data, axis=0), 0) + if ranked._mask is not nomask: + ranked = ma.mask_cols(ranked) + ranked = ranked.compressed().reshape(k,-1).view(ndarray) + else: + ranked = ranked._data + (k,n) = ranked.shape + # Ties correction + repeats = [find_repeats(row) for row in ranked.T] + ties = np.array([y for x, y in repeats if x.size > 0]) + tie_correction = 1 - (ties**3-ties).sum()/float(n*(k**3-k)) + + ssbg = np.sum((ranked.sum(-1) - n*(k+1)/2.)**2) + chisq = ssbg * 12./(n*k*(k+1)) * 1./tie_correction + + return FriedmanchisquareResult(chisq, + distributions.chi2.sf(chisq, k-1)) + + +BrunnerMunzelResult = namedtuple('BrunnerMunzelResult', ('statistic', 'pvalue')) + + +def brunnermunzel(x, y, alternative="two-sided", distribution="t"): + """ + Compute the Brunner-Munzel test on samples x and y. + + Any missing values in `x` and/or `y` are discarded. + + The Brunner-Munzel test is a nonparametric test of the null hypothesis that + when values are taken one by one from each group, the probabilities of + getting large values in both groups are equal. + Unlike the Wilcoxon-Mann-Whitney's U test, this does not require the + assumption of equivariance of two groups. Note that this does not assume + the distributions are same. This test works on two independent samples, + which may have different sizes. + + Parameters + ---------- + x, y : array_like + Array of samples, should be one-dimensional. + alternative : 'less', 'two-sided', or 'greater', optional + Whether to get the p-value for the one-sided hypothesis ('less' + or 'greater') or for the two-sided hypothesis ('two-sided'). + Defaults value is 'two-sided' . + distribution : 't' or 'normal', optional + Whether to get the p-value by t-distribution or by standard normal + distribution. + Defaults value is 't' . + + Returns + ------- + statistic : float + The Brunner-Munzer W statistic. + pvalue : float + p-value assuming a t distribution. One-sided or + two-sided, depending on the choice of `alternative` and `distribution`. + + See Also + -------- + mannwhitneyu : Mann-Whitney rank test on two samples. + + Notes + ----- + For more details on `brunnermunzel`, see `scipy.stats.brunnermunzel`. + + Examples + -------- + >>> from scipy.stats.mstats import brunnermunzel + >>> import numpy as np + >>> x1 = [1, 2, np.nan, np.nan, 1, 1, 1, 1, 1, 1, 2, 4, 1, 1] + >>> x2 = [3, 3, 4, 3, 1, 2, 3, 1, 1, 5, 4] + >>> brunnermunzel(x1, x2) + BrunnerMunzelResult(statistic=1.4723186918922935, pvalue=0.15479415300426624) # may vary + + """ # noqa: E501 + x = ma.asarray(x).compressed().view(ndarray) + y = ma.asarray(y).compressed().view(ndarray) + nx = len(x) + ny = len(y) + if nx == 0 or ny == 0: + return BrunnerMunzelResult(np.nan, np.nan) + rankc = rankdata(np.concatenate((x,y))) + rankcx = rankc[0:nx] + rankcy = rankc[nx:nx+ny] + rankcx_mean = np.mean(rankcx) + rankcy_mean = np.mean(rankcy) + rankx = rankdata(x) + ranky = rankdata(y) + rankx_mean = np.mean(rankx) + ranky_mean = np.mean(ranky) + + Sx = np.sum(np.power(rankcx - rankx - rankcx_mean + rankx_mean, 2.0)) + Sx /= nx - 1 + Sy = np.sum(np.power(rankcy - ranky - rankcy_mean + ranky_mean, 2.0)) + Sy /= ny - 1 + + wbfn = nx * ny * (rankcy_mean - rankcx_mean) + wbfn /= (nx + ny) * np.sqrt(nx * Sx + ny * Sy) + + if distribution == "t": + df_numer = np.power(nx * Sx + ny * Sy, 2.0) + df_denom = np.power(nx * Sx, 2.0) / (nx - 1) + df_denom += np.power(ny * Sy, 2.0) / (ny - 1) + df = df_numer / df_denom + p = distributions.t.cdf(wbfn, df) + elif distribution == "normal": + p = distributions.norm.cdf(wbfn) + else: + raise ValueError( + "distribution should be 't' or 'normal'") + + if alternative == "greater": + pass + elif alternative == "less": + p = 1 - p + elif alternative == "two-sided": + p = 2 * np.min([p, 1-p]) + else: + raise ValueError( + "alternative should be 'less', 'greater' or 'two-sided'") + + return BrunnerMunzelResult(wbfn, p) diff --git a/v12x/pylib/scipy/stats/_mstats_extras.py b/v12x/pylib/scipy/stats/_mstats_extras.py new file mode 100644 index 0000000000000000000000000000000000000000..d7b5484fb930edadb0538576e0cb736c82957e7a --- /dev/null +++ b/v12x/pylib/scipy/stats/_mstats_extras.py @@ -0,0 +1,520 @@ +""" +Additional statistics functions with support for masked arrays. + +""" + +# Original author (2007): Pierre GF Gerard-Marchant + + +__all__ = ['compare_medians_ms', + 'hdquantiles', 'hdmedian', 'hdquantiles_sd', + 'idealfourths', + 'median_cihs','mjci','mquantiles_cimj', + 'rsh', + 'trimmed_mean_ci',] + + +import numpy as np +from numpy import float64, ndarray + +import numpy.ma as ma +from numpy.ma import MaskedArray + +from . import _mstats_basic as mstats + +from scipy.stats.distributions import norm, beta, t, binom + + +def hdquantiles(data, prob=(.25, .5, .75), axis=None, var=False,): + """ + Computes quantile estimates with the Harrell-Davis method. + + The quantile estimates are calculated as a weighted linear combination + of order statistics. + + Parameters + ---------- + data : array_like + Data array. + prob : sequence, optional + Sequence of probabilities at which to compute the quantiles. + axis : int or None, optional + Axis along which to compute the quantiles. If None, use a flattened + array. + var : bool, optional + Whether to return the variance of the estimate. + + Returns + ------- + hdquantiles : MaskedArray + A (p,) array of quantiles (if `var` is False), or a (2,p) array of + quantiles and variances (if `var` is True), where ``p`` is the + number of quantiles. + + See Also + -------- + hdquantiles_sd + + Examples + -------- + >>> import numpy as np + >>> from scipy.stats.mstats import hdquantiles + >>> + >>> # Sample data + >>> data = np.array([1.2, 2.5, 3.7, 4.0, 5.1, 6.3, 7.0, 8.2, 9.4]) + >>> + >>> # Probabilities at which to compute quantiles + >>> probabilities = [0.25, 0.5, 0.75] + >>> + >>> # Compute Harrell-Davis quantile estimates + >>> quantile_estimates = hdquantiles(data, prob=probabilities) + >>> + >>> # Display the quantile estimates + >>> for i, quantile in enumerate(probabilities): + ... print(f"{int(quantile * 100)}th percentile: {quantile_estimates[i]}") + 25th percentile: 3.1505820231763066 # may vary + 50th percentile: 5.194344084883956 + 75th percentile: 7.430626414674935 + + """ + def _hd_1D(data,prob,var): + "Computes the HD quantiles for a 1D array. Returns nan for invalid data." + xsorted = np.squeeze(np.sort(data.compressed().view(ndarray))) + # Don't use length here, in case we have a numpy scalar + n = xsorted.size + + hd = np.empty((2,len(prob)), float64) + if n < 2: + hd.flat = np.nan + if var: + return hd + return hd[0] + + v = np.arange(n+1) / float(n) + betacdf = beta.cdf + for (i,p) in enumerate(prob): + _w = betacdf(v, (n+1)*p, (n+1)*(1-p)) + w = _w[1:] - _w[:-1] + hd_mean = np.dot(w, xsorted) + hd[0,i] = hd_mean + # + hd[1,i] = np.dot(w, (xsorted-hd_mean)**2) + # + hd[0, prob == 0] = xsorted[0] + hd[0, prob == 1] = xsorted[-1] + if var: + hd[1, prob == 0] = hd[1, prob == 1] = np.nan + return hd + return hd[0] + # Initialization & checks + data = ma.array(data, copy=False, dtype=float64) + p = np.atleast_1d(np.asarray(prob)) + # Computes quantiles along axis (or globally) + if (axis is None) or (data.ndim == 1): + result = _hd_1D(data, p, var) + else: + if data.ndim > 2: + raise ValueError(f"Array 'data' must be at most two dimensional, " + f"but got data.ndim = {data.ndim}") + result = ma.apply_along_axis(_hd_1D, axis, data, p, var) + + return ma.fix_invalid(result, copy=False) + + +def hdmedian(data, axis=-1, var=False): + """ + Returns the Harrell-Davis estimate of the median along the given axis. + + Parameters + ---------- + data : ndarray + Data array. + axis : int, optional + Axis along which to compute the quantiles. If None, use a flattened + array. + var : bool, optional + Whether to return the variance of the estimate. + + Returns + ------- + hdmedian : MaskedArray + The median values. If ``var=True``, the variance is returned inside + the masked array. E.g. for a 1-D array the shape change from (1,) to + (2,). + + """ + result = hdquantiles(data,[0.5], axis=axis, var=var) + return result.squeeze() + + +def hdquantiles_sd(data, prob=(.25, .5, .75), axis=None): + """ + The standard error of the Harrell-Davis quantile estimates by jackknife. + + Parameters + ---------- + data : array_like + Data array. + prob : sequence, optional + Sequence of quantiles to compute. + axis : int, optional + Axis along which to compute the quantiles. If None, use a flattened + array. + + Returns + ------- + hdquantiles_sd : MaskedArray + Standard error of the Harrell-Davis quantile estimates. + + See Also + -------- + hdquantiles + + """ + def _hdsd_1D(data, prob): + "Computes the std error for 1D arrays." + xsorted = np.sort(data.compressed()) + n = len(xsorted) + + hdsd = np.empty(len(prob), float64) + if n < 2: + hdsd.flat = np.nan + + vv = np.arange(n) / float(n-1) + betacdf = beta.cdf + + for (i,p) in enumerate(prob): + _w = betacdf(vv, n*p, n*(1-p)) + w = _w[1:] - _w[:-1] + # cumulative sum of weights and data points if + # ith point is left out for jackknife + mx_ = np.zeros_like(xsorted) + mx_[1:] = np.cumsum(w * xsorted[:-1]) + # similar but from the right + mx_[:-1] += np.cumsum(w[::-1] * xsorted[:0:-1])[::-1] + hdsd[i] = np.sqrt(mx_.var() * (n - 1)) + return hdsd + + # Initialization & checks + data = ma.array(data, copy=False, dtype=float64) + p = np.atleast_1d(np.asarray(prob)) + # Computes quantiles along axis (or globally) + if (axis is None): + result = _hdsd_1D(data, p) + else: + if data.ndim > 2: + raise ValueError(f"Array 'data' must be at most two dimensional, " + f"but got data.ndim = {data.ndim}") + result = ma.apply_along_axis(_hdsd_1D, axis, data, p) + + return ma.fix_invalid(result, copy=False).ravel() + + +def trimmed_mean_ci(data, limits=(0.2,0.2), inclusive=(True,True), + alpha=0.05, axis=None): + """ + Selected confidence interval of the trimmed mean along the given axis. + + Parameters + ---------- + data : array_like + Input data. + limits : {None, tuple}, optional + None or a two item tuple. + Tuple of the percentages to cut on each side of the array, with respect + to the number of unmasked data, as floats between 0. and 1. If ``n`` + is the number of unmasked data before trimming, then + (``n * limits[0]``)th smallest data and (``n * limits[1]``)th + largest data are masked. The total number of unmasked data after + trimming is ``n * (1. - sum(limits))``. + The value of one limit can be set to None to indicate an open interval. + + Defaults to (0.2, 0.2). + inclusive : (2,) tuple of boolean, optional + If relative==False, tuple indicating whether values exactly equal to + the absolute limits are allowed. + If relative==True, tuple indicating whether the number of data being + masked on each side should be rounded (True) or truncated (False). + + Defaults to (True, True). + alpha : float, optional + Confidence level of the intervals. + + Defaults to 0.05. + axis : int, optional + Axis along which to cut. If None, uses a flattened version of `data`. + + Defaults to None. + + Returns + ------- + trimmed_mean_ci : (2,) ndarray + The lower and upper confidence intervals of the trimmed data. + + """ + data = ma.array(data, copy=False) + trimmed = mstats.trimr(data, limits=limits, inclusive=inclusive, axis=axis) + tmean = trimmed.mean(axis) + tstde = mstats.trimmed_stde(data,limits=limits,inclusive=inclusive,axis=axis) + df = trimmed.count(axis) - 1 + tppf = t.ppf(1-alpha/2.,df) + return np.array((tmean - tppf*tstde, tmean+tppf*tstde)) + + +def mjci(data, prob=(0.25, 0.5, 0.75), axis=None): + """ + Returns the Maritz-Jarrett estimators of the standard error of selected + experimental quantiles of the data. + + Parameters + ---------- + data : ndarray + Data array. + prob : sequence, optional + Sequence of quantiles to compute. + axis : int or None, optional + Axis along which to compute the quantiles. If None, use a flattened + array. + + """ # numpydoc ignore=RT01 + def _mjci_1D(data, p): + data = np.sort(data.compressed()) + n = data.size + prob = (np.array(p) * n + 0.5).astype(int) + betacdf = beta.cdf + + mj = np.empty(len(prob), float64) + x = np.arange(1,n+1, dtype=float64) / n + y = x - 1./n + for (i,m) in enumerate(prob): + W = betacdf(x,m-1,n-m) - betacdf(y,m-1,n-m) + C1 = np.dot(W,data) + C2 = np.dot(W,data**2) + mj[i] = np.sqrt(C2 - C1**2) + return mj + + data = ma.array(data, copy=False) + if data.ndim > 2: + raise ValueError(f"Array 'data' must be at most two dimensional, " + f"but got data.ndim = {data.ndim}") + + p = np.atleast_1d(np.asarray(prob)) + # Computes quantiles along axis (or globally) + if (axis is None): + return _mjci_1D(data, p) + else: + return ma.apply_along_axis(_mjci_1D, axis, data, p) + + +def mquantiles_cimj(data, prob=(0.25, 0.50, 0.75), alpha=0.05, axis=None): + """ + Computes the alpha confidence interval for the selected quantiles of the + data, with Maritz-Jarrett estimators. + + Parameters + ---------- + data : ndarray + Data array. + prob : sequence, optional + Sequence of quantiles to compute. + alpha : float, optional + Confidence level of the intervals. + axis : int or None, optional + Axis along which to compute the quantiles. + If None, use a flattened array. + + Returns + ------- + ci_lower : ndarray + The lower boundaries of the confidence interval. Of the same length as + `prob`. + ci_upper : ndarray + The upper boundaries of the confidence interval. Of the same length as + `prob`. + + """ + alpha = min(alpha, 1 - alpha) + z = norm.ppf(1 - alpha/2.) + xq = mstats.mquantiles(data, prob, alphap=0, betap=0, axis=axis) + smj = mjci(data, prob, axis=axis) + return (xq - z * smj, xq + z * smj) + + +def median_cihs(data, alpha=0.05, axis=None): + """ + Computes the alpha-level confidence interval for the median of the data. + + Uses the Hettmasperger-Sheather method. + + Parameters + ---------- + data : array_like + Input data. Masked values are discarded. The input should be 1D only, + or `axis` should be set to None. + alpha : float, optional + Confidence level of the intervals. + axis : int or None, optional + Axis along which to compute the quantiles. If None, use a flattened + array. + + Returns + ------- + median_cihs + Alpha level confidence interval. + + """ + def _cihs_1D(data, alpha): + data = np.sort(data.compressed()) + n = len(data) + alpha = min(alpha, 1-alpha) + k = int(binom._ppf(alpha/2., n, 0.5)) + gk = binom.cdf(n-k,n,0.5) - binom.cdf(k-1,n,0.5) + if gk < 1-alpha: + k -= 1 + gk = binom.cdf(n-k,n,0.5) - binom.cdf(k-1,n,0.5) + gkk = binom.cdf(n-k-1,n,0.5) - binom.cdf(k,n,0.5) + I = (gk - 1 + alpha)/(gk - gkk) + lambd = (n-k) * I / float(k + (n-2*k)*I) + lims = (lambd*data[k] + (1-lambd)*data[k-1], + lambd*data[n-k-1] + (1-lambd)*data[n-k]) + return lims + data = ma.array(data, copy=False) + # Computes quantiles along axis (or globally) + if (axis is None): + result = _cihs_1D(data, alpha) + else: + if data.ndim > 2: + raise ValueError(f"Array 'data' must be at most two dimensional, " + f"but got data.ndim = {data.ndim}") + result = ma.apply_along_axis(_cihs_1D, axis, data, alpha) + + return result + + +def compare_medians_ms(group_1, group_2, axis=None): + """ + Compares the medians from two independent groups along the given axis. + + The comparison is performed using the McKean-Schrader estimate of the + standard error of the medians. + + Parameters + ---------- + group_1 : array_like + First dataset. Has to be of size >=7. + group_2 : array_like + Second dataset. Has to be of size >=7. + axis : int, optional + Axis along which the medians are estimated. If None, the arrays are + flattened. If `axis` is not None, then `group_1` and `group_2` + should have the same shape. + + Returns + ------- + compare_medians_ms : {float, ndarray} + If `axis` is None, then returns a float, otherwise returns a 1-D + ndarray of floats with a length equal to the length of `group_1` + along `axis`. + + References + ---------- + .. [1] McKean, Joseph W., and Ronald M. Schrader. "A comparison of methods + for studentizing the sample median." Communications in + Statistics-Simulation and Computation 13.6 (1984): 751-773. + + Examples + -------- + + >>> from scipy import stats + >>> a = [1, 2, 3, 4, 5, 6, 7] + >>> b = [8, 9, 10, 11, 12, 13, 14] + >>> stats.mstats.compare_medians_ms(a, b, axis=None) + 1.0693225866553746e-05 + + The function is vectorized to compute along a given axis. + + >>> import numpy as np + >>> rng = np.random.default_rng() + >>> x = rng.random(size=(3, 7)) + >>> y = rng.random(size=(3, 8)) + >>> stats.mstats.compare_medians_ms(x, y, axis=1) + array([0.36908985, 0.36092538, 0.2765313 ]) + """ + (med_1, med_2) = (ma.median(group_1,axis=axis), ma.median(group_2,axis=axis)) + (std_1, std_2) = (mstats.stde_median(group_1, axis=axis), + mstats.stde_median(group_2, axis=axis)) + W = np.abs(med_1 - med_2) / ma.sqrt(std_1**2 + std_2**2) + return 1 - norm.cdf(W) + + +def idealfourths(data, axis=None): + """ + Returns an estimate of the lower and upper quartiles. + + Uses the ideal fourths algorithm. + + Parameters + ---------- + data : array_like + Input array. + axis : int, optional + Axis along which the quartiles are estimated. If None, the arrays are + flattened. + + Returns + ------- + idealfourths : {list of floats, masked array} + Returns the two internal values that divide `data` into four parts + using the ideal fourths algorithm either along the flattened array + (if `axis` is None) or along `axis` of `data`. + + """ + def _idf(data): + x = data.compressed() + n = len(x) + if n < 3: + return [np.nan,np.nan] + (j,h) = divmod(n/4. + 5/12.,1) + j = int(j) + qlo = (1-h)*x[j-1] + h*x[j] + k = n - j + qup = (1-h)*x[k] + h*x[k-1] + return [qlo, qup] + data = ma.sort(data, axis=axis).view(MaskedArray) + if (axis is None): + return _idf(data) + else: + return ma.apply_along_axis(_idf, axis, data) + + +def rsh(data, points=None): + """ + Evaluates Rosenblatt's shifted histogram estimators for each data point. + + Rosenblatt's estimator is a centered finite-difference approximation to the + derivative of the empirical cumulative distribution function. + + Parameters + ---------- + data : sequence + Input data, should be 1-D. Masked values are ignored. + points : sequence or None, optional + Sequence of points where to evaluate Rosenblatt shifted histogram. + If None, use the data. + + """ # numpydoc ignore=RT01 + data = ma.array(data, copy=False) + if points is None: + points = data + else: + points = np.atleast_1d(np.asarray(points)) + + if data.ndim != 1: + raise AttributeError("The input array should be 1D only !") + + n = data.count() + r = idealfourths(data, axis=None) + h = 1.2 * (r[-1]-r[0]) / n**(1./5) + nhi = (data[:,None] <= points[None,:] + h).sum(0) + nlo = (data[:,None] < points[None,:] - h).sum(0) + return (nhi-nlo) / (2.*n*h) diff --git a/v12x/pylib/scipy/stats/_multicomp.py b/v12x/pylib/scipy/stats/_multicomp.py new file mode 100644 index 0000000000000000000000000000000000000000..f7251be0912025e81afee9169e781e9b16ff5ce7 --- /dev/null +++ b/v12x/pylib/scipy/stats/_multicomp.py @@ -0,0 +1,457 @@ +import warnings +from collections.abc import Sequence +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Literal + +import numpy as np + +from scipy import stats +from scipy.optimize import minimize_scalar +from scipy.stats._common import ConfidenceInterval +from scipy.stats._qmc import check_random_state +from scipy.stats._stats_py import _var +from scipy._lib._array_api import xp_capabilities +from scipy._lib._util import _transition_to_rng, DecimalNumber, SeedType + + +if TYPE_CHECKING: + import numpy.typing as npt + + +__all__ = [ + 'dunnett' +] + + +@dataclass +class DunnettResult: + """Result object returned by `scipy.stats.dunnett`. + + Attributes + ---------- + statistic : float ndarray + The computed statistic of the test for each comparison. The element + at index ``i`` is the statistic for the comparison between + groups ``i`` and the control. + pvalue : float ndarray + The computed p-value of the test for each comparison. The element + at index ``i`` is the p-value for the comparison between + group ``i`` and the control. + """ + statistic: np.ndarray + pvalue: np.ndarray + _alternative: Literal['two-sided', 'less', 'greater'] = field(repr=False) + _rho: np.ndarray = field(repr=False) + _df: int = field(repr=False) + _std: float = field(repr=False) + _mean_samples: np.ndarray = field(repr=False) + _mean_control: np.ndarray = field(repr=False) + _n_samples: np.ndarray = field(repr=False) + _n_control: int = field(repr=False) + _rng: SeedType = field(repr=False) + _ci: ConfidenceInterval | None = field(default=None, repr=False) + _ci_cl: DecimalNumber | None = field(default=None, repr=False) + + @classmethod + def __class_getitem__(cls, arg, /): + # generic type compatibility with scipy-stubs + from types import GenericAlias + return GenericAlias(cls, arg) + + def __str__(self): + # Note: `__str__` prints the confidence intervals from the most + # recent call to `confidence_interval`. If it has not been called, + # it will be called with the default CL of .95. + if self._ci is None: + self.confidence_interval(confidence_level=.95) + s = ( + "Dunnett's test" + f" ({self._ci_cl*100:.1f}% Confidence Interval)\n" + "Comparison Statistic p-value Lower CI Upper CI\n" + ) + for i in range(self.pvalue.size): + s += (f" (Sample {i} - Control) {self.statistic[i]:>10.3f}" + f"{self.pvalue[i]:>10.3f}" + f"{self._ci.low[i]:>10.3f}" + f"{self._ci.high[i]:>10.3f}\n") + + return s + + def _allowance( + self, confidence_level: DecimalNumber = 0.95, tol: DecimalNumber = 1e-3 + ) -> float: + """Allowance. + + It is the quantity to add/subtract from the observed difference + between the means of observed groups and the mean of the control + group. The result gives confidence limits. + + Parameters + ---------- + confidence_level : float, optional + Confidence level for the computed confidence interval. + Default is .95. + tol : float, optional + A tolerance for numerical optimization: the allowance will produce + a confidence within ``10*tol*(1 - confidence_level)`` of the + specified level, or a warning will be emitted. Tight tolerances + may be impractical due to noisy evaluation of the objective. + Default is 1e-3. + + Returns + ------- + allowance : float + Allowance around the mean. + """ + alpha = 1 - confidence_level + + def pvalue_from_stat(statistic): + statistic = np.array(statistic) + sf = _pvalue_dunnett( + rho=self._rho, df=self._df, + statistic=statistic, alternative=self._alternative, + rng=self._rng + ) + return abs(sf - alpha)/alpha + + # Evaluation of `pvalue_from_stat` is noisy due to the use of RQMC to + # evaluate `multivariate_t.cdf`. `minimize_scalar` is not designed + # to tolerate a noisy objective function and may fail to find the + # minimum accurately. We mitigate this possibility with the validation + # step below, but implementation of a noise-tolerant root finder or + # minimizer would be a welcome enhancement. See gh-18150. + res = minimize_scalar(pvalue_from_stat, method='brent', tol=tol) + critical_value = res.x + + # validation + # tol*10 because tol=1e-3 means we tolerate a 1% change at most + if res.success is False or res.fun >= tol*10: + warnings.warn( + "Computation of the confidence interval did not converge to " + "the desired level. The confidence level corresponding with " + f"the returned interval is approximately {alpha*(1+res.fun)}.", + stacklevel=3 + ) + + # From [1] p. 1101 between (1) and (3) + allowance = critical_value*self._std*np.sqrt( + 1/self._n_samples + 1/self._n_control + ) + return abs(allowance) + + def confidence_interval( + self, confidence_level: DecimalNumber = 0.95 + ) -> ConfidenceInterval: + """Compute the confidence interval for the specified confidence level. + + Parameters + ---------- + confidence_level : float, optional + Confidence level for the computed confidence interval. + Default is .95. + + Returns + ------- + ci : ``ConfidenceInterval`` object + The object has attributes ``low`` and ``high`` that hold the + lower and upper bounds of the confidence intervals for each + comparison. The high and low values are accessible for each + comparison at index ``i`` for each group ``i``. + + """ + # check to see if the supplied confidence level matches that of the + # previously computed CI. + if (self._ci is not None) and (confidence_level == self._ci_cl): + return self._ci + + if not (0 < confidence_level < 1): + raise ValueError("Confidence level must be between 0 and 1.") + + allowance = self._allowance(confidence_level=confidence_level) + diff_means = self._mean_samples - self._mean_control + + low = diff_means-allowance + high = diff_means+allowance + + if self._alternative == 'greater': + high = [np.inf] * len(diff_means) + elif self._alternative == 'less': + low = [-np.inf] * len(diff_means) + + self._ci_cl = confidence_level + self._ci = ConfidenceInterval( + low=low, + high=high + ) + return self._ci + + +@xp_capabilities(np_only=True) +@_transition_to_rng('random_state', replace_doc=False) +def dunnett( + *samples: "npt.ArrayLike", # noqa: D417 + control: "npt.ArrayLike", + alternative: Literal['two-sided', 'less', 'greater'] = "two-sided", + rng: SeedType = None +) -> DunnettResult: + """Dunnett's test: multiple comparisons of means against a control group. + + This is an implementation of Dunnett's original, single-step test as + described in [1]_. + + Parameters + ---------- + *samples : 1D array_like + The sample measurements for each experimental group. + control : 1D array_like + The sample measurements for the control group. + alternative : {'two-sided', 'less', 'greater'}, optional + Defines the alternative hypothesis. + + The null hypothesis is that the means of the distributions underlying + the samples and control are equal. The following alternative + hypotheses are available (default is 'two-sided'): + + * 'two-sided': the means of the distributions underlying the samples + and control are unequal. + * 'less': the means of the distributions underlying the samples + are less than the mean of the distribution underlying the control. + * 'greater': the means of the distributions underlying the + samples are greater than the mean of the distribution underlying + the control. + rng : `numpy.random.Generator`, optional + Pseudorandom number generator state. When `rng` is None, a new + `numpy.random.Generator` is created using entropy from the + operating system. Types other than `numpy.random.Generator` are + passed to `numpy.random.default_rng` to instantiate a ``Generator``. + + .. versionchanged:: 1.15.0 + + As part of the `SPEC-007 `_ + transition from use of `numpy.random.RandomState` to + `numpy.random.Generator`, this keyword was changed from `random_state` to + `rng`. For an interim period, both keywords will continue to work, although + only one may be specified at a time. After the interim period, function + calls using the `random_state` keyword will emit warnings. Following a + deprecation period, the `random_state` keyword will be removed. + + Returns + ------- + res : `~scipy.stats._result_classes.DunnettResult` + An object containing attributes: + + statistic : float ndarray + The computed statistic of the test for each comparison. The element + at index ``i`` is the statistic for the comparison between + groups ``i`` and the control. + pvalue : float ndarray + The computed p-value of the test for each comparison. The element + at index ``i`` is the p-value for the comparison between + group ``i`` and the control. + + And the following method: + + confidence_interval(confidence_level=0.95) : + Compute the difference in means of the groups + with the control +- the allowance. + + See Also + -------- + tukey_hsd : performs pairwise comparison of means. + :ref:`hypothesis_dunnett` : Extended example + + Notes + ----- + Like the independent-sample t-test, Dunnett's test [1]_ is used to make + inferences about the means of distributions from which samples were drawn. + However, when multiple t-tests are performed at a fixed significance level, + the "family-wise error rate" - the probability of incorrectly rejecting the + null hypothesis in at least one test - will exceed the significance level. + Dunnett's test is designed to perform multiple comparisons while + controlling the family-wise error rate. + + Dunnett's test compares the means of multiple experimental groups + against a single control group. Tukey's Honestly Significant Difference Test + is another multiple-comparison test that controls the family-wise error + rate, but `tukey_hsd` performs *all* pairwise comparisons between groups. + When pairwise comparisons between experimental groups are not needed, + Dunnett's test is preferable due to its higher power. + + The use of this test relies on several assumptions. + + 1. The observations are independent within and among groups. + 2. The observations within each group are normally distributed. + 3. The distributions from which the samples are drawn have the same finite + variance. + + References + ---------- + .. [1] Dunnett, Charles W. (1955) "A Multiple Comparison Procedure for + Comparing Several Treatments with a Control." Journal of the American + Statistical Association, 50:272, 1096-1121, + :doi:`10.1080/01621459.1955.10501294` + .. [2] Thomson, M. L., & Short, M. D. (1969). Mucociliary function in + health, chronic obstructive airway disease, and asbestosis. Journal + of applied physiology, 26(5), 535-539. + :doi:`10.1152/jappl.1969.26.5.535` + + Examples + -------- + We'll use data from [2]_, Table 1. The null hypothesis is that the means of + the distributions underlying the samples and control are equal. + + First, we test that the means of the distributions underlying the samples + and control are unequal (``alternative='two-sided'``, the default). + + >>> import numpy as np + >>> from scipy.stats import dunnett + >>> samples = [[3.8, 2.7, 4.0, 2.4], [2.8, 3.4, 3.7, 2.2, 2.0]] + >>> control = [2.9, 3.0, 2.5, 2.6, 3.2] + >>> res = dunnett(*samples, control=control) + >>> res.statistic + array([ 0.90874545, -0.05007117]) + >>> res.pvalue + array([0.58325114, 0.99819341]) + + Now, we test that the means of the distributions underlying the samples are + greater than the mean of the distribution underlying the control. + + >>> res = dunnett(*samples, control=control, alternative='greater') + >>> res.statistic + array([ 0.90874545, -0.05007117]) + >>> res.pvalue + array([0.30230596, 0.69115597]) + + For a more detailed example, see :ref:`hypothesis_dunnett`. + """ + samples_, control_, rng = _iv_dunnett( + samples=samples, control=control, + alternative=alternative, rng=rng + ) + + rho, df, n_group, n_samples, n_control = _params_dunnett( + samples=samples_, control=control_ + ) + + statistic, std, mean_control, mean_samples = _statistic_dunnett( + samples_, control_, df, n_samples, n_control + ) + + pvalue = _pvalue_dunnett( + rho=rho, df=df, statistic=statistic, alternative=alternative, rng=rng + ) + + return DunnettResult( + statistic=statistic, pvalue=pvalue, + _alternative=alternative, + _rho=rho, _df=df, _std=std, + _mean_samples=mean_samples, + _mean_control=mean_control, + _n_samples=n_samples, + _n_control=n_control, + _rng=rng + ) + + +def _iv_dunnett( + samples: Sequence["npt.ArrayLike"], + control: "npt.ArrayLike", + alternative: Literal['two-sided', 'less', 'greater'], + rng: SeedType +) -> tuple[list[np.ndarray], np.ndarray, SeedType]: + """Input validation for Dunnett's test.""" + rng = check_random_state(rng) + + if alternative not in {'two-sided', 'less', 'greater'}: + raise ValueError( + "alternative must be 'less', 'greater' or 'two-sided'" + ) + + ndim_msg = "Control and samples groups must be 1D arrays" + n_obs_msg = "Control and samples groups must have at least 1 observation" + + control = np.asarray(control) + samples_ = [np.asarray(sample) for sample in samples] + + # samples checks + samples_control: list[np.ndarray] = samples_ + [control] + for sample in samples_control: + if sample.ndim > 1: + raise ValueError(ndim_msg) + + if sample.size < 1: + raise ValueError(n_obs_msg) + + return samples_, control, rng + + +def _params_dunnett( + samples: list[np.ndarray], control: np.ndarray +) -> tuple[np.ndarray, int, int, np.ndarray, int]: + """Specific parameters for Dunnett's test. + + Degree of freedom is the number of observations minus the number of groups + including the control. + """ + n_samples = np.array([sample.size for sample in samples]) + + # From [1] p. 1100 d.f. = (sum N)-(p+1) + n_sample = n_samples.sum() + n_control = control.size + n = n_sample + n_control + n_groups = len(samples) + df = n - n_groups - 1 + + # From [1] p. 1103 rho_ij = 1/sqrt((N0/Ni+1)(N0/Nj+1)) + rho = n_control/n_samples + 1 + rho = 1/np.sqrt(rho[:, None] * rho[None, :]) + np.fill_diagonal(rho, 1) + + return rho, df, n_groups, n_samples, n_control + + +def _statistic_dunnett( + samples: list[np.ndarray], control: np.ndarray, df: int, + n_samples: np.ndarray, n_control: int +) -> tuple[np.ndarray, float, np.ndarray, np.ndarray]: + """Statistic of Dunnett's test. + + Computation based on the original single-step test from [1]. + """ + mean_control = np.mean(control) + mean_samples = np.array([np.mean(sample) for sample in samples]) + all_samples = [control] + samples + all_means = np.concatenate([[mean_control], mean_samples]) + + # Variance estimate s^2 from [1] Eq. 1 + s2 = np.sum([_var(sample, mean=mean)*sample.size + for sample, mean in zip(all_samples, all_means)]) / df + std = np.sqrt(s2) + + # z score inferred from [1] unlabeled equation after Eq. 1 + z = (mean_samples - mean_control) / np.sqrt(1/n_samples + 1/n_control) + + return z / std, std, mean_control, mean_samples + + +def _pvalue_dunnett( + rho: np.ndarray, df: int, statistic: np.ndarray, + alternative: Literal['two-sided', 'less', 'greater'], + rng: SeedType = None +) -> np.ndarray: + """pvalue from the multivariate t-distribution. + + Critical values come from the multivariate student-t distribution. + """ + statistic = statistic.reshape(-1, 1) + + mvt = stats.multivariate_t(shape=rho, df=df, seed=rng) + if alternative == "two-sided": + statistic = abs(statistic) + pvalue = 1 - mvt.cdf(statistic, lower_limit=-statistic) + elif alternative == "greater": + pvalue = 1 - mvt.cdf(statistic, lower_limit=-np.inf) + else: + pvalue = 1 - mvt.cdf(np.inf, lower_limit=statistic) + + return np.atleast_1d(pvalue) diff --git a/v12x/pylib/scipy/stats/_multivariate.py b/v12x/pylib/scipy/stats/_multivariate.py new file mode 100644 index 0000000000000000000000000000000000000000..51e5c8dc8f5272381d17e8edc33662a635628bba --- /dev/null +++ b/v12x/pylib/scipy/stats/_multivariate.py @@ -0,0 +1,8054 @@ +# +# Author: Joris Vankerschaver 2013 +# +import math +import threading +import types +import numpy as np +import scipy.linalg +from scipy._lib import doccer +from scipy.special import (gammaln, psi, multigammaln, xlogy, entr, betaln, + ive, loggamma) +from scipy import special +import scipy._external.array_api_extra as xpx +from scipy._lib._util import check_random_state +from scipy.linalg.blas import get_blas_funcs +from ._continuous_distns import norm, invgamma +from ._discrete_distns import binom +from . import _covariance, _rcont +from ._qmvnt import _qmvt, _qmvn, _qauto +from ._morestats import directional_stats +from scipy.optimize import root_scalar + +__all__ = ['multivariate_normal', + 'matrix_normal', + 'dirichlet', + 'dirichlet_multinomial', + 'wishart', + 'invwishart', + 'multinomial', + 'special_ortho_group', + 'ortho_group', + 'random_correlation', + 'unitary_group', + 'multivariate_t', + 'multivariate_hypergeom', + 'random_table', + 'uniform_direction', + 'vonmises_fisher', + 'normal_inverse_gamma', + 'matrix_t'] + +_LOG_2PI = np.log(2 * np.pi) +_LOG_2 = np.log(2) +_LOG_PI = np.log(np.pi) +MVN_LOCK = threading.Lock() + + +_doc_random_state = """\ +seed : {None, int, np.random.RandomState, np.random.Generator}, optional + Used for drawing random variates. + If `seed` is `None`, the `~np.random.RandomState` singleton is used. + If `seed` is an int, a new ``RandomState`` instance is used, seeded + with seed. + If `seed` is already a ``RandomState`` or ``Generator`` instance, + then that object is used. + Default is `None`. +""" + + +def _squeeze_output(out): + """ + Remove single-dimensional entries from array and convert to scalar, + if necessary. + """ + out = out.squeeze() + if out.ndim == 0: + out = out[()] + return out + + +def _eigvalsh_to_eps(spectrum, cond=None, rcond=None): + """Determine which eigenvalues are "small" given the spectrum. + + This is for compatibility across various linear algebra functions + that should agree about whether or not a Hermitian matrix is numerically + singular and what is its numerical matrix rank. + This is designed to be compatible with scipy.linalg.pinvh. + + Parameters + ---------- + spectrum : 1d ndarray + Array of eigenvalues of a Hermitian matrix. + cond, rcond : float, optional + Cutoff for small eigenvalues. + Singular values smaller than rcond * largest_eigenvalue are + considered zero. + If None or -1, suitable machine precision is used. + + Returns + ------- + eps : float + Magnitude cutoff for numerical negligibility. + + """ + if rcond is not None: + cond = rcond + if cond in [None, -1]: + t = spectrum.dtype.char.lower() + factor = {'f': 1E3, 'd': 1E6} + cond = factor[t] * np.finfo(t).eps + eps = cond * np.max(abs(spectrum)) + return eps + + +def _pinv_1d(v, eps=1e-5): + """A helper function for computing the pseudoinverse. + + Parameters + ---------- + v : iterable of numbers + This may be thought of as a vector of eigenvalues or singular values. + eps : float + Values with magnitude no greater than eps are considered negligible. + + Returns + ------- + v_pinv : 1d float ndarray + A vector of pseudo-inverted numbers. + + """ + return np.array([0 if abs(x) <= eps else 1/x for x in v], dtype=float) + + +def _validate_marginal_input(dimensions, multivariate_dims): + """Determine if input dimensions can be marginalized. + + Parameters + ---------- + dimensions : float, ndarray + Input dimensions to be marginalized + multivariate_dims : int + Number of dimensions of multivariate distribution. + + Returns + ------- + dims : ndarray + Array of indices to marginalize + """ + dims = np.copy(dimensions) + dims = np.atleast_1d(dims) + + if len(dims) == 0: + msg = "Cannot marginalize all dimensions." + raise ValueError(msg) + + if not np.issubdtype(dims.dtype, np.integer): + msg = ("Elements of `dimensions` must be integers - the indices " + "of the marginal variables being retained.") + raise ValueError(msg) + + original_dims = np.copy(dims) + + dims[dims < 0] += multivariate_dims + + if len(np.unique(dims)) != len(dims): + msg = "All elements of `dimensions` must be unique." + raise ValueError(msg) + + i_invalid = (dims < 0) | (dims >= multivariate_dims) + if np.any(i_invalid): + msg = (f"Dimensions {original_dims[i_invalid]} are invalid " + f"for a distribution in {multivariate_dims} dimensions.") + raise ValueError(msg) + + return dims + + +class _PSD: + """ + Compute coordinated functions of a symmetric positive semidefinite matrix. + + This class addresses two issues. Firstly it allows the pseudoinverse, + the logarithm of the pseudo-determinant, and the rank of the matrix + to be computed using one call to eigh instead of three. + Secondly it allows these functions to be computed in a way + that gives mutually compatible results. + All of the functions are computed with a common understanding as to + which of the eigenvalues are to be considered negligibly small. + The functions are designed to coordinate with scipy.linalg.pinvh() + but not necessarily with np.linalg.det() or with np.linalg.matrix_rank(). + + Parameters + ---------- + M : array_like + Symmetric positive semidefinite matrix (2-D). + cond, rcond : float, optional + Cutoff for small eigenvalues. + Singular values smaller than rcond * largest_eigenvalue are + considered zero. + If None or -1, suitable machine precision is used. + lower : bool, optional + Whether the pertinent array data is taken from the lower + or upper triangle of M. (Default: lower) + check_finite : bool, optional + Whether to check that the input matrices contain only finite + numbers. Disabling may give a performance gain, but may result + in problems (crashes, non-termination) if the inputs do contain + infinities or NaNs. + allow_singular : bool, optional + Whether to allow a singular matrix. (Default: True) + + Notes + ----- + The arguments are similar to those of scipy.linalg.pinvh(). + + """ + + def __init__(self, M, cond=None, rcond=None, lower=True, + check_finite=True, allow_singular=True): + self._M = np.asarray(M) + + # Compute the symmetric eigendecomposition. + # Note that eigh takes care of array conversion, chkfinite, + # and assertion that the matrix is square. + s, u = scipy.linalg.eigh(M, lower=lower, check_finite=check_finite) + + eps = _eigvalsh_to_eps(s, cond, rcond) + if np.min(s) < -eps: + msg = "The input matrix must be symmetric positive semidefinite." + raise ValueError(msg) + d = s[s > eps] + if len(d) < len(s) and not allow_singular: + msg = ("When `allow_singular is False`, the input matrix must be " + "symmetric positive definite.") + raise np.linalg.LinAlgError(msg) + s_pinv = _pinv_1d(s, eps) + U = np.multiply(u, np.sqrt(s_pinv)) + + # Save the eigenvector basis, and tolerance for testing support + self.eps = 1e3*eps + self.V = u[:, s <= eps] + + # Initialize the eagerly precomputed attributes. + self.rank = len(d) + self.U = U + self.log_pdet = np.sum(np.log(d)) + + # Initialize attributes to be lazily computed. + self._pinv = None + + def _support_mask(self, x): + """ + Check whether x lies in the support of the distribution. + """ + residual = np.linalg.norm(x @ self.V, axis=-1) + in_support = residual < self.eps + return in_support + + @property + def pinv(self): + if self._pinv is None: + self._pinv = np.dot(self.U, self.U.T) + return self._pinv + + +class multi_rv_generic: + """ + Class which encapsulates common functionality between all multivariate + distributions. + """ + + def __init__(self, seed=None): + super().__init__() + self._random_state = check_random_state(seed) + + @property + def random_state(self): + """ Get or set the Generator object for generating random variates. + + If `seed` is None (or `np.random`), the `numpy.random.RandomState` + singleton is used. + If `seed` is an int, a new ``RandomState`` instance is used, + seeded with `seed`. + If `seed` is already a ``Generator`` or ``RandomState`` instance then + that instance is used. + + """ + return self._random_state + + @random_state.setter + def random_state(self, seed): + self._random_state = check_random_state(seed) + + def _get_random_state(self, random_state): + if random_state is not None: + return check_random_state(random_state) + else: + return self._random_state + + +class multi_rv_frozen: + """ + Class which encapsulates common functionality between all frozen + multivariate distributions. + """ + + # generic type compatibility with scipy-stubs + __class_getitem__: classmethod = classmethod(types.GenericAlias) + + @property + def random_state(self): + return self._dist._random_state + + @random_state.setter + def random_state(self, seed): + self._dist._random_state = check_random_state(seed) + + +_mvn_doc_default_callparams = """\ +mean : array_like, default: ``[0]`` + Mean of the distribution. +cov : array_like or `Covariance`, default: ``[1]`` + Symmetric positive (semi)definite covariance matrix of the distribution. +allow_singular : bool, default: ``False`` + Whether to allow a singular covariance matrix. This is ignored if `cov` is + a `Covariance` object. +""" + +_mvn_doc_callparams_note = """\ +Setting the parameter `mean` to `None` is equivalent to having `mean` +be the zero-vector. The parameter `cov` can be a scalar, in which case +the covariance matrix is the identity times that value, a vector of +diagonal entries for the covariance matrix, a two-dimensional array_like, +or a `Covariance` object. +""" + +_mvn_doc_frozen_callparams = "" + +_mvn_doc_frozen_callparams_note = """\ +See class definition for a detailed description of parameters.""" + +mvn_docdict_params = { + '_mvn_doc_default_callparams': _mvn_doc_default_callparams, + '_mvn_doc_callparams_note': _mvn_doc_callparams_note, + '_doc_random_state': _doc_random_state +} + +mvn_docdict_noparams = { + '_mvn_doc_default_callparams': _mvn_doc_frozen_callparams, + '_mvn_doc_callparams_note': _mvn_doc_frozen_callparams_note, + '_doc_random_state': _doc_random_state +} + + +class multivariate_normal_gen(multi_rv_generic): + r""" + A multivariate normal random variable. + + The `mean` keyword specifies the mean. The `cov` keyword specifies the + covariance matrix. + + Parameters + ---------- + %(_mvn_doc_default_callparams)s + %(_doc_random_state)s + + Methods + ------- + pdf(x, mean=None, cov=1, allow_singular=False) + Probability density function. + logpdf(x, mean=None, cov=1, allow_singular=False) + Log of the probability density function. + cdf(x, mean=None, cov=1, allow_singular=False, maxpts=1000000*dim, abseps=1e-5, releps=1e-5, lower_limit=None) + Cumulative distribution function. + logcdf(x, mean=None, cov=1, allow_singular=False, maxpts=1000000*dim, abseps=1e-5, releps=1e-5) + Log of the cumulative distribution function. + rvs(mean=None, cov=1, size=1, random_state=None) + Draw random samples from a multivariate normal distribution. + entropy(mean=None, cov=1) + Compute the differential entropy of the multivariate normal. + marginal(dimensions, mean=None, cov=1, allow_singular=False) + Return a marginal multivariate normal distribution. + fit(x, fix_mean=None, fix_cov=None) + Fit a multivariate normal distribution to data. + + Notes + ----- + %(_mvn_doc_callparams_note)s + + The covariance matrix `cov` may be an instance of a subclass of + `Covariance`, e.g. `scipy.stats.CovViaPrecision`. If so, `allow_singular` + is ignored. + + Otherwise, `cov` must be a symmetric positive semidefinite + matrix when `allow_singular` is True; it must be (strictly) positive + definite when `allow_singular` is False. + Symmetry is not checked; only the lower triangular portion is used. + The determinant and inverse of `cov` are computed + as the pseudo-determinant and pseudo-inverse, respectively, so + that `cov` does not need to have full rank. + + The probability density function for `multivariate_normal` is + + .. math:: + + f(x) = \frac{1}{\sqrt{(2 \pi)^k \det \Sigma}} + \exp\left( -\frac{1}{2} (x - \mu)^T \Sigma^{-1} (x - \mu) \right), + + where :math:`\mu` is the mean, :math:`\Sigma` the covariance matrix, + :math:`k` the rank of :math:`\Sigma`. In case of singular :math:`\Sigma`, + SciPy extends this definition according to [1]_. + + .. versionadded:: 0.14.0 + + References + ---------- + .. [1] Multivariate Normal Distribution - Degenerate Case, Wikipedia, + https://en.wikipedia.org/wiki/Multivariate_normal_distribution#Degenerate_case + + Examples + -------- + >>> import numpy as np + >>> import matplotlib.pyplot as plt + >>> from scipy.stats import multivariate_normal + + >>> x = np.linspace(0, 5, 10, endpoint=False) + >>> y = multivariate_normal.pdf(x, mean=2.5, cov=0.5); y + array([ 0.00108914, 0.01033349, 0.05946514, 0.20755375, 0.43939129, + 0.56418958, 0.43939129, 0.20755375, 0.05946514, 0.01033349]) + >>> fig1 = plt.figure() + >>> ax = fig1.add_subplot(111) + >>> ax.plot(x, y) + >>> plt.show() + + Alternatively, the object may be called (as a function) to fix the mean + and covariance parameters, returning a "frozen" multivariate normal + random variable: + + >>> rv = multivariate_normal(mean=None, cov=1, allow_singular=False) + >>> # Frozen object with the same methods but holding the given + >>> # mean and covariance fixed. + + The input quantiles can be any shape of array, as long as the last + axis labels the components. This allows us for instance to + display the frozen pdf for a non-isotropic random variable in 2D as + follows: + + >>> x, y = np.mgrid[-1:1:.01, -1:1:.01] + >>> pos = np.dstack((x, y)) + >>> rv = multivariate_normal([0.5, -0.2], [[2.0, 0.3], [0.3, 0.5]]) + >>> fig2 = plt.figure() + >>> ax2 = fig2.add_subplot(111) + >>> ax2.contourf(x, y, rv.pdf(pos)) + """ # noqa: E501 + + def __init__(self, seed=None): + super().__init__(seed) + self.__doc__ = doccer.docformat(self.__doc__, mvn_docdict_params) + + def __call__(self, mean=None, cov=1, allow_singular=False, seed=None, **kwds): + """Create a frozen multivariate normal distribution. + + See `multivariate_normal_frozen` for more information. + """ + return multivariate_normal_frozen(mean, cov, + allow_singular=allow_singular, + seed=seed, **kwds) + + def _process_parameters(self, mean, cov, allow_singular=True): + """ + Infer dimensionality from mean or covariance matrix, ensure that + mean and covariance are full vector resp. matrix. + """ + if isinstance(cov, _covariance.Covariance): + return self._process_parameters_Covariance(mean, cov) + else: + # Before `Covariance` classes were introduced, + # `multivariate_normal` accepted plain arrays as `cov` and used the + # following input validation. To avoid disturbing the behavior of + # `multivariate_normal` when plain arrays are used, we use the + # original input validation here. + dim, mean, cov = self._process_parameters_psd(None, mean, cov) + # After input validation, some methods then processed the arrays + # with a `_PSD` object and used that to perform computation. + # To avoid branching statements in each method depending on whether + # `cov` is an array or `Covariance` object, we always process the + # array with `_PSD`, and then use wrapper that satisfies the + # `Covariance` interface, `CovViaPSD`. + psd = _PSD(cov, allow_singular=allow_singular) + cov_object = _covariance.CovViaPSD(psd) + return dim, mean, cov_object + + def _process_parameters_Covariance(self, mean, cov): + dim = cov.shape[-1] + mean = np.array([0.]) if mean is None else mean + message = (f"`cov` represents a covariance matrix in {dim} dimensions," + f"and so `mean` must be broadcastable to shape {(dim,)}") + try: + mean = np.broadcast_to(mean, dim) + except ValueError as e: + raise ValueError(message) from e + return dim, mean, cov + + def _process_parameters_psd(self, dim, mean, cov): + # Try to infer dimensionality + if dim is None: + if mean is None: + if cov is None: + dim = 1 + else: + cov = np.asarray(cov, dtype=float) + if cov.ndim < 2: + dim = 1 + else: + dim = cov.shape[0] + else: + mean = np.asarray(mean, dtype=float) + dim = mean.size + else: + if not np.isscalar(dim): + raise ValueError("Dimension of random variable must be " + "a scalar.") + + # Check input sizes and return full arrays for mean and cov if + # necessary + if mean is None: + mean = np.zeros(dim) + mean = np.asarray(mean, dtype=float) + + if cov is None: + cov = 1.0 + cov = np.asarray(cov, dtype=float) + + if dim == 1: + mean = mean.reshape(1) + cov = cov.reshape(1, 1) + + if mean.ndim != 1 or mean.shape[0] != dim: + raise ValueError(f"Array 'mean' must be a vector of length {dim}.") + if cov.ndim == 0: + cov = cov * np.eye(dim) + elif cov.ndim == 1: + cov = np.diag(cov) + elif cov.ndim == 2 and cov.shape != (dim, dim): + rows, cols = cov.shape + if rows != cols: + msg = ("Array 'cov' must be square if it is two dimensional," + f" but cov.shape = {str(cov.shape)}.") + else: + msg = (f"Dimension mismatch: array 'cov' is of shape {cov.shape}, " + f"but 'mean' is a vector of length {len(mean)}.") + raise ValueError(msg) + elif cov.ndim > 2: + raise ValueError(f"Array 'cov' must be at most two-dimensional, " + f"but cov.ndim = {cov.ndim}") + + return dim, mean, cov + + def _process_quantiles(self, x, dim): + """ + Adjust quantiles array so that last axis labels the components of + each data point. + """ + x = np.asarray(x, dtype=float) + + if x.ndim == 0: + x = x[np.newaxis] + elif x.ndim == 1: + if dim == 1: + x = x[:, np.newaxis] + else: + x = x[np.newaxis, :] + + return x + + def _logpdf(self, x, mean, cov_object): + """Log of the multivariate normal probability density function. + + Parameters + ---------- + x : ndarray + Points at which to evaluate the log of the probability + density function + mean : ndarray + Mean of the distribution + cov_object : Covariance + An object representing the Covariance matrix + + Notes + ----- + As this function does no argument checking, it should not be + called directly; use 'logpdf' instead. + + """ + log_det_cov, rank = cov_object.log_pdet, cov_object.rank + dev = x - mean + if dev.ndim > 1: + log_det_cov = log_det_cov[..., np.newaxis] + rank = rank[..., np.newaxis] + maha = np.sum(np.square(cov_object.whiten(dev)), axis=-1) + return -0.5 * (rank * _LOG_2PI + log_det_cov + maha) + + def logpdf(self, x, mean=None, cov=1, allow_singular=False): + """Log of the multivariate normal probability density function. + + Parameters + ---------- + x : array_like + Quantiles, with the last axis of `x` denoting the components. + %(_mvn_doc_default_callparams)s + + Returns + ------- + pdf : ndarray or scalar + Log of the probability density function evaluated at `x` + + Notes + ----- + %(_mvn_doc_callparams_note)s + + """ + params = self._process_parameters(mean, cov, allow_singular) + dim, mean, cov_object = params + x = self._process_quantiles(x, dim) + out = self._logpdf(x, mean, cov_object) + if np.any(cov_object.rank < dim): + out_of_bounds = ~cov_object._support_mask(x-mean) + out[out_of_bounds] = -np.inf + return _squeeze_output(out) + + def pdf(self, x, mean=None, cov=1, allow_singular=False): + """Multivariate normal probability density function. + + Parameters + ---------- + x : array_like + Quantiles, with the last axis of `x` denoting the components. + %(_mvn_doc_default_callparams)s + + Returns + ------- + pdf : ndarray or scalar + Probability density function evaluated at `x` + + Notes + ----- + %(_mvn_doc_callparams_note)s + + """ + params = self._process_parameters(mean, cov, allow_singular) + dim, mean, cov_object = params + x = self._process_quantiles(x, dim) + out = np.exp(self._logpdf(x, mean, cov_object)) + if np.any(cov_object.rank < dim): + out_of_bounds = ~cov_object._support_mask(x-mean) + out[out_of_bounds] = 0.0 + return _squeeze_output(out) + + def _cdf(self, x, mean, cov, maxpts, abseps, releps, lower_limit, rng): + """Multivariate normal cumulative distribution function. + + Parameters + ---------- + x : ndarray + Points at which to evaluate the cumulative distribution function. + mean : ndarray + Mean of the distribution + cov : array_like + Covariance matrix of the distribution + maxpts : int + The maximum number of points to use for integration + abseps : float + Absolute error tolerance + releps : float + Relative error tolerance + lower_limit : array_like, optional + Lower limit of integration of the cumulative distribution function. + Default is negative infinity. Must be broadcastable with `x`. + rng : Generator + an instance of ``np.random.Generator``, which is used internally + for QMC integration. + + Notes + ----- + As this function does no argument checking, it should not be + called directly; use 'cdf' instead. + + + .. versionadded:: 1.0.0 + + """ + lower = (np.full(mean.shape, -np.inf) + if lower_limit is None else lower_limit) + # In 2d, _mvn.mvnun accepts input in which `lower` bound elements + # are greater than `x`. Not so in other dimensions. Fix this by + # ensuring that lower bounds are indeed lower when passed, then + # set signs of resulting CDF manually. + b, a = np.broadcast_arrays(x, lower) + b, a = b - mean, a - mean # _qmvn only accepts zero mean + i_swap = b < a + signs = (-1)**(i_swap.sum(axis=-1)) # odd # of swaps -> negative + a, b = a.copy(), b.copy() + a[i_swap], b[i_swap] = b[i_swap], a[i_swap] + n = x.shape[-1] + limits = np.concatenate((a, b), axis=-1) + + # qmvn expects 1-d arguments, so process points sequentially + # XXX: if cov.ndim == 2 and limits.ndim == 1, can avoid apply_along_axis + def func1d(limits): + # res0 = _qmvn(maxpts, cov, limits[:n], limits[n:], rng)[0] + res = _qauto(_qmvn, cov, limits[:n], limits[n:], + rng, error=abseps, limit=maxpts, n_batches=10) + return np.squeeze(res[0]) + + out = np.apply_along_axis(func1d, -1, limits) * signs + return _squeeze_output(out) + + def logcdf(self, x, mean=None, cov=1, allow_singular=False, maxpts=None, + abseps=1e-5, releps=1e-5, *, lower_limit=None, rng=None): + """Log of the multivariate normal cumulative distribution function. + + Parameters + ---------- + x : array_like + Quantiles, with the last axis of `x` denoting the components. + %(_mvn_doc_default_callparams)s + maxpts : int, optional + The maximum number of points to use for integration + (default ``1000000*dim``) + abseps : float, optional + Absolute error tolerance (default 1e-5) + releps : float, optional + Relative error tolerance (default 1e-5) + lower_limit : array_like, optional + Lower limit of integration of the cumulative distribution function. + Default is negative infinity. Must be broadcastable with `x`. + rng : Generator, optional + an instance of ``np.random.Generator``, which is used internally + for QMC integration. + + Returns + ------- + cdf : ndarray or scalar + Log of the cumulative distribution function evaluated at `x` + + Notes + ----- + %(_mvn_doc_callparams_note)s + + .. versionadded:: 1.0.0 + + """ + params = self._process_parameters(mean, cov, allow_singular) + dim, mean, cov_object = params + cov = cov_object.covariance + x = self._process_quantiles(x, dim) + if not maxpts: + maxpts = 1000000 * dim + + rng = self._get_random_state(rng) + cdf = self._cdf(x, mean, cov, maxpts, abseps, releps, lower_limit, rng) + # the log of a negative real is complex, and cdf can be negative + # if lower limit is greater than upper limit + cdf = cdf + 0j if np.any(cdf < 0) else cdf + out = np.log(cdf) + return out + + def cdf(self, x, mean=None, cov=1, allow_singular=False, maxpts=None, + abseps=1e-5, releps=1e-5, *, lower_limit=None, rng=None): + """Multivariate normal cumulative distribution function. + + Parameters + ---------- + x : array_like + Quantiles, with the last axis of `x` denoting the components. + %(_mvn_doc_default_callparams)s + maxpts : int, optional + The maximum number of points to use for integration + (default ``1000000*dim``) + abseps : float, optional + Absolute error tolerance (default 1e-5) + releps : float, optional + Relative error tolerance (default 1e-5) + lower_limit : array_like, optional + Lower limit of integration of the cumulative distribution function. + Default is negative infinity. Must be broadcastable with `x`. + rng : Generator, optional + an instance of ``np.random.Generator``, which is used internally + for QMC integration. + + Returns + ------- + cdf : ndarray or scalar + Cumulative distribution function evaluated at `x` + + Notes + ----- + %(_mvn_doc_callparams_note)s + + .. versionadded:: 1.0.0 + + """ + params = self._process_parameters(mean, cov, allow_singular) + dim, mean, cov_object = params + cov = cov_object.covariance + x = self._process_quantiles(x, dim) + if not maxpts: + maxpts = 1000000 * dim + rng = self._get_random_state(rng) + out = self._cdf(x, mean, cov, maxpts, abseps, releps, lower_limit, rng) + return out + + def rvs(self, mean=None, cov=1, size=1, random_state=None): + """Draw random samples from a multivariate normal distribution. + + Parameters + ---------- + %(_mvn_doc_default_callparams)s + size : int, optional + Number of samples to draw (default 1). + %(_doc_random_state)s + + Returns + ------- + rvs : ndarray or scalar + Random variates of size (`size`, `N`), where `N` is the + dimension of the random variable. + + Notes + ----- + %(_mvn_doc_callparams_note)s + + """ + dim, mean, cov_object = self._process_parameters(mean, cov) + random_state = self._get_random_state(random_state) + + if isinstance(cov_object, _covariance.CovViaPSD): + cov = cov_object.covariance + out = random_state.multivariate_normal(mean, cov, size) + out = _squeeze_output(out) + else: + size = size or tuple() + if not np.iterable(size): + size = (size,) + shape = tuple(size) + (cov_object.shape[-1],) + x = random_state.normal(size=shape) + out = mean + cov_object.colorize(x) + return out + + def entropy(self, mean=None, cov=1): + """Compute the differential entropy of the multivariate normal. + + Parameters + ---------- + %(_mvn_doc_default_callparams)s + + Returns + ------- + h : scalar + Entropy of the multivariate normal distribution + + Notes + ----- + %(_mvn_doc_callparams_note)s + + """ + dim, mean, cov_object = self._process_parameters(mean, cov) + return 0.5 * (cov_object.rank * (_LOG_2PI + 1) + cov_object.log_pdet) + + def fit(self, x, fix_mean=None, fix_cov=None): + """Fit a multivariate normal distribution to data. + + Parameters + ---------- + x : ndarray (m, n) + Data the distribution is fitted to. Must have two axes. + The first axis of length `m` represents the number of vectors + the distribution is fitted to. The second axis of length `n` + determines the dimensionality of the fitted distribution. + fix_mean : ndarray(n, ) + Fixed mean vector. Must have length `n`. + fix_cov: ndarray (n, n) + Fixed covariance matrix. Must have shape ``(n, n)``. + + Returns + ------- + mean : ndarray (n, ) + Maximum likelihood estimate of the mean vector + cov : ndarray (n, n) + Maximum likelihood estimate of the covariance matrix + + """ + # input validation for data to be fitted + x = np.asarray(x) + if x.ndim != 2: + raise ValueError("`x` must be two-dimensional.") + + n_vectors, dim = x.shape + + # parameter estimation + # reference: https://home.ttic.edu/~shubhendu/Slides/Estimation.pdf + if fix_mean is not None: + # input validation for `fix_mean` + fix_mean = np.atleast_1d(fix_mean) + if fix_mean.shape != (dim, ): + msg = ("`fix_mean` must be a one-dimensional array the same " + "length as the dimensionality of the vectors `x`.") + raise ValueError(msg) + mean = fix_mean + else: + mean = x.mean(axis=0) + + if fix_cov is not None: + # input validation for `fix_cov` + fix_cov = np.atleast_2d(fix_cov) + # validate shape + if fix_cov.shape != (dim, dim): + msg = ("`fix_cov` must be a two-dimensional square array " + "of same side length as the dimensionality of the " + "vectors `x`.") + raise ValueError(msg) + # validate positive semidefiniteness + # a trimmed down copy from _PSD + s, u = scipy.linalg.eigh(fix_cov, lower=True, check_finite=True) + eps = _eigvalsh_to_eps(s) + if np.min(s) < -eps: + msg = "`fix_cov` must be symmetric positive semidefinite." + raise ValueError(msg) + cov = fix_cov + else: + centered_data = x - mean + cov = centered_data.T @ centered_data / n_vectors + return mean, cov + + def marginal(self, dimensions, mean=None, cov=1, allow_singular=False): + """Return a marginal multivariate normal distribution. + + Parameters + ---------- + dimensions : int or 1-d array_like + The dimensions of the multivariate distribution corresponding + with the marginal variables, that is, the indices of the dimensions + that are being retained. The other dimensions are marginalized out. + %(_mvn_doc_default_callparams)s + + Returns + ------- + marginal_multivariate_normal : multivariate_normal_frozen + An object representing the marginal distribution. + + Notes + ----- + %(_mvn_doc_callparams_note)s + """ + params = self._process_parameters(mean, cov, allow_singular) + n, mean, cov_object = params + dims = _validate_marginal_input(dimensions, n) + + mean = mean[dims] + cov = cov_object.covariance[np.ix_(dims, dims)] + + return multivariate_normal_frozen(mean, cov, allow_singular) + +multivariate_normal = multivariate_normal_gen() + + +class multivariate_normal_frozen(multi_rv_frozen): + __class_getitem__ = None # pyrefly:ignore[bad-assignment] + + def __init__(self, mean=None, cov=1, allow_singular=False, seed=None, + maxpts=None, abseps=1e-5, releps=1e-5): + """Create a frozen multivariate normal distribution. + + Parameters + ---------- + mean : array_like, default: ``[0]`` + Mean of the distribution. + cov : array_like, default: ``[1]`` + Symmetric positive (semi)definite covariance matrix of the + distribution. + allow_singular : bool, default: ``False`` + Whether to allow a singular covariance matrix. + seed : {None, int, `numpy.random.Generator`, `numpy.random.RandomState`}, optional + If `seed` is None (or `np.random`), the `numpy.random.RandomState` + singleton is used. + If `seed` is an int, a new ``RandomState`` instance is used, + seeded with `seed`. + If `seed` is already a ``Generator`` or ``RandomState`` instance + then that instance is used. + maxpts : int, optional + The maximum number of points to use for integration of the + cumulative distribution function (default ``1000000*dim``) + abseps : float, optional + Absolute error tolerance for the cumulative distribution function + (default 1e-5) + releps : float, optional + Relative error tolerance for the cumulative distribution function + (default 1e-5) + + Examples + -------- + When called with the default parameters, this will create a 1D random + variable with mean 0 and covariance 1: + + >>> from scipy.stats import multivariate_normal + >>> r = multivariate_normal() + >>> r.mean + array([ 0.]) + >>> r.cov + array([[1.]]) + + """ # numpy/numpydoc#87 # noqa: E501 + self._dist = multivariate_normal_gen(seed) + self.dim, self.mean, self.cov_object = ( + self._dist._process_parameters(mean, cov, allow_singular)) + self.allow_singular = allow_singular or self.cov_object._allow_singular + if not maxpts: + maxpts = 1000000 * self.dim + self.maxpts = maxpts + self.abseps = abseps + self.releps = releps + + @property + def cov(self): + return self.cov_object.covariance + + def logpdf(self, x): + x = self._dist._process_quantiles(x, self.dim) + out = self._dist._logpdf(x, self.mean, self.cov_object) + if np.any(self.cov_object.rank < self.dim): + out_of_bounds = ~self.cov_object._support_mask(x-self.mean) + out[out_of_bounds] = -np.inf + return _squeeze_output(out) + + def pdf(self, x): + return np.exp(self.logpdf(x)) + + def logcdf(self, x, *, lower_limit=None, rng=None): + cdf = self.cdf(x, lower_limit=lower_limit, rng=rng) + # the log of a negative real is complex, and cdf can be negative + # if lower limit is greater than upper limit + cdf = cdf + 0j if np.any(cdf < 0) else cdf + out = np.log(cdf) + return out + + def cdf(self, x, *, lower_limit=None, rng=None): + x = self._dist._process_quantiles(x, self.dim) + rng = self._dist._get_random_state(rng) + out = self._dist._cdf(x, self.mean, self.cov_object.covariance, + self.maxpts, self.abseps, self.releps, + lower_limit, rng) + return _squeeze_output(out) + + def rvs(self, size=1, random_state=None): + return self._dist.rvs(self.mean, self.cov_object, size, random_state) + + def entropy(self): + """Computes the differential entropy of the multivariate normal. + + Returns + ------- + h : scalar + Entropy of the multivariate normal distribution + + """ + log_pdet = self.cov_object.log_pdet + rank = self.cov_object.rank + return 0.5 * (rank * (_LOG_2PI + 1) + log_pdet) + + def marginal(self, dimensions): + return self._dist.marginal(dimensions, self.mean, + self.cov_object, self.allow_singular) + +# Set frozen generator docstrings from corresponding docstrings in +# multivariate_normal_gen and fill in default strings in class docstrings +for name in ['logpdf', 'pdf', 'logcdf', 'cdf', 'rvs']: + method = multivariate_normal_gen.__dict__[name] + method_frozen = multivariate_normal_frozen.__dict__[name] + method_frozen.__doc__ = doccer.docformat(method.__doc__, + mvn_docdict_noparams) + method.__doc__ = doccer.docformat(method.__doc__, mvn_docdict_params) + +_matnorm_doc_default_callparams = """\ +mean : array_like, optional + Mean of the distribution (default: `None`) +rowcov : array_like, optional + Among-row covariance matrix of the distribution (default: ``1``) +colcov : array_like, optional + Among-column covariance matrix of the distribution (default: ``1``) +""" + +_matnorm_doc_callparams_note = """\ +If `mean` is set to `None` then a matrix of zeros is used for the mean. +The dimensions of this matrix are inferred from the shape of `rowcov` and +`colcov`, if these are provided, or set to ``1`` if ambiguous. + +`rowcov` and `colcov` can be two-dimensional array_likes specifying the +covariance matrices directly. Alternatively, a one-dimensional array will +be be interpreted as the entries of a diagonal matrix, and a scalar or +zero-dimensional array will be interpreted as this value times the +identity matrix. +""" + +_matnorm_doc_frozen_callparams = "" + +_matnorm_doc_frozen_callparams_note = """\ +See class definition for a detailed description of parameters.""" + +matnorm_docdict_params = { + '_matnorm_doc_default_callparams': _matnorm_doc_default_callparams, + '_matnorm_doc_callparams_note': _matnorm_doc_callparams_note, + '_doc_random_state': _doc_random_state +} + +matnorm_docdict_noparams = { + '_matnorm_doc_default_callparams': _matnorm_doc_frozen_callparams, + '_matnorm_doc_callparams_note': _matnorm_doc_frozen_callparams_note, + '_doc_random_state': _doc_random_state +} + + +class matrix_normal_gen(multi_rv_generic): + r""" + A matrix normal random variable. + + The `mean` keyword specifies the mean. The `rowcov` keyword specifies the + among-row covariance matrix. The 'colcov' keyword specifies the + among-column covariance matrix. + + Parameters + ---------- + %(_matnorm_doc_default_callparams)s + %(_doc_random_state)s + + Methods + ------- + pdf(X, mean=None, rowcov=1, colcov=1) + Probability density function. + logpdf(X, mean=None, rowcov=1, colcov=1) + Log of the probability density function. + rvs(mean=None, rowcov=1, colcov=1, size=1, random_state=None) + Draw random samples. + entropy(rowcol=1, colcov=1) + Differential entropy. + + Notes + ----- + %(_matnorm_doc_callparams_note)s + + The covariance matrices specified by `rowcov` and `colcov` must be + (symmetric) positive definite. If the samples in `X` are + :math:`m \times n`, then `rowcov` must be :math:`m \times m` and + `colcov` must be :math:`n \times n`. `mean` must be the same shape as `X`. + + The probability density function for `matrix_normal` is + + .. math:: + + f(X) = (2 \pi)^{-\frac{mn}{2}}|U|^{-\frac{n}{2}} |V|^{-\frac{m}{2}} + \exp\left( -\frac{1}{2} \mathrm{Tr}\left[ U^{-1} (X-M) V^{-1} + (X-M)^T \right] \right), + + where :math:`M` is the mean, :math:`U` the among-row covariance matrix, + :math:`V` the among-column covariance matrix. + + The `allow_singular` behaviour of the `multivariate_normal` + distribution is not currently supported. Covariance matrices must be + full rank. + + The `matrix_normal` distribution is closely related to the + `multivariate_normal` distribution. Specifically, :math:`\mathrm{Vec}(X)` + (the vector formed by concatenating the columns of :math:`X`) has a + multivariate normal distribution with mean :math:`\mathrm{Vec}(M)` + and covariance :math:`V \otimes U` (where :math:`\otimes` is the Kronecker + product). Sampling and pdf evaluation are + :math:`\mathcal{O}(m^3 + n^3 + m^2 n + m n^2)` for the matrix normal, but + :math:`\mathcal{O}(m^3 n^3)` for the equivalent multivariate normal, + making this equivalent form algorithmically inefficient. + + .. versionadded:: 0.17.0 + + References + ---------- + .. [1] "Matrix normal distribution", Wikipedia, + https://en.wikipedia.org/wiki/Matrix_normal_distribution + + Examples + -------- + + >>> import numpy as np + >>> from scipy.stats import matrix_normal + + >>> M = np.arange(6).reshape(3,2); M + array([[0, 1], + [2, 3], + [4, 5]]) + >>> U = np.diag([1,2,3]); U + array([[1, 0, 0], + [0, 2, 0], + [0, 0, 3]]) + >>> V = 0.3*np.identity(2); V + array([[ 0.3, 0. ], + [ 0. , 0.3]]) + >>> X = M + 0.1; X + array([[ 0.1, 1.1], + [ 2.1, 3.1], + [ 4.1, 5.1]]) + >>> matrix_normal.pdf(X, mean=M, rowcov=U, colcov=V) + 0.023410202050005054 + + >>> # Equivalent multivariate normal + >>> from scipy.stats import multivariate_normal + >>> vectorised_X = X.T.flatten() + >>> equiv_mean = M.T.flatten() + >>> equiv_cov = np.kron(V,U) + >>> multivariate_normal.pdf(vectorised_X, mean=equiv_mean, cov=equiv_cov) + 0.023410202050005054 + + Alternatively, the object may be called (as a function) to fix the mean + and covariance parameters, returning a "frozen" matrix normal + random variable: + + >>> rv = matrix_normal(mean=None, rowcov=1, colcov=1) + >>> # Frozen object with the same methods but holding the given + >>> # mean and covariance fixed. + """ + + def __init__(self, seed=None): + super().__init__(seed) + self.__doc__ = doccer.docformat(self.__doc__, matnorm_docdict_params) + + def __call__(self, mean=None, rowcov=1, colcov=1, seed=None): + """Create a frozen matrix normal distribution. + + See `matrix_normal_frozen` for more information. + + """ + return matrix_normal_frozen(mean, rowcov, colcov, seed=seed) + + def _process_parameters(self, mean, rowcov, colcov): + """ + Infer dimensionality from mean or covariance matrices. Handle + defaults. Ensure compatible dimensions. + """ + + # Process mean + if mean is not None: + mean = np.asarray(mean, dtype=float) + meanshape = mean.shape + if len(meanshape) != 2: + raise ValueError("Array `mean` must be two dimensional.") + if np.any(meanshape == 0): + raise ValueError("Array `mean` has invalid shape.") + + # Process among-row covariance + rowcov = np.asarray(rowcov, dtype=float) + if rowcov.ndim == 0: + if mean is not None: + rowcov = rowcov * np.identity(meanshape[0]) + else: + rowcov = rowcov * np.identity(1) + elif rowcov.ndim == 1: + rowcov = np.diag(rowcov) + rowshape = rowcov.shape + if len(rowshape) != 2: + raise ValueError("`rowcov` must be a scalar or a 2D array.") + if rowshape[0] != rowshape[1]: + raise ValueError("Array `rowcov` must be square.") + if rowshape[0] == 0: + raise ValueError("Array `rowcov` has invalid shape.") + numrows = rowshape[0] + + # Process among-column covariance + colcov = np.asarray(colcov, dtype=float) + if colcov.ndim == 0: + if mean is not None: + colcov = colcov * np.identity(meanshape[1]) + else: + colcov = colcov * np.identity(1) + elif colcov.ndim == 1: + colcov = np.diag(colcov) + colshape = colcov.shape + if len(colshape) != 2: + raise ValueError("`colcov` must be a scalar or a 2D array.") + if colshape[0] != colshape[1]: + raise ValueError("Array `colcov` must be square.") + if colshape[0] == 0: + raise ValueError("Array `colcov` has invalid shape.") + numcols = colshape[0] + + # Ensure mean and covariances compatible + if mean is not None: + if meanshape[0] != numrows: + raise ValueError("Arrays `mean` and `rowcov` must have the " + "same number of rows.") + if meanshape[1] != numcols: + raise ValueError("Arrays `mean` and `colcov` must have the " + "same number of columns.") + else: + mean = np.zeros((numrows, numcols)) + + dims = (numrows, numcols) + + return dims, mean, rowcov, colcov + + def _process_quantiles(self, X, dims): + """ + Adjust quantiles array so that last two axes labels the components of + each data point. + """ + X = np.asarray(X, dtype=float) + if X.ndim == 2: + X = X[np.newaxis, :] + if X.shape[-2:] != dims: + raise ValueError("The shape of array `X` is not compatible " + "with the distribution parameters.") + return X + + def _logpdf(self, dims, X, mean, row_prec_rt, log_det_rowcov, + col_prec_rt, log_det_colcov): + """Log of the matrix normal probability density function. + + Parameters + ---------- + dims : tuple + Dimensions of the matrix variates + X : ndarray + Points at which to evaluate the log of the probability + density function + mean : ndarray + Mean of the distribution + row_prec_rt : ndarray + A decomposition such that np.dot(row_prec_rt, row_prec_rt.T) + is the inverse of the among-row covariance matrix + log_det_rowcov : float + Logarithm of the determinant of the among-row covariance matrix + col_prec_rt : ndarray + A decomposition such that np.dot(col_prec_rt, col_prec_rt.T) + is the inverse of the among-column covariance matrix + log_det_colcov : float + Logarithm of the determinant of the among-column covariance matrix + + Notes + ----- + As this function does no argument checking, it should not be + called directly; use 'logpdf' instead. + + """ + numrows, numcols = dims + roll_dev = np.moveaxis(X-mean, -1, 0) + scale_dev = np.tensordot(col_prec_rt.T, + np.dot(roll_dev, row_prec_rt), 1) + maha = np.sum(np.sum(np.square(scale_dev), axis=-1), axis=0) + return -0.5 * (numrows*numcols*_LOG_2PI + numcols*log_det_rowcov + + numrows*log_det_colcov + maha) + + def logpdf(self, X, mean=None, rowcov=1, colcov=1): + """Log of the matrix normal probability density function. + + Parameters + ---------- + X : array_like + Quantiles, with the last two axes of `X` denoting the components. + %(_matnorm_doc_default_callparams)s + + Returns + ------- + logpdf : ndarray + Log of the probability density function evaluated at `X` + + Notes + ----- + %(_matnorm_doc_callparams_note)s + + """ + dims, mean, rowcov, colcov = self._process_parameters(mean, rowcov, + colcov) + X = self._process_quantiles(X, dims) + rowpsd = _PSD(rowcov, allow_singular=False) + colpsd = _PSD(colcov, allow_singular=False) + out = self._logpdf(dims, X, mean, rowpsd.U, rowpsd.log_pdet, colpsd.U, + colpsd.log_pdet) + return _squeeze_output(out) + + def pdf(self, X, mean=None, rowcov=1, colcov=1): + """Matrix normal probability density function. + + Parameters + ---------- + X : array_like + Quantiles, with the last two axes of `X` denoting the components. + %(_matnorm_doc_default_callparams)s + + Returns + ------- + pdf : ndarray + Probability density function evaluated at `X` + + Notes + ----- + %(_matnorm_doc_callparams_note)s + + """ + return np.exp(self.logpdf(X, mean, rowcov, colcov)) + + def rvs(self, mean=None, rowcov=1, colcov=1, size=1, random_state=None): + """Draw random samples from a matrix normal distribution. + + Parameters + ---------- + %(_matnorm_doc_default_callparams)s + size : int, optional + Number of samples to draw (default 1). + %(_doc_random_state)s + + Returns + ------- + rvs : ndarray or scalar + Random variates of size (`size`, `dims`), where `dims` is the + dimension of the random matrices. + + Notes + ----- + %(_matnorm_doc_callparams_note)s + + """ + size = int(size) + dims, mean, rowcov, colcov = self._process_parameters(mean, rowcov, + colcov) + rowchol = scipy.linalg.cholesky(rowcov, lower=True) + colchol = scipy.linalg.cholesky(colcov, lower=True) + random_state = self._get_random_state(random_state) + # We aren't generating standard normal variates with size=(size, + # dims[0], dims[1]) directly to ensure random variates remain backwards + # compatible. See https://github.com/scipy/scipy/pull/12312 for more + # details. + std_norm = random_state.standard_normal( + size=(dims[1], size, dims[0]) + ).transpose(1, 2, 0) + out = mean + np.einsum('jp,ipq,kq->ijk', + rowchol, std_norm, colchol, + optimize=True) + if size == 1: + out = out.reshape(mean.shape) + return out + + def entropy(self, rowcov=1, colcov=1): + """Log of the matrix normal probability density function. + + Parameters + ---------- + rowcov : array_like, optional + Among-row covariance matrix of the distribution (default: ``1``) + colcov : array_like, optional + Among-column covariance matrix of the distribution (default: ``1``) + + Returns + ------- + entropy : float + Entropy of the distribution + + Notes + ----- + %(_matnorm_doc_callparams_note)s + + """ + dummy_mean = np.zeros((rowcov.shape[0], colcov.shape[0])) + dims, _, rowcov, colcov = self._process_parameters(dummy_mean, + rowcov, + colcov) + rowpsd = _PSD(rowcov, allow_singular=False) + colpsd = _PSD(colcov, allow_singular=False) + + return self._entropy(dims, rowpsd.log_pdet, colpsd.log_pdet) + + def _entropy(self, dims, row_cov_logdet, col_cov_logdet): + n, p = dims + return (0.5 * n * p * (1 + _LOG_2PI) + 0.5 * p * row_cov_logdet + + 0.5 * n * col_cov_logdet) + + +matrix_normal = matrix_normal_gen() + + +class matrix_normal_frozen(multi_rv_frozen): + """ + Create a frozen matrix normal distribution. + + Parameters + ---------- + %(_matnorm_doc_default_callparams)s + seed : {None, int, `numpy.random.Generator`, `numpy.random.RandomState`}, optional + If `seed` is `None` the `~np.random.RandomState` singleton is used. + If `seed` is an int, a new ``RandomState`` instance is used, seeded + with seed. + If `seed` is already a ``RandomState`` or ``Generator`` instance, + then that object is used. + Default is `None`. + + Examples + -------- + >>> import numpy as np + >>> from scipy.stats import matrix_normal + + >>> distn = matrix_normal(mean=np.zeros((3,3))) + >>> X = distn.rvs(); X + array([[-0.02976962, 0.93339138, -0.09663178], + [ 0.67405524, 0.28250467, -0.93308929], + [-0.31144782, 0.74535536, 1.30412916]]) + >>> distn.pdf(X) + 2.5160642368346784e-05 + >>> distn.logpdf(X) + -10.590229595124615 + """ + __class_getitem__ = None # pyrefly:ignore[bad-assignment] + + def __init__(self, mean=None, rowcov=1, colcov=1, seed=None): + self._dist = matrix_normal_gen(seed) + self.dims, self.mean, self.rowcov, self.colcov = \ + self._dist._process_parameters(mean, rowcov, colcov) + self.rowpsd = _PSD(self.rowcov, allow_singular=False) + self.colpsd = _PSD(self.colcov, allow_singular=False) + + def logpdf(self, X): + X = self._dist._process_quantiles(X, self.dims) + out = self._dist._logpdf(self.dims, X, self.mean, self.rowpsd.U, + self.rowpsd.log_pdet, self.colpsd.U, + self.colpsd.log_pdet) + return _squeeze_output(out) + + def pdf(self, X): + return np.exp(self.logpdf(X)) + + def rvs(self, size=1, random_state=None): + return self._dist.rvs(self.mean, self.rowcov, self.colcov, size, + random_state) + + def entropy(self): + return self._dist._entropy(self.dims, self.rowpsd.log_pdet, + self.colpsd.log_pdet) + +# Set frozen generator docstrings from corresponding docstrings in +# matrix_normal_gen and fill in default strings in class docstrings +for name in ['logpdf', 'pdf', 'rvs', 'entropy']: + method = matrix_normal_gen.__dict__[name] + method_frozen = matrix_normal_frozen.__dict__[name] + method_frozen.__doc__ = doccer.docformat(method.__doc__, + matnorm_docdict_noparams) + method.__doc__ = doccer.docformat(method.__doc__, matnorm_docdict_params) + + +_matt_doc_default_callparams = """\ +mean : array_like, optional + Mean of the distribution (default: `None`) +row_spread : array_like, optional + Row-wise 2nd order raw central moment matrix (default: ``1``) +col_spread : array_like, optional + Column-wise 2nd order raw central moment matrix (default: ``1``) +df : scalar, optional + Degrees of freedom (default: ``1``) +""" + +_matt_doc_callparams_note = """\ +If `mean` is set to `None` then a matrix of zeros is used for the mean. +The dimensions of this matrix are inferred from the shape of `row_spread` and +`col_spread`, if these are provided, or set to ``1`` if ambiguous. + +`row_spread` and `col_spread` can be two-dimensional array_likes specifying the +spread matrices directly. Alternatively, a one-dimensional array will +be be interpreted as the entries of a diagonal matrix, and a scalar or +zero-dimensional array will be interpreted as this value times the +identity matrix. +""" + +_matt_doc_frozen_callparams = "" + +_matt_doc_frozen_callparams_note = """\ +See class definition for a detailed description of parameters.""" + +matrix_t_docdict_params = { + "_matt_doc_default_callparams": _matt_doc_default_callparams, + "_matt_doc_callparams_note": _matt_doc_callparams_note, + "_doc_random_state": _doc_random_state, +} + +matrix_t_docdict_noparams = { + "_matt_doc_default_callparams": _matt_doc_frozen_callparams, + "_matt_doc_callparams_note": _matt_doc_frozen_callparams_note, + "_doc_random_state": _doc_random_state, +} + + +class matrix_t_gen(multi_rv_generic): + r""" + A matrix t-random variable. + + The `mean` keyword specifies the mean. + The `row_spread` keyword specifies the row-wise spread matrix. + The `col_spread` keyword specifies the column-wise spread matrix. + + Parameters + ---------- + %(_matt_doc_default_callparams)s + %(_doc_random_state)s + + Methods + ------- + pdf(x, mean=None, row_spread=None, col_spread=None) + Probability density function. + logpdf(x, mean=None, row_spread=None, col_spread=None) + Log of the probability density function. + rvs(mean=None, row_spread=1, col_spread=1, df=1, size=1, random_state=None) + Draw random samples. + + Notes + ----- + %(_matt_doc_callparams_note)s + + The spread matrices specified by `row_spread` and `col_spread` must be + (symmetric) positive definite. If the samples in `X` have shape `(m,n)` + then `row_spread` must have shape `(m,m)` and `col_spread` must have shape `(n,n)`. + Spread matrices must be full rank. + + The probability density function for `matrix_t` is + + .. math:: + + f(X \vert \mathrm{M}, \Sigma, \Omega, \mathrm{df}) = + \frac{ + \Gamma_n \left( + \frac{\mathrm{df} + m + n - 1}{2} + \right) + \left( + \det \left( + I_n + (X - \mathrm{M})^T \Sigma^{-1} (X - \mathrm{M}) \Omega^{-1} + \right) + \right)^{ -\frac{\mathrm{df} + m + n - 1}{2} } + }{ + \Gamma_n \left( + \frac{\mathrm{df} + n - 1}{2} + \right) + \pi^{mn / 2} + \left( \det \Sigma \right)^{n/2} + \left( \det \Omega \right)^{m/2} + } + + or, alternatively, + + .. math:: + + f(X \vert \mathrm{M}, \Sigma, \Omega, \mathrm{df}) = + \frac{ + \Gamma_m \left( + \frac{\mathrm{df} + m + n - 1}{2} + \right) + \left( + \det \left( + I_m + \Sigma^{-1} (X - \mathrm{M}) \Omega^{-1} (X - \mathrm{M})^T + \right) + \right)^{ -\frac{\mathrm{df} + m + n - 1}{2} } + }{ + \Gamma_m \left( + \frac{\mathrm{df} + n - 1}{2} + \right) + \pi^{mn / 2} + \left( \det \Sigma \right)^{n/2} + \left( \det \Omega \right)^{m/2} + } + + where :math:`\mathrm{M}` is the mean, + :math:`\Sigma` is the row-wise spread matrix, + :math:`\Omega` is the column-wise matrix, + :math:`\mathrm{df}` is the degrees of freedom, + and :math:`\Gamma_n` is the multivariate gamma function. + + These equivalent formulations come from the identity + + .. math:: + + \det\left( I_m + A B \right) = \det\left( I_n + B A \right) + + for :math:`m \times n` arrays :math:`A` and :math:`B^T` + and the fact that + :math:`\gamma_n(\mathrm{df} + m) / \gamma_n(\mathrm{df})` + is equal to + :math:`\gamma_m(\mathrm{df} + n) / \gamma_m(\mathrm{df})`, + where + + .. math:: + + \gamma_m(\mathrm{df}) = 2^{m(m-1)/2} + \Gamma_m\left( (\mathrm{df} + m - 1) / 2 \right) + + denotes a normalized multivariate gamma function. + + When :math:`\mathrm{df} = 1` this distribution is known as the matrix + variate Cauchy. + + .. versionadded:: 1.17.0 + + References + ---------- + .. [1] Gupta, A.K., & Nagar, D.K. (2000). Matrix Variate Distributions (1st ed.). + Chapman and Hall/CRC. + + Examples + -------- + + >>> import numpy as np + >>> from scipy.stats import matrix_t + >>> M = np.arange(6).reshape(3,2) + >>> M + array([[0, 1], + [2, 3], + [4, 5]]) + >>> Sigma = np.diag([1,2,3]) + >>> Sigma + array([[1, 0, 0], + [0, 2, 0], + [0, 0, 3]]) + >>> Omega = 0.3*np.identity(2) + >>> Omega + array([[ 0.3, 0. ], + [ 0. , 0.3]]) + >>> X = M + 0.1 + >>> X + array([[ 0.1, 1.1], + [ 2.1, 3.1], + [ 4.1, 5.1]]) + >>> df = 3 + >>> matrix_t.pdf(X, mean=M, row_spread=Sigma, col_spread=Omega, df=df) + 0.9972880280135796 + + Alternatively, the object may be called (as a function) to fix the mean + and spread parameters, returning a "frozen" matrix t + random variable: + + >>> rv = matrix_t(mean=None, row_spread=1, col_spread=1, df=1) + >>> # Frozen object with the same methods but holding the given + >>> # mean and spreads and degrees of freedom fixed. + """ + + def __init__(self, seed=None): + super().__init__(seed) + self.__doc__ = scipy._lib.doccer.docformat( + self.__doc__, matrix_t_docdict_params + ) + + def __call__(self, mean=None, row_spread=1, col_spread=1, df=None, seed=None): + """Create a frozen matrix t distribution. + + See `matrix_t_frozen` for more information. + """ + return matrix_t_frozen(mean, row_spread, col_spread, df, seed) + + def _process_parameters(self, mean, row_spread, col_spread, df): + """ + Infer dimensionality from mean or covariance matrices. + Handle defaults. Ensure conformality. + + Parameters + ---------- + mean : ndarray, shape (m,n) + Mean of the distribution + row_spread : ndarray, shape (m,m) + Row-wise spread matrix + col_spread : ndarray, shape (n,n) + Column-wise spread matrix + df : float + Degrees of freedom + """ + + # Process mean + if mean is not None: + mean = np.asarray(mean, dtype=float) + meanshape = mean.shape + if 0 in meanshape: + raise ValueError("Array `mean` has invalid shape.") + if len(meanshape) != 2: + raise ValueError("Array `mean` must be 2D.") + + # Process row-wise spread + row_spread = np.asarray(row_spread, dtype=float) + if row_spread.ndim == 0: + if mean is not None: + row_spread = row_spread * np.identity(meanshape[0]) + else: + row_spread = row_spread * np.identity(1) + elif row_spread.ndim == 1: + row_spread = np.diag(row_spread) + rowshape = row_spread.shape + if 0 in rowshape: + raise ValueError("Array `row_spread` has invalid shape.") + if len(rowshape) != 2: + raise ValueError("Array `row_spread` must be a scalar or a 2D array.") + if rowshape[0] != rowshape[1]: + raise ValueError("Array `row_spread` must be square.") + numrows = rowshape[0] + + # Process column-wise spread + col_spread = np.asarray(col_spread, dtype=float) + if col_spread.ndim == 0: + if mean is not None: + col_spread = col_spread * np.identity(meanshape[1]) + else: + col_spread = col_spread * np.identity(1) + elif col_spread.ndim == 1: + col_spread = np.diag(col_spread) + colshape = col_spread.shape + if 0 in colshape: + raise ValueError("Array `col_spread` has invalid shape.") + if len(colshape) != 2: + raise ValueError("Array `col_spread` must be a scalar or a 2D array.") + if colshape[0] != colshape[1]: + raise ValueError("Array `col_spread` must be square.") + numcols = colshape[0] + + # Ensure mean and spreads are conformal + if mean is not None: + if meanshape[0] != numrows: + raise ValueError( + "Arrays `mean` and `row_spread` must have the same number of rows." + ) + if meanshape[1] != numcols: + raise ValueError( + "Arrays `mean` and `col_spread` must have the same number " + "of columns." + ) + else: + mean = np.zeros((numrows, numcols)) + + dims = (numrows, numcols) + + if df is None: + df = 1 # default to matrix variate Cauchy + elif not np.isscalar(df): + raise ValueError("Degrees of freedom must be a scalar.") + elif df <= 0: + raise ValueError("Degrees of freedom must be positive.") + + return dims, mean, row_spread, col_spread, df + + def _process_quantiles(self, X, dims): + """ + Adjust quantiles array so that last two axes labels the component of + each data point. + """ + X = np.asarray(X, dtype=float) + if X.ndim == 2: + X = X[np.newaxis, :] + if X.shape[-2:] != dims: + raise ValueError( + "The shape of array `X` is not conformal with " + "the distribution parameters." + ) + return X + + def _logpdf( + self, + dims, + X, + mean, + df, + invrow_spread, + invcol_spread, + logdetrow_spread, + logdetcol_spread, + ): + """ + Log of the matrix t probability density function. + + Parameters + ---------- + dims : tuple + Dimensions of the matrix variates + X : ndarray, shape (m,n) (equal to `dims`) + Points at which to evaluate the log of the probability density function + mean : ndarray, shape (m,n) + Mean of the distribution + df : float + Degrees-of-freedom parameter + invrow_spread : ndarray, shape (m,m) + Inverse of the row-wise spread matrix + invcol_spread : ndarray, shape (n,n) + Inverse of the column-wise spread matrix + logdetrow_spread : float + Log-determinant of the row-wise spread matrix + detcol_spread : float + Log-determinant of the column-wise spread matrix + + Notes + ----- + As this function does no argument checking, it should not be + called directly; use `logpdf` instead. + """ + m, n = dims + X_shape = X.shape + if X.ndim > 3: + X = X.reshape(-1, m, n) + X_centered = X - mean[np.newaxis, ...] + det_arg = np.identity(n) + np.einsum( + "nij,njk,nkl,nlp->nip", + X_centered.transpose(0, 2, 1), + invrow_spread[np.newaxis, ...], + X_centered, + invcol_spread[np.newaxis, ...], + optimize=True, + ) + _, logdet = np.linalg.slogdet(det_arg) + log_d_mn = -((df + m + n - 1) / 2) * logdet + log_f_mn = ( + scipy.special.multigammaln((df + m + n - 1) / 2, n) + - scipy.special.multigammaln((df + n - 1) / 2, n) + - (m * n / 2) * _LOG_PI + - (n / 2) * logdetrow_spread + - (m / 2) * logdetcol_spread + ) + retval = log_d_mn + log_f_mn + if len(X_shape) > 3: + retval = retval.reshape(X_shape[:-2]) + return retval + + def logpdf(self, X, mean=None, row_spread=1, col_spread=1, df=1): + """Log of the matrix normal probability density function. + + Parameters + ---------- + X : array_like + Quantiles, with the last two axes of `X` denoting the components. + %(_matt_doc_default_callparams)s + + Returns + ------- + logpdf : ndarray + Log of the probability density function evaluated at `X` + + Notes + ----- + %(_matt_doc_callparams_note)s + + Examples + ------- + + >>> import numpy as np + >>> from scipy.stats import matrix_t + >>> M = np.arange(6).reshape(3,2); M + array([[0, 1], + [2, 3], + [4, 5]]) + >>> Sigma = np.diag([1,2,3]); Sigma + array([[1, 0, 0], + [0, 2, 0], + [0, 0, 3]]) + >>> Omega = 0.3*np.identity(2); Omega + array([[ 0.3, 0. ], + [ 0. , 0.3]]) + >>> X = M + 0.1; X + array([[ 0.1, 1.1], + [ 2.1, 3.1], + [ 4.1, 5.1]]) + >>> df = 3; df + 3 + >>> matrix_t.logpdf(X, mean=M, row_spread=Sigma, col_spread=Omega, df=df) + -0.002715656044664061 + """ + dims, mean, row_spread, col_spread, df = self._process_parameters( + mean, row_spread, col_spread, df + ) + X = self._process_quantiles(X, dims) + rowpsd = _PSD(row_spread, allow_singular=False) + colpsd = _PSD(col_spread, allow_singular=False) + invrow_spread = rowpsd.pinv + invcol_spread = colpsd.pinv + logdetrow_spread = rowpsd.log_pdet + logdetcol_spread = colpsd.log_pdet + out = self._logpdf( + dims, + X, + mean, + df, + invrow_spread, + invcol_spread, + logdetrow_spread, + logdetcol_spread, + ) + return _squeeze_output(out) + + def pdf(self, X, mean=None, row_spread=1, col_spread=1, df=1): + """Matrix t probability density function. + + Parameters + ---------- + X : array_like + Quantiles, with the last two axes of `X` denoting the components. + %(_matt_doc_default_callparams)s + + Returns + ------- + pdf : ndarray + Probability density function evaluated at `X` + + Notes + ----- + %(_matt_doc_callparams_note)s + + Examples + -------- + + >>> import numpy as np + >>> from scipy.stats import matrix_t + >>> M = np.arange(6).reshape(3,2); M + array([[0, 1], + [2, 3], + [4, 5]]) + >>> Sigma = np.diag([1,2,3]); Sigma + array([[1, 0, 0], + [0, 2, 0], + [0, 0, 3]]) + >>> Omega = 0.3*np.identity(2); Omega + array([[ 0.3, 0. ], + [ 0. , 0.3]]) + >>> X = M + 0.1; X + array([[ 0.1, 1.1], + [ 2.1, 3.1], + [ 4.1, 5.1]]) + >>> df = 3; df + 3 + >>> matrix_t.logpdf(X, mean=M, row_spread=Sigma, col_spread=Omega, df=df) + 0.9972880280135796 + """ + return np.exp(self.logpdf(X, mean, row_spread, col_spread, df)) + + def rvs( + self, mean=None, row_spread=1, col_spread=1, df=1, size=1, random_state=None + ) -> np.ndarray: + """Draw random samples from a matrix t distribution. + + Parameters + ---------- + %(_matt_doc_default_callparams)s + size : int, optional + Number of samples to draw (default 1). + %(_doc_random_state)s + + Returns + ------- + rvs : ndarray or scalar + Random variates of size (`size`, `dims`), where `dims` is the + dimension of the random matrices. + + Notes + ----- + %(_matt_doc_callparams_note)s + + This method takes advantage of the two equivalent expressions of the + probability density function. It samples a Cholesky factor of a + random variate of the appropriate inverse Wishart distribution using + the smaller of the row/column dimensions. + """ + size = int(size) + dims, mean, row_spread, col_spread, df = self._process_parameters( + mean, row_spread, col_spread, df + ) + random_state = self._get_random_state(random_state) + # see scipy.stats.matrix_normal.rvs + std_norm = random_state.standard_normal( + size=(dims[1], size, dims[0]) + ).transpose(1, 2, 0) + if dims[0] <= dims[1]: + rowchol = _cholesky_invwishart_rvs(df, row_spread, size, random_state) + colchol = scipy.linalg.cholesky(col_spread, lower=True)[np.newaxis, ...] + else: + rowchol = scipy.linalg.cholesky(row_spread, lower=True)[np.newaxis, ...] + colchol = _cholesky_invwishart_rvs(df, col_spread, size, random_state) + t_raw = np.einsum("ijp,ipq,ikq->ijk", rowchol, std_norm, colchol, optimize=True) + t_centered = mean[np.newaxis, ...] + t_raw + if size == 1: + t_centered = t_centered.reshape(mean.shape) + return t_centered + + +matrix_t = matrix_t_gen() + + +class matrix_t_frozen: + def __init__(self, mean, row_spread, col_spread, df, seed=None): + self._dist = matrix_t_gen(seed) + self.dims, self.mean, self.row_spread, self.col_spread, self.df = ( + self._dist._process_parameters(mean, row_spread, col_spread, df) + ) + self._random_state = np.random.RandomState(seed) + self.rowpsd = _PSD(self.row_spread, allow_singular=False) + self.colpsd = _PSD(self.col_spread, allow_singular=False) + + def logpdf(self, X): + X = self._dist._process_quantiles(X, self.dims) + rowpsd = _PSD(self.row_spread, allow_singular=False) + colpsd = _PSD(self.col_spread, allow_singular=False) + invrow_spread = rowpsd.pinv + invcol_spread = colpsd.pinv + logdetrow_spread = rowpsd.log_pdet + logdetcol_spread = colpsd.log_pdet + out = self._dist._logpdf( + self.dims, + X, + self.mean, + self.df, + invrow_spread, + invcol_spread, + logdetrow_spread, + logdetcol_spread, + ) + return _squeeze_output(out) + + def pdf(self, X): + return np.exp(self.logpdf(X)) + + def rvs(self, size=1, random_state=None): + return self._dist.rvs( + self.mean, self.row_spread, self.col_spread, self.df, size, random_state + ) + + +# Set frozen generator docstrings from corresponding docstrings in +# matrix_t_gen and fill in default strings in class docstrings +for name in ["logpdf", "pdf", "rvs"]: + method = matrix_t_gen.__dict__[name] + method_frozen = matrix_t_frozen.__dict__[name] + method_frozen.__doc__ = scipy._lib.doccer.docformat( + method.__doc__, matrix_t_docdict_noparams + ) + method.__doc__ = scipy._lib.doccer.docformat( + method.__doc__, matrix_t_docdict_params + ) + + + +def _cholesky_invwishart_rvs( + df: float, scale: np.ndarray, size: int, random_state: np.random.Generator +) -> np.ndarray: + r"""Samples the lower Cholesky factor of a matrix following an inverse + Wishart distribution. + + Notes + ----- + Intended to be used *as a step in the process* for computing random variates + of a matrix t distribution :math:`\mathcal{T}_{m,n}` by appealing to its + alternative form as a matrix mixture + .. math:: + \mathcal{T}_{m,n}( \mathrm{df}, \mathrm{M}, \Sigma, \Omega ) + = \mathcal{N}_{m,n}( + \mathrm{M}, + \mathcal{W}^{-1}_m(\mathrm{df} + m - 1, \Sigma), + \Omega + ) + = \mathcal{N}_{m,n}( + \mathrm{M}, + \Sigma, + \mathcal{W}^{-1}_n(\mathrm{df} - n + 1, \Omega) + ) + where :math:`\mathcal{N}_{m,n}` is a matrix normal distribution + and :math:`\mathcal{W}^{-1}_d` is an inverse Wishart distribution. + Accordingly, the degrees of freedom adjustment + :math:`\mathrm{df} \to \mathrm{df} + d - 1` + occurrs in the scope of this function. + """ + df_iw = df + scale.shape[0] - 1 + iw_samples = scipy.stats.invwishart.rvs(df_iw, scale, size, random_state) + if size == 1: + iw_samples = iw_samples[np.newaxis, ...] + chol_samples = np.empty_like(iw_samples) + for idx in range(size): + chol_samples[idx] = scipy.linalg.cholesky( + iw_samples[idx], lower=True, check_finite=False + ).reshape(iw_samples.shape[1:]) + return chol_samples.reshape((size, *scale.shape)) + + +_dirichlet_doc_default_callparams = """\ +alpha : array_like + The concentration parameters. The number of entries determines the + dimensionality of the distribution. +""" +_dirichlet_doc_frozen_callparams = "" + +_dirichlet_doc_frozen_callparams_note = """\ +See class definition for a detailed description of parameters.""" + +dirichlet_docdict_params = { + '_dirichlet_doc_default_callparams': _dirichlet_doc_default_callparams, + '_doc_random_state': _doc_random_state +} + +dirichlet_docdict_noparams = { + '_dirichlet_doc_default_callparams': _dirichlet_doc_frozen_callparams, + '_doc_random_state': _doc_random_state +} + + +def _dirichlet_check_parameters(alpha): + alpha = np.asarray(alpha) + if np.min(alpha) <= 0: + raise ValueError("All parameters must be greater than 0") + elif alpha.ndim != 1: + raise ValueError("Parameter vector 'a' must be one dimensional, " + f"but a.shape = {alpha.shape}.") + return alpha + + +def _dirichlet_check_input(alpha, x): + x = np.asarray(x) + + if x.shape[0] + 1 != alpha.shape[0] and x.shape[0] != alpha.shape[0]: + raise ValueError("Vector 'x' must have either the same number " + "of entries as, or one entry fewer than, " + f"parameter vector 'a', but alpha.shape = {alpha.shape} " + f"and x.shape = {x.shape}.") + + if x.shape[0] != alpha.shape[0]: + xk = np.array([1 - np.sum(x, 0)]) + if xk.ndim == 1: + x = np.append(x, xk) + elif xk.ndim == 2: + x = np.vstack((x, xk)) + else: + raise ValueError("The input must be one dimensional or a two " + "dimensional matrix containing the entries.") + + if np.min(x) < 0: + raise ValueError("Each entry in 'x' must be greater than or equal " + "to zero.") + + if np.max(x) > 1: + raise ValueError("Each entry in 'x' must be smaller or equal one.") + + # Check x_i > 0 or alpha_i > 1 + xeq0 = (x == 0) + alphalt1 = (alpha < 1) + if x.shape != alpha.shape: + alphalt1 = np.repeat(alphalt1, x.shape[-1], axis=-1).reshape(x.shape) + chk = np.logical_and(xeq0, alphalt1) + + if np.sum(chk): + raise ValueError("Each entry in 'x' must be greater than zero if its " + "alpha is less than one.") + + if (np.abs(np.sum(x, 0) - 1.0) > 10e-10).any(): + raise ValueError("The input vector 'x' must lie within the normal " + f"simplex. but np.sum(x, 0) = {np.sum(x, 0)}.") + + return x + + +def _lnB(alpha): + r"""Internal helper function to compute the log of the useful quotient. + + .. math:: + + B(\alpha) = \frac{\prod_{i=1}{K}\Gamma(\alpha_i)} + {\Gamma\left(\sum_{i=1}^{K} \alpha_i \right)} + + Parameters + ---------- + %(_dirichlet_doc_default_callparams)s + + Returns + ------- + B : scalar + Helper quotient, internal use only + + """ + return np.sum(gammaln(alpha)) - gammaln(np.sum(alpha)) + + +class dirichlet_gen(multi_rv_generic): + r""" + A Dirichlet random variable. + + The ``alpha`` keyword specifies the concentration parameters of the + distribution. + + .. versionadded:: 0.15.0 + + Parameters + ---------- + %(_dirichlet_doc_default_callparams)s + %(_doc_random_state)s + + Methods + ------- + pdf(x, alpha) + Probability density function. + logpdf(x, alpha) + Log of the probability density function. + rvs(alpha, size=1, random_state=None) + Draw random samples from a Dirichlet distribution. + mean(alpha) + The mean of the Dirichlet distribution + var(alpha) + The variance of the Dirichlet distribution + cov(alpha) + The covariance of the Dirichlet distribution + entropy(alpha) + Compute the differential entropy of the Dirichlet distribution. + + Notes + ----- + Each :math:`\alpha` entry must be positive. The distribution has only + support on the simplex defined by + + .. math:: + \sum_{i=1}^{K} x_i = 1 + + where :math:`0 < x_i < 1`. + + If the quantiles don't lie within the simplex, a ValueError is raised. + + The probability density function for `dirichlet` is + + .. math:: + + f(x) = \frac{1}{\mathrm{B}(\boldsymbol\alpha)} \prod_{i=1}^K x_i^{\alpha_i - 1} + + where + + .. math:: + + \mathrm{B}(\boldsymbol\alpha) = \frac{\prod_{i=1}^K \Gamma(\alpha_i)} + {\Gamma\bigl(\sum_{i=1}^K \alpha_i\bigr)} + + and :math:`\boldsymbol\alpha=(\alpha_1,\ldots,\alpha_K)`, the + concentration parameters and :math:`K` is the dimension of the space + where :math:`x` takes values. + + Note that the `dirichlet` interface is somewhat inconsistent. + The array returned by the rvs function is transposed + with respect to the format expected by the pdf and logpdf. + + Examples + -------- + >>> import numpy as np + >>> from scipy.stats import dirichlet + + Generate a dirichlet random variable + + >>> quantiles = np.array([0.2, 0.2, 0.6]) # specify quantiles + >>> alpha = np.array([0.4, 5, 15]) # specify concentration parameters + >>> dirichlet.pdf(quantiles, alpha) + 0.2843831684937255 + + The same PDF but following a log scale + + >>> dirichlet.logpdf(quantiles, alpha) + -1.2574327653159187 + + Once we specify the dirichlet distribution + we can then calculate quantities of interest + + >>> dirichlet.mean(alpha) # get the mean of the distribution + array([0.01960784, 0.24509804, 0.73529412]) + >>> dirichlet.var(alpha) # get variance + array([0.00089829, 0.00864603, 0.00909517]) + >>> dirichlet.entropy(alpha) # calculate the differential entropy + -4.3280162474082715 + + We can also return random samples from the distribution + + >>> dirichlet.rvs(alpha, size=1, random_state=1) + array([[0.00766178, 0.24670518, 0.74563305]]) + >>> dirichlet.rvs(alpha, size=2, random_state=2) + array([[0.01639427, 0.1292273 , 0.85437844], + [0.00156917, 0.19033695, 0.80809388]]) + + Alternatively, the object may be called (as a function) to fix + concentration parameters, returning a "frozen" Dirichlet + random variable: + + >>> rv = dirichlet(alpha) + >>> # Frozen object with the same methods but holding the given + >>> # concentration parameters fixed. + """ + + def __init__(self, seed=None): + super().__init__(seed) + self.__doc__ = doccer.docformat(self.__doc__, dirichlet_docdict_params) + + def __call__(self, alpha, seed=None): + return dirichlet_frozen(alpha, seed=seed) + + def _logpdf(self, x, alpha): + """Log of the Dirichlet probability density function. + + Parameters + ---------- + x : ndarray + Points at which to evaluate the log of the probability + density function + %(_dirichlet_doc_default_callparams)s + + Notes + ----- + As this function does no argument checking, it should not be + called directly; use 'logpdf' instead. + + """ + lnB = _lnB(alpha) + return - lnB + np.sum((xlogy(alpha - 1, x.T)).T, 0) + + def logpdf(self, x, alpha): + """Log of the Dirichlet probability density function. + + Parameters + ---------- + x : array_like + Quantiles, with the last axis of `x` denoting the components. + %(_dirichlet_doc_default_callparams)s + + Returns + ------- + pdf : ndarray or scalar + Log of the probability density function evaluated at `x`. + + """ + alpha = _dirichlet_check_parameters(alpha) + x = _dirichlet_check_input(alpha, x) + + out = self._logpdf(x, alpha) + return _squeeze_output(out) + + def pdf(self, x, alpha): + """The Dirichlet probability density function. + + Parameters + ---------- + x : array_like + Quantiles, with the last axis of `x` denoting the components. + %(_dirichlet_doc_default_callparams)s + + Returns + ------- + pdf : ndarray or scalar + The probability density function evaluated at `x`. + + """ + alpha = _dirichlet_check_parameters(alpha) + x = _dirichlet_check_input(alpha, x) + + out = np.exp(self._logpdf(x, alpha)) + return _squeeze_output(out) + + def mean(self, alpha): + """Mean of the Dirichlet distribution. + + Parameters + ---------- + %(_dirichlet_doc_default_callparams)s + + Returns + ------- + mu : ndarray or scalar + Mean of the Dirichlet distribution. + + """ + alpha = _dirichlet_check_parameters(alpha) + + out = alpha / (np.sum(alpha)) + return _squeeze_output(out) + + def var(self, alpha): + """Variance of the Dirichlet distribution. + + Parameters + ---------- + %(_dirichlet_doc_default_callparams)s + + Returns + ------- + v : ndarray or scalar + Variance of the Dirichlet distribution. + + """ + + alpha = _dirichlet_check_parameters(alpha) + + alpha0 = np.sum(alpha) + out = (alpha * (alpha0 - alpha)) / ((alpha0 * alpha0) * (alpha0 + 1)) + return _squeeze_output(out) + + def cov(self, alpha): + """Covariance matrix of the Dirichlet distribution. + + Parameters + ---------- + %(_dirichlet_doc_default_callparams)s + + Returns + ------- + cov : ndarray + The covariance matrix of the distribution. + """ + + alpha = _dirichlet_check_parameters(alpha) + alpha0 = np.sum(alpha) + a = alpha / alpha0 + + cov = (np.diag(a) - np.outer(a, a)) / (alpha0 + 1) + return _squeeze_output(cov) + + def entropy(self, alpha): + """ + Differential entropy of the Dirichlet distribution. + + Parameters + ---------- + %(_dirichlet_doc_default_callparams)s + + Returns + ------- + h : scalar + Entropy of the Dirichlet distribution + + """ + + alpha = _dirichlet_check_parameters(alpha) + + alpha0 = np.sum(alpha) + lnB = _lnB(alpha) + K = alpha.shape[0] + + out = lnB + (alpha0 - K) * scipy.special.psi(alpha0) - np.sum( + (alpha - 1) * scipy.special.psi(alpha)) + return _squeeze_output(out) + + def rvs(self, alpha, size=1, random_state=None): + """ + Draw random samples from a Dirichlet distribution. + + Parameters + ---------- + %(_dirichlet_doc_default_callparams)s + size : int, optional + Number of samples to draw (default 1). + %(_doc_random_state)s + + Returns + ------- + rvs : ndarray or scalar + Random variates of size (`size`, `N`), where `N` is the + dimension of the random variable. + + """ + alpha = _dirichlet_check_parameters(alpha) + random_state = self._get_random_state(random_state) + return random_state.dirichlet(alpha, size=size) + + +dirichlet = dirichlet_gen() + + +class dirichlet_frozen(multi_rv_frozen): + __class_getitem__ = None # pyrefly:ignore[bad-assignment] + + def __init__(self, alpha, seed=None): + self.alpha = _dirichlet_check_parameters(alpha) + self._dist = dirichlet_gen(seed) + + def logpdf(self, x): + return self._dist.logpdf(x, self.alpha) + + def pdf(self, x): + return self._dist.pdf(x, self.alpha) + + def mean(self): + return self._dist.mean(self.alpha) + + def var(self): + return self._dist.var(self.alpha) + + def cov(self): + return self._dist.cov(self.alpha) + + def entropy(self): + return self._dist.entropy(self.alpha) + + def rvs(self, size=1, random_state=None): + return self._dist.rvs(self.alpha, size, random_state) + + +# Set frozen generator docstrings from corresponding docstrings in +# multivariate_normal_gen and fill in default strings in class docstrings +for name in ['logpdf', 'pdf', 'rvs', 'mean', 'var', 'cov', 'entropy']: + method = dirichlet_gen.__dict__[name] + method_frozen = dirichlet_frozen.__dict__[name] + method_frozen.__doc__ = doccer.docformat( + method.__doc__, dirichlet_docdict_noparams) + method.__doc__ = doccer.docformat(method.__doc__, dirichlet_docdict_params) + + +_wishart_doc_default_callparams = """\ +df : int + Degrees of freedom, must be greater than or equal to dimension of the + scale matrix +scale : array_like + Symmetric positive definite scale matrix of the distribution +""" + +_wishart_doc_callparams_note = "" + +_wishart_doc_frozen_callparams = "" + +_wishart_doc_frozen_callparams_note = """\ +See class definition for a detailed description of parameters.""" + +wishart_docdict_params = { + '_doc_default_callparams': _wishart_doc_default_callparams, + '_doc_callparams_note': _wishart_doc_callparams_note, + '_doc_random_state': _doc_random_state +} + +wishart_docdict_noparams = { + '_doc_default_callparams': _wishart_doc_frozen_callparams, + '_doc_callparams_note': _wishart_doc_frozen_callparams_note, + '_doc_random_state': _doc_random_state +} + + +class wishart_gen(multi_rv_generic): + r""" + A Wishart random variable. + + The `df` keyword specifies the degrees of freedom. The `scale` keyword + specifies the scale matrix, which must be symmetric and positive definite. + In this context, the scale matrix is often interpreted in terms of a + multivariate normal precision matrix (the inverse of the covariance + matrix). These arguments must satisfy the relationship + ``df > scale.ndim - 1``, but see notes on using the `rvs` method with + ``df < scale.ndim``. + + Parameters + ---------- + %(_doc_default_callparams)s + %(_doc_random_state)s + + Methods + ------- + pdf(x, df, scale) + Probability density function. + logpdf(x, df, scale) + Log of the probability density function. + rvs(df, scale, size=1, random_state=None) + Draw random samples from a Wishart distribution. + entropy() + Compute the differential entropy of the Wishart distribution. + + Raises + ------ + scipy.linalg.LinAlgError + If the scale matrix `scale` is not positive definite. + + See Also + -------- + invwishart, chi2 + + Notes + ----- + %(_doc_callparams_note)s + + The scale matrix `scale` must be a symmetric positive definite + matrix. Singular matrices, including the symmetric positive semi-definite + case, are not supported. Symmetry is not checked; only the lower triangular + portion is used. + + The Wishart distribution is often denoted + + .. math:: + + W_p(\nu, \Sigma) + + where :math:`\nu` is the degrees of freedom and :math:`\Sigma` is the + :math:`p \times p` scale matrix. + + The probability density function for `wishart` has support over positive + definite matrices :math:`S`; if :math:`S \sim W_p(\nu, \Sigma)`, then + its PDF is given by: + + .. math:: + + f(S) = \frac{|S|^{\frac{\nu - p - 1}{2}}}{2^{ \frac{\nu p}{2} } + |\Sigma|^\frac{\nu}{2} \Gamma_p \left ( \frac{\nu}{2} \right )} + \exp\left( -tr(\Sigma^{-1} S) / 2 \right) + + If :math:`S \sim W_p(\nu, \Sigma)` (Wishart) then + :math:`S^{-1} \sim W_p^{-1}(\nu, \Sigma^{-1})` (inverse Wishart). + + If the scale matrix is 1-dimensional and equal to one, then the Wishart + distribution :math:`W_1(\nu, 1)` collapses to the :math:`\chi^2(\nu)` + distribution. + + The algorithm [2]_ implemented by the `rvs` method may + produce numerically singular matrices with :math:`p - 1 < \nu < p`; the + user may wish to check for this condition and generate replacement samples + as necessary. + + + .. versionadded:: 0.16.0 + + References + ---------- + .. [1] M.L. Eaton, "Multivariate Statistics: A Vector Space Approach", + Wiley, 1983. + .. [2] W.B. Smith and R.R. Hocking, "Algorithm AS 53: Wishart Variate + Generator", Applied Statistics, vol. 21, pp. 341-345, 1972. + + Examples + -------- + >>> import numpy as np + >>> import matplotlib.pyplot as plt + >>> from scipy.stats import wishart, chi2 + >>> x = np.linspace(1e-5, 8, 100) + >>> w = wishart.pdf(x, df=3, scale=1); w[:5] + array([ 0.00126156, 0.10892176, 0.14793434, 0.17400548, 0.1929669 ]) + >>> c = chi2.pdf(x, 3); c[:5] + array([ 0.00126156, 0.10892176, 0.14793434, 0.17400548, 0.1929669 ]) + >>> plt.plot(x, w) + >>> plt.show() + + The input quantiles can be any shape of array, as long as the last + axis labels the components. + + Alternatively, the object may be called (as a function) to fix the degrees + of freedom and scale parameters, returning a "frozen" Wishart random + variable: + + >>> rv = wishart(df=1, scale=1) + >>> # Frozen object with the same methods but holding the given + >>> # degrees of freedom and scale fixed. + """ + + def __init__(self, seed=None): + super().__init__(seed) + self.__doc__ = doccer.docformat(self.__doc__, wishart_docdict_params) + + def __call__(self, df=None, scale=None, seed=None): + """Create a frozen Wishart distribution. + + See `wishart_frozen` for more information. + """ + return wishart_frozen(df, scale, seed) + + def _process_parameters(self, df, scale): + if scale is None: + scale = 1.0 + scale = np.asarray(scale, dtype=float) + + if scale.ndim == 0: + scale = scale[np.newaxis, np.newaxis] + elif scale.ndim == 1: + scale = np.diag(scale) + elif scale.ndim == 2 and not scale.shape[0] == scale.shape[1]: + raise ValueError("Array 'scale' must be square if it is two dimensional," + f" but scale.scale = {str(scale.shape)}.") + elif scale.ndim > 2: + raise ValueError(f"Array 'scale' must be at most two-dimensional, " + f"but scale.ndim = {scale.ndim}") + + dim = scale.shape[0] + + if df is None: + df = dim + elif not np.isscalar(df): + raise ValueError("Degrees of freedom must be a scalar.") + elif df <= dim - 1: + raise ValueError("Degrees of freedom must be greater than the " + "dimension of scale matrix minus 1.") + + return dim, df, scale + + def _process_quantiles(self, x, dim): + """ + Adjust quantiles array so that last axis labels the components of + each data point. + """ + x = np.asarray(x, dtype=float) + + if x.ndim == 0: + x = x * np.eye(dim)[:, :, np.newaxis] + if x.ndim == 1: + if dim == 1: + x = x[np.newaxis, np.newaxis, :] + else: + x = np.diag(x)[:, :, np.newaxis] + elif x.ndim == 2: + if not x.shape[0] == x.shape[1]: + raise ValueError( + "Quantiles must be square if they are two dimensional," + f" but x.shape = {str(x.shape)}.") + x = x[:, :, np.newaxis] + elif x.ndim == 3: + if not x.shape[0] == x.shape[1]: + raise ValueError( + "Quantiles must be square in the first two dimensions " + f"if they are three dimensional, but x.shape = {str(x.shape)}.") + elif x.ndim > 3: + raise ValueError(f"Quantiles must be at most two-dimensional with an " + f"additional dimension for multiple components, " + f"but x.ndim = {x.ndim}") + + # Now we have 3-dim array; should have shape [dim, dim, *] + if not x.shape[0:2] == (dim, dim): + raise ValueError('Quantiles have incompatible dimensions: should' + f' be {(dim, dim)}, got {x.shape[0:2]}.') + + return x + + def _process_size(self, size): + size = np.asarray(size) + + if size.ndim == 0: + size = size[np.newaxis] + elif size.ndim > 1: + raise ValueError('Size must be an integer or tuple of integers;' + ' thus must have dimension <= 1.' + f' Got size.ndim = {str(tuple(size))}') + n = size.prod() + shape = tuple(size) + + return n, shape + + def _logpdf(self, x, dim, df, scale, log_det_scale, C): + """Log of the Wishart probability density function. + + Parameters + ---------- + x : ndarray + Points at which to evaluate the log of the probability + density function + dim : int + Dimension of the scale matrix + df : int + Degrees of freedom + scale : ndarray + Scale matrix + log_det_scale : float + Logarithm of the determinant of the scale matrix + C : ndarray + Cholesky factorization of the scale matrix, lower triangular. + + Notes + ----- + As this function does no argument checking, it should not be + called directly; use 'logpdf' instead. + + """ + # log determinant of x + # Note: x has components along the last axis, so that x.T has + # components alone the 0-th axis. Then since det(A) = det(A'), this + # gives us a 1-dim vector of determinants + + # Retrieve tr(scale^{-1} x) + log_det_x = np.empty(x.shape[-1]) + scale_inv_x = np.empty(x.shape) + tr_scale_inv_x = np.empty(x.shape[-1]) + for i in range(x.shape[-1]): + _, log_det_x[i] = self._cholesky_logdet(x[:, :, i]) + scale_inv_x[:, :, i] = scipy.linalg.cho_solve((C, True), x[:, :, i]) + tr_scale_inv_x[i] = scale_inv_x[:, :, i].trace() + + # Log PDF + out = ((0.5 * (df - dim - 1) * log_det_x - 0.5 * tr_scale_inv_x) - + (0.5 * df * dim * _LOG_2 + 0.5 * df * log_det_scale + + multigammaln(0.5*df, dim))) + + return out + + def logpdf(self, x, df, scale): + """Log of the Wishart probability density function. + + Parameters + ---------- + x : array_like + Quantiles, with the last axis of `x` denoting the components. + Each quantile must be a symmetric positive definite matrix. + %(_doc_default_callparams)s + + Returns + ------- + pdf : ndarray + Log of the probability density function evaluated at `x` + + Notes + ----- + %(_doc_callparams_note)s + + """ + dim, df, scale = self._process_parameters(df, scale) + x = self._process_quantiles(x, dim) + + # Cholesky decomposition of scale, get log(det(scale)) + C, log_det_scale = self._cholesky_logdet(scale) + + out = self._logpdf(x, dim, df, scale, log_det_scale, C) + return _squeeze_output(out) + + def pdf(self, x, df, scale): + """Wishart probability density function. + + Parameters + ---------- + x : array_like + Quantiles, with the last axis of `x` denoting the components. + Each quantile must be a symmetric positive definite matrix. + %(_doc_default_callparams)s + + Returns + ------- + pdf : ndarray + Probability density function evaluated at `x` + + Notes + ----- + %(_doc_callparams_note)s + + """ + return np.exp(self.logpdf(x, df, scale)) + + def _mean(self, dim, df, scale): + """Mean of the Wishart distribution. + + Parameters + ---------- + dim : int + Dimension of the scale matrix + %(_doc_default_callparams)s + + Notes + ----- + As this function does no argument checking, it should not be + called directly; use 'mean' instead. + + """ + return df * scale + + def mean(self, df, scale): + """Mean of the Wishart distribution. + + Parameters + ---------- + %(_doc_default_callparams)s + + Returns + ------- + mean : float + The mean of the distribution + """ + dim, df, scale = self._process_parameters(df, scale) + out = self._mean(dim, df, scale) + return _squeeze_output(out) + + def _mode(self, dim, df, scale): + """Mode of the Wishart distribution. + + Parameters + ---------- + dim : int + Dimension of the scale matrix + %(_doc_default_callparams)s + + Notes + ----- + As this function does no argument checking, it should not be + called directly; use 'mode' instead. + + """ + if df >= dim + 1: + out = (df-dim-1) * scale + else: + out = None + return out + + def mode(self, df, scale): + """Mode of the Wishart distribution + + Only valid if the degrees of freedom are greater than the dimension of + the scale matrix. + + Parameters + ---------- + %(_doc_default_callparams)s + + Returns + ------- + mode : float or None + The Mode of the distribution + """ + dim, df, scale = self._process_parameters(df, scale) + out = self._mode(dim, df, scale) + return _squeeze_output(out) if out is not None else out + + def _var(self, dim, df, scale): + """Variance of the Wishart distribution. + + Parameters + ---------- + dim : int + Dimension of the scale matrix + %(_doc_default_callparams)s + + Notes + ----- + As this function does no argument checking, it should not be + called directly; use 'var' instead. + + """ + var = scale**2 + diag = scale.diagonal() # 1 x dim array + var += np.outer(diag, diag) + var *= df + return var + + def var(self, df, scale): + """Variance of the Wishart distribution. + + Parameters + ---------- + %(_doc_default_callparams)s + + Returns + ------- + var : float + The variance of the distribution + """ + dim, df, scale = self._process_parameters(df, scale) + out = self._var(dim, df, scale) + return _squeeze_output(out) + + def _standard_rvs(self, n, shape, dim, df, random_state): + """ + Parameters + ---------- + n : int + Number of variates to generate + shape : iterable + Shape of the variates to generate + dim : int + Dimension of the scale matrix + df : int + Degrees of freedom + random_state : {None, int, `numpy.random.Generator`, + `numpy.random.RandomState`}, optional + + If `seed` is None (or `np.random`), the `numpy.random.RandomState` + singleton is used. + If `seed` is an int, a new ``RandomState`` instance is used, + seeded with `seed`. + If `seed` is already a ``Generator`` or ``RandomState`` instance + then that instance is used. + + Notes + ----- + As this function does no argument checking, it should not be + called directly; use 'rvs' instead. + + """ + # Random normal variates for off-diagonal elements + n_tril = dim * (dim-1) // 2 + covariances = random_state.normal( + size=n*n_tril).reshape(shape+(n_tril,)) + + # Random chi-square variates for diagonal elements + variances = (np.r_[[random_state.chisquare(df-(i+1)+1, size=n)**0.5 + for i in range(dim)]].reshape((dim,) + + shape[::-1]).T) + + # Create the A matri(ces) - lower triangular + A = np.zeros(shape + (dim, dim)) + + # Input the covariances + size_idx = tuple([slice(None, None, None)]*len(shape)) + tril_idx = np.tril_indices(dim, k=-1) + A[size_idx + tril_idx] = covariances + + # Input the variances + diag_idx = np.diag_indices(dim) + A[size_idx + diag_idx] = variances + + return A + + def _rvs(self, n, shape, dim, df, C, random_state): + """Draw random samples from a Wishart distribution. + + Parameters + ---------- + n : int + Number of variates to generate + shape : iterable + Shape of the variates to generate + dim : int + Dimension of the scale matrix + df : int + Degrees of freedom + C : ndarray + Cholesky factorization of the scale matrix, lower triangular. + %(_doc_random_state)s + + Notes + ----- + As this function does no argument checking, it should not be + called directly; use 'rvs' instead. + + """ + random_state = self._get_random_state(random_state) + # Calculate the matrices A, which are actually lower triangular + # Cholesky factorizations of a matrix B such that B ~ W(df, I) + A = self._standard_rvs(n, shape, dim, df, random_state) + + # Calculate SA = C A A' C', where SA ~ W(df, scale) + # Note: this is the product of a (lower) (lower) (lower)' (lower)' + # or, denoting B = AA', it is C B C' where C is the lower + # triangular Cholesky factorization of the scale matrix. + # this appears to conflict with the instructions in [1]_, which + # suggest that it should be D' B D where D is the lower + # triangular factorization of the scale matrix. However, it is + # meant to refer to the Bartlett (1933) representation of a + # Wishart random variate as L A A' L' where L is lower triangular + # so it appears that understanding D' to be upper triangular + # is either a typo in or misreading of [1]_. + for index in np.ndindex(shape): + CA = np.dot(C, A[index]) + A[index] = np.dot(CA, CA.T) + + return A + + def rvs(self, df, scale, size=1, random_state=None): + """Draw random samples from a Wishart distribution. + + Parameters + ---------- + %(_doc_default_callparams)s + size : int or iterable of integers, optional + Number of samples to draw (default 1). + %(_doc_random_state)s + + Returns + ------- + rvs : ndarray + Random variates of shape (`size`) + (``dim``, ``dim``), where + ``dim`` is the dimension of the scale matrix. + + Notes + ----- + %(_doc_callparams_note)s + + """ + n, shape = self._process_size(size) + dim, df, scale = self._process_parameters(df, scale) + + # Cholesky decomposition of scale + C = scipy.linalg.cholesky(scale, lower=True) + + out = self._rvs(n, shape, dim, df, C, random_state) + + return _squeeze_output(out) + + def _entropy(self, dim, df, log_det_scale): + """Compute the differential entropy of the Wishart. + + Parameters + ---------- + dim : int + Dimension of the scale matrix + df : int + Degrees of freedom + log_det_scale : float + Logarithm of the determinant of the scale matrix + + Notes + ----- + As this function does no argument checking, it should not be + called directly; use 'entropy' instead. + + """ + return ( + 0.5 * (dim+1) * log_det_scale + + 0.5 * dim * (dim+1) * _LOG_2 + + multigammaln(0.5*df, dim) - + 0.5 * (df - dim - 1) * np.sum( + [psi(0.5*(df + 1 - (i+1))) for i in range(dim)] + ) + + 0.5 * df * dim + ) + + def entropy(self, df, scale): + """Compute the differential entropy of the Wishart. + + Parameters + ---------- + %(_doc_default_callparams)s + + Returns + ------- + h : scalar + Entropy of the Wishart distribution + + Notes + ----- + %(_doc_callparams_note)s + + """ + dim, df, scale = self._process_parameters(df, scale) + _, log_det_scale = self._cholesky_logdet(scale) + return self._entropy(dim, df, log_det_scale) + + def _cholesky_logdet(self, scale): + """Compute Cholesky decomposition and determine (log(det(scale)). + + Parameters + ---------- + scale : ndarray + Scale matrix. + + Returns + ------- + c_decomp : ndarray + The Cholesky decomposition of `scale`. + logdet : scalar + The log of the determinant of `scale`. + + Notes + ----- + This computation of ``logdet`` is equivalent to + ``np.linalg.slogdet(scale)``. It is ~2x faster though. + + """ + c_decomp = scipy.linalg.cholesky(scale, lower=True) + logdet = 2 * np.sum(np.log(c_decomp.diagonal())) + return c_decomp, logdet + + +wishart = wishart_gen() + + +class wishart_frozen(multi_rv_frozen): + """Create a frozen Wishart distribution. + + Parameters + ---------- + df : array_like + Degrees of freedom of the distribution + scale : array_like + Scale matrix of the distribution + seed : {None, int, `numpy.random.Generator`, `numpy.random.RandomState`}, optional + If `seed` is None (or `np.random`), the `numpy.random.RandomState` + singleton is used. + If `seed` is an int, a new ``RandomState`` instance is used, + seeded with `seed`. + If `seed` is already a ``Generator`` or ``RandomState`` instance then + that instance is used. + + """ + __class_getitem__ = None # pyrefly:ignore[bad-assignment] + + def __init__(self, df, scale, seed=None): + self._dist = wishart_gen(seed) + self.dim, self.df, self.scale = self._dist._process_parameters( + df, scale) + self.C, self.log_det_scale = self._dist._cholesky_logdet(self.scale) + + def logpdf(self, x): + x = self._dist._process_quantiles(x, self.dim) + + out = self._dist._logpdf(x, self.dim, self.df, self.scale, + self.log_det_scale, self.C) + return _squeeze_output(out) + + def pdf(self, x): + return np.exp(self.logpdf(x)) + + def mean(self): + out = self._dist._mean(self.dim, self.df, self.scale) + return _squeeze_output(out) + + def mode(self): + out = self._dist._mode(self.dim, self.df, self.scale) + return _squeeze_output(out) if out is not None else out + + def var(self): + out = self._dist._var(self.dim, self.df, self.scale) + return _squeeze_output(out) + + def rvs(self, size=1, random_state=None): + n, shape = self._dist._process_size(size) + out = self._dist._rvs(n, shape, self.dim, self.df, + self.C, random_state) + return _squeeze_output(out) + + def entropy(self): + return self._dist._entropy(self.dim, self.df, self.log_det_scale) + + +# Set frozen generator docstrings from corresponding docstrings in +# Wishart and fill in default strings in class docstrings +for name in ['logpdf', 'pdf', 'mean', 'mode', 'var', 'rvs', 'entropy']: + method = wishart_gen.__dict__[name] + method_frozen = wishart_frozen.__dict__[name] + method_frozen.__doc__ = doccer.docformat( + method.__doc__, wishart_docdict_noparams) + method.__doc__ = doccer.docformat(method.__doc__, wishart_docdict_params) + + +class invwishart_gen(wishart_gen): + r""" + An inverse Wishart random variable. + + The `df` keyword specifies the degrees of freedom. The `scale` keyword + specifies the scale matrix, which must be symmetric and positive definite. + In this context, the scale matrix is often interpreted in terms of a + multivariate normal covariance matrix. + + Parameters + ---------- + %(_doc_default_callparams)s + %(_doc_random_state)s + + Methods + ------- + pdf(x, df, scale) + Probability density function. + logpdf(x, df, scale) + Log of the probability density function. + rvs(df, scale, size=1, random_state=None) + Draw random samples from an inverse Wishart distribution. + entropy(df, scale) + Differential entropy of the distribution. + + Raises + ------ + scipy.linalg.LinAlgError + If the scale matrix `scale` is not positive definite. + + See Also + -------- + wishart + + Notes + ----- + %(_doc_callparams_note)s + + The scale matrix `scale` must be a symmetric positive definite + matrix. Singular matrices, including the symmetric positive semi-definite + case, are not supported. Symmetry is not checked; only the lower triangular + portion is used. + + The inverse Wishart distribution is often denoted + + .. math:: + + W_p^{-1}(\nu, \Psi) + + where :math:`\nu` is the degrees of freedom and :math:`\Psi` is the + :math:`p \times p` scale matrix. + + The probability density function for `invwishart` has support over positive + definite matrices :math:`S`; if :math:`S \sim W^{-1}_p(\nu, \Sigma)`, + then its PDF is given by: + + .. math:: + + f(S) = \frac{|\Sigma|^\frac{\nu}{2}}{2^{ \frac{\nu p}{2} } + |S|^{\frac{\nu + p + 1}{2}} \Gamma_p \left(\frac{\nu}{2} \right)} + \exp\left( -tr(\Sigma S^{-1}) / 2 \right) + + If :math:`S \sim W_p^{-1}(\nu, \Psi)` (inverse Wishart) then + :math:`S^{-1} \sim W_p(\nu, \Psi^{-1})` (Wishart). + + If the scale matrix is 1-dimensional and equal to one, then the inverse + Wishart distribution :math:`W_1(\nu, 1)` collapses to the + inverse Gamma distribution with parameters shape = :math:`\frac{\nu}{2}` + and scale = :math:`\frac{1}{2}`. + + Instead of inverting a randomly generated Wishart matrix as described in [2], + here the algorithm in [4] is used to directly generate a random inverse-Wishart + matrix without inversion. + + .. versionadded:: 0.16.0 + + References + ---------- + .. [1] M.L. Eaton, "Multivariate Statistics: A Vector Space Approach", + Wiley, 1983. + .. [2] M.C. Jones, "Generating Inverse Wishart Matrices", Communications + in Statistics - Simulation and Computation, vol. 14.2, pp.511-514, + 1985. + .. [3] Gupta, M. and Srivastava, S. "Parametric Bayesian Estimation of + Differential Entropy and Relative Entropy". Entropy 12, 818 - 843. + 2010. + .. [4] S.D. Axen, "Efficiently generating inverse-Wishart matrices and + their Cholesky factors", :arXiv:`2310.15884v1`. 2023. + + Examples + -------- + >>> import numpy as np + >>> import matplotlib.pyplot as plt + >>> from scipy.stats import invwishart, invgamma + >>> x = np.linspace(0.01, 1, 100) + >>> iw = invwishart.pdf(x, df=6, scale=1) + >>> iw[:3] + array([ 1.20546865e-15, 5.42497807e-06, 4.45813929e-03]) + >>> ig = invgamma.pdf(x, 6/2., scale=1./2) + >>> ig[:3] + array([ 1.20546865e-15, 5.42497807e-06, 4.45813929e-03]) + >>> plt.plot(x, iw) + >>> plt.show() + + The input quantiles can be any shape of array, as long as the last + axis labels the components. + + Alternatively, the object may be called (as a function) to fix the degrees + of freedom and scale parameters, returning a "frozen" inverse Wishart + random variable: + + >>> rv = invwishart(df=1, scale=1) + >>> # Frozen object with the same methods but holding the given + >>> # degrees of freedom and scale fixed. + """ + + def __init__(self, seed=None): + super().__init__(seed) + self.__doc__ = doccer.docformat(self.__doc__, wishart_docdict_params) + + def __call__(self, df=None, scale=None, seed=None): + """Create a frozen inverse Wishart distribution. + + See `invwishart_frozen` for more information. + + """ + return invwishart_frozen(df, scale, seed) + + def _logpdf(self, x, dim, df, log_det_scale, C): + """Log of the inverse Wishart probability density function. + + Parameters + ---------- + x : ndarray + Points at which to evaluate the log of the probability + density function. + dim : int + Dimension of the scale matrix + df : int + Degrees of freedom + log_det_scale : float + Logarithm of the determinant of the scale matrix + C : ndarray + Cholesky factorization of the scale matrix, lower triangular. + + Notes + ----- + As this function does no argument checking, it should not be + called directly; use 'logpdf' instead. + + """ + # Retrieve tr(scale x^{-1}) + log_det_x = np.empty(x.shape[-1]) + tr_scale_x_inv = np.empty(x.shape[-1]) + trsm = get_blas_funcs(('trsm'), (x,)) + if dim > 1: + for i in range(x.shape[-1]): + Cx, log_det_x[i] = self._cholesky_logdet(x[:, :, i]) + A = trsm(1., Cx, C, side=0, lower=True) + tr_scale_x_inv[i] = np.linalg.norm(A)**2 + else: + log_det_x[:] = np.log(x[0, 0]) + tr_scale_x_inv[:] = C[0, 0]**2 / x[0, 0] + + # Log PDF + out = ((0.5 * df * log_det_scale - 0.5 * tr_scale_x_inv) - + (0.5 * df * dim * _LOG_2 + 0.5 * (df + dim + 1) * log_det_x) - + multigammaln(0.5*df, dim)) + + return out + + def logpdf(self, x, df, scale): + """Log of the inverse Wishart probability density function. + + Parameters + ---------- + x : array_like + Quantiles, with the last axis of `x` denoting the components. + Each quantile must be a symmetric positive definite matrix. + %(_doc_default_callparams)s + + Returns + ------- + pdf : ndarray + Log of the probability density function evaluated at `x` + + Notes + ----- + %(_doc_callparams_note)s + + """ + dim, df, scale = self._process_parameters(df, scale) + x = self._process_quantiles(x, dim) + C, log_det_scale = self._cholesky_logdet(scale) + out = self._logpdf(x, dim, df, log_det_scale, C) + return _squeeze_output(out) + + def pdf(self, x, df, scale): + """Inverse Wishart probability density function. + + Parameters + ---------- + x : array_like + Quantiles, with the last axis of `x` denoting the components. + Each quantile must be a symmetric positive definite matrix. + %(_doc_default_callparams)s + + Returns + ------- + pdf : ndarray + Probability density function evaluated at `x` + + Notes + ----- + %(_doc_callparams_note)s + + """ + return np.exp(self.logpdf(x, df, scale)) + + def _mean(self, dim, df, scale): + """Mean of the inverse Wishart distribution. + + Parameters + ---------- + dim : int + Dimension of the scale matrix + %(_doc_default_callparams)s + + Notes + ----- + As this function does no argument checking, it should not be + called directly; use 'mean' instead. + + """ + if df > dim + 1: + out = scale / (df - dim - 1) + else: + out = None + return out + + def mean(self, df, scale): + """Mean of the inverse Wishart distribution. + + Only valid if the degrees of freedom are greater than the dimension of + the scale matrix plus one. + + Parameters + ---------- + %(_doc_default_callparams)s + + Returns + ------- + mean : float or None + The mean of the distribution + + """ + dim, df, scale = self._process_parameters(df, scale) + out = self._mean(dim, df, scale) + return _squeeze_output(out) if out is not None else out + + def _mode(self, dim, df, scale): + """Mode of the inverse Wishart distribution. + + Parameters + ---------- + dim : int + Dimension of the scale matrix + %(_doc_default_callparams)s + + Notes + ----- + As this function does no argument checking, it should not be + called directly; use 'mode' instead. + + """ + return scale / (df + dim + 1) + + def mode(self, df, scale): + """Mode of the inverse Wishart distribution. + + Parameters + ---------- + %(_doc_default_callparams)s + + Returns + ------- + mode : float + The Mode of the distribution + + """ + dim, df, scale = self._process_parameters(df, scale) + out = self._mode(dim, df, scale) + return _squeeze_output(out) + + def _var(self, dim, df, scale): + """Variance of the inverse Wishart distribution. + + Parameters + ---------- + dim : int + Dimension of the scale matrix + %(_doc_default_callparams)s + + Notes + ----- + As this function does no argument checking, it should not be + called directly; use 'var' instead. + + """ + if df > dim + 3: + var = (df - dim + 1) * scale**2 + diag = scale.diagonal() # 1 x dim array + var += (df - dim - 1) * np.outer(diag, diag) + var /= (df - dim) * (df - dim - 1)**2 * (df - dim - 3) + else: + var = None + return var + + def var(self, df, scale): + """Variance of the inverse Wishart distribution. + + Only valid if the degrees of freedom are greater than the dimension of + the scale matrix plus three. + + Parameters + ---------- + %(_doc_default_callparams)s + + Returns + ------- + var : float + The variance of the distribution + """ + dim, df, scale = self._process_parameters(df, scale) + out = self._var(dim, df, scale) + return _squeeze_output(out) if out is not None else out + + def _inv_standard_rvs(self, n, shape, dim, df, random_state): + """ + Parameters + ---------- + n : int + Number of variates to generate + shape : iterable + Shape of the variates to generate + dim : int + Dimension of the scale matrix + df : int + Degrees of freedom + random_state : {None, int, `numpy.random.Generator`, + `numpy.random.RandomState`}, optional + + If `seed` is None (or `np.random`), the `numpy.random.RandomState` + singleton is used. + If `seed` is an int, a new ``RandomState`` instance is used, + seeded with `seed`. + If `seed` is already a ``Generator`` or ``RandomState`` instance + then that instance is used. + + Returns + ------- + A : ndarray + Random variates of shape (`shape`) + (``dim``, ``dim``). + Each slice `A[..., :, :]` is lower-triangular, and its + inverse is the lower Cholesky factor of a draw from + `invwishart(df, np.eye(dim))`. + + Notes + ----- + As this function does no argument checking, it should not be + called directly; use 'rvs' instead. + + """ + A = np.zeros(shape + (dim, dim)) + + # Random normal variates for off-diagonal elements + tri_rows, tri_cols = np.tril_indices(dim, k=-1) + n_tril = dim * (dim-1) // 2 + A[..., tri_rows, tri_cols] = random_state.normal( + size=(*shape, n_tril), + ) + + # Random chi variates for diagonal elements + rows = np.arange(dim) + chi_dfs = (df - dim + 1) + rows + A[..., rows, rows] = random_state.chisquare( + df=chi_dfs, size=(*shape, dim), + )**0.5 + + return A + + def _rvs(self, n, shape, dim, df, C, random_state): + """Draw random samples from an inverse Wishart distribution. + + Parameters + ---------- + n : int + Number of variates to generate + shape : iterable + Shape of the variates to generate + dim : int + Dimension of the scale matrix + df : int + Degrees of freedom + C : ndarray + Cholesky factorization of the scale matrix, lower triangular. + %(_doc_random_state)s + + Notes + ----- + As this function does no argument checking, it should not be + called directly; use 'rvs' instead. + + """ + random_state = self._get_random_state(random_state) + # Get random draws A such that inv(A) ~ iW(df, I) + A = self._inv_standard_rvs(n, shape, dim, df, random_state) + + # Calculate SA = (CA)'^{-1} (CA)^{-1} ~ iW(df, scale) + trsm = get_blas_funcs(('trsm'), (A,)) + trmm = get_blas_funcs(('trmm'), (A,)) + + for index in np.ndindex(A.shape[:-2]): + if dim > 1: + # Calculate CA + # Get CA = C A^{-1} via triangular solver + CA = trsm(1., A[index], C, side=1, lower=True) + # get SA + A[index] = trmm(1., CA, CA, side=1, lower=True, trans_a=True) + else: + A[index][0, 0] = (C[0, 0] / A[index][0, 0])**2 + + return A + + def rvs(self, df, scale, size=1, random_state=None): + """Draw random samples from an inverse Wishart distribution. + + Parameters + ---------- + %(_doc_default_callparams)s + size : int or iterable of integers, optional + Number of samples to draw (default 1). + %(_doc_random_state)s + + Returns + ------- + rvs : ndarray + Random variates of shape (`size`) + (``dim``, ``dim``), where + ``dim`` is the dimension of the scale matrix. + + Notes + ----- + %(_doc_callparams_note)s + + """ + n, shape = self._process_size(size) + dim, df, scale = self._process_parameters(df, scale) + + # Cholesky decomposition of scale + C = scipy.linalg.cholesky(scale, lower=True) + + out = self._rvs(n, shape, dim, df, C, random_state) + + return _squeeze_output(out) + + def _entropy(self, dim, df, log_det_scale): + # reference: eq. (17) from ref. 3 + psi_eval_points = [0.5 * (df - dim + i) for i in range(1, dim + 1)] + psi_eval_points = np.asarray(psi_eval_points) + return multigammaln(0.5 * df, dim) + 0.5 * dim * df + \ + 0.5 * (dim + 1) * (log_det_scale - _LOG_2) - \ + 0.5 * (df + dim + 1) * \ + psi(psi_eval_points, out=psi_eval_points).sum() + + def entropy(self, df, scale): + dim, df, scale = self._process_parameters(df, scale) + _, log_det_scale = self._cholesky_logdet(scale) + return self._entropy(dim, df, log_det_scale) + + +invwishart = invwishart_gen() + + +class invwishart_frozen(multi_rv_frozen): + __class_getitem__ = None # pyrefly:ignore[bad-assignment] + + def __init__(self, df, scale, seed=None): + """Create a frozen inverse Wishart distribution. + + Parameters + ---------- + df : array_like + Degrees of freedom of the distribution + scale : array_like + Scale matrix of the distribution + seed : {None, int, `numpy.random.Generator`}, optional + If `seed` is None the `numpy.random.Generator` singleton is used. + If `seed` is an int, a new ``Generator`` instance is used, + seeded with `seed`. + If `seed` is already a ``Generator`` instance then that instance is + used. + + """ + self._dist = invwishart_gen(seed) + self.dim, self.df, self.scale = self._dist._process_parameters( + df, scale + ) + + # Get the determinant via Cholesky factorization + self.C = scipy.linalg.cholesky(self.scale, lower=True) + self.log_det_scale = 2 * np.sum(np.log(self.C.diagonal())) + + def logpdf(self, x): + x = self._dist._process_quantiles(x, self.dim) + out = self._dist._logpdf(x, self.dim, self.df, + self.log_det_scale, self.C) + return _squeeze_output(out) + + def pdf(self, x): + return np.exp(self.logpdf(x)) + + def mean(self): + out = self._dist._mean(self.dim, self.df, self.scale) + return _squeeze_output(out) if out is not None else out + + def mode(self): + out = self._dist._mode(self.dim, self.df, self.scale) + return _squeeze_output(out) + + def var(self): + out = self._dist._var(self.dim, self.df, self.scale) + return _squeeze_output(out) if out is not None else out + + def rvs(self, size=1, random_state=None): + n, shape = self._dist._process_size(size) + + out = self._dist._rvs(n, shape, self.dim, self.df, + self.C, random_state) + + return _squeeze_output(out) + + def entropy(self): + return self._dist._entropy(self.dim, self.df, self.log_det_scale) + + +# Set frozen generator docstrings from corresponding docstrings in +# inverse Wishart and fill in default strings in class docstrings +for name in ['logpdf', 'pdf', 'mean', 'mode', 'var', 'rvs']: + method = invwishart_gen.__dict__[name] + method_frozen = wishart_frozen.__dict__[name] + method_frozen.__doc__ = doccer.docformat( + method.__doc__, wishart_docdict_noparams) + method.__doc__ = doccer.docformat(method.__doc__, wishart_docdict_params) + +_multinomial_doc_default_callparams = """\ +n : int + Number of trials +p : array_like + Probability of a trial falling into each category; should sum to 1 +""" + +_multinomial_doc_callparams_note = """\ +`n` should be a nonnegative integer. Each element of `p` should be in the +interval :math:`[0,1]` and the elements should sum to 1. If they do not sum to +1, the last element of the `p` array is not used and is replaced with the +remaining probability left over from the earlier elements. +""" + +_multinomial_doc_frozen_callparams = "" + +_multinomial_doc_frozen_callparams_note = """\ +See class definition for a detailed description of parameters.""" + +multinomial_docdict_params = { + '_doc_default_callparams': _multinomial_doc_default_callparams, + '_doc_callparams_note': _multinomial_doc_callparams_note, + '_doc_random_state': _doc_random_state +} + +multinomial_docdict_noparams = { + '_doc_default_callparams': _multinomial_doc_frozen_callparams, + '_doc_callparams_note': _multinomial_doc_frozen_callparams_note, + '_doc_random_state': _doc_random_state +} + + +class multinomial_gen(multi_rv_generic): + r""" + A multinomial random variable. + + Parameters + ---------- + %(_doc_default_callparams)s + %(_doc_random_state)s + + Methods + ------- + pmf(x, n, p) + Probability mass function. + logpmf(x, n, p) + Log of the probability mass function. + rvs(n, p, size=1, random_state=None) + Draw random samples from a multinomial distribution. + entropy(n, p) + Compute the entropy of the multinomial distribution. + cov(n, p) + Compute the covariance matrix of the multinomial distribution. + + See Also + -------- + scipy.stats.binom : The binomial distribution. + numpy.random.Generator.multinomial : Sampling from the multinomial distribution. + scipy.stats.multivariate_hypergeom : + The multivariate hypergeometric distribution. + + Notes + ----- + %(_doc_callparams_note)s + + The probability mass function for `multinomial` is + + .. math:: + + f(x) = \frac{n!}{x_1! \cdots x_k!} p_1^{x_1} \cdots p_k^{x_k}, + + supported on :math:`x=(x_1, \ldots, x_k)` where each :math:`x_i` is a + nonnegative integer and their sum is :math:`n`. + + .. versionadded:: 0.19.0 + + Examples + -------- + + >>> from scipy.stats import multinomial + >>> rv = multinomial(8, [0.3, 0.2, 0.5]) + >>> rv.pmf([1, 3, 4]) + 0.042000000000000072 + + The multinomial distribution for :math:`k=2` is identical to the + corresponding binomial distribution (tiny numerical differences + notwithstanding): + + >>> from scipy.stats import binom + >>> multinomial.pmf([3, 4], n=7, p=[0.4, 0.6]) + 0.29030399999999973 + >>> binom.pmf(3, 7, 0.4) + 0.29030400000000012 + + The functions ``pmf``, ``logpmf``, ``entropy``, and ``cov`` support + broadcasting, under the convention that the vector parameters (``x`` and + ``p``) are interpreted as if each row along the last axis is a single + object. For instance: + + >>> multinomial.pmf([[3, 4], [3, 5]], n=[7, 8], p=[.3, .7]) + array([0.2268945, 0.25412184]) + + Here, ``x.shape == (2, 2)``, ``n.shape == (2,)``, and ``p.shape == (2,)``, + but following the rules mentioned above they behave as if the rows + ``[3, 4]`` and ``[3, 5]`` in ``x`` and ``[.3, .7]`` in ``p`` were a single + object, and as if we had ``x.shape = (2,)``, ``n.shape = (2,)``, and + ``p.shape = ()``. To obtain the individual elements without broadcasting, + we would do this: + + >>> multinomial.pmf([3, 4], n=7, p=[.3, .7]) + 0.2268945 + >>> multinomial.pmf([3, 5], 8, p=[.3, .7]) + 0.25412184 + + This broadcasting also works for ``cov``, where the output objects are + square matrices of size ``p.shape[-1]``. For example: + + >>> multinomial.cov([4, 5], [[.3, .7], [.4, .6]]) + array([[[ 0.84, -0.84], + [-0.84, 0.84]], + [[ 1.2 , -1.2 ], + [-1.2 , 1.2 ]]]) + + In this example, ``n.shape == (2,)`` and ``p.shape == (2, 2)``, and + following the rules above, these broadcast as if ``p.shape == (2,)``. + Thus the result should also be of shape ``(2,)``, but since each output is + a :math:`2 \times 2` matrix, the result in fact has shape ``(2, 2, 2)``, + where ``result[0]`` is equal to ``multinomial.cov(n=4, p=[.3, .7])`` and + ``result[1]`` is equal to ``multinomial.cov(n=5, p=[.4, .6])``. + + Alternatively, the object may be called (as a function) to fix the `n` and + `p` parameters, returning a "frozen" multinomial random variable: + + >>> rv = multinomial(n=7, p=[.3, .7]) + >>> # Frozen object with the same methods but holding the given + >>> # degrees of freedom and scale fixed. + """ + + def __init__(self, seed=None): + super().__init__(seed) + self.__doc__ = \ + doccer.docformat(self.__doc__, multinomial_docdict_params) + + def __call__(self, n, p, seed=None): + """Create a frozen multinomial distribution. + + See `multinomial_frozen` for more information. + """ + return multinomial_frozen(n, p, seed) + + def _process_parameters(self, n, p): + """Returns: n_, p_, npcond. + + n_ and p_ are arrays of the correct shape; npcond is a boolean array + flagging values out of the domain. + """ + eps = np.finfo(np.result_type(np.asarray(p), np.float32)).eps * 10 + p = np.array(p, dtype=np.float64, copy=True) + + # true for bad p + pcond = np.abs(1 - p.sum(axis=-1)) > eps + pcond |= np.any(p < 0, axis=-1) + pcond |= np.any(p > 1, axis=-1) + + n = np.array(n, dtype=int, copy=True) + + # true for bad n + ncond = n < 0 + + return n, p, ncond | pcond + + def _process_quantiles(self, x, n, p): + """Returns: x_, xcond. + + x_ is an int array; xcond is a boolean array flagging values out of the + domain. + """ + xx = np.asarray(x, dtype=int) + + if xx.ndim == 0: + raise ValueError("x must be an array.") + + if xx.size != 0 and not xx.shape[-1] == p.shape[-1]: + raise ValueError(f"Size of each quantile should be size of p: " + f"received {xx.shape[-1]}, but expected " + f"{p.shape[-1]}.") + + # true for x out of the domain + cond = np.any(xx != x, axis=-1) + cond |= np.any(xx < 0, axis=-1) + cond = cond | (np.sum(xx, axis=-1) != n) + + return xx, cond + + def _checkresult(self, result, cond, bad_value): + result = np.asarray(result) + + if cond.ndim != 0: + result[cond] = bad_value + elif cond: + if result.ndim == 0: + return bad_value + result[...] = bad_value + return result + + def _logpmf(self, x, n, p): + return gammaln(n+1) + np.sum(xlogy(x, p) - gammaln(x+1), axis=-1) + + def logpmf(self, x, n, p): + """Log of the Multinomial probability mass function. + + Parameters + ---------- + x : array_like + Quantiles, with the last axis of `x` denoting the components. + %(_doc_default_callparams)s + + Returns + ------- + logpmf : ndarray or scalar + Log of the probability mass function evaluated at `x` + + Notes + ----- + %(_doc_callparams_note)s + """ + n, p, npcond = self._process_parameters(n, p) + x, xcond = self._process_quantiles(x, n, p) + + result = self._logpmf(x, n, p) + + # replace values for which x was out of the domain; broadcast + # xcond to the right shape + xcond_ = xcond | np.zeros(npcond.shape, dtype=np.bool_) + result = self._checkresult(result, xcond_, -np.inf) + + # replace values bad for n or p; broadcast npcond to the right shape + npcond_ = npcond | np.zeros(xcond.shape, dtype=np.bool_) + return self._checkresult(result, npcond_, np.nan) + + def pmf(self, x, n, p): + """Multinomial probability mass function. + + Parameters + ---------- + x : array_like + Quantiles, with the last axis of `x` denoting the components. + %(_doc_default_callparams)s + + Returns + ------- + pmf : ndarray or scalar + Probability density function evaluated at `x` + + Notes + ----- + %(_doc_callparams_note)s + """ + return np.exp(self.logpmf(x, n, p)) + + def mean(self, n, p): + """Mean of the Multinomial distribution. + + Parameters + ---------- + %(_doc_default_callparams)s + + Returns + ------- + mean : float + The mean of the distribution + """ + n, p, npcond = self._process_parameters(n, p) + result = n[..., np.newaxis]*p + return self._checkresult(result, npcond, np.nan) + + def cov(self, n, p): + """Covariance matrix of the multinomial distribution. + + Parameters + ---------- + %(_doc_default_callparams)s + + Returns + ------- + cov : ndarray + The covariance matrix of the distribution + """ + n, p, npcond = self._process_parameters(n, p) + + nn = n[..., np.newaxis, np.newaxis] + result = nn * np.einsum('...j,...k->...jk', -p, p) + + # change the diagonal + for i in range(p.shape[-1]): + result[..., i, i] += n*p[..., i] + + return self._checkresult(result, npcond, np.nan) + + def entropy(self, n, p): + r"""Compute the entropy of the multinomial distribution. + + The entropy is computed using this expression: + + .. math:: + + f(x) = - \log n! - n\sum_{i=1}^k p_i \log p_i + + \sum_{i=1}^k \sum_{x=0}^n \binom n x p_i^x(1-p_i)^{n-x} \log x! + + Parameters + ---------- + %(_doc_default_callparams)s + + Returns + ------- + h : scalar + Entropy of the Multinomial distribution + + Notes + ----- + %(_doc_callparams_note)s + """ + n, p, npcond = self._process_parameters(n, p) + + x = np.r_[1:np.max(n)+1] + + term1 = n*np.sum(entr(p), axis=-1) + term1 -= gammaln(n+1) + + n = n[..., np.newaxis] + new_axes_needed = max(p.ndim, n.ndim) - x.ndim + 1 + new_shape = x.shape + (1,)*new_axes_needed + x = x.reshape(new_shape) + + term2 = np.sum(binom.pmf(x, n, p)*gammaln(x+1), + axis=(-1, -1-new_axes_needed)) + + return self._checkresult(term1 + term2, npcond, np.nan) + + def rvs(self, n, p, size=None, random_state=None): + """Draw random samples from a Multinomial distribution. + + Parameters + ---------- + %(_doc_default_callparams)s + size : int or iterable of integers, optional + Number of samples to draw (default 1). + %(_doc_random_state)s + + Returns + ------- + rvs : ndarray or scalar + Random variates of shape (`size`, `len(p)`) + + Notes + ----- + %(_doc_callparams_note)s + """ + n, p, npcond = self._process_parameters(n, p) + if np.any(npcond): + message = ("`multinomial.rvs` requires `n > 0`, `(p > 0).all()`, and" + "`p.sum() == 1`.") + raise ValueError(message) + random_state = self._get_random_state(random_state) + return random_state.multinomial(n, p, size) + + +multinomial = multinomial_gen() + + +class multinomial_frozen(multi_rv_frozen): + r"""Create a frozen Multinomial distribution. + + Parameters + ---------- + n : int + number of trials + p: array_like + probability of a trial falling into each category; should sum to 1 + seed : {None, int, `numpy.random.Generator`, `numpy.random.RandomState`}, optional + If `seed` is None (or `np.random`), the `numpy.random.RandomState` + singleton is used. + If `seed` is an int, a new ``RandomState`` instance is used, + seeded with `seed`. + If `seed` is already a ``Generator`` or ``RandomState`` instance then + that instance is used. + """ + def __init__(self, n, p, seed=None): + self._dist = multinomial_gen(seed) + self.n, self.p, self.npcond = self._dist._process_parameters(n, p) + + # monkey patch self._dist + def _process_parameters(n, p): + return self.n, self.p, self.npcond + + self._dist._process_parameters = _process_parameters + + def logpmf(self, x): + return self._dist.logpmf(x, self.n, self.p) + + def pmf(self, x): + return self._dist.pmf(x, self.n, self.p) + + def mean(self): + return self._dist.mean(self.n, self.p) + + def cov(self): + return self._dist.cov(self.n, self.p) + + def entropy(self): + return self._dist.entropy(self.n, self.p) + + def rvs(self, size=1, random_state=None): + return self._dist.rvs(self.n, self.p, size, random_state) + + +# Set frozen generator docstrings from corresponding docstrings in +# multinomial and fill in default strings in class docstrings +for name in ['logpmf', 'pmf', 'mean', 'cov', 'rvs']: + method = multinomial_gen.__dict__[name] + method_frozen = multinomial_frozen.__dict__[name] + method_frozen.__doc__ = doccer.docformat( + method.__doc__, multinomial_docdict_noparams) + method.__doc__ = doccer.docformat(method.__doc__, + multinomial_docdict_params) + + +class special_ortho_group_gen(multi_rv_generic): + r""" + A Special Orthogonal matrix (SO(N)) random variable. + + Return a random rotation matrix, drawn from the Haar distribution + (the only uniform distribution on SO(N)) with a determinant of +1. + + The `dim` keyword specifies the dimension N. + + Parameters + ---------- + dim : scalar + Dimension of matrices + seed : {None, int, np.random.RandomState, np.random.Generator}, optional + Used for drawing random variates. + If `seed` is `None`, the `~np.random.RandomState` singleton is used. + If `seed` is an int, a new ``RandomState`` instance is used, seeded + with seed. + If `seed` is already a ``RandomState`` or ``Generator`` instance, + then that object is used. + Default is `None`. + + Methods + ------- + rvs(dim=None, size=1, random_state=None) + Draw random samples from SO(N). + + See Also + -------- + ortho_group, scipy.spatial.transform.Rotation.random + + Notes + ----- + The ``rvs`` method returns a random rotation matrix drawn from the Haar + distribution, the only uniform distribution on SO(N). The algorithm generates + a Haar-distributed orthogonal matrix in O(N) using the ``rvs`` method of + `ortho_group`, then adjusts the matrix to ensure that the determinant is +1. + + For a random rotation in three dimensions, see + `scipy.spatial.transform.Rotation.random`. + + Examples + -------- + >>> import numpy as np + >>> from scipy.stats import special_ortho_group + >>> x = special_ortho_group.rvs(3) + + >>> np.dot(x, x.T) + array([[ 1.00000000e+00, 1.13231364e-17, -2.86852790e-16], + [ 1.13231364e-17, 1.00000000e+00, -1.46845020e-16], + [ -2.86852790e-16, -1.46845020e-16, 1.00000000e+00]]) + + >>> import scipy.linalg + >>> scipy.linalg.det(x) + 1.0 + + This generates one random matrix from SO(3). It is orthogonal and + has a determinant of 1. + + Alternatively, the object may be called (as a function) to fix the `dim` + parameter, returning a "frozen" special_ortho_group random variable: + + >>> rv = special_ortho_group(5) + >>> # Frozen object with the same methods but holding the + >>> # dimension parameter fixed. + """ + + def __init__(self, seed=None): + super().__init__(seed) + self.__doc__ = doccer.docformat(self.__doc__) + + def __call__(self, dim=None, seed=None): + """Create a frozen SO(N) distribution. + + See `special_ortho_group_frozen` for more information. + """ + return special_ortho_group_frozen(dim, seed=seed) + + def _process_parameters(self, dim): + """Dimension N must be specified; it cannot be inferred.""" + if dim is None or not np.isscalar(dim) or dim < 0 or dim != int(dim): + raise ValueError("""Dimension of rotation must be specified, + and must be a scalar nonnegative integer.""") + + return dim + + def rvs(self, dim, size=1, random_state=None): + """Draw random samples from SO(N). + + Parameters + ---------- + dim : int + Dimension of rotation space (N). + size : int, optional + Number of samples to draw (default 1). + + Returns + ------- + rvs : ndarray or scalar + Random size N-dimensional matrices, dimension (size, dim, dim) + + """ + random_state = self._get_random_state(random_state) + + q = ortho_group.rvs(dim, size, random_state) + dets = np.linalg.det(q) + if dim: + q[..., 0, :] /= dets[..., np.newaxis] + return q + +special_ortho_group = special_ortho_group_gen() + + +class special_ortho_group_frozen(multi_rv_frozen): + __class_getitem__ = None # pyrefly:ignore[bad-assignment] + + def __init__(self, dim=None, seed=None): + """Create a frozen SO(N) distribution. + + Parameters + ---------- + dim : scalar + Dimension of matrices + seed : {None, int, `numpy.random.Generator`, `numpy.random.RandomState`}, optional + If `seed` is None (or `np.random`), the `numpy.random.RandomState` + singleton is used. + If `seed` is an int, a new ``RandomState`` instance is used, + seeded with `seed`. + If `seed` is already a ``Generator`` or ``RandomState`` instance + then that instance is used. + + Examples + -------- + >>> from scipy.stats import special_ortho_group + >>> g = special_ortho_group(5) + >>> x = g.rvs() + + """ # numpy/numpydoc#87 # noqa: E501 + self._dist = special_ortho_group_gen(seed) + self.dim = self._dist._process_parameters(dim) + + def rvs(self, size=1, random_state=None): + return self._dist.rvs(self.dim, size, random_state) + + +class ortho_group_gen(multi_rv_generic): + r""" + An Orthogonal matrix (O(N)) random variable. + + Return a random orthogonal matrix, drawn from the O(N) Haar + distribution (the only uniform distribution on O(N)). + + The `dim` keyword specifies the dimension N. + + Parameters + ---------- + dim : scalar + Dimension of matrices + seed : {None, int, np.random.RandomState, np.random.Generator}, optional + Used for drawing random variates. + If `seed` is `None`, the `~np.random.RandomState` singleton is used. + If `seed` is an int, a new ``RandomState`` instance is used, seeded + with seed. + If `seed` is already a ``RandomState`` or ``Generator`` instance, + then that object is used. + Default is `None`. + + Methods + ------- + rvs(dim=None, size=1, random_state=None) + Draw random samples from O(N). + + See Also + -------- + special_ortho_group + + Notes + ----- + This class is closely related to `special_ortho_group`. + + Some care is taken to avoid numerical error, as per the paper by Mezzadri. + + References + ---------- + .. [1] F. Mezzadri, "How to generate random matrices from the classical + compact groups", :arXiv:`math-ph/0609050v2`. + + Examples + -------- + >>> import numpy as np + >>> from scipy.stats import ortho_group + >>> x = ortho_group.rvs(3) + + >>> np.dot(x, x.T) + array([[ 1.00000000e+00, 1.13231364e-17, -2.86852790e-16], + [ 1.13231364e-17, 1.00000000e+00, -1.46845020e-16], + [ -2.86852790e-16, -1.46845020e-16, 1.00000000e+00]]) + + >>> import scipy.linalg + >>> np.fabs(scipy.linalg.det(x)) + 1.0 + + This generates one random matrix from O(3). It is orthogonal and + has a determinant of +1 or -1. + + Alternatively, the object may be called (as a function) to fix the `dim` + parameter, returning a "frozen" ortho_group random variable: + + >>> rv = ortho_group(5) + >>> # Frozen object with the same methods but holding the + >>> # dimension parameter fixed. + """ + + def __init__(self, seed=None): + super().__init__(seed) + self.__doc__ = doccer.docformat(self.__doc__) + + def __call__(self, dim=None, seed=None): + """Create a frozen O(N) distribution. + + See `ortho_group_frozen` for more information. + """ + return ortho_group_frozen(dim, seed=seed) + + def _process_parameters(self, dim): + """Dimension N must be specified; it cannot be inferred.""" + if dim is None or not np.isscalar(dim) or dim < 0 or dim != int(dim): + raise ValueError("Dimension of rotation must be specified," + "and must be a scalar nonnegative integer.") + + return dim + + def rvs(self, dim, size=1, random_state=None): + """Draw random samples from O(N). + + Parameters + ---------- + dim : int + Dimension of rotation space (N). + size : int, optional + Number of samples to draw (default 1). + + Returns + ------- + rvs : ndarray or scalar + Random size N-dimensional matrices, dimension (size, dim, dim) + + """ + random_state = self._get_random_state(random_state) + + size = int(size) + + dim = self._process_parameters(dim) + + size = (size,) if size > 1 else () + z = random_state.normal(size=size + (dim, dim)) + q, r = np.linalg.qr(z) + # The last two dimensions are the rows and columns of R matrices. + # Extract the diagonals. Note that this eliminates a dimension. + d = r.diagonal(offset=0, axis1=-2, axis2=-1) + # Add back a dimension for proper broadcasting: we're dividing + # each row of each R matrix by the diagonal of the R matrix. + q *= (d/abs(d))[..., np.newaxis, :] # to broadcast properly + return q + + +ortho_group = ortho_group_gen() + + +class ortho_group_frozen(multi_rv_frozen): + __class_getitem__ = None # pyrefly:ignore[bad-assignment] + + def __init__(self, dim=None, seed=None): + """Create a frozen O(N) distribution. + + Parameters + ---------- + dim : scalar + Dimension of matrices + seed : {None, int, `numpy.random.Generator`, `numpy.random.RandomState`}, optional + If `seed` is None (or `np.random`), the `numpy.random.RandomState` + singleton is used. + If `seed` is an int, a new ``RandomState`` instance is used, + seeded with `seed`. + If `seed` is already a ``Generator`` or ``RandomState`` instance + then that instance is used. + + Examples + -------- + >>> from scipy.stats import ortho_group + >>> g = ortho_group(5) + >>> x = g.rvs() + + """ # numpy/numpydoc#87 # noqa: E501 + self._dist = ortho_group_gen(seed) + self.dim = self._dist._process_parameters(dim) + + def rvs(self, size=1, random_state=None): + return self._dist.rvs(self.dim, size, random_state) + + +class random_correlation_gen(multi_rv_generic): + r""" + A random correlation matrix. + + Return a random correlation matrix, given a vector of eigenvalues. + The returned matrix is symmetric positive semidefinite with unit diagonal. + + The `eigs` keyword specifies the eigenvalues of the correlation matrix, + and implies the dimension. + + Parameters + ---------- + eigs : 1d ndarray + Eigenvalues of correlation matrix. All eigenvalues need to be non-negative and + need to sum to the number of eigenvalues. + seed : {None, int, `numpy.random.Generator`, `numpy.random.RandomState`}, optional + If `seed` is None (or `np.random`), the `numpy.random.RandomState` + singleton is used. + If `seed` is an int, a new ``RandomState`` instance is used, + seeded with `seed`. + If `seed` is already a ``Generator`` or ``RandomState`` instance + then that instance is used. + tol : float, optional + Tolerance for input parameter checks + diag_tol : float, optional + Tolerance for deviation of the diagonal of the resulting + matrix. Default: 1e-7 + + Methods + ------- + rvs(eigs=None, random_state=None) + Draw random correlation matrices, all with eigenvalues eigs. + + Returns + ------- + rvs : ndarray or scalar + Random size N-dimensional matrices, dimension (size, dim, dim), + each having eigenvalues eigs. + + Raises + ------ + RuntimeError + Floating point error prevented generating a valid correlation + matrix. + + Notes + ----- + + Generates a random correlation matrix following a numerically stable + algorithm spelled out by Davies & Higham. This algorithm uses a single O(N) + similarity transformation to construct a symmetric positive semi-definite + matrix, and applies a series of Givens rotations to scale it to have ones + on the diagonal. + + References + ---------- + + .. [1] Davies, Philip I; Higham, Nicholas J; "Numerically stable generation + of correlation matrices and their factors", BIT 2000, Vol. 40, + No. 4, pp. 640 651 + + Examples + -------- + >>> import numpy as np + >>> from scipy.stats import random_correlation + >>> rng = np.random.default_rng() + >>> x = random_correlation.rvs((.5, .8, 1.2, 1.5), random_state=rng) + >>> x + array([[ 1. , -0.02423399, 0.03130519, 0.4946965 ], + [-0.02423399, 1. , 0.20334736, 0.04039817], + [ 0.03130519, 0.20334736, 1. , 0.02694275], + [ 0.4946965 , 0.04039817, 0.02694275, 1. ]]) + >>> import scipy.linalg + >>> e, v = scipy.linalg.eigh(x) + >>> e + array([ 0.5, 0.8, 1.2, 1.5]) + """ + + def __init__(self, seed=None): + super().__init__(seed) + self.__doc__ = doccer.docformat(self.__doc__) + + def __call__(self, eigs, seed=None, tol=1e-13, diag_tol=1e-7): + """Create a frozen random correlation matrix. + + See `random_correlation_frozen` for more information. + """ + return random_correlation_frozen(eigs, seed=seed, tol=tol, + diag_tol=diag_tol) + + def _process_parameters(self, eigs, tol): + eigs = np.asarray(eigs, dtype=float) + dim = eigs.size + + if eigs.ndim != 1 or eigs.shape[0] != dim or dim <= 1: + raise ValueError("Array 'eigs' must be a vector of length " + "greater than 1.") + + if np.fabs(np.sum(eigs) - dim) > tol: + raise ValueError("Sum of eigenvalues must equal dimensionality.") + + for x in eigs: + if x < -tol: + raise ValueError("All eigenvalues must be non-negative.") + + return dim, eigs + + def _givens_to_1(self, aii, ajj, aij): + """Computes a 2x2 Givens matrix to put 1's on the diagonal. + + The input matrix is a 2x2 symmetric matrix M = [ aii aij ; aij ajj ]. + + The output matrix g is a 2x2 anti-symmetric matrix of the form + [ c s ; -s c ]; the elements c and s are returned. + + Applying the output matrix to the input matrix (as b=g.T M g) + results in a matrix with bii=1, provided tr(M) - det(M) >= 1 + and floating point issues do not occur. Otherwise, some other + valid rotation is returned. When tr(M)==2, also bjj=1. + + """ + aiid = aii - 1. + ajjd = ajj - 1. + + if ajjd == 0: + # ajj==1, so swap aii and ajj to avoid division by zero + return 0., 1. + + dd = math.sqrt(max(aij**2 - aiid*ajjd, 0)) + + # The choice of t should be chosen to avoid cancellation [1] + t = (aij + math.copysign(dd, aij)) / ajjd + c = 1. / math.sqrt(1. + t*t) + if c == 0: + # Underflow + s = 1.0 + else: + s = c*t + return c, s + + def _to_corr(self, m): + """ + Given a psd matrix m, rotate to put one's on the diagonal, turning it + into a correlation matrix. This also requires the trace equal the + dimensionality. Note: modifies input matrix + """ + # Check requirements for in-place Givens + if not (m.flags.c_contiguous and m.dtype == np.float64 and + m.shape[0] == m.shape[1]): + raise ValueError() + + drot = get_blas_funcs('rot', dtype=np.float64, ilp64='preferred') + + d = m.shape[0] + for i in range(d-1): + if m[i, i] == 1: + continue + elif m[i, i] > 1: + for j in range(i+1, d): + if m[j, j] < 1: + break + else: + for j in range(i+1, d): + if m[j, j] > 1: + break + + c, s = self._givens_to_1(m[i, i], m[j, j], m[i, j]) + + # Use BLAS to apply Givens rotations in-place. Equivalent to: + # g = np.eye(d) + # g[i, i] = g[j,j] = c + # g[j, i] = -s; g[i, j] = s + # m = np.dot(g.T, np.dot(m, g)) + mv = m.ravel() + drot(mv, mv, c, -s, n=d, + offx=i*d, incx=1, offy=j*d, incy=1, + overwrite_x=True, overwrite_y=True) + drot(mv, mv, c, -s, n=d, + offx=i, incx=d, offy=j, incy=d, + overwrite_x=True, overwrite_y=True) + + return m + + def rvs(self, eigs, random_state=None, tol=1e-13, diag_tol=1e-7): + """Draw random correlation matrices. + + Parameters + ---------- + eigs : 1d ndarray + Eigenvalues of correlation matrix + tol : float, optional + Tolerance for input parameter checks + diag_tol : float, optional + Tolerance for deviation of the diagonal of the resulting + matrix. Default: 1e-7 + + Raises + ------ + RuntimeError + Floating point error prevented generating a valid correlation + matrix. + + Returns + ------- + rvs : ndarray or scalar + Random size N-dimensional matrices, dimension (size, dim, dim), + each having eigenvalues eigs. + + """ + dim, eigs = self._process_parameters(eigs, tol=tol) + + random_state = self._get_random_state(random_state) + + m = ortho_group.rvs(dim, random_state=random_state) + m = np.dot(np.dot(m, np.diag(eigs)), m.T) # Set the trace of m + m = self._to_corr(m) # Carefully rotate to unit diagonal + + # Check diagonal + if abs(m.diagonal() - 1).max() > diag_tol: + raise RuntimeError("Failed to generate a valid correlation matrix") + + return m + + +random_correlation = random_correlation_gen() + + +class random_correlation_frozen(multi_rv_frozen): + __class_getitem__ = None # pyrefly:ignore[bad-assignment] + + def __init__(self, eigs, seed=None, tol=1e-13, diag_tol=1e-7): + """Create a frozen random correlation matrix distribution. + + Parameters + ---------- + eigs : 1d ndarray + Eigenvalues of correlation matrix + seed : {None, int, `numpy.random.Generator`, `numpy.random.RandomState`}, optional + If `seed` is None (or `np.random`), the `numpy.random.RandomState` + singleton is used. + If `seed` is an int, a new ``RandomState`` instance is used, + seeded with `seed`. + If `seed` is already a ``Generator`` or ``RandomState`` instance + then that instance is used. + tol : float, optional + Tolerance for input parameter checks + diag_tol : float, optional + Tolerance for deviation of the diagonal of the resulting + matrix. Default: 1e-7 + + Raises + ------ + RuntimeError + Floating point error prevented generating a valid correlation + matrix. + + Returns + ------- + rvs : ndarray or scalar + Random size N-dimensional matrices, dimension (size, dim, dim), + each having eigenvalues eigs. + """ # numpy/numpydoc#87 # noqa: E501 + + self._dist = random_correlation_gen(seed) + self.tol = tol + self.diag_tol = diag_tol + _, self.eigs = self._dist._process_parameters(eigs, tol=self.tol) + + def rvs(self, random_state=None): + return self._dist.rvs(self.eigs, random_state=random_state, + tol=self.tol, diag_tol=self.diag_tol) + + +class unitary_group_gen(multi_rv_generic): + r""" + A matrix-valued U(N) random variable. + + Return a random unitary matrix. + + The `dim` keyword specifies the dimension N. + + Parameters + ---------- + dim : scalar + Dimension of matrices. + seed : {None, int, np.random.RandomState, np.random.Generator}, optional + Used for drawing random variates. + If `seed` is `None`, the `~np.random.RandomState` singleton is used. + If `seed` is an int, a new ``RandomState`` instance is used, seeded + with seed. + If `seed` is already a ``RandomState`` or ``Generator`` instance, + then that object is used. + Default is `None`. + + Methods + ------- + rvs(dim=None, size=1, random_state=None) + Draw random samples from U(N). + + See Also + -------- + ortho_group + + Notes + ----- + This class is similar to `ortho_group`. + + References + ---------- + .. [1] F. Mezzadri, "How to generate random matrices from the classical + compact groups", :arXiv:`math-ph/0609050v2`. + + Examples + -------- + >>> import numpy as np + >>> from scipy.stats import unitary_group + >>> x = unitary_group.rvs(3) + + >>> np.dot(x, x.conj().T) + array([[ 1.00000000e+00, 1.13231364e-17, -2.86852790e-16], + [ 1.13231364e-17, 1.00000000e+00, -1.46845020e-16], + [ -2.86852790e-16, -1.46845020e-16, 1.00000000e+00]]) # may vary + + This generates one random matrix from U(3). The dot product confirms that + it is unitary up to machine precision. + + Alternatively, the object may be called (as a function) to fix the `dim` + parameter, return a "frozen" unitary_group random variable: + + >>> rv = unitary_group(5) + """ + + def __init__(self, seed=None): + super().__init__(seed) + self.__doc__ = doccer.docformat(self.__doc__) + + def __call__(self, dim=None, seed=None): + """Create a frozen (U(N)) n-dimensional unitary matrix distribution. + + See `unitary_group_frozen` for more information. + """ + return unitary_group_frozen(dim, seed=seed) + + def _process_parameters(self, dim): + """Dimension N must be specified; it cannot be inferred.""" + if dim is None or not np.isscalar(dim) or dim < 0 or dim != int(dim): + raise ValueError("Dimension of rotation must be specified," + "and must be a scalar nonnegative integer.") + + return dim + + def rvs(self, dim, size=1, random_state=None): + """Draw random samples from U(N). + + Parameters + ---------- + dim : int + Dimension of space (N). + size : int, optional + Number of samples to draw (default 1). + + Returns + ------- + rvs : ndarray or scalar + Random size N-dimensional matrices, dimension (size, dim, dim) + + """ + random_state = self._get_random_state(random_state) + + size = int(size) + + dim = self._process_parameters(dim) + + size = (size,) if size > 1 else () + z = 1/math.sqrt(2)*(random_state.normal(size=size + (dim, dim)) + + 1j*random_state.normal(size=size + (dim, dim))) + q, r = np.linalg.qr(z) + # The last two dimensions are the rows and columns of R matrices. + # Extract the diagonals. Note that this eliminates a dimension. + d = r.diagonal(offset=0, axis1=-2, axis2=-1) + # Add back a dimension for proper broadcasting: we're dividing + # each row of each R matrix by the diagonal of the R matrix. + q *= (d/abs(d))[..., np.newaxis, :] # to broadcast properly + return q + + +unitary_group = unitary_group_gen() + + +class unitary_group_frozen(multi_rv_frozen): + __class_getitem__ = None # pyrefly:ignore[bad-assignment] + + def __init__(self, dim=None, seed=None): + """Create a frozen (U(N)) n-dimensional unitary matrix distribution. + + Parameters + ---------- + dim : scalar + Dimension of matrices + seed : {None, int, `numpy.random.Generator`, `numpy.random.RandomState`}, optional + If `seed` is None (or `np.random`), the `numpy.random.RandomState` + singleton is used. + If `seed` is an int, a new ``RandomState`` instance is used, + seeded with `seed`. + If `seed` is already a ``Generator`` or ``RandomState`` instance + then that instance is used. + + Examples + -------- + >>> from scipy.stats import unitary_group + >>> x = unitary_group(3) + >>> x.rvs() + + """ # numpy/numpydoc#87 # noqa: E501 + self._dist = unitary_group_gen(seed) + self.dim = self._dist._process_parameters(dim) + + def rvs(self, size=1, random_state=None): + return self._dist.rvs(self.dim, size, random_state) + + +_mvt_doc_default_callparams = """\ +loc : array_like, optional + Location of the distribution. (default ``0``) +shape : array_like, optional + Positive semidefinite matrix of the distribution. (default ``1``) +df : float, optional + Degrees of freedom of the distribution; must be greater than zero. + If ``np.inf`` then results are multivariate normal. The default is ``1``. +allow_singular : bool, optional + Whether to allow a singular matrix. (default ``False``) +""" + +_mvt_doc_callparams_note = """\ +Setting the parameter `loc` to ``None`` is equivalent to having `loc` +be the zero-vector. The parameter `shape` can be a scalar, in which case +the shape matrix is the identity times that value, a vector of +diagonal entries for the shape matrix, or a two-dimensional array_like. +""" + +_mvt_doc_frozen_callparams_note = """\ +See class definition for a detailed description of parameters.""" + +mvt_docdict_params = { + '_mvt_doc_default_callparams': _mvt_doc_default_callparams, + '_mvt_doc_callparams_note': _mvt_doc_callparams_note, + '_doc_random_state': _doc_random_state +} + +mvt_docdict_noparams = { + '_mvt_doc_default_callparams': "", + '_mvt_doc_callparams_note': _mvt_doc_frozen_callparams_note, + '_doc_random_state': _doc_random_state +} + + +class multivariate_t_gen(multi_rv_generic): + r""" + A multivariate t-distributed random variable. + + The `loc` parameter specifies the location. The `shape` parameter specifies + the positive semidefinite shape matrix. The `df` parameter specifies the + degrees of freedom. + + In addition to calling the methods below, the object itself may be called + as a function to fix the location, shape matrix, and degrees of freedom + parameters, returning a "frozen" multivariate t-distribution random. + + Parameters + ---------- + %(_mvt_doc_default_callparams)s + %(_doc_random_state)s + + Methods + ------- + pdf(x, loc=None, shape=1, df=1, allow_singular=False) + Probability density function. + logpdf(x, loc=None, shape=1, df=1, allow_singular=False) + Log of the probability density function. + cdf(x, loc=None, shape=1, df=1, allow_singular=False, *, + maxpts=None, lower_limit=None, random_state=None) + Cumulative distribution function. + rvs(loc=None, shape=1, df=1, size=1, random_state=None) + Draw random samples from a multivariate t-distribution. + entropy(loc=None, shape=1, df=1) + Differential entropy of a multivariate t-distribution. + marginal(dimensions, loc=None, shape=1, df=1, allow_singular=False) + Return a marginal multivariate t-distribution. + + Notes + ----- + %(_mvt_doc_callparams_note)s + The matrix `shape` must be a (symmetric) positive semidefinite matrix. The + determinant and inverse of `shape` are computed as the pseudo-determinant + and pseudo-inverse, respectively, so that `shape` does not need to have + full rank. + + The probability density function for `multivariate_t` is + + .. math:: + + f(x) = \frac{\Gamma((\nu + p)/2)}{\Gamma(\nu/2)\nu^{p/2}\pi^{p/2}|\Sigma|^{1/2}} + \left[1 + \frac{1}{\nu} (\mathbf{x} - \boldsymbol{\mu})^{\top} + \boldsymbol{\Sigma}^{-1} + (\mathbf{x} - \boldsymbol{\mu}) \right]^{-(\nu + p)/2}, + + where :math:`p` is the dimension of :math:`\mathbf{x}`, + :math:`\boldsymbol{\mu}` is the :math:`p`-dimensional location, + :math:`\boldsymbol{\Sigma}` the :math:`p \times p`-dimensional shape + matrix, and :math:`\nu` is the degrees of freedom. + + .. versionadded:: 1.6.0 + + References + ---------- + .. [1] Arellano-Valle et al. "Shannon Entropy and Mutual Information for + Multivariate Skew-Elliptical Distributions". Scandinavian Journal + of Statistics. Vol. 40, issue 1. + + Examples + -------- + The object may be called (as a function) to fix the `loc`, `shape`, + `df`, and `allow_singular` parameters, returning a "frozen" + multivariate_t random variable: + + >>> import numpy as np + >>> from scipy.stats import multivariate_t + >>> rv = multivariate_t([1.0, -0.5], [[2.1, 0.3], [0.3, 1.5]], df=2) + >>> # Frozen object with the same methods but holding the given location, + >>> # scale, and degrees of freedom fixed. + + Create a contour plot of the PDF. + + >>> import matplotlib.pyplot as plt + >>> x, y = np.mgrid[-1:3:.01, -2:1.5:.01] + >>> pos = np.dstack((x, y)) + >>> fig, ax = plt.subplots(1, 1) + >>> ax.set_aspect('equal') + >>> plt.contourf(x, y, rv.pdf(pos)) + """ + + def __init__(self, seed=None): + """Initialize a multivariate t-distributed random variable. + + Parameters + ---------- + seed : Random state. + + """ + super().__init__(seed) + self.__doc__ = doccer.docformat(self.__doc__, mvt_docdict_params) + self._random_state = check_random_state(seed) + + def __call__(self, loc=None, shape=1, df=1, allow_singular=False, + seed=None): + """Create a frozen multivariate t-distribution. + + See `multivariate_t_frozen` for parameters. + """ + if df == np.inf: + return multivariate_normal_frozen(mean=loc, cov=shape, + allow_singular=allow_singular, + seed=seed) + return multivariate_t_frozen(loc=loc, shape=shape, df=df, + allow_singular=allow_singular, seed=seed) + + def pdf(self, x, loc=None, shape=1, df=1, allow_singular=False): + """Multivariate t-distribution probability density function. + + Parameters + ---------- + x : array_like + Points at which to evaluate the probability density function. + %(_mvt_doc_default_callparams)s + + Returns + ------- + pdf : Probability density function evaluated at `x`. + + Examples + -------- + >>> from scipy.stats import multivariate_t + >>> x = [0.4, 5] + >>> loc = [0, 1] + >>> shape = [[1, 0.1], [0.1, 1]] + >>> df = 7 + >>> multivariate_t.pdf(x, loc, shape, df) + 0.00075713 + + """ + dim, loc, shape, df = self._process_parameters(loc, shape, df) + x = self._process_quantiles(x, dim) + shape_info = _PSD(shape, allow_singular=allow_singular) + logpdf = self._logpdf(x, loc, shape_info.U, shape_info.log_pdet, df, + dim, shape_info.rank) + return np.exp(logpdf) + + def logpdf(self, x, loc=None, shape=1, df=1): + """Log of the multivariate t-distribution probability density function. + + Parameters + ---------- + x : array_like + Points at which to evaluate the log of the probability density + function. + %(_mvt_doc_default_callparams)s + + Returns + ------- + logpdf : Log of the probability density function evaluated at `x`. + + Examples + -------- + >>> from scipy.stats import multivariate_t + >>> x = [0.4, 5] + >>> loc = [0, 1] + >>> shape = [[1, 0.1], [0.1, 1]] + >>> df = 7 + >>> multivariate_t.logpdf(x, loc, shape, df) + -7.1859802 + + See Also + -------- + pdf : Probability density function. + + """ + dim, loc, shape, df = self._process_parameters(loc, shape, df) + x = self._process_quantiles(x, dim) + shape_info = _PSD(shape) + cov_object = _covariance.CovViaPSD(shape_info) + + return self._logpdf(x, loc, shape_info.U, shape_info.log_pdet, df, dim, + shape_info.rank, cov_object) + + def _logpdf(self, x, loc, prec_U, log_pdet, df, dim, rank, cov_object=None): + """Utility method `pdf`, `logpdf` for parameters. + + Parameters + ---------- + x : ndarray + Points at which to evaluate the log of the probability density + function. + loc : ndarray + Location of the distribution. + prec_U : ndarray + A decomposition such that `np.dot(prec_U, prec_U.T)` is the inverse + of the shape matrix. + log_pdet : float + Logarithm of the determinant of the shape matrix. + df : float + Degrees of freedom of the distribution. + dim : int + Dimension of the quantiles x. + rank : int + Rank of the shape matrix. + + Notes + ----- + As this function does no argument checking, it should not be called + directly; use 'logpdf' instead. + + """ + if df == np.inf: + return multivariate_normal._logpdf(x, loc, cov_object) + + dev = x - loc + maha = np.square(np.dot(dev, prec_U)).sum(axis=-1) + + t = 0.5 * (df + dim) + A = gammaln(t) + B = gammaln(0.5 * df) + C = dim/2. * np.log(df * np.pi) + D = 0.5 * log_pdet + E = -t * np.log(1 + (1./df) * maha) + + return _squeeze_output(A - B - C - D + E) + + def _cdf(self, x, loc, shape, df, dim, maxpts=None, lower_limit=None, + random_state=None): + + # All of this - random state validation, maxpts, apply_along_axis, + # etc. needs to go in this private method unless we want + # frozen distribution's `cdf` method to duplicate it or call `cdf`, + # which would require re-processing parameters + if random_state is not None: + rng = check_random_state(random_state) + else: + rng = self._random_state + + if not maxpts: + maxpts = 1000 * dim + + x = self._process_quantiles(x, dim) + lower_limit = (np.full(loc.shape, -np.inf) + if lower_limit is None else lower_limit) + + # remove the mean + x, lower_limit = x - loc, lower_limit - loc + + b, a = np.broadcast_arrays(x, lower_limit) + i_swap = b < a + signs = (-1)**(i_swap.sum(axis=-1)) # odd # of swaps -> negative + a, b = a.copy(), b.copy() + a[i_swap], b[i_swap] = b[i_swap], a[i_swap] + n = x.shape[-1] + limits = np.concatenate((a, b), axis=-1) + + def func1d(limits): + a, b = limits[:n], limits[n:] + return _qmvt(maxpts, df, shape, a, b, rng)[0] + + res = np.apply_along_axis(func1d, -1, limits) * signs + # Fixing the output shape for existing distributions is a separate + # issue. For now, let's keep this consistent with pdf. + return _squeeze_output(res) + + def cdf(self, x, loc=None, shape=1, df=1, allow_singular=False, *, + maxpts=None, lower_limit=None, random_state=None): + """Multivariate t-distribution cumulative distribution function. + + Parameters + ---------- + x : array_like + Points at which to evaluate the cumulative distribution function. + %(_mvt_doc_default_callparams)s + maxpts : int, optional + Maximum number of points to use for integration. The default is + 1000 times the number of dimensions. + lower_limit : array_like, optional + Lower limit of integration of the cumulative distribution function. + Default is negative infinity. Must be broadcastable with `x`. + %(_doc_random_state)s + + Returns + ------- + cdf : ndarray or scalar + Cumulative distribution function evaluated at `x`. + + Examples + -------- + >>> from scipy.stats import multivariate_t + >>> x = [0.4, 5] + >>> loc = [0, 1] + >>> shape = [[1, 0.1], [0.1, 1]] + >>> df = 7 + >>> multivariate_t.cdf(x, loc, shape, df) + 0.64798491 + + """ + dim, loc, shape, df = self._process_parameters(loc, shape, df) + shape = _PSD(shape, allow_singular=allow_singular)._M + + return self._cdf(x, loc, shape, df, dim, maxpts, + lower_limit, random_state) + + def _entropy(self, dim, df=1, shape=1): + if df == np.inf: + return multivariate_normal(None, cov=shape).entropy() + + shape_info = _PSD(shape) + shape_term = 0.5 * shape_info.log_pdet + + def regular(dim, df): + halfsum = 0.5 * (dim + df) + half_df = 0.5 * df + return ( + -gammaln(halfsum) + gammaln(half_df) + + 0.5 * dim * np.log(df * np.pi) + halfsum + * (psi(halfsum) - psi(half_df)) + + shape_term + ) + + def asymptotic(dim, df): + # Formula from Wolfram Alpha: + # "asymptotic expansion -gammaln((m+d)/2) + gammaln(d/2) + (m*log(d*pi))/2 + # + ((m+d)/2) * (digamma((m+d)/2) - digamma(d/2))" + return ( + dim * norm._entropy() + dim / df + - dim * (dim - 2) * df**-2.0 / 4 + + dim**2 * (dim - 2) * df**-3.0 / 6 + + dim * (-3 * dim**3 + 8 * dim**2 - 8) * df**-4.0 / 24 + + dim**2 * (3 * dim**3 - 10 * dim**2 + 16) * df**-5.0 / 30 + + shape_term + )[()] + + # preserves ~12 digits accuracy up to at least `dim=1e5`. See gh-18465. + threshold = dim * 100 * 4 / (np.log(dim) + 1) + return xpx.apply_where(df >= threshold, (dim, df), asymptotic, regular) + + def entropy(self, loc=None, shape=1, df=1): + """Calculate the differential entropy of a multivariate + t-distribution. + + Parameters + ---------- + %(_mvt_doc_default_callparams)s + + Returns + ------- + h : float + Differential entropy + + """ + dim, loc, shape, df = self._process_parameters(None, shape, df) + return self._entropy(dim, df, shape) + + def rvs(self, loc=None, shape=1, df=1, size=1, random_state=None): + """Draw random samples from a multivariate t-distribution. + + Parameters + ---------- + %(_mvt_doc_default_callparams)s + size : int, optional + Number of samples to draw (default 1). + %(_doc_random_state)s + + Returns + ------- + rvs : ndarray or scalar + Random variates of size (`size`, `P`), where `P` is the + dimension of the random variable. + + Examples + -------- + >>> from scipy.stats import multivariate_t + >>> x = [0.4, 5] + >>> loc = [0, 1] + >>> shape = [[1, 0.1], [0.1, 1]] + >>> df = 7 + >>> multivariate_t.rvs(loc, shape, df) + array([[0.93477495, 3.00408716]]) + + """ + # For implementation details, see equation (3): + # + # Hofert, "On Sampling from the Multivariatet Distribution", 2013 + # http://rjournal.github.io/archive/2013-2/hofert.pdf + # + dim, loc, shape, df = self._process_parameters(loc, shape, df) + if random_state is not None: + rng = check_random_state(random_state) + else: + rng = self._random_state + + if np.isinf(df): + x = np.ones(size) + else: + x = rng.chisquare(df, size=size) / df + + z = rng.multivariate_normal(np.zeros(dim), shape, size=size) + samples = loc + z / np.sqrt(x)[..., None] + return _squeeze_output(samples) + + def _process_quantiles(self, x, dim): + """ + Adjust quantiles array so that last axis labels the components of + each data point. + """ + x = np.asarray(x, dtype=float) + if x.ndim == 0: + x = x[np.newaxis] + elif x.ndim == 1: + if dim == 1: + x = x[:, np.newaxis] + else: + x = x[np.newaxis, :] + return x + + def _process_parameters(self, loc, shape, df): + """ + Infer dimensionality from location array and shape matrix, handle + defaults, and ensure compatible dimensions. + """ + if loc is None and shape is None: + loc = np.asarray(0, dtype=float) + shape = np.asarray(1, dtype=float) + dim = 1 + elif loc is None: + shape = np.asarray(shape, dtype=float) + if shape.ndim < 2: + dim = 1 + else: + dim = shape.shape[0] + loc = np.zeros(dim) + elif shape is None: + loc = np.asarray(loc, dtype=float) + dim = loc.size + shape = np.eye(dim) + else: + shape = np.asarray(shape, dtype=float) + loc = np.asarray(loc, dtype=float) + dim = loc.size + + if dim == 1: + loc = loc.reshape(1) + shape = shape.reshape(1, 1) + + if loc.ndim != 1 or loc.shape[0] != dim: + raise ValueError(f"Array 'loc' must be a vector of length {dim}.") + if shape.ndim == 0: + shape = shape * np.eye(dim) + elif shape.ndim == 1: + shape = np.diag(shape) + elif shape.ndim == 2 and shape.shape != (dim, dim): + rows, cols = shape.shape + if rows != cols: + msg = ("Array 'cov' must be square if it is two dimensional," + f" but cov.shape = {str(shape.shape)}.") + else: + msg = ("Dimension mismatch: array 'cov' is of shape %s," + " but 'loc' is a vector of length %d.") + msg = msg % (str(shape.shape), len(loc)) + raise ValueError(msg) + elif shape.ndim > 2: + raise ValueError(f"Array 'cov' must be at most two-dimensional, " + f"but cov.ndim = {shape.ndim}") + + # Process degrees of freedom. + if df is None: + df = 1 + elif df <= 0: + raise ValueError("'df' must be greater than zero.") + elif np.isnan(df): + raise ValueError("'df' is 'nan' but must be greater than zero or 'np.inf'.") + + return dim, loc, shape, df + + def marginal(self, dimensions, loc=None, shape=1, df=1, allow_singular=False): + """Return a marginal multivariate t-distribution. + + Parameters + ---------- + dimensions : int or 1-d array_like + The dimensions of the multivariate t corresponding + with the marginal variables, that is, the indices of the dimensions + that are being retained. The other dimensions are marginalized out. + %(_mvt_doc_default_callparams)s + + Returns + ------- + marginal_multivariate_t : multivariate_t_frozen + An object representing the marginal t-distribution. + + Notes + ----- + %(_mvt_doc_frozen_callparams_note)s + """ + params = self._process_parameters(loc, shape, df) + n, loc, shape, df = params + dims = _validate_marginal_input(dimensions, n) + + loc = loc[dims] + shape = shape[np.ix_(dims, dims)] + + return multivariate_t_frozen(loc, shape, df, allow_singular) + + +class multivariate_t_frozen(multi_rv_frozen): + __class_getitem__ = None # pyrefly:ignore[bad-assignment] + + def __init__(self, loc=None, shape=1, df=1, allow_singular=False, + seed=None): + """Create a frozen multivariate t distribution. + + Parameters + ---------- + %(_mvt_doc_default_callparams)s + + Examples + -------- + >>> import numpy as np + >>> from scipy.stats import multivariate_t + >>> loc = np.zeros(3) + >>> shape = np.eye(3) + >>> df = 10 + >>> dist = multivariate_t(loc, shape, df) + >>> dist.rvs() + array([[ 0.81412036, -1.53612361, 0.42199647]]) + >>> dist.pdf([1, 1, 1]) + array([0.01237803]) + + """ + self._dist = multivariate_t_gen(seed) + dim, loc, shape, df = self._dist._process_parameters(loc, shape, df) + self.dim, self.loc, self.shape, self.df = dim, loc, shape, df + self.shape_info = _PSD(shape, allow_singular=allow_singular) + self.allow_singular = allow_singular + + def logpdf(self, x): + x = self._dist._process_quantiles(x, self.dim) + U = self.shape_info.U + log_pdet = self.shape_info.log_pdet + return self._dist._logpdf(x, self.loc, U, log_pdet, self.df, self.dim, + self.shape_info.rank) + + def cdf(self, x, *, maxpts=None, lower_limit=None, random_state=None): + x = self._dist._process_quantiles(x, self.dim) + return self._dist._cdf(x, self.loc, self.shape, self.df, self.dim, + maxpts, lower_limit, random_state) + + def pdf(self, x): + return np.exp(self.logpdf(x)) + + def rvs(self, size=1, random_state=None): + return self._dist.rvs(loc=self.loc, + shape=self.shape, + df=self.df, + size=size, + random_state=random_state) + + def entropy(self): + return self._dist._entropy(self.dim, self.df, self.shape) + + def marginal(self, dimensions): + return self._dist.marginal(dimensions, self.loc, + self.shape, self.df, self.allow_singular) + +multivariate_t = multivariate_t_gen() + + +# Set frozen generator docstrings from corresponding docstrings in +# multivariate_t_gen and fill in default strings in class docstrings +for name in ['logpdf', 'pdf', 'rvs', 'cdf', 'entropy']: + method = multivariate_t_gen.__dict__[name] + method_frozen = multivariate_t_frozen.__dict__[name] + method_frozen.__doc__ = doccer.docformat(method.__doc__, + mvt_docdict_noparams) + method.__doc__ = doccer.docformat(method.__doc__, mvt_docdict_params) + + +_mhg_doc_default_callparams = """\ +m : array_like + The number of each type of object in the population. + That is, :math:`m[i]` is the number of objects of + type :math:`i`. +n : array_like + The number of samples taken from the population. +""" + +_mhg_doc_callparams_note = """\ +`m` must be an array of positive integers. If the quantile +:math:`i` contains values out of the range :math:`[0, m_i]` +where :math:`m_i` is the number of objects of type :math:`i` +in the population or if the parameters are inconsistent with one +another (e.g. ``x.sum() != n``), methods return the appropriate +value (e.g. ``0`` for ``pmf``). If `m` or `n` contain negative +values, the result will contain ``nan`` there. +""" + +_mhg_doc_frozen_callparams = "" + +_mhg_doc_frozen_callparams_note = """\ +See class definition for a detailed description of parameters.""" + +mhg_docdict_params = { + '_doc_default_callparams': _mhg_doc_default_callparams, + '_doc_callparams_note': _mhg_doc_callparams_note, + '_doc_random_state': _doc_random_state +} + +mhg_docdict_noparams = { + '_doc_default_callparams': _mhg_doc_frozen_callparams, + '_doc_callparams_note': _mhg_doc_frozen_callparams_note, + '_doc_random_state': _doc_random_state +} + + +class multivariate_hypergeom_gen(multi_rv_generic): + r""" + A multivariate hypergeometric random variable. + + Parameters + ---------- + %(_doc_default_callparams)s + %(_doc_random_state)s + + Methods + ------- + pmf(x, m, n) + Probability mass function. + logpmf(x, m, n) + Log of the probability mass function. + rvs(m, n, size=1, random_state=None) + Draw random samples from a multivariate hypergeometric + distribution. + mean(m, n) + Mean of the multivariate hypergeometric distribution. + var(m, n) + Variance of the multivariate hypergeometric distribution. + cov(m, n) + Compute the covariance matrix of the multivariate + hypergeometric distribution. + + See Also + -------- + scipy.stats.hypergeom : The hypergeometric distribution. + scipy.stats.multinomial : The multinomial distribution. + + Notes + ----- + %(_doc_callparams_note)s + + The probability mass function for `multivariate_hypergeom` is + + .. math:: + + P(X_1 = x_1, X_2 = x_2, \ldots, X_k = x_k) = \frac{\binom{m_1}{x_1} + \binom{m_2}{x_2} \cdots \binom{m_k}{x_k}}{\binom{M}{n}}, \\ \quad + (x_1, x_2, \ldots, x_k) \in \mathbb{N}^k \text{ with } + \sum_{i=1}^k x_i = n + + where :math:`m_i` are the number of objects of type :math:`i`, :math:`M` + is the total number of objects in the population (sum of all the + :math:`m_i`), and :math:`n` is the size of the sample to be taken + from the population. + + .. versionadded:: 1.6.0 + + References + ---------- + .. [1] The Multivariate Hypergeometric Distribution, + http://www.randomservices.org/random/urn/MultiHypergeometric.html + .. [2] Thomas J. Sargent and John Stachurski, 2020, + Multivariate Hypergeometric Distribution + https://python.quantecon.org/multi_hyper.html + + Examples + -------- + To evaluate the probability mass function of the multivariate + hypergeometric distribution, with a dichotomous population of size + :math:`10` and :math:`20`, at a sample of size :math:`12` with + :math:`8` objects of the first type and :math:`4` objects of the + second type, use: + + >>> from scipy.stats import multivariate_hypergeom + >>> multivariate_hypergeom.pmf(x=[8, 4], m=[10, 20], n=12) + 0.0025207176631464523 + + The `multivariate_hypergeom` distribution is identical to the + corresponding `hypergeom` distribution (tiny numerical differences + notwithstanding) when only two types (good and bad) of objects + are present in the population as in the example above. Consider + another example for a comparison with the hypergeometric distribution: + + >>> from scipy.stats import hypergeom + >>> multivariate_hypergeom.pmf(x=[3, 1], m=[10, 5], n=4) + 0.4395604395604395 + >>> hypergeom.pmf(k=3, M=15, n=4, N=10) + 0.43956043956044005 + + The functions ``pmf``, ``logpmf``, ``mean``, ``var``, ``cov``, and ``rvs`` + support broadcasting, under the convention that the vector parameters + (``x``, ``m``, and ``n``) are interpreted as if each row along the last + axis is a single object. For instance, we can combine the previous two + calls to `multivariate_hypergeom` as + + >>> multivariate_hypergeom.pmf(x=[[8, 4], [3, 1]], m=[[10, 20], [10, 5]], + ... n=[12, 4]) + array([0.00252072, 0.43956044]) + + This broadcasting also works for ``cov``, where the output objects are + square matrices of size ``m.shape[-1]``. For example: + + >>> multivariate_hypergeom.cov(m=[[7, 9], [10, 15]], n=[8, 12]) + array([[[ 1.05, -1.05], + [-1.05, 1.05]], + [[ 1.56, -1.56], + [-1.56, 1.56]]]) + + That is, ``result[0]`` is equal to + ``multivariate_hypergeom.cov(m=[7, 9], n=8)`` and ``result[1]`` is equal + to ``multivariate_hypergeom.cov(m=[10, 15], n=12)``. + + Alternatively, the object may be called (as a function) to fix the `m` + and `n` parameters, returning a "frozen" multivariate hypergeometric + random variable. + + >>> rv = multivariate_hypergeom(m=[10, 20], n=12) + >>> rv.pmf(x=[8, 4]) + 0.0025207176631464523 + """ + def __init__(self, seed=None): + super().__init__(seed) + self.__doc__ = doccer.docformat(self.__doc__, mhg_docdict_params) + + def __call__(self, m, n, seed=None): + """Create a frozen multivariate_hypergeom distribution. + + See `multivariate_hypergeom_frozen` for more information. + """ + return multivariate_hypergeom_frozen(m, n, seed=seed) + + def _process_parameters(self, m, n): + m = np.asarray(m) + n = np.asarray(n) + if m.size == 0: + m = m.astype(int) + if n.size == 0: + n = n.astype(int) + if not np.issubdtype(m.dtype, np.integer): + raise TypeError("'m' must an array of integers.") + if not np.issubdtype(n.dtype, np.integer): + raise TypeError("'n' must an array of integers.") + if m.ndim == 0: + raise ValueError("'m' must be an array with" + " at least one dimension.") + + # check for empty arrays + if m.size != 0: + n = n[..., np.newaxis] + + m, n = np.broadcast_arrays(m, n) + + # check for empty arrays + if m.size != 0: + n = n[..., 0] + + mcond = m < 0 + + M = m.sum(axis=-1) + + ncond = (n < 0) | (n > M) + return M, m, n, mcond, ncond, np.any(mcond, axis=-1) | ncond + + def _process_quantiles(self, x, M, m, n): + x = np.asarray(x) + if not np.issubdtype(x.dtype, np.integer): + raise TypeError("'x' must an array of integers.") + if x.ndim == 0: + raise ValueError("'x' must be an array with" + " at least one dimension.") + if not x.shape[-1] == m.shape[-1]: + raise ValueError(f"Size of each quantile must be size of 'm': " + f"received {x.shape[-1]}, " + f"but expected {m.shape[-1]}.") + + # check for empty arrays + if m.size != 0: + n = n[..., np.newaxis] + M = M[..., np.newaxis] + + x, m, n, M = np.broadcast_arrays(x, m, n, M) + + # check for empty arrays + if m.size != 0: + n, M = n[..., 0], M[..., 0] + + xcond = (x < 0) | (x > m) + return (x, M, m, n, xcond, + np.any(xcond, axis=-1) | (x.sum(axis=-1) != n)) + + def _checkresult(self, result, cond, bad_value): + result = np.asarray(result) + if cond.ndim != 0: + result[cond] = bad_value + elif cond: + return bad_value + if result.ndim == 0: + return result[()] + return result + + def _logpmf(self, x, M, m, n, mxcond, ncond): + # This equation of the pmf comes from the relation, + # n combine r = beta(n+1, 1) / beta(r+1, n-r+1) + num = np.zeros_like(m, dtype=np.float64) + den = np.zeros_like(n, dtype=np.float64) + m, x = m[~mxcond], x[~mxcond] + M, n = M[~ncond], n[~ncond] + num[~mxcond] = (betaln(m+1, 1) - betaln(x+1, m-x+1)) + den[~ncond] = (betaln(M+1, 1) - betaln(n+1, M-n+1)) + num[mxcond] = np.nan + den[ncond] = np.nan + num = num.sum(axis=-1) + return num - den + + def logpmf(self, x, m, n): + """Log of the multivariate hypergeometric probability mass function. + + Parameters + ---------- + x : array_like + Quantiles, with the last axis of `x` denoting the components. + %(_doc_default_callparams)s + + Returns + ------- + logpmf : ndarray or scalar + Log of the probability mass function evaluated at `x` + + Notes + ----- + %(_doc_callparams_note)s + """ + M, m, n, mcond, ncond, mncond = self._process_parameters(m, n) + (x, M, m, n, xcond, + xcond_reduced) = self._process_quantiles(x, M, m, n) + mxcond = mcond | xcond + ncond = ncond | np.zeros(n.shape, dtype=np.bool_) + + result = self._logpmf(x, M, m, n, mxcond, ncond) + + # replace values for which x was out of the domain; broadcast + # xcond to the right shape + xcond_ = xcond_reduced | np.zeros(mncond.shape, dtype=np.bool_) + result = self._checkresult(result, xcond_, -np.inf) + + # replace values bad for n or m; broadcast + # mncond to the right shape + mncond_ = mncond | np.zeros(xcond_reduced.shape, dtype=np.bool_) + return self._checkresult(result, mncond_, np.nan) + + def pmf(self, x, m, n): + """Multivariate hypergeometric probability mass function. + + Parameters + ---------- + x : array_like + Quantiles, with the last axis of `x` denoting the components. + %(_doc_default_callparams)s + + Returns + ------- + pmf : ndarray or scalar + Probability density function evaluated at `x` + + Notes + ----- + %(_doc_callparams_note)s + """ + out = np.exp(self.logpmf(x, m, n)) + return out + + def mean(self, m, n): + """Mean of the multivariate hypergeometric distribution. + + Parameters + ---------- + %(_doc_default_callparams)s + + Returns + ------- + mean : array_like or scalar + The mean of the distribution + """ + M, m, n, _, _, mncond = self._process_parameters(m, n) + # check for empty arrays + if m.size != 0: + M, n = M[..., np.newaxis], n[..., np.newaxis] + cond = (M == 0) + M = np.ma.masked_array(M, mask=cond) + mu = n*(m/M) + if m.size != 0: + mncond = (mncond[..., np.newaxis] | + np.zeros(mu.shape, dtype=np.bool_)) + return self._checkresult(mu, mncond, np.nan) + + def var(self, m, n): + """Variance of the multivariate hypergeometric distribution. + + Parameters + ---------- + %(_doc_default_callparams)s + + Returns + ------- + array_like + The variances of the components of the distribution. This is + the diagonal of the covariance matrix of the distribution + """ + M, m, n, _, _, mncond = self._process_parameters(m, n) + # check for empty arrays + if m.size != 0: + M, n = M[..., np.newaxis], n[..., np.newaxis] + cond = (M == 0) & (M-1 == 0) + M = np.ma.masked_array(M, mask=cond) + output = n * m/M * (M-m)/M * (M-n)/(M-1) + if m.size != 0: + mncond = (mncond[..., np.newaxis] | + np.zeros(output.shape, dtype=np.bool_)) + return self._checkresult(output, mncond, np.nan) + + def cov(self, m, n): + """Covariance matrix of the multivariate hypergeometric distribution. + + Parameters + ---------- + %(_doc_default_callparams)s + + Returns + ------- + cov : array_like + The covariance matrix of the distribution + """ + # see [1]_ for the formula and [2]_ for implementation + # cov( x_i,x_j ) = -n * (M-n)/(M-1) * (K_i*K_j) / (M**2) + M, m, n, _, _, mncond = self._process_parameters(m, n) + # check for empty arrays + if m.size != 0: + M = M[..., np.newaxis, np.newaxis] + n = n[..., np.newaxis, np.newaxis] + cond = (M == 0) & (M-1 == 0) + M = np.ma.masked_array(M, mask=cond) + output = (-n * (M-n)/(M-1) * + np.einsum("...i,...j->...ij", m, m) / (M**2)) + # check for empty arrays + if m.size != 0: + M, n = M[..., 0, 0], n[..., 0, 0] + cond = cond[..., 0, 0] + dim = m.shape[-1] + # diagonal entries need to be computed differently + for i in range(dim): + output[..., i, i] = (n * (M-n) * m[..., i]*(M-m[..., i])) + output[..., i, i] = output[..., i, i] / (M-1) + output[..., i, i] = output[..., i, i] / (M**2) + if m.size != 0: + mncond = (mncond[..., np.newaxis, np.newaxis] | + np.zeros(output.shape, dtype=np.bool_)) + return self._checkresult(output, mncond, np.nan) + + def rvs(self, m, n, size=None, random_state=None): + """Draw random samples from a multivariate hypergeometric distribution. + + Parameters + ---------- + %(_doc_default_callparams)s + size : int or iterable of integers, optional + Number of samples to draw. Default is ``None``, in which case a + single variate is returned as an array with shape ``m.shape``. + %(_doc_random_state)s + + Returns + ------- + rvs : array_like + Random variates of shape ``size`` or ``m.shape`` + (if ``size=None``). + + Notes + ----- + %(_doc_callparams_note)s + + Also note that NumPy's `multivariate_hypergeometric` sampler is not + used as it doesn't support broadcasting. + """ + M, m, n, _, _, _ = self._process_parameters(m, n) + + random_state = self._get_random_state(random_state) + + if size is not None and isinstance(size, int): + size = (size, ) + + if size is None: + rvs = np.empty(m.shape, dtype=m.dtype) + else: + rvs = np.empty(size + (m.shape[-1], ), dtype=m.dtype) + rem = M + + # This sampler has been taken from numpy gh-13794 + # https://github.com/numpy/numpy/pull/13794 + for c in range(m.shape[-1] - 1): + rem = rem - m[..., c] + n0mask = n == 0 + rvs[..., c] = (~n0mask * + random_state.hypergeometric(m[..., c], + rem + n0mask, + n + n0mask, + size=size)) + n = n - rvs[..., c] + rvs[..., m.shape[-1] - 1] = n + + return rvs + + +multivariate_hypergeom = multivariate_hypergeom_gen() + + +class multivariate_hypergeom_frozen(multi_rv_frozen): + def __init__(self, m, n, seed=None): + self._dist = multivariate_hypergeom_gen(seed) + (self.M, self.m, self.n, + self.mcond, self.ncond, + self.mncond) = self._dist._process_parameters(m, n) + + # monkey patch self._dist + def _process_parameters(m, n): + return (self.M, self.m, self.n, + self.mcond, self.ncond, + self.mncond) + self._dist._process_parameters = _process_parameters + + def logpmf(self, x): + return self._dist.logpmf(x, self.m, self.n) + + def pmf(self, x): + return self._dist.pmf(x, self.m, self.n) + + def mean(self): + return self._dist.mean(self.m, self.n) + + def var(self): + return self._dist.var(self.m, self.n) + + def cov(self): + return self._dist.cov(self.m, self.n) + + def rvs(self, size=1, random_state=None): + return self._dist.rvs(self.m, self.n, + size=size, + random_state=random_state) + + +# Set frozen generator docstrings from corresponding docstrings in +# multivariate_hypergeom and fill in default strings in class docstrings +for name in ['logpmf', 'pmf', 'mean', 'var', 'cov', 'rvs']: + method = multivariate_hypergeom_gen.__dict__[name] + method_frozen = multivariate_hypergeom_frozen.__dict__[name] + method_frozen.__doc__ = doccer.docformat( + method.__doc__, mhg_docdict_noparams) + method.__doc__ = doccer.docformat(method.__doc__, + mhg_docdict_params) + + +class random_table_gen(multi_rv_generic): + r""" + Contingency tables from independent samples with fixed marginal sums. + + This is the distribution of random tables with given row and column vector + sums. This distribution represents the set of random tables under the null + hypothesis that rows and columns are independent. It is used in hypothesis + tests of independence. + + Because of assumed independence, the expected frequency of each table + element can be computed from the row and column sums, so that the + distribution is completely determined by these two vectors. + + Parameters + ---------- + %(_doc_row_col)s + %(_doc_random_state)s + + Methods + ------- + logpmf(x) + Log-probability of table `x` to occur in the distribution. + pmf(x) + Probability of table `x` to occur in the distribution. + mean(row, col) + Mean table. + rvs(row, col, size=None, method=None, random_state=None) + Draw random tables with given row and column vector sums. + + Notes + ----- + %(_doc_row_col_note)s + + Random elements from the distribution are generated either with Boyett's + [1]_ or Patefield's algorithm [2]_. Boyett's algorithm has + O(N) time and space complexity, where N is the total sum of entries in the + table. Patefield's algorithm has O(K x log(N)) time complexity, where K is + the number of cells in the table and requires only a small constant work + space. By default, the `rvs` method selects the fastest algorithm based on + the input, but you can specify the algorithm with the keyword `method`. + Allowed values are "boyett" and "patefield". + + .. versionadded:: 1.10.0 + + References + ---------- + .. [1] J. Boyett, AS 144 Appl. Statist. 28 (1979) 329-332 + .. [2] W.M. Patefield, AS 159 Appl. Statist. 30 (1981) 91-97 + + Examples + -------- + >>> from scipy.stats import random_table + + >>> row = [1, 5] + >>> col = [2, 3, 1] + >>> random_table.mean(row, col) + array([[0.33333333, 0.5 , 0.16666667], + [1.66666667, 2.5 , 0.83333333]]) + + Alternatively, the object may be called (as a function) to fix the row + and column vector sums, returning a "frozen" distribution. + + >>> dist = random_table(row, col) + >>> dist.rvs(random_state=123) + array([[1, 0, 0], + [1, 3, 1]]) + """ + + def __init__(self, seed=None): + super().__init__(seed) + + def __call__(self, row, col, *, seed=None): + """Create a frozen distribution of tables with given marginals. + + See `random_table_frozen` for more information. + """ + return random_table_frozen(row, col, seed=seed) + + def logpmf(self, x, row, col): + """Log-probability of table to occur in the distribution. + + Parameters + ---------- + %(_doc_x)s + %(_doc_row_col)s + + Returns + ------- + logpmf : ndarray or scalar + Log of the probability mass function evaluated at `x`. + + Notes + ----- + %(_doc_row_col_note)s + + If row and column marginals of `x` do not match `row` and `col`, + negative infinity is returned. + + Examples + -------- + >>> from scipy.stats import random_table + >>> import numpy as np + + >>> x = [[1, 5, 1], [2, 3, 1]] + >>> row = np.sum(x, axis=1) + >>> col = np.sum(x, axis=0) + >>> random_table.logpmf(x, row, col) + -1.6306401200847027 + + Alternatively, the object may be called (as a function) to fix the row + and column vector sums, returning a "frozen" distribution. + + >>> d = random_table(row, col) + >>> d.logpmf(x) + -1.6306401200847027 + """ + r, c, n = self._process_parameters(row, col) + x = np.asarray(x) + + if x.ndim < 2: + raise ValueError("`x` must be at least two-dimensional") + + dtype_is_int = np.issubdtype(x.dtype, np.integer) + with np.errstate(invalid='ignore'): + if not dtype_is_int and not np.all(x.astype(int) == x): + raise ValueError("`x` must contain only integral values") + + # x does not contain NaN if we arrive here + if np.any(x < 0): + raise ValueError("`x` must contain only non-negative values") + + r2 = np.sum(x, axis=-1) + c2 = np.sum(x, axis=-2) + + if r2.shape[-1] != len(r): + raise ValueError("shape of `x` must agree with `row`") + + if c2.shape[-1] != len(c): + raise ValueError("shape of `x` must agree with `col`") + + res = np.empty(x.shape[:-2]) + + mask = np.all(r2 == r, axis=-1) & np.all(c2 == c, axis=-1) + + def lnfac(x): + return gammaln(x + 1) + + res[mask] = (np.sum(lnfac(r), axis=-1) + np.sum(lnfac(c), axis=-1) + - lnfac(n) - np.sum(lnfac(x[mask]), axis=(-1, -2))) + res[~mask] = -np.inf + + return res[()] + + def pmf(self, x, row, col): + """Probability of table to occur in the distribution. + + Parameters + ---------- + %(_doc_x)s + %(_doc_row_col)s + + Returns + ------- + pmf : ndarray or scalar + Probability mass function evaluated at `x`. + + Notes + ----- + %(_doc_row_col_note)s + + If row and column marginals of `x` do not match `row` and `col`, + zero is returned. + + Examples + -------- + >>> from scipy.stats import random_table + >>> import numpy as np + + >>> x = [[1, 5, 1], [2, 3, 1]] + >>> row = np.sum(x, axis=1) + >>> col = np.sum(x, axis=0) + >>> random_table.pmf(x, row, col) + 0.19580419580419592 + + Alternatively, the object may be called (as a function) to fix the row + and column vector sums, returning a "frozen" distribution. + + >>> d = random_table(row, col) + >>> d.pmf(x) + 0.19580419580419592 + """ + return np.exp(self.logpmf(x, row, col)) + + def mean(self, row, col): + """Mean of distribution of conditional tables. + %(_doc_mean_params)s + + Returns + ------- + mean: ndarray + Mean of the distribution. + + Notes + ----- + %(_doc_row_col_note)s + + Examples + -------- + >>> from scipy.stats import random_table + + >>> row = [1, 5] + >>> col = [2, 3, 1] + >>> random_table.mean(row, col) + array([[0.33333333, 0.5 , 0.16666667], + [1.66666667, 2.5 , 0.83333333]]) + + Alternatively, the object may be called (as a function) to fix the row + and column vector sums, returning a "frozen" distribution. + + >>> d = random_table(row, col) + >>> d.mean() + array([[0.33333333, 0.5 , 0.16666667], + [1.66666667, 2.5 , 0.83333333]]) + """ + r, c, n = self._process_parameters(row, col) + return np.outer(r, c) / n + + def rvs(self, row, col, *, size=None, method=None, random_state=None): + """Draw random tables with fixed column and row marginals. + + Parameters + ---------- + %(_doc_row_col)s + size : int, optional + Number of samples to draw (default 1). + method : str, optional + Which method to use, "boyett" or "patefield". If None (default), + selects the fastest method for this input. + %(_doc_random_state)s + + Returns + ------- + rvs : ndarray + Random 2D tables of shape (`size`, `len(row)`, `len(col)`). + + Notes + ----- + %(_doc_row_col_note)s + + Examples + -------- + >>> from scipy.stats import random_table + + >>> row = [1, 5] + >>> col = [2, 3, 1] + >>> random_table.rvs(row, col, random_state=123) + array([[1., 0., 0.], + [1., 3., 1.]]) + + Alternatively, the object may be called (as a function) to fix the row + and column vector sums, returning a "frozen" distribution. + + >>> d = random_table(row, col) + >>> d.rvs(random_state=123) + array([[1., 0., 0.], + [1., 3., 1.]]) + """ + r, c, n = self._process_parameters(row, col) + size, shape = self._process_size_shape(size, r, c) + + random_state = self._get_random_state(random_state) + meth = self._process_rvs_method(method, r, c, n) + + return meth(r, c, n, size, random_state).reshape(shape) + + @staticmethod + def _process_parameters(row, col): + """ + Check that row and column vectors are one-dimensional, that they do + not contain negative or non-integer entries, and that the sums over + both vectors are equal. + """ + r = np.array(row, dtype=np.int64, copy=True) + c = np.array(col, dtype=np.int64, copy=True) + + if np.ndim(r) != 1: + raise ValueError("`row` must be one-dimensional") + if np.ndim(c) != 1: + raise ValueError("`col` must be one-dimensional") + + if np.any(r < 0): + raise ValueError("each element of `row` must be non-negative") + if np.any(c < 0): + raise ValueError("each element of `col` must be non-negative") + + n = np.sum(r) + if n != np.sum(c): + raise ValueError("sums over `row` and `col` must be equal") + + if not np.all(r == np.asarray(row)): + raise ValueError("each element of `row` must be an integer") + if not np.all(c == np.asarray(col)): + raise ValueError("each element of `col` must be an integer") + + return r, c, n + + @staticmethod + def _process_size_shape(size, r, c): + """ + Compute the number of samples to be drawn and the shape of the output + """ + shape = (len(r), len(c)) + + if size is None: + return 1, shape + + size = np.atleast_1d(size) + if not np.issubdtype(size.dtype, np.integer) or np.any(size < 0): + raise ValueError("`size` must be a non-negative integer or `None`") + + return np.prod(size), tuple(size) + shape + + @classmethod + def _process_rvs_method(cls, method, r, c, n): + known_methods = { + None: cls._rvs_select(r, c, n), + "boyett": cls._rvs_boyett, + "patefield": cls._rvs_patefield, + } + try: + return known_methods[method] + except KeyError: + raise ValueError(f"'{method}' not recognized, " + f"must be one of {set(known_methods)}") + + @classmethod + def _rvs_select(cls, r, c, n): + fac = 1.0 # benchmarks show that this value is about 1 + k = len(r) * len(c) # number of cells + # n + 1 guards against failure if n == 0 + if n > fac * np.log(n + 1) * k: + return cls._rvs_patefield + return cls._rvs_boyett + + @staticmethod + def _rvs_boyett(row, col, ntot, size, random_state): + return _rcont.rvs_rcont1(row, col, ntot, size, random_state) + + @staticmethod + def _rvs_patefield(row, col, ntot, size, random_state): + return _rcont.rvs_rcont2(row, col, ntot, size, random_state) + + +random_table = random_table_gen() + + +class random_table_frozen(multi_rv_frozen): + __class_getitem__ = None # pyrefly:ignore[bad-assignment] + + def __init__(self, row, col, *, seed=None): + self._dist = random_table_gen(seed) + self._params = self._dist._process_parameters(row, col) + + # monkey patch self._dist + def _process_parameters(r, c): + return self._params + self._dist._process_parameters = _process_parameters + + def logpmf(self, x): + return self._dist.logpmf(x, None, None) + + def pmf(self, x): + return self._dist.pmf(x, None, None) + + def mean(self): + return self._dist.mean(None, None) + + def rvs(self, size=None, method=None, random_state=None): + # optimisations are possible here + return self._dist.rvs(None, None, size=size, method=method, + random_state=random_state) + + +_ctab_doc_row_col = """\ +row : array_like + Sum of table entries in each row. +col : array_like + Sum of table entries in each column.""" + +_ctab_doc_x = """\ +x : array-like + Two-dimensional table of non-negative integers, or a + multi-dimensional array with the last two dimensions + corresponding with the tables.""" + +_ctab_doc_row_col_note = """\ +The row and column vectors must be one-dimensional, not empty, +and each sum up to the same value. They cannot contain negative +or noninteger entries.""" + +_ctab_doc_mean_params = f""" +Parameters +---------- +{_ctab_doc_row_col}""" + +_ctab_doc_row_col_note_frozen = """\ +See class definition for a detailed description of parameters.""" + +_ctab_docdict = { + "_doc_random_state": _doc_random_state, + "_doc_row_col": _ctab_doc_row_col, + "_doc_x": _ctab_doc_x, + "_doc_mean_params": _ctab_doc_mean_params, + "_doc_row_col_note": _ctab_doc_row_col_note, +} + +_ctab_docdict_frozen = _ctab_docdict.copy() +_ctab_docdict_frozen.update({ + "_doc_row_col": "", + "_doc_mean_params": "", + "_doc_row_col_note": _ctab_doc_row_col_note_frozen, +}) + + +def _docfill(obj, docdict, template=None): + obj.__doc__ = doccer.docformat(template or obj.__doc__, docdict) + + +# Set frozen generator docstrings from corresponding docstrings in +# random_table and fill in default strings in class docstrings +_docfill(random_table_gen, _ctab_docdict) +for name in ['logpmf', 'pmf', 'mean', 'rvs']: + method = random_table_gen.__dict__[name] + method_frozen = random_table_frozen.__dict__[name] + _docfill(method_frozen, _ctab_docdict_frozen, method.__doc__) + _docfill(method, _ctab_docdict) + + +class uniform_direction_gen(multi_rv_generic): + r""" + A vector-valued uniform direction. + + Return a random direction (unit vector). The `dim` keyword specifies + the dimensionality of the space. + + Parameters + ---------- + dim : scalar + Dimension of directions. + seed : {None, int, `numpy.random.Generator`, + `numpy.random.RandomState`}, optional + + Used for drawing random variates. + If `seed` is `None`, the `~np.random.RandomState` singleton is used. + If `seed` is an int, a new ``RandomState`` instance is used, seeded + with seed. + If `seed` is already a ``RandomState`` or ``Generator`` instance, + then that object is used. + Default is `None`. + + Methods + ------- + rvs(dim=None, size=1, random_state=None) + Draw random directions. + + Notes + ----- + This distribution generates unit vectors uniformly distributed on + the surface of a hypersphere. These can be interpreted as random + directions. + For example, if `dim` is 3, 3D vectors from the surface of :math:`S^2` + will be sampled. + + References + ---------- + .. [1] Marsaglia, G. (1972). "Choosing a Point from the Surface of a + Sphere". Annals of Mathematical Statistics. 43 (2): 645-646. + + Examples + -------- + >>> import numpy as np + >>> from scipy.stats import uniform_direction + >>> x = uniform_direction.rvs(3) + >>> np.linalg.norm(x) + 1. + + This generates one random direction, a vector on the surface of + :math:`S^2`. + + Alternatively, the object may be called (as a function) to return a frozen + distribution with fixed `dim` parameter. Here, + we create a `uniform_direction` with ``dim=3`` and draw 5 observations. + The samples are then arranged in an array of shape 5x3. + + >>> rng = np.random.default_rng() + >>> uniform_sphere_dist = uniform_direction(3) + >>> unit_vectors = uniform_sphere_dist.rvs(5, random_state=rng) + >>> unit_vectors + array([[ 0.56688642, -0.1332634 , -0.81294566], + [-0.427126 , -0.74779278, 0.50830044], + [ 0.3793989 , 0.92346629, 0.05715323], + [ 0.36428383, -0.92449076, -0.11231259], + [-0.27733285, 0.94410968, -0.17816678]]) + """ + + def __init__(self, seed=None): + super().__init__(seed) + self.__doc__ = doccer.docformat(self.__doc__) + + def __call__(self, dim=None, seed=None): + """Create a frozen n-dimensional uniform direction distribution. + + See `uniform_direction` for more information. + """ + return uniform_direction_frozen(dim, seed=seed) + + def _process_parameters(self, dim): + """Dimension N must be specified; it cannot be inferred.""" + if dim is None or not np.isscalar(dim) or dim < 1 or dim != int(dim): + raise ValueError("Dimension of vector must be specified, " + "and must be an integer greater than 0.") + + return int(dim) + + def rvs(self, dim, size=None, random_state=None): + """Draw random samples from S(N-1). + + Parameters + ---------- + dim : int + Dimension of space (N). + size : int or tuple of ints, optional + Given a shape of, for example, (m,n,k), m*n*k samples are + generated, and packed in an m-by-n-by-k arrangement. + Because each sample is N-dimensional, the output shape + is (m,n,k,N). If no shape is specified, a single (N-D) + sample is returned. + random_state : {None, int, `numpy.random.Generator`, + `numpy.random.RandomState`}, optional + + Pseudorandom number generator state used to generate resamples. + + If `random_state` is ``None`` (or `np.random`), the + `numpy.random.RandomState` singleton is used. + If `random_state` is an int, a new ``RandomState`` instance is + used, seeded with `random_state`. + If `random_state` is already a ``Generator`` or ``RandomState`` + instance then that instance is used. + + Returns + ------- + rvs : ndarray + Random direction vectors + + """ + random_state = self._get_random_state(random_state) + if size is None: + size = np.array([], dtype=int) + size = np.atleast_1d(size) + + dim = self._process_parameters(dim) + + samples = _sample_uniform_direction(dim, size, random_state) + return samples + + +uniform_direction = uniform_direction_gen() + + +class uniform_direction_frozen(multi_rv_frozen): + def __init__(self, dim=None, seed=None): + """Create a frozen n-dimensional uniform direction distribution. + + Parameters + ---------- + dim : int + Dimension of matrices + seed : {None, int, `numpy.random.Generator`, + `numpy.random.RandomState`}, optional + + If `seed` is None (or `np.random`), the `numpy.random.RandomState` + singleton is used. + If `seed` is an int, a new ``RandomState`` instance is used, + seeded with `seed`. + If `seed` is already a ``Generator`` or ``RandomState`` instance + then that instance is used. + + Examples + -------- + >>> from scipy.stats import uniform_direction + >>> x = uniform_direction(3) + >>> x.rvs() + + """ + self._dist = uniform_direction_gen(seed) + self.dim = self._dist._process_parameters(dim) + + def rvs(self, size=None, random_state=None): + return self._dist.rvs(self.dim, size, random_state) + + +def _sample_uniform_direction(dim, size, random_state): + """ + Private method to generate uniform directions + Reference: Marsaglia, G. (1972). "Choosing a Point from the Surface of a + Sphere". Annals of Mathematical Statistics. 43 (2): 645-646. + """ + samples_shape = np.append(size, dim) + samples = random_state.standard_normal(samples_shape) + samples /= np.linalg.norm(samples, axis=-1, keepdims=True) + return samples + + +_dirichlet_mn_doc_default_callparams = """\ +alpha : array_like + The concentration parameters. The number of entries along the last axis + determines the dimensionality of the distribution. Each entry must be + strictly positive. +n : int or array_like + The number of trials. Each element must be a non-negative integer. +""" + +_dirichlet_mn_doc_frozen_callparams = "" + +_dirichlet_mn_doc_frozen_callparams_note = """\ +See class definition for a detailed description of parameters.""" + +dirichlet_mn_docdict_params = { + '_dirichlet_mn_doc_default_callparams': _dirichlet_mn_doc_default_callparams, + '_doc_random_state': _doc_random_state +} + +dirichlet_mn_docdict_noparams = { + '_dirichlet_mn_doc_default_callparams': _dirichlet_mn_doc_frozen_callparams, + '_doc_random_state': _doc_random_state +} + + +def _dirichlet_multinomial_check_parameters(alpha, n, x=None): + + alpha = np.asarray(alpha) + n = np.asarray(n) + + if x is not None: + # Ensure that `x` and `alpha` are arrays. If the shapes are + # incompatible, NumPy will raise an appropriate error. + try: + x, alpha = np.broadcast_arrays(x, alpha) + except ValueError as e: + msg = "`x` and `alpha` must be broadcastable." + raise ValueError(msg) from e + + x_int = np.floor(x) + if np.any(x < 0) or np.any(x != x_int): + raise ValueError("`x` must contain only non-negative integers.") + x = x_int + + if np.any(alpha <= 0): + raise ValueError("`alpha` must contain only positive values.") + + n_int = np.floor(n) + if np.any(n < 0) or np.any(n != n_int): + raise ValueError("`n` must be a non-negative integer.") + n = n_int + + sum_alpha = np.sum(alpha, axis=-1) + sum_alpha, n = np.broadcast_arrays(sum_alpha, n) + + return (alpha, sum_alpha, n) if x is None else (alpha, sum_alpha, n, x) + + +class dirichlet_multinomial_gen(multi_rv_generic): + r""" + A Dirichlet multinomial random variable. + + The Dirichlet multinomial distribution is a compound probability + distribution: it is the multinomial distribution with number of trials + `n` and class probabilities ``p`` randomly sampled from a Dirichlet + distribution with concentration parameters ``alpha``. + + Parameters + ---------- + %(_dirichlet_mn_doc_default_callparams)s + %(_doc_random_state)s + + Methods + ------- + logpmf(x, alpha, n): + Log of the probability mass function. + pmf(x, alpha, n): + Probability mass function. + mean(alpha, n): + Mean of the Dirichlet multinomial distribution. + var(alpha, n): + Variance of the Dirichlet multinomial distribution. + cov(alpha, n): + The covariance of the Dirichlet multinomial distribution. + + See Also + -------- + scipy.stats.dirichlet : The dirichlet distribution. + scipy.stats.multinomial : The multinomial distribution. + + References + ---------- + .. [1] Dirichlet-multinomial distribution, Wikipedia, + https://www.wikipedia.org/wiki/Dirichlet-multinomial_distribution + + Examples + -------- + >>> from scipy.stats import dirichlet_multinomial + + Get the PMF + + >>> n = 6 # number of trials + >>> alpha = [3, 4, 5] # concentration parameters + >>> x = [1, 2, 3] # counts + >>> dirichlet_multinomial.pmf(x, alpha, n) + 0.08484162895927604 + + If the sum of category counts does not equal the number of trials, + the probability mass is zero. + + >>> dirichlet_multinomial.pmf(x, alpha, n=7) + 0.0 + + Get the log of the PMF + + >>> dirichlet_multinomial.logpmf(x, alpha, n) + -2.4669689491013327 + + Get the mean + + >>> dirichlet_multinomial.mean(alpha, n) + array([1.5, 2. , 2.5]) + + Get the variance + + >>> dirichlet_multinomial.var(alpha, n) + array([1.55769231, 1.84615385, 2.01923077]) + + Get the covariance + + >>> dirichlet_multinomial.cov(alpha, n) + array([[ 1.55769231, -0.69230769, -0.86538462], + [-0.69230769, 1.84615385, -1.15384615], + [-0.86538462, -1.15384615, 2.01923077]]) + + Alternatively, the object may be called (as a function) to fix the + `alpha` and `n` parameters, returning a "frozen" Dirichlet multinomial + random variable. + + >>> dm = dirichlet_multinomial(alpha, n) + >>> dm.pmf(x) + 0.08484162895927579 + + All methods are fully vectorized. Each element of `x` and `alpha` is + a vector (along the last axis), each element of `n` is an + integer (scalar), and the result is computed element-wise. + + >>> x = [[1, 2, 3], [4, 5, 6]] + >>> alpha = [[1, 2, 3], [4, 5, 6]] + >>> n = [6, 15] + >>> dirichlet_multinomial.pmf(x, alpha, n) + array([0.06493506, 0.02626937]) + + >>> dirichlet_multinomial.cov(alpha, n).shape # both covariance matrices + (2, 3, 3) + + Broadcasting according to standard NumPy conventions is supported. Here, + we have four sets of concentration parameters (each a two element vector) + for each of three numbers of trials (each a scalar). + + >>> alpha = [[3, 4], [4, 5], [5, 6], [6, 7]] + >>> n = [[6], [7], [8]] + >>> dirichlet_multinomial.mean(alpha, n).shape + (3, 4, 2) + """ + def __init__(self, seed=None): + super().__init__(seed) + self.__doc__ = doccer.docformat(self.__doc__, + dirichlet_mn_docdict_params) + + def __call__(self, alpha, n, seed=None): + return dirichlet_multinomial_frozen(alpha, n, seed=seed) + + def logpmf(self, x, alpha, n): + """The log of the probability mass function. + + Parameters + ---------- + x: ndarray + Category counts (non-negative integers). Must be broadcastable + with shape parameter ``alpha``. If multidimensional, the last axis + must correspond with the categories. + %(_dirichlet_mn_doc_default_callparams)s + + Returns + ------- + out: ndarray or scalar + Log of the probability mass function. + + """ + + a, Sa, n, x = _dirichlet_multinomial_check_parameters(alpha, n, x) + + out = np.asarray(loggamma(Sa) + loggamma(n + 1) - loggamma(n + Sa)) + out += (loggamma(x + a) - (loggamma(a) + loggamma(x + 1))).sum(axis=-1) + np.place(out, n != x.sum(axis=-1), -np.inf) + return out[()] + + def pmf(self, x, alpha, n): + """Probability mass function for a Dirichlet multinomial distribution. + + Parameters + ---------- + x: ndarray + Category counts (non-negative integers). Must be broadcastable + with shape parameter ``alpha``. If multidimensional, the last axis + must correspond with the categories. + %(_dirichlet_mn_doc_default_callparams)s + + Returns + ------- + out: ndarray or scalar + Probability mass function. + + """ + return np.exp(self.logpmf(x, alpha, n)) + + def mean(self, alpha, n): + """Mean of a Dirichlet multinomial distribution. + + Parameters + ---------- + %(_dirichlet_mn_doc_default_callparams)s + + Returns + ------- + out: ndarray + Mean of a Dirichlet multinomial distribution. + + """ + a, Sa, n = _dirichlet_multinomial_check_parameters(alpha, n) + n, Sa = n[..., np.newaxis], Sa[..., np.newaxis] + return n * a / Sa + + def var(self, alpha, n): + """The variance of the Dirichlet multinomial distribution. + + Parameters + ---------- + %(_dirichlet_mn_doc_default_callparams)s + + Returns + ------- + out: array_like + The variances of the components of the distribution. This is + the diagonal of the covariance matrix of the distribution. + + """ + a, Sa, n = _dirichlet_multinomial_check_parameters(alpha, n) + n, Sa = n[..., np.newaxis], Sa[..., np.newaxis] + return n * a / Sa * (1 - a/Sa) * (n + Sa) / (1 + Sa) + + def cov(self, alpha, n): + """Covariance matrix of a Dirichlet multinomial distribution. + + Parameters + ---------- + %(_dirichlet_mn_doc_default_callparams)s + + Returns + ------- + out : array_like + The covariance matrix of the distribution. + + """ + a, Sa, n = _dirichlet_multinomial_check_parameters(alpha, n) + var = dirichlet_multinomial.var(a, n) + + n, Sa = n[..., np.newaxis, np.newaxis], Sa[..., np.newaxis, np.newaxis] + aiaj = a[..., :, np.newaxis] * a[..., np.newaxis, :] + cov = -n * aiaj / Sa ** 2 * (n + Sa) / (1 + Sa) + + ii = np.arange(cov.shape[-1]) + cov[..., ii, ii] = var + return cov + + +dirichlet_multinomial = dirichlet_multinomial_gen() + + +class dirichlet_multinomial_frozen(multi_rv_frozen): + def __init__(self, alpha, n, seed=None): + alpha, Sa, n = _dirichlet_multinomial_check_parameters(alpha, n) + self.alpha = alpha + self.n = n + self._dist = dirichlet_multinomial_gen(seed) + + def logpmf(self, x): + return self._dist.logpmf(x, self.alpha, self.n) + + def pmf(self, x): + return self._dist.pmf(x, self.alpha, self.n) + + def mean(self): + return self._dist.mean(self.alpha, self.n) + + def var(self): + return self._dist.var(self.alpha, self.n) + + def cov(self): + return self._dist.cov(self.alpha, self.n) + + +# Set frozen generator docstrings from corresponding docstrings in +# dirichlet_multinomial and fill in default strings in class docstrings. +for name in ['logpmf', 'pmf', 'mean', 'var', 'cov']: + method = dirichlet_multinomial_gen.__dict__[name] + method_frozen = dirichlet_multinomial_frozen.__dict__[name] + method_frozen.__doc__ = doccer.docformat( + method.__doc__, dirichlet_mn_docdict_noparams) + method.__doc__ = doccer.docformat(method.__doc__, + dirichlet_mn_docdict_params) + + +class vonmises_fisher_gen(multi_rv_generic): + r""" + A von Mises-Fisher variable. + + The `mu` keyword specifies the mean direction vector. The `kappa` keyword + specifies the concentration parameter. + + Parameters + ---------- + mu : array_like + Mean direction of the distribution. Must be a one-dimensional unit + vector of norm 1. + kappa : float + Concentration parameter. Must be positive. + seed : {None, int, np.random.RandomState, np.random.Generator}, optional + Used for drawing random variates. + If `seed` is `None`, the `~np.random.RandomState` singleton is used. + If `seed` is an int, a new ``RandomState`` instance is used, seeded + with seed. + If `seed` is already a ``RandomState`` or ``Generator`` instance, + then that object is used. + Default is `None`. + + Methods + ------- + pdf(x, mu=None, kappa=1) + Probability density function. + logpdf(x, mu=None, kappa=1) + Log of the probability density function. + rvs(mu=None, kappa=1, size=1, random_state=None) + Draw random samples from a von Mises-Fisher distribution. + entropy(mu=None, kappa=1) + Compute the differential entropy of the von Mises-Fisher distribution. + fit(data) + Fit a von Mises-Fisher distribution to data. + + See Also + -------- + scipy.stats.vonmises : Von-Mises Fisher distribution in 2D on a circle + uniform_direction : uniform distribution on the surface of a hypersphere + + Notes + ----- + The von Mises-Fisher distribution is a directional distribution on the + surface of the unit hypersphere. The probability density + function of a unit vector :math:`\mathbf{x}` is + + .. math:: + + f(\mathbf{x}) = \frac{\kappa^{d/2-1}}{(2\pi)^{d/2}I_{d/2-1}(\kappa)} + \exp\left(\kappa \mathbf{\mu}^T\mathbf{x}\right), + + where :math:`\mathbf{\mu}` is the mean direction, :math:`\kappa` the + concentration parameter, :math:`d` the dimension and :math:`I` the + modified Bessel function of the first kind. As :math:`\mu` represents + a direction, it must be a unit vector or in other words, a point + on the hypersphere: :math:`\mathbf{\mu}\in S^{d-1}`. :math:`\kappa` is a + concentration parameter, which means that it must be positive + (:math:`\kappa>0`) and that the distribution becomes more narrow with + increasing :math:`\kappa`. In that sense, the reciprocal value + :math:`1/\kappa` resembles the variance parameter of the normal + distribution. + + The von Mises-Fisher distribution often serves as an analogue of the + normal distribution on the sphere. Intuitively, for unit vectors, a + useful distance measure is given by the angle :math:`\alpha` between + them. This is exactly what the scalar product + :math:`\mathbf{\mu}^T\mathbf{x}=\cos(\alpha)` in the + von Mises-Fisher probability density function describes: the angle + between the mean direction :math:`\mathbf{\mu}` and the vector + :math:`\mathbf{x}`. The larger the angle between them, the smaller the + probability to observe :math:`\mathbf{x}` for this particular mean + direction :math:`\mathbf{\mu}`. + + In dimensions 2 and 3, specialized algorithms are used for fast sampling + [2]_, [3]_. For dimensions of 4 or higher the rejection sampling algorithm + described in [4]_ is utilized. This implementation is partially based on + the geomstats package [5]_, [6]_. + + .. versionadded:: 1.11 + + References + ---------- + .. [1] Von Mises-Fisher distribution, Wikipedia, + https://en.wikipedia.org/wiki/Von_Mises%E2%80%93Fisher_distribution + .. [2] Mardia, K., and Jupp, P. Directional statistics. Wiley, 2000. + .. [3] J. Wenzel. Numerically stable sampling of the von Mises Fisher + distribution on S2. + https://www.mitsuba-renderer.org/~wenzel/files/vmf.pdf + .. [4] Wood, A. Simulation of the von mises fisher distribution. + Communications in statistics-simulation and computation 23, + 1 (1994), 157-164. :doi:`10.1080/03610919408813161`. + .. [5] geomstats, Github. MIT License. Accessed: 06.01.2023. + https://github.com/geomstats/geomstats + .. [6] Miolane, N. et al. Geomstats: A Python Package for Riemannian + Geometry in Machine Learning. Journal of Machine Learning Research + 21 (2020). http://jmlr.org/papers/v21/19-027.html + + Examples + -------- + **Visualization of the probability density** + + Plot the probability density in three dimensions for increasing + concentration parameter. The density is calculated by the ``pdf`` + method. + + >>> import numpy as np + >>> import matplotlib.pyplot as plt + >>> from scipy.stats import vonmises_fisher + >>> from matplotlib.colors import Normalize + >>> n_grid = 100 + >>> u = np.linspace(0, np.pi, n_grid) + >>> v = np.linspace(0, 2 * np.pi, n_grid) + >>> u_grid, v_grid = np.meshgrid(u, v) + >>> vertices = np.stack([np.cos(v_grid) * np.sin(u_grid), + ... np.sin(v_grid) * np.sin(u_grid), + ... np.cos(u_grid)], + ... axis=2) + >>> x = np.outer(np.cos(v), np.sin(u)) + >>> y = np.outer(np.sin(v), np.sin(u)) + >>> z = np.outer(np.ones_like(u), np.cos(u)) + >>> def plot_vmf_density(ax, x, y, z, vertices, mu, kappa): + ... vmf = vonmises_fisher(mu, kappa) + ... pdf_values = vmf.pdf(vertices) + ... pdfnorm = Normalize(vmin=pdf_values.min(), vmax=pdf_values.max()) + ... ax.plot_surface(x, y, z, rstride=1, cstride=1, + ... facecolors=plt.cm.viridis(pdfnorm(pdf_values)), + ... linewidth=0) + ... ax.set_aspect('equal') + ... ax.view_init(azim=-130, elev=0) + ... ax.axis('off') + ... ax.set_title(rf"$\kappa={kappa}$") + >>> fig, axes = plt.subplots(nrows=1, ncols=3, figsize=(9, 4), + ... subplot_kw={"projection": "3d"}) + >>> left, middle, right = axes + >>> mu = np.array([-np.sqrt(0.5), -np.sqrt(0.5), 0]) + >>> plot_vmf_density(left, x, y, z, vertices, mu, 5) + >>> plot_vmf_density(middle, x, y, z, vertices, mu, 20) + >>> plot_vmf_density(right, x, y, z, vertices, mu, 100) + >>> plt.subplots_adjust(top=1, bottom=0.0, left=0.0, right=1.0, wspace=0.) + >>> plt.show() + + As we increase the concentration parameter, the points are getting more + clustered together around the mean direction. + + **Sampling** + + Draw 5 samples from the distribution using the ``rvs`` method resulting + in a 5x3 array. + + >>> rng = np.random.default_rng() + >>> mu = np.array([0, 0, 1]) + >>> samples = vonmises_fisher(mu, 20).rvs(5, random_state=rng) + >>> samples + array([[ 0.3884594 , -0.32482588, 0.86231516], + [ 0.00611366, -0.09878289, 0.99509023], + [-0.04154772, -0.01637135, 0.99900239], + [-0.14613735, 0.12553507, 0.98126695], + [-0.04429884, -0.23474054, 0.97104814]]) + + These samples are unit vectors on the sphere :math:`S^2`. To verify, + let us calculate their euclidean norms: + + >>> np.linalg.norm(samples, axis=1) + array([1., 1., 1., 1., 1.]) + + Plot 20 observations drawn from the von Mises-Fisher distribution for + increasing concentration parameter :math:`\kappa`. The red dot highlights + the mean direction :math:`\mu`. + + >>> def plot_vmf_samples(ax, x, y, z, mu, kappa): + ... vmf = vonmises_fisher(mu, kappa) + ... samples = vmf.rvs(20) + ... ax.plot_surface(x, y, z, rstride=1, cstride=1, linewidth=0, + ... alpha=0.2) + ... ax.scatter(samples[:, 0], samples[:, 1], samples[:, 2], c='k', s=5) + ... ax.scatter(mu[0], mu[1], mu[2], c='r', s=30) + ... ax.set_aspect('equal') + ... ax.view_init(azim=-130, elev=0) + ... ax.axis('off') + ... ax.set_title(rf"$\kappa={kappa}$") + >>> mu = np.array([-np.sqrt(0.5), -np.sqrt(0.5), 0]) + >>> fig, axes = plt.subplots(nrows=1, ncols=3, + ... subplot_kw={"projection": "3d"}, + ... figsize=(9, 4)) + >>> left, middle, right = axes + >>> plot_vmf_samples(left, x, y, z, mu, 5) + >>> plot_vmf_samples(middle, x, y, z, mu, 20) + >>> plot_vmf_samples(right, x, y, z, mu, 100) + >>> plt.subplots_adjust(top=1, bottom=0.0, left=0.0, + ... right=1.0, wspace=0.) + >>> plt.show() + + The plots show that with increasing concentration :math:`\kappa` the + resulting samples are centered more closely around the mean direction. + + **Fitting the distribution parameters** + + The distribution can be fitted to data using the ``fit`` method returning + the estimated parameters. As a toy example let's fit the distribution to + samples drawn from a known von Mises-Fisher distribution. + + >>> mu, kappa = np.array([0, 0, 1]), 20 + >>> samples = vonmises_fisher(mu, kappa).rvs(1000, random_state=rng) + >>> mu_fit, kappa_fit = vonmises_fisher.fit(samples) + >>> mu_fit, kappa_fit + (array([0.01126519, 0.01044501, 0.99988199]), 19.306398751730995) + + We see that the estimated parameters `mu_fit` and `kappa_fit` are + very close to the ground truth parameters. + """ + def __init__(self, seed=None): + super().__init__(seed) + + def __call__(self, mu=None, kappa=1, seed=None): + """Create a frozen von Mises-Fisher distribution. + + See `vonmises_fisher_frozen` for more information. + """ + return vonmises_fisher_frozen(mu, kappa, seed=seed) + + def _process_parameters(self, mu, kappa): + """ + Infer dimensionality from mu and ensure that mu is a one-dimensional + unit vector and kappa positive. + """ + mu = np.asarray(mu) + if mu.ndim > 1: + raise ValueError("'mu' must have one-dimensional shape.") + if not np.allclose(np.linalg.norm(mu), 1.): + raise ValueError("'mu' must be a unit vector of norm 1.") + if not mu.size > 1: + raise ValueError("'mu' must have at least two entries.") + kappa_error_msg = "'kappa' must be a positive scalar." + if not np.isscalar(kappa) or kappa < 0: + raise ValueError(kappa_error_msg) + if float(kappa) == 0.: + raise ValueError("For 'kappa=0' the von Mises-Fisher distribution " + "becomes the uniform distribution on the sphere " + "surface. Consider using " + "'scipy.stats.uniform_direction' instead.") + dim = mu.size + + return dim, mu, kappa + + def _check_data_vs_dist(self, x, dim): + if x.shape[-1] != dim: + raise ValueError("The dimensionality of the last axis of 'x' must " + "match the dimensionality of the " + "von Mises Fisher distribution.") + if not np.allclose(np.linalg.norm(x, axis=-1), 1.): + msg = "'x' must be unit vectors of norm 1 along last dimension." + raise ValueError(msg) + + def _log_norm_factor(self, dim, kappa): + # normalization factor is given by + # c = kappa**(dim/2-1)/((2*pi)**(dim/2)*I[dim/2-1](kappa)) + # = kappa**(dim/2-1)*exp(-kappa) / + # ((2*pi)**(dim/2)*I[dim/2-1](kappa)*exp(-kappa) + # = kappa**(dim/2-1)*exp(-kappa) / + # ((2*pi)**(dim/2)*ive[dim/2-1](kappa) + # Then the log is given by + # log c = 1/2*(dim -1)*log(kappa) - kappa - -1/2*dim*ln(2*pi) - + # ive[dim/2-1](kappa) + halfdim = 0.5 * dim + return (0.5 * (dim - 2)*np.log(kappa) - halfdim * _LOG_2PI - + np.log(ive(halfdim - 1, kappa)) - kappa) + + def _logpdf(self, x, dim, mu, kappa): + """Log of the von Mises-Fisher probability density function. + + As this function does no argument checking, it should not be + called directly; use 'logpdf' instead. + + """ + x = np.asarray(x) + self._check_data_vs_dist(x, dim) + dotproducts = np.einsum('i,...i->...', mu, x) + return self._log_norm_factor(dim, kappa) + kappa * dotproducts + + def logpdf(self, x, mu=None, kappa=1): + """Log of the von Mises-Fisher probability density function. + + Parameters + ---------- + x : array_like + Points at which to evaluate the log of the probability + density function. The last axis of `x` must correspond + to unit vectors of the same dimensionality as the distribution. + mu : array_like, default: None + Mean direction of the distribution. Must be a one-dimensional unit + vector of norm 1. + kappa : float, default: 1 + Concentration parameter. Must be positive. + + Returns + ------- + logpdf : ndarray or scalar + Log of the probability density function evaluated at `x`. + + """ + dim, mu, kappa = self._process_parameters(mu, kappa) + return self._logpdf(x, dim, mu, kappa) + + def pdf(self, x, mu=None, kappa=1): + """Von Mises-Fisher probability density function. + + Parameters + ---------- + x : array_like + Points at which to evaluate the probability + density function. The last axis of `x` must correspond + to unit vectors of the same dimensionality as the distribution. + mu : array_like + Mean direction of the distribution. Must be a one-dimensional unit + vector of norm 1. + kappa : float + Concentration parameter. Must be positive. + + Returns + ------- + pdf : ndarray or scalar + Probability density function evaluated at `x`. + + """ + dim, mu, kappa = self._process_parameters(mu, kappa) + return np.exp(self._logpdf(x, dim, mu, kappa)) + + def _rvs_2d(self, mu, kappa, size, random_state): + """ + In 2D, the von Mises-Fisher distribution reduces to the + von Mises distribution which can be efficiently sampled by numpy. + This method is much faster than the general rejection + sampling based algorithm. + + """ + mean_angle = np.arctan2(mu[1], mu[0]) + angle_samples = random_state.vonmises(mean_angle, kappa, size=size) + samples = np.stack([np.cos(angle_samples), np.sin(angle_samples)], + axis=-1) + return samples + + def _rvs_3d(self, kappa, size, random_state): + """ + Generate samples from a von Mises-Fisher distribution + with mu = [1, 0, 0] and kappa. Samples then have to be + rotated towards the desired mean direction mu. + This method is much faster than the general rejection + sampling based algorithm. + Reference: https://www.mitsuba-renderer.org/~wenzel/files/vmf.pdf + + """ + if size is None: + sample_size = 1 + else: + sample_size = size + + # compute x coordinate acc. to equation from section 3.1 + x = random_state.random(sample_size) + x = 1. + np.log(x + (1. - x) * np.exp(-2 * kappa))/kappa + + # (y, z) are random 2D vectors that only have to be + # normalized accordingly. Then (x, y z) follow a VMF distribution + temp = np.sqrt(1. - np.square(x)) + uniformcircle = _sample_uniform_direction(2, sample_size, random_state) + samples = np.stack([x, temp * uniformcircle[..., 0], + temp * uniformcircle[..., 1]], + axis=-1) + if size is None: + samples = np.squeeze(samples) + return samples + + def _rejection_sampling(self, dim, kappa, size, random_state): + """ + Generate samples from an n-dimensional von Mises-Fisher distribution + with mu = [1, 0, ..., 0] and kappa via rejection sampling. + Samples then have to be rotated towards the desired mean direction mu. + Reference: :doi:`10.1080/03610919408813161`. + """ + dim_minus_one = dim - 1 + # calculate number of requested samples + if size is not None: + if not np.iterable(size): + size = (size, ) + n_samples = math.prod(size) + else: + n_samples = 1 + # calculate envelope for rejection sampler (eq. 4) + sqrt = np.sqrt(4 * kappa ** 2. + dim_minus_one ** 2) + envelop_param = (-2 * kappa + sqrt) / dim_minus_one + if envelop_param == 0: + # the regular formula suffers from loss of precision for high + # kappa. This can only be detected by checking for 0 here. + # Workaround: expansion for sqrt variable + # https://www.wolframalpha.com/input?i=sqrt%284*x%5E2%2Bd%5E2%29 + # e = (-2 * k + sqrt(k**2 + d**2)) / d + # ~ (-2 * k + 2 * k + d**2/(4 * k) - d**4/(64 * k**3)) / d + # = d/(4 * k) - d**3/(64 * k**3) + envelop_param = (dim_minus_one/4 * kappa**-1. + - dim_minus_one**3/64 * kappa**-3.) + # reference step 0 + node = (1. - envelop_param) / (1. + envelop_param) + # t = ln(1 - ((1-x)/(1+x))**2) + # = ln(4 * x / (1+x)**2) + # = ln(4) + ln(x) - 2*log1p(x) + correction = (kappa * node + dim_minus_one + * (np.log(4) + np.log(envelop_param) + - 2 * np.log1p(envelop_param))) + n_accepted = 0 + x = np.zeros((n_samples, )) + halfdim = 0.5 * dim_minus_one + # main loop + while n_accepted < n_samples: + # generate candidates acc. to reference step 1 + sym_beta = random_state.beta(halfdim, halfdim, + size=n_samples - n_accepted) + coord_x = (1 - (1 + envelop_param) * sym_beta) / ( + 1 - (1 - envelop_param) * sym_beta) + # accept or reject: reference step 2 + # reformulation for numerical stability: + # t = ln(1 - (1-x)/(1+x) * y) + # = ln((1 + x - y +x*y)/(1 +x)) + accept_tol = random_state.random(n_samples - n_accepted) + criterion = ( + kappa * coord_x + + dim_minus_one * (np.log((1 + envelop_param - coord_x + + coord_x * envelop_param) / (1 + envelop_param))) + - correction) > np.log(accept_tol) + accepted_iter = np.sum(criterion) + x[n_accepted:n_accepted + accepted_iter] = coord_x[criterion] + n_accepted += accepted_iter + # concatenate x and remaining coordinates: step 3 + coord_rest = _sample_uniform_direction(dim_minus_one, n_accepted, + random_state) + coord_rest = np.einsum( + '...,...i->...i', np.sqrt(1 - x ** 2), coord_rest) + samples = np.concatenate([x[..., None], coord_rest], axis=1) + # reshape output to (size, dim) + if size is not None: + samples = samples.reshape(size + (dim, )) + else: + samples = np.squeeze(samples) + return samples + + def _rotate_samples(self, samples, mu, dim): + """A QR decomposition is used to find the rotation that maps the + north pole (1, 0,...,0) to the vector mu. This rotation is then + applied to all samples. + + Parameters + ---------- + samples: array_like, shape = [..., n] + mu : array-like, shape=[n, ] + Point to parametrise the rotation. + + Returns + ------- + samples : rotated samples + + """ + base_point = np.zeros((dim, )) + base_point[0] = 1. + embedded = np.concatenate([mu[None, :], np.zeros((dim - 1, dim))]) + rotmatrix, _ = np.linalg.qr(np.transpose(embedded)) + if np.allclose(np.matmul(rotmatrix, base_point[:, None])[:, 0], mu): + rotsign = 1 + else: + rotsign = -1 + + # apply rotation + samples = np.einsum('ij,...j->...i', rotmatrix, samples) * rotsign + return samples + + def _rvs(self, dim, mu, kappa, size, random_state): + if dim == 2: + samples = self._rvs_2d(mu, kappa, size, random_state) + elif dim == 3: + samples = self._rvs_3d(kappa, size, random_state) + else: + samples = self._rejection_sampling(dim, kappa, size, + random_state) + + if dim != 2: + samples = self._rotate_samples(samples, mu, dim) + return samples + + def rvs(self, mu=None, kappa=1, size=1, random_state=None): + """Draw random samples from a von Mises-Fisher distribution. + + Parameters + ---------- + mu : array_like + Mean direction of the distribution. Must be a one-dimensional unit + vector of norm 1. + kappa : float + Concentration parameter. Must be positive. + size : int or tuple of ints, optional + Given a shape of, for example, (m,n,k), m*n*k samples are + generated, and packed in an m-by-n-by-k arrangement. + Because each sample is N-dimensional, the output shape + is (m,n,k,N). If no shape is specified, a single (N-D) + sample is returned. + random_state : {None, int, np.random.RandomState, np.random.Generator}, + optional + Used for drawing random variates. + If `seed` is `None`, the `~np.random.RandomState` singleton is used. + If `seed` is an int, a new ``RandomState`` instance is used, seeded + with seed. + If `seed` is already a ``RandomState`` or ``Generator`` instance, + then that object is used. + Default is `None`. + + Returns + ------- + rvs : ndarray + Random variates of shape (`size`, `N`), where `N` is the + dimension of the distribution. + + """ + dim, mu, kappa = self._process_parameters(mu, kappa) + random_state = self._get_random_state(random_state) + samples = self._rvs(dim, mu, kappa, size, random_state) + return samples + + def _entropy(self, dim, kappa): + halfdim = 0.5 * dim + return (-self._log_norm_factor(dim, kappa) - kappa * + ive(halfdim, kappa) / ive(halfdim - 1, kappa)) + + def entropy(self, mu=None, kappa=1): + """Compute the differential entropy of the von Mises-Fisher + distribution. + + Parameters + ---------- + mu : array_like, default: None + Mean direction of the distribution. Must be a one-dimensional unit + vector of norm 1. + kappa : float, default: 1 + Concentration parameter. Must be positive. + + Returns + ------- + h : scalar + Entropy of the von Mises-Fisher distribution. + + """ + dim, _, kappa = self._process_parameters(mu, kappa) + return self._entropy(dim, kappa) + + def fit(self, x): + """Fit the von Mises-Fisher distribution to data. + + Parameters + ---------- + x : array-like + Data the distribution is fitted to. Must be two dimensional. + The second axis of `x` must be unit vectors of norm 1 and + determine the dimensionality of the fitted + von Mises-Fisher distribution. + + Returns + ------- + mu : ndarray + Estimated mean direction. + kappa : float + Estimated concentration parameter. + + """ + # validate input data + x = np.asarray(x) + if x.ndim != 2: + raise ValueError("'x' must be two dimensional.") + if not np.allclose(np.linalg.norm(x, axis=-1), 1.): + msg = "'x' must be unit vectors of norm 1 along last dimension." + raise ValueError(msg) + dim = x.shape[-1] + + # mu is simply the directional mean + dirstats = directional_stats(x) + mu = dirstats.mean_direction + r = dirstats.mean_resultant_length + + # kappa is the solution to the equation: + # r = I[dim/2](kappa) / I[dim/2 -1](kappa) + # = I[dim/2](kappa) * exp(-kappa) / I[dim/2 -1](kappa) * exp(-kappa) + # = ive(dim/2, kappa) / ive(dim/2 -1, kappa) + + halfdim = 0.5 * dim + + def solve_for_kappa(kappa): + bessel_vals = ive([halfdim, halfdim - 1], kappa) + return bessel_vals[0]/bessel_vals[1] - r + + root_res = root_scalar(solve_for_kappa, method="brentq", + bracket=(1e-8, 1e9)) + kappa = root_res.root + return mu, kappa + + +vonmises_fisher = vonmises_fisher_gen() + + +class vonmises_fisher_frozen(multi_rv_frozen): + def __init__(self, mu=None, kappa=1, seed=None): + """Create a frozen von Mises-Fisher distribution. + + Parameters + ---------- + mu : array_like, default: None + Mean direction of the distribution. + kappa : float, default: 1 + Concentration parameter. Must be positive. + seed : {None, int, `numpy.random.Generator`, + `numpy.random.RandomState`}, optional + If `seed` is None (or `np.random`), the `numpy.random.RandomState` + singleton is used. + If `seed` is an int, a new ``RandomState`` instance is used, + seeded with `seed`. + If `seed` is already a ``Generator`` or ``RandomState`` instance + then that instance is used. + + """ + self._dist = vonmises_fisher_gen(seed) + self.dim, self.mu, self.kappa = ( + self._dist._process_parameters(mu, kappa) + ) + + def logpdf(self, x): + """ + Parameters + ---------- + x : array_like + Points at which to evaluate the log of the probability + density function. The last axis of `x` must correspond + to unit vectors of the same dimensionality as the distribution. + + Returns + ------- + logpdf : ndarray or scalar + Log of probability density function evaluated at `x`. + + """ + return self._dist._logpdf(x, self.dim, self.mu, self.kappa) + + def pdf(self, x): + """ + Parameters + ---------- + x : array_like + Points at which to evaluate the log of the probability + density function. The last axis of `x` must correspond + to unit vectors of the same dimensionality as the distribution. + + Returns + ------- + pdf : ndarray or scalar + Probability density function evaluated at `x`. + + """ + return np.exp(self.logpdf(x)) + + def rvs(self, size=1, random_state=None): + """Draw random variates from the Von Mises-Fisher distribution. + + Parameters + ---------- + size : int or tuple of ints, optional + Given a shape of, for example, (m,n,k), m*n*k samples are + generated, and packed in an m-by-n-by-k arrangement. + Because each sample is N-dimensional, the output shape + is (m,n,k,N). If no shape is specified, a single (N-D) + sample is returned. + random_state : {None, int, `numpy.random.Generator`, + `numpy.random.RandomState`}, optional + If `seed` is None (or `np.random`), the `numpy.random.RandomState` + singleton is used. + If `seed` is an int, a new ``RandomState`` instance is used, + seeded with `seed`. + If `seed` is already a ``Generator`` or ``RandomState`` instance + then that instance is used. + + Returns + ------- + rvs : ndarray or scalar + Random variates of size (`size`, `N`), where `N` is the + dimension of the distribution. + + """ + random_state = self._dist._get_random_state(random_state) + return self._dist._rvs(self.dim, self.mu, self.kappa, size, + random_state) + + def entropy(self): + """ + Calculate the differential entropy of the von Mises-Fisher + distribution. + + Returns + ------- + h: float + Entropy of the Von Mises-Fisher distribution. + + """ + return self._dist._entropy(self.dim, self.kappa) + + +class normal_inverse_gamma_gen(multi_rv_generic): + r""" + Normal-inverse-gamma distribution. + + The normal-inverse-gamma distribution is the conjugate prior of a normal + distribution with unknown mean and variance. + + Parameters + ---------- + mu, lmbda, a, b : array_like + Shape parameters of the distribution. See notes. + seed : {None, int, np.random.RandomState, np.random.Generator}, optional + Used for drawing random variates. + If `seed` is `None`, the `~np.random.RandomState` singleton is used. + If `seed` is an int, a new ``RandomState`` instance is used, seeded + with seed. + If `seed` is already a ``RandomState`` or ``Generator`` instance, + then that object is used. + Default is `None`. + + Methods + ------- + pdf(x, s2, mu=0, lmbda=1, a=1, b=1) + Probability density function. + logpdf(x, s2, mu=0, lmbda=1, a=1, b=1) + Log of the probability density function. + mean(mu=0, lmbda=1, a=1, b=1) + Distribution mean. + var(mu=0, lmbda=1, a=1, b=1) + Distribution variance. + rvs(mu=0, lmbda=1, a=1, b=1, size=None, random_state=None) + Draw random samples. + + See Also + -------- + norm + invgamma + + Notes + ----- + + The probability density function of `normal_inverse_gamma` is: + + .. math:: + + f(x, \sigma^2; \mu, \lambda, \alpha, \beta) = + \frac{\sqrt{\lambda}}{\sqrt{2 \pi \sigma^2}} + \frac{\beta^\alpha}{\Gamma(\alpha)} + \left( \frac{1}{\sigma^2} \right)^{\alpha + 1} + \exp \left(- \frac{2 \beta + \lambda (x - \mu)^2} {2 \sigma^2} \right) + + where all parameters are real and finite, and :math:`\sigma^2 > 0`, + :math:`\lambda > 0`, :math:`\alpha > 0`, and :math:`\beta > 0`. + + Methods ``normal_inverse_gamma.pdf`` and ``normal_inverse_gamma.logpdf`` + accept `x` and `s2` for arguments :math:`x` and :math:`\sigma^2`. + All methods accept `mu`, `lmbda`, `a`, and `b` for shape parameters + :math:`\mu`, :math:`\lambda`, :math:`\alpha`, and :math:`\beta`, + respectively. + + .. versionadded:: 1.15 + + References + ---------- + .. [1] Normal-inverse-gamma distribution, Wikipedia, + https://en.wikipedia.org/wiki/Normal-inverse-gamma_distribution + + Examples + -------- + Suppose we wish to investigate the relationship between the + normal-inverse-gamma distribution and the inverse gamma distribution. + + >>> import numpy as np + >>> from scipy import stats + >>> import matplotlib.pyplot as plt + >>> rng = np.random.default_rng(527484872345) + >>> mu, lmbda, a, b = 0, 1, 20, 20 + >>> norm_inv_gamma = stats.normal_inverse_gamma(mu, lmbda, a, b) + >>> inv_gamma = stats.invgamma(a, scale=b) + + One approach is to compare the distribution of the `s2` elements of + random variates against the PDF of an inverse gamma distribution. + + >>> _, s2 = norm_inv_gamma.rvs(size=10000, random_state=rng) + >>> bins = np.linspace(s2.min(), s2.max(), 50) + >>> plt.hist(s2, bins=bins, density=True, label='Frequency density') + >>> s2 = np.linspace(s2.min(), s2.max(), 300) + >>> plt.plot(s2, inv_gamma.pdf(s2), label='PDF') + >>> plt.xlabel(r'$\sigma^2$') + >>> plt.ylabel('Frequency density / PMF') + >>> plt.show() + + Similarly, we can compare the marginal distribution of `s2` against + an inverse gamma distribution. + + >>> from scipy.integrate import quad_vec + >>> from scipy import integrate + >>> s2 = np.linspace(0.5, 3, 6) + >>> res = quad_vec(lambda x: norm_inv_gamma.pdf(x, s2), -np.inf, np.inf)[0] + >>> np.allclose(res, inv_gamma.pdf(s2)) + True + + The sample mean is comparable to the mean of the distribution. + + >>> x, s2 = norm_inv_gamma.rvs(size=10000, random_state=rng) + >>> x.mean(), s2.mean() + (np.float64(-0.005254750127304425), np.float64(1.050438111436508)) + >>> norm_inv_gamma.mean() + (np.float64(0.0), np.float64(1.0526315789473684)) + + Similarly, for the variance: + + >>> x.var(ddof=1), s2.var(ddof=1) + (np.float64(1.0546150578185023), np.float64(0.061829865266330754)) + >>> norm_inv_gamma.var() + (np.float64(1.0526315789473684), np.float64(0.061557402277623886)) + """ + def rvs(self, mu=0, lmbda=1, a=1, b=1, size=None, random_state=None): + """Draw random samples from the distribution. + + Parameters + ---------- + mu, lmbda, a, b : array_like, optional + Shape parameters. `lmbda`, `a`, and `b` must be greater + than zero. + size : int or tuple of ints, optional + Shape of samples to draw. + random_state : {None, int, np.random.RandomState, np.random.Generator}, optional + Used for drawing random variates. + If `random_state` is `None`, the `~np.random.RandomState` singleton is used. + If `random_state` is an int, a new ``RandomState`` instance is used, seeded + with `random_state`. + If `random_state` is already a ``RandomState`` or ``Generator`` instance, + then that object is used. + Default is `None`. + + Returns + ------- + x, s2 : ndarray + Random variates. + + """ + random_state = self._get_random_state(random_state) + s2 = invgamma(a, scale=b).rvs(size=size, random_state=random_state) + scale = (s2 / lmbda)**0.5 + x = norm(loc=mu, scale=scale).rvs(size=size, random_state=random_state) + dtype = np.result_type(1.0, mu, lmbda, a, b) + return x.astype(dtype), s2.astype(dtype) + + def _logpdf(self, x, s2, mu, lmbda, a, b): + t1 = 0.5 * (np.log(lmbda) - np.log(2 * np.pi * s2)) + t2 = a*np.log(b) - special.gammaln(a).astype(a.dtype) + t3 = -(a + 1) * np.log(s2) + t4 = -(2*b + lmbda*(x - mu)**2) / (2*s2) + return t1 + t2 + t3 + t4 + + def logpdf(self, x, s2, mu=0, lmbda=1, a=1, b=1): + """Log of the probability density function. + + Parameters + ---------- + x, s2 : array_like + Arguments. `s2` must be greater than zero. + mu, lmbda, a, b : array_like, optional + Shape parameters. `lmbda`, `a`, and `b` must be greater + than zero. + + Returns + ------- + logpdf : ndarray or scalar + Log of the probability density function. + + """ + invalid, args = self._process_parameters_pdf(x, s2, mu, lmbda, a, b) + s2 = args[1] + # Keep it simple for now; lazyselect later, perhaps. + with np.errstate(all='ignore'): + logpdf = np.asarray(self._logpdf(*args)) + logpdf[s2 <= 0] = -np.inf + logpdf[invalid] = np.nan + return logpdf[()] + + def _pdf(self, x, s2, mu, lmbda, a, b): + t1 = np.sqrt(lmbda / (2 * np.pi * s2)) + t2 = b**a / special.gamma(a).astype(a.dtype) + t3 = (1 / s2)**(a + 1) + t4 = np.exp(-(2*b + lmbda*(x - mu)**2) / (2*s2)) + return t1 * t2 * t3 * t4 + + def pdf(self, x, s2, mu=0, lmbda=1, a=1, b=1): + """The probability density function. + + Parameters + ---------- + x, s2 : array_like + Arguments. `s2` must be greater than zero. + mu, lmbda, a, b : array_like, optional + Shape parameters. `lmbda`, `a`, and `b` must be greater + than zero. + + Returns + ------- + logpdf : ndarray or scalar + The probability density function. + + """ + invalid, args = self._process_parameters_pdf(x, s2, mu, lmbda, a, b) + s2 = args[1] + # Keep it simple for now; lazyselect later, perhaps. + with np.errstate(all='ignore'): + pdf = np.asarray(self._pdf(*args)) + pdf[s2 <= 0] = 0 + pdf[invalid] = np.nan + return pdf[()] + + def mean(self, mu=0, lmbda=1, a=1, b=1): + """The mean of the distribution. + + Parameters + ---------- + mu, lmbda, a, b : array_like, optional + Shape parameters. `lmbda` and `b` must be greater + than zero, and `a` must be greater than one. + + Returns + ------- + x, s2 : ndarray + The mean of the distribution. + + """ + invalid, args = self._process_shapes(mu, lmbda, a, b) + mu, lmbda, a, b = args + invalid |= ~(a > 1) + mean_x = np.asarray(mu).copy() + mean_s2 = np.asarray(b / (a - 1)) + mean_x[invalid] = np.nan + mean_s2[invalid] = np.nan + return mean_x[()], mean_s2[()] + + def var(self, mu=0, lmbda=1, a=1, b=1): + """The variance of the distribution. + + Parameters + ---------- + mu, lmbda, a, b : array_like, optional + Shape parameters. `lmbda` and `b` must be greater + than zero, and `a` must be greater than two. + + Returns + ------- + x, s2 : ndarray + The variance of the distribution. + + """ + invalid, args = self._process_shapes(mu, lmbda, a, b) + mu, lmbda, a, b = args + invalid_x = invalid | ~(a > 1) + invalid_s2 = invalid | ~(a > 2) + var_x = b / ((a - 1) * lmbda) + var_s2 = b**2 / ((a - 1)**2 * (a - 2)) + var_x, var_s2 = np.asarray(var_x), np.asarray(var_s2) + var_x[invalid_x] = np.nan + var_s2[invalid_s2] = np.nan + return var_x[()], var_s2[()] + + def _process_parameters_pdf(self, x, s2, mu, lmbda, a, b): + args = np.broadcast_arrays(x, s2, mu, lmbda, a, b) + dtype = np.result_type(1.0, *(arg.dtype for arg in args)) + args = [arg.astype(dtype, copy=False) for arg in args] + x, s2, mu, lmbda, a, b = args + invalid = ~((lmbda > 0) & (a > 0) & (b > 0)) + return invalid, args + + def _process_shapes(self, mu, lmbda, a, b): + args = np.broadcast_arrays(mu, lmbda, a, b) + dtype = np.result_type(1.0, *(arg.dtype for arg in args)) + args = [arg.astype(dtype, copy=False) for arg in args] + mu, lmbda, a, b = args + invalid = ~((lmbda > 0) & (a > 0) & (b > 0)) + return invalid, args + + def __call__(self, mu=0, lmbda=1, a=1, b=1, seed=None): + return normal_inverse_gamma_frozen(mu, lmbda, a, b, seed=seed) + + +normal_inverse_gamma = normal_inverse_gamma_gen() + + +class normal_inverse_gamma_frozen(multi_rv_frozen): + + def __init__(self, mu=0, lmbda=1, a=1, b=1, seed=None): + self._dist = normal_inverse_gamma_gen(seed) + self._shapes = mu, lmbda, a, b + + def logpdf(self, x, s2): + return self._dist.logpdf(x, s2, *self._shapes) + + def pdf(self, x, s2): + return self._dist.pdf(x, s2, *self._shapes) + + def mean(self): + return self._dist.mean(*self._shapes) + + def var(self): + return self._dist.var(*self._shapes) + + def rvs(self, size=None, random_state=None): + return self._dist.rvs(*self._shapes, size=size, random_state=random_state) + + +# Set frozen generator docstrings from corresponding docstrings in +# normal_inverse_gamma_gen and fill in default strings in class docstrings +for name in ['logpdf', 'pdf', 'mean', 'var', 'rvs']: + method = normal_inverse_gamma_gen.__dict__[name] + method_frozen = normal_inverse_gamma_frozen.__dict__[name] + method_frozen.__doc__ = doccer.docformat(method.__doc__, + mvn_docdict_noparams) + method.__doc__ = doccer.docformat(method.__doc__, mvn_docdict_params) diff --git a/v12x/pylib/scipy/stats/_new_distributions.py b/v12x/pylib/scipy/stats/_new_distributions.py new file mode 100644 index 0000000000000000000000000000000000000000..8393e7718b7300cc4aed4f620f31c8d44c0848ce --- /dev/null +++ b/v12x/pylib/scipy/stats/_new_distributions.py @@ -0,0 +1,551 @@ +import sys + +import numpy as np +from numpy import inf + +from scipy._external import array_api_extra as xpx +from scipy import special +from scipy.special import _ufuncs as scu +from scipy.stats._distribution_infrastructure import ( + ContinuousDistribution, DiscreteDistribution, _RealInterval, _IntegerInterval, + _RealParameter, _Parameterization, _combine_docs) + +__all__ = ['Normal', 'Logistic', 'Uniform', 'Binomial'] + + +class Normal(ContinuousDistribution): + r"""Normal distribution with prescribed mean and standard deviation. + + The probability density function of the normal distribution is: + + .. math:: + + f(x) = \frac{1}{\sigma \sqrt{2 \pi}} \exp { + \left( -\frac{1}{2}\left( \frac{x - \mu}{\sigma} \right)^2 \right)} + + """ + # `ShiftedScaledDistribution` allows this to be generated automatically from + # an instance of `StandardNormal`, but the normal distribution is so frequently + # used that it's worth a bit of code duplication to get better performance. + _mu_domain = _RealInterval(endpoints=(-inf, inf)) + _sigma_domain = _RealInterval(endpoints=(0, inf)) + _x_support = _RealInterval(endpoints=(-inf, inf)) + + _mu_param = _RealParameter('mu', symbol=r'\mu', domain=_mu_domain, + typical=(-1, 1)) + _sigma_param = _RealParameter('sigma', symbol=r'\sigma', domain=_sigma_domain, + typical=(0.5, 1.5)) + _x_param = _RealParameter('x', domain=_x_support, typical=(-1, 1)) + + _parameterizations = [_Parameterization(_mu_param, _sigma_param)] + + _variable = _x_param + _normalization = 1/np.sqrt(2*np.pi) + _log_normalization = np.log(2*np.pi)/2 + + def __new__(cls, mu=None, sigma=None, **kwargs): + if mu is None and sigma is None: + return super().__new__(StandardNormal) + return super().__new__(cls) + + def __init__(self, *, mu=0., sigma=1., **kwargs): + super().__init__(mu=mu, sigma=sigma, **kwargs) + + def _logpdf_formula(self, x, *, mu, sigma, **kwargs): + return StandardNormal._logpdf_formula(self, (x - mu)/sigma) - np.log(sigma) + + def _pdf_formula(self, x, *, mu, sigma, **kwargs): + return StandardNormal._pdf_formula(self, (x - mu)/sigma) / sigma + + def _logcdf_formula(self, x, *, mu, sigma, **kwargs): + return StandardNormal._logcdf_formula(self, (x - mu)/sigma) + + def _cdf_formula(self, x, *, mu, sigma, **kwargs): + return StandardNormal._cdf_formula(self, (x - mu)/sigma) + + def _logccdf_formula(self, x, *, mu, sigma, **kwargs): + return StandardNormal._logccdf_formula(self, (x - mu)/sigma) + + def _ccdf_formula(self, x, *, mu, sigma, **kwargs): + return StandardNormal._ccdf_formula(self, (x - mu)/sigma) + + def _icdf_formula(self, x, *, mu, sigma, **kwargs): + return StandardNormal._icdf_formula(self, x) * sigma + mu + + def _ilogcdf_formula(self, x, *, mu, sigma, **kwargs): + return StandardNormal._ilogcdf_formula(self, x) * sigma + mu + + def _iccdf_formula(self, x, *, mu, sigma, **kwargs): + return StandardNormal._iccdf_formula(self, x) * sigma + mu + + def _ilogccdf_formula(self, x, *, mu, sigma, **kwargs): + return StandardNormal._ilogccdf_formula(self, x) * sigma + mu + + def _entropy_formula(self, *, mu, sigma, **kwargs): + return StandardNormal._entropy_formula(self) + np.log(abs(sigma)) + + def _logentropy_formula(self, *, mu, sigma, **kwargs): + lH0 = StandardNormal._logentropy_formula(self) + with np.errstate(divide='ignore'): + # sigma = 1 -> log(sigma) = 0 -> log(log(sigma)) = -inf + # Silence the unnecessary runtime warning + lls = np.log(np.log(abs(sigma))+0j) + return special.logsumexp(np.broadcast_arrays(lH0, lls), axis=0) + + def _median_formula(self, *, mu, sigma, **kwargs): + return mu + + def _mode_formula(self, *, mu, sigma, **kwargs): + return mu + + def _moment_raw_formula(self, order, *, mu, sigma, **kwargs): + if order == 0: + return np.ones_like(mu) + elif order == 1: + return mu + else: + return None + _moment_raw_formula.orders = [0, 1] # type: ignore[attr-defined] + + def _moment_central_formula(self, order, *, mu, sigma, **kwargs): + if order == 0: + return np.ones_like(mu) + elif order % 2: + return np.zeros_like(mu) + else: + # exact is faster (and obviously more accurate) for reasonable orders + return sigma**order * special.factorial2(int(order) - 1, exact=True) + + def _lmoment_formula(self, order, *, mu, sigma, **kwargs): + lscale = sigma / np.sqrt(np.pi) + lkurtosis = 30*np.arctan(np.sqrt(2))/np.pi - 9 + lmoments = {1: mu, 2: lscale, 3: 0, 4: lkurtosis * lscale} + return lmoments.get(order, None) + + def _sample_formula(self, full_shape, rng, *, mu, sigma, **kwargs): + return rng.normal(loc=mu, scale=sigma, size=full_shape)[()] + + +def _log_diff(log_p, log_q): + return special.logsumexp([log_p, log_q+np.pi*1j], axis=0) + + +class StandardNormal(Normal): + r"""Standard normal distribution. + + The probability density function of the standard normal distribution is: + + .. math:: + + f(x) = \frac{1}{\sqrt{2 \pi}} \exp \left( -\frac{1}{2} x^2 \right) + + """ + _x_support = _RealInterval(endpoints=(-inf, inf)) + _x_param = _RealParameter('x', domain=_x_support, typical=(-5, 5)) + _variable = _x_param + _parameterizations = [] + _normalization = 1/np.sqrt(2*np.pi) + _log_normalization = np.log(2*np.pi)/2 + mu = np.float64(0.) + sigma = np.float64(1.) + + def __init__(self, **kwargs): + ContinuousDistribution.__init__(self, **kwargs) + + def _logpdf_formula(self, x, **kwargs): + return -(self._log_normalization + x**2/2) + + def _pdf_formula(self, x, **kwargs): + return self._normalization * np.exp(-x**2/2) + + def _logcdf_formula(self, x, **kwargs): + return special.log_ndtr(x) + + def _cdf_formula(self, x, **kwargs): + return special.ndtr(x) + + def _logccdf_formula(self, x, **kwargs): + return special.log_ndtr(-x) + + def _ccdf_formula(self, x, **kwargs): + return special.ndtr(-x) + + def _icdf_formula(self, x, **kwargs): + return special.ndtri(x) + + def _ilogcdf_formula(self, x, **kwargs): + return special.ndtri_exp(x) + + def _iccdf_formula(self, x, **kwargs): + return -special.ndtri(x) + + def _ilogccdf_formula(self, x, **kwargs): + return -special.ndtri_exp(x) + + def _entropy_formula(self, **kwargs): + return (1 + np.log(2*np.pi))/2 + + def _logentropy_formula(self, **kwargs): + return np.log1p(np.log(2*np.pi)) - np.log(2) + + def _median_formula(self, **kwargs): + return 0 + + def _mode_formula(self, **kwargs): + return 0 + + def _moment_raw_formula(self, order, **kwargs): + raw_moments = {0: 1, 1: 0, 2: 1, 3: 0, 4: 3, 5: 0} + return raw_moments.get(order, None) + + def _moment_central_formula(self, order, **kwargs): + return self._moment_raw_formula(order, **kwargs) + + def _moment_standardized_formula(self, order, **kwargs): + return self._moment_raw_formula(order, **kwargs) + + def _lmoment_formula(self, order, **kwargs): + return super()._lmoment_formula(order, mu=0., sigma=1., **kwargs) + + def _sample_formula(self, full_shape, rng, **kwargs): + return rng.normal(size=full_shape)[()] + + +class Logistic(ContinuousDistribution): + r"""Standard logistic distribution. + + The probability density function of the standard logistic distribution is: + + .. math:: + + f(x) = \frac{1}{\left( e^{x / 2} + e^{-x / 2} \right)^2} + + """ + _x_support = _RealInterval(endpoints=(-inf, inf)) + _variable = _x_param = _RealParameter('x', domain=_x_support, typical=(-9, 9)) + _parameterizations = () # type:ignore[assignment] + + _scale = np.pi / np.sqrt(3) + + def _logpdf_formula(self, x, **kwargs): + y = -np.abs(x) + return y - 2 * special.log1p(np.exp(y)) + + def _pdf_formula(self, x, **kwargs): + # f(x) = sech(x / 2)**2 / 4 + return (.5 / np.cosh(x / 2))**2 + + def _logcdf_formula(self, x, **kwargs): + return special.log_expit(x) + + def _cdf_formula(self, x, **kwargs): + return special.expit(x) + + def _logccdf_formula(self, x, **kwargs): + return special.log_expit(-x) + + def _ccdf_formula(self, x, **kwargs): + return special.expit(-x) + + def _icdf_formula(self, x, **kwargs): + return special.logit(x) + + def _iccdf_formula(self, x, **kwargs): + return -special.logit(x) + + def _entropy_formula(self, **kwargs): + return 2.0 + + def _logentropy_formula(self, **kwargs): + return np.log(2) + + def _median_formula(self, **kwargs): + return 0 + + def _mode_formula(self, **kwargs): + return 0 + + def _moment_raw_formula(self, order, **kwargs): + n = int(order) + if n % 2: + return 0.0 + return np.pi**n * abs((2**n - 2) * float(special.bernoulli(n)[-1])) + + def _moment_central_formula(self, order, **kwargs): + return self._moment_raw_formula(order, **kwargs) + + def _moment_standardized_formula(self, order, **kwargs): + return self._moment_raw_formula(order, **kwargs) / self._scale**order + + def _lmoment_formula(self, order, **kwargs): + lmoments = {1: 0, 2: 1, 3: 0, 4: 1/6} + return lmoments.get(order, None) + + def _sample_formula(self, full_shape, rng, **kwargs): + return rng.logistic(size=full_shape)[()] + + +# currently for testing only +class _LogUniform(ContinuousDistribution): + r"""Log-uniform distribution. + + The probability density function of the log-uniform distribution is: + + .. math:: + + f(x; a, b) = \frac{1} + {x (\log(b) - \log(a))} + + If :math:`\log(X)` is a random variable that follows a uniform distribution + between :math:`\log(a)` and :math:`\log(b)`, then :math:`X` is log-uniformly + distributed with shape parameters :math:`a` and :math:`b`. + + """ + + _a_domain = _RealInterval(endpoints=(0, inf)) + _b_domain = _RealInterval(endpoints=('a', inf)) + _log_a_domain = _RealInterval(endpoints=(-inf, inf)) + _log_b_domain = _RealInterval(endpoints=('log_a', inf)) + _x_support = _RealInterval(endpoints=('a', 'b'), inclusive=(True, True)) + + _a_param = _RealParameter('a', domain=_a_domain, typical=(1e-3, 0.9)) + _b_param = _RealParameter('b', domain=_b_domain, typical=(1.1, 1e3)) + _log_a_param = _RealParameter('log_a', symbol=r'\log(a)', + domain=_log_a_domain, typical=(-3, -0.1)) + _log_b_param = _RealParameter('log_b', symbol=r'\log(b)', + domain=_log_b_domain, typical=(0.1, 3)) + _x_param = _RealParameter('x', domain=_x_support, typical=('a', 'b')) + + _b_domain.define_parameters(_a_param) + _log_b_domain.define_parameters(_log_a_param) + _x_support.define_parameters(_a_param, _b_param) + + _parameterizations = [_Parameterization(_log_a_param, _log_b_param), + _Parameterization(_a_param, _b_param)] + _variable = _x_param + + def __init__(self, *, a=None, b=None, log_a=None, log_b=None, **kwargs): + super().__init__(a=a, b=b, log_a=log_a, log_b=log_b, **kwargs) + + def _process_parameters(self, a=None, b=None, log_a=None, log_b=None, **kwargs): + a = np.exp(log_a) if a is None else a + b = np.exp(log_b) if b is None else b + log_a = np.log(a) if log_a is None else log_a + log_b = np.log(b) if log_b is None else log_b + kwargs.update(dict(a=a, b=b, log_a=log_a, log_b=log_b)) + return kwargs + + def _pdf_formula(self, x, *, log_a, log_b, **kwargs): + return ((log_b - log_a)*x)**-1 + + def _moment_raw_formula(self, order, log_a, log_b, **kwargs): + if order == 0: + return self._one + t1 = self._one / (log_b - log_a) / order + t2 = np.real(np.exp(_log_diff(order * log_b, order * log_a))) + return t1 * t2 + + +class Uniform(ContinuousDistribution): + r"""Uniform distribution. + + The probability density function of the uniform distribution is: + + .. math:: + + f(x; a, b) = \frac{1} + {b - a} + + """ + + _a_domain = _RealInterval(endpoints=(-inf, inf)) + _b_domain = _RealInterval(endpoints=('a', inf)) + _x_support = _RealInterval(endpoints=('a', 'b'), inclusive=(True, True)) + + _a_param = _RealParameter('a', domain=_a_domain, typical=(1e-3, 0.9)) + _b_param = _RealParameter('b', domain=_b_domain, typical=(1.1, 1e3)) + _x_param = _RealParameter('x', domain=_x_support, typical=('a', 'b')) + + _b_domain.define_parameters(_a_param) + _x_support.define_parameters(_a_param, _b_param) + + _parameterizations = [_Parameterization(_a_param, _b_param)] + _variable = _x_param + + def __init__(self, *, a=None, b=None, **kwargs): + super().__init__(a=a, b=b, **kwargs) + + def _process_parameters(self, a=None, b=None, ab=None, **kwargs): + ab = b - a + kwargs.update(dict(a=a, b=b, ab=ab)) + return kwargs + + def _logpdf_formula(self, x, *, ab, **kwargs): + return np.where(np.isnan(x), np.nan, -np.log(ab)) + + def _pdf_formula(self, x, *, ab, **kwargs): + return np.where(np.isnan(x), np.nan, 1/ab) + + def _logcdf_formula(self, x, *, a, ab, **kwargs): + with np.errstate(divide='ignore'): + return np.log(x - a) - np.log(ab) + + def _cdf_formula(self, x, *, a, ab, **kwargs): + return (x - a) / ab + + def _logccdf_formula(self, x, *, b, ab, **kwargs): + with np.errstate(divide='ignore'): + return np.log(b - x) - np.log(ab) + + def _ccdf_formula(self, x, *, b, ab, **kwargs): + return (b - x) / ab + + def _icdf_formula(self, p, *, a, ab, **kwargs): + return a + ab*p + + def _iccdf_formula(self, p, *, b, ab, **kwargs): + return b - ab*p + + def _entropy_formula(self, *, ab, **kwargs): + return np.log(ab) + + def _mode_formula(self, *, a, b, ab, **kwargs): + return a + 0.5*ab + + def _median_formula(self, *, a, b, ab, **kwargs): + return a + 0.5*ab + + def _moment_raw_formula(self, order, a, b, ab, **kwargs): + np1 = order + 1 + return (b**np1 - a**np1) / (np1 * ab) + + def _moment_central_formula(self, order, ab, **kwargs): + return ab**2/12 if order == 2 else None + + _moment_central_formula.orders = [2] # type: ignore[attr-defined] + + def _lmoment_formula(self, order, *, a, b, ab, **kwargs): + lmoments = {1: 0.5*(a + b), 2: ab / 6} + return lmoments.get(order, np.zeros_like(ab)) + + def _sample_formula(self, full_shape, rng, a, b, ab, **kwargs): + try: + return rng.uniform(a, b, size=full_shape)[()] + except OverflowError: # happens when there are NaNs + return rng.uniform(0, 1, size=full_shape)*ab + a + + +class _Gamma(ContinuousDistribution): + # Gamma distribution for testing only + _a_domain = _RealInterval(endpoints=(0, inf)) + _x_support = _RealInterval(endpoints=(0, inf), inclusive=(False, False)) + + _a_param = _RealParameter('a', domain=_a_domain, typical=(0.1, 10)) + _x_param = _RealParameter('x', domain=_x_support, typical=(0.1, 10)) + + _parameterizations = [_Parameterization(_a_param)] + _variable = _x_param + + def _pdf_formula(self, x, *, a, **kwargs): + return x ** (a - 1) * np.exp(-x) / special.gamma(a) + + +class Binomial(DiscreteDistribution): + r"""Binomial distribution with prescribed success probability and number of trials + + The probability density function of the binomial distribution is: + + .. math:: + + f(x) = {n \choose x} p^x (1 - p)^{n-x} + + """ + _n_domain = _IntegerInterval(endpoints=(0, inf), inclusive=(False, False)) + _p_domain = _RealInterval(endpoints=(0, 1), inclusive=(False, False)) + _x_support = _IntegerInterval(endpoints=(0, 'n'), inclusive=(True, True)) + + _n_param = _RealParameter('n', domain=_n_domain, typical=(10, 20)) + _p_param = _RealParameter('p', domain=_p_domain, typical=(0.25, 0.75)) + _x_param = _RealParameter('x', domain=_x_support, typical=(0, 10)) + + _parameterizations = [_Parameterization(_n_param, _p_param)] + _variable = _x_param + + def __init__(self, *, n, p, **kwargs): + super().__init__(n=n, p=p, **kwargs) + + def _pmf_formula(self, x, *, n, p, **kwargs): + return scu._binom_pmf(x, n, p) + + def _logpmf_formula(self, x, *, n, p, **kwargs): + # This implementation is from the ``scipy.stats.binom`` and could be improved + # by using a more numerically sound implementation of the absolute value of + # the binomial coefficient. + combiln = ( + special.gammaln(n+1) - (special.gammaln(x+1) + special.gammaln(n-x+1)) + ) + return combiln + special.xlogy(x, p) + special.xlog1py(n-x, -p) + + def _cdf_formula(self, x, *, n, p, **kwargs): + return scu._binom_cdf(x, n, p) + + def _logcdf_formula(self, x, *, n, p, **kwargs): + # todo: add this strategy to infrastructure more generally, but allow dist + # author to specify threshold other than median in case median is expensive + median = self._icdf_formula(0.5, n=n, p=p) + return xpx.apply_where(x < median, (x, n, p), + lambda *args: np.log(scu._binom_cdf(*args)), + lambda *args: np.log1p(-scu._binom_sf(*args)) + ) + + def _ccdf_formula(self, x, *, n, p, **kwargs): + return scu._binom_sf(x, n, p) + + def _logccdf_formula(self, x, *, n, p, **kwargs): + median = self._icdf_formula(0.5, n=n, p=p) + return xpx.apply_where(x < median, (x, n, p), + lambda *args: np.log1p(-scu._binom_cdf(*args)), + lambda *args: np.log(scu._binom_sf(*args)) + ) + + def _icdf_formula(self, x, *, n, p, **kwargs): + return scu._binom_ppf(x, n, p) + + def _iccdf_formula(self, x, *, n, p, **kwargs): + return scu._binom_isf(x, n, p) + + def _mode_formula(self, *, n, p, **kwargs): + # https://en.wikipedia.org/wiki/Binomial_distribution#Mode + mode = np.floor((n+1)*p) + mode = np.where(p == 1, mode - 1, mode) + return mode[()] + + def _moment_raw_formula(self, order, *, n, p, **kwargs): + # https://en.wikipedia.org/wiki/Binomial_distribution#Higher_moments + if order == 1: + return n*p + if order == 2: + return n*p*(1 - p + n*p) + return None + _moment_raw_formula.orders = [1, 2] # type: ignore[attr-defined] + + def _moment_central_formula(self, order, *, n, p, **kwargs): + # https://en.wikipedia.org/wiki/Binomial_distribution#Higher_moments + if order == 1: + return np.zeros_like(n) + if order == 2: + return n*p*(1 - p) + if order == 3: + return n*p*(1 - p)*(1 - 2*p) + if order == 4: + return n*p*(1 - p)*(1 + (3*n - 6)*p*(1 - p)) + return None + _moment_central_formula.orders = [1, 2, 3, 4] # type: ignore[attr-defined] + + +# Distribution classes need only define the summary and beginning of the extended +# summary portion of the class documentation. All other documentation, including +# examples, is generated automatically. +_module = sys.modules[__name__].__dict__ +for dist_name in __all__: + _module[dist_name].__doc__ = _combine_docs(_module[dist_name]) diff --git a/v12x/pylib/scipy/stats/_odds_ratio.py b/v12x/pylib/scipy/stats/_odds_ratio.py new file mode 100644 index 0000000000000000000000000000000000000000..bc593f5adc9a700c618c721b6c37c801809d868b --- /dev/null +++ b/v12x/pylib/scipy/stats/_odds_ratio.py @@ -0,0 +1,466 @@ +import numpy as np + +from scipy.special import ndtri +from scipy.optimize import brentq +from ._discrete_distns import nchypergeom_fisher +from ._common import ConfidenceInterval + + +def _sample_odds_ratio(table): + """ + Given a table [[a, b], [c, d]], compute a*d/(b*c). + + Return nan if the numerator and denominator are 0. + Return inf if just the denominator is 0. + """ + # table must be a 2x2 numpy array. + if table[1, 0] > 0 and table[0, 1] > 0: + oddsratio = table[0, 0] * table[1, 1] / (table[1, 0] * table[0, 1]) + elif table[0, 0] == 0 or table[1, 1] == 0: + oddsratio = np.nan + else: + oddsratio = np.inf + return oddsratio + + +def _solve(func): + """ + Solve func(nc) = 0. func must be an increasing function. + """ + # We could just as well call the variable `x` instead of `nc`, but we + # always call this function with functions for which nc (the noncentrality + # parameter) is the variable for which we are solving. + nc = 1.0 + value = func(nc) + if value == 0: + return nc + + # Multiplicative factor by which to increase or decrease nc when + # searching for a bracketing interval. + factor = 2.0 + # Find a bracketing interval. + if value > 0: + nc /= factor + while func(nc) > 0: + nc /= factor + lo = nc + hi = factor*nc + else: + nc *= factor + while func(nc) < 0: + nc *= factor + lo = nc/factor + hi = nc + + # lo and hi bracket the solution for nc. + nc = brentq(func, lo, hi, xtol=1e-13) + return nc + + +def _nc_hypergeom_mean_inverse(x, M, n, N): + """ + For the given noncentral hypergeometric parameters x, M, n,and N + (table[0,0], total, row 0 sum and column 0 sum, resp., of a 2x2 + contingency table), find the noncentrality parameter of Fisher's + noncentral hypergeometric distribution whose mean is x. + """ + nc = _solve(lambda nc: nchypergeom_fisher.mean(M, n, N, nc) - x) + return nc + + +def _hypergeom_params_from_table(table): + # The notation M, n and N is consistent with stats.hypergeom and + # stats.nchypergeom_fisher. + x = table[0, 0] + M = table.sum() + n = table[0].sum() + N = table[:, 0].sum() + return x, M, n, N + + +def _ci_upper(table, alpha): + """ + Compute the upper end of the confidence interval. + """ + if _sample_odds_ratio(table) == np.inf: + return np.inf + + x, M, n, N = _hypergeom_params_from_table(table) + + # nchypergeom_fisher.cdf is a decreasing function of nc, so we negate + # it in the lambda expression. + nc = _solve(lambda nc: -nchypergeom_fisher.cdf(x, M, n, N, nc) + alpha) + return nc + + +def _ci_lower(table, alpha): + """ + Compute the lower end of the confidence interval. + """ + if _sample_odds_ratio(table) == 0: + return 0 + + x, M, n, N = _hypergeom_params_from_table(table) + + nc = _solve(lambda nc: nchypergeom_fisher.sf(x - 1, M, n, N, nc) - alpha) + return nc + + +def _conditional_oddsratio(table): + """ + Conditional MLE of the odds ratio for the 2x2 contingency table. + """ + x, M, n, N = _hypergeom_params_from_table(table) + # Get the bounds of the support. The support of the noncentral + # hypergeometric distribution with parameters M, n, and N is the same + # for all values of the noncentrality parameter, so we can use 1 here. + lo, hi = nchypergeom_fisher.support(M, n, N, 1) + + # Check if x is at one of the extremes of the support. If so, we know + # the odds ratio is either 0 or inf. + if x == lo: + # x is at the low end of the support. + return 0 + if x == hi: + # x is at the high end of the support. + return np.inf + + nc = _nc_hypergeom_mean_inverse(x, M, n, N) + return nc + + +def _conditional_oddsratio_ci(table, confidence_level=0.95, + alternative='two-sided'): + """ + Conditional exact confidence interval for the odds ratio. + """ + if alternative == 'two-sided': + alpha = 0.5*(1 - confidence_level) + lower = _ci_lower(table, alpha) + upper = _ci_upper(table, alpha) + elif alternative == 'less': + lower = 0.0 + upper = _ci_upper(table, 1 - confidence_level) + else: + # alternative == 'greater' + lower = _ci_lower(table, 1 - confidence_level) + upper = np.inf + + return lower, upper + + +def _sample_odds_ratio_ci(table, confidence_level=0.95, + alternative='two-sided'): + oddsratio = _sample_odds_ratio(table) + log_or = np.log(oddsratio) + se = np.sqrt((1/table).sum()) + if alternative == 'less': + z = ndtri(confidence_level) + loglow = -np.inf + loghigh = log_or + z*se + elif alternative == 'greater': + z = ndtri(confidence_level) + loglow = log_or - z*se + loghigh = np.inf + else: + # alternative is 'two-sided' + z = ndtri(0.5*confidence_level + 0.5) + loglow = log_or - z*se + loghigh = log_or + z*se + + return np.exp(loglow), np.exp(loghigh) + + +class OddsRatioResult: + """ + Result of `scipy.stats.contingency.odds_ratio`. See the + docstring for `odds_ratio` for more details. + + Attributes + ---------- + statistic : float + The computed odds ratio. + + * If `kind` is ``'sample'``, this is sample (or unconditional) + estimate, given by + ``table[0, 0]*table[1, 1]/(table[0, 1]*table[1, 0])``. + * If `kind` is ``'conditional'``, this is the conditional + maximum likelihood estimate for the odds ratio. It is + the noncentrality parameter of Fisher's noncentral + hypergeometric distribution with the same hypergeometric + parameters as `table` and whose mean is ``table[0, 0]``. + + Methods + ------- + confidence_interval : + Confidence interval for the odds ratio. + """ + + def __init__(self, _table, _kind, statistic): + # for now, no need to make _table and _kind public, since this sort of + # information is returned in very few `scipy.stats` results + self._table = _table + self._kind = _kind + self.statistic = statistic + + def __repr__(self): + return f"OddsRatioResult(statistic={self.statistic})" + + def confidence_interval(self, confidence_level=0.95, + alternative='two-sided'): + """ + Confidence interval for the odds ratio. + + Parameters + ---------- + confidence_level: float + Desired confidence level for the confidence interval. + The value must be given as a fraction between 0 and 1. + Default is 0.95 (meaning 95%). + + alternative : {'two-sided', 'less', 'greater'}, optional + The alternative hypothesis of the hypothesis test to which the + confidence interval corresponds. That is, suppose the null + hypothesis is that the true odds ratio equals ``OR`` and the + confidence interval is ``(low, high)``. Then the following options + for `alternative` are available (default is 'two-sided'): + + * 'two-sided': the true odds ratio is not equal to ``OR``. There + is evidence against the null hypothesis at the chosen + `confidence_level` if ``high < OR`` or ``low > OR``. + * 'less': the true odds ratio is less than ``OR``. The ``low`` end + of the confidence interval is 0, and there is evidence against + the null hypothesis at the chosen `confidence_level` if + ``high < OR``. + * 'greater': the true odds ratio is greater than ``OR``. The + ``high`` end of the confidence interval is ``np.inf``, and there + is evidence against the null hypothesis at the chosen + `confidence_level` if ``low > OR``. + + Returns + ------- + ci : ``ConfidenceInterval`` instance + The confidence interval, represented as an object with + attributes ``low`` and ``high``. + + Notes + ----- + When `kind` is ``'conditional'``, the limits of the confidence + interval are the conditional "exact confidence limits" as described + by Fisher [1]_. The conditional odds ratio and confidence interval are + also discussed in Section 4.1.2 of the text by Sahai and Khurshid [2]_. + + When `kind` is ``'sample'``, the confidence interval is computed + under the assumption that the logarithm of the odds ratio is normally + distributed with standard error given by:: + + se = sqrt(1/a + 1/b + 1/c + 1/d) + + where ``a``, ``b``, ``c`` and ``d`` are the elements of the + contingency table. (See, for example, [2]_, section 3.1.3.2, + or [3]_, section 2.3.3). + + References + ---------- + .. [1] R. A. Fisher (1935), The logic of inductive inference, + Journal of the Royal Statistical Society, Vol. 98, No. 1, + pp. 39-82. + .. [2] H. Sahai and A. Khurshid (1996), Statistics in Epidemiology: + Methods, Techniques, and Applications, CRC Press LLC, Boca + Raton, Florida. + .. [3] Alan Agresti, An Introduction to Categorical Data Analysis + (second edition), Wiley, Hoboken, NJ, USA (2007). + """ + if alternative not in ['two-sided', 'less', 'greater']: + raise ValueError("`alternative` must be 'two-sided', 'less' or " + "'greater'.") + + if confidence_level < 0 or confidence_level > 1: + raise ValueError('confidence_level must be between 0 and 1') + + if self._kind == 'conditional': + ci = self._conditional_odds_ratio_ci(confidence_level, alternative) + else: + ci = self._sample_odds_ratio_ci(confidence_level, alternative) + return ci + + def _conditional_odds_ratio_ci(self, confidence_level=0.95, + alternative='two-sided'): + """ + Confidence interval for the conditional odds ratio. + """ + + table = self._table + if 0 in table.sum(axis=0) or 0 in table.sum(axis=1): + # If both values in a row or column are zero, the p-value is 1, + # the odds ratio is NaN and the confidence interval is (0, inf). + ci = (0, np.inf) + else: + ci = _conditional_oddsratio_ci(table, + confidence_level=confidence_level, + alternative=alternative) + return ConfidenceInterval(low=ci[0], high=ci[1]) + + def _sample_odds_ratio_ci(self, confidence_level=0.95, + alternative='two-sided'): + """ + Confidence interval for the sample odds ratio. + """ + if confidence_level < 0 or confidence_level > 1: + raise ValueError('confidence_level must be between 0 and 1') + + table = self._table + if 0 in table.sum(axis=0) or 0 in table.sum(axis=1): + # If both values in a row or column are zero, the p-value is 1, + # the odds ratio is NaN and the confidence interval is (0, inf). + ci = (0, np.inf) + else: + ci = _sample_odds_ratio_ci(table, + confidence_level=confidence_level, + alternative=alternative) + return ConfidenceInterval(low=ci[0], high=ci[1]) + + +def odds_ratio(table, *, kind='conditional'): + r""" + Compute the odds ratio for a 2x2 contingency table. + + Parameters + ---------- + table : array_like of ints + A 2x2 contingency table. Elements must be non-negative integers. + kind : str, optional + Which kind of odds ratio to compute, either the sample + odds ratio (``kind='sample'``) or the conditional odds ratio + (``kind='conditional'``). Default is ``'conditional'``. + + Returns + ------- + result : `~scipy.stats._result_classes.OddsRatioResult` instance + The returned object has two computed attributes: + + statistic : float + * If `kind` is ``'sample'``, this is sample (or unconditional) + estimate, given by + ``table[0, 0]*table[1, 1]/(table[0, 1]*table[1, 0])``. + * If `kind` is ``'conditional'``, this is the conditional + maximum likelihood estimate for the odds ratio. It is + the noncentrality parameter of Fisher's noncentral + hypergeometric distribution with the same hypergeometric + parameters as `table` and whose mean is ``table[0, 0]``. + + The object has the method `confidence_interval` that computes + the confidence interval of the odds ratio. + + See Also + -------- + scipy.stats.fisher_exact + relative_risk + :ref:`hypothesis_odds_ratio` : Extended example + + Notes + ----- + The conditional odds ratio was discussed by Fisher (see "Example 1" + of [1]_). Texts that cover the odds ratio include [2]_ and [3]_. + + .. versionadded:: 1.10.0 + + References + ---------- + .. [1] R. A. Fisher (1935), The logic of inductive inference, + Journal of the Royal Statistical Society, Vol. 98, No. 1, + pp. 39-82. + .. [2] Breslow NE, Day NE (1980). Statistical methods in cancer research. + Volume I - The analysis of case-control studies. IARC Sci Publ. + (32):5-338. PMID: 7216345. (See section 4.2.) + .. [3] H. Sahai and A. Khurshid (1996), Statistics in Epidemiology: + Methods, Techniques, and Applications, CRC Press LLC, Boca + Raton, Florida. + + Examples + -------- + In epidemiology, individuals are classified as "exposed" or + "unexposed" to some factor or treatment. If the occurrence of some + illness is under study, those who have the illness are often + classified as "cases", and those without it are "noncases". The + counts of the occurrences of these classes gives a contingency + table:: + + exposed unexposed + cases a b + noncases c d + + The sample odds ratio may be written ``(a/c) / (b/d)``. ``a/c`` can + be interpreted as the odds of a case occurring in the exposed group, + and ``b/d`` as the odds of a case occurring in the unexposed group. + The sample odds ratio is the ratio of these odds. If the odds ratio + is greater than 1, it suggests that there is a positive association + between being exposed and being a case. + + Interchanging the rows or columns of the contingency table inverts + the odds ratio, so it is important to understand the meaning of labels + given to the rows and columns of the table when interpreting the + odds ratio. + + Consider a hypothetical example where it is hypothesized that exposure to a + certain chemical is associated with increased occurrence of a certain + disease. Suppose we have the following table for a collection of 410 people:: + + exposed unexposed + cases 7 15 + noncases 58 472 + + The question we ask is "Is exposure to the chemical associated with + increased risk of the disease?" + + Compute the odds ratio: + + >>> from scipy.stats.contingency import odds_ratio + >>> res = odds_ratio([[7, 15], [58, 472]]) + >>> res.statistic + 3.7836687705553493 + + For this sample, the odds of getting the disease for those who have been + exposed to the chemical are almost 3.8 times that of those who have not been + exposed. + + We can compute the 95% confidence interval for the odds ratio: + + >>> res.confidence_interval(confidence_level=0.95) + ConfidenceInterval(low=1.2514829132266785, high=10.363493716701269) + + The 95% confidence interval for the conditional odds ratio is approximately + (1.25, 10.4). + + For a more detailed example, see :ref:`hypothesis_odds_ratio`. + """ + if kind not in ['conditional', 'sample']: + raise ValueError("`kind` must be 'conditional' or 'sample'.") + + c = np.asarray(table) + + if c.shape != (2, 2): + raise ValueError(f"Invalid shape {c.shape}. The input `table` must be " + "of shape (2, 2).") + + if not np.issubdtype(c.dtype, np.integer): + raise ValueError("`table` must be an array of integers, but got " + f"type {c.dtype}") + c = c.astype(np.int64) + + if np.any(c < 0): + raise ValueError("All values in `table` must be nonnegative.") + + if 0 in c.sum(axis=0) or 0 in c.sum(axis=1): + # If both values in a row or column are zero, the p-value is NaN and + # the odds ratio is NaN. + result = OddsRatioResult(_table=c, _kind=kind, statistic=np.nan) + return result + + if kind == 'sample': + oddsratio = _sample_odds_ratio(c) + else: # kind is 'conditional' + oddsratio = _conditional_oddsratio(c) + + result = OddsRatioResult(_table=c, _kind=kind, statistic=oddsratio) + return result diff --git a/v12x/pylib/scipy/stats/_page_trend_test.py b/v12x/pylib/scipy/stats/_page_trend_test.py new file mode 100644 index 0000000000000000000000000000000000000000..27f8dda8c043eafb0c0380f0cd3269afd30046e9 --- /dev/null +++ b/v12x/pylib/scipy/stats/_page_trend_test.py @@ -0,0 +1,488 @@ +from dataclasses import dataclass +from itertools import permutations +import math +import threading + +import numpy as np + +from ._continuous_distns import norm +from scipy._lib._array_api import xp_capabilities +import scipy.stats + + +@dataclass +class PageTrendTestResult: + statistic: float + pvalue: float + method: str + + +@xp_capabilities(np_only=True) +def page_trend_test(data, ranked=False, predicted_ranks=None, method='auto'): + r""" + Perform Page's Test, a measure of trend in observations between treatments. + + Page's Test (also known as Page's :math:`L` test) is useful when: + + * there are :math:`n \geq 3` treatments, + * :math:`m \geq 2` subjects are observed for each treatment, and + * the observations are hypothesized to have a particular order. + + Specifically, the test considers the null hypothesis that + + .. math:: + + m_1 = m_2 = m_3 \cdots = m_n, + + where :math:`m_j` is the mean of the observed quantity under treatment + :math:`j`, against the alternative hypothesis that + + .. math:: + + m_1 \leq m_2 \leq m_3 \leq \cdots \leq m_n, + + where at least one inequality is strict. + + As noted by [4]_, Page's :math:`L` test has greater statistical power than + the Friedman test against the alternative that there is a difference in + trend, as Friedman's test only considers a difference in the means of the + observations without considering their order. Whereas Spearman :math:`\rho` + considers the correlation between the ranked observations of two variables + (e.g. the airspeed velocity of a swallow vs. the weight of the coconut it + carries), Page's :math:`L` is concerned with a trend in an observation + (e.g. the airspeed velocity of a swallow) across several distinct + treatments (e.g. carrying each of five coconuts of different weight) even + as the observation is repeated with multiple subjects (e.g. one European + swallow and one African swallow). + + Parameters + ---------- + data : array-like + A :math:`m \times n` array; the element in row :math:`i` and + column :math:`j` is the observation corresponding with subject + :math:`i` and treatment :math:`j`. By default, the columns are + assumed to be arranged in order of increasing predicted mean. + + ranked : bool, optional + By default, `data` is assumed to be observations rather than ranks; + it will be ranked with `scipy.stats.rankdata` along ``axis=1``. If + `data` is provided in the form of ranks, pass argument ``True``. + + predicted_ranks : array-like, optional + The predicted ranks of the column means. If not specified, + the columns are assumed to be arranged in order of increasing + predicted mean, so the default `predicted_ranks` are + :math:`[1, 2, \dots, n-1, n]`. + + method : {'auto', 'asymptotic', 'exact'}, optional + Selects the method used to calculate the *p*-value. The following + options are available. + + * 'auto': selects between 'exact' and 'asymptotic' to + achieve reasonably accurate results in reasonable time (default) + * 'asymptotic': compares the standardized test statistic against + the normal distribution + * 'exact': computes the exact *p*-value by comparing the observed + :math:`L` statistic against those realized by all possible + permutations of ranks (under the null hypothesis that each + permutation is equally likely) + + Returns + ------- + res : PageTrendTestResult + An object containing attributes: + + statistic : float + Page's :math:`L` test statistic. + pvalue : float + The associated *p*-value + method : {'asymptotic', 'exact'} + The method used to compute the *p*-value + + See Also + -------- + rankdata, friedmanchisquare, spearmanr + + Notes + ----- + As noted in [1]_, "the :math:`n` 'treatments' could just as well represent + :math:`n` objects or events or performances or persons or trials ranked." + Similarly, the :math:`m` 'subjects' could equally stand for :math:`m` + "groupings by ability or some other control variable, or judges doing + the ranking, or random replications of some other sort." + + The procedure for calculating the :math:`L` statistic, adapted from + [1]_, is: + + 1. "Predetermine with careful logic the appropriate hypotheses + concerning the predicted ordering of the experimental results. + If no reasonable basis for ordering any treatments is known, the + :math:`L` test is not appropriate." + 2. "As in other experiments, determine at what level of confidence + you will reject the null hypothesis that there is no agreement of + experimental results with the monotonic hypothesis." + 3. "Cast the experimental material into a two-way table of :math:`n` + columns (treatments, objects ranked, conditions) and :math:`m` + rows (subjects, replication groups, levels of control variables)." + 4. "When experimental observations are recorded, rank them across each + row", e.g. ``ranks = scipy.stats.rankdata(data, axis=1)``. + 5. "Add the ranks in each column", e.g. + ``colsums = np.sum(ranks, axis=0)``. + 6. "Multiply each sum of ranks by the predicted rank for that same + column", e.g. ``products = predicted_ranks * colsums``. + 7. "Sum all such products", e.g. ``L = products.sum()``. + + [1]_ continues by suggesting use of the standardized statistic + + .. math:: + + \chi_L^2 = \frac{\left[12L-3mn(n+1)^2\right]^2}{mn^2(n^2-1)(n+1)} + + "which is distributed approximately as chi-square with 1 degree of + freedom. The ordinary use of :math:`\chi^2` tables would be + equivalent to a two-sided test of agreement. If a one-sided test + is desired, *as will almost always be the case*, the probability + discovered in the chi-square table should be *halved*." + + However, this standardized statistic does not distinguish between the + observed values being well correlated with the predicted ranks and being + _anti_-correlated with the predicted ranks. Instead, we follow [2]_ + and calculate the standardized statistic + + .. math:: + + \Lambda = \frac{L - E_0}{\sqrt{V_0}}, + + where :math:`E_0 = \frac{1}{4} mn(n+1)^2` and + :math:`V_0 = \frac{1}{144} mn^2(n+1)(n^2-1)`, "which is asymptotically + normal under the null hypothesis". + + The *p*-value for ``method='exact'`` is generated by comparing the observed + value of :math:`L` against the :math:`L` values generated for all + :math:`(n!)^m` possible permutations of ranks. The calculation is performed + using the recursive method of [5]. + + The *p*-values are not adjusted for the possibility of ties. When + ties are present, the reported ``'exact'`` *p*-values may be somewhat + larger (i.e. more conservative) than the true *p*-value [2]_. The + ``'asymptotic'`` *p*-values, however, tend to be smaller (i.e. less + conservative) than the ``'exact'`` *p*-values. + + References + ---------- + .. [1] Ellis Batten Page, "Ordered hypotheses for multiple treatments: + a significant test for linear ranks", *Journal of the American + Statistical Association* 58(301), p. 216--230, 1963. + + .. [2] Markus Neuhauser, *Nonparametric Statistical Test: A computational + approach*, CRC Press, p. 150--152, 2012. + + .. [3] Statext LLC, "Page's L Trend Test - Easy Statistics", *Statext - + Statistics Study*, https://www.statext.com/practice/PageTrendTest03.php, + Accessed July 12, 2020. + + .. [4] "Page's Trend Test", *Wikipedia*, WikimediaFoundation, + https://en.wikipedia.org/wiki/Page%27s_trend_test, + Accessed July 12, 2020. + + .. [5] Robert E. Odeh, "The exact distribution of Page's L-statistic in + the two-way layout", *Communications in Statistics - Simulation and + Computation*, 6(1), p. 49--61, 1977. + + Examples + -------- + We use the example from [3]_: 10 students are asked to rate three + teaching methods - tutorial, lecture, and seminar - on a scale of 1-5, + with 1 being the lowest and 5 being the highest. We have decided that + a confidence level of 99% is required to reject the null hypothesis in + favor of our alternative: that the seminar will have the highest ratings + and the tutorial will have the lowest. Initially, the data have been + tabulated with each row representing an individual student's ratings of + the three methods in the following order: tutorial, lecture, seminar. + + >>> table = [[3, 4, 3], + ... [2, 2, 4], + ... [3, 3, 5], + ... [1, 3, 2], + ... [2, 3, 2], + ... [2, 4, 5], + ... [1, 2, 4], + ... [3, 4, 4], + ... [2, 4, 5], + ... [1, 3, 4]] + + Because the tutorial is hypothesized to have the lowest ratings, the + column corresponding with tutorial rankings should be first; the seminar + is hypothesized to have the highest ratings, so its column should be last. + Since the columns are already arranged in this order of increasing + predicted mean, we can pass the table directly into `page_trend_test`. + + >>> from scipy.stats import page_trend_test + >>> res = page_trend_test(table) + >>> res + PageTrendTestResult(statistic=133.5, pvalue=0.0018191161948127822, + method='exact') + + This *p*-value indicates that there is a 0.1819% chance that + the :math:`L` statistic would reach such an extreme value under the null + hypothesis. Because 0.1819% is less than 1%, we have evidence to reject + the null hypothesis in favor of our alternative at a 99% confidence level. + + The value of the :math:`L` statistic is 133.5. To check this manually, + we rank the data such that high scores correspond with high ranks, settling + ties with an average rank: + + >>> from scipy.stats import rankdata + >>> ranks = rankdata(table, axis=1) + >>> ranks + array([[1.5, 3. , 1.5], + [1.5, 1.5, 3. ], + [1.5, 1.5, 3. ], + [1. , 3. , 2. ], + [1.5, 3. , 1.5], + [1. , 2. , 3. ], + [1. , 2. , 3. ], + [1. , 2.5, 2.5], + [1. , 2. , 3. ], + [1. , 2. , 3. ]]) + + We add the ranks within each column, multiply the sums by the + predicted ranks, and sum the products. + + >>> import numpy as np + >>> m, n = ranks.shape + >>> predicted_ranks = np.arange(1, n+1) + >>> L = (predicted_ranks * np.sum(ranks, axis=0)).sum() + >>> res.statistic == L + True + + As presented in [3]_, the asymptotic approximation of the *p*-value is the + survival function of the normal distribution evaluated at the standardized + test statistic: + + >>> from scipy.stats import norm + >>> E0 = (m*n*(n+1)**2)/4 + >>> V0 = (m*n**2*(n+1)*(n**2-1))/144 + >>> Lambda = (L-E0)/np.sqrt(V0) + >>> p = norm.sf(Lambda) + >>> p + 0.0012693433690751756 + + This does not precisely match the *p*-value reported by `page_trend_test` + above. The asymptotic distribution is not very accurate, nor conservative, + for :math:`m \leq 12` and :math:`n \leq 8`, so `page_trend_test` chose to + use ``method='exact'`` based on the dimensions of the table and the + recommendations in Page's original paper [1]_. To override + `page_trend_test`'s choice, provide the `method` argument. + + >>> res = page_trend_test(table, method="asymptotic") + >>> res + PageTrendTestResult(statistic=133.5, pvalue=0.0012693433690751756, + method='asymptotic') + + If the data are already ranked, we can pass in the ``ranks`` instead of + the ``table`` to save computation time. + + >>> res = page_trend_test(ranks, # ranks of data + ... ranked=True, # data is already ranked + ... ) + >>> res + PageTrendTestResult(statistic=133.5, pvalue=0.0018191161948127822, + method='exact') + + Suppose the raw data had been tabulated in an order different from the + order of predicted means, say lecture, seminar, tutorial. + + >>> table = np.asarray(table)[:, [1, 2, 0]] + + Since the arrangement of this table is not consistent with the assumed + ordering, we can either rearrange the table or provide the + `predicted_ranks`. Remembering that the lecture is predicted + to have the middle rank, the seminar the highest, and tutorial the lowest, + we pass: + + >>> res = page_trend_test(table, # data as originally tabulated + ... predicted_ranks=[2, 3, 1], # our predicted order + ... ) + >>> res + PageTrendTestResult(statistic=133.5, pvalue=0.0018191161948127822, + method='exact') + + """ + if not hasattr(_pagel_state, 'state'): + _pagel_state.state = _PageL() + + # Possible values of the method parameter and the corresponding function + # used to evaluate the p value + methods = {"asymptotic": _l_p_asymptotic, + "exact": _l_p_exact, + "auto": None} + if method not in methods: + raise ValueError(f"`method` must be in {set(methods)}") + + ranks = np.asarray(data) + if ranks.ndim != 2: # TODO: relax this to accept 3d arrays? + raise ValueError("`data` must be a 2d array.") + + m, n = ranks.shape + if m < 2 or n < 3: + raise ValueError("Page's L is only appropriate for data with two " + "or more rows and three or more columns.") + + if np.any(np.isnan(data)): + raise ValueError("`data` contains NaNs, which cannot be ranked " + "meaningfully") + + # ensure NumPy array and rank the data if it's not already ranked + if ranked: + # Only a basic check on whether data is ranked. Checking that the data + # is properly ranked could take as much time as ranking it. + if not (ranks.min() >= 1 and ranks.max() <= ranks.shape[1]): + raise ValueError("`data` is not properly ranked. Rank the data or " + "pass `ranked=False`.") + else: + ranks = scipy.stats.rankdata(data, axis=-1) + + # generate predicted ranks if not provided, ensure valid NumPy array + if predicted_ranks is None: + predicted_ranks = np.arange(1, n+1) + else: + predicted_ranks = np.asarray(predicted_ranks) + if (predicted_ranks.ndim < 1 or + (set(predicted_ranks) != set(range(1, n+1)) or + len(predicted_ranks) != n)): + raise ValueError(f"`predicted_ranks` must include each integer " + f"from 1 to {n} (the number of columns in " + f"`data`) exactly once.") + + if not isinstance(ranked, bool): + raise TypeError("`ranked` must be boolean.") + + # Calculate the L statistic + L = _l_vectorized(ranks, predicted_ranks) + + # Calculate the p-value + if method == "auto": + method = _choose_method(ranks) + p_fun = methods[method] # get the function corresponding with the method + p = p_fun(L, m, n) + + page_result = PageTrendTestResult(statistic=L, pvalue=p, method=method) + return page_result + + +def _choose_method(ranks): + '''Choose method for computing p-value automatically''' + m, n = ranks.shape + if n > 8 or (m > 12 and n > 3) or m > 20: # as in [1], [4] + method = "asymptotic" + else: + method = "exact" + return method + + +def _l_vectorized(ranks, predicted_ranks): + '''Calculate's Page's L statistic for each page of a 3d array''' + colsums = ranks.sum(axis=-2, keepdims=True) + products = predicted_ranks * colsums + Ls = products.sum(axis=-1) + Ls = Ls[0] if Ls.size == 1 else Ls.ravel() + return Ls + + +def _l_p_asymptotic(L, m, n): + '''Calculate the p-value of Page's L from the asymptotic distribution''' + # Using [1] as a reference, the asymptotic p-value would be calculated as: + # chi_L = (12*L - 3*m*n*(n+1)**2)**2/(m*n**2*(n**2-1)*(n+1)) + # p = chi2.sf(chi_L, df=1, loc=0, scale=1)/2 + # but this is insensitive to the direction of the hypothesized ranking + + # See [2] page 151 + E0 = (m*n*(n+1)**2)/4 + V0 = (m*n**2*(n+1)*(n**2-1))/144 + Lambda = (L-E0)/np.sqrt(V0) + # This is a one-sided "greater" test - calculate the probability that the + # L statistic under H0 would be greater than the observed L statistic + p = norm.sf(Lambda) + return p + + +def _l_p_exact(L, m, n): + '''Calculate the p-value of Page's L exactly''' + # [1] uses m, n; [5] uses n, k. + # Switch convention here because exact calculation code references [5]. + L, n, k = int(L), int(m), int(n) + _pagel_state.state.set_k(k) + return _pagel_state.state.sf(L, n) + + +class _PageL: + '''Maintains state between `page_trend_test` executions''' + + def __init__(self): + '''Lightweight initialization''' + self.all_pmfs = {} + + def set_k(self, k): + '''Calculate lower and upper limits of L for single row''' + self.k = k + # See [5] top of page 52 + self.a, self.b = (k*(k+1)*(k+2))//6, (k*(k+1)*(2*k+1))//6 + + def sf(self, l, n): + '''Survival function of Page's L statistic''' + ps = [self.pmf(l, n) for l in range(l, n*self.b + 1)] + return np.sum(ps) + + def p_l_k_1(self): + '''Relative frequency of each L value over all possible single rows''' + + # See [5] Equation (6) + ranks = range(1, self.k+1) + # generate all possible rows of length k + rank_perms = np.array(list(permutations(ranks))) + # compute Page's L for all possible rows + Ls = (ranks*rank_perms).sum(axis=1) + # count occurrences of each L value + counts = np.histogram(Ls, np.arange(self.a-0.5, self.b+1.5))[0] + # factorial(k) is number of possible permutations + return counts/math.factorial(self.k) + + def pmf(self, l, n): + '''Recursive function to evaluate p(l, k, n); see [5] Equation 1''' + + if n not in self.all_pmfs: + self.all_pmfs[n] = {} + if self.k not in self.all_pmfs[n]: + self.all_pmfs[n][self.k] = {} + + # Cache results to avoid repeating calculation. Initially this was + # written with lru_cache, but this seems faster? Also, we could add + # an option to save this for future lookup. + if l in self.all_pmfs[n][self.k]: + return self.all_pmfs[n][self.k][l] + + if n == 1: + ps = self.p_l_k_1() # [5] Equation 6 + ls = range(self.a, self.b+1) + # not fast, but we'll only be here once + self.all_pmfs[n][self.k] = {l: p for l, p in zip(ls, ps)} + return self.all_pmfs[n][self.k][l] + + p = 0 + low = max(l-(n-1)*self.b, self.a) # [5] Equation 2 + high = min(l-(n-1)*self.a, self.b) + + # [5] Equation 1 + for t in range(low, high+1): + p1 = self.pmf(l-t, n-1) + p2 = self.pmf(t, 1) + p += p1*p2 + self.all_pmfs[n][self.k][l] = p + return p + + +# Maintain state for faster repeat calls to page_trend_test w/ method='exact' +# _PageL() is calculated once per thread and stored as an attribute on +# this thread-local variable inside page_trend_test(). +_pagel_state = threading.local() diff --git a/v12x/pylib/scipy/stats/_probability_distribution.py b/v12x/pylib/scipy/stats/_probability_distribution.py new file mode 100644 index 0000000000000000000000000000000000000000..71ed8f48920f3031a88cf1285fae25e4bb28eefa --- /dev/null +++ b/v12x/pylib/scipy/stats/_probability_distribution.py @@ -0,0 +1,2091 @@ +# Temporary file separated from _distribution_infrastructure.py +# to simplify the diff during PR review. +from abc import ABC, abstractmethod +from types import GenericAlias + +class _ProbabilityDistribution(ABC): + + # generic type compatibility with scipy-stubs + __class_getitem__: classmethod = classmethod(GenericAlias) + + @abstractmethod + def support(self): + r"""Support of the random variable + + The support of a random variable is set of all possible outcomes; + i.e., the subset of the domain of argument :math:`x` for which + the probability density function :math:`f(x)` is nonzero. + + This function returns lower and upper bounds of the support. + + Returns + ------- + out : tuple of Array + The lower and upper bounds of the support. + + See Also + -------- + pdf + + References + ---------- + .. [1] Support (mathematics), *Wikipedia*, + https://en.wikipedia.org/wiki/Support_(mathematics) + + Notes + ----- + Suppose a continuous probability distribution has support ``(l, r)``. + The following table summarizes the value returned by several + methods when the argument is outside the support. + + +----------------+---------------------+---------------------+ + | Method | Value for ``x < l`` | Value for ``x > r`` | + +================+=====================+=====================+ + | ``pdf(x)`` | 0 | 0 | + +----------------+---------------------+---------------------+ + | ``logpdf(x)`` | -inf | -inf | + +----------------+---------------------+---------------------+ + | ``cdf(x)`` | 0 | 1 | + +----------------+---------------------+---------------------+ + | ``logcdf(x)`` | -inf | 0 | + +----------------+---------------------+---------------------+ + | ``ccdf(x)`` | 1 | 0 | + +----------------+---------------------+---------------------+ + | ``logccdf(x)`` | 0 | -inf | + +----------------+---------------------+---------------------+ + + For discrete distributions, the same table is applicable with + ``pmf`` and ``logpmf`` substituted for ``pdf`` and ``logpdf``. + + For the ``cdf`` and related methods of continuous distributions, the + inequality need not be strict; i.e. the tabulated value is returned + when the method is evaluated *at* the corresponding boundary. + + The following table summarizes the value returned by the inverse + methods for arguments ``0`` and ``1``, whether the distribution + is continuous or discrete. + + +-------------+-----------+-----------+ + | Method | ``x = 0`` | ``x = 1`` | + +=============+===========+===========+ + | ``icdf(x)`` | ``l`` | ``r`` | + +-------------+-----------+-----------+ + | ``icdf(x)`` | ``r`` | ``l`` | + +-------------+-----------+-----------+ + + For the inverse log-functions, the same values are returned + for ``x = log(0)`` and ``x = log(1)``. All inverse functions return + ``nan`` when evaluated at an argument outside the domain ``0`` to ``1``. + + Examples + -------- + Instantiate a distribution with the desired parameters: + + >>> from scipy import stats + >>> X = stats.Uniform(a=-0.5, b=0.5) + + Retrieve the support of the distribution: + + >>> X.support() + (-0.5, 0.5) + + For a distribution with infinite support, + + >>> X = stats.Normal() + >>> X.support() + (-inf, inf) + + Due to underflow, the numerical value returned by the PDF may be zero + even for arguments within the support, even if the true value is + nonzero. In such cases, the log-PDF may be useful. + + >>> X.pdf([-100., 100.]) + array([0., 0.]) + >>> X.logpdf([-100., 100.]) + array([-5000.91893853, -5000.91893853]) + + Use cases for the log-CDF and related methods are analogous. + + """ + raise NotImplementedError() + + @abstractmethod + def sample(self, shape, *, method, rng): + r"""Random sample from the distribution. + + Parameters + ---------- + shape : tuple of ints, default: () + The shape of the sample to draw. If the parameters of the distribution + underlying the random variable are arrays of shape ``param_shape``, + the output array will be of shape ``shape + param_shape``. + method : {None, 'formula', 'inverse_transform'} + The strategy used to produce the sample. By default (``None``), + the infrastructure chooses between the following options, + listed in order of precedence. + + - ``'formula'``: an implementation specific to the distribution + - ``'inverse_transform'``: generate a uniformly distributed sample and + return the inverse CDF at these arguments. + + Not all `method` options are available for all distributions. + If the selected `method` is not available, a `NotImplementedError`` + will be raised. + rng : `numpy.random.Generator` or `scipy.stats.qmc.QMCEngine`, optional + Pseudo- or quasi-random number generator state. When `rng` is None, + a new `numpy.random.Generator` is created using entropy from the + operating system. Types other than `numpy.random.Generator` and + `scipy.stats.qmc.QMCEngine` are passed to `numpy.random.default_rng` + to instantiate a ``Generator``. + + If `rng` is an instance of `scipy.stats.qmc.QMCEngine` configured to use + scrambling and `shape` is not empty, then each slice along the zeroth + axis of the result is a "quasi-independent", low-discrepancy sequence; + that is, they are distinct sequences that can be treated as statistically + independent for most practical purposes. Separate calls to `sample` + produce new quasi-independent, low-discrepancy sequences. The dimensionality + (``d``) of the provided ``QMCEngine`` is ignored; new instances of the same + class and configuration options (``scramble``, ``optimization``, and + ``bits``, where applicable) will be created with ``d=1``. + + References + ---------- + .. [1] Sampling (statistics), *Wikipedia*, + https://en.wikipedia.org/wiki/Sampling_(statistics) + + Examples + -------- + Instantiate a distribution with the desired parameters: + + >>> import numpy as np + >>> from scipy import stats + >>> X = stats.Uniform(a=0., b=1.) + + Generate a pseudorandom sample: + + >>> x = X.sample((1000, 1)) + >>> octiles = (np.arange(8) + 1) / 8 + >>> np.count_nonzero(x <= octiles, axis=0) + array([ 148, 263, 387, 516, 636, 751, 865, 1000]) # may vary + + >>> X = stats.Uniform(a=np.zeros((3, 1)), b=np.ones(2)) + >>> X.a.shape, + (3, 2) + >>> x = X.sample(shape=(5, 4)) + >>> x.shape + (5, 4, 3, 2) + + """ + raise NotImplementedError() + + @abstractmethod + def moment(self, order, kind, *, method): + r"""Raw, central, or standard moment of positive integer order. + + In terms of probability density function :math:`f(x)` and support + :math:`\chi`, the "raw" moment (about the origin) of order :math:`n` of + a continuous random variable :math:`X` is: + + .. math:: + + \mu'_n(X) = \int_{\chi} x^n f(x) dx + + The "central" moment is the raw moment taken about the mean, + :math:`\mu = \mu'_1`: + + .. math:: + + \mu_n(X) = \int_{\chi} (x - \mu) ^n f(x) dx + + The "standardized" moment is the central moment normalized by the + :math:`n^\text{th}` power of the standard deviation + :math:`\sigma = \sqrt{\mu_2}` to produce a scale invariant quantity: + + .. math:: + + \tilde{\mu}_n(X) = \frac{\mu_n(X)} + {\sigma^n} + + The definitions for discrete random variables are analogous, with + sums over the support replacing the integrals. + + Parameters + ---------- + order : int + The integer order of the moment; i.e. :math:`n` in the formulae above. + kind : {'raw', 'central', 'standardized'} + Whether to return the raw (default), central, or standardized moment + defined above. + method : {None, 'formula', 'general', 'transform', 'normalize', 'quadrature', 'cache'} + The strategy used to evaluate the moment. By default (``None``), + the infrastructure chooses between the following options, + listed in order of precedence. + + - ``'cache'``: use the value of the moment most recently calculated + via another method + - ``'formula'``: use a formula for the moment itself + - ``'general'``: use a general result that is true for all distributions + with finite moments; for instance, the zeroth raw moment is + identically 1 + - ``'transform'``: transform a raw moment to a central moment or + vice versa (see Notes) + - ``'normalize'``: normalize a central moment to get a standardized + or vice versa + - ``'quadrature'``: numerically integrate (or, in the discrete case, sum) + according to the definition + + Not all `method` options are available for all orders, kinds, and + distributions. If the selected `method` is not available, a + ``NotImplementedError`` will be raised. + + Returns + ------- + out : array + The moment of the random variable of the specified order and kind. + + See Also + -------- + pdf + mean + variance + standard_deviation + skewness + kurtosis + lmoment + + Notes + ----- + Not all distributions have finite moments of all orders; moments of some + orders may be undefined or infinite. If a formula for the moment is not + specifically implemented for the chosen distribution, SciPy will attempt + to compute the moment via a generic method, which may yield a finite + result where none exists. This is not a critical bug, but an opportunity + for an enhancement. + + The definition of a raw moment in the summary is specific to the raw moment + about the origin. The raw moment about any point :math:`a` is: + + .. math:: + + E[(X-a)^n] = \int_{\chi} (x-a)^n f(x) dx + + In this notation, a raw moment about the origin is :math:`\mu'_n = E[x^n]`, + and a central moment is :math:`\mu_n = E[(x-\mu)^n]`, where :math:`\mu` + is the first raw moment; i.e. the mean. + + The ``'transform'`` method takes advantage of the following relationships + between moments taken about different points :math:`a` and :math:`b`. + + .. math:: + + E[(X-b)^n] = \sum_{i=0}^n E[(X-a)^i] {n \choose i} (a - b)^{n-i} + + For instance, to transform the raw moment to the central moment, we let + :math:`b = \mu` and :math:`a = 0`. + + The distribution infrastructure provides flexibility for distribution + authors to implement separate formulas for raw moments, central moments, + and standardized moments of any order. By default, the moment of the + desired order and kind is evaluated from the formula if such a formula + is available; if not, the infrastructure uses any formulas that are + available rather than resorting directly to numerical integration. + For instance, if formulas for the first three raw moments are + available and the third standardized moments is desired, the + infrastructure will evaluate the raw moments and perform the transforms + and standardization required. The decision tree is somewhat complex, + but the strategy for obtaining a moment of a given order and kind + (possibly as an intermediate step due to the recursive nature of the + transform formula above) roughly follows this order of priority: + + #. Use cache (if order of same moment and kind has been calculated) + #. Use formula (if available) + #. Transform between raw and central moment and/or normalize to convert + between central and standardized moments (if efficient) + #. Use a generic result true for most distributions (if available) + #. Use quadrature + + References + ---------- + .. [1] Moment, *Wikipedia*, + https://en.wikipedia.org/wiki/Moment_(mathematics) + + Examples + -------- + Instantiate a distribution with the desired parameters: + + >>> from scipy import stats + >>> X = stats.Normal(mu=1., sigma=2.) + + Evaluate the first raw moment: + + >>> X.moment(order=1, kind='raw') + 1.0 + >>> X.moment(order=1, kind='raw') == X.mean() == X.mu + True + + Evaluate the second central moment: + + >>> X.moment(order=2, kind='central') + 4.0 + >>> X.moment(order=2, kind='central') == X.variance() == X.sigma**2 + True + + Evaluate the fourth standardized moment: + + >>> X.moment(order=4, kind='standardized') + 3.0 + >>> X.moment(order=4, kind='standardized') == X.kurtosis(convention='non-excess') + True + + """ # noqa:E501 + raise NotImplementedError() + + @abstractmethod + def lmoment(self, order, *, standardize, method): + r"""L-moment or L-moment ratio of positive integer order. + + The L-moment of order :math:`n` of a continuous random variable :math:`X` is: + + .. math:: + + \lambda_n(X) = \frac{1}{n} \sum_{k=0}^{n-1} (-1)^{k} {{n-1}\choose k} E[X_{(n-k)}] + + where :math:`X_{(1)}, \dots, X_{(r)}, \dots, X_{(n)}` are the order statistics + of an independent sample of size :math:`n`. + + The L-moment can also be expressed in terms of the random variable's inverse + cumulative distribution function :math:`F^{-1}` and the shifted Legendre + polynomial :math:`\widetilde{P}_{n-1}`: + + .. math:: + + \lambda_n(X) = \int_0^1 F^{-1}(p) \widetilde{P}_{n-1}(p) dp + + For order :math:`n \geq 3`, the "standardized" L-moment, known as the L-moment + ratio, is the L-moment normalized by the L-moment of order 2, resulting in a + scale invariant quantity: + + .. math:: + + \tau_n(X) = \frac{\lambda_n(X)} + {\lambda_2(X)} + + Parameters + ---------- + order : int + The positive integer order of the L-moment; i.e. :math:`n` in the formulae + above. + standardize : bool, default: True + Whether to return L-moment ratios for orders 3 and higher. L-moment ratios + are analogous to standardized conventional moments: they are the + non-standardized L-moments divided by the L-moment of order 2. + method : {None, 'formula', 'general', 'order_statistics', 'quadrature_icdf', 'cache'} + The strategy used to evaluate the L-moment. By default (``None``), + the infrastructure chooses between the following options, + listed in order of precedence. + + - ``'cache'``: use the value of the L-moment most recently calculated + via another method + - ``'formula'``: use a formula specific to the distribution. + - ``'general'``: use a general result that is true for all distributions + with finite L-moments; for instance, the first L-moment is identically + equal to the mean. + - ``'quadrature_icdf'``: numerically integrate according to the definition + in terms of the inverse cumulative distribution function. + - ``'order_statistics'``: compute according to the definition in terms of + order statistics. + + Not all `method` options are available for all orders and distributions. + If the selected `method` is not available, a ``NotImplementedError`` will + be raised. + + Returns + ------- + out : array + The L-moment of the random variable of the specified order. + + See Also + -------- + moment + order_statistic + + Notes + ----- + L-moments are only defined for distributions with finite mean. If a formula for + the L-moment is not specifically implemented for the chosen distribution, SciPy + will attempt to compute the moment via a generic method, which may yield a + finite result where none exists. This is not a critical bug, but an opportunity + for an enhancement. + + SciPy offers only basic capabilities for working with L-moments. For more advanced + features, consider the ``lmo`` package [2]_. + + References + ---------- + .. [1] L-moment, *Wikipedia*, + https://en.wikipedia.org/wiki/L-moment + .. [2] @jorenham, *Lmo*, https://github.com/jorenham/Lmo/ + + Examples + -------- + Instantiate a distribution with the desired parameters: + + >>> import numpy as np + >>> from scipy import stats + >>> X = stats.Normal(mu=1., sigma=2.) + + Evaluate the first L-moment: + + >>> X.lmoment(order=1) + 1.0 + >>> X.lmoment(order=1) == X.mean() == X.mu + True + + Evaluate the second L-moment: + + >>> X.lmoment(order=2) + np.float64(1.1283791670955123) + >>> np.allclose(X.lmoment(order=2), X.sigma / np.sqrt(np.pi)) + True + + Evaluate the fourth L-moment ratio, that is, the L-kurtosis: + + >>> X.lmoment(order=4) + np.float64(0.12260171954089069) + >>> X.lmoment(order=4) == X.lmoment(order=4, standardize=False) / X.lmoment(order=2) + True + + """ # noqa:E501 + raise NotImplementedError() + + @abstractmethod + def mean(self, *, method): + r"""Mean (raw first moment about the origin) + + Parameters + ---------- + method : {None, 'formula', 'transform', 'quadrature', 'cache'} + Method used to calculate the raw first moment. Not + all methods are available for all distributions. See + `moment` for details. + + See Also + -------- + moment + median + mode + + References + ---------- + .. [1] Mean, *Wikipedia*, + https://en.wikipedia.org/wiki/Mean#Mean_of_a_probability_distribution + + Examples + -------- + Instantiate a distribution with the desired parameters: + + >>> from scipy import stats + >>> X = stats.Normal(mu=1., sigma=2.) + + Evaluate the variance: + + >>> X.mean() + 1.0 + >>> X.mean() == X.moment(order=1, kind='raw') == X.mu + True + + """ + raise NotImplementedError() + + @abstractmethod + def median(self, *, method): + r"""Median (50th percentile) + + If a continuous random variable :math:`X` has probability :math:`0.5` of + taking on a value less than :math:`m`, then :math:`m` is the median. + + More generally, a median is a value :math:`m` for which: + + .. math:: + + P(X ≤ m) ≤ 0.5 ≥ P(X ≥ m) + + For discrete random variables, the median may not be unique, in which + case the smallest value satisfying the definition is reported. + + Parameters + ---------- + method : {None, 'formula', 'icdf'} + The strategy used to evaluate the median. + By default (``None``), the infrastructure chooses between the + following options, listed in order of precedence. + + - ``'formula'``: use a formula for the median + - ``'icdf'``: evaluate the inverse CDF of 0.5 + + Not all `method` options are available for all distributions. + If the selected `method` is not available, a ``NotImplementedError`` + will be raised. + + Returns + ------- + out : array + The median + + See Also + -------- + mean + mode + icdf + + References + ---------- + .. [1] Median, *Wikipedia*, + https://en.wikipedia.org/wiki/Median#Probability_distributions + + Examples + -------- + Instantiate a distribution with the desired parameters: + + >>> from scipy import stats + >>> X = stats.Uniform(a=0., b=10.) + + Compute the median: + + >>> X.median() + np.float64(5.0) + >>> X.median() == X.icdf(0.5) == X.iccdf(0.5) + True + + """ + raise NotImplementedError() + + @abstractmethod + def mode(self, *, method): + r"""Mode (most likely value) + + Informally, the mode is a value that a random variable has the highest + probability (density) of assuming. That is, the mode is the element of + the support :math:`\chi` that maximizes the probability density (or mass, + for discrete random variables) function :math:`f(x)`: + + .. math:: + + \text{mode} = \arg\max_{x \in \chi} f(x) + + Parameters + ---------- + method : {None, 'formula', 'optimization'} + The strategy used to evaluate the mode. + By default (``None``), the infrastructure chooses between the + following options, listed in order of precedence. + + - ``'formula'``: use a formula for the median + - ``'optimization'``: numerically maximize the PDF/PMF + + Not all `method` options are available for all distributions. + If the selected `method` is not available, a ``NotImplementedError`` + will be raised. + + Returns + ------- + out : array + The mode + + See Also + -------- + mean + median + pdf + + Notes + ----- + For some distributions + + #. the mode is not unique (e.g. the uniform distribution); + #. the PDF has one or more singularities, and it is debatable whether + a singularity is considered to be in the domain and called the mode + (e.g. the gamma distribution with shape parameter less than 1); and/or + #. the probability density function may have one or more local maxima + that are not a global maximum (e.g. mixture distributions). + + In such cases, `mode` will + + #. return a single value, + #. consider the mode to occur at a singularity, and/or + #. return a local maximum which may or may not be a global maximum. + + If a formula for the mode is not specifically implemented for the + chosen distribution, SciPy will attempt to compute the mode + numerically, which may not meet the user's preferred definition of a + mode. In such cases, the user is encouraged to subclass the + distribution and override ``mode``. + + References + ---------- + .. [1] Mode (statistics), *Wikipedia*, + https://en.wikipedia.org/wiki/Mode_(statistics) + + Examples + -------- + Instantiate a distribution with the desired parameters: + + >>> from scipy import stats + >>> X = stats.Normal(mu=1., sigma=2.) + + Evaluate the mode: + + >>> X.mode() + 1.0 + + If the mode is not uniquely defined, ``mode`` nonetheless returns a + single value. + + >>> X = stats.Uniform(a=0., b=1.) + >>> X.mode() + 0.5 + + If this choice does not satisfy your requirements, subclass the + distribution and override ``mode``: + + >>> class BetterUniform(stats.Uniform): + ... def mode(self): + ... return self.b + >>> X = BetterUniform(a=0., b=1.) + >>> X.mode() + 1.0 + + """ + raise NotImplementedError() + + @abstractmethod + def variance(self, *, method): + r"""Variance (central second moment) + + Parameters + ---------- + method : {None, 'formula', 'transform', 'normalize', 'quadrature', 'cache'} + Method used to calculate the central second moment. Not + all methods are available for all distributions. See + `moment` for details. + + See Also + -------- + moment + standard_deviation + mean + + References + ---------- + .. [1] Variance, *Wikipedia*, + https://en.wikipedia.org/wiki/Variance#Absolutely_continuous_random_variable + + Examples + -------- + Instantiate a distribution with the desired parameters: + + >>> from scipy import stats + >>> X = stats.Normal(mu=1., sigma=2.) + + Evaluate the variance: + + >>> X.variance() + 4.0 + >>> X.variance() == X.moment(order=2, kind='central') == X.sigma**2 + True + + """ + raise NotImplementedError() + + @abstractmethod + def standard_deviation(self, *, method): + r"""Standard deviation (square root of the second central moment) + + Parameters + ---------- + method : {None, 'formula', 'transform', 'normalize', 'quadrature', 'cache'} + Method used to calculate the central second moment. Not + all methods are available for all distributions. See + `moment` for details. + + See Also + -------- + variance + mean + moment + + References + ---------- + .. [1] Standard deviation, *Wikipedia*, + https://en.wikipedia.org/wiki/Standard_deviation#Definition_of_population_values + + Examples + -------- + Instantiate a distribution with the desired parameters: + + >>> from scipy import stats + >>> X = stats.Normal(mu=1., sigma=2.) + + Evaluate the standard deviation: + + >>> X.standard_deviation() + 2.0 + >>> X.standard_deviation() == X.moment(order=2, kind='central')**0.5 == X.sigma + True + + """ + raise NotImplementedError() + + @abstractmethod + def skewness(self, *, method): + r"""Skewness (standardized third moment) + + Parameters + ---------- + method : {None, 'formula', 'general', 'transform', 'normalize', 'cache'} + Method used to calculate the standardized third moment. Not + all methods are available for all distributions. See + `moment` for details. + + See Also + -------- + moment + mean + variance + + References + ---------- + .. [1] Skewness, *Wikipedia*, + https://en.wikipedia.org/wiki/Skewness + + Examples + -------- + Instantiate a distribution with the desired parameters: + + >>> from scipy import stats + >>> X = stats.Normal(mu=1., sigma=2.) + + Evaluate the skewness: + + >>> X.skewness() + 0.0 + >>> X.skewness() == X.moment(order=3, kind='standardized') + True + + """ + raise NotImplementedError() + + @abstractmethod + def kurtosis(self, *, method): + r"""Kurtosis (standardized fourth moment) + + By default, this is the standardized fourth moment, also known as the + "non-excess" or "Pearson" kurtosis (e.g. the kurtosis of the normal + distribution is 3). The "excess" or "Fisher" kurtosis (the standardized + fourth moment minus 3) is available via the `convention` parameter. + + Parameters + ---------- + method : {None, 'formula', 'general', 'transform', 'normalize', 'cache'} + Method used to calculate the standardized fourth moment. Not + all methods are available for all distributions. See + `moment` for details. + convention : {'non-excess', 'excess'} + Two distinct conventions are available: + + - ``'non-excess'``: the standardized fourth moment (Pearson's kurtosis) + - ``'excess'``: the standardized fourth moment minus 3 (Fisher's kurtosis) + + The default is ``'non-excess'``. + + See Also + -------- + moment + mean + variance + + References + ---------- + .. [1] Kurtosis, *Wikipedia*, + https://en.wikipedia.org/wiki/Kurtosis + + Examples + -------- + Instantiate a distribution with the desired parameters: + + >>> from scipy import stats + >>> X = stats.Normal(mu=1., sigma=2.) + + Evaluate the kurtosis: + + >>> X.kurtosis() + 3.0 + >>> (X.kurtosis() + ... == X.kurtosis(convention='excess') + 3. + ... == X.moment(order=4, kind='standardized')) + True + + """ + raise NotImplementedError() + + @abstractmethod + def pdf(self, x, /, *, method): + r"""Probability density function + + The probability density function ("PDF"), denoted :math:`f(x)`, is the + probability *per unit length* that the random variable will assume the + value :math:`x`. Mathematically, it can be defined as the derivative + of the cumulative distribution function :math:`F(x)`: + + .. math:: + + f(x) = \frac{d}{dx} F(x) + + `pdf` accepts `x` for :math:`x`. + + Parameters + ---------- + x : array_like + The argument of the PDF. + method : {None, 'formula', 'logexp'} + The strategy used to evaluate the PDF. By default (``None``), the + infrastructure chooses between the following options, listed in + order of precedence. + + - ``'formula'``: use a formula for the PDF itself + - ``'logexp'``: evaluate the log-PDF and exponentiate + + Not all `method` options are available for all distributions. + If the selected `method` is not available, a ``NotImplementedError`` + will be raised. + + Returns + ------- + out : array + The PDF evaluated at the argument `x`. + + See Also + -------- + cdf + logpdf + + Notes + ----- + Suppose a continuous probability distribution has support :math:`[l, r]`. + By definition of the support, the PDF evaluates to its minimum value + of :math:`0` outside the support; i.e. for :math:`x < l` or + :math:`x > r`. The maximum of the PDF may be less than or greater than + :math:`1`; since the value is a probability *density*, only its integral + over the support must equal :math:`1`. + + For discrete distributions, `pdf` returns ``inf`` at supported points + and ``0`` elsewhere. + + References + ---------- + .. [1] Probability density function, *Wikipedia*, + https://en.wikipedia.org/wiki/Probability_density_function + + Examples + -------- + Instantiate a distribution with the desired parameters: + + >>> from scipy import stats + >>> X = stats.Uniform(a=-1., b=1.) + + Evaluate the PDF at the desired argument: + + >>> X.pdf(0.25) + 0.5 + + """ + raise NotImplementedError() + + @abstractmethod + def logpdf(self, x, /, *, method): + r"""Log of the probability density function + + The probability density function ("PDF"), denoted :math:`f(x)`, is the + probability *per unit length* that the random variable will assume the + value :math:`x`. Mathematically, it can be defined as the derivative + of the cumulative distribution function :math:`F(x)`: + + .. math:: + + f(x) = \frac{d}{dx} F(x) + + `logpdf` computes the logarithm of the probability density function + ("log-PDF"), :math:`\log(f(x))`, but it may be numerically favorable + compared to the naive implementation (computing :math:`f(x)` and + taking the logarithm). + + `logpdf` accepts `x` for :math:`x`. + + Parameters + ---------- + x : array_like + The argument of the log-PDF. + method : {None, 'formula', 'logexp'} + The strategy used to evaluate the log-PDF. By default (``None``), the + infrastructure chooses between the following options, listed in order + of precedence. + + - ``'formula'``: use a formula for the log-PDF itself + - ``'logexp'``: evaluate the PDF and takes its logarithm + + Not all `method` options are available for all distributions. + If the selected `method` is not available, a ``NotImplementedError`` + will be raised. + + Returns + ------- + out : array + The log-PDF evaluated at the argument `x`. + + See Also + -------- + pdf + logcdf + + Notes + ----- + Suppose a continuous probability distribution has support :math:`[l, r]`. + By definition of the support, the log-PDF evaluates to its minimum value + of :math:`-\infty` (i.e. :math:`\log(0)`) outside the support; i.e. for + :math:`x < l` or :math:`x > r`. The maximum of the log-PDF may be less + than or greater than :math:`\log(1) = 0` because the maximum of the PDF + can be any positive real. + + For distributions with infinite support, it is common for `pdf` to return + a value of ``0`` when the argument is theoretically within the support; + this can occur because the true value of the PDF is too small to be + represented by the chosen dtype. The log-PDF, however, will often be finite + (not ``-inf``) over a much larger domain. Consequently, it may be preferred + to work with the logarithms of probabilities and probability densities to + avoid underflow. + + For discrete distributions, `logpdf` returns ``inf`` at supported points and + ``-inf`` (``log(0)``) elsewhere. + + References + ---------- + .. [1] Probability density function, *Wikipedia*, + https://en.wikipedia.org/wiki/Probability_density_function + + Examples + -------- + Instantiate a distribution with the desired parameters: + + >>> import numpy as np + >>> from scipy import stats + >>> X = stats.Uniform(a=-1.0, b=1.0) + + Evaluate the log-PDF at the desired argument: + + >>> X.logpdf(0.5) + -0.6931471805599453 + >>> np.allclose(X.logpdf(0.5), np.log(X.pdf(0.5))) + True + + """ + raise NotImplementedError() + + def pmf(self, x, /, *, method=None): + r"""Probability mass function + + The probability mass function ("PMF"), denoted :math:`f(x)`, is the + probability that the random variable :math:`X` will assume the value :math:`x`. + + .. math:: + + f(x) = P(X = x) + + `pmf` accepts `x` for :math:`x`. + + Parameters + ---------- + x : array_like + The argument of the PMF. + method : {None, 'formula', 'logexp'} + The strategy used to evaluate the PMF. By default (``None``), the + infrastructure chooses between the following options, listed in + order of precedence. + + - ``'formula'``: use a formula for the PMF itself + - ``'logexp'``: evaluate the log-PMF and exponentiate + + Not all `method` options are available for all distributions. + If the selected `method` is not available, a ``NotImplementedError`` + will be raised. + + Returns + ------- + out : array + The PMF evaluated at the argument `x`. + + See Also + -------- + cdf + logpmf + + Notes + ----- + Suppose a discrete probability distribution has support over the integers + :math:`{l, l+1, ..., r-1, r}`. + By definition of the support, the PMF evaluates to its minimum value + of :math:`0` for non-integral :math:`x` and for :math:`x` outside the support; + i.e. for :math:`x < l` or :math:`x > r`. + + For continuous distributions, `pmf` returns ``0`` at all real arguments. + + References + ---------- + .. [1] Probability mass function, *Wikipedia*, + https://en.wikipedia.org/wiki/Probability_mass_function + + Examples + -------- + Instantiate a distribution with the desired parameters: + + >>> from scipy import stats + >>> X = stats.Binomial(n=10, p=0.5) + + Evaluate the PMF at the desired argument: + + >>> X.pmf(5) + np.float64(0.24609375) + + """ + raise NotImplementedError() + + def logpmf(self, x, /, *, method=None): + r"""Log of the probability mass function + + The probability mass function ("PMF"), denoted :math:`f(x)`, is the + probability that the random variable :math:`X` will assume the value :math:`x`. + + .. math:: + + f(x) = \frac{d}{dx} F(x) + + `logpmf` computes the logarithm of the probability mass function + ("log-PMF"), :math:`\log(f(x))`, but it may be numerically favorable + compared to the naive implementation (computing :math:`f(x)` and + taking the logarithm). + + `logpmf` accepts `x` for :math:`x`. + + Parameters + ---------- + x : array_like + The argument of the log-PMF. + method : {None, 'formula', 'logexp'} + The strategy used to evaluate the log-PMF. By default (``None``), the + infrastructure chooses between the following options, listed in order + of precedence. + + - ``'formula'``: use a formula for the log-PMF itself + - ``'logexp'``: evaluate the PMF and takes its logarithm + + Not all `method` options are available for all distributions. + If the selected `method` is not available, a ``NotImplementedError`` + will be raised. + + Returns + ------- + out : array + The log-PMF evaluated at the argument `x`. + + See Also + -------- + pmf + logcdf + + Notes + ----- + Suppose a discrete probability distribution has support over the integers + :math:`{l, l+1, ..., r-1, r}`. + By definition of the support, the log-PMF evaluates to its minimum value + of :math:`-\infty` (i.e. :math:`\log(0)`) for non-integral :math:`x` and + for :math:`x` outside the support; i.e. for :math:`x < l` or :math:`x > r`. + + For distributions with infinite support, it is common for `pmf` to return + a value of ``0`` when the argument is theoretically within the support; + this can occur because the true value of the PMF is too small to be + represented by the chosen dtype. The log-PMF, however, will often be finite + (not ``-inf``) over a much larger domain. Consequently, it may be preferred + to work with the logarithms of probabilities and probability densities to + avoid underflow. + + References + ---------- + .. [1] Probability density function, *Wikipedia*, + https://en.wikipedia.org/wiki/Probability_density_function + + Examples + -------- + Instantiate a distribution with the desired parameters: + + >>> import numpy as np + >>> from scipy import stats + >>> X = stats.Binomial(n=10, p=0.5) + + Evaluate the log-PMF at the desired argument: + + >>> X.logpmf(5) + np.float64(-1.4020427180880297) + >>> np.allclose(X.logpmf(5), np.log(X.pmf(5))) + True + + """ + raise NotImplementedError() + + @abstractmethod + def cdf(self, x, y, /, *, method): + r"""Cumulative distribution function + + The cumulative distribution function ("CDF"), denoted :math:`F(x)`, is + the probability the random variable :math:`X` will assume a value + less than or equal to :math:`x`: + + .. math:: + + F(x) = P(X ≤ x) + + A two-argument variant of this function is also defined as the + probability the random variable :math:`X` will assume a value between + :math:`x` and :math:`y`. + + .. math:: + + F(x, y) = P(x ≤ X ≤ y) + + `cdf` accepts `x` for :math:`x` and `y` for :math:`y`. + + Parameters + ---------- + x, y : array_like + The arguments of the CDF. `x` is required; `y` is optional. + method : {None, 'formula', 'logexp', 'complement', 'quadrature', 'subtraction'} + The strategy used to evaluate the CDF. + By default (``None``), the one-argument form of the function + chooses between the following options, listed in order of precedence. + + - ``'formula'``: use a formula for the CDF itself + - ``'logexp'``: evaluate the log-CDF and exponentiate + - ``'complement'``: evaluate the CCDF and take the complement + - ``'quadrature'``: numerically integrate the PDF (or, in the discrete + case, sum the PMF) + + In place of ``'complement'``, the two-argument form accepts: + + - ``'subtraction'``: compute the CDF at each argument and take + the difference. + + Not all `method` options are available for all distributions. + If the selected `method` is not available, a ``NotImplementedError`` + will be raised. + + Returns + ------- + out : array + The CDF evaluated at the provided argument(s). + + See Also + -------- + logcdf + ccdf + + Notes + ----- + Suppose a continuous probability distribution has support :math:`[l, r]`. + The CDF :math:`F(x)` is related to the probability density function + :math:`f(x)` by: + + .. math:: + + F(x) = \int_l^x f(u) du + + The two argument version is: + + .. math:: + + F(x, y) = \int_x^y f(u) du = F(y) - F(x) + + The CDF evaluates to its minimum value of :math:`0` for :math:`x ≤ l` + and its maximum value of :math:`1` for :math:`x ≥ r`. + + Suppose a discrete probability distribution has support :math:`[l, r]`. + The CDF :math:`F(x)` is related to the probability mass function + :math:`f(x)` by: + + .. math:: + + F(x) = \sum_{u=l}^{\lfloor x \rfloor} f(u) + + The CDF evaluates to its minimum value of :math:`0` for :math:`x < l` + and its maximum value of :math:`1` for :math:`x ≥ r`. + + The CDF is also known simply as the "distribution function". + + References + ---------- + .. [1] Cumulative distribution function, *Wikipedia*, + https://en.wikipedia.org/wiki/Cumulative_distribution_function + + Examples + -------- + Instantiate a distribution with the desired parameters: + + >>> from scipy import stats + >>> X = stats.Uniform(a=-0.5, b=0.5) + + Evaluate the CDF at the desired argument: + + >>> X.cdf(0.25) + 0.75 + + Evaluate the cumulative probability between two arguments: + + >>> X.cdf(-0.25, 0.25) == X.cdf(0.25) - X.cdf(-0.25) + True + + """ # noqa: E501 + raise NotImplementedError() + + @abstractmethod + def icdf(self, p, /, *, method): + r"""Inverse of the cumulative distribution function. + + For monotonic continuous distributions, the inverse of the cumulative + distribution function ("inverse CDF"), denoted :math:`F^{-1}(p)`, is the + argument :math:`x` for which the cumulative distribution function + :math:`F(x)` evaluates to :math:`p`. + + .. math:: + + F^{-1}(p) = x \quad \text{s.t.} \quad F(x) = p + + When a strict "inverse" of the cumulative distribution function does not + exist (e.g. discrete random variables), the "inverse CDF" is defined by + convention as the smallest value within the support :math:`\chi` for which + :math:`F(x)` is at least :math:`p`. + + .. math:: + + F^{-1}(p) = \min_\chi \quad \text{s.t.} \quad F(x) ≥ p + + `icdf` accepts `p` for :math:`p \in [0, 1]`. + + Parameters + ---------- + p : array_like + The argument of the inverse CDF. + method : {None, 'formula', 'complement', 'inversion'} + The strategy used to evaluate the inverse CDF. + By default (``None``), the infrastructure chooses between the + following options, listed in order of precedence. + + - ``'formula'``: use a formula for the inverse CDF itself + - ``'complement'``: evaluate the inverse CCDF at the + complement of `p` + - ``'inversion'``: solve numerically for the argument at which the + CDF is equal to `p` + + Not all `method` options are available for all distributions. + If the selected `method` is not available, a ``NotImplementedError`` + will be raised. + + Returns + ------- + out : array + The inverse CDF evaluated at the provided argument. + + See Also + -------- + cdf + ilogcdf + + Notes + ----- + Suppose a probability distribution has support :math:`[l, r]`. The + inverse CDF returns its minimum value of :math:`l` at :math:`p = 0` + and its maximum value of :math:`r` at :math:`p = 1`. Because the CDF + has range :math:`[0, 1]`, the inverse CDF is only defined on the + domain :math:`[0, 1]`; for :math:`p < 0` and :math:`p > 1`, `icdf` + returns ``nan``. + + The inverse CDF is also known as the quantile function, percentile function, + and percent-point function. + + References + ---------- + .. [1] Quantile function, *Wikipedia*, + https://en.wikipedia.org/wiki/Quantile_function + + Examples + -------- + Instantiate a distribution with the desired parameters: + + >>> import numpy as np + >>> from scipy import stats + >>> X = stats.Uniform(a=-0.5, b=0.5) + + Evaluate the inverse CDF at the desired argument: + + >>> X.icdf(0.25) + -0.25 + >>> np.allclose(X.cdf(X.icdf(0.25)), 0.25) + True + + This function returns NaN when the argument is outside the domain. + + >>> X.icdf([-0.1, 0, 1, 1.1]) + array([ nan, -0.5, 0.5, nan]) + + """ + raise NotImplementedError() + + @abstractmethod + def ccdf(self, x, y, /, *, method): + r"""Complementary cumulative distribution function + + The complementary cumulative distribution function ("CCDF"), denoted + :math:`G(x)`, is the complement of the cumulative distribution function + :math:`F(x)`; i.e., probability the random variable :math:`X` will + assume a value greater than :math:`x`: + + .. math:: + + G(x) = 1 - F(x) = P(X > x) + + A two-argument variant of this function is: + + .. math:: + + G(x, y) = 1 - F(x, y) = P(X < x \text{ or } X > y) + + `ccdf` accepts `x` for :math:`x` and `y` for :math:`y`. + + Parameters + ---------- + x, y : array_like + The arguments of the CCDF. `x` is required; `y` is optional. + method : {None, 'formula', 'logexp', 'complement', 'quadrature', 'addition'} + The strategy used to evaluate the CCDF. + By default (``None``), the infrastructure chooses between the + following options, listed in order of precedence. + + - ``'formula'``: use a formula for the CCDF itself + - ``'logexp'``: evaluate the log-CCDF and exponentiate + - ``'complement'``: evaluate the CDF and take the complement + - ``'quadrature'``: numerically integrate the PDF (or, in the discrete + case, sum the PMF) + + The two-argument form chooses between: + + - ``'formula'``: use a formula for the CCDF itself + - ``'addition'``: compute the CDF at `x` and the CCDF at `y`, then add + + Not all `method` options are available for all distributions. + If the selected `method` is not available, a ``NotImplementedError`` + will be raised. + + Returns + ------- + out : array + The CCDF evaluated at the provided argument(s). + + See Also + -------- + cdf + logccdf + + Notes + ----- + Suppose a continuous probability distribution has support :math:`[l, r]`. + The CCDF :math:`G(x)` is related to the probability density function + :math:`f(x)` by: + + .. math:: + + G(x) = \int_x^r f(u) du + + The two argument version is: + + .. math:: + + G(x, y) = \int_l^x f(u) du + \int_y^r f(u) du + + The CCDF returns its minimum value of :math:`0` for :math:`x ≥ r` + and its maximum value of :math:`1` for :math:`x ≤ l`. + + Suppose a discrete probability distribution has support :math:`[l, r]`. + The CCDF :math:`G(x)` is related to the probability mass function + :math:`f(x)` by: + + .. math:: + + G(x) = \sum_{u=\lfloor x + 1 \rfloor}^{r} f(u) + + The CCDF evaluates to its minimum value of :math:`0` for :math:`x ≥ r` + and its maximum value of :math:`1` for :math:`x < l`. + + The CCDF is also known as the "survival function". + + References + ---------- + .. [1] Cumulative distribution function, *Wikipedia*, + https://en.wikipedia.org/wiki/Cumulative_distribution_function#Derived_functions + + Examples + -------- + Instantiate a distribution with the desired parameters: + + >>> import numpy as np + >>> from scipy import stats + >>> X = stats.Uniform(a=-0.5, b=0.5) + + Evaluate the CCDF at the desired argument: + + >>> X.ccdf(0.25) + 0.25 + >>> np.allclose(X.ccdf(0.25), 1-X.cdf(0.25)) + True + + Evaluate the complement of the cumulative probability between two arguments: + + >>> X.ccdf(-0.25, 0.25) == X.cdf(-0.25) + X.ccdf(0.25) + True + + """ # noqa: E501 + raise NotImplementedError() + + @abstractmethod + def iccdf(self, p, /, *, method): + r"""Inverse complementary cumulative distribution function. + + The inverse complementary cumulative distribution function ("inverse CCDF"), + denoted :math:`G^{-1}(p)`, is the argument :math:`x` for which the + complementary cumulative distribution function :math:`G(x)` evaluates to + :math:`p`. + + .. math:: + + G^{-1}(p) = x \quad \text{s.t.} \quad G(x) = p + + When a strict "inverse" of the complementary cumulative distribution function + does not exist (e.g. discrete random variables), the "inverse CCDF" is defined + by convention as the smallest value within the support :math:`\chi` for which + :math:`G(x)` is no greater than :math:`p`. + + .. math:: + + G^{-1}(p) = \min_\chi \quad \text{s.t.} \quad G(x) ≤ p + + `iccdf` accepts `p` for :math:`p \in [0, 1]`. + + Parameters + ---------- + p : array_like + The argument of the inverse CCDF. + method : {None, 'formula', 'complement', 'inversion'} + The strategy used to evaluate the inverse CCDF. + By default (``None``), the infrastructure chooses between the + following options, listed in order of precedence. + + - ``'formula'``: use a formula for the inverse CCDF itself + - ``'complement'``: evaluate the inverse CDF at the + complement of `p` + - ``'inversion'``: solve numerically for the argument at which the + CCDF is equal to `p` + + Not all `method` options are available for all distributions. + If the selected `method` is not available, a ``NotImplementedError`` + will be raised. + + Returns + ------- + out : array + The inverse CCDF evaluated at the provided argument. + + Notes + ----- + Suppose a probability distribution has support :math:`[l, r]`. The + inverse CCDF returns its minimum value of :math:`l` at :math:`p = 1` + and its maximum value of :math:`r` at :math:`p = 0`. Because the CCDF + has range :math:`[0, 1]`, the inverse CCDF is only defined on the + domain :math:`[0, 1]`; for :math:`p < 0` and :math:`p > 1`, ``iccdf`` + returns ``nan``. + + See Also + -------- + icdf + ilogccdf + + Examples + -------- + Instantiate a distribution with the desired parameters: + + >>> import numpy as np + >>> from scipy import stats + >>> X = stats.Uniform(a=-0.5, b=0.5) + + Evaluate the inverse CCDF at the desired argument: + + >>> X.iccdf(0.25) + 0.25 + >>> np.allclose(X.iccdf(0.25), X.icdf(1-0.25)) + True + + This function returns NaN when the argument is outside the domain. + + >>> X.iccdf([-0.1, 0, 1, 1.1]) + array([ nan, 0.5, -0.5, nan]) + + """ + raise NotImplementedError() + + @abstractmethod + def logcdf(self, x, y, /, *, method): + r"""Log of the cumulative distribution function + + The cumulative distribution function ("CDF"), denoted :math:`F(x)`, is + the probability the random variable :math:`X` will assume a value + less than or equal to :math:`x`: + + .. math:: + + F(x) = P(X ≤ x) + + A two-argument variant of this function is also defined as the + probability the random variable :math:`X` will assume a value between + :math:`x` and :math:`y`. + + .. math:: + + F(x, y) = P(x ≤ X ≤ y) + + `logcdf` computes the logarithm of the cumulative distribution function + ("log-CDF"), :math:`\log(F(x))`/:math:`\log(F(x, y))`, but it may be + numerically favorable compared to the naive implementation (computing + the CDF and taking the logarithm). + + `logcdf` accepts `x` for :math:`x` and `y` for :math:`y`. + + Parameters + ---------- + x, y : array_like + The arguments of the log-CDF. `x` is required; `y` is optional. + method : {None, 'formula', 'logexp', 'complement', 'quadrature', 'subtraction'} + The strategy used to evaluate the log-CDF. + By default (``None``), the one-argument form of the function + chooses between the following options, listed in order of precedence. + + - ``'formula'``: use a formula for the log-CDF itself + - ``'logexp'``: evaluate the CDF and take the logarithm + - ``'complement'``: evaluate the log-CCDF and take the + logarithmic complement (see Notes) + - ``'quadrature'``: numerically log-integrate the log-PDF (or, in the + discrete case, log-sum the log-PMF) + + In place of ``'complement'``, the two-argument form accepts: + + - ``'subtraction'``: compute the log-CDF at each argument and take + the logarithmic difference (see Notes) + + Not all `method` options are available for all distributions. + If the selected `method` is not available, a ``NotImplementedError`` + will be raised. + + Returns + ------- + out : array + The log-CDF evaluated at the provided argument(s). + + See Also + -------- + cdf + logccdf + + Notes + ----- + Suppose a continuous probability distribution has support :math:`[l, r]`. + The log-CDF evaluates to its minimum value of :math:`\log(0) = -\infty` + for :math:`x ≤ l` and its maximum value of :math:`\log(1) = 0` for + :math:`x ≥ r`. An analogous statement can be made for discrete distributions, + but the inequality governing the minimum value is strict. + + For distributions with infinite support, it is common for + `cdf` to return a value of ``0`` when the argument + is theoretically within the support; this can occur because the true value + of the CDF is too small to be represented by the chosen dtype. `logcdf`, + however, will often return a finite (not ``-inf``) result over a much larger + domain. Similarly, `logcdf` may provided a strictly negative result with + arguments for which `cdf` would return ``1.0``. Consequently, it may be + preferred to work with the logarithms of probabilities to avoid underflow + and related limitations of floating point numbers. + + The "logarithmic complement" of a number :math:`z` is mathematically + equivalent to :math:`\log(1-\exp(z))`, but it is computed to avoid loss + of precision when :math:`\exp(z)` is nearly :math:`0` or :math:`1`. + Similarly, the term "logarithmic difference" of :math:`w` and :math:`z` + is used here to mean :math:`\log(\exp(w)-\exp(z))`. + + If ``y < x``, the CDF is negative, and therefore the log-CCDF + is complex with imaginary part :math:`\pi`. For + consistency, the result of this function always has complex dtype + when `y` is provided, regardless of the value of the imaginary part. + + References + ---------- + .. [1] Cumulative distribution function, *Wikipedia*, + https://en.wikipedia.org/wiki/Cumulative_distribution_function + + Examples + -------- + Instantiate a distribution with the desired parameters: + + >>> import numpy as np + >>> from scipy import stats + >>> X = stats.Uniform(a=-0.5, b=0.5) + + Evaluate the log-CDF at the desired argument: + + >>> X.logcdf(0.25) + -0.287682072451781 + >>> np.allclose(X.logcdf(0.), np.log(X.cdf(0.))) + True + + """ # noqa: E501 + raise NotImplementedError() + + @abstractmethod + def ilogcdf(self, logp, /, *, method): + r"""Inverse of the logarithm of the cumulative distribution function. + + The inverse of the logarithm of the cumulative distribution function + ("inverse log-CDF") is the argument :math:`x` for which the logarithm + of the cumulative distribution function :math:`\log(F(x))` evaluates + to :math:`\log(p)`. + + Mathematically, it is equivalent to :math:`F^{-1}(\exp(y))`, where + :math:`y = \log(p)`, but it may be numerically favorable compared to + the naive implementation (computing :math:`p = \exp(y)`, then + :math:`F^{-1}(p)`). + + `ilogcdf` accepts `logp` for :math:`\log(p) ≤ 0`. + + Parameters + ---------- + logp : array_like + The argument of the inverse log-CDF. + method : {None, 'formula', 'complement', 'inversion'} + The strategy used to evaluate the inverse log-CDF. + By default (``None``), the infrastructure chooses between the + following options, listed in order of precedence. + + - ``'formula'``: use a formula for the inverse log-CDF itself + - ``'complement'``: evaluate the inverse log-CCDF at the + logarithmic complement of `logp` (see Notes) + - ``'inversion'``: solve numerically for the argument at which the + log-CDF is equal to `logp` + + Not all `method` options are available for all distributions. + If the selected `method` is not available, a ``NotImplementedError`` + will be raised. + + Returns + ------- + out : array + The inverse log-CDF evaluated at the provided argument. + + See Also + -------- + icdf + logcdf + + Notes + ----- + Suppose a probability distribution has support :math:`[l, r]`. + The inverse log-CDF returns its minimum value of :math:`l` at + :math:`\log(p) = \log(0) = -\infty` and its maximum value of :math:`r` at + :math:`\log(p) = \log(1) = 0`. Because the log-CDF has range + :math:`[-\infty, 0]`, the inverse log-CDF is only defined on the + negative reals; for :math:`\log(p) > 0`, `ilogcdf` returns ``nan``. + + Occasionally, it is needed to find the argument of the CDF for which + the resulting probability is very close to ``0`` or ``1`` - too close to + represent accurately with floating point arithmetic. In many cases, + however, the *logarithm* of this resulting probability may be + represented in floating point arithmetic, in which case this function + may be used to find the argument of the CDF for which the *logarithm* + of the resulting probability is :math:`y = \log(p)`. + + The "logarithmic complement" of a number :math:`z` is mathematically + equivalent to :math:`\log(1-\exp(z))`, but it is computed to avoid loss + of precision when :math:`\exp(z)` is nearly :math:`0` or :math:`1`. + + Examples + -------- + Instantiate a distribution with the desired parameters: + + >>> import numpy as np + >>> from scipy import stats + >>> X = stats.Uniform(a=-0.5, b=0.5) + + Evaluate the inverse log-CDF at the desired argument: + + >>> X.ilogcdf(-0.25) + 0.2788007830714034 + >>> np.allclose(X.ilogcdf(-0.25), X.icdf(np.exp(-0.25))) + True + + """ + raise NotImplementedError() + + @abstractmethod + def logccdf(self, x, y, /, *, method): + r"""Log of the complementary cumulative distribution function + + The complementary cumulative distribution function ("CCDF"), denoted + :math:`G(x)` is the complement of the cumulative distribution function + :math:`F(x)`; i.e., probability the random variable :math:`X` will + assume a value greater than :math:`x`: + + .. math:: + + G(x) = 1 - F(x) = P(X > x) + + A two-argument variant of this function is: + + .. math:: + + G(x, y) = 1 - F(x, y) = P(X < x \quad \text{or} \quad X > y) + + `logccdf` computes the logarithm of the complementary cumulative + distribution function ("log-CCDF"), :math:`\log(G(x))`/:math:`\log(G(x, y))`, + but it may be numerically favorable compared to the naive implementation + (computing the CDF and taking the logarithm). + + `logccdf` accepts `x` for :math:`x` and `y` for :math:`y`. + + Parameters + ---------- + x, y : array_like + The arguments of the log-CCDF. `x` is required; `y` is optional. + method : {None, 'formula', 'logexp', 'complement', 'quadrature', 'addition'} + The strategy used to evaluate the log-CCDF. + By default (``None``), the one-argument form of the function + chooses between the following options, listed in order of precedence. + + - ``'formula'``: use a formula for the log CCDF itself + - ``'logexp'``: evaluate the CCDF and take the logarithm + - ``'complement'``: evaluate the log-CDF and take the + logarithmic complement (see Notes) + - ``'quadrature'``: numerically log-integrate the log-PDF (or, in the + discrete case, log-sum the log-PMF) + + The two-argument form chooses between: + + - ``'formula'``: use a formula for the log CCDF itself + - ``'addition'``: compute the log-CDF at `x` and the log-CCDF at `y`, + then take the logarithmic sum (see Notes) + + Not all `method` options are available for all distributions. + If the selected `method` is not available, a ``NotImplementedError`` + will be raised. + + Returns + ------- + out : array + The log-CCDF evaluated at the provided argument(s). + + See Also + -------- + ccdf + logcdf + + Notes + ----- + Suppose a continuous probability distribution has support :math:`[l, r]`. + The log-CCDF returns its minimum value of :math:`\log(0)=-\infty` for + :math:`x ≥ r` and its maximum value of :math:`\log(1) = 0` for + :math:`x ≤ l`. An analogous statement can be made for discrete distributions, + but the inequality governing the maximum value is strict. + + For distributions with infinite support, it is common for + `ccdf` to return a value of ``0`` when the argument + is theoretically within the support; this can occur because the true value + of the CCDF is too small to be represented by the chosen dtype. The log + of the CCDF, however, will often be finite (not ``-inf``) over a much larger + domain. Similarly, `logccdf` may provided a strictly negative result with + arguments for which `ccdf` would return ``1.0``. Consequently, it may be + preferred to work with the logarithms of probabilities to avoid underflow + and related limitations of floating point numbers. + + The "logarithmic complement" of a number :math:`z` is mathematically + equivalent to :math:`\log(1-\exp(z))`, but it is computed to avoid loss + of precision when :math:`\exp(z)` is nearly :math:`0` or :math:`1`. + Similarly, the term "logarithmic sum" of :math:`w` and :math:`z` + is used here to mean the :math:`\log(\exp(w)+\exp(z))`, AKA + :math:`\text{LogSumExp}(w, z)`. + + References + ---------- + .. [1] Cumulative distribution function, *Wikipedia*, + https://en.wikipedia.org/wiki/Cumulative_distribution_function#Derived_functions + + Examples + -------- + Instantiate a distribution with the desired parameters: + + >>> import numpy as np + >>> from scipy import stats + >>> X = stats.Uniform(a=-0.5, b=0.5) + + Evaluate the log-CCDF at the desired argument: + + >>> X.logccdf(0.25) + -1.3862943611198906 + >>> np.allclose(X.logccdf(0.), np.log(X.ccdf(0.))) + True + + """ # noqa: E501 + raise NotImplementedError() + + @abstractmethod + def ilogccdf(self, logp, /, *, method): + r"""Inverse of the log of the complementary cumulative distribution function. + + The inverse of the logarithm of the complementary cumulative distribution + function ("inverse log-CCDF") is the argument :math:`x` for which the logarithm + of the complementary cumulative distribution function :math:`\log(G(x))` + evaluates to :math:`\log(p)`. + + Mathematically, it is equivalent to :math:`G^{-1}(\exp(y))`, where + :math:`y = \log(p)`, but it may be numerically favorable compared to the naive + implementation (computing :math:`p = \exp(y)`, then :math:`G^{-1}(p)`). + + `ilogccdf` accepts `logp` for :math:`\log(p) ≤ 0`. + + Parameters + ---------- + x : array_like + The argument of the inverse log-CCDF. + method : {None, 'formula', 'complement', 'inversion'} + The strategy used to evaluate the inverse log-CCDF. + By default (``None``), the infrastructure chooses between the + following options, listed in order of precedence. + + - ``'formula'``: use a formula for the inverse log-CCDF itself + - ``'complement'``: evaluate the inverse log-CDF at the + logarithmic complement of `x` (see Notes) + - ``'inversion'``: solve numerically for the argument at which the + log-CCDF is equal to `x` + + Not all `method` options are available for all distributions. + If the selected `method` is not available, a ``NotImplementedError`` + will be raised. + + Returns + ------- + out : array + The inverse log-CCDF evaluated at the provided argument. + + Notes + ----- + Suppose a probability distribution has support :math:`[l, r]`. The + inverse log-CCDF returns its minimum value of :math:`l` at + :math:`\log(p) = \log(1) = 0` and its maximum value of :math:`r` at + :math:`\log(p) = \log(0) = -\infty`. Because the log-CCDF has range + :math:`[-\infty, 0]`, the inverse log-CDF is only defined on the + negative reals; for :math:`\log(p) > 0`, `ilogccdf` returns ``nan``. + + Occasionally, it is needed to find the argument of the CCDF for which + the resulting probability is very close to ``0`` or ``1`` - too close to + represent accurately with floating point arithmetic. In many cases, + however, the *logarithm* of this resulting probability may be + represented in floating point arithmetic, in which case this function + may be used to find the argument of the CCDF for which the *logarithm* + of the resulting probability is :math:`y = \log(p)`. + + The "logarithmic complement" of a number :math:`z` is mathematically + equivalent to :math:`\log(1-\exp(z))`, but it is computed to avoid loss + of precision when :math:`\exp(z)` is nearly :math:`0` or :math:`1`. + + See Also + -------- + iccdf + ilogccdf + + Examples + -------- + Instantiate a distribution with the desired parameters: + + >>> import numpy as np + >>> from scipy import stats + >>> X = stats.Uniform(a=-0.5, b=0.5) + + Evaluate the inverse log-CCDF at the desired argument: + + >>> X.ilogccdf(-0.25) + -0.2788007830714034 + >>> np.allclose(X.ilogccdf(-0.25), X.iccdf(np.exp(-0.25))) + True + + """ + raise NotImplementedError() + + @abstractmethod + def logentropy(self, *, method): + r"""Logarithm of the differential entropy + + In terms of probability density function :math:`f(x)` and support + :math:`\chi`, the differential entropy (or simply "entropy") of a + continuous random variable :math:`X` is: + + .. math:: + + h(X) = - \int_{\chi} f(x) \log f(x) dx + + The definition for a discrete random variable is analogous, with the PMF + replacing the PDF and a sum over the support replacing the integral. + + `logentropy` computes the logarithm of the differential entropy + ("log-entropy"), :math:`\log(h(X))`, but it may be numerically favorable + compared to the naive implementation (computing :math:`h(X)` then + taking the logarithm). + + Parameters + ---------- + method : {None, 'formula', 'logexp', 'quadrature} + The strategy used to evaluate the log-entropy. By default + (``None``), the infrastructure chooses between the following options, + listed in order of precedence. + + - ``'formula'``: use a formula for the log-entropy itself + - ``'logexp'``: evaluate the entropy and take the logarithm + - ``'quadrature'``: numerically log-integrate (or, in the discrete + case, log-sum) the logarithm of the entropy integrand (summand) + + Not all `method` options are available for all distributions. + If the selected `method` is not available, a ``NotImplementedError`` + will be raised. + + Returns + ------- + out : array + The log-entropy. + + See Also + -------- + entropy + logpdf + + Notes + ----- + The differential entropy of a continuous distribution can be negative. + In this case, the log-entropy is complex with imaginary part :math:`\pi`. + For consistency, the result of this function always has complex dtype, + regardless of the value of the imaginary part. + + References + ---------- + .. [1] Differential entropy, *Wikipedia*, + https://en.wikipedia.org/wiki/Differential_entropy + + Examples + -------- + Instantiate a distribution with the desired parameters: + + >>> import numpy as np + >>> from scipy import stats + >>> X = stats.Uniform(a=-1., b=1.) + + Evaluate the log-entropy: + + >>> X.logentropy() + (-0.3665129205816642+0j) + >>> np.allclose(np.exp(X.logentropy()), X.entropy()) + True + + For a random variable with negative entropy, the log-entropy has an + imaginary part equal to `np.pi`. + + >>> X = stats.Uniform(a=-.1, b=.1) + >>> X.entropy(), X.logentropy() + (-1.6094379124341007, (0.4758849953271105+3.141592653589793j)) + + """ + raise NotImplementedError() + + @abstractmethod + def entropy(self, *, method): + r"""Differential entropy + + In terms of probability density function :math:`f(x)` and support + :math:`\chi`, the differential entropy (or simply "entropy") of a + continuous random variable :math:`X` is: + + .. math:: + + h(X) = - \int_{\chi} f(x) \log f(x) dx + + The definition for a discrete random variable is analogous, with the + PMF replacing the PDF and a sum over the support replacing the integral. + + Parameters + ---------- + method : {None, 'formula', 'logexp', 'quadrature'} + The strategy used to evaluate the entropy. By default (``None``), + the infrastructure chooses between the following options, listed + in order of precedence. + + - ``'formula'``: use a formula for the entropy itself + - ``'logexp'``: evaluate the log-entropy and exponentiate + - ``'quadrature'``: numerically integrate (or, in the discrete + case, sum) the entropy integrand (summand) + + + Not all `method` options are available for all distributions. + If the selected `method` is not available, a ``NotImplementedError`` + will be raised. + + Returns + ------- + out : array + The entropy of the random variable. + + See Also + -------- + logentropy + pdf + + Notes + ----- + This function calculates the entropy using the natural logarithm; i.e. + the logarithm with base :math:`e`. Consequently, the value is expressed + in (dimensionless) "units" of nats. To convert the entropy to different + units (i.e. corresponding with a different base), divide the result by + the natural logarithm of the desired base. + + References + ---------- + .. [1] Differential entropy, *Wikipedia*, + https://en.wikipedia.org/wiki/Differential_entropy + + Examples + -------- + Instantiate a distribution with the desired parameters: + + >>> from scipy import stats + >>> X = stats.Uniform(a=-1., b=1.) + + Evaluate the entropy: + + >>> X.entropy() + 0.6931471805599454 + + """ + raise NotImplementedError() diff --git a/v12x/pylib/scipy/stats/_qmc.py b/v12x/pylib/scipy/stats/_qmc.py new file mode 100644 index 0000000000000000000000000000000000000000..3358bc68222035ddc02a5f7977041bd313f3d0a9 --- /dev/null +++ b/v12x/pylib/scipy/stats/_qmc.py @@ -0,0 +1,2951 @@ +"""Quasi-Monte Carlo engines and helpers.""" +import copy +import math +import numbers +import os +import warnings +from abc import ABC, abstractmethod +from functools import partial +from typing import ( + ClassVar, + Literal, + overload, + TYPE_CHECKING, +) +from collections.abc import Callable + +import numpy as np + +from scipy._lib._util import DecimalNumber, GeneratorType, IntNumber, SeedType + +if TYPE_CHECKING: + import numpy.typing as npt + +import scipy.stats as stats +from scipy._lib._util import rng_integers, _rng_spawn, _transition_to_rng +from scipy.sparse.csgraph import minimum_spanning_tree +from scipy.spatial import distance, Voronoi +from scipy.special import gammainc +from ._sobol import ( + _initialize_v, _cscramble, _fill_p_cumulative, _draw, _fast_forward, + _categorize, _MAXDIM +) +from ._qmc_cy import ( + _cy_wrapper_centered_discrepancy, + _cy_wrapper_wrap_around_discrepancy, + _cy_wrapper_mixture_discrepancy, + _cy_wrapper_l2_star_discrepancy, + _cy_wrapper_update_discrepancy, + _cy_van_der_corput_scrambled, + _cy_van_der_corput, +) + + +__all__ = ['scale', 'discrepancy', 'geometric_discrepancy', 'update_discrepancy', + 'QMCEngine', 'Sobol', 'Halton', 'LatinHypercube', 'PoissonDisk', + 'MultinomialQMC', 'MultivariateNormalQMC'] + + +@overload +def check_random_state(seed: IntNumber | None = ...) -> np.random.Generator: + ... + + +@overload +def check_random_state(seed: GeneratorType) -> GeneratorType: + ... + + +# Based on scipy._lib._util.check_random_state +# This is going to be removed at the end of the SPEC 7 transition, +# so I'll just leave the argument name `seed` alone +def check_random_state(seed=None): + """Turn `seed` into a `numpy.random.Generator` instance. + + Parameters + ---------- + seed : {None, int, `numpy.random.Generator`, `numpy.random.RandomState`}, optional + If `seed` is an int or None, a new `numpy.random.Generator` is + created using ``np.random.default_rng(seed)``. + If `seed` is already a ``Generator`` or ``RandomState`` instance, then + the provided instance is used. + + Returns + ------- + seed : {`numpy.random.Generator`, `numpy.random.RandomState`} + Random number generator. + + """ + if seed is None or isinstance(seed, numbers.Integral | np.integer): + return np.random.default_rng(seed) + elif isinstance(seed, np.random.RandomState | np.random.Generator): + return seed + else: + raise ValueError(f'{seed!r} cannot be used to seed a' + ' numpy.random.Generator instance') + + +def scale( + sample: "npt.ArrayLike", + l_bounds: "npt.ArrayLike", + u_bounds: "npt.ArrayLike", + *, + reverse: bool = False +) -> np.ndarray: + r"""Sample scaling from unit hypercube to different bounds. + + To convert a sample from :math:`[0, 1)` to :math:`[a, b), b>a`, + with :math:`a` the lower bounds and :math:`b` the upper bounds. + The following transformation is used: + + .. math:: + + (b - a) \cdot \text{sample} + a + + Parameters + ---------- + sample : array_like (n, d) + Sample to scale. + l_bounds, u_bounds : array_like (d,) + Lower and upper bounds (resp. :math:`a`, :math:`b`) of transformed + data. If `reverse` is True, range of the original data to transform + to the unit hypercube. + reverse : bool, optional + Reverse the transformation from different bounds to the unit hypercube. + Default is False. + + Returns + ------- + sample : array_like (n, d) + Scaled sample. + + Examples + -------- + Transform 3 samples in the unit hypercube to bounds: + + >>> from scipy.stats import qmc + >>> l_bounds = [-2, 0] + >>> u_bounds = [6, 5] + >>> sample = [[0.5 , 0.75], + ... [0.5 , 0.5], + ... [0.75, 0.25]] + >>> sample_scaled = qmc.scale(sample, l_bounds, u_bounds) + >>> sample_scaled + array([[2. , 3.75], + [2. , 2.5 ], + [4. , 1.25]]) + + And convert back to the unit hypercube: + + >>> sample_ = qmc.scale(sample_scaled, l_bounds, u_bounds, reverse=True) + >>> sample_ + array([[0.5 , 0.75], + [0.5 , 0.5 ], + [0.75, 0.25]]) + + """ + sample = np.asarray(sample) + + # Checking bounds and sample + if not sample.ndim == 2: + raise ValueError('Sample is not a 2D array') + + lower, upper = _validate_bounds( + l_bounds=l_bounds, u_bounds=u_bounds, d=sample.shape[1] + ) + + if not reverse: + # Checking that sample is within the hypercube + if (sample.max() > 1.) or (sample.min() < 0.): + raise ValueError('Sample is not in unit hypercube') + + return sample * (upper - lower) + lower + else: + # Checking that sample is within the bounds + if not (np.all(sample >= lower) and np.all(sample <= upper)): + raise ValueError('Sample is out of bounds') + + return (sample - lower) / (upper - lower) + + +def _ensure_in_unit_hypercube(sample: "npt.ArrayLike") -> np.ndarray: + """Ensure that sample is a 2D array and is within a unit hypercube + + Parameters + ---------- + sample : array_like (n, d) + A 2D array of points. + + Returns + ------- + np.ndarray + The array interpretation of the input sample + + Raises + ------ + ValueError + If the input is not a 2D array or contains points outside of + a unit hypercube. + """ + sample = np.asarray(sample, dtype=np.float64, order="C") + + if not sample.ndim == 2: + raise ValueError("Sample is not a 2D array") + + if (sample.max() > 1.) or (sample.min() < 0.): + raise ValueError("Sample is not in unit hypercube") + + return sample + + +def discrepancy( + sample: "npt.ArrayLike", + *, + iterative: bool = False, + method: Literal["CD", "WD", "MD", "L2-star"] = "CD", + workers: IntNumber = 1) -> float: + """Discrepancy of a given sample. + + Parameters + ---------- + sample : array_like (n, d) + The sample to compute the discrepancy from. + iterative : bool, optional + Must be False if not using it for updating the discrepancy. + Default is False. Refer to the notes for more details. + method : str, optional + Type of discrepancy, can be ``CD``, ``WD``, ``MD`` or ``L2-star``. + Refer to the notes for more details. Default is ``CD``. + workers : int, optional + Number of workers to use for parallel processing. If -1 is given all + CPU threads are used. Default is 1. + + Returns + ------- + discrepancy : float + Discrepancy. + + See Also + -------- + geometric_discrepancy + + Notes + ----- + The discrepancy is a uniformity criterion used to assess the space filling + of a number of samples in a hypercube. A discrepancy quantifies the + distance between the continuous uniform distribution on a hypercube and the + discrete uniform distribution on :math:`n` distinct sample points. + + The lower the value is, the better the coverage of the parameter space is. + + For a collection of subsets of the hypercube, the discrepancy is the + difference between the fraction of sample points in one of those + subsets and the volume of that subset. There are different definitions of + discrepancy corresponding to different collections of subsets. Some + versions take a root mean square difference over subsets instead of + a maximum. + + A measure of uniformity is reasonable if it satisfies the following + criteria [1]_: + + 1. It is invariant under permuting factors and/or runs. + 2. It is invariant under rotation of the coordinates. + 3. It can measure not only uniformity of the sample over the hypercube, + but also the projection uniformity of the sample over non-empty + subset of lower dimension hypercubes. + 4. There is some reasonable geometric meaning. + 5. It is easy to compute. + 6. It satisfies the Koksma-Hlawka-like inequality. + 7. It is consistent with other criteria in experimental design. + + Four methods are available: + + * ``CD``: Centered Discrepancy - subspace involves a corner of the + hypercube + * ``WD``: Wrap-around Discrepancy - subspace can wrap around bounds + * ``MD``: Mixture Discrepancy - mix between CD/WD covering more criteria + * ``L2-star``: L2-star discrepancy - like CD BUT variant to rotation + + Methods ``CD``, ``WD``, and ``MD`` implement the right hand side of equations + 9, 10, and 18 of [2]_, respectively; the square root is not taken. On the + other hand, ``L2-star`` computes the quantity given by equation 10 of + [3]_ as implemented by subsequent equations; the square root is taken. + + Lastly, using ``iterative=True``, it is possible to compute the + discrepancy as if we had :math:`n+1` samples. This is useful if we want + to add a point to a sampling and check the candidate which would give the + lowest discrepancy. Then you could just update the discrepancy with + each candidate using `update_discrepancy`. This method is faster than + computing the discrepancy for a large number of candidates. + + References + ---------- + .. [1] Fang et al. "Design and modeling for computer experiments". + Computer Science and Data Analysis Series, 2006. + .. [2] Zhou Y.-D. et al. "Mixture discrepancy for quasi-random point sets." + Journal of Complexity, 29 (3-4) , pp. 283-301, 2013. + .. [3] T. T. Warnock. "Computational investigations of low discrepancy + point sets." Applications of Number Theory to Numerical + Analysis, Academic Press, pp. 319-343, 1972. + + Examples + -------- + Calculate the quality of the sample using the discrepancy: + + >>> import numpy as np + >>> from scipy.stats import qmc + >>> space = np.array([[1, 3], [2, 6], [3, 2], [4, 5], [5, 1], [6, 4]]) + >>> l_bounds = [0.5, 0.5] + >>> u_bounds = [6.5, 6.5] + >>> space = qmc.scale(space, l_bounds, u_bounds, reverse=True) + >>> space + array([[0.08333333, 0.41666667], + [0.25 , 0.91666667], + [0.41666667, 0.25 ], + [0.58333333, 0.75 ], + [0.75 , 0.08333333], + [0.91666667, 0.58333333]]) + >>> qmc.discrepancy(space) + 0.008142039609053464 + + We can also compute iteratively the ``CD`` discrepancy by using + ``iterative=True``. + + >>> disc_init = qmc.discrepancy(space[:-1], iterative=True) + >>> disc_init + 0.04769081147119336 + >>> qmc.update_discrepancy(space[-1], space[:-1], disc_init) + 0.008142039609053513 + + """ + sample = _ensure_in_unit_hypercube(sample) + + workers = _validate_workers(workers) + + methods = { + "CD": _cy_wrapper_centered_discrepancy, + "WD": _cy_wrapper_wrap_around_discrepancy, + "MD": _cy_wrapper_mixture_discrepancy, + "L2-star": _cy_wrapper_l2_star_discrepancy, + } + + if method in methods: + return methods[method](sample, iterative, workers=workers) + else: + raise ValueError(f"{method!r} is not a valid method. It must be one of" + f" {set(methods)!r}") + + +def geometric_discrepancy( + sample: "npt.ArrayLike", + method: Literal["mindist", "mst"] = "mindist", + metric: str = "euclidean") -> float: + """Discrepancy of a given sample based on its geometric properties. + + Parameters + ---------- + sample : array_like (n, d) + The sample to compute the discrepancy from. + method : {"mindist", "mst"}, optional + The method to use. One of ``mindist`` for minimum distance (default) + or ``mst`` for minimum spanning tree. + metric : str or callable, optional + The distance metric to use. See the documentation + for `scipy.spatial.distance.pdist` for the available metrics and + the default. + + Returns + ------- + discrepancy : float + Discrepancy (higher values correspond to greater sample uniformity). + + See Also + -------- + discrepancy + + Notes + ----- + The discrepancy can serve as a simple measure of quality of a random sample. + This measure is based on the geometric properties of the distribution of points + in the sample, such as the minimum distance between any pair of points, or + the mean edge length in a minimum spanning tree. + + The higher the value is, the better the coverage of the parameter space is. + Note that this is different from `scipy.stats.qmc.discrepancy`, where lower + values correspond to higher quality of the sample. + + Also note that when comparing different sampling strategies using this function, + the sample size must be kept constant. + + It is possible to calculate two metrics from the minimum spanning tree: + the mean edge length and the standard deviation of edges lengths. Using + both metrics offers a better picture of uniformity than either metric alone, + with higher mean and lower standard deviation being preferable (see [1]_ + for a brief discussion). This function currently only calculates the mean + edge length. + + References + ---------- + .. [1] Franco J. et al. "Minimum Spanning Tree: A new approach to assess the quality + of the design of computer experiments." Chemometrics and Intelligent Laboratory + Systems, 97 (2), pp. 164-169, 2009. + + Examples + -------- + Calculate the quality of the sample using the minimum euclidean distance + (the defaults): + + >>> import numpy as np + >>> from scipy.stats import qmc + >>> rng = np.random.default_rng(191468432622931918890291693003068437394) + >>> sample = qmc.LatinHypercube(d=2, rng=rng).random(50) + >>> qmc.geometric_discrepancy(sample) + 0.03708161435687876 + + Calculate the quality using the mean edge length in the minimum + spanning tree: + + >>> qmc.geometric_discrepancy(sample, method='mst') + 0.1105149978798376 + + Display the minimum spanning tree and the points with + the smallest distance: + + >>> import matplotlib.pyplot as plt + >>> from matplotlib.lines import Line2D + >>> from scipy.sparse.csgraph import minimum_spanning_tree + >>> from scipy.spatial.distance import pdist, squareform + >>> dist = pdist(sample) + >>> mst = minimum_spanning_tree(squareform(dist)) + >>> edges = np.where(mst.toarray() > 0) + >>> edges = np.asarray(edges).T + >>> min_dist = np.min(dist) + >>> min_idx = np.argwhere(squareform(dist) == min_dist)[0] + >>> fig, ax = plt.subplots(figsize=(10, 5)) + >>> _ = ax.set(aspect='equal', xlabel=r'$x_1$', ylabel=r'$x_2$', + ... xlim=[0, 1], ylim=[0, 1]) + >>> for edge in edges: + ... ax.plot(sample[edge, 0], sample[edge, 1], c='k') + >>> ax.scatter(sample[:, 0], sample[:, 1]) + >>> ax.add_patch(plt.Circle(sample[min_idx[0]], min_dist, color='red', fill=False)) + >>> markers = [ + ... Line2D([0], [0], marker='o', lw=0, label='Sample points'), + ... Line2D([0], [0], color='k', label='Minimum spanning tree'), + ... Line2D([0], [0], marker='o', lw=0, markerfacecolor='w', markeredgecolor='r', + ... label='Minimum point-to-point distance'), + ... ] + >>> ax.legend(handles=markers, loc='center left', bbox_to_anchor=(1, 0.5)); + >>> plt.show() + + """ + sample = _ensure_in_unit_hypercube(sample) + if sample.shape[0] < 2: + raise ValueError("Sample must contain at least two points") + + distances = distance.pdist(sample, metric=metric) # type: ignore[call-overload] + + if np.any(distances == 0.0): + warnings.warn("Sample contains duplicate points.", stacklevel=2) + + if method == "mindist": + return np.min(distances[distances.nonzero()]) + elif method == "mst": + fully_connected_graph = distance.squareform(distances) + mst = minimum_spanning_tree(fully_connected_graph) + distances = mst[mst.nonzero()] + # TODO consider returning both the mean and the standard deviation + # see [1] for a discussion + return np.mean(distances) + else: + raise ValueError(f"{method!r} is not a valid method. " + f"It must be one of {{'mindist', 'mst'}}") + + +def update_discrepancy( + x_new: "npt.ArrayLike", + sample: "npt.ArrayLike", + initial_disc: DecimalNumber) -> float: + """Update the centered discrepancy with a new sample. + + Parameters + ---------- + x_new : array_like (1, d) + The new sample to add in `sample`. + sample : array_like (n, d) + The initial sample. + initial_disc : float + Centered discrepancy of the `sample`. + + Returns + ------- + discrepancy : float + Centered discrepancy of the sample composed of `x_new` and `sample`. + + Examples + -------- + We can also compute iteratively the discrepancy by using + ``iterative=True``. + + >>> import numpy as np + >>> from scipy.stats import qmc + >>> space = np.array([[1, 3], [2, 6], [3, 2], [4, 5], [5, 1], [6, 4]]) + >>> l_bounds = [0.5, 0.5] + >>> u_bounds = [6.5, 6.5] + >>> space = qmc.scale(space, l_bounds, u_bounds, reverse=True) + >>> disc_init = qmc.discrepancy(space[:-1], iterative=True) + >>> disc_init + 0.04769081147119336 + >>> qmc.update_discrepancy(space[-1], space[:-1], disc_init) + 0.008142039609053513 + + """ + sample = np.asarray(sample, dtype=np.float64, order="C") + x_new = np.asarray(x_new, dtype=np.float64, order="C") + + # Checking that sample is within the hypercube and 2D + if not sample.ndim == 2: + raise ValueError('Sample is not a 2D array') + + if (sample.max() > 1.) or (sample.min() < 0.): + raise ValueError('Sample is not in unit hypercube') + + # Checking that x_new is within the hypercube and 1D + if not x_new.ndim == 1: + raise ValueError('x_new is not a 1D array') + + if not (np.all(x_new >= 0) and np.all(x_new <= 1)): + raise ValueError('x_new is not in unit hypercube') + + if x_new.shape[0] != sample.shape[1]: + raise ValueError("x_new and sample must be broadcastable") + + return _cy_wrapper_update_discrepancy(x_new, sample, initial_disc) + + +def _perturb_discrepancy(sample: np.ndarray, i1: int, i2: int, k: int, + disc: float): + """Centered discrepancy after an elementary perturbation of a LHS. + + An elementary perturbation consists of an exchange of coordinates between + two points: ``sample[i1, k] <-> sample[i2, k]``. By construction, + this operation conserves the LHS properties. + + Parameters + ---------- + sample : array_like (n, d) + The sample (before permutation) to compute the discrepancy from. + i1 : int + The first line of the elementary permutation. + i2 : int + The second line of the elementary permutation. + k : int + The column of the elementary permutation. + disc : float + Centered discrepancy of the design before permutation. + + Returns + ------- + discrepancy : float + Centered discrepancy of the design after permutation. + + References + ---------- + .. [1] Jin et al. "An efficient algorithm for constructing optimal design + of computer experiments", Journal of Statistical Planning and + Inference, 2005. + + """ + n = sample.shape[0] + + z_ij = sample - 0.5 + + # Eq (19) + c_i1j = (1. / n ** 2. + * np.prod(0.5 * (2. + abs(z_ij[i1, :]) + + abs(z_ij) - abs(z_ij[i1, :] - z_ij)), axis=1)) + c_i2j = (1. / n ** 2. + * np.prod(0.5 * (2. + abs(z_ij[i2, :]) + + abs(z_ij) - abs(z_ij[i2, :] - z_ij)), axis=1)) + + # Eq (20) + c_i1i1 = (1. / n ** 2 * np.prod(1 + abs(z_ij[i1, :])) + - 2. / n * np.prod(1. + 0.5 * abs(z_ij[i1, :]) + - 0.5 * z_ij[i1, :] ** 2)) + c_i2i2 = (1. / n ** 2 * np.prod(1 + abs(z_ij[i2, :])) + - 2. / n * np.prod(1. + 0.5 * abs(z_ij[i2, :]) + - 0.5 * z_ij[i2, :] ** 2)) + + # Eq (22), typo in the article in the denominator i2 -> i1 + num = (2 + abs(z_ij[i2, k]) + abs(z_ij[:, k]) + - abs(z_ij[i2, k] - z_ij[:, k])) + denum = (2 + abs(z_ij[i1, k]) + abs(z_ij[:, k]) + - abs(z_ij[i1, k] - z_ij[:, k])) + gamma = num / denum + + # Eq (23) + c_p_i1j = gamma * c_i1j + # Eq (24) + c_p_i2j = c_i2j / gamma + + alpha = (1 + abs(z_ij[i2, k])) / (1 + abs(z_ij[i1, k])) + beta = (2 - abs(z_ij[i2, k])) / (2 - abs(z_ij[i1, k])) + + g_i1 = np.prod(1. + abs(z_ij[i1, :])) + g_i2 = np.prod(1. + abs(z_ij[i2, :])) + h_i1 = np.prod(1. + 0.5 * abs(z_ij[i1, :]) - 0.5 * (z_ij[i1, :] ** 2)) + h_i2 = np.prod(1. + 0.5 * abs(z_ij[i2, :]) - 0.5 * (z_ij[i2, :] ** 2)) + + # Eq (25), typo in the article g is missing + c_p_i1i1 = ((g_i1 * alpha) / (n ** 2) - 2. * alpha * beta * h_i1 / n) + # Eq (26), typo in the article n ** 2 + c_p_i2i2 = ((g_i2 / ((n ** 2) * alpha)) - (2. * h_i2 / (n * alpha * beta))) + + # Eq (26) + sum_ = c_p_i1j - c_i1j + c_p_i2j - c_i2j + + mask = np.ones(n, dtype=bool) + mask[[i1, i2]] = False + sum_ = sum(sum_[mask]) + + disc_ep = (disc + c_p_i1i1 - c_i1i1 + c_p_i2i2 - c_i2i2 + 2 * sum_) + + return disc_ep + + +def primes_from_2_to(n: int) -> np.ndarray: + """Prime numbers from 2 to *n*. + + Parameters + ---------- + n : int + Sup bound with ``n >= 6``. + + Returns + ------- + primes : list(int) + Primes in ``2 <= p < n``. + + Notes + ----- + Taken from [1]_ by P.T. Roy, written consent given on 23.04.2021 + by the original author, Bruno Astrolino, for free use in SciPy under + the 3-clause BSD. + + References + ---------- + .. [1] `StackOverflow `_. + + """ + sieve = np.ones(n // 3 + (n % 6 == 2), dtype=bool) + for i in range(1, int(n ** 0.5) // 3 + 1): + k = 3 * i + 1 | 1 + sieve[k * k // 3::2 * k] = False + sieve[k * (k - 2 * (i & 1) + 4) // 3::2 * k] = False + return np.r_[2, 3, ((3 * np.nonzero(sieve)[0][1:] + 1) | 1)] + + +def n_primes(n: IntNumber) -> list[int] | np.ndarray: + """List of the n-first prime numbers. + + Parameters + ---------- + n : int + Number of prime numbers wanted. + + Returns + ------- + primes : list(int) + List of primes. + + """ + primes = [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, + 61, 67, 71, 73, 79, 83, 89, 97, 101, 103, 107, 109, 113, 127, + 131, 137, 139, 149, 151, 157, 163, 167, 173, 179, 181, 191, 193, + 197, 199, 211, 223, 227, 229, 233, 239, 241, 251, 257, 263, 269, + 271, 277, 281, 283, 293, 307, 311, 313, 317, 331, 337, 347, 349, + 353, 359, 367, 373, 379, 383, 389, 397, 401, 409, 419, 421, 431, + 433, 439, 443, 449, 457, 461, 463, 467, 479, 487, 491, 499, 503, + 509, 521, 523, 541, 547, 557, 563, 569, 571, 577, 587, 593, 599, + 601, 607, 613, 617, 619, 631, 641, 643, 647, 653, 659, 661, 673, + 677, 683, 691, 701, 709, 719, 727, 733, 739, 743, 751, 757, 761, + 769, 773, 787, 797, 809, 811, 821, 823, 827, 829, 839, 853, 857, + 859, 863, 877, 881, 883, 887, 907, 911, 919, 929, 937, 941, 947, + 953, 967, 971, 977, 983, 991, 997][:n] + + if len(primes) < n: + big_number = 2000 + while 'Not enough primes': + primes = primes_from_2_to(big_number)[:n] # type: ignore + if len(primes) == n: + break + big_number += 1000 + + return primes + + +def _van_der_corput_permutations( + base: IntNumber, *, rng: SeedType = None +) -> np.ndarray: + """Permutations for scrambling a Van der Corput sequence. + + Parameters + ---------- + base : int + Base of the sequence. + rng : `numpy.random.Generator`, optional + Pseudorandom number generator state. When `rng` is None, a new + `numpy.random.Generator` is created using entropy from the + operating system. Types other than `numpy.random.Generator` are + passed to `numpy.random.default_rng` to instantiate a ``Generator``. + + .. versionchanged:: 1.15.0 + + As part of the `SPEC-007 `_ + transition from use of `numpy.random.RandomState` to + `numpy.random.Generator`, this keyword was changed from `seed` to + `rng`. During the transition, the behavior documented above is not + accurate; see `check_random_state` for actual behavior. After the + transition, this admonition can be removed. + + Returns + ------- + permutations : array_like + Permutation indices. + + Notes + ----- + In Algorithm 1 of Owen 2017, a permutation of `np.arange(base)` is + created for each positive integer `k` such that ``1 - base**-k < 1`` + using floating-point arithmetic. For double precision floats, the + condition ``1 - base**-k < 1`` can also be written as ``base**-k > + 2**-54``, which makes it more apparent how many permutations we need + to create. + """ + rng = check_random_state(rng) + count = math.ceil(54 / math.log2(base)) - 1 + permutations = np.repeat(np.arange(base)[None], count, axis=0) + for perm in permutations: + rng.shuffle(perm) + + return permutations + + +def van_der_corput( + n: IntNumber, + base: IntNumber = 2, + *, + start_index: IntNumber = 0, + scramble: bool = False, + permutations: "npt.ArrayLike | None" = None, + rng: SeedType = None, + workers: IntNumber = 1) -> np.ndarray: + """Van der Corput sequence. + + Pseudo-random number generator based on a b-adic expansion. + + Scrambling uses permutations of the remainders (see [1]_). Multiple + permutations are applied to construct a point. The sequence of + permutations has to be the same for all points of the sequence. + + Parameters + ---------- + n : int + Number of element of the sequence. + base : int, optional + Base of the sequence. Default is 2. + start_index : int, optional + Index to start the sequence from. Default is 0. + scramble : bool, optional + If True, use Owen scrambling. Otherwise no scrambling is done. + Default is True. + permutations : array_like, optional + Permutations used for scrambling. + rng : `numpy.random.Generator`, optional + Pseudorandom number generator state. When `rng` is None, a new + `numpy.random.Generator` is created using entropy from the + operating system. Types other than `numpy.random.Generator` are + passed to `numpy.random.default_rng` to instantiate a ``Generator``. + workers : int, optional + Number of workers to use for parallel processing. If -1 is + given all CPU threads are used. Default is 1. + + Returns + ------- + sequence : list (n,) + Sequence of Van der Corput. + + References + ---------- + .. [1] A. B. Owen. "A randomized Halton algorithm in R", + :arxiv:`1706.02808`, 2017. + + """ + if base < 2: + raise ValueError("'base' must be at least 2") + + if scramble: + if permutations is None: + permutations = _van_der_corput_permutations( + base=base, rng=rng + ) + else: + permutations = np.asarray(permutations) + + permutations = permutations.astype(np.int64) + return _cy_van_der_corput_scrambled(n, base, start_index, + permutations, workers) + + else: + return _cy_van_der_corput(n, base, start_index, workers) + + +class QMCEngine(ABC): + """A generic Quasi-Monte Carlo sampler class meant for subclassing. + + QMCEngine is a base class to construct a specific Quasi-Monte Carlo + sampler. It cannot be used directly as a sampler. + + Parameters + ---------- + d : int + Dimension of the parameter space. + optimization : {None, "random-cd", "lloyd"}, optional + Whether to use an optimization scheme to improve the quality after + sampling. Note that this is a post-processing step that does not + guarantee that all properties of the sample will be conserved. + Default is None. + + * ``random-cd``: random permutations of coordinates to lower the + centered discrepancy. The best sample based on the centered + discrepancy is constantly updated. Centered discrepancy-based + sampling shows better space-filling robustness toward 2D and 3D + subprojections compared to using other discrepancy measures. + * ``lloyd``: Perturb samples using a modified Lloyd-Max algorithm. + The process converges to equally spaced samples. + + .. versionadded:: 1.10.0 + + rng : `numpy.random.Generator`, optional + Pseudorandom number generator state. When `rng` is None, a new + `numpy.random.Generator` is created using entropy from the + operating system. Types other than `numpy.random.Generator` are + passed to `numpy.random.default_rng` to instantiate a ``Generator``. + + .. versionchanged:: 1.15.0 + + As part of the `SPEC-007 `_ + transition from use of `numpy.random.RandomState` to + `numpy.random.Generator`, this keyword was changed from `seed` to + `rng`. For an interim period, both keywords will continue to work, although + only one may be specified at a time. After the interim period, function + calls using the `seed` keyword will emit warnings. Following a + deprecation period, the `seed` keyword will be removed. + + Notes + ----- + By convention samples are distributed over the half-open interval + ``[0, 1)``. Instances of the class can access the attributes: ``d`` for + the dimension; and ``rng`` for the random number generator. + + **Subclassing** + + When subclassing `QMCEngine` to create a new sampler, ``__init__`` and + ``random`` must be redefined. + + * ``__init__(d, rng=None)``: at least fix the dimension. If the sampler + does not take advantage of a ``rng`` (deterministic methods like + Halton), this parameter can be omitted. + * ``_random(n, *, workers=1)``: draw ``n`` from the engine. ``workers`` + is used for parallelism. See `Halton` for example. + + Optionally, two other methods can be overwritten by subclasses: + + * ``reset``: Reset the engine to its original state. + * ``fast_forward``: If the sequence is deterministic (like Halton + sequence), then ``fast_forward(n)`` is skipping the ``n`` first draw. + + Examples + -------- + To create a random sampler based on ``np.random.random``, we would do the + following: + + >>> from scipy.stats import qmc + >>> class RandomEngine(qmc.QMCEngine): + ... def __init__(self, d, rng=None): + ... super().__init__(d=d, rng=rng) + ... + ... + ... def _random(self, n=1, *, workers=1): + ... return self.rng.random((n, self.d)) + ... + ... + ... def reset(self): + ... super().__init__(d=self.d, rng=self.rng_seed) + ... return self + ... + ... + ... def fast_forward(self, n): + ... self.random(n) + ... return self + + After subclassing `QMCEngine` to define the sampling strategy we want to + use, we can create an instance to sample from. + + >>> engine = RandomEngine(2) + >>> engine.random(5) + array([[0.22733602, 0.31675834], # random + [0.79736546, 0.67625467], + [0.39110955, 0.33281393], + [0.59830875, 0.18673419], + [0.67275604, 0.94180287]]) + + We can also reset the state of the generator and resample again. + + >>> _ = engine.reset() + >>> engine.random(5) + array([[0.22733602, 0.31675834], # random + [0.79736546, 0.67625467], + [0.39110955, 0.33281393], + [0.59830875, 0.18673419], + [0.67275604, 0.94180287]]) + + """ + + @abstractmethod + @_transition_to_rng('seed', replace_doc=False) + def __init__( + self, + d: IntNumber, + *, + optimization: Literal["random-cd", "lloyd"] | None = None, + rng: SeedType = None + ) -> None: + self._initialize(d, optimization=optimization, rng=rng) + + # During SPEC 7 transition: + # `__init__` has to be wrapped with @_transition_to_rng decorator + # because it is public. Subclasses previously called `__init__` + # directly, but this was problematic because arguments passed to + # subclass `__init__` as `seed` would get passed to superclass + # `__init__` as `rng`, rejecting `RandomState` arguments. + def _initialize( + self, + d: IntNumber, + *, + optimization: Literal["random-cd", "lloyd"] | None = None, + rng: SeedType = None + ) -> None: + if not np.issubdtype(type(d), np.integer) or d < 0: + raise ValueError('d must be a non-negative integer value') + + self.d = d + + if isinstance(rng, np.random.Generator): + # Spawn a Generator that we can own and reset. + self.rng = _rng_spawn(rng, 1)[0] + else: + # Create our instance of Generator, does not need spawning + # Also catch RandomState which cannot be spawned + self.rng = check_random_state(rng) + self.rng_seed = copy.deepcopy(self.rng) + + self.num_generated = 0 + + config = { + # random-cd + "n_nochange": 100, + "n_iters": 10_000, + "rng": self.rng, + + # lloyd + "tol": 1e-5, + "maxiter": 10, + "qhull_options": None, + } + self._optimization = optimization + self.optimization_method = _select_optimizer(optimization, config) + + @abstractmethod + def _random( + self, n: IntNumber = 1, *, workers: IntNumber = 1 + ) -> np.ndarray: + ... + + def random( + self, n: IntNumber = 1, *, workers: IntNumber = 1 + ) -> np.ndarray: + """Draw `n` in the half-open interval ``[0, 1)``. + + Parameters + ---------- + n : int, optional + Number of samples to generate in the parameter space. + Default is 1. + workers : int, optional + Only supported with `Halton`. + Number of workers to use for parallel processing. If -1 is + given all CPU threads are used. Default is 1. It becomes faster + than one worker for `n` greater than :math:`10^3`. + + Returns + ------- + sample : array_like (n, d) + QMC sample. + + """ + sample = self._random(n, workers=workers) + if self.optimization_method is not None: + sample = self.optimization_method(sample) + + self.num_generated += n + return sample + + def integers( + self, + l_bounds: "npt.ArrayLike", + *, + u_bounds: "npt.ArrayLike | None" = None, + n: IntNumber = 1, + endpoint: bool = False, + workers: IntNumber = 1 + ) -> np.ndarray: + r""" + Draw `n` integers from `l_bounds` (inclusive) to `u_bounds` + (exclusive), or if endpoint=True, `l_bounds` (inclusive) to + `u_bounds` (inclusive). + + Parameters + ---------- + l_bounds : int or array-like of ints + Lowest (signed) integers to be drawn (unless ``u_bounds=None``, + in which case this parameter is 0 and this value is used for + `u_bounds`). + u_bounds : int or array-like of ints, optional + If provided, one above the largest (signed) integer to be drawn + (see above for behavior if ``u_bounds=None``). + If array-like, must contain integer values. + n : int, optional + Number of samples to generate in the parameter space. + Default is 1. + endpoint : bool, optional + If true, sample from the interval ``[l_bounds, u_bounds]`` instead + of the default ``[l_bounds, u_bounds)``. Defaults is False. + workers : int, optional + Number of workers to use for parallel processing. If -1 is + given all CPU threads are used. Only supported when using `Halton` + Default is 1. + + Returns + ------- + sample : array_like (n, d) + QMC sample. + + Notes + ----- + It is safe to just use the same ``[0, 1)`` to integer mapping + with QMC that you would use with MC. You still get unbiasedness, + a strong law of large numbers, an asymptotically infinite variance + reduction and a finite sample variance bound. + + To convert a sample from :math:`[0, 1)` to :math:`[a, b), b>a`, + with :math:`a` the lower bounds and :math:`b` the upper bounds, + the following transformation is used: + + .. math:: + + \text{floor}((b - a) \cdot \text{sample} + a) + + """ + if u_bounds is None: + u_bounds = l_bounds + l_bounds = 0 + + u_bounds = np.atleast_1d(u_bounds) + l_bounds = np.atleast_1d(l_bounds) + + if endpoint: + u_bounds = u_bounds + 1 + + if (not np.issubdtype(l_bounds.dtype, np.integer) or + not np.issubdtype(u_bounds.dtype, np.integer)): + message = ("'u_bounds' and 'l_bounds' must be integers or" + " array-like of integers") + raise ValueError(message) + + if isinstance(self, Halton): + sample = self.random(n=n, workers=workers) + else: + sample = self.random(n=n) + + sample = scale(sample, l_bounds=l_bounds, u_bounds=u_bounds) + sample = np.floor(sample).astype(np.int64) + + return sample + + def reset(self) -> "QMCEngine": + """Reset the engine to base state. + + Returns + ------- + engine : QMCEngine + Engine reset to its base state. + + """ + rng = copy.deepcopy(self.rng_seed) + self.rng = check_random_state(rng) + self.num_generated = 0 + return self + + def fast_forward(self, n: IntNumber) -> "QMCEngine": + """Fast-forward the sequence by `n` positions. + + Parameters + ---------- + n : int + Number of points to skip in the sequence. + + Returns + ------- + engine : QMCEngine + Engine reset to its base state. + + """ + self.random(n=n) + return self + + +class Halton(QMCEngine): + """Halton sequence. + + Pseudo-random number generator that generalize the Van der Corput sequence + for multiple dimensions. The Halton sequence uses the base-two Van der + Corput sequence for the first dimension, base-three for its second and + base-:math:`p` for its :math:`n`-dimension, with :math:`p` the + :math:`n`'th prime. + + Parameters + ---------- + d : int + Dimension of the parameter space. + scramble : bool, optional + If True, use random scrambling from [2]_. Otherwise no scrambling + is done. + Default is True. + optimization : {None, "random-cd", "lloyd"}, optional + Whether to use an optimization scheme to improve the quality after + sampling. Note that this is a post-processing step that does not + guarantee that all properties of the sample will be conserved. + Default is None. + + * ``random-cd``: random permutations of coordinates to lower the + centered discrepancy. The best sample based on the centered + discrepancy is constantly updated. Centered discrepancy-based + sampling shows better space-filling robustness toward 2D and 3D + subprojections compared to using other discrepancy measures. + * ``lloyd``: Perturb samples using a modified Lloyd-Max algorithm. + The process converges to equally spaced samples. + + .. versionadded:: 1.10.0 + + rng : `numpy.random.Generator`, optional + Pseudorandom number generator state. When `rng` is None, a new + `numpy.random.Generator` is created using entropy from the + operating system. Types other than `numpy.random.Generator` are + passed to `numpy.random.default_rng` to instantiate a ``Generator``. + + .. versionchanged:: 1.15.0 + + As part of the `SPEC-007 `_ + transition from use of `numpy.random.RandomState` to + `numpy.random.Generator`, this keyword was changed from `seed` to + `rng`. For an interim period, both keywords will continue to work, although + only one may be specified at a time. After the interim period, function + calls using the `seed` keyword will emit warnings. Following a + deprecation period, the `seed` keyword will be removed. + + Notes + ----- + The Halton sequence has severe striping artifacts for even modestly + large dimensions. These can be ameliorated by scrambling. Scrambling + also supports replication-based error estimates and extends + applicability to unbounded integrands. + + References + ---------- + .. [1] Halton, "On the efficiency of certain quasi-random sequences of + points in evaluating multi-dimensional integrals", Numerische + Mathematik, 1960. + .. [2] A. B. Owen. "A randomized Halton algorithm in R", + :arxiv:`1706.02808`, 2017. + + Examples + -------- + Generate samples from a low discrepancy sequence of Halton. + + >>> from scipy.stats import qmc + >>> sampler = qmc.Halton(d=2, scramble=False) + >>> sample = sampler.random(n=5) + >>> sample + array([[0. , 0. ], + [0.5 , 0.33333333], + [0.25 , 0.66666667], + [0.75 , 0.11111111], + [0.125 , 0.44444444]]) + + Compute the quality of the sample using the discrepancy criterion. + + >>> qmc.discrepancy(sample) + 0.088893711419753 + + If some wants to continue an existing design, extra points can be obtained + by calling again `random`. Alternatively, you can skip some points like: + + >>> _ = sampler.fast_forward(5) + >>> sample_continued = sampler.random(n=5) + >>> sample_continued + array([[0.3125 , 0.37037037], + [0.8125 , 0.7037037 ], + [0.1875 , 0.14814815], + [0.6875 , 0.48148148], + [0.4375 , 0.81481481]]) + + Finally, samples can be scaled to bounds. + + >>> l_bounds = [0, 2] + >>> u_bounds = [10, 5] + >>> qmc.scale(sample_continued, l_bounds, u_bounds) + array([[3.125 , 3.11111111], + [8.125 , 4.11111111], + [1.875 , 2.44444444], + [6.875 , 3.44444444], + [4.375 , 4.44444444]]) + + """ + @_transition_to_rng('seed', replace_doc=False) + def __init__( + self, d: IntNumber, *, scramble: bool = True, + optimization: Literal["random-cd", "lloyd"] | None = None, + rng: SeedType = None + ) -> None: + # Used in `scipy.integrate.qmc_quad` + self._init_quad = {'d': d, 'scramble': True, + 'optimization': optimization} + super()._initialize(d=d, optimization=optimization, rng=rng) + + # important to have ``type(bdim) == int`` for performance reason + self.base = [int(bdim) for bdim in n_primes(d)] + self.scramble = scramble + + self._initialize_permutations() + + def _initialize_permutations(self) -> None: + """Initialize permutations for all Van der Corput sequences. + + Permutations are only needed for scrambling. + """ + self._permutations: list = [None] * len(self.base) + if self.scramble: + for i, bdim in enumerate(self.base): + permutations = _van_der_corput_permutations( + base=bdim, rng=self.rng + ) + + self._permutations[i] = permutations + + def _random( + self, n: IntNumber = 1, *, workers: IntNumber = 1 + ) -> np.ndarray: + """Draw `n` in the half-open interval ``[0, 1)``. + + Parameters + ---------- + n : int, optional + Number of samples to generate in the parameter space. Default is 1. + workers : int, optional + Number of workers to use for parallel processing. If -1 is + given all CPU threads are used. Default is 1. It becomes faster + than one worker for `n` greater than :math:`10^3`. + + Returns + ------- + sample : array_like (n, d) + QMC sample. + + """ + workers = _validate_workers(workers) + # Generate a sample using a Van der Corput sequence per dimension. + sample = [van_der_corput(n, bdim, start_index=self.num_generated, + scramble=self.scramble, + permutations=self._permutations[i], + workers=workers) + for i, bdim in enumerate(self.base)] + + return np.array(sample).T.reshape(n, self.d) + + +class LatinHypercube(QMCEngine): + r"""Latin hypercube sampling (LHS). + + A Latin hypercube sample [1]_ generates :math:`n` points in + :math:`[0,1)^{d}`. Each univariate marginal distribution is stratified, + placing exactly one point in :math:`[j/n, (j+1)/n)` for + :math:`j=0,1,...,n-1`. They are still applicable when :math:`n << d`. + + Parameters + ---------- + d : int + Dimension of the parameter space. + scramble : bool, optional + When False, center samples within cells of a multi-dimensional grid. + Otherwise, samples are randomly placed within cells of the grid. + + .. note:: + Setting ``scramble=False`` does not ensure deterministic output. + For that, use the `rng` parameter. + + Default is True. + + .. versionadded:: 1.10.0 + + optimization : {None, "random-cd", "lloyd"}, optional + Whether to use an optimization scheme to improve the quality after + sampling. Note that this is a post-processing step that does not + guarantee that all properties of the sample will be conserved. + Default is None. + + * ``random-cd``: random permutations of coordinates to lower the + centered discrepancy. The best sample based on the centered + discrepancy is constantly updated. Centered discrepancy-based + sampling shows better space-filling robustness toward 2D and 3D + subprojections compared to using other discrepancy measures. + * ``lloyd``: Perturb samples using a modified Lloyd-Max algorithm. + The process converges to equally spaced samples. + + .. versionadded:: 1.8.0 + .. versionchanged:: 1.10.0 + Add ``lloyd``. + + strength : {1, 2}, optional + Strength of the LHS. ``strength=1`` produces a plain LHS while + ``strength=2`` produces an orthogonal array based LHS of strength 2 + [7]_, [8]_. In that case, only ``n=p**2`` points can be sampled, + with ``p`` a prime number. It also constrains ``d <= p + 1``. + Default is 1. + + .. versionadded:: 1.8.0 + + rng : `numpy.random.Generator`, optional + Pseudorandom number generator state. When `rng` is None, a new + `numpy.random.Generator` is created using entropy from the + operating system. Types other than `numpy.random.Generator` are + passed to `numpy.random.default_rng` to instantiate a ``Generator``. + + .. versionchanged:: 1.15.0 + + As part of the `SPEC-007 `_ + transition from use of `numpy.random.RandomState` to + `numpy.random.Generator`, this keyword was changed from `seed` to + `rng`. For an interim period, both keywords will continue to work, although + only one may be specified at a time. After the interim period, function + calls using the `seed` keyword will emit warnings. Following a + deprecation period, the `seed` keyword will be removed. + + See Also + -------- + :ref:`quasi-monte-carlo` + + Notes + ----- + + When LHS is used for integrating a function :math:`f` over :math:`n`, + LHS is extremely effective on integrands that are nearly additive [2]_. + With a LHS of :math:`n` points, the variance of the integral is always + lower than plain MC on :math:`n-1` points [3]_. There is a central limit + theorem for LHS on the mean and variance of the integral [4]_, but not + necessarily for optimized LHS due to the randomization. + + :math:`A` is called an orthogonal array of strength :math:`t` if in each + n-row-by-t-column submatrix of :math:`A`: all :math:`p^t` possible + distinct rows occur the same number of times. The elements of :math:`A` + are in the set :math:`\{0, 1, ..., p-1\}`, also called symbols. + The constraint that :math:`p` must be a prime number is to allow modular + arithmetic. Increasing strength adds some symmetry to the sub-projections + of a sample. With strength 2, samples are symmetric along the diagonals of + 2D sub-projections. This may be undesirable, but on the other hand, the + sample dispersion is improved. + + Strength 1 (plain LHS) brings an advantage over strength 0 (MC) and + strength 2 is a useful increment over strength 1. Going to strength 3 is + a smaller increment and scrambled QMC like Sobol', Halton are more + performant [7]_. + + To create a LHS of strength 2, the orthogonal array :math:`A` is + randomized by applying a random, bijective map of the set of symbols onto + itself. For example, in column 0, all 0s might become 2; in column 1, + all 0s might become 1, etc. + Then, for each column :math:`i` and symbol :math:`j`, we add a plain, + one-dimensional LHS of size :math:`p` to the subarray where + :math:`A^i = j`. The resulting matrix is finally divided by :math:`p`. + + References + ---------- + .. [1] Mckay et al., "A Comparison of Three Methods for Selecting Values + of Input Variables in the Analysis of Output from a Computer Code." + Technometrics, 1979. + .. [2] M. Stein, "Large sample properties of simulations using Latin + hypercube sampling." Technometrics 29, no. 2: 143-151, 1987. + .. [3] A. B. Owen, "Monte Carlo variance of scrambled net quadrature." + SIAM Journal on Numerical Analysis 34, no. 5: 1884-1910, 1997 + .. [4] Loh, W.-L. "On Latin hypercube sampling." The annals of statistics + 24, no. 5: 2058-2080, 1996. + .. [5] Fang et al. "Design and modeling for computer experiments". + Computer Science and Data Analysis Series, 2006. + .. [6] Damblin et al., "Numerical studies of space filling designs: + optimization of Latin Hypercube Samples and subprojection properties." + Journal of Simulation, 2013. + .. [7] A. B. Owen , "Orthogonal arrays for computer experiments, + integration and visualization." Statistica Sinica, 1992. + .. [8] B. Tang, "Orthogonal Array-Based Latin Hypercubes." + Journal of the American Statistical Association, 1993. + .. [9] Seaholm, Susan K. et al. (1988). Latin hypercube sampling and the + sensitivity analysis of a Monte Carlo epidemic model. Int J Biomed + Comput, 23(1-2), 97-112. :doi:`10.1016/0020-7101(88)90067-0` + + Examples + -------- + Generate samples from a Latin hypercube generator. + + >>> from scipy.stats import qmc + >>> sampler = qmc.LatinHypercube(d=2) + >>> sample = sampler.random(n=5) + >>> sample + array([[0.1545328 , 0.53664833], # random + [0.84052691, 0.06474907], + [0.52177809, 0.93343721], + [0.68033825, 0.36265316], + [0.26544879, 0.61163943]]) + + Compute the quality of the sample using the discrepancy criterion. + + >>> qmc.discrepancy(sample) + 0.0196... # random + + Samples can be scaled to bounds. + + >>> l_bounds = [0, 2] + >>> u_bounds = [10, 5] + >>> qmc.scale(sample, l_bounds, u_bounds) + array([[1.54532796, 3.609945 ], # random + [8.40526909, 2.1942472 ], + [5.2177809 , 4.80031164], + [6.80338249, 3.08795949], + [2.65448791, 3.83491828]]) + + Below are other examples showing alternative ways to construct LHS with + even better coverage of the space. + + Using a base LHS as a baseline. + + >>> sampler = qmc.LatinHypercube(d=2) + >>> sample = sampler.random(n=5) + >>> qmc.discrepancy(sample) + 0.0196... # random + + Use the `optimization` keyword argument to produce a LHS with + lower discrepancy at higher computational cost. + + >>> sampler = qmc.LatinHypercube(d=2, optimization="random-cd") + >>> sample = sampler.random(n=5) + >>> qmc.discrepancy(sample) + 0.0176... # random + + Use the `strength` keyword argument to produce an orthogonal array based + LHS of strength 2. In this case, the number of sample points must be the + square of a prime number. + + >>> sampler = qmc.LatinHypercube(d=2, strength=2) + >>> sample = sampler.random(n=9) + >>> qmc.discrepancy(sample) + 0.00526... # random + + Options could be combined to produce an optimized centered + orthogonal array based LHS. After optimization, the result would not + be guaranteed to be of strength 2. + + **Real-world example** + + In [9]_, a Latin Hypercube sampling (LHS) strategy was used to sample a + parameter space to study the importance of each parameter of an epidemic + model. Such analysis is also called a sensitivity analysis. + + Since the dimensionality of the problem is high (6), it is computationally + expensive to cover the space. When numerical experiments are costly, QMC + enables analysis that may not be possible if using a grid. + + The six parameters of the model represented the probability of illness, + the probability of withdrawal, and four contact probabilities. The + authors assumed uniform distributions for all parameters and generated + 50 samples. + + Using `scipy.stats.qmc.LatinHypercube` to replicate the protocol, + the first step is to create a sample in the unit hypercube: + + >>> from scipy.stats import qmc + >>> sampler = qmc.LatinHypercube(d=6) + >>> sample = sampler.random(n=50) + + Then the sample can be scaled to the appropriate bounds: + + >>> l_bounds = [0.000125, 0.01, 0.0025, 0.05, 0.47, 0.7] + >>> u_bounds = [0.000375, 0.03, 0.0075, 0.15, 0.87, 0.9] + >>> sample_scaled = qmc.scale(sample, l_bounds, u_bounds) + + Such a sample was used to run the model 50 times, and a polynomial + response surface was constructed. This allowed the authors to study the + relative importance of each parameter across the range of possibilities + of every other parameter. + + In this computer experiment, they showed a 14-fold reduction in the + number of samples required to maintain an error below 2% on their + response surface when compared to a grid sampling. + + """ + + @_transition_to_rng('seed', replace_doc=False) + def __init__( + self, d: IntNumber, *, + scramble: bool = True, + strength: int = 1, + optimization: Literal["random-cd", "lloyd"] | None = None, + rng: SeedType = None + ) -> None: + # Used in `scipy.integrate.qmc_quad` + self._init_quad = {'d': d, 'scramble': True, 'strength': strength, + 'optimization': optimization} + super()._initialize(d=d, rng=rng, optimization=optimization) + self.scramble = scramble + + lhs_method_strength = { + 1: self._random_lhs, + 2: self._random_oa_lhs + } + + try: + self.lhs_method: Callable = lhs_method_strength[strength] + except KeyError as exc: + message = (f"{strength!r} is not a valid strength. It must be one" + f" of {set(lhs_method_strength)!r}") + raise ValueError(message) from exc + + def _random( + self, n: IntNumber = 1, *, workers: IntNumber = 1 + ) -> np.ndarray: + lhs = self.lhs_method(n) + return lhs + + def _random_lhs(self, n: IntNumber = 1) -> np.ndarray: + """Base LHS algorithm.""" + if not self.scramble: + samples: np.ndarray | float = 0.5 + else: + samples = self.rng.uniform(size=(n, self.d)) + + perms = np.tile(np.arange(1, n + 1), (int(self.d), 1)) + for i in range(self.d): + self.rng.shuffle(perms[i, :]) + perms = perms.T + + samples = (perms - samples) / n + return samples + + def _random_oa_lhs(self, n: IntNumber = 4) -> np.ndarray: + """Orthogonal array based LHS of strength 2.""" + p = np.sqrt(n).astype(int) + n_row = p**2 + n_col = p + 1 + + primes = primes_from_2_to(p + 1) + if p not in primes or n != n_row: + raise ValueError( + "n is not the square of a prime number. Close" + f" values are {primes[-2:]**2}" + ) + if self.d > p + 1: + raise ValueError("n is too small for d. Must be n > (d-1)**2") + + oa_sample = np.zeros(shape=(n_row, n_col), dtype=int) + + # OA of strength 2 + arrays = np.tile(np.arange(p), (2, 1)) + oa_sample[:, :2] = np.stack(np.meshgrid(*arrays), + axis=-1).reshape(-1, 2) + for p_ in range(1, p): + oa_sample[:, 2+p_-1] = np.mod(oa_sample[:, 0] + + p_*oa_sample[:, 1], p) + + # scramble the OA + oa_sample_ = np.empty(shape=(n_row, n_col), dtype=int) + for j in range(n_col): + perms = self.rng.permutation(p) + oa_sample_[:, j] = perms[oa_sample[:, j]] + + oa_sample = oa_sample_ + # following is making a scrambled OA into an OA-LHS + oa_lhs_sample = np.zeros(shape=(n_row, n_col)) + lhs_engine = LatinHypercube(d=1, scramble=self.scramble, strength=1, + rng=self.rng) # type: QMCEngine + for j in range(n_col): + for k in range(p): + idx = oa_sample[:, j] == k + lhs = lhs_engine.random(p).flatten() + oa_lhs_sample[:, j][idx] = lhs + oa_sample[:, j][idx] + + oa_lhs_sample /= p + + return oa_lhs_sample[:, :self.d] + + +class Sobol(QMCEngine): + """Engine for generating (scrambled) Sobol' sequences. + + Sobol' sequences are low-discrepancy, quasi-random numbers. Points + can be drawn using two methods: + + * `random_base2`: safely draw :math:`n=2^m` points. This method + guarantees the balance properties of the sequence. + * `random`: draw an arbitrary number of points from the + sequence. See warning below. + + Parameters + ---------- + d : int + Dimensionality of the sequence. Max dimensionality is 21201. + scramble : bool, optional + If True, use LMS+shift scrambling. Otherwise, no scrambling is done. + Default is True. + bits : int, optional + Number of bits of the generator. Control the maximum number of points + that can be generated, which is ``2**bits``. Maximal value is 64. + It does not correspond to the return type, which is always + ``np.float64`` to prevent points from repeating themselves. + Default is None, which for backward compatibility, corresponds to 30. + + .. versionadded:: 1.9.0 + optimization : {None, "random-cd", "lloyd"}, optional + Whether to use an optimization scheme to improve the quality after + sampling. Note that this is a post-processing step that does not + guarantee that all properties of the sample will be conserved. + Default is None. + + * ``random-cd``: random permutations of coordinates to lower the + centered discrepancy. The best sample based on the centered + discrepancy is constantly updated. Centered discrepancy-based + sampling shows better space-filling robustness toward 2D and 3D + subprojections compared to using other discrepancy measures. + * ``lloyd``: Perturb samples using a modified Lloyd-Max algorithm. + The process converges to equally spaced samples. + + .. versionadded:: 1.10.0 + + rng : `numpy.random.Generator`, optional + Pseudorandom number generator state. When `rng` is None, a new + `numpy.random.Generator` is created using entropy from the + operating system. Types other than `numpy.random.Generator` are + passed to `numpy.random.default_rng` to instantiate a ``Generator``. + + .. versionchanged:: 1.15.0 + + As part of the `SPEC-007 `_ + transition from use of `numpy.random.RandomState` to + `numpy.random.Generator`, this keyword was changed from `seed` to + `rng`. For an interim period, both keywords will continue to work, although + only one may be specified at a time. After the interim period, function + calls using the `seed` keyword will emit warnings. Following a + deprecation period, the `seed` keyword will be removed. + + Notes + ----- + Sobol' sequences [1]_ provide :math:`n=2^m` low discrepancy points in + :math:`[0,1)^{d}`. Scrambling them [3]_ makes them suitable for singular + integrands, provides a means of error estimation, and can improve their + rate of convergence. The scrambling strategy which is implemented is a + (left) linear matrix scramble (LMS) followed by a digital random shift + (LMS+shift) [2]_. + + There are many versions of Sobol' sequences depending on their + 'direction numbers'. This code uses direction numbers from [4]_. Hence, + the maximum number of dimension is 21201. The direction numbers have been + precomputed with search criterion 6 and can be retrieved at + https://web.maths.unsw.edu.au/~fkuo/sobol/. + + .. warning:: + + Sobol' sequences are a quadrature rule and they lose their balance + properties if one uses a sample size that is not a power of 2, or skips + the first point, or thins the sequence [5]_. + + If :math:`n=2^m` points are not enough then one should take :math:`2^M` + points for :math:`M>m`. When scrambling, the number R of independent + replicates does not have to be a power of 2. + + Sobol' sequences are generated to some number :math:`B` of bits. + After :math:`2^B` points have been generated, the sequence would + repeat. Hence, an error is raised. + The number of bits can be controlled with the parameter `bits`. + + References + ---------- + .. [1] I. M. Sobol', "The distribution of points in a cube and the accurate + evaluation of integrals." Zh. Vychisl. Mat. i Mat. Phys., 7:784-802, + 1967. + .. [2] J. Matousek, "On the L2-discrepancy for anchored boxes." + J. of Complexity 14, 527-556, 1998. + .. [3] Art B. Owen, "Scrambling Sobol and Niederreiter-Xing points." + Journal of Complexity, 14(4):466-489, December 1998. + .. [4] S. Joe and F. Y. Kuo, "Constructing sobol sequences with better + two-dimensional projections." SIAM Journal on Scientific Computing, + 30(5):2635-2654, 2008. + .. [5] Art B. Owen, "On dropping the first Sobol' point." + :arxiv:`2008.08051`, 2020. + + Examples + -------- + Generate samples from a low discrepancy sequence of Sobol'. + + >>> from scipy.stats import qmc + >>> sampler = qmc.Sobol(d=2, scramble=False) + >>> sample = sampler.random_base2(m=3) + >>> sample + array([[0. , 0. ], + [0.5 , 0.5 ], + [0.75 , 0.25 ], + [0.25 , 0.75 ], + [0.375, 0.375], + [0.875, 0.875], + [0.625, 0.125], + [0.125, 0.625]]) + + Compute the quality of the sample using the discrepancy criterion. + + >>> qmc.discrepancy(sample) + 0.013882107204860938 + + To continue an existing design, extra points can be obtained + by calling again `random_base2`. Alternatively, you can skip some + points like: + + >>> _ = sampler.reset() + >>> _ = sampler.fast_forward(4) + >>> sample_continued = sampler.random_base2(m=2) + >>> sample_continued + array([[0.375, 0.375], + [0.875, 0.875], + [0.625, 0.125], + [0.125, 0.625]]) + + Finally, samples can be scaled to bounds. + + >>> l_bounds = [0, 2] + >>> u_bounds = [10, 5] + >>> qmc.scale(sample_continued, l_bounds, u_bounds) + array([[3.75 , 3.125], + [8.75 , 4.625], + [6.25 , 2.375], + [1.25 , 3.875]]) + + """ + + MAXDIM: ClassVar[int] = _MAXDIM + + @_transition_to_rng('seed', replace_doc=False) + def __init__( + self, d: IntNumber, *, scramble: bool = True, + bits: IntNumber | None = None, rng: SeedType = None, + optimization: Literal["random-cd", "lloyd"] | None = None + ) -> None: + # Used in `scipy.integrate.qmc_quad` + self._init_quad = {'d': d, 'scramble': True, 'bits': bits, + 'optimization': optimization} + + super()._initialize(d=d, optimization=optimization, rng=rng) + if d > self.MAXDIM: + raise ValueError( + f"Maximum supported dimensionality is {self.MAXDIM}." + ) + + self.bits = bits + self.dtype_i: type + self.scramble = scramble + + if self.bits is None: + self.bits = 30 + + if self.bits <= 32: + self.dtype_i = np.uint32 + elif 32 < self.bits <= 64: + self.dtype_i = np.uint64 + else: + raise ValueError("Maximum supported 'bits' is 64") + + self.maxn = 2**self.bits + + # v is d x maxbit matrix + self._sv: np.ndarray = np.zeros((d, self.bits), dtype=self.dtype_i) + _initialize_v(self._sv, dim=d, bits=self.bits) + + if not scramble: + self._shift: np.ndarray = np.zeros(d, dtype=self.dtype_i) + else: + # scramble self._shift and self._sv + self._scramble() + + self._quasi = self._shift.copy() + + # normalization constant with the largest possible number + # calculate in Python to not overflow int with 2**64 + self._scale = 1.0 / 2 ** self.bits + + self._first_point = (self._quasi * self._scale).reshape(1, -1) + # explicit casting to float64 + self._first_point = self._first_point.astype(np.float64) + + def _scramble(self) -> None: + """Scramble the sequence using LMS+shift.""" + # Generate shift vector + self._shift = np.dot( + rng_integers(self.rng, 2, size=(self.d, self.bits), + dtype=self.dtype_i), + # pyrefly:ignore[no-matching-overload] + 2 ** np.arange(self.bits, dtype=self.dtype_i), + ) + # Generate lower triangular matrices (stacked across dimensions) + ltm = np.tril(rng_integers(self.rng, 2, + size=(self.d, self.bits, self.bits), + dtype=self.dtype_i)) + _cscramble( + dim=self.d, bits=self.bits, # type: ignore[arg-type] + ltm=ltm, sv=self._sv + ) + + def _random( + self, n: IntNumber = 1, *, workers: IntNumber = 1 + ) -> np.ndarray: + """Draw next point(s) in the Sobol' sequence. + + Parameters + ---------- + n : int, optional + Number of samples to generate in the parameter space. Default is 1. + + Returns + ------- + sample : array_like (n, d) + Sobol' sample. + + """ + sample: np.ndarray = np.empty((n, self.d), dtype=np.float64) + + if n == 0: + return sample + + total_n = self.num_generated + n + if total_n > self.maxn: + msg = ( + f"At most 2**{self.bits}={self.maxn} distinct points can be " + f"generated. {self.num_generated} points have been previously " + f"generated, then: n={self.num_generated}+{n}={total_n}. " + ) + if self.bits != 64: + msg += "Consider increasing `bits`." + raise ValueError(msg) + + if self.num_generated == 0: + # verify n is 2**n + if not (n & (n - 1) == 0): + warnings.warn("The balance properties of Sobol' points require" + " n to be a power of 2.", stacklevel=3) + + if n == 1: + sample = self._first_point + else: + _draw( + n=n - 1, num_gen=self.num_generated, dim=self.d, + scale=self._scale, sv=self._sv, quasi=self._quasi, + sample=sample + ) + sample = np.concatenate( + [self._first_point, sample] + )[:n] + else: + _draw( + n=n, num_gen=self.num_generated - 1, dim=self.d, + scale=self._scale, sv=self._sv, quasi=self._quasi, + sample=sample + ) + + return sample + + def random_base2(self, m: IntNumber) -> np.ndarray: + """Draw point(s) from the Sobol' sequence. + + This function draws :math:`n=2^m` points in the parameter space + ensuring the balance properties of the sequence. + + Parameters + ---------- + m : int + Logarithm in base 2 of the number of samples; i.e., n = 2^m. + + Returns + ------- + sample : array_like (n, d) + Sobol' sample. + + """ + n = 2 ** m + + total_n = self.num_generated + n + if not (total_n & (total_n - 1) == 0): + raise ValueError('The balance properties of Sobol\' points require ' + f'n to be a power of 2. {self.num_generated} points ' + 'have been previously generated, then: ' + f'n={self.num_generated}+2**{m}={total_n}. ' + 'If you still want to do this, the function ' + '\'Sobol.random()\' can be used.' + ) + + return self.random(n) + + def reset(self) -> "Sobol": + """Reset the engine to base state. + + Returns + ------- + engine : Sobol + Engine reset to its base state. + + """ + super().reset() + self._quasi = self._shift.copy() + return self + + def fast_forward(self, n: IntNumber) -> "Sobol": + """Fast-forward the sequence by `n` positions. + + Parameters + ---------- + n : int + Number of points to skip in the sequence. + + Returns + ------- + engine : Sobol + The fast-forwarded engine. + + """ + if self.num_generated == 0: + _fast_forward( + n=n - 1, num_gen=self.num_generated, dim=self.d, + sv=self._sv, quasi=self._quasi + ) + else: + _fast_forward( + n=n, num_gen=self.num_generated - 1, dim=self.d, + sv=self._sv, quasi=self._quasi + ) + self.num_generated += n + return self + + +class PoissonDisk(QMCEngine): + """Poisson disk sampling. + + Parameters + ---------- + d : int + Dimension of the parameter space. + radius : float + Minimal distance to keep between points when sampling new candidates. + hypersphere : {"volume", "surface"}, optional + Sampling strategy to generate potential candidates to be added in the + final sample. Default is "volume". + + * ``volume``: original Bridson algorithm as described in [1]_. + New candidates are sampled *within* the hypersphere. + * ``surface``: only sample the surface of the hypersphere. + ncandidates : int + Number of candidates to sample per iteration. More candidates result + in a denser sampling as more candidates can be accepted per iteration. + optimization : {None, "random-cd", "lloyd"}, optional + Whether to use an optimization scheme to improve the quality after + sampling. Note that this is a post-processing step that does not + guarantee that all properties of the sample will be conserved. + Default is None. + + * ``random-cd``: random permutations of coordinates to lower the + centered discrepancy. The best sample based on the centered + discrepancy is constantly updated. Centered discrepancy-based + sampling shows better space-filling robustness toward 2D and 3D + subprojections compared to using other discrepancy measures. + * ``lloyd``: Perturb samples using a modified Lloyd-Max algorithm. + The process converges to equally spaced samples. + + .. versionadded:: 1.10.0 + + rng : `numpy.random.Generator`, optional + Pseudorandom number generator state. When `rng` is None, a new + `numpy.random.Generator` is created using entropy from the + operating system. Types other than `numpy.random.Generator` are + passed to `numpy.random.default_rng` to instantiate a ``Generator``. + + .. versionchanged:: 1.15.0 + + As part of the `SPEC-007 `_ + transition from use of `numpy.random.RandomState` to + `numpy.random.Generator`, this keyword was changed from `seed` to + `rng`. For an interim period, both keywords will continue to work, although + only one may be specified at a time. After the interim period, function + calls using the `seed` keyword will emit warnings. Following a + deprecation period, the `seed` keyword will be removed. + + l_bounds, u_bounds : array_like (d,) + Lower and upper bounds of target sample data. + + Notes + ----- + Poisson disk sampling is an iterative sampling strategy. Starting from + a seed sample, `ncandidates` are sampled in the hypersphere + surrounding the seed. Candidates below a certain `radius` or outside the + domain are rejected. New samples are added in a pool of sample seed. The + process stops when the pool is empty or when the number of required + samples is reached. + + The maximum number of point that a sample can contain is directly linked + to the `radius`. As the dimension of the space increases, a higher radius + spreads the points further and help overcome the curse of dimensionality. + See the :ref:`quasi monte carlo tutorial ` for more + details. + + .. warning:: + + The algorithm is more suitable for low dimensions and sampling size + due to its iterative nature and memory requirements. + Selecting a small radius with a high dimension would + mean that the space could contain more samples than using lower + dimension or a bigger radius. + + Some code taken from [2]_, written consent given on 31.03.2021 + by the original author, Shamis, for free use in SciPy under + the 3-clause BSD. + + References + ---------- + .. [1] Robert Bridson, "Fast Poisson Disk Sampling in Arbitrary + Dimensions." SIGGRAPH, 2007. + .. [2] `StackOverflow `__. + + Examples + -------- + Generate a 2D sample using a `radius` of 0.2. + + >>> import numpy as np + >>> import matplotlib.pyplot as plt + >>> from matplotlib.collections import PatchCollection + >>> from scipy.stats import qmc + >>> + >>> rng = np.random.default_rng() + >>> radius = 0.2 + >>> engine = qmc.PoissonDisk(d=2, radius=radius, rng=rng) + >>> sample = engine.random(20) + + Visualizing the 2D sample and showing that no points are closer than + `radius`. ``radius/2`` is used to visualize non-intersecting circles. + If two samples are exactly at `radius` from each other, then their circle + of radius ``radius/2`` will touch. + + >>> fig, ax = plt.subplots() + >>> _ = ax.scatter(sample[:, 0], sample[:, 1]) + >>> circles = [plt.Circle((xi, yi), radius=radius/2, fill=False) + ... for xi, yi in sample] + >>> collection = PatchCollection(circles, match_original=True) + >>> ax.add_collection(collection) + >>> _ = ax.set(aspect='equal', xlabel=r'$x_1$', ylabel=r'$x_2$', + ... xlim=[0, 1], ylim=[0, 1]) + >>> plt.show() + + Such visualization can be seen as circle packing: how many circle can + we put in the space. It is an np-hard problem. The method `fill_space` + can be used to add samples until no more samples can be added. This is + a hard problem and parameters may need to be adjusted manually. Beware of + the dimension: as the dimensionality increases, the number of samples + required to fill the space increases exponentially + (curse-of-dimensionality). + + """ + + @_transition_to_rng('seed', replace_doc=False) + def __init__( + self, + d: IntNumber, + *, + radius: DecimalNumber = 0.05, + hypersphere: Literal["volume", "surface"] = "volume", + ncandidates: IntNumber = 30, + optimization: Literal["random-cd", "lloyd"] | None = None, + rng: SeedType = None, + l_bounds: "npt.ArrayLike | None" = None, + u_bounds: "npt.ArrayLike | None" = None, + ) -> None: + # Used in `scipy.integrate.qmc_quad` + self._init_quad = {'d': d, 'radius': radius, + 'hypersphere': hypersphere, + 'ncandidates': ncandidates, + 'optimization': optimization} + super()._initialize(d=d, optimization=optimization, rng=rng) + + hypersphere_sample = { + "volume": self._hypersphere_volume_sample, + "surface": self._hypersphere_surface_sample + } + + try: + self.hypersphere_method = hypersphere_sample[hypersphere] + except KeyError as exc: + message = ( + f"{hypersphere!r} is not a valid hypersphere sampling" + f" method. It must be one of {set(hypersphere_sample)!r}") + raise ValueError(message) from exc + + # size of the sphere from which the samples are drawn relative to the + # size of a disk (radius) + # for the surface sampler, all new points are almost exactly 1 radius + # away from at least one existing sample +eps to avoid rejection + self.radius_factor = 2 if hypersphere == "volume" else 1.001 + self.radius = radius + self.radius_squared = self.radius**2 + + # sample to generate per iteration in the hypersphere around center + self.ncandidates = ncandidates + + if u_bounds is None: + u_bounds = np.ones(d) + if l_bounds is None: + l_bounds = np.zeros(d) + self.l_bounds, self.u_bounds = _validate_bounds( + l_bounds=l_bounds, u_bounds=u_bounds, d=int(d) + ) + + with np.errstate(divide='ignore'): + self.cell_size = self.radius / np.sqrt(self.d) + self.grid_size = ( + np.ceil((self.u_bounds - self.l_bounds) / self.cell_size) + ).astype(int) + + self._initialize_grid_pool() + + def _initialize_grid_pool(self): + """Sampling pool and sample grid.""" + self.sample_pool = [] + # Positions of cells + # n-dim value for each grid cell + self.sample_grid = np.empty( + np.append(self.grid_size, self.d), + dtype=np.float32 + ) + # Initialise empty cells with NaNs + self.sample_grid.fill(np.nan) + + def _random( + self, n: IntNumber = 1, *, workers: IntNumber = 1 + ) -> np.ndarray: + """Draw `n` in the interval ``[l_bounds, u_bounds]``. + + Note that it can return fewer samples if the space is full. + See the note section of the class. + + Parameters + ---------- + n : int, optional + Number of samples to generate in the parameter space. Default is 1. + + Returns + ------- + sample : array_like (n, d) + QMC sample. + + """ + if n == 0 or self.d == 0: + return np.empty((n, self.d)) + + def in_limits(sample: np.ndarray) -> bool: + for i in range(self.d): + if (sample[i] > self.u_bounds[i] or sample[i] < self.l_bounds[i]): + return False + return True + + def in_neighborhood(candidate: np.ndarray, n: int = 2) -> bool: + """ + Check if there are samples closer than ``radius_squared`` to the + `candidate` sample. + """ + indices = ((candidate - self.l_bounds) / self.cell_size).astype(int) + ind_min = np.maximum(indices - n, 0) + ind_max = np.minimum(indices + n + 1, self.grid_size) + + # Check if the center cell is empty + if not np.isnan(self.sample_grid[tuple(indices)][0]): + return True + + a = [slice(ind_min[i], ind_max[i]) for i in range(self.d)] + + # guards against: invalid value encountered in less as we are + # comparing with nan and returns False. Which is wanted. + with np.errstate(invalid='ignore'): + if np.any( + np.sum( + np.square(candidate - self.sample_grid[tuple(a)]), + axis=self.d + ) < self.radius_squared + ): + return True + + return False + + def add_sample(candidate: np.ndarray) -> None: + self.sample_pool.append(candidate) + indices = ((candidate - self.l_bounds) / self.cell_size).astype(int) + self.sample_grid[tuple(indices)] = candidate + curr_sample.append(candidate) + + curr_sample: list[np.ndarray] = [] + + if len(self.sample_pool) == 0: + # the pool is being initialized with a single random sample + add_sample(self.rng.uniform(self.l_bounds, self.u_bounds)) + num_drawn = 1 + else: + num_drawn = 0 + + # exhaust sample pool to have up to n sample + while len(self.sample_pool) and num_drawn < n: + # select a sample from the available pool + idx_center = rng_integers(self.rng, len(self.sample_pool)) + center = self.sample_pool[idx_center] + del self.sample_pool[idx_center] + + # generate candidates around the center sample + candidates = self.hypersphere_method( + center, self.radius * self.radius_factor, self.ncandidates + ) + + # keep candidates that satisfy some conditions + for candidate in candidates: + if in_limits(candidate) and not in_neighborhood(candidate): + add_sample(candidate) + + num_drawn += 1 + if num_drawn >= n: + break + + self.num_generated += num_drawn + return np.array(curr_sample) + + def fill_space(self) -> np.ndarray: + """Draw ``n`` samples in the interval ``[l_bounds, u_bounds]``. + + Unlike `random`, this method will try to add points until + the space is full. Depending on ``candidates`` (and to a lesser extent + other parameters), some empty areas can still be present in the sample. + + .. warning:: + + This can be extremely slow in high dimensions or if the + ``radius`` is very small-with respect to the dimensionality. + + Returns + ------- + sample : array_like (n, d) + QMC sample. + + """ + return self.random(np.inf) # type: ignore[arg-type] + + def reset(self) -> "PoissonDisk": + """Reset the engine to base state. + + Returns + ------- + engine : PoissonDisk + Engine reset to its base state. + + """ + super().reset() + self._initialize_grid_pool() + return self + + def _hypersphere_volume_sample( + self, center: np.ndarray, radius: DecimalNumber, + candidates: IntNumber = 1 + ) -> np.ndarray: + """Uniform sampling within hypersphere.""" + # should remove samples within r/2 + x = self.rng.standard_normal(size=(candidates, self.d)) + ssq = np.sum(x**2, axis=1) + fr = radius * gammainc(self.d/2, ssq/2)**(1/self.d) / np.sqrt(ssq) + fr_tiled = np.tile(fr.reshape(-1, 1), (1, int(self.d))) + p = center + np.multiply(x, fr_tiled) + return p + + def _hypersphere_surface_sample( + self, center: np.ndarray, radius: DecimalNumber, + candidates: IntNumber = 1 + ) -> np.ndarray: + """Uniform sampling on the hypersphere's surface.""" + vec = self.rng.standard_normal(size=(candidates, self.d)) + vec /= np.linalg.norm(vec, axis=1)[:, None] + p = center + np.multiply(vec, radius) + return p + + +class MultivariateNormalQMC: + r"""QMC sampling from a multivariate Normal :math:`N(\mu, \Sigma)`. + + Parameters + ---------- + mean : array_like (d,) + The mean vector. Where ``d`` is the dimension. + cov : array_like (d, d), optional + The covariance matrix. If omitted, use `cov_root` instead. + If both `cov` and `cov_root` are omitted, use the identity matrix. + cov_root : array_like (d, d'), optional + A root decomposition of the covariance matrix, where ``d'`` may be less + than ``d`` if the covariance is not full rank. If omitted, use `cov`. + inv_transform : bool, optional + If True, use inverse transform instead of Box-Muller. Default is True. + engine : QMCEngine, optional + Quasi-Monte Carlo engine sampler. If None, `Sobol` is used. + rng : `numpy.random.Generator`, optional + Pseudorandom number generator state. When `rng` is None, a new + `numpy.random.Generator` is created using entropy from the + operating system. Types other than `numpy.random.Generator` are + passed to `numpy.random.default_rng` to instantiate a ``Generator``. + + .. versionchanged:: 1.15.0 + + As part of the `SPEC-007 `_ + transition from use of `numpy.random.RandomState` to + `numpy.random.Generator`, this keyword was changed from `seed` to + `rng`. For an interim period, both keywords will continue to work, although + only one may be specified at a time. After the interim period, function + calls using the `seed` keyword will emit warnings. Following a + deprecation period, the `seed` keyword will be removed. + + Examples + -------- + >>> import matplotlib.pyplot as plt + >>> from scipy.stats import qmc + >>> dist = qmc.MultivariateNormalQMC(mean=[0, 5], cov=[[1, 0], [0, 1]]) + >>> sample = dist.random(512) + >>> _ = plt.scatter(sample[:, 0], sample[:, 1]) + >>> plt.show() + + """ + + @_transition_to_rng('seed', replace_doc=False) + def __init__( + self, + mean: "npt.ArrayLike", + cov: "npt.ArrayLike | None" = None, + *, + cov_root: "npt.ArrayLike | None" = None, + inv_transform: bool = True, + engine: QMCEngine | None = None, + rng: SeedType = None, + ) -> None: + mean = np.asarray(np.atleast_1d(mean)) + d = mean.shape[0] + if cov is not None: + # covariance matrix provided + cov = np.asarray(np.atleast_2d(cov)) + # check for square/symmetric cov matrix and mean vector has the + # same d + if not mean.shape[0] == cov.shape[0]: + raise ValueError("Dimension mismatch between mean and " + "covariance.") + if not np.allclose(cov, cov.transpose()): + raise ValueError("Covariance matrix is not symmetric.") + # compute Cholesky decomp; if it fails, do the eigen decomposition + try: + cov_root = np.linalg.cholesky(cov).transpose() + except np.linalg.LinAlgError: + eigval, eigvec = np.linalg.eigh(cov) + if not np.all(eigval >= -1.0e-8): + raise ValueError("Covariance matrix not PSD.") + eigval = np.clip(eigval, 0.0, None) + cov_root = (eigvec * np.sqrt(eigval)).transpose() + elif cov_root is not None: + # root decomposition provided + cov_root = np.atleast_2d(cov_root) + if not mean.shape[0] == cov_root.shape[0]: + raise ValueError("Dimension mismatch between mean and " + "covariance.") + else: + # corresponds to identity covariance matrix + cov_root = None + + self._inv_transform = inv_transform + + if not inv_transform: + # to apply Box-Muller, we need an even number of dimensions + engine_dim = 2 * math.ceil(d / 2) + else: + engine_dim = d + if engine is None: + # Need this during SPEC 7 transition to prevent `RandomState` + # from being passed via `rng`. + kwarg = "seed" if isinstance(rng, np.random.RandomState) else "rng" + kwargs = {kwarg: rng} + self.engine: QMCEngine = Sobol( + d=engine_dim, scramble=True, bits=30, **kwargs + ) + elif isinstance(engine, QMCEngine): + if engine.d != engine_dim: + raise ValueError("Dimension of `engine` must be consistent" + " with dimensions of mean and covariance." + " If `inv_transform` is False, it must be" + " an even number.") + self.engine = engine + else: + raise ValueError("`engine` must be an instance of " + "`scipy.stats.qmc.QMCEngine` or `None`.") + + self._mean = mean + self._corr_matrix = cov_root + + self._d = d + + def random(self, n: IntNumber = 1) -> np.ndarray: + """Draw `n` QMC samples from the multivariate Normal. + + Parameters + ---------- + n : int, optional + Number of samples to generate in the parameter space. Default is 1. + + Returns + ------- + sample : array_like (n, d) + Sample. + + """ + base_samples = self._standard_normal_samples(n) + return self._correlate(base_samples) + + def _correlate(self, base_samples: np.ndarray) -> np.ndarray: + if self._corr_matrix is not None: + return base_samples @ self._corr_matrix + self._mean + else: + # avoid multiplying with identity here + return base_samples + self._mean + + def _standard_normal_samples(self, n: IntNumber = 1) -> np.ndarray: + """Draw `n` QMC samples from the standard Normal :math:`N(0, I_d)`. + + Parameters + ---------- + n : int, optional + Number of samples to generate in the parameter space. Default is 1. + + Returns + ------- + sample : array_like (n, d) + Sample. + + """ + # get base samples + samples = self.engine.random(n) + if self._inv_transform: + # apply inverse transform + # (values to close to 0/1 result in inf values) + return stats.norm.ppf(0.5 + (1 - 1e-10) * (samples - 0.5)) # type: ignore[attr-defined] # noqa: E501 + else: + # apply Box-Muller transform (note: indexes starting from 1) + even = np.arange(0, samples.shape[-1], 2) + Rs = np.sqrt(-2 * np.log(samples[:, even])) + thetas = 2 * math.pi * samples[:, 1 + even] + cos = np.cos(thetas) + sin = np.sin(thetas) + transf_samples = np.stack([Rs * cos, Rs * sin], + -1).reshape(n, -1) + # make sure we only return the number of dimension requested + return transf_samples[:, : self._d] + + +class MultinomialQMC: + r"""QMC sampling from a multinomial distribution. + + Parameters + ---------- + pvals : array_like (k,) + Vector of probabilities of size ``k``, where ``k`` is the number + of categories. Elements must be non-negative and sum to 1. + n_trials : int + Number of trials. + engine : QMCEngine, optional + Quasi-Monte Carlo engine sampler. If None, `Sobol` is used. + rng : `numpy.random.Generator`, optional + Pseudorandom number generator state. When `rng` is None, a new + `numpy.random.Generator` is created using entropy from the + operating system. Types other than `numpy.random.Generator` are + passed to `numpy.random.default_rng` to instantiate a ``Generator``. + + .. versionchanged:: 1.15.0 + + As part of the `SPEC-007 `_ + transition from use of `numpy.random.RandomState` to + `numpy.random.Generator`, this keyword was changed from `seed` to + `rng`. For an interim period, both keywords will continue to work, although + only one may be specified at a time. After the interim period, function + calls using the `seed` keyword will emit warnings. Following a + deprecation period, the `seed` keyword will be removed. + + Examples + -------- + Let's define 3 categories and for a given sample, the sum of the trials + of each category is 8. The number of trials per category is determined + by the `pvals` associated to each category. + Then, we sample this distribution 64 times. + + >>> import matplotlib.pyplot as plt + >>> from scipy.stats import qmc + >>> dist = qmc.MultinomialQMC( + ... pvals=[0.2, 0.4, 0.4], n_trials=10, engine=qmc.Halton(d=1) + ... ) + >>> sample = dist.random(64) + + We can plot the sample and verify that the median of number of trials + for each category is following the `pvals`. That would be + ``pvals * n_trials = [2, 4, 4]``. + + >>> fig, ax = plt.subplots() + >>> ax.yaxis.get_major_locator().set_params(integer=True) + >>> _ = ax.boxplot(sample) + >>> ax.set(xlabel="Categories", ylabel="Trials") + >>> plt.show() + + """ + + @_transition_to_rng('seed', replace_doc=False) + def __init__( + self, + pvals: "npt.ArrayLike", + n_trials: IntNumber, + *, + engine: QMCEngine | None = None, + rng: SeedType = None, + ) -> None: + self.pvals = np.atleast_1d(np.asarray(pvals)) + if np.min(pvals) < 0: + raise ValueError('Elements of pvals must be non-negative.') + if not np.isclose(np.sum(pvals), 1): + raise ValueError('Elements of pvals must sum to 1.') + self.n_trials = n_trials + if engine is None: + # Need this during SPEC 7 transition to prevent `RandomState` + # from being passed via `rng`. + kwarg = "seed" if isinstance(rng, np.random.RandomState) else "rng" + kwargs = {kwarg: rng} + self.engine: QMCEngine = Sobol(d=1, scramble=True, bits=30, **kwargs) + elif isinstance(engine, QMCEngine): + if engine.d != 1: + raise ValueError("Dimension of `engine` must be 1.") + self.engine = engine + else: + raise ValueError("`engine` must be an instance of " + "`scipy.stats.qmc.QMCEngine` or `None`.") + + def random(self, n: IntNumber = 1) -> np.ndarray: + """Draw `n` QMC samples from the multinomial distribution. + + Parameters + ---------- + n : int, optional + Number of samples to generate in the parameter space. Default is 1. + + Returns + ------- + samples : array_like (n, pvals) + Sample. + + """ + sample = np.empty((n, len(self.pvals))) + for i in range(n): + base_draws = self.engine.random(self.n_trials).ravel() + p_cumulative = np.empty_like(self.pvals, dtype=float) + _fill_p_cumulative(np.array(self.pvals, dtype=float), p_cumulative) + sample_ = np.zeros_like(self.pvals, dtype=np.intp) + _categorize(base_draws, p_cumulative, sample_) + sample[i] = sample_ + return sample + + +def _select_optimizer( + optimization: Literal["random-cd", "lloyd"] | None, config: dict +) -> Callable | None: + """A factory for optimization methods.""" + optimization_method: dict[str, Callable] = { + "random-cd": _random_cd, + "lloyd": _lloyd_centroidal_voronoi_tessellation + } + + optimizer: partial | None + if optimization is not None: + try: + optimizer_ = optimization_method[optimization.lower()] + except KeyError as exc: + message = (f"{optimization!r} is not a valid optimization" + f" method. It must be one of" + f" {set(optimization_method)!r}") + raise ValueError(message) from exc + + # config + optimizer = partial(optimizer_, **config) + else: + optimizer = None + + return optimizer + + +def _random_cd( + best_sample: np.ndarray, n_iters: int, n_nochange: int, rng: GeneratorType, + **kwargs: dict +) -> np.ndarray: + """Optimal LHS on CD. + + Create a base LHS and do random permutations of coordinates to + lower the centered discrepancy. + Because it starts with a normal LHS, it also works with the + `scramble` keyword argument. + + Two stopping criterion are used to stop the algorithm: at most, + `n_iters` iterations are performed; or if there is no improvement + for `n_nochange` consecutive iterations. + """ + del kwargs # only use keywords which are defined, needed by factory + + n, d = best_sample.shape + + if d == 0 or n == 0: + return np.empty((n, d)) + + if d == 1 or n == 1: + # discrepancy measures are invariant under permuting factors and runs + return best_sample + + best_disc = discrepancy(best_sample) + + bounds = ([0, d - 1], + [0, n - 1], + [0, n - 1]) + + n_nochange_ = 0 + n_iters_ = 0 + while n_nochange_ < n_nochange and n_iters_ < n_iters: + n_iters_ += 1 + + col = rng_integers(rng, *bounds[0], endpoint=True) # type: ignore[misc] + row_1 = rng_integers(rng, *bounds[1], endpoint=True) # type: ignore[misc] + row_2 = rng_integers(rng, *bounds[2], endpoint=True) # type: ignore[misc] + disc = _perturb_discrepancy(best_sample, + row_1, row_2, col, + best_disc) + if disc < best_disc: + best_sample[row_1, col], best_sample[row_2, col] = ( + best_sample[row_2, col], best_sample[row_1, col]) + + best_disc = disc + n_nochange_ = 0 + else: + n_nochange_ += 1 + + return best_sample + + +def _l1_norm(sample: np.ndarray) -> float: + return distance.pdist(sample, 'cityblock').min() + + +def _lloyd_iteration( + sample: np.ndarray, + decay: float, + qhull_options: str +) -> np.ndarray: + """Lloyd-Max algorithm iteration. + + Based on the implementation of Stéfan van der Walt: + + https://github.com/stefanv/lloyd + + which is: + + Copyright (c) 2021-04-21 Stéfan van der Walt + https://github.com/stefanv/lloyd + MIT License + + Parameters + ---------- + sample : array_like (n, d) + The sample to iterate on. + decay : float + Relaxation decay. A positive value would move the samples toward + their centroid, and negative value would move them away. + 1 would move the samples to their centroid. + qhull_options : str + Additional options to pass to Qhull. See Qhull manual + for details. (Default: "Qbb Qc Qz Qj Qx" for ndim > 4 and + "Qbb Qc Qz Qj" otherwise.) + + Returns + ------- + sample : array_like (n, d) + The sample after an iteration of Lloyd's algorithm. + + """ + new_sample = np.empty_like(sample) + + voronoi = Voronoi(sample, qhull_options=qhull_options) + + for ii, idx in enumerate(voronoi.point_region): + # the region is a series of indices into self.voronoi.vertices + # remove samples at infinity, designated by index -1 + region = [i for i in voronoi.regions[idx] if i != -1] + + # get the vertices for this region + verts = voronoi.vertices[region] + + # clipping would be wrong, we need to intersect + # verts = np.clip(verts, 0, 1) + + # move samples towards centroids: + # Centroid in n-D is the mean for uniformly distributed nodes + # of a geometry. + centroid = np.mean(verts, axis=0) + new_sample[ii] = sample[ii] + (centroid - sample[ii]) * decay + + # only update sample to centroid within the region + is_valid = np.all(np.logical_and(new_sample >= 0, new_sample <= 1), axis=1) + sample[is_valid] = new_sample[is_valid] + + return sample + + +def _lloyd_centroidal_voronoi_tessellation( + sample: "npt.ArrayLike", + *, + tol: DecimalNumber = 1e-5, + maxiter: IntNumber = 10, + qhull_options: str | None = None, + **kwargs: dict +) -> np.ndarray: + """Approximate Centroidal Voronoi Tessellation. + + Perturb samples in N-dimensions using Lloyd-Max algorithm. + + Parameters + ---------- + sample : array_like (n, d) + The sample to iterate on. With ``n`` the number of samples and ``d`` + the dimension. Samples must be in :math:`[0, 1]^d`, with ``d>=2``. + tol : float, optional + Tolerance for termination. If the min of the L1-norm over the samples + changes less than `tol`, it stops the algorithm. Default is 1e-5. + maxiter : int, optional + Maximum number of iterations. It will stop the algorithm even if + `tol` is above the threshold. + Too many iterations tend to cluster the samples as a hypersphere. + Default is 10. + qhull_options : str, optional + Additional options to pass to Qhull. See Qhull manual + for details. (Default: "Qbb Qc Qz Qj Qx" for ndim > 4 and + "Qbb Qc Qz Qj" otherwise.) + + Returns + ------- + sample : array_like (n, d) + The sample after being processed by Lloyd-Max algorithm. + + Notes + ----- + Lloyd-Max algorithm is an iterative process with the purpose of improving + the dispersion of samples. For given sample: (i) compute a Voronoi + Tessellation; (ii) find the centroid of each Voronoi cell; (iii) move the + samples toward the centroid of their respective cell. See [1]_, [2]_. + + A relaxation factor is used to control how fast samples can move at each + iteration. This factor is starting at 2 and ending at 1 after `maxiter` + following an exponential decay. + + The process converges to equally spaced samples. It implies that measures + like the discrepancy could suffer from too many iterations. On the other + hand, L1 and L2 distances should improve. This is especially true with + QMC methods which tend to favor the discrepancy over other criteria. + + .. note:: + + The current implementation does not intersect the Voronoi Tessellation + with the boundaries. This implies that for a low number of samples, + empirically below 20, no Voronoi cell is touching the boundaries. + Hence, samples cannot be moved close to the boundaries. + + Further improvements could consider the samples at infinity so that + all boundaries are segments of some Voronoi cells. This would fix + the computation of the centroid position. + + .. warning:: + + The Voronoi Tessellation step is expensive and quickly becomes + intractable with dimensions as low as 10 even for a sample + of size as low as 1000. + + .. versionadded:: 1.9.0 + + References + ---------- + .. [1] Lloyd. "Least Squares Quantization in PCM". + IEEE Transactions on Information Theory, 1982. + .. [2] Max J. "Quantizing for minimum distortion". + IEEE Transactions on Information Theory, 1960. + + Examples + -------- + >>> import numpy as np + >>> from scipy.spatial import distance + >>> from scipy.stats._qmc import _lloyd_centroidal_voronoi_tessellation + >>> rng = np.random.default_rng() + >>> sample = rng.random((128, 2)) + + .. note:: + + The samples need to be in :math:`[0, 1]^d`. `scipy.stats.qmc.scale` + can be used to scale the samples from their + original bounds to :math:`[0, 1]^d`. And back to their original bounds. + + Compute the quality of the sample using the L1 criterion. + + >>> def l1_norm(sample): + ... return distance.pdist(sample, 'cityblock').min() + + >>> l1_norm(sample) + 0.00161... # random + + Now process the sample using Lloyd's algorithm and check the improvement + on the L1. The value should increase. + + >>> sample = _lloyd_centroidal_voronoi_tessellation(sample) + >>> l1_norm(sample) + 0.0278... # random + + """ + del kwargs # only use keywords which are defined, needed by factory + + sample = np.asarray(sample).copy() + + if not sample.ndim == 2: + raise ValueError('`sample` is not a 2D array') + + if not sample.shape[1] >= 2: + raise ValueError('`sample` dimension is not >= 2') + + # Checking that sample is within the hypercube + if (sample.max() > 1.) or (sample.min() < 0.): + raise ValueError('`sample` is not in unit hypercube') + + if qhull_options is None: + qhull_options = 'Qbb Qc Qz QJ' + + if sample.shape[1] >= 5: + qhull_options += ' Qx' + + # Fit an exponential to be 2 at 0 and 1 at `maxiter`. + # The decay is used for relaxation. + # analytical solution for y=exp(-maxiter/x) - 0.1 + root = -maxiter / np.log(0.1) + decay = [np.exp(-x / root)+0.9 for x in range(maxiter)] + + l1_old = _l1_norm(sample=sample) + for i in range(maxiter): + sample = _lloyd_iteration( + sample=sample, decay=decay[i], + qhull_options=qhull_options, + ) + + l1_new = _l1_norm(sample=sample) + + if abs(l1_new - l1_old) < tol: + break + else: + l1_old = l1_new + + return sample + + +def _validate_workers(workers: IntNumber = 1) -> IntNumber: + """Validate `workers` based on platform and value. + + Parameters + ---------- + workers : int, optional + Number of workers to use for parallel processing. If -1 is + given all CPU threads are used. Default is 1. + + Returns + ------- + Workers : int + Number of CPU used by the algorithm + + """ + workers = int(workers) + if workers == -1: + workers = os.cpu_count() # type: ignore[assignment] + if workers is None: + raise NotImplementedError( + "Cannot determine the number of cpus using os.cpu_count(), " + "cannot use -1 for the number of workers" + ) + elif workers <= 0: + raise ValueError(f"Invalid number of workers: {workers}, must be -1 " + "or > 0") + + return workers + + +def _validate_bounds( + l_bounds: "npt.ArrayLike", u_bounds: "npt.ArrayLike", d: int +) -> tuple[np.ndarray, np.ndarray]: + """Bounds input validation. + + Parameters + ---------- + l_bounds, u_bounds : array_like (d,) + Lower and upper bounds. + d : int + Dimension to use for broadcasting. + + Returns + ------- + l_bounds, u_bounds : array_like (d,) + Lower and upper bounds. + + """ + try: + lower = np.broadcast_to(l_bounds, d) + upper = np.broadcast_to(u_bounds, d) + except ValueError as exc: + msg = ("'l_bounds' and 'u_bounds' must be broadcastable and respect" + " the sample dimension") + raise ValueError(msg) from exc + + if not np.all(lower < upper): + raise ValueError("Bounds are not consistent 'l_bounds' < 'u_bounds'") + + return lower, upper diff --git a/v12x/pylib/scipy/stats/_qmvnt.py b/v12x/pylib/scipy/stats/_qmvnt.py new file mode 100644 index 0000000000000000000000000000000000000000..9f942da89e0a03c456b4ee0011b711ba217b4024 --- /dev/null +++ b/v12x/pylib/scipy/stats/_qmvnt.py @@ -0,0 +1,475 @@ +# Integration of multivariate normal and t distributions. + +# Adapted from the MATLAB original implementations by Dr. Alan Genz. + +# http://www.math.wsu.edu/faculty/genz/software/software.html + +# Copyright (C) 2013, Alan Genz, All rights reserved. +# Python implementation is copyright (C) 2022, Robert Kern, All rights +# reserved. + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided the following conditions are met: +# 1. Redistributions of source code must retain the above copyright +# notice, this list of conditions and the following disclaimer. +# 2. Redistributions in binary form must reproduce the above copyright +# notice, this list of conditions and the following disclaimer in +# the documentation and/or other materials provided with the +# distribution. +# 3. The contributor name(s) may not be used to endorse or promote +# products derived from this software without specific prior +# written permission. +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS +# FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE +# COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, +# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, +# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS +# OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND +# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR +# TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + + +import math +import numpy as np + +from scipy.fft import fft, ifft +from scipy.special import ndtr as phi, ndtri as phinv +from scipy.special._ufuncs import _bivariate_normal_cdf as _bvnu +from scipy.stats._qmc import primes_from_2_to + +from ._qmvnt_cy import _qmvn_inner, _qmvt_inner + + +def _factorize_int(n): + """Return a sorted list of the unique prime factors of a positive integer. + """ + # NOTE: There are lots faster ways to do this, but this isn't terrible. + factors = set() + for p in primes_from_2_to(int(np.sqrt(n)) + 1): + while not (n % p): + factors.add(p) + n //= p + if n == 1: + break + if n != 1: + factors.add(n) + return sorted(factors) + + +def _primitive_root(p): + """Compute a primitive root of the prime number `p`. + + Used in the CBC lattice construction. + + References + ---------- + .. [1] https://en.wikipedia.org/wiki/Primitive_root_modulo_n + """ + # p is prime + pm = p - 1 + factors = _factorize_int(pm) + n = len(factors) + r = 2 + k = 0 + while k < n: + d = pm // factors[k] + # pow() doesn't like numpy scalar types. + rd = pow(int(r), int(d), int(p)) + if rd == 1: + r += 1 + k = 0 + else: + k += 1 + return r + + +def _cbc_lattice(n_dim, n_qmc_samples): + """Compute a QMC lattice generator using a Fast CBC construction. + + Parameters + ---------- + n_dim : int > 0 + The number of dimensions for the lattice. + n_qmc_samples : int > 0 + The desired number of QMC samples. This will be rounded down to the + nearest prime to enable the CBC construction. + + Returns + ------- + q : float array : shape=(n_dim,) + The lattice generator vector. All values are in the open interval + ``(0, 1)``. + actual_n_qmc_samples : int + The prime number of QMC samples that must be used with this lattice, + no more, no less. + + References + ---------- + .. [1] Nuyens, D. and Cools, R. "Fast Component-by-Component Construction, + a Reprise for Different Kernels", In H. Niederreiter and D. Talay, + editors, Monte-Carlo and Quasi-Monte Carlo Methods 2004, + Springer-Verlag, 2006, 371-385. + """ + # Round down to the nearest prime number. + primes = primes_from_2_to(n_qmc_samples + 1) + n_qmc_samples = primes[-1] + + bt = np.ones(n_dim) + gm = np.hstack([1.0, 0.8 ** np.arange(n_dim - 1)]) + q = 1 + w = 0 + z = np.arange(1, n_dim + 1) + m = (n_qmc_samples - 1) // 2 + g = _primitive_root(n_qmc_samples) + # Slightly faster way to compute perm[j] = pow(g, j, n_qmc_samples) + # Shame that we don't have modulo pow() implemented as a ufunc. + perm = np.ones(m, dtype=int) + for j in range(m - 1): + perm[j + 1] = (g * perm[j]) % n_qmc_samples + perm = np.minimum(n_qmc_samples - perm, perm) + pn = perm / n_qmc_samples + c = pn * pn - pn + 1.0 / 6 + fc = fft(c) + for s in range(1, n_dim): + reordered = np.hstack([ + c[:w+1][::-1], + c[w+1:m][::-1], + ]) + q = q * (bt[s-1] + gm[s-1] * reordered) + w = ifft(fc * fft(q)).real.argmin() + z[s] = perm[w] + q = z / n_qmc_samples + return q, n_qmc_samples + + +def _qauto(func, covar, low, high, rng, error=1e-3, limit=10_000, **kwds): + """Automatically rerun the integration to get the required error bound. + + Parameters + ---------- + func : callable + Either :func:`_qmvn` or :func:`_qmvt`. + covar, low, high : array + As specified in :func:`_qmvn` and :func:`_qmvt`. + rng : Generator, optional + default_rng(), yada, yada + error : float > 0 + The desired error bound. + limit : int > 0: + The rough limit of the number of integration points to consider. The + integration will stop looping once this limit has been *exceeded*. + **kwds : + Other keyword arguments to pass to `func`. When using :func:`_qmvt`, be + sure to include ``nu=`` as one of these. + + Returns + ------- + prob : float + The estimated probability mass within the bounds. + est_error : float + 3 times the standard error of the batch estimates. + n_samples : int + The number of integration points actually used. + """ + n = len(covar) + n_samples = 0 + if n == 1: + prob = phi(high / covar**0.5) - phi(low / covar**0.5) + # More or less + est_error = 1e-15 + elif n == 2: + prob = _bvn(low, high, covar) + est_error = 1e-15 + else: + mi = min(limit, n * 1000) + prob = 0.0 + est_error = 1.0 + ei = 0.0 + while est_error > error and n_samples < limit: + mi = round(np.sqrt(2) * mi) + pi, ei, ni = func(mi, covar, low, high, rng=rng, **kwds) + n_samples += ni + wt = 1.0 / (1 + (ei / est_error)**2) + prob += wt * (pi - prob) + est_error = np.sqrt(wt) * ei + return prob, est_error, n_samples + + +def _qmvn(m, covar, low, high, rng, lattice='cbc', n_batches=10): + """Multivariate normal integration over box bounds. + + Parameters + ---------- + m : int > n_batches + The number of points to sample. This number will be divided into + `n_batches` batches that apply random offsets of the sampling lattice + for each batch in order to estimate the error. + covar : (n, n) float array + Possibly singular, positive semidefinite symmetric covariance matrix. + low, high : (n,) float array + The low and high integration bounds. + rng : Generator, optional + default_rng(), yada, yada + lattice : 'cbc' or callable + The type of lattice rule to use to construct the integration points. + n_batches : int > 0, optional + The number of QMC batches to apply. + + Returns + ------- + prob : float + The estimated probability mass within the bounds. + est_error : float + 3 times the standard error of the batch estimates. + """ + cho, lo, hi = _permuted_cholesky(covar, low, high) + if not cho.flags.c_contiguous: + # qmvn_inner expects contiguous buffers + cho = cho.copy() + + n = cho.shape[0] + q, n_qmc_samples = _cbc_lattice(n - 1, max(m // n_batches, 1)) + rndm = rng.random(size=(n_batches, n)) + + prob, est_error, n_samples = _qmvn_inner( + q, rndm, int(n_qmc_samples), int(n_batches), cho, lo, hi + ) + return prob, est_error, n_samples + + +# Note: this function is not currently used or tested by any SciPy code. It is +# included in this file to facilitate the resolution of gh-8367, gh-16142, and +# possibly gh-14286, but must be reviewed and tested before use. +def _mvn_qmc_integrand(covar, low, high, use_tent=False): + """Transform the multivariate normal integration into a QMC integrand over + a unit hypercube. + + The dimensionality of the resulting hypercube integration domain is one + less than the dimensionality of the original integrand. Note that this + transformation subsumes the integration bounds in order to account for + infinite bounds. The QMC integration one does with the returned integrand + should be on the unit hypercube. + + Parameters + ---------- + covar : (n, n) float array + Possibly singular, positive semidefinite symmetric covariance matrix. + low, high : (n,) float array + The low and high integration bounds. + use_tent : bool, optional + If True, then use tent periodization. Only helpful for lattice rules. + + Returns + ------- + integrand : Callable[[NDArray], NDArray] + The QMC-integrable integrand. It takes an + ``(n_qmc_samples, ndim_integrand)`` array of QMC samples in the unit + hypercube and returns the ``(n_qmc_samples,)`` evaluations of at these + QMC points. + ndim_integrand : int + The dimensionality of the integrand. Equal to ``n-1``. + """ + cho, lo, hi = _permuted_cholesky(covar, low, high) + n = cho.shape[0] + ndim_integrand = n - 1 + ct = cho[0, 0] + c = phi(lo[0] / ct) + d = phi(hi[0] / ct) + ci = c + dci = d - ci + + def integrand(*zs): + ndim_qmc = len(zs) + n_qmc_samples = len(np.atleast_1d(zs[0])) + assert ndim_qmc == ndim_integrand + y = np.zeros((ndim_qmc, n_qmc_samples)) + c = np.full(n_qmc_samples, ci) + dc = np.full(n_qmc_samples, dci) + pv = dc.copy() + for i in range(1, n): + if use_tent: + # Tent periodization transform. + x = abs(2 * zs[i-1] - 1) + else: + x = zs[i-1] + y[i - 1, :] = phinv(c + x * dc) + s = cho[i, :i] @ y[:i, :] + ct = cho[i, i] + c = phi((lo[i] - s) / ct) + d = phi((hi[i] - s) / ct) + dc = d - c + pv = pv * dc + return pv + + return integrand, ndim_integrand + + +def _qmvt(m, nu, covar, low, high, rng, lattice='cbc', n_batches=10): + """Multivariate t integration over box bounds. + + Parameters + ---------- + m : int > n_batches + The number of points to sample. This number will be divided into + `n_batches` batches that apply random offsets of the sampling lattice + for each batch in order to estimate the error. + nu : float >= 0 + The shape parameter of the multivariate t distribution. + covar : (n, n) float array + Possibly singular, positive semidefinite symmetric covariance matrix. + low, high : (n,) float array + The low and high integration bounds. + rng : Generator, optional + default_rng(), yada, yada + lattice : 'cbc' or callable + The type of lattice rule to use to construct the integration points. + n_batches : int > 0, optional + The number of QMC batches to apply. + + Returns + ------- + prob : float + The estimated probability mass within the bounds. + est_error : float + 3 times the standard error of the batch estimates. + n_samples : int + The number of samples actually used. + """ + sn = max(1.0, np.sqrt(nu)) + low = np.asarray(low, dtype=np.float64) + high = np.asarray(high, dtype=np.float64) + cho, lo, hi = _permuted_cholesky(covar, low / sn, high / sn) + n = cho.shape[0] + q, n_qmc_samples = _cbc_lattice(n, max(m // n_batches, 1)) + rndm = rng.random(size=(n_batches, n)) + prob, est_error, n_samples = _qmvt_inner( + q, rndm, int(n_qmc_samples), int(n_batches), cho, lo, hi, float(nu) + ) + return prob, est_error, n_samples + + +def _permuted_cholesky(covar, low, high, tol=1e-10): + """Compute a scaled, permuted Cholesky factor, with integration bounds. + + The scaling and permuting of the dimensions accomplishes part of the + transformation of the original integration problem into a more numerically + tractable form. The lower-triangular Cholesky factor will then be used in + the subsequent integration. The integration bounds will be scaled and + permuted as well. + + Parameters + ---------- + covar : (n, n) float array + Possibly singular, positive semidefinite symmetric covariance matrix. + low, high : (n,) float array + The low and high integration bounds. + tol : float, optional + The singularity tolerance. + + Returns + ------- + cho : (n, n) float array + Lower Cholesky factor, scaled and permuted. + new_low, new_high : (n,) float array + The scaled and permuted low and high integration bounds. + """ + # Make copies for outputting. + cho = np.array(covar, dtype=np.float64) + new_lo = np.array(low, dtype=np.float64) + new_hi = np.array(high, dtype=np.float64) + n = cho.shape[0] + if cho.shape != (n, n): + raise ValueError("expected a square symmetric array") + if new_lo.shape != (n,) or new_hi.shape != (n,): + raise ValueError( + "expected integration boundaries the same dimensions " + "as the covariance matrix" + ) + # Scale by the sqrt of the diagonal. + dc = np.sqrt(np.maximum(np.diag(cho), 0.0)) + # But don't divide by 0. + dc[dc == 0.0] = 1.0 + new_lo /= dc + new_hi /= dc + cho /= dc + cho /= dc[:, np.newaxis] + + y = np.zeros(n) + sqtp = np.sqrt(2 * np.pi) + for k in range(n): + epk = (k + 1) * tol + im = k + ck = 0.0 + dem = 1.0 + s = 0.0 + lo_m = 0.0 + hi_m = 0.0 + for i in range(k, n): + if cho[i, i] > tol: + ci = np.sqrt(cho[i, i]) + if i > 0: + s = cho[i, :k] @ y[:k] + lo_i = (new_lo[i] - s) / ci + hi_i = (new_hi[i] - s) / ci + de = phi(hi_i) - phi(lo_i) + if de <= dem: + ck = ci + dem = de + lo_m = lo_i + hi_m = hi_i + im = i + if im > k: + # Swap im and k + cho[im, im] = cho[k, k] + _swap_slices(cho, np.s_[im, :k], np.s_[k, :k]) + _swap_slices(cho, np.s_[im + 1:, im], np.s_[im + 1:, k]) + _swap_slices(cho, np.s_[k + 1:im, k], np.s_[im, k + 1:im]) + _swap_slices(new_lo, k, im) + _swap_slices(new_hi, k, im) + if ck > epk: + cho[k, k] = ck + cho[k, k + 1:] = 0.0 + for i in range(k + 1, n): + cho[i, k] /= ck + cho[i, k + 1:i + 1] -= cho[i, k] * cho[k + 1:i + 1, k] + if abs(dem) > tol: + y[k] = ((np.exp(-lo_m * lo_m / 2) - np.exp(-hi_m * hi_m / 2)) / + (sqtp * dem)) + else: + y[k] = (lo_m + hi_m) / 2 + if lo_m < -10: + y[k] = hi_m + elif hi_m > 10: + y[k] = lo_m + cho[k, :k + 1] /= ck + new_lo[k] /= ck + new_hi[k] /= ck + else: + cho[k:, k] = 0.0 + y[k] = (new_lo[k] + new_hi[k]) / 2 + return cho, new_lo, new_hi + + +def _swap_slices(x, slc1, slc2): + t = x[slc1].copy() + x[slc1] = x[slc2].copy() + x[slc2] = t + + +def _bvn(a, b, A): + # covariance matrix is written [[s1**2, rho*s1*s2], [rho*s1*s2, s2**2]] + # e.g. https://en.wikipedia.org/wiki/Multivariate_normal_distribution + # therefore, s12 = rho*s1*s2 -> rho = s12/(s1*s2) + s1 = math.sqrt(A[0, 0]) + s2 = math.sqrt(A[1, 1]) + s12 = A[0, 1] + r = s12 / (s1 * s2) + # the x and y coordinates seem to be normalized by the standard devs + xl, xu = a[0] / s1, b[0] / s1 + yl, yu = a[1] / s2, b[1] / s2 + p = _bvnu(xl, yl, r) - _bvnu(xu, yl, r) - _bvnu(xl, yu, r) + _bvnu(xu, yu, r) + p = max( 0., min( p, 1. ) ) + return p diff --git a/v12x/pylib/scipy/stats/_quantile.py b/v12x/pylib/scipy/stats/_quantile.py new file mode 100644 index 0000000000000000000000000000000000000000..cd3756f23138ae04cfabb7b47dfaa24a47c4c10a --- /dev/null +++ b/v12x/pylib/scipy/stats/_quantile.py @@ -0,0 +1,816 @@ +import math +import numpy as np +from scipy.special import betainc +from scipy._lib._array_api import ( + xp_capabilities, + xp_ravel, + array_namespace, + xp_promote, + xp_device, + xp_size, + _count_nonmasked, + is_torch, + is_lazy_array, +) +import scipy._external.array_api_extra as xpx +from scipy.stats._axis_nan_policy import _broadcast_arrays, _contains_nan + + +def _quantile_iv(x, p, method, axis, nan_policy, keepdims, weights, fun='quantile'): + xp = array_namespace(x, p, weights) + + if fun == "quantile": + methods = {'inverted_cdf', 'averaged_inverted_cdf', 'closest_observation', + 'hazen', 'interpolated_inverted_cdf', 'linear', + 'median_unbiased', 'normal_unbiased', 'weibull', + 'harrell-davis', '_lower', '_midpoint', '_higher', '_nearest', + 'round_outward', 'round_inward', 'round_nearest'} + allowed_types = 'real floating' + def mask_fun(p): return (p > 1) | (p < 0) | xp.isnan(p) + var2_name = 'p' + var2_type_msg = '`p` must have real floating dtype.' + else: + methods = set(_estimated_cdf_methods) + allowed_types = ('integral', 'real floating') + mask_fun = xp.isnan + var2_name = 'y' + var2_type_msg = '`y` must have real dtype.' + + if not xp.isdtype(xp.asarray(x).dtype, ('integral', 'real floating')): + raise ValueError("`x` must have real dtype.") + + if not xp.isdtype(xp.asarray(p).dtype, allowed_types): + raise ValueError(var2_type_msg) + + if not (weights is None + or xp.isdtype(xp.asarray(weights).dtype, ('integral', 'real floating'))): + raise ValueError("`weights` must have real dtype.") + + x, p, weights = xp_promote(x, p, weights, force_floating=True, xp=xp) + p = xp.asarray(p, device=xp_device(x)) + dtype = x.dtype + + axis_none = axis is None + ndim = max(x.ndim, p.ndim) + if axis_none: + x = xp_ravel(x) + p = xp_ravel(p) + axis = 0 + elif np.iterable(axis) or int(axis) != axis: + message = "`axis` must be an integer or None." + raise ValueError(message) + elif (axis >= ndim) or (axis < -ndim): + message = "`axis` is not compatible with the shapes of the inputs." + raise ValueError(message) + axis = int(axis) + + if method not in methods: + message = f"`method` must be one of {methods}" + raise ValueError(message) + + no_weights = {'_lower', '_midpoint', '_higher', '_nearest', 'harrell-davis', + 'round_nearest', 'round_inward', 'round_outward'} + if weights is not None and method in no_weights: + message = f"`method='{method}'` does not support `weights`." + raise ValueError(message) + + contains_nans = _contains_nan(x, nan_policy, xp_omit_okay=True, xp=xp) + + if keepdims not in {None, True, False}: + message = "If specified, `keepdims` must be True or False." + raise ValueError(message) + + # If data has length zero along `axis`, the result will be an array of NaNs just + # as if the data had length 1 along axis and were filled with NaNs. This is treated + # naturally below whether `nan_policy` is `'propagate'` or `'omit'`. + if x.shape[axis] == 0 and fun == 'quantile': + shape = list(x.shape) + shape[axis] = 1 + x = xp.full(shape, xp.nan, dtype=dtype, device=xp_device(x)) + + if weights is None: + y = xp.sort(x, axis=axis, stable=False) + y, p = _broadcast_arrays((y, p), axis=axis) + n_zero_weight = None + else: + x, weights = xp.broadcast_arrays(x, weights) + i_zero_weight = (weights == 0) + n_zero_weight = xp.count_nonzero(i_zero_weight, axis=axis, keepdims=True) + x = xpx.at(x)[i_zero_weight].set(xp.inf, copy=True) + i_y = xp.argsort(x, axis=axis, stable=False) + y = xp.take_along_axis(x, i_y, axis=axis) + weights = xp.take_along_axis(weights, i_y, axis=axis) + y, p, weights, i_y, n_zero_weight = _broadcast_arrays( + (y, p, weights, i_y, n_zero_weight), axis=axis) + + if (keepdims is False) and (p.shape[axis] != 1): + message = ("`keepdims` may be False only if the length of " + f"`{var2_name}` along `axis` is 1.") + raise ValueError(message) + keepdims = (p.shape[axis] != 1) if keepdims is None else keepdims + + y = xp.moveaxis(y, axis, -1) + p = xp.moveaxis(p, axis, -1) + weights = weights if weights is None else xp.moveaxis(weights, axis, -1) + n_zero_weight = (n_zero_weight if n_zero_weight is None + else xp.moveaxis(n_zero_weight, axis, -1)) + + n = _count_nonmasked(y, -1, xp=xp, keepdims=True) + n = n if n_zero_weight is None else n - n_zero_weight + + nan_out = None + if is_lazy_array(y) or contains_nans: + nans = xp.isnan(y) + + # Note that if length along `axis` were 0 to begin with, + # it is now length 1 and filled with NaNs. + if nan_policy == 'propagate': + nan_out = xp.any(nans, axis=-1, keepdims=True) + else: # 'omit' + n_int = n - xp.count_nonzero(nans, axis=-1, keepdims=True) + n = xp.astype(n_int, dtype) + # NaNs are produced only if slice is empty after removing NaNs + nan_out = n == 0 + n = xpx.at(n, nan_out).set(y.shape[-1]) # avoids pytorch/pytorch#146211 + + if (is_lazy_array(nans) or xp.any(nans)) and (method == 'harrell-davis'): + y = xp.asarray(y, copy=True) # ensure writable + y = xpx.at(y, nans).set(0) # any non-nan will prevent NaN from propagating + if is_lazy_array(nan_out) or xp.any(nan_out): + y = xp.asarray(y, copy=True) # ensure writable + y = xpx.at(y, xp.broadcast_to(nan_out, y.shape)).set(xp.nan) + + n = xp.asarray(n, dtype=dtype, device=xp_device(y)) + + # apparently xpx.at is accepting nan_out as a mask even though it doesn't have the + # same number of dimensions as `y`, yet it still appears to work correctly? + # should refactor for clarity, and p_mask that gets returned here should probably + # get a different name - `nan_out` would be more appropriate! + p_mask = mask_fun(p) if nan_out is None else mask_fun(p) | nan_out + if is_lazy_array(p_mask) or xp.any(p_mask): + p = xp.where(p_mask, 0.5, p) # these get NaN-ed out at the end + + return (y, p, method, axis, nan_policy, keepdims, + n, axis_none, ndim, p_mask, weights, xp) + + +@xp_capabilities(skip_backends=[("dask.array", "No take_along_axis yet.")], + marray=True, + extra_note=("Use of `weights` is incompatible with MArray. " + "``method='harrell-davis'`` is CPU-only.")) +def quantile(x, p, *, method='linear', axis=0, nan_policy='propagate', keepdims=None, + weights=None): + """ + Compute the p-th quantile of the data along the specified axis. + + Parameters + ---------- + x : array_like of real numbers + Data array. + p : array_like of float + Probability or sequence of probabilities of the quantiles to compute. + Values must be between 0 and 1 (inclusive). + While `numpy.quantile` can only compute quantiles according to the Cartesian + product of the first two arguments, this function enables calculation of + quantiles at different probabilities for each axis slice by following + broadcasting rules like those of `scipy.stats` reducing functions. + See `axis`, `keepdims`, and the examples. + method : str, default: 'linear' + The method to use for estimating the quantile. + The available options, numbered as they appear in [1]_, are: + + 1. 'inverted_cdf' + 2. 'averaged_inverted_cdf' + 3. 'closest_observation' + 4. 'interpolated_inverted_cdf' + 5. 'hazen' + 6. 'weibull' + 7. 'linear' (default) + 8. 'median_unbiased' + 9. 'normal_unbiased' + + 'harrell-davis' is also available to compute the quantile estimate + according to [2]_. + + 'round_outward', 'round_inward', and 'round_nearest' are available for use + in trimming and winsorizing data. + + See Notes for details. + axis : int or None, default: 0 + Axis along which the quantiles are computed. + ``None`` ravels both `x` and `p` before performing the calculation, + without checking whether the original shapes were compatible. + As in other `scipy.stats` functions, a positive integer `axis` is resolved + after prepending 1s to the shape of `x` or `p` as needed until the two arrays + have the same dimensionality. When providing `x` and `p` with different + dimensionality, consider using negative `axis` integers for clarity. + nan_policy : str, default: 'propagate' + Defines how to handle NaNs in the input data `x`. + + - ``propagate``: if a NaN is present in the axis slice (e.g. row) along + which the statistic is computed, the corresponding slice of the output + will contain NaN(s). + - ``omit``: NaNs will be omitted when performing the calculation. + If insufficient data remains in the axis slice along which the + statistic is computed, the corresponding slice of the output will + contain NaN(s). + - ``raise``: if a NaN is present, a ``ValueError`` will be raised. + + If NaNs are present in `p`, a ``ValueError`` will be raised. + keepdims : bool, optional + Consider the case in which `x` is 1-D and `p` is a scalar: the quantile + is a reducing statistic, and the default behavior is to return a scalar. + If `keepdims` is set to True, the axis will not be reduced away, and the + result will be a 1-D array with one element. + + The general case is more subtle, since multiple quantiles may be + requested for each axis-slice of `x`. For instance, if both `x` and `p` + are 1-D and ``p.size > 1``, no axis can be reduced away; there must be an + axis to contain the number of quantiles given by ``p.size``. Therefore: + + - By default, the axis will be reduced away if possible (i.e. if there is + exactly one element of `p` per axis-slice of `x`). + - If `keepdims` is set to True, the axis will not be reduced away. + - If `keepdims` is set to False, the axis will be reduced away + if possible, and an error will be raised otherwise. + + weights : array_like of finite, non-negative real numbers + Frequency weights; e.g., for counting number weights, + ``quantile(x, p, weights=weights)`` is equivalent to + ``quantile(np.repeat(x, weights), p)``. Values other than finite counting + numbers are accepted, but may not have valid statistical interpretations. + Not compatible with ``method='harrell-davis'`` or those that begin with + ``'round_'``. + + Returns + ------- + quantile : scalar or ndarray + The resulting quantile(s). The dtype is the result dtype of `x` and `p`. + + See Also + -------- + numpy.quantile + estimated_cdf + :ref:`outliers` + + Notes + ----- + Given a sample `x` from an underlying distribution, `quantile` provides a + nonparametric estimate of the inverse cumulative distribution function. + + By default, this is done by interpolating between adjacent elements in + ``y``, a sorted copy of `x`:: + + (1-g)*y[j] + g*y[j+1] + + where the index ``j`` and coefficient ``g`` are the integral and + fractional components of ``p * (n-1)``, and ``n`` is the number of + elements in the sample. + + This is a special case of Equation 1 of H&F [1]_. More generally, + + - ``j = (p*n + m - 1) // 1``, and + - ``g = (p*n + m - 1) % 1``, + + where ``m`` may be defined according to several different conventions. + The preferred convention may be selected using the ``method`` parameter: + + =============================== =============== =============== + ``method`` number in H&F ``m`` + =============================== =============== =============== + ``interpolated_inverted_cdf`` 4 ``0`` + ``hazen`` 5 ``1/2`` + ``weibull`` 6 ``p`` + ``linear`` (default) 7 ``1 - p`` + ``median_unbiased`` 8 ``p/3 + 1/3`` + ``normal_unbiased`` 9 ``p/4 + 3/8`` + =============================== =============== =============== + + Note that indices ``j`` and ``j + 1`` are clipped to the range ``0`` to + ``n - 1`` when the results of the formula would be outside the allowed + range of non-negative indices. When ``j`` is clipped to zero, ``g`` is + set to zero as well. The ``-1`` in the formulas for ``j`` and ``g`` + accounts for Python's 0-based indexing. + + The table above includes only the estimators from [1]_ that are continuous + functions of probability `p` (estimators 4-9). SciPy also provides the + three discontinuous estimators from [1]_ (estimators 1-3), where ``j`` is + defined as above, ``m`` is defined as follows, and ``g`` is ``0`` when + ``index = p*n + m - 1`` is less than ``0`` and otherwise is defined below. + + 1. ``inverted_cdf``: ``m = 0`` and ``g = int(index - j > 0)`` + 2. ``averaged_inverted_cdf``: ``m = 0`` and + ``g = (1 + int(index - j > 0)) / 2`` + 3. ``closest_observation``: ``m = -1/2`` and + ``g = 1 - int((index == j) & (j%2 == 1))`` + + Note that for methods ``inverted_cdf`` and ``averaged_inverted_cdf``, only the + relative proportions of tied observations (and relative weights) affect the + results; for all other methods, the total number of observations (and absolute + weights) matter. + + A different strategy for computing quantiles from [2]_, ``method='harrell-davis'``, + uses a weighted combination of all elements. The weights are computed as: + + .. math:: + + w_{n, i} = I_{i/n}(a, b) - I_{(i - 1)/n}(a, b) + + where :math:`n` is the number of elements in the sample, + :math:`i` are the indices :math:`1, 2, ..., n-1, n` of the sorted elements, + :math:`a = p (n + 1)`, :math:`b = (1 - p)(n + 1)`, + :math:`p` is the probability of the quantile, and + :math:`I` is the regularized, lower incomplete beta function + (`scipy.special.betainc`). + + ``method='round_nearest'`` is equivalent to indexing ``y[j]``, where:: + + j = int(np.round(p*n) if p < 0.5 else np.round(n*p - 1)) + + This is useful when winsorizing data: replacing ``p*n`` of the most extreme + observations with the next most extreme observation. ``method='round_outward'`` + adjusts the direction of rounding to winsorize fewer elements:: + + j = int(np.floor(p*n) if p < 0.5 else np.ceil(n*p - 1)) + + and ``method='round_inward'`` rounds to winsorize more elements:: + + j = int(np.ceil(p*n) if p < 0.5 else np.floor(n*p - 1)) + + These methods are also useful for trimming data: removing ``p*n`` of the most + extreme observations. See :ref:`outliers` for example applications. + + Examples + -------- + >>> import numpy as np + >>> from scipy import stats + >>> x = np.asarray([[10, 8, 7, 5, 4], + ... [0, 1, 2, 3, 5]]) + + Take the median of each row. + + >>> stats.quantile(x, 0.5, axis=-1) + array([7., 2.]) + + Take a different quantile for each row. + + >>> stats.quantile(x, [[0.25], [0.75]], axis=-1, keepdims=True) + array([[5.], + [3.]]) + + Take multiple quantiles for each row. + + >>> stats.quantile(x, [0.25, 0.75], axis=-1) + array([[5., 8.], + [1., 3.]]) + + Take different quantiles for each row. + + >>> p = np.asarray([[0.25, 0.75], + ... [0.5, 1.0]]) + >>> stats.quantile(x, p, axis=-1) + array([[5., 8.], + [2., 5.]]) + + Take different quantiles for each column. + + >>> stats.quantile(x.T, p.T, axis=0) + array([[5., 2.], + [8., 5.]]) + + References + ---------- + .. [1] R. J. Hyndman and Y. Fan, + "Sample quantiles in statistical packages," + The American Statistician, 50(4), pp. 361-365, 1996 + .. [2] Harrell, Frank E., and C. E. Davis. + "A new distribution-free quantile estimator." + Biometrika 69.3 (1982): 635-640. + + """ + # Input validation / standardization + + temp = _quantile_iv(x, p, method, axis, nan_policy, keepdims, weights) + (y, p, method, axis, nan_policy, keepdims, + n, axis_none, ndim, p_mask, weights, xp) = temp + + if method in {'inverted_cdf', 'averaged_inverted_cdf', 'closest_observation', + 'hazen', 'interpolated_inverted_cdf', 'linear', + 'median_unbiased', 'normal_unbiased', 'weibull'}: + res = _quantile_hf(y, p, n, method, weights, xp) + elif method in {'harrell-davis'}: + res = _quantile_hd(y, p, n, xp) + elif method in {'_lower', '_midpoint', '_higher', '_nearest'}: + res = _quantile_bc(y, p, n, method, xp) + else: # method.startswith('round'): + res = _quantile_winsor(y, p, n, method, xp) + + return _post_quantile(res, p_mask, axis, axis_none, ndim, keepdims, xp) + + +def _post_quantile(res, p_mask, axis, axis_none, ndim, keepdims, xp): + res = xpx.at(res, p_mask).set(xp.nan) + + # Reshape per axis/keepdims + if axis_none and keepdims: + shape = (1,)*(ndim - 1) + res.shape + res = xp.reshape(res, shape) + axis = -1 + + res = xp.moveaxis(res, -1, axis) + + if not keepdims: + res = xp.squeeze(res, axis=axis) + + return res[()] if res.ndim == 0 else res + + +def _quantile_hf(y, p, n, method, weights, xp): + ms = dict(inverted_cdf=0, averaged_inverted_cdf=0, closest_observation=-0.5, + interpolated_inverted_cdf=0, hazen=0.5, weibull=p, linear=1 - p, + median_unbiased=p/3 + 1/3, normal_unbiased=p/4 + 3/8) + m = ms[method] + + if weights is None: + jg = p * n + m + jp1 = jg // 1 + j = jp1 - 1 + else: + cumulative_weights = xp.cumulative_sum(weights, axis=-1) + n_int = xp.asarray(n, dtype=xp.int64) + n_int = xp.broadcast_to(n_int, cumulative_weights.shape[:-1] + (1,)) + total_weight = xp.take_along_axis(cumulative_weights, n_int-1, axis=-1) + jg = p * total_weight + m + jp1 = _xp_searchsorted(cumulative_weights, jg, side='right') + j = _xp_searchsorted(cumulative_weights, jg-1, side='right') + j, jp1 = xp.astype(j, y.dtype), xp.astype(jp1, y.dtype) + + g = jg % 1 + if method == 'inverted_cdf': + g = xp.astype((g > 0), jg.dtype) + elif method == 'averaged_inverted_cdf': + g = (1 + xp.astype((g > 0), jg.dtype)) / 2 + elif method == 'closest_observation': + g = (1 - xp.astype((g == 0) & (j % 2 == 1), jg.dtype)) + if method in {'inverted_cdf', 'averaged_inverted_cdf', 'closest_observation'}: + g = xp.asarray(g) + g = xpx.at(g, jg < 0).set(0) + + g = xpx.at(g)[j < 0].set(0) + j = xp.clip(j, 0., n - 1) + jp1 = xp.clip(jp1, 0., n - 1) + + return ((1 - g) * xp.take_along_axis(y, xp.astype(j, xp.int64), axis=-1) + + g * xp.take_along_axis(y, xp.astype(jp1, xp.int64), axis=-1)) + + +def _quantile_hd(y, p, n, xp): + # RE axis handling: We need to perform a reducing operation over rows of `y` for + # each element in the corresponding row of `p` (a la Cartesian product). Strategy: + # move rows of `p` to an axis at the front that is orthogonal to all the rest, + # perform the reducing operating over the last axis, then move the front axis back + # to the end. + p = xp.moveaxis(p, -1, 0)[..., xp.newaxis] + a = p * (n + 1) + b = (1 - p) * (n + 1) + i = xp.arange(y.shape[-1] + 1, dtype=y.dtype, device=xp_device(y)) + w = betainc(a, b, i / n) + w = w[..., 1:] - w[..., :-1] + w = xpx.at(w, xp.isnan(w)).set(0) + res = xp.vecdot(w, y, axis=-1) + return xp.moveaxis(res, 0, -1) + + +def _quantile_winsor(y, p, n, method, xp): + ops = dict(round_outward=(xp.floor, xp.ceil), + round_inward=(xp.ceil, xp.floor), + round_nearest=(xp.round, xp.round)) + op_left, op_right = ops[method] + j = xp.where(p < 0.5, op_left(p*n), op_right(n*p - 1)) + return xp.take_along_axis(y, xp.astype(j, xp.int64), axis=-1) + + +def _quantile_bc(y, p, n, method, xp): + # Methods retained for backward compatibility. NumPy documentation is not + # quite right about what these methods do: if `p * (n - 1)` is integral, + # that is used as the index. See numpy/numpy#28910. + ij = p * (n - 1) + if method == '_midpoint': + return (xp.take_along_axis(y, xp.astype(xp.floor(ij), xp.int64), axis=-1) + + xp.take_along_axis(y, xp.astype(xp.ceil(ij), xp.int64), axis=-1)) / 2 + elif method == '_lower': + k = xp.floor(ij) + elif method == '_higher': + k = xp.ceil(ij) + elif method == '_nearest': + k = xp.round(ij) + return xp.take_along_axis(y, xp.astype(k, xp.int64), axis=-1) + + +@xp_capabilities(skip_backends=[("dask.array", "No take_along_axis yet.")]) +def _xp_searchsorted(x, y, *, side='left', xp=None): + # Vectorized xp.searchsorted. Search is performed along last axis. The shape of the + # output is that of `y`, broadcasting the batch dimensions with those of `x` if + # necessary. + xp = array_namespace(x, y) if xp is None else xp + xp_default_int = xp.asarray(1).dtype + y_0d = xp.asarray(y).ndim == 0 + x, y = _broadcast_arrays((x, y), axis=-1, xp=xp) + x_1d = x.ndim <= 1 + + if x_1d or is_torch(xp): + y = xp.reshape(y, ()) if (y_0d and x_1d) else y + out = xp.searchsorted(x, y, side=side) + out = xp.astype(out, xp_default_int, copy=False) + return out + + a = xp.full(y.shape, 0, device=xp_device(x)) + + if x.shape[-1] == 0: + return a + + n = xp.count_nonzero(~xp.isnan(x), axis=-1, keepdims=True) + b = xp.broadcast_to(n, y.shape) + + compare = xp.less_equal if side == 'left' else xp.less + + # while xp.any(b - a > 1): + # refactored to for loop with ~log2(n) iterations for JAX JIT + for i in range(int(math.log2(x.shape[-1])) + 1): + c = (a + b) // 2 + x0 = xp.take_along_axis(x, c, axis=-1) + j = compare(y, x0) + b = xp.where(j, c, b) + a = xp.where(j, a, c) + + out = xp.where(compare(y, xp.min(x, axis=-1, keepdims=True)), 0, b) + out = xp.where(xp.isnan(y), x.shape[-1], out) if side == 'right' else out + out = xp.astype(out, xp_default_int, copy=False) + return out + + +@xp_capabilities(skip_backends=[("dask.array", "No take_along_axis yet.")], + marray=True) +def estimated_cdf(x, y, *, method='linear', + axis=0, nan_policy='propagate', keepdims=None): + """ + Estimate the CDF of the distribution underlying a sample. + + Parameters + ---------- + x : array_like of real numbers + Data array. + y : array_like of real numbers + Points at which to evaluate the estimated cdf. + See `axis`, `keepdims`, and the examples for broadcasting behavior. + method : str, default: 'linear' + The method to use for estimating the cumulative distribution function. + The available options, numbered as they appear in [1]_, are: + + 1. 'inverted_cdf' (AKA *the* empirical cumulative distribution function) + 2. 'averaged_inverted_cdf' + 3. 'closest_observation' + 4. 'interpolated_inverted_cdf' + 5. 'hazen' + 6. 'weibull' + 7. 'linear' (default) + 8. 'median_unbiased' + 9. 'normal_unbiased' + + See Notes for details. + axis : int or None, default: 0 + Axis along which samples in `x` are given in ND case. + ``None`` ravels both `x` and `y` before performing the calculation, + without checking whether the original shapes were compatible. + As in other `scipy.stats` functions, a positive integer `axis` is resolved + after prepending 1s to the shape of `x` or `y` as needed until the two arrays + have the same dimensionality. When providing `x` and `y` with different + dimensionality, consider using negative `axis` integers for clarity. + nan_policy : str, default: 'propagate' + Defines how to handle NaNs in the input data `x`. + + - ``propagate``: if a NaN is present in the axis slice (e.g. row) along + which the statistic is computed, the corresponding slice of the output + will contain NaN(s). + - ``omit``: NaNs will be omitted when performing the calculation. + If insufficient data remains in the axis slice along which the + statistic is computed, the corresponding slice of the output will + contain NaN(s). + - ``raise``: if a NaN is present, a ``ValueError`` will be raised. + + If NaNs are present in `y`, the corresponding entries in the output will be NaN. + keepdims : bool, optional + Consider the case in which `x` is 1-D and `y` is a scalar: the estimated + cumulative distribution function at a point is a reducing statistic, and the + default behavior is to return a scalar. + If `keepdims` is set to True, the axis will not be reduced away, and the + result will be a 1-D array with one element. + + The general case is more subtle, since multiple values of `y` may be + specified for each axis-slice of `x`. For instance, if both `x` and `y` + are 1-D and ``y.size > 1``, no axis can be reduced away; there must be an + axis to contain the number of values given by ``y.size``. Therefore: + + - By default, the axis will be reduced away if possible (i.e. if there is + exactly one element of `y` per axis-slice of `x`). + - If `keepdims` is set to True, the axis will not be reduced away. + - If `keepdims` is set to False, the axis will be reduced away + if possible, and an error will be raised otherwise. + + Returns + ------- + probability : scalar or ndarray + The resulting probabilities(s). The dtype is the result dtype of `x`, `y`, and + the Python ``float`` type. + + See Also + -------- + quantile + + Notes + ----- + Given a sample `x` from an underlying distribution, `estimated_cdf` provides a + nonparametric estimate of the cumulative distribution function. + + By default, this is done by computing the "fractional index" ``p`` at which ``y`` + would appear within ``z``, a sorted copy of `x`:: + + p = 1 / (n - 1) * (j + ( y - z[j]) + / (z[j+1] - z[j])) + + where the index ``j`` is that of the largest element of ``z`` that does not exceed + ``y``, and ``n`` is the number of elements in the sample. Note that if ``y`` is an + element of ``z``, then ``j`` is the index such that ``y = z[j]``, and the formula + simplifies to the intuitive ``j / (n - 1)``. The full formula linearly interpolates + between ``j / (n - 1)`` and ``(j + 1) / (n - 1)``. + + This is a special case of the more general:: + + p = (j + (y - z[j]) / (z[j+1] - z[j] + 1 - m) / n + + where ``m`` may be defined according to several different conventions. + The preferred convention may be selected using the ``method`` parameter: + + =============================== =============== =============== + ``method`` number in H&F ``m`` + =============================== =============== =============== + ``interpolated_inverted_cdf`` 4 ``0`` + ``hazen`` 5 ``1/2`` + ``weibull`` 6 ``p`` + ``linear`` (default) 7 ``1 - p`` + ``median_unbiased`` 8 ``p/3 + 1/3`` + ``normal_unbiased`` 9 ``p/4 + 3/8`` + =============================== =============== =============== + + Note that indices ``j`` and ``j + 1`` are clipped to the range ``0`` to + ``n - 1`` when the results of the formula would be outside the allowed + range of non-negative indices, and resulting ``p`` is clipped to the range + ``0`` to ``1``. + + The table above includes only the estimators from [1]_ that are continuous + functions of probability `p` (estimators 4-9). SciPy also provides the + three discontinuous estimators from [1]_ (estimators 1-3), where ``j`` is + defined as above, ``m`` is defined as follows, and ``g`` is ``0`` when + ``index = p*n + m - 1`` is less than ``0`` and otherwise is defined below. + + 1. ``inverted_cdf``: ``m = 0`` and ``g = int(index - j > 0)`` + 2. ``averaged_inverted_cdf``: ``m = 0`` and + ``g = (1 + int(index - j > 0)) / 2`` + 3. ``closest_observation``: ``m = -1/2`` and + ``g = 1 - int((index == j) & (j%2 == 1))`` + + When all the data in ``x`` are unique, `estimated_cdf` and `quantile` are are + inverses of one another within a certain domain. + Although `quantile` with ``method='linear'`` is invertible over the whole domain + of ``p`` from ``0`` to ``1``, this is not true of other methods. + + References + ---------- + .. [1] R. J. Hyndman and Y. Fan, + "Sample quantiles in statistical packages," + The American Statistician, 50(4), pp. 361-365, 1996 + + Examples + -------- + >>> import numpy as np + >>> from scipy import stats + >>> x = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10] + + Estimate the cumulative distribution function of one sample at a single point. + + >>> stats.estimated_cdf(x, 5, axis=-1) + np.float64(0.5) + + Estimate the cumulative distribution function of one sample at two points. + + >>> stats.estimated_cdf(x, [2.5, 7.5], axis=-1) + array([0.25, 0.75]) + + Estimate the cumulative distribution function of two samples at different points. + + >>> x = np.stack((np.arange(0, 11), np.arange(10, 21))) + >>> stats.estimated_cdf(x, [[2.5], [17.5]], axis=-1, keepdims=True) + array([[0.25], + [0.75]]) + + Estimate the cumulative distribution function at many points for each of two + samples. + + >>> import matplotlib.pyplot as plt + >>> rng = np.random.default_rng(6110515095) + >>> x = stats.Normal(mu=[-1, 1]).sample(10000) + >>> y = np.linspace(-4, 4, 5000)[:, np.newaxis] + >>> p = stats.estimated_cdf(x, y, axis=0) + >>> plt.plot(y, p) + >>> plt.show() + + Note that the `quantile` and `estimated_cdf` functions are inverses of one another + within a certain domain. + + >>> p = np.linspace(0, 1, 300) + >>> x = rng.standard_normal(300) + >>> y = stats.quantile(x, p) + >>> p2 = stats.estimated_cdf(x, y) + >>> np.testing.assert_allclose(p2, p) + >>> y2 = stats.quantile(x, p2) + >>> np.testing.assert_allclose(y2, y) + + However, the domain over which `quantile` can be inverted by `estimated_cdf` depends + on the `method` used. This is most noticeable when there are few observations in the + sample. + + >>> x = np.asarray([0, 1]) + >>> y_linear = stats.quantile(x, p, method='linear') + >>> y_weibull = stats.quantile(x, p, method='weibull') + >>> y_iicdf = stats.quantile(x, p, method='interpolated_inverted_cdf') + >>> plt.plot(p, y_linear, p, y_weibull, p, y_iicdf) + >>> plt.legend(['linear', 'weibull', 'iicdf']) + >>> plt.xlabel('p') + >>> plt.ylabel('y = quantile(x, p)') + >>> plt.show() + + For example, in the case above, `quantile` is only invertible from + ``p = 0.5`` to ``p = 1.0`` with ``method = 'interpolated_inverted_cdf'``. This is a + fundamental characteristic of the methods, not a shortcoming of `estimated_cdf`. + + """ + temp = _quantile_iv(x, y, method, axis, nan_policy, keepdims, weights=None, + fun='estimated_cdf') + x, y, method, axis, nan_policy, keepdims, n, axis_none, ndim, y_mask, _, xp = temp + + if xp_size(x) == 0: + res = xp.full_like(y, xp.nan) + else: + res = _estimated_cdf_hf(x, y, n, method, xp) + + return _post_quantile(res, y_mask, axis, axis_none, ndim, keepdims, xp) + + +_estimated_cdf_discontinuous_methods = dict( + inverted_cdf=0.0, + averaged_inverted_cdf=0.0, + closest_observation=0.5, +) + + +_estimated_cdf_continuous_methods = dict( + interpolated_inverted_cdf=(0, 1), + hazen=(0.5, 0.5), + weibull=(0, 0), + linear=(1, 1), + median_unbiased=(1 / 3, 1 / 3), + normal_unbiased=(3 / 8, 3 / 8), +) + + +_estimated_cdf_methods = (set(_estimated_cdf_continuous_methods).union( + set(_estimated_cdf_discontinuous_methods))) + + +def _estimated_cdf_hf(x, y, n, method, xp): + n_int = xp.astype(n, xp.int64) + j_max = n_int - 1 + j_min = xp.minimum(j_max, xp.asarray(1, dtype=j_max.dtype)) + jp1 = _xp_searchsorted(x, y, side='right') + + if method in _estimated_cdf_discontinuous_methods: + dp = _estimated_cdf_discontinuous_methods[method] + p = (xp.astype(jp1, x.dtype)+dp)/n + + else: + jp1 = xp.clip(jp1, j_min, j_max) + j = xp.clip(jp1-1, 0) + xj = xp.take_along_axis(x, j, axis=-1) + xjp1 = xp.take_along_axis(x, jp1, axis=-1) + with np.errstate(divide='ignore', invalid='ignore'): # refactor to apply_where? + delta = xp.where((xjp1 > xj) & xp.isfinite(xj), (y - xj) / (xjp1 - xj), 1.) + + a, b = _estimated_cdf_continuous_methods[method] + p = (xp.astype(jp1, x.dtype) + delta - a) / (n + 1 - a - b) + + xmin = x[..., :1] + xmax = (x[..., -1:] if n.shape == () else xp.take_along_axis(x, j_max)) + p = xpx.at(p)[y < xmin].set(0.) + p = xpx.at(p)[y > xmax].set(1.) + + return xp.clip(p, 0., 1.) diff --git a/v12x/pylib/scipy/stats/_rcont/__init__.py b/v12x/pylib/scipy/stats/_rcont/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..25728ead5f7ea277e6e94c359d5a5603b99eeb38 --- /dev/null +++ b/v12x/pylib/scipy/stats/_rcont/__init__.py @@ -0,0 +1,4 @@ +# +from .rcont import rvs_rcont1, rvs_rcont2 + +__all__ = ["rvs_rcont1", "rvs_rcont2"] diff --git a/v12x/pylib/scipy/stats/_relative_risk.py b/v12x/pylib/scipy/stats/_relative_risk.py new file mode 100644 index 0000000000000000000000000000000000000000..51525fd28adb37c72b12106450e4178c786091b2 --- /dev/null +++ b/v12x/pylib/scipy/stats/_relative_risk.py @@ -0,0 +1,263 @@ +import operator +from dataclasses import dataclass +import numpy as np +from scipy.special import ndtri +from ._common import ConfidenceInterval + + +def _validate_int(n, bound, name): + msg = f'{name} must be an integer not less than {bound}, but got {n!r}' + try: + n = operator.index(n) + except TypeError: + raise TypeError(msg) from None + if n < bound: + raise ValueError(msg) + return n + + +@dataclass +class RelativeRiskResult: + """ + Result of `scipy.stats.contingency.relative_risk`. + + Attributes + ---------- + relative_risk : float + This is:: + + (exposed_cases/exposed_total) / (control_cases/control_total) + + exposed_cases : int + The number of "cases" (i.e. occurrence of disease or other event + of interest) among the sample of "exposed" individuals. + exposed_total : int + The total number of "exposed" individuals in the sample. + control_cases : int + The number of "cases" among the sample of "control" or non-exposed + individuals. + control_total : int + The total number of "control" individuals in the sample. + + Methods + ------- + confidence_interval : + Compute the confidence interval for the relative risk estimate. + """ + + relative_risk: float + exposed_cases: int + exposed_total: int + control_cases: int + control_total: int + + def confidence_interval(self, confidence_level=0.95): + """ + Compute the confidence interval for the relative risk. + + The confidence interval is computed using the Katz method + (i.e. "Method C" of [1]_; see also [2]_, section 3.1.2). + + Parameters + ---------- + confidence_level : float, optional + The confidence level to use for the confidence interval. + Default is 0.95. + + Returns + ------- + ci : ConfidenceInterval instance + The return value is an object with attributes ``low`` and + ``high`` that hold the confidence interval. + + References + ---------- + .. [1] D. Katz, J. Baptista, S. P. Azen and M. C. Pike, "Obtaining + confidence intervals for the risk ratio in cohort studies", + Biometrics, 34, 469-474 (1978). + .. [2] Hardeo Sahai and Anwer Khurshid, Statistics in Epidemiology, + CRC Press LLC, Boca Raton, FL, USA (1996). + + + Examples + -------- + >>> from scipy.stats.contingency import relative_risk + >>> result = relative_risk(exposed_cases=10, exposed_total=75, + ... control_cases=12, control_total=225) + >>> result.relative_risk + 2.5 + >>> result.confidence_interval() + ConfidenceInterval(low=1.1261564003469628, high=5.549850800541033) + """ + if not 0 <= confidence_level <= 1: + raise ValueError('confidence_level must be in the interval ' + '[0, 1].') + + # Handle edge cases where either exposed_cases or control_cases + # is zero. We follow the convention of the R function riskratio + # from the epitools library. + if self.exposed_cases == 0 and self.control_cases == 0: + # relative risk is nan. + return ConfidenceInterval(low=np.nan, high=np.nan) + elif self.exposed_cases == 0: + # relative risk is 0. + return ConfidenceInterval(low=0.0, high=np.nan) + elif self.control_cases == 0: + # relative risk is inf + return ConfidenceInterval(low=np.nan, high=np.inf) + + alpha = 1 - confidence_level + z = ndtri(1 - alpha/2) + rr = self.relative_risk + + # Estimate of the variance of log(rr) is + # var(log(rr)) = 1/exposed_cases - 1/exposed_total + + # 1/control_cases - 1/control_total + # and the standard error is the square root of that. + se = np.sqrt(1/self.exposed_cases - 1/self.exposed_total + + 1/self.control_cases - 1/self.control_total) + delta = z*se + katz_lo = rr*np.exp(-delta) + katz_hi = rr*np.exp(delta) + return ConfidenceInterval(low=katz_lo, high=katz_hi) + + +def relative_risk(exposed_cases, exposed_total, control_cases, control_total): + """ + Compute the relative risk (also known as the risk ratio). + + This function computes the relative risk associated with a 2x2 + contingency table ([1]_, section 2.2.3; [2]_, section 3.1.2). Instead + of accepting a table as an argument, the individual numbers that are + used to compute the relative risk are given as separate parameters. + This is to avoid the ambiguity of which row or column of the contingency + table corresponds to the "exposed" cases and which corresponds to the + "control" cases. Unlike, say, the odds ratio, the relative risk is not + invariant under an interchange of the rows or columns. + + Parameters + ---------- + exposed_cases : nonnegative int + The number of "cases" (i.e. occurrence of disease or other event + of interest) among the sample of "exposed" individuals. + exposed_total : positive int + The total number of "exposed" individuals in the sample. + control_cases : nonnegative int + The number of "cases" among the sample of "control" or non-exposed + individuals. + control_total : positive int + The total number of "control" individuals in the sample. + + Returns + ------- + result : instance of `~scipy.stats._result_classes.RelativeRiskResult` + The object has the float attribute ``relative_risk``, which is:: + + rr = (exposed_cases/exposed_total) / (control_cases/control_total) + + The object also has the method ``confidence_interval`` to compute + the confidence interval of the relative risk for a given confidence + level. + + See Also + -------- + odds_ratio + + Notes + ----- + The R package epitools has the function `riskratio`, which accepts + a table with the following layout:: + + disease=0 disease=1 + exposed=0 (ref) n00 n01 + exposed=1 n10 n11 + + With a 2x2 table in the above format, the estimate of the CI is + computed by `riskratio` when the argument method="wald" is given, + or with the function `riskratio.wald`. + + For example, in a test of the incidence of lung cancer among a + sample of smokers and nonsmokers, the "exposed" category would + correspond to "is a smoker" and the "disease" category would + correspond to "has or had lung cancer". + + To pass the same data to ``relative_risk``, use:: + + relative_risk(n11, n10 + n11, n01, n00 + n01) + + .. versionadded:: 1.7.0 + + References + ---------- + .. [1] Alan Agresti, An Introduction to Categorical Data Analysis + (second edition), Wiley, Hoboken, NJ, USA (2007). + .. [2] Hardeo Sahai and Anwer Khurshid, Statistics in Epidemiology, + CRC Press LLC, Boca Raton, FL, USA (1996). + + Examples + -------- + >>> from scipy.stats.contingency import relative_risk + + This example is from Example 3.1 of [2]_. The results of a heart + disease study are summarized in the following table:: + + High CAT Low CAT Total + -------- ------- ----- + CHD 27 44 71 + No CHD 95 443 538 + + Total 122 487 609 + + CHD is coronary heart disease, and CAT refers to the level of + circulating catecholamine. CAT is the "exposure" variable, and + high CAT is the "exposed" category. So the data from the table + to be passed to ``relative_risk`` is:: + + exposed_cases = 27 + exposed_total = 122 + control_cases = 44 + control_total = 487 + + >>> result = relative_risk(27, 122, 44, 487) + >>> result.relative_risk + 2.4495156482861398 + + Find the confidence interval for the relative risk. + + >>> result.confidence_interval(confidence_level=0.95) + ConfidenceInterval(low=1.5836990926700116, high=3.7886786315466354) + + The interval does not contain 1, so the data supports the statement + that high CAT is associated with greater risk of CHD. + """ + # Relative risk is a trivial calculation. The nontrivial part is in the + # `confidence_interval` method of the RelativeRiskResult class. + + exposed_cases = _validate_int(exposed_cases, 0, "exposed_cases") + exposed_total = _validate_int(exposed_total, 1, "exposed_total") + control_cases = _validate_int(control_cases, 0, "control_cases") + control_total = _validate_int(control_total, 1, "control_total") + + if exposed_cases > exposed_total: + raise ValueError('exposed_cases must not exceed exposed_total.') + if control_cases > control_total: + raise ValueError('control_cases must not exceed control_total.') + + if exposed_cases == 0 and control_cases == 0: + # relative risk is 0/0. + rr = np.nan + elif exposed_cases == 0: + # relative risk is 0/nonzero + rr = 0.0 + elif control_cases == 0: + # relative risk is nonzero/0. + rr = np.inf + else: + p1 = exposed_cases / exposed_total + p2 = control_cases / control_total + rr = p1 / p2 + return RelativeRiskResult(relative_risk=rr, + exposed_cases=exposed_cases, + exposed_total=exposed_total, + control_cases=control_cases, + control_total=control_total) diff --git a/v12x/pylib/scipy/stats/_resampling.py b/v12x/pylib/scipy/stats/_resampling.py new file mode 100644 index 0000000000000000000000000000000000000000..3977e3718809ff6772dbee247c506f1257c9a624 --- /dev/null +++ b/v12x/pylib/scipy/stats/_resampling.py @@ -0,0 +1,2425 @@ +import warnings +import numpy as np +from itertools import combinations, permutations, product, accumulate +from collections.abc import Sequence +from dataclasses import dataclass, field +import inspect +import math + +from scipy._lib._util import (check_random_state, _rename_parameter, rng_integers, + _transition_to_rng) +from scipy._lib._array_api import ( + array_namespace, + is_numpy, + is_array_api_strict, + xp_capabilities, + xp_result_type, + xp_size, + xp_device, + xp_swapaxes, + is_lazy_array, +) +from scipy._external import array_api_extra as xpx +from scipy.special import ndtr, ndtri +from scipy import stats + +from ._common import ConfidenceInterval +from ._axis_nan_policy import _broadcast_concatenate, _broadcast_arrays +from ._warnings_errors import DegenerateDataWarning + +__all__ = ['bootstrap', 'monte_carlo_test', 'permutation_test'] + + +def _vectorize_statistic(statistic): + """Vectorize an n-sample statistic""" + # This is a little cleaner than np.nditer at the expense of some data + # copying: concatenate samples together, then use np.apply_along_axis + def stat_nd(*data, axis=0): + lengths = [sample.shape[axis] for sample in data] + split_indices = np.cumsum(lengths)[:-1] + z = _broadcast_concatenate(data, axis) + + # move working axis to position 0 so that new dimensions in the output + # of `statistic` are _prepended_. ("This axis is removed, and replaced + # with new dimensions...") + z = np.moveaxis(z, axis, 0) + + def stat_1d(z): + data = np.split(z, split_indices) + return statistic(*data) + + return np.apply_along_axis(stat_1d, 0, z)[()] + return stat_nd + + +def _jackknife_resample(sample, batch=None, *, xp): + """Jackknife resample the sample. Only one-sample stats for now.""" + n = sample.shape[-1] + batch_nominal = batch or n + + for k in range(0, n, batch_nominal): + # col_start:col_end are the observations to remove + batch_actual = min(batch_nominal, n-k) + + # jackknife - each row leaves out one observation + j = np.ones((batch_actual, n), dtype=bool) + np.fill_diagonal(j[:, k:k+batch_actual], False) + i = np.arange(n) + i = np.broadcast_to(i, (batch_actual, n)) + i = i[j].reshape((batch_actual, n-1)) + + i = xp.asarray(i, device=xp_device(sample)) + yield _get_from_last_axis(sample, i, xp=xp) + + +def _bootstrap_resample(sample, n_resamples=None, rng=None, *, xp): + """Bootstrap resample the sample.""" + n = sample.shape[-1] + + # bootstrap - each row is a random resample of original observations + i = rng_integers(rng, 0, n, (n_resamples, n)) + + i = xp.asarray(i, device=xp_device(sample)) + return _get_from_last_axis(sample, i, xp=xp) + + +def _get_from_last_axis(sample, i, xp): + if not is_array_api_strict(xp): + return sample[..., i] + + # Equivalent to `sample[..., i]` as used in `bootstrap`. Assumes i.ndim <=2. + if i.ndim == 2: + sample = xp.expand_dims(sample, axis=-2) + sample, i = _broadcast_arrays((sample, i), axis=-1, xp=xp) + return xp.take_along_axis(sample, i, axis=-1) + + +def _percentile_of_score(a, score, axis, xp): + """Vectorized, simplified `scipy.stats.percentileofscore`. + Uses logic of the 'mean' value of percentileofscore's kind parameter. + + Unlike `stats.percentileofscore`, the percentile returned is a fraction + in [0, 1]. + """ + B = a.shape[axis] + nonzeros = (xp.count_nonzero(a < score, axis=axis) + + xp.count_nonzero(a <= score, axis=axis)) + return xp.astype(nonzeros, score.dtype) / (2 * B) + + +def _bca_interval(data, statistic, axis, alpha, theta_hat_b, batch, xp): + """Bias-corrected and accelerated interval.""" + # closely follows [1] 14.3 and 15.4 (Eq. 15.36) + + # calculate z0_hat + theta_hat = xp.expand_dims(statistic(*data, axis=axis), axis=-1) + percentile = _percentile_of_score(theta_hat_b, theta_hat, axis=-1, xp=xp) + percentile = xp.asarray(percentile, dtype=theta_hat.dtype, + device=xp_device(theta_hat)) + z0_hat = ndtri(percentile) + + # calculate a_hat + theta_hat_ji = [] # j is for sample of data, i is for jackknife resample + for j, sample in enumerate(data): + # _jackknife_resample will add an axis prior to the last axis that + # corresponds with the different jackknife resamples. Do the same for + # each sample of the data to ensure broadcastability. We need to + # create a copy of the list containing the samples anyway, so do this + # in the loop to simplify the code. This is not the bottleneck... + samples = [xp.expand_dims(sample, axis=-2) for sample in data] + theta_hat_i = [] + for jackknife_sample in _jackknife_resample(sample, batch, xp=xp): + samples[j] = jackknife_sample + broadcasted = _broadcast_arrays(samples, axis=-1, xp=xp) + theta_hat_i.append(statistic(*broadcasted, axis=-1)) + theta_hat_ji.append(theta_hat_i) + + theta_hat_ji = [xp.concat(theta_hat_i, axis=-1) + for theta_hat_i in theta_hat_ji] + + # CuPy array dtype is unstable in 1.13.6 when divided by large Python int: + # (cp.ones(1, dtype=cp.float32)/10000).dtype # float32 + # (cp.ones(1, dtype=cp.float32)/100000).dtype # float64 + # Solution for now is to make the divisor a Python float + # Looks like this will be fixed in 1.14.0 + n_j = [float(theta_hat_i.shape[-1]) for theta_hat_i in theta_hat_ji] + + theta_hat_j_dot = [xp.mean(theta_hat_i, axis=-1, keepdims=True) + for theta_hat_i in theta_hat_ji] + + U_ji = [(n - 1) * (theta_hat_dot - theta_hat_i) + for theta_hat_dot, theta_hat_i, n + in zip(theta_hat_j_dot, theta_hat_ji, n_j)] + + nums = [xp.sum(U_i**3, axis=-1)/n**3 for U_i, n in zip(U_ji, n_j)] + dens = [xp.sum(U_i**2, axis=-1)/n**2 for U_i, n in zip(U_ji, n_j)] + a_hat = 1/6 * sum(nums) / sum(dens)**(3/2) + + # calculate alpha_1, alpha_2 + # `float` because dtype of ndtri output (float64) should not promote other vals + z_alpha = float(ndtri(alpha)) + z_1alpha = -z_alpha + num1 = z0_hat + z_alpha + alpha_1 = ndtr(z0_hat + num1/(1 - a_hat*num1)) + num2 = z0_hat + z_1alpha + alpha_2 = ndtr(z0_hat + num2/(1 - a_hat*num2)) + return alpha_1, alpha_2, a_hat # return a_hat for testing + + +def _bootstrap_iv(data, statistic, vectorized, paired, axis, confidence_level, + alternative, n_resamples, batch, method, bootstrap_result, + rng): + """Input validation and standardization for `bootstrap`.""" + xp = array_namespace(*data) + + if vectorized not in {True, False, None}: + raise ValueError("`vectorized` must be `True`, `False`, or `None`.") + + if vectorized is None: + vectorized = 'axis' in inspect.signature(statistic).parameters + + if not vectorized: + if not is_numpy(xp): + message = (f"When using array library {xp.__name__}, `func` must be " + "vectorized and accept argument `axis`.") + raise TypeError(message) + + statistic = _vectorize_statistic(statistic) + + axis_int = int(axis) + if axis != axis_int: + raise ValueError("`axis` must be an integer.") + + n_samples = 0 + try: + n_samples = len(data) + except TypeError: + raise ValueError("`data` must be a sequence of samples.") + + if n_samples == 0: + raise ValueError("`data` must contain at least one sample.") + + data = _broadcast_arrays(data, axis_int, xp=xp) + + data_iv = [] + for sample in data: + if sample.shape[axis_int] <= 1: + raise ValueError("each sample in `data` must contain two or more " + "observations along `axis`.") + sample = xp.moveaxis(sample, axis_int, -1) + data_iv.append(sample) + + if paired not in {True, False}: + raise ValueError("`paired` must be `True` or `False`.") + + if paired: + n = data_iv[0].shape[-1] + for sample in data_iv[1:]: + if sample.shape[-1] != n: + message = ("When `paired is True`, all samples must have the " + "same length along `axis`") + raise ValueError(message) + + # to generate the bootstrap distribution for paired-sample statistics, + # resample the indices of the observations + def statistic(i, axis=-1, data=data_iv, unpaired_statistic=statistic): + # data = [sample[..., i] for sample in data] + data = [_get_from_last_axis(sample, i, xp=xp) for sample in data] + return unpaired_statistic(*data, axis=axis) + + data_iv = [xp.arange(n)] + + confidence_level_float = float(confidence_level) + + alternative = alternative.lower() + alternatives = {'two-sided', 'less', 'greater'} + if alternative not in alternatives: + raise ValueError(f"`alternative` must be one of {alternatives}") + + n_resamples_int = int(n_resamples) + if n_resamples != n_resamples_int or n_resamples_int < 0: + raise ValueError("`n_resamples` must be a non-negative integer.") + + if batch is None: + batch_iv = batch + else: + batch_iv = int(batch) + if batch != batch_iv or batch_iv <= 0: + raise ValueError("`batch` must be a positive integer or None.") + + methods = {'percentile', 'basic', 'bca'} + method = method.lower() + if method not in methods: + raise ValueError(f"`method` must be in {methods}") + + message = "`bootstrap_result` must have attribute `bootstrap_distribution'" + if (bootstrap_result is not None + and not hasattr(bootstrap_result, "bootstrap_distribution")): + raise ValueError(message) + + message = ("Either `bootstrap_result.bootstrap_distribution.size` or " + "`n_resamples` must be positive.") + if ((not bootstrap_result or + not xp_size(bootstrap_result.bootstrap_distribution)) + and n_resamples_int == 0): + raise ValueError(message) + + rng = check_random_state(rng) + + return (data_iv, statistic, vectorized, paired, axis_int, + confidence_level_float, alternative, n_resamples_int, batch_iv, + method, bootstrap_result, rng, xp) + + +@dataclass +class BootstrapResult: + """Result object returned by `scipy.stats.bootstrap`. + + Attributes + ---------- + confidence_interval : ConfidenceInterval + The bootstrap confidence interval as an instance of + `collections.namedtuple` with attributes `low` and `high`. + bootstrap_distribution : ndarray + The bootstrap distribution, that is, the value of `statistic` for + each resample. The last dimension corresponds with the resamples + (e.g. ``res.bootstrap_distribution.shape[-1] == n_resamples``). + standard_error : float or ndarray + The bootstrap standard error, that is, the sample standard + deviation of the bootstrap distribution. + + """ + confidence_interval: ConfidenceInterval + bootstrap_distribution: np.ndarray + standard_error: float | np.ndarray + + +@xp_capabilities(skip_backends=[("dask.array", "Dask doesn't have take_along_axis.")], + jax_jit=False) # a few failed assertions - not sure what's going on +@_transition_to_rng('random_state') +def bootstrap(data, statistic, *, n_resamples=9999, batch=None, + vectorized=None, paired=False, axis=0, confidence_level=0.95, + alternative='two-sided', method='BCa', bootstrap_result=None, + rng=None): + r""" + Compute a two-sided bootstrap confidence interval of a statistic. + + When `method` is ``'percentile'`` and `alternative` is ``'two-sided'``, + a bootstrap confidence interval is computed according to the following + procedure. + + 1. Resample the data: for each sample in `data` and for each of + `n_resamples`, take a random sample of the original sample + (with replacement) of the same size as the original sample. + + 2. Compute the bootstrap distribution of the statistic: for each set of + resamples, compute the test statistic. + + 3. Determine the confidence interval: find the interval of the bootstrap + distribution that is + + - symmetric about the median and + - contains `confidence_level` of the resampled statistic values. + + While the ``'percentile'`` method is the most intuitive, it is rarely + used in practice. Two more common methods are available, ``'basic'`` + ('reverse percentile') and ``'BCa'`` ('bias-corrected and accelerated'); + they differ in how step 3 is performed. + + If the samples in `data` are taken at random from their respective + distributions :math:`n` times, the confidence interval returned by + `bootstrap` will contain the true value of the statistic for those + distributions approximately `confidence_level`:math:`\, \times \, n` times. + + Parameters + ---------- + data : sequence of array-like + Each element of `data` is a sample containing scalar observations from an + underlying distribution. Elements of `data` must be broadcastable to the + same shape (with the possible exception of the dimension specified by `axis`). + statistic : callable + Statistic for which the confidence interval is to be calculated. + `statistic` must be a callable that accepts ``len(data)`` samples + as separate arguments and returns the resulting statistic. + If `vectorized` is set ``True``, + `statistic` must also accept a keyword argument `axis` and be + vectorized to compute the statistic along the provided `axis`. + n_resamples : int, default: ``9999`` + The number of resamples performed to form the bootstrap distribution + of the statistic. + batch : int, optional + The number of resamples to process in each vectorized call to + `statistic`. Memory usage is O( `batch` * ``n`` ), where ``n`` is the + sample size. Default is ``None``, in which case ``batch = n_resamples`` + (or ``batch = max(n_resamples, n)`` for ``method='BCa'``). + vectorized : bool, optional + If `vectorized` is set ``False``, `statistic` will not be passed + keyword argument `axis` and is expected to calculate the statistic + only for 1D samples. If ``True``, `statistic` will be passed keyword + argument `axis` and is expected to calculate the statistic along `axis` + when passed an ND sample array. If ``None`` (default), `vectorized` + will be set ``True`` if ``axis`` is a parameter of `statistic`. Use of + a vectorized statistic typically reduces computation time. + paired : bool, default: ``False`` + Whether the statistic treats corresponding elements of the samples + in `data` as paired. If True, `bootstrap` resamples an array of + *indices* and uses the same indices for all arrays in `data`; otherwise, + `bootstrap` independently resamples the elements of each array. + axis : int, default: ``0`` + The axis of the samples in `data` along which the `statistic` is + calculated. + confidence_level : float, default: ``0.95`` + The confidence level of the confidence interval. + alternative : {'two-sided', 'less', 'greater'}, default: ``'two-sided'`` + Choose ``'two-sided'`` (default) for a two-sided confidence interval, + ``'less'`` for a one-sided confidence interval with the lower bound + at ``-np.inf``, and ``'greater'`` for a one-sided confidence interval + with the upper bound at ``np.inf``. The other bound of the one-sided + confidence intervals is the same as that of a two-sided confidence + interval with `confidence_level` twice as far from 1.0; e.g. the upper + bound of a 95% ``'less'`` confidence interval is the same as the upper + bound of a 90% ``'two-sided'`` confidence interval. + method : {'percentile', 'basic', 'bca'}, default: ``'BCa'`` + Whether to return the 'percentile' bootstrap confidence interval + (``'percentile'``), the 'basic' (AKA 'reverse') bootstrap confidence + interval (``'basic'``), or the bias-corrected and accelerated bootstrap + confidence interval (``'BCa'``). + bootstrap_result : BootstrapResult, optional + Provide the result object returned by a previous call to `bootstrap` + to include the previous bootstrap distribution in the new bootstrap + distribution. This can be used, for example, to change + `confidence_level`, change `method`, or see the effect of performing + additional resampling without repeating computations. + rng : `numpy.random.Generator`, optional + Pseudorandom number generator state. When `rng` is None, a new + `numpy.random.Generator` is created using entropy from the + operating system. Types other than `numpy.random.Generator` are + passed to `numpy.random.default_rng` to instantiate a ``Generator``. + + Returns + ------- + res : BootstrapResult + An object with attributes: + + confidence_interval : ConfidenceInterval + The bootstrap confidence interval as an instance of + `collections.namedtuple` with attributes `low` and `high`. + bootstrap_distribution : ndarray + The bootstrap distribution, that is, the value of `statistic` for + each resample. The last dimension corresponds with the resamples + (e.g. ``res.bootstrap_distribution.shape[-1] == n_resamples``). + standard_error : float or ndarray + The bootstrap standard error, that is, the sample standard + deviation of the bootstrap distribution. + + Warns + ----- + `~scipy.stats.DegenerateDataWarning` + Generated when ``method='BCa'`` and the bootstrap distribution is + degenerate (e.g. all elements are identical). + + Notes + ----- + Elements of the confidence interval may be NaN for ``method='BCa'`` if + the bootstrap distribution is degenerate (e.g. all elements are identical). + In this case, consider using another `method` or inspecting `data` for + indications that other analysis may be more appropriate (e.g. all + observations are identical). + + References + ---------- + .. [1] B. Efron and R. J. Tibshirani, An Introduction to the Bootstrap, + Chapman & Hall/CRC, Boca Raton, FL, USA (1993) + .. [2] Nathaniel E. Helwig, "Bootstrap Confidence Intervals", + http://users.stat.umn.edu/~helwig/notes/bootci-Notes.pdf + .. [3] Bootstrapping (statistics), Wikipedia, + https://en.wikipedia.org/wiki/Bootstrapping_%28statistics%29 + + Examples + -------- + Suppose we have sampled data from an unknown distribution. + + >>> import numpy as np + >>> rng = np.random.default_rng() + >>> from scipy.stats import norm + >>> dist = norm(loc=2, scale=4) # our "unknown" distribution + >>> data = dist.rvs(size=100, random_state=rng) + + We are interested in the standard deviation of the distribution. + + >>> std_true = dist.std() # the true value of the statistic + >>> print(std_true) + 4.0 + >>> std_sample = np.std(data) # the sample statistic + >>> print(std_sample) + 3.9460644295563863 + + The bootstrap is used to approximate the variability we would expect if we + were to repeatedly sample from the unknown distribution and calculate the + statistic of the sample each time. It does this by repeatedly resampling + values *from the original sample* with replacement and calculating the + statistic of each resample. This results in a "bootstrap distribution" of + the statistic. + + >>> import matplotlib.pyplot as plt + >>> from scipy.stats import bootstrap + >>> data = (data,) # samples must be in a sequence + >>> res = bootstrap(data, np.std, confidence_level=0.9, rng=rng) + >>> fig, ax = plt.subplots() + >>> ax.hist(res.bootstrap_distribution, bins=25) + >>> ax.set_title('Bootstrap Distribution') + >>> ax.set_xlabel('statistic value') + >>> ax.set_ylabel('frequency') + >>> plt.show() + + The standard error quantifies this variability. It is calculated as the + standard deviation of the bootstrap distribution. + + >>> res.standard_error + 0.24427002125829136 + >>> res.standard_error == np.std(res.bootstrap_distribution, ddof=1) + True + + The bootstrap distribution of the statistic is often approximately normal + with scale equal to the standard error. + + >>> x = np.linspace(3, 5) + >>> pdf = norm.pdf(x, loc=std_sample, scale=res.standard_error) + >>> fig, ax = plt.subplots() + >>> ax.hist(res.bootstrap_distribution, bins=25, density=True) + >>> ax.plot(x, pdf) + >>> ax.set_title('Normal Approximation of the Bootstrap Distribution') + >>> ax.set_xlabel('statistic value') + >>> ax.set_ylabel('pdf') + >>> plt.show() + + This suggests that we could construct a 90% confidence interval on the + statistic based on quantiles of this normal distribution. + + >>> norm.interval(0.9, loc=std_sample, scale=res.standard_error) + (3.5442759991341726, 4.3478528599786) + + Due to central limit theorem, this normal approximation is accurate for a + variety of statistics and distributions underlying the samples; however, + the approximation is not reliable in all cases. Because `bootstrap` is + designed to work with arbitrary underlying distributions and statistics, + it uses more advanced techniques to generate an accurate confidence + interval. + + >>> print(res.confidence_interval) + ConfidenceInterval(low=3.57655333533867, high=4.382043696342881) + + If we sample from the original distribution 100 times and form a bootstrap + confidence interval for each sample, the confidence interval + contains the true value of the statistic approximately 90% of the time. + + >>> n_trials = 100 + >>> ci_contains_true_std = 0 + >>> for i in range(n_trials): + ... data = (dist.rvs(size=100, random_state=rng),) + ... res = bootstrap(data, np.std, confidence_level=0.9, + ... n_resamples=999, rng=rng) + ... ci = res.confidence_interval + ... if ci[0] < std_true < ci[1]: + ... ci_contains_true_std += 1 + >>> print(ci_contains_true_std) + 88 + + Rather than writing a loop, we can also determine the confidence intervals + for all 100 samples at once. + + >>> data = (dist.rvs(size=(n_trials, 100), random_state=rng),) + >>> res = bootstrap(data, np.std, axis=-1, confidence_level=0.9, + ... n_resamples=999, rng=rng) + >>> ci_l, ci_u = res.confidence_interval + + Here, `ci_l` and `ci_u` contain the confidence interval for each of the + ``n_trials = 100`` samples. + + >>> print(ci_l[:5]) + [3.86401283 3.33304394 3.52474647 3.54160981 3.80569252] + >>> print(ci_u[:5]) + [4.80217409 4.18143252 4.39734707 4.37549713 4.72843584] + + And again, approximately 90% contain the true value, ``std_true = 4``. + + >>> print(np.sum((ci_l < std_true) & (std_true < ci_u))) + 93 + + `bootstrap` can also be used to estimate confidence intervals of + multi-sample statistics. For example, to get a confidence interval + for the difference between means, we write a function that accepts + two sample arguments and returns only the statistic. The use of the + ``axis`` argument ensures that all mean calculations are perform in + a single vectorized call, which is faster than looping over pairs + of resamples in Python. + + >>> def my_statistic(sample1, sample2, axis=-1): + ... mean1 = np.mean(sample1, axis=axis) + ... mean2 = np.mean(sample2, axis=axis) + ... return mean1 - mean2 + + Here, we use the 'percentile' method with the default 95% confidence level. + + >>> sample1 = norm.rvs(scale=1, size=100, random_state=rng) + >>> sample2 = norm.rvs(scale=2, size=100, random_state=rng) + >>> data = (sample1, sample2) + >>> res = bootstrap(data, my_statistic, method='basic', rng=rng) + >>> print(my_statistic(sample1, sample2)) + 0.16661030792089523 + >>> print(res.confidence_interval) + ConfidenceInterval(low=-0.29087973240818693, high=0.6371338699912273) + + The bootstrap estimate of the standard error is also available. + + >>> print(res.standard_error) + 0.238323948262459 + + Paired-sample statistics work, too. For example, consider the Pearson + correlation coefficient. + + >>> from scipy.stats import pearsonr + >>> n = 100 + >>> x = np.linspace(0, 10, n) + >>> y = x + rng.uniform(size=n) + >>> print(pearsonr(x, y)[0]) # element 0 is the statistic + 0.9954306665125647 + + We wrap `pearsonr` so that it returns only the statistic, ensuring + that we use the `axis` argument because it is available. + + >>> def my_statistic(x, y, axis=-1): + ... return pearsonr(x, y, axis=axis)[0] + + We call `bootstrap` using ``paired=True``. + + >>> res = bootstrap((x, y), my_statistic, paired=True, rng=rng) + >>> print(res.confidence_interval) + ConfidenceInterval(low=0.9941504301315878, high=0.996377412215445) + + The result object can be passed back into `bootstrap` to perform additional + resampling: + + >>> len(res.bootstrap_distribution) + 9999 + >>> res = bootstrap((x, y), my_statistic, paired=True, + ... n_resamples=1000, rng=rng, + ... bootstrap_result=res) + >>> len(res.bootstrap_distribution) + 10999 + + or to change the confidence interval options: + + >>> res2 = bootstrap((x, y), my_statistic, paired=True, + ... n_resamples=0, rng=rng, bootstrap_result=res, + ... method='percentile', confidence_level=0.9) + >>> np.testing.assert_equal(res2.bootstrap_distribution, + ... res.bootstrap_distribution) + >>> res.confidence_interval + ConfidenceInterval(low=0.9941574828235082, high=0.9963781698210212) + + without repeating computation of the original bootstrap distribution. + + """ + # Input validation + args = _bootstrap_iv(data, statistic, vectorized, paired, axis, + confidence_level, alternative, n_resamples, batch, + method, bootstrap_result, rng) + (data, statistic, vectorized, paired, axis, confidence_level, + alternative, n_resamples, batch, method, bootstrap_result, + rng, xp) = args + + theta_hat_b = ([] if bootstrap_result is None + else [bootstrap_result.bootstrap_distribution]) + + batch_nominal = batch or n_resamples or 1 + + for k in range(0, n_resamples, batch_nominal): + batch_actual = min(batch_nominal, n_resamples-k) + # Generate resamples + resampled_data = [] + for sample in data: + resample = _bootstrap_resample(sample, n_resamples=batch_actual, + rng=rng, xp=xp) + resampled_data.append(resample) + + # Compute bootstrap distribution of statistic + theta_hat_b.append(statistic(*resampled_data, axis=-1)) + theta_hat_b = xp.concat(theta_hat_b, axis=-1) + + # Calculate percentile interval + alpha = ((1 - confidence_level)/2 if alternative == 'two-sided' + else (1 - confidence_level)) + if method == 'bca': + interval = _bca_interval(data, statistic, axis=-1, alpha=alpha, + theta_hat_b=theta_hat_b, batch=batch, xp=xp)[:2] + else: + alpha = xp.asarray(alpha, dtype=theta_hat_b.dtype, + device=xp_device(theta_hat_b)) + interval = alpha, 1 - alpha + + # Calculate confidence interval of statistic + interval = xp.stack(interval, axis=-1) + ci = stats.quantile(theta_hat_b, interval, axis=-1) + if not is_lazy_array(ci) and xp.any(xp.isnan(ci)): + msg = ( + "The BCa confidence interval cannot be calculated. " + "This problem is known to occur when the distribution " + "is degenerate or the statistic is np.min." + ) + warnings.warn(DegenerateDataWarning(msg), stacklevel=2) + + ci_l = ci[..., 0] + ci_u = ci[..., 1] + + if method == 'basic': # see [3] + theta_hat = statistic(*data, axis=-1) + ci_l, ci_u = 2*theta_hat - ci_u, 2*theta_hat - ci_l + + if alternative == 'less': + ci_l = xp.full_like(ci_l, -xp.inf) + elif alternative == 'greater': + ci_u = xp.full_like(ci_u, xp.inf) + + standard_error = xp.std(theta_hat_b, correction=1, axis=-1) + + ci_l = ci_l[()] if ci_l.ndim == 0 else ci_l + ci_u = ci_u[()] if ci_u.ndim == 0 else ci_u + standard_error = standard_error[()] if standard_error.ndim == 0 else standard_error + + return BootstrapResult(confidence_interval=ConfidenceInterval(ci_l, ci_u), + bootstrap_distribution=theta_hat_b, + standard_error=standard_error) + + +def _monte_carlo_test_iv(data, rvs, statistic, vectorized, n_resamples, + batch, alternative, axis): + """Input validation for `monte_carlo_test`.""" + axis_int = int(axis) + if axis != axis_int: + raise ValueError("`axis` must be an integer.") + + if vectorized not in {True, False, None}: + raise ValueError("`vectorized` must be `True`, `False`, or `None`.") + + if not isinstance(rvs, Sequence): + rvs = (rvs,) + data = (data,) + for rvs_i in rvs: + if not callable(rvs_i): + raise TypeError("`rvs` must be callable or sequence of callables.") + + # At this point, `data` should be a sequence + # If it isn't, the user passed a sequence for `rvs` but not `data` + message = "If `rvs` is a sequence, `len(rvs)` must equal `len(data)`." + try: + len(data) + except TypeError as e: + raise ValueError(message) from e + if not len(rvs) == len(data): + raise ValueError(message) + + if not callable(statistic): + raise TypeError("`statistic` must be callable.") + + if vectorized is None: + try: + signature = inspect.signature(statistic).parameters + except ValueError as e: + message = (f"Signature inspection of {statistic=} failed; " + "pass `vectorize` explicitly.") + raise ValueError(message) from e + vectorized = 'axis' in signature + + xp = array_namespace(*data) + dtype = xp_result_type(*data, force_floating=True, xp=xp) + + if not vectorized: + if is_numpy(xp): + statistic_vectorized = _vectorize_statistic(statistic) + else: + message = ("`statistic` must be vectorized (i.e. support an `axis` " + f"argument) when `data` contains {xp.__name__} arrays.") + raise ValueError(message) + else: + statistic_vectorized = statistic + + data = _broadcast_arrays(data, axis, xp=xp) + data_iv = [] + for sample in data: + sample = xp.broadcast_to(sample, (1,)) if sample.ndim == 0 else sample + sample = xp.moveaxis(sample, axis_int, -1) + data_iv.append(sample) + + n_resamples_int = int(n_resamples) + if n_resamples != n_resamples_int or n_resamples_int <= 0: + raise ValueError("`n_resamples` must be a positive integer.") + + if batch is None: + batch_iv = batch + else: + batch_iv = int(batch) + if batch != batch_iv or batch_iv <= 0: + raise ValueError("`batch` must be a positive integer or None.") + + alternatives = {'two-sided', 'greater', 'less'} + alternative = alternative.lower() + if alternative not in alternatives: + raise ValueError(f"`alternative` must be in {alternatives}") + + return (data_iv, rvs, statistic_vectorized, vectorized, n_resamples_int, + batch_iv, alternative, axis_int, dtype, xp) + + +@dataclass +class MonteCarloTestResult: + """Result object returned by `scipy.stats.monte_carlo_test`. + + Attributes + ---------- + statistic : float or ndarray + The observed test statistic of the sample. + pvalue : float or ndarray + The p-value for the given alternative. + null_distribution : ndarray + The values of the test statistic generated under the null + hypothesis. + """ + statistic: float | np.ndarray + pvalue: float | np.ndarray + null_distribution: np.ndarray + + +@xp_capabilities() +@_rename_parameter('sample', 'data') +def monte_carlo_test(data, rvs, statistic, *, vectorized=None, + n_resamples=9999, batch=None, alternative="two-sided", + axis=0): + r"""Perform a Monte Carlo hypothesis test. + + `data` contains a sample or a sequence of one or more samples. `rvs` + specifies the distribution(s) of the sample(s) in `data` under the null + hypothesis. The value of `statistic` for the given `data` is compared + against a Monte Carlo null distribution: the value of the statistic for + each of `n_resamples` sets of samples generated using `rvs`. This gives + the p-value, the probability of observing such an extreme value of the + test statistic under the null hypothesis. + + Parameters + ---------- + data : array-like or sequence of array-like + An array or sequence of arrays of observations. + rvs : callable or tuple of callables + A callable or sequence of callables that generates random variates + under the null hypothesis. Each element of `rvs` must be a callable + that accepts keyword argument ``size`` (e.g. ``rvs(size=(m, n))``) and + returns an N-d array sample of that shape. If `rvs` is a sequence, the + number of callables in `rvs` must match the number of samples in + `data`, i.e. ``len(rvs) == len(data)``. If `rvs` is a single callable, + `data` is treated as a single sample. + statistic : callable + Statistic for which the p-value of the hypothesis test is to be + calculated. `statistic` must be a callable that accepts a sample + (e.g. ``statistic(sample)``) or ``len(rvs)`` separate samples (e.g. + ``statistic(samples1, sample2)`` if `rvs` contains two callables and + `data` contains two samples) and returns the resulting statistic. + If `vectorized` is set ``True``, `statistic` must also accept a keyword + argument `axis` and be vectorized to compute the statistic along the + provided `axis` of the samples in `data`. + vectorized : bool, optional + If `vectorized` is set ``False``, `statistic` will not be passed + keyword argument `axis` and is expected to calculate the statistic + only for 1D samples. If ``True``, `statistic` will be passed keyword + argument `axis` and is expected to calculate the statistic along `axis` + when passed ND sample arrays. If ``None`` (default), `vectorized` + will be set ``True`` if ``axis`` is a parameter of `statistic`. Use of + a vectorized statistic typically reduces computation time. + n_resamples : int, default: 9999 + Number of samples drawn from each of the callables of `rvs`. + Equivalently, the number statistic values under the null hypothesis + used as the Monte Carlo null distribution. + batch : int, optional + The number of Monte Carlo samples to process in each call to + `statistic`. Memory usage is O( `batch` * ``sample.size[axis]`` ). Default + is ``None``, in which case `batch` equals `n_resamples`. + alternative : {'two-sided', 'less', 'greater'} + The alternative hypothesis for which the p-value is calculated. + For each alternative, the p-value is defined as follows. + + - ``'greater'`` : the percentage of the null distribution that is + greater than or equal to the observed value of the test statistic. + - ``'less'`` : the percentage of the null distribution that is + less than or equal to the observed value of the test statistic. + - ``'two-sided'`` : twice the smaller of the p-values above. + + axis : int, default: 0 + The axis of `data` (or each sample within `data`) over which to + calculate the statistic. + + Returns + ------- + res : MonteCarloTestResult + An object with attributes: + + statistic : float or ndarray + The test statistic of the observed `data`. + pvalue : float or ndarray + The p-value for the given alternative. + null_distribution : ndarray + The values of the test statistic generated under the null + hypothesis. + + .. warning:: + The p-value is calculated by counting the elements of the null + distribution that are as extreme or more extreme than the observed + value of the statistic. Due to the use of finite precision arithmetic, + some statistic functions return numerically distinct values when the + theoretical values would be exactly equal. In some cases, this could + lead to a large error in the calculated p-value. `monte_carlo_test` + guards against this by considering elements in the null distribution + that are "close" (within a relative tolerance of 100 times the + floating point epsilon of inexact dtypes) to the observed + value of the test statistic as equal to the observed value of the + test statistic. However, the user is advised to inspect the null + distribution to assess whether this method of comparison is + appropriate, and if not, calculate the p-value manually. + + References + ---------- + + .. [1] B. Phipson and G. K. Smyth. "Permutation P-values Should Never Be + Zero: Calculating Exact P-values When Permutations Are Randomly Drawn." + Statistical Applications in Genetics and Molecular Biology 9.1 (2010). + + Examples + -------- + + Suppose we wish to test whether a small sample has been drawn from a normal + distribution. We decide that we will use the skew of the sample as a + test statistic, and we will consider a p-value of 0.05 to be statistically + significant. + + >>> import numpy as np + >>> from scipy import stats + >>> def statistic(x, axis): + ... return stats.skew(x, axis) + + After collecting our data, we calculate the observed value of the test + statistic. + + >>> rng = np.random.default_rng() + >>> x = stats.skewnorm.rvs(a=1, size=50, random_state=rng) + >>> statistic(x, axis=0) + 0.12457412450240658 + + To determine the probability of observing such an extreme value of the + skewness by chance if the sample were drawn from the normal distribution, + we can perform a Monte Carlo hypothesis test. The test will draw many + samples at random from their normal distribution, calculate the skewness + of each sample, and compare our original skewness against this + distribution to determine an approximate p-value. + + >>> from scipy.stats import monte_carlo_test + >>> # because our statistic is vectorized, we pass `vectorized=True` + >>> rvs = lambda size: stats.norm.rvs(size=size, random_state=rng) + >>> res = monte_carlo_test(x, rvs, statistic, vectorized=True) + >>> print(res.statistic) + 0.12457412450240658 + >>> print(res.pvalue) + 0.7012 + + The probability of obtaining a test statistic less than or equal to the + observed value under the null hypothesis is ~70%. This is greater than + our chosen threshold of 5%, so we cannot consider this to be significant + evidence against the null hypothesis. + + Note that this p-value essentially matches that of + `scipy.stats.skewtest`, which relies on an asymptotic distribution of a + test statistic based on the sample skewness. + + >>> stats.skewtest(x).pvalue + 0.6892046027110614 + + This asymptotic approximation is not valid for small sample sizes, but + `monte_carlo_test` can be used with samples of any size. + + >>> x = stats.skewnorm.rvs(a=1, size=7, random_state=rng) + >>> # stats.skewtest(x) would produce an error due to small sample + >>> res = monte_carlo_test(x, rvs, statistic, vectorized=True) + + The Monte Carlo distribution of the test statistic is provided for + further investigation. + + >>> import matplotlib.pyplot as plt + >>> fig, ax = plt.subplots() + >>> ax.hist(res.null_distribution, bins=50) + >>> ax.set_title("Monte Carlo distribution of test statistic") + >>> ax.set_xlabel("Value of Statistic") + >>> ax.set_ylabel("Frequency") + >>> plt.show() + + """ + args = _monte_carlo_test_iv(data, rvs, statistic, vectorized, + n_resamples, batch, alternative, axis) + (data, rvs, statistic, vectorized, n_resamples, + batch, alternative, axis, dtype, xp) = args + + # Some statistics return plain floats; ensure they're at least a NumPy float + observed = xp.asarray(statistic(*data, axis=-1)) + observed = observed[()] if observed.ndim == 0 else observed + + n_observations = [sample.shape[-1] for sample in data] + batch_nominal = batch or n_resamples + null_distribution = [] + for k in range(0, n_resamples, batch_nominal): + batch_actual = min(batch_nominal, n_resamples - k) + resamples = [rvs_i(size=(batch_actual, n_observations_i)) + for rvs_i, n_observations_i in zip(rvs, n_observations)] + null_distribution.append(statistic(*resamples, axis=-1)) + null_distribution = xp.concat(null_distribution) + null_distribution = xp.reshape(null_distribution, (-1,) + (1,)*observed.ndim) + + # relative tolerance for detecting numerically distinct but + # theoretically equal values in the null distribution + eps = (0 if not xp.isdtype(observed.dtype, ('real floating')) + else xp.finfo(observed.dtype).eps*100) + gamma = xp.abs(eps * observed) + + def less(null_distribution, observed): + cmps = null_distribution <= observed + gamma + cmps = xp.asarray(cmps, dtype=dtype) + pvalues = (xp.sum(cmps, axis=0, dtype=dtype) + 1.) / (n_resamples + 1.) + return pvalues + + def greater(null_distribution, observed): + cmps = null_distribution >= observed - gamma + cmps = xp.asarray(cmps, dtype=dtype) + pvalues = (xp.sum(cmps, axis=0, dtype=dtype) + 1.) / (n_resamples + 1.) + return pvalues + + def two_sided(null_distribution, observed): + pvalues_less = less(null_distribution, observed) + pvalues_greater = greater(null_distribution, observed) + pvalues = xp.minimum(pvalues_less, pvalues_greater) * 2 + return pvalues + + compare = {"less": less, + "greater": greater, + "two-sided": two_sided} + + pvalues = compare[alternative](null_distribution, observed) + pvalues = xp.clip(pvalues, 0., 1.) + + return MonteCarloTestResult(observed, pvalues, null_distribution) + + +@dataclass +class PowerResult: + """Result object returned by `scipy.stats.power`. + + Attributes + ---------- + power : float or ndarray + The estimated power. + pvalues : float or ndarray + The simulated p-values. + """ + power: float | np.ndarray + pvalues: float | np.ndarray + + +def _wrap_kwargs(fun): + """Wrap callable to accept arbitrary kwargs and ignore unused ones""" + + try: + keys = set(inspect.signature(fun).parameters.keys()) + except ValueError: + # NumPy Generator methods can't be inspected + keys = {'size'} + + # Set keys=keys/fun=fun to avoid late binding gotcha + def wrapped_rvs_i(*args, keys=keys, fun=fun, **all_kwargs): + kwargs = {key: val for key, val in all_kwargs.items() + if key in keys} + return fun(*args, **kwargs) + return wrapped_rvs_i + + +def _power_iv(rvs, test, n_observations, significance, vectorized, + n_resamples, batch, kwargs): + """Input validation for `monte_carlo_test`.""" + if vectorized not in {True, False, None}: + raise ValueError("`vectorized` must be `True`, `False`, or `None`.") + + if not isinstance(rvs, Sequence): + rvs = (rvs,) + n_observations = (n_observations,) + for rvs_i in rvs: + if not callable(rvs_i): + raise TypeError("`rvs` must be callable or sequence of callables.") + + if not len(rvs) == len(n_observations): + message = ("If `rvs` is a sequence, `len(rvs)` " + "must equal `len(n_observations)`.") + raise ValueError(message) + + kwargs = dict() if kwargs is None else kwargs + if not isinstance(kwargs, dict): + raise TypeError("`kwargs` must be a dictionary that maps keywords to arrays.") + + vals = kwargs.values() + keys = kwargs.keys() + + xp = array_namespace(*n_observations, significance, *vals) + + significance = xp.asarray(significance) + if not xp.isdtype(significance.dtype, "real floating"): + raise ValueError("`significance` must be of floating point dtype.") + + if is_lazy_array(significance): + significance = xp.where((significance < 0.) | (significance > 1.), + xp.nan, significance) + elif xp.min(significance) < 0. or xp.max(significance) > 1.: + raise ValueError("All elements of `significance` must be between 0. and 1.") + + # Wrap callables to ignore unused keyword arguments + wrapped_rvs = [_wrap_kwargs(rvs_i) for rvs_i in rvs] + + # Broadcast, then ravel nobs/kwarg combinations. In the end, + # `nobs` and `vals` have shape (# of combinations, number of variables) + # todo: find a better way to do this without combining arrays + tmp = xp.stack(xp.broadcast_arrays(*n_observations, *vals)) + shape = tmp.shape + if tmp.ndim == 1: + tmp = xp.expand_dims(tmp, axis=0) + else: + tmp = xp.reshape(tmp, (shape[0], -1)).T + nobs, vals = tmp[:, :len(rvs)], tmp[:, len(rvs):] + integer_dtype = xp_result_type(*n_observations, xp=xp) + nobs = xp.astype(nobs, integer_dtype) + + if not callable(test): + raise TypeError("`test` must be callable.") + + if vectorized is None: + vectorized = 'axis' in inspect.signature(test).parameters + + test_vectorized = test + if not vectorized: + if not is_numpy(xp): + message = (f"When using array library {xp.__name__}, `test` must be " + "be vectorized and accept argument `axis`.") + raise TypeError(message) + + test_vectorized = _vectorize_statistic(test) + + # Wrap `test` function to ignore unused kwargs + test_vectorized = _wrap_kwargs(test_vectorized) + + n_resamples_int = int(n_resamples) + if n_resamples != n_resamples_int or n_resamples_int <= 0: + raise ValueError("`n_resamples` must be a positive integer.") + + if batch is None: + batch_iv = batch + else: + batch_iv = int(batch) + if batch != batch_iv or batch_iv <= 0: + raise ValueError("`batch` must be a positive integer or None.") + + return (wrapped_rvs, test_vectorized, nobs, significance, vectorized, + n_resamples_int, batch_iv, vals, keys, shape[1:], xp) + + +@xp_capabilities(skip_backends=[('dask.array', 'just because')], + jax_jit=False) # some problem with batch looping +def power(test, rvs, n_observations, *, significance=0.01, vectorized=None, + n_resamples=10000, batch=None, kwargs=None): + r"""Simulate the power of a hypothesis test under an alternative hypothesis. + + Parameters + ---------- + test : callable + Hypothesis test for which the power is to be simulated. + `test` must be a callable that accepts a sample (e.g. ``test(sample)``) + or ``len(rvs)`` separate samples (e.g. ``test(samples1, sample2)`` if + `rvs` contains two callables and `n_observations` contains two values) + and returns the p-value of the test. + If `vectorized` is set to ``True``, `test` must also accept a keyword + argument `axis` and be vectorized to perform the test along the + provided `axis` of the samples. + Any callable from `scipy.stats` with an `axis` argument that returns an + object with a `pvalue` attribute is also acceptable. + rvs : callable or tuple of callables + A callable or sequence of callables that generate(s) random variates + under the alternative hypothesis. Each element of `rvs` must accept + keyword argument ``size`` (e.g. ``rvs(size=(m, n))``) and return an + N-d array of that shape. If `rvs` is a sequence, the number of callables + in `rvs` must match the number of elements of `n_observations`, i.e. + ``len(rvs) == len(n_observations)``. If `rvs` is a single callable, + `n_observations` is treated as a single element. + n_observations : tuple of ints or tuple of int arrays + If a sequence of ints, each is the sizes of a sample to be passed to `test`. + If a sequence of integer arrays, the power is simulated for each + set of corresponding sample sizes. See Examples. + significance : float or array_like of floats, default: 0.01 + The threshold for significance; i.e., the p-value below which the + hypothesis test results will be considered as evidence against the null + hypothesis. Equivalently, the acceptable rate of Type I error under + the null hypothesis. If an array, the power is simulated for each + significance threshold. + vectorized : bool, optional + If `vectorized` is set to ``False``, `test` will not be passed keyword + argument `axis` and is expected to perform the test only for 1D samples. + If ``True``, `test` will be passed keyword argument `axis` and is + expected to perform the test along `axis` when passed N-D sample arrays. + If ``None`` (default), `vectorized` will be set ``True`` if ``axis`` is + a parameter of `test`. Use of a vectorized test typically reduces + computation time. + n_resamples : int, default: 10000 + Number of samples drawn from each of the callables of `rvs`. + Equivalently, the number tests performed under the alternative + hypothesis to approximate the power. + batch : int, optional + The number of samples to process in each call to `test`. Memory usage is + proportional to the product of `batch` and the largest sample size. Default + is ``None``, in which case `batch` equals `n_resamples`. + kwargs : dict, optional + Keyword arguments to be passed to `rvs` and/or `test` callables. + Introspection is used to determine which keyword arguments may be + passed to each callable. + The value corresponding with each keyword must be an array. + Arrays must be broadcastable with one another and with each array in + `n_observations`. The power is simulated for each set of corresponding + sample sizes and arguments. See Examples. + + Returns + ------- + res : PowerResult + An object with attributes: + + power : float or ndarray + The estimated power against the alternative. + pvalues : ndarray + The p-values observed under the alternative hypothesis. + + Notes + ----- + The power is simulated as follows: + + - Draw many random samples (or sets of samples), each of the size(s) + specified by `n_observations`, under the alternative specified by + `rvs`. + - For each sample (or set of samples), compute the p-value according to + `test`. These p-values are recorded in the ``pvalues`` attribute of + the result object. + - Compute the proportion of p-values that are less than the `significance` + level. This is the power recorded in the ``power`` attribute of the + result object. + + Suppose that `significance` is an array with shape ``shape1``, the elements + of `kwargs` and `n_observations` are mutually broadcastable to shape ``shape2``, + and `test` returns an array of p-values of shape ``shape3``. Then the result + object ``power`` attribute will be of shape ``shape1 + shape2 + shape3``, and + the ``pvalues`` attribute will be of shape ``shape2 + shape3 + (n_resamples,)``. + + Examples + -------- + Suppose we wish to simulate the power of the independent sample t-test + under the following conditions: + + - The first sample has 10 observations drawn from a normal distribution + with mean 0. + - The second sample has 12 observations drawn from a normal distribution + with mean 1.0. + - The threshold on p-values for significance is 0.05. + + >>> import numpy as np + >>> from scipy import stats + >>> rng = np.random.default_rng(2549598345528) + >>> + >>> test = stats.ttest_ind + >>> n_observations = (10, 12) + >>> rvs1 = rng.normal + >>> rvs2 = lambda size: rng.normal(loc=1, size=size) + >>> rvs = (rvs1, rvs2) + >>> res = stats.power(test, rvs, n_observations, significance=0.05) + >>> res.power + 0.6116 + + With samples of size 10 and 12, respectively, the power of the t-test + with a significance threshold of 0.05 is approximately 60% under the chosen + alternative. We can investigate the effect of sample size on the power + by passing sample size arrays. + + >>> import matplotlib.pyplot as plt + >>> nobs_x = np.arange(5, 21) + >>> nobs_y = nobs_x + >>> n_observations = (nobs_x, nobs_y) + >>> res = stats.power(test, rvs, n_observations, significance=0.05) + >>> ax = plt.subplot() + >>> ax.plot(nobs_x, res.power) + >>> ax.set_xlabel('Sample Size') + >>> ax.set_ylabel('Simulated Power') + >>> ax.set_title('Simulated Power of `ttest_ind` with Equal Sample Sizes') + >>> plt.show() + + Alternatively, we can investigate the impact that effect size has on the power. + In this case, the effect size is the location of the distribution underlying + the second sample. + + >>> n_observations = (10, 12) + >>> loc = np.linspace(0, 1, 20) + >>> rvs2 = lambda size, loc: rng.normal(loc=loc, size=size) + >>> rvs = (rvs1, rvs2) + >>> res = stats.power(test, rvs, n_observations, significance=0.05, + ... kwargs={'loc': loc}) + >>> ax = plt.subplot() + >>> ax.plot(loc, res.power) + >>> ax.set_xlabel('Effect Size') + >>> ax.set_ylabel('Simulated Power') + >>> ax.set_title('Simulated Power of `ttest_ind`, Varying Effect Size') + >>> plt.show() + + We can also use `power` to estimate the Type I error rate (also referred to by the + ambiguous term "size") of a test and assess whether it matches the nominal level. + For example, the null hypothesis of `jarque_bera` is that the sample was drawn from + a distribution with the same skewness and kurtosis as the normal distribution. To + estimate the Type I error rate, we can consider the null hypothesis to be a true + *alternative* hypothesis and calculate the power. + + >>> test = stats.jarque_bera + >>> n_observations = 10 + >>> rvs = rng.normal + >>> significance = np.linspace(0.0001, 0.1, 1000) + >>> res = stats.power(test, rvs, n_observations, significance=significance) + >>> size = res.power + + As shown below, the Type I error rate of the test is far below the nominal level + for such a small sample, as mentioned in its documentation. + + >>> ax = plt.subplot() + >>> ax.plot(significance, size) + >>> ax.plot([0, 0.1], [0, 0.1], '--') + >>> ax.set_xlabel('nominal significance level') + >>> ax.set_ylabel('estimated test size (Type I error rate)') + >>> ax.set_title('Estimated test size vs nominal significance level') + >>> ax.set_aspect('equal', 'box') + >>> ax.legend(('`ttest_1samp`', 'ideal test')) + >>> plt.show() + + As one might expect from such a conservative test, the power is quite low with + respect to some alternatives. For example, the power of the test under the + alternative that the sample was drawn from the Laplace distribution may not + be much greater than the Type I error rate. + + >>> rvs = rng.laplace + >>> significance = np.linspace(0.0001, 0.1, 1000) + >>> res = stats.power(test, rvs, n_observations, significance=0.05) + >>> print(res.power) + 0.0587 + + This is not a mistake in SciPy's implementation; it is simply due to the fact + that the null distribution of the test statistic is derived under the assumption + that the sample size is large (i.e. approaches infinity), and this asymptotic + approximation is not accurate for small samples. In such cases, resampling + and Monte Carlo methods (e.g. `permutation_test`, `goodness_of_fit`, + `monte_carlo_test`) may be more appropriate. + + """ + tmp = _power_iv(rvs, test, n_observations, significance, + vectorized, n_resamples, batch, kwargs) + (rvs, test, nobs, significance, + vectorized, n_resamples, batch, args, kwds, shape, xp) = tmp + + batch_nominal = batch or n_resamples + pvalues = [] # results of various nobs/kwargs combinations + for i in range(nobs.shape[0]): + nobs_i, args_i = nobs[i, ...], args[i, ...] + kwargs_i = dict(zip(kwds, args_i)) + pvalues_i = [] # results of batches; fixed nobs/kwargs combination + for k in range(0, n_resamples, batch_nominal): + batch_actual = min(batch_nominal, n_resamples - k) + resamples = [rvs_j(size=(batch_actual, int(nobs_ij)), **kwargs_i) + for rvs_j, nobs_ij in zip(rvs, nobs_i)] + res = test(*resamples, **kwargs_i, axis=-1) + p = getattr(res, 'pvalue', res) + pvalues_i.append(p) + # Concatenate results from batches + pvalues_i = xp.concat(pvalues_i, axis=-1) + pvalues.append(pvalues_i) + # `test` can return result with array of p-values + shape += pvalues_i.shape[:-1] + # Concatenate results from various nobs/kwargs combinations + pvalues = xp.concat(pvalues, axis=0) + # nobs/kwargs arrays were raveled to single axis; unravel + pvalues = xp.reshape(pvalues, shape + (-1,)) + if significance.ndim > 0: + newdims = tuple(range(significance.ndim, pvalues.ndim + significance.ndim)) + significance = xpx.expand_dims(significance, axis=newdims) + + float_dtype = xp_result_type(significance, pvalues, xp=xp) + powers = xp.mean(xp.astype(pvalues < significance, float_dtype), axis=-1) + + return PowerResult(power=powers, pvalues=pvalues) + + +@dataclass +class PermutationTestResult: + """Result object returned by `scipy.stats.permutation_test`. + + Attributes + ---------- + statistic : float or ndarray + The observed test statistic of the data. + pvalue : float or ndarray + The p-value for the given alternative. + null_distribution : ndarray + The values of the test statistic generated under the null + hypothesis. + """ + statistic: float | np.ndarray + pvalue: float | np.ndarray + null_distribution: np.ndarray + + +def _all_partitions_concatenated(ns): + """ + Generate all partitions of indices of groups of given sizes, concatenated + + `ns` is an iterable of ints. + """ + def all_partitions(z, n): + for c in combinations(z, n): + x0 = set(c) + x1 = z - x0 + yield [x0, x1] + + def all_partitions_n(z, ns): + if len(ns) == 0: + yield [z] + return + for c in all_partitions(z, ns[0]): + for d in all_partitions_n(c[1], ns[1:]): + yield c[0:1] + d + + z = set(range(np.sum(ns))) + for partitioning in all_partitions_n(z, ns[:]): + x = np.concatenate([list(partition) + for partition in partitioning]).astype(int) + yield x + + +def _batch_generator(iterable, batch): + """A generator that yields batches of elements from an iterable""" + iterator = iter(iterable) + if batch <= 0: + raise ValueError("`batch` must be positive.") + z = [item for i, item in zip(range(batch), iterator)] + while z: # we don't want StopIteration without yielding an empty list + yield z + z = [item for i, item in zip(range(batch), iterator)] + + +def _pairings_permutations_gen(n_permutations, n_samples, n_obs_sample, batch, + rng): + # Returns a generator that yields arrays of size + # `(batch, n_samples, n_obs_sample)`. + # Each row is an independent permutation of indices 0 to `n_obs_sample`. + batch = min(batch, n_permutations) + + if hasattr(rng, 'permuted'): + def batched_perm_generator(): + indices = np.arange(n_obs_sample) + indices = np.tile(indices, (batch, n_samples, 1)) + for k in range(0, n_permutations, batch): + batch_actual = min(batch, n_permutations-k) + # Don't permute in place, otherwise results depend on `batch` + permuted_indices = rng.permuted(indices, axis=-1) + yield permuted_indices[:batch_actual] + else: # RandomState and early Generators don't have `permuted` + def batched_perm_generator(): + for k in range(0, n_permutations, batch): + batch_actual = min(batch, n_permutations-k) + size = (batch_actual, n_samples, n_obs_sample) + x = rng.random(size=size) + yield np.argsort(x, axis=-1)[:batch_actual] + + return batched_perm_generator() + + +def _calculate_null_both(data, statistic, n_permutations, batch, + rng=None, *, xp): + """ + Calculate null distribution for independent sample tests. + """ + # compute number of permutations + # (distinct partitions of data into samples of these sizes) + n_obs_i = [sample.shape[-1] for sample in data] # observations per sample + n_obs_ic = list(accumulate(n_obs_i, initial=0)) + n_obs = n_obs_ic[-1] # total number of observations + n_max = math.prod([math.comb(n, k) for n, k in zip(n_obs_ic[1:], n_obs_ic[:-1])]) + + # perm_generator is an iterator that produces permutations of indices + # from 0 to n_obs. We'll concatenate the samples, use these indices to + # permute the data, then split the samples apart again. + if n_permutations >= n_max: + exact_test = True + n_permutations = n_max + perm_generator = _all_partitions_concatenated(n_obs_i) + else: + exact_test = False + # Neither RandomState.permutation nor Generator.permutation + # can permute axis-slices independently. If this feature is + # added in the future, batches of the desired size should be + # generated in a single call. + perm_generator = (rng.permutation(n_obs) + for i in range(n_permutations)) + + batch = batch or int(n_permutations) + null_distribution = [] + + # First, concatenate all the samples. In batches, permute samples with + # indices produced by the `perm_generator`, split them into new samples of + # the original sizes, compute the statistic for each batch, and add these + # statistic values to the null distribution. + data = xp.concat(data, axis=-1) + for indices in _batch_generator(perm_generator, batch=batch): + # Creating a tensor from a list of numpy.ndarrays is extremely slow... + indices = np.asarray(indices) + indices = xp.asarray(indices) + + # `indices` is 2D: each row is a permutation of the indices. + # We use it to index `data` along its last axis, which corresponds + # with observations. + # After indexing, the second to last axis of `data_batch` corresponds + # with permutations, and the last axis corresponds with observations. + # data_batch = data[..., indices] + data_batch = _get_from_last_axis(data, indices, xp=xp) + + # Move the permutation axis to the front: we'll concatenate a list + # of batched statistic values along this zeroth axis to form the + # null distribution. + data_batch = xp.moveaxis(data_batch, -2, 0) + # data_batch = np.split(data_batch, n_obs_ic[:-1], axis=-1) + data_batch = [data_batch[..., i:j] for i, j in zip(n_obs_ic[:-1], n_obs_ic[1:])] + null_distribution.append(statistic(*data_batch, axis=-1)) + null_distribution = xp.concat(null_distribution, axis=0) + + return null_distribution, n_permutations, exact_test + + +def _calculate_null_pairings(data, statistic, n_permutations, batch, + rng=None, *, xp): + """ + Calculate null distribution for association tests. + """ + n_samples = len(data) + + # compute number of permutations (factorial(n) permutations of each sample) + n_obs_sample = data[0].shape[-1] # observations per sample; same for each + n_max = math.factorial(n_obs_sample)**n_samples + + # `perm_generator` is an iterator that produces a list of permutations of + # indices from 0 to n_obs_sample, one for each sample. + if n_permutations >= n_max: + exact_test = True + n_permutations = n_max + batch = batch or int(n_permutations) + # Cartesian product of the sets of all permutations of indices + perm_generator = product(*(permutations(range(n_obs_sample)) + for i in range(n_samples))) + batched_perm_generator = _batch_generator(perm_generator, batch=batch) + else: + exact_test = False + batch = batch or int(n_permutations) + # Separate random permutations of indices for each sample. + # Again, it would be nice if RandomState/Generator.permutation + # could permute each axis-slice separately. + args = n_permutations, n_samples, n_obs_sample, batch, rng + batched_perm_generator = _pairings_permutations_gen(*args) + + null_distribution = [] + + for indices in batched_perm_generator: + indices = xp.asarray(indices) + + # `indices` is 3D: the zeroth axis is for permutations, the next is + # for samples, and the last is for observations. Swap the first two + # to make the zeroth axis correspond with samples, as it does for + # `data`. + indices = xp_swapaxes(indices, 0, 1, xp=xp) + + # When we're done, `data_batch` will be a list of length `n_samples`. + # Each element will be a batch of random permutations of one sample. + # The zeroth axis of each batch will correspond with permutations, + # and the last will correspond with observations. (This makes it + # easy to pass into `statistic`.) + data_batch = [None]*n_samples + for i in range(n_samples): + # data_batch[i] = data[i][..., indices[i]] + data_batch[i] = _get_from_last_axis(data[i], indices[i, ...], xp=xp) + data_batch[i] = xp.moveaxis(data_batch[i], -2, 0) + + null_distribution.append(statistic(*data_batch, axis=-1)) + null_distribution = xp.concat(null_distribution, axis=0) + + return null_distribution, n_permutations, exact_test + + +def _calculate_null_samples(data, statistic, n_permutations, batch, + rng=None, *, xp): + """ + Calculate null distribution for paired-sample tests. + """ + n_samples = len(data) + + # By convention, the meaning of the "samples" permutations type for + # data with only one sample is to flip the sign of the observations. + # Achieve this by adding a second sample - the negative of the original. + if n_samples == 1: + data = [data[0], -data[0]] + + # The "samples" permutation strategy is the same as the "pairings" + # strategy except the roles of samples and observations are flipped. + # So swap these axes, then we'll use the function for the "pairings" + # strategy to do all the work! + data = xp.stack(data, axis=0) + data = xp_swapaxes(data, 0, -1, xp=xp) + + # (Of course, the user's statistic doesn't know what we've done here, + # so we need to pass it what it's expecting.) + def statistic_wrapped(*data, axis): + # can we do this without converting back and forth between + # array and list? + data = xp.stack(data, axis=0) + data = xp_swapaxes(data, 0, -1, xp=xp) + if n_samples == 1: + data = data[0:1, ...] + data = [data[i, ...] for i in range(data.shape[0])] + return statistic(*data, axis=axis) + + data = [data[i, ...] for i in range(data.shape[0])] + return _calculate_null_pairings(data, statistic_wrapped, n_permutations, + batch, rng, xp=xp) + + +def _permutation_test_iv(data, statistic, permutation_type, vectorized, + n_resamples, batch, alternative, axis, rng): + """Input validation for `permutation_test`.""" + axis_int = int(axis) + if axis != axis_int: + raise ValueError("`axis` must be an integer.") + + permutation_types = {'samples', 'pairings', 'independent'} + permutation_type = permutation_type.lower() + if permutation_type not in permutation_types: + raise ValueError(f"`permutation_type` must be in {permutation_types}.") + + if vectorized not in {True, False, None}: + raise ValueError("`vectorized` must be `True`, `False`, or `None`.") + + if vectorized is None: + vectorized = 'axis' in inspect.signature(statistic).parameters + + message = "`data` must be a tuple containing at least two samples" + try: + if len(data) < 2 and permutation_type == 'independent': + raise ValueError(message) + except TypeError: + raise TypeError(message) + + xp = array_namespace(*data) + + if not vectorized: + if not is_numpy(xp): + message = (f"When using array library {xp.__name__}, `func` must be " + "vectorized and accept argument `axis`.") + raise TypeError(message) + + statistic = _vectorize_statistic(statistic) + + data = _broadcast_arrays(data, axis, xp=xp) + data_iv = [] + for sample in data: + sample = xpx.atleast_nd(sample, ndim=1) + if sample.shape[axis] <= 1: + raise ValueError("each sample in `data` must contain two or more " + "observations along `axis`.") + sample = xp.moveaxis(sample, axis_int, -1) + data_iv.append(sample) + + n_resamples_int = (int(n_resamples) if not math.isinf(n_resamples) + else xp.inf) + if n_resamples != n_resamples_int or n_resamples_int <= 0: + raise ValueError("`n_resamples` must be a positive integer.") + + if batch is None: + batch_iv = batch + else: + batch_iv = int(batch) + if batch != batch_iv or batch_iv <= 0: + raise ValueError("`batch` must be a positive integer or None.") + + alternatives = {'two-sided', 'greater', 'less'} + alternative = alternative.lower() + if alternative not in alternatives: + raise ValueError(f"`alternative` must be in {alternatives}") + + rng = check_random_state(rng) + + float_dtype = xp_result_type(*data_iv, force_floating=True, xp=xp) + + return (data_iv, statistic, permutation_type, vectorized, n_resamples_int, + batch_iv, alternative, axis_int, rng, float_dtype, xp) + + +@xp_capabilities(skip_backends=[('dask.array', "lacks required indexing capabilities")]) +@_transition_to_rng('random_state') +def permutation_test(data, statistic, *, permutation_type='independent', + vectorized=None, n_resamples=9999, batch=None, + alternative="two-sided", axis=0, rng=None): + r""" + Performs a permutation test of a given statistic on provided data. + + For independent sample statistics, the null hypothesis is that the data are + randomly sampled from the same distribution. + For paired sample statistics, two null hypothesis can be tested: + that the data are paired at random or that the data are assigned to samples + at random. + + Parameters + ---------- + data : iterable of array-like + Contains the samples, each of which is an array of observations. + Dimensions of sample arrays must be compatible for broadcasting except + along `axis`. + statistic : callable + Statistic for which the p-value of the hypothesis test is to be + calculated. `statistic` must be a callable that accepts samples + as separate arguments (e.g. ``statistic(*data)``) and returns the + resulting statistic. + If `vectorized` is set ``True``, `statistic` must also accept a keyword + argument `axis` and be vectorized to compute the statistic along the + provided `axis` of the sample arrays. + permutation_type : {'independent', 'samples', 'pairings'}, optional + The type of permutations to be performed, in accordance with the + null hypothesis. The first two permutation types are for paired sample + statistics, in which all samples contain the same number of + observations and observations with corresponding indices along `axis` + are considered to be paired; the third is for independent sample + statistics. + + - ``'samples'`` : observations are assigned to different samples + but remain paired with the same observations from other samples. + This permutation type is appropriate for paired sample hypothesis + tests such as the Wilcoxon signed-rank test and the paired t-test. + - ``'pairings'`` : observations are paired with different observations, + but they remain within the same sample. This permutation type is + appropriate for association/correlation tests with statistics such + as Spearman's :math:`\rho`, Kendall's :math:`\tau`, and Pearson's + :math:`r`. + - ``'independent'`` (default) : observations are assigned to different + samples. Samples may contain different numbers of observations. This + permutation type is appropriate for independent sample hypothesis + tests such as the Mann-Whitney :math:`U` test and the independent + sample t-test. + + Please see the Notes section below for more detailed descriptions + of the permutation types. + + vectorized : bool, optional + If `vectorized` is set ``False``, `statistic` will not be passed + keyword argument `axis` and is expected to calculate the statistic + only for 1D samples. If ``True``, `statistic` will be passed keyword + argument `axis` and is expected to calculate the statistic along `axis` + when passed an ND sample array. If ``None`` (default), `vectorized` + will be set ``True`` if ``axis`` is a parameter of `statistic`. Use + of a vectorized statistic typically reduces computation time. + n_resamples : int or np.inf, default: 9999 + Number of random permutations (resamples) used to approximate the null + distribution. If greater than or equal to the number of distinct + permutations, the exact null distribution will be computed. + Note that the number of distinct permutations grows very rapidly with + the sizes of samples, so exact tests are feasible only for very small + data sets. + batch : int, optional + The number of permutations to process in each call to `statistic`. + Memory usage is O( `batch` * ``n`` ), where ``n`` is the total size + of all samples, regardless of the value of `vectorized`. Default is + ``None``, in which case ``batch`` is the number of permutations. + alternative : {'two-sided', 'less', 'greater'}, optional + The alternative hypothesis for which the p-value is calculated. + For each alternative, the p-value is defined for exact tests as + follows. + + - ``'greater'`` : the percentage of the null distribution that is + greater than or equal to the observed value of the test statistic. + - ``'less'`` : the percentage of the null distribution that is + less than or equal to the observed value of the test statistic. + - ``'two-sided'`` (default) : twice the smaller of the p-values above. + + Note that p-values for randomized tests are calculated according to the + conservative (over-estimated) approximation suggested in [2]_ and [3]_ + rather than the unbiased estimator suggested in [4]_. That is, when + calculating the proportion of the randomized null distribution that is + as extreme as the observed value of the test statistic, the values in + the numerator and denominator are both increased by one. An + interpretation of this adjustment is that the observed value of the + test statistic is always included as an element of the randomized + null distribution. + The convention used for two-sided p-values is not universal; + the observed test statistic and null distribution are returned in + case a different definition is preferred. + + axis : int, default: 0 + The axis of the (broadcasted) samples over which to calculate the + statistic. If samples have a different number of dimensions, + singleton dimensions are prepended to samples with fewer dimensions + before `axis` is considered. + rng : `numpy.random.Generator`, optional + Pseudorandom number generator state. When `rng` is None, a new + `numpy.random.Generator` is created using entropy from the + operating system. Types other than `numpy.random.Generator` are + passed to `numpy.random.default_rng` to instantiate a ``Generator``. + + Returns + ------- + res : PermutationTestResult + An object with attributes: + + statistic : float or ndarray + The observed test statistic of the data. + pvalue : float or ndarray + The p-value for the given alternative. + null_distribution : ndarray + The values of the test statistic generated under the null + hypothesis. + + Notes + ----- + + The three types of permutation tests supported by this function are + described below. + + **Unpaired statistics** (``permutation_type='independent'``): + + The null hypothesis associated with this permutation type is that all + observations are sampled from the same underlying distribution and that + they have been assigned to one of the samples at random. + + Suppose ``data`` contains two samples; e.g. ``a, b = data``. + When ``1 < n_resamples < binom(n, k)``, where + + * ``k`` is the number of observations in ``a``, + * ``n`` is the total number of observations in ``a`` and ``b``, and + * ``binom(n, k)`` is the binomial coefficient (``n`` choose ``k``), + + the data are pooled (concatenated), randomly assigned to either the first + or second sample, and the statistic is calculated. This process is + performed repeatedly, `permutation` times, generating a distribution of the + statistic under the null hypothesis. The statistic of the original + data is compared to this distribution to determine the p-value. + + When ``n_resamples >= binom(n, k)``, an exact test is performed: the data + are *partitioned* between the samples in each distinct way exactly once, + and the exact null distribution is formed. + Note that for a given partitioning of the data between the samples, + only one ordering/permutation of the data *within* each sample is + considered. For statistics that do not depend on the order of the data + within samples, this dramatically reduces computational cost without + affecting the shape of the null distribution (because the frequency/count + of each value is affected by the same factor). + + For ``a = [a1, a2, a3, a4]`` and ``b = [b1, b2, b3]``, an example of this + permutation type is ``x = [b3, a1, a2, b2]`` and ``y = [a4, b1, a3]``. + Because only one ordering/permutation of the data *within* each sample + is considered in an exact test, a resampling like ``x = [b3, a1, b2, a2]`` + and ``y = [a4, a3, b1]`` would *not* be considered distinct from the + example above. + + ``permutation_type='independent'`` does not support one-sample statistics, + but it can be applied to statistics with more than two samples. In this + case, if ``n`` is an array of the number of observations within each + sample, the number of distinct partitions is:: + + np.prod([binom(sum(n[i:]), sum(n[i+1:])) for i in range(len(n)-1)]) + + **Paired statistics, permute pairings** (``permutation_type='pairings'``): + + The null hypothesis associated with this permutation type is that + observations within each sample are drawn from the same underlying + distribution and that pairings with elements of other samples are + assigned at random. + + Suppose ``data`` contains only one sample; e.g. ``a, = data``, and we + wish to consider all possible pairings of elements of ``a`` with elements + of a second sample, ``b``. Let ``n`` be the number of observations in + ``a``, which must also equal the number of observations in ``b``. + + When ``1 < n_resamples < factorial(n)``, the elements of ``a`` are + randomly permuted. The user-supplied statistic accepts one data argument, + say ``a_perm``, and calculates the statistic considering ``a_perm`` and + ``b``. This process is performed repeatedly, `permutation` times, + generating a distribution of the statistic under the null hypothesis. + The statistic of the original data is compared to this distribution to + determine the p-value. + + When ``n_resamples >= factorial(n)``, an exact test is performed: + ``a`` is permuted in each distinct way exactly once. Therefore, the + `statistic` is computed for each unique pairing of samples between ``a`` + and ``b`` exactly once. + + For ``a = [a1, a2, a3]`` and ``b = [b1, b2, b3]``, an example of this + permutation type is ``a_perm = [a3, a1, a2]`` while ``b`` is left + in its original order. + + ``permutation_type='pairings'`` supports ``data`` containing any number + of samples, each of which must contain the same number of observations. + All samples provided in ``data`` are permuted *independently*. Therefore, + if ``m`` is the number of samples and ``n`` is the number of observations + within each sample, then the number of permutations in an exact test is:: + + factorial(n)**m + + Note that if a two-sample statistic, for example, does not inherently + depend on the order in which observations are provided - only on the + *pairings* of observations - then only one of the two samples should be + provided in ``data``. This dramatically reduces computational cost without + affecting the shape of the null distribution (because the frequency/count + of each value is affected by the same factor). + + **Paired statistics, permute samples** (``permutation_type='samples'``): + + The null hypothesis associated with this permutation type is that + observations within each pair are drawn from the same underlying + distribution and that the sample to which they are assigned is random. + + Suppose ``data`` contains two samples; e.g. ``a, b = data``. + Let ``n`` be the number of observations in ``a``, which must also equal + the number of observations in ``b``. + + When ``1 < n_resamples < 2**n``, the elements of ``a`` are ``b`` are + randomly swapped between samples (maintaining their pairings) and the + statistic is calculated. This process is performed repeatedly, + `permutation` times, generating a distribution of the statistic under the + null hypothesis. The statistic of the original data is compared to this + distribution to determine the p-value. + + When ``n_resamples >= 2**n``, an exact test is performed: the observations + are assigned to the two samples in each distinct way (while maintaining + pairings) exactly once. + + For ``a = [a1, a2, a3]`` and ``b = [b1, b2, b3]``, an example of this + permutation type is ``x = [b1, a2, b3]`` and ``y = [a1, b2, a3]``. + + ``permutation_type='samples'`` supports ``data`` containing any number + of samples, each of which must contain the same number of observations. + If ``data`` contains more than one sample, paired observations within + ``data`` are exchanged between samples *independently*. Therefore, if ``m`` + is the number of samples and ``n`` is the number of observations within + each sample, then the number of permutations in an exact test is:: + + factorial(m)**n + + Several paired-sample statistical tests, such as the Wilcoxon signed rank + test and paired-sample t-test, can be performed considering only the + *difference* between two paired elements. Accordingly, if ``data`` contains + only one sample, then the null distribution is formed by independently + changing the *sign* of each observation. + + .. warning:: + The p-value is calculated by counting the elements of the null + distribution that are as extreme or more extreme than the observed + value of the statistic. Due to the use of finite precision arithmetic, + some statistic functions return numerically distinct values when the + theoretical values would be exactly equal. In some cases, this could + lead to a large error in the calculated p-value. `permutation_test` + guards against this by considering elements in the null distribution + that are "close" (within a relative tolerance of 100 times the + floating point epsilon of inexact dtypes) to the observed + value of the test statistic as equal to the observed value of the + test statistic. However, the user is advised to inspect the null + distribution to assess whether this method of comparison is + appropriate, and if not, calculate the p-value manually. See example + below. + + References + ---------- + + .. [1] R. A. Fisher. The Design of Experiments, 6th Ed (1951). + .. [2] B. Phipson and G. K. Smyth. "Permutation P-values Should Never Be + Zero: Calculating Exact P-values When Permutations Are Randomly Drawn." + Statistical Applications in Genetics and Molecular Biology 9.1 (2010). + .. [3] M. D. Ernst. "Permutation Methods: A Basis for Exact Inference". + Statistical Science (2004). + .. [4] B. Efron and R. J. Tibshirani. An Introduction to the Bootstrap + (1993). + + Examples + -------- + + Suppose we wish to test whether two samples are drawn from the same + distribution. Assume that the underlying distributions are unknown to us, + and that before observing the data, we hypothesized that the mean of the + first sample would be less than that of the second sample. We decide that + we will use the difference between the sample means as a test statistic, + and we will consider a p-value of 0.05 to be statistically significant. + + For efficiency, we write the function defining the test statistic in a + vectorized fashion: the samples ``x`` and ``y`` can be ND arrays, and the + statistic will be calculated for each axis-slice along `axis`. + + >>> import numpy as np + >>> def statistic(x, y, axis): + ... return np.mean(x, axis=axis) - np.mean(y, axis=axis) + + After collecting our data, we calculate the observed value of the test + statistic. + + >>> from scipy.stats import norm + >>> rng = np.random.default_rng() + >>> x = norm.rvs(size=5, random_state=rng) + >>> y = norm.rvs(size=6, loc = 3, random_state=rng) + >>> statistic(x, y, 0) + -3.5411688580987266 + + Indeed, the test statistic is negative, suggesting that the true mean of + the distribution underlying ``x`` is less than that of the distribution + underlying ``y``. To determine the probability of this occurring by chance + if the two samples were drawn from the same distribution, we perform + a permutation test. + + >>> from scipy.stats import permutation_test + >>> # because our statistic is vectorized, we pass `vectorized=True` + >>> # `n_resamples=np.inf` indicates that an exact test is to be performed + >>> res = permutation_test((x, y), statistic, vectorized=True, + ... n_resamples=np.inf, alternative='less') + >>> print(res.statistic) + -3.5411688580987266 + >>> print(res.pvalue) + 0.004329004329004329 + + The probability of obtaining a test statistic less than or equal to the + observed value under the null hypothesis is 0.4329%. This is less than our + chosen threshold of 5%, so we consider this to be significant evidence + against the null hypothesis in favor of the alternative. + + Because the size of the samples above was small, `permutation_test` could + perform an exact test. For larger samples, we resort to a randomized + permutation test. + + >>> x = norm.rvs(size=100, random_state=rng) + >>> y = norm.rvs(size=120, loc=0.2, random_state=rng) + >>> res = permutation_test((x, y), statistic, n_resamples=9999, + ... vectorized=True, alternative='less', + ... rng=rng) + >>> print(res.statistic) + -0.4230459671240913 + >>> print(res.pvalue) + 0.0015 + + The approximate probability of obtaining a test statistic less than or + equal to the observed value under the null hypothesis is 0.0225%. This is + again less than our chosen threshold of 5%, so again we have significant + evidence to reject the null hypothesis in favor of the alternative. + + For large samples and number of permutations, the result is comparable to + that of the corresponding asymptotic test, the independent sample t-test. + + >>> from scipy.stats import ttest_ind + >>> res_asymptotic = ttest_ind(x, y, alternative='less') + >>> print(res_asymptotic.pvalue) + 0.0014669545224902675 + + The permutation distribution of the test statistic is provided for + further investigation. + + >>> import matplotlib.pyplot as plt + >>> plt.hist(res.null_distribution, bins=50) + >>> plt.title("Permutation distribution of test statistic") + >>> plt.xlabel("Value of Statistic") + >>> plt.ylabel("Frequency") + >>> plt.show() + + Inspection of the null distribution is essential if the statistic suffers + from inaccuracy due to limited machine precision. Consider the following + case: + + >>> from scipy.stats import pearsonr + >>> x = [1, 2, 4, 3] + >>> y = [2, 4, 6, 8] + >>> def statistic(x, y, axis=-1): + ... return pearsonr(x, y, axis=axis).statistic + >>> res = permutation_test((x, y), statistic, vectorized=True, + ... permutation_type='pairings', + ... alternative='greater') + >>> r, pvalue, null = res.statistic, res.pvalue, res.null_distribution + + In this case, some elements of the null distribution differ from the + observed value of the correlation coefficient ``r`` due to numerical noise. + We manually inspect the elements of the null distribution that are nearly + the same as the observed value of the test statistic. + + >>> r + 0.7999999999999999 + >>> unique = np.unique(null) + >>> unique + array([-1. , -1. , -0.8, -0.8, -0.8, -0.6, -0.4, -0.4, -0.2, -0.2, -0.2, + 0. , 0.2, 0.2, 0.2, 0.4, 0.4, 0.6, 0.8, 0.8, 0.8, 1. , + 1. ]) # may vary + >>> unique[np.isclose(r, unique)].tolist() + [0.7999999999999998, 0.7999999999999999, 0.8] # may vary + + If `permutation_test` were to perform the comparison naively, the + elements of the null distribution with value ``0.7999999999999998`` would + not be considered as extreme or more extreme as the observed value of the + statistic, so the calculated p-value would be too small. + + >>> incorrect_pvalue = np.count_nonzero(null >= r) / len(null) + >>> incorrect_pvalue + 0.14583333333333334 # may vary + + Instead, `permutation_test` treats elements of the null distribution that + are within ``max(1e-14, abs(r)*1e-14)`` of the observed value of the + statistic ``r`` to be equal to ``r``. + + >>> correct_pvalue = np.count_nonzero(null >= r - 1e-14) / len(null) + >>> correct_pvalue + 0.16666666666666666 + >>> res.pvalue == correct_pvalue + True + + This method of comparison is expected to be accurate in most practical + situations, but the user is advised to assess this by inspecting the + elements of the null distribution that are close to the observed value + of the statistic. Also, consider the use of statistics that can be + calculated using exact arithmetic (e.g. integer statistics). + + """ + args = _permutation_test_iv(data, statistic, permutation_type, vectorized, + n_resamples, batch, alternative, axis, + rng) + (data, statistic, permutation_type, vectorized, n_resamples, batch, + alternative, axis, rng, float_dtype, xp) = args + + observed = statistic(*data, axis=-1) + + null_calculators = {"pairings": _calculate_null_pairings, + "samples": _calculate_null_samples, + "independent": _calculate_null_both} + null_calculator_args = (data, statistic, n_resamples, + batch, rng) + calculate_null = null_calculators[permutation_type] + null_distribution, n_resamples, exact_test = ( + calculate_null(*null_calculator_args, xp=xp)) + + # See References [2] and [3] + adjustment = 0 if exact_test else 1 + + # relative tolerance for detecting numerically distinct but + # theoretically equal values in the null distribution + eps = (0 if not xp.isdtype(observed.dtype, 'real floating') + else xp.finfo(observed.dtype).eps*100) + gamma = xp.abs(eps * observed) + + def less(null_distribution, observed): + cmps = null_distribution <= observed + gamma + count = xp.count_nonzero(cmps, axis=0) + adjustment + pvalues = xp.astype(count, float_dtype) / (n_resamples + adjustment) + return pvalues + + def greater(null_distribution, observed): + cmps = null_distribution >= observed - gamma + count = xp.count_nonzero(cmps, axis=0) + adjustment + pvalues = xp.astype(count, float_dtype) / (n_resamples + adjustment) + return pvalues + + def two_sided(null_distribution, observed): + pvalues_less = less(null_distribution, observed) + pvalues_greater = greater(null_distribution, observed) + pvalues = xp.minimum(pvalues_less, pvalues_greater) * 2 + return pvalues + + compare = {"less": less, + "greater": greater, + "two-sided": two_sided} + + pvalues = compare[alternative](null_distribution, observed) + pvalues = xpx.at(pvalues)[xp.isnan(observed)].set(xp.nan) + pvalues = xp.clip(pvalues, 0., 1.) + + return PermutationTestResult(observed, pvalues, null_distribution) + + +@dataclass +class ResamplingMethod: + """Configuration information for a statistical resampling method. + + Instances of this class can be passed into the `method` parameter of some + hypothesis test functions to perform a resampling or Monte Carlo version + of the hypothesis test. + + Attributes + ---------- + n_resamples : int + The number of resamples to perform or Monte Carlo samples to draw. + batch : int, optional + The number of resamples to process in each vectorized call to + the statistic. Batch sizes >>1 tend to be faster when the statistic + is vectorized, but memory usage scales linearly with the batch size. + Default is ``None``, which processes all resamples in a single batch. + + """ + n_resamples: int = 9999 + batch: int = None # type: ignore[assignment] + + +@dataclass +class MonteCarloMethod(ResamplingMethod): + """Configuration information for a Monte Carlo hypothesis test. + + Instances of this class can be passed into the `method` parameter of some + hypothesis test functions to perform a Monte Carlo version of the + hypothesis tests. + + Parameters + ---------- + n_resamples : int, optional + The number of Monte Carlo samples to draw. Default is 9999. + batch : int, optional + The number of Monte Carlo samples to process in each vectorized call to + the statistic. Batch sizes >>1 tend to be faster when the statistic + is vectorized, but memory usage scales linearly with the batch size. + Default is ``None``, which processes all samples in a single batch. + rvs : callable or tuple of callables, optional + A callable or sequence of callables that generates random variates + under the null hypothesis. Each element of `rvs` must be a callable + that accepts keyword argument ``size`` (e.g. ``rvs(size=(m, n))``) and + returns an N-d array sample of that shape. If `rvs` is a sequence, the + number of callables in `rvs` must match the number of samples passed + to the hypothesis test in which the `MonteCarloMethod` is used. Default + is ``None``, in which case the hypothesis test function chooses values + to match the standard version of the hypothesis test. For example, + the null hypothesis of `scipy.stats.pearsonr` is typically that the + samples are drawn from the standard normal distribution, so + ``rvs = (rng.normal, rng.normal)`` where + ``rng = np.random.default_rng()``. + rng : `numpy.random.Generator`, optional + Pseudorandom number generator state. When `rng` is None, a new + `numpy.random.Generator` is created using entropy from the + operating system. Types other than `numpy.random.Generator` are + passed to `numpy.random.default_rng` to instantiate a ``Generator``. + + """ + rvs: object = None + rng: object = None + + def __init__(self, n_resamples=9999, batch=None, rvs=None, rng=None): + if (rvs is not None) and (rng is not None): + message = 'Use of `rvs` and `rng` are mutually exclusive.' + raise ValueError(message) + + self.n_resamples = n_resamples + self.batch = batch + self.rvs = rvs + self.rng = rng + + def _asdict(self): + # `dataclasses.asdict` deepcopies; we don't want that. + return dict(n_resamples=self.n_resamples, batch=self.batch, + rvs=self.rvs, rng=self.rng) + + +_rs_deprecation = ("Use of attribute `random_state` is deprecated and replaced by " + "`rng`. Support for `random_state` will be removed in SciPy 1.19.0. " + "To silence this warning and ensure consistent behavior in SciPy " + "1.19.0, control the RNG using attribute `rng`. Values set using " + "attribute `rng` will be validated by `np.random.default_rng`, so " + "the behavior corresponding with a given value may change compared " + "to use of `random_state`. For example, 1) `None` will result in " + "unpredictable random numbers, 2) an integer will result in a " + "different stream of random numbers, (with the same distribution), " + "and 3) `np.random` or `RandomState` instances will result in an " + "error. See the documentation of `default_rng` for more " + "information.") + + +@dataclass +class PermutationMethod(ResamplingMethod): + """Configuration information for a permutation hypothesis test. + + Instances of this class can be passed into the `method` parameter of some + hypothesis test functions to perform a permutation version of the + hypothesis tests. + + Parameters + ---------- + n_resamples : int, optional + The number of resamples to perform. Default is 9999. + batch : int, optional + The number of resamples to process in each vectorized call to + the statistic. Batch sizes >>1 tend to be faster when the statistic + is vectorized, but memory usage scales linearly with the batch size. + Default is ``None``, which processes all resamples in a single batch. + rng : `numpy.random.Generator`, optional + Pseudorandom number generator used to perform resampling. + + If `rng` is passed by keyword to the initializer or the `rng` attribute is used + directly, types other than `numpy.random.Generator` are passed to + `numpy.random.default_rng` to instantiate a ``Generator`` before use. + If `rng` is already a ``Generator`` instance, then the provided instance is + used. Specify `rng` for repeatable behavior. + + If this argument is passed by position, if `random_state` is passed by keyword + into the initializer, or if the `random_state` attribute is used directly, + legacy behavior for `random_state` applies: + + - If `random_state` is None (or `numpy.random`), the `numpy.random.RandomState` + singleton is used. + - If `random_state` is an int, a new ``RandomState`` instance is used, + seeded with `random_state`. + - If `random_state` is already a ``Generator`` or ``RandomState`` instance then + that instance is used. + + .. versionchanged:: 1.15.0 + + As part of the `SPEC-007 `_ + transition from use of `numpy.random.RandomState` to + `numpy.random.Generator`, this attribute name was changed from + `random_state` to `rng`. For an interim period, both names will continue to + work, although only one may be specified at a time. After the interim + period, uses of `random_state` will emit warnings. The behavior of both + `random_state` and `rng` are outlined above, but only `rng` should be used + in new code. + + """ + rng: object # type: ignore[misc] + _rng: object = field(init=False, repr=False, default=None) # type: ignore[assignment] + + @property + def random_state(self): + # Uncomment in SciPy 1.17.0 + # warnings.warn(_rs_deprecation, DeprecationWarning, stacklevel=2) + return self._random_state + + @random_state.setter + def random_state(self, val): + # Uncomment in SciPy 1.17.0 + # warnings.warn(_rs_deprecation, DeprecationWarning, stacklevel=2) + self._random_state = val + + @property # type: ignore[no-redef] + def rng(self): # noqa: F811 + return self._rng + + def __init__(self, n_resamples=9999, batch=None, random_state=None, *, rng=None): + # Uncomment in SciPy 1.17.0 + # warnings.warn(_rs_deprecation.replace('attribute', 'argument'), + # DeprecationWarning, stacklevel=2) + self._rng = rng + self._random_state = random_state + super().__init__(n_resamples=n_resamples, batch=batch) + + def _asdict(self): + # `dataclasses.asdict` deepcopies; we don't want that. + d = dict(n_resamples=self.n_resamples, batch=self.batch) + if self.rng is not None: + d['rng'] = self.rng + if self.random_state is not None: + d['random_state'] = self.random_state + return d + + +@dataclass +class BootstrapMethod(ResamplingMethod): + """Configuration information for a bootstrap confidence interval. + + Instances of this class can be passed into the `method` parameter of some + confidence interval methods to generate a bootstrap confidence interval. + + Parameters + ---------- + n_resamples : int, optional + The number of resamples to perform. Default is 9999. + batch : int, optional + The number of resamples to process in each vectorized call to + the statistic. Batch sizes >>1 tend to be faster when the statistic + is vectorized, but memory usage scales linearly with the batch size. + Default is ``None``, which processes all resamples in a single batch. + rng : `numpy.random.Generator`, optional + Pseudorandom number generator used to perform resampling. + + If `rng` is passed by keyword to the initializer or the `rng` attribute is used + directly, types other than `numpy.random.Generator` are passed to + `numpy.random.default_rng` to instantiate a ``Generator`` before use. + If `rng` is already a ``Generator`` instance, then the provided instance is + used. Specify `rng` for repeatable behavior. + + If this argument is passed by position, if `random_state` is passed by keyword + into the initializer, or if the `random_state` attribute is used directly, + legacy behavior for `random_state` applies: + + - If `random_state` is None (or `numpy.random`), the `numpy.random.RandomState` + singleton is used. + - If `random_state` is an int, a new ``RandomState`` instance is used, + seeded with `random_state`. + - If `random_state` is already a ``Generator`` or ``RandomState`` instance then + that instance is used. + + .. versionchanged:: 1.15.0 + + As part of the `SPEC-007 `_ + transition from use of `numpy.random.RandomState` to + `numpy.random.Generator`, this attribute name was changed from + `random_state` to `rng`. For an interim period, both names will continue to + work, although only one may be specified at a time. After the interim + period, uses of `random_state` will emit warnings. The behavior of both + `random_state` and `rng` are outlined above, but only `rng` should be used + in new code. + + method : {'BCa', 'percentile', 'basic'} + Whether to use the 'percentile' bootstrap ('percentile'), the 'basic' + (AKA 'reverse') bootstrap ('basic'), or the bias-corrected and + accelerated bootstrap ('BCa', default). + + """ + rng: object # type: ignore[misc] + _rng: object = field(init=False, repr=False, default=None) # type: ignore[assignment] + method: str = 'BCa' + + @property + def random_state(self): + # Uncomment in SciPy 1.17.0 + # warnings.warn(_rs_deprecation, DeprecationWarning, stacklevel=2) + return self._random_state + + @random_state.setter + def random_state(self, val): + # Uncomment in SciPy 1.17.0 + # warnings.warn(_rs_deprecation, DeprecationWarning, stacklevel=2) + self._random_state = val + + @property # type: ignore[no-redef] + def rng(self): # noqa: F811 + return self._rng + + def __init__(self, n_resamples=9999, batch=None, random_state=None, + method='BCa', *, rng=None): + # Uncomment in SciPy 1.17.0 + # warnings.warn(_rs_deprecation.replace('attribute', 'argument'), + # DeprecationWarning, stacklevel=2) + self._rng = rng # don't validate with `default_rng` + self._random_state = random_state + self.method = method + super().__init__(n_resamples=n_resamples, batch=batch) + + def _asdict(self): + # `dataclasses.asdict` deepcopies; we don't want that. + d = dict(n_resamples=self.n_resamples, batch=self.batch, + method=self.method) + if self.rng is not None: + d['rng'] = self.rng + if self.random_state is not None: + d['random_state'] = self.random_state + return d diff --git a/v12x/pylib/scipy/stats/_result_classes.py b/v12x/pylib/scipy/stats/_result_classes.py new file mode 100644 index 0000000000000000000000000000000000000000..975af9310efb0c9a414439fd8d531fb95c988951 --- /dev/null +++ b/v12x/pylib/scipy/stats/_result_classes.py @@ -0,0 +1,40 @@ +# This module exists only to allow Sphinx to generate docs +# for the result objects returned by some functions in stats +# _without_ adding them to the main stats documentation page. + +""" +Result classes +-------------- + +.. currentmodule:: scipy.stats._result_classes + +.. autosummary:: + :toctree: generated/ + + RelativeRiskResult + BinomTestResult + TukeyHSDResult + DunnettResult + PearsonRResult + FitResult + OddsRatioResult + TtestResult + ECDFResult + EmpiricalDistributionFunction + +""" + +__all__ = ['BinomTestResult', 'RelativeRiskResult', 'TukeyHSDResult', + 'PearsonRResult', 'FitResult', 'OddsRatioResult', + 'TtestResult', 'DunnettResult', 'ECDFResult', + 'EmpiricalDistributionFunction'] + + +from ._binomtest import BinomTestResult +from ._odds_ratio import OddsRatioResult +from ._relative_risk import RelativeRiskResult +from ._hypotests import TukeyHSDResult +from ._multicomp import DunnettResult +from ._stats_py import PearsonRResult, TtestResult +from ._fit import FitResult +from ._survival import ECDFResult, EmpiricalDistributionFunction diff --git a/v12x/pylib/scipy/stats/_sampling.py b/v12x/pylib/scipy/stats/_sampling.py new file mode 100644 index 0000000000000000000000000000000000000000..3c0abd14bb9e7f14bb3ae98cb13bcb2b6d8144c7 --- /dev/null +++ b/v12x/pylib/scipy/stats/_sampling.py @@ -0,0 +1,1319 @@ +import math +import numbers +import numpy as np +from scipy import stats +from scipy import special as sc +from ._qmc import (check_random_state as check_random_state_qmc, + Halton, QMCEngine) +from ._unuran.unuran_wrapper import NumericalInversePolynomial +from scipy._lib._util import check_random_state + + +__all__ = ['FastGeneratorInversion', 'RatioUniforms'] + + +# define pdfs and other helper functions to create the generators + +def argus_pdf(x, chi): + # approach follows Baumgarten/Hoermann: Generating ARGUS random variates + # for chi > 5, use relationship of the ARGUS distribution to Gamma(1.5) + if chi <= 5: + y = 1 - x * x + return x * math.sqrt(y) * math.exp(-0.5 * chi**2 * y) + return math.sqrt(x) * math.exp(-x) + + +def argus_gamma_trf(x, chi): + if chi <= 5: + return x + return np.sqrt(1.0 - 2 * x / chi**2) + + +def argus_gamma_inv_trf(x, chi): + if chi <= 5: + return x + return 0.5 * chi**2 * (1 - x**2) + + +def betaprime_pdf(x, a, b): + if x > 0: + logf = (a - 1) * math.log(x) - (a + b) * math.log1p(x) - sc.betaln(a, b) + return math.exp(logf) + else: + # return pdf at x == 0 separately to avoid runtime warnings + if a > 1: + return 0 + elif a < 1: + return np.inf + else: + return 1 / sc.beta(a, b) + + +def beta_valid_params(a, b): + return (min(a, b) >= 0.1) and (max(a, b) <= 700) + + +def gamma_pdf(x, a): + if x > 0: + return math.exp(-math.lgamma(a) + (a - 1.0) * math.log(x) - x) + else: + return 0 if a >= 1 else np.inf + + +def invgamma_pdf(x, a): + if x > 0: + return math.exp(-(a + 1.0) * math.log(x) - math.lgamma(a) - 1 / x) + else: + return 0 if a >= 1 else np.inf + + +def burr_pdf(x, cc, dd): + # note: we use np.exp instead of math.exp, otherwise an overflow + # error can occur in the setup, e.g., for parameters + # 1.89128135, 0.30195177, see test test_burr_overflow + if x > 0: + lx = math.log(x) + return np.exp(-(cc + 1) * lx - (dd + 1) * math.log1p(np.exp(-cc * lx))) + else: + return 0 + + +def burr12_pdf(x, cc, dd): + if x > 0: + lx = math.log(x) + logterm = math.log1p(math.exp(cc * lx)) + return math.exp((cc - 1) * lx - (dd + 1) * logterm + math.log(cc * dd)) + else: + return 0 + + +def chi_pdf(x, a): + if x > 0: + return math.exp( + (a - 1) * math.log(x) + - 0.5 * (x * x) + - (a / 2 - 1) * math.log(2) + - math.lgamma(0.5 * a) + ) + else: + return 0 if a >= 1 else np.inf + + +def chi2_pdf(x, df): + if x > 0: + return math.exp( + (df / 2 - 1) * math.log(x) + - 0.5 * x + - (df / 2) * math.log(2) + - math.lgamma(0.5 * df) + ) + else: + return 0 if df >= 1 else np.inf + + +def alpha_pdf(x, a): + if x > 0: + return math.exp(-2.0 * math.log(x) - 0.5 * (a - 1.0 / x) ** 2) + return 0.0 + + +def bradford_pdf(x, c): + if 0 <= x <= 1: + return 1.0 / (1.0 + c * x) + return 0.0 + + +def crystalball_pdf(x, b, m): + if x > -b: + return math.exp(-0.5 * x * x) + return math.exp(m * math.log(m / b) - 0.5 * b * b - m * math.log(m / b - b - x)) + + +def weibull_min_pdf(x, c): + if x > 0: + return c * math.exp((c - 1) * math.log(x) - x**c) + return 0.0 + + +def weibull_max_pdf(x, c): + if x < 0: + return c * math.exp((c - 1) * math.log(-x) - ((-x) ** c)) + return 0.0 + + +def invweibull_pdf(x, c): + if x > 0: + return c * math.exp(-(c + 1) * math.log(x) - x ** (-c)) + return 0.0 + + +def wald_pdf(x): + if x > 0: + return math.exp(-((x - 1) ** 2) / (2 * x)) / math.sqrt(x**3) + return 0.0 + + +def geninvgauss_mode(p, b): + if p > 1: # equivalent mode formulas numerical more stable versions + return (math.sqrt((1 - p) ** 2 + b**2) - (1 - p)) / b + return b / (math.sqrt((1 - p) ** 2 + b**2) + (1 - p)) + + +def geninvgauss_pdf(x, p, b): + m = geninvgauss_mode(p, b) + lfm = (p - 1) * math.log(m) - 0.5 * b * (m + 1 / m) + if x > 0: + return math.exp((p - 1) * math.log(x) - 0.5 * b * (x + 1 / x) - lfm) + return 0.0 + + +def invgauss_mode(mu): + return 1.0 / (math.sqrt(1.5 * 1.5 + 1 / (mu * mu)) + 1.5) + + +def invgauss_pdf(x, mu): + m = invgauss_mode(mu) + lfm = -1.5 * math.log(m) - (m - mu) ** 2 / (2 * m * mu**2) + if x > 0: + return math.exp(-1.5 * math.log(x) - (x - mu) ** 2 / (2 * x * mu**2) - lfm) + return 0.0 + + +def powerlaw_pdf(x, a): + if x > 0: + return x ** (a - 1) + return 0.0 + + +# Define a dictionary: for a given distribution (keys), another dictionary +# (values) specifies the parameters for NumericalInversePolynomial (PINV). +# The keys of the latter dictionary are: +# - pdf: the pdf of the distribution (callable). The signature of the pdf +# is float -> float (i.e., the function does not have to be vectorized). +# If possible, functions like log or exp from the module math should be +# preferred over functions from numpy since the PINV setup will be faster +# in that case. +# - check_pinv_params: callable f that returns true if the shape parameters +# (args) are recommended parameters for PINV (i.e., the u-error does +# not exceed the default tolerance) +# - center: scalar if the center does not depend on args, otherwise +# callable that returns the center as a function of the shape parameters +# - rvs_transform: a callable that can be used to transform the rvs that +# are distributed according to the pdf to the target distribution +# (as an example, see the entry for the beta distribution) +# - rvs_transform_inv: the inverse of rvs_transform (it is required +# for the transformed ppf) +# - mirror_uniform: bool or a callable that returns true or false +# depending on the shape parameters. If True, the ppf is applied +# to 1-u instead of u to generate rvs, where u is a uniform rv. +# While both u and 1-u are uniform, it can be required to use 1-u +# to compute the u-error correctly. This is only relevant for the argus +# distribution. +# The only required keys are "pdf" and "check_pinv_params". +# All other keys are optional. + +PINV_CONFIG = { + "alpha": { + "pdf": alpha_pdf, + "check_pinv_params": lambda a: 1.0e-11 <= a < 2.1e5, + "center": lambda a: 0.25 * (math.sqrt(a * a + 8.0) - a), + }, + "anglit": { + "pdf": lambda x: math.cos(2 * x) + 1.0e-13, + # +1.e-13 is necessary, otherwise PINV has strange problems as + # f(upper border) is very close to 0 + "center": 0, + }, + "argus": { + "pdf": argus_pdf, + "center": lambda chi: 0.7 if chi <= 5 else 0.5, + "check_pinv_params": lambda chi: 1e-20 < chi < 901, + "rvs_transform": argus_gamma_trf, + "rvs_transform_inv": argus_gamma_inv_trf, + "mirror_uniform": lambda chi: chi > 5, + }, + "beta": { + "pdf": betaprime_pdf, + "center": lambda a, b: max(0.1, (a - 1) / (b + 1)), + "check_pinv_params": beta_valid_params, + "rvs_transform": lambda x, *args: x / (1 + x), + "rvs_transform_inv": lambda x, *args: x / (1 - x) if x < 1 else np.inf, + }, + "betaprime": { + "pdf": betaprime_pdf, + "center": lambda a, b: max(0.1, (a - 1) / (b + 1)), + "check_pinv_params": beta_valid_params, + }, + "bradford": { + "pdf": bradford_pdf, + "check_pinv_params": lambda a: 1.0e-6 <= a <= 1e9, + "center": 0.5, + }, + "burr": { + "pdf": burr_pdf, + "center": lambda a, b: (2 ** (1 / b) - 1) ** (-1 / a), + "check_pinv_params": lambda a, b: (min(a, b) >= 0.3) and (max(a, b) <= 50), + }, + "burr12": { + "pdf": burr12_pdf, + "center": lambda a, b: (2 ** (1 / b) - 1) ** (1 / a), + "check_pinv_params": lambda a, b: (min(a, b) >= 0.2) and (max(a, b) <= 50), + }, + "cauchy": { + "pdf": lambda x: 1 / (1 + (x * x)), + "center": 0, + }, + "chi": { + "pdf": chi_pdf, + "check_pinv_params": lambda df: 0.05 <= df <= 1.0e6, + "center": lambda a: math.sqrt(a), + }, + "chi2": { + "pdf": chi2_pdf, + "check_pinv_params": lambda df: 0.07 <= df <= 1e6, + "center": lambda a: a, + }, + "cosine": { + "pdf": lambda x: 1 + math.cos(x), + "center": 0, + }, + "crystalball": { + "pdf": crystalball_pdf, + "check_pinv_params": lambda b, m: (0.01 <= b <= 5.5) + and (1.1 <= m <= 75.1), + "center": 0.0, + }, + "expon": { + "pdf": lambda x: math.exp(-x), + "center": 1.0, + }, + "gamma": { + "pdf": gamma_pdf, + "check_pinv_params": lambda a: 0.04 <= a <= 1e6, + "center": lambda a: a, + }, + "gennorm": { + "pdf": lambda x, b: math.exp(-abs(x) ** b), + "check_pinv_params": lambda b: 0.081 <= b <= 45.0, + "center": 0.0, + }, + "geninvgauss": { + "pdf": geninvgauss_pdf, + "check_pinv_params": lambda p, b: (abs(p) <= 1200.0) + and (1.0e-10 <= b <= 1200.0), + "center": geninvgauss_mode, + }, + "gumbel_l": { + "pdf": lambda x: math.exp(x - math.exp(x)), + "center": -0.6, + }, + "gumbel_r": { + "pdf": lambda x: math.exp(-x - math.exp(-x)), + "center": 0.6, + }, + "hypsecant": { + "pdf": lambda x: 1.0 / (math.exp(x) + math.exp(-x)), + "center": 0.0, + }, + "invgamma": { + "pdf": invgamma_pdf, + "check_pinv_params": lambda a: 0.04 <= a <= 1e6, + "center": lambda a: 1 / a, + }, + "invgauss": { + "pdf": invgauss_pdf, + "check_pinv_params": lambda mu: 1.0e-10 <= mu <= 1.0e9, + "center": invgauss_mode, + }, + "invweibull": { + "pdf": invweibull_pdf, + "check_pinv_params": lambda a: 0.12 <= a <= 512, + "center": 1.0, + }, + "laplace": { + "pdf": lambda x: math.exp(-abs(x)), + "center": 0.0, + }, + "logistic": { + "pdf": lambda x: math.exp(-x) / (1 + math.exp(-x)) ** 2, + "center": 0.0, + }, + "maxwell": { + "pdf": lambda x: x * x * math.exp(-0.5 * x * x), + "center": 1.41421, + }, + "moyal": { + "pdf": lambda x: math.exp(-(x + math.exp(-x)) / 2), + "center": 1.2, + }, + "norm": { + "pdf": lambda x: math.exp(-x * x / 2), + "center": 0.0, + }, + "pareto": { + "pdf": lambda x, b: x ** -(b + 1), + "center": lambda b: b / (b - 1) if b > 2 else 1.5, + "check_pinv_params": lambda b: 0.08 <= b <= 400000, + }, + "powerlaw": { + "pdf": powerlaw_pdf, + "center": 1.0, + "check_pinv_params": lambda a: 0.06 <= a <= 1.0e5, + }, + "t": { + "pdf": lambda x, df: (1 + x * x / df) ** (-0.5 * (df + 1)), + "check_pinv_params": lambda a: 0.07 <= a <= 1e6, + "center": 0.0, + }, + "rayleigh": { + "pdf": lambda x: x * math.exp(-0.5 * (x * x)), + "center": 1.0, + }, + "semicircular": { + "pdf": lambda x: math.sqrt(1.0 - (x * x)), + "center": 0, + }, + "wald": { + "pdf": wald_pdf, + "center": 1.0, + }, + "weibull_max": { + "pdf": weibull_max_pdf, + "check_pinv_params": lambda a: 0.25 <= a <= 512, + "center": -1.0, + }, + "weibull_min": { + "pdf": weibull_min_pdf, + "check_pinv_params": lambda a: 0.25 <= a <= 512, + "center": 1.0, + }, +} + + +def _validate_qmc_input(qmc_engine, d, seed): + # Input validation for `qmc_engine` and `d` + # Error messages for invalid `d` are raised by QMCEngine + # we could probably use a stats.qmc.check_qrandom_state + if isinstance(qmc_engine, QMCEngine): + if d is not None and qmc_engine.d != d: + message = "`d` must be consistent with dimension of `qmc_engine`." + raise ValueError(message) + d = qmc_engine.d if d is None else d + elif qmc_engine is None: + d = 1 if d is None else d + qmc_engine = Halton(d, seed=seed) + else: + message = ( + "`qmc_engine` must be an instance of " + "`scipy.stats.qmc.QMCEngine` or `None`." + ) + raise ValueError(message) + + return qmc_engine, d + + +class CustomDistPINV: + def __init__(self, pdf, args): + self._pdf = lambda x: pdf(x, *args) + + def pdf(self, x): + return self._pdf(x) + + +class FastGeneratorInversion: + """ + Fast sampling by numerical inversion of the CDF for a large class of + continuous distributions in `scipy.stats`. + + Parameters + ---------- + dist : rv_frozen object + Frozen distribution object from `scipy.stats`. The list of supported + distributions can be found in the Notes section. The shape parameters, + `loc` and `scale` used to create the distributions must be scalars. + For example, for the Gamma distribution with shape parameter `p`, + `p` has to be a float, and for the beta distribution with shape + parameters (a, b), both a and b have to be floats. + domain : tuple of floats, optional + If one wishes to sample from a truncated/conditional distribution, + the domain has to be specified. + The default is None. In that case, the random variates are not + truncated, and the domain is inferred from the support of the + distribution. + ignore_shape_range : bool, optional + If False, shape parameters that are outside of the valid range + of values to ensure that the numerical accuracy (see Notes) is + high, raise a ValueError. If True, any shape parameters that are valid + for the distribution are accepted. This can be useful for testing. + The default is False. + random_state : {None, int, `numpy.random.Generator`, + `numpy.random.RandomState`}, optional + + A NumPy random number generator or seed for the underlying NumPy + random number generator used to generate the stream of uniform + random numbers. + If `random_state` is None, it uses ``self.random_state``. + If `random_state` is an int, + ``np.random.default_rng(random_state)`` is used. + If `random_state` is already a ``Generator`` or ``RandomState`` + instance then that instance is used. + + Attributes + ---------- + loc : float + The location parameter. + random_state : {`numpy.random.Generator`, `numpy.random.RandomState`} + The random state used in relevant methods like `rvs` (unless + another `random_state` is passed as an argument to these methods). + scale : float + The scale parameter. + + Methods + ------- + cdf + evaluate_error + ppf + qrvs + rvs + support + + Notes + ----- + The class creates an object for continuous distributions specified + by `dist`. The method `rvs` uses a generator from + `scipy.stats.sampling` that is created when the object is instantiated. + In addition, the methods `qrvs` and `ppf` are added. + `qrvs` generate samples based on quasi-random numbers from + `scipy.stats.qmc`. `ppf` is the PPF based on the + numerical inversion method in [1]_ (`NumericalInversePolynomial`) that is + used to generate random variates. + + Supported distributions (`distname`) are: + ``alpha``, ``anglit``, ``argus``, ``beta``, ``betaprime``, ``bradford``, + ``burr``, ``burr12``, ``cauchy``, ``chi``, ``chi2``, ``cosine``, + ``crystalball``, ``expon``, ``gamma``, ``gennorm``, ``geninvgauss``, + ``gumbel_l``, ``gumbel_r``, ``hypsecant``, ``invgamma``, ``invgauss``, + ``invweibull``, ``laplace``, ``logistic``, ``maxwell``, ``moyal``, + ``norm``, ``pareto``, ``powerlaw``, ``t``, ``rayleigh``, ``semicircular``, + ``wald``, ``weibull_max``, ``weibull_min``. + + `rvs` relies on the accuracy of the numerical inversion. If very extreme + shape parameters are used, the numerical inversion might not work. However, + for all implemented distributions, the admissible shape parameters have + been tested, and an error will be raised if the user supplies values + outside of the allowed range. The u-error should not exceed 1e-10 for all + valid parameters. Note that warnings might be raised even if parameters + are within the valid range when the object is instantiated. + To check numerical accuracy, the method `evaluate_error` can be used. + + Note that all implemented distributions are also part of `scipy.stats`, and + the object created by `FastGeneratorInversion` relies on methods like + `ppf`, `cdf` and `pdf` from `rv_frozen`. The main benefit of using this + class can be summarized as follows: Once the generator to sample random + variates is created in the setup step, sampling and evaluation of + the PPF using `ppf` are very fast, + and performance is essentially independent of the distribution. Therefore, + a substantial speed-up can be achieved for many distributions if large + numbers of random variates are required. It is important to know that this + fast sampling is achieved by inversion of the CDF. Thus, one uniform + random variate is transformed into a non-uniform variate, which is an + advantage for several simulation methods, e.g., when + the variance reduction methods of common random variates or + antithetic variates are be used ([2]_). + + In addition, inversion makes it possible to + + - to use a QMC generator from `scipy.stats.qmc` (method `qrvs`), + - to generate random variates truncated to an interval. For example, if + + one aims to sample standard normal random variates from + the interval (2, 4), this can be easily achieved by using the parameter + `domain`. + + The location and scale that are initially defined by `dist` + can be reset without having to rerun the setup + step to create the generator that is used for sampling. The relation + of the distribution `Y` with `loc` and `scale` to the standard + distribution `X` (i.e., ``loc=0`` and ``scale=1``) is given by + ``Y = loc + scale * X``. + + References + ---------- + .. [1] Derflinger, Gerhard, Wolfgang Hörmann, and Josef Leydold. + "Random variate generation by numerical inversion when only the + density is known." ACM Transactions on Modeling and Computer + Simulation (TOMACS) 20.4 (2010): 1-25. + .. [2] Hörmann, Wolfgang, Josef Leydold and Gerhard Derflinger. + "Automatic nonuniform random number generation." + Springer, 2004. + + Examples + -------- + >>> import numpy as np + >>> from scipy import stats + >>> from scipy.stats.sampling import FastGeneratorInversion + + Let's start with a simple example to illustrate the main features: + + >>> gamma_frozen = stats.gamma(1.5) + >>> gamma_dist = FastGeneratorInversion(gamma_frozen) + >>> r = gamma_dist.rvs(size=1000) + + The mean should be approximately equal to the shape parameter 1.5: + + >>> r.mean() + 1.52423591130436 # may vary + + Similarly, we can draw a sample based on quasi-random numbers: + + >>> r = gamma_dist.qrvs(size=1000) + >>> r.mean() + 1.4996639255942914 # may vary + + Compare the PPF against approximation `ppf`. + + >>> q = [0.001, 0.2, 0.5, 0.8, 0.999] + >>> np.max(np.abs(gamma_frozen.ppf(q) - gamma_dist.ppf(q))) + 4.313394796895409e-08 + + To confirm that the numerical inversion is accurate, we evaluate the + approximation error (u-error), which should be below 1e-10 (for more + details, refer to the documentation of `evaluate_error`): + + >>> gamma_dist.evaluate_error() + (7.446320551265581e-11, nan) # may vary + + Note that the location and scale can be changed without instantiating a + new generator: + + >>> gamma_dist.loc = 2 + >>> gamma_dist.scale = 3 + >>> r = gamma_dist.rvs(size=1000) + + The mean should be approximately 2 + 3*1.5 = 6.5. + + >>> r.mean() + 6.399549295242894 # may vary + + Let us also illustrate how truncation can be applied: + + >>> trunc_norm = FastGeneratorInversion(stats.norm(), domain=(3, 4)) + >>> r = trunc_norm.rvs(size=1000) + >>> 3 < r.min() < r.max() < 4 + True + + Check the mean: + + >>> r.mean() + 3.250433367078603 # may vary + + >>> stats.norm.expect(lb=3, ub=4, conditional=True) + 3.260454285589997 + + In this particular, case, `scipy.stats.truncnorm` could also be used to + generate truncated normal random variates. + + """ + + def __init__( + self, + dist, + *, + domain=None, + ignore_shape_range=False, + random_state=None, + ): + + if isinstance(dist, stats.distributions.rv_frozen): + distname = dist.dist.name + if distname not in PINV_CONFIG.keys(): + raise ValueError( + f"Distribution '{distname}' is not supported." + f"It must be one of {list(PINV_CONFIG.keys())}" + ) + else: + raise ValueError("`dist` must be a frozen distribution object") + + loc = dist.kwds.get("loc", 0) + scale = dist.kwds.get("scale", 1) + args = dist.args + if not np.isscalar(loc): + raise ValueError("loc must be scalar.") + if not np.isscalar(scale): + raise ValueError("scale must be scalar.") + + self._frozendist = getattr(stats, distname)( + *args, + loc=loc, + scale=scale, + ) + self._distname = distname + + nargs = np.broadcast_arrays(args)[0].size + nargs_expected = self._frozendist.dist.numargs + if nargs != nargs_expected: + raise ValueError( + f"Each of the {nargs_expected} shape parameters must be a " + f"scalar, but {nargs} values are provided." + ) + + self.random_state = random_state + + if domain is None: + self._domain = self._frozendist.support() + self._p_lower = 0.0 + self._p_domain = 1.0 + else: + self._domain = domain + self._p_lower = self._frozendist.cdf(self._domain[0]) + _p_domain = self._frozendist.cdf(self._domain[1]) - self._p_lower + self._p_domain = _p_domain + self._set_domain_adj() + self._ignore_shape_range = ignore_shape_range + + # the domain to be passed to NumericalInversePolynomial + # define a separate variable since in case of a transformation, + # domain_pinv will not be the same as self._domain + self._domain_pinv = self._domain + + # get information about the distribution from the config to set up + # the generator + dist = self._process_config(distname, args) + + if self._rvs_transform_inv is not None: + d0 = self._rvs_transform_inv(self._domain[0], *args) + d1 = self._rvs_transform_inv(self._domain[1], *args) + if d0 > d1: + # swap values if transformation if decreasing + d0, d1 = d1, d0 + # only update _domain_pinv and not _domain + # _domain refers to the original distribution, _domain_pinv + # to the transformed distribution + self._domain_pinv = d0, d1 + + # self._center has been set by the call self._process_config + # check if self._center is inside the transformed domain + # _domain_pinv, otherwise move it to the endpoint that is closer + if self._center is not None: + if self._center < self._domain_pinv[0]: + self._center = self._domain_pinv[0] + elif self._center > self._domain_pinv[1]: + self._center = self._domain_pinv[1] + + self._rng = NumericalInversePolynomial( + dist, + random_state=self.random_state, + domain=self._domain_pinv, + center=self._center, + ) + + @property + def random_state(self): + return self._random_state + + @random_state.setter + def random_state(self, random_state): + self._random_state = check_random_state_qmc(random_state) + + @property + def loc(self): + return self._frozendist.kwds.get("loc", 0) + + @loc.setter + def loc(self, loc): + if not np.isscalar(loc): + raise ValueError("loc must be scalar.") + self._frozendist.kwds["loc"] = loc + # update the adjusted domain that depends on loc and scale + self._set_domain_adj() + + @property + def scale(self): + return self._frozendist.kwds.get("scale", 0) + + @scale.setter + def scale(self, scale): + if not np.isscalar(scale): + raise ValueError("scale must be scalar.") + self._frozendist.kwds["scale"] = scale + # update the adjusted domain that depends on loc and scale + self._set_domain_adj() + + def _set_domain_adj(self): + """ Adjust the domain based on loc and scale. """ + loc = self.loc + scale = self.scale + lb = self._domain[0] * scale + loc + ub = self._domain[1] * scale + loc + self._domain_adj = (lb, ub) + + def _process_config(self, distname, args): + cfg = PINV_CONFIG[distname] + if "check_pinv_params" in cfg: + if not self._ignore_shape_range: + if not cfg["check_pinv_params"](*args): + msg = ("No generator is defined for the shape parameters " + f"{args}. Use ignore_shape_range to proceed " + "with the selected values.") + raise ValueError(msg) + + if "center" in cfg.keys(): + if not np.isscalar(cfg["center"]): + self._center = cfg["center"](*args) + else: + self._center = cfg["center"] + else: + self._center = None + self._rvs_transform = cfg.get("rvs_transform", None) + self._rvs_transform_inv = cfg.get("rvs_transform_inv", None) + _mirror_uniform = cfg.get("mirror_uniform", None) + if _mirror_uniform is None: + self._mirror_uniform = False + else: + self._mirror_uniform = _mirror_uniform(*args) + + return CustomDistPINV(cfg["pdf"], args) + + def rvs(self, size=None): + """ + Sample from the distribution by inversion. + + Parameters + ---------- + size : int or tuple, optional + The shape of samples. Default is ``None`` in which case a scalar + sample is returned. + + Returns + ------- + rvs : array_like + A NumPy array of random variates. + + Notes + ----- + Random variates are generated by numerical inversion of the CDF, i.e., + `ppf` computed by `NumericalInversePolynomial` when the class + is instantiated. Note that the + default ``rvs`` method of the rv_continuous class is + overwritten. Hence, a different stream of random numbers is generated + even if the same seed is used. + """ + # note: we cannot use self._rng.rvs directly in case + # self._mirror_uniform is true + u = self.random_state.uniform(size=size) + if self._mirror_uniform: + u = 1 - u + r = self._rng.ppf(u) + if self._rvs_transform is not None: + r = self._rvs_transform(r, *self._frozendist.args) + return self.loc + self.scale * r + + def ppf(self, q): + """ + Very fast PPF (inverse CDF) of the distribution which + is a very close approximation of the exact PPF values. + + Parameters + ---------- + q : array_like + Array with probabilities. + + Returns + ------- + ppf : array_like + Quantiles corresponding to the values in `q`. + + Notes + ----- + The evaluation of the PPF is very fast but it may have a large + relative error in the far tails. The numerical precision of the PPF + is controlled by the q-error, that is, + ``max |q - CDF(PPF(q))|`` where the max is taken over points in + the interval [0,1], see `evaluate_error`. + + Note that this PPF is designed to generate random samples. + """ + q = np.asarray(q) + if self._mirror_uniform: + x = self._rng.ppf(1 - q) + else: + x = self._rng.ppf(q) + if self._rvs_transform is not None: + x = self._rvs_transform(x, *self._frozendist.args) + return self.scale * x + self.loc + + def qrvs(self, size=None, d=None, qmc_engine=None): + """ + Quasi-random variates of the given distribution. + + The `qmc_engine` is used to draw uniform quasi-random variates, and + these are converted to quasi-random variates of the given distribution + using inverse transform sampling. + + Parameters + ---------- + size : int, tuple of ints, or None; optional + Defines shape of random variates array. Default is ``None``. + d : int or None, optional + Defines dimension of uniform quasi-random variates to be + transformed. Default is ``None``. + qmc_engine : scipy.stats.qmc.QMCEngine(d=1), optional + Defines the object to use for drawing + quasi-random variates. Default is ``None``, which uses + `scipy.stats.qmc.Halton(1)`. + + Returns + ------- + rvs : ndarray or scalar + Quasi-random variates. See Notes for shape information. + + Notes + ----- + The shape of the output array depends on `size`, `d`, and `qmc_engine`. + The intent is for the interface to be natural, but the detailed rules + to achieve this are complicated. + + - If `qmc_engine` is ``None``, a `scipy.stats.qmc.Halton` instance is + created with dimension `d`. If `d` is not provided, ``d=1``. + - If `qmc_engine` is not ``None`` and `d` is ``None``, `d` is + determined from the dimension of the `qmc_engine`. + - If `qmc_engine` is not ``None`` and `d` is not ``None`` but the + dimensions are inconsistent, a ``ValueError`` is raised. + - After `d` is determined according to the rules above, the output + shape is ``tuple_shape + d_shape``, where: + + - ``tuple_shape = tuple()`` if `size` is ``None``, + - ``tuple_shape = (size,)`` if `size` is an ``int``, + - ``tuple_shape = size`` if `size` is a sequence, + - ``d_shape = tuple()`` if `d` is ``None`` or `d` is 1, and + - ``d_shape = (d,)`` if `d` is greater than 1. + + The elements of the returned array are part of a low-discrepancy + sequence. If `d` is 1, this means that none of the samples are truly + independent. If `d` > 1, each slice ``rvs[..., i]`` will be of a + quasi-independent sequence; see `scipy.stats.qmc.QMCEngine` for + details. Note that when `d` > 1, the samples returned are still those + of the provided univariate distribution, not a multivariate + generalization of that distribution. + + """ + qmc_engine, d = _validate_qmc_input(qmc_engine, d, self.random_state) + # mainly copied from unuran_wrapper.pyx.templ + # `rvs` is flexible about whether `size` is an int or tuple, so this + # should be, too. + try: + if size is None: + tuple_size = (1,) + else: + tuple_size = tuple(size) + except TypeError: + tuple_size = (size,) + # we do not use rng.qrvs directly since we need to be + # able to apply the ppf to 1 - u + N = 1 if size is None else np.prod(size) + u = qmc_engine.random(N) + if self._mirror_uniform: + u = 1 - u + qrvs = self._ppf(u) + if self._rvs_transform is not None: + qrvs = self._rvs_transform(qrvs, *self._frozendist.args) + if size is None: + qrvs = qrvs.squeeze()[()] + else: + if d == 1: + qrvs = qrvs.reshape(tuple_size) + else: + qrvs = qrvs.reshape(tuple_size + (d,)) + return self.loc + self.scale * qrvs + + def evaluate_error(self, size=100000, random_state=None, x_error=False): + """ + Evaluate the numerical accuracy of the inversion (u- and x-error). + + Parameters + ---------- + size : int, optional + The number of random points over which the error is estimated. + Default is ``100000``. + random_state : {None, int, `numpy.random.Generator`, + `numpy.random.RandomState`}, optional + + A NumPy random number generator or seed for the underlying NumPy + random number generator used to generate the stream of uniform + random numbers. + If `random_state` is None, use ``self.random_state``. + If `random_state` is an int, + ``np.random.default_rng(random_state)`` is used. + If `random_state` is already a ``Generator`` or ``RandomState`` + instance then that instance is used. + x_error : bool + If ``True``, return the x-error, as per the notes section. If ``False``, + returns ``NaN``. + + Returns + ------- + u_error, x_error : tuple of floats + A NumPy array of random variates. + + Notes + ----- + The numerical precision of the inverse CDF `ppf` is controlled by + the u-error. It is computed as follows: + ``max |u - CDF(PPF(u))|`` where the max is taken `size` random + points in the interval [0,1]. `random_state` determines the random + sample. Note that if `ppf` was exact, the u-error would be zero. + + The x-error measures the direct distance between the exact PPF + and `ppf`. If ``x_error`` is set to ``True`, it is + computed as the maximum of the minimum of the relative and absolute + x-error: + ``max(min(x_error_abs[i], x_error_rel[i]))`` where + ``x_error_abs[i] = |PPF(u[i]) - PPF_fast(u[i])|``, + ``x_error_rel[i] = max |(PPF(u[i]) - PPF_fast(u[i])) / PPF(u[i])|``. + Note that it is important to consider the relative x-error in the case + that ``PPF(u)`` is close to zero or very large. + + By default, only the u-error is evaluated and the x-error is set to + ``np.nan``. Note that the evaluation of the x-error will be very slow + if the implementation of the PPF is slow. + + Further information about these error measures can be found in [1]_. + + References + ---------- + .. [1] Derflinger, Gerhard, Wolfgang Hörmann, and Josef Leydold. + "Random variate generation by numerical inversion when only the + density is known." ACM Transactions on Modeling and Computer + Simulation (TOMACS) 20.4 (2010): 1-25. + + Examples + -------- + + >>> import numpy as np + >>> from scipy import stats + >>> from scipy.stats.sampling import FastGeneratorInversion + + Create an object for the normal distribution: + + >>> d_norm_frozen = stats.norm() + >>> d_norm = FastGeneratorInversion(d_norm_frozen) + + To confirm that the numerical inversion is accurate, we evaluate the + approximation error (u-error and x-error). + + >>> u_error, x_error = d_norm.evaluate_error(x_error=True) + + The u-error should be below 1e-10: + + >>> u_error + 8.785783212061915e-11 # may vary + + Compare the PPF against approximation `ppf`: + + >>> q = [0.001, 0.2, 0.4, 0.6, 0.8, 0.999] + >>> diff = np.abs(d_norm_frozen.ppf(q) - d_norm.ppf(q)) + >>> x_error_abs = np.max(diff) + >>> x_error_abs + 1.2937954707581412e-08 + + This is the absolute x-error evaluated at the points q. The relative + error is given by + + >>> x_error_rel = np.max(diff / np.abs(d_norm_frozen.ppf(q))) + >>> x_error_rel + 4.186725600453555e-09 + + The x_error computed above is derived in a very similar way over a + much larger set of random values q. At each value q[i], the minimum + of the relative and absolute error is taken. The final value is then + derived as the maximum of these values. In our example, we get the + following value: + + >>> x_error + 4.507068014335139e-07 # may vary + + """ + if not isinstance(size, numbers.Integral | np.integer): + raise ValueError("size must be an integer.") + # urng will be used to draw the samples for testing the error + # it must not interfere with self.random_state. therefore, do not + # call self.rvs, but draw uniform random numbers and apply + # self.ppf (note: like in rvs, consider self._mirror_uniform) + urng = check_random_state_qmc(random_state) + u = urng.uniform(size=size) + if self._mirror_uniform: + u = 1 - u + x = self.ppf(u) + uerr = np.max(np.abs(self._cdf(x) - u)) + if not x_error: + return uerr, np.nan + ppf_u = self._ppf(u) + x_error_abs = np.abs(self.ppf(u)-ppf_u) + x_error_rel = x_error_abs / np.abs(ppf_u) + x_error_combined = np.array([x_error_abs, x_error_rel]).min(axis=0) + return uerr, np.max(x_error_combined) + + def support(self): + """Support of the distribution. + + Returns + ------- + a, b : float + end-points of the distribution's support. + + Notes + ----- + + Note that the support of the distribution depends on `loc`, + `scale` and `domain`. + + Examples + -------- + + >>> from scipy import stats + >>> from scipy.stats.sampling import FastGeneratorInversion + + Define a truncated normal distribution: + + >>> d_norm = FastGeneratorInversion(stats.norm(), domain=(0, 1)) + >>> d_norm.support() + (0, 1) + + Shift the distribution: + + >>> d_norm.loc = 2.5 + >>> d_norm.support() + (2.5, 3.5) + + """ + return self._domain_adj + + def _cdf(self, x): + """Cumulative distribution function (CDF) + + Parameters + ---------- + x : array_like + The values where the CDF is evaluated + + Returns + ------- + y : ndarray + CDF evaluated at x + + """ + y = self._frozendist.cdf(x) + if self._p_domain == 1.0: + return y + return np.clip((y - self._p_lower) / self._p_domain, 0, 1) + + def _ppf(self, q): + """Percent point function (inverse of `cdf`) + + Parameters + ---------- + q : array_like + lower tail probability + + Returns + ------- + x : array_like + quantile corresponding to the lower tail probability q. + + """ + if self._p_domain == 1.0: + return self._frozendist.ppf(q) + x = self._frozendist.ppf(self._p_domain * np.array(q) + self._p_lower) + return np.clip(x, self._domain_adj[0], self._domain_adj[1]) + + +class RatioUniforms: + """ + Generate random samples from a probability density function using the + ratio-of-uniforms method. + + Parameters + ---------- + pdf : callable + A function with signature `pdf(x)` that is proportional to the + probability density function of the distribution. + umax : float + The upper bound of the bounding rectangle in the u-direction. + vmin : float + The lower bound of the bounding rectangle in the v-direction. + vmax : float + The upper bound of the bounding rectangle in the v-direction. + c : float, optional + Shift parameter of ratio-of-uniforms method, see Notes. Default is 0. + random_state : {None, int, `numpy.random.Generator`, + `numpy.random.RandomState`}, optional + + If `seed` is None (or `np.random`), the `numpy.random.RandomState` + singleton is used. + If `seed` is an int, a new ``RandomState`` instance is used, + seeded with `seed`. + If `seed` is already a ``Generator`` or ``RandomState`` instance then + that instance is used. + + Methods + ------- + rvs + + Notes + ----- + Given a univariate probability density function `pdf` and a constant `c`, + define the set ``A = {(u, v) : 0 < u <= sqrt(pdf(v/u + c))}``. + If ``(U, V)`` is a random vector uniformly distributed over ``A``, + then ``V/U + c`` follows a distribution according to `pdf`. + + The above result (see [1]_, [2]_) can be used to sample random variables + using only the PDF, i.e. no inversion of the CDF is required. Typical + choices of `c` are zero or the mode of `pdf`. The set ``A`` is a subset of + the rectangle ``R = [0, umax] x [vmin, vmax]`` where + + - ``umax = sup sqrt(pdf(x))`` + - ``vmin = inf (x - c) sqrt(pdf(x))`` + - ``vmax = sup (x - c) sqrt(pdf(x))`` + + In particular, these values are finite if `pdf` is bounded and + ``x**2 * pdf(x)`` is bounded (i.e. subquadratic tails). + One can generate ``(U, V)`` uniformly on ``R`` and return + ``V/U + c`` if ``(U, V)`` are also in ``A`` which can be directly + verified. + + The algorithm is not changed if one replaces `pdf` by k * `pdf` for any + constant k > 0. Thus, it is often convenient to work with a function + that is proportional to the probability density function by dropping + unnecessary normalization factors. + + Intuitively, the method works well if ``A`` fills up most of the + enclosing rectangle such that the probability is high that ``(U, V)`` + lies in ``A`` whenever it lies in ``R`` as the number of required + iterations becomes too large otherwise. To be more precise, note that + the expected number of iterations to draw ``(U, V)`` uniformly + distributed on ``R`` such that ``(U, V)`` is also in ``A`` is given by + the ratio ``area(R) / area(A) = 2 * umax * (vmax - vmin) / area(pdf)``, + where `area(pdf)` is the integral of `pdf` (which is equal to one if the + probability density function is used but can take on other values if a + function proportional to the density is used). The equality holds since + the area of ``A`` is equal to ``0.5 * area(pdf)`` (Theorem 7.1 in [1]_). + If the sampling fails to generate a single random variate after 50000 + iterations (i.e. not a single draw is in ``A``), an exception is raised. + + If the bounding rectangle is not correctly specified (i.e. if it does not + contain ``A``), the algorithm samples from a distribution different from + the one given by `pdf`. It is therefore recommended to perform a + test such as `~scipy.stats.kstest` as a check. + + References + ---------- + .. [1] L. Devroye, "Non-Uniform Random Variate Generation", + Springer-Verlag, 1986. + + .. [2] W. Hoermann and J. Leydold, "Generating generalized inverse Gaussian + random variates", Statistics and Computing, 24(4), p. 547--557, 2014. + + .. [3] A.J. Kinderman and J.F. Monahan, "Computer Generation of Random + Variables Using the Ratio of Uniform Deviates", + ACM Transactions on Mathematical Software, 3(3), p. 257--260, 1977. + + Examples + -------- + >>> import numpy as np + >>> from scipy import stats + + >>> from scipy.stats.sampling import RatioUniforms + >>> rng = np.random.default_rng() + + Simulate normally distributed random variables. It is easy to compute the + bounding rectangle explicitly in that case. For simplicity, we drop the + normalization factor of the density. + + >>> f = lambda x: np.exp(-x**2 / 2) + >>> v = np.sqrt(f(np.sqrt(2))) * np.sqrt(2) + >>> umax = np.sqrt(f(0)) + >>> gen = RatioUniforms(f, umax=umax, vmin=-v, vmax=v, random_state=rng) + >>> r = gen.rvs(size=2500) + + The K-S test confirms that the random variates are indeed normally + distributed (normality is not rejected at 5% significance level): + + >>> stats.kstest(r, 'norm')[1] + 0.250634764150542 + + The exponential distribution provides another example where the bounding + rectangle can be determined explicitly. + + >>> gen = RatioUniforms(lambda x: np.exp(-x), umax=1, vmin=0, + ... vmax=2*np.exp(-1), random_state=rng) + >>> r = gen.rvs(1000) + >>> stats.kstest(r, 'expon')[1] + 0.21121052054580314 + + """ + + def __init__(self, pdf, *, umax, vmin, vmax, c=0, random_state=None): + if vmin >= vmax: + raise ValueError("vmin must be smaller than vmax.") + + if umax <= 0: + raise ValueError("umax must be positive.") + + self._pdf = pdf + self._umax = umax + self._vmin = vmin + self._vmax = vmax + self._c = c + self._rng = check_random_state(random_state) + + def rvs(self, size=1): + """Sampling of random variates. + + Parameters + ---------- + size : int or tuple of ints, optional + Number of random variates to be generated (default is 1). + + Returns + ------- + rvs : ndarray + The random variates distributed according to the probability + distribution defined by the pdf. + + """ + size1d = tuple(np.atleast_1d(size)) + N = np.prod(size1d) # number of rvs needed, reshape upon return + + # start sampling using ratio of uniforms method + x = np.zeros(N) + simulated, i = 0, 1 + + # loop until N rvs have been generated: expected runtime is finite. + # to avoid infinite loop, raise exception if not a single rv has been + # generated after 50000 tries. even if the expected number of iterations + # is 1000, the probability of this event is (1-1/1000)**50000 + # which is of order 10e-22 + while simulated < N: + k = N - simulated + # simulate uniform rvs on [0, umax] and [vmin, vmax] + u1 = self._umax * self._rng.uniform(size=k) + v1 = self._rng.uniform(self._vmin, self._vmax, size=k) + # apply rejection method + rvs = v1 / u1 + self._c + accept = (u1**2 <= self._pdf(rvs)) + num_accept = np.sum(accept) + if num_accept > 0: + x[simulated:(simulated + num_accept)] = rvs[accept] + simulated += num_accept + + if (simulated == 0) and (i*N >= 50000): + msg = ( + f"Not a single random variate could be generated in {i*N} " + "attempts. The ratio of uniforms method does not appear " + "to work for the provided parameters. Please check the " + "pdf and the bounds." + ) + raise RuntimeError(msg) + i += 1 + + return np.reshape(x, size1d) diff --git a/v12x/pylib/scipy/stats/_sensitivity_analysis.py b/v12x/pylib/scipy/stats/_sensitivity_analysis.py new file mode 100644 index 0000000000000000000000000000000000000000..529c1e6120a63ff65313928ee87b4ee1e881a73c --- /dev/null +++ b/v12x/pylib/scipy/stats/_sensitivity_analysis.py @@ -0,0 +1,718 @@ +import inspect +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any +from collections.abc import Callable + +import numpy as np + +from scipy.stats._common import ConfidenceInterval +from scipy.stats._qmc import check_random_state +from scipy.stats._resampling import BootstrapResult +from scipy.stats import qmc, bootstrap +from scipy._lib._array_api import xp_capabilities +from scipy._lib._util import _transition_to_rng + + +if TYPE_CHECKING: + import numpy.typing as npt + from scipy._lib._util import DecimalNumber, IntNumber + + +__all__ = [ + 'sobol_indices' +] + + +def f_ishigami(x: "npt.ArrayLike") -> "npt.NDArray[np.inexact[Any]]": + r"""Ishigami function. + + .. math:: + + Y(\mathbf{x}) = \sin x_1 + 7 \sin^2 x_2 + 0.1 x_3^4 \sin x_1 + + with :math:`\mathbf{x} \in [-\pi, \pi]^3`. + + Parameters + ---------- + x : array_like ([x1, x2, x3], n) + + Returns + ------- + f : array_like (n,) + Function evaluation. + + References + ---------- + .. [1] Ishigami, T. and T. Homma. "An importance quantification technique + in uncertainty analysis for computer models." IEEE, + :doi:`10.1109/ISUMA.1990.151285`, 1990. + """ + x = np.atleast_2d(x) + f_eval = ( + np.sin(x[0]) + + 7 * np.sin(x[1])**2 + + 0.1 * (x[2]**4) * np.sin(x[0]) + ) + return f_eval + + +def sample_A_B( + n, + dists, + rng=None +): + """Sample two matrices A and B. + + Uses a Sobol' sequence with 2`d` columns to have 2 uncorrelated matrices. + This is more efficient than using 2 random draw of Sobol'. + See sec. 5 from [1]_. + + Output shape is (d, n). + + References + ---------- + .. [1] Saltelli, A., P. Annoni, I. Azzini, F. Campolongo, M. Ratto, and + S. Tarantola. "Variance based sensitivity analysis of model + output. Design and estimator for the total sensitivity index." + Computer Physics Communications, 181(2):259-270, + :doi:`10.1016/j.cpc.2009.09.018`, 2010. + """ + d = len(dists) + A_B = qmc.Sobol(d=2*d, seed=rng, bits=64).random(n).T + A_B = A_B.reshape(2, d, -1) + try: + for d_, dist in enumerate(dists): + A_B[:, d_] = dist.ppf(A_B[:, d_]) + except AttributeError as exc: + message = "Each distribution in `dists` must have method `ppf`." + raise ValueError(message) from exc + return A_B + + +def sample_AB(A: np.ndarray, B: np.ndarray) -> np.ndarray: + """AB matrix. + + AB: rows of B into A. Shape (d, d, n). + + - Copy A into d "pages" + - In the first page, replace 1st rows of A with 1st row of B. + ... + - In the dth page, replace dth row of A with dth row of B. + - return the stack of pages + + """ + d, n = A.shape + AB = np.tile(A, (d, 1, 1)) + i = np.arange(d) + AB[i, i] = B[i] + return AB + + +def saltelli_2010( + f_A: np.ndarray, f_B: np.ndarray, f_AB: np.ndarray +) -> tuple[np.ndarray, np.ndarray]: + r"""Saltelli2010 formulation. + + .. math:: + + S_i = \frac{1}{N} \sum_{j=1}^N + f(\mathbf{B})_j (f(\mathbf{AB}^{(i)})_j - f(\mathbf{A})_j) + + .. math:: + + S_{T_i} = \frac{1}{N} \sum_{j=1}^N + (f(\mathbf{A})_j - f(\mathbf{AB}^{(i)})_j)^2 + + Parameters + ---------- + f_A, f_B : array_like (s, n) + Function values at A and B, respectively + f_AB : array_like (d, s, n) + Function values at each of the AB pages + + Returns + ------- + s, st : array_like (s, d) + First order and total order Sobol' indices. + + References + ---------- + .. [1] Saltelli, A., P. Annoni, I. Azzini, F. Campolongo, M. Ratto, and + S. Tarantola. "Variance based sensitivity analysis of model + output. Design and estimator for the total sensitivity index." + Computer Physics Communications, 181(2):259-270, + :doi:`10.1016/j.cpc.2009.09.018`, 2010. + """ + # Empirical variance calculated using output from A and B which are + # independent. Output of AB is not independent and cannot be used + var = np.var([f_A, f_B], axis=(0, -1)) + + # We divide by the variance to have a ratio of variance + # this leads to eq. 2 + s = np.mean(f_B * (f_AB - f_A), axis=-1) / var # Table 2 (b) + st = 0.5 * np.mean((f_A - f_AB) ** 2, axis=-1) / var # Table 2 (f) + + return s.T, st.T + + +@dataclass +class BootstrapSobolResult: + first_order: BootstrapResult + total_order: BootstrapResult + + +@dataclass +class SobolResult: + first_order: np.ndarray + total_order: np.ndarray + _indices_method: Callable + _f_A: np.ndarray + _f_B: np.ndarray + _f_AB: np.ndarray + _A: np.ndarray | None = None + _B: np.ndarray | None = None + _AB: np.ndarray | None = None + _bootstrap_result: BootstrapResult | None = None + + def bootstrap( + self, + confidence_level: "DecimalNumber" = 0.95, + n_resamples: "IntNumber" = 999 + ) -> BootstrapSobolResult: + """Bootstrap Sobol' indices to provide confidence intervals. + + Parameters + ---------- + confidence_level : float, default: ``0.95`` + The confidence level of the confidence intervals. + n_resamples : int, default: ``999`` + The number of resamples performed to form the bootstrap + distribution of the indices. + + Returns + ------- + res : BootstrapSobolResult + Bootstrap result containing the confidence intervals and the + bootstrap distribution of the indices. + + An object with attributes: + + first_order : BootstrapResult + Bootstrap result of the first order indices. + total_order : BootstrapResult + Bootstrap result of the total order indices. + See `BootstrapResult` for more details. + + """ + def statistic(idx): + f_A_ = self._f_A[:, idx] + f_B_ = self._f_B[:, idx] + f_AB_ = self._f_AB[..., idx] + return self._indices_method(f_A_, f_B_, f_AB_) + + n = self._f_A.shape[1] + + res = bootstrap( + [np.arange(n)], statistic=statistic, method="BCa", + n_resamples=n_resamples, + confidence_level=confidence_level, + bootstrap_result=self._bootstrap_result + ) + self._bootstrap_result = res + + first_order = BootstrapResult( + confidence_interval=ConfidenceInterval( + res.confidence_interval.low[0], res.confidence_interval.high[0] + ), + bootstrap_distribution=res.bootstrap_distribution[0], + standard_error=res.standard_error[0], + ) + total_order = BootstrapResult( + confidence_interval=ConfidenceInterval( + res.confidence_interval.low[1], res.confidence_interval.high[1] + ), + bootstrap_distribution=res.bootstrap_distribution[1], + standard_error=res.standard_error[1], + ) + + return BootstrapSobolResult( + first_order=first_order, total_order=total_order + ) + + +@xp_capabilities(np_only=True) +@_transition_to_rng('random_state', replace_doc=False) +def sobol_indices( + *, + func, + n, + dists=None, + method='saltelli_2010', + rng=None +): + r"""Global sensitivity indices of Sobol'. + + Parameters + ---------- + func : callable or dict(str, array_like) + If `func` is a callable, function to compute the Sobol' indices from. + Its signature must be:: + + func(x: ArrayLike) -> ArrayLike + + with ``x`` of shape ``(d, n)`` and output of shape ``(s, n)`` where: + + - ``d`` is the input dimensionality of `func` + (number of input variables), + - ``s`` is the output dimensionality of `func` + (number of output variables), and + - ``n`` is the number of samples (see `n` below). + + Function evaluation values must be finite. + + If `func` is a dictionary, contains the function evaluations from three + different arrays. Keys must be: ``f_A``, ``f_B`` and ``f_AB``. + ``f_A`` and ``f_B`` should have a shape ``(s, n)`` and ``f_AB`` + should have a shape ``(d, s, n)``. + This is an advanced feature and misuse can lead to wrong analysis. + n : int + Number of samples used to generate the matrices ``A`` and ``B``. + Must be a power of 2. The total number of points at which `func` is + evaluated will be ``n*(d+2)``. + dists : list(distributions), optional + List of each parameter's distribution. The distribution of parameters + depends on the application and should be carefully chosen. + Parameters are assumed to be independently distributed, meaning there + is no constraint nor relationship between their values. + + Distributions must be an instance of a class with a ``ppf`` + method. + + Must be specified if `func` is a callable, and ignored otherwise. + method : Callable or str, default: 'saltelli_2010' + Method used to compute the first and total Sobol' indices. + + If a callable, its signature must be:: + + func(f_A: np.ndarray, f_B: np.ndarray, f_AB: np.ndarray) + -> tuple[np.ndarray, np.ndarray] + + with ``f_A, f_B`` of shape ``(s, n)`` and ``f_AB`` of shape + ``(d, s, n)``. + These arrays contain the function evaluations from three different sets + of samples. + The output is a tuple of the first and total indices with + shape ``(s, d)``. + This is an advanced feature and misuse can lead to wrong analysis. + rng : `numpy.random.Generator`, optional + Pseudorandom number generator state. When `rng` is None, a new + `numpy.random.Generator` is created using entropy from the + operating system. Types other than `numpy.random.Generator` are + passed to `numpy.random.default_rng` to instantiate a ``Generator``. + + .. versionchanged:: 1.15.0 + + As part of the `SPEC-007 `_ + transition from use of `numpy.random.RandomState` to + `numpy.random.Generator`, this keyword was changed from `random_state` to + `rng`. For an interim period, both keywords will continue to work, although + only one may be specified at a time. After the interim period, function + calls using the `random_state` keyword will emit warnings. Following a + deprecation period, the `random_state` keyword will be removed. + + Returns + ------- + res : SobolResult + An object with attributes: + + first_order : ndarray of shape (s, d) + First order Sobol' indices. + total_order : ndarray of shape (s, d) + Total order Sobol' indices. + + And method: + + bootstrap(confidence_level: float, n_resamples: int) + -> BootstrapSobolResult + + A method providing confidence intervals on the indices. + See `scipy.stats.bootstrap` for more details. + + The bootstrapping is done on both first and total order indices, + and they are available in `BootstrapSobolResult` as attributes + ``first_order`` and ``total_order``. + + Notes + ----- + The Sobol' method [1]_, [2]_ is a variance-based Sensitivity Analysis which + obtains the contribution of each parameter to the variance of the + quantities of interest (QoIs; i.e., the outputs of `func`). + Respective contributions can be used to rank the parameters and + also gauge the complexity of the model by computing the + model's effective (or mean) dimension. + + .. note:: + + Parameters are assumed to be independently distributed. Each + parameter can still follow any distribution. In fact, the distribution + is very important and should match the real distribution of the + parameters. + + It uses a functional decomposition of the variance of the function to + explore + + .. math:: + + \mathbb{V}(Y) = \sum_{i}^{d} \mathbb{V}_i (Y) + \sum_{i= 2**12``. The more complex the model is, + the more samples will be needed. + + Even for a purely additive model, the indices may not sum to 1 due + to numerical noise. + + References + ---------- + .. [1] Sobol, I. M.. "Sensitivity analysis for nonlinear mathematical + models." Mathematical Modeling and Computational Experiment, 1:407-414, + 1993. + .. [2] Sobol, I. M. (2001). "Global sensitivity indices for nonlinear + mathematical models and their Monte Carlo estimates." Mathematics + and Computers in Simulation, 55(1-3):271-280, + :doi:`10.1016/S0378-4754(00)00270-6`, 2001. + .. [3] Saltelli, A. "Making best use of model evaluations to + compute sensitivity indices." Computer Physics Communications, + 145(2):280-297, :doi:`10.1016/S0010-4655(02)00280-1`, 2002. + .. [4] Saltelli, A., M. Ratto, T. Andres, F. Campolongo, J. Cariboni, + D. Gatelli, M. Saisana, and S. Tarantola. "Global Sensitivity Analysis. + The Primer." 2007. + .. [5] Saltelli, A., P. Annoni, I. Azzini, F. Campolongo, M. Ratto, and + S. Tarantola. "Variance based sensitivity analysis of model + output. Design and estimator for the total sensitivity index." + Computer Physics Communications, 181(2):259-270, + :doi:`10.1016/j.cpc.2009.09.018`, 2010. + .. [6] Ishigami, T. and T. Homma. "An importance quantification technique + in uncertainty analysis for computer models." IEEE, + :doi:`10.1109/ISUMA.1990.151285`, 1990. + + Examples + -------- + The following is an example with the Ishigami function [6]_ + + .. math:: + + Y(\mathbf{x}) = \sin x_1 + 7 \sin^2 x_2 + 0.1 x_3^4 \sin x_1, + + with :math:`\mathbf{x} \in [-\pi, \pi]^3`. This function exhibits strong + non-linearity and non-monotonicity. + + Remember, Sobol' indices assumes that samples are independently + distributed. In this case we use a uniform distribution on each marginals. + + >>> import numpy as np + >>> from scipy.stats import sobol_indices, uniform + >>> rng = np.random.default_rng() + >>> def f_ishigami(x): + ... f_eval = ( + ... np.sin(x[0]) + ... + 7 * np.sin(x[1])**2 + ... + 0.1 * (x[2]**4) * np.sin(x[0]) + ... ) + ... return f_eval + >>> indices = sobol_indices( + ... func=f_ishigami, n=1024, + ... dists=[ + ... uniform(loc=-np.pi, scale=2*np.pi), + ... uniform(loc=-np.pi, scale=2*np.pi), + ... uniform(loc=-np.pi, scale=2*np.pi) + ... ], + ... rng=rng + ... ) + >>> indices.first_order + array([0.31637954, 0.43781162, 0.00318825]) + >>> indices.total_order + array([0.56122127, 0.44287857, 0.24229595]) + + Confidence interval can be obtained using bootstrapping. + + >>> boot = indices.bootstrap() + + Then, this information can be easily visualized. + + >>> import matplotlib.pyplot as plt + >>> fig, axs = plt.subplots(1, 2, figsize=(9, 4)) + >>> _ = axs[0].errorbar( + ... [1, 2, 3], indices.first_order, fmt='o', + ... yerr=[ + ... indices.first_order - boot.first_order.confidence_interval.low, + ... boot.first_order.confidence_interval.high - indices.first_order + ... ], + ... ) + >>> axs[0].set_ylabel("First order Sobol' indices") + >>> axs[0].set_xlabel('Input parameters') + >>> axs[0].set_xticks([1, 2, 3]) + >>> _ = axs[1].errorbar( + ... [1, 2, 3], indices.total_order, fmt='o', + ... yerr=[ + ... indices.total_order - boot.total_order.confidence_interval.low, + ... boot.total_order.confidence_interval.high - indices.total_order + ... ], + ... ) + >>> axs[1].set_ylabel("Total order Sobol' indices") + >>> axs[1].set_xlabel('Input parameters') + >>> axs[1].set_xticks([1, 2, 3]) + >>> plt.tight_layout() + >>> plt.show() + + .. note:: + + By default, `scipy.stats.uniform` has support ``[0, 1]``. + Using the parameters ``loc`` and ``scale``, one obtains the uniform + distribution on ``[loc, loc + scale]``. + + This result is particularly interesting because the first order index + :math:`S_{x_3} = 0` whereas its total order is :math:`S_{T_{x_3}} = 0.244`. + This means that higher order interactions with :math:`x_3` are responsible + for the difference. Almost 25% of the observed variance + on the QoI is due to the correlations between :math:`x_3` and :math:`x_1`, + although :math:`x_3` by itself has no impact on the QoI. + + The following gives a visual explanation of Sobol' indices on this + function. Let's generate 1024 samples in :math:`[-\pi, \pi]^3` and + calculate the value of the output. + + >>> from scipy.stats import qmc + >>> n_dim = 3 + >>> p_labels = ['$x_1$', '$x_2$', '$x_3$'] + >>> sample = qmc.Sobol(d=n_dim, seed=rng).random(1024) + >>> sample = qmc.scale( + ... sample=sample, + ... l_bounds=[-np.pi, -np.pi, -np.pi], + ... u_bounds=[np.pi, np.pi, np.pi] + ... ) + >>> output = f_ishigami(sample.T) + + Now we can do scatter plots of the output with respect to each parameter. + This gives a visual way to understand how each parameter impacts the + output of the function. + + >>> fig, ax = plt.subplots(1, n_dim, figsize=(12, 4)) + >>> for i in range(n_dim): + ... xi = sample[:, i] + ... ax[i].scatter(xi, output, marker='+') + ... ax[i].set_xlabel(p_labels[i]) + >>> ax[0].set_ylabel('Y') + >>> plt.tight_layout() + >>> plt.show() + + Now Sobol' goes a step further: + by conditioning the output value by given values of the parameter + (black lines), the conditional output mean is computed. It corresponds to + the term :math:`\mathbb{E}(Y|x_i)`. Taking the variance of this term gives + the numerator of the Sobol' indices. + + >>> mini = np.min(output) + >>> maxi = np.max(output) + >>> n_bins = 10 + >>> bins = np.linspace(-np.pi, np.pi, num=n_bins, endpoint=False) + >>> dx = bins[1] - bins[0] + >>> fig, ax = plt.subplots(1, n_dim, figsize=(12, 4)) + >>> for i in range(n_dim): + ... xi = sample[:, i] + ... ax[i].scatter(xi, output, marker='+') + ... ax[i].set_xlabel(p_labels[i]) + ... for bin_ in bins: + ... idx = np.where((bin_ <= xi) & (xi <= bin_ + dx)) + ... xi_ = xi[idx] + ... y_ = output[idx] + ... ave_y_ = np.mean(y_) + ... ax[i].plot([bin_ + dx/2] * 2, [mini, maxi], c='k') + ... ax[i].scatter(bin_ + dx/2, ave_y_, c='r') + >>> ax[0].set_ylabel('Y') + >>> plt.tight_layout() + >>> plt.show() + + Looking at :math:`x_3`, the variance + of the mean is zero leading to :math:`S_{x_3} = 0`. But we can further + observe that the variance of the output is not constant along the parameter + values of :math:`x_3`. This heteroscedasticity is explained by higher order + interactions. Moreover, a heteroscedasticity is also noticeable on + :math:`x_1` leading to an interaction between :math:`x_3` and :math:`x_1`. + On :math:`x_2`, the variance seems to be constant and thus null interaction + with this parameter can be supposed. + + This case is fairly simple to analyse visually---although it is only a + qualitative analysis. Nevertheless, when the number of input parameters + increases such analysis becomes unrealistic as it would be difficult to + conclude on high-order terms. Hence the benefit of using Sobol' indices. + + """ + rng = check_random_state(rng) + + n_ = int(n) + if not (n_ & (n_ - 1) == 0) or n != n_: + raise ValueError( + "The balance properties of Sobol' points require 'n' " + "to be a power of 2." + ) + n = n_ + + if not callable(method): + indices_methods = { + "saltelli_2010": saltelli_2010, + } + try: + method = method.lower() # type: ignore[assignment] + indices_method_ = indices_methods[method] + except KeyError as exc: + message = ( + f"{method!r} is not a valid 'method'. It must be one of" + f" {set(indices_methods)!r} or a callable." + ) + raise ValueError(message) from exc + else: + indices_method_ = method + sig = inspect.signature(indices_method_) + + if set(sig.parameters) != {'f_A', 'f_B', 'f_AB'}: + message = ( + "If 'method' is a callable, it must have the following" + f" signature: {inspect.signature(saltelli_2010)}" + ) + raise ValueError(message) + + def indices_method(f_A, f_B, f_AB): + """Wrap indices method to ensure proper output dimension. + + 1D when single output, 2D otherwise. + """ + return np.squeeze(indices_method_(f_A=f_A, f_B=f_B, f_AB=f_AB)) + + if callable(func): + if dists is None: + raise ValueError( + "'dists' must be defined when 'func' is a callable." + ) + + def wrapped_func(x): + return np.atleast_2d(func(x)) + + A, B = sample_A_B(n=n, dists=dists, rng=rng) + AB = sample_AB(A=A, B=B) + + f_A = wrapped_func(A) + + if f_A.shape[1] != n: + raise ValueError( + "'func' output should have a shape ``(s, -1)`` with ``s`` " + "the number of output." + ) + + def funcAB(AB): + d, d, n = AB.shape + AB = np.moveaxis(AB, 0, -1).reshape(d, n*d) + f_AB = wrapped_func(AB) + return np.moveaxis(f_AB.reshape((-1, n, d)), -1, 0) + + f_B = wrapped_func(B) + f_AB = funcAB(AB) + else: + message = ( + "When 'func' is a dictionary, it must contain the following " + "keys: 'f_A', 'f_B' and 'f_AB'." + "'f_A' and 'f_B' should have a shape ``(s, n)`` and 'f_AB' " + "should have a shape ``(d, s, n)``." + ) + try: + f_A, f_B, f_AB = map(lambda arr: arr.copy(), np.atleast_2d( + func['f_A'], func['f_B'], func['f_AB'] + )) + except KeyError as exc: + raise ValueError(message) from exc + + if f_A.shape[1] != n or f_A.shape != f_B.shape or \ + f_AB.shape == f_A.shape or f_AB.shape[-1] % n != 0: + raise ValueError(message) + + # Normalization by mean + # Sobol', I. and Levitan, Y. L. (1999). On the use of variance reducing + # multipliers in monte carlo computations of a global sensitivity index. + # Computer Physics Communications, 117(1) :52-61. + mean = np.mean([f_A, f_B], axis=(0, -1)).reshape(-1, 1) + f_A -= mean + f_B -= mean + f_AB -= mean + + # Compute indices + # Filter warnings for constant output as var = 0 + with np.errstate(divide='ignore', invalid='ignore'): + first_order, total_order = indices_method(f_A=f_A, f_B=f_B, f_AB=f_AB) + + # null variance means null indices + first_order[~np.isfinite(first_order)] = 0 + total_order[~np.isfinite(total_order)] = 0 + + res = dict( + first_order=first_order, + total_order=total_order, + _indices_method=indices_method, + _f_A=f_A, + _f_B=f_B, + _f_AB=f_AB + ) + + if callable(func): + res.update( + dict( + _A=A, + _B=B, + _AB=AB, + ) + ) + + return SobolResult(**res) diff --git a/v12x/pylib/scipy/stats/_stats_mstats_common.py b/v12x/pylib/scipy/stats/_stats_mstats_common.py new file mode 100644 index 0000000000000000000000000000000000000000..c272fcc1a31a6f296cb54c28859ea51b3ca03362 --- /dev/null +++ b/v12x/pylib/scipy/stats/_stats_mstats_common.py @@ -0,0 +1,34 @@ +import numpy as np +from .._lib._bunch import _make_tuple_bunch + +__all__ = ['_find_repeats'] + +# This is not a namedtuple for backwards compatibility. See PR #12983 +TheilslopesResult = _make_tuple_bunch('TheilslopesResult', + ['slope', 'intercept', + 'low_slope', 'high_slope']) +SiegelslopesResult = _make_tuple_bunch('SiegelslopesResult', + ['slope', 'intercept']) + + +def _n_samples_optional_x(kwargs): + return 2 if kwargs.get('x', None) is not None else 1 + + +def _find_repeats(arr): + # This function assumes it may clobber its input. + if len(arr) == 0: + return np.array(0, np.float64), np.array(0, np.intp) + + # XXX This cast was previously needed for the Fortran implementation, + # should we ditch it? + arr = np.asarray(arr, np.float64).ravel() + arr.sort() + + # Taken from NumPy 1.9's np.unique. + change = np.concatenate(([True], arr[1:] != arr[:-1])) + unique = arr[change] + change_idx = np.concatenate(np.nonzero(change) + ([arr.size],)) + freq = np.diff(change_idx) + atleast2 = freq > 1 + return unique[atleast2], freq[atleast2] diff --git a/v12x/pylib/scipy/stats/_stats_py.py b/v12x/pylib/scipy/stats/_stats_py.py new file mode 100644 index 0000000000000000000000000000000000000000..bf10af5738d8c2fa4ec7ea33cf0deca9e9bc2a19 --- /dev/null +++ b/v12x/pylib/scipy/stats/_stats_py.py @@ -0,0 +1,11250 @@ +# Copyright 2002 Gary Strangman. All rights reserved +# Copyright 2002-2016 The SciPy Developers +# +# The original code from Gary Strangman was heavily adapted for +# use in SciPy by Travis Oliphant. The original code came with the +# following disclaimer: +# +# This software is provided "as-is". There are no expressed or implied +# warranties of any kind, including, but not limited to, the warranties +# of merchantability and fitness for a given application. In no event +# shall Gary Strangman be liable for any direct, indirect, incidental, +# special, exemplary or consequential damages (including, but not limited +# to, loss of use, data or profits, or business interruption) however +# caused and on any theory of liability, whether in contract, strict +# liability or tort (including negligence or otherwise) arising in any way +# out of the use of this software, even if advised of the possibility of +# such damage. + +""" +A collection of basic statistical functions for Python. + +References +---------- +.. [CRCProbStat2000] Zwillinger, D. and Kokoska, S. (2000). CRC Standard + Probability and Statistics Tables and Formulae. Chapman & Hall: New + York. 2000. + +""" +import math +import itertools +import operator +import warnings +from collections import namedtuple +from collections.abc import Sequence + +import numpy as np +from numpy import array, asarray, ma + +from scipy import sparse +from scipy.spatial.distance import cdist + +from scipy.optimize import milp, LinearConstraint +from scipy.optimize.elementwise import find_root +from scipy._lib._util import _get_nan, _rename_parameter, _contains_nan + +import scipy.special as special +# Import unused here but needs to stay until end of deprecation periode +# See https://github.com/scipy/scipy/issues/15765#issuecomment-1875564522 +from scipy import linalg # noqa: F401 +from . import distributions +from . import _mstats_basic as mstats_basic + +from ._stats import _kendall_dis, _toint64, _weightedrankedtau + +from dataclasses import dataclass, field +from ._stats_pythran import _compute_outer_prob_inside_method +from ._resampling import (MonteCarloMethod, PermutationMethod, BootstrapMethod, + monte_carlo_test, permutation_test, bootstrap,) +from ._axis_nan_policy import (_axis_nan_policy_factory, _broadcast_shapes, + _broadcast_arrays, + _broadcast_array_shapes_remove_axis, SmallSampleWarning, + too_small_1d_not_omit, too_small_1d_omit, + too_small_nd_not_omit, too_small_nd_omit) +from ._binomtest import _binary_search_for_binom_tst as _binary_search, _SimpleBinomial +from scipy._lib._bunch import _make_tuple_bunch +from scipy import stats +from scipy._lib._array_api import ( + _asarray, + array_namespace, + is_lazy_array, + is_dask, + is_numpy, + is_cupy, + is_jax, + is_marray, + xp_size, + xp_vector_norm, + xp_promote, + xp_result_type, + xp_capabilities, + xp_ravel, + _count_nonmasked, + _share_masks, + _masked_apply, + xp_swapaxes, + xp_device, +) +import scipy._external.array_api_extra as xpx +from scipy.stats._quantile import _xp_searchsorted + + +# Functions/classes in other files should be added in `__init__.py`, not here +__all__ = ['gmean', 'hmean', 'pmean', 'mode', 'tmean', 'tvar', + 'tmin', 'tmax', 'tstd', 'tsem', 'moment', + 'skew', 'kurtosis', 'describe', 'skewtest', 'kurtosistest', + 'normaltest', 'jarque_bera', + 'scoreatpercentile', 'percentileofscore', + 'cumfreq', 'relfreq', 'obrientransform', + 'sem', 'zmap', 'zscore', 'gzscore', 'iqr', 'gstd', + 'median_abs_deviation', + 'sigmaclip', 'trimboth', 'trim1', 'trim_mean', + 'f_oneway', 'pearsonr', 'fisher_exact', + 'spearmanr', 'pointbiserialr', + 'kendalltau', 'weightedtau', + 'linregress', 'ttest_1samp', + 'ttest_ind', 'ttest_ind_from_stats', 'ttest_rel', + 'kstest', 'ks_1samp', 'ks_2samp', + 'chisquare', 'power_divergence', + 'tiecorrect', 'ranksums', 'kruskal', 'friedmanchisquare', + 'rankdata', 'combine_pvalues', 'quantile_test', + 'wasserstein_distance', 'wasserstein_distance_nd', 'energy_distance', + 'brunnermunzel', 'alexandergovern', + 'expectile', 'lmoment'] + + +def _chk_asarray(a, axis, *, xp=None): + if xp is None: + xp = array_namespace(a) + + if axis is None: + a = xp.reshape(a, (-1,)) + outaxis = 0 + else: + a = xp.asarray(a) + outaxis = axis + + if a.ndim == 0: + a = xp.reshape(a, (-1,)) + + return a, outaxis + + +SignificanceResult = _make_tuple_bunch('SignificanceResult', + ['statistic', 'pvalue'], []) +# Let's call a SignificanceResult with legacy :correlation" attribute a +# "CorrelationResult". Don't add to `extra_field_names`- shouldn't be in repr. + + +def _pack_CorrelationResult(statistic, pvalue, correlation): + res = SignificanceResult(statistic, pvalue) + res.correlation = correlation + return res + + +def _unpack_CorrelationResult(res, _): + return res.statistic, res.pvalue, res.correlation + + +# note that `weights` are paired with `x` +@xp_capabilities(marray=True) +@_axis_nan_policy_factory( + lambda x: x, n_samples=1, n_outputs=1, too_small=0, paired=True, + result_to_tuple=lambda x, _: (x,), kwd_samples=['weights']) +def gmean(a, axis=0, dtype=None, weights=None): + r"""Compute the weighted geometric mean along the specified axis. + + The weighted geometric mean of the array :math:`a_i` associated to weights + :math:`w_i` is: + + .. math:: + + \exp \left( \frac{ \sum_{i=1}^n w_i \ln a_i }{ \sum_{i=1}^n w_i } + \right) \, , + + and, with equal weights, it gives: + + .. math:: + + \sqrt[n]{ \prod_{i=1}^n a_i } \, . + + Parameters + ---------- + a : array_like + Input array or object that can be converted to an array. + axis : int or None, optional + Axis along which the geometric mean is computed. Default is 0. + If None, compute over the whole array `a`. + dtype : dtype, optional + (Floating point) dtype to which numerical arguments are cast before performing + the calculation. If `dtype` is not specified, it defaults to the floating point + result dtype of the numerical arguments. + weights : array_like, optional + The `weights` array must be broadcastable to the same shape as `a`. + Default is None, which gives each value a weight of 1.0. + + Returns + ------- + gmean : ndarray + See `dtype` parameter above. + + See Also + -------- + numpy.mean : Arithmetic average + numpy.average : Weighted average + hmean : Harmonic mean + + Notes + ----- + The sample geometric mean is the exponential of the mean of the natural + logarithms of the observations. + Negative observations will produce NaNs in the output because the *natural* + logarithm (as opposed to the *complex* logarithm) is defined only for + non-negative reals. + + References + ---------- + .. [1] "Weighted Geometric Mean", *Wikipedia*, + https://en.wikipedia.org/wiki/Weighted_geometric_mean. + .. [2] Grossman, J., Grossman, M., Katz, R., "Averages: A New Approach", + Archimedes Foundation, 1983 + + Examples + -------- + >>> from scipy.stats import gmean + >>> gmean([1, 4]) + 2.0 + >>> gmean([1, 2, 3, 4, 5, 6, 7]) + 3.3800151591412964 + >>> gmean([1, 4, 7], weights=[3, 1, 3]) + 2.80668351922014 + + """ + xp = array_namespace(a, weights) + dtype = (xp_result_type(a, weights, force_floating=True, xp=xp) + if dtype is None else dtype) + a = xp.asarray(a, dtype=dtype) + + if weights is not None: + weights = xp.asarray(weights, dtype=dtype) + + with np.errstate(divide='ignore'): + log_a = xp.log(a) + + return xp.exp(_xp_mean(log_a, axis=axis, weights=weights)) + + +@xp_capabilities(marray=True) +@_axis_nan_policy_factory( + lambda x: x, n_samples=1, n_outputs=1, too_small=0, paired=True, + result_to_tuple=lambda x, _: (x,), kwd_samples=['weights']) +def hmean(a, axis=0, dtype=None, *, weights=None): + r"""Calculate the weighted harmonic mean along the specified axis. + + The weighted harmonic mean of the array :math:`a_i` associated to weights + :math:`w_i` is: + + .. math:: + + \frac{ \sum_{i=1}^n w_i }{ \sum_{i=1}^n \frac{w_i}{a_i} } \, , + + and, with equal weights, it gives: + + .. math:: + + \frac{ n }{ \sum_{i=1}^n \frac{1}{a_i} } \, . + + Parameters + ---------- + a : array_like + Input array, masked array or object that can be converted to an array. + axis : int or None, optional + Axis along which the harmonic mean is computed. Default is 0. + If None, compute over the whole array `a`. + dtype : dtype, optional + (Floating point) dtype to which numerical arguments are cast before performing + the calculation. If `dtype` is not specified, it defaults to the floating point + result dtype of the numerical arguments. + weights : array_like, optional + The weights array can either be 1-D (in which case its length must be + the size of `a` along the given `axis`) or of the same shape as `a`. + Default is None, which gives each value a weight of 1.0. + + .. versionadded:: 1.9 + + Returns + ------- + hmean : ndarray + See `dtype` parameter above. + + See Also + -------- + numpy.mean : Arithmetic average + numpy.average : Weighted average + gmean : Geometric mean + + Notes + ----- + The sample harmonic mean is the reciprocal of the mean of the reciprocals + of the observations. + + The harmonic mean is computed over a single dimension of the input + array, axis=0 by default, or all values in the array if axis=None. + float64 intermediate and return values are used for integer inputs. + + The harmonic mean is only defined if all observations are non-negative; + otherwise, the result is NaN. + + References + ---------- + .. [1] "Weighted Harmonic Mean", *Wikipedia*, + https://en.wikipedia.org/wiki/Harmonic_mean#Weighted_harmonic_mean + .. [2] Ferger, F., "The nature and use of the harmonic mean", Journal of + the American Statistical Association, vol. 26, pp. 36-40, 1931 + + Examples + -------- + >>> from scipy.stats import hmean + >>> hmean([1, 4]) + 1.6000000000000001 + >>> hmean([1, 2, 3, 4, 5, 6, 7]) + 2.6997245179063363 + >>> hmean([1, 4, 7], weights=[3, 1, 3]) + 1.9029126213592233 + + """ + xp = array_namespace(a, weights) + dtype = (xp_result_type(a, weights, force_floating=True, xp=xp) + if dtype is None else dtype) + a = xp.asarray(a, dtype=dtype) + + if weights is not None: + weights = xp.asarray(weights, dtype=dtype) + + negative_mask = a < 0 + a = xp.where(negative_mask, xp.nan, a) + + if not is_lazy_array(negative_mask) and xp.any(negative_mask): + message = ("The harmonic mean is only defined if all elements are " + "non-negative; otherwise, the result is NaN.") + warnings.warn(message, RuntimeWarning, stacklevel=2) + + with np.errstate(divide='ignore'): + return 1.0 / _xp_mean(1.0 / a, axis=axis, weights=weights) + + +@xp_capabilities(marray=True) +@_axis_nan_policy_factory( + lambda x: x, n_samples=1, n_outputs=1, too_small=0, paired=True, + result_to_tuple=lambda x, _: (x,), kwd_samples=['weights']) +def pmean(a, p, *, axis=0, dtype=None, weights=None): + r"""Calculate the weighted power mean along the specified axis. + + The weighted power mean of the array :math:`a_i` associated to weights + :math:`w_i` is: + + .. math:: + + \left( \frac{ \sum_{i=1}^n w_i a_i^p }{ \sum_{i=1}^n w_i } + \right)^{ 1 / p } \, , + + and, with equal weights, it gives: + + .. math:: + + \left( \frac{ 1 }{ n } \sum_{i=1}^n a_i^p \right)^{ 1 / p } \, . + + When ``p=0``, it returns the geometric mean. + + This mean is also called generalized mean or Hölder mean, and must not be + confused with the Kolmogorov generalized mean, also called + quasi-arithmetic mean or generalized f-mean [3]_. + + Parameters + ---------- + a : array_like + Input array, masked array or object that can be converted to an array. + p : int or float + Exponent. Must be finite. + axis : int or None, optional + Axis along which the power mean is computed. Default is 0. + If None, compute over the whole array `a`. + dtype : dtype, optional + (Floating point) dtype to which numerical arguments are cast before performing + the calculation. If `dtype` is not specified, it defaults to the floating point + result dtype of the numerical arguments. + weights : array_like, optional + The weights array can either be 1-D (in which case its length must be + the size of `a` along the given `axis`) or of the same shape as `a`. + Default is None, which gives each value a weight of 1.0. + + Returns + ------- + pmean : ndarray, see `dtype` parameter above. + Output array containing the power mean values. + + See Also + -------- + numpy.average : Weighted average + gmean : Geometric mean + hmean : Harmonic mean + + Notes + ----- + The power mean is computed over a single dimension of the input + array, ``axis=0`` by default, or all values in the array if ``axis=None``. + float64 intermediate and return values are used for integer inputs. + + The power mean is only defined if all observations are non-negative; + otherwise, the result is NaN. + + .. versionadded:: 1.9 + + References + ---------- + .. [1] "Generalized Mean", *Wikipedia*, + https://en.wikipedia.org/wiki/Generalized_mean + .. [2] Norris, N., "Convexity properties of generalized mean value + functions", The Annals of Mathematical Statistics, vol. 8, + pp. 118-120, 1937 + .. [3] Bullen, P.S., Handbook of Means and Their Inequalities, 2003 + + Examples + -------- + >>> from scipy.stats import pmean, hmean, gmean + >>> pmean([1, 4], 1.3) + 2.639372938300652 + >>> pmean([1, 2, 3, 4, 5, 6, 7], 1.3) + 4.157111214492084 + >>> pmean([1, 4, 7], -2, weights=[3, 1, 3]) + 1.4969684896631954 + + For p=-1, power mean is equal to harmonic mean: + + >>> pmean([1, 4, 7], -1, weights=[3, 1, 3]) + 1.9029126213592233 + >>> hmean([1, 4, 7], weights=[3, 1, 3]) + 1.9029126213592233 + + For p=0, power mean is defined as the geometric mean: + + >>> pmean([1, 4, 7], 0, weights=[3, 1, 3]) + 2.80668351922014 + >>> gmean([1, 4, 7], weights=[3, 1, 3]) + 2.80668351922014 + + """ + if not isinstance(p, int | float): + raise ValueError("Power mean only defined for exponent of type int or " + "float.") + if p == 0: + return gmean(a, axis=axis, dtype=dtype, weights=weights) + elif math.isinf(p): + message = "Power mean only implemented for finite `p`" + raise NotImplementedError(message) + + xp = array_namespace(a, weights) + dtype = (xp_result_type(a, p, weights, force_floating=True, xp=xp) + if dtype is None else dtype) + a = xp.asarray(a, dtype=dtype) + + if weights is not None: + weights = xp.asarray(weights, dtype=dtype) + + negative_mask = a < 0 + a = xp.where(negative_mask, np.nan, a) + + if not is_lazy_array(negative_mask) and xp.any(negative_mask): + message = ("The power mean is only defined if all elements are " + "non-negative; otherwise, the result is NaN.") + warnings.warn(message, RuntimeWarning, stacklevel=2) + + # Linearized approximation for small p to avoid numerical issues; see gh-23407 + p_threshold = 2e-6 + if abs(p) <= p_threshold: + return _linearized_pmean(a, p, axis=axis, weights=weights, xp=xp) + + with np.errstate(divide='ignore', invalid='ignore'): + return _xp_mean(a**float(p), axis=axis, weights=weights)**(1/p) + + +ModeResult = namedtuple('ModeResult', ('mode', 'count')) + + +def _mode_result(mode, count): + # When a slice is empty, `_axis_nan_policy` automatically produces + # NaN for `mode` and `count`. This is a reasonable convention for `mode`, + # but `count` should not be NaN; it should be zero. + xp = array_namespace(mode, count) + i = xp.isnan(count) + if i.shape == (): + count = xp.asarray(0, dtype=count.dtype)[()] if i else count + else: + count = xpx.at(count)[i].set(0) + if is_numpy(xp) and mode.ndim > 0: + # When a) xp is NumPy, b) array is multidimensional, c) NaNs are present, + # and d) nan_policy='omit', `_axis_nan_policy` decorator uses + # `np.apply_along_axis` to compute mode/count along each axis-slice separately. + # The result is a single array with the result dtype of `mode` and `count`. + # We need to change `count` back to an integer. + count = xp.astype(count, xp.asarray(0).dtype) + return ModeResult(mode, count) + + +@xp_capabilities(skip_backends=[('dask.array', "can't compute chunk size"), + ('cupy', "data-apis/array-api-compat#312")], + jax_jit=False, # would delegate, but jax-ml/jax#34486 + marray=True, + extra_note=("Values in MArray input must be real numbers " + "that cast safely to floating point.")) +@_axis_nan_policy_factory(_mode_result, override={'nan_propagation': False}) +def mode(a, axis=0, nan_policy='propagate', keepdims=False): + r"""Return an array of the modal (most common) value in the passed array. + + If there is more than one such value, only one is returned. + The bin-count for the modal bins is also returned. + + Parameters + ---------- + a : array_like + Numeric, n-dimensional array of which to find mode(s). + axis : int or None, optional + Axis along which to operate. Default is 0. If None, compute over + the whole array `a`. + nan_policy : {'propagate', 'raise', 'omit'}, optional + Defines how to handle when input contains nan. + The following options are available (default is 'propagate'): + + * 'propagate': treats nan as it would treat any other value + * 'raise': throws an error + * 'omit': performs the calculations ignoring nan values + + keepdims : bool, optional + If set to ``False``, the `axis` over which the statistic is taken + is consumed (eliminated from the output array). If set to ``True``, + the `axis` is retained with size one, and the result will broadcast + correctly against the input array. + + Returns + ------- + mode : ndarray + Array of modal values. + count : ndarray + Array of counts for each mode. + + Notes + ----- + The mode is calculated using `numpy.unique`. + In NumPy versions 1.21 and after, all NaNs - even those with different + binary representations - are treated as equivalent and counted as separate + instances of the same value. + + By convention, the mode of an empty array is NaN, and the associated count + is zero. + + Examples + -------- + >>> import numpy as np + >>> a = np.array([[3, 0, 3, 7], + ... [3, 2, 6, 2], + ... [1, 7, 2, 8], + ... [3, 0, 6, 1], + ... [3, 2, 5, 5]]) + >>> from scipy import stats + >>> stats.mode(a, keepdims=True) + ModeResult(mode=array([[3, 0, 6, 1]]), count=array([[4, 2, 2, 1]])) + + To get mode of whole array, specify ``axis=None``: + + >>> stats.mode(a, axis=None, keepdims=True) + ModeResult(mode=[[3]], count=[[5]]) + >>> stats.mode(a, axis=None, keepdims=False) + ModeResult(mode=3, count=5) + + """ + xp = array_namespace(a) + + # `axis`, `nan_policy`, and `keepdims` are handled by `_axis_nan_policy` + if not xp.isdtype(a.dtype, 'numeric'): + message = ("Argument `a` is not recognized as numeric. " + "Support for input that cannot be coerced to a numeric " + "array was deprecated in SciPy 1.9.0 and removed in SciPy " + "1.11.0. Please consider `np.unique`.") + raise TypeError(message) + + if xp_size(a) == 0: + NaN = _get_nan(a, xp=xp) + return ModeResult(*xp.asarray([NaN, 0], dtype=NaN.dtype)) + + if a.ndim == 1 and not is_marray(xp): + vals, cnts = xp.unique_counts(a) + # in contrast with np.unique, `unique_counts` treats all NaNs as distinct, + # but we have always grouped them. Replace `cnts` corresponding with NaNs + # with the number of NaNs. + mask = xp.isnan(vals) + cnts = xpx.at(cnts)[mask].set(xp.count_nonzero(mask)) + modes, counts = vals[xp.argmax(cnts)], xp.max(cnts) + default_int = xp.asarray(1).dtype # fail slow CI job failed - incorrect dtype + counts = xp.astype(counts, default_int, copy=False) + modes = modes[()] if modes.ndim == 0 else modes + counts = counts[()] if counts.ndim == 0 else counts + return ModeResult(modes, counts) + + # `axis` is always -1 after the `_axis_nan_policy` decorator + y = xp.sort(a, axis=-1) + + # Get boolean array of elements that are different from the previous element + if is_marray(xp): + # Treat masked elements as NaNs and *don't* count NaNs as equal. Two downsides: + # - integer arrays are converted to floating point, so some very large distinct + # integers will no longer be distinct. + # - NaNs in the original data *will* be treated as distinct. + y_data = xp_promote(y.data, force_floating=True, xp=y._xp) + y_data = y._xp.where(y.mask, y._xp.nan, y_data) + i = xp.concat([xp.ones(y.shape[:-1] + (1,), dtype=xp.bool), + (y_data[..., :-1] != y_data[..., 1:])], axis=-1) + else: + i = xp.concat([xp.ones(y.shape[:-1] + (1,), dtype=xp.bool), + (y[..., :-1] != y[..., 1:]) & ~xp.isnan(y[..., :-1])], axis=-1) + # Get linear integer indices of these elements in a raveled array + indices = xp.arange(xp_size(y), device=xp_device(y))[xp_ravel(i)] + # The difference between integer indices is the number of repeats + append = xp.full(indices.shape[:-1] + (1,), xp_size(y), dtype=indices.dtype) + counts = xp.diff(indices, append=append) + # Now we form an array of `counts` corresponding with each element of `y`... + counts = xp.reshape(xp.repeat(counts, counts), y.shape) + # ... so we can get the argmax of *each slice* separately. + k = xp.argmax(counts, axis=-1, keepdims=True) + # Extract the corresponding element/count, and eliminate the reduced dimension + modes = xp.take_along_axis(y, k, axis=-1)[..., 0] + counts = xp.take_along_axis(counts, k, axis=-1)[..., 0] + modes = modes[()] if modes.ndim == 0 else modes + counts = counts[()] if counts.ndim == 0 else counts + return ModeResult(modes, counts) + + +def _put_val_to_limits(a, limits, inclusive, val=np.nan, xp=None): + """Replace elements outside limits with a value. + + This is primarily a utility function. + + Parameters + ---------- + a : array + limits : (float or None, float or None) + A tuple consisting of the (lower limit, upper limit). Elements in the + input array less than the lower limit or greater than the upper limit + will be replaced with `val`. None implies no limit. + inclusive : (bool, bool) + A tuple consisting of the (lower flag, upper flag). These flags + determine whether values exactly equal to lower or upper are allowed. + val : float, default: NaN + The value with which extreme elements of the array are replaced. + + """ + xp = array_namespace(a) if xp is None else xp + mask = xp.zeros_like(a, dtype=xp.bool) + if limits is None: + return a, mask + lower_limit, upper_limit = limits + lower_include, upper_include = inclusive + if lower_limit is not None: + mask = mask | ((a < lower_limit) if lower_include else a <= lower_limit) + if upper_limit is not None: + mask = mask | ((a > upper_limit) if upper_include else a >= upper_limit) + lazy = is_lazy_array(mask) + if not lazy and xp.all(mask): + raise ValueError("No array values within given limits") + if lazy or xp.any(mask): + a = xp.where(mask, val, a) + return a, mask + + +@xp_capabilities(marray=True) +@_axis_nan_policy_factory( + lambda x: x, n_outputs=1, default_axis=None, + result_to_tuple=lambda x, _: (x,) +) +def tmean(a, limits=None, inclusive=(True, True), axis=None): + """Compute the trimmed mean. + + This function finds the arithmetic mean of given values, ignoring values + outside the given `limits`. + + Parameters + ---------- + a : array_like + Array of values. + limits : None or (lower limit, upper limit), optional + Values in the input array less than the lower limit or greater than the + upper limit will be ignored. When limits is None (default), then all + values are used. Either of the limit values in the tuple can also be + None representing a half-open interval. + inclusive : (bool, bool), optional + A tuple consisting of the (lower flag, upper flag). These flags + determine whether values exactly equal to the lower or upper limits + are included. The default value is (True, True). + axis : int or None, optional + Axis along which to compute test. Default is None. + + Returns + ------- + tmean : ndarray + Trimmed mean. + + See Also + -------- + trim_mean : Returns mean after trimming a proportion from both tails. + + Examples + -------- + >>> import numpy as np + >>> from scipy import stats + >>> x = np.arange(20) + >>> stats.tmean(x) + 9.5 + >>> stats.tmean(x, (3,17)) + 10.0 + + """ + xp = array_namespace(a) + a, mask = _put_val_to_limits(a, limits, inclusive, val=0., xp=xp) + # explicit dtype specification required due to data-apis/array-api-compat#152 + sum = xp.sum(a, axis=axis, dtype=a.dtype) + n = xp.sum(xp.asarray(~mask, dtype=a.dtype, device=xp_device(a)), axis=axis, + dtype=a.dtype) + mean = xpx.apply_where(n != 0, (sum, n), operator.truediv, fill_value=xp.nan) + return mean[()] if mean.ndim == 0 else mean + + +@xp_capabilities(marray=True) +@_axis_nan_policy_factory( + lambda x: x, n_outputs=1, result_to_tuple=lambda x, _: (x,) +) +def tvar(a, limits=None, inclusive=(True, True), axis=0, ddof=1): + """Compute the trimmed variance. + + This function computes the sample variance of an array of values, + while ignoring values which are outside of given `limits`. + + Parameters + ---------- + a : array_like + Array of values. + limits : None or (lower limit, upper limit), optional + Values in the input array less than the lower limit or greater than the + upper limit will be ignored. When limits is None, then all values are + used. Either of the limit values in the tuple can also be None + representing a half-open interval. The default value is None. + inclusive : (bool, bool), optional + A tuple consisting of the (lower flag, upper flag). These flags + determine whether values exactly equal to the lower or upper limits + are included. The default value is (True, True). + axis : int or None, optional + Axis along which to operate. Default is 0. If None, compute over the + whole array `a`. + ddof : int, optional + Delta degrees of freedom. Default is 1. + + Returns + ------- + tvar : float + Trimmed variance. + + Notes + ----- + `tvar` computes the unbiased sample variance, i.e. it uses a correction + factor ``n / (n - 1)``. + + Examples + -------- + >>> import numpy as np + >>> from scipy import stats + >>> x = np.arange(20) + >>> stats.tvar(x) + 35.0 + >>> stats.tvar(x, (3,17)) + 20.0 + + """ + xp = array_namespace(a) + a, _ = _put_val_to_limits(a, limits, inclusive, xp=xp) + with warnings.catch_warnings(): + warnings.simplefilter("ignore", SmallSampleWarning) + # Currently, this behaves like nan_policy='omit' for alternative array + # backends, but nan_policy='propagate' will be handled for other backends + # by the axis_nan_policy decorator shortly. + return _xp_var(a, correction=ddof, axis=axis, nan_policy='omit', xp=xp) + + +@xp_capabilities(marray=True) +@_axis_nan_policy_factory( + lambda x: x, n_outputs=1, result_to_tuple=lambda x, _: (x,) +) +def tmin(a, lowerlimit=None, axis=0, inclusive=True, nan_policy='propagate'): + """Compute the trimmed minimum. + + This function finds the minimum value of an array `a` along the + specified axis, but only considering values greater than a specified + lower limit. + + Parameters + ---------- + a : array_like + Array of values. + lowerlimit : None or float, optional + Values in the input array less than the given limit will be ignored. + When lowerlimit is None, then all values are used. The default value + is None. + axis : int or None, optional + Axis along which to operate. Default is 0. If None, compute over the + whole array `a`. + inclusive : {True, False}, optional + This flag determines whether values exactly equal to the lower limit + are included. The default value is True. + + Returns + ------- + tmin : float, int or ndarray + Trimmed minimum. + + Examples + -------- + >>> import numpy as np + >>> from scipy import stats + >>> x = np.arange(20) + >>> stats.tmin(x) + 0 + + >>> stats.tmin(x, 13) + 13 + + >>> stats.tmin(x, 13, inclusive=False) + 14 + + """ + xp = array_namespace(a) + + max_ = xp.iinfo(a.dtype).max if xp.isdtype(a.dtype, 'integral') else xp.inf + a, mask = _put_val_to_limits(a, (lowerlimit, None), (inclusive, None), + val=max_, xp=xp) + + res = xp.min(a, axis=axis) + invalid = xp.all(mask, axis=axis) # All elements are below lowerlimit + + # For eager backends, output dtype is data-dependent + if is_lazy_array(invalid) or xp.any(invalid): + # Possible loss of precision for int types + res = xp_promote(res, force_floating=True, xp=xp) + res = xp.where(invalid, xp.nan, res) + + return res[()] if res.ndim == 0 else res + + +@xp_capabilities(marray=True) +@_axis_nan_policy_factory( + lambda x: x, n_outputs=1, result_to_tuple=lambda x, _: (x,) +) +def tmax(a, upperlimit=None, axis=0, inclusive=True, nan_policy='propagate'): + """Compute the trimmed maximum. + + This function computes the maximum value of an array along a given axis, + while ignoring values larger than a specified upper limit. + + Parameters + ---------- + a : array_like + Array of values. + upperlimit : None or float, optional + Values in the input array greater than the given limit will be ignored. + When upperlimit is None, then all values are used. The default value + is None. + axis : int or None, optional + Axis along which to operate. Default is 0. If None, compute over the + whole array `a`. + inclusive : {True, False}, optional + This flag determines whether values exactly equal to the upper limit + are included. The default value is True. + + Returns + ------- + tmax : float, int or ndarray + Trimmed maximum. + + Examples + -------- + >>> import numpy as np + >>> from scipy import stats + >>> x = np.arange(20) + >>> stats.tmax(x) + 19 + + >>> stats.tmax(x, 13) + 13 + + >>> stats.tmax(x, 13, inclusive=False) + 12 + + """ + xp = array_namespace(a) + + min_ = xp.iinfo(a.dtype).min if xp.isdtype(a.dtype, 'integral') else -xp.inf + a, mask = _put_val_to_limits(a, (None, upperlimit), (None, inclusive), + val=min_, xp=xp) + + res = xp.max(a, axis=axis) + invalid = xp.all(mask, axis=axis) # All elements are above upperlimit + + # For eager backends, output dtype is data-dependent + if is_lazy_array(invalid) or xp.any(invalid): + # Possible loss of precision for int types + res = xp_promote(res, force_floating=True, xp=xp) + res = xp.where(invalid, xp.nan, res) + + return res[()] if res.ndim == 0 else res + + +@xp_capabilities(marray=True) +@_axis_nan_policy_factory( + lambda x: x, n_outputs=1, result_to_tuple=lambda x, _: (x,) +) +def tstd(a, limits=None, inclusive=(True, True), axis=0, ddof=1): + """Compute the trimmed sample standard deviation. + + This function finds the sample standard deviation of given values, + ignoring values outside the given `limits`. + + Parameters + ---------- + a : array_like + Array of values. + limits : None or (lower limit, upper limit), optional + Values in the input array less than the lower limit or greater than the + upper limit will be ignored. When limits is None, then all values are + used. Either of the limit values in the tuple can also be None + representing a half-open interval. The default value is None. + inclusive : (bool, bool), optional + A tuple consisting of the (lower flag, upper flag). These flags + determine whether values exactly equal to the lower or upper limits + are included. The default value is (True, True). + axis : int or None, optional + Axis along which to operate. Default is 0. If None, compute over the + whole array `a`. + ddof : int, optional + Delta degrees of freedom. Default is 1. + + Returns + ------- + tstd : float + Trimmed sample standard deviation. + + Notes + ----- + `tstd` computes the unbiased sample standard deviation, i.e. it uses a + correction factor ``n / (n - 1)``. + + Examples + -------- + >>> import numpy as np + >>> from scipy import stats + >>> x = np.arange(20) + >>> stats.tstd(x) + 5.9160797830996161 + >>> stats.tstd(x, (3,17)) + 4.4721359549995796 + + """ + return tvar(a, limits, inclusive, axis, ddof, _no_deco=True)**0.5 + + +@xp_capabilities(marray=True) +@_axis_nan_policy_factory( + lambda x: x, n_outputs=1, result_to_tuple=lambda x, _: (x,) +) +def tsem(a, limits=None, inclusive=(True, True), axis=0, ddof=1): + """Compute the trimmed standard error of the mean. + + This function finds the standard error of the mean for given + values, ignoring values outside the given `limits`. + + Parameters + ---------- + a : array_like + Array of values. + limits : None or (lower limit, upper limit), optional + Values in the input array less than the lower limit or greater than the + upper limit will be ignored. When limits is None, then all values are + used. Either of the limit values in the tuple can also be None + representing a half-open interval. The default value is None. + inclusive : (bool, bool), optional + A tuple consisting of the (lower flag, upper flag). These flags + determine whether values exactly equal to the lower or upper limits + are included. The default value is (True, True). + axis : int or None, optional + Axis along which to operate. Default is 0. If None, compute over the + whole array `a`. + ddof : int, optional + Delta degrees of freedom. Default is 1. + + Returns + ------- + tsem : float + Trimmed standard error of the mean. + + Notes + ----- + `tsem` uses unbiased sample standard deviation, i.e. it uses a + correction factor ``n / (n - 1)``. + + Examples + -------- + >>> import numpy as np + >>> from scipy import stats + >>> x = np.arange(20) + >>> stats.tsem(x) + 1.3228756555322954 + >>> stats.tsem(x, (3,17)) + 1.1547005383792515 + + """ + xp = array_namespace(a) + a, _ = _put_val_to_limits(a, limits, inclusive, xp=xp) + + with warnings.catch_warnings(): + warnings.simplefilter("ignore", SmallSampleWarning) + # Currently, this behaves like nan_policy='omit' for alternative array + # backends, but nan_policy='propagate' will be handled for other backends + # by the axis_nan_policy decorator shortly. + sd = _xp_var(a, correction=ddof, axis=axis, nan_policy='omit', xp=xp)**0.5 + + not_nan = xp.astype(~xp.isnan(a), a.dtype) + n_obs = xp.sum(not_nan, axis=axis, dtype=sd.dtype) + return sd / n_obs**0.5 + + +##################################### +# MOMENTS # +##################################### + + +def _moment_outputs(kwds, default_order=1): + order = np.atleast_1d(kwds.get('order', default_order)) + message = "`order` must be a scalar or a non-empty 1D array." + if order.size == 0 or order.ndim > 1: + raise ValueError(message) + return len(order) + + +def _moment_result_object(*args): + if len(args) == 1: + return args[0] + xp = array_namespace(*args) + return xp.stack(args) + + +# When `order` is array-like with size > 1, moment produces an *array* +# rather than a tuple, but the zeroth dimension is to be treated like +# separate outputs. It is important to make the distinction between +# separate outputs when adding the reduced axes back (`keepdims=True`). +def _moment_tuple(x, n_out): + return tuple(x[i, ...] for i in range(x.shape[0])) if n_out > 1 else (x,) + + +# `moment` fits into the `_axis_nan_policy` pattern, but it is a bit unusual +# because the number of outputs is variable. Specifically, +# `result_to_tuple=lambda x: (x,)` may be surprising for a function that +# can produce more than one output, but it is intended here. +# When `moment is called to produce the output: +# - `result_to_tuple` packs the returned array into a single-element tuple, +# - `_moment_result_object` extracts and returns that single element. +# However, when the input array is empty, `moment` is never called. Instead, +# - `_check_empty_inputs` is used to produce an empty array with the +# appropriate dimensions. +# - A list comprehension creates the appropriate number of copies of this +# array, depending on `n_outputs`. +# - This list - which may have multiple elements - is passed into +# `_moment_result_object`. +# - If there is a single output, `_moment_result_object` extracts and returns +# the single output from the list. +# - If there are multiple outputs, and therefore multiple elements in the list, +# `_moment_result_object` converts the list of arrays to a single array and +# returns it. +# Currently, this leads to a slight inconsistency: when the input array is +# empty, there is no distinction between the `moment` function being called +# with parameter `order=1` and `order=[1]`; the latter *should* produce +# the same as the former but with a singleton zeroth dimension. +@xp_capabilities(skip_backends=[('dask.array', 'needs axis_nan_policy decorator')], + marray=True) +@_rename_parameter('moment', 'order') +@_axis_nan_policy_factory( # noqa: E302 + _moment_result_object, n_samples=1, result_to_tuple=_moment_tuple, + n_outputs=_moment_outputs +) +def moment(a, order=1, axis=0, nan_policy='propagate', *, center=None): + r"""Calculate the nth moment about the mean for a sample. + + A moment is a specific quantitative measure of the shape of a set of + points. It is often used to calculate coefficients of skewness and kurtosis + due to its close relationship with them. + + Parameters + ---------- + a : array_like + Input array. + order : int or 1-D array_like of ints, optional + Order of central moment that is returned. Default is 1. + axis : int or None, optional + Axis along which the central moment is computed. Default is 0. + If None, compute over the whole array `a`. + nan_policy : {'propagate', 'raise', 'omit'}, optional + Defines how to handle when input contains nan. + The following options are available (default is 'propagate'): + + * 'propagate': returns nan + * 'raise': throws an error + * 'omit': performs the calculations ignoring nan values + + center : float or None, optional + The point about which moments are taken. This can be the sample mean, + the origin, or any other be point. If `None` (default) compute the + center as the sample mean. + + Returns + ------- + n-th moment about the `center` : ndarray or float + The appropriate moment along the given axis or over all values if axis + is None. The denominator for the moment calculation is the number of + observations, no degrees of freedom correction is done. + + See Also + -------- + kurtosis, skew, describe + + Notes + ----- + The k-th moment of a data sample is: + + .. math:: + + m_k = \frac{1}{n} \sum_{i = 1}^n (x_i - c)^k + + Where `n` is the number of samples, and `c` is the center around which the + moment is calculated. This function uses exponentiation by squares [1]_ for + efficiency. + + Note that, if `a` is an empty array (``a.size == 0``), array `moment` with + one element (`moment.size == 1`) is treated the same as scalar `moment` + (``np.isscalar(moment)``). This might produce arrays of unexpected shape. + + References + ---------- + .. [1] https://eli.thegreenplace.net/2009/03/21/efficient-integer-exponentiation-algorithms + + Examples + -------- + >>> from scipy.stats import moment + >>> moment([1, 2, 3, 4, 5], order=1) + 0.0 + >>> moment([1, 2, 3, 4, 5], order=2) + 2.0 + + """ + xp = array_namespace(a) + a, center, order = xp_promote(a, center, order, force_floating=True, xp=xp) + + if not is_lazy_array(order) and xp.any(order != xp.round(order)): + raise ValueError("All elements of `order` must be integral.") + + # _axis_nan_policy decorator ensures that axis=-1 + if order.ndim > 0: + order = xp.reshape(order, (-1,) + (1,)*a.ndim) + return _moment(a, order, axis=-1, center=center) + else: + res = _moment(a, order, axis=-1, center=center) + return res[()] if res.ndim == 0 else res + + +def _demean(a, mean, axis, *, xp, precision_warning=True): + # subtracts `mean` from `a` and returns the result, + # warning if there is catastrophic cancellation. `mean` + # must be the mean of `a` along axis with `keepdims=True`. + # Used in e.g. `_moment`, `_zscore`, `_xp_var`. See gh-15905. + a_zero_mean = a - mean + + if (xp_size(a_zero_mean) == 0 or not precision_warning + or is_lazy_array(a_zero_mean)): + return a_zero_mean + + eps = xp.finfo(mean.dtype).eps * 10 + + with np.errstate(divide='ignore', invalid='ignore'): + rel_diff = xp.max(xp.abs(a_zero_mean), axis=axis, + keepdims=True) / xp.abs(mean) + + n = _count_nonmasked(a, axis, xp=xp) + with np.errstate(invalid='ignore'): + device = {'device': xp_device(a)} + precision_loss = xp.any(xp.asarray(rel_diff < eps, **device) + & xp.asarray(n > 1, **device)) + + if precision_loss: + message = ("Precision loss occurred in moment calculation due to " + "catastrophic cancellation. This occurs when the data " + "are nearly identical. Results may be unreliable.") + warnings.warn(message, RuntimeWarning, stacklevel=5) + return a_zero_mean + + +def _moment(a, order, axis, *, center=None, xp=None): + """Vectorized calculation of raw moment about specified center + + When `center` is None, the mean is computed and used as the center; + otherwise, the provided value is used as the center. + """ + xp = array_namespace(a) if xp is None else xp + + order = xp.asarray(order, dtype=a.dtype, device=xp_device(a)) + order_0 = order == 0 + order_1 = (order == 1) & (center is None) + center = xp.mean(a, axis=axis, keepdims=True) if center is None else center + a_zero_mean = _demean(a, center, axis, xp=xp) + res = xp.mean(a_zero_mean**order, axis=axis, keepdims=True) + if a.shape[-1] > 0 and (is_lazy_array(res) + or xp.any(order_0) or xp.any(order_1)): + res = xp.where(order_0, xp.ones_like(res), res) + res = xp.where(order_1, xp.zeros_like(res), res) + + return xp.squeeze(res, axis=axis) + + +def _var(x, axis=0, ddof=0, mean=None, xp=None): + # Calculate variance of sample, warning if precision is lost + xp = array_namespace(x) if xp is None else xp + var = _moment(x, 2, axis, center=mean, xp=xp) + if ddof != 0: + n = _count_nonmasked(x, axis, xp=xp) + n = xp.asarray(n, dtype=x.dtype, device=xp_device(x)) + var *= (n / (n-ddof)) # to avoid error on division by zero + return var + + +@xp_capabilities(marray=True) +@_axis_nan_policy_factory( + lambda x: x, result_to_tuple=lambda x, _: (x,), n_outputs=1 +) +# nan_policy handled by `_axis_nan_policy`, but needs to be left +# in signature to preserve use as a positional argument +def skew(a, axis=0, bias=True, nan_policy='propagate'): + r"""Compute the sample skewness of a data set. + + For normally distributed data, the skewness should be about zero. For + unimodal continuous distributions, a skewness value greater than zero means + that there is more weight in the right tail of the distribution. The + function `skewtest` can be used to determine if the skewness value + is close enough to zero, statistically speaking. + + Parameters + ---------- + a : ndarray + Input array. + axis : int or None, optional + Axis along which skewness is calculated. Default is 0. + If None, compute over the whole array `a`. + bias : bool, optional + If False, then the calculations are corrected for statistical bias. + nan_policy : {'propagate', 'raise', 'omit'}, optional + Defines how to handle when input contains nan. + The following options are available (default is 'propagate'): + + * 'propagate': returns nan + * 'raise': throws an error + * 'omit': performs the calculations ignoring nan values + + Returns + ------- + skewness : ndarray + The skewness of values along an axis, returning NaN where all values + are equal. + + Notes + ----- + The sample skewness is computed as the Fisher-Pearson coefficient + of skewness, i.e. + + .. math:: + + g_1=\frac{m_3}{m_2^{3/2}} + + where + + .. math:: + + m_i=\frac{1}{N}\sum_{n=1}^N(x[n]-\bar{x})^i + + is the biased sample :math:`i\texttt{th}` central moment, and + :math:`\bar{x}` is + the sample mean. If ``bias`` is False, the calculations are + corrected for bias and the value computed is the adjusted + Fisher-Pearson standardized moment coefficient, i.e. + + .. math:: + + G_1=\frac{k_3}{k_2^{3/2}}= + \frac{\sqrt{N(N-1)}}{N-2}\frac{m_3}{m_2^{3/2}}. + + References + ---------- + .. [1] Zwillinger, D. and Kokoska, S. (2000). CRC Standard + Probability and Statistics Tables and Formulae. Chapman & Hall: New + York. 2000. + Section 2.2.24.1 + + Examples + -------- + >>> from scipy.stats import skew + >>> skew([1, 2, 3, 4, 5]) + 0.0 + >>> skew([2, 8, 0, 4, 1, 9, 9, 0]) + 0.2650554122698573 + + """ + xp = array_namespace(a) + a, axis = _chk_asarray(a, axis, xp=xp) + n = _count_nonmasked(a, axis, xp=xp) + + mean = xp.mean(a, axis=axis, keepdims=True) + mean_reduced = xp.squeeze(mean, axis=axis) # needed later + m2 = _moment(a, 2, axis, center=mean, xp=xp) + m3 = _moment(a, 3, axis, center=mean, xp=xp) + with np.errstate(all='ignore'): + eps = xp.finfo(m2.dtype).eps + zero = m2 <= (eps * mean_reduced)**2 + vals = xp.where(zero, xp.nan, m3 / m2**1.5) + if not bias: + can_correct = ~zero & (n > 2) + if is_lazy_array(can_correct) or xp.any(can_correct): + nval = ((n - 1.0) * n)**0.5 / (n - 2.0) * m3 / m2**1.5 + vals = xp.where(can_correct, nval, vals) + + return vals[()] if vals.ndim == 0 else vals + + +@xp_capabilities(marray=True) +@_axis_nan_policy_factory( + lambda x: x, result_to_tuple=lambda x, _: (x,), n_outputs=1 +) +# nan_policy handled by `_axis_nan_policy`, but needs to be left +# in signature to preserve use as a positional argument +def kurtosis(a, axis=0, fisher=True, bias=True, nan_policy='propagate'): + """Compute the kurtosis (Fisher or Pearson) of a dataset. + + Kurtosis is the fourth central moment divided by the square of the + variance. If Fisher's definition is used, then 3.0 is subtracted from + the result to give 0.0 for a normal distribution. + + If bias is False then the kurtosis is calculated using k statistics to + eliminate bias coming from biased moment estimators + + Use `kurtosistest` to see if result is close enough to normal. + + Parameters + ---------- + a : array + Data for which the kurtosis is calculated. + axis : int or None, optional + Axis along which the kurtosis is calculated. Default is 0. + If None, compute over the whole array `a`. + fisher : bool, optional + If True, Fisher's definition is used (normal ==> 0.0). If False, + Pearson's definition is used (normal ==> 3.0). + bias : bool, optional + If False, then the calculations are corrected for statistical bias. + nan_policy : {'propagate', 'raise', 'omit'}, optional + Defines how to handle when input contains nan. 'propagate' returns nan, + 'raise' throws an error, 'omit' performs the calculations ignoring nan + values. Default is 'propagate'. + + Returns + ------- + kurtosis : array + The kurtosis of values along an axis, returning NaN where all values + are equal. + + References + ---------- + .. [1] Zwillinger, D. and Kokoska, S. (2000). CRC Standard + Probability and Statistics Tables and Formulae. Chapman & Hall: New + York. 2000. + + Examples + -------- + In Fisher's definition, the kurtosis of the normal distribution is zero. + In the following example, the kurtosis is close to zero, because it was + calculated from the dataset, not from the continuous distribution. + + >>> import numpy as np + >>> from scipy.stats import norm, kurtosis + >>> data = norm.rvs(size=1000, random_state=3) + >>> kurtosis(data) + -0.06928694200380558 + + The distribution with a higher kurtosis has a heavier tail. + The zero valued kurtosis of the normal distribution in Fisher's definition + can serve as a reference point. + + >>> import matplotlib.pyplot as plt + >>> import scipy.stats as stats + >>> from scipy.stats import kurtosis + + >>> x = np.linspace(-5, 5, 100) + >>> ax = plt.subplot() + >>> distnames = ['laplace', 'norm', 'uniform'] + + >>> for distname in distnames: + ... if distname == 'uniform': + ... dist = getattr(stats, distname)(loc=-2, scale=4) + ... else: + ... dist = getattr(stats, distname) + ... data = dist.rvs(size=1000) + ... kur = kurtosis(data, fisher=True) + ... y = dist.pdf(x) + ... ax.plot(x, y, label="{}, {}".format(distname, round(kur, 3))) + ... ax.legend() + + The Laplace distribution has a heavier tail than the normal distribution. + The uniform distribution (which has negative kurtosis) has the thinnest + tail. + + """ + xp = array_namespace(a) + a, axis = _chk_asarray(a, axis, xp=xp) + + n = _count_nonmasked(a, axis, xp=xp) + mean = xp.mean(a, axis=axis, keepdims=True) + mean_reduced = xp.squeeze(mean, axis=axis) # needed later + m2 = _moment(a, 2, axis, center=mean, xp=xp) + m4 = _moment(a, 4, axis, center=mean, xp=xp) + with np.errstate(all='ignore'): + zero = m2 <= (xp.finfo(m2.dtype).eps * mean_reduced)**2 + vals = xp.where(zero, xp.nan, m4 / m2**2.0) + + if not bias: + can_correct = ~zero & (n > 3) + if is_lazy_array(can_correct) or xp.any(can_correct): + nval = 1.0/(n-2)/(n-3) * ((n**2-1.0)*m4/m2**2.0 - 3*(n-1)**2.0) + vals = xp.where(can_correct, nval + 3.0, vals) + + vals = vals - 3 if fisher else vals + return vals[()] if vals.ndim == 0 else vals + + +DescribeResult = namedtuple('DescribeResult', + ('nobs', 'minmax', 'mean', 'variance', 'skewness', + 'kurtosis')) + + +@xp_capabilities(marray=True) +def describe(a, axis=0, ddof=1, bias=True, nan_policy='propagate'): + """Compute several descriptive statistics of the passed array. + + Parameters + ---------- + a : array_like + Input data. + axis : int or None, optional + Axis along which statistics are calculated. Default is 0. + If None, compute over the whole array `a`. + ddof : int, optional + Delta degrees of freedom (only for variance). Default is 1. + bias : bool, optional + If False, then the skewness and kurtosis calculations are corrected + for statistical bias. + nan_policy : {'propagate', 'raise', 'omit'}, optional + Defines how to handle when input contains nan. + The following options are available (default is 'propagate'): + + * 'propagate': returns nan + * 'raise': throws an error + * 'omit': performs the calculations ignoring nan values + + Returns + ------- + nobs : int or ndarray of ints + Number of observations (length of data along `axis`). + When 'omit' is chosen as nan_policy, the length along each axis + slice is counted separately. + minmax: tuple of ndarrays or floats + Minimum and maximum value of `a` along the given axis. + mean : ndarray or float + Arithmetic mean of `a` along the given axis. + variance : ndarray or float + Unbiased variance of `a` along the given axis; denominator is number + of observations minus one. + skewness : ndarray or float + Skewness of `a` along the given axis, based on moment calculations + with denominator equal to the number of observations, i.e. no degrees + of freedom correction. + kurtosis : ndarray or float + Kurtosis (Fisher) of `a` along the given axis. The kurtosis is + normalized so that it is zero for the normal distribution. No + degrees of freedom are used. + + Raises + ------ + ValueError + If size of `a` is 0. + + See Also + -------- + skew, kurtosis + + Examples + -------- + >>> import numpy as np + >>> from scipy import stats + >>> a = np.arange(10) + >>> stats.describe(a) + DescribeResult(nobs=10, minmax=(0, 9), mean=4.5, + variance=9.166666666666666, skewness=0.0, + kurtosis=-1.2242424242424244) + >>> b = [[1, 2], [3, 4]] + >>> stats.describe(b) + DescribeResult(nobs=2, minmax=(array([1, 2]), array([3, 4])), + mean=array([2., 3.]), variance=array([2., 2.]), + skewness=array([0., 0.]), kurtosis=array([-2., -2.])) + + """ + xp = array_namespace(a) + a, axis = _chk_asarray(a, axis, xp=xp) + + contains_nan = _contains_nan(a, nan_policy) + + # Test nan_policy before the implicit call to bool(contains_nan) + # to avoid raising on lazy xps on the default nan_policy='propagate' + if nan_policy == 'omit' and contains_nan: + # only NumPy gets here; `_contains_nan` raises error for the rest + a = ma.masked_invalid(a) + return mstats_basic.describe(a, axis, ddof, bias) + + if xp_size(a) == 0: + raise ValueError("The input must not be empty.") + + # use xp.astype when data-apis/array-api-compat#226 is resolved + n = xp.asarray(_count_nonmasked(a, axis, xp=xp), dtype=xp.int64, + device=xp_device(a)) + n = n[()] if n.ndim == 0 else n + mm = (xp.min(a, axis=axis), xp.max(a, axis=axis)) + a = xp_promote(a, force_floating=True, xp=xp) + m = xp.mean(a, axis=axis) + v = _var(a, axis=axis, ddof=ddof, xp=xp) + v = v[()] if v.ndim == 0 else v + sk = skew(a, axis, bias=bias) + kurt = kurtosis(a, axis, bias=bias) + + return DescribeResult(n, mm, m, v, sk, kurt) + +##################################### +# NORMALITY TESTS # +##################################### + + +def _get_pvalue(statistic, distribution, alternative, symmetric=True, xp=None): + """Get p-value given the statistic, (continuous) distribution, and alternative""" + xp = array_namespace(statistic) if xp is None else xp + + if alternative == 'less': + pvalue = distribution.cdf(statistic) + elif alternative == 'greater': + pvalue = distribution.sf(statistic) + elif alternative == 'two-sided': + pvalue = 2 * (distribution.sf(xp.abs(statistic)) if symmetric + else xp.minimum(distribution.cdf(statistic), + distribution.sf(statistic))) + else: + message = "`alternative` must be 'less', 'greater', or 'two-sided'." + raise ValueError(message) + + return pvalue + + +SkewtestResult = namedtuple('SkewtestResult', ('statistic', 'pvalue')) + + +@xp_capabilities(marray=True) +@_axis_nan_policy_factory(SkewtestResult, n_samples=1, too_small=7) +# nan_policy handled by `_axis_nan_policy`, but needs to be left +# in signature to preserve use as a positional argument +def skewtest(a, axis=0, nan_policy='propagate', alternative='two-sided'): + r"""Test whether the skew is different from the normal distribution. + + This function tests the null hypothesis that the skewness of + the population that the sample was drawn from is the same + as that of a corresponding normal distribution. + + Parameters + ---------- + a : array + The data to be tested. Must contain at least eight observations. + axis : int or None, optional + Axis along which statistics are calculated. Default is 0. + If None, compute over the whole array `a`. + nan_policy : {'propagate', 'raise', 'omit'}, optional + Defines how to handle when input contains nan. + The following options are available (default is 'propagate'): + + * 'propagate': returns nan + * 'raise': throws an error + * 'omit': performs the calculations ignoring nan values + + alternative : {'two-sided', 'less', 'greater'}, optional + Defines the alternative hypothesis. Default is 'two-sided'. + The following options are available: + + * 'two-sided': the skewness of the distribution underlying the sample + is different from that of the normal distribution (i.e. 0) + * 'less': the skewness of the distribution underlying the sample + is less than that of the normal distribution + * 'greater': the skewness of the distribution underlying the sample + is greater than that of the normal distribution + + .. versionadded:: 1.7.0 + + Returns + ------- + statistic : float + The computed z-score for this test. + pvalue : float + The p-value for the hypothesis test. + + See Also + -------- + :ref:`hypothesis_skewtest` : Extended example + + Notes + ----- + The sample size must be at least 8. + + References + ---------- + .. [1] R. B. D'Agostino, A. J. Belanger and R. B. D'Agostino Jr., + "A suggestion for using powerful and informative tests of + normality", American Statistician 44, pp. 316-321, 1990. + + Examples + -------- + + >>> from scipy.stats import skewtest + >>> skewtest([1, 2, 3, 4, 5, 6, 7, 8]) + SkewtestResult(statistic=1.0108048609177787, pvalue=0.3121098361421897) + >>> skewtest([2, 8, 0, 4, 1, 9, 9, 0]) + SkewtestResult(statistic=0.44626385374196975, pvalue=0.6554066631275459) + >>> skewtest([1, 2, 3, 4, 5, 6, 7, 8000]) + SkewtestResult(statistic=3.571773510360407, pvalue=0.0003545719905823133) + >>> skewtest([100, 100, 100, 100, 100, 100, 100, 101]) + SkewtestResult(statistic=3.5717766638478072, pvalue=0.000354567720281634) + >>> skewtest([1, 2, 3, 4, 5, 6, 7, 8], alternative='less') + SkewtestResult(statistic=1.0108048609177787, pvalue=0.8439450819289052) + >>> skewtest([1, 2, 3, 4, 5, 6, 7, 8], alternative='greater') + SkewtestResult(statistic=1.0108048609177787, pvalue=0.15605491807109484) + + For a more detailed example, see :ref:`hypothesis_skewtest`. + """ + xp = array_namespace(a) + a, axis = _chk_asarray(a, axis, xp=xp) + + b2 = skew(a, axis, _no_deco=True) + + n = xp.asarray(_count_nonmasked(a, axis), dtype=b2.dtype, device=xp_device(a)) + n = xpx.at(n, n < 8).set(xp.nan) + + with np.errstate(divide='ignore', invalid='ignore'): + y = b2 * xp.sqrt(((n + 1) * (n + 3)) / (6.0 * (n - 2))) + beta2 = (3.0 * (n**2 + 27*n - 70) * (n+1) * (n+3) / + ((n-2.0) * (n+5) * (n+7) * (n+9))) + W2 = -1 + xp.sqrt(2 * (beta2 - 1)) + delta = 1 / xp.sqrt(0.5 * xp.log(W2)) + alpha = xp.sqrt(2.0 / (W2 - 1)) + y = xp.where(y == 0, 1., y) + Z = delta * xp.log(y / alpha + xp.sqrt((y / alpha)**2 + 1)) + pvalue = _get_pvalue(Z, _SimpleNormal(), alternative, xp=xp) + + Z = Z[()] if Z.ndim == 0 else Z + pvalue = pvalue[()] if pvalue.ndim == 0 else pvalue + return SkewtestResult(Z, pvalue) + + +KurtosistestResult = namedtuple('KurtosistestResult', ('statistic', 'pvalue')) + + +@xp_capabilities(marray=True) +@_axis_nan_policy_factory(KurtosistestResult, n_samples=1, too_small=4) +def kurtosistest(a, axis=0, nan_policy='propagate', alternative='two-sided'): + r"""Test whether a dataset has normal kurtosis. + + This function tests the null hypothesis that the kurtosis + of the population from which the sample was drawn is that + of the normal distribution. + + Parameters + ---------- + a : array + Array of the sample data. Must contain at least five observations. + axis : int or None, optional + Axis along which to compute test. Default is 0. If None, + compute over the whole array `a`. + nan_policy : {'propagate', 'raise', 'omit'}, optional + Defines how to handle when input contains nan. + The following options are available (default is 'propagate'): + + * 'propagate': returns nan + * 'raise': throws an error + * 'omit': performs the calculations ignoring nan values + + alternative : {'two-sided', 'less', 'greater'}, optional + Defines the alternative hypothesis. + The following options are available (default is 'two-sided'): + + * 'two-sided': the kurtosis of the distribution underlying the sample + is different from that of the normal distribution + * 'less': the kurtosis of the distribution underlying the sample + is less than that of the normal distribution + * 'greater': the kurtosis of the distribution underlying the sample + is greater than that of the normal distribution + + .. versionadded:: 1.7.0 + + Returns + ------- + statistic : float + The computed z-score for this test. + pvalue : float + The p-value for the hypothesis test. + + See Also + -------- + :ref:`hypothesis_kurtosistest` : Extended example + + Notes + ----- + Valid only for n>20. This function uses the method described in [1]_. + + References + ---------- + .. [1] F. J. Anscombe, W. J. Glynn, "Distribution of the kurtosis + statistic b2 for normal samples", Biometrika, vol. 70, pp. 227-234, 1983. + + Examples + -------- + + >>> import numpy as np + >>> from scipy.stats import kurtosistest + >>> kurtosistest(list(range(20))) + KurtosistestResult(statistic=-1.7058104152122062, pvalue=0.08804338332528348) + >>> kurtosistest(list(range(20)), alternative='less') + KurtosistestResult(statistic=-1.7058104152122062, pvalue=0.04402169166264174) + >>> kurtosistest(list(range(20)), alternative='greater') + KurtosistestResult(statistic=-1.7058104152122062, pvalue=0.9559783083373583) + >>> rng = np.random.default_rng() + >>> s = rng.normal(0, 1, 1000) + >>> kurtosistest(s) + KurtosistestResult(statistic=-1.475047944490622, pvalue=0.14019965402996987) + + For a more detailed example, see :ref:`hypothesis_kurtosistest`. + """ + xp = array_namespace(a) + a, axis = _chk_asarray(a, axis, xp=xp) + + b2 = kurtosis(a, axis, fisher=False, _no_deco=True) + + n = xp.asarray(_count_nonmasked(a, axis), dtype=b2.dtype, device=xp_device(a)) + n = xpx.at(n, n < 5).set(xp.nan) + + E = 3.0*(n-1) / (n+1) + varb2 = 24.0*n*(n-2)*(n-3) / ((n+1)*(n+1.)*(n+3)*(n+5)) # [1]_ Eq. 1 + x = (b2-E) / varb2**0.5 # [1]_ Eq. 4 + # [1]_ Eq. 2: + sqrtbeta1 = 6.0*(n*n-5*n+2)/((n+7)*(n+9)) * ((6.0*(n+3)*(n+5)) + / (n*(n-2)*(n-3)))**0.5 + # [1]_ Eq. 3: + A = 6.0 + 8.0/sqrtbeta1 * (2.0/sqrtbeta1 + (1+4.0/(sqrtbeta1**2))**0.5) + term1 = 1 - 2/(9.0*A) + denom = 1 + x * (2/(A-4.0))**0.5 + term2 = xp.sign(denom) * xp.where(denom == 0.0, xp.nan, + ((1-2.0/A)/xp.abs(denom))**(1/3)) + if not is_lazy_array(denom) and xp.any(denom == 0): + msg = ("Test statistic not defined in some cases due to division by " + "zero. Return nan in that case...") + warnings.warn(msg, RuntimeWarning, stacklevel=2) + + Z = (term1 - term2) / (2/(9.0*A))**0.5 # [1]_ Eq. 5 + pvalue = _get_pvalue(Z, _SimpleNormal(), alternative, xp=xp) + + Z = Z[()] if Z.ndim == 0 else Z + pvalue = pvalue[()] if pvalue.ndim == 0 else pvalue + return KurtosistestResult(Z, pvalue) + + +NormaltestResult = namedtuple('NormaltestResult', ('statistic', 'pvalue')) + + +@xp_capabilities(marray=True) +@_axis_nan_policy_factory(NormaltestResult, n_samples=1, too_small=7) +def normaltest(a, axis=0, nan_policy='propagate'): + r"""Test whether a sample differs from a normal distribution. + + This function tests the null hypothesis that a sample comes + from a normal distribution. It is based on D'Agostino and + Pearson's [1]_, [2]_ test that combines skew and kurtosis to + produce an omnibus test of normality. + + Parameters + ---------- + a : array_like + The array containing the sample to be tested. Must contain + at least eight observations. + axis : int or None, optional + Axis along which to compute test. Default is 0. If None, + compute over the whole array `a`. + nan_policy : {'propagate', 'raise', 'omit'}, optional + Defines how to handle when input contains nan. + The following options are available (default is 'propagate'): + + * 'propagate': returns nan + * 'raise': throws an error + * 'omit': performs the calculations ignoring nan values + + Returns + ------- + statistic : float or array + ``s^2 + k^2``, where ``s`` is the z-score returned by `skewtest` and + ``k`` is the z-score returned by `kurtosistest`. + pvalue : float or array + A 2-sided chi squared probability for the hypothesis test. + + See Also + -------- + :ref:`hypothesis_normaltest` : Extended example + + References + ---------- + .. [1] D'Agostino, R. B. (1971), "An omnibus test of normality for + moderate and large sample size", Biometrika, 58, 341-348 + .. [2] D'Agostino, R. and Pearson, E. S. (1973), "Tests for departure from + normality", Biometrika, 60, 613-622 + + Examples + -------- + + >>> import numpy as np + >>> from scipy import stats + >>> rng = np.random.default_rng() + >>> pts = 1000 + >>> a = rng.normal(0, 1, size=pts) + >>> b = rng.normal(2, 1, size=pts) + >>> x = np.concatenate((a, b)) + >>> res = stats.normaltest(x) + >>> res.statistic + 53.619... # random + >>> res.pvalue + 2.273917413209226e-12 # random + + For a more detailed example, see :ref:`hypothesis_normaltest`. + """ + xp = array_namespace(a) + + s, _ = skewtest(a, axis, _no_deco=True) + k, _ = kurtosistest(a, axis, _no_deco=True) + statistic = s*s + k*k + + chi2 = _SimpleChi2(xp.asarray(2., dtype=statistic.dtype, device=xp_device(a))) + pvalue = _get_pvalue(statistic, chi2, alternative='greater', symmetric=False, xp=xp) + + statistic = statistic[()] if statistic.ndim == 0 else statistic + pvalue = pvalue[()] if pvalue.ndim == 0 else pvalue + + return NormaltestResult(statistic, pvalue) + + +@xp_capabilities(marray=True) +@_axis_nan_policy_factory(SignificanceResult, default_axis=None) +def jarque_bera(x, *, axis=None): + r"""Perform the Jarque-Bera goodness of fit test on sample data. + + The Jarque-Bera test tests whether the sample data has the skewness and + kurtosis matching a normal distribution. + + Note that this test only works for a large enough number of data samples + (>2000) as the test statistic asymptotically has a Chi-squared distribution + with 2 degrees of freedom. + + Parameters + ---------- + x : array_like + Observations of a random variable. + axis : int or None, default: 0 + If an int, the axis of the input along which to compute the statistic. + The statistic of each axis-slice (e.g. row) of the input will appear in + a corresponding element of the output. + If ``None``, the input will be raveled before computing the statistic. + + Returns + ------- + result : SignificanceResult + An object with the following attributes: + + statistic : float + The test statistic. + pvalue : float + The p-value for the hypothesis test. + + See Also + -------- + :ref:`hypothesis_jarque_bera` : Extended example + + References + ---------- + .. [1] Jarque, C. and Bera, A. (1980) "Efficient tests for normality, + homoscedasticity and serial independence of regression residuals", + 6 Econometric Letters 255-259. + + Examples + -------- + + >>> import numpy as np + >>> from scipy import stats + >>> rng = np.random.default_rng() + >>> x = rng.normal(0, 1, 100000) + >>> jarque_bera_test = stats.jarque_bera(x) + >>> jarque_bera_test + Jarque_beraResult(statistic=3.3415184718131554, pvalue=0.18810419594996775) + >>> jarque_bera_test.statistic + 3.3415184718131554 + >>> jarque_bera_test.pvalue + 0.18810419594996775 + + For a more detailed example, see :ref:`hypothesis_jarque_bera`. + """ + xp = array_namespace(x) + x, axis = _chk_asarray(x, axis, xp=xp) + + mu = _xp_mean(x, axis=axis, keepdims=True) + diffx = x - mu + s = skew(diffx, axis=axis, _no_deco=True) + k = kurtosis(diffx, axis=axis, _no_deco=True) + + n = xp.asarray(_count_nonmasked(x, axis), dtype=mu.dtype, device=xp_device(x)) + statistic = n / 6 * (s**2 + k**2 / 4) + + chi2 = _SimpleChi2(xp.asarray(2., dtype=mu.dtype, device=xp_device(x))) + pvalue = _get_pvalue(statistic, chi2, alternative='greater', symmetric=False, xp=xp) + + statistic = statistic[()] if statistic.ndim == 0 else statistic + pvalue = pvalue[()] if pvalue.ndim == 0 else pvalue + + return SignificanceResult(statistic, pvalue) + + +##################################### +# FREQUENCY FUNCTIONS # +##################################### + + +@xp_capabilities(np_only=True) +def scoreatpercentile(a, per, limit=(), interpolation_method='fraction', + axis=None): + """Calculate the score at a given percentile of the input sequence. + + For example, the score at ``per=50`` is the median. If the desired quantile + lies between two data points, we interpolate between them, according to + the value of `interpolation`. If the parameter `limit` is provided, it + should be a tuple (lower, upper) of two values. + + Parameters + ---------- + a : array_like + A 1-D array of values from which to extract score. + per : array_like + Percentile(s) at which to extract score. Values should be in range + [0,100]. + limit : tuple, optional + Tuple of two scalars, the lower and upper limits within which to + compute the percentile. Values of `a` outside + this (closed) interval will be ignored. + interpolation_method : {'fraction', 'lower', 'higher'}, optional + Specifies the interpolation method to use, + when the desired quantile lies between two data points `i` and `j` + The following options are available (default is 'fraction'): + + * 'fraction': ``i + (j - i) * fraction`` where ``fraction`` is the + fractional part of the index surrounded by ``i`` and ``j`` + * 'lower': ``i`` + * 'higher': ``j`` + + axis : int, optional + Axis along which the percentiles are computed. Default is None. If + None, compute over the whole array `a`. + + Returns + ------- + score : float or ndarray + Score at percentile(s). + + See Also + -------- + percentileofscore, numpy.percentile + + Notes + ----- + This function will become obsolete in the future. + For NumPy 1.9 and higher, `numpy.percentile` provides all the functionality + that `scoreatpercentile` provides. And it's significantly faster. + Therefore it's recommended to use `numpy.percentile` for users that have + numpy >= 1.9. + + Examples + -------- + >>> import numpy as np + >>> from scipy import stats + >>> a = np.arange(100) + >>> stats.scoreatpercentile(a, 50) + 49.5 + + """ + # adapted from NumPy's percentile function. When we require numpy >= 1.8, + # the implementation of this function can be replaced by np.percentile. + a = np.asarray(a) + if a.size == 0: + # empty array, return nan(s) with shape matching `per` + if np.isscalar(per): + return np.nan + else: + return np.full(np.asarray(per).shape, np.nan, dtype=np.float64) + + if limit: + a = a[(limit[0] <= a) & (a <= limit[1])] + + sorted_ = np.sort(a, axis=axis) + if axis is None: + axis = 0 + + return _compute_qth_percentile(sorted_, per, interpolation_method, axis) + + +# handle sequence of per's without calling sort multiple times +def _compute_qth_percentile(sorted_, per, interpolation_method, axis): + if not np.isscalar(per): + score = [_compute_qth_percentile(sorted_, i, + interpolation_method, axis) + for i in per] + return np.array(score) + + if not (0 <= per <= 100): + raise ValueError("percentile must be in the range [0, 100]") + + indexer = [slice(None)] * sorted_.ndim + idx = per / 100. * (sorted_.shape[axis] - 1) + + if int(idx) != idx: + # round fractional indices according to interpolation method + if interpolation_method == 'lower': + idx = int(np.floor(idx)) + elif interpolation_method == 'higher': + idx = int(np.ceil(idx)) + elif interpolation_method == 'fraction': + pass # keep idx as fraction and interpolate + else: + raise ValueError("interpolation_method can only be 'fraction', " + "'lower' or 'higher'") + + i = int(idx) + if i == idx: + indexer[axis] = slice(i, i + 1) + weights = array(1) + sumval = 1.0 + else: + indexer[axis] = slice(i, i + 2) + j = i + 1 + weights = array([(j - idx), (idx - i)], float) + wshape = [1] * sorted_.ndim + wshape[axis] = 2 + weights = weights.reshape(wshape) + sumval = weights.sum() + + # Use np.add.reduce (== np.sum but a little faster) to coerce data type + return np.add.reduce(sorted_[tuple(indexer)] * weights, axis=axis) / sumval + + +@xp_capabilities(np_only=True) +def percentileofscore(a, score, kind='rank', nan_policy='propagate'): + """Compute the percentile rank of a score relative to a list of scores. + + A `percentileofscore` of, for example, 80% means that 80% of the + scores in `a` are below the given score. In the case of gaps or + ties, the exact definition depends on the optional keyword, `kind`. + + Parameters + ---------- + a : array_like + A 1-D array to which `score` is compared. + score : float or array_like + A float score or array of scores for which to compute the percentile(s). + kind : {'rank', 'weak', 'strict', 'mean'}, optional + Specifies the interpretation of the resulting score. + The following options are available (default is 'rank'): + + * 'rank': Average percentage ranking of score. In case of multiple + matches, average the percentage rankings of all matching scores. + * 'weak': This kind corresponds to the definition of a cumulative + distribution function. A percentileofscore of 80% means that 80% + of values are less than or equal to the provided score. + * 'strict': Similar to "weak", except that only values that are + strictly less than the given score are counted. + * 'mean': The average of the "weak" and "strict" scores, often used + in testing. See https://en.wikipedia.org/wiki/Percentile_rank + + nan_policy : {'propagate', 'raise', 'omit'}, optional + Specifies how to treat `nan` values in `a`. + The following options are available (default is 'propagate'): + + * 'propagate': returns nan (for each value in `score`). + * 'raise': throws an error + * 'omit': performs the calculations ignoring nan values + + Returns + ------- + pcos : float or array-like + Percentile-position(s) of `score` (0-100) relative to `a`. + + See Also + -------- + numpy.percentile + scipy.stats.scoreatpercentile, scipy.stats.rankdata + + Examples + -------- + Three-quarters of the given values lie below a given score: + + >>> import numpy as np + >>> from scipy import stats + >>> stats.percentileofscore([1, 2, 3, 4], 3) + 75.0 + + With multiple matches, note how the scores of the two matches, 0.6 + and 0.8 respectively, are averaged: + + >>> stats.percentileofscore([1, 2, 3, 3, 4], 3) + 70.0 + + Only 2/5 values are strictly less than 3: + + >>> stats.percentileofscore([1, 2, 3, 3, 4], 3, kind='strict') + 40.0 + + But 4/5 values are less than or equal to 3: + + >>> stats.percentileofscore([1, 2, 3, 3, 4], 3, kind='weak') + 80.0 + + The average between the weak and the strict scores is: + + >>> stats.percentileofscore([1, 2, 3, 3, 4], 3, kind='mean') + 60.0 + + Score arrays (of any dimensionality) are supported: + + >>> stats.percentileofscore([1, 2, 3, 3, 4], [2, 3]) + array([40., 70.]) + + The inputs can be infinite: + + >>> stats.percentileofscore([-np.inf, 0, 1, np.inf], [1, 2, np.inf]) + array([75., 75., 100.]) + + If `a` is empty, then the resulting percentiles are all `nan`: + + >>> stats.percentileofscore([], [1, 2]) + array([nan, nan]) + """ + + a = np.asarray(a) + score = np.asarray(score) + + if a.ndim != 1: + raise ValueError("`a` must be 1-dimensional.") + + n = len(a) + + # Nan treatment + cna = _contains_nan(a, nan_policy) + cns = _contains_nan(score, nan_policy) + + if cns: + # If a score is nan, then the output should be nan + # (also if nan_policy is "omit", because it only applies to `a`) + score = ma.masked_where(np.isnan(score), score) + + if cna: + if nan_policy == "omit": + # Don't count nans + a = ma.masked_where(np.isnan(a), a) + n = a.count() + + if nan_policy == "propagate": + # All outputs should be nans + n = 0 + + # Cannot compare to empty list ==> nan + if n == 0: + perct = np.full_like(score, np.nan, dtype=np.float64) + + else: + # Prepare broadcasting + score = score[..., None] + + def count(x): + return np.count_nonzero(x, -1) + + # Main computations/logic + if kind == 'rank': + left = count(a < score) + right = count(a <= score) + plus1 = left < right + perct = (left + right + plus1) * (50.0 / n) + elif kind == 'strict': + perct = count(a < score) * (100.0 / n) + elif kind == 'weak': + perct = count(a <= score) * (100.0 / n) + elif kind == 'mean': + left = count(a < score) + right = count(a <= score) + perct = (left + right) * (50.0 / n) + else: + raise ValueError( + "kind can only be 'rank', 'strict', 'weak' or 'mean'") + + # Re-insert nan values + perct = ma.filled(perct, np.nan) + + if perct.ndim == 0: + return perct[()] + return perct + + +HistogramResult = namedtuple('HistogramResult', + ('count', 'lowerlimit', 'binsize', 'extrapoints')) + + +def _histogram(a, numbins=10, defaultlimits=None, weights=None, *, + density=False, cumulative=False): + """Create a histogram. + + Separate the range into several bins and return the number of instances + in each bin. + + Parameters + ---------- + a : array_like + Array of scores which will be put into bins. + numbins : int, optional + The number of bins to use for the histogram. Default is 10. + defaultlimits : tuple (lower, upper), optional + The lower and upper values for the range of the histogram. + If no value is given, a range slightly larger than the range of the + values in a is used. Specifically ``(a.min() - s, a.max() + s)``, + where ``s = (1/2)(a.max() - a.min()) / (numbins - 1)``. + weights : array_like, optional + The weights for each value in `a`. Default is None, which gives each + value a weight of 1.0 + + Returns + ------- + count : ndarray + Number of points (or sum of weights) in each bin. + lowerlimit : float + Lowest value of histogram, the lower limit of the first bin. + binsize : float + The size of the bins (all bins have the same size). + extrapoints : int + The number of points outside the range of the histogram. + + See Also + -------- + numpy.histogram + + Notes + ----- + This histogram is based on numpy's histogram but has a larger range by + default if default limits is not set. + + """ + xp = array_namespace(a) + a = xp_ravel(a) + a, weights = xp_promote(a, weights, force_floating=True, xp=xp) + + if defaultlimits is None: + if xp_size(a) == 0: + # handle empty arrays. Undetermined range, so use 0-1. + defaultlimits = xp.asarray(0., dtype=a.dtype), xp.asarray(1., dtype=a.dtype) + else: + # no range given, so use values in `a` + data_min = xp.min(a) + data_max = xp.max(a) + # Have bins extend past min and max values slightly + s = (data_max - data_min) / (2. * (numbins - 1.)) + defaultlimits = (data_min - s, data_max + s) + else: + if not (np.iterable(defaultlimits) and len(defaultlimits)==2 + and defaultlimits[0] < defaultlimits[1]): + message = ('If specified, `defaultreallimits` must be given as an iterable ' + 'in the order (lower limit, upper limit).') + raise ValueError(message) + if not (xp.isdtype(a.dtype, 'real floating') + and (weights is None or xp.isdtype(weights.dtype, 'real floating'))): + message = '`a` and (if specified) `weights` must have real dtype.' + raise ValueError(message) + if weights is not None and not is_lazy_array(weights) and xp.any(weights < 0): + message = 'All `weights` must be non-negative.' + raise ValueError(message) + + bin_edges = xp.linspace(*defaultlimits, numbins+1, dtype=a.dtype) + if weights is None: + indices = xp.searchsorted(xp.sort(a), bin_edges, side='left') + hist = xp.diff(indices) + n_right_limit = xp.count_nonzero(a == bin_edges[-1]) + hist = xpx.at(hist)[-1].add(xp.astype(n_right_limit, hist.dtype)) + else: + i = xp.argsort(a) + weights, a = weights[i], a[i] + cumulative_weights = xp.cumulative_sum(weights, include_initial=True) + indices = xp.searchsorted(a, bin_edges, side='left') + hist = xp.diff(xp.take_along_axis(cumulative_weights, indices, axis=-1)) + hist = xpx.at(hist)[-1].add(xp.sum(xp.where(a == bin_edges[-1], weights, 0.))) + + # fixed width for bins is assumed, as numpy's histogram gives + # fixed width bins for int values for 'bins' + binsize = bin_edges[1] - bin_edges[0] + # calculate number of extra points + binnedpoints = (xp.sum(hist) if weights is None + else xp.count_nonzero((bin_edges[0] <= a) & (a <= bin_edges[-1]))) + extrapoints = a.shape[0] - binnedpoints + + lowerlimit = xp.asarray(defaultlimits[0], dtype=a.dtype)[()] + + hist = xp.asarray(hist, dtype=a.dtype) + if density: + hist = hist / (a.shape[0] if weights is None else cumulative_weights[-1]) + if cumulative: + hist = xp.cumulative_sum(hist, axis=0) + + return HistogramResult(hist, lowerlimit, binsize, extrapoints) + + +CumfreqResult = namedtuple('CumfreqResult', + ('cumcount', 'lowerlimit', 'binsize', + 'extrapoints')) + + +@xp_capabilities(skip_backends=[('dask.array', 'fails tests')]) +def cumfreq(a, numbins=10, defaultreallimits=None, weights=None): + """Return a cumulative frequency histogram, using the histogram function. + + A cumulative histogram is a mapping that counts the cumulative number of + observations in all of the bins up to the specified bin. + + Parameters + ---------- + a : array_like + Input array. + numbins : int, optional + The number of bins to use for the histogram. Default is 10. + defaultreallimits : tuple (lower, upper), optional + The lower and upper values for the range of the histogram. + If no value is given, a range slightly larger than the range of the + values in `a` is used. Specifically ``(a.min() - s, a.max() + s)``, + where ``s = (1/2)(a.max() - a.min()) / (numbins - 1)``. + weights : array_like, optional + The weights for each value in `a`. Default is None, which gives each + value a weight of 1.0 + + Returns + ------- + cumcount : ndarray + Binned values of cumulative frequency. + lowerlimit : float + Lower real limit + binsize : float + Width of each bin. + extrapoints : int + Extra points. + + Examples + -------- + >>> import numpy as np + >>> import matplotlib.pyplot as plt + >>> from scipy import stats + >>> rng = np.random.default_rng() + >>> x = [1, 4, 2, 1, 3, 1] + >>> res = stats.cumfreq(x, numbins=4, defaultreallimits=(1.5, 5)) + >>> res.cumcount + array([ 1., 2., 3., 3.]) + >>> res.extrapoints + 3 + + Create a normal distribution with 1000 random values + + >>> samples = stats.norm.rvs(size=1000, random_state=rng) + + Calculate cumulative frequencies + + >>> res = stats.cumfreq(samples, numbins=25) + + Calculate space of values for x + + >>> x = res.lowerlimit + np.linspace(0, res.binsize*res.cumcount.size, + ... res.cumcount.size + 1) + + Plot histogram and cumulative histogram + + >>> fig = plt.figure(figsize=(10, 4)) + >>> ax1 = fig.add_subplot(1, 2, 1) + >>> ax2 = fig.add_subplot(1, 2, 2) + >>> ax1.hist(samples, bins=25) + >>> ax1.set_title('Histogram') + >>> ax2.bar(x[:-1], res.cumcount, width=res.binsize, align='edge') + >>> ax2.set_title('Cumulative histogram') + >>> ax2.set_xlim([x.min(), x.max()]) + + >>> plt.show() + + """ + h, l, b, e = _histogram(a, numbins, defaultreallimits, weights=weights, + cumulative=True) + return CumfreqResult(h, l, b, e) + + +RelfreqResult = namedtuple('RelfreqResult', + ('frequency', 'lowerlimit', 'binsize', + 'extrapoints')) + + +@xp_capabilities(skip_backends=[('dask.array', 'fails tests')]) +def relfreq(a, numbins=10, defaultreallimits=None, weights=None): + """Return a relative frequency histogram, using the histogram function. + + A relative frequency histogram is a mapping of the number of + observations in each of the bins relative to the total of observations. + + Parameters + ---------- + a : array_like + Input array. + numbins : int, optional + The number of bins to use for the histogram. Default is 10. + defaultreallimits : tuple (lower, upper), optional + The lower and upper values for the range of the histogram. + If no value is given, a range slightly larger than the range of the + values in a is used. Specifically ``(a.min() - s, a.max() + s)``, + where ``s = (1/2)(a.max() - a.min()) / (numbins - 1)``. + weights : array_like, optional + The weights for each value in `a`. Default is None, which gives each + value a weight of 1.0 + + Returns + ------- + frequency : ndarray + Binned values of relative frequency. + lowerlimit : float + Lower real limit. + binsize : float + Width of each bin. + extrapoints : int + Extra points. + + Examples + -------- + >>> import numpy as np + >>> import matplotlib.pyplot as plt + >>> from scipy import stats + >>> rng = np.random.default_rng() + >>> a = np.array([2, 4, 1, 2, 3, 2]) + >>> res = stats.relfreq(a, numbins=4) + >>> res.frequency + array([ 0.16666667, 0.5 , 0.16666667, 0.16666667]) + >>> np.sum(res.frequency) # relative frequencies should add up to 1 + 1.0 + + Create a normal distribution with 1000 random values + + >>> samples = stats.norm.rvs(size=1000, random_state=rng) + + Calculate relative frequencies + + >>> res = stats.relfreq(samples, numbins=25) + + Calculate space of values for x + + >>> x = res.lowerlimit + np.linspace(0, res.binsize*res.frequency.size, + ... res.frequency.size) + + Plot relative frequency histogram + + >>> fig = plt.figure(figsize=(5, 4)) + >>> ax = fig.add_subplot(1, 1, 1) + >>> ax.bar(x, res.frequency, width=res.binsize) + >>> ax.set_title('Relative frequency histogram') + >>> ax.set_xlim([x.min(), x.max()]) + + >>> plt.show() + + """ + h, l, b, e = _histogram(a, numbins, defaultreallimits, weights=weights, + density=True) + return RelfreqResult(h, l, b, e) + + +##################################### +# VARIABILITY FUNCTIONS # +##################################### + +@xp_capabilities() +def obrientransform(*samples, nan_policy='propagate'): + """Compute the O'Brien transform on input data (any number of arrays). + + Used to test for homogeneity of variance prior to running one-way stats. + Each array in ``*samples`` is one level of a factor. + Significant results of `f_oneway` on the transformed data suggest that the + variances of the underlying distributions are unequal. + See Maxwell and Delaney [1]_, p.112. + + Parameters + ---------- + *samples : array_like + Any number of arrays. + nan_policy : {'propagate', 'omit', 'raise'} + Defines how to handle input NaNs. + + - ``propagate``: if a NaN is present in a sample, all elements of the + transformed sample will be NaN. + - ``omit``: NaNs will be omitted when computing reducing statistics for the + transform, but NaNs in the sample will remain NaNs in the transformed sample. + - ``raise``: if a NaN is present, a ``ValueError`` will be raised. + + Returns + ------- + obrientransform : tuple of arrays + Transformed arrays for use in ANOVA. + + Raises + ------ + ValueError + If the mean of the transformed data is not equal to the original + variance, indicating a lack of convergence in the O'Brien transform. + + References + ---------- + .. [1] S. E. Maxwell and H. D. Delaney, "Designing Experiments and + Analyzing Data: A Model Comparison Perspective", Wadsworth, 1990. + + Examples + -------- + We'll test the following data sets for differences in their variance. + + >>> x = [10, 11, 13, 9, 7, 12, 12, 9, 10] + >>> y = [13, 21, 5, 10, 8, 14, 10, 12, 7, 15] + + Apply the O'Brien transform to the data. + + >>> from scipy.stats import obrientransform + >>> tx, ty = obrientransform(x, y) + + Use `scipy.stats.f_oneway` to apply a one-way ANOVA test to the + transformed data. + + >>> from scipy.stats import f_oneway + >>> F, p = f_oneway(tx, ty) + >>> p + 0.1314139477040335 + + If we require that ``p < 0.05`` for significance, we cannot conclude + that the variances are different. + + """ + xp = array_namespace(*samples) + n_samples = len(samples) + samples = xp_promote(*samples, force_floating=True, xp=xp) + samples = (samples,) if n_samples == 1 else samples + return tuple(_xp_obrientransform_one_sample(sample, xp=xp, nan_policy=nan_policy) + for sample in samples) + + +def _xp_obrientransform_one_sample(a, *, xp, nan_policy): + _contains_nan(a, nan_policy, xp_omit_okay=True) # handle `nan_policy='raise'` + n = xp.asarray(xp.count_nonzero(~xp.isnan(a)), dtype=a.dtype) + mu = _xp_mean(a, nan_policy=nan_policy) + sq = (a - mu)**2 + sumsq = _xp_mean(sq, nan_policy=nan_policy) * n + return ((n - 1.5) * n * sq - 0.5 * sumsq) / ((n - 1) * (n - 2)) + + +@xp_capabilities(marray=True) +@_axis_nan_policy_factory( + lambda x: x, result_to_tuple=lambda x, _: (x,), n_outputs=1, too_small=1 +) +def sem(a, axis=0, ddof=1, nan_policy='propagate'): + """Compute standard error of the mean. + + Calculate the standard error of the mean (or standard error of + measurement) of the values in the input array. + + Parameters + ---------- + a : array_like + An array containing the values for which the standard error is + returned. Must contain at least two observations. + axis : int or None, optional + Axis along which to operate. Default is 0. If None, compute over + the whole array `a`. + ddof : int, optional + Delta degrees-of-freedom. How many degrees of freedom to adjust + for bias in limited samples relative to the population estimate + of variance. Defaults to 1. + nan_policy : {'propagate', 'raise', 'omit'}, optional + Defines how to handle when input contains nan. + The following options are available (default is 'propagate'): + + * 'propagate': returns nan + * 'raise': throws an error + * 'omit': performs the calculations ignoring nan values + + Returns + ------- + s : ndarray or float + The standard error of the mean in the sample(s), along the input axis. + + Notes + ----- + The default value for `ddof` is different to the default (0) used by other + ddof containing routines, such as np.std and np.nanstd. + + Examples + -------- + Find standard error along the first axis: + + >>> import numpy as np + >>> from scipy import stats + >>> a = np.arange(20).reshape(5,4) + >>> stats.sem(a) + array([ 2.8284, 2.8284, 2.8284, 2.8284]) + + Find standard error across the whole array, using n degrees of freedom: + + >>> stats.sem(a, axis=None, ddof=0) + 1.2893796958227628 + + """ + xp = array_namespace(a) + if axis is None: + a = xp.reshape(a, (-1,)) + axis = 0 + a = xpx.atleast_nd(xp.asarray(a), ndim=1, xp=xp) + n = _count_nonmasked(a, axis, xp=xp) + s = xp.std(a, axis=axis, correction=ddof) / n**0.5 + return s + + +def _isconst(x): + """ + Check if all values in x are the same. nans are ignored. + + x must be a 1d array. + + The return value is a 1d array with length 1, so it can be used + in np.apply_along_axis. + """ + y = x[~np.isnan(x)] + if y.size == 0: + return np.array([True]) + else: + return (y[0] == y).all(keepdims=True) + + +@xp_capabilities(marray=True) +def zscore(a, axis=0, ddof=0, nan_policy='propagate'): + """ + Compute the z score. + + Compute the z score of each value in the sample, relative to the + sample mean and standard deviation. + + Parameters + ---------- + a : array_like + An array like object containing the sample data. + axis : int or None, optional + Axis along which to operate. Default is 0. If None, compute over + the whole array `a`. + ddof : int, optional + Degrees of freedom correction in the calculation of the + standard deviation. Default is 0. + nan_policy : {'propagate', 'raise', 'omit'}, optional + Defines how to handle when input contains nan. 'propagate' returns nan, + 'raise' throws an error, 'omit' performs the calculations ignoring nan + values. Default is 'propagate'. Note that when the value is 'omit', + nans in the input also propagate to the output, but they do not affect + the z-scores computed for the non-nan values. + + Returns + ------- + zscore : array_like + The z-scores, standardized by mean and standard deviation of + input array `a`. + + See Also + -------- + numpy.mean : Arithmetic average + numpy.std : Arithmetic standard deviation + scipy.stats.gzscore : Geometric standard score + + Notes + ----- + This function preserves ndarray subclasses, and works also with + matrices and masked arrays (it uses `asanyarray` instead of + `asarray` for parameters). + + References + ---------- + .. [1] "Standard score", *Wikipedia*, + https://en.wikipedia.org/wiki/Standard_score. + .. [2] Huck, S. W., Cross, T. L., Clark, S. B, "Overcoming misconceptions + about Z-scores", Teaching Statistics, vol. 8, pp. 38-40, 1986 + + Examples + -------- + >>> import numpy as np + >>> a = np.array([ 0.7972, 0.0767, 0.4383, 0.7866, 0.8091, + ... 0.1954, 0.6307, 0.6599, 0.1065, 0.0508]) + >>> from scipy import stats + >>> stats.zscore(a) + array([ 1.1273, -1.247 , -0.0552, 1.0923, 1.1664, -0.8559, 0.5786, + 0.6748, -1.1488, -1.3324]) + + Computing along a specified axis, using n-1 degrees of freedom + (``ddof=1``) to calculate the standard deviation: + + >>> b = np.array([[ 0.3148, 0.0478, 0.6243, 0.4608], + ... [ 0.7149, 0.0775, 0.6072, 0.9656], + ... [ 0.6341, 0.1403, 0.9759, 0.4064], + ... [ 0.5918, 0.6948, 0.904 , 0.3721], + ... [ 0.0921, 0.2481, 0.1188, 0.1366]]) + >>> stats.zscore(b, axis=1, ddof=1) + array([[-0.19264823, -1.28415119, 1.07259584, 0.40420358], + [ 0.33048416, -1.37380874, 0.04251374, 1.00081084], + [ 0.26796377, -1.12598418, 1.23283094, -0.37481053], + [-0.22095197, 0.24468594, 1.19042819, -1.21416216], + [-0.82780366, 1.4457416 , -0.43867764, -0.1792603 ]]) + + An example with ``nan_policy='omit'``: + + >>> x = np.array([[25.11, 30.10, np.nan, 32.02, 43.15], + ... [14.95, 16.06, 121.25, 94.35, 29.81]]) + >>> stats.zscore(x, axis=1, nan_policy='omit') + array([[-1.13490897, -0.37830299, nan, -0.08718406, 1.60039602], + [-0.91611681, -0.89090508, 1.4983032 , 0.88731639, -0.5785977 ]]) + """ + return zmap(a, a, axis=axis, ddof=ddof, nan_policy=nan_policy) + + +@xp_capabilities(marray=True) +def gzscore(a, *, axis=0, ddof=0, nan_policy='propagate'): + """ + Compute the geometric standard score. + + Compute the geometric z score of each strictly positive value in the + sample, relative to the geometric mean and standard deviation. + Mathematically the geometric z score can be evaluated as:: + + gzscore = log(a/gmu) / log(gsigma) + + where ``gmu`` (resp. ``gsigma``) is the geometric mean (resp. standard + deviation). + + Parameters + ---------- + a : array_like + Sample data. + axis : int or None, optional + Axis along which to operate. Default is 0. If None, compute over + the whole array `a`. + ddof : int, optional + Degrees of freedom correction in the calculation of the + standard deviation. Default is 0. + nan_policy : {'propagate', 'raise', 'omit'}, optional + Defines how to handle when input contains nan. 'propagate' returns nan, + 'raise' throws an error, 'omit' performs the calculations ignoring nan + values. Default is 'propagate'. Note that when the value is 'omit', + nans in the input also propagate to the output, but they do not affect + the geometric z scores computed for the non-nan values. + + Returns + ------- + gzscore : array_like + The geometric z scores, standardized by geometric mean and geometric + standard deviation of input array `a`. + + See Also + -------- + gmean : Geometric mean + gstd : Geometric standard deviation + zscore : Standard score + + Notes + ----- + This function preserves ndarray subclasses, and works also with + matrices and masked arrays (it uses ``asanyarray`` instead of + ``asarray`` for parameters). + + .. versionadded:: 1.8 + + References + ---------- + .. [1] "Geometric standard score", *Wikipedia*, + https://en.wikipedia.org/wiki/Geometric_standard_deviation#Geometric_standard_score. + + Examples + -------- + Draw samples from a log-normal distribution: + + >>> import numpy as np + >>> from scipy.stats import zscore, gzscore + >>> import matplotlib.pyplot as plt + + >>> rng = np.random.default_rng() + >>> mu, sigma = 3., 1. # mean and standard deviation + >>> x = rng.lognormal(mu, sigma, size=500) + + Display the histogram of the samples: + + >>> fig, ax = plt.subplots() + >>> ax.hist(x, 50) + >>> plt.show() + + Display the histogram of the samples standardized by the classical zscore. + Distribution is rescaled but its shape is unchanged. + + >>> fig, ax = plt.subplots() + >>> ax.hist(zscore(x), 50) + >>> plt.show() + + Demonstrate that the distribution of geometric zscores is rescaled and + quasinormal: + + >>> fig, ax = plt.subplots() + >>> ax.hist(gzscore(x), 50) + >>> plt.show() + + """ + xp = array_namespace(a) + a = xp_promote(a, force_floating=True, xp=xp) + log = ma.log if isinstance(a, ma.MaskedArray) else xp.log + return zscore(log(a), axis=axis, ddof=ddof, nan_policy=nan_policy) + + +@xp_capabilities(marray=True) +def zmap(scores, compare, axis=0, ddof=0, nan_policy='propagate'): + """ + Calculate the relative z-scores. + + Return an array of z-scores, i.e., scores that are standardized to + zero mean and unit variance, where mean and variance are calculated + from the comparison array. + + Parameters + ---------- + scores : array_like + The input for which z-scores are calculated. + compare : array_like + The input from which the mean and standard deviation of the + normalization are taken; assumed to have the same dimension as + `scores`. + axis : int or None, optional + Axis over which mean and variance of `compare` are calculated. + Default is 0. If None, compute over the whole array `scores`. + ddof : int, optional + Degrees of freedom correction in the calculation of the + standard deviation. Default is 0. + nan_policy : {'propagate', 'raise', 'omit'}, optional + Defines how to handle the occurrence of nans in `compare`. + 'propagate' returns nan, 'raise' raises an exception, 'omit' + performs the calculations ignoring nan values. Default is + 'propagate'. Note that when the value is 'omit', nans in `scores` + also propagate to the output, but they do not affect the z-scores + computed for the non-nan values. + + Returns + ------- + zscore : array_like + Z-scores, in the same shape as `scores`. + + Notes + ----- + This function preserves ndarray subclasses, and works also with + matrices and masked arrays (it uses `asanyarray` instead of + `asarray` for parameters). + + Examples + -------- + >>> from scipy.stats import zmap + >>> a = [0.5, 2.0, 2.5, 3] + >>> b = [0, 1, 2, 3, 4] + >>> zmap(a, b) + array([-1.06066017, 0. , 0.35355339, 0.70710678]) + + """ + # The docstring explicitly states that it preserves subclasses. + # Let's table deprecating that and just get the array API version + # working. + + like_zscore = (scores is compare) + xp = array_namespace(scores, compare) + scores, compare = xp_promote(scores, compare, force_floating=True, xp=xp) + + with warnings.catch_warnings(): + if like_zscore: # zscore should not emit SmallSampleWarning + warnings.simplefilter('ignore', SmallSampleWarning) + + mn = _xp_mean(compare, axis=axis, keepdims=True, nan_policy=nan_policy) + std = _xp_var(compare, axis=axis, correction=ddof, + keepdims=True, nan_policy=nan_policy)**0.5 + + with np.errstate(invalid='ignore', divide='ignore'): + z = _demean(scores, mn, axis, xp=xp, precision_warning=False) / std + + # If we know that scores and compare are identical, we can infer that + # some slices should have NaNs. + if like_zscore: + eps = xp.finfo(z.dtype).eps + zero = std <= xp.abs(eps * mn) + zero = xp.broadcast_to(zero, z.shape) + z = xpx.at(z, zero).set(xp.nan) + + return z + + +@xp_capabilities(marray=True) +def gstd(a, axis=0, ddof=1, *, keepdims=False, nan_policy='propagate'): + r""" + Calculate the geometric standard deviation of an array. + + The geometric standard deviation describes the spread of a set of numbers + where the geometric mean is preferred. It is a multiplicative factor, and + so a dimensionless quantity. + + It is defined as the exponential of the standard deviation of the + natural logarithms of the observations. + + Parameters + ---------- + a : array_like + An array containing finite, strictly positive, real numbers. + axis : int, tuple or None, optional + Axis along which to operate. Default is 0. If None, compute over + the whole array `a`. + ddof : int, optional + Degree of freedom correction in the calculation of the + geometric standard deviation. Default is 1. + keepdims : bool, optional + If this is set to ``True``, the axes which are reduced are left + in the result as dimensions with length one. With this option, + the result will broadcast correctly against the input array. + nan_policy : {'propagate', 'omit', 'raise'}, default: 'propagate' + Defines how to handle input NaNs. + + - ``propagate``: if a NaN is present in the axis slice (e.g. row) along + which the statistic is computed, the corresponding entry of the output + will be NaN. + - ``omit``: NaNs will be omitted when performing the calculation. + If insufficient data remains in the axis slice along which the + statistic is computed, the corresponding entry of the output will be + NaN. + - ``raise``: if a NaN is present, a ``ValueError`` will be raised. + + Returns + ------- + gstd : ndarray or float + An array of the geometric standard deviation. If `axis` is None or `a` + is a 1d array a float is returned. + + See Also + -------- + gmean : Geometric mean + numpy.std : Standard deviation + gzscore : Geometric standard score + + Notes + ----- + Mathematically, the sample geometric standard deviation :math:`s_G` can be + defined in terms of the natural logarithms of the observations + :math:`y_i = \log(x_i)`: + + .. math:: + + s_G = \exp(s), \quad s = \sqrt{\frac{1}{n - d} \sum_{i=1}^n (y_i - \bar y)^2} + + where :math:`n` is the number of observations, :math:`d` is the adjustment `ddof` + to the degrees of freedom, and :math:`\bar y` denotes the mean of the natural + logarithms of the observations. Note that the default ``ddof=1`` is different from + the default value used by similar functions, such as `numpy.std` and `numpy.var`. + + When an observation is infinite, the geometric standard deviation is + NaN (undefined). Non-positive observations will also produce NaNs in the + output because the *natural* logarithm (as opposed to the *complex* + logarithm) is defined and finite only for positive reals. + The geometric standard deviation is sometimes confused with the exponential + of the standard deviation, ``exp(std(a))``. Instead, the geometric standard + deviation is ``exp(std(log(a)))``. + + References + ---------- + .. [1] "Geometric standard deviation", *Wikipedia*, + https://en.wikipedia.org/wiki/Geometric_standard_deviation. + .. [2] Kirkwood, T. B., "Geometric means and measures of dispersion", + Biometrics, vol. 35, pp. 908-909, 1979 + + Examples + -------- + Find the geometric standard deviation of a log-normally distributed sample. + Note that the standard deviation of the distribution is one; on a + log scale this evaluates to approximately ``exp(1)``. + + >>> import numpy as np + >>> from scipy.stats import gstd + >>> rng = np.random.default_rng() + >>> sample = rng.lognormal(mean=0, sigma=1, size=1000) + >>> gstd(sample) + 2.810010162475324 + + Compute the geometric standard deviation of a multidimensional array and + of a given axis. + + >>> a = np.arange(1, 25).reshape(2, 3, 4) + >>> gstd(a, axis=None) + 2.2944076136018947 + >>> gstd(a, axis=2) + array([[1.82424757, 1.22436866, 1.13183117], + [1.09348306, 1.07244798, 1.05914985]]) + >>> gstd(a, axis=(1,2)) + array([2.12939215, 1.22120169]) + + """ + xp = array_namespace(a) + a = xp_promote(a, force_floating=True, xp=xp) + + kwargs = dict(axis=axis, correction=ddof, keepdims=keepdims, nan_policy=nan_policy) + with np.errstate(invalid='ignore', divide='ignore'): + res = xp.exp(_xp_var(xp.log(a), **kwargs)**0.5) + + if not is_lazy_array(a) and xp.any(a <= 0): + message = ("The geometric standard deviation is only defined if all elements " + "are greater than or equal to zero; otherwise, the result is NaN.") + warnings.warn(message, RuntimeWarning, stacklevel=2) + + return res + +# Private dictionary initialized only once at module level +# See https://en.wikipedia.org/wiki/Robust_measures_of_scale +_scale_conversions = {'normal': float(special.erfinv(0.5) * 2.0 * math.sqrt(2.0))} + + +@xp_capabilities(skip_backends=[('dask.array', 'no quantile (take_along_axis)')], + marray=True) +@_axis_nan_policy_factory( + lambda x: x, result_to_tuple=lambda x, _: (x,), n_outputs=1, + default_axis=None, override={'nan_propagation': False} +) +def iqr(x, axis=None, rng=(25, 75), scale=1.0, nan_policy='propagate', + interpolation='linear', keepdims=False): + r""" + Compute the interquartile range of the data along the specified axis. + + The interquartile range (IQR) is the difference between the 75th and + 25th percentile of the data. It is a measure of the dispersion + similar to standard deviation or variance, but is much more robust + against outliers [2]_. + + The ``rng`` parameter allows this function to compute other + percentile ranges than the actual IQR. For example, setting + ``rng=(0, 100)`` is equivalent to `numpy.ptp`. + + The IQR of an empty array is `np.nan`. + + .. versionadded:: 0.18.0 + + Parameters + ---------- + x : array_like + Input array or object that can be converted to an array. + axis : int or sequence of int, optional + Axis along which the range is computed. The default is to + compute the IQR for the entire array. + rng : Two-element sequence containing floats in range of [0,100] optional + Percentiles over which to compute the range. Each must be + between 0 and 100, inclusive. The default is the true IQR: + ``(25, 75)``. The order of the elements is not important. + scale : scalar or str or array_like of reals, optional + The numerical value of scale will be divided out of the final + result. The following string value is also recognized: + + * 'normal' : Scale by + :math:`2 \sqrt{2} erf^{-1}(\frac{1}{2}) \approx 1.349`. + + The default is 1.0. + Array-like `scale` of real dtype is also allowed, as long + as it broadcasts correctly to the output such that + ``out / scale`` is a valid operation. The output dimensions + depend on the input array, `x`, the `axis` argument, and the + `keepdims` flag. + nan_policy : {'propagate', 'raise', 'omit'}, optional + Defines how to handle when input contains nan. + The following options are available (default is 'propagate'): + + * 'propagate': returns nan + * 'raise': throws an error + * 'omit': performs the calculations ignoring nan values + + interpolation : str, optional + + Specifies the interpolation method to use when the percentile + boundaries lie between two data points ``i`` and ``j``. + The following options are available (default is 'linear'): + + * 'linear': ``i + (j - i)*fraction``, where ``fraction`` is the + fractional part of the index surrounded by ``i`` and ``j``. + * 'lower': ``i``. + * 'higher': ``j``. + * 'nearest': ``i`` or ``j`` whichever is nearest. + * 'midpoint': ``(i + j)/2``. + + For NumPy >= 1.22.0, the additional options provided by the ``method`` + keyword of `numpy.percentile` are also valid. + + keepdims : bool, optional + If this is set to True, the reduced axes are left in the + result as dimensions with size one. With this option, the result + will broadcast correctly against the original array `x`. + + Returns + ------- + iqr : scalar or ndarray + If ``axis=None``, a scalar is returned. If the input contains + integers or floats of smaller precision than ``np.float64``, then the + output data-type is ``np.float64``. Otherwise, the output data-type is + the same as that of the input. + + See Also + -------- + numpy.std, numpy.var + + References + ---------- + .. [1] "Interquartile range" https://en.wikipedia.org/wiki/Interquartile_range + .. [2] "Robust measures of scale" https://en.wikipedia.org/wiki/Robust_measures_of_scale + .. [3] "Quantile" https://en.wikipedia.org/wiki/Quantile + + Examples + -------- + >>> import numpy as np + >>> from scipy.stats import iqr + >>> x = np.array([[10, 7, 4], [3, 2, 1]]) + >>> x + array([[10, 7, 4], + [ 3, 2, 1]]) + >>> iqr(x) + 4.0 + >>> iqr(x, axis=0) + array([ 3.5, 2.5, 1.5]) + >>> iqr(x, axis=1) + array([ 3., 1.]) + >>> iqr(x, axis=1, keepdims=True) + array([[ 3.], + [ 1.]]) + + """ + xp = array_namespace(x) + + # An error may be raised here, so fail-fast, before doing lengthy + # computations, even though `scale` is not used until later + if isinstance(scale, str): + scale_key = scale.lower() + if scale_key not in _scale_conversions: + raise ValueError(f"{scale} not a valid scale for `iqr`") + scale = _scale_conversions[scale_key] + + if len(rng) != 2: + raise TypeError("`rng` must be a two element sequence.") + + if np.isnan(rng).any(): # OK to use NumPy; this shouldn't be an array + raise ValueError("`rng` must not contain NaNs.") + + rng = (rng[0]/100, rng[1]/100) if rng[0] < rng[1] else (rng[1]/100, rng[0]/100) + + if rng[0] < 0 or rng[1] > 1: + raise ValueError("Elements of `rng` must be in the range [0, 100].") + + if interpolation in {'lower', 'midpoint', 'higher', 'nearest'}: + interpolation = '_' + interpolation + + rng = xp.asarray(rng, dtype=xp_result_type(x, force_floating=True, xp=xp)) + pct = stats.quantile(x, rng, axis=-1, method=interpolation, keepdims=True) + out = pct[..., 1:2] - pct[..., 0:1] + + if scale != 1.0: + out /= scale + + out = out if keepdims else xp.squeeze(out, axis=-1) + return out[()] if out.ndim == 0 else out + + +def _mad_1d(x, center, nan_policy): + # Median absolute deviation for 1-d array x. + # This is a helper function for `median_abs_deviation`; it assumes its + # arguments have been validated already. In particular, x must be a + # 1-d numpy array, center must be callable, and if nan_policy is not + # 'propagate', it is assumed to be 'omit', because 'raise' is handled + # in `median_abs_deviation`. + # No warning is generated if x is empty or all nan. + isnan = np.isnan(x) + if isnan.any(): + if nan_policy == 'propagate': + return np.nan + x = x[~isnan] + if x.size == 0: + # MAD of an empty array is nan. + return np.nan + # Edge cases have been handled, so do the basic MAD calculation. + med = center(x) + mad = np.median(np.abs(x - med)) + return mad + + +@xp_capabilities(skip_backends=[('dask.array', 'not supported by `quantile`')], + marray=True) +@_axis_nan_policy_factory( + lambda x: x, result_to_tuple=lambda x, _: (x,), n_outputs=1, default_axis=0 +) +def median_abs_deviation(x, axis=0, center=None, scale=1.0, + nan_policy='propagate'): + r""" + Compute the median absolute deviation of the data along the given axis. + + The median absolute deviation (MAD, [1]_) computes the median over the + absolute deviations from the median. It is a measure of dispersion + similar to the standard deviation but more robust to outliers [2]_. + + The MAD of an empty array is ``np.nan``. + + .. versionadded:: 1.5.0 + + Parameters + ---------- + x : array_like + Input array or object that can be converted to an array. + axis : int or None, optional + Axis along which the range is computed. Default is 0. If None, compute + the MAD over the entire array. + center : callable, optional + A function that will return the central value. The default is to use + np.median. Any user defined function used will need to have the + function signature ``func(arr, axis)``. + scale : scalar or str, optional + The numerical value of scale will be divided out of the final + result. The default is 1.0. The string "normal" is also accepted, + and results in `scale` being the inverse of the standard normal + quantile function at 0.75, which is approximately 0.67449. + Array-like scale is also allowed, as long as it broadcasts correctly + to the output such that ``out / scale`` is a valid operation. The + output dimensions depend on the input array, `x`, and the `axis` + argument. + nan_policy : {'propagate', 'raise', 'omit'}, optional + Defines how to handle when input contains nan. + The following options are available (default is 'propagate'): + + * 'propagate': returns nan + * 'raise': throws an error + * 'omit': performs the calculations ignoring nan values + + Returns + ------- + mad : scalar or ndarray + If ``axis=None``, a scalar is returned. If the input contains + integers or floats of smaller precision than ``np.float64``, then the + output data-type is ``np.float64``. Otherwise, the output data-type is + the same as that of the input. + + See Also + -------- + numpy.std, numpy.var, numpy.median, scipy.stats.iqr, scipy.stats.tmean, + scipy.stats.tstd, scipy.stats.tvar + + Notes + ----- + The `center` argument only affects the calculation of the central value + around which the MAD is calculated. That is, passing in ``center=np.mean`` + will calculate the MAD around the mean - it will not calculate the *mean* + absolute deviation. + + The input array may contain `inf`, but if `center` returns `inf`, the + corresponding MAD for that data will be `nan`. + + References + ---------- + .. [1] "Median absolute deviation", + https://en.wikipedia.org/wiki/Median_absolute_deviation + .. [2] "Robust measures of scale", + https://en.wikipedia.org/wiki/Robust_measures_of_scale + + Examples + -------- + When comparing the behavior of `median_abs_deviation` with ``np.std``, + the latter is affected when we change a single value of an array to have an + outlier value while the MAD hardly changes: + + >>> import numpy as np + >>> from scipy import stats + >>> x = stats.norm.rvs(size=100, scale=1, random_state=123456) + >>> x.std() + 0.9973906394005013 + >>> stats.median_abs_deviation(x) + 0.82832610097857 + >>> x[0] = 345.6 + >>> x.std() + 34.42304872314415 + >>> stats.median_abs_deviation(x) + 0.8323442311590675 + + Axis handling example: + + >>> x = np.array([[10, 7, 4], [3, 2, 1]]) + >>> x + array([[10, 7, 4], + [ 3, 2, 1]]) + >>> stats.median_abs_deviation(x) + array([3.5, 2.5, 1.5]) + >>> stats.median_abs_deviation(x, axis=None) + 2.0 + + Scale normal example: + + >>> x = stats.norm.rvs(size=1000000, scale=2, random_state=123456) + >>> stats.median_abs_deviation(x) + 1.3487398527041636 + >>> stats.median_abs_deviation(x, scale='normal') + 1.9996446978061115 + + """ + xp = array_namespace(x) + xp_median = (xp.median if is_numpy(xp) + else lambda x, axis: stats.quantile(x, 0.5, axis=axis)) + center = (xp_median if center is None else center) + + if not callable(center): + raise TypeError("The argument 'center' must be callable. The given " + f"value {repr(center)} is not callable.") + + # An error may be raised here, so fail-fast, before doing lengthy + # computations, even though `scale` is not used until later + if isinstance(scale, str): + if scale.lower() == 'normal': + scale = 0.6744897501960817 # special.ndtri(0.75) + else: + raise ValueError(f"{scale} is not a valid scale value.") + + # Wrap the call to center() in expand_dims() so it acts like + # keepdims=True was used. + med = xp.expand_dims(center(x, axis=axis), axis=axis) + mad = xp_median(xp.abs(x - med), axis=axis) + + return mad / scale + + +##################################### +# TRIMMING FUNCTIONS # +##################################### + + +SigmaclipResult = namedtuple('SigmaclipResult', ('clipped', 'lower', 'upper')) + + +@xp_capabilities(skip_backends=[('dask.array', "doesn't know array size")], + jax_jit=False) +def sigmaclip(a, low=4., high=4., *, nan_policy='propagate'): + """Perform iterative sigma-clipping of array elements. + + Starting from the full sample, all elements outside the critical range are + removed, i.e. all elements of the input array `c` that satisfy either of + the following conditions:: + + c < mean(c) - std(c)*low + c > mean(c) + std(c)*high + + The iteration continues with the updated sample until no + elements are outside the (updated) range. + + Parameters + ---------- + a : array_like + Data array, will be raveled if not 1-D. + low : float, optional + Lower bound factor of sigma clipping. Default is 4. + high : float, optional + Upper bound factor of sigma clipping. Default is 4. + nan_policy : {'propagate', 'omit', 'raise'} + Defines how to handle input NaNs. + + - ``propagate``: if a NaN is present in the input, the clipped array will be + empty, and the upper and lower thresholds will be NaN. + - ``omit``: NaNs will be omitted when performing the calculation. + - ``raise``: if a NaN is present, a ``ValueError`` will be raised. + + Returns + ------- + clipped : ndarray + Input array with clipped elements removed. + lower : float + Lower threshold value use for clipping. + upper : float + Upper threshold value use for clipping. + + Notes + ----- + This function iteratively *removes* observations. Once observations are + removed, they are not re-added in subsequent iterations. Consequently, + although it is often the case that ``clipped`` is identical to + ``a[(a >= lower) & (a <= upper)]``, this property is not guaranteed to be + satisfied; ``clipped`` may have fewer elements. + + Examples + -------- + >>> import numpy as np + >>> from scipy.stats import sigmaclip + >>> a = np.concatenate((np.linspace(9.5, 10.5, 31), + ... np.linspace(0, 20, 5))) + >>> fact = 1.5 + >>> c, low, upp = sigmaclip(a, fact, fact) + >>> c + array([ 9.96666667, 10. , 10.03333333, 10. ]) + >>> c.var(), c.std() + (0.00055555555555555165, 0.023570226039551501) + >>> low, c.mean() - fact*c.std(), c.min() + (9.9646446609406727, 9.9646446609406727, 9.9666666666666668) + >>> upp, c.mean() + fact*c.std(), c.max() + (10.035355339059327, 10.035355339059327, 10.033333333333333) + + >>> a = np.concatenate((np.linspace(9.5, 10.5, 11), + ... np.linspace(-100, -50, 3))) + >>> c, low, upp = sigmaclip(a, 1.8, 1.8) + >>> (c == np.linspace(9.5, 10.5, 11)).all() + True + + """ + xp = array_namespace(a) + c = xp_ravel(xp.asarray(a)) + contains_nan = _contains_nan(c, nan_policy, xp_omit_okay=True) + if contains_nan: + if nan_policy == 'propagate': + NaN = _get_nan(c, xp=xp) + clipped = xp.empty_like(c[0:0]) + return SigmaclipResult(clipped, NaN, NaN) + elif nan_policy == 'omit': + c = c[~xp.isnan(c)] + + delta = 1 + while delta: + c_std = xp.std(c) + c_mean = xp.mean(c) + size = xp_size(c) + critlower = c_mean - c_std * low + critupper = c_mean + c_std * high + c = c[(c >= critlower) & (c <= critupper)] + delta = size - xp_size(c) + + return SigmaclipResult(c, critlower, critupper) + + +@xp_capabilities(np_only=True) +def trimboth(a, proportiontocut, axis=0): + """Slice off a proportion of items from both ends of an array. + + Slice off the passed proportion of items from both ends of the passed + array (i.e., with `proportiontocut` = 0.1, slices leftmost 10% **and** + rightmost 10% of scores). The trimmed values are the lowest and + highest ones. + Slice off less if proportion results in a non-integer slice index (i.e. + conservatively slices off `proportiontocut`). + + Parameters + ---------- + a : array_like + Data to trim. + proportiontocut : float + Proportion (in range 0-1) of total data set to trim of each end. + axis : int or None, optional + Axis along which to trim data. Default is 0. If None, compute over + the whole array `a`. + + Returns + ------- + out : ndarray + Trimmed version of array `a`. The order of the trimmed content + is undefined. + + See Also + -------- + trim_mean + + Examples + -------- + Create an array of 10 values and trim 10% of those values from each end: + + >>> import numpy as np + >>> from scipy import stats + >>> a = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] + >>> stats.trimboth(a, 0.1) + array([1, 3, 2, 4, 5, 6, 7, 8]) + + Note that the elements of the input array are trimmed by value, but the + output array is not necessarily sorted. + + The proportion to trim is rounded down to the nearest integer. For + instance, trimming 25% of the values from each end of an array of 10 + values will return an array of 6 values: + + >>> b = np.arange(10) + >>> stats.trimboth(b, 1/4).shape + (6,) + + Multidimensional arrays can be trimmed along any axis or across the entire + array: + + >>> c = [2, 4, 6, 8, 0, 1, 3, 5, 7, 9] + >>> d = np.array([a, b, c]) + >>> stats.trimboth(d, 0.4, axis=0).shape + (1, 10) + >>> stats.trimboth(d, 0.4, axis=1).shape + (3, 2) + >>> stats.trimboth(d, 0.4, axis=None).shape + (6,) + + """ + a = np.asarray(a) + + if a.size == 0: + return a + + if axis is None: + a = a.ravel() + axis = 0 + + nobs = a.shape[axis] + lowercut = int(proportiontocut * nobs) + uppercut = nobs - lowercut + if (lowercut >= uppercut): + raise ValueError("Proportion too big.") + + atmp = np.partition(a, (lowercut, uppercut - 1), axis) + + sl = [slice(None)] * atmp.ndim + sl[axis] = slice(lowercut, uppercut) + return atmp[tuple(sl)] + + +@xp_capabilities(np_only=True) +def trim1(a, proportiontocut, tail='right', axis=0): + """Slice off a proportion from ONE end of the passed array distribution. + + If `proportiontocut` = 0.1, slices off 'leftmost' or 'rightmost' + 10% of scores. The lowest or highest values are trimmed (depending on + the tail). + Slice off less if proportion results in a non-integer slice index + (i.e. conservatively slices off `proportiontocut` ). + + Parameters + ---------- + a : array_like + Input array. + proportiontocut : float + Fraction to cut off of 'left' or 'right' of distribution. + tail : {'left', 'right'}, optional + Defaults to 'right'. + axis : int or None, optional + Axis along which to trim data. Default is 0. If None, compute over + the whole array `a`. + + Returns + ------- + trim1 : ndarray + Trimmed version of array `a`. The order of the trimmed content is + undefined. + + Examples + -------- + Create an array of 10 values and trim 20% of its lowest values: + + >>> import numpy as np + >>> from scipy import stats + >>> a = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] + >>> stats.trim1(a, 0.2, 'left') + array([2, 4, 3, 5, 6, 7, 8, 9]) + + Note that the elements of the input array are trimmed by value, but the + output array is not necessarily sorted. + + The proportion to trim is rounded down to the nearest integer. For + instance, trimming 25% of the values from an array of 10 values will + return an array of 8 values: + + >>> b = np.arange(10) + >>> stats.trim1(b, 1/4).shape + (8,) + + Multidimensional arrays can be trimmed along any axis or across the entire + array: + + >>> c = [2, 4, 6, 8, 0, 1, 3, 5, 7, 9] + >>> d = np.array([a, b, c]) + >>> stats.trim1(d, 0.8, axis=0).shape + (1, 10) + >>> stats.trim1(d, 0.8, axis=1).shape + (3, 2) + >>> stats.trim1(d, 0.8, axis=None).shape + (6,) + + """ + a = np.asarray(a) + if axis is None: + a = a.ravel() + axis = 0 + + nobs = a.shape[axis] + + # avoid possible corner case + if proportiontocut >= 1: + return [] + + if tail.lower() == 'right': + lowercut = 0 + uppercut = nobs - int(proportiontocut * nobs) + + elif tail.lower() == 'left': + lowercut = int(proportiontocut * nobs) + uppercut = nobs + + atmp = np.partition(a, (lowercut, uppercut - 1), axis) + + sl = [slice(None)] * atmp.ndim + sl[axis] = slice(lowercut, uppercut) + return atmp[tuple(sl)] + + +@xp_capabilities(marray=True) +@_axis_nan_policy_factory(lambda x: x, result_to_tuple=lambda x, _: (x,), n_outputs=1) +def trim_mean(a, proportiontocut, axis=0): + """Return mean of array after trimming a specified fraction of extreme values. + + Removes the specified proportion of elements from *each* end of the + sorted array, then computes the mean of the remaining elements. + + Parameters + ---------- + a : array_like + Input array. + proportiontocut : float + Fraction of the most positive and most negative elements to remove. + When the specified proportion does not result in an integer number of + elements, the number of elements to trim is rounded down. + axis : int or None, default: 0 + Axis along which the trimmed means are computed. + If None, compute over the raveled array. + + Returns + ------- + trim_mean : ndarray + Mean of trimmed array. + + See Also + -------- + trimboth : Remove a proportion of elements from each end of an array. + tmean : Compute the mean after trimming values outside specified limits. + + Notes + ----- + For 1-D array `a`, `trim_mean` is approximately equivalent to the following + calculation:: + + import numpy as np + a = np.sort(a) + m = int(proportiontocut * len(a)) + np.mean(a[m: len(a) - m]) + + Examples + -------- + >>> import numpy as np + >>> from scipy import stats + >>> x = [1, 2, 3, 5] + >>> stats.trim_mean(x, 0.25) + 2.5 + + When the specified proportion does not result in an integer number of + elements, the number of elements to trim is rounded down. + + >>> stats.trim_mean(x, 0.24999) == np.mean(x) + True + + Use `axis` to specify the axis along which the calculation is performed. + + >>> x2 = [[1, 2, 3, 5], + ... [10, 20, 30, 50]] + >>> stats.trim_mean(x2, 0.25) + array([ 5.5, 11. , 16.5, 27.5]) + >>> stats.trim_mean(x2, 0.25, axis=1) + array([ 2.5, 25. ]) + + """ + xp = array_namespace(a) + + a = xp.asarray(a) + + if xp_size(a) == 0: + return _get_nan(a, xp=xp) + + if axis is None: + a = xp_ravel(a) + axis = 0 + + nobs = _count_nonmasked(a, axis=axis, keepdims=True, xp=xp) + lowercut = proportiontocut * nobs + nobs, lowercut = ((nobs, int(lowercut)) if not is_marray(xp) + else (xp.astype(nobs, xp.int64), xp.astype(lowercut, xp.int64))) + uppercut = nobs - lowercut + if not is_marray(xp) and (lowercut > uppercut): + raise ValueError("Proportion too big.") + + atmp = (np.partition(a, (lowercut, uppercut - 1), axis) if is_numpy(xp) + else xp.sort(a, axis=axis)) + + if is_marray(xp): + indices = xp.arange(a.shape[-1]) # axis_nan_policy decorator -> axis=-1 + mask = (indices < lowercut) | (indices >= (nobs - lowercut)) | atmp.mask + trimmed = xp.asarray(atmp.data, mask=mask.data) + else: + sl = [slice(None)] * atmp.ndim + sl[axis] = slice(lowercut, uppercut) + trimmed = atmp[tuple(sl)] + + trimmed = xp_promote(trimmed, force_floating=True, xp=xp) + return xp.mean(trimmed, axis=axis) + + +F_onewayResult = namedtuple('F_onewayResult', ('statistic', 'pvalue')) + + +def _f_oneway_is_too_small(samples, kwargs=None, axis=-1): + message = f"At least two samples are required; got {len(samples)}." + if len(samples) < 2: + raise TypeError(message) + + # Check this after forming alldata, so shape errors are detected + # and reported before checking for 0 length inputs. + if any(sample.shape[axis] == 0 for sample in samples): + return True + + # Must have at least one group with length greater than 1. + if all(sample.shape[axis] == 1 for sample in samples): + msg = ('all input arrays have length 1. f_oneway requires that at ' + 'least one input has length greater than 1.') + warnings.warn(SmallSampleWarning(msg), stacklevel=2) + return True + + return False + + +# Torch GPU need fdtrc +@xp_capabilities(cpu_only=True, exceptions=['cupy'], marray=True) +@_axis_nan_policy_factory( + F_onewayResult, n_samples=None, too_small=_f_oneway_is_too_small) +def f_oneway(*samples, axis=0, equal_var=True): + """Perform one-way ANOVA. + + The one-way ANOVA tests the null hypothesis that two or more groups have + the same population mean. The test is applied to samples from two or + more groups, possibly with differing sizes. + + Parameters + ---------- + *samples : array_like + The sample measurements for each group. There must be at least + two arguments. If the arrays are multidimensional, then all the + dimensions of the array must be the same except for `axis`. + axis : int, optional + Axis of the input arrays along which the test is applied. + Default is 0. + equal_var : bool, optional + If True (default), perform a standard one-way ANOVA test that + assumes equal population variances [2]_. + If False, perform Welch's ANOVA test, which does not assume + equal population variances [4]_. + + .. versionadded:: 1.16.0 + + Returns + ------- + statistic : float + The computed F statistic of the test. + pvalue : float + The associated p-value from the F distribution. + + Warns + ----- + `~scipy.stats.ConstantInputWarning` + Emitted if all values within each of the input arrays are identical. + In this case the F statistic is either infinite or isn't defined, + so ``np.inf`` or ``np.nan`` is returned. + + RuntimeWarning + Emitted if the length of any input array is 0, or if all the input + arrays have length 1. ``np.nan`` is returned for the F statistic + and the p-value in these cases. + + Notes + ----- + The ANOVA test has important assumptions that must be satisfied in order + for the associated p-value to be valid. + + 1. The samples are independent. + 2. Each sample is from a normally distributed population. + 3. The population standard deviations of the groups are all equal. This + property is known as homoscedasticity. + + If these assumptions are not true for a given set of data, it may still + be possible to use the Kruskal-Wallis H-test (`scipy.stats.kruskal`) or + the Alexander-Govern test (`scipy.stats.alexandergovern`) although with + some loss of power. + + The length of each group must be at least one, and there must be at + least one group with length greater than one. If these conditions + are not satisfied, a warning is generated and (``np.nan``, ``np.nan``) + is returned. + + If all values in each group are identical, and there exist at least two + groups with different values, the function generates a warning and + returns (``np.inf``, 0). + + If all values in all groups are the same, function generates a warning + and returns (``np.nan``, ``np.nan``). + + The algorithm is from Heiman [2]_, pp.394-7. + + References + ---------- + .. [1] R. Lowry, "Concepts and Applications of Inferential Statistics", + Chapter 14, 2014, http://vassarstats.net/textbook/ + + .. [2] G.W. Heiman, "Understanding research methods and statistics: An + integrated introduction for psychology", Houghton, Mifflin and + Company, 2001. + + .. [3] J.H. McDonald, "Handbook of Biological Statistics", + One-way ANOVA, 2014. + http://www.biostathandbook.com/onewayanova.html + + .. [4] B. L. Welch, "On the Comparison of Several Mean Values: + An Alternative Approach", Biometrika, vol. 38, no. 3/4, + pp. 330-336, 1951. :doi:`10.2307/2332579`. + + .. [5] J.H. McDonald, R. Seed and R.K. Koehn, "Allozymes and + morphometric characters of three species of Mytilus in + the Northern and Southern Hemispheres", + Marine Biology, vol. 111, pp. 323-333, 1991. + + Examples + -------- + >>> import numpy as np + >>> from scipy.stats import f_oneway + + Here are some data [3]_ on a shell measurement (the length of the anterior + adductor muscle scar, standardized by dividing by length) in the mussel + Mytilus trossulus from five locations: Tillamook, Oregon; Newport, Oregon; + Petersburg, Alaska; Magadan, Russia; and Tvarminne, Finland, taken from a + much larger data set used in [5]_. + + >>> tillamook = [0.0571, 0.0813, 0.0831, 0.0976, 0.0817, 0.0859, 0.0735, + ... 0.0659, 0.0923, 0.0836] + >>> newport = [0.0873, 0.0662, 0.0672, 0.0819, 0.0749, 0.0649, 0.0835, + ... 0.0725] + >>> petersburg = [0.0974, 0.1352, 0.0817, 0.1016, 0.0968, 0.1064, 0.105] + >>> magadan = [0.1033, 0.0915, 0.0781, 0.0685, 0.0677, 0.0697, 0.0764, + ... 0.0689] + >>> tvarminne = [0.0703, 0.1026, 0.0956, 0.0973, 0.1039, 0.1045] + >>> f_oneway(tillamook, newport, petersburg, magadan, tvarminne) + F_onewayResult(statistic=7.121019471642447, pvalue=0.0002812242314534544) + + `f_oneway` accepts multidimensional input arrays. When the inputs + are multidimensional and `axis` is not given, the test is performed + along the first axis of the input arrays. For the following data, the + test is performed three times, once for each column. + + >>> a = np.array([[9.87, 9.03, 6.81], + ... [7.18, 8.35, 7.00], + ... [8.39, 7.58, 7.68], + ... [7.45, 6.33, 9.35], + ... [6.41, 7.10, 9.33], + ... [8.00, 8.24, 8.44]]) + >>> b = np.array([[6.35, 7.30, 7.16], + ... [6.65, 6.68, 7.63], + ... [5.72, 7.73, 6.72], + ... [7.01, 9.19, 7.41], + ... [7.75, 7.87, 8.30], + ... [6.90, 7.97, 6.97]]) + >>> c = np.array([[3.31, 8.77, 1.01], + ... [8.25, 3.24, 3.62], + ... [6.32, 8.81, 5.19], + ... [7.48, 8.83, 8.91], + ... [8.59, 6.01, 6.07], + ... [3.07, 9.72, 7.48]]) + >>> F = f_oneway(a, b, c) + >>> F.statistic + array([1.75676344, 0.03701228, 3.76439349]) + >>> F.pvalue + array([0.20630784, 0.96375203, 0.04733157]) + + Welch ANOVA will be performed if `equal_var` is False. + + """ + xp = array_namespace(*samples) + samples = xp_promote(*samples, force_floating=True, xp=xp) + + if len(samples) < 2: + raise TypeError('at least two inputs are required;' + f' got {len(samples)}.') + + # ANOVA on N groups, each in its own array + num_groups = len(samples) + + # axis is guaranteed to be -1 by the _axis_nan_policy decorator + alldata = xp.concat(samples, axis=-1) + bign = _count_nonmasked(alldata, axis=-1, xp=xp) + + # Check if the inputs are too small (for testing _axis_nan_policy decorator) + if _f_oneway_is_too_small(samples): + NaN = _get_nan(*samples, xp=xp) + return F_onewayResult(NaN, NaN) + + # Check if all values within each group are identical, and if the common + # value in at least one group is different from that in another group. + # Based on https://github.com/scipy/scipy/issues/11669 + + # If axis=0, say, and the groups have shape (n0, ...), (n1, ...), ..., + # then is_const is a boolean array with shape (num_groups, ...). + # It is True if the values within the groups along the axis slice are + # identical. In the typical case where each input array is 1-d, is_const is + # a 1-d array with length num_groups. + is_const = xp.concat([xp.all(xp.diff(sample, axis=-1) == 0, axis=-1, keepdims=True) + for sample in samples], axis=-1) + + # all_const is a boolean array with shape (...) (see previous comment). + # It is True if the values within each group along the axis slice are + # the same (e.g. [[3, 3, 3], [5, 5, 5, 5], [4, 4, 4]]). + all_const = xp.all(is_const, axis=-1) + + # all_same_const is True if all the values in the groups along the axis=0 + # slice are the same (e.g. [[3, 3, 3], [3, 3, 3, 3], [3, 3, 3]]). + all_same_const = xp.all(xp.diff(alldata, axis=-1) == 0, axis=-1) + + if not isinstance(equal_var, bool): + raise TypeError("Expected a boolean value for 'equal_var'") + + if equal_var: + # Determine the mean of the data, and subtract that from all inputs to a + # variance (via sum_of_sq / sq_of_sum) calculation. Variance is invariant + # to a shift in location, and centering all data around zero vastly + # improves numerical stability. + offset = xp.mean(alldata, axis=-1, keepdims=True) + alldata = alldata - offset + + normalized_ss = xp.sum(alldata, axis=-1)**2. / bign + + sstot = xp.vecdot(alldata, alldata, axis=-1) - normalized_ss + + ssbn = 0 + for sample in samples: + smo_ss = xp.sum(sample - offset, axis=-1)**2. + ssbn = ssbn + smo_ss / _count_nonmasked(sample, axis=-1, xp=xp) + + # Naming: variables ending in bn/b are for "between treatments", wn/w are + # for "within treatments" + ssbn = ssbn - normalized_ss + sswn = sstot - ssbn + dfbn = num_groups - 1 + dfwn = bign - num_groups + msb = ssbn / dfbn + msw = sswn / dfwn + with np.errstate(divide='ignore', invalid='ignore'): + f = msb / msw + dfn, dfd = dfbn, dfwn + + else: + # calculate basic statistics for each sample + # Beginning of second paragraph [4] page 1: + # "As a particular case $y_t$ may be the means ... of samples + y_t = xp.stack([xp.mean(sample, axis=-1) for sample in samples]) + # "... of $n_t$ observations..." + if is_marray(xp): + n_t = xp.stack([_count_nonmasked(sample, axis=-1, xp=xp) + for sample in samples]) + n_t = xp.asarray(n_t, dtype=n_t.dtype) + else: + n_t = xp.asarray([sample.shape[-1] for sample in samples], dtype=y_t.dtype) + n_t = xp.reshape(n_t, (-1,) + (1,) * (y_t.ndim - 1)) + # "... from $k$ different normal populations..." + k = len(samples) + # "The separate samples provide estimates $s_t^2$ of the $\sigma_t^2$." + s_t2 = xp.stack([xp.var(sample, axis=-1, correction=1) for sample in samples]) + + # calculate weight by number of data and variance + # "we have $\lambda_t = 1 / n_t$ ... where w_t = 1 / {\lambda_t s_t^2}$" + w_t = n_t / s_t2 + # sum of w_t + s_w_t = xp.sum(w_t, axis=0) + + # calculate adjusted grand mean + # "... and $\hat{y} = \sum w_t y_t / \sum w_t$. When all..." + axis_zero = -w_t.ndim + y_hat = xp.vecdot(w_t, y_t, axis=axis_zero) / xp.sum(w_t, axis=0) + + # adjust f statistic + # ref.[4] p.334 eq.29 + numerator = xp.vecdot(w_t, (y_t - y_hat)**2, axis=axis_zero) / (k - 1) + denominator = ( + 1 + 2 * (k - 2) / (k**2 - 1) * + xp.vecdot(1 / (n_t - 1), (1 - w_t / s_w_t)**2, axis=axis_zero) + ) + f = numerator / denominator + + # degree of freedom 1 + # ref.[4] p.334 eq.30 + hat_f1 = k - 1 + + # adjusted degree of freedom 2 + # ref.[4] p.334 eq.30 + hat_f2 = ( + (k**2 - 1) / + (3 * xp.vecdot(1 / (n_t - 1), (1 - w_t / s_w_t)**2, axis=axis_zero)) + ) + + dfn, dfd = hat_f1, hat_f2 + + # Fix any f values that should be inf or nan because the corresponding + # inputs were constant. + f = xpx.at(f)[all_const].set(xp.inf) + f = xpx.at(f)[all_same_const].set(xp.nan) + + # calculate p value + # ref.[4] p.334 eq.28 + prob = special.fdtrc(dfn, dfd, f) + + f, prob = (f[()], prob[()]) if f.ndim == 0 else (f, prob) + return F_onewayResult(f, prob) + + +@dataclass +class AlexanderGovernResult: + statistic: float + pvalue: float + + +@xp_capabilities(marray=True) +@_axis_nan_policy_factory( + AlexanderGovernResult, n_samples=None, + result_to_tuple=lambda x, _: (x.statistic, x.pvalue), + too_small=1 +) +def alexandergovern(*samples, nan_policy='propagate', axis=0): + """Performs the Alexander Govern test. + + The Alexander-Govern approximation tests the equality of k independent + means in the face of heterogeneity of variance. The test is applied to + samples from two or more groups, possibly with differing sizes. + + Parameters + ---------- + *samples : array_like + The sample measurements for each group. There must be at least + two samples, and each sample must contain at least two observations. + nan_policy : {'propagate', 'raise', 'omit'}, optional + Defines how to handle when input contains nan. + The following options are available (default is 'propagate'): + + * 'propagate': returns nan + * 'raise': throws an error + * 'omit': performs the calculations ignoring nan values + + Returns + ------- + res : AlexanderGovernResult + An object with attributes: + + statistic : float + The computed A statistic of the test. + pvalue : float + The associated p-value from the chi-squared distribution. + + Warns + ----- + `~scipy.stats.ConstantInputWarning` + Raised if an input is a constant array. The statistic is not defined + in this case, so ``np.nan`` is returned. + + See Also + -------- + f_oneway : one-way ANOVA + + Notes + ----- + The use of this test relies on several assumptions. + + 1. The samples are independent. + 2. Each sample is from a normally distributed population. + 3. Unlike `f_oneway`, this test does not assume on homoscedasticity, + instead relaxing the assumption of equal variances. + + Input samples must be finite, one dimensional, and with size greater than + one. + + References + ---------- + .. [1] Alexander, Ralph A., and Diane M. Govern. "A New and Simpler + Approximation for ANOVA under Variance Heterogeneity." Journal + of Educational Statistics, vol. 19, no. 2, 1994, pp. 91-101. + https://www.jstor.org/stable/1165140 + + Examples + -------- + >>> from scipy.stats import alexandergovern + + Here are some data on annual percentage rate of interest charged on + new car loans at nine of the largest banks in four American cities + taken from the National Institute of Standards and Technology's + ANOVA dataset. + + We use `alexandergovern` to test the null hypothesis that all cities + have the same mean APR against the alternative that the cities do not + all have the same mean APR. We decide that a significance level of 5% + is required to reject the null hypothesis in favor of the alternative. + + >>> atlanta = [13.75, 13.75, 13.5, 13.5, 13.0, 13.0, 13.0, 12.75, 12.5] + >>> chicago = [14.25, 13.0, 12.75, 12.5, 12.5, 12.4, 12.3, 11.9, 11.9] + >>> houston = [14.0, 14.0, 13.51, 13.5, 13.5, 13.25, 13.0, 12.5, 12.5] + >>> memphis = [15.0, 14.0, 13.75, 13.59, 13.25, 12.97, 12.5, 12.25, + ... 11.89] + >>> alexandergovern(atlanta, chicago, houston, memphis) + AlexanderGovernResult(statistic=4.65087071883494, + pvalue=0.19922132490385214) + + The p-value is 0.1992, indicating a nearly 20% chance of observing + such an extreme value of the test statistic under the null hypothesis. + This exceeds 5%, so we do not reject the null hypothesis in favor of + the alternative. + + """ + xp = array_namespace(*samples) + samples = _alexandergovern_input_validation(samples, nan_policy, axis, xp=xp) + + # The following formula numbers reference the equation described on + # page 92 by Alexander, Govern. Formulas 5, 6, and 7 describe other + # tests that serve as the basis for equation (8) but are not needed + # to perform the test. + + # precalculate mean and length of each sample + lengths = [_count_nonmasked(sample, axis=-1, xp=xp) for sample in samples] + means = xp.stack([_xp_mean(sample, axis=-1) for sample in samples]) + + # (1) determine standard error of the mean for each sample + se2 = [(_xp_var(sample, correction=1, axis=-1) / length) + for sample, length in zip(samples, lengths)] + standard_errors_squared = xp.stack(se2) + standard_errors = standard_errors_squared**0.5 + + # Special case: statistic is NaN when variance is zero + eps = xp.finfo(standard_errors.dtype).eps + zero = standard_errors <= xp.abs(eps * means) + NaN = xp.asarray(xp.nan, dtype=standard_errors.dtype) + standard_errors = xp.where(zero, NaN, standard_errors) + + # (2) define a weight for each sample + inv_sq_se = 1 / standard_errors_squared + weights = inv_sq_se / xp.sum(inv_sq_se, axis=0, keepdims=True) + + # (3) determine variance-weighted estimate of the common mean + # Consider replacing with vecdot when data-apis/array-api#910 is resolved + var_w = xp.sum(weights * means, axis=0, keepdims=True) + + # (4) determine one-sample t statistic for each group + t_stats = _demean(means, var_w, axis=0, xp=xp) / standard_errors + + # calculate parameters to be used in transformation + if is_marray(xp): + v = xp.stack(lengths) - 1 + else: + v = xp.asarray(lengths, dtype=t_stats.dtype) - 1 + # align along 0th axis, which corresponds with separate samples + v = xp.reshape(v, (-1,) + (1,)*(t_stats.ndim-1)) + a = v - .5 + b = 48 * a**2 + c = (a * xp.log(1 + (t_stats ** 2)/v))**.5 + + # (8) perform a normalizing transformation on t statistic + z = (c + ((c**3 + 3*c)/b) - + ((4*c**7 + 33*c**5 + 240*c**3 + 855*c) / + (b**2*10 + 8*b*c**4 + 1000*b))) + + # (9) calculate statistic + A = xp.vecdot(z, z, axis=-z.ndim) + A = A[()] if A.ndim == 0 else A # data-apis/array-api-compat#355 + + # "[the p value is determined from] central chi-square random deviates + # with k - 1 degrees of freedom". Alexander, Govern (94) + df = xp.asarray(len(samples) - 1, dtype=A.dtype) + chi2 = _SimpleChi2(df) + p = _get_pvalue(A, chi2, alternative='greater', symmetric=False, xp=xp) + return AlexanderGovernResult(A, p) + + +def _alexandergovern_input_validation(samples, nan_policy, axis, xp): + if len(samples) < 2: + raise TypeError(f"2 or more inputs required, got {len(samples)}") + + for sample in samples: + if sample.shape[axis] <= 1: + raise ValueError("Input sample size must be greater than one.") + + samples = [xp.moveaxis(sample, axis, -1) for sample in samples] + + return samples + + +def _pearsonr_fisher_ci(r, n, confidence_level, alternative): + """ + Compute the confidence interval for Pearson's R. + + Fisher's transformation is used to compute the confidence interval + (https://en.wikipedia.org/wiki/Fisher_transformation). + """ + xp = array_namespace(r) + + ones = xp.ones_like(r) + n = xp.asarray(n, dtype=r.dtype, device=xp_device(r)) + confidence_level = xp.asarray(confidence_level, dtype=r.dtype, device=xp_device(r)) + + with np.errstate(divide='ignore', invalid='ignore'): + zr = xp.atanh(r) + se = xp.sqrt(1 / (n - 3)) + + if alternative == "two-sided": + h = special.ndtri(0.5 + confidence_level/2) + zlo = zr - h*se + zhi = zr + h*se + rlo = xp.tanh(zlo) + rhi = xp.tanh(zhi) + elif alternative == "less": + h = special.ndtri(confidence_level) + zhi = zr + h*se + rhi = xp.tanh(zhi) + rlo = -ones + else: + # alternative == "greater": + h = special.ndtri(confidence_level) + zlo = zr - h*se + rlo = xp.tanh(zlo) + rhi = ones + + mask = (n <= 3) + if mask.ndim == 0: + # This is Array API legal, but Dask doesn't like it. + mask = xp.broadcast_to(mask, rlo.shape) + + rlo = xpx.at(rlo)[mask].set(-1) + rhi = xpx.at(rhi)[mask].set(1) + + rlo = rlo[()] if rlo.ndim == 0 else rlo + rhi = rhi[()] if rhi.ndim == 0 else rhi + return ConfidenceInterval(low=rlo, high=rhi) + + +def _pearsonr_bootstrap_ci(confidence_level, method, x, y, alternative, axis): + """ + Compute the confidence interval for Pearson's R using the bootstrap. + """ + def statistic(x, y, axis): + statistic, _ = pearsonr(x, y, axis=axis) + return statistic + + res = bootstrap((x, y), statistic, confidence_level=confidence_level, axis=axis, + paired=True, alternative=alternative, **method._asdict()) + # for one-sided confidence intervals, bootstrap gives +/- inf on one side + res.confidence_interval = np.clip(res.confidence_interval, -1, 1) + + return ConfidenceInterval(*res.confidence_interval) + + +ConfidenceInterval = namedtuple('ConfidenceInterval', ['low', 'high']) + +PearsonRResultBase = _make_tuple_bunch('PearsonRResultBase', + ['statistic', 'pvalue'], []) + + +class PearsonRResult(PearsonRResultBase): + """ + Result of `scipy.stats.pearsonr` + + Attributes + ---------- + statistic : float + Pearson product-moment correlation coefficient. + pvalue : float + The p-value associated with the chosen alternative. + + Methods + ------- + confidence_interval + Computes the confidence interval of the correlation + coefficient `statistic` for the given confidence level. + + """ + def __init__(self, statistic, pvalue, alternative, n, x, y, axis): + super().__init__(statistic, pvalue) + self._alternative = alternative + self._n = n + self._x = x + self._y = y + self._axis = axis + + # add alias for consistency with other correlation functions + self.correlation = statistic + + def confidence_interval(self, confidence_level=0.95, method=None): + """ + The confidence interval for the correlation coefficient. + + Compute the confidence interval for the correlation coefficient + ``statistic`` with the given confidence level. + + If `method` is not provided, + The confidence interval is computed using the Fisher transformation + F(r) = arctanh(r) [1]_. When the sample pairs are drawn from a + bivariate normal distribution, F(r) approximately follows a normal + distribution with standard error ``1/sqrt(n - 3)``, where ``n`` is the + length of the original samples along the calculation axis. When + ``n <= 3``, this approximation does not yield a finite, real standard + error, so we define the confidence interval to be -1 to 1. + + If `method` is an instance of `BootstrapMethod`, the confidence + interval is computed using `scipy.stats.bootstrap` with the provided + configuration options and other appropriate settings. In some cases, + confidence limits may be NaN due to a degenerate resample, and this is + typical for very small samples (~6 observations). + + Parameters + ---------- + confidence_level : float + The confidence level for the calculation of the correlation + coefficient confidence interval. Default is 0.95. + + method : BootstrapMethod, optional + Defines the method used to compute the confidence interval. See + method description for details. + + .. versionadded:: 1.11.0 + + Returns + ------- + ci : namedtuple + The confidence interval is returned in a ``namedtuple`` with + fields `low` and `high`. + + References + ---------- + .. [1] "Pearson correlation coefficient", Wikipedia, + https://en.wikipedia.org/wiki/Pearson_correlation_coefficient + """ + if isinstance(method, BootstrapMethod): + xp = array_namespace(self._x) + message = ('`method` must be `None` if `pearsonr` ' + 'arguments were not NumPy arrays.') + if not is_numpy(xp): + raise ValueError(message) + + ci = _pearsonr_bootstrap_ci(confidence_level, method, self._x, self._y, + self._alternative, self._axis) + elif method is None: + ci = _pearsonr_fisher_ci(self.statistic, self._n, confidence_level, + self._alternative) + else: + message = ('`method` must be an instance of `BootstrapMethod` ' + 'or None.') + raise ValueError(message) + return ci + + +# Missing special.betainc on torch +@xp_capabilities(cpu_only=True, exceptions=['cupy', 'jax.numpy'], marray=True, + extra_note='Only the default `method` is compatible with MArray input.') +def pearsonr(x, y, *, alternative='two-sided', method=None, axis=0): + r""" + Pearson correlation coefficient and p-value for testing non-correlation. + + The Pearson correlation coefficient [1]_ measures the linear relationship + between two datasets. Like other correlation + coefficients, this one varies between -1 and +1 with 0 implying no + correlation. Correlations of -1 or +1 imply an exact linear relationship. + Positive correlations imply that as x increases, so does y. Negative + correlations imply that as x increases, y decreases. + + This function also performs a test of the null hypothesis that the + distributions underlying the samples are uncorrelated and normally + distributed. (See Kowalski [3]_ + for a discussion of the effects of non-normality of the input on the + distribution of the correlation coefficient.) + The p-value roughly indicates the probability of an uncorrelated system + producing datasets that have a Pearson correlation at least as extreme + as the one computed from these datasets. + + Parameters + ---------- + x : array_like + Input array. + y : array_like + Input array. + alternative : {'two-sided', 'greater', 'less'}, optional + Defines the alternative hypothesis. Default is 'two-sided'. + The following options are available: + + * 'two-sided': the correlation is nonzero + * 'less': the correlation is negative (less than zero) + * 'greater': the correlation is positive (greater than zero) + + .. versionadded:: 1.9.0 + method : ResamplingMethod, optional + Defines the method used to compute the p-value. If `method` is an + instance of `PermutationMethod`/`MonteCarloMethod`, the p-value is + computed using + `scipy.stats.permutation_test`/`scipy.stats.monte_carlo_test` with the + provided configuration options and other appropriate settings. + Otherwise, the p-value is computed as documented in the notes. + + .. versionadded:: 1.11.0 + axis : int or None, default + Axis along which to perform the calculation. Default is 0. + If None, ravel both arrays before performing the calculation. + + .. versionadded:: 1.14.0 + + Returns + ------- + result : `~scipy.stats._result_classes.PearsonRResult` + An object with the following attributes: + + statistic : float + Pearson product-moment correlation coefficient. + pvalue : float + The p-value associated with the chosen alternative. + + The object has the following method: + + confidence_interval(confidence_level, method) + This computes the confidence interval of the correlation + coefficient `statistic` for the given confidence level. + The confidence interval is returned in a ``namedtuple`` with + fields `low` and `high`. If `method` is not provided, the + confidence interval is computed using the Fisher transformation + [1]_. If `method` is an instance of `BootstrapMethod`, the + confidence interval is computed using `scipy.stats.bootstrap` with + the provided configuration options and other appropriate settings. + In some cases, confidence limits may be NaN due to a degenerate + resample, and this is typical for very small samples (~6 + observations). + + Raises + ------ + ValueError + If `x` and `y` do not have length at least 2. + + Warns + ----- + `~scipy.stats.ConstantInputWarning` + Raised if an input is a constant array. The correlation coefficient + is not defined in this case, so ``np.nan`` is returned. + + `~scipy.stats.NearConstantInputWarning` + Raised if an input is "nearly" constant. The array ``x`` is considered + nearly constant if ``norm(x - mean(x)) < 1e-13 * abs(mean(x))``. + Numerical errors in the calculation ``x - mean(x)`` in this case might + result in an inaccurate calculation of r. + + See Also + -------- + spearmanr : Spearman rank-order correlation coefficient. + kendalltau : Kendall's tau, a correlation measure for ordinal data. + :ref:`hypothesis_pearsonr` : Extended example + + Notes + ----- + The correlation coefficient is calculated as follows: + + .. math:: + + r = \frac{\sum (x - m_x) (y - m_y)} + {\sqrt{\sum (x - m_x)^2 \sum (y - m_y)^2}} + + where :math:`m_x` is the mean of the vector x and :math:`m_y` is + the mean of the vector y. + + Under the assumption that x and y are drawn from + independent normal distributions (so the population correlation coefficient + is 0), the probability density function of the sample correlation + coefficient r is ([1]_, [2]_): + + .. math:: + f(r) = \frac{{(1-r^2)}^{n/2-2}}{\mathrm{B}(\frac{1}{2},\frac{n}{2}-1)} + + where n is the number of samples, and B is the beta function. This + is sometimes referred to as the exact distribution of r. This is + the distribution that is used in `pearsonr` to compute the p-value when + the `method` parameter is left at its default value (None). + The distribution is a beta distribution on the interval [-1, 1], + with equal shape parameters a = b = n/2 - 1. In terms of SciPy's + implementation of the beta distribution, the distribution of r is:: + + dist = scipy.stats.beta(n/2 - 1, n/2 - 1, loc=-1, scale=2) + + The default p-value returned by `pearsonr` is a two-sided p-value. For a + given sample with correlation coefficient r, the p-value is + the probability that abs(r') of a random sample x' and y' drawn from + the population with zero correlation would be greater than or equal + to abs(r). In terms of the object ``dist`` shown above, the p-value + for a given r and length n can be computed as:: + + p = 2*dist.cdf(-abs(r)) + + When n is 2, the above continuous distribution is not well-defined. + One can interpret the limit of the beta distribution as the shape + parameters a and b approach a = b = 0 as a discrete distribution with + equal probability masses at r = 1 and r = -1. More directly, one + can observe that, given the data x = [x1, x2] and y = [y1, y2], and + assuming x1 != x2 and y1 != y2, the only possible values for r are 1 + and -1. Because abs(r') for any sample x' and y' with length 2 will + be 1, the two-sided p-value for a sample of length 2 is always 1. + + For backwards compatibility, the object that is returned also behaves + like a tuple of length two that holds the statistic and the p-value. + + References + ---------- + .. [1] "Pearson correlation coefficient", Wikipedia, + https://en.wikipedia.org/wiki/Pearson_correlation_coefficient + .. [2] Student, "Probable error of a correlation coefficient", + Biometrika, Volume 6, Issue 2-3, 1 September 1908, pp. 302-310. + .. [3] C. J. Kowalski, "On the Effects of Non-Normality on the Distribution + of the Sample Product-Moment Correlation Coefficient" + Journal of the Royal Statistical Society. Series C (Applied + Statistics), Vol. 21, No. 1 (1972), pp. 1-12. + + Examples + -------- + >>> import numpy as np + >>> from scipy import stats + >>> x, y = [1, 2, 3, 4, 5, 6, 7], [10, 9, 2.5, 6, 4, 3, 2] + >>> res = stats.pearsonr(x, y) + >>> res + PearsonRResult(statistic=-0.828503883588428, pvalue=0.021280260007523286) + + To perform an exact permutation version of the test: + + >>> rng = np.random.default_rng(7796654889291491997) + >>> method = stats.PermutationMethod(n_resamples=np.inf, random_state=rng) + >>> stats.pearsonr(x, y, method=method) + PearsonRResult(statistic=-0.828503883588428, pvalue=0.028174603174603175) + + To perform the test under the null hypothesis that the data were drawn from + *uniform* distributions: + + >>> method = stats.MonteCarloMethod(rvs=(rng.uniform, rng.uniform)) + >>> stats.pearsonr(x, y, method=method) + PearsonRResult(statistic=-0.828503883588428, pvalue=0.0188) + + To produce an asymptotic 90% confidence interval: + + >>> res.confidence_interval(confidence_level=0.9) + ConfidenceInterval(low=-0.9644331982722841, high=-0.3460237473272273) + + And for a bootstrap confidence interval: + + >>> method = stats.BootstrapMethod(method='BCa', rng=rng) + >>> res.confidence_interval(confidence_level=0.9, method=method) + ConfidenceInterval(low=-0.9983163756488651, high=-0.22771001702132443) # may vary + + If N-dimensional arrays are provided, multiple tests are performed in a + single call according to the same conventions as most `scipy.stats` functions: + + >>> rng = np.random.default_rng(2348246935601934321) + >>> x = rng.standard_normal((8, 15)) + >>> y = rng.standard_normal((8, 15)) + >>> stats.pearsonr(x, y, axis=0).statistic.shape # between corresponding columns + (15,) + >>> stats.pearsonr(x, y, axis=1).statistic.shape # between corresponding rows + (8,) + + To perform all pairwise comparisons between slices of the arrays, + use standard NumPy broadcasting techniques. For instance, to compute the + correlation between all pairs of rows: + + >>> stats.pearsonr(x[:, np.newaxis, :], y, axis=-1).statistic.shape + (8, 8) + + There is a linear dependence between x and y if y = a + b*x + e, where + a,b are constants and e is a random error term, assumed to be independent + of x. For simplicity, assume that x is standard normal, a=0, b=1 and let + e follow a normal distribution with mean zero and standard deviation s>0. + + >>> rng = np.random.default_rng() + >>> s = 0.5 + >>> x = stats.norm.rvs(size=500, random_state=rng) + >>> e = stats.norm.rvs(scale=s, size=500, random_state=rng) + >>> y = x + e + >>> stats.pearsonr(x, y).statistic + 0.9001942438244763 + + This should be close to the exact value given by + + >>> 1/np.sqrt(1 + s**2) + 0.8944271909999159 + + For s=0.5, we observe a high level of correlation. In general, a large + variance of the noise reduces the correlation, while the correlation + approaches one as the variance of the error goes to zero. + + It is important to keep in mind that no correlation does not imply + independence unless (x, y) is jointly normal. Correlation can even be zero + when there is a very simple dependence structure: if X follows a + standard normal distribution, let y = abs(x). Note that the correlation + between x and y is zero. Indeed, since the expectation of x is zero, + cov(x, y) = E[x*y]. By definition, this equals E[x*abs(x)] which is zero + by symmetry. The following lines of code illustrate this observation: + + >>> y = np.abs(x) + >>> stats.pearsonr(x, y) + PearsonRResult(statistic=-0.05444919272687482, pvalue=0.22422294836207743) + + A non-zero correlation coefficient can be misleading. For example, if X has + a standard normal distribution, define y = x if x < 0 and y = 0 otherwise. + A simple calculation shows that corr(x, y) = sqrt(2/Pi) = 0.797..., + implying a high level of correlation: + + >>> y = np.where(x < 0, x, 0) + >>> stats.pearsonr(x, y) + PearsonRResult(statistic=0.861985781588, pvalue=4.813432002751103e-149) + + This is unintuitive since there is no dependence of x and y if x is larger + than zero which happens in about half of the cases if we sample x and y. + + For a more detailed example, see :ref:`hypothesis_pearsonr`. + + """ + xp = array_namespace(x, y) + x, y = xp_promote(x, y, force_floating=True, xp=xp) + dtype = x.dtype + + if not is_numpy(xp) and method is not None: + method = 'invalid' + + if axis is None: + x = xp.reshape(x, (-1,)) + y = xp.reshape(y, (-1,)) + axis = -1 + + axis_int = int(axis) + if axis_int != axis: + raise ValueError('`axis` must be an integer.') + axis = axis_int + + try: + np.broadcast_shapes(x.shape, y.shape) + # For consistency with other `stats` functions, we need to + # match the dimensionalities before looking at `axis`. + # (Note: this is not the NEP 5 / gufunc order of operations; + # see TestPearsonr::test_different_dimensionality for more information.) + ndim = max(x.ndim, y.ndim) + x = xp.reshape(x, (1,) * (ndim - x.ndim) + x.shape) + y = xp.reshape(y, (1,) * (ndim - y.ndim) + y.shape) + + except (ValueError, RuntimeError) as e: + message = '`x` and `y` must be broadcastable.' + raise ValueError(message) from e + + if x.shape[axis] != y.shape[axis]: + raise ValueError('`x` and `y` must have the same length along `axis`.') + + if x.shape[axis] < 2: + raise ValueError('`x` and `y` must have length at least 2.') + + x, y = _share_masks(x, y, xp=xp) + n = xp.asarray(_count_nonmasked(x, axis=axis), dtype=x.dtype) + + x = xp.moveaxis(x, axis, -1) + y = xp.moveaxis(y, axis, -1) + axis = -1 + + if xp.isdtype(dtype, "complex floating"): + raise ValueError('This function does not support complex data') + + x = xp.astype(x, dtype, copy=False) + y = xp.astype(y, dtype, copy=False) + threshold = xp.finfo(dtype).eps ** 0.75 + + # If an input is constant, the correlation coefficient is not defined. + if is_marray(xp): + # sort to ensure that we are comparing to a non-masked element + const_x = xp.all(x == xp.sort(x, axis=-1)[..., 0:1], axis=-1) + const_y = xp.all(y == xp.sort(y, axis=-1)[..., 0:1], axis=-1) + else: + const_x = xp.all(x == x[..., 0:1], axis=-1) + const_y = xp.all(y == y[..., 0:1], axis=-1) + const_xy = const_x | const_y + + any_const_xy = xp.any(const_xy) + lazy = is_lazy_array(const_xy) + if not lazy and any_const_xy: + msg = ("An input array is constant; the correlation coefficient " + "is not defined.") + warnings.warn(stats.ConstantInputWarning(msg), stacklevel=2) + if lazy or any_const_xy: + x = xp.where(const_x[..., xp.newaxis], xp.nan, x) + y = xp.where(const_y[..., xp.newaxis], xp.nan, y) + + if isinstance(method, PermutationMethod): + def statistic(y, axis): + statistic, _ = pearsonr(x, y, axis=axis, alternative=alternative) + return statistic + + res = permutation_test((y,), statistic, permutation_type='pairings', + axis=axis, alternative=alternative, **method._asdict()) + + return PearsonRResult(statistic=res.statistic, pvalue=res.pvalue, n=n, + alternative=alternative, x=x, y=y, axis=axis) + elif isinstance(method, MonteCarloMethod): + def statistic(x, y, axis): + statistic, _ = pearsonr(x, y, axis=axis, alternative=alternative) + return statistic + + # `monte_carlo_test` accepts an `rvs` tuple of callables, not an `rng` + # If the user specified an `rng`, replace it with the appropriate callables + method = method._asdict() + if (rng := method.pop('rng', None)) is not None: # goo-goo g'joob + rng = np.random.default_rng(rng) + method['rvs'] = rng.normal, rng.normal + + res = monte_carlo_test((x, y,), statistic=statistic, axis=axis, + alternative=alternative, **method) + + return PearsonRResult(statistic=res.statistic, pvalue=res.pvalue, n=n, + alternative=alternative, x=x, y=y, axis=axis) + elif method == 'invalid': + message = '`method` must be `None` if arguments are not NumPy arrays.' + raise ValueError(message) + elif method is not None: + message = ('`method` must be an instance of `PermutationMethod`, ' + '`MonteCarloMethod`, or None.') + raise ValueError(message) + + xmean = xp.mean(x, axis=axis, keepdims=True) + ymean = xp.mean(y, axis=axis, keepdims=True) + xm = x - xmean + ym = y - ymean + + # scipy.linalg.norm(xm) avoids premature overflow when xm is e.g. + # [-5e210, 5e210, 3e200, -3e200] + # but not when `axis` is provided, so scale manually. scipy.linalg.norm + # also raises an error with NaN input rather than returning NaN, so + # use np.linalg.norm. + xmax = xp.max(xp.abs(xm), axis=axis, keepdims=True) + ymax = xp.max(xp.abs(ym), axis=axis, keepdims=True) + with np.errstate(invalid='ignore', divide='ignore'): + normxm = xmax * xp_vector_norm(xm/xmax, axis=axis, keepdims=True) + normym = ymax * xp_vector_norm(ym/ymax, axis=axis, keepdims=True) + + if not lazy: + nconst_x = xp.any(normxm < threshold*xp.abs(xmean), axis=axis) + nconst_y = xp.any(normym < threshold*xp.abs(ymean), axis=axis) + nconst_xy = nconst_x | nconst_y + if xp.any(nconst_xy & (~const_xy)): + # If all the values in x (likewise y) are very close to the mean, + # the loss of precision that occurs in the subtraction xm = x - xmean + # might result in large errors in r. + msg = ("An input array is nearly constant; the computed " + "correlation coefficient may be inaccurate.") + warnings.warn(stats.NearConstantInputWarning(msg), stacklevel=2) + + with np.errstate(invalid='ignore', divide='ignore'): + r = xp.vecdot(xm / normxm, ym / normym, axis=axis) + + # Presumably, if abs(r) > 1, then it is only some small artifact of + # floating point arithmetic. + r = xp.clip(r, -1., 1.) + r = xpx.at(r, const_xy).set(xp.nan) + + # As explained in the docstring, the distribution of `r` under the null + # hypothesis is the beta distribution on (-1, 1) with a = b = n/2 - 1. + ab = xp.asarray(n/2 - 1, dtype=dtype, device=xp_device(x)) + dist = _SimpleBeta(ab, ab, loc=-1, scale=2) + pvalue = _get_pvalue(r, dist, alternative, xp=xp) + + mask = (n == 2) # return exactly 1.0 or -1.0 values for n == 2 case as promised + # data-apis/array-api-extra#196 + mxp = array_namespace(r._meta) if is_dask(xp) else xp + def special_case(r): + return mxp.where(mxp.isnan(r), mxp.nan, mxp.ones_like(r)) + r = xpx.apply_where(mask, r, mxp.round, fill_value=r) + pvalue = xpx.apply_where(mask, (r,), special_case, fill_value=pvalue) + + r = r[()] if r.ndim == 0 else r + pvalue = pvalue[()] if pvalue.ndim == 0 else pvalue + return PearsonRResult(statistic=r, pvalue=pvalue, n=n, + alternative=alternative, x=x, y=y, axis=axis) + + +@xp_capabilities(np_only=True) +def fisher_exact(table, alternative=None, *, method=None): + """Perform a Fisher exact test on a contingency table. + + For a 2x2 table, + the null hypothesis is that the true odds ratio of the populations + underlying the observations is one, and the observations were sampled + from these populations under a condition: the marginals of the + resulting table must equal those of the observed table. + The statistic is the unconditional maximum likelihood estimate of the odds + ratio, and the p-value is the probability under the null hypothesis of + obtaining a table at least as extreme as the one that was actually + observed. + + For other table sizes, or if `method` is provided, the null hypothesis + is that the rows and columns of the tables have fixed sums and are + independent; i.e., the table was sampled from a `scipy.stats.random_table` + distribution with the observed marginals. The statistic is the + probability mass of this distribution evaluated at `table`, and the + p-value is the percentage of the population of tables with statistic at + least as extreme (small) as that of `table`. There is only one alternative + hypothesis available: the rows and columns are not independent. + + There are other possible choices of statistic and two-sided + p-value definition associated with Fisher's exact test; please see the + Notes for more information. + + Parameters + ---------- + table : array_like of ints + A contingency table. Elements must be non-negative integers. + alternative : {'two-sided', 'less', 'greater'}, optional + Defines the alternative hypothesis for 2x2 tables; unused for other + table sizes. + The following options are available (default is 'two-sided'): + + * 'two-sided': the odds ratio of the underlying population is not one + * 'less': the odds ratio of the underlying population is less than one + * 'greater': the odds ratio of the underlying population is greater + than one + + See the Notes for more details. + + method : ResamplingMethod, optional + Defines the method used to compute the p-value. + If `method` is an instance of `PermutationMethod`/`MonteCarloMethod`, + the p-value is computed using + `scipy.stats.permutation_test`/`scipy.stats.monte_carlo_test` with the + provided configuration options and other appropriate settings. + Note that if `method` is an instance of `MonteCarloMethod`, the ``rvs`` + attribute must be left unspecified; Monte Carlo samples are always drawn + using the ``rvs`` method of `scipy.stats.random_table`. + Otherwise, the p-value is computed as documented in the notes. + + .. versionadded:: 1.15.0 + + Returns + ------- + res : SignificanceResult + An object containing attributes: + + statistic : float + For a 2x2 table with default `method`, this is the odds ratio - the + prior odds ratio not a posterior estimate. In all other cases, this + is the probability density of obtaining the observed table under the + null hypothesis of independence with marginals fixed. + pvalue : float + The probability under the null hypothesis of obtaining a + table at least as extreme as the one that was actually observed. + + Raises + ------ + ValueError + If `table` is not two-dimensional or has negative entries. + + See Also + -------- + chi2_contingency : Chi-square test of independence of variables in a + contingency table. This can be used as an alternative to + `fisher_exact` when the numbers in the table are large. + contingency.odds_ratio : Compute the odds ratio (sample or conditional + MLE) for a 2x2 contingency table. + barnard_exact : Barnard's exact test, which is a more powerful alternative + than Fisher's exact test for 2x2 contingency tables. + boschloo_exact : Boschloo's exact test, which is a more powerful + alternative than Fisher's exact test for 2x2 contingency tables. + :ref:`hypothesis_fisher_exact` : Extended example + + Notes + ----- + *Null hypothesis and p-values* + + The null hypothesis is that the true odds ratio of the populations + underlying the observations is one, and the observations were sampled at + random from these populations under a condition: the marginals of the + resulting table must equal those of the observed table. Equivalently, + the null hypothesis is that the input table is from the hypergeometric + distribution with parameters (as used in `hypergeom`) + ``M = a + b + c + d``, ``n = a + b`` and ``N = a + c``, where the + input table is ``[[a, b], [c, d]]``. This distribution has support + ``max(0, N + n - M) <= x <= min(N, n)``, or, in terms of the values + in the input table, ``min(0, a - d) <= x <= a + min(b, c)``. ``x`` + can be interpreted as the upper-left element of a 2x2 table, so the + tables in the distribution have form:: + + [ x n - x ] + [N - x M - (n + N) + x] + + For example, if:: + + table = [6 2] + [1 4] + + then the support is ``2 <= x <= 7``, and the tables in the distribution + are:: + + [2 6] [3 5] [4 4] [5 3] [6 2] [7 1] + [5 0] [4 1] [3 2] [2 3] [1 4] [0 5] + + The probability of each table is given by the hypergeometric distribution + ``hypergeom.pmf(x, M, n, N)``. For this example, these are (rounded to + three significant digits):: + + x 2 3 4 5 6 7 + p 0.0163 0.163 0.408 0.326 0.0816 0.00466 + + These can be computed with:: + + >>> import numpy as np + >>> from scipy.stats import hypergeom + >>> table = np.array([[6, 2], [1, 4]]) + >>> M = table.sum() + >>> n = table[0].sum() + >>> N = table[:, 0].sum() + >>> start, end = hypergeom.support(M, n, N) + >>> hypergeom.pmf(np.arange(start, end+1), M, n, N) + array([0.01631702, 0.16317016, 0.40792541, 0.32634033, 0.08158508, + 0.004662 ]) + + The two-sided p-value is the probability that, under the null hypothesis, + a random table would have a probability equal to or less than the + probability of the input table. For our example, the probability of + the input table (where ``x = 6``) is 0.0816. The x values where the + probability does not exceed this are 2, 6 and 7, so the two-sided p-value + is ``0.0163 + 0.0816 + 0.00466 ~= 0.10256``:: + + >>> from scipy.stats import fisher_exact + >>> res = fisher_exact(table, alternative='two-sided') + >>> res.pvalue + 0.10256410256410257 + + The one-sided p-value for ``alternative='greater'`` is the probability + that a random table has ``x >= a``, which in our example is ``x >= 6``, + or ``0.0816 + 0.00466 ~= 0.08626``:: + + >>> res = fisher_exact(table, alternative='greater') + >>> res.pvalue + 0.08624708624708627 + + This is equivalent to computing the survival function of the + distribution at ``x = 5`` (one less than ``x`` from the input table, + because we want to include the probability of ``x = 6`` in the sum):: + + >>> hypergeom.sf(5, M, n, N) + 0.08624708624708627 + + For ``alternative='less'``, the one-sided p-value is the probability + that a random table has ``x <= a``, (i.e. ``x <= 6`` in our example), + or ``0.0163 + 0.163 + 0.408 + 0.326 + 0.0816 ~= 0.9949``:: + + >>> res = fisher_exact(table, alternative='less') + >>> res.pvalue + 0.9953379953379957 + + This is equivalent to computing the cumulative distribution function + of the distribution at ``x = 6``: + + >>> hypergeom.cdf(6, M, n, N) + 0.9953379953379957 + + *Odds ratio* + + The calculated odds ratio is different from the value computed by the + R function ``fisher.test``. This implementation returns the "sample" + or "unconditional" maximum likelihood estimate, while ``fisher.test`` + in R uses the conditional maximum likelihood estimate. To compute the + conditional maximum likelihood estimate of the odds ratio, use + `scipy.stats.contingency.odds_ratio`. + + References + ---------- + .. [1] Fisher, Sir Ronald A, "The Design of Experiments: + Mathematics of a Lady Tasting Tea." ISBN 978-0-486-41151-4, 1935. + .. [2] "Fisher's exact test", + https://en.wikipedia.org/wiki/Fisher's_exact_test + + Examples + -------- + + >>> from scipy.stats import fisher_exact + >>> res = fisher_exact([[8, 2], [1, 5]]) + >>> res.statistic + 20.0 + >>> res.pvalue + 0.034965034965034975 + + For tables with shape other than ``(2, 2)``, provide an instance of + `scipy.stats.MonteCarloMethod` or `scipy.stats.PermutationMethod` for the + `method` parameter: + + >>> import numpy as np + >>> from scipy.stats import MonteCarloMethod + >>> rng = np.random.default_rng(4507195762371367) + >>> method = MonteCarloMethod(rng=rng) + >>> fisher_exact([[8, 2, 3], [1, 5, 4]], method=method) + SignificanceResult(statistic=np.float64(0.005782), pvalue=np.float64(0.0603)) + + For a more detailed example, see :ref:`hypothesis_fisher_exact`. + """ + hypergeom = distributions.hypergeom + # int32 is not enough for the algorithm + c = np.asarray(table, dtype=np.int64) + if not c.ndim == 2: + raise ValueError("The input `table` must have two dimensions.") + + if np.any(c < 0): + raise ValueError("All values in `table` must be nonnegative.") + + if not c.shape == (2, 2) or method is not None: + return _fisher_exact_rxc(c, alternative, method) + alternative = 'two-sided' if alternative is None else alternative + + if 0 in c.sum(axis=0) or 0 in c.sum(axis=1): + # If both values in a row or column are zero, the p-value is 1 and + # the odds ratio is NaN. + return SignificanceResult(np.nan, 1.0) + + if c[1, 0] > 0 and c[0, 1] > 0: + oddsratio = c[0, 0] * c[1, 1] / (c[1, 0] * c[0, 1]) + else: + oddsratio = np.inf + + n1 = c[0, 0] + c[0, 1] + n2 = c[1, 0] + c[1, 1] + n = c[0, 0] + c[1, 0] + + def pmf(x): + return hypergeom.pmf(x, n1 + n2, n1, n) + + if alternative == 'less': + pvalue = hypergeom.cdf(c[0, 0], n1 + n2, n1, n) + elif alternative == 'greater': + # Same formula as the 'less' case, but with the second column. + pvalue = hypergeom.cdf(c[0, 1], n1 + n2, n1, c[0, 1] + c[1, 1]) + elif alternative == 'two-sided': + mode = int((n + 1) * (n1 + 1) / (n1 + n2 + 2)) + pexact = hypergeom.pmf(c[0, 0], n1 + n2, n1, n) + pmode = hypergeom.pmf(mode, n1 + n2, n1, n) + + epsilon = 1e-14 + gamma = 1 + epsilon + + if np.abs(pexact - pmode) / np.maximum(pexact, pmode) <= epsilon: + return SignificanceResult(oddsratio, 1.) + + elif c[0, 0] < mode: + plower = hypergeom.cdf(c[0, 0], n1 + n2, n1, n) + if hypergeom.pmf(n, n1 + n2, n1, n) > pexact * gamma: + return SignificanceResult(oddsratio, plower) + + guess = _binary_search(lambda x: -pmf(x), -pexact * gamma, mode, n, xp=np) + pvalue = plower + hypergeom.sf(guess, n1 + n2, n1, n) + else: + pupper = hypergeom.sf(c[0, 0] - 1, n1 + n2, n1, n) + if hypergeom.pmf(0, n1 + n2, n1, n) > pexact * gamma: + return SignificanceResult(oddsratio, pupper) + + guess = _binary_search(pmf, pexact * gamma, 0, mode, xp=np) + pvalue = pupper + hypergeom.cdf(guess, n1 + n2, n1, n) + else: + msg = "`alternative` should be one of {'two-sided', 'less', 'greater'}" + raise ValueError(msg) + + pvalue = min(pvalue, 1.0) + + return SignificanceResult(oddsratio, pvalue) + + +def _fisher_exact_rxc(table, alternative, method): + if alternative is not None: + message = ('`alternative` must be the default (None) unless ' + '`table` has shape `(2, 2)` and `method is None`.') + raise ValueError(message) + + if table.size == 0: + raise ValueError("`table` must have at least one row and one column.") + + if table.shape[0] == 1 or table.shape[1] == 1 or np.all(table == 0): + # Only one such table with those marginals + return SignificanceResult(1.0, 1.0) + + if method is None: + method = stats.MonteCarloMethod() + + if isinstance(method, stats.PermutationMethod): + res = _fisher_exact_permutation_method(table, method) + elif isinstance(method, stats.MonteCarloMethod): + res = _fisher_exact_monte_carlo_method(table, method) + else: + message = (f'`{method=}` not recognized; if provided, `method` must be an ' + 'instance of `PermutationMethod` or `MonteCarloMethod`.') + raise ValueError(message) + + return SignificanceResult(np.clip(res.statistic, None, 1.0), res.pvalue) + + +def _fisher_exact_permutation_method(table, method): + x, y = _untabulate(table) + colsums = np.sum(table, axis=0) + rowsums = np.sum(table, axis=1) + X = stats.random_table(rowsums, colsums) + + # `permutation_test` with `permutation_type='pairings' permutes the order of `x`, + # which pairs observations in `x` with different observations in `y`. + def statistic(x): + # crosstab the resample and compute the statistic + table = stats.contingency.crosstab(x, y)[1] + return X.pmf(table) + + # tables with *smaller* probability mass are considered to be more extreme + return stats.permutation_test((x,), statistic, permutation_type='pairings', + alternative='less', **method._asdict()) + + +def _fisher_exact_monte_carlo_method(table, method): + method = method._asdict() + + if method.pop('rvs', None) is not None: + message = ('If the `method` argument of `fisher_exact` is an ' + 'instance of `MonteCarloMethod`, its `rvs` attribute ' + 'must be unspecified. Use the `MonteCarloMethod` `rng` argument ' + 'to control the random state.') + raise ValueError(message) + rng = np.random.default_rng(method.pop('rng', None)) + + # `random_table.rvs` produces random contingency tables with the given marginals + # under the null hypothesis of independence + shape = table.shape + colsums = np.sum(table, axis=0) + rowsums = np.sum(table, axis=1) + totsum = np.sum(table) + X = stats.random_table(rowsums, colsums, seed=rng) + + def rvs(size): + n_resamples = size[0] + return X.rvs(size=n_resamples).reshape(size) + + # axis signals to `monte_carlo_test` that statistic is vectorized, but we know + # how it will pass the table(s), so we don't need to use `axis` explicitly. + def statistic(table, axis): + shape_ = (-1,) + shape if table.size > totsum else shape + return X.pmf(table.reshape(shape_)) + + # tables with *smaller* probability mass are considered to be more extreme + return stats.monte_carlo_test(table.ravel(), rvs, statistic, + alternative='less', **method) + + +def _untabulate(table): + # converts a contingency table to paired samples indicating the + # correspondence between row and column indices + r, c = table.shape + x, y = [], [] + for i in range(r): + for j in range(c): + x.append([i] * table[i, j]) + y.append([j] * table[i, j]) + return np.concatenate(x), np.concatenate(y) + + +@xp_capabilities(np_only=True) +def spearmanr(a, b=None, axis=0, nan_policy='propagate', + alternative='two-sided'): + r"""Calculate a Spearman correlation coefficient with associated p-value. + + The Spearman rank-order correlation coefficient is a nonparametric measure + of the monotonicity of the relationship between two datasets. + Like other correlation coefficients, + this one varies between -1 and +1 with 0 implying no correlation. + Correlations of -1 or +1 imply an exact monotonic relationship. Positive + correlations imply that as x increases, so does y. Negative correlations + imply that as x increases, y decreases. + + The p-value roughly indicates the probability of an uncorrelated system + producing datasets that have a Spearman correlation at least as extreme + as the one computed from these datasets. Although calculation of the + p-value does not make strong assumptions about the distributions underlying + the samples, it is only accurate for very large samples (>500 + observations). For smaller sample sizes, consider a permutation test (see + Examples section below). + + Parameters + ---------- + a, b : 1D or 2D array_like, b is optional + One or two 1-D or 2-D arrays containing multiple variables and + observations. When these are 1-D, each represents a vector of + observations of a single variable. For the behavior in the 2-D case, + see under ``axis``, below. + Both arrays need to have the same length in the ``axis`` dimension. + axis : int or None, optional + If axis=0 (default), then each column represents a variable, with + observations in the rows. If axis=1, the relationship is transposed: + each row represents a variable, while the columns contain observations. + If axis=None, then both arrays will be raveled. + nan_policy : {'propagate', 'raise', 'omit'}, optional + Defines how to handle when input contains nan. + The following options are available (default is 'propagate'): + + * 'propagate': returns nan + * 'raise': throws an error + * 'omit': performs the calculations ignoring nan values + + alternative : {'two-sided', 'less', 'greater'}, optional + Defines the alternative hypothesis. Default is 'two-sided'. + The following options are available: + + * 'two-sided': the correlation is nonzero + * 'less': the correlation is negative (less than zero) + * 'greater': the correlation is positive (greater than zero) + + .. versionadded:: 1.7.0 + + Returns + ------- + res : SignificanceResult + An object containing attributes: + + statistic : float or ndarray (2-D square) + Spearman correlation matrix or correlation coefficient (if only 2 + variables are given as parameters). Correlation matrix is square + with length equal to total number of variables (columns or rows) in + ``a`` and ``b`` combined. + pvalue : float + The p-value for a hypothesis test whose null hypothesis + is that two samples have no ordinal correlation. See + `alternative` above for alternative hypotheses. `pvalue` has the + same shape as `statistic`. + + Raises + ------ + ValueError + If `axis` is not 0, 1 or None, or if the number of dimensions of `a` + is greater than 2, or if `b` is None and the number of dimensions of + `a` is less than 2. + + Warns + ----- + `~scipy.stats.ConstantInputWarning` + Raised if an input is a constant array. The correlation coefficient + is not defined in this case, so ``np.nan`` is returned. + + See Also + -------- + :ref:`hypothesis_spearmanr` : Extended example + + References + ---------- + .. [1] Zwillinger, D. and Kokoska, S. (2000). CRC Standard + Probability and Statistics Tables and Formulae. Chapman & Hall: New + York. 2000. + Section 14.7 + .. [2] Kendall, M. G. and Stuart, A. (1973). + The Advanced Theory of Statistics, Volume 2: Inference and Relationship. + Griffin. 1973. + Section 31.18 + + Examples + -------- + + >>> import numpy as np + >>> from scipy import stats + >>> res = stats.spearmanr([1, 2, 3, 4, 5], [5, 6, 7, 8, 7]) + >>> res.statistic + 0.8207826816681233 + >>> res.pvalue + 0.08858700531354381 + + >>> rng = np.random.default_rng() + >>> x2n = rng.standard_normal((100, 2)) + >>> y2n = rng.standard_normal((100, 2)) + >>> res = stats.spearmanr(x2n) + >>> res.statistic, res.pvalue + (-0.07960396039603959, 0.4311168705769747) + + >>> res = stats.spearmanr(x2n[:, 0], x2n[:, 1]) + >>> res.statistic, res.pvalue + (-0.07960396039603959, 0.4311168705769747) + + >>> res = stats.spearmanr(x2n, y2n) + >>> res.statistic + array([[ 1. , -0.07960396, -0.08314431, 0.09662166], + [-0.07960396, 1. , -0.14448245, 0.16738074], + [-0.08314431, -0.14448245, 1. , 0.03234323], + [ 0.09662166, 0.16738074, 0.03234323, 1. ]]) + >>> res.pvalue + array([[0. , 0.43111687, 0.41084066, 0.33891628], + [0.43111687, 0. , 0.15151618, 0.09600687], + [0.41084066, 0.15151618, 0. , 0.74938561], + [0.33891628, 0.09600687, 0.74938561, 0. ]]) + + >>> res = stats.spearmanr(x2n.T, y2n.T, axis=1) + >>> res.statistic + array([[ 1. , -0.07960396, -0.08314431, 0.09662166], + [-0.07960396, 1. , -0.14448245, 0.16738074], + [-0.08314431, -0.14448245, 1. , 0.03234323], + [ 0.09662166, 0.16738074, 0.03234323, 1. ]]) + + >>> res = stats.spearmanr(x2n, y2n, axis=None) + >>> res.statistic, res.pvalue + (0.044981624540613524, 0.5270803651336189) + + >>> res = stats.spearmanr(x2n.ravel(), y2n.ravel()) + >>> res.statistic, res.pvalue + (0.044981624540613524, 0.5270803651336189) + + >>> rng = np.random.default_rng() + >>> xint = rng.integers(10, size=(100, 2)) + >>> res = stats.spearmanr(xint) + >>> res.statistic, res.pvalue + (0.09800224850707953, 0.3320271757932076) + + For small samples, consider performing a permutation test instead of + relying on the asymptotic p-value. Note that to calculate the null + distribution of the statistic (for all possibly pairings between + observations in sample ``x`` and ``y``), only one of the two inputs needs + to be permuted. + + >>> x = [1.76405235, 0.40015721, 0.97873798, + ... 2.2408932, 1.86755799, -0.97727788] + >>> y = [2.71414076, 0.2488, 0.87551913, + ... 2.6514917, 2.01160156, 0.47699563] + + >>> def statistic(x): # permute only `x` + ... return stats.spearmanr(x, y).statistic + >>> res_exact = stats.permutation_test((x,), statistic, + ... permutation_type='pairings') + >>> res_asymptotic = stats.spearmanr(x, y) + >>> res_exact.pvalue, res_asymptotic.pvalue # asymptotic pvalue is too low + (0.10277777777777777, 0.07239650145772594) + + For a more detailed example, see :ref:`hypothesis_spearmanr`. + """ + if axis is not None and axis > 1: + raise ValueError("spearmanr only handles 1-D or 2-D arrays, " + f"supplied axis argument {axis}, please use only " + "values 0, 1 or None for axis") + + a, axisout = _chk_asarray(a, axis) + if a.ndim > 2: + raise ValueError("spearmanr only handles 1-D or 2-D arrays") + + if b is None: + if a.ndim < 2: + raise ValueError("`spearmanr` needs at least 2 " + "variables to compare") + else: + # Concatenate a and b, so that we now only have to handle the case + # of a 2-D `a`. + b, _ = _chk_asarray(b, axis) + if axisout == 0: + a = np.column_stack((a, b)) + else: + a = np.vstack((a, b)) + + n_vars = a.shape[1 - axisout] + n_obs = a.shape[axisout] + if n_obs <= 1: + # Handle empty arrays or single observations. + res = SignificanceResult(np.nan, np.nan) + res.correlation = np.nan + return res + + warn_msg = ("An input array is constant; the correlation coefficient " + "is not defined.") + + constant_axis = False + if axisout == 0: + constant_columns = np.all(a == a[0, :], axis=0) + constant_axis = np.any(constant_columns) + if constant_axis: + # If an input is constant, the correlation coefficient + # is not defined. + warnings.warn(stats.ConstantInputWarning(warn_msg), stacklevel=2) + else: # case when axisout == 1 b/c a is 2 dim only + constant_rows = np.all(a.T == a.T[0, :], axis=0) + constant_axis = np.any(constant_rows) + if constant_axis: + # If an input is constant, the correlation coefficient + # is not defined. + warnings.warn(stats.ConstantInputWarning(warn_msg), stacklevel=2) + + a_contains_nan = _contains_nan(a, nan_policy) + variable_has_nan = np.zeros(n_vars, dtype=bool) + if a_contains_nan: + if nan_policy == 'omit': + return mstats_basic.spearmanr(a, axis=axis, nan_policy=nan_policy, + alternative=alternative) + elif nan_policy == 'propagate': + if a.ndim == 1 or n_vars <= 2: + res = SignificanceResult(np.nan, np.nan) + res.correlation = np.nan + return res + else: + # Keep track of variables with NaNs, set the outputs to NaN + # only for those variables + variable_has_nan = np.isnan(a).any(axis=axisout) + + a_ranked = np.apply_along_axis(rankdata, axisout, a) + + if constant_axis: + with np.errstate(invalid='ignore'): + rs = np.corrcoef(a_ranked, rowvar=axisout) + else: + rs = np.corrcoef(a_ranked, rowvar=axisout) + + dof = n_obs - 2 # degrees of freedom + + # rs can have elements equal to 1, so avoid zero division warnings + with np.errstate(divide='ignore'): + # clip the small negative values possibly caused by rounding + # errors before taking the square root + t = rs * np.sqrt((dof/((rs+1.0)*(1.0-rs))).clip(0)) + + dist = _SimpleStudentT(dof) + prob = _get_pvalue(t, dist, alternative, xp=np) + + # For backwards compatibility, return scalars when comparing 2 columns + if rs.shape == (2, 2): + res = SignificanceResult(rs[1, 0], prob[1, 0]) + res.correlation = rs[1, 0] + return res + else: + rs[variable_has_nan, :] = np.nan + rs[:, variable_has_nan] = np.nan + res = SignificanceResult(rs[()], prob[()]) + res.correlation = rs + return res + + +@xp_capabilities(cpu_only=True, exceptions=['cupy', 'jax.numpy'], marray=True) +@_axis_nan_policy_factory(_pack_CorrelationResult, n_samples=2, + result_to_tuple=_unpack_CorrelationResult, paired=True, + too_small=1, n_outputs=3) +def pointbiserialr(x, y, *, axis=0): + r"""Calculate a point biserial correlation coefficient and its p-value. + + The point biserial correlation is used to measure the relationship + between a binary variable, x, and a continuous variable, y. Like other + correlation coefficients, this one varies between -1 and +1 with 0 + implying no correlation. Correlations of -1 or +1 imply a determinative + relationship. + + This function may be computed using a shortcut formula but produces the + same result as `pearsonr`. + + Parameters + ---------- + x : array_like of bools + Input array. + y : array_like + Input array. + axis : int or None, default: 0 + Axis along which to perform the calculation. Default is 0. + If None, ravel both arrays before performing the calculation. + + Returns + ------- + res: SignificanceResult + An object containing attributes: + + statistic : float + The R value. + pvalue : float + The two-sided p-value. + + Notes + ----- + `pointbiserialr` uses a t-test with ``n-1`` degrees of freedom. + It is equivalent to `pearsonr`. + + The value of the point-biserial correlation can be calculated from: + + .. math:: + + r_{pb} = \frac{\overline{Y_1} - \overline{Y_0}} + {s_y} + \sqrt{\frac{N_0 N_1} + {N (N - 1)}} + + Where :math:`\overline{Y_{0}}` and :math:`\overline{Y_{1}}` are means + of the metric observations coded 0 and 1 respectively; :math:`N_{0}` and + :math:`N_{1}` are number of observations coded 0 and 1 respectively; + :math:`N` is the total number of observations and :math:`s_{y}` is the + standard deviation of all the metric observations. + + A value of :math:`r_{pb}` that is significantly different from zero is + completely equivalent to a significant difference in means between the two + groups. Thus, an independent groups t Test with :math:`N-2` degrees of + freedom may be used to test whether :math:`r_{pb}` is nonzero. The + relation between the t-statistic for comparing two independent groups and + :math:`r_{pb}` is given by: + + .. math:: + + t = \sqrt{N - 2}\frac{r_{pb}}{\sqrt{1 - r^{2}_{pb}}} + + References + ---------- + .. [1] J. Lev, "The Point Biserial Coefficient of Correlation", Ann. Math. + Statist., Vol. 20, no.1, pp. 125-126, 1949. + + .. [2] R.F. Tate, "Correlation Between a Discrete and a Continuous + Variable. Point-Biserial Correlation.", Ann. Math. Statist., Vol. 25, + np. 3, pp. 603-607, 1954. + + .. [3] D. Kornbrot "Point Biserial Correlation", In Wiley StatsRef: + Statistics Reference Online (eds N. Balakrishnan, et al.), 2014. + :doi:`10.1002/9781118445112.stat06227` + + Examples + -------- + >>> import numpy as np + >>> from scipy import stats + >>> a = np.array([0, 0, 0, 1, 1, 1, 1]) + >>> b = np.arange(7) + >>> stats.pointbiserialr(a, b) + (0.8660254037844386, 0.011724811003954652) + >>> stats.pearsonr(a, b) + (0.86602540378443871, 0.011724811003954626) + >>> np.corrcoef(a, b) + array([[ 1. , 0.8660254], + [ 0.8660254, 1. ]]) + + """ + rpb, prob = pearsonr(x, y, axis=axis) + # create result object with alias for backward compatibility + res = SignificanceResult(rpb, prob) + res.correlation = rpb + return res + + +@xp_capabilities(cpu_only=True, jax_jit=False, allow_dask_compute=2, + marray=True, extra_note=("`nan_policy='propagate'` is " + "incompatible with MArray input.")) +def kendalltau(x, y, *, nan_policy='propagate', method='auto', variant='b', + alternative='two-sided', axis=None, keepdims=False): + r"""Calculate Kendall's tau, a correlation measure for ordinal data. + + Kendall's tau is a measure of the correspondence between two rankings. + Values close to 1 indicate strong agreement, and values close to -1 + indicate strong disagreement. This implements two variants of Kendall's + tau: tau-b (the default) and tau-c (also known as Stuart's tau-c). These + differ only in how they are normalized to lie within the range -1 to 1; + the hypothesis tests (their p-values) are identical. Kendall's original + tau-a is not implemented separately because both tau-b and tau-c reduce + to tau-a in the absence of ties. + + Although a naive implementation has O(n^2) complexity, this implementation + uses a Fenwick tree to do the computation in O(n log(n)) complexity. + + Parameters + ---------- + x, y : array_like + Arrays of rankings, of the same shape. If arrays are not 1-D, they + will be flattened to 1-D. + nan_policy : {'propagate', 'raise', 'omit'}, optional + Defines how to handle when input contains nan. + The following options are available (default is 'propagate'): + + * 'propagate': returns nan + * 'raise': throws an error + * 'omit': performs the calculations ignoring nan values + + method : {'auto', 'asymptotic', 'exact'}, optional + Defines which method is used to calculate the p-value [5]_. + The following options are available (default is 'auto'): + + * 'auto': selects the appropriate method based on a trade-off + between speed and accuracy + * 'asymptotic': uses a normal approximation valid for large samples + * 'exact': computes the exact p-value, but can only be used if no ties + are present. As the sample size increases, the 'exact' computation + time may grow and the result may lose some precision. + + variant : {'b', 'c'}, optional + Defines which variant of Kendall's tau is returned. Default is 'b'. + alternative : {'two-sided', 'less', 'greater'}, optional + Defines the alternative hypothesis. Default is 'two-sided'. + The following options are available: + + * 'two-sided': the rank correlation is nonzero + * 'less': the rank correlation is negative (less than zero) + * 'greater': the rank correlation is positive (greater than zero) + axis : int or None, default: None + If an int, the axis of the input along which to compute the statistic. + The statistic of each axis-slice (e.g. row) of the input will appear in a + corresponding element of the output. + If ``None``, the input will be raveled before computing the statistic. + keepdims : bool, optional + If this is set to ``True``, the axes which are reduced are left + in the result as dimensions with length one. With this option, + the result will broadcast correctly against the input array. + + Returns + ------- + res : SignificanceResult + An object containing attributes: + + statistic : float + The tau statistic. + pvalue : float + The p-value for a hypothesis test whose null hypothesis is + an absence of association, tau = 0. + + Raises + ------ + ValueError + If `nan_policy` is 'omit' and `variant` is not 'b' or + if `method` is 'exact' and there are ties between `x` and `y`. + + See Also + -------- + spearmanr : Calculates a Spearman rank-order correlation coefficient. + theilslopes : Computes the Theil-Sen estimator for a set of points (x, y). + weightedtau : Computes a weighted version of Kendall's tau. + :ref:`hypothesis_kendalltau` : Extended example + + Notes + ----- + The definition of Kendall's tau that is used is [2]_:: + + tau_b = (P - Q) / sqrt((P + Q + T) * (P + Q + U)) + + tau_c = 2 (P - Q) / (n**2 * (m - 1) / m) + + where P is the number of concordant pairs, Q the number of discordant + pairs, T the number of tied pairs only in `x`, and U the number of tied pairs only + in `y`. If a tie occurs for the same pair in both `x` and `y`, it is not + added to either T or U. n is the total number of samples, and m is the + number of unique values in either `x` or `y`, whichever is smaller. + + References + ---------- + .. [1] Maurice G. Kendall, "A New Measure of Rank Correlation", Biometrika + Vol. 30, No. 1/2, pp. 81-93, 1938. + .. [2] Maurice G. Kendall, "The treatment of ties in ranking problems", + Biometrika Vol. 33, No. 3, pp. 239-251. 1945. + .. [3] Gottfried E. Noether, "Elements of Nonparametric Statistics", John + Wiley & Sons, 1967. + .. [4] Peter M. Fenwick, "A new data structure for cumulative frequency + tables", Software: Practice and Experience, Vol. 24, No. 3, + pp. 327-336, 1994. + .. [5] Maurice G. Kendall, "Rank Correlation Methods" (4th Edition), + Charles Griffin & Co., 1970. + + Examples + -------- + + >>> from scipy import stats + >>> x1 = [12, 2, 1, 12, 2] + >>> x2 = [1, 4, 7, 1, 0] + >>> res = stats.kendalltau(x1, x2) + >>> res.statistic + -0.47140452079103173 + >>> res.pvalue + 0.2827454599327748 + + For a more detailed example, see :ref:`hypothesis_kendalltau`. + """ + xp = array_namespace(x, y) + x, y = xp_promote(x, y, force_floating=True, xp=xp) + dtype = x.dtype + if is_marray(xp): + mask_x, mask_y = np.asarray(x.mask), np.asarray(y.mask) + x, y = np.asarray(x.data).copy(), np.asarray(y.data).copy() + unmasked_nans = np.any(np.isnan(x[~mask_x])) | np.any(np.isnan(y[~mask_y])) + if unmasked_nans: + if nan_policy == 'propagate': + message = "`nan_policy='propagate'` is incompatible with MArray input." + elif nan_policy == 'raise': + message = "The input contains nan values." + raise ValueError(message) + nan_policy = 'omit' + x[mask_x], y[mask_y] = np.nan, np.nan + else: + x, y = _asarray(x, subok=True, xp=np), _asarray(y, subok=True, xp=np) + res = _kendalltau(x, y, nan_policy=nan_policy, method=method, variant=variant, + alternative=alternative, axis=axis, keepdims=keepdims) + vals = res.statistic, res.pvalue, res.statistic + vals = (xp.asarray(val, dtype=dtype)[()] if val.ndim == 0 + else xp.asarray(val, dtype=dtype) for val in vals) + return _pack_CorrelationResult(*vals) + + +@_axis_nan_policy_factory(_pack_CorrelationResult, default_axis=None, n_samples=2, + result_to_tuple=_unpack_CorrelationResult, paired=True, + too_small=1, n_outputs=3) +def _kendalltau(x, y, *, method='auto', variant='b', alternative='two-sided'): + x = np.asarray(x).ravel() + y = np.asarray(y).ravel() + + if x.size != y.size: + raise ValueError("Array shapes are incompatible for broadcasting.") + elif not x.size or not y.size: + # Return NaN if arrays are empty + NaN = _get_nan(x, y) + res = SignificanceResult(NaN, NaN) + res.correlation = NaN + return res + + def count_rank_tie(ranks): + cnt = np.bincount(ranks).astype('int64', copy=False) + cnt = cnt[cnt > 1] + # Python ints to avoid overflow down the line + return (int((cnt * (cnt - 1) // 2).sum()), + int((cnt * (cnt - 1.) * (cnt - 2)).sum()), + int((cnt * (cnt - 1.) * (2*cnt + 5)).sum())) + + size = x.size + perm = np.argsort(y) # sort on y and convert y to dense ranks + x, y = x[perm], y[perm] + y = np.r_[True, y[1:] != y[:-1]].cumsum(dtype=np.intp) + + # stable sort on x and convert x to dense ranks + perm = np.argsort(x, kind='mergesort') + x, y = x[perm], y[perm] + x = np.r_[True, x[1:] != x[:-1]].cumsum(dtype=np.intp) + + dis = _kendall_dis(x, y) # discordant pairs + + obs = np.r_[True, (x[1:] != x[:-1]) | (y[1:] != y[:-1]), True] + cnt = np.diff(np.nonzero(obs)[0]).astype('int64', copy=False) + + ntie = int((cnt * (cnt - 1) // 2).sum()) # joint ties + xtie, x0, x1 = count_rank_tie(x) # ties in x, stats + ytie, y0, y1 = count_rank_tie(y) # ties in y, stats + + tot = (size * (size - 1)) // 2 + + if xtie == tot or ytie == tot: + NaN = _get_nan(x, y) + res = SignificanceResult(NaN, NaN) + res.correlation = NaN + return res + + # Note that tot = con + dis + (xtie - ntie) + (ytie - ntie) + ntie + # = con + dis + xtie + ytie - ntie + con_minus_dis = tot - xtie - ytie + ntie - 2 * dis + if variant == 'b': + tau = con_minus_dis / np.sqrt(tot - xtie) / np.sqrt(tot - ytie) + elif variant == 'c': + minclasses = min(len(set(x)), len(set(y))) + tau = 2*con_minus_dis / (size**2 * (minclasses-1)/minclasses) + else: + raise ValueError(f"Unknown `variant` '{variant}' specified. " + "`variant` must be 'b' or 'c'.") + + # Limit range to fix computational errors + tau = np.minimum(1., max(-1., tau)) + + # The p-value calculation is the same for all variants since the p-value + # depends only on con_minus_dis. + if method == 'exact' and (xtie != 0 or ytie != 0): + raise ValueError("Ties found; exact method cannot be used.") + + if method == 'auto': + if (xtie == 0 and ytie == 0) and (size <= 33 or + min(dis, tot-dis) <= 1): + method = 'exact' + else: + method = 'asymptotic' + + if xtie == 0 and ytie == 0 and method == 'exact': + pvalue = mstats_basic._kendall_p_exact(size, tot-dis, alternative) + elif method == 'asymptotic': + # con_minus_dis is approx normally distributed with this variance [3]_ + m = size * (size - 1.) + var = ((m * (2*size + 5) - x1 - y1) / 18 + + (2 * xtie * ytie) / m + x0 * y0 / (9 * m * (size - 2))) + z = con_minus_dis / np.sqrt(var) + pvalue = _get_pvalue(z, _SimpleNormal(), alternative, xp=np) + else: + raise ValueError(f"Unknown `method` '{method}' specified. `method` must be" + "'auto', 'exact' or 'asymptotic'.") + + # create result object with alias for backward compatibility + res = SignificanceResult(tau[()], pvalue[()]) + res.correlation = tau[()] + return res + + +def _weightedtau_n_samples(kwargs): + rank = kwargs.get('rank', False) + return 2 if (isinstance(rank, bool) or rank is None) else 3 + + +@xp_capabilities(np_only=True) +@_axis_nan_policy_factory(_pack_CorrelationResult, default_axis=None, + n_samples=_weightedtau_n_samples, + result_to_tuple=_unpack_CorrelationResult, paired=True, + too_small=1, n_outputs=3, override={'nan_propagation': False}) +def weightedtau(x, y, rank=True, weigher=None, additive=True): + r"""Compute a weighted version of Kendall's :math:`\tau`. + + The weighted :math:`\tau` is a weighted version of Kendall's + :math:`\tau` in which exchanges of high weight are more influential than + exchanges of low weight. The default parameters compute the additive + hyperbolic version of the index, :math:`\tau_\mathrm h`, which has + been shown to provide the best balance between important and + unimportant elements [1]_. + + The weighting is defined by means of a rank array, which assigns a + nonnegative rank to each element (higher importance ranks being + associated with smaller values, e.g., 0 is the highest possible rank), + and a weigher function, which assigns a weight based on the rank to + each element. The weight of an exchange is then the sum or the product + of the weights of the ranks of the exchanged elements. The default + parameters compute :math:`\tau_\mathrm h`: an exchange between + elements with rank :math:`r` and :math:`s` (starting from zero) has + weight :math:`1/(r+1) + 1/(s+1)`. + + Specifying a rank array is meaningful only if you have in mind an + external criterion of importance. If, as it usually happens, you do + not have in mind a specific rank, the weighted :math:`\tau` is + defined by averaging the values obtained using the decreasing + lexicographical rank by (`x`, `y`) and by (`y`, `x`). This is the + behavior with default parameters. Note that the convention used + here for ranking (lower values imply higher importance) is opposite + to that used by other SciPy statistical functions. + + Parameters + ---------- + x, y : array_like + Arrays of scores, of the same shape. If arrays are not 1-D, they will + be flattened to 1-D. + rank : array_like of ints or bool, optional + A nonnegative rank assigned to each element. If it is None, the + decreasing lexicographical rank by (`x`, `y`) will be used: elements of + higher rank will be those with larger `x`-values, using `y`-values to + break ties (in particular, swapping `x` and `y` will give a different + result). If it is False, the element indices will be used + directly as ranks. The default is True, in which case this + function returns the average of the values obtained using the + decreasing lexicographical rank by (`x`, `y`) and by (`y`, `x`). + weigher : callable, optional + The weigher function. Must map nonnegative integers (zero + representing the most important element) to a nonnegative weight. + The default, None, provides hyperbolic weighing, that is, + rank :math:`r` is mapped to weight :math:`1/(r+1)`. + additive : bool, optional + If True, the weight of an exchange is computed by adding the + weights of the ranks of the exchanged elements; otherwise, the weights + are multiplied. The default is True. + + Returns + ------- + res: SignificanceResult + An object containing attributes: + + statistic : float + The weighted :math:`\tau` correlation index. + pvalue : float + Presently ``np.nan``, as the null distribution of the statistic is + unknown (even in the additive hyperbolic case). + + See Also + -------- + kendalltau : Calculates Kendall's tau. + spearmanr : Calculates a Spearman rank-order correlation coefficient. + theilslopes : Computes the Theil-Sen estimator for a set of points (x, y). + + Notes + ----- + This function uses an :math:`O(n \log n)`, mergesort-based algorithm + [1]_ that is a weighted extension of Knight's algorithm for Kendall's + :math:`\tau` [2]_. It can compute Shieh's weighted :math:`\tau` [3]_ + between rankings without ties (i.e., permutations) by setting + `additive` and `rank` to False, as the definition given in [1]_ is a + generalization of Shieh's. + + NaNs are considered the smallest possible score. + + .. versionadded:: 0.19.0 + + References + ---------- + .. [1] Sebastiano Vigna, "A weighted correlation index for rankings with + ties", Proceedings of the 24th international conference on World + Wide Web, pp. 1166-1176, ACM, 2015. + .. [2] W.R. Knight, "A Computer Method for Calculating Kendall's Tau with + Ungrouped Data", Journal of the American Statistical Association, + Vol. 61, No. 314, Part 1, pp. 436-439, 1966. + .. [3] Grace S. Shieh. "A weighted Kendall's tau statistic", Statistics & + Probability Letters, Vol. 39, No. 1, pp. 17-24, 1998. + + Examples + -------- + >>> import numpy as np + >>> from scipy import stats + >>> x = [12, 2, 1, 12, 2] + >>> y = [1, 4, 7, 1, 0] + >>> res = stats.weightedtau(x, y) + >>> res.statistic + -0.56694968153682723 + >>> res.pvalue + nan + >>> res = stats.weightedtau(x, y, additive=False) + >>> res.statistic + -0.62205716951801038 + + NaNs are considered the smallest possible score: + + >>> x = [12, 2, 1, 12, 2] + >>> y = [1, 4, 7, 1, np.nan] + >>> res = stats.weightedtau(x, y) + >>> res.statistic + -0.56694968153682723 + + This is exactly Kendall's tau: + + >>> x = [12, 2, 1, 12, 2] + >>> y = [1, 4, 7, 1, 0] + >>> res = stats.weightedtau(x, y, weigher=lambda x: 1) + >>> res.statistic + -0.47140452079103173 + + >>> x = [12, 2, 1, 12, 2] + >>> y = [1, 4, 7, 1, 0] + >>> stats.weightedtau(x, y, rank=None) + SignificanceResult(statistic=-0.4157652301037516, pvalue=nan) + >>> stats.weightedtau(y, x, rank=None) + SignificanceResult(statistic=-0.7181341329699028, pvalue=nan) + + """ + x = np.asarray(x).ravel() + y = np.asarray(y).ravel() + NaN = _get_nan(x, y) + + if x.size != y.size: + raise ValueError("Array shapes are incompatible for broadcasting.") + if not x.size: + # Return NaN if arrays are empty + res = SignificanceResult(NaN, NaN) + res.correlation = NaN + return res + + # If there are NaNs we apply _toint64() + if np.isnan(np.sum(x)): + x = _toint64(x) + if np.isnan(np.sum(y)): + y = _toint64(y) + + # Reduce to ranks unsupported types + if x.dtype != y.dtype: + if x.dtype != np.int64: + x = _toint64(x) + if y.dtype != np.int64: + y = _toint64(y) + else: + if x.dtype not in (np.int32, np.int64, np.float32, np.float64): + x = _toint64(x) + y = _toint64(y) + + if rank is True: + tau = np.asarray( + _weightedrankedtau(x, y, None, weigher, additive) + + _weightedrankedtau(y, x, None, weigher, additive) + )[()] / 2 + res = SignificanceResult(tau, NaN) + res.correlation = tau + return res + + if rank is False: + rank = np.arange(x.size, dtype=np.intp) + elif rank is not None: + rank = np.asarray(rank).ravel() + rank = _toint64(rank).astype(np.intp) + if rank.size != x.size: + raise ValueError( + "All inputs to `weightedtau` must be of the same size, " + f"found x-size {x.size} and rank-size {rank.size}" + ) + + tau = np.asarray(_weightedrankedtau(x, y, rank, weigher, additive))[()] + res = SignificanceResult(tau, NaN) + res.correlation = tau + return res + + +##################################### +# INFERENTIAL STATISTICS # +##################################### + +TtestResultBase = _make_tuple_bunch('TtestResultBase', + ['statistic', 'pvalue'], ['df']) + + +class TtestResult(TtestResultBase): + """ + Result of a t-test. + + See the documentation of the particular t-test function for more + information about the definition of the statistic and meaning of + the confidence interval. + + Attributes + ---------- + statistic : float or array + The t-statistic of the sample. + pvalue : float or array + The p-value associated with the given alternative. + df : float or array + The number of degrees of freedom used in calculation of the + t-statistic; this is one less than the size of the sample + (``a.shape[axis]-1`` if there are no masked elements or omitted NaNs). + + Methods + ------- + confidence_interval + Computes a confidence interval around the population statistic + for the given confidence level. + The confidence interval is returned in a ``namedtuple`` with + fields `low` and `high`. + + """ + + def __init__(self, statistic, pvalue, df, # public + alternative, standard_error, estimate, # private + statistic_np=None, xp=None): # private + super().__init__(statistic, pvalue, df=df) + self._alternative = alternative + self._standard_error = standard_error # denominator of t-statistic + self._estimate = estimate # point estimate of sample mean + self._statistic_np = statistic if statistic_np is None else statistic_np + self._dtype = statistic.dtype + self._xp = array_namespace(statistic, pvalue) if xp is None else xp + + + def confidence_interval(self, confidence_level=0.95): + """ + Parameters + ---------- + confidence_level : float + The confidence level for the calculation of the population mean + confidence interval. Default is 0.95. + + Returns + ------- + ci : namedtuple + The confidence interval is returned in a ``namedtuple`` with + fields `low` and `high`. + + """ + low, high = _t_confidence_interval(self.df, self._statistic_np, + confidence_level, self._alternative, + self._dtype, self._xp) + low = low * self._standard_error + self._estimate + high = high * self._standard_error + self._estimate + return ConfidenceInterval(low=low, high=high) + + +def pack_TtestResult(statistic, pvalue, df, alternative, standard_error, + estimate): + xp = array_namespace(statistic, pvalue) + # Due to behavior of `_axis_nan_policy` decorator, `alternative` can be any number + # of dimensions, but there is at most one unique non-NaN value. + # `_xp_mean` with `nan_policy='omit'` is a JIT-compatible way to extract it. + alternative = xp.asarray(alternative) + alternative = (_xp_mean(alternative, axis=None, nan_policy='omit', warn=False) + if xp_size(alternative) != 0 else xp.nan) + return TtestResult(statistic, pvalue, df=df, alternative=alternative, + standard_error=standard_error, estimate=estimate) + + +def unpack_TtestResult(res, _): + return (res.statistic, res.pvalue, res.df, res._alternative, + res._standard_error, res._estimate) + + +@xp_capabilities(cpu_only=True, exceptions=["cupy", "jax.numpy"], marray=True, + extra_note=("The ``confidence_interval`` method of the output object " + "is incompatible with JAX JIT.")) +@_axis_nan_policy_factory(pack_TtestResult, default_axis=0, n_samples=2, + result_to_tuple=unpack_TtestResult, n_outputs=6) +# nan_policy handled by `_axis_nan_policy`, but needs to be left +# in signature to preserve use as a positional argument +def ttest_1samp(a, popmean, axis=0, nan_policy="propagate", alternative="two-sided"): + """Calculate the T-test for the mean of ONE group of scores. + + This is a test for the null hypothesis that the expected value + (mean) of a sample of independent observations `a` is equal to the given + population mean, `popmean`. + + Parameters + ---------- + a : array_like + Sample observations. + popmean : float or array_like + Expected value in null hypothesis. If array_like, then its length along + `axis` must equal 1, and it must otherwise be broadcastable with `a`. + axis : int or None, optional + Axis along which to compute test; default is 0. If None, compute over + the whole array `a`. + nan_policy : {'propagate', 'raise', 'omit'}, optional + Defines how to handle when input contains nan. + The following options are available (default is 'propagate'): + + * 'propagate': returns nan + * 'raise': throws an error + * 'omit': performs the calculations ignoring nan values + + alternative : {'two-sided', 'less', 'greater'}, optional + Defines the alternative hypothesis. + The following options are available (default is 'two-sided'): + + * 'two-sided': the mean of the underlying distribution of the sample + is different than the given population mean (`popmean`) + * 'less': the mean of the underlying distribution of the sample is + less than the given population mean (`popmean`) + * 'greater': the mean of the underlying distribution of the sample is + greater than the given population mean (`popmean`) + + Returns + ------- + result : `~scipy.stats._result_classes.TtestResult` + An object with the following attributes: + + statistic : float or array + The t-statistic. + pvalue : float or array + The p-value associated with the given alternative. + df : float or array + The number of degrees of freedom used in calculation of the + t-statistic; this is one less than the size of the sample + (``a.shape[axis]``). + + .. versionadded:: 1.10.0 + + The object also has the following method: + + confidence_interval(confidence_level=0.95) + Computes a confidence interval around the population + mean for the given confidence level. + The confidence interval is returned in a ``namedtuple`` with + fields `low` and `high`. + + .. versionadded:: 1.10.0 + + Notes + ----- + The statistic is calculated as ``(np.mean(a) - popmean)/se``, where + ``se`` is the standard error. Therefore, the statistic will be positive + when the sample mean is greater than the population mean and negative when + the sample mean is less than the population mean. + + Examples + -------- + Suppose we wish to test the null hypothesis that the mean of a population + is equal to 0.5. We choose a confidence level of 99%; that is, we will + reject the null hypothesis in favor of the alternative if the p-value is + less than 0.01. + + When testing random variates from the standard uniform distribution, which + has a mean of 0.5, we expect the data to be consistent with the null + hypothesis most of the time. + + >>> import numpy as np + >>> from scipy import stats + >>> rng = np.random.default_rng() + >>> rvs = stats.uniform.rvs(size=50, random_state=rng) + >>> stats.ttest_1samp(rvs, popmean=0.5) + TtestResult(statistic=2.456308468440, pvalue=0.017628209047638, df=49) + + As expected, the p-value of 0.017 is not below our threshold of 0.01, so + we cannot reject the null hypothesis. + + When testing data from the standard *normal* distribution, which has a mean + of 0, we would expect the null hypothesis to be rejected. + + >>> rvs = stats.norm.rvs(size=50, random_state=rng) + >>> stats.ttest_1samp(rvs, popmean=0.5) + TtestResult(statistic=-7.433605518875, pvalue=1.416760157221e-09, df=49) + + Indeed, the p-value is lower than our threshold of 0.01, so we reject the + null hypothesis in favor of the default "two-sided" alternative: the mean + of the population is *not* equal to 0.5. + + However, suppose we were to test the null hypothesis against the + one-sided alternative that the mean of the population is *greater* than + 0.5. Since the mean of the standard normal is less than 0.5, we would not + expect the null hypothesis to be rejected. + + >>> stats.ttest_1samp(rvs, popmean=0.5, alternative='greater') + TtestResult(statistic=-7.433605518875, pvalue=0.99999999929, df=49) + + Unsurprisingly, with a p-value greater than our threshold, we would not + reject the null hypothesis. + + Note that when working with a confidence level of 99%, a true null + hypothesis will be rejected approximately 1% of the time. + + >>> rvs = stats.uniform.rvs(size=(100, 50), random_state=rng) + >>> res = stats.ttest_1samp(rvs, popmean=0.5, axis=1) + >>> np.sum(res.pvalue < 0.01) + 1 + + Indeed, even though all 100 samples above were drawn from the standard + uniform distribution, which *does* have a population mean of 0.5, we would + mistakenly reject the null hypothesis for one of them. + + `ttest_1samp` can also compute a confidence interval around the population + mean. + + >>> rvs = stats.norm.rvs(size=50, random_state=rng) + >>> res = stats.ttest_1samp(rvs, popmean=0) + >>> ci = res.confidence_interval(confidence_level=0.95) + >>> ci + ConfidenceInterval(low=-0.3193887540880017, high=0.2898583388980972) + + The bounds of the 95% confidence interval are the + minimum and maximum values of the parameter `popmean` for which the + p-value of the test would be 0.05. + + >>> res = stats.ttest_1samp(rvs, popmean=ci.low) + >>> np.testing.assert_allclose(res.pvalue, 0.05) + >>> res = stats.ttest_1samp(rvs, popmean=ci.high) + >>> np.testing.assert_allclose(res.pvalue, 0.05) + + Under certain assumptions about the population from which a sample + is drawn, the confidence interval with confidence level 95% is expected + to contain the true population mean in 95% of sample replications. + + >>> rvs = stats.norm.rvs(size=(50, 1000), loc=1, random_state=rng) + >>> res = stats.ttest_1samp(rvs, popmean=0) + >>> ci = res.confidence_interval() + >>> contains_pop_mean = (ci.low < 1) & (ci.high > 1) + >>> contains_pop_mean.sum() + 953 + + """ + xp = array_namespace(a) + a, popmean = xp_promote(a, popmean, force_floating=True, xp=xp) + a, axis = _chk_asarray(a, axis, xp=xp) + + n = _count_nonmasked(a, axis) + df = n - 1 + + if a.shape[axis] == 0: + # This is really only needed for *testing* _axis_nan_policy decorator + # It won't happen when the decorator is used. + NaN = _get_nan(a) + return TtestResult(NaN, NaN, df=NaN, alternative=NaN, + standard_error=NaN, estimate=NaN) + + mean = xp.mean(a, axis=axis) + try: + popmean = xp.asarray(popmean) + popmean = xp.squeeze(popmean, axis=axis) if popmean.ndim > 0 else popmean + except ValueError as e: + raise ValueError("`popmean.shape[axis]` must equal 1.") from e + d = mean - popmean + v = _var(a, axis=axis, ddof=1) + denom = xp.sqrt(v / n) + + with np.errstate(divide='ignore', invalid='ignore'): + t = xp.divide(d, denom) + t = t[()] if t.ndim == 0 else t + + dist = _SimpleStudentT(xp.asarray(df, dtype=t.dtype, device=xp_device(a))) + prob = _get_pvalue(t, dist, alternative, xp=xp) + prob = prob[()] if prob.ndim == 0 else prob + + # when nan_policy='omit', `df` can be different for different axis-slices + df = xp.broadcast_to(xp.asarray(df, device=xp_device(a)), t.shape) + df = df[()] if df.ndim == 0 else df + # _axis_nan_policy decorator doesn't play well with strings + alternative_num = {"less": -1, "two-sided": 0, "greater": 1}[alternative] + return TtestResult(t, prob, df=df, alternative=alternative_num, + standard_error=denom, estimate=mean, + statistic_np=xp.asarray(t), xp=xp) + + +def _t_confidence_interval(df, t, confidence_level, alternative, dtype=None, xp=None): + # Input validation on `alternative` is already done + # We just need IV on confidence_level + dtype = t.dtype if dtype is None else dtype + xp = array_namespace(t) if xp is None else xp + + if confidence_level < 0 or confidence_level > 1: + message = "`confidence_level` must be a number between 0 and 1." + raise ValueError(message) + + confidence_level = xp.asarray(confidence_level, dtype=dtype, device=xp_device(t)) + inf = xp.asarray(xp.inf, dtype=dtype) + + if alternative < 0: # 'less' + p = confidence_level + low, high = xp.broadcast_arrays(-inf, special.stdtrit(df, p)) + elif alternative > 0: # 'greater' + p = 1 - confidence_level + low, high = xp.broadcast_arrays(special.stdtrit(df, p), inf) + elif alternative == 0: # 'two-sided' + tail_probability = (1 - confidence_level)/2 + p = xp.stack((tail_probability, 1-tail_probability)) + # axis of p must be the zeroth and orthogonal to all the rest + p = xp.reshape(p, tuple([2] + [1]*xp.asarray(df, device=xp_device(t)).ndim)) + ci = special.stdtrit(df, p) + low, high = ci[0, ...], ci[1, ...] + else: # alternative is NaN when input is empty (see _axis_nan_policy) + nan = xp.asarray(xp.nan, device=xp_device(t)) + p, nans = xp.broadcast_arrays(t, nan) + low, high = nans, nans + + low = xp.asarray(low, dtype=dtype) + low = low[()] if low.ndim == 0 else low + high = xp.asarray(high, dtype=dtype) + high = high[()] if high.ndim == 0 else high + return low, high + + +def _ttest_ind_from_stats(mean1, mean2, denom, df, alternative, xp=None): + xp = array_namespace(mean1, mean2, denom) if xp is None else xp + + d = mean1 - mean2 + with np.errstate(divide='ignore', invalid='ignore'): + t = xp.divide(d, denom) + + dist = _SimpleStudentT(xp.asarray(df, dtype=t.dtype, device=xp_device(t))) + prob = _get_pvalue(t, dist, alternative, xp=xp) + prob = prob[()] if prob.ndim == 0 else prob + + t = t[()] if t.ndim == 0 else t + prob = prob[()] if prob.ndim == 0 else prob + return t, prob + + +def _unequal_var_ttest_denom(v1, n1, v2, n2, xp=None): + xp = array_namespace(v1, v2) if xp is None else xp + vn1 = v1 / n1 + vn2 = v2 / n2 + with np.errstate(divide='ignore', invalid='ignore'): + df = (vn1 + vn2)**2 / (vn1**2 / (n1 - 1) + vn2**2 / (n2 - 1)) + + # If df is undefined, variances are zero (assumes n1 > 0 & n2 > 0). + # Hence it doesn't matter what df is as long as it's not NaN. + df = xp.where(xp.isnan(df), 1., df) + denom = xp.sqrt(vn1 + vn2) + return df, denom + + +def _equal_var_ttest_denom(v1, n1, v2, n2, xp=None): + xp = array_namespace(v1, v2) if xp is None else xp + + # If there is a single observation in one sample, this formula for pooled + # variance breaks down because the variance of that sample is undefined. + # The pooled variance is still defined, though, because the (n-1) in the + # numerator should cancel with the (n-1) in the denominator, leaving only + # the sum of squared differences from the mean: zero. + v1 = xp.where(xp.asarray(n1 == 1), 0., v1) + v2 = xp.where(xp.asarray(n2 == 1), 0., v2) + + df = n1 + n2 - 2.0 + svar = ((n1 - 1) * v1 + (n2 - 1) * v2) / df + denom = xp.sqrt(svar * (1.0 / n1 + 1.0 / n2)) + df = xp.asarray(df, dtype=denom.dtype) + return df, denom + + +Ttest_indResult = namedtuple('Ttest_indResult', ('statistic', 'pvalue')) + + +@xp_capabilities(cpu_only=True, exceptions=["cupy", "jax.numpy"], marray=True) +def ttest_ind_from_stats(mean1, std1, nobs1, mean2, std2, nobs2, + equal_var=True, alternative="two-sided"): + r""" + T-test for means of two independent samples from descriptive statistics. + + This is a test for the null hypothesis that two independent + samples have identical average (expected) values. + + Parameters + ---------- + mean1 : array_like + The mean(s) of sample 1. + std1 : array_like + The corrected sample standard deviation of sample 1 (i.e. ``ddof=1``). + nobs1 : array_like + The number(s) of observations of sample 1. + mean2 : array_like + The mean(s) of sample 2. + std2 : array_like + The corrected sample standard deviation of sample 2 (i.e. ``ddof=1``). + nobs2 : array_like + The number(s) of observations of sample 2. + equal_var : bool, optional + If True (default), perform a standard independent 2 sample test + that assumes equal population variances [1]_. + If False, perform Welch's t-test, which does not assume equal + population variance [2]_. + alternative : {'two-sided', 'less', 'greater'}, optional + Defines the alternative hypothesis. + The following options are available (default is 'two-sided'): + + * 'two-sided': the means of the distributions are unequal. + * 'less': the mean of the first distribution is less than the + mean of the second distribution. + * 'greater': the mean of the first distribution is greater than the + mean of the second distribution. + + .. versionadded:: 1.6.0 + + Returns + ------- + statistic : float or array + The calculated t-statistics. + pvalue : float or array + The two-tailed p-value. + + See Also + -------- + scipy.stats.ttest_ind + + Notes + ----- + The statistic is calculated as ``(mean1 - mean2)/se``, where ``se`` is the + standard error. Therefore, the statistic will be positive when `mean1` is + greater than `mean2` and negative when `mean1` is less than `mean2`. + + This method does not check whether any of the elements of `std1` or `std2` + are negative. If any elements of the `std1` or `std2` parameters are + negative in a call to this method, this method will return the same result + as if it were passed ``numpy.abs(std1)`` and ``numpy.abs(std2)``, + respectively, instead; no exceptions or warnings will be emitted. + + References + ---------- + .. [1] https://en.wikipedia.org/wiki/T-test#Independent_two-sample_t-test + + .. [2] https://en.wikipedia.org/wiki/Welch%27s_t-test + + Examples + -------- + Suppose we have the summary data for two samples, as follows (with the + Sample Variance being the corrected sample variance):: + + Sample Sample + Size Mean Variance + Sample 1 13 15.0 87.5 + Sample 2 11 12.0 39.0 + + Apply the t-test to this data (with the assumption that the population + variances are equal): + + >>> import numpy as np + >>> from scipy.stats import ttest_ind_from_stats + >>> ttest_ind_from_stats(mean1=15.0, std1=np.sqrt(87.5), nobs1=13, + ... mean2=12.0, std2=np.sqrt(39.0), nobs2=11) + Ttest_indResult(statistic=0.9051358093310269, pvalue=0.3751996797581487) + + For comparison, here is the data from which those summary statistics + were taken. With this data, we can compute the same result using + `scipy.stats.ttest_ind`: + + >>> a = np.array([1, 3, 4, 6, 11, 13, 15, 19, 22, 24, 25, 26, 26]) + >>> b = np.array([2, 4, 6, 9, 11, 13, 14, 15, 18, 19, 21]) + >>> from scipy.stats import ttest_ind + >>> ttest_ind(a, b) + TtestResult(statistic=0.905135809331027, + pvalue=0.3751996797581486, + df=22.0) + + Suppose we instead have binary data and would like to apply a t-test to + compare the proportion of 1s in two independent groups:: + + Number of Sample Sample + Size ones Mean Variance + Sample 1 150 30 0.2 0.161073 + Sample 2 200 45 0.225 0.175251 + + The sample mean :math:`\hat{p}` is the proportion of ones in the sample + and the variance for a binary observation is estimated by + :math:`\hat{p}(1-\hat{p})`. + + >>> ttest_ind_from_stats(mean1=0.2, std1=np.sqrt(0.161073), nobs1=150, + ... mean2=0.225, std2=np.sqrt(0.175251), nobs2=200) + Ttest_indResult(statistic=-0.5627187905196761, pvalue=0.5739887114209541) + + For comparison, we could compute the t statistic and p-value using + arrays of 0s and 1s and `scipy.stat.ttest_ind`, as above. + + >>> group1 = np.array([1]*30 + [0]*(150-30)) + >>> group2 = np.array([1]*45 + [0]*(200-45)) + >>> ttest_ind(group1, group2) + TtestResult(statistic=-0.5627179589855622, + pvalue=0.573989277115258, + df=348.0) + + """ + xp = array_namespace(mean1, std1, mean2, std2) + + mean1, std1, nobs1, mean2, std2, nobs2 = xp_promote( + mean1, std1, nobs1, mean2, std2, nobs2, force_floating=True, xp=xp) + + if equal_var: + df, denom = _equal_var_ttest_denom(std1**2, nobs1, std2**2, nobs2, xp=xp) + else: + df, denom = _unequal_var_ttest_denom(std1**2, nobs1, std2**2, nobs2, xp=xp) + + res = _ttest_ind_from_stats(mean1, mean2, denom, df, alternative) + return Ttest_indResult(*res) + + +@xp_capabilities(cpu_only=True, exceptions=["cupy", "jax.numpy"], marray=True, + extra_note=("`trim` and `method` are incompatible with MArray input." + "The ``confidence_interval`` method of the output object " + "is incompatible with JAX JIT.")) +@_axis_nan_policy_factory(pack_TtestResult, default_axis=0, n_samples=2, + result_to_tuple=unpack_TtestResult, n_outputs=6) +def ttest_ind(a, b, *, axis=0, equal_var=True, nan_policy='propagate', + alternative="two-sided", trim=0, method=None): + """ + Calculate the T-test for the means of *two independent* samples of scores. + + This is a test for the null hypothesis that 2 independent samples + have identical average (expected) values. This test assumes that the + populations have identical variances by default. + + Parameters + ---------- + a, b : array_like + The arrays must have the same shape, except in the dimension + corresponding to `axis` (the first, by default). + axis : int or None, optional + Axis along which to compute test. If None, compute over the whole + arrays, `a`, and `b`. + equal_var : bool, optional + If True (default), perform a standard independent 2 sample test + that assumes equal population variances [1]_. + If False, perform Welch's t-test, which does not assume equal + population variance [2]_. + + .. versionadded:: 0.11.0 + + nan_policy : {'propagate', 'raise', 'omit'}, optional + Defines how to handle when input contains nan. + The following options are available (default is 'propagate'): + + * 'propagate': returns nan + * 'raise': throws an error + * 'omit': performs the calculations ignoring nan values + + The 'omit' option is not currently available for one-sided asymptotic tests. + + alternative : {'two-sided', 'less', 'greater'}, optional + Defines the alternative hypothesis. + The following options are available (default is 'two-sided'): + + * 'two-sided': the means of the distributions underlying the samples + are unequal. + * 'less': the mean of the distribution underlying the first sample + is less than the mean of the distribution underlying the second + sample. + * 'greater': the mean of the distribution underlying the first + sample is greater than the mean of the distribution underlying + the second sample. + + trim : float, optional + If nonzero, performs a trimmed (Yuen's) t-test. + Defines the fraction of elements to be trimmed from each end of the + input samples. If 0 (default), no elements will be trimmed from either + side. The number of trimmed elements from each tail is the floor of the + trim times the number of elements. Valid range is [0, .5). + method : ResamplingMethod, optional + Defines the method used to compute the p-value. If `method` is an + instance of `PermutationMethod`/`MonteCarloMethod`, the p-value is + computed using + `scipy.stats.permutation_test`/`scipy.stats.monte_carlo_test` with the + provided configuration options and other appropriate settings. + Otherwise, the p-value is computed by comparing the test statistic + against a theoretical t-distribution. + + .. versionadded:: 1.15.0 + + Returns + ------- + result : `~scipy.stats._result_classes.TtestResult` + An object with the following attributes: + + statistic : float or ndarray + The t-statistic. + pvalue : float or ndarray + The p-value associated with the given alternative. + df : float or ndarray + The number of degrees of freedom used in calculation of the + t-statistic. + + .. versionadded:: 1.11.0 + + The object also has the following method: + + confidence_interval(confidence_level=0.95) + Computes a confidence interval around the difference in + population means for the given confidence level. + The confidence interval is returned in a ``namedtuple`` with + fields ``low`` and ``high``. + + .. versionadded:: 1.11.0 + + Notes + ----- + Suppose we observe two independent samples, e.g. flower petal lengths, and + we are considering whether the two samples were drawn from the same + population (e.g. the same species of flower or two species with similar + petal characteristics) or two different populations. + + The t-test quantifies the difference between the arithmetic means + of the two samples. The p-value quantifies the probability of observing + as or more extreme values assuming the null hypothesis, that the + samples are drawn from populations with the same population means, is true. + A p-value larger than a chosen threshold (e.g. 5% or 1%) indicates that + our observation is not so unlikely to have occurred by chance. Therefore, + we do not reject the null hypothesis of equal population means. + If the p-value is smaller than our threshold, then we have evidence + against the null hypothesis of equal population means. + + By default, the p-value is determined by comparing the t-statistic of the + observed data against a theoretical t-distribution. + + It is also possible to compute the test statistic using a permutation test by + passing ``method=scipy.stats.PermutationMethod(n_resamples=permutations)``, + where ``permutations`` is the desired number of "permutations" to use in + forming the null distribution. When ``1 < permutations < binom(n, k)``, where + + * ``k`` is the number of observations in `a`, + * ``n`` is the total number of observations in `a` and `b`, and + * ``binom(n, k)`` is the binomial coefficient (``n`` choose ``k``), + + the data are pooled (concatenated), randomly assigned to either group `a` + or `b`, and the t-statistic is calculated. This process is performed + repeatedly (``permutations`` times), generating a distribution of the + t-statistic under the null hypothesis, and the t-statistic of the observed + data is compared to this distribution to determine the p-value. + Specifically, the p-value reported is the "achieved significance level" + (ASL) as defined in 4.4 of [3]_. Note that there are other ways of + estimating p-values using randomized permutation tests; for other + options, see the more general `permutation_test`. + + When ``permutations >= binom(n, k)``, an exact test is performed: the data + are partitioned between the groups in each distinct way exactly once. + + The permutation test can be computationally expensive and not necessarily + more accurate than the analytical test, but it does not make strong + assumptions about the shape of the underlying distribution. + + Use of trimming is commonly referred to as the trimmed t-test. At times + called Yuen's t-test, this is an extension of Welch's t-test, with the + difference being the use of winsorized means in calculation of the variance + and the trimmed sample size in calculation of the statistic. Trimming is + recommended if the underlying distribution is long-tailed or contaminated + with outliers [4]_. + + The statistic is calculated as ``(np.mean(a) - np.mean(b))/se``, where + ``se`` is the standard error. Therefore, the statistic will be positive + when the sample mean of `a` is greater than the sample mean of `b` and + negative when the sample mean of `a` is less than the sample mean of + `b`. + + References + ---------- + .. [1] https://en.wikipedia.org/wiki/T-test#Independent_two-sample_t-test + + .. [2] https://en.wikipedia.org/wiki/Welch%27s_t-test + + .. [3] B. Efron and T. Hastie. + Computer Age Statistical Inference. (2016). + + .. [4] Yuen, Karen K. "The Two-Sample Trimmed t for Unequal Population + Variances." Biometrika, vol. 61, no. 1, 1974, pp. 165-170. + https://www.jstor.org/stable/2334299 + + .. [5] Yuen, Karen K., and W. J. Dixon. "The Approximate Behaviour and + Performance of the Two-Sample Trimmed t." Biometrika, vol. 60, + no. 2, 1973, pp. 369-374. + https://www.jstor.org/stable/2334550 + + Examples + -------- + >>> import numpy as np + >>> from scipy import stats + >>> rng = np.random.default_rng() + + Test with sample with identical means: + + >>> rvs1 = stats.norm.rvs(loc=5, scale=10, size=500, random_state=rng) + >>> rvs2 = stats.norm.rvs(loc=5, scale=10, size=500, random_state=rng) + >>> stats.ttest_ind(rvs1, rvs2) + TtestResult(statistic=-0.4390847099199348, + pvalue=0.6606952038870015, + df=998.0) + >>> stats.ttest_ind(rvs1, rvs2, equal_var=False) + TtestResult(statistic=-0.4390847099199348, + pvalue=0.6606952553131064, + df=997.4602304121448) + + `ttest_ind` underestimates p for unequal variances: + + >>> rvs3 = stats.norm.rvs(loc=5, scale=20, size=500, random_state=rng) + >>> stats.ttest_ind(rvs1, rvs3) + TtestResult(statistic=-1.6370984482905417, + pvalue=0.1019251574705033, + df=998.0) + >>> stats.ttest_ind(rvs1, rvs3, equal_var=False) + TtestResult(statistic=-1.637098448290542, + pvalue=0.10202110497954867, + df=765.1098655246868) + + When ``n1 != n2``, the equal variance t-statistic is no longer equal to the + unequal variance t-statistic: + + >>> rvs4 = stats.norm.rvs(loc=5, scale=20, size=100, random_state=rng) + >>> stats.ttest_ind(rvs1, rvs4) + TtestResult(statistic=-1.9481646859513422, + pvalue=0.05186270935842703, + df=598.0) + >>> stats.ttest_ind(rvs1, rvs4, equal_var=False) + TtestResult(statistic=-1.3146566100751664, + pvalue=0.1913495266513811, + df=110.41349083985212) + + T-test with different means, variance, and n: + + >>> rvs5 = stats.norm.rvs(loc=8, scale=20, size=100, random_state=rng) + >>> stats.ttest_ind(rvs1, rvs5) + TtestResult(statistic=-2.8415950600298774, + pvalue=0.0046418707568707885, + df=598.0) + >>> stats.ttest_ind(rvs1, rvs5, equal_var=False) + TtestResult(statistic=-1.8686598649188084, + pvalue=0.06434714193919686, + df=109.32167496550137) + + Take these two samples, one of which has an extreme tail. + + >>> a = (56, 128.6, 12, 123.8, 64.34, 78, 763.3) + >>> b = (1.1, 2.9, 4.2) + + Use the `trim` keyword to perform a trimmed (Yuen) t-test. For example, + using 20% trimming, ``trim=.2``, the test will reduce the impact of one + (``np.floor(trim*len(a))``) element from each tail of sample `a`. It will + have no effect on sample `b` because ``np.floor(trim*len(b))`` is 0. + + >>> stats.ttest_ind(a, b, trim=.2) + TtestResult(statistic=3.4463884028073513, + pvalue=0.01369338726499547, + df=6.0) + """ + xp = array_namespace(a, b) + + a, b = xp_promote(a, b, force_floating=True, xp=xp) + + if axis is None: + a, b, axis = xp_ravel(a), xp_ravel(b), 0 + + if not (0 <= trim < .5): + raise ValueError("Trimming percentage should be 0 <= `trim` < .5.") + + if not isinstance(method, PermutationMethod | MonteCarloMethod | None): + message = ("`method` must be an instance of `PermutationMethod`, an instance " + "of `MonteCarloMethod`, or None (default).") + raise ValueError(message) + + result_shape = _broadcast_array_shapes_remove_axis((a, b), axis=axis) + NaN = _get_nan(a, b, shape=result_shape, xp=xp) + if xp_size(a) == 0 or xp_size(b) == 0: + return TtestResult(NaN, NaN, df=NaN, alternative=NaN, + standard_error=NaN, estimate=NaN) + + alternative_nums = {"less": -1, "two-sided": 0, "greater": 1} + + n1 = _count_nonmasked(a, axis) + n2 = _count_nonmasked(b, axis) + + if trim == 0: + with np.errstate(divide='ignore', invalid='ignore'): + v1 = _var(a, axis, ddof=1, xp=xp) + v2 = _var(b, axis, ddof=1, xp=xp) + + m1 = xp.mean(a, axis=axis) + m2 = xp.mean(b, axis=axis) + else: + v1, m1, n1 = _ttest_trim_var_mean_len(a, trim, axis, xp=xp) + v2, m2, n2 = _ttest_trim_var_mean_len(b, trim, axis, xp=xp) + + if equal_var: + df, denom = _equal_var_ttest_denom(v1, n1, v2, n2, xp=xp) + else: + df, denom = _unequal_var_ttest_denom(v1, n1, v2, n2, xp=xp) + + if method is None: + t, prob = _ttest_ind_from_stats(m1, m2, denom, df, alternative) + else: + # nan_policy is taken care of by axis_nan_policy decorator + ttest_kwargs = dict(equal_var=equal_var, trim=trim) + t, prob = _ttest_resampling(a, b, axis, alternative, + ttest_kwargs, method, xp=xp) + + # when nan_policy='omit', `df` can be different for different axis-slices + df = xp.broadcast_to(df, t.shape) + df = df[()] if df.ndim ==0 else df + estimate = m1 - m2 + + return TtestResult(t, prob, df=df, alternative=alternative_nums[alternative], + standard_error=denom, estimate=estimate) + + +def _ttest_resampling(x, y, axis, alternative, ttest_kwargs, method, *, xp): + def statistic(x, y, axis): + x, y = xp.asarray(x), xp.asarray(y) + return ttest_ind(x, y, axis=axis, **ttest_kwargs).statistic + + test = (permutation_test if isinstance(method, PermutationMethod) + else monte_carlo_test) + method = method._asdict() + + if test is monte_carlo_test: + # `monte_carlo_test` accepts an `rvs` tuple of callables, not an `rng` + # If the user specified an `rng`, replace it with the default callables + if (rng := method.pop('rng', None)) is not None: + rng = np.random.default_rng(rng) + method['rvs'] = rng.normal, rng.normal + + res = test((x, y,), statistic=statistic, axis=axis, + alternative=alternative, **method) + + return res.statistic, res.pvalue + + +def _ttest_trim_var_mean_len(a, trim, axis, *, xp): + """Variance, mean, and length of winsorized input along specified axis""" + # for use with `ttest_ind` when trimming. + # further calculations in this test assume that the inputs are sorted. + # From [4] Section 1 "Let x_1, ..., x_n be n ordered observations..." + a = xp.sort(a, axis=axis) + + # `g` is the number of elements to be replaced on each tail, converted + # from a percentage amount of trimming + n = a.shape[axis] + g = int(n * trim) + + # Calculate the Winsorized variance of the input samples according to + # specified `g` + v = _calculate_winsorized_variance(a, g, axis, xp=xp) + + # the total number of elements in the trimmed samples + n -= 2 * g + + # calculate the g-times trimmed mean, as defined in [4] (1-1) + m = trim_mean(a, trim, axis=axis) + return v, m, n + + +def _calculate_winsorized_variance(a, g, axis, *, xp): + """Calculates g-times winsorized variance along specified axis""" + # it is expected that the input `a` is sorted along the correct axis + if g == 0: + return _var(a, ddof=1, axis=axis, xp=xp) + # move the intended axis to the end that way it is easier to manipulate + a_win = xp.moveaxis(a, axis, -1) + + # save where NaNs are for later use. + nans_indices = xp.any(xp.isnan(a_win), axis=-1) + + # Winsorization and variance calculation are done in one step in [4] + # (1-3), but here winsorization is done first; replace the left and + # right sides with the repeating value. This can be see in effect in ( + # 1-3) in [4], where the leftmost and rightmost tails are replaced with + # `(g + 1) * x_{g + 1}` on the left and `(g + 1) * x_{n - g}` on the + # right. Zero-indexing turns `g + 1` to `g`, and `n - g` to `- g - 1` in + # array indexing. + a_win = xpx.at(a_win)[..., :g].set(a_win[..., g:g+1]) + a_win = xpx.at(a_win)[..., -g:].set(a_win[..., -g - 1:-g]) + + # Determine the variance. In [4], the degrees of freedom is expressed as + # `h - 1`, where `h = n - 2g` (unnumbered equations in Section 1, end of + # page 369, beginning of page 370). This is converted to NumPy's format, + # `n - ddof` for use with `np.var`. The result is converted to an + # array to accommodate indexing later. + var_win = xp.asarray(_var(a_win, ddof=(2 * g + 1), axis=-1, xp=xp)) + + # with `nan_policy='propagate'`, NaNs may be completely trimmed out + # because they were sorted into the tail of the array. In these cases, + # replace computed variances with `np.nan`. + var_win = xpx.at(var_win)[nans_indices].set(xp.nan) + return var_win + + +@xp_capabilities(cpu_only=True, exceptions=["cupy", "jax.numpy"], marray=True, + extra_note=("The ``confidence_interval`` method of the output object" + "is incompatible with JAX JIT.")) +@_axis_nan_policy_factory(pack_TtestResult, default_axis=0, n_samples=2, + result_to_tuple=unpack_TtestResult, n_outputs=6, + paired=True) +def ttest_rel(a, b, axis=0, nan_policy='propagate', alternative="two-sided"): + """Calculate the t-test on TWO RELATED samples of scores, a and b. + + This is a test for the null hypothesis that two related or + repeated samples have identical average (expected) values. + + Parameters + ---------- + a, b : array_like + The arrays must have the same shape. + axis : int or None, optional + Axis along which to compute test. If None, compute over the whole + arrays, `a`, and `b`. + nan_policy : {'propagate', 'raise', 'omit'}, optional + Defines how to handle when input contains nan. + The following options are available (default is 'propagate'): + + * 'propagate': returns nan + * 'raise': throws an error + * 'omit': performs the calculations ignoring nan values + + alternative : {'two-sided', 'less', 'greater'}, optional + Defines the alternative hypothesis. + The following options are available (default is 'two-sided'): + + * 'two-sided': the means of the distributions underlying the samples + are unequal. + * 'less': the mean of the distribution underlying the first sample + is less than the mean of the distribution underlying the second + sample. + * 'greater': the mean of the distribution underlying the first + sample is greater than the mean of the distribution underlying + the second sample. + + .. versionadded:: 1.6.0 + + Returns + ------- + result : `~scipy.stats._result_classes.TtestResult` + An object with the following attributes: + + statistic : float or array + The t-statistic. + pvalue : float or array + The p-value associated with the given alternative. + df : float or array + The number of degrees of freedom used in calculation of the + t-statistic; this is one less than the size of the sample + (``a.shape[axis]``). + + .. versionadded:: 1.10.0 + + The object also has the following method: + + confidence_interval(confidence_level=0.95) + Computes a confidence interval around the difference in + population means for the given confidence level. + The confidence interval is returned in a ``namedtuple`` with + fields `low` and `high`. + + .. versionadded:: 1.10.0 + + Notes + ----- + Examples for use are scores of the same set of student in + different exams, or repeated sampling from the same units. The + test measures whether the average score differs significantly + across samples (e.g. exams). If we observe a large p-value, for + example greater than 0.05 or 0.1 then we cannot reject the null + hypothesis of identical average scores. If the p-value is smaller + than the threshold, e.g. 1%, 5% or 10%, then we reject the null + hypothesis of equal averages. Small p-values are associated with + large t-statistics. + + The t-statistic is calculated as ``np.mean(a - b)/se``, where ``se`` is the + standard error. Therefore, the t-statistic will be positive when the sample + mean of ``a - b`` is greater than zero and negative when the sample mean of + ``a - b`` is less than zero. + + References + ---------- + https://en.wikipedia.org/wiki/T-test#Dependent_t-test_for_paired_samples + + Examples + -------- + >>> import numpy as np + >>> from scipy import stats + >>> rng = np.random.default_rng() + + >>> rvs1 = stats.norm.rvs(loc=5, scale=10, size=500, random_state=rng) + >>> rvs2 = (stats.norm.rvs(loc=5, scale=10, size=500, random_state=rng) + ... + stats.norm.rvs(scale=0.2, size=500, random_state=rng)) + >>> stats.ttest_rel(rvs1, rvs2) + TtestResult(statistic=-0.4549717054410304, pvalue=0.6493274702088672, df=499) + >>> rvs3 = (stats.norm.rvs(loc=8, scale=10, size=500, random_state=rng) + ... + stats.norm.rvs(scale=0.2, size=500, random_state=rng)) + >>> stats.ttest_rel(rvs1, rvs3) + TtestResult(statistic=-5.879467544540889, pvalue=7.540777129099917e-09, df=499) + + """ + return ttest_1samp(a - b, popmean=0., axis=axis, alternative=alternative, + _no_deco=True) + + +# Map from names to lambda_ values used in power_divergence(). +_power_div_lambda_names = { + "pearson": 1, + "log-likelihood": 0, + "freeman-tukey": -0.5, + "mod-log-likelihood": -1, + "neyman": -2, + "cressie-read": 2/3, +} + + +Power_divergenceResult = namedtuple('Power_divergenceResult', + ('statistic', 'pvalue')) + + +def _pd_nsamples(kwargs): + return 2 if kwargs.get('f_exp', None) is not None else 1 + + +@xp_capabilities(marray=True) +@_axis_nan_policy_factory(Power_divergenceResult, paired=True, n_samples=_pd_nsamples, + too_small=-1) +def power_divergence(f_obs, f_exp=None, ddof=0, axis=0, lambda_=None): + """Cressie-Read power divergence statistic and goodness of fit test. + + This function tests the null hypothesis that the categorical data + has the given frequencies, using the Cressie-Read power divergence + statistic. + + Parameters + ---------- + f_obs : array_like + Observed frequencies in each category. + f_exp : array_like, optional + Expected frequencies in each category. By default the categories are + assumed to be equally likely. + ddof : int, optional + "Delta degrees of freedom": adjustment to the degrees of freedom + for the p-value. The p-value is computed using a chi-squared + distribution with ``k - 1 - ddof`` degrees of freedom, where `k` + is the number of observed frequencies. The default value of `ddof` + is 0. + axis : int or None, optional + The axis of the broadcast result of `f_obs` and `f_exp` along which to + apply the test. If axis is None, all values in `f_obs` are treated + as a single data set. Default is 0. + lambda_ : float or str, optional + The power in the Cressie-Read power divergence statistic. The default + is 1. For convenience, `lambda_` may be assigned one of the following + strings, in which case the corresponding numerical value is used: + + * ``"pearson"`` (value 1) + Pearson's chi-squared statistic. In this case, the function is + equivalent to `chisquare`. + * ``"log-likelihood"`` (value 0) + Log-likelihood ratio. Also known as the G-test [3]_. + * ``"freeman-tukey"`` (value -1/2) + Freeman-Tukey statistic. + * ``"mod-log-likelihood"`` (value -1) + Modified log-likelihood ratio. + * ``"neyman"`` (value -2) + Neyman's statistic. + * ``"cressie-read"`` (value 2/3) + The power recommended in [5]_. + + Returns + ------- + res: Power_divergenceResult + An object containing attributes: + + statistic : float or ndarray + The Cressie-Read power divergence test statistic. The value is + a float if `axis` is None or if` `f_obs` and `f_exp` are 1-D. + pvalue : float or ndarray + The p-value of the test. The value is a float if `ddof` and the + return value `stat` are scalars. + + See Also + -------- + chisquare + + Notes + ----- + This test is invalid when the observed or expected frequencies in each + category are too small. A typical rule is that all of the observed + and expected frequencies should be at least 5. + + Also, the sum of the observed and expected frequencies must be the same + for the test to be valid; `power_divergence` raises an error if the sums + do not agree within a relative tolerance of ``eps**0.5``, where ``eps`` + is the precision of the input dtype. + + When `lambda_` is less than zero, the formula for the statistic involves + dividing by `f_obs`, so a warning or error may be generated if any value + in `f_obs` is 0. + + Similarly, a warning or error may be generated if any value in `f_exp` is + zero when `lambda_` >= 0. + + The default degrees of freedom, k-1, are for the case when no parameters + of the distribution are estimated. If p parameters are estimated by + efficient maximum likelihood then the correct degrees of freedom are + k-1-p. If the parameters are estimated in a different way, then the + dof can be between k-1-p and k-1. However, it is also possible that + the asymptotic distribution is not a chisquare, in which case this + test is not appropriate. + + References + ---------- + .. [1] Lowry, Richard. "Concepts and Applications of Inferential + Statistics". Chapter 8. + https://web.archive.org/web/20171015035606/http://faculty.vassar.edu/lowry/ch8pt1.html + .. [2] "Chi-squared test", https://en.wikipedia.org/wiki/Chi-squared_test + .. [3] "G-test", https://en.wikipedia.org/wiki/G-test + .. [4] Sokal, R. R. and Rohlf, F. J. "Biometry: the principles and + practice of statistics in biological research", New York: Freeman + (1981) + .. [5] Cressie, N. and Read, T. R. C., "Multinomial Goodness-of-Fit + Tests", J. Royal Stat. Soc. Series B, Vol. 46, No. 3 (1984), + pp. 440-464. + + Examples + -------- + (See `chisquare` for more examples.) + + When just `f_obs` is given, it is assumed that the expected frequencies + are uniform and given by the mean of the observed frequencies. Here we + perform a G-test (i.e. use the log-likelihood ratio statistic): + + >>> import numpy as np + >>> from scipy.stats import power_divergence + >>> power_divergence([16, 18, 16, 14, 12, 12], lambda_='log-likelihood') + (2.006573162632538, 0.84823476779463769) + + The expected frequencies can be given with the `f_exp` argument: + + >>> power_divergence([16, 18, 16, 14, 12, 12], + ... f_exp=[16, 16, 16, 16, 16, 8], + ... lambda_='log-likelihood') + (3.3281031458963746, 0.6495419288047497) + + When `f_obs` is 2-D, by default the test is applied to each column. + + >>> obs = np.array([[16, 18, 16, 14, 12, 12], [32, 24, 16, 28, 20, 24]]).T + >>> obs.shape + (6, 2) + >>> power_divergence(obs, lambda_="log-likelihood") + (array([ 2.00657316, 6.77634498]), array([ 0.84823477, 0.23781225])) + + By setting ``axis=None``, the test is applied to all data in the array, + which is equivalent to applying the test to the flattened array. + + >>> power_divergence(obs, axis=None) + (23.31034482758621, 0.015975692534127565) + >>> power_divergence(obs.ravel()) + (23.31034482758621, 0.015975692534127565) + + `ddof` is the change to make to the default degrees of freedom. + + >>> power_divergence([16, 18, 16, 14, 12, 12], ddof=1) + (2.0, 0.73575888234288467) + + The calculation of the p-values is done by broadcasting the + test statistic with `ddof`. + + >>> power_divergence([16, 18, 16, 14, 12, 12], ddof=[0,1,2]) + (2.0, array([ 0.84914504, 0.73575888, 0.5724067 ])) + + `f_obs` and `f_exp` are also broadcast. In the following, `f_obs` has + shape (6,) and `f_exp` has shape (2, 6), so the result of broadcasting + `f_obs` and `f_exp` has shape (2, 6). To compute the desired chi-squared + statistics, we must use ``axis=1``: + + >>> power_divergence([16, 18, 16, 14, 12, 12], + ... f_exp=[[16, 16, 16, 16, 16, 8], + ... [8, 20, 20, 16, 12, 12]], + ... axis=1) + (array([ 3.5 , 9.25]), array([ 0.62338763, 0.09949846])) + + """ + return _power_divergence(f_obs, f_exp=f_exp, ddof=ddof, axis=axis, lambda_=lambda_) + + +def _power_divergence(f_obs, f_exp, ddof, axis, lambda_, sum_check=True): + xp = array_namespace(f_obs, f_exp, ddof) + f_obs, f_exp, ddof = xp_promote(f_obs, f_exp, ddof, + force_floating=True, xp=xp) + + # Convert the input argument `lambda_` to a numerical value. + if isinstance(lambda_, str): + if lambda_ not in _power_div_lambda_names: + names = repr(list(_power_div_lambda_names.keys()))[1:-1] + raise ValueError(f"invalid string for lambda_: {lambda_!r}. " + f"Valid strings are {names}") + lambda_ = _power_div_lambda_names[lambda_] + elif lambda_ is None: + lambda_ = 1 + + if f_exp is not None: + # not sure why we force to float64, but not going to touch it + f_obs_float = xp.asarray(f_obs, dtype=xp.float64) + bshape = _broadcast_shapes((f_obs_float.shape, f_exp.shape)) + f_obs_float = xp.broadcast_to(f_obs_float, bshape) + f_exp = xp.broadcast_to(f_exp, bshape) + f_obs_float, f_exp = _share_masks(f_obs_float, f_exp, xp=xp) + + if sum_check: + dtype_res = xp.result_type(f_obs.dtype, f_exp.dtype) + rtol = xp.finfo(dtype_res).eps**0.5 # to pass existing tests + with np.errstate(invalid='ignore'): + f_obs_sum = xp.sum(f_obs_float, axis=axis, keepdims=True) + f_exp_sum = xp.sum(f_exp, axis=axis, keepdims=True) + relative_diff = (xp.abs(f_obs_sum - f_exp_sum) / + xp.minimum(f_obs_sum, f_exp_sum)) + diff_gt_tol = xp.any(relative_diff > rtol, axis=axis, keepdims=True) + + if not is_lazy_array(diff_gt_tol) and xp.any(diff_gt_tol): + msg = (f"For each axis slice, the sum of the observed " + f"frequencies must agree with the sum of the " + f"expected frequencies to a relative tolerance " + f"of {rtol}, but the percent differences are:\n" + f"{relative_diff}") + raise ValueError(msg) + elif is_lazy_array(diff_gt_tol): + diff_gt_tol = xp.broadcast_to(diff_gt_tol, f_obs.shape) + f_obs = xpx.at(f_obs)[diff_gt_tol].set(xp.nan) + + else: + # Avoid warnings with the edge case of a data set with length 0 + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + f_exp = xp.mean(f_obs, axis=axis, keepdims=True) + + # `terms` is the array of terms that are summed along `axis` to create + # the test statistic. We use some specialized code for a few special + # cases of lambda_. + if lambda_ == 1: + # Pearson's chi-squared statistic + terms = (f_obs - f_exp)**2 / f_exp + elif lambda_ == 0: + # Log-likelihood ratio (i.e. G-test) + terms = 2.0 * special.xlogy(f_obs, f_obs / f_exp) + elif lambda_ == -1: + # Modified log-likelihood ratio + terms = 2.0 * special.xlogy(f_exp, f_exp / f_obs) + else: + # General Cressie-Read power divergence. + terms = f_obs * ((f_obs / f_exp)**lambda_ - 1) + terms /= 0.5 * lambda_ * (lambda_ + 1) + + stat = xp.sum(terms, axis=axis) + + num_obs = xp.asarray(_count_nonmasked(terms, axis), device=xp_device(terms), + dtype=f_obs.dtype) + + df = num_obs - 1 - ddof + chi2 = _SimpleChi2(df) + pvalue = _get_pvalue(stat, chi2 , alternative='greater', symmetric=False, xp=xp) + + stat = stat[()] if stat.ndim == 0 else stat + pvalue = pvalue[()] if pvalue.ndim == 0 else pvalue + + return Power_divergenceResult(stat, pvalue) + + + +@xp_capabilities(marray=True) +@_axis_nan_policy_factory(Power_divergenceResult, paired=True, n_samples=_pd_nsamples, + too_small=-1) +def chisquare(f_obs, f_exp=None, ddof=0, axis=0, *, sum_check=True): + """Perform Pearson's chi-squared test. + + Pearson's chi-squared test [1]_ is a goodness-of-fit test for a multinomial + distribution with given probabilities; that is, it assesses the null hypothesis + that the observed frequencies (counts) are obtained by independent + sampling of *N* observations from a categorical distribution with given + expected frequencies. + + Parameters + ---------- + f_obs : array_like + Observed frequencies in each category. + f_exp : array_like, optional + Expected frequencies in each category. By default, the categories are + assumed to be equally likely. + ddof : int, optional + "Delta degrees of freedom": adjustment to the degrees of freedom + for the p-value. The p-value is computed using a chi-squared + distribution with ``k - 1 - ddof`` degrees of freedom, where ``k`` + is the number of categories. The default value of `ddof` is 0. + axis : int or None, optional + The axis of the broadcast result of `f_obs` and `f_exp` along which to + apply the test. If axis is None, all values in `f_obs` are treated + as a single data set. Default is 0. + sum_check : bool, optional + Whether to perform a check that ``sum(f_obs) - sum(f_exp) == 0``. If True, + (default) raise an error (or, for lazy backends, return NaN) when the relative + difference exceeds the square root of the precision of the data type. + See Notes for rationale and possible exceptions. + + Returns + ------- + res: Power_divergenceResult + An object containing attributes: + + statistic : float or ndarray + The chi-squared test statistic. The value is a float if `axis` is + None or `f_obs` and `f_exp` are 1-D. + pvalue : float or ndarray + The p-value of the test. The value is a float if `ddof` and the + result attribute `statistic` are scalars. + + See Also + -------- + scipy.stats.power_divergence + scipy.stats.fisher_exact : Fisher exact test on a 2x2 contingency table. + scipy.stats.barnard_exact : An unconditional exact test. An alternative + to chi-squared test for small sample sizes. + :ref:`hypothesis_chisquare` : Extended example + + Notes + ----- + This test is invalid when the observed or expected frequencies in each + category are too small. A typical rule is that all of the observed + and expected frequencies should be at least 5. According to [2]_, the + total number of observations is recommended to be greater than 13, + otherwise exact tests (such as Barnard's Exact test) should be used + because they do not overreject. + + The default degrees of freedom, k-1, are for the case when no parameters + of the distribution are estimated. If p parameters are estimated by + efficient maximum likelihood then the correct degrees of freedom are + k-1-p. If the parameters are estimated in a different way, then the + dof can be between k-1-p and k-1. However, it is also possible that + the asymptotic distribution is not chi-square, in which case this test + is not appropriate. + + For Pearson's chi-squared test, the total observed and expected counts must match + for the p-value to accurately reflect the probability of observing such an extreme + value of the statistic under the null hypothesis. + This function may be used to perform other statistical tests that do not require + the total counts to be equal. For instance, to test the null hypothesis that + ``f_obs[i]`` is Poisson-distributed with expectation ``f_exp[i]``, set ``ddof=-1`` + and ``sum_check=False``. This test follows from the fact that a Poisson random + variable with mean and variance ``f_exp[i]`` is approximately normal with the + same mean and variance; the chi-squared statistic standardizes, squares, and sums + the observations; and the sum of ``n`` squared standard normal variables follows + the chi-squared distribution with ``n`` degrees of freedom. + + References + ---------- + .. [1] "Pearson's chi-squared test". + *Wikipedia*. https://en.wikipedia.org/wiki/Pearson%27s_chi-squared_test + .. [2] Pearson, Karl. "On the criterion that a given system of deviations from the probable + in the case of a correlated system of variables is such that it can be reasonably + supposed to have arisen from random sampling", Philosophical Magazine. Series 5. 50 + (1900), pp. 157-175. + + Examples + -------- + When only the mandatory `f_obs` argument is given, it is assumed that the + expected frequencies are uniform and given by the mean of the observed + frequencies: + + >>> import numpy as np + >>> from scipy.stats import chisquare + >>> chisquare([16, 18, 16, 14, 12, 12]) + Power_divergenceResult(statistic=2.0, pvalue=0.84914503608460956) + + The optional `f_exp` argument gives the expected frequencies. + + >>> chisquare([16, 18, 16, 14, 12, 12], f_exp=[16, 16, 16, 16, 16, 8]) + Power_divergenceResult(statistic=3.5, pvalue=0.62338762774958223) + + When `f_obs` is 2-D, by default the test is applied to each column. + + >>> obs = np.array([[16, 18, 16, 14, 12, 12], [32, 24, 16, 28, 20, 24]]).T + >>> obs.shape + (6, 2) + >>> chisquare(obs) + Power_divergenceResult(statistic=array([2. , 6.66666667]), pvalue=array([0.84914504, 0.24663415])) + + By setting ``axis=None``, the test is applied to all data in the array, + which is equivalent to applying the test to the flattened array. + + >>> chisquare(obs, axis=None) + Power_divergenceResult(statistic=23.31034482758621, pvalue=0.015975692534127565) + >>> chisquare(obs.ravel()) + Power_divergenceResult(statistic=23.310344827586206, pvalue=0.01597569253412758) + + `ddof` is the change to make to the default degrees of freedom. + + >>> chisquare([16, 18, 16, 14, 12, 12], ddof=1) + Power_divergenceResult(statistic=2.0, pvalue=0.7357588823428847) + + The calculation of the p-values is done by broadcasting the + chi-squared statistic with `ddof`. + + >>> chisquare([16, 18, 16, 14, 12, 12], ddof=[0, 1, 2]) + Power_divergenceResult(statistic=2.0, pvalue=array([0.84914504, 0.73575888, 0.5724067 ])) + + `f_obs` and `f_exp` are also broadcast. In the following, `f_obs` has + shape (6,) and `f_exp` has shape (2, 6), so the result of broadcasting + `f_obs` and `f_exp` has shape (2, 6). To compute the desired chi-squared + statistics, we use ``axis=1``: + + >>> chisquare([16, 18, 16, 14, 12, 12], + ... f_exp=[[16, 16, 16, 16, 16, 8], [8, 20, 20, 16, 12, 12]], + ... axis=1) + Power_divergenceResult(statistic=array([3.5 , 9.25]), pvalue=array([0.62338763, 0.09949846])) + + For a more detailed example, see :ref:`hypothesis_chisquare`. + """ # noqa: E501 + return _power_divergence(f_obs, f_exp=f_exp, ddof=ddof, axis=axis, + lambda_="pearson", sum_check=sum_check) + + +KstestResult = _make_tuple_bunch('KstestResult', ['statistic', 'pvalue'], + ['statistic_location', 'statistic_sign']) + + +def _compute_d(cdfvals, x, sign, xp=None): + """Computes D+/D- as used in the Kolmogorov-Smirnov test. + + Vectorized along the last axis. + + Parameters + ---------- + cdfvals : array_like + Sorted array of CDF values between 0 and 1 + x: array_like + Sorted array of the stochastic variable itself + sign: int + Indicates whether to compute D+ (+1) or D- (-1). + + Returns + ------- + D : float or array + The maximum distance of the CDF values below/above (D+/D-) Uniform(0, 1). + loc_max : float or array + The location at which the maximum is reached. + """ + xp = array_namespace(cdfvals, x) if xp is None else xp + length = cdfvals.shape[-1] + n = _count_nonmasked(cdfvals, axis=-1, keepdims=True) + D = (xp.arange(1.0, length + 1, dtype=x.dtype) / n - cdfvals if sign == +1 + else (cdfvals - xp.arange(0.0, length, dtype=x.dtype) / n)) + amax = xp.argmax(D, axis=-1, keepdims=True) + loc_max = xp.squeeze(xp.take_along_axis(x, amax, axis=-1), axis=-1) + D = xp.squeeze(xp.take_along_axis(D, amax, axis=-1), axis=-1) + return D[()] if D.ndim == 0 else D, loc_max[()] if loc_max.ndim == 0 else loc_max + + +def _tuple_to_KstestResult(statistic, pvalue, + statistic_location, statistic_sign): + return KstestResult(statistic, pvalue, + statistic_location=statistic_location, + statistic_sign=statistic_sign) + + +def _KstestResult_to_tuple(res, _): + return *res, res.statistic_location, res.statistic_sign + + +# JAX JIT / Torch GPU / CuPy need null distribution +@xp_capabilities(cpu_only=True, jax_jit=False, marray=True, + skip_backends=[('dask.array', 'needs take_along_axis')]) +@_axis_nan_policy_factory(_tuple_to_KstestResult, n_samples=1, n_outputs=4, + result_to_tuple=_KstestResult_to_tuple) +@_rename_parameter("mode", "method") +def ks_1samp(x, cdf, args=(), alternative='two-sided', method='auto', *, axis=0): + """ + Performs the one-sample Kolmogorov-Smirnov test for goodness of fit. + + This test compares the underlying distribution F(x) of a sample + against a given continuous distribution G(x). See Notes for a description + of the available null and alternative hypotheses. + + Parameters + ---------- + x : array_like + a 1-D array of observations of iid random variables. + cdf : callable + callable used to calculate the cdf. + args : tuple, sequence, optional + Distribution parameters, used with `cdf`. + alternative : {'two-sided', 'less', 'greater'}, optional + Defines the null and alternative hypotheses. Default is 'two-sided'. + Please see explanations in the Notes below. + method : {'auto', 'exact', 'approx', 'asymp'}, optional + Defines the distribution used for calculating the p-value. + The following options are available (default is 'auto'): + + * 'auto' : selects one of the other options. + * 'exact' : uses the exact distribution of test statistic. + * 'approx' : approximates the two-sided probability with twice + the one-sided probability + * 'asymp': uses asymptotic distribution of test statistic + + axis : int or tuple of ints, default: 0 + If an int or tuple of ints, the axis or axes of the input along which + to compute the statistic. The statistic of each axis-slice (e.g. row) + of the input will appear in a corresponding element of the output. + If ``None``, the input will be raveled before computing the statistic. + + Returns + ------- + res: KstestResult + An object containing attributes: + + statistic : float + KS test statistic, either D+, D-, or D (the maximum of the two) + pvalue : float + One-tailed or two-tailed p-value. + statistic_location : float + Value of `x` corresponding with the KS statistic; i.e., the + distance between the empirical distribution function and the + hypothesized cumulative distribution function is measured at this + observation. + statistic_sign : int + +1 if the KS statistic is the maximum positive difference between + the empirical distribution function and the hypothesized cumulative + distribution function (D+); -1 if the KS statistic is the maximum + negative difference (D-). + + + See Also + -------- + ks_2samp, kstest + + Notes + ----- + There are three options for the null and corresponding alternative + hypothesis that can be selected using the `alternative` parameter. + + - `two-sided`: The null hypothesis is that the two distributions are + identical, F(x)=G(x) for all x; the alternative is that they are not + identical. + + - `less`: The null hypothesis is that F(x) >= G(x) for all x; the + alternative is that F(x) < G(x) for at least one x. + + - `greater`: The null hypothesis is that F(x) <= G(x) for all x; the + alternative is that F(x) > G(x) for at least one x. + + Note that the alternative hypotheses describe the *CDFs* of the + underlying distributions, not the observed values. For example, + suppose x1 ~ F and x2 ~ G. If F(x) > G(x) for all x, the values in + x1 tend to be less than those in x2. + + Examples + -------- + Suppose we wish to test the null hypothesis that a sample is distributed + according to the standard normal. + We choose a confidence level of 95%; that is, we will reject the null + hypothesis in favor of the alternative if the p-value is less than 0.05. + + When testing uniformly distributed data, we would expect the + null hypothesis to be rejected. + + >>> import numpy as np + >>> from scipy import stats + >>> rng = np.random.default_rng() + >>> stats.ks_1samp(stats.uniform.rvs(size=100, random_state=rng), + ... stats.norm.cdf) + KstestResult(statistic=0.5001899973268688, + pvalue=1.1616392184763533e-23, + statistic_location=0.00047625268963724654, + statistic_sign=-1) + + Indeed, the p-value is lower than our threshold of 0.05, so we reject the + null hypothesis in favor of the default "two-sided" alternative: the data + are *not* distributed according to the standard normal. + + When testing random variates from the standard normal distribution, we + expect the data to be consistent with the null hypothesis most of the time. + + >>> x = stats.norm.rvs(size=100, random_state=rng) + >>> stats.ks_1samp(x, stats.norm.cdf) + KstestResult(statistic=0.05345882212970396, + pvalue=0.9227159037744717, + statistic_location=-1.2451343873745018, + statistic_sign=1) + + As expected, the p-value of 0.92 is not below our threshold of 0.05, so + we cannot reject the null hypothesis. + + Suppose, however, that the random variates are distributed according to + a normal distribution that is shifted toward greater values. In this case, + the cumulative density function (CDF) of the underlying distribution tends + to be *less* than the CDF of the standard normal. Therefore, we would + expect the null hypothesis to be rejected with ``alternative='less'``: + + >>> x = stats.norm.rvs(size=100, loc=0.5, random_state=rng) + >>> stats.ks_1samp(x, stats.norm.cdf, alternative='less') + KstestResult(statistic=0.17482387821055168, + pvalue=0.001913921057766743, + statistic_location=0.3713830565352756, + statistic_sign=-1) + + and indeed, with p-value smaller than our threshold, we reject the null + hypothesis in favor of the alternative. + + """ + # `_axis_nan_policy` decorator ensures `axis=-1` + xp = array_namespace(x) + mode = method + if mode not in ['auto', 'exact', 'approx', 'asymp']: + raise ValueError( + f"Invalid value for method: {mode!r}. " + f"Must be one of 'auto', 'exact', 'approx', 'asymp'." + ) + + alternative = {'t': 'two-sided', 'g': 'greater', 'l': 'less'}.get( + alternative.lower()[0], alternative) + if alternative not in ['two-sided', 'greater', 'less']: + raise ValueError(f"Unexpected value {alternative=}") + + x = xp.sort(x, axis=-1) + x = xp_promote(x, force_floating=True, xp=xp) + N = _count_nonmasked(x, axis=-1, xp=xp) + cdfvals = _masked_apply(cdf, args=(x, *args), xp=xp) + + ones = xp.ones(x.shape[:-1], dtype=xp.int8) + ones = ones[()] if ones.ndim == 0 else ones + + if alternative == 'greater': + Dplus, d_location = _compute_d(cdfvals, x, +1) + pvalue = _masked_apply(distributions.ksone.sf, args=(Dplus, N), xp=xp) + pvalue = xp.asarray(pvalue, dtype=x.dtype) + pvalue = pvalue[()] if pvalue.ndim == 0 else pvalue + Dplus = xp.asarray(Dplus) if is_marray(xp) else Dplus + return KstestResult(Dplus, pvalue, + statistic_location=d_location, + statistic_sign=ones) + + if alternative == 'less': + Dminus, d_location = _compute_d(cdfvals, x, -1) + pvalue = _masked_apply(distributions.ksone.sf, args=(Dminus, N), xp=xp) + pvalue = xp.asarray(pvalue, dtype=x.dtype) + pvalue = pvalue[()] if pvalue.ndim == 0 else pvalue + Dminus = xp.asarray(Dminus) if is_marray(xp) else Dminus + return KstestResult(Dminus, pvalue, + statistic_location=d_location, + statistic_sign=-ones) + + # alternative == 'two-sided': + Dplus, dplus_location = _compute_d(cdfvals, x, +1) + Dminus, dminus_location = _compute_d(cdfvals, x, -1) + i_plus = Dplus > Dminus + D = xp.where(i_plus, Dplus, Dminus) + d_location = xp.where(i_plus, dplus_location, dminus_location) + d_sign = xp.where(i_plus, ones, -ones) + if D.ndim == 0: + D, d_location, d_sign = D[()], d_location[()], d_sign[()] + + if mode == 'auto': # Always select exact + mode = 'exact' + if mode == 'exact': + prob = _masked_apply(distributions.kstwo.sf, args=(D, N), xp=xp) + elif mode == 'asymp': + prob = _masked_apply(distributions.kstwobign.sf, args=(D * N**0.5,), xp=xp) + else: + # mode == 'approx' + prob = 2 * _masked_apply(distributions.ksone.sf, args=(D, N), xp=xp) + prob = xp.clip(xp.asarray(prob, dtype=x.dtype), 0., 1.) + return KstestResult(D, prob, + statistic_location=d_location, + statistic_sign=d_sign) + + +Ks_2sampResult = KstestResult + + +def _compute_prob_outside_square(n, h): + """ + Compute the proportion of paths that pass outside the two diagonal lines. + + Parameters + ---------- + n : int + n > 0 + h : int + 0 <= h <= n + + Returns + ------- + p : float + The proportion of paths that pass outside the lines x-y = +/-h. + + """ + # Compute Pr(D_{n,n} >= h/n) + # Prob = 2 * ( binom(2n, n-h) - binom(2n, n-2a) + binom(2n, n-3a) - ... ) + # / binom(2n, n) + # This formulation exhibits subtractive cancellation. + # Instead divide each term by binom(2n, n), then factor common terms + # and use a Horner-like algorithm + # P = 2 * A0 * (1 - A1*(1 - A2*(1 - A3*(1 - A4*(...))))) + + P = 0.0 + k = int(np.floor(n / h)) + while k >= 0: + p1 = 1.0 + # Each of the Ai terms has numerator and denominator with + # h simple terms. + for j in range(h): + p1 = (n - k * h - j) * p1 / (n + k * h + j + 1) + P = p1 * (1.0 - P) + k -= 1 + return 2 * P + + +def _count_paths_outside_method(m, n, g, h): + """Count the number of paths that pass outside the specified diagonal. + + Parameters + ---------- + m : int + m > 0 + n : int + n > 0 + g : int + g is greatest common divisor of m and n + h : int + 0 <= h <= lcm(m,n) + + Returns + ------- + p : float + The number of paths that go low. + The calculation may overflow - check for a finite answer. + + Notes + ----- + Count the integer lattice paths from (0, 0) to (m, n), which at some + point (x, y) along the path, satisfy: + m*y <= n*x - h*g + The paths make steps of size +1 in either positive x or positive y + directions. + + We generally follow Hodges' treatment of Drion/Gnedenko/Korolyuk. + Hodges, J.L. Jr., + "The Significance Probability of the Smirnov Two-Sample Test," + Arkiv fiur Matematik, 3, No. 43 (1958), 469-86. + + """ + # Compute #paths which stay lower than x/m-y/n = h/lcm(m,n) + # B(x, y) = #{paths from (0,0) to (x,y) without + # previously crossing the boundary} + # = binom(x, y) - #{paths which already reached the boundary} + # Multiply by the number of path extensions going from (x, y) to (m, n) + # Sum. + + # Probability is symmetrical in m, n. Computation below assumes m >= n. + if m < n: + m, n = n, m + mg = m // g + ng = n // g + + # Not every x needs to be considered. + # xj holds the list of x values to be checked. + # Wherever n*x/m + ng*h crosses an integer + lxj = n + (mg-h)//mg + xj = [(h + mg * j + ng-1)//ng for j in range(lxj)] + # B is an array just holding a few values of B(x,y), the ones needed. + # B[j] == B(x_j, j) + if lxj == 0: + return special.binom(m + n, n) + B = np.zeros(lxj) + B[0] = 1 + # Compute the B(x, y) terms + for j in range(1, lxj): + Bj = special.binom(xj[j] + j, j) + for i in range(j): + bin = special.binom(xj[j] - xj[i] + j - i, j-i) + Bj -= bin * B[i] + B[j] = Bj + # Compute the number of path extensions... + num_paths = 0 + for j in range(lxj): + bin = special.binom((m-xj[j]) + (n - j), n-j) + term = B[j] * bin + num_paths += term + return num_paths + + +def _attempt_exact_2kssamp(n1, n2, g, d, alternative): + """Attempts to compute the exact 2sample probability. + + n1, n2 are the sample sizes + g is the gcd(n1, n2) + d is the computed max difference in ECDFs + + Returns (success, d, probability) + """ + lcm = (n1 // g) * n2 + h = int(np.round(d * lcm)) + d = h * 1.0 / lcm + if h == 0: + return True, d, 1.0 + saw_fp_error, prob = False, np.nan + try: + with np.errstate(invalid="raise", over="raise"): + if alternative == 'two-sided': + if n1 == n2: + prob = _compute_prob_outside_square(n1, h) + else: + prob = _compute_outer_prob_inside_method(n1, n2, g, h) + else: + if n1 == n2: + # prob = binom(2n, n-h) / binom(2n, n) + # Evaluating in that form incurs roundoff errors + # from special.binom. Instead calculate directly + jrange = np.arange(h) + prob = np.prod((n1 - jrange) / (n1 + jrange + 1.0)) + else: + with np.errstate(over='raise'): + num_paths = _count_paths_outside_method(n1, n2, g, h) + bin = special.binom(n1 + n2, n1) + if num_paths > bin or np.isinf(bin): + saw_fp_error = True + else: + prob = num_paths / bin + + except (FloatingPointError, OverflowError): + saw_fp_error = True + + if saw_fp_error: + return False, d, np.nan + if not (0 <= prob <= 1): + return False, d, prob + return True, d, prob + + +@xp_capabilities(skip_backends=[('dask.array', 'no rankdata')], + jax_jit=False, cpu_only=True, # null distribution is NumPy-only + marray=True) +@_axis_nan_policy_factory(_tuple_to_KstestResult, n_samples=2, n_outputs=4, + result_to_tuple=_KstestResult_to_tuple) +@_rename_parameter("mode", "method") +def ks_2samp(data1, data2, alternative='two-sided', method='auto', *, axis=0): + """ + Performs the two-sample Kolmogorov-Smirnov test for goodness of fit. + + This test compares the underlying continuous distributions F(x) and G(x) + of two independent samples. See Notes for a description of the available + null and alternative hypotheses. + + Parameters + ---------- + data1, data2 : array_like, 1-Dimensional + Two arrays of sample observations assumed to be drawn from a continuous + distribution, sample sizes can be different. + alternative : {'two-sided', 'less', 'greater'}, optional + Defines the null and alternative hypotheses. Default is 'two-sided'. + Please see explanations in the Notes below. + method : {'auto', 'exact', 'asymp'}, optional + Defines the method used for calculating the p-value. + The following options are available (default is 'auto'): + + * 'auto' : use 'exact' for small size arrays, 'asymp' for large + * 'exact' : use exact distribution of test statistic + * 'asymp' : use asymptotic distribution of test statistic + + axis : int or tuple of ints, default: 0 + If an int or tuple of ints, the axis or axes of the input along which + to compute the statistic. The statistic of each axis-slice (e.g. row) + of the input will appear in a corresponding element of the output. + If ``None``, the input will be raveled before computing the statistic. + + Returns + ------- + res: KstestResult + An object containing attributes: + + statistic : float + KS test statistic. + pvalue : float + One-tailed or two-tailed p-value. + statistic_location : float + Value from `data1` or `data2` corresponding with the KS statistic; + i.e., the distance between the empirical distribution functions is + measured at this observation. + statistic_sign : int + +1 if the empirical distribution function of `data1` exceeds + the empirical distribution function of `data2` at + `statistic_location`, otherwise -1. + + See Also + -------- + kstest, ks_1samp, epps_singleton_2samp, anderson_ksamp + + Notes + ----- + There are three options for the null and corresponding alternative + hypothesis that can be selected using the `alternative` parameter. + + - `less`: The null hypothesis is that F(x) >= G(x) for all x; the + alternative is that F(x) < G(x) for at least one x. The statistic + is the magnitude of the minimum (most negative) difference between the + empirical distribution functions of the samples. + + - `greater`: The null hypothesis is that F(x) <= G(x) for all x; the + alternative is that F(x) > G(x) for at least one x. The statistic + is the maximum (most positive) difference between the empirical + distribution functions of the samples. + + - `two-sided`: The null hypothesis is that the two distributions are + identical, F(x)=G(x) for all x; the alternative is that they are not + identical. The statistic is the maximum absolute difference between the + empirical distribution functions of the samples. + + Note that the alternative hypotheses describe the *CDFs* of the + underlying distributions, not the observed values of the data. For example, + suppose x1 ~ F and x2 ~ G. If F(x) > G(x) for all x, the values in + x1 tend to be less than those in x2. + + If the KS statistic is large, then the p-value will be small, and this may + be taken as evidence against the null hypothesis in favor of the + alternative. + + If ``method='exact'``, `ks_2samp` attempts to compute an exact p-value, + that is, the probability under the null hypothesis of obtaining a test + statistic value as extreme as the value computed from the data. + If ``method='asymp'``, the asymptotic Kolmogorov-Smirnov distribution is + used to compute an approximate p-value. + If ``method='auto'``, an exact p-value computation is attempted if both + sample sizes are less than 10000; otherwise, the asymptotic method is used. + In any case, if an exact p-value calculation is attempted and fails, a + warning will be emitted, and the asymptotic p-value will be returned. + + The 'two-sided' 'exact' computation computes the complementary probability + and then subtracts from 1. As such, the minimum probability it can return + is about 1e-16. While the algorithm itself is exact, numerical + errors may accumulate for large sample sizes. It is most suited to + situations in which one of the sample sizes is only a few thousand. + + We generally follow Hodges' treatment of Drion/Gnedenko/Korolyuk [1]_. + + References + ---------- + .. [1] Hodges, J.L. Jr., "The Significance Probability of the Smirnov + Two-Sample Test," Arkiv fiur Matematik, 3, No. 43 (1958), 469-486. + + Examples + -------- + Suppose we wish to test the null hypothesis that two samples were drawn + from the same distribution. + We choose a confidence level of 95%; that is, we will reject the null + hypothesis in favor of the alternative if the p-value is less than 0.05. + + If the first sample were drawn from a uniform distribution and the second + were drawn from the standard normal, we would expect the null hypothesis + to be rejected. + + >>> import numpy as np + >>> from scipy import stats + >>> rng = np.random.default_rng() + >>> sample1 = stats.uniform.rvs(size=100, random_state=rng) + >>> sample2 = stats.norm.rvs(size=110, random_state=rng) + >>> stats.ks_2samp(sample1, sample2) + KstestResult(statistic=0.5454545454545454, + pvalue=7.37417839555191e-15, + statistic_location=-0.014071496412861274, + statistic_sign=-1) + + + Indeed, the p-value is lower than our threshold of 0.05, so we reject the + null hypothesis in favor of the default "two-sided" alternative: the data + were *not* drawn from the same distribution. + + When both samples are drawn from the same distribution, we expect the data + to be consistent with the null hypothesis most of the time. + + >>> sample1 = stats.norm.rvs(size=105, random_state=rng) + >>> sample2 = stats.norm.rvs(size=95, random_state=rng) + >>> stats.ks_2samp(sample1, sample2) + KstestResult(statistic=0.10927318295739348, + pvalue=0.5438289009927495, + statistic_location=-0.1670157701848795, + statistic_sign=-1) + + As expected, the p-value of 0.54 is not below our threshold of 0.05, so + we cannot reject the null hypothesis. + + Suppose, however, that the first sample were drawn from + a normal distribution shifted toward greater values. In this case, + the cumulative density function (CDF) of the underlying distribution tends + to be *less* than the CDF underlying the second sample. Therefore, we would + expect the null hypothesis to be rejected with ``alternative='less'``: + + >>> sample1 = stats.norm.rvs(size=105, loc=0.5, random_state=rng) + >>> stats.ks_2samp(sample1, sample2, alternative='less') + KstestResult(statistic=0.4055137844611529, + pvalue=3.5474563068855554e-08, + statistic_location=-0.13249370614972575, + statistic_sign=-1) + + and indeed, with p-value smaller than our threshold, we reject the null + hypothesis in favor of the alternative. + + """ + # because of the _axis_nan_policy decorator, we can assume the arrays + # are broadcasted and `axis=-1` + mode = method + + if mode not in ['auto', 'exact', 'asymp']: + raise ValueError(f'Invalid value for mode: {mode}') + alternative = {'t': 'two-sided', 'g': 'greater', 'l': 'less'}.get( + alternative.lower()[0], alternative) + if alternative not in ['two-sided', 'less', 'greater']: + raise ValueError(f'Invalid value for alternative: {alternative}') + MAX_AUTO_N = 10000 # 'auto' will attempt to be exact if n1,n2 <= MAX_AUTO_N + + xp = array_namespace(data1, data2) + data1 = xp.sort(data1, axis=-1) + data2 = xp.sort(data2, axis=-1) + n1 = data1.shape[-1] + n2 = data2.shape[-1] + if min(n1, n2) == 0: + raise ValueError('Data passed to ks_2samp must not be empty') + + data_all = xp.concat((data1, data2), axis=-1) + batch_shape = data_all.shape[:-1] + dtype = xp_result_type(data1, data2, force_floating=True, xp=xp) + + if is_marray(xp): + # Previously, we used this algorithm for all backends: + n1 = xp.astype(_count_nonmasked(data1, axis=-1), dtype) + n2 = xp.astype(_count_nonmasked(data2, axis=-1), dtype) + cdf1 = xp.astype(_xp_searchsorted(data1, data_all, side='right'), dtype) + cdf2 = xp.astype(_xp_searchsorted(data2, data_all, side='right'), dtype) + cddiffs = cdf1/n1[..., xp.newaxis] - cdf2/n2[..., xp.newaxis] + # but the switch from `xp.searchsorted` to `_xp_searchsorted` would come with a + # pretty steep performance hit, so the algorithm below is faster. + # When `xp.searchsorted` is vectorized, we can revert. + else: + n = n1 + n2 + + # We want the ECDF of each sample evaluated at *all* the points in the pooled + # sample. The values each ECDF can assume are given by: + cdf1_vals = xp.broadcast_to(xp.linspace(0, 1, n1 + 1, dtype=dtype), + batch_shape + (n1 + 1,)) + cdf2_vals = xp.broadcast_to(xp.linspace(0, 1, n2 + 1, dtype=dtype), + batch_shape + (n2 + 1,)) + # Now we "just" need to know how many times each of these values *will* be + # assumed when we evaluate the ECDFs at all points in the pooled sample. + # These counts are given by the differences between consecutive ("min" or "max") + # ranks corresponding with the observations in the (sorted) samples. + ranks, data_all, _ = _rankdata(data_all, method='min', return_sorted=True) + ranks = xp.astype(ranks, xp.asarray(1).dtype) # default int type + one = xp.ones((*ranks.shape[:-1], 1), dtype=ranks.dtype, + device=xp_device(ranks)) + cdf1_counts = xp.diff(ranks[..., :n1], prepend=one, append=n + one, axis=-1) + cdf2_counts = xp.diff(ranks[..., -n2:], prepend=one, append=n + one, axis=-1) + # Repeat isn't vectorized - in general, this would produce a ragged array. + # However, in our case, the sum of repeats for each slice is the same, so we can + # do a vectorized repeat by raveling, repeating, then restoring the shape. + cdf1 = xp.repeat(xp_ravel(cdf1_vals), xp_ravel(cdf1_counts), axis=-1) + cdf2 = xp.repeat(xp_ravel(cdf2_vals), xp_ravel(cdf2_counts), axis=-1) + cddiffs = xp.reshape(cdf1 - cdf2, ranks.shape[:-1] + (-1,)) + + # Identify the location of the statistic + argminS = xp.argmin(cddiffs, axis=-1, keepdims=True) + argmaxS = xp.argmax(cddiffs, axis=-1, keepdims=True) + loc_minS = xp.squeeze(xp.take_along_axis(data_all, argminS, axis=-1), axis=-1) + loc_maxS = xp.squeeze(xp.take_along_axis(data_all, argmaxS, axis=-1), axis=-1) + + # Ensure sign of minS is not negative. + minS = -xp.squeeze(xp.take_along_axis(cddiffs, argminS, axis=-1), axis=-1) + maxS = xp.squeeze(xp.take_along_axis(cddiffs, argmaxS, axis=-1), axis=-1) + minS = xp.clip(minS, 0., 1.) + + if alternative == 'less': + selector = xp.ones(minS.shape, dtype=xp.bool) + elif alternative == 'two-sided': + selector = minS > maxS + else: + selector = xp.zeros(minS.shape, dtype=xp.bool) + + d = xp.where(selector, minS, maxS) + d_location = xp.where(selector, loc_minS, loc_maxS) + one = xp.asarray(1, dtype=xp.int8) + d_sign = xp.where(selector, -one, one) + + if is_marray(xp): + d = d.data # converted to NumPy below + n1, n2 = np.asarray(n1.data, dtype=int), np.asarray(n2.data, dtype=int) + prob = _ks_2samp_prob(np.asarray(d), n1, n2, mode, MAX_AUTO_N, alternative) + dtype = xp_result_type(data1, data2, force_floating=True, xp=xp) + prob = xp.asarray(prob, dtype=dtype) + d = xp.asarray(d, dtype=dtype) + if d.ndim == 0: + d, prob, d_location, d_sign = d[()], prob[()], d_location[()], d_sign[()] + return KstestResult(d, prob, statistic_location=d_location, statistic_sign=d_sign) + + +@np.vectorize(excluded={3, 4, 5}, otypes=[np.float64]) +def _ks_2samp_prob(d, n1, n2, mode, MAX_AUTO_N, alternative): + # math.gcd used to return an `int`, but now that these return arrays, NumPy can + # return an int32. Unless we want the pythranized functions to be compiled for + # int32, this needs to produce an int64. + g = np.gcd(n1, n2).astype(np.int64) + n1g = n1 // g + n2g = n2 // g + prob = -np.inf + if mode == 'auto': + mode = 'exact' if max(n1, n2) <= MAX_AUTO_N else 'asymp' + elif mode == 'exact': + # If lcm(n1, n2) is too big, switch from exact to asymp + if n1g >= np.iinfo(np.int32).max / n2g: + mode = 'asymp' + warnings.warn( + f"Exact ks_2samp calculation not possible with samples sizes " + f"{n1} and {n2}. Switching to 'asymp'.", RuntimeWarning, + stacklevel=3) + + if mode == 'exact': + success, d, prob = _attempt_exact_2kssamp(n1, n2, g, d, alternative) + if not success: + mode = 'asymp' + warnings.warn(f"ks_2samp: Exact calculation unsuccessful. " + f"Switching to method={mode}.", RuntimeWarning, + stacklevel=3) + + if mode == 'asymp': + # The product n1*n2 is large. Use Smirnov's asymptotic formula. + # Ensure float to avoid overflow in multiplication + # sorted because the one-sided formula is not symmetric in n1, n2 + m, n = sorted([float(n1), float(n2)], reverse=True) + en = m * n / (m + n) + if alternative == 'two-sided': + prob = distributions.kstwo.sf(d, np.round(en)) + else: + z = np.sqrt(en) * d + # Use Hodges' suggested approximation Eqn 5.3 + # Requires m to be the larger of (n1, n2) + expt = -2 * z**2 - 2 * z * (m + 2*n)/np.sqrt(m*n*(m+n))/3.0 + prob = np.exp(expt) + + prob = np.clip(prob, 0, 1) + return prob + + +def _parse_kstest_args(data1, data2, args, N): + # kstest allows many different variations of arguments. + # Pull out the parsing into a separate function + # (xvals, yvals, ) # 2sample + # (xvals, cdf function,..) + # (xvals, name of distribution, ...) + # (name of distribution, name of distribution, ...) + + # Returns xvals, yvals, cdf + # where cdf is a cdf function, or None + # and yvals is either an array_like of values, or None + # and xvals is array_like. + rvsfunc, cdf = None, None + if isinstance(data1, str): + rvsfunc = getattr(distributions, data1).rvs + elif callable(data1): + rvsfunc = data1 + + if isinstance(data2, str): + special_distributions = {'norm': special.ndtr} + cdf = special_distributions.get(data2, getattr(distributions, data2).cdf) + data2 = None + elif callable(data2): + cdf = data2 + data2 = None + + xp = array_namespace(data1, data2, *args) + data1 = xp.sort(rvsfunc(*args, size=N) if rvsfunc else data1) + return data1, data2, cdf + + +def _kstest_n_samples(kwargs): + cdf = kwargs['cdf'] + return 1 if (isinstance(cdf, str) or callable(cdf)) else 2 + + +@xp_capabilities(skip_backends=[('dask.array', 'no rankdata')], + jax_jit=False, cpu_only=True, # see ks_1samp/ks_2samp + marray=True) +@_axis_nan_policy_factory(_tuple_to_KstestResult, n_samples=_kstest_n_samples, + n_outputs=4, result_to_tuple=_KstestResult_to_tuple) +@_rename_parameter("mode", "method") +def kstest(rvs, cdf, args=(), N=20, alternative='two-sided', method='auto', *, + axis=0): + """ + Performs the (one-sample or two-sample) Kolmogorov-Smirnov test for + goodness of fit. + + The one-sample test compares the underlying distribution F(x) of a sample + against a given distribution G(x). The two-sample test compares the + underlying distributions of two independent samples. Both tests are valid + only for continuous distributions. + + Parameters + ---------- + rvs : str, array_like, or callable + If an array, it should be a 1-D array of observations of random + variables. + If a callable, it should be a function to generate random variables; + it is required to have a keyword argument `size`. + If a string, it should be the name of a distribution in `scipy.stats`, + which will be used to generate random variables. + cdf : str, array_like or callable + If array_like, it should be a 1-D array of observations of random + variables, and the two-sample test is performed + (and rvs must be array_like). + If a callable, that callable is used to calculate the cdf. + If a string, it should be the name of a distribution in `scipy.stats`, + which will be used as the cdf function. + args : tuple, sequence, optional + Distribution parameters, used if `rvs` or `cdf` are strings or + callables. + N : int, optional + Sample size if `rvs` is string or callable. Default is 20. + alternative : {'two-sided', 'less', 'greater'}, optional + Defines the null and alternative hypotheses. Default is 'two-sided'. + Please see explanations in the Notes below. + method : {'auto', 'exact', 'approx', 'asymp'}, optional + Defines the distribution used for calculating the p-value. + The following options are available (default is 'auto'): + + * 'auto' : selects one of the other options. + * 'exact' : uses the exact distribution of test statistic. + * 'approx' : approximates the two-sided probability with twice the + one-sided probability + * 'asymp': uses asymptotic distribution of test statistic + + Returns + ------- + res: KstestResult + An object containing attributes: + + statistic : float + KS test statistic, either D+, D-, or D (the maximum of the two) + pvalue : float + One-tailed or two-tailed p-value. + statistic_location : float + In a one-sample test, this is the value of `rvs` + corresponding with the KS statistic; i.e., the distance between + the empirical distribution function and the hypothesized cumulative + distribution function is measured at this observation. + + In a two-sample test, this is the value from `rvs` or `cdf` + corresponding with the KS statistic; i.e., the distance between + the empirical distribution functions is measured at this + observation. + statistic_sign : int + In a one-sample test, this is +1 if the KS statistic is the + maximum positive difference between the empirical distribution + function and the hypothesized cumulative distribution function + (D+); it is -1 if the KS statistic is the maximum negative + difference (D-). + + In a two-sample test, this is +1 if the empirical distribution + function of `rvs` exceeds the empirical distribution + function of `cdf` at `statistic_location`, otherwise -1. + + See Also + -------- + ks_1samp, ks_2samp + + Notes + ----- + There are three options for the null and corresponding alternative + hypothesis that can be selected using the `alternative` parameter. + + - `two-sided`: The null hypothesis is that the two distributions are + identical, F(x)=G(x) for all x; the alternative is that they are not + identical. + + - `less`: The null hypothesis is that F(x) >= G(x) for all x; the + alternative is that F(x) < G(x) for at least one x. + + - `greater`: The null hypothesis is that F(x) <= G(x) for all x; the + alternative is that F(x) > G(x) for at least one x. + + Note that the alternative hypotheses describe the *CDFs* of the + underlying distributions, not the observed values. For example, + suppose x1 ~ F and x2 ~ G. If F(x) > G(x) for all x, the values in + x1 tend to be less than those in x2. + + + Examples + -------- + Suppose we wish to test the null hypothesis that a sample is distributed + according to the standard normal. + We choose a confidence level of 95%; that is, we will reject the null + hypothesis in favor of the alternative if the p-value is less than 0.05. + + When testing uniformly distributed data, we would expect the + null hypothesis to be rejected. + + >>> import numpy as np + >>> from scipy import stats + >>> rng = np.random.default_rng() + >>> stats.kstest(stats.uniform.rvs(size=100, random_state=rng), + ... stats.norm.cdf) + KstestResult(statistic=0.5001899973268688, + pvalue=1.1616392184763533e-23, + statistic_location=0.00047625268963724654, + statistic_sign=-1) + + Indeed, the p-value is lower than our threshold of 0.05, so we reject the + null hypothesis in favor of the default "two-sided" alternative: the data + are *not* distributed according to the standard normal. + + When testing random variates from the standard normal distribution, we + expect the data to be consistent with the null hypothesis most of the time. + + >>> x = stats.norm.rvs(size=100, random_state=rng) + >>> stats.kstest(x, stats.norm.cdf) + KstestResult(statistic=0.05345882212970396, + pvalue=0.9227159037744717, + statistic_location=-1.2451343873745018, + statistic_sign=1) + + + As expected, the p-value of 0.92 is not below our threshold of 0.05, so + we cannot reject the null hypothesis. + + Suppose, however, that the random variates are distributed according to + a normal distribution that is shifted toward greater values. In this case, + the cumulative density function (CDF) of the underlying distribution tends + to be *less* than the CDF of the standard normal. Therefore, we would + expect the null hypothesis to be rejected with ``alternative='less'``: + + >>> x = stats.norm.rvs(size=100, loc=0.5, random_state=rng) + >>> stats.kstest(x, stats.norm.cdf, alternative='less') + KstestResult(statistic=0.17482387821055168, + pvalue=0.001913921057766743, + statistic_location=0.3713830565352756, + statistic_sign=-1) + + and indeed, with p-value smaller than our threshold, we reject the null + hypothesis in favor of the alternative. + + For convenience, the previous test can be performed using the name of the + distribution as the second argument. + + >>> stats.kstest(x, "norm", alternative='less') + KstestResult(statistic=0.17482387821055168, + pvalue=0.001913921057766743, + statistic_location=0.3713830565352756, + statistic_sign=-1) + + The examples above have all been one-sample tests identical to those + performed by `ks_1samp`. Note that `kstest` can also perform two-sample + tests identical to those performed by `ks_2samp`. For example, when two + samples are drawn from the same distribution, we expect the data to be + consistent with the null hypothesis most of the time. + + >>> sample1 = stats.laplace.rvs(size=105, random_state=rng) + >>> sample2 = stats.laplace.rvs(size=95, random_state=rng) + >>> stats.kstest(sample1, sample2) + KstestResult(statistic=0.11779448621553884, + pvalue=0.4494256912629795, + statistic_location=0.6138814275424155, + statistic_sign=1) + + As expected, the p-value of 0.45 is not below our threshold of 0.05, so + we cannot reject the null hypothesis. + + """ + # to not break compatibility with existing code + if alternative == 'two_sided': + alternative = 'two-sided' + if alternative not in ['two-sided', 'greater', 'less']: + raise ValueError(f"Unexpected alternative: {alternative}") + xvals, yvals, cdf = _parse_kstest_args(rvs, cdf, args, N) + if cdf: + return ks_1samp(xvals, cdf, args=args, alternative=alternative, + method=method, axis=axis, _no_deco=True) + return ks_2samp(xvals, yvals, alternative=alternative, method=method, + axis=axis, _no_deco=True) + + +@xp_capabilities(np_only=True) +def tiecorrect(rankvals): + """Tie correction factor for Mann-Whitney U and Kruskal-Wallis H tests. + + Parameters + ---------- + rankvals : array_like + A 1-D sequence of ranks. Typically this will be the array + returned by `~scipy.stats.rankdata`. + + Returns + ------- + factor : float + Correction factor for U or H. + + See Also + -------- + rankdata : Assign ranks to the data + mannwhitneyu : Mann-Whitney rank test + kruskal : Kruskal-Wallis H test + + References + ---------- + .. [1] Siegel, S. (1956) Nonparametric Statistics for the Behavioral + Sciences. New York: McGraw-Hill. + + Examples + -------- + >>> from scipy.stats import tiecorrect, rankdata + >>> tiecorrect([1, 2.5, 2.5, 4]) + 0.9 + >>> ranks = rankdata([1, 3, 2, 4, 5, 7, 2, 8, 4]) + >>> ranks + array([ 1. , 4. , 2.5, 5.5, 7. , 8. , 2.5, 9. , 5.5]) + >>> tiecorrect(ranks) + 0.9833333333333333 + + """ + arr = np.sort(rankvals) + idx = np.nonzero(np.r_[True, arr[1:] != arr[:-1], True])[0] + cnt = np.diff(idx).astype(np.float64) + + size = np.float64(arr.size) + return 1.0 if size < 2 else 1.0 - (cnt**3 - cnt).sum() / (size**3 - size) + + +RanksumsResult = namedtuple('RanksumsResult', ('statistic', 'pvalue')) + + +@xp_capabilities(out_of_scope=True) # legacy +@_axis_nan_policy_factory(RanksumsResult, n_samples=2) +def ranksums(x, y, alternative='two-sided'): + """Compute the Wilcoxon rank-sum statistic for two samples. + + .. legacy:: function + + This function includes a subset of the features of `mannwhitneyu`. + Its statistic and p-value can be reproduced using `mannwhitneyu` with + options ``use_continuity=False`` and ``method='asymptotic'``. + Prefer `mannwhitneyu` for use in new code. + + The Wilcoxon rank-sum test tests the null hypothesis that two sets + of measurements are drawn from the same distribution. The alternative + hypothesis is that values in one sample are more likely to be + larger than the values in the other sample. + + This test should be used to compare two samples from continuous + distributions. It does not handle ties between measurements + in x and y. For tie-handling and an optional continuity correction + see `scipy.stats.mannwhitneyu`. + + Parameters + ---------- + x, y : array_like + The data from the two samples. + alternative : {'two-sided', 'less', 'greater'}, optional + Defines the alternative hypothesis. Default is 'two-sided'. + The following options are available: + + * 'two-sided': one of the distributions (underlying `x` or `y`) is + stochastically greater than the other. + * 'less': the distribution underlying `x` is stochastically less + than the distribution underlying `y`. + * 'greater': the distribution underlying `x` is stochastically greater + than the distribution underlying `y`. + + .. versionadded:: 1.7.0 + + Returns + ------- + statistic : float + The test statistic under the large-sample approximation that the + rank sum statistic is normally distributed. + pvalue : float + The p-value of the test. + + See Also + -------- + scipy.stats.mannwhitneyu + + References + ---------- + .. [1] https://en.wikipedia.org/wiki/Wilcoxon_rank-sum_test + + Examples + -------- + We can test the hypothesis that two independent unequal-sized samples are + drawn from the same distribution with computing the Wilcoxon rank-sum + statistic. + + >>> import numpy as np + >>> from scipy import stats + >>> rng = np.random.default_rng(6004253334) + >>> sample1 = rng.uniform(-1, 1, 200) + >>> sample2 = rng.uniform(-0.5, 1.5, 300) # a shifted distribution + >>> stats.ranksums(sample1, sample2) + RanksumsResult(statistic=np.float64(-7.576201867), + pvalue=np.float64(3.5581802537469756e-14)) + >>> stats.ranksums(sample1, sample2, alternative='less') + RanksumsResult(statistic=np.float64(-7.576201867066302), + pvalue=np.float64(1.7790901268734878e-14)) + >>> stats.ranksums(sample1, sample2, alternative='greater') + RanksumsResult(statistic=np.float64(-7.576201867066302), + pvalue=np.float64(0.9999999999999822)) + + The p-value of less than ``0.05`` indicates that this test rejects the + hypothesis at the 5% significance level. + + Note that identical hypothesis tests (and more accurate ones) can be performed + using `mannwhitneyu`. Prefer `mannwhitneyu` for use in new code. + + >>> res = stats.mannwhitneyu(sample1, sample2, alternative='greater', + ... method='asymptotic', use_continuity = False) + >>> res.zstatistic + np.float64(-7.576201867066302) + >>> res.pvalue + np.float64(0.9999999999999822) + """ + x, y = map(np.asarray, (x, y)) + n1 = len(x) + n2 = len(y) + alldata = np.concatenate((x, y)) + ranked = rankdata(alldata) + x = ranked[:n1] + s = np.sum(x, axis=0) + expected = n1 * (n1+n2+1) / 2.0 + z = (s - expected) / np.sqrt(n1*n2*(n1+n2+1)/12.0) + pvalue = _get_pvalue(z, _SimpleNormal(), alternative, xp=np) + + return RanksumsResult(z[()], pvalue[()]) + + +KruskalResult = namedtuple('KruskalResult', ('statistic', 'pvalue')) + + +@xp_capabilities(skip_backends=[('dask.array', 'no rankdata')], + marray=True) +@_axis_nan_policy_factory(KruskalResult, n_samples=None) +def kruskal(*samples, nan_policy='propagate', axis=0): + """Compute the Kruskal-Wallis H-test for independent samples. + + The Kruskal-Wallis H-test tests the null hypothesis that the population + median of all of the groups are equal. It is a non-parametric version of + ANOVA. The test works on 2 or more independent samples, which may have + different sizes. Note that rejecting the null hypothesis does not + indicate which of the groups differs. Post hoc comparisons between + groups are required to determine which groups are different. + + Parameters + ---------- + *samples : array_like + Two or more arrays with the sample measurements can be given as + arguments. Samples must be one-dimensional. + nan_policy : {'propagate', 'raise', 'omit'}, optional + Defines how to handle when input contains nan. + The following options are available (default is 'propagate'): + + * 'propagate': returns nan + * 'raise': throws an error + * 'omit': performs the calculations ignoring nan values + + axis : int or tuple of ints, default: 0 + If an int or tuple of ints, the axis or axes of the input along which + to compute the statistic. The statistic of each axis-slice (e.g. row) + of the input will appear in a corresponding element of the output. + If ``None``, the input will be raveled before computing the statistic. + + Returns + ------- + statistic : float + The Kruskal-Wallis H statistic, corrected for ties. + pvalue : float + The p-value for the test using the assumption that H has a chi + square distribution. The p-value returned is the survival function of + the chi square distribution evaluated at H. + + See Also + -------- + f_oneway : 1-way ANOVA. + mannwhitneyu : Mann-Whitney rank test on two samples. + friedmanchisquare : Friedman test for repeated measurements. + + Notes + ----- + Due to the assumption that H has a chi square distribution, the number + of samples in each group must not be too small. A typical rule is + that each sample must have at least 5 measurements. + + References + ---------- + .. [1] W. H. Kruskal & W. W. Wallis, "Use of Ranks in + One-Criterion Variance Analysis", Journal of the American Statistical + Association, Vol. 47, Issue 260, pp. 583-621, 1952. + .. [2] https://en.wikipedia.org/wiki/Kruskal-Wallis_one-way_analysis_of_variance + + Examples + -------- + >>> from scipy import stats + >>> x = [1, 3, 5, 7, 9] + >>> y = [2, 4, 6, 8, 10] + >>> stats.kruskal(x, y) + KruskalResult(statistic=0.2727272727272734, pvalue=0.6015081344405895) + + >>> x = [1, 1, 1] + >>> y = [2, 2, 2] + >>> z = [2, 2] + >>> stats.kruskal(x, y, z) + KruskalResult(statistic=7.0, pvalue=0.0301973834223185) + + """ + xp = array_namespace(*samples) + samples = xp_promote(*samples, force_floating=True, xp=xp) + + num_groups = len(samples) + if num_groups < 2: + raise ValueError("Need at least two groups in stats.kruskal()") + + lengths = [sample.shape[-1] for sample in samples] + if any(lengths) < 1: # Only needed for `test_axis_nan_policy` + raise ValueError("Inputs must not be empty.") + + alldata = xp.concat(samples, axis=-1) + ranked, _, t = _rankdata(alldata, method='average', return_ties=True) + counts = [xp.asarray(_count_nonmasked(sample, -1), dtype=t.dtype) + for sample in samples] + totaln = sum(counts) + ties = 1 - xp.sum(t**3 - t, axis=-1) / (totaln**3 - totaln) # tiecorrect(ranked) + + # Compute sum^2/count for each group and sum + j = list(itertools.accumulate(lengths, initial=0)) + ssbn = sum(xp.sum(ranked[..., j[i]:j[i + 1]], axis=-1)**2 / counts[i] + for i in range(num_groups)) + + h = 12.0 / (totaln * (totaln + 1)) * ssbn - 3 * (totaln + 1) + df = xp.asarray(num_groups - 1, dtype=h.dtype) + h /= ties + + chi2 = _SimpleChi2(df) + pvalue = _get_pvalue(h, chi2, alternative='greater', symmetric=False, xp=np) + return KruskalResult(h, pvalue) + + +FriedmanchisquareResult = namedtuple('FriedmanchisquareResult', + ('statistic', 'pvalue')) + + +@xp_capabilities(skip_backends=[("dask.array", "no rankdata")], + marray=True) +@_axis_nan_policy_factory(FriedmanchisquareResult, n_samples=None, paired=True) +def friedmanchisquare(*samples, axis=0): + """Compute the Friedman test for repeated samples. + + The Friedman test tests the null hypothesis that repeated samples of + the same individuals have the same distribution. It is often used + to test for consistency among samples obtained in different ways. + For example, if two sampling techniques are used on the same set of + individuals, the Friedman test can be used to determine if the two + sampling techniques are consistent. + + Parameters + ---------- + *samples : array_like + Arrays of observations. All of the arrays must have the same number + of elements. At least three samples must be given. + axis : int or tuple of ints, default: 0 + If an int or tuple of ints, the axis or axes of the input along which + to compute the statistic. The statistic of each axis-slice (e.g. row) + of the input will appear in a corresponding element of the output. + If ``None``, the input will be raveled before computing the statistic. + + Returns + ------- + statistic : float + The test statistic, correcting for ties. + pvalue : float + The associated p-value assuming that the test statistic has a chi + squared distribution. + + See Also + -------- + :ref:`hypothesis_friedmanchisquare` : Extended example + + Notes + ----- + Due to the assumption that the test statistic has a chi squared + distribution, the p-value is only reliable for n > 10 and more than + 6 repeated samples. + + References + ---------- + .. [1] https://en.wikipedia.org/wiki/Friedman_test + .. [2] Demsar, J. (2006). Statistical comparisons of classifiers over + multiple data sets. Journal of Machine Learning Research, 7, 1-30. + + Examples + -------- + + >>> import numpy as np + >>> rng = np.random.default_rng(seed=18) + >>> x = rng.random((6, 10)) + >>> from scipy.stats import friedmanchisquare + >>> res = friedmanchisquare(x[0], x[1], x[2], x[3], x[4], x[5]) + >>> res.statistic, res.pvalue + (11.428571428571416, 0.043514520866727614) + + The p-value is less than 0.05; however, as noted above, the results may not + be reliable since we have a small number of repeated samples. + + For a more detailed example, see :ref:`hypothesis_friedmanchisquare`. + """ + k = len(samples) + if k < 3: + raise ValueError('At least 3 samples must be given ' + f'for Friedman test, got {k}.') + + xp = array_namespace(*samples) + samples = xp_promote(*samples, force_floating=True, xp=xp) + dtype = samples[0].dtype + samples = _share_masks(*samples, xp=xp) # paired-sample test + + n = samples[0].shape[-1] + if n == 0: # only for `test_axis_nan_policy`; user doesn't see this + raise ValueError("One or more sample arguments is too small.") + + # Rank data + # axis-slices are aligned with axis -1 by decorator; stack puts samples along axis 0 + # The transpose flips this so we can work with axis-slices along -1. This is a + # reducing statistic, so both axes 0 and -1 are consumed. + data = xp_swapaxes(xp.stack(samples), 0, -1) + data, _, t = _rankdata(data, method='average', return_ties=True) + + # Handle ties + ties = xp.sum(t * (t*t - 1), axis=(0, -1)) + count = xp.asarray(_count_nonmasked(samples[0], axis=-1), dtype=ties.dtype) + c = 1 - ties / (k*(k*k - 1)*count) + + ssbn = xp.sum(xp.sum(data, axis=0)**2, axis=-1) + statistic = (12.0 / (k*count*(k+1)) * ssbn - 3*count*(k+1)) / c + + chi2 = _SimpleChi2(xp.asarray(k - 1, dtype=dtype)) + pvalue = _get_pvalue(statistic, chi2, alternative='greater', symmetric=False, xp=xp) + return FriedmanchisquareResult(statistic, pvalue) + + +BrunnerMunzelResult = namedtuple('BrunnerMunzelResult', + ('statistic', 'pvalue')) + + +@xp_capabilities(cpu_only=True, # torch GPU can't use `stdtr` + marray=True, + skip_backends=[('dask.array', 'needs rankdata')]) +@_axis_nan_policy_factory(BrunnerMunzelResult, n_samples=2) +def brunnermunzel(x, y, alternative="two-sided", distribution="t", + nan_policy='propagate', *, axis=0): + """Compute the Brunner-Munzel test on samples x and y. + + The Brunner-Munzel test is a nonparametric test of the null hypothesis that + when values are taken one by one from each group, the probabilities of + getting large values in both groups are equal. + Unlike the Wilcoxon-Mann-Whitney's U test, this does not require the + assumption of equivariance of two groups. Note that this does not assume + the distributions are same. This test works on two independent samples, + which may have different sizes. + + Parameters + ---------- + x, y : array_like + Array of samples, should be one-dimensional. + alternative : {'two-sided', 'less', 'greater'}, optional + Defines the alternative hypothesis. + The following options are available (default is 'two-sided'): + + * 'two-sided' + * 'less': one-sided + * 'greater': one-sided + + distribution : {'t', 'normal'}, optional + Defines how to get the p-value. + The following options are available (default is 't'): + + * 't': get the p-value by t-distribution + * 'normal': get the p-value by standard normal distribution. + + nan_policy : {'propagate', 'raise', 'omit'}, optional + Defines how to handle when input contains nan. + The following options are available (default is 'propagate'): + + * 'propagate': returns nan + * 'raise': throws an error + * 'omit': performs the calculations ignoring nan values + + axis : int or None, default=0 + If an int, the axis of the input along which to compute the statistic. + The statistic of each axis-slice (e.g. row) of the input will appear + in a corresponding element of the output. If None, the input will be + raveled before computing the statistic. + + Returns + ------- + statistic : float + The Brunner-Munzer W statistic. + pvalue : float + p-value assuming an t distribution. One-sided or + two-sided, depending on the choice of `alternative` and `distribution`. + + See Also + -------- + mannwhitneyu : Mann-Whitney rank test on two samples. + + Notes + ----- + Brunner and Munzel recommended to estimate the p-value by t-distribution + when the size of data is 50 or less. If the size is lower than 10, it would + be better to use permuted Brunner Munzel test (see [2]_). + + References + ---------- + .. [1] Brunner, E. and Munzel, U. "The nonparametric Benhrens-Fisher + problem: Asymptotic theory and a small-sample approximation". + Biometrical Journal. Vol. 42(2000): 17-25. + .. [2] Neubert, K. and Brunner, E. "A studentized permutation test for the + non-parametric Behrens-Fisher problem". Computational Statistics and + Data Analysis. Vol. 51(2007): 5192-5204. + + Examples + -------- + >>> from scipy import stats + >>> x1 = [1,2,1,1,1,1,1,1,1,1,2,4,1,1] + >>> x2 = [3,3,4,3,1,2,3,1,1,5,4] + >>> w, p_value = stats.brunnermunzel(x1, x2) + >>> w + 3.1374674823029505 + >>> p_value + 0.0057862086661515377 + + """ + # _axis_nan_policy decorator ensures we can work along the last axis + xp = array_namespace(x, y) + length_x = x.shape[-1] + nx = _count_nonmasked(x, axis=-1) + ny = _count_nonmasked(y, axis=-1) + + rankc = rankdata(xp.concat((x, y), axis=axis), axis=-1) + rankcx = rankc[..., :length_x] + rankcy = rankc[..., length_x:] + rankcx_mean = xp.mean(rankcx, axis=-1, keepdims=True) + rankcy_mean = xp.mean(rankcy, axis=-1, keepdims=True) + rankx = rankdata(x, axis=-1) + ranky = rankdata(y, axis=-1) + rankx_mean = xp.mean(rankx, axis=-1, keepdims=True) + ranky_mean = xp.mean(ranky, axis=-1, keepdims=True) + + temp_x = rankcx - rankx - rankcx_mean + rankx_mean + Sx = xp.vecdot(temp_x, temp_x, axis=-1) + Sx /= nx - 1 + temp_y = rankcy - ranky - rankcy_mean + ranky_mean + Sy = xp.vecdot(temp_y, temp_y, axis=-1) + Sy /= ny - 1 + + rankcx_mean = xp.squeeze(rankcx_mean, axis=-1) + rankcy_mean = xp.squeeze(rankcy_mean, axis=-1) + wbfn = nx * ny * (rankcy_mean - rankcx_mean) + wbfn /= (nx + ny) * xp.sqrt(nx * Sx + ny * Sy) + + if distribution == "t": + df_numer = xp.pow(nx * Sx + ny * Sy, 2.0) + df_denom = xp.pow(nx * Sx, 2.0) / (nx - 1) + df_denom += xp.pow(ny * Sy, 2.0) / (ny - 1) + df = df_numer / df_denom + + if not is_lazy_array(df_numer) and not is_lazy_array(df_denom) and ( + xp.any(df_numer == 0) and xp.any(df_denom == 0)): + message = ("p-value cannot be estimated with `distribution='t' " + "because degrees of freedom parameter is undefined " + "(0/0). Try using `distribution='normal'") + warnings.warn(message, RuntimeWarning, stacklevel=2) + + distribution = _SimpleStudentT(df) + elif distribution == "normal": + distribution = _SimpleNormal() + else: + raise ValueError( + "distribution should be 't' or 'normal'") + + p = _get_pvalue(-wbfn, distribution, alternative, xp=xp) + + return BrunnerMunzelResult(wbfn, p) + + +@xp_capabilities(cpu_only=True, exceptions=['cupy', 'jax.numpy'], + reason='Delegation for `special.stdtr` only implemented for CuPy and JAX.', + marray=True) +@_axis_nan_policy_factory(SignificanceResult, kwd_samples=['weights'], paired=True) +def combine_pvalues(pvalues, method='fisher', weights=None, *, axis=0): + """ + Combine p-values from independent tests that bear upon the same hypothesis. + + These methods are intended only for combining p-values from hypothesis + tests based upon continuous distributions. + + Each method assumes that under the null hypothesis, the p-values are + sampled independently and uniformly from the interval [0, 1]. A test + statistic (different for each method) is computed and a combined + p-value is calculated based upon the distribution of this test statistic + under the null hypothesis. + + Parameters + ---------- + pvalues : array_like + Array of p-values assumed to come from independent tests based on + continuous distributions. + method : {'fisher', 'pearson', 'tippett', 'stouffer', 'mudholkar_george'} + + Name of method to use to combine p-values. + + The available methods are (see Notes for details): + + * 'fisher': Fisher's method (Fisher's combined probability test) + * 'pearson': Pearson's method + * 'mudholkar_george': Mudholkar's and George's method + * 'tippett': Tippett's method + * 'stouffer': Stouffer's Z-score method + + weights : array_like, optional + Optional array of weights used only for Stouffer's Z-score method. + Ignored by other methods. + + Returns + ------- + res : SignificanceResult + An object containing attributes: + + statistic : float + The statistic calculated by the specified method. + pvalue : float + The combined p-value. + + Examples + -------- + Suppose we wish to combine p-values from four independent tests + of the same null hypothesis using Fisher's method (default). + + >>> from scipy.stats import combine_pvalues + >>> pvalues = [0.1, 0.05, 0.02, 0.3] + >>> combine_pvalues(pvalues) + SignificanceResult(statistic=20.828626352604235, pvalue=0.007616871850449092) + + When the individual p-values carry different weights, consider Stouffer's + method. + + >>> weights = [1, 2, 3, 4] + >>> res = combine_pvalues(pvalues, method='stouffer', weights=weights) + >>> res.pvalue + 0.009578891494533616 + + Notes + ----- + If this function is applied to tests with a discrete statistics such as + any rank test or contingency-table test, it will yield systematically + wrong results, e.g. Fisher's method will systematically overestimate the + p-value [1]_. This problem becomes less severe for large sample sizes + when the discrete distributions become approximately continuous. + + The differences between the methods can be best illustrated by their + statistics and what aspects of a combination of p-values they emphasise + when considering significance [2]_. For example, methods emphasising large + p-values are more sensitive to strong false and true negatives; conversely + methods focussing on small p-values are sensitive to positives. + + * The statistics of Fisher's method (also known as Fisher's combined + probability test) [3]_ is :math:`-2\\sum_i \\log(p_i)`, which is + equivalent (as a test statistics) to the product of individual p-values: + :math:`\\prod_i p_i`. Under the null hypothesis, this statistics follows + a :math:`\\chi^2` distribution. This method emphasises small p-values. + * Pearson's method uses :math:`-2\\sum_i\\log(1-p_i)`, which is equivalent + to :math:`\\prod_i \\frac{1}{1-p_i}` [2]_. + It thus emphasises large p-values. + * Mudholkar and George compromise between Fisher's and Pearson's method by + averaging their statistics [4]_. Their method emphasises extreme + p-values, both close to 1 and 0. + * Stouffer's method [5]_ uses Z-scores and the statistic: + :math:`\\sum_i \\Phi^{-1} (p_i)`, where :math:`\\Phi` is the CDF of the + standard normal distribution. The advantage of this method is that it is + straightforward to introduce weights, which can make Stouffer's method + more powerful than Fisher's method when the p-values are from studies + of different size [6]_ [7]_. + * Tippett's method uses the smallest p-value as a statistic. + (Mind that this minimum is not the combined p-value.) + + Fisher's method may be extended to combine p-values from dependent tests + [8]_. Extensions such as Brown's method and Kost's method are not currently + implemented. + + .. versionadded:: 0.15.0 + + References + ---------- + .. [1] Kincaid, W. M., "The Combination of Tests Based on Discrete + Distributions." Journal of the American Statistical Association 57, + no. 297 (1962), 10-19. + .. [2] Heard, N. and Rubin-Delanchey, P. "Choosing between methods of + combining p-values." Biometrika 105.1 (2018): 239-246. + .. [3] https://en.wikipedia.org/wiki/Fisher%27s_method + .. [4] George, E. O., and G. S. Mudholkar. "On the convolution of logistic + random variables." Metrika 30.1 (1983): 1-13. + .. [5] https://en.wikipedia.org/wiki/Fisher%27s_method#Relation_to_Stouffer.27s_Z-score_method + .. [6] Whitlock, M. C. "Combining probability from independent tests: the + weighted Z-method is superior to Fisher's approach." Journal of + Evolutionary Biology 18, no. 5 (2005): 1368-1373. + .. [7] Zaykin, Dmitri V. "Optimally weighted Z-test is a powerful method + for combining probabilities in meta-analysis." Journal of + Evolutionary Biology 24, no. 8 (2011): 1836-1841. + .. [8] https://en.wikipedia.org/wiki/Extensions_of_Fisher%27s_method + + """ + xp = array_namespace(pvalues, weights) + pvalues, weights = xp_promote(pvalues, weights, broadcast=True, + force_floating=True, xp=xp) + + if xp_size(pvalues) == 0: + # This is really only needed for *testing* _axis_nan_policy decorator + # It won't happen when the decorator is used. + NaN = _get_nan(pvalues) + return SignificanceResult(NaN, NaN) + + n = _count_nonmasked(pvalues, axis) + n = xp.asarray(n, dtype=pvalues.dtype, device=xp_device(pvalues)) + + if method == 'fisher': + statistic = -2 * xp.sum(xp.log(pvalues), axis=axis) + chi2 = _SimpleChi2(2*n) + pval = _get_pvalue(statistic, chi2, alternative='greater', + symmetric=False, xp=xp) + elif method == 'pearson': + statistic = 2 * xp.sum(xp.log1p(-pvalues), axis=axis) + chi2 = _SimpleChi2(2*n) + pval = _get_pvalue(-statistic, chi2, alternative='less', symmetric=False, xp=xp) + elif method == 'mudholkar_george': + normalizing_factor = xp.sqrt(3/n)/xp.pi + statistic = (-xp.sum(xp.log(pvalues), axis=axis) + + xp.sum(xp.log1p(-pvalues), axis=axis)) + nu = 5*n + 4 + approx_factor = xp.sqrt(nu / (nu - 2)) + t = _SimpleStudentT(nu) + pval = _get_pvalue(statistic * normalizing_factor * approx_factor, t, + alternative="greater", xp=xp) + elif method == 'tippett': + statistic = xp.min(pvalues, axis=axis) + beta = _SimpleBeta(xp.ones_like(n), n) + pval = _get_pvalue(statistic, beta, alternative='less', symmetric=False, xp=xp) + elif method == 'stouffer': + if weights is None: + weights = xp.ones_like(pvalues, dtype=pvalues.dtype) + pvalues, weights = _share_masks(pvalues, weights, xp=xp) + + norm = _SimpleNormal() + Zi = norm.isf(pvalues) + # Consider `vecdot` when data-apis/array-api#910 is resolved + statistic = (xp.sum(weights * Zi, axis=axis) + / xp_vector_norm(weights, axis=axis)) + pval = _get_pvalue(statistic, norm, alternative="greater", xp=xp) + + else: + raise ValueError( + f"Invalid method {method!r}. Valid methods are 'fisher', " + "'pearson', 'mudholkar_george', 'tippett', and 'stouffer'" + ) + + return SignificanceResult(statistic, pval) + + +@dataclass +class QuantileTestResult: + r""" + Result of `scipy.stats.quantile_test`. + + Attributes + ---------- + statistic: float + The statistic used to calculate the p-value; either ``T1``, the + number of observations less than or equal to the hypothesized quantile, + or ``T2``, the number of observations strictly less than the + hypothesized quantile. Two test statistics are required to handle the + possibility the data was generated from a discrete or mixed + distribution. + + statistic_type : int + ``1`` or ``2`` depending on which of ``T1`` or ``T2`` was used to + calculate the p-value respectively. ``T1`` corresponds to the + ``"greater"`` alternative hypothesis and ``T2`` to the ``"less"``. For + the ``"two-sided"`` case, the statistic type that leads to smallest + p-value is used. For significant tests, ``statistic_type = 1`` means + there is evidence that the population quantile is significantly greater + than the hypothesized value and ``statistic_type = 2`` means there is + evidence that it is significantly less than the hypothesized value. + + pvalue : float + The p-value of the hypothesis test. + """ + statistic: float + statistic_type: int + pvalue: float + _alternative: list[str] = field(repr=False) + _x : np.ndarray = field(repr=False) + _p : float = field(repr=False) + _statistic : float = field(repr=False) + _statistic_type : int = field(repr=False) + _pvalue : float = field(repr=False) + _axis : int = field(repr=False) + _axis_none : bool = field(repr=False) + _keepdims : bool = field(repr=False) + _ndim : int = field(repr=False) + _nan_out : bool = field(repr=False) + _xp : bool = field(repr=False) + + + def confidence_interval(self, confidence_level=0.95): + """ + Compute the confidence interval of the quantile. + + Parameters + ---------- + confidence_level : float, default: 0.95 + Confidence level for the computed confidence interval + of the quantile. Default is 0.95. + + Returns + ------- + ci : ``ConfidenceInterval`` object + The object has attributes ``low`` and ``high`` that hold the + lower and upper bounds of the confidence interval. + + Examples + -------- + >>> import numpy as np + >>> import scipy.stats as stats + >>> p = 0.75 # quantile of interest + >>> q = 0 # hypothesized value of the quantile + >>> x = np.exp(np.arange(0, 1.01, 0.01)) + >>> res = stats.quantile_test(x, q=q, p=p, alternative='less') + >>> lb, ub = res.confidence_interval() + >>> lb, ub + (-inf, 2.293318740264183) + >>> res = stats.quantile_test(x, q=q, p=p, alternative='two-sided') + >>> lb, ub = res.confidence_interval(0.9) + >>> lb, ub + (1.9542373206359396, 2.293318740264183) + """ + + xp = self._xp + alternative = self._alternative + p = self._p + x = xp.sort(self._x, axis=-1) + n = x.shape[-1] + bd = _SimpleBinomial(n=n, p=p) + shape = p.shape + + if confidence_level <= 0 or confidence_level >= 1: + message = "`confidence_level` must be a number between 0 and 1." + raise ValueError(message) + + if n == 0: + zeros = xp.zeros(p.shape, dtype=x.dtype) + low, high = zeros, zeros + elif alternative == 'less': + p = 1 - confidence_level + low = xp.full(shape, -xp.inf) + high_index = xp.astype(bd.isf(p), xp.int64) + valid_index = high_index < n + high_index = xpx.at(high_index)[~valid_index].set(0) + x_high = xp.take_along_axis(x, high_index, axis=-1) + high = xp.where(valid_index, x_high, xp.nan) + elif alternative == 'greater': + p = 1 - confidence_level + low_index = xp.astype(bd.ppf(p), xp.int64) - 1 + valid_index = low_index >= 0 + low_index = xpx.at(low_index)[~valid_index].set(0) + x_low = xp.take_along_axis(x, low_index, axis=-1) + low = xp.where(valid_index, x_low, xp.nan) + high = xp.full(shape, xp.inf) + elif alternative == 'two-sided': + p = (1 - confidence_level) / 2 + low_index = xp.astype(bd.ppf(p), xp.int64) - 1 + valid_index = low_index >= 0 + low_index = xpx.at(low_index)[~valid_index].set(0) + x_low = xp.take_along_axis(x, low_index, axis=-1) + low = xp.where(valid_index, x_low, xp.nan) + high_index = xp.astype(bd.isf(p), xp.int64) + valid_index = high_index < n + high_index = xpx.at(high_index)[~valid_index].set(0) + x_high = xp.take_along_axis(x, high_index, axis=-1) + high = xp.where(valid_index, x_high, xp.nan) + + args = (self._axis, self._axis_none, self._keepdims, + self._ndim, self._nan_out, self._xp) + low = _quantile_test_postprocess(low, *args) + high = _quantile_test_postprocess(high, *args) + return ConfidenceInterval(low, high) + + +def quantile_test_iv(x, q, p, alternative, axis, keepdims): + xp = array_namespace(x, q, p) + dtype = xp_result_type(x, q, p, force_floating=True, xp=xp) + + x = xpx.atleast_nd(x, ndim=1, xp=xp) + message = '`x` must be an array of numbers.' + if not xp.isdtype(x.dtype, 'numeric'): + raise ValueError(message) + + q = xp.asarray(q) + message = "`q` must be a scalar or array of numbers." + if not xp.isdtype(q.dtype, 'numeric'): + raise ValueError(message) + + p = xp.asarray(p) + message = "`p` must be a scalar or array of floats." + if not xp.isdtype(p.dtype, 'real floating'): + raise ValueError(message) + + axis_none = axis is None + ndim = max(x.ndim, p.ndim) + if axis_none: + x = xp_ravel(x) + q = xp_ravel(q) + p = xp_ravel(p) + axis = 0 + elif np.iterable(axis) or int(axis) != axis: + message = "`axis` must be an integer or None." + raise ValueError(message) + elif (axis >= ndim) or (axis < -ndim): + message = "`axis` is not compatible with the shapes of the inputs." + raise ValueError(message) + axis = int(axis) + + alternatives = {'two-sided', 'less', 'greater'} + message = f"`alternative` must be one of {alternatives}" + if alternative not in alternatives: + raise ValueError(message) + + if keepdims not in {None, True, False}: + message = "If specified, `keepdims` must be True or False." + raise ValueError(message) + + x = xp.sort(x, axis=axis) + q, p = xp.broadcast_arrays(q, p) # length along axis matters for q and p + x, q, p = _broadcast_arrays((x, q, p), axis=axis, xp=xp) + + if (keepdims is False) and (p.shape[axis] != 1): + message = ("`keepdims` may be False only if `p` and `q` are scalars or the " + "length of `p` and `q` along `axis` is 1.") + raise ValueError(message) + keepdims = (p.shape[axis] != 1) if keepdims is None else keepdims + + x = xp.moveaxis(x, axis, -1) + q = xp.moveaxis(q, axis, -1) + p = xp.moveaxis(p, axis, -1) + + nan_out = ((p >= 1) | (p <= 0) | xp.isnan(p) | xp.isnan(q) + | xp.any(xp.isnan(x), axis=-1, keepdims=True)) + if is_lazy_array(nan_out) or xp.any(nan_out): + # These get NaN-ed out at the end. In the meantime, we want them + # to pass through calculations without warnings or errors + q = xpx.at(q, nan_out).set(0.0, copy=True) + p = xpx.at(p, nan_out).set(0.5, copy=True) + + return x, q, p, alternative, axis, keepdims, axis_none, ndim, nan_out, dtype, xp + + +def _quantile_test_postprocess(res, axis, axis_none, keepdims, ndim, nan_out, xp): + # Reshape per axis/keepdims + + res = xpx.at(res)[nan_out].set(xp.nan) + + if axis_none and keepdims: + shape = (1,)*(ndim - 1) + res.shape + res = xp.reshape(res, shape) + axis = -1 + + res = xp.moveaxis(res, -1, axis) + + if not keepdims: + res = xp.squeeze(res, axis=axis) + return res[()] if res.ndim == 0 else res + + +@xp_capabilities(skip_backends=[("dask.array", "no `take_along_axis`")], cpu_only=True, + reason="binomial distribution ufuncs only available for NumPy") +def quantile_test(x, *, q=0.0, p=0.5, alternative='two-sided', axis=0, keepdims=None): + r""" + Perform a quantile test and compute a confidence interval of the quantile. + + This function tests the null hypothesis that `q` is the value of the + quantile associated with probability `p` of the population underlying + sample `x`. For example, with default parameters, it tests that the + median of the population underlying `x` is zero. The function returns an + object including the test statistic, a p-value, and a method for computing + the confidence interval around the quantile. + + Parameters + ---------- + x : array_like + A one-dimensional sample. + q : float, default: 0 + The hypothesized value of the quantile. + p : float, default: 0.5 + The probability associated with the quantile; i.e. the proportion of + the population less than `q` is `p`. Must be strictly between 0 and 1. + alternative : {'two-sided', 'less', 'greater'}, optional + Defines the alternative hypothesis. + The following options are available (default is 'two-sided'): + + * 'two-sided': the quantile associated with the probability `p` + is not `q`. + * 'less': the quantile associated with the probability `p` is less + than `q`. + * 'greater': the quantile associated with the probability `p` is + greater than `q`. + axis : int or None, default: 0 + Axis of `x` along which the test is performed. + ``None`` ravels `x`, `q`, and `p` before performing the calculation, + without checking whether the original shapes were compatible. + As in other `scipy.stats` functions, a positive integer `axis` is resolved + after prepending 1s to the shape of `x`, `q`, and `p` as needed until the arrays + have the same dimensionality. When providing `x`, `q`, and `p` with different + dimensionality, consider using negative `axis` integers for clarity. + keepdims : bool, optional + Consider the case in which `x` is 1-D and `p` and `q` are scalars: the test + computes a reducing statistic, and the default behavior is to return a scalar. + If `keepdims` is set to True, the axis will not be reduced away, and the + result will be a 1-D array with one element. + + The general case is more subtle, since multiple tests may be + requested for each axis-slice of `x`. For instance, if `x`, `q`, and `p` + are 1-D and ``p.size == q.size > 1``, no axis can be reduced away; there must be + an axis to contain the results. Therefore: + + - By default, the axis will be reduced away if possible (i.e. if there is + exactly one element of `p` and `q` per axis-slice of `x`). + - If `keepdims` is set to True, the axis will not be reduced away. + - If `keepdims` is set to False, the axis will be reduced away + if possible, and an error will be raised otherwise. + + Returns + ------- + result : QuantileTestResult + An object with the following attributes: + + statistic : float + One of two test statistics that may be used in the quantile test. + The first test statistic, ``T1``, is the number of observations in + `x` that are less than or equal to the hypothesized quantile + `q`. The second test statistic, ``T2``, is the number of + observations in `x` that are strictly less than the hypothesized + quantile `q`. + + When ``alternative = 'greater'``, ``T1`` is used to calculate the + p-value and ``statistic`` is set to ``T1``. + + When ``alternative = 'less'``, ``T2`` is used to calculate the + p-value and ``statistic`` is set to ``T2``. + + When ``alternative = 'two-sided'``, both ``T1`` and ``T2`` are + considered, and the one that leads to the smallest p-value is used. + + statistic_type : float + Either ``1.`` or ``2.`` depending on which of ``T1`` or ``T2`` was + used to calculate the p-value. + + pvalue : float + The p-value associated with the given alternative. + + The object also has the following method: + + confidence_interval(confidence_level=0.95) + Computes a confidence interval around the the + population quantile associated with the probability `p`. The + confidence interval is returned in a ``namedtuple`` with + fields `low` and `high`. Values are `nan` when there are + not enough observations to compute the confidence interval at + the desired confidence. + + Notes + ----- + This test and its method for computing confidence intervals are + non-parametric. They are valid if and only if the observations are i.i.d. + + The implementation of the test follows Conover [1]_. Two test statistics + are considered. + + ``T1``: The number of observations in `x` less than or equal to `q`. + + ``T1 = (x <= q).sum()`` + + ``T2``: The number of observations in `x` strictly less than `q`. + + ``T2 = (x < q).sum()`` + + The use of two test statistics is necessary to handle the possibility that + `x` was generated from a discrete or mixed distribution. + + The null hypothesis for the test is: + + H0: The :math:`p^{\mathrm{th}}` population quantile is `q`. + + and the null distribution for each test statistic is + :math:`\mathrm{binom}\left(n, p\right)`. When ``alternative='less'``, + the alternative hypothesis is: + + H1: The :math:`p^{\mathrm{th}}` population quantile is less than `q`. + + and the p-value is the probability that the binomial random variable + + .. math:: + Y \sim \mathrm{binom}\left(n, p\right) + + is greater than or equal to the observed value ``T2``. + + When ``alternative='greater'``, the alternative hypothesis is: + + H1: The :math:`p^{\mathrm{th}}` population quantile is greater than `q` + + and the p-value is the probability that the binomial random variable Y + is less than or equal to the observed value ``T1``. + + When ``alternative='two-sided'``, the alternative hypothesis is + + H1: `q` is not the :math:`p^{\mathrm{th}}` population quantile. + + and the p-value is twice the smaller of the p-values for the ``'less'`` + and ``'greater'`` cases. Both of these p-values can exceed 0.5 for the same + data, so the value is clipped into the interval :math:`[0, 1]`. + + The approach for confidence intervals is attributed to Thompson [2]_ and + later proven to be applicable to any i.i.d. sample [3]_. The + computation is based on the observation that the probability of a quantile + :math:`q` to be larger than any observations :math:`x_m (1\leq m \leq N)` + can be computed as + + .. math:: + + \mathbb{P}(x_m \leq q) = 1 - \sum_{k=0}^{m-1} \binom{N}{k} + q^k(1-q)^{N-k} + + By default, confidence intervals are computed for a 95% confidence level. + A common interpretation of a 95% confidence intervals is that if i.i.d. + samples are drawn repeatedly from the same population and confidence + intervals are formed each time, the confidence interval will contain the + true value of the specified quantile in approximately 95% of trials. + + A similar function is available in the QuantileNPCI R package [4]_. The + foundation is the same, but it computes the confidence interval bounds by + doing interpolations between the sample values, whereas this function uses + only sample values as bounds. Thus, ``quantile_test.confidence_interval`` + returns more conservative intervals (i.e., larger). + + The same computation of confidence intervals for quantiles is included in + the confintr package [5]_. + + Two-sided confidence intervals are not guaranteed to be optimal; i.e., + there may exist a tighter interval that may contain the quantile of + interest with probability larger than the confidence level. + Without further assumption on the nature of the + underlying distribution, the one-sided intervals are optimally tight. + + References + ---------- + .. [1] W. J. Conover. Practical Nonparametric Statistics, 3rd Ed. 1999. + .. [2] W. R. Thompson, "On Confidence Ranges for the Median and Other + Expectation Distributions for Populations of Unknown Distribution + Form," The Annals of Mathematical Statistics, vol. 7, no. 3, + pp. 122-128, 1936, Accessed: Sep. 18, 2019. [Online]. Available: + https://www.jstor.org/stable/2957563. + .. [3] H. A. David and H. N. Nagaraja, "Order Statistics in Nonparametric + Inference" in Order Statistics, John Wiley & Sons, Ltd, 2005, pp. + 159-170. :doi:`10.1002/0471722162.ch7`. + .. [4] N. Hutson, A. Hutson, L. Yan, "QuantileNPCI: Nonparametric + Confidence Intervals for Quantiles," R package, + https://cran.r-project.org/package=QuantileNPCI + .. [5] M. Mayer, "confintr: Confidence Intervals," R package, + https://cran.r-project.org/package=confintr + + + Examples + -------- + + Suppose we wish to test the null hypothesis that the median of a population + is equal to 0.5. We choose a confidence level of 99%; that is, we will + reject the null hypothesis in favor of the alternative if the p-value is + less than 0.01. + + When testing random variates from the standard uniform distribution, which + has a median of 0.5, we expect the data to be consistent with the null + hypothesis most of the time. + + >>> import numpy as np + >>> from scipy import stats + >>> rng = np.random.default_rng(6981396440634228121) + >>> rvs = stats.uniform.rvs(size=100, random_state=rng) + >>> stats.quantile_test(rvs, q=0.5, p=0.5) + QuantileTestResult(statistic=45, statistic_type=1, pvalue=0.36820161732669576) + + As expected, the p-value is not below our threshold of 0.01, so + we cannot reject the null hypothesis. + + When testing data from the standard *normal* distribution, which has a + median of 0, we would expect the null hypothesis to be rejected. + + >>> rvs = stats.norm.rvs(size=100, random_state=rng) + >>> stats.quantile_test(rvs, q=0.5, p=0.5) + QuantileTestResult(statistic=67, statistic_type=2, pvalue=0.0008737198369123724) + + Indeed, the p-value is lower than our threshold of 0.01, so we reject the + null hypothesis in favor of the default "two-sided" alternative: the median + of the population is *not* equal to 0.5. + + However, suppose we were to test the null hypothesis against the + one-sided alternative that the median of the population is *greater* than + 0.5. Since the median of the standard normal is less than 0.5, we would not + expect the null hypothesis to be rejected. + + >>> stats.quantile_test(rvs, q=0.5, p=0.5, alternative='greater') + QuantileTestResult(statistic=67, statistic_type=1, pvalue=0.9997956114162866) + + Unsurprisingly, with a p-value greater than our threshold, we would not + reject the null hypothesis in favor of the chosen alternative. + + The quantile test can be used for any quantile, not only the median. For + example, we can test whether the third quartile of the distribution + underlying the sample is greater than 0.6. + + >>> rvs = stats.uniform.rvs(size=100, random_state=rng) + >>> stats.quantile_test(rvs, q=0.6, p=0.75, alternative='greater') + QuantileTestResult(statistic=64, statistic_type=1, pvalue=0.00940696592998271) + + The p-value is lower than the threshold. We reject the null hypothesis in + favor of the alternative: the third quartile of the distribution underlying + our sample is greater than 0.6. + + `quantile_test` can also compute confidence intervals for any quantile. + + >>> rvs = stats.norm.rvs(size=100, random_state=rng) + >>> res = stats.quantile_test(rvs, q=0.6, p=0.75) + >>> ci = res.confidence_interval(confidence_level=0.95) + >>> ci + ConfidenceInterval(low=0.284491604437432, high=0.8912531024914844) + + When testing a one-sided alternative, the confidence interval contains + all observations such that if passed as `q`, the p-value of the + test would be greater than 0.05, and therefore the null hypothesis + would not be rejected. For example: + + >>> rvs.sort() + >>> q, p, alpha = 0.6, 0.75, 0.95 + >>> res = stats.quantile_test(rvs, q=q, p=p, alternative='less') + >>> ci = res.confidence_interval(confidence_level=alpha) + >>> for x in rvs[rvs <= ci.high]: + ... res = stats.quantile_test(rvs, q=x, p=p, alternative='less') + ... assert res.pvalue > 1-alpha + >>> for x in rvs[rvs > ci.high]: + ... res = stats.quantile_test(rvs, q=x, p=p, alternative='less') + ... assert res.pvalue < 1-alpha + + Also, if a 95% confidence interval is repeatedly generated for random + samples, the confidence interval will contain the true quantile value in + approximately 95% of replications. + + >>> dist = stats.rayleigh() # our "unknown" distribution + >>> p = 0.2 + >>> true_stat = dist.ppf(p) # the true value of the statistic + >>> n_trials = 1000 + >>> quantile_ci_contains_true_stat = 0 + >>> for i in range(n_trials): + ... data = dist.rvs(size=100, random_state=rng) + ... res = stats.quantile_test(data, p=p) + ... ci = res.confidence_interval(0.95) + ... if ci[0] < true_stat < ci[1]: + ... quantile_ci_contains_true_stat += 1 + >>> quantile_ci_contains_true_stat >= 950 + True + + This works with any distribution and any quantile, as long as the observations + are i.i.d. + """ + # Implementation carefully follows [1] 3.2 + # "H0: the p*th quantile of X is x*" + # To facilitate comparison with [1], we'll use variable names that + # best match Conover's notation + X, x_star, p_star, H1, axis, keepdims, axis_none, ndim, nan_out, dtype, xp = ( + quantile_test_iv(x, q, p, alternative, axis, keepdims)) + # `axis` is the original `axis`; the working axis is -1. + + # "We will use two test statistics in this test. Let T1 equal " + # "the number of observations less than or equal to x*, and " + # "let T2 equal the number of observations less than x*." + if X.shape[-1] > 0: + T1 = _xp_searchsorted(X, x_star, side='right', xp=xp) + T2 = _xp_searchsorted(X, x_star, side='left', xp=xp) + T1, T2 = xp.astype(T1, dtype), xp.astype(T2, dtype) + else: + nan_out = xp.ones_like(nan_out) + T = xp.zeros(x_star.shape, dtype=dtype) + T1, T2 = T, T + + # "The null distribution of the test statistics T1 and T2 is " + # "the binomial distribution, with parameters n = sample size, and " + # "p = p* as given in the null hypothesis.... Y has the binomial " + # "distribution with parameters n and p*." + n = X.shape[-1] + Y = _SimpleBinomial(n=n, p=p_star) + + # "H1: the p* population quantile is less than x*" + if H1 == 'less': + # "The p-value is the probability that a binomial random variable Y " + # "is greater than *or equal to* the observed value of T2...using p=p*" + pvalue = Y.sf(T2-1) # Y.pmf(T2) + Y.sf(T2) + statistic = xp.broadcast_to(T2, pvalue.shape) + statistic_type = xp.full_like(statistic, 2.) + # "H1: the p* population quantile is greater than x*" + elif H1 == 'greater': + # "The p-value is the probability that a binomial random variable Y " + # "is less than or equal to the observed value of T1... using p = p*" + pvalue = Y.cdf(T1) + statistic = xp.broadcast_to(T1, pvalue.shape) + statistic_type = xp.full_like(statistic, 1.) + # "H1: x* is not the p*th population quantile" + elif H1 == 'two-sided': + # "The p-value is twice the smaller of the probabilities that a + # binomial random variable Y is less than or equal to the observed + # value of T1 or greater than or equal to the observed value of T2 + # using p=p*." + # Note: both one-sided p-values can exceed 0.5 for the same data, so + # `clip` + p_min = xp.minimum(Y.cdf(T1), Y.sf(T2 - 1)) + pvalue = xp.clip(2*p_min, 0., 1.) + i2 = Y.cdf(T1) > Y.sf(T2 - 1) + statistic = xp.where(i2, T2, T1) + statistic_type = xp.where(i2, xp.full_like(T2, 2.), xp.full_like(T1, 1.)) + + statistic = xp.astype(statistic, dtype) + statistic_type = xp.astype(statistic_type, dtype) + pvalue = xp.astype(pvalue, dtype) + X = xp.astype(X, dtype) + + _statistic, _statistic_type, _pvalue = statistic, statistic_type, pvalue + args = axis, axis_none, keepdims, ndim, nan_out, xp + statistic = _quantile_test_postprocess(statistic, *args) + statistic_type = _quantile_test_postprocess(statistic_type, *args) + pvalue = _quantile_test_postprocess(pvalue, *args) + + return QuantileTestResult( + statistic=statistic, + statistic_type=statistic_type, + pvalue=pvalue, + _alternative=H1, + _x=X, + _p=p_star, + _statistic=statistic, + _statistic_type=_statistic_type, + _pvalue=_pvalue, + _axis=axis, + _axis_none=axis_none, + _keepdims=keepdims, + _ndim=ndim, + _nan_out=nan_out, + _xp=xp, + ) + + +##################################### +# STATISTICAL DISTANCES # +##################################### + + +@xp_capabilities(np_only=True) +def wasserstein_distance_nd(u_values, v_values, u_weights=None, v_weights=None): + r""" + Compute the Wasserstein-1 distance between two N-D discrete distributions. + + The Wasserstein distance, also called the Earth mover's distance or the + optimal transport distance, is a similarity metric between two probability + distributions [1]_. In the discrete case, the Wasserstein distance can be + understood as the cost of an optimal transport plan to convert one + distribution into the other. The cost is calculated as the product of the + amount of probability mass being moved and the distance it is being moved. + A brief and intuitive introduction can be found at [2]_. + + .. versionadded:: 1.13.0 + + Parameters + ---------- + u_values : 2d array_like + A sample from a probability distribution or the support (set of all + possible values) of a probability distribution. Each element along + axis 0 is an observation or possible value, and axis 1 represents the + dimensionality of the distribution; i.e., each row is a vector + observation or possible value. + + v_values : 2d array_like + A sample from or the support of a second distribution. + + u_weights, v_weights : 1d array_like, optional + Weights or counts corresponding with the sample or probability masses + corresponding with the support values. Sum of elements must be positive + and finite. If unspecified, each value is assigned the same weight. + + Returns + ------- + distance : float + The computed distance between the distributions. + + Notes + ----- + Given two probability mass functions, :math:`u` + and :math:`v`, the first Wasserstein distance between the distributions + using the Euclidean norm is: + + .. math:: + + l_1 (u, v) = \inf_{\pi \in \Gamma (u, v)} \int \| x-y \|_2 \mathrm{d} \pi (x, y) + + where :math:`\Gamma (u, v)` is the set of (probability) distributions on + :math:`\mathbb{R}^n \times \mathbb{R}^n` whose marginals are :math:`u` and + :math:`v` on the first and second factors respectively. For a given value + :math:`x`, :math:`u(x)` gives the probability of :math:`u` at position + :math:`x`, and the same for :math:`v(x)`. + + This is also called the optimal transport problem or the Monge problem. + Let the finite point sets :math:`\{x_i\}` and :math:`\{y_j\}` denote + the support set of probability mass function :math:`u` and :math:`v` + respectively. The Monge problem can be expressed as follows, + + Let :math:`\Gamma` denote the transport plan, :math:`D` denote the + distance matrix and, + + .. math:: + + x = \text{vec}(\Gamma) \\ + c = \text{vec}(D) \\ + b = \begin{bmatrix} + u\\ + v\\ + \end{bmatrix} + + The :math:`\text{vec}()` function denotes the Vectorization function + that transforms a matrix into a column vector by vertically stacking + the columns of the matrix. + The transport plan :math:`\Gamma` is a matrix :math:`[\gamma_{ij}]` in + which :math:`\gamma_{ij}` is a positive value representing the amount of + probability mass transported from :math:`u(x_i)` to :math:`v(y_i)`. + Summing over the rows of :math:`\Gamma` should give the source distribution + :math:`u` : :math:`\sum_j \gamma_{ij} = u(x_i)` holds for all :math:`i` + and summing over the columns of :math:`\Gamma` should give the target + distribution :math:`v`: :math:`\sum_i \gamma_{ij} = v(y_j)` holds for all + :math:`j`. + The distance matrix :math:`D` is a matrix :math:`[d_{ij}]`, in which + :math:`d_{ij} = d(x_i, y_j)`. + + Given :math:`\Gamma`, :math:`D`, :math:`b`, the Monge problem can be + transformed into a linear programming problem by + taking :math:`A x = b` as constraints and :math:`z = c^T x` as minimization + target (sum of costs) , where matrix :math:`A` has the form + + .. math:: + + \begin{array} {rrrr|rrrr|r|rrrr} + 1 & 1 & \dots & 1 & 0 & 0 & \dots & 0 & \dots & 0 & 0 & \dots & + 0 \cr + 0 & 0 & \dots & 0 & 1 & 1 & \dots & 1 & \dots & 0 & 0 &\dots & + 0 \cr + \vdots & \vdots & \ddots & \vdots & \vdots & \vdots & \ddots + & \vdots & \vdots & \vdots & \vdots & \ddots & \vdots \cr + 0 & 0 & \dots & 0 & 0 & 0 & \dots & 0 & \dots & 1 & 1 & \dots & + 1 \cr \hline + + 1 & 0 & \dots & 0 & 1 & 0 & \dots & \dots & \dots & 1 & 0 & \dots & + 0 \cr + 0 & 1 & \dots & 0 & 0 & 1 & \dots & \dots & \dots & 0 & 1 & \dots & + 0 \cr + \vdots & \vdots & \ddots & \vdots & \vdots & \vdots & \ddots & + \vdots & \vdots & \vdots & \vdots & \ddots & \vdots \cr + 0 & 0 & \dots & 1 & 0 & 0 & \dots & 1 & \dots & 0 & 0 & \dots & 1 + \end{array} + + By solving the dual form of the above linear programming problem (with + solution :math:`y^*`), the Wasserstein distance :math:`l_1 (u, v)` can + be computed as :math:`b^T y^*`. + + The above solution is inspired by Vincent Herrmann's blog [3]_ . For a + more thorough explanation, see [4]_ . + + The input distributions can be empirical, therefore coming from samples + whose values are effectively inputs of the function, or they can be seen as + generalized functions, in which case they are weighted sums of Dirac delta + functions located at the specified values. + + References + ---------- + .. [1] "Wasserstein metric", + https://en.wikipedia.org/wiki/Wasserstein_metric + .. [2] Lili Weng, "What is Wasserstein distance?", Lil'log, + https://lilianweng.github.io/posts/2017-08-20-gan/#what-is-wasserstein-distance. + .. [3] Hermann, Vincent. "Wasserstein GAN and the Kantorovich-Rubinstein + Duality". https://vincentherrmann.github.io/blog/wasserstein/. + .. [4] Peyré, Gabriel, and Marco Cuturi. "Computational optimal + transport." Center for Research in Economics and Statistics + Working Papers 2017-86 (2017). + + See Also + -------- + wasserstein_distance: Compute the Wasserstein-1 distance between two + 1D discrete distributions. + + Examples + -------- + Compute the Wasserstein distance between two three-dimensional samples, + each with two observations. + + >>> from scipy.stats import wasserstein_distance_nd + >>> wasserstein_distance_nd([[0, 2, 3], [1, 2, 5]], [[3, 2, 3], [4, 2, 5]]) + 3.0 + + Compute the Wasserstein distance between two two-dimensional distributions + with three and two weighted observations, respectively. + + >>> wasserstein_distance_nd([[0, 2.75], [2, 209.3], [0, 0]], + ... [[0.2, 0.322], [4.5, 25.1808]], + ... [0.4, 5.2, 0.114], [0.8, 1.5]) + 174.15840245217169 + """ + m, n = len(u_values), len(v_values) + u_values = asarray(u_values) + v_values = asarray(v_values) + + if u_values.ndim > 2 or v_values.ndim > 2: + raise ValueError('Invalid input values. The inputs must have either ' + 'one or two dimensions.') + # if dimensions are not equal throw error + if u_values.ndim != v_values.ndim: + raise ValueError('Invalid input values. Dimensions of inputs must be ' + 'equal.') + # if data is 1D then call the cdf_distance function + if u_values.ndim == 1 and v_values.ndim == 1: + return _cdf_distance(1, u_values, v_values, u_weights, v_weights) + + u_values, u_weights = _validate_distribution(u_values, u_weights) + v_values, v_weights = _validate_distribution(v_values, v_weights) + # if number of columns is not equal throw error + if u_values.shape[1] != v_values.shape[1]: + raise ValueError('Invalid input values. If two-dimensional, ' + '`u_values` and `v_values` must have the same ' + 'number of columns.') + + # if data contains np.inf then return inf or nan + if np.any(np.isinf(u_values)) ^ np.any(np.isinf(v_values)): + return np.inf + elif np.any(np.isinf(u_values)) and np.any(np.isinf(v_values)): + return np.nan + + # create constraints + A_upper_part = sparse.block_diag((sparse.coo_array(np.ones((1, n))), ) * m) + A_lower_part = sparse.hstack((sparse.eye_array(n), ) * m) + # sparse constraint matrix of size (m + n)*(m * n) + A = sparse.vstack((A_upper_part, A_lower_part)) + + # get cost matrix + D = cdist(u_values, v_values) + cost = D.ravel() + + # create the minimization target + p_u = np.full(m, 1/m) if u_weights is None else u_weights/np.sum(u_weights) + p_v = np.full(n, 1/n) if v_weights is None else v_weights/np.sum(v_weights) + b = np.concatenate((p_u, p_v), axis=0) + + # solving LP + constraints = LinearConstraint(A=A.T, ub=cost) + opt_res = milp(c=-b, constraints=constraints, bounds=(-np.inf, np.inf)) + return -opt_res.fun + + +@xp_capabilities(np_only=True) +def wasserstein_distance(u_values, v_values, u_weights=None, v_weights=None): + r""" + Compute the Wasserstein-1 distance between two 1D discrete distributions. + + The Wasserstein distance, also called the Earth mover's distance or the + optimal transport distance, is a similarity metric between two probability + distributions [1]_. In the discrete case, the Wasserstein distance can be + understood as the cost of an optimal transport plan to convert one + distribution into the other. The cost is calculated as the product of the + amount of probability mass being moved and the distance it is being moved. + A brief and intuitive introduction can be found at [2]_. + + .. versionadded:: 1.0.0 + + Parameters + ---------- + u_values : 1d array_like + A sample from a probability distribution or the support (set of all + possible values) of a probability distribution. Each element is an + observation or possible value. + + v_values : 1d array_like + A sample from or the support of a second distribution. + + u_weights, v_weights : 1d array_like, optional + Weights or counts corresponding with the sample or probability masses + corresponding with the support values. Sum of elements must be positive + and finite. If unspecified, each value is assigned the same weight. + + Returns + ------- + distance : float + The computed distance between the distributions. + + Notes + ----- + Given two 1D probability mass functions, :math:`u` and :math:`v`, the first + Wasserstein distance between the distributions is: + + .. math:: + + l_1 (u, v) = \inf_{\pi \in \Gamma (u, v)} \int_{\mathbb{R} \times + \mathbb{R}} |x-y| \mathrm{d} \pi (x, y) + + where :math:`\Gamma (u, v)` is the set of (probability) distributions on + :math:`\mathbb{R} \times \mathbb{R}` whose marginals are :math:`u` and + :math:`v` on the first and second factors respectively. For a given value + :math:`x`, :math:`u(x)` gives the probability of :math:`u` at position + :math:`x`, and the same for :math:`v(x)`. + + If :math:`U` and :math:`V` are the respective CDFs of :math:`u` and + :math:`v`, this distance also equals to: + + .. math:: + + l_1(u, v) = \int_{-\infty}^{+\infty} |U-V| + + See [3]_ for a proof of the equivalence of both definitions. + + The input distributions can be empirical, therefore coming from samples + whose values are effectively inputs of the function, or they can be seen as + generalized functions, in which case they are weighted sums of Dirac delta + functions located at the specified values. + + References + ---------- + .. [1] "Wasserstein metric", https://en.wikipedia.org/wiki/Wasserstein_metric + .. [2] Lili Weng, "What is Wasserstein distance?", Lil'log, + https://lilianweng.github.io/posts/2017-08-20-gan/#what-is-wasserstein-distance. + .. [3] Ramdas, Garcia, Cuturi "On Wasserstein Two Sample Testing and Related + Families of Nonparametric Tests" (2015). :arXiv:`1509.02237`. + + See Also + -------- + wasserstein_distance_nd: Compute the Wasserstein-1 distance between two N-D + discrete distributions. + + Examples + -------- + >>> from scipy.stats import wasserstein_distance + >>> wasserstein_distance([0, 1, 3], [5, 6, 8]) + 5.0 + >>> wasserstein_distance([0, 1], [0, 1], [3, 1], [2, 2]) + 0.25 + >>> wasserstein_distance([3.4, 3.9, 7.5, 7.8], [4.5, 1.4], + ... [1.4, 0.9, 3.1, 7.2], [3.2, 3.5]) + 4.0781331438047861 + + """ + return _cdf_distance(1, u_values, v_values, u_weights, v_weights) + + +@xp_capabilities(np_only=True) +def energy_distance(u_values, v_values, u_weights=None, v_weights=None): + r"""Compute the energy distance between two 1D distributions. + + .. versionadded:: 1.0.0 + + Parameters + ---------- + u_values, v_values : array_like + Values observed in the (empirical) distribution. + u_weights, v_weights : array_like, optional + Weight for each value. If unspecified, each value is assigned the same + weight. + `u_weights` (resp. `v_weights`) must have the same length as + `u_values` (resp. `v_values`). If the weight sum differs from 1, it + must still be positive and finite so that the weights can be normalized + to sum to 1. + + Returns + ------- + distance : float + The computed distance between the distributions. + + Notes + ----- + The energy distance between two distributions :math:`u` and :math:`v`, whose + respective CDFs are :math:`U` and :math:`V`, equals to: + + .. math:: + + D(u, v) = \left( 2\mathbb E|X - Y| - \mathbb E|X - X'| - + \mathbb E|Y - Y'| \right)^{1/2} + + where :math:`X` and :math:`X'` (resp. :math:`Y` and :math:`Y'`) are + independent random variables whose probability distribution is :math:`u` + (resp. :math:`v`). + + Sometimes the square of this quantity is referred to as the "energy + distance" (e.g. in [2]_, [4]_), but as noted in [1]_ and [3]_, only the + definition above satisfies the axioms of a distance function (metric). + + As shown in [2]_, for one-dimensional real-valued variables, the energy + distance is linked to the non-distribution-free version of the Cramér-von + Mises distance: + + .. math:: + + D(u, v) = \sqrt{2} l_2(u, v) = \left( 2 \int_{-\infty}^{+\infty} (U-V)^2 + \right)^{1/2} + + Note that the common Cramér-von Mises criterion uses the distribution-free + version of the distance. See [2]_ (section 2), for more details about both + versions of the distance. + + The input distributions can be empirical, therefore coming from samples + whose values are effectively inputs of the function, or they can be seen as + generalized functions, in which case they are weighted sums of Dirac delta + functions located at the specified values. + + References + ---------- + .. [1] Rizzo, Szekely "Energy distance." Wiley Interdisciplinary Reviews: + Computational Statistics, 8(1):27-38 (2015). + .. [2] Szekely "E-statistics: The energy of statistical samples." Bowling + Green State University, Department of Mathematics and Statistics, + Technical Report 02-16 (2002). + .. [3] "Energy distance", https://en.wikipedia.org/wiki/Energy_distance + .. [4] Bellemare, Danihelka, Dabney, Mohamed, Lakshminarayanan, Hoyer, + Munos "The Cramer Distance as a Solution to Biased Wasserstein + Gradients" (2017). :arXiv:`1705.10743`. + + Examples + -------- + >>> from scipy.stats import energy_distance + >>> energy_distance([0], [2]) + 2.0000000000000004 + >>> energy_distance([0, 8], [0, 8], [3, 1], [2, 2]) + 1.0000000000000002 + >>> energy_distance([0.7, 7.4, 2.4, 6.8], [1.4, 8. ], + ... [2.1, 4.2, 7.4, 8. ], [7.6, 8.8]) + 0.88003340976158217 + + """ + return np.sqrt(2) * _cdf_distance(2, u_values, v_values, + u_weights, v_weights) + + +def _cdf_distance(p, u_values, v_values, u_weights=None, v_weights=None): + r""" + Compute, between two one-dimensional distributions :math:`u` and + :math:`v`, whose respective CDFs are :math:`U` and :math:`V`, the + statistical distance that is defined as: + + .. math:: + + l_p(u, v) = \left( \int_{-\infty}^{+\infty} |U-V|^p \right)^{1/p} + + p is a positive parameter; p = 1 gives the Wasserstein distance, p = 2 + gives the energy distance. + + Parameters + ---------- + u_values, v_values : array_like + Values observed in the (empirical) distribution. + u_weights, v_weights : array_like, optional + Weight for each value. If unspecified, each value is assigned the same + weight. + `u_weights` (resp. `v_weights`) must have the same length as + `u_values` (resp. `v_values`). If the weight sum differs from 1, it + must still be positive and finite so that the weights can be normalized + to sum to 1. + + Returns + ------- + distance : float + The computed distance between the distributions. + + Notes + ----- + The input distributions can be empirical, therefore coming from samples + whose values are effectively inputs of the function, or they can be seen as + generalized functions, in which case they are weighted sums of Dirac delta + functions located at the specified values. + + References + ---------- + .. [1] Bellemare, Danihelka, Dabney, Mohamed, Lakshminarayanan, Hoyer, + Munos "The Cramer Distance as a Solution to Biased Wasserstein + Gradients" (2017). :arXiv:`1705.10743`. + + """ + u_values, u_weights = _validate_distribution(u_values, u_weights) + v_values, v_weights = _validate_distribution(v_values, v_weights) + + u_sorter = np.argsort(u_values) + v_sorter = np.argsort(v_values) + + all_values = np.concatenate((u_values, v_values)) + all_values.sort(kind='mergesort') + + # Compute the differences between pairs of successive values of u and v. + deltas = np.diff(all_values) + + # Get the respective positions of the values of u and v among the values of + # both distributions. + u_cdf_indices = u_values[u_sorter].searchsorted(all_values[:-1], 'right') + v_cdf_indices = v_values[v_sorter].searchsorted(all_values[:-1], 'right') + + # Calculate the CDFs of u and v using their weights, if specified. + if u_weights is None: + u_cdf = u_cdf_indices / u_values.size + else: + u_sorted_cumweights = np.concatenate(([0], + np.cumsum(u_weights[u_sorter]))) + u_cdf = u_sorted_cumweights[u_cdf_indices] / u_sorted_cumweights[-1] + + if v_weights is None: + v_cdf = v_cdf_indices / v_values.size + else: + v_sorted_cumweights = np.concatenate(([0], + np.cumsum(v_weights[v_sorter]))) + v_cdf = v_sorted_cumweights[v_cdf_indices] / v_sorted_cumweights[-1] + + # Compute the value of the integral based on the CDFs. + # If p = 1 or p = 2, we avoid using np.power, which introduces an overhead + # of about 15%. + if p == 1: + return np.vecdot(np.abs(u_cdf - v_cdf), deltas) + if p == 2: + return np.sqrt(np.vecdot(np.square(u_cdf - v_cdf), deltas)) + return np.power(np.vecdot(np.power(np.abs(u_cdf - v_cdf), p), deltas), 1/p) + + +def _validate_distribution(values, weights): + """ + Validate the values and weights from a distribution input of `cdf_distance` + and return them as ndarray objects. + + Parameters + ---------- + values : array_like + Values observed in the (empirical) distribution. + weights : array_like + Weight for each value. + + Returns + ------- + values : ndarray + Values as ndarray. + weights : ndarray + Weights as ndarray. + + """ + # Validate the value array. + values = np.asarray(values, dtype=float) + if len(values) == 0: + raise ValueError("Distribution can't be empty.") + + # Validate the weight array, if specified. + if weights is not None: + weights = np.asarray(weights, dtype=float) + if len(weights) != len(values): + raise ValueError('Value and weight array-likes for the same ' + 'empirical distribution must be of the same size.') + if np.any(weights < 0): + raise ValueError('All weights must be non-negative.') + if not 0 < np.sum(weights) < np.inf: + raise ValueError('Weight array-like sum must be positive and ' + 'finite. Set as None for an equal distribution of ' + 'weight.') + + return values, weights + + return values, None + + +@xp_capabilities(skip_backends=[("dask.array", "no `take_along_axis`")], + marray=True) +def rankdata(a, method='average', *, axis=None, nan_policy='propagate'): + """Assign ranks to data, dealing with ties appropriately. + + By default (``axis=None``), the data array is first flattened, and a flat + array of ranks is returned. Separately reshape the rank array to the + shape of the data array if desired (see Examples). + + Ranks begin at 1. The `method` argument controls how ranks are assigned + to equal values. See [1]_ for further discussion of ranking methods. + + Parameters + ---------- + a : array_like + The array of values to be ranked. + method : {'average', 'min', 'max', 'dense', 'ordinal'}, optional + The method used to assign ranks to tied elements. + The following methods are available (default is 'average'): + + * 'average': The average of the ranks that would have been assigned to + all the tied values is assigned to each value. + * 'min': The minimum of the ranks that would have been assigned to all + the tied values is assigned to each value. (This is also + referred to as "competition" ranking.) + * 'max': The maximum of the ranks that would have been assigned to all + the tied values is assigned to each value. + * 'dense': Like 'min', but the rank of the next highest element is + assigned the rank immediately after those assigned to the tied + elements. + * 'ordinal': All values are given a distinct rank, corresponding to + the order that the values occur in `a`. + + axis : {None, int}, optional + Axis along which to perform the ranking. If ``None``, the data array + is first flattened. + nan_policy : {'propagate', 'omit', 'raise'}, optional + Defines how to handle when input contains nan. + The following options are available (default is 'propagate'): + + * 'propagate': propagates nans through the rank calculation + * 'omit': performs the calculations ignoring nan values + * 'raise': raises an error + + .. note:: + + When `nan_policy` is 'propagate', the output is an array of *all* + nans because ranks relative to nans in the input are undefined. + When `nan_policy` is 'omit', nans in `a` are ignored when ranking + the other values, and the corresponding locations of the output + are nan. + + .. versionadded:: 1.10 + + Returns + ------- + ranks : ndarray + An array of size equal to the size of `a`, containing rank + scores. The dtype is the result dtype of `a` and a Python float. + + References + ---------- + .. [1] "Ranking", https://en.wikipedia.org/wiki/Ranking + + Examples + -------- + >>> import numpy as np + >>> from scipy.stats import rankdata + >>> rankdata([0, 2, 3, 2]) + array([1. , 2.5, 4. , 2.5]) + >>> rankdata([0, 2, 3, 2], method='min') + array([1., 2., 4., 2.]) + >>> rankdata([0, 2, 3, 2], method='max') + array([1., 3., 4., 3.]) + >>> rankdata([0, 2, 3, 2], method='dense') + array([1., 2., 3., 2.]) + >>> rankdata([0, 2, 3, 2], method='ordinal') + array([1., 2., 4., 3.]) + >>> rankdata([[0, 2], [3, 2]]).reshape(2, 2) + array([[1. , 2.5], + [4. , 2.5]]) + >>> rankdata([[0, 2, 2], [3, 2, 5]], axis=1) + array([[1. , 2.5, 2.5], + [2. , 1. , 3. ]]) + >>> rankdata([0, 2, 3, np.nan, -2, np.nan], nan_policy="propagate") + array([nan, nan, nan, nan, nan, nan]) + >>> rankdata([0, 2, 3, np.nan, -2, np.nan], nan_policy="omit") + array([ 2., 3., 4., nan, 1., nan]) + + """ + methods = ('average', 'min', 'max', 'dense', 'ordinal') + if method not in methods: + raise ValueError(f'unknown method "{method}"') + + xp = array_namespace(a) + x = xp.asarray(a) + + if axis is None: + x = xp_ravel(x) + axis = -1 + + if xp_size(x) == 0: + dtype = xp_result_type(x, force_floating=True, xp=xp) + return xp.empty_like(x, dtype=dtype) + + contains_nan = _contains_nan(x, nan_policy) + + x = xp_swapaxes(x, axis, -1, xp=xp) + ranks, _, _ = _rankdata(x, method, xp=xp) + + # JIT won't allow use of `contains_nan` for control flow here, so we always have to + # run this with JIT. + if is_lazy_array(x) or contains_nan: + i_nan = (xp.isnan(x) if nan_policy == 'omit' + else xp.any(xp.isnan(x), axis=-1, keepdims=True)) + i_nan = xp.broadcast_to(i_nan, ranks.shape) + ranks = xpx.at(ranks)[i_nan].set(xp.nan) + + ranks = xp_swapaxes(ranks, axis, -1, xp=xp) + return ranks + + +def _order_ranks(ranks, j, *, xp): + # Reorder ascending order `ranks` according to `j` + xp = array_namespace(ranks) if xp is None else xp + if is_numpy(xp) or is_cupy(xp): + ordered_ranks = xp.empty(j.shape, dtype=ranks.dtype) + xp.put_along_axis(ordered_ranks, j, ranks, axis=-1) + else: + # `put_along_axis` not in array API (data-apis/array-api#177) + # so argsort the argsort and take_along_axis... + j_inv = xp.argsort(j, axis=-1, stable=True) + ordered_ranks = xp.take_along_axis(ranks, j_inv, axis=-1) + return ordered_ranks + + +def _rankdata(x, method, return_sorted=False, return_ties=False, xp=None): + # Rank data `x` by desired `method`. For methods other than 'ordinal': + # - `return_sorted=True` ensures that the second output is sorted `x` + # - `return_ties=True` ensures that the third output is tie data + # Otherwise, the second/third output will be None. + xp = array_namespace(x) if xp is None else xp + dtype = xp_result_type(x, force_floating=True, xp=xp) + + if is_jax(xp): + import jax.scipy.stats as jax_stats + ranks = jax_stats.rankdata(x, method=method, axis=-1) + ranks = xp.astype(ranks, dtype, copy=False) + y = xp.sort(x, axis=-1) if (return_ties or return_sorted) else None + t = None + if return_ties: + max_ranks = jax_stats.rankdata(y, method='max', axis=-1) + t = xp.diff(max_ranks, axis=-1, prepend=0) + t = xp.astype(t, dtype, copy=False) + return ranks, y, t + + if is_marray(xp): + data, mask = x.data, x.mask + uxp = array_namespace(data) + data = uxp.where(mask, uxp.nan, data) + ranks, sorted, ties = _rankdata(data, method, return_sorted, return_ties, uxp) + ranks = xp.asarray(ranks, mask=mask) + mask_sorted = uxp.isnan(sorted) if (return_sorted or return_ties) else None + sorted = xp.asarray(sorted, mask=mask_sorted) if return_sorted else None + ties = xp.asarray(ties, mask=mask_sorted) if return_ties else None + return ranks, sorted, ties + + shape = x.shape + + # Get sort order + j = xp.argsort(x, axis=-1, stable=True) + ordinal_ranks = xp.broadcast_to(xp.arange(1, shape[-1]+1, dtype=dtype), shape) + + # Ordinal ranks is very easy because ties don't matter. We're done. + if method == 'ordinal': + ranks = _order_ranks(ordinal_ranks, j, xp=xp) + return (ranks, None, None) + + # Sort array + y = xp.take_along_axis(x, j, axis=-1) + # Logical indices of unique elements + i = xp.concat([xp.ones(shape[:-1] + (1,), dtype=xp.bool), + y[..., :-1] != y[..., 1:]], axis=-1) + + # Integer indices of unique elements + indices = xp.arange(xp_size(y))[xp.reshape(i, (-1,))] # i gets raveled + # Counts of unique elements + counts = xp.diff(indices, append=xp.asarray([xp_size(y)], dtype=indices.dtype)) + + # Compute `'min'`, `'max'`, and `'mid'` ranks of unique elements + if method == 'min': + ranks = ordinal_ranks[i] + elif method == 'max': + ranks = ordinal_ranks[i] + xp.astype(counts, dtype) - 1 + elif method == 'average': + # array API doesn't promote integers to floats + ranks = ordinal_ranks[i] + (xp.astype(counts, dtype) - 1)/2 + elif method == 'dense': + ranks = xp.cumulative_sum(xp.astype(i, dtype, copy=False), axis=-1)[i] + + if not is_cupy(xp): + ranks = xp.reshape(xp.repeat(ranks, counts), shape) + else: + # workaround for cupy.repeat not accepting arrays for repeats + # (ref https://github.com/cupy/cupy/issues/3849). + # The cumulative sum over i will + # increment every time a new unique rank appears, giving the correct + # indices to replicate xp.repeats. + ranks = xp.reshape( + ranks[xp.cumulative_sum(xp.astype(xp.reshape(i, (-1,)), xp.int64)) - 1], + shape, + ) + ranks = _order_ranks(ranks, j, xp=xp) + + t = None + if return_ties: + # Tie information is returned in a format that is useful to functions that + # rely on this (private) function. Example: + # >>> x = np.asarray([3, 2, 1, 2, 2, 2, 1]) + # >>> _, t = _rankdata(x, 'average', return_ties=True) + # >>> t # array([2., 0., 4., 0., 0., 0., 1.]) # two 1s, four 2s, and one 3 + # Unlike ranks, tie counts are *not* reordered to correspond with the order of + # the input; e.g. the number of appearances of the lowest rank element comes + # first. This is a useful format because: + # - The shape of the result is the shape of the input. Different slices can + # have different numbers of tied elements but not result in a ragged array. + # - Functions that use `t` usually don't need to which each element of the + # original array is associated with each tie count; they perform a reduction + # over the tie counts onnly. The tie counts are naturally computed in a + # sorted order, so this does not unnecessarily reorder them. + # - One exception is `wilcoxon`, which needs the number of zeros. Zeros always + # have the lowest rank, so it is easy to find them at the zeroth index. + t = xp.zeros(shape, dtype=dtype) + t = xpx.at(t)[i].set(xp.astype(counts, dtype, copy=False)) + + return ranks, y, t + + +@xp_capabilities(marray=True, + skip_backends=[('dask.array', 'uses `optimize.elementwise.find_root`'), + ('array_api_strict', 'uses `optimize.elementwise.find_root`'), + ('jax.numpy', 'uses `optimize.elementwise.find_root`')]) +@_axis_nan_policy_factory( + lambda x: x, n_outputs=1, result_to_tuple=lambda x, _: (x,), + kwd_samples=['weights'], paired=True +) +def expectile(a, alpha=0.5, *, weights=None, axis=None): + r"""Compute the expectile at the specified level. + + Expectiles are a generalization of the expectation in the same way as + quantiles are a generalization of the median. The expectile at level + `alpha = 0.5` is the mean (average). See Notes for more details. + + Parameters + ---------- + a : array_like + Array containing numbers whose expectile is desired. + alpha : float, default: 0.5 + The level of the expectile; e.g., ``alpha=0.5`` gives the mean. + weights : array_like, optional + An array of weights associated with the values in `a`. + The `weights` must be broadcastable to the same shape as `a`. + Default is None, which gives each value a weight of 1.0. + An integer valued weight element acts like repeating the corresponding + observation in `a` that many times. See Notes for more details. + axis : int or tuple of ints, default: None + If an int or tuple of ints, the axis or axes of the input along which + to compute the statistic. The statistic of each axis-slice (e.g. row) + of the input will appear in a corresponding element of the output. + If ``None``, the input will be raveled before computing the statistic. + + Returns + ------- + expectile : ndarray + The empirical expectile at level `alpha`. + + See Also + -------- + numpy.mean : Arithmetic average + numpy.quantile : Quantile + + Notes + ----- + In general, the expectile at level :math:`\alpha` of a random variable + :math:`X` with cumulative distribution function (CDF) :math:`F` is given + by the unique solution :math:`t` of: + + .. math:: + + \alpha E((X - t)_+) = (1 - \alpha) E((t - X)_+) \,. + + Here, :math:`(x)_+ = \max(0, x)` is the positive part of :math:`x`. + This equation can be equivalently written as: + + .. math:: + + \alpha \int_t^\infty (x - t)\mathrm{d}F(x) + = (1 - \alpha) \int_{-\infty}^t (t - x)\mathrm{d}F(x) \,. + + The empirical expectile at level :math:`\alpha` (`alpha`) of a sample + :math:`a_i` (the array `a`) is defined by plugging in the empirical CDF of + `a`. Given sample or case weights :math:`w` (the array `weights`), it + reads :math:`F_a(x) = \frac{1}{\sum_i w_i} \sum_i w_i 1_{a_i \leq x}` + with indicator function :math:`1_{A}`. This leads to the definition of the + empirical expectile at level `alpha` as the unique solution :math:`t` of: + + .. math:: + + \alpha \sum_{i=1}^n w_i (a_i - t)_+ = + (1 - \alpha) \sum_{i=1}^n w_i (t - a_i)_+ \,. + + For :math:`\alpha=0.5`, this simplifies to the weighted average. + Furthermore, the larger :math:`\alpha`, the larger the value of the + expectile. + + As a final remark, the expectile at level :math:`\alpha` can also be + written as a minimization problem. One often used choice is + + .. math:: + + \operatorname{argmin}_t + E(\lvert 1_{t\geq X} - \alpha\rvert(t - X)^2) \,. + + References + ---------- + .. [1] W. K. Newey and J. L. Powell (1987), "Asymmetric Least Squares + Estimation and Testing," Econometrica, 55, 819-847. + .. [2] T. Gneiting (2009). "Making and Evaluating Point Forecasts," + Journal of the American Statistical Association, 106, 746 - 762. + :doi:`10.48550/arXiv.0912.0902` + + Examples + -------- + >>> import numpy as np + >>> from scipy.stats import expectile + >>> a = [1, 4, 2, -1] + >>> expectile(a, alpha=0.5) == np.mean(a) + True + >>> expectile(a, alpha=0.2) + 0.42857142857142855 + >>> expectile(a, alpha=0.8) + 2.5714285714285716 + >>> weights = [1, 3, 1, 1] + >>> expectile(a, alpha=0.8, weights=weights) + 3.3333333333333335 + """ + xp = array_namespace(a, weights) + a, weights = xp_promote(a, weights, force_floating=True, xp=xp) + + if alpha < 0 or alpha > 1: + raise ValueError( + "The expectile level alpha must be in the range [0, 1]." + ) + + if a.shape[-1] == 0: # Only for *testing* _axis_nan_policy decorator + return _get_nan(a, xp=xp) + + # for simplicity, ensure that shape is (batch size, sample size) and at least 2d + shape = a.shape # remember original shape, though, to ensure correct output shape + # `reshape` can't infer shape element -1 when array is size 0, so take `prod` + a = xp.reshape(a, (math.prod(shape[:-1]), shape[-1])) + weights = weights if weights is None else xp.reshape(weights, (-1, shape[-1])) + + # This is the empirical equivalent of Eq. (13) with identification + # function from Table 9 (omitting a factor of 2) in [2] (their y is our + # data a, their x is our t) + def first_order(t, i): + ai = a[i, ...].T + wi = weights if weights is None else weights[i, ...].T + return _xp_mean(xp.abs(xp.astype(ai <= t, ai.dtype) - alpha) * (t - ai), + weights=wi, axis=0) + + # axis is -1 per the decorator. Use mean to produce NaNs of the correct shape/type. + x0 = xp.min(a, axis=-1) + x1 = xp.max(a, axis=-1) + + # Note that the expectile is the unique solution, so no worries about + # finding a wrong root. + i = xp.arange(a.shape[0]) + res = find_root(first_order, (x0, x1), args=(i,)) + res = xp.reshape(res.x, shape[:-1]) + return res[()] if res.ndim == 0 else res + + +def _lmoment_iv(sample, order, axis, sorted, standardize, xp): + # input validation/standardization for `lmoment` + sample = xp_promote(sample, force_floating=True, xp=xp) + + message = "`sample` must be an array of real numbers." + if not xp.isdtype(sample.dtype, "real floating"): + raise ValueError(message) + + message = "`order` must be a scalar or a non-empty array of positive integers." + order = xp.arange(1, 5) if order is None else xp.asarray(order) + if (not xp.isdtype(order.dtype, "integral") or order.size == 0 or order.ndim > 1 + or (not is_lazy_array(order) and xp.any(order <= 0))): + raise ValueError(message) + + # input validation of non-array types can still be performed with NumPy + axis = np.asarray(axis)[()] + message = "`axis` must be an integer." + if not np.issubdtype(axis.dtype, np.integer) or axis.ndim != 0: + raise ValueError(message) + axis = int(axis) + + sorted = np.asarray(sorted)[()] + message = "`sorted` must be True or False." + if not np.issubdtype(sorted.dtype, np.bool_) or sorted.ndim != 0: + raise ValueError(message) + sorted = bool(sorted) + + standardize = np.asarray(standardize)[()] + message = "`standardize` must be True or False." + if not np.issubdtype(standardize.dtype, np.bool_) or standardize.ndim != 0: + raise ValueError(message) + standardize = bool(standardize) + + sample = xp.moveaxis(sample, axis, -1) + sample = xp.sort(sample, axis=-1) if not sorted else sample + + return sample, order, axis, sorted, standardize + + +def _br(x, *, r=0, xp): + n = x.shape[-1] + x = xp.expand_dims(x, axis=-2) + x = xp.broadcast_to(x, x.shape[:-2] + (r.shape[0], n)) + x = xp.triu(x) + j = xp.arange(n, dtype=x.dtype) + n = xp.asarray(n, dtype=x.dtype)[()] + binom_j_r = xpx.lazy_apply(special.binom, j, r[:, xp.newaxis]) + binom_nm1_r = xpx.lazy_apply(special.binom, n-1, r) + return xp.vecdot(binom_j_r, x, axis=-1) / binom_nm1_r / n + + +def _prk(r, k): + # Writen to match [1] Equation 27 closely to facilitate review. + # This does not protect against overflow, so improvements to + # robustness would be a welcome follow-up. + binom_r_k = xpx.lazy_apply(special.binom, r, k) + binom_rpk_k = xpx.lazy_apply(special.binom, r+k, k) + return (-1)**(r-k)*binom_r_k*binom_rpk_k + + +@xp_capabilities(skip_backends=[('dask.array', "too many issues")], + cpu_only=True, # torch doesn't have `binom` + exceptions=('cupy', 'jax.numpy'), + extra_note="Lazy backends do not support `order` greater than 4.") +@_axis_nan_policy_factory( # noqa: E302 + _moment_result_object, n_samples=1, result_to_tuple=_moment_tuple, + n_outputs=lambda kwds: _moment_outputs(kwds, [1, 2, 3, 4]) +) +def lmoment(sample, order=None, *, axis=0, sorted=False, standardize=True): + r"""Compute L-moments of a sample from a continuous distribution. + + The L-moments of a probability distribution are summary statistics with + uses similar to those of conventional moments, but they are defined in + terms of the expected values of order statistics. + Sample L-moments are defined analogously to population L-moments, and + they can serve as estimators of population L-moments. They tend to be less + sensitive to extreme observations than conventional moments. + + Parameters + ---------- + sample : array_like + The real-valued sample whose L-moments are desired. + order : array_like, optional + The (positive integer) orders of the desired L-moments. + Must be a scalar or non-empty 1D array. Default is [1, 2, 3, 4]. + axis : int or None, default=0 + If an int, the axis of the input along which to compute the statistic. + The statistic of each axis-slice (e.g. row) of the input will appear + in a corresponding element of the output. If None, the input will be + raveled before computing the statistic. + sorted : bool, default=False + Whether `sample` is already sorted in increasing order along `axis`. + If False (default), `sample` will be sorted. + standardize : bool, default=True + Whether to return L-moment ratios for orders 3 and higher. + L-moment ratios are analogous to standardized conventional + moments: they are the non-standardized L-moments divided + by the L-moment of order 2. + + Returns + ------- + lmoments : ndarray + The sample L-moments of order `order`. + + See Also + -------- + moment + + Notes + ----- + SciPy offers only basic capabilities for working with L-moments. For more advanced + features, consider the ``lmo`` package [4]_. + + References + ---------- + .. [1] D. Bilkova. "L-Moments and TL-Moments as an Alternative Tool of + Statistical Data Analysis". Journal of Applied Mathematics and + Physics. 2014. :doi:`10.4236/jamp.2014.210104` + .. [2] J. R. M. Hosking. "L-Moments: Analysis and Estimation of Distributions + Using Linear Combinations of Order Statistics". Journal of the Royal + Statistical Society. 1990. :doi:`10.1111/j.2517-6161.1990.tb01775.x` + .. [3] "L-moment". *Wikipedia*. https://en.wikipedia.org/wiki/L-moment. + .. [4] @jorenham, *Lmo*, https://github.com/jorenham/Lmo/ + + Examples + -------- + >>> import numpy as np + >>> from scipy import stats + >>> rng = np.random.default_rng(328458568356392) + >>> sample = rng.exponential(size=100000) + >>> stats.lmoment(sample) + array([1.00124272, 0.50111437, 0.3340092 , 0.16755338]) + + Note that the first four standardized population L-moments of the standard + exponential distribution are 1, 1/2, 1/3, and 1/6; the sample L-moments + provide reasonable estimates. + + """ + xp = array_namespace(sample) + args = _lmoment_iv(sample, order, axis, sorted, standardize, xp=xp) + sample, order, axis, sorted, standardize = args + + n_moments = 4 if is_lazy_array(order) else int(xp.max(order)) + k = xp.arange(n_moments, dtype=sample.dtype) + prk = _prk(xpx.expand_dims(k, axis=tuple(range(1, sample.ndim+1))), k) + bk = _br(sample, r=k, xp=xp) + + n = sample.shape[-1] + if n < bk.shape[-1]: + bk = xpx.at(bk)[..., n:].set(0) # remove NaNs due to n_moments > n + + lmoms = xp.vecdot(prk, bk, axis=-1) + if standardize and n_moments > 2: + lmoms = xpx.at(lmoms)[2:, ...].divide(lmoms[1, ...]) + + if n < lmoms.shape[0]: + lmoms = xpx.at(lmoms)[n:, ...].set(xp.nan) # add NaNs where appropriate + # return lmoms[order-1] # strict can't handle fancy indexing plus ellipses + return xp.take(lmoms, order - 1, axis=0) if order.ndim == 1 else lmoms[order - 1] + + +LinregressResult = _make_tuple_bunch('LinregressResult', + ['slope', 'intercept', 'rvalue', + 'pvalue', 'stderr'], + extra_field_names=['intercept_stderr']) + + +def _pack_LinregressResult(slope, intercept, rvalue, pvalue, stderr, intercept_stderr): + return LinregressResult(slope, intercept, rvalue, pvalue, stderr, + intercept_stderr=intercept_stderr) + + +def _unpack_LinregressResult(res, _): + return tuple(res) + (res.intercept_stderr,) + + +@xp_capabilities(skip_backends=[('dask.array', '`dtype` inference failed....')], + cpu_only=True, exceptions=['cupy', 'jax.numpy'], + marray=True) +@_axis_nan_policy_factory(_pack_LinregressResult, n_samples=2, + result_to_tuple=_unpack_LinregressResult, paired=True, + too_small=1, n_outputs=6) +def linregress(x, y, alternative='two-sided', *, axis=0): + """ + Calculate a linear least-squares regression for two sets of measurements. + + Parameters + ---------- + x, y : array_like + Two sets of measurements. Both arrays should have the same length N. + alternative : {'two-sided', 'less', 'greater'}, optional + Defines the alternative hypothesis. Default is 'two-sided'. + The following options are available: + + * 'two-sided': the slope of the regression line is nonzero + * 'less': the slope of the regression line is less than zero + * 'greater': the slope of the regression line is greater than zero + + .. versionadded:: 1.7.0 + axis : int or None, default: 0 + If an int, the axis of the input along which to compute the statistic. + The statistic of each axis-slice (e.g. row) of the input will appear in a + corresponding element of the output. + If ``None``, the input will be raveled before computing the statistic. + + Returns + ------- + result : ``LinregressResult`` instance + The return value is an object with the following attributes: + + slope : float + Slope of the regression line. + intercept : float + Intercept of the regression line. + rvalue : float + The Pearson correlation coefficient. The square of ``rvalue`` + is equal to the coefficient of determination. + pvalue : float + The p-value for a hypothesis test whose null hypothesis is + that the slope is zero, using Wald Test with t-distribution of + the test statistic. See `alternative` above for alternative + hypotheses. + stderr : float + Standard error of the estimated slope (gradient), under the + assumption of residual normality. + intercept_stderr : float + Standard error of the estimated intercept, under the assumption + of residual normality. + + See Also + -------- + scipy.optimize.curve_fit : + Use non-linear least squares to fit a function to data. + scipy.optimize.leastsq : + Minimize the sum of squares of a set of equations. + + Notes + ----- + For compatibility with older versions of SciPy, the return value acts + like a ``namedtuple`` of length 5, with fields ``slope``, ``intercept``, + ``rvalue``, ``pvalue`` and ``stderr``, so one can continue to write:: + + slope, intercept, r, p, se = linregress(x, y) + + With that style, however, the standard error of the intercept is not + available. To have access to all the computed values, including the + standard error of the intercept, use the return value as an object + with attributes, e.g.:: + + result = linregress(x, y) + print(result.intercept, result.intercept_stderr) + + Examples + -------- + >>> import numpy as np + >>> import matplotlib.pyplot as plt + >>> from scipy import stats + >>> rng = np.random.default_rng() + + Generate some data: + + >>> x = rng.random(10) + >>> y = 1.6*x + rng.random(10) + + Perform the linear regression: + + >>> res = stats.linregress(x, y) + + Coefficient of determination (R-squared): + + >>> print(f"R-squared: {res.rvalue**2:.6f}") + R-squared: 0.717533 + + Plot the data along with the fitted line: + + >>> plt.plot(x, y, 'o', label='original data') + >>> plt.plot(x, res.intercept + res.slope*x, 'r', label='fitted line') + >>> plt.legend() + >>> plt.show() + + Calculate 95% confidence interval on slope and intercept: + + >>> # Two-sided inverse Students t-distribution + >>> # p - probability, df - degrees of freedom + >>> from scipy.stats import t + >>> tinv = lambda p, df: abs(t.ppf(p/2, df)) + + >>> ts = tinv(0.05, len(x)-2) + >>> print(f"slope (95%): {res.slope:.6f} +/- {ts*res.stderr:.6f}") + slope (95%): 1.453392 +/- 0.743465 + >>> print(f"intercept (95%): {res.intercept:.6f}" + ... f" +/- {ts*res.intercept_stderr:.6f}") + intercept (95%): 0.616950 +/- 0.544475 + + """ + xp = array_namespace(x, y) + x, y = xp_promote(x, y, force_floating=True, xp=xp) + + TINY = 1.0e-20 + + # _axis_nan_policy decorator ensures that `axis=-1` + x, y = _share_masks(x, y, xp=xp) + n = _count_nonmasked(x, axis=-1, keepdims=False, xp=xp) + xmean = xp.mean(x, axis=-1, keepdims=True) + ymean = xp.mean(y, axis=-1, keepdims=True) + + # Average sums of square differences from the mean + # ssxm = mean( (x-mean(x))^2 ) + # ssxym = mean( (x-mean(x)) * (y-mean(y)) ) + x_ = _demean(x, xmean, axis=-1, xp=xp) + y_ = _demean(y, ymean, axis=-1, xp=xp, precision_warning=False) + xmean = xp.squeeze(xmean, axis=-1) + ymean = xp.squeeze(ymean, axis=-1) + + ssxm = xp.vecdot(x_, x_, axis=-1) / n + ssym = xp.vecdot(y_, y_, axis=-1) / n + ssxym = xp.vecdot(x_, y_, axis=-1) / n + + # R-value + # r = ssxym / sqrt( ssxm * ssym ) + degenerate = (ssxm == 0.0) | (ssym == 0.0) + NaN = xp.asarray(xp.nan, dtype=ssxym.dtype) + r = xpx.apply_where( + ~degenerate, + (ssxym, ssxm, ssym), + lambda ssxym, ssxm, ssym: xp.clip(ssxym / xp.sqrt(ssxm * ssym), -1.0, 1.0), + lambda ssxym, ssxm, ssym: xp.where(ssxym==0, NaN, 0.0) + ) + + slope = ssxym / ssxm + intercept = ymean - slope*xmean + with np.errstate(invalid='ignore', divide='ignore'): + df = n - 2 # Number of degrees of freedom + # n-2 degrees of freedom because 2 has been used up + # to estimate the mean and standard deviation + t = r * xp.sqrt(df / ((1.0 - r + TINY)*(1.0 + r + TINY))) + + dist = _SimpleStudentT(xp.asarray(df, dtype=t.dtype)) + prob = _get_pvalue(t, dist, alternative, xp=xp) + prob = prob[()] if prob.ndim == 0 else prob + + slope_stderr = xp.sqrt((1 - r**2) * ssym / ssxm / df) + + # Also calculate the standard error of the intercept + # The following relationship is used: + # ssxm = mean( (x-mean(x))^2 ) + # = ssx - sx*sx + # = mean( x^2 ) - mean(x)^2 + intercept_stderr = slope_stderr * xp.sqrt(ssxm + xmean**2) + + outputs = slope, intercept, r, prob, slope_stderr, intercept_stderr + outputs = (output[()] if output.ndim == 0 else output for output in outputs) + slope, intercept, r, prob, slope_stderr, intercept_stderr = outputs + + return LinregressResult(slope=slope, intercept=intercept, rvalue=r, + pvalue=prob, stderr=slope_stderr, + intercept_stderr=intercept_stderr) + +def _linearized_pmean(a, p, *, axis=None, weights=None, xp=None): + # pmean linearized as a function of p about p = 0; see gh-23407 + M0 = gmean(a, axis=axis, weights=weights) + + loga = xp.log(a) + + ln_avg = _xp_mean(loga, axis=axis, weights=weights) + ln2_avg = _xp_mean(loga * loga, axis=axis, weights=weights) + + return M0 * (1 + 0.5 * p * (ln2_avg - ln_avg * ln_avg)) + + +@xp_capabilities() +def _xp_mean(x, /, *, axis=None, weights=None, keepdims=False, nan_policy='propagate', + dtype=None, warn=True, xp=None): + r"""Compute the arithmetic mean along the specified axis. + + Parameters + ---------- + x : real array + Array containing real numbers whose mean is desired. + axis : int or tuple of ints, default: None + If an int or tuple of ints, the axis or axes of the input along which + to compute the statistic. The statistic of each axis-slice (e.g. row) + of the input will appear in a corresponding element of the output. + If ``None``, the input will be raveled before computing the statistic. + weights : real array, optional + If specified, an array of weights associated with the values in `x`; + otherwise ``1``. If `weights` and `x` do not have the same shape, the + arrays will be broadcasted before performing the calculation. See + Notes for details. + keepdims : bool, optional + If this is set to ``True``, the axes which are reduced are left + in the result as dimensions with length one. With this option, + the result will broadcast correctly against the input array. + nan_policy : {'propagate', 'omit', 'raise'}, default: 'propagate' + Defines how to handle input NaNs. + + - ``propagate``: if a NaN is present in the axis slice (e.g. row) along + which the statistic is computed, the corresponding entry of the output + will be NaN. + - ``omit``: NaNs will be omitted when performing the calculation. + If insufficient data remains in the axis slice along which the + statistic is computed, the corresponding entry of the output will be + NaN. + - ``raise``: if a NaN is present, a ``ValueError`` will be raised. + + dtype : dtype, optional + Type to use in computing the mean. For integer inputs, the default is + the default float type of the array library; for floating point inputs, + the dtype is that of the input. + + Returns + ------- + out : array + The mean of each slice + + Notes + ----- + Let :math:`x_i` represent element :math:`i` of data `x` and let :math:`w_i` + represent the corresponding element of `weights` after broadcasting. Then the + (weighted) mean :math:`\bar{x}_w` is given by: + + .. math:: + + \bar{x}_w = \frac{ \sum_{i=0}^{n-1} w_i x_i } + { \sum_{i=0}^{n-1} w_i } + + where :math:`n` is the number of elements along a slice. Note that this simplifies + to the familiar :math:`(\sum_i x_i) / n` when the weights are all ``1`` (default). + + The behavior of this function with respect to weights is somewhat different + from that of `np.average`. For instance, + `np.average` raises an error when `axis` is not specified and the shapes of `x` + and the `weights` array are not the same; `xp_mean` simply broadcasts the two. + Also, `np.average` raises an error when weights sum to zero along a slice; + `xp_mean` computes the appropriate result. The intent is for this function's + interface to be consistent with the rest of `scipy.stats`. + + Note that according to the formula, including NaNs with zero weights is not + the same as *omitting* NaNs with ``nan_policy='omit'``; in the former case, + the NaNs will continue to propagate through the calculation whereas in the + latter case, the NaNs are excluded entirely. + + """ + # ensure that `x` and `weights` are array-API compatible arrays of identical shape + xp = array_namespace(x) if xp is None else xp + x = _asarray(x, dtype=dtype, subok=True) + weights = xp.asarray(weights, dtype=dtype) if weights is not None else weights + + # to ensure that this matches the behavior of decorated functions when one of the + # arguments has size zero, it's easiest to call a similar decorated function. + if is_numpy(xp) and (xp_size(x) == 0 + or (weights is not None and xp_size(weights) == 0)): + return gmean(x, weights=weights, axis=axis, keepdims=keepdims) + + x, weights = xp_promote(x, weights, broadcast=True, force_floating=True, xp=xp) + if weights is not None: + x, weights = _share_masks(x, weights, xp=xp) + + # handle the special case of zero-sized arrays + message = (too_small_1d_not_omit if (x.ndim == 1 or axis is None) + else too_small_nd_not_omit) + if xp_size(x) == 0: + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + res = xp.mean(x, axis=axis, keepdims=keepdims) + if warn and xp_size(res) != 0: + warnings.warn(message, SmallSampleWarning, stacklevel=2) + return res + + contains_nan = _contains_nan(x, nan_policy, xp_omit_okay=True, xp=xp) + if weights is not None: + contains_nan_w = _contains_nan(weights, nan_policy, xp_omit_okay=True, xp=xp) + contains_nan = contains_nan | contains_nan_w + + # Handle `nan_policy='omit'` by giving zero weight to NaNs, whether they + # appear in `x` or `weights`. Emit warning if there is an all-NaN slice. + # Test nan_policy before the implicit call to bool(contains_nan) + # to avoid raising on lazy xps on the default nan_policy='propagate' + lazy = is_lazy_array(x) + if nan_policy == 'omit' and (lazy or contains_nan): + nan_mask = xp.isnan(x) + if weights is not None: + nan_mask |= xp.isnan(weights) + if warn and not lazy and xp.any(xp.all(nan_mask, axis=axis)): + message = (too_small_1d_omit if (x.ndim == 1 or axis is None) + else too_small_nd_omit) + warnings.warn(message, SmallSampleWarning, stacklevel=2) + weights = xp.ones_like(x) if weights is None else weights + x = xp.where(nan_mask, 0., x) + weights = xp.where(nan_mask, 0., weights) + + # Perform the mean calculation itself + if weights is None: + return xp.mean(x, axis=axis, keepdims=keepdims) + + # consider using `vecdot` if `axis` tuple support is added (data-apis/array-api#910) + norm = xp.sum(weights, axis=axis) + wsum = xp.sum(x * weights, axis=axis) + with np.errstate(divide='ignore', invalid='ignore'): + res = wsum/norm + + # Respect `keepdims` and convert NumPy 0-D arrays to scalars + if keepdims: + + if axis is None: + final_shape = (1,) * len(x.shape) + else: + # axis can be a scalar or sequence + axes = (axis,) if not isinstance(axis, Sequence) else axis + final_shape = list(x.shape) + for i in axes: + final_shape[i] = 1 + + res = xp.reshape(res, tuple(final_shape)) + + return res[()] if res.ndim == 0 else res + + +@xp_capabilities() +def _xp_var(x, /, *, axis=None, correction=0, keepdims=False, nan_policy='propagate', + dtype=None, xp=None): + # an array-api compatible function for variance with scipy.stats interface + # and features (e.g. `nan_policy`). + xp = array_namespace(x) if xp is None else xp + x = _asarray(x, subok=True) + + # use `_xp_mean` instead of `xp.var` for desired warning behavior + # it would be nice to combine this with `_var`, which uses `_moment` + # and therefore warns when precision is lost, but that does not support + # `axis` tuples or keepdims. Eventually, `_axis_nan_policy` will simplify + # `axis` tuples and implement `keepdims` for non-NumPy arrays; then it will + # be easy. + kwargs = dict(axis=axis, nan_policy=nan_policy, dtype=dtype, xp=xp) + mean = _xp_mean(x, keepdims=True, **kwargs) + x = _asarray(x, dtype=mean.dtype, subok=True) + x_mean = _demean(x, mean, axis, xp=xp) + x_mean_conj = (xp.conj(x_mean) if xp.isdtype(x_mean.dtype, 'complex floating') + else x_mean) # crossref data-apis/array-api#824 + var = _xp_mean(x_mean * x_mean_conj, keepdims=keepdims, **kwargs) + + if correction != 0: + n = _count_nonmasked(x, axis, xp=xp, keepdims=keepdims) + # Or two lines with ternaries : ) + # axis = range(x.ndim) if axis is None else axis + # n = math.prod(x.shape[i] for i in axis) if iterable(axis) else x.shape[axis] + + n = xp.asarray(n, dtype=var.dtype, device=xp_device(x)) + + if nan_policy == 'omit': + nan_mask = xp.astype(xp.isnan(x), var.dtype) + n = n - xp.sum(nan_mask, axis=axis, keepdims=keepdims) + + # Produce NaNs silently when n - correction <= 0 + nc = n - correction + factor = xpx.apply_where(nc > 0, (n, nc), operator.truediv, fill_value=xp.nan) + var *= factor + + return var[()] if var.ndim == 0 else var + + +class _SimpleNormal: + # A very simple, array-API compatible normal distribution for use in + # hypothesis tests. May be replaced by new infrastructure Normal + # distribution in due time. + + def cdf(self, x): + return special.ndtr(x) + + def sf(self, x): + return special.ndtr(-x) + + def isf(self, x): + return -special.ndtri(x) + + +class _SimpleChi2: + # A very simple, array-API compatible chi-squared distribution for use in + # hypothesis tests. May be replaced by new infrastructure chi-squared + # distribution in due time. + def __init__(self, df): + self.df = df + + def cdf(self, x): + return special.chdtr(self.df, x) + + def sf(self, x): + return special.chdtrc(self.df, x) + + +class _SimpleBeta: + # A very simple, array-API compatible beta distribution for use in + # hypothesis tests. May be replaced by new infrastructure beta + # distribution in due time. + def __init__(self, a, b, *, loc=None, scale=None): + self.a = a + self.b = b + self.loc = loc + self.scale = scale + + def cdf(self, x): + if self.loc is not None or self.scale is not None: + loc = 0 if self.loc is None else self.loc + scale = 1 if self.scale is None else self.scale + return special.betainc(self.a, self.b, (x - loc)/scale) + return special.betainc(self.a, self.b, x) + + def sf(self, x): + if self.loc is not None or self.scale is not None: + loc = 0 if self.loc is None else self.loc + scale = 1 if self.scale is None else self.scale + return special.betaincc(self.a, self.b, (x - loc)/scale) + return special.betaincc(self.a, self.b, x) + + +class _SimpleStudentT: + # A very simple, array-API compatible t distribution for use in + # hypothesis tests. May be replaced by new infrastructure t + # distribution in due time. + def __init__(self, df): + self.df = df + + def cdf(self, t): + return special.stdtr(self.df, t) + + def sf(self, t): + return special.stdtr(self.df, -t) + + +class _SimpleF: + # A very simple, array-API compatible F distribution for use in + # hypothesis tests. + def __init__(self, dfn, dfd): + self.dfn = dfn + self.dfd = dfd + + def cdf(self, x): + return special.fdtr(self.dfn, self.dfd, x) + + def sf(self, x): + return special.fdtrc(self.dfn, self.dfd, x) diff --git a/v12x/pylib/scipy/stats/_survival.py b/v12x/pylib/scipy/stats/_survival.py new file mode 100644 index 0000000000000000000000000000000000000000..70a84806a95ae82b2968bb9cf5ee37f48808c59c --- /dev/null +++ b/v12x/pylib/scipy/stats/_survival.py @@ -0,0 +1,686 @@ +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Literal +import warnings + +import numpy as np +from scipy import special, interpolate, stats +from scipy._lib._array_api import xp_capabilities +from scipy.stats._censored_data import CensoredData +from scipy.stats._common import ConfidenceInterval + +if TYPE_CHECKING: + import numpy.typing as npt + + +__all__ = ['ecdf', 'logrank'] + + +@dataclass +class EmpiricalDistributionFunction: + """An empirical distribution function produced by `scipy.stats.ecdf` + + Attributes + ---------- + quantiles : ndarray + The unique values of the sample from which the + `EmpiricalDistributionFunction` was estimated. + probabilities : ndarray + The point estimates of the cumulative distribution function (CDF) or + its complement, the survival function (SF), corresponding with + `quantiles`. + """ + quantiles: np.ndarray + probabilities: np.ndarray + # Exclude these from __str__ + _n: np.ndarray = field(repr=False) # number "at risk" + _d: np.ndarray = field(repr=False) # number of "deaths" + _sf: np.ndarray = field(repr=False) # survival function for var estimate + _kind: str = field(repr=False) # type of function: "cdf" or "sf" + + def __init__(self, q, p, n, d, kind): + self.probabilities = p + self.quantiles = q + self._n = n + self._d = d + self._sf = p if kind == 'sf' else 1 - p + self._kind = kind + + f0 = 1 if kind == 'sf' else 0 # leftmost function value + f1 = 1 - f0 + # fill_value can't handle edge cases at infinity + x = np.insert(q, [0, len(q)], [-np.inf, np.inf]) + y = np.insert(p, [0, len(p)], [f0, f1]) + # `or` conditions handle the case of empty x, points + self._f = interpolate.interp1d(x, y, kind='previous', + assume_sorted=True) + + def evaluate(self, x): + """Evaluate the empirical CDF/SF function at the input. + + Parameters + ---------- + x : ndarray + Argument to the CDF/SF + + Returns + ------- + y : ndarray + The CDF/SF evaluated at the input + """ + return self._f(x) + + def plot(self, ax=None, **matplotlib_kwargs): + """Plot the empirical distribution function + + Available only if ``matplotlib`` is installed. + + Parameters + ---------- + ax : matplotlib.axes.Axes + Axes object to draw the plot onto, otherwise uses the current Axes. + + **matplotlib_kwargs : dict, optional + Keyword arguments passed directly to `matplotlib.axes.Axes.step`. + Unless overridden, ``where='post'``. + + Returns + ------- + lines : list of `matplotlib.lines.Line2D` + Objects representing the plotted data + """ + try: + import matplotlib # noqa: F401 + except ModuleNotFoundError as exc: + message = "matplotlib must be installed to use method `plot`." + raise ModuleNotFoundError(message) from exc + + if ax is None: + import matplotlib.pyplot as plt + ax = plt.gca() + + kwargs = {'where': 'post'} + kwargs.update(matplotlib_kwargs) + + delta = np.ptp(self.quantiles)*0.05 # how far past sample edge to plot + q = self.quantiles + q = [q[0] - delta] + list(q) + [q[-1] + delta] + + return ax.step(q, self.evaluate(q), **kwargs) + + def confidence_interval(self, confidence_level=0.95, *, method='linear'): + """Compute a confidence interval around the CDF/SF point estimate + + Parameters + ---------- + confidence_level : float, default: 0.95 + Confidence level for the computed confidence interval + + method : str, {"linear", "log-log"} + Method used to compute the confidence interval. Options are + "linear" for the conventional Greenwood confidence interval + (default) and "log-log" for the "exponential Greenwood", + log-negative-log-transformed confidence interval. + + Returns + ------- + ci : ``ConfidenceInterval`` + An object with attributes ``low`` and ``high``, instances of + `~scipy.stats._result_classes.EmpiricalDistributionFunction` that + represent the lower and upper bounds (respectively) of the + confidence interval. + + Notes + ----- + Confidence intervals are computed according to the Greenwood formula + (``method='linear'``) or the more recent "exponential Greenwood" + formula (``method='log-log'``) as described in [1]_. The conventional + Greenwood formula can result in lower confidence limits less than 0 + and upper confidence limits greater than 1; these are clipped to the + unit interval. NaNs may be produced by either method; these are + features of the formulas. + + References + ---------- + .. [1] Sawyer, Stanley. "The Greenwood and Exponential Greenwood + Confidence Intervals in Survival Analysis." + https://www.math.wustl.edu/~sawyer/handouts/greenwood.pdf + + """ + message = ("Confidence interval bounds do not implement a " + "`confidence_interval` method.") + if self._n is None: + raise NotImplementedError(message) + + methods = {'linear': self._linear_ci, + 'log-log': self._loglog_ci} + + message = f"`method` must be one of {set(methods)}." + if method.lower() not in methods: + raise ValueError(message) + + message = "`confidence_level` must be a scalar between 0 and 1." + confidence_level = np.asarray(confidence_level)[()] + if confidence_level.shape or not (0 <= confidence_level <= 1): + raise ValueError(message) + + method_fun = methods[method.lower()] + low, high = method_fun(confidence_level) + + message = ("The confidence interval is undefined at some observations." + " This is a feature of the mathematical formula used, not" + " an error in its implementation.") + if np.any(np.isnan(low) | np.isnan(high)): + warnings.warn(message, RuntimeWarning, stacklevel=2) + + low, high = np.clip(low, 0, 1), np.clip(high, 0, 1) + low = EmpiricalDistributionFunction(self.quantiles, low, None, None, + self._kind) + high = EmpiricalDistributionFunction(self.quantiles, high, None, None, + self._kind) + return ConfidenceInterval(low, high) + + def _linear_ci(self, confidence_level): + sf, d, n = self._sf, self._d, self._n + # When n == d, Greenwood's formula divides by zero. + # When s != 0, this can be ignored: var == inf, and CI is [0, 1] + # When s == 0, this results in NaNs. Produce an informative warning. + with np.errstate(divide='ignore', invalid='ignore'): + var = sf ** 2 * np.cumsum(d / (n * (n - d))) + + se = np.sqrt(var) + z = special.ndtri(1 / 2 + confidence_level / 2) + + z_se = z * se + low = self.probabilities - z_se + high = self.probabilities + z_se + + return low, high + + def _loglog_ci(self, confidence_level): + sf, d, n = self._sf, self._d, self._n + + with np.errstate(divide='ignore', invalid='ignore'): + var = 1 / np.log(sf) ** 2 * np.cumsum(d / (n * (n - d))) + + se = np.sqrt(var) + z = special.ndtri(1 / 2 + confidence_level / 2) + + with np.errstate(divide='ignore'): + lnl_points = np.log(-np.log(sf)) + + z_se = z * se + low = np.exp(-np.exp(lnl_points + z_se)) + high = np.exp(-np.exp(lnl_points - z_se)) + if self._kind == "cdf": + low, high = 1-high, 1-low + + return low, high + + +@dataclass +class ECDFResult: + """ Result object returned by `scipy.stats.ecdf` + + Attributes + ---------- + cdf : `~scipy.stats._result_classes.EmpiricalDistributionFunction` + An object representing the empirical cumulative distribution function. + sf : `~scipy.stats._result_classes.EmpiricalDistributionFunction` + An object representing the complement of the empirical cumulative + distribution function. + """ + cdf: EmpiricalDistributionFunction + sf: EmpiricalDistributionFunction + + def __init__(self, q, cdf, sf, n, d): + self.cdf = EmpiricalDistributionFunction(q, cdf, n, d, "cdf") + self.sf = EmpiricalDistributionFunction(q, sf, n, d, "sf") + + +def _iv_CensoredData( + sample: "npt.ArrayLike | CensoredData", param_name: str = "sample" +) -> CensoredData: + """Attempt to convert `sample` to `CensoredData`.""" + if not isinstance(sample, CensoredData): + try: # takes care of input standardization/validation + sample = CensoredData(uncensored=sample) + except ValueError as e: + message = str(e).replace('uncensored', param_name) + raise type(e)(message) from e + return sample + + +@xp_capabilities(np_only=True) +def ecdf(sample: "npt.ArrayLike | CensoredData") -> ECDFResult: + """Empirical cumulative distribution function of a sample. + + The empirical cumulative distribution function (ECDF) is a step function + estimate of the CDF of the distribution underlying a sample. This function + returns objects representing both the empirical distribution function and + its complement, the empirical survival function. + + Parameters + ---------- + sample : 1D array_like or `scipy.stats.CensoredData` + Besides array_like, instances of `scipy.stats.CensoredData` containing + uncensored and right-censored observations are supported. Currently, + other instances of `scipy.stats.CensoredData` will result in a + ``NotImplementedError``. + + Returns + ------- + res : `~scipy.stats._result_classes.ECDFResult` + An object with the following attributes. + + cdf : `~scipy.stats._result_classes.EmpiricalDistributionFunction` + An object representing the empirical cumulative distribution + function. + sf : `~scipy.stats._result_classes.EmpiricalDistributionFunction` + An object representing the empirical survival function. + + The `cdf` and `sf` attributes themselves have the following attributes. + + quantiles : ndarray + The unique values in the sample that defines the empirical CDF/SF. + probabilities : ndarray + The point estimates of the probabilities corresponding with + `quantiles`. + + And the following methods: + + evaluate(x) : + Evaluate the CDF/SF at the argument. + + plot(ax) : + Plot the CDF/SF on the provided axes. + + confidence_interval(confidence_level=0.95) : + Compute the confidence interval around the CDF/SF at the values in + `quantiles`. + + Notes + ----- + When each observation of the sample is a precise measurement, the ECDF + steps up by ``1/len(sample)`` at each of the observations [1]_. + + When observations are lower bounds, upper bounds, or both upper and lower + bounds, the data is said to be "censored", and `sample` may be provided as + an instance of `scipy.stats.CensoredData`. + + For right-censored data, the ECDF is given by the Kaplan-Meier estimator + [2]_; other forms of censoring are not supported at this time. + + Confidence intervals are computed according to the Greenwood formula or the + more recent "Exponential Greenwood" formula as described in [4]_. + + References + ---------- + .. [1] Conover, William Jay. Practical nonparametric statistics. Vol. 350. + John Wiley & Sons, 1999. + + .. [2] Kaplan, Edward L., and Paul Meier. "Nonparametric estimation from + incomplete observations." Journal of the American statistical + association 53.282 (1958): 457-481. + + .. [3] Goel, Manish Kumar, Pardeep Khanna, and Jugal Kishore. + "Understanding survival analysis: Kaplan-Meier estimate." + International journal of Ayurveda research 1.4 (2010): 274. + + .. [4] Sawyer, Stanley. "The Greenwood and Exponential Greenwood Confidence + Intervals in Survival Analysis." + https://www.math.wustl.edu/~sawyer/handouts/greenwood.pdf + + Examples + -------- + **Uncensored Data** + + As in the example from [1]_ page 79, five boys were selected at random from + those in a single high school. Their one-mile run times were recorded as + follows. + + >>> sample = [6.23, 5.58, 7.06, 6.42, 5.20] # one-mile run times (minutes) + + The empirical distribution function, which approximates the distribution + function of one-mile run times of the population from which the boys were + sampled, is calculated as follows. + + >>> from scipy import stats + >>> res = stats.ecdf(sample) + >>> res.cdf.quantiles + array([5.2 , 5.58, 6.23, 6.42, 7.06]) + >>> res.cdf.probabilities + array([0.2, 0.4, 0.6, 0.8, 1. ]) + + To plot the result as a step function: + + >>> import matplotlib.pyplot as plt + >>> ax = plt.subplot() + >>> res.cdf.plot(ax) + >>> ax.set_xlabel('One-Mile Run Time (minutes)') + >>> ax.set_ylabel('Empirical CDF') + >>> plt.show() + + **Right-censored Data** + + As in the example from [1]_ page 91, the lives of ten car fanbelts were + tested. Five tests concluded because the fanbelt being tested broke, but + the remaining tests concluded for other reasons (e.g. the study ran out of + funding, but the fanbelt was still functional). The mileage driven + with the fanbelts were recorded as follows. + + >>> broken = [77, 47, 81, 56, 80] # in thousands of miles driven + >>> unbroken = [62, 60, 43, 71, 37] + + Precise survival times of the fanbelts that were still functional at the + end of the tests are unknown, but they are known to exceed the values + recorded in ``unbroken``. Therefore, these observations are said to be + "right-censored", and the data is represented using + `scipy.stats.CensoredData`. + + >>> sample = stats.CensoredData(uncensored=broken, right=unbroken) + + The empirical survival function is calculated as follows. + + >>> res = stats.ecdf(sample) + >>> res.sf.quantiles + array([37., 43., 47., 56., 60., 62., 71., 77., 80., 81.]) + >>> res.sf.probabilities + array([1. , 1. , 0.875, 0.75 , 0.75 , 0.75 , 0.75 , 0.5 , 0.25 , 0. ]) + + To plot the result as a step function: + + >>> ax = plt.subplot() + >>> res.sf.plot(ax) + >>> ax.set_xlabel('Fanbelt Survival Time (thousands of miles)') + >>> ax.set_ylabel('Empirical SF') + >>> plt.show() + + """ + sample = _iv_CensoredData(sample) + + if sample.num_censored() == 0: + res = _ecdf_uncensored(sample._uncensor()) + elif sample.num_censored() == sample._right.size: + res = _ecdf_right_censored(sample) + else: + # Support additional censoring options in follow-up PRs + message = ("Currently, only uncensored and right-censored data is " + "supported.") + raise NotImplementedError(message) + + t, cdf, sf, n, d = res + return ECDFResult(t, cdf, sf, n, d) + + +def _ecdf_uncensored(sample): + sample = np.sort(sample) + x, counts = np.unique(sample, return_counts=True) + + # [1].81 "the fraction of [observations] that are less than or equal to x + events = np.cumsum(counts) + n = sample.size + cdf = events / n + + # [1].89 "the relative frequency of the sample that exceeds x in value" + sf = 1 - cdf + + at_risk = np.concatenate(([n], n - events[:-1])) + return x, cdf, sf, at_risk, counts + + +def _ecdf_right_censored(sample): + # It is conventional to discuss right-censored data in terms of + # "survival time", "death", and "loss" (e.g. [2]). We'll use that + # terminology here. + # This implementation was influenced by the references cited and also + # https://www.youtube.com/watch?v=lxoWsVco_iM + # https://en.wikipedia.org/wiki/Kaplan%E2%80%93Meier_estimator + # In retrospect it is probably most easily compared against [3]. + # Ultimately, the data needs to be sorted, so this implementation is + # written to avoid a separate call to `unique` after sorting. In hope of + # better performance on large datasets, it also computes survival + # probabilities at unique times only rather than at each observation. + tod = sample._uncensored # time of "death" + tol = sample._right # time of "loss" + times = np.concatenate((tod, tol)) + died = np.asarray([1]*tod.size + [0]*tol.size) + + # sort by times + i = np.argsort(times) + times = times[i] + died = died[i] + at_risk = np.arange(times.size, 0, -1) + + # logical indices of unique times + j = np.diff(times, prepend=-np.inf, append=np.inf) > 0 + j_l = j[:-1] # first instances of unique times + j_r = j[1:] # last instances of unique times + + # get number at risk and deaths at each unique time + t = times[j_l] # unique times + n = at_risk[j_l] # number at risk at each unique time + cd = np.cumsum(died)[j_r] # cumulative deaths up to/including unique times + d = np.diff(cd, prepend=0) # deaths at each unique time + + # compute survival function + sf = np.cumprod((n - d) / n) + cdf = 1 - sf + return t, cdf, sf, n, d + + +@dataclass +class LogRankResult: + """Result object returned by `scipy.stats.logrank`. + + Attributes + ---------- + statistic : float ndarray + The computed statistic (defined below). Its magnitude is the + square root of the magnitude returned by most other logrank test + implementations. + pvalue : float ndarray + The computed p-value of the test. + """ + statistic: np.ndarray + pvalue: np.ndarray + + +@xp_capabilities(np_only=True) +def logrank( + x: "npt.ArrayLike | CensoredData", + y: "npt.ArrayLike | CensoredData", + alternative: Literal['two-sided', 'less', 'greater'] = "two-sided" +) -> LogRankResult: + r"""Compare the survival distributions of two samples via the logrank test. + + Parameters + ---------- + x, y : array_like or CensoredData + Samples to compare based on their empirical survival functions. + alternative : {'two-sided', 'less', 'greater'}, optional + Defines the alternative hypothesis. + + The null hypothesis is that the survival distributions of the two + groups, say *X* and *Y*, are identical. + + The following alternative hypotheses [4]_ are available (default is + 'two-sided'): + + * 'two-sided': the survival distributions of the two groups are not + identical. + * 'less': survival of group *X* is favored: the group *X* failure rate + function is less than the group *Y* failure rate function at some + times. + * 'greater': survival of group *Y* is favored: the group *X* failure + rate function is greater than the group *Y* failure rate function at + some times. + + Returns + ------- + res : `~scipy.stats._result_classes.LogRankResult` + An object containing attributes: + + statistic : float ndarray + The computed statistic (defined below). Its magnitude is the + square root of the magnitude returned by most other logrank test + implementations. + pvalue : float ndarray + The computed p-value of the test. + + See Also + -------- + scipy.stats.ecdf + + Notes + ----- + The logrank test [1]_ compares the observed number of events to + the expected number of events under the null hypothesis that the two + samples were drawn from the same distribution. The statistic is + + .. math:: + + Z_i = \frac{\sum_{j=1}^J(O_{i,j}-E_{i,j})}{\sqrt{\sum_{j=1}^J V_{i,j}}} + \rightarrow \mathcal{N}(0,1) + + where + + .. math:: + + E_{i,j} = O_j \frac{N_{i,j}}{N_j}, + \qquad + V_{i,j} = E_{i,j} \left(\frac{N_j-O_j}{N_j}\right) + \left(\frac{N_j-N_{i,j}}{N_j-1}\right), + + :math:`i` denotes the group (i.e. it may assume values :math:`x` or + :math:`y`, or it may be omitted to refer to the combined sample) + :math:`j` denotes the time (at which an event occurred), + :math:`N` is the number of subjects at risk just before an event occurred, + and :math:`O` is the observed number of events at that time. + + The ``statistic`` :math:`Z_x` returned by `logrank` is the (signed) square + root of the statistic returned by many other implementations. Under the + null hypothesis, :math:`Z_x**2` is asymptotically distributed according to + the chi-squared distribution with one degree of freedom. Consequently, + :math:`Z_x` is asymptotically distributed according to the standard normal + distribution. The advantage of using :math:`Z_x` is that the sign + information (i.e. whether the observed number of events tends to be less + than or greater than the number expected under the null hypothesis) is + preserved, allowing `scipy.stats.logrank` to offer one-sided alternative + hypotheses. + + References + ---------- + .. [1] Mantel N. "Evaluation of survival data and two new rank order + statistics arising in its consideration." + Cancer Chemotherapy Reports, 50(3):163-170, PMID: 5910392, 1966 + .. [2] Bland, Altman, "The logrank test", BMJ, 328:1073, + :doi:`10.1136/bmj.328.7447.1073`, 2004 + .. [3] "Logrank test", Wikipedia, + https://en.wikipedia.org/wiki/Logrank_test + .. [4] Brown, Mark. "On the choice of variance for the log rank test." + Biometrika 71.1 (1984): 65-74. + .. [5] Klein, John P., and Melvin L. Moeschberger. Survival analysis: + techniques for censored and truncated data. Vol. 1230. New York: + Springer, 2003. + + Examples + -------- + Reference [2]_ compared the survival times of patients with two different + types of recurrent malignant gliomas. The samples below record the time + (number of weeks) for which each patient participated in the study. The + `scipy.stats.CensoredData` class is used because the data is + right-censored: the uncensored observations correspond with observed deaths + whereas the censored observations correspond with the patient leaving the + study for another reason. + + >>> from scipy import stats + >>> x = stats.CensoredData( + ... uncensored=[6, 13, 21, 30, 37, 38, 49, 50, + ... 63, 79, 86, 98, 202, 219], + ... right=[31, 47, 80, 82, 82, 149] + ... ) + >>> y = stats.CensoredData( + ... uncensored=[10, 10, 12, 13, 14, 15, 16, 17, 18, 20, 24, 24, + ... 25, 28,30, 33, 35, 37, 40, 40, 46, 48, 76, 81, + ... 82, 91, 112, 181], + ... right=[34, 40, 70] + ... ) + + We can calculate and visualize the empirical survival functions + of both groups as follows. + + >>> import numpy as np + >>> import matplotlib.pyplot as plt + >>> ax = plt.subplot() + >>> ecdf_x = stats.ecdf(x) + >>> ecdf_x.sf.plot(ax, label='Astrocytoma') + >>> ecdf_y = stats.ecdf(y) + >>> ecdf_y.sf.plot(ax, label='Glioblastoma') + >>> ax.set_xlabel('Time to death (weeks)') + >>> ax.set_ylabel('Empirical SF') + >>> plt.legend() + >>> plt.show() + + Visual inspection of the empirical survival functions suggests that the + survival times tend to be different between the two groups. To formally + assess whether the difference is significant at the 1% level, we use the + logrank test. + + >>> res = stats.logrank(x=x, y=y) + >>> res.statistic + -2.73799 + >>> res.pvalue + 0.00618 + + The p-value is less than 1%, so we can consider the data to be evidence + against the null hypothesis in favor of the alternative that there is a + difference between the two survival functions. + + """ + # Input validation. `alternative` IV handled in `_get_pvalue` below. + x = _iv_CensoredData(sample=x, param_name='x') + y = _iv_CensoredData(sample=y, param_name='y') + + # Combined sample. (Under H0, the two groups are identical.) + xy = CensoredData( + uncensored=np.concatenate((x._uncensored, y._uncensored)), + right=np.concatenate((x._right, y._right)) + ) + + # Extract data from the combined sample + res = ecdf(xy) + idx = res.sf._d.astype(bool) # indices of observed events + times_xy = res.sf.quantiles[idx] # unique times of observed events + at_risk_xy = res.sf._n[idx] # combined number of subjects at risk + deaths_xy = res.sf._d[idx] # combined number of events + + # Get the number at risk within each sample. + # First compute the number at risk in group X at each of the `times_xy`. + # Could use `interpolate_1d`, but this is more compact. + res_x = ecdf(x) + i = np.searchsorted(res_x.sf.quantiles, times_xy) + at_risk_x = np.append(res_x.sf._n, 0)[i] # 0 at risk after last time + # Subtract from the combined number at risk to get number at risk in Y + at_risk_y = at_risk_xy - at_risk_x + + # Compute the variance. + num = at_risk_x * at_risk_y * deaths_xy * (at_risk_xy - deaths_xy) + den = at_risk_xy**2 * (at_risk_xy - 1) + # Note: when `at_risk_xy == 1`, we would have `at_risk_xy - 1 == 0` in the + # numerator and denominator. Simplifying the fraction symbolically, we + # would always find the overall quotient to be zero, so don't compute it. + i = at_risk_xy > 1 + sum_var = np.sum(num[i]/den[i]) + + # Get the observed and expected number of deaths in group X + n_died_x = x._uncensored.size + sum_exp_deaths_x = np.sum(at_risk_x * (deaths_xy/at_risk_xy)) + + # Compute the statistic. This is the square root of that in references. + statistic = (n_died_x - sum_exp_deaths_x)/np.sqrt(sum_var) + + # Equivalent to chi2(df=1).sf(statistic**2) when alternative='two-sided' + norm = stats._stats_py._SimpleNormal() + pvalue = stats._stats_py._get_pvalue(statistic, norm, alternative, xp=np) + + return LogRankResult(statistic=statistic[()], pvalue=pvalue[()]) diff --git a/v12x/pylib/scipy/stats/_tukeylambda_stats.py b/v12x/pylib/scipy/stats/_tukeylambda_stats.py new file mode 100644 index 0000000000000000000000000000000000000000..d3cda54fa2801f7659419b3c9c68d15ed80fc13e --- /dev/null +++ b/v12x/pylib/scipy/stats/_tukeylambda_stats.py @@ -0,0 +1,199 @@ +import numpy as np +from numpy import poly1d +from scipy.special import beta + + +# The following code was used to generate the Pade coefficients for the +# Tukey Lambda variance function. Version 0.17 of mpmath was used. +#--------------------------------------------------------------------------- +# import mpmath as mp +# +# mp.mp.dps = 60 +# +# one = mp.mpf(1) +# two = mp.mpf(2) +# +# def mpvar(lam): +# if lam == 0: +# v = mp.pi**2 / three +# else: +# v = (two / lam**2) * (one / (one + two*lam) - +# mp.beta(lam + one, lam + one)) +# return v +# +# t = mp.taylor(mpvar, 0, 8) +# p, q = mp.pade(t, 4, 4) +# print("p =", [mp.fp.mpf(c) for c in p]) +# print("q =", [mp.fp.mpf(c) for c in q]) +#--------------------------------------------------------------------------- + +# Pade coefficients for the Tukey Lambda variance function. +_tukeylambda_var_pc = [3.289868133696453, 0.7306125098871127, + -0.5370742306855439, 0.17292046290190008, + -0.02371146284628187] +_tukeylambda_var_qc = [1.0, 3.683605511659861, 4.184152498888124, + 1.7660926747377275, 0.2643989311168465] + +# numpy.poly1d instances for the numerator and denominator of the +# Pade approximation to the Tukey Lambda variance. +_tukeylambda_var_p = poly1d(_tukeylambda_var_pc[::-1]) +_tukeylambda_var_q = poly1d(_tukeylambda_var_qc[::-1]) + + +def tukeylambda_variance(lam): + """Variance of the Tukey Lambda distribution. + + Parameters + ---------- + lam : array_like + The lambda values at which to compute the variance. + + Returns + ------- + v : ndarray + The variance. For lam < -0.5, the variance is not defined, so + np.nan is returned. For lam = 0.5, np.inf is returned. + + Notes + ----- + In an interval around lambda=0, this function uses the [4,4] Pade + approximation to compute the variance. Otherwise it uses the standard + formula (https://en.wikipedia.org/wiki/Tukey_lambda_distribution). The + Pade approximation is used because the standard formula has a removable + discontinuity at lambda = 0, and does not produce accurate numerical + results near lambda = 0. + """ + lam = np.asarray(lam) + shp = lam.shape + lam = np.atleast_1d(lam).astype(np.float64) + + # For absolute values of lam less than threshold, use the Pade + # approximation. + threshold = 0.075 + + # Play games with masks to implement the conditional evaluation of + # the distribution. + # lambda < -0.5: var = nan + low_mask = lam < -0.5 + # lambda == -0.5: var = inf + neghalf_mask = lam == -0.5 + # abs(lambda) < threshold: use Pade approximation + small_mask = np.abs(lam) < threshold + # else the "regular" case: use the explicit formula. + reg_mask = ~(low_mask | neghalf_mask | small_mask) + + # Get the 'lam' values for the cases where they are needed. + small = lam[small_mask] + reg = lam[reg_mask] + + # Compute the function for each case. + v = np.empty_like(lam) + v[low_mask] = np.nan + v[neghalf_mask] = np.inf + if small.size > 0: + # Use the Pade approximation near lambda = 0. + v[small_mask] = _tukeylambda_var_p(small) / _tukeylambda_var_q(small) + if reg.size > 0: + v[reg_mask] = (2.0 / reg**2) * (1.0 / (1.0 + 2 * reg) - + beta(reg + 1, reg + 1)) + v = v.reshape(shp) + return v + + +# The following code was used to generate the Pade coefficients for the +# Tukey Lambda kurtosis function. Version 0.17 of mpmath was used. +#--------------------------------------------------------------------------- +# import mpmath as mp +# +# mp.mp.dps = 60 +# +# one = mp.mpf(1) +# two = mp.mpf(2) +# three = mp.mpf(3) +# four = mp.mpf(4) +# +# def mpkurt(lam): +# if lam == 0: +# k = mp.mpf(6)/5 +# else: +# numer = (one/(four*lam+one) - four*mp.beta(three*lam+one, lam+one) + +# three*mp.beta(two*lam+one, two*lam+one)) +# denom = two*(one/(two*lam+one) - mp.beta(lam+one,lam+one))**2 +# k = numer / denom - three +# return k +# +# # There is a bug in mpmath 0.17: when we use the 'method' keyword of the +# # taylor function and we request a degree 9 Taylor polynomial, we actually +# # get degree 8. +# t = mp.taylor(mpkurt, 0, 9, method='quad', radius=0.01) +# t = [mp.chop(c, tol=1e-15) for c in t] +# p, q = mp.pade(t, 4, 4) +# print("p =", [mp.fp.mpf(c) for c in p]) +# print("q =", [mp.fp.mpf(c) for c in q]) +#--------------------------------------------------------------------------- + +# Pade coefficients for the Tukey Lambda kurtosis function. +_tukeylambda_kurt_pc = [1.2, -5.853465139719495, -22.653447381131077, + 0.20601184383406815, 4.59796302262789] +_tukeylambda_kurt_qc = [1.0, 7.171149192233599, 12.96663094361842, + 0.43075235247853005, -2.789746758009912] + +# numpy.poly1d instances for the numerator and denominator of the +# Pade approximation to the Tukey Lambda kurtosis. +_tukeylambda_kurt_p = poly1d(_tukeylambda_kurt_pc[::-1]) +_tukeylambda_kurt_q = poly1d(_tukeylambda_kurt_qc[::-1]) + + +def tukeylambda_kurtosis(lam): + """Kurtosis of the Tukey Lambda distribution. + + Parameters + ---------- + lam : array_like + The lambda values at which to compute the variance. + + Returns + ------- + v : ndarray + The variance. For lam < -0.25, the variance is not defined, so + np.nan is returned. For lam = 0.25, np.inf is returned. + + """ + lam = np.asarray(lam) + shp = lam.shape + lam = np.atleast_1d(lam).astype(np.float64) + + # For absolute values of lam less than threshold, use the Pade + # approximation. + threshold = 0.055 + + # Use masks to implement the conditional evaluation of the kurtosis. + # lambda < -0.25: kurtosis = nan + low_mask = lam < -0.25 + # lambda == -0.25: kurtosis = inf + negqrtr_mask = lam == -0.25 + # lambda near 0: use Pade approximation + small_mask = np.abs(lam) < threshold + # else the "regular" case: use the explicit formula. + reg_mask = ~(low_mask | negqrtr_mask | small_mask) + + # Get the 'lam' values for the cases where they are needed. + small = lam[small_mask] + reg = lam[reg_mask] + + # Compute the function for each case. + k = np.empty_like(lam) + k[low_mask] = np.nan + k[negqrtr_mask] = np.inf + if small.size > 0: + k[small_mask] = _tukeylambda_kurt_p(small) / _tukeylambda_kurt_q(small) + if reg.size > 0: + numer = (1.0 / (4 * reg + 1) - 4 * beta(3 * reg + 1, reg + 1) + + 3 * beta(2 * reg + 1, 2 * reg + 1)) + denom = 2 * (1.0/(2 * reg + 1) - beta(reg + 1, reg + 1))**2 + k[reg_mask] = numer / denom - 3 + + # The return value will be a numpy array; resetting the shape ensures that + # if `lam` was a scalar, the return value is a 0-d array. + k = k.reshape(shp) + return k diff --git a/v12x/pylib/scipy/stats/_unuran/__init__.py b/v12x/pylib/scipy/stats/_unuran/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/v12x/pylib/scipy/stats/_variation.py b/v12x/pylib/scipy/stats/_variation.py new file mode 100644 index 0000000000000000000000000000000000000000..c991a70158a226a1b0fb6a1a0fce6f1ec4a67400 --- /dev/null +++ b/v12x/pylib/scipy/stats/_variation.py @@ -0,0 +1,133 @@ +import warnings +import numpy as np + +from scipy._lib._array_api import ( + array_namespace, + xp_capabilities, + xp_device, + _count_nonmasked, + xp_promote, +) +import scipy._external.array_api_extra as xpx + +from ._axis_nan_policy import _axis_nan_policy_factory + + +@xp_capabilities(marray=True) +@_axis_nan_policy_factory( + lambda x: x, n_outputs=1, result_to_tuple=lambda x, _: (x,) +) +def variation(a, axis=0, nan_policy='propagate', ddof=0, *, keepdims=False): + """ + Compute the coefficient of variation. + + The coefficient of variation is the standard deviation divided by the + mean. This function is equivalent to:: + + np.std(x, axis=axis, ddof=ddof) / np.mean(x) + + The default for ``ddof`` is 0, but many definitions of the coefficient + of variation use the square root of the unbiased sample variance + for the sample standard deviation, which corresponds to ``ddof=1``. + + The function does not take the absolute value of the mean of the data, + so the return value is negative if the mean is negative. + + Parameters + ---------- + a : array_like + Input array. + axis : int or None, optional + Axis along which to calculate the coefficient of variation. + Default is 0. If None, compute over the whole array `a`. + nan_policy : {'propagate', 'raise', 'omit'}, optional + Defines how to handle when input contains ``nan``. + The following options are available: + + * 'propagate': return ``nan`` + * 'raise': raise an exception + * 'omit': perform the calculation with ``nan`` values omitted + + The default is 'propagate'. + ddof : int, optional + Gives the "Delta Degrees Of Freedom" used when computing the + standard deviation. The divisor used in the calculation of the + standard deviation is ``N - ddof``, where ``N`` is the number of + elements. `ddof` must be less than ``N``; if it isn't, the result + will be ``nan`` or ``inf``, depending on ``N`` and the values in + the array. By default `ddof` is zero for backwards compatibility, + but it is recommended to use ``ddof=1`` to ensure that the sample + standard deviation is computed as the square root of the unbiased + sample variance. + + Returns + ------- + variation : ndarray + The calculated variation along the requested axis. + + Notes + ----- + There are several edge cases that are handled without generating a + warning: + + * If both the mean and the standard deviation are zero, ``nan`` + is returned. + * If the mean is zero and the standard deviation is nonzero, ``inf`` + is returned. + * If the input has length zero (either because the array has zero + length, or all the input values are ``nan`` and ``nan_policy`` is + ``'omit'``), ``nan`` is returned. + * If the input contains ``inf``, ``nan`` is returned. + + References + ---------- + .. [1] Zwillinger, D. and Kokoska, S. (2000). CRC Standard + Probability and Statistics Tables and Formulae. Chapman & Hall: New + York. 2000. + + Examples + -------- + >>> import numpy as np + >>> from scipy.stats import variation + >>> variation([1, 2, 3, 4, 5], ddof=1) + 0.5270462766947299 + + Compute the variation along a given dimension of an array that contains + a few ``nan`` values: + + >>> x = np.array([[ 10.0, np.nan, 11.0, 19.0, 23.0, 29.0, 98.0], + ... [ 29.0, 30.0, 32.0, 33.0, 35.0, 56.0, 57.0], + ... [np.nan, np.nan, 12.0, 13.0, 16.0, 16.0, 17.0]]) + >>> variation(x, axis=1, ddof=1, nan_policy='omit') + array([1.05109361, 0.31428986, 0.146483 ]) + + """ + xp = array_namespace(a) + a = xp_promote(a, force_floating=True, xp=xp) + + # `nan_policy` and `keepdims` are handled by `_axis_nan_policy` + if axis is None: + a = xp.reshape(a, (-1,)) + axis = 0 + + n = xp.asarray(_count_nonmasked(a, axis=axis), dtype=a.dtype, device=xp_device(a)) + + with (np.errstate(divide='ignore', invalid='ignore'), warnings.catch_warnings()): + warnings.simplefilter("ignore") + mean_a = xp.mean(a, axis=axis) + std_a = xp.std(a, axis=axis) + correction = (n / (n - ddof))**0.5 # we may need uncorrected std below + result = std_a * correction / mean_a + + def special_case(std_a, mean_a): + # xref data-apis/array-api-extra#196 + mxp = array_namespace(std_a, mean_a) + # `_xp_inf` is a workaround for torch.copysign not accepting a scalar yet, + # xref data-apis/array-api-compat#271 + _xp_inf = mxp.asarray(mxp.inf, dtype=mean_a.dtype, device=xp_device(mean_a)) + return mxp.where(std_a > 0, mxp.copysign(_xp_inf, mean_a), mxp.nan) + + result = xpx.apply_where((ddof == n), (std_a, mean_a), + special_case, fill_value=result) + + return result[()] if result.ndim == 0 else result diff --git a/v12x/pylib/scipy/stats/_warnings_errors.py b/v12x/pylib/scipy/stats/_warnings_errors.py new file mode 100644 index 0000000000000000000000000000000000000000..57bc8fcabd7a0f382960046cd253bbc74d4b76e4 --- /dev/null +++ b/v12x/pylib/scipy/stats/_warnings_errors.py @@ -0,0 +1,62 @@ +# Warnings + + +class DegenerateDataWarning(RuntimeWarning): + """Warns when data is degenerate and results may not be reliable. + + Parameters + ---------- + msg : str, optional + Warning message. + """ + def __init__(self, msg=None): + if msg is None: + msg = ("Degenerate data encountered; results may not be reliable.") + self.args = (msg,) + + +class ConstantInputWarning(DegenerateDataWarning): + """Warns when all values in data are exactly equal. + + Parameters + ---------- + msg : str, optional + Warning message. + """ + def __init__(self, msg=None): + if msg is None: + msg = ("All values in data are exactly equal; " + "results may not be reliable.") + self.args = (msg,) + + +class NearConstantInputWarning(DegenerateDataWarning): + """Warns when all values in data are nearly equal. + + Parameters + ---------- + msg : str, optional + Warning message. + """ + def __init__(self, msg=None): + if msg is None: + msg = ("All values in data are nearly equal; " + "results may not be reliable.") + self.args = (msg,) + + +# Errors + + +class FitError(RuntimeError): + """Represents an error condition when fitting a distribution to data. + + Parameters + ---------- + msg : str, optional + Error message. + """ + def __init__(self, msg=None): + if msg is None: + msg = ("An error occurred when fitting a distribution to data.") + self.args = (msg,) diff --git a/v12x/pylib/scipy/stats/_wilcoxon.py b/v12x/pylib/scipy/stats/_wilcoxon.py new file mode 100644 index 0000000000000000000000000000000000000000..a4437ba01a80deb0a09d36b4f3fca54ae5da960e --- /dev/null +++ b/v12x/pylib/scipy/stats/_wilcoxon.py @@ -0,0 +1,284 @@ +import numpy as np + +from scipy import stats +from ._stats_py import _get_pvalue, _rankdata, _SimpleNormal +from . import _morestats +from ._axis_nan_policy import _broadcast_arrays +from ._hypotests import _get_wilcoxon_distr +from scipy._lib._util import _get_nan +from scipy._lib._array_api import (array_namespace, xp_promote, xp_size, is_jax, + is_marray, _count_nonmasked) +import scipy._external.array_api_extra as xpx + + +class WilcoxonDistribution: + + def __init__(self, n): + n = np.asarray(n).astype(int, copy=False) + self.n = n + self._dists = {ni: _get_wilcoxon_distr(ni) for ni in np.unique(n)} + + def _cdf1(self, k, n): + pmfs = self._dists[n] + return pmfs[:k + 1].sum() + + def _cdf(self, k, n): + return np.vectorize(self._cdf1, otypes=[float])(k, n) + + def _sf1(self, k, n): + pmfs = self._dists[n] + return pmfs[k:].sum() + + def _sf(self, k, n): + return np.vectorize(self._sf1, otypes=[float])(k, n) + + def mean(self): + return self.n * (self.n + 1) / 4 + + def _prep(self, k): + k = np.asarray(k).astype(int, copy=False) + mn = self.mean() + out = np.empty(k.shape, dtype=np.float64) + return k, mn, out + + def cdf(self, k): + k, mn, out = self._prep(k) + return xpx.apply_where( + k <= mn, (k, self.n), + self._cdf, + lambda k, n: 1 - self._sf(k+1, n))[()] + + def sf(self, k): + k, mn, out = self._prep(k) + return xpx.apply_where( + k <= mn, (k, self.n), + self._sf, + lambda k, n: 1 - self._cdf(k-1, n))[()] + + +def _wilcoxon_iv(x, y, zero_method, correction, alternative, method, axis): + xp = array_namespace(x, y) + x, y = xp_promote(x, y, force_floating=True, xp=xp) + + axis = np.asarray(axis)[()] # OK to use NumPy for input validation + message = "`axis` must be an integer." + if not np.issubdtype(axis.dtype, np.integer) or axis.ndim != 0: + raise ValueError(message) + axis = int(axis) + + message = '`axis` must be compatible with the shape(s) of `x` (and `y`)' + AxisError = getattr(np, 'AxisError', None) or np.exceptions.AxisError + try: + if y is None: + d = x + else: + x, y = _broadcast_arrays((x, y), axis=axis, xp=xp) + d = x - y + d = xp.moveaxis(d, axis, -1) + except AxisError as e: + raise AxisError(message) from e + + message = "`x` and `y` must have the same length along `axis`." + if y is not None and x.shape[axis] != y.shape[axis]: + raise ValueError(message) + + message = "`x` (and `y`, if provided) must be an array of real numbers." + if not xp.isdtype(d.dtype, "real floating"): + raise ValueError(message) + + zero_method = str(zero_method).lower() + zero_methods = {"wilcox", "pratt", "zsplit"} + message = f"`zero_method` must be one of {zero_methods}." + if zero_method not in zero_methods: + raise ValueError(message) + + corrections = {True, False} + message = f"`correction` must be one of {corrections}." + if correction not in corrections: + raise ValueError(message) + + alternative = str(alternative).lower() + alternatives = {"two-sided", "less", "greater"} + message = f"`alternative` must be one of {alternatives}." + if alternative not in alternatives: + raise ValueError(message) + + if not isinstance(method, stats.PermutationMethod): + methods = {"auto", "asymptotic", "exact"} + message = (f"`method` must be one of {methods} or " + "an instance of `stats.PermutationMethod`.") + if method not in methods: + raise ValueError(message) + output_z = True if method == 'asymptotic' else False + + if is_jax(xp) and str(method) in {"auto", "exact"}: + message = ("When using `wilcoxon` with `jax.numpy` arrays, `method` must be " + "either 'asymptotic' or an instance of `stats.PermutationMethod`.") + raise ValueError(message) + + if is_marray(xp) and (method != "asymptotic" or zero_method != 'zsplit'): + message = ("Only `method='asymptotic'`/`zero_method='zsplit'` is compatible " + "with MArrays.") + raise ValueError(message) + + # For small samples, we decide later whether to perform an exact test or a + # permutation test. The reason is that the presence of ties is not + # known at the input validation stage. + n_zero = xp.count_nonzero(d == 0, axis=None) + if method == "auto" and d.shape[-1] > 50: + method = "asymptotic" + + return d, zero_method, correction, alternative, method, axis, output_z, n_zero, xp + + +def _wilcoxon_statistic(d, method, zero_method='wilcox', *, xp): + dtype = d.dtype + i_zeros = (d == 0) + + if zero_method == 'wilcox': + # Wilcoxon's method for treating zeros was to remove them from + # the calculation. We do this by replacing 0s with NaNs, which + # are ignored anyway. + # Copy required for array-api-strict. See data-apis/array-api-extra#506. + d = xpx.at(d)[i_zeros].set(xp.nan, copy=True) + + i_nan = xp.isnan(d) + n_nan = xp.astype(xp.count_nonzero(i_nan, axis=-1), dtype) + count = _count_nonmasked(d, axis=-1) - n_nan + + r, _, t = _rankdata(xp.abs(d), 'average', return_ties=True, xp=xp) + r, t = xp.astype(r, dtype, copy=False), xp.astype(t, dtype, copy=False) + + r_plus = xp.sum(xp.astype(d > 0, dtype) * r, axis=-1) + r_minus = xp.sum(xp.astype(d < 0, dtype) * r, axis=-1) + + has_ties = xp.any(t == 0) + + if zero_method == "zsplit": + # The "zero-split" method for treating zeros is to add half their contribution + # to r_plus and half to r_minus. + # See gh-2263 for the origin of this method. + r_zero_2 = xp.sum(xp.astype(i_zeros, dtype) * r, axis=-1) / 2 + r_plus = xpx.at(r_plus)[...].add(r_zero_2) + r_minus = xpx.at(r_minus)[...].add(r_zero_2) + + mn = count * (count + 1.) * 0.25 + se = count * (count + 1.) * (2. * count + 1.) + + if zero_method == "pratt": + # Pratt's method for treating zeros was just to modify the z-statistic. + + # normal approximation needs to be adjusted, see Cureton (1967) + n_zero = xp.astype(xp.count_nonzero(i_zeros, axis=-1), dtype) + mn = xpx.at(mn)[...].subtract(n_zero * (n_zero + 1.) * 0.25) + se = xpx.at(se)[...].subtract(n_zero * (n_zero + 1.) * (2. * n_zero + 1.)) + + # zeros are not to be included in tie-correction. + # any tie counts corresponding with zeros are in the 0th column + # t[xp.any(i_zeros, axis=-1), 0] = 0 + t_i_zeros = xp.zeros_like(i_zeros) + t_i_zeros = xpx.at(t_i_zeros)[..., 0].set(xp.any(i_zeros, axis=-1)) + t = xpx.at(t)[t_i_zeros].set(0.) + + tie_correct = xp.sum(t**3 - t, axis=-1) + se = xp.sqrt((se - tie_correct/2) / 24) + + # se = 0 means that no non-zero values are left in d. we only need z + # if method is asymptotic. however, if method="auto", the switch to + # asymptotic might only happen after the statistic is calculated, so z + # needs to be computed. in all other cases, avoid division by zero warning + # (z is not needed anyways) + if method in ["asymptotic", "auto"]: + z = (r_plus - mn) / se + else: + z = xp.nan + + return r_plus, r_minus, se, z, count, has_ties + + +def _correction_sign(z, alternative, xp): + if alternative == 'greater': + return 1 + elif alternative == 'less': + return -1 + else: + return xp.sign(z) + + +def _wilcoxon_nd(x, y=None, zero_method='wilcox', correction=True, + alternative='two-sided', method='auto', axis=0): + + temp = _wilcoxon_iv(x, y, zero_method, correction, alternative, method, axis) + d, zero_method, correction, alternative, method, axis, output_z, n_zero, xp = temp + + if xp_size(d) == 0: + NaN = _get_nan(d, xp=xp) + res = _morestats.WilcoxonResult(statistic=NaN, pvalue=NaN) + if method == 'asymptotic': + res.zstatistic = NaN + return res + + r_plus, r_minus, se, z, count, has_ties = _wilcoxon_statistic( + d, method, zero_method, xp=xp + ) + + # we only know if there are ties after computing the statistic and not + # at the input validation stage. if the original method was auto and + # the decision was to use an exact test, we override this to + # a permutation test now (since method='exact' is not exact in the + # presence of ties) + if method == "auto": + if not (has_ties or n_zero > 0): + method = "exact" + elif d.shape[-1] <= 13: + # the possible outcomes to be simulated by the permutation test + # are 2**n, where n is the sample size. + # if n <= 13, the p-value is deterministic since 2**13 is less + # than 9999, the default number of n_resamples + method = stats.PermutationMethod() + else: + # if there are ties and the sample size is too large to + # run a deterministic permutation test, fall back to asymptotic + method = "asymptotic" + + if method == 'asymptotic': + if correction: + sign = _correction_sign(z, alternative, xp=xp) + z = xpx.at(z)[...].subtract(sign * 0.5 / se) + p = _get_pvalue(z, _SimpleNormal(), alternative, xp=xp) + elif method == 'exact': + dist = WilcoxonDistribution(count) + # The null distribution in `dist` is exact only if there are no ties + # or zeros. If there are ties or zeros, the statistic can be non- + # integral, but the null distribution is only defined for integral + # values of the statistic. Therefore, we're conservative: round + # non-integral statistic up before computing CDF and down before + # computing SF. This preserves symmetry w.r.t. alternatives and + # order of the input arguments. See gh-19872. + r_plus_np = np.asarray(r_plus) + if alternative == 'less': + p = dist.cdf(np.ceil(r_plus_np)) + elif alternative == 'greater': + p = dist.sf(np.floor(r_plus_np)) + else: + p = 2 * np.minimum(dist.sf(np.floor(r_plus_np)), + dist.cdf(np.ceil(r_plus_np))) + p = np.clip(p, 0, 1) + p = xp.asarray(p, dtype=d.dtype) + else: # `PermutationMethod` instance (already validated) + p = stats.permutation_test( + # permutation_test always uses `axis=-1` as `_wilcoxon_statistic` assumes + (d,), lambda d, axis: _wilcoxon_statistic(d, method, zero_method, xp=xp)[0], + permutation_type='samples', **method._asdict(), + alternative=alternative, axis=-1).pvalue + + # for backward compatibility... + statistic = xp.minimum(r_plus, r_minus) if alternative=='two-sided' else r_plus + z = -xp.abs(z) if (alternative == 'two-sided' and method == 'asymptotic') else z + + statistic = statistic[()] if statistic.ndim == 0 else statistic + p = p[()] if p.ndim == 0 else p + res = _morestats.WilcoxonResult(statistic=statistic, pvalue=p) + if output_z: + res.zstatistic = z[()] if z.ndim == 0 else z + return res diff --git a/v12x/pylib/scipy/stats/biasedurn.py b/v12x/pylib/scipy/stats/biasedurn.py new file mode 100644 index 0000000000000000000000000000000000000000..2b1c9f1cf2b3b39acdcaeda97a90e8c11c589d89 --- /dev/null +++ b/v12x/pylib/scipy/stats/biasedurn.py @@ -0,0 +1,16 @@ +# This file is not meant for public use and will be removed in SciPy v2.0.0. + +from scipy._lib.deprecation import _sub_module_deprecation + + +__all__: list[str] = [] + + +def __dir__(): + return __all__ + + +def __getattr__(name): + return _sub_module_deprecation(sub_package="stats", module="biasedurn", + private_modules=["_biasedurn"], all=__all__, + attribute=name) diff --git a/v12x/pylib/scipy/stats/contingency.py b/v12x/pylib/scipy/stats/contingency.py new file mode 100644 index 0000000000000000000000000000000000000000..c45112f3f1d7ae346681732d75943c75e6ef3fad --- /dev/null +++ b/v12x/pylib/scipy/stats/contingency.py @@ -0,0 +1,526 @@ +""" +Contingency table functions (:mod:`scipy.stats.contingency`) +============================================================ + +Functions for creating and analyzing contingency tables. + +.. currentmodule:: scipy.stats.contingency + +.. autosummary:: + :toctree: generated/ + + chi2_contingency + relative_risk + odds_ratio + crosstab + association + + expected_freq + margins + +""" + + +from functools import reduce +import math +import numpy as np +from ._stats_py import power_divergence, _untabulate +from ._relative_risk import relative_risk +from ._crosstab import crosstab +from ._odds_ratio import odds_ratio +from scipy._lib._array_api import xp_capabilities +from scipy._lib._bunch import _make_tuple_bunch +from scipy import stats + + +__all__ = ['margins', 'expected_freq', 'chi2_contingency', 'crosstab', + 'association', 'relative_risk', 'odds_ratio'] + + +@xp_capabilities(np_only=True) +def margins(a): + """Return a list of the marginal sums of the array `a`. + + Parameters + ---------- + a : ndarray + The array for which to compute the marginal sums. + + Returns + ------- + margsums : list of ndarrays + A list of length `a.ndim`. `margsums[k]` is the result + of summing `a` over all axes except `k`; it has the same + number of dimensions as `a`, but the length of each axis + except axis `k` will be 1. + + Examples + -------- + >>> import numpy as np + >>> from scipy.stats.contingency import margins + + >>> a = np.arange(12).reshape(2, 6) + >>> a + array([[ 0, 1, 2, 3, 4, 5], + [ 6, 7, 8, 9, 10, 11]]) + >>> m0, m1 = margins(a) + >>> m0 + array([[15], + [51]]) + >>> m1 + array([[ 6, 8, 10, 12, 14, 16]]) + + >>> b = np.arange(24).reshape(2,3,4) + >>> m0, m1, m2 = margins(b) + >>> m0 + array([[[ 66]], + [[210]]]) + >>> m1 + array([[[ 60], + [ 92], + [124]]]) + >>> m2 + array([[[60, 66, 72, 78]]]) + """ + margsums = [] + ranged = list(range(a.ndim)) + for k in ranged: + marg = np.apply_over_axes(np.sum, a, [j for j in ranged if j != k]) + margsums.append(marg) + return margsums + + +@xp_capabilities(np_only=True) +def expected_freq(observed): + """ + Compute the expected frequencies from a contingency table. + + Given an n-dimensional contingency table of observed frequencies, + compute the expected frequencies for the table based on the marginal + sums under the assumption that the groups associated with each + dimension are independent. + + Parameters + ---------- + observed : array_like + The table of observed frequencies. (While this function can handle + a 1-D array, that case is trivial. Generally `observed` is at + least 2-D.) + + Returns + ------- + expected : ndarray of float64 + The expected frequencies, based on the marginal sums of the table. + Same shape as `observed`. + + Examples + -------- + >>> import numpy as np + >>> from scipy.stats.contingency import expected_freq + >>> observed = np.array([[10, 10, 20],[20, 20, 20]]) + >>> expected_freq(observed) + array([[ 12., 12., 16.], + [ 18., 18., 24.]]) + + """ + # Typically `observed` is an integer array. If `observed` has a large + # number of dimensions or holds large values, some of the following + # computations may overflow, so we first switch to floating point. + observed = np.asarray(observed, dtype=np.float64) + + # Create a list of the marginal sums. + margsums = margins(observed) + + # Create the array of expected frequencies. The shapes of the + # marginal sums returned by apply_over_axes() are just what we + # need for broadcasting in the following product. + d = observed.ndim + expected = reduce(np.multiply, margsums) / observed.sum() ** (d - 1) + return expected + + +Chi2ContingencyResult = _make_tuple_bunch( + 'Chi2ContingencyResult', + ['statistic', 'pvalue', 'dof', 'expected_freq'], [] +) + + +@xp_capabilities(np_only=True) +def chi2_contingency(observed, correction=True, lambda_=None, *, method=None): + """Chi-square test of independence of variables in a contingency table. + + This function computes the chi-square statistic and p-value for the + hypothesis test of independence of the observed frequencies in the + contingency table [1]_ `observed`. The expected frequencies are computed + based on the marginal sums under the assumption of independence; see + `scipy.stats.contingency.expected_freq`. The number of degrees of + freedom is (expressed using numpy functions and attributes):: + + dof = observed.size - sum(observed.shape) + observed.ndim - 1 + + + Parameters + ---------- + observed : array_like + The contingency table. The table contains the observed frequencies + (i.e. number of occurrences) in each category. In the two-dimensional + case, the table is often described as an "R x C table". + correction : bool, optional + If True, *and* the degrees of freedom is 1, apply Yates' correction + for continuity. The effect of the correction is to adjust each + observed value by 0.5 towards the corresponding expected value. + lambda_ : float or str, optional + By default, the statistic computed in this test is Pearson's + chi-squared statistic [2]_. `lambda_` allows a statistic from the + Cressie-Read power divergence family [3]_ to be used instead. See + `scipy.stats.power_divergence` for details. + method : ResamplingMethod, optional + Defines the method used to compute the p-value. Compatible only with + ``correction=False``, default `lambda_`, and two-way tables. + If `method` is an instance of `scipy.stats.PermutationMethod` or + `scipy.stats.MonteCarloMethod`, the p-value is computed using + `scipy.stats.permutation_test`/`scipy.stats.monte_carlo_test` with the + provided configuration options and other appropriate settings. + Otherwise, the p-value is computed as documented in the notes. + Note that if `method` is an instance of `scipy.stats.MonteCarloMethod`, the + ``rvs`` attribute must be left unspecified; Monte Carlo samples are always drawn + using the ``rvs`` method of `scipy.stats.random_table`. + + .. versionadded:: 1.15.0 + + + Returns + ------- + res : Chi2ContingencyResult + An object containing attributes: + + statistic : float + The test statistic. + pvalue : float + The p-value of the test. + dof : int + The degrees of freedom. NaN if `method` is not ``None``. + expected_freq : ndarray, same shape as `observed` + The expected frequencies, based on the marginal sums of the table. + + See Also + -------- + scipy.stats.contingency.expected_freq + scipy.stats.fisher_exact + scipy.stats.chisquare + scipy.stats.power_divergence + scipy.stats.barnard_exact + scipy.stats.boschloo_exact + :ref:`hypothesis_chi2_contingency` : Extended example + + Notes + ----- + An often quoted guideline for the validity of this calculation is that + the test should be used only if the observed and expected frequencies + in each cell are at least 5. + + This is a test for the independence of different categories of a + population. The test is only meaningful when the dimension of + `observed` is two or more. Applying the test to a one-dimensional + table will always result in `expected` equal to `observed` and a + chi-square statistic equal to 0. + + This function does not handle masked arrays, because the calculation + does not make sense with missing values. + + Like `scipy.stats.chisquare`, this function computes a chi-square + statistic; the convenience this function provides is to figure out the + expected frequencies and degrees of freedom from the given contingency + table. If these were already known, and if the Yates' correction was not + required, one could use `scipy.stats.chisquare`. That is, if one calls:: + + res = chi2_contingency(obs, correction=False) + + then the following is true:: + + (res.statistic, res.pvalue) == stats.chisquare(obs.ravel(), + f_exp=ex.ravel(), + ddof=obs.size - 1 - dof) + + The `lambda_` argument was added in version 0.13.0 of scipy. + + References + ---------- + .. [1] "Contingency table", + https://en.wikipedia.org/wiki/Contingency_table + .. [2] "Pearson's chi-squared test", + https://en.wikipedia.org/wiki/Pearson%27s_chi-squared_test + .. [3] Cressie, N. and Read, T. R. C., "Multinomial Goodness-of-Fit + Tests", J. Royal Stat. Soc. Series B, Vol. 46, No. 3 (1984), + pp. 440-464. + + Examples + -------- + A two-way example (2 x 3): + + >>> import numpy as np + >>> from scipy.stats import chi2_contingency + >>> obs = np.array([[10, 10, 20], [20, 20, 20]]) + >>> res = chi2_contingency(obs) + >>> res.statistic + 2.7777777777777777 + >>> res.pvalue + 0.24935220877729619 + >>> res.dof + 2 + >>> res.expected_freq + array([[ 12., 12., 16.], + [ 18., 18., 24.]]) + + Perform the test using the log-likelihood ratio (i.e. the "G-test") + instead of Pearson's chi-squared statistic. + + >>> res = chi2_contingency(obs, lambda_="log-likelihood") + >>> res.statistic + 2.7688587616781319 + >>> res.pvalue + 0.25046668010954165 + + A four-way example (2 x 2 x 2 x 2): + + >>> obs = np.array( + ... [[[[12, 17], + ... [11, 16]], + ... [[11, 12], + ... [15, 16]]], + ... [[[23, 15], + ... [30, 22]], + ... [[14, 17], + ... [15, 16]]]]) + >>> res = chi2_contingency(obs) + >>> res.statistic + 8.7584514426741897 + >>> res.pvalue + 0.64417725029295503 + + When the sum of the elements in a two-way table is small, the p-value + produced by the default asymptotic approximation may be inaccurate. + Consider passing a `PermutationMethod` or `MonteCarloMethod` as the + `method` parameter with `correction=False`. + + >>> from scipy.stats import PermutationMethod + >>> obs = np.asarray([[12, 3], + ... [17, 16]]) + >>> res = chi2_contingency(obs, correction=False) + >>> ref = chi2_contingency(obs, correction=False, method=PermutationMethod()) + >>> res.pvalue, ref.pvalue + (0.0614122539870913, 0.1074) # may vary + + For a more detailed example, see :ref:`hypothesis_chi2_contingency`. + + """ + observed = np.asarray(observed) + if np.any(observed < 0): + raise ValueError("All values in `observed` must be nonnegative.") + if observed.size == 0: + raise ValueError("No data; `observed` has size 0.") + + expected = expected_freq(observed) + if np.any(expected == 0): + # Include one of the positions where expected is zero in + # the exception message. + zeropos = list(zip(*np.nonzero(expected == 0)))[0] + raise ValueError("The internally computed table of expected " + f"frequencies has a zero element at {zeropos}.") + + if method is not None: + return _chi2_resampling_methods(observed, expected, correction, lambda_, method) + + # The degrees of freedom + dof = expected.size - sum(expected.shape) + expected.ndim - 1 + + if dof == 0: + # Degenerate case; this occurs when `observed` is 1D (or, more + # generally, when it has only one nontrivial dimension). In this + # case, we also have observed == expected, so chi2 is 0. + chi2 = 0.0 + p = 1.0 + else: + if dof == 1 and correction: + # Adjust `observed` according to Yates' correction for continuity. + # Magnitude of correction no bigger than difference; see gh-13875 + diff = expected - observed + direction = np.sign(diff) + magnitude = np.minimum(0.5, np.abs(diff)) + observed = observed + magnitude * direction + + chi2, p = power_divergence(observed, expected, + ddof=observed.size - 1 - dof, axis=None, + lambda_=lambda_) + + return Chi2ContingencyResult(chi2, p, dof, expected) + + +def _chi2_resampling_methods(observed, expected, correction, lambda_, method): + + if observed.ndim != 2: + message = 'Use of `method` is only compatible with two-way tables.' + raise ValueError(message) + + if correction: + message = f'`{correction=}` is not compatible with `{method=}.`' + raise ValueError(message) + + if lambda_ is not None: + message = f'`{lambda_=}` is not compatible with `{method=}.`' + raise ValueError(message) + + if isinstance(method, stats.PermutationMethod): + res = _chi2_permutation_method(observed, expected, method) + elif isinstance(method, stats.MonteCarloMethod): + res = _chi2_monte_carlo_method(observed, expected, method) + else: + message = (f'`{method=}` not recognized; if provided, `method` must be an ' + 'instance of `PermutationMethod` or `MonteCarloMethod`.') + raise ValueError(message) + + return Chi2ContingencyResult(res.statistic, res.pvalue, np.nan, expected) + + +def _chi2_permutation_method(observed, expected, method): + x, y = _untabulate(observed) + # `permutation_test` with `permutation_type='pairings' permutes the order of `x`, + # which pairs observations in `x` with different observations in `y`. + def statistic(x): + # crosstab the resample and compute the statistic + table = crosstab(x, y)[1] + return np.sum((table - expected)**2/expected) + + return stats.permutation_test((x,), statistic, permutation_type='pairings', + alternative='greater', **method._asdict()) + + +def _chi2_monte_carlo_method(observed, expected, method): + method = method._asdict() + + if method.pop('rvs', None) is not None: + message = ('If the `method` argument of `chi2_contingency` is an ' + 'instance of `MonteCarloMethod`, its `rvs` attribute ' + 'must be unspecified. Use the `MonteCarloMethod` `rng` argument ' + 'to control the random state.') + raise ValueError(message) + rng = np.random.default_rng(method.pop('rng', None)) + + # `random_table.rvs` produces random contingency tables with the given marginals + # under the null hypothesis of independence + rowsums, colsums = stats.contingency.margins(observed) + X = stats.random_table(rowsums.ravel(), colsums.ravel(), seed=rng) + def rvs(size): + n_resamples = size[0] + return X.rvs(size=n_resamples).reshape(size) + + expected = expected.ravel() + def statistic(table, axis): + return np.sum((table - expected)**2/expected, axis=axis) + + return stats.monte_carlo_test(observed.ravel(), rvs, statistic, + alternative='greater', **method) + + +@xp_capabilities(np_only=True) +def association(observed, method="cramer", correction=False, lambda_=None): + """Calculates degree of association between two nominal variables. + + The function provides the option for computing one of three measures of + association between two nominal variables from the data given in a 2d + contingency table: Tschuprow's T, Pearson's Contingency Coefficient + and Cramer's V. + + Parameters + ---------- + observed : array-like + The array of observed values + method : {"cramer", "tschuprow", "pearson"} (default = "cramer") + The association test statistic. + correction : bool, optional + Inherited from `scipy.stats.contingency.chi2_contingency()` + lambda_ : float or str, optional + Inherited from `scipy.stats.contingency.chi2_contingency()` + + Returns + ------- + statistic : float + Value of the test statistic + + Notes + ----- + Cramer's V, Tschuprow's T and Pearson's Contingency Coefficient, all + measure the degree to which two nominal or ordinal variables are related, + or the level of their association. This differs from correlation, although + many often mistakenly consider them equivalent. Correlation measures in + what way two variables are related, whereas, association measures how + related the variables are. As such, association does not subsume + independent variables, and is rather a test of independence. A value of + 1.0 indicates perfect association, and 0.0 means the variables have no + association. + + Both the Cramer's V and Tschuprow's T are extensions of the phi + coefficient. Moreover, due to the close relationship between the + Cramer's V and Tschuprow's T the returned values can often be similar + or even equivalent. They are likely to diverge more as the array shape + diverges from a 2x2. + + References + ---------- + .. [1] "Tschuprow's T", + https://en.wikipedia.org/wiki/Tschuprow's_T + .. [2] Tschuprow, A. A. (1939) + Principles of the Mathematical Theory of Correlation; + translated by M. Kantorowitsch. W. Hodge & Co. + .. [3] "Cramer's V", https://en.wikipedia.org/wiki/Cramer's_V + .. [4] "Nominal Association: Phi and Cramer's V", + http://www.people.vcu.edu/~pdattalo/702SuppRead/MeasAssoc/NominalAssoc.html + .. [5] Gingrich, Paul, "Association Between Variables", + http://uregina.ca/~gingrich/ch11a.pdf + + Examples + -------- + An example with a 4x2 contingency table: + + >>> import numpy as np + >>> from scipy.stats.contingency import association + >>> obs4x2 = np.array([[100, 150], [203, 322], [420, 700], [320, 210]]) + + Pearson's contingency coefficient + + >>> association(obs4x2, method="pearson") + 0.18303298140595667 + + Cramer's V + + >>> association(obs4x2, method="cramer") + 0.18617813077483678 + + Tschuprow's T + + >>> association(obs4x2, method="tschuprow") + 0.14146478765062995 + """ + arr = np.asarray(observed) + if not np.issubdtype(arr.dtype, np.integer): + raise ValueError("`observed` must be an integer array.") + + if len(arr.shape) != 2: + raise ValueError("method only accepts 2d arrays") + + chi2_stat = chi2_contingency(arr, correction=correction, + lambda_=lambda_) + + phi2 = chi2_stat.statistic / arr.sum() + n_rows, n_cols = arr.shape + if method == "cramer": + value = phi2 / min(n_cols - 1, n_rows - 1) + elif method == "tschuprow": + value = phi2 / math.sqrt((n_rows - 1) * (n_cols - 1)) + elif method == 'pearson': + value = phi2 / (1 + phi2) + else: + raise ValueError("Invalid argument value: 'method' argument must " + "be 'cramer', 'tschuprow', or 'pearson'") + + return math.sqrt(value) diff --git a/v12x/pylib/scipy/stats/distributions.py b/v12x/pylib/scipy/stats/distributions.py new file mode 100644 index 0000000000000000000000000000000000000000..ac9c37aa98c9545b2616c8d32e8f676d8d49289e --- /dev/null +++ b/v12x/pylib/scipy/stats/distributions.py @@ -0,0 +1,24 @@ +# +# Author: Travis Oliphant 2002-2011 with contributions from +# SciPy Developers 2004-2011 +# +# NOTE: To look at history using `git blame`, use `git blame -M -C -C` +# instead of `git blame -Lxxx,+x`. +# +from ._distn_infrastructure import (rv_discrete, rv_continuous, rv_frozen) # noqa: F401 + +from . import _continuous_distns +from . import _discrete_distns + +from ._continuous_distns import * # noqa: F403 +from ._levy_stable import levy_stable +from ._discrete_distns import * # noqa: F403 +from ._entropy import entropy + +# For backwards compatibility e.g. pymc expects distributions.__all__. +__all__ = ['rv_discrete', 'rv_continuous', 'rv_histogram', 'entropy'] # noqa: F405 + +# Add only the distribution names, not the *_gen names. +__all__ += _continuous_distns._distn_names +__all__ += ['levy_stable'] +__all__ += _discrete_distns._distn_names diff --git a/v12x/pylib/scipy/stats/kde.py b/v12x/pylib/scipy/stats/kde.py new file mode 100644 index 0000000000000000000000000000000000000000..4401da5a30f4452ab394232d3928493d0e3b77ec --- /dev/null +++ b/v12x/pylib/scipy/stats/kde.py @@ -0,0 +1,18 @@ +# This file is not meant for public use and will be removed in SciPy v2.0.0. +# Use the `scipy.stats` namespace for importing the functions +# included below. + +from scipy._lib.deprecation import _sub_module_deprecation + + +__all__ = ["gaussian_kde"] # noqa: F822 + + +def __dir__(): + return __all__ + + +def __getattr__(name): + return _sub_module_deprecation(sub_package="stats", module="kde", + private_modules=["_kde"], all=__all__, + attribute=name) diff --git a/v12x/pylib/scipy/stats/morestats.py b/v12x/pylib/scipy/stats/morestats.py new file mode 100644 index 0000000000000000000000000000000000000000..76040ea0ca5251cabbdda7fb7936af4b281b3212 --- /dev/null +++ b/v12x/pylib/scipy/stats/morestats.py @@ -0,0 +1,27 @@ +# This file is not meant for public use and will be removed in SciPy v2.0.0. +# Use the `scipy.stats` namespace for importing the functions +# included below. + +from scipy._lib.deprecation import _sub_module_deprecation + + +__all__ = [ # noqa: F822 + 'mvsdist', + 'bayes_mvs', 'kstat', 'kstatvar', 'probplot', 'ppcc_max', 'ppcc_plot', + 'boxcox_llf', 'boxcox', 'boxcox_normmax', 'boxcox_normplot', + 'shapiro', 'anderson', 'ansari', 'bartlett', 'levene', + 'fligner', 'mood', 'wilcoxon', 'median_test', + 'circmean', 'circvar', 'circstd', 'anderson_ksamp', + 'yeojohnson_llf', 'yeojohnson', 'yeojohnson_normmax', + 'yeojohnson_normplot', 'chi2_contingency', 'distributions', +] + + +def __dir__(): + return __all__ + + +def __getattr__(name): + return _sub_module_deprecation(sub_package="stats", module="morestats", + private_modules=["_morestats"], all=__all__, + attribute=name) diff --git a/v12x/pylib/scipy/stats/mstats.py b/v12x/pylib/scipy/stats/mstats.py new file mode 100644 index 0000000000000000000000000000000000000000..88016af71803dc5c4ebadba168f22cdcd8273dbb --- /dev/null +++ b/v12x/pylib/scipy/stats/mstats.py @@ -0,0 +1,140 @@ +""" +=================================================================== +Statistical functions for masked arrays (:mod:`scipy.stats.mstats`) +=================================================================== + +.. currentmodule:: scipy.stats.mstats + +This module contains a large number of statistical functions that can +be used with masked arrays. + +Most of these functions are similar to those in `scipy.stats` but might +have small differences in the API or in the algorithm used. Since this +is a relatively new package, some API changes are still possible. + +Summary statistics +================== + +.. autosummary:: + :toctree: generated/ + + describe + gmean + hmean + kurtosis + mode + mquantiles + hdmedian + hdquantiles + hdquantiles_sd + idealfourths + plotting_positions + meppf + moment + skew + tmean + tvar + tmin + tmax + tsem + variation + find_repeats + sem + trimmed_mean + trimmed_mean_ci + trimmed_std + trimmed_var + +Frequency statistics +==================== + +.. autosummary:: + :toctree: generated/ + + scoreatpercentile + +Correlation functions +===================== + +.. autosummary:: + :toctree: generated/ + + f_oneway + pearsonr + spearmanr + pointbiserialr + kendalltau + kendalltau_seasonal + linregress + siegelslopes + theilslopes + sen_seasonal_slopes + +Statistical tests +================= + +.. autosummary:: + :toctree: generated/ + + ttest_1samp + ttest_onesamp + ttest_ind + ttest_rel + chisquare + kstest + ks_2samp + ks_1samp + ks_twosamp + mannwhitneyu + rankdata + kruskal + kruskalwallis + friedmanchisquare + brunnermunzel + skewtest + kurtosistest + normaltest + +Transformations +=============== + +.. autosummary:: + :toctree: generated/ + + obrientransform + trim + trima + trimmed_stde + trimr + trimtail + trimboth + winsorize + zmap + zscore + +Other +===== + +.. autosummary:: + :toctree: generated/ + + argstoarray + count_tied_groups + msign + compare_medians_ms + median_cihs + mjci + mquantiles_cimj + rsh + +""" +from . import _mstats_basic +from . import _mstats_extras +from ._mstats_basic import * # noqa: F403 +from ._mstats_extras import * # noqa: F403 +# Functions that support masked array input in stats but need to be kept in the +# mstats namespace for backwards compatibility: +from scipy.stats import gmean, hmean, zmap, zscore, chisquare + +__all__ = _mstats_basic.__all__ + _mstats_extras.__all__ +__all__ += ['gmean', 'hmean', 'zmap', 'zscore', 'chisquare'] diff --git a/v12x/pylib/scipy/stats/mstats_basic.py b/v12x/pylib/scipy/stats/mstats_basic.py new file mode 100644 index 0000000000000000000000000000000000000000..19cc67a6acdfa054ffa2b29b6e774dd7aafda263 --- /dev/null +++ b/v12x/pylib/scipy/stats/mstats_basic.py @@ -0,0 +1,42 @@ +# This file is not meant for public use and will be removed in SciPy v2.0.0. +# Use the `scipy.stats` namespace for importing the functions +# included below. + +from scipy._lib.deprecation import _sub_module_deprecation + + +__all__ = [ # noqa: F822 + 'argstoarray', + 'count_tied_groups', + 'describe', + 'f_oneway', 'find_repeats','friedmanchisquare', + 'kendalltau','kendalltau_seasonal','kruskal','kruskalwallis', + 'ks_twosamp', 'ks_2samp', 'kurtosis', 'kurtosistest', + 'ks_1samp', 'kstest', + 'linregress', + 'mannwhitneyu', 'meppf','mode','moment','mquantiles','msign', + 'normaltest', + 'obrientransform', + 'pearsonr','plotting_positions','pointbiserialr', + 'rankdata', + 'scoreatpercentile','sem', + 'sen_seasonal_slopes','skew','skewtest','spearmanr', + 'siegelslopes', 'theilslopes', + 'tmax','tmean','tmin','trim','trimboth', + 'trimtail','trima','trimr','trimmed_mean','trimmed_std', + 'trimmed_stde','trimmed_var','tsem','ttest_1samp','ttest_onesamp', + 'ttest_ind','ttest_rel','tvar', + 'variation', + 'winsorize', + 'brunnermunzel', +] + + +def __dir__(): + return __all__ + + +def __getattr__(name): + return _sub_module_deprecation(sub_package="stats", module="mstats_basic", + private_modules=["_mstats_basic"], all=__all__, + attribute=name, correct_module="mstats") diff --git a/v12x/pylib/scipy/stats/mstats_extras.py b/v12x/pylib/scipy/stats/mstats_extras.py new file mode 100644 index 0000000000000000000000000000000000000000..fec695329cf2c2d58a4918cc99e209c0650c3ea6 --- /dev/null +++ b/v12x/pylib/scipy/stats/mstats_extras.py @@ -0,0 +1,25 @@ +# This file is not meant for public use and will be removed in SciPy v2.0.0. +# Use the `scipy.stats` namespace for importing the functions +# included below. + +from scipy._lib.deprecation import _sub_module_deprecation + + +__all__ = [ # noqa: F822 + 'compare_medians_ms', + 'hdquantiles', 'hdmedian', 'hdquantiles_sd', + 'idealfourths', + 'median_cihs','mjci','mquantiles_cimj', + 'rsh', + 'trimmed_mean_ci', +] + + +def __dir__(): + return __all__ + + +def __getattr__(name): + return _sub_module_deprecation(sub_package="stats", module="mstats_extras", + private_modules=["_mstats_extras"], all=__all__, + attribute=name, correct_module="mstats") diff --git a/v12x/pylib/scipy/stats/mvn.py b/v12x/pylib/scipy/stats/mvn.py new file mode 100644 index 0000000000000000000000000000000000000000..65da9e20f6a4e6d24c1cb206c59821730fb6ab83 --- /dev/null +++ b/v12x/pylib/scipy/stats/mvn.py @@ -0,0 +1,17 @@ +# This file is not meant for public use and will be removed in SciPy v2.0.0. +# Use the `scipy.stats` namespace for importing the functions +# included below. + +from scipy._lib.deprecation import _sub_module_deprecation + +__all__: list[str] = [] + + +def __dir__(): + return __all__ + + +def __getattr__(name): + return _sub_module_deprecation(sub_package="stats", module="mvn", + private_modules=["_mvn"], all=__all__, + attribute=name) diff --git a/v12x/pylib/scipy/stats/qmc.py b/v12x/pylib/scipy/stats/qmc.py new file mode 100644 index 0000000000000000000000000000000000000000..a8a08343cf4c759938b31c29e32aaa644bf6e0fd --- /dev/null +++ b/v12x/pylib/scipy/stats/qmc.py @@ -0,0 +1,236 @@ +r""" +==================================================== +Quasi-Monte Carlo submodule (:mod:`scipy.stats.qmc`) +==================================================== + +.. currentmodule:: scipy.stats.qmc + +This module provides Quasi-Monte Carlo generators and associated helper +functions. + + +Quasi-Monte Carlo +================= + +Engines +------- + +.. autosummary:: + :toctree: generated/ + + QMCEngine + Sobol + Halton + LatinHypercube + PoissonDisk + MultinomialQMC + MultivariateNormalQMC + +Helpers +------- + +.. autosummary:: + :toctree: generated/ + + discrepancy + geometric_discrepancy + update_discrepancy + scale + + +Introduction to Quasi-Monte Carlo +================================= + +Quasi-Monte Carlo (QMC) methods [1]_, [2]_, [3]_ provide an +:math:`n \times d` array of numbers in :math:`[0,1]`. They can be used in +place of :math:`n` points from the :math:`U[0,1]^{d}` distribution. Compared to +random points, QMC points are designed to have fewer gaps and clumps. This is +quantified by discrepancy measures [4]_. From the Koksma-Hlawka +inequality [5]_ we know that low discrepancy reduces a bound on +integration error. Averaging a function :math:`f` over :math:`n` QMC points +can achieve an integration error close to :math:`O(n^{-1})` for well +behaved functions [2]_. + +Most QMC constructions are designed for special values of :math:`n` +such as powers of 2 or large primes. Changing the sample +size by even one can degrade their performance, even their +rate of convergence [6]_. For instance :math:`n=100` points may give less +accuracy than :math:`n=64` if the method was designed for :math:`n=2^m`. + +Some QMC constructions are extensible in :math:`n`: we can find +another special sample size :math:`n' > n` and often an infinite +sequence of increasing special sample sizes. Some QMC +constructions are extensible in :math:`d`: we can increase the dimension, +possibly to some upper bound, and typically without requiring +special values of :math:`d`. Some QMC methods are extensible in +both :math:`n` and :math:`d`. + +QMC points are deterministic. That makes it hard to estimate the accuracy of +integrals estimated by averages over QMC points. Randomized QMC (RQMC) [7]_ +points are constructed so that each point is individually :math:`U[0,1]^{d}` +while collectively the :math:`n` points retain their low discrepancy. +One can make :math:`R` independent replications of RQMC points to +see how stable a computation is. From :math:`R` independent values, +a t-test (or bootstrap t-test [8]_) then gives approximate confidence +intervals on the mean value. Some RQMC methods produce a +root mean squared error that is actually :math:`o(1/n)` and smaller than +the rate seen in unrandomized QMC. An intuitive explanation is +that the error is a sum of many small ones and random errors +cancel in a way that deterministic ones do not. RQMC also +has advantages on integrands that are singular or, for other +reasons, fail to be Riemann integrable. + +(R)QMC cannot beat Bahkvalov's curse of dimension (see [9]_). For +any random or deterministic method, there are worst case functions +that will give it poor performance in high dimensions. A worst +case function for QMC might be 0 at all n points but very +large elsewhere. Worst case analyses get very pessimistic +in high dimensions. (R)QMC can bring a great improvement over +MC when the functions on which it is used are not worst case. +For instance (R)QMC can be especially effective on integrands +that are well approximated by sums of functions of +some small number of their input variables at a time [10]_, [11]_. +That property is often a surprising finding about those functions. + +Also, to see an improvement over IID MC, (R)QMC requires a bit of smoothness of +the integrand, roughly the mixed first order derivative in each direction, +:math:`\partial^d f/\partial x_1 \cdots \partial x_d`, must be integral. +For instance, a function that is 1 inside the hypersphere and 0 outside of it +has infinite variation in the sense of Hardy and Krause for any dimension +:math:`d = 2`. + +Scrambled nets are a kind of RQMC that have some valuable robustness +properties [12]_. If the integrand is square integrable, they give variance +:math:`var_{SNET} = o(1/n)`. There is a finite upper bound on +:math:`var_{SNET} / var_{MC}` that holds simultaneously for every square +integrable integrand. Scrambled nets satisfy a strong law of large numbers +for :math:`f` in :math:`L^p` when :math:`p>1`. In some +special cases there is a central limit theorem [13]_. For smooth enough +integrands they can achieve RMSE nearly :math:`O(n^{-3})`. See [12]_ +for references about these properties. + +The main kinds of QMC methods are lattice rules [14]_ and digital +nets and sequences [2]_, [15]_. The theories meet up in polynomial +lattice rules [16]_ which can produce digital nets. Lattice rules +require some form of search for good constructions. For digital +nets there are widely used default constructions. + +The most widely used QMC methods are Sobol' sequences [17]_. +These are digital nets. They are extensible in both :math:`n` and :math:`d`. +They can be scrambled. The special sample sizes are powers +of 2. Another popular method are Halton sequences [18]_. +The constructions resemble those of digital nets. The earlier +dimensions have much better equidistribution properties than +later ones. There are essentially no special sample sizes. +They are not thought to be as accurate as Sobol' sequences. +They can be scrambled. The nets of Faure [19]_ are also widely +used. All dimensions are equally good, but the special sample +sizes grow rapidly with dimension :math:`d`. They can be scrambled. +The nets of Niederreiter and Xing [20]_ have the best asymptotic +properties but have not shown good empirical performance [21]_. + +Higher order digital nets are formed by a digit interleaving process +in the digits of the constructed points. They can achieve higher +levels of asymptotic accuracy given higher smoothness conditions on :math:`f` +and they can be scrambled [22]_. There is little or no empirical work +showing the improved rate to be attained. + +Using QMC is like using the entire period of a small random +number generator. The constructions are similar and so +therefore are the computational costs [23]_. + +(R)QMC is sometimes improved by passing the points through +a baker's transformation (tent function) prior to using them. +That function has the form :math:`1-2|x-1/2|`. As :math:`x` goes from 0 to +1, this function goes from 0 to 1 and then back. It is very +useful to produce a periodic function for lattice rules [14]_, +and sometimes it improves the convergence rate [24]_. + +It is not straightforward to apply QMC methods to Markov +chain Monte Carlo (MCMC). We can think of MCMC as using +:math:`n=1` point in :math:`[0,1]^{d}` for very large :math:`d`, with +ergodic results corresponding to :math:`d \to \infty`. One proposal is +in [25]_ and under strong conditions an improved rate of convergence +has been shown [26]_. + +Returning to Sobol' points: there are many versions depending +on what are called direction numbers. Those are the result of +searches and are tabulated. A very widely used set of direction +numbers come from [27]_. It is extensible in dimension up to +:math:`d=21201`. + +References +---------- +.. [1] Owen, Art B. "Monte Carlo Book: the Quasi-Monte Carlo parts." 2019. +.. [2] Niederreiter, Harald. "Random number generation and quasi-Monte Carlo + methods." Society for Industrial and Applied Mathematics, 1992. +.. [3] Dick, Josef, Frances Y. Kuo, and Ian H. Sloan. "High-dimensional + integration: the quasi-Monte Carlo way." Acta Numerica no. 22: 133, 2013. +.. [4] Aho, A. V., C. Aistleitner, T. Anderson, K. Appel, V. Arnol'd, N. + Aronszajn, D. Asotsky et al. "W. Chen et al.(eds.), "A Panorama of + Discrepancy Theory", Sringer International Publishing, + Switzerland: 679, 2014. +.. [5] Hickernell, Fred J. "Koksma-Hlawka Inequality." Wiley StatsRef: + Statistics Reference Online, 2014. +.. [6] Owen, Art B. "On dropping the first Sobol' point." :arxiv:`2008.08051`, + 2020. +.. [7] L'Ecuyer, Pierre, and Christiane Lemieux. "Recent advances in randomized + quasi-Monte Carlo methods." In Modeling uncertainty, pp. 419-474. Springer, + New York, NY, 2002. +.. [8] DiCiccio, Thomas J., and Bradley Efron. "Bootstrap confidence + intervals." Statistical science: 189-212, 1996. +.. [9] Dimov, Ivan T. "Monte Carlo methods for applied scientists." World + Scientific, 2008. +.. [10] Caflisch, Russel E., William J. Morokoff, and Art B. Owen. "Valuation + of mortgage backed securities using Brownian bridges to reduce effective + dimension." Journal of Computational Finance: no. 1 27-46, 1997. +.. [11] Sloan, Ian H., and Henryk Wozniakowski. "When are quasi-Monte Carlo + algorithms efficient for high dimensional integrals?." Journal of Complexity + 14, no. 1 (1998): 1-33. +.. [12] Owen, Art B., and Daniel Rudolf, "A strong law of large numbers for + scrambled net integration." SIAM Review, to appear. +.. [13] Loh, Wei-Liem. "On the asymptotic distribution of scrambled net + quadrature." The Annals of Statistics 31, no. 4: 1282-1324, 2003. +.. [14] Sloan, Ian H. and S. Joe. "Lattice methods for multiple integration." + Oxford University Press, 1994. +.. [15] Dick, Josef, and Friedrich Pillichshammer. "Digital nets and sequences: + discrepancy theory and quasi-Monte Carlo integration." Cambridge University + Press, 2010. +.. [16] Dick, Josef, F. Kuo, Friedrich Pillichshammer, and I. Sloan. + "Construction algorithms for polynomial lattice rules for multivariate + integration." Mathematics of computation 74, no. 252: 1895-1921, 2005. +.. [17] Sobol', Il'ya Meerovich. "On the distribution of points in a cube and + the approximate evaluation of integrals." Zhurnal Vychislitel'noi Matematiki + i Matematicheskoi Fiziki 7, no. 4: 784-802, 1967. +.. [18] Halton, John H. "On the efficiency of certain quasi-random sequences of + points in evaluating multi-dimensional integrals." Numerische Mathematik 2, + no. 1: 84-90, 1960. +.. [19] Faure, Henri. "Discrepance de suites associees a un systeme de + numeration (en dimension s)." Acta arithmetica 41, no. 4: 337-351, 1982. +.. [20] Niederreiter, Harold, and Chaoping Xing. "Low-discrepancy sequences and + global function fields with many rational places." Finite Fields and their + applications 2, no. 3: 241-273, 1996. +.. [21] Hong, Hee Sun, and Fred J. Hickernell. "Algorithm 823: Implementing + scrambled digital sequences." ACM Transactions on Mathematical Software + (TOMS) 29, no. 2: 95-109, 2003. +.. [22] Dick, Josef. "Higher order scrambled digital nets achieve the optimal + rate of the root mean square error for smooth integrands." The Annals of + Statistics 39, no. 3: 1372-1398, 2011. +.. [23] Niederreiter, Harald. "Multidimensional numerical integration using + pseudorandom numbers." In Stochastic Programming 84 Part I, pp. 17-38. + Springer, Berlin, Heidelberg, 1986. +.. [24] Hickernell, Fred J. "Obtaining O (N-2+e) Convergence for Lattice + Quadrature Rules." In Monte Carlo and Quasi-Monte Carlo Methods 2000, + pp. 274-289. Springer, Berlin, Heidelberg, 2002. +.. [25] Owen, Art B., and Seth D. Tribble. "A quasi-Monte Carlo Metropolis + algorithm." Proceedings of the National Academy of Sciences 102, + no. 25: 8844-8849, 2005. +.. [26] Chen, Su. "Consistency and convergence rate of Markov chain quasi Monte + Carlo with examples." PhD diss., Stanford University, 2011. +.. [27] Joe, Stephen, and Frances Y. Kuo. "Constructing Sobol sequences with + better two-dimensional projections." SIAM Journal on Scientific Computing + 30, no. 5: 2635-2654, 2008. + +""" +from ._qmc import * # noqa: F403 +from ._qmc import __all__ # noqa: F401 diff --git a/v12x/pylib/scipy/stats/sampling.py b/v12x/pylib/scipy/stats/sampling.py new file mode 100644 index 0000000000000000000000000000000000000000..12174d9dfb3cb93fa33811ed4b5d233817512e36 --- /dev/null +++ b/v12x/pylib/scipy/stats/sampling.py @@ -0,0 +1,73 @@ +""" +====================================================== +Random Number Generators (:mod:`scipy.stats.sampling`) +====================================================== + +.. currentmodule:: scipy.stats.sampling + +This module contains a collection of random number generators to sample +from univariate continuous and discrete distributions. It uses the +implementation of a C library called "UNU.RAN". The only exception is +RatioUniforms, which is a pure Python implementation of the +Ratio-of-Uniforms method. + +Generators Wrapped +================== + +For continuous distributions +---------------------------- + +.. autosummary:: + :toctree: generated/ + + NumericalInverseHermite + NumericalInversePolynomial + TransformedDensityRejection + SimpleRatioUniforms + RatioUniforms + +For discrete distributions +-------------------------- + +.. autosummary:: + :toctree: generated/ + + DiscreteAliasUrn + DiscreteGuideTable + +Warnings / Errors used in :mod:`scipy.stats.sampling` +----------------------------------------------------- + +.. autosummary:: + :toctree: generated/ + + UNURANError + + +Generators for pre-defined distributions +======================================== + +To easily apply the above methods for some of the continuous distributions +in :mod:`scipy.stats`, the following functionality can be used: + +.. autosummary:: + :toctree: generated/ + + FastGeneratorInversion + +""" +from ._sampling import FastGeneratorInversion, RatioUniforms # noqa: F401 +from ._unuran.unuran_wrapper import ( # noqa: F401 + TransformedDensityRejection, + DiscreteAliasUrn, + DiscreteGuideTable, + NumericalInversePolynomial, + NumericalInverseHermite, + SimpleRatioUniforms, + UNURANError +) + +__all__ = ["NumericalInverseHermite", "NumericalInversePolynomial", + "TransformedDensityRejection", "SimpleRatioUniforms", + "RatioUniforms", "DiscreteAliasUrn", "DiscreteGuideTable", + "UNURANError", "FastGeneratorInversion"] diff --git a/v12x/pylib/scipy/stats/stats.py b/v12x/pylib/scipy/stats/stats.py new file mode 100644 index 0000000000000000000000000000000000000000..d0589a9907d3cd181739135a3e6fd1d14d18dc50 --- /dev/null +++ b/v12x/pylib/scipy/stats/stats.py @@ -0,0 +1,43 @@ +# This file is not meant for public use and will be removed in SciPy v2.0.0. +# Use the `scipy.stats` namespace for importing the functions +# included below. + +from scipy._lib.deprecation import _sub_module_deprecation + + +__all__ = [ # noqa: F822 + 'gmean', 'hmean', 'pmean', 'mode', 'tmean', 'tvar', + 'tmin', 'tmax', 'tstd', 'tsem', 'moment', + 'skew', 'kurtosis', 'describe', 'skewtest', 'kurtosistest', + 'normaltest', 'jarque_bera', + 'scoreatpercentile', 'percentileofscore', + 'cumfreq', 'relfreq', 'obrientransform', + 'sem', 'zmap', 'zscore', 'gzscore', 'iqr', 'gstd', + 'median_abs_deviation', + 'sigmaclip', 'trimboth', 'trim1', 'trim_mean', + 'f_oneway', + 'pearsonr', 'fisher_exact', + 'spearmanr', 'pointbiserialr', + 'kendalltau', 'weightedtau', 'multiscale_graphcorr', + 'linregress', 'siegelslopes', 'theilslopes', 'ttest_1samp', + 'ttest_ind', 'ttest_ind_from_stats', 'ttest_rel', + 'kstest', 'ks_1samp', 'ks_2samp', + 'chisquare', 'power_divergence', + 'tiecorrect', 'ranksums', 'kruskal', 'friedmanchisquare', + 'rankdata', + 'combine_pvalues', 'wasserstein_distance', 'energy_distance', + 'brunnermunzel', 'alexandergovern', 'distributions', + 'mstats_basic', +] + + +def __dir__(): + return __all__ + + +def __getattr__(name): + return _sub_module_deprecation( + sub_package="stats", module="stats", + private_modules=["_stats_py", "_mgc", "_correlation"], + all=__all__, + attribute=name) diff --git a/v12x/pylib/scipy/stats/tests/__init__.py b/v12x/pylib/scipy/stats/tests/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/v12x/pylib/scipy/stats/tests/common_tests.py b/v12x/pylib/scipy/stats/tests/common_tests.py new file mode 100644 index 0000000000000000000000000000000000000000..981bfc3061e3158116a9a2c9f448eb1eb9ff2c61 --- /dev/null +++ b/v12x/pylib/scipy/stats/tests/common_tests.py @@ -0,0 +1,347 @@ +import pickle + +import numpy as np +import numpy.testing as npt +from numpy.testing import assert_allclose, assert_equal +from pytest import raises as assert_raises + +import numpy.ma.testutils as ma_npt + +from scipy._lib._util import getfullargspec_no_self as _getfullargspec +from scipy._lib._array_api_no_0d import xp_assert_equal +from scipy import stats + + +def check_named_results(res, attributes, ma=False, xp=None): + for i, attr in enumerate(attributes): + if ma: + ma_npt.assert_equal(res[i], getattr(res, attr)) + elif xp is not None: + xp_assert_equal(res[i], getattr(res, attr)) + else: + npt.assert_equal(res[i], getattr(res, attr)) + + +def check_normalization(distfn, args, distname): + norm_moment = distfn.moment(0, *args) + npt.assert_allclose(norm_moment, 1.0) + + if distname == "rv_histogram_instance": + atol, rtol = 1e-5, 0 + else: + atol, rtol = 1e-7, 1e-7 + + normalization_expect = distfn.expect(lambda x: 1, args=args) + npt.assert_allclose(normalization_expect, 1.0, atol=atol, rtol=rtol, + err_msg=distname, verbose=True) + + _a, _b = distfn.support(*args) + normalization_cdf = distfn.cdf(_b, *args) + npt.assert_allclose(normalization_cdf, 1.0) + + +def check_moment(distfn, arg, m, v, msg): + m1 = distfn.moment(1, *arg) + m2 = distfn.moment(2, *arg) + if not np.isinf(m): + npt.assert_almost_equal(m1, m, decimal=10, + err_msg=msg + ' - 1st moment') + else: # or np.isnan(m1), + npt.assert_(np.isinf(m1), + msg + f' - 1st moment -infinite, m1={str(m1)}') + + if not np.isinf(v): + npt.assert_almost_equal(m2 - m1 * m1, v, decimal=10, + err_msg=msg + ' - 2ndt moment') + else: # or np.isnan(m2), + npt.assert_(np.isinf(m2), msg + f' - 2nd moment -infinite, {m2=}') + + +def check_mean_expect(distfn, arg, m, msg): + if np.isfinite(m): + m1 = distfn.expect(lambda x: x, arg) + npt.assert_almost_equal(m1, m, decimal=5, + err_msg=msg + ' - 1st moment (expect)') + + +def check_var_expect(distfn, arg, m, v, msg): + dist_looser_tolerances = {"rv_histogram_instance" , "ksone"} + kwargs = {'rtol': 5e-6} if msg in dist_looser_tolerances else {} + if np.isfinite(v): + m2 = distfn.expect(lambda x: x*x, arg) + npt.assert_allclose(m2, v + m*m, **kwargs) + + +def check_skew_expect(distfn, arg, m, v, s, msg): + if np.isfinite(s): + m3e = distfn.expect(lambda x: np.power(x-m, 3), arg) + npt.assert_almost_equal(m3e, s * np.power(v, 1.5), + decimal=5, err_msg=msg + ' - skew') + else: + npt.assert_(np.isnan(s)) + + +def check_kurt_expect(distfn, arg, m, v, k, msg): + if np.isfinite(k): + m4e = distfn.expect(lambda x: np.power(x-m, 4), arg) + npt.assert_allclose(m4e, (k + 3.) * np.power(v, 2), + atol=1e-5, rtol=1e-5, + err_msg=msg + ' - kurtosis') + elif not np.isposinf(k): + npt.assert_(np.isnan(k)) + + +def check_munp_expect(dist, args, msg): + # If _munp is overridden, test a higher moment. (Before gh-18634, some + # distributions had issues with moments 5 and higher.) + if dist._munp.__func__ != stats.rv_continuous._munp: + res = dist.moment(5, *args) # shouldn't raise an error + ref = dist.expect(lambda x: x ** 5, args, lb=-np.inf, ub=np.inf) + if not np.isfinite(res): # could be valid; automated test can't know + return + # loose tolerance, mostly to see whether _munp returns *something* + assert_allclose(res, ref, atol=1e-10, rtol=1e-4, + err_msg=msg + ' - higher moment / _munp') + + +def check_entropy(distfn, arg, msg): + ent = distfn.entropy(*arg) + npt.assert_(not np.isnan(ent), msg + 'test Entropy is nan') + + +def check_private_entropy(distfn, args, superclass): + # compare a generic _entropy with the distribution-specific implementation + npt.assert_allclose(distfn._entropy(*args), + superclass._entropy(distfn, *args)) + + +def check_entropy_vect_scale(distfn, arg): + # check 2-d + sc = np.asarray([[1, 2], [3, 4]]) + v_ent = distfn.entropy(*arg, scale=sc) + s_ent = [distfn.entropy(*arg, scale=s) for s in sc.ravel()] + s_ent = np.asarray(s_ent).reshape(v_ent.shape) + assert_allclose(v_ent, s_ent, atol=1e-14) + + # check invalid value, check cast + sc = [1, 2, -3] + v_ent = distfn.entropy(*arg, scale=sc) + s_ent = [distfn.entropy(*arg, scale=s) for s in sc] + s_ent = np.asarray(s_ent).reshape(v_ent.shape) + assert_allclose(v_ent, s_ent, atol=1e-14) + + +def check_edge_support(distfn, args): + # Make sure that x=self.a and self.b are handled correctly. + x = distfn.support(*args) + if isinstance(distfn, stats.rv_discrete): + x = x[0]-1, x[1] + + npt.assert_equal(distfn.cdf(x, *args), [0.0, 1.0]) + npt.assert_equal(distfn.sf(x, *args), [1.0, 0.0]) + + if distfn.name not in ('skellam', 'dlaplace'): + # with a = -inf, log(0) generates warnings + npt.assert_equal(distfn.logcdf(x, *args), [-np.inf, 0.0]) + npt.assert_equal(distfn.logsf(x, *args), [0.0, -np.inf]) + + npt.assert_equal(distfn.ppf([0.0, 1.0], *args), x) + npt.assert_equal(distfn.isf([0.0, 1.0], *args), x[::-1]) + + # out-of-bounds for isf & ppf + npt.assert_(np.isnan(distfn.isf([-1, 2], *args)).all()) + npt.assert_(np.isnan(distfn.ppf([-1, 2], *args)).all()) + + +def check_named_args(distfn, x, shape_args, defaults, meths): + ## Check calling w/ named arguments. + + # check consistency of shapes, numargs and _parse signature + signature = _getfullargspec(distfn._parse_args) + npt.assert_(signature.varargs is None) + npt.assert_(signature.varkw is None) + npt.assert_(not signature.kwonlyargs) + npt.assert_(list(signature.defaults) == list(defaults)) + + shape_argnames = signature.args[:-len(defaults)] # a, b, loc=0, scale=1 + if distfn.shapes: + shapes_ = distfn.shapes.replace(',', ' ').split() + else: + shapes_ = '' + npt.assert_(len(shapes_) == distfn.numargs) + npt.assert_(len(shapes_) == len(shape_argnames)) + + # check calling w/ named arguments + shape_args = list(shape_args) + + vals = [meth(x, *shape_args) for meth in meths] + npt.assert_(np.all(np.isfinite(vals))) + + names, a, k = shape_argnames[:], shape_args[:], {} + while names: + k.update({names.pop(): a.pop()}) + v = [meth(x, *a, **k) for meth in meths] + npt.assert_array_equal(vals, v) + if 'n' not in k.keys(): + # `n` is first parameter of moment(), so can't be used as named arg + npt.assert_equal(distfn.moment(1, *a, **k), + distfn.moment(1, *shape_args)) + + # unknown arguments should not go through: + k.update({'kaboom': 42}) + assert_raises(TypeError, distfn.cdf, x, **k) + + +def check_random_state_property(distfn, args): + # check the random_state attribute of a distribution *instance* + + # baseline: this relies on the global state + np.random.seed(1234) # valid use of np.random.seed + distfn.random_state = None + r0 = distfn.rvs(*args, size=8) + + # use an explicit instance-level random_state + distfn.random_state = 1234 + r1 = distfn.rvs(*args, size=8) + npt.assert_equal(r0, r1) + + distfn.random_state = np.random.RandomState(1234) + r2 = distfn.rvs(*args, size=8) + npt.assert_equal(r0, r2) + + # check that np.random.Generator can be used (numpy >= 1.17) + if hasattr(np.random, 'default_rng'): + # obtain a np.random.Generator object + rng = np.random.default_rng(1234) + distfn.rvs(*args, size=1, random_state=rng) + + # can override the instance-level random_state for an individual .rvs call + distfn.random_state = 2 + orig_state = distfn.random_state.get_state() + + r3 = distfn.rvs(*args, size=8, random_state=np.random.RandomState(1234)) + npt.assert_equal(r0, r3) + + # ... and that does not alter the instance-level random_state! + npt.assert_equal(distfn.random_state.get_state(), orig_state) + + +def check_meth_dtype(distfn, arg, meths): + q0 = [0.25, 0.5, 0.75] + x0 = distfn.ppf(q0, *arg) + x_cast = [x0.astype(tp) for tp in (np.long, np.float16, np.float32, + np.float64)] + + for x in x_cast: + # casting may have clipped the values, exclude those + distfn._argcheck(*arg) + x = x[(distfn.a < x) & (x < distfn.b)] + for meth in meths: + val = meth(x, *arg) + npt.assert_(val.dtype == np.float64) + + +def check_ppf_dtype(distfn, arg): + q0 = np.asarray([0.25, 0.5, 0.75]) + q_cast = [q0.astype(tp) for tp in (np.float16, np.float32, np.float64)] + for q in q_cast: + for meth in [distfn.ppf, distfn.isf]: + val = meth(q, *arg) + npt.assert_(val.dtype == np.float64) + + +def check_cmplx_deriv(distfn, arg): + # Distributions allow complex arguments. + def deriv(f, x, *arg): + x = np.asarray(x) + h = 1e-10 + return (f(x + h*1j, *arg)/h).imag + + x0 = distfn.ppf([0.25, 0.51, 0.75], *arg) + x_cast = [x0.astype(tp) for tp in (np.long, np.float16, np.float32, + np.float64)] + + for x in x_cast: + # casting may have clipped the values, exclude those + distfn._argcheck(*arg) + x = x[(distfn.a < x) & (x < distfn.b)] + + pdf, cdf, sf = distfn.pdf(x, *arg), distfn.cdf(x, *arg), distfn.sf(x, *arg) + assert_allclose(deriv(distfn.cdf, x, *arg), pdf, rtol=1e-5) + assert_allclose(deriv(distfn.logcdf, x, *arg), pdf/cdf, rtol=1e-5) + + assert_allclose(deriv(distfn.sf, x, *arg), -pdf, rtol=1e-5) + assert_allclose(deriv(distfn.logsf, x, *arg), -pdf/sf, rtol=1e-5) + + assert_allclose(deriv(distfn.logpdf, x, *arg), + deriv(distfn.pdf, x, *arg) / distfn.pdf(x, *arg), + rtol=1e-5) + + +def check_pickling(distfn, args): + # check that a distribution instance pickles and unpickles + # pay special attention to the random_state property + + # save the random_state (restore later) + rndm = distfn.random_state + + # check unfrozen + distfn.random_state = 1234 + distfn.rvs(*args, size=8) + s = pickle.dumps(distfn) + r0 = distfn.rvs(*args, size=8) + + unpickled = pickle.loads(s) + r1 = unpickled.rvs(*args, size=8) + npt.assert_equal(r0, r1) + + # also smoke test some methods + medians = [distfn.ppf(0.5, *args), unpickled.ppf(0.5, *args)] + npt.assert_equal(medians[0], medians[1]) + npt.assert_equal(distfn.cdf(medians[0], *args), + unpickled.cdf(medians[1], *args)) + + # check frozen pickling/unpickling with rvs + frozen_dist = distfn(*args) + pkl = pickle.dumps(frozen_dist) + unpickled = pickle.loads(pkl) + + r0 = frozen_dist.rvs(size=8) + r1 = unpickled.rvs(size=8) + npt.assert_equal(r0, r1) + + # check pickling/unpickling of .fit method + if hasattr(distfn, "fit"): + fit_function = distfn.fit + pickled_fit_function = pickle.dumps(fit_function) + unpickled_fit_function = pickle.loads(pickled_fit_function) + assert fit_function.__name__ == unpickled_fit_function.__name__ == "fit" + + # restore the random_state + distfn.random_state = rndm + + +def check_freezing(distfn, args): + # regression test for gh-11089: freezing a distribution fails + # if loc and/or scale are specified + if isinstance(distfn, stats.rv_continuous): + locscale = {'loc': 1, 'scale': 2} + else: + locscale = {'loc': 1} + + rv = distfn(*args, **locscale) + assert rv.a == distfn(*args).a + assert rv.b == distfn(*args).b + + +def check_rvs_broadcast(distfunc, distname, allargs, shape, shape_only, otype): + rng = np.random.RandomState(123) + sample = distfunc.rvs(*allargs, random_state=rng) + assert_equal(sample.shape, shape, f"{distname}: rvs failed to broadcast") + if not shape_only: + rvs = np.vectorize( + lambda *allargs: distfunc.rvs(*allargs, random_state=rng), + otypes=otype) + rng = np.random.RandomState(123) + expected = rvs(*allargs) + assert_allclose(sample, expected, rtol=1e-13) diff --git a/v12x/pylib/scipy/stats/tests/test_axis_nan_policy.py b/v12x/pylib/scipy/stats/tests/test_axis_nan_policy.py new file mode 100644 index 0000000000000000000000000000000000000000..2b7f1bccccbe41d1cd40dd7479159f31282be151 --- /dev/null +++ b/v12x/pylib/scipy/stats/tests/test_axis_nan_policy.py @@ -0,0 +1,1495 @@ +# Many scipy.stats functions support `axis` and `nan_policy` parameters. +# When the two are combined, it can be tricky to get all the behavior just +# right. This file contains a suite of common tests for scipy.stats functions +# that support `axis` and `nan_policy` and additional tests for some associated +# functions in stats._util. + +from itertools import product, combinations_with_replacement, permutations +import os +import re +import pickle +import pytest +import warnings + +import numpy as np +from numpy.testing import assert_allclose, assert_equal +from scipy import stats, special +from scipy.stats._axis_nan_policy import (_masked_arrays_2_sentinel_arrays, + SmallSampleWarning, + too_small_nd_omit, too_small_nd_not_omit, + too_small_1d_omit, too_small_1d_not_omit) +from scipy._lib._util import AxisError +from scipy._lib._array_api import make_xp_test_case +from scipy.conftest import skip_xp_invalid_arg + + +SCIPY_XSLOW = int(os.environ.get('SCIPY_XSLOW', '0')) + + +def _using_accelerate(): + config = np.show_config('dicts') + return config['Build Dependencies']['blas']['name'].lower() == 'accelerate' + + +RTOL = 1e-6 if _using_accelerate() else 1e-15 + + +tolerance_overrides = {stats.epps_singleton_2samp: 1e-10} + + +def unpack_ttest_result(res): + low, high = res.confidence_interval() + return (res.statistic, res.pvalue, res.df, res._standard_error, + res._estimate, low, high) + + +def _get_ttest_ci(ttest): + # get a function that returns the CI bounds of provided `ttest` + def ttest_ci(*args, **kwargs): + res = ttest(*args, **kwargs) + return res.confidence_interval() + return ttest_ci + + +def xp_mean_1samp(*args, **kwargs): + kwargs.pop('_no_deco', None) + return stats._stats_py._xp_mean(*args, **kwargs) + + +def xp_mean_2samp(*args, **kwargs): + kwargs.pop('_no_deco', None) + weights = args[1] + return stats._stats_py._xp_mean(args[0], *args[2:], weights=weights, **kwargs) + + +def xp_var(*args, **kwargs): + kwargs.pop('_no_deco', None) + return stats._stats_py._xp_var(*args, **kwargs) + + +def gstd(*args, **kwargs): + kwargs.pop('_no_deco', None) + return stats.gstd(*args, **kwargs) + + +def combine_pvalues_weighted(*args, **kwargs): + return stats.combine_pvalues(args[0], *args[2:], weights=args[1], + method='stouffer', **kwargs) + + +def weightedtau_weighted(x, y, rank, **kwargs): + axis = kwargs.get('axis', 0) + nan_policy = kwargs.get('nan_policy', 'propagate') + rank = stats.rankdata(rank, axis=axis, nan_policy=nan_policy) + return stats.weightedtau(x, y, rank, **kwargs) + + +def boxcox_llf(data, lmb, axis=0, _no_deco=False, **kwargs): + if _no_deco: + return stats._morestats._boxcox_llf(data, lmb=lmb, axis=axis) + return stats.boxcox_llf(lmb, data, axis=axis, **kwargs) + + +def yeojohnson_llf(data, lmb, axis=0, _no_deco=False, **kwargs): + if _no_deco: + return stats._morestats._yeojohnson_llf(data, lmb=lmb, axis=axis) + return stats.yeojohnson_llf(lmb, data, axis=axis, **kwargs) + + +def kendalltau(*args, _no_deco=False, **kwargs): + if _no_deco: + return stats._stats_py._kendalltau(*args, _no_deco=_no_deco, **kwargs) + return stats.kendalltau(*args, **kwargs) + + +axis_nan_policy_cases = [ + # function, args, kwds, number of samples, number of outputs, + # ... paired, unpacker function + # args, kwds typically aren't needed; just showing that they work + (stats.kruskal, tuple(), dict(), 3, 2, False, None), # 4 samples is slow + (stats.ranksums, ('less',), dict(), 2, 2, False, None), + (stats.mannwhitneyu, tuple(), {'method': 'asymptotic'}, 2, 3, False, + lambda res: (res.statistic, res.pvalue, res.zstatistic)), + (stats.mannwhitneyu, tuple(), {'method': 'auto'}, 2, 2, False, None), + (stats.wilcoxon, ('pratt',), {'mode': 'auto'}, 2, 2, True, + lambda res: (res.statistic, res.pvalue)), + (stats.wilcoxon, tuple(), dict(), 1, 2, True, + lambda res: (res.statistic, res.pvalue)), + (stats.wilcoxon, tuple(), {'method': 'asymptotic'}, 1, 3, True, + lambda res: (res.statistic, res.pvalue, res.zstatistic)), + (stats.gmean, tuple(), dict(), 1, 1, False, lambda x: (x,)), + (stats.hmean, tuple(), dict(), 1, 1, False, lambda x: (x,)), + (stats.pmean, (1.42,), dict(), 1, 1, False, lambda x: (x,)), + (stats.trim_mean, (0.05,), dict(), 1, 1, False, lambda x: (x,)), + (stats.sem, tuple(), dict(), 1, 1, False, lambda x: (x,)), + (stats.iqr, tuple(), dict(), 1, 1, False, lambda x: (x,)), + (stats.kurtosis, tuple(), dict(), 1, 1, False, lambda x: (x,)), + (stats.skew, tuple(), dict(), 1, 1, False, lambda x: (x,)), + (stats.kstat, tuple(), dict(), 1, 1, False, lambda x: (x,)), + (stats.kstatvar, tuple(), dict(), 1, 1, False, lambda x: (x,)), + (stats.moment, tuple(), dict(), 1, 1, False, lambda x: (x,)), + (stats.moment, tuple(), dict(order=[1, 2]), 1, 2, False, None), + (stats.jarque_bera, tuple(), dict(), 1, 2, False, None), + (stats.ttest_1samp, (np.array([0]),), dict(), 1, 7, False, + unpack_ttest_result), + (stats.ttest_rel, tuple(), dict(), 2, 7, True, unpack_ttest_result), + (stats.ttest_ind, tuple(), dict(), 2, 7, False, unpack_ttest_result), + (_get_ttest_ci(stats.ttest_1samp), (0,), dict(), 1, 2, False, None), + (_get_ttest_ci(stats.ttest_rel), tuple(), dict(), 2, 2, True, None), + (_get_ttest_ci(stats.ttest_ind), tuple(), dict(), 2, 2, False, None), + (stats.mode, tuple(), dict(), 1, 2, True, lambda x: (x.mode, x.count)), + (stats.differential_entropy, tuple(), dict(), 1, 1, False, lambda x: (x,)), + (stats.variation, tuple(), dict(), 1, 1, False, lambda x: (x,)), + (stats.friedmanchisquare, tuple(), dict(), 3, 2, True, None), + (stats.brunnermunzel, tuple(), dict(distribution='normal'), 2, 2, False, None), + (stats.mood, tuple(), {}, 2, 2, False, None), + (stats.shapiro, tuple(), {}, 1, 2, False, None), + (stats.ks_1samp, (special.ndtr,), dict(), 1, 4, False, + lambda res: (*res, res.statistic_location, res.statistic_sign)), + (stats.ks_1samp, (special.ndtr,), dict(alternative='greater'), 1, 4, False, + lambda res: (*res, res.statistic_location, res.statistic_sign)), + (stats.ks_1samp, (special.ndtr,), dict(alternative='less'), 1, 4, False, + lambda res: (*res, res.statistic_location, res.statistic_sign)), + (stats.ks_2samp, tuple(), dict(), 2, 4, False, + lambda res: (*res, res.statistic_location, res.statistic_sign)), + (stats.kstest, (special.ndtr,), dict(), 1, 4, False, + lambda res: (*res, res.statistic_location, res.statistic_sign)), + (stats.kstest, tuple(), dict(), 2, 4, False, + lambda res: (*res, res.statistic_location, res.statistic_sign)), + (stats.levene, tuple(), {}, 2, 2, False, None), + (stats.fligner, tuple(), {'center': 'trimmed', 'proportiontocut': 0.01}, + 2, 2, False, None), + (stats.ansari, tuple(), {}, 2, 2, False, None), + (stats.entropy, tuple(), dict(), 1, 1, False, lambda x: (x,)), + (stats.entropy, tuple(), dict(), 2, 1, True, lambda x: (x,)), + (stats.skewtest, tuple(), dict(), 1, 2, False, None), + (stats.kurtosistest, tuple(), dict(), 1, 2, False, None), + (stats.normaltest, tuple(), dict(), 1, 2, False, None), + (stats.cramervonmises, (special.ndtr,), dict(), 1, 2, False, + lambda res: (res.statistic, res.pvalue)), + (stats.cramervonmises_2samp, tuple(), dict(method='asymptotic'), 2, 2, False, + lambda res: (res.statistic, res.pvalue)), + (stats.epps_singleton_2samp, tuple(), dict(), 2, 2, False, None), + (stats.bartlett, tuple(), {}, 2, 2, False, None), + (stats.tmean, tuple(), {}, 1, 1, False, lambda x: (x,)), + (stats.tvar, tuple(), {}, 1, 1, False, lambda x: (x,)), + (stats.tmin, tuple(), {}, 1, 1, False, lambda x: (x,)), + (stats.tmax, tuple(), {}, 1, 1, False, lambda x: (x,)), + (stats.tstd, tuple(), {}, 1, 1, False, lambda x: (x,)), + (stats.tsem, tuple(), {}, 1, 1, False, lambda x: (x,)), + (stats.circmean, tuple(), dict(), 1, 1, False, lambda x: (x,)), + (stats.circvar, tuple(), dict(), 1, 1, False, lambda x: (x,)), + (stats.circstd, tuple(), dict(), 1, 1, False, lambda x: (x,)), + (stats.f_oneway, tuple(), {}, 2, 2, False, None), + (stats.f_oneway, tuple(), {'equal_var': False}, 2, 2, False, None), + (stats.alexandergovern, tuple(), {}, 2, 2, False, + lambda res: (res.statistic, res.pvalue)), + (stats.combine_pvalues, tuple(), {}, 1, 2, False, None), + (stats.lmoment, tuple(), dict(), 1, 4, False, lambda x: tuple(x)), + (combine_pvalues_weighted, tuple(), {}, 2, 2, True, None), + (xp_mean_1samp, tuple(), dict(), 1, 1, False, lambda x: (x,)), + (xp_mean_2samp, tuple(), dict(), 2, 1, True, lambda x: (x,)), + (xp_var, tuple(), dict(), 1, 1, False, lambda x: (x,)), + (stats.chatterjeexi, tuple(), dict(), 2, 2, True, + lambda res: (res.statistic, res.pvalue)), + (stats.spearmanrho, tuple(), dict(), 2, 2, True, + lambda res: (res.statistic, res.pvalue)), + (stats.pointbiserialr, tuple(), dict(), 2, 3, True, + lambda res: (res.statistic, res.pvalue, res.correlation)), + (kendalltau, tuple(), dict(), 2, 3, True, + lambda res: (res.statistic, res.pvalue, res.correlation)), + (stats.weightedtau, tuple(), dict(), 2, 3, True, + lambda res: (res.statistic, res.pvalue, res.correlation)), + (weightedtau_weighted, tuple(), dict(), 3, 3, True, + lambda res: (res.statistic, res.pvalue, res.correlation)), + (stats.linregress, tuple(), dict(), 2, 6, True, + lambda res: tuple(res) + (res.intercept_stderr,)), + (stats.theilslopes, tuple(), dict(), 2, 4, True, tuple), + (stats.theilslopes, tuple(), dict(), 1, 4, True, tuple), + (stats.theilslopes, tuple(), dict(method='joint'), 2, 4, True, tuple), + (stats.siegelslopes, tuple(), dict(), 2, 2, True, tuple), + (stats.siegelslopes, tuple(), dict(), 1, 2, True, tuple), + (stats.siegelslopes, tuple(), dict(method='hierarchical'), 2, 2, True, tuple), + (gstd, tuple(), dict(), 1, 1, False, lambda x: (x,)), + (stats.power_divergence, tuple(), dict(), 1, 2, False, None), + (stats.chisquare, tuple(), dict(), 1, 2, False, None), + (stats.median_abs_deviation, tuple(), dict(), 1, 1, False, lambda x: (x,)), + (boxcox_llf, tuple(), dict(lmb=1.5), 1, 1, False, lambda x: (x,)), + (yeojohnson_llf, tuple(), dict(lmb=1.5), 1, 1, False, lambda x: (x,)), + (stats.expectile, (0.4,), dict(), 1, 1, False, lambda x: (x,)), +] + +# If the message is one of those expected, put nans in +# appropriate places of `statistics` and `pvalues` +too_small_messages = {"Degrees of freedom <= 0 for slice", + "x and y should have at least 5 elements", + "Data must be at least length 3", + "The sample must contain at least two", + "x and y must contain at least two", + "division by zero", + "Mean of empty slice", + "Data passed to ks_2samp must not be empty", + "Not enough test observations", + "Not enough other observations", + "Not enough observations.", + "At least one observation is required", + "zero-size array to reduction operation maximum", + "`x` and `y` must be of nonzero size.", + "The exact distribution of the Wilcoxon test", + "Data input must not be empty", + "Window length (0) must be positive and less", + "Window length (1) must be positive and less", + "Window length (2) must be positive and less", + "`skewtest` requires at least", + "`kurtosistest` requires at least", + "attempt to get argmax of an empty sequence", + "No array values within given limits", + "Input sample size must be greater than one.", + "At least one slice along `axis` has zero length", + "One or more sample arguments is too small", + "invalid value encountered", + "divide by zero encountered", + "`x` and `y` must have length at least 2.", + "Inputs must not be empty.", + "All `x` coordinates are identical.", +} + +# If the message is one of these, results of the function may be inaccurate, +# but NaNs are not to be placed +inaccuracy_messages = {"Precision loss occurred in moment calculation", + "Sample size too small for normal approximation."} + +# For some functions, nan_policy='propagate' should not just return NaNs +override_propagate_funcs = {stats.mode, weightedtau_weighted, stats.weightedtau} + +# For some functions, empty arrays produce non-NaN results +empty_special_case_funcs = {stats.entropy, stats.chisquare, stats.power_divergence} + +# Some functions don't follow the usual "too small" warning rules +too_small_special_case_funcs = {stats.entropy, stats.chisquare, stats.power_divergence} + +def _mixed_data_generator(n_samples, n_repetitions, axis, rng, + paired=False): + # generate random samples to check the response of hypothesis tests to + # samples with different (but broadcastable) shapes and various + # nan patterns (e.g. all nans, some nans, no nans) along axis-slices + + data = [] + for i in range(n_samples): + n_patterns = 6 # number of distinct nan patterns + n_obs = 20 if paired else 20 + i # observations per axis-slice + x = np.ones((n_repetitions, n_patterns, n_obs)) * np.nan + + for j in range(n_repetitions): + samples = x[j, :, :] + + # case 0: axis-slice with all nans (0 reals) + # cases 1-3: axis-slice with 1-3 reals (the rest nans) + # case 4: axis-slice with mostly (all but two) reals + # case 5: axis slice with all reals + for k, n_reals in enumerate([0, 1, 2, 3, n_obs-2, n_obs]): + # for cases 1-3, need paired nansw to be in the same place + indices = rng.permutation(n_obs)[:n_reals] + samples[k, indices] = rng.random(size=n_reals) + + # permute the axis-slices just to show that order doesn't matter + samples[:] = rng.permutation(samples, axis=0) + + # For multi-sample tests, we want to test broadcasting and check + # that nan policy works correctly for each nan pattern for each input. + # This takes care of both simultaneously. + new_shape = [n_repetitions] + [1]*n_samples + [n_obs] + new_shape[1 + i] = 6 + x = x.reshape(new_shape) + + x = np.moveaxis(x, -1, axis) + data.append(x) + return data + + +def _homogeneous_data_generator(n_samples, n_repetitions, axis, rng, + paired=False, all_nans=True): + # generate random samples to check the response of hypothesis tests to + # samples with different (but broadcastable) shapes and homogeneous + # data (all nans or all finite) + data = [] + for i in range(n_samples): + n_obs = 20 if paired else 20 + i # observations per axis-slice + shape = [n_repetitions] + [1]*n_samples + [n_obs] + shape[1 + i] = 2 + x = np.ones(shape) * np.nan if all_nans else rng.random(shape) + x = np.moveaxis(x, -1, axis) + data.append(x) + return data + + +def nan_policy_1d(hypotest, data1d, unpacker, *args, n_outputs=2, + nan_policy='raise', paired=False, _no_deco=True, **kwds): + # Reference implementation for how `nan_policy` should work for 1d samples + + if nan_policy == 'raise': + for sample in data1d: + if np.any(np.isnan(sample)): + raise ValueError("The input contains nan values") + + elif (nan_policy == 'propagate' + and hypotest not in override_propagate_funcs): + # For all hypothesis tests tested, returning nans is the right thing. + # But many hypothesis tests don't propagate correctly (e.g. they treat + # np.nan the same as np.inf, which doesn't make sense when ranks are + # involved) so override that behavior here. + for sample in data1d: + if np.any(np.isnan(sample)): + return np.full(n_outputs, np.nan) + + elif nan_policy == 'omit': + # manually omit nans (or pairs in which at least one element is nan) + if not paired: + data1d = [sample[~np.isnan(sample)] for sample in data1d] + else: + nan_mask = np.isnan(data1d[0]) + for sample in data1d[1:]: + nan_mask = np.logical_or(nan_mask, np.isnan(sample)) + data1d = [sample[~nan_mask] for sample in data1d] + + return unpacker(hypotest(*data1d, *args, _no_deco=_no_deco, **kwds)) + + +# These three warnings are intentional +# For `wilcoxon` when the sample size < 50 +@pytest.mark.filterwarnings('ignore:Sample size too small for normal:UserWarning') +# `kurtosistest` and `normaltest` when sample size < 20 +@pytest.mark.filterwarnings('ignore:`kurtosistest` p-value may be:UserWarning') +# `foneway` +@pytest.mark.filterwarnings('ignore:all input arrays have length 1.:RuntimeWarning') + +# The rest of these may or may not be desirable. They need further investigation +# to determine whether the function's decorator should define `too_small. +# `bartlett`, `tvar`, `tstd`, `tsem` +@pytest.mark.filterwarnings('ignore:Degrees of freedom <= 0 for slice:RuntimeWarning') +# kstat, kstatvar, ttest_1samp, ttest_rel, ttest_ind, ttest_ci, brunnermunzel +# mood, levene, fligner, bartlett +@pytest.mark.filterwarnings('ignore:Invalid value encountered in:RuntimeWarning') +# kstatvar, ttest_1samp, ttest_rel, ttest_ci, brunnermunzel, levene, bartlett +@pytest.mark.filterwarnings('ignore:divide by zero encountered:RuntimeWarning') + +@pytest.mark.parametrize(("hypotest", "args", "kwds", "n_samples", "n_outputs", + "paired", "unpacker"), axis_nan_policy_cases) +@pytest.mark.parametrize(("nan_policy"), ("propagate", "omit", "raise")) +@pytest.mark.parametrize(("axis"), (1,)) +@pytest.mark.parametrize(("data_generator"), ("mixed",)) +def test_axis_nan_policy_fast(hypotest, args, kwds, n_samples, n_outputs, + paired, unpacker, nan_policy, axis, + data_generator): + if hypotest in {stats.kruskal} and not SCIPY_XSLOW: + pytest.skip("Too slow.") + _axis_nan_policy_test(hypotest, args, kwds, n_samples, n_outputs, paired, + unpacker, nan_policy, axis, data_generator) + + +if SCIPY_XSLOW: + # Takes O(1 min) to run, and even skipping with the `xslow` decorator takes + # about 3 sec because this is >3,000 tests. So ensure pytest doesn't see + # them at all unless `SCIPY_XSLOW` is defined. + + # These three warnings are intentional + # For `wilcoxon` when the sample size < 50 + @pytest.mark.filterwarnings('ignore:Sample size too small for normal:UserWarning') + # `kurtosistest` and `normaltest` when sample size < 20 + @pytest.mark.filterwarnings('ignore:`kurtosistest` p-value may be:UserWarning') + # `foneway` + @pytest.mark.filterwarnings('ignore:all input arrays have length 1.:RuntimeWarning') + + # The rest of these may or may not be desirable. They need further investigation + # to determine whether the function's decorator should define `too_small. + # `bartlett`, `tvar`, `tstd`, `tsem` + @pytest.mark.filterwarnings('ignore:Degrees of freedom <= 0 for:RuntimeWarning') + # kstat, kstatvar, ttest_1samp, ttest_rel, ttest_ind, ttest_ci, brunnermunzel + # mood, levene, fligner, bartlett + @pytest.mark.filterwarnings('ignore:Invalid value encountered in:RuntimeWarning') + # kstatvar, ttest_1samp, ttest_rel, ttest_ci, brunnermunzel, levene, bartlett + @pytest.mark.filterwarnings('ignore:divide by zero encountered:RuntimeWarning') + + @pytest.mark.parametrize(("hypotest", "args", "kwds", "n_samples", "n_outputs", + "paired", "unpacker"), axis_nan_policy_cases) + @pytest.mark.parametrize(("nan_policy"), ("propagate", "omit", "raise")) + @pytest.mark.parametrize(("axis"), list(range(-3, 3))) + @pytest.mark.parametrize(("data_generator"), + ("all_nans", "all_finite", "mixed")) + def test_axis_nan_policy_full(hypotest, args, kwds, n_samples, n_outputs, + paired, unpacker, nan_policy, axis, + data_generator): + _axis_nan_policy_test(hypotest, args, kwds, n_samples, n_outputs, paired, + unpacker, nan_policy, axis, data_generator) + + +def _axis_nan_policy_test(hypotest, args, kwds, n_samples, n_outputs, paired, + unpacker, nan_policy, axis, data_generator): + # Tests the 1D and vectorized behavior of hypothesis tests against a + # reference implementation (nan_policy_1d with np.ndenumerate) + + # Some hypothesis tests return a non-iterable that needs an `unpacker` to + # extract the statistic and p-value. For those that don't: + if not unpacker: + def unpacker(res): + return res + + rng = np.random.default_rng(0) + + # Generate multi-dimensional test data with all important combinations + # of patterns of nans along `axis` + n_repetitions = 3 # number of repetitions of each pattern + data_gen_kwds = {'n_samples': n_samples, 'n_repetitions': n_repetitions, + 'axis': axis, 'rng': rng, 'paired': paired} + if data_generator == 'mixed': + inherent_size = 6 # number of distinct types of patterns + data = _mixed_data_generator(**data_gen_kwds) + elif data_generator == 'all_nans': + inherent_size = 2 # hard-coded in _homogeneous_data_generator + data_gen_kwds['all_nans'] = True + data = _homogeneous_data_generator(**data_gen_kwds) + elif data_generator == 'all_finite': + inherent_size = 2 # hard-coded in _homogeneous_data_generator + data_gen_kwds['all_nans'] = False + data = _homogeneous_data_generator(**data_gen_kwds) + + output_shape = [n_repetitions] + [inherent_size]*n_samples + + # To generate reference behavior to compare against, loop over the axis- + # slices in data. Make indexing easier by moving `axis` to the end and + # broadcasting all samples to the same shape. + data_b = [np.moveaxis(sample, axis, -1) for sample in data] + data_b = [np.broadcast_to(sample, output_shape + [sample.shape[-1]]) + for sample in data_b] + res_1d = np.zeros(output_shape + [n_outputs]) + + for i, _ in np.ndenumerate(np.zeros(output_shape)): + data1d = [sample[i] for sample in data_b] + contains_nan = any([np.isnan(sample).any() for sample in data1d]) + + # Take care of `nan_policy='raise'`. + # Afterward, the 1D part of the test is over + message = "The input contains nan values" + if nan_policy == 'raise' and contains_nan: + with pytest.raises(ValueError, match=message): + nan_policy_1d(hypotest, data1d, unpacker, *args, + n_outputs=n_outputs, + nan_policy=nan_policy, + paired=paired, _no_deco=True, **kwds) + + with pytest.raises(ValueError, match=message): + hypotest(*data1d, *args, nan_policy=nan_policy, **kwds) + + continue + + # Take care of `nan_policy='propagate'` and `nan_policy='omit'` + + # Get results of simple reference implementation + try: + res_1da = nan_policy_1d(hypotest, data1d, unpacker, *args, + n_outputs=n_outputs, + nan_policy=nan_policy, + paired=paired, _no_deco=True, **kwds) + except (ValueError, RuntimeWarning, ZeroDivisionError, UserWarning) as ea: + ea_str = str(ea) + if any([str(ea_str).startswith(msg) for msg in too_small_messages]): + res_1da = np.full(n_outputs, np.nan) + else: + raise + + # Get results of public function with 1D slices + # Should warn for all slices + if (nan_policy == 'omit' and data_generator == "all_nans" + and hypotest not in too_small_special_case_funcs): + with pytest.warns(SmallSampleWarning, match=too_small_1d_omit): + res = hypotest(*data1d, *args, nan_policy=nan_policy, **kwds) + # warning depends on slice + elif (nan_policy == 'omit' and data_generator == "mixed" + and hypotest not in too_small_special_case_funcs): + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", too_small_1d_omit, SmallSampleWarning) + res = hypotest(*data1d, *args, nan_policy=nan_policy, **kwds) + # shouldn't complain if there are no NaNs + else: + res = hypotest(*data1d, *args, nan_policy=nan_policy, **kwds) + res_1db = unpacker(res) + + assert_allclose(res_1db, res_1da, rtol=RTOL) + res_1d[i] = res_1db + + res_1d = np.moveaxis(res_1d, -1, 0) + + # Perform a vectorized call to the hypothesis test. + + # If `nan_policy == 'raise'`, check that it raises the appropriate error. + # Test is done, so return + if nan_policy == 'raise' and not data_generator == "all_finite": + message = 'The input contains nan values' + with pytest.raises(ValueError, match=message): + hypotest(*data, axis=axis, nan_policy=nan_policy, *args, **kwds) + return + + # If `nan_policy == 'omit', we might be left with a small sample. + # Check for the appropriate warning. + if (nan_policy == 'omit' and data_generator in {"all_nans", "mixed"} + and hypotest not in too_small_special_case_funcs): + with pytest.warns(SmallSampleWarning, match=too_small_nd_omit): + res = hypotest(*data, axis=axis, nan_policy=nan_policy, *args, **kwds) + else: # otherwise, there should be no warning + res = hypotest(*data, axis=axis, nan_policy=nan_policy, *args, **kwds) + + # Compare against the output against looping over 1D slices + res_nd = unpacker(res) + + rtol = max(tolerance_overrides.get(hypotest, RTOL), RTOL) + assert_allclose(res_nd, res_1d, rtol=rtol) + +# nan should not raise an exception in np.mean() +# but does on some mips64el systems, triggering failure in some test cases +# see https://github.com/scipy/scipy/issues/22360 +# and https://github.com/numpy/numpy/issues/23158 +def skip_nan_unexpected_exception(): + try: + # should not raise an exception + with np.errstate(all='raise'): + x = np.asarray([1, 2, np.nan]) + np.mean(x) + except Exception as e: + pytest.skip(f"nan raises unexpected {e.__class__.__name__} in numpy") + +@pytest.mark.parametrize(("hypotest", "args", "kwds", "n_samples", "n_outputs", + "paired", "unpacker"), axis_nan_policy_cases) +@pytest.mark.parametrize(("nan_policy"), ("propagate", "omit", "raise")) +@pytest.mark.parametrize(("data_generator"), + ("all_nans", "all_finite", "mixed", "empty")) +def test_axis_nan_policy_axis_is_None(hypotest, args, kwds, n_samples, + n_outputs, paired, unpacker, nan_policy, + data_generator): + # check for correct behavior when `axis=None` + if not unpacker: + def unpacker(res): + return res + + # skip if nan emits unexpected RuntimeWarning in np.mean() + # seen on mips64el, https://github.com/scipy/scipy/issues/22360 + # Only affects nan_policy "mixed-propagate" for selected hypotests + if data_generator=="mixed" and nan_policy=="propagate": + # only skip affected hypotests + if hypotest.__name__ in ["iqr", "ttest_ci", + "xp_mean_1samp", "xp_mean_2samp", "xp_var", + "weightedtau", "weightedtau_weighted"]: + skip_nan_unexpected_exception() + # all_nans-propagate-ttest_ci is also affected, via scalar multiply + if (data_generator=="all_nans" and nan_policy=="propagate" + and hypotest.__name__=="ttest_ci"): + skip_nan_unexpected_exception() + # mixed-omit-xp_var is also affected, via subtract + if (data_generator=="mixed" and nan_policy=="omit" + and hypotest.__name__=="xp_var"): + skip_nan_unexpected_exception() + + rng = np.random.default_rng(0) + + if data_generator == "empty": + data = [rng.random((2, 0)) for i in range(n_samples)] + else: + data = [rng.random((2, 20)) for i in range(n_samples)] + + if data_generator == "mixed": + masks = [rng.random((2, 20)) > 0.9 for i in range(n_samples)] + for sample, mask in zip(data, masks): + sample[mask] = np.nan + elif data_generator == "all_nans": + data = [sample * np.nan for sample in data] + + data_raveled = [sample.ravel() for sample in data] + + if nan_policy == 'raise' and data_generator not in {"all_finite", "empty"}: + message = 'The input contains nan values' + + # check for correct behavior whether or not data is 1d to begin with + with pytest.raises(ValueError, match=message): + hypotest(*data, axis=None, nan_policy=nan_policy, + *args, **kwds) + with pytest.raises(ValueError, match=message): + hypotest(*data_raveled, axis=None, nan_policy=nan_policy, + *args, **kwds) + + return + + # behavior of reference implementation with 1d input, public function with 1d + # input, and public function with Nd input and `axis=None` should be consistent. + # This means: + # - If the reference version raises an error or emits a warning, it's because + # the sample is too small, so check that the public function emits an + # appropriate "too small" warning + # - Any results returned by the three versions should be the same. + with warnings.catch_warnings(): # treat warnings as errors + warnings.simplefilter("error") + + ea_str, eb_str, ec_str = None, None, None + try: + res1da = nan_policy_1d(hypotest, data_raveled, unpacker, *args, + n_outputs=n_outputs, nan_policy=nan_policy, + paired=paired, _no_deco=True, **kwds) + except (RuntimeWarning, ValueError, ZeroDivisionError, UserWarning) as ea: + res1da = None + ea_str = str(ea) + + try: + res1db = hypotest(*data_raveled, *args, nan_policy=nan_policy, **kwds) + except SmallSampleWarning as eb: + eb_str = str(eb) + + try: + res1dc = hypotest(*data, *args, axis=None, nan_policy=nan_policy, **kwds) + except SmallSampleWarning as ec: + ec_str = str(ec) + + if ea_str or eb_str or ec_str: # *if* there is some sort of error or warning + # If the reference implemented generated an error or warning, make sure the + # message was one of the expected "too small" messages. Note that some + # functions don't complain at all without the decorator; that's OK, too. + ok_msg = any([str(ea_str).startswith(msg) for msg in too_small_messages]) + assert (ea_str is None) or ok_msg + + # make sure the wrapped function emits the *intended* warning + desired_warnings = {too_small_1d_omit, too_small_1d_not_omit} + assert str(eb_str) in desired_warnings + assert str(ec_str) in desired_warnings + + with warnings.catch_warnings(): # ignore warnings to get return value + warnings.simplefilter("ignore") + res1db = hypotest(*data_raveled, *args, nan_policy=nan_policy, **kwds) + res1dc = hypotest(*data, *args, axis=None, nan_policy=nan_policy, **kwds) + + # Make sure any results returned by reference/public function are identical + # and all attributes are *NumPy* scalars + res1db, res1dc = unpacker(res1db), unpacker(res1dc) + # changed from 1e-15 solely to appease macosx-x86_64+Accelerate + assert_allclose(res1dc, res1db, rtol=2e-14) + all_results = list(res1db) + list(res1dc) + + if res1da is not None: + # changed from 1e-15 solely to appease macosx-x86_64+Accelerate + assert_allclose(res1db, res1da, rtol=2e-14) + all_results += list(res1da) + + for item in all_results: + assert np.issubdtype(item.dtype, np.number) + assert np.isscalar(item) + + +# Test keepdims for: +# - Axis negative, positive, None, and tuple +# - 1D with no NaNs +# - 1D with NaN propagation +# - Zero-sized output +# We're working on making `stats` quieter, but that's not what this test +# is about. For now, we expect all sorts of warnings here due to small samples. +@pytest.mark.filterwarnings('ignore::UserWarning') +@pytest.mark.filterwarnings('ignore::RuntimeWarning') +@pytest.mark.parametrize("nan_policy", ("omit", "propagate")) +@pytest.mark.parametrize(("hypotest", "args", "kwds", "n_samples", "n_outputs", + "paired", "unpacker"), axis_nan_policy_cases) +@pytest.mark.parametrize( + ("sample_shape", "axis_cases"), + (((2, 3, 3, 4), (None, 0, -1, (0, 2), (1, -1), (3, 1, 2, 0))), + ((10, ), (0, -1)), + ((20, 0), (0, 1))) +) +def test_keepdims(hypotest, args, kwds, n_samples, n_outputs, paired, unpacker, + sample_shape, axis_cases, nan_policy): + small_sample_raises = {stats.skewtest, stats.kurtosistest, stats.normaltest, + stats.differential_entropy, stats.epps_singleton_2samp, + stats.shapiro} + if sample_shape == (2, 3, 3, 4) and hypotest in small_sample_raises: + pytest.skip("Sample too small; test raises error.") + if hypotest in {weightedtau_weighted}: + pytest.skip("`rankdata` used in testing doesn't support axis tuple.") + # test if keepdims parameter works correctly + if not unpacker: + def unpacker(res): + return res + rng = np.random.default_rng(0) + data = [rng.random(sample_shape) for _ in range(n_samples)] + nan_data = [sample.copy() for sample in data] + nan_mask = [rng.random(sample_shape) < 0.2 for _ in range(n_samples)] + for sample, mask in zip(nan_data, nan_mask): + sample[mask] = np.nan + for axis in axis_cases: + expected_shape = list(sample_shape) + if axis is None: + expected_shape = np.ones(len(sample_shape)) + else: + if isinstance(axis, int): + expected_shape[axis] = 1 + else: + for ax in axis: + expected_shape[ax] = 1 + expected_shape = tuple(expected_shape) + res = unpacker(hypotest(*data, *args, axis=axis, keepdims=True, + **kwds)) + res_base = unpacker(hypotest(*data, *args, axis=axis, keepdims=False, + **kwds)) + nan_res = unpacker(hypotest(*nan_data, *args, axis=axis, + keepdims=True, nan_policy=nan_policy, + **kwds)) + nan_res_base = unpacker(hypotest(*nan_data, *args, axis=axis, + keepdims=False, + nan_policy=nan_policy, **kwds)) + for r, r_base, rn, rn_base in zip(res, res_base, nan_res, + nan_res_base): + assert r.shape == expected_shape + r = np.squeeze(r, axis=axis) + assert_allclose(r, r_base, atol=1e-16) + assert rn.shape == expected_shape + rn = np.squeeze(rn, axis=axis) + # ideally assert_equal, but `combine_pvalues` failed on 32-bit build + assert_allclose(rn, rn_base, atol=1e-16) + + +@pytest.mark.parametrize(("fun", "nsamp"), + [(stats.kstat, 1), + (stats.kstatvar, 1)]) +def test_hypotest_back_compat_no_axis(fun, nsamp): + m, n = 8, 9 + + rng = np.random.default_rng(0) + x = rng.random((nsamp, m, n)) + res = fun(*x) + res2 = fun(*x, _no_deco=True) + res3 = fun([xi.ravel() for xi in x]) + assert_equal(res, res2) + assert_equal(res, res3) + + +@pytest.mark.parametrize(("axis"), (0, 1, 2)) +def test_axis_nan_policy_decorated_positional_axis(axis): + # Test for correct behavior of function decorated with + # _axis_nan_policy_decorator whether `axis` is provided as positional or + # keyword argument + + shape = (8, 9, 10) + rng = np.random.default_rng(0) + x = rng.random(shape) + y = rng.random(shape) + res1 = stats.mannwhitneyu(x, y, True, 'two-sided', axis) + res2 = stats.mannwhitneyu(x, y, True, 'two-sided', axis=axis) + assert_equal(res1, res2) + + message = "mannwhitneyu() got multiple values for argument 'axis'" + with pytest.raises(TypeError, match=re.escape(message)): + stats.mannwhitneyu(x, y, True, 'two-sided', axis, axis=axis) + + +def test_axis_nan_policy_decorated_positional_args(): + # Test for correct behavior of function decorated with + # _axis_nan_policy_decorator when function accepts *args + + shape = (3, 8, 9, 10) + rng = np.random.default_rng(0) + x = rng.random(shape) + x[0, 0, 0, 0] = np.nan + stats.kruskal(*x) + + message = "kruskal() got an unexpected keyword argument 'samples'" + with pytest.raises(TypeError, match=re.escape(message)): + stats.kruskal(samples=x) + + with pytest.raises(TypeError, match=re.escape(message)): + stats.kruskal(*x, samples=x) + + +def test_axis_nan_policy_decorated_keyword_samples(): + # Test for correct behavior of function decorated with + # _axis_nan_policy_decorator whether samples are provided as positional or + # keyword arguments + + shape = (2, 8, 9, 10) + rng = np.random.default_rng(0) + x = rng.random(shape) + x[0, 0, 0, 0] = np.nan + res1 = stats.mannwhitneyu(*x) + res2 = stats.mannwhitneyu(x=x[0], y=x[1]) + assert_equal(res1, res2) + + message = "mannwhitneyu() got multiple values for argument" + with pytest.raises(TypeError, match=re.escape(message)): + stats.mannwhitneyu(*x, x=x[0], y=x[1]) + + +@pytest.mark.parametrize(("hypotest", "args", "kwds", "n_samples", "n_outputs", + "paired", "unpacker"), axis_nan_policy_cases) +def test_axis_nan_policy_decorated_pickled(hypotest, args, kwds, n_samples, + n_outputs, paired, unpacker): + if "ttest_ci" in hypotest.__name__: + pytest.skip("Can't pickle functions defined within functions.") + + rng = np.random.default_rng(0) + + # Some hypothesis tests return a non-iterable that needs an `unpacker` to + # extract the statistic and p-value. For those that don't: + if not unpacker: + def unpacker(res): + return res + + data = rng.uniform(size=(n_samples, 2, 30)) + pickled_hypotest = pickle.dumps(hypotest) + unpickled_hypotest = pickle.loads(pickled_hypotest) + res1 = unpacker(hypotest(*data, *args, axis=-1, **kwds)) + res2 = unpacker(unpickled_hypotest(*data, *args, axis=-1, **kwds)) + assert_allclose(res1, res2, rtol=1e-12) + + +def test_check_empty_inputs(): + # Test that _check_empty_inputs is doing its job, at least for single- + # sample inputs. (Multi-sample functionality is tested below.) + # If the input sample is not empty, it should return None. + # If the input sample is empty, it should return an array of NaNs or an + # empty array of appropriate shape. np.mean is used as a reference for the + # output because, like the statistics calculated by these functions, + # it works along and "consumes" `axis` but preserves the other axes. + for i in range(5): + for combo in combinations_with_replacement([0, 1, 2], i): + for axis in range(len(combo)): + samples = (np.zeros(combo),) + output = stats._axis_nan_policy._check_empty_inputs(samples, + axis) + if output is not None: + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", "Mean of empty slice", RuntimeWarning) + warnings.filterwarnings( + "ignore", "invalid value encountered", RuntimeWarning) + reference = samples[0].mean(axis=axis) + np.testing.assert_equal(output, reference) + + +def _check_arrays_broadcastable(arrays, axis): + # https://numpy.org/doc/stable/user/basics.broadcasting.html + # "When operating on two arrays, NumPy compares their shapes element-wise. + # It starts with the trailing (i.e. rightmost) dimensions and works its + # way left. + # Two dimensions are compatible when + # 1. they are equal, or + # 2. one of them is 1 + # ... + # Arrays do not need to have the same number of dimensions." + # (Clarification: if the arrays are compatible according to the criteria + # above and an array runs out of dimensions, it is still compatible.) + # Below, we follow the rules above except ignoring `axis` + + n_dims = max([arr.ndim for arr in arrays]) + if axis is not None: + # convert to negative axis + axis = (-n_dims + axis) if axis >= 0 else axis + + for dim in range(1, n_dims+1): # we'll index from -1 to -n_dims, inclusive + if -dim == axis: + continue # ignore lengths along `axis` + + dim_lengths = set() + for arr in arrays: + if dim <= arr.ndim and arr.shape[-dim] != 1: + dim_lengths.add(arr.shape[-dim]) + + if len(dim_lengths) > 1: + return False + return True + + +@pytest.mark.slow +@pytest.mark.parametrize(("hypotest", "args", "kwds", "n_samples", "n_outputs", + "paired", "unpacker"), axis_nan_policy_cases) +def test_empty(hypotest, args, kwds, n_samples, n_outputs, paired, unpacker): + # test for correct output shape when at least one input is empty + if hypotest in {stats.kruskal, stats.friedmanchisquare} and not SCIPY_XSLOW: + pytest.skip("Too slow.") + + if hypotest in override_propagate_funcs: + reason = "Doesn't follow the usual pattern. Tested separately." + pytest.skip(reason=reason) + + if unpacker is None: + unpacker = lambda res: (res[0], res[1]) # noqa: E731 + + def small_data_generator(n_samples, n_dims): + + def small_sample_generator(n_dims): + # return all possible "small" arrays in up to n_dim dimensions + for i in n_dims: + # "small" means with size along dimension either 0 or 1 + for combo in combinations_with_replacement([0, 1, 2], i): + yield np.zeros(combo) + + # yield all possible combinations of small samples + gens = [small_sample_generator(n_dims) for i in range(n_samples)] + yield from product(*gens) + + n_dims = [1, 2, 3] + for samples in small_data_generator(n_samples, n_dims): + + # this test is only for arrays of zero size + if not any(sample.size == 0 for sample in samples): + continue + + max_axis = max(sample.ndim for sample in samples) + + # need to test for all valid values of `axis` parameter, too + for axis in range(-max_axis, max_axis): + + try: + # After broadcasting, all arrays are the same shape, so + # the shape of the output should be the same as a single- + # sample statistic. Use np.mean as a reference. + concat = stats._axis_nan_policy._broadcast_concatenate(samples, axis, + paired=paired) + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", "Mean of empty slice", RuntimeWarning) + warnings.filterwarnings( + "ignore", "invalid value encountered", RuntimeWarning) + expected = np.mean(concat, axis=axis) * np.nan + mask = np.isnan(expected) + expected = [np.asarray(expected.copy()) for i in range(n_outputs)] + + if hypotest in empty_special_case_funcs: + empty_val = hypotest(*([[]]*len(samples)), *args, **kwds) + empty_val = list(unpacker(empty_val)) + for i in range(n_outputs): + expected[i][mask] = empty_val[i] + + if expected[0].size and hypotest not in too_small_special_case_funcs: + message = (too_small_1d_not_omit if max_axis == 1 + else too_small_nd_not_omit) + with pytest.warns(SmallSampleWarning, match=message): + res = hypotest(*samples, *args, axis=axis, **kwds) + else: + with warnings.catch_warnings(): + # f_oneway special case + msg = "all input arrays have length 1" + warnings.filterwarnings("ignore", msg, SmallSampleWarning) + res = hypotest(*samples, *args, axis=axis, **kwds) + res = unpacker(res) + + for i in range(n_outputs): + assert_equal(res[i], expected[i]) + + except ValueError: + # confirm that the arrays truly are not broadcastable + assert not _check_arrays_broadcastable(samples, + None if paired else axis) + + # confirm that _both_ `_broadcast_concatenate` and `hypotest` + # produce this information. + message = "Array shapes are incompatible for broadcasting." + with pytest.raises(ValueError, match=message): + stats._axis_nan_policy._broadcast_concatenate(samples, axis, paired) + with pytest.raises(ValueError, match=message): + hypotest(*samples, *args, axis=axis, **kwds) + + +def paired_non_broadcastable_cases(): + for case in axis_nan_policy_cases: + hypotest, args, kwds, n_samples, n_outputs, paired, unpacker = case + if n_samples == 1: # broadcasting only needed with >1 sample + continue + yield case + + +@pytest.mark.parametrize("axis", [0, 1]) +@pytest.mark.parametrize(("hypotest", "args", "kwds", "n_samples", "n_outputs", + "paired", "unpacker"), + list(paired_non_broadcastable_cases())) +def test_non_broadcastable(hypotest, args, kwds, n_samples, n_outputs, paired, + unpacker, axis): + # test for correct error message when shapes are not broadcastable + rng = np.random.default_rng(91359824598245) + get_samples = True + while get_samples: + samples = [rng.random(size=rng.integers(2, 100, size=2)) + for i in range(n_samples)] + # if samples are broadcastable, try again + get_samples = _check_arrays_broadcastable(samples, axis=axis) + + message = "Array shapes are incompatible for broadcasting." + with pytest.raises(ValueError, match=message): + hypotest(*samples, *args, axis=axis, **kwds) + + if not paired: # there's another test for paired-sample statistics + return + + # Previously, paired sample statistics did not raise an error + # message when the shapes were broadcastable except along `axis` + # https://github.com/scipy/scipy/pull/19578#pullrequestreview-1766857165 + shape = rng.integers(2, 10, size=2) + most_samples = [rng.random(size=shape) for i in range(n_samples-1)] + shape = list(shape) + shape[axis] += 1 + other_sample = rng.random(size=shape) + with pytest.raises(ValueError, match=message): + hypotest(other_sample, *most_samples, *args, axis=axis, **kwds) + +def test_masked_array_2_sentinel_array(): + # prepare arrays + rng = np.random.default_rng(805715284) + A = rng.random((10, 11, 12)) + B = rng.random(12) + mask = A < 0.5 + A = np.ma.masked_array(A, mask) + + # set arbitrary elements to special values + # (these values might have been considered for use as sentinel values) + max_float = np.finfo(np.float64).max + max_float2 = np.nextafter(max_float, -np.inf) + max_float3 = np.nextafter(max_float2, -np.inf) + A[3, 4, 1] = np.nan + A[4, 5, 2] = np.inf + A[5, 6, 3] = max_float + B[8] = np.nan + B[7] = np.inf + B[6] = max_float2 + + # convert masked A to array with sentinel value, don't modify B + out_arrays, sentinel = _masked_arrays_2_sentinel_arrays([A, B]) + A_out, B_out = out_arrays + + # check that good sentinel value was chosen (according to intended logic) + assert (sentinel != max_float) and (sentinel != max_float2) + assert sentinel == max_float3 + + # check that output arrays are as intended + A_reference = A.data + A_reference[A.mask] = sentinel + np.testing.assert_array_equal(A_out, A_reference) + assert B_out is B + + +@skip_xp_invalid_arg +def test_masked_dtype(): + # When _masked_arrays_2_sentinel_arrays was first added, it always + # upcast the arrays to np.float64. After gh16662, check expected promotion + # and that the expected sentinel is found. + + # these are important because the max of the promoted dtype is the first + # candidate to be the sentinel value + max16 = np.iinfo(np.int16).max + max128c = np.finfo(np.complex128).max + + # a is a regular array, b has masked elements, and c has no masked elements + a = np.array([1, 2, max16], dtype=np.int16) + b = np.ma.array([1, 2, 1], dtype=np.int8, mask=[0, 1, 0]) + c = np.ma.array([1, 2, 1], dtype=np.complex128, mask=[0, 0, 0]) + + # check integer masked -> sentinel conversion + out_arrays, sentinel = _masked_arrays_2_sentinel_arrays([a, b]) + a_out, b_out = out_arrays + assert sentinel == max16-1 # not max16 because max16 was in the data + assert b_out.dtype == np.int16 # check expected promotion + assert_allclose(b_out, [b[0], sentinel, b[-1]]) # check sentinel placement + assert a_out is a # not a masked array, so left untouched + assert not isinstance(b_out, np.ma.MaskedArray) # b became regular array + + # similarly with complex + out_arrays, sentinel = _masked_arrays_2_sentinel_arrays([b, c]) + b_out, c_out = out_arrays + assert sentinel == max128c # max128c was not in the data + assert b_out.dtype == np.complex128 # b got promoted + assert_allclose(b_out, [b[0], sentinel, b[-1]]) # check sentinel placement + assert not isinstance(b_out, np.ma.MaskedArray) # b became regular array + assert not isinstance(c_out, np.ma.MaskedArray) # c became regular array + + # Also, check edge case when a sentinel value cannot be found in the data + min8, max8 = np.iinfo(np.int8).min, np.iinfo(np.int8).max + a = np.arange(min8, max8+1, dtype=np.int8) # use all possible values + mask1 = np.zeros_like(a, dtype=bool) + mask0 = np.zeros_like(a, dtype=bool) + + # a masked value can be used as the sentinel + mask1[1] = True + a1 = np.ma.array(a, mask=mask1) + out_arrays, sentinel = _masked_arrays_2_sentinel_arrays([a1]) + assert sentinel == min8+1 + + # unless it's the smallest possible; skipped for simiplicity (see code) + mask0[0] = True + a0 = np.ma.array(a, mask=mask0) + message = "This function replaces masked elements with sentinel..." + with pytest.raises(ValueError, match=message): + _masked_arrays_2_sentinel_arrays([a0]) + + # test that dtype is preserved in functions + a = np.ma.array([1, 2, 3], mask=[0, 1, 0], dtype=np.float32) + assert stats.gmean(a).dtype == np.float32 + + +@skip_xp_invalid_arg +def test_masked_stat_1d(): + # basic test of _axis_nan_policy_factory with 1D masked sample + males = [19, 22, 16, 29, 24] + females = [20, 11, 17, 12] + res = stats.mannwhitneyu(males, females) + + # same result when extra nan is omitted + females2 = [20, 11, 17, np.nan, 12] + res2 = stats.mannwhitneyu(males, females2, nan_policy='omit') + np.testing.assert_array_equal(res2, res) + + # same result when extra element is masked + females3 = [20, 11, 17, 1000, 12] + mask3 = [False, False, False, True, False] + females3 = np.ma.masked_array(females3, mask=mask3) + res3 = stats.mannwhitneyu(males, females3) + np.testing.assert_array_equal(res3, res) + + # same result when extra nan is omitted and additional element is masked + females4 = [20, 11, 17, np.nan, 1000, 12] + mask4 = [False, False, False, False, True, False] + females4 = np.ma.masked_array(females4, mask=mask4) + res4 = stats.mannwhitneyu(males, females4, nan_policy='omit') + np.testing.assert_array_equal(res4, res) + + # same result when extra elements, including nan, are masked + females5 = [20, 11, 17, np.nan, 1000, 12] + mask5 = [False, False, False, True, True, False] + females5 = np.ma.masked_array(females5, mask=mask5) + res5 = stats.mannwhitneyu(males, females5, nan_policy='propagate') + res6 = stats.mannwhitneyu(males, females5, nan_policy='raise') + np.testing.assert_array_equal(res5, res) + np.testing.assert_array_equal(res6, res) + + +@pytest.mark.filterwarnings('ignore:After omitting NaNs...') +@pytest.mark.filterwarnings('ignore:One or more axis-slices of one...') +@skip_xp_invalid_arg +@pytest.mark.parametrize(("axis"), list(range(-3, 3))) +def test_masked_stat_3d(axis): + # basic test of _axis_nan_policy_factory with 3D masked sample + rng = np.random.default_rng(3679428403) + a = rng.random((3, 4, 5)) + b = rng.random((4, 5)) + c = rng.random((4, 1)) + + mask_a = a < 0.1 + mask_c = [False, False, False, True] + a_masked = np.ma.masked_array(a, mask=mask_a) + c_masked = np.ma.masked_array(c, mask=mask_c) + + a_nans = a.copy() + a_nans[mask_a] = np.nan + c_nans = c.copy() + c_nans[mask_c] = np.nan + + res = stats.kruskal(a_nans, b, c_nans, nan_policy='omit', axis=axis) + res2 = stats.kruskal(a_masked, b, c_masked, axis=axis) + np.testing.assert_array_equal(res, res2) + + +@pytest.mark.filterwarnings('ignore:After omitting NaNs...') +@pytest.mark.filterwarnings('ignore:One or more axis-slices of one...') +@skip_xp_invalid_arg +def test_mixed_mask_nan_1(): + # targeted test of _axis_nan_policy_factory with 2D masked sample: + # omitting samples with masks and nan_policy='omit' are equivalent + # also checks paired-sample sentinel value removal + m, n = 3, 20 + axis = -1 + + rng = np.random.RandomState(0) + a = rng.rand(m, n) + b = rng.rand(m, n) + mask_a1 = rng.rand(m, n) < 0.2 + mask_a2 = rng.rand(m, n) < 0.1 + mask_b1 = rng.rand(m, n) < 0.15 + mask_b2 = rng.rand(m, n) < 0.15 + mask_a1[2, :] = True + + a_nans = a.copy() + b_nans = b.copy() + a_nans[mask_a1 | mask_a2] = np.nan + b_nans[mask_b1 | mask_b2] = np.nan + + a_masked1 = np.ma.masked_array(a, mask=mask_a1) + b_masked1 = np.ma.masked_array(b, mask=mask_b1) + a_masked1[mask_a2] = np.nan + b_masked1[mask_b2] = np.nan + + a_masked2 = np.ma.masked_array(a, mask=mask_a2) + b_masked2 = np.ma.masked_array(b, mask=mask_b2) + a_masked2[mask_a1] = np.nan + b_masked2[mask_b1] = np.nan + + a_masked3 = np.ma.masked_array(a, mask=(mask_a1 | mask_a2)) + b_masked3 = np.ma.masked_array(b, mask=(mask_b1 | mask_b2)) + + res = stats.wilcoxon(a_nans, b_nans, nan_policy='omit', axis=axis) + res1 = stats.wilcoxon(a_masked1, b_masked1, nan_policy='omit', axis=axis) + res2 = stats.wilcoxon(a_masked2, b_masked2, nan_policy='omit', axis=axis) + res3 = stats.wilcoxon(a_masked3, b_masked3, nan_policy='raise', axis=axis) + res4 = stats.wilcoxon(a_masked3, b_masked3, + nan_policy='propagate', axis=axis) + + np.testing.assert_array_equal(res1, res) + np.testing.assert_array_equal(res2, res) + np.testing.assert_array_equal(res3, res) + np.testing.assert_array_equal(res4, res) + + +@pytest.mark.filterwarnings('ignore:After omitting NaNs...') +@pytest.mark.filterwarnings('ignore:One or more axis-slices of one...') +@skip_xp_invalid_arg +def test_mixed_mask_nan_2(): + # targeted test of _axis_nan_policy_factory with 2D masked sample: + # check for expected interaction between masks and nans + + # Cases here are + # [mixed nan/mask, all nans, all masked, + # unmasked nan, masked nan, unmasked non-nan] + a = [[1, np.nan, 2], [np.nan, np.nan, np.nan], [1, 2, 3], + [1, np.nan, 3], [1, np.nan, 3], [1, 2, 3]] + mask = [[1, 0, 1], [0, 0, 0], [1, 1, 1], + [0, 0, 0], [0, 1, 0], [0, 0, 0]] + a_masked = np.ma.masked_array(a, mask=mask) + b = [[4, 5, 6]] + ref1 = stats.ranksums([1, 3], [4, 5, 6]) + ref2 = stats.ranksums([1, 2, 3], [4, 5, 6]) + + # nan_policy = 'omit' + # all elements are removed from first three rows + # middle element is removed from fourth and fifth rows + # no elements removed from last row + res = stats.ranksums(a_masked, b, nan_policy='omit', axis=-1) + stat_ref = [np.nan, np.nan, np.nan, + ref1.statistic, ref1.statistic, ref2.statistic] + p_ref = [np.nan, np.nan, np.nan, + ref1.pvalue, ref1.pvalue, ref2.pvalue] + np.testing.assert_array_equal(res.statistic, stat_ref) + np.testing.assert_array_equal(res.pvalue, p_ref) + + # nan_policy = 'propagate' + # nans propagate in first, second, and fourth row + # all elements are removed by mask from third row + # middle element is removed from fifth row + # no elements removed from last row + res = stats.ranksums(a_masked, b, nan_policy='propagate', axis=-1) + stat_ref = [np.nan, np.nan, np.nan, + np.nan, ref1.statistic, ref2.statistic] + p_ref = [np.nan, np.nan, np.nan, + np.nan, ref1.pvalue, ref2.pvalue] + np.testing.assert_array_equal(res.statistic, stat_ref) + np.testing.assert_array_equal(res.pvalue, p_ref) + + +def test_axis_None_vs_tuple(): + # `axis` `None` should be equivalent to tuple with all axes + shape = (3, 8, 9, 10) + rng = np.random.default_rng(0) + x = rng.random(shape) + res = stats.kruskal(*x, axis=None) + res2 = stats.kruskal(*x, axis=(0, 1, 2)) + np.testing.assert_array_equal(res, res2) + + +def test_axis_None_vs_tuple_with_broadcasting(): + # `axis` `None` should be equivalent to tuple with all axes, + # which should be equivalent to raveling the arrays before passing them + rng = np.random.default_rng(0) + x = rng.random((5, 1)) + y = rng.random((1, 5)) + x2, y2 = np.broadcast_arrays(x, y) + + res0 = stats.mannwhitneyu(x.ravel(), y.ravel()) + res1 = stats.mannwhitneyu(x, y, axis=None) + res2 = stats.mannwhitneyu(x, y, axis=(0, 1)) + res3 = stats.mannwhitneyu(x2.ravel(), y2.ravel()) + + assert res1 == res0 + assert res2 == res0 + assert res3 != res0 + + +@pytest.mark.parametrize(("axis"), + list(permutations(range(-3, 3), 2)) + [(-4, 1)]) +def test_other_axis_tuples(axis): + # Check that _axis_nan_policy_factory treats all `axis` tuples as expected + rng = np.random.default_rng(0) + shape_x = (4, 5, 6) + shape_y = (1, 6) + x = rng.random(shape_x) + y = rng.random(shape_y) + axis_original = axis + + # convert axis elements to positive + axis = tuple([(i if i >= 0 else 3 + i) for i in axis]) + axis = sorted(axis) + + if len(set(axis)) != len(axis): + message = "`axis` must contain only distinct elements" + with pytest.raises(AxisError, match=re.escape(message)): + stats.mannwhitneyu(x, y, axis=axis_original) + return + + if axis[0] < 0 or axis[-1] > 2: + message = "`axis` is out of bounds for array of dimension 3" + with pytest.raises(AxisError, match=re.escape(message)): + stats.mannwhitneyu(x, y, axis=axis_original) + return + + res = stats.mannwhitneyu(x, y, axis=axis_original) + + # reference behavior + not_axis = {0, 1, 2} - set(axis) # which axis is not part of `axis` + not_axis = next(iter(not_axis)) # take it out of the set + + x2 = x + shape_y_broadcasted = [1, 1, 6] + shape_y_broadcasted[not_axis] = shape_x[not_axis] + y2 = np.broadcast_to(y, shape_y_broadcasted) + + m = x2.shape[not_axis] + x2 = np.moveaxis(x2, axis, (1, 2)) + y2 = np.moveaxis(y2, axis, (1, 2)) + x2 = np.reshape(x2, (m, -1)) + y2 = np.reshape(y2, (m, -1)) + res2 = stats.mannwhitneyu(x2, y2, axis=1) + + np.testing.assert_array_equal(res, res2) + + +@pytest.mark.filterwarnings('ignore:After omitting NaNs...') +@pytest.mark.filterwarnings('ignore:One or more axis-slices of one...') +@skip_xp_invalid_arg +@pytest.mark.parametrize( + ("weighted_fun_name, unpacker"), + [ + ("gmean", lambda x: x), + ("hmean", lambda x: x), + ("pmean", lambda x: x), + ("combine_pvalues", lambda x: (x.pvalue, x.statistic)), + ], +) +def test_mean_mixed_mask_nan_weights(weighted_fun_name, unpacker): + # targeted test of _axis_nan_policy_factory with 2D masked sample: + # omitting samples with masks and nan_policy='omit' are equivalent + # also checks paired-sample sentinel value removal + + if weighted_fun_name == 'pmean': + def weighted_fun(a, **kwargs): + return stats.pmean(a, p=0.42, **kwargs) + else: + weighted_fun = getattr(stats, weighted_fun_name) + + def func(*args, **kwargs): + return unpacker(weighted_fun(*args, **kwargs)) + + m, n = 3, 20 + axis = -1 + + rng = np.random.default_rng(6541968121) + a = rng.uniform(size=(m, n)) + b = rng.uniform(size=(m, n)) + mask_a1 = rng.uniform(size=(m, n)) < 0.2 + mask_a2 = rng.uniform(size=(m, n)) < 0.1 + mask_b1 = rng.uniform(size=(m, n)) < 0.15 + mask_b2 = rng.uniform(size=(m, n)) < 0.15 + mask_a1[2, :] = True + + a_nans = a.copy() + b_nans = b.copy() + a_nans[mask_a1 | mask_a2] = np.nan + b_nans[mask_b1 | mask_b2] = np.nan + + a_masked1 = np.ma.masked_array(a, mask=mask_a1) + b_masked1 = np.ma.masked_array(b, mask=mask_b1) + a_masked1[mask_a2] = np.nan + b_masked1[mask_b2] = np.nan + + a_masked2 = np.ma.masked_array(a, mask=mask_a2) + b_masked2 = np.ma.masked_array(b, mask=mask_b2) + a_masked2[mask_a1] = np.nan + b_masked2[mask_b1] = np.nan + + a_masked3 = np.ma.masked_array(a, mask=(mask_a1 | mask_a2)) + b_masked3 = np.ma.masked_array(b, mask=(mask_b1 | mask_b2)) + + with warnings.catch_warnings(): + message = 'invalid value encountered' + warnings.filterwarnings("ignore", message, RuntimeWarning) + res = func(a_nans, weights=b_nans, nan_policy="omit", axis=axis) + res1 = func(a_masked1, weights=b_masked1, nan_policy="omit", axis=axis) + res2 = func(a_masked2, weights=b_masked2, nan_policy="omit", axis=axis) + res3 = func(a_masked3, weights=b_masked3, nan_policy="raise", axis=axis) + res4 = func(a_masked3, weights=b_masked3, nan_policy="propagate", axis=axis) + + np.testing.assert_array_equal(res1, res) + np.testing.assert_array_equal(res2, res) + np.testing.assert_array_equal(res3, res) + np.testing.assert_array_equal(res4, res) + + +def test_raise_invalid_args_g17713(): + # other cases are handled in: + # test_axis_nan_policy_decorated_positional_axis - multiple values for arg + # test_axis_nan_policy_decorated_positional_args - unexpected kwd arg + message = "got an unexpected keyword argument" + with pytest.raises(TypeError, match=message): + stats.gmean([1, 2, 3], invalid_arg=True) + + message = " got multiple values for argument" + with pytest.raises(TypeError, match=message): + stats.gmean([1, 2, 3], a=True) + + message = "missing 1 required positional argument" + with pytest.raises(TypeError, match=message): + stats.gmean() + + message = "takes from 1 to 4 positional arguments but 5 were given" + with pytest.raises(TypeError, match=message): + stats.gmean([1, 2, 3], 0, float, [1, 1, 1], 10) + + +@pytest.mark.parametrize('dtype', [np.int16, np.float32, np.complex128]) +def test_array_like_input(dtype): + # Check that `_axis_nan_policy`-decorated functions work with custom + # containers that are coercible to numeric arrays + + class ArrLike: + def __init__(self, x, dtype): + self._x = x + self._dtype = dtype + + def __array__(self, dtype=None, copy=None): + return np.asarray(x, dtype=self._dtype) + + def __iter__(self): + # I don't know of a canonical way to determine whether an object should + # be coerced to a NumPy array or not. Currently, `xp_promote` checks + # whether they are iterable, and if so uses `_asarray` with whatever + # `xp` is. So for this to get coerced, it needs to be iterable. + return iter(self._x) + + x = [1]*2 + [3, 4, 5] + res = stats.mode(ArrLike(x, dtype=dtype)) + assert res.mode == 1 + assert res.count == 2 + + +@make_xp_test_case(stats._axis_nan_policy._broadcast_concatenate) +def test__broadcast_concatenate(xp): + # test that _broadcast_concatenate properly broadcasts arrays along all + # axes except `axis`, then concatenates along axis + rng = np.random.default_rng(7544340069) + a = rng.random((5, 4, 4, 3, 1, 6)) + b = rng.random((4, 1, 8, 2, 6)) + arrays = (xp.asarray(a), xp.asarray(b)) + c = stats._axis_nan_policy._broadcast_concatenate(arrays, axis=-3, xp=xp) + # broadcast manually as an independent check + a = np.tile(a, (1, 1, 1, 1, 2, 1)) + b = np.tile(b[None, ...], (5, 1, 4, 1, 1, 1)) + for index in product(*(range(i) for i in c.shape)): + i, j, k, l, m, n = index + if l < a.shape[-3]: + assert a[i, j, k, l, m, n] == c[i, j, k, l, m, n] + else: + assert b[i, j, k, l - a.shape[-3], m, n] == c[i, j, k, l, m, n] diff --git a/v12x/pylib/scipy/stats/tests/test_binned_statistic.py b/v12x/pylib/scipy/stats/tests/test_binned_statistic.py new file mode 100644 index 0000000000000000000000000000000000000000..a1fc499123f8f0eed07a767f2d9e361a84452bcb --- /dev/null +++ b/v12x/pylib/scipy/stats/tests/test_binned_statistic.py @@ -0,0 +1,602 @@ +import numpy as np +from numpy.testing import assert_allclose +import pytest +from pytest import raises as assert_raises +from scipy.stats import (binned_statistic, binned_statistic_2d, + binned_statistic_dd) +from scipy._lib._util import check_random_state + +from .common_tests import check_named_results + + +class TestBinnedStatistic: + + @classmethod + def setup_class(cls): + rng = check_random_state(9865) + cls.x = rng.uniform(size=100) + cls.y = rng.uniform(size=100) + cls.v = rng.uniform(size=100) + cls.X = rng.uniform(size=(100, 3)) + cls.w = rng.uniform(size=100) + cls.u = rng.uniform(size=100) + 1e6 + + def test_1d_count(self): + x = self.x + v = self.v + + count1, edges1, bc = binned_statistic(x, v, 'count', bins=10) + count2, edges2 = np.histogram(x, bins=10) + + assert_allclose(count1, count2) + assert_allclose(edges1, edges2) + + def test_gh5927(self): + # smoke test for gh5927 - binned_statistic was using `is` for string + # comparison + x = self.x + v = self.v + statistics = ['mean', 'median', 'count', 'sum'] + for statistic in statistics: + binned_statistic(x, v, statistic, bins=10) + + def test_big_number_std(self): + # tests for numerical stability of std calculation + # see issue gh-10126 for more + x = self.x + u = self.u + stat1, edges1, bc = binned_statistic(x, u, 'std', bins=10) + stat2, edges2, bc = binned_statistic(x, u, np.std, bins=10) + + assert_allclose(stat1, stat2) + + def test_empty_bins_std(self): + # tests that std returns gives nan for empty bins + x = self.x + u = self.u + print(binned_statistic(x, u, 'count', bins=1000)) + stat1, edges1, bc = binned_statistic(x, u, 'std', bins=1000) + stat2, edges2, bc = binned_statistic(x, u, np.std, bins=1000) + + assert_allclose(stat1, stat2) + + def test_1d_result_attributes(self): + x = self.x + v = self.v + + res = binned_statistic(x, v, 'count', bins=10) + attributes = ('statistic', 'bin_edges', 'binnumber') + check_named_results(res, attributes) + + def test_1d_sum(self): + x = self.x + v = self.v + + sum1, edges1, bc = binned_statistic(x, v, 'sum', bins=10) + sum2, edges2 = np.histogram(x, bins=10, weights=v) + + assert_allclose(sum1, sum2) + assert_allclose(edges1, edges2) + + def test_1d_mean(self): + x = self.x + v = self.v + + stat1, edges1, bc = binned_statistic(x, v, 'mean', bins=10) + stat2, edges2, bc = binned_statistic(x, v, np.mean, bins=10) + + assert_allclose(stat1, stat2) + assert_allclose(edges1, edges2) + + def test_1d_std(self): + x = self.x + v = self.v + + stat1, edges1, bc = binned_statistic(x, v, 'std', bins=10) + stat2, edges2, bc = binned_statistic(x, v, np.std, bins=10) + + assert_allclose(stat1, stat2) + assert_allclose(edges1, edges2) + + def test_1d_min(self): + x = self.x + v = self.v + + stat1, edges1, bc = binned_statistic(x, v, 'min', bins=10) + stat2, edges2, bc = binned_statistic(x, v, np.min, bins=10) + + assert_allclose(stat1, stat2) + assert_allclose(edges1, edges2) + + def test_1d_max(self): + x = self.x + v = self.v + + stat1, edges1, bc = binned_statistic(x, v, 'max', bins=10) + stat2, edges2, bc = binned_statistic(x, v, np.max, bins=10) + + assert_allclose(stat1, stat2) + assert_allclose(edges1, edges2) + + def test_1d_median(self): + x = self.x + v = self.v + + stat1, edges1, bc = binned_statistic(x, v, 'median', bins=10) + stat2, edges2, bc = binned_statistic(x, v, np.median, bins=10) + + assert_allclose(stat1, stat2) + assert_allclose(edges1, edges2) + + def test_1d_bincode(self): + x = self.x[:20] + v = self.v[:20] + + count1, edges1, bc = binned_statistic(x, v, 'count', bins=3) + bc2 = np.array([3, 2, 1, 3, 2, 3, 3, 3, 3, 1, 1, 3, 3, 1, 2, 3, 1, + 1, 2, 1]) + + bcount = [(bc == i).sum() for i in np.unique(bc)] + + assert_allclose(bc, bc2) + assert_allclose(bcount, count1) + + def test_1d_range_keyword(self): + # Regression test for gh-3063, range can be (min, max) or [(min, max)] + rng = np.random.default_rng(6823616729) + x = np.arange(30) + data = rng.random(30) + + mean, bins, _ = binned_statistic(x[:15], data[:15]) + mean_range, bins_range, _ = binned_statistic(x, data, range=[(0, 14)]) + mean_range2, bins_range2, _ = binned_statistic(x, data, range=(0, 14)) + + assert_allclose(mean, mean_range) + assert_allclose(bins, bins_range) + assert_allclose(mean, mean_range2) + assert_allclose(bins, bins_range2) + + def test_1d_multi_values(self): + x = self.x + v = self.v + w = self.w + + stat1v, edges1v, bc1v = binned_statistic(x, v, 'mean', bins=10) + stat1w, edges1w, bc1w = binned_statistic(x, w, 'mean', bins=10) + stat2, edges2, bc2 = binned_statistic(x, [v, w], 'mean', bins=10) + + assert_allclose(stat2[0], stat1v) + assert_allclose(stat2[1], stat1w) + assert_allclose(edges1v, edges2) + assert_allclose(bc1v, bc2) + + def test_2d_count(self): + x = self.x + y = self.y + v = self.v + + count1, binx1, biny1, bc = binned_statistic_2d( + x, y, v, 'count', bins=5) + count2, binx2, biny2 = np.histogram2d(x, y, bins=5) + + assert_allclose(count1, count2) + assert_allclose(binx1, binx2) + assert_allclose(biny1, biny2) + + def test_2d_result_attributes(self): + x = self.x + y = self.y + v = self.v + + res = binned_statistic_2d(x, y, v, 'count', bins=5) + attributes = ('statistic', 'x_edge', 'y_edge', 'binnumber') + check_named_results(res, attributes) + + def test_2d_sum(self): + x = self.x + y = self.y + v = self.v + + sum1, binx1, biny1, bc = binned_statistic_2d(x, y, v, 'sum', bins=5) + sum2, binx2, biny2 = np.histogram2d(x, y, bins=5, weights=v) + + assert_allclose(sum1, sum2) + assert_allclose(binx1, binx2) + assert_allclose(biny1, biny2) + + def test_2d_mean(self): + x = self.x + y = self.y + v = self.v + + stat1, binx1, biny1, bc = binned_statistic_2d(x, y, v, 'mean', bins=5) + stat2, binx2, biny2, bc = binned_statistic_2d(x, y, v, np.mean, bins=5) + + assert_allclose(stat1, stat2) + assert_allclose(binx1, binx2) + assert_allclose(biny1, biny2) + + def test_2d_mean_unicode(self): + x = self.x + y = self.y + v = self.v + stat1, binx1, biny1, bc = binned_statistic_2d( + x, y, v, 'mean', bins=5) + stat2, binx2, biny2, bc = binned_statistic_2d(x, y, v, np.mean, bins=5) + assert_allclose(stat1, stat2) + assert_allclose(binx1, binx2) + assert_allclose(biny1, biny2) + + def test_2d_std(self): + x = self.x + y = self.y + v = self.v + + stat1, binx1, biny1, bc = binned_statistic_2d(x, y, v, 'std', bins=5) + stat2, binx2, biny2, bc = binned_statistic_2d(x, y, v, np.std, bins=5) + + assert_allclose(stat1, stat2) + assert_allclose(binx1, binx2) + assert_allclose(biny1, biny2) + + def test_2d_min(self): + x = self.x + y = self.y + v = self.v + + stat1, binx1, biny1, bc = binned_statistic_2d(x, y, v, 'min', bins=5) + stat2, binx2, biny2, bc = binned_statistic_2d(x, y, v, np.min, bins=5) + + assert_allclose(stat1, stat2) + assert_allclose(binx1, binx2) + assert_allclose(biny1, biny2) + + def test_2d_max(self): + x = self.x + y = self.y + v = self.v + + stat1, binx1, biny1, bc = binned_statistic_2d(x, y, v, 'max', bins=5) + stat2, binx2, biny2, bc = binned_statistic_2d(x, y, v, np.max, bins=5) + + assert_allclose(stat1, stat2) + assert_allclose(binx1, binx2) + assert_allclose(biny1, biny2) + + def test_2d_median(self): + x = self.x + y = self.y + v = self.v + + stat1, binx1, biny1, bc = binned_statistic_2d( + x, y, v, 'median', bins=5) + stat2, binx2, biny2, bc = binned_statistic_2d( + x, y, v, np.median, bins=5) + + assert_allclose(stat1, stat2) + assert_allclose(binx1, binx2) + assert_allclose(biny1, biny2) + + def test_2d_bincode(self): + x = self.x[:20] + y = self.y[:20] + v = self.v[:20] + + count1, binx1, biny1, bc = binned_statistic_2d( + x, y, v, 'count', bins=3) + bc2 = np.array([17, 11, 6, 16, 11, 17, 18, 17, 17, 7, 6, 18, 16, + 6, 11, 16, 6, 6, 11, 8]) + + bcount = [(bc == i).sum() for i in np.unique(bc)] + + assert_allclose(bc, bc2) + count1adj = count1[count1.nonzero()] + assert_allclose(bcount, count1adj) + + def test_2d_multi_values(self): + x = self.x + y = self.y + v = self.v + w = self.w + + stat1v, binx1v, biny1v, bc1v = binned_statistic_2d( + x, y, v, 'mean', bins=8) + stat1w, binx1w, biny1w, bc1w = binned_statistic_2d( + x, y, w, 'mean', bins=8) + stat2, binx2, biny2, bc2 = binned_statistic_2d( + x, y, [v, w], 'mean', bins=8) + + assert_allclose(stat2[0], stat1v) + assert_allclose(stat2[1], stat1w) + assert_allclose(binx1v, binx2) + assert_allclose(biny1w, biny2) + assert_allclose(bc1v, bc2) + + def test_2d_binnumbers_unraveled(self): + x = self.x + y = self.y + v = self.v + + stat, edgesx, bcx = binned_statistic(x, v, 'mean', bins=20) + stat, edgesy, bcy = binned_statistic(y, v, 'mean', bins=10) + + stat2, edgesx2, edgesy2, bc2 = binned_statistic_2d( + x, y, v, 'mean', bins=(20, 10), expand_binnumbers=True) + + bcx3 = np.searchsorted(edgesx, x, side='right') + bcy3 = np.searchsorted(edgesy, y, side='right') + + # `numpy.searchsorted` is non-inclusive on right-edge, compensate + bcx3[x == x.max()] -= 1 + bcy3[y == y.max()] -= 1 + + assert_allclose(bcx, bc2[0]) + assert_allclose(bcy, bc2[1]) + assert_allclose(bcx3, bc2[0]) + assert_allclose(bcy3, bc2[1]) + + def test_dd_count(self): + X = self.X + v = self.v + + count1, edges1, bc = binned_statistic_dd(X, v, 'count', bins=3) + count2, edges2 = np.histogramdd(X, bins=3) + + assert_allclose(count1, count2) + assert_allclose(edges1, edges2) + + def test_dd_result_attributes(self): + X = self.X + v = self.v + + res = binned_statistic_dd(X, v, 'count', bins=3) + attributes = ('statistic', 'bin_edges', 'binnumber') + check_named_results(res, attributes) + + def test_dd_sum(self): + X = self.X + v = self.v + + sum1, edges1, bc = binned_statistic_dd(X, v, 'sum', bins=3) + sum2, edges2 = np.histogramdd(X, bins=3, weights=v) + sum3, edges3, bc = binned_statistic_dd(X, v, np.sum, bins=3) + + assert_allclose(sum1, sum2) + assert_allclose(edges1, edges2) + assert_allclose(sum1, sum3) + assert_allclose(edges1, edges3) + + def test_dd_mean(self): + X = self.X + v = self.v + + stat1, edges1, bc = binned_statistic_dd(X, v, 'mean', bins=3) + stat2, edges2, bc = binned_statistic_dd(X, v, np.mean, bins=3) + + assert_allclose(stat1, stat2) + assert_allclose(edges1, edges2) + + def test_dd_std(self): + X = self.X + v = self.v + + stat1, edges1, bc = binned_statistic_dd(X, v, 'std', bins=3) + stat2, edges2, bc = binned_statistic_dd(X, v, np.std, bins=3) + + assert_allclose(stat1, stat2) + assert_allclose(edges1, edges2) + + def test_dd_min(self): + X = self.X + v = self.v + + stat1, edges1, bc = binned_statistic_dd(X, v, 'min', bins=3) + stat2, edges2, bc = binned_statistic_dd(X, v, np.min, bins=3) + + assert_allclose(stat1, stat2) + assert_allclose(edges1, edges2) + + def test_dd_max(self): + X = self.X + v = self.v + + stat1, edges1, bc = binned_statistic_dd(X, v, 'max', bins=3) + stat2, edges2, bc = binned_statistic_dd(X, v, np.max, bins=3) + + assert_allclose(stat1, stat2) + assert_allclose(edges1, edges2) + + def test_dd_median(self): + X = self.X + v = self.v + + stat1, edges1, bc = binned_statistic_dd(X, v, 'median', bins=3) + stat2, edges2, bc = binned_statistic_dd(X, v, np.median, bins=3) + + assert_allclose(stat1, stat2) + assert_allclose(edges1, edges2) + + def test_dd_bincode(self): + X = self.X[:20] + v = self.v[:20] + + count1, edges1, bc = binned_statistic_dd(X, v, 'count', bins=3) + bc2 = np.array([63, 33, 86, 83, 88, 67, 57, 33, 42, 41, 82, 83, 92, + 32, 36, 91, 43, 87, 81, 81]) + + bcount = [(bc == i).sum() for i in np.unique(bc)] + + assert_allclose(bc, bc2) + count1adj = count1[count1.nonzero()] + assert_allclose(bcount, count1adj) + + def test_dd_multi_values(self): + X = self.X + v = self.v + w = self.w + + for stat in ["count", "sum", "mean", "std", "min", "max", "median", + np.std]: + stat1v, edges1v, bc1v = binned_statistic_dd(X, v, stat, bins=8) + stat1w, edges1w, bc1w = binned_statistic_dd(X, w, stat, bins=8) + stat2, edges2, bc2 = binned_statistic_dd(X, [v, w], stat, bins=8) + assert_allclose(stat2[0], stat1v) + assert_allclose(stat2[1], stat1w) + assert_allclose(edges1v, edges2) + assert_allclose(edges1w, edges2) + assert_allclose(bc1v, bc2) + + def test_dd_binnumbers_unraveled(self): + X = self.X + v = self.v + + stat, edgesx, bcx = binned_statistic(X[:, 0], v, 'mean', bins=15) + stat, edgesy, bcy = binned_statistic(X[:, 1], v, 'mean', bins=20) + stat, edgesz, bcz = binned_statistic(X[:, 2], v, 'mean', bins=10) + + stat2, edges2, bc2 = binned_statistic_dd( + X, v, 'mean', bins=(15, 20, 10), expand_binnumbers=True) + + assert_allclose(bcx, bc2[0]) + assert_allclose(bcy, bc2[1]) + assert_allclose(bcz, bc2[2]) + + def test_dd_binned_statistic_result(self): + # NOTE: tests the reuse of bin_edges from previous call + rng = np.random.default_rng(8111360615) + x = rng.random((10000, 3)) + v = rng.random(10000) + bins = np.linspace(0, 1, 10) + bins = (bins, bins, bins) + + result = binned_statistic_dd(x, v, 'mean', bins=bins) + stat = result.statistic + + result = binned_statistic_dd(x, v, 'mean', + binned_statistic_result=result) + stat2 = result.statistic + + assert_allclose(stat, stat2) + + def test_dd_zero_dedges(self): + rng = np.random.default_rng(1132724173) + x = rng.random((10000, 3)) + v = rng.random(10000) + bins = np.linspace(0, 1, 10) + bins = np.append(bins, 1) + bins = (bins, bins, bins) + with assert_raises(ValueError, match='difference is numerically 0'): + binned_statistic_dd(x, v, 'mean', bins=bins) + + def test_dd_range_errors(self): + # Test that descriptive exceptions are raised as appropriate for bad + # values of the `range` argument. (See gh-12996) + with assert_raises(ValueError, + match='In range, start must be <= stop'): + binned_statistic_dd([self.y], self.v, + range=[[1, 0]]) + with assert_raises( + ValueError, + match='In dimension 1 of range, start must be <= stop'): + binned_statistic_dd([self.x, self.y], self.v, + range=[[1, 0], [0, 1]]) + with assert_raises( + ValueError, + match='In dimension 2 of range, start must be <= stop'): + binned_statistic_dd([self.x, self.y], self.v, + range=[[0, 1], [1, 0]]) + with assert_raises( + ValueError, + match='range given for 1 dimensions; 2 required'): + binned_statistic_dd([self.x, self.y], self.v, + range=[[0, 1]]) + + def test_binned_statistic_float32(self): + X = np.array([0, 0.42358226], dtype=np.float32) + stat, _, _ = binned_statistic(X, None, 'count', bins=5) + assert_allclose(stat, np.array([1, 0, 0, 0, 1], dtype=np.float64)) + + def test_gh14332(self): + # Test the wrong output when the `sample` is close to bin edge + x = [] + size = 20 + for i in range(size): + x += [1-0.1**i] + + bins = np.linspace(0,1,11) + sum1, edges1, bc = binned_statistic_dd(x, np.ones(len(x)), + bins=[bins], statistic='sum') + sum2, edges2 = np.histogram(x, bins=bins) + + assert_allclose(sum1, sum2) + assert_allclose(edges1[0], edges2) + + @pytest.mark.parametrize("dtype", [np.float64, np.complex128]) + @pytest.mark.parametrize("statistic", [np.mean, np.median, np.sum, np.std, + np.min, np.max, 'count', + lambda x: (x**2).sum(), + lambda x: (x**2).sum() * 1j]) + def test_dd_all(self, dtype, statistic): + def ref_statistic(x): + return len(x) if statistic == 'count' else statistic(x) + + rng = np.random.default_rng(3704743126639371) + n = 10 + x = rng.random(size=n) + i = x >= 0.5 + v = rng.random(size=n) + if dtype is np.complex128: + v = v + rng.random(size=n)*1j + + stat, _, _ = binned_statistic_dd(x, v, statistic, bins=2) + ref = np.array([ref_statistic(v[~i]), ref_statistic(v[i])]) + assert_allclose(stat, ref) + assert stat.dtype == np.result_type(ref.dtype, np.float64) + + @pytest.mark.parametrize('bins', [10, (10, 10, 10), [np.arange(10)]*3]) + def test_nan_gh17154(self, bins): + # gh-17154 reported that `binned_statistic_dd` raised an error when `samples` + # contains NaNs and argued that this was an inappropriate default. gh-23751 + # decided that it was indeed an appropriate choice, and went further, ensuring + # that the same error would be raised regardless of how `bins` is specified. + + # Generate a samples (with NaNs) and values + rng = np.random.default_rng(529847529874529829) + samples = rng.standard_normal(size=(100, 3)) + values = rng.standard_normal(100) + i = rng.random(samples.shape) < 0.05 + samples[i] = np.nan + + message = "The sample on which `values` is to be binned..." + with pytest.raises(ValueError, match=message): + binned_statistic_dd(samples, values) + + def test_inf_gh17154(self): + # While addressing gh-17154, it seemed that `inf`s should not be ignored + # entirely; rather, values corresponding with `inf`s can always be in outlier + # bins. + + # Generate a samples (with infs) and values + rng = np.random.default_rng(529847529874529829) + samples = rng.standard_normal(size=(100, 3)) + values = rng.standard_normal(100) + i_pinf = rng.random(samples.shape) < 0.05 + samples[i_pinf] = np.inf + i_ninf = rng.random(samples.shape) < 0.05 + samples[i_ninf] = -np.inf + + # Compute results + res = binned_statistic_dd(samples, values) + + # Compute reference values (with bins computed ignoring infs) + i_inf = i_pinf | i_ninf + temp = samples.copy() + temp[i_inf] = 0 + low = np.min(temp, axis=0) + high = np.max(temp, axis=0) + ref = binned_statistic_dd(samples, values, range=list(zip(low, high))) + + # Compare results against reference values + np.testing.assert_allclose(res[0], ref[0]) + np.testing.assert_allclose(res[1], ref[1]) + np.testing.assert_allclose(res[2], ref[2]) diff --git a/v12x/pylib/scipy/stats/tests/test_censored_data.py b/v12x/pylib/scipy/stats/tests/test_censored_data.py new file mode 100644 index 0000000000000000000000000000000000000000..ae71dcfaccf899645051287f8944131ec48a1eee --- /dev/null +++ b/v12x/pylib/scipy/stats/tests/test_censored_data.py @@ -0,0 +1,152 @@ +# Tests for the CensoredData class. + +import pytest +import numpy as np +from numpy.testing import assert_equal, assert_array_equal +from scipy.stats import CensoredData + + +class TestCensoredData: + + def test_basic(self): + uncensored = [1] + left = [0] + right = [2, 5] + interval = [[2, 3]] + data = CensoredData(uncensored, left=left, right=right, + interval=interval) + assert_equal(data._uncensored, uncensored) + assert_equal(data._left, left) + assert_equal(data._right, right) + assert_equal(data._interval, interval) + + udata = data._uncensor() + assert_equal(udata, np.concatenate((uncensored, left, right, + np.mean(interval, axis=1)))) + + def test_right_censored(self): + x = np.array([0, 3, 2.5]) + is_censored = np.array([0, 1, 0], dtype=bool) + data = CensoredData.right_censored(x, is_censored) + assert_equal(data._uncensored, x[~is_censored]) + assert_equal(data._right, x[is_censored]) + assert_equal(data._left, []) + assert_equal(data._interval, np.empty((0, 2))) + + def test_left_censored(self): + x = np.array([0, 3, 2.5]) + is_censored = np.array([0, 1, 0], dtype=bool) + data = CensoredData.left_censored(x, is_censored) + assert_equal(data._uncensored, x[~is_censored]) + assert_equal(data._left, x[is_censored]) + assert_equal(data._right, []) + assert_equal(data._interval, np.empty((0, 2))) + + def test_interval_censored_basic(self): + a = [0.5, 2.0, 3.0, 5.5] + b = [1.0, 2.5, 3.5, 7.0] + data = CensoredData.interval_censored(low=a, high=b) + assert_array_equal(data._interval, np.array(list(zip(a, b)))) + assert data._uncensored.shape == (0,) + assert data._left.shape == (0,) + assert data._right.shape == (0,) + + def test_interval_censored_mixed(self): + # This is actually a mix of uncensored, left-censored, right-censored + # and interval-censored data. Check that when the `interval_censored` + # class method is used, the data is correctly separated into the + # appropriate arrays. + a = [0.5, -np.inf, -13.0, 2.0, 1.0, 10.0, -1.0] + b = [0.5, 2500.0, np.inf, 3.0, 1.0, 11.0, np.inf] + data = CensoredData.interval_censored(low=a, high=b) + assert_array_equal(data._interval, [[2.0, 3.0], [10.0, 11.0]]) + assert_array_equal(data._uncensored, [0.5, 1.0]) + assert_array_equal(data._left, [2500.0]) + assert_array_equal(data._right, [-13.0, -1.0]) + + def test_interval_to_other_types(self): + # The interval parameter can represent uncensored and + # left- or right-censored data. Test the conversion of such + # an example to the canonical form in which the different + # types have been split into the separate arrays. + interval = np.array([[0, 1], # interval-censored + [2, 2], # not censored + [3, 3], # not censored + [9, np.inf], # right-censored + [8, np.inf], # right-censored + [-np.inf, 0], # left-censored + [1, 2]]) # interval-censored + data = CensoredData(interval=interval) + assert_equal(data._uncensored, [2, 3]) + assert_equal(data._left, [0]) + assert_equal(data._right, [9, 8]) + assert_equal(data._interval, [[0, 1], [1, 2]]) + + def test_empty_arrays(self): + data = CensoredData(uncensored=[], left=[], right=[], interval=[]) + assert data._uncensored.shape == (0,) + assert data._left.shape == (0,) + assert data._right.shape == (0,) + assert data._interval.shape == (0, 2) + assert len(data) == 0 + + def test_invalid_constructor_args(self): + with pytest.raises(ValueError, match='must be a one-dimensional'): + CensoredData(uncensored=[[1, 2, 3]]) + with pytest.raises(ValueError, match='must be a one-dimensional'): + CensoredData(left=[[1, 2, 3]]) + with pytest.raises(ValueError, match='must be a one-dimensional'): + CensoredData(right=[[1, 2, 3]]) + with pytest.raises(ValueError, match='must be a two-dimensional'): + CensoredData(interval=[[1, 2, 3]]) + + with pytest.raises(ValueError, match='must not contain nan'): + CensoredData(uncensored=[1, np.nan, 2]) + with pytest.raises(ValueError, match='must not contain nan'): + CensoredData(left=[1, np.nan, 2]) + with pytest.raises(ValueError, match='must not contain nan'): + CensoredData(right=[1, np.nan, 2]) + with pytest.raises(ValueError, match='must not contain nan'): + CensoredData(interval=[[1, np.nan], [2, 3]]) + + with pytest.raises(ValueError, + match='both values must not be infinite'): + CensoredData(interval=[[1, 3], [2, 9], [np.inf, np.inf]]) + + with pytest.raises(ValueError, + match='left value must not exceed the right'): + CensoredData(interval=[[1, 0], [2, 2]]) + + @pytest.mark.parametrize('func', [CensoredData.left_censored, + CensoredData.right_censored]) + def test_invalid_left_right_censored_args(self, func): + with pytest.raises(ValueError, + match='`x` must be one-dimensional'): + func([[1, 2, 3]], [0, 1, 1]) + with pytest.raises(ValueError, + match='`censored` must be one-dimensional'): + func([1, 2, 3], [[0, 1, 1]]) + with pytest.raises(ValueError, match='`x` must not contain'): + func([1, 2, np.nan], [0, 1, 1]) + with pytest.raises(ValueError, match='must have the same length'): + func([1, 2, 3], [0, 0, 1, 1]) + + def test_invalid_censored_args(self): + with pytest.raises(ValueError, + match='`low` must be a one-dimensional'): + CensoredData.interval_censored(low=[[3]], high=[4, 5]) + with pytest.raises(ValueError, + match='`high` must be a one-dimensional'): + CensoredData.interval_censored(low=[3], high=[[4, 5]]) + with pytest.raises(ValueError, match='`low` must not contain'): + CensoredData.interval_censored([1, 2, np.nan], [0, 1, 1]) + with pytest.raises(ValueError, match='must have the same length'): + CensoredData.interval_censored([1, 2, 3], [0, 0, 1, 1]) + + def test_count_censored(self): + x = [1, 2, 3] + # data1 has no censored data. + data1 = CensoredData(x) + assert data1.num_censored() == 0 + data2 = CensoredData(uncensored=[2.5], left=[10], interval=[[0, 1]]) + assert data2.num_censored() == 2 diff --git a/v12x/pylib/scipy/stats/tests/test_contingency.py b/v12x/pylib/scipy/stats/tests/test_contingency.py new file mode 100644 index 0000000000000000000000000000000000000000..aa83e05d7da6b3d21d5515a332344c2347f0d7eb --- /dev/null +++ b/v12x/pylib/scipy/stats/tests/test_contingency.py @@ -0,0 +1,294 @@ +import numpy as np +from numpy.testing import (assert_equal, assert_array_equal, + assert_array_almost_equal, assert_approx_equal, + assert_allclose) +import pytest +from pytest import raises as assert_raises +from scipy import stats +from scipy.special import xlogy +from scipy.stats.contingency import (margins, expected_freq, + chi2_contingency, association) + + +def test_margins(): + a = np.array([1]) + m = margins(a) + assert_equal(len(m), 1) + m0 = m[0] + assert_array_equal(m0, np.array([1])) + + a = np.array([[1]]) + m0, m1 = margins(a) + expected0 = np.array([[1]]) + expected1 = np.array([[1]]) + assert_array_equal(m0, expected0) + assert_array_equal(m1, expected1) + + a = np.arange(12).reshape(2, 6) + m0, m1 = margins(a) + expected0 = np.array([[15], [51]]) + expected1 = np.array([[6, 8, 10, 12, 14, 16]]) + assert_array_equal(m0, expected0) + assert_array_equal(m1, expected1) + + a = np.arange(24).reshape(2, 3, 4) + m0, m1, m2 = margins(a) + expected0 = np.array([[[66]], [[210]]]) + expected1 = np.array([[[60], [92], [124]]]) + expected2 = np.array([[[60, 66, 72, 78]]]) + assert_array_equal(m0, expected0) + assert_array_equal(m1, expected1) + assert_array_equal(m2, expected2) + + +def test_expected_freq(): + assert_array_equal(expected_freq([1]), np.array([1.0])) + + observed = np.array([[[2, 0], [0, 2]], [[0, 2], [2, 0]], [[1, 1], [1, 1]]]) + e = expected_freq(observed) + assert_array_equal(e, np.ones_like(observed)) + + observed = np.array([[10, 10, 20], [20, 20, 20]]) + e = expected_freq(observed) + correct = np.array([[12., 12., 16.], [18., 18., 24.]]) + assert_array_almost_equal(e, correct) + + +class TestChi2Contingency: + def test_chi2_contingency_trivial(self): + # Some very simple tests for chi2_contingency. + + # A trivial case + obs = np.array([[1, 2], [1, 2]]) + chi2, p, dof, expected = chi2_contingency(obs, correction=False) + assert_equal(chi2, 0.0) + assert_equal(p, 1.0) + assert_equal(dof, 1) + assert_array_equal(obs, expected) + + # A *really* trivial case: 1-D data. + obs = np.array([1, 2, 3]) + chi2, p, dof, expected = chi2_contingency(obs, correction=False) + assert_equal(chi2, 0.0) + assert_equal(p, 1.0) + assert_equal(dof, 0) + assert_array_equal(obs, expected) + + def test_chi2_contingency_R(self): + # Some test cases that were computed independently, using R. + + # Rcode = \ + # """ + # # Data vector. + # data <- c( + # 12, 34, 23, 4, 47, 11, + # 35, 31, 11, 34, 10, 18, + # 12, 32, 9, 18, 13, 19, + # 12, 12, 14, 9, 33, 25 + # ) + # + # # Create factor tags:r=rows, c=columns, t=tiers + # r <- factor(gl(4, 2*3, 2*3*4, labels=c("r1", "r2", "r3", "r4"))) + # c <- factor(gl(3, 1, 2*3*4, labels=c("c1", "c2", "c3"))) + # t <- factor(gl(2, 3, 2*3*4, labels=c("t1", "t2"))) + # + # # 3-way Chi squared test of independence + # s = summary(xtabs(data~r+c+t)) + # print(s) + # """ + # Routput = \ + # """ + # Call: xtabs(formula = data ~ r + c + t) + # Number of cases in table: 478 + # Number of factors: 3 + # Test for independence of all factors: + # Chisq = 102.17, df = 17, p-value = 3.514e-14 + # """ + obs = np.array( + [[[12, 34, 23], + [35, 31, 11], + [12, 32, 9], + [12, 12, 14]], + [[4, 47, 11], + [34, 10, 18], + [18, 13, 19], + [9, 33, 25]]]) + chi2, p, dof, expected = chi2_contingency(obs) + assert_approx_equal(chi2, 102.17, significant=5) + assert_approx_equal(p, 3.514e-14, significant=4) + assert_equal(dof, 17) + + # Rcode = \ + # """ + # # Data vector. + # data <- c( + # # + # 12, 17, + # 11, 16, + # # + # 11, 12, + # 15, 16, + # # + # 23, 15, + # 30, 22, + # # + # 14, 17, + # 15, 16 + # ) + # + # # Create factor tags:r=rows, c=columns, d=depths(?), t=tiers + # r <- factor(gl(2, 2, 2*2*2*2, labels=c("r1", "r2"))) + # c <- factor(gl(2, 1, 2*2*2*2, labels=c("c1", "c2"))) + # d <- factor(gl(2, 4, 2*2*2*2, labels=c("d1", "d2"))) + # t <- factor(gl(2, 8, 2*2*2*2, labels=c("t1", "t2"))) + # + # # 4-way Chi squared test of independence + # s = summary(xtabs(data~r+c+d+t)) + # print(s) + # """ + # Routput = \ + # """ + # Call: xtabs(formula = data ~ r + c + d + t) + # Number of cases in table: 262 + # Number of factors: 4 + # Test for independence of all factors: + # Chisq = 8.758, df = 11, p-value = 0.6442 + # """ + obs = np.array( + [[[[12, 17], + [11, 16]], + [[11, 12], + [15, 16]]], + [[[23, 15], + [30, 22]], + [[14, 17], + [15, 16]]]]) + chi2, p, dof, expected = chi2_contingency(obs) + assert_approx_equal(chi2, 8.758, significant=4) + assert_approx_equal(p, 0.6442, significant=4) + assert_equal(dof, 11) + + def test_chi2_contingency_g(self): + c = np.array([[15, 60], [15, 90]]) + g, p, dof, e = chi2_contingency(c, lambda_='log-likelihood', + correction=False) + assert_allclose(g, 2*xlogy(c, c/e).sum()) + + g, p, dof, e = chi2_contingency(c, lambda_='log-likelihood', + correction=True) + c_corr = c + np.array([[-0.5, 0.5], [0.5, -0.5]]) + assert_allclose(g, 2*xlogy(c_corr, c_corr/e).sum()) + + c = np.array([[10, 12, 10], [12, 10, 10]]) + g, p, dof, e = chi2_contingency(c, lambda_='log-likelihood') + assert_allclose(g, 2*xlogy(c, c/e).sum()) + + def test_chi2_contingency_bad_args(self): + # Test that "bad" inputs raise a ValueError. + + # Negative value in the array of observed frequencies. + obs = np.array([[-1, 10], [1, 2]]) + assert_raises(ValueError, chi2_contingency, obs) + + # The zeros in this will result in zeros in the array + # of expected frequencies. + obs = np.array([[0, 1], [0, 1]]) + assert_raises(ValueError, chi2_contingency, obs) + + # A degenerate case: `observed` has size 0. + obs = np.empty((0, 8)) + assert_raises(ValueError, chi2_contingency, obs) + + def test_chi2_contingency_yates_gh13875(self): + # Magnitude of Yates' continuity correction should not exceed difference + # between expected and observed value of the statistic; see gh-13875 + observed = np.array([[1573, 3], [4, 0]]) + p = chi2_contingency(observed)[1] + assert_allclose(p, 1, rtol=1e-12) + + @pytest.mark.parametrize("correction", [False, True]) + def test_result(self, correction): + obs = np.array([[1, 2], [1, 2]]) + res = chi2_contingency(obs, correction=correction) + assert_equal((res.statistic, res.pvalue, res.dof, res.expected_freq), res) + + @pytest.mark.slow + def test_exact_permutation(self): + table = np.arange(4).reshape(2, 2) + ref_statistic = chi2_contingency(table, correction=False).statistic + ref_pvalue = stats.fisher_exact(table).pvalue + method = stats.PermutationMethod(n_resamples=50000) + res = chi2_contingency(table, correction=False, method=method) + assert_equal(res.statistic, ref_statistic) + assert_allclose(res.pvalue, ref_pvalue, rtol=1e-15) + + @pytest.mark.slow + @pytest.mark.parametrize('method', (stats.PermutationMethod, + stats.MonteCarloMethod)) + def test_resampling_randomized(self, method): + rng = np.random.default_rng(2592340925) + # need to have big sum for asymptotic approximation to be good + rows = [300, 1000, 800] + cols = [200, 400, 800, 700] + table = stats.random_table(rows, cols, seed=rng).rvs() + res = chi2_contingency(table, correction=False, method=method(rng=rng)) + ref = chi2_contingency(table, correction=False) + assert_equal(res.statistic, ref.statistic) + assert_allclose(res.pvalue, ref.pvalue, atol=5e-3) + assert_equal(res.dof, np.nan) + assert_equal(res.expected_freq, ref.expected_freq) + + def test_resampling_invalid_args(self): + table = np.arange(8).reshape(2, 2, 2) + + method = stats.PermutationMethod() + message = "Use of `method` is only compatible with two-way tables." + with pytest.raises(ValueError, match=message): + chi2_contingency(table, correction=False, method=method) + + table = np.arange(4).reshape(2, 2) + + method = stats.PermutationMethod() + message = "`correction=True` is not compatible with..." + with pytest.raises(ValueError, match=message): + chi2_contingency(table, method=method) + + method = stats.MonteCarloMethod() + message = "`lambda_=2` is not compatible with..." + with pytest.raises(ValueError, match=message): + chi2_contingency(table, correction=False, lambda_=2, method=method) + + method = 'herring' + message = "`method='herring'` not recognized; if provided, `method`..." + with pytest.raises(ValueError, match=message): + chi2_contingency(table, correction=False, method=method) + + method = stats.MonteCarloMethod(rvs=stats.norm.rvs) + message = "If the `method` argument of `chi2_contingency` is..." + with pytest.raises(ValueError, match=message): + chi2_contingency(table, correction=False, method=method) + + +def test_bad_association_args(): + # Invalid Test Statistic + assert_raises(ValueError, association, [[1, 2], [3, 4]], "X") + # Invalid array shape + assert_raises(ValueError, association, [[[1, 2]], [[3, 4]]], "cramer") + # chi2_contingency exception + assert_raises(ValueError, association, [[-1, 10], [1, 2]], 'cramer') + # Invalid Array Item Data Type + assert_raises(ValueError, association, + np.array([[1, 2], ["dd", 4]], dtype=object), 'cramer') + + +@pytest.mark.parametrize('stat, expected', + [('cramer', 0.09222412010290792), + ('tschuprow', 0.0775509319944633), + ('pearson', 0.12932925727138758)]) +def test_assoc(stat, expected): + # 2d Array + obs1 = np.array([[12, 13, 14, 15, 16], + [17, 16, 18, 19, 11], + [9, 15, 14, 12, 11]]) + a = association(observed=obs1, method=stat) + assert_allclose(a, expected) diff --git a/v12x/pylib/scipy/stats/tests/test_continued_fraction.py b/v12x/pylib/scipy/stats/tests/test_continued_fraction.py new file mode 100644 index 0000000000000000000000000000000000000000..e6911aa4956cc66a95eeb1e87aedb04782e2440b --- /dev/null +++ b/v12x/pylib/scipy/stats/tests/test_continued_fraction.py @@ -0,0 +1,168 @@ +import math + +import pytest +import numpy as np + +from scipy._lib._array_api import array_namespace +from scipy._lib._array_api_no_0d import xp_assert_close, xp_assert_less, xp_assert_equal +from scipy.stats._continued_fraction import _continued_fraction + + +@pytest.mark.skip_xp_backends('array_api_strict', reason='No fancy indexing assignment') +@pytest.mark.skip_xp_backends('jax.numpy', reason="Don't support mutation") +# dask doesn't like lines like this +# n = int(xp.real(xp_ravel(n))[0]) +# (at some point in here the shape becomes nan) +@pytest.mark.skip_xp_backends('dask.array', reason="dask has issues with the shapes") +@pytest.mark.uses_xp_capabilities(False, reason="private") +class TestContinuedFraction: + rng = np.random.default_rng(5895448232066142650) + p = rng.uniform(1, 10, size=10) + + def a1(self, n, x=1.5): + if n == 0: + y = 0*x + elif n == 1: + y = x + else: + y = -x**2 + return y + + def b1(self, n, x=1.5): + if n == 0: + y = 0*x + else: + one = x/x # gets array of correct type, dtype, and shape + y = one * (2*n - 1) + return y + + def log_a1(self, n, x): + xp = array_namespace(x) + if n == 0: + y = xp.full_like(x, -xp.asarray(math.inf, dtype=x.dtype)) + elif n == 1: + y = xp.log(x) + else: + y = 2 * xp.log(x) + math.pi * 1j + return y + + def log_b1(self, n, x): + xp = array_namespace(x) + if n == 0: + y = xp.full_like(x, -xp.asarray(math.inf, dtype=x.dtype)) + else: + one = x - x # gets array of correct type, dtype, and shape + y = one + math.log(2 * n - 1) + return y + + def test_input_validation(self, xp): + a1 = self.a1 + b1 = self.b1 + + message = '`a` and `b` must be callable.' + with pytest.raises(ValueError, match=message): + _continued_fraction(1, b1) + with pytest.raises(ValueError, match=message): + _continued_fraction(a1, 1) + + message = r'`eps` and `tiny` must be \(or represent the logarithm of\)...' + with pytest.raises(ValueError, match=message): + _continued_fraction(a1, b1, tolerances={'eps': -10}) + with pytest.raises(ValueError, match=message): + _continued_fraction(a1, b1, tolerances={'eps': np.nan}) + with pytest.raises(ValueError, match=message): + _continued_fraction(a1, b1, tolerances={'eps': 1+1j}, log=True) + with pytest.raises(ValueError, match=message): + _continued_fraction(a1, b1, tolerances={'tiny': 0}) + with pytest.raises(ValueError, match=message): + _continued_fraction(a1, b1, tolerances={'tiny': np.inf}) + with pytest.raises(ValueError, match=message): + _continued_fraction(a1, b1, tolerances={'tiny': np.inf}, log=True) + # this should not raise + kwargs = dict(args=xp.asarray(1.5+0j), log=True, maxiter=0) + _continued_fraction(a1, b1, tolerances={'eps': -10}, **kwargs) + _continued_fraction(a1, b1, tolerances={'tiny': -10}, **kwargs) + + message = '`maxiter` must be a non-negative integer.' + with pytest.raises(ValueError, match=message): + _continued_fraction(a1, b1, maxiter=-1) + + message = '`log` must be boolean.' + with pytest.raises(ValueError, match=message): + _continued_fraction(a1, b1, log=2) + + @pytest.mark.parametrize('dtype', ['float32', 'float64', 'complex64', 'complex128']) + @pytest.mark.parametrize('shape', [(), (1,), (3,), (3, 2)]) + def test_basic(self, shape, dtype, xp): + np_dtype = getattr(np, dtype) + xp_dtype = getattr(xp, dtype) + rng = np.random.default_rng(2435908729190400) + + x = rng.random(shape).astype(np_dtype) + x = x + rng.random(shape).astype(np_dtype)*1j if dtype.startswith('c') else x + x = xp.asarray(x, dtype=xp_dtype) + + res = _continued_fraction(self.a1, self.b1, args=(x,)) + ref = xp.tan(x) + xp_assert_close(res.f, ref) + + @pytest.mark.skip_xp_backends('torch', reason='pytorch/pytorch#136063') + @pytest.mark.parametrize('dtype', ['float32', 'float64']) + @pytest.mark.parametrize('shape', [(), (1,), (3,), (3, 2)]) + def test_log(self, shape, dtype, xp): + np_dtype = getattr(np, dtype) + rng = np.random.default_rng(2435908729190400) + x = rng.random(shape).astype(np_dtype) + x = xp.asarray(x) + + res = _continued_fraction(self.log_a1, self.log_b1, args=(x + 0j,), log=True) + ref = xp.tan(x) + xp_assert_close(xp.exp(xp.real(res.f)), ref) + + def test_maxiter(self, xp): + rng = np.random.default_rng(2435908729190400) + x = xp.asarray(rng.random(), dtype=xp.float64) + ref = xp.tan(x) + + res1 = _continued_fraction(self.a1, self.b1, args=(x,), maxiter=3) + assert res1.nit == 3 + + res2 = _continued_fraction(self.a1, self.b1, args=(x,), maxiter=6) + assert res2.nit == 6 + + xp_assert_less(xp.abs(res2.f - ref), xp.abs(res1.f - ref)) + + def test_eps(self, xp): + x = xp.asarray(1.5, dtype=xp.float64) # x = 1.5 is the default defined above + ref = xp.tan(x) + res1 = _continued_fraction(self.a1, self.b1, args=(x,), + tolerances={'eps': 1e-6}) + res2 = _continued_fraction(self.a1, self.b1, args=(x,)) + xp_assert_less(res1.nit, res2.nit) + xp_assert_less(xp.abs(res2.f - ref), xp.abs(res1.f - ref)) + + def test_feval(self, xp): + def a(n, x): + a.nfev += 1 + return n * x + + def b(n, x): + b.nfev += 1 + return n * x + + a.nfev, b.nfev = 0, 0 + + res = _continued_fraction(a, b, args=(xp.asarray(1.),)) + assert res.nfev == a.nfev == b.nfev == res.nit + 1 + + def test_status(self, xp): + x = xp.asarray([1, 10, np.nan], dtype=xp.float64) + res = _continued_fraction(self.a1, self.b1, args=(x,), maxiter=15) + xp_assert_equal(res.success, xp.asarray([True, False, False])) + xp_assert_equal(res.status, xp.asarray([0, -2, -3], dtype=xp.int32)) + + def test_special_cases(self, xp): + one = xp.asarray(1) + res = _continued_fraction(lambda x: one, lambda x: one, maxiter=0) + xp_assert_close(res.f, xp.asarray(1.)) + assert res.nit == res.nfev - 1 == 0 diff --git a/v12x/pylib/scipy/stats/tests/test_continuous.py b/v12x/pylib/scipy/stats/tests/test_continuous.py new file mode 100644 index 0000000000000000000000000000000000000000..898b7bcdf6b5ae4d72404fe98a3fd32ea8a90564 --- /dev/null +++ b/v12x/pylib/scipy/stats/tests/test_continuous.py @@ -0,0 +1,2328 @@ +import itertools as it +import os +import pickle +from copy import deepcopy + +import numpy as np +from numpy import inf +import pytest +from numpy.testing import assert_allclose, assert_equal +from hypothesis import strategies, given, reproduce_failure, settings # noqa: F401 +import hypothesis.extra.numpy as npst + +from scipy import special +from scipy import stats +from scipy.stats._fit import _kolmogorov_smirnov +from scipy.stats._ksstats import kolmogn +from scipy.stats import qmc +from scipy.stats._distr_params import distcont, distdiscrete +from scipy.stats._distribution_infrastructure import ( + _Domain, _RealInterval, _Parameter, _Parameterization, _RealParameter, + ContinuousDistribution, ShiftedScaledDistribution, _fiinfo, + _generate_domain_support, Mixture, _logexpxmexpy) +from scipy.stats._new_distributions import StandardNormal, _LogUniform, _Gamma +from scipy.stats._new_distributions import DiscreteDistribution +from scipy.stats import Normal, Logistic, Uniform, Binomial + + +class Test_RealInterval: + rng = np.random.default_rng(349849812549824) + + def test_iv(self): + domain = _RealInterval(endpoints=('a', 'b')) + message = "The endpoints of the distribution are defined..." + with pytest.raises(TypeError, match=message): + domain.get_numerical_endpoints(dict) + + @pytest.mark.parametrize('x', [rng.uniform(10, 10, size=(2, 3, 4)), + -np.inf, np.pi]) + def test_contains_simple(self, x): + # Test `contains` when endpoints are defined by constants + a, b = -np.inf, np.pi + domain = _RealInterval(endpoints=(a, b), inclusive=(False, True)) + assert_equal(domain.contains(x), (a < x) & (x <= b)) + + @pytest.mark.slow + @given(shapes=npst.mutually_broadcastable_shapes(num_shapes=3, min_side=0), + inclusive_a=strategies.booleans(), + inclusive_b=strategies.booleans(), + data=strategies.data()) + def test_contains(self, shapes, inclusive_a, inclusive_b, data): + # Test `contains` when endpoints are defined by parameters + input_shapes, result_shape = shapes + shape_a, shape_b, shape_x = input_shapes + + # Without defining min and max values, I spent forever trying to set + # up a valid test without overflows or similar just drawing arrays. + a_elements = dict(allow_nan=False, allow_infinity=False, + min_value=-1e3, max_value=1) + b_elements = dict(allow_nan=False, allow_infinity=False, + min_value=2, max_value=1e3) + a = data.draw(npst.arrays(npst.floating_dtypes(), + shape_a, elements=a_elements)) + b = data.draw(npst.arrays(npst.floating_dtypes(), + shape_b, elements=b_elements)) + # ensure some points are to the left, some to the right, and some + # are exactly on the boundary + d = b - a + x = np.concatenate([np.linspace(a-d, a, 10), + np.linspace(a, b, 10), + np.linspace(b, b+d, 10)]) + # Domain is defined by two parameters, 'a' and 'b' + domain = _RealInterval(endpoints=('a', 'b'), + inclusive=(inclusive_a, inclusive_b)) + domain.define_parameters(_RealParameter('a', domain=_RealInterval()), + _RealParameter('b', domain=_RealInterval())) + # Check that domain and string evaluation give the same result + res = domain.contains(x, dict(a=a, b=b)) + + # Apparently, `np.float16([2]) < np.float32(2.0009766)` is False + # but `np.float16([2]) < np.float32([2.0009766])` is True + # dtype = np.result_type(a.dtype, b.dtype, x.dtype) + # a, b, x = a.astype(dtype), b.astype(dtype), x.astype(dtype) + # unclear whether we should be careful about this, since it will be + # fixed with NEP50. Just do what makes the test pass. + left_comparison = '<=' if inclusive_a else '<' + right_comparison = '<=' if inclusive_b else '<' + ref = eval(f'(a {left_comparison} x) & (x {right_comparison} b)') + assert_equal(res, ref) + + @pytest.mark.parametrize("inclusive", list(it.product([True, False], repeat=2))) + @pytest.mark.parametrize("a,b", [(0, 1), (3, 1)]) + def test_contains_function_endpoints(self, inclusive, a, b): + # Test `contains` when endpoints are defined by functions. + endpoints = (lambda a, b: (a - b) / 2, lambda a, b: (a + b) / 2) + domain = _RealInterval(endpoints=endpoints, inclusive=inclusive) + x = np.asarray([(a - 2*b)/2, (a - b)/2, a/2, (a + b)/2, (a + 2*b)/2]) + res = domain.contains(x, dict(a=a, b=b)) + + numerical_endpoints = ((a - b) / 2, (a + b) / 2) + assert numerical_endpoints == domain.get_numerical_endpoints(dict(a=a, b=b)) + alpha, beta = numerical_endpoints + + above_left = alpha <= x if inclusive[0] else alpha < x + below_right = x <= beta if inclusive[1] else x < beta + ref = above_left & below_right + assert_equal(res, ref) + + + @pytest.mark.parametrize('case', [ + (-np.inf, np.pi, False, True, r"(-\infty, \pi]"), + ('a', 5, True, False, "[a, 5)") + ]) + def test_str(self, case): + domain = _RealInterval(endpoints=case[:2], inclusive=case[2:4]) + assert str(domain) == case[4] + + @pytest.mark.slow + @given(a=strategies.one_of( + strategies.decimals(allow_nan=False), + strategies.characters(whitelist_categories="L"), # type: ignore[arg-type] + strategies.sampled_from(list(_Domain.symbols))), + b=strategies.one_of( + strategies.decimals(allow_nan=False), + strategies.characters(whitelist_categories="L"), # type: ignore[arg-type] + strategies.sampled_from(list(_Domain.symbols))), + inclusive_a=strategies.booleans(), + inclusive_b=strategies.booleans(), + ) + def test_str2(self, a, b, inclusive_a, inclusive_b): + # I wrote this independently from the implementation of __str__, but + # I imagine it looks pretty similar to __str__. + a = _Domain.symbols.get(a, a) + b = _Domain.symbols.get(b, b) + left_bracket = '[' if inclusive_a else '(' + right_bracket = ']' if inclusive_b else ')' + domain = _RealInterval(endpoints=(a, b), + inclusive=(inclusive_a, inclusive_b)) + ref = f"{left_bracket}{a}, {b}{right_bracket}" + assert str(domain) == ref + + def test_symbols_gh22137(self): + # `symbols` was accidentally shared between instances originally + # Check that this is no longer the case + domain1 = _RealInterval(endpoints=(0, 1)) + domain2 = _RealInterval(endpoints=(0, 1)) + assert domain1.symbols is not domain2.symbols + + +def draw_distribution_from_family(family, data, rng, proportions, min_side=0): + # If the distribution has parameters, choose a parameterization and + # draw broadcastable shapes for the parameter arrays. + n_parameterizations = family._num_parameterizations() + if n_parameterizations > 0: + i = data.draw(strategies.integers(0, max_value=n_parameterizations-1)) + n_parameters = family._num_parameters(i) + shapes, result_shape = data.draw( + npst.mutually_broadcastable_shapes(num_shapes=n_parameters, + min_side=min_side)) + dist = family._draw(shapes, rng=rng, proportions=proportions, + i_parameterization=i) + else: + dist = family._draw(rng=rng) + result_shape = tuple() + + # Draw a broadcastable shape for the arguments, and draw values for the + # arguments. + x_shape = data.draw(npst.broadcastable_shapes(result_shape, + min_side=min_side)) + x = dist._variable.draw(x_shape, parameter_values=dist._parameters, + proportions=proportions, rng=rng, region='typical') + x_result_shape = np.broadcast_shapes(x_shape, result_shape) + y_shape = data.draw(npst.broadcastable_shapes(x_result_shape, + min_side=min_side)) + y = dist._variable.draw(y_shape, parameter_values=dist._parameters, + proportions=proportions, rng=rng, region='typical') + xy_result_shape = np.broadcast_shapes(y_shape, x_result_shape) + p_domain = _RealInterval((0, 1), (True, True)) + p_var = _RealParameter('p', domain=p_domain) + p = p_var.draw(x_shape, proportions=proportions, rng=rng) + with np.errstate(divide='ignore', invalid='ignore'): + logp = np.log(p) + + return dist, x, y, p, logp, result_shape, x_result_shape, xy_result_shape + + +continuous_families = [ + StandardNormal, + Normal, + Logistic, + Uniform, + _LogUniform +] + +discrete_families = [ + Binomial, +] + +families = continuous_families + discrete_families + + +class TestDistributions: + @pytest.mark.fail_slow(60) # need to break up check_moment_funcs + @settings(max_examples=20) + @pytest.mark.parametrize('family', families) + @given(data=strategies.data(), seed=strategies.integers(min_value=0)) + def test_support_moments_sample(self, family, data, seed): + rng = np.random.default_rng(seed) + + # relative proportions of valid, endpoint, out of bounds, and NaN params + proportions = (0.7, 0.1, 0.1, 0.1) + tmp = draw_distribution_from_family(family, data, rng, proportions) + dist, x, y, p, logp, result_shape, x_result_shape, xy_result_shape = tmp + sample_shape = data.draw(npst.array_shapes(min_dims=0, min_side=0, + max_side=20)) + + with np.errstate(invalid='ignore', divide='ignore'): + check_support(dist) + check_moment_funcs(dist, result_shape) # this needs to get split up + check_lmoment_funcs(dist, result_shape) + check_sample_shape_NaNs(dist, 'sample', sample_shape, result_shape, rng) + qrng = qmc.Halton(d=1, seed=rng) + check_sample_shape_NaNs(dist, 'sample', sample_shape, result_shape, qrng) + + @pytest.mark.fail_slow(10) + @pytest.mark.parametrize('family', families) + @pytest.mark.parametrize('func, methods, arg', + [('entropy', {'log/exp', 'quadrature'}, None), + ('logentropy', {'log/exp', 'quadrature'}, None), + ('median', {'icdf'}, None), + ('mode', {'optimization'}, None), + ('mean', {'cache'}, None), + ('variance', {'cache'}, None), + ('skewness', {'cache'}, None), + ('kurtosis', {'cache'}, None), + ('pdf', {'log/exp'}, 'x'), + ('logpdf', {'log/exp'}, 'x'), + ('logcdf', {'log/exp', 'complement', 'quadrature'}, 'x'), + ('cdf', {'log/exp', 'complement', 'quadrature'}, 'x'), + ('logccdf', {'log/exp', 'complement', 'quadrature'}, 'x'), + ('ccdf', {'log/exp', 'complement', 'quadrature'}, 'x'), + ('ilogccdf', {'complement', 'inversion'}, 'logp'), + ('iccdf', {'complement', 'inversion'}, 'p'), + ]) + @settings(max_examples=20) + @given(data=strategies.data(), seed=strategies.integers(min_value=0)) + def test_funcs(self, family, data, seed, func, methods, arg): + if family == Uniform and func == 'mode': + pytest.skip("Mode is not unique; `method`s disagree.") + + rng = np.random.default_rng(seed) + + # relative proportions of valid, endpoint, out of bounds, and NaN params + proportions = (0.7, 0.1, 0.1, 0.1) + tmp = draw_distribution_from_family(family, data, rng, proportions) + dist, x, y, p, logp, result_shape, x_result_shape, xy_result_shape = tmp + + args = {'x': x, 'p': p, 'logp': p} + with np.errstate(invalid='ignore', divide='ignore', over='ignore'): + if arg is None: + check_dist_func(dist, func, None, result_shape, methods) + elif arg in args: + check_dist_func(dist, func, args[arg], x_result_shape, methods) + + if func == 'variance': + assert_allclose(dist.standard_deviation()**2, dist.variance()) + + # invalid and divide are to be expected; maybe look into over + with np.errstate(invalid='ignore', divide='ignore', over='ignore'): + if not isinstance(dist, ShiftedScaledDistribution): + if func == 'cdf': + methods = {'quadrature'} + check_cdf2(dist, False, x, y, xy_result_shape, methods) + check_cdf2(dist, True, x, y, xy_result_shape, methods) + elif func == 'ccdf': + methods = {'addition'} + check_ccdf2(dist, False, x, y, xy_result_shape, methods) + check_ccdf2(dist, True, x, y, xy_result_shape, methods) + + def test_plot(self): + try: + import matplotlib.pyplot as plt + except ImportError: + return + + X = Uniform(a=0., b=1.) + ax = X.plot() + assert ax == plt.gca() + + @pytest.mark.parametrize('method_name', ['cdf', 'ccdf']) + def test_complement_safe(self, method_name): + X = stats.Normal(mu=1, sigma=2) + X.tol = 1e-12 + p = np.asarray([1e-4, 1e-3]) + func = getattr(X, method_name) + ifunc = getattr(X, 'i'+method_name) + x = ifunc(p, method='formula') + p1 = func(x, method='complement_safe') + p2 = func(x, method='complement') + assert_equal(p1[1], p2[1]) + assert p1[0] != p2[0] + assert_allclose(p1[0], p[0], rtol=X.tol) + + @pytest.mark.parametrize('method_name', ['cdf', 'ccdf']) + def test_icomplement_safe(self, method_name): + X = stats.Normal(mu=1, sigma=2) + X.tol = 1e-12 + p = np.asarray([1e-4, 1e-3]) + func = getattr(X, method_name) + ifunc = getattr(X, 'i'+method_name) + x1 = ifunc(p, method='complement_safe') + x2 = ifunc(p, method='complement') + assert_equal(x1[1], x2[1]) + assert x1[0] != x2[0] + assert_allclose(func(x1[0]), p[0], rtol=X.tol) + + def test_subtraction_safe(self): + X = stats.Normal() + X.tol = 1e-12 + + # Regular subtraction is fine in either tail (and of course, across tails) + x = [-11, -10, 10, 11] + y = [-10, -11, 11, 10] + p0 = X.cdf(x, y, method='quadrature') + p1 = X.cdf(x, y, method='subtraction_safe') + p2 = X.cdf(x, y, method='subtraction') + assert_equal(p2, p1) + assert_allclose(p1, p0, rtol=X.tol) + + # Safe subtraction is needed in special cases + x = np.asarray([-1e-20, -1e-21, 1e-20, 1e-21, -1e-20]) + y = np.asarray([-1e-21, -1e-20, 1e-21, 1e-20, 1e-20]) + + + p0 = X.pdf(0)*(y-x) + p1 = X.cdf(x, y, method='subtraction_safe') + p2 = X.cdf(x, y, method='subtraction') + assert_equal(p2, 0) + assert_allclose(p1, p0, rtol=X.tol) + + def test_logentropy_safe(self): + # simulate an `entropy` calculation over/underflowing with extreme parameters + class _Normal(stats.Normal): + def _entropy_formula(self, **params): + out = np.asarray(super()._entropy_formula(**params)) + out[0] = 0 + out[-1] = np.inf + return out + + X = _Normal(sigma=[1, 2, 3]) + with np.errstate(divide='ignore'): + res1 = X.logentropy(method='logexp_safe') + res2 = X.logentropy(method='logexp') + ref = X.logentropy(method='quadrature') + i_fl = [0, -1] # first and last + assert np.isinf(res2[i_fl]).all() + assert res1[1] == res2[1] + # quadrature happens to be perfectly accurate on some platforms + # assert res1[1] != ref[1] + assert_equal(res1[i_fl], ref[i_fl]) + + def test_logcdf2_safe(self): + # test what happens when 2-arg `cdf` underflows + X = stats.Normal(sigma=[1, 2, 3]) + x = [-301, 1, 300] + y = [-300, 2, 301] + with np.errstate(divide='ignore'): + res1 = X.logcdf(x, y, method='logexp_safe') + res2 = X.logcdf(x, y, method='logexp') + ref = X.logcdf(x, y, method='quadrature') + i_fl = [0, -1] # first and last + assert np.isinf(res2[i_fl]).all() + assert res1[1] == res2[1] + # quadrature happens to be perfectly accurate on some platforms + # assert res1[1] != ref[1] + assert_equal(res1[i_fl], ref[i_fl]) + + @pytest.mark.parametrize('method_name', ['logcdf', 'logccdf']) + def test_logexp_safe(self, method_name): + # test what happens when `cdf`/`ccdf` underflows + X = stats.Normal(sigma=2) + x = [-301, 1] if method_name == 'logcdf' else [301, 1] + func = getattr(X, method_name) + with np.errstate(divide='ignore'): + res1 = func(x, method='logexp_safe') + res2 = func(x, method='logexp') + ref = func(x, method='quadrature') + assert res1[0] == ref[0] + assert res1[0] != res2[0] + assert res1[1] == res2[1] + assert res1[1] != ref[1] + +def check_sample_shape_NaNs(dist, fname, sample_shape, result_shape, rng): + full_shape = sample_shape + result_shape + if fname == 'sample': + sample_method = dist.sample + + methods = {'inverse_transform'} + if dist._overrides(f'_{fname}_formula') and not isinstance(rng, qmc.QMCEngine): + methods.add('formula') + + for method in methods: + res = sample_method(sample_shape, method=method, rng=rng) + valid_parameters = np.broadcast_to(get_valid_parameters(dist), + res.shape) + assert_equal(res.shape, full_shape) + np.testing.assert_equal(res.dtype, dist._dtype) + + if full_shape == (): + # NumPy random makes a distinction between a 0d array and a scalar. + # In stats, we consistently turn 0d arrays into scalars, so + # maintain that behavior here. (With Array API arrays, this will + # change.) + assert np.isscalar(res) + assert np.all(np.isfinite(res[valid_parameters])) + assert_equal(res[~valid_parameters], np.nan) + + sample1 = sample_method(sample_shape, method=method, rng=42) + sample2 = sample_method(sample_shape, method=method, rng=42) + if not isinstance(dist, DiscreteDistribution): + # The idea is that it's very unlikely that the random sample + # for a randomly chosen seed will match that for seed 42, + # but it is not so unlikely if `dist` is a discrete distribution. + assert not np.any(np.equal(res, sample1)) + assert_equal(sample1, sample2) + + +def check_support(dist): + a, b = dist.support() + check_nans_and_edges(dist, 'support', None, a) + check_nans_and_edges(dist, 'support', None, b) + assert a.shape == dist._shape + assert b.shape == dist._shape + assert a.dtype == dist._dtype + assert b.dtype == dist._dtype + + +def check_dist_func(dist, fname, arg, result_shape, methods): + # Check that all computation methods of all distribution functions agree + # with one another, effectively testing the correctness of the generic + # computation methods and confirming the consistency of specific + # distributions with their pdf/logpdf. + + args = tuple() if arg is None else (arg,) + methods = methods.copy() + + if "cache" in methods: + # If "cache" is specified before the value has been evaluated, it + # raises an error. After the value is evaluated, it will succeed. + with pytest.raises(NotImplementedError): + getattr(dist, fname)(*args, method="cache") + + ref = getattr(dist, fname)(*args) + check_nans_and_edges(dist, fname, arg, ref) + + # Remove this after fixing `draw` + tol_override = {'atol': 1e-15} + # Mean can be 0, which makes logmean -inf. + if fname in {'logmean', 'mean', 'logskewness', 'skewness'}: + tol_override = {'atol': 1e-15} + elif fname in {'mode'}: + # can only expect about half of machine precision for optimization + # because math + tol_override = {'atol': 1e-6} + elif fname in {'logcdf'}: # gh-22276 + tol_override = {'rtol': 2e-7} + + if dist._overrides(f'_{fname}_formula'): + methods.add('formula') + + np.testing.assert_equal(ref.shape, result_shape) + # Until we convert to array API, let's do the familiar thing: + # 0d things are scalars, not arrays + if result_shape == tuple(): + assert np.isscalar(ref) + + for method in methods: + res = getattr(dist, fname)(*args, method=method) + if 'log' in fname: + np.testing.assert_allclose(np.exp(res), np.exp(ref), + **tol_override) + else: + np.testing.assert_allclose(res, ref, **tol_override) + + # for now, make sure dtypes are consistent; later, we can check whether + # they are correct. + np.testing.assert_equal(res.dtype, ref.dtype) + np.testing.assert_equal(res.shape, result_shape) + if result_shape == tuple(): + assert np.isscalar(res) + +def check_cdf2(dist, log, x, y, result_shape, methods): + # Specialized test for 2-arg cdf since the interface is a bit different + # from the other methods. Here, we'll use 1-arg cdf as a reference, and + # since we have already checked 1-arg cdf in `check_nans_and_edges`, this + # checks the equivalent of both `check_dist_func` and + # `check_nans_and_edges`. + methods = methods.copy() + + if log: + if dist._overrides('_logcdf2_formula'): + methods.add('formula') + if dist._overrides('_logcdf_formula') or dist._overrides('_logccdf_formula'): + methods.add('subtraction') + if (dist._overrides('_cdf_formula') + or dist._overrides('_ccdf_formula')): + methods.add('log/exp') + else: + if dist._overrides('_cdf2_formula'): + methods.add('formula') + if dist._overrides('_cdf_formula') or dist._overrides('_ccdf_formula'): + methods.add('subtraction') + if (dist._overrides('_logcdf_formula') + or dist._overrides('_logccdf_formula')): + methods.add('log/exp') + + ref = dist.cdf(y) - dist.cdf(x) + np.testing.assert_equal(ref.shape, result_shape) + + if result_shape == tuple(): + assert np.isscalar(ref) + + for method in methods: + if isinstance(dist, DiscreteDistribution): + message = ("Two argument cdf functions are currently only supported for " + "continuous distributions.") + with pytest.raises(NotImplementedError, match=message): + res = (np.exp(dist.logcdf(x, y, method=method)) if log + else dist.cdf(x, y, method=method)) + continue + res = (np.exp(dist.logcdf(x, y, method=method)) if log + else dist.cdf(x, y, method=method)) + np.testing.assert_allclose(res, ref, atol=1e-14) + if log: + np.testing.assert_equal(res.dtype, (ref + 0j).dtype) + else: + np.testing.assert_equal(res.dtype, ref.dtype) + np.testing.assert_equal(res.shape, result_shape) + if result_shape == tuple(): + assert np.isscalar(res) + + +def check_ccdf2(dist, log, x, y, result_shape, methods): + # Specialized test for 2-arg ccdf since the interface is a bit different + # from the other methods. Could be combined with check_cdf2 above, but + # writing it separately is simpler. + methods = methods.copy() + + if dist._overrides(f'_{"log" if log else ""}ccdf2_formula'): + methods.add('formula') + + ref = dist.cdf(x) + dist.ccdf(y) + np.testing.assert_equal(ref.shape, result_shape) + + if result_shape == tuple(): + assert np.isscalar(ref) + + for method in methods: + message = ("Two argument cdf functions are currently only supported for " + "continuous distributions.") + if isinstance(dist, DiscreteDistribution): + with pytest.raises(NotImplementedError, match=message): + res = (np.exp(dist.logccdf(x, y, method=method)) if log + else dist.ccdf(x, y, method=method)) + continue + res = (np.exp(dist.logccdf(x, y, method=method)) if log + else dist.ccdf(x, y, method=method)) + np.testing.assert_allclose(res, ref, atol=1e-14) + np.testing.assert_equal(res.dtype, ref.dtype) + np.testing.assert_equal(res.shape, result_shape) + if result_shape == tuple(): + assert np.isscalar(res) + + +def check_nans_and_edges(dist, fname, arg, res): + + valid_parameters = get_valid_parameters(dist) + if fname in {'icdf', 'iccdf'}: + arg_domain = _RealInterval(endpoints=(0, 1), inclusive=(True, True)) + elif fname in {'ilogcdf', 'ilogccdf'}: + arg_domain = _RealInterval(endpoints=(-inf, 0), inclusive=(True, True)) + else: + arg_domain = dist._variable.domain + + classified_args = classify_arg(dist, arg, arg_domain) + valid_parameters, *classified_args = np.broadcast_arrays(valid_parameters, + *classified_args) + valid_arg, endpoint_arg, outside_arg, nan_arg = classified_args + all_valid = valid_arg & valid_parameters + + # Check NaN pattern and edge cases + assert_equal(res[~valid_parameters], np.nan) + assert_equal(res[nan_arg], np.nan) + + a, b = dist.support() + a = np.broadcast_to(a, res.shape) + b = np.broadcast_to(b, res.shape) + + outside_arg_minus = (outside_arg == -1) & valid_parameters + outside_arg_plus = (outside_arg == 1) & valid_parameters + endpoint_arg_minus = (endpoint_arg == -1) & valid_parameters + endpoint_arg_plus = (endpoint_arg == 1) & valid_parameters + + is_discrete = isinstance(dist, DiscreteDistribution) + # Writing this independently of how the are set in the distribution + # infrastructure. That is very compact; this is very verbose. + if fname in {'logpdf'}: + assert_equal(res[outside_arg_minus], -np.inf) + assert_equal(res[outside_arg_plus], -np.inf) + ref = -np.inf if not is_discrete else np.inf + assert_equal(res[endpoint_arg_minus & ~valid_arg], ref) + assert_equal(res[endpoint_arg_plus & ~valid_arg], ref) + elif fname in {'pdf'}: + assert_equal(res[outside_arg_minus], 0) + assert_equal(res[outside_arg_plus], 0) + ref = 0 if not is_discrete else np.inf + assert_equal(res[endpoint_arg_minus & ~valid_arg], ref) + assert_equal(res[endpoint_arg_plus & ~valid_arg], ref) + elif fname in {'logcdf'} and not is_discrete: + assert_equal(res[outside_arg_minus], -inf) + assert_equal(res[outside_arg_plus], 0) + assert_equal(res[endpoint_arg_minus], -inf) + assert_equal(res[endpoint_arg_plus], 0) + elif fname in {'cdf'} and not is_discrete: + assert_equal(res[outside_arg_minus], 0) + assert_equal(res[outside_arg_plus], 1) + assert_equal(res[endpoint_arg_minus], 0) + assert_equal(res[endpoint_arg_plus], 1) + elif fname in {'logccdf'} and not is_discrete: + assert_equal(res[outside_arg_minus], 0) + assert_equal(res[outside_arg_plus], -inf) + assert_equal(res[endpoint_arg_minus], 0) + assert_equal(res[endpoint_arg_plus], -inf) + elif fname in {'ccdf'} and not is_discrete: + assert_equal(res[outside_arg_minus], 1) + assert_equal(res[outside_arg_plus], 0) + assert_equal(res[endpoint_arg_minus], 1) + assert_equal(res[endpoint_arg_plus], 0) + elif fname in {'ilogcdf', 'icdf'} and not is_discrete: + assert_equal(res[outside_arg == -1], np.nan) + assert_equal(res[outside_arg == 1], np.nan) + assert_equal(res[endpoint_arg == -1], a[endpoint_arg == -1]) + assert_equal(res[endpoint_arg == 1], b[endpoint_arg == 1]) + elif fname in {'ilogccdf', 'iccdf'} and not is_discrete: + assert_equal(res[outside_arg == -1], np.nan) + assert_equal(res[outside_arg == 1], np.nan) + assert_equal(res[endpoint_arg == -1], b[endpoint_arg == -1]) + assert_equal(res[endpoint_arg == 1], a[endpoint_arg == 1]) + + exclude = {'logmean', 'mean', 'logskewness', 'skewness', 'support'} + if isinstance(dist, DiscreteDistribution): + exclude.update({'pdf', 'logpdf'}) + + if fname not in exclude: + assert np.isfinite(res[all_valid & (endpoint_arg == 0)]).all() + + +def check_moment_funcs(dist, result_shape): + # Check that all computation methods of all distribution functions agree + # with one another, effectively testing the correctness of the generic + # computation methods and confirming the consistency of specific + # distributions with their pdf/logpdf. + + atol = 1e-9 # make this tighter (e.g. 1e-13) after fixing `draw` + + def check(order, kind, method=None, ref=None, success=True): + if success: + res = dist.moment(order, kind, method=method) + assert_allclose(res, ref, atol=atol*10**order) + assert res.shape == ref.shape + else: + with pytest.raises(NotImplementedError): + dist.moment(order, kind, method=method) + + def has_formula(order, kind): + formula_name = f'_moment_{kind}_formula' + overrides = dist._overrides(formula_name) + if not overrides: + return False + formula = getattr(dist, formula_name) + orders = getattr(formula, 'orders', set(range(6))) + return order in orders + + dist.reset_cache() + + ### Check Raw Moments ### + for i in range(6): + check(i, 'raw', 'cache', success=False) # not cached yet + ref = dist.moment(i, 'raw', method='quadrature') + check_nans_and_edges(dist, 'moment', None, ref) + assert ref.shape == result_shape + check(i, 'raw','cache', ref, success=True) # cached now + check(i, 'raw', 'formula', ref, success=has_formula(i, 'raw')) + check(i, 'raw', 'general', ref, success=(i == 0)) + if dist.__class__ == stats.Normal: + check(i, 'raw', 'quadrature_icdf', ref, success=True) + + + # Clearing caches to better check their behavior + dist.reset_cache() + + # If we have central or standard moment formulas, or if there are + # values in their cache, we can use method='transform' + dist.moment(0, 'central') # build up the cache + dist.moment(1, 'central') + for i in range(2, 6): + ref = dist.moment(i, 'raw', method='quadrature') + check(i, 'raw', 'transform', ref, + success=has_formula(i, 'central') or has_formula(i, 'standardized')) + dist.moment(i, 'central') # build up the cache + check(i, 'raw', 'transform', ref) + + dist.reset_cache() + + ### Check Central Moments ### + + for i in range(6): + check(i, 'central', 'cache', success=False) + ref = dist.moment(i, 'central', method='quadrature') + assert ref.shape == result_shape + check(i, 'central', 'cache', ref, success=True) + check(i, 'central', 'formula', ref, success=has_formula(i, 'central')) + check(i, 'central', 'general', ref, success=i <= 1) + if dist.__class__ == stats.Normal: + check(i, 'central', 'quadrature_icdf', ref, success=True) + if not (dist.__class__ == stats.Uniform and i == 5): + # Quadrature is not super accurate for 5th central moment when the + # support is really big. Skip this one failing test. We need to come + # up with a better system of skipping individual failures w/ hypothesis. + check(i, 'central', 'transform', ref, + success=has_formula(i, 'raw') or (i <= 1)) + if not has_formula(i, 'raw'): + dist.moment(i, 'raw') + check(i, 'central', 'transform', ref) + + variance = dist.variance() + dist.reset_cache() + + # If we have standard moment formulas, or if there are + # values in their cache, we can use method='normalize' + dist.moment(0, 'standardized') # build up the cache + dist.moment(1, 'standardized') + dist.moment(2, 'standardized') + for i in range(3, 6): + ref = dist.moment(i, 'central', method='quadrature') + check(i, 'central', 'normalize', ref, + success=has_formula(i, 'standardized') and not np.any(variance == 0)) + dist.moment(i, 'standardized') # build up the cache + check(i, 'central', 'normalize', ref, success=not np.any(variance == 0)) + + ### Check Standardized Moments ### + + var = dist.moment(2, 'central', method='quadrature') + dist.reset_cache() + + for i in range(6): + check(i, 'standardized', 'cache', success=False) + ref = dist.moment(i, 'central', method='quadrature') / var ** (i / 2) + assert ref.shape == result_shape + check(i, 'standardized', 'formula', ref, + success=has_formula(i, 'standardized')) + check(i, 'standardized', 'general', ref, success=i <= 2) + check(i, 'standardized', 'normalize', ref) + + dist.reset_cache() + + +def check_lmoment_funcs(dist, result_shape): + # Perform consistency check for L-moments similar to check_moment_funcs above + + if not isinstance(dist, ContinuousDistribution): + message = "L-moments are currently available only for continuous..." + with pytest.raises(NotImplementedError, match=message): + dist.lmoment(1) + return + + atol = 2e-9 + + def check(order, standardize=False, method=None, ref=None, success=True): + if success: + res = dist.lmoment(order, standardize=standardize, method=method) + assert_allclose(res, ref, atol=atol) + assert res.shape == ref.shape + else: + with pytest.raises(NotImplementedError): + dist.lmoment(order, standardize=standardize, method=method) + + ### Check L-Moments ### + + standardize = False + for i in range(1, 6): + check(i, standardize, 'cache', success=standardize) # not cached yet + ref = dist.lmoment(i, standardize=standardize, method='order_statistics') + check_nans_and_edges(dist, 'lmoment', None, ref) + assert ref.shape == result_shape + check(i, standardize, 'cache', ref, success=True) # cached now + check(i, standardize, 'formula', ref, + success=dist._overrides('_lmoment_formula') + and (i < 5 or dist.__class__.__name__ == "Uniform")) + check(i, standardize, 'general', ref, success=(i == 1)) + if dist._overrides('_icdf_formula'): + check(i, standardize, 'quadrature_icdf', ref, success=True) + + standardize=True + for i in range(3, 6): + ref = dist.lmoment(i, standardize=standardize, method='order_statistics') + assert ref.shape == result_shape + check(i, standardize, 'formula', ref, + success=dist._overrides('_lmoment_formula') + and (i < 5 or dist.__class__.__name__ == "Uniform")) + check(i, standardize, 'general', ref, success=False) + if dist._overrides('_icdf_formula'): + check(i, standardize, 'quadrature_icdf', ref, success=True) + + +@pytest.mark.parametrize('family', (Normal,)) +@pytest.mark.parametrize('x_shape', [tuple(), (2, 3)]) +@pytest.mark.parametrize('dist_shape', [tuple(), (4, 1)]) +@pytest.mark.parametrize('fname', ['sample']) +@pytest.mark.parametrize('rng_type', [np.random.Generator, qmc.Halton, qmc.Sobol]) +def test_sample_against_cdf(family, dist_shape, x_shape, fname, rng_type): + rng = np.random.default_rng(842582438235635) + num_parameters = family._num_parameters() + + if dist_shape and num_parameters == 0: + pytest.skip("Distribution can't have a shape without parameters.") + + dist = family._draw(dist_shape, rng) + + n = 1024 + sample_size = (n,) + x_shape + sample_array_shape = sample_size + dist_shape + + if fname == 'sample': + sample_method = dist.sample + + if rng_type != np.random.Generator: + rng = rng_type(d=1, seed=rng) + x = sample_method(sample_size, rng=rng) + assert x.shape == sample_array_shape + + # probably should give `axis` argument to ks_1samp, review that separately + statistic = _kolmogorov_smirnov(dist, x, axis=0) + pvalue = kolmogn(x.shape[0], statistic, cdf=False) + p_threshold = 0.01 + num_pvalues = pvalue.size + num_small_pvalues = np.sum(pvalue < p_threshold) + assert num_small_pvalues < p_threshold * num_pvalues + + +def get_valid_parameters(dist): + # Given a distribution, return a logical array that is true where all + # distribution parameters are within their respective domains. The code + # here is probably quite similar to that used to form the `_invalid` + # attribute of the distribution, but this was written about a week later + # without referring to that code, so it is a somewhat independent check. + + # Get all parameter values and `_Parameter` objects + parameter_values = dist._parameters + parameters = {} + for parameterization in dist._parameterizations: + parameters.update(parameterization.parameters) + + all_valid = np.ones(dist._shape, dtype=bool) + for name, value in parameter_values.items(): + if name not in parameters: # cached value not part of parameterization + continue + parameter = parameters[name] + + # Check that the numerical endpoints and inclusivity attribute + # agree with the `contains` method about which parameter values are + # within the domain. + a, b = parameter.domain.get_numerical_endpoints( + parameter_values=parameter_values) + a_included, b_included = parameter.domain.inclusive + valid = (a <= value) if a_included else a < value + valid &= (value <= b) if b_included else value < b + assert_equal(valid, parameter.domain.contains( + value, parameter_values=parameter_values)) + + # Form `all_valid` mask that is True where *all* parameters are valid + all_valid &= valid + + # Check that the `all_valid` mask formed here is the complement of the + # `dist._invalid` mask stored by the infrastructure + assert_equal(~all_valid, dist._invalid) + + return all_valid + +def classify_arg(dist, arg, arg_domain): + if arg is None: + valid_args = np.ones(dist._shape, dtype=bool) + endpoint_args = np.zeros(dist._shape, dtype=bool) + outside_args = np.zeros(dist._shape, dtype=bool) + nan_args = np.zeros(dist._shape, dtype=bool) + return valid_args, endpoint_args, outside_args, nan_args + + a, b = arg_domain.get_numerical_endpoints( + parameter_values=dist._parameters) + + a, b, arg = np.broadcast_arrays(a, b, arg) + a_included, b_included = arg_domain.inclusive + + inside = (a <= arg) if a_included else a < arg + inside &= (arg <= b) if b_included else arg < b + # TODO: add `supported` method and check here + on = np.zeros(a.shape, dtype=int) + on[a == arg] = -1 + on[b == arg] = 1 + outside = np.zeros(a.shape, dtype=int) + outside[(arg < a) if a_included else arg <= a] = -1 + outside[(b < arg) if b_included else b <= arg] = 1 + nan = np.isnan(arg) + + return inside, on, outside, nan + + +def test_input_validation(): + class Test(ContinuousDistribution): + _variable = _RealParameter('x', domain=_RealInterval()) + + message = ("The `Test` distribution family does not accept parameters, " + "but parameters `{'a'}` were provided.") + with pytest.raises(ValueError, match=message): + Test(a=1, ) + + message = "Attribute `tol` of `Test` must be a positive float, if specified." + with pytest.raises(ValueError, match=message): + Test(tol=np.asarray([])) + with pytest.raises(ValueError, match=message): + Test(tol=[1, 2, 3]) + with pytest.raises(ValueError, match=message): + Test(tol=np.nan) + with pytest.raises(ValueError, match=message): + Test(tol=-1) + + message = ("Argument `order` of `Test.moment` must be a " + "finite integer greater than or equal to 0.") + with pytest.raises(ValueError, match=message): + Test().moment(-1) + with pytest.raises(ValueError, match=message): + Test().moment(np.inf) + + message = ("Argument `order` of `Test.lmoment` must be a " + "finite integer greater than or equal to 1.") + with pytest.raises(ValueError, match=message): + Test().lmoment(0) + + message = "Argument `kind` of `Test.moment` must be one of..." + with pytest.raises(ValueError, match=message): + Test().moment(2, kind='coconut') + + class Test2(ContinuousDistribution): + _p1 = _RealParameter('c', domain=_RealInterval()) + _p2 = _RealParameter('d', domain=_RealInterval()) + _parameterizations = [_Parameterization(_p1, _p2)] + _variable = _RealParameter('x', domain=_RealInterval()) + + message = ("The provided parameters `{a}` do not match a supported " + "parameterization of the `Test2` distribution family.") + with pytest.raises(ValueError, match=message): + Test2(a=1) + + message = ("The `Test2` distribution family requires parameters, but none " + "were provided.") + with pytest.raises(ValueError, match=message): + Test2() + + message = ("The parameters `{c, d}` provided to the `Test2` " + "distribution family cannot be broadcast to the same shape.") + with pytest.raises(ValueError, match=message): + Test2(c=[1, 2], d=[1, 2, 3]) + + message = ("The argument provided to `Test2.pdf` cannot be be broadcast to " + "the same shape as the distribution parameters.") + with pytest.raises(ValueError, match=message): + dist = Test2(c=[1, 2, 3], d=[1, 2, 3]) + dist.pdf([1, 2]) + + message = "Parameter `c` must be of real dtype." + with pytest.raises(TypeError, match=message): + Test2(c=[1, object()], d=[1, 2]) + + message = "Parameter `convention` of `Test2.kurtosis` must be one of..." + with pytest.raises(ValueError, match=message): + dist = Test2(c=[1, 2, 3], d=[1, 2, 3]) + dist.kurtosis(convention='coconut') + + +def test_rng_deepcopy_pickle(): + # test behavior of `rng` attribute and copy behavior + kwargs = dict(a=[-1, 2], b=10) + dist1 = Uniform(**kwargs) + dist2 = deepcopy(dist1) + dist3 = pickle.loads(pickle.dumps(dist1)) + + res1, res2, res3 = dist1.sample(), dist2.sample(), dist3.sample() + assert np.all(res2 != res1) + assert np.all(res3 != res1) + + res1, res2, res3 = dist1.sample(rng=42), dist2.sample(rng=42), dist3.sample(rng=42) + assert np.all(res2 == res1) + assert np.all(res3 == res1) + + +class TestAttributes: + def test_cache_policy(self): + dist = StandardNormal(cache_policy="no_cache") + # make error message more appropriate + message = "`StandardNormal` does not provide an accurate implementation of the " + with pytest.raises(NotImplementedError, match=message): + dist.mean(method='cache') + mean = dist.mean() + with pytest.raises(NotImplementedError, match=message): + dist.mean(method='cache') + + # add to enum + dist.cache_policy = None + with pytest.raises(NotImplementedError, match=message): + dist.mean(method='cache') + mean = dist.mean() # method is 'formula' by default + cached_mean = dist.mean(method='cache') + assert_equal(cached_mean, mean) + + # cache is not cleared when used in a TransformedDistribution + dist2 = stats.exp(dist) + assert_equal(dist.mean(method='cache'), cached_mean) + dist2.reset_cache() + assert_equal(dist.mean(method='cache'), cached_mean) + + # cache is overridden by latest evaluation + quadrature_mean = dist.mean(method='quadrature') + cached_mean = dist.mean(method='cache') + assert_equal(cached_mean, quadrature_mean) + assert not np.all(mean == quadrature_mean) + + # We can turn the cache off, and it won't change, but the old cache is + # still available + dist.cache_policy = "no_cache" + mean = dist.mean(method='formula') + cached_mean = dist.mean(method='cache') + assert_equal(cached_mean, quadrature_mean) + assert not np.all(mean == quadrature_mean) + + dist.reset_cache() + with pytest.raises(NotImplementedError, match=message): + dist.mean(method='cache') + + message = "Attribute `cache_policy` of `StandardNormal`..." + with pytest.raises(ValueError, match=message): + dist.cache_policy = "invalid" + + def test_tol(self): + x = 3. + X = stats.Normal() + + message = "Attribute `tol` of `StandardNormal` must..." + with pytest.raises(ValueError, match=message): + X.tol = -1. + with pytest.raises(ValueError, match=message): + X.tol = (0.1,) + with pytest.raises(ValueError, match=message): + X.tol = np.nan + + X1 = stats.Normal(tol=1e-1) + X2 = stats.Normal(tol=1e-12) + ref = X.cdf(x) + res1 = X1.cdf(x, method='quadrature') + res2 = X2.cdf(x, method='quadrature') + assert_allclose(res1, ref, rtol=X1.tol) + assert_allclose(res2, ref, rtol=X2.tol) + assert abs(res1 - ref) > abs(res2 - ref) + + p = 0.99 + X1.tol, X2.tol = X2.tol, X1.tol + ref = X.icdf(p) + res1 = X1.icdf(p, method='inversion') + res2 = X2.icdf(p, method='inversion') + assert_allclose(res1, ref, rtol=X1.tol) + assert_allclose(res2, ref, rtol=X2.tol) + assert abs(res2 - ref) > abs(res1 - ref) + + def test_iv_policy(self): + X = Uniform(a=0, b=1) + assert X.pdf(2) == 0 + + X.validation_policy = 'skip_all' + assert X.pdf(np.asarray(2.)) == 1 + + # Tests _set_invalid_nan + a, b = np.asarray(1.), np.asarray(0.) # invalid parameters + X = Uniform(a=a, b=b, validation_policy='skip_all') + assert X.pdf(np.asarray(2.)) == -1 + + # Tests _set_invalid_nan_property + class MyUniform(Uniform): + def _entropy_formula(self, *args, **kwargs): + return 'incorrect' + + def _moment_raw_formula(self, order, **params): + return 'incorrect' + + X = MyUniform(a=a, b=b, validation_policy='skip_all') + assert X.entropy() == 'incorrect' + + # Tests _validate_order_kind + assert X.moment(kind='raw', order=-1) == 'incorrect' + + # Test input validation + message = "Attribute `validation_policy` of `MyUniform`..." + with pytest.raises(ValueError, match=message): + X.validation_policy = "invalid" + + def test_shapes(self): + X = stats.Normal(mu=1, sigma=2) + Y = stats.Normal(mu=[2], sigma=3) + + # Check that attributes are available as expected + assert X.mu == 1 + assert X.sigma == 2 + assert Y.mu[0] == 2 + assert Y.sigma[0] == 3 + + # Trying to set an attribute raises + # message depends on Python version + with pytest.raises(AttributeError): + X.mu = 2 + + # Trying to mutate an attribute really mutates a copy + Y.mu[0] = 10 + assert Y.mu[0] == 2 + + +class TestMakeDistribution: + @pytest.mark.parametrize('i, distdata', list(enumerate(distcont + distdiscrete))) + def test_rv_generic(self, i, distdata): + distname = distdata[0] + + slow = {'argus', 'exponpow', 'exponweib', 'genexpon', 'gompertz', + 'johnsonsb', 'kappa4', 'ksone', 'kstwo', 'kstwobign', 'norminvgauss', + 'powerlognorm', 'powernorm', 'recipinvgauss', 'studentized_range', + 'vonmises_line'} # continuous + if not int(os.environ.get('SCIPY_XSLOW', '0')) and distname in slow: + pytest.skip('Skipping as XSLOW') + + if distname in { # skip these distributions + 'levy_stable', # private methods seem to require >= 1d args + 'vonmises', # circular distribution; shouldn't work + 'poisson_binom', # vector shape parameter + 'hypergeom', # distribution functions need interpolation + 'nchypergeom_fisher', # distribution functions need interpolation + 'nchypergeom_wallenius', # distribution functions need interpolation + }: + return + + # skip single test, mostly due to slight disagreement + custom_tolerances = {'ksone': 1e-5, 'kstwo': 1e-5} # discontinuous PDF + skip_entropy = {'kstwobign', 'pearson3'} # tolerance issue + skip_skewness = {'exponpow', 'ksone', 'nchypergeom_wallenius'} # tolerance + skip_kurtosis = {'chi', 'exponpow', 'invgamma', # tolerance + 'johnsonsb', 'ksone', 'kstwo', # tolerance + 'nchypergeom_wallenius'} # tolerance + skip_logccdf = {'arcsine', 'skewcauchy', 'trapezoid', 'triang', # tolerance + 'dlaplace'} # slight tolerance issue, but only for array shape + skip_raw = {2: {'alpha', 'foldcauchy', 'halfcauchy', 'levy', 'levy_l'}, + 3: {'pareto'}, # stats.pareto is just wrong + 4: {'invgamma'}} # tolerance issue + skip_standardized = {'exponpow', 'ksone'} # tolerances + + dist = getattr(stats, distname) + params = dict(zip(dist.shapes.split(', '), distdata[1])) if dist.shapes else {} + rng = np.random.default_rng(7548723590230982) + CustomDistribution = stats.make_distribution(dist) + params = {key: np.asarray([val, val]) for key, val in params.items()} + X = CustomDistribution(**params) + Y = dist(**params) + x = X.sample(shape=10, rng=rng) + p = X.cdf(x) + rtol = custom_tolerances.get(distname, 1e-7) + atol = 1e-12 + + with np.errstate(divide='ignore', invalid='ignore'): + m, v, s, k = Y.stats('mvsk') + assert_allclose(X.support(), Y.support()) + if distname not in skip_entropy: + assert_allclose(X.entropy(), Y.entropy(), rtol=rtol) + if isinstance(Y, stats.rv_discrete): + # some continuous distributions have trouble with `logentropy` because + # it uses complex numbers + assert_allclose(np.exp(X.logentropy()), Y.entropy(), rtol=rtol) + assert_allclose(X.median(), Y.median(), rtol=rtol) + assert_allclose(X.mean(), m, rtol=rtol, atol=atol) + assert_allclose(X.variance(), v, rtol=rtol, atol=atol) + if distname not in skip_skewness: + assert_allclose(X.skewness(), s, rtol=rtol, atol=atol) + if distname not in skip_kurtosis: + assert_allclose(X.kurtosis(convention='excess'), k, + rtol=rtol, atol=atol) + if isinstance(dist, stats.rv_continuous): + assert_allclose(X.logpdf(x), Y.logpdf(x), rtol=rtol) + assert_allclose(X.pdf(x), Y.pdf(x), rtol=rtol) + else: + assert_allclose(X.logpmf(x), Y.logpmf(x), rtol=rtol) + assert_allclose(X.pmf(x), Y.pmf(x), rtol=rtol) + assert_allclose(X.logcdf(x), Y.logcdf(x), rtol=rtol) + assert_allclose(X.cdf(x), Y.cdf(x), rtol=rtol) + if distname not in skip_logccdf: + assert_allclose(X.logccdf(x), Y.logsf(x), rtol=rtol) + assert_allclose(X.ccdf(x), Y.sf(x), rtol=rtol) + + # old infrastructure convention for ppf(p=0) and isf(p=1) is different than + # new infrastructure. Adjust reference values accordingly. + a, _ = Y.support() + a, p = np.broadcast_arrays(a, p) + ref_ppf = Y.ppf(p) + ref_ppf[p == 0] = a[p == 0] + ref_isf = Y.isf(p) + ref_isf[p == 1] = a[p == 1] + + assert_allclose(X.icdf(p), ref_ppf, rtol=rtol) + assert_allclose(X.iccdf(p), ref_isf, rtol=rtol) + + for order in range(5): + if distname not in skip_raw.get(order, {}): + assert_allclose(X.moment(order, kind='raw'), + Y.moment(order), rtol=rtol, atol=atol) + for order in range(3, 4): + if distname not in skip_standardized: + assert_allclose(X.moment(order, kind='standardized'), + Y.stats('mvsk'[order-1]), rtol=rtol, atol=atol) + if isinstance(dist, stats.rv_continuous): + # For discrete distributions, these won't agree at the far left end + # of the support, and the new infrastructure is slow there (for now). + seed = 845298245687345 + assert_allclose(X.sample(shape=10, rng=seed), + Y.rvs(size=p.shape, + random_state=np.random.default_rng(seed)), + rtol=rtol) + + def test_custom(self): + rng = np.random.default_rng(7548723590230982) + + class MyLogUniform: + @property + def __make_distribution_version__(self): + return "1.16.0" + + @property + def parameters(self): + return {'a': {'endpoints': (0, np.inf), 'inclusive': (False, False)}, + 'b': {'endpoints': ('a', np.inf), 'inclusive': (False, False)}} + + @property + def support(self): + return {'endpoints': ('a', 'b')} + + def pdf(self, x, a, b): + return 1 / (x * (np.log(b) - np.log(a))) + + def sample(self, shape, *, a, b, rng=None): + p = rng.uniform(size=shape) + return np.exp(np.log(a) + p * (np.log(b) - np.log(a))) + + def moment(self, order, kind='raw', *, a, b): + if order == 1 and kind == 'raw': + # quadrature is perfectly accurate here; add 1e-10 error so we + # can tell the difference between the two + return (b - a) / np.log(b/a) + 1e-10 + + def lmoment(self, order, *, a, b): + s = np.log(b/a) + l1 = (b - a) / s + l2 = ((s - 2) * b + (s + 2) * a) / s**2 + l3 = ((s**2 - 6*s + 12) * b - (s**2 + 6*s + 12) * a) / s ** 3 + l4 = ((s**3 - 12*s**2 + 60*s - 120) * b + + (s**3 + 12*s**2 + 60*s + 120) * a) / s ** 4 + ls = {1: l1, 2: l2, 3: l3, 4: l4} + return ls.get(int(order), None) + + LogUniform = stats.make_distribution(MyLogUniform()) + + X = LogUniform(a=1., b=np.e) + Y = stats.exp(Uniform(a=0., b=1.)) + + assert str(X) == f"MyLogUniform(a=1.0, b={np.e})" + assert repr(X) == f"MyLogUniform(a=np.float64(1.0), b=np.float64({np.e}))" + + x = X.sample(shape=10, rng=rng) + p = X.cdf(x) + + assert_allclose(X.support(), Y.support()) + assert_allclose(X.entropy(), Y.entropy()) + assert_allclose(X.median(), Y.median()) + assert_allclose(X.logpdf(x), Y.logpdf(x)) + assert_allclose(X.pdf(x), Y.pdf(x)) + assert_allclose(X.logcdf(x), Y.logcdf(x)) + assert_allclose(X.cdf(x), Y.cdf(x)) + assert_allclose(X.logccdf(x), Y.logccdf(x)) + assert_allclose(X.ccdf(x), Y.ccdf(x)) + assert_allclose(X.icdf(p), Y.icdf(p)) + assert_allclose(X.iccdf(p), Y.iccdf(p)) + for kind in ['raw', 'central', 'standardized']: + for order in range(5): + assert_allclose(X.moment(order, kind=kind), + Y.moment(order, kind=kind)) + for standardize in [False, True]: + for order in range(1, 5): + if standardize and order < 3: + continue + assert_allclose(X.lmoment(order, standardize=standardize), + Y.lmoment(order, standardize=standardize)) + + # Confirm that the `sample` and `moment` methods are overriden as expected + sample_formula = X.sample(shape=10, rng=0, method='formula') + sample_inverse = X.sample(shape=10, rng=0, method='inverse_transform') + assert_allclose(sample_formula, sample_inverse) + assert not np.all(sample_formula == sample_inverse) + + assert_allclose(X.mean(method='formula'), X.mean(method='quadrature')) + assert not X.mean(method='formula') == X.mean(method='quadrature') + + assert_allclose(X.lmoment(method='formula'), + X.lmoment(method='quadrature_icdf')) + assert not X.lmoment(method='formula') == X.lmoment(method='quadrature_icdf') + + # pdf and cdf formulas below can warn on boundary of support in some cases. + # See https://github.com/scipy/scipy/pull/22560#discussion_r1962763840. + @pytest.mark.slow + @pytest.mark.filterwarnings("ignore::RuntimeWarning") + @pytest.mark.parametrize("c", [-1, 0, 1, np.asarray([-2.1, -1., 0., 1., 2.1])]) + def test_custom_variable_support(self, c): + rng = np.random.default_rng(7548723590230982) + + class MyGenExtreme: + @property + def __make_distribution_version__(self): + return "1.16.0" + + @property + def parameters(self): + return { + 'c': {'endpoints': (-np.inf, np.inf), 'inclusive': (False, False)}, + 'mu': {'endpoints': (-np.inf, np.inf), 'inclusive': (False, False)}, + 'sigma': {'endpoints': (0, np.inf), 'inclusive': (False, False)} + } + + @property + def support(self): + def left(*, c, mu, sigma): + c, mu, sigma = np.broadcast_arrays(c, mu, sigma) + result = np.empty_like(c) + result[c >= 0] = -np.inf + result[c < 0] = mu[c < 0] + sigma[c < 0] / c[c < 0] + return result[()] + + def right(*, c, mu, sigma): + c, mu, sigma = np.broadcast_arrays(c, mu, sigma) + result = np.empty_like(c) + result[c <= 0] = np.inf + result[c > 0] = mu[c > 0] + sigma[c > 0] / c[c > 0] + return result[()] + + return {"endpoints": (left, right), "inclusive": (False, False)} + + def pdf(self, x, *, c, mu, sigma): + x, c, mu, sigma = np.broadcast_arrays(x, c, mu, sigma) + t = np.empty_like(x) + mask = (c == 0) + t[mask] = np.exp(-(x[mask] - mu[mask])/sigma[mask]) + t[~mask] = ( + 1 - c[~mask]*(x[~mask] - mu[~mask])/sigma[~mask] + )**(1/c[~mask]) + result = 1/sigma * t**(1 - c)*np.exp(-t) + return result[()] + + def cdf(self, x, *, c, mu, sigma): + x, c, mu, sigma = np.broadcast_arrays(x, c, mu, sigma) + t = np.empty_like(x) + mask = (c == 0) + t[mask] = np.exp(-(x[mask] - mu[mask])/sigma[mask]) + t[~mask] = ( + 1 - c[~mask]*(x[~mask] - mu[~mask])/sigma[~mask] + )**(1/c[~mask]) + return np.exp(-t)[()] + + GenExtreme1 = stats.make_distribution(MyGenExtreme()) + GenExtreme2 = stats.make_distribution(stats.genextreme) + + X1 = GenExtreme1(c=c, mu=0, sigma=1) + X2 = GenExtreme2(c=c) + + x = X1.sample(shape=10, rng=rng) + p = X1.cdf(x) + + assert_allclose(X1.support(), X2.support()) + assert_allclose(X1.entropy(), X2.entropy(), rtol=5e-6) + assert_allclose(X1.median(), X2.median()) + assert_allclose(X1.logpdf(x), X2.logpdf(x)) + assert_allclose(X1.pdf(x), X2.pdf(x)) + assert_allclose(X1.logcdf(x), X2.logcdf(x)) + assert_allclose(X1.cdf(x), X2.cdf(x)) + assert_allclose(X1.logccdf(x), X2.logccdf(x)) + assert_allclose(X1.ccdf(x), X2.ccdf(x)) + assert_allclose(X1.icdf(p), X2.icdf(p)) + assert_allclose(X1.iccdf(p), X2.iccdf(p)) + + @pytest.mark.slow + @pytest.mark.parametrize("a", [0.5, np.asarray([0.5, 1.0, 2.0, 4.0, 8.0])]) + @pytest.mark.parametrize("b", [0.5, np.asarray([0.5, 1.0, 2.0, 4.0, 8.0])]) + def test_custom_multiple_parameterizations(self, a, b): + rng = np.random.default_rng(7548723590230982) + class MyBeta: + @property + def __make_distribution_version__(self): + return "1.16.0" + + @property + def parameters(self): + return ( + {"a": (0, np.inf), "b": (0, np.inf)}, + {"mu": (0, 1), "nu": (0, np.inf)}, + ) + + def process_parameters(self, a=None, b=None, mu=None, nu=None): + if a is not None and b is not None and mu is None and nu is None: + nu = a + b + mu = a / nu + else: + a = mu * nu + b = nu - a + return {"a": a, "b": b, "mu": mu, "nu": nu} + + @property + def support(self): + return {'endpoints': (0, 1)} + + def pdf(self, x, a, b, mu, nu): + return special._ufuncs._beta_pdf(x, a, b) + + def cdf(self, x, a, b, mu, nu): + return special.betainc(a, b, x) + + Beta = stats.make_distribution(stats.beta) + MyBeta = stats.make_distribution(MyBeta()) + + mu = a / (a + b) + nu = a + b + + X = MyBeta(a=a, b=b) + Y = MyBeta(mu=mu, nu=nu) + Z = Beta(a=a, b=b) + + x = Z.sample(shape=10, rng=rng) + p = Z.cdf(x) + + assert_allclose(X.support(), Z.support()) + assert_allclose(X.median(), Z.median()) + assert_allclose(X.pdf(x), Z.pdf(x)) + assert_allclose(X.cdf(x), Z.cdf(x)) + assert_allclose(X.ccdf(x), Z.ccdf(x)) + assert_allclose(X.icdf(p), Z.icdf(p)) + assert_allclose(X.iccdf(p), Z.iccdf(p)) + + assert_allclose(Y.support(), Z.support()) + assert_allclose(Y.median(), Z.median()) + assert_allclose(Y.pdf(x), Z.pdf(x)) + assert_allclose(Y.cdf(x), Z.cdf(x)) + assert_allclose(Y.ccdf(x), Z.ccdf(x)) + assert_allclose(Y.icdf(p), Z.icdf(p)) + assert_allclose(Y.iccdf(p), Z.iccdf(p)) + + def test_input_validation(self): + message = '`levy_stable` is not supported.' + with pytest.raises(NotImplementedError, match=message): + stats.make_distribution(stats.levy_stable) + + message = '`vonmises` is not supported.' + with pytest.raises(NotImplementedError, match=message): + stats.make_distribution(stats.vonmises) + + message = "The argument must be an instance of..." + with pytest.raises(ValueError, match=message): + stats.make_distribution(object()) + + def test_repr_str_docs(self): + from scipy.stats._distribution_infrastructure import _distribution_names + for dist in _distribution_names.keys(): + assert hasattr(stats, dist) + + dist = stats.make_distribution(stats.gamma) + assert str(dist(a=2)) == "Gamma(a=2.0)" + assert repr(dist(a=2)) == "Gamma(a=np.float64(2.0))" + assert 'Gamma' in dist.__doc__ + + dist = stats.make_distribution(stats.halfgennorm) + assert str(dist(beta=2)) == "HalfGeneralizedNormal(beta=2.0)" + assert repr(dist(beta=2)) == "HalfGeneralizedNormal(beta=np.float64(2.0))" + assert 'HalfGeneralizedNormal' in dist.__doc__ + + @pytest.mark.slow # just in case + @settings(max_examples=20) # no need for more + @given(data=strategies.data()) + def test_draw_distribution(self, data): + # `draw_distribution_from_family` is a private function right now, but we will + # want that functionality to be public someday. It was broken for custom + # distributions because the `typical` parameter of the support was ignored. + # Check that this is resolved. + rng = np.random.default_rng(8465652168548465121) + u_typical = tuple(np.sort(rng.standard_normal(2))) + s_typical = tuple(np.sort(rng.random(2)*2)) + x_typical = tuple(np.sort(rng.standard_normal(2))) + + class MyNormal: + __make_distribution_version__ = "1.16.0" + parameters = {'u': {'endpoints': (-np.inf, np.inf), 'typical': u_typical}, + 's': {'endpoints': (0, np.inf), 'typical': s_typical}} + support = {'endpoints': (-np.inf, np.inf), 'typical': x_typical} + + def pdf(self, x, u, s): + return 1 / np.sqrt(2*np.pi) / s * np.exp(-((x-u)/s)**2/2) + + family = stats.make_distribution(MyNormal()) + proportions = (1.0, 0., 0., 0.) + tmp = draw_distribution_from_family(family, data, rng, proportions, min_side=1) + dist, x, y, p, logp, result_shape, x_result_shape, xy_result_shape = tmp + assert u_typical[0] < np.min(dist.u) and np.max(dist.u) < u_typical[1] + assert s_typical[0] < np.min(dist.s) and np.max(dist.s) < s_typical[1] + assert x_typical[0] < np.min(x) and np.max(x) < x_typical[1] + + @pytest.mark.parametrize('p', [None, ()]) + def test_no_parameters(self, p): + # To create a distribution without parameters, it is natural to try an empty + # dictionary (since a dictionary with entries is valid if there is only + # one parameterization), but that didn't work; only an empty tuple (no + # parameterizations) worked originally. Check that this is resolved. + class MyStandardNormal: + __make_distribution_version__ = "1.16.0" + parameters = {} if p is None else p + support = {'endpoints': (-np.inf, np.inf)} + + def pdf(self, x): + return 1 / np.sqrt(2*np.pi) * np.exp(-x**2/2) + + StandardNormal = stats.make_distribution(MyStandardNormal()) + X = StandardNormal() + assert X.support() == (-np.inf, np.inf) + + +class TestTransforms: + + def test_ContinuousDistribution_only(self): + X = stats.Binomial(n=10, p=0.5) + # This is applied at the top level TransformedDistribution, + # so testing one subclass is enough + message = "Transformations are currently only supported for continuous RVs." + with pytest.raises(NotImplementedError, match=message): + stats.exp(X) + + def test_truncate(self): + rng = np.random.default_rng(81345982345826) + lb = rng.random((3, 1)) + ub = rng.random((3, 1)) + lb, ub = np.minimum(lb, ub), np.maximum(lb, ub) + + Y = stats.truncate(Normal(), lb=lb, ub=ub) + Y0 = stats.truncnorm(lb, ub) + + y = Y0.rvs((3, 10), random_state=rng) + p = Y0.cdf(y) + + assert_allclose(Y.logentropy(), np.log(Y0.entropy() + 0j)) + assert_allclose(Y.entropy(), Y0.entropy()) + assert_allclose(Y.median(), Y0.ppf(0.5)) + assert_allclose(Y.mean(), Y0.mean()) + assert_allclose(Y.variance(), Y0.var()) + assert_allclose(Y.standard_deviation(), np.sqrt(Y0.var())) + assert_allclose(Y.skewness(), Y0.stats('s')) + assert_allclose(Y.kurtosis(), Y0.stats('k') + 3) + assert_allclose(Y.support(), Y0.support()) + assert_allclose(Y.pdf(y), Y0.pdf(y)) + assert_allclose(Y.cdf(y), Y0.cdf(y)) + assert_allclose(Y.ccdf(y), Y0.sf(y)) + assert_allclose(Y.icdf(p), Y0.ppf(p)) + assert_allclose(Y.iccdf(p), Y0.isf(p)) + assert_allclose(Y.logpdf(y), Y0.logpdf(y)) + assert_allclose(Y.logcdf(y), Y0.logcdf(y)) + assert_allclose(Y.logccdf(y), Y0.logsf(y)) + assert_allclose(Y.ilogcdf(np.log(p)), Y0.ppf(p)) + assert_allclose(Y.ilogccdf(np.log(p)), Y0.isf(p)) + sample = Y.sample(10) + assert np.all((sample > lb) & (sample < ub)) + + @pytest.mark.fail_slow(10) + @given(data=strategies.data(), seed=strategies.integers(min_value=0)) + def test_loc_scale(self, data, seed): + # Need tests with negative scale + rng = np.random.default_rng(seed) + + class TransformedNormal(ShiftedScaledDistribution): + def __init__(self, *args, **kwargs): + super().__init__(StandardNormal(), *args, **kwargs) + + tmp = draw_distribution_from_family( + TransformedNormal, data, rng, proportions=(1, 0, 0, 0), min_side=1) + dist, x, y, p, logp, result_shape, x_result_shape, xy_result_shape = tmp + + loc = dist.loc + scale = dist.scale + dist0 = StandardNormal() + dist_ref = stats.norm(loc=loc, scale=scale) + dist_ref_lmoment = Normal(mu=loc, sigma=scale) + + x0 = (x - loc) / scale + y0 = (y - loc) / scale + + a, b = dist.support() + a0, b0 = dist0.support() + assert_allclose(a, a0 + loc) + assert_allclose(b, b0 + loc) + + with np.errstate(invalid='ignore', divide='ignore'): + assert_allclose(np.exp(dist.logentropy()), dist.entropy()) + assert_allclose(dist.entropy(), dist_ref.entropy()) + assert_allclose(dist.median(), dist0.median() + loc) + assert_allclose(dist.mode(), dist0.mode() + loc) + assert_allclose(dist.mean(), dist0.mean() + loc) + assert_allclose(dist.variance(), dist0.variance() * scale**2) + assert_allclose(dist.standard_deviation(), dist.variance()**0.5) + assert_allclose(dist.skewness(), dist0.skewness() * np.sign(scale)) + assert_allclose(dist.kurtosis(), dist0.kurtosis()) + assert_allclose(dist.logpdf(x), dist0.logpdf(x0) - np.log(scale)) + assert_allclose(dist.pdf(x), dist0.pdf(x0) / scale) + assert_allclose(dist.logcdf(x), dist0.logcdf(x0)) + assert_allclose(dist.cdf(x), dist0.cdf(x0)) + assert_allclose(dist.logccdf(x), dist0.logccdf(x0)) + assert_allclose(dist.ccdf(x), dist0.ccdf(x0)) + assert_allclose(dist.logcdf(x, y), dist0.logcdf(x0, y0)) + assert_allclose(dist.cdf(x, y), dist0.cdf(x0, y0)) + assert_allclose(dist.logccdf(x, y), dist0.logccdf(x0, y0)) + assert_allclose(dist.ccdf(x, y), dist0.ccdf(x0, y0)) + assert_allclose(dist.ilogcdf(logp), dist0.ilogcdf(logp)*scale + loc) + assert_allclose(dist.icdf(p), dist0.icdf(p)*scale + loc) + assert_allclose(dist.ilogccdf(logp), dist0.ilogccdf(logp)*scale + loc) + assert_allclose(dist.iccdf(p), dist0.iccdf(p)*scale + loc) + for i in range(1, 5): + assert_allclose(dist.moment(i, 'raw'), dist_ref.moment(i)) + assert_allclose(dist.moment(i, 'central'), + dist0.moment(i, 'central') * scale**i) + assert_allclose(dist.moment(i, 'standardized'), + dist0.moment(i, 'standardized') * np.sign(scale)**i) + for i in range(1, 5): + assert_allclose(dist.lmoment(i), dist_ref_lmoment.lmoment(i), atol=1e-8) + if i >= 3: + assert_allclose(dist.lmoment(i, standardize=True), + dist_ref_lmoment.lmoment(i, standardize=True), + atol=1e-8) + + # Transform back to the original distribution using all arithmetic + # operations; check that it behaves as expected. + dist = (dist - 2*loc) + loc + dist = dist/scale**2 * scale + z = np.zeros(dist._shape) # compact broadcasting + + a, b = dist.support() + a0, b0 = dist0.support() + assert_allclose(a, a0 + z) + assert_allclose(b, b0 + z) + + with np.errstate(invalid='ignore', divide='ignore'): + assert_allclose(dist.logentropy(), dist0.logentropy() + z) + assert_allclose(dist.entropy(), dist0.entropy() + z) + assert_allclose(dist.median(), dist0.median() + z) + assert_allclose(dist.mode(), dist0.mode() + z) + assert_allclose(dist.mean(), dist0.mean() + z) + assert_allclose(dist.variance(), dist0.variance() + z) + assert_allclose(dist.standard_deviation(), dist0.standard_deviation() + z) + assert_allclose(dist.skewness(), dist0.skewness() + z) + assert_allclose(dist.kurtosis(), dist0.kurtosis() + z) + assert_allclose(dist.logpdf(x), dist0.logpdf(x)+z) + assert_allclose(dist.pdf(x), dist0.pdf(x) + z) + assert_allclose(dist.logcdf(x), dist0.logcdf(x) + z) + assert_allclose(dist.cdf(x), dist0.cdf(x) + z) + assert_allclose(dist.logccdf(x), dist0.logccdf(x) + z) + assert_allclose(dist.ccdf(x), dist0.ccdf(x) + z) + assert_allclose(dist.ilogcdf(logp), dist0.ilogcdf(logp) + z) + assert_allclose(dist.icdf(p), dist0.icdf(p) + z) + assert_allclose(dist.ilogccdf(logp), dist0.ilogccdf(logp) + z) + assert_allclose(dist.iccdf(p), dist0.iccdf(p) + z) + for i in range(1, 5): + assert_allclose(dist.moment(i, 'raw'), dist0.moment(i, 'raw')) + assert_allclose(dist.moment(i, 'central'), dist0.moment(i, 'central')) + assert_allclose(dist.moment(i, 'standardized'), + dist0.moment(i, 'standardized')) + for i in range(1, 5): + assert_allclose(dist.lmoment(i), dist0.lmoment(i)) + if i >= 3: + assert_allclose(dist.lmoment(i, standardize=True), + dist0.lmoment(i, standardize=True)) + + # These are tough to compare because of the way the shape works + # rng = np.random.default_rng(seed) + # rng0 = np.random.default_rng(seed) + # assert_allclose(dist.sample(x_result_shape, rng=rng), + # dist0.sample(x_result_shape, rng=rng0) * scale + loc) + # Should also try to test fit, plot? + + @pytest.mark.fail_slow(5) + @pytest.mark.parametrize('exp_pow', ['exp', 'pow']) + def test_exp_pow(self, exp_pow): + rng = np.random.default_rng(81345982345826) + mu = rng.random((3, 1)) + sigma = rng.random((3, 1)) + + X = Normal()*sigma + mu + if exp_pow == 'exp': + Y = stats.exp(X) + else: + Y = np.e ** X + Y0 = stats.lognorm(sigma, scale=np.exp(mu)) + + y = Y0.rvs((3, 10), random_state=rng) + p = Y0.cdf(y) + + assert_allclose(Y.logentropy(), np.log(Y0.entropy())) + assert_allclose(Y.entropy(), Y0.entropy()) + assert_allclose(Y.median(), Y0.ppf(0.5)) + assert_allclose(Y.mean(), Y0.mean()) + assert_allclose(Y.variance(), Y0.var()) + assert_allclose(Y.standard_deviation(), np.sqrt(Y0.var())) + assert_allclose(Y.skewness(), Y0.stats('s')) + assert_allclose(Y.kurtosis(), Y0.stats('k') + 3) + assert_allclose(Y.support(), Y0.support()) + assert_allclose(Y.pdf(y), Y0.pdf(y)) + assert_allclose(Y.cdf(y), Y0.cdf(y)) + assert_allclose(Y.ccdf(y), Y0.sf(y)) + assert_allclose(Y.icdf(p), Y0.ppf(p)) + assert_allclose(Y.iccdf(p), Y0.isf(p)) + assert_allclose(Y.logpdf(y), Y0.logpdf(y)) + assert_allclose(Y.logcdf(y), Y0.logcdf(y)) + assert_allclose(Y.logccdf(y), Y0.logsf(y)) + assert_allclose(Y.ilogcdf(np.log(p)), Y0.ppf(p)) + assert_allclose(Y.ilogccdf(np.log(p)), Y0.isf(p)) + seed = 3984593485 + assert_allclose(Y.sample(rng=seed), np.exp(X.sample(rng=seed))) + + + @pytest.mark.fail_slow(10) + @pytest.mark.parametrize('scale', [1, 2, -1]) + @pytest.mark.xfail_on_32bit("`scale=-1` fails on 32-bit; needs investigation") + def test_reciprocal(self, scale): + rng = np.random.default_rng(81345982345826) + a = rng.random((3, 1)) + + # Separate sign from scale. It's easy to scale the resulting + # RV with negative scale; we want to test the ability to divide + # by a RV with negative support + sign, scale = np.sign(scale), abs(scale) + + # Reference distribution + InvGamma = stats.make_distribution(stats.invgamma) + Y0 = sign * scale * InvGamma(a=a) + + # Test distribution + X = _Gamma(a=a) if sign > 0 else -_Gamma(a=a) + Y = scale / X + + y = Y0.sample(shape=(3, 10), rng=rng) + p = Y0.cdf(y) + logp = np.log(p) + + assert_allclose(Y.logentropy(), np.log(Y0.entropy())) + assert_allclose(Y.entropy(), Y0.entropy()) + assert_allclose(Y.median(), Y0.median()) + # moments are not finite + assert_allclose(Y.support(), Y0.support()) + assert_allclose(Y.pdf(y), Y0.pdf(y)) + assert_allclose(Y.cdf(y), Y0.cdf(y)) + assert_allclose(Y.ccdf(y), Y0.ccdf(y)) + assert_allclose(Y.icdf(p), Y0.icdf(p)) + assert_allclose(Y.iccdf(p), Y0.iccdf(p)) + assert_allclose(Y.logpdf(y), Y0.logpdf(y)) + assert_allclose(Y.logcdf(y), Y0.logcdf(y)) + assert_allclose(Y.logccdf(y), Y0.logccdf(y)) + with np.errstate(divide='ignore', invalid='ignore'): + assert_allclose(Y.ilogcdf(logp), Y0.ilogcdf(logp)) + assert_allclose(Y.ilogccdf(logp), Y0.ilogccdf(logp)) + seed = 3984593485 + assert_allclose(Y.sample(rng=seed), scale/(X.sample(rng=seed))) + + @pytest.mark.fail_slow(5) + def test_log(self): + rng = np.random.default_rng(81345982345826) + a = rng.random((3, 1)) + + X = _Gamma(a=a) + Y0 = stats.loggamma(a) + Y = stats.log(X) + y = Y0.rvs((3, 10), random_state=rng) + p = Y0.cdf(y) + + assert_allclose(Y.logentropy(), np.log(Y0.entropy())) + assert_allclose(Y.entropy(), Y0.entropy()) + assert_allclose(Y.median(), Y0.ppf(0.5)) + assert_allclose(Y.mean(), Y0.mean()) + assert_allclose(Y.variance(), Y0.var()) + assert_allclose(Y.standard_deviation(), np.sqrt(Y0.var())) + assert_allclose(Y.skewness(), Y0.stats('s')) + assert_allclose(Y.kurtosis(), Y0.stats('k') + 3) + assert_allclose(Y.support(), Y0.support()) + assert_allclose(Y.pdf(y), Y0.pdf(y)) + assert_allclose(Y.cdf(y), Y0.cdf(y)) + assert_allclose(Y.ccdf(y), Y0.sf(y)) + assert_allclose(Y.icdf(p), Y0.ppf(p)) + assert_allclose(Y.iccdf(p), Y0.isf(p)) + assert_allclose(Y.logpdf(y), Y0.logpdf(y)) + assert_allclose(Y.logcdf(y), Y0.logcdf(y)) + assert_allclose(Y.logccdf(y), Y0.logsf(y)) + with np.errstate(invalid='ignore'): + assert_allclose(Y.ilogcdf(np.log(p)), Y0.ppf(p)) + assert_allclose(Y.ilogccdf(np.log(p)), Y0.isf(p)) + seed = 3984593485 + assert_allclose(Y.sample(rng=seed), np.log(X.sample(rng=seed))) + + def test_monotonic_transforms(self): + # Some tests of monotonic transforms that are better to be grouped or + # don't fit well above + + X = Uniform(a=1, b=2) + X_str = "Uniform(a=1.0, b=2.0)" + + assert str(stats.log(X)) == f"log({X_str})" + assert str(1 / X) == f"1/({X_str})" + assert str(stats.exp(X)) == f"exp({X_str})" + + X = Uniform(a=-1, b=2) + message = "Division by a random variable is only implemented when the..." + with pytest.raises(NotImplementedError, match=message): + 1 / X + message = "The logarithm of a random variable is only implemented when the..." + with pytest.raises(NotImplementedError, match=message): + stats.log(X) + message = "Raising an argument to the power of a random variable is only..." + with pytest.raises(NotImplementedError, match=message): + (-2) ** X + with pytest.raises(NotImplementedError, match=message): + 1 ** X + with pytest.raises(NotImplementedError, match=message): + [0.5, 1.5] ** X + + message = "Raising a random variable to the power of an argument is only" + with pytest.raises(NotImplementedError, match=message): + X ** (-2) + with pytest.raises(NotImplementedError, match=message): + X ** 0 + with pytest.raises(NotImplementedError, match=message): + X ** [0.5, 1.5] + + def test_arithmetic_operators(self): + rng = np.random.default_rng(2348923495832349834) + + a, b, loc, scale = 0.294, 1.34, 0.57, 1.16 + + x = rng.uniform(-3, 3, 100) + Y = _LogUniform(a=a, b=b) + + X = scale*Y + loc + assert_allclose(X.cdf(x), Y.cdf((x - loc) / scale)) + X = loc + Y*scale + assert_allclose(X.cdf(x), Y.cdf((x - loc) / scale)) + + X = Y/scale - loc + assert_allclose(X.cdf(x), Y.cdf((x + loc) * scale)) + X = loc -_LogUniform(a=a, b=b)/scale + assert_allclose(X.cdf(x), Y.ccdf((-x + loc)*scale)) + + def test_abs(self): + rng = np.random.default_rng(81345982345826) + loc = rng.random((3, 1)) + + Y = stats.abs(Normal() + loc) + Y0 = stats.foldnorm(loc) + + y = Y0.rvs((3, 10), random_state=rng) + p = Y0.cdf(y) + + assert_allclose(Y.logentropy(), np.log(Y0.entropy() + 0j)) + assert_allclose(Y.entropy(), Y0.entropy()) + assert_allclose(Y.median(), Y0.ppf(0.5)) + assert_allclose(Y.mean(), Y0.mean()) + assert_allclose(Y.variance(), Y0.var()) + assert_allclose(Y.standard_deviation(), np.sqrt(Y0.var())) + assert_allclose(Y.skewness(), Y0.stats('s')) + assert_allclose(Y.kurtosis(), Y0.stats('k') + 3) + assert_allclose(Y.support(), Y0.support()) + assert_allclose(Y.pdf(y), Y0.pdf(y)) + assert_allclose(Y.cdf(y), Y0.cdf(y)) + assert_allclose(Y.ccdf(y), Y0.sf(y)) + assert_allclose(Y.icdf(p), Y0.ppf(p)) + assert_allclose(Y.iccdf(p), Y0.isf(p)) + assert_allclose(Y.logpdf(y), Y0.logpdf(y)) + assert_allclose(Y.logcdf(y), Y0.logcdf(y)) + assert_allclose(Y.logccdf(y), Y0.logsf(y)) + assert_allclose(Y.ilogcdf(np.log(p)), Y0.ppf(p)) + assert_allclose(Y.ilogccdf(np.log(p)), Y0.isf(p)) + sample = Y.sample(10) + assert np.all(sample > 0) + + def test_abs_finite_support(self): + # The original implementation of `FoldedDistribution` might evaluate + # the private distribution methods outside the support. Check that this + # is resolved. + Weibull = stats.make_distribution(stats.weibull_min) + X = Weibull(c=2) + Y = abs(-X) + assert_equal(X.logpdf(1), Y.logpdf(1)) + assert_equal(X.pdf(1), Y.pdf(1)) + assert_equal(X.logcdf(1), Y.logcdf(1)) + assert_equal(X.cdf(1), Y.cdf(1)) + assert_equal(X.logccdf(1), Y.logccdf(1)) + assert_equal(X.ccdf(1), Y.ccdf(1)) + + def test_pow(self): + rng = np.random.default_rng(81345982345826) + + Y = Normal()**2 + Y0 = stats.chi2(df=1) + + y = Y0.rvs(10, random_state=rng) + p = Y0.cdf(y) + + assert_allclose(Y.logentropy(), np.log(Y0.entropy() + 0j), rtol=1e-6) + assert_allclose(Y.entropy(), Y0.entropy(), rtol=1e-6) + assert_allclose(Y.median(), Y0.median()) + assert_allclose(Y.mean(), Y0.mean()) + assert_allclose(Y.variance(), Y0.var()) + assert_allclose(Y.standard_deviation(), np.sqrt(Y0.var())) + assert_allclose(Y.skewness(), Y0.stats('s')) + assert_allclose(Y.kurtosis(), Y0.stats('k') + 3) + assert_allclose(Y.support(), Y0.support()) + assert_allclose(Y.pdf(y), Y0.pdf(y)) + assert_allclose(Y.cdf(y), Y0.cdf(y)) + assert_allclose(Y.ccdf(y), Y0.sf(y)) + assert_allclose(Y.icdf(p), Y0.ppf(p)) + assert_allclose(Y.iccdf(p), Y0.isf(p)) + assert_allclose(Y.logpdf(y), Y0.logpdf(y)) + assert_allclose(Y.logcdf(y), Y0.logcdf(y)) + assert_allclose(Y.logccdf(y), Y0.logsf(y)) + assert_allclose(Y.ilogcdf(np.log(p)), Y0.ppf(p)) + assert_allclose(Y.ilogccdf(np.log(p)), Y0.isf(p)) + sample = Y.sample(10) + assert np.all(sample > 0) + +class TestOrderStatistic: + @pytest.mark.fail_slow(20) # Moments require integration + def test_order_statistic(self): + rng = np.random.default_rng(7546349802439582) + X = Uniform(a=0, b=1) + n = 5 + r = np.asarray([[1], [3], [5]]) + Y = stats.order_statistic(X, n=n, r=r) + Y0 = stats.beta(r, n + 1 - r) + + y = Y0.rvs((3, 10), random_state=rng) + p = Y0.cdf(y) + + # log methods need some attention before merge + assert_allclose(np.exp(Y.logentropy()), Y0.entropy()) + assert_allclose(Y.entropy(), Y0.entropy()) + assert_allclose(Y.mean(), Y0.mean()) + assert_allclose(Y.variance(), Y0.var()) + assert_allclose(Y.skewness(), Y0.stats('s'), atol=1e-15) + assert_allclose(Y.kurtosis(), Y0.stats('k') + 3, atol=1e-15) + assert_allclose(Y.median(), Y0.ppf(0.5)) + assert_allclose(Y.support(), Y0.support()) + assert_allclose(Y.pdf(y), Y0.pdf(y)) + assert_allclose(Y.cdf(y, method='formula'), Y.cdf(y, method='quadrature')) + assert_allclose(Y.ccdf(y, method='formula'), Y.ccdf(y, method='quadrature')) + assert_allclose(Y.icdf(p, method='formula'), Y.icdf(p, method='inversion')) + assert_allclose(Y.iccdf(p, method='formula'), Y.iccdf(p, method='inversion')) + assert_allclose(Y.logpdf(y), Y0.logpdf(y)) + assert_allclose(Y.logcdf(y), Y0.logcdf(y)) + assert_allclose(Y.logccdf(y), Y0.logsf(y)) + with np.errstate(invalid='ignore', divide='ignore'): + assert_allclose(Y.ilogcdf(np.log(p),), Y0.ppf(p)) + assert_allclose(Y.ilogccdf(np.log(p)), Y0.isf(p)) + + message = "`r` and `n` must contain only positive integers." + with pytest.raises(ValueError, match=message): + stats.order_statistic(X, n=n, r=-1) + with pytest.raises(ValueError, match=message): + stats.order_statistic(X, n=-1, r=r) + with pytest.raises(ValueError, match=message): + stats.order_statistic(X, n=n, r=1.5) + with pytest.raises(ValueError, match=message): + stats.order_statistic(X, n=1.5, r=r) + + def test_support_gh22037(self): + # During review of gh-22037, it was noted that the `support` of + # an `OrderStatisticDistribution` returned incorrect results; + # this was resolved by overriding `_support`. + Uniform = stats.make_distribution(stats.uniform) + X = Uniform() + Y = X*5 + 2 + Z = stats.order_statistic(Y, r=3, n=5) + assert_allclose(Z.support(), Y.support()) + + def test_composition_gh22037(self): + # During review of gh-22037, it was noted that an error was + # raised when creating an `OrderStatisticDistribution` from + # a `TruncatedDistribution`. This was resolved by overriding + # `_update_parameters`. + Normal = stats.make_distribution(stats.norm) + TruncatedNormal = stats.make_distribution(stats.truncnorm) + a, b = [-2, -1], 1 + r, n = 3, [[4], [5]] + x = [[[-0.3]], [[0.1]]] + X1 = Normal() + Y1 = stats.truncate(X1, a, b) + Z1 = stats.order_statistic(Y1, r=r, n=n) + X2 = TruncatedNormal(a=a, b=b) + Z2 = stats.order_statistic(X2, r=r, n=n) + np.testing.assert_allclose(Z1.cdf(x), Z2.cdf(x)) + + +class TestFullCoverage: + # Adds tests just to get to 100% test coverage; this way it's more obvious + # if new lines are untested. + def test_Domain(self): + with pytest.raises(NotImplementedError): + _Domain.contains(None, 1.) + with pytest.raises(NotImplementedError): + _Domain.get_numerical_endpoints(None, 1.) + with pytest.raises(NotImplementedError): + _Domain.__str__(None) + + def test_Parameter(self): + with pytest.raises(NotImplementedError): + _Parameter.validate(None, 1.) + + @pytest.mark.parametrize(("dtype_in", "dtype_out"), + [(np.float16, np.float16), + (np.int16, np.float64)]) + def test_RealParameter_uncommon_dtypes(self, dtype_in, dtype_out): + domain = _RealInterval((-1, 1)) + parameter = _RealParameter('x', domain=domain) + + x = np.asarray([0.5, 2.5], dtype=dtype_in) + arr, dtype, valid = parameter.validate(x, parameter_values={}) + assert_equal(arr, x) + assert dtype == dtype_out + assert_equal(valid, [True, False]) + + def test_ContinuousDistribution_set_invalid_nan(self): + # Exercise code paths when formula returns wrong shape and dtype + # We could consider making this raise an error to force authors + # to return the right shape and dytpe, but this would need to be + # configurable. + class TestDist(ContinuousDistribution): + _variable = _RealParameter('x', domain=_RealInterval(endpoints=(0., 1.))) + def _logpdf_formula(self, x, *args, **kwargs): + return 0 + + X = TestDist() + dtype = np.float32 + X._dtype = dtype + x = np.asarray([0.5], dtype=dtype) + assert X.logpdf(x).dtype == dtype + + def test_fiinfo(self): + assert _fiinfo(np.float64(1.)).max == np.finfo(np.float64).max + assert _fiinfo(np.int64(1)).max == np.iinfo(np.int64).max + + def test_generate_domain_support(self): + msg = _generate_domain_support(StandardNormal) + assert "accepts no distribution parameters" in msg + + msg = _generate_domain_support(Normal) + assert "accepts one parameterization" in msg + + msg = _generate_domain_support(_LogUniform) + assert "accepts two parameterizations" in msg + + def test_ContinuousDistribution__repr__(self): + X = Uniform(a=0, b=1) + assert repr(X) == "Uniform(a=np.float64(0.0), b=np.float64(1.0))" + assert repr(X*3 + 2) == ( + "np.float64(3.0)*Uniform(a=np.float64(0.0), b=np.float64(1.0))" + " + np.float64(2.0)" + ) + + X = Uniform(a=np.zeros(4), b=1) + assert repr(X) == "Uniform(a=array([0., 0., 0., 0.]), b=1)" + + X = Uniform(a=np.zeros(4, dtype=np.float32), b=np.ones(4, dtype=np.float32)) + assert repr(X) == ( + "Uniform(a=array([0., 0., 0., 0.], dtype=float32)," + " b=array([1., 1., 1., 1.], dtype=float32))" + ) + + +class TestReprs: + U = Uniform(a=0, b=1) + V = Uniform(a=np.float32(0.0), b=np.float32(1.0)) + X = Normal(mu=-1, sigma=1) + Y = Normal(mu=1, sigma=1) + Z = Normal(mu=np.zeros(1000), sigma=1) + + @pytest.mark.parametrize( + "dist", + [ + U, + U - np.array([1.0, 2.0]), + V, + np.ones(2, dtype=np.float32)*V + np.zeros(2, dtype=np.float64), + 3*U + 2, + U**4, + (3*U + 2)**4, + (3*U + 2)**3, + 2**U, + 2**(3*U + 1), + 1 / (1 + U), + stats.order_statistic(U, r=3, n=5), + stats.truncate(U, 0.2, 0.8), + stats.Mixture([X, Y], weights=[0.3, 0.7]), + abs(U), + stats.exp(U), + stats.log(1 + U), + np.array([1.0, 2.0])*U + np.array([2.0, 3.0]), + ] + ) + def test_executable(self, dist): + # Test that reprs actually evaluate to proper distribution + # provided relevant imports are made. + from numpy import array # noqa: F401 + from numpy import float32 # noqa: F401 + from scipy.stats import abs, exp, log, order_statistic, truncate # noqa: F401 + from scipy.stats import Mixture, Normal # noqa: F401 + from scipy.stats._new_distributions import Uniform # noqa: F401 + new_dist = eval(repr(dist)) + # A basic check that the distributions are the same + sample1 = dist.sample(shape=10, rng=1234) + sample2 = new_dist.sample(shape=10, rng=1234) + assert_equal(sample1, sample2) + assert sample1.dtype is sample2.dtype + + @pytest.mark.parametrize( + "dist", + [ + Z, + np.full(1000, 2.0) * X + 1.0, + 2.0 * X + np.full(1000, 1.0), + np.full(1000, 2.0) * X + 1.0, + stats.truncate(Z, -1, 1), + stats.truncate(Z, -np.ones(1000), np.ones(1000)), + stats.order_statistic(X, r=np.arange(1, 1000), n=1000), + Z**2, + 1.0 / (1 + stats.exp(Z)), + 2**Z, + ] + ) + def test_not_too_long(self, dist): + # Tests that array summarization is working to ensure reprs aren't too long. + # None of the reprs above will be executable. + assert len(repr(dist)) < 250 + + +class MixedDist(ContinuousDistribution): + _variable = _RealParameter('x', domain=_RealInterval(endpoints=(-np.inf, np.inf))) + def _pdf_formula(self, x, *args, **kwargs): + return (0.4 * 1/(1.1 * np.sqrt(2*np.pi)) * np.exp(-0.5*((x+0.25)/1.1)**2) + + 0.6 * 1/(0.9 * np.sqrt(2*np.pi)) * np.exp(-0.5*((x-0.5)/0.9)**2)) + + +class TestMixture: + def test_input_validation(self): + message = "`components` must contain at least one random variable." + with pytest.raises(ValueError, match=message): + Mixture([]) + + message = "Each element of `components` must be an instance..." + with pytest.raises(ValueError, match=message): + Mixture((1, 2, 3)) + + message = "All elements of `components` must have scalar shapes." + with pytest.raises(ValueError, match=message): + Mixture([Normal(mu=[1, 2]), Normal()]) + + message = "`components` and `weights` must have the same length." + with pytest.raises(ValueError, match=message): + Mixture([Normal()], weights=[0.5, 0.5]) + + message = "`weights` must have floating point dtype." + with pytest.raises(ValueError, match=message): + Mixture([Normal()], weights=[1]) + + message = "`weights` must have floating point dtype." + with pytest.raises(ValueError, match=message): + Mixture([Normal()], weights=[1]) + + message = "`weights` must sum to 1.0." + with pytest.raises(ValueError, match=message): + Mixture([Normal(), Normal()], weights=[0.5, 1.0]) + + message = "All `weights` must be non-negative." + with pytest.raises(ValueError, match=message): + Mixture([Normal(), Normal()], weights=[1.5, -0.5]) + + @pytest.mark.parametrize('shape', [(), (10,)]) + def test_basic(self, shape): + rng = np.random.default_rng(582348972387243524) + X = Mixture((Normal(mu=-0.25, sigma=1.1), Normal(mu=0.5, sigma=0.9)), + weights=(0.4, 0.6)) + Y = MixedDist() + x = rng.random(shape) + + def assert_allclose(res, ref, **kwargs): + if shape == (): + assert np.isscalar(res) + np.testing.assert_allclose(res, ref, **kwargs) + + assert_allclose(X.logentropy(), Y.logentropy()) + assert_allclose(X.entropy(), Y.entropy()) + assert_allclose(X.mode(), Y.mode()) + assert_allclose(X.median(), Y.median()) + assert_allclose(X.mean(), Y.mean()) + assert_allclose(X.variance(), Y.variance()) + assert_allclose(X.standard_deviation(), Y.standard_deviation()) + assert_allclose(X.skewness(), Y.skewness()) + assert_allclose(X.kurtosis(), Y.kurtosis()) + assert_allclose(X.logpdf(x), Y.logpdf(x)) + assert_allclose(X.pdf(x), Y.pdf(x)) + assert_allclose(X.logcdf(x), Y.logcdf(x)) + assert_allclose(X.cdf(x), Y.cdf(x)) + assert_allclose(X.logccdf(x), Y.logccdf(x)) + assert_allclose(X.ccdf(x), Y.ccdf(x)) + assert_allclose(X.ilogcdf(x), Y.ilogcdf(x)) + assert_allclose(X.icdf(x), Y.icdf(x)) + assert_allclose(X.ilogccdf(x), Y.ilogccdf(x)) + assert_allclose(X.iccdf(x), Y.iccdf(x)) + for kind in ['raw', 'central', 'standardized']: + for order in range(5): + assert_allclose(X.moment(order, kind=kind), + Y.moment(order, kind=kind), + atol=1e-15) + message = "L-moments are not currently available..." + with pytest.raises(NotImplementedError, match=message): + X.lmoment(1) + + # weak test of `sample` + shape = (10, 20, 5) + y = X.sample(shape, rng=rng) + assert y.shape == shape + assert stats.ks_1samp(y.ravel(), X.cdf).pvalue > 0.05 + + def test_default_weights(self): + a = 1.1 + Gamma = stats.make_distribution(stats.gamma) + X = Gamma(a=a) + Y = stats.Mixture((X, -X)) + x = np.linspace(-4, 4, 300) + assert_allclose(Y.pdf(x), stats.dgamma(a=a).pdf(x)) + + def test_properties(self): + components = [Normal(mu=-0.25, sigma=1.1), Normal(mu=0.5, sigma=0.9)] + weights = (0.4, 0.6) + X = Mixture(components, weights=weights) + + # Replacing properties doesn't work + # Different version of Python have different messages + with pytest.raises(AttributeError): + X.components = 10 + with pytest.raises(AttributeError): + X.weights = 10 + + # Mutation doesn't work + X.components[0] = components[1] + assert X.components[0] == components[0] + X.weights[0] = weights[1] + assert X.weights[0] == weights[0] + + def test_inverse(self): + # Originally, inverse relied on the mean to start the bracket search. + # This didn't work for distributions with non-finite mean. Check that + # this is resolved. + rng = np.random.default_rng(24358934657854237863456) + Cauchy = stats.make_distribution(stats.cauchy) + X0 = Cauchy() + X = stats.Mixture([X0, X0]) + p = rng.random(size=10) + np.testing.assert_allclose(X.icdf(p), X0.icdf(p)) + np.testing.assert_allclose(X.iccdf(p), X0.iccdf(p)) + np.testing.assert_allclose(X.ilogcdf(p), X0.ilogcdf(p)) + np.testing.assert_allclose(X.ilogccdf(p), X0.ilogccdf(p)) + + +def test_zipfian_distribution_wrapper(): + # Regression test for gh-23678: calling the cdf method at the end + # point of the Zipfian distribution would generate a warning. + Zipfian = stats.make_distribution(stats.zipfian) + zdist = Zipfian(a=0.75, n=15) + # This should not generate any warnings. + assert_equal(zdist.cdf(15), 1.0) + + +class Test_logexpxmexpy: + # Regression tests for some cases that the simplest version of `_logexpmexpy` + # did not originally handle correctly. + + def test_x_equals_y(self): + # Test x - x == 0 in log-space + x = np.asarray(2.) + assert_equal(_logexpxmexpy(x, x), -np.inf) + + def test_y_neg_inf(self): + # Test x - 0 == x in log-space + x, y = np.asarray(2.), np.asarray(-inf) + assert_equal(_logexpxmexpy(x, y), x) + + def test_nan(self): + # operations involving NaNs should not produce warnings + x = np.asarray(np.nan) + assert_equal(_logexpxmexpy(x, x), x) + + +def test_gh_25180(): + lu = _LogUniform(log_a=np.float64(-0.28915434544814245), + log_b=np.float64(0.3085224670197856)) + actual = lu.mode(method='optimization') + assert np.isfinite(actual) diff --git a/v12x/pylib/scipy/stats/tests/test_continuous_basic.py b/v12x/pylib/scipy/stats/tests/test_continuous_basic.py new file mode 100644 index 0000000000000000000000000000000000000000..f2a8ee50c83d8149b8a551cfcd068683c2babe3f --- /dev/null +++ b/v12x/pylib/scipy/stats/tests/test_continuous_basic.py @@ -0,0 +1,1065 @@ +import sys +import warnings +import platform + +import numpy as np +import numpy.testing as npt +import pytest +from pytest import raises as assert_raises +from scipy.integrate import IntegrationWarning +import itertools + +from scipy import stats +from .common_tests import (check_normalization, check_moment, + check_mean_expect, + check_var_expect, check_skew_expect, + check_kurt_expect, check_entropy, + check_private_entropy, check_entropy_vect_scale, + check_edge_support, check_named_args, + check_random_state_property, + check_meth_dtype, check_ppf_dtype, + check_cmplx_deriv, + check_pickling, check_rvs_broadcast, + check_freezing, check_munp_expect,) +from scipy.stats._distr_params import distcont +from scipy.stats._distn_infrastructure import rv_continuous_frozen + +""" +Test all continuous distributions. + +Parameters were chosen for those distributions that pass the +Kolmogorov-Smirnov test. This provides safe parameters for each +distributions so that we can perform further testing of class methods. + +These tests currently check only/mostly for serious errors and exceptions, +not for numerically exact results. +""" + +# Note that you need to add new distributions you want tested +# to _distr_params + +DECIMAL = 5 # specify the precision of the tests # increased from 0 to 5 +_IS_32BIT = (sys.maxsize < 2**32) + +# Sets of tests to skip. +# Entries sorted by speed (very slow to slow). +# xslow took > 1s; slow took > 0.5s + +xslow_test_cont_basic = {'studentized_range', 'kstwo', 'ksone', 'vonmises', 'kappa4', + 'recipinvgauss', 'vonmises_line', 'gausshyper', + 'rel_breitwigner', 'norminvgauss'} +slow_test_cont_basic = {'crystalball', 'powerlognorm', 'pearson3'} + +# test_moments is already marked slow +xslow_test_moments = {'studentized_range', 'ksone', 'vonmises', 'vonmises_line', + 'recipinvgauss', 'kstwo', 'kappa4'} + +slow_fit_mle = {'exponweib', 'genexpon', 'genhyperbolic', 'johnsonsb', + 'kappa4', 'powerlognorm', 'tukeylambda'} +xslow_fit_mle = {'gausshyper', 'ncf', 'ncx2', 'recipinvgauss', 'vonmises_line'} +xfail_fit_mle = {'ksone', 'kstwo', 'truncpareto', 'irwinhall'} +skip_fit_mle = {'levy_stable', 'studentized_range'} # far too slow (>10min) +slow_fit_mm = {'chi2', 'expon', 'lognorm', 'loguniform', 'powerlaw', 'reciprocal'} +xslow_fit_mm = {'argus', 'beta', 'exponpow', 'gausshyper', 'gengamma', + 'genhalflogistic', 'geninvgauss', 'gompertz', + 'johnsonsb', 'kstwobign', 'ncx2', 'norminvgauss', 'trapezoid', + 'truncnorm', 'truncweibull_min', 'wrapcauchy'} +xfail_fit_mm = {'alpha', 'betaprime', 'bradford', 'burr', 'burr12', 'cauchy', + 'crystalball', 'dpareto_lognorm', 'exponweib', 'f', 'fisk', + 'foldcauchy', 'genextreme', 'genpareto', 'halfcauchy', 'invgamma', + 'irwinhall', 'jf_skew_t', 'johnsonsu', 'kappa3', 'kappa4', 'landau', + 'levy', 'levy_l', 'loglaplace', 'lomax', 'mielke', 'ncf', 'nct', + 'pareto', 'powerlognorm', 'powernorm', 'rel_breitwigner', + 'skewcauchy', 't', 'truncexpon', 'truncpareto', + 'tukeylambda', 'vonmises', 'vonmises_line'} +skip_fit_mm = {'genexpon', 'genhyperbolic', 'ksone', 'kstwo', 'levy_stable', + 'recipinvgauss', 'studentized_range'} # far too slow (>10min) + +# These distributions fail the complex derivative test below. +# Here 'fail' mean produce wrong results and/or raise exceptions, depending +# on the implementation details of corresponding special functions. +# cf https://github.com/scipy/scipy/pull/4979 for a discussion. +fails_cmplx = {'argus', 'beta', 'betaprime', 'cauchy', 'chi', 'chi2', 'cosine', + 'dgamma', 'dpareto_lognorm', 'dweibull', 'erlang', 'f', 'foldcauchy', + 'gamma', 'gausshyper', 'gengamma', 'genhyperbolic', + 'geninvgauss', 'gennorm', 'genpareto', + 'halfcauchy', 'halfgennorm', 'invgamma', 'irwinhall', 'jf_skew_t', + 'ksone', 'kstwo', 'kstwobign', 'landau', 'levy_l', 'loggamma', + 'logistic', 'loguniform', 'maxwell', 'nakagami', + 'ncf', 'nct', 'ncx2', 'norminvgauss', 'pearson3', + 'powerlaw', 'rdist', 'reciprocal', 'rice', + 'skewnorm', 't', 'truncpareto', 'truncweibull_min', + 'tukeylambda', 'vonmises', 'vonmises_line', + 'rv_histogram_instance', 'truncnorm', 'studentized_range', + 'johnsonsb', 'halflogistic', 'rel_breitwigner'} + +# Slow test_method_with_lists +slow_with_lists = {'studentized_range'} + + +# rv_histogram instances, with uniform and non-uniform bins; +# stored as (dist, arg) tuples for cases_test_cont_basic +# and cases_test_moments. +histogram_test_instances = [] +case1 = {'a': [1, 2, 2, 3, 3, 3, 4, 4, 4, 4, 5, 5, 5, 5, 5, 6, + 6, 6, 6, 7, 7, 7, 8, 8, 9], 'bins': 8} # equal width bins +case2 = {'a': [1, 1], 'bins': [0, 1, 10]} # unequal width bins +for case, density in itertools.product([case1, case2], [True, False]): + _hist = np.histogram(**case, density=density) + _rv_hist = stats.rv_histogram(_hist, density=density) + histogram_test_instances.append((_rv_hist, tuple())) + + +def cases_test_cont_basic(): + for distname, arg in distcont[:] + histogram_test_instances: + if distname == 'levy_stable': # fails; tested separately + continue + if distname in slow_test_cont_basic: + yield pytest.param(distname, arg, marks=pytest.mark.slow) + elif distname in xslow_test_cont_basic: + yield pytest.param(distname, arg, marks=pytest.mark.xslow) + else: + yield distname, arg + + +@pytest.mark.parametrize('distname,arg', list(cases_test_cont_basic())) +@pytest.mark.parametrize('sn', [500]) +def test_cont_basic(distname, arg, sn, num_parallel_threads): + try: + distfn = getattr(stats, distname) + except TypeError: + distfn = distname + distname = 'rv_histogram_instance' + + rng = np.random.default_rng(7654565) + rvs = distfn.rvs(size=sn, *arg, random_state=rng) + m, v = distfn.stats(*arg) + + if distname not in {'laplace_asymmetric'}: + # TODO: multiple checks in this function are not robust, tweaking the + # seed above will make different distributions fail. + check_sample_meanvar_(m, v, rvs, rng) + check_cdf_ppf(distfn, arg, distname) + check_sf_isf(distfn, arg, distname) + check_cdf_sf(distfn, arg, distname) + check_ppf_isf(distfn, arg, distname) + check_pdf(distfn, arg, distname) + check_pdf_logpdf(distfn, arg, distname) + check_pdf_logpdf_at_endpoints(distfn, arg, distname) + check_cdf_logcdf(distfn, arg, distname) + check_sf_logsf(distfn, arg, distname) + check_ppf_broadcast(distfn, arg, distname) + + alpha = 0.01 + if distname == 'rv_histogram_instance': + check_distribution_rvs(distfn.cdf, arg, alpha, rvs) + elif distname != 'geninvgauss': + # skip kstest for geninvgauss since cdf is too slow; see test for + # rv generation in TestGenInvGauss in test_distributions.py + check_distribution_rvs(distname, arg, alpha, rvs) + + locscale_defaults = (0, 1) + meths = [distfn.pdf, distfn.logpdf, distfn.cdf, distfn.logcdf, + distfn.logsf] + # make sure arguments are within support + spec_x = {'weibull_max': -0.5, 'levy_l': -0.5, + 'pareto': 1.5, 'truncpareto': 3.2, 'tukeylambda': 0.3, + 'rv_histogram_instance': 5.0} + x = spec_x.get(distname, 0.5) + if distname == 'invweibull': + arg = (1,) + elif distname == 'ksone': + arg = (3,) + + check_named_args(distfn, x, arg, locscale_defaults, meths) + if num_parallel_threads == 1: + check_random_state_property(distfn, arg) + + if distname in ['rel_breitwigner'] and _IS_32BIT: + # gh18414 + pytest.skip("fails on Linux 32-bit") + else: + check_pickling(distfn, arg) + check_freezing(distfn, arg) + + # Entropy + if distname not in ['kstwobign', 'kstwo', 'ncf']: + check_entropy(distfn, arg, distname) + + if distfn.numargs == 0: + check_vecentropy(distfn, arg) + + if (distfn.__class__._entropy != stats.rv_continuous._entropy + and distname != 'vonmises'): + check_private_entropy(distfn, arg, stats.rv_continuous) + + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", "The occurrence of roundoff error", IntegrationWarning) + warnings.filterwarnings("ignore", "Extremely bad integrand", IntegrationWarning) + warnings.filterwarnings("ignore", "invalid value", RuntimeWarning) + check_entropy_vect_scale(distfn, arg) + + check_retrieving_support(distfn, arg) + check_edge_support(distfn, arg) + + check_meth_dtype(distfn, arg, meths) + check_ppf_dtype(distfn, arg) + + # complex special functions known to fail on sparc64 (gh-22577) + if distname not in fails_cmplx and platform.machine() != 'sparc64': + check_cmplx_deriv(distfn, arg) + + if distname != 'truncnorm': + check_ppf_private(distfn, arg, distname) + + +def cases_test_cont_basic_fit(): + slow = pytest.mark.slow + xslow = pytest.mark.xslow + fail = pytest.mark.skip(reason="Test fails and may be slow.") + skip = pytest.mark.skip(reason="Test too slow to run to completion (>10m).") + + for distname, arg in distcont[:] + histogram_test_instances: + for method in ["MLE", "MM"]: + for fix_args in [True, False]: + if method == 'MLE' and distname in slow_fit_mle: + yield pytest.param(distname, arg, method, fix_args, marks=slow) + continue + if method == 'MLE' and distname in xslow_fit_mle: + yield pytest.param(distname, arg, method, fix_args, marks=xslow) + continue + if method == 'MLE' and distname in xfail_fit_mle: + yield pytest.param(distname, arg, method, fix_args, marks=fail) + continue + if method == 'MLE' and distname in skip_fit_mle: + yield pytest.param(distname, arg, method, fix_args, marks=skip) + continue + if method == 'MM' and distname in slow_fit_mm: + yield pytest.param(distname, arg, method, fix_args, marks=slow) + continue + if method == 'MM' and distname in xslow_fit_mm: + yield pytest.param(distname, arg, method, fix_args, marks=xslow) + continue + if method == 'MM' and distname in xfail_fit_mm: + yield pytest.param(distname, arg, method, fix_args, marks=fail) + continue + if method == 'MM' and distname in skip_fit_mm: + yield pytest.param(distname, arg, method, fix_args, marks=skip) + continue + + yield distname, arg, method, fix_args + + +def test_cont_basic_fit_cases(): + # Distribution names should not be in multiple MLE or MM sets + assert (len(xslow_fit_mle.union(xfail_fit_mle).union(skip_fit_mle)) == + len(xslow_fit_mle) + len(xfail_fit_mle) + len(skip_fit_mle)) + assert (len(xslow_fit_mm.union(xfail_fit_mm).union(skip_fit_mm)) == + len(xslow_fit_mm) + len(xfail_fit_mm) + len(skip_fit_mm)) + + +@pytest.mark.parametrize('distname, arg, method, fix_args', + list(cases_test_cont_basic_fit())) +@pytest.mark.parametrize('n_fit_samples', [200]) +def test_cont_basic_fit(distname, arg, n_fit_samples, method, fix_args): + try: + distfn = getattr(stats, distname) + except TypeError: + distfn = distname + + rng = np.random.RandomState(765456) + rvs = distfn.rvs(size=n_fit_samples, *arg, random_state=rng) + if fix_args: + check_fit_args_fix(distfn, arg, rvs, method) + else: + check_fit_args(distfn, arg, rvs, method) + +@pytest.mark.parametrize('distname,arg', list(cases_test_cont_basic())) +def test_rvs_scalar(distname, arg): + # rvs should return a scalar when given scalar arguments (gh-12428) + try: + distfn = getattr(stats, distname) + except TypeError: + distfn = distname + distname = 'rv_histogram_instance' + + assert np.isscalar(distfn.rvs(*arg)) + assert np.isscalar(distfn.rvs(*arg, size=())) + assert np.isscalar(distfn.rvs(*arg, size=None)) + + +@pytest.mark.thread_unsafe(reason="global rng") +def test_levy_stable_random_state_property(): + # levy_stable only implements rvs(), so it is skipped in the + # main loop in test_cont_basic(). Here we apply just the test + # check_random_state_property to levy_stable. + check_random_state_property(stats.levy_stable, (0.5, 0.1)) + + +def cases_test_moments(): + fail_normalization = set() + fail_higher = {'ncf'} + fail_moment = {'johnsonsu'} # generic `munp` is inaccurate for johnsonsu + + for distname, arg in distcont[:] + histogram_test_instances: + if distname == 'levy_stable': + continue + + if distname in xslow_test_moments: + yield pytest.param(distname, arg, True, True, True, True, + marks=pytest.mark.xslow(reason="too slow")) + continue + + cond1 = distname not in fail_normalization + cond2 = distname not in fail_higher + cond3 = distname not in fail_moment + + marks = list() + # Currently unused, `marks` can be used to add a timeout to a test of + # a specific distribution. For example, this shows how a timeout could + # be added for the 'skewnorm' distribution: + # + # marks = list() + # if distname == 'skewnorm': + # marks.append(pytest.mark.timeout(300)) + + yield pytest.param(distname, arg, cond1, cond2, cond3, + False, marks=marks) + + if not cond1 or not cond2 or not cond3: + # Run the distributions that have issues twice, once skipping the + # not_ok parts, once with the not_ok parts but marked as knownfail + yield pytest.param(distname, arg, True, True, True, True, + marks=[pytest.mark.xfail] + marks) + + +@pytest.mark.slow +@pytest.mark.parametrize('distname,arg,normalization_ok,higher_ok,moment_ok,' + 'is_xfailing', + list(cases_test_moments())) +def test_moments(distname, arg, normalization_ok, higher_ok, moment_ok, + is_xfailing): + try: + distfn = getattr(stats, distname) + except TypeError: + distfn = distname + distname = 'rv_histogram_instance' + + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", + "The integral is probably divergent, or slowly convergent.", + IntegrationWarning, + ) + warnings.filterwarnings( + "ignore", + "The maximum number of subdivisions.", + IntegrationWarning + ) + warnings.filterwarnings( + "ignore", + "The algorithm does not converge.", + IntegrationWarning + ) + + if is_xfailing: + warnings.simplefilter("ignore", IntegrationWarning) + + m, v, s, k = distfn.stats(*arg, moments='mvsk') + + with np.errstate(all="ignore"): + if normalization_ok: + check_normalization(distfn, arg, distname) + + if higher_ok: + check_mean_expect(distfn, arg, m, distname) + check_skew_expect(distfn, arg, m, v, s, distname) + check_var_expect(distfn, arg, m, v, distname) + check_kurt_expect(distfn, arg, m, v, k, distname) + check_munp_expect(distfn, arg, distname) + + check_loc_scale(distfn, arg, m, v, distname) + + if moment_ok: + check_moment(distfn, arg, m, v, distname) + + +@pytest.mark.parametrize('dist,shape_args', distcont) +def test_rvs_broadcast(dist, shape_args): + if dist in ['gausshyper', 'studentized_range']: + pytest.skip("too slow") + + if dist in ['rel_breitwigner'] and _IS_32BIT: + # gh18414 + pytest.skip("fails on Linux 32-bit") + + # If shape_only is True, it means the _rvs method of the + # distribution uses more than one random number to generate a random + # variate. That means the result of using rvs with broadcasting or + # with a nontrivial size will not necessarily be the same as using the + # numpy.vectorize'd version of rvs(), so we can only compare the shapes + # of the results, not the values. + # Whether or not a distribution is in the following list is an + # implementation detail of the distribution, not a requirement. If + # the implementation the rvs() method of a distribution changes, this + # test might also have to be changed. + shape_only = dist in ['argus', 'betaprime', 'dgamma', 'dpareto_lognorm', 'dweibull', + 'exponnorm', 'genhyperbolic', 'geninvgauss', 'landau', + 'levy_stable', 'nct', 'norminvgauss', 'rice', + 'skewnorm', 'semicircular', 'gennorm', 'loggamma'] + + distfunc = getattr(stats, dist) + loc = np.zeros(2) + scale = np.ones((3, 1)) + nargs = distfunc.numargs + allargs = [] + bshape = [3, 2] + # Generate shape parameter arguments... + for k in range(nargs): + shp = (k + 4,) + (1,)*(k + 2) + allargs.append(shape_args[k]*np.ones(shp)) + bshape.insert(0, k + 4) + allargs.extend([loc, scale]) + # bshape holds the expected shape when loc, scale, and the shape + # parameters are all broadcast together. + + check_rvs_broadcast(distfunc, dist, allargs, bshape, shape_only, 'd') + + +# Expected values of the SF, CDF, PDF were computed using +# mpmath with mpmath.mp.dps = 50 and output at 20: +# +# def ks(x, n): +# x = mpmath.mpf(x) +# logp = -mpmath.power(6.0*n*x+1.0, 2)/18.0/n +# sf, cdf = mpmath.exp(logp), -mpmath.expm1(logp) +# pdf = (6.0*n*x+1.0) * 2 * sf/3 +# print(mpmath.nstr(sf, 20), mpmath.nstr(cdf, 20), mpmath.nstr(pdf, 20)) +# +# Tests use 1/n < x < 1-1/n and n > 1e6 to use the asymptotic computation. +# Larger x has a smaller sf. +@pytest.mark.parametrize('x,n,sf,cdf,pdf,rtol', + [(2.0e-5, 1000000000, + 0.44932297307934442379, 0.55067702692065557621, + 35946.137394996276407, 5e-15), + (2.0e-9, 1000000000, + 0.99999999061111115519, 9.3888888448132728224e-9, + 8.6666665852962971765, 5e-14), + (5.0e-4, 1000000000, + 7.1222019433090374624e-218, 1.0, + 1.4244408634752704094e-211, 5e-14)]) +def test_gh17775_regression(x, n, sf, cdf, pdf, rtol): + # Regression test for gh-17775. In scipy 1.9.3 and earlier, + # these test would fail. + # + # KS one asymptotic sf ~ e^(-(6nx+1)^2 / 18n) + # Given a large 32-bit integer n, 6n will overflow in the c implementation. + # Example of broken behaviour: + # ksone.sf(2.0e-5, 1000000000) == 0.9374359693473666 + ks = stats.ksone + vals = np.array([ks.sf(x, n), ks.cdf(x, n), ks.pdf(x, n)]) + expected = np.array([sf, cdf, pdf]) + npt.assert_allclose(vals, expected, rtol=rtol) + # The sf+cdf must sum to 1.0. + npt.assert_equal(vals[0] + vals[1], 1.0) + # Check inverting the (potentially very small) sf (uses a lower tolerance) + npt.assert_allclose([ks.isf(sf, n)], [x], rtol=1e-8) + + +def test_rvs_gh2069_regression(): + # Regression tests for gh-2069. In scipy 0.17 and earlier, + # these tests would fail. + # + # A typical example of the broken behavior: + # >>> norm.rvs(loc=np.zeros(5), scale=np.ones(5)) + # array([-2.49613705, -2.49613705, -2.49613705, -2.49613705, -2.49613705]) + rng = np.random.RandomState(123) + vals = stats.norm.rvs(loc=np.zeros(5), scale=1, random_state=rng) + d = np.diff(vals) + npt.assert_(np.all(d != 0), "All the values are equal, but they shouldn't be!") + vals = stats.norm.rvs(loc=0, scale=np.ones(5), random_state=rng) + d = np.diff(vals) + npt.assert_(np.all(d != 0), "All the values are equal, but they shouldn't be!") + vals = stats.norm.rvs(loc=np.zeros(5), scale=np.ones(5), random_state=rng) + d = np.diff(vals) + npt.assert_(np.all(d != 0), "All the values are equal, but they shouldn't be!") + vals = stats.norm.rvs(loc=np.array([[0], [0]]), scale=np.ones(5), + random_state=rng) + d = np.diff(vals.ravel()) + npt.assert_(np.all(d != 0), "All the values are equal, but they shouldn't be!") + + assert_raises(ValueError, stats.norm.rvs, [[0, 0], [0, 0]], + [[1, 1], [1, 1]], 1) + assert_raises(ValueError, stats.gamma.rvs, [2, 3, 4, 5], 0, 1, (2, 2)) + assert_raises(ValueError, stats.gamma.rvs, [1, 1, 1, 1], [0, 0, 0, 0], + [[1], [2]], (4,)) + + +def test_nomodify_gh9900_regression(): + # Regression test for gh-9990 + # Prior to gh-9990, calls to stats.truncnorm._cdf() use what ever was + # set inside the stats.truncnorm instance during stats.truncnorm.cdf(). + # This could cause issues with multi-threaded code. + # Since then, the calls to cdf() are not permitted to modify the global + # stats.truncnorm instance. + tn = stats.truncnorm + # Use the right-half truncated normal + # Check that the cdf and _cdf return the same result. + npt.assert_almost_equal(tn.cdf(1, 0, np.inf), + 0.6826894921370859) + npt.assert_almost_equal(tn._cdf([1], [0], [np.inf]), + 0.6826894921370859) + + # Now use the left-half truncated normal + npt.assert_almost_equal(tn.cdf(-1, -np.inf, 0), + 0.31731050786291415) + npt.assert_almost_equal(tn._cdf([-1], [-np.inf], [0]), + 0.31731050786291415) + + # Check that the right-half truncated normal _cdf hasn't changed + npt.assert_almost_equal(tn._cdf([1], [0], [np.inf]), + 0.6826894921370859) # Not 1.6826894921370859 + npt.assert_almost_equal(tn.cdf(1, 0, np.inf), + 0.6826894921370859) + + # Check that the left-half truncated normal _cdf hasn't changed + npt.assert_almost_equal(tn._cdf([-1], [-np.inf], [0]), + 0.31731050786291415) # Not -0.6826894921370859 + npt.assert_almost_equal(tn.cdf(1, -np.inf, 0), + 1) # Not 1.6826894921370859 + npt.assert_almost_equal(tn.cdf(-1, -np.inf, 0), + 0.31731050786291415) # Not -0.6826894921370859 + + +def test_broadcast_gh9990_regression(): + # Regression test for gh-9990 + # The x-value 7 only lies within the support of 4 of the supplied + # distributions. Prior to 9990, one array passed to + # stats.reciprocal._cdf would have 4 elements, but an array + # previously stored by stats.reciprocal_argcheck() would have 6, leading + # to a broadcast error. + a = np.array([1, 2, 3, 4, 5, 6]) + b = np.array([8, 16, 1, 32, 1, 48]) + ans = [stats.reciprocal.cdf(7, _a, _b) for _a, _b in zip(a,b)] + npt.assert_array_almost_equal(stats.reciprocal.cdf(7, a, b), ans) + + ans = [stats.reciprocal.cdf(1, _a, _b) for _a, _b in zip(a,b)] + npt.assert_array_almost_equal(stats.reciprocal.cdf(1, a, b), ans) + + ans = [stats.reciprocal.cdf(_a, _a, _b) for _a, _b in zip(a,b)] + npt.assert_array_almost_equal(stats.reciprocal.cdf(a, a, b), ans) + + ans = [stats.reciprocal.cdf(_b, _a, _b) for _a, _b in zip(a,b)] + npt.assert_array_almost_equal(stats.reciprocal.cdf(b, a, b), ans) + + +def test_broadcast_gh7933_regression(): + # Check broadcast works + stats.truncnorm.logpdf( + np.array([3.0, 2.0, 1.0]), + a=(1.5 - np.array([6.0, 5.0, 4.0])) / 3.0, + b=np.inf, + loc=np.array([6.0, 5.0, 4.0]), + scale=3.0 + ) + + +def test_gh2002_regression(): + # Add a check that broadcast works in situations where only some + # x-values are compatible with some of the shape arguments. + x = np.r_[-2:2:101j] + a = np.r_[-np.ones(50), np.ones(51)] + expected = [stats.truncnorm.pdf(_x, _a, np.inf) for _x, _a in zip(x, a)] + ans = stats.truncnorm.pdf(x, a, np.inf) + npt.assert_array_almost_equal(ans, expected) + + +def test_gh1320_regression(): + # Check that the first example from gh-1320 now works. + c = 2.62 + stats.genextreme.ppf(0.5, np.array([[c], [c + 0.5]])) + # The other examples in gh-1320 appear to have stopped working + # some time ago. + # ans = stats.genextreme.moment(2, np.array([c, c + 0.5])) + # expected = np.array([25.50105963, 115.11191437]) + # stats.genextreme.moment(5, np.array([[c], [c + 0.5]])) + # stats.genextreme.moment(5, np.array([c, c + 0.5])) + + +def test_method_of_moments(): + # example from https://en.wikipedia.org/wiki/Method_of_moments_(statistics) + x = [0, 0, 0, 0, 1] + a = 1/5 - 2*np.sqrt(3)/5 + b = 1/5 + 2*np.sqrt(3)/5 + # force use of method of moments (uniform.fit is overridden) + loc, scale = super(type(stats.uniform), stats.uniform).fit(x, method="MM") + npt.assert_almost_equal(loc, a, decimal=4) + npt.assert_almost_equal(loc+scale, b, decimal=4) + + +def check_sample_meanvar_(popmean, popvar, sample, rng): + if np.isfinite(popmean): + check_sample_mean(sample, popmean) + if np.isfinite(popvar): + check_sample_var(sample, popvar, rng) + + +def check_sample_mean(sample, popmean): + # Checks for unlikely difference between sample mean and population mean + prob = stats.ttest_1samp(sample, popmean).pvalue + assert prob > 0.01 + + +def check_sample_var(sample, popvar, rng): + # check that population mean lies within the CI bootstrapped from the + # sample. This used to be a chi-squared test for variance, but there were + # too many false positives + res = stats.bootstrap( + (sample,), + lambda x, axis: x.var(ddof=1, axis=axis), + confidence_level=0.995, + rng=rng, + ) + conf = res.confidence_interval + low, high = conf.low, conf.high + assert low <= popvar <= high + + +def check_cdf_ppf(distfn, arg, msg): + values = [0.001, 0.5, 0.999] + npt.assert_almost_equal(distfn.cdf(distfn.ppf(values, *arg), *arg), + values, decimal=DECIMAL, err_msg=msg + + ' - cdf-ppf roundtrip') + + +def check_sf_isf(distfn, arg, msg): + npt.assert_almost_equal(distfn.sf(distfn.isf([0.1, 0.5, 0.9], *arg), *arg), + [0.1, 0.5, 0.9], decimal=DECIMAL, err_msg=msg + + ' - sf-isf roundtrip') + + +def check_cdf_sf(distfn, arg, msg): + npt.assert_almost_equal(distfn.cdf([0.1, 0.9], *arg), + 1.0 - distfn.sf([0.1, 0.9], *arg), + decimal=DECIMAL, err_msg=msg + + ' - cdf-sf relationship') + + +def check_ppf_isf(distfn, arg, msg): + p = np.array([0.1, 0.9]) + npt.assert_almost_equal(distfn.isf(p, *arg), distfn.ppf(1-p, *arg), + decimal=DECIMAL, err_msg=msg + + ' - ppf-isf relationship') + + +def check_pdf(distfn, arg, msg): + # compares pdf at median with numerical derivative of cdf + median = distfn.ppf(0.5, *arg) + eps = 1e-6 + pdfv = distfn.pdf(median, *arg) + if (pdfv < 1e-4) or (pdfv > 1e4): + # avoid checking a case where pdf is close to zero or + # huge (singularity) + median = median + 0.1 + pdfv = distfn.pdf(median, *arg) + cdfdiff = (distfn.cdf(median + eps, *arg) - + distfn.cdf(median - eps, *arg))/eps/2.0 + # replace with better diff and better test (more points), + # actually, this works pretty well + msg += ' - cdf-pdf relationship' + npt.assert_almost_equal(pdfv, cdfdiff, decimal=DECIMAL, err_msg=msg) + + +def check_pdf_logpdf(distfn, args, msg): + # compares pdf at several points with the log of the pdf + points = np.array([0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8]) + vals = distfn.ppf(points, *args) + vals = vals[np.isfinite(vals)] + pdf = distfn.pdf(vals, *args) + logpdf = distfn.logpdf(vals, *args) + pdf = pdf[(pdf != 0) & np.isfinite(pdf)] + logpdf = logpdf[np.isfinite(logpdf)] + msg += " - logpdf-log(pdf) relationship" + npt.assert_almost_equal(np.log(pdf), logpdf, decimal=7, err_msg=msg) + + +def check_pdf_logpdf_at_endpoints(distfn, args, msg): + # compares pdf with the log of the pdf at the (finite) end points + points = np.array([0, 1]) + vals = distfn.ppf(points, *args) + vals = vals[np.isfinite(vals)] + pdf = distfn.pdf(vals, *args) + logpdf = distfn.logpdf(vals, *args) + pdf = pdf[(pdf != 0) & np.isfinite(pdf)] + logpdf = logpdf[np.isfinite(logpdf)] + msg += " - logpdf-log(pdf) relationship" + npt.assert_almost_equal(np.log(pdf), logpdf, decimal=7, err_msg=msg) + + +def check_sf_logsf(distfn, args, msg): + # compares sf at several points with the log of the sf + points = np.array([0.0, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 1.0]) + vals = distfn.ppf(points, *args) + vals = vals[np.isfinite(vals)] + sf = distfn.sf(vals, *args) + logsf = distfn.logsf(vals, *args) + sf = sf[sf != 0] + logsf = logsf[np.isfinite(logsf)] + msg += " - logsf-log(sf) relationship" + npt.assert_almost_equal(np.log(sf), logsf, decimal=7, err_msg=msg) + + +def check_cdf_logcdf(distfn, args, msg): + # compares cdf at several points with the log of the cdf + points = np.array([0, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 1.0]) + vals = distfn.ppf(points, *args) + vals = vals[np.isfinite(vals)] + cdf = distfn.cdf(vals, *args) + logcdf = distfn.logcdf(vals, *args) + cdf = cdf[cdf != 0] + logcdf = logcdf[np.isfinite(logcdf)] + msg += " - logcdf-log(cdf) relationship" + npt.assert_almost_equal(np.log(cdf), logcdf, decimal=7, err_msg=msg) + + +def check_ppf_broadcast(distfn, arg, msg): + # compares ppf for multiple argsets. + num_repeats = 5 + args = [] * num_repeats + if arg: + args = [np.array([_] * num_repeats) for _ in arg] + + median = distfn.ppf(0.5, *arg) + medians = distfn.ppf(0.5, *args) + msg += " - ppf multiple" + npt.assert_almost_equal(medians, [median] * num_repeats, decimal=7, err_msg=msg) + + +def check_distribution_rvs(dist, args, alpha, rvs): + # dist is either a cdf function or name of a distribution in scipy.stats. + # args are the args for scipy.stats.dist(*args) + # alpha is a significance level, ~0.01 + # rvs is array_like of random variables + # test from scipy.stats.tests + # this version reuses existing random variables + D, pval = stats.kstest(rvs, dist, args=args, N=1000) + if (pval < alpha): + # The rvs passed in failed the K-S test, which _could_ happen + # but is unlikely if alpha is small enough. + # Repeat the test with a new sample of rvs. + # Generate 1000 rvs, perform a K-S test that the new sample of rvs + # are distributed according to the distribution. + D, pval = stats.kstest(dist, dist, args=args, N=1000) + npt.assert_(pval > alpha, "D = " + str(D) + "; pval = " + str(pval) + + "; alpha = " + str(alpha) + "\nargs = " + str(args)) + + +def check_vecentropy(distfn, args): + npt.assert_equal(distfn.vecentropy(*args), distfn._entropy(*args)) + + +def check_loc_scale(distfn, arg, m, v, msg): + # Make `loc` and `scale` arrays to catch bugs like gh-13580 where + # `loc` and `scale` arrays improperly broadcast with shapes. + loc, scale = np.array([10.0, 20.0]), np.array([10.0, 20.0]) + mt, vt = distfn.stats(*arg, loc=loc, scale=scale) + npt.assert_allclose(m*scale + loc, mt) + npt.assert_allclose(v*scale*scale, vt) + + +def check_ppf_private(distfn, arg, msg): + # fails by design for truncnorm self.nb not defined + ppfs = distfn._ppf(np.array([0.1, 0.5, 0.9]), *arg) + npt.assert_(not np.any(np.isnan(ppfs)), msg + 'ppf private is nan') + + +def check_retrieving_support(distfn, args): + loc, scale = 1, 2 + supp = distfn.support(*args) + supp_loc_scale = distfn.support(*args, loc=loc, scale=scale) + npt.assert_almost_equal(np.array(supp)*scale + loc, + np.array(supp_loc_scale)) + + +def check_fit_args(distfn, arg, rvs, method): + with np.errstate(all='ignore'), warnings.catch_warnings(): + warnings.filterwarnings("ignore", category=RuntimeWarning, + message="The shape parameter of the erlang") + warnings.filterwarnings("ignore", category=RuntimeWarning, + message="floating point number truncated") + vals = distfn.fit(rvs, method=method) + vals2 = distfn.fit(rvs, optimizer='powell', method=method) + # Only check the length of the return; accuracy tested in test_fit.py + npt.assert_(len(vals) == 2+len(arg)) + npt.assert_(len(vals2) == 2+len(arg)) + + +def check_fit_args_fix(distfn, arg, rvs, method): + with np.errstate(all='ignore'), warnings.catch_warnings(): + warnings.filterwarnings("ignore", category=RuntimeWarning, + message="The shape parameter of the erlang") + + vals = distfn.fit(rvs, floc=0, method=method) + vals2 = distfn.fit(rvs, fscale=1, method=method) + npt.assert_(len(vals) == 2+len(arg)) + npt.assert_(vals[-2] == 0) + npt.assert_(vals2[-1] == 1) + npt.assert_(len(vals2) == 2+len(arg)) + if len(arg) > 0: + vals3 = distfn.fit(rvs, f0=arg[0], method=method) + npt.assert_(len(vals3) == 2+len(arg)) + npt.assert_(vals3[0] == arg[0]) + if len(arg) > 1: + vals4 = distfn.fit(rvs, f1=arg[1], method=method) + npt.assert_(len(vals4) == 2+len(arg)) + npt.assert_(vals4[1] == arg[1]) + if len(arg) > 2: + vals5 = distfn.fit(rvs, f2=arg[2], method=method) + npt.assert_(len(vals5) == 2+len(arg)) + npt.assert_(vals5[2] == arg[2]) + + +def cases_test_methods_with_lists(): + for distname, arg in distcont: + if distname in slow_with_lists: + yield pytest.param(distname, arg, marks=pytest.mark.slow) + else: + yield distname, arg + + +@pytest.mark.parametrize('method', ['pdf', 'logpdf', 'cdf', 'logcdf', + 'sf', 'logsf', 'ppf', 'isf']) +@pytest.mark.parametrize('distname, args', list(cases_test_methods_with_lists())) +def test_methods_with_lists(method, distname, args): + # Test that the continuous distributions can accept Python lists + # as arguments. + dist = getattr(stats, distname) + f = getattr(dist, method) + if distname == 'invweibull' and method.startswith('log'): + x = [1.5, 2] + else: + x = [0.1, 0.2] + + shape2 = [[a]*2 for a in args] + loc = [0, 0.1] + scale = [1, 1.01] + result = f(x, *shape2, loc=loc, scale=scale) + npt.assert_allclose(result, + [f(*v) for v in zip(x, *shape2, loc, scale)], + rtol=1e-14, atol=5e-14) + + +def test_burr_fisk_moment_gh13234_regression(): + vals0 = stats.burr.moment(1, 5, 4) + assert isinstance(vals0, float) + + vals1 = stats.fisk.moment(1, 8) + assert isinstance(vals1, float) + + +def test_moments_with_array_gh12192_regression(): + # array loc and scalar scale + vals0 = stats.norm.moment(order=1, loc=np.array([1, 2, 3]), scale=1) + expected0 = np.array([1., 2., 3.]) + npt.assert_equal(vals0, expected0) + + # array loc and invalid scalar scale + vals1 = stats.norm.moment(order=1, loc=np.array([1, 2, 3]), scale=-1) + expected1 = np.array([np.nan, np.nan, np.nan]) + npt.assert_equal(vals1, expected1) + + # array loc and array scale with invalid entries + vals2 = stats.norm.moment(order=1, loc=np.array([1, 2, 3]), + scale=[-3, 1, 0]) + expected2 = np.array([np.nan, 2., np.nan]) + npt.assert_equal(vals2, expected2) + + # (loc == 0) & (scale < 0) + vals3 = stats.norm.moment(order=2, loc=0, scale=-4) + expected3 = np.nan + npt.assert_equal(vals3, expected3) + assert isinstance(vals3, expected3.__class__) + + # array loc with 0 entries and scale with invalid entries + vals4 = stats.norm.moment(order=2, loc=[1, 0, 2], scale=[3, -4, -5]) + expected4 = np.array([10., np.nan, np.nan]) + npt.assert_equal(vals4, expected4) + + # all(loc == 0) & (array scale with invalid entries) + vals5 = stats.norm.moment(order=2, loc=[0, 0, 0], scale=[5., -2, 100.]) + expected5 = np.array([25., np.nan, 10000.]) + npt.assert_equal(vals5, expected5) + + # all( (loc == 0) & (scale < 0) ) + vals6 = stats.norm.moment(order=2, loc=[0, 0, 0], scale=[-5., -2, -100.]) + expected6 = np.array([np.nan, np.nan, np.nan]) + npt.assert_equal(vals6, expected6) + + # scalar args, loc, and scale + vals7 = stats.chi.moment(order=2, df=1, loc=0, scale=0) + expected7 = np.nan + npt.assert_equal(vals7, expected7) + assert isinstance(vals7, expected7.__class__) + + # array args, scalar loc, and scalar scale + vals8 = stats.chi.moment(order=2, df=[1, 2, 3], loc=0, scale=0) + expected8 = np.array([np.nan, np.nan, np.nan]) + npt.assert_equal(vals8, expected8) + + # array args, array loc, and array scale + vals9 = stats.chi.moment(order=2, df=[1, 2, 3], loc=[1., 0., 2.], + scale=[1., -3., 0.]) + expected9 = np.array([3.59576912, np.nan, np.nan]) + npt.assert_allclose(vals9, expected9, rtol=1e-8) + + # (n > 4), all(loc != 0), and all(scale != 0) + vals10 = stats.norm.moment(5, [1., 2.], [1., 2.]) + expected10 = np.array([26., 832.]) + npt.assert_allclose(vals10, expected10, rtol=1e-13) + + # test broadcasting and more + a = [-1.1, 0, 1, 2.2, np.pi] + b = [-1.1, 0, 1, 2.2, np.pi] + loc = [-1.1, 0, np.sqrt(2)] + scale = [-2.1, 0, 1, 2.2, np.pi] + + a = np.array(a).reshape((-1, 1, 1, 1)) + b = np.array(b).reshape((-1, 1, 1)) + loc = np.array(loc).reshape((-1, 1)) + scale = np.array(scale) + + vals11 = stats.beta.moment(order=2, a=a, b=b, loc=loc, scale=scale) + + a, b, loc, scale = np.broadcast_arrays(a, b, loc, scale) + + for i in np.ndenumerate(a): + with np.errstate(invalid='ignore', divide='ignore'): + i = i[0] # just get the index + # check against same function with scalar input + expected = stats.beta.moment(order=2, a=a[i], b=b[i], + loc=loc[i], scale=scale[i]) + np.testing.assert_equal(vals11[i], expected) + + +def test_broadcasting_in_moments_gh12192_regression(): + vals0 = stats.norm.moment(order=1, loc=np.array([1, 2, 3]), scale=[[1]]) + expected0 = np.array([[1., 2., 3.]]) + npt.assert_equal(vals0, expected0) + assert vals0.shape == expected0.shape + + vals1 = stats.norm.moment(order=1, loc=np.array([[1], [2], [3]]), + scale=[1, 2, 3]) + expected1 = np.array([[1., 1., 1.], [2., 2., 2.], [3., 3., 3.]]) + npt.assert_equal(vals1, expected1) + assert vals1.shape == expected1.shape + + vals2 = stats.chi.moment(order=1, df=[1., 2., 3.], loc=0., scale=1.) + expected2 = np.array([0.79788456, 1.25331414, 1.59576912]) + npt.assert_allclose(vals2, expected2, rtol=1e-8) + assert vals2.shape == expected2.shape + + vals3 = stats.chi.moment(order=1, df=[[1.], [2.], [3.]], loc=[0., 1., 2.], + scale=[-1., 0., 3.]) + expected3 = np.array([[np.nan, np.nan, 4.39365368], + [np.nan, np.nan, 5.75994241], + [np.nan, np.nan, 6.78730736]]) + npt.assert_allclose(vals3, expected3, rtol=1e-8) + assert vals3.shape == expected3.shape + + +@pytest.mark.slow +def test_kappa3_array_gh13582(): + # https://github.com/scipy/scipy/pull/15140#issuecomment-994958241 + shapes = [0.5, 1.5, 2.5, 3.5, 4.5] + moments = 'mvsk' + res = np.array([[stats.kappa3.stats(shape, moments=moment) + for shape in shapes] for moment in moments]) + res2 = np.array(stats.kappa3.stats(shapes, moments=moments)) + npt.assert_allclose(res, res2) + + +@pytest.mark.xslow +def test_kappa4_array_gh13582(): + h = np.array([-0.5, 2.5, 3.5, 4.5, -3]) + k = np.array([-0.5, 1, -1.5, 0, 3.5]) + moments = 'mvsk' + res = np.array([[stats.kappa4.stats(h[i], k[i], moments=moment) + for i in range(5)] for moment in moments]) + res2 = np.array(stats.kappa4.stats(h, k, moments=moments)) + npt.assert_allclose(res, res2) + + # https://github.com/scipy/scipy/pull/15250#discussion_r775112913 + h = np.array([-1, -1/4, -1/4, 1, -1, 0]) + k = np.array([1, 1, 1/2, -1/3, -1, 0]) + res = np.array([[stats.kappa4.stats(h[i], k[i], moments=moment) + for i in range(6)] for moment in moments]) + res2 = np.array(stats.kappa4.stats(h, k, moments=moments)) + npt.assert_allclose(res, res2) + + # https://github.com/scipy/scipy/pull/15250#discussion_r775115021 + h = np.array([-1, -0.5, 1]) + k = np.array([-1, -0.5, 0, 1])[:, None] + res2 = np.array(stats.kappa4.stats(h, k, moments=moments)) + assert res2.shape == (4, 4, 3) + + +def test_frozen_attributes(monkeypatch): + # gh-14827 reported that all frozen distributions had both pmf and pdf + # attributes; continuous should have pdf and discrete should have pmf. + message = "'rv_continuous_frozen' object has no attribute" + with pytest.raises(AttributeError, match=message): + stats.norm().pmf + with pytest.raises(AttributeError, match=message): + stats.norm().logpmf + monkeypatch.setattr(stats.norm, "pmf", "herring", raising=False) + frozen_norm = stats.norm() + assert isinstance(frozen_norm, rv_continuous_frozen) + assert not hasattr(frozen_norm, "pmf") + + +def test_skewnorm_pdf_gh16038(): + rng = np.random.default_rng(0) + x, a = -np.inf, 0 + npt.assert_equal(stats.skewnorm.pdf(x, a), stats.norm.pdf(x)) + x, a = rng.random(size=(3, 3)), rng.random(size=(3, 3)) + mask = rng.random(size=(3, 3)) < 0.5 + a[mask] = 0 + x_norm = x[mask] + res = stats.skewnorm.pdf(x, a) + npt.assert_equal(res[mask], stats.norm.pdf(x_norm)) + npt.assert_equal(res[~mask], stats.skewnorm.pdf(x[~mask], a[~mask])) + + +# for scalar input, these functions should return scalar output +scalar_out = [['rvs', []], ['pdf', [0]], ['logpdf', [0]], ['cdf', [0]], + ['logcdf', [0]], ['sf', [0]], ['logsf', [0]], ['ppf', [0]], + ['isf', [0]], ['moment', [1]], ['entropy', []], ['expect', []], + ['median', []], ['mean', []], ['std', []], ['var', []]] +scalars_out = [['interval', [0.95]], ['support', []], ['stats', ['mv']]] + + +@pytest.mark.parametrize('case', scalar_out + scalars_out) +def test_scalar_for_scalar(case): + # Some rv_continuous functions returned 0d array instead of NumPy scalar + # Guard against regression + method_name, args = case + method = getattr(stats.norm(), method_name) + res = method(*args) + if case in scalar_out: + assert isinstance(res, np.number) + else: + assert isinstance(res[0], np.number) + assert isinstance(res[1], np.number) + + +def test_scalar_for_scalar2(): + # test methods that are not attributes of frozen distributions + res = stats.norm.fit([1, 2, 3]) + assert isinstance(res[0], np.number) + assert isinstance(res[1], np.number) + res = stats.norm.fit_loc_scale([1, 2, 3]) + assert isinstance(res[0], np.number) + assert isinstance(res[1], np.number) + res = stats.norm.nnlf((0, 1), [1, 2, 3]) + assert isinstance(res, np.number) diff --git a/v12x/pylib/scipy/stats/tests/test_continuous_fit_censored.py b/v12x/pylib/scipy/stats/tests/test_continuous_fit_censored.py new file mode 100644 index 0000000000000000000000000000000000000000..4508b49712e5bc8975bf2f9b2681ccc6504b0ae0 --- /dev/null +++ b/v12x/pylib/scipy/stats/tests/test_continuous_fit_censored.py @@ -0,0 +1,683 @@ +# Tests for fitting specific distributions to censored data. + +import numpy as np +from numpy.testing import assert_allclose + +from scipy.optimize import fmin +from scipy.stats import (CensoredData, beta, cauchy, chi2, expon, gamma, + gumbel_l, gumbel_r, invgauss, invweibull, laplace, + logistic, lognorm, nct, ncx2, norm, weibull_max, + weibull_min) + + +# In some tests, we'll use this optimizer for improved accuracy. +def optimizer(func, x0, args=(), disp=0): + return fmin(func, x0, args=args, disp=disp, xtol=1e-12, ftol=1e-12) + + +def test_beta(): + """ + Test fitting beta shape parameters to interval-censored data. + + Calculation in R: + + > library(fitdistrplus) + > data <- data.frame(left=c(0.10, 0.50, 0.75, 0.80), + + right=c(0.20, 0.55, 0.90, 0.95)) + > result = fitdistcens(data, 'beta', control=list(reltol=1e-14)) + + > result + Fitting of the distribution ' beta ' on censored data by maximum likelihood + Parameters: + estimate + shape1 1.419941 + shape2 1.027066 + > result$sd + shape1 shape2 + 0.9914177 0.6866565 + """ + data = CensoredData(interval=[[0.10, 0.20], + [0.50, 0.55], + [0.75, 0.90], + [0.80, 0.95]]) + + # For this test, fit only the shape parameters; loc and scale are fixed. + a, b, loc, scale = beta.fit(data, floc=0, fscale=1, optimizer=optimizer) + + assert_allclose(a, 1.419941, rtol=5e-6) + assert_allclose(b, 1.027066, rtol=5e-6) + assert loc == 0 + assert scale == 1 + + +def test_cauchy_right_censored(): + """ + Test fitting the Cauchy distribution to right-censored data. + + Calculation in R, with two values not censored [1, 10] and + one right-censored value [30]. + + > library(fitdistrplus) + > data <- data.frame(left=c(1, 10, 30), right=c(1, 10, NA)) + > result = fitdistcens(data, 'cauchy', control=list(reltol=1e-14)) + > result + Fitting of the distribution ' cauchy ' on censored data by maximum + likelihood + Parameters: + estimate + location 7.100001 + scale 7.455866 + """ + data = CensoredData(uncensored=[1, 10], right=[30]) + loc, scale = cauchy.fit(data, optimizer=optimizer) + assert_allclose(loc, 7.10001, rtol=5e-6) + assert_allclose(scale, 7.455866, rtol=5e-6) + + +def test_cauchy_mixed(): + """ + Test fitting the Cauchy distribution to data with mixed censoring. + + Calculation in R, with: + * two values not censored [1, 10], + * one left-censored [1], + * one right-censored [30], and + * one interval-censored [[4, 8]]. + + > library(fitdistrplus) + > data <- data.frame(left=c(NA, 1, 4, 10, 30), right=c(1, 1, 8, 10, NA)) + > result = fitdistcens(data, 'cauchy', control=list(reltol=1e-14)) + > result + Fitting of the distribution ' cauchy ' on censored data by maximum + likelihood + Parameters: + estimate + location 4.605150 + scale 5.900852 + """ + data = CensoredData(uncensored=[1, 10], left=[1], right=[30], + interval=[[4, 8]]) + loc, scale = cauchy.fit(data, optimizer=optimizer) + assert_allclose(loc, 4.605150, rtol=5e-6) + assert_allclose(scale, 5.900852, rtol=5e-6) + + +def test_chi2_mixed(): + """ + Test fitting just the shape parameter (df) of chi2 to mixed data. + + Calculation in R, with: + * two values not censored [1, 10], + * one left-censored [1], + * one right-censored [30], and + * one interval-censored [[4, 8]]. + + > library(fitdistrplus) + > data <- data.frame(left=c(NA, 1, 4, 10, 30), right=c(1, 1, 8, 10, NA)) + > result = fitdistcens(data, 'chisq', control=list(reltol=1e-14)) + > result + Fitting of the distribution ' chisq ' on censored data by maximum + likelihood + Parameters: + estimate + df 5.060329 + """ + data = CensoredData(uncensored=[1, 10], left=[1], right=[30], + interval=[[4, 8]]) + df, loc, scale = chi2.fit(data, floc=0, fscale=1, optimizer=optimizer) + assert_allclose(df, 5.060329, rtol=5e-6) + assert loc == 0 + assert scale == 1 + + +def test_expon_right_censored(): + """ + For the exponential distribution with loc=0, the exact solution for + fitting n uncensored points x[0]...x[n-1] and m right-censored points + x[n]..x[n+m-1] is + + scale = sum(x)/n + + That is, divide the sum of all the values (not censored and + right-censored) by the number of uncensored values. (See, for example, + https://en.wikipedia.org/wiki/Censoring_(statistics)#Likelihood.) + + The second derivative of the log-likelihood function is + + n/scale**2 - 2*sum(x)/scale**3 + + from which the estimate of the standard error can be computed. + + ----- + + Calculation in R, for reference only. The R results are not + used in the test. + + > library(fitdistrplus) + > dexps <- function(x, scale) { + + return(dexp(x, 1/scale)) + + } + > pexps <- function(q, scale) { + + return(pexp(q, 1/scale)) + + } + > left <- c(1, 2.5, 3, 6, 7.5, 10, 12, 12, 14.5, 15, + + 16, 16, 20, 20, 21, 22) + > right <- c(1, 2.5, 3, 6, 7.5, 10, 12, 12, 14.5, 15, + + NA, NA, NA, NA, NA, NA) + > result = fitdistcens(data, 'exps', start=list(scale=mean(data$left)), + + control=list(reltol=1e-14)) + > result + Fitting of the distribution ' exps ' on censored data by maximum likelihood + Parameters: + estimate + scale 19.85 + > result$sd + scale + 6.277119 + """ + # This data has 10 uncensored values and 6 right-censored values. + obs = [1, 2.5, 3, 6, 7.5, 10, 12, 12, 14.5, 15, 16, 16, 20, 20, 21, 22] + cens = [False]*10 + [True]*6 + data = CensoredData.right_censored(obs, cens) + + loc, scale = expon.fit(data, floc=0, optimizer=optimizer) + + assert loc == 0 + # Use the analytical solution to compute the expected value. This + # is the sum of the observed values divided by the number of uncensored + # values. + n = len(data) - data.num_censored() + total = data._uncensored.sum() + data._right.sum() + expected = total / n + assert_allclose(scale, expected, 1e-8) + + +def test_gamma_right_censored(): + """ + Fit gamma shape and scale to data with one right-censored value. + + Calculation in R: + + > library(fitdistrplus) + > data <- data.frame(left=c(2.5, 2.9, 3.8, 9.1, 9.3, 12.0, 23.0, 25.0), + + right=c(2.5, 2.9, 3.8, 9.1, 9.3, 12.0, 23.0, NA)) + > result = fitdistcens(data, 'gamma', start=list(shape=1, scale=10), + + control=list(reltol=1e-13)) + > result + Fitting of the distribution ' gamma ' on censored data by maximum + likelihood + Parameters: + estimate + shape 1.447623 + scale 8.360197 + > result$sd + shape scale + 0.7053086 5.1016531 + """ + # The last value is right-censored. + x = CensoredData.right_censored([2.5, 2.9, 3.8, 9.1, 9.3, 12.0, 23.0, + 25.0], + [0]*7 + [1]) + + a, loc, scale = gamma.fit(x, floc=0, optimizer=optimizer) + + assert_allclose(a, 1.447623, rtol=5e-6) + assert loc == 0 + assert_allclose(scale, 8.360197, rtol=5e-6) + + +def test_gumbel(): + """ + Fit gumbel_l and gumbel_r to censored data. + + This R calculation should match gumbel_r. + + > library(evd) + > library(fitdistrplus) + > data = data.frame(left=c(0, 2, 3, 9, 10, 10), + + right=c(1, 2, 3, 9, NA, NA)) + > result = fitdistcens(data, 'gumbel', + + control=list(reltol=1e-14), + + start=list(loc=4, scale=5)) + > result + Fitting of the distribution ' gumbel ' on censored data by maximum + likelihood + Parameters: + estimate + loc 4.487853 + scale 4.843640 + """ + # First value is interval-censored. Last two are right-censored. + uncensored = np.array([2, 3, 9]) + right = np.array([10, 10]) + interval = np.array([[0, 1]]) + data = CensoredData(uncensored, right=right, interval=interval) + loc, scale = gumbel_r.fit(data, optimizer=optimizer) + assert_allclose(loc, 4.487853, rtol=5e-6) + assert_allclose(scale, 4.843640, rtol=5e-6) + + # Negate the data and reverse the intervals, and test with gumbel_l. + data2 = CensoredData(-uncensored, left=-right, + interval=-interval[:, ::-1]) + # Fitting gumbel_l to data2 should give the same result as above, but + # with loc negated. + loc2, scale2 = gumbel_l.fit(data2, optimizer=optimizer) + assert_allclose(loc2, -4.487853, rtol=5e-6) + assert_allclose(scale2, 4.843640, rtol=5e-6) + + +def test_invgauss(): + """ + Fit just the shape parameter of invgauss to data with one value + left-censored and one value right-censored. + + Calculation in R; using a fixed dispersion parameter amounts to fixing + the scale to be 1. + + > library(statmod) + > library(fitdistrplus) + > left <- c(NA, 0.4813096, 0.5571880, 0.5132463, 0.3801414, 0.5904386, + + 0.4822340, 0.3478597, 3, 0.7191797, 1.5810902, 0.4442299) + > right <- c(0.15, 0.4813096, 0.5571880, 0.5132463, 0.3801414, 0.5904386, + + 0.4822340, 0.3478597, NA, 0.7191797, 1.5810902, 0.4442299) + > data <- data.frame(left=left, right=right) + > result = fitdistcens(data, 'invgauss', control=list(reltol=1e-12), + + fix.arg=list(dispersion=1), start=list(mean=3)) + > result + Fitting of the distribution ' invgauss ' on censored data by maximum + likelihood + Parameters: + estimate + mean 0.853469 + Fixed parameters: + value + dispersion 1 + > result$sd + mean + 0.247636 + + Here's the R calculation with the dispersion as a free parameter to + be fit. + + > result = fitdistcens(data, 'invgauss', control=list(reltol=1e-12), + + start=list(mean=3, dispersion=1)) + > result + Fitting of the distribution ' invgauss ' on censored data by maximum + likelihood + Parameters: + estimate + mean 0.8699819 + dispersion 1.2261362 + + The parametrization of the inverse Gaussian distribution in the + `statmod` package is not the same as in SciPy (see + https://arxiv.org/abs/1603.06687 + for details). The translation from R to SciPy is + + scale = 1/dispersion + mu = mean * dispersion + + > 1/result$estimate['dispersion'] # 1/dispersion + dispersion + 0.8155701 + > result$estimate['mean'] * result$estimate['dispersion'] + mean + 1.066716 + + Those last two values are the SciPy scale and shape parameters. + """ + # One point is left-censored, and one is right-censored. + x = [0.4813096, 0.5571880, 0.5132463, 0.3801414, + 0.5904386, 0.4822340, 0.3478597, 0.7191797, + 1.5810902, 0.4442299] + data = CensoredData(uncensored=x, left=[0.15], right=[3]) + + # Fit only the shape parameter. + mu, loc, scale = invgauss.fit(data, floc=0, fscale=1, optimizer=optimizer) + + assert_allclose(mu, 0.853469, rtol=5e-5) + assert loc == 0 + assert scale == 1 + + # Fit the shape and scale. + mu, loc, scale = invgauss.fit(data, floc=0, optimizer=optimizer) + + assert_allclose(mu, 1.066716, rtol=5e-5) + assert loc == 0 + assert_allclose(scale, 0.8155701, rtol=5e-5) + + +def test_invweibull(): + """ + Fit invweibull to censored data. + + Here is the calculation in R. The 'frechet' distribution from the evd + package matches SciPy's invweibull distribution. The `loc` parameter + is fixed at 0. + + > library(evd) + > library(fitdistrplus) + > data = data.frame(left=c(0, 2, 3, 9, 10, 10), + + right=c(1, 2, 3, 9, NA, NA)) + > result = fitdistcens(data, 'frechet', + + control=list(reltol=1e-14), + + start=list(loc=4, scale=5)) + > result + Fitting of the distribution ' frechet ' on censored data by maximum + likelihood + Parameters: + estimate + scale 2.7902200 + shape 0.6379845 + Fixed parameters: + value + loc 0 + """ + # In the R data, the first value is interval-censored, and the last + # two are right-censored. The rest are not censored. + data = CensoredData(uncensored=[2, 3, 9], right=[10, 10], + interval=[[0, 1]]) + c, loc, scale = invweibull.fit(data, floc=0, optimizer=optimizer) + assert_allclose(c, 0.6379845, rtol=5e-6) + assert loc == 0 + assert_allclose(scale, 2.7902200, rtol=5e-6) + + +def test_laplace(): + """ + Fir the Laplace distribution to left- and right-censored data. + + Calculation in R: + + > library(fitdistrplus) + > dlaplace <- function(x, location=0, scale=1) { + + return(0.5*exp(-abs((x - location)/scale))/scale) + + } + > plaplace <- function(q, location=0, scale=1) { + + z <- (q - location)/scale + + s <- sign(z) + + f <- -s*0.5*exp(-abs(z)) + (s+1)/2 + + return(f) + + } + > left <- c(NA, -41.564, 50.0, 15.7384, 50.0, 10.0452, -2.0684, + + -19.5399, 50.0, 9.0005, 27.1227, 4.3113, -3.7372, + + 25.3111, 14.7987, 34.0887, 50.0, 42.8496, 18.5862, + + 32.8921, 9.0448, -27.4591, NA, 19.5083, -9.7199) + > right <- c(-50.0, -41.564, NA, 15.7384, NA, 10.0452, -2.0684, + + -19.5399, NA, 9.0005, 27.1227, 4.3113, -3.7372, + + 25.3111, 14.7987, 34.0887, NA, 42.8496, 18.5862, + + 32.8921, 9.0448, -27.4591, -50.0, 19.5083, -9.7199) + > data <- data.frame(left=left, right=right) + > result <- fitdistcens(data, 'laplace', start=list(location=10, scale=10), + + control=list(reltol=1e-13)) + > result + Fitting of the distribution ' laplace ' on censored data by maximum + likelihood + Parameters: + estimate + location 14.79870 + scale 30.93601 + > result$sd + location scale + 0.1758864 7.0972125 + """ + # The value -50 is left-censored, and the value 50 is right-censored. + obs = np.array([-50.0, -41.564, 50.0, 15.7384, 50.0, 10.0452, -2.0684, + -19.5399, 50.0, 9.0005, 27.1227, 4.3113, -3.7372, + 25.3111, 14.7987, 34.0887, 50.0, 42.8496, 18.5862, + 32.8921, 9.0448, -27.4591, -50.0, 19.5083, -9.7199]) + x = obs[(obs != -50.0) & (obs != 50)] + left = obs[obs == -50.0] + right = obs[obs == 50.0] + data = CensoredData(uncensored=x, left=left, right=right) + loc, scale = laplace.fit(data, loc=10, scale=10, optimizer=optimizer) + assert_allclose(loc, 14.79870, rtol=5e-6) + assert_allclose(scale, 30.93601, rtol=5e-6) + + +def test_logistic(): + """ + Fit the logistic distribution to left-censored data. + + Calculation in R: + > library(fitdistrplus) + > left = c(13.5401, 37.4235, 11.906 , 13.998 , NA , 0.4023, NA , + + 10.9044, 21.0629, 9.6985, NA , 12.9016, 39.164 , 34.6396, + + NA , 20.3665, 16.5889, 18.0952, 45.3818, 35.3306, 8.4949, + + 3.4041, NA , 7.2828, 37.1265, 6.5969, 17.6868, 17.4977, + + 16.3391, 36.0541) + > right = c(13.5401, 37.4235, 11.906 , 13.998 , 0. , 0.4023, 0. , + + 10.9044, 21.0629, 9.6985, 0. , 12.9016, 39.164 , 34.6396, + + 0. , 20.3665, 16.5889, 18.0952, 45.3818, 35.3306, 8.4949, + + 3.4041, 0. , 7.2828, 37.1265, 6.5969, 17.6868, 17.4977, + + 16.3391, 36.0541) + > data = data.frame(left=left, right=right) + > result = fitdistcens(data, 'logis', control=list(reltol=1e-14)) + > result + Fitting of the distribution ' logis ' on censored data by maximum + likelihood + Parameters: + estimate + location 14.633459 + scale 9.232736 + > result$sd + location scale + 2.931505 1.546879 + """ + # Values that are zero are left-censored; the true values are less than 0. + x = np.array([13.5401, 37.4235, 11.906, 13.998, 0.0, 0.4023, 0.0, 10.9044, + 21.0629, 9.6985, 0.0, 12.9016, 39.164, 34.6396, 0.0, 20.3665, + 16.5889, 18.0952, 45.3818, 35.3306, 8.4949, 3.4041, 0.0, + 7.2828, 37.1265, 6.5969, 17.6868, 17.4977, 16.3391, + 36.0541]) + data = CensoredData.left_censored(x, censored=(x == 0)) + loc, scale = logistic.fit(data, optimizer=optimizer) + assert_allclose(loc, 14.633459, rtol=5e-7) + assert_allclose(scale, 9.232736, rtol=5e-6) + + +def test_lognorm(): + """ + Ref: https://math.montana.edu/jobo/st528/documents/relc.pdf + + The data is the locomotive control time to failure example that starts + on page 8. That's the 8th page in the PDF; the page number shown in + the text is 270). + The document includes SAS output for the data. + """ + # These are the uncensored measurements. There are also 59 right-censored + # measurements where the lower bound is 135. + miles_to_fail = [22.5, 37.5, 46.0, 48.5, 51.5, 53.0, 54.5, 57.5, 66.5, + 68.0, 69.5, 76.5, 77.0, 78.5, 80.0, 81.5, 82.0, 83.0, + 84.0, 91.5, 93.5, 102.5, 107.0, 108.5, 112.5, 113.5, + 116.0, 117.0, 118.5, 119.0, 120.0, 122.5, 123.0, 127.5, + 131.0, 132.5, 134.0] + + data = CensoredData.right_censored(miles_to_fail + [135]*59, + [0]*len(miles_to_fail) + [1]*59) + sigma, loc, scale = lognorm.fit(data, floc=0) + + assert loc == 0 + # Convert the lognorm parameters to the mu and sigma of the underlying + # normal distribution. + mu = np.log(scale) + # The expected results are from the 17th page of the PDF document + # (labeled page 279), in the SAS output on the right side of the page. + assert_allclose(mu, 5.1169, rtol=5e-4) + assert_allclose(sigma, 0.7055, rtol=5e-3) + + +def test_nct(): + """ + Test fitting the noncentral t distribution to censored data. + + Calculation in R: + + > library(fitdistrplus) + > data <- data.frame(left=c(1, 2, 3, 5, 8, 10, 25, 25), + + right=c(1, 2, 3, 5, 8, 10, NA, NA)) + > result = fitdistcens(data, 't', control=list(reltol=1e-14), + + start=list(df=1, ncp=2)) + > result + Fitting of the distribution ' t ' on censored data by maximum likelihood + Parameters: + estimate + df 0.5432336 + ncp 2.8893565 + + """ + data = CensoredData.right_censored([1, 2, 3, 5, 8, 10, 25, 25], + [0, 0, 0, 0, 0, 0, 1, 1]) + # Fit just the shape parameter df and nc; loc and scale are fixed. + with np.errstate(over='ignore'): # remove context when gh-14901 is closed + df, nc, loc, scale = nct.fit(data, floc=0, fscale=1, + optimizer=optimizer) + assert_allclose(df, 0.5432336, rtol=5e-6) + assert_allclose(nc, 2.8893565, rtol=5e-6) + assert loc == 0 + assert scale == 1 + + +def test_ncx2(): + """ + Test fitting the shape parameters (df, ncp) of ncx2 to mixed data. + + Calculation in R, with + * 5 not censored values [2.7, 0.2, 6.5, 0.4, 0.1], + * 1 interval-censored value [[0.6, 1.0]], and + * 2 right-censored values [8, 8]. + + > library(fitdistrplus) + > data <- data.frame(left=c(2.7, 0.2, 6.5, 0.4, 0.1, 0.6, 8, 8), + + right=c(2.7, 0.2, 6.5, 0.4, 0.1, 1.0, NA, NA)) + > result = fitdistcens(data, 'chisq', control=list(reltol=1e-14), + + start=list(df=1, ncp=2)) + > result + Fitting of the distribution ' chisq ' on censored data by maximum + likelihood + Parameters: + estimate + df 1.052871 + ncp 2.362934 + """ + data = CensoredData(uncensored=[2.7, 0.2, 6.5, 0.4, 0.1], right=[8, 8], + interval=[[0.6, 1.0]]) + with np.errstate(over='ignore'): # remove context when gh-14901 is closed + df, ncp, loc, scale = ncx2.fit(data, floc=0, fscale=1, + optimizer=optimizer) + assert_allclose(df, 1.052871, rtol=5e-6) + assert_allclose(ncp, 2.362934, rtol=5e-6) + assert loc == 0 + assert scale == 1 + + +def test_norm(): + """ + Test fitting the normal distribution to interval-censored data. + + Calculation in R: + + > library(fitdistrplus) + > data <- data.frame(left=c(0.10, 0.50, 0.75, 0.80), + + right=c(0.20, 0.55, 0.90, 0.95)) + > result = fitdistcens(data, 'norm', control=list(reltol=1e-14)) + + > result + Fitting of the distribution ' norm ' on censored data by maximum likelihood + Parameters: + estimate + mean 0.5919990 + sd 0.2868042 + > result$sd + mean sd + 0.1444432 0.1029451 + """ + data = CensoredData(interval=[[0.10, 0.20], + [0.50, 0.55], + [0.75, 0.90], + [0.80, 0.95]]) + + loc, scale = norm.fit(data, optimizer=optimizer) + + assert_allclose(loc, 0.5919990, rtol=5e-6) + assert_allclose(scale, 0.2868042, rtol=5e-6) + + +def test_weibull_censored1(): + # Ref: http://www.ams.sunysb.edu/~zhu/ams588/Lecture_3_likelihood.pdf + + # Survival times; '*' indicates right-censored. + s = "3,5,6*,8,10*,11*,15,20*,22,23,27*,29,32,35,40,26,28,33*,21,24*" + + times, cens = zip(*[(float(t[0]), len(t) == 2) + for t in [w.split('*') for w in s.split(',')]]) + data = CensoredData.right_censored(times, cens) + + c, loc, scale = weibull_min.fit(data, floc=0) + + # Expected values are from the reference. + assert_allclose(c, 2.149, rtol=1e-3) + assert loc == 0 + assert_allclose(scale, 28.99, rtol=1e-3) + + # Flip the sign of the data, and make the censored values + # left-censored. We should get the same parameters when we fit + # weibull_max to the flipped data. + data2 = CensoredData.left_censored(-np.array(times), cens) + + c2, loc2, scale2 = weibull_max.fit(data2, floc=0) + + assert_allclose(c2, 2.149, rtol=1e-3) + assert loc2 == 0 + assert_allclose(scale2, 28.99, rtol=1e-3) + + +def test_weibull_min_sas1(): + # Data and SAS results from + # https://support.sas.com/documentation/cdl/en/qcug/63922/HTML/default/ + # viewer.htm#qcug_reliability_sect004.htm + + text = """ + 450 0 460 1 1150 0 1150 0 1560 1 + 1600 0 1660 1 1850 1 1850 1 1850 1 + 1850 1 1850 1 2030 1 2030 1 2030 1 + 2070 0 2070 0 2080 0 2200 1 3000 1 + 3000 1 3000 1 3000 1 3100 0 3200 1 + 3450 0 3750 1 3750 1 4150 1 4150 1 + 4150 1 4150 1 4300 1 4300 1 4300 1 + 4300 1 4600 0 4850 1 4850 1 4850 1 + 4850 1 5000 1 5000 1 5000 1 6100 1 + 6100 0 6100 1 6100 1 6300 1 6450 1 + 6450 1 6700 1 7450 1 7800 1 7800 1 + 8100 1 8100 1 8200 1 8500 1 8500 1 + 8500 1 8750 1 8750 0 8750 1 9400 1 + 9900 1 10100 1 10100 1 10100 1 11500 1 + """ + + life, cens = np.array([int(w) for w in text.split()]).reshape(-1, 2).T + life = life/1000.0 + + data = CensoredData.right_censored(life, cens) + + c, loc, scale = weibull_min.fit(data, floc=0, optimizer=optimizer) + assert_allclose(c, 1.0584, rtol=1e-4) + assert_allclose(scale, 26.2968, rtol=1e-5) + assert loc == 0 + + +def test_weibull_min_sas2(): + # http://support.sas.com/documentation/cdl/en/ormpug/67517/HTML/default/ + # viewer.htm#ormpug_nlpsolver_examples06.htm + + # The last two values are right-censored. + days = np.array([143, 164, 188, 188, 190, 192, 206, 209, 213, 216, 220, + 227, 230, 234, 246, 265, 304, 216, 244]) + + data = CensoredData.right_censored(days, [0]*(len(days) - 2) + [1]*2) + + c, loc, scale = weibull_min.fit(data, 1, loc=100, scale=100, + optimizer=optimizer) + + assert_allclose(c, 2.7112, rtol=5e-4) + assert_allclose(loc, 122.03, rtol=5e-4) + assert_allclose(scale, 108.37, rtol=5e-4) diff --git a/v12x/pylib/scipy/stats/tests/test_correlation.py b/v12x/pylib/scipy/stats/tests/test_correlation.py new file mode 100644 index 0000000000000000000000000000000000000000..fdc8abb068dd92b71aeb3081a895045c24faebe0 --- /dev/null +++ b/v12x/pylib/scipy/stats/tests/test_correlation.py @@ -0,0 +1,423 @@ +import pytest +import numpy as np +from numpy.testing import assert_allclose + +from scipy.conftest import skip_xp_invalid_arg +import scipy._external.array_api_extra as xpx +from scipy._lib._array_api import (make_xp_test_case, xp_default_dtype, is_jax, is_cupy, + eager_warns, xp_result_type, is_array_api_strict) +from scipy._lib._array_api_no_0d import xp_assert_close, xp_assert_equal +from scipy import stats +from scipy.stats._axis_nan_policy import SmallSampleWarning + + +lazy_xp_modules = [stats] + + +@make_xp_test_case(stats.chatterjeexi) +class TestChatterjeeXi: + @pytest.mark.parametrize('case', [ + dict(y_cont=True, statistic=-0.303030303030303, pvalue=0.9351329808526656), + dict(y_cont=False, statistic=0.07407407407407396, pvalue=0.3709859367123997)]) + @pytest.mark.parametrize('dtype', ['float32', 'float64', None]) + def test_against_R_XICOR(self, case, dtype, xp): + # Test against R package XICOR, e.g. + # library(XICOR) + # options(digits=16) + # x = c(0.11027287231363914, 0.8154770102474279, 0.7073943466920335, + # 0.6651317324378386, 0.6905752850115503, 0.06115250587536558, + # 0.5209906494474178, 0.3155763519785274, 0.18405731803625924, + # 0.8613557911541495) + # y = c(0.8402081904493103, 0.5946972833914318, 0.23481606164114155, + # 0.49754786197715384, 0.9146460831206026, 0.5848057749217579, + # 0.7620801065573549, 0.31410063302647495, 0.7935620302236199, + # 0.5423085761365468) + # xicor(x, y, ties=FALSE, pvalue=TRUE) + dtype = xp_default_dtype(xp) if dtype is None else getattr(xp, dtype) + rng = np.random.default_rng(25982435982346983) + x = rng.random(size=10) + y = (rng.random(size=10) if case['y_cont'] + else rng.integers(0, 5, size=10)) + + x, y = xp.asarray(x, dtype=dtype), xp.asarray(y, dtype=dtype) + res = stats.chatterjeexi(x, y, y_continuous=case['y_cont']) + + xp_assert_close(res.statistic, xp.asarray(case['statistic'], dtype=dtype)) + xp_assert_close(res.pvalue, xp.asarray(case['pvalue'], dtype=dtype)) + + @pytest.mark.parametrize('y_continuous', (False, True)) + def test_permutation_asymptotic(self, y_continuous): + # XICOR doesn't seem to perform the permutation test as advertised, so + # compare the result of a permutation test against an asymptotic test. + rng = np.random.default_rng(2524579827426) + n = np.floor(rng.uniform(100, 150)).astype(int) + shape = (2, n) + x = rng.random(size=shape) + y = (rng.random(size=shape) if y_continuous + else rng.integers(0, 10, size=shape)) + method = stats.PermutationMethod(rng=rng) + res = stats.chatterjeexi(x, y, method=method, + y_continuous=y_continuous, axis=-1) + ref = stats.chatterjeexi(x, y, y_continuous=y_continuous, axis=-1) + np.testing.assert_allclose(res.statistic, ref.statistic, rtol=1e-15) + np.testing.assert_allclose(res.pvalue, ref.pvalue, rtol=2e-2) + + def test_input_validation(self, xp): + rng = np.random.default_rng(25932435798274926) + x, y = rng.random(size=(2, 10)) + x, y = xp.asarray(x), xp.asarray(y) + + message = 'Array shapes are incompatible for broadcasting.|Incompatible shapes' + with pytest.raises((ValueError, TypeError), match=message): + stats.chatterjeexi(x, y[:-1]) + + if not (is_jax(xp) or is_cupy(xp)): + # jax misses out on some input validation from _axis_nan_policy decorator + # This also fails with cupy for reasons that need to be investigated. + message = '...axis 10 is out of bounds for array...|out of range' + with pytest.raises((ValueError, IndexError), match=message): + stats.chatterjeexi(x, y, axis=10) + + message = '`y_continuous` must be boolean.' + with pytest.raises(ValueError, match=message): + stats.chatterjeexi(x, y, y_continuous='a herring') + + message = "`method` must be 'asymptotic' or" + with pytest.raises(ValueError, match=message): + stats.chatterjeexi(x, y, method='ekki ekii') + + def test_special_cases(self, xp): + message = 'One or more sample arguments is too small...' + with pytest.warns(SmallSampleWarning, match=message): + res = stats.chatterjeexi(xp.asarray([1]), xp.asarray([2])) + + assert xp.isnan(res.statistic) + assert xp.isnan(res.pvalue) + + +@make_xp_test_case(stats.spearmanrho) +class TestSpearmanRho: + @pytest.mark.parametrize('alternative, statistic, pvalue', [ + ('two-sided', -0.2727272727272727, 0.4458383415428), + ('greater', -0.2727272727272727, 0.7770808292286), + ('less', -0.2727272727272727, 0.2229191707714), + ]) + @pytest.mark.parametrize('dtype', ['float32', 'float64', None]) + def test_against_R_cor_test(self, alternative, statistic, pvalue, dtype, xp): + # Test against R cor.test, e.g. + # options(digits=16) + # x = c(0.11027287231363914, 0.8154770102474279, 0.7073943466920335, + # 0.6651317324378386, 0.6905752850115503, 0.06115250587536558, + # 0.5209906494474178, 0.3155763519785274, 0.18405731803625924, + # 0.8613557911541495) + # y = c(0.8402081904493103, 0.5946972833914318, 0.23481606164114155, + # 0.49754786197715384, 0.9146460831206026, 0.5848057749217579, + # 0.7620801065573549, 0.31410063302647495, 0.7935620302236199, + # 0.5423085761365468) + # cor.test(x, y, method='spearman', alternative='t', exact=FALSE) + dtype = xp_default_dtype(xp) if dtype is None else getattr(xp, dtype) + rng = np.random.default_rng(25982435982346983) + x = rng.random(size=10) + y = rng.random(size=10) + + x, y = xp.asarray(x, dtype=dtype), xp.asarray(y, dtype=dtype) + res = stats.spearmanrho(x, y, alternative=alternative) + + xp_assert_close(res.statistic, xp.asarray(statistic, dtype=dtype)) + xp_assert_close(res.pvalue, xp.asarray(pvalue, dtype=dtype)) + + + @pytest.mark.parametrize('alternative, statistic, pvalue', [ + ('two-sided', -0.4857142857142857, 0.3555555555556), + ('greater', -0.4857142857142857, 0.8513888888889), + ('less', -0.4857142857142857, 0.1777777777778), + ]) + def test_against_R_cor_test_exact(self, alternative, statistic, pvalue): + xp = np # will test for multiple backends when gh-23772 merges + # Test against R cor.test exact=TRUE, e.g. + # options(digits=16) + # x = c(0.11027287231363914, 0.8154770102474279, 0.7073943466920335, + # 0.6651317324378386, 0.6905752850115503, 0.06115250587536558) + # y = c(0.5209906494474178, 0.3155763519785274, 0.18405731803625924, + # 0.8613557911541495, 0.8402081904493103, 0.5946972833914318) + # cor.test(x, y, method='spearman', alternative='t', exact=TRUE) + rng = np.random.default_rng(25982435982346983) + x = rng.random(size=6) + y = rng.random(size=6) + x, y = xp.asarray(x), xp.asarray(y) + + method = stats.PermutationMethod() + res = stats.spearmanrho(x, y, method=method, alternative=alternative) + + xp_assert_close(res.statistic, xp.asarray(statistic)) + xp_assert_close(res.pvalue, xp.asarray(pvalue)) + + + @pytest.mark.parametrize('alternative', ('two-sided', 'greater', 'less')) + @pytest.mark.parametrize('n', [9, 99, 999]) + def test_against_scipy_spearmanr(self, alternative, n, xp): + rng = np.random.default_rng(5982435982346983) + x = rng.integers(n//2, size=n) + y = rng.integers(n//2, size=n) + dtype = xp_default_dtype(xp) + + ref = stats.spearmanr(x, y, alternative=alternative) + ref_statistic = xp.asarray(ref.statistic, dtype=dtype) + ref_pvalue = xp.asarray(ref.pvalue, dtype=dtype) + + x = xp.asarray(x, dtype=dtype) + y = xp.asarray(y, dtype=dtype) + res = stats.spearmanrho(x, y, alternative=alternative) + + xp_assert_close(res.statistic, ref_statistic) + xp_assert_close(res.pvalue, ref_pvalue) + + + @pytest.mark.parametrize('axis', [-1, 0, 1]) + def test_other_backends_nd(self, axis, xp): + # NumPy n-d behavior is tested by `test_axis_nan_policy`; + # check other backends against it. + rng = np.random.default_rng(8243598346983259) + shape = (8, 9, 10) + x = rng.standard_normal(size=shape) + y = rng.standard_normal(size=shape) + res = stats.spearmanrho(xp.asarray(x), xp.asarray(y), axis=axis) + ref = stats.spearmanrho(x, y, axis=axis) + xp_assert_close(res.statistic, xp.asarray(ref.statistic), atol=1e-16) + xp_assert_close(res.pvalue, xp.asarray(ref.pvalue), atol=1e-16) + + + def test_input_validation(self, xp): + rng = np.random.default_rng(25932435798274926) + x, y = rng.random(size=(2, 10)) + x, y = xp.asarray(x), xp.asarray(y) + + msg = 'incompatible for broadcasting|Incompatible shapes|must be broadcastable' + with pytest.raises((ValueError, TypeError), match=msg): + stats.spearmanrho(x, y[:-1]) + + if not is_jax(xp): + # jax misses out on some input validation from _axis_nan_policy decorator + message = '...axis 10 is out of bounds for array...|out of range' + with pytest.raises((ValueError, IndexError), match=message): + stats.spearmanrho(x, y, axis=10) + + message = "`alternative` must be 'less', 'greater', or 'two-sided'." + with pytest.raises(ValueError, match=message): + stats.spearmanrho(x, y, alternative='alternative') + + message = ("`method` must be...") + with pytest.raises(ValueError, match=message): + stats.spearmanrho(x, y, method='method') + + def test_special_cases(self, xp): + def check_nan(res): + assert xp.isnan(res.statistic) + assert xp.isnan(res.pvalue) + + message = 'One or more sample arguments is too small...' + with eager_warns(SmallSampleWarning, match=message, xp=xp): + res = stats.spearmanrho(xp.asarray([1]), xp.asarray([2])) + check_nan(res) + + x = xp.asarray([1, 1, 1, 1, 1]) + y = xp.asarray([1, 2, 3, 4, 5]) + message = 'An input array is constant; the correlation coefficient...' + with eager_warns(stats.ConstantInputWarning, match=message, xp=xp): + res = stats.spearmanrho(x, y) + check_nan(res) + with eager_warns(stats.ConstantInputWarning, match=message, xp=xp): + res = stats.spearmanrho(y, x) + check_nan(res) + + +class RobustSlopesTest: + # 'iv, mask, warnings, consistency, gh19678, mstats' + def test_input_validation(self, xp): + pfun = getattr(stats, self.pfun) + other_method = 'joint' if pfun == stats.theilslopes else 'hierarchical' + msg = (f"method must be either '{other_method}' or 'separate'. " + "'joint_separate' is invalid.") + with pytest.raises(ValueError, match=msg): + pfun(xp.asarray([0, 1, 1]), method='joint_separate') + + @skip_xp_invalid_arg + @pytest.mark.parametrize("method", ['separate', 'other']) + def test_mask(self, method): + pfun = getattr(stats, self.pfun) + if method == 'other': + method = 'joint' if self.pfun =='theilslopes' else 'hierarchical' + + # Test for correct masking. + mask = np.asarray([False, False, True, False]) + y = np.ma.array([0, 1, 100, 1], mask=mask) + res = pfun(y, method=method) + ref = pfun(y[~mask], method=method) + + assert_allclose(res.slope, ref.slope) + assert_allclose(res.intercept, ref.intercept) + + def test_degenerate(self, xp): + # Test with degenerate input; see gh-15943 + pfun = getattr(stats, self.pfun) + res = pfun(xp.asarray([0, 1]), xp.asarray([0, 0])) + assert xp.all(xp.isnan(xp.stack(res))) + res = pfun(xp.asarray([0, 0, 0]), xp.asarray([0, 1, 0])) + xp_assert_equal(res[0], xp.asarray(0.)) + xp_assert_equal(res[1], xp.asarray(0.)) + if pfun == stats.theilslopes: + xp_assert_equal(res[2], xp.asarray(xp.nan)) + xp_assert_equal(res[3], xp.asarray(xp.nan)) + + def test_namedtuple_consistency(self, xp): + """ + Simple test to ensure tuple backwards-compatibility of the returned object. + """ + pfun = getattr(stats, self.pfun) + + y = xp.asarray([1, 2, 4]) + x = xp.asarray([4, 6, 8]) + + result = pfun(y, x) + + # note both returned values are distinct here + xp_assert_equal(result[0], result.slope) + xp_assert_equal(result[1], result.intercept) + if pfun == stats.theilslopes: + xp_assert_equal(result[2], result.low_slope) + xp_assert_equal(result[3], result.high_slope) + + def test_gh19678_uint8(self, xp): + # `theilslopes` returned unexpected results when `y` was an unsigned type. + # Check that this is resolved. + pfun = getattr(stats, self.pfun) + rng = np.random.default_rng(2549824598234528) + y = xp.asarray(rng.integers(0, 255, size=10, dtype=np.uint8)) + res = pfun(y, y) + xp_assert_close(res.slope, xp.asarray(1.)) + + @pytest.mark.parametrize("method", ['separate', 'other']) + @pytest.mark.parametrize("case_number", [0, 1, 2, 3]) + def test_against_mstats(self, method, case_number, xp): + pfun = getattr(stats, self.pfun) + if method == 'other': + method = 'joint' if self.pfun =='theilslopes' else 'hierarchical' + + rng = np.random.default_rng(349824598234528554) + match case_number: + case 0: # no x + x = None + y = rng.random(100) + case 1: # no ties + x = rng.random(100) + y = rng.random(100) + case 2: # no x ties + x = rng.random(100) + y = rng.integers(50, size=100) + case 3: # ties in x and y + x = rng.integers(25, size=100) + y = rng.integers(50, size=100) + + ref = pfun(y, x, method=method) + + y, x = xp.asarray(y), x if x is None else xp.asarray(x) + if (is_array_api_strict(xp) and xp.isdtype(y.dtype, 'integral') + and xp.isdtype(x.dtype, 'real floating')): + # array API strict doesn't allow mixed type promotion, so promote for it + y = xp.astype(y, xp.float64) + + dtype = xp_result_type(y, x, force_floating=True, xp=xp) + + res = pfun(y, x, method=method, axis=-1) + + assert res.slope != 0 # ensure randomly generated test case isn't trivial + xp_assert_close(res.slope, xp.asarray(ref.slope, dtype=dtype)) + xp_assert_close(res.intercept, xp.asarray(ref.intercept, dtype=dtype)) + if pfun == stats.theilslopes: + xp_assert_close(res.low_slope, xp.asarray(ref.low_slope, dtype=dtype)) + xp_assert_close(res.high_slope, xp.asarray(ref.high_slope, dtype=dtype)) + + +@make_xp_test_case(stats.theilslopes) +class TestTheilslopes(RobustSlopesTest): + pfun = 'theilslopes' + def test_theilslopes(self, xp): + # Test for basic slope and intercept. + y = xp.asarray([0, 1, 1]) + slope, intercept, lower, upper = stats.theilslopes(y) + xp_assert_close(slope, xp.asarray(0.5)) + xp_assert_close(intercept, xp.asarray(0.5)) + + slope, intercept, lower, upper = stats.theilslopes(y, method='joint') + xp_assert_close(slope, xp.asarray(0.5)) + xp_assert_close(intercept, xp.asarray(0.0)) + + # Test of confidence intervals from example in Sen (1968). + x = xp.asarray([1, 2, 3, 4, 10, 12, 18]) + y = xp.asarray([9, 15, 19, 20, 45, 55, 78]) + slope, intercept, lower, upper = stats.theilslopes(y, x, 0.07) + xp_assert_close(slope, xp.asarray(4.0)) + xp_assert_close(intercept, xp.asarray(4.0)) + xp_assert_close(upper, xp.asarray(4.38), rtol=5e-3) + xp_assert_close(lower, xp.asarray(3.71), rtol=5e-3) + + slope, intercept, lower, upper = stats.theilslopes(y, x, 0.07, + method='joint') + xp_assert_close(slope, xp.asarray(4.0)) + xp_assert_close(intercept, xp.asarray(6.0)) + xp_assert_close(upper, xp.asarray(4.38), rtol=5e-3) + xp_assert_close(lower, xp.asarray(3.71), rtol=5e-3) + + +@make_xp_test_case(stats.siegelslopes) +class TestSiegelslopes(RobustSlopesTest): + pfun = 'siegelslopes' + def test_siegelslopes(self, xp): + # method should be exact for straight line + y = 2 * xp.arange(10.) + 0.5 + slope, intercept = stats.siegelslopes(y) + xp_assert_close(slope, xp.asarray(2.0)) + xp_assert_close(intercept, xp.asarray(0.5)) + slope, intercept = stats.siegelslopes(y, method='separate') + xp_assert_close(slope, xp.asarray(2.0)) + xp_assert_close(intercept, xp.asarray(0.5)) + + x = 2 * xp.arange(10.) + y = 5 * x - 3.0 + slope, intercept = stats.siegelslopes(y, x) + xp_assert_close(slope, xp.asarray(5.0)) + xp_assert_close(intercept, xp.asarray(-3.0)) + slope, intercept = stats.siegelslopes(y, x, method='separate') + xp_assert_close(slope, xp.asarray(5.0)) + xp_assert_close(intercept, xp.asarray(-3.0)) + + # method is robust to outliers: breakdown point of 50% + y = xpx.at(y)[:4].set(1000.) + xp_assert_close(slope, xp.asarray(5.0)) + xp_assert_close(intercept, xp.asarray(-3.0)) + + # if there are no outliers, results should be comparable to linregress + x = np.arange(10.) + y = -2.3 + 0.3 * x + stats.norm.rvs(size=10, random_state=231) + slope_ols, intercept_ols, _, _, _ = stats.linregress(x, y) + + y, x = xp.asarray(y), xp.asarray(x) + slope, intercept = stats.siegelslopes(y, x) + xp_assert_close(slope, xp.asarray(slope_ols), rtol=0.1) + xp_assert_close(intercept, xp.asarray(intercept_ols), rtol=0.1) + + slope, intercept = stats.siegelslopes(y, x, method='separate') + xp_assert_close(slope, xp.asarray(slope_ols), rtol=0.1) + xp_assert_close(intercept, xp.asarray(intercept_ols), rtol=0.1) + + def test_method(self, xp): + # Test that distinguishes between the methods. Reference values generated with + # SciPy 1.17 (results unchanged since at least 1.12, so assumed "correct") to + # detect *changes*. + y = xp.asarray([2, 1, 3, 0, 4]) + slope, intercept = stats.siegelslopes(y, method='separate') + xp_assert_close(slope, xp.asarray(0.25)) + xp_assert_close(intercept, xp.asarray(1.75)) + slope, intercept = stats.siegelslopes(y, method='hierarchical') + xp_assert_close(slope, xp.asarray(0.25)) # always the same between the methods + xp_assert_close(intercept, xp.asarray(2.0)) diff --git a/v12x/pylib/scipy/stats/tests/test_crosstab.py b/v12x/pylib/scipy/stats/tests/test_crosstab.py new file mode 100644 index 0000000000000000000000000000000000000000..42dc525045b6e455c0bf498adcac5faabe2106ff --- /dev/null +++ b/v12x/pylib/scipy/stats/tests/test_crosstab.py @@ -0,0 +1,118 @@ +import pytest +import numpy as np +from numpy.testing import assert_array_equal, assert_equal +from scipy.stats.contingency import crosstab + + +@pytest.mark.parametrize('sparse', [False, True]) +def test_crosstab_basic(sparse): + a = [0, 0, 9, 9, 0, 0, 9] + b = [2, 1, 3, 1, 2, 3, 3] + expected_avals = [0, 9] + expected_bvals = [1, 2, 3] + expected_count = np.array([[1, 2, 1], + [1, 0, 2]]) + (avals, bvals), count = crosstab(a, b, sparse=sparse) + assert_array_equal(avals, expected_avals) + assert_array_equal(bvals, expected_bvals) + if sparse: + assert_array_equal(count.toarray(), expected_count) + else: + assert_array_equal(count, expected_count) + + +@pytest.mark.parametrize('sparse', [False, True]) +def test_crosstab_basic_1d(sparse): + # Verify that a single input sequence works as expected. + x = [1, 2, 3, 1, 2, 3, 3] + expected_xvals = [1, 2, 3] + expected_count = np.array([2, 2, 3]) + (xvals,), count = crosstab(x, sparse=sparse) + assert_array_equal(xvals, expected_xvals) + if sparse: + assert_array_equal(count.toarray(), expected_count) + else: + assert_array_equal(count, expected_count) + + +@pytest.mark.parametrize('sparse', [False, True]) +def test_crosstab_basic_3d(sparse): + # Verify the function for three input sequences. + a = 'a' + b = 'b' + x = [0, 0, 9, 9, 0, 0, 9, 9] + y = [a, a, a, a, b, b, b, a] + z = [1, 2, 3, 1, 2, 3, 3, 1] + expected_xvals = [0, 9] + expected_yvals = [a, b] + expected_zvals = [1, 2, 3] + expected_count = np.array([[[1, 1, 0], + [0, 1, 1]], + [[2, 0, 1], + [0, 0, 1]]]) + (xvals, yvals, zvals), count = crosstab(x, y, z, sparse=sparse) + assert_array_equal(xvals, expected_xvals) + assert_array_equal(yvals, expected_yvals) + assert_array_equal(zvals, expected_zvals) + if sparse: + assert_array_equal(count.toarray(), expected_count) + else: + assert_array_equal(count, expected_count) + + +@pytest.mark.parametrize('sparse', [False, True]) +def test_crosstab_levels(sparse): + a = [0, 0, 9, 9, 0, 0, 9] + b = [1, 2, 3, 1, 2, 3, 3] + expected_avals = [0, 9] + expected_bvals = [0, 1, 2, 3] + expected_count = np.array([[0, 1, 2, 1], + [0, 1, 0, 2]]) + (avals, bvals), count = crosstab(a, b, levels=[None, [0, 1, 2, 3]], + sparse=sparse) + assert_array_equal(avals, expected_avals) + assert_array_equal(bvals, expected_bvals) + if sparse: + assert_array_equal(count.toarray(), expected_count) + else: + assert_array_equal(count, expected_count) + + +@pytest.mark.parametrize('sparse', [False, True]) +def test_crosstab_extra_levels(sparse): + # The pair of values (-1, 3) will be ignored, because we explicitly + # request the counted `a` values to be [0, 9]. + a = [0, 0, 9, 9, 0, 0, 9, -1] + b = [1, 2, 3, 1, 2, 3, 3, 3] + expected_avals = [0, 9] + expected_bvals = [0, 1, 2, 3] + expected_count = np.array([[0, 1, 2, 1], + [0, 1, 0, 2]]) + (avals, bvals), count = crosstab(a, b, levels=[[0, 9], [0, 1, 2, 3]], + sparse=sparse) + assert_array_equal(avals, expected_avals) + assert_array_equal(bvals, expected_bvals) + if sparse: + assert_array_equal(count.toarray(), expected_count) + else: + assert_array_equal(count, expected_count) + + +def test_validation_at_least_one(): + with pytest.raises(TypeError, match='At least one'): + crosstab() + + +def test_validation_same_lengths(): + with pytest.raises(ValueError, match='must have the same length'): + crosstab([1, 2], [1, 2, 3, 4]) + + +def test_validation_len_levels_matches_args(): + with pytest.raises(ValueError, match='number of input sequences'): + crosstab([0, 1, 1], [8, 8, 9], levels=([0, 1, 2, 3],)) + + +def test_result(): + res = crosstab([0, 1], [1, 2]) + assert_equal((res.elements, res.count), res) diff --git a/v12x/pylib/scipy/stats/tests/test_device_dtype.py b/v12x/pylib/scipy/stats/tests/test_device_dtype.py new file mode 100644 index 0000000000000000000000000000000000000000..c20bfc1104a7c3b0dcd2ee6484aed9169380eeef --- /dev/null +++ b/v12x/pylib/scipy/stats/tests/test_device_dtype.py @@ -0,0 +1,248 @@ +import pytest +import numpy as np +from scipy import stats + +from scipy._lib._array_api import xp_device, is_array_api_strict, is_torch +from scipy.stats._stats_py import _xp_mean, _xp_var + +skip_xp_backends = pytest.mark.skip_xp_backends + +dtypes = ['float32', 'float64'] + + +def get_arrays(n_arrays, *, dtype=np.float64, xp=np, shape=(30,), device=None, + seed=84912165484321): + rng = np.random.default_rng(seed) + + datas = [] + for i in range(n_arrays): + data = 10*rng.random(size=shape) + if xp.isdtype(dtype, 'complex floating'): + data = data * 10j*rng.standard_normal(size=shape) + data = xp.asarray(data, dtype=dtype, device=device) + datas.append(data) + + return datas + + +@pytest.mark.parametrize('fun, kwargs', [(stats.gmean, {}), + (stats.hmean, {}), + (stats.pmean, {'p': 2}), + (_xp_mean, {}), + ]) +@pytest.mark.parametrize('dtype', dtypes) +@pytest.mark.uses_xp_capabilities(False, reason="not used in this test yet") +def test_xmean(fun, kwargs, dtype, xp, devices): + dtype = getattr(xp, dtype) + for device in devices: + x, weights = get_arrays(2, device=device, dtype=dtype, xp=xp) + res = fun(x, weights=weights, **kwargs) + assert xp_device(res) == xp_device(x) + assert x.dtype == dtype + + +@skip_xp_backends('array_api_strict', + reason="special functions don't work with 'device1'") +@pytest.mark.parametrize('nargs', [1, 2]) +@pytest.mark.parametrize('dtype', dtypes) +@pytest.mark.uses_xp_capabilities(False, reason="not used in this test yet") +def test_entropy(nargs, dtype, xp, devices): + dtype = getattr(xp, dtype) + for device in devices: + args = get_arrays(nargs, device=device, dtype=dtype, xp=xp) + res = stats.entropy(*args) + assert xp_device(res) == xp_device(args[0]) + assert res.dtype == dtype + + +@pytest.mark.parametrize('dtype', dtypes) +@pytest.mark.uses_xp_capabilities(False, reason="not used in this test yet") +def test_directional_stats(dtype, xp, devices): + dtype = getattr(xp, dtype) + for device in devices: + x = get_arrays(1, shape=(30, 3), device=device, dtype=dtype, xp=xp)[0] + res = stats.directional_stats(x) + assert xp_device(res.mean_direction) == xp_device(x) + assert res.mean_direction.dtype == dtype + assert xp_device(res.mean_resultant_length) == xp_device(x) + assert res.mean_resultant_length.dtype == dtype + + +@pytest.mark.parametrize('method', [ + "vasicek", + "van es", + pytest.param( + "correa", + marks=[ + skip_xp_backends("array_api_strict", reason="Invalid fancy indexing"), + skip_xp_backends("dask.array", reason="Invalid fancy indexing"), + ], + ), + "ebrahimi", + "auto", +]) +@pytest.mark.parametrize('dtype', dtypes) +@pytest.mark.uses_xp_capabilities(False, reason="not used in this test yet") +def test_differential_entropy(method, dtype, xp, devices): + dtype = getattr(xp, dtype) + for device in devices: + values = get_arrays(1, device=device, dtype=dtype, xp=xp)[0] + res = stats.differential_entropy(values, method=method) + assert xp_device(res) == xp_device(values) + assert values.dtype == dtype + + +@skip_xp_backends('dask.array', reason='no take_along_axis') +@pytest.mark.parametrize('method', ['inverted_cdf', 'averaged_inverted_cdf', + 'closest_observation', 'interpolated_inverted_cdf', + 'hazen', 'weibull', 'linear', 'median_unbiased', + 'normal_unbiased', 'harrell-davis']) +@pytest.mark.parametrize('dtype', dtypes) +@pytest.mark.uses_xp_capabilities(False, reason="not used in this test yet") +def test_quantile(method, dtype, xp, devices): + if (is_array_api_strict(xp) or is_torch(xp)) and method == 'harrell-davis': + pytest.skip("'harrell-davis' not currently supported on GPU.") + + dtype = getattr(xp, dtype) + for device in devices: + values = get_arrays(1, device=device, dtype=dtype, xp=xp)[0] + res = stats.quantile(values, 0.5, method=method) + assert xp_device(res) == xp_device(values) + assert values.dtype == dtype + + +@pytest.mark.parametrize('dtype', dtypes) +@pytest.mark.uses_xp_capabilities(False, reason="not used in this test yet") +def test_boxcox_llf(dtype, xp, devices): + dtype = getattr(xp, dtype) + for device in devices: + data = get_arrays(1, device=device, dtype=dtype, xp=xp)[0] + res = stats.boxcox_llf(1, data) + assert xp_device(res) == xp_device(data) + assert data.dtype == dtype + + +@pytest.mark.parametrize('fun, kwargs', + [(stats.moment, {'order': 2}), + (stats.skew, {}), + (stats.skew, {'bias': False}), + (stats.kurtosis, {}), + (stats.kurtosis, {'bias': False}), + (stats.sem, {}), + (stats.kstat, {'n': 1}), + (stats.kstat, {'n': 2}), + (stats.kstat, {'n': 3}), + (stats.kstat, {'n': 4}), + (stats.kstatvar, {'n': 1}), + (stats.kstatvar, {'n': 2}), + (stats.circmean, {}), + (stats.circvar, {}), + (stats.circstd, {}), + (_xp_var, {}), + (stats.gstd, {}), + (stats.variation, {}), + (stats.tmean, {'limits': (0.1, 0.9)}), + (stats.tvar, {'limits': (0.1, 0.9)}), + (stats.tmin, {'lowerlimit': 0.1}), + (stats.tmax, {'upperlimit': 0.9}), + (stats.tstd, {'limits': (0.1, 0.9)}), + (stats.tsem, {'limits': (0.1, 0.9)}), + ]) +@pytest.mark.parametrize('dtype', dtypes) +@pytest.mark.uses_xp_capabilities(False, reason="not used in this test yet") +def test_one_in_one_out(fun, kwargs, dtype, xp, devices): + dtype = getattr(xp, dtype) + for device in devices: + array = get_arrays(1, device=device, dtype=dtype, xp=xp)[0] + res = fun(array, **kwargs) + assert xp_device(res) == xp_device(array) + assert res.dtype == dtype + + +@pytest.mark.parametrize('dtype', dtypes) +@pytest.mark.uses_xp_capabilities(False, reason="not used in this test yet") +def test_describe(dtype, xp, devices): + dtype = getattr(xp, dtype) + for device in devices: + array = get_arrays(1, shape=10, dtype=dtype, device=device, xp=xp)[0] + res = stats.describe(array, axis=-1) + + assert xp_device(res.nobs) == xp_device(array) + assert xp_device(res.minmax[0]) == xp_device(array) + assert xp_device(res.minmax[1]) == xp_device(array) + assert xp_device(res.variance) == xp_device(array) + assert xp_device(res.skewness) == xp_device(array) + assert xp_device(res.kurtosis) == xp_device(array) + + assert res.minmax[0].dtype == dtype + assert res.minmax[1].dtype == dtype + assert res.variance.dtype == dtype + assert res.skewness.dtype == dtype + assert res.kurtosis.dtype == dtype + + +@pytest.mark.parametrize('fun', [stats.zscore, stats.gzscore, stats.zmap]) +@pytest.mark.parametrize('dtype', dtypes) +@pytest.mark.uses_xp_capabilities(False, reason="not used in this test yet") +def test_zscore(fun, dtype, xp, devices): + dtype = getattr(xp, dtype) + for device in devices: + n = 2 if (fun == stats.zmap) else 1 + arrays = get_arrays(n, device=device, dtype=dtype, xp=xp) + res = fun(*arrays) + assert xp_device(res) == xp_device(arrays[0]) + assert res.dtype == dtype + + +@skip_xp_backends('array_api_strict', reason="special/_support_alternative_backends") +@skip_xp_backends(cpu_only=True, exceptions=['cupy', 'jax.numpy']) +@pytest.mark.parametrize('f_name', ['ttest_1samp', 'ttest_rel', 'ttest_ind', 'skewtest', + 'kurtosistest', 'normaltest', 'jarque_bera', + 'bartlett', 'pearsonr', 'chisquare', + 'power_divergence']) +@pytest.mark.parametrize('dtype', dtypes) +@pytest.mark.uses_xp_capabilities(False, reason="not used in this test yet") +def test_hypothesis_tests(f_name, dtype, xp, devices): + dtype = getattr(xp, dtype) + for device in devices: + f = getattr(stats, f_name) + + n = 2 if f_name in {'ttest_1samp', 'ttest_rel', 'ttest_ind', 'bartlett', + 'pearsonr', 'chisquare'} else 1 + + arrays = get_arrays(n, xp=xp, dtype=dtype, device=device) + if f_name == 'ttest_1samp': + arrays[1] = xp.mean(arrays[1]) + if f_name == 'chisquare': + arrays[1] = xp.sum(arrays[0]) * arrays[1] / xp.sum(arrays[1]) + + res = f(*arrays) + assert xp_device(res.statistic) == xp_device(arrays[0]) + assert xp_device(res.pvalue) == xp_device(arrays[0]) + assert res.statistic.dtype == dtype + assert res.pvalue.dtype == dtype + + if f_name in {'ttest_1samp', 'ttest_rel', 'ttest_ind', 'pearsonr'}: + res_ci = res.confidence_interval() + assert xp_device(res_ci.low) == xp_device(arrays[0]) + assert xp_device(res_ci.high) == xp_device(arrays[0]) + assert res_ci.low.dtype == dtype + assert res_ci.high.dtype == dtype + + +@skip_xp_backends('array_api_strict', + reason="special functions don't work with 'device1'") +@skip_xp_backends(cpu_only=True, exceptions=['cupy', 'jax.numpy']) +@pytest.mark.parametrize('method', ['fisher', 'pearson', 'tippett', 'stouffer', + 'mudholkar_george']) +@pytest.mark.parametrize('dtype', dtypes) +@pytest.mark.uses_xp_capabilities(False, reason="not used in this test yet") +def test_combine_pvalues(method, dtype, xp, devices): + dtype = getattr(xp, dtype) + for device in devices: + pvalues = get_arrays(1, xp=xp, dtype=dtype, device=device)[0] / 10 + res = stats.combine_pvalues(pvalues, method=method) + assert xp_device(res.statistic) == xp_device(pvalues) + assert xp_device(res.pvalue) == xp_device(pvalues) + assert res.statistic.dtype == dtype + assert res.pvalue.dtype == dtype diff --git a/v12x/pylib/scipy/stats/tests/test_discrete_basic.py b/v12x/pylib/scipy/stats/tests/test_discrete_basic.py new file mode 100644 index 0000000000000000000000000000000000000000..8e67c60a21d193ee3f304b80f6660110fb95b177 --- /dev/null +++ b/v12x/pylib/scipy/stats/tests/test_discrete_basic.py @@ -0,0 +1,625 @@ +import warnings + +import numpy.testing as npt +from numpy.testing import assert_allclose + +import numpy as np +import pytest + +from scipy import stats +from scipy.special import _ufuncs +from .common_tests import (check_normalization, check_moment, + check_mean_expect, + check_var_expect, check_skew_expect, + check_kurt_expect, check_entropy, + check_private_entropy, check_edge_support, + check_named_args, check_random_state_property, + check_pickling, check_rvs_broadcast, + check_freezing,) +from scipy.stats._distr_params import distdiscrete, invdistdiscrete +from scipy.stats._distn_infrastructure import rv_discrete_frozen + +vals = ([1, 2, 3, 4], [0.1, 0.2, 0.3, 0.4]) +distdiscrete += [[stats.rv_discrete(values=vals), ()]] + +# For these distributions, test_discrete_basic only runs with test mode full +distslow = {'nhypergeom'} + +# Override number of ULPs adjustment for `check_cdf_ppf` +roundtrip_cdf_ppf_exceptions = {'nbinom': 30} + +def cases_test_discrete_basic(): + seen = set() + for distname, arg in distdiscrete: + if distname in distslow: + yield pytest.param(distname, arg, distname, marks=pytest.mark.slow) + else: + yield distname, arg, distname not in seen + seen.add(distname) + + +@pytest.mark.parametrize('distname,arg,first_case', list(cases_test_discrete_basic())) +def test_discrete_basic(distname, arg, first_case, num_parallel_threads): + if (isinstance(distname, str) and distname.startswith('nchypergeom') + and num_parallel_threads > 1): + pytest.skip(reason='nchypergeom has a global random generator') + + try: + distfn = getattr(stats, distname) + except TypeError: + distfn = distname + distname = 'sample distribution' + rng = np.random.RandomState(9765456) + rvs = distfn.rvs(*arg, size=2000, random_state=rng) + supp = np.unique(rvs) + m, v = distfn.stats(*arg) + check_cdf_ppf(distfn, arg, supp, distname + ' cdf_ppf') + + check_pmf_cdf(distfn, arg, distname) + check_oth(distfn, arg, supp, distname + ' oth') + check_edge_support(distfn, arg) + + alpha = 0.01 + check_discrete_chisquare(distfn, arg, rvs, alpha, + distname + ' chisquare') + + if first_case: + locscale_defaults = (0,) + meths = [distfn.pmf, distfn.logpmf, distfn.cdf, distfn.logcdf, + distfn.logsf] + # make sure arguments are within support + # for some distributions, this needs to be overridden + spec_k = {'randint': 11, 'hypergeom': 4, 'bernoulli': 0, + 'nchypergeom_wallenius': 6} + k = spec_k.get(distname, 1) + check_named_args(distfn, k, arg, locscale_defaults, meths) + if distname != 'sample distribution': + check_scale_docstring(distfn) + if num_parallel_threads == 1: + check_random_state_property(distfn, arg) + if distname not in {'poisson_binom'}: # can't be pickled + check_pickling(distfn, arg) + check_freezing(distfn, arg) + + # Entropy + check_entropy(distfn, arg, distname) + if distfn.__class__._entropy != stats.rv_discrete._entropy: + check_private_entropy(distfn, arg, stats.rv_discrete) + + +@pytest.mark.parametrize('distname,arg', distdiscrete) +def test_moments(distname, arg): + try: + distfn = getattr(stats, distname) + except TypeError: + distfn = distname + distname = 'sample distribution' + m, v, s, k = distfn.stats(*arg, moments='mvsk') + check_normalization(distfn, arg, distname) + + # compare `stats` and `moment` methods + check_moment(distfn, arg, m, v, distname) + check_mean_expect(distfn, arg, m, distname) + check_var_expect(distfn, arg, m, v, distname) + check_skew_expect(distfn, arg, m, v, s, distname) + with warnings.catch_warnings(): + if distname in ['zipf', 'betanbinom']: + warnings.simplefilter("ignore", RuntimeWarning) + check_kurt_expect(distfn, arg, m, v, k, distname) + + # frozen distr moments + check_moment_frozen(distfn, arg, m, 1) + check_moment_frozen(distfn, arg, v+m*m, 2) + + +@pytest.mark.parametrize('dist,shape_args', distdiscrete) +def test_rvs_broadcast(dist, shape_args): + # If shape_only is True, it means the _rvs method of the + # distribution uses more than one random number to generate a random + # variate. That means the result of using rvs with broadcasting or + # with a nontrivial size will not necessarily be the same as using the + # numpy.vectorize'd version of rvs(), so we can only compare the shapes + # of the results, not the values. + # Whether or not a distribution is in the following list is an + # implementation detail of the distribution, not a requirement. If + # the implementation the rvs() method of a distribution changes, this + # test might also have to be changed. + shape_only = dist in ['betabinom', 'betanbinom', 'skellam', 'yulesimon', + 'dlaplace', 'nchypergeom_fisher', + 'nchypergeom_wallenius', 'poisson_binom'] + try: + distfunc = getattr(stats, dist) + except TypeError: + distfunc = dist + dist = f'rv_discrete(values=({dist.xk!r}, {dist.pk!r}))' + loc = np.zeros(2) + nargs = distfunc.numargs + allargs = [] + bshape = [] + + if dist == 'poisson_binom': + # normal rules apply except the last axis of `p` is ignored + p = np.full((3, 1, 10), 0.5) + allargs = (p, loc) + bshape = (3, 2) + check_rvs_broadcast(distfunc, dist, allargs, + bshape, shape_only, [np.dtype(int)]) + return + + # Generate shape parameter arguments... + for k in range(nargs): + shp = (k + 3,) + (1,)*(k + 1) + param_val = shape_args[k] + allargs.append(np.full(shp, param_val)) + bshape.insert(0, shp[0]) + allargs.append(loc) + bshape.append(loc.size) + # bshape holds the expected shape when loc, scale, and the shape + # parameters are all broadcast together. + check_rvs_broadcast( + distfunc, dist, allargs, bshape, shape_only, [np.dtype(int)] + ) + + +@pytest.mark.parametrize('dist,args', distdiscrete) +def test_ppf_with_loc(dist, args): + try: + distfn = getattr(stats, dist) + except TypeError: + distfn = dist + #check with a negative, no and positive relocation. + rng = np.random.default_rng(5108587887) + + re_locs = [rng.integers(-10, -1), 0, rng.integers(1, 10)] + _a, _b = distfn.support(*args) + for loc in re_locs: + npt.assert_array_equal( + [_a-1+loc, _b+loc], + [distfn.ppf(0.0, *args, loc=loc), distfn.ppf(1.0, *args, loc=loc)] + ) + + +@pytest.mark.parametrize('dist, args', distdiscrete) +def test_isf_with_loc(dist, args): + try: + distfn = getattr(stats, dist) + except TypeError: + distfn = dist + # check with a negative, no and positive relocation. + rng = np.random.default_rng(4030503535) + re_locs = [rng.integers(-10, -1), 0, rng.integers(1, 10)] + _a, _b = distfn.support(*args) + for loc in re_locs: + expected = _b + loc, _a - 1 + loc + res = distfn.isf(0., *args, loc=loc), distfn.isf(1., *args, loc=loc) + npt.assert_array_equal(expected, res) + # test broadcasting behaviour + re_locs = [rng.integers(-10, -1, size=(5, 3)), + np.zeros((5, 3)), + rng.integers(1, 10, size=(5, 3))] + _a, _b = distfn.support(*args) + for loc in re_locs: + expected = _b + loc, _a - 1 + loc + res = distfn.isf(0., *args, loc=loc), distfn.isf(1., *args, loc=loc) + npt.assert_array_equal(expected, res) + + +def check_cdf_ppf(distfn, arg, supp, msg): + # supp is assumed to be an array of integers in the support of distfn + # (but not necessarily all the integers in the support). + # This test assumes that the PMF of any value in the support of the + # distribution is greater than 1e-8. + + # cdf is a step function, and ppf(q) = min{k : cdf(k) >= q, k integer} + cdf_supp = distfn.cdf(supp, *arg) + # In very rare cases, the finite precision calculation of ppf(cdf(supp)) + # can produce an array in which an element is off by one. We nudge the + # CDF values down by a few ULPs help to avoid this. + n_ulps = roundtrip_cdf_ppf_exceptions.get(distfn.name, 15) + cdf_supp0 = cdf_supp - n_ulps*np.spacing(cdf_supp) + npt.assert_array_equal(distfn.ppf(cdf_supp0, *arg), + supp, msg + '-roundtrip') + # Repeat the same calculation, but with the CDF values decreased by 1e-8. + npt.assert_array_equal(distfn.ppf(distfn.cdf(supp, *arg) - 1e-8, *arg), + supp, msg + '-roundtrip') + + if not hasattr(distfn, 'xk'): + _a, _b = distfn.support(*arg) + supp1 = supp[supp < _b] + npt.assert_array_equal(distfn.ppf(distfn.cdf(supp1, *arg) + 1e-8, *arg), + supp1 + distfn.inc, msg + ' ppf-cdf-next') + + +def check_pmf_cdf(distfn, arg, distname): + if hasattr(distfn, 'xk'): + index = distfn.xk + else: + startind = int(distfn.ppf(0.01, *arg) - 1) + index = list(range(startind, startind + 10)) + cdfs = distfn.cdf(index, *arg) + pmfs_cum = distfn.pmf(index, *arg).cumsum() + + atol, rtol = 1e-10, 1e-10 + if distname == 'skellam': # ncx2 accuracy + atol, rtol = 1e-5, 1e-5 + npt.assert_allclose(cdfs - cdfs[0], pmfs_cum - pmfs_cum[0], + atol=atol, rtol=rtol) + + # also check that pmf at non-integral k is zero + k = np.asarray(index) + k_shifted = k[:-1] + np.diff(k)/2 + npt.assert_equal(distfn.pmf(k_shifted, *arg), 0) + + # better check frozen distributions, and also when loc != 0 + loc = 0.5 + dist = distfn(loc=loc, *arg) + npt.assert_allclose(dist.pmf(k[1:] + loc), np.diff(dist.cdf(k + loc))) + npt.assert_equal(dist.pmf(k_shifted + loc), 0) + + +def check_moment_frozen(distfn, arg, m, k): + npt.assert_allclose(distfn(*arg).moment(k), m, + atol=1e-10, rtol=1e-10) + + +def check_oth(distfn, arg, supp, msg): + # checking other methods of distfn + npt.assert_allclose(distfn.sf(supp, *arg), 1. - distfn.cdf(supp, *arg), + atol=1e-10, rtol=1e-10) + + q = np.linspace(0.01, 0.99, 20) + npt.assert_allclose(distfn.isf(q, *arg), distfn.ppf(1. - q, *arg), + atol=1e-10, rtol=1e-10) + + median_sf = distfn.isf(0.5, *arg) + npt.assert_(distfn.sf(median_sf - 1, *arg) > 0.5) + npt.assert_(distfn.cdf(median_sf + 1, *arg) > 0.5) + + +def check_discrete_chisquare(distfn, arg, rvs, alpha, msg): + """Perform chisquare test for random sample of a discrete distribution + + Parameters + ---------- + distname : str + name of distribution function + arg : sequence + parameters of distribution + alpha : float + significance level, threshold for p-value + + Returns + ------- + result : bool + 0 if test passes, 1 if test fails + + """ + wsupp = 0.05 + + # construct intervals with minimum mass `wsupp`. + # intervals are left-half-open as in a cdf difference + _a, _b = distfn.support(*arg) + lo = int(max(_a, -1000)) + high = int(min(_b, 1000)) + 1 + distsupport = range(lo, high) + last = 0 + distsupp = [lo] + distmass = [] + for ii in distsupport: + current = distfn.cdf(ii, *arg) + if current - last >= wsupp - 1e-14: + distsupp.append(ii) + distmass.append(current - last) + last = current + if current > (1 - wsupp): + break + if distsupp[-1] < _b: + distsupp.append(_b) + distmass.append(1 - last) + distsupp = np.array(distsupp) + distmass = np.array(distmass) + + # convert intervals to right-half-open as required by histogram + histsupp = distsupp + 1e-8 + histsupp[0] = _a + + # find sample frequencies and perform chisquare test + freq, hsupp = np.histogram(rvs, histsupp) + chis, pval = stats.chisquare(np.array(freq), len(rvs)*distmass) + + npt.assert_( + pval > alpha, + f'chisquare - test for {msg} at arg = {str(arg)} with pval = {str(pval)}' + ) + + +def check_scale_docstring(distfn): + if distfn.__doc__ is not None: + # Docstrings can be stripped if interpreter is run with -OO + npt.assert_('scale' not in distfn.__doc__) + + +@pytest.mark.parametrize('method', ['pmf', 'logpmf', 'cdf', 'logcdf', + 'sf', 'logsf', 'ppf', 'isf']) +@pytest.mark.parametrize('distname, args', distdiscrete) +def test_methods_with_lists(method, distname, args): + # Test that the discrete distributions can accept Python lists + # as arguments. + try: + dist = getattr(stats, distname) + except TypeError: + return + dist_method = getattr(dist, method) + if method in ['ppf', 'isf']: + z = [0.1, 0.2] + else: + z = [0, 1] + p2 = [[p]*2 for p in args] + loc = [0, 1] + result = dist_method(z, *p2, loc=loc) + npt.assert_allclose(result, + [dist_method(*v) for v in zip(z, *p2, loc)], + rtol=1e-15, atol=1e-15) + + +@pytest.mark.parametrize('distname, args', invdistdiscrete) +def test_cdf_gh13280_regression(distname, args): + # Test for nan output when shape parameters are invalid + dist = getattr(stats, distname) + x = np.arange(-2, 15) + vals = dist.cdf(x, *args) + expected = np.nan + npt.assert_equal(vals, expected) + + +def cases_test_discrete_integer_shapes(): + # distributions parameters that are only allowed to be integral when + # fitting, but are allowed to be real as input to PDF, etc. + integrality_exceptions = {'nbinom': {'n'}, 'betanbinom': {'n'}} + + seen = set() + for distname, shapes in distdiscrete: + if distname in seen: + continue + seen.add(distname) + + try: + dist = getattr(stats, distname) + except TypeError: + continue + + shape_info = dist._shape_info() + + for i, shape in enumerate(shape_info): + if (shape.name in integrality_exceptions.get(distname, set()) or + not shape.integrality): + continue + + yield distname, shape.name, shapes + + +@pytest.mark.parametrize('distname, shapename, shapes', + list(cases_test_discrete_integer_shapes())) +def test_integer_shapes(distname, shapename, shapes): + dist = getattr(stats, distname) + shape_info = dist._shape_info() + shape_names = [shape.name for shape in shape_info] + i = shape_names.index(shapename) # this element of params must be integral + + shapes_copy = list(shapes) + + valid_shape = shapes[i] + invalid_shape = valid_shape - 0.5 # arbitrary non-integral value + new_valid_shape = valid_shape - 1 + shapes_copy[i] = [[valid_shape], [invalid_shape], [new_valid_shape]] + + a, b = dist.support(*shapes) + x = np.round(np.linspace(a, b, 5)) + + pmf = dist.pmf(x, *shapes_copy) + assert not np.any(np.isnan(pmf[0, :])) + assert np.all(np.isnan(pmf[1, :])) + assert not np.any(np.isnan(pmf[2, :])) + + +def test_frozen_attributes(monkeypatch): + # gh-14827 reported that all frozen distributions had both pmf and pdf + # attributes; continuous should have pdf and discrete should have pmf. + message = "'rv_discrete_frozen' object has no attribute" + with pytest.raises(AttributeError, match=message): + stats.binom(10, 0.5).pdf + with pytest.raises(AttributeError, match=message): + stats.binom(10, 0.5).logpdf + monkeypatch.setattr(stats.binom, "pdf", "herring", raising=False) + frozen_binom = stats.binom(10, 0.5) + assert isinstance(frozen_binom, rv_discrete_frozen) + assert not hasattr(frozen_binom, "pdf") + + +@pytest.mark.parametrize('distname, shapes', distdiscrete) +def test_interval(distname, shapes): + # gh-11026 reported that `interval` returns incorrect values when + # `confidence=1`. The values were not incorrect, but it was not intuitive + # that the left end of the interval should extend beyond the support of the + # distribution. Confirm that this is the behavior for all distributions. + if isinstance(distname, str): + dist = getattr(stats, distname) + else: + dist = distname + a, b = dist.support(*shapes) + npt.assert_equal(dist.ppf([0, 1], *shapes), (a-1, b)) + npt.assert_equal(dist.isf([1, 0], *shapes), (a-1, b)) + npt.assert_equal(dist.interval(1, *shapes), (a-1, b)) + + +@pytest.mark.xfail_on_32bit("Sensible to machine precision") +def test_rv_sample(): + # Thoroughly test rv_sample and check that gh-3758 is resolved + + # Generate a random discrete distribution + rng = np.random.default_rng(98430143469) + xk = np.sort(rng.random(10) * 10) + pk = rng.random(10) + pk /= np.sum(pk) + dist = stats.rv_discrete(values=(xk, pk)) + + # Generate points to the left and right of xk + xk_left = (np.array([0] + xk[:-1].tolist()) + xk)/2 + xk_right = (np.array(xk[1:].tolist() + [xk[-1]+1]) + xk)/2 + + # Generate points to the left and right of cdf + cdf2 = np.cumsum(pk) + cdf2_left = (np.array([0] + cdf2[:-1].tolist()) + cdf2)/2 + cdf2_right = (np.array(cdf2[1:].tolist() + [1]) + cdf2)/2 + + # support - leftmost and rightmost xk + a, b = dist.support() + assert_allclose(a, xk[0]) + assert_allclose(b, xk[-1]) + + # pmf - supported only on the xk + assert_allclose(dist.pmf(xk), pk) + assert_allclose(dist.pmf(xk_right), 0) + assert_allclose(dist.pmf(xk_left), 0) + + # logpmf is log of the pmf; log(0) = -np.inf + with np.errstate(divide='ignore'): + assert_allclose(dist.logpmf(xk), np.log(pk)) + assert_allclose(dist.logpmf(xk_right), -np.inf) + assert_allclose(dist.logpmf(xk_left), -np.inf) + + # cdf - the cumulative sum of the pmf + assert_allclose(dist.cdf(xk), cdf2) + assert_allclose(dist.cdf(xk_right), cdf2) + assert_allclose(dist.cdf(xk_left), [0]+cdf2[:-1].tolist()) + + with np.errstate(divide='ignore'): + assert_allclose(dist.logcdf(xk), np.log(dist.cdf(xk)), + atol=1e-15) + assert_allclose(dist.logcdf(xk_right), np.log(dist.cdf(xk_right)), + atol=1e-15) + assert_allclose(dist.logcdf(xk_left), np.log(dist.cdf(xk_left)), + atol=1e-15) + + # sf is 1-cdf + assert_allclose(dist.sf(xk), 1-dist.cdf(xk)) + assert_allclose(dist.sf(xk_right), 1-dist.cdf(xk_right)) + assert_allclose(dist.sf(xk_left), 1-dist.cdf(xk_left)) + + with np.errstate(divide='ignore'): + assert_allclose(dist.logsf(xk), np.log(dist.sf(xk)), + atol=1e-15) + assert_allclose(dist.logsf(xk_right), np.log(dist.sf(xk_right)), + atol=1e-15) + assert_allclose(dist.logsf(xk_left), np.log(dist.sf(xk_left)), + atol=1e-15) + + # ppf + assert_allclose(dist.ppf(cdf2), xk) + assert_allclose(dist.ppf(cdf2_left), xk) + assert_allclose(dist.ppf(cdf2_right)[:-1], xk[1:]) + assert_allclose(dist.ppf(0), a - 1) + assert_allclose(dist.ppf(1), b) + + # isf + sf2 = dist.sf(xk) + assert_allclose(dist.isf(sf2), xk) + assert_allclose(dist.isf(1-cdf2_left), dist.ppf(cdf2_left)) + assert_allclose(dist.isf(1-cdf2_right), dist.ppf(cdf2_right)) + assert_allclose(dist.isf(0), b) + assert_allclose(dist.isf(1), a - 1) + + # interval is (ppf(alpha/2), isf(alpha/2)) + ps = np.linspace(0.01, 0.99, 10) + int2 = dist.ppf(ps/2), dist.isf(ps/2) + assert_allclose(dist.interval(1-ps), int2) + assert_allclose(dist.interval(0), dist.median()) + assert_allclose(dist.interval(1), (a-1, b)) + + # median is simply ppf(0.5) + med2 = dist.ppf(0.5) + assert_allclose(dist.median(), med2) + + # all four stats (mean, var, skew, and kurtosis) from the definitions + mean2 = np.sum(xk*pk) + var2 = np.sum((xk - mean2)**2 * pk) + skew2 = np.sum((xk - mean2)**3 * pk) / var2**(3/2) + kurt2 = np.sum((xk - mean2)**4 * pk) / var2**2 - 3 + assert_allclose(dist.mean(), mean2) + assert_allclose(dist.std(), np.sqrt(var2)) + assert_allclose(dist.var(), var2) + assert_allclose(dist.stats(moments='mvsk'), (mean2, var2, skew2, kurt2)) + + # noncentral moment against definition + mom3 = np.sum((xk**3) * pk) + assert_allclose(dist.moment(3), mom3) + + # expect - check against moments + assert_allclose(dist.expect(lambda x: 1), 1) + assert_allclose(dist.expect(), mean2) + assert_allclose(dist.expect(lambda x: x**3), mom3) + + # entropy is the negative of the expected value of log(p) + with np.errstate(divide='ignore'): + assert_allclose(-dist.expect(lambda x: dist.logpmf(x)), dist.entropy()) + + # RVS is just ppf of uniform random variates + rng = np.random.default_rng(98430143469) + rvs = dist.rvs(size=100, random_state=rng) + rng = np.random.default_rng(98430143469) + rvs0 = dist.ppf(rng.random(size=100)) + assert_allclose(rvs, rvs0) + +def test__pmf_float_input(): + # gh-21272 + # test that `rvs()` can be computed when `_pmf` requires float input + + class rv_exponential(stats.rv_discrete): + def _pmf(self, i): + return (2/3)*3**(1 - i) + + rv = rv_exponential(a=0.0, b=float('inf')) + rvs = rv.rvs(random_state=42) # should not crash due to integer input to `_pmf` + assert_allclose(rvs, 0) + + +def test_gh18919_ppf_array_args(): + # gh-18919 reported incorrect results for ppf and isf of discrete distributions when + # arguments were arrays and first argument (`q`) had elements at the boundaries of + # the support. + q = [[0.5, 1.0, 0.5], + [1.0, 0.5, 1.0], + [0.5, 1.0, 0.5]] + + n = [[45, 46, 47], + [48, 49, 50], + [51, 52, 53]] + + p = 0.5 + + ref = _ufuncs._binom_ppf(q, n, p) + res = stats.binom.ppf(q, n, p) + np.testing.assert_allclose(res, ref) + +@pytest.mark.parametrize("dist", [stats.binom, stats.boltzmann]) +def test_gh18919_ppf_isf_array_args2(dist): + # a more general version of the test above. Requires that arguments are broadcasted + # by the infrastructure. + rng = np.random.default_rng(34873457824358729823) + q = rng.random(size=(30, 1, 1, 1)) + n = rng.integers(10, 30, size=(10, 1, 1)) + p = rng.random(size=(4, 1)) + loc = rng.integers(5, size=(3,)) + + q[rng.random(size=30) > 0.7] = 0 + q[rng.random(size=30) > 0.7] = 1 + + args = (q, n, p) if dist == stats.binom else (q, p, n) + + res = dist.ppf(*args, loc=loc) + ref = np.vectorize(dist.ppf)(*args) + loc + np.testing.assert_allclose(res, ref) + + res = dist.isf(*args, loc=loc) + ref = np.vectorize(dist.isf)(*args) + loc + np.testing.assert_allclose(res, ref) diff --git a/v12x/pylib/scipy/stats/tests/test_discrete_distns.py b/v12x/pylib/scipy/stats/tests/test_discrete_distns.py new file mode 100644 index 0000000000000000000000000000000000000000..8ba8da84e2dca928426a3b5283d85f5cf29be07d --- /dev/null +++ b/v12x/pylib/scipy/stats/tests/test_discrete_distns.py @@ -0,0 +1,747 @@ +import pytest +import itertools +import warnings + +from scipy import stats +from scipy.stats import (betabinom, betanbinom, hypergeom, nhypergeom, + bernoulli, boltzmann, skellam, zipf, zipfian, binom, + nbinom, nchypergeom_fisher, nchypergeom_wallenius, + randint, poisson_binom) + +import numpy as np +from numpy.testing import ( + assert_almost_equal, assert_equal, assert_allclose +) +from scipy.special import binom as special_binom +from scipy.optimize import root_scalar +from scipy.integrate import quad + + +# The expected values were computed with Wolfram Alpha, using +# the expression CDF[HypergeometricDistribution[N, n, M], k]. +@pytest.mark.parametrize('k, M, n, N, expected, rtol', + [(3, 10, 4, 5, + 0.9761904761904762, 1e-15), + (107, 10000, 3000, 215, + 0.9999999997226765, 1e-15), + (10, 10000, 3000, 215, + 2.681682217692179e-21, 5e-11)]) +def test_hypergeom_cdf(k, M, n, N, expected, rtol): + p = hypergeom.cdf(k, M, n, N) + assert_allclose(p, expected, rtol=rtol) + + +# The expected values were computed with Wolfram Alpha, using +# the expression SurvivalFunction[HypergeometricDistribution[N, n, M], k]. +@pytest.mark.parametrize('k, M, n, N, expected, rtol', + [(25, 10000, 3000, 215, + 0.9999999999052958, 1e-15), + (125, 10000, 3000, 215, + 1.4416781705752128e-18, 5e-11)]) +def test_hypergeom_sf(k, M, n, N, expected, rtol): + p = hypergeom.sf(k, M, n, N) + assert_allclose(p, expected, rtol=rtol) + + +def test_hypergeom_logpmf(): + # symmetries test + # f(k,N,K,n) = f(n-k,N,N-K,n) = f(K-k,N,K,N-n) = f(k,N,n,K) + k = 5 + N = 50 + K = 10 + n = 5 + logpmf1 = hypergeom.logpmf(k, N, K, n) + logpmf2 = hypergeom.logpmf(n - k, N, N - K, n) + logpmf3 = hypergeom.logpmf(K - k, N, K, N - n) + logpmf4 = hypergeom.logpmf(k, N, n, K) + assert_almost_equal(logpmf1, logpmf2, decimal=12) + assert_almost_equal(logpmf1, logpmf3, decimal=12) + assert_almost_equal(logpmf1, logpmf4, decimal=12) + + # test related distribution + # Bernoulli distribution if n = 1 + k = 1 + N = 10 + K = 7 + n = 1 + hypergeom_logpmf = hypergeom.logpmf(k, N, K, n) + bernoulli_logpmf = bernoulli.logpmf(k, K/N) + assert_almost_equal(hypergeom_logpmf, bernoulli_logpmf, decimal=12) + + +def test_nhypergeom_pmf(): + # test with hypergeom + M, n, r = 45, 13, 8 + k = 6 + NHG = nhypergeom.pmf(k, M, n, r) + HG = hypergeom.pmf(k, M, n, k+r-1) * (M - n - (r-1)) / (M - (k+r-1)) + assert_allclose(HG, NHG, rtol=1e-10) + + +def test_nhypergeom_pmfcdf(): + # test pmf and cdf with arbitrary values. + M = 8 + n = 3 + r = 4 + support = np.arange(n+1) + pmf = nhypergeom.pmf(support, M, n, r) + cdf = nhypergeom.cdf(support, M, n, r) + assert_allclose(pmf, [1/14, 3/14, 5/14, 5/14], rtol=1e-13) + assert_allclose(cdf, [1/14, 4/14, 9/14, 1.0], rtol=1e-13) + + +def test_nhypergeom_r0(): + # test with `r = 0`. + M = 10 + n = 3 + r = 0 + pmf = nhypergeom.pmf([[0, 1, 2, 0], [1, 2, 0, 3]], M, n, r) + assert_allclose(pmf, [[1, 0, 0, 1], [0, 0, 1, 0]], rtol=1e-13) + + +def test_nhypergeom_rvs_shape(): + # Check that when given a size with more dimensions than the + # dimensions of the broadcast parameters, rvs returns an array + # with the correct shape. + x = nhypergeom.rvs(22, [7, 8, 9], [[12], [13]], size=(5, 1, 2, 3)) + assert x.shape == (5, 1, 2, 3) + + +def test_nhypergeom_accuracy(): + # Check that nhypergeom.rvs post-gh-13431 gives the same values as + # inverse transform sampling + rng = np.random.RandomState(0) + x = nhypergeom.rvs(22, 7, 11, size=100, random_state=rng) + rng = np.random.RandomState(0) + p = rng.uniform(size=100) + y = nhypergeom.ppf(p, 22, 7, 11) + assert_equal(x, y) + + +def test_boltzmann_upper_bound(): + k = np.arange(-3, 5) + + N = 1 + p = boltzmann.pmf(k, 0.123, N) + expected = k == 0 + assert_equal(p, expected) + + lam = np.log(2) + N = 3 + p = boltzmann.pmf(k, lam, N) + expected = [0, 0, 0, 4/7, 2/7, 1/7, 0, 0] + assert_allclose(p, expected, rtol=1e-13) + + c = boltzmann.cdf(k, lam, N) + expected = [0, 0, 0, 4/7, 6/7, 1, 1, 1] + assert_allclose(c, expected, rtol=1e-13) + + +def test_betabinom_a_and_b_unity(): + # test limiting case that betabinom(n, 1, 1) is a discrete uniform + # distribution from 0 to n + n = 20 + k = np.arange(n + 1) + p = betabinom(n, 1, 1).pmf(k) + expected = np.repeat(1 / (n + 1), n + 1) + assert_almost_equal(p, expected) + + +@pytest.mark.parametrize('dtypes', list(itertools.product(*[(int, float)]*3))) +def test_betabinom_stats_a_and_b_integers_gh18026(dtypes): + # gh-18026 reported that `betabinom` kurtosis calculation fails when some + # parameters are integers. Check that this is resolved. + n_type, a_type, b_type = dtypes + n, a, b = n_type(10), a_type(2), b_type(3) + assert_allclose(betabinom.stats(n, a, b, moments='k'), -0.6904761904761907) + + +def test_betabinom_bernoulli(): + # test limiting case that betabinom(1, a, b) = bernoulli(a / (a + b)) + a = 2.3 + b = 0.63 + k = np.arange(2) + p = betabinom(1, a, b).pmf(k) + expected = bernoulli(a / (a + b)).pmf(k) + assert_almost_equal(p, expected) + + +def test_issue_10317(): + alpha, n, p = 0.9, 10, 1 + assert_equal(nbinom.interval(confidence=alpha, n=n, p=p), (0, 0)) + + +def test_issue_11134(): + alpha, n, p = 0.95, 10, 0 + assert_equal(binom.interval(confidence=alpha, n=n, p=p), (0, 0)) + + +def test_issue_7406(): + rng = np.random.default_rng(4763112764) + assert_equal(binom.ppf(rng.random(10), 0, 0.5), 0) + + # Also check that endpoints (q=0, q=1) are correct + assert_equal(binom.ppf(0, 0, 0.5), -1) + assert_equal(binom.ppf(1, 0, 0.5), 0) + + +def test_issue_5122(): + rng = np.random.default_rng(8312492117) + p = 0 + n = rng.integers(100, size=10) + + x = 0 + ppf = binom.ppf(x, n, p) + assert_equal(ppf, -1) + + x = np.linspace(0.01, 0.99, 10) + ppf = binom.ppf(x, n, p) + assert_equal(ppf, 0) + + x = 1 + ppf = binom.ppf(x, n, p) + assert_equal(ppf, n) + + +def test_issue_1603(): + assert_equal(binom(1000, np.logspace(-3, -100)).ppf(0.01), 0) + + +def test_issue_5503(): + p = 0.5 + x = np.logspace(3, 14, 12) + assert_allclose(binom.cdf(x, 2*x, p), 0.5, atol=1e-2) + + +@pytest.mark.parametrize('x, n, p, cdf_desired', [ + (300, 1000, 3/10, 0.51559351981411995636), + (3000, 10000, 3/10, 0.50493298381929698016), + (30000, 100000, 3/10, 0.50156000591726422864), + (300000, 1000000, 3/10, 0.50049331906666960038), + (3000000, 10000000, 3/10, 0.50015600124585261196), + (30000000, 100000000, 3/10, 0.50004933192735230102), + (30010000, 100000000, 3/10, 0.98545384016570790717), + (29990000, 100000000, 3/10, 0.01455017177985268670), + (29950000, 100000000, 3/10, 5.02250963487432024943e-28), + (300_000_000, 1_000_000_000, 3/10, 0.50001560012523938), + (3_000_000_000, 10_000_000_000, 3/10, 0.50000493319275) +]) +def test_issue_5503pt2(x, n, p, cdf_desired): + assert_allclose(binom.cdf(x, n, p), cdf_desired) + + +def test_issue_5503pt3(): + # From Wolfram Alpha: CDF[BinomialDistribution[1e12, 1e-12], 2] + assert_allclose(binom.cdf(2, 10**12, 10**-12), 0.91969860292869777384) + + +def test_issue_6682(): + # Reference value from R: + # options(digits=16) + # print(pnbinom(250, 50, 32/63, lower.tail=FALSE)) + assert_allclose(nbinom.sf(250, 50, 32./63.), 1.460458510976452e-35) + + +def test_issue_19747(): + # test that negative k does not raise an error in nbinom.logcdf + result = nbinom.logcdf([5, -1, 1], 5, 0.5) + reference = [-0.47313352, -np.inf, -2.21297293] + assert_allclose(result, reference) + + +def test_boost_divide_by_zero_issue_15101(): + n = 1000 + p = 0.01 + k = 996 + assert_allclose(binom.pmf(k, n, p), 0.0) + + +def test_skellam_gh11474(): + # test issue reported in gh-11474 caused by `cdfchn` + mu = [1, 10, 100, 1000, 5000, 5050, 5100, 5250, 6000] + cdf = skellam.cdf(0, mu, mu) + # generated in R + # library(skellam) + # options(digits = 16) + # mu = c(1, 10, 100, 1000, 5000, 5050, 5100, 5250, 6000) + # pskellam(0, mu, mu, TRUE) + cdf_expected = [0.6542541612768356, 0.5448901559424127, 0.5141135799745580, + 0.5044605891382528, 0.5019947363350450, 0.5019848365953181, + 0.5019750827993392, 0.5019466621805060, 0.5018209330219539] + assert_allclose(cdf, cdf_expected) + + +class TestZipfian: + def test_zipfian_asymptotic(self): + # test limiting case that zipfian(a, n) -> zipf(a) as n-> oo + a = 6.5 + N = 10000000 + k = np.arange(1, 21) + assert_allclose(zipfian.pmf(k, a, N), zipf.pmf(k, a)) + assert_allclose(zipfian.cdf(k, a, N), zipf.cdf(k, a)) + assert_allclose(zipfian.sf(k, a, N), zipf.sf(k, a)) + assert_allclose(zipfian.stats(a, N, moments='msvk'), + zipf.stats(a, moments='msvk')) + + def test_zipfian_continuity(self): + # test that zipfian(0.999999, n) ~ zipfian(1.000001, n) + # (a = 1 switches between methods of calculating harmonic sum) + alt1, agt1 = 0.99999999, 1.00000001 + N = 30 + k = np.arange(1, N + 1) + assert_allclose(zipfian.pmf(k, alt1, N), zipfian.pmf(k, agt1, N), + rtol=5e-7) + assert_allclose(zipfian.cdf(k, alt1, N), zipfian.cdf(k, agt1, N), + rtol=5e-7) + assert_allclose(zipfian.sf(k, alt1, N), zipfian.sf(k, agt1, N), + rtol=5e-7) + assert_allclose(zipfian.stats(alt1, N, moments='msvk'), + zipfian.stats(agt1, N, moments='msvk'), rtol=5e-7) + + def test_zipfian_R(self): + # test against R VGAM package + # library(VGAM) + # k <- c(13, 16, 1, 4, 4, 8, 10, 19, 5, 7) + # a <- c(1.56712977, 3.72656295, 5.77665117, 9.12168729, 5.79977172, + # 4.92784796, 9.36078764, 4.3739616 , 7.48171872, 4.6824154) + # n <- c(70, 80, 48, 65, 83, 89, 50, 30, 20, 20) + # pmf <- dzipf(k, N = n, shape = a) + # cdf <- pzipf(k, N = n, shape = a) + # print(pmf) + # print(cdf) + rng = np.random.RandomState(0) + k = rng.randint(1, 20, size=10) + a = rng.rand(10)*10 + 1 + n = rng.randint(1, 100, size=10) + pmf = [8.076972e-03, 2.950214e-05, 9.799333e-01, 3.216601e-06, + 3.158895e-04, 3.412497e-05, 4.350472e-10, 2.405773e-06, + 5.860662e-06, 1.053948e-04] + cdf = [0.8964133, 0.9998666, 0.9799333, 0.9999995, 0.9998584, + 0.9999458, 1.0000000, 0.9999920, 0.9999977, 0.9998498] + # skip the first point; zipUC is not accurate for low a, n + assert_allclose(zipfian.pmf(k, a, n)[1:], pmf[1:], rtol=1e-6) + assert_allclose(zipfian.cdf(k, a, n)[1:], cdf[1:], rtol=5e-5) + + rng = np.random.RandomState(0) + naive_tests = np.vstack((np.logspace(-2, 1, 10), + rng.randint(2, 40, 10))).T + + @pytest.mark.parametrize("a, n", list(naive_tests)) + def test_zipfian_naive(self, a, n): + # test against bare-bones implementation + + @np.vectorize + def Hns(n, s): + """Naive implementation of harmonic sum""" + return (1/np.arange(1, n+1)**s).sum() + + @np.vectorize + def pzip(k, a, n): + """Naive implementation of zipfian pmf""" + if k < 1 or k > n: + return 0. + else: + return 1 / k**a / Hns(n, a) + + k = np.arange(n+1) + pmf = pzip(k, a, n) + cdf = np.cumsum(pmf) + mean = np.average(k, weights=pmf) + var = np.average((k - mean)**2, weights=pmf) + std = var**0.5 + skew = np.average(((k-mean)/std)**3, weights=pmf) + kurtosis = np.average(((k-mean)/std)**4, weights=pmf) - 3 + assert_allclose(zipfian.pmf(k, a, n), pmf) + assert_allclose(zipfian.cdf(k, a, n), cdf) + assert_allclose(zipfian.stats(a, n, moments="mvsk"), + [mean, var, skew, kurtosis]) + + def test_pmf_integer_k(self): + k = np.arange(0, 1000) + k_int32 = k.astype(np.int32) + dist = zipfian(111, 22) + pmf = dist.pmf(k) + pmf_k_int32 = dist.pmf(k_int32) + assert_equal(pmf, pmf_k_int32) + + # Reference values were computed with mpmath. + @pytest.mark.parametrize( + 'k, a, n, ref', + [(3, 1 + 1e-12, 10, 0.11380571738244807), + (995, 1 + 1e-9, 1000, 0.0001342634472310051)] + ) + def test_pmf_against_mpmath(self, k, a, n, ref): + p = zipfian.pmf(k, a, n) + assert_allclose(p, ref, 5e-16) + + # Reference values were computed with mpmath. + @pytest.mark.parametrize( + 'k, a, n, ref', + [(4990, 1.25, 5000, 5.780138225335147e-05), + (9998, 3.5, 10000, 1.775352757966365e-14), + (50000, 3.5, 100000, 5.227803621367486e-13), + (10, 6.0, 100, 1.523108153557902e-06), + (95, 6.0, 100, 5.572438078308601e-12)] + ) + def test_sf_against_mpmath(self, k, a, n, ref): + sf = zipfian.sf(k, a, n) + assert_allclose(sf, ref, rtol=8e-15) + + # Reference values were computed with mpmath. + @pytest.mark.parametrize( + 'a, n, ref', + [(1 + 1e-9, 100, 19.27756356649707), + (1 + 1e-7, 100000, 8271.194454731552), + (1.001, 3, 1.6360225521183804)] + ) + def test_mean_against_mpmath(self, a, n, ref): + m = zipfian.mean(a, n) + assert_allclose(m, ref, rtol=8e-15) + + def test_ridiculously_large_n(self): + # This should return nan with no errors or warnings. + a = 2.5 + n = 1e100 + p = zipfian.pmf(10, a, n) + assert_equal(p, np.nan) + + +class TestNCH: + def setup_method(self): + rng = np.random.default_rng(7162434334) + shape = (2, 4, 3) + max_m = 100 + m1 = rng.integers(1, max_m, size=shape) # red balls + m2 = rng.integers(1, max_m, size=shape) # white balls + N = m1 + m2 # total balls + n = randint.rvs(0, N, size=N.shape, random_state=rng) # number of draws + xl = np.maximum(0, n-m2) # lower bound of support + xu = np.minimum(n, m1) # upper bound of support + x = randint.rvs(xl, xu, size=xl.shape, random_state=rng) + odds = rng.random(x.shape)*2 + self.x, self.N, self.m1, self.n, self.odds = x, N, m1, n, odds + + # test output is more readable when function names (strings) are passed + @pytest.mark.parametrize('dist_name', + ['nchypergeom_fisher', 'nchypergeom_wallenius']) + def test_nch_hypergeom(self, dist_name): + # Both noncentral hypergeometric distributions reduce to the + # hypergeometric distribution when odds = 1 + dists = {'nchypergeom_fisher': nchypergeom_fisher, + 'nchypergeom_wallenius': nchypergeom_wallenius} + dist = dists[dist_name] + x, N, m1, n = self.x, self.N, self.m1, self.n + assert_allclose(dist.pmf(x, N, m1, n, odds=1), + hypergeom.pmf(x, N, m1, n)) + + def test_nchypergeom_fisher_naive(self): + # test against a very simple implementation + x, N, m1, n, odds = self.x, self.N, self.m1, self.n, self.odds + + @np.vectorize + def pmf_mean_var(x, N, m1, n, w): + # simple implementation of nchypergeom_fisher pmf + m2 = N - m1 + xl = np.maximum(0, n-m2) + xu = np.minimum(n, m1) + + def f(x): + t1 = special_binom(m1, x) + t2 = special_binom(m2, n - x) + return t1 * t2 * w**x + + def P(k): + return sum(f(y)*y**k for y in range(xl, xu + 1)) + + P0 = P(0) + P1 = P(1) + P2 = P(2) + pmf = f(x) / P0 + mean = P1 / P0 + var = P2 / P0 - (P1 / P0)**2 + return pmf, mean, var + + pmf, mean, var = pmf_mean_var(x, N, m1, n, odds) + assert_allclose(nchypergeom_fisher.pmf(x, N, m1, n, odds), pmf) + assert_allclose(nchypergeom_fisher.stats(N, m1, n, odds, moments='m'), + mean) + assert_allclose(nchypergeom_fisher.stats(N, m1, n, odds, moments='v'), + var) + + def test_nchypergeom_wallenius_naive(self): + # test against a very simple implementation + + rng = np.random.RandomState(2) + shape = (2, 4, 3) + max_m = 100 + m1 = rng.randint(1, max_m, size=shape) + m2 = rng.randint(1, max_m, size=shape) + N = m1 + m2 + n = randint.rvs(0, N, size=N.shape, random_state=rng) + xl = np.maximum(0, n-m2) + xu = np.minimum(n, m1) + x = randint.rvs(xl, xu, size=xl.shape, random_state=rng) + w = rng.rand(*x.shape)*2 + + def support(N, m1, n, w): + m2 = N - m1 + xl = np.maximum(0, n-m2) + xu = np.minimum(n, m1) + return xl, xu + + @np.vectorize + def mean(N, m1, n, w): + m2 = N - m1 + xl, xu = support(N, m1, n, w) + + def fun(u): + return u/m1 + (1 - (n-u)/m2)**w - 1 + + return root_scalar(fun, bracket=(xl, xu)).root + + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", category=RuntimeWarning, + message="invalid value encountered in mean") + assert_allclose(nchypergeom_wallenius.mean(N, m1, n, w), + mean(N, m1, n, w), rtol=2e-2) + + @np.vectorize + def variance(N, m1, n, w): + m2 = N - m1 + u = mean(N, m1, n, w) + a = u * (m1 - u) + b = (n-u)*(u + m2 - n) + return N*a*b / ((N-1) * (m1*b + m2*a)) + + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", category=RuntimeWarning, + message="invalid value encountered in mean") + assert_allclose( + nchypergeom_wallenius.stats(N, m1, n, w, moments='v'), + variance(N, m1, n, w), + rtol=5e-2 + ) + + @np.vectorize + def pmf(x, N, m1, n, w): + m2 = N - m1 + xl, xu = support(N, m1, n, w) + + def integrand(t): + D = w*(m1 - x) + (m2 - (n-x)) + res = (1-t**(w/D))**x * (1-t**(1/D))**(n-x) + return res + + def f(x): + t1 = special_binom(m1, x) + t2 = special_binom(m2, n - x) + the_integral = quad(integrand, 0, 1, + epsrel=1e-16, epsabs=1e-16) + return t1 * t2 * the_integral[0] + + return f(x) + + pmf0 = pmf(x, N, m1, n, w) + pmf1 = nchypergeom_wallenius.pmf(x, N, m1, n, w) + + atol, rtol = 1e-6, 1e-6 + i = np.abs(pmf1 - pmf0) < atol + rtol*np.abs(pmf0) + assert i.sum() > np.prod(shape) / 2 # works at least half the time + + # for those that fail, discredit the naive implementation + for N, m1, n, w in zip(N[~i], m1[~i], n[~i], w[~i]): + # get the support + m2 = N - m1 + xl, xu = support(N, m1, n, w) + x = np.arange(xl, xu + 1) + + # calculate sum of pmf over the support + # the naive implementation is very wrong in these cases + assert pmf(x, N, m1, n, w).sum() < .5 + assert_allclose(nchypergeom_wallenius.pmf(x, N, m1, n, w).sum(), 1) + + def test_wallenius_against_mpmath(self): + # precompute data with mpmath since naive implementation above + # is not reliable. See source code in gh-13330. + M = 50 + n = 30 + N = 20 + odds = 2.25 + # Expected results, computed with mpmath. + sup = np.arange(21) + pmf = np.array([3.699003068656875e-20, + 5.89398584245431e-17, + 2.1594437742911123e-14, + 3.221458044649955e-12, + 2.4658279241205077e-10, + 1.0965862603981212e-08, + 3.057890479665704e-07, + 5.622818831643761e-06, + 7.056482841531681e-05, + 0.000618899425358671, + 0.003854172932571669, + 0.01720592676256026, + 0.05528844897093792, + 0.12772363313574242, + 0.21065898367825722, + 0.24465958845359234, + 0.1955114898110033, + 0.10355390084949237, + 0.03414490375225675, + 0.006231989845775931, + 0.0004715577304677075]) + mean = 14.808018384813426 + var = 2.6085975877923717 + + # nchypergeom_wallenius.pmf returns 0 for pmf(0) and pmf(1), and pmf(2) + # has only three digits of accuracy (~ 2.1511e-14). + assert_allclose(nchypergeom_wallenius.pmf(sup, M, n, N, odds), pmf, + rtol=1e-13, atol=1e-13) + assert_allclose(nchypergeom_wallenius.mean(M, n, N, odds), + mean, rtol=1e-13) + assert_allclose(nchypergeom_wallenius.var(M, n, N, odds), + var, rtol=1e-11) + + @pytest.mark.parametrize('dist_name', + ['nchypergeom_fisher', 'nchypergeom_wallenius']) + def test_rvs_shape(self, dist_name): + # Check that when given a size with more dimensions than the + # dimensions of the broadcast parameters, rvs returns an array + # with the correct shape. + dists = {'nchypergeom_fisher': nchypergeom_fisher, + 'nchypergeom_wallenius': nchypergeom_wallenius} + dist = dists[dist_name] + x = dist.rvs(50, 30, [[10], [20]], [0.5, 1.0, 2.0], size=(5, 1, 2, 3)) + assert x.shape == (5, 1, 2, 3) + + +@pytest.mark.parametrize("mu, q, expected", + [[10, 120, -1.240089881791596e-38], + [1500, 0, -86.61466680572661]]) +def test_nbinom_11465(mu, q, expected): + # test nbinom.logcdf at extreme tails + size = 20 + n, p = size, size/(size+mu) + # In R: + # options(digits=16) + # pnbinom(mu=10, size=20, q=120, log.p=TRUE) + assert_allclose(nbinom.logcdf(q, n, p), expected) + + +def test_gh_17146(): + # Check that discrete distributions return PMF of zero at non-integral x. + # See gh-17146. + x = np.linspace(0, 1, 11) + p = 0.8 + pmf = bernoulli(p).pmf(x) + i = (x % 1 == 0) + assert_allclose(pmf[-1], p) + assert_allclose(pmf[0], 1-p) + assert_equal(pmf[~i], 0) + + +class TestBetaNBinom: + @pytest.mark.parametrize('x, n, a, b, ref', + [[5, 5e6, 5, 20, 1.1520944824139114e-107], + [100, 50, 5, 20, 0.002855762954310226], + [10000, 1000, 5, 20, 1.9648515726019154e-05]]) + def test_betanbinom_pmf(self, x, n, a, b, ref): + # test that PMF stays accurate in the distribution tails + # reference values computed with mpmath + # from mpmath import mp + # mp.dps = 500 + # def betanbinom_pmf(k, n, a, b): + # k = mp.mpf(k) + # a = mp.mpf(a) + # b = mp.mpf(b) + # n = mp.mpf(n) + # return float(mp.binomial(n + k - mp.one, k) + # * mp.beta(a + n, b + k) / mp.beta(a, b)) + assert_allclose(betanbinom.pmf(x, n, a, b), ref, rtol=1e-10) + + + @pytest.mark.parametrize('n, a, b, ref', + [[10000, 5000, 50, 0.12841520515722202], + [10, 9, 9, 7.9224400871459695], + [100, 1000, 10, 1.5849602176622748]]) + def test_betanbinom_kurtosis(self, n, a, b, ref): + # reference values were computed via mpmath + # from mpmath import mp + # def kurtosis_betanegbinom(n, a, b): + # n = mp.mpf(n) + # a = mp.mpf(a) + # b = mp.mpf(b) + # four = mp.mpf(4.) + # mean = n * b / (a - mp.one) + # var = (n * b * (n + a - 1.) * (a + b - 1.) + # / ((a - 2.) * (a - 1.)**2.)) + # def f(k): + # return (mp.binomial(n + k - mp.one, k) + # * mp.beta(a + n, b + k) / mp.beta(a, b) + # * (k - mean)**four) + # fourth_moment = mp.nsum(f, [0, mp.inf]) + # return float(fourth_moment/var**2 - 3.) + assert_allclose(betanbinom.stats(n, a, b, moments="k"), + ref, rtol=3e-15) + + +class TestZipf: + def test_gh20692(self): + # test that int32 data for k generates same output as double + k = np.arange(0, 1000) + k_int32 = k.astype(np.int32) + dist = zipf(9) + pmf = dist.pmf(k) + pmf_k_int32 = dist.pmf(k_int32) + assert_equal(pmf, pmf_k_int32) + + +def test_gh20048(): + # gh-20048 reported an infinite loop in _drv2_ppfsingle + # check that the one identified is resolved + class test_dist_gen(stats.rv_discrete): + def _cdf(self, k): + return min(k / 100, 0.99) + + test_dist = test_dist_gen(b=np.inf) + + message = "Arguments that bracket..." + with pytest.raises(RuntimeError, match=message): + test_dist.ppf(0.999) + + +class TestPoissonBinomial: + def test_pmf(self): + # Test pmf against R `poisbinom` to confirm that this is indeed the Poisson + # binomial distribution. Consistency of other methods and all other behavior + # should be covered by generic tests. (If not, please add a generic test.) + # Like many other distributions, no special attempt is made to be more + # accurate than the usual formulas provide, so we use default tolerances. + # + # library(poisbinom) + # options(digits=16) + # k = c(0, 1, 2, 3, 4) + # p = c(0.9480654803913988, 0.052428488100509374, + # 0.25863527358887417, 0.057764076043633206) + # dpoisbinom(k, p) + rng = np.random.default_rng(259823598254) + n = rng.integers(10) # 4 + k = np.arange(n + 1) + p = rng.random(n) # [0.9480654803913988, 0.052428488100509374, + # 0.25863527358887417, 0.057764076043633206] + res = poisson_binom.pmf(k, p) + ref = [0.0343763443678060318, 0.6435428452689714307, 0.2936345519235536994, + 0.0277036647503902354, 0.0007425936892786034] + assert_allclose(res, ref) + + +class TestRandInt: + def test_gh19759(self): + # test zero PMF values within the support reported by gh-19759 + a = -354 + max_range = abs(a) + all_b_1 = [a + 2 ** 31 + i for i in range(max_range)] + res = randint.pmf(325, a, all_b_1) + assert (res > 0).all() + ref = 1 / (np.asarray(all_b_1, dtype=np.float64) - a) + assert_allclose(res, ref) diff --git a/v12x/pylib/scipy/stats/tests/test_distributions.py b/v12x/pylib/scipy/stats/tests/test_distributions.py new file mode 100644 index 0000000000000000000000000000000000000000..a355a702807438358bd88c00de026008055dbc51 --- /dev/null +++ b/v12x/pylib/scipy/stats/tests/test_distributions.py @@ -0,0 +1,10606 @@ +""" +Test functions for stats module +""" +import warnings +import re +import sys +import pickle +import threading +from pathlib import Path +import os +import json +import platform + +from numpy.testing import (assert_equal, assert_array_equal, + assert_almost_equal, assert_array_almost_equal, + assert_allclose, assert_, + assert_array_less, assert_array_max_ulp) +import pytest +from pytest import raises as assert_raises + +import numpy as np +from numpy import typecodes, array +from numpy.lib.recfunctions import rec_append_fields +from scipy import special +from scipy._lib._util import check_random_state +from scipy.integrate import (IntegrationWarning, quad, trapezoid, + cumulative_trapezoid) +import scipy.stats as stats +from scipy.stats._distn_infrastructure import argsreduce +from scipy.stats._constants import _XMAX +import scipy.stats.distributions + +from scipy.special import xlogy, polygamma, entr +from scipy.stats._distr_params import distcont, invdistcont +from .test_discrete_basic import distdiscrete, invdistdiscrete +from scipy.stats._continuous_distns import FitDataError, _argus_phi +from scipy.optimize import root, fmin, differential_evolution +from itertools import product + +# python -OO strips docstrings +DOCSTRINGS_STRIPPED = sys.flags.optimize > 1 + +# Failing on macOS 11, Intel CPUs. See gh-14901 +MACOS_INTEL = (sys.platform == 'darwin') and (platform.machine() == 'x86_64') + + +# distributions to skip while testing the fix for the support method +# introduced in gh-13294. These distributions are skipped as they +# always return a non-nan support for every parametrization. +skip_test_support_gh13294_regression = ['tukeylambda', 'pearson3'] + + +def _assert_hasattr(a, b, msg=None): + if msg is None: + msg = f'{a} does not have attribute {b}' + assert_(hasattr(a, b), msg=msg) + + +def test_api_regression(): + # https://github.com/scipy/scipy/issues/3802 + _assert_hasattr(scipy.stats.distributions, 'f_gen') + + +def test_distributions_submodule(): + actual = set(scipy.stats.distributions.__all__) + continuous = [dist[0] for dist in distcont] # continuous dist names + discrete = [dist[0] for dist in distdiscrete] # discrete dist names + other = ['rv_discrete', 'rv_continuous', 'rv_histogram', + 'entropy'] + expected = continuous + discrete + other + + # need to remove, e.g., + # + expected = set(filter(lambda s: not str(s).startswith('<'), expected)) + + assert actual == expected + + +class TestVonMises: + def setup_method(self): + self.rng = np.random.default_rng(6320571663) + + @pytest.mark.parametrize('k', [0.1, 1, 101]) + @pytest.mark.parametrize('x', [0, 1, np.pi, 10, 100]) + def test_vonmises_periodic(self, k, x): + def check_vonmises_pdf_periodic(k, L, s, x): + vm = stats.vonmises(k, loc=L, scale=s) + assert_almost_equal(vm.pdf(x), vm.pdf(x % (2 * np.pi * s))) + + def check_vonmises_cdf_periodic(k, L, s, x): + vm = stats.vonmises(k, loc=L, scale=s) + assert_almost_equal(vm.cdf(x) % 1, + vm.cdf(x % (2 * np.pi * s)) % 1) + + check_vonmises_pdf_periodic(k, 0, 1, x) + check_vonmises_pdf_periodic(k, 1, 1, x) + check_vonmises_pdf_periodic(k, 0, 10, x) + + check_vonmises_cdf_periodic(k, 0, 1, x) + check_vonmises_cdf_periodic(k, 1, 1, x) + check_vonmises_cdf_periodic(k, 0, 10, x) + + def test_vonmises_line_support(self): + assert_equal(stats.vonmises_line.a, -np.pi) + assert_equal(stats.vonmises_line.b, np.pi) + + def test_vonmises_numerical(self): + vm = stats.vonmises(800) + assert_almost_equal(vm.cdf(0), 0.5) + + # Expected values of the vonmises PDF were computed using + # mpmath with 50 digits of precision: + # + # def vmpdf_mp(x, kappa): + # x = mpmath.mpf(x) + # kappa = mpmath.mpf(kappa) + # num = mpmath.exp(kappa*mpmath.cos(x)) + # den = 2 * mpmath.pi * mpmath.besseli(0, kappa) + # return num/den + + @pytest.mark.parametrize('x, kappa, expected_pdf', + [(0.1, 0.01, 0.16074242744907072), + (0.1, 25.0, 1.7515464099118245), + (0.1, 800, 0.2073272544458798), + (2.0, 0.01, 0.15849003875385817), + (2.0, 25.0, 8.356882934278192e-16), + (2.0, 800, 0.0)]) + def test_vonmises_pdf(self, x, kappa, expected_pdf): + pdf = stats.vonmises.pdf(x, kappa) + assert_allclose(pdf, expected_pdf, rtol=1e-15) + + # Expected values of the vonmises entropy were computed using + # mpmath with 50 digits of precision: + # + # def vonmises_entropy(kappa): + # kappa = mpmath.mpf(kappa) + # return (-kappa * mpmath.besseli(1, kappa) / + # mpmath.besseli(0, kappa) + mpmath.log(2 * mpmath.pi * + # mpmath.besseli(0, kappa))) + # >>> float(vonmises_entropy(kappa)) + + @pytest.mark.parametrize('kappa, expected_entropy', + [(1, 1.6274014590199897), + (5, 0.6756431570114528), + (100, -0.8811275441649473), + (1000, -2.03468891852547), + (2000, -2.3813876496587847)]) + def test_vonmises_entropy(self, kappa, expected_entropy): + entropy = stats.vonmises.entropy(kappa) + assert_allclose(entropy, expected_entropy, rtol=1e-13) + + def test_vonmises_rvs_gh4598(self): + # check that random variates wrap around as discussed in gh-4598 + seed = 30899520 + rng1 = np.random.default_rng(seed) + rng2 = np.random.default_rng(seed) + rng3 = np.random.default_rng(seed) + rvs1 = stats.vonmises(1, loc=0, scale=1).rvs(random_state=rng1) + rvs2 = stats.vonmises(1, loc=2*np.pi, scale=1).rvs(random_state=rng2) + rvs3 = stats.vonmises(1, loc=0, + scale=(2*np.pi/abs(rvs1)+1)).rvs(random_state=rng3) + assert_allclose(rvs1, rvs2, atol=1e-15) + assert_allclose(rvs1, rvs3, atol=1e-15) + + # Expected values of the vonmises LOGPDF were computed + # using wolfram alpha: + # kappa * cos(x) - log(2*pi*I0(kappa)) + @pytest.mark.parametrize('x, kappa, expected_logpdf', + [(0.1, 0.01, -1.8279520246003170), + (0.1, 25.0, 0.5604990605420549), + (0.1, 800, -1.5734567947337514), + (2.0, 0.01, -1.8420635346185686), + (2.0, 25.0, -34.7182759850871489), + (2.0, 800, -1130.4942582548682739)]) + def test_vonmises_logpdf(self, x, kappa, expected_logpdf): + logpdf = stats.vonmises.logpdf(x, kappa) + assert_allclose(logpdf, expected_logpdf, rtol=1e-15) + + def test_vonmises_expect(self): + """ + Test that the vonmises expectation values are + computed correctly. This test checks that the + numeric integration estimates the correct normalization + (1) and mean angle (loc). These expectations are + independent of the chosen 2pi interval. + """ + rng = np.random.default_rng(6762668991392531563) + + loc, kappa, lb = rng.random(3) * 10 + res = stats.vonmises(loc=loc, kappa=kappa).expect(lambda x: 1) + assert_allclose(res, 1) + assert np.issubdtype(res.dtype, np.floating) + + bounds = lb, lb + 2 * np.pi + res = stats.vonmises(loc=loc, kappa=kappa).expect(lambda x: 1, *bounds) + assert_allclose(res, 1) + assert np.issubdtype(res.dtype, np.floating) + + bounds = lb, lb + 2 * np.pi + res = stats.vonmises(loc=loc, kappa=kappa).expect(lambda x: np.exp(1j*x), + *bounds, complex_func=1) + assert_allclose(np.angle(res), loc % (2*np.pi)) + assert np.issubdtype(res.dtype, np.complexfloating) + + @pytest.mark.xslow + @pytest.mark.parametrize("rvs_loc", [0, 2]) + @pytest.mark.parametrize("rvs_shape", [1, 100, 1e8]) + @pytest.mark.parametrize('fix_loc', [True, False]) + @pytest.mark.parametrize('fix_shape', [True, False]) + def test_fit_MLE_comp_optimizer(self, rvs_loc, rvs_shape, + fix_loc, fix_shape): + if fix_shape and fix_loc: + pytest.skip("Nothing to fit.") + + rng = np.random.default_rng(6762668991392531563) + data = stats.vonmises.rvs(rvs_shape, size=1000, loc=rvs_loc, + random_state=rng) + + kwds = {'fscale': 1} + if fix_loc: + kwds['floc'] = rvs_loc + if fix_shape: + kwds['f0'] = rvs_shape + + _assert_less_or_close_loglike(stats.vonmises, data, + stats.vonmises.nnlf, **kwds) + + @pytest.mark.slow + def test_vonmises_fit_bad_floc(self): + data = [-0.92923506, -0.32498224, 0.13054989, -0.97252014, 2.79658071, + -0.89110948, 1.22520295, 1.44398065, 2.49163859, 1.50315096, + 3.05437696, -2.73126329, -3.06272048, 1.64647173, 1.94509247, + -1.14328023, 0.8499056, 2.36714682, -1.6823179, -0.88359996] + data = np.asarray(data) + loc = -0.5 * np.pi + kappa_fit, loc_fit, scale_fit = stats.vonmises.fit(data, floc=loc) + assert kappa_fit == np.finfo(float).tiny + _assert_less_or_close_loglike(stats.vonmises, data, + stats.vonmises.nnlf, fscale=1, floc=loc) + + @pytest.mark.parametrize('sign', [-1, 1]) + def test_vonmises_fit_unwrapped_data(self, sign): + rng = np.random.default_rng(6762668991392531563) + data = stats.vonmises(loc=sign*0.5*np.pi, kappa=10).rvs(100000, + random_state=rng) + shifted_data = data + 4*np.pi + kappa_fit, loc_fit, scale_fit = stats.vonmises.fit(data) + kappa_fit_shifted, loc_fit_shifted, _ = stats.vonmises.fit(shifted_data) + assert_allclose(loc_fit, loc_fit_shifted) + assert_allclose(kappa_fit, kappa_fit_shifted) + assert scale_fit == 1 + assert -np.pi < loc_fit < np.pi + + def test_vonmises_kappa_0_gh18166(self): + # Check that kappa = 0 is supported. + dist = stats.vonmises(0) + assert_allclose(dist.pdf(0), 1 / (2 * np.pi), rtol=1e-15) + assert_allclose(dist.cdf(np.pi/2), 0.75, rtol=1e-15) + assert_allclose(dist.sf(-np.pi/2), 0.75, rtol=1e-15) + assert_allclose(dist.ppf(0.9), np.pi*0.8, rtol=1e-15) + assert_allclose(dist.mean(), 0, atol=1e-15) + assert_allclose(dist.expect(), 0, atol=1e-15) + assert np.all(np.abs(dist.rvs(size=10, random_state=self.rng)) <= np.pi) + + def test_vonmises_fit_equal_data(self): + # When all data are equal, expect kappa = 1e16. + kappa, loc, scale = stats.vonmises.fit([0]) + assert kappa == 1e16 and loc == 0 and scale == 1 + + def test_vonmises_fit_bounds(self): + # For certain input data, the root bracket is violated numerically. + # Test that this situation is handled. The input data below are + # crafted to trigger the bound violation for the current choice of + # bounds and the specific way the bounds and the objective function + # are computed. + + # Test that no exception is raised when the lower bound is violated. + scipy.stats.vonmises.fit([0, 3.7e-08], floc=0) + + # Test that no exception is raised when the upper bound is violated. + scipy.stats.vonmises.fit([np.pi/2*(1-4.86e-9)], floc=0) + + +def _assert_less_or_close_loglike(dist, data, func=None, maybe_identical=False, + **kwds): + """ + This utility function checks that the negative log-likelihood function + (or `func`) of the result computed using dist.fit() is less than or equal + to the result computed using the generic fit method. Because of + normal numerical imprecision, the "equality" check is made using + `np.allclose` with a relative tolerance of 1e-15. + """ + if func is None: + func = dist.nnlf + + mle_analytical = dist.fit(data, **kwds) + numerical_opt = super(type(dist), dist).fit(data, **kwds) + + # Sanity check that the analytical MLE is actually executed. + # Due to floating point arithmetic, the generic MLE is unlikely + # to produce the exact same result as the analytical MLE. + if not maybe_identical: + assert np.any(mle_analytical != numerical_opt) + + ll_mle_analytical = func(mle_analytical, data) + ll_numerical_opt = func(numerical_opt, data) + assert (ll_mle_analytical <= ll_numerical_opt or + np.allclose(ll_mle_analytical, ll_numerical_opt, rtol=1e-15)) + + # Ideally we'd check that shapes are correctly fixed, too, but that is + # complicated by the many ways of fixing them (e.g. f0, fix_a, fa). + if 'floc' in kwds: + assert mle_analytical[-2] == kwds['floc'] + if 'fscale' in kwds: + assert mle_analytical[-1] == kwds['fscale'] + + +def assert_fit_warnings(dist): + param = ['floc', 'fscale'] + if dist.shapes: + nshapes = len(dist.shapes.split(",")) + param += ['f0', 'f1', 'f2'][:nshapes] + all_fixed = dict(zip(param, np.arange(len(param)))) + data = [1, 2, 3] + with pytest.raises(RuntimeError, + match="All parameters fixed. There is nothing " + "to optimize."): + dist.fit(data, **all_fixed) + with pytest.raises(ValueError, + match="The data contains non-finite values"): + dist.fit([np.nan]) + with pytest.raises(ValueError, + match="The data contains non-finite values"): + dist.fit([np.inf]) + with pytest.raises(TypeError, match="Unknown keyword arguments:"): + dist.fit(data, extra_keyword=2) + with pytest.raises(TypeError, match="Too many positional arguments."): + dist.fit(data, *[1]*(len(param) - 1)) + + +@pytest.mark.parametrize('dist', + ['alpha', 'betaprime', + 'fatiguelife', 'invgamma', 'invgauss', 'invweibull', + 'johnsonsb', 'levy', 'levy_l', 'lognorm', 'gibrat', + 'powerlognorm', 'rayleigh', 'wald']) +def test_support(dist): + """gh-6235""" + dct = dict(distcont) + args = dct[dist] + + dist = getattr(stats, dist) + + assert_almost_equal(dist.pdf(dist.a, *args), 0) + assert_equal(dist.logpdf(dist.a, *args), -np.inf) + assert_almost_equal(dist.pdf(dist.b, *args), 0) + assert_equal(dist.logpdf(dist.b, *args), -np.inf) + + +class TestRandInt: + def setup_method(self): + self.rng = np.random.default_rng(8826485737) + + def test_rvs(self): + vals = stats.randint.rvs(5, 30, size=100, random_state=self.rng) + assert_(np.all(vals < 30) & np.all(vals >= 5)) + assert_(len(vals) == 100) + vals = stats.randint.rvs(5, 30, size=(2, 50), random_state=self.rng) + assert_(np.shape(vals) == (2, 50)) + assert_(vals.dtype.char in typecodes['AllInteger']) + val = stats.randint.rvs(15, 46, random_state=self.rng) + assert_((val >= 15) & (val < 46)) + assert_(isinstance(val, np.ScalarType), msg=repr(type(val))) + val = stats.randint(15, 46).rvs(3, random_state=self.rng) + assert_(val.dtype.char in typecodes['AllInteger']) + + def test_pdf(self): + k = np.r_[0:36] + out = np.where((k >= 5) & (k < 30), 1.0/(30-5), 0) + vals = stats.randint.pmf(k, 5, 30) + assert_array_almost_equal(vals, out) + + def test_cdf(self): + x = np.linspace(0, 36, 100) + k = np.floor(x) + out = np.select([k >= 30, k >= 5], [1.0, (k-5.0+1)/(30-5.0)], 0) + vals = stats.randint.cdf(x, 5, 30) + assert_array_almost_equal(vals, out, decimal=12) + + +class TestBinom: + def setup_method(self): + self.rng = np.random.default_rng(1778595878) + + def test_rvs(self): + vals = stats.binom.rvs(10, 0.75, size=(2, 50), random_state=self.rng) + assert_(np.all(vals >= 0) & np.all(vals <= 10)) + assert_(np.shape(vals) == (2, 50)) + assert_(vals.dtype.char in typecodes['AllInteger']) + val = stats.binom.rvs(10, 0.75, random_state=self.rng) + assert_(isinstance(val, int)) + val = stats.binom(10, 0.75).rvs(3, random_state=self.rng) + assert_(isinstance(val, np.ndarray)) + assert_(val.dtype.char in typecodes['AllInteger']) + + def test_pmf(self): + # regression test for Ticket #1842 + vals1 = stats.binom.pmf(100, 100, 1) + vals2 = stats.binom.pmf(0, 100, 0) + assert_allclose(vals1, 1.0, rtol=1e-15, atol=0) + assert_allclose(vals2, 1.0, rtol=1e-15, atol=0) + + def test_entropy(self): + # Basic entropy tests. + b = stats.binom(2, 0.5) + expected_p = np.array([0.25, 0.5, 0.25]) + expected_h = -sum(xlogy(expected_p, expected_p)) + h = b.entropy() + assert_allclose(h, expected_h) + + b = stats.binom(2, 0.0) + h = b.entropy() + assert_equal(h, 0.0) + + b = stats.binom(2, 1.0) + h = b.entropy() + assert_equal(h, 0.0) + + def test_warns_p0(self): + # no spurious warnings are generated for p=0; gh-3817 + with warnings.catch_warnings(): + warnings.simplefilter("error", RuntimeWarning) + assert_equal(stats.binom(n=2, p=0).mean(), 0) + assert_equal(stats.binom(n=2, p=0).std(), 0) + + def test_ppf_p1(self): + # Check that gh-17388 is resolved: PPF == n when p = 1 + n = 4 + assert stats.binom.ppf(q=0.3, n=n, p=1.0) == n + + def test_pmf_poisson(self): + # Check that gh-17146 is resolved: binom -> poisson + n = 1541096362225563.0 + p = 1.0477878413173978e-18 + x = np.arange(3) + res = stats.binom.pmf(x, n=n, p=p) + ref = stats.poisson.pmf(x, n * p) + assert_allclose(res, ref, atol=1e-16) + + def test_pmf_cdf(self): + # Check that gh-17809 is resolved: binom.pmf(0) ~ binom.cdf(0) + n = 25.0 * 10 ** 21 + p = 1.0 * 10 ** -21 + r = 0 + res = stats.binom.pmf(r, n, p) + ref = stats.binom.cdf(r, n, p) + assert_allclose(res, ref, atol=1e-16) + + def test_pmf_gh15101(self): + # Check that gh-15101 is resolved (no divide warnings when p~1, n~oo) + res = stats.binom.pmf(3, 2000, 0.999) + assert_allclose(res, 0, atol=1e-16) + + +class TestArcsine: + + def test_endpoints(self): + # Regression test for gh-13697. The following calculation + # should not generate a warning. + p = stats.arcsine.pdf([0, 1]) + assert_equal(p, [np.inf, np.inf]) + + +class TestBernoulli: + def setup_method(self): + self.rng = np.random.default_rng(7836792223) + + def test_rvs(self): + vals = stats.bernoulli.rvs(0.75, size=(2, 50), random_state=self.rng) + assert_(np.all(vals >= 0) & np.all(vals <= 1)) + assert_(np.shape(vals) == (2, 50)) + assert_(vals.dtype.char in typecodes['AllInteger']) + val = stats.bernoulli.rvs(0.75, random_state=self.rng) + assert_(isinstance(val, int)) + val = stats.bernoulli(0.75).rvs(3, random_state=self.rng) + assert_(isinstance(val, np.ndarray)) + assert_(val.dtype.char in typecodes['AllInteger']) + + def test_entropy(self): + # Simple tests of entropy. + b = stats.bernoulli(0.25) + expected_h = -0.25*np.log(0.25) - 0.75*np.log(0.75) + h = b.entropy() + assert_allclose(h, expected_h) + + b = stats.bernoulli(0.0) + h = b.entropy() + assert_equal(h, 0.0) + + b = stats.bernoulli(1.0) + h = b.entropy() + assert_equal(h, 0.0) + + +class TestBradford: + # gh-6216 + def test_cdf_ppf(self): + c = 0.1 + x = np.logspace(-20, -4) + q = stats.bradford.cdf(x, c) + xx = stats.bradford.ppf(q, c) + assert_allclose(x, xx) + + +class TestCauchy: + + def test_pdf_no_overflow_warning(self): + # The argument is large enough that x**2 will overflow to + # infinity and 1/(1 + x**2) will be 0. This should not + # trigger a warning. + p = stats.cauchy.pdf(1e200) + assert p == 0.0 + + # Reference values were computed with mpmath. + @pytest.mark.parametrize( + 'x, ref', + [(0.0, -1.1447298858494002), + (5e-324, -1.1447298858494002), + (1e-34, -1.1447298858494002), + (2.2e-16, -1.1447298858494002), + (2e-8, -1.1447298858494006), + (5e-4, -1.144730135849369), + (0.1, -1.1546802167025683), + (1.5, -2.3233848821910463), + (2e18, -85.42408759475494), + (1e200, -922.1787670834676), + (_XMAX, -1420.7101556726175)]) + def test_logpdf(self, x, ref): + logp = stats.cauchy.logpdf([x, -x]) + assert_allclose(logp, [ref, ref], rtol=1e-15) + + # Reference values were computed with mpmath. + @pytest.mark.parametrize( + 'x, ref', + [(-5e15, 6.366197723675814e-17), + (-5, 0.06283295818900118), + (-1, 0.25), + (0, 0.5), + (1, 0.75), + (5, 0.9371670418109989), + (5e15, 0.9999999999999999)] + ) + @pytest.mark.parametrize( + 'method, sgn', + [(stats.cauchy.cdf, 1), + (stats.cauchy.sf, -1)] + ) + def test_cdf_sf(self, x, ref, method, sgn): + p = method(sgn*x) + assert_allclose(p, ref, rtol=1e-15) + + # Reference values were computed with mpmath. + @pytest.mark.parametrize('x, ref', + [(4e250, -7.957747154594767e-252), + (1e25, -3.1830988618379063e-26), + (10.0, -0.03223967552667532), + (0.0, -0.6931471805599453), + (-10.0, -3.4506339556469654), + (-7e45, -106.70696921963678), + (-3e225, -520.3249880981778)]) + def test_logcdf_logsf(self, x, ref): + logcdf = stats.cauchy.logcdf(x) + assert_allclose(logcdf, ref, rtol=5e-15) + logsf = stats.cauchy.logsf(-x) + assert_allclose(logsf, ref, rtol=5e-15) + + # Reference values were computed with mpmath. + @pytest.mark.parametrize( + 'p, ref', + [(1e-20, -3.1830988618379067e+19), + (1e-9, -318309886.1837906), + (0.25, -1.0), + (0.50, 0.0), + (0.75, 1.0), + (0.999999, 318309.88617359026), + (0.999999999999, 318316927901.77966)] + ) + @pytest.mark.parametrize( + 'method, sgn', + [(stats.cauchy.ppf, 1), + (stats.cauchy.isf, -1)]) + def test_ppf_isf(self, p, ref, method, sgn): + x = sgn*method(p) + assert_allclose(x, ref, rtol=1e-15) + + +class TestChi: + + # "Exact" value of chi.sf(10, 4), as computed by Wolfram Alpha with + # 1 - CDF[ChiDistribution[4], 10] + CHI_SF_10_4 = 9.83662422461598e-21 + # "Exact" value of chi.mean(df=1000) as computed by Wolfram Alpha with + # Mean[ChiDistribution[1000]] + CHI_MEAN_1000 = 31.614871896980 + + def test_sf(self): + s = stats.chi.sf(10, 4) + assert_allclose(s, self.CHI_SF_10_4, rtol=1e-15) + + def test_isf(self): + x = stats.chi.isf(self.CHI_SF_10_4, 4) + assert_allclose(x, 10, rtol=1e-15) + + def test_logcdf(self): + x = 10.0 + df = 15 + logcdf = stats.chi.logcdf(x, df) + # Reference value computed with mpath. + assert_allclose(logcdf, -1.304704343625153e-14, rtol=5e-15) + + def test_logsf(self): + x = 0.01 + df = 15 + logsf = stats.chi.logsf(x, df) + # Reference value computed with mpath. + assert_allclose(logsf, -3.936060782678026e-37, rtol=5e-15) + + # reference value for 1e14 was computed via mpmath + # from mpmath import mp + # mp.dps = 500 + # df = mp.mpf(1e14) + # float(mp.rf(mp.mpf(0.5) * df, mp.mpf(0.5)) * mp.sqrt(2.)) + + @pytest.mark.parametrize('df, ref', + [(1e3, CHI_MEAN_1000), + (1e14, 9999999.999999976)]) + def test_mean(self, df, ref): + assert_allclose(stats.chi.mean(df), ref, rtol=1e-12) + + # Entropy references values were computed with the following mpmath code + # from mpmath import mp + # mp.dps = 50 + # def chi_entropy_mpmath(df): + # df = mp.mpf(df) + # half_df = 0.5 * df + # entropy = mp.log(mp.gamma(half_df)) + 0.5 * \ + # (df - mp.log(2) - (df - mp.one) * mp.digamma(half_df)) + # return float(entropy) + + @pytest.mark.parametrize('df, ref', + [(1e-4, -9989.7316027504), + (1, 0.7257913526447274), + (1e3, 1.0721981095025448), + (1e10, 1.0723649429080335), + (1e100, 1.0723649429247002)]) + def test_entropy(self, df, ref): + assert_allclose(stats.chi(df).entropy(), ref, rtol=1e-15) + + +class TestCrystalBall: + + def test_pdf(self): + """ + All values are calculated using the independent implementation of the + ROOT framework (see https://root.cern.ch/). + Corresponding ROOT code is given in the comments. + """ + X = np.linspace(-5.0, 5.0, 21)[:-1] + + # for (double x = -5.0; x < 5.0; x += 0.5) { + # cout << setprecision(16) + # << ROOT::Math::crystalball_pdf(x, 1.0, 2.0, 1.0) + # << ", "; + # } + calculated = stats.crystalball.pdf(X, beta=1.0, m=2.0) + expected = np.array([0.02028666423671257, 0.02414280702550917, + 0.02921279650086611, 0.03606518086526679, + 0.04564499453260328, 0.05961795204258388, + 0.08114665694685029, 0.1168511860034644, + 0.1825799781304131, 0.2656523006609301, + 0.3010234935475763, 0.2656523006609301, + 0.1825799781304131, 0.09772801991305094, + 0.0407390997601359, 0.01322604925508607, + 0.003344068947749631, 0.0006584862184997063, + 0.0001009821322058648, 1.206059579124873e-05]) + assert_allclose(expected, calculated, rtol=1e-14) + + # for (double x = -5.0; x < 5.0; x += 0.5) { + # cout << setprecision(16) + # << ROOT::Math::crystalball_pdf(x, 2.0, 3.0, 1.0) + # << ", "; + # } + calculated = stats.crystalball.pdf(X, beta=2.0, m=3.0) + expected = np.array([0.00196480373120913, 0.0027975428126005, + 0.004175923965164595, 0.006631212592830816, + 0.01145873536041165, 0.022380342500804, + 0.05304970074264653, 0.1272596164638828, + 0.237752264003024, 0.3459275029304401, + 0.3919872148188981, 0.3459275029304401, + 0.237752264003024, 0.1272596164638828, + 0.05304970074264653, 0.01722271623872227, + 0.004354584612458383, 0.0008574685508575863, + 0.000131497061187334, 1.570508433595375e-05]) + assert_allclose(expected, calculated, rtol=1e-14) + + # for (double x = -5.0; x < 5.0; x += 0.5) { + # cout << setprecision(16) + # << ROOT::Math::crystalball_pdf(x, 2.0, 3.0, 2.0, 0.5) + # << ", "; + # } + calculated = stats.crystalball.pdf(X, beta=2.0, m=3.0, loc=0.5, scale=2.0) + expected = np.array([0.007859214924836521, 0.011190171250402, + 0.01670369586065838, 0.02652485037132326, + 0.04238659020399594, 0.06362980823194138, + 0.08973241216601403, 0.118876132001512, + 0.1479437366093383, 0.17296375146522, + 0.1899635180461471, 0.1959936074094491, + 0.1899635180461471, 0.17296375146522, + 0.1479437366093383, 0.118876132001512, + 0.08973241216601403, 0.06362980823194138, + 0.04238659020399594, 0.02652485037132326]) + assert_allclose(expected, calculated, rtol=1e-14) + + def test_cdf(self): + """ + All values are calculated using the independent implementation of the + ROOT framework (see https://root.cern.ch/). + Corresponding ROOT code is given in the comments. + """ + X = np.linspace(-5.0, 5.0, 21)[:-1] + + # for (double x = -5.0; x < 5.0; x += 0.5) { + # cout << setprecision(16) + # << ROOT::Math::crystalball_cdf(x, 1.0, 2.0, 1.0) + # << ", "; + # } + calculated = stats.crystalball.cdf(X, beta=1.0, m=2.0) + expected = np.array([0.1217199854202754, 0.1327854386403005, + 0.1460639825043305, 0.1622933138937006, + 0.1825799781304132, 0.2086628321490436, + 0.2434399708405509, 0.292127965008661, + 0.3651599562608263, 0.4782542338198316, + 0.6227229998727213, 0.7671917659256111, + 0.8802860434846165, 0.9495903590367718, + 0.9828337969321823, 0.9953144721881936, + 0.9989814290402977, 0.9998244687978383, + 0.9999761023377818, 0.9999974362721522]) + assert_allclose(expected, calculated, rtol=1e-13) + + # for (double x = -5.0; x < 5.0; x += 0.5) { + # cout << setprecision(16) + # << ROOT::Math::crystalball_cdf(x, 2.0, 3.0, 1.0) + # << ", "; + # } + calculated = stats.crystalball.cdf(X, beta=2.0, m=3.0) + expected = np.array([0.004420808395220632, 0.005595085625200946, + 0.007307866939038177, 0.009946818889246312, + 0.01432341920051472, 0.02238034250080412, + 0.03978727555698502, 0.08307626432678494, + 0.1733230597116304, 0.3205923321191123, + 0.508716882020547, 0.6968414319219818, + 0.8441107043294638, 0.934357499714309, + 0.9776464884841091, 0.9938985925142876, + 0.9986736357721329, 0.9997714265214375, + 0.9999688809071239, 0.9999966615611068]) + assert_allclose(expected, calculated, rtol=1e-13) + + # for (double x = -5.0; x < 5.0; x += 0.5) { + # cout << setprecision(16) + # << ROOT::Math::crystalball_cdf(x, 2.0, 3.0, 2.0, 0.5); + # << ", "; + # } + calculated = stats.crystalball.cdf(X, beta=2.0, m=3.0, loc=0.5, scale=2.0) + expected = np.array([0.0176832335808822, 0.02238034250080412, + 0.02923146775615237, 0.03978727555698502, + 0.05679453901646225, 0.08307626432678494, + 0.1212416644828466, 0.1733230597116304, + 0.2401101486313661, 0.3205923321191123, + 0.4117313791289429, 0.508716882020547, + 0.6057023849121512, 0.6968414319219818, + 0.7773236154097279, 0.8441107043294638, + 0.8961920995582476, 0.934357499714309, + 0.9606392250246318, 0.9776464884841091]) + assert_allclose(expected, calculated, rtol=1e-13) + + # Reference value computed with ROOT, e.g. + # cout << setprecision(16) + # << ROOT::Math::crystalball_cdf_c(12.0, 1.0, 2.0, 1.0) + # << endl; + @pytest.mark.parametrize( + 'x, beta, m, rootref', + [(12.0, 1.0, 2.0, 1.340451684048897e-33), + (9.0, 4.0, 1.25, 1.12843537145273e-19), + (20, 0.1, 1.001, 6.929038716892384e-93), + (-4.5, 2.0, 3.0, 0.9944049143747991), + (-30.0, 0.5, 5.0, 0.9976994814571858), + (-1e50, 1.5, 1.1, 0.9999951099570382)] + ) + def test_sf(self, x, beta, m, rootref): + sf = stats.crystalball.sf(x, beta=beta, m=m) + assert_allclose(sf, rootref, rtol=1e-13) + + def test_moments(self): + """ + All values are calculated using the pdf formula and the integrate function + of Mathematica + """ + # The Last two (alpha, n) pairs test the special case n == alpha**2 + beta = np.array([2.0, 1.0, 3.0, 2.0, 3.0]) + m = np.array([3.0, 3.0, 2.0, 4.0, 9.0]) + + # The distribution should be correctly normalised + expected_0th_moment = np.array([1.0, 1.0, 1.0, 1.0, 1.0]) + calculated_0th_moment = stats.crystalball._munp(0, beta, m) + assert_allclose(expected_0th_moment, calculated_0th_moment, rtol=0.001) + + # calculated using wolframalpha.com + # e.g. for beta = 2 and m = 3 we calculate the norm like this: + # integrate exp(-x^2/2) from -2 to infinity + + # integrate (3/2)^3*exp(-2^2/2)*(3/2-2-x)^(-3) from -infinity to -2 + norm = np.array([2.5511, 3.01873, 2.51065, 2.53983, 2.507410455]) + + a = np.array([-0.21992, -3.03265, np.inf, -0.135335, -0.003174]) + expected_1th_moment = a / norm + calculated_1th_moment = stats.crystalball._munp(1, beta, m) + assert_allclose(expected_1th_moment, calculated_1th_moment, rtol=0.001) + + a = np.array([np.inf, np.inf, np.inf, 3.2616, 2.519908]) + expected_2th_moment = a / norm + calculated_2th_moment = stats.crystalball._munp(2, beta, m) + assert_allclose(expected_2th_moment, calculated_2th_moment, rtol=0.001) + + a = np.array([np.inf, np.inf, np.inf, np.inf, -0.0577668]) + expected_3th_moment = a / norm + calculated_3th_moment = stats.crystalball._munp(3, beta, m) + assert_allclose(expected_3th_moment, calculated_3th_moment, rtol=0.001) + + a = np.array([np.inf, np.inf, np.inf, np.inf, 7.78468]) + expected_4th_moment = a / norm + calculated_4th_moment = stats.crystalball._munp(4, beta, m) + assert_allclose(expected_4th_moment, calculated_4th_moment, rtol=0.001) + + a = np.array([np.inf, np.inf, np.inf, np.inf, -1.31086]) + expected_5th_moment = a / norm + calculated_5th_moment = stats.crystalball._munp(5, beta, m) + assert_allclose(expected_5th_moment, calculated_5th_moment, rtol=0.001) + + def test_entropy(self): + # regression test for gh-13602 + cb = stats.crystalball(2, 3) + res1 = cb.entropy() + # -20000 and 30 are negative and positive infinity, respectively + lo, hi, N = -20000, 30, 200000 + x = np.linspace(lo, hi, N) + res2 = trapezoid(entr(cb.pdf(x)), x) + assert_allclose(res1, res2, rtol=1e-7) + + +class TestNBinom: + def setup_method(self): + self.rng = np.random.default_rng(5861367021) + + def test_rvs(self): + vals = stats.nbinom.rvs(10, 0.75, size=(2, 50), random_state=self.rng) + assert_(np.all(vals >= 0)) + assert_(np.shape(vals) == (2, 50)) + assert_(vals.dtype.char in typecodes['AllInteger']) + val = stats.nbinom.rvs(10, 0.75, random_state=self.rng) + assert_(isinstance(val, int)) + val = stats.nbinom(10, 0.75).rvs(3, random_state=self.rng) + assert_(isinstance(val, np.ndarray)) + assert_(val.dtype.char in typecodes['AllInteger']) + + def test_pmf(self): + # regression test for ticket 1779 + assert_allclose(np.exp(stats.nbinom.logpmf(700, 721, 0.52)), + stats.nbinom.pmf(700, 721, 0.52)) + # logpmf(0,1,1) shouldn't return nan (regression test for gh-4029) + val = scipy.stats.nbinom.logpmf(0, 1, 1) + assert_equal(val, 0) + + def test_logcdf_gh16159(self): + # check that gh16159 is resolved. + vals = stats.nbinom.logcdf([0, 5, 0, 5], n=4.8, p=0.45) + ref = np.log(stats.nbinom.cdf([0, 5, 0, 5], n=4.8, p=0.45)) + assert_allclose(vals, ref) + + +class TestGenInvGauss: + def setup_method(self): + self.rng = np.random.default_rng(6473281180) + + @pytest.mark.slow + def test_rvs_with_mode_shift(self): + # ratio_unif w/ mode shift + gig = stats.geninvgauss(2.3, 1.5) + _, p = stats.kstest(gig.rvs(size=1500, random_state=self.rng), gig.cdf) + assert_equal(p > 0.05, True) + + @pytest.mark.slow + def test_rvs_without_mode_shift(self): + # ratio_unif w/o mode shift + gig = stats.geninvgauss(0.9, 0.75) + _, p = stats.kstest(gig.rvs(size=1500, random_state=self.rng), gig.cdf) + assert_equal(p > 0.05, True) + + @pytest.mark.slow + def test_rvs_new_method(self): + # new algorithm of Hoermann / Leydold + gig = stats.geninvgauss(0.1, 0.2) + _, p = stats.kstest(gig.rvs(size=1500, random_state=self.rng), gig.cdf) + assert_equal(p > 0.05, True) + + @pytest.mark.slow + def test_rvs_p_zero(self): + def my_ks_check(p, b): + gig = stats.geninvgauss(p, b) + rvs = gig.rvs(size=1500, random_state=self.rng) + return stats.kstest(rvs, gig.cdf)[1] > 0.05 + # boundary cases when p = 0 + assert_equal(my_ks_check(0, 0.2), True) # new algo + assert_equal(my_ks_check(0, 0.9), True) # ratio_unif w/o shift + assert_equal(my_ks_check(0, 1.5), True) # ratio_unif with shift + + def test_rvs_negative_p(self): + # if p negative, return inverse + assert_equal( + stats.geninvgauss(-1.5, 2).rvs(size=10, random_state=1234), + 1 / stats.geninvgauss(1.5, 2).rvs(size=10, random_state=1234)) + + def test_invgauss(self): + # test that invgauss is special case + ig = stats.geninvgauss.rvs(size=1500, p=-0.5, b=1, random_state=1464878613) + assert_equal(stats.kstest(ig, 'invgauss', args=[1])[1] > 0.15, True) + # test pdf and cdf + mu, x = 100, np.linspace(0.01, 1, 10) + pdf_ig = stats.geninvgauss.pdf(x, p=-0.5, b=1 / mu, scale=mu) + assert_allclose(pdf_ig, stats.invgauss(mu).pdf(x)) + cdf_ig = stats.geninvgauss.cdf(x, p=-0.5, b=1 / mu, scale=mu) + assert_allclose(cdf_ig, stats.invgauss(mu).cdf(x)) + + def test_pdf_R(self): + # test against R package GIGrvg + # x <- seq(0.01, 5, length.out = 10) + # GIGrvg::dgig(x, 0.5, 1, 1) + vals_R = np.array([2.081176820e-21, 4.488660034e-01, 3.747774338e-01, + 2.693297528e-01, 1.905637275e-01, 1.351476913e-01, + 9.636538981e-02, 6.909040154e-02, 4.978006801e-02, + 3.602084467e-02]) + x = np.linspace(0.01, 5, 10) + assert_allclose(vals_R, stats.geninvgauss.pdf(x, 0.5, 1)) + + def test_pdf_zero(self): + # pdf at 0 is 0, needs special treatment to avoid 1/x in pdf + assert_equal(stats.geninvgauss.pdf(0, 0.5, 0.5), 0) + # if x is large and p is moderate, make sure that pdf does not + # overflow because of x**(p-1); exp(-b*x) forces pdf to zero + assert_equal(stats.geninvgauss.pdf(2e6, 50, 2), 0) + + +class TestGenHyperbolic: + + def test_pdf_r(self): + # test against R package GeneralizedHyperbolic + # x <- seq(-10, 10, length.out = 10) + # GeneralizedHyperbolic::dghyp( + # x = x, lambda = 2, alpha = 2, beta = 1, delta = 1.5, mu = 0.5 + # ) + vals_R = np.array([ + 2.94895678275316e-13, 1.75746848647696e-10, 9.48149804073045e-08, + 4.17862521692026e-05, 0.0103947630463822, 0.240864958986839, + 0.162833527161649, 0.0374609592899472, 0.00634894847327781, + 0.000941920705790324 + ]) + + lmbda, alpha, beta = 2, 2, 1 + mu, delta = 0.5, 1.5 + args = (lmbda, alpha*delta, beta*delta) + + gh = stats.genhyperbolic(*args, loc=mu, scale=delta) + x = np.linspace(-10, 10, 10) + + assert_allclose(gh.pdf(x), vals_R, atol=0, rtol=1e-13) + + def test_cdf_r(self): + # test against R package GeneralizedHyperbolic + # q <- seq(-10, 10, length.out = 10) + # GeneralizedHyperbolic::pghyp( + # q = q, lambda = 2, alpha = 2, beta = 1, delta = 1.5, mu = 0.5 + # ) + vals_R = np.array([ + 1.01881590921421e-13, 6.13697274983578e-11, 3.37504977637992e-08, + 1.55258698166181e-05, 0.00447005453832497, 0.228935323956347, + 0.755759458895243, 0.953061062884484, 0.992598013917513, + 0.998942646586662 + ]) + + lmbda, alpha, beta = 2, 2, 1 + mu, delta = 0.5, 1.5 + args = (lmbda, alpha*delta, beta*delta) + + gh = stats.genhyperbolic(*args, loc=mu, scale=delta) + x = np.linspace(-10, 10, 10) + + assert_allclose(gh.cdf(x), vals_R, atol=0, rtol=1e-6) + + # The reference values were computed by implementing the PDF with mpmath + # and integrating it with mp.quad. The values were computed with + # mp.dps=250, and then again with mp.dps=400 to ensure the full 64 bit + # precision was computed. + @pytest.mark.parametrize( + 'x, p, a, b, loc, scale, ref', + [(-15, 2, 3, 1.5, 0.5, 1.5, 4.770036428808252e-20), + (-15, 10, 1.5, 0.25, 1, 5, 0.03282964575089294), + (-15, 10, 1.5, 1.375, 0, 1, 3.3711159600215594e-23), + (-15, 0.125, 1.5, 1.49995, 0, 1, 4.729401428898605e-23), + (-1, 0.125, 1.5, 1.49995, 0, 1, 0.0003565725914786859), + (5, -0.125, 1.5, 1.49995, 0, 1, 0.2600651974023352), + (5, -0.125, 1000, 999, 0, 1, 5.923270556517253e-28), + (20, -0.125, 1000, 999, 0, 1, 0.23452293711665634), + (40, -0.125, 1000, 999, 0, 1, 0.9999648749561968), + (60, -0.125, 1000, 999, 0, 1, 0.9999999999975475)] + ) + def test_cdf_mpmath(self, x, p, a, b, loc, scale, ref): + cdf = stats.genhyperbolic.cdf(x, p, a, b, loc=loc, scale=scale) + assert_allclose(cdf, ref, rtol=5e-12) + + # The reference values were computed by implementing the PDF with mpmath + # and integrating it with mp.quad. The values were computed with + # mp.dps=250, and then again with mp.dps=400 to ensure the full 64 bit + # precision was computed. + @pytest.mark.parametrize( + 'x, p, a, b, loc, scale, ref', + [(0, 1e-6, 12, -1, 0, 1, 0.38520358671350524), + (-1, 3, 2.5, 2.375, 1, 3, 0.9999901774267577), + (-20, 3, 2.5, 2.375, 1, 3, 1.0), + (25, 2, 3, 1.5, 0.5, 1.5, 8.593419916523976e-10), + (300, 10, 1.5, 0.25, 1, 5, 6.137415609872158e-24), + (60, -0.125, 1000, 999, 0, 1, 2.4524915075944173e-12), + (75, -0.125, 1000, 999, 0, 1, 2.9435194886214633e-18)] + ) + def test_sf_mpmath(self, x, p, a, b, loc, scale, ref): + sf = stats.genhyperbolic.sf(x, p, a, b, loc=loc, scale=scale) + assert_allclose(sf, ref, rtol=5e-12) + + def test_moments_r(self): + # test against R package GeneralizedHyperbolic + # sapply(1:4, + # function(x) GeneralizedHyperbolic::ghypMom( + # order = x, lambda = 2, alpha = 2, + # beta = 1, delta = 1.5, mu = 0.5, + # momType = 'raw') + # ) + + vals_R = [2.36848366948115, 8.4739346779246, + 37.8870502710066, 205.76608511485] + + lmbda, alpha, beta = 2, 2, 1 + mu, delta = 0.5, 1.5 + args = (lmbda, alpha*delta, beta*delta) + + vals_us = [ + stats.genhyperbolic(*args, loc=mu, scale=delta).moment(i) + for i in range(1, 5) + ] + + assert_allclose(vals_us, vals_R, atol=0, rtol=1e-13) + + def test_rvs(self): + # Kolmogorov-Smirnov test to ensure alignment + # of analytical and empirical cdfs + + lmbda, alpha, beta = 2, 2, 1 + mu, delta = 0.5, 1.5 + args = (lmbda, alpha*delta, beta*delta) + + gh = stats.genhyperbolic(*args, loc=mu, scale=delta) + _, p = stats.kstest(gh.rvs(size=1500, random_state=1234), gh.cdf) + + assert_equal(p > 0.05, True) + + def test_pdf_t(self): + # Test Against T-Student with 1 - 30 df + df = np.linspace(1, 30, 10) + + # in principle alpha should be zero in practice for big lmbdas + # alpha cannot be too small else pdf does not integrate + alpha, beta = np.float_power(df, 2)*np.finfo(np.float32).eps, 0 + mu, delta = 0, np.sqrt(df) + args = (-df/2, alpha, beta) + + gh = stats.genhyperbolic(*args, loc=mu, scale=delta) + x = np.linspace(gh.ppf(0.01), gh.ppf(0.99), 50)[:, np.newaxis] + + assert_allclose( + gh.pdf(x), stats.t.pdf(x, df), + atol=0, rtol=1e-6 + ) + + def test_pdf_cauchy(self): + # Test Against Cauchy distribution + + # in principle alpha should be zero in practice for big lmbdas + # alpha cannot be too small else pdf does not integrate + lmbda, alpha, beta = -0.5, np.finfo(np.float32).eps, 0 + mu, delta = 0, 1 + args = (lmbda, alpha, beta) + + gh = stats.genhyperbolic(*args, loc=mu, scale=delta) + x = np.linspace(gh.ppf(0.01), gh.ppf(0.99), 50)[:, np.newaxis] + + assert_allclose( + gh.pdf(x), stats.cauchy.pdf(x), + atol=0, rtol=1e-6 + ) + + def test_pdf_laplace(self): + # Test Against Laplace with location param [-10, 10] + loc = np.linspace(-10, 10, 10) + + # in principle delta should be zero in practice for big loc delta + # cannot be too small else pdf does not integrate + delta = np.finfo(np.float32).eps + + lmbda, alpha, beta = 1, 1, 0 + args = (lmbda, alpha*delta, beta*delta) + + # ppf does not integrate for scale < 5e-4 + # therefore using simple linspace to define the support + gh = stats.genhyperbolic(*args, loc=loc, scale=delta) + x = np.linspace(-20, 20, 50)[:, np.newaxis] + + assert_allclose( + gh.pdf(x), stats.laplace.pdf(x, loc=loc, scale=1), + atol=0, rtol=1e-11 + ) + + def test_pdf_norminvgauss(self): + # Test Against NIG with varying alpha/beta/delta/mu + + alpha, beta, delta, mu = ( + np.linspace(1, 20, 10), + np.linspace(0, 19, 10)*np.float_power(-1, range(10)), + np.linspace(1, 1, 10), + np.linspace(-100, 100, 10) + ) + + lmbda = - 0.5 + args = (lmbda, alpha * delta, beta * delta) + + gh = stats.genhyperbolic(*args, loc=mu, scale=delta) + x = np.linspace(gh.ppf(0.01), gh.ppf(0.99), 50)[:, np.newaxis] + + assert_allclose( + gh.pdf(x), stats.norminvgauss.pdf( + x, a=alpha, b=beta, loc=mu, scale=delta), + atol=0, rtol=1e-13 + ) + + +class TestHypSecant: + + # Reference values were computed with the mpmath expression + # float((2/mp.pi)*mp.atan(mp.exp(-x))) + # and mp.dps = 50. + @pytest.mark.parametrize('x, reference', + [(30, 5.957247804324683e-14), + (50, 1.2278802891647964e-22)]) + def test_sf(self, x, reference): + sf = stats.hypsecant.sf(x) + assert_allclose(sf, reference, rtol=5e-15) + + # Reference values were computed with the mpmath expression + # float(-mp.log(mp.tan((mp.pi/2)*p))) + # and mp.dps = 50. + @pytest.mark.parametrize('p, reference', + [(1e-6, 13.363927852673998), + (1e-12, 27.179438410639094)]) + def test_isf(self, p, reference): + x = stats.hypsecant.isf(p) + assert_allclose(x, reference, rtol=5e-15) + + def test_logcdf_logsf(self): + x = 50.0 + # Reference value was computed with mpmath. + ref = -1.2278802891647964e-22 + logcdf = stats.hypsecant.logcdf(x) + assert_allclose(logcdf, ref, rtol=5e-15) + logsf = stats.hypsecant.logsf(-x) + assert_allclose(logsf, ref, rtol=5e-15) + + +class TestNormInvGauss: + def test_cdf_R(self): + # test pdf and cdf vals against R + # require("GeneralizedHyperbolic") + # x_test <- c(-7, -5, 0, 8, 15) + # r_cdf <- GeneralizedHyperbolic::pnig(x_test, mu = 0, a = 1, b = 0.5) + # r_pdf <- GeneralizedHyperbolic::dnig(x_test, mu = 0, a = 1, b = 0.5) + r_cdf = np.array([8.034920282e-07, 2.512671945e-05, 3.186661051e-01, + 9.988650664e-01, 9.999848769e-01]) + x_test = np.array([-7, -5, 0, 8, 15]) + vals_cdf = stats.norminvgauss.cdf(x_test, a=1, b=0.5) + assert_allclose(vals_cdf, r_cdf, atol=1e-9) + + def test_pdf_R(self): + # values from R as defined in test_cdf_R + r_pdf = np.array([1.359600783e-06, 4.413878805e-05, 4.555014266e-01, + 7.450485342e-04, 8.917889931e-06]) + x_test = np.array([-7, -5, 0, 8, 15]) + vals_pdf = stats.norminvgauss.pdf(x_test, a=1, b=0.5) + assert_allclose(vals_pdf, r_pdf, atol=1e-9) + + @pytest.mark.parametrize('x, a, b, sf, rtol', + [(-1, 1, 0, 0.8759652211005315, 1e-13), + (25, 1, 0, 1.1318690184042579e-13, 1e-4), + (1, 5, -1.5, 0.002066711134653577, 1e-12), + (10, 5, -1.5, 2.308435233930669e-29, 1e-9)]) + def test_sf_isf_mpmath(self, x, a, b, sf, rtol): + # Reference data generated with `reference_distributions.NormInvGauss`, + # e.g. `NormInvGauss(alpha=1, beta=0).sf(-1)` with mp.dps = 50 + s = stats.norminvgauss.sf(x, a, b) + assert_allclose(s, sf, rtol=rtol) + i = stats.norminvgauss.isf(sf, a, b) + assert_allclose(i, x, rtol=rtol) + + def test_sf_isf_mpmath_vectorized(self): + x = [-1, 25] + a = [1, 1] + b = 0 + sf = [0.8759652211005315, 1.1318690184042579e-13] # see previous test + s = stats.norminvgauss.sf(x, a, b) + assert_allclose(s, sf, rtol=1e-13, atol=1e-16) + i = stats.norminvgauss.isf(sf, a, b) + # Not perfect, but better than it was. See gh-13338. + assert_allclose(i, x, rtol=1e-6) + + def test_gh8718(self): + # Add test that gh-13338 resolved gh-8718 + dst = stats.norminvgauss(1, 0) + x = np.arange(0, 20, 2) + sf = dst.sf(x) + isf = dst.isf(sf) + assert_allclose(isf, x) + + def test_stats(self): + a, b = 1, 0.5 + gamma = np.sqrt(a**2 - b**2) + v_stats = (b / gamma, a**2 / gamma**3, 3.0 * b / (a * np.sqrt(gamma)), + 3.0 * (1 + 4 * b**2 / a**2) / gamma) + assert_equal(v_stats, stats.norminvgauss.stats(a, b, moments='mvsk')) + + def test_ppf(self): + a, b = 1, 0.5 + x_test = np.array([0.001, 0.5, 0.999]) + vals = stats.norminvgauss.ppf(x_test, a, b) + assert_allclose(x_test, stats.norminvgauss.cdf(vals, a, b)) + + @pytest.mark.parametrize('a', [0.1, 1.0, 10.0]) + @pytest.mark.parametrize('b_a', [-0.9, -0.5, -0.1, 0.1, 0.5, 0.9]) + def test_gh23196_tails_cdf_sf(self, a, b_a): + # gh-23196 reported that the norminvgauss CDF would drop to zero in the far + # right tail. The SF had a similar problem, dropping to zero at the far left. + # Check that this is resolved by checking that values in the deep tails are + # close to 1.0. + b = b_a * a + res = stats.norminvgauss(a=a, b=b).cdf(300) + np.testing.assert_allclose(res, 1, atol=1e-2) + res = stats.norminvgauss(a=a, b=b).sf(-300) + np.testing.assert_allclose(res, 1, atol=1e-2) + + @pytest.mark.parametrize("x, a, b, ref", [ + (100, 1.0, 0.0, 1.0), # originally reported in gh-23196 + (1500, 0.1, 0.09, 0.9999999999418577), # suggested in gh-24567 + ]) + def test_gh23196_tails_cdf_sf_accuracy(self, x, a, b, ref): + # Spot check the accuracy of the CDF/SF in the deep tails. + # library(GeneralizedHyperbolic) + # options(digits=16) + # pnig(1500, alpha=0.1, beta=0.09) + rtol = 1e-14 + np.testing.assert_allclose(stats.norminvgauss(a=a, b=b).cdf(x), ref, rtol=rtol) + np.testing.assert_allclose(stats.norminvgauss(a=a, b=-b).sf(-x), ref, rtol=rtol) + + +class TestGeom: + def setup_method(self): + self.rng = np.random.default_rng(7672986002) + + def test_rvs(self): + vals = stats.geom.rvs(0.75, size=(2, 50), random_state=self.rng) + assert_(np.all(vals >= 0)) + assert_(np.shape(vals) == (2, 50)) + assert_(vals.dtype.char in typecodes['AllInteger']) + val = stats.geom.rvs(0.75, random_state=self.rng) + assert_(isinstance(val, int)) + val = stats.geom(0.75).rvs(3, random_state=self.rng) + assert_(isinstance(val, np.ndarray)) + assert_(val.dtype.char in typecodes['AllInteger']) + + def test_rvs_9313(self): + # previously, RVS were converted to `np.int32` on some platforms, + # causing overflow for moderately large integer output (gh-9313). + # Check that this is resolved to the extent possible w/ `np.int64`. + rvs = stats.geom.rvs(np.exp(-35), size=5, random_state=self.rng) + assert rvs.dtype == np.int64 + assert np.all(rvs > np.iinfo(np.int32).max) + + def test_pmf(self): + vals = stats.geom.pmf([1, 2, 3], 0.5) + assert_array_almost_equal(vals, [0.5, 0.25, 0.125]) + + def test_logpmf(self): + # regression test for ticket 1793 + vals1 = np.log(stats.geom.pmf([1, 2, 3], 0.5)) + vals2 = stats.geom.logpmf([1, 2, 3], 0.5) + assert_allclose(vals1, vals2, rtol=1e-15, atol=0) + + # regression test for gh-4028 + val = stats.geom.logpmf(1, 1) + assert_equal(val, 0.0) + + def test_cdf_sf(self): + vals = stats.geom.cdf([1, 2, 3], 0.5) + vals_sf = stats.geom.sf([1, 2, 3], 0.5) + expected = array([0.5, 0.75, 0.875]) + assert_array_almost_equal(vals, expected) + assert_array_almost_equal(vals_sf, 1-expected) + + def test_logcdf_logsf(self): + vals = stats.geom.logcdf([1, 2, 3], 0.5) + vals_sf = stats.geom.logsf([1, 2, 3], 0.5) + expected = array([0.5, 0.75, 0.875]) + assert_array_almost_equal(vals, np.log(expected)) + assert_array_almost_equal(vals_sf, np.log1p(-expected)) + + def test_ppf(self): + vals = stats.geom.ppf([0.5, 0.75, 0.875], 0.5) + expected = array([1.0, 2.0, 3.0]) + assert_array_almost_equal(vals, expected) + + def test_ppf_underflow(self): + # this should not underflow + assert_allclose(stats.geom.ppf(1e-20, 1e-20), 1.0, atol=1e-14) + + def test_entropy_gh18226(self): + # gh-18226 reported that `geom.entropy` produced a warning and + # inaccurate output for small p. Check that this is resolved. + h = stats.geom(0.0146).entropy() + assert_allclose(h, 5.219397961962308, rtol=1e-15) + + def test_rvs_gh18372(self): + # gh-18372 reported that `geom.rvs` could produce negative numbers, + # with `RandomState` PRNG, but the support is positive integers. + # Check that this is resolved. + random_state = np.random.RandomState(294582935) + assert (stats.geom.rvs(1e-30, size=10, random_state=random_state) > 0).all() + + +class TestPlanck: + + def test_sf(self): + vals = stats.planck.sf([1, 2, 3], 5.) + expected = array([4.5399929762484854e-05, + 3.0590232050182579e-07, + 2.0611536224385579e-09]) + assert_array_almost_equal(vals, expected) + + def test_logsf(self): + vals = stats.planck.logsf([1000., 2000., 3000.], 1000.) + expected = array([-1001000., -2001000., -3001000.]) + assert_array_almost_equal(vals, expected) + + +class TestGennorm: + def test_laplace(self): + # test against Laplace (special case for beta=1) + points = [1, 2, 3] + pdf1 = stats.gennorm.pdf(points, 1) + pdf2 = stats.laplace.pdf(points) + assert_almost_equal(pdf1, pdf2) + + def test_norm(self): + # test against normal (special case for beta=2) + points = [1, 2, 3] + pdf1 = stats.gennorm.pdf(points, 2) + pdf2 = stats.norm.pdf(points, scale=2**-.5) + assert_almost_equal(pdf1, pdf2) + + def test_rvs(self): + # 0 < beta < 1 + dist = stats.gennorm(0.5) + rng = np.random.default_rng(2204049394) + rvs = dist.rvs(size=1000, random_state=rng) + assert stats.kstest(rvs, dist.cdf).pvalue > 0.1 + # beta = 1 + dist = stats.gennorm(1) + rvs = dist.rvs(size=1000, random_state=rng) + rvs_laplace = stats.laplace.rvs(size=1000, random_state=rng) + assert stats.ks_2samp(rvs, rvs_laplace).pvalue > 0.1 + # beta = 2 + dist = stats.gennorm(2) + dist.random_state = rng + rvs = dist.rvs(size=1000, random_state=rng) + rvs_norm = stats.norm.rvs(scale=1/2**0.5, size=1000, random_state=rng) + assert stats.ks_2samp(rvs, rvs_norm).pvalue > 0.1 + + def test_rvs_broadcasting(self): + dist = stats.gennorm([[0.5, 1.], [2., 5.]]) + rng = np.random.default_rng(2204049394) + rvs = dist.rvs(size=[1000, 2, 2], random_state=rng) + assert stats.kstest(rvs[:, 0, 0], stats.gennorm(0.5).cdf)[1] > 0.1 + assert stats.kstest(rvs[:, 0, 1], stats.gennorm(1.0).cdf)[1] > 0.1 + assert stats.kstest(rvs[:, 1, 0], stats.gennorm(2.0).cdf)[1] > 0.1 + assert stats.kstest(rvs[:, 1, 1], stats.gennorm(5.0).cdf)[1] > 0.1 + + +class TestGibrat: + + # sfx is sf(x). The values were computed with mpmath: + # + # from mpmath import mp + # mp.dps = 100 + # def gibrat_sf(x): + # return 1 - mp.ncdf(mp.log(x)) + # + # E.g. + # + # >>> float(gibrat_sf(1.5)) + # 0.3425678305148459 + # + @pytest.mark.parametrize('x, sfx', [(1.5, 0.3425678305148459), + (5000, 8.173334352522493e-18)]) + def test_sf_isf(self, x, sfx): + assert_allclose(stats.gibrat.sf(x), sfx, rtol=2e-14) + assert_allclose(stats.gibrat.isf(sfx), x, rtol=2e-14) + + +class TestGompertz: + + def test_gompertz_accuracy(self): + # Regression test for gh-4031 + p = stats.gompertz.ppf(stats.gompertz.cdf(1e-100, 1), 1) + assert_allclose(p, 1e-100) + + # sfx is sf(x). The values were computed with mpmath: + # + # from mpmath import mp + # mp.dps = 100 + # def gompertz_sf(x, c): + # return mp.exp(-c*mp.expm1(x)) + # + # E.g. + # + # >>> float(gompertz_sf(1, 2.5)) + # 0.013626967146253437 + # + @pytest.mark.parametrize('x, c, sfx', [(1, 2.5, 0.013626967146253437), + (3, 2.5, 1.8973243273704087e-21), + (0.05, 5, 0.7738668242570479), + (2.25, 5, 3.707795833465481e-19)]) + def test_sf_isf(self, x, c, sfx): + assert_allclose(stats.gompertz.sf(x, c), sfx, rtol=1e-14) + assert_allclose(stats.gompertz.isf(sfx, c), x, rtol=1e-14) + + def test_logcdf(self): + x = 8.0 + c = 0.1 + # Reference value computed with mpmath. + ref = -3.820049516821143e-130 + logcdf = stats.gompertz.logcdf(x, c) + assert_allclose(logcdf, ref, rtol=5e-15) + + def test_logsf(self): + x = 3e-80 + c = 12 + # Reference value computed with mpmath. + ref = -3.6e-79 + logsf = stats.gompertz.logsf(x, c) + assert_allclose(logsf, ref, rtol=5e-15) + + # reference values were computed with mpmath + # from mpmath import mp + # mp.dps = 100 + # def gompertz_entropy(c): + # c = mp.mpf(c) + # return float(mp.one - mp.log(c) - mp.exp(c)*mp.e1(c)) + + @pytest.mark.parametrize('c, ref', [(1e-4, 1.5762523017634573), + (1, 0.4036526376768059), + (1000, -5.908754280976161), + (1e10, -22.025850930040455)]) + def test_entropy(self, c, ref): + assert_allclose(stats.gompertz.entropy(c), ref, rtol=1e-14) + + +class TestFoldNorm: + + # reference values were computed with mpmath with 50 digits of precision + # from mpmath import mp + # mp.dps = 50 + # mp.mpf(0.5) * (mp.erf((x - c)/mp.sqrt(2)) + mp.erf((x + c)/mp.sqrt(2))) + + @pytest.mark.parametrize('x, c, ref', [(1e-4, 1e-8, 7.978845594730578e-05), + (1e-4, 1e-4, 7.97884555483635e-05)]) + def test_cdf(self, x, c, ref): + assert_allclose(stats.foldnorm.cdf(x, c), ref, rtol=1e-15) + + +class TestHalfNorm: + + # sfx is sf(x). The values were computed with mpmath: + # + # from mpmath import mp + # mp.dps = 100 + # def halfnorm_sf(x): + # return 2*(1 - mp.ncdf(x)) + # + # E.g. + # + # >>> float(halfnorm_sf(1)) + # 0.3173105078629141 + # + @pytest.mark.parametrize('x, sfx', [(1, 0.3173105078629141), + (10, 1.523970604832105e-23)]) + def test_sf_isf(self, x, sfx): + assert_allclose(stats.halfnorm.sf(x), sfx, rtol=1e-14) + assert_allclose(stats.halfnorm.isf(sfx), x, rtol=1e-14) + + # reference values were computed via mpmath + # from mpmath import mp + # mp.dps = 100 + # def halfnorm_cdf_mpmath(x): + # x = mp.mpf(x) + # return float(mp.erf(x/mp.sqrt(2.))) + + @pytest.mark.parametrize('x, ref', [(1e-40, 7.978845608028653e-41), + (1e-18, 7.978845608028654e-19), + (8, 0.9999999999999988)]) + def test_cdf(self, x, ref): + assert_allclose(stats.halfnorm.cdf(x), ref, rtol=1e-15) + + @pytest.mark.parametrize("rvs_loc", [1e-5, 1e10]) + @pytest.mark.parametrize("rvs_scale", [1e-2, 100, 1e8]) + @pytest.mark.parametrize('fix_loc', [True, False]) + @pytest.mark.parametrize('fix_scale', [True, False]) + def test_fit_MLE_comp_optimizer(self, rvs_loc, rvs_scale, + fix_loc, fix_scale): + + rng = np.random.default_rng(6762668991392531563) + data = stats.halfnorm.rvs(loc=rvs_loc, scale=rvs_scale, size=1000, + random_state=rng) + + if fix_loc and fix_scale: + error_msg = ("All parameters fixed. There is nothing to " + "optimize.") + with pytest.raises(RuntimeError, match=error_msg): + stats.halflogistic.fit(data, floc=rvs_loc, fscale=rvs_scale) + return + + kwds = {} + if fix_loc: + kwds['floc'] = rvs_loc + if fix_scale: + kwds['fscale'] = rvs_scale + + # Numerical result may equal analytical result if the initial guess + # computed from moment condition is already optimal. + _assert_less_or_close_loglike(stats.halfnorm, data, **kwds, + maybe_identical=True) + + def test_fit_error(self): + # `floc` bigger than the minimal data point + with pytest.raises(FitDataError): + stats.halfnorm.fit([1, 2, 3], floc=2) + + +class TestHalfCauchy: + + @pytest.mark.parametrize("rvs_loc", [1e-5, 1e10]) + @pytest.mark.parametrize("rvs_scale", [1e-2, 1e8]) + @pytest.mark.parametrize('fix_loc', [True, False]) + @pytest.mark.parametrize('fix_scale', [True, False]) + def test_fit_MLE_comp_optimizer(self, rvs_loc, rvs_scale, + fix_loc, fix_scale): + + rng = np.random.default_rng(6762668991392531563) + data = stats.halfnorm.rvs(loc=rvs_loc, scale=rvs_scale, size=1000, + random_state=rng) + + if fix_loc and fix_scale: + error_msg = ("All parameters fixed. There is nothing to " + "optimize.") + with pytest.raises(RuntimeError, match=error_msg): + stats.halfcauchy.fit(data, floc=rvs_loc, fscale=rvs_scale) + return + + kwds = {} + if fix_loc: + kwds['floc'] = rvs_loc + if fix_scale: + kwds['fscale'] = rvs_scale + + _assert_less_or_close_loglike(stats.halfcauchy, data, **kwds) + + def test_fit_error(self): + # `floc` bigger than the minimal data point + with pytest.raises(FitDataError): + stats.halfcauchy.fit([1, 2, 3], floc=2) + + +class TestHalfLogistic: + # survival function reference values were computed with mpmath + # from mpmath import mp + # mp.dps = 50 + # def sf_mpmath(x): + # x = mp.mpf(x) + # return float(mp.mpf(2.)/(mp.exp(x) + mp.one)) + + @pytest.mark.parametrize('x, ref', [(100, 7.440151952041672e-44), + (200, 2.767793053473475e-87)]) + def test_sf(self, x, ref): + assert_allclose(stats.halflogistic.sf(x), ref, rtol=1e-15) + + # inverse survival function reference values were computed with mpmath + # from mpmath import mp + # mp.dps = 200 + # def isf_mpmath(x): + # halfx = mp.mpf(x)/2 + # return float(-mp.log(halfx/(mp.one - halfx))) + + @pytest.mark.parametrize('q, ref', [(7.440151952041672e-44, 100), + (2.767793053473475e-87, 200), + (1-1e-9, 1.999999943436137e-09), + (1-1e-15, 1.9984014443252818e-15)]) + def test_isf(self, q, ref): + assert_allclose(stats.halflogistic.isf(q), ref, rtol=1e-15) + + def test_logcdf(self): + x = 30.0 + # Reference value computed with mpmath. + ref = -1.871524593768035e-13 + logcdf = stats.halflogistic.logcdf(x) + assert_allclose(logcdf, ref, rtol=5e-15) + + def test_logsf(self): + x = 2e-14 + # Reference value computed with mpmath. + ref = -1.000000000000005e-14 + logsf = stats.halflogistic.logsf(x) + assert_allclose(logsf, ref, rtol=5e-15) + + @pytest.mark.parametrize("rvs_loc", [1e-5, 1e10]) + @pytest.mark.parametrize("rvs_scale", [1e-2, 100, 1e8]) + @pytest.mark.parametrize('fix_loc', [True, False]) + @pytest.mark.parametrize('fix_scale', [True, False]) + def test_fit_MLE_comp_optimizer(self, rvs_loc, rvs_scale, + fix_loc, fix_scale): + + rng = np.random.default_rng(6762668991392531563) + data = stats.halflogistic.rvs(loc=rvs_loc, scale=rvs_scale, size=1000, + random_state=rng) + + kwds = {} + if fix_loc and fix_scale: + error_msg = ("All parameters fixed. There is nothing to " + "optimize.") + with pytest.raises(RuntimeError, match=error_msg): + stats.halflogistic.fit(data, floc=rvs_loc, fscale=rvs_scale) + return + + if fix_loc: + kwds['floc'] = rvs_loc + if fix_scale: + kwds['fscale'] = rvs_scale + + # Numerical result may equal analytical result if the initial guess + # computed from moment condition is already optimal. + _assert_less_or_close_loglike(stats.halflogistic, data, **kwds, + maybe_identical=True) + + def test_fit_bad_floc(self): + msg = r" Maximum likelihood estimation with 'halflogistic' requires" + with assert_raises(FitDataError, match=msg): + stats.halflogistic.fit([0, 2, 4], floc=1) + + +class TestHalfgennorm: + def test_expon(self): + # test against exponential (special case for beta=1) + points = [1, 2, 3] + pdf1 = stats.halfgennorm.pdf(points, 1) + pdf2 = stats.expon.pdf(points) + assert_almost_equal(pdf1, pdf2) + + def test_halfnorm(self): + # test against half normal (special case for beta=2) + points = [1, 2, 3] + pdf1 = stats.halfgennorm.pdf(points, 2) + pdf2 = stats.halfnorm.pdf(points, scale=2**-.5) + assert_almost_equal(pdf1, pdf2) + + def test_gennorm(self): + # test against generalized normal + points = [1, 2, 3] + pdf1 = stats.halfgennorm.pdf(points, .497324) + pdf2 = stats.gennorm.pdf(points, .497324) + assert_almost_equal(pdf1, 2*pdf2) + + # Reference values (e.g. for n=3, beta=1.25) + # from mpmath import mp + # mp.dps = 100 + # n = 3 + # beta = mp.mpf(1.25) + # ref = float(mp.rf(1/beta, n/beta)) + @pytest.mark.parametrize('n, beta, ref', + [(0, 2.25, 1), + (1, 0.5, 6.0), + (2, 12.5, 0.3155489903298712), + (3, 0.125, 1.631515605987683e+30)]) + def test_munp(self, n, beta, ref): + munp = stats.halfgennorm._munp(n, beta) + assert_allclose(munp, ref, rtol=5e-15) + + # Reference values (e.g. for beta=0.02) + # from mpmath import mp + # mp.dps = 100 + # beta = mp.mpf(0.02) + # ref = float(mp.rf(1/beta, 1/beta)) + @pytest.mark.parametrize('beta, ref', + [(0.02, 1.5342593781274747e+93), + (0.125, 259459200.0), + (0.5, 6.0), + (10.0, 0.48256057149575027), + (125.0, 0.49777435261046765), + (15000.0, 0.49998076533051666)]) + def test_mean(self, beta, ref): + m = stats.halfgennorm.mean(beta) + assert_allclose(m, ref, rtol=1e-14) + + # Reference values (e.g. for beta=0.02) + # from mpmath import mp + # mp.dps = 100 + # beta = mp.mpf(0.02) + # mu = mp.rf(1/beta, 1/beta) + # # var = E(X**2) - E(X)**2: + # ref = float(mp.rf(1/beta, 2/beta) - mu**2) + @pytest.mark.parametrize('beta, ref', + [(0.02, 6.261772482175519e+197), + (0.125, 5.062049324107776e+18), + (0.5, 84.0), + (10.0, 0.08159018176717249), + (125.0, 0.0826270885971827), + (15000.0, 0.08332692433644018)]) + def test_var(self, beta, ref): + v = stats.halfgennorm.var(beta) + assert_allclose(v, ref, rtol=1.5e-14) + + +class TestLaplaceasymmetric: + def test_laplace(self): + # test against Laplace (special case for kappa=1) + points = np.array([1, 2, 3]) + pdf1 = stats.laplace_asymmetric.pdf(points, 1) + pdf2 = stats.laplace.pdf(points) + assert_allclose(pdf1, pdf2) + + def test_asymmetric_laplace_pdf(self): + # test asymmetric Laplace + points = np.array([1, 2, 3]) + kappa = 2 + kapinv = 1/kappa + pdf1 = stats.laplace_asymmetric.pdf(points, kappa) + pdf2 = stats.laplace_asymmetric.pdf(points*(kappa**2), kapinv) + assert_allclose(pdf1, pdf2) + + def test_asymmetric_laplace_log_10_16(self): + # test asymmetric Laplace + points = np.array([-np.log(16), np.log(10)]) + kappa = 2 + pdf1 = stats.laplace_asymmetric.pdf(points, kappa) + cdf1 = stats.laplace_asymmetric.cdf(points, kappa) + sf1 = stats.laplace_asymmetric.sf(points, kappa) + pdf2 = np.array([1/10, 1/250]) + cdf2 = np.array([1/5, 1 - 1/500]) + sf2 = np.array([4/5, 1/500]) + ppf1 = stats.laplace_asymmetric.ppf(cdf2, kappa) + ppf2 = points + isf1 = stats.laplace_asymmetric.isf(sf2, kappa) + isf2 = points + assert_allclose(np.concatenate((pdf1, cdf1, sf1, ppf1, isf1)), + np.concatenate((pdf2, cdf2, sf2, ppf2, isf2))) + + +class TestTruncnorm: + def setup_method(self): + self.rng = np.random.default_rng(3255963201) + + @pytest.mark.parametrize("a, b, ref", + [(0, 100, 0.7257913526447274), + (0.6, 0.7, -2.3027610681852573), + (1e-06, 2e-06, -13.815510557964274)]) + def test_entropy(self, a, b, ref): + # All reference values were calculated with mpmath: + # import numpy as np + # from mpmath import mp + # mp.dps = 50 + # def entropy_trun(a, b): + # a, b = mp.mpf(a), mp.mpf(b) + # Z = mp.ncdf(b) - mp.ncdf(a) + # + # def pdf(x): + # return mp.npdf(x) / Z + # + # res = -mp.quad(lambda t: pdf(t) * mp.log(pdf(t)), [a, b]) + # return np.float64(res) + assert_allclose(stats.truncnorm.entropy(a, b), ref, rtol=1e-10) + + @pytest.mark.parametrize("a, b, ref", + [(1e-11, 10000000000.0, 0.725791352640738), + (1e-100, 1e+100, 0.7257913526447274), + (-1e-100, 1e+100, 0.7257913526447274), + (-1e+100, 1e+100, 1.4189385332046727)]) + def test_extreme_entropy(self, a, b, ref): + # The reference values were calculated with mpmath + # import numpy as np + # from mpmath import mp + # mp.dps = 50 + # def trunc_norm_entropy(a, b): + # a, b = mp.mpf(a), mp.mpf(b) + # Z = mp.ncdf(b) - mp.ncdf(a) + # A = mp.log(mp.sqrt(2 * mp.pi * mp.e) * Z) + # B = (a * mp.npdf(a) - b * mp.npdf(b)) / (2 * Z) + # return np.float64(A + B) + assert_allclose(stats.truncnorm.entropy(a, b), ref, rtol=1e-14) + + def test_ppf_ticket1131(self): + vals = stats.truncnorm.ppf([-0.5, 0, 1e-4, 0.5, 1-1e-4, 1, 2], -1., 1., + loc=[3]*7, scale=2) + expected = np.array([np.nan, 1, 1.00056419, 3, 4.99943581, 5, np.nan]) + assert_array_almost_equal(vals, expected) + + def test_isf_ticket1131(self): + vals = stats.truncnorm.isf([-0.5, 0, 1e-4, 0.5, 1-1e-4, 1, 2], -1., 1., + loc=[3]*7, scale=2) + expected = np.array([np.nan, 5, 4.99943581, 3, 1.00056419, 1, np.nan]) + assert_array_almost_equal(vals, expected) + + def test_gh_2477_small_values(self): + # Check a case that worked in the original issue. + low, high = -11, -10 + x = stats.truncnorm.rvs(low, high, 0, 1, size=10, random_state=self.rng) + assert_(low < x.min() < x.max() < high) + # Check a case that failed in the original issue. + low, high = 10, 11 + x = stats.truncnorm.rvs(low, high, 0, 1, size=10, random_state=self.rng) + assert_(low < x.min() < x.max() < high) + + def test_gh_2477_large_values(self): + # Check a case that used to fail because of extreme tailness. + low, high = 100, 101 + x = stats.truncnorm.rvs(low, high, 0, 1, size=10, random_state=self.rng) + assert_(low <= x.min() <= x.max() <= high), str([low, high, x]) + + # Check some additional extreme tails + low, high = 1000, 1001 + x = stats.truncnorm.rvs(low, high, 0, 1, size=10, random_state=self.rng) + assert_(low < x.min() < x.max() < high) + + low, high = 10000, 10001 + x = stats.truncnorm.rvs(low, high, 0, 1, size=10, random_state=self.rng) + assert_(low < x.min() < x.max() < high) + + low, high = -10001, -10000 + x = stats.truncnorm.rvs(low, high, 0, 1, size=10, random_state=self.rng) + assert_(low < x.min() < x.max() < high) + + def test_gh_9403_nontail_values(self): + for low, high in [[3, 4], [-4, -3]]: + xvals = np.array([-np.inf, low, high, np.inf]) + xmid = (high+low)/2.0 + cdfs = stats.truncnorm.cdf(xvals, low, high) + sfs = stats.truncnorm.sf(xvals, low, high) + pdfs = stats.truncnorm.pdf(xvals, low, high) + expected_cdfs = np.array([0, 0, 1, 1]) + expected_sfs = np.array([1.0, 1.0, 0.0, 0.0]) + expected_pdfs = np.array([0, 3.3619772, 0.1015229, 0]) + if low < 0: + expected_pdfs = np.array([0, 0.1015229, 3.3619772, 0]) + assert_almost_equal(cdfs, expected_cdfs) + assert_almost_equal(sfs, expected_sfs) + assert_almost_equal(pdfs, expected_pdfs) + assert_almost_equal(np.log(expected_pdfs[1]/expected_pdfs[2]), + low + 0.5) + pvals = np.array([0, 0.5, 1.0]) + ppfs = stats.truncnorm.ppf(pvals, low, high) + expected_ppfs = np.array([low, np.sign(low)*3.1984741, high]) + assert_almost_equal(ppfs, expected_ppfs) + + if low < 0: + assert_almost_equal(stats.truncnorm.sf(xmid, low, high), + 0.8475544278436675) + assert_almost_equal(stats.truncnorm.cdf(xmid, low, high), + 0.1524455721563326) + else: + assert_almost_equal(stats.truncnorm.cdf(xmid, low, high), + 0.8475544278436675) + assert_almost_equal(stats.truncnorm.sf(xmid, low, high), + 0.1524455721563326) + pdf = stats.truncnorm.pdf(xmid, low, high) + assert_almost_equal(np.log(pdf/expected_pdfs[2]), (xmid+0.25)/2) + + def test_gh_9403_medium_tail_values(self): + for low, high in [[39, 40], [-40, -39]]: + xvals = np.array([-np.inf, low, high, np.inf]) + xmid = (high+low)/2.0 + cdfs = stats.truncnorm.cdf(xvals, low, high) + sfs = stats.truncnorm.sf(xvals, low, high) + pdfs = stats.truncnorm.pdf(xvals, low, high) + expected_cdfs = np.array([0, 0, 1, 1]) + expected_sfs = np.array([1.0, 1.0, 0.0, 0.0]) + expected_pdfs = np.array([0, 3.90256074e+01, 2.73349092e-16, 0]) + if low < 0: + expected_pdfs = np.array([0, 2.73349092e-16, + 3.90256074e+01, 0]) + assert_almost_equal(cdfs, expected_cdfs) + assert_almost_equal(sfs, expected_sfs) + assert_almost_equal(pdfs, expected_pdfs) + assert_almost_equal(np.log(expected_pdfs[1]/expected_pdfs[2]), + low + 0.5) + pvals = np.array([0, 0.5, 1.0]) + ppfs = stats.truncnorm.ppf(pvals, low, high) + expected_ppfs = np.array([low, np.sign(low)*39.01775731, high]) + assert_almost_equal(ppfs, expected_ppfs) + cdfs = stats.truncnorm.cdf(ppfs, low, high) + assert_almost_equal(cdfs, pvals) + + if low < 0: + assert_almost_equal(stats.truncnorm.sf(xmid, low, high), + 0.9999999970389126) + assert_almost_equal(stats.truncnorm.cdf(xmid, low, high), + 2.961048103554866e-09) + else: + assert_almost_equal(stats.truncnorm.cdf(xmid, low, high), + 0.9999999970389126) + assert_almost_equal(stats.truncnorm.sf(xmid, low, high), + 2.961048103554866e-09) + pdf = stats.truncnorm.pdf(xmid, low, high) + assert_almost_equal(np.log(pdf/expected_pdfs[2]), (xmid+0.25)/2) + + xvals = np.linspace(low, high, 11) + xvals2 = -xvals[::-1] + assert_almost_equal(stats.truncnorm.cdf(xvals, low, high), + stats.truncnorm.sf(xvals2, -high, -low)[::-1]) + assert_almost_equal(stats.truncnorm.sf(xvals, low, high), + stats.truncnorm.cdf(xvals2, -high, -low)[::-1]) + assert_almost_equal(stats.truncnorm.pdf(xvals, low, high), + stats.truncnorm.pdf(xvals2, -high, -low)[::-1]) + + def test_cdf_tail_15110_14753(self): + # Check accuracy issues reported in gh-14753 and gh-155110 + # Ground truth values calculated using Wolfram Alpha, e.g. + # (CDF[NormalDistribution[0,1],83/10]-CDF[NormalDistribution[0,1],8])/ + # (1 - CDF[NormalDistribution[0,1],8]) + assert_allclose(stats.truncnorm(13., 15.).cdf(14.), + 0.9999987259565643) + assert_allclose(stats.truncnorm(8, np.inf).cdf(8.3), + 0.9163220907327540) + + # Test data for the truncnorm stats() method. + # The data in each row is: + # a, b, mean, variance, skewness, excess kurtosis. Generated using + # https://gist.github.com/WarrenWeckesser/636b537ee889679227d53543d333a720 + _truncnorm_stats_data = [ + [-30, 30, + 0.0, 1.0, 0.0, 0.0], + [-10, 10, + 0.0, 1.0, 0.0, -1.4927521335810455e-19], + [-3, 3, + 0.0, 0.9733369246625415, 0.0, -0.17111443639774404], + [-2, 2, + 0.0, 0.7737413035499232, 0.0, -0.6344632828703505], + [0, np.inf, + 0.7978845608028654, + 0.3633802276324187, + 0.995271746431156, + 0.8691773036059741], + [-np.inf, 0, + -0.7978845608028654, + 0.3633802276324187, + -0.995271746431156, + 0.8691773036059741], + [-1, 3, + 0.282786110727154, + 0.6161417353578293, + 0.5393018494027877, + -0.20582065135274694], + [-3, 1, + -0.282786110727154, + 0.6161417353578293, + -0.5393018494027877, + -0.20582065135274694], + [-10, -9, + -9.108456288012409, + 0.011448805821636248, + -1.8985607290949496, + 5.0733461105025075], + ] + _truncnorm_stats_data = np.array(_truncnorm_stats_data) + + @pytest.mark.parametrize("case", _truncnorm_stats_data) + def test_moments(self, case): + a, b, m0, v0, s0, k0 = case + m, v, s, k = stats.truncnorm.stats(a, b, moments='mvsk') + assert_allclose([m, v, s, k], [m0, v0, s0, k0], atol=1e-17) + + def test_9902_moments(self): + m, v = stats.truncnorm.stats(0, np.inf, moments='mv') + assert_almost_equal(m, 0.79788456) + assert_almost_equal(v, 0.36338023) + + def test_gh_1489_trac_962_rvs(self): + # Check the original example. + low, high = 10, 15 + x = stats.truncnorm.rvs(low, high, 0, 1, size=10, random_state=self.rng) + assert_(low < x.min() < x.max() < high) + + def test_gh_11299_rvs(self): + # Arose from investigating gh-11299 + # Test multiple shape parameters simultaneously. + low = [-10, 10, -np.inf, -5, -np.inf, -np.inf, -45, -45, 40, -10, 40] + high = [-5, 11, 5, np.inf, 40, -40, 40, -40, 45, np.inf, np.inf] + x = stats.truncnorm.rvs(low, high, size=(5, len(low)), random_state=self.rng) + assert np.shape(x) == (5, len(low)) + assert_(np.all(low <= x.min(axis=0))) + assert_(np.all(x.max(axis=0) <= high)) + + def test_rvs_Generator(self): + # check that rvs can use a Generator + if hasattr(np.random, "default_rng"): + stats.truncnorm.rvs(-10, -5, size=5, random_state=self.rng) + + def test_logcdf_gh17064(self): + # regression test for gh-17064 - avoid roundoff error for logcdfs ~0 + a = np.array([-np.inf, -np.inf, -8, -np.inf, 10]) + b = np.array([np.inf, np.inf, 8, 10, np.inf]) + x = np.array([10, 7.5, 7.5, 9, 20]) + expected = [-7.619853024160525e-24, -3.190891672910947e-14, + -3.128682067168231e-14, -1.1285122074235991e-19, + -3.61374964828753e-66] + assert_allclose(stats.truncnorm(a, b).logcdf(x), expected) + assert_allclose(stats.truncnorm(-b, -a).logsf(-x), expected) + + def test_moments_gh18634(self): + # gh-18634 reported that moments 5 and higher didn't work; check that + # this is resolved + res = stats.truncnorm(-2, 3).moment(5) + # From Mathematica: + # Moment[TruncatedDistribution[{-2, 3}, NormalDistribution[]], 5] + ref = 1.645309620208361 + assert_allclose(res, ref) + + +class TestGenLogistic: + + # Expected values computed with mpmath with 50 digits of precision. + @pytest.mark.parametrize('x, expected', [(-1000, -1499.5945348918917), + (-125, -187.09453489189184), + (0, -1.3274028432916989), + (100, -99.59453489189184), + (1000, -999.5945348918918)]) + def test_logpdf(self, x, expected): + c = 1.5 + logp = stats.genlogistic.logpdf(x, c) + assert_allclose(logp, expected, rtol=1e-13) + + # Expected values computed with mpmath with 50 digits of precision + # from mpmath import mp + # mp.dps = 50 + # def entropy_mp(c): + # c = mp.mpf(c) + # return float(-mp.log(c)+mp.one+mp.digamma(c + mp.one) + mp.euler) + + @pytest.mark.parametrize('c, ref', [(1e-100, 231.25850929940458), + (1e-4, 10.21050485336338), + (1e8, 1.577215669901533), + (1e100, 1.5772156649015328)]) + def test_entropy(self, c, ref): + assert_allclose(stats.genlogistic.entropy(c), ref, rtol=5e-15) + + # Expected values computed with mpmath with 50 digits of precision + # from mpmath import mp + # mp.dps = 1000 + # + # def genlogistic_cdf_mp(x, c): + # x = mp.mpf(x) + # c = mp.mpf(c) + # return (mp.one + mp.exp(-x)) ** (-c) + # + # def genlogistic_sf_mp(x, c): + # return mp.one - genlogistic_cdf_mp(x, c) + # + # x, c, ref = 100, 0.02, -7.440151952041672e-466 + # print(float(mp.log(genlogistic_cdf_mp(x, c)))) + # ppf/isf reference values generated by passing in `ref` (`q` is produced) + + @pytest.mark.parametrize('x, c, ref', [(200, 10, 1.3838965267367375e-86), + (500, 20, 1.424915281348257e-216)]) + def test_sf(self, x, c, ref): + assert_allclose(stats.genlogistic.sf(x, c), ref, rtol=1e-14) + + @pytest.mark.parametrize('q, c, ref', [(0.01, 200, 9.898441467379765), + (0.001, 2, 7.600152115573173)]) + def test_isf(self, q, c, ref): + assert_allclose(stats.genlogistic.isf(q, c), ref, rtol=5e-16) + + @pytest.mark.parametrize('q, c, ref', [(0.5, 200, 5.6630969187064615), + (0.99, 20, 7.595630231412436)]) + def test_ppf(self, q, c, ref): + assert_allclose(stats.genlogistic.ppf(q, c), ref, rtol=5e-16) + + @pytest.mark.parametrize('x, c, ref', [(100, 0.02, -7.440151952041672e-46), + (50, 20, -3.857499695927835e-21)]) + def test_logcdf(self, x, c, ref): + assert_allclose(stats.genlogistic.logcdf(x, c), ref, rtol=1e-15) + + +class TestHypergeom: + def setup_method(self): + self.rng = np.random.default_rng(1765545342) + + def test_rvs(self): + vals = stats.hypergeom.rvs(20, 10, 3, size=(2, 50), random_state=self.rng) + assert np.all(vals >= 0) & np.all(vals <= 3) + assert np.shape(vals) == (2, 50) + assert vals.dtype.char in typecodes['AllInteger'] + val = stats.hypergeom.rvs(20, 3, 10, random_state=self.rng) + assert isinstance(val, int) + val = stats.hypergeom(20, 3, 10).rvs(3, random_state=self.rng) + assert isinstance(val, np.ndarray) + assert val.dtype.char in typecodes['AllInteger'] + + def test_precision(self): + # comparison number from mpmath + M = 2500 + n = 50 + N = 500 + tot = M + good = n + hgpmf = stats.hypergeom.pmf(2, tot, good, N) + assert_almost_equal(hgpmf, 0.0010114963068932233, 11) + + def test_args(self): + # test correct output for corner cases of arguments + # see gh-2325 + assert_almost_equal(stats.hypergeom.pmf(0, 2, 1, 0), 1.0, 11) + assert_almost_equal(stats.hypergeom.pmf(1, 2, 1, 0), 0.0, 11) + + assert_almost_equal(stats.hypergeom.pmf(0, 2, 0, 2), 1.0, 11) + assert_almost_equal(stats.hypergeom.pmf(1, 2, 1, 0), 0.0, 11) + + def test_cdf_above_one(self): + # for some values of parameters, hypergeom cdf was >1, see gh-2238 + assert_(0 <= stats.hypergeom.cdf(30, 13397950, 4363, 12390) <= 1.0) + + def test_precision2(self): + # Test hypergeom precision for large numbers. See #1218. + # Results compared with those from R. + oranges = 9.9e4 + pears = 1.1e5 + fruits_eaten = np.array([3, 3.8, 3.9, 4, 4.1, 4.2, 5]) * 1e4 + quantile = 2e4 + res = [stats.hypergeom.sf(quantile, oranges + pears, oranges, eaten) + for eaten in fruits_eaten] + expected = np.array([0, 1.904153e-114, 2.752693e-66, 4.931217e-32, + 8.265601e-11, 0.1237904, 1]) + assert_allclose(res, expected, atol=0, rtol=5e-7) + + # Test with array_like first argument + quantiles = [1.9e4, 2e4, 2.1e4, 2.15e4] + res2 = stats.hypergeom.sf(quantiles, oranges + pears, oranges, 4.2e4) + expected2 = [1, 0.1237904, 6.511452e-34, 3.277667e-69] + assert_allclose(res2, expected2, atol=0, rtol=5e-7) + + def test_entropy(self): + # Simple tests of entropy. + hg = stats.hypergeom(4, 1, 1) + h = hg.entropy() + expected_p = np.array([0.75, 0.25]) + expected_h = -np.sum(xlogy(expected_p, expected_p)) + assert_allclose(h, expected_h) + + hg = stats.hypergeom(1, 1, 1) + h = hg.entropy() + assert_equal(h, 0.0) + + def test_logsf(self): + # Test logsf for very large numbers. See issue #4982 + # Results compare with those from R (v3.2.0): + # phyper(k, n, M-n, N, lower.tail=FALSE, log.p=TRUE) + # -2239.771 + + k = 1e4 + M = 1e7 + n = 1e6 + N = 5e4 + + result = stats.hypergeom.logsf(k, M, n, N) + expected = -2239.771 # From R + assert_almost_equal(result, expected, decimal=3) + + k = 1 + M = 1600 + n = 600 + N = 300 + + result = stats.hypergeom.logsf(k, M, n, N) + expected = -2.566567e-68 # From R + assert_almost_equal(result, expected, decimal=15) + + def test_logcdf(self): + # Test logcdf for very large numbers. See issue #8692 + # Results compare with those from R (v3.3.2): + # phyper(k, n, M-n, N, lower.tail=TRUE, log.p=TRUE) + # -5273.335 + + k = 1 + M = 1e7 + n = 1e6 + N = 5e4 + + result = stats.hypergeom.logcdf(k, M, n, N) + expected = -5273.335 # From R + assert_almost_equal(result, expected, decimal=3) + + # Same example as in issue #8692 + k = 40 + M = 1600 + n = 50 + N = 300 + + result = stats.hypergeom.logcdf(k, M, n, N) + expected = -7.565148879229e-23 # From R + assert_almost_equal(result, expected, decimal=15) + + k = 125 + M = 1600 + n = 250 + N = 500 + + result = stats.hypergeom.logcdf(k, M, n, N) + expected = -4.242688e-12 # From R + assert_almost_equal(result, expected, decimal=15) + + # test broadcasting robustness based on reviewer + # concerns in PR 9603; using an array version of + # the example from issue #8692 + k = np.array([40, 40, 40]) + M = 1600 + n = 50 + N = 300 + + result = stats.hypergeom.logcdf(k, M, n, N) + expected = np.full(3, -7.565148879229e-23) # filled from R result + assert_almost_equal(result, expected, decimal=15) + + def test_mean_gh18511(self): + # gh-18511 reported that the `mean` was incorrect for large arguments; + # check that this is resolved + M = 390_000 + n = 370_000 + N = 12_000 + + hm = stats.hypergeom.mean(M, n, N) + rm = n / M * N + assert_allclose(hm, rm) + + @pytest.mark.xslow + def test_sf_gh18506(self): + # gh-18506 reported that `sf` was incorrect for large population; + # check that this is resolved + n = 10 + N = 10**5 + i = np.arange(5, 15) + population_size = 10.**i + p = stats.hypergeom.sf(n - 1, population_size, N, n) + assert np.all(p > 0) + assert np.all(np.diff(p) < 0) + + +class TestLoggamma: + def setup_method(self): + self.rng = np.random.default_rng(8638464332) + + # Expected cdf values were computed with mpmath. For given x and c, + # x = mpmath.mpf(x) + # c = mpmath.mpf(c) + # cdf = mpmath.gammainc(c, 0, mpmath.exp(x), + # regularized=True) + @pytest.mark.parametrize('x, c, cdf', + [(1, 2, 0.7546378854206702), + (-1, 14, 6.768116452566383e-18), + (-745.1, 0.001, 0.4749605142005238), + (-800, 0.001, 0.44958802911019136), + (-725, 0.1, 3.4301205868273265e-32), + (-740, 0.75, 1.0074360436599631e-241)]) + def test_cdf_ppf(self, x, c, cdf): + p = stats.loggamma.cdf(x, c) + assert_allclose(p, cdf, rtol=1e-13) + y = stats.loggamma.ppf(cdf, c) + assert_allclose(y, x, rtol=1e-13) + + # Expected sf values were computed with mpmath. For given x and c, + # x = mpmath.mpf(x) + # c = mpmath.mpf(c) + # sf = mpmath.gammainc(c, mpmath.exp(x), mpmath.inf, + # regularized=True) + @pytest.mark.parametrize('x, c, sf', + [(4, 1.5, 1.6341528919488565e-23), + (6, 100, 8.23836829202024e-74), + (-800, 0.001, 0.5504119708898086), + (-743, 0.0025, 0.8437131370024089)]) + def test_sf_isf(self, x, c, sf): + s = stats.loggamma.sf(x, c) + assert_allclose(s, sf, rtol=1e-13) + y = stats.loggamma.isf(sf, c) + assert_allclose(y, x, rtol=1e-13) + + def test_logpdf(self): + # Test logpdf with x=-500, c=2. ln(gamma(2)) = 0, and + # exp(-500) ~= 7e-218, which is far smaller than the ULP + # of c*x=-1000, so logpdf(-500, 2) = c*x - exp(x) - ln(gamma(2)) + # should give -1000.0. + lp = stats.loggamma.logpdf(-500, 2) + assert_allclose(lp, -1000.0, rtol=1e-14) + + def test_logcdf(self): + x = 4.0 + c = 4.5 + logcdf = stats.loggamma.logcdf(x, c) + # Reference value computed with mpmath. + ref = -2.1429747073164531e-19 + assert_allclose(logcdf, ref, rtol=5e-15) + + def test_logsf(self): + x = -25.0 + c = 3.5 + logsf = stats.loggamma.logsf(x, c) + # Reference value computed with mpmath. + ref = -8.58200139319556e-40 + assert_allclose(logsf, ref, rtol=5e-15) + + def test_stats(self): + # The following precomputed values are from the table in section 2.2 + # of "A Statistical Study of Log-Gamma Distribution", by Ping Shing + # Chan (thesis, McMaster University, 1993). + table = np.array([ + # c, mean, var, skew, exc. kurt. + 0.5, -1.9635, 4.9348, -1.5351, 4.0000, + 1.0, -0.5772, 1.6449, -1.1395, 2.4000, + 12.0, 2.4427, 0.0869, -0.2946, 0.1735, + ]).reshape(-1, 5) + for c, mean, var, skew, kurt in table: + computed = stats.loggamma.stats(c, moments='msvk') + assert_array_almost_equal(computed, [mean, var, skew, kurt], + decimal=4) + + @pytest.mark.parametrize('c', [0.1, 0.001]) + def test_rvs(self, c): + # Regression test for gh-11094. + x = stats.loggamma.rvs(c, size=100000, random_state=self.rng) + # Before gh-11094 was fixed, the case with c=0.001 would + # generate many -inf values. + assert np.isfinite(x).all() + # Crude statistical test. About half the values should be + # less than the median and half greater than the median. + med = stats.loggamma.median(c) + btest = stats.binomtest(np.count_nonzero(x < med), len(x)) + ci = btest.proportion_ci(confidence_level=0.999) + assert ci.low < 0.5 < ci.high + + @pytest.mark.parametrize("c, ref", + [(1e-8, 19.420680753952364), + (1, 1.5772156649015328), + (1e4, -3.186214986116763), + (1e10, -10.093986931748889), + (1e100, -113.71031611649761)]) + def test_entropy(self, c, ref): + + # Reference values were calculated with mpmath + # from mpmath import mp + # mp.dps = 500 + # def loggamma_entropy_mpmath(c): + # c = mp.mpf(c) + # return float(mp.log(mp.gamma(c)) + c * (mp.one - mp.digamma(c))) + + assert_allclose(stats.loggamma.entropy(c), ref, rtol=1e-14) + + +class TestJohnsonsu: + # reference values were computed via mpmath + # from mpmath import mp + # mp.dps = 50 + # def johnsonsu_sf(x, a, b): + # x = mp.mpf(x) + # a = mp.mpf(a) + # b = mp.mpf(b) + # return float(mp.ncdf(-(a + b * mp.log(x + mp.sqrt(x*x + 1))))) + # Order is x, a, b, sf, isf tol + # (Can't expect full precision when the ISF input is very nearly 1) + cases = [(-500, 1, 1, 0.9999999982660072, 1e-8), + (2000, 1, 1, 7.426351000595343e-21, 5e-14), + (100000, 1, 1, 4.046923979269977e-40, 5e-14)] + + @pytest.mark.parametrize("case", cases) + def test_sf_isf(self, case): + x, a, b, sf, tol = case + assert_allclose(stats.johnsonsu.sf(x, a, b), sf, rtol=5e-14) + assert_allclose(stats.johnsonsu.isf(sf, a, b), x, rtol=tol) + + +class TestJohnsonb: + # reference values were computed via mpmath + # from mpmath import mp + # mp.dps = 50 + # def johnsonb_sf(x, a, b): + # x = mp.mpf(x) + # a = mp.mpf(a) + # b = mp.mpf(b) + # return float(mp.ncdf(-(a + b * mp.log(x/(mp.one - x))))) + # Order is x, a, b, sf, isf atol + # (Can't expect full precision when the ISF input is very nearly 1) + cases = [(1e-4, 1, 1, 0.9999999999999999, 1e-7), + (0.9999, 1, 1, 8.921114313932308e-25, 5e-14), + (0.999999, 1, 1, 5.815197487181902e-50, 5e-14)] + + @pytest.mark.parametrize("case", cases) + def test_sf_isf(self, case): + x, a, b, sf, tol = case + assert_allclose(stats.johnsonsb.sf(x, a, b), sf, rtol=5e-14) + assert_allclose(stats.johnsonsb.isf(sf, a, b), x, atol=tol) + + +class TestLogistic: + def setup_method(self): + self.rng = np.random.default_rng(2807014525) + + # gh-6226 + def test_cdf_ppf(self): + x = np.linspace(-20, 20) + y = stats.logistic.cdf(x) + xx = stats.logistic.ppf(y) + assert_allclose(x, xx) + + def test_sf_isf(self): + x = np.linspace(-20, 20) + y = stats.logistic.sf(x) + xx = stats.logistic.isf(y) + assert_allclose(x, xx) + + def test_extreme_values(self): + # p is chosen so that 1 - (1 - p) == p in double precision + p = 9.992007221626409e-16 + desired = 34.53957599234088 + assert_allclose(stats.logistic.ppf(1 - p), desired) + assert_allclose(stats.logistic.isf(p), desired) + + def test_logpdf_basic(self): + logp = stats.logistic.logpdf([-15, 0, 10]) + # Expected values computed with mpmath with 50 digits of precision. + expected = [-15.000000611804547, + -1.3862943611198906, + -10.000090797798434] + assert_allclose(logp, expected, rtol=1e-13) + + def test_logpdf_extreme_values(self): + logp = stats.logistic.logpdf([800, -800]) + # For such large arguments, logpdf(x) = -abs(x) when computed + # with 64 bit floating point. + assert_equal(logp, [-800, -800]) + + @pytest.mark.parametrize("loc_rvs,scale_rvs", [(0.4484955, 0.10216821), + (0.62918191, 0.74367064)]) + def test_fit(self, loc_rvs, scale_rvs): + data = stats.logistic.rvs(size=100, loc=loc_rvs, scale=scale_rvs, + random_state=self.rng) + + # test that result of fit method is the same as optimization + def func(input, data): + a, b = input + n = len(data) + x1 = np.sum(np.exp((data - a) / b) / + (1 + np.exp((data - a) / b))) - n / 2 + x2 = np.sum(((data - a) / b) * + ((np.exp((data - a) / b) - 1) / + (np.exp((data - a) / b) + 1))) - n + return x1, x2 + + expected_solution = root(func, stats.logistic._fitstart(data), args=( + data,)).x + fit_method = stats.logistic.fit(data) + + # other than computational variances, the fit method and the solution + # to this system of equations are equal + assert_allclose(fit_method, expected_solution, atol=1e-30) + + def test_fit_comp_optimizer(self): + data = stats.logistic.rvs(size=100, loc=0.5, scale=2, random_state=self.rng) + _assert_less_or_close_loglike(stats.logistic, data) + _assert_less_or_close_loglike(stats.logistic, data, floc=1) + _assert_less_or_close_loglike(stats.logistic, data, fscale=1) + + @pytest.mark.parametrize('testlogcdf', [True, False]) + def test_logcdfsf_tails(self, testlogcdf): + # Test either logcdf or logsf. By symmetry, we can use the same + # expected values for both by switching the sign of x for logsf. + x = np.array([-10000, -800, 17, 50, 500]) + if testlogcdf: + y = stats.logistic.logcdf(x) + else: + y = stats.logistic.logsf(-x) + # The expected values were computed with mpmath. + expected = [-10000.0, -800.0, -4.139937633089748e-08, + -1.9287498479639178e-22, -7.124576406741286e-218] + assert_allclose(y, expected, rtol=2e-15) + + def test_fit_gh_18176(self): + # logistic.fit returned `scale < 0` for this data. Check that this has + # been fixed. + data = np.array([-459, 37, 43, 45, 45, 48, 54, 55, 58] + + [59] * 3 + [61] * 9) + # If scale were negative, NLLF would be infinite, so this would fail + _assert_less_or_close_loglike(stats.logistic, data) + + +class TestLogser: + def setup_method(self): + self.rng = np.random.default_rng(187505461) + + def test_rvs(self): + vals = stats.logser.rvs(0.75, size=(2, 50), random_state=self.rng) + assert np.all(vals >= 1) + assert np.shape(vals) == (2, 50) + assert vals.dtype.char in typecodes['AllInteger'] + val = stats.logser.rvs(0.75, random_state=self.rng) + assert isinstance(val, int) + val = stats.logser(0.75).rvs(3, random_state=self.rng) + assert isinstance(val, np.ndarray) + assert val.dtype.char in typecodes['AllInteger'] + + def test_pmf_small_p(self): + m = stats.logser.pmf(4, 1e-20) + # The expected value was computed using mpmath: + # >>> import mpmath + # >>> mpmath.mp.dps = 64 + # >>> k = 4 + # >>> p = mpmath.mpf('1e-20') + # >>> float(-(p**k)/k/mpmath.log(1-p)) + # 2.5e-61 + # It is also clear from noticing that for very small p, + # log(1-p) is approximately -p, and the formula becomes + # p**(k-1) / k + assert_allclose(m, 2.5e-61) + + def test_mean_small_p(self): + m = stats.logser.mean(1e-8) + # The expected mean was computed using mpmath: + # >>> import mpmath + # >>> mpmath.dps = 60 + # >>> p = mpmath.mpf('1e-8') + # >>> float(-p / ((1 - p)*mpmath.log(1 - p))) + # 1.000000005 + assert_allclose(m, 1.000000005) + + def test_sf(self): + p = [[0.5], [1e-5], [1 - 1e-5]] + k = [1, 10, 100, 1000] + # Reference values from Wolfram Alpha, e.g. + # SurvivalFunction[LogSeriesDistribution[99999/100000], 1000] + # 0.35068668662799737584735036958139157462633608106500173019897861351038286634 + ref = [[0.2786524795555183, 0.00011876901682721189, + 1.1159788564768581e-32, 1.3437300083506688e-304], + [5.000008333375e-06, 9.090946969973778e-52, 0, 0], + [0.9131419722083134, 0.745601735620566, + 0.5495169511115096, 0.3506866866279974]] + res = stats.logser.sf(k, p) + np.testing.assert_allclose(res, ref, atol=1e-300) + + +class TestGumbel_r_l: + @pytest.fixture(scope='function') + def rng(self): + return np.random.default_rng(1234) + + @pytest.mark.parametrize("dist", [stats.gumbel_r, stats.gumbel_l]) + @pytest.mark.parametrize("loc_rvs", [-1, 0, 1]) + @pytest.mark.parametrize("scale_rvs", [.1, 1, 5]) + @pytest.mark.parametrize('fix_loc, fix_scale', + ([True, False], [False, True])) + def test_fit_comp_optimizer(self, dist, loc_rvs, scale_rvs, + fix_loc, fix_scale, rng): + data = dist.rvs(size=100, loc=loc_rvs, scale=scale_rvs, + random_state=rng) + + kwds = dict() + # the fixed location and scales are arbitrarily modified to not be + # close to the true value. + if fix_loc: + kwds['floc'] = loc_rvs * 2 + if fix_scale: + kwds['fscale'] = scale_rvs * 2 + + # test that the gumbel_* fit method is better than super method + _assert_less_or_close_loglike(dist, data, **kwds) + + @pytest.mark.parametrize("dist, sgn", [(stats.gumbel_r, 1), + (stats.gumbel_l, -1)]) + def test_fit(self, dist, sgn): + z = sgn*np.array([3, 3, 3, 3, 3, 3, 3, 3.00000001]) + loc, scale = dist.fit(z) + # The expected values were computed with mpmath with 60 digits + # of precision. + assert_allclose(loc, sgn*3.0000000001667906) + assert_allclose(scale, 1.2495222465145514e-09, rtol=1e-6) + + +class TestPareto: + def test_stats(self): + # Check the stats() method with some simple values. Also check + # that the calculations do not trigger RuntimeWarnings. + with warnings.catch_warnings(): + warnings.simplefilter("error", RuntimeWarning) + + m, v, s, k = stats.pareto.stats(0.5, moments='mvsk') + assert_equal(m, np.inf) + assert_equal(v, np.inf) + assert_equal(s, np.nan) + assert_equal(k, np.nan) + + m, v, s, k = stats.pareto.stats(1.0, moments='mvsk') + assert_equal(m, np.inf) + assert_equal(v, np.inf) + assert_equal(s, np.nan) + assert_equal(k, np.nan) + + m, v, s, k = stats.pareto.stats(1.5, moments='mvsk') + assert_equal(m, 3.0) + assert_equal(v, np.inf) + assert_equal(s, np.nan) + assert_equal(k, np.nan) + + m, v, s, k = stats.pareto.stats(2.0, moments='mvsk') + assert_equal(m, 2.0) + assert_equal(v, np.inf) + assert_equal(s, np.nan) + assert_equal(k, np.nan) + + m, v, s, k = stats.pareto.stats(2.5, moments='mvsk') + assert_allclose(m, 2.5 / 1.5) + assert_allclose(v, 2.5 / (1.5*1.5*0.5)) + assert_equal(s, np.nan) + assert_equal(k, np.nan) + + m, v, s, k = stats.pareto.stats(3.0, moments='mvsk') + assert_allclose(m, 1.5) + assert_allclose(v, 0.75) + assert_equal(s, np.nan) + assert_equal(k, np.nan) + + m, v, s, k = stats.pareto.stats(3.5, moments='mvsk') + assert_allclose(m, 3.5 / 2.5) + assert_allclose(v, 3.5 / (2.5*2.5*1.5)) + assert_allclose(s, (2*4.5/0.5)*np.sqrt(1.5/3.5)) + assert_equal(k, np.nan) + + m, v, s, k = stats.pareto.stats(4.0, moments='mvsk') + assert_allclose(m, 4.0 / 3.0) + assert_allclose(v, 4.0 / 18.0) + assert_allclose(s, 2*(1+4.0)/(4.0-3) * np.sqrt((4.0-2)/4.0)) + assert_equal(k, np.nan) + + m, v, s, k = stats.pareto.stats(4.5, moments='mvsk') + assert_allclose(m, 4.5 / 3.5) + assert_allclose(v, 4.5 / (3.5*3.5*2.5)) + assert_allclose(s, (2*5.5/1.5) * np.sqrt(2.5/4.5)) + assert_allclose(k, 6*(4.5**3 + 4.5**2 - 6*4.5 - 2)/(4.5*1.5*0.5)) + + def test_sf(self): + x = 1e9 + b = 2 + scale = 1.5 + p = stats.pareto.sf(x, b, loc=0, scale=scale) + expected = (scale/x)**b # 2.25e-18 + assert_allclose(p, expected) + + @pytest.mark.filterwarnings("ignore:invalid value encountered in " + "double_scalars") + @pytest.mark.parametrize("rvs_shape", [1, 2]) + @pytest.mark.parametrize("rvs_loc", [0, 2]) + @pytest.mark.parametrize("rvs_scale", [1, 5]) + def test_fit(self, rvs_shape, rvs_loc, rvs_scale): + rng = np.random.default_rng(1234) + data = stats.pareto.rvs(size=100, b=rvs_shape, scale=rvs_scale, + loc=rvs_loc, random_state=rng) + + # shape can still be fixed with multiple names + shape_mle_analytical1 = stats.pareto.fit(data, floc=0, f0=1.04)[0] + shape_mle_analytical2 = stats.pareto.fit(data, floc=0, fix_b=1.04)[0] + shape_mle_analytical3 = stats.pareto.fit(data, floc=0, fb=1.04)[0] + assert (shape_mle_analytical1 == shape_mle_analytical2 == + shape_mle_analytical3 == 1.04) + + # data can be shifted with changes to `loc` + data = stats.pareto.rvs(size=100, b=rvs_shape, scale=rvs_scale, + loc=(rvs_loc + 2), random_state=rng) + shape_mle_a, loc_mle_a, scale_mle_a = stats.pareto.fit(data, floc=2) + assert_equal(scale_mle_a + 2, data.min()) + + data_shift = data - 2 + ndata = data_shift.shape[0] + assert_equal(shape_mle_a, + ndata / np.sum(np.log(data_shift/data_shift.min()))) + assert_equal(loc_mle_a, 2) + + @pytest.mark.parametrize("rvs_shape", [.1, 2]) + @pytest.mark.parametrize("rvs_loc", [0, 2]) + @pytest.mark.parametrize("rvs_scale", [1, 5]) + @pytest.mark.parametrize('fix_shape, fix_loc, fix_scale', + [p for p in product([True, False], repeat=3) + if False in p]) + @np.errstate(invalid="ignore") + def test_fit_MLE_comp_optimizer(self, rvs_shape, rvs_loc, rvs_scale, + fix_shape, fix_loc, fix_scale): + rng = np.random.default_rng(1234) + data = stats.pareto.rvs(size=100, b=rvs_shape, scale=rvs_scale, + loc=rvs_loc, random_state=rng) + + kwds = {} + if fix_shape: + kwds['f0'] = rvs_shape + if fix_loc: + kwds['floc'] = rvs_loc + if fix_scale: + kwds['fscale'] = rvs_scale + + _assert_less_or_close_loglike(stats.pareto, data, **kwds) + + @np.errstate(invalid="ignore") + def test_fit_known_bad_seed(self): + # Tests a known seed and set of parameters that would produce a result + # would violate the support of Pareto if the fit method did not check + # the constraint `fscale + floc < min(data)`. + shape, location, scale = 1, 0, 1 + data = stats.pareto.rvs(shape, location, scale, size=100, + random_state=np.random.default_rng(2535619)) + _assert_less_or_close_loglike(stats.pareto, data) + + def test_fit_warnings(self): + assert_fit_warnings(stats.pareto) + # `floc` that causes invalid negative data + assert_raises(FitDataError, stats.pareto.fit, [1, 2, 3], floc=2) + # `floc` and `fscale` combination causes invalid data + assert_raises(FitDataError, stats.pareto.fit, [5, 2, 3], floc=1, + fscale=3) + + def test_negative_data(self): + rng = np.random.default_rng(1234) + data = stats.pareto.rvs(loc=-130, b=1, size=100, random_state=rng) + assert_array_less(data, 0) + # The purpose of this test is to make sure that no runtime warnings are + # raised for all negative data, not the output of the fit method. Other + # methods test the output but have to silence warnings from the super + # method. + _ = stats.pareto.fit(data) + + +class TestGenpareto: + def test_ab(self): + # c >= 0: a, b = [0, inf] + for c in [1., 0.]: + c = np.asarray(c) + a, b = stats.genpareto._get_support(c) + assert_equal(a, 0.) + assert_(np.isposinf(b)) + + # c < 0: a=0, b=1/|c| + c = np.asarray(-2.) + a, b = stats.genpareto._get_support(c) + assert_allclose([a, b], [0., 0.5]) + + def test_c0(self): + # with c=0, genpareto reduces to the exponential distribution + # rv = stats.genpareto(c=0.) + rv = stats.genpareto(c=0.) + x = np.linspace(0, 10., 30) + assert_allclose(rv.pdf(x), stats.expon.pdf(x)) + assert_allclose(rv.cdf(x), stats.expon.cdf(x)) + assert_allclose(rv.sf(x), stats.expon.sf(x)) + + q = np.linspace(0., 1., 10) + assert_allclose(rv.ppf(q), stats.expon.ppf(q)) + + def test_cm1(self): + # with c=-1, genpareto reduces to the uniform distr on [0, 1] + rv = stats.genpareto(c=-1.) + x = np.linspace(0, 10., 30) + assert_allclose(rv.pdf(x), stats.uniform.pdf(x)) + assert_allclose(rv.cdf(x), stats.uniform.cdf(x)) + assert_allclose(rv.sf(x), stats.uniform.sf(x)) + + q = np.linspace(0., 1., 10) + assert_allclose(rv.ppf(q), stats.uniform.ppf(q)) + + # logpdf(1., c=-1) should be zero + assert_allclose(rv.logpdf(1), 0) + + def test_x_inf(self): + # make sure x=inf is handled gracefully + rv = stats.genpareto(c=0.1) + assert_allclose([rv.pdf(np.inf), rv.cdf(np.inf)], [0., 1.]) + assert_(np.isneginf(rv.logpdf(np.inf))) + + rv = stats.genpareto(c=0.) + assert_allclose([rv.pdf(np.inf), rv.cdf(np.inf)], [0., 1.]) + assert_(np.isneginf(rv.logpdf(np.inf))) + + rv = stats.genpareto(c=-1.) + assert_allclose([rv.pdf(np.inf), rv.cdf(np.inf)], [0., 1.]) + assert_(np.isneginf(rv.logpdf(np.inf))) + + def test_c_continuity(self): + # pdf is continuous at c=0, -1 + x = np.linspace(0, 10, 30) + for c in [0, -1]: + pdf0 = stats.genpareto.pdf(x, c) + for dc in [1e-14, -1e-14]: + pdfc = stats.genpareto.pdf(x, c + dc) + assert_allclose(pdf0, pdfc, atol=1e-12) + + cdf0 = stats.genpareto.cdf(x, c) + for dc in [1e-14, 1e-14]: + cdfc = stats.genpareto.cdf(x, c + dc) + assert_allclose(cdf0, cdfc, atol=1e-12) + + def test_c_continuity_ppf(self): + q = np.r_[np.logspace(1e-12, 0.01, base=0.1), + np.linspace(0.01, 1, 30, endpoint=False), + 1. - np.logspace(1e-12, 0.01, base=0.1)] + for c in [0., -1.]: + ppf0 = stats.genpareto.ppf(q, c) + for dc in [1e-14, -1e-14]: + ppfc = stats.genpareto.ppf(q, c + dc) + assert_allclose(ppf0, ppfc, atol=1e-12) + + def test_c_continuity_isf(self): + q = np.r_[np.logspace(1e-12, 0.01, base=0.1), + np.linspace(0.01, 1, 30, endpoint=False), + 1. - np.logspace(1e-12, 0.01, base=0.1)] + for c in [0., -1.]: + isf0 = stats.genpareto.isf(q, c) + for dc in [1e-14, -1e-14]: + isfc = stats.genpareto.isf(q, c + dc) + assert_allclose(isf0, isfc, atol=1e-12) + + def test_cdf_ppf_roundtrip(self): + # this should pass with machine precision. hat tip @pbrod + q = np.r_[np.logspace(1e-12, 0.01, base=0.1), + np.linspace(0.01, 1, 30, endpoint=False), + 1. - np.logspace(1e-12, 0.01, base=0.1)] + for c in [1e-8, -1e-18, 1e-15, -1e-15]: + assert_allclose(stats.genpareto.cdf(stats.genpareto.ppf(q, c), c), + q, atol=1e-15) + + def test_logsf(self): + logp = stats.genpareto.logsf(1e10, .01, 0, 1) + assert_allclose(logp, -1842.0680753952365) + + # Values in 'expected_stats' are + # [mean, variance, skewness, excess kurtosis]. + @pytest.mark.parametrize( + 'c, expected_stats', + [(0, [1, 1, 2, 6]), + (1/4, [4/3, 32/9, 10/np.sqrt(2), np.nan]), + (1/9, [9/8, (81/64)*(9/7), (10/9)*np.sqrt(7), 754/45]), + (-1, [1/2, 1/12, 0, -6/5])]) + def test_stats(self, c, expected_stats): + result = stats.genpareto.stats(c, moments='mvsk') + assert_allclose(result, expected_stats, rtol=1e-13, atol=1e-15) + + def test_var(self): + # Regression test for gh-11168. + v = stats.genpareto.var(1e-8) + assert_allclose(v, 1.000000040000001, rtol=1e-13) + + +class TestPearson3: + def setup_method(self): + self.rng = np.random.default_rng(2775995570) + + def test_rvs(self): + vals = stats.pearson3.rvs(0.1, size=(2, 50), random_state=self.rng) + assert np.shape(vals) == (2, 50) + assert vals.dtype.char in typecodes['AllFloat'] + val = stats.pearson3.rvs(0.5, random_state=self.rng) + assert isinstance(val, float) + val = stats.pearson3(0.5).rvs(3, random_state=self.rng) + assert isinstance(val, np.ndarray) + assert val.dtype.char in typecodes['AllFloat'] + assert len(val) == 3 + + def test_pdf(self): + vals = stats.pearson3.pdf(2, [0.0, 0.1, 0.2]) + assert_allclose(vals, np.array([0.05399097, 0.05555481, 0.05670246]), + atol=1e-6) + vals = stats.pearson3.pdf(-3, 0.1) + assert_allclose(vals, np.array([0.00313791]), atol=1e-6) + vals = stats.pearson3.pdf([-3, -2, -1, 0, 1], 0.1) + assert_allclose(vals, np.array([0.00313791, 0.05192304, 0.25028092, + 0.39885918, 0.23413173]), atol=1e-6) + + def test_cdf(self): + vals = stats.pearson3.cdf(2, [0.0, 0.1, 0.2]) + assert_allclose(vals, np.array([0.97724987, 0.97462004, 0.97213626]), + atol=1e-6) + vals = stats.pearson3.cdf(-3, 0.1) + assert_allclose(vals, [0.00082256], atol=1e-6) + vals = stats.pearson3.cdf([-3, -2, -1, 0, 1], 0.1) + assert_allclose(vals, [8.22563821e-04, 1.99860448e-02, 1.58550710e-01, + 5.06649130e-01, 8.41442111e-01], atol=1e-6) + + def test_negative_cdf_bug_11186(self): + # incorrect CDFs for negative skews in gh-11186; fixed in gh-12640 + # Also check vectorization w/ negative, zero, and positive skews + skews = [-3, -1, 0, 0.5] + x_eval = 0.5 + neg_inf = -30 # avoid RuntimeWarning caused by np.log(0) + cdfs = stats.pearson3.cdf(x_eval, skews) + int_pdfs = [quad(stats.pearson3(skew).pdf, neg_inf, x_eval)[0] + for skew in skews] + assert_allclose(cdfs, int_pdfs) + + def test_return_array_bug_11746(self): + # pearson3.moment was returning size 0 or 1 array instead of float + # The first moment is equal to the loc, which defaults to zero + moment = stats.pearson3.moment(1, 2) + assert_equal(moment, 0) + assert isinstance(moment, np.number) + + moment = stats.pearson3.moment(1, 0.000001) + assert_equal(moment, 0) + assert isinstance(moment, np.number) + + def test_ppf_bug_17050(self): + # incorrect PPF for negative skews were reported in gh-17050 + # Check that this is fixed (even in the array case) + skews = [-3, -1, 0, 0.5] + x_eval = 0.5 + res = stats.pearson3.ppf(stats.pearson3.cdf(x_eval, skews), skews) + assert_allclose(res, x_eval) + + # Negation of the skew flips the distribution about the origin, so + # the following should hold + skew = np.array([[-0.5], [1.5]]) + x = np.linspace(-2, 2) + assert_allclose(stats.pearson3.pdf(x, skew), + stats.pearson3.pdf(-x, -skew)) + assert_allclose(stats.pearson3.cdf(x, skew), + stats.pearson3.sf(-x, -skew)) + assert_allclose(stats.pearson3.ppf(x, skew), + -stats.pearson3.isf(x, -skew)) + + def test_sf(self): + # reference values were computed via the reference distribution, e.g. + # mp.dps = 50; Pearson3(skew=skew).sf(x). Check positive, negative, + # and zero skew due to branching. + skew = [0.1, 0.5, 1.0, -0.1] + x = [5.0, 10.0, 50.0, 8.0] + ref = [1.64721926440872e-06, 8.271911573556123e-11, + 1.3149506021756343e-40, 2.763057937820296e-21] + assert_allclose(stats.pearson3.sf(x, skew), ref, rtol=2e-14) + assert_allclose(stats.pearson3.sf(x, 0), stats.norm.sf(x), rtol=2e-14) + + +class TestKappa4: + def test_cdf_genpareto(self): + # h = 1 and k != 0 is generalized Pareto + x = [0.0, 0.1, 0.2, 0.5] + h = 1.0 + for k in [-1.9, -1.0, -0.5, -0.2, -0.1, 0.1, 0.2, 0.5, 1.0, + 1.9]: + vals = stats.kappa4.cdf(x, h, k) + # shape parameter is opposite what is expected + vals_comp = stats.genpareto.cdf(x, -k) + assert_allclose(vals, vals_comp) + + def test_cdf_genextreme(self): + # h = 0 and k != 0 is generalized extreme value + x = np.linspace(-5, 5, 10) + h = 0.0 + k = np.linspace(-3, 3, 10) + vals = stats.kappa4.cdf(x, h, k) + vals_comp = stats.genextreme.cdf(x, k) + assert_allclose(vals, vals_comp) + + def test_cdf_expon(self): + # h = 1 and k = 0 is exponential + x = np.linspace(0, 10, 10) + h = 1.0 + k = 0.0 + vals = stats.kappa4.cdf(x, h, k) + vals_comp = stats.expon.cdf(x) + assert_allclose(vals, vals_comp) + + def test_cdf_gumbel_r(self): + # h = 0 and k = 0 is gumbel_r + x = np.linspace(-5, 5, 10) + h = 0.0 + k = 0.0 + vals = stats.kappa4.cdf(x, h, k) + vals_comp = stats.gumbel_r.cdf(x) + assert_allclose(vals, vals_comp) + + def test_cdf_logistic(self): + # h = -1 and k = 0 is logistic + x = np.linspace(-5, 5, 10) + h = -1.0 + k = 0.0 + vals = stats.kappa4.cdf(x, h, k) + vals_comp = stats.logistic.cdf(x) + assert_allclose(vals, vals_comp) + + def test_cdf_uniform(self): + # h = 1 and k = 1 is uniform + x = np.linspace(-5, 5, 10) + h = 1.0 + k = 1.0 + vals = stats.kappa4.cdf(x, h, k) + vals_comp = stats.uniform.cdf(x) + assert_allclose(vals, vals_comp) + + def test_integers_ctor(self): + # regression test for gh-7416: _argcheck fails for integer h and k + # in numpy 1.12 + stats.kappa4(1, 2) + + +class TestPoisson: + def setup_method(self): + self.rng = np.random.default_rng(9799340796) + + def test_pmf_basic(self): + # Basic case + ln2 = np.log(2) + vals = stats.poisson.pmf([0, 1, 2], ln2) + expected = [0.5, ln2/2, ln2**2/4] + assert_allclose(vals, expected) + + def test_mu0(self): + # Edge case: mu=0 + vals = stats.poisson.pmf([0, 1, 2], 0) + expected = [1, 0, 0] + assert_array_equal(vals, expected) + + interval = stats.poisson.interval(0.95, 0) + assert_equal(interval, (0, 0)) + + def test_rvs(self): + vals = stats.poisson.rvs(0.5, size=(2, 50), random_state=self.rng) + assert np.all(vals >= 0) + assert np.shape(vals) == (2, 50) + assert vals.dtype.char in typecodes['AllInteger'] + val = stats.poisson.rvs(0.5, random_state=self.rng) + assert isinstance(val, int) + val = stats.poisson(0.5).rvs(3, random_state=self.rng) + assert isinstance(val, np.ndarray) + assert val.dtype.char in typecodes['AllInteger'] + + def test_stats(self): + mu = 16.0 + result = stats.poisson.stats(mu, moments='mvsk') + assert_allclose(result, [mu, mu, np.sqrt(1.0/mu), 1.0/mu]) + + mu = np.array([0.0, 1.0, 2.0]) + result = stats.poisson.stats(mu, moments='mvsk') + expected = (mu, mu, [np.inf, 1, 1/np.sqrt(2)], [np.inf, 1, 0.5]) + assert_allclose(result, expected) + + +class TestKSTwo: + + def test_cdf(self): + for n in [1, 2, 3, 10, 100, 1000]: + # Test x-values: + # 0, 1/2n, where the cdf should be 0 + # 1/n, where the cdf should be n!/n^n + # 0.5, where the cdf should match ksone.cdf + # 1-1/n, where cdf = 1-2/n^n + # 1, where cdf == 1 + # (E.g. Exact values given by Eqn 1 in Simard / L'Ecuyer) + x = np.array([0, 0.5/n, 1/n, 0.5, 1-1.0/n, 1]) + v1 = (1.0/n)**n + lg = scipy.special.gammaln(n+1) + elg = (np.exp(lg) if v1 != 0 else 0) + expected = np.array([0, 0, v1 * elg, + 1 - 2*stats.ksone.sf(0.5, n), + max(1 - 2*v1, 0.0), + 1.0]) + vals_cdf = stats.kstwo.cdf(x, n) + assert_allclose(vals_cdf, expected) + + def test_sf(self): + x = np.linspace(0, 1, 11) + for n in [1, 2, 3, 10, 100, 1000]: + # Same x values as in test_cdf, and use sf = 1 - cdf + x = np.array([0, 0.5/n, 1/n, 0.5, 1-1.0/n, 1]) + v1 = (1.0/n)**n + lg = scipy.special.gammaln(n+1) + elg = (np.exp(lg) if v1 != 0 else 0) + expected = np.array([1.0, 1.0, + 1 - v1 * elg, + 2*stats.ksone.sf(0.5, n), + min(2*v1, 1.0), 0]) + vals_sf = stats.kstwo.sf(x, n) + assert_allclose(vals_sf, expected) + + def test_cdf_sqrtn(self): + # For fixed a, cdf(a/sqrt(n), n) -> kstwobign(a) as n->infinity + # cdf(a/sqrt(n), n) is an increasing function of n (and a) + # Check that the function is indeed increasing (allowing for some + # small floating point and algorithm differences.) + x = np.linspace(0, 2, 11)[1:] + ns = [50, 100, 200, 400, 1000, 2000] + for _x in x: + xn = _x / np.sqrt(ns) + probs = stats.kstwo.cdf(xn, ns) + diffs = np.diff(probs) + assert_array_less(diffs, 1e-8) + + def test_cdf_sf(self): + x = np.linspace(0, 1, 11) + for n in [1, 2, 3, 10, 100, 1000]: + vals_cdf = stats.kstwo.cdf(x, n) + vals_sf = stats.kstwo.sf(x, n) + assert_array_almost_equal(vals_cdf, 1 - vals_sf) + + def test_cdf_sf_sqrtn(self): + x = np.linspace(0, 1, 11) + for n in [1, 2, 3, 10, 100, 1000]: + xn = x / np.sqrt(n) + vals_cdf = stats.kstwo.cdf(xn, n) + vals_sf = stats.kstwo.sf(xn, n) + assert_array_almost_equal(vals_cdf, 1 - vals_sf) + + def test_ppf_of_cdf(self): + x = np.linspace(0, 1, 11) + for n in [1, 2, 3, 10, 100, 1000]: + xn = x[x > 0.5/n] + vals_cdf = stats.kstwo.cdf(xn, n) + # CDFs close to 1 are better dealt with using the SF + cond = (0 < vals_cdf) & (vals_cdf < 0.99) + vals = stats.kstwo.ppf(vals_cdf, n) + assert_allclose(vals[cond], xn[cond], rtol=1e-4) + + def test_isf_of_sf(self): + x = np.linspace(0, 1, 11) + for n in [1, 2, 3, 10, 100, 1000]: + xn = x[x > 0.5/n] + vals_isf = stats.kstwo.isf(xn, n) + cond = (0 < vals_isf) & (vals_isf < 1.0) + vals = stats.kstwo.sf(vals_isf, n) + assert_allclose(vals[cond], xn[cond], rtol=1e-4) + + def test_ppf_of_cdf_sqrtn(self): + x = np.linspace(0, 1, 11) + for n in [1, 2, 3, 10, 100, 1000]: + xn = (x / np.sqrt(n))[x > 0.5/n] + vals_cdf = stats.kstwo.cdf(xn, n) + cond = (0 < vals_cdf) & (vals_cdf < 1.0) + vals = stats.kstwo.ppf(vals_cdf, n) + assert_allclose(vals[cond], xn[cond]) + + def test_isf_of_sf_sqrtn(self): + x = np.linspace(0, 1, 11) + for n in [1, 2, 3, 10, 100, 1000]: + xn = (x / np.sqrt(n))[x > 0.5/n] + vals_sf = stats.kstwo.sf(xn, n) + # SFs close to 1 are better dealt with using the CDF + cond = (0 < vals_sf) & (vals_sf < 0.95) + vals = stats.kstwo.isf(vals_sf, n) + assert_allclose(vals[cond], xn[cond]) + + def test_ppf(self): + probs = np.linspace(0, 1, 11)[1:] + for n in [1, 2, 3, 10, 100, 1000]: + xn = stats.kstwo.ppf(probs, n) + vals_cdf = stats.kstwo.cdf(xn, n) + assert_allclose(vals_cdf, probs) + + def test_simard_lecuyer_table1(self): + # Compute the cdf for values near the mean of the distribution. + # The mean u ~ log(2)*sqrt(pi/(2n)) + # Compute for x in [u/4, u/3, u/2, u, 2u, 3u] + # This is the computation of Table 1 of Simard, R., L'Ecuyer, P. (2011) + # "Computing the Two-Sided Kolmogorov-Smirnov Distribution". + # Except that the values below are not from the published table, but + # were generated using an independent SageMath implementation of + # Durbin's algorithm (with the exponentiation and scaling of + # Marsaglia/Tsang/Wang's version) using 500 bit arithmetic. + # Some of the values in the published table have relative + # errors greater than 1e-4. + ns = [10, 50, 100, 200, 500, 1000] + ratios = np.array([1.0/4, 1.0/3, 1.0/2, 1, 2, 3]) + expected = np.array([ + [1.92155292e-08, 5.72933228e-05, 2.15233226e-02, 6.31566589e-01, + 9.97685592e-01, 9.99999942e-01], + [2.28096224e-09, 1.99142563e-05, 1.42617934e-02, 5.95345542e-01, + 9.96177701e-01, 9.99998662e-01], + [1.00201886e-09, 1.32673079e-05, 1.24608594e-02, 5.86163220e-01, + 9.95866877e-01, 9.99998240e-01], + [4.93313022e-10, 9.52658029e-06, 1.12123138e-02, 5.79486872e-01, + 9.95661824e-01, 9.99997964e-01], + [2.37049293e-10, 6.85002458e-06, 1.01309221e-02, 5.73427224e-01, + 9.95491207e-01, 9.99997750e-01], + [1.56990874e-10, 5.71738276e-06, 9.59725430e-03, 5.70322692e-01, + 9.95409545e-01, 9.99997657e-01] + ]) + for idx, n in enumerate(ns): + x = ratios * np.log(2) * np.sqrt(np.pi/2/n) + vals_cdf = stats.kstwo.cdf(x, n) + assert_allclose(vals_cdf, expected[idx], rtol=1e-5) + + +class TestZipf: + def setup_method(self): + self.rng = np.random.default_rng(2444103536) + + def test_rvs(self): + vals = stats.zipf.rvs(1.5, size=(2, 50), random_state=self.rng) + assert np.all(vals >= 1) + assert np.shape(vals) == (2, 50) + assert vals.dtype.char in typecodes['AllInteger'] + val = stats.zipf.rvs(1.5, random_state=self.rng) + assert isinstance(val, int) + val = stats.zipf(1.5).rvs(3, random_state=self.rng) + assert isinstance(val, np.ndarray) + assert val.dtype.char in typecodes['AllInteger'] + + def test_moments(self): + # n-th moment is finite iff a > n + 1 + m, v = stats.zipf.stats(a=2.8) + assert_(np.isfinite(m)) + assert_equal(v, np.inf) + + s, k = stats.zipf.stats(a=4.8, moments='sk') + assert_(not np.isfinite([s, k]).all()) + + +class TestDLaplace: + def setup_method(self): + self.rng = np.random.default_rng(460403075) + + def test_rvs(self): + vals = stats.dlaplace.rvs(1.5, size=(2, 50), random_state=self.rng) + assert np.shape(vals) == (2, 50) + assert vals.dtype.char in typecodes['AllInteger'] + val = stats.dlaplace.rvs(1.5, random_state=self.rng) + assert isinstance(val, int) + val = stats.dlaplace(1.5).rvs(3, random_state=self.rng) + assert isinstance(val, np.ndarray) + assert val.dtype.char in typecodes['AllInteger'] + assert stats.dlaplace.rvs(0.8, random_state=self.rng) is not None + + def test_stats(self): + # compare the explicit formulas w/ direct summation using pmf + a = 1. + dl = stats.dlaplace(a) + m, v, s, k = dl.stats('mvsk') + + N = 37 + xx = np.arange(-N, N+1) + pp = dl.pmf(xx) + m2, m4 = np.sum(pp*xx**2), np.sum(pp*xx**4) + assert_equal((m, s), (0, 0)) + assert_allclose((v, k), (m2, m4/m2**2 - 3.), atol=1e-14, rtol=1e-8) + + def test_stats2(self): + a = np.log(2.) + dl = stats.dlaplace(a) + m, v, s, k = dl.stats('mvsk') + assert_equal((m, s), (0., 0.)) + assert_allclose((v, k), (4., 3.25)) + + +class TestInvgauss: + def setup_method(self): + self.rng = np.random.default_rng(5422839947) + + @pytest.mark.parametrize("rvs_mu,rvs_loc,rvs_scale", + [(2, 0, 1), (4.635, 4.362, 6.303)]) + def test_fit(self, rvs_mu, rvs_loc, rvs_scale): + data = stats.invgauss.rvs(size=100, mu=rvs_mu, + loc=rvs_loc, scale=rvs_scale, random_state=self.rng) + # Analytical MLEs are calculated with formula when `floc` is fixed + mu, loc, scale = stats.invgauss.fit(data, floc=rvs_loc) + + data = data - rvs_loc + mu_temp = np.mean(data) + scale_mle = len(data) / (np.sum(data**(-1) - mu_temp**(-1))) + mu_mle = mu_temp/scale_mle + + # `mu` and `scale` match analytical formula + assert_allclose(mu_mle, mu, atol=1e-15, rtol=1e-15) + assert_allclose(scale_mle, scale, atol=1e-15, rtol=1e-15) + assert_equal(loc, rvs_loc) + data = stats.invgauss.rvs(size=100, mu=rvs_mu, + loc=rvs_loc, scale=rvs_scale, random_state=self.rng) + # fixed parameters are returned + mu, loc, scale = stats.invgauss.fit(data, floc=rvs_loc - 1, + fscale=rvs_scale + 1) + assert_equal(rvs_scale + 1, scale) + assert_equal(rvs_loc - 1, loc) + + # shape can still be fixed with multiple names + shape_mle1 = stats.invgauss.fit(data, fmu=1.04)[0] + shape_mle2 = stats.invgauss.fit(data, fix_mu=1.04)[0] + shape_mle3 = stats.invgauss.fit(data, f0=1.04)[0] + assert shape_mle1 == shape_mle2 == shape_mle3 == 1.04 + + @pytest.mark.parametrize("rvs_mu,rvs_loc,rvs_scale", + [(2, 0, 1), (6.311, 3.225, 4.520)]) + def test_fit_MLE_comp_optimizer(self, rvs_mu, rvs_loc, rvs_scale): + rng = np.random.RandomState(1234) + data = stats.invgauss.rvs(size=100, mu=rvs_mu, + loc=rvs_loc, scale=rvs_scale, random_state=rng) + + super_fit = super(type(stats.invgauss), stats.invgauss).fit + # fitting without `floc` uses superclass fit method + super_fitted = super_fit(data) + invgauss_fit = stats.invgauss.fit(data) + assert_equal(super_fitted, invgauss_fit) + + # fitting with `fmu` is uses superclass fit method + super_fitted = super_fit(data, floc=0, fmu=2) + invgauss_fit = stats.invgauss.fit(data, floc=0, fmu=2) + assert_equal(super_fitted, invgauss_fit) + + # fixed `floc` uses analytical formula and provides better fit than + # super method + _assert_less_or_close_loglike(stats.invgauss, data, floc=rvs_loc) + + # fixed `floc` not resulting in invalid data < 0 uses analytical + # formulas and provides a better fit than the super method + assert np.all((data - (rvs_loc - 1)) > 0) + _assert_less_or_close_loglike(stats.invgauss, data, floc=rvs_loc - 1) + + # fixed `floc` to an arbitrary number, 0, still provides a better fit + # than the super method + _assert_less_or_close_loglike(stats.invgauss, data, floc=0) + + # fixed `fscale` to an arbitrary number still provides a better fit + # than the super method + _assert_less_or_close_loglike(stats.invgauss, data, floc=rvs_loc, + fscale=self.rng.random(1)[0]) + + def test_fit_raise_errors(self): + assert_fit_warnings(stats.invgauss) + # FitDataError is raised when negative invalid data + with pytest.raises(FitDataError): + stats.invgauss.fit([1, 2, 3], floc=2) + + def test_cdf_sf(self): + # Regression tests for gh-13614. + # Ground truth from R's statmod library (pinvgauss), e.g. + # library(statmod) + # options(digits=15) + # mu = c(4.17022005e-04, 7.20324493e-03, 1.14374817e-06, + # 3.02332573e-03, 1.46755891e-03) + # print(pinvgauss(5, mu, 1)) + + # make sure a finite value is returned when mu is very small. see + # GH-13614 + mu = [4.17022005e-04, 7.20324493e-03, 1.14374817e-06, + 3.02332573e-03, 1.46755891e-03] + expected = [1, 1, 1, 1, 1] + actual = stats.invgauss.cdf(0.4, mu=mu) + assert_equal(expected, actual) + + # test if the function can distinguish small left/right tail + # probabilities from zero. + cdf_actual = stats.invgauss.cdf(0.001, mu=1.05) + assert_allclose(cdf_actual, 4.65246506892667e-219) + sf_actual = stats.invgauss.sf(110, mu=1.05) + assert_allclose(sf_actual, 4.12851625944048e-25) + + # test if x does not cause numerical issues when mu is very small + # and x is close to mu in value. + + # slightly smaller than mu + actual = stats.invgauss.cdf(0.00009, 0.0001) + assert_allclose(actual, 2.9458022894924e-26) + + # slightly bigger than mu + actual = stats.invgauss.cdf(0.000102, 0.0001) + assert_allclose(actual, 0.976445540507925) + + def test_logcdf_logsf(self): + # Regression tests for improvements made in gh-13616. + # Ground truth from R's statmod library (pinvgauss), e.g. + # library(statmod) + # options(digits=15) + # print(pinvgauss(0.001, 1.05, 1, log.p=TRUE, lower.tail=FALSE)) + + # test if logcdf and logsf can compute values too small to + # be represented on the unlogged scale. See: gh-13616 + logcdf = stats.invgauss.logcdf(0.0001, mu=1.05) + assert_allclose(logcdf, -5003.87872590367) + logcdf = stats.invgauss.logcdf(110, 1.05) + assert_allclose(logcdf, -4.12851625944087e-25) + logsf = stats.invgauss.logsf(0.001, mu=1.05) + assert_allclose(logsf, -4.65246506892676e-219) + logsf = stats.invgauss.logsf(110, 1.05) + assert_allclose(logsf, -56.1467092416426) + + # from mpmath import mp + # mp.dps = 100 + # mu = mp.mpf(1e-2) + # ref = (1/2 * mp.log(2 * mp.pi * mp.e * mu**3) + # - 3/2* mp.exp(2/mu) * mp.e1(2/mu)) + @pytest.mark.parametrize("mu, ref", [(2e-8, -25.172361826883957), + (1e-3, -8.943444010642972), + (1e-2, -5.4962796152622335), + (1e8, 3.3244822568873476), + (1e100, 3.32448280139689)]) + def test_entropy(self, mu, ref): + assert_allclose(stats.invgauss.entropy(mu), ref, rtol=5e-14) + + def test_mu_inf_gh13666(self): + # invgauss methods should return correct result when mu=inf + # invgauss as mu -> oo is invgamma with shape and scale 0.5; + # see gh-13666 and gh-22496 + dist = stats.invgauss(mu=np.inf) + dist0 = stats.invgamma(0.5, scale=0.5) + x, p = 1., 0.5 + assert_allclose(dist.logpdf(x), dist0.logpdf(x)) + assert_allclose(dist.pdf(x), dist0.pdf(x)) + assert_allclose(dist.logcdf(x), dist0.logcdf(x)) + assert_allclose(dist.cdf(x), dist0.cdf(x)) + assert_allclose(dist.logsf(x), dist0.logsf(x)) + assert_allclose(dist.sf(x), dist0.sf(x)) + assert_allclose(dist.ppf(p), dist0.ppf(p)) + assert_allclose(dist.isf(p), dist0.isf(p)) + + +class TestLandau: + @pytest.mark.parametrize('name', ['pdf', 'cdf', 'sf', 'ppf', 'isf']) + def test_landau_levy_agreement(self, name): + # Test PDF to confirm that this is the Landau distribution + # Test other methods with tighter tolerance than generic tests + # Levy entropy is slow and inaccurate, and RVS is tested generically + if name in {'ppf', 'isf'}: + x = np.linspace(0.1, 0.9, 5), + else: + x = np.linspace(-2, 5, 10), + + landau_method = getattr(stats.landau, name) + levy_method = getattr(stats.levy_stable, name) + res = landau_method(*x) + ref = levy_method(*x, 1, 1) + assert_allclose(res, ref, rtol=1e-14) + + def test_moments(self): + # I would test these against Levy above, but Levy says variance is infinite. + assert_equal(stats.landau.stats(moments='mvsk'), (np.nan,)*4) + assert_equal(stats.landau.moment(5), np.nan) + + +class TestLaplace: + @pytest.mark.parametrize("rvs_loc", [-5, 0, 1, 2]) + @pytest.mark.parametrize("rvs_scale", [1, 2, 3, 10]) + def test_fit(self, rvs_loc, rvs_scale): + # tests that various inputs follow expected behavior + # for a variety of `loc` and `scale`. + rng = np.random.RandomState(1234) + data = stats.laplace.rvs(size=100, loc=rvs_loc, scale=rvs_scale, + random_state=rng) + + # MLE estimates are given by + loc_mle = np.median(data) + scale_mle = np.sum(np.abs(data - loc_mle)) / len(data) + + # standard outputs should match analytical MLE formulas + loc, scale = stats.laplace.fit(data) + assert_allclose(loc, loc_mle, atol=1e-15, rtol=1e-15) + assert_allclose(scale, scale_mle, atol=1e-15, rtol=1e-15) + + # fixed parameter should use analytical formula for other + loc, scale = stats.laplace.fit(data, floc=loc_mle) + assert_allclose(scale, scale_mle, atol=1e-15, rtol=1e-15) + loc, scale = stats.laplace.fit(data, fscale=scale_mle) + assert_allclose(loc, loc_mle) + + # test with non-mle fixed parameter + # create scale with non-median loc + loc = rvs_loc * 2 + scale_mle = np.sum(np.abs(data - loc)) / len(data) + + # fixed loc to non median, scale should match + # scale calculation with modified loc + loc, scale = stats.laplace.fit(data, floc=loc) + assert_equal(scale_mle, scale) + + # fixed scale created with non median loc, + # loc output should still be the data median. + loc, scale = stats.laplace.fit(data, fscale=scale_mle) + assert_equal(loc_mle, loc) + + # error raised when both `floc` and `fscale` are fixed + assert_raises(RuntimeError, stats.laplace.fit, data, floc=loc_mle, + fscale=scale_mle) + + # error is raised with non-finite values + assert_raises(ValueError, stats.laplace.fit, [np.nan]) + assert_raises(ValueError, stats.laplace.fit, [np.inf]) + + @pytest.mark.parametrize("rvs_loc,rvs_scale", [(-5, 10), + (10, 5), + (0.5, 0.2)]) + def test_fit_MLE_comp_optimizer(self, rvs_loc, rvs_scale): + rng = np.random.RandomState(1234) + data = stats.laplace.rvs(size=1000, loc=rvs_loc, scale=rvs_scale, + random_state=rng) + + # the log-likelihood function for laplace is given by + def ll(loc, scale, data): + return -1 * (- (len(data)) * np.log(2*scale) - + (1/scale)*np.sum(np.abs(data - loc))) + + # test that the objective function result of the analytical MLEs is + # less than or equal to that of the numerically optimized estimate + loc, scale = stats.laplace.fit(data) + loc_opt, scale_opt = super(type(stats.laplace), + stats.laplace).fit(data) + ll_mle = ll(loc, scale, data) + ll_opt = ll(loc_opt, scale_opt, data) + assert ll_mle < ll_opt or np.allclose(ll_mle, ll_opt, + atol=1e-15, rtol=1e-15) + + def test_fit_simple_non_random_data(self): + data = np.array([1.0, 1.0, 3.0, 5.0, 8.0, 14.0]) + # with `floc` fixed to 6, scale should be 4. + loc, scale = stats.laplace.fit(data, floc=6) + assert_allclose(scale, 4, atol=1e-15, rtol=1e-15) + # with `fscale` fixed to 6, loc should be 4. + loc, scale = stats.laplace.fit(data, fscale=6) + assert_allclose(loc, 4, atol=1e-15, rtol=1e-15) + + def test_sf_cdf_extremes(self): + # These calculations should not generate warnings. + x = 1000 + p0 = stats.laplace.cdf(-x) + # The exact value is smaller than can be represented with + # 64 bit floating point, so the expected result is 0. + assert p0 == 0.0 + # The closest 64 bit floating point representation of the + # exact value is 1.0. + p1 = stats.laplace.cdf(x) + assert p1 == 1.0 + + p0 = stats.laplace.sf(x) + # The exact value is smaller than can be represented with + # 64 bit floating point, so the expected result is 0. + assert p0 == 0.0 + # The closest 64 bit floating point representation of the + # exact value is 1.0. + p1 = stats.laplace.sf(-x) + assert p1 == 1.0 + + def test_sf(self): + x = 200 + p = stats.laplace.sf(x) + assert_allclose(p, np.exp(-x)/2, rtol=1e-13) + + def test_isf(self): + p = 1e-25 + x = stats.laplace.isf(p) + assert_allclose(x, -np.log(2*p), rtol=1e-13) + + def test_logcdf_logsf(self): + x = 40 + # Reference value computed with mpmath. + ref = -2.1241771276457944e-18 + logcdf = stats.laplace.logcdf(x) + assert_allclose(logcdf, ref) + logsf = stats.laplace.logsf(-x) + assert_allclose(logsf, ref, rtol=5e-15) + + +class TestLogLaplace: + + def test_sf(self): + # reference values were computed via the reference distribution, e.g. + # mp.dps = 100; LogLaplace(c=c).sf(x). + c = np.array([2.0, 3.0, 5.0]) + x = np.array([1e-5, 1e10, 1e15]) + ref = [0.99999999995, 5e-31, 5e-76] + assert_allclose(stats.loglaplace.sf(x, c), ref, rtol=1e-15) + + def test_isf(self): + # reference values were computed via the reference distribution, e.g. + # mp.dps = 100; LogLaplace(c=c).isf(q). + c = 3.25 + q = [0.8, 0.1, 1e-10, 1e-20, 1e-40] + ref = [0.7543222539245642, 1.6408455124660906, 964.4916294395846, + 1151387.578354072, 1640845512466.0906] + assert_allclose(stats.loglaplace.isf(q, c), ref, rtol=1e-14) + + @pytest.mark.parametrize('r', [1, 2, 3, 4]) + def test_moments_stats(self, r): + mom = 'mvsk'[r - 1] + c = np.arange(0.5, r + 0.5, 0.5) + + # r-th non-central moment is infinite if |r| >= c. + assert_allclose(stats.loglaplace.moment(r, c), np.inf) + + # r-th non-central moment is non-finite (inf or nan) if r >= c. + assert not np.any(np.isfinite(stats.loglaplace.stats(c, moments=mom))) + + @pytest.mark.parametrize("c", [0.5, 1.0, 2.0]) + @pytest.mark.parametrize("loc, scale", [(-1.2, 3.45)]) + @pytest.mark.parametrize("fix_c", [True, False]) + @pytest.mark.parametrize("fix_scale", [True, False]) + def test_fit_analytic_mle(self, c, loc, scale, fix_c, fix_scale): + # Test that the analytical MLE produces no worse result than the + # generic (numerical) MLE. + + rng = np.random.default_rng(6762668991392531563) + data = stats.loglaplace.rvs(c, loc=loc, scale=scale, size=100, + random_state=rng) + + kwds = {'floc': loc} + if fix_c: + kwds['fc'] = c + if fix_scale: + kwds['fscale'] = scale + nfree = 3 - len(kwds) + + if nfree == 0: + error_msg = "All parameters fixed. There is nothing to optimize." + with pytest.raises((RuntimeError, ValueError), match=error_msg): + stats.loglaplace.fit(data, **kwds) + return + + _assert_less_or_close_loglike(stats.loglaplace, data, **kwds) + + +class TestPowerlaw: + + # In the following data, `sf` was computed with mpmath. + @pytest.mark.parametrize('x, a, sf', + [(0.25, 2.0, 0.9375), + (0.99609375, 1/256, 1.528855235208108e-05)]) + def test_sf(self, x, a, sf): + assert_allclose(stats.powerlaw.sf(x, a), sf, rtol=1e-15) + + @pytest.fixture(scope='function') + def rng(self): + return np.random.default_rng(1234) + + @pytest.mark.parametrize("rvs_shape", [.1, .5, .75, 1, 2]) + @pytest.mark.parametrize("rvs_loc", [-1, 0, 1]) + @pytest.mark.parametrize("rvs_scale", [.1, 1, 5]) + @pytest.mark.parametrize('fix_shape, fix_loc, fix_scale', + [p for p in product([True, False], repeat=3) + if False in p]) + def test_fit_MLE_comp_optimizer(self, rvs_shape, rvs_loc, rvs_scale, + fix_shape, fix_loc, fix_scale, rng): + data = stats.powerlaw.rvs(size=250, a=rvs_shape, loc=rvs_loc, + scale=rvs_scale, random_state=rng) + + kwds = dict() + if fix_shape: + kwds['f0'] = rvs_shape + if fix_loc: + kwds['floc'] = np.nextafter(data.min(), -np.inf) + if fix_scale: + kwds['fscale'] = rvs_scale + + # Numerical result may equal analytical result if some code path + # of the analytical routine makes use of numerical optimization. + _assert_less_or_close_loglike(stats.powerlaw, data, **kwds, + maybe_identical=True) + + def test_problem_case(self): + # An observed problem with the test method indicated that some fixed + # scale values could cause bad results, this is now corrected. + a = 2.50002862645130604506 + location = 0.0 + scale = 35.249023299873095 + + data = stats.powerlaw.rvs(a=a, loc=location, scale=scale, size=100, + random_state=np.random.default_rng(5)) + + kwds = {'fscale': np.ptp(data) * 2} + + _assert_less_or_close_loglike(stats.powerlaw, data, **kwds) + + def test_fit_warnings(self): + assert_fit_warnings(stats.powerlaw) + # test for error when `fscale + floc <= np.max(data)` is not satisfied + msg = r" Maximum likelihood estimation with 'powerlaw' requires" + with assert_raises(FitDataError, match=msg): + stats.powerlaw.fit([1, 2, 4], floc=0, fscale=3) + + # test for error when `data - floc >= 0` is not satisfied + msg = r" Maximum likelihood estimation with 'powerlaw' requires" + with assert_raises(FitDataError, match=msg): + stats.powerlaw.fit([1, 2, 4], floc=2) + + # test for fixed location not less than `min(data)`. + msg = r" Maximum likelihood estimation with 'powerlaw' requires" + with assert_raises(FitDataError, match=msg): + stats.powerlaw.fit([1, 2, 4], floc=1) + + # test for when fixed scale is less than or equal to range of data + msg = r"Negative or zero `fscale` is outside" + with assert_raises(ValueError, match=msg): + stats.powerlaw.fit([1, 2, 4], fscale=-3) + + # test for when fixed scale is less than or equal to range of data + msg = r"`fscale` must be greater than the range of data." + with assert_raises(ValueError, match=msg): + stats.powerlaw.fit([1, 2, 4], fscale=3) + + def test_minimum_data_zero_gh17801(self): + # gh-17801 reported an overflow error when the minimum value of the + # data is zero. Check that this problem is resolved. + data = [0, 1, 2, 2, 3, 3, 3, 3, 4, 4, 5, 6] + dist = stats.powerlaw + with np.errstate(over='ignore'): + _assert_less_or_close_loglike(dist, data) + + +class TestPowerLogNorm: + + # reference values were computed via mpmath + # from mpmath import mp + # mp.dps = 80 + # def powerlognorm_sf_mp(x, c, s): + # x = mp.mpf(x) + # c = mp.mpf(c) + # s = mp.mpf(s) + # return mp.ncdf(-mp.log(x) / s)**c + # + # def powerlognormal_cdf_mp(x, c, s): + # return mp.one - powerlognorm_sf_mp(x, c, s) + # + # x, c, s = 100, 20, 1 + # print(float(powerlognorm_sf_mp(x, c, s))) + + @pytest.mark.parametrize("x, c, s, ref", + [(100, 20, 1, 1.9057100820561928e-114), + (1e-3, 20, 1, 0.9999999999507617), + (1e-3, 0.02, 1, 0.9999999999999508), + (1e22, 0.02, 1, 6.50744044621611e-12)]) + def test_sf(self, x, c, s, ref): + assert_allclose(stats.powerlognorm.sf(x, c, s), ref, rtol=1e-13) + + # reference values were computed via mpmath using the survival + # function above (passing in `ref` and getting `q`). + @pytest.mark.parametrize("q, c, s, ref", + [(0.9999999587870905, 0.02, 1, 0.01), + (6.690376686108851e-233, 20, 1, 1000)]) + def test_isf(self, q, c, s, ref): + assert_allclose(stats.powerlognorm.isf(q, c, s), ref, rtol=5e-11) + + @pytest.mark.parametrize("x, c, s, ref", + [(1e25, 0.02, 1, 0.9999999999999963), + (1e-6, 0.02, 1, 2.054921078040843e-45), + (1e-6, 200, 1, 2.0549210780408428e-41), + (0.3, 200, 1, 0.9999999999713368)]) + def test_cdf(self, x, c, s, ref): + assert_allclose(stats.powerlognorm.cdf(x, c, s), ref, rtol=3e-14) + + # reference values were computed via mpmath + # from mpmath import mp + # mp.dps = 50 + # def powerlognorm_pdf_mpmath(x, c, s): + # x = mp.mpf(x) + # c = mp.mpf(c) + # s = mp.mpf(s) + # res = (c/(x * s) * mp.npdf(mp.log(x)/s) * + # mp.ncdf(-mp.log(x)/s)**(c - mp.one)) + # return float(res) + + @pytest.mark.parametrize("x, c, s, ref", + [(1e22, 0.02, 1, 6.5954987852335016e-34), + (1e20, 1e-3, 1, 1.588073750563988e-22), + (1e40, 1e-3, 1, 1.3179391812506349e-43)]) + def test_pdf(self, x, c, s, ref): + assert_allclose(stats.powerlognorm.pdf(x, c, s), ref, rtol=3e-12) + + +class TestPowerNorm: + + # survival function references were computed with mpmath via + # from mpmath import mp + # x = mp.mpf(x) + # c = mp.mpf(x) + # float(mp.ncdf(-x)**c) + + @pytest.mark.parametrize("x, c, ref", + [(9, 1, 1.1285884059538405e-19), + (20, 2, 7.582445786569958e-178), + (100, 0.02, 3.330957891903866e-44), + (200, 0.01, 1.3004759092324774e-87)]) + def test_sf(self, x, c, ref): + assert_allclose(stats.powernorm.sf(x, c), ref, rtol=1e-13) + + # inverse survival function references were computed with mpmath via + # from mpmath import mp + # def isf_mp(q, c): + # q = mp.mpf(q) + # c = mp.mpf(c) + # arg = q**(mp.one / c) + # return float(-mp.sqrt(2) * mp.erfinv(mp.mpf(2.) * arg - mp.one)) + + @pytest.mark.parametrize("q, c, ref", + [(1e-5, 20, -0.15690800666514138), + (0.99999, 100, -5.19933666203545), + (0.9999, 0.02, -2.576676052143387), + (5e-2, 0.02, 17.089518110222244), + (1e-18, 2, 5.9978070150076865), + (1e-50, 5, 6.361340902404057)]) + def test_isf(self, q, c, ref): + assert_allclose(stats.powernorm.isf(q, c), ref, rtol=5e-12) + + # CDF reference values were computed with mpmath via + # from mpmath import mp + # def cdf_mp(x, c): + # x = mp.mpf(x) + # c = mp.mpf(c) + # return float(mp.one - mp.ncdf(-x)**c) + + @pytest.mark.parametrize("x, c, ref", + [(-12, 9, 1.598833900869911e-32), + (2, 9, 0.9999999999999983), + (-20, 9, 2.4782617067456103e-88), + (-5, 0.02, 5.733032242841443e-09), + (-20, 0.02, 5.507248237212467e-91)]) + def test_cdf(self, x, c, ref): + assert_allclose(stats.powernorm.cdf(x, c), ref, rtol=5e-14) + + +class TestInvGamma: + def test_invgamma_inf_gh_1866(self): + # invgamma's moments are only finite for a>n + # specific numbers checked w/ boost 1.54 + with warnings.catch_warnings(): + warnings.simplefilter('error', RuntimeWarning) + mvsk = stats.invgamma.stats(a=19.31, moments='mvsk') + expected = [0.05461496450, 0.0001723162534, 1.020362676, + 2.055616582] + assert_allclose(mvsk, expected) + + a = [1.1, 3.1, 5.6] + mvsk = stats.invgamma.stats(a=a, moments='mvsk') + expected = ([10., 0.476190476, 0.2173913043], # mmm + [np.inf, 0.2061430632, 0.01312749422], # vvv + [np.nan, 41.95235392, 2.919025532], # sss + [np.nan, np.nan, 24.51923076]) # kkk + for x, y in zip(mvsk, expected): + assert_almost_equal(x, y) + + def test_cdf_ppf(self): + # gh-6245 + x = np.logspace(-2.6, 0) + y = stats.invgamma.cdf(x, 1) + xx = stats.invgamma.ppf(y, 1) + assert_allclose(x, xx) + + def test_sf_isf(self): + # gh-6245 + if sys.maxsize > 2**32: + x = np.logspace(2, 100) + else: + # Invgamme roundtrip on 32-bit systems has relative accuracy + # ~1e-15 until x=1e+15, and becomes inf above x=1e+18 + x = np.logspace(2, 18) + + y = stats.invgamma.sf(x, 1) + xx = stats.invgamma.isf(y, 1) + assert_allclose(x, xx, rtol=1.0) + + def test_logcdf(self): + x = 1e7 + a = 2.25 + # Reference value computed with mpmath. + ref = -6.97567687425534e-17 + logcdf = stats.invgamma.logcdf(x, a) + assert_allclose(logcdf, ref, rtol=5e-15) + + def test_logsf(self): + x = 0.01 + a = 3.5 + # Reference value computed with mpmath. + ref = -1.147781224014262e-39 + logsf = stats.invgamma.logsf(x, a) + assert_allclose(logsf, ref, rtol=5e-15) + + @pytest.mark.parametrize("a, ref", + [(100000000.0, -26.21208257605721), + (1e+100, -343.9688254159022)]) + def test_large_entropy(self, a, ref): + # The reference values were calculated with mpmath: + # from mpmath import mp + # mp.dps = 500 + + # def invgamma_entropy(a): + # a = mp.mpf(a) + # h = a + mp.loggamma(a) - (mp.one + a) * mp.digamma(a) + # return float(h) + assert_allclose(stats.invgamma.entropy(a), ref, rtol=1e-15) + + +class TestF: + def test_endpoints(self): + # Compute the pdf at the left endpoint dst.a. + data = [[stats.f, (2, 1), 1.0]] + for _f, _args, _correct in data: + ans = _f.pdf(_f.a, *_args) + + ans = [_f.pdf(_f.a, *_args) for _f, _args, _ in data] + correct = [_correct_ for _f, _args, _correct_ in data] + assert_array_almost_equal(ans, correct) + + def test_f_moments(self): + # n-th moment of F distributions is only finite for n < dfd / 2 + m, v, s, k = stats.f.stats(11, 6.5, moments='mvsk') + assert_(np.isfinite(m)) + assert_(np.isfinite(v)) + assert_(np.isfinite(s)) + assert_(not np.isfinite(k)) + + def test_moments_warnings(self): + # no warnings should be generated for dfd = 2, 4, 6, 8 (div by zero) + with warnings.catch_warnings(): + warnings.simplefilter('error', RuntimeWarning) + stats.f.stats(dfn=[11]*4, dfd=[2, 4, 6, 8], moments='mvsk') + + def test_stats_broadcast(self): + dfn = np.array([[3], [11]]) + dfd = np.array([11, 12]) + m, v, s, k = stats.f.stats(dfn=dfn, dfd=dfd, moments='mvsk') + m2 = [dfd / (dfd - 2)]*2 + assert_allclose(m, m2) + v2 = 2 * dfd**2 * (dfn + dfd - 2) / dfn / (dfd - 2)**2 / (dfd - 4) + assert_allclose(v, v2) + s2 = ((2*dfn + dfd - 2) * np.sqrt(8*(dfd - 4)) / + ((dfd - 6) * np.sqrt(dfn*(dfn + dfd - 2)))) + assert_allclose(s, s2) + k2num = 12 * (dfn * (5*dfd - 22) * (dfn + dfd - 2) + + (dfd - 4) * (dfd - 2)**2) + k2den = dfn * (dfd - 6) * (dfd - 8) * (dfn + dfd - 2) + k2 = k2num / k2den + assert_allclose(k, k2) + + +class TestStudentT: + def test_rvgeneric_std(self): + # Regression test for #1191 + assert_array_almost_equal(stats.t.std([5, 6]), [1.29099445, 1.22474487]) + + def test_moments_t(self): + # regression test for #8786 + assert_equal(stats.t.stats(df=1, moments='mvsk'), + (np.inf, np.nan, np.nan, np.nan)) + assert_equal(stats.t.stats(df=1.01, moments='mvsk'), + (0.0, np.inf, np.nan, np.nan)) + assert_equal(stats.t.stats(df=2, moments='mvsk'), + (0.0, np.inf, np.nan, np.nan)) + assert_equal(stats.t.stats(df=2.01, moments='mvsk'), + (0.0, 2.01/(2.01-2.0), np.nan, np.inf)) + assert_equal(stats.t.stats(df=3, moments='sk'), (np.nan, np.inf)) + assert_equal(stats.t.stats(df=3.01, moments='sk'), (0.0, np.inf)) + assert_equal(stats.t.stats(df=4, moments='sk'), (0.0, np.inf)) + assert_allclose(stats.t.stats(df=4.01, moments='sk'), (0.0, 6.0/(4.01 - 4.0)), + rtol=1e-14) + + def test_t_entropy(self): + df = [1, 2, 25, 100] + # Expected values were computed with mpmath. + expected = [2.5310242469692907, 1.9602792291600821, + 1.459327578078393, 1.4289633653182439] + assert_allclose(stats.t.entropy(df), expected, rtol=1e-13) + + @pytest.mark.parametrize("v, ref", + [(100, 1.4289633653182439), + (1e+100, 1.4189385332046727)]) + def test_t_extreme_entropy(self, v, ref): + # Reference values were calculated with mpmath: + # from mpmath import mp + # mp.dps = 500 + # + # def t_entropy(v): + # v = mp.mpf(v) + # C = (v + mp.one) / 2 + # A = C * (mp.digamma(C) - mp.digamma(v / 2)) + # B = 0.5 * mp.log(v) + mp.log(mp.beta(v / 2, mp.one / 2)) + # h = A + B + # return float(h) + assert_allclose(stats.t.entropy(v), ref, rtol=1e-14) + + @pytest.mark.parametrize("methname", ["pdf", "logpdf", "cdf", + "ppf", "sf", "isf"]) + @pytest.mark.parametrize("df_infmask", [[0, 0], [1, 1], [0, 1], + [[0, 1, 0], [1, 1, 1]], + [[1, 0], [0, 1]], + [[0], [1]]]) + def test_t_inf_df(self, methname, df_infmask): + df_infmask = np.asarray(df_infmask, dtype=bool) + rng = np.random.default_rng(5442451539) + df = rng.uniform(0, 10, size=df_infmask.shape) + x = rng.standard_normal(df_infmask.shape) + df[df_infmask] = np.inf + t_dist = stats.t(df=df, loc=3, scale=1) + t_dist_ref = stats.t(df=df[~df_infmask], loc=3, scale=1) + norm_dist = stats.norm(loc=3, scale=1) + t_meth = getattr(t_dist, methname) + t_meth_ref = getattr(t_dist_ref, methname) + norm_meth = getattr(norm_dist, methname) + res = t_meth(x) + assert_allclose(res[df_infmask], norm_meth(x[df_infmask]), rtol=5e-15) + assert_equal(res[~df_infmask], t_meth_ref(x[~df_infmask])) + + @pytest.mark.parametrize("df_infmask", [[0, 0], [1, 1], [0, 1], + [[0, 1, 0], [1, 1, 1]], + [[1, 0], [0, 1]], + [[0], [1]]]) + def test_t_inf_df_stats_entropy(self, df_infmask): + df_infmask = np.asarray(df_infmask, dtype=bool) + rng = np.random.default_rng(5442451539) + df = rng.uniform(0, 10, size=df_infmask.shape) + df[df_infmask] = np.inf + res = stats.t.stats(df=df, loc=3, scale=1, moments='mvsk') + res_ex_inf = stats.norm.stats(loc=3, scale=1, moments='mvsk') + res_ex_noinf = stats.t.stats(df=df[~df_infmask], loc=3, scale=1, + moments='mvsk') + for i in range(4): + assert_equal(res[i][df_infmask], res_ex_inf[i]) + assert_equal(res[i][~df_infmask], res_ex_noinf[i]) + + res = stats.t.entropy(df=df, loc=3, scale=1) + res_ex_inf = stats.norm.entropy(loc=3, scale=1) + res_ex_noinf = stats.t.entropy(df=df[~df_infmask], loc=3, scale=1) + assert_equal(res[df_infmask], res_ex_inf) + assert_equal(res[~df_infmask], res_ex_noinf) + + def test_logpdf_pdf(self): + # reference values were computed via the reference distribution, e.g. + # mp.dps = 500; StudentT(df=df).logpdf(x), StudentT(df=df).pdf(x) + x = [1, 1e3, 10, 1] + df = [1e100, 1e50, 1e20, 1] + logpdf_ref = [-1.4189385332046727, -500000.9189385332, + -50.918938533204674, -1.8378770664093456] + pdf_ref = [0.24197072451914334, 0, + 7.69459862670642e-23, 0.15915494309189535] + assert_allclose(stats.t.logpdf(x, df), logpdf_ref, rtol=1e-14) + assert_allclose(stats.t.pdf(x, df), pdf_ref, rtol=1e-14) + + # Reference values were computed with mpmath, and double-checked with + # Wolfram Alpha. + @pytest.mark.parametrize('x, df, ref', + [(-75.0, 15, -46.76036184546812), + (0, 15, -0.6931471805599453), + (75.0, 15, -4.9230344937641665e-21)]) + def test_logcdf_logsf(self, x, df, ref): + logcdf = stats.t.logcdf(x, df) + assert_allclose(logcdf, ref, rtol=5e-15) + # The reference value is logcdf(x, df) == logsf(-x, df). + logsf = stats.t.logsf(-x, df) + assert_allclose(logsf, ref, rtol=5e-15) + + +class TestRvDiscrete: + def setup_method(self): + self.rng = np.random.default_rng(333348228) + + def test_rvs(self): + states = [-1, 0, 1, 2, 3, 4] + probability = [0.0, 0.3, 0.4, 0.0, 0.3, 0.0] + samples = 1000 + r = stats.rv_discrete(name='sample', values=(states, probability)) + x = r.rvs(size=samples, random_state=self.rng) + assert isinstance(x, np.ndarray) + + for s, p in zip(states, probability): + assert abs(sum(x == s)/float(samples) - p) < 0.05 + + x = r.rvs(random_state=self.rng) + assert np.issubdtype(type(x), np.integer) + + def test_entropy(self): + # Basic tests of entropy. + pvals = np.array([0.25, 0.45, 0.3]) + p = stats.rv_discrete(values=([0, 1, 2], pvals)) + expected_h = -sum(xlogy(pvals, pvals)) + h = p.entropy() + assert_allclose(h, expected_h) + + p = stats.rv_discrete(values=([0, 1, 2], [1.0, 0, 0])) + h = p.entropy() + assert_equal(h, 0.0) + + def test_pmf(self): + xk = [1, 2, 4] + pk = [0.5, 0.3, 0.2] + rv = stats.rv_discrete(values=(xk, pk)) + + x = [[1., 4.], + [3., 2]] + assert_allclose(rv.pmf(x), + [[0.5, 0.2], + [0., 0.3]], atol=1e-14) + + def test_cdf(self): + xk = [1, 2, 4] + pk = [0.5, 0.3, 0.2] + rv = stats.rv_discrete(values=(xk, pk)) + + x_values = [-2, 1., 1.1, 1.5, 2.0, 3.0, 4, 5] + expected = [0, 0.5, 0.5, 0.5, 0.8, 0.8, 1, 1] + assert_allclose(rv.cdf(x_values), expected, atol=1e-14) + + # also check scalar arguments + assert_allclose([rv.cdf(xx) for xx in x_values], + expected, atol=1e-14) + + def test_ppf(self): + xk = [1, 2, 4] + pk = [0.5, 0.3, 0.2] + rv = stats.rv_discrete(values=(xk, pk)) + + q_values = [0.1, 0.5, 0.6, 0.8, 0.9, 1.] + expected = [1, 1, 2, 2, 4, 4] + assert_allclose(rv.ppf(q_values), expected, atol=1e-14) + + # also check scalar arguments + assert_allclose([rv.ppf(q) for q in q_values], + expected, atol=1e-14) + + def test_cdf_ppf_next(self): + # copied and special cased from test_discrete_basic + vals = ([1, 2, 4, 7, 8], [0.1, 0.2, 0.3, 0.3, 0.1]) + rv = stats.rv_discrete(values=vals) + + assert_array_equal(rv.ppf(rv.cdf(rv.xk[:-1]) + 1e-8), + rv.xk[1:]) + + def test_multidimension(self): + xk = np.arange(12).reshape((3, 4)) + pk = np.array([[0.1, 0.1, 0.15, 0.05], + [0.1, 0.1, 0.05, 0.05], + [0.1, 0.1, 0.05, 0.05]]) + rv = stats.rv_discrete(values=(xk, pk)) + + assert_allclose(rv.expect(), np.sum(rv.xk * rv.pk), atol=1e-14) + + def test_bad_input(self): + xk = [1, 2, 3] + pk = [0.5, 0.5] + assert_raises(ValueError, stats.rv_discrete, **dict(values=(xk, pk))) + + pk = [1, 2, 3] + assert_raises(ValueError, stats.rv_discrete, **dict(values=(xk, pk))) + + xk = [1, 2, 3] + pk = [0.5, 1.2, -0.7] + assert_raises(ValueError, stats.rv_discrete, **dict(values=(xk, pk))) + + xk = [1, 2, 3, 4, 5] + pk = [0.3, 0.3, 0.3, 0.3, -0.2] + assert_raises(ValueError, stats.rv_discrete, **dict(values=(xk, pk))) + + xk = [1, 1] + pk = [0.5, 0.5] + assert_raises(ValueError, stats.rv_discrete, **dict(values=(xk, pk))) + + def test_shape_rv_sample(self): + # tests added for gh-9565 + + # mismatch of 2d inputs + xk, pk = np.arange(4).reshape((2, 2)), np.full((2, 3), 1/6) + assert_raises(ValueError, stats.rv_discrete, **dict(values=(xk, pk))) + + # same number of elements, but shapes not compatible + xk, pk = np.arange(6).reshape((3, 2)), np.full((2, 3), 1/6) + assert_raises(ValueError, stats.rv_discrete, **dict(values=(xk, pk))) + + # same shapes => no error + xk, pk = np.arange(6).reshape((3, 2)), np.full((3, 2), 1/6) + assert_equal(stats.rv_discrete(values=(xk, pk)).pmf(0), 1/6) + + def test_expect1(self): + xk = [1, 2, 4, 6, 7, 11] + pk = [0.1, 0.2, 0.2, 0.2, 0.2, 0.1] + rv = stats.rv_discrete(values=(xk, pk)) + + assert_allclose(rv.expect(), np.sum(rv.xk * rv.pk), atol=1e-14) + + def test_expect2(self): + # rv_sample should override _expect. Bug report from + # https://stackoverflow.com/questions/63199792 + y = [200.0, 300.0, 400.0, 500.0, 600.0, 700.0, 800.0, 900.0, 1000.0, + 1100.0, 1200.0, 1300.0, 1400.0, 1500.0, 1600.0, 1700.0, 1800.0, + 1900.0, 2000.0, 2100.0, 2200.0, 2300.0, 2400.0, 2500.0, 2600.0, + 2700.0, 2800.0, 2900.0, 3000.0, 3100.0, 3200.0, 3300.0, 3400.0, + 3500.0, 3600.0, 3700.0, 3800.0, 3900.0, 4000.0, 4100.0, 4200.0, + 4300.0, 4400.0, 4500.0, 4600.0, 4700.0, 4800.0] + + py = [0.0004, 0.0, 0.0033, 0.006500000000000001, 0.0, 0.0, + 0.004399999999999999, 0.6862, 0.0, 0.0, 0.0, + 0.00019999999999997797, 0.0006000000000000449, + 0.024499999999999966, 0.006400000000000072, + 0.0043999999999999595, 0.019499999999999962, + 0.03770000000000007, 0.01759999999999995, 0.015199999999999991, + 0.018100000000000005, 0.04500000000000004, 0.0025999999999999357, + 0.0, 0.0041000000000001036, 0.005999999999999894, + 0.0042000000000000925, 0.0050000000000000044, + 0.0041999999999999815, 0.0004999999999999449, + 0.009199999999999986, 0.008200000000000096, + 0.0, 0.0, 0.0046999999999999265, 0.0019000000000000128, + 0.0006000000000000449, 0.02510000000000001, 0.0, + 0.007199999999999984, 0.0, 0.012699999999999934, 0.0, 0.0, + 0.008199999999999985, 0.005600000000000049, 0.0] + + rv = stats.rv_discrete(values=(y, py)) + + # check the mean + assert_allclose(rv.expect(), rv.mean(), atol=1e-14) + assert_allclose(rv.expect(), + sum(v * w for v, w in zip(y, py)), atol=1e-14) + + # also check the second moment + assert_allclose(rv.expect(lambda x: x**2), + sum(v**2 * w for v, w in zip(y, py)), atol=1e-14) + + +class TestSkewCauchy: + def test_cauchy(self): + x = np.linspace(-5, 5, 100) + assert_array_almost_equal(stats.skewcauchy.pdf(x, a=0), + stats.cauchy.pdf(x)) + assert_array_almost_equal(stats.skewcauchy.cdf(x, a=0), + stats.cauchy.cdf(x)) + assert_array_almost_equal(stats.skewcauchy.ppf(x, a=0), + stats.cauchy.ppf(x)) + + def test_skewcauchy_R(self): + # options(digits=16) + # library(sgt) + # # lmbda, x contain the values generated for a, x below + # lmbda <- c(0.0976270078546495, 0.430378732744839, 0.2055267521432877, + # 0.0897663659937937, -0.15269040132219, 0.2917882261333122, + # -0.12482557747462, 0.7835460015641595, 0.9273255210020589, + # -0.2331169623484446) + # x <- c(2.917250380826646, 0.2889491975290444, 0.6804456109393229, + # 4.25596638292661, -4.289639418021131, -4.1287070029845925, + # -4.797816025596743, 3.32619845547938, 2.7815675094985046, + # 3.700121482468191) + # pdf = dsgt(x, mu=0, lambda=lambda, sigma=1, q=1/2, mean.cent=FALSE, + # var.adj = sqrt(2)) + # cdf = psgt(x, mu=0, lambda=lambda, sigma=1, q=1/2, mean.cent=FALSE, + # var.adj = sqrt(2)) + # qsgt(cdf, mu=0, lambda=lambda, sigma=1, q=1/2, mean.cent=FALSE, + # var.adj = sqrt(2)) + + rng = np.random.RandomState(0) + a = rng.rand(10) * 2 - 1 + x = rng.rand(10) * 10 - 5 + pdf = [0.039473975217333909, 0.305829714049903223, 0.24140158118994162, + 0.019585772402693054, 0.021436553695989482, 0.00909817103867518, + 0.01658423410016873, 0.071083288030394126, 0.103250045941454524, + 0.013110230778426242] + cdf = [0.87426677718213752, 0.37556468910780882, 0.59442096496538066, + 0.91304659850890202, 0.09631964100300605, 0.03829624330921733, + 0.08245240578402535, 0.72057062945510386, 0.62826415852515449, + 0.95011308463898292] + assert_allclose(stats.skewcauchy.pdf(x, a), pdf) + assert_allclose(stats.skewcauchy.cdf(x, a), cdf) + assert_allclose(stats.skewcauchy.ppf(cdf, a), x) + + +class TestJFSkewT: + def test_compare_t(self): + # Verify that jf_skew_t with a=b recovers the t distribution with 2a + # degrees of freedom + a = b = 5 + df = a * 2 + x = [-1.0, 0.0, 1.0, 2.0] + q = [0.0, 0.1, 0.25, 0.75, 0.90, 1.0] + + jf = stats.jf_skew_t(a, b) + t = stats.t(df) + + assert_allclose(jf.pdf(x), t.pdf(x)) + assert_allclose(jf.cdf(x), t.cdf(x)) + assert_allclose(jf.ppf(q), t.ppf(q)) + assert_allclose(jf.stats('mvsk'), t.stats('mvsk')) + + @pytest.fixture + def gamlss_pdf_data(self): + """Sample data points computed using the `ST5` distribution from the + GAMLSS package in R. The pdf has been calculated for (a,b)=(2,3), + (a,b)=(8,4), and (a,b)=(12,13) for x in `np.linspace(-10, 10, 41)`. + + N.B. the `ST5` distribution in R uses an alternative parameterization + in terms of nu and tau, where: + - nu = (a - b) / (a * b * (a + b)) ** 0.5 + - tau = 2 / (a + b) + """ + data = np.load( + Path(__file__).parent / "data/jf_skew_t_gamlss_pdf_data.npy" + ) + return np.rec.fromarrays(data, names="x,pdf,a,b") + + @pytest.mark.parametrize("a,b", [(2, 3), (8, 4), (12, 13)]) + def test_compare_with_gamlss_r(self, gamlss_pdf_data, a, b): + """Compare the pdf with a table of reference values. The table of + reference values was produced using R, where the Jones and Faddy skew + t distribution is available in the GAMLSS package as `ST5`. + """ + data = gamlss_pdf_data[ + (gamlss_pdf_data["a"] == a) & (gamlss_pdf_data["b"] == b) + ] + x, pdf = data["x"], data["pdf"] + assert_allclose(pdf, stats.jf_skew_t(a, b).pdf(x), rtol=1e-12) + + +# Test data for TestSkewNorm.test_noncentral_moments() +# The expected noncentral moments were computed by Wolfram Alpha. +# In Wolfram Alpha, enter +# SkewNormalDistribution[0, 1, a] moment +# with `a` replaced by the desired shape parameter. In the results, there +# should be a table of the first four moments. Click on "More" to get more +# moments. The expected moments start with the first moment (order = 1). +_skewnorm_noncentral_moments = [ + (2, [2*np.sqrt(2/(5*np.pi)), + 1, + 22/5*np.sqrt(2/(5*np.pi)), + 3, + 446/25*np.sqrt(2/(5*np.pi)), + 15, + 2682/25*np.sqrt(2/(5*np.pi)), + 105, + 107322/125*np.sqrt(2/(5*np.pi))]), + (0.1, [np.sqrt(2/(101*np.pi)), + 1, + 302/101*np.sqrt(2/(101*np.pi)), + 3, + (152008*np.sqrt(2/(101*np.pi)))/10201, + 15, + (107116848*np.sqrt(2/(101*np.pi)))/1030301, + 105, + (97050413184*np.sqrt(2/(101*np.pi)))/104060401]), + (-3, [-3/np.sqrt(5*np.pi), + 1, + -63/(10*np.sqrt(5*np.pi)), + 3, + -2529/(100*np.sqrt(5*np.pi)), + 15, + -30357/(200*np.sqrt(5*np.pi)), + 105, + -2428623/(2000*np.sqrt(5*np.pi)), + 945, + -242862867/(20000*np.sqrt(5*np.pi)), + 10395, + -29143550277/(200000*np.sqrt(5*np.pi)), + 135135]), +] + + +class TestSkewNorm: + def setup_method(self): + self.rng = check_random_state(1234) + + def test_normal(self): + # When the skewness is 0 the distribution is normal + x = np.linspace(-5, 5, 100) + assert_array_almost_equal(stats.skewnorm.pdf(x, a=0), + stats.norm.pdf(x)) + + def test_rvs(self): + rng = check_random_state(1234) + shape = (3, 4, 5) + x = stats.skewnorm.rvs(a=0.75, size=shape, random_state=rng) + assert_equal(shape, x.shape) + + x = stats.skewnorm.rvs(a=-3, size=shape, random_state=rng) + assert_equal(shape, x.shape) + + def test_moments(self): + rng = check_random_state(1234) + X = stats.skewnorm.rvs(a=4, size=int(1e6), loc=5, scale=2, + random_state=rng) + expected = [np.mean(X), np.var(X), stats.skew(X), stats.kurtosis(X)] + computed = stats.skewnorm.stats(a=4, loc=5, scale=2, moments='mvsk') + assert_array_almost_equal(computed, expected, decimal=2) + + X = stats.skewnorm.rvs(a=-4, size=int(1e6), loc=5, scale=2, + random_state=rng) + expected = [np.mean(X), np.var(X), stats.skew(X), stats.kurtosis(X)] + computed = stats.skewnorm.stats(a=-4, loc=5, scale=2, moments='mvsk') + assert_array_almost_equal(computed, expected, decimal=2) + + def test_pdf_large_x(self): + # Triples are [x, a, logpdf(x, a)]. These values were computed + # using Log[PDF[SkewNormalDistribution[0, 1, a], x]] in Wolfram Alpha. + logpdfvals = [ + [40, -1, -1604.834233366398515598970], + [40, -1/2, -1004.142946723741991369168], + [40, 0, -800.9189385332046727417803], + [40, 1/2, -800.2257913526447274323631], + [-40, -1/2, -800.2257913526447274323631], + [-2, 1e7, -2.000000000000199559727173e14], + [2, -1e7, -2.000000000000199559727173e14], + ] + for x, a, logpdfval in logpdfvals: + logp = stats.skewnorm.logpdf(x, a) + assert_allclose(logp, logpdfval, rtol=1e-8) + + def test_cdf_large_x(self): + # Regression test for gh-7746. + # The x values are large enough that the closest 64 bit floating + # point representation of the exact CDF is 1.0. + p = stats.skewnorm.cdf([10, 20, 30], -1) + assert_allclose(p, np.ones(3), rtol=1e-14) + p = stats.skewnorm.cdf(25, 2.5) + assert_allclose(p, 1.0, rtol=1e-14) + + def test_cdf_sf_small_values(self): + # Triples are [x, a, cdf(x, a)]. These values were computed + # using CDF[SkewNormalDistribution[0, 1, a], x] in Wolfram Alpha. + cdfvals = [ + [-8, 1, 3.870035046664392611e-31], + [-4, 2, 8.1298399188811398e-21], + [-2, 5, 1.55326826787106273e-26], + [-9, -1, 2.257176811907681295e-19], + [-10, -4, 1.523970604832105213e-23], + ] + for x, a, cdfval in cdfvals: + p = stats.skewnorm.cdf(x, a) + assert_allclose(p, cdfval, rtol=1e-8) + # For the skew normal distribution, sf(-x, -a) = cdf(x, a). + p = stats.skewnorm.sf(-x, -a) + assert_allclose(p, cdfval, rtol=1e-8) + + @pytest.mark.parametrize('a, moments', _skewnorm_noncentral_moments) + def test_noncentral_moments(self, a, moments): + for order, expected in enumerate(moments, start=1): + mom = stats.skewnorm.moment(order, a) + assert_allclose(mom, expected, rtol=1e-14) + + def test_fit(self): + rng = np.random.default_rng(4609813989115202851) + + a, loc, scale = -2, 3.5, 0.5 # arbitrary, valid parameters + dist = stats.skewnorm(a, loc, scale) + rvs = dist.rvs(size=100, random_state=rng) + + # test that MLE still honors guesses and fixed parameters + a2, loc2, scale2 = stats.skewnorm.fit(rvs, -1.5, floc=3) + a3, loc3, scale3 = stats.skewnorm.fit(rvs, -1.6, floc=3) + assert loc2 == loc3 == 3 # fixed parameter is respected + assert a2 != a3 # different guess -> (slightly) different outcome + # quality of fit is tested elsewhere + + # test that MoM honors fixed parameters, accepts (but ignores) guesses + a4, loc4, scale4 = stats.skewnorm.fit(rvs, 3, fscale=3, method='mm') + assert scale4 == 3 + # because scale was fixed, only the mean and skewness will be matched + dist4 = stats.skewnorm(a4, loc4, scale4) + res = dist4.stats(moments='ms') + ref = np.mean(rvs), stats.skew(rvs) + assert_allclose(res, ref) + + # Test behavior when skew of data is beyond maximum of skewnorm + rvs2 = stats.pareto.rvs(1, size=100, random_state=rng) + + # MLE still works + res = stats.skewnorm.fit(rvs2) + assert np.all(np.isfinite(res)) + + # MoM fits variance and skewness + a5, loc5, scale5 = stats.skewnorm.fit(rvs2, method='mm') + assert np.isinf(a5) + # distribution infrastructure doesn't allow infinite shape parameters + # into _stats; it just bypasses it and produces NaNs. Calculate + # moments manually. + m, v = np.mean(rvs2), np.var(rvs2) + assert_allclose(m, loc5 + scale5 * np.sqrt(2/np.pi)) + assert_allclose(v, scale5**2 * (1 - 2 / np.pi)) + + # test that MLE and MoM behave as expected under sign changes + a6p, loc6p, scale6p = stats.skewnorm.fit(rvs, method='mle') + a6m, loc6m, scale6m = stats.skewnorm.fit(-rvs, method='mle') + assert_allclose([a6m, loc6m, scale6m], [-a6p, -loc6p, scale6p]) + a7p, loc7p, scale7p = stats.skewnorm.fit(rvs, method='mm') + a7m, loc7m, scale7m = stats.skewnorm.fit(-rvs, method='mm') + assert_allclose([a7m, loc7m, scale7m], [-a7p, -loc7p, scale7p]) + + def test_fit_gh19332(self): + # When the skewness of the data was high, `skewnorm.fit` fell back on + # generic `fit` behavior with a bad guess of the skewness parameter. + # Test that this is improved; `skewnorm.fit` is now better at finding + # the global optimum when the sample is highly skewed. See gh-19332. + x = np.array([-5, -1, 1 / 100_000] + 12 * [1] + [5]) + + params = stats.skewnorm.fit(x) + res = stats.skewnorm.nnlf(params, x) + + # Compare overridden fit against generic fit. + # res should be about 32.01, and generic fit is worse at 32.64. + # In case the generic fit improves, remove this assertion (see gh-19333). + params_super = stats.skewnorm.fit(x, superfit=True) + ref = stats.skewnorm.nnlf(params_super, x) + assert res < ref - 0.5 + + # Compare overridden fit against stats.fit + rng = np.random.default_rng(9842356982345693637) + bounds = {'a': (-5, 5), 'loc': (-10, 10), 'scale': (1e-16, 10)} + + def optimizer(fun, bounds): + return differential_evolution(fun, bounds, rng=rng) + + fit_result = stats.fit(stats.skewnorm, x, bounds, optimizer=optimizer) + np.testing.assert_allclose(params, fit_result.params, rtol=1e-4) + + def test_ppf(self): + # gh-20124 reported that Boost's ppf was wrong for high skewness + # Reference value was calculated using + # N[InverseCDF[SkewNormalDistribution[0, 1, 500], 1/100], 14] in Wolfram Alpha. + assert_allclose(stats.skewnorm.ppf(0.01, 500), 0.012533469508013, rtol=1e-13) + + +class TestExpon: + def test_zero(self): + assert_equal(stats.expon.pdf(0), 1) + + def test_tail(self): # Regression test for ticket 807 + assert_equal(stats.expon.cdf(1e-18), 1e-18) + assert_equal(stats.expon.isf(stats.expon.sf(40)), 40) + + def test_nan_raises_error(self): + # see gh-issue 10300 + x = np.array([1.6483, 2.7169, 2.4667, 1.1791, 3.5433, np.nan]) + assert_raises(ValueError, stats.expon.fit, x) + + def test_inf_raises_error(self): + # see gh-issue 10300 + x = np.array([1.6483, 2.7169, 2.4667, 1.1791, 3.5433, np.inf]) + assert_raises(ValueError, stats.expon.fit, x) + + +class TestNorm: + def test_nan_raises_error(self): + # see gh-issue 10300 + x = np.array([1.6483, 2.7169, 2.4667, 1.1791, 3.5433, np.nan]) + assert_raises(ValueError, stats.norm.fit, x) + + def test_inf_raises_error(self): + # see gh-issue 10300 + x = np.array([1.6483, 2.7169, 2.4667, 1.1791, 3.5433, np.inf]) + assert_raises(ValueError, stats.norm.fit, x) + + def test_bad_keyword_arg(self): + x = [1, 2, 3] + assert_raises(TypeError, stats.norm.fit, x, plate="shrimp") + + @pytest.mark.parametrize('loc', [0, 1]) + def test_delta_cdf(self, loc): + # The expected value is computed with mpmath: + # >>> import mpmath + # >>> mpmath.mp.dps = 60 + # >>> float(mpmath.ncdf(12) - mpmath.ncdf(11)) + # 1.910641809677555e-28 + expected = 1.910641809677555e-28 + delta = stats.norm._delta_cdf(11+loc, 12+loc, loc=loc) + assert_allclose(delta, expected, rtol=1e-13) + delta = stats.norm._delta_cdf(-(12+loc), -(11+loc), loc=-loc) + assert_allclose(delta, expected, rtol=1e-13) + + +class TestUniform: + """gh-10300""" + def test_nan_raises_error(self): + # see gh-issue 10300 + x = np.array([1.6483, 2.7169, 2.4667, 1.1791, 3.5433, np.nan]) + assert_raises(ValueError, stats.uniform.fit, x) + + def test_inf_raises_error(self): + # see gh-issue 10300 + x = np.array([1.6483, 2.7169, 2.4667, 1.1791, 3.5433, np.inf]) + assert_raises(ValueError, stats.uniform.fit, x) + + +class TestExponNorm: + def test_moments(self): + # Some moment test cases based on non-loc/scaled formula + def get_moms(lam, sig, mu): + # See wikipedia for these formulae + # where it is listed as an exponentially modified gaussian + opK2 = 1.0 + 1 / (lam*sig)**2 + exp_skew = 2 / (lam * sig)**3 * opK2**(-1.5) + exp_kurt = 6.0 * (1 + (lam * sig)**2)**(-2) + return [mu + 1/lam, sig*sig + 1.0/(lam*lam), exp_skew, exp_kurt] + + mu, sig, lam = 0, 1, 1 + K = 1.0 / (lam * sig) + sts = stats.exponnorm.stats(K, loc=mu, scale=sig, moments='mvsk') + assert_almost_equal(sts, get_moms(lam, sig, mu)) + mu, sig, lam = -3, 2, 0.1 + K = 1.0 / (lam * sig) + sts = stats.exponnorm.stats(K, loc=mu, scale=sig, moments='mvsk') + assert_almost_equal(sts, get_moms(lam, sig, mu)) + mu, sig, lam = 0, 3, 1 + K = 1.0 / (lam * sig) + sts = stats.exponnorm.stats(K, loc=mu, scale=sig, moments='mvsk') + assert_almost_equal(sts, get_moms(lam, sig, mu)) + mu, sig, lam = -5, 11, 3.5 + K = 1.0 / (lam * sig) + sts = stats.exponnorm.stats(K, loc=mu, scale=sig, moments='mvsk') + assert_almost_equal(sts, get_moms(lam, sig, mu)) + + def test_nan_raises_error(self): + # see gh-issue 10300 + x = np.array([1.6483, 2.7169, 2.4667, 1.1791, 3.5433, np.nan]) + assert_raises(ValueError, stats.exponnorm.fit, x, floc=0, fscale=1) + + def test_inf_raises_error(self): + # see gh-issue 10300 + x = np.array([1.6483, 2.7169, 2.4667, 1.1791, 3.5433, np.inf]) + assert_raises(ValueError, stats.exponnorm.fit, x, floc=0, fscale=1) + + def test_extremes_x(self): + # Test for extreme values against overflows + assert_almost_equal(stats.exponnorm.pdf(-900, 1), 0.0) + assert_almost_equal(stats.exponnorm.pdf(+900, 1), 0.0) + assert_almost_equal(stats.exponnorm.pdf(-900, 0.01), 0.0) + assert_almost_equal(stats.exponnorm.pdf(+900, 0.01), 0.0) + + # Expected values for the PDF were computed with mpmath, with + # the following function, and with mpmath.mp.dps = 50. + # + # def exponnorm_stdpdf(x, K): + # x = mpmath.mpf(x) + # K = mpmath.mpf(K) + # t1 = mpmath.exp(1/(2*K**2) - x/K) + # erfcarg = -(x - 1/K)/mpmath.sqrt(2) + # t2 = mpmath.erfc(erfcarg) + # return t1 * t2 / (2*K) + # + @pytest.mark.parametrize('x, K, expected', + [(20, 0.01, 6.90010764753618e-88), + (1, 0.01, 0.24438994313247364), + (-1, 0.01, 0.23955149623472075), + (-20, 0.01, 4.6004708690125477e-88), + (10, 1, 7.48518298877006e-05), + (10, 10000, 9.990005048283775e-05)]) + def test_std_pdf(self, x, K, expected): + assert_allclose(stats.exponnorm.pdf(x, K), expected, rtol=5e-12) + + # Expected values for the CDF were computed with mpmath using + # the following function and with mpmath.mp.dps = 60: + # + # def mp_exponnorm_cdf(x, K, loc=0, scale=1): + # x = mpmath.mpf(x) + # K = mpmath.mpf(K) + # loc = mpmath.mpf(loc) + # scale = mpmath.mpf(scale) + # z = (x - loc)/scale + # return (mpmath.ncdf(z) + # - mpmath.exp((1/(2*K) - z)/K)*mpmath.ncdf(z - 1/K)) + # + @pytest.mark.parametrize('x, K, scale, expected', + [[0, 0.01, 1, 0.4960109760186432], + [-5, 0.005, 1, 2.7939945412195734e-07], + [-1e4, 0.01, 100, 0.0], + [-1e4, 0.01, 1000, 6.920401854427357e-24], + [5, 0.001, 1, 0.9999997118542392]]) + def test_cdf_small_K(self, x, K, scale, expected): + p = stats.exponnorm.cdf(x, K, scale=scale) + if expected == 0.0: + assert p == 0.0 + else: + assert_allclose(p, expected, rtol=1e-13) + + # Expected values for the SF were computed with mpmath using + # the following function and with mpmath.mp.dps = 60: + # + # def mp_exponnorm_sf(x, K, loc=0, scale=1): + # x = mpmath.mpf(x) + # K = mpmath.mpf(K) + # loc = mpmath.mpf(loc) + # scale = mpmath.mpf(scale) + # z = (x - loc)/scale + # return (mpmath.ncdf(-z) + # + mpmath.exp((1/(2*K) - z)/K)*mpmath.ncdf(z - 1/K)) + # + @pytest.mark.parametrize('x, K, scale, expected', + [[10, 0.01, 1, 8.474702916146657e-24], + [2, 0.005, 1, 0.02302280664231312], + [5, 0.005, 0.5, 8.024820681931086e-24], + [10, 0.005, 0.5, 3.0603340062892486e-89], + [20, 0.005, 0.5, 0.0], + [-3, 0.001, 1, 0.9986545205566117]]) + def test_sf_small_K(self, x, K, scale, expected): + p = stats.exponnorm.sf(x, K, scale=scale) + if expected == 0.0: + assert p == 0.0 + else: + assert_allclose(p, expected, rtol=5e-13) + + @pytest.mark.parametrize('K', [1e-9, 1e-12, 1e-15]) + def test_extremely_small_K(self, K): + # see gh-issue 24551 + # test whether distribution approaches standard normal for K -> 0 + x = np.array([-10.5, -3.0, -1.0, 1.0, 3.0, 10.5]) + dist = stats.exponnorm(K) + dist_norm = stats.norm() + assert_allclose(dist.logpdf(x), dist_norm.logpdf(x)) + assert_allclose(dist.cdf(x), dist_norm.cdf(x)) + assert_allclose(dist.sf(x), dist_norm.sf(x)) + +class TestGenExpon: + def test_pdf_unity_area(self): + from scipy.integrate import simpson + # PDF should integrate to one + p = stats.genexpon.pdf(np.arange(0, 10, 0.01), 0.5, 0.5, 2.0) + assert_almost_equal(simpson(p, dx=0.01), 1, 1) + + def test_cdf_bounds(self): + # CDF should always be positive + cdf = stats.genexpon.cdf(np.arange(0, 10, 0.01), 0.5, 0.5, 2.0) + assert np.all((0 <= cdf) & (cdf <= 1)) + + # The values of p in the following data were computed with mpmath. + # E.g. the script + # from mpmath import mp + # mp.dps = 80 + # x = mp.mpf('15.0') + # a = mp.mpf('1.0') + # b = mp.mpf('2.0') + # c = mp.mpf('1.5') + # print(float(mp.exp((-a-b)*x + (b/c)*-mp.expm1(-c*x)))) + # prints + # 1.0859444834514553e-19 + @pytest.mark.parametrize('x, p, a, b, c', + [(15, 1.0859444834514553e-19, 1, 2, 1.5), + (0.25, 0.7609068232534623, 0.5, 2, 3), + (0.25, 0.09026661397565876, 9.5, 2, 0.5), + (0.01, 0.9753038265071597, 2.5, 0.25, 0.5), + (3.25, 0.0001962824553094492, 2.5, 0.25, 0.5), + (0.125, 0.9508674287164001, 0.25, 5, 0.5)]) + def test_sf_isf(self, x, p, a, b, c): + sf = stats.genexpon.sf(x, a, b, c) + assert_allclose(sf, p, rtol=2e-14) + isf = stats.genexpon.isf(p, a, b, c) + assert_allclose(isf, x, rtol=2e-14) + + # The values of p in the following data were computed with mpmath. + @pytest.mark.parametrize('x, p, a, b, c', + [(0.25, 0.2390931767465377, 0.5, 2, 3), + (0.25, 0.9097333860243412, 9.5, 2, 0.5), + (0.01, 0.0246961734928403, 2.5, 0.25, 0.5), + (3.25, 0.9998037175446906, 2.5, 0.25, 0.5), + (0.125, 0.04913257128359998, 0.25, 5, 0.5)]) + def test_cdf_ppf(self, x, p, a, b, c): + cdf = stats.genexpon.cdf(x, a, b, c) + assert_allclose(cdf, p, rtol=2e-14) + ppf = stats.genexpon.ppf(p, a, b, c) + assert_allclose(ppf, x, rtol=2e-14) + + +class TestTruncexpon: + + def test_sf_isf(self): + # reference values were computed via the reference distribution, e.g. + # mp.dps = 50; TruncExpon(b=b).sf(x) + b = [20, 100] + x = [19.999999, 99.999999] + ref = [2.0611546593828472e-15, 3.7200778266671455e-50] + assert_allclose(stats.truncexpon.sf(x, b), ref, rtol=1.5e-10) + assert_allclose(stats.truncexpon.isf(ref, b), x, rtol=1e-12) + + +class TestExponpow: + def test_tail(self): + assert_almost_equal(stats.exponpow.cdf(1e-10, 2.), 1e-20) + assert_almost_equal(stats.exponpow.isf(stats.exponpow.sf(5, .8), .8), + 5) + + +class TestSkellam: + def test_pmf(self): + # comparison to R + k = np.arange(-10, 15) + mu1, mu2 = 10, 5 + skpmfR = np.array( + [4.2254582961926893e-005, 1.1404838449648488e-004, + 2.8979625801752660e-004, 6.9177078182101231e-004, + 1.5480716105844708e-003, 3.2412274963433889e-003, + 6.3373707175123292e-003, 1.1552351566696643e-002, + 1.9606152375042644e-002, 3.0947164083410337e-002, + 4.5401737566767360e-002, 6.1894328166820688e-002, + 7.8424609500170578e-002, 9.2418812533573133e-002, + 1.0139793148019728e-001, 1.0371927988298846e-001, + 9.9076583077406091e-002, 8.8546660073089561e-002, + 7.4187842052486810e-002, 5.8392772862200251e-002, + 4.3268692953013159e-002, 3.0248159818374226e-002, + 1.9991434305603021e-002, 1.2516877303301180e-002, + 7.4389876226229707e-003]) + + assert_almost_equal(stats.skellam.pmf(k, mu1, mu2), skpmfR, decimal=15) + + def test_cdf(self): + # comparison to R, only 5 decimals + k = np.arange(-10, 15) + mu1, mu2 = 10, 5 + skcdfR = np.array( + [6.4061475386192104e-005, 1.7810985988267694e-004, + 4.6790611790020336e-004, 1.1596768997212152e-003, + 2.7077485103056847e-003, 5.9489760066490718e-003, + 1.2286346724161398e-002, 2.3838698290858034e-002, + 4.3444850665900668e-002, 7.4392014749310995e-002, + 1.1979375231607835e-001, 1.8168808048289900e-001, + 2.6011268998306952e-001, 3.5253150251664261e-001, + 4.5392943399683988e-001, 5.5764871387982828e-001, + 6.5672529695723436e-001, 7.4527195703032389e-001, + 8.1945979908281064e-001, 8.7785257194501087e-001, + 9.2112126489802404e-001, 9.5136942471639818e-001, + 9.7136085902200120e-001, 9.8387773632530240e-001, + 9.9131672394792536e-001]) + + assert_almost_equal(stats.skellam.cdf(k, mu1, mu2), skcdfR, decimal=5) + + def test_extreme_mu2(self): + # check that crash reported by gh-17916 large mu2 is resolved + x, mu1, mu2 = 0, 1, 4820232647677555.0 + assert_allclose(stats.skellam.pmf(x, mu1, mu2), 0, atol=1e-16) + assert_allclose(stats.skellam.cdf(x, mu1, mu2), 1, atol=1e-16) + + +class TestLognorm: + def test_pdf(self): + # Regression test for Ticket #1471: avoid nan with 0/0 situation + # Also make sure there are no warnings at x=0, cf gh-5202 + with warnings.catch_warnings(): + warnings.simplefilter('error', RuntimeWarning) + pdf = stats.lognorm.pdf([0, 0.5, 1], 1) + assert_array_almost_equal(pdf, [0.0, 0.62749608, 0.39894228]) + + def test_logcdf(self): + # Regression test for gh-5940: sf et al would underflow too early + x2, mu, sigma = 201.68, 195, 0.149 + assert_allclose(stats.lognorm.sf(x2-mu, s=sigma), + stats.norm.sf(np.log(x2-mu)/sigma)) + assert_allclose(stats.lognorm.logsf(x2-mu, s=sigma), + stats.norm.logsf(np.log(x2-mu)/sigma)) + + @pytest.fixture(scope='function') + def rng(self): + return np.random.default_rng(1234) + + @pytest.mark.parametrize("rvs_shape", [.1, 2]) + @pytest.mark.parametrize("rvs_loc", [-2, 0, 2]) + @pytest.mark.parametrize("rvs_scale", [.2, 1, 5]) + @pytest.mark.parametrize('fix_shape, fix_loc, fix_scale', + [e for e in product((False, True), repeat=3) + if False in e]) + @np.errstate(invalid="ignore") + def test_fit_MLE_comp_optimizer(self, rvs_shape, rvs_loc, rvs_scale, + fix_shape, fix_loc, fix_scale, rng): + data = stats.lognorm.rvs(size=100, s=rvs_shape, scale=rvs_scale, + loc=rvs_loc, random_state=rng) + + kwds = {} + if fix_shape: + kwds['f0'] = rvs_shape + if fix_loc: + kwds['floc'] = rvs_loc + if fix_scale: + kwds['fscale'] = rvs_scale + + # Numerical result may equal analytical result if some code path + # of the analytical routine makes use of numerical optimization. + _assert_less_or_close_loglike(stats.lognorm, data, **kwds, + maybe_identical=True) + + def test_isf(self): + # reference values were computed via the reference distribution, e.g. + # mp.dps = 100; + # LogNormal(s=s).isf(q=0.1, guess=0) + # LogNormal(s=s).isf(q=2e-10, guess=100) + s = 0.954 + q = [0.1, 2e-10, 5e-20, 6e-40] + ref = [3.3960065375794937, 390.07632793595974, 5830.5020828128445, + 287872.84087457904] + assert_allclose(stats.lognorm.isf(q, s), ref, rtol=1e-14) + + +class TestBeta: + def test_logpdf(self): + # Regression test for Ticket #1326: avoid nan with 0*log(0) situation + logpdf = stats.beta.logpdf(0, 1, 0.5) + assert_almost_equal(logpdf, -0.69314718056) + logpdf = stats.beta.logpdf(0, 0.5, 1) + assert_almost_equal(logpdf, np.inf) + + def test_logpdf_ticket_1866(self): + alpha, beta = 267, 1472 + x = np.array([0.2, 0.5, 0.6]) + b = stats.beta(alpha, beta) + assert_allclose(b.logpdf(x).sum(), -1201.699061824062) + assert_allclose(b.pdf(x), np.exp(b.logpdf(x))) + + def test_fit_bad_keyword_args(self): + x = [0.1, 0.5, 0.6] + assert_raises(TypeError, stats.beta.fit, x, floc=0, fscale=1, + plate="shrimp") + + def test_fit_duplicated_fixed_parameter(self): + # At most one of 'f0', 'fa' or 'fix_a' can be given to the fit method. + # More than one raises a ValueError. + x = [0.1, 0.5, 0.6] + assert_raises(ValueError, stats.beta.fit, x, fa=0.5, fix_a=0.5) + + @pytest.mark.skipif(MACOS_INTEL, reason="Overflow, see gh-14901") + def test_issue_12635(self): + # Confirm that Boost's beta distribution resolves gh-12635. + # Check against R: + # options(digits=16) + # p = 0.9999999999997369 + # a = 75.0 + # b = 66334470.0 + # print(qbeta(p, a, b)) + p, a, b = 0.9999999999997369, 75.0, 66334470.0 + assert_allclose(stats.beta.ppf(p, a, b), 2.343620802982393e-06) + + @pytest.mark.skipif(MACOS_INTEL, reason="Overflow, see gh-14901") + def test_issue_12794(self): + # Confirm that Boost's beta distribution resolves gh-12794. + # Check against R. + # options(digits=16) + # p = 1e-11 + # count_list = c(10,100,1000) + # print(qbeta(1-p, count_list + 1, 100000 - count_list)) + inv_R = np.array([0.0004944464889611935, + 0.0018360586912635726, + 0.0122663919942518351]) + count_list = np.array([10, 100, 1000]) + p = 1e-11 + inv = stats.beta.isf(p, count_list + 1, 100000 - count_list) + assert_allclose(inv, inv_R) + res = stats.beta.sf(inv, count_list + 1, 100000 - count_list) + assert_allclose(res, p) + + @pytest.mark.skipif(MACOS_INTEL, reason="Overflow, see gh-14901") + def test_issue_12796(self): + # Confirm that Boost's beta distribution succeeds in the case + # of gh-12796 + alpha_2 = 5e-6 + count_ = np.arange(1, 20) + nobs = 100000 + q, a, b = 1 - alpha_2, count_ + 1, nobs - count_ + inv = stats.beta.ppf(q, a, b) + res = stats.beta.cdf(inv, a, b) + assert_allclose(res, 1 - alpha_2) + + def test_endpoints(self): + # Confirm that boost's beta distribution returns inf at x=1 + # when b<1 + a, b = 1, 0.5 + assert_equal(stats.beta.pdf(1, a, b), np.inf) + + # Confirm that boost's beta distribution returns inf at x=0 + # when a<1 + a, b = 0.2, 3 + assert_equal(stats.beta.pdf(0, a, b), np.inf) + + # Confirm that boost's beta distribution returns 5 at x=0 + # when a=1, b=5 + a, b = 1, 5 + assert_equal(stats.beta.pdf(0, a, b), 5) + assert_equal(stats.beta.pdf(1e-310, a, b), 5) + + # Confirm that boost's beta distribution returns 5 at x=1 + # when a=5, b=1 + a, b = 5, 1 + assert_equal(stats.beta.pdf(1, a, b), 5) + assert_equal(stats.beta.pdf(1-1e-310, a, b), 5) + + def test_boost_eval_issue_14606(self): + q, a, b = 0.995, 1.0e11, 1.0e13 + stats.beta.ppf(q, a, b) + + @pytest.mark.parametrize('method', [stats.beta.ppf, stats.beta.isf]) + @pytest.mark.parametrize('a, b', [(1e-310, 12.5), (12.5, 1e-310)]) + def test_beta_ppf_with_subnormal_a_b(self, method, a, b): + # Regression test for gh-17444: beta.ppf(p, a, b) and beta.isf(p, a, b) + # would result in a segmentation fault if either a or b was subnormal. + p = 0.9 + # Depending on the version of Boost that we have vendored and + # our setting of the Boost double promotion policy, the call + # `stats.beta.ppf(p, a, b)` might raise an OverflowError or + # return a value. We'll accept either behavior (and not care about + # the value), because our goal here is to verify that the call does + # not trigger a segmentation fault. + try: + method(p, a, b) + except OverflowError: + # The OverflowError exception occurs with Boost 1.80 or earlier + # when Boost's double promotion policy is false; see + # https://github.com/boostorg/math/issues/882 + # and + # https://github.com/boostorg/math/pull/883 + # Once we have vendored the fixed version of Boost, we can drop + # this try-except wrapper and just call the function. + pass + + # Reference values computed with mpmath. + @pytest.mark.parametrize('x, a, b, ref', + [(0.999, 1.5, 2.5, -6.439838145196121e-08), + (2e-9, 3.25, 2.5, -63.13030939685114)]) + def test_logcdf(self, x, a, b, ref): + logcdf = stats.beta.logcdf(x, a, b) + assert_allclose(logcdf, ref, rtol=5e-15) + + # Reference values computed with mpmath. + @pytest.mark.parametrize('x, a, b, ref', + [(2e-9, 1.5, 2.5, -3.0368535131140806e-13), + (0.998, 3.25, 2.5, -13.309796070871489)]) + def test_logsf(self, x, a, b, ref): + logsf = stats.beta.logsf(x, a, b) + assert_allclose(logsf, ref, 5e-15) + + # entropy accuracy was confirmed using the following mpmath function + # from mpmath import mp + # mp.dps = 50 + # def beta_entropy_mpmath(a, b): + # a = mp.mpf(a) + # b = mp.mpf(b) + # entropy = mp.log(mp.beta(a, b)) - (a - 1) * mp.digamma(a) -\ + # (b - 1) * mp.digamma(b) + (a + b -2) * mp.digamma(a + b) + # return float(entropy) + + @pytest.mark.parametrize('a, b, ref', + [(0.5, 0.5, -0.24156447527049044), + (0.001, 1, -992.0922447210179), + (1, 10000, -8.210440371976183), + (100000, 100000, -5.377247470132859)]) + def test_entropy(self, a, b, ref): + assert_allclose(stats.beta(a, b).entropy(), ref) + + @pytest.mark.parametrize( + "a, b, ref, tol", + [ + (1, 10, -1.4025850929940458, 1e-14), + (10, 20, -1.0567887388936708, 1e-13), + (4e6, 4e6+20, -7.221686009678741, 1e-9), + (5e6, 5e6+10, -7.333257022834638, 1e-8), + (1e10, 1e10+20, -11.133707703130474, 1e-11), + (1e50, 1e50+20, -57.185409562486385, 1e-15), + (2, 1e10, -21.448635265288925, 1e-11), + (2, 1e20, -44.47448619497938, 1e-14), + (2, 1e50, -113.55203898480075, 1e-14), + (5, 1e10, -20.87226777401971, 1e-10), + (5, 1e20, -43.89811870326017, 1e-14), + (5, 1e50, -112.97567149308153, 1e-14), + (10, 1e10, -20.489796752909477, 1e-9), + (10, 1e20, -43.51564768139993, 1e-14), + (10, 1e50, -112.59320047122131, 1e-14), + (1e20, 2, -44.47448619497938, 1e-14), + (1e20, 5, -43.89811870326017, 1e-14), + (1e50, 10, -112.59320047122131, 1e-14), + ] + ) + def test_extreme_entropy(self, a, b, ref, tol): + # Reference values were calculated with mpmath: + # from mpmath import mp + # mp.dps = 500 + # + # def beta_entropy_mpmath(a, b): + # a = mp.mpf(a) + # b = mp.mpf(b) + # entropy = ( + # mp.log(mp.beta(a, b)) - (a - 1) * mp.digamma(a) + # - (b - 1) * mp.digamma(b) + (a + b - 2) * mp.digamma(a + b) + # ) + # return float(entropy) + assert_allclose(stats.beta(a, b).entropy(), ref, rtol=tol) + + def test_entropy_broadcasting(self): + # gh-23127 reported that the entropy method of the beta + # distribution did not broadcast correctly. + Beta = stats.make_distribution(stats.beta) + a = np.asarray([5e6, 100, 1e9, 10]) + b = np.asarray([5e6, 1e9, 100, 20]) + res = Beta(a=a, b=b).entropy() + ref = np.asarray([Beta(a=a[0], b=b[0]).entropy(), + Beta(a=a[1], b=b[1]).entropy(), + Beta(a=a[2], b=b[2]).entropy(), + Beta(a=a[3], b=b[3]).entropy()]) + assert_allclose(res, ref) + + +class TestBetaPrime: + # the test values are used in test_cdf_gh_17631 / test_ppf_gh_17631 + # They are computed with mpmath. Example: + # from mpmath import mp + # mp.dps = 50 + # a, b = mp.mpf(0.05), mp.mpf(0.1) + # x = mp.mpf(1e22) + # float(mp.betainc(a, b, 0.0, x/(1+x), regularized=True)) + # note: we use the values computed by the cdf to test whether + # ppf(cdf(x)) == x (up to a small tolerance) + # since the ppf can be very sensitive to small variations of the input, + # it can be required to generate the test case for the ppf separately, + # see self.test_ppf + cdf_vals = [ + (1e22, 100.0, 0.05, 0.8973027435427167), + (1e10, 100.0, 0.05, 0.5911548582766262), + (1e8, 0.05, 0.1, 0.9467768090820048), + (1e8, 100.0, 0.05, 0.4852944858726726), + (1e-10, 0.05, 0.1, 0.21238845427095), + (1e-10, 1.5, 1.5, 1.697652726007973e-15), + (1e-10, 0.05, 100.0, 0.40884514172337383), + (1e-22, 0.05, 0.1, 0.053349567649287326), + (1e-22, 1.5, 1.5, 1.6976527263135503e-33), + (1e-22, 0.05, 100.0, 0.10269725645728331), + (1e-100, 0.05, 0.1, 6.7163126421919795e-06), + (1e-100, 1.5, 1.5, 1.6976527263135503e-150), + (1e-100, 0.05, 100.0, 1.2928818587561651e-05), + ] + + def test_logpdf(self): + alpha, beta = 267, 1472 + x = np.array([0.2, 0.5, 0.6]) + b = stats.betaprime(alpha, beta) + assert_(np.isfinite(b.logpdf(x)).all()) + assert_allclose(b.pdf(x), np.exp(b.logpdf(x))) + + def test_cdf(self): + # regression test for gh-4030: Implementation of + # scipy.stats.betaprime.cdf() + x = stats.betaprime.cdf(0, 0.2, 0.3) + assert_equal(x, 0.0) + + alpha, beta = 267, 1472 + x = np.array([0.2, 0.5, 0.6]) + cdfs = stats.betaprime.cdf(x, alpha, beta) + assert_(np.isfinite(cdfs).all()) + + # check the new cdf implementation vs generic one: + gen_cdf = stats.rv_continuous._cdf_single + cdfs_g = [gen_cdf(stats.betaprime, val, alpha, beta) for val in x] + assert_allclose(cdfs, cdfs_g, atol=0, rtol=2e-12) + + # The expected values for test_ppf() were computed with mpmath, e.g. + # + # from mpmath import mp + # mp.dps = 125 + # p = 0.01 + # a, b = 1.25, 2.5 + # x = mp.findroot(lambda t: mp.betainc(a, b, x1=0, x2=t/(1+t), + # regularized=True) - p, + # x0=(0.01, 0.011), method='secant') + # print(float(x)) + # + # prints + # + # 0.01080162700956614 + # + @pytest.mark.parametrize( + 'p, a, b, expected', + [(0.010, 1.25, 2.5, 0.01080162700956614), + (1e-12, 1.25, 2.5, 1.0610141996279122e-10), + (1e-18, 1.25, 2.5, 1.6815941817974941e-15), + (1e-17, 0.25, 7.0, 1.0179194531881782e-69), + (0.375, 0.25, 7.0, 0.002036820346115211), + (0.9978811466052919, 0.05, 0.1, 1.0000000000001218e22),] + ) + def test_ppf(self, p, a, b, expected): + x = stats.betaprime.ppf(p, a, b) + assert_allclose(x, expected, rtol=1e-14) + + @pytest.mark.parametrize('x, a, b, p', cdf_vals) + def test_ppf_gh_17631(self, x, a, b, p): + assert_allclose(stats.betaprime.ppf(p, a, b), x, rtol=2e-14) + + def test__ppf(self): + # Verify that _ppf supports scalar arrays. + a = np.array(1.0) + b = np.array(1.0) + p = np.array(0.5) + assert_allclose(stats.betaprime._ppf(p, a, b), 1.0, rtol=5e-16) + + @pytest.mark.parametrize( + 'x, a, b, expected', + cdf_vals + [ + (1e10, 1.5, 1.5, 0.9999999999999983), + (1e10, 0.05, 0.1, 0.9664184367890859), + (1e22, 0.05, 0.1, 0.9978811466052919), + ]) + def test_cdf_gh_17631(self, x, a, b, expected): + assert_allclose(stats.betaprime.cdf(x, a, b), expected, rtol=1e-14) + + @pytest.mark.parametrize( + 'x, a, b, expected', + [(1e50, 0.05, 0.1, 0.9999966641709545), + (1e50, 100.0, 0.05, 0.995925162631006)]) + def test_cdf_extreme_tails(self, x, a, b, expected): + # for even more extreme values, we only get a few correct digits + # results are still < 1 + y = stats.betaprime.cdf(x, a, b) + assert y < 1.0 + assert_allclose(y, expected, rtol=2e-5) + + def test_sf(self): + # reference values were computed via the reference distribution, + # e.g. + # mp.dps = 50 + # a, b = 5, 3 + # x = 1e10 + # BetaPrime(a=a, b=b).sf(x); returns 3.4999999979e-29 + a = [5, 4, 2, 0.05, 0.05, 0.05, 0.05, 100.0, 100.0, 0.05, 0.05, + 0.05, 1.5, 1.5] + b = [3, 2, 1, 0.1, 0.1, 0.1, 0.1, 0.05, 0.05, 100.0, 100.0, + 100.0, 1.5, 1.5] + x = [1e10, 1e20, 1e30, 1e22, 1e-10, 1e-22, 1e-100, 1e22, 1e10, + 1e-10, 1e-22, 1e-100, 1e10, 1e-10] + ref = [3.4999999979e-29, 9.999999999994357e-40, 1.9999999999999998e-30, + 0.0021188533947081017, 0.78761154572905, 0.9466504323507127, + 0.9999932836873578, 0.10269725645728331, 0.40884514172337383, + 0.5911548582766262, 0.8973027435427167, 0.9999870711814124, + 1.6976527260079727e-15, 0.9999999999999983] + sf_values = stats.betaprime.sf(x, a, b) + assert_allclose(sf_values, ref, rtol=1e-12) + + def test_logcdf(self): + x = 800 + a = 0.5 + b = 5.0 + ref = -7.467307556554531e-16 + logcdf = stats.betaprime.logcdf(x, a, b) + assert_allclose(logcdf, ref, rtol=5e-15) + + def test_logsf(self): + x = 1e-8 + a = 4.5 + b = 0.5 + ref = -2.5868992866500915e-37 + logsf = stats.betaprime.logsf(x, a, b) + assert_allclose(logsf, ref, rtol=5e-15) + + + def test_fit_stats_gh18274(self): + # gh-18274 reported spurious warning emitted when fitting `betaprime` + # to data. Some of these were emitted by stats, too. Check that the + # warnings are no longer emitted. + stats.betaprime.fit([0.1, 0.25, 0.3, 1.2, 1.6], floc=0, fscale=1) + stats.betaprime(a=1, b=1).stats('mvsk') + + def test_moment_gh18634(self): + # Testing for gh-18634 revealed that `betaprime` raised a + # NotImplementedError for higher moments. Check that this is + # resolved. Parameters are arbitrary but lie on either side of the + # moment order (5) to test both branches of `xpx.apply_where`. + # Reference values produced with Mathematica, e.g. + # `Moment[BetaPrimeDistribution[2,7],5]` + ref = [np.inf, 0.867096912929055] + res = stats.betaprime(2, [4.2, 7.1]).moment(5) + assert_allclose(res, ref) + + +class TestGamma: + def test_pdf(self): + # a few test cases to compare with R + pdf = stats.gamma.pdf(90, 394, scale=1./5) + assert_almost_equal(pdf, 0.002312341) + + pdf = stats.gamma.pdf(3, 10, scale=1./5) + assert_almost_equal(pdf, 0.1620358) + + def test_logpdf(self): + # Regression test for Ticket #1326: cornercase avoid nan with 0*log(0) + # situation + logpdf = stats.gamma.logpdf(0, 1) + assert_almost_equal(logpdf, 0) + + def test_fit_bad_keyword_args(self): + x = [0.1, 0.5, 0.6] + assert_raises(TypeError, stats.gamma.fit, x, floc=0, plate="shrimp") + + def test_isf(self): + # Test cases for when the probability is very small. See gh-13664. + # The expected values can be checked with mpmath. With mpmath, + # the survival function sf(x, k) can be computed as + # + # mpmath.gammainc(k, x, mpmath.inf, regularized=True) + # + # Here we have: + # + # >>> mpmath.mp.dps = 60 + # >>> float(mpmath.gammainc(1, 39.14394658089878, mpmath.inf, + # ... regularized=True)) + # 9.99999999999999e-18 + # >>> float(mpmath.gammainc(100, 330.6557590436547, mpmath.inf, + # regularized=True)) + # 1.000000000000028e-50 + # + assert np.isclose(stats.gamma.isf(1e-17, 1), + 39.14394658089878, atol=1e-14) + assert np.isclose(stats.gamma.isf(1e-50, 100), + 330.6557590436547, atol=1e-13) + + def test_logcdf(self): + x = 80 + a = 7 + ref = -7.096510270453943e-27 + logcdf = stats.gamma.logcdf(x, a) + assert_allclose(logcdf, ref, rtol=5e-15) + + def test_logsf(self): + x = 0.001 + a = 3.0 + ref = -1.6654171666664883e-10 + logsf = stats.gamma.logsf(x, a) + assert_allclose(logsf, ref, rtol=5e-15) + + @pytest.mark.parametrize('scale', [1.0, 5.0]) + def test_delta_cdf(self, scale): + # Expected value computed with mpmath: + # + # >>> import mpmath + # >>> mpmath.mp.dps = 150 + # >>> cdf1 = mpmath.gammainc(3, 0, 245, regularized=True) + # >>> cdf2 = mpmath.gammainc(3, 0, 250, regularized=True) + # >>> float(cdf2 - cdf1) + # 1.1902609356171962e-102 + # + delta = stats.gamma._delta_cdf(scale*245, scale*250, 3, scale=scale) + assert_allclose(delta, 1.1902609356171962e-102, rtol=1e-13) + + @pytest.mark.parametrize('a, ref, rtol', + [(1e-4, -9990.366610819761, 1e-15), + (2, 1.5772156649015328, 1e-15), + (100, 3.7181819485047463, 1e-13), + (1e4, 6.024075385026086, 1e-15), + (1e18, 22.142204370151084, 1e-15), + (1e100, 116.54819318290696, 1e-15)]) + def test_entropy(self, a, ref, rtol): + # expected value computed with mpmath: + # from mpmath import mp + # mp.dps = 500 + # def gamma_entropy_reference(x): + # x = mp.mpf(x) + # return float(mp.digamma(x) * (mp.one - x) + x + mp.loggamma(x)) + + assert_allclose(stats.gamma.entropy(a), ref, rtol=rtol) + + @pytest.mark.parametrize("a", [1e-2, 1, 1e2]) + @pytest.mark.parametrize("loc", [1e-2, 0, 1e2]) + @pytest.mark.parametrize('scale', [1e-2, 1, 1e2]) + @pytest.mark.parametrize('fix_a', [True, False]) + @pytest.mark.parametrize('fix_loc', [True, False]) + @pytest.mark.parametrize('fix_scale', [True, False]) + def test_fit_mm(self, a, loc, scale, fix_a, fix_loc, fix_scale): + rng = np.random.default_rng(6762668991392531563) + data = stats.gamma.rvs(a, loc=loc, scale=scale, size=100, + random_state=rng) + + kwds = {} + if fix_a: + kwds['fa'] = a + if fix_loc: + kwds['floc'] = loc + if fix_scale: + kwds['fscale'] = scale + nfree = 3 - len(kwds) + + if nfree == 0: + error_msg = "All parameters fixed. There is nothing to optimize." + with pytest.raises(ValueError, match=error_msg): + stats.gamma.fit(data, method='mm', **kwds) + return + + theta = stats.gamma.fit(data, method='mm', **kwds) + dist = stats.gamma(*theta) + if nfree >= 1: + assert_allclose(dist.mean(), np.mean(data)) + if nfree >= 2: + assert_allclose(dist.moment(2), np.mean(data**2)) + if nfree >= 3: + assert_allclose(dist.moment(3), np.mean(data**3)) + + +def test_pdf_overflow_gh19616(): + # Confirm that gh19616 (intermediate over/underflows in PDF) is resolved + # Reference value from R GeneralizedHyperbolic library + # library(GeneralizedHyperbolic) + # options(digits=16) + # jitter = 1e-3 + # dnig(1, a=2**0.5 / jitter**2, b=1 / jitter**2) + jitter = 1e-3 + Z = stats.norminvgauss(2**0.5 / jitter**2, 1 / jitter**2, loc=0, scale=1) + assert_allclose(Z.pdf(1.0), 282.0948446666433) + + +class TestDgamma: + def test_pdf(self): + rng = np.random.default_rng(3791303244302340058) + size = 10 # number of points to check + x = rng.normal(scale=10, size=size) + a = rng.uniform(high=10, size=size) + res = stats.dgamma.pdf(x, a) + ref = stats.gamma.pdf(np.abs(x), a) / 2 + assert_allclose(res, ref) + + dist = stats.dgamma(a) + # There was an intermittent failure with assert_equal on Linux - 32 bit + assert_allclose(dist.pdf(x), res, rtol=5e-16) + + # mpmath was used to compute the expected values. + # For x < 0, cdf(x, a) is mp.gammainc(a, -x, mp.inf, regularized=True)/2 + # For x > 0, cdf(x, a) is (1 + mp.gammainc(a, 0, x, regularized=True))/2 + # E.g. + # from mpmath import mp + # mp.dps = 50 + # print(float(mp.gammainc(1, 20, mp.inf, regularized=True)/2)) + # prints + # 1.030576811219279e-09 + @pytest.mark.parametrize('x, a, expected', + [(-20, 1, 1.030576811219279e-09), + (-40, 1, 2.1241771276457944e-18), + (-50, 5, 2.7248509914602648e-17), + (-25, 0.125, 5.333071920958156e-14), + (5, 1, 0.9966310265004573)]) + def test_cdf_ppf_sf_isf_tail(self, x, a, expected): + cdf = stats.dgamma.cdf(x, a) + assert_allclose(cdf, expected, rtol=5e-15) + ppf = stats.dgamma.ppf(expected, a) + assert_allclose(ppf, x, rtol=5e-15) + sf = stats.dgamma.sf(-x, a) + assert_allclose(sf, expected, rtol=5e-15) + isf = stats.dgamma.isf(expected, a) + assert_allclose(isf, -x, rtol=5e-15) + + @pytest.mark.parametrize("a, ref", + [(1.5, 2.0541199559354117), + (1.3, 1.9357296377121247), + (1.1, 1.7856502333412134)]) + def test_entropy(self, a, ref): + # The reference values were calculated with mpmath: + # def entropy_dgamma(a): + # def pdf(x): + # A = mp.one / (mp.mpf(2.) * mp.gamma(a)) + # B = mp.fabs(x) ** (a - mp.one) + # C = mp.exp(-mp.fabs(x)) + # h = A * B * C + # return h + # + # return -mp.quad(lambda t: pdf(t) * mp.log(pdf(t)), + # [-mp.inf, mp.inf]) + assert_allclose(stats.dgamma.entropy(a), ref, rtol=1e-14) + + @pytest.mark.parametrize("a, ref", + [(1e-100, -1e+100), + (1e-10, -9999999975.858217), + (1e-5, -99987.37111657023), + (1e4, 6.717222565586032), + (1000000000000000.0, 19.38147391121996), + (1e+100, 117.2413403634669)]) + def test_entropy_entreme_values(self, a, ref): + # The reference values were calculated with mpmath: + # from mpmath import mp + # mp.dps = 500 + # def second_dgamma(a): + # a = mp.mpf(a) + # x_1 = a + mp.log(2) + mp.loggamma(a) + # x_2 = (mp.one - a) * mp.digamma(a) + # h = x_1 + x_2 + # return h + assert_allclose(stats.dgamma.entropy(a), ref, rtol=1e-10) + + def test_entropy_array_input(self): + x = np.array([1, 5, 1e20, 1e-5]) + y = stats.dgamma.entropy(x) + for i in range(len(y)): + assert y[i] == stats.dgamma.entropy(x[i]) + + +class TestChi2: + + # regression tests after precision improvements, ticket:1041, not verified + def test_precision(self): + assert_almost_equal(stats.chi2.pdf(1000, 1000), 8.919133934753128e-003, + decimal=14) + assert_almost_equal(stats.chi2.pdf(100, 100), 0.028162503162596778, + decimal=14) + + # Reference values computed with mpmath. + @pytest.mark.parametrize( + 'x, df, ref', + [(750.0, 3, -3.0172957781136564e-162), + (120.0, 15, -1.8924849646375648e-18), + (15.0, 13, -0.36723446372517876)] + ) + def test_logcdf(self, x, df, ref): + logcdf = stats.chi2.logcdf(x, df) + assert_allclose(logcdf, ref, rtol=5e-15) + + # Reference values computed with mpmath. + @pytest.mark.parametrize( + 'x, df, ref', + [(1e-4, 15, -3.936060782678026e-37), + (1.5, 40, -6.384797888313517e-22), + (3.0, 10, -0.018750635779926784)] + ) + def test_logsf(self, x, df, ref): + logsf = stats.chi2.logsf(x, df) + assert_allclose(logsf, ref, rtol=5e-15) + + def test_ppf(self): + # Expected values computed with mpmath. + df = 4.8 + x = stats.chi2.ppf(2e-47, df) + assert_allclose(x, 1.098472479575179840604902808e-19, rtol=1e-10) + x = stats.chi2.ppf(0.5, df) + assert_allclose(x, 4.15231407598589358660093156, rtol=1e-10) + + df = 13 + x = stats.chi2.ppf(2e-77, df) + assert_allclose(x, 1.0106330688195199050507943e-11, rtol=1e-10) + x = stats.chi2.ppf(0.1, df) + assert_allclose(x, 7.041504580095461859307179763, rtol=1e-10) + + # Entropy references values were computed with the following mpmath code + # from mpmath import mp + # mp.dps = 50 + # def chisq_entropy_mpmath(df): + # df = mp.mpf(df) + # half_df = 0.5 * df + # entropy = (half_df + mp.log(2) + mp.log(mp.gamma(half_df)) + + # (mp.one - half_df) * mp.digamma(half_df)) + # return float(entropy) + + @pytest.mark.parametrize('df, ref', + [(1e-4, -19988.980448690163), + (1, 0.7837571104739337), + (100, 4.061397128938114), + (251, 4.525577254045129), + (1e15, 19.034900320939986)]) + def test_entropy(self, df, ref): + assert_allclose(stats.chi2(df).entropy(), ref, rtol=1e-13) + + def test_regression_ticket_1326(self): + # adjust to avoid nan with 0*log(0) + assert_almost_equal(stats.chi2.pdf(0.0, 2), 0.5, 14) + + +class TestGumbelL: + def setup_method(self): + self.rng = np.random.default_rng(3922444007) + + # gh-6228 + def test_cdf_ppf(self): + x = np.linspace(-100, -4) + y = stats.gumbel_l.cdf(x) + xx = stats.gumbel_l.ppf(y) + assert_allclose(x, xx) + + def test_logcdf_logsf(self): + x = np.linspace(-100, -4) + y = stats.gumbel_l.logcdf(x) + z = stats.gumbel_l.logsf(x) + u = np.exp(y) + v = -special.expm1(z) + assert_allclose(u, v) + + def test_sf_isf(self): + x = np.linspace(-20, 5) + y = stats.gumbel_l.sf(x) + xx = stats.gumbel_l.isf(y) + assert_allclose(x, xx) + + @pytest.mark.parametrize('loc', [-1, 1]) + def test_fit_fixed_param(self, loc): + # ensure fixed location is correctly reflected from `gumbel_r.fit` + # See comments at end of gh-12737. + data = stats.gumbel_l.rvs(size=100, loc=loc, random_state=self.rng) + fitted_loc, _ = stats.gumbel_l.fit(data, floc=loc) + assert_equal(fitted_loc, loc) + + +class TestGumbelR: + + def test_sf(self): + # Expected value computed with mpmath: + # >>> import mpmath + # >>> mpmath.mp.dps = 40 + # >>> float(mpmath.mp.one - mpmath.exp(-mpmath.exp(-50))) + # 1.9287498479639178e-22 + assert_allclose(stats.gumbel_r.sf(50), 1.9287498479639178e-22, + rtol=1e-14) + + def test_isf(self): + # Expected value computed with mpmath: + # >>> import mpmath + # >>> mpmath.mp.dps = 40 + # >>> float(-mpmath.log(-mpmath.log(mpmath.mp.one - 1e-17))) + # 39.14394658089878 + assert_allclose(stats.gumbel_r.isf(1e-17), 39.14394658089878, + rtol=1e-14) + + +class TestLevyStable: + def setup_method(self): + self.rng = np.random.default_rng(7195199371) + + @pytest.fixture(autouse=True) + def reset_levy_stable_params(self): + """Setup default parameters for levy_stable generator""" + stats.levy_stable.parameterization = "S1" + stats.levy_stable.cdf_default_method = "piecewise" + stats.levy_stable.pdf_default_method = "piecewise" + stats.levy_stable.quad_eps = stats._levy_stable._QUAD_EPS + + @pytest.fixture + def nolan_pdf_sample_data(self): + """Sample data points for pdf computed with Nolan's stablec + + See - http://fs2.american.edu/jpnolan/www/stable/stable.html + + There's a known limitation of Nolan's executable for alpha < 0.2. + + The data table loaded below is generated from Nolan's stablec + with the following parameter space: + + alpha = 0.1, 0.2, ..., 2.0 + beta = -1.0, -0.9, ..., 1.0 + p = 0.01, 0.05, 0.1, 0.25, 0.35, 0.5, + and the equivalent for the right tail + + Typically inputs for stablec: + + stablec.exe << + 1 # pdf + 1 # Nolan S equivalent to S0 in scipy + .25,2,.25 # alpha + -1,-1,0 # beta + -10,10,1 # x + 1,0 # gamma, delta + 2 # output file + """ + data = np.load( + Path(__file__).parent / + 'data/levy_stable/stable-Z1-pdf-sample-data.npy' + ) + data = np.rec.fromarrays(data.T, names='x,p,alpha,beta,pct') + return data + + @pytest.fixture + def nolan_cdf_sample_data(self): + """Sample data points for cdf computed with Nolan's stablec + + See - http://fs2.american.edu/jpnolan/www/stable/stable.html + + There's a known limitation of Nolan's executable for alpha < 0.2. + + The data table loaded below is generated from Nolan's stablec + with the following parameter space: + + alpha = 0.1, 0.2, ..., 2.0 + beta = -1.0, -0.9, ..., 1.0 + p = 0.01, 0.05, 0.1, 0.25, 0.35, 0.5, + + and the equivalent for the right tail + + Ideally, Nolan's output for CDF values should match the percentile + from where they have been sampled from. Even more so as we extract + percentile x positions from stablec too. However, we note at places + Nolan's stablec will produce absolute errors in order of 1e-5. We + compare against his calculations here. In future, once we less + reliant on Nolan's paper we might switch to comparing directly at + percentiles (those x values being produced from some alternative + means). + + Typically inputs for stablec: + + stablec.exe << + 2 # cdf + 1 # Nolan S equivalent to S0 in scipy + .25,2,.25 # alpha + -1,-1,0 # beta + -10,10,1 # x + 1,0 # gamma, delta + 2 # output file + """ + data = np.load( + Path(__file__).parent / + 'data/levy_stable/stable-Z1-cdf-sample-data.npy' + ) + data = np.rec.fromarrays(data.T, names='x,p,alpha,beta,pct') + return data + + @pytest.fixture + def nolan_loc_scale_sample_data(self): + """Sample data where loc, scale are different from 0, 1 + + Data extracted in similar way to pdf/cdf above using + Nolan's stablec but set to an arbitrary location scale of + (2, 3) for various important parameters alpha, beta and for + parameterisations S0 and S1. + """ + data = np.load( + Path(__file__).parent / + 'data/levy_stable/stable-loc-scale-sample-data.npy' + ) + return data + + @pytest.mark.slow + @pytest.mark.parametrize( + "sample_size", [ + pytest.param(50), pytest.param(1500, marks=pytest.mark.slow) + ] + ) + @pytest.mark.parametrize("parameterization", ["S0", "S1"]) + @pytest.mark.parametrize( + "alpha,beta", [(1.0, 0), (1.0, -0.5), (1.5, 0), (1.9, 0.5)] + ) + @pytest.mark.parametrize("gamma,delta", [(1, 0), (3, 2)]) + def test_rvs( + self, + parameterization, + alpha, + beta, + gamma, + delta, + sample_size, + ): + stats.levy_stable.parameterization = parameterization + ls = stats.levy_stable( + alpha=alpha, beta=beta, scale=gamma, loc=delta + ) + _, p = stats.kstest( + ls.rvs(size=sample_size, random_state=self.rng), ls.cdf + ) + assert p > 0.05 + + @pytest.mark.xslow + @pytest.mark.parametrize('beta', [0.5, 1]) + def test_rvs_alpha1(self, beta): + """Additional test cases for rvs for alpha equal to 1.""" + alpha = 1.0 + loc = 0.5 + scale = 1.5 + x = stats.levy_stable.rvs(alpha, beta, loc=loc, scale=scale, + size=5000, random_state=self.rng) + stat, p = stats.kstest(x, 'levy_stable', + args=(alpha, beta, loc, scale)) + assert p > 0.01 + + def test_fit(self): + # construct data to have percentiles that match + # example in McCulloch 1986. + x = [ + -.05413, -.05413, 0., 0., 0., 0., .00533, .00533, .00533, .00533, + .00533, .03354, .03354, .03354, .03354, .03354, .05309, .05309, + .05309, .05309, .05309 + ] + alpha1, beta1, loc1, scale1 = stats.levy_stable._fitstart(x) + assert_allclose(alpha1, 1.48, rtol=0, atol=0.01) + assert_almost_equal(beta1, -.22, 2) + assert_almost_equal(scale1, 0.01717, 4) + assert_almost_equal( + loc1, 0.00233, 2 + ) # to 2 dps due to rounding error in McCulloch86 + + # cover alpha=2 scenario + x2 = x + [.05309, .05309, .05309, .05309, .05309] + alpha2, beta2, loc2, scale2 = stats.levy_stable._fitstart(x2) + assert_equal(alpha2, 2) + assert_equal(beta2, -1) + assert_almost_equal(scale2, .02503, 4) + assert_almost_equal(loc2, .03354, 4) + + @pytest.mark.xfail(reason="Unknown problem with fitstart.") + @pytest.mark.parametrize( + "alpha,beta,delta,gamma", + [ + (1.5, 0.4, 2, 3), + (1.0, 0.4, 2, 3), + ] + ) + @pytest.mark.parametrize( + "parametrization", ["S0", "S1"] + ) + def test_fit_rvs(self, alpha, beta, delta, gamma, parametrization): + """Test that fit agrees with rvs for each parametrization.""" + stats.levy_stable.parametrization = parametrization + data = stats.levy_stable.rvs( + alpha, beta, loc=delta, scale=gamma, size=10000, random_state=self.rng + ) + fit = stats.levy_stable._fitstart(data) + alpha_obs, beta_obs, delta_obs, gamma_obs = fit + assert_allclose( + [alpha, beta, delta, gamma], + [alpha_obs, beta_obs, delta_obs, gamma_obs], + rtol=0.01, + ) + + def test_fit_beta_flip(self): + # Confirm that sign of beta affects loc, not alpha or scale. + x = np.array([1, 1, 3, 3, 10, 10, 10, 30, 30, 100, 100]) + alpha1, beta1, loc1, scale1 = stats.levy_stable._fitstart(x) + alpha2, beta2, loc2, scale2 = stats.levy_stable._fitstart(-x) + assert_equal(beta1, 1) + assert loc1 != 0 + assert_almost_equal(alpha2, alpha1) + assert_almost_equal(beta2, -beta1) + assert_almost_equal(loc2, -loc1) + assert_almost_equal(scale2, scale1) + + def test_fit_delta_shift(self): + # Confirm that loc slides up and down if data shifts. + SHIFT = 1 + x = np.array([1, 1, 3, 3, 10, 10, 10, 30, 30, 100, 100]) + alpha1, beta1, loc1, scale1 = stats.levy_stable._fitstart(-x) + alpha2, beta2, loc2, scale2 = stats.levy_stable._fitstart(-x + SHIFT) + assert_almost_equal(alpha2, alpha1) + assert_almost_equal(beta2, beta1) + assert_almost_equal(loc2, loc1 + SHIFT) + assert_almost_equal(scale2, scale1) + + def test_fit_loc_extrap(self): + # Confirm that loc goes out of sample for alpha close to 1. + x = [1, 1, 3, 3, 10, 10, 10, 30, 30, 140, 140] + alpha1, beta1, loc1, scale1 = stats.levy_stable._fitstart(x) + assert alpha1 < 1, f"Expected alpha < 1, got {alpha1}" + assert loc1 < min(x), f"Expected loc < {min(x)}, got {loc1}" + + x2 = [1, 1, 3, 3, 10, 10, 10, 30, 30, 130, 130] + alpha2, beta2, loc2, scale2 = stats.levy_stable._fitstart(x2) + assert alpha2 > 1, f"Expected alpha > 1, got {alpha2}" + assert loc2 > max(x2), f"Expected loc > {max(x2)}, got {loc2}" + + @pytest.mark.slow + @pytest.mark.parametrize( + "pct_range,alpha_range,beta_range", [ + pytest.param( + [.01, .5, .99], + [.1, 1, 2], + [-1, 0, .8], + ), + pytest.param( + [.01, .05, .5, .95, .99], + [.1, .5, 1, 1.5, 2], + [-.9, -.5, 0, .3, .6, 1], + marks=pytest.mark.slow + ), + pytest.param( + [.01, .05, .1, .25, .35, .5, .65, .75, .9, .95, .99], + np.linspace(0.1, 2, 20), + np.linspace(-1, 1, 21), + marks=pytest.mark.xslow, + ), + ] + ) + def test_pdf_nolan_samples( + self, nolan_pdf_sample_data, pct_range, alpha_range, beta_range + ): + """Test pdf values against Nolan's stablec.exe output""" + data = nolan_pdf_sample_data + + # some tests break on linux 32 bit + uname = platform.uname() + is_linux_32 = uname.system == 'Linux' and uname.machine == 'i686' + platform_desc = "/".join( + [uname.system, uname.machine, uname.processor]) + + # fmt: off + # There are a number of cases which fail on some but not all platforms. + # These are excluded by the filters below. TODO: Rewrite tests so that + # the now filtered out test cases are still run but marked in pytest as + # expected to fail. + tests = [ + [ + 'dni', 1e-7, lambda r: ( + np.isin(r['pct'], pct_range) & + np.isin(r['alpha'], alpha_range) & + np.isin(r['beta'], beta_range) & + ~( + ( + (r['beta'] == 0) & + (r['pct'] == 0.5) + ) | + ( + (r['beta'] >= 0.9) & + (r['alpha'] >= 1.6) & + (r['pct'] == 0.5) + ) | + ( + (r['alpha'] <= 0.4) & + np.isin(r['pct'], [.01, .99]) + ) | + ( + (r['alpha'] <= 0.3) & + np.isin(r['pct'], [.05, .95]) + ) | + ( + (r['alpha'] <= 0.2) & + np.isin(r['pct'], [.1, .9]) + ) | + ( + (r['alpha'] == 0.1) & + np.isin(r['pct'], [.25, .75]) & + np.isin(np.abs(r['beta']), [.5, .6, .7]) + ) | + ( + (r['alpha'] == 0.1) & + np.isin(r['pct'], [.5]) & + np.isin(np.abs(r['beta']), [.1]) + ) | + ( + (r['alpha'] == 0.1) & + np.isin(r['pct'], [.35, .65]) & + np.isin(np.abs(r['beta']), [-.4, -.3, .3, .4, .5]) + ) | + ( + (r['alpha'] == 0.2) & + (r['beta'] == 0.5) & + (r['pct'] == 0.25) + ) | + ( + (r['alpha'] == 0.2) & + (r['beta'] == -0.3) & + (r['pct'] == 0.65) + ) | + ( + (r['alpha'] == 0.2) & + (r['beta'] == 0.3) & + (r['pct'] == 0.35) + ) | + ( + (r['alpha'] == 1.) & + np.isin(r['pct'], [.5]) & + np.isin(np.abs(r['beta']), [.1, .2, .3, .4]) + ) | + ( + (r['alpha'] == 1.) & + np.isin(r['pct'], [.35, .65]) & + np.isin(np.abs(r['beta']), [.8, .9, 1.]) + ) | + ( + (r['alpha'] == 1.) & + np.isin(r['pct'], [.01, .99]) & + np.isin(np.abs(r['beta']), [-.1, .1]) + ) | + # various points ok but too sparse to list + (r['alpha'] >= 1.1) + ) + ) + ], + # piecewise generally good accuracy + [ + 'piecewise', 1e-11, lambda r: ( + np.isin(r['pct'], pct_range) & + np.isin(r['alpha'], alpha_range) & + np.isin(r['beta'], beta_range) & + (r['alpha'] > 0.2) & + (r['alpha'] != 1.) + ) + ], + # for alpha = 1. for linux 32 bit optimize.bisect + # has some issues for .01 and .99 percentile + [ + 'piecewise', 1e-11, lambda r: ( + (r['alpha'] == 1.) & + (not is_linux_32) & + np.isin(r['pct'], pct_range) & + (1. in alpha_range) & + np.isin(r['beta'], beta_range) + ) + ], + # for small alpha very slightly reduced accuracy + [ + 'piecewise', 2.5e-10, lambda r: ( + np.isin(r['pct'], pct_range) & + np.isin(r['alpha'], alpha_range) & + np.isin(r['beta'], beta_range) & + (r['alpha'] <= 0.2) + ) + ], + # fft accuracy reduces as alpha decreases + [ + 'fft-simpson', 1e-5, lambda r: ( + (r['alpha'] >= 1.9) & + np.isin(r['pct'], pct_range) & + np.isin(r['alpha'], alpha_range) & + np.isin(r['beta'], beta_range) + ), + ], + [ + 'fft-simpson', 1e-6, lambda r: ( + np.isin(r['pct'], pct_range) & + np.isin(r['alpha'], alpha_range) & + np.isin(r['beta'], beta_range) & + (r['alpha'] > 1) & + (r['alpha'] < 1.9) + ) + ], + # fft relative errors for alpha < 1, will raise if enabled + # ['fft-simpson', 1e-4, lambda r: r['alpha'] == 0.9], + # ['fft-simpson', 1e-3, lambda r: r['alpha'] == 0.8], + # ['fft-simpson', 1e-2, lambda r: r['alpha'] == 0.7], + # ['fft-simpson', 1e-1, lambda r: r['alpha'] == 0.6], + ] + # fmt: on + for ix, (default_method, rtol, + filter_func) in enumerate(tests): + stats.levy_stable.pdf_default_method = default_method + subdata = data[filter_func(data) + ] if filter_func is not None else data + msg = "Density calculations experimental for FFT method" + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", msg, RuntimeWarning) + # occurs in FFT methods only + p = stats.levy_stable.pdf( + subdata['x'], + subdata['alpha'], + subdata['beta'], + scale=1, + loc=0 + ) + with np.errstate(over="ignore"): + subdata2 = rec_append_fields( + subdata, + ['calc', 'abserr', 'relerr'], + [ + p, + np.abs(p - subdata['p']), + np.abs(p - subdata['p']) / np.abs(subdata['p']) + ] + ) + failures = subdata2[ + (subdata2['relerr'] >= rtol) | + np.isnan(p) + ] + message = ( + f"pdf test {ix} failed with method '{default_method}' " + f"[platform: {platform_desc}]\n{failures.dtype.names}\n{failures}" + ) + assert_allclose( + p, + subdata['p'], + rtol, + err_msg=message, + verbose=False + ) + + @pytest.mark.parametrize( + "pct_range,alpha_range,beta_range", [ + pytest.param( + [.01, .5, .99], + [.1, 1, 2], + [-1, 0, .8], + ), + pytest.param( + [.01, .05, .5, .95, .99], + [.1, .5, 1, 1.5, 2], + [-.9, -.5, 0, .3, .6, 1], + marks=pytest.mark.slow + ), + pytest.param( + [.01, .05, .1, .25, .35, .5, .65, .75, .9, .95, .99], + np.linspace(0.1, 2, 20), + np.linspace(-1, 1, 21), + marks=pytest.mark.xslow, + ), + ] + ) + def test_cdf_nolan_samples( + self, nolan_cdf_sample_data, pct_range, alpha_range, beta_range + ): + """ Test cdf values against Nolan's stablec.exe output.""" + data = nolan_cdf_sample_data + tests = [ + # piecewise generally good accuracy + [ + 'piecewise', 2e-12, lambda r: ( + np.isin(r['pct'], pct_range) & + np.isin(r['alpha'], alpha_range) & + np.isin(r['beta'], beta_range) & + ~( + ( + (r['alpha'] == 1.) & + np.isin(r['beta'], [-0.3, -0.2, -0.1]) & + (r['pct'] == 0.01) + ) | + ( + (r['alpha'] == 1.) & + np.isin(r['beta'], [0.1, 0.2, 0.3]) & + (r['pct'] == 0.99) + ) + ) + ) + ], + # for some points with alpha=1, Nolan's STABLE clearly + # loses accuracy + [ + 'piecewise', 5e-2, lambda r: ( + np.isin(r['pct'], pct_range) & + np.isin(r['alpha'], alpha_range) & + np.isin(r['beta'], beta_range) & + ( + (r['alpha'] == 1.) & + np.isin(r['beta'], [-0.3, -0.2, -0.1]) & + (r['pct'] == 0.01) + ) | + ( + (r['alpha'] == 1.) & + np.isin(r['beta'], [0.1, 0.2, 0.3]) & + (r['pct'] == 0.99) + ) + ) + ], + # fft accuracy poor, very poor alpha < 1 + [ + 'fft-simpson', 1e-5, lambda r: ( + np.isin(r['pct'], pct_range) & + np.isin(r['alpha'], alpha_range) & + np.isin(r['beta'], beta_range) & + (r['alpha'] > 1.7) + ) + ], + [ + 'fft-simpson', 1e-4, lambda r: ( + np.isin(r['pct'], pct_range) & + np.isin(r['alpha'], alpha_range) & + np.isin(r['beta'], beta_range) & + (r['alpha'] > 1.5) & + (r['alpha'] <= 1.7) + ) + ], + [ + 'fft-simpson', 1e-3, lambda r: ( + np.isin(r['pct'], pct_range) & + np.isin(r['alpha'], alpha_range) & + np.isin(r['beta'], beta_range) & + (r['alpha'] > 1.3) & + (r['alpha'] <= 1.5) + ) + ], + [ + 'fft-simpson', 1e-2, lambda r: ( + np.isin(r['pct'], pct_range) & + np.isin(r['alpha'], alpha_range) & + np.isin(r['beta'], beta_range) & + (r['alpha'] > 1.0) & + (r['alpha'] <= 1.3) + ) + ], + ] + for ix, (default_method, rtol, + filter_func) in enumerate(tests): + stats.levy_stable.cdf_default_method = default_method + subdata = data[filter_func(data) + ] if filter_func is not None else data + with warnings.catch_warnings(): + warnings.filterwarnings( + 'ignore', + ('Cumulative density calculations experimental for FFT' + ' method. Use piecewise method instead.'), + RuntimeWarning) + p = stats.levy_stable.cdf( + subdata['x'], + subdata['alpha'], + subdata['beta'], + scale=1, + loc=0 + ) + with np.errstate(over="ignore"): + subdata2 = rec_append_fields( + subdata, + ['calc', 'abserr', 'relerr'], + [ + p, + np.abs(p - subdata['p']), + np.abs(p - subdata['p']) / np.abs(subdata['p']) + ] + ) + failures = subdata2[ + (subdata2['relerr'] >= rtol) | + np.isnan(p) + ] + message = (f"cdf test {ix} failed with method '{default_method}'\n" + f"{failures.dtype.names}\n{failures}") + assert_allclose( + p, + subdata['p'], + rtol, + err_msg=message, + verbose=False + ) + + @pytest.mark.parametrize("param", [0, 1]) + @pytest.mark.parametrize("case", ["pdf", "cdf"]) + def test_location_scale( + self, nolan_loc_scale_sample_data, param, case + ): + """Tests for pdf and cdf where loc, scale are different from 0, 1 + """ + + uname = platform.uname() + is_linux_32 = uname.system == 'Linux' and "32bit" in platform.architecture()[0] + # Test seems to be unstable (see gh-17839 for a bug report on Debian + # i386), so skip it. + if is_linux_32 and case == 'pdf': + pytest.skip("Test unstable on some platforms; see gh-17839, 17859") + + data = nolan_loc_scale_sample_data + # We only test against piecewise as location/scale transforms + # are same for other methods. + stats.levy_stable.cdf_default_method = "piecewise" + stats.levy_stable.pdf_default_method = "piecewise" + + subdata = data[data["param"] == param] + stats.levy_stable.parameterization = f"S{param}" + + assert case in ["pdf", "cdf"] + function = ( + stats.levy_stable.pdf if case == "pdf" else stats.levy_stable.cdf + ) + + v1 = function( + subdata['x'], subdata['alpha'], subdata['beta'], scale=2, loc=3 + ) + assert_allclose(v1, subdata[case], 1e-5) + + @pytest.mark.parametrize( + "method,decimal_places", + [ + ['dni', 4], + ['piecewise', 4], + ] + ) + def test_pdf_alpha_equals_one_beta_non_zero(self, method, decimal_places): + """ sample points extracted from Tables and Graphs of Stable + Probability Density Functions - Donald R Holt - 1973 - p 187. + """ + xs = np.array( + [0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4, 4] + ) + density = np.array( + [ + .3183, .3096, .2925, .2622, .1591, .1587, .1599, .1635, .0637, + .0729, .0812, .0955, .0318, .0390, .0458, .0586, .0187, .0236, + .0285, .0384 + ] + ) + betas = np.array( + [ + 0, .25, .5, 1, 0, .25, .5, 1, 0, .25, .5, 1, 0, .25, .5, 1, 0, + .25, .5, 1 + ] + ) + with np.errstate(all='ignore'), warnings.catch_warnings(): + warnings.filterwarnings("ignore", + category=RuntimeWarning, + message="Density calculation unstable.*" + ) + stats.levy_stable.pdf_default_method = method + # stats.levy_stable.fft_grid_spacing = 0.0001 + pdf = stats.levy_stable.pdf(xs, 1, betas, scale=1, loc=0) + assert_almost_equal( + pdf, density, decimal_places, method + ) + + @pytest.mark.parametrize( + "params,expected", + [ + [(1.48, -.22, 0, 1), (0, np.inf, np.nan, np.nan)], + [(2, .9, 10, 1.5), (10, 4.5, 0, 0)] + ] + ) + def test_stats(self, params, expected): + observed = stats.levy_stable.stats( + params[0], params[1], loc=params[2], scale=params[3], + moments='mvsk' + ) + assert_almost_equal(observed, expected) + + @pytest.mark.parametrize('alpha', [0.25, 0.5, 0.75]) + @pytest.mark.parametrize( + 'function,beta,points,expected', + [ + ( + stats.levy_stable.cdf, + 1.0, + np.linspace(-25, 0, 10), + 0.0, + ), + ( + stats.levy_stable.pdf, + 1.0, + np.linspace(-25, 0, 10), + 0.0, + ), + ( + stats.levy_stable.cdf, + -1.0, + np.linspace(0, 25, 10), + 1.0, + ), + ( + stats.levy_stable.pdf, + -1.0, + np.linspace(0, 25, 10), + 0.0, + ) + ] + ) + def test_distribution_outside_support( + self, alpha, function, beta, points, expected + ): + """Ensure the pdf/cdf routines do not return nan outside support. + + This distribution's support becomes truncated in a few special cases: + support is [mu, infty) if alpha < 1 and beta = 1 + support is (-infty, mu] if alpha < 1 and beta = -1 + Otherwise, the support is all reals. Here, mu is zero by default. + """ + assert 0 < alpha < 1 + assert_almost_equal( + function(points, alpha=alpha, beta=beta), + np.full(len(points), expected) + ) + + @pytest.mark.parametrize( + 'x,alpha,beta,expected', + # Reference values from Matlab + # format long + # alphas = [1.7720732804618808, 1.9217001522410235, 1.5654806051633634, + # 1.7420803447784388, 1.5748002527689913]; + # betas = [0.5059373136902996, -0.8779442746685926, -0.4016220341911392, + # -0.38180029468259247, -0.25200194914153684]; + # x0s = [0, 1e-4, -1e-4]; + # for x0 = x0s + # disp("x0 = " + x0) + # for ii = 1:5 + # alpha = alphas(ii); + # beta = betas(ii); + # pd = makedist('Stable','alpha',alpha,'beta',beta,'gam',1,'delta',0); + # % we need to adjust x. It is the same as x = 0 In scipy. + # x = x0 - beta * tan(pi * alpha / 2); + # disp(pd.pdf(x)) + # end + # end + [ + (0, 1.7720732804618808, 0.5059373136902996, 0.278932636798268), + (0, 1.9217001522410235, -0.8779442746685926, 0.281054757202316), + (0, 1.5654806051633634, -0.4016220341911392, 0.271282133194204), + (0, 1.7420803447784388, -0.38180029468259247, 0.280202199244247), + (0, 1.5748002527689913, -0.25200194914153684, 0.280136576218665), + ] + ) + def test_x_equal_zeta( + self, x, alpha, beta, expected + ): + """Test pdf for x equal to zeta. + + With S1 parametrization: x0 = x + zeta if alpha != 1 So, for x = 0, x0 + will be close to zeta. + + When case "x equal zeta" is not handled properly and quad_eps is not + low enough: - pdf may be less than 0 - logpdf is nan + + The points from the parametrize block are found randomly so that PDF is + less than 0. + + Reference values taken from MATLAB + https://www.mathworks.com/help/stats/stable-distribution.html + """ + stats.levy_stable.quad_eps = 1.2e-11 + + assert_almost_equal( + stats.levy_stable.pdf(x, alpha=alpha, beta=beta), + expected, + ) + + @pytest.mark.xfail + @pytest.mark.parametrize( + # See comment for test_x_equal_zeta for script for reference values + 'x,alpha,beta,expected', + [ + (1e-4, 1.7720732804618808, 0.5059373136902996, 0.278929165340670), + (1e-4, 1.9217001522410235, -0.8779442746685926, 0.281056564327953), + (1e-4, 1.5654806051633634, -0.4016220341911392, 0.271252432161167), + (1e-4, 1.7420803447784388, -0.38180029468259247, 0.280205311264134), + (1e-4, 1.5748002527689913, -0.25200194914153684, 0.280140965235426), + (-1e-4, 1.7720732804618808, 0.5059373136902996, 0.278936106741754), + (-1e-4, 1.9217001522410235, -0.8779442746685926, 0.281052948629429), + (-1e-4, 1.5654806051633634, -0.4016220341911392, 0.271275394392385), + (-1e-4, 1.7420803447784388, -0.38180029468259247, 0.280199085645099), + (-1e-4, 1.5748002527689913, -0.25200194914153684, 0.280132185432842), + ] + ) + def test_x_near_zeta( + self, x, alpha, beta, expected + ): + """Test pdf for x near zeta. + + With S1 parametrization: x0 = x + zeta if alpha != 1 So, for x = 0, x0 + will be close to zeta. + + When case "x near zeta" is not handled properly and quad_eps is not + low enough: - pdf may be less than 0 - logpdf is nan + + The points from the parametrize block are found randomly so that PDF is + less than 0. + + Reference values taken from MATLAB + https://www.mathworks.com/help/stats/stable-distribution.html + """ + stats.levy_stable.quad_eps = 1.2e-11 + + assert_almost_equal( + stats.levy_stable.pdf(x, alpha=alpha, beta=beta), + expected, + ) + + @pytest.fixture + def levy_stable_lock(self): + return threading.Lock() + + def test_frozen_parameterization_gh20821(self, levy_stable_lock): + # gh-20821 reported that frozen distributions ignore the parameterization. + # Check that this is resolved and that the frozen distribution's + # parameterization can be changed independently of stats.levy_stable + rng = np.random.default_rng + shapes = dict(alpha=1.9, beta=0.1, loc=0.0, scale=1.0) + unfrozen = stats.levy_stable + frozen = stats.levy_stable(**shapes) + + with levy_stable_lock: + unfrozen.parameterization = "S0" + frozen.parameterization = "S1" + unfrozen_a = unfrozen.rvs(**shapes, size=10, random_state=rng(329823498)) + frozen_a = frozen.rvs(size=10, random_state=rng(329823498)) + assert not np.any(frozen_a == unfrozen_a) + + unfrozen.parameterization = "S1" + frozen.parameterization = "S0" + unfrozen_b = unfrozen.rvs(**shapes, size=10, random_state=rng(329823498)) + frozen_b = frozen.rvs(size=10, random_state=rng(329823498)) + assert_equal(frozen_b, unfrozen_a) + assert_equal(unfrozen_b, frozen_a) + + def test_frozen_parameterization_gh20821b(self, levy_stable_lock): + # Check that the parameterization of the frozen distribution is that of + # the unfrozen distribution at the time of freezing + rng = np.random.default_rng + shapes = dict(alpha=1.9, beta=0.1, loc=0.0, scale=1.0) + unfrozen = stats.levy_stable + + with levy_stable_lock: + unfrozen.parameterization = "S0" + frozen = stats.levy_stable(**shapes) + unfrozen_a = unfrozen.rvs(**shapes, size=10, random_state=rng(329823498)) + frozen_a = frozen.rvs(size=10, random_state=rng(329823498)) + assert_equal(frozen_a, unfrozen_a) + + unfrozen.parameterization = "S1" + frozen = stats.levy_stable(**shapes) + unfrozen_b = unfrozen.rvs(**shapes, size=10, random_state=rng(329823498)) + frozen_b = frozen.rvs(size=10, random_state=rng(329823498)) + assert_equal(frozen_b, unfrozen_b) + + @pytest.mark.parametrize( + "param, alpha, beta, xs, expected", + [ + (0, 1.6, 1, [-15, -14, -13], [-187.1444263, -157.1411663, -130.3526688]), + (1, 1.6, 1, [-15, -14, -13], [-165.0203129, -137.3671717, -112.8223545]), + (0, 1.4, -1, [13, 14, 15], [-367.1943293, -464.7049565, -579.4399286]), + (1, 1.4, -1, [13, 14, 15], [-258.1216497, -334.6139857, -426.0724094]), + ], + ) + @pytest.mark.xfail_on_32bit("`bisect` fails to recognize valid bracket") + def test_logpdf_tail(self, param, alpha, beta, xs, expected): + # gh-20636 + # Expected values computed from Mathematica. Mathematica command is: + # Log[PDF[StableDistribution[0, 1.6, 1, 0, 1], -15.0]] // InputForm + # Test levy_stable.logpdf accuracy in the tails when beta=1 or beta=-1 + # for both S0/S1 parameterizations. + # Testing logpdf instead of pdf as more compact. + stats.levy_stable.parameterization = f"S{param}" + logpdfs = stats.levy_stable(alpha=alpha, beta=beta).logpdf(xs) + assert_allclose(logpdfs, expected, rtol=1e-5) + +class TestArrayArgument: # test for ticket:992 + def setup_method(self): + self.rng = np.random.default_rng(7556981556) + + def test_noexception(self): + rvs = stats.norm.rvs(loc=(np.arange(5)), scale=np.ones(5), + size=(10, 5), random_state=self.rng) + assert_equal(rvs.shape, (10, 5)) + + +class TestDocstring: + def test_docstrings(self): + # See ticket #761 + if stats.rayleigh.__doc__ is not None: + assert_("rayleigh" in stats.rayleigh.__doc__.lower()) + if stats.bernoulli.__doc__ is not None: + assert_("bernoulli" in stats.bernoulli.__doc__.lower()) + + def test_no_name_arg(self): + # If name is not given, construction shouldn't fail. See #1508. + stats.rv_continuous() + stats.rv_discrete() + + +def test_args_reduce(): + a = array([1, 3, 2, 1, 2, 3, 3]) + b, c = argsreduce(a > 1, a, 2) + + assert_array_equal(b, [3, 2, 2, 3, 3]) + assert_array_equal(c, [2]) + + b, c = argsreduce(2 > 1, a, 2) + assert_array_equal(b, a) + assert_array_equal(c, [2] * np.size(a)) + + b, c = argsreduce(a > 0, a, 2) + assert_array_equal(b, a) + assert_array_equal(c, [2] * np.size(a)) + + +class TestFitMethod: + # fitting assumes continuous parameters + skip = ['ncf', 'ksone', 'kstwo', 'irwinhall'] + + def setup_method(self): + self.rng = np.random.default_rng(4522425749) + + # skip these b/c deprecated, or only loc and scale arguments + fitSkipNonFinite = ['expon', 'norm', 'uniform', 'irwinhall'] + + @pytest.mark.parametrize('dist,args', distcont) + def test_fit_w_non_finite_data_values(self, dist, args): + """gh-10300""" + if dist in self.fitSkipNonFinite: + pytest.skip(f"{dist} fit known to fail or deprecated") + x = np.array([1.6483, 2.7169, 2.4667, 1.1791, 3.5433, np.nan]) + y = np.array([1.6483, 2.7169, 2.4667, 1.1791, 3.5433, np.inf]) + distfunc = getattr(stats, dist) + assert_raises(ValueError, distfunc.fit, x, fscale=1) + assert_raises(ValueError, distfunc.fit, y, fscale=1) + + def test_fix_fit_2args_lognorm(self): + # Regression test for #1551. + with np.errstate(all='ignore'): + x = stats.lognorm.rvs(0.25, 0., 20.0, size=20, random_state=self.rng) + expected_shape = np.sqrt(((np.log(x) - np.log(20))**2).mean()) + assert_allclose(np.array(stats.lognorm.fit(x, floc=0, fscale=20)), + [expected_shape, 0, 20], atol=1e-8) + + def test_fix_fit_norm(self): + x = np.arange(1, 6) + + loc, scale = stats.norm.fit(x) + assert_almost_equal(loc, 3) + assert_almost_equal(scale, np.sqrt(2)) + + loc, scale = stats.norm.fit(x, floc=2) + assert_equal(loc, 2) + assert_equal(scale, np.sqrt(3)) + + loc, scale = stats.norm.fit(x, fscale=2) + assert_almost_equal(loc, 3) + assert_equal(scale, 2) + + def test_fix_fit_gamma(self): + x = np.arange(1, 6) + meanlog = np.log(x).mean() + + # A basic test of gamma.fit with floc=0. + floc = 0 + a, loc, scale = stats.gamma.fit(x, floc=floc) + s = np.log(x.mean()) - meanlog + assert_almost_equal(np.log(a) - special.digamma(a), s, decimal=5) + assert_equal(loc, floc) + assert_almost_equal(scale, x.mean()/a, decimal=8) + + # Regression tests for gh-2514. + # The problem was that if `floc=0` was given, any other fixed + # parameters were ignored. + f0 = 1 + floc = 0 + a, loc, scale = stats.gamma.fit(x, f0=f0, floc=floc) + assert_equal(a, f0) + assert_equal(loc, floc) + assert_almost_equal(scale, x.mean()/a, decimal=8) + + f0 = 2 + floc = 0 + a, loc, scale = stats.gamma.fit(x, f0=f0, floc=floc) + assert_equal(a, f0) + assert_equal(loc, floc) + assert_almost_equal(scale, x.mean()/a, decimal=8) + + # loc and scale fixed. + floc = 0 + fscale = 2 + a, loc, scale = stats.gamma.fit(x, floc=floc, fscale=fscale) + assert_equal(loc, floc) + assert_equal(scale, fscale) + c = meanlog - np.log(fscale) + assert_almost_equal(special.digamma(a), c) + + def test_fix_fit_beta(self): + # Test beta.fit when both floc and fscale are given. + + def mlefunc(a, b, x): + # Zeros of this function are critical points of + # the maximum likelihood function. + n = len(x) + s1 = np.log(x).sum() + s2 = np.log(1-x).sum() + psiab = special.psi(a + b) + func = [s1 - n * (-psiab + special.psi(a)), + s2 - n * (-psiab + special.psi(b))] + return func + + # Basic test with floc and fscale given. + x = np.array([0.125, 0.25, 0.5]) + a, b, loc, scale = stats.beta.fit(x, floc=0, fscale=1) + assert_equal(loc, 0) + assert_equal(scale, 1) + assert_allclose(mlefunc(a, b, x), [0, 0], atol=1e-6) + + # Basic test with f0, floc and fscale given. + # This is also a regression test for gh-2514. + x = np.array([0.125, 0.25, 0.5]) + a, b, loc, scale = stats.beta.fit(x, f0=2, floc=0, fscale=1) + assert_equal(a, 2) + assert_equal(loc, 0) + assert_equal(scale, 1) + da, db = mlefunc(a, b, x) + assert_allclose(db, 0, atol=1e-5) + + # Same floc and fscale values as above, but reverse the data + # and fix b (f1). + x2 = 1 - x + a2, b2, loc2, scale2 = stats.beta.fit(x2, f1=2, floc=0, fscale=1) + assert_equal(b2, 2) + assert_equal(loc2, 0) + assert_equal(scale2, 1) + da, db = mlefunc(a2, b2, x2) + assert_allclose(da, 0, atol=1e-5) + # a2 of this test should equal b from above. + assert_almost_equal(a2, b) + + # Check for detection of data out of bounds when floc and fscale + # are given. + assert_raises(ValueError, stats.beta.fit, x, floc=0.5, fscale=1) + y = np.array([0, .5, 1]) + assert_raises(ValueError, stats.beta.fit, y, floc=0, fscale=1) + assert_raises(ValueError, stats.beta.fit, y, floc=0, fscale=1, f0=2) + assert_raises(ValueError, stats.beta.fit, y, floc=0, fscale=1, f1=2) + + # Check that attempting to fix all the parameters raises a ValueError. + assert_raises(ValueError, stats.beta.fit, y, f0=0, f1=1, + floc=2, fscale=3) + + def test_expon_fit(self): + x = np.array([2, 2, 4, 4, 4, 4, 4, 8]) + + loc, scale = stats.expon.fit(x) + assert_equal(loc, 2) # x.min() + assert_equal(scale, 2) # x.mean() - x.min() + + loc, scale = stats.expon.fit(x, fscale=3) + assert_equal(loc, 2) # x.min() + assert_equal(scale, 3) # fscale + + loc, scale = stats.expon.fit(x, floc=0) + assert_equal(loc, 0) # floc + assert_equal(scale, 4) # x.mean() - loc + + def test_lognorm_fit(self): + x = np.array([1.5, 3, 10, 15, 23, 59]) + lnxm1 = np.log(x - 1) + + shape, loc, scale = stats.lognorm.fit(x, floc=1) + assert_allclose(shape, lnxm1.std(), rtol=1e-12) + assert_equal(loc, 1) + assert_allclose(scale, np.exp(lnxm1.mean()), rtol=1e-12) + + shape, loc, scale = stats.lognorm.fit(x, floc=1, fscale=6) + assert_allclose(shape, np.sqrt(((lnxm1 - np.log(6))**2).mean()), + rtol=1e-12) + assert_equal(loc, 1) + assert_equal(scale, 6) + + shape, loc, scale = stats.lognorm.fit(x, floc=1, fix_s=0.75) + assert_equal(shape, 0.75) + assert_equal(loc, 1) + assert_allclose(scale, np.exp(lnxm1.mean()), rtol=1e-12) + + def test_uniform_fit(self): + x = np.array([1.0, 1.1, 1.2, 9.0]) + + loc, scale = stats.uniform.fit(x) + assert_equal(loc, x.min()) + assert_equal(scale, np.ptp(x)) + + loc, scale = stats.uniform.fit(x, floc=0) + assert_equal(loc, 0) + assert_equal(scale, x.max()) + + loc, scale = stats.uniform.fit(x, fscale=10) + assert_equal(loc, 0) + assert_equal(scale, 10) + + assert_raises(ValueError, stats.uniform.fit, x, floc=2.0) + assert_raises(ValueError, stats.uniform.fit, x, fscale=5.0) + + @pytest.mark.xslow + @pytest.mark.parametrize("method", ["MLE", "MM"]) + def test_fshapes(self, method): + # take a beta distribution, with shapes='a, b', and make sure that + # fa is equivalent to f0, and fb is equivalent to f1 + a, b = 3., 4. + x = stats.beta.rvs(a, b, size=100, random_state=self.rng) + res_1 = stats.beta.fit(x, f0=3., method=method) + res_2 = stats.beta.fit(x, fa=3., method=method) + assert_allclose(res_1, res_2, atol=1e-12, rtol=1e-12) + + res_2 = stats.beta.fit(x, fix_a=3., method=method) + assert_allclose(res_1, res_2, atol=1e-12, rtol=1e-12) + + res_3 = stats.beta.fit(x, f1=4., method=method) + res_4 = stats.beta.fit(x, fb=4., method=method) + assert_allclose(res_3, res_4, atol=1e-12, rtol=1e-12) + + res_4 = stats.beta.fit(x, fix_b=4., method=method) + assert_allclose(res_3, res_4, atol=1e-12, rtol=1e-12) + + # cannot specify both positional and named args at the same time + assert_raises(ValueError, stats.beta.fit, x, fa=1, f0=2, method=method) + + # check that attempting to fix all parameters raises a ValueError + assert_raises(ValueError, stats.beta.fit, x, fa=0, f1=1, + floc=2, fscale=3, method=method) + + # check that specifying floc, fscale and fshapes works for + # beta and gamma which override the generic fit method + res_5 = stats.beta.fit(x, fa=3., floc=0, fscale=1, method=method) + aa, bb, ll, ss = res_5 + assert_equal([aa, ll, ss], [3., 0, 1]) + + # gamma distribution + a = 3. + data = stats.gamma.rvs(a, size=100, random_state=self.rng) + aa, ll, ss = stats.gamma.fit(data, fa=a, method=method) + assert_equal(aa, a) + + @pytest.mark.parametrize("method", ["MLE", "MM"]) + def test_extra_params(self, method): + # unknown parameters should raise rather than be silently ignored + dist = stats.exponnorm + data = dist.rvs(K=2, size=100, random_state=self.rng) + dct = dict(enikibeniki=-101) + assert_raises(TypeError, dist.fit, data, **dct, method=method) + + +class TestFrozen: + # Test that a frozen distribution gives the same results as the original + # object. + # + # Only tested for the normal distribution (with loc and scale specified) + # and for the gamma distribution (with a shape parameter specified). + def test_norm(self): + dist = stats.norm + frozen = stats.norm(loc=10.0, scale=3.0) + + result_f = frozen.pdf(20.0) + result = dist.pdf(20.0, loc=10.0, scale=3.0) + assert_equal(result_f, result) + + result_f = frozen.cdf(20.0) + result = dist.cdf(20.0, loc=10.0, scale=3.0) + assert_equal(result_f, result) + + result_f = frozen.ppf(0.25) + result = dist.ppf(0.25, loc=10.0, scale=3.0) + assert_equal(result_f, result) + + result_f = frozen.isf(0.25) + result = dist.isf(0.25, loc=10.0, scale=3.0) + assert_equal(result_f, result) + + result_f = frozen.sf(10.0) + result = dist.sf(10.0, loc=10.0, scale=3.0) + assert_equal(result_f, result) + + result_f = frozen.median() + result = dist.median(loc=10.0, scale=3.0) + assert_equal(result_f, result) + + result_f = frozen.mean() + result = dist.mean(loc=10.0, scale=3.0) + assert_equal(result_f, result) + + result_f = frozen.var() + result = dist.var(loc=10.0, scale=3.0) + assert_equal(result_f, result) + + result_f = frozen.std() + result = dist.std(loc=10.0, scale=3.0) + assert_equal(result_f, result) + + result_f = frozen.entropy() + result = dist.entropy(loc=10.0, scale=3.0) + assert_equal(result_f, result) + + result_f = frozen.moment(2) + result = dist.moment(2, loc=10.0, scale=3.0) + assert_equal(result_f, result) + + assert_equal(frozen.a, dist.a) + assert_equal(frozen.b, dist.b) + + def test_gamma(self): + a = 2.0 + dist = stats.gamma + frozen = stats.gamma(a) + + result_f = frozen.pdf(20.0) + result = dist.pdf(20.0, a) + assert_equal(result_f, result) + + result_f = frozen.cdf(20.0) + result = dist.cdf(20.0, a) + assert_equal(result_f, result) + + result_f = frozen.ppf(0.25) + result = dist.ppf(0.25, a) + assert_equal(result_f, result) + + result_f = frozen.isf(0.25) + result = dist.isf(0.25, a) + assert_equal(result_f, result) + + result_f = frozen.sf(10.0) + result = dist.sf(10.0, a) + assert_equal(result_f, result) + + result_f = frozen.median() + result = dist.median(a) + assert_equal(result_f, result) + + result_f = frozen.mean() + result = dist.mean(a) + assert_equal(result_f, result) + + result_f = frozen.var() + result = dist.var(a) + assert_equal(result_f, result) + + result_f = frozen.std() + result = dist.std(a) + assert_equal(result_f, result) + + result_f = frozen.entropy() + result = dist.entropy(a) + assert_equal(result_f, result) + + result_f = frozen.moment(2) + result = dist.moment(2, a) + assert_equal(result_f, result) + + assert_equal(frozen.a, frozen.dist.a) + assert_equal(frozen.b, frozen.dist.b) + + def test_regression_ticket_1293(self): + # Create a frozen distribution. + frozen = stats.lognorm(1) + # Call one of its methods that does not take any keyword arguments. + m1 = frozen.moment(2) + # Now call a method that takes a keyword argument. + frozen.stats(moments='mvsk') + # Call moment(2) again. + # After calling stats(), the following was raising an exception. + # So this test passes if the following does not raise an exception. + m2 = frozen.moment(2) + # The following should also be true, of course. But it is not + # the focus of this test. + assert_equal(m1, m2) + + def test_ab(self): + # test that the support of a frozen distribution + # (i) remains frozen even if it changes for the original one + # (ii) is actually correct if the shape parameters are such that + # the values of [a, b] are not the default [0, inf] + # take a genpareto as an example where the support + # depends on the value of the shape parameter: + # for c > 0: a, b = 0, inf + # for c < 0: a, b = 0, -1/c + + c = -0.1 + rv = stats.genpareto(c=c) + a, b = rv.dist._get_support(c) + assert_equal([a, b], [0., 10.]) + + c = 0.1 + stats.genpareto.pdf(0, c=c) + assert_equal(rv.dist._get_support(c), [0, np.inf]) + + c = -0.1 + rv = stats.genpareto(c=c) + a, b = rv.dist._get_support(c) + assert_equal([a, b], [0., 10.]) + + c = 0.1 + stats.genpareto.pdf(0, c) # this should NOT change genpareto.b + assert_equal((rv.dist.a, rv.dist.b), stats.genpareto._get_support(c)) + + rv1 = stats.genpareto(c=0.1) + assert_(rv1.dist is not rv.dist) + + # c >= 0: a, b = [0, inf] + for c in [1., 0.]: + c = np.asarray(c) + rv = stats.genpareto(c=c) + a, b = rv.a, rv.b + assert_equal(a, 0.) + assert_(np.isposinf(b)) + + # c < 0: a=0, b=1/|c| + c = np.asarray(-2.) + a, b = stats.genpareto._get_support(c) + assert_allclose([a, b], [0., 0.5]) + + def test_rv_frozen_in_namespace(self): + # Regression test for gh-3522 + assert_(hasattr(stats.distributions, 'rv_frozen')) + + def test_random_state(self): + # only check that the random_state attribute exists, + frozen = stats.norm() + assert_(hasattr(frozen, 'random_state')) + + # ... that it can be set, + frozen.random_state = 42 + assert_equal(frozen.random_state.get_state(), + np.random.RandomState(42).get_state()) + + # ... and that .rvs method accepts it as an argument + rndm = np.random.RandomState(1234) + frozen.rvs(size=8, random_state=rndm) + + def test_pickling(self): + # test that a frozen instance pickles and unpickles + # (this method is a clone of common_tests.check_pickling) + beta = stats.beta(2.3098496451481823, 0.62687954300963677) + poiss = stats.poisson(3.) + sample = stats.rv_discrete(values=([0, 1, 2, 3], + [0.1, 0.2, 0.3, 0.4])) + + for distfn in [beta, poiss, sample]: + distfn.random_state = 1234 + distfn.rvs(size=8) + s = pickle.dumps(distfn) + r0 = distfn.rvs(size=8) + + unpickled = pickle.loads(s) + r1 = unpickled.rvs(size=8) + assert_equal(r0, r1) + + # also smoke test some methods + medians = [distfn.ppf(0.5), unpickled.ppf(0.5)] + assert_equal(medians[0], medians[1]) + assert_equal(distfn.cdf(medians[0]), + unpickled.cdf(medians[1])) + + def test_expect(self): + # smoke test the expect method of the frozen distribution + # only take a gamma w/loc and scale and poisson with loc specified + def func(x): + return x + + gm = stats.gamma(a=2, loc=3, scale=4) + with np.errstate(invalid="ignore", divide="ignore"): + gm_val = gm.expect(func, lb=1, ub=2, conditional=True) + gamma_val = stats.gamma.expect(func, args=(2,), loc=3, scale=4, + lb=1, ub=2, conditional=True) + assert_allclose(gm_val, gamma_val) + + p = stats.poisson(3, loc=4) + p_val = p.expect(func) + poisson_val = stats.poisson.expect(func, args=(3,), loc=4) + assert_allclose(p_val, poisson_val) + + +class TestExpect: + # Test for expect method. + # + # Uses normal distribution and beta distribution for finite bounds, and + # hypergeom for discrete distribution with finite support + def test_norm(self): + v = stats.norm.expect(lambda x: (x-5)*(x-5), loc=5, scale=2) + assert_almost_equal(v, 4, decimal=14) + + m = stats.norm.expect(lambda x: (x), loc=5, scale=2) + assert_almost_equal(m, 5, decimal=14) + + lb = stats.norm.ppf(0.05, loc=5, scale=2) + ub = stats.norm.ppf(0.95, loc=5, scale=2) + prob90 = stats.norm.expect(lambda x: 1, loc=5, scale=2, lb=lb, ub=ub) + assert_almost_equal(prob90, 0.9, decimal=14) + + prob90c = stats.norm.expect(lambda x: 1, loc=5, scale=2, lb=lb, ub=ub, + conditional=True) + assert_almost_equal(prob90c, 1., decimal=14) + + def test_beta(self): + # case with finite support interval + v = stats.beta.expect(lambda x: (x-19/3.)*(x-19/3.), args=(10, 5), + loc=5, scale=2) + assert_almost_equal(v, 1./18., decimal=13) + + m = stats.beta.expect(lambda x: x, args=(10, 5), loc=5., scale=2.) + assert_almost_equal(m, 19/3., decimal=13) + + ub = stats.beta.ppf(0.95, 10, 10, loc=5, scale=2) + lb = stats.beta.ppf(0.05, 10, 10, loc=5, scale=2) + prob90 = stats.beta.expect(lambda x: 1., args=(10, 10), loc=5., + scale=2., lb=lb, ub=ub, conditional=False) + assert_almost_equal(prob90, 0.9, decimal=13) + + prob90c = stats.beta.expect(lambda x: 1, args=(10, 10), loc=5, + scale=2, lb=lb, ub=ub, conditional=True) + assert_almost_equal(prob90c, 1., decimal=13) + + def test_hypergeom(self): + # test case with finite bounds + + # without specifying bounds + m_true, v_true = stats.hypergeom.stats(20, 10, 8, loc=5.) + m = stats.hypergeom.expect(lambda x: x, args=(20, 10, 8), loc=5.) + assert_almost_equal(m, m_true, decimal=13) + + v = stats.hypergeom.expect(lambda x: (x-9.)**2, args=(20, 10, 8), + loc=5.) + assert_almost_equal(v, v_true, decimal=14) + + # with bounds, bounds equal to shifted support + v_bounds = stats.hypergeom.expect(lambda x: (x-9.)**2, + args=(20, 10, 8), + loc=5., lb=5, ub=13) + assert_almost_equal(v_bounds, v_true, decimal=14) + + # drop boundary points + prob_true = 1-stats.hypergeom.pmf([5, 13], 20, 10, 8, loc=5).sum() + prob_bounds = stats.hypergeom.expect(lambda x: 1, args=(20, 10, 8), + loc=5., lb=6, ub=12) + assert_almost_equal(prob_bounds, prob_true, decimal=13) + + # conditional + prob_bc = stats.hypergeom.expect(lambda x: 1, args=(20, 10, 8), loc=5., + lb=6, ub=12, conditional=True) + assert_almost_equal(prob_bc, 1, decimal=14) + + # check simple integral + prob_b = stats.hypergeom.expect(lambda x: 1, args=(20, 10, 8), + lb=0, ub=8) + assert_almost_equal(prob_b, 1, decimal=13) + + def test_poisson(self): + # poisson, use lower bound only + prob_bounds = stats.poisson.expect(lambda x: 1, args=(2,), lb=3, + conditional=False) + prob_b_true = 1-stats.poisson.cdf(2, 2) + assert_almost_equal(prob_bounds, prob_b_true, decimal=14) + + prob_lb = stats.poisson.expect(lambda x: 1, args=(2,), lb=2, + conditional=True) + assert_almost_equal(prob_lb, 1, decimal=14) + + def test_genhalflogistic(self): + # genhalflogistic, changes upper bound of support in _argcheck + # regression test for gh-2622 + halflog = stats.genhalflogistic + # check consistency when calling expect twice with the same input + res1 = halflog.expect(args=(1.5,)) + halflog.expect(args=(0.5,)) + res2 = halflog.expect(args=(1.5,)) + assert_almost_equal(res1, res2, decimal=14) + + def test_rice_overflow(self): + # rice.pdf(999, 0.74) was inf since special.i0 silently overflows + # check that using i0e fixes it + assert_(np.isfinite(stats.rice.pdf(999, 0.74))) + + assert_(np.isfinite(stats.rice.expect(lambda x: 1, args=(0.74,)))) + assert_(np.isfinite(stats.rice.expect(lambda x: 2, args=(0.74,)))) + assert_(np.isfinite(stats.rice.expect(lambda x: 3, args=(0.74,)))) + + def test_logser(self): + # test a discrete distribution with infinite support and loc + p, loc = 0.3, 3 + res_0 = stats.logser.expect(lambda k: k, args=(p,)) + # check against the correct answer (sum of a geom series) + assert_allclose(res_0, + p / (p - 1.) / np.log(1. - p), atol=1e-15) + + # now check it with `loc` + res_l = stats.logser.expect(lambda k: k, args=(p,), loc=loc) + assert_allclose(res_l, res_0 + loc, atol=1e-15) + + def test_skellam(self): + # Use a discrete distribution w/ bi-infinite support. Compute two first + # moments and compare to known values (cf skellam.stats) + p1, p2 = 18, 22 + m1 = stats.skellam.expect(lambda x: x, args=(p1, p2)) + m2 = stats.skellam.expect(lambda x: x**2, args=(p1, p2)) + assert_allclose(m1, p1 - p2, atol=1e-12) + assert_allclose(m2 - m1**2, p1 + p2, atol=1e-12) + + def test_randint(self): + # Use a discrete distribution w/ parameter-dependent support, which + # is larger than the default chunksize + lo, hi = 0, 113 + res = stats.randint.expect(lambda x: x, (lo, hi)) + assert_allclose(res, + sum(_ for _ in range(lo, hi)) / (hi - lo), atol=1e-15) + + def test_zipf(self): + # Test that there is no infinite loop even if the sum diverges + with pytest.warns(RuntimeWarning): + stats.zipf.expect(lambda x: x**2, (2,)) + + def test_discrete_kwds(self): + # check that discrete expect accepts keywords to control the summation + n0 = stats.poisson.expect(lambda x: 1, args=(2,)) + n1 = stats.poisson.expect(lambda x: 1, args=(2,), + maxcount=1001, chunksize=32, tolerance=1e-8) + assert_almost_equal(n0, n1, decimal=14) + + def test_moment(self): + # test the .moment() method: compute a higher moment and compare to + # a known value + def poiss_moment5(mu): + return mu**5 + 10*mu**4 + 25*mu**3 + 15*mu**2 + mu + + for mu in [5, 7]: + m5 = stats.poisson.moment(5, mu) + assert_allclose(m5, poiss_moment5(mu), rtol=1e-10) + + def test_challenging_cases_gh8928(self): + # Several cases where `expect` failed to produce a correct result were + # reported in gh-8928. Check that these cases have been resolved. + assert_allclose(stats.norm.expect(loc=36, scale=1.0), 36) + assert_allclose(stats.norm.expect(loc=40, scale=1.0), 40) + assert_allclose(stats.norm.expect(loc=10, scale=0.1), 10) + assert_allclose(stats.gamma.expect(args=(148,)), 148) + assert_allclose(stats.logistic.expect(loc=85), 85) + + def test_lb_ub_gh15855(self): + # Make sure changes to `expect` made in gh15855 treat lb/ub correctly + dist = stats.uniform + ref = dist.mean(loc=10, scale=5) # 12.5 + # moment over whole distribution + assert_allclose(dist.expect(loc=10, scale=5), ref) + # moment over whole distribution, lb and ub outside of support + assert_allclose(dist.expect(loc=10, scale=5, lb=9, ub=16), ref) + # moment over 60% of distribution, [lb, ub] centered within support + assert_allclose(dist.expect(loc=10, scale=5, lb=11, ub=14), ref*0.6) + # moment over truncated distribution, essentially + assert_allclose(dist.expect(loc=10, scale=5, lb=11, ub=14, + conditional=True), ref) + # moment over 40% of distribution, [lb, ub] not centered within support + assert_allclose(dist.expect(loc=10, scale=5, lb=11, ub=13), 12*0.4) + # moment with lb > ub + assert_allclose(dist.expect(loc=10, scale=5, lb=13, ub=11), -12*0.4) + # moment with lb > ub, conditional + assert_allclose(dist.expect(loc=10, scale=5, lb=13, ub=11, + conditional=True), 12) + + +class TestNct: + def test_nc_parameter(self): + # Parameter values c<=0 were not enabled (gh-2402). + # For negative values c and for c=0 results of rv.cdf(0) below were nan + rv = stats.nct(5, 0) + assert_equal(rv.cdf(0), 0.5) + rv = stats.nct(5, -1) + assert_almost_equal(rv.cdf(0), 0.841344746069, decimal=10) + + def test_broadcasting(self): + res = stats.nct.pdf(5, np.arange(4, 7)[:, None], + np.linspace(0.1, 1, 4)) + expected = array([[0.00321886, 0.00557466, 0.00918418, 0.01442997], + [0.00217142, 0.00395366, 0.00683888, 0.01126276], + [0.00153078, 0.00291093, 0.00525206, 0.00900815]]) + assert_allclose(res, expected, rtol=1e-5) + + def test_variance_gh_issue_2401(self): + # Computation of the variance of a non-central t-distribution resulted + # in a TypeError: ufunc 'isinf' not supported for the input types, + # and the inputs could not be safely coerced to any supported types + # according to the casting rule 'safe' + rv = stats.nct(4, 0) + assert_equal(rv.var(), 2.0) + + def test_nct_inf_moments(self): + # n-th moment of nct only exists for df > n + m, v, s, k = stats.nct.stats(df=0.9, nc=0.3, moments='mvsk') + assert_equal([m, v, s, k], [np.nan, np.nan, np.nan, np.nan]) + + m, v, s, k = stats.nct.stats(df=1.9, nc=0.3, moments='mvsk') + assert_(np.isfinite(m)) + assert_equal([v, s, k], [np.nan, np.nan, np.nan]) + + m, v, s, k = stats.nct.stats(df=3.1, nc=0.3, moments='mvsk') + assert_(np.isfinite([m, v, s]).all()) + assert_equal(k, np.nan) + + def test_nct_stats_large_df_values(self): + # previously gamma function was used which lost precision at df=345 + # cf. https://github.com/scipy/scipy/issues/12919 for details + nct_mean_df_1000 = stats.nct.mean(1000, 2) + nct_stats_df_1000 = stats.nct.stats(1000, 2) + # These expected values were computed with mpmath. They were also + # verified with the Wolfram Alpha expressions: + # Mean[NoncentralStudentTDistribution[1000, 2]] + # Var[NoncentralStudentTDistribution[1000, 2]] + expected_stats_df_1000 = [2.0015015641422464, 1.0040115288163005] + assert_allclose(nct_mean_df_1000, expected_stats_df_1000[0], + rtol=1e-10) + assert_allclose(nct_stats_df_1000, expected_stats_df_1000, + rtol=1e-10) + # and a bigger df value + nct_mean = stats.nct.mean(100000, 2) + nct_stats = stats.nct.stats(100000, 2) + # These expected values were computed with mpmath. + expected_stats = [2.0000150001562518, 1.0000400011500288] + assert_allclose(nct_mean, expected_stats[0], rtol=1e-10) + assert_allclose(nct_stats, expected_stats, rtol=1e-9) + + def test_cdf_large_nc(self): + # gh-17916 reported a crash with large `nc` values + assert_allclose(stats.nct.cdf(2, 2, float(2**16)), 0) + + # PDF reference values were computed with mpmath + # with 100 digits of precision + + # def nct_pdf(x, df, nc): + # x = mp.mpf(x) + # n = mp.mpf(df) + # nc = mp.mpf(nc) + + # x2 = x*x + # ncx2 = nc*nc*x2 + # fac1 = n + x2 + # trm1 = (n/2.*mp.log(n) + mp.loggamma(n + mp.one) + # - (n * mp.log(2.) + nc*nc/2 + (n/2)*mp.log(fac1) + # + mp.loggamma(n/2))) + # Px = mp.exp(trm1) + # valF = ncx2 / (2*fac1) + # trm1 = (mp.sqrt(2)*nc*x*mp.hyp1f1(n/2+1, 1.5, valF) + # / (fac1*mp.gamma((n+1)/2))) + # trm2 = (mp.hyp1f1((n+1)/2, 0.5, valF) + # / (mp.sqrt(fac1)*mp.gamma(n/2 + mp.one))) + # Px *= trm1+trm2 + # return float(Px) + + @pytest.mark.parametrize("x, df, nc, expected", [ + (10000, 10, 16, 3.394646922945872e-30), + (-10, 8, 16, 4.282769500264159e-70) + ]) + def test_pdf_large_nc(self, x, df, nc, expected): + # gh-#20693 reported zero values for large `nc` values + assert_allclose(stats.nct.pdf(x, df, nc), expected, rtol=1e-12) + + +class TestRecipInvGauss: + + def test_pdf_endpoint(self): + p = stats.recipinvgauss.pdf(0, 0.6) + assert p == 0.0 + + def test_logpdf_endpoint(self): + logp = stats.recipinvgauss.logpdf(0, 0.6) + assert logp == -np.inf + + def test_cdf_small_x(self): + # The expected value was computer with mpmath: + # + # import mpmath + # + # mpmath.mp.dps = 100 + # + # def recipinvgauss_cdf_mp(x, mu): + # x = mpmath.mpf(x) + # mu = mpmath.mpf(mu) + # trm1 = 1/mu - x + # trm2 = 1/mu + x + # isqx = 1/mpmath.sqrt(x) + # return (mpmath.ncdf(-isqx*trm1) + # - mpmath.exp(2/mu)*mpmath.ncdf(-isqx*trm2)) + # + p = stats.recipinvgauss.cdf(0.05, 0.5) + expected = 6.590396159501331e-20 + assert_allclose(p, expected, rtol=1e-14) + + def test_sf_large_x(self): + # The expected value was computed with mpmath; see test_cdf_small. + p = stats.recipinvgauss.sf(80, 0.5) + expected = 2.699819200556787e-18 + assert_allclose(p, expected, 5e-15) + + +class TestRice: + def setup_method(self): + self.rng = np.random.default_rng(666822542) + + def test_rice_zero_b(self): + # rice distribution should work with b=0, cf gh-2164 + x = [0.2, 1., 5.] + assert_(np.isfinite(stats.rice.pdf(x, b=0.)).all()) + assert_(np.isfinite(stats.rice.logpdf(x, b=0.)).all()) + assert_(np.isfinite(stats.rice.cdf(x, b=0.)).all()) + assert_(np.isfinite(stats.rice.logcdf(x, b=0.)).all()) + + q = [0.1, 0.1, 0.5, 0.9] + assert_(np.isfinite(stats.rice.ppf(q, b=0.)).all()) + + mvsk = stats.rice.stats(0, moments='mvsk') + assert_(np.isfinite(mvsk).all()) + + # furthermore, pdf is continuous as b\to 0 + # rice.pdf(x, b\to 0) = x exp(-x^2/2) + O(b^2) + # see e.g. Abramovich & Stegun 9.6.7 & 9.6.10 + b = 1e-8 + assert_allclose(stats.rice.pdf(x, 0), stats.rice.pdf(x, b), + atol=b, rtol=0) + + def test_rice_rvs(self): + rvs = stats.rice.rvs + assert_equal(rvs(b=3., random_state=self.rng).size, 1) + assert_equal(rvs(b=3., size=(3, 5), random_state=self.rng).shape, (3, 5)) + + def test_rice_gh9836(self): + # test that gh-9836 is resolved; previously jumped to 1 at the end + + cdf = stats.rice.cdf(np.arange(10, 160, 10), np.arange(10, 160, 10)) + # Generated in R + # library(VGAM) + # options(digits=16) + # x = seq(10, 150, 10) + # print(price(x, sigma=1, vee=x)) + cdf_exp = [0.4800278103504522, 0.4900233218590353, 0.4933500379379548, + 0.4950128317658719, 0.4960103776798502, 0.4966753655438764, + 0.4971503395812474, 0.4975065620443196, 0.4977836197921638, + 0.4980052636649550, 0.4981866072661382, 0.4983377260666599, + 0.4984655952615694, 0.4985751970541413, 0.4986701850071265] + assert_allclose(cdf, cdf_exp) + + probabilities = np.arange(0.1, 1, 0.1) + ppf = stats.rice.ppf(probabilities, 500/4, scale=4) + # Generated in R + # library(VGAM) + # options(digits=16) + # p = seq(0.1, .9, by = .1) + # print(qrice(p, vee = 500, sigma = 4)) + ppf_exp = [494.8898762347361, 496.6495690858350, 497.9184315188069, + 499.0026277378915, 500.0159999146250, 501.0293721352668, + 502.1135684981884, 503.3824312270405, 505.1421247157822] + assert_allclose(ppf, ppf_exp) + + ppf = scipy.stats.rice.ppf(0.5, np.arange(10, 150, 10)) + # Generated in R + # library(VGAM) + # options(digits=16) + # b <- seq(10, 140, 10) + # print(qrice(0.5, vee = b, sigma = 1)) + ppf_exp = [10.04995862522287, 20.02499480078302, 30.01666512465732, + 40.01249934924363, 50.00999966676032, 60.00833314046875, + 70.00714273568241, 80.00624991862573, 90.00555549840364, + 100.00499995833597, 110.00454542324384, 120.00416664255323, + 130.00384613488120, 140.00357141338748] + assert_allclose(ppf, ppf_exp) + + +class TestErlang: + def setup_method(self): + self.rng = np.random.default_rng(2792245532) + + def test_erlang_runtimewarning(self): + # erlang should generate a RuntimeWarning if a non-integer + # shape parameter is used. + with warnings.catch_warnings(): + warnings.simplefilter("error", RuntimeWarning) + + # The non-integer shape parameter 1.3 should trigger a + # RuntimeWarning + assert_raises(RuntimeWarning, stats.erlang.rvs, 1.3, loc=0, + scale=1, size=4, random_state=self.rng) + + # Calling the fit method with `f0` set to an integer should + # *not* trigger a RuntimeWarning. It should return the same + # values as gamma.fit(...). + data = [0.5, 1.0, 2.0, 4.0] + result_erlang = stats.erlang.fit(data, f0=1) + result_gamma = stats.gamma.fit(data, f0=1) + assert_allclose(result_erlang, result_gamma, rtol=1e-3) + + def test_gh_pr_10949_argcheck(self): + assert_equal(stats.erlang.pdf(0.5, a=[1, -1]), + stats.gamma.pdf(0.5, a=[1, -1])) + + +class TestRayleigh: + def setup_method(self): + self.rng = np.random.default_rng(7186715712) + + # gh-6227 + def test_logpdf(self): + y = stats.rayleigh.logpdf(50) + assert_allclose(y, -1246.0879769945718) + + def test_logsf(self): + y = stats.rayleigh.logsf(50) + assert_allclose(y, -1250) + + @pytest.mark.parametrize("rvs_loc,rvs_scale", [(0.85373171, 0.86932204), + (0.20558821, 0.61621008)]) + def test_fit(self, rvs_loc, rvs_scale): + data = stats.rayleigh.rvs(size=250, loc=rvs_loc, + scale=rvs_scale, random_state=self.rng) + + def scale_mle(data, floc): + return (np.sum((data - floc) ** 2) / (2 * len(data))) ** .5 + + # when `floc` is provided, `scale` is found with an analytical formula + scale_expect = scale_mle(data, rvs_loc) + loc, scale = stats.rayleigh.fit(data, floc=rvs_loc) + assert_equal(loc, rvs_loc) + assert_equal(scale, scale_expect) + + # when `fscale` is fixed, superclass fit is used to determine `loc`. + loc, scale = stats.rayleigh.fit(data, fscale=.6) + assert_equal(scale, .6) + + # with both parameters free, one dimensional optimization is done + # over a new function that takes into account the dependent relation + # of `scale` to `loc`. + loc, scale = stats.rayleigh.fit(data) + # test that `scale` is defined by its relation to `loc` + assert_equal(scale, scale_mle(data, loc)) + + @pytest.mark.parametrize("rvs_loc,rvs_scale", [[0.74, 0.01], + [0.08464463, 0.12069025]]) + def test_fit_comparison_super_method(self, rvs_loc, rvs_scale): + # test that the objective function result of the analytical MLEs is + # less than or equal to that of the numerically optimized estimate + data = stats.rayleigh.rvs(size=250, loc=rvs_loc, + scale=rvs_scale, random_state=self.rng) + _assert_less_or_close_loglike(stats.rayleigh, data) + + def test_fit_warnings(self): + assert_fit_warnings(stats.rayleigh) + + def test_fit_gh17088(self): + # `rayleigh.fit` could return a location that was inconsistent with + # the data. See gh-17088. + rng = np.random.default_rng(456) + loc, scale, size = 50, 600, 500 + rvs = stats.rayleigh.rvs(loc, scale, size=size, random_state=rng) + loc_fit, _ = stats.rayleigh.fit(rvs) + assert loc_fit < np.min(rvs) + loc_fit, scale_fit = stats.rayleigh.fit(rvs, fscale=scale) + assert loc_fit < np.min(rvs) + assert scale_fit == scale + + +class TestExponWeib: + + def test_pdf_logpdf(self): + # Regression test for gh-3508. + x = 0.1 + a = 1.0 + c = 100.0 + p = stats.exponweib.pdf(x, a, c) + logp = stats.exponweib.logpdf(x, a, c) + # Expected values were computed with mpmath. + assert_allclose([p, logp], + [1.0000000000000054e-97, -223.35075402042244]) + + def test_a_is_1(self): + # For issue gh-3508. + # Check that when a=1, the pdf and logpdf methods of exponweib are the + # same as those of weibull_min. + x = np.logspace(-4, -1, 4) + a = 1 + c = 100 + + p = stats.exponweib.pdf(x, a, c) + expected = stats.weibull_min.pdf(x, c) + assert_allclose(p, expected) + + logp = stats.exponweib.logpdf(x, a, c) + expected = stats.weibull_min.logpdf(x, c) + assert_allclose(logp, expected) + + def test_a_is_1_c_is_1(self): + # When a = 1 and c = 1, the distribution is exponential. + x = np.logspace(-8, 1, 10) + a = 1 + c = 1 + + p = stats.exponweib.pdf(x, a, c) + expected = stats.expon.pdf(x) + assert_allclose(p, expected) + + logp = stats.exponweib.logpdf(x, a, c) + expected = stats.expon.logpdf(x) + assert_allclose(logp, expected) + + # Reference values were computed with mpmath, e.g: + # + # from mpmath import mp + # + # def mp_sf(x, a, c): + # x = mp.mpf(x) + # a = mp.mpf(a) + # c = mp.mpf(c) + # return -mp.powm1(-mp.expm1(-x**c)), a) + # + # mp.dps = 100 + # print(float(mp_sf(1, 2.5, 0.75))) + # + # prints + # + # 0.6823127476985246 + # + @pytest.mark.parametrize( + 'x, a, c, ref', + [(1, 2.5, 0.75, 0.6823127476985246), + (50, 2.5, 0.75, 1.7056666054719663e-08), + (125, 2.5, 0.75, 1.4534393150714602e-16), + (250, 2.5, 0.75, 1.2391389689773512e-27), + (250, 0.03125, 0.75, 1.548923711221689e-29), + (3, 0.03125, 3.0, 5.873527551689983e-14), + (2e80, 10.0, 0.02, 2.9449084156902135e-17)] + ) + def test_sf(self, x, a, c, ref): + sf = stats.exponweib.sf(x, a, c) + assert_allclose(sf, ref, rtol=1e-14) + + # Reference values were computed with mpmath, e.g. + # + # from mpmath import mp + # + # def mp_isf(p, a, c): + # p = mp.mpf(p) + # a = mp.mpf(a) + # c = mp.mpf(c) + # return (-mp.log(-mp.expm1(mp.log1p(-p)/a)))**(1/c) + # + # mp.dps = 100 + # print(float(mp_isf(0.25, 2.5, 0.75))) + # + # prints + # + # 2.8946008178158924 + # + @pytest.mark.parametrize( + 'p, a, c, ref', + [(0.25, 2.5, 0.75, 2.8946008178158924), + (3e-16, 2.5, 0.75, 121.77966713102938), + (1e-12, 1, 2, 5.256521769756932), + (2e-13, 0.03125, 3, 2.953915059484589), + (5e-14, 10.0, 0.02, 7.57094886384687e+75)] + ) + def test_isf(self, p, a, c, ref): + isf = stats.exponweib.isf(p, a, c) + assert_allclose(isf, ref, rtol=5e-14) + + # Reference values computed with mpmath. + @pytest.mark.parametrize('x, a, c, ref', + [(2, 3, 8, -1.9848783170128456e-111), + (1000, 0.5, 0.75, -2.946296827524972e-78)]) + def test_logcdf(self, x, a, c, ref): + logcdf = stats.exponweib.logcdf(x, a, c) + assert_allclose(logcdf, ref, rtol=5e-15) + + # Reference values computed with mpmath. + @pytest.mark.parametrize('x, a, c, ref', + [(1e-65, 1.5, 1.25, -1.333521432163324e-122), + (2e-10, 2, 10, -1.0485760000000007e-194)]) + def test_logsf(self, x, a, c, ref): + logsf = stats.exponweib.logsf(x, a, c) + assert_allclose(logsf, ref, rtol=5e-15) + + +class TestFatigueLife: + + def test_sf_tail(self): + # Expected value computed with mpmath: + # import mpmath + # mpmath.mp.dps = 80 + # x = mpmath.mpf(800.0) + # c = mpmath.mpf(2.5) + # s = float(1 - mpmath.ncdf(1/c * (mpmath.sqrt(x) + # - 1/mpmath.sqrt(x)))) + # print(s) + # Output: + # 6.593376447038406e-30 + s = stats.fatiguelife.sf(800.0, 2.5) + assert_allclose(s, 6.593376447038406e-30, rtol=1e-13) + + def test_isf_tail(self): + # See test_sf_tail for the mpmath code. + p = 6.593376447038406e-30 + q = stats.fatiguelife.isf(p, 2.5) + assert_allclose(q, 800.0, rtol=1e-13) + + +class TestWeibull: + + def test_logpdf(self): + # gh-6217 + y = stats.weibull_min.logpdf(0, 1) + assert_equal(y, 0) + + def test_with_maxima_distrib(self): + # Tests for weibull_min and weibull_max. + # The expected values were computed using the symbolic algebra + # program 'maxima' with the package 'distrib', which has + # 'pdf_weibull' and 'cdf_weibull'. The mapping between the + # scipy and maxima functions is as follows: + # ----------------------------------------------------------------- + # scipy maxima + # --------------------------------- ------------------------------ + # weibull_min.pdf(x, a, scale=b) pdf_weibull(x, a, b) + # weibull_min.logpdf(x, a, scale=b) log(pdf_weibull(x, a, b)) + # weibull_min.cdf(x, a, scale=b) cdf_weibull(x, a, b) + # weibull_min.logcdf(x, a, scale=b) log(cdf_weibull(x, a, b)) + # weibull_min.sf(x, a, scale=b) 1 - cdf_weibull(x, a, b) + # weibull_min.logsf(x, a, scale=b) log(1 - cdf_weibull(x, a, b)) + # + # weibull_max.pdf(x, a, scale=b) pdf_weibull(-x, a, b) + # weibull_max.logpdf(x, a, scale=b) log(pdf_weibull(-x, a, b)) + # weibull_max.cdf(x, a, scale=b) 1 - cdf_weibull(-x, a, b) + # weibull_max.logcdf(x, a, scale=b) log(1 - cdf_weibull(-x, a, b)) + # weibull_max.sf(x, a, scale=b) cdf_weibull(-x, a, b) + # weibull_max.logsf(x, a, scale=b) log(cdf_weibull(-x, a, b)) + # ----------------------------------------------------------------- + x = 1.5 + a = 2.0 + b = 3.0 + + # weibull_min + + p = stats.weibull_min.pdf(x, a, scale=b) + assert_allclose(p, np.exp(-0.25)/3) + + lp = stats.weibull_min.logpdf(x, a, scale=b) + assert_allclose(lp, -0.25 - np.log(3)) + + c = stats.weibull_min.cdf(x, a, scale=b) + assert_allclose(c, -special.expm1(-0.25)) + + lc = stats.weibull_min.logcdf(x, a, scale=b) + assert_allclose(lc, np.log(-special.expm1(-0.25))) + + s = stats.weibull_min.sf(x, a, scale=b) + assert_allclose(s, np.exp(-0.25)) + + ls = stats.weibull_min.logsf(x, a, scale=b) + assert_allclose(ls, -0.25) + + # Also test using a large value x, for which computing the survival + # function using the CDF would result in 0. + s = stats.weibull_min.sf(30, 2, scale=3) + assert_allclose(s, np.exp(-100)) + + ls = stats.weibull_min.logsf(30, 2, scale=3) + assert_allclose(ls, -100) + + # weibull_max + x = -1.5 + + p = stats.weibull_max.pdf(x, a, scale=b) + assert_allclose(p, np.exp(-0.25)/3) + + lp = stats.weibull_max.logpdf(x, a, scale=b) + assert_allclose(lp, -0.25 - np.log(3)) + + c = stats.weibull_max.cdf(x, a, scale=b) + assert_allclose(c, np.exp(-0.25)) + + lc = stats.weibull_max.logcdf(x, a, scale=b) + assert_allclose(lc, -0.25) + + s = stats.weibull_max.sf(x, a, scale=b) + assert_allclose(s, -special.expm1(-0.25)) + + ls = stats.weibull_max.logsf(x, a, scale=b) + assert_allclose(ls, np.log(-special.expm1(-0.25))) + + # Also test using a value of x close to 0, for which computing the + # survival function using the CDF would result in 0. + s = stats.weibull_max.sf(-1e-9, 2, scale=3) + assert_allclose(s, -special.expm1(-1/9000000000000000000)) + + ls = stats.weibull_max.logsf(-1e-9, 2, scale=3) + assert_allclose(ls, np.log(-special.expm1(-1/9000000000000000000))) + + @pytest.mark.parametrize('scale', [1.0, 0.1]) + def test_delta_cdf(self, scale): + # Expected value computed with mpmath: + # + # def weibull_min_sf(x, k, scale): + # x = mpmath.mpf(x) + # k = mpmath.mpf(k) + # scale =mpmath.mpf(scale) + # return mpmath.exp(-(x/scale)**k) + # + # >>> import mpmath + # >>> mpmath.mp.dps = 60 + # >>> sf1 = weibull_min_sf(7.5, 3, 1) + # >>> sf2 = weibull_min_sf(8.0, 3, 1) + # >>> float(sf1 - sf2) + # 6.053624060118734e-184 + # + delta = stats.weibull_min._delta_cdf(scale*7.5, scale*8, 3, + scale=scale) + assert_allclose(delta, 6.053624060118734e-184) + + def test_fit_min(self): + rng = np.random.default_rng(5985959307161735394) + + c, loc, scale = 2, 3.5, 0.5 # arbitrary, valid parameters + dist = stats.weibull_min(c, loc, scale) + rvs = dist.rvs(size=100, random_state=rng) + + # test that MLE still honors guesses and fixed parameters + c2, loc2, scale2 = stats.weibull_min.fit(rvs, 1.5, floc=3) + c3, loc3, scale3 = stats.weibull_min.fit(rvs, 1.6, floc=3) + assert loc2 == loc3 == 3 # fixed parameter is respected + assert c2 != c3 # different guess -> (slightly) different outcome + # quality of fit is tested elsewhere + + # test that MoM honors fixed parameters, accepts (but ignores) guesses + c4, loc4, scale4 = stats.weibull_min.fit(rvs, 3, fscale=3, method='mm') + assert scale4 == 3 + # because scale was fixed, only the mean and skewness will be matched + dist4 = stats.weibull_min(c4, loc4, scale4) + res = dist4.stats(moments='ms') + ref = np.mean(rvs), stats.skew(rvs) + assert_allclose(res, ref) + + # reference values were computed via mpmath + # from mpmath import mp + # def weibull_sf_mpmath(x, c): + # x = mp.mpf(x) + # c = mp.mpf(c) + # return float(mp.exp(-x**c)) + + @pytest.mark.parametrize('x, c, ref', [(50, 1, 1.9287498479639178e-22), + (1000, 0.8, + 8.131269637872743e-110)]) + def test_sf_isf(self, x, c, ref): + assert_allclose(stats.weibull_min.sf(x, c), ref, rtol=5e-14) + assert_allclose(stats.weibull_min.isf(ref, c), x, rtol=5e-14) + + +class TestDweibull: + def test_entropy(self): + # Test that dweibull entropy follows that of weibull_min. + # (Generic tests check that the dweibull entropy is consistent + # with its PDF. As for accuracy, dweibull entropy should be just + # as accurate as weibull_min entropy. Checks of accuracy against + # a reference need only be applied to the fundamental distribution - + # weibull_min.) + rng = np.random.default_rng(8486259129157041777) + c = 10**rng.normal(scale=100, size=10) + res = stats.dweibull.entropy(c) + ref = stats.weibull_min.entropy(c) - np.log(0.5) + assert_allclose(res, ref, rtol=1e-15) + + def test_sf(self): + # test that for positive values the dweibull survival function is half + # the weibull_min survival function + rng = np.random.default_rng(8486259129157041777) + c = 10**rng.normal(scale=1, size=10) + x = 10 * rng.uniform() + res = stats.dweibull.sf(x, c) + ref = 0.5 * stats.weibull_min.sf(x, c) + assert_allclose(res, ref, rtol=1e-15) + + +class TestTruncWeibull: + + def test_pdf_bounds(self): + # test bounds + y = stats.truncweibull_min.pdf([0.1, 2.0], 2.0, 0.11, 1.99) + assert_equal(y, [0.0, 0.0]) + + def test_logpdf(self): + y = stats.truncweibull_min.logpdf(2.0, 1.0, 2.0, np.inf) + assert_equal(y, 0.0) + + # hand calculation + y = stats.truncweibull_min.logpdf(2.0, 1.0, 2.0, 4.0) + assert_allclose(y, 0.14541345786885884) + + def test_ppf_bounds(self): + # test bounds + y = stats.truncweibull_min.ppf([0.0, 1.0], 2.0, 0.1, 2.0) + assert_equal(y, [0.1, 2.0]) + + def test_cdf_to_ppf(self): + q = [0., 0.1, .25, 0.50, 0.75, 0.90, 1.] + x = stats.truncweibull_min.ppf(q, 2., 0., 3.) + q_out = stats.truncweibull_min.cdf(x, 2., 0., 3.) + assert_allclose(q, q_out) + + def test_sf_to_isf(self): + q = [0., 0.1, .25, 0.50, 0.75, 0.90, 1.] + x = stats.truncweibull_min.isf(q, 2., 0., 3.) + q_out = stats.truncweibull_min.sf(x, 2., 0., 3.) + assert_allclose(q, q_out) + + def test_munp(self): + c = 2. + a = 1. + b = 3. + + def xnpdf(x, n): + return x**n*stats.truncweibull_min.pdf(x, c, a, b) + + m0 = stats.truncweibull_min.moment(0, c, a, b) + assert_equal(m0, 1.) + + m1 = stats.truncweibull_min.moment(1, c, a, b) + m1_expected, _ = quad(lambda x: xnpdf(x, 1), a, b) + assert_allclose(m1, m1_expected) + + m2 = stats.truncweibull_min.moment(2, c, a, b) + m2_expected, _ = quad(lambda x: xnpdf(x, 2), a, b) + assert_allclose(m2, m2_expected) + + m3 = stats.truncweibull_min.moment(3, c, a, b) + m3_expected, _ = quad(lambda x: xnpdf(x, 3), a, b) + assert_allclose(m3, m3_expected) + + m4 = stats.truncweibull_min.moment(4, c, a, b) + m4_expected, _ = quad(lambda x: xnpdf(x, 4), a, b) + assert_allclose(m4, m4_expected) + + def test_reference_values(self): + a = 1. + b = 3. + c = 2. + x_med = np.sqrt(1 - np.log(0.5 + np.exp(-(8. + np.log(2.))))) + + cdf = stats.truncweibull_min.cdf(x_med, c, a, b) + assert_allclose(cdf, 0.5) + + lc = stats.truncweibull_min.logcdf(x_med, c, a, b) + assert_allclose(lc, -np.log(2.)) + + ppf = stats.truncweibull_min.ppf(0.5, c, a, b) + assert_allclose(ppf, x_med) + + sf = stats.truncweibull_min.sf(x_med, c, a, b) + assert_allclose(sf, 0.5) + + ls = stats.truncweibull_min.logsf(x_med, c, a, b) + assert_allclose(ls, -np.log(2.)) + + isf = stats.truncweibull_min.isf(0.5, c, a, b) + assert_allclose(isf, x_med) + + def test_compare_weibull_min(self): + # Verify that the truncweibull_min distribution gives the same results + # as the original weibull_min + x = 1.5 + c = 2.0 + a = 0.0 + b = np.inf + scale = 3.0 + + p = stats.weibull_min.pdf(x, c, scale=scale) + p_trunc = stats.truncweibull_min.pdf(x, c, a, b, scale=scale) + assert_allclose(p, p_trunc) + + lp = stats.weibull_min.logpdf(x, c, scale=scale) + lp_trunc = stats.truncweibull_min.logpdf(x, c, a, b, scale=scale) + assert_allclose(lp, lp_trunc) + + cdf = stats.weibull_min.cdf(x, c, scale=scale) + cdf_trunc = stats.truncweibull_min.cdf(x, c, a, b, scale=scale) + assert_allclose(cdf, cdf_trunc) + + lc = stats.weibull_min.logcdf(x, c, scale=scale) + lc_trunc = stats.truncweibull_min.logcdf(x, c, a, b, scale=scale) + assert_allclose(lc, lc_trunc) + + s = stats.weibull_min.sf(x, c, scale=scale) + s_trunc = stats.truncweibull_min.sf(x, c, a, b, scale=scale) + assert_allclose(s, s_trunc) + + ls = stats.weibull_min.logsf(x, c, scale=scale) + ls_trunc = stats.truncweibull_min.logsf(x, c, a, b, scale=scale) + assert_allclose(ls, ls_trunc) + + # # Also test using a large value x, for which computing the survival + # # function using the CDF would result in 0. + s = stats.truncweibull_min.sf(30, 2, a, b, scale=3) + assert_allclose(s, np.exp(-100)) + + ls = stats.truncweibull_min.logsf(30, 2, a, b, scale=3) + assert_allclose(ls, -100) + + def test_compare_weibull_min2(self): + # Verify that the truncweibull_min distribution PDF and CDF results + # are the same as those calculated from truncating weibull_min + c, a, b = 2.5, 0.25, 1.25 + x = np.linspace(a, b, 100) + + pdf1 = stats.truncweibull_min.pdf(x, c, a, b) + cdf1 = stats.truncweibull_min.cdf(x, c, a, b) + + norm = stats.weibull_min.cdf(b, c) - stats.weibull_min.cdf(a, c) + pdf2 = stats.weibull_min.pdf(x, c) / norm + cdf2 = (stats.weibull_min.cdf(x, c) - stats.weibull_min.cdf(a, c))/norm + + np.testing.assert_allclose(pdf1, pdf2) + np.testing.assert_allclose(cdf1, cdf2) + + +class TestRdist: + def test_rdist_cdf_gh1285(self): + # check workaround in rdist._cdf for issue gh-1285. + distfn = stats.rdist + values = [0.001, 0.5, 0.999] + assert_almost_equal(distfn.cdf(distfn.ppf(values, 541.0), 541.0), + values, decimal=5) + + def test_rdist_beta(self): + # rdist is a special case of stats.beta + x = np.linspace(-0.99, 0.99, 10) + c = 2.7 + assert_almost_equal(0.5*stats.beta(c/2, c/2).pdf((x + 1)/2), + stats.rdist(c).pdf(x)) + + # reference values were computed via mpmath + # from mpmath import mp + # mp.dps = 200 + # def rdist_sf_mpmath(x, c): + # x = mp.mpf(x) + # c = mp.mpf(c) + # return float(mp.betainc(c/2, c/2, (x+1)/2, mp.one, regularized=True)) + @pytest.mark.parametrize( + "x, c, ref", + [ + (0.0001, 541, 0.49907251345565845), + (0.1, 241, 0.06000788166249205), + (0.5, 441, 1.0655898106047832e-29), + (0.8, 341, 6.025478373732215e-78), + ] + ) + def test_rdist_sf(self, x, c, ref): + assert_allclose(stats.rdist.sf(x, c), ref, rtol=5e-14) + + +class TestTrapezoid: + def test_reduces_to_triang(self): + modes = [0, 0.3, 0.5, 1] + for mode in modes: + x = [0, mode, 1] + assert_almost_equal(stats.trapezoid.pdf(x, mode, mode), + stats.triang.pdf(x, mode)) + assert_almost_equal(stats.trapezoid.cdf(x, mode, mode), + stats.triang.cdf(x, mode)) + + def test_reduces_to_uniform(self): + x = np.linspace(0, 1, 10) + assert_almost_equal(stats.trapezoid.pdf(x, 0, 1), stats.uniform.pdf(x)) + assert_almost_equal(stats.trapezoid.cdf(x, 0, 1), stats.uniform.cdf(x)) + + def test_cases(self): + # edge cases + assert_almost_equal(stats.trapezoid.pdf(0, 0, 0), 2) + assert_almost_equal(stats.trapezoid.pdf(1, 1, 1), 2) + assert_almost_equal(stats.trapezoid.pdf(0.5, 0, 0.8), + 1.11111111111111111) + assert_almost_equal(stats.trapezoid.pdf(0.5, 0.2, 1.0), + 1.11111111111111111) + + # straightforward case + assert_almost_equal(stats.trapezoid.pdf(0.1, 0.2, 0.8), 0.625) + assert_almost_equal(stats.trapezoid.pdf(0.5, 0.2, 0.8), 1.25) + assert_almost_equal(stats.trapezoid.pdf(0.9, 0.2, 0.8), 0.625) + + assert_almost_equal(stats.trapezoid.cdf(0.1, 0.2, 0.8), 0.03125) + assert_almost_equal(stats.trapezoid.cdf(0.2, 0.2, 0.8), 0.125) + assert_almost_equal(stats.trapezoid.cdf(0.5, 0.2, 0.8), 0.5) + assert_almost_equal(stats.trapezoid.cdf(0.9, 0.2, 0.8), 0.96875) + assert_almost_equal(stats.trapezoid.cdf(1.0, 0.2, 0.8), 1.0) + + def test_moments_and_entropy(self): + # issue #11795: improve precision of trapezoid stats + # Apply formulas from Wikipedia for the following parameters: + a, b, c, d = -3, -1, 2, 3 # => 1/3, 5/6, -3, 6 + p1, p2, loc, scale = (b-a) / (d-a), (c-a) / (d-a), a, d-a + h = 2 / (d+c-b-a) + + def moment(n): + return (h * ((d**(n+2) - c**(n+2)) / (d-c) + - (b**(n+2) - a**(n+2)) / (b-a)) / + (n+1) / (n+2)) + + mean = moment(1) + var = moment(2) - mean**2 + entropy = 0.5 * (d-c+b-a) / (d+c-b-a) + np.log(0.5 * (d+c-b-a)) + assert_almost_equal(stats.trapezoid.mean(p1, p2, loc, scale), + mean, decimal=13) + assert_almost_equal(stats.trapezoid.var(p1, p2, loc, scale), + var, decimal=13) + assert_almost_equal(stats.trapezoid.entropy(p1, p2, loc, scale), + entropy, decimal=13) + + # Check boundary cases where scipy d=0 or d=1. + assert_almost_equal(stats.trapezoid.mean(0, 0, -3, 6), -1, decimal=13) + assert_almost_equal(stats.trapezoid.mean(0, 1, -3, 6), 0, decimal=13) + assert_almost_equal(stats.trapezoid.var(0, 1, -3, 6), 3, decimal=13) + + def test_trapezoid_vect(self): + # test that array-valued shapes and arguments are handled + c = np.array([0.1, 0.2, 0.3]) + d = np.array([0.5, 0.6])[:, None] + x = np.array([0.15, 0.25, 0.9]) + v = stats.trapezoid.pdf(x, c, d) + + cc, dd, xx = np.broadcast_arrays(c, d, x) + + res = np.empty(xx.size, dtype=xx.dtype) + ind = np.arange(xx.size) + for i, x1, c1, d1 in zip(ind, xx.ravel(), cc.ravel(), dd.ravel()): + res[i] = stats.trapezoid.pdf(x1, c1, d1) + + assert_allclose(v, res.reshape(v.shape), atol=1e-15) + + # Check that the stats() method supports vector arguments. + v = np.asarray(stats.trapezoid.stats(c, d, moments="mvsk")) + cc, dd = np.broadcast_arrays(c, d) + res = np.empty((cc.size, 4)) # 4 stats returned per value + ind = np.arange(cc.size) + for i, c1, d1 in zip(ind, cc.ravel(), dd.ravel()): + res[i] = stats.trapezoid.stats(c1, d1, moments="mvsk") + + assert_allclose(v, res.T.reshape(v.shape), atol=1e-15) + + def test_trapezoid_fit_convergence_gh23503(self): + # gh-23503 reported that trapezoid.fit would consistently converge to a + # triangular distribution unless starting values were provided. Check that this + # is resolved. + + # Generate test data from a trapezoidal distribution + rng = np.random.default_rng(23842359234598263956) + true_args = 0.3, 0.7, -1, 2 + true_dist = stats.trapezoid(*true_args) + x = true_dist.rvs(1000, random_state=rng) + + # fit to data + fitted_args = stats.trapezoid.fit(x) + + # Should not converge to triangular distribution (c=d=1) + fitted_c, fitted_d = fitted_args[:2] + assert not np.allclose(fitted_c, 1, atol=0.1) + assert not np.allclose(fitted_d, 1, atol=0.1) + + # objective function is better than with true values of parameters + true_llf = stats.trapezoid.nnlf(true_args, x) + fitted_llf = stats.trapezoid.nnlf(fitted_args, x) + assert fitted_llf < true_llf + + +class TestTriang: + def test_edge_cases(self): + with np.errstate(all='raise'): + assert_equal(stats.triang.pdf(0, 0), 2.) + assert_equal(stats.triang.pdf(0.5, 0), 1.) + assert_equal(stats.triang.pdf(1, 0), 0.) + + assert_equal(stats.triang.pdf(0, 1), 0) + assert_equal(stats.triang.pdf(0.5, 1), 1.) + assert_equal(stats.triang.pdf(1, 1), 2) + + assert_equal(stats.triang.cdf(0., 0.), 0.) + assert_equal(stats.triang.cdf(0.5, 0.), 0.75) + assert_equal(stats.triang.cdf(1.0, 0.), 1.0) + + assert_equal(stats.triang.cdf(0., 1.), 0.) + assert_equal(stats.triang.cdf(0.5, 1.), 0.25) + assert_equal(stats.triang.cdf(1., 1.), 1) + + +class TestMaxwell: + + # reference values were computed with wolfram alpha + # erfc(x/sqrt(2)) + sqrt(2/pi) * x * e^(-x^2/2) + + @pytest.mark.parametrize("x, ref", + [(20, 2.2138865931011177e-86), + (0.01, 0.999999734046458435)]) + def test_sf(self, x, ref): + assert_allclose(stats.maxwell.sf(x), ref, rtol=1e-14) + + # reference values were computed with wolfram alpha + # sqrt(2) * sqrt(Q^(-1)(3/2, q)) + + @pytest.mark.parametrize("q, ref", + [(0.001, 4.033142223656157022), + (0.9999847412109375, 0.0385743284050381), + (2**-55, 8.95564974719481)]) + def test_isf(self, q, ref): + assert_allclose(stats.maxwell.isf(q), ref, rtol=1e-15) + + def test_logcdf(self): + # Reference value computed with mpmath. + ref = -1.8729310110194814e-17 + logcdf = stats.maxwell.logcdf(9) + assert_allclose(logcdf, ref, rtol=5e-15) + + def test_logsf(self): + # Reference value computed with mpmath. + ref = -2.6596152026762177e-25 + logsf = stats.maxwell.logsf(1e-8) + assert_allclose(logsf, ref, rtol=5e-15) + + +class TestMielke: + def test_moments(self): + k, s = 4.642, 0.597 + # n-th moment exists only if n < s + assert_equal(stats.mielke(k, s).moment(1), np.inf) + assert_equal(stats.mielke(k, 1.0).moment(1), np.inf) + assert_(np.isfinite(stats.mielke(k, 1.01).moment(1))) + + def test_burr_equivalence(self): + x = np.linspace(0.01, 100, 50) + k, s = 2.45, 5.32 + assert_allclose(stats.burr.pdf(x, s, k/s), stats.mielke.pdf(x, k, s)) + + +class TestBurr: + def test_endpoints_7491(self): + # gh-7491 + # Compute the pdf at the left endpoint dst.a. + data = [ + [stats.fisk, (1,), 1], + [stats.burr, (0.5, 2), 1], + [stats.burr, (1, 1), 1], + [stats.burr, (2, 0.5), 1], + [stats.burr12, (1, 0.5), 0.5], + [stats.burr12, (1, 1), 1.0], + [stats.burr12, (1, 2), 2.0]] + + ans = [_f.pdf(_f.a, *_args) for _f, _args, _ in data] + correct = [_correct_ for _f, _args, _correct_ in data] + assert_array_almost_equal(ans, correct) + + ans = [_f.logpdf(_f.a, *_args) for _f, _args, _ in data] + correct = [np.log(_correct_) for _f, _args, _correct_ in data] + assert_array_almost_equal(ans, correct) + + def test_burr_stats_9544(self): + # gh-9544. Test from gh-9978 + c, d = 5.0, 3 + mean, variance = stats.burr(c, d).stats() + # mean = sc.beta(3 + 1/5, 1. - 1/5) * 3 = 1.4110263... + # var = sc.beta(3 + 2 / 5, 1. - 2 / 5) * 3 - + # (sc.beta(3 + 1 / 5, 1. - 1 / 5) * 3) ** 2 + mean_hc, variance_hc = 1.4110263183925857, 0.22879948026191643 + assert_allclose(mean, mean_hc) + assert_allclose(variance, variance_hc) + + def test_burr_nan_mean_var_9544(self): + # gh-9544. Test from gh-9978 + c, d = 0.5, 3 + mean, variance = stats.burr(c, d).stats() + assert_(np.isnan(mean)) + assert_(np.isnan(variance)) + c, d = 1.5, 3 + mean, variance = stats.burr(c, d).stats() + assert_(np.isfinite(mean)) + assert_(np.isnan(variance)) + + c, d = 0.5, 3 + e1, e2, e3, e4 = stats.burr._munp(np.array([1, 2, 3, 4]), c, d) + assert_(np.isnan(e1)) + assert_(np.isnan(e2)) + assert_(np.isnan(e3)) + assert_(np.isnan(e4)) + c, d = 1.5, 3 + e1, e2, e3, e4 = stats.burr._munp([1, 2, 3, 4], c, d) + assert_(np.isfinite(e1)) + assert_(np.isnan(e2)) + assert_(np.isnan(e3)) + assert_(np.isnan(e4)) + c, d = 2.5, 3 + e1, e2, e3, e4 = stats.burr._munp([1, 2, 3, 4], c, d) + assert_(np.isfinite(e1)) + assert_(np.isfinite(e2)) + assert_(np.isnan(e3)) + assert_(np.isnan(e4)) + c, d = 3.5, 3 + e1, e2, e3, e4 = stats.burr._munp([1, 2, 3, 4], c, d) + assert_(np.isfinite(e1)) + assert_(np.isfinite(e2)) + assert_(np.isfinite(e3)) + assert_(np.isnan(e4)) + c, d = 4.5, 3 + e1, e2, e3, e4 = stats.burr._munp([1, 2, 3, 4], c, d) + assert_(np.isfinite(e1)) + assert_(np.isfinite(e2)) + assert_(np.isfinite(e3)) + assert_(np.isfinite(e4)) + + def test_burr_isf(self): + # reference values were computed via the reference distribution, e.g. + # mp.dps = 100 + # Burr(c=5, d=3).isf([0.1, 1e-10, 1e-20, 1e-40]) + c, d = 5.0, 3.0 + q = [0.1, 1e-10, 1e-20, 1e-40] + ref = [1.9469686558286508, 124.57309395989076, 12457.309396155173, + 124573093.96155174] + assert_allclose(stats.burr.isf(q, c, d), ref, rtol=1e-14) + + +class TestBurr12: + + @pytest.mark.parametrize('scale, expected', + [(1.0, 2.3283064359965952e-170), + (3.5, 5.987114417447875e-153)]) + def test_delta_cdf(self, scale, expected): + # Expected value computed with mpmath: + # + # def burr12sf(x, c, d, scale): + # x = mpmath.mpf(x) + # c = mpmath.mpf(c) + # d = mpmath.mpf(d) + # scale = mpmath.mpf(scale) + # return (mpmath.mp.one + (x/scale)**c)**(-d) + # + # >>> import mpmath + # >>> mpmath.mp.dps = 60 + # >>> float(burr12sf(2e5, 4, 8, 1) - burr12sf(4e5, 4, 8, 1)) + # 2.3283064359965952e-170 + # >>> float(burr12sf(2e5, 4, 8, 3.5) - burr12sf(4e5, 4, 8, 3.5)) + # 5.987114417447875e-153 + # + delta = stats.burr12._delta_cdf(2e5, 4e5, 4, 8, scale=scale) + assert_allclose(delta, expected, rtol=1e-13) + + def test_moments_edge(self): + # gh-18838 reported that burr12 moments could be invalid; see above. + # Check that this is resolved in an edge case where c*d == n, and + # compare the results against those produced by Mathematica, e.g. + # `SinghMaddalaDistribution[2, 2, 1]` at Wolfram Alpha. + c, d = 2, 2 + mean = np.pi/4 + var = 1 - np.pi**2/16 + skew = np.pi**3/(32*var**1.5) + kurtosis = np.nan + ref = [mean, var, skew, kurtosis] + res = stats.burr12(c, d).stats('mvsk') + assert_allclose(res, ref, rtol=1e-14) + + # Reference values were computed with mpmath using mp.dps = 80 + # and then cast to float. + @pytest.mark.parametrize( + 'p, c, d, ref', + [(1e-12, 20, 0.5, 15.848931924611135), + (1e-19, 20, 0.5, 79.43282347242815), + (1e-12, 0.25, 35, 2.0888618213462466), + (1e-80, 0.25, 35, 1360930951.7972188)] + ) + def test_isf_near_zero(self, p, c, d, ref): + x = stats.burr12.isf(p, c, d) + assert_allclose(x, ref, rtol=1e-14) + + +class TestStudentizedRange: + # For alpha = .05, .01, and .001, and for each value of + # v = [1, 3, 10, 20, 120, inf], a Q was picked from each table for + # k = [2, 8, 14, 20]. + + # these arrays are written with `k` as column, and `v` as rows. + # Q values are taken from table 3: + # https://www.jstor.org/stable/2237810 + q05 = [17.97, 45.40, 54.33, 59.56, + 4.501, 8.853, 10.35, 11.24, + 3.151, 5.305, 6.028, 6.467, + 2.950, 4.768, 5.357, 5.714, + 2.800, 4.363, 4.842, 5.126, + 2.772, 4.286, 4.743, 5.012] + q01 = [90.03, 227.2, 271.8, 298.0, + 8.261, 15.64, 18.22, 19.77, + 4.482, 6.875, 7.712, 8.226, + 4.024, 5.839, 6.450, 6.823, + 3.702, 5.118, 5.562, 5.827, + 3.643, 4.987, 5.400, 5.645] + q001 = [900.3, 2272, 2718, 2980, + 18.28, 34.12, 39.69, 43.05, + 6.487, 9.352, 10.39, 11.03, + 5.444, 7.313, 7.966, 8.370, + 4.772, 6.039, 6.448, 6.695, + 4.654, 5.823, 6.191, 6.411] + qs = np.concatenate((q05, q01, q001)) + ps = [.95, .99, .999] + vs = [1, 3, 10, 20, 120, np.inf] + ks = [2, 8, 14, 20] + + data = list(zip(product(ps, vs, ks), qs)) + + # A small selection of large-v cases generated with R's `ptukey` + # Each case is in the format (q, k, v, r_result) + r_data = [ + (0.1, 3, 9001, 0.002752818526842), + (1, 10, 1000, 0.000526142388912), + (1, 3, np.inf, 0.240712641229283), + (4, 3, np.inf, 0.987012338626815), + (1, 10, np.inf, 0.000519869467083), + ] + + @pytest.mark.slow + def test_cdf_against_tables(self): + for pvk, q in self.data: + p_expected, v, k = pvk + res_p = stats.studentized_range.cdf(q, k, v) + assert_allclose(res_p, p_expected, rtol=1e-4) + + @pytest.mark.xslow + def test_ppf_against_tables(self): + for pvk, q_expected in self.data: + p, v, k = pvk + res_q = stats.studentized_range.ppf(p, k, v) + assert_allclose(res_q, q_expected, rtol=5e-4) + + path_prefix = os.path.dirname(__file__) + relative_path = "data/studentized_range_mpmath_ref.json" + with open(os.path.join(path_prefix, relative_path)) as file: + pregenerated_data = json.load(file) + + @pytest.mark.parametrize("case_result", pregenerated_data["cdf_data"]) + def test_cdf_against_mp(self, case_result): + src_case = case_result["src_case"] + mp_result = case_result["mp_result"] + qkv = src_case["q"], src_case["k"], src_case["v"] + res = stats.studentized_range.cdf(*qkv) + + assert_allclose(res, mp_result, + atol=src_case["expected_atol"], + rtol=src_case["expected_rtol"]) + + @pytest.mark.parametrize("case_result", pregenerated_data["pdf_data"]) + def test_pdf_against_mp(self, case_result): + src_case = case_result["src_case"] + mp_result = case_result["mp_result"] + qkv = src_case["q"], src_case["k"], src_case["v"] + res = stats.studentized_range.pdf(*qkv) + + assert_allclose(res, mp_result, + atol=src_case["expected_atol"], + rtol=src_case["expected_rtol"]) + + @pytest.mark.xslow + @pytest.mark.xfail_on_32bit("intermittent RuntimeWarning: invalid value.") + @pytest.mark.parametrize("case_result", pregenerated_data["moment_data"]) + def test_moment_against_mp(self, case_result): + src_case = case_result["src_case"] + mp_result = case_result["mp_result"] + mkv = src_case["m"], src_case["k"], src_case["v"] + + # Silence invalid value encountered warnings. Actual problems will be + # caught by the result comparison. + with np.errstate(invalid='ignore'): + res = stats.studentized_range.moment(*mkv) + + assert_allclose(res, mp_result, + atol=src_case["expected_atol"], + rtol=src_case["expected_rtol"]) + + @pytest.mark.slow + def test_pdf_integration(self): + k, v = 3, 10 + # Test whether PDF integration is 1 like it should be. + res = quad(stats.studentized_range.pdf, 0, np.inf, args=(k, v)) + assert_allclose(res[0], 1) + + @pytest.mark.xslow + def test_pdf_against_cdf(self): + k, v = 3, 10 + + # Test whether the integrated PDF matches the CDF using cumulative + # integration. Use a small step size to reduce error due to the + # summation. This is slow, but tests the results well. + x = np.arange(0, 10, step=0.01) + + y_cdf = stats.studentized_range.cdf(x, k, v)[1:] + y_pdf_raw = stats.studentized_range.pdf(x, k, v) + y_pdf_cumulative = cumulative_trapezoid(y_pdf_raw, x) + + # Because of error caused by the summation, use a relatively large rtol + assert_allclose(y_pdf_cumulative, y_cdf, rtol=1e-4) + + @pytest.mark.parametrize("r_case_result", r_data) + def test_cdf_against_r(self, r_case_result): + # Test large `v` values using R + q, k, v, r_res = r_case_result + with np.errstate(invalid='ignore'): + res = stats.studentized_range.cdf(q, k, v) + assert_allclose(res, r_res) + + @pytest.mark.xslow + @pytest.mark.xfail_on_32bit("intermittent RuntimeWarning: invalid value.") + def test_moment_vectorization(self): + # Test moment broadcasting. Calls `_munp` directly because + # `rv_continuous.moment` is broken at time of writing. See gh-12192 + + # Silence invalid value encountered warnings. Actual problems will be + # caught by the result comparison. + with np.errstate(invalid='ignore'): + m = stats.studentized_range._munp([1, 2], [4, 5], [10, 11]) + + assert_allclose(m.shape, (2,)) + + with pytest.raises(ValueError, match="...could not be broadcast..."): + stats.studentized_range._munp(1, [4, 5], [10, 11, 12]) + + @pytest.mark.xslow + def test_fitstart_valid(self): + with warnings.catch_warnings(), np.errstate(invalid="ignore"): + # the integration warning message may differ + warnings.simplefilter("ignore", IntegrationWarning) + k, df, _, _ = stats.studentized_range._fitstart([1, 2, 3]) + assert_(stats.studentized_range._argcheck(k, df)) + + def test_infinite_df(self): + # Check that the CDF and PDF infinite and normal integrators + # roughly match for a high df case + res = stats.studentized_range.pdf(3, 10, np.inf) + res_finite = stats.studentized_range.pdf(3, 10, 99999) + assert_allclose(res, res_finite, atol=1e-4, rtol=1e-4) + + res = stats.studentized_range.cdf(3, 10, np.inf) + res_finite = stats.studentized_range.cdf(3, 10, 99999) + assert_allclose(res, res_finite, atol=1e-4, rtol=1e-4) + + def test_df_cutoff(self): + # Test that the CDF and PDF properly switch integrators at df=100,000. + # The infinite integrator should be different enough that it fails + # an allclose assertion. Also sanity check that using the same + # integrator does pass the allclose with a 1-df difference, which + # should be tiny. + + res = stats.studentized_range.pdf(3, 10, 100000) + res_finite = stats.studentized_range.pdf(3, 10, 99999) + res_sanity = stats.studentized_range.pdf(3, 10, 99998) + assert_raises(AssertionError, assert_allclose, res, res_finite, + atol=1e-6, rtol=1e-6) + assert_allclose(res_finite, res_sanity, atol=1e-6, rtol=1e-6) + + res = stats.studentized_range.cdf(3, 10, 100000) + res_finite = stats.studentized_range.cdf(3, 10, 99999) + res_sanity = stats.studentized_range.cdf(3, 10, 99998) + assert_raises(AssertionError, assert_allclose, res, res_finite, + atol=1e-6, rtol=1e-6) + assert_allclose(res_finite, res_sanity, atol=1e-6, rtol=1e-6) + + def test_clipping(self): + # The result of this computation was -9.9253938401489e-14 on some + # systems. The correct result is very nearly zero, but should not be + # negative. + q, k, v = 34.6413996195345746, 3, 339 + p = stats.studentized_range.sf(q, k, v) + assert_allclose(p, 0, atol=1e-10) + assert p >= 0 + + +class TestTukeyLambda: + + @pytest.mark.parametrize( + 'lam', + [0.0, -1.0, -2.0, np.array([[-1.0], [0.0], [-2.0]])] + ) + def test_pdf_nonpositive_lambda(self, lam): + # Make sure that Tukey-Lambda distribution correctly handles + # non-positive lambdas. + # This is a crude test--it just checks that all the PDF values + # are finite and greater than 0. + x = np.linspace(-5.0, 5.0, 101) + p = stats.tukeylambda.pdf(x, lam) + assert np.isfinite(p).all() + assert (p > 0.0).all() + + def test_pdf_mixed_lambda(self): + # Another crude test of the behavior of the PDF method. + x = np.linspace(-5.0, 5.0, 101) + lam = np.array([[-1.0], [0.0], [2.0]]) + p = stats.tukeylambda.pdf(x, lam) + assert np.isfinite(p).all() + # For p[0] and p[1], where lam <= 0, the support is (-inf, inf), + # so the PDF should be nonzero everywhere (assuming we aren't so + # far in the tails that we get underflow). + assert (p[:2] > 0.0).all() + # For p[2], where lam=2.0, the support is [-0.5, 0.5], so in pdf(x), + # some values should be positive and some should be 0. + assert (p[2] > 0.0).any() + assert (p[2] == 0.0).any() + + def test_support(self): + lam = np.array([-1.75, -0.5, 0.0, 0.25, 0.5, 2.0]) + a, b = stats.tukeylambda.support(lam) + expected_b = np.array([np.inf, np.inf, np.inf, 4, 2, 0.5]) + assert_equal(b, expected_b) + assert_equal(a, -expected_b) + + def test_pdf_support_boundary(self): + # Verify that tukeylambda.pdf() doesn't generate a + # warning when evaluated at the bounds of the support. + # For lam=0.5, the support is (-2, 2). + p = stats.tukeylambda.pdf([-2.0, 2.0], 0.5) + assert_equal(p, [0.0, 0.0]) + + def test_tukeylambda_stats_ticket_1545(self): + # Some test for the variance and kurtosis of the Tukey Lambda distr. + # See test_tukeylamdba_stats.py for more tests. + + mv = stats.tukeylambda.stats(0, moments='mvsk') + # Known exact values: + expected = [0, np.pi**2/3, 0, 1.2] + assert_almost_equal(mv, expected, decimal=10) + + mv = stats.tukeylambda.stats(3.13, moments='mvsk') + # 'expected' computed with mpmath. + expected = [0, 0.0269220858861465102, 0, -0.898062386219224104] + assert_almost_equal(mv, expected, decimal=10) + + mv = stats.tukeylambda.stats(0.14, moments='mvsk') + # 'expected' computed with mpmath. + expected = [0, 2.11029702221450250, 0, -0.02708377353223019456] + assert_almost_equal(mv, expected, decimal=10) + + @pytest.mark.xfail(reason="Precision in tail hasn't been fixed in xsf") + def test_large_argument_ticket_21370(self): + # Check that survival function goes to 0 as x gets large + x = [1e4, 1e6, 1e8, 1e10, 1e12, 1e24] + lam = -0.5 + sf = stats.tukeylambda.sf(x, lam) + # Check that the sf is 0 for large x + assert sf[-1] == 0.0 + # Check that the sf is always decreasing except when it is 0 + assert_array_less(np.diff(sf[sf > 0]), 0) + + cdf = stats.tukeylambda.cdf(x, lam) + assert cdf[-1] == 1.0 + # Check that the cdf is always increasing except when it saturates + # to 1.0 + assert np.all(np.diff(cdf[cdf < 1]) > 0) + +class TestLevy: + + def test_levy_cdf_ppf(self): + # Test levy.cdf, including small arguments. + x = np.array([1000, 1.0, 0.5, 0.1, 0.01, 0.001]) + + # Expected values were calculated separately with mpmath. + # E.g. + # >>> mpmath.mp.dps = 100 + # >>> x = mpmath.mp.mpf('0.01') + # >>> cdf = mpmath.erfc(mpmath.sqrt(1/(2*x))) + expected = np.array([0.9747728793699604, + 0.3173105078629141, + 0.1572992070502851, + 0.0015654022580025495, + 1.523970604832105e-23, + 1.795832784800726e-219]) + + y = stats.levy.cdf(x) + assert_allclose(y, expected, rtol=1e-10) + + # ppf(expected) should get us back to x. + xx = stats.levy.ppf(expected) + assert_allclose(xx, x, rtol=1e-13) + + def test_levy_sf(self): + # Large values, far into the tail of the distribution. + x = np.array([1e15, 1e25, 1e35, 1e50]) + # Expected values were calculated with mpmath. + expected = np.array([2.5231325220201597e-08, + 2.52313252202016e-13, + 2.52313252202016e-18, + 7.978845608028653e-26]) + y = stats.levy.sf(x) + assert_allclose(y, expected, rtol=1e-14) + + # The expected values for levy.isf(p) were calculated with mpmath. + # For loc=0 and scale=1, the inverse SF can be computed with + # + # import mpmath + # + # def levy_invsf(p): + # return 1/(2*mpmath.erfinv(p)**2) + # + # For example, with mpmath.mp.dps set to 60, float(levy_invsf(1e-20)) + # returns 6.366197723675814e+39. + # + @pytest.mark.parametrize('p, expected_isf', + [(1e-20, 6.366197723675814e+39), + (1e-8, 6366197723675813.0), + (0.375, 4.185810119346273), + (0.875, 0.42489442055310134), + (0.999, 0.09235685880262713), + (0.9999999962747097, 0.028766845244146945)]) + def test_levy_isf(self, p, expected_isf): + x = stats.levy.isf(p) + assert_allclose(x, expected_isf, atol=5e-15) + + def test_levy_logcdf(self): + x = 1e50 + ref = -7.978845608028653e-26 + logcdf = stats.levy.logcdf(x) + assert_allclose(logcdf, ref, rtol=5e-15) + + def test_levy_logsf(self): + x = 5e-3 + ref = -2.0884875837625492e-45 + logsf = stats.levy.logsf(x) + assert_allclose(logsf, ref, rtol=5e-15) + + +def test_540_567(): + # test for nan returned in tickets 540, 567 + assert_almost_equal(stats.norm.cdf(-1.7624320982), 0.03899815971089126, + decimal=10, err_msg='test_540_567') + assert_almost_equal(stats.norm.cdf(-1.7624320983), 0.038998159702449846, + decimal=10, err_msg='test_540_567') + assert_almost_equal(stats.norm.cdf(1.38629436112, loc=0.950273420309, + scale=0.204423758009), + 0.98353464004309321, + decimal=10, err_msg='test_540_567') + + +@pytest.mark.skipif(DOCSTRINGS_STRIPPED, reason="docstrings stripped") +def test_regression_ticket_1421(): + assert_('pdf(x, mu, loc=0, scale=1)' not in stats.poisson.__doc__) + assert_('pmf(x,' in stats.poisson.__doc__) + + +def test_nan_arguments_gh_issue_1362(): + with np.errstate(invalid='ignore'): + assert_(np.isnan(stats.t.logcdf(1, np.nan))) + assert_(np.isnan(stats.t.cdf(1, np.nan))) + assert_(np.isnan(stats.t.logsf(1, np.nan))) + assert_(np.isnan(stats.t.sf(1, np.nan))) + assert_(np.isnan(stats.t.pdf(1, np.nan))) + assert_(np.isnan(stats.t.logpdf(1, np.nan))) + assert_(np.isnan(stats.t.ppf(1, np.nan))) + assert_(np.isnan(stats.t.isf(1, np.nan))) + + assert_(np.isnan(stats.bernoulli.logcdf(np.nan, 0.5))) + assert_(np.isnan(stats.bernoulli.cdf(np.nan, 0.5))) + assert_(np.isnan(stats.bernoulli.logsf(np.nan, 0.5))) + assert_(np.isnan(stats.bernoulli.sf(np.nan, 0.5))) + assert_(np.isnan(stats.bernoulli.pmf(np.nan, 0.5))) + assert_(np.isnan(stats.bernoulli.logpmf(np.nan, 0.5))) + assert_(np.isnan(stats.bernoulli.ppf(np.nan, 0.5))) + assert_(np.isnan(stats.bernoulli.isf(np.nan, 0.5))) + + +def test_frozen_fit_ticket_1536(): + rng = np.random.default_rng(5678) + true = np.array([0.25, 0., 0.5]) + x = stats.lognorm.rvs(true[0], true[1], true[2], size=100, random_state=rng) + + with np.errstate(divide='ignore'): + params = np.array(stats.lognorm.fit(x, floc=0.)) + + assert_almost_equal(params, true, decimal=2) + + params = np.array(stats.lognorm.fit(x, fscale=0.5, loc=0)) + assert_almost_equal(params, true, decimal=2) + + params = np.array(stats.lognorm.fit(x, f0=0.25, loc=0)) + assert_almost_equal(params, true, decimal=2) + + params = np.array(stats.lognorm.fit(x, f0=0.25, floc=0)) + assert_almost_equal(params, true, decimal=2) + + rng = np.random.default_rng(5678) + loc = 1 + floc = 0.9 + x = stats.norm.rvs(loc, 2., size=100, random_state=rng) + params = np.array(stats.norm.fit(x, floc=floc)) + expected = np.array([floc, np.sqrt(((x-floc)**2).mean())]) + assert_almost_equal(params, expected, decimal=4) + + +def test_regression_ticket_1530(): + # Check the starting value works for Cauchy distribution fit. + rng = np.random.default_rng(654321) + rvs = stats.cauchy.rvs(size=100, random_state=rng) + params = stats.cauchy.fit(rvs) + expected = (0.045, 1.142) + assert_almost_equal(params, expected, decimal=1) + + +def test_gh_pr_4806(): + # Check starting values for Cauchy distribution fit. + rng = np.random.RandomState(1234) + x = rng.randn(42) + for offset in 10000.0, 1222333444.0: + loc, scale = stats.cauchy.fit(x + offset) + assert_allclose(loc, offset, atol=1.0) + assert_allclose(scale, 0.6, atol=1.0) + + +def test_poisson_logpmf_ticket_1436(): + assert_(np.isfinite(stats.poisson.logpmf(1500, 200))) + + +def test_powerlaw_stats(): + """Test the powerlaw stats function. + + This unit test is also a regression test for ticket 1548. + + The exact values are: + mean: + mu = a / (a + 1) + variance: + sigma**2 = a / ((a + 2) * (a + 1) ** 2) + skewness: + One formula (see https://en.wikipedia.org/wiki/Skewness) is + gamma_1 = (E[X**3] - 3*mu*E[X**2] + 2*mu**3) / sigma**3 + A short calculation shows that E[X**k] is a / (a + k), so gamma_1 + can be implemented as + n = a/(a+3) - 3*(a/(a+1))*a/(a+2) + 2*(a/(a+1))**3 + d = sqrt(a/((a+2)*(a+1)**2)) ** 3 + gamma_1 = n/d + Either by simplifying, or by a direct calculation of mu_3 / sigma**3, + one gets the more concise formula: + gamma_1 = -2.0 * ((a - 1) / (a + 3)) * sqrt((a + 2) / a) + kurtosis: (See https://en.wikipedia.org/wiki/Kurtosis) + The excess kurtosis is + gamma_2 = mu_4 / sigma**4 - 3 + A bit of calculus and algebra (sympy helps) shows that + mu_4 = 3*a*(3*a**2 - a + 2) / ((a+1)**4 * (a+2) * (a+3) * (a+4)) + so + gamma_2 = 3*(3*a**2 - a + 2) * (a+2) / (a*(a+3)*(a+4)) - 3 + which can be rearranged to + gamma_2 = 6 * (a**3 - a**2 - 6*a + 2) / (a*(a+3)*(a+4)) + """ + cases = [(1.0, (0.5, 1./12, 0.0, -1.2)), + (2.0, (2./3, 2./36, -0.56568542494924734, -0.6))] + for a, exact_mvsk in cases: + mvsk = stats.powerlaw.stats(a, moments="mvsk") + assert_array_almost_equal(mvsk, exact_mvsk) + + +def test_powerlaw_edge(): + # Regression test for gh-3986. + p = stats.powerlaw.logpdf(0, 1) + assert_equal(p, 0.0) + + +def test_exponpow_edge(): + # Regression test for gh-3982. + p = stats.exponpow.logpdf(0, 1) + assert_equal(p, 0.0) + + # Check pdf and logpdf at x = 0 for other values of b. + p = stats.exponpow.pdf(0, [0.25, 1.0, 1.5]) + assert_equal(p, [np.inf, 1.0, 0.0]) + p = stats.exponpow.logpdf(0, [0.25, 1.0, 1.5]) + assert_equal(p, [np.inf, 0.0, -np.inf]) + + +class TestGenGamma: + def test_gengamma_edge(self): + # Regression test for gh-3985. + p = stats.gengamma.pdf(0, 1, 1) + assert_equal(p, 1.0) + + @pytest.mark.parametrize("a, c, ref, tol", + [(1500000.0, 1, 8.529426144018633, 1e-15), + (1e+30, 1, 35.95771492811536, 1e-15), + (1e+100, 1, 116.54819318290696, 1e-15), + (3e3, 1, 5.422011196659015, 1e-13), + (3e6, -1e100, -236.29663213396054, 1e-15), + (3e60, 1e-100, 1.3925371786831085e+102, 1e-15)]) + def test_gengamma_extreme_entropy(self, a, c, ref, tol): + # The reference values were calculated with mpmath: + # from mpmath import mp + # mp.dps = 500 + # + # def gen_entropy(a, c): + # a, c = mp.mpf(a), mp.mpf(c) + # val = mp.digamma(a) + # h = (a * (mp.one - val) + val/c + mp.loggamma(a) - mp.log(abs(c))) + # return float(h) + assert_allclose(stats.gengamma.entropy(a, c), ref, rtol=tol) + + def test_gengamma_endpoint_with_neg_c(self): + p = stats.gengamma.pdf(0, 1, -1) + assert p == 0.0 + logp = stats.gengamma.logpdf(0, 1, -1) + assert logp == -np.inf + + def test_gengamma_munp(self): + # Regression tests for gh-4724. + p = stats.gengamma._munp(-2, 200, 1.) + assert_almost_equal(p, 1./199/198) + + p = stats.gengamma._munp(-2, 10, 1.) + assert_almost_equal(p, 1./9/8) + + def test_gengamma_logpdf_broadcasting_gh24574(self): + # gh-24574 reported a broadcasting error when `x` included 0s. + assert_allclose(stats.gengamma.logpdf([0, 1, 1], 1, -1), [-np.inf, -1, -1]) + + +def test_ksone_fit_freeze(): + # Regression test for ticket #1638. + d = np.array( + [-0.18879233, 0.15734249, 0.18695107, 0.27908787, -0.248649, + -0.2171497, 0.12233512, 0.15126419, 0.03119282, 0.4365294, + 0.08930393, -0.23509903, 0.28231224, -0.09974875, -0.25196048, + 0.11102028, 0.1427649, 0.10176452, 0.18754054, 0.25826724, + 0.05988819, 0.0531668, 0.21906056, 0.32106729, 0.2117662, + 0.10886442, 0.09375789, 0.24583286, -0.22968366, -0.07842391, + -0.31195432, -0.21271196, 0.1114243, -0.13293002, 0.01331725, + -0.04330977, -0.09485776, -0.28434547, 0.22245721, -0.18518199, + -0.10943985, -0.35243174, 0.06897665, -0.03553363, -0.0701746, + -0.06037974, 0.37670779, -0.21684405]) + + with np.errstate(invalid='ignore'): + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", + "The maximum number of subdivisions .50. has been achieved.", + IntegrationWarning, + ) + warnings.filterwarnings( + "ignore", + "floating point number truncated to an integer", + RuntimeWarning, + ) + stats.ksone.fit(d) + + +@pytest.mark.xfail(platform.machine() == "sparc64", + reason="complex special functions known to fail on sparc64 (gh-22577)") +def test_norm_logcdf(): + # Test precision of the logcdf of the normal distribution. + # This precision was enhanced in ticket 1614. + x = -np.asarray(list(range(0, 120, 4))) + # Values from R + expected = [-0.69314718, -10.36010149, -35.01343716, -75.41067300, + -131.69539607, -203.91715537, -292.09872100, -396.25241451, + -516.38564863, -652.50322759, -804.60844201, -972.70364403, + -1156.79057310, -1356.87055173, -1572.94460885, -1805.01356068, + -2053.07806561, -2317.13866238, -2597.19579746, -2893.24984493, + -3205.30112136, -3533.34989701, -3877.39640444, -4237.44084522, + -4613.48339520, -5005.52420869, -5413.56342187, -5837.60115548, + -6277.63751711, -6733.67260303] + + assert_allclose(stats.norm().logcdf(x), expected, atol=1e-8) + + # also test the complex-valued code path + assert_allclose(stats.norm().logcdf(x + 1e-14j).real, expected, atol=1e-8) + + # test the accuracy: d(logcdf)/dx = pdf / cdf \equiv exp(logpdf - logcdf) + deriv = (stats.norm.logcdf(x + 1e-10j)/1e-10).imag + deriv_expected = np.exp(stats.norm.logpdf(x) - stats.norm.logcdf(x)) + assert_allclose(deriv, deriv_expected, atol=1e-10) + + +def test_levy_l_sf(): + # Test levy_l.sf for small arguments. + x = np.array([-0.016, -0.01, -0.005, -0.0015]) + # Expected values were calculated with mpmath. + expected = np.array([2.6644463892359302e-15, + 1.523970604832107e-23, + 2.0884875837625492e-45, + 5.302850374626878e-147]) + y = stats.levy_l.sf(x) + assert_allclose(y, expected, rtol=1e-13) + + +def test_levy_l_isf(): + # Test roundtrip sf(isf(p)), including a small input value. + p = np.array([3.0e-15, 0.25, 0.99]) + x = stats.levy_l.isf(p) + q = stats.levy_l.sf(x) + assert_allclose(q, p, rtol=5e-14) + + +def test_hypergeom_interval_1802(): + # these two had endless loops + assert_equal(stats.hypergeom.interval(.95, 187601, 43192, 757), + (152.0, 197.0)) + assert_equal(stats.hypergeom.interval(.945, 187601, 43192, 757), + (152.0, 197.0)) + # this was working also before + assert_equal(stats.hypergeom.interval(.94, 187601, 43192, 757), + (153.0, 196.0)) + + # degenerate case .a == .b + assert_equal(stats.hypergeom.ppf(0.02, 100, 100, 8), 8) + assert_equal(stats.hypergeom.ppf(1, 100, 100, 8), 8) + + +def test_distribution_too_many_args(): + rng = np.random.default_rng(5976604568) + # Check that a TypeError is raised when too many args are given to a method + # Regression test for ticket 1815. + x = np.linspace(0.1, 0.7, num=5) + assert_raises(TypeError, stats.gamma.pdf, x, 2, 3, loc=1.0) + assert_raises(TypeError, stats.gamma.pdf, x, 2, 3, 4, loc=1.0) + assert_raises(TypeError, stats.gamma.pdf, x, 2, 3, 4, 5) + assert_raises(TypeError, stats.gamma.pdf, x, 2, 3, loc=1.0, scale=0.5) + assert_raises(TypeError, stats.gamma.rvs, 2., 3, loc=1.0, scale=0.5, + random_state=rng) + assert_raises(TypeError, stats.gamma.cdf, x, 2., 3, loc=1.0, scale=0.5) + assert_raises(TypeError, stats.gamma.ppf, x, 2., 3, loc=1.0, scale=0.5) + assert_raises(TypeError, stats.gamma.stats, 2., 3, loc=1.0, scale=0.5) + assert_raises(TypeError, stats.gamma.entropy, 2., 3, loc=1.0, scale=0.5) + assert_raises(TypeError, stats.gamma.fit, x, 2., 3, loc=1.0, scale=0.5) + + # These should not give errors + stats.gamma.pdf(x, 2, 3) # loc=3 + stats.gamma.pdf(x, 2, 3, 4) # loc=3, scale=4 + stats.gamma.stats(2., 3) + stats.gamma.stats(2., 3, 4) + stats.gamma.stats(2., 3, 4, 'mv') + stats.gamma.rvs(2., 3, 4, 5, random_state=rng) + stats.gamma.fit(stats.gamma.rvs(2., size=7, random_state=rng), 2.) + + # Also for a discrete distribution + stats.geom.pmf(x, 2, loc=3) # no error, loc=3 + assert_raises(TypeError, stats.geom.pmf, x, 2, 3, 4) + assert_raises(TypeError, stats.geom.pmf, x, 2, 3, loc=4) + + # And for distributions with 0, 2 and 3 args respectively + assert_raises(TypeError, stats.expon.pdf, x, 3, loc=1.0) + assert_raises(TypeError, stats.exponweib.pdf, x, 3, 4, 5, loc=1.0) + assert_raises(TypeError, stats.exponweib.pdf, x, 3, 4, 5, 0.1, 0.1) + assert_raises(TypeError, stats.ncf.pdf, x, 3, 4, 5, 6, loc=1.0) + assert_raises(TypeError, stats.ncf.pdf, x, 3, 4, 5, 6, 1.0, scale=0.5) + stats.ncf.pdf(x, 3, 4, 5, 6, 1.0) # 3 args, plus loc/scale + + +def test_ncx2_tails_ticket_955(): + # Trac #955 -- check that the cdf computed by special functions + # matches the integrated pdf + a = stats.ncx2.cdf(np.arange(20, 25, 0.2), 2, 1.07458615e+02) + b = stats.ncx2._cdfvec(np.arange(20, 25, 0.2), 2, 1.07458615e+02) + assert_allclose(a, b, rtol=1e-3, atol=0) + + +def test_ncx2_tails_pdf(): + # ncx2.pdf does not return nans in extreme tails(example from gh-1577) + # NB: this is to check that nan_to_num is not needed in ncx2.pdf + with warnings.catch_warnings(): + warnings.simplefilter('error', RuntimeWarning) + assert_equal(stats.ncx2.pdf(1, np.arange(340, 350), 2), 0) + logval = stats.ncx2.logpdf(1, np.arange(340, 350), 2) + + assert_(np.isneginf(logval).all()) + + # Verify logpdf has extended precision when pdf underflows to 0 + with warnings.catch_warnings(): + warnings.simplefilter('error', RuntimeWarning) + assert_equal(stats.ncx2.pdf(10000, 3, 12), 0) + assert_allclose(stats.ncx2.logpdf(10000, 3, 12), -4662.444377524883) + + +@pytest.mark.parametrize('method, expected', [ + ('cdf', np.array([2.497951336e-09, 3.437288941e-10])), + ('pdf', np.array([1.238579980e-07, 1.710041145e-08])), + ('logpdf', np.array([-15.90413011, -17.88416331])), + ('ppf', np.array([4.865182052, 7.017182271])) +]) +def test_ncx2_zero_nc(method, expected): + # gh-5441 + # ncx2 with nc=0 is identical to chi2 + # Comparison to R (v3.5.1) + # > options(digits=10) + # > pchisq(0.1, df=10, ncp=c(0,4)) + # > dchisq(0.1, df=10, ncp=c(0,4)) + # > dchisq(0.1, df=10, ncp=c(0,4), log=TRUE) + # > qchisq(0.1, df=10, ncp=c(0,4)) + + result = getattr(stats.ncx2, method)(0.1, nc=[0, 4], df=10) + assert_allclose(result, expected, atol=1e-15) + + +def test_ncx2_zero_nc_rvs(): + # gh-5441 + # ncx2 with nc=0 is identical to chi2 + result = stats.ncx2.rvs(df=10, nc=0, random_state=1) + expected = stats.chi2.rvs(df=10, random_state=1) + assert_allclose(result, expected, atol=1e-15) + + +def test_ncx2_gh12731(): + # test that gh-12731 is resolved; previously these were all 0.5 + nc = 10**np.arange(5, 10) + assert_equal(stats.ncx2.cdf(1e4, df=1, nc=nc), 0) + + +def test_ncx2_gh8665(): + # test that gh-8665 is resolved; previously this tended to nonzero value + x = np.array([4.99515382e+00, 1.07617327e+01, 2.31854502e+01, + 4.99515382e+01, 1.07617327e+02, 2.31854502e+02, + 4.99515382e+02, 1.07617327e+03, 2.31854502e+03, + 4.99515382e+03, 1.07617327e+04, 2.31854502e+04, + 4.99515382e+04]) + nu, lam = 20, 499.51538166556196 + + sf = stats.ncx2.sf(x, df=nu, nc=lam) + # computed in R. Couldn't find a survival function implementation + # options(digits=16) + # x <- c(4.99515382e+00, 1.07617327e+01, 2.31854502e+01, 4.99515382e+01, + # 1.07617327e+02, 2.31854502e+02, 4.99515382e+02, 1.07617327e+03, + # 2.31854502e+03, 4.99515382e+03, 1.07617327e+04, 2.31854502e+04, + # 4.99515382e+04) + # nu <- 20 + # lam <- 499.51538166556196 + # 1 - pchisq(x, df = nu, ncp = lam) + sf_expected = [1.0000000000000000, 1.0000000000000000, 1.0000000000000000, + 1.0000000000000000, 1.0000000000000000, 0.9999999999999888, + 0.6646525582135460, 0.0000000000000000, 0.0000000000000000, + 0.0000000000000000, 0.0000000000000000, 0.0000000000000000, + 0.0000000000000000] + assert_allclose(sf, sf_expected, atol=1e-12) + + +def test_ncx2_gh11777(): + # regression test for gh-11777: + # At high values of degrees of freedom df, ensure the pdf of ncx2 does + # not get clipped to zero when the non-centrality parameter is + # sufficiently less than df + df = 6700 + nc = 5300 + x = np.linspace(stats.ncx2.ppf(0.001, df, nc), + stats.ncx2.ppf(0.999, df, nc), num=10000) + ncx2_pdf = stats.ncx2.pdf(x, df, nc) + gauss_approx = stats.norm.pdf(x, df + nc, np.sqrt(2 * df + 4 * nc)) + # use huge tolerance as we're only looking for obvious discrepancy + assert_allclose(ncx2_pdf, gauss_approx, atol=1e-4) + + +# Expected values for foldnorm.sf were computed with mpmath: +# +# from mpmath import mp +# mp.dps = 60 +# def foldcauchy_sf(x, c): +# x = mp.mpf(x) +# c = mp.mpf(c) +# return mp.one - (mp.atan(x - c) + mp.atan(x + c))/mp.pi +# +# E.g. +# +# >>> float(foldcauchy_sf(2, 1)) +# 0.35241638234956674 +# +@pytest.mark.parametrize('x, c, expected', + [(2, 1, 0.35241638234956674), + (2, 2, 0.5779791303773694), + (1e13, 1, 6.366197723675813e-14), + (2e16, 1, 3.183098861837907e-17), + (1e13, 2e11, 6.368745221764519e-14), + (0.125, 200, 0.999998010612169)]) +def test_foldcauchy_sf(x, c, expected): + sf = stats.foldcauchy.sf(x, c) + assert_allclose(sf, expected, 2e-15) + + +# The same mpmath code shown in the comments above test_foldcauchy_sf() +# is used to create these expected values. +@pytest.mark.parametrize('x, expected', + [(2, 0.2951672353008665), + (1e13, 6.366197723675813e-14), + (2e16, 3.183098861837907e-17), + (5e80, 1.2732395447351629e-81)]) +def test_halfcauchy_sf(x, expected): + sf = stats.halfcauchy.sf(x) + assert_allclose(sf, expected, 2e-15) + + +# Expected value computed with mpmath: +# expected = mp.cot(mp.pi*p/2) +@pytest.mark.parametrize('p, expected', + [(0.9999995, 7.853981633329977e-07), + (0.975, 0.039290107007669675), + (0.5, 1.0), + (0.01, 63.65674116287158), + (1e-14, 63661977236758.13), + (5e-80, 1.2732395447351627e+79)]) +def test_halfcauchy_isf(p, expected): + x = stats.halfcauchy.isf(p) + assert_allclose(x, expected) + + +def test_foldnorm_zero(): + # Parameter value c=0 was not enabled, see gh-2399. + rv = stats.foldnorm(0, scale=1) + assert_equal(rv.cdf(0), 0) # rv.cdf(0) previously resulted in: nan + + +# Expected values for foldnorm.sf were computed with mpmath: +# +# from mpmath import mp +# mp.dps = 60 +# def foldnorm_sf(x, c): +# x = mp.mpf(x) +# c = mp.mpf(c) +# return mp.ncdf(-x+c) + mp.ncdf(-x-c) +# +# E.g. +# +# >>> float(foldnorm_sf(2, 1)) +# 0.16000515196308715 +# +@pytest.mark.parametrize('x, c, expected', + [(2, 1, 0.16000515196308715), + (20, 1, 8.527223952630977e-81), + (10, 15, 0.9999997133484281), + (25, 15, 7.619853024160525e-24)]) +def test_foldnorm_sf(x, c, expected): + sf = stats.foldnorm.sf(x, c) + assert_allclose(sf, expected, 1e-14) + + +def test_stats_shapes_argcheck(): + # stats method was failing for vector shapes if some of the values + # were outside of the allowed range, see gh-2678 + mv3 = stats.invgamma.stats([0.0, 0.5, 1.0], 1, 0.5) # 0 is not a legal `a` + mv2 = stats.invgamma.stats([0.5, 1.0], 1, 0.5) + mv2_augmented = tuple(np.r_[np.nan, _] for _ in mv2) + assert_equal(mv2_augmented, mv3) + + # -1 is not a legal shape parameter + mv3 = stats.lognorm.stats([2, 2.4, -1]) + mv2 = stats.lognorm.stats([2, 2.4]) + mv2_augmented = tuple(np.r_[_, np.nan] for _ in mv2) + assert_equal(mv2_augmented, mv3) + + # FIXME: this is only a quick-and-dirty test of a quick-and-dirty bugfix. + # stats method with multiple shape parameters is not properly vectorized + # anyway, so some distributions may or may not fail. + + +# Test subclassing distributions w/ explicit shapes + +class _distr_gen(stats.rv_continuous): + def _pdf(self, x, a): + return 42 + + +class _distr2_gen(stats.rv_continuous): + def _cdf(self, x, a): + return 42 * a + x + + +class _distr3_gen(stats.rv_continuous): + def _pdf(self, x, a, b): + return a + b + + def _cdf(self, x, a): + # Different # of shape params from _pdf, to be able to check that + # inspection catches the inconsistency. + return 42 * a + x + + +class _distr6_gen(stats.rv_continuous): + # Two shape parameters (both _pdf and _cdf defined, consistent shapes.) + def _pdf(self, x, a, b): + return a*x + b + + def _cdf(self, x, a, b): + return 42 * a + x + + +class TestSubclassingExplicitShapes: + # Construct a distribution w/ explicit shapes parameter and test it. + + def test_correct_shapes(self): + dummy_distr = _distr_gen(name='dummy', shapes='a') + assert_equal(dummy_distr.pdf(1, a=1), 42) + + def test_wrong_shapes_1(self): + dummy_distr = _distr_gen(name='dummy', shapes='A') + assert_raises(TypeError, dummy_distr.pdf, 1, **dict(a=1)) + + def test_wrong_shapes_2(self): + dummy_distr = _distr_gen(name='dummy', shapes='a, b, c') + dct = dict(a=1, b=2, c=3) + assert_raises(TypeError, dummy_distr.pdf, 1, **dct) + + def test_shapes_string(self): + # shapes must be a string + dct = dict(name='dummy', shapes=42) + assert_raises(TypeError, _distr_gen, **dct) + + def test_shapes_identifiers_1(self): + # shapes must be a comma-separated list of valid python identifiers + dct = dict(name='dummy', shapes='(!)') + assert_raises(SyntaxError, _distr_gen, **dct) + + def test_shapes_identifiers_2(self): + dct = dict(name='dummy', shapes='4chan') + assert_raises(SyntaxError, _distr_gen, **dct) + + def test_shapes_identifiers_3(self): + dct = dict(name='dummy', shapes='m(fti)') + assert_raises(SyntaxError, _distr_gen, **dct) + + def test_shapes_identifiers_nodefaults(self): + dct = dict(name='dummy', shapes='a=2') + assert_raises(SyntaxError, _distr_gen, **dct) + + def test_shapes_args(self): + dct = dict(name='dummy', shapes='*args') + assert_raises(SyntaxError, _distr_gen, **dct) + + def test_shapes_kwargs(self): + dct = dict(name='dummy', shapes='**kwargs') + assert_raises(SyntaxError, _distr_gen, **dct) + + def test_shapes_keywords(self): + # python keywords cannot be used for shape parameters + dct = dict(name='dummy', shapes='a, b, c, lambda') + assert_raises(SyntaxError, _distr_gen, **dct) + + def test_shapes_signature(self): + # test explicit shapes which agree w/ the signature of _pdf + class _dist_gen(stats.rv_continuous): + def _pdf(self, x, a): + return stats.norm._pdf(x) * a + + dist = _dist_gen(shapes='a') + assert_equal(dist.pdf(0.5, a=2), stats.norm.pdf(0.5)*2) + + def test_shapes_signature_inconsistent(self): + # test explicit shapes which do not agree w/ the signature of _pdf + class _dist_gen(stats.rv_continuous): + def _pdf(self, x, a): + return stats.norm._pdf(x) * a + + dist = _dist_gen(shapes='a, b') + assert_raises(TypeError, dist.pdf, 0.5, **dict(a=1, b=2)) + + def test_star_args(self): + # test _pdf with only starargs + # NB: **kwargs of pdf will never reach _pdf + class _dist_gen(stats.rv_continuous): + def _pdf(self, x, *args): + extra_kwarg = args[0] + return stats.norm._pdf(x) * extra_kwarg + + dist = _dist_gen(shapes='extra_kwarg') + assert_equal(dist.pdf(0.5, extra_kwarg=33), stats.norm.pdf(0.5)*33) + assert_equal(dist.pdf(0.5, 33), stats.norm.pdf(0.5)*33) + assert_raises(TypeError, dist.pdf, 0.5, **dict(xxx=33)) + + def test_star_args_2(self): + # test _pdf with named & starargs + # NB: **kwargs of pdf will never reach _pdf + class _dist_gen(stats.rv_continuous): + def _pdf(self, x, offset, *args): + extra_kwarg = args[0] + return stats.norm._pdf(x) * extra_kwarg + offset + + dist = _dist_gen(shapes='offset, extra_kwarg') + assert_equal(dist.pdf(0.5, offset=111, extra_kwarg=33), + stats.norm.pdf(0.5)*33 + 111) + assert_equal(dist.pdf(0.5, 111, 33), + stats.norm.pdf(0.5)*33 + 111) + + def test_extra_kwarg(self): + # **kwargs to _pdf are ignored. + # this is a limitation of the framework (_pdf(x, *goodargs)) + class _distr_gen(stats.rv_continuous): + def _pdf(self, x, *args, **kwargs): + # _pdf should handle *args, **kwargs itself. Here "handling" + # is ignoring *args and looking for ``extra_kwarg`` and using + # that. + extra_kwarg = kwargs.pop('extra_kwarg', 1) + return stats.norm._pdf(x) * extra_kwarg + + dist = _distr_gen(shapes='extra_kwarg') + assert_equal(dist.pdf(1, extra_kwarg=3), stats.norm.pdf(1)) + + def test_shapes_empty_string(self): + # shapes='' is equivalent to shapes=None + class _dist_gen(stats.rv_continuous): + def _pdf(self, x): + return stats.norm.pdf(x) + + dist = _dist_gen(shapes='') + assert_equal(dist.pdf(0.5), stats.norm.pdf(0.5)) + + +class TestSubclassingNoShapes: + # Construct a distribution w/o explicit shapes parameter and test it. + + def test_only__pdf(self): + dummy_distr = _distr_gen(name='dummy') + assert_equal(dummy_distr.pdf(1, a=1), 42) + + def test_only__cdf(self): + # _pdf is determined from _cdf by taking numerical derivative + dummy_distr = _distr2_gen(name='dummy') + assert_almost_equal(dummy_distr.pdf(1, a=1), 1) + + @pytest.mark.skipif(DOCSTRINGS_STRIPPED, reason="docstring stripped") + def test_signature_inspection(self): + # check that _pdf signature inspection works correctly, and is used in + # the class docstring + dummy_distr = _distr_gen(name='dummy') + assert_equal(dummy_distr.numargs, 1) + assert_equal(dummy_distr.shapes, 'a') + res = re.findall(r'logpdf\(x, a, loc=0, scale=1\)', + dummy_distr.__doc__) + assert_(len(res) == 1) + + @pytest.mark.skipif(DOCSTRINGS_STRIPPED, reason="docstring stripped") + def test_signature_inspection_2args(self): + # same for 2 shape params and both _pdf and _cdf defined + dummy_distr = _distr6_gen(name='dummy') + assert_equal(dummy_distr.numargs, 2) + assert_equal(dummy_distr.shapes, 'a, b') + res = re.findall(r'logpdf\(x, a, b, loc=0, scale=1\)', + dummy_distr.__doc__) + assert_(len(res) == 1) + + def test_signature_inspection_2args_incorrect_shapes(self): + # both _pdf and _cdf defined, but shapes are inconsistent: raises + assert_raises(TypeError, _distr3_gen, name='dummy') + + def test_defaults_raise(self): + # default arguments should raise + class _dist_gen(stats.rv_continuous): + def _pdf(self, x, a=42): + return 42 + assert_raises(TypeError, _dist_gen, **dict(name='dummy')) + + def test_starargs_raise(self): + # without explicit shapes, *args are not allowed + class _dist_gen(stats.rv_continuous): + def _pdf(self, x, a, *args): + return 42 + assert_raises(TypeError, _dist_gen, **dict(name='dummy')) + + def test_kwargs_raise(self): + # without explicit shapes, **kwargs are not allowed + class _dist_gen(stats.rv_continuous): + def _pdf(self, x, a, **kwargs): + return 42 + assert_raises(TypeError, _dist_gen, **dict(name='dummy')) + + +@pytest.mark.skipif(DOCSTRINGS_STRIPPED, reason="docstring stripped") +def test_docstrings(): + badones = [r',\s*,', r'\(\s*,', r'^\s*:'] + for distname in stats.__all__: + dist = getattr(stats, distname) + if isinstance(dist, (stats.rv_discrete | stats.rv_continuous)): + for regex in badones: + assert_(re.search(regex, dist.__doc__) is None) + + +def test_infinite_input(): + assert_almost_equal(stats.skellam.sf(np.inf, 10, 11), 0) + assert_almost_equal(stats.ncx2._cdf(np.inf, 8, 0.1), 1) + + +def test_lomax_accuracy(): + # regression test for gh-4033 + p = stats.lomax.ppf(stats.lomax.cdf(1e-100, 1), 1) + assert_allclose(p, 1e-100) + + +def test_truncexpon_accuracy(): + # regression test for gh-4035 + p = stats.truncexpon.ppf(stats.truncexpon.cdf(1e-100, 1), 1) + assert_allclose(p, 1e-100) + + +def test_rayleigh_accuracy(): + # regression test for gh-4034 + p = stats.rayleigh.isf(stats.rayleigh.sf(9, 1), 1) + assert_almost_equal(p, 9.0, decimal=15) + + +def test_genextreme_give_no_warnings(): + """regression test for gh-6219""" + + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + + stats.genextreme.cdf(.5, 0) + stats.genextreme.pdf(.5, 0) + stats.genextreme.ppf(.5, 0) + stats.genextreme.logpdf(-np.inf, 0.0) + number_of_warnings_thrown = len(w) + assert_equal(number_of_warnings_thrown, 0) + + +def test_moments_gh22400(): + # Regression test for gh-22400 + # Check for correct results at c=0 with no warnings. While we're at it, + # check that NaN and sufficiently negative input produce NaNs, and output + # with `c=1` also agrees with reference values. + res = np.asarray(stats.genextreme.stats([0.0, np.nan, 1, -1.5], moments='mvsk')) + + # Reference values for c=0 (Wikipedia) + mean = np.euler_gamma + var = np.pi**2 / 6 + skew = 12 * np.sqrt(6) * special.zeta(3) / np.pi**3 + kurt = 12 / 5 + ref_0 = [mean, var, skew, kurt] + ref_1 = ref_3 = [np.nan]*4 + ref_2 = [0, 1, -2, 6] # Wolfram Alpha, MaxStableDistribution[0, 1, -1] + + assert_allclose(res[:, 0], ref_0, rtol=1e-14) + assert_equal(res[:, 1], ref_1) + assert_allclose(res[:, 2], ref_2, rtol=1e-14) + assert_equal(res[:, 3], ref_3) + + +def test_genextreme_entropy(): + # regression test for gh-5181 + euler_gamma = 0.5772156649015329 + + h = stats.genextreme.entropy(-1.0) + assert_allclose(h, 2*euler_gamma + 1, rtol=1e-14) + + h = stats.genextreme.entropy(0) + assert_allclose(h, euler_gamma + 1, rtol=1e-14) + + h = stats.genextreme.entropy(1.0) + assert_equal(h, 1) + + h = stats.genextreme.entropy(-2.0, scale=10) + assert_allclose(h, euler_gamma*3 + np.log(10) + 1, rtol=1e-14) + + h = stats.genextreme.entropy(10) + assert_allclose(h, -9*euler_gamma + 1, rtol=1e-14) + + h = stats.genextreme.entropy(-10) + assert_allclose(h, 11*euler_gamma + 1, rtol=1e-14) + + +def test_genextreme_sf_isf(): + # Expected values were computed using mpmath: + # + # import mpmath + # + # def mp_genextreme_sf(x, xi, mu=0, sigma=1): + # # Formula from wikipedia, which has a sign convention for xi that + # # is the opposite of scipy's shape parameter. + # if xi != 0: + # t = mpmath.power(1 + ((x - mu)/sigma)*xi, -1/xi) + # else: + # t = mpmath.exp(-(x - mu)/sigma) + # return 1 - mpmath.exp(-t) + # + # >>> mpmath.mp.dps = 1000 + # >>> s = mp_genextreme_sf(mpmath.mp.mpf("1e8"), mpmath.mp.mpf("0.125")) + # >>> float(s) + # 1.6777205262585625e-57 + # >>> s = mp_genextreme_sf(mpmath.mp.mpf("7.98"), mpmath.mp.mpf("-0.125")) + # >>> float(s) + # 1.52587890625e-21 + # >>> s = mp_genextreme_sf(mpmath.mp.mpf("7.98"), mpmath.mp.mpf("0")) + # >>> float(s) + # 0.00034218086528426593 + + x = 1e8 + s = stats.genextreme.sf(x, -0.125) + assert_allclose(s, 1.6777205262585625e-57) + x2 = stats.genextreme.isf(s, -0.125) + assert_allclose(x2, x) + + x = 7.98 + s = stats.genextreme.sf(x, 0.125) + assert_allclose(s, 1.52587890625e-21) + x2 = stats.genextreme.isf(s, 0.125) + assert_allclose(x2, x) + + x = 7.98 + s = stats.genextreme.sf(x, 0) + assert_allclose(s, 0.00034218086528426593) + x2 = stats.genextreme.isf(s, 0) + assert_allclose(x2, x) + + +def test_burr12_ppf_small_arg(): + prob = 1e-16 + quantile = stats.burr12.ppf(prob, 2, 3) + # The expected quantile was computed using mpmath: + # >>> import mpmath + # >>> mpmath.mp.dps = 100 + # >>> prob = mpmath.mpf('1e-16') + # >>> c = mpmath.mpf(2) + # >>> d = mpmath.mpf(3) + # >>> float(((1-prob)**(-1/d) - 1)**(1/c)) + # 5.7735026918962575e-09 + assert_allclose(quantile, 5.7735026918962575e-09) + + +def test_invweibull_fit(): + """ + Test fitting invweibull to data. + + Here is a the same calculation in R: + + > library(evd) + > library(fitdistrplus) + > x = c(1, 1.25, 2, 2.5, 2.8, 3, 3.8, 4, 5, 8, 10, 12, 64, 99) + > result = fitdist(x, 'frechet', control=list(reltol=1e-13), + + fix.arg=list(loc=0), start=list(shape=2, scale=3)) + > result + Fitting of the distribution ' frechet ' by maximum likelihood + Parameters: + estimate Std. Error + shape 1.048482 0.2261815 + scale 3.099456 0.8292887 + Fixed parameters: + value + loc 0 + + """ + + def optimizer(func, x0, args=(), disp=0): + return fmin(func, x0, args=args, disp=disp, xtol=1e-12, ftol=1e-12) + + x = np.array([1, 1.25, 2, 2.5, 2.8, 3, 3.8, 4, 5, 8, 10, 12, 64, 99]) + c, loc, scale = stats.invweibull.fit(x, floc=0, optimizer=optimizer) + assert_allclose(c, 1.048482, rtol=5e-6) + assert loc == 0 + assert_allclose(scale, 3.099456, rtol=5e-6) + + +# Expected values were computed with mpmath. +@pytest.mark.parametrize('x, c, expected', + [(3, 1.5, 0.175064510070713299327), + (2000, 1.5, 1.11802773877318715787e-5), + (2000, 9.25, 2.92060308832269637092e-31), + (1e15, 1.5, 3.16227766016837933199884e-23)]) +def test_invweibull_sf(x, c, expected): + computed = stats.invweibull.sf(x, c) + assert_allclose(computed, expected, rtol=1e-15) + + +# Expected values were computed with mpmath. +@pytest.mark.parametrize('p, c, expected', + [(0.5, 2.5, 1.15789669836468183976), + (3e-18, 5, 3195.77171838060906447)]) +def test_invweibull_isf(p, c, expected): + computed = stats.invweibull.isf(p, c) + assert_allclose(computed, expected, rtol=1e-15) + + +@pytest.mark.parametrize( + 'df1,df2,x', + [(2, 2, [-0.5, 0.2, 1.0, 2.3]), + (4, 11, [-0.5, 0.2, 1.0, 2.3]), + (7, 17, [1, 2, 3, 4, 5])] +) +def test_ncf_edge_case(df1, df2, x): + # Test for edge case described in gh-11660. + # Non-central Fisher distribution when nc = 0 + # should be the same as Fisher distribution. + nc = 0 + expected_cdf = stats.f.cdf(x, df1, df2) + calculated_cdf = stats.ncf.cdf(x, df1, df2, nc) + assert_allclose(expected_cdf, calculated_cdf, rtol=1e-14) + + # when ncf_gen._skip_pdf will be used instead of generic pdf, + # this additional test will be useful. + expected_pdf = stats.f.pdf(x, df1, df2) + calculated_pdf = stats.ncf.pdf(x, df1, df2, nc) + assert_allclose(expected_pdf, calculated_pdf, rtol=1e-6) + + +def test_ncf_variance(): + # Regression test for gh-10658 (incorrect variance formula for ncf). + # The correct value of ncf.var(2, 6, 4), 42.75, can be verified with, for + # example, Wolfram Alpha with the expression + # Variance[NoncentralFRatioDistribution[2, 6, 4]] + # or with the implementation of the noncentral F distribution in the C++ + # library Boost. + v = stats.ncf.var(2, 6, 4) + assert_allclose(v, 42.75, rtol=1e-14) + + +def test_ncf_cdf_spotcheck(): + # Regression test for gh-15582 testing against values from R/MATLAB + # Generate check_val from R or MATLAB as follows: + # R: pf(20, df1 = 6, df2 = 33, ncp = 30.4) = 0.998921 + # MATLAB: ncfcdf(20, 6, 33, 30.4) = 0.998921 + scipy_val = stats.ncf.cdf(20, 6, 33, 30.4) + check_val = 0.998921 + assert_allclose(check_val, np.round(scipy_val, decimals=6)) + + +def test_ncf_ppf_issue_17026(): + # Regression test for gh-17026 + x = np.linspace(0, 1, 600) + x[0] = 1e-16 + par = (0.1, 2, 5, 0, 1) + q = stats.ncf.ppf(x, *par) + q0 = [stats.ncf.ppf(xi, *par) for xi in x] + assert_allclose(q, q0) + + +class TestHistogram: + def setup_method(self): + self.rng = np.random.default_rng(6561174822) + + # We have 8 bins + # [1,2), [2,3), [3,4), [4,5), [5,6), [6,7), [7,8), [8,9) + # But actually np.histogram will put the last 9 also in the [8,9) bin! + # Therefore there is a slight difference below for the last bin, from + # what you might have expected. + histogram = np.histogram([1, 2, 2, 3, 3, 3, 4, 4, 4, 4, 5, 5, 5, 5, 5, + 6, 6, 6, 6, 7, 7, 7, 8, 8, 9], bins=8) + self.template = stats.rv_histogram(histogram) + + data = stats.norm.rvs(loc=1.0, scale=2.5, size=10000, random_state=self.rng) + norm_histogram = np.histogram(data, bins=50) + self.norm_template = stats.rv_histogram(norm_histogram) + + def test_pdf(self): + values = np.array([0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0, 4.5, + 5.0, 5.5, 6.0, 6.5, 7.0, 7.5, 8.0, 8.5, 9.0, 9.5]) + pdf_values = np.asarray([0.0/25.0, 0.0/25.0, 1.0/25.0, 1.0/25.0, + 2.0/25.0, 2.0/25.0, 3.0/25.0, 3.0/25.0, + 4.0/25.0, 4.0/25.0, 5.0/25.0, 5.0/25.0, + 4.0/25.0, 4.0/25.0, 3.0/25.0, 3.0/25.0, + 3.0/25.0, 3.0/25.0, 0.0/25.0, 0.0/25.0]) + assert_allclose(self.template.pdf(values), pdf_values) + + # Test explicitly the corner cases: + # As stated above the pdf in the bin [8,9) is greater than + # one would naively expect because np.histogram putted the 9 + # into the [8,9) bin. + assert_almost_equal(self.template.pdf(8.0), 3.0/25.0) + assert_almost_equal(self.template.pdf(8.5), 3.0/25.0) + # 9 is outside our defined bins [8,9) hence the pdf is already 0 + # for a continuous distribution this is fine, because a single value + # does not have a finite probability! + assert_almost_equal(self.template.pdf(9.0), 0.0/25.0) + assert_almost_equal(self.template.pdf(10.0), 0.0/25.0) + + x = np.linspace(-2, 2, 10) + assert_allclose(self.norm_template.pdf(x), + stats.norm.pdf(x, loc=1.0, scale=2.5), rtol=0.1) + + def test_cdf_ppf(self): + values = np.array([0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0, 4.5, + 5.0, 5.5, 6.0, 6.5, 7.0, 7.5, 8.0, 8.5, 9.0, 9.5]) + cdf_values = np.asarray([0.0/25.0, 0.0/25.0, 0.0/25.0, 0.5/25.0, + 1.0/25.0, 2.0/25.0, 3.0/25.0, 4.5/25.0, + 6.0/25.0, 8.0/25.0, 10.0/25.0, 12.5/25.0, + 15.0/25.0, 17.0/25.0, 19.0/25.0, 20.5/25.0, + 22.0/25.0, 23.5/25.0, 25.0/25.0, 25.0/25.0]) + assert_allclose(self.template.cdf(values), cdf_values) + # First three and last two values in cdf_value are not unique + assert_allclose(self.template.ppf(cdf_values[2:-1]), values[2:-1]) + + # Test of cdf and ppf are inverse functions + x = np.linspace(1.0, 9.0, 100) + assert_allclose(self.template.ppf(self.template.cdf(x)), x) + x = np.linspace(0.0, 1.0, 100) + assert_allclose(self.template.cdf(self.template.ppf(x)), x) + + x = np.linspace(-2, 2, 10) + assert_allclose(self.norm_template.cdf(x), + stats.norm.cdf(x, loc=1.0, scale=2.5), rtol=0.1) + + def test_rvs(self): + N = 10000 + sample = self.template.rvs(size=N, random_state=self.rng) + assert_equal(np.sum(sample < 1.0), 0.0) + assert_allclose(np.sum(sample <= 2.0), 1.0/25.0 * N, rtol=0.2) + assert_allclose(np.sum(sample <= 2.5), 2.0/25.0 * N, rtol=0.2) + assert_allclose(np.sum(sample <= 3.0), 3.0/25.0 * N, rtol=0.1) + assert_allclose(np.sum(sample <= 3.5), 4.5/25.0 * N, rtol=0.1) + assert_allclose(np.sum(sample <= 4.0), 6.0/25.0 * N, rtol=0.1) + assert_allclose(np.sum(sample <= 4.5), 8.0/25.0 * N, rtol=0.1) + assert_allclose(np.sum(sample <= 5.0), 10.0/25.0 * N, rtol=0.05) + assert_allclose(np.sum(sample <= 5.5), 12.5/25.0 * N, rtol=0.05) + assert_allclose(np.sum(sample <= 6.0), 15.0/25.0 * N, rtol=0.05) + assert_allclose(np.sum(sample <= 6.5), 17.0/25.0 * N, rtol=0.05) + assert_allclose(np.sum(sample <= 7.0), 19.0/25.0 * N, rtol=0.05) + assert_allclose(np.sum(sample <= 7.5), 20.5/25.0 * N, rtol=0.05) + assert_allclose(np.sum(sample <= 8.0), 22.0/25.0 * N, rtol=0.05) + assert_allclose(np.sum(sample <= 8.5), 23.5/25.0 * N, rtol=0.05) + assert_allclose(np.sum(sample <= 9.0), 25.0/25.0 * N, rtol=0.05) + assert_allclose(np.sum(sample <= 9.0), 25.0/25.0 * N, rtol=0.05) + assert_equal(np.sum(sample > 9.0), 0.0) + + def test_munp(self): + for n in range(4): + assert_allclose(self.norm_template._munp(n), + stats.norm(1.0, 2.5).moment(n), rtol=0.05) + + def test_entropy(self): + assert_allclose(self.norm_template.entropy(), + stats.norm.entropy(loc=1.0, scale=2.5), rtol=0.05) + + +def test_histogram_non_uniform(): + # Tests rv_histogram works even for non-uniform bin widths + counts, bins = ([1, 1], [0, 1, 1001]) + + dist = stats.rv_histogram((counts, bins), density=False) + np.testing.assert_allclose(dist.pdf([0.5, 200]), [0.5, 0.0005]) + assert dist.median() == 1 + + dist = stats.rv_histogram((counts, bins), density=True) + np.testing.assert_allclose(dist.pdf([0.5, 200]), 1/1001) + assert dist.median() == 1001/2 + + # Omitting density produces a warning for non-uniform bins... + message = "Bin widths are not constant. Assuming..." + with pytest.warns(RuntimeWarning, match=message): + dist = stats.rv_histogram((counts, bins)) + assert dist.median() == 1001/2 # default is like `density=True` + + # ... but not for uniform bins + dist = stats.rv_histogram((counts, [0, 1, 2])) + assert dist.median() == 1 + + +class TestLogUniform: + def test_alias(self): + # This test makes sure that "reciprocal" and "loguniform" are + # aliases of the same distribution and that both are log-uniform + rng = np.random.default_rng(98643218961) + rv = stats.loguniform(10 ** -3, 10 ** 0) + rvs = rv.rvs(size=10000, random_state=rng) + + rng = np.random.default_rng(98643218961) + rv2 = stats.reciprocal(10 ** -3, 10 ** 0) + rvs2 = rv2.rvs(size=10000, random_state=rng) + + assert_allclose(rvs2, rvs) + + vals, _ = np.histogram(np.log10(rvs), bins=10) + assert 900 <= vals.min() <= vals.max() <= 1100 + assert np.abs(np.median(vals) - 1000) <= 10 + + @pytest.mark.parametrize("method", ['mle', 'mm']) + def test_fit_override(self, method): + # loguniform is overparameterized, so check that fit override enforces + # scale=1 unless fscale is provided by the user + rng = np.random.default_rng(98643218961) + rvs = stats.loguniform.rvs(0.1, 1, size=1000, random_state=rng) + + a, b, loc, scale = stats.loguniform.fit(rvs, method=method) + assert scale == 1 + + a, b, loc, scale = stats.loguniform.fit(rvs, fscale=2, method=method) + assert scale == 2 + + def test_overflow(self): + # original formulation had overflow issues; check that this is resolved + # Extensive accuracy tests elsewhere, no need to test all methods + rng = np.random.default_rng(7136519550773909093) + a, b = 1e-200, 1e200 + dist = stats.loguniform(a, b) + + # test roundtrip error + cdf = rng.uniform(0, 1, size=1000) + assert_allclose(dist.cdf(dist.ppf(cdf)), cdf) + rvs = dist.rvs(size=1000, random_state=rng) + assert_allclose(dist.ppf(dist.cdf(rvs)), rvs) + + # test a property of the pdf (and that there is no overflow) + x = 10.**np.arange(-200, 200) + pdf = dist.pdf(x) # no overflow + assert_allclose(pdf[:-1]/pdf[1:], 10) + + # check munp against wikipedia reference + mean = (b - a)/(np.log(b) - np.log(a)) + assert_allclose(dist.mean(), mean) + + +class TestArgus: + def test_argus_rvs_large_chi(self): + # test that the algorithm can handle large values of chi + x = stats.argus.rvs(50, size=500, random_state=325) + assert_almost_equal(stats.argus(50).mean(), x.mean(), decimal=4) + + @pytest.mark.parametrize('chi, random_state', [ + [0.1, 325], # chi <= 0.5: rejection method case 1 + [1.3, 155], # 0.5 < chi <= 1.8: rejection method case 2 + [3.5, 135] # chi > 1.8: transform conditional Gamma distribution + ]) + def test_rvs(self, chi, random_state): + x = stats.argus.rvs(chi, size=500, random_state=random_state) + _, p = stats.kstest(x, "argus", (chi, )) + assert_(p > 0.05) + + @pytest.mark.parametrize('chi', [1e-9, 1e-6]) + def test_rvs_small_chi(self, chi): + # test for gh-11699 => rejection method case 1 can even handle chi=0 + # the CDF of the distribution for chi=0 is 1 - (1 - x**2)**(3/2) + # test rvs against distribution of limit chi=0 + r = stats.argus.rvs(chi, size=500, random_state=890981) + _, p = stats.kstest(r, lambda x: 1 - (1 - x**2)**(3/2)) + assert_(p > 0.05) + + # Expected values were computed with mpmath. + @pytest.mark.parametrize('chi, expected_mean', + [(1, 0.6187026683551835), + (10, 0.984805536783744), + (40, 0.9990617659702923), + (60, 0.9995831885165300), + (99, 0.9998469348663028)]) + def test_mean(self, chi, expected_mean): + m = stats.argus.mean(chi, scale=1) + assert_allclose(m, expected_mean, rtol=1e-13) + + # Expected values were computed with mpmath. + @pytest.mark.parametrize('chi, expected_var, rtol', + [(1, 0.05215651254197807, 1e-13), + (10, 0.00015805472008165595, 1e-11), + (40, 5.877763210262901e-07, 1e-8), + (60, 1.1590179389611416e-07, 1e-8), + (99, 1.5623277006064666e-08, 1e-8)]) + def test_var(self, chi, expected_var, rtol): + v = stats.argus.var(chi, scale=1) + assert_allclose(v, expected_var, rtol=rtol) + + # Expected values were computed with mpmath (code: see gh-13370). + @pytest.mark.parametrize('chi, expected, rtol', + [(0.9, 0.07646314974436118, 1e-14), + (0.5, 0.015429797891863365, 1e-14), + (0.1, 0.0001325825293278049, 1e-14), + (0.01, 1.3297677078224565e-07, 1e-15), + (1e-3, 1.3298072023958999e-10, 1e-14), + (1e-4, 1.3298075973486862e-13, 1e-14), + (1e-6, 1.32980760133771e-19, 1e-14), + (1e-9, 1.329807601338109e-28, 1e-15)]) + def test_argus_phi_small_chi(self, chi, expected, rtol): + assert_allclose(_argus_phi(chi), expected, rtol=rtol) + + # Expected values were computed with mpmath (code: see gh-13370). + @pytest.mark.parametrize( + 'chi, expected', + [(0.5, (0.28414073302940573, 1.2742227939992954, 1.2381254688255896)), + (0.2, (0.296172952995264, 1.2951290588110516, 1.1865767100877576)), + (0.1, (0.29791447523536274, 1.29806307956989, 1.1793168289857412)), + (0.01, (0.2984904104866452, 1.2990283628160553, 1.1769268414080531)), + (1e-3, (0.298496172925224, 1.2990380082487925, 1.176902956021053)), + (1e-4, (0.29849623054991836, 1.2990381047023793, 1.1769027171686324)), + (1e-6, (0.2984962311319278, 1.2990381056765605, 1.1769027147562232)), + (1e-9, (0.298496231131986, 1.299038105676658, 1.1769027147559818))]) + def test_pdf_small_chi(self, chi, expected): + x = np.array([0.1, 0.5, 0.9]) + assert_allclose(stats.argus.pdf(x, chi), expected, rtol=1e-13) + + # Expected values were computed with mpmath (code: see gh-13370). + @pytest.mark.parametrize( + 'chi, expected', + [(0.5, (0.9857660526895221, 0.6616565930168475, 0.08796070398429937)), + (0.2, (0.9851555052359501, 0.6514666238985464, 0.08362690023746594)), + (0.1, (0.9850670974995661, 0.6500061310508574, 0.08302050640683846)), + (0.01, (0.9850378582451867, 0.6495239242251358, 0.08282109244852445)), + (1e-3, (0.9850375656906663, 0.6495191015522573, 0.08281910005231098)), + (1e-4, (0.9850375627651049, 0.6495190533254682, 0.08281908012852317)), + (1e-6, (0.9850375627355568, 0.6495190528383777, 0.08281907992729293)), + (1e-9, (0.9850375627355538, 0.649519052838329, 0.0828190799272728))]) + def test_sf_small_chi(self, chi, expected): + x = np.array([0.1, 0.5, 0.9]) + assert_allclose(stats.argus.sf(x, chi), expected, rtol=1e-14) + + # Expected values were computed with mpmath. + @pytest.mark.parametrize( + 'x, chi, expected', + [(0.9999999, 0.25, 9.113252974162428e-11), + (0.9999999, 3.0, 6.616650419714568e-10), + (0.999999999, 2.5, 4.130195911418939e-13), + (0.999999999, 10.0, 2.3788319094393724e-11)]) + def test_sf_near_1(self, x, chi, expected): + sf = stats.argus.sf(x, chi) + assert_allclose(sf, expected, rtol=5e-15) + + # Expected values were computed with mpmath (code: see gh-13370). + @pytest.mark.parametrize( + 'chi, expected', + [(0.5, (0.0142339473104779, 0.3383434069831524, 0.9120392960157007)), + (0.2, (0.014844494764049919, 0.34853337610145363, 0.916373099762534)), + (0.1, (0.014932902500433911, 0.34999386894914264, 0.9169794935931616)), + (0.01, (0.014962141754813293, 0.35047607577486417, 0.9171789075514756)), + (1e-3, (0.01496243430933372, 0.35048089844774266, 0.917180899947689)), + (1e-4, (0.014962437234895118, 0.3504809466745317, 0.9171809198714769)), + (1e-6, (0.01496243726444329, 0.3504809471616223, 0.9171809200727071)), + (1e-9, (0.014962437264446245, 0.350480947161671, 0.9171809200727272))]) + def test_cdf_small_chi(self, chi, expected): + x = np.array([0.1, 0.5, 0.9]) + assert_allclose(stats.argus.cdf(x, chi), expected, rtol=1e-12) + + # Expected values were computed with mpmath (code: see gh-13370). + @pytest.mark.parametrize( + 'chi, expected, rtol', + [(0.5, (0.5964284712757741, 0.052890651988588604), 1e-12), + (0.101, (0.5893490968089076, 0.053017469847275685), 1e-11), + (0.1, (0.5893431757009437, 0.05301755449499372), 1e-13), + (0.01, (0.5890515677940915, 0.05302167905837031), 1e-13), + (1e-3, (0.5890486520005177, 0.053021719862088104), 1e-13), + (1e-4, (0.5890486228426105, 0.0530217202700811), 1e-13), + (1e-6, (0.5890486225481156, 0.05302172027420182), 1e-13), + (1e-9, (0.5890486225480862, 0.05302172027420224), 1e-13)]) + def test_stats_small_chi(self, chi, expected, rtol): + val = stats.argus.stats(chi, moments='mv') + assert_allclose(val, expected, rtol=rtol) + + +class TestNakagami: + def setup_method(self): + self.rng = np.random.default_rng(9626839526) + + def test_logpdf(self): + # Test nakagami logpdf for an input where the PDF is smaller + # than can be represented with 64 bit floating point. + # The expected value of logpdf was computed with mpmath: + # + # def logpdf(x, nu): + # x = mpmath.mpf(x) + # nu = mpmath.mpf(nu) + # return (mpmath.log(2) + nu*mpmath.log(nu) - + # mpmath.loggamma(nu) + (2*nu - 1)*mpmath.log(x) - + # nu*x**2) + # + nu = 2.5 + x = 25 + logp = stats.nakagami.logpdf(x, nu) + assert_allclose(logp, -1546.9253055607549) + + def test_sf_isf(self): + # Test nakagami sf and isf when the survival function + # value is very small. + # The expected value of the survival function was computed + # with mpmath: + # + # def sf(x, nu): + # x = mpmath.mpf(x) + # nu = mpmath.mpf(nu) + # return mpmath.gammainc(nu, nu*x*x, regularized=True) + # + nu = 2.5 + x0 = 5.0 + sf = stats.nakagami.sf(x0, nu) + assert_allclose(sf, 2.736273158588307e-25, rtol=1e-13) + # Check round trip back to x0. + x1 = stats.nakagami.isf(sf, nu) + assert_allclose(x1, x0, rtol=1e-13) + + def test_logcdf(self): + x = 8 + nu = 0.5 + # Reference value computed with mpmath. + ref = -1.2441921148543576e-15 + logcdf = stats.nakagami.logcdf(x, nu) + assert_allclose(logcdf, ref, rtol=5e-15) + + def test_logsf(self): + x = 0.05 + nu = 12 + # Reference value computed with mpmath. + ref = -1.0791764722337046e-27 + logsf = stats.nakagami.logsf(x, nu) + assert_allclose(logsf, ref, rtol=5e-15) + + @pytest.mark.parametrize("m, ref", + [(5, -0.097341814372152), + (0.5, 0.7257913526447274), + (10, -0.43426184310934907)]) + def test_entropy(self, m, ref): + # from sympy import * + # from mpmath import mp + # import numpy as np + # v, x = symbols('v, x', real=True, positive=True) + # pdf = 2 * v ** v / gamma(v) * x ** (2 * v - 1) * exp(-v * x ** 2) + # h = simplify(simplify(integrate(-pdf * log(pdf), (x, 0, oo)))) + # entropy = lambdify(v, h, 'mpmath') + # mp.dps = 200 + # nu = 5 + # ref = np.float64(entropy(mp.mpf(nu))) + # print(ref) + assert_allclose(stats.nakagami.entropy(m), ref, rtol=1.1e-14) + + @pytest.mark.parametrize("m, ref", + [(1e-100, -5.0e+99), (1e-10, -4999999965.442979), + (9.999e6, -7.333206478668433), (1.001e7, -7.3337562313259825), + (1e10, -10.787134112333835), (1e100, -114.40346329705756)]) + def test_extreme_nu(self, m, ref): + assert_allclose(stats.nakagami.entropy(m), ref) + + def test_entropy_overflow(self): + assert np.isfinite(stats.nakagami._entropy(1e100)) + assert np.isfinite(stats.nakagami._entropy(1e-100)) + + @pytest.mark.parametrize("nu, ref", + [(1e10, 0.9999999999875), + (1e3, 0.9998750078173821), + (1e-10, 1.772453850659802e-05)]) + def test_mean(self, nu, ref): + # reference values were computed with mpmath + # from mpmath import mp + # mp.dps = 500 + # nu = mp.mpf(1e10) + # float(mp.rf(nu, mp.mpf(0.5))/mp.sqrt(nu)) + assert_allclose(stats.nakagami.mean(nu), ref, rtol=1e-12) + + @pytest.mark.xfail(reason="Fit of nakagami not reliable, see gh-10908.") + @pytest.mark.parametrize('nu', [1.6, 2.5, 3.9]) + @pytest.mark.parametrize('loc', [25.0, 10, 35]) + @pytest.mark.parametrize('scale', [13, 5, 20]) + def test_fit(self, nu, loc, scale): + # Regression test for gh-13396 (21/27 cases failed previously) + # The first tuple of the parameters' values is discussed in gh-10908 + N = 100 + samples = stats.nakagami.rvs(size=N, nu=nu, loc=loc, + scale=scale, random_state=self.rng) + nu_est, loc_est, scale_est = stats.nakagami.fit(samples) + assert_allclose(nu_est, nu, rtol=0.2) + assert_allclose(loc_est, loc, rtol=0.2) + assert_allclose(scale_est, scale, rtol=0.2) + + def dlogl_dnu(nu, loc, scale): + return ((-2*nu + 1) * np.sum(1/(samples - loc)) + + 2*nu/scale**2 * np.sum(samples - loc)) + + def dlogl_dloc(nu, loc, scale): + return (N * (1 + np.log(nu) - polygamma(0, nu)) + + 2 * np.sum(np.log((samples - loc) / scale)) + - np.sum(((samples - loc) / scale)**2)) + + def dlogl_dscale(nu, loc, scale): + return (- 2 * N * nu / scale + + 2 * nu / scale ** 3 * np.sum((samples - loc) ** 2)) + + assert_allclose(dlogl_dnu(nu_est, loc_est, scale_est), 0, atol=1e-3) + assert_allclose(dlogl_dloc(nu_est, loc_est, scale_est), 0, atol=1e-3) + assert_allclose(dlogl_dscale(nu_est, loc_est, scale_est), 0, atol=1e-3) + + @pytest.mark.parametrize('loc', [25.0, 10, 35]) + @pytest.mark.parametrize('scale', [13, 5, 20]) + def test_fit_nu(self, loc, scale): + # For nu = 0.5, we have analytical values for + # the MLE of the loc and the scale + nu = 0.5 + n = 100 + samples = stats.nakagami.rvs(size=n, nu=nu, loc=loc, + scale=scale, random_state=self.rng) + nu_est, loc_est, scale_est = stats.nakagami.fit(samples, f0=nu) + + # Analytical values + loc_theo = np.min(samples) + scale_theo = np.sqrt(np.mean((samples - loc_est) ** 2)) + + assert_allclose(nu_est, nu, rtol=1e-7) + assert_allclose(loc_est, loc_theo, rtol=1e-7) + assert_allclose(scale_est, scale_theo, rtol=1e-7) + + +class TestWrapCauchy: + def setup_method(self): + self.rng = np.random.default_rng(2439107123) + + def test_cdf_shape_broadcasting(self): + # Regression test for gh-13791. + # Check that wrapcauchy.cdf broadcasts the shape parameter + # correctly. + c = np.array([[0.03, 0.25], [0.5, 0.75]]) + x = np.array([[1.0], [4.0]]) + p = stats.wrapcauchy.cdf(x, c) + assert p.shape == (2, 2) + scalar_values = [stats.wrapcauchy.cdf(x1, c1) + for (x1, c1) in np.nditer((x, c))] + assert_allclose(p.ravel(), scalar_values, rtol=1e-13) + + def test_cdf_center(self): + p = stats.wrapcauchy.cdf(np.pi, 0.03) + assert_allclose(p, 0.5, rtol=1e-14) + + def test_cdf(self): + x1 = 1.0 # less than pi + x2 = 4.0 # greater than pi + c = 0.75 + p = stats.wrapcauchy.cdf([x1, x2], c) + cr = (1 + c)/(1 - c) + assert_allclose(p[0], np.arctan(cr*np.tan(x1/2))/np.pi) + assert_allclose(p[1], 1 - np.arctan(cr*np.tan(np.pi - x2/2))/np.pi) + + @pytest.mark.parametrize('c', [1e-10, 1e-1]) + @pytest.mark.parametrize('loc', [-100, -2*np.pi, -np.pi, 0, np.pi, 2*np.pi, 100]) + @pytest.mark.parametrize('scale', [1e-10, 1, 1e10]) + def test_rvs_lie_on_circle(self, c, loc, scale): + # Check that the random variates lie in range [0, 2*pi] + x = stats.wrapcauchy.rvs(c=c, loc=loc, scale=scale, + size=1000, random_state=self.rng) + assert np.all(x >= 0) + assert np.all(x <= 2 * np.pi) + + +def test_rvs_no_size_error(): + # _rvs methods must have parameter `size`; see gh-11394 + class rvs_no_size_gen(stats.rv_continuous): + def _rvs(self): + return 1 + + rvs_no_size = rvs_no_size_gen(name='rvs_no_size') + rng = np.random.default_rng(1334886239) + with assert_raises(TypeError, match=r"_rvs\(\) got (an|\d) unexpected"): + rvs_no_size.rvs(random_state=rng) + + +@pytest.mark.parametrize('distname, args', invdistdiscrete + invdistcont) +def test_support_gh13294_regression(distname, args): + if distname in skip_test_support_gh13294_regression: + pytest.skip(f"skipping test for the support method for " + f"distribution {distname}.") + dist = getattr(stats, distname) + # test support method with invalid arguments + if isinstance(dist, stats.rv_continuous): + # test with valid scale + if len(args) != 0: + a0, b0 = dist.support(*args) + assert_equal(a0, np.nan) + assert_equal(b0, np.nan) + # test with invalid scale + # For some distributions, that take no parameters, + # the case of only invalid scale occurs and hence, + # it is implicitly tested in this test case. + loc1, scale1 = 0, -1 + a1, b1 = dist.support(*args, loc1, scale1) + assert_equal(a1, np.nan) + assert_equal(b1, np.nan) + else: + a, b = dist.support(*args) + assert_equal(a, np.nan) + assert_equal(b, np.nan) + + +def test_support_broadcasting_gh13294_regression(): + a0, b0 = stats.norm.support([0, 0, 0, 1], [1, 1, 1, -1]) + ex_a0 = np.array([-np.inf, -np.inf, -np.inf, np.nan]) + ex_b0 = np.array([np.inf, np.inf, np.inf, np.nan]) + assert_equal(a0, ex_a0) + assert_equal(b0, ex_b0) + assert a0.shape == ex_a0.shape + assert b0.shape == ex_b0.shape + + a1, b1 = stats.norm.support([], []) + ex_a1, ex_b1 = np.array([]), np.array([]) + assert_equal(a1, ex_a1) + assert_equal(b1, ex_b1) + assert a1.shape == ex_a1.shape + assert b1.shape == ex_b1.shape + + a2, b2 = stats.norm.support([0, 0, 0, 1], [-1]) + ex_a2 = np.array(4*[np.nan]) + ex_b2 = np.array(4*[np.nan]) + assert_equal(a2, ex_a2) + assert_equal(b2, ex_b2) + assert a2.shape == ex_a2.shape + assert b2.shape == ex_b2.shape + + +def test_stats_broadcasting_gh14953_regression(): + # test case in gh14953 + loc = [0., 0.] + scale = [[1.], [2.], [3.]] + assert_equal(stats.norm.var(loc, scale), [[1., 1.], [4., 4.], [9., 9.]]) + # test some edge cases + loc = np.empty((0, )) + scale = np.empty((1, 0)) + assert stats.norm.var(loc, scale).shape == (1, 0) + + +# Check a few values of the cosine distribution's cdf, sf, ppf and +# isf methods. Expected values were computed with mpmath. + +@pytest.mark.parametrize('x, expected', + [(-3.14159, 4.956444476505336e-19), + (3.14, 0.9999999998928399)]) +def test_cosine_cdf_sf(x, expected): + assert_allclose(stats.cosine.cdf(x), expected) + assert_allclose(stats.cosine.sf(-x), expected) + + +@pytest.mark.parametrize('p, expected', + [(1e-6, -3.1080612413765905), + (1e-17, -3.141585429601399), + (0.975, 2.1447547020964923)]) +def test_cosine_ppf_isf(p, expected): + assert_allclose(stats.cosine.ppf(p), expected) + assert_allclose(stats.cosine.isf(p), -expected) + + +def test_cosine_logpdf_endpoints(): + logp = stats.cosine.logpdf([-np.pi, np.pi]) + # reference value calculated using mpmath assuming `np.cos(-1)` is four + # floating point numbers too high. See gh-18382. + assert_array_less(logp, -37.18838327496655) + + +def test_distr_params_lists(): + # distribution objects are extra distributions added in + # test_discrete_basic. All other distributions are strings (names) + # and so we only choose those to compare whether both lists match. + discrete_distnames = {name for name, _ in distdiscrete + if isinstance(name, str)} + invdiscrete_distnames = {name for name, _ in invdistdiscrete} + assert discrete_distnames == invdiscrete_distnames + + cont_distnames = {name for name, _ in distcont} + invcont_distnames = {name for name, _ in invdistcont} + assert cont_distnames == invcont_distnames + + +def test_moment_order_4(): + # gh-13655 reported that if a distribution has a `_stats` method that + # accepts the `moments` parameter, then if the distribution's `moment` + # method is called with `order=4`, the faster/more accurate`_stats` gets + # called, but the results aren't used, and the generic `_munp` method is + # called to calculate the moment anyway. This tests that the issue has + # been fixed. + # stats.skewnorm._stats accepts the `moments` keyword + stats.skewnorm._stats(a=0, moments='k') # no failure = has `moments` + # When `moment` is called, `_stats` is used, so the moment is very accurate + # (exactly equal to Pearson's kurtosis of the normal distribution, 3) + assert stats.skewnorm.moment(order=4, a=0) == 3.0 + # At the time of gh-13655, skewnorm._munp() used the generic method + # to compute its result, which was inefficient and not very accurate. + # At that time, the following assertion would fail. skewnorm._munp() + # has since been made more accurate and efficient, so now this test + # is expected to pass. + assert stats.skewnorm._munp(4, 0) == 3.0 + + +class TestRelativisticBW: + @pytest.fixture + def ROOT_pdf_sample_data(self): + """Sample data points for pdf computed with CERN's ROOT + + See - https://root.cern/ + + Uses ROOT.TMath.BreitWignerRelativistic, available in ROOT + versions 6.27+ + + pdf calculated for Z0 Boson, W Boson, and Higgs Boson for + x in `np.linspace(0, 200, 401)`. + """ + data = np.load( + Path(__file__).parent / + 'data/rel_breitwigner_pdf_sample_data_ROOT.npy' + ) + data = np.rec.fromarrays(data.T, names='x,pdf,rho,gamma') + return data + + @pytest.mark.parametrize( + "rho,gamma,rtol", [ + (36.545206797050334, 2.4952, 5e-14), # Z0 Boson + (38.55107913669065, 2.085, 1e-14), # W Boson + (96292.3076923077, 0.0013, 5e-13), # Higgs Boson + ] + ) + def test_pdf_against_ROOT(self, ROOT_pdf_sample_data, rho, gamma, rtol): + data = ROOT_pdf_sample_data[ + (ROOT_pdf_sample_data['rho'] == rho) + & (ROOT_pdf_sample_data['gamma'] == gamma) + ] + x, pdf = data['x'], data['pdf'] + assert_allclose( + pdf, stats.rel_breitwigner.pdf(x, rho, scale=gamma), rtol=rtol + ) + + @pytest.mark.parametrize("rho, Gamma, rtol", [ + (36.545206797050334, 2.4952, 5e-13), # Z0 Boson + (38.55107913669065, 2.085, 5e-13), # W Boson + (96292.3076923077, 0.0013, 5e-10), # Higgs Boson + ] + ) + def test_pdf_against_simple_implementation(self, rho, Gamma, rtol): + # reference implementation straight from formulas on Wikipedia [1] + def pdf(E, M, Gamma): + gamma = np.sqrt(M**2 * (M**2 + Gamma**2)) + k = (2 * np.sqrt(2) * M * Gamma * gamma + / (np.pi * np.sqrt(M**2 + gamma))) + return k / ((E**2 - M**2)**2 + M**2*Gamma**2) + # get reasonable values at which to evaluate the CDF + p = np.linspace(0.05, 0.95, 10) + x = stats.rel_breitwigner.ppf(p, rho, scale=Gamma) + res = stats.rel_breitwigner.pdf(x, rho, scale=Gamma) + ref = pdf(x, rho*Gamma, Gamma) + assert_allclose(res, ref, rtol=rtol) + + @pytest.mark.xslow + @pytest.mark.parametrize( + "rho,gamma", [ + pytest.param( + 36.545206797050334, 2.4952, marks=pytest.mark.slow + ), # Z0 Boson + pytest.param( + 38.55107913669065, 2.085, marks=pytest.mark.xslow + ), # W Boson + pytest.param( + 96292.3076923077, 0.0013, marks=pytest.mark.xslow + ), # Higgs Boson + ] + ) + def test_fit_floc(self, rho, gamma): + """Tests fit for cases where floc is set. + + `rel_breitwigner` has special handling for these cases. + """ + seed = 6936804688480013683 + rng = np.random.default_rng(seed) + data = stats.rel_breitwigner.rvs( + rho, scale=gamma, size=1000, random_state=rng + ) + fit = stats.rel_breitwigner.fit(data, floc=0) + assert_allclose((fit[0], fit[2]), (rho, gamma), rtol=2e-1) + assert fit[1] == 0 + # Check again with fscale set. + fit = stats.rel_breitwigner.fit(data, floc=0, fscale=gamma) + assert_allclose(fit[0], rho, rtol=1e-2) + assert (fit[1], fit[2]) == (0, gamma) + + +class TestJohnsonSU: + @pytest.mark.parametrize("case", [ # a, b, loc, scale, m1, m2, g1, g2 + (-0.01, 1.1, 0.02, 0.0001, 0.02000137427557091, + 2.1112742956578063e-08, 0.05989781342460999, 20.36324408592951-3), + (2.554395574161155, 2.2482281679651965, 0, 1, -1.54215386737391, + 0.7629882028469993, -1.256656139406788, 6.303058419339775-3)]) + def test_moment_gh18071(self, case): + # gh-18071 reported an IntegrationWarning emitted by johnsonsu.stats + # Check that the warning is no longer emitted and that the values + # are accurate compared against results from Mathematica. + # Reference values from Mathematica, e.g. + # Mean[JohnsonDistribution["SU",-0.01, 1.1, 0.02, 0.0001]] + res = stats.johnsonsu.stats(*case[:4], moments='mvsk') + assert_allclose(res, case[4:], rtol=1e-14) + + +class TestTruncPareto: + def test_pdf(self): + # PDF is that of the truncated pareto distribution + b, c = 1.8, 5.3 + x = np.linspace(1.8, 5.3) + res = stats.truncpareto(b, c).pdf(x) + ref = stats.pareto(b).pdf(x) / stats.pareto(b).cdf(c) + assert_allclose(res, ref) + + def test_pdf_negative(self): + # truncpareto is equivalent to more general powerlaw from gh-23648 + # exponent of truncpareto is negative in this case + a, xmin, xmax = 4, 3, 5 + x = np.linspace(xmin, xmax) + + # compute reference using PDF from gh-23648 + C = a / (xmax ** a - xmin ** a) + ref = C * x ** (a - 1) + + # compute using `truncpareto` with negative exponent + b = -a + c = xmax / xmin + scale = xmin + loc = 0 + X = stats.truncpareto(b, c, loc, scale) + + assert_allclose(X.pdf(x), ref) + assert_allclose(X.logpdf(x), np.log(X.pdf(x))) + # indexing avoids RuntimeWarning with `np.log(0)` + assert_allclose(X.logcdf(x[1:]), np.log(X.cdf(x[1:]))) + assert_allclose(X.logsf(x[:-1]), np.log(X.sf(x[:-1]))) + + @pytest.mark.parametrize('fix_loc', [True, False]) + @pytest.mark.parametrize('fix_scale', [True, False]) + @pytest.mark.parametrize('fix_b', [True, False]) + @pytest.mark.parametrize('fix_c', [True, False]) + def test_fit(self, fix_loc, fix_scale, fix_b, fix_c): + + rng = np.random.default_rng(6747363148258237171) + b, c, loc, scale = 1.8, 5.3, 1, 2.5 + dist = stats.truncpareto(b, c, loc=loc, scale=scale) + data = dist.rvs(size=500, random_state=rng) + + kwds = {} + if fix_loc: + kwds['floc'] = loc + if fix_scale: + kwds['fscale'] = scale + if fix_b: + kwds['f0'] = b + if fix_c: + kwds['f1'] = c + + if fix_loc and fix_scale and fix_b and fix_c: + message = "All parameters fixed. There is nothing to optimize." + with pytest.raises(RuntimeError, match=message): + stats.truncpareto.fit(data, **kwds) + else: + _assert_less_or_close_loglike(stats.truncpareto, data, **kwds) + + +class TestKappa3: + def test_sf(self): + # During development of gh-18822, we found that the override of + # kappa3.sf could experience overflow where the version in main did + # not. Check that this does not happen in final implementation. + sf0 = 1 - stats.kappa3.cdf(0.5, 1e5) + sf1 = stats.kappa3.sf(0.5, 1e5) + assert_allclose(sf1, sf0) + + +class TestIrwinHall: + unif = stats.uniform(0, 1) + ih1 = stats.irwinhall(1) + ih10 = stats.irwinhall(10) + + def test_stats_ih10(self): + # from Wolfram Alpha "mean variance skew kurtosis UniformSumDistribution[10]" + # W|A uses Pearson's definition of kurtosis so subtract 3 + # should be exact integer division converted to fp64, without any further ops + assert_array_max_ulp(self.ih10.stats('mvsk'), (5, 10/12, 0, -3/25)) + + def test_moments_ih10(self): + # from Wolfram Alpha "values moments UniformSumDistribution[10]" + # algo should use integer division converted to fp64, without any further ops + # so these should be precise to the ulpm if not exact + vals = [5, 155 / 6, 275 / 2, 752, 12650 / 3, + 677465 / 28, 567325 / 4, + 15266213 / 18, 10333565 / 2] + moments = [self.ih10.moment(n+1) for n in range(len(vals))] + assert_array_max_ulp(moments, vals) + # also from Wolfram Alpha "50th moment UniformSumDistribution[10]" + m50 = self.ih10.moment(50) + m50_exact = 17453002755350010529309685557285098151740985685/4862 + assert_array_max_ulp(m50, m50_exact) + + def test_pdf_ih1_unif(self): + # IH(1) PDF is by definition U(0,1) + # we should be too, but differences in floating point eval order happen + # it's unclear if we can get down to the single ulp for doubles unless + # quads are used we're within 6-10 ulps otherwise (across sf/cdf/pdf) + # which is pretty good + + pts = np.linspace(0, 1, 100) + pdf_unif = self.unif.pdf(pts) + pdf_ih1 = self.ih1.pdf(pts) + assert_array_max_ulp(pdf_ih1, pdf_unif, maxulp=10) + + def test_pdf_ih2_triangle(self): + # IH(2) PDF is a triangle + ih2 = stats.irwinhall(2) + npts = 101 + pts = np.linspace(0, 2, npts) + expected = np.linspace(0, 2, npts) + expected[(npts + 1) // 2:] = 2 - expected[(npts + 1) // 2:] + pdf_ih2 = ih2.pdf(pts) + assert_array_max_ulp(pdf_ih2, expected, maxulp=10) + + def test_cdf_ih1_unif(self): + # CDF of IH(1) should be identical to uniform + pts = np.linspace(0, 1, 100) + cdf_unif = self.unif.cdf(pts) + cdf_ih1 = self.ih1.cdf(pts) + + assert_array_max_ulp(cdf_ih1, cdf_unif, maxulp=10) + + def test_cdf(self): + # CDF of IH is symmetric so CDF should be 0.5 at n/2 + n = np.arange(1, 10) + ih = stats.irwinhall(n) + ih_cdf = ih.cdf(n / 2) + exact = np.repeat(1/2, len(n)) + # should be identically 1/2 but fp order of eval differences happen + assert_array_max_ulp(ih_cdf, exact, maxulp=10) + + def test_cdf_ih10_exact(self): + # from Wolfram Alpha "values CDF[UniformSumDistribution[10], x] x=0 to x=10" + # symmetric about n/2, i.e., cdf[n-x] = 1-cdf[x] = sf[x] + vals = [0, 1 / 3628800, 169 / 604800, 24427 / 1814400, + 252023 / 1814400, 1 / 2, 1562377 / 1814400, + 1789973 / 1814400, 604631 / 604800, + 3628799 / 3628800, 1] + + # essentially a test of bspline evaluation + # this and the other ones are mostly to detect regressions + assert_array_max_ulp(self.ih10.cdf(np.arange(11)), vals, maxulp=10) + + assert_array_max_ulp(self.ih10.cdf(1/10), 1/36288000000000000, maxulp=10) + ref = 36287999999999999/36288000000000000 + assert_array_max_ulp(self.ih10.cdf(99/10), ref, maxulp=10) + + def test_pdf_ih10_exact(self): + # from Wolfram Alpha "values PDF[UniformSumDistribution[10], x] x=0 to x=10" + # symmetric about n/2 = 5 + vals = [0, 1 / 362880, 251 / 181440, 913 / 22680, 44117 / 181440] + vals += [15619 / 36288] + vals[::-1] + assert_array_max_ulp(self.ih10.pdf(np.arange(11)), vals, maxulp=10) + + def test_sf_ih10_exact(self): + assert_allclose(self.ih10.sf(np.arange(11)), 1 - self.ih10.cdf(np.arange(11))) + # from Wolfram Alpha "SurvivalFunction[UniformSumDistribution[10],x] at x=1/10" + # and symmetry about n/2 = 5 + # W|A returns 1 for CDF @ x=9.9 + ref = 36287999999999999/36288000000000000 + assert_array_max_ulp(self.ih10.sf(1/10), ref, maxulp=10) + + +class TestDParetoLognorm: + def test_against_R(self): + # Test against R implementation in `distributionsrd` + # library(distributionsrd) + # options(digits=16) + # x = 1.1 + # b = 2 + # a = 1.5 + # m = 3 + # s = 1.2 + # ddoubleparetolognormal(x, b, a, m, s) + # pdoubleparetolognormal(x, b, a, m, s) + x, m, s, a, b = 1.1, 3, 1.2, 1.5, 2 + dist = stats.dpareto_lognorm(m, s, a, b) + np.testing.assert_allclose(dist.pdf(x), 0.02490187219085912) + np.testing.assert_allclose(dist.cdf(x), 0.01664024173822796) + + +# Cases are (distribution name, log10 of smallest probability mass to test, +# log10 of the complement of the largest probability mass to test, atol, +# rtol). None uses default values. +@pytest.mark.parametrize("case", [("kappa3", None, None, None, None), + ("loglaplace", None, None, None, None), + ("lognorm", None, None, None, None), + ("lomax", None, None, None, None), + ("pareto", None, None, None, None),]) +def test_sf_isf_overrides(case): + # Test that SF is the inverse of ISF. Supplements + # `test_continuous_basic.check_sf_isf` for distributions with overridden + # `sf` and `isf` methods. + distname, lp1, lp2, atol, rtol = case + + lpm = np.log10(0.5) # log10 of the probability mass at the median + lp1 = lp1 or -290 + lp2 = lp2 or -14 + atol = atol or 0 + rtol = rtol or 1e-12 + dist = getattr(stats, distname) + params = dict(distcont)[distname] + dist_frozen = dist(*params) + + # Test (very deep) right tail to median. We can benchmark with random + # (loguniform) points, but strictly logspaced points are fine for tests. + ref = np.logspace(lp1, lpm) + res = dist_frozen.sf(dist_frozen.isf(ref)) + assert_allclose(res, ref, atol=atol, rtol=rtol) + + # test median to left tail + ref = 1 - np.logspace(lp2, lpm, 20) + res = dist_frozen.sf(dist_frozen.isf(ref)) + assert_allclose(res, ref, atol=atol, rtol=rtol) diff --git a/v12x/pylib/scipy/stats/tests/test_entropy.py b/v12x/pylib/scipy/stats/tests/test_entropy.py new file mode 100644 index 0000000000000000000000000000000000000000..1302d21a8eedf9804176df9d9096a86b325d3558 --- /dev/null +++ b/v12x/pylib/scipy/stats/tests/test_entropy.py @@ -0,0 +1,317 @@ +import math +import pytest +from pytest import raises as assert_raises + +import numpy as np + +from scipy import stats +from scipy.stats import norm, expon # type: ignore[attr-defined] +from scipy._lib._array_api import make_xp_test_case +from scipy._lib._array_api_no_0d import (xp_assert_close, xp_assert_equal, + xp_assert_less) + +lazy_xp_modules = [stats] +skip_xp_backends = pytest.mark.skip_xp_backends + +@make_xp_test_case(stats.entropy) +class TestEntropy: + def test_entropy_positive(self, xp): + # See ticket #497 + pk = xp.asarray([0.5, 0.2, 0.3]) + qk = xp.asarray([0.1, 0.25, 0.65]) + eself = stats.entropy(pk, pk) + edouble = stats.entropy(pk, qk) + xp_assert_close(eself, xp.asarray(0.), atol=1e-16) + xp_assert_less(-edouble, xp.asarray(0.)) + + def test_entropy_base(self, xp): + pk = xp.ones(16) + S = stats.entropy(pk, base=2.) + xp_assert_less(xp.abs(S - 4.), xp.asarray(1.e-5)) + + qk = xp.ones(16) + qk = xp.where(xp.arange(16) < 8, 2., qk) + S = stats.entropy(pk, qk) + S2 = stats.entropy(pk, qk, base=2.) + xp_assert_less(xp.abs(S/S2 - math.log(2.)), xp.asarray(1.e-5)) + + def test_entropy_zero(self, xp): + # Test for PR-479 + x = xp.asarray([0., 1., 2.]) + xp_assert_close(stats.entropy(x), + xp.asarray(0.63651416829481278)) + + def test_entropy_2d(self, xp): + pk = xp.asarray([[0.1, 0.2], [0.6, 0.3], [0.3, 0.5]]) + qk = xp.asarray([[0.2, 0.1], [0.3, 0.6], [0.5, 0.3]]) + xp_assert_close(stats.entropy(pk, qk), + xp.asarray([0.1933259, 0.18609809])) + + def test_entropy_2d_zero(self, xp): + pk = xp.asarray([[0.1, 0.2], [0.6, 0.3], [0.3, 0.5]]) + qk = xp.asarray([[0.0, 0.1], [0.3, 0.6], [0.5, 0.3]]) + xp_assert_close(stats.entropy(pk, qk), + xp.asarray([xp.inf, 0.18609809])) + + pk = xp.asarray([[0.0, 0.2], [0.6, 0.3], [0.3, 0.5]]) + xp_assert_close(stats.entropy(pk, qk), + xp.asarray([0.17403988, 0.18609809])) + + def test_entropy_base_2d_nondefault_axis(self, xp): + pk = xp.asarray([[0.1, 0.2], [0.6, 0.3], [0.3, 0.5]]) + xp_assert_close(stats.entropy(pk, axis=1), + xp.asarray([0.63651417, 0.63651417, 0.66156324])) + + def test_entropy_2d_nondefault_axis(self, xp): + pk = xp.asarray([[0.1, 0.2], [0.6, 0.3], [0.3, 0.5]]) + qk = xp.asarray([[0.2, 0.1], [0.3, 0.6], [0.5, 0.3]]) + xp_assert_close(stats.entropy(pk, qk, axis=1), + xp.asarray([0.23104906, 0.23104906, 0.12770641])) + + def test_entropy_raises_value_error(self, xp): + pk = xp.asarray([[0.1, 0.2], [0.6, 0.3], [0.3, 0.5]]) + qk = xp.asarray([[0.1, 0.2], [0.6, 0.3]]) + message = "Array shapes are incompatible for broadcasting." + with pytest.raises(ValueError, match=message): + stats.entropy(pk, qk) + + def test_base_entropy_with_axis_0_is_equal_to_default(self, xp): + pk = xp.asarray([[0.1, 0.2], [0.6, 0.3], [0.3, 0.5]]) + xp_assert_close(stats.entropy(pk, axis=0), + stats.entropy(pk)) + + def test_entropy_with_axis_0_is_equal_to_default(self, xp): + pk = xp.asarray([[0.1, 0.2], [0.6, 0.3], [0.3, 0.5]]) + qk = xp.asarray([[0.2, 0.1], [0.3, 0.6], [0.5, 0.3]]) + xp_assert_close(stats.entropy(pk, qk, axis=0), + stats.entropy(pk, qk)) + + def test_base_entropy_transposed(self, xp): + pk = xp.asarray([[0.1, 0.2], [0.6, 0.3], [0.3, 0.5]]) + xp_assert_close(stats.entropy(pk.T), + stats.entropy(pk, axis=1)) + + def test_entropy_transposed(self, xp): + pk = xp.asarray([[0.1, 0.2], [0.6, 0.3], [0.3, 0.5]]) + qk = xp.asarray([[0.2, 0.1], [0.3, 0.6], [0.5, 0.3]]) + xp_assert_close(stats.entropy(pk.T, qk.T), + stats.entropy(pk, qk, axis=1)) + + def test_entropy_broadcasting(self, xp): + rng = np.random.default_rng(74187315492831452) + x = xp.asarray(rng.random(3)) + y = xp.asarray(rng.random((2, 1))) + res = stats.entropy(x, y, axis=-1) + xp_assert_close(res[0], stats.entropy(x, y[0, ...])) + xp_assert_close(res[1], stats.entropy(x, y[1, ...])) + + def test_entropy_shape_mismatch(self, xp): + x = xp.ones((10, 1, 12)) + y = xp.ones((11, 2)) + message = "Array shapes are incompatible for broadcasting." + with pytest.raises(ValueError, match=message): + stats.entropy(x, y) + + def test_input_validation(self, xp): + x = xp.ones(10) + message = "`base` must be a positive number." + with pytest.raises(ValueError, match=message): + stats.entropy(x, base=-2) + + +@make_xp_test_case(stats.differential_entropy) +class TestDifferentialEntropy: + """ + Vasicek results are compared with the R package vsgoftest. + + # library(vsgoftest) + # + # samp <- c() + # entropy.estimate(x = samp, window = ) + + """ + methods = pytest.mark.parametrize('method', [ + "vasicek", + "van es", + pytest.param( + "correa", + marks=[ + skip_xp_backends("array_api_strict", reason="Invalid fancy indexing"), + skip_xp_backends("dask.array", reason="Invalid fancy indexing"), + ], + ), + "ebrahimi", + ]) + + def test_differential_entropy_vasicek(self, xp): + + random_state = np.random.RandomState(0) + values = random_state.standard_normal(100) + values = xp.asarray(values.tolist()) + + entropy = stats.differential_entropy(values, method='vasicek') + xp_assert_close(entropy, xp.asarray(1.342551187000946)) + + entropy = stats.differential_entropy(values, window_length=1, + method='vasicek') + xp_assert_close(entropy, xp.asarray(1.122044177725947)) + + entropy = stats.differential_entropy(values, window_length=8, + method='vasicek') + xp_assert_close(entropy, xp.asarray(1.349401487550325)) + + def test_differential_entropy_vasicek_2d_nondefault_axis(self, xp): + random_state = np.random.RandomState(0) + values = random_state.standard_normal((3, 100)) + values = xp.asarray(values.tolist()) + + entropy = stats.differential_entropy(values, axis=1, method='vasicek') + ref = xp.asarray([1.342551187000946, 1.341825903922332, 1.293774601883585]) + xp_assert_close(entropy, ref) + + entropy = stats.differential_entropy(values, axis=1, window_length=1, + method='vasicek') + ref = xp.asarray([1.122044177725947, 1.10294413850758, 1.129615790292772]) + xp_assert_close(entropy, ref) + + entropy = stats.differential_entropy(values, axis=1, window_length=8, + method='vasicek') + ref = xp.asarray([1.349401487550325, 1.338514126301301, 1.292331889365405]) + xp_assert_close(entropy, ref) + + + def test_differential_entropy_raises_value_error(self, xp): + random_state = np.random.RandomState(0) + values = random_state.standard_normal((3, 100)) + values = xp.asarray(values.tolist()) + + error_str = ( + r"Window length \({window_length}\) must be positive and less " + r"than half the sample size \({sample_size}\)." + ) + + sample_size = values.shape[1] + + for window_length in {-1, 0, sample_size//2, sample_size}: + + formatted_error_str = error_str.format( + window_length=window_length, + sample_size=sample_size, + ) + + with assert_raises(ValueError, match=formatted_error_str): + stats.differential_entropy( + values, + window_length=window_length, + axis=1, + ) + + def test_base_differential_entropy_with_axis_0_is_equal_to_default(self, xp): + random_state = np.random.RandomState(0) + values = random_state.standard_normal((100, 3)) + values = xp.asarray(values.tolist()) + + entropy = stats.differential_entropy(values, axis=0) + default_entropy = stats.differential_entropy(values) + xp_assert_close(entropy, default_entropy) + + def test_base_differential_entropy_transposed(self, xp): + random_state = np.random.RandomState(0) + values = random_state.standard_normal((3, 100)) + values = xp.asarray(values.tolist()) + + xp_assert_close( + stats.differential_entropy(values.T), + stats.differential_entropy(values, axis=1), + ) + + def test_input_validation(self, xp): + x = np.ones(10) + x = xp.asarray(x.tolist()) + + message = "`base` must be a positive number or `None`." + with pytest.raises(ValueError, match=message): + stats.differential_entropy(x, base=-2) + + message = "`method` must be one of..." + with pytest.raises(ValueError, match=message): + stats.differential_entropy(x, method='ekki-ekki') + + def test_window_length_is_none(self, xp): + rng = np.random.default_rng(358923459826738562) + x = xp.asarray(rng.random(size=10)) + ref = stats.differential_entropy(x) + res = stats.differential_entropy(x, window_length=None) + xp_assert_close(res, ref, rtol=0.005) + + @methods + def test_consistency(self, method, xp): + # test that method is a consistent estimator + n = 10000 if method == 'correa' else 1000000 + rvs = stats.norm.rvs(size=n, random_state=0) + rvs = xp.asarray(rvs.tolist()) + expected = xp.asarray(float(stats.norm.entropy())) + res = stats.differential_entropy(rvs, method=method) + xp_assert_close(res, expected, rtol=0.005) + + # values from differential_entropy reference [6], table 1, n=50, m=7 + norm_rmse_std_cases = { # method: (RMSE, STD) + 'vasicek': (0.198, 0.109), + 'van es': (0.212, 0.110), + 'correa': (0.135, 0.112), + 'ebrahimi': (0.128, 0.109) + } + + # values from differential_entropy reference [6], table 2, n=50, m=7 + expon_rmse_std_cases = { # method: (RMSE, STD) + 'vasicek': (0.194, 0.148), + 'van es': (0.179, 0.149), + 'correa': (0.155, 0.152), + 'ebrahimi': (0.151, 0.148) + } + + rmse_std_cases = {norm: norm_rmse_std_cases, + expon: expon_rmse_std_cases} + + @methods + @pytest.mark.parametrize('dist', [norm, expon]) + def test_rmse_std(self, method, dist, xp): + # test that RMSE and standard deviation of estimators matches values + # given in differential_entropy reference [6]. Incidentally, also + # tests vectorization. + reps, n, m = 10000, 50, 7 + expected = self.rmse_std_cases[dist][method] + rmse_expected, std_expected = xp.asarray(expected[0]), xp.asarray(expected[1]) + rvs = dist.rvs(size=(reps, n), random_state=0) + rvs = xp.asarray(rvs.tolist()) + true_entropy = xp.asarray(float(dist.entropy())) + res = stats.differential_entropy(rvs, window_length=m, + method=method, axis=-1) + xp_assert_close(xp.sqrt(xp.mean((res - true_entropy)**2)), + rmse_expected, atol=0.005) + xp_assert_close(xp.std(res, correction=0), std_expected, atol=0.002) + + @pytest.mark.parametrize('n, method', [ + (8, 'van es'), + (12, 'ebrahimi'), + (1001, 'vasicek') + ]) + def test_method_auto(self, n, method, xp): + rvs = stats.norm.rvs(size=(n,), random_state=0) + rvs = xp.asarray(rvs.tolist()) + res1 = stats.differential_entropy(rvs) + res2 = stats.differential_entropy(rvs, method=method) + xp_assert_equal(res1, res2) + + @methods + @pytest.mark.parametrize('dtype', [None, 'float32', 'float64']) + def test_dtypes_gh21192(self, method, dtype, xp): + # gh-21192 noted a change in the output of method='ebrahimi' + # with integer input. Check that the output is consistent regardless + # of input dtype. + x = [1, 1, 2, 3, 3, 4, 5, 5, 6, 7, 8, 9, 10, 11] + dtype_in = getattr(xp, str(dtype), None) + dtype_out = getattr(xp, str(dtype), xp.asarray(1.).dtype) + res = stats.differential_entropy(xp.asarray(x, dtype=dtype_in), method=method) + ref = stats.differential_entropy(xp.asarray(x, dtype=xp.float64), method=method) + xp_assert_close(res, xp.asarray(ref, dtype=dtype_out)[()]) diff --git a/v12x/pylib/scipy/stats/tests/test_fast_gen_inversion.py b/v12x/pylib/scipy/stats/tests/test_fast_gen_inversion.py new file mode 100644 index 0000000000000000000000000000000000000000..dc57dd17d4562909f0a21b23a6ff2768c7373e3a --- /dev/null +++ b/v12x/pylib/scipy/stats/tests/test_fast_gen_inversion.py @@ -0,0 +1,434 @@ +import pytest +import warnings +import numpy as np +from numpy.testing import assert_array_equal, assert_allclose +from copy import deepcopy +from scipy.stats.sampling import FastGeneratorInversion +from scipy import stats +from scipy._lib._testutils import IS_MUSL + + +def test_bad_args(): + # loc and scale must be scalar + with pytest.raises(ValueError, match="loc must be scalar"): + FastGeneratorInversion(stats.norm(loc=(1.2, 1.3))) + with pytest.raises(ValueError, match="scale must be scalar"): + FastGeneratorInversion(stats.norm(scale=[1.5, 5.7])) + + with pytest.raises(ValueError, match="'test' cannot be used to seed"): + FastGeneratorInversion(stats.norm(), random_state="test") + + msg = "Each of the 1 shape parameters must be a scalar" + with pytest.raises(ValueError, match=msg): + FastGeneratorInversion(stats.gamma([1.3, 2.5])) + + with pytest.raises(ValueError, match="`dist` must be a frozen"): + FastGeneratorInversion("xy") + + with pytest.raises(ValueError, match="Distribution 'truncnorm' is not"): + FastGeneratorInversion(stats.truncnorm(1.3, 4.5)) + + +def test_random_state(): + # fixed seed + gen = FastGeneratorInversion(stats.norm(), random_state=68734509) + x1 = gen.rvs(size=10) + gen.random_state = 68734509 + x2 = gen.rvs(size=10) + assert_array_equal(x1, x2) + + # Generator + urng = np.random.default_rng(20375857) + gen = FastGeneratorInversion(stats.norm(), random_state=urng) + x1 = gen.rvs(size=10) + gen.random_state = np.random.default_rng(20375857) + x2 = gen.rvs(size=10) + assert_array_equal(x1, x2) + + # RandomState + urng = np.random.RandomState(2364) + gen = FastGeneratorInversion(stats.norm(), random_state=urng) + x1 = gen.rvs(size=10) + gen.random_state = np.random.RandomState(2364) + x2 = gen.rvs(size=10) + assert_array_equal(x1, x2) + + # if evaluate_error is called, it must not interfere with the random_state + # used by rvs + gen = FastGeneratorInversion(stats.norm(), random_state=68734509) + x1 = gen.rvs(size=10) + _ = gen.evaluate_error(size=5) # this will generate 5 uniform rvs + x2 = gen.rvs(size=10) + gen.random_state = 68734509 + x3 = gen.rvs(size=20) + assert_array_equal(x2, x3[10:]) + + +dists_with_params = [ + ("alpha", (3.5,)), + ("anglit", ()), + ("argus", (3.5,)), + ("argus", (5.1,)), + ("beta", (1.5, 0.9)), + ("cosine", ()), + ("betaprime", (2.5, 3.3)), + ("bradford", (1.2,)), + ("burr", (1.3, 2.4)), + ("burr12", (0.7, 1.2)), + ("cauchy", ()), + ("chi2", (3.5,)), + ("chi", (4.5,)), + ("crystalball", (0.7, 1.2)), + ("expon", ()), + ("gamma", (1.5,)), + ("gennorm", (2.7,)), + ("gumbel_l", ()), + ("gumbel_r", ()), + ("hypsecant", ()), + ("invgauss", (3.1,)), + ("invweibull", (1.5,)), + ("laplace", ()), + ("logistic", ()), + ("maxwell", ()), + ("moyal", ()), + ("norm", ()), + ("pareto", (1.3,)), + ("powerlaw", (7.6,)), + ("rayleigh", ()), + ("semicircular", ()), + ("t", (5.7,)), + ("wald", ()), + ("weibull_max", (2.4,)), + ("weibull_min", (1.2,)), +] + + +@pytest.mark.parametrize(("distname, args"), dists_with_params) +def test_rvs_and_ppf(distname, args): + # check sample against rvs generated by rv_continuous + urng = np.random.default_rng(9807324628097097) + rng1 = getattr(stats, distname)(*args) + rvs1 = rng1.rvs(size=500, random_state=urng) + rng2 = FastGeneratorInversion(rng1, random_state=urng) + rvs2 = rng2.rvs(size=500) + assert stats.cramervonmises_2samp(rvs1, rvs2).pvalue > 0.01 + + # check ppf + q = [0.001, 0.1, 0.5, 0.9, 0.999] + assert_allclose(rng1.ppf(q), rng2.ppf(q), atol=1e-10) + + +@pytest.mark.parametrize(("distname, args"), dists_with_params) +def test_u_error(distname, args): + # check sample against rvs generated by rv_continuous + dist = getattr(stats, distname)(*args) + with warnings.catch_warnings(): + # filter the warnings thrown by UNU.RAN + warnings.simplefilter("ignore", RuntimeWarning) + rng = FastGeneratorInversion(dist) + u_error, x_error = rng.evaluate_error( + size=10_000, random_state=9807324628097097, x_error=False + ) + assert u_error <= 1e-10 + + +@pytest.mark.xslow +@pytest.mark.xfail(reason="geninvgauss CDF is not accurate") +def test_geninvgauss_uerror(): + dist = stats.geninvgauss(3.2, 1.5) + rng = FastGeneratorInversion(dist) + err = rng.evaluate_error(size=10_000, random_state=67982) + assert err[0] < 1e-10 + + +# TODO: add more distributions +@pytest.mark.skipif(IS_MUSL, reason="Hits RecursionError, see gh-23172") +@pytest.mark.fail_slow(5) +@pytest.mark.parametrize(("distname, args"), [("beta", (0.11, 0.11))]) +def test_error_extreme_params(distname, args): + # take extreme parameters where u-error might not be below the tolerance + # due to limitations of floating point arithmetic + with warnings.catch_warnings(): + # filter the warnings thrown by UNU.RAN for such extreme parameters + warnings.simplefilter("ignore", RuntimeWarning) + dist = getattr(stats, distname)(*args) + rng = FastGeneratorInversion(dist) + u_error, x_error = rng.evaluate_error( + size=10_000, random_state=980732462809709732623, x_error=True + ) + if u_error >= 2.5 * 1e-10: + assert x_error < 1e-9 + + +def test_evaluate_error_inputs(): + gen = FastGeneratorInversion(stats.norm()) + with pytest.raises(ValueError, match="size must be an integer"): + gen.evaluate_error(size=3.5) + with pytest.raises(ValueError, match="size must be an integer"): + gen.evaluate_error(size=(3, 3)) + + +def test_rvs_ppf_loc_scale(): + loc, scale = 3.5, 2.3 + dist = stats.norm(loc=loc, scale=scale) + rng = FastGeneratorInversion(dist, random_state=1234) + r = rng.rvs(size=1000) + r_rescaled = (r - loc) / scale + assert stats.cramervonmises(r_rescaled, "norm").pvalue > 0.01 + q = [0.001, 0.1, 0.5, 0.9, 0.999] + assert_allclose(rng._ppf(q), rng.ppf(q), atol=1e-10) + + +def test_domain(): + # only a basic check that the domain argument is passed to the + # UNU.RAN generators + rng = FastGeneratorInversion(stats.norm(), domain=(-1, 1)) + r = rng.rvs(size=100) + assert -1 <= r.min() < r.max() <= 1 + + # if loc and scale are used, new domain is loc + scale*domain + loc, scale = 3.5, 1.3 + dist = stats.norm(loc=loc, scale=scale) + rng = FastGeneratorInversion(dist, domain=(-1.5, 2)) + r = rng.rvs(size=100) + lb, ub = loc - scale * 1.5, loc + scale * 2 + assert lb <= r.min() < r.max() <= ub + + +@pytest.mark.parametrize(("distname, args, expected"), + [("beta", (3.5, 2.5), (0, 1)), + ("norm", (), (-np.inf, np.inf))]) +def test_support(distname, args, expected): + # test that the support is updated if truncation and loc/scale are applied + # use beta distribution since it is a transformed betaprime distribution, + # so it is important that the correct support is considered + # (i.e., the support of beta is (0,1), while betaprime is (0, inf)) + dist = getattr(stats, distname)(*args) + rng = FastGeneratorInversion(dist) + assert_array_equal(rng.support(), expected) + rng.loc = 1 + rng.scale = 2 + assert_array_equal(rng.support(), 1 + 2*np.array(expected)) + + +@pytest.mark.parametrize(("distname, args"), + [("beta", (3.5, 2.5)), ("norm", ())]) +def test_support_truncation(distname, args): + # similar test for truncation + dist = getattr(stats, distname)(*args) + rng = FastGeneratorInversion(dist, domain=(0.5, 0.7)) + assert_array_equal(rng.support(), (0.5, 0.7)) + rng.loc = 1 + rng.scale = 2 + assert_array_equal(rng.support(), (1 + 2 * 0.5, 1 + 2 * 0.7)) + + +def test_domain_shift_truncation(): + # center of norm is zero, it should be shifted to the left endpoint of + # domain. if this was not the case, PINV in UNURAN would raise a warning + # as the center is not inside the domain + with warnings.catch_warnings(): + warnings.simplefilter("error") + rng = FastGeneratorInversion(stats.norm(), domain=(1, 2)) + r = rng.rvs(size=100) + assert 1 <= r.min() < r.max() <= 2 + + +def test_non_rvs_methods_with_domain(): + # as a first step, compare truncated normal against stats.truncnorm + rng = FastGeneratorInversion(stats.norm(), domain=(2.3, 3.2)) + trunc_norm = stats.truncnorm(2.3, 3.2) + # take values that are inside and outside the domain + x = (2.0, 2.4, 3.0, 3.4) + p = (0.01, 0.5, 0.99) + assert_allclose(rng._cdf(x), trunc_norm.cdf(x)) + assert_allclose(rng._ppf(p), trunc_norm.ppf(p)) + loc, scale = 2, 3 + rng.loc = 2 + rng.scale = 3 + trunc_norm = stats.truncnorm(2.3, 3.2, loc=loc, scale=scale) + x = np.array(x) * scale + loc + assert_allclose(rng._cdf(x), trunc_norm.cdf(x)) + assert_allclose(rng._ppf(p), trunc_norm.ppf(p)) + + # do another sanity check with beta distribution + # in that case, it is important to use the correct domain since beta + # is a transformation of betaprime which has a different support + rng = FastGeneratorInversion(stats.beta(2.5, 3.5), domain=(0.3, 0.7)) + rng.loc = 2 + rng.scale = 2.5 + # the support is 2.75, , 3.75 (2 + 2.5 * 0.3, 2 + 2.5 * 0.7) + assert_array_equal(rng.support(), (2.75, 3.75)) + x = np.array([2.74, 2.76, 3.74, 3.76]) + # the cdf needs to be zero outside of the domain + y_cdf = rng._cdf(x) + assert_array_equal((y_cdf[0], y_cdf[3]), (0, 1)) + assert np.min(y_cdf[1:3]) > 0 + # ppf needs to map 0 and 1 to the boundaries + assert_allclose(rng._ppf(y_cdf), (2.75, 2.76, 3.74, 3.75)) + + +def test_non_rvs_methods_without_domain(): + norm_dist = stats.norm() + rng = FastGeneratorInversion(norm_dist) + x = np.linspace(-3, 3, num=10) + p = (0.01, 0.5, 0.99) + assert_allclose(rng._cdf(x), norm_dist.cdf(x)) + assert_allclose(rng._ppf(p), norm_dist.ppf(p)) + loc, scale = 0.5, 1.3 + rng.loc = loc + rng.scale = scale + norm_dist = stats.norm(loc=loc, scale=scale) + assert_allclose(rng._cdf(x), norm_dist.cdf(x)) + assert_allclose(rng._ppf(p), norm_dist.ppf(p)) + +@pytest.mark.parametrize(("domain, x"), + [(None, 0.5), + ((0, 1), 0.5), + ((0, 1), 1.5)]) +def test_scalar_inputs(domain, x): + """ pdf, cdf etc should map scalar values to scalars. check with and + w/o domain since domain impacts pdf, cdf etc + Take x inside and outside of domain """ + rng = FastGeneratorInversion(stats.norm(), domain=domain) + assert np.isscalar(rng._cdf(x)) + assert np.isscalar(rng._ppf(0.5)) + + +def test_domain_argus_large_chi(): + # for large chi, the Gamma distribution is used and the domain has to be + # transformed. this is a test to ensure that the transformation works + chi, lb, ub = 5.5, 0.25, 0.75 + rng = FastGeneratorInversion(stats.argus(chi), domain=(lb, ub)) + rng.random_state = 4574 + r = rng.rvs(size=500) + assert lb <= r.min() < r.max() <= ub + # perform goodness of fit test with conditional cdf + cdf = stats.argus(chi).cdf + prob = cdf(ub) - cdf(lb) + assert stats.cramervonmises(r, lambda x: cdf(x) / prob).pvalue > 0.05 + + +def test_setting_loc_scale(): + rng = FastGeneratorInversion(stats.norm(), random_state=765765864) + r1 = rng.rvs(size=1000) + rng.loc = 3.0 + rng.scale = 2.5 + r2 = rng.rvs(1000) + # rescaled r2 should be again standard normal + assert stats.cramervonmises_2samp(r1, (r2 - 3) / 2.5).pvalue > 0.05 + # reset values to default loc=0, scale=1 + rng.loc = 0 + rng.scale = 1 + r2 = rng.rvs(1000) + assert stats.cramervonmises_2samp(r1, r2).pvalue > 0.05 + + +def test_ignore_shape_range(): + msg = "No generator is defined for the shape parameters" + with pytest.raises(ValueError, match=msg): + rng = FastGeneratorInversion(stats.t(0.03)) + rng = FastGeneratorInversion(stats.t(0.03), ignore_shape_range=True) + # we can ignore the recommended range of shape parameters + # but u-error can be expected to be too large in that case + u_err, _ = rng.evaluate_error(size=1000, random_state=234) + assert u_err >= 1e-6 + +@pytest.mark.xfail_on_32bit( + "NumericalInversePolynomial.qrvs fails for Win 32-bit" +) +class TestQRVS: + def test_input_validation(self): + gen = FastGeneratorInversion(stats.norm()) + + match = "`qmc_engine` must be an instance of..." + with pytest.raises(ValueError, match=match): + gen.qrvs(qmc_engine=0) + + match = "`d` must be consistent with dimension of `qmc_engine`." + with pytest.raises(ValueError, match=match): + gen.qrvs(d=3, qmc_engine=stats.qmc.Halton(2)) + + qrngs = [None, stats.qmc.Sobol(1, seed=0), stats.qmc.Halton(3, seed=0)] + # `size=None` should not add anything to the shape, `size=1` should + sizes = [ + (None, tuple()), + (1, (1,)), + (4, (4,)), + ((4,), (4,)), + ((2, 4), (2, 4)), + ] + # Neither `d=None` nor `d=1` should add anything to the shape + ds = [(None, tuple()), (1, tuple()), (3, (3,))] + + @pytest.mark.parametrize("qrng", qrngs) + @pytest.mark.parametrize("size_in, size_out", sizes) + @pytest.mark.parametrize("d_in, d_out", ds) + def test_QRVS_shape_consistency(self, qrng, size_in, size_out, + d_in, d_out): + gen = FastGeneratorInversion(stats.norm()) + + # If d and qrng.d are inconsistent, an error is raised + if d_in is not None and qrng is not None and qrng.d != d_in: + match = "`d` must be consistent with dimension of `qmc_engine`." + with pytest.raises(ValueError, match=match): + gen.qrvs(size_in, d=d_in, qmc_engine=qrng) + return + + # Sometimes d is really determined by qrng + if d_in is None and qrng is not None and qrng.d != 1: + d_out = (qrng.d,) + + shape_expected = size_out + d_out + + qrng2 = deepcopy(qrng) + qrvs = gen.qrvs(size=size_in, d=d_in, qmc_engine=qrng) + if size_in is not None: + assert qrvs.shape == shape_expected + + if qrng2 is not None: + uniform = qrng2.random(np.prod(size_in) or 1) + qrvs2 = stats.norm.ppf(uniform).reshape(shape_expected) + assert_allclose(qrvs, qrvs2, atol=1e-12) + + def test_QRVS_size_tuple(self): + # QMCEngine samples are always of shape (n, d). When `size` is a tuple, + # we set `n = prod(size)` in the call to qmc_engine.random, transform + # the sample, and reshape it to the final dimensions. When we reshape, + # we need to be careful, because the _columns_ of the sample returned + # by a QMCEngine are "independent"-ish, but the elements within the + # columns are not. We need to make sure that this doesn't get mixed up + # by reshaping: qrvs[..., i] should remain "independent"-ish of + # qrvs[..., i+1], but the elements within qrvs[..., i] should be + # transformed from the same low-discrepancy sequence. + + gen = FastGeneratorInversion(stats.norm()) + + size = (3, 4) + d = 5 + qrng = stats.qmc.Halton(d, seed=0) + qrng2 = stats.qmc.Halton(d, seed=0) + + uniform = qrng2.random(np.prod(size)) + + qrvs = gen.qrvs(size=size, d=d, qmc_engine=qrng) + qrvs2 = stats.norm.ppf(uniform) + + for i in range(d): + sample = qrvs[..., i] + sample2 = qrvs2[:, i].reshape(size) + assert_allclose(sample, sample2, atol=1e-12) + + +def test_burr_overflow(): + # this case leads to an overflow error if math.exp is used + # in the definition of the burr pdf instead of np.exp + # a direct implementation of the PDF as x**(-c-1) / (1+x**(-c))**(d+1) + # also leads to an overflow error in the setup + args = (1.89128135, 0.30195177) + with warnings.catch_warnings(): + # filter potential overflow warning + warnings.simplefilter("ignore", RuntimeWarning) + gen = FastGeneratorInversion(stats.burr(*args)) + u_error, _ = gen.evaluate_error(random_state=4326) + assert u_error <= 1e-10 diff --git a/v12x/pylib/scipy/stats/tests/test_fit.py b/v12x/pylib/scipy/stats/tests/test_fit.py new file mode 100644 index 0000000000000000000000000000000000000000..e77cb931308948fb48b080ba180e20b7b9e9aefb --- /dev/null +++ b/v12x/pylib/scipy/stats/tests/test_fit.py @@ -0,0 +1,1092 @@ +import os +import warnings + +import numpy as np +from numpy.testing import assert_allclose, assert_equal +import pytest +from scipy import stats +from scipy.optimize import differential_evolution + +from .test_continuous_basic import distcont +from scipy.stats._distn_infrastructure import FitError +from scipy.stats._distr_params import distdiscrete +from scipy.stats import goodness_of_fit + + +# this is not a proper statistical test for convergence, but only +# verifies that the estimate and true values don't differ by too much + +fit_sizes = [1000, 5000, 10000] # sample sizes to try + +thresh_percent = 0.25 # percent of true parameters for fail cut-off +thresh_min = 0.75 # minimum difference estimate - true to fail test + +mle_failing_fits = [ + 'dpareto_lognorm', + 'gausshyper', + 'genexpon', + 'gengamma', + 'irwinhall', + 'kappa4', + 'ksone', + 'kstwo', + 'ncf', + 'ncx2', + 'truncexpon', + 'tukeylambda', + 'vonmises', + 'levy_stable', + 'truncpareto', + 'truncweibull_min', + 'studentized_range', +] + +# these pass but are XSLOW (>1s) +mle_Xslow_fits = ['betaprime', 'crystalball', 'exponweib', 'f', 'geninvgauss', + 'jf_skew_t', 'nct', 'recipinvgauss', 'rel_breitwigner', + 'vonmises_line'] + +# The MLE fit method of these distributions doesn't perform well when all +# parameters are fit, so test them with the location fixed at 0. +mle_use_floc0 = [ + 'burr', + 'chi', + 'chi2', + 'mielke', + 'pearson3', + 'genhalflogistic', + 'rdist', + 'pareto', + 'powerlaw', # distfn.nnlf(est2, rvs) > distfn.nnlf(est1, rvs) otherwise + 'powerlognorm', + 'wrapcauchy', + 'rel_breitwigner', +] + +mm_failing_fits = ['alpha', 'betaprime', 'burr', 'burr12', 'cauchy', 'chi', + 'chi2', 'crystalball', 'dgamma', 'dpareto_lognorm', 'dweibull', + 'f', 'fatiguelife', 'fisk', 'foldcauchy', 'genextreme', + 'gengamma', 'genhyperbolic', 'gennorm', 'genpareto', + 'halfcauchy', 'invgamma', 'invweibull', 'irwinhall', 'jf_skew_t', + 'johnsonsu', 'kappa3', 'ksone', 'kstwo', 'landau', 'levy', 'levy_l', + 'levy_stable', 'loglaplace', 'lomax', 'mielke', 'nakagami', + 'ncf', 'nct', 'ncx2', 'pareto', 'powerlognorm', 'powernorm', + 'rel_breitwigner', 'skewcauchy', 't', 'triang', + 'truncpareto', 'truncweibull_min', 'tukeylambda', + 'studentized_range'] + +# not sure if these fail, but they caused my patience to fail +mm_XXslow_fits = ['argus', 'exponpow', 'exponweib', 'gausshyper', 'genexpon', + 'genhalflogistic', 'gompertz', 'johnsonsb', + 'kappa4', 'kstwobign', 'recipinvgauss', + 'truncexpon', 'vonmises', 'vonmises_line'] + +# these pass but are XSLOW (>1s) +mm_Xslow_fits = ['wrapcauchy'] + +failing_fits = {"MM": mm_failing_fits + mm_XXslow_fits, "MLE": mle_failing_fits} +xslow_fits = {"MM": mm_Xslow_fits, "MLE": mle_Xslow_fits} +fail_interval_censored = {"truncpareto"} + +# Don't run the fit test on these: +skip_fit = [ + 'erlang', # Subclass of gamma, generates a warning. + 'genhyperbolic', 'norminvgauss', # too slow +] + + +def cases_test_cont_fit(): + # this tests the closeness of the estimated parameters to the true + # parameters with fit method of continuous distributions + # Note: is slow, some distributions don't converge with sample + # size <= 10000 + for distname, arg in distcont: + if distname not in skip_fit: + yield distname, arg + + +@pytest.mark.slow +@pytest.mark.parametrize('distname,arg', list(cases_test_cont_fit())) +@pytest.mark.parametrize('method', ["MLE", "MM"]) +def test_cont_fit(distname, arg, method): + run_xfail = int(os.getenv('SCIPY_XFAIL', default=False)) + run_xslow = int(os.getenv('SCIPY_XSLOW', default=False)) + + if distname in failing_fits[method] and not run_xfail: + # The generic `fit` method can't be expected to work perfectly for all + # distributions, data, and guesses. Some failures are expected. + msg = "Failure expected; set environment variable SCIPY_XFAIL=1 to run." + pytest.xfail(msg) + + if distname in xslow_fits[method] and not run_xslow: + msg = "Very slow; set environment variable SCIPY_XSLOW=1 to run." + pytest.skip(msg) + + distfn = getattr(stats, distname) + + truearg = np.hstack([arg, [0.0, 1.0]]) + diffthreshold = np.max(np.vstack([truearg*thresh_percent, + np.full(distfn.numargs+2, thresh_min)]), + 0) + + for fit_size in fit_sizes: + # Note that if a fit succeeds, the other fit_sizes are skipped + rng = np.random.default_rng(1234) + + with np.errstate(all='ignore'): + rvs = distfn.rvs(size=fit_size, *arg, random_state=rng) + if method == 'MLE' and distfn.name in mle_use_floc0: + kwds = {'floc': 0} + else: + kwds = {} + # start with default values + est = distfn.fit(rvs, method=method, **kwds) + if method == 'MLE': + # Trivial test of the use of CensoredData. The fit() method + # will check that data contains no actual censored data, and + # do a regular uncensored fit. + data1 = stats.CensoredData(rvs) + est1 = distfn.fit(data1, **kwds) + msg = ('Different results fitting uncensored data wrapped as' + f' CensoredData: {distfn.name}: est={est} est1={est1}') + assert_allclose(est1, est, rtol=1e-10, err_msg=msg) + if method == 'MLE' and distname not in fail_interval_censored: + # Convert the first `nic` values in rvs to interval-censored + # values. The interval is small, so est2 should be close to + # est. + nic = 15 + interval = np.column_stack((rvs, rvs)) + interval[:nic, 0] *= 0.99 + interval[:nic, 1] *= 1.01 + interval.sort(axis=1) + data2 = stats.CensoredData(interval=interval) + est2 = distfn.fit(data2, **kwds) + msg = ('Different results fitting interval-censored' + f' data: {distfn.name}: est={est} est2={est2}') + assert_allclose(est2, est, rtol=0.05, err_msg=msg) + + diff = est - truearg + + # threshold for location + diffthreshold[-2] = np.max([np.abs(rvs.mean())*thresh_percent, + thresh_min]) + + if np.any(np.isnan(est)): + raise AssertionError('nan returned in fit') + else: + if np.all(np.abs(diff) <= diffthreshold): + break + else: + txt = f'parameter: {str(truearg)}\n' + txt += f'estimated: {str(est)}\n' + txt += f'diff : {str(diff)}\n' + raise AssertionError(f'fit not very good in {distfn.name}\n' + txt) + + +def _check_loc_scale_mle_fit(name, data, desired, atol=None): + d = getattr(stats, name) + actual = d.fit(data)[-2:] + assert_allclose(actual, desired, atol=atol, + err_msg=f'poor mle fit of (loc, scale) in {name}') + + +def test_non_default_loc_scale_mle_fit(): + data = np.array([1.01, 1.78, 1.78, 1.78, 1.88, 1.88, 1.88, 2.00]) + _check_loc_scale_mle_fit('uniform', data, [1.01, 0.99], 1e-3) + _check_loc_scale_mle_fit('expon', data, [1.01, 0.73875], 1e-3) + + +def test_expon_fit(): + """gh-6167""" + data = [0, 0, 0, 0, 2, 2, 2, 2] + phat = stats.expon.fit(data, floc=0) + assert_allclose(phat, [0, 1.0], atol=1e-3) + + +def test_fit_error(): + data = np.concatenate([np.zeros(29), np.ones(21)]) + message = "Optimization converged to parameters that are..." + with pytest.raises(FitError, match=message), \ + pytest.warns(RuntimeWarning): + stats.beta.fit(data) + + +@pytest.mark.parametrize("dist, params", + [(stats.norm, (0.5, 2.5)), # type: ignore[attr-defined] + (stats.binom, (10, 0.3, 2))]) # type: ignore[attr-defined] +def test_nnlf_and_related_methods(dist, params): + rng = np.random.default_rng(983459824) + + if hasattr(dist, 'pdf'): + logpxf = dist.logpdf + else: + logpxf = dist.logpmf + + x = dist.rvs(*params, size=100, random_state=rng) + ref = -logpxf(x, *params).sum() + res1 = dist.nnlf(params, x) + res2 = dist._penalized_nnlf(params, x) + assert_allclose(res1, ref) + assert_allclose(res2, ref) + + +def cases_test_fit_mle(): + # These fail default test or hang + skip_basic_fit = {'argus', 'irwinhall', 'foldnorm', 'truncpareto', + 'truncweibull_min', 'ksone', 'levy_stable', + 'studentized_range', 'kstwo', + 'beta', 'nakagami', 'truncnorm', # don't meet tolerance + 'poisson_binom'} # vector-valued shape parameter + + # Please keep this list in alphabetical order... + slow_basic_fit = {'alpha', 'arcsine', 'betaprime', 'binom', 'bradford', 'burr12', + 'chi', 'crystalball', 'dweibull', 'erlang', 'exponnorm', + 'exponpow', 'f', 'fatiguelife', 'fisk', 'foldcauchy', 'gamma', + 'genexpon', 'genextreme', 'gennorm', 'genpareto', + 'gompertz', 'invgamma', 'invgauss', 'invweibull', + 'jf_skew_t', 'johnsonsb', 'johnsonsu', 'kappa3', + 'kstwobign', 'loglaplace', 'lognorm', 'lomax', 'mielke', + 'nbinom', 'norminvgauss', + 'pareto', 'pearson3', 'powerlaw', 'powernorm', + 'randint', 'rdist', 'recipinvgauss', 'rice', 'skewnorm', + 't', 'uniform', 'weibull_max', 'weibull_min', 'wrapcauchy', + 'zipfian'} + + # Please keep this list in alphabetical order... + xslow_basic_fit = {'betabinom', 'betanbinom', 'burr', 'dpareto_lognorm', + 'exponweib', 'gausshyper', 'gengamma', 'genhalflogistic', + 'genhyperbolic', 'geninvgauss', + 'hypergeom', 'kappa4', 'loguniform', + 'ncf', 'nchypergeom_fisher', 'nchypergeom_wallenius', + 'nct', 'ncx2', 'nhypergeom', + 'powerlognorm', 'reciprocal', 'rel_breitwigner', + 'skellam', 'trapezoid', 'triang', + 'tukeylambda', 'vonmises'} + + for dist in dict(distdiscrete + distcont): + if dist in skip_basic_fit or not isinstance(dist, str): + reason = "tested separately" + yield pytest.param(dist, marks=pytest.mark.skip(reason=reason)) + elif dist in slow_basic_fit: + reason = "too slow (>= 0.25s)" + yield pytest.param(dist, marks=pytest.mark.slow(reason=reason)) + elif dist in xslow_basic_fit: + reason = "too slow (>= 1.0s)" + yield pytest.param(dist, marks=pytest.mark.xslow(reason=reason)) + else: + yield dist + + +def cases_test_fit_mse(): + # the first four are so slow that I'm not sure whether they would pass + skip_basic_fit = {'levy_stable', 'studentized_range', 'ksone', 'skewnorm', + 'irwinhall', # hangs + 'norminvgauss', # super slow (~1 hr) but passes + 'kstwo', # very slow (~25 min) but passes + 'geninvgauss', # quite slow (~4 minutes) but passes + 'gausshyper', 'genhyperbolic', # integration warnings + 'tukeylambda', # close, but doesn't meet tolerance + 'vonmises', # can have negative CDF; doesn't play nice + 'arcsine', 'argus', 'powerlaw', 'rdist', # don't meet tolerance + 'poisson_binom', # vector-valued shape parameter + } + + # Please keep this list in alphabetical order... + slow_basic_fit = {'alpha', 'anglit', 'betabinom', 'bradford', + 'chi', 'chi2', 'crystalball', 'dweibull', + 'erlang', 'exponnorm', 'exponpow', 'exponweib', + 'fatiguelife', 'fisk', 'foldcauchy', 'foldnorm', + 'gamma', 'genexpon', 'genextreme', 'genhalflogistic', + 'genlogistic', 'genpareto', 'gompertz', + 'hypergeom', 'invweibull', + 'johnsonsu', 'kappa3', 'kstwobign', + 'laplace_asymmetric', 'loggamma', 'loglaplace', + 'lognorm', 'lomax', + 'maxwell', 'nhypergeom', + 'pareto', 'powernorm', 'randint', 'recipinvgauss', + 'semicircular', + 't', 'triang', 'truncexpon', 'truncpareto', + 'uniform', + 'wald', 'weibull_max', 'weibull_min', 'wrapcauchy', + 'zipfian'} + + # Please keep this list in alphabetical order... + xslow_basic_fit = {'argus', 'beta', 'betaprime', 'burr', 'burr12', + 'dgamma', 'dpareto_lognorm', 'f', 'gengamma', 'gennorm', + 'invgamma', 'invgauss', 'jf_skew_t', + 'johnsonsb', 'kappa4', 'loguniform', 'mielke', + 'nakagami', 'ncf', 'nchypergeom_fisher', + 'nchypergeom_wallenius', 'nct', 'ncx2', + 'pearson3', 'powerlognorm', + 'reciprocal', 'rel_breitwigner', 'rice', + 'trapezoid', 'truncnorm', 'truncweibull_min', + 'vonmises_line'} + + warns_basic_fit = {'skellam'} # can remove mark after gh-14901 is resolved + + for dist in dict(distdiscrete + distcont): + if dist in skip_basic_fit or not isinstance(dist, str): + reason = "Fails. Oh well." + yield pytest.param(dist, marks=pytest.mark.skip(reason=reason)) + elif dist in slow_basic_fit: + reason = "too slow (>= 0.25s)" + yield pytest.param(dist, marks=pytest.mark.slow(reason=reason)) + elif dist in xslow_basic_fit: + reason = "too slow (>= 1.0s)" + yield pytest.param(dist, marks=pytest.mark.xslow(reason=reason)) + elif dist in warns_basic_fit: + mark = pytest.mark.filterwarnings('ignore::RuntimeWarning') + yield pytest.param(dist, marks=mark) + else: + yield dist + + +def cases_test_fitstart(): + for distname, shapes in dict(distcont).items(): + if (not isinstance(distname, str) or + distname in {'studentized_range', 'recipinvgauss'}): # slow + continue + yield distname, shapes + + +@pytest.mark.parametrize('distname, shapes', list(cases_test_fitstart())) +def test_fitstart(distname, shapes): + dist = getattr(stats, distname) + rng = np.random.default_rng(216342614) + data = rng.random(10) + + with np.errstate(invalid='ignore', divide='ignore'): # irrelevant to test + guess = dist._fitstart(data) + + assert dist._argcheck(*guess[:-2]) + + +def assert_nlff_less_or_close(dist, data, params1, params0, rtol=1e-7, atol=0, + nlff_name='nnlf'): + nlff = getattr(dist, nlff_name) + nlff1 = nlff(params1, data) + nlff0 = nlff(params0, data) + if not (nlff1 < nlff0): + np.testing.assert_allclose(nlff1, nlff0, rtol=rtol, atol=atol) + + +class TestFit: + dist = stats.binom # type: ignore[attr-defined] + seed = 654634816187 + rng = np.random.default_rng(seed) + data = stats.binom.rvs(5, 0.5, size=100, random_state=rng) # type: ignore[attr-defined] # noqa: E501 + shape_bounds_a = [(1, 10), (0, 1)] + shape_bounds_d = {'n': (1, 10), 'p': (0, 1)} + atol = 5e-2 + rtol = 1e-2 + tols = {'atol': atol, 'rtol': rtol} + + def opt(self, *args, rng=1, **kwds): + return differential_evolution(*args, rng=rng, **kwds) + + def test_dist_iv(self): + message = "`dist` must be an instance of..." + with pytest.raises(ValueError, match=message): + stats.fit(10, self.data, self.shape_bounds_a) + + def test_data_iv(self): + message = "`data` must be exactly one-dimensional." + with pytest.raises(ValueError, match=message): + stats.fit(self.dist, [[1, 2, 3]], self.shape_bounds_a) + + message = "All elements of `data` must be finite numbers." + with pytest.raises(ValueError, match=message): + stats.fit(self.dist, [1, 2, 3, np.nan], self.shape_bounds_a) + with pytest.raises(ValueError, match=message): + stats.fit(self.dist, [1, 2, 3, np.inf], self.shape_bounds_a) + with pytest.raises(ValueError, match=message): + stats.fit(self.dist, ['1', '2', '3'], self.shape_bounds_a) + + def test_bounds_iv(self): + message = "Bounds provided for the following unrecognized..." + shape_bounds = {'n': (1, 10), 'p': (0, 1), '1': (0, 10)} + with pytest.warns(RuntimeWarning, match=message): + stats.fit(self.dist, self.data, shape_bounds) + + message = "Each element of a `bounds` sequence must be a tuple..." + shape_bounds = [(1, 10, 3), (0, 1)] + with pytest.raises(ValueError, match=message): + stats.fit(self.dist, self.data, shape_bounds) + + message = "Each element of `bounds` must be a tuple specifying..." + shape_bounds = [(1, 10, 3), (0, 1, 0.5)] + with pytest.raises(ValueError, match=message): + stats.fit(self.dist, self.data, shape_bounds) + shape_bounds = [1, 0] + with pytest.raises(ValueError, match=message): + stats.fit(self.dist, self.data, shape_bounds) + + message = "A `bounds` sequence must contain at least 2 elements..." + shape_bounds = [(1, 10)] + with pytest.raises(ValueError, match=message): + stats.fit(self.dist, self.data, shape_bounds) + + message = "A `bounds` sequence may not contain more than 3 elements..." + bounds = [(1, 10), (1, 10), (1, 10), (1, 10)] + with pytest.raises(ValueError, match=message): + stats.fit(self.dist, self.data, bounds) + + message = "There are no values for `p` on the interval..." + shape_bounds = {'n': (1, 10), 'p': (1, 0)} + with pytest.raises(ValueError, match=message): + stats.fit(self.dist, self.data, shape_bounds) + + message = "There are no values for `n` on the interval..." + shape_bounds = [(10, 1), (0, 1)] + with pytest.raises(ValueError, match=message): + stats.fit(self.dist, self.data, shape_bounds) + + message = "There are no integer values for `n` on the interval..." + shape_bounds = [(1.4, 1.6), (0, 1)] + with pytest.raises(ValueError, match=message): + stats.fit(self.dist, self.data, shape_bounds) + + message = "The intersection of user-provided bounds for `n`" + with pytest.raises(ValueError, match=message): + stats.fit(self.dist, self.data) + shape_bounds = [(-np.inf, np.inf), (0, 1)] + with pytest.raises(ValueError, match=message): + stats.fit(self.dist, self.data, shape_bounds) + + def test_guess_iv(self): + message = "Guesses provided for the following unrecognized..." + guess = {'n': 1, 'p': 0.5, '1': 255} + with pytest.warns(RuntimeWarning, match=message): + stats.fit(self.dist, self.data, self.shape_bounds_d, guess=guess) + + message = "Each element of `guess` must be a scalar..." + guess = {'n': 1, 'p': 'hi'} + with pytest.raises(ValueError, match=message): + stats.fit(self.dist, self.data, self.shape_bounds_d, guess=guess) + guess = [1, 'f'] + with pytest.raises(ValueError, match=message): + stats.fit(self.dist, self.data, self.shape_bounds_d, guess=guess) + guess = [[1, 2]] + with pytest.raises(ValueError, match=message): + stats.fit(self.dist, self.data, self.shape_bounds_d, guess=guess) + + message = "A `guess` sequence must contain at least 2..." + guess = [1] + with pytest.raises(ValueError, match=message): + stats.fit(self.dist, self.data, self.shape_bounds_d, guess=guess) + + message = "A `guess` sequence may not contain more than 3..." + guess = [1, 2, 3, 4] + with pytest.raises(ValueError, match=message): + stats.fit(self.dist, self.data, self.shape_bounds_d, guess=guess) + + message = "Guess for parameter `n` rounded.*|Guess for parameter `p` clipped.*" + guess = {'n': 4.5, 'p': -0.5} + with pytest.warns(RuntimeWarning, match=message): + stats.fit(self.dist, self.data, self.shape_bounds_d, guess=guess) + + message = "Guess for parameter `loc` rounded..." + guess = [5, 0.5, 0.5] + with pytest.warns(RuntimeWarning, match=message): + stats.fit(self.dist, self.data, self.shape_bounds_d, guess=guess) + + message = "Guess for parameter `p` clipped..." + guess = {'n': 5, 'p': -0.5} + with pytest.warns(RuntimeWarning, match=message): + stats.fit(self.dist, self.data, self.shape_bounds_d, guess=guess) + + message = "Guess for parameter `loc` clipped..." + guess = [5, 0.5, 1] + with pytest.warns(RuntimeWarning, match=message): + stats.fit(self.dist, self.data, self.shape_bounds_d, guess=guess) + + def basic_fit_test(self, dist_name, method, rng=1): + + N = 5000 + dist_data = dict(distcont + distdiscrete) + rng = np.random.default_rng(self.seed) + dist = getattr(stats, dist_name) + shapes = np.array(dist_data[dist_name]) + bounds = np.empty((len(shapes) + 2, 2), dtype=np.float64) + bounds[:-2, 0] = shapes/10.**np.sign(shapes) + bounds[:-2, 1] = shapes*10.**np.sign(shapes) + bounds[-2] = (0, 10) + bounds[-1] = (1e-16, 10) + loc = rng.uniform(*bounds[-2]) + scale = rng.uniform(*bounds[-1]) + ref = list(dist_data[dist_name]) + [loc, scale] + + if getattr(dist, 'pmf', False): + ref = ref[:-1] + ref[-1] = np.floor(loc) + data = dist.rvs(*ref, size=N, random_state=rng) + bounds = bounds[:-1] + if getattr(dist, 'pdf', False): + data = dist.rvs(*ref, size=N, random_state=rng) + + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", "overflow encountered", RuntimeWarning) + res = stats.fit(dist, data, bounds, method=method, + optimizer=self.opt) + + nlff_names = {'mle': 'nnlf', 'mse': '_penalized_nlpsf'} + nlff_name = nlff_names[method] + assert_nlff_less_or_close(dist, data, res.params, ref, **self.tols, + nlff_name=nlff_name) + + @pytest.mark.parametrize("dist_name", list(cases_test_fit_mle())) + def test_basic_fit_mle(self, dist_name): + self.basic_fit_test(dist_name, "mle", rng=5) + + @pytest.mark.parametrize("dist_name", list(cases_test_fit_mse())) + def test_basic_fit_mse(self, dist_name): + self.basic_fit_test(dist_name, "mse", rng=2) + + @pytest.mark.slow + def test_arcsine(self): + # Can't guarantee that all distributions will fit all data with + # arbitrary bounds. This distribution just happens to fail above. + # Try something slightly different. + N = 1000 + rng = np.random.default_rng(self.seed) + dist = stats.arcsine + shapes = (1., 2.) + data = dist.rvs(*shapes, size=N, random_state=rng) + shape_bounds = {'loc': (0.1, 10), 'scale': (0.1, 10)} + res = stats.fit(dist, data, shape_bounds, method='mse', optimizer=self.opt) + assert_nlff_less_or_close(dist, data, res.params, shapes, + nlff_name='_penalized_nlpsf', **self.tols) + + @pytest.mark.parametrize("method", ('mle', 'mse')) + def test_argus(self, method): + # Can't guarantee that all distributions will fit all data with + # arbitrary bounds. This distribution just happens to fail above. + # Try something slightly different. + N = 1000 + rng = np.random.default_rng(self.seed) + dist = stats.argus + shapes = (1., 2., 3.) + data = dist.rvs(*shapes, size=N, random_state=rng) + shape_bounds = {'chi': (0.1, 10), 'loc': (0.1, 10), 'scale': (0.1, 10)} + res = stats.fit(dist, data, shape_bounds, optimizer=self.opt, method=method) + nlff_name = {'mle': 'nnlf', 'mse': '_penalized_nlpsf'}[method] + assert_nlff_less_or_close(dist, data, res.params, shapes, **self.tols, + nlff_name=nlff_name) + + @pytest.mark.xslow + def test_beta(self): + # Can't guarantee that all distributions will fit all data with + # arbitrary bounds. This distribution just happens to fail above. + # Try something slightly different. + N = 1000 + rng = np.random.default_rng(self.seed) + dist = stats.beta + shapes = (2.3098496451481823, 0.62687954300963677, 1., 2.) + data = dist.rvs(*shapes, size=N, random_state=rng) + shape_bounds = {'a': (0.1, 10), 'b':(0.1, 10), + 'loc': (0.1, 10), 'scale': (0.1, 10)} + res = stats.fit(dist, data, shape_bounds, method='mle', optimizer=self.opt) + assert_nlff_less_or_close(dist, data, res.params, shapes, + nlff_name='nnlf', **self.tols) + + def test_foldnorm(self): + # Can't guarantee that all distributions will fit all data with + # arbitrary bounds. This distribution just happens to fail above. + # Try something slightly different. + N = 1000 + rng = np.random.default_rng(self.seed) + dist = stats.foldnorm + shapes = (1.952125337355587, 2., 3.) + data = dist.rvs(*shapes, size=N, random_state=rng) + shape_bounds = {'c': (0.1, 10), 'loc': (0.1, 10), 'scale': (0.1, 10)} + res = stats.fit(dist, data, shape_bounds, optimizer=self.opt) + + assert_nlff_less_or_close(dist, data, res.params, shapes, **self.tols) + + def test_nakagami(self): + # Can't guarantee that all distributions will fit all data with + # arbitrary bounds. This distribution just happens to fail above. + # Try something slightly different. + N = 1000 + rng = np.random.default_rng(self.seed) + dist = stats.nakagami + shapes = (4.9673794866666237, 1., 2.) + data = dist.rvs(*shapes, size=N, random_state=rng) + shape_bounds = {'nu':(0.1, 10), 'loc': (0.1, 10), 'scale': (0.1, 10)} + res = stats.fit(dist, data, shape_bounds, method='mle', optimizer=self.opt) + assert_nlff_less_or_close(dist, data, res.params, shapes, + nlff_name='nnlf', **self.tols) + + @pytest.mark.slow + def test_powerlaw(self): + # Can't guarantee that all distributions will fit all data with + # arbitrary bounds. This distribution just happens to fail above. + # Try something slightly different. + N = 1000 + rng = np.random.default_rng(self.seed) + dist = stats.powerlaw + shapes = (1.6591133289905851, 1., 2.) + data = dist.rvs(*shapes, size=N, random_state=rng) + shape_bounds = {'a': (0.1, 10), 'loc': (0.1, 10), 'scale': (0.1, 10)} + res = stats.fit(dist, data, shape_bounds, method='mse', optimizer=self.opt) + assert_nlff_less_or_close(dist, data, res.params, shapes, + nlff_name='_penalized_nlpsf', **self.tols) + + def test_truncpareto(self): + # Can't guarantee that all distributions will fit all data with + # arbitrary bounds. This distribution just happens to fail above. + # Try something slightly different. + N = 1000 + rng = np.random.default_rng(self.seed) + dist = stats.truncpareto + shapes = (1.8, 5.3, 2.3, 4.1) + data = dist.rvs(*shapes, size=N, random_state=rng) + shape_bounds = [(0.1, 10)]*4 + res = stats.fit(dist, data, shape_bounds, optimizer=self.opt) + + assert_nlff_less_or_close(dist, data, res.params, shapes, **self.tols) + + @pytest.mark.slow + def test_truncweibull_min(self): + # Can't guarantee that all distributions will fit all data with + # arbitrary bounds. This distribution just happens to fail above. + # Try something slightly different. + N = 1000 + rng = np.random.default_rng(self.seed) + dist = stats.truncweibull_min + shapes = (2.5, 0.25, 1.75, 2., 3.) + data = dist.rvs(*shapes, size=N, random_state=rng) + shape_bounds = [(0.1, 10)]*5 + res = stats.fit(dist, data, shape_bounds, optimizer=self.opt) + + assert_nlff_less_or_close(dist, data, res.params, shapes, **self.tols) + + def test_missing_shape_bounds(self): + # some distributions have a small domain w.r.t. a parameter, e.g. + # $p \in [0, 1]$ for binomial distribution + # User does not need to provide these because the intersection of the + # user's bounds (none) and the distribution's domain is finite + N = 1000 + rng = np.random.default_rng(self.seed) + + dist = stats.binom + n, p, loc = 10, 0.65, 0 + data = dist.rvs(n, p, loc=loc, size=N, random_state=rng) + shape_bounds = {'n': np.array([0, 20])} # check arrays are OK, too + res = stats.fit(dist, data, shape_bounds, optimizer=self.opt) + assert_allclose(res.params, (n, p, loc), **self.tols) + + dist = stats.bernoulli + p, loc = 0.314159, 0 + data = dist.rvs(p, loc=loc, size=N, random_state=rng) + res = stats.fit(dist, data, optimizer=self.opt) + assert_allclose(res.params, (p, loc), **self.tols) + + def test_fit_only_loc_scale(self): + # fit only loc + N = 5000 + rng = np.random.default_rng(self.seed) + + dist = stats.norm + loc, scale = 1.5, 1 + data = dist.rvs(loc=loc, size=N, random_state=rng) + loc_bounds = (0, 5) + bounds = {'loc': loc_bounds} + res = stats.fit(dist, data, bounds, optimizer=self.opt) + assert_allclose(res.params, (loc, scale), **self.tols) + + # fit only scale + loc, scale = 0, 2.5 + data = dist.rvs(scale=scale, size=N, random_state=rng) + scale_bounds = (0.01, 5) + bounds = {'scale': scale_bounds} + res = stats.fit(dist, data, bounds, optimizer=self.opt) + assert_allclose(res.params, (loc, scale), **self.tols) + + # fit only loc and scale + dist = stats.norm + loc, scale = 1.5, 2.5 + data = dist.rvs(loc=loc, scale=scale, size=N, random_state=rng) + bounds = {'loc': loc_bounds, 'scale': scale_bounds} + res = stats.fit(dist, data, bounds, optimizer=self.opt) + assert_allclose(res.params, (loc, scale), **self.tols) + + def test_everything_fixed(self): + N = 5000 + rng = np.random.default_rng(self.seed) + + dist = stats.norm + loc, scale = 1.5, 2.5 + data = dist.rvs(loc=loc, scale=scale, size=N, random_state=rng) + + # loc, scale fixed to 0, 1 by default + res = stats.fit(dist, data) + assert_allclose(res.params, (0, 1), **self.tols) + + # loc, scale explicitly fixed + bounds = {'loc': (loc, loc), 'scale': (scale, scale)} + res = stats.fit(dist, data, bounds) + assert_allclose(res.params, (loc, scale), **self.tols) + + # `n` gets fixed during polishing + dist = stats.binom + n, p, loc = 10, 0.65, 0 + data = dist.rvs(n, p, loc=loc, size=N, random_state=rng) + shape_bounds = {'n': (0, 20), 'p': (0.65, 0.65)} + res = stats.fit(dist, data, shape_bounds, optimizer=self.opt) + assert_allclose(res.params, (n, p, loc), **self.tols) + + def test_failure(self): + N = 5000 + rng = np.random.default_rng(self.seed) + + dist = stats.nbinom + shapes = (5, 0.5) + data = dist.rvs(*shapes, size=N, random_state=rng) + + assert data.min() == 0 + # With lower bounds on location at 0.5, likelihood is zero + bounds = [(0, 30), (0, 1), (0.5, 10)] + res = stats.fit(dist, data, bounds) + message = "Optimization converged to parameter values that are" + assert res.message.startswith(message) + assert res.success is False + + @pytest.mark.xslow + def test_guess(self): + # Test that guess helps DE find the desired solution + N = 2000 + # With some seeds, `fit` doesn't need a guess + rng = np.random.default_rng(196390444561) + dist = stats.nhypergeom + params = (20, 7, 12, 0) + bounds = [(2, 200), (0.7, 70), (1.2, 120), (0, 10)] + + data = dist.rvs(*params, size=N, random_state=rng) + + res = stats.fit(dist, data, bounds, optimizer=self.opt) + assert not np.allclose(res.params, params, **self.tols) + + res = stats.fit(dist, data, bounds, guess=params, optimizer=self.opt) + assert_allclose(res.params, params, **self.tols) + + def test_mse_accuracy_1(self): + # Test maximum spacing estimation against example from Wikipedia + # https://en.wikipedia.org/wiki/Maximum_spacing_estimation#Examples + data = [2, 4] + dist = stats.expon + bounds = {'loc': (0, 0), 'scale': (1e-8, 10)} + res_mle = stats.fit(dist, data, bounds=bounds, method='mle') + assert_allclose(res_mle.params.scale, 3, atol=1e-3) + res_mse = stats.fit(dist, data, bounds=bounds, method='mse') + assert_allclose(res_mse.params.scale, 3.915, atol=1e-3) + + def test_mse_accuracy_2(self): + # Test maximum spacing estimation against example from Wikipedia + # https://en.wikipedia.org/wiki/Maximum_spacing_estimation#Examples + rng = np.random.default_rng(9843212616816518964) + + dist = stats.uniform + n = 10 + data = dist(3, 6).rvs(size=n, random_state=rng) + bounds = {'loc': (0, 10), 'scale': (1e-8, 10)} + res = stats.fit(dist, data, bounds=bounds, method='mse') + # (loc=3.608118420015416, scale=5.509323262055043) + + x = np.sort(data) + a = (n*x[0] - x[-1])/(n - 1) + b = (n*x[-1] - x[0])/(n - 1) + ref = a, b-a # (3.6081133632151503, 5.509328130317254) + assert_allclose(res.params, ref, rtol=1e-4) + + +# Data from Matlab: https://www.mathworks.com/help/stats/lillietest.html +examgrades = [65, 61, 81, 88, 69, 89, 55, 84, 86, 84, 71, 81, 84, 81, 78, 67, + 96, 66, 73, 75, 59, 71, 69, 63, 79, 76, 63, 85, 87, 88, 80, 71, + 65, 84, 71, 75, 81, 79, 64, 65, 84, 77, 70, 75, 84, 75, 73, 92, + 90, 79, 80, 71, 73, 71, 58, 79, 73, 64, 77, 82, 81, 59, 54, 82, + 57, 79, 79, 73, 74, 82, 63, 64, 73, 69, 87, 68, 81, 73, 83, 73, + 80, 73, 73, 71, 66, 78, 64, 74, 68, 67, 75, 75, 80, 85, 74, 76, + 80, 77, 93, 70, 86, 80, 81, 83, 68, 60, 85, 64, 74, 82, 81, 77, + 66, 85, 75, 81, 69, 60, 83, 72] + + +class TestGoodnessOfFit: + + def test_gof_iv(self): + dist = stats.norm + x = [1, 2, 3] + + message = r"`dist` must be a \(non-frozen\) instance of..." + with pytest.raises(TypeError, match=message): + goodness_of_fit(stats.norm(), x) + + message = "`data` must be a one-dimensional array of numbers." + with pytest.raises(ValueError, match=message): + goodness_of_fit(dist, [[1, 2, 3]]) + + message = "`statistic` must be one of..." + with pytest.raises(ValueError, match=message): + goodness_of_fit(dist, x, statistic='mm') + + message = "`n_mc_samples` must be an integer." + with pytest.raises(TypeError, match=message): + goodness_of_fit(dist, x, n_mc_samples=1000.5) + + message = "SeedSequence expects int or sequence" + with pytest.raises(TypeError, match=message): + goodness_of_fit(dist, x, rng='herring') + + def test_against_ks(self): + rng = np.random.default_rng(8517426291317196949) + x = examgrades + known_params = {'loc': np.mean(x), 'scale': np.std(x, ddof=1)} + res = goodness_of_fit(stats.norm, x, known_params=known_params, + statistic='ks', rng=rng) + ref = stats.kstest(x, stats.norm(**known_params).cdf, method='exact') + assert_allclose(res.statistic, ref.statistic) # ~0.0848 + assert_allclose(res.pvalue, ref.pvalue, atol=5e-3) # ~0.335 + + def test_against_lilliefors(self): + rng = np.random.default_rng(2291803665717442724) + x = examgrades + # preserve use of old random_state during SPEC 7 transition + res = goodness_of_fit(stats.norm, x, statistic='ks', random_state=rng) + known_params = {'loc': np.mean(x), 'scale': np.std(x, ddof=1)} + ref = stats.kstest(x, stats.norm(**known_params).cdf, method='exact') + assert_allclose(res.statistic, ref.statistic) # ~0.0848 + assert_allclose(res.pvalue, 0.0348, atol=5e-3) + + def test_against_cvm(self): + rng = np.random.default_rng(8674330857509546614) + x = examgrades + known_params = {'loc': np.mean(x), 'scale': np.std(x, ddof=1)} + res = goodness_of_fit(stats.norm, x, known_params=known_params, + statistic='cvm', rng=rng) + ref = stats.cramervonmises(x, stats.norm(**known_params).cdf) + assert_allclose(res.statistic, ref.statistic) # ~0.090 + assert_allclose(res.pvalue, ref.pvalue, atol=5e-3) # ~0.636 + + def test_against_anderson_case_0(self): + # "Case 0" is where loc and scale are known [1] + rng = np.random.default_rng(7384539336846690410) + x = np.arange(1, 101) + # loc that produced critical value of statistic found w/ root_scalar + known_params = {'loc': 45.01575354024957, 'scale': 30} + res = goodness_of_fit(stats.norm, x, known_params=known_params, + statistic='ad', rng=rng) + assert_allclose(res.statistic, 2.492) # See [1] Table 1A 1.0 + assert_allclose(res.pvalue, 0.05, atol=5e-3) + + def test_against_anderson_case_1(self): + # "Case 1" is where scale is known and loc is fit [1] + rng = np.random.default_rng(5040212485680146248) + x = np.arange(1, 101) + # scale that produced critical value of statistic found w/ root_scalar + known_params = {'scale': 29.957112639101933} + res = goodness_of_fit(stats.norm, x, known_params=known_params, + statistic='ad', rng=rng) + assert_allclose(res.statistic, 0.908) # See [1] Table 1B 1.1 + assert_allclose(res.pvalue, 0.1, atol=5e-3) + + def test_against_anderson_case_2(self): + # "Case 2" is where loc is known and scale is fit [1] + rng = np.random.default_rng(726693985720914083) + x = np.arange(1, 101) + # loc that produced critical value of statistic found w/ root_scalar + known_params = {'loc': 44.5680212261933} + res = goodness_of_fit(stats.norm, x, known_params=known_params, + statistic='ad', rng=rng) + assert_allclose(res.statistic, 2.904) # See [1] Table 1B 1.2 + assert_allclose(res.pvalue, 0.025, atol=5e-3) + + def test_against_anderson_case_3(self): + # "Case 3" is where both loc and scale are fit [1] + rng = np.random.default_rng(6763691329830218206) + # c that produced critical value of statistic found w/ root_scalar + x = stats.skewnorm.rvs(1.4477847789132101, loc=1, scale=2, size=100, + random_state=rng) + res = goodness_of_fit(stats.norm, x, statistic='ad', rng=rng) + assert_allclose(res.statistic, 0.559) # See [1] Table 1B 1.2 + assert_allclose(res.pvalue, 0.15, atol=5e-3) + + @pytest.mark.xslow + def test_against_anderson_gumbel_r(self): + rng = np.random.default_rng(7302761058217743) + # c that produced critical value of statistic found w/ root_scalar + x = stats.genextreme(0.051896837188595134, loc=0.5, + scale=1.5).rvs(size=1000, random_state=rng) + res = goodness_of_fit(stats.gumbel_r, x, statistic='ad', rng=rng) + ref = stats.anderson(x, dist='gumbel_r', method='interpolate') + assert_allclose(res.statistic, ref.statistic) + assert_allclose(res.pvalue, ref.pvalue, atol=5e-3) + + def test_against_filliben_norm(self): + # Test against `stats.fit` ref. [7] Section 8 "Example" + rng = np.random.default_rng(8024266430745011915) + y = [6, 1, -4, 8, -2, 5, 0] + known_params = {'loc': 0, 'scale': 1} + res = stats.goodness_of_fit(stats.norm, y, known_params=known_params, + statistic="filliben", rng=rng) + # Slight discrepancy presumably due to roundoff in Filliben's + # calculation. Using exact order statistic medians instead of + # Filliben's approximation doesn't account for it. + assert_allclose(res.statistic, 0.98538, atol=1e-4) + assert 0.75 < res.pvalue < 0.9 + + # Using R's ppcc library: + # library(ppcc) + # options(digits=16) + # x < - c(6, 1, -4, 8, -2, 5, 0) + # set.seed(100) + # ppccTest(x, "qnorm", ppos="Filliben") + # Discrepancy with + assert_allclose(res.statistic, 0.98540957187084, rtol=2e-5) + assert_allclose(res.pvalue, 0.8875, rtol=2e-3) + + def test_filliben_property(self): + # Filliben's statistic should be independent of data location and scale + rng = np.random.default_rng(8535677809395478813) + x = rng.normal(loc=10, scale=0.5, size=100) + res = stats.goodness_of_fit(stats.norm, x, + statistic="filliben", rng=rng) + known_params = {'loc': 0, 'scale': 1} + ref = stats.goodness_of_fit(stats.norm, x, known_params=known_params, + statistic="filliben", rng=rng) + assert_allclose(res.statistic, ref.statistic, rtol=1e-15) + + @pytest.mark.parametrize('case', [(25, [.928, .937, .950, .958, .966]), + (50, [.959, .965, .972, .977, .981]), + (95, [.977, .979, .983, .986, .989])]) + def test_against_filliben_norm_table(self, case): + # Test against `stats.fit` ref. [7] Table 1 + rng = np.random.default_rng(504569995557928957) + n, ref = case + x = rng.random(n) + known_params = {'loc': 0, 'scale': 1} + res = stats.goodness_of_fit(stats.norm, x, known_params=known_params, + statistic="filliben", rng=rng) + percentiles = np.array([0.005, 0.01, 0.025, 0.05, 0.1]) + res = stats.scoreatpercentile(res.null_distribution, percentiles*100) + assert_allclose(res, ref, atol=2e-3) + + @pytest.mark.xslow + @pytest.mark.parametrize('case', [(5, 0.95772790260469, 0.4755), + (6, 0.95398832257958, 0.3848), + (7, 0.9432692889277, 0.2328)]) + def test_against_ppcc(self, case): + # Test against R ppcc, e.g. + # library(ppcc) + # options(digits=16) + # x < - c(0.52325412, 1.06907699, -0.36084066, 0.15305959, 0.99093194) + # set.seed(100) + # ppccTest(x, "qrayleigh", ppos="Filliben") + n, ref_statistic, ref_pvalue = case + rng = np.random.default_rng(7777775561439803116) + x = rng.normal(size=n) + res = stats.goodness_of_fit(stats.rayleigh, x, statistic="filliben", + rng=rng) + assert_allclose(res.statistic, ref_statistic, rtol=1e-4) + assert_allclose(res.pvalue, ref_pvalue, atol=1.5e-2) + + def test_params_effects(self): + # Ensure that `guessed_params`, `fit_params`, and `known_params` have + # the intended effects. + rng = np.random.default_rng(9121950977643805391) + x = stats.skewnorm.rvs(-5.044559778383153, loc=1, scale=2, size=50, + random_state=rng) + + # Show that `guessed_params` don't fit to the guess, + # but `fit_params` and `known_params` respect the provided fit + guessed_params = {'c': 13.4} + fit_params = {'scale': 13.73} + known_params = {'loc': -13.85} + rng = np.random.default_rng(9121950977643805391) + res1 = goodness_of_fit(stats.weibull_min, x, n_mc_samples=2, + guessed_params=guessed_params, + fit_params=fit_params, + known_params=known_params, rng=rng) + assert not np.allclose(res1.fit_result.params.c, 13.4) + assert_equal(res1.fit_result.params.scale, 13.73) + assert_equal(res1.fit_result.params.loc, -13.85) + + # Show that changing the guess changes the parameter that gets fit, + # and it changes the null distribution + guessed_params = {'c': 2} + rng = np.random.default_rng(9121950977643805391) + res2 = goodness_of_fit(stats.weibull_min, x, n_mc_samples=2, + guessed_params=guessed_params, + fit_params=fit_params, + known_params=known_params, rng=rng) + assert not np.allclose(res2.fit_result.params.c, + res1.fit_result.params.c, rtol=1e-8) + assert not np.allclose(res2.null_distribution, + res1.null_distribution, rtol=1e-8) + assert_equal(res2.fit_result.params.scale, 13.73) + assert_equal(res2.fit_result.params.loc, -13.85) + + # If we set all parameters as fit_params and known_params, + # they're all fixed to those values, but the null distribution + # varies. + fit_params = {'c': 13.4, 'scale': 13.73} + rng = np.random.default_rng(9121950977643805391) + res3 = goodness_of_fit(stats.weibull_min, x, n_mc_samples=2, + guessed_params=guessed_params, + fit_params=fit_params, + known_params=known_params, rng=rng) + assert_equal(res3.fit_result.params.c, 13.4) + assert_equal(res3.fit_result.params.scale, 13.73) + assert_equal(res3.fit_result.params.loc, -13.85) + assert not np.allclose(res3.null_distribution, res1.null_distribution) + + def test_custom_statistic(self): + # Test support for custom statistic function. + + # References: + # [1] Pyke, R. (1965). "Spacings". Journal of the Royal Statistical + # Society: Series B (Methodological), 27(3): 395-436. + # [2] Burrows, P. M. (1979). "Selected Percentage Points of + # Greenwood's Statistics". Journal of the Royal Statistical + # Society. Series A (General), 142(2): 256-258. + + # Use the Greenwood statistic for illustration; see [1, p.402]. + def greenwood(dist, data, *, axis): + x = np.sort(data, axis=axis) + y = dist.cdf(x) + d = np.diff(y, axis=axis, prepend=0, append=1) + return np.sum(d ** 2, axis=axis) + + # Run the Monte Carlo test with sample size = 5 on a fully specified + # null distribution, and compare the simulated quantiles to the exact + # ones given in [2, Table 1, column (n = 5)]. + rng = np.random.default_rng(9121950977643805391) + data = stats.expon.rvs(size=5, random_state=rng) + result = goodness_of_fit(stats.expon, data, + known_params={'loc': 0, 'scale': 1}, + statistic=greenwood, rng=rng) + p = [.01, .05, .1, .2, .3, .4, .5, .6, .7, .8, .9, .95, .99] + exact_quantiles = [ + .183863, .199403, .210088, .226040, .239947, .253677, .268422, + .285293, .306002, .334447, .382972, .432049, .547468] + simulated_quantiles = np.quantile(result.null_distribution, p) + assert_allclose(simulated_quantiles, exact_quantiles, atol=0.005) + +class TestFitResult: + def test_plot_iv(self): + rng = np.random.default_rng(1769658657308472721) + data = stats.norm.rvs(0, 1, size=100, random_state=rng) + + def optimizer(*args, **kwargs): + return differential_evolution(*args, **kwargs, rng=rng) + + bounds = [(0, 30), (0, 1)] + res = stats.fit(stats.norm, data, bounds, optimizer=optimizer) + try: + import matplotlib # noqa: F401 + message = r"`plot_type` must be one of \{'..." + with pytest.raises(ValueError, match=message): + res.plot(plot_type='llama') + except (ModuleNotFoundError, ImportError): + message = r"matplotlib must be installed to use method `plot`." + with pytest.raises(ModuleNotFoundError, match=message): + res.plot(plot_type='llama') diff --git a/v12x/pylib/scipy/stats/tests/test_generation/reference_distribution_infrastructure_tests.py b/v12x/pylib/scipy/stats/tests/test_generation/reference_distribution_infrastructure_tests.py new file mode 100644 index 0000000000000000000000000000000000000000..9f44d24da8804b8bf9c84685f6892dd11723d933 --- /dev/null +++ b/v12x/pylib/scipy/stats/tests/test_generation/reference_distribution_infrastructure_tests.py @@ -0,0 +1,90 @@ +# Note: this file is to be run locally, not on CI. It is only for tests of the +# reference distribution *infrastructure*; unit tests of SciPy distributions +# do not go here. + +import numpy as np +import pytest + +from scipy import stats +from numpy.testing import assert_allclose +import scipy.stats.tests.test_generation.reference_distributions as rd +import mpmath +from mpmath import mp + +def test_basic(): + + message = "`mpmath.mp.dps <= 15`. Set a higher precision..." + with pytest.raises(RuntimeError, match=message): + rd.Normal() + + mpmath.dps = 20 + message = "`mpmath.dps` has been assigned. This is not intended usage..." + with pytest.raises(RuntimeError, match=message): + rd.Normal() + del mpmath.dps + + mp.dps = 20 # high enough to pass, not unreasonably slow + + # Basic tests of the mpmath distribution infrastructure using a SciPy + # distribution as a reference. The intent is just to make sure that the + # implementations do not have *mistakes* and that broadcasting is working + # as expected. The accuracy is what it is. + + rng = np.random.default_rng(6716188855217730280) + + x = rng.random(size=3) + a = rng.random(size=(2, 1)) + rtol = 1e-15 + + dist = rd.SkewNormal(a=a) + dist_ref = stats.skewnorm(a) + + assert_allclose(dist.pdf(x), dist_ref.pdf(x), rtol=rtol) + assert_allclose(dist.cdf(x), dist_ref.cdf(x), rtol=rtol) + assert_allclose(dist.sf(x), dist_ref.sf(x), rtol=rtol) + assert_allclose(dist.ppf(x), dist_ref.ppf(x), rtol=rtol) + assert_allclose(dist.isf(x), dist_ref.isf(x), rtol=rtol) + assert_allclose(dist.logpdf(x), dist_ref.logpdf(x), rtol=rtol) + assert_allclose(dist.logcdf(x), dist_ref.logcdf(x), rtol=rtol) + assert_allclose(dist.logsf(x), dist_ref.logsf(x), rtol=rtol) + assert_allclose(dist.support(), dist_ref.support(), rtol=rtol) + assert_allclose(dist.entropy(), dist_ref.entropy(), rtol=rtol) + assert_allclose(dist.mean(), dist_ref.mean(), rtol=rtol) + assert_allclose(dist.var(), dist_ref.var(), rtol=rtol) + assert_allclose(dist.skew(), dist_ref.stats('s'), rtol=rtol) + assert_allclose(dist.kurtosis(), dist_ref.stats('k'), rtol=rtol) + + +def test_complementary_method_use(): + # Show that complementary methods are used as expected. + # E.g., use 1 - CDF to compute SF if CDF is overridden but SF is not + + mp.dps = 50 + x = np.linspace(0, 1, 10) + class MyDist(rd.ReferenceDistribution): + def _cdf(self, x): + return x + + dist = MyDist() + assert_allclose(dist.sf(x), 1 - dist.cdf(x)) + + class MyDist(rd.ReferenceDistribution): + def _sf(self, x): + return 1-x + + dist = MyDist() + assert_allclose(dist.cdf(x), 1 - dist.sf(x)) + + class MyDist(rd.ReferenceDistribution): + def _ppf(self, x, guess): + return x + + dist = MyDist() + assert_allclose(dist.isf(x), dist.ppf(1-x)) + + class MyDist(rd.ReferenceDistribution): + def _isf(self, x, guess): + return 1-x + + dist = MyDist() + assert_allclose(dist.ppf(x), dist.isf(1-x)) diff --git a/v12x/pylib/scipy/stats/tests/test_generation/reference_distributions.py b/v12x/pylib/scipy/stats/tests/test_generation/reference_distributions.py new file mode 100644 index 0000000000000000000000000000000000000000..037b16a2d0c88ddf75416fa1324d010d35e0b363 --- /dev/null +++ b/v12x/pylib/scipy/stats/tests/test_generation/reference_distributions.py @@ -0,0 +1,430 @@ +import numpy as np +import mpmath +from mpmath import mp + + +class ReferenceDistribution: + """Minimalist distribution infrastructure for generating reference data. + + The purpose is to generate reference values for unit tests of SciPy + distribution accuracy and robustness. + + Handles array input with standard broadcasting rules, and method + implementations are easily compared against their mathematical definitions. + No attempt is made to handle edge cases or be fast, and arbitrary precision + arithmetic is trusted for accuracy rather than making the method + implementations "smart". + + Notes + ----- + + In this infrastructure, distributions families are classes, and + fully-specified distributions (i.e. with definite values of all family + parameters) are instances of these classes. Typically, the public methods + accept as input only the argument at which the at which the function is to + be evaluated. Unlike SciPy distributions, they never accept values of + distribution family shape, location, or scale parameters. A few + other parameters are noteworthy: + + - All methods accept `dtype` to control the output data type. The default + is `np.float64`, but `object` or `mp.mpf` may be + specified to output the full `mpf`. + - `ppf`/`isf` accept a `guess` because they use a scalar rootfinder + to invert the `cdf`/`sf`. This is passed directly into the `x0` method + of `mpmath.findroot`; see its documentation for details. + - moment accepts `order`, an integer that specifies the order of the (raw) + moment, and `center`, which is the value about which the moment is + taken. The default is to calculate the mean and use it to calculate + central moments; passing ``0`` results in a noncentral moment. For + efficiency, the mean can be passed explicitly if it is already known. + + Follow the example of SkewNormal to generate new reference distributions, + overriding only `__init__` and `_pdf`*. Use the reference distributions to + generate reference values for unit tests of SciPy distribution method + precision and robustness (e.g. for extreme arguments). If the a SciPy + methods implementation is independent and yet the output matches reference + values generated with this infrastructure, it is unlikely that the SciPy + and reference values are both inaccurate. + + * If the SciPy output *doesn't* match and the cause appears to be + inaccuracy of the reference values (e.g. due to numerical issues that + mpmath's arbitrary precision arithmetic doesn't handle), then it may be + appropriate to override a method of the reference distribution rather than + relying on the generic implementation. Otherwise, hesitate to override + methods: the generic implementations are mathematically correct and easy + to verify, whereas an override introduces many possibilities of mistakes, + requires more time to write, and requires more time to review. + + In general, do not create custom unit tests to ensure that + SciPy distribution methods are *correct* (in the sense of being consistent + with the rest of the distribution methods); generic tests take care of + that. + """ + + def __init__(self, **kwargs): + try: + if mpmath.dps is not None: + message = ("`mpmath.dps` has been assigned. This is not " + "intended usage; instead, assign the desired " + "precision to `mpmath.mp.dps` (e.g. `from mpmath " + "as mp; mp.dps = 50.") + raise RuntimeError(message) + except AttributeError: + mpmath.dps = None + + if mp.dps <= 15: + message = ("`mpmath.mp.dps <= 15`. Set a higher precision (e.g." + "`50`) to use this distribution.") + raise RuntimeError(message) + + self._params = {key:self._make_mpf_array(val) + for key, val in kwargs.items()} + + def _make_mpf_array(self, x): + shape = np.shape(x) + x = np.asarray(x, dtype=np.float64).ravel() + return np.asarray([mp.mpf(xi) for xi in x]).reshape(shape)[()] + + def _pdf(self, x): + raise NotImplementedError("_pdf must be overridden.") + + def _cdf(self, x, **kwargs): + if ((self._cdf.__func__ is ReferenceDistribution._cdf) + and (self._sf.__func__ is not ReferenceDistribution._sf)): + return mp.one - self._sf(x, **kwargs) + + a, b = self._support(**kwargs) + res = mp.quad(lambda x: self._pdf(x, **kwargs), (a, x)) + res = res if res < 0.5 else mp.one - self._sf(x, **kwargs) + return res + + def _sf(self, x, **kwargs): + if ((self._sf.__func__ is ReferenceDistribution._sf) + and (self._cdf.__func__ is not ReferenceDistribution._cdf)): + return mp.one - self._cdf(x, **kwargs) + + a, b = self._support(**kwargs) + res = mp.quad(lambda x: self._pdf(x, **kwargs), (x, b)) + res = res if res < 0.5 else mp.one - self._cdf(x, **kwargs) + return res + + def _ppf(self, p, guess=0, **kwargs): + if ((self._ppf.__func__ is ReferenceDistribution._ppf) + and (self._isf.__func__ is not ReferenceDistribution._isf)): + return self._isf(mp.one - p, guess, **kwargs) + def f(x): + return self._cdf(x, **kwargs) - p + return mp.findroot(f, guess) + + def _isf(self, p, guess=0, **kwargs): + if ((self._isf.__func__ is ReferenceDistribution._isf) + and (self._ppf.__func__ is not ReferenceDistribution._ppf)): + return self._ppf(mp.one - p, guess, **kwargs) + def f(x): + return self._sf(x, **kwargs) - p + return mp.findroot(f, guess) + + def _logpdf(self, x, **kwargs): + return mp.log(self._pdf(x, **kwargs)) + + def _logcdf(self, x, **kwargs): + return mp.log(self._cdf(x, **kwargs)) + + def _logsf(self, x, **kwargs): + return mp.log(self._sf(x, **kwargs)) + + def _support(self, **kwargs): + return -mp.inf, mp.inf + + def _entropy(self, **kwargs): + def integrand(x): + logpdf = self._logpdf(x, **kwargs) + pdf = mp.exp(logpdf) + return -pdf*logpdf + + a, b = self._support(**kwargs) + return mp.quad(integrand, (a, b)) + + def _mean(self, **kwargs): + return self._moment(order=1, center=0, **kwargs) + + def _var(self, **kwargs): + mu = self._mean(**kwargs) + return self._moment(order=2, center=mu, **kwargs) + + def _skew(self, **kwargs): + mu = self._mean(**kwargs) + u2 = self._moment(order=2, center=mu, **kwargs) + sigma = mp.sqrt(u2) + u3 = self._moment(order=3, center=mu, **kwargs) + return u3 / sigma**3 + + def _kurtosis(self, **kwargs): + mu = self._mean(**kwargs) + u2 = self._moment(order=2, center=mu, **kwargs) + u4 = self._moment(order=4, center=mu, **kwargs) + return u4 / u2**2 - 3 + + def _moment(self, order, center, **kwargs): + def integrand(x): + return self._pdf(x, **kwargs)*(x - center)**order + + if center is None: + center = self._mean(**kwargs) + + a, b = self._support(**kwargs) + return mp.quad(integrand, (a, b)) + + def pdf(self, x, dtype=np.float64): + fun = np.vectorize(self._pdf) + x = self._make_mpf_array(x) + res = fun(x, **self._params) + return np.asarray(res, dtype=dtype)[()] + + def cdf(self, x, dtype=np.float64): + fun = np.vectorize(self._cdf) + x = self._make_mpf_array(x) + res = fun(x, **self._params) + return np.asarray(res, dtype=dtype)[()] + + def sf(self, x, dtype=np.float64): + fun = np.vectorize(self._sf) + x = self._make_mpf_array(x) + res = fun(x, **self._params) + return np.asarray(res, dtype=dtype)[()] + + def ppf(self, x, guess=0, dtype=np.float64): + fun = np.vectorize(self._ppf, excluded={1}) # don't vectorize guess + x = self._make_mpf_array(x) + res = fun(x, guess, **self._params) + return np.asarray(res, dtype=dtype)[()] + + def isf(self, x, guess=0, dtype=np.float64): + fun = np.vectorize(self._isf, excluded={1}) # don't vectorize guess + x = self._make_mpf_array(x) + res = fun(x, guess, **self._params) + return np.asarray(res, dtype=dtype)[()] + + def logpdf(self, x, dtype=np.float64): + fun = np.vectorize(self._logpdf) + x = self._make_mpf_array(x) + res = fun(x, **self._params) + return np.asarray(res, dtype=dtype)[()] + + def logcdf(self, x, dtype=np.float64): + fun = np.vectorize(self._logcdf) + x = self._make_mpf_array(x) + res = fun(x, **self._params) + return np.asarray(res, dtype=dtype)[()] + + def logsf(self, x, dtype=np.float64): + fun = np.vectorize(self._logsf) + x = self._make_mpf_array(x) + res = fun(x, **self._params) + return np.asarray(res, dtype=dtype)[()] + + def support(self, dtype=np.float64): + fun = np.vectorize(self._support) + res = fun(**self._params) + return np.asarray(res, dtype=dtype)[()] + + def entropy(self, dtype=np.float64): + fun = np.vectorize(self._entropy) + res = fun(**self._params) + return np.asarray(res, dtype=dtype)[()] + + def mean(self, dtype=np.float64): + fun = np.vectorize(self._mean) + res = fun(**self._params) + return np.asarray(res, dtype=dtype)[()] + + def var(self, dtype=np.float64): + fun = np.vectorize(self._var) + res = fun(**self._params) + return np.asarray(res, dtype=dtype)[()] + + def skew(self, dtype=np.float64): + fun = np.vectorize(self._skew) + res = fun(**self._params) + return np.asarray(res, dtype=dtype)[()] + + def kurtosis(self, dtype=np.float64): + fun = np.vectorize(self._kurtosis) + res = fun(**self._params) + return np.asarray(res, dtype=dtype)[()] + + def moment(self, order, center=None, dtype=np.float64): + fun = np.vectorize(self._moment) + order = self._make_mpf_array(order) + res = fun(order, **self._params) + return np.asarray(res, dtype=dtype)[()] + + +class SkewNormal(ReferenceDistribution): + """Reference implementation of the SkewNormal distribution. + + Follow the example here to generate new reference distributions. + Use the reference distributions to generate reference values of + distributions functions. For now, copy-paste the output into unit + tests. Full code to generate reference values does not need to be + included as a comment in the test; just refer to the reference + distribution used and the settings (e.g. mp.dps=50). + """ + + def __init__(self, *, a): + # Overriding __init__ is not necessary, but it allows IDEs to hint at + # shape parameters. All parameters are keyword only to avoid the + # ambiguity inherent in positional arguments. The infrastructure does + # not take care of location and scale; nonetheless, assume standard + # location and scale. Typically, there is no need to test the SciPy + # distribution infrastructure's treatment of location and scale + # separately for a specific distribution. + super().__init__(a=a) + + def _support(self, a): + # Override _support if the support of the distribution is a subset of + # the real line + return -mp.inf, mp.inf + + def _pdf(self, x, a): + # Write PDFs following a scholarly reference as closely as possible. + # Trust mpmath for the accuracy, and don't worry about speed. What's + # important is the ease of verifying the PDF against the reference. If + # the result is inaccurate, it will not match SciPy's output (whether + # SciPy is accurate or not). If this occurs, try increasing dps before + # implementing a numerically favorable (but presumably more complex) + # implementation. + return 2 * mp.npdf(x) * mp.ncdf(a * x) + + # Avoid overriding other methods unless the generic implementation is + # believed to be inaccurate (e.g. due to numerical difficulties) or it is + # too slow. Why? Less code to write, less code to review, and a guarantee + # that there is no *mistake* in the implementation (e.g. wrong formula). + + +class BetaPrime(ReferenceDistribution): + + def __init__(self, *, a, b): + super().__init__(a=a, b=b) + + def _support(self, **kwargs): + return mp.zero, mp.inf + + def _logpdf(self, x, a, b): + return (a - mp.one)*mp.log(x) - (a + b)*mp.log1p(x) - mp.log(mp.beta(a, b)) + + def _pdf(self, x, a, b): + return mp.exp(self._logpdf(x=x, a=a, b=b)) + + def _sf(self, x, a, b): + return 1.0 - mp.betainc(a, b, 0, x/(1+x), regularized=True) + + +class Burr(ReferenceDistribution): + + def __init__(self, *, c, d): + super().__init__(c=c, d=d) + + def _support(self, c, d): + return 0, mp.inf + + def _pdf(self, x, c, d): + return c * d * x ** (-c - 1) * (1 + x ** (-c)) ** (-d - 1) + + def _ppf(self, p, guess, c, d): + return (p**(-1.0/d) - 1)**(-1.0/c) + + +class LogLaplace(ReferenceDistribution): + + def __init__(self, *, c): + super().__init__(c=c) + + def _support(self, c): + return 0, mp.inf + + def _pdf(self, x, c): + if x < mp.one: + return c / 2 * x**(c - mp.one) + else: + return c / 2 * x**(-c - mp.one) + + def _ppf(self, q, guess, c): + if q < 0.5: + return (2.0 * q)**(mp.one / c) + else: + return (2 * (mp.one - q))**(-mp.one / c) + + +class LogNormal(ReferenceDistribution): + + def __init__(self, *, s): + super().__init__(s=s) + + def _support(self, s): + return 0, mp.inf + + def _pdf(self, x, s): + return ( + mp.one / (s * x * mp.sqrt(2 * mp.pi)) + * mp.exp(-mp.one / 2 * (mp.log(x) / s)**2) + ) + + def _cdf(self, x, s): + return mp.ncdf(mp.log(x) / s) + + +class Normal(ReferenceDistribution): + + def _pdf(self, x): + return mp.npdf(x) + + +class NormInvGauss(ReferenceDistribution): + + def __init__(self, *, alpha, beta): + super().__init__(alpha=alpha, beta=beta) + + def _pdf(self, x, alpha, beta): + # Implemented as described in https://www.jstor.org/stable/4616433 + # Equations 2.1 - 2.3 + q = mp.sqrt(1 + x**2) + a = mp.pi**-1 * alpha * mp.exp(mp.sqrt(alpha**2 - beta**2)) + return a * q**-1 * mp.besselk(1, alpha*q) * mp.exp(beta*x) + + +class Pearson3(ReferenceDistribution): + def __init__(self, *, skew): + super().__init__(skew=skew) + + def _pdf(self, x, skew): + b = 2 / skew + a = b**2 + c = -b + res = abs(b)/mp.gamma(a) * (b*(x-c))**(a-1) * mp.exp(-b*(x-c)) + return res if abs(res.real) == res else 0 + + +class StudentT(ReferenceDistribution): + + def __init(self, *, df): + super().__init__(df=df) + + def _pdf(self, x, df): + return (mp.gamma((df + mp.one)/2)/(mp.sqrt(df * mp.pi) * mp.gamma(df/2)) + * (mp.one + x*x/df)**(-(df + mp.one)/2)) + + +class TruncExpon(ReferenceDistribution): + + def __init__(self, *, b): + super().__init__(b=b) + + def _support(self, b): + return 0, b + + def _pdf(self, x, b): + return -mp.exp(-x)/mp.expm1(-b) + + def _sf(self, x, b): + return (mp.exp(-b) - mp.exp(-x))/mp.expm1(-b) diff --git a/v12x/pylib/scipy/stats/tests/test_generation/studentized_range_mpmath_ref.py b/v12x/pylib/scipy/stats/tests/test_generation/studentized_range_mpmath_ref.py new file mode 100644 index 0000000000000000000000000000000000000000..2f67a3453e08561cfab7f2d85b34a58d7e345992 --- /dev/null +++ b/v12x/pylib/scipy/stats/tests/test_generation/studentized_range_mpmath_ref.py @@ -0,0 +1,252 @@ +# To run this script, run +# `python studentized_range_mpmath_ref.py` +# in the "scipy/stats/tests/" directory + +# This script generates a JSON file "./data/studentized_range_mpmath_ref.json" +# that is used to compare the accuracy of `studentized_range` functions against +# precise (20 DOP) results generated using `mpmath`. + +# Equations in this file have been taken from +# https://en.wikipedia.org/wiki/Studentized_range_distribution +# and have been checked against the following reference: +# Lund, R. E., and J. R. Lund. "Algorithm AS 190: Probabilities and +# Upper Quantiles for the Studentized Range." Journal of the Royal +# Statistical Society. Series C (Applied Statistics), vol. 32, no. 2, +# 1983, pp. 204-210. JSTOR, www.jstor.org/stable/2347300. Accessed 18 +# Feb. 2021. + +# Note: I would have preferred to use pickle rather than JSON, but - +# due to security concerns - decided against it. +import itertools +from collections import namedtuple +import json +import time + +import os +from multiprocessing import Pool, cpu_count + +from mpmath import gamma, pi, sqrt, quad, inf, mpf, mp +from mpmath import npdf as phi +from mpmath import ncdf as Phi + +results_filepath = "data/studentized_range_mpmath_ref.json" +num_pools = max(cpu_count() - 1, 1) + +MPResult = namedtuple("MPResult", ["src_case", "mp_result"]) + +CdfCase = namedtuple("CdfCase", + ["q", "k", "v", "expected_atol", "expected_rtol"]) + +MomentCase = namedtuple("MomentCase", + ["m", "k", "v", "expected_atol", "expected_rtol"]) + +# Load previously generated JSON results, or init a new dict if none exist +if os.path.isfile(results_filepath): + res_dict = json.load(open(results_filepath)) +else: + res_dict = dict() + +# Frame out data structure. Store data with the function type as a top level +# key to allow future expansion +res_dict["COMMENT"] = ("!!!!!! THIS FILE WAS AUTOGENERATED BY RUNNING " + "`python studentized_range_mpmath_ref.py` !!!!!!") +res_dict.setdefault("cdf_data", []) +res_dict.setdefault("pdf_data", []) +res_dict.setdefault("moment_data", []) + +general_atol, general_rtol = 1e-11, 1e-11 + +mp.dps = 24 + +cp_q = [0.1, 1, 4, 10] +cp_k = [3, 10, 20] +cp_nu = [3, 10, 20, 50, 100, 120] + +cdf_pdf_cases = [ + CdfCase(*case, + general_atol, + general_rtol) + for case in + itertools.product(cp_q, cp_k, cp_nu) +] + +mom_atol, mom_rtol = 1e-9, 1e-9 +# These are EXTREMELY slow - Multiple days each in worst case. +moment_cases = [ + MomentCase(i, 3, 10, mom_atol, mom_rtol) + for i in range(5) +] + + +def write_data(): + """Writes the current res_dict to the target JSON file""" + with open(results_filepath, mode="w") as f: + json.dump(res_dict, f, indent=2) + + +def to_dict(named_tuple): + """Converts a namedtuple to a dict""" + return dict(named_tuple._asdict()) + + +def mp_res_to_dict(mp_result): + """Formats an MPResult namedtuple into a dict for JSON dumping""" + return { + "src_case": to_dict(mp_result.src_case), + + # np assert can't handle mpf, so take the accuracy hit here. + "mp_result": float(mp_result.mp_result) + } + + +def cdf_mp(q, k, nu): + """Straightforward implementation of studentized range CDF""" + q, k, nu = mpf(q), mpf(k), mpf(nu) + + def inner(s, z): + return phi(z) * (Phi(z + q * s) - Phi(z)) ** (k - 1) + + def outer(s, z): + return s ** (nu - 1) * phi(sqrt(nu) * s) * inner(s, z) + + def whole(s, z): + return (sqrt(2 * pi) * k * nu ** (nu / 2) + / (gamma(nu / 2) * 2 ** (nu / 2 - 1)) * outer(s, z)) + + res = quad(whole, [0, inf], [-inf, inf], + method="gauss-legendre", maxdegree=10) + return res + + +def pdf_mp(q, k, nu): + """Straightforward implementation of studentized range PDF""" + q, k, nu = mpf(q), mpf(k), mpf(nu) + + def inner(s, z): + return phi(z + q * s) * phi(z) * (Phi(z + q * s) - Phi(z)) ** (k - 2) + + def outer(s, z): + return s ** nu * phi(sqrt(nu) * s) * inner(s, z) + + def whole(s, z): + return (sqrt(2 * pi) * k * (k - 1) * nu ** (nu / 2) + / (gamma(nu / 2) * 2 ** (nu / 2 - 1)) * outer(s, z)) + + res = quad(whole, [0, inf], [-inf, inf], + method="gauss-legendre", maxdegree=10) + return res + + +def moment_mp(m, k, nu): + """Implementation of the studentized range moment""" + m, k, nu = mpf(m), mpf(k), mpf(nu) + + def inner(q, s, z): + return phi(z + q * s) * phi(z) * (Phi(z + q * s) - Phi(z)) ** (k - 2) + + def outer(q, s, z): + return s ** nu * phi(sqrt(nu) * s) * inner(q, s, z) + + def pdf(q, s, z): + return (sqrt(2 * pi) * k * (k - 1) * nu ** (nu / 2) + / (gamma(nu / 2) * 2 ** (nu / 2 - 1)) * outer(q, s, z)) + + def whole(q, s, z): + return q ** m * pdf(q, s, z) + + res = quad(whole, [0, inf], [0, inf], [-inf, inf], + method="gauss-legendre", maxdegree=10) + return res + + +def result_exists(set_key, case): + """Searches the results dict for a result in the set that matches a case. + Returns True if such a case exists.""" + if set_key not in res_dict: + raise ValueError(f"{set_key} not present in data structure!") + + case_dict = to_dict(case) + existing_res = list(filter( + lambda res: res["src_case"] == case_dict, # dict comparison + res_dict[set_key])) + + return len(existing_res) > 0 + + +def run(case, run_lambda, set_key, index=0, total_cases=0): + """Runs the single passed case, returning an mp dictionary and index""" + t_start = time.perf_counter() + + res = run_lambda(case) + + print(f"Finished {index + 1}/{total_cases} in batch. " + f"(Took {time.perf_counter() - t_start}s)") + + return index, set_key, mp_res_to_dict(MPResult(case, res)) + + +def write_result(res): + """A callback for completed jobs. Inserts and writes a calculated result + to file.""" + index, set_key, result_dict = res + res_dict[set_key].insert(index, result_dict) + write_data() + + +def run_cases(cases, run_lambda, set_key): + """Runs an array of cases and writes to file""" + # Generate jobs to run from cases that do not have a result in + # the previously loaded JSON. + job_arg = [(case, run_lambda, set_key, index, len(cases)) + for index, case in enumerate(cases) + if not result_exists(set_key, case)] + + print(f"{len(cases) - len(job_arg)}/{len(cases)} cases won't be " + f"calculated because their results already exist.") + + jobs = [] + pool = Pool(num_pools) + + # Run all using multiprocess + for case in job_arg: + jobs.append(pool.apply_async(run, args=case, callback=write_result)) + + pool.close() + pool.join() + + +def run_pdf(case): + return pdf_mp(case.q, case.k, case.v) + + +def run_cdf(case): + return cdf_mp(case.q, case.k, case.v) + + +def run_moment(case): + return moment_mp(case.m, case.k, case.v) + + +def main(): + t_start = time.perf_counter() + + total_cases = 2 * len(cdf_pdf_cases) + len(moment_cases) + print(f"Processing {total_cases} test cases") + + print(f"Running 1st batch ({len(cdf_pdf_cases)} PDF cases). " + f"These take about 30s each.") + run_cases(cdf_pdf_cases, run_pdf, "pdf_data") + + print(f"Running 2nd batch ({len(cdf_pdf_cases)} CDF cases). " + f"These take about 30s each.") + run_cases(cdf_pdf_cases, run_cdf, "cdf_data") + + print(f"Running 3rd batch ({len(moment_cases)} moment cases). " + f"These take about anywhere from a few hours to days each.") + run_cases(moment_cases, run_moment, "moment_data") + + print(f"Test data generated in {time.perf_counter() - t_start}s") + + +if __name__ == "__main__": + main() diff --git a/v12x/pylib/scipy/stats/tests/test_hypotests.py b/v12x/pylib/scipy/stats/tests/test_hypotests.py new file mode 100644 index 0000000000000000000000000000000000000000..a5762185f7b8b4363e469c3974b20106fad7cb0b --- /dev/null +++ b/v12x/pylib/scipy/stats/tests/test_hypotests.py @@ -0,0 +1,2248 @@ +from itertools import product + +import numpy as np +import functools +import pytest +from numpy.testing import assert_, assert_equal, assert_allclose +from pytest import raises as assert_raises + +from scipy import stats, special +from scipy.stats import distributions +from scipy.stats._hypotests import (epps_singleton_2samp, cramervonmises, + _cdf_cvm, cramervonmises_2samp, + _pval_cvm_2samp_exact, barnard_exact, + boschloo_exact) +from scipy.stats._mannwhitneyu import mannwhitneyu, _mwu_state, _MWU +from scipy._lib._testutils import _TestPythranFunc +from scipy._external import array_api_extra as xpx +from scipy._lib._array_api import (make_xp_test_case, xp_default_dtype, is_numpy, + eager_warns, xp_ravel, is_jax) +from scipy._lib._array_api_no_0d import xp_assert_equal, xp_assert_close +from scipy.stats._axis_nan_policy import SmallSampleWarning, too_small_1d_not_omit + +lazy_xp_modules = [stats] +skip_xp_backends = pytest.mark.skip_xp_backends + + +@make_xp_test_case(epps_singleton_2samp) +class TestEppsSingleton: + def test_epps_singleton_iv(self, xp): + # partial input validation test relevant to the addition of JIT support + # TODO: complete input validation tests + x = xp.asarray([-0.35, 2.55, 1.73, 0.73, 0.35]) + y = xp.asarray([-1.15, -0.15, 2.48, 3.25, 3.71]) + t = xp.asarray([0.4, -0.4]) + if is_jax(xp): + w, p = epps_singleton_2samp(x, y, t=t) + xp_assert_equal(w, xp.asarray(xp.nan)) + xp_assert_equal(p, xp.asarray(xp.nan)) + else: + message = "t must contain positive elements only." + with pytest.raises(ValueError, match=message): + epps_singleton_2samp(x, y, t=t) + + @pytest.mark.parametrize('dtype', [None, 'float32', 'float64']) + def test_statistic_1(self, dtype, xp): + # first example in Goerg & Kaiser, also in original paper of + # Epps & Singleton. Note: values do not match exactly, the + # value of the interquartile range varies depending on how + # quantiles are computed + dtype = xp_default_dtype(xp) if dtype is None else getattr(xp, dtype) + x = xp.asarray([-0.35, 2.55, 1.73, 0.73, 0.35, + 2.69, 0.46, -0.94, -0.37, 12.07], dtype=dtype) + y = xp.asarray([-1.15, -0.15, 2.48, 3.25, 3.71, + 4.29, 5.00, 7.74, 8.38, 8.60], dtype=dtype) + w, p = epps_singleton_2samp(x, y) + xp_assert_close(w, xp.asarray(15.14, dtype=dtype), atol=0.03) + xp_assert_close(p, xp.asarray(0.00442, dtype=dtype), atol=0.0001) + + def test_statistic_2(self, xp): + # second example in Goerg & Kaiser, again not a perfect match + x = xp.asarray([0, 1, 2, 2, 2, 2, 3, 3, 3, 3, 4, + 5, 5, 5, 5, 6, 10, 10, 10, 10]) + y = xp.asarray([10, 4, 0, 5, 10, 10, 0, 5, 6, 7, + 10, 3, 1, 7, 0, 8, 1, 5, 8, 10]) + w, p = epps_singleton_2samp(x, y) + xp_assert_close(w, xp.asarray(8.900), atol=1e-3) + xp_assert_close(p, xp.asarray(0.06364), atol=5e-5) + + def test_epps_singleton_array_like(self): # only relevant for NumPy + x, y = np.arange(30), np.arange(28) + + w1, p1 = epps_singleton_2samp(list(x), list(y)) + w2, p2 = epps_singleton_2samp(tuple(x), tuple(y)) + w3, p3 = epps_singleton_2samp(x, y) + + assert_(w1 == w2 == w3) + assert_(p1 == p2 == p3) + + def test_epps_singleton_size(self, xp): + # warns if sample contains fewer than 5 elements + x, y = xp.asarray([1, 2, 3, 4]), xp.arange(10) + with eager_warns(SmallSampleWarning, match=too_small_1d_not_omit, xp=xp): + res = epps_singleton_2samp(x, y) + xp_assert_equal(res.statistic, xp.asarray(xp.nan)) + xp_assert_equal(res.pvalue, xp.asarray(xp.nan)) + + def test_epps_singleton_nonfinite(self, xp): + rng = np.random.default_rng(83249872384543) + x = rng.random(size=(10, 11)) + y = rng.random(size=(10, 12)) + i = np.asarray([1, 4, 9]) # arbitrary rows + + w_ref, p_ref = epps_singleton_2samp(x, y, axis=-1) + w_ref[i] = np.nan + p_ref[i] = np.nan + + x[i[0], 0] = np.nan + x[i[1], 1] = np.inf + y[i[2], 2] = -np.inf + + x, y = xp.asarray(x), xp.asarray(y) + w_res, p_res = epps_singleton_2samp(x, y, axis=-1) + xp_assert_close(w_res, xp.asarray(w_ref)) + xp_assert_close(p_res, xp.asarray(p_ref)) + + +@make_xp_test_case(stats.cramervonmises) +class TestCvm: + # the expected values of the cdfs are taken from Table 1 in + # Csorgo / Faraway: The Exact and Asymptotic Distribution of + # Cramér-von Mises Statistics, 1996. + @pytest.mark.parametrize('n, x, ref', [ + (4, [0.02983, 0.04111, 0.12331, 0.94251], [0.01, 0.05, 0.5, 0.999]), + (10, [0.02657, 0.03830, 0.12068, 0.56643], [0.01, 0.05, 0.5, 0.975]), + (1000, [0.02481, 0.03658, 0.11889, 1.16120], [0.01, 0.05, 0.5, 0.999]), + (None, [0.02480, 0.03656, 0.11888, 1.16204], [0.01, 0.05, 0.5, 0.999]), + ]) + def test_cdf_ref(self, n, x, ref, xp): + xp_assert_close(_cdf_cvm(xp.asarray(x), n), xp.asarray(ref), atol=1e-4) + + def test_cdf_support(self, xp): + # cdf has support on [1/(12*n), n/3] + xp_assert_equal(_cdf_cvm(xp.asarray([1/(12*533), 533/3]), 533), + xp.asarray([0., 1.])) + xp_assert_equal(_cdf_cvm(xp.asarray([1/(12*(27 + 1)), (27 + 1)/3]), 27), + xp.asarray([0., 1.])) + + def test_cdf_large_n(self, xp): + # test that asymptotic cdf and cdf for large samples are close + x = xp.asarray([0.02480, 0.03656, 0.11888, 1.16204, 100]) + xp_assert_close(_cdf_cvm(x, 10000), _cdf_cvm(x), atol=1e-4) + + def test_large_x(self, xp): + # for large values of x and n, the series used to compute the cdf + # converges slowly. + # this leads to bug in R package goftest and MAPLE code that is + # the basis of the implementation in scipy + # note: cdf = 1 for x >= 1000/3 and n = 1000 + x = xp.asarray(333.3, dtype=xp.float64) + assert (0.99999 < _cdf_cvm(x, 1000) < 1.0) + assert (0.99999 < _cdf_cvm(x) < 1.0) + + def test_low_p(self, xp): + # _cdf_cvm can return values larger than 1. In that case, we just + # return a p-value of zero. + n = 12 + res = cramervonmises(xp.ones(n)*0.8, special.ndtr) + assert _cdf_cvm(res.statistic, n, xp=xp) > 1.0 + xp_assert_equal(res.pvalue, xp.asarray(0.)) + + @pytest.mark.parametrize('x', [(), [1.5]]) + def test_invalid_input(self, x, xp): + with pytest.warns(SmallSampleWarning, match=too_small_1d_not_omit): + res = cramervonmises(xp.asarray(x), special.ndtr) + xp_assert_equal(res.statistic, xp.asarray(xp.nan)) + xp_assert_equal(res.pvalue, xp.asarray(xp.nan)) + + @pytest.mark.parametrize('dtype', [None, 'float32', 'float64']) + def test_values_R(self, dtype, xp): + dtype = xp_default_dtype(xp) if dtype is None else getattr(xp, dtype) + # compared against R package goftest, version 1.1.1 + # library(goftest) + # options(digits=16) + # cvm.test(c(-1.7, 2, 0, 1.3, 4, 0.1, 0.6), "pnorm") + res = cramervonmises(xp.asarray([-1.7, 2, 0, 1.3, 4, 0.1, 0.6], dtype=dtype), + special.ndtr) + xp_assert_close(res.statistic, xp.asarray(0.28815604599198, dtype=dtype)) + xp_assert_close(res.pvalue, xp.asarray( 0.1453465252039, dtype=dtype)) + + # cvm.test(c(-1.7, 2, 0, 1.3, 4, 0.1, 0.6), "pnorm", mean = 3, sd = 1.5) + res = cramervonmises(xp.asarray([-1.7, 2, 0, 1.3, 4, 0.1, 0.6], dtype=dtype), + lambda x: special.ndtr((x - 3)/1.5)) + xp_assert_close(res.statistic, xp.asarray(0.94266847977072, dtype=dtype)) + xp_assert_close(res.pvalue, xp.asarray(0.002026416728467, dtype=dtype)) + + # cvm.test(c(1, 2, 5, 1.4, 0.14, 11, 13, 0.9, 7.5), "pexp") + res = cramervonmises( + xp.asarray([1, 2, 5, 1.4, 0.14, 11, 13, 0.9, 7.5], dtype=dtype), + lambda x: -xp.expm1(-x)) + xp_assert_close(res.statistic, xp.asarray(0.84218540922393, dtype=dtype)) + xp_assert_close(res.pvalue, xp.asarray(0.004433406063014, dtype=dtype)) + + def test_str_cdf(self, xp): + if not is_numpy(xp): + message = "`cdf` must be a callable if `rvs` is a non-NumPy array." + with pytest.raises(ValueError, match=message): + cramervonmises(xp.asarray([1, 2, 3]), "beta") + return + + x, args = np.arange(5), (1.4, 0.7) + ref = cramervonmises(x, distributions.expon.cdf) + res = cramervonmises(x, "expon") + assert_equal((res.statistic, res.pvalue), (ref.statistic, ref.pvalue)) + + ref = cramervonmises(x, distributions.beta.cdf, args) + res = cramervonmises(x, "beta", args) + assert_equal((res.statistic, res.pvalue), (ref.statistic, ref.pvalue)) + + +@make_xp_test_case(stats.mannwhitneyu) +class TestMannWhitneyU: + + # All magic numbers are from R wilcox.test unless otherwise specified + # https://rdrr.io/r/stats/wilcox.test.html + + # --- Test Input Validation --- + + def get_zstatistic_from_pvalue(self, p, U, n1, n2, alternative): + if alternative == 'less': + z = special.ndtri(p) + elif alternative == 'greater': + z = -special.ndtri(p) + else: + z = special.ndtri(p/2) if U < n1*n2/2 else -special.ndtri(p/2) + return z + + def test_empty(self, xp): + x = xp.asarray([1, 2]) # generic, valid inputs + y = xp.asarray([3, 4]) + empty = xp.asarray([], dtype=x.dtype) + nan = xp.asarray(xp.nan) + + with pytest.warns(SmallSampleWarning, match=too_small_1d_not_omit): + res = mannwhitneyu(x, empty) + xp_assert_close(res.statistic, nan) + xp_assert_close(res.pvalue, nan) + + with pytest.warns(SmallSampleWarning, match=too_small_1d_not_omit): + res = mannwhitneyu(empty, y) + xp_assert_close(res.statistic, nan) + xp_assert_close(res.pvalue, nan) + + with pytest.warns(SmallSampleWarning, match=too_small_1d_not_omit): + res = mannwhitneyu(empty, empty) + xp_assert_close(res.statistic, nan) + xp_assert_close(res.pvalue, nan) + + def test_input_validation(self, xp): + x = xp.asarray([1, 2]) # generic, valid inputs + y = xp.asarray([3, 4]) + with assert_raises(ValueError, match="`use_continuity` must be one"): + mannwhitneyu(x, y, use_continuity='ekki') + with assert_raises(ValueError, match="`alternative` must be one of"): + mannwhitneyu(x, y, alternative='ekki') + with assert_raises(ValueError, match="`axis` must be an integer"): + mannwhitneyu(x, y, axis=1.5) + with assert_raises(ValueError, match="`method` must be one of"): + mannwhitneyu(x, y, method='ekki') + + def test_auto(self, xp): + # Test that default method ('auto') chooses intended method + if is_jax(xp): + message = "`method='auto'` is incompatible with JAX arrays." + with pytest.raises(ValueError, match=message): + mannwhitneyu(xp.arange(10), xp.arange(10)) + return + + rng = np.random.default_rng(923823782530925934) + n = 8 # threshold to switch from exact to asymptotic + + # both inputs are smaller than threshold; should use exact + x = xp.asarray(rng.random(n-1)) + y = xp.asarray(rng.random(n-1)) + auto = mannwhitneyu(x, y) + asymptotic = mannwhitneyu(x, y, method='asymptotic') + exact = mannwhitneyu(x, y, method='exact') + assert auto.pvalue == exact.pvalue + assert auto.pvalue != asymptotic.pvalue + + # one input is smaller than threshold; should use exact + x = xp.asarray(rng.random(n-1)) + y = xp.asarray(rng.random(n+1)) + auto = mannwhitneyu(x, y) + asymptotic = mannwhitneyu(x, y, method='asymptotic') + exact = mannwhitneyu(x, y, method='exact') + assert auto.pvalue == exact.pvalue + assert auto.pvalue != asymptotic.pvalue + + # other input is smaller than threshold; should use exact + auto = mannwhitneyu(y, x) + asymptotic = mannwhitneyu(x, y, method='asymptotic') + exact = mannwhitneyu(x, y, method='exact') + assert auto.pvalue == exact.pvalue + assert auto.pvalue != asymptotic.pvalue + + # both inputs are larger than threshold; should use asymptotic + x = xp.asarray(rng.random(n+1)) + y = xp.asarray(rng.random(n+1)) + auto = mannwhitneyu(x, y) + asymptotic = mannwhitneyu(x, y, method='asymptotic') + exact = mannwhitneyu(x, y, method='exact') + assert auto.pvalue != exact.pvalue + assert auto.pvalue == asymptotic.pvalue + + # both inputs are smaller than threshold, but there is a tie + # should use asymptotic + x = xp.asarray(rng.random(n-1)) + y = xp.asarray(rng.random(n-1)) + y = xpx.at(y)[3].set(x[3]) + auto = mannwhitneyu(x, y) + asymptotic = mannwhitneyu(x, y, method='asymptotic') + exact = mannwhitneyu(x, y, method='exact') + assert auto.pvalue != exact.pvalue + assert auto.pvalue == asymptotic.pvalue + + # --- Test Basic Functionality --- + + x = [210.052110, 110.190630, 307.918612] + y = [436.08811482466416, 416.37397329768191, 179.96975939463582, + 197.8118754228619, 34.038757281225756, 138.54220550921517, + 128.7769351470246, 265.92721427951852, 275.6617533155341, + 592.34083395416258, 448.73177590617018, 300.61495185038905, + 187.97508449019588] + + # This test was written for mann_whitney_u in gh-4933. + # Originally, the p-values for alternatives were swapped; + # this has been corrected and the tests have been refactored for + # compactness, but otherwise the tests are unchanged. + # R code for comparison, e.g.: + # options(digits = 16) + # x = c(210.052110, 110.190630, 307.918612) + # y = c(436.08811482466416, 416.37397329768191, 179.96975939463582, + # 197.8118754228619, 34.038757281225756, 138.54220550921517, + # 128.7769351470246, 265.92721427951852, 275.6617533155341, + # 592.34083395416258, 448.73177590617018, 300.61495185038905, + # 187.97508449019588) + # wilcox.test(x, y, alternative="g", exact=TRUE) + cases_basic = [[{"alternative": 'two-sided', "method": "asymptotic"}, + (16., 0.6865041817876)], + [{"alternative": 'less', "method": "asymptotic"}, + (16., 0.3432520908938)], + [{"alternative": 'greater', "method": "asymptotic"}, + (16., 0.7047591913255)], + [{"alternative": 'two-sided', "method": "exact"}, + (16., 0.7035714285714)], + [{"alternative": 'less', "method": "exact"}, + (16., 0.3517857142857)], + [{"alternative": 'greater', "method": "exact"}, + (16., 0.6946428571429)]] + + @pytest.mark.parametrize(("kwds", "expected"), cases_basic) + @pytest.mark.parametrize("dtype", [None, 'float32', 'float64']) + def test_basic(self, kwds, expected, dtype, xp): + dtype = xp_default_dtype(xp) if dtype is None else getattr(xp, dtype) + x, y = xp.asarray(self.x, dtype=dtype), xp.asarray(self.y, dtype=dtype) + res = mannwhitneyu(x, y, **kwds) + xp_assert_close(res.statistic, xp.asarray(expected[0], dtype=dtype)) + xp_assert_close(res.pvalue, xp.asarray(expected[1], dtype=dtype)) + if kwds['method'] == 'asymptotic': + z = self.get_zstatistic_from_pvalue(res.pvalue, res.statistic, + x.shape[0], y.shape[0], + alternative=kwds['alternative']) + xp_assert_close(res.zstatistic, xp.asarray(z, dtype=dtype)) + + cases_continuity = [[{"alternative": 'two-sided', "use_continuity": True}, + (23., 0.6865041817876)], + [{"alternative": 'less', "use_continuity": True}, + (23., 0.7047591913255)], + [{"alternative": 'greater', "use_continuity": True}, + (23., 0.3432520908938)], + [{"alternative": 'two-sided', "use_continuity": False}, + (23., 0.6377328900502)], + [{"alternative": 'less', "use_continuity": False}, + (23., 0.6811335549749)], + [{"alternative": 'greater', "use_continuity": False}, + (23., 0.3188664450251)]] + + @pytest.mark.parametrize(("kwds", "expected"), cases_continuity) + def test_continuity(self, kwds, expected, xp): + # When x and y are interchanged, less and greater p-values should + # swap (compare to above). This wouldn't happen if the continuity + # correction were applied in the wrong direction. Note that less and + # greater p-values do not sum to 1 when continuity correction is on, + # which is what we'd expect. Also check that results match R when + # continuity correction is turned off. + # Note that method='asymptotic' -> exact=FALSE + # and use_continuity=False -> correct=FALSE, e.g.: + # wilcox.test(x, y, alternative="t", exact=FALSE, correct=FALSE) + x, y = xp.asarray(self.x), xp.asarray(self.y) + res = mannwhitneyu(y, x, method='asymptotic', **kwds) + xp_assert_close(res.statistic, xp.asarray(expected[0])) + xp_assert_close(res.pvalue, xp.asarray(expected[1])) + z = self.get_zstatistic_from_pvalue(res.pvalue, res.statistic, + x.shape[0], y.shape[0], + alternative=kwds['alternative']) + xp_assert_close(res.zstatistic, xp.asarray(z)) + + def test_tie_correct(self, xp): + # Test tie correction against R's wilcox.test + # options(digits = 16) + # x = c(1, 2, 3, 4) + # y = c(1, 2, 3, 4, 5) + # wilcox.test(x, y, exact=FALSE) + x = xp.asarray([1., 2., 3., 4.]) + y0 = xp.asarray([1., 2., 3., 4., 5.]) + dy = xp.asarray([0., 1., 0., 1., 0.])*0.01 + dy2 = xp.asarray([0., 0., 1., 0., 0.])*0.01 + y = xp.stack([y0-0.01, y0-dy, y0-dy2, y0, y0+dy2, y0+dy, y0+0.01]) + res = mannwhitneyu(x, y, axis=-1, method="asymptotic") + U_expected = [10, 9, 8.5, 8, 7.5, 7, 6] + p_expected = [1, 0.9017048037317, 0.804080657472, 0.7086240584439, + 0.6197963884941, 0.5368784563079, 0.3912672792826] + xp_assert_equal(res.statistic, xp.asarray(U_expected)) + xp_assert_close(res.pvalue, xp.asarray(p_expected)) + + # --- Test Exact Distribution of U --- + + # These are tabulated values of the CDF of the exact distribution of + # the test statistic from pg 52 of reference [1] (Mann-Whitney Original) + pn3 = {1: [0.25, 0.5, 0.75], 2: [0.1, 0.2, 0.4, 0.6], + 3: [0.05, .1, 0.2, 0.35, 0.5, 0.65]} + pn4 = {1: [0.2, 0.4, 0.6], 2: [0.067, 0.133, 0.267, 0.4, 0.6], + 3: [0.028, 0.057, 0.114, 0.2, .314, 0.429, 0.571], + 4: [0.014, 0.029, 0.057, 0.1, 0.171, 0.243, 0.343, 0.443, 0.557]} + pm5 = {1: [0.167, 0.333, 0.5, 0.667], + 2: [0.047, 0.095, 0.19, 0.286, 0.429, 0.571], + 3: [0.018, 0.036, 0.071, 0.125, 0.196, 0.286, 0.393, 0.5, 0.607], + 4: [0.008, 0.016, 0.032, 0.056, 0.095, 0.143, + 0.206, 0.278, 0.365, 0.452, 0.548], + 5: [0.004, 0.008, 0.016, 0.028, 0.048, 0.075, 0.111, + 0.155, 0.21, 0.274, 0.345, .421, 0.5, 0.579]} + pm6 = {1: [0.143, 0.286, 0.428, 0.571], + 2: [0.036, 0.071, 0.143, 0.214, 0.321, 0.429, 0.571], + 3: [0.012, 0.024, 0.048, 0.083, 0.131, + 0.19, 0.274, 0.357, 0.452, 0.548], + 4: [0.005, 0.01, 0.019, 0.033, 0.057, 0.086, 0.129, + 0.176, 0.238, 0.305, 0.381, 0.457, 0.543], # the last element + # of the previous list, 0.543, has been modified from 0.545; + # I assume it was a typo + 5: [0.002, 0.004, 0.009, 0.015, 0.026, 0.041, 0.063, 0.089, + 0.123, 0.165, 0.214, 0.268, 0.331, 0.396, 0.465, 0.535], + 6: [0.001, 0.002, 0.004, 0.008, 0.013, 0.021, 0.032, 0.047, + 0.066, 0.09, 0.12, 0.155, 0.197, 0.242, 0.294, 0.350, + 0.409, 0.469, 0.531]} + + def test_exact_distribution(self): + # I considered parametrize. I decided against it. + setattr(_mwu_state, 's', _MWU(0, 0)) + + p_tables = {3: self.pn3, 4: self.pn4, 5: self.pm5, 6: self.pm6} + for n, table in p_tables.items(): + for m, p in table.items(): + # check p-value against table + u = np.arange(0, len(p)) + _mwu_state.s.set_shapes(m, n) + assert_allclose(_mwu_state.s.cdf(k=u), p, atol=1e-3) + + # check identity CDF + SF - PMF = 1 + # ( In this implementation, SF(U) includes PMF(U) ) + u2 = np.arange(0, m*n+1) + assert_allclose(_mwu_state.s.cdf(k=u2) + + _mwu_state.s.sf(k=u2) + - _mwu_state.s.pmf(k=u2), 1) + + # check symmetry about mean of U, i.e. pmf(U) = pmf(m*n-U) + pmf = _mwu_state.s.pmf(k=u2) + assert_allclose(pmf, pmf[::-1]) + + # check symmetry w.r.t. interchange of m, n + _mwu_state.s.set_shapes(n, m) + pmf2 = _mwu_state.s.pmf(k=u2) + assert_allclose(pmf, pmf2) + + def test_asymptotic_behavior(self, xp): + rng = np.random.default_rng(12543) + + # for small samples, the asymptotic test is not very accurate + x = xp.asarray(rng.random(5)) + y = xp.asarray(rng.random(5)) + res1 = mannwhitneyu(x, y, method="exact") + res2 = mannwhitneyu(x, y, method="asymptotic") + assert res1.statistic == res2.statistic + assert xp.abs(res1.pvalue - res2.pvalue) > 1e-2 + + # for large samples, they agree reasonably well + x = xp.asarray(rng.random(40)) + y = xp.asarray(rng.random(40)) + res1 = mannwhitneyu(x, y, method="exact") + res2 = mannwhitneyu(x, y, method="asymptotic") + assert res1.statistic == res2.statistic + assert xp.abs(res1.pvalue - res2.pvalue) < 1e-3 + + # --- Test Corner Cases --- + + def test_exact_U_equals_mean(self, xp): + # Test U == m*n/2 with exact method + # Without special treatment, two-sided p-value > 1 because both + # one-sided p-values are > 0.5 + x, y = xp.asarray([1., 2., 3.]), xp.asarray([1.5, 2.5]) + res_l = mannwhitneyu(x, y, alternative="less", method="exact") + res_g = mannwhitneyu(x, y, alternative="greater", method="exact") + xp_assert_equal(res_l.pvalue, res_g.pvalue) + assert res_l.pvalue > 0.5 + + res = mannwhitneyu(x, y, alternative="two-sided", method="exact") + xp_assert_equal(res.statistic, xp.asarray(3.)) + xp_assert_equal(res.pvalue, xp.asarray(1.)) + # U == m*n/2 for asymptotic case tested in test_gh_2118 + # The reason it's tricky for the asymptotic test has to do with + # continuity correction. + + cases_scalar = [[{"alternative": 'two-sided', "method": "asymptotic"}, + (0., 1.)], + [{"alternative": 'less', "method": "asymptotic"}, + (0., 0.5)], + [{"alternative": 'greater', "method": "asymptotic"}, + (0., 0.977249868052)], + [{"alternative": 'two-sided', "method": "exact"}, (0., 1)], + [{"alternative": 'less', "method": "exact"}, (0., 0.5)], + [{"alternative": 'greater', "method": "exact"}, (0., 1)]] + + @pytest.mark.parametrize(("kwds", "result"), cases_scalar) + def test_scalar_data(self, kwds, result): # not important to preserve w/ array API + # just making sure scalars work + assert_allclose(mannwhitneyu(1, 2, **kwds), result) + + def test_equal_scalar_data(self): # not important to preserve w/ array API + # when two scalars are equal, there is an -0.5/0 in the asymptotic + # approximation. R gives pvalue=1.0 for alternatives 'less' and + # 'greater' but NA for 'two-sided'. + assert_equal(mannwhitneyu(1, 1, method="exact"), (0.5, 1)) + assert_equal(mannwhitneyu(1, 1, method="asymptotic", alternative='less'), + (0.5, 1)) + assert_equal(mannwhitneyu(1, 1, method="asymptotic", alternative='greater'), + (0.5, 1)) + assert_equal(mannwhitneyu(1, 1, method="asymptotic"), + (0.5, np.nan)) + + # without continuity correction, this definitely becomes 0/0, which really + # is undefined + assert_equal(mannwhitneyu(1, 1, method="asymptotic", alternative='less', + use_continuity=False), (0.5, np.nan)) + assert_equal(mannwhitneyu(1, 1, method="asymptotic", alternative='greater', + use_continuity=False), (0.5, np.nan)) + assert_equal(mannwhitneyu(1, 1, method="asymptotic", + use_continuity=False), (0.5, np.nan)) + + # --- Test Enhancements / Bug Reports --- + + @pytest.mark.parametrize("method", ["asymptotic", "exact"]) + def test_gh_12837_11113(self, method, xp): + # Test that behavior for broadcastable nd arrays is appropriate: + # output shape is correct and all values are equal to when the test + # is performed on one pair of samples at a time. + # Tests that gh-12837 and gh-11113 (requests for n-d input) + # are resolved + rng = np.random.default_rng(6083743794) + + # arrays are broadcastable except for axis = -3 + axis = -3 + m, n = 7, 10 # sample sizes + x = rng.random((m, 3, 8)) + y = rng.random((6, n, 1, 8)) + 0.1 + res = mannwhitneyu(xp.asarray(x), xp.asarray(y), method=method, axis=axis) + + shape = (6, 3, 8) # appropriate shape of outputs, given inputs + assert res.pvalue.shape == shape + assert res.statistic.shape == shape + + # move axis of test to end for simplicity + x, y = np.moveaxis(x, axis, -1), np.moveaxis(y, axis, -1) + + x = x[None, ...] # give x a zeroth dimension + assert x.ndim == y.ndim + + x = np.broadcast_to(x, shape + (m,)) + y = np.broadcast_to(y, shape + (n,)) + assert x.shape[:-1] == shape + assert y.shape[:-1] == shape + + # loop over pairs of samples + statistics = np.zeros(shape) + pvalues = np.zeros(shape) + for indices in product(*[range(i) for i in shape]): + xi = x[indices] + yi = y[indices] + temp = mannwhitneyu(xi, yi, method=method) + statistics[indices] = temp.statistic + pvalues[indices] = temp.pvalue + + xp_assert_close(res.pvalue, xp.asarray(pvalues), atol=1e-16) + xp_assert_close(res.statistic, xp.asarray(statistics), atol=1e-16) + + def test_gh_11355(self, xp): + # Test for correct behavior with NaN/Inf in input + x = [1, 2, 3, 4] + y = [3, 6, 7, 8, 9, 3, 2, 1, 4, 4, 5] + res1 = mannwhitneyu(xp.asarray(x), xp.asarray(y), method='asymptotic') + + # Inf is not a problem. This is a rank test, and it's the largest value + x[0] = 1. # ensure floating point + y[4] = np.inf + res2 = mannwhitneyu(xp.asarray(x), xp.asarray(y), method='asymptotic') + + xp_assert_equal(res1.statistic, res2.statistic) + xp_assert_equal(res1.pvalue, res2.pvalue) + + @skip_xp_backends("jax.numpy", reason="needs jax-ml/jax#34490") + def test_gh11355_nan(self, xp): + # NaNs should propagate by default. + x = [1., 2., 3., 4.] + y = [3, 6, 7, np.nan, 9, 3, 2, 1, 4, 4, 5] + res3 = mannwhitneyu(xp.asarray(x), xp.asarray(y)) + xp_assert_equal(res3.statistic, xp.asarray(xp.nan)) + xp_assert_equal(res3.pvalue, xp.asarray(xp.nan)) + + cases_11355 = [([1., 2, 3, 4], + [3, 6, 7, 8, np.inf, 3, 2, 1, 4, 4, 5], + 10., 0.1297704873477), + ([1., 2, 3, 4], + [3, 6, 7, 8, np.inf, np.inf, 2, 1, 4, 4, 5], + 8.5, 0.08735617507695), + ([1, 2, np.inf, 4], + [3, 6, 7, 8, np.inf, 3, 2, 1, 4, 4, 5], + 17.5, 0.5988856695752), + ([1, 2, np.inf, 4], + [3, 6, 7, 8, np.inf, np.inf, 2, 1, 4, 4, 5], + 16., 0.4687165824462), + ([1, np.inf, np.inf, 4], + [3, 6, 7, 8, np.inf, np.inf, 2, 1, 4, 4, 5], + 24.5, 0.7912517950119)] + + @pytest.mark.parametrize(("x", "y", "statistic", "pvalue"), cases_11355) + def test_gh_11355b(self, x, y, statistic, pvalue, xp): + # Test for correct behavior with NaN/Inf in input + res = mannwhitneyu(xp.asarray(x), xp.asarray(y), method='asymptotic') + xp_assert_close(res.statistic, xp.asarray(statistic), atol=1e-12) + xp_assert_close(res.pvalue, xp.asarray(pvalue), atol=1e-12) + + cases_9184 = [[True, "less", "asymptotic", 0.900775348204], + [True, "greater", "asymptotic", 0.1223118025635], + [True, "two-sided", "asymptotic", 0.244623605127], + [False, "less", "asymptotic", 0.8896643190401], + [False, "greater", "asymptotic", 0.1103356809599], + [False, "two-sided", "asymptotic", 0.2206713619198], + [True, "less", "exact", 0.8967698967699], + [True, "greater", "exact", 0.1272061272061], + [True, "two-sided", "exact", 0.2544122544123]] + + @pytest.mark.parametrize(("use_continuity", "alternative", + "method", "pvalue_exp"), cases_9184) + def test_gh_9184(self, use_continuity, alternative, method, pvalue_exp, xp): + # gh-9184 might be considered a doc-only bug. Please see the + # documentation to confirm that mannwhitneyu correctly notes + # that the output statistic is that of the first sample (x). In any + # case, check the case provided there against output from R. + # R code: + # options(digits=16) + # x <- c(0.80, 0.83, 1.89, 1.04, 1.45, 1.38, 1.91, 1.64, 0.73, 1.46) + # y <- c(1.15, 0.88, 0.90, 0.74, 1.21) + # wilcox.test(x, y, alternative = "less", exact = FALSE) + # wilcox.test(x, y, alternative = "greater", exact = FALSE) + # wilcox.test(x, y, alternative = "two.sided", exact = FALSE) + # wilcox.test(x, y, alternative = "less", exact = FALSE, + # correct=FALSE) + # wilcox.test(x, y, alternative = "greater", exact = FALSE, + # correct=FALSE) + # wilcox.test(x, y, alternative = "two.sided", exact = FALSE, + # correct=FALSE) + # wilcox.test(x, y, alternative = "less", exact = TRUE) + # wilcox.test(x, y, alternative = "greater", exact = TRUE) + # wilcox.test(x, y, alternative = "two.sided", exact = TRUE) + statistic_exp = 35. + x = xp.asarray([0.80, 0.83, 1.89, 1.04, 1.45, 1.38, 1.91, 1.64, 0.73, 1.46]) + y = xp.asarray([1.15, 0.88, 0.90, 0.74, 1.21]) + res = mannwhitneyu(x, y, use_continuity=use_continuity, + alternative=alternative, method=method) + xp_assert_equal(res.statistic, xp.asarray(statistic_exp)) + xp_assert_close(res.pvalue, xp.asarray(pvalue_exp)) + + @skip_xp_backends("jax.numpy", reason="needs jax-ml/jax#34490") + def test_gh_4067(self, xp): + # Test for correct behavior with all NaN input - default is propagate + nan = xp.asarray(xp.nan) + a = xp.stack([nan, nan, nan, nan, nan]) + b = xp.stack([nan, nan, nan, nan, nan]) + res = mannwhitneyu(a, b) + xp_assert_equal(res.statistic, xp.asarray(nan)) + xp_assert_equal(res.pvalue, nan) + + # All cases checked against R wilcox.test, e.g. + # options(digits=16) + # x = c(1, 2, 3) + # y = c(1.5, 2.5) + # wilcox.test(x, y, exact=FALSE, alternative='less') + + cases_2118 = [[[1., 2., 3.], [1.5, 2.5], "greater", (3., 0.6135850036578)], + [[1., 2., 3.], [1.5, 2.5], "less", (3., 0.6135850036578)], + [[1., 2., 3.], [1.5, 2.5], "two-sided", (3., 1.0)], + [[1, 2, 3], [2], "greater", (1.5, 0.681324055883)], + [[1, 2, 3], [2], "less", (1.5, 0.681324055883)], + [[1, 2, 3], [2], "two-sided", (1.5, 1.)], + [[1, 2], [1, 2], "greater", (2., 0.667497228949)], + [[1, 2], [1, 2], "less", (2., 0.667497228949)], + [[1, 2], [1, 2], "two-sided", (2., 1.)]] + + @pytest.mark.parametrize(["x", "y", "alternative", "expected"], cases_2118) + def test_gh_2118(self, x, y, alternative, expected, xp): + # test cases in which U == m*n/2 when method is asymptotic + # applying continuity correction could result in p-value > 1 + res = mannwhitneyu(xp.asarray(x), xp.asarray(y), use_continuity=True, + alternative=alternative, method="asymptotic") + rtol = 1e-6 if xp_default_dtype(xp) == xp.float32 else 1e-12 + xp_assert_close(res.statistic, xp.asarray(expected[0]), rtol=rtol) + xp_assert_close(res.pvalue, xp.asarray(expected[1]), rtol=rtol) + + def test_gh19692_smaller_table(self, xp): + # In gh-19692, we noted that the shape of the cache used in calculating + # p-values was dependent on the order of the inputs because the sample + # sizes n1 and n2 changed. This was indicative of unnecessary cache + # growth and redundant calculation. Check that this is resolved. + rng = np.random.default_rng(7600451795963068007) + m, n = 5, 11 + x = xp.asarray(rng.random(size=m).tolist()) + y = xp.asarray(rng.random(size=n).tolist()) + + setattr(_mwu_state, 's', _MWU(0, 0)) + _mwu_state.s.reset() # reset cache + + res = stats.mannwhitneyu(x, y, method='exact') + shape = _mwu_state.s.configurations.shape + assert shape[-1] == min(res.statistic, m*n - res.statistic) + 1 + stats.mannwhitneyu(y, x, method='exact') + assert shape == _mwu_state.s.configurations.shape # same with reversed sizes + + # Also, we weren't exploiting the symmetry of the null distribution + # to its full potential. Ensure that the null distribution is not + # evaluated explicitly for `k > m*n/2`. + _mwu_state.s.reset() # reset cache + stats.mannwhitneyu(x, 0*y, method='exact', alternative='greater') + shape = _mwu_state.s.configurations.shape + assert shape[-1] == 1 # k is smallest possible + stats.mannwhitneyu(0*x, y, method='exact', alternative='greater') + assert shape == _mwu_state.s.configurations.shape + + @pytest.mark.parametrize('alternative', ['less', 'greater', 'two-sided']) + def test_permutation_method(self, alternative, xp): + rng = np.random.default_rng(7600451795963068007) + x = xp.asarray(rng.random(size=(2, 5))) + y = xp.asarray(rng.random(size=(2, 6))) + res = stats.mannwhitneyu(x, y, method=stats.PermutationMethod(), + alternative=alternative, axis=1) + res2 = stats.mannwhitneyu(x, y, method='exact', + alternative=alternative, axis=1) + xp_assert_close(res.statistic, res2.statistic, rtol=1e-15) + xp_assert_close(res.pvalue, res2.pvalue, rtol=1e-15) + + # Old tests moved from test_stats. Source of magic numbers unknown. + X = [19.8958398126694, 19.5452691647182, 19.0577309166425, 21.716543054589, + 20.3269502208702, 20.0009273294025, 19.3440043632957, 20.4216806548105, + 19.0649894736528, 18.7808043120398, 19.3680942943298, 19.4848044069953, + 20.7514611265663, 19.0894948874598, 19.4975522356628, 18.9971170734274, + 20.3239606288208, 20.6921298083835, 19.0724259532507, 18.9825187935021, + 19.5144462609601, 19.8256857844223, 20.5174677102032, 21.1122407995892, + 17.9490854922535, 18.2847521114727, 20.1072217648826, 18.6439891962179, + 20.4970638083542, 19.5567594734914] + + Y = [19.2790668029091, 16.993808441865, 18.5416338448258, 17.2634018833575, + 19.1577183624616, 18.5119655377495, 18.6068455037221, 18.8358343362655, + 19.0366413269742, 18.1135025515417, 19.2201873866958, 17.8344909022841, + 18.2894380745856, 18.6661374133922, 19.9688601693252, 16.0672254617636, + 19.00596360572, 19.201561539032, 19.0487501090183, 19.0847908674356] + + rtol = 1e-14 + + def test_mannwhitneyu_one_sided(self, xp): + X, Y = xp.asarray(self.X), xp.asarray(self.Y) + kwargs = dict(method='asymptotic') + u1, p1 = stats.mannwhitneyu(X, Y, alternative='less', **kwargs) + u2, p2 = stats.mannwhitneyu(Y, X, alternative='greater', **kwargs) + u3, p3 = stats.mannwhitneyu(X, Y, alternative='greater', **kwargs) + u4, p4 = stats.mannwhitneyu(Y, X, alternative='less', **kwargs) + + xp_assert_equal(p1, p2) + xp_assert_equal(p3, p4) + assert p1 != p3 + xp_assert_equal(u1, xp.asarray(498.)) + xp_assert_equal(u2, xp.asarray(102.)) + xp_assert_equal(u3, xp.asarray(498.)) + xp_assert_equal(u4, xp.asarray(102.)) + assert_allclose(p1, xp.asarray(0.999957683256589), rtol=self.rtol) + rtol = self.rtol if X.dtype == xp.float64 else 5e-4 + assert_allclose(p3, xp.asarray(4.5941632666275e-05), rtol=rtol, atol=1e-16) + + def test_mannwhitneyu_two_sided(self, xp): + X, Y = xp.asarray(self.X), xp.asarray(self.Y) + kwargs = dict(alternative='two-sided', method='asymptotic') + u1, p1 = stats.mannwhitneyu(X, Y, **kwargs) + u2, p2 = stats.mannwhitneyu(Y, X, **kwargs) + + xp_assert_equal(p1, p2) + xp_assert_equal(u1, xp.asarray(498.)) + xp_assert_equal(u2, xp.asarray(102.)) + rtol = self.rtol if X.dtype == xp.float64 else 5e-4 + xp_assert_close(p1, xp.asarray(9.188326533255e-05), rtol=rtol, atol=1e-16) + + def test_mannwhitneyu_no_correct_one_sided(self, xp): + X, Y = xp.asarray(self.X), xp.asarray(self.Y) + kwargs = dict(use_continuity=False, method='asymptotic') + u1, p1 = stats.mannwhitneyu(X, Y, alternative='less', **kwargs) + u2, p2 = stats.mannwhitneyu(Y, X, alternative='greater', **kwargs) + u3, p3 = stats.mannwhitneyu(X, Y, alternative='greater', **kwargs) + u4, p4 = stats.mannwhitneyu(Y, X, alternative='less', **kwargs) + + xp_assert_equal(p1, p2) + xp_assert_equal(p3, p4) + assert p1 != p3 + xp_assert_equal(u1, xp.asarray(498.)) + xp_assert_equal(u2, xp.asarray(102.)) + xp_assert_equal(u3, xp.asarray(498.)) + xp_assert_equal(u4, xp.asarray(102.)) + rtol = self.rtol if X.dtype == xp.float64 else 5e-4 + xp_assert_close(p1, xp.asarray(0.999955905990004), rtol=rtol, atol=1e-16) + xp_assert_close(p3, xp.asarray(4.40940099958089e-05), rtol=rtol, atol=1e-16) + + def test_mannwhitneyu_no_correct_two_sided(self, xp): + X, Y = xp.asarray(self.X), xp.asarray(self.Y) + kwargs = dict(use_continuity=False, method='asymptotic') + u1, p1 = stats.mannwhitneyu(X, Y, **kwargs) + u2, p2 = stats.mannwhitneyu(Y, X, **kwargs) + + xp_assert_equal(p1, p2) + xp_assert_equal(u1, xp.asarray(498.)) + xp_assert_equal(u2, xp.asarray(102.)) + rtol = self.rtol if X.dtype == xp.float64 else 5e-4 + xp_assert_close(p1, xp.asarray(8.81880199916178e-05), rtol=rtol, atol=1e-16) + + def test_mannwhitneyu_ones(self, xp): + # test for gh-1428 + x = xp.asarray([1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., + 1., 1., 1., 1., 1., 1., 1., 2., 1., 1., 1., 1., 1., 1., + 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., + 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., + 1., 1., 1., 1., 1., 1., 1., 2., 1., 1., 1., 1., 1., 1., + 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., + 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 2., + 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., + 1., 1., 2., 1., 1., 1., 1., 2., 1., 1., 2., 1., 1., 2., + 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., + 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 2., 1., + 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., + 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., + 1., 1., 1., 1., 1., 1., 1., 2., 1., 1., 1., 1., 1., 1., + 1., 1., 1., 2., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., + 1., 1., 1., 1., 1., 1., 1., 1., 3., 1., 1., 1., 1., 1., + 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., + 1., 1., 1., 1., 1., 1.]) + + y = xp.asarray([1., 1., 1., 1., 1., 1., 1., 2., 1., 2., 1., 1., 1., 1., + 2., 1., 1., 1., 2., 1., 1., 1., 1., 1., 2., 1., 1., 3., + 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 2., 1., 2., 1., + 1., 1., 1., 1., 1., 2., 1., 1., 1., 1., 1., 1., 1., 1., + 1., 1., 1., 1., 1., 1., 1., 2., 1., 1., 1., 1., 1., 2., + 2., 1., 1., 2., 1., 1., 2., 1., 2., 1., 1., 1., 1., 2., + 2., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., + 1., 2., 1., 1., 1., 1., 1., 2., 2., 2., 1., 1., 1., 1., + 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., + 2., 1., 1., 2., 1., 1., 1., 1., 2., 1., 1., 1., 1., 1., + 1., 1., 1., 1., 1., 1., 1., 2., 1., 1., 1., 2., 1., 1., + 1., 1., 1., 1.]) + + # p-value from R, e.g. wilcox.test(x, y, alternative="g") + kwargs = dict(method='asymptotic') + res = stats.mannwhitneyu(x, y, alternative='less', **kwargs) + xp_assert_close(res.statistic, xp.asarray(16980.5)) + xp_assert_close(res.pvalue, xp.asarray(2.8214327656317373e-5)) + res = stats.mannwhitneyu(x, y, alternative='greater', **kwargs) + xp_assert_close(res.statistic, xp.asarray(16980.5)) + xp_assert_close(res.pvalue, xp.asarray(0.9999719954296)) + res = stats.mannwhitneyu(x, y, alternative='two-sided', **kwargs) + xp_assert_close(res.statistic, xp.asarray(16980.5)) + xp_assert_close(res.pvalue, xp.asarray(5.642865531266e-5)) + + +class TestSomersD(_TestPythranFunc): + def setup_method(self): + self.dtypes = self.ALL_INTEGER + self.ALL_FLOAT + self.arguments = {0: (np.arange(10), + self.ALL_INTEGER + self.ALL_FLOAT), + 1: (np.arange(10), + self.ALL_INTEGER + self.ALL_FLOAT)} + input_array = [self.arguments[idx][0] for idx in self.arguments] + # In this case, self.partialfunc can simply be stats.somersd, + # since `alternative` is an optional argument. If it is required, + # we can use functools.partial to freeze the value, because + # we only mainly test various array inputs, not str, etc. + self.partialfunc = functools.partial(stats.somersd, + alternative='two-sided') + self.expected = self.partialfunc(*input_array) + + def pythranfunc(self, *args): + res = self.partialfunc(*args) + assert_allclose(res.statistic, self.expected.statistic, atol=1e-15) + assert_allclose(res.pvalue, self.expected.pvalue, atol=1e-15) + + def test_pythranfunc_keywords(self): + # Not specifying the optional keyword args + table = [[27, 25, 14, 7, 0], [7, 14, 18, 35, 12], [1, 3, 2, 7, 17]] + res1 = stats.somersd(table) + # Specifying the optional keyword args with default value + optional_args = self.get_optional_args(stats.somersd) + res2 = stats.somersd(table, **optional_args) + # Check if the results are the same in two cases + assert_allclose(res1.statistic, res2.statistic, atol=1e-15) + assert_allclose(res1.pvalue, res2.pvalue, atol=1e-15) + + def test_like_kendalltau(self): + # All tests correspond with one in test_stats.py `test_kendalltau` + + # case without ties, con-dis equal zero + x = [5, 2, 1, 3, 6, 4, 7, 8] + y = [5, 2, 6, 3, 1, 8, 7, 4] + # Cross-check with result from SAS FREQ: + expected = (0.000000000000000, 1.000000000000000) + res = stats.somersd(x, y) + assert_allclose(res.statistic, expected[0], atol=1e-15) + assert_allclose(res.pvalue, expected[1], atol=1e-15) + + # case without ties, con-dis equal zero + x = [0, 5, 2, 1, 3, 6, 4, 7, 8] + y = [5, 2, 0, 6, 3, 1, 8, 7, 4] + # Cross-check with result from SAS FREQ: + expected = (0.000000000000000, 1.000000000000000) + res = stats.somersd(x, y) + assert_allclose(res.statistic, expected[0], atol=1e-15) + assert_allclose(res.pvalue, expected[1], atol=1e-15) + + # case without ties, con-dis close to zero + x = [5, 2, 1, 3, 6, 4, 7] + y = [5, 2, 6, 3, 1, 7, 4] + # Cross-check with result from SAS FREQ: + expected = (-0.142857142857140, 0.630326953157670) + res = stats.somersd(x, y) + assert_allclose(res.statistic, expected[0], atol=1e-15) + assert_allclose(res.pvalue, expected[1], atol=1e-15) + + # simple case without ties + x = np.arange(10) + y = np.arange(10) + # Cross-check with result from SAS FREQ: + # SAS p value is not provided. + expected = (1.000000000000000, 0) + res = stats.somersd(x, y) + assert_allclose(res.statistic, expected[0], atol=1e-15) + assert_allclose(res.pvalue, expected[1], atol=1e-15) + + # swap a couple values and a couple more + x = np.arange(10) + y = np.array([0, 2, 1, 3, 4, 6, 5, 7, 8, 9]) + # Cross-check with result from SAS FREQ: + expected = (0.911111111111110, 0.000000000000000) + res = stats.somersd(x, y) + assert_allclose(res.statistic, expected[0], atol=1e-15) + assert_allclose(res.pvalue, expected[1], atol=1e-15) + + # same in opposite direction + x = np.arange(10) + y = np.arange(10)[::-1] + # Cross-check with result from SAS FREQ: + # SAS p value is not provided. + expected = (-1.000000000000000, 0) + res = stats.somersd(x, y) + assert_allclose(res.statistic, expected[0], atol=1e-15) + assert_allclose(res.pvalue, expected[1], atol=1e-15) + + # swap a couple values and a couple more + x = np.arange(10) + y = np.array([9, 7, 8, 6, 5, 3, 4, 2, 1, 0]) + # Cross-check with result from SAS FREQ: + expected = (-0.9111111111111111, 0.000000000000000) + res = stats.somersd(x, y) + assert_allclose(res.statistic, expected[0], atol=1e-15) + assert_allclose(res.pvalue, expected[1], atol=1e-15) + + # with some ties + x1 = [12, 2, 1, 12, 2] + x2 = [1, 4, 7, 1, 0] + # Cross-check with result from SAS FREQ: + expected = (-0.500000000000000, 0.304901788178780) + res = stats.somersd(x1, x2) + assert_allclose(res.statistic, expected[0], atol=1e-15) + assert_allclose(res.pvalue, expected[1], atol=1e-15) + + # with only ties in one or both inputs + # SAS will not produce an output for these: + # NOTE: No statistics are computed for x * y because x has fewer + # than 2 nonmissing levels. + # WARNING: No OUTPUT data set is produced for this table because a + # row or column variable has fewer than 2 nonmissing levels and no + # statistics are computed. + + res = stats.somersd([2, 2, 2], [2, 2, 2]) + assert_allclose(res.statistic, np.nan) + assert_allclose(res.pvalue, np.nan) + + res = stats.somersd([2, 0, 2], [2, 2, 2]) + assert_allclose(res.statistic, np.nan) + assert_allclose(res.pvalue, np.nan) + + res = stats.somersd([2, 2, 2], [2, 0, 2]) + assert_allclose(res.statistic, np.nan) + assert_allclose(res.pvalue, np.nan) + + res = stats.somersd([0], [0]) + assert_allclose(res.statistic, np.nan) + assert_allclose(res.pvalue, np.nan) + + # empty arrays provided as input + res = stats.somersd([], []) + assert_allclose(res.statistic, np.nan) + assert_allclose(res.pvalue, np.nan) + + # test unequal length inputs + x = np.arange(10.) + y = np.arange(20.) + assert_raises(ValueError, stats.somersd, x, y) + + def test_asymmetry(self): + # test that somersd is asymmetric w.r.t. input order and that + # convention is as described: first input is row variable & independent + # data is from Wikipedia: + # https://en.wikipedia.org/wiki/Somers%27_D + # but currently that example contradicts itself - it says X is + # independent yet take D_XY + + x = [1, 1, 1, 2, 2, 2, 2, 2, 3, 3, 1, 2, + 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3] + y = [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, + 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2] + # Cross-check with result from SAS FREQ: + d_cr = 0.272727272727270 + d_rc = 0.342857142857140 + p = 0.092891940883700 # same p-value for either direction + res = stats.somersd(x, y) + assert_allclose(res.statistic, d_cr, atol=1e-15) + assert_allclose(res.pvalue, p, atol=1e-4) + assert_equal(res.table.shape, (3, 2)) + res = stats.somersd(y, x) + assert_allclose(res.statistic, d_rc, atol=1e-15) + assert_allclose(res.pvalue, p, atol=1e-15) + assert_equal(res.table.shape, (2, 3)) + + def test_somers_original(self): + # test against Somers' original paper [1] + + # Table 5A + # Somers' convention was column IV + table = np.array([[8, 2], [6, 5], [3, 4], [1, 3], [2, 3]]) + # Our convention (and that of SAS FREQ) is row IV + table = table.T + dyx = 129/340 + assert_allclose(stats.somersd(table).statistic, dyx) + + # table 7A - d_yx = 1 + table = np.array([[25, 0], [85, 0], [0, 30]]) + dxy, dyx = 3300/5425, 3300/3300 + assert_allclose(stats.somersd(table).statistic, dxy) + assert_allclose(stats.somersd(table.T).statistic, dyx) + + # table 7B - d_yx < 0 + table = np.array([[25, 0], [0, 30], [85, 0]]) + dyx = -1800/3300 + assert_allclose(stats.somersd(table.T).statistic, dyx) + + def test_contingency_table_with_zero_rows_cols(self): + # test that zero rows/cols in contingency table don't affect result + + N = 100 + shape = 4, 6 + size = np.prod(shape) + + rng = np.random.RandomState(0) + s = stats.multinomial.rvs(N, p=np.ones(size)/size, + random_state=rng).reshape(shape) + res = stats.somersd(s) + + s2 = np.insert(s, 2, np.zeros(shape[1]), axis=0) + res2 = stats.somersd(s2) + + s3 = np.insert(s, 2, np.zeros(shape[0]), axis=1) + res3 = stats.somersd(s3) + + s4 = np.insert(s2, 2, np.zeros(shape[0]+1), axis=1) + res4 = stats.somersd(s4) + + # Cross-check with result from SAS FREQ: + assert_allclose(res.statistic, -0.116981132075470, atol=1e-15) + assert_allclose(res.statistic, res2.statistic) + assert_allclose(res.statistic, res3.statistic) + assert_allclose(res.statistic, res4.statistic) + + assert_allclose(res.pvalue, 0.156376448188150, atol=1e-15) + assert_allclose(res.pvalue, res2.pvalue) + assert_allclose(res.pvalue, res3.pvalue) + assert_allclose(res.pvalue, res4.pvalue) + + def test_invalid_contingency_tables(self): + N = 100 + shape = 4, 6 + size = np.prod(shape) + + rng = np.random.default_rng(0) + # start with a valid contingency table + s = stats.multinomial.rvs(N, p=np.ones(size)/size, + random_state=rng).reshape(shape) + + s5 = s - 2 + message = "All elements of the contingency table must be non-negative" + with assert_raises(ValueError, match=message): + stats.somersd(s5) + + s6 = s + 0.01 + message = "All elements of the contingency table must be integer" + with assert_raises(ValueError, match=message): + stats.somersd(s6) + + message = ("At least two elements of the contingency " + "table must be nonzero.") + with assert_raises(ValueError, match=message): + stats.somersd([[]]) + + with assert_raises(ValueError, match=message): + stats.somersd([[1]]) + + s7 = np.zeros((3, 3)) + with assert_raises(ValueError, match=message): + stats.somersd(s7) + + s7[0, 1] = 1 + with assert_raises(ValueError, match=message): + stats.somersd(s7) + + def test_only_ranks_matter(self): + # only ranks of input data should matter + x = [1, 2, 3] + x2 = [-1, 2.1, np.inf] + y = [3, 2, 1] + y2 = [0, -0.5, -np.inf] + res = stats.somersd(x, y) + res2 = stats.somersd(x2, y2) + assert_equal(res.statistic, res2.statistic) + assert_equal(res.pvalue, res2.pvalue) + + def test_contingency_table_return(self): + # check that contingency table is returned + x = np.arange(10) + y = np.arange(10) + res = stats.somersd(x, y) + assert_equal(res.table, np.eye(10)) + + def test_somersd_alternative(self): + # Test alternative parameter, asymptotic method (due to tie) + + # Based on scipy.stats.test_stats.TestCorrSpearman2::test_alternative + x1 = [1, 2, 3, 4, 5] + x2 = [5, 6, 7, 8, 7] + + # strong positive correlation + expected = stats.somersd(x1, x2, alternative="two-sided") + assert expected.statistic > 0 + + # rank correlation > 0 -> large "less" p-value + res = stats.somersd(x1, x2, alternative="less") + assert_equal(res.statistic, expected.statistic) + assert_allclose(res.pvalue, 1 - (expected.pvalue / 2)) + + # rank correlation > 0 -> small "greater" p-value + res = stats.somersd(x1, x2, alternative="greater") + assert_equal(res.statistic, expected.statistic) + assert_allclose(res.pvalue, expected.pvalue / 2) + + # reverse the direction of rank correlation + x2.reverse() + + # strong negative correlation + expected = stats.somersd(x1, x2, alternative="two-sided") + assert expected.statistic < 0 + + # rank correlation < 0 -> large "greater" p-value + res = stats.somersd(x1, x2, alternative="greater") + assert_equal(res.statistic, expected.statistic) + assert_allclose(res.pvalue, 1 - (expected.pvalue / 2)) + + # rank correlation < 0 -> small "less" p-value + res = stats.somersd(x1, x2, alternative="less") + assert_equal(res.statistic, expected.statistic) + assert_allclose(res.pvalue, expected.pvalue / 2) + + with pytest.raises(ValueError, match="`alternative` must be..."): + stats.somersd(x1, x2, alternative="ekki-ekki") + + @pytest.mark.parametrize("positive_correlation", (False, True)) + def test_somersd_perfect_correlation(self, positive_correlation): + # Before the addition of `alternative`, perfect correlation was + # treated as a special case. Now it is treated like any other case, but + # make sure there are no divide by zero warnings or associated errors + + x1 = np.arange(10) + x2 = x1 if positive_correlation else np.flip(x1) + expected_statistic = 1 if positive_correlation else -1 + + # perfect correlation -> small "two-sided" p-value (0) + res = stats.somersd(x1, x2, alternative="two-sided") + assert res.statistic == expected_statistic + assert res.pvalue == 0 + + # rank correlation > 0 -> large "less" p-value (1) + res = stats.somersd(x1, x2, alternative="less") + assert res.statistic == expected_statistic + assert res.pvalue == (1 if positive_correlation else 0) + + # rank correlation > 0 -> small "greater" p-value (0) + res = stats.somersd(x1, x2, alternative="greater") + assert res.statistic == expected_statistic + assert res.pvalue == (0 if positive_correlation else 1) + + def test_somersd_large_inputs_gh18132(self): + # Test that large inputs where potential overflows could occur give + # the expected output. This is tested in the case of binary inputs. + # See gh-18126. + + # generate lists of random classes 1-2 (binary) + classes = [1, 2] + n_samples = 10 ** 6 + rng = np.random.default_rng(6889320191) + x = rng.choice(classes, n_samples) + y = rng.choice(classes, n_samples) + + # get value to compare with: sklearn output + # from sklearn import metrics + # val_auc_sklearn = metrics.roc_auc_score(x, y) + # # convert to the Gini coefficient (Gini = (AUC*2)-1) + # val_sklearn = 2 * val_auc_sklearn - 1 + val_sklearn = 0.000624401938730923 + + # calculate the Somers' D statistic, which should be equal to the + # result of val_sklearn until approximately machine precision + val_scipy = stats.somersd(x, y).statistic + assert_allclose(val_sklearn, val_scipy, atol=1e-15) + + +class TestBarnardExact: + """Some tests to show that barnard_exact() works correctly.""" + + @pytest.mark.parametrize( + "input_sample,expected", + [ + ([[43, 40], [10, 39]], (3.555406779643, 0.000362832367)), + ([[100, 2], [1000, 5]], (-1.776382925679, 0.135126970878)), + ([[2, 7], [8, 2]], (-2.518474945157, 0.019210815430)), + ([[5, 1], [10, 10]], (1.449486150679, 0.156277546306)), + ([[5, 15], [20, 20]], (-1.851640199545, 0.066363501421)), + ([[5, 16], [20, 25]], (-1.609639949352, 0.116984852192)), + ([[10, 5], [10, 1]], (-1.449486150679, 0.177536588915)), + ([[5, 0], [1, 4]], (2.581988897472, 0.013671875000)), + ([[0, 1], [3, 2]], (-1.095445115010, 0.509667991877)), + ([[0, 2], [6, 4]], (-1.549193338483, 0.197019618792)), + ([[2, 7], [8, 2]], (-2.518474945157, 0.019210815430)), + ], + ) + def test_precise(self, input_sample, expected): + """The expected values have been generated by R, using a resolution + for the nuisance parameter of 1e-6 : + ```R + library(Barnard) + options(digits=10) + barnard.test(43, 40, 10, 39, dp=1e-6, pooled=TRUE) + ``` + """ + res = barnard_exact(input_sample) + statistic, pvalue = res.statistic, res.pvalue + assert_allclose([statistic, pvalue], expected) + + @pytest.mark.parametrize( + "input_sample,expected", + [ + ([[43, 40], [10, 39]], (3.920362887717, 0.000289470662)), + ([[100, 2], [1000, 5]], (-1.139432816087, 0.950272080594)), + ([[2, 7], [8, 2]], (-3.079373904042, 0.020172119141)), + ([[5, 1], [10, 10]], (1.622375939458, 0.150599922226)), + ([[5, 15], [20, 20]], (-1.974771239528, 0.063038448651)), + ([[5, 16], [20, 25]], (-1.722122973346, 0.133329494287)), + ([[10, 5], [10, 1]], (-1.765469659009, 0.250566655215)), + ([[5, 0], [1, 4]], (5.477225575052, 0.007812500000)), + ([[0, 1], [3, 2]], (-1.224744871392, 0.509667991877)), + ([[0, 2], [6, 4]], (-1.732050807569, 0.197019618792)), + ([[2, 7], [8, 2]], (-3.079373904042, 0.020172119141)), + ], + ) + def test_pooled_param(self, input_sample, expected): + """The expected values have been generated by R, using a resolution + for the nuisance parameter of 1e-6 : + ```R + library(Barnard) + options(digits=10) + barnard.test(43, 40, 10, 39, dp=1e-6, pooled=FALSE) + ``` + """ + res = barnard_exact(input_sample, pooled=False) + statistic, pvalue = res.statistic, res.pvalue + assert_allclose([statistic, pvalue], expected) + + def test_raises(self): + # test we raise an error for wrong input number of nuisances. + error_msg = ( + "Number of points `n` must be strictly positive, found 0" + ) + with assert_raises(ValueError, match=error_msg): + barnard_exact([[1, 2], [3, 4]], n=0) + + # test we raise an error for wrong shape of input. + error_msg = "The input `table` must be of shape \\(2, 2\\)." + with assert_raises(ValueError, match=error_msg): + barnard_exact(np.arange(6).reshape(2, 3)) + + # Test all values must be positives + error_msg = "All values in `table` must be nonnegative." + with assert_raises(ValueError, match=error_msg): + barnard_exact([[-1, 2], [3, 4]]) + + # Test value error on wrong alternative param + error_msg = ( + "`alternative` should be one of {'two-sided', 'less', 'greater'}," + " found .*" + ) + with assert_raises(ValueError, match=error_msg): + barnard_exact([[1, 2], [3, 4]], "not-correct") + + @pytest.mark.parametrize( + "input_sample,expected", + [ + ([[0, 0], [4, 3]], (1.0, 0)), + ], + ) + def test_edge_cases(self, input_sample, expected): + res = barnard_exact(input_sample) + statistic, pvalue = res.statistic, res.pvalue + assert_equal(pvalue, expected[0]) + assert_equal(statistic, expected[1]) + + @pytest.mark.parametrize( + "input_sample,expected", + [ + ([[0, 5], [0, 10]], (1.0, np.nan)), + ([[5, 0], [10, 0]], (1.0, np.nan)), + ], + ) + def test_row_or_col_zero(self, input_sample, expected): + res = barnard_exact(input_sample) + statistic, pvalue = res.statistic, res.pvalue + assert_equal(pvalue, expected[0]) + assert_equal(statistic, expected[1]) + + @pytest.mark.parametrize( + "input_sample,expected", + [ + ([[2, 7], [8, 2]], (-2.518474945157, 0.009886140845)), + ([[7, 200], [300, 8]], (-21.320036698460, 0.0)), + ([[21, 28], [1957, 6]], (-30.489638143953, 0.0)), + ], + ) + @pytest.mark.parametrize("alternative", ["greater", "less"]) + def test_less_greater(self, input_sample, expected, alternative): + """ + "The expected values have been generated by R, using a resolution + for the nuisance parameter of 1e-6 : + ```R + library(Barnard) + options(digits=10) + a = barnard.test(2, 7, 8, 2, dp=1e-6, pooled=TRUE) + a$p.value[1] + ``` + In this test, we are using the "one-sided" return value `a$p.value[1]` + to test our pvalue. + """ + expected_stat, less_pvalue_expect = expected + + if alternative == "greater": + input_sample = np.array(input_sample)[:, ::-1] + expected_stat = -expected_stat + + res = barnard_exact(input_sample, alternative=alternative) + statistic, pvalue = res.statistic, res.pvalue + assert_allclose( + [statistic, pvalue], [expected_stat, less_pvalue_expect], atol=1e-7 + ) + + +class TestBoschlooExact: + """Some tests to show that boschloo_exact() works correctly.""" + + ATOL = 1e-7 + + @pytest.mark.parametrize( + "input_sample,expected", + [ + ([[2, 7], [8, 2]], (0.01852173, 0.009886142)), + ([[5, 1], [10, 10]], (0.9782609, 0.9450994)), + ([[5, 16], [20, 25]], (0.08913823, 0.05827348)), + ([[10, 5], [10, 1]], (0.1652174, 0.08565611)), + ([[5, 0], [1, 4]], (1, 1)), + ([[0, 1], [3, 2]], (0.5, 0.34375)), + ([[2, 7], [8, 2]], (0.01852173, 0.009886142)), + ([[7, 12], [8, 3]], (0.06406797, 0.03410916)), + ([[10, 24], [25, 37]], (0.2009359, 0.1512882)), + ], + ) + def test_less(self, input_sample, expected): + """The expected values have been generated by R, using a resolution + for the nuisance parameter of 1e-8 : + ```R + library(Exact) + options(digits=10) + data <- matrix(c(43, 10, 40, 39), 2, 2, byrow=TRUE) + a = exact.test(data, method="Boschloo", alternative="less", + tsmethod="central", np.interval=TRUE, beta=1e-8) + ``` + """ + res = boschloo_exact(input_sample, alternative="less") + statistic, pvalue = res.statistic, res.pvalue + assert_allclose([statistic, pvalue], expected, atol=self.ATOL) + + @pytest.mark.parametrize( + "input_sample,expected", + [ + ([[43, 40], [10, 39]], (0.0002875544, 0.0001615562)), + ([[2, 7], [8, 2]], (0.9990149, 0.9918327)), + ([[5, 1], [10, 10]], (0.1652174, 0.09008534)), + ([[5, 15], [20, 20]], (0.9849087, 0.9706997)), + ([[5, 16], [20, 25]], (0.972349, 0.9524124)), + ([[5, 0], [1, 4]], (0.02380952, 0.006865367)), + ([[0, 1], [3, 2]], (1, 1)), + ([[0, 2], [6, 4]], (1, 1)), + ([[2, 7], [8, 2]], (0.9990149, 0.9918327)), + ([[7, 12], [8, 3]], (0.9895302, 0.9771215)), + ([[10, 24], [25, 37]], (0.9012936, 0.8633275)), + ], + ) + def test_greater(self, input_sample, expected): + """The expected values have been generated by R, using a resolution + for the nuisance parameter of 1e-8 : + ```R + library(Exact) + options(digits=10) + data <- matrix(c(43, 10, 40, 39), 2, 2, byrow=TRUE) + a = exact.test(data, method="Boschloo", alternative="greater", + tsmethod="central", np.interval=TRUE, beta=1e-8) + ``` + """ + res = boschloo_exact(input_sample, alternative="greater") + statistic, pvalue = res.statistic, res.pvalue + assert_allclose([statistic, pvalue], expected, atol=self.ATOL) + + @pytest.mark.parametrize( + "input_sample,expected", + [ + ([[43, 40], [10, 39]], (0.0002875544, 0.0003231115)), + ([[2, 7], [8, 2]], (0.01852173, 0.01977228)), + ([[5, 1], [10, 10]], (0.1652174, 0.1801707)), + ([[5, 16], [20, 25]], (0.08913823, 0.116547)), + ([[5, 0], [1, 4]], (0.02380952, 0.01373073)), + ([[0, 1], [3, 2]], (0.5, 0.6875)), + ([[2, 7], [8, 2]], (0.01852173, 0.01977228)), + ([[7, 12], [8, 3]], (0.06406797, 0.06821831)), + ], + ) + def test_two_sided(self, input_sample, expected): + """The expected values have been generated by R, using a resolution + for the nuisance parameter of 1e-8 : + ```R + library(Exact) + options(digits=10) + data <- matrix(c(43, 10, 40, 39), 2, 2, byrow=TRUE) + a = exact.test(data, method="Boschloo", alternative="two.sided", + tsmethod="central", np.interval=TRUE, beta=1e-8) + ``` + """ + res = boschloo_exact(input_sample, alternative="two-sided", n=64) + # Need n = 64 for python 32-bit + statistic, pvalue = res.statistic, res.pvalue + assert_allclose([statistic, pvalue], expected, atol=self.ATOL) + + def test_raises(self): + # test we raise an error for wrong input number of nuisances. + error_msg = ( + "Number of points `n` must be strictly positive, found 0" + ) + with assert_raises(ValueError, match=error_msg): + boschloo_exact([[1, 2], [3, 4]], n=0) + + # test we raise an error for wrong shape of input. + error_msg = "The input `table` must be of shape \\(2, 2\\)." + with assert_raises(ValueError, match=error_msg): + boschloo_exact(np.arange(6).reshape(2, 3)) + + # Test all values must be positives + error_msg = "All values in `table` must be nonnegative." + with assert_raises(ValueError, match=error_msg): + boschloo_exact([[-1, 2], [3, 4]]) + + # Test value error on wrong alternative param + error_msg = ( + r"`alternative` should be one of \('two-sided', 'less', " + r"'greater'\), found .*" + ) + with assert_raises(ValueError, match=error_msg): + boschloo_exact([[1, 2], [3, 4]], "not-correct") + + @pytest.mark.parametrize( + "input_sample,expected", + [ + ([[0, 5], [0, 10]], (np.nan, np.nan)), + ([[5, 0], [10, 0]], (np.nan, np.nan)), + ], + ) + def test_row_or_col_zero(self, input_sample, expected): + res = boschloo_exact(input_sample) + statistic, pvalue = res.statistic, res.pvalue + assert_equal(pvalue, expected[0]) + assert_equal(statistic, expected[1]) + + def test_two_sided_gt_1(self): + # Check that returned p-value does not exceed 1 even when twice + # the minimum of the one-sided p-values does. See gh-15345. + tbl = [[1, 1], [13, 12]] + pl = boschloo_exact(tbl, alternative='less').pvalue + pg = boschloo_exact(tbl, alternative='greater').pvalue + assert 2*min(pl, pg) > 1 + pt = boschloo_exact(tbl, alternative='two-sided').pvalue + assert pt == 1.0 + + @pytest.mark.parametrize("alternative", ("less", "greater")) + def test_against_fisher_exact(self, alternative): + # Check that the statistic of `boschloo_exact` is the same as the + # p-value of `fisher_exact` (for one-sided tests). See gh-15345. + tbl = [[2, 7], [8, 2]] + boschloo_stat = boschloo_exact(tbl, alternative=alternative).statistic + fisher_p = stats.fisher_exact(tbl, alternative=alternative)[1] + assert_allclose(boschloo_stat, fisher_p) + + +@make_xp_test_case(cramervonmises_2samp) +class TestCvm_2samp: + @pytest.mark.parametrize('args', [([], np.arange(5)), + (np.arange(5), [1])]) + def test_too_small_input(self, args, xp): + args = (xp.asarray(arg, dtype=xp_default_dtype(xp)) for arg in args) + with eager_warns(SmallSampleWarning, match=too_small_1d_not_omit, xp=xp): + res = cramervonmises_2samp(*args) + xp_assert_equal(res.statistic, xp.asarray(xp.nan)) + xp_assert_equal(res.pvalue, xp.asarray(xp.nan)) + + def test_invalid_input(self, xp): + y = xp.arange(5) + msg = 'method must be either auto, exact or asymptotic' + with pytest.raises(ValueError, match=msg): + cramervonmises_2samp(y, y, 'xyz') + + def test_list_input(self): # list input only relevant for NumPy + x = [2, 3, 4, 7, 6] + y = [0.2, 0.7, 12, 18] + r1 = cramervonmises_2samp(x, y) + r2 = cramervonmises_2samp(np.array(x), np.array(y)) + assert_equal((r1.statistic, r1.pvalue), (r2.statistic, r2.pvalue)) + + @pytest.mark.parametrize('dtype', [None, 'float32', 'float64']) + def test_example_conover(self, dtype, xp): + # Example 2 in Section 6.2 of W.J. Conover: Practical Nonparametric + # Statistics, 1971. + dtype = xp_default_dtype(xp) if dtype is None else getattr(xp, dtype) + x = xp.asarray([7.6, 8.4, 8.6, 8.7, 9.3, 9.9, 10.1, 10.6, 11.2], dtype=dtype) + y = xp.asarray([5.2, 5.7, 5.9, 6.5, 6.8, 8.2, 9.1, 9.8, + 10.8, 11.3, 11.5, 12.3, 12.5, 13.4, 14.6], dtype=dtype) + r = cramervonmises_2samp(x, y) + xp_assert_close(r.statistic, xp.asarray(0.262, dtype=dtype), atol=1e-3) + xp_assert_close(r.pvalue, xp.asarray(.18, dtype=dtype), atol=1e-2) + + @pytest.mark.parametrize('statistic, m, n, pval', + [(710, 5, 6, 48./462), + (1897, 7, 7, 117./1716), + (576, 4, 6, 2./210), + (1764, 6, 7, 2./1716)]) + def test_exact_pvalue(self, statistic, m, n, pval): # only implemented w/ NumPy + # the exact values are taken from Anderson: On the distribution of the + # two-sample Cramer-von-Mises criterion, 1962. + # The values are taken from Table 2, 3, 4 and 5 + assert_equal(_pval_cvm_2samp_exact(statistic, m, n), pval) + + @pytest.mark.xslow + def test_large_sample(self, xp): + # for large samples, the statistic U gets very large + # do a sanity check that p-value is not 0, 1 or nan + rng = np.random.default_rng(4367) + x = distributions.norm.rvs(size=1000000, random_state=rng) + y = distributions.norm.rvs(size=900000, random_state=rng) + x, y = xp.asarray(x), xp.asarray(y) + r = cramervonmises_2samp(x, y) + assert 0 < r.pvalue < 1 + r = cramervonmises_2samp(x, y+0.1) + assert 0 < r.pvalue < 1 + + def test_exact_vs_asymptotic(self, xp): + rng = np.random.RandomState(0) + x = xp.asarray(rng.random(7)) + y = xp.asarray(rng.random(8)) + r1 = cramervonmises_2samp(x, y, method='exact') + r2 = cramervonmises_2samp(x, y, method='asymptotic') + xp_assert_equal(r1.statistic, r2.statistic) + xp_assert_close(r1.pvalue, r2.pvalue, atol=1e-2) + + def test_method_auto(self, xp): + x = xp.arange(20.) + y = xp.asarray([0.5, 4.7, 13.1]) + r1 = cramervonmises_2samp(x, y, method='exact') + r2 = cramervonmises_2samp(x, y, method='auto') + xp_assert_equal(r1.pvalue, r2.pvalue) + # switch to asymptotic if one sample has more than 20 observations + x = xp.arange(21.) + r1 = cramervonmises_2samp(x, y, method='asymptotic') + r2 = cramervonmises_2samp(x, y, method='auto') + xp_assert_equal(r1.pvalue, r2.pvalue) + + def test_same_input(self, xp): + # make sure trivial edge case can be handled + # note that _cdf_cvm_inf(0) = nan. implementation avoids nan by + # returning pvalue=1 for very small values of the statistic + x = xp.arange(15) + res = cramervonmises_2samp(x, x) + xp_assert_equal(res.statistic, xp.asarray(0.)) + xp_assert_equal(res.pvalue, xp.asarray(1.)) + # check exact p-value + res = cramervonmises_2samp(x[:4], x[:4]) + xp_assert_equal(res.statistic, xp.asarray(0.)) + xp_assert_equal(res.pvalue, xp.asarray(1.)) + + +class TestTukeyHSD: + + data_same_size = ([24.5, 23.5, 26.4, 27.1, 29.9], + [28.4, 34.2, 29.5, 32.2, 30.1], + [26.1, 28.3, 24.3, 26.2, 27.8]) + data_diff_size = ([24.5, 23.5, 26.28, 26.4, 27.1, 29.9, 30.1, 30.1], + [28.4, 34.2, 29.5, 32.2, 30.1], + [26.1, 28.3, 24.3, 26.2, 27.8]) + extreme_size = ([24.5, 23.5, 26.4], + [28.4, 34.2, 29.5, 32.2, 30.1, 28.4, 34.2, 29.5, 32.2, + 30.1], + [26.1, 28.3, 24.3, 26.2, 27.8]) + + sas_same_size = """ + Comparison LowerCL Difference UpperCL Significance + 2 - 3 0.6908830568 4.34 7.989116943 1 + 2 - 1 0.9508830568 4.6 8.249116943 1 + 3 - 2 -7.989116943 -4.34 -0.6908830568 1 + 3 - 1 -3.389116943 0.26 3.909116943 0 + 1 - 2 -8.249116943 -4.6 -0.9508830568 1 + 1 - 3 -3.909116943 -0.26 3.389116943 0 + """ + + sas_diff_size = """ + Comparison LowerCL Difference UpperCL Significance + 2 - 1 0.2679292645 3.645 7.022070736 1 + 2 - 3 0.5934764007 4.34 8.086523599 1 + 1 - 2 -7.022070736 -3.645 -0.2679292645 1 + 1 - 3 -2.682070736 0.695 4.072070736 0 + 3 - 2 -8.086523599 -4.34 -0.5934764007 1 + 3 - 1 -4.072070736 -0.695 2.682070736 0 + """ + + sas_extreme = """ + Comparison LowerCL Difference UpperCL Significance + 2 - 3 1.561605075 4.34 7.118394925 1 + 2 - 1 2.740784879 6.08 9.419215121 1 + 3 - 2 -7.118394925 -4.34 -1.561605075 1 + 3 - 1 -1.964526566 1.74 5.444526566 0 + 1 - 2 -9.419215121 -6.08 -2.740784879 1 + 1 - 3 -5.444526566 -1.74 1.964526566 0 + """ + + @pytest.mark.parametrize("data,res_expect_str,atol", + ((data_same_size, sas_same_size, 1e-4), + (data_diff_size, sas_diff_size, 1e-4), + (extreme_size, sas_extreme, 1e-10), + ), + ids=["equal size sample", + "unequal sample size", + "extreme sample size differences"]) + def test_compare_sas(self, data, res_expect_str, atol): + ''' + SAS code used to generate results for each sample: + DATA ACHE; + INPUT BRAND RELIEF; + CARDS; + 1 24.5 + ... + 3 27.8 + ; + ods graphics on; ODS RTF;ODS LISTING CLOSE; + PROC ANOVA DATA=ACHE; + CLASS BRAND; + MODEL RELIEF=BRAND; + MEANS BRAND/TUKEY CLDIFF; + TITLE 'COMPARE RELIEF ACROSS MEDICINES - ANOVA EXAMPLE'; + ods output CLDiffs =tc; + proc print data=tc; + format LowerCL 17.16 UpperCL 17.16 Difference 17.16; + title "Output with many digits"; + RUN; + QUIT; + ODS RTF close; + ODS LISTING; + ''' + res_expect = np.asarray(res_expect_str.replace(" - ", " ").split()[5:], + dtype=float).reshape((6, 6)) + res_tukey = stats.tukey_hsd(*data) + conf = res_tukey.confidence_interval() + # loop over the comparisons + for i, j, l, s, h, sig in res_expect: + i, j = int(i) - 1, int(j) - 1 + assert_allclose(conf.low[i, j], l, atol=atol) + assert_allclose(res_tukey.statistic[i, j], s, atol=atol) + assert_allclose(conf.high[i, j], h, atol=atol) + assert_allclose((res_tukey.pvalue[i, j] <= .05), sig == 1) + + matlab_sm_siz = """ + 1 2 -8.2491590248597 -4.6 -0.9508409751403 0.0144483269098 + 1 3 -3.9091590248597 -0.26 3.3891590248597 0.9803107240900 + 2 3 0.6908409751403 4.34 7.9891590248597 0.0203311368795 + """ + + matlab_diff_sz = """ + 1 2 -7.02207069748501 -3.645 -0.26792930251500 0.03371498443080 + 1 3 -2.68207069748500 0.695 4.07207069748500 0.85572267328807 + 2 3 0.59347644287720 4.34 8.08652355712281 0.02259047020620 + """ + + @pytest.mark.parametrize("data,res_expect_str,atol", + ((data_same_size, matlab_sm_siz, 1e-12), + (data_diff_size, matlab_diff_sz, 1e-7)), + ids=["equal size sample", + "unequal size sample"]) + def test_compare_matlab(self, data, res_expect_str, atol): + """ + vals = [24.5, 23.5, 26.4, 27.1, 29.9, 28.4, 34.2, 29.5, 32.2, 30.1, + 26.1, 28.3, 24.3, 26.2, 27.8] + names = {'zero', 'zero', 'zero', 'zero', 'zero', 'one', 'one', 'one', + 'one', 'one', 'two', 'two', 'two', 'two', 'two'} + [p,t,stats] = anova1(vals,names,"off"); + [c,m,h,nms] = multcompare(stats, "CType","hsd"); + """ + res_expect = np.asarray(res_expect_str.split(), + dtype=float).reshape((3, 6)) + res_tukey = stats.tukey_hsd(*data) + conf = res_tukey.confidence_interval() + # loop over the comparisons + for i, j, l, s, h, p in res_expect: + i, j = int(i) - 1, int(j) - 1 + assert_allclose(conf.low[i, j], l, atol=atol) + assert_allclose(res_tukey.statistic[i, j], s, atol=atol) + assert_allclose(conf.high[i, j], h, atol=atol) + assert_allclose(res_tukey.pvalue[i, j], p, atol=atol) + + def test_compare_r(self): + """ + Testing against results and p-values from R: + from: https://www.rdocumentation.org/packages/stats/versions/3.6.2/ + topics/TukeyHSD + > require(graphics) + > summary(fm1 <- aov(breaks ~ tension, data = warpbreaks)) + > TukeyHSD(fm1, "tension", ordered = TRUE) + > plot(TukeyHSD(fm1, "tension")) + Tukey multiple comparisons of means + 95% family-wise confidence level + factor levels have been ordered + Fit: aov(formula = breaks ~ tension, data = warpbreaks) + $tension + """ + str_res = """ + diff lwr upr p adj + 2 - 3 4.722222 -4.8376022 14.28205 0.4630831 + 1 - 3 14.722222 5.1623978 24.28205 0.0014315 + 1 - 2 10.000000 0.4401756 19.55982 0.0384598 + """ + res_expect = np.asarray(str_res.replace(" - ", " ").split()[5:], + dtype=float).reshape((3, 6)) + data = ([26, 30, 54, 25, 70, 52, 51, 26, 67, + 27, 14, 29, 19, 29, 31, 41, 20, 44], + [18, 21, 29, 17, 12, 18, 35, 30, 36, + 42, 26, 19, 16, 39, 28, 21, 39, 29], + [36, 21, 24, 18, 10, 43, 28, 15, 26, + 20, 21, 24, 17, 13, 15, 15, 16, 28]) + + res_tukey = stats.tukey_hsd(*data) + conf = res_tukey.confidence_interval() + # loop over the comparisons + for i, j, s, l, h, p in res_expect: + i, j = int(i) - 1, int(j) - 1 + # atols are set to the number of digits present in the r result. + assert_allclose(conf.low[i, j], l, atol=1e-7) + assert_allclose(res_tukey.statistic[i, j], s, atol=1e-6) + assert_allclose(conf.high[i, j], h, atol=1e-5) + assert_allclose(res_tukey.pvalue[i, j], p, atol=1e-7) + + def test_engineering_stat_handbook(self): + ''' + Example sourced from: + https://www.itl.nist.gov/div898/handbook/prc/section4/prc471.htm + ''' + group1 = [6.9, 5.4, 5.8, 4.6, 4.0] + group2 = [8.3, 6.8, 7.8, 9.2, 6.5] + group3 = [8.0, 10.5, 8.1, 6.9, 9.3] + group4 = [5.8, 3.8, 6.1, 5.6, 6.2] + res = stats.tukey_hsd(group1, group2, group3, group4) + conf = res.confidence_interval() + lower = np.asarray([ + [0, 0, 0, -2.25], + [.29, 0, -2.93, .13], + [1.13, 0, 0, .97], + [0, 0, 0, 0]]) + upper = np.asarray([ + [0, 0, 0, 1.93], + [4.47, 0, 1.25, 4.31], + [5.31, 0, 0, 5.15], + [0, 0, 0, 0]]) + + for (i, j) in [(1, 0), (2, 0), (0, 3), (1, 2), (2, 3)]: + assert_allclose(conf.low[i, j], lower[i, j], atol=1e-2) + assert_allclose(conf.high[i, j], upper[i, j], atol=1e-2) + + def test_rand_symm(self): + # test some expected identities of the results + rng = np.random.default_rng(2699550179) + data = rng.random((3, 100)) + res = stats.tukey_hsd(*data) + conf = res.confidence_interval() + # the confidence intervals should be negated symmetric of each other + assert_equal(conf.low, -conf.high.T) + # the `high` and `low` center diagonals should be the same since the + # mean difference in a self comparison is 0. + assert_equal(np.diagonal(conf.high), conf.high[0, 0]) + assert_equal(np.diagonal(conf.low), conf.low[0, 0]) + # statistic array should be antisymmetric with zeros on the diagonal + assert_equal(res.statistic, -res.statistic.T) + assert_equal(np.diagonal(res.statistic), 0) + # p-values should be symmetric and 1 when compared to itself + assert_equal(res.pvalue, res.pvalue.T) + assert_equal(np.diagonal(res.pvalue), 1) + + def test_no_inf(self): + with assert_raises(ValueError, match="...must be finite."): + stats.tukey_hsd([1, 2, 3], [2, np.inf], [6, 7, 3]) + + def test_is_1d(self): + with assert_raises(ValueError, match="...must be one-dimensional"): + stats.tukey_hsd([[1, 2], [2, 3]], [2, 5], [5, 23, 6]) + + def test_no_empty(self): + with assert_raises(ValueError, match="...must be greater than one"): + stats.tukey_hsd([], [2, 5], [4, 5, 6]) + + def test_equal_var_input_validation(self): + msg = "Expected a boolean value for 'equal_var'" + with assert_raises(TypeError, match=msg): + stats.tukey_hsd([1, 2, 3], [2, 5], [6, 7], equal_var="False") + + @pytest.mark.parametrize("nargs", (0, 1)) + def test_not_enough_treatments(self, nargs): + with assert_raises(ValueError, match="...more than 1 treatment."): + stats.tukey_hsd(*([[23, 7, 3]] * nargs)) + + @pytest.mark.parametrize("cl", [-.5, 0, 1, 2]) + def test_conf_level_invalid(self, cl): + with assert_raises(ValueError, match="must be between 0 and 1"): + r = stats.tukey_hsd([23, 7, 3], [3, 4], [9, 4]) + r.confidence_interval(cl) + + def test_2_args_ttest(self): + # that with 2 treatments the `pvalue` is equal to that of `ttest_ind` + res_tukey = stats.tukey_hsd(*self.data_diff_size[:2]) + res_ttest = stats.ttest_ind(*self.data_diff_size[:2]) + assert_allclose(res_ttest.pvalue, res_tukey.pvalue[0, 1]) + assert_allclose(res_ttest.pvalue, res_tukey.pvalue[1, 0]) + + +class TestGamesHowell: + # data with unequal variances + data_same_size = ([24., 23., 31., 51.], + [34., 18., 18., 26.], + [17., 68., 59., 7.]) + + data_diff_size = ([30., 23., 51.], + [-81., 71., -27., 63.], + [42., 11., 29., 19., 50.], + [23., 22., 20., 18., 9.]) + + spss_same_size = """ + Mean Diff Lower Bound Upper Bound Sig + 0 - 1 8.25000000 -16.5492749527311 33.0492749527311 0.558733632413559 + 0 - 2 -5.50000000 -63.6702454316458 52.6702454316458 0.941147750599221 + 1 - 2 -13.7500000 -74.3174374251372 46.8174374251372 0.682983914946841 + """ + + spss_diff_size = """ + Mean Diff Lower Bound Upper Bound Sig + 0 - 1 28.16666667 -141.985416377670 198.318749711003 0.8727542747886180 + 0 - 2 4.466666667 -37.2830676783904 46.2164010117237 0.9752628408671710 + 0 - 3 16.26666667 -35.0933112382470 67.6266445715803 0.4262506151302880 + 1 - 2 -23.70000000 -195.315617201249 147.915617201249 0.9148950609000590 + 1 - 3 -11.90000000 -188.105478728519 164.305478728519 0.9861432250093960 + 2 - 3 11.80000000 -16.2894857524254 39.8894857524254 0.4755344436335670 + """ + + @pytest.mark.xslow + @pytest.mark.parametrize("data, res_expect_str", + ((data_same_size, spss_same_size), + (data_diff_size, spss_diff_size)), + ids=["equal size sample", + "unequal sample size"]) + def test_compare_spss(self, data, res_expect_str): + """ + DATA LIST LIST /Group (F1.0) Value (F8.2). + BEGIN DATA + 0 24 + 0 23 + 0 31 + 0 51 + 1 34 + 1 18 + 1 18 + 1 26 + 2 17 + 2 68 + 2 59 + 2 7 + END DATA. + + ONEWAY Value BY Group + /MISSING ANALYSIS + /POSTHOC=GH ALPHA(0.05). + """ + res_expect = np.asarray( + res_expect_str.replace(" - ", " ").split()[7:], + dtype=float).reshape(-1, 6) + res_games = stats.tukey_hsd(*data, equal_var=False) + conf = res_games.confidence_interval() + # loop over the comparisons + for i, j, s, l, h, p in res_expect: + i, j = int(i), int(j) + assert_allclose(res_games.statistic[i, j], s, atol=1e-8) + assert_allclose(res_games.pvalue[i, j], p, atol=1e-8) + assert_allclose(conf.low[i, j], l, atol=1e-6) + assert_allclose(conf.high[i, j], h, atol=1e-5) + + r_same_size = """ + q value Pr(>|q|) + 1 - 0 == 0 -1.5467805948856344 0.55873362851759 + 2 - 0 == 0 0.4726721776628535 0.94114775035993 + 2 - 1 == 0 1.246837541297872 0.68298393799782 + """ + + r_diff_size = """ + q value Pr(>|q|) + 1 - 0 == 0 -1.0589317485313876 0.87275427357438 + 2 - 0 == 0 -0.5716222106144833 0.97526284087419 + 3 - 0 == 0 -2.6209678382077000 0.42625067714691 + 2 - 1 == 0 0.8971899898179028 0.91489506061850 + 3 - 1 == 0 0.4579447210555352 0.98614322544695 + 3 - 2 == 0 -2.198800177874794 0.47553444364614 + """ + + @pytest.mark.parametrize("data, res_expect_str", + ((data_same_size, r_same_size), + (data_diff_size, r_diff_size)), + ids=["equal size sample", + "unequal sample size"]) + def test_compare_r(self, data, res_expect_str): + """ + games-howell is provided by PMCMRplus package + https://search.r-project.org/CRAN/refmans/PMCMRplus/html/gamesHowellTest.html + > library("PMCMRplus") + > options(digits=16) + > table = data.frame( + values = c(24., 23., 31., 51., 34., 18., 18., 26., 17., 68., 59., 7.), + groups = c("0", "0", "0", "0", "1", "1", "1", "1", "2", "2", "2", "2") + ) + > table$groups = as.factor(table$groups) + > fit <-aov(values ~ groups, table) + > res = gamesHowellTest(fit) + > summary(res) + """ + res_expect = np.asarray( + res_expect_str.replace(" - ", " ") + .replace(" == ", " ").split()[3:], + dtype=float).reshape(-1, 5) + res_games = stats.tukey_hsd(*data, equal_var=False) + # loop over the comparisons + # note confidence intervals are not provided by PMCMRplus + for j, i, _, _, p in res_expect: + i, j = int(i), int(j) + assert_allclose(res_games.pvalue[i, j], p, atol=1e-7) + + # Data validation test has been covered by TestTukeyHSD + # like empty, 1d, inf, and lack of tretments + # because games_howell leverage _tukey_hsd_iv() + +class TestPoissonMeansTest: + @pytest.mark.parametrize("c1, n1, c2, n2, p_expect", ( + # example from [1], 6. Illustrative examples: Example 1 + [0, 100, 3, 100, 0.0884], + [2, 100, 6, 100, 0.1749] + )) + def test_paper_examples(self, c1, n1, c2, n2, p_expect): + res = stats.poisson_means_test(c1, n1, c2, n2) + assert_allclose(res.pvalue, p_expect, atol=1e-4) + + @pytest.mark.parametrize("c1, n1, c2, n2, p_expect, alt, d", ( + # These test cases are produced by the wrapped fortran code from the + # original authors. Using a slightly modified version of this fortran, + # found here, https://github.com/nolanbconaway/poisson-etest, + # additional tests were created. + [20, 10, 20, 10, 0.9999997568929630, 'two-sided', 0], + [10, 10, 10, 10, 0.9999998403241203, 'two-sided', 0], + [50, 15, 1, 1, 0.09920321053409643, 'two-sided', .05], + [3, 100, 20, 300, 0.12202725450896404, 'two-sided', 0], + [3, 12, 4, 20, 0.40416087318539173, 'greater', 0], + [4, 20, 3, 100, 0.008053640402974236, 'greater', 0], + # publishing paper does not include a `less` alternative, + # so it was calculated with switched argument order and + # alternative="greater" + [4, 20, 3, 10, 0.3083216325432898, 'less', 0], + [1, 1, 50, 15, 0.09322998607245102, 'less', 0] + )) + def test_fortran_authors(self, c1, n1, c2, n2, p_expect, alt, d): + res = stats.poisson_means_test(c1, n1, c2, n2, alternative=alt, diff=d) + assert_allclose(res.pvalue, p_expect, atol=2e-6, rtol=1e-16) + + def test_different_results(self): + # The implementation in Fortran is known to break down at higher + # counts and observations, so we expect different results. By + # inspection we can infer the p-value to be near one. + count1, count2 = 10000, 10000 + nobs1, nobs2 = 10000, 10000 + res = stats.poisson_means_test(count1, nobs1, count2, nobs2) + assert_allclose(res.pvalue, 1) + + def test_less_than_zero_lambda_hat2(self): + # demonstrates behavior that fixes a known fault from original Fortran. + # p-value should clearly be near one. + count1, count2 = 0, 0 + nobs1, nobs2 = 1, 1 + res = stats.poisson_means_test(count1, nobs1, count2, nobs2) + assert_allclose(res.pvalue, 1) + + def test_input_validation(self): + count1, count2 = 0, 0 + nobs1, nobs2 = 1, 1 + + # test non-integral events + message = '`k1` and `k2` must be integers.' + with assert_raises(TypeError, match=message): + stats.poisson_means_test(.7, nobs1, count2, nobs2) + with assert_raises(TypeError, match=message): + stats.poisson_means_test(count1, nobs1, .7, nobs2) + + # test negative events + message = '`k1` and `k2` must be greater than or equal to 0.' + with assert_raises(ValueError, match=message): + stats.poisson_means_test(-1, nobs1, count2, nobs2) + with assert_raises(ValueError, match=message): + stats.poisson_means_test(count1, nobs1, -1, nobs2) + + # test negative sample size + message = '`n1` and `n2` must be greater than 0.' + with assert_raises(ValueError, match=message): + stats.poisson_means_test(count1, -1, count2, nobs2) + with assert_raises(ValueError, match=message): + stats.poisson_means_test(count1, nobs1, count2, -1) + + # test negative difference + message = 'diff must be greater than or equal to 0.' + with assert_raises(ValueError, match=message): + stats.poisson_means_test(count1, nobs1, count2, nobs2, diff=-1) + + # test invalid alternative + message = 'Alternative must be one of ...' + with assert_raises(ValueError, match=message): + stats.poisson_means_test(1, 2, 1, 2, alternative='error') + + +@make_xp_test_case(stats.bws_test) +class TestBWSTest: + + def test_bws_input_validation(self, xp): + rng = np.random.default_rng(4571775098104213308) + + x, y = rng.random(size=(2, 7)) + x, y = xp.asarray(x), xp.asarray(y) + + message = '`axis` must be an integer or None.' + with pytest.raises(ValueError, match=message): + stats.bws_test(x, y, axis=1.5) + + message = '`axis` is not compatible with the shapes of the inputs.' + with pytest.raises(ValueError, match=message): + stats.bws_test(x, y, axis=3) + + message = '`x` and `y` must contain...' + with pytest.raises(ValueError, match=message): + stats.bws_test(x, xp.asarray([], dtype=x.dtype)) + + message = 'alternative` must be one of...' + with pytest.raises(ValueError, match=message): + stats.bws_test(x, y, alternative='ekki-ekki') + + message = 'method` must be an instance of...' + with pytest.raises(ValueError, match=message): + stats.bws_test(x, y, method=42) + + @pytest.mark.parametrize('dtype', [None, 'float32', 'float64']) + def test_against_published_reference(self, dtype, xp): + # Test against Example 2 in bws_test Reference [1], pg 9 + # https://link.springer.com/content/pdf/10.1007/BF02762032.pdf + dtype = xp_default_dtype(xp) if dtype is None else getattr(xp, dtype) + x = xp.asarray([1, 2, 3, 4, 6, 7, 8], dtype=dtype) + y = xp.asarray([5, 9, 10, 11, 12, 13, 14], dtype=dtype) + res = stats.bws_test(x, y, alternative='two-sided') + xp_assert_close(res.statistic, xp.asarray(5.132, dtype=dtype), atol=1e-3) + xp_assert_close(res.pvalue, xp.asarray(10/3432, dtype=dtype)) + + @pytest.mark.parametrize(('alternative', 'statistic', 'pvalue'), + [('two-sided', 1.7510204081633, 0.1264422777777), + ('less', -1.7510204081633, 0.05754662004662), + ('greater', -1.7510204081633, 0.9424533799534)]) + def test_against_R(self, alternative, statistic, pvalue, xp): + # Test against R library BWStest function bws_test + # library(BWStest) + # options(digits=16) + # x = c(...) + # y = c(...) + # bws_test(x, y, alternative='two.sided') + rng = np.random.default_rng(4571775098104213308) + x, y = rng.random(size=(2, 7)) + x, y = xp.asarray(x, dtype=xp.float64), xp.asarray(y, dtype=xp.float64) + res = stats.bws_test(x, y, alternative=alternative) + statistic = xp.asarray(statistic, dtype=xp.float64) + pvalue = xp.asarray(pvalue, dtype=xp.float64) + xp_assert_close(res.statistic, statistic, rtol=1e-13) + xp_assert_close(res.pvalue, pvalue, atol=1e-2, rtol=1e-1) + + @pytest.mark.parametrize(('alternative', 'statistic', 'pvalue'), + [('two-sided', 1.142629265891, 0.2903950180801), + ('less', 0.99629665877411, 0.8545660222131), + ('greater', 0.99629665877411, 0.1454339777869)]) + def test_against_R_imbalanced(self, alternative, statistic, pvalue, xp): + # Test against R library BWStest function bws_test + # library(BWStest) + # options(digits=16) + # x = c(...) + # y = c(...) + # bws_test(x, y, alternative='two.sided') + rng = np.random.default_rng(5429015622386364034) + x = rng.random(size=9) + y = rng.random(size=8) + x, y = xp.asarray(x), xp.asarray(y) + res = stats.bws_test(x, y, alternative=alternative) + x, y = xp.asarray(x, dtype=xp.float64), xp.asarray(y, dtype=xp.float64) + res = stats.bws_test(x, y, alternative=alternative) + statistic = xp.asarray(statistic, dtype=xp.float64) + pvalue = xp.asarray(pvalue, dtype=xp.float64) + xp_assert_close(res.statistic, statistic, rtol=1e-13) + xp_assert_close(res.pvalue, pvalue, atol=1e-2, rtol=1e-1) + + def test_method(self, xp): + # Test that `method` parameter has the desired effect + rng = np.random.default_rng(1520514347193347862) + x, y = rng.random(size=(2, 10)) + x, y = xp.asarray(x), xp.asarray(y) + + rng = np.random.default_rng(1520514347193347862) + method = stats.PermutationMethod(n_resamples=10, rng=rng) + res1 = stats.bws_test(x, y, method=method) + + assert res1.null_distribution.shape[0] == 10 + + rng = np.random.default_rng(1520514347193347862) + method = stats.PermutationMethod(n_resamples=10, rng=rng) + res2 = stats.bws_test(x, y, method=method) + + xp_assert_close(res1.null_distribution, res2.null_distribution) + + rng = np.random.default_rng(5205143471933478621) + method = stats.PermutationMethod(n_resamples=10, rng=rng) + res3 = stats.bws_test(x, y, method=method) + + assert xp.all(xp.abs(res3.null_distribution - res1.null_distribution) > 0.1) + + def test_directions(self, xp): + # Sanity check of the sign of the one-sided statistic + rng = np.random.default_rng(1520514347193347862) + x = xp.asarray(rng.random(size=5)) + y = x - 1. + + pmin = xp.asarray(1 / 252, dtype=xp.float64) + pmax = xp.asarray(1, dtype=xp.float64) + + res = stats.bws_test(x, y, alternative='greater') + assert res.statistic > 0 + xp_assert_equal(res.pvalue, pmin) + + res = stats.bws_test(x, y, alternative='less') + assert res.statistic > 0 + xp_assert_equal(res.pvalue, pmax) + + res = stats.bws_test(y, x, alternative='less') + assert res.statistic < 0 + xp_assert_equal(res.pvalue, pmin) + + res = stats.bws_test(y, x, alternative='greater') + assert res.statistic < 0 + xp_assert_equal(res.pvalue, pmax) + + @pytest.mark.parametrize('axis', [0, 1]) + def test_nd(self, axis, xp): + rng = np.random.default_rng(1520514347193347862) + z = rng.random(size=(2, 3, 5)) + + x, y = z + res = stats.bws_test(xp.asarray(x), xp.asarray(y), axis=axis) + + x, y = np.moveaxis(z, axis+1, -1) + ref_statistic, ref_pvalue = [], [] + for x_, y_ in zip(x, y): + ref = stats.bws_test(xp.asarray(x_), xp.asarray(y_)) + ref_statistic.append(ref.statistic) + ref_pvalue.append(ref.pvalue) + ref_statistic, ref_pvalue = xp.stack(ref_statistic), xp.stack(ref_pvalue) + + xp_assert_close(res.statistic, ref_statistic) + xp_assert_close(res.pvalue, ref_pvalue) + + def test_axis_None(self, xp): + rng = np.random.default_rng(5205143471933478621) + x, y = rng.random(size=(2, 2, 3)) + x, y = xp.asarray(x), xp.asarray(y) + res = stats.bws_test(x, y, axis=None) + ref = stats.bws_test(xp_ravel(x), xp_ravel(y)) + xp_assert_close(res.statistic, ref.statistic) + xp_assert_close(res.pvalue, ref.pvalue) diff --git a/v12x/pylib/scipy/stats/tests/test_kdeoth.py b/v12x/pylib/scipy/stats/tests/test_kdeoth.py new file mode 100644 index 0000000000000000000000000000000000000000..104d31f9ec0c5d93ae89246d8625a14694e64990 --- /dev/null +++ b/v12x/pylib/scipy/stats/tests/test_kdeoth.py @@ -0,0 +1,678 @@ +from scipy import stats, linalg, integrate +import numpy as np +from numpy.testing import (assert_almost_equal, assert_, assert_equal, + assert_array_almost_equal, + assert_array_almost_equal_nulp, assert_allclose) +import pytest +from pytest import raises as assert_raises + + +def test_kde_1d(): + #some basic tests comparing to normal distribution + rng = np.random.default_rng(8765678) + n_basesample = 500 + xn = rng.normal(0, 1, n_basesample) + xnmean = xn.mean() + xnstd = xn.std(ddof=1) + + # get kde for original sample + gkde = stats.gaussian_kde(xn) + + # evaluate the density function for the kde for some points + xx = np.asarray([0.1, 0.5, 0.9]) + loc, scale = gkde.dataset, np.sqrt(gkde.covariance) + assert_allclose( + gkde(xx), + stats.norm.pdf(xx[:, None], loc=loc, scale=scale).sum(axis=-1) / gkde.n, + rtol=5e-14 + ) + + xs = np.linspace(-7, 7, 501) + kdepdf = gkde.evaluate(xs) + normpdf = stats.norm.pdf(xs, loc=xnmean, scale=xnstd) + intervall = xs[1] - xs[0] + + assert_(np.sum((kdepdf - normpdf)**2)*intervall < 0.01) + prob1 = gkde.integrate_box_1d(xnmean, np.inf) + prob2 = gkde.integrate_box_1d(-np.inf, xnmean) + assert_almost_equal(prob1, 0.5, decimal=1) + assert_almost_equal(prob2, 0.5, decimal=1) + assert_almost_equal(gkde.integrate_box(xnmean, np.inf), prob1, decimal=13) + assert_almost_equal(gkde.integrate_box(-np.inf, xnmean), prob2, decimal=13) + + assert_almost_equal(gkde.integrate_kde(gkde), + (kdepdf**2).sum()*intervall, decimal=2) + assert_almost_equal(gkde.integrate_gaussian(xnmean, xnstd**2), + (kdepdf*normpdf).sum()*intervall, decimal=2) + + +def test_kde_1d_weighted(): + #some basic tests comparing to normal distribution + rng = np.random.default_rng(8765678) + n_basesample = 500 + xn = rng.normal(0, 1, n_basesample) + wn = rng.random(n_basesample) + xnmean = np.average(xn, weights=wn) + xnstd = np.sqrt(np.average((xn-xnmean)**2, weights=wn)) + + # get kde for original sample + gkde = stats.gaussian_kde(xn, weights=wn) + + # evaluate the density function for the kde for some points + # evaluate the density function for the kde for some points + xx = np.asarray([0.1, 0.5, 0.9]) + loc, scale = gkde.dataset, np.sqrt(gkde.covariance) + + pdf = stats.norm.pdf + assert_allclose( + gkde(xx), + np.sum(pdf(xx[:, None], loc=loc, scale=scale) * gkde.weights, axis=-1), + rtol=5e-14 + ) + + xs = np.linspace(-7, 7, 501) + kdepdf = gkde.evaluate(xs) + normpdf = stats.norm.pdf(xs, loc=xnmean, scale=xnstd) + intervall = xs[1] - xs[0] + + assert_(np.sum((kdepdf - normpdf)**2)*intervall < 0.01) + prob1 = gkde.integrate_box_1d(xnmean, np.inf) + prob2 = gkde.integrate_box_1d(-np.inf, xnmean) + assert_almost_equal(prob1, 0.5, decimal=1) + assert_almost_equal(prob2, 0.5, decimal=1) + assert_almost_equal(gkde.integrate_box(xnmean, np.inf), prob1, decimal=13) + assert_almost_equal(gkde.integrate_box(-np.inf, xnmean), prob2, decimal=13) + + assert_almost_equal(gkde.integrate_kde(gkde), + (kdepdf**2).sum()*intervall, decimal=2) + assert_almost_equal(gkde.integrate_gaussian(xnmean, xnstd**2), + (kdepdf*normpdf).sum()*intervall, decimal=2) + + +@pytest.mark.parametrize("n_basesample", + [ + 20, + pytest.param(500, marks=[pytest.mark.xslow]) + ] +) +def test_kde_2d(n_basesample): + #some basic tests comparing to normal distribution + rng = np.random.default_rng(8765678) + + mean = np.array([1.0, 3.0]) + covariance = np.array([[1.0, 2.0], [2.0, 6.0]]) + + # Need transpose (shape (2, 500)) for kde + xn = rng.multivariate_normal(mean, covariance, size=n_basesample).T + + # get kde for original sample + gkde = stats.gaussian_kde(xn) + + # evaluate vs multivariate normal, using the KDE definition + xx = np.asarray([[1, 2], [3, 4], [5, 6]]) + arg = xx[:, None, :] - gkde.dataset.T + pdf = stats.multivariate_normal.pdf + assert_allclose( + gkde(xx.T), + pdf(arg, cov=gkde.covariance).sum(axis=-1) / gkde.n, + rtol=5e-14 + ) + + # ... and cdf + cdf = stats.multivariate_normal.cdf + lo, hi = [-1, -2], [0, 0] + lo_, hi_ = lo - gkde.dataset.T, hi - gkde.dataset.T + assert_allclose( + gkde.integrate_box(lo, hi, rng=rng), + cdf(hi_, lower_limit=lo_, cov=gkde.covariance, rng=rng).sum(axis=-1) / gkde.n, + rtol=5e-7 + ) + + # evaluate the density function for the kde for some points + x, y = np.mgrid[-7:7:500j, -7:7:500j] + grid_coords = np.vstack([x.ravel(), y.ravel()]) + kdepdf = gkde.evaluate(grid_coords) + kdepdf = kdepdf.reshape(500, 500) + + normpdf = stats.multivariate_normal.pdf(np.dstack([x, y]), + mean=mean, cov=covariance) + intervall = y.ravel()[1] - y.ravel()[0] + + assert_(np.sum((kdepdf - normpdf)**2) * (intervall**2) < 0.01) + + small = -1e100 + large = 1e100 + prob1 = gkde.integrate_box([small, mean[1]], [large, large], rng=rng) + prob2 = gkde.integrate_box([small, small], [large, mean[1]], rng=rng) + + assert_almost_equal(prob1, 0.5, decimal=1) + assert_almost_equal(prob2, 0.5, decimal=1) + assert_almost_equal(gkde.integrate_kde(gkde), + (kdepdf**2).sum()*(intervall**2), decimal=2) + assert_almost_equal(gkde.integrate_gaussian(mean, covariance), + (kdepdf*normpdf).sum()*(intervall**2), decimal=2) + + +@pytest.mark.parametrize("n_basesample", + [ + 20, + pytest.param(500, marks=[pytest.mark.xslow]) + ] +) +def test_kde_2d_weighted(n_basesample): + #some basic tests comparing to normal distribution + rng = np.random.RandomState(8765678) + + mean = np.array([1.0, 3.0]) + covariance = np.array([[1.0, 2.0], [2.0, 6.0]]) + + # Need transpose (shape (2, 500)) for kde + xn = rng.multivariate_normal(mean, covariance, size=n_basesample).T + wn = rng.rand(n_basesample) + + # get kde for original sample + gkde = stats.gaussian_kde(xn, weights=wn) + + + # evaluate vs multivariate normal, using the kde definition + xx = np.asarray([[1, 2], [3, 4], [5, 6]]) + arg = xx[:, None, :] - gkde.dataset.T + pdf = stats.multivariate_normal.pdf + assert_allclose( + gkde(xx.T), + np.sum(pdf(arg, cov=gkde.covariance) * gkde.weights, axis=-1), + rtol=6e-14 + ) + + # ... and cdf + cdf = stats.multivariate_normal.cdf + lo, hi = [-1, -2], [0, 0] + lo_, hi_ = lo - gkde.dataset.T, hi - gkde.dataset.T + assert_allclose( + gkde.integrate_box(lo, hi, rng=rng), + np.sum(cdf(hi_, lower_limit=lo_, cov=gkde.covariance, rng=rng) * + gkde.weights, axis=-1), + rtol=5e-6 + ) + + # evaluate the density function for the kde for some points + x, y = np.mgrid[-7:7:500j, -7:7:500j] + grid_coords = np.vstack([x.ravel(), y.ravel()]) + kdepdf = gkde.evaluate(grid_coords) + kdepdf = kdepdf.reshape(500, 500) + + normpdf = stats.multivariate_normal.pdf(np.dstack([x, y]), + mean=mean, cov=covariance) + intervall = y.ravel()[1] - y.ravel()[0] + + assert_(np.sum((kdepdf - normpdf)**2) * (intervall**2) < 0.01) + + small = -1e100 + large = 1e100 + prob1 = gkde.integrate_box([small, mean[1]], [large, large], rng=rng) + prob2 = gkde.integrate_box([small, small], [large, mean[1]], rng=rng) + + assert_almost_equal(prob1, 0.5, decimal=1) + assert_almost_equal(prob2, 0.5, decimal=1) + assert_almost_equal(gkde.integrate_kde(gkde), + (kdepdf**2).sum()*(intervall**2), decimal=2) + assert_almost_equal(gkde.integrate_gaussian(mean, covariance), + (kdepdf*normpdf).sum()*(intervall**2), decimal=2) + + +def test_kde_bandwidth_method(): + def scotts_factor(kde_obj): + """Same as default, just check that it works.""" + return np.power(kde_obj.n, -1./(kde_obj.d+4)) + + rng = np.random.default_rng(8765678) + n_basesample = 50 + xn = rng.normal(0, 1, n_basesample) + + # Default + gkde = stats.gaussian_kde(xn) + # Supply a callable + gkde2 = stats.gaussian_kde(xn, bw_method=scotts_factor) + # Supply a scalar + gkde3 = stats.gaussian_kde(xn, bw_method=gkde.factor) + + xs = np.linspace(-7,7,51) + kdepdf = gkde.evaluate(xs) + kdepdf2 = gkde2.evaluate(xs) + assert_almost_equal(kdepdf, kdepdf2) + kdepdf3 = gkde3.evaluate(xs) + assert_almost_equal(kdepdf, kdepdf3) + + assert_raises(ValueError, stats.gaussian_kde, xn, bw_method='wrongstring') + + +def test_kde_bandwidth_method_weighted(): + def scotts_factor(kde_obj): + """Same as default, just check that it works.""" + return np.power(kde_obj.neff, -1./(kde_obj.d+4)) + + rng = np.random.default_rng(8765678) + n_basesample = 50 + xn = rng.normal(0, 1, n_basesample) + + # Default + gkde = stats.gaussian_kde(xn) + # Supply a callable + gkde2 = stats.gaussian_kde(xn, bw_method=scotts_factor) + # Supply a scalar + gkde3 = stats.gaussian_kde(xn, bw_method=gkde.factor) + + xs = np.linspace(-7,7,51) + kdepdf = gkde.evaluate(xs) + kdepdf2 = gkde2.evaluate(xs) + assert_almost_equal(kdepdf, kdepdf2) + kdepdf3 = gkde3.evaluate(xs) + assert_almost_equal(kdepdf, kdepdf3) + + assert_raises(ValueError, stats.gaussian_kde, xn, bw_method='wrongstring') + + +# Subclasses that should stay working (extracted from various sources). +# Unfortunately the earlier design of gaussian_kde made it necessary for users +# to create these kinds of subclasses, or call _compute_covariance() directly. + +class _kde_subclass1(stats.gaussian_kde): + def __init__(self, dataset): + self.dataset = np.atleast_2d(dataset) + self.d, self.n = self.dataset.shape + self.covariance_factor = self.scotts_factor + self._compute_covariance() + + +class _kde_subclass2(stats.gaussian_kde): + def __init__(self, dataset): + self.covariance_factor = self.scotts_factor + super().__init__(dataset) + + +class _kde_subclass4(stats.gaussian_kde): + def covariance_factor(self): + return 0.5 * self.silverman_factor() + + +def test_gaussian_kde_subclassing(): + x1 = np.array([-7, -5, 1, 4, 5], dtype=float) + xs = np.linspace(-10, 10, num=50) + + # gaussian_kde itself + kde = stats.gaussian_kde(x1) + ys = kde(xs) + + # subclass 1 + kde1 = _kde_subclass1(x1) + y1 = kde1(xs) + assert_array_almost_equal_nulp(ys, y1, nulp=10) + + # subclass 2 + kde2 = _kde_subclass2(x1) + y2 = kde2(xs) + assert_array_almost_equal_nulp(ys, y2, nulp=10) + + # subclass 3 was removed because we have no obligation to maintain support + # for user invocation of private methods + + # subclass 4 + kde4 = _kde_subclass4(x1) + y4 = kde4(x1) + y_expected = [0.06292987, 0.06346938, 0.05860291, 0.08657652, 0.07904017] + + assert_array_almost_equal(y_expected, y4, decimal=6) + + # Not a subclass, but check for use of _compute_covariance() + kde5 = kde + kde5.covariance_factor = lambda: kde.factor + kde5._compute_covariance() + y5 = kde5(xs) + assert_array_almost_equal_nulp(ys, y5, nulp=10) + + +def test_gaussian_kde_covariance_caching(): + x1 = np.array([-7, -5, 1, 4, 5], dtype=float) + xs = np.linspace(-10, 10, num=5) + # These expected values are from scipy 0.10, before some changes to + # gaussian_kde. They were not compared with any external reference. + y_expected = [0.02463386, 0.04689208, 0.05395444, 0.05337754, 0.01664475] + + # Set the bandwidth, then reset it to the default. + kde = stats.gaussian_kde(x1) + kde.set_bandwidth(bw_method=0.5) + kde.set_bandwidth(bw_method='scott') + y2 = kde(xs) + + assert_array_almost_equal(y_expected, y2, decimal=7) + + +def test_gaussian_kde_monkeypatch(): + """Ugly, but people may rely on this. See scipy pull request 123, + specifically the linked ML thread "Width of the Gaussian in stats.kde". + If it is necessary to break this later on, that is to be discussed on ML. + """ + x1 = np.array([-7, -5, 1, 4, 5], dtype=float) + xs = np.linspace(-10, 10, num=50) + + # The old monkeypatched version to get at Silverman's Rule. + kde = stats.gaussian_kde(x1) + kde.covariance_factor = kde.silverman_factor + kde._compute_covariance() + y1 = kde(xs) + + # The new saner version. + kde2 = stats.gaussian_kde(x1, bw_method='silverman') + y2 = kde2(xs) + + assert_array_almost_equal_nulp(y1, y2, nulp=10) + + +def test_kde_integer_input(): + """Regression test for #1181.""" + x1 = np.arange(5) + kde = stats.gaussian_kde(x1) + y_expected = [0.13480721, 0.18222869, 0.19514935, 0.18222869, 0.13480721] + assert_array_almost_equal(kde(x1), y_expected, decimal=6) + + +_ftypes = ['float32', 'float64', 'float96', 'float128', 'int32', 'int64'] + + +@pytest.mark.filterwarnings("ignore::DeprecationWarning") +@pytest.mark.parametrize("bw_type", _ftypes + ["scott", "silverman"]) +@pytest.mark.parametrize("dtype", _ftypes) +def test_kde_output_dtype(dtype, bw_type): + # Check whether the datatypes are available + dtype = getattr(np, dtype, None) + + if bw_type in ["scott", "silverman"]: + bw = bw_type + else: + bw_type = getattr(np, bw_type, None) + bw = bw_type(3) if bw_type else None + + if any(dt is None for dt in [dtype, bw]): + pytest.skip() + + weights = np.arange(5, dtype=dtype) + dataset = np.arange(5, dtype=dtype) + k = stats.gaussian_kde(dataset, bw_method=bw, weights=weights) + points = np.arange(5, dtype=dtype) + result = k(points) + # weights are always cast to float64 + assert result.dtype == np.result_type(dataset, points, np.float64(weights), + k.factor) + + +def test_pdf_logpdf_validation(): + rng = np.random.default_rng(64202298293133848336925499069837723291) + xn = rng.standard_normal((2, 10)) + gkde = stats.gaussian_kde(xn) + xs = rng.standard_normal((3, 10)) + + msg = "points have dimension 3, dataset has dimension 2" + with pytest.raises(ValueError, match=msg): + gkde.logpdf(xs) + + +def test_pdf_logpdf(): + rng = np.random.default_rng(1) + n_basesample = 50 + xn = rng.normal(0, 1, n_basesample) + + # Default + gkde = stats.gaussian_kde(xn) + + xs = np.linspace(-15, 12, 25) + pdf = gkde.evaluate(xs) + pdf2 = gkde.pdf(xs) + assert_almost_equal(pdf, pdf2, decimal=12) + + logpdf = np.log(pdf) + logpdf2 = gkde.logpdf(xs) + assert_almost_equal(logpdf, logpdf2, decimal=12) + + # There are more points than data + gkde = stats.gaussian_kde(xs) + pdf = np.log(gkde.evaluate(xn)) + pdf2 = gkde.logpdf(xn) + assert_almost_equal(pdf, pdf2, decimal=12) + + +def test_pdf_logpdf_weighted(): + rng = np.random.default_rng(1) + n_basesample = 50 + xn = rng.normal(0, 1, n_basesample) + wn = rng.random(n_basesample) + + # Default + gkde = stats.gaussian_kde(xn, weights=wn) + + xs = np.linspace(-15, 12, 25) + pdf = gkde.evaluate(xs) + pdf2 = gkde.pdf(xs) + assert_almost_equal(pdf, pdf2, decimal=12) + + logpdf = np.log(pdf) + logpdf2 = gkde.logpdf(xs) + assert_almost_equal(logpdf, logpdf2, decimal=12) + + # There are more points than data + rng = np.random.default_rng(4531935345) + gkde = stats.gaussian_kde(xs, weights=rng.random(len(xs))) + pdf = np.log(gkde.evaluate(xn)) + pdf2 = gkde.logpdf(xn) + assert_almost_equal(pdf, pdf2, decimal=12) + + +def test_marginal_1_axis(): + rng = np.random.default_rng(6111799263660870475) + n_data = 50 + n_dim = 10 + dataset = rng.normal(size=(n_dim, n_data)) + points = rng.normal(size=(n_dim, 3)) + + dimensions = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9]) # dimensions to keep + + kde = stats.gaussian_kde(dataset) + marginal = kde.marginal(dimensions) + pdf = marginal.pdf(points[dimensions]) + + def marginal_pdf_single(point): + def f(x): + x = np.concatenate(([x], point[dimensions])) + return kde.pdf(x)[0] + return integrate.quad(f, -np.inf, np.inf)[0] + + def marginal_pdf(points): + return np.apply_along_axis(marginal_pdf_single, axis=0, arr=points) + + ref = marginal_pdf(points) + + assert_allclose(pdf, ref, rtol=1e-6) + + +@pytest.mark.xslow +def test_marginal_2_axis(): + rng = np.random.default_rng(6111799263660870475) + n_data = 30 + n_dim = 4 + dataset = rng.normal(size=(n_dim, n_data)) + points = rng.normal(size=(n_dim, 3)) + + dimensions = np.array([1, 3]) # dimensions to keep + + kde = stats.gaussian_kde(dataset) + marginal = kde.marginal(dimensions) + pdf = marginal.pdf(points[dimensions]) + + def marginal_pdf(points): + def marginal_pdf_single(point): + def f(y, x): + w, z = point[dimensions] + x = np.array([x, w, y, z]) + return kde.pdf(x)[0] + return integrate.dblquad(f, -np.inf, np.inf, -np.inf, np.inf)[0] + + return np.apply_along_axis(marginal_pdf_single, axis=0, arr=points) + + ref = marginal_pdf(points) + + assert_allclose(pdf, ref, rtol=1e-6) + + +def test_marginal_iv(): + # test input validation + rng = np.random.default_rng(6111799263660870475) + n_data = 30 + n_dim = 4 + dataset = rng.normal(size=(n_dim, n_data)) + points = rng.normal(size=(n_dim, 3)) + + kde = stats.gaussian_kde(dataset) + + # check that positive and negative indices are equivalent + dimensions1 = [-1, 1] + marginal1 = kde.marginal(dimensions1) + pdf1 = marginal1.pdf(points[dimensions1]) + + dimensions2 = [3, -3] + marginal2 = kde.marginal(dimensions2) + pdf2 = marginal2.pdf(points[dimensions2]) + + assert_equal(pdf1, pdf2) + + # IV for non-integer dimensions + message = "Elements of `dimensions` must be integers..." + with pytest.raises(ValueError, match=message): + kde.marginal([1, 2.5]) + + # IV for uniqueness + message = "All elements of `dimensions` must be unique." + with pytest.raises(ValueError, match=message): + kde.marginal([1, 2, 2]) + + # IV for non-integer dimensions + message = (r"Dimensions \[-5 6\] are invalid for a distribution in 4...") + with pytest.raises(ValueError, match=message): + kde.marginal([1, -5, 6]) + + +@pytest.mark.xslow +def test_logpdf_overflow(): + # regression test for gh-12988; testing against linalg instability for + # very high dimensionality kde + rng = np.random.default_rng(1) + n_dimensions = 2500 + n_samples = 5000 + xn = np.array([rng.normal(0, 1, n_samples) + (n) for n in range( + 0, n_dimensions)]) + + # Default + gkde = stats.gaussian_kde(xn) + + logpdf = gkde.logpdf(np.arange(0, n_dimensions)) + np.testing.assert_equal(np.isneginf(logpdf[0]), False) + np.testing.assert_equal(np.isnan(logpdf[0]), False) + + +def test_weights_intact(): + # regression test for gh-9709: weights are not modified + rng = np.random.default_rng(12345) + vals = rng.lognormal(size=100) + weights = rng.choice([1.0, 10.0, 100], size=vals.size) + orig_weights = weights.copy() + + stats.gaussian_kde(np.log10(vals), weights=weights) + assert_allclose(weights, orig_weights, atol=1e-14, rtol=1e-14) + + +def test_weights_integer(): + # integer weights are OK, cf gh-9709 (comment) + values = [0.2, 13.5, 21.0, 75.0, 99.0] + weights = [1, 2, 4, 8, 16] # a list of integers + pdf_i = stats.gaussian_kde(values, weights=weights) + pdf_f = stats.gaussian_kde(values, weights=np.float64(weights)) + + xn = [0.3, 11, 88] + assert_allclose(pdf_i.evaluate(xn), + pdf_f.evaluate(xn), atol=1e-14, rtol=1e-14) + + +def test_seed(): + # Test the seed option of the resample method + def test_seed_sub(gkde_trail): + n_sample = 200 + # The results should be different without using seed + samp1 = gkde_trail.resample(n_sample) + samp2 = gkde_trail.resample(n_sample) + assert_raises( + AssertionError, assert_allclose, samp1, samp2, atol=1e-13 + ) + # Use integer seed + seed = 831 + samp1 = gkde_trail.resample(n_sample, seed=seed) + samp2 = gkde_trail.resample(n_sample, seed=seed) + assert_allclose(samp1, samp2, atol=1e-13) + # Use RandomState + rstate1 = np.random.RandomState(seed=138) + samp1 = gkde_trail.resample(n_sample, seed=rstate1) + rstate2 = np.random.RandomState(seed=138) + samp2 = gkde_trail.resample(n_sample, seed=rstate2) + assert_allclose(samp1, samp2, atol=1e-13) + + # check that np.random.Generator can be used (numpy >= 1.17) + if hasattr(np.random, 'default_rng'): + # obtain a np.random.Generator object + rng = np.random.default_rng(1234) + gkde_trail.resample(n_sample, seed=rng) + + rng = np.random.default_rng(8765678) + n_basesample = 500 + wn = rng.random(n_basesample) + # Test 1D case + xn_1d = rng.normal(0, 1, n_basesample) + + gkde_1d = stats.gaussian_kde(xn_1d) + test_seed_sub(gkde_1d) + gkde_1d_weighted = stats.gaussian_kde(xn_1d, weights=wn) + test_seed_sub(gkde_1d_weighted) + + # Test 2D case + mean = np.array([1.0, 3.0]) + covariance = np.array([[1.0, 2.0], [2.0, 6.0]]) + xn_2d = rng.multivariate_normal(mean, covariance, size=n_basesample).T + + gkde_2d = stats.gaussian_kde(xn_2d) + test_seed_sub(gkde_2d) + gkde_2d_weighted = stats.gaussian_kde(xn_2d, weights=wn) + test_seed_sub(gkde_2d_weighted) + + +def test_singular_data_covariance_gh10205(): + # When the data lie in a lower-dimensional subspace and this causes + # and exception, check that the error message is informative. + rng = np.random.default_rng(2321583144339784787) + mu = np.array([1, 10, 20]) + sigma = np.array([[4, 10, 0], [10, 25, 0], [0, 0, 100]]) + data = rng.multivariate_normal(mu, sigma, 1000) + try: # doesn't raise any error on some platforms, and that's OK + stats.gaussian_kde(data.T) + except linalg.LinAlgError: + msg = "The data appears to lie in a lower-dimensional subspace..." + with assert_raises(linalg.LinAlgError, match=msg): + stats.gaussian_kde(data.T) + + +def test_fewer_points_than_dimensions_gh17436(): + # When the number of points is fewer than the number of dimensions, the + # the covariance matrix would be singular, and the exception tested in + # test_singular_data_covariance_gh10205 would occur. However, sometimes + # this occurs when the user passes in the transpose of what `gaussian_kde` + # expects. This can result in a huge covariance matrix, so bail early. + rng = np.random.default_rng(2046127537594925772) + rvs = rng.multivariate_normal(np.zeros(3), np.eye(3), size=5) + message = "Number of dimensions is greater than number of samples..." + with pytest.raises(ValueError, match=message): + stats.gaussian_kde(rvs) diff --git a/v12x/pylib/scipy/stats/tests/test_marray.py b/v12x/pylib/scipy/stats/tests/test_marray.py new file mode 100644 index 0000000000000000000000000000000000000000..cbae7321f6acdf1222655ebac391dc874821a66f --- /dev/null +++ b/v12x/pylib/scipy/stats/tests/test_marray.py @@ -0,0 +1,613 @@ +import numpy as np +import pytest +from scipy import stats + +from scipy._lib._array_api import xp_assert_close, xp_assert_equal, _count_nonmasked +from scipy._lib._array_api import make_xp_pytest_param, make_xp_test_case +from scipy._lib._array_api import SCIPY_ARRAY_API, is_torch +from scipy.stats._stats_py import _xp_mean, _xp_var + + +marray = pytest.importorskip('marray') +pytestmark = [ + pytest.mark.skipif( + not SCIPY_ARRAY_API, + reason=( + "special function dispatch to marray required for these tests" + " is hidden behind SCIPY_ARRAY_API flag." + ), + ), +] + +skip_backend = pytest.mark.skip_xp_backends + +def get_arrays(n_arrays, *, dtype='float64', xp=np, shape=(7, 8), all_unique=True, + seed=84912165484321): + mxp = marray._get_namespace(xp) + rng = np.random.default_rng(seed) + + datas, masks = [], [] + for i in range(n_arrays): + data = (rng.random(size=shape) if all_unique + else rng.integers(np.min(shape) // 2, size=shape)) + if dtype.startswith('complex'): + data = 10*data * 10j*rng.standard_normal(size=shape) + data = data.astype(dtype) + datas.append(data) + mask = rng.random(size=shape) > 0.75 + masks.append(mask) + + marrays = [] + nan_arrays = [] + for array, mask in zip(datas, masks): + marrays.append(mxp.asarray(array, mask=mask)) + nan_array = array.copy() + nan_array[mask] = xp.nan + nan_arrays.append(nan_array) + + return mxp, marrays, nan_arrays + + +@skip_backend('jax.numpy', reason="JAX doesn't allow item assignment.") +@pytest.mark.parametrize('fun, kwargs', [ + make_xp_pytest_param(stats.gmean, {}), + make_xp_pytest_param(stats.hmean, {}), + make_xp_pytest_param(stats.pmean, {'p': 2}), + make_xp_pytest_param(stats.expectile, {'alpha': 0.4}), + make_xp_pytest_param(_xp_mean, {}), +]) +@pytest.mark.parametrize('weighted', [False, True]) +@pytest.mark.parametrize('axis', [0, 1, None]) +def test_one_sample_weighted_reducing(fun, kwargs, weighted, axis, xp): + mxp, marrays, narrays = get_arrays(2, xp=xp) + mweights = marrays[1] if weighted else None + nweights = narrays[1] if weighted else None + res = fun(marrays[0], weights=mweights, axis=axis, **kwargs) + ref = fun(narrays[0], weights=nweights, nan_policy='omit', axis=axis, **kwargs) + xp_assert_close(res.data, xp.asarray(ref)) + + +@skip_backend('dask.array', reason='Arrays need `device` attribute: dask/dask#11711') +@skip_backend('jax.numpy', reason="JAX doesn't allow item assignment.") +@pytest.mark.parametrize('fun, kwargs', + [make_xp_pytest_param(stats.moment, {'order': 2}), + make_xp_pytest_param(stats.moment, {'center': 0.5}), + make_xp_pytest_param(stats.skew, {}), + make_xp_pytest_param(stats.skew, {'bias': False}), + make_xp_pytest_param(stats.kurtosis, {}), + make_xp_pytest_param(stats.kurtosis, {'fisher': False, 'bias': False}), + make_xp_pytest_param(stats.sem, {}), + make_xp_pytest_param(stats.sem, {'ddof': 2}), + make_xp_pytest_param(stats.kstat, {'n': 1}), + make_xp_pytest_param(stats.kstat, {'n': 2}), + make_xp_pytest_param(stats.kstat, {'n': 3}), + make_xp_pytest_param(stats.kstat, {'n': 4}), + make_xp_pytest_param(stats.kstatvar, {'n': 1}), + make_xp_pytest_param(stats.kstatvar, {'n': 2}), + make_xp_pytest_param(stats.circmean, {}), + make_xp_pytest_param(stats.circmean, {'low': 0, 'high': 1}), + make_xp_pytest_param(stats.circvar, {}), + make_xp_pytest_param(stats.circvar, {'low': 0, 'high': 1}), + make_xp_pytest_param(stats.circstd, {}), + make_xp_pytest_param(stats.circstd, {'low': 0, 'high': 1, 'normalize':True}), + make_xp_pytest_param(stats.gstd, {}), + make_xp_pytest_param(stats.gstd, {'ddof': 2}), + make_xp_pytest_param(_xp_var, {}), + make_xp_pytest_param(_xp_var, {'correction': 1, 'keepdims': True}), + make_xp_pytest_param(stats.variation, {}), + make_xp_pytest_param(stats.variation, {'ddof': 1}), + make_xp_pytest_param(stats.tmean, dict(limits=(0.1, 0.9))), + make_xp_pytest_param(stats.tvar, dict(limits=(0.1, 0.9), inclusive=(False, True))), + make_xp_pytest_param(stats.tstd, dict(limits=(0.1, 0.9), inclusive=(True, False))), + make_xp_pytest_param(stats.tsem, dict(limits=(0.1, 0.9), inclusive=(False,)*2)), + make_xp_pytest_param(stats.tmin, dict(lowerlimit=0.5, inclusive=True)), + make_xp_pytest_param(stats.tmax, dict(upperlimit=0.5, inclusive=False)), + make_xp_pytest_param(stats.iqr, {'interpolation': 'nearest'}), + make_xp_pytest_param(stats.iqr, {'rng': (10, 90), 'scale': 'normal'}), + make_xp_pytest_param(stats.median_abs_deviation, {}), + make_xp_pytest_param(stats.median_abs_deviation, {'scale': 'normal'}), + make_xp_pytest_param(stats.trim_mean, {'proportiontocut': 0.1}), + ]) +@pytest.mark.parametrize('axis', [0, 1, None]) +def test_one_sample_reducing(fun, kwargs, axis, xp): + mxp, marrays, narrays = get_arrays(1, xp=xp) + kwargs = dict(axis=axis) | kwargs + res = fun(marrays[0], **kwargs) + ref = fun(narrays[0], nan_policy='omit', **kwargs) + xp_assert_close(res.data, xp.asarray(ref)) + + +@make_xp_test_case(stats.mode) +@pytest.mark.parametrize('dtype', ['int32', 'float64']) +@pytest.mark.parametrize('shape', [10, (7, 8)]) +@pytest.mark.parametrize('axis', [0, 1, None]) +def test_mode(dtype, shape, axis, xp): + mxp, marrays, narrays = get_arrays(1, shape=shape, all_unique=False, xp=xp) + res = stats.mode(mxp.astype(marrays[0], getattr(mxp, dtype))) + ref = stats.mode(*narrays, nan_policy='omit') + xp_assert_close(res.mode.data, xp.asarray(ref.mode.astype(dtype))) + xp_assert_close(res.count.data, xp.asarray(ref.count)) + + +@make_xp_test_case(stats.describe) +@skip_backend('dask.array', reason='Arrays need `device` attribute: dask/dask#11711') +@skip_backend('jax.numpy', reason="JAX doesn't allow item assignment.") +@pytest.mark.parametrize('axis', [0, 1, None]) +@pytest.mark.parametrize('kwargs', [{}, {'bias': False}]) +def test_describe(axis, kwargs, xp): + mxp, marrays, narrays = get_arrays(1, xp=xp) + kwargs = dict(axis=axis) | kwargs + res = stats.describe(marrays[0], **kwargs) + ref = stats.describe(narrays[0], nan_policy='omit', **kwargs) + xp_assert_close(res.nobs.data, xp.asarray(ref.nobs)) + xp_assert_close(res.minmax[0].data, xp.asarray(ref.minmax[0].data)) + xp_assert_close(res.minmax[1].data, xp.asarray(ref.minmax[1].data)) + # copy reference arrays due to torch complaint about non-writeable buffer + xp_assert_close(res.variance.data, + xp.asarray(np.asarray(ref.variance.data, copy=True))) + xp_assert_close(res.skewness.data, + xp.asarray(np.asarray(ref.skewness.data, copy=True))) + xp_assert_close(res.kurtosis.data, + xp.asarray(np.asarray(ref.kurtosis.data, copy=True))) + + +@skip_backend('dask.array', reason='shape (nan,) in _xp_var when ddof=1') +@skip_backend('jax.numpy', reason="JAX doesn't allow item assignment.") +@pytest.mark.parametrize('fun', [make_xp_pytest_param(stats.zscore), + make_xp_pytest_param(stats.gzscore), + make_xp_pytest_param(stats.zmap)]) +@pytest.mark.parametrize('ddof', [0, 1]) +@pytest.mark.parametrize('axis', [0, 1, None]) +def test_one_sample_non_reducing(fun, ddof, axis, xp): + mxp, marrays, narrays = (get_arrays(2, xp=xp) if fun == stats.zmap + else get_arrays(1, xp=xp)) + res = fun(*marrays, ddof=ddof, axis=axis) + ref = xp.asarray(fun(*narrays, ddof=ddof, nan_policy='omit', axis=axis)) + xp_assert_close(res.data[~res.mask], ref[~xp.isnan(ref)]) + xp_assert_equal(res.mask, marrays[0].mask) + + +@skip_backend('dask.array', reason='Arrays need `device` attribute: dask/dask#11711') +@skip_backend('jax.numpy', reason="JAX doesn't allow item assignment.") +@pytest.mark.parametrize('f', [make_xp_pytest_param(stats.ttest_1samp), + make_xp_pytest_param(stats.ttest_rel), + make_xp_pytest_param(stats.ttest_ind)]) +@pytest.mark.parametrize('alternative', ['less', 'greater', 'two-sided']) +@pytest.mark.parametrize('axis', [0, 1, None]) +def test_ttests(f, alternative, axis, xp): + f_name = f.__name__ + mxp, marrays, narrays = get_arrays(2, xp=xp) + if f_name == 'ttest_1samp': + marrays[1] = mxp.mean(marrays[1], axis=axis, keepdims=axis is not None) + narrays[1] = np.nanmean(narrays[1], axis=axis, keepdims=axis is not None) + res = f(*marrays, alternative=alternative, axis=axis) + ref = f(*narrays, alternative=alternative, nan_policy='omit', axis=axis) + xp_assert_close(res.statistic.data, xp.asarray(ref.statistic)) + xp_assert_close(res.pvalue.data, xp.asarray(ref.pvalue)) + res_ci = res.confidence_interval() + ref_ci = ref.confidence_interval() + xp_assert_close(res_ci.low.data, xp.asarray(ref_ci.low)) + xp_assert_close(res_ci.high.data, xp.asarray(ref_ci.high)) + + +@skip_backend('dask.array', reason='Arrays need `device` attribute: dask/dask#11711') +@skip_backend('jax.numpy', reason="JAX doesn't allow item assignment.") +@pytest.mark.filterwarnings("ignore::scipy.stats._axis_nan_policy.SmallSampleWarning") +@pytest.mark.parametrize('f, args', [ + make_xp_pytest_param(stats.skewtest, tuple()), + make_xp_pytest_param(stats.kurtosistest, tuple()), + make_xp_pytest_param(stats.normaltest, tuple()), + make_xp_pytest_param(stats.jarque_bera, tuple()), + make_xp_pytest_param(stats.cramervonmises, (stats.norm.cdf,)), # type:ignore[attr-defined] +]) +@pytest.mark.parametrize('alternative', ['less', 'greater', 'two-sided']) +@pytest.mark.parametrize('axis', [0, 1, None]) +def test_goodness_of_fit(f, args, alternative, axis, xp): + mxp, marrays, narrays = get_arrays(1, xp=xp, shape=(10, 11)) + + if f in {stats.skewtest, stats.kurtosistest}: + kwds = {'alternative': alternative} + else: + if alternative != 'greater': + pytest.skip(f'str({f.__name__} does not support multiple alternatives.') + kwds = {} + + res = f(*marrays, *args, **kwds, axis=axis) + ref = f(*narrays, *args, **kwds, nan_policy='omit', axis=axis) + + xp_assert_close(res.statistic.data, xp.asarray(ref.statistic)) + xp_assert_close(res.pvalue.data, xp.asarray(ref.pvalue), atol=1e-15) + + +@skip_backend('dask.array', reason='Arrays need `device` attribute: dask/dask#11711') +@skip_backend('jax.numpy', reason="JAX doesn't allow item assignment.") +@pytest.mark.parametrize('f, lambda_', [ + make_xp_pytest_param(stats.chisquare, {}), + make_xp_pytest_param(stats.power_divergence, {'lambda_': 'pearson'}), + make_xp_pytest_param(stats.power_divergence, {'lambda_': 'log-likelihood'}), + make_xp_pytest_param(stats.power_divergence, {'lambda_': 'freeman-tukey'}), + make_xp_pytest_param(stats.power_divergence, + {'lambda_': 'mod-log-likelihood'}), + make_xp_pytest_param(stats.power_divergence, {'lambda_': 'neyman'}), + make_xp_pytest_param(stats.power_divergence, {'lambda_': 'cressie-read'}), +]) +@pytest.mark.parametrize('ddof', [0, 1]) +@pytest.mark.parametrize('axis', [0, 1, None]) +def test_power_divergence_chisquare(f, lambda_, ddof, axis, xp): + mxp, marrays, narrays = get_arrays(2, xp=xp, shape=(5, 6)) + + kwargs = dict(axis=axis, ddof=ddof) + + # test 1-arg + res = f(marrays[0], **lambda_, **kwargs) + ref = stats.power_divergence(narrays[0], nan_policy='omit', **lambda_, **kwargs) + + xp_assert_close(res.statistic.data, xp.asarray(ref[0])) + xp_assert_close(res.pvalue.data, xp.asarray(ref[1])) + + # test 2-arg + common_mask = np.isnan(narrays[0]) | np.isnan(narrays[1]) + normalize = (np.nansum(narrays[1] * ~common_mask, axis=axis, keepdims=True) + / np.nansum(narrays[0] * ~common_mask, axis=axis, keepdims=True)) + marrays[0] *= xp.asarray(normalize) + narrays[0] *= normalize + + res = f(*marrays, **lambda_, **kwargs) + ref = stats.power_divergence(*narrays, nan_policy='omit', **lambda_, **kwargs) + + xp_assert_close(res.statistic.data, xp.asarray(ref[0])) + xp_assert_close(res.pvalue.data, xp.asarray(ref[1])) + + +@make_xp_test_case(stats.combine_pvalues) +@skip_backend('jax.numpy', reason="JAX doesn't allow item assignment.") +@pytest.mark.parametrize('method', ['fisher', 'pearson', 'mudholkar_george', + 'tippett', 'stouffer']) +@pytest.mark.parametrize('axis', [0, 1, None]) +def test_combine_pvalues(method, axis, xp): + mxp, marrays, narrays = get_arrays(2, xp=xp, shape=(10, 11)) + + kwargs = dict(method=method, axis=axis) + res = stats.combine_pvalues(marrays[0], **kwargs) + ref = stats.combine_pvalues(narrays[0], nan_policy='omit', **kwargs) + + xp_assert_close(res.statistic.data, xp.asarray(ref.statistic)) + xp_assert_close(res.pvalue.data, xp.asarray(ref.pvalue)) + + if method != 'stouffer': + return + + # test method='stouffer' with weights + res = stats.combine_pvalues(marrays[0], weights=marrays[1], **kwargs) + ref = stats.combine_pvalues(narrays[0], weights=narrays[1], + nan_policy='omit', **kwargs) + + xp_assert_close(res.statistic.data, xp.asarray(ref.statistic)) + xp_assert_close(res.pvalue.data, xp.asarray(ref.pvalue)) + + +@make_xp_test_case(stats.ttest_ind_from_stats) +@skip_backend('jax.numpy', reason="JAX doesn't allow item assignment.") +@pytest.mark.parametrize("alternative", ['less', 'greater', 'two-sided']) +@pytest.mark.parametrize("equal_var", [False, True]) +def test_ttest_ind_from_stats(alternative, equal_var, xp): + shape = (10, 11) + mxp, marrays, narrays = get_arrays(6, xp=xp, shape=shape) + mask = np.sum(np.stack([np.isnan(arg) for arg in narrays]), axis=0).astype(bool) + narrays = [arg[~mask] for arg in narrays] + marrays[2], marrays[5] = marrays[2] * 100, marrays[5] * 100 + narrays[2], narrays[5] = narrays[2] * 100, narrays[5] * 100 + + kwargs = dict(alternative=alternative, equal_var=equal_var) + res = stats.ttest_ind_from_stats(*marrays, **kwargs) + ref = stats.ttest_ind_from_stats(*narrays, **kwargs) + + mask = xp.asarray(mask) + assert xp.any(mask) and xp.any(~mask) + xp_assert_close(res.statistic.data[~mask], xp.asarray(ref.statistic)) + xp_assert_close(res.pvalue.data[~mask], xp.asarray(ref.pvalue)) + xp_assert_close(res.statistic.mask, mask) + xp_assert_close(res.pvalue.mask, mask) + assert res.statistic.shape == shape + assert res.pvalue.shape == shape + + +def test_length_nonmasked_marray_iterable_axis_raises(): + xp = marray._get_namespace(np) + + data = [[1.0, 2.0], [3.0, 4.0]] + mask = [[False, False], [True, False]] + marr = xp.asarray(data, mask=mask) + + # Axis tuples are not currently supported for MArray input. + # This test can be removed after support is added. + with pytest.raises(NotImplementedError, + match="`axis` must be an integer or None for use with `MArray`"): + _count_nonmasked(marr, axis=(0, 1), xp=xp) + + +@make_xp_test_case(stats.directional_stats) +@skip_backend('jax.numpy', reason="JAX doesn't allow item assignment.") +@pytest.mark.filterwarnings("ignore::RuntimeWarning") # mdhaber/marray#120 +@pytest.mark.parametrize('normalize', [False, True]) +@pytest.mark.parametrize('axis', [0, 1]) +def test_directional_stats(normalize, axis, xp): + mxp, marrays, narrays = get_arrays(1, shape=(19, 20, 3), xp=xp) + res = stats.directional_stats(*marrays, normalize=normalize, axis=axis) + + x, = narrays + if axis == 0: + x = np.swapaxes(x, 0, 1) + + for i in range(x.shape[0]): + xi = x[i] + xi = xi[~np.any(np.isnan(xi), axis=-1), :] + ref = stats.directional_stats(xi, normalize=normalize) + xp_assert_close(res.mean_direction.data[i, ...], + xp.asarray(ref.mean_direction)) + xp_assert_close(res.mean_resultant_length.data[i], + xp.asarray(ref.mean_resultant_length)[()]) + assert not xp.any(res.mean_direction.mask) + assert not xp.any(res.mean_resultant_length.mask) + + +@make_xp_test_case(stats.wilcoxon) +@skip_backend('jax.numpy', reason="JAX doesn't allow item assignment.") +@pytest.mark.parametrize('n_samples', [1, 2]) +@pytest.mark.parametrize('zero_method', ['zsplit']) # TODO: add 'wilcox', 'pratt' +@pytest.mark.parametrize('correction', [False, True]) +@pytest.mark.parametrize('alternative', ['less', 'greater', 'two-sided']) +@pytest.mark.parametrize('method', ['asymptotic']) # TODO: add 'exact', 'auto' +@pytest.mark.parametrize('axis', [0, 1, None]) +def test_wilcoxon(n_samples, zero_method, correction, alternative, method, axis, xp): + mxp, marrays, narrays = get_arrays(n_samples, xp=xp, seed=84912165484322) + kwargs = dict(zero_method=zero_method, correction=correction, + alternative=alternative, method=method) + res = stats.wilcoxon(*marrays, axis=axis, **kwargs) + ref = stats.wilcoxon(*narrays, nan_policy='omit', axis=axis, **kwargs) + xp_assert_close(res.statistic.data, xp.asarray(ref.statistic)) + xp_assert_close(res.pvalue.data, xp.asarray(ref.pvalue)) + + +@make_xp_test_case(stats.mannwhitneyu) +@skip_backend('jax.numpy', reason="JAX doesn't allow item assignment.") +@pytest.mark.parametrize('use_continuity', [False, True]) +@pytest.mark.parametrize('alternative', ['less', 'greater', 'two-sided']) +@pytest.mark.parametrize('method', ['asymptotic']) # TODO: add 'exact', 'auto' +@pytest.mark.parametrize('axis', [0, 1, None]) +def test_mannwhitneyu(use_continuity, alternative, method, axis, xp): + mxp, marrays, narrays = get_arrays(2, xp=xp, seed=84912165484322) + kwargs = dict(use_continuity=use_continuity, alternative=alternative, method=method) + res = stats.mannwhitneyu(*marrays, axis=axis, **kwargs) + ref = stats.mannwhitneyu(*narrays, nan_policy='omit', axis=axis, **kwargs) + xp_assert_close(res.statistic.data, xp.asarray(ref.statistic)) + xp_assert_close(res.pvalue.data, xp.asarray(ref.pvalue)) + + +@skip_backend('jax.numpy', reason="JAX doesn't allow item assignment.") +@pytest.mark.parametrize('fun', [make_xp_pytest_param(stats.ks_1samp), + make_xp_pytest_param(stats.kstest)]) +@pytest.mark.parametrize('method', ['exact', 'approx', 'asymp']) # auto == exact +@pytest.mark.parametrize('alternative', ['less', 'greater', 'two-sided']) +@pytest.mark.parametrize('axis', [0, 1, None]) +def test_ks_1samp(fun, method, alternative, axis, xp): + mxp, marrays, narrays = get_arrays(1, xp=xp, seed=84912165484322) + kwargs = dict(method=method, alternative=alternative, axis=axis) + res = fun(*marrays, stats.norm.cdf, **kwargs) + ref = stats.ks_1samp(*narrays, stats.norm.cdf, nan_policy='omit', **kwargs) + xp_assert_close(res.statistic.data, xp.asarray(ref.statistic)) + xp_assert_close(res.pvalue.data, xp.asarray(ref.pvalue)) + xp_assert_equal(res.statistic_location.data, xp.asarray(ref.statistic_location)) + xp_assert_equal(res.statistic_sign.data, + xp.asarray(ref.statistic_sign, dtype=xp.int8)) + + +@skip_backend('jax.numpy', reason="JAX doesn't allow item assignment.") +@pytest.mark.parametrize('fun, kwargs', [ + make_xp_pytest_param(stats.brunnermunzel, {}), + make_xp_pytest_param(stats.brunnermunzel, {'distribution': 'normal'}), + make_xp_pytest_param(stats.cramervonmises_2samp, {'method': 'exact'}), + # make_xp_pytest_param(stats.cramervonmises_2samp, {}), # TODO: add methods +]) +@pytest.mark.parametrize('axis', [0, 1, None]) +def test_two_sample_tests(fun, kwargs, axis, xp): + if fun == stats.cramervonmises_2samp and axis is None: + pytest.skip("Sample too large for exact method.") + mxp, marrays, narrays = get_arrays(2, xp=xp, seed=84912165484322) + res = fun(*marrays, axis=axis, **kwargs) + ref = fun(*narrays, nan_policy='omit', axis=axis, **kwargs) + xp_assert_close(res.statistic.data, xp.asarray(ref.statistic)) + xp_assert_close(res.pvalue.data, xp.asarray(ref.pvalue)) + + +@make_xp_test_case(stats.ks_2samp) +@skip_backend('jax.numpy', reason="JAX doesn't allow item assignment.") +@pytest.mark.parametrize('fun', [make_xp_pytest_param(stats.ks_2samp), + make_xp_pytest_param(stats.kstest)]) +@pytest.mark.parametrize('method', ['exact', 'asymp', 'auto']) # auto == exact +@pytest.mark.parametrize('alternative', ['less', 'greater', 'two-sided']) +@pytest.mark.parametrize('axis', [0, 1, None]) +def test_ks_2samp(fun, method, alternative, axis, xp): + mxp, marrays, narrays = get_arrays(2, xp=xp, seed=84912165484322) + kwargs = dict(method=method, alternative=alternative, axis=axis) + res = fun(*marrays, **kwargs) + ref = stats.ks_2samp(*narrays, nan_policy='omit', **kwargs) + xp_assert_close(res.statistic.data, xp.asarray(ref.statistic)) + xp_assert_close(res.pvalue.data, xp.asarray(ref.pvalue)) + # with this random data, there often multiple locations where the statistic assumes + # the most extreme value, so we can't expect these to always match + # xp_assert_equal(res.statistic_location.data, xp.asarray(ref.statistic_location)) + xp_assert_equal(res.statistic_sign.data, + xp.asarray(ref.statistic_sign, dtype=xp.int8)) + + +@make_xp_test_case(stats.kendalltau) +class TestKendallTau: + @pytest.mark.parametrize('method', ['asymptotic', 'exact']) + @pytest.mark.parametrize('variant', ['b', 'c']) + @pytest.mark.parametrize('alternative', ['less', 'greater', 'two-sided']) + @pytest.mark.parametrize('axis', [0, 1, None]) + def test_omit_masked_elements(self, method, variant, alternative, axis, xp): + mxp, marrays, narrays = get_arrays(2, shape=(17, 18), + xp=xp, seed=84912165484322) + kwargs = dict(method=method, variant=variant, alternative=alternative) + res = stats.kendalltau(*marrays, **kwargs, axis=axis) + ref = stats.kendalltau(*narrays, **kwargs, nan_policy='omit', axis=axis) + xp_assert_close(res.statistic.data, xp.asarray(ref.statistic)) + xp_assert_close(res.pvalue.data, xp.asarray(ref.pvalue)) + + @pytest.mark.parametrize('nan_policy, message', [ + ('propagate', "`nan_policy='propagate'` is incompatible with MArray input."), + ('raise', 'The input contains nan values.') + ]) + def test_input_validation(self, nan_policy, message, xp): + mxp = marray._get_namespace(xp) + x = mxp.asarray([1, 2, 3, xp.nan]) + with pytest.raises(ValueError, match=message): + stats.kendalltau(x, x, nan_policy=nan_policy) + + +@skip_backend('dask.array', reason='Arrays need `device` attribute: dask/dask#11711') +@skip_backend('jax.numpy', reason="JAX doesn't allow item assignment.") +@pytest.mark.parametrize('fun, kwargs', [ + make_xp_pytest_param(stats.bartlett, {}), + make_xp_pytest_param(stats.alexandergovern, {}), + make_xp_pytest_param(stats.levene, {'center': 'median'}), + make_xp_pytest_param(stats.levene, {'center': 'mean'}), + make_xp_pytest_param(stats.levene, {'center': 'trimmed'}), + make_xp_pytest_param(stats.f_oneway, {'equal_var': True}), + make_xp_pytest_param(stats.f_oneway, {'equal_var': False}), + make_xp_pytest_param(stats.kruskal, {}), + make_xp_pytest_param(stats.fligner, {'center': 'median'}), + make_xp_pytest_param(stats.fligner, {'center': 'mean'}), + make_xp_pytest_param(stats.fligner, {'center': 'trimmed'}), +]) +@pytest.mark.parametrize('axis', [0, 1, None]) +def test_k_sample_tests(fun, kwargs, axis, xp): + mxp, marrays, narrays = get_arrays(3, xp=xp) + res = fun(*marrays, axis=axis, **kwargs) + ref = fun(*narrays, nan_policy='omit', axis=axis, **kwargs) + xp_assert_close(res.statistic.data, xp.asarray(ref.statistic)) + xp_assert_close(res.pvalue.data, xp.asarray(ref.pvalue)) + + +@skip_backend('jax.numpy', reason="JAX currently incompatible with marray") +@pytest.mark.parametrize('fun, kwargs', [ + make_xp_pytest_param(stats.friedmanchisquare, {}), +]) +@pytest.mark.parametrize('axis', [0, 1, None]) +def test_k_sample_paired_tests(fun, kwargs, axis, xp): + mxp, marrays, narrays = get_arrays(3, shape=(8, 9), xp=xp) + res = fun(*marrays, axis=axis, **kwargs) + ref = fun(*narrays, nan_policy='omit', axis=axis, **kwargs) + xp_assert_close(res.statistic.data, xp.asarray(ref.statistic)) + xp_assert_close(res.pvalue.data, xp.asarray(ref.pvalue)) + + +@skip_backend('dask.array', reason='Arrays need `device` attribute: dask/dask#11711') +@skip_backend('jax.numpy', reason="JAX doesn't allow item assignment.") +@pytest.mark.parametrize('f', [ + make_xp_pytest_param(stats.pearsonr), + make_xp_pytest_param(stats.pointbiserialr), + make_xp_pytest_param(stats.spearmanrho), +]) +@pytest.mark.parametrize('alternative', ['less', 'greater', 'two-sided']) +@pytest.mark.parametrize('axis', [0, 1, None]) +def test_correlation(f, alternative, axis, xp): + mxp, marrays, narrays = get_arrays(2, xp=xp) + + kwargs = {} if f == stats.pointbiserialr else {'alternative': alternative} + res = f(*marrays, **kwargs, axis=axis) + + # `pearsonr` does not have `axis_nan_policy`, so do this manually + x, y = narrays + if axis == 0: + x, y = x.T, y.T + elif axis is None: + x, y = x.ravel()[np.newaxis, :], y.ravel()[np.newaxis, :] + + for i in range(x.shape[0]): + xi, yi = x[i, ...], y[i, ...] + i = () if axis is None else i + + mask = np.isnan(xi) | np.isnan(yi) + ref = f(xi[~mask], yi[~mask], **kwargs) + + atol = 1e-7 if (is_torch(xp) and f == stats.spearmanrho) else 0. + xp_assert_close(res.statistic.data[i], xp.asarray(ref.statistic)[()], atol=atol) + xp_assert_close(res.pvalue.data[i], xp.asarray(ref.pvalue)[()], atol=atol) + + if f == stats.pearsonr: + res_ci_low, res_ci_high = res.confidence_interval() + ref_ci_low, ref_ci_high = ref.confidence_interval() + xp_assert_close(res_ci_low.data[i], xp.asarray(ref_ci_low)[()]) + xp_assert_close(res_ci_high.data[i], xp.asarray(ref_ci_high)[()]) + + +@skip_backend('jax.numpy', reason="JAX doesn't allow item assignment.") +@skip_backend('array_api_strict', reason="issue with clip.") +@pytest.mark.parametrize('f, method', [ + make_xp_pytest_param(stats.siegelslopes, {'method':'hierarchical'}), + make_xp_pytest_param(stats.siegelslopes, {'method':'separate'}), + make_xp_pytest_param(stats.theilslopes, {'method': 'joint'}), + make_xp_pytest_param(stats.theilslopes, {'method': 'separate'}), + make_xp_pytest_param(stats.linregress, {'alternative': 'less'}), + make_xp_pytest_param(stats.linregress, {'alternative': 'greater'}), + make_xp_pytest_param(stats.linregress, {'alternative': 'two-sided'}), +]) +@pytest.mark.parametrize('axis', [0, 1, None]) +def test_slopes_intercepts(f, method, axis, xp): + mxp, marrays, narrays = get_arrays(2, shape=(19, 20), seed=84912165484320, xp=xp) + res = f(*marrays, axis=axis, **method) + ref = f(*narrays, nan_policy='omit', axis=axis, **method) + + xp_assert_close(res.slope.data, xp.asarray(ref.slope)) + xp_assert_close(res.intercept.data, xp.asarray(ref.intercept)) + + if f == stats.theilslopes: + xp_assert_close(res.low_slope.data, xp.asarray(ref.low_slope)) + xp_assert_close(res.high_slope.data, xp.asarray(ref.high_slope)) + elif f == stats.linregress: + xp_assert_close(res.rvalue.data, xp.asarray(ref.rvalue)) + xp_assert_close(res.pvalue.data, xp.asarray(ref.pvalue)) + xp_assert_close(res.stderr.data, xp.asarray(ref.stderr)) + xp_assert_close(res.intercept_stderr.data, xp.asarray(ref.intercept_stderr)) + + +@make_xp_test_case(stats.entropy) +@skip_backend('dask.array', reason="Dask doesn't like the /= when base is not None.") +@skip_backend('jax.numpy', reason="JAX doesn't allow item assignment.") +@pytest.mark.parametrize('qk', [False, True]) +@pytest.mark.parametrize('base', [None, 2]) +@pytest.mark.parametrize('axis', [0, 1, None]) +def test_entropy(qk, base, axis, xp): + mxp, marrays, narrays = get_arrays(2 if qk else 1, xp=xp) + res = stats.entropy(*marrays, base=base, axis=axis) + ref = stats.entropy(*narrays, base=base, nan_policy='omit', axis=axis) + xp_assert_close(res.data, xp.asarray(ref)) + + +@make_xp_test_case(stats.rankdata) +@skip_backend('jax.numpy', reason="JAX currently incompatible with marray") +@pytest.mark.parametrize('method', ['average', 'min', 'max', 'dense', 'ordinal']) +@pytest.mark.parametrize('axis', [0, 1, None]) +def test_rankdata(method, axis, xp): + mxp, marrays, narrays = get_arrays(1, xp=xp) + res = stats.rankdata(*marrays, method=method, axis=axis) + ref = stats.rankdata(*narrays, method=method, nan_policy='omit', axis=axis) + xp_assert_close(res.data[~res.mask], xp.asarray(ref[~np.isnan(ref)])) + xp_assert_close(res.mask, xp.asarray(np.isnan(ref))) + + +@make_xp_test_case(stats.obrientransform) +@skip_backend('jax.numpy', reason="JAX currently incompatible with marray") +@pytest.mark.parametrize('dtype', ['float32', 'float64']) +@pytest.mark.parametrize('n_arrays', [1, 3]) +def test_obrientransform(dtype, n_arrays, xp): + # obrientransform is not yet vectorized + mxp, marrays, narrays = get_arrays(n_arrays, dtype=dtype, shape=(25,), xp=xp) + res = stats.obrientransform(*marrays) + ref = stats.obrientransform(*narrays, nan_policy='omit') + for res_i, ref_i in zip(res, ref): + xp_assert_close(res_i.data[~res_i.mask], + xp.asarray(ref_i[~np.isnan(ref_i)], dtype=getattr(xp, dtype))) diff --git a/v12x/pylib/scipy/stats/tests/test_mgc.py b/v12x/pylib/scipy/stats/tests/test_mgc.py new file mode 100644 index 0000000000000000000000000000000000000000..fd4d27d7c1647b0ab72cd668ccd0509f02b4684a --- /dev/null +++ b/v12x/pylib/scipy/stats/tests/test_mgc.py @@ -0,0 +1,233 @@ +import pytest +from pytest import raises as assert_raises, warns as assert_warns + +import numpy as np +from numpy.testing import assert_approx_equal, assert_allclose, assert_equal + +from scipy.spatial.distance import cdist +from scipy import stats + +class TestMGCErrorWarnings: + """ Tests errors and warnings derived from MGC. + """ + def test_error_notndarray(self): + # raises error if x or y is not an ndarray + x = np.arange(20) + y = [5] * 20 + assert_raises(ValueError, stats.multiscale_graphcorr, x, y) + assert_raises(ValueError, stats.multiscale_graphcorr, y, x) + + def test_error_shape(self): + # raises error if number of samples different (n) + x = np.arange(100).reshape(25, 4) + y = x.reshape(10, 10) + assert_raises(ValueError, stats.multiscale_graphcorr, x, y) + + def test_error_lowsamples(self): + # raises error if samples are low (< 3) + x = np.arange(3) + y = np.arange(3) + assert_raises(ValueError, stats.multiscale_graphcorr, x, y) + + def test_error_nans(self): + # raises error if inputs contain NaNs + x = np.arange(20, dtype=float) + x[0] = np.nan + assert_raises(ValueError, stats.multiscale_graphcorr, x, x) + + y = np.arange(20) + assert_raises(ValueError, stats.multiscale_graphcorr, x, y) + + def test_error_wrongdisttype(self): + # raises error if metric is not a function + x = np.arange(20) + compute_distance = 0 + assert_raises(ValueError, stats.multiscale_graphcorr, x, x, + compute_distance=compute_distance) + + @pytest.mark.parametrize("reps", [ + -1, # reps is negative + '1', # reps is not integer + ]) + def test_error_reps(self, reps): + # raises error if reps is negative + x = np.arange(20) + assert_raises(ValueError, stats.multiscale_graphcorr, x, x, reps=reps) + + def test_warns_reps(self): + # raises warning when reps is less than 1000 + x = np.arange(20) + reps = 100 + assert_warns(RuntimeWarning, stats.multiscale_graphcorr, x, x, reps=reps) + + def test_error_infty(self): + # raises error if input contains infinities + x = np.arange(20) + y = np.ones(20) * np.inf + assert_raises(ValueError, stats.multiscale_graphcorr, x, y) + + +class TestMGCStat: + """ Test validity of MGC test statistic + """ + def setup_method(self): + self.rng = np.random.default_rng(1266219746) + + def _simulations(self, samps=100, dims=1, sim_type="", rng=None): + rng = rng or self.rng + # linear simulation + if sim_type == "linear": + x = rng.uniform(-1, 1, size=(samps, 1)) + y = x + 0.3 * rng.random(size=(x.size, 1)) + + # spiral simulation + elif sim_type == "nonlinear": + unif = np.array(self.rng.uniform(0, 5, size=(samps, 1))) + x = unif * np.cos(np.pi * unif) + y = (unif * np.sin(np.pi * unif) + + 0.4*self.rng.random(size=(x.size, 1))) + + # independence (tests type I simulation) + elif sim_type == "independence": + u = self.rng.standard_normal(size=(samps, 1)) + v = self.rng.standard_normal(size=(samps, 1)) + u_2 = self.rng.binomial(1, p=0.5, size=(samps, 1)) + v_2 = self.rng.binomial(1, p=0.5, size=(samps, 1)) + x = u/3 + 2*u_2 - 1 + y = v/3 + 2*v_2 - 1 + + # raises error if not approved sim_type + else: + raise ValueError("sim_type must be linear, nonlinear, or " + "independence") + + # add dimensions of noise for higher dimensions + if dims > 1: + dims_noise = self.rng.standard_normal(size=(samps, dims-1)) + x = np.concatenate((x, dims_noise), axis=1) + + return x, y + + @pytest.mark.xslow + @pytest.mark.parametrize("sim_type, obs_stat, obs_pvalue", [ + # Reference values produced by `multiscale_graphcorr` -> + # this only guards against unintended changes + ("linear", 0.970137188683, 0.000999000999), # test linear simulation + ("nonlinear", 0.149843048137, 0.000999000999), # test spiral simulation + ("independence", -0.003965353120, 0.4675324675) # test independence simulation + ]) + def test_oned(self, sim_type, obs_stat, obs_pvalue): + # generate x and y + rng = np.random.default_rng(8157117705) + x, y = self._simulations(samps=100, dims=1, sim_type=sim_type, rng=rng) + + # test stat and pvalue + stat, pvalue, _ = stats.multiscale_graphcorr(x, y, random_state=rng) + assert_allclose(stat, obs_stat) + assert_allclose(pvalue, obs_pvalue) + + @pytest.mark.xslow + @pytest.mark.parametrize("sim_type, obs_stat, obs_pvalue", [ + # Reference values produced by `multiscale_graphcorr` -> + # this only guards against unintended changes + ("linear", 0.193063672972, 0.000999000999), # test linear simulation + ("nonlinear", 0.010042844939, 0.216783216783), # test spiral simulation + ]) + def test_fived(self, sim_type, obs_stat, obs_pvalue): + # generate x and y + rng = np.random.default_rng(8157117705) + x, y = self._simulations(samps=100, dims=5, sim_type=sim_type, rng=rng) + + # test stat and pvalue + stat, pvalue, _ = stats.multiscale_graphcorr(x, y, random_state=rng) + assert_allclose(stat, obs_stat) + assert_allclose(pvalue, obs_pvalue) + + @pytest.mark.xslow + def test_twosamp(self): + # generate x and y + x = self.rng.binomial(100, 0.5, size=(100, 5)) + y = self.rng.standard_normal(size=(80, 5)) + + # test stat and pvalue + stat, pvalue, _ = stats.multiscale_graphcorr(x, y, random_state=self.rng) + assert_approx_equal(stat, 1.0, significant=1) + assert_approx_equal(pvalue, 0.001, significant=1) + + # generate x and y + y = self.rng.standard_normal(size=(100, 5)) + + # test stat and pvalue + stat, pvalue, _ = stats.multiscale_graphcorr(x, y, is_twosamp=True, + random_state=self.rng) + assert_approx_equal(stat, 1.0, significant=1) + assert_approx_equal(pvalue, 0.001, significant=1) + + @pytest.mark.xslow + def test_workers(self): + # generate x and y + x, y = self._simulations(samps=100, dims=1, sim_type="linear") + + # test stat and pvalue + stat, pvalue, _ = stats.multiscale_graphcorr(x, y, workers=2, + random_state=self.rng) + assert_approx_equal(stat, 0.97, significant=1) + assert_approx_equal(pvalue, 0.001, significant=1) + + @pytest.mark.xslow + def test_random_state(self): + # generate x and y + x, y = self._simulations(samps=100, dims=1, sim_type="linear") + + # test stat and pvalue + stat, pvalue, _ = stats.multiscale_graphcorr(x, y, random_state=1) + assert_approx_equal(stat, 0.97, significant=1) + assert_approx_equal(pvalue, 0.001, significant=1) + + @pytest.mark.xslow + def test_dist_perm(self): + # generate x and y + x, y = self._simulations(samps=100, dims=1, sim_type="nonlinear") + distx = cdist(x, x, metric="euclidean") + disty = cdist(y, y, metric="euclidean") + + stat_dist, pvalue_dist, _ = stats.multiscale_graphcorr(distx, disty, + compute_distance=None, + random_state=1) + assert_approx_equal(stat_dist, 0.163, significant=1) + assert_approx_equal(pvalue_dist, 0.001, significant=1) + + @pytest.mark.fail_slow(20) # all other tests are XSLOW; we need at least one to run + @pytest.mark.slow + def test_pvalue_literature(self): + # generate x and y + x, y = self._simulations(samps=100, dims=1, sim_type="linear") + + # test stat and pvalue + _, pvalue, _ = stats.multiscale_graphcorr(x, y, random_state=1) + assert_allclose(pvalue, 1/1001) + + def test_zero_distance_matrix_index_error(self): + # regression test for zero distance + # matrix IndexError issue; see gh-19769 + n = 6 + x = np.arange(n).reshape(-1, 1) + + # valid distance matrix + distx = cdist(x, x) + + # degenerate all-zero distance matrix + disty = np.zeros((n, n)) + + # test stat and pvalue + res = stats.multiscale_graphcorr(distx, disty, compute_distance=None) + assert_allclose(res.statistic, 0.0) + assert_allclose(res.pvalue, 1.0) + + @pytest.mark.xslow + def test_alias(self): + # generate x and y + x, y = self._simulations(samps=100, dims=1, sim_type="linear") + + res = stats.multiscale_graphcorr(x, y, random_state=1) + assert_equal(res.stat, res.statistic) diff --git a/v12x/pylib/scipy/stats/tests/test_morestats.py b/v12x/pylib/scipy/stats/tests/test_morestats.py new file mode 100644 index 0000000000000000000000000000000000000000..070be8f03d96c62e08364f4e573d7cf2c4705dc2 --- /dev/null +++ b/v12x/pylib/scipy/stats/tests/test_morestats.py @@ -0,0 +1,3916 @@ +# Author: Travis Oliphant, 2002 +# +# Further enhancements and tests added by numerous SciPy developers. +# +import math +import re +import sys +import warnings +from functools import partial + +import numpy as np +from numpy.random import RandomState +from numpy.testing import (assert_array_equal, assert_almost_equal, + assert_array_less, assert_array_almost_equal, + assert_, assert_allclose, assert_equal) +import pytest +from pytest import raises as assert_raises + +from scipy import optimize, stats, special +from scipy.stats._morestats import (_abw_state, _get_As_weibull, _Avals_weibull, + _yeojohnson_transform) +from .common_tests import check_named_results +from .._hypotests import _get_wilcoxon_distr, _get_wilcoxon_distr2 +from scipy.stats._ansari_swilk_statistics import swilk +from scipy.stats._binomtest import _binary_search_for_binom_tst +from scipy.stats._distr_params import distcont +from scipy.stats._axis_nan_policy import (SmallSampleWarning, too_small_nd_omit, + too_small_1d_omit, too_small_1d_not_omit) + +import scipy._external.array_api_extra as xpx +from scipy._lib._array_api import (is_torch, make_xp_test_case, eager_warns, xp_ravel, + is_numpy, xp_default_dtype, is_array_api_strict, + is_jax, is_lazy_array) +from scipy._lib._array_api_no_0d import ( + xp_assert_close, + xp_assert_equal, + xp_assert_less, +) + +lazy_xp_modules = [stats] +skip_xp_backends = pytest.mark.skip_xp_backends + +distcont = dict(distcont) # type: ignore + +# Matplotlib is not a scipy dependency but is optionally used in probplot, so +# check if it's available +try: + import matplotlib + matplotlib.rcParams['backend'] = 'Agg' + import matplotlib.pyplot as plt + have_matplotlib = True +except Exception: + have_matplotlib = False + + +# test data gear.dat from NIST for Levene and Bartlett test +# https://www.itl.nist.gov/div898/handbook/eda/section3/eda3581.htm +g1 = [1.006, 0.996, 0.998, 1.000, 0.992, 0.993, 1.002, 0.999, 0.994, 1.000] +g2 = [0.998, 1.006, 1.000, 1.002, 0.997, 0.998, 0.996, 1.000, 1.006, 0.988] +g3 = [0.991, 0.987, 0.997, 0.999, 0.995, 0.994, 1.000, 0.999, 0.996, 0.996] +g4 = [1.005, 1.002, 0.994, 1.000, 0.995, 0.994, 0.998, 0.996, 1.002, 0.996] +g5 = [0.998, 0.998, 0.982, 0.990, 1.002, 0.984, 0.996, 0.993, 0.980, 0.996] +g6 = [1.009, 1.013, 1.009, 0.997, 0.988, 1.002, 0.995, 0.998, 0.981, 0.996] +g7 = [0.990, 1.004, 0.996, 1.001, 0.998, 1.000, 1.018, 1.010, 0.996, 1.002] +g8 = [0.998, 1.000, 1.006, 1.000, 1.002, 0.996, 0.998, 0.996, 1.002, 1.006] +g9 = [1.002, 0.998, 0.996, 0.995, 0.996, 1.004, 1.004, 0.998, 0.999, 0.991] +g10 = [0.991, 0.995, 0.984, 0.994, 0.997, 0.997, 0.991, 0.998, 1.004, 0.997] + + +# The loggamma RVS stream is changing due to gh-13349; this version +# preserves the old stream so that tests don't change. +def _old_loggamma_rvs(*args, **kwargs): + return np.log(stats.gamma.rvs(*args, **kwargs)) + + +class TestBayes_mvs: + def test_basic(self): + # Expected values in this test simply taken from the function. For + # some checks regarding correctness of implementation, see review in + # gh-674 + data = [6, 9, 12, 7, 8, 8, 13] + mean, var, std = stats.bayes_mvs(data) + assert_almost_equal(mean.statistic, 9.0) + assert_allclose(mean.minmax, (7.103650222492964, 10.896349777507034), + rtol=1e-6) + + assert_almost_equal(var.statistic, 10.0) + assert_allclose(var.minmax, (3.1767242068607087, 24.45910381334018), + rtol=1e-09) + + assert_almost_equal(std.statistic, 2.9724954732045084, decimal=14) + assert_allclose(std.minmax, (1.7823367265645145, 4.9456146050146312), + rtol=1e-14) + + def test_empty_input(self): + assert_raises(ValueError, stats.bayes_mvs, []) + + def test_result_attributes(self): + x = np.arange(15) + attributes = ('statistic', 'minmax') + res = stats.bayes_mvs(x) + + for i in res: + check_named_results(i, attributes) + + +class TestMvsdist: + def test_basic(self): + data = [6, 9, 12, 7, 8, 8, 13] + mean, var, std = stats.mvsdist(data) + assert_almost_equal(mean.mean(), 9.0) + assert_allclose(mean.interval(0.9), (7.103650222492964, + 10.896349777507034), rtol=1e-14) + + assert_almost_equal(var.mean(), 10.0) + assert_allclose(var.interval(0.9), (3.1767242068607087, + 24.45910381334018), rtol=1e-09) + + assert_almost_equal(std.mean(), 2.9724954732045084, decimal=14) + assert_allclose(std.interval(0.9), (1.7823367265645145, + 4.9456146050146312), rtol=1e-14) + + def test_empty_input(self): + assert_raises(ValueError, stats.mvsdist, []) + + def test_bad_arg(self): + # Raise ValueError if fewer than two data points are given. + data = [1] + assert_raises(ValueError, stats.mvsdist, data) + + def test_warns(self): + # regression test for gh-5270 + # make sure there are no spurious divide-by-zero warnings + with warnings.catch_warnings(): + warnings.simplefilter('error', RuntimeWarning) + [x.mean() for x in stats.mvsdist([1, 2, 3])] + [x.mean() for x in stats.mvsdist([1, 2, 3, 4, 5])] + + +@make_xp_test_case(stats.shapiro) +class TestShapiro: + def test_basic(self, xp): + # reference values generated using R shapiro.test, e.g. + # options(digits=16) + # x = c(0.11, 7.87, 4.61, 10.14, 7.95, 3.14, 0.46, + # 4.43, 0.21, 4.75, 0.71, 1.52, 3.24, + # 0.93, 0.42, 4.97, 9.53, 4.55, 0.47, 6.66) + # shapiro.test(x) + x1 = [0.11, 7.87, 4.61, 10.14, 7.95, 3.14, 0.46, + 4.43, 0.21, 4.75, 0.71, 1.52, 3.24, + 0.93, 0.42, 4.97, 9.53, 4.55, 0.47, 6.66] + shapiro_test = stats.shapiro(xp.asarray(x1)) + w, pw = shapiro_test + xp_assert_close(w, xp.asarray(0.90047287931756)) + xp_assert_close(shapiro_test.statistic, xp.asarray(0.90047287931756)) + xp_assert_close(pw, xp.asarray(0.04208957522226)) + xp_assert_close(shapiro_test.pvalue, xp.asarray(0.04208957522226)) + + x2 = [1.36, 1.14, 2.92, 2.55, 1.46, 1.06, 5.27, -1.11, + 3.48, 1.10, 0.88, -0.51, 1.46, 0.52, 6.20, 1.69, + 0.08, 3.67, 2.81, 3.49] + shapiro_test = stats.shapiro(xp.asarray(x2)) + xp_assert_close(shapiro_test.statistic, xp.asarray(0.95902694603234)) + xp_assert_close(shapiro_test.pvalue, xp.asarray(0.5245979304707)) + + x3 = stats.norm.rvs(loc=5, scale=3, size=100, random_state=12345678) + shapiro_test = stats.shapiro(xp.asarray(x3.tolist())) + xp_assert_close(shapiro_test.statistic, xp.asarray(0.97728027037175)) + xp_assert_close(shapiro_test.pvalue, xp.asarray(0.08143656270016)) + + @pytest.mark.parametrize('dtype', [None, 'float32', 'float64']) + def test_basic2(self, dtype, xp): + # # Data from original paper; additional precision from R shapiro.test + x4 = [0.139, 0.157, 0.175, 0.256, 0.344, 0.413, 0.503, 0.577, 0.614, + 0.655, 0.954, 1.392, 1.557, 1.648, 1.690, 1.994, 2.174, 2.206, + 3.245, 3.510, 3.571, 4.354, 4.980, 6.084, 8.351] + dtype = {'dtype': dtype if dtype is None else getattr(xp, dtype)} + W_expected = xp.asarray(0.83466627531817, **dtype) + p_expected = xp.asarray(0.000913490481813, **dtype) + shapiro_test = stats.shapiro(xp.asarray(x4, **dtype)) + xp_assert_close(shapiro_test.statistic, xp.asarray(W_expected, **dtype)) + xp_assert_close(shapiro_test.pvalue, xp.asarray(p_expected, **dtype)) + + @skip_xp_backends('jax.numpy', reason='lazy backend -> limited input validation') + @skip_xp_backends('dask.array', reason='lazy backend -> limited input validation') + @pytest.mark.parametrize('x', ([], [1], [1, 2])) + def test_not_enough_values(self, x, xp): + with eager_warns(SmallSampleWarning, match=too_small_1d_not_omit, xp=xp): + res = stats.shapiro(xp.asarray(x)) + xp_assert_equal(res.statistic, xp.asarray(xp.nan)) + xp_assert_equal(res.pvalue, xp.asarray(xp.nan)) + + def test_nan_input(self, xp): + x = xp.arange(10.) + x = xpx.at(x)[9].set(xp.nan) + + shapiro_test = stats.shapiro(xp.asarray(x)) + xp_assert_equal(shapiro_test.statistic, xp.asarray(xp.nan)) + xp_assert_equal(shapiro_test.pvalue, xp.asarray(xp.nan)) + + def test_gh14462(self, xp): + # shapiro is theoretically location-invariant, but when the magnitude + # of the values is much greater than the variance, there can be + # numerical issues. Fixed by subtracting median from the data. + # See gh-14462. + trans_val, maxlog = stats.boxcox([122500, 474400, 110400]) + dtype = xp.float64 # we need some precision to work with + res = stats.shapiro(xp.asarray(trans_val, dtype=dtype)) + + # Reference from R: + # options(digits=16) + # x = c(0.00000000e+00, 3.39996924e-08, -6.35166875e-09) + # shapiro.test(x) + ref_statistic = xp.asarray(0.86468431705371, dtype=dtype) + ref_pvalue = xp.asarray(0.2805581751566, dtype=dtype) + + xp_assert_close(res.statistic, ref_statistic, rtol=5e-7) + xp_assert_close(res.pvalue, ref_pvalue, rtol=5e-7) + + def test_length_3_gh18322(self, xp): + # gh-18322 reported that the p-value could be negative for input of + # length 3. Check that this is resolved. + res = stats.shapiro(xp.asarray([0.6931471805599453, 0.0, 0.0])) + assert res.pvalue >= 0 + + # R `shapiro.test` doesn't produce an accurate p-value in the case + # above. Check that the formula used in `stats.shapiro` is not wrong. + # options(digits=16) + # x = c(-0.7746653110021126, -0.4344432067942129, 1.8157053280290931) + # shapiro.test(x) + x = [-0.7746653110021126, -0.4344432067942129, 1.8157053280290931] + res = stats.shapiro(xp.asarray(x)) + xp_assert_close(res.statistic, xp.asarray(0.84658770645509)) + xp_assert_close(res.pvalue, xp.asarray(0.2313666489882)) + + @pytest.mark.parametrize('n', [3, 4, 5, 8, 11, 20, 100, 200, 500, 1000, 2000, 5000]) + def test_against_as181(self, n, xp): + rng = np.random.default_rng(n) + x = rng.standard_normal(n) + + a = np.zeros(n // 2, dtype=np.float64) + y = np.sort(x) - x[n // 2] + ref_statistic, ref_pvalue, _ = swilk(y, a, 0) + + x = xp.asarray(x) + res = stats.shapiro(x) + ref_statistic = xp.asarray(ref_statistic, dtype=x.dtype) + ref_pvalue = xp.asarray(ref_pvalue, dtype=x.dtype) + + xp_assert_close(res.statistic, ref_statistic, rtol=5e-7) + xp_assert_close(res.pvalue, ref_pvalue, rtol=5e-7) + + +@pytest.mark.filterwarnings("ignore: As of SciPy 1.17: FutureWarning") +class TestAnderson: + def test_normal(self): + rs = RandomState(1234567890) + x1 = rs.standard_exponential(size=50) + x2 = rs.standard_normal(size=50) + A, crit, sig = stats.anderson(x1) + assert_array_less(crit[:-1], A) + A, crit, sig = stats.anderson(x2) + assert_array_less(A, crit[-2:]) + + v = np.ones(10) + v[0] = 0 + A, crit, sig = stats.anderson(v) + # The expected statistic 3.208057 was computed independently of scipy. + # For example, in R: + # > library(nortest) + # > v <- rep(1, 10) + # > v[1] <- 0 + # > result <- ad.test(v) + # > result$statistic + # A + # 3.208057 + assert_allclose(A, 3.208057) + + def test_expon(self): + rs = RandomState(1234567890) + x1 = rs.standard_exponential(size=50) + x2 = rs.standard_normal(size=50) + A, crit, sig = stats.anderson(x1, 'expon') + assert_array_less(A, crit[-2:]) + with np.errstate(all='ignore'): + A, crit, sig = stats.anderson(x2, 'expon') + assert_(A > crit[-1]) + + def test_gumbel(self): + # Regression test for gh-6306. Before that issue was fixed, + # this case would return a2=inf. + v = np.ones(100) + v[0] = 0.0 + a2, crit, sig = stats.anderson(v, 'gumbel') + # A brief reimplementation of the calculation of the statistic. + n = len(v) + xbar, s = stats.gumbel_l.fit(v) + logcdf = stats.gumbel_l.logcdf(v, xbar, s) + logsf = stats.gumbel_l.logsf(v, xbar, s) + i = np.arange(1, n+1) + expected_a2 = -n - np.mean((2*i - 1) * (logcdf + logsf[::-1])) + + assert_allclose(a2, expected_a2) + + def test_bad_arg(self): + assert_raises(ValueError, stats.anderson, [1], dist='plate_of_shrimp') + + def test_result_attributes(self): + rs = RandomState(1234567890) + x = rs.standard_exponential(size=50) + res = stats.anderson(x) + attributes = ('statistic', 'critical_values', 'significance_level') + check_named_results(res, attributes) + + def test_gumbel_l(self): + # gh-2592, gh-6337 + # Adds support to 'gumbel_r' and 'gumbel_l' as valid inputs for dist. + rs = RandomState(1234567890) + x = rs.gumbel(size=100) + A1, crit1, sig1 = stats.anderson(x, 'gumbel') + A2, crit2, sig2 = stats.anderson(x, 'gumbel_l') + + assert_allclose(A2, A1) + + def test_gumbel_r(self): + # gh-2592, gh-6337 + # Adds support to 'gumbel_r' and 'gumbel_l' as valid inputs for dist. + rs = RandomState(1234567890) + x1 = rs.gumbel(size=100) + x2 = np.ones(100) + # A constant array is a degenerate case and breaks gumbel_r.fit, so + # change one value in x2. + x2[0] = 0.996 + A1, crit1, sig1 = stats.anderson(x1, 'gumbel_r') + A2, crit2, sig2 = stats.anderson(x2, 'gumbel_r') + + assert_array_less(A1, crit1[-2:]) + assert_(A2 > crit2[-1]) + + def test_weibull_min_case_A(self): + # data and reference values from `anderson` reference [7] + x = np.array([225, 171, 198, 189, 189, 135, 162, 135, 117, 162]) + res = stats.anderson(x, 'weibull_min') + m, loc, scale = res.fit_result.params + assert_allclose((m, loc, scale), (2.38, 99.02, 78.23), rtol=2e-3) + assert_allclose(res.statistic, 0.260, rtol=1e-3) + assert res.statistic < res.critical_values[0] + + c = 1 / m # ~0.42 + assert_allclose(c, 1/2.38, rtol=2e-3) + # interpolate between rows for c=0.4 and c=0.45, indices -3 and -2 + As40 = _Avals_weibull[-3] + As45 = _Avals_weibull[-2] + As_ref = As40 + (c - 0.4)/(0.45 - 0.4) * (As45 - As40) + # atol=1e-3 because results are rounded up to the next third decimal + assert np.all(res.critical_values > As_ref) + assert_allclose(res.critical_values, As_ref, atol=1e-3) + + def test_weibull_min_case_B(self): + # From `anderson` reference [7] + x = np.array([74, 57, 48, 29, 502, 12, 70, 21, + 29, 386, 59, 27, 153, 26, 326]) + message = "Maximum likelihood estimation has converged to " + with pytest.raises(ValueError, match=message): + stats.anderson(x, 'weibull_min') + + def test_weibull_warning_error(self): + # Check for warning message when there are too few observations + # This is also an example in which an error occurs during fitting + x = -np.array([225, 75, 57, 168, 107, 12, 61, 43, 29]) + wmessage = "Critical values of the test statistic are given for the..." + emessage = "An error occurred while fitting the Weibull distribution..." + wcontext = pytest.warns(UserWarning, match=wmessage) + econtext = pytest.raises(ValueError, match=emessage) + with wcontext, econtext: + stats.anderson(x, 'weibull_min') + + @pytest.mark.parametrize('distname', + ['norm', 'expon', 'gumbel_l', 'extreme1', + 'gumbel', 'gumbel_r', 'logistic', 'weibull_min']) + def test_anderson_fit_params(self, distname): + # check that anderson now returns a FitResult + rng = np.random.default_rng(330691555377792039) + real_distname = ('gumbel_l' if distname in {'extreme1', 'gumbel'} + else distname) + dist = getattr(stats, real_distname) + params = distcont[real_distname] + x = dist.rvs(*params, size=1000, random_state=rng) + res = stats.anderson(x, distname) + assert res.fit_result.success + + def test_anderson_weibull_As(self): + m = 1 # "when mi < 2, so that c > 0.5, the last line...should be used" + assert_equal(_get_As_weibull(1/m), _Avals_weibull[-1]) + m = np.inf + assert_equal(_get_As_weibull(1/m), _Avals_weibull[0]) + + +class TestAndersonMethod: + def test_warning(self): + message = "As of SciPy 1.17, users..." + with pytest.warns(FutureWarning, match=message): + stats.anderson([1, 2, 3], 'norm') + + def test_method_input_validation(self): + message = "`method` must be either..." + with pytest.raises(ValueError, match=message): + stats.anderson([1, 2, 3], 'norm', method='ekki-ekki') + + def test_monte_carlo_method(self): + rng = np.random.default_rng(94982389149239) + + message = "The `rvs` attribute..." + with pytest.warns(UserWarning, match=message): + method = stats.MonteCarloMethod(rvs=rng.random) + stats.anderson([1, 2, 3], 'norm', method=method) + + message = "The `batch` attribute..." + with pytest.warns(UserWarning, match=message): + method = stats.MonteCarloMethod(batch=10) + stats.anderson([1, 2, 3], 'norm', method=method) + + method = stats.MonteCarloMethod(n_resamples=9, rng=rng) + res = stats.anderson([1, 2, 3], 'norm', method=method) + ten_p = res.pvalue * 10 + # p-value will always be divisible by n_resamples + 1 + assert np.round(ten_p) == ten_p + + method = stats.MonteCarloMethod(rng=np.random.default_rng(23495984827)) + ref = stats.anderson([1, 2, 3, 4, 5], 'norm', method=method) + method = stats.MonteCarloMethod(rng=np.random.default_rng(23495984827)) + res = stats.anderson([1, 2, 3, 4, 5], 'norm', method=method) + assert res.pvalue == ref.pvalue # same random state -> same p-value + method = stats.MonteCarloMethod(rng=np.random.default_rng(23495984828)) + res = stats.anderson([1, 2, 3, 4, 5], 'norm', method=method) + assert res.pvalue != ref.pvalue # different random state -> different p-value + + @pytest.mark.parametrize('dist_name, seed', + [('norm', 4202165767275), + ('expon', 9094400417269), + pytest.param('logistic', 3776634590070, marks=pytest.mark.xslow), + pytest.param('gumbel_l', 7966588969335, marks=pytest.mark.xslow), + pytest.param('gumbel_r', 1886450383828, marks=pytest.mark.xslow)]) + def test_method_consistency(self, dist_name, seed): + dist = getattr(stats, dist_name) + rng = np.random.default_rng(seed) + x = dist.rvs(size=50, random_state=rng) + ref = stats.anderson(x, dist_name, method='interpolate') + res = stats.anderson(x, dist_name, method=stats.MonteCarloMethod(rng=rng)) + np.testing.assert_allclose(res.statistic, ref.statistic) + np.testing.assert_allclose(res.pvalue, ref.pvalue, atol=0.005) + + @pytest.mark.parametrize('dist_name', + ['norm', 'expon', 'logistic', 'gumbel_l', 'gumbel_r', 'weibull_min']) + def test_interpolate_saturation(self, dist_name): + dist = getattr(stats, dist_name) + rng = np.random.default_rng(4202165767276) + args = (3.5,) if dist_name == 'weibull_min' else tuple() + x = dist.rvs(*args, size=50, random_state=rng) + + with pytest.warns(FutureWarning): + res = stats.anderson(x, dist_name) + pvalues = (1 - np.asarray(res.significance_level) if dist_name == 'weibull_min' + else np.asarray(res.significance_level) / 100) + pvalue_min = np.min(pvalues) + pvalue_max = np.max(pvalues) + statistic_min = np.min(res.critical_values) + statistic_max = np.max(res.critical_values) + + # data drawn from distribution -> low statistic / high p-value + res = stats.anderson(x, dist_name, method='interpolate') + assert res.statistic < statistic_min + assert res.pvalue == pvalue_max + + # data not from distribution -> high statistic / low p-value + res = stats.anderson(rng.random(size=50), dist_name, method='interpolate') + assert res.statistic > statistic_max + assert res.pvalue == pvalue_min + + +@pytest.mark.filterwarnings("ignore:Parameter `variant`...:UserWarning") +class TestAndersonKSamp: + def test_example1a(self): + # Example data from Scholz & Stephens (1987), originally + # published in Lehmann (1995, Nonparametrics, Statistical + # Methods Based on Ranks, p. 309) + # Pass a mixture of lists and arrays + t1 = [38.7, 41.5, 43.8, 44.5, 45.5, 46.0, 47.7, 58.0] + t2 = np.array([39.2, 39.3, 39.7, 41.4, 41.8, 42.9, 43.3, 45.8]) + t3 = np.array([34.0, 35.0, 39.0, 40.0, 43.0, 43.0, 44.0, 45.0]) + t4 = np.array([34.0, 34.8, 34.8, 35.4, 37.2, 37.8, 41.2, 42.8]) + + Tk, tm, p = stats.anderson_ksamp((t1, t2, t3, t4), midrank=False) + + assert_almost_equal(Tk, 4.449, 3) + assert_array_almost_equal([0.4985, 1.3237, 1.9158, 2.4930, 3.2459], + tm[0:5], 4) + assert_allclose(p, 0.0021, atol=0.00025) + + def test_example1b(self): + # Example data from Scholz & Stephens (1987), originally + # published in Lehmann (1995, Nonparametrics, Statistical + # Methods Based on Ranks, p. 309) + # Pass arrays + t1 = np.array([38.7, 41.5, 43.8, 44.5, 45.5, 46.0, 47.7, 58.0]) + t2 = np.array([39.2, 39.3, 39.7, 41.4, 41.8, 42.9, 43.3, 45.8]) + t3 = np.array([34.0, 35.0, 39.0, 40.0, 43.0, 43.0, 44.0, 45.0]) + t4 = np.array([34.0, 34.8, 34.8, 35.4, 37.2, 37.8, 41.2, 42.8]) + Tk, tm, p = stats.anderson_ksamp((t1, t2, t3, t4), midrank=True) + + assert_almost_equal(Tk, 4.480, 3) + assert_array_almost_equal([0.4985, 1.3237, 1.9158, 2.4930, 3.2459], + tm[0:5], 4) + assert_allclose(p, 0.0020, atol=0.00025) + + @pytest.mark.xslow + def test_example2a(self): + # Example data taken from an earlier technical report of + # Scholz and Stephens + # Pass lists instead of arrays + t1 = [194, 15, 41, 29, 33, 181] + t2 = [413, 14, 58, 37, 100, 65, 9, 169, 447, 184, 36, 201, 118] + t3 = [34, 31, 18, 18, 67, 57, 62, 7, 22, 34] + t4 = [90, 10, 60, 186, 61, 49, 14, 24, 56, 20, 79, 84, 44, 59, 29, + 118, 25, 156, 310, 76, 26, 44, 23, 62] + t5 = [130, 208, 70, 101, 208] + t6 = [74, 57, 48, 29, 502, 12, 70, 21, 29, 386, 59, 27] + t7 = [55, 320, 56, 104, 220, 239, 47, 246, 176, 182, 33] + t8 = [23, 261, 87, 7, 120, 14, 62, 47, 225, 71, 246, 21, 42, 20, 5, + 12, 120, 11, 3, 14, 71, 11, 14, 11, 16, 90, 1, 16, 52, 95] + t9 = [97, 51, 11, 4, 141, 18, 142, 68, 77, 80, 1, 16, 106, 206, 82, + 54, 31, 216, 46, 111, 39, 63, 18, 191, 18, 163, 24] + t10 = [50, 44, 102, 72, 22, 39, 3, 15, 197, 188, 79, 88, 46, 5, 5, 36, + 22, 139, 210, 97, 30, 23, 13, 14] + t11 = [359, 9, 12, 270, 603, 3, 104, 2, 438] + t12 = [50, 254, 5, 283, 35, 12] + t13 = [487, 18, 100, 7, 98, 5, 85, 91, 43, 230, 3, 130] + t14 = [102, 209, 14, 57, 54, 32, 67, 59, 134, 152, 27, 14, 230, 66, + 61, 34] + + samples = (t1, t2, t3, t4, t5, t6, t7, t8, t9, t10, t11, t12, t13, t14) + Tk, tm, p = stats.anderson_ksamp(samples, midrank=False) + assert_almost_equal(Tk, 3.288, 3) + assert_array_almost_equal([0.5990, 1.3269, 1.8052, 2.2486, 2.8009], + tm[0:5], 4) + assert_allclose(p, 0.0041, atol=0.00025) + + rng = np.random.default_rng(6989860141921615054) + method = stats.PermutationMethod(n_resamples=9999, rng=rng) + res = stats.anderson_ksamp(samples, midrank=False, method=method) + assert_array_equal(res.statistic, Tk) + assert_array_equal(res.critical_values, tm) + assert_allclose(res.pvalue, p, atol=6e-4) + + def test_example2b(self): + # Example data taken from an earlier technical report of + # Scholz and Stephens + t1 = [194, 15, 41, 29, 33, 181] + t2 = [413, 14, 58, 37, 100, 65, 9, 169, 447, 184, 36, 201, 118] + t3 = [34, 31, 18, 18, 67, 57, 62, 7, 22, 34] + t4 = [90, 10, 60, 186, 61, 49, 14, 24, 56, 20, 79, 84, 44, 59, 29, + 118, 25, 156, 310, 76, 26, 44, 23, 62] + t5 = [130, 208, 70, 101, 208] + t6 = [74, 57, 48, 29, 502, 12, 70, 21, 29, 386, 59, 27] + t7 = [55, 320, 56, 104, 220, 239, 47, 246, 176, 182, 33] + t8 = [23, 261, 87, 7, 120, 14, 62, 47, 225, 71, 246, 21, 42, 20, 5, + 12, 120, 11, 3, 14, 71, 11, 14, 11, 16, 90, 1, 16, 52, 95] + t9 = [97, 51, 11, 4, 141, 18, 142, 68, 77, 80, 1, 16, 106, 206, 82, + 54, 31, 216, 46, 111, 39, 63, 18, 191, 18, 163, 24] + t10 = [50, 44, 102, 72, 22, 39, 3, 15, 197, 188, 79, 88, 46, 5, 5, 36, + 22, 139, 210, 97, 30, 23, 13, 14] + t11 = [359, 9, 12, 270, 603, 3, 104, 2, 438] + t12 = [50, 254, 5, 283, 35, 12] + t13 = [487, 18, 100, 7, 98, 5, 85, 91, 43, 230, 3, 130] + t14 = [102, 209, 14, 57, 54, 32, 67, 59, 134, 152, 27, 14, 230, 66, + 61, 34] + + Tk, tm, p = stats.anderson_ksamp((t1, t2, t3, t4, t5, t6, t7, t8, + t9, t10, t11, t12, t13, t14), + midrank=True) + + assert_almost_equal(Tk, 3.294, 3) + assert_array_almost_equal([0.5990, 1.3269, 1.8052, 2.2486, 2.8009], + tm[0:5], 4) + assert_allclose(p, 0.0041, atol=0.00025) + + def test_R_kSamples(self): + # test values generates with R package kSamples + # package version 1.2-6 (2017-06-14) + # r1 = 1:100 + # continuous case (no ties) --> version 1 + # res <- kSamples::ad.test(r1, r1 + 40.5) + # res$ad[1, "T.AD"] # 41.105 + # res$ad[1, " asympt. P-value"] # 5.8399e-18 + # + # discrete case (ties allowed) --> version 2 (here: midrank=True) + # res$ad[2, "T.AD"] # 41.235 + # + # res <- kSamples::ad.test(r1, r1 + .5) + # res$ad[1, "T.AD"] # -1.2824 + # res$ad[1, " asympt. P-value"] # 1 + # res$ad[2, "T.AD"] # -1.2944 + # + # res <- kSamples::ad.test(r1, r1 + 7.5) + # res$ad[1, "T.AD"] # 1.4923 + # res$ad[1, " asympt. P-value"] # 0.077501 + # + # res <- kSamples::ad.test(r1, r1 + 6) + # res$ad[2, "T.AD"] # 0.63892 + # res$ad[2, " asympt. P-value"] # 0.17981 + # + # res <- kSamples::ad.test(r1, r1 + 11.5) + # res$ad[1, "T.AD"] # 4.5042 + # res$ad[1, " asympt. P-value"] # 0.00545 + # + # res <- kSamples::ad.test(r1, r1 + 13.5) + # res$ad[1, "T.AD"] # 6.2982 + # res$ad[1, " asympt. P-value"] # 0.00118 + + x1 = np.linspace(1, 100, 100) + # test case: different distributions;p-value floored at 0.001 + # test case for issue #5493 / #8536 + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", 'p-value floored', UserWarning) + s, _, p = stats.anderson_ksamp([x1, x1 + 40.5], midrank=False) + assert_almost_equal(s, 41.105, 3) + assert_equal(p, 0.001) + + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", 'p-value floored', UserWarning) + s, _, p = stats.anderson_ksamp([x1, x1 + 40.5]) + assert_almost_equal(s, 41.235, 3) + assert_equal(p, 0.001) + + # test case: similar distributions --> p-value capped at 0.25 + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", 'p-value capped', UserWarning) + s, _, p = stats.anderson_ksamp([x1, x1 + .5], midrank=False) + assert_almost_equal(s, -1.2824, 4) + assert_equal(p, 0.25) + + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", 'p-value capped', UserWarning) + s, _, p = stats.anderson_ksamp([x1, x1 + .5]) + assert_almost_equal(s, -1.2944, 4) + assert_equal(p, 0.25) + + # test case: check interpolated p-value in [0.01, 0.25] (no ties) + s, _, p = stats.anderson_ksamp([x1, x1 + 7.5], midrank=False) + assert_almost_equal(s, 1.4923, 4) + assert_allclose(p, 0.0775, atol=0.005, rtol=0) + + # test case: check interpolated p-value in [0.01, 0.25] (w/ ties) + s, _, p = stats.anderson_ksamp([x1, x1 + 6]) + assert_almost_equal(s, 0.6389, 4) + assert_allclose(p, 0.1798, atol=0.005, rtol=0) + + # test extended critical values for p=0.001 and p=0.005 + s, _, p = stats.anderson_ksamp([x1, x1 + 11.5], midrank=False) + assert_almost_equal(s, 4.5042, 4) + assert_allclose(p, 0.00545, atol=0.0005, rtol=0) + + s, _, p = stats.anderson_ksamp([x1, x1 + 13.5], midrank=False) + assert_almost_equal(s, 6.2982, 4) + assert_allclose(p, 0.00118, atol=0.0001, rtol=0) + + def test_not_enough_samples(self): + assert_raises(ValueError, stats.anderson_ksamp, np.ones(5)) + + def test_no_distinct_observations(self): + assert_raises(ValueError, stats.anderson_ksamp, + (np.ones(5), np.ones(5))) + + def test_empty_sample(self): + assert_raises(ValueError, stats.anderson_ksamp, (np.ones(5), [])) + + def test_result_attributes(self): + # Pass a mixture of lists and arrays + t1 = [38.7, 41.5, 43.8, 44.5, 45.5, 46.0, 47.7, 58.0] + t2 = np.array([39.2, 39.3, 39.7, 41.4, 41.8, 42.9, 43.3, 45.8]) + res = stats.anderson_ksamp((t1, t2), midrank=False) + + attributes = ('statistic', 'critical_values', 'significance_level') + check_named_results(res, attributes) + + assert_equal(res.significance_level, res.pvalue) + + +class TestAndersonKSampVariant: + def test_variant_values(self): + x = [1, 2, 2, 3, 4, 5] + y = [1, 2, 3, 4, 4, 5, 6, 6, 6, 7] + message = "Parameter `variant` has been introduced..." + with pytest.warns(UserWarning, match=message): + ref = stats.anderson_ksamp((x, y)) + assert len(ref) == 3 and hasattr(ref, 'critical_values') + + with pytest.warns(UserWarning, match=message): + res = stats.anderson_ksamp((x, y), midrank=True) + assert_equal(res.statistic, ref.statistic) + assert_equal(res.pvalue, ref.pvalue) + assert len(res) == 3 and hasattr(res, 'critical_values') + + with pytest.warns(UserWarning, match=message): + res = stats.anderson_ksamp((x, y), midrank=False, variant='midrank') + assert_equal(res.statistic, ref.statistic) + assert_equal(res.pvalue, ref.pvalue) + assert not hasattr(res, 'critical_values') + + res = stats.anderson_ksamp((x, y), variant='midrank') + assert_equal(res.statistic, ref.statistic) + assert_equal(res.pvalue, ref.pvalue) + assert not hasattr(res, 'critical_values') + + with pytest.warns(UserWarning, match=message): + ref = stats.anderson_ksamp((x, y), midrank=False) + assert len(ref) == 3 and hasattr(ref, 'critical_values') + + with pytest.warns(UserWarning, match=message): + res = stats.anderson_ksamp((x, y), midrank=True, variant='right') + assert_equal(res.statistic, ref.statistic) + assert_equal(res.pvalue, ref.pvalue) + assert not hasattr(res, 'critical_values') + + res = stats.anderson_ksamp((x, y), variant='right') + assert_equal(res.statistic, ref.statistic) + assert_equal(res.pvalue, ref.pvalue) + assert not hasattr(res, 'critical_values') + + def test_variant_input_validation(self): + x = np.arange(10) + message = "`variant` must be one of 'midrank', 'right', or 'continuous'." + with pytest.raises(ValueError, match=message): + stats.anderson_ksamp((x, x), variant='Camelot') + + @pytest.mark.parametrize('n_samples', [2, 3]) + def test_variant_continuous(self, n_samples): + rng = np.random.default_rng(20182053007) + samples = rng.random((n_samples, 15)) + 0.1*np.arange(n_samples)[:, np.newaxis] + ref = stats.anderson_ksamp(samples, variant='right') + res = stats.anderson_ksamp(samples, variant='continuous') + assert_allclose(res.statistic, ref.statistic) + assert_allclose(res.pvalue, ref.pvalue) + + +@make_xp_test_case(stats.ansari) +class TestAnsari: + + def test_small(self, xp): + x = xp.asarray([1, 2, 3, 3, 4]) + y = xp.asarray([3, 2, 6, 1, 6, 1, 4, 1]) + W, pval = stats.ansari(x, y, method='asymptotic') + xp_assert_close(W, xp.asarray(23.5)) + xp_assert_close(pval, xp.asarray(0.13499256881897437)) + + def test_approx(self, xp): + ramsay = xp.asarray([111, 107, 100, 99, 102, 106, 109, 108, 104, 99, + 101, 96, 97, 102, 107, 113, 116, 113, 110, 98]) + parekh = xp.asarray([107, 108, 106, 98, 105, 103, 110, 105, 104, + 100, 96, 108, 103, 104, 114, 114, 113, 108, + 106, 99]) + + W, pval = stats.ansari(ramsay, parekh, method='asymptotic') + xp_assert_close(W, xp.asarray(185.5)) + xp_assert_close(pval, xp.asarray(0.18145819972867083)) + + def test_exact(self, xp): + x, y = xp.asarray([1, 2, 3, 4]), xp.asarray([15, 5, 20, 8, 10, 12]) + W, pval = stats.ansari(x, y, method='exact') + xp_assert_close(W, xp.asarray(10.0)) + xp_assert_close(pval, xp.asarray(0.533333333333333333)) + + @pytest.mark.parametrize('args', [([], [1.]), ([1.], [])]) + def test_bad_arg(self, args, xp): + args = [xp.asarray(arg) for arg in args] + with pytest.warns(SmallSampleWarning, match=too_small_1d_not_omit): + res = stats.ansari(*args) + xp_assert_equal(res.statistic, xp.asarray(xp.nan)) + xp_assert_equal(res.pvalue, xp.asarray(xp.nan)) + + def test_bad_alternative(self, xp): + # invalid value for alternative must raise a ValueError + x1 = xp.asarray([1, 2, 3, 4]) + x2 = xp.asarray([5, 6, 7, 8]) + match = "'alternative' must be 'two-sided'" + with assert_raises(ValueError, match=match): + stats.ansari(x1, x2, alternative='foo') + + def test_alternative_exact(self, xp): + x1 = xp.asarray([-5, 1, 5, 10, 15, 20, 25.]) # high scale, loc=10 + x2 = xp.asarray([7.5, 8.5, 9.5, 10.5, 11.5, 12.5]) # low scale, loc=10 + # ratio of scales is greater than 1. So, the + # p-value must be high when `alternative='less'` + # and low when `alternative='greater'`. + statistic, pval = stats.ansari(x1, x2, method='exact') + pval_l = stats.ansari(x1, x2, alternative='less', method='exact').pvalue + pval_g = stats.ansari(x1, x2, alternative='greater', method='exact').pvalue + assert pval_l > 0.95 + assert pval_g < 0.05 # level of significance. + # also check if the p-values sum up to 1 plus the probability + # mass under the calculated statistic. + prob = _abw_state.a.pmf(float(statistic), x1.shape[0], x2.shape[0]) + prob = xp.asarray(float(prob)) + xp_assert_close(pval_g + pval_l, 1 + prob, atol=1e-12) + # also check if one of the one-sided p-value equals half the + # two-sided p-value and the other one-sided p-value is its + # compliment. + xp_assert_close(pval_g, pval/2, atol=1e-12) + xp_assert_close(pval_l, 1+prob-pval/2, atol=1e-12) + # sanity check. The result should flip if + # we exchange x and y. + pval_l_reverse = stats.ansari(x2, x1, alternative='less', + method='exact').pvalue + pval_g_reverse = stats.ansari(x2, x1, alternative='greater', + method='exact').pvalue + assert pval_l_reverse < 0.05 + assert pval_g_reverse > 0.95 + + @pytest.mark.parametrize( + 'x, y, alternative, expected', + # the tests are designed in such a way that the + # if else statement in ansari test for exact + # mode is covered. + [([1, 2, 3, 4], [5, 6, 7, 8], 'less', 0.6285714285714), + ([1, 2, 3, 4], [5, 6, 7, 8], 'greater', 0.6285714285714), + ([1, 2, 3], [4, 5, 6, 7, 8], 'less', 0.8928571428571), + ([1, 2, 3], [4, 5, 6, 7, 8], 'greater', 0.2857142857143), + ([1, 2, 3, 4, 5], [6, 7, 8], 'less', 0.2857142857143), + ([1, 2, 3, 4, 5], [6, 7, 8], 'greater', 0.8928571428571)] + ) + def test_alternative_exact_with_R(self, x, y, alternative, expected, xp): + # testing with R on arbitrary data + # Sample R code used for the third test case above: + # ```R + # > options(digits=16) + # > x <- c(1,2,3) + # > y <- c(4,5,6,7,8) + # > ansari.test(x, y, alternative='less', exact=TRUE) + # + # Ansari-Bradley test + # + # data: x and y + # AB = 6, p-value = 0.8928571428571 + # alternative hypothesis: true ratio of scales is less than 1 + # + # ``` + x, y = xp.asarray(x), xp.asarray(y) + pval = stats.ansari(x, y, alternative=alternative, method='exact').pvalue + xp_assert_close(pval, xp.asarray(expected), atol=1e-12) + + def test_alternative_approx(self, xp): + # intuitive tests for approximation + x1 = xp.asarray(stats.norm.rvs(0, 5, size=100, random_state=123)) + x2 = xp.asarray(stats.norm.rvs(0, 2, size=100, random_state=123)) + # for m > 55 or n > 55, the test should automatically + # switch to approximation. + pval_l = stats.ansari(x1, x2, alternative='less', method='asymptotic').pvalue + pval_g = stats.ansari(x1, x2, alternative='greater', method='asymptotic').pvalue + xp_assert_close(pval_l, xp.asarray(1.0, dtype=xp.float64), atol=1e-12) + xp_assert_close(pval_g, xp.asarray(0.0, dtype=xp.float64), atol=1e-12) + # also check if one of the one-sided p-value equals half the + # two-sided p-value and the other one-sided p-value is its + # compliment. + x1 = xp.asarray(stats.norm.rvs(0, 2, size=60, random_state=123)) + x2 = xp.asarray(stats.norm.rvs(0, 1.5, size=60, random_state=123)) + pval = stats.ansari(x1, x2, method='asymptotic').pvalue + pval_l = stats.ansari(x1, x2, alternative='less', method='asymptotic').pvalue + pval_g = stats.ansari(x1, x2, alternative='greater', method='asymptotic').pvalue + xp_assert_close(pval_g, pval/2, atol=1e-12) + xp_assert_close(pval_l, 1-pval/2, atol=1e-12) + + @pytest.mark.parametrize('dtype', [None, 'float32', 'float64']) + @pytest.mark.parametrize('n', [10, 100]) # affects code path + @pytest.mark.parametrize('ties', [False, True]) # affects code path + def test_dtypes(self, dtype, n, ties, xp): + if is_jax(xp) and dtype != 'float64': + pytest.xfail("p-value calculation works natively with 'float64 only") + dtype = xp_default_dtype(xp) if dtype is None else getattr(xp, dtype) + rng = np.random.default_rng(78587342806484) + x, y = rng.integers(6, size=(2, n)) if ties else rng.random(size=(2, n)) + method = {'method': 'exact' if ((n <= 55) and not ties) else 'asymptotic'} + ref = stats.ansari(x, y, **method) + res = stats.ansari(xp.asarray(x, dtype=dtype), xp.asarray(y, dtype=dtype), + **method) + xp_assert_close(res.statistic, xp.asarray(ref.statistic, dtype=dtype)) + xp_assert_close(res.pvalue, xp.asarray(ref.pvalue, dtype=dtype)) + + @pytest.mark.parametrize('alternative', ['less', 'greater', 'two-sided']) + def test_permutation_method(self, alternative, xp): + # test that permutation test can reproduce results of method='exact', + # but can also be configured to perform randomized test + rng = np.random.default_rng(6123643029) + x = xp.asarray(rng.random(7).tolist()) + y = xp.asarray(rng.random(6).tolist()) + res = stats.ansari(x, y, alternative=alternative, + method=stats.PermutationMethod(rng=rng)) + ref = stats.ansari(x, y, alternative=alternative, method='exact') + xp_assert_close(res.statistic, xp.asarray(ref.statistic)) + xp_assert_close(res.pvalue, xp.asarray(ref.pvalue)) + + if not is_lazy_array(x): # randomized permutation test has side-effects + res = stats.ansari(x, y, alternative=alternative, + method=stats.PermutationMethod(rng=rng, n_resamples=99)) + xp_assert_close(res.statistic, xp.asarray(ref.statistic)) + assert res.pvalue != ref.pvalue + assert res.pvalue == xp.round(res.pvalue*100)/100 # 99+1 in denominator + + def test_method(self, xp): + # test that `method` selection works and 'auto' behaves as expected + # expected p-values computed with `ansari.test` (R) + # e.g. `ansari.test(x, y, exact=False)` + rng = np.random.default_rng(6123643029) + + if is_jax(xp): + message = "`method` must be 'exact', 'asymptotic', or an instance..." + with pytest.raises(ValueError, match=message): + stats.ansari(xp.zeros(4), xp.ones(5)) + return + + # default is exact, but asymptotic can be selected + x = xp.asarray(rng.random(54).tolist()) + y = xp.asarray(rng.random(54).tolist()) + res1 = stats.ansari(x, y) + res2 = stats.ansari(x, y, method='exact') + res3 = stats.ansari(x, y, method='asymptotic') + xp_assert_close(res1.pvalue, res2.pvalue) + xp_assert_close(res3.pvalue, xp.asarray(0.01259674075933)) + + # default is asymptotic, but exact can be selected + x = xp.asarray(rng.random(55).tolist()) + y = xp.asarray(rng.random(55).tolist()) + res1 = stats.ansari(x, y) + res2 = stats.ansari(x, y, method='asymptotic') + res3 = stats.ansari(x, y, method='exact') + xp_assert_close(res1.pvalue, res2.pvalue) + xp_assert_close(res3.pvalue, xp.asarray(0.7613489076855)) + + # default is asymptotic because of ties + x = xp.asarray(rng.integers(10, size=20).tolist()) + y = xp.asarray(rng.integers(10, size=20).tolist()) + res1 = stats.ansari(x, y) + res2 = stats.ansari(x, y, method='asymptotic') + xp_assert_close(res1.pvalue, res2.pvalue) + + +@make_xp_test_case(stats.bartlett) +class TestBartlett: + def test_data(self, xp): + # https://www.itl.nist.gov/div898/handbook/eda/section3/eda357.htm + args = [g1, g2, g3, g4, g5, g6, g7, g8, g9, g10] + args = [xp.asarray(arg) for arg in args] + T, pval = stats.bartlett(*args) + xp_assert_close(T, xp.asarray(20.78587342806484)) + xp_assert_close(pval, xp.asarray(0.0136358632781)) + + def test_too_few_args(self, xp): + message = "Must enter at least two input sample vectors." + with pytest.raises(ValueError, match=message): + stats.bartlett(xp.asarray([1.])) + + def test_result_attributes(self, xp): + args = [g1, g2, g3, g4, g5, g6, g7, g8, g9, g10] + args = [xp.asarray(arg) for arg in args] + res = stats.bartlett(*args) + attributes = ('statistic', 'pvalue') + check_named_results(res, attributes, xp=xp) + + @pytest.mark.filterwarnings("ignore:invalid value encountered") # Dask + def test_empty_arg(self, xp): + args = (g1, g2, g3, g4, g5, g6, g7, g8, g9, g10, []) + args = [xp.asarray(arg) for arg in args] + + with eager_warns(SmallSampleWarning, match=too_small_1d_not_omit, xp=xp): + res = stats.bartlett(*args) + + NaN = xp.asarray(xp.nan) + xp_assert_equal(res.statistic, NaN) + xp_assert_equal(res.pvalue, NaN) + + def test_negative_pvalue_gh21152(self, xp): + a = xp.asarray([10.1, 10.2, 10.3, 10.4], dtype=xp.float32) + b = xp.asarray([10.15, 10.25, 10.35, 10.45], dtype=xp.float32) + c = xp.asarray([10.05, 10.15, 10.25, 10.35], dtype=xp.float32) + res = stats.bartlett(a, b, c) + assert xp.all(res.statistic >= 0) + + +@make_xp_test_case(stats.levene) +class TestLevene: + + def test_data(self, xp): + # https://www.itl.nist.gov/div898/handbook/eda/section3/eda35a.htm + args = [g1, g2, g3, g4, g5, g6, g7, g8, g9, g10] + args = [xp.asarray(arg) for arg in args] + W, pval = stats.levene(*args) + xp_assert_close(W, xp.asarray(1.7059176930008939)) + xp_assert_close(pval, xp.asarray(0.0990829755522)) + + def test_mean(self, xp): + # numbers from R: leveneTest in package car + args = [g1, g2, g3, g4, g5, g6, g7, g8, g9, g10] + args = [xp.asarray(arg) for arg in args] + W, pval = stats.levene(*args, center="mean") + xp_assert_close(W, xp.asarray(2.15945985647285)) + xp_assert_close(pval, xp.asarray(0.032236826559783)) + + def test_trimmed1(self, xp): + # Test that center='trimmed' gives the same result as center='mean' + # when proportiontocut=0. + args = (xp.asarray(g1), xp.asarray(g2), xp.asarray(g3)) + W1, pval1 = stats.levene(*args, center='mean') + W2, pval2 = stats.levene(*args, center='trimmed', proportiontocut=0.0) + xp_assert_close(W1, W2) + xp_assert_close(pval1, pval2) + + def test_trimmed2(self, xp): + # numbers from R: leveneTest in package car + args = [g1, g2, g3, g4, g5, g6, g7, g8, g9, g10] + args = [xp.asarray(arg) for arg in args] + W, pval = stats.levene(*args, center="trimmed", proportiontocut=0.25) + xp_assert_close(W, xp.asarray(2.07712845686874)) + xp_assert_close(pval, xp.asarray(0.0397269688035377)) + + def test_equal_mean_median(self, xp): + x = np.linspace(-1, 1, 21) + rng = np.random.default_rng(4058827756) + x2 = rng.permutation(x) + y = x**3 + x, x2, y = xp.asarray(x), xp.asarray(x2), xp.asarray(y) + W1, pval1 = stats.levene(x, y, center='mean') + W2, pval2 = stats.levene(x2, y, center='median') + xp_assert_close(W1, W2) + xp_assert_close(pval1, pval2) + + def test_bad_center_value(self, xp): + x = xp.linspace(-1, 1, 21) + message = "center must be 'mean', 'median' or 'trimmed'." + with pytest.raises(ValueError, match=message): + stats.levene(x, x, center='trim') + + def test_too_few_args(self, xp): + message = "Must provide at least two samples." + with pytest.raises(ValueError, match=message): + stats.levene(xp.asarray([1])) + + def test_result_attributes(self, xp): + args = [g1, g2, g3, g4, g5, g6, g7, g8, g9, g10] + args = [xp.asarray(arg) for arg in args] + res = stats.levene(*args) + attributes = ('statistic', 'pvalue') + check_named_results(res, attributes, xp=xp) + + +@make_xp_test_case(stats.binomtest) +class TestBinomTest: + """Tests for stats.binomtest.""" + + # Expected results here are from R binom.test, e.g. + # options(digits=16) + # binom.test(484, 967, p=0.48) + @skip_xp_backends("jax.numpy", reason="'two-sided' is incompatible with JAX") + @pytest.mark.parametrize("k, n, p, ref, rtol", + # aarch64 observed rtol: 1.5e-11 + [(10079999, 21000000, 0.48, 1.0, 1e-10), + (10079990, 21000000, 0.48, 0.9966892187965, 1e-10), + (10080009, 21000000, 0.48, 0.9970377203856, 1e-10), + (10080017, 21000000, 0.48, 0.9940754817328, 1e-9)]) + def test_two_sided_pvalues1(self, k, n, p, ref, rtol, xp): + # `tol` could be stricter on most architectures, but the value + # here is limited by accuracy of `binom.cdf` for large inputs on + # Linux_Python_37_32bit_full and aarch64 + dtype = xp.float64 # because this tests precision + res = stats.binomtest(k, n, xp.asarray(p, dtype=dtype)) + xp_assert_close(res.pvalue, xp.asarray(ref, dtype=dtype), rtol=rtol) + + @skip_xp_backends("jax.numpy", reason="'two-sided' is incompatible with JAX") + @pytest.mark.parametrize("k, n, p, ref, rtol", + # no aarch64 failure with 1e-15, preemptive bump + [(9, 21, 0.48, 0.6689672431939, 1e-10), + (4, 21, 0.48, 0.0081395634521, 1e-10), + (11, 21, 0.48, 0.8278629664608, 1e-10), + (7, 21, 0.48, 0.1966772901718, 1e-10), + (3, 10, 0.5, 0.34375, 1e-10), + (2, 2, 0.4, 0.16, 1e-10), + (2, 4, 0.3, 0.5884, 1e-10)]) + def test_two_sided_pvalues2(self, k, n, p, ref, rtol, xp): + dtype = xp.float64 # because this tests precision + res = stats.binomtest(k, n, xp.asarray(p, dtype=dtype)) + xp_assert_close(res.pvalue, xp.asarray(ref, dtype=dtype), rtol=rtol) + + @skip_xp_backends("jax.numpy", reason="'two-sided' is incompatible with JAX") + @pytest.mark.parametrize("k, n, p, ref, rtol", + # no aarch64 failure with 1e-15, preemptive bump + [(484, 967, 0.5, 1.0, 1e-10), + (3, 47, 3/47, 1.0, 1e-10), + (13, 46, 13/46, 1.0, 1e-10), + (15, 44, 15/44, 1.0, 1e-10), + (7, 13, 0.5, 1.0, 1e-10), + (6, 11, 0.5, 1.0, 1e-10)]) + def test_edge_cases(self, k, n, p, ref, rtol, xp): + dtype = xp.float64 + res = stats.binomtest(k, n, xp.asarray(p, dtype=dtype)) + xp_assert_close(res.pvalue, xp.asarray(ref, dtype=dtype), rtol=rtol) + + # tests private function, so leave NumPy-only + def test_binary_srch_for_binom_tst(self): + # Test that old behavior of binomtest is maintained + # by the new binary search method in cases where d + # exactly equals the input on one side. + n = 10 + p = 0.5 + k = 3 + # First test for the case where k > mode of PMF + i = np.arange(np.ceil(p * n), n+1) + d = stats.binom.pmf(k, n, p) + # Old way of calculating y, probably consistent with R. + y1 = np.sum(stats.binom.pmf(i, n, p) <= d, axis=0) + # New way with binary search. + ix = _binary_search_for_binom_tst(lambda x1: + -stats.binom.pmf(x1, n, p), + -d, np.ceil(p * n), n, xp=np) + y2 = n - ix + int(d == stats.binom.pmf(ix, n, p)) + assert_allclose(y1, y2, rtol=1e-9) + # Now test for the other side. + k = 7 + i = np.arange(np.floor(p * n) + 1) + d = stats.binom.pmf(k, n, p) + # Old way of calculating y. + y1 = np.sum(stats.binom.pmf(i, n, p) <= d, axis=0) + # New way with binary search. + ix = _binary_search_for_binom_tst(lambda x1: + stats.binom.pmf(x1, n, p), + d, 0, np.floor(p * n), xp=np) + y2 = ix + 1 + assert_allclose(y1, y2, rtol=1e-9) + + # Expected results here are from R 3.6.2 binom.test + @skip_xp_backends("array_api_strict", reason="needs find_root") + @skip_xp_backends("jax.numpy", reason="needs find_root") + @pytest.mark.parametrize('alternative, pval, ci_low, ci_high', + [('less', 0.148831050443, + 0.0, 0.2772002496709138), + ('greater', 0.9004695898947, + 0.1366613252458672, 1.0), + ('two-sided', 0.2983720970096, + 0.1266555521019559, 0.2918426890886281)]) + def test_confidence_intervals1(self, alternative, pval, ci_low, ci_high, xp): + dtype = xp.float64 + res = stats.binomtest(20, n=100, p=xp.asarray(0.25, dtype=dtype), + alternative=alternative) + xp_assert_close(res.pvalue, xp.asarray(pval, dtype=dtype), rtol=1e-12) + xp_assert_equal(res.statistic, xp.asarray(0.2, dtype=dtype)) + ci = res.proportion_ci(confidence_level=xp.asarray(0.95, dtype=dtype)) + xp_assert_close(ci.low, xp.asarray(ci_low, dtype=dtype), rtol=1e-12) + xp_assert_close(ci.high, xp.asarray(ci_high, dtype=dtype), rtol=1e-12) + + # Expected results here are from R 3.6.2 binom.test. + @skip_xp_backends("array_api_strict", reason="needs find_root") + @skip_xp_backends("jax.numpy", reason="needs find_root") + @pytest.mark.parametrize('alternative, pval, ci_low, ci_high', + [('less', + 0.005656361, 0.0, 0.1872093), + ('greater', + 0.9987146, 0.008860761, 1.0), + ('two-sided', + 0.01191714, 0.006872485, 0.202706269)]) + def test_confidence_intervals2(self, alternative, pval, ci_low, ci_high, xp): + dtype = xp.float64 + res = stats.binomtest(3, n=50, p=xp.asarray(0.2, dtype=dtype), + alternative=alternative) + xp_assert_close(res.pvalue, xp.asarray(pval, dtype=dtype), rtol=1e-6) + xp_assert_equal(res.statistic, xp.asarray(0.06, dtype=dtype)) + ci = res.proportion_ci(confidence_level=xp.asarray(0.99, dtype=dtype)) + xp_assert_close(ci.low, xp.asarray(ci_low, dtype=dtype), rtol=1e-6) + xp_assert_close(ci.high, xp.asarray(ci_high, dtype=dtype), rtol=1e-6) + + # Expected results here are from R 3.6.2 binom.test. + @skip_xp_backends("array_api_strict", reason="needs find_root") + @skip_xp_backends("jax.numpy", reason="needs find_root") + @pytest.mark.parametrize('alternative, pval, ci_high', + [('less', 0.05631351, 0.2588656), + ('greater', 1.0, 1.0), + ('two-sided', 0.07604122, 0.3084971)]) + def test_confidence_interval_exact_k0(self, alternative, pval, ci_high, xp): + # Test with k=0, n = 10. + dtype = xp.float64 + res = stats.binomtest(0, n=10, p=xp.asarray(0.25, dtype=dtype), + alternative=alternative) + xp_assert_close(res.pvalue, xp.asarray(pval, dtype=dtype), rtol=1e-6) + ci = res.proportion_ci(confidence_level=xp.asarray(0.95, dtype=dtype)) + xp_assert_equal(ci.low, xp.asarray(0.0, dtype=dtype)) + xp_assert_close(ci.high, xp.asarray(ci_high, dtype=dtype), rtol=1e-6) + + # Expected results here are from R 3.6.2 binom.test. + @skip_xp_backends("array_api_strict", reason="needs find_root") + @skip_xp_backends("jax.numpy", reason="needs find_root") + @pytest.mark.parametrize('alternative, pval, ci_low', + [('less', 1.0, 0.0), + ('greater', 9.536743e-07, 0.7411344), + ('two-sided', 9.536743e-07, 0.6915029)]) + def test_confidence_interval_exact_k_is_n(self, alternative, pval, ci_low, xp): + # Test with k = n = 10. + dtype = xp.float64 + res = stats.binomtest(10, n=10, p=xp.asarray(0.25, dtype=dtype), + alternative=alternative) + xp_assert_close(res.pvalue, xp.asarray(pval, dtype=dtype), rtol=1e-6) + ci = res.proportion_ci(confidence_level=xp.asarray(0.95, dtype=dtype)) + xp_assert_equal(ci.high, xp.asarray(1.0, dtype=dtype)) + xp_assert_close(ci.low, xp.asarray(ci_low, dtype=dtype), rtol=1e-6) + + # Expected results are from the prop.test function in R 3.6.2. + @skip_xp_backends("jax.numpy", reason="'two-sided' is incompatible with JAX") + @pytest.mark.parametrize( + 'k, alternative, corr, conf, ci_low, ci_high', + [[3, 'two-sided', True, 0.95, 0.08094782, 0.64632928], + [3, 'two-sided', True, 0.99, 0.0586329, 0.7169416], + [3, 'two-sided', False, 0.95, 0.1077913, 0.6032219], + [3, 'two-sided', False, 0.99, 0.07956632, 0.6799753], + [3, 'less', True, 0.95, 0.0, 0.6043476], + [3, 'less', True, 0.99, 0.0, 0.6901811], + [3, 'less', False, 0.95, 0.0, 0.5583002], + [3, 'less', False, 0.99, 0.0, 0.6507187], + [3, 'greater', True, 0.95, 0.09644904, 1.0], + [3, 'greater', True, 0.99, 0.06659141, 1.0], + [3, 'greater', False, 0.95, 0.1268766, 1.0], + [3, 'greater', False, 0.99, 0.08974147, 1.0], + + [0, 'two-sided', True, 0.95, 0.0, 0.3445372], + [0, 'two-sided', False, 0.95, 0.0, 0.2775328], + [0, 'less', True, 0.95, 0.0, 0.2847374], + [0, 'less', False, 0.95, 0.0, 0.212942], + [0, 'greater', True, 0.95, 0.0, 1.0], + [0, 'greater', False, 0.95, 0.0, 1.0], + + [10, 'two-sided', True, 0.95, 0.6554628, 1.0], + [10, 'two-sided', False, 0.95, 0.7224672, 1.0], + [10, 'less', True, 0.95, 0.0, 1.0], + [10, 'less', False, 0.95, 0.0, 1.0], + [10, 'greater', True, 0.95, 0.7152626, 1.0], + [10, 'greater', False, 0.95, 0.787058, 1.0]] + ) + def test_ci_wilson_method(self, k, alternative, corr, conf, + ci_low, ci_high, xp): + dtype = xp.float64 + res = stats.binomtest(k, n=10, p=xp.asarray(0.1, dtype=dtype), + alternative=alternative) + if corr: + method = 'wilsoncc' + else: + method = 'wilson' + conf = xp.asarray(conf, dtype=dtype) + ci = res.proportion_ci(confidence_level=conf, method=method) + xp_assert_close(ci.low, xp.asarray(ci_low, dtype=dtype), rtol=1e-6) + xp_assert_close(ci.high, xp.asarray(ci_high, dtype=dtype), rtol=1e-6) + + @skip_xp_backends("jax.numpy", reason="'two-sided' is incompatible with JAX") + def test_estimate_equals_hypothesized_prop(self, xp): + # Test the special case where the estimated proportion equals + # the hypothesized proportion. When alternative is 'two-sided', + # the p-value is 1. + res = stats.binomtest(4, 16, xp.asarray(0.25)) + xp_assert_equal(res.statistic, xp.asarray(0.25)) + xp_assert_equal(res.pvalue, xp.asarray(1.0)) + + def test_invalid_confidence_level(self, xp): + res = stats.binomtest(3, n=10, p=xp.asarray(0.1), alternative='less') + message = r"confidence_level \(-1\) must be in the interval" + with pytest.raises(ValueError, match=message): + res.proportion_ci(confidence_level=-1) + + def test_invalid_ci_method(self, xp): + res = stats.binomtest(3, n=10, p=xp.asarray(0.1), alternative='greater') + with pytest.raises(ValueError, match=r"method \('plate of shrimp'\) must be"): + res.proportion_ci(method="plate of shrimp") + + def test_invalid_alternative(self, xp): + with pytest.raises(ValueError, match=r"alternative \('ekki'\) not..."): + stats.binomtest(3, n=10, p=xp.asarray(0.1), alternative='ekki') + + def test_alias(self, xp): + res = stats.binomtest(3, n=10, p=xp.asarray(0.1), alternative='greater') + xp_assert_equal(res.proportion_estimate, res.statistic) + + @pytest.mark.skipif(sys.maxsize <= 2**32, reason="32-bit does not overflow") + def test_boost_overflow_raises(self): # NumPy-specific + # Boost.Math error policy should raise exceptions in Python + with pytest.raises(OverflowError, match='Error in function...'): + stats.binomtest(5, 6, p=sys.float_info.min) + + @pytest.mark.parametrize("k, n, p", + [(-1, 10, 0.5), (11, 10, 0.5), (5.5, 10, 0.5), (np.nan, 10, 0.5), + (0, 0, 0.5), (5, 10.5, 0.5), (5, np.nan, 0.5), + (5, 10, -0.1), (5, 10, 1.1), (5, 10, np.nan)]) + def test_invalid(self, k, n, p, xp): + res = stats.binomtest(k, n, xp.asarray(p), alternative='less') + xp_assert_equal(res.statistic, xp.asarray(np.nan)) + xp_assert_equal(res.pvalue, xp.asarray(np.nan)) + + @pytest.mark.parametrize("alternative", ['less', 'greater', 'two-sided']) + @pytest.mark.parametrize("method", ['exact', 'wilson', 'wilsoncc']) + def test_scalar_0d_in_scalar_out(self, alternative, method): + # Tests that scalar / 0-D NumPy array input are equivalent and that both + # result in NumPy scalars (as confirmed by xp_assert_equal). + res = stats.binomtest(3, 11, 0.4, alternative=alternative) + ref = stats.binomtest(np.asarray(3), np.asarray(11), np.asarray(0.4), + alternative=alternative) + + xp_assert_equal(res.statistic, ref.statistic) + xp_assert_equal(res.pvalue, ref.pvalue) + xp_assert_equal(res.n, ref.n) + xp_assert_equal(res.k, ref.k) + + res = res.proportion_ci(method=method) + ref = ref.proportion_ci(method=method) + + xp_assert_equal(res.low, ref.low) + xp_assert_equal(res.high, ref.high) + + @pytest.mark.parametrize("dtype", [None, 'float32', 'float64']) + @pytest.mark.parametrize("alternative", ['less', 'greater', 'two-sided']) + @pytest.mark.parametrize("method", ['exact', 'wilson', 'wilsoncc']) + def test_dtype(self, dtype, alternative, method, xp): + # Tests that output dtype is as expected + dtype = dtype if dtype is None else getattr(xp, dtype) + k = xp.asarray(3, dtype=dtype) + n = xp.asarray(11, dtype=dtype) + p = xp.asarray(0.4, dtype=dtype) + if is_jax(xp) and alternative=='two-sided': + message = "`alternative='two-sided'` is incompatible with JAX arrays." + with pytest.raises(ValueError, match=message): + stats.binomtest(k, n, p) + return + + res = stats.binomtest(k, n, p, alternative=alternative) + ref = stats.binomtest(3, 11, 0.4, alternative=alternative) + + xp_assert_close(res.statistic, xp.asarray(float(ref.statistic), dtype=dtype)) + xp_assert_close(res.pvalue, xp.asarray(float(ref.pvalue), dtype=dtype)) + + if res.statistic.ndim == 0: + if not is_lazy_array(res.n): + assert isinstance(res.n, int) + assert isinstance(res.k, int) + assert res.n == ref.n + assert res.k == ref.k + else: + xp_assert_close(res.n, xp.asarray(int(ref.n))) + xp_assert_close(res.k, xp.asarray(int(ref.k))) + + if (is_array_api_strict(xp) or is_jax(xp)) and method == 'exact': + # array API strict and JAX don't support exact CI right now + return + + confidence_level = xp.asarray(0.95, dtype=dtype) + res = res.proportion_ci(method=method, confidence_level=confidence_level) + ref = ref.proportion_ci(method=method, confidence_level=0.95) + + xp_assert_close(res.low, xp.asarray(float(ref.low), dtype=dtype)) + xp_assert_close(res.high, xp.asarray(float(ref.high), dtype=dtype)) + + @pytest.mark.parametrize("shape", [(), (7, 8, 9)]) + @pytest.mark.parametrize("alternative", ["less", "greater", "two-sided"]) + @pytest.mark.parametrize("method", ["exact", "wilson", "wilsoncc"]) + def test_ndarray(self, shape, alternative, method, xp): + rng = np.random.default_rng(2150248640) + k = rng.integers(-1, 11, size=shape) + n = rng.integers(-1, 11, size=shape) + p = rng.uniform(-0.1, 1.1, size=shape) + + if is_jax(xp) and alternative=='two-sided': + pytest.skip("`alternative='two-sided'` is incompatible with JAX arrays.") + + res = stats.binomtest(xp.asarray(k), xp.asarray(n), xp.asarray(p), + alternative=alternative) + + @np.vectorize(excluded='alternative') + def binomtest_1d(k, n, p, alternative): + ref = stats.binomtest(k, n, p, alternative=alternative) + ci = ref.proportion_ci(method=method) + return ref.k, ref.n, ref.statistic, ref.pvalue, ci.low, ci.high + + ref_k, ref_n, ref_statistic, ref_pvalue, ci_low, ci_high = binomtest_1d( + k, n, p, alternative) + xp_assert_close(res.k, xp.asarray(ref_k)) + xp_assert_close(res.n, xp.asarray(ref_n)) + xp_assert_close(res.statistic, xp.asarray(ref_statistic)) + xp_assert_close(res.pvalue, xp.asarray(ref_pvalue)) + + if (is_array_api_strict(xp) or is_jax(xp)) and method == 'exact': + # array API strict and JAX don't support exact CI right now + return + + ci = res.proportion_ci(method=method) + xp_assert_close(ci.low, xp.asarray(ci_low)) + xp_assert_close(ci.high, xp.asarray(ci_high)) + + @skip_xp_backends('jax.numpy', reason="'two-sided' alternative needs root finder") + def test_binomtest(self, xp): + # precision tests compared to R for ticket:986 + pp = xp.concat((xp.linspace(0.1, 0.2, 5), + xp.linspace(0.45, 0.65, 5), + xp.linspace(0.85, 0.95, 5))) + n = 501 + x = 450 + results = [0.0, 0.0, 1.0159969301994141e-304, + 2.9752418572150531e-275, 7.7668382922535275e-250, + 2.3381250925167094e-099, 7.8284591587323951e-081, + 9.9155947819961383e-065, 2.8729390725176308e-050, + 1.7175066298388421e-037, 0.0021070691951093692, + 0.12044570587262322, 0.88154763174802508, 0.027120993063129286, + 2.6102587134694721e-006] + + xp_assert_close(stats.binomtest(x, n, pp).pvalue, xp.asarray(results)) + + xp_assert_close(stats.binomtest(50, 100, xp.asarray(0.1)).pvalue, + xp.asarray(5.8320387857343647e-024)) + + @skip_xp_backends('jax.numpy', reason="'two-sided' alternative needs root finder") + def test_binomtest2(self, xp): + # test added for issue #2384 + res2 = [ + [1.0, 1.0], + [0.5, 1.0, 0.5], + [0.25, 1.00, 1.00, 0.25], + [0.125, 0.625, 1.000, 0.625, 0.125], + [0.0625, 0.3750, 1.0000, 1.0000, 0.3750, 0.0625], + [0.03125, 0.21875, 0.68750, 1.00000, 0.68750, 0.21875, 0.03125], + [0.015625, 0.125000, 0.453125, 1.000000, 1.000000, 0.453125, 0.125000, + 0.015625], + [0.0078125, 0.0703125, 0.2890625, 0.7265625, 1.0000000, 0.7265625, + 0.2890625, 0.0703125, 0.0078125], + [0.00390625, 0.03906250, 0.17968750, 0.50781250, 1.00000000, + 1.00000000, 0.50781250, 0.17968750, 0.03906250, 0.00390625], + [0.001953125, 0.021484375, 0.109375000, 0.343750000, 0.753906250, + 1.000000000, 0.753906250, 0.343750000, 0.109375000, 0.021484375, + 0.001953125] + ] + for k in range(1, 11): + v = xp.arange(k+1, dtype=xp.float64) + res1 = stats.binomtest(v, k, 0.5).pvalue + xp_assert_close(res1, xp.asarray(res2[k-1], dtype=xp.float64)) + + @skip_xp_backends('jax.numpy', reason="'two-sided' alternative needs root finder") + def test_binomtest3(self, xp): + # test added for issue #2384 + # test when x == n*p and neighbors + v = xp.arange(1., 11.)[:, xp.newaxis] + k = xp.arange(2., 11.) + shape = (v.shape[0], k.shape[0]) + + res3 = stats.binomtest(v, v*k, 1./k).pvalue + xp_assert_close(res3, xp.ones(shape)) + + # > bt=c() + # > for(i in as.single(1:10)) { + # + for(k in as.single(2:10)) { + # + bt = c(bt, binom.test(i-1, k*i,(1/k))$p.value); + # + print(c(i+1, k*i,(1/k))) + # + } + # + } + binom_testm1 = xp.asarray([ + 0.5, 0.5555555555555556, 0.578125, 0.5904000000000003, + 0.5981224279835393, 0.603430543396034, 0.607304096221924, + 0.610255656871054, 0.612579511000001, 0.625, 0.670781893004115, + 0.68853759765625, 0.6980101120000006, 0.703906431368616, + 0.70793209416498, 0.7108561134173507, 0.713076544331419, + 0.714820192935702, 0.6875, 0.7268709038256367, 0.7418963909149174, + 0.74986110468096, 0.7548015520398076, 0.7581671424768577, + 0.760607984787832, 0.762459425024199, 0.7639120677676575, 0.7265625, + 0.761553963657302, 0.774800934828818, 0.7818005980538996, + 0.78613491480358, 0.789084353140195, 0.7912217659828884, + 0.79284214559524, 0.794112956558801, 0.75390625, 0.7856929451142176, + 0.7976688481430754, 0.8039848974727624, 0.807891868948366, + 0.8105487660137676, 0.812473307174702, 0.8139318233591120, + 0.815075399104785, 0.7744140625, 0.8037322594985427, + 0.814742863657656, 0.8205425178645808, 0.8241275984172285, + 0.8265645374416, 0.8283292196088257, 0.829666291102775, + 0.8307144686362666, 0.7905273437499996, 0.8178712053954738, + 0.828116983756619, 0.833508948940494, 0.8368403871552892, + 0.839104213210105, 0.840743186196171, 0.84198481438049, + 0.8429580531563676, 0.803619384765625, 0.829338573944648, + 0.8389591907548646, 0.84401876783902, 0.84714369697889, + 0.8492667010581667, 0.850803474598719, 0.851967542858308, + 0.8528799045949524, 0.8145294189453126, 0.838881732845347, + 0.847979024541911, 0.852760894015685, 0.8557134656773457, + 0.8577190131799202, 0.85917058278431, 0.860270010472127, + 0.861131648404582, 0.823802947998047, 0.846984756807511, + 0.855635653643743, 0.860180994825685, 0.86298688573253, + 0.864892525675245, 0.866271647085603, 0.867316125625004, + 0.8681346531755114 + ]) + + # > bt=c() + # > for(i in as.single(1:10)) { + # + for(k in as.single(2:10)) { + # + bt = c(bt, binom.test(i+1, k*i,(1/k))$p.value); + # + print(c(i+1, k*i,(1/k))) + # + } + # + } + + binom_testp1 = xp.asarray([ + 0.5, 0.259259259259259, 0.26171875, 0.26272, 0.2632244513031551, + 0.2635138663069203, 0.2636951804161073, 0.2638162407564354, + 0.2639010709000002, 0.625, 0.4074074074074074, 0.42156982421875, + 0.4295746560000003, 0.43473045988554, 0.4383309503172684, + 0.4409884859402103, 0.4430309389962837, 0.444649849401104, 0.6875, + 0.4927602499618962, 0.5096031427383425, 0.5189636628480, + 0.5249280070771274, 0.5290623300865124, 0.5320974248125793, + 0.5344204730474308, 0.536255847400756, 0.7265625, 0.5496019313526808, + 0.5669248746708034, 0.576436455045805, 0.5824538812831795, + 0.5866053321547824, 0.589642781414643, 0.5919618019300193, + 0.593790427805202, 0.75390625, 0.590868349763505, 0.607983393277209, + 0.617303847446822, 0.623172512167948, 0.627208862156123, + 0.6301556891501057, 0.632401894928977, 0.6341708982290303, + 0.7744140625, 0.622562037497196, 0.639236102912278, 0.648263335014579, + 0.65392850011132, 0.657816519817211, 0.660650782947676, + 0.662808780346311, 0.6645068560246006, 0.7905273437499996, + 0.6478843304312477, 0.6640468318879372, 0.6727589686071775, + 0.6782129857784873, 0.681950188903695, 0.684671508668418, + 0.686741824999918, 0.688369886732168, 0.803619384765625, + 0.668716055304315, 0.684360013879534, 0.6927642396829181, + 0.6980155964704895, 0.701609591890657, 0.7042244320992127, + 0.7062125081341817, 0.707775152962577, 0.8145294189453126, + 0.686243374488305, 0.7013873696358975, 0.709501223328243, + 0.714563595144314, 0.718024953392931, 0.7205416252126137, + 0.722454130389843, 0.723956813292035, 0.823802947998047, + 0.701255953767043, 0.715928221686075, 0.723772209289768, + 0.7286603031173616, 0.7319999279787631, 0.7344267920995765, + 0.736270323773157, 0.737718376096348 + ]) + + k, v = xp.asarray(k, dtype=xp.float64), xp.asarray(v, dtype=xp.float64) + binom_testp1 = xp.reshape(xp.asarray(binom_testp1, dtype=xp.float64), shape) + binom_testm1 = xp.reshape(xp.asarray(binom_testm1, dtype=xp.float64), shape) + + res4_p1 = stats.binomtest(v+1, v*k, 1./k).pvalue + res4_m1 = stats.binomtest(v-1, v*k, 1./k).pvalue + + xp_assert_close(res4_p1, binom_testp1) + xp_assert_close(res4_m1, binom_testm1) + + +@make_xp_test_case(stats.fligner) +class TestFligner: + + def _perturb(self, g, rng=124987234782812): + # g arrays have ties to which statistic is very sensitive; break them + rng = np.random.default_rng(rng) + return (np.asarray(g) + 1e-10 * rng.standard_normal(len(g))).tolist() + + @pytest.mark.parametrize('dtype', [None, 'float32', 'float64']) + def test_data(self, dtype, xp): + # numbers from R: fligner.test in package stats + dtype = xp_default_dtype(xp) if dtype is None else getattr(xp, dtype) + x1 = xp.arange(5, dtype=dtype) + res = stats.fligner(x1, x1**2) + ref = (xp.asarray(3.2282229927203536, dtype=dtype), + xp.asarray(0.072379187848207877, dtype=dtype)) + xp_assert_close(res[0], ref[0]) + xp_assert_close(res[1], ref[1]) + + def test_trimmed1(self, xp): + # Perturb input to break ties in the transformed data + # See https://github.com/scipy/scipy/pull/8042 for more details + rng = np.random.default_rng(9952379681) + g1_ = xp.asarray(self._perturb(g1, rng=rng)) + g2_ = xp.asarray(self._perturb(g2, rng=rng)) + g3_ = xp.asarray(self._perturb(g3, rng=rng)) + # Test that center='trimmed' gives the same result as center='mean' + # when proportiontocut=0. + Xsq1, pval1 = stats.fligner(g1_, g2_, g3_, center='mean') + Xsq2, pval2 = stats.fligner(g1_, g2_, g3_, center='trimmed', + proportiontocut=0.0) + xp_assert_close(Xsq1, Xsq2) + xp_assert_close(pval1, pval2) + + @skip_xp_backends(np_only=True, reason="inconsistent tie-breaking across backends") + def test_trimmed_nonregression(self, xp): + # This is a non-regression test + # Expected results are *not* from an external gold standard, + # we're just making sure the results remain consistent + # in the future in case of changes + args = [g1, g2, g3, g4, g5, g6, g7, g8, g9, g10] + args = [xp.asarray(arg, dtype=xp.float64) for arg in args] + W, pval = stats.fligner(*args, center="trimmed", proportiontocut=0.25) + xp_assert_close(W, xp.asarray(15.953569890010614), rtol=5e-14) + xp_assert_close(pval, xp.asarray(0.06785752327432863), rtol=5e-14) + + def test_trimmed_consistency(self, xp): + # Tests for consistency across multiple backends when ties are broken + rng = np.random.default_rng(4839206199) + args = [g1, g2, g3, g4, g5, g6, g7, g8, g9, g10] + args = [self._perturb(arg, rng=rng) for arg in args] + ref = stats.fligner(*args, center="trimmed", proportiontocut=0.25) + args = [xp.asarray(arg, dtype=xp.float64) for arg in args] + res = stats.fligner(*args, center="trimmed", proportiontocut=0.25) + xp_assert_close(res.statistic, xp.asarray(ref.statistic), rtol=5e-14) + xp_assert_close(res.pvalue, xp.asarray(ref.pvalue), rtol=5e-14) + + def test_bad_center_value(self, xp): + x = xp.linspace(-1, 1, 21) + message = "center must be 'mean', 'median' or 'trimmed'." + with pytest.raises(ValueError, match=message): + stats.fligner(x, x, center='trim') + + def test_bad_num_args(self, xp): + # Too few args raises ValueError. + message = "Must provide at least two samples." + with pytest.raises(ValueError, match=message): + stats.fligner(xp.asarray([1, 2])) + + def test_empty_arg(self, xp): + x = xp.arange(5.) + with pytest.warns(SmallSampleWarning, match=too_small_1d_not_omit): + res = stats.fligner(x, x**2, xp.asarray([])) + xp_assert_equal(res.statistic, xp.asarray(xp.nan)) + xp_assert_equal(res.pvalue, xp.asarray(xp.nan)) + + +def mood_cases_with_ties(): + # Generate random `x` and `y` arrays with ties both between and within the + # samples. Expected results are (statistic, pvalue) from SAS. + expected_results = [(-1.76658511464992, .0386488678399305), + (-.694031428192304, .2438312498647250), + (-1.15093525352151, .1248794365836150)] + seeds = [23453254, 1298352315, 987234597] + for si, seed in enumerate(seeds): + rng = np.random.default_rng(seed) + xy = rng.random(100) + # Generate random indices to make ties + tie_ind = rng.integers(low=0, high=99, size=5) + # Generate a random number of ties for each index. + num_ties_per_ind = rng.integers(low=1, high=5, size=5) + # At each `tie_ind`, mark the next `n` indices equal to that value. + for i, n in zip(tie_ind, num_ties_per_ind): + for j in range(i + 1, i + n): + xy[j] = xy[i] + # scramble order of xy before splitting into `x, y` + rng.shuffle(xy) + x, y = np.split(xy, 2) + yield x, y, 'less', *expected_results[si] + + +@make_xp_test_case(stats.mood) +class TestMood: + @pytest.mark.parametrize("x,y,alternative,stat_expect,p_expect", + list(mood_cases_with_ties())) + def test_against_SAS(self, x, y, alternative, stat_expect, p_expect, xp): + """ + Example code used to generate SAS output: + DATA myData; + INPUT X Y; + CARDS; + 1 0 + 1 1 + 1 2 + 1 3 + 1 4 + 2 0 + 2 1 + 2 4 + 2 9 + 2 16 + ods graphics on; + proc npar1way mood data=myData ; + class X; + ods output MoodTest=mt; + proc contents data=mt; + proc print data=mt; + format Prob1 17.16 Prob2 17.16 Statistic 17.16 Z 17.16 ; + title "Mood Two-Sample Test"; + proc print data=myData; + title "Data for above results"; + run; + """ + x, y = xp.asarray(x.tolist()), xp.asarray(y.tolist()) + statistic, pvalue = stats.mood(x, y, alternative=alternative) + xp_assert_close(statistic, xp.asarray(stat_expect), atol=1e-16) + xp_assert_close(pvalue, xp.asarray(p_expect), atol=1e-16) + + @pytest.mark.parametrize("dtype", [None, 'float32', 'float64']) + @pytest.mark.parametrize("alternative, expected", + [('two-sided', (1.019938533549930, + .3077576129778760)), + ('less', (1.019938533549930, + 1 - .1538788064889380)), + ('greater', (1.019938533549930, + .1538788064889380))]) + def test_against_SAS_2(self, dtype, alternative, expected, xp): + # Code to run in SAS in above function + dtype = xp_default_dtype(xp) if dtype is None else getattr(xp, dtype) + x = [111, 107, 100, 99, 102, 106, 109, 108, 104, 99, + 101, 96, 97, 102, 107, 113, 116, 113, 110, 98] + y = [107, 108, 106, 98, 105, 103, 110, 105, 104, 100, + 96, 108, 103, 104, 114, 114, 113, 108, 106, 99] + x, y = xp.asarray(x, dtype=dtype), xp.asarray(y, dtype=dtype) + res = stats.mood(x, y, alternative=alternative) + xp_assert_close(res.statistic, xp.asarray(expected[0], dtype=dtype)) + xp_assert_close(res.pvalue, xp.asarray(expected[1], dtype=dtype)) + + def test_mood_order_of_args(self, xp): + # z should change sign when the order of arguments changes, pvalue + # should not change + rng = np.random.default_rng(4058827756) + x1 = xp.asarray(rng.standard_normal((10, 1))) + x2 = xp.asarray(rng.standard_normal((15, 1))) + z1, p1 = stats.mood(x1, x2) + z2, p2 = stats.mood(x2, x1) + xp_assert_close(z2, -z1) + xp_assert_close(p2, p1) + + def test_mood_with_axis_none(self, xp): + # Test with axis = None, compare with results from R + x1 = [-0.626453810742332, 0.183643324222082, -0.835628612410047, + 1.59528080213779, 0.329507771815361, -0.820468384118015, + 0.487429052428485, 0.738324705129217, 0.575781351653492, + -0.305388387156356, 1.51178116845085, 0.389843236411431, + -0.621240580541804, -2.2146998871775, 1.12493091814311, + -0.0449336090152309, -0.0161902630989461, 0.943836210685299, + 0.821221195098089, 0.593901321217509] + + x2 = [-0.896914546624981, 0.184849184646742, 1.58784533120882, + -1.13037567424629, -0.0802517565509893, 0.132420284381094, + 0.707954729271733, -0.23969802417184, 1.98447393665293, + -0.138787012119665, 0.417650750792556, 0.981752777463662, + -0.392695355503813, -1.03966897694891, 1.78222896030858, + -2.31106908460517, 0.878604580921265, 0.035806718015226, + 1.01282869212708, 0.432265154539617, 2.09081920524915, + -1.19992581964387, 1.58963820029007, 1.95465164222325, + 0.00493777682814261, -2.45170638784613, 0.477237302613617, + -0.596558168631403, 0.792203270299649, 0.289636710177348] + + x1 = xp.reshape(xp.asarray(x1), (10, 2)) + x2 = xp.reshape(xp.asarray(x2), (15, 2)) + res = stats.mood(x1, x2, axis=None) + xp_assert_close(res.statistic, xp.asarray(-1.31716607555)) + xp_assert_close(res.pvalue, xp.asarray(0.18778296257)) + + @pytest.mark.parametrize('rng_method, args', [('standard_normal', tuple()), + ('integers', (8,))]) + def test_mood_2d(self, rng_method, args, xp): + # Test if the results of mood test in 2-D vectorized call are consistent + # with result when looping over the slices. + ny = 5 + rng = np.random.default_rng(3602349075) + rng_method = getattr(rng, rng_method) + x1 = rng_method(*args, size=(10, ny)) + x2 = rng_method(*args, size=(15, ny)) + dtype = xp.float64 + res = stats.mood(xp.asarray(x1, dtype=dtype), xp.asarray(x2, dtype=dtype)) + + for j in range(ny): + ref = stats.mood(x1[:, j], x2[:, j]) + xp_assert_close(res.statistic[j], xp.asarray(ref.statistic)) + xp_assert_close(res.pvalue[j], xp.asarray(ref.pvalue)) + + # inverse order of dimensions + x1 = x1.transpose() + x2 = x2.transpose() + res = stats.mood(xp.asarray(x1, dtype=dtype), xp.asarray(x2, dtype=dtype), + axis=1) + + for i in range(ny): + # check axis handling is self consistent + ref = stats.mood(x1[i, :], x2[i, :]) + xp_assert_close(res.statistic[i], xp.asarray(ref.statistic), atol=1e-14) + xp_assert_close(res.pvalue[i], xp.asarray(ref.pvalue)) + + @pytest.mark.parametrize('rng_method, args', [('standard_normal', tuple()), + ('integers', (8,))]) + def test_mood_3d(self, rng_method, args, xp): + shape = (10, 5, 6) + rng = np.random.default_rng(3602349075) + rng_method = getattr(rng, rng_method) + x1 = xp.asarray(rng_method(*args, size=shape)) + x2 = xp.asarray(rng_method(*args, size=shape)) + + for axis in range(3): + res = stats.mood(x1, x2, axis=axis) + # Tests that result for 3-D arrays is equal to that for the + # same calculation on a set of 1-D arrays taken from the + # 3-D array + axes_idx = ([1, 2], [0, 2], [0, 1]) # the two axes != axis + for i in range(shape[axes_idx[axis][0]]): + for j in range(shape[axes_idx[axis][1]]): + if axis == 0: + slice1 = x1[:, i, j] + slice2 = x2[:, i, j] + elif axis == 1: + slice1 = x1[i, :, j] + slice2 = x2[i, :, j] + else: + slice1 = x1[i, j, :] + slice2 = x2[i, j, :] + + ref = stats.mood(slice1, slice2) + xp_assert_close(res.statistic[i, j], ref.statistic, atol=1e-15) + xp_assert_close(res.pvalue[i, j], ref.pvalue) + + def test_mood_bad_arg(self, xp): + # Warns when the sum of the lengths of the args is less than 3 + with eager_warns(SmallSampleWarning, match=too_small_1d_not_omit, xp=xp): + res = stats.mood(xp.asarray([1.]), xp.asarray([])) + xp_assert_equal(res.statistic, xp.asarray(np.nan)) + xp_assert_equal(res.pvalue, xp.asarray(np.nan)) + + @pytest.mark.parametrize("dtype", [None, 'float32', 'float64']) + def test_mood_alternative(self, dtype, xp): + dtype = xp_default_dtype(xp) if dtype is None else getattr(xp, dtype) + + rng = np.random.RandomState(0) + x = stats.norm.rvs(scale=0.75, size=100, random_state=rng) + y = stats.norm.rvs(scale=1.25, size=100, random_state=rng) + x, y = xp.asarray(x, dtype=dtype), xp.asarray(y, dtype=dtype) + + stat1, p1 = stats.mood(x, y, alternative='two-sided') + stat2, p2 = stats.mood(x, y, alternative='less') + stat3, p3 = stats.mood(x, y, alternative='greater') + + assert stat1 == stat2 == stat3 + xp_assert_close(p1, xp.asarray(0., dtype=dtype), atol=1e-7) + xp_assert_close(p2, xp.asarray(p1/2, dtype=dtype)) + xp_assert_close(p3, xp.asarray(1 - p1/2, dtype=dtype)) + + with pytest.raises(ValueError, match="`alternative` must be..."): + stats.mood(x, y, alternative='ekki-ekki') + + +class TestProbplot: + + def test_basic(self): + x = stats.norm.rvs(size=20, random_state=12345) + osm, osr = stats.probplot(x, fit=False) + osm_expected = [-1.8241636, -1.38768012, -1.11829229, -0.91222575, + -0.73908135, -0.5857176, -0.44506467, -0.31273668, + -0.18568928, -0.06158146, 0.06158146, 0.18568928, + 0.31273668, 0.44506467, 0.5857176, 0.73908135, + 0.91222575, 1.11829229, 1.38768012, 1.8241636] + assert_allclose(osr, np.sort(x)) + assert_allclose(osm, osm_expected) + + res, res_fit = stats.probplot(x, fit=True) + res_fit_expected = [1.05361841, 0.31297795, 0.98741609] + assert_allclose(res_fit, res_fit_expected) + + def test_sparams_keyword(self): + x = stats.norm.rvs(size=100, random_state=123456) + # Check that None, () and 0 (loc=0, for normal distribution) all work + # and give the same results + osm1, osr1 = stats.probplot(x, sparams=None, fit=False) + osm2, osr2 = stats.probplot(x, sparams=0, fit=False) + osm3, osr3 = stats.probplot(x, sparams=(), fit=False) + assert_allclose(osm1, osm2) + assert_allclose(osm1, osm3) + assert_allclose(osr1, osr2) + assert_allclose(osr1, osr3) + # Check giving (loc, scale) params for normal distribution + osm, osr = stats.probplot(x, sparams=(), fit=False) + + def test_dist_keyword(self): + x = stats.norm.rvs(size=20, random_state=12345) + osm1, osr1 = stats.probplot(x, fit=False, dist='t', sparams=(3,)) + osm2, osr2 = stats.probplot(x, fit=False, dist=stats.t, sparams=(3,)) + assert_allclose(osm1, osm2) + assert_allclose(osr1, osr2) + + assert_raises(ValueError, stats.probplot, x, dist='wrong-dist-name') + assert_raises(AttributeError, stats.probplot, x, dist=[]) + + class custom_dist: + """Some class that looks just enough like a distribution.""" + def ppf(self, q): + return stats.norm.ppf(q, loc=2) + + osm1, osr1 = stats.probplot(x, sparams=(2,), fit=False) + osm2, osr2 = stats.probplot(x, dist=custom_dist(), fit=False) + assert_allclose(osm1, osm2) + assert_allclose(osr1, osr2) + + @pytest.mark.skipif(not have_matplotlib, reason="no matplotlib") + def test_plot_kwarg(self): + fig = plt.figure() + fig.add_subplot(111) + x = stats.t.rvs(3, size=100, random_state=7654321) + res1, fitres1 = stats.probplot(x, plot=plt) + plt.close() + res2, fitres2 = stats.probplot(x, plot=None) + res3 = stats.probplot(x, fit=False, plot=plt) + plt.close() + res4 = stats.probplot(x, fit=False, plot=None) + # Check that results are consistent between combinations of `fit` and + # `plot` keywords. + assert_(len(res1) == len(res2) == len(res3) == len(res4) == 2) + assert_allclose(res1, res2) + assert_allclose(res1, res3) + assert_allclose(res1, res4) + assert_allclose(fitres1, fitres2) + + # Check that a Matplotlib Axes object is accepted + fig = plt.figure() + ax = fig.add_subplot(111) + stats.probplot(x, fit=False, plot=ax) + plt.close() + + def test_probplot_bad_args(self): + # Raise ValueError when given an invalid distribution. + assert_raises(ValueError, stats.probplot, [1], dist="plate_of_shrimp") + + def test_empty(self): + assert_equal(stats.probplot([], fit=False), + (np.array([]), np.array([]))) + assert_equal(stats.probplot([], fit=True), + ((np.array([]), np.array([])), + (np.nan, np.nan, 0.0))) + + def test_array_of_size_one(self): + message = "One or more sample arguments is too small..." + with (np.errstate(invalid='ignore'), + pytest.warns(SmallSampleWarning, match=message)): + assert_equal(stats.probplot([1], fit=True), + ((np.array([0.]), np.array([1])), + (np.nan, np.nan, np.nan))) + + +@make_xp_test_case(stats.wilcoxon) +class TestWilcoxon: + def test_wilcoxon_bad_arg(self, xp): + # Raise ValueError when two args of different lengths are given or + # zero_method is unknown. + x = xp.asarray([1, 2]) + d = xp.asarray([1]*10) + message = "`zero_method` must be one of..." + with pytest.raises(ValueError, match=message): + stats.wilcoxon(x, x, "dummy...") + message = "`alternative` must be one of" + with pytest.raises(ValueError, match=message): + stats.wilcoxon(x, x, alternative="dummy") + message = "`method` must be one of..." + with pytest.raises(ValueError, match=message): + stats.wilcoxon(d, method="xyz") + + def test_zero_diff(self, xp): + x = xp.arange(20) + # pratt and wilcox do not work if x - y == 0 and method == "asymptotic" + # => warning may be emitted and p-value is nan + with np.errstate(invalid="ignore"): + w, p = stats.wilcoxon(x, x, "wilcox", method="asymptotic") + xp_assert_equal(w, xp.asarray(0.0)) + xp_assert_equal(p, xp.asarray(xp.nan)) + w, p = stats.wilcoxon(x, x, "pratt", method="asymptotic") + xp_assert_equal(w, xp.asarray(0.0)) + xp_assert_equal(p, xp.asarray(xp.nan)) + # ranksum is n*(n+1)/2, split in half if zero_method == "zsplit" + w, p = stats.wilcoxon(x, x, "zsplit", method="asymptotic") + xp_assert_equal(w, xp.asarray(20*21/4)) + xp_assert_equal(p, xp.asarray(1.0)) + + def test_pratt(self, xp): + # regression test for gh-6805: p-value matches value from R package + # coin (wilcoxsign_test) reported in the issue + x = xp.asarray([1, 2, 3, 4]) + y = xp.asarray([1, 2, 3, 5]) + res = stats.wilcoxon(x, y, zero_method="pratt", method="asymptotic", + correction=False) + xp_assert_close(res.statistic, xp.asarray(0.0)) + xp_assert_close(res.pvalue, xp.asarray(0.31731050786291415)) + + def test_wilcoxon_arg_type(self): + # Should be able to accept list as arguments. + # Address issue 6070. + arr = [1, 2, 3, 0, -1, 3, 1, 2, 1, 1, 2] + + _ = stats.wilcoxon(arr, zero_method="pratt", method="asymptotic") + _ = stats.wilcoxon(arr, zero_method="zsplit", method="asymptotic") + _ = stats.wilcoxon(arr, zero_method="wilcox", method="asymptotic") + + def test_accuracy_wilcoxon(self, xp): + freq = [1, 4, 16, 15, 8, 4, 5, 1, 2] + nums = range(-4, 5) + x = xp.asarray(np.concatenate([[u] * v for u, v in zip(nums, freq)])) + y = xp.zeros_like(x) + + T, p = stats.wilcoxon(x, y, "pratt", method="asymptotic", + correction=False) + xp_assert_close(T, xp.asarray(423.)) + xp_assert_close(p, xp.asarray(0.0031724568006762576)) + + T, p = stats.wilcoxon(x, y, "zsplit", method="asymptotic", + correction=False) + xp_assert_equal(T, xp.asarray(441.)) + xp_assert_close(p, xp.asarray(0.0032145343172473055)) + + T, p = stats.wilcoxon(x, y, "wilcox", method="asymptotic", + correction=False) + xp_assert_equal(T, xp.asarray(327.)) + xp_assert_close(p, xp.asarray(0.00641346115861)) + + # Test the 'correction' option, using values computed in R with: + # > options(digits=16) + # > wilcox.test(x, y, paired=TRUE, exact=FALSE, correct={FALSE,TRUE}) + x = xp.asarray([120, 114, 181, 188, 180, 146, 121, 191, 132, 113, 127, 112]) + y = xp.asarray([133, 143, 119, 189, 112, 199, 198, 113, 115, 121, 142, 187]) + T, p = stats.wilcoxon(x, y, correction=False, method="asymptotic") + xp_assert_equal(T, xp.asarray(34.)) + xp_assert_close(p, xp.asarray(0.6948866023724735)) + T, p = stats.wilcoxon(x, y, correction=True, method="asymptotic") + xp_assert_equal(T, xp.asarray(34.)) + xp_assert_close(p, xp.asarray(0.7240816609153895)) + + @pytest.mark.parametrize("kwarg", + [{"method": "approx"}, {"mode": "approx"}, {"mode": "asymptotic"}]) + def test_approx_mode(self, kwarg, xp): + # Check that `mode` is still an alias of keyword `method`, + # and `"approx"` is still an alias of argument `"asymptotic"` + x = xp.asarray([3, 5, 23, 7, 243, 58, 98, 2, 8, -3, 9, 11]) + y = xp.asarray([2, -2, 1, 23, 0, 5, 12, 18, 99, 12, 17, 27]) + res = stats.wilcoxon(x, y, "wilcox", **kwarg) + ref = stats.wilcoxon(x, y, "wilcox", method="asymptotic") + xp_assert_equal(res.statistic, ref.statistic) + xp_assert_equal(res.pvalue, ref.pvalue) + + def test_wilcoxon_result_attributes(self, xp): + x = xp.asarray([120, 114, 181, 188, 180, 146, 121, 191, 132, 113, 127, 112]) + y = xp.asarray([133, 143, 119, 189, 112, 199, 198, 113, 115, 121, 142, 187]) + res = stats.wilcoxon(x, y, correction=False, method="asymptotic") + attributes = ('statistic', 'pvalue') + check_named_results(res, attributes, xp=xp) + + def test_wilcoxon_has_zstatistic(self, xp): + rng = np.random.default_rng(89426135444) + x, y = rng.random(15), rng.random(15) + x, y = xp.asarray(x), xp.asarray(y) + + res = stats.wilcoxon(x, y, method="asymptotic") + ref = special.ndtri(res.pvalue/2) + xp_assert_close(res.zstatistic, ref) + + if is_jax(xp): + return + + res = stats.wilcoxon(x, y, method="exact") + assert not hasattr(res, 'zstatistic') + + res = stats.wilcoxon(x, y) + assert not hasattr(res, 'zstatistic') + + def test_wilcoxon_tie(self, xp): + # Regression test for gh-2391. + # Corresponding R code is: + # > result = wilcox.test(rep(0.1, 10), exact=FALSE, correct=FALSE) + # > result$p.value + # [1] 0.001565402258002551 + # > result = wilcox.test(rep(0.1, 10), exact=FALSE, correct=TRUE) + # > result$p.value + # [1] 0.00190419504300439 + d = xp.asarray([0.1] * 10) + expected_stat = xp.asarray(0.0) + + stat, p = stats.wilcoxon(d, method="asymptotic", correction=False) + expected_p = xp.asarray(0.001565402258002551) + xp_assert_equal(stat, expected_stat) + xp_assert_close(p, expected_p) + + stat, p = stats.wilcoxon(d, method="asymptotic", correction=True) + expected_p = xp.asarray(0.00190419504300439) + xp_assert_equal(stat, expected_stat) + xp_assert_close(p, expected_p) + + @pytest.mark.parametrize('dtype', [None, 'float32', 'float64']) + def test_onesided(self, dtype, xp): + # tested against "R version 4.0.3 (2020-10-10)" + # + # x <- c(125, 115, 130, 140, 140, 115, 140, 125, 140, 135) + # y <- c(110, 122, 125, 120, 140, 124, 123, 137, 135, 145) + # cfg <- list(x = x, y = y, paired = TRUE, exact = FALSE) + # do.call(wilcox.test, c(cfg, list(alternative = "less", correct = FALSE))) + # do.call(wilcox.test, c(cfg, list(alternative = "less", correct = TRUE))) + # do.call(wilcox.test, c(cfg, list(alternative = "greater", correct = FALSE))) + # do.call(wilcox.test, c(cfg, list(alternative = "greater", correct = TRUE))) + + dtype = dtype if dtype is None else getattr(xp, dtype) + x = xp.asarray([125, 115, 130, 140, 140, 115, 140, 125, 140, 135], dtype=dtype) + y = xp.asarray([110, 122, 125, 120, 140, 124, 123, 137, 135, 145], dtype=dtype) + + w_ref = xp.asarray(27.0, dtype=dtype) + w, p = stats.wilcoxon(x, y, alternative="less", method="asymptotic", + correction=False) + xp_assert_equal(w, w_ref) + xp_assert_close(p, xp.asarray(0.7031847042787, dtype=dtype)) + + w, p = stats.wilcoxon(x, y, alternative="less", correction=True, + method="asymptotic") + xp_assert_equal(w, w_ref) + xp_assert_close(p, xp.asarray(0.72336564289, dtype=dtype)) + + w, p = stats.wilcoxon(x, y, alternative="greater", + method="asymptotic", correction=False) + xp_assert_equal(w, w_ref) + xp_assert_close(p, xp.asarray(0.2968152957213, dtype=dtype)) + + w, p = stats.wilcoxon(x, y, alternative="greater", + correction=True, method="asymptotic") + xp_assert_equal(w, w_ref) + xp_assert_close(p, xp.asarray(0.3176446594176, dtype=dtype)) + + def test_exact_basic(self): # exact distribution is NumPy-only + for n in range(1, 51): + pmf1 = _get_wilcoxon_distr(n) + pmf2 = _get_wilcoxon_distr2(n) + assert_equal(n*(n+1)/2 + 1, len(pmf1)) + assert_equal(sum(pmf1), 1) + assert_array_almost_equal(pmf1, pmf2) + + @skip_xp_backends("jax.numpy", reason="`method='exact'` is incompatible with JAX") + @pytest.mark.parametrize('dtype', [None, 'float32', 'float64']) + def test_exact_pval(self, dtype, xp): + # expected values computed with "R version 4.0.3 (2020-10-10)" + # options(digits=16) + # x < - c(1.81, 0.82, 1.56, -0.48, 0.81, 1.28, -1.04, 0.23, -0.75, 0.14) + # y < - c(0.71, 0.65, -0.2, 0.85, -1.1, -0.45, -0.84, -0.24, -0.68, -0.76) + # result = wilcox.test(x - y, exact=TRUE, alternative='l', correct=TRUE) + # result$p.value + dtype = dtype if dtype is None else getattr(xp, dtype) + x = xp.asarray([1.81, 0.82, 1.56, -0.48, 0.81, 1.28, -1.04, 0.23, + -0.75, 0.14], dtype=dtype) + y = xp.asarray([0.71, 0.65, -0.2, 0.85, -1.1, -0.45, -0.84, -0.24, + -0.68, -0.76], dtype=dtype) + _, p = stats.wilcoxon(x, y, alternative="two-sided", method="exact") + xp_assert_close(p, xp.asarray(0.10546875, dtype=dtype)) + _, p = stats.wilcoxon(x, y, alternative="less", method="exact") + xp_assert_close(p, xp.asarray(0.95800781256, dtype=dtype)) + _, p = stats.wilcoxon(x, y, alternative="greater", method="exact") + xp_assert_close(p, xp.asarray(0.052734375, dtype=dtype)) + + x = xp.arange(0., 20., dtype=dtype) + 0.5 + y = xp.arange(20., 0., -1., dtype=dtype) + _, p = stats.wilcoxon(x, y, alternative="two-sided", method="exact") + xp_assert_close(p, xp.asarray(0.8694877624511719, dtype=dtype)) + _, p = stats.wilcoxon(x, y, alternative="less", method="exact") + xp_assert_close(p, xp.asarray(0.4347438812255859, dtype=dtype)) + _, p = stats.wilcoxon(x, y, alternative="greater", method="exact") + xp_assert_close(p, xp.asarray(0.5795888900756836, dtype=dtype)) + + # These inputs were chosen to give a W statistic that is either the + # center of the distribution (when the length of the support is odd), or + # the value to the left of the center (when the length of the support is + # even). Also, the numbers are chosen so that the W statistic is the + # sum of the positive values. + + @skip_xp_backends("jax.numpy", reason="`method='exact'` is incompatible with JAX") + @pytest.mark.parametrize('x', [[-1, -2, 3], + [-1, 2, -3, -4, 5], + [-1, -2, 3, -4, -5, -6, 7, 8]]) + def test_exact_p_1(self, x, xp): + w, p = stats.wilcoxon(xp.asarray(x)) + x = np.array(x) + wtrue = x[x > 0].sum() + xp_assert_equal(w, xp.asarray(float(wtrue))) + xp_assert_equal(p, xp.asarray(1.0)) + + def test_auto(self, xp): + # auto default to exact if there are no ties and n <= 50 + x = xp.arange(0., 50.) + 0.5 + y = xp.arange(50., 0., -1.) + + if is_jax(xp): # + message = "When using `wilcoxon` with `jax.numpy` arrays..." + with pytest.raises(ValueError, match=message): + stats.wilcoxon(x, y, method="auto") + with pytest.raises(ValueError, match=message): + stats.wilcoxon(x, y, method="exact") + return + + xp_assert_equal(stats.wilcoxon(x, y).pvalue, + stats.wilcoxon(x, y, method="exact").pvalue) + + # n <= 50: if there are zeros in d = x-y, use PermutationMethod + pm = stats.PermutationMethod() + d = np.arange(-2, 5) + w, p = stats.wilcoxon(d) + # rerunning the test gives the same results since n_resamples + # is large enough to get deterministic results if n <= 13 + # so we do not need to use a seed. to avoid longer runtimes of the + # test, use n=7 only. For n=13, see test_auto_permutation_edge_case + assert_equal((w, p), stats.wilcoxon(d, method=pm)) + + # for larger vectors (n > 13) with ties/zeros, use asymptotic test + d = xp.arange(-5, 9) # zero + _, p = stats.wilcoxon(d) + xp_assert_equal(stats.wilcoxon(d, method="asymptotic").pvalue, xp.asarray(p)) + + d = xpx.at(d)[d == 0].set(1) # tie + _, p = stats.wilcoxon(d) + xp_assert_equal(stats.wilcoxon(d, method="asymptotic").pvalue, xp.asarray(p)) + + # use approximation for samples > 50 + d = xp.arange(1, 52) + _, p = stats.wilcoxon(d) + xp_assert_equal(stats.wilcoxon(d, method="asymptotic").pvalue, xp.asarray(p)) + + @pytest.mark.xslow + @skip_xp_backends("jax.numpy", reason="lazy -> limited `method` choices") + def test_auto_permutation_edge_case(self, xp): + # Check that `PermutationMethod()` is used and results are deterministic when + # `method='auto'`, there are zeros or ties in `d = x-y`, and `len(d) <= 13`. + d = np.arange(-5, 8) # zero + res = stats.wilcoxon(xp.asarray(d)) + # generated with stats.wilcoxon(d, method=PermutationMethod()) + w, p = xp.asarray([27.5, 0.3955078125]) + xp_assert_equal(res.statistic, w) + xp_assert_equal(res.pvalue, p) + + d[d == 0] = 1 # tie + res = stats.wilcoxon(xp.asarray(d)) + # generated with stats.wilcoxon(d, method=PermutationMethod()) + w, p = xp.asarray([32, 0.3779296875]) + xp_assert_equal(res.statistic, w) + xp_assert_equal(res.pvalue, p) + + @pytest.mark.parametrize('size', [3, 5, 10]) + def test_permutation_method(self, size, xp): + rng = np.random.default_rng(92348034828501345) + x = rng.random(size=size).tolist() + res = stats.wilcoxon(xp.asarray(x), method=stats.PermutationMethod()) + ref = stats.wilcoxon(x, method='exact') # all backends test against NumPy + dtype = xp_default_dtype(xp) + xp_assert_equal(res.statistic, xp.asarray(ref.statistic, dtype=dtype)) + xp_assert_equal(res.pvalue, xp.asarray(ref.pvalue, dtype=dtype)) + + x = xp.asarray(rng.random(size=size*10)) + rng = np.random.default_rng(59234803482850134) + pm = stats.PermutationMethod(n_resamples=99, rng=rng) + ref = stats.wilcoxon(x, method=pm) + # preserve use of old random_state during SPEC 7 transition + rng = np.random.default_rng(59234803482850134) + pm = stats.PermutationMethod(n_resamples=99, random_state=rng) + res = stats.wilcoxon(x, method=pm) + + xp_assert_equal(xp.round(res.pvalue*100)/100, res.pvalue) # n_resamples used + xp_assert_equal(res.pvalue, ref.pvalue) # rng/random_state used + + @skip_xp_backends("jax.numpy", reason="`method='auto'` is incompatible with JAX") + def test_method_auto_nan_propagate_ND_length_gt_50_gh20591(self, xp): + # When method!='asymptotic', nan_policy='propagate', and a slice of + # a >1 dimensional array input contained NaN, the result object of + # `wilcoxon` could (under yet other conditions) return `zstatistic` + # for some slices but not others. This resulted in an error because + # `apply_along_axis` would have to create a ragged array. + # Check that this is resolved. + rng = np.random.default_rng(235889269872456) + A = rng.normal(size=(51, 2)) # length along slice > exact threshold + A[5, 1] = np.nan + A = xp.asarray(A) + res = stats.wilcoxon(A) + ref = stats.wilcoxon(A, method='asymptotic') + xp_assert_close(res.statistic, ref.statistic) + xp_assert_close(res.pvalue, ref.pvalue) + assert hasattr(ref, 'zstatistic') + assert not hasattr(res, 'zstatistic') + + @pytest.mark.parametrize('method', ['exact', 'asymptotic']) + def test_symmetry_gh19872_gh20752(self, method, xp): + if is_jax(xp) and method == 'exact': + pytest.skip("`method='exact'` is incompatible with JAX") + # Check that one-sided exact tests obey required symmetry. Bug reported + # in gh-19872 and again in gh-20752; example from gh-19872 is more concise: + var1 = xp.asarray([62, 66, 61, 68, 74, 62, 68, 62, 55, 59]) + var2 = xp.asarray([71, 71, 69, 61, 75, 71, 77, 72, 62, 65]) + ref = stats.wilcoxon(var1, var2, alternative='less', method=method) + res = stats.wilcoxon(var2, var1, alternative='greater', method=method) + max_statistic = var1.shape[0] * (var1.shape[0] + 1) / 2 + assert int(res.statistic) != res.statistic + xp_assert_close(max_statistic - res.statistic, ref.statistic) + xp_assert_close(res.pvalue, ref.pvalue) + + @pytest.mark.parametrize("method", ('exact', stats.PermutationMethod())) + def test_all_zeros_exact(self, method, xp): + if is_jax(xp) and method == 'exact': + pytest.skip("`method='exact'` is incompatible with JAX") + # previously, this raised a RuntimeWarning when calculating Z, even + # when the Z value was not needed. Confirm that this no longer + # occurs when `method` is 'exact' or a `PermutationMethod`. + res = stats.wilcoxon(xp.zeros(5), method=method) + xp_assert_close(res.statistic, xp.asarray(0.)) + xp_assert_close(res.pvalue, xp.asarray(1.)) + + def test_wilcoxon_axis_broadcasting_errors_gh22051(self, xp): + # In previous versions of SciPy, `wilcoxon` gave an incorrect error + # message when `AxisError` was not found in the base NumPy namespace. + # Check that this is resolved with and without the ANP decorator. + x = xp.asarray([1, 2, 3]) + y = xp.asarray([4, 5, 6]) + + message = "Array shapes are incompatible for broadcasting." + with pytest.raises(ValueError, match=message): + stats.wilcoxon(x, y[:-1]) + + if not is_numpy(xp): + return # hard to generalize the rest + + message = "operands could not be broadcast together with..." + with pytest.raises(ValueError, match=message): + stats.wilcoxon(x, y[:-1], _no_deco=True) + + AxisError = getattr(np, 'AxisError', None) or np.exceptions.AxisError + message = "source: axis 3 is out of bounds for array of dimension 1" + with pytest.raises(AxisError, match=message): + stats.wilcoxon(x, y, axis=3) + + message = "`axis` must be compatible with the shape..." + with pytest.raises(AxisError, match=message): + stats.wilcoxon(x, y, axis=3, _no_deco=True) + + +# data for k-statistics tests from +# https://cran.r-project.org/web/packages/kStatistics/kStatistics.pdf +# see nKS "Examples" +x_kstat = [16.34, 10.76, 11.84, 13.55, 15.85, 18.20, 7.51, 10.22, 12.52, 14.68, + 16.08, 19.43, 8.12, 11.20, 12.95, 14.77, 16.83, 19.80, 8.55, 11.58, + 12.10, 15.02, 16.83, 16.98, 19.92, 9.47, 11.68, 13.41, 15.35, 19.11] + + +@make_xp_test_case(stats.kstat) +class TestKstat: + def test_moments_normal_distribution(self, xp): + rng = np.random.RandomState(32149) + data = xp.asarray(rng.randn(12345), dtype=xp.float64) + moments = xp.stack([stats.kstat(data, n) for n in [1, 2, 3, 4]]) + + expected = xp.asarray([0.011315, 1.017931, 0.05811052, 0.0754134], + dtype=data.dtype) + xp_assert_close(moments, expected, rtol=1e-4) + + # test equivalence with `stats.moment` + m1 = stats.moment(data, order=1) + m2 = stats.moment(data, order=2) + m3 = stats.moment(data, order=3) + xp_assert_close(xp.stack((m1, m2, m3)), expected[:-1], atol=0.02, rtol=1e-2) + + @pytest.mark.filterwarnings("ignore:invalid value encountered") # Dask + def test_empty_input(self, xp): + with eager_warns(SmallSampleWarning, match=too_small_1d_not_omit, xp=xp): + res = stats.kstat(xp.asarray([])) + xp_assert_equal(res, xp.asarray(xp.nan)) + + def test_nan_input(self, xp): + data = xp.arange(10.) + data = xp.where(data == 6, xp.nan, data) + + xp_assert_equal(stats.kstat(data), xp.asarray(xp.nan)) + + @pytest.mark.parametrize('n', [0, 4.001]) + def test_kstat_bad_arg(self, n, xp): + # Raise ValueError if n > 4 or n < 1. + data = xp.arange(10) + message = 'k-statistics only supported for 1<=n<=4' + with pytest.raises(ValueError, match=message): + stats.kstat(data, n=n) + + @pytest.mark.parametrize('case', [(1, 14.02166666666667), + (2, 12.65006954022974), + (3, -1.447059503280798), + (4, -141.6682291883626)]) + def test_against_R(self, case, xp): + # Test against reference values computed with R kStatistics, e.g. + # options(digits=16) + # library(kStatistics) + # data <-c (16.34, 10.76, 11.84, 13.55, 15.85, 18.20, 7.51, 10.22, + # 12.52, 14.68, 16.08, 19.43, 8.12, 11.20, 12.95, 14.77, + # 16.83, 19.80, 8.55, 11.58, 12.10, 15.02, 16.83, 16.98, + # 19.92, 9.47, 11.68, 13.41, 15.35, 19.11) + # nKS(4, data) + n, ref = case + res = stats.kstat(xp.asarray(x_kstat), n) + xp_assert_close(res, xp.asarray(ref)) + + +@make_xp_test_case(stats.kstatvar) +class TestKstatVar: + @pytest.mark.filterwarnings("ignore:invalid value encountered") # Dask + def test_empty_input(self, xp): + x = xp.asarray([]) + with eager_warns(SmallSampleWarning, match=too_small_1d_not_omit, xp=xp): + res = stats.kstatvar(x) + xp_assert_equal(res, xp.asarray(xp.nan)) + + def test_nan_input(self, xp): + data = xp.arange(10.) + data = xp.where(data == 6, xp.nan, data) + + xp_assert_equal(stats.kstat(data), xp.asarray(xp.nan)) + + def test_bad_arg(self, xp): + # Raise ValueError is n is not 1 or 2. + data = xp.asarray([1]) + n = 10 + message = 'Only n=1 or n=2 supported.' + with pytest.raises(ValueError, match=message): + stats.kstatvar(data, n=n) + + def test_against_R_mathworld(self, xp): + # Test against reference values computed using formulas exactly as + # they appear at https://mathworld.wolfram.com/k-Statistic.html + # This is *really* similar to how they appear in the implementation, + # but that could change, and this should not. + n = len(x_kstat) + k2 = 12.65006954022974 # see source code in TestKstat + k4 = -141.6682291883626 + + res = stats.kstatvar(xp.asarray(x_kstat), 1) + ref = k2 / n + xp_assert_close(res, xp.asarray(ref)) + + res = stats.kstatvar(xp.asarray(x_kstat), 2) + # *unbiased estimator* for var(k2) + ref = (2*k2**2*n + (n-1)*k4) / (n * (n+1)) + xp_assert_close(res, xp.asarray(ref)) + + +class TestPpccPlot: + def setup_method(self): + self.x = _old_loggamma_rvs(5, size=500, random_state=7654321) + 5 + + def test_basic(self): + N = 5 + svals, ppcc = stats.ppcc_plot(self.x, -10, 10, N=N) + ppcc_expected = [0.21139644, 0.21384059, 0.98766719, 0.97980182, + 0.93519298] + assert_allclose(svals, np.linspace(-10, 10, num=N)) + assert_allclose(ppcc, ppcc_expected) + + def test_dist(self): + # Test that we can specify distributions both by name and as objects. + svals1, ppcc1 = stats.ppcc_plot(self.x, -10, 10, dist='tukeylambda') + svals2, ppcc2 = stats.ppcc_plot(self.x, -10, 10, + dist=stats.tukeylambda) + assert_allclose(svals1, svals2, rtol=1e-20) + assert_allclose(ppcc1, ppcc2, rtol=1e-20) + # Test that 'tukeylambda' is the default dist + svals3, ppcc3 = stats.ppcc_plot(self.x, -10, 10) + assert_allclose(svals1, svals3, rtol=1e-20) + assert_allclose(ppcc1, ppcc3, rtol=1e-20) + + @pytest.mark.skipif(not have_matplotlib, reason="no matplotlib") + def test_plot_kwarg(self): + # Check with the matplotlib.pyplot module + fig = plt.figure() + ax = fig.add_subplot(111) + stats.ppcc_plot(self.x, -20, 20, plot=plt) + fig.delaxes(ax) + + # Check that a Matplotlib Axes object is accepted + ax = fig.add_subplot(111) + stats.ppcc_plot(self.x, -20, 20, plot=ax) + plt.close() + + def test_invalid_inputs(self): + # `b` has to be larger than `a` + assert_raises(ValueError, stats.ppcc_plot, self.x, 1, 0) + + # Raise ValueError when given an invalid distribution. + assert_raises(ValueError, stats.ppcc_plot, [1, 2, 3], 0, 1, + dist="plate_of_shrimp") + + def test_empty(self): + # For consistency with probplot return for one empty array, + # ppcc contains all zeros and svals is the same as for normal array + # input. + svals, ppcc = stats.ppcc_plot([], 0, 1) + assert_allclose(svals, np.linspace(0, 1, num=80)) + assert_allclose(ppcc, np.zeros(80, dtype=float)) + + +class TestPpccMax: + def test_ppcc_max_bad_arg(self): + # Raise ValueError when given an invalid distribution. + data = [1] + assert_raises(ValueError, stats.ppcc_max, data, dist="plate_of_shrimp") + + def test_ppcc_max_basic(self): + x = stats.tukeylambda.rvs(-0.7, loc=2, scale=0.5, size=10000, + random_state=1234567) + 1e4 + assert_almost_equal(stats.ppcc_max(x), -0.71215366521264145, decimal=7) + + def test_dist(self): + x = stats.tukeylambda.rvs(-0.7, loc=2, scale=0.5, size=10000, + random_state=1234567) + 1e4 + + # Test that we can specify distributions both by name and as objects. + max1 = stats.ppcc_max(x, dist='tukeylambda') + max2 = stats.ppcc_max(x, dist=stats.tukeylambda) + assert_almost_equal(max1, -0.71215366521264145, decimal=5) + assert_almost_equal(max2, -0.71215366521264145, decimal=5) + + # Test that 'tukeylambda' is the default dist + max3 = stats.ppcc_max(x) + assert_almost_equal(max3, -0.71215366521264145, decimal=5) + + def test_brack(self): + x = stats.tukeylambda.rvs(-0.7, loc=2, scale=0.5, size=10000, + random_state=1234567) + 1e4 + assert_raises(ValueError, stats.ppcc_max, x, brack=(0.0, 1.0, 0.5)) + + assert_almost_equal(stats.ppcc_max(x, brack=(0, 1)), + -0.71215366521264145, decimal=7) + + assert_almost_equal(stats.ppcc_max(x, brack=(-2, 2)), + -0.71215366521264145, decimal=7) + + +@make_xp_test_case(stats.boxcox_llf) +class TestBoxcox_llf: + + @pytest.mark.parametrize("dtype", ["float32", "float64"]) + def test_basic(self, dtype, xp): + dt = getattr(xp, dtype) + x = stats.norm.rvs(size=10000, loc=10, random_state=54321) + lmbda = 1 + llf = stats.boxcox_llf(lmbda, xp.asarray(x, dtype=dt)) + llf_expected = -x.size / 2. * np.log(np.sum(x.std()**2)) + xp_assert_close(llf, xp.asarray(llf_expected, dtype=dt)) + + @skip_xp_backends(np_only=True, + reason='array-likes only accepted for NumPy backend.') + def test_array_like(self, xp): + x = stats.norm.rvs(size=100, loc=10, random_state=54321) + lmbda = 1 + llf = stats.boxcox_llf(lmbda, x) + llf2 = stats.boxcox_llf(lmbda, list(x)) + xp_assert_close(llf, llf2, rtol=1e-12) + + def test_2d_input(self, xp): + # Note: boxcox_llf() was already working with 2-D input (sort of), so + # keep it like that. boxcox() doesn't work with 2-D input though, due + # to brent() returning a scalar. + x = stats.norm.rvs(size=100, loc=10, random_state=54321) + lmbda = 1 + llf = stats.boxcox_llf(lmbda, x) + llf2 = stats.boxcox_llf(lmbda, np.vstack([x, x]).T) + xp_assert_close(xp.asarray([llf, llf]), xp.asarray(llf2), rtol=1e-12) + + def test_empty(self, xp): + message = "One or more sample arguments is too small..." + with eager_warns(SmallSampleWarning, match=message, xp=xp): + assert xp.isnan(xp.asarray(stats.boxcox_llf(1, xp.asarray([])))) + + def test_gh_6873(self, xp): + # Regression test for gh-6873. + # This example was taken from gh-7534, a duplicate of gh-6873. + data = xp.asarray([198.0, 233.0, 233.0, 392.0]) + llf = stats.boxcox_llf(-8, data) + # The expected value was computed with mpmath. + xp_assert_close(llf, xp.asarray(-17.93934208579061)) + + def test_instability_gh20021(self, xp): + data = xp.asarray([2003, 1950, 1997, 2000, 2009], dtype=xp.float64) + llf = stats.boxcox_llf(1e-8, data) + # The expected value was computed with mpsci, set mpmath.mp.dps=100 + # expect float64 output for integer input + xp_assert_close(llf, xp.asarray(-15.32401272869016598, dtype=xp.float64), + rtol=5e-7) # bumped tolerance from 1e-7 for Accelerate + + def test_axis(self, xp): + data = xp.asarray([[100, 200], [300, 400]]) + llf_axis_0 = stats.boxcox_llf(1, data, axis=0) + llf_0 = xp.stack([ + stats.boxcox_llf(1, data[:, 0]), + stats.boxcox_llf(1, data[:, 1]), + ]) + xp_assert_close(llf_axis_0, llf_0) + llf_axis_1 = stats.boxcox_llf(1, data, axis=1) + llf_1 = xp.stack([ + stats.boxcox_llf(1, data[0, :]), + stats.boxcox_llf(1, data[1, :]), + ]) + xp_assert_close(llf_axis_1, llf_1) + + +# This is the data from GitHub user Qukaiyi, given as an example +# of a data set that caused boxcox to fail. +_boxcox_data = [ + 15957, 112079, 1039553, 711775, 173111, 307382, 183155, 53366, 760875, + 207500, 160045, 473714, 40194, 440319, 133261, 265444, 155590, 36660, + 904939, 55108, 138391, 339146, 458053, 63324, 1377727, 1342632, 41575, + 68685, 172755, 63323, 368161, 199695, 538214, 167760, 388610, 398855, + 1001873, 364591, 1320518, 194060, 194324, 2318551, 196114, 64225, 272000, + 198668, 123585, 86420, 1925556, 695798, 88664, 46199, 759135, 28051, + 345094, 1977752, 51778, 82746, 638126, 2560910, 45830, 140576, 1603787, + 57371, 548730, 5343629, 2298913, 998813, 2156812, 423966, 68350, 145237, + 131935, 1600305, 342359, 111398, 1409144, 281007, 60314, 242004, 113418, + 246211, 61940, 95858, 957805, 40909, 307955, 174159, 124278, 241193, + 872614, 304180, 146719, 64361, 87478, 509360, 167169, 933479, 620561, + 483333, 97416, 143518, 286905, 597837, 2556043, 89065, 69944, 196858, + 88883, 49379, 916265, 1527392, 626954, 54415, 89013, 2883386, 106096, + 402697, 45578, 349852, 140379, 34648, 757343, 1305442, 2054757, 121232, + 606048, 101492, 51426, 1820833, 83412, 136349, 1379924, 505977, 1303486, + 95853, 146451, 285422, 2205423, 259020, 45864, 684547, 182014, 784334, + 174793, 563068, 170745, 1195531, 63337, 71833, 199978, 2330904, 227335, + 898280, 75294, 2011361, 116771, 157489, 807147, 1321443, 1148635, 2456524, + 81839, 1228251, 97488, 1051892, 75397, 3009923, 2732230, 90923, 39735, + 132433, 225033, 337555, 1204092, 686588, 1062402, 40362, 1361829, 1497217, + 150074, 551459, 2019128, 39581, 45349, 1117187, 87845, 1877288, 164448, + 10338362, 24942, 64737, 769946, 2469124, 2366997, 259124, 2667585, 29175, + 56250, 74450, 96697, 5920978, 838375, 225914, 119494, 206004, 430907, + 244083, 219495, 322239, 407426, 618748, 2087536, 2242124, 4736149, 124624, + 406305, 240921, 2675273, 4425340, 821457, 578467, 28040, 348943, 48795, + 145531, 52110, 1645730, 1768364, 348363, 85042, 2673847, 81935, 169075, + 367733, 135474, 383327, 1207018, 93481, 5934183, 352190, 636533, 145870, + 55659, 146215, 73191, 248681, 376907, 1606620, 169381, 81164, 246390, + 236093, 885778, 335969, 49266, 381430, 307437, 350077, 34346, 49340, + 84715, 527120, 40163, 46898, 4609439, 617038, 2239574, 159905, 118337, + 120357, 430778, 3799158, 3516745, 54198, 2970796, 729239, 97848, 6317375, + 887345, 58198, 88111, 867595, 210136, 1572103, 1420760, 574046, 845988, + 509743, 397927, 1119016, 189955, 3883644, 291051, 126467, 1239907, 2556229, + 411058, 657444, 2025234, 1211368, 93151, 577594, 4842264, 1531713, 305084, + 479251, 20591, 1466166, 137417, 897756, 594767, 3606337, 32844, 82426, + 1294831, 57174, 290167, 322066, 813146, 5671804, 4425684, 895607, 450598, + 1048958, 232844, 56871, 46113, 70366, 701618, 97739, 157113, 865047, + 194810, 1501615, 1765727, 38125, 2733376, 40642, 437590, 127337, 106310, + 4167579, 665303, 809250, 1210317, 45750, 1853687, 348954, 156786, 90793, + 1885504, 281501, 3902273, 359546, 797540, 623508, 3672775, 55330, 648221, + 266831, 90030, 7118372, 735521, 1009925, 283901, 806005, 2434897, 94321, + 309571, 4213597, 2213280, 120339, 64403, 8155209, 1686948, 4327743, + 1868312, 135670, 3189615, 1569446, 706058, 58056, 2438625, 520619, 105201, + 141961, 179990, 1351440, 3148662, 2804457, 2760144, 70775, 33807, 1926518, + 2362142, 186761, 240941, 97860, 1040429, 1431035, 78892, 484039, 57845, + 724126, 3166209, 175913, 159211, 1182095, 86734, 1921472, 513546, 326016, + 1891609 +] + + +class TransformTest: + def test_nan_policy(self): + transform = getattr(stats, self.transform) + normmax = getattr(stats, f"{self.transform}_normmax") + special_transform = getattr(special, self.transform, _yeojohnson_transform) + + rng = np.random.default_rng(5908294582456923546) + x = rng.random(10) + x[3] = np.nan + + # lmb provided + lmb = 1.0 + + with pytest.raises(ValueError, match="The input contains nan values"): + transform(x, lmb, nan_policy='raise') + + res = transform(x, lmb, nan_policy='propagate') + np.testing.assert_allclose(res, special_transform(x, 1.)) + + res = transform(x, lmb, nan_policy='omit') + np.testing.assert_allclose(res, transform(x, lmb, nan_policy='propagate')) + + # lmb not provided + with pytest.raises(ValueError, match="The input contains nan values"): + transform(x, nan_policy='raise') + + res = transform(x, nan_policy='propagate') + np.testing.assert_equal(res[0], special_transform(x, np.nan)) + np.testing.assert_equal(res[1], np.nan) + + kwarg = {'method': 'mle'} if transform == stats.boxcox else {} + res = transform(x, nan_policy='omit') + lmb = normmax(x, nan_policy='omit', **kwarg) + ref = transform(x, lmb) + np.testing.assert_allclose(res[0], ref) + np.testing.assert_allclose(res[1], lmb) + + +class TestBoxcox(TransformTest): + transform = 'boxcox' + + def test_fixed_lmbda(self): + x = _old_loggamma_rvs(5, size=50, random_state=12345) + 5 + xt = stats.boxcox(x, lmbda=1) + assert_allclose(xt, x - 1) + xt = stats.boxcox(x, lmbda=-1) + assert_allclose(xt, 1 - 1/x) + + xt = stats.boxcox(x, lmbda=0) + assert_allclose(xt, np.log(x)) + + # Also test that array_like input works + xt = stats.boxcox(list(x), lmbda=0) + assert_allclose(xt, np.log(x)) + + # test that constant input is accepted; see gh-12225 + xt = stats.boxcox(np.ones(10), 2) + assert_equal(xt, np.zeros(10)) + + def test_lmbda_None(self): + # Start from normal rv's, do inverse transform to check that + # optimization function gets close to the right answer. + lmbda = 2.5 + x = stats.norm.rvs(loc=10, size=50000, random_state=1245) + x_inv = (x * lmbda + 1)**(-lmbda) + xt, maxlog = stats.boxcox(x_inv) + + assert_almost_equal(maxlog, -1 / lmbda, decimal=2) + + def test_alpha(self): + rng = np.random.RandomState(1234) + x = _old_loggamma_rvs(5, size=50, random_state=rng) + 5 + + # Some regular values for alpha, on a small sample size + _, _, interval = stats.boxcox(x, alpha=0.75) + assert_allclose(interval, [4.004485780226041, 5.138756355035744]) + _, _, interval = stats.boxcox(x, alpha=0.05) + assert_allclose(interval, [1.2138178554857557, 8.209033272375663]) + + # Try some extreme values, see we don't hit the N=500 limit + x = _old_loggamma_rvs(7, size=500, random_state=rng) + 15 + _, _, interval = stats.boxcox(x, alpha=0.001) + assert_allclose(interval, [0.3988867, 11.40553131]) + _, _, interval = stats.boxcox(x, alpha=0.999) + assert_allclose(interval, [5.83316246, 5.83735292]) + + def test_boxcox_bad_arg(self): + # Raise ValueError if any data value is negative. + x = np.array([-1, 2]) + assert_raises(ValueError, stats.boxcox, x) + # Raise ValueError if data is constant. + assert_raises(ValueError, stats.boxcox, np.array([1])) + # Raise ValueError if data is not 1-dimensional. + assert_raises(ValueError, stats.boxcox, np.array([[1], [2]])) + + def test_empty(self): + assert_(stats.boxcox([]).shape == (0,)) + + def test_gh_6873(self): + # Regression test for gh-6873. + y, lam = stats.boxcox(_boxcox_data) + # The expected value of lam was computed with the function + # powerTransform in the R library 'car'. I trust that value + # to only about five significant digits. + assert_allclose(lam, -0.051654, rtol=1e-5) + + @pytest.mark.parametrize("bounds", [(-1, 1), (1.1, 2), (-2, -1.1)]) + def test_bounded_optimizer_within_bounds(self, bounds): + # Define custom optimizer with bounds. + def optimizer(fun): + return optimize.minimize_scalar(fun, bounds=bounds, + method="bounded") + + _, lmbda = stats.boxcox(_boxcox_data, lmbda=None, optimizer=optimizer) + assert bounds[0] < lmbda < bounds[1] + + def test_bounded_optimizer_against_unbounded_optimizer(self): + # Test whether setting bounds on optimizer excludes solution from + # unbounded optimizer. + + # Get unbounded solution. + _, lmbda = stats.boxcox(_boxcox_data, lmbda=None) + + # Set tolerance and bounds around solution. + bounds = (lmbda + 0.1, lmbda + 1) + options = {'xatol': 1e-12} + + def optimizer(fun): + return optimize.minimize_scalar(fun, bounds=bounds, + method="bounded", options=options) + + # Check bounded solution. Lower bound should be active. + _, lmbda_bounded = stats.boxcox(_boxcox_data, lmbda=None, + optimizer=optimizer) + assert lmbda_bounded != lmbda + assert_allclose(lmbda_bounded, bounds[0]) + + @pytest.mark.parametrize("optimizer", ["str", (1, 2), 0.1]) + def test_bad_optimizer_type_raises_error(self, optimizer): + # Check if error is raised if string, tuple or float is passed + with pytest.raises(ValueError, match="`optimizer` must be a callable"): + stats.boxcox(_boxcox_data, lmbda=None, optimizer=optimizer) + + def test_bad_optimizer_value_raises_error(self): + # Check if error is raised if `optimizer` function does not return + # `OptimizeResult` object + + # Define test function that always returns 1 + def optimizer(fun): + return 1 + + message = "return an object containing the optimal `lmbda`" + with pytest.raises(ValueError, match=message): + stats.boxcox(_boxcox_data, lmbda=None, optimizer=optimizer) + + @pytest.mark.parametrize( + "bad_x", [np.array([1, -42, 12345.6]), np.array([np.inf, 42, 1])] + ) + def test_negative_x_value_raises_error(self, bad_x): + """Test boxcox_normmax raises ValueError if x contains non-positive values.""" + message = "only positive, finite, real numbers" + with pytest.raises(ValueError, match=message): + stats.boxcox_normmax(bad_x) + + @pytest.mark.parametrize('x', [ + # Attempt to trigger overflow in power expressions. + np.array([2003.0, 1950.0, 1997.0, 2000.0, 2009.0, + 2009.0, 1980.0, 1999.0, 2007.0, 1991.0]), + # Attempt to trigger overflow with a large optimal lambda. + np.array([2003.0, 1950.0, 1997.0, 2000.0, 2009.0]), + # Attempt to trigger overflow with large data. + np.array([2003.0e200, 1950.0e200, 1997.0e200, 2000.0e200, 2009.0e200]) + ]) + def test_overflow(self, x): + with pytest.warns(UserWarning, match="The optimal lambda is"): + xt_bc, lam_bc = stats.boxcox(x) + assert np.all(np.isfinite(xt_bc)) + + def test_confidence_interval_nan_policy(self): + # other outputs with `nan_policy` tested in superclass + rng = np.random.default_rng(5908294582456923546) + x = rng.random(10) + x[3] = np.nan + alpha = 0.5 + + with pytest.raises(ValueError, match="The input contains nan values"): + stats.boxcox(x, alpha=alpha, nan_policy='raise') + + res = stats.boxcox(x, alpha=alpha, nan_policy='propagate')[2] + np.testing.assert_equal(res, (np.nan, np.nan)) + + res = stats.boxcox(x, alpha=alpha, nan_policy='omit')[2] + ref = stats.boxcox(x[~np.isnan(x)], alpha=alpha)[2] + np.testing.assert_allclose(res, ref) + + +class NormmaxTest: + def test_nan_policy(self): + transform_normmax = getattr(stats, self.transform_normmax) + rng = np.random.default_rng(5908294582456923546) + x = rng.random(10) + x[3] = np.nan + + with pytest.raises(ValueError, match="The input contains nan values"): + transform_normmax(x, nan_policy='raise') + + res = transform_normmax(x, nan_policy='propagate') + np.testing.assert_equal(res, np.nan) + + res = transform_normmax(x, nan_policy='omit') + ref = transform_normmax(x[~np.isnan(x)]) + np.testing.assert_allclose(res, ref) + + +class TestBoxcoxNormmax(NormmaxTest): + transform_normmax = 'boxcox_normmax' + + def setup_method(self): + self.x = _old_loggamma_rvs(5, size=50, random_state=12345) + 5 + + def test_pearsonr(self): + maxlog = stats.boxcox_normmax(self.x) + assert_allclose(maxlog, 1.804465, rtol=1e-6) + + def test_mle(self): + maxlog = stats.boxcox_normmax(self.x, method='mle') + assert_allclose(maxlog, 1.758101, rtol=1e-6) + + # Check that boxcox() uses 'mle' + _, maxlog_boxcox = stats.boxcox(self.x) + assert_allclose(maxlog_boxcox, maxlog) + + def test_all(self): + maxlog_all = stats.boxcox_normmax(self.x, method='all') + assert_allclose(maxlog_all, [1.804465, 1.758101], rtol=1e-6) + + @pytest.mark.parametrize("method", ["mle", "pearsonr", "all"]) + @pytest.mark.parametrize("bounds", [(-1, 1), (1.1, 2), (-2, -1.1)]) + def test_bounded_optimizer_within_bounds(self, method, bounds): + + def optimizer(fun): + return optimize.minimize_scalar(fun, bounds=bounds, + method="bounded") + + maxlog = stats.boxcox_normmax(self.x, method=method, + optimizer=optimizer) + assert np.all(bounds[0] < maxlog) + assert np.all(maxlog < bounds[1]) + + @pytest.mark.slow + def test_user_defined_optimizer(self): + # tests an optimizer that is not based on scipy.optimize.minimize + lmbda = stats.boxcox_normmax(self.x) + lmbda_rounded = np.round(lmbda, 5) + lmbda_range = np.linspace(lmbda_rounded-0.01, lmbda_rounded+0.01, 1001) + + class MyResult: + pass + + def optimizer(fun): + # brute force minimum over the range + objs = [] + for lmbda in lmbda_range: + objs.append(fun(lmbda)) + res = MyResult() + res.x = lmbda_range[np.argmin(objs)] + return res + + lmbda2 = stats.boxcox_normmax(self.x, optimizer=optimizer) + assert lmbda2 != lmbda # not identical + assert_allclose(lmbda2, lmbda, 1e-5) # but as close as it should be + + def test_user_defined_optimizer_and_brack_raises_error(self): + optimizer = optimize.minimize_scalar + + # Using default `brack=None` with user-defined `optimizer` works as + # expected. + stats.boxcox_normmax(self.x, brack=None, optimizer=optimizer) + + # Using user-defined `brack` with user-defined `optimizer` is expected + # to throw an error. Instead, users should specify + # optimizer-specific parameters in the optimizer function itself. + with pytest.raises(ValueError, match="`brack` must be None if " + "`optimizer` is given"): + + stats.boxcox_normmax(self.x, brack=(-2.0, 2.0), + optimizer=optimizer) + + @pytest.mark.parametrize( + 'x', ([2003.0, 1950.0, 1997.0, 2000.0, 2009.0], + [0.50000471, 0.50004979, 0.50005902, 0.50009312, 0.50001632])) + def test_overflow(self, x): + message = "The optimal lambda is..." + with pytest.warns(UserWarning, match=message): + lmbda = stats.boxcox_normmax(x, method='mle') + assert np.isfinite(special.boxcox(x, lmbda)).all() + # 10000 is safety factor used in boxcox_normmax + ymax = np.finfo(np.float64).max / 10000 + x_treme = np.max(x) if lmbda > 0 else np.min(x) + y_extreme = special.boxcox(x_treme, lmbda) + assert_allclose(y_extreme, ymax * np.sign(lmbda)) + + def test_negative_ymax(self): + with pytest.raises(ValueError, match="`ymax` must be strictly positive"): + stats.boxcox_normmax(self.x, ymax=-1) + + @pytest.mark.parametrize("x", [ + # positive overflow in float64 + np.array([2003.0, 1950.0, 1997.0, 2000.0, 2009.0], + dtype=np.float64), + # negative overflow in float64 + np.array([0.50000471, 0.50004979, 0.50005902, 0.50009312, 0.50001632], + dtype=np.float64), + # positive overflow in float32 + np.array([200.3, 195.0, 199.7, 200.0, 200.9], + dtype=np.float32), + # negative overflow in float32 + np.array([2e-30, 1e-30, 1e-30, 1e-30, 1e-30, 1e-30], + dtype=np.float32), + ]) + @pytest.mark.parametrize("ymax", [1e10, 1e30, None]) + # TODO: add method "pearsonr" after fix overflow issue + @pytest.mark.parametrize("method", ["mle"]) + def test_user_defined_ymax_input_float64_32(self, x, ymax, method): + # Test the maximum of the transformed data close to ymax + with pytest.warns(UserWarning, match="The optimal lambda is"): + kwarg = {'ymax': ymax} if ymax is not None else {} + lmb = stats.boxcox_normmax(x, method=method, **kwarg) + x_treme = [np.min(x), np.max(x)] + ymax_res = max(abs(stats.boxcox(x_treme, lmb))) + if ymax is None: + # 10000 is safety factor used in boxcox_normmax + ymax = np.finfo(x.dtype).max / 10000 + assert_allclose(ymax, ymax_res, rtol=1e-5) + + @pytest.mark.parametrize("x", [ + # positive overflow in float32 but not float64 + [200.3, 195.0, 199.7, 200.0, 200.9], + # negative overflow in float32 but not float64 + [2e-30, 1e-30, 1e-30, 1e-30, 1e-30, 1e-30], + ]) + # TODO: add method "pearsonr" after fix overflow issue + @pytest.mark.parametrize("method", ["mle"]) + def test_user_defined_ymax_inf(self, x, method): + x_32 = np.asarray(x, dtype=np.float32) + x_64 = np.asarray(x, dtype=np.float64) + + # assert overflow with float32 but not float64 + with pytest.warns(UserWarning, match="The optimal lambda is"): + stats.boxcox_normmax(x_32, method=method) + stats.boxcox_normmax(x_64, method=method) + + # compute the true optimal lambda then compare them + lmb_32 = stats.boxcox_normmax(x_32, ymax=np.inf, method=method) + lmb_64 = stats.boxcox_normmax(x_64, ymax=np.inf, method=method) + assert_allclose(lmb_32, lmb_64, rtol=1e-2) + + def test_nan_policy_all(self): + # `nan_policy` with other `method`s tested in superclass + rng = np.random.default_rng(5908294582456923546) + x = rng.random(10) + x[3] = np.nan + + with pytest.raises(ValueError, match="The input contains nan values"): + stats.boxcox_normmax(x, method='all', nan_policy='raise') + + res = stats.boxcox_normmax(x, method='all', nan_policy='propagate') + np.testing.assert_equal(res, (np.nan, np.nan)) + + res = stats.boxcox_normmax(x, method='all', nan_policy='omit') + ref = stats.boxcox_normmax(x[~np.isnan(x)], method='all') + np.testing.assert_allclose(res, ref, rtol=2e-7) + + +class TestBoxcoxNormplot: + def setup_method(self): + self.x = _old_loggamma_rvs(5, size=500, random_state=7654321) + 5 + + def test_basic(self): + N = 5 + lmbdas, ppcc = stats.boxcox_normplot(self.x, -10, 10, N=N) + ppcc_expected = [0.57783375, 0.83610988, 0.97524311, 0.99756057, + 0.95843297] + assert_allclose(lmbdas, np.linspace(-10, 10, num=N)) + assert_allclose(ppcc, ppcc_expected) + + @pytest.mark.skipif(not have_matplotlib, reason="no matplotlib") + def test_plot_kwarg(self): + # Check with the matplotlib.pyplot module + fig = plt.figure() + ax = fig.add_subplot(111) + stats.boxcox_normplot(self.x, -20, 20, plot=plt) + fig.delaxes(ax) + + # Check that a Matplotlib Axes object is accepted + ax = fig.add_subplot(111) + stats.boxcox_normplot(self.x, -20, 20, plot=ax) + plt.close() + + def test_invalid_inputs(self): + # `lb` has to be larger than `la` + assert_raises(ValueError, stats.boxcox_normplot, self.x, 1, 0) + # `x` can not contain negative values + assert_raises(ValueError, stats.boxcox_normplot, [-1, 1], 0, 1) + + def test_empty(self): + assert_(stats.boxcox_normplot([], 0, 1).size == 0) + + +@make_xp_test_case(stats.yeojohnson_llf) +class TestYeojohnson_llf: + + def test_array_like(self): + # array_like not applicable with SCIPY_ARRAY_API=1 + x = stats.norm.rvs(size=100, loc=0, random_state=54321) + lmbda = 1 + llf = stats.yeojohnson_llf(lmbda, x) + llf2 = stats.yeojohnson_llf(lmbda, list(x)) + assert_allclose(llf, llf2, rtol=1e-12) + + def test_2d_input(self, xp): + x = stats.norm.rvs(size=100, loc=10, random_state=54321) + x = xp.asarray(x) + lmbda = 1 + ref = stats.yeojohnson_llf(lmbda, x) + res = stats.yeojohnson_llf(lmbda, xp.stack([x, x]).T) + xp_assert_close(res, xp.stack((ref, ref)), rtol=1e-12) + + def test_empty(self, xp): + message = "One or more sample arguments is too small..." + with eager_warns(SmallSampleWarning, match=message, xp=xp): + assert xp.isnan(stats.yeojohnson_llf(1, xp.asarray([]))) + + @skip_xp_backends('jax.numpy', reason="JIT can't special-case for all >0 / <0") + def test_gh24172(self, xp): + # Test all negative and all positive data + data = xp.asarray([10, 10, 10, 9.9], dtype=xp.float64) + # The expected value was computed with mpsci, set mpmath.mp.dps=100 + # 1. Test precision loss + xp_assert_close(stats.yeojohnson_llf(-14, data), + xp.asarray(12.418595783381280, dtype=xp.float64), rtol=1e-7) + xp_assert_close(stats.yeojohnson_llf(16, -data), + xp.asarray(12.418595783381280, dtype=xp.float64), rtol=1e-7) + # 2. Test overflow + xp_assert_close(stats.yeojohnson_llf(-20, data), + xp.asarray(12.360966379200528, dtype=xp.float64), rtol=1e-7) + xp_assert_close(stats.yeojohnson_llf(20, -data), + xp.asarray(12.380287243698629, dtype=xp.float64), rtol=1e-7) + + @pytest.mark.parametrize("lmb, ref", + [(-1., 6.757930444097559), + (0., 11.90755408530727), + (1., 14.179952211901629), + (2., 13.040657538043835)]) + def test_against_reference(self, lmb, ref, xp): + # Reference values generated with mpsci, e.g. + # import numpy as np + # from mpmath import mp + # from mpsci.stats import yeojohnson_llf + # mp.dps = 100 + # rng = np.random.default_rng(78435782834992002) + # x = rng.uniform(-1, 1, size=25) + # float(yeojohnson_llf(-1, x)) # 6.757930444097559 + rng = np.random.default_rng(78435782834992002) + x = rng.uniform(-1, 1, size=25) + res = stats.yeojohnson_llf(lmb, xp.asarray(x.tolist())) + xp_assert_close(res, xp.asarray(ref)) + + +class TestYeojohnson(TransformTest): + transform = 'yeojohnson' + + def test_fixed_lmbda(self): + rng = np.random.RandomState(12345) + + # Test positive input + x = _old_loggamma_rvs(5, size=50, random_state=rng) + 5 + assert np.all(x > 0) + xt = stats.yeojohnson(x, lmbda=1) + assert_allclose(xt, x) + xt = stats.yeojohnson(x, lmbda=-1) + assert_allclose(xt, 1 - 1 / (x + 1)) + xt = stats.yeojohnson(x, lmbda=0) + assert_allclose(xt, np.log(x + 1)) + xt = stats.yeojohnson(x, lmbda=1) + assert_allclose(xt, x) + + # Test negative input + x = _old_loggamma_rvs(5, size=50, random_state=rng) - 5 + assert np.all(x < 0) + xt = stats.yeojohnson(x, lmbda=2) + assert_allclose(xt, -np.log(-x + 1)) + xt = stats.yeojohnson(x, lmbda=1) + assert_allclose(xt, x) + xt = stats.yeojohnson(x, lmbda=3) + assert_allclose(xt, 1 / (-x + 1) - 1) + + # test both positive and negative input + x = _old_loggamma_rvs(5, size=50, random_state=rng) - 2 + assert not np.all(x < 0) + assert not np.all(x >= 0) + pos = x >= 0 + xt = stats.yeojohnson(x, lmbda=1) + assert_allclose(xt[pos], x[pos]) + xt = stats.yeojohnson(x, lmbda=-1) + assert_allclose(xt[pos], 1 - 1 / (x[pos] + 1)) + xt = stats.yeojohnson(x, lmbda=0) + assert_allclose(xt[pos], np.log(x[pos] + 1)) + xt = stats.yeojohnson(x, lmbda=1) + assert_allclose(xt[pos], x[pos]) + + neg = ~pos + xt = stats.yeojohnson(x, lmbda=2) + assert_allclose(xt[neg], -np.log(-x[neg] + 1)) + xt = stats.yeojohnson(x, lmbda=1) + assert_allclose(xt[neg], x[neg]) + xt = stats.yeojohnson(x, lmbda=3) + assert_allclose(xt[neg], 1 / (-x[neg] + 1) - 1) + + @pytest.mark.parametrize('lmbda', [0, .1, .5, 2]) + def test_lmbda_None(self, lmbda): + # Start from normal rv's, do inverse transform to check that + # optimization function gets close to the right answer. + + def _inverse_transform(x, lmbda): + x_inv = np.zeros(x.shape, dtype=x.dtype) + pos = x >= 0 + + # when x >= 0 + if abs(lmbda) < np.spacing(1.): + x_inv[pos] = np.exp(x[pos]) - 1 + else: # lmbda != 0 + x_inv[pos] = np.power(x[pos] * lmbda + 1, 1 / lmbda) - 1 + + # when x < 0 + if abs(lmbda - 2) > np.spacing(1.): + x_inv[~pos] = 1 - np.power(-(2 - lmbda) * x[~pos] + 1, + 1 / (2 - lmbda)) + else: # lmbda == 2 + x_inv[~pos] = 1 - np.exp(-x[~pos]) + + return x_inv + + n_samples = 20000 + rng = np.random.RandomState(1234567) + x = rng.normal(loc=0, scale=1, size=(n_samples)) + + x_inv = _inverse_transform(x, lmbda) + xt, maxlog = stats.yeojohnson(x_inv) + + assert_allclose(maxlog, lmbda, atol=1e-2) + + assert_almost_equal(0, np.linalg.norm(x - xt) / n_samples, decimal=2) + assert_almost_equal(0, xt.mean(), decimal=1) + assert_almost_equal(1, xt.std(), decimal=1) + + def test_empty(self): + assert_(stats.yeojohnson([]).shape == (0,)) + + def test_array_like(self): + x = stats.norm.rvs(size=100, loc=0, random_state=54321) + xt1, _ = stats.yeojohnson(x) + xt2, _ = stats.yeojohnson(list(x)) + assert_allclose(xt1, xt2, rtol=1e-12) + + @pytest.mark.parametrize('dtype', [np.complex64, np.complex128]) + def test_input_dtype_complex(self, dtype): + x = np.arange(6, dtype=dtype) + err_msg = ('Yeo-Johnson transformation is not defined for complex ' + 'numbers.') + with pytest.raises(ValueError, match=err_msg): + stats.yeojohnson(x) + + @pytest.mark.parametrize('dtype', [np.int8, np.uint8, np.int16, np.int32]) + def test_input_dtype_integer(self, dtype): + x_int = np.arange(8, dtype=dtype) + x_float = np.arange(8, dtype=np.float64) + xt_int, lmbda_int = stats.yeojohnson(x_int) + xt_float, lmbda_float = stats.yeojohnson(x_float) + assert_allclose(xt_int, xt_float, rtol=1e-7) + assert_allclose(lmbda_int, lmbda_float, rtol=1e-7) + + def test_input_high_variance(self): + # non-regression test for gh-10821 + x = np.array([3251637.22, 620695.44, 11642969.00, 2223468.22, + 85307500.00, 16494389.89, 917215.88, 11642969.00, + 2145773.87, 4962000.00, 620695.44, 651234.50, + 1907876.71, 4053297.88, 3251637.22, 3259103.08, + 9547969.00, 20631286.23, 12807072.08, 2383819.84, + 90114500.00, 17209575.46, 12852969.00, 2414609.99, + 2170368.23]) + xt_yeo, lam_yeo = stats.yeojohnson(x) + xt_box, lam_box = stats.boxcox(x + 1) + assert_allclose(xt_yeo, xt_box, rtol=1e-6) + assert_allclose(lam_yeo, lam_box, rtol=1e-6) + + @pytest.mark.parametrize('x', [ + np.array([1.0, float("inf"), 2.0]), + np.array([1.0, -float("inf"), 2.0]), + np.array([-1.0, float("inf"), -float("inf"), 1.0]) + ]) + def test_nonfinite_input(self, x): + with pytest.raises(ValueError, match='Yeo-Johnson input must be finite'): + xt_yeo, lam_yeo = stats.yeojohnson(x) + + @pytest.mark.parametrize('x', [ + # Attempt to trigger overflow in power expressions. + np.array([2003.0, 1950.0, 1997.0, 2000.0, 2009.0, + 2009.0, 1980.0, 1999.0, 2007.0, 1991.0]), + # Attempt to trigger overflow with a large optimal lambda. + np.array([2003.0, 1950.0, 1997.0, 2000.0, 2009.0]), + # Attempt to trigger overflow with large data. + np.array([2003.0e200, 1950.0e200, 1997.0e200, 2000.0e200, 2009.0e200]) + ]) + def test_overflow(self, x): + # non-regression test for gh-18389 + + def optimizer(fun, lam_yeo): + out = optimize.fminbound(fun, -lam_yeo, lam_yeo, xtol=1.48e-08) + result = optimize.OptimizeResult() + result.x = out + return result + + with np.errstate(all="raise"): + xt_yeo, lam_yeo = stats.yeojohnson(x) + xt_box, lam_box = stats.boxcox( + x + 1, optimizer=partial(optimizer, lam_yeo=lam_yeo)) + assert np.isfinite(np.var(xt_yeo)) + assert np.isfinite(np.var(xt_box)) + assert_allclose(lam_yeo, lam_box, rtol=1e-6) + assert_allclose(xt_yeo, xt_box, rtol=1e-4) + + @pytest.mark.parametrize('x', [ + np.array([2003.0, 1950.0, 1997.0, 2000.0, 2009.0, + 2009.0, 1980.0, 1999.0, 2007.0, 1991.0]), + np.array([2003.0, 1950.0, 1997.0, 2000.0, 2009.0]) + ]) + @pytest.mark.parametrize('scale', [1, 1e-12, 1e-32, 1e-150, 1e32, 1e200]) + @pytest.mark.parametrize('sign', [1, -1]) + def test_overflow_underflow_signed_data(self, x, scale, sign): + # non-regression test for gh-18389 + with np.errstate(all="raise"): + xt_yeo, lam_yeo = stats.yeojohnson(sign * x * scale) + assert np.all(np.sign(sign * x) == np.sign(xt_yeo)) + assert np.isfinite(lam_yeo) + assert np.isfinite(np.var(xt_yeo)) + + @pytest.mark.parametrize('x', [ + np.array([0, 1, 2, 3]), + np.array([0, -1, 2, -3]), + np.array([0, 0, 0]) + ]) + @pytest.mark.parametrize('sign', [1, -1]) + @pytest.mark.parametrize('brack', [None, (-2, 2)]) + def test_integer_signed_data(self, x, sign, brack): + with np.errstate(all="raise"): + x_int = sign * x + x_float = x_int.astype(np.float64) + lam_yeo_int = stats.yeojohnson_normmax(x_int, brack=brack) + xt_yeo_int = stats.yeojohnson(x_int, lmbda=lam_yeo_int) + lam_yeo_float = stats.yeojohnson_normmax(x_float, brack=brack) + xt_yeo_float = stats.yeojohnson(x_float, lmbda=lam_yeo_float) + assert np.all(np.sign(x_int) == np.sign(xt_yeo_int)) + assert np.isfinite(lam_yeo_int) + assert np.isfinite(np.var(xt_yeo_int)) + assert lam_yeo_int == lam_yeo_float + assert np.all(xt_yeo_int == xt_yeo_float) + + +class TestYeojohnsonNormmax(NormmaxTest): + transform_normmax = 'yeojohnson_normmax' + + def setup_method(self): + self.x = _old_loggamma_rvs(5, size=50, random_state=12345) + 5 + + def test_mle(self): + maxlog = stats.yeojohnson_normmax(self.x) + assert_allclose(maxlog, 1.876393, rtol=1e-6) + + def test_darwin_example(self): + # test from original paper "A new family of power transformations to + # improve normality or symmetry" by Yeo and Johnson. + x = [6.1, -8.4, 1.0, 2.0, 0.7, 2.9, 3.5, 5.1, 1.8, 3.6, 7.0, 3.0, 9.3, + 7.5, -6.0] + lmbda = stats.yeojohnson_normmax(x) + assert np.allclose(lmbda, 1.305, atol=1e-3) + + +@make_xp_test_case(stats.circmean, stats.circvar, stats.circstd) +class TestCircFuncs: + # In gh-5747, the R package `circular` was used to calculate reference + # values for the circular variance, e.g.: + # library(circular) + # options(digits=16) + # x = c(0, 2*pi/3, 5*pi/3) + # var.circular(x) + @pytest.mark.parametrize("test_func,expected", + [(stats.circmean, 0.167690146), + (stats.circvar, 0.006455174000787767), + (stats.circstd, 6.520702116)]) + def test_circfuncs(self, test_func, expected, xp): + x = xp.asarray([355., 5., 2., 359., 10., 350.]) + xp_assert_close(test_func(x, high=360), xp.asarray(expected)) + + def test_circfuncs_small(self, xp): + # Default tolerances won't work here because the reference values + # are approximations. Ensure all array types work in float64 to + # avoid needing separate float32 and float64 tolerances. + x = xp.asarray([20, 21, 22, 18, 19, 20.5, 19.2], dtype=xp.float64) + M1 = xp.mean(x) + M2 = stats.circmean(x, high=360) + xp_assert_close(M2, M1, rtol=1e-5) + + V1 = xp.var(x*xp.pi/180, correction=0) + # for small variations, circvar is approximately half the + # linear variance + V1 = V1 / 2. + V2 = stats.circvar(x, high=360) + xp_assert_close(V2, V1, rtol=1e-4) + + S1 = xp.std(x, correction=0) + S2 = stats.circstd(x, high=360) + xp_assert_close(S2, S1, rtol=1e-4) + + @pytest.mark.parametrize("test_func, numpy_func", + [(stats.circmean, np.mean), + (stats.circvar, np.var), + (stats.circstd, np.std)]) + def test_circfuncs_close(self, test_func, numpy_func, xp): + # circfuncs should handle very similar inputs (gh-12740) + x = np.asarray([0.12675364631578953] * 10 + [0.12675365920187928] * 100) + circstat = test_func(xp.asarray(x)) + normal = xp.asarray(numpy_func(x)) + xp_assert_close(circstat, normal, atol=2e-8) + + @pytest.mark.parametrize('circfunc', [stats.circmean, + stats.circvar, + stats.circstd]) + def test_circmean_axis(self, circfunc, xp): + x = xp.asarray([[355, 5, 2, 359, 10, 350], + [351, 7, 4, 352, 9, 349], + [357, 9, 8, 358, 4, 356.]]) + res = circfunc(x, high=360) + ref = circfunc(xp.reshape(x, (-1,)), high=360) + xp_assert_close(res, xp.asarray(ref)) + + res = circfunc(x, high=360, axis=1) + ref = [circfunc(x[i, :], high=360) for i in range(x.shape[0])] + xp_assert_close(res, xp.stack(ref)) + + res = circfunc(x, high=360, axis=0) + ref = [circfunc(x[:, i], high=360) for i in range(x.shape[1])] + xp_assert_close(res, xp.stack(ref)) + + @pytest.mark.parametrize("test_func,expected", + [(stats.circmean, 0.167690146), + (stats.circvar, 0.006455174270186603), + (stats.circstd, 6.520702116)]) + def test_circfuncs_array_like(self, test_func, expected, xp): + x = xp.asarray([355, 5, 2, 359, 10, 350.]) + xp_assert_close(test_func(x, high=360), xp.asarray(expected)) + + @pytest.mark.parametrize("test_func", [stats.circmean, stats.circvar, + stats.circstd]) + def test_empty(self, test_func, xp): + dtype = xp.float64 + x = xp.asarray([], dtype=dtype) + with eager_warns(SmallSampleWarning, match=too_small_1d_not_omit, xp=xp): + res = test_func(x) + + xp_assert_equal(res, xp.asarray(xp.nan, dtype=dtype)) + + @pytest.mark.parametrize("test_func", [stats.circmean, stats.circvar, + stats.circstd]) + def test_nan_propagate(self, test_func, xp): + x = xp.asarray([355, 5, 2, 359, 10, 350, np.nan]) + xp_assert_equal(test_func(x, high=360), xp.asarray(xp.nan)) + + @pytest.mark.parametrize("test_func,expected", + [(stats.circmean, + {None: np.nan, 0: 355.66582264, 1: 0.28725053}), + (stats.circvar, + {None: np.nan, + 0: 0.002570671054089924, + 1: 0.005545914017677123}), + (stats.circstd, + {None: np.nan, 0: 4.11093193, 1: 6.04265394})]) + def test_nan_propagate_array(self, test_func, expected, xp): + x = xp.asarray([[355, 5, 2, 359, 10, 350, 1], + [351, 7, 4, 352, 9, 349, np.nan], + [1, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan]]) + for axis in expected.keys(): + out = test_func(x, high=360, axis=axis) + if axis is None: + xp_assert_equal(out, xp.asarray(xp.nan)) + else: + xp_assert_close(out[0], xp.asarray(expected[axis])) + xp_assert_equal(out[1:], xp.full_like(out[1:], xp.nan)) + + def test_circmean_scalar(self, xp): + x = xp.asarray(1.)[()] + M1 = x + M2 = stats.circmean(x) + xp_assert_close(M2, M1, rtol=1e-5) + + def test_circmean_range(self, xp): + # regression test for gh-6420: circmean(..., high, low) must be + # between `high` and `low` + m = stats.circmean(xp.arange(0, 2, 0.1), xp.pi, -xp.pi) + xp_assert_less(m, xp.asarray(xp.pi)) + xp_assert_less(-m, xp.asarray(xp.pi)) + + def test_circfuncs_uint8(self, xp): + # regression test for gh-7255: overflow when working with + # numpy uint8 data type + x = xp.asarray([150, 10], dtype=xp.uint8) + xp_assert_close(stats.circmean(x, high=180), xp.asarray(170.0)) + xp_assert_close(stats.circvar(x, high=180), xp.asarray(0.2339555554617)) + xp_assert_close(stats.circstd(x, high=180), xp.asarray(20.91551378)) + + def test_circstd_zero(self, xp): + # circstd() of a single number should return positive zero. + y = stats.circstd(xp.asarray([0])) + assert math.copysign(1.0, y) == 1.0 + + def test_circmean_accuracy_tiny_input(self, xp): + # For tiny x such that sin(x) == x and cos(x) == 1.0 numerically, + # circmean(x) should return x because atan2(sin(x), cos(x)) == x. + # This test verifies this. + # + # The purpose of this test is not to show that circmean() is + # accurate in the last digit for certain input, because this is + # neither guaranteed not particularly useful. Rather, it is a + # "white-box" sanity check that no undue loss of precision is + # introduced by conversion between (high - low) and (2 * pi). + + x = xp.linspace(1e-9, 6e-9, 50) + assert xp.all(xp.sin(x) == x) and xp.all(xp.cos(x) == 1.0) + + m = (x * (2 * xp.pi) / (2 * xp.pi)) != x + assert xp.any(m) + x = x[m] + + y = stats.circmean(x[:, None], axis=1) + assert xp.all(y == x) + + def test_circmean_accuracy_huge_input(self, xp): + # White-box test that circmean() does not introduce undue loss of + # numerical accuracy by eagerly rotating the input. This is detected + # by supplying a huge input x such that (x - low) == x numerically. + x = xp.asarray(1e17, dtype=xp.float64) + y = math.atan2(xp.sin(x), xp.cos(x)) # -2.6584887370946806 + expected = xp.asarray(y, dtype=xp.float64) + actual = stats.circmean(x, high=xp.pi, low=-xp.pi) + xp_assert_close(actual, expected, rtol=1e-15, atol=0.0) + + +class TestCircFuncsNanPolicy: + # `nan_policy` is implemented by the `_axis_nan_policy` decorator, which is + # not yet array-API compatible. When it is array-API compatible, the generic + # tests run on every function will be much stronger than these, so these + # will not be necessary. So I don't see a need to make these array-API compatible; + # when the time comes, they can just be removed. + @pytest.mark.parametrize("test_func,expected", + [(stats.circmean, + {None: 359.4178026893944, + 0: np.array([353.0, 6.0, 3.0, 355.5, 9.5, + 349.5]), + 1: np.array([0.16769015, 358.66510252])}), + (stats.circvar, + {None: 0.008396678483192477, + 0: np.array([1.9997969, 0.4999873, 0.4999873, + 6.1230956, 0.1249992, 0.1249992] + )*(np.pi/180)**2, + 1: np.array([0.006455174270186603, + 0.01016767581393285])}), + (stats.circstd, + {None: 7.440570778057074, + 0: np.array([2.00020313, 1.00002539, 1.00002539, + 3.50108929, 0.50000317, + 0.50000317]), + 1: np.array([6.52070212, 8.19138093])})]) + def test_nan_omit_array(self, test_func, expected): + x = np.array([[355, 5, 2, 359, 10, 350, np.nan], + [351, 7, 4, 352, 9, 349, np.nan], + [np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan]]) + for axis in expected.keys(): + if axis is None: + out = test_func(x, high=360, nan_policy='omit', axis=axis) + assert_allclose(out, expected[axis], rtol=1e-7) + else: + with pytest.warns(SmallSampleWarning, match=too_small_nd_omit): + out = test_func(x, high=360, nan_policy='omit', axis=axis) + assert_allclose(out[:-1], expected[axis], rtol=1e-7) + assert_(np.isnan(out[-1])) + + @pytest.mark.parametrize("test_func,expected", + [(stats.circmean, 0.167690146), + (stats.circvar, 0.006455174270186603), + (stats.circstd, 6.520702116)]) + def test_nan_omit(self, test_func, expected): + x = [355, 5, 2, 359, 10, 350, np.nan] + assert_allclose(test_func(x, high=360, nan_policy='omit'), + expected, rtol=1e-7) + + @pytest.mark.parametrize("test_func", [stats.circmean, stats.circvar, + stats.circstd]) + def test_nan_omit_all(self, test_func): + x = [np.nan, np.nan, np.nan, np.nan, np.nan] + with pytest.warns(SmallSampleWarning, match=too_small_1d_omit): + assert_(np.isnan(test_func(x, nan_policy='omit'))) + + @pytest.mark.parametrize("test_func", [stats.circmean, stats.circvar, + stats.circstd]) + def test_nan_omit_all_axis(self, test_func): + with pytest.warns(SmallSampleWarning, match=too_small_nd_omit): + x = np.array([[np.nan, np.nan, np.nan, np.nan, np.nan], + [np.nan, np.nan, np.nan, np.nan, np.nan]]) + out = test_func(x, nan_policy='omit', axis=1) + assert_(np.isnan(out).all()) + assert_(len(out) == 2) + + @pytest.mark.parametrize("x", + [[355, 5, 2, 359, 10, 350, np.nan], + np.array([[355, 5, 2, 359, 10, 350, np.nan], + [351, 7, 4, 352, np.nan, 9, 349]])]) + @pytest.mark.parametrize("test_func", [stats.circmean, stats.circvar, + stats.circstd]) + def test_nan_raise(self, test_func, x): + assert_raises(ValueError, test_func, x, high=360, nan_policy='raise') + + @pytest.mark.parametrize("x", + [[355, 5, 2, 359, 10, 350, np.nan], + np.array([[355, 5, 2, 359, 10, 350, np.nan], + [351, 7, 4, 352, np.nan, 9, 349]])]) + @pytest.mark.parametrize("test_func", [stats.circmean, stats.circvar, + stats.circstd]) + def test_bad_nan_policy(self, test_func, x): + assert_raises(ValueError, test_func, x, high=360, nan_policy='foobar') + + +class TestMedianTest: + + def test_bad_n_samples(self): + # median_test requires at least two samples. + assert_raises(ValueError, stats.median_test, [1, 2, 3]) + + def test_empty_sample(self): + # Each sample must contain at least one value. + assert_raises(ValueError, stats.median_test, [], [1, 2, 3]) + + def test_empty_when_ties_ignored(self): + # The grand median is 1, and all values in the first argument are + # equal to the grand median. With ties="ignore", those values are + # ignored, which results in the first sample being (in effect) empty. + # This should raise a ValueError. + assert_raises(ValueError, stats.median_test, + [1, 1, 1, 1], [2, 0, 1], [2, 0], ties="ignore") + + def test_empty_contingency_row(self): + # The grand median is 1, and with the default ties="below", all the + # values in the samples are counted as being below the grand median. + # This would result a row of zeros in the contingency table, which is + # an error. + assert_raises(ValueError, stats.median_test, [1, 1, 1], [1, 1, 1]) + + # With ties="above", all the values are counted as above the + # grand median. + assert_raises(ValueError, stats.median_test, [1, 1, 1], [1, 1, 1], + ties="above") + + def test_bad_ties(self): + assert_raises(ValueError, stats.median_test, [1, 2, 3], [4, 5], + ties="foo") + + def test_bad_nan_policy(self): + assert_raises(ValueError, stats.median_test, [1, 2, 3], [4, 5], + nan_policy='foobar') + + def test_bad_keyword(self): + assert_raises(TypeError, stats.median_test, [1, 2, 3], [4, 5], + foo="foo") + + def test_simple(self): + x = [1, 2, 3] + y = [1, 2, 3] + stat, p, med, tbl = stats.median_test(x, y) + + # The median is floating point, but this equality test should be safe. + assert_equal(med, 2.0) + + assert_array_equal(tbl, [[1, 1], [2, 2]]) + + # The expected values of the contingency table equal the contingency + # table, so the statistic should be 0 and the p-value should be 1. + assert_equal(stat, 0) + assert_equal(p, 1) + + def test_ties_options(self): + # Test the contingency table calculation. + x = [1, 2, 3, 4] + y = [5, 6] + z = [7, 8, 9] + # grand median is 5. + + # Default 'ties' option is "below". + stat, p, m, tbl = stats.median_test(x, y, z) + assert_equal(m, 5) + assert_equal(tbl, [[0, 1, 3], [4, 1, 0]]) + + stat, p, m, tbl = stats.median_test(x, y, z, ties="ignore") + assert_equal(m, 5) + assert_equal(tbl, [[0, 1, 3], [4, 0, 0]]) + + stat, p, m, tbl = stats.median_test(x, y, z, ties="above") + assert_equal(m, 5) + assert_equal(tbl, [[0, 2, 3], [4, 0, 0]]) + + def test_nan_policy_options(self): + x = [1, 2, np.nan] + y = [4, 5, 6] + mt1 = stats.median_test(x, y, nan_policy='propagate') + s, p, m, t = stats.median_test(x, y, nan_policy='omit') + + assert_equal(mt1, (np.nan, np.nan, np.nan, None)) + assert_allclose(s, 0.31250000000000006) + assert_allclose(p, 0.57615012203057869) + assert_equal(m, 4.0) + assert_equal(t, np.array([[0, 2], [2, 1]])) + assert_raises(ValueError, stats.median_test, x, y, nan_policy='raise') + + def test_basic(self): + # median_test calls chi2_contingency to compute the test statistic + # and p-value. Make sure it hasn't screwed up the call... + + x = [1, 2, 3, 4, 5] + y = [2, 4, 6, 8] + + stat, p, m, tbl = stats.median_test(x, y) + assert_equal(m, 4) + assert_equal(tbl, [[1, 2], [4, 2]]) + + exp_stat, exp_p, dof, e = stats.chi2_contingency(tbl) + assert_allclose(stat, exp_stat) + assert_allclose(p, exp_p) + + stat, p, m, tbl = stats.median_test(x, y, lambda_=0) + assert_equal(m, 4) + assert_equal(tbl, [[1, 2], [4, 2]]) + + exp_stat, exp_p, dof, e = stats.chi2_contingency(tbl, lambda_=0) + assert_allclose(stat, exp_stat) + assert_allclose(p, exp_p) + + stat, p, m, tbl = stats.median_test(x, y, correction=False) + assert_equal(m, 4) + assert_equal(tbl, [[1, 2], [4, 2]]) + + exp_stat, exp_p, dof, e = stats.chi2_contingency(tbl, correction=False) + assert_allclose(stat, exp_stat) + assert_allclose(p, exp_p) + + @pytest.mark.parametrize("correction", [False, True]) + def test_result(self, correction): + x = [1, 2, 3] + y = [1, 2, 3] + + res = stats.median_test(x, y, correction=correction) + assert_equal((res.statistic, res.pvalue, res.median, res.table), res) + + +@make_xp_test_case(stats.directional_stats) +class TestDirectionalStats: + # Reference implementations are not available + def test_directional_stats_correctness(self, xp): + # Data from Fisher: Dispersion on a sphere, 1953 and + # Mardia and Jupp, Directional Statistics. + decl = -np.deg2rad(np.array([343.2, 62., 36.9, 27., 359., + 5.7, 50.4, 357.6, 44.])) + incl = -np.deg2rad(np.array([66.1, 68.7, 70.1, 82.1, 79.5, + 73., 69.3, 58.8, 51.4])) + data = np.stack((np.cos(incl) * np.cos(decl), + np.cos(incl) * np.sin(decl), + np.sin(incl)), + axis=1) + + decl = xp.asarray(decl.tolist()) + incl = xp.asarray(incl.tolist()) + data = xp.asarray(data.tolist()) + + dirstats = stats.directional_stats(data) + directional_mean = dirstats.mean_direction + + reference_mean = xp.asarray([0.2984, -0.1346, -0.9449]) + xp_assert_close(directional_mean, reference_mean, atol=1e-4) + + @pytest.mark.parametrize('angles, ref', [ + ([-np.pi/2, np.pi/2], 1.), + ([0, 2 * np.pi], 0.) + ]) + def test_directional_stats_2d_special_cases(self, angles, ref, xp): + angles = xp.asarray(angles) + ref = xp.asarray(ref) + data = xp.stack([xp.cos(angles), xp.sin(angles)], axis=1) + res = 1 - stats.directional_stats(data).mean_resultant_length + xp_assert_close(res, ref) + + def test_directional_stats_2d(self, xp): + # Test that for circular data directional_stats + # yields the same result as circmean/circvar + rng = np.random.default_rng(0xec9a6899d5a2830e0d1af479dbe1fd0c) + testdata = xp.asarray(2 * xp.pi * rng.random((1000, ))) + testdata_vector = xp.stack((xp.cos(testdata), + xp.sin(testdata)), + axis=1) + dirstats = stats.directional_stats(testdata_vector) + directional_mean = dirstats.mean_direction + directional_mean_angle = xp.atan2(directional_mean[1], directional_mean[0]) + directional_mean_angle = directional_mean_angle % (2 * xp.pi) + circmean = stats.circmean(testdata) + xp_assert_close(directional_mean_angle, circmean) + + directional_var = 1. - dirstats.mean_resultant_length + circular_var = stats.circvar(testdata) + xp_assert_close(directional_var, circular_var) + + def test_directional_mean_higher_dim(self, xp): + # test that directional_stats works for higher dimensions + # here a 4D array is reduced over axis = 2 + data = xp.asarray([[0.8660254, 0.5, 0.], + [0.8660254, -0.5, 0.]]) + full_array = xp.asarray(xp.tile(data, (2, 2, 2, 1))) + expected = xp.asarray([[[1., 0., 0.], + [1., 0., 0.]], + [[1., 0., 0.], + [1., 0., 0.]]]) + dirstats = stats.directional_stats(full_array, axis=2) + xp_assert_close(dirstats.mean_direction, expected) + + @skip_xp_backends(np_only=True, reason='checking array-like input') + def test_directional_stats_list_ndarray_input(self, xp): + # test that list and numpy array inputs yield same results + data = [[0.8660254, 0.5, 0.], [0.8660254, -0.5, 0]] + data_array = xp.asarray(data, dtype=xp.float64) + ref = stats.directional_stats(data) + res = stats.directional_stats(data_array) + xp_assert_close(res.mean_direction, + xp.asarray(ref.mean_direction)) + xp_assert_close(res.mean_resultant_length, + xp.asarray(res.mean_resultant_length)) + + def test_directional_stats_1d_error(self, xp): + # test that one-dimensional data raises ValueError + data = xp.ones((5, )) + message = (r"samples must at least be two-dimensional. " + r"Instead samples has shape: (5,)") + with pytest.raises(ValueError, match=re.escape(message)): + stats.directional_stats(data) + + @pytest.mark.parametrize("dtype", ["float32", "float64"]) + def test_directional_stats_normalize(self, dtype, xp): + # test that directional stats calculations yield same results + # for unnormalized input with normalize=True and normalized + # input with normalize=False + data = np.array([[0.8660254, 0.5, 0.], + [1.7320508, -1., 0.]], dtype=dtype) + res = stats.directional_stats(xp.asarray(data), normalize=True) + normalized_data = data / np.linalg.norm(data, axis=-1, + keepdims=True) + ref = stats.directional_stats(normalized_data, normalize=False) + xp_assert_close(res.mean_direction, + xp.asarray(ref.mean_direction)) + xp_assert_close(res.mean_resultant_length, + xp.asarray(ref.mean_resultant_length)) + + +@make_xp_test_case(stats.false_discovery_control) +class TestFDRControl: + @skip_xp_backends(eager_only=True) + def test_input_validation_ps(self, xp): + message = "All values in `ps` must lie between 0. and 1." + with pytest.raises(ValueError, match=message): + stats.false_discovery_control(xp.asarray([-1, 0.5, 0.7])) + with pytest.raises(ValueError, match=message): + stats.false_discovery_control(xp.asarray([0.5, 0.7, 2])) + with pytest.raises(ValueError, match=message): + stats.false_discovery_control(xp.asarray([0.5, 0.7, xp.nan])) + + def test_input_validation(self, xp): + message = "`ps` must contain only real numbers." + with pytest.raises(ValueError, match=message): + stats.false_discovery_control(xp.asarray([1+1j, 2+2j])) + + message = "Unrecognized `method` 'YAK'" + with pytest.raises(ValueError, match=message): + stats.false_discovery_control(xp.asarray([0.5, 0.7, 0.9]), method='YAK') + + message = "`axis` must be an integer or `None`" + with pytest.raises(ValueError, match=message): + stats.false_discovery_control(xp.asarray([0.5, 0.7, 0.9]), axis=1.5) + with pytest.raises(ValueError, match=message): + stats.false_discovery_control(xp.asarray([0.5, 0.7, 0.9]), axis=(1, 2)) + + def test_against_TileStats(self, xp): + # See reference [3] of false_discovery_control + ps = xp.asarray([0.005, 0.009, 0.019, 0.022, 0.051, 0.101, 0.361, 0.387]) + res = stats.false_discovery_control(ps) + ref = xp.asarray([0.036, 0.036, 0.044, 0.044, 0.082, 0.135, 0.387, 0.387]) + xp_assert_close(res, ref, atol=1e-3) + + @pytest.mark.parametrize("case", + [([0.24617028, 0.01140030, 0.05652047, 0.06841983, + 0.07989886, 0.01841490, 0.17540784, 0.06841983, + 0.06841983, 0.25464082], 'bh'), + ([0.72102493, 0.03339112, 0.16554665, 0.20039952, + 0.23402122, 0.05393666, 0.51376399, 0.20039952, + 0.20039952, 0.74583488], 'by')]) + def test_against_R(self, case, xp): + # Test against p.adjust, e.g. + # p = c(0.22155325, 0.00114003,..., 0.0364813 , 0.25464082) + # p.adjust(p, "BY") + ref, method = case + rng = np.random.default_rng(6134137338861652935) + ps = stats.loguniform.rvs(1e-3, 0.5, size=10, random_state=rng).tolist() + ps[3] = ps[7] # force a tie + res = stats.false_discovery_control(xp.asarray(ps), method=method) + xp_assert_close(res, xp.asarray(ref), atol=1e-6) + + def test_axis_None(self, xp): + rng = np.random.default_rng(6134137338861652935) + ps = stats.loguniform.rvs(1e-3, 0.5, size=(3, 4, 5), random_state=rng) + ps = xp.asarray(ps) + res = stats.false_discovery_control(ps, axis=None) + ref = stats.false_discovery_control(xp_ravel(ps)) + xp_assert_equal(res, ref) + + @pytest.mark.parametrize("axis", [0, 1, -1]) + def test_axis(self, axis, xp): + rng = np.random.default_rng(6134137338861652935) + ps = stats.loguniform.rvs(1e-3, 0.5, size=(3, 4, 5), random_state=rng) + res = stats.false_discovery_control(xp.asarray(ps), axis=axis) + ref = np.apply_along_axis(stats.false_discovery_control, axis, ps) + xp_assert_close(res, xp.asarray(ref)) # torch isn't *equal* + + def test_edge_cases(self, xp): + ps = xp.asarray([0.25]) + xp_assert_equal(stats.false_discovery_control(ps), ps) + + ps = xp.asarray([]) + xp_assert_equal(stats.false_discovery_control(ps), ps) + + if is_numpy(xp): + xp_assert_equal(stats.false_discovery_control(0.25), 0.25) + + +@pytest.mark.uses_xp_capabilities(False, reason="not used in this test yet") +class TestCommonAxis: + # More thorough testing of `axis` in `test_axis_nan_policy`, + # but those tests aren't run with array API yet. This class + # is in `test_morestats` instead of `test_axis_nan_policy` + # because there is no reason to run `test_axis_nan_policy` + # with the array API CI job right now. + + @pytest.mark.parametrize('case', [(stats.sem, {}), + (stats.kstat, {'n': 4}), + (stats.kstat, {'n': 2}), + (stats.variation, {})]) + def test_axis(self, case, xp): + if is_torch(xp) and case[0] == stats.variation: + pytest.xfail(reason="copysign doesn't accept scalar array-api-compat#271") + fun, kwargs = case + rng = np.random.default_rng(24598245982345) + x = xp.asarray(rng.random((6, 7))) + + res = fun(x, **kwargs, axis=0) + ref = xp.stack([fun(x[:, i], **kwargs) for i in range(x.shape[1])]) + xp_assert_close(res, ref) + + res = fun(x, **kwargs, axis=1) + ref = xp.stack([fun(x[i, :], **kwargs) for i in range(x.shape[0])]) + xp_assert_close(res, ref) + + res = fun(x, **kwargs, axis=None) + ref = fun(xp.reshape(x, (-1,)), **kwargs) + xp_assert_close(res, ref) diff --git a/v12x/pylib/scipy/stats/tests/test_mstats_basic.py b/v12x/pylib/scipy/stats/tests/test_mstats_basic.py new file mode 100644 index 0000000000000000000000000000000000000000..475ca50f7dd001d4d04326cc04a258ff739c3154 --- /dev/null +++ b/v12x/pylib/scipy/stats/tests/test_mstats_basic.py @@ -0,0 +1,2146 @@ +""" +Tests for the stats.mstats module (support for masked arrays) +""" +import warnings +import platform + +import numpy as np +from numpy import nan +import numpy.ma as ma +from numpy.ma import masked, nomask + +import scipy.stats.mstats as mstats +from scipy import stats +from .common_tests import check_named_results +import pytest +from pytest import raises as assert_raises +from numpy.ma.testutils import (assert_equal, assert_almost_equal, + assert_array_almost_equal, + assert_array_almost_equal_nulp, assert_, + assert_allclose, assert_array_equal) +from scipy.stats import _mstats_basic, _stats_py +from scipy.conftest import skip_xp_invalid_arg +from scipy.stats._axis_nan_policy import SmallSampleWarning, too_small_1d_not_omit + +class TestMquantiles: + def test_mquantiles_limit_keyword(self): + # Regression test for Trac ticket #867 + data = np.array([[6., 7., 1.], + [47., 15., 2.], + [49., 36., 3.], + [15., 39., 4.], + [42., 40., -999.], + [41., 41., -999.], + [7., -999., -999.], + [39., -999., -999.], + [43., -999., -999.], + [40., -999., -999.], + [36., -999., -999.]]) + desired = [[19.2, 14.6, 1.45], + [40.0, 37.5, 2.5], + [42.8, 40.05, 3.55]] + quants = mstats.mquantiles(data, axis=0, limit=(0, 50)) + assert_almost_equal(quants, desired) + + +def check_equal_gmean(array_like, desired, axis=None, dtype=None, rtol=1e-7): + # Note this doesn't test when axis is not specified + x = mstats.gmean(array_like, axis=axis, dtype=dtype) + assert_allclose(x, desired, rtol=rtol) + assert_equal(x.dtype, dtype) + + +def check_equal_hmean(array_like, desired, axis=None, dtype=None, rtol=1e-7): + x = stats.hmean(array_like, axis=axis, dtype=dtype) + assert_allclose(x, desired, rtol=rtol) + assert_equal(x.dtype, dtype) + + +@skip_xp_invalid_arg +class TestGeoMean: + def test_1d(self): + a = [1, 2, 3, 4] + desired = np.power(1*2*3*4, 1./4.) + check_equal_gmean(a, desired, rtol=1e-14) + + def test_1d_ma(self): + # Test a 1d masked array + a = ma.array([10, 20, 30, 40, 50, 60, 70, 80, 90, 100]) + desired = 45.2872868812 + check_equal_gmean(a, desired) + + a = ma.array([1, 2, 3, 4], mask=[0, 0, 0, 1]) + desired = np.power(1*2*3, 1./3.) + check_equal_gmean(a, desired, rtol=1e-14) + + def test_1d_ma_value(self): + # Test a 1d masked array with a masked value + a = np.ma.array([10, 20, 30, 40, 50, 60, 70, 80, 90, 100], + mask=[0, 0, 0, 0, 0, 0, 0, 0, 0, 1]) + desired = 41.4716627439 + check_equal_gmean(a, desired) + + def test_1d_ma0(self): + # Test a 1d masked array with zero element + a = np.ma.array([10, 20, 30, 40, 50, 60, 70, 80, 90, 0]) + desired = 0 + check_equal_gmean(a, desired) + + def test_1d_ma_inf(self): + # Test a 1d masked array with negative element + a = np.ma.array([10, 20, 30, 40, 50, 60, 70, 80, 90, -1]) + desired = np.nan + with np.errstate(invalid='ignore'): + check_equal_gmean(a, desired) + + @pytest.mark.skipif(not hasattr(np, 'float96'), + reason='cannot find float96 so skipping') + def test_1d_float96(self): + a = ma.array([1, 2, 3, 4], mask=[0, 0, 0, 1]) + desired_dt = np.power(1*2*3, 1./3.).astype(np.float96) + check_equal_gmean(a, desired_dt, dtype=np.float96, rtol=1e-14) + + def test_2d_ma(self): + a = ma.array([[1, 2, 3, 4], [1, 2, 3, 4], [1, 2, 3, 4]], + mask=[[0, 0, 0, 0], [1, 0, 0, 1], [0, 1, 1, 0]]) + desired = np.array([1, 2, 3, 4]) + check_equal_gmean(a, desired, axis=0, rtol=1e-14) + + desired = ma.array([np.power(1*2*3*4, 1./4.), + np.power(2*3, 1./2.), + np.power(1*4, 1./2.)]) + check_equal_gmean(a, desired, axis=-1, rtol=1e-14) + + # Test a 2d masked array + a = [[10, 20, 30, 40], [50, 60, 70, 80], [90, 100, 110, 120]] + desired = 52.8885199 + check_equal_gmean(np.ma.array(a), desired) + + +@skip_xp_invalid_arg +class TestHarMean: + def test_1d(self): + a = ma.array([1, 2, 3, 4], mask=[0, 0, 0, 1]) + desired = 3. / (1./1 + 1./2 + 1./3) + check_equal_hmean(a, desired, rtol=1e-14) + + a = np.ma.array([10, 20, 30, 40, 50, 60, 70, 80, 90, 100]) + desired = 34.1417152147 + check_equal_hmean(a, desired) + + a = np.ma.array([10, 20, 30, 40, 50, 60, 70, 80, 90, 100], + mask=[0, 0, 0, 0, 0, 0, 0, 0, 0, 1]) + desired = 31.8137186141 + check_equal_hmean(a, desired) + + @pytest.mark.skipif(not hasattr(np, 'float96'), + reason='cannot find float96 so skipping') + def test_1d_float96(self): + a = ma.array([1, 2, 3, 4], mask=[0, 0, 0, 1]) + desired_dt = np.asarray(3. / (1./1 + 1./2 + 1./3), dtype=np.float96) + check_equal_hmean(a, desired_dt, dtype=np.float96) + + def test_2d(self): + a = ma.array([[1, 2, 3, 4], [1, 2, 3, 4], [1, 2, 3, 4]], + mask=[[0, 0, 0, 0], [1, 0, 0, 1], [0, 1, 1, 0]]) + desired = ma.array([1, 2, 3, 4]) + check_equal_hmean(a, desired, axis=0, rtol=1e-14) + + desired = [4./(1/1.+1/2.+1/3.+1/4.), 2./(1/2.+1/3.), 2./(1/1.+1/4.)] + check_equal_hmean(a, desired, axis=-1, rtol=1e-14) + + a = [[10, 20, 30, 40], [50, 60, 70, 80], [90, 100, 110, 120]] + desired = 38.6696271841 + check_equal_hmean(np.ma.array(a), desired) + + +class TestRanking: + def test_ranking(self): + x = ma.array([0,1,1,1,2,3,4,5,5,6,]) + assert_almost_equal(mstats.rankdata(x), + [1,3,3,3,5,6,7,8.5,8.5,10]) + x[[3,4]] = masked + assert_almost_equal(mstats.rankdata(x), + [1,2.5,2.5,0,0,4,5,6.5,6.5,8]) + assert_almost_equal(mstats.rankdata(x, use_missing=True), + [1,2.5,2.5,4.5,4.5,4,5,6.5,6.5,8]) + x = ma.array([0,1,5,1,2,4,3,5,1,6,]) + assert_almost_equal(mstats.rankdata(x), + [1,3,8.5,3,5,7,6,8.5,3,10]) + x = ma.array([[0,1,1,1,2], [3,4,5,5,6,]]) + assert_almost_equal(mstats.rankdata(x), + [[1,3,3,3,5], [6,7,8.5,8.5,10]]) + assert_almost_equal(mstats.rankdata(x, axis=1), + [[1,3,3,3,5], [1,2,3.5,3.5,5]]) + assert_almost_equal(mstats.rankdata(x,axis=0), + [[1,1,1,1,1], [2,2,2,2,2,]]) + + +class TestCorr: + def test_pearsonr(self): + # Tests some computations of Pearson's r + x = ma.arange(10) + with warnings.catch_warnings(): + # The tests in this context are edge cases, with perfect + # correlation or anticorrelation, or totally masked data. + # None of these should trigger a RuntimeWarning. + warnings.simplefilter("error", RuntimeWarning) + + assert_almost_equal(mstats.pearsonr(x, x)[0], 1.0) + assert_almost_equal(mstats.pearsonr(x, x[::-1])[0], -1.0) + + x = ma.array(x, mask=True) + pr = mstats.pearsonr(x, x) + assert_(pr[0] is masked) + assert_(pr[1] is masked) + + x1 = ma.array([-1.0, 0.0, 1.0]) + y1 = ma.array([0, 0, 3]) + r, p = mstats.pearsonr(x1, y1) + assert_almost_equal(r, np.sqrt(3)/2) + assert_almost_equal(p, 1.0/3) + + # (x2, y2) have the same unmasked data as (x1, y1). + mask = [False, False, False, True] + x2 = ma.array([-1.0, 0.0, 1.0, 99.0], mask=mask) + y2 = ma.array([0, 0, 3, -1], mask=mask) + r, p = mstats.pearsonr(x2, y2) + assert_almost_equal(r, np.sqrt(3)/2) + assert_almost_equal(p, 1.0/3) + + def test_pearsonr_misaligned_mask(self): + mx = np.ma.masked_array([1, 2, 3, 4, 5, 6], mask=[0, 1, 0, 0, 0, 0]) + my = np.ma.masked_array([9, 8, 7, 6, 5, 9], mask=[0, 0, 1, 0, 0, 0]) + x = np.array([1, 4, 5, 6]) + y = np.array([9, 6, 5, 9]) + mr, mp = mstats.pearsonr(mx, my) + r, p = stats.pearsonr(x, y) + assert_equal(mr, r) + assert_equal(mp, p) + + def test_spearmanr(self): + # Tests some computations of Spearman's rho + (x, y) = ([5.05,6.75,3.21,2.66], [1.65,2.64,2.64,6.95]) + assert_almost_equal(mstats.spearmanr(x,y)[0], -0.6324555) + (x, y) = ([5.05,6.75,3.21,2.66,np.nan],[1.65,2.64,2.64,6.95,np.nan]) + (x, y) = (ma.fix_invalid(x), ma.fix_invalid(y)) + assert_almost_equal(mstats.spearmanr(x,y)[0], -0.6324555) + + x = [2.0, 47.4, 42.0, 10.8, 60.1, 1.7, 64.0, 63.1, + 1.0, 1.4, 7.9, 0.3, 3.9, 0.3, 6.7] + y = [22.6, 8.3, 44.4, 11.9, 24.6, 0.6, 5.7, 41.6, + 0.0, 0.6, 6.7, 3.8, 1.0, 1.2, 1.4] + assert_almost_equal(mstats.spearmanr(x,y)[0], 0.6887299) + x = [2.0, 47.4, 42.0, 10.8, 60.1, 1.7, 64.0, 63.1, + 1.0, 1.4, 7.9, 0.3, 3.9, 0.3, 6.7, np.nan] + y = [22.6, 8.3, 44.4, 11.9, 24.6, 0.6, 5.7, 41.6, + 0.0, 0.6, 6.7, 3.8, 1.0, 1.2, 1.4, np.nan] + (x, y) = (ma.fix_invalid(x), ma.fix_invalid(y)) + assert_almost_equal(mstats.spearmanr(x,y)[0], 0.6887299) + # Next test is to make sure calculation uses sufficient precision. + # The denominator's value is ~n^3 and used to be represented as an + # int. 2000**3 > 2**32 so these arrays would cause overflow on + # some machines. + x = list(range(2000)) + y = list(range(2000)) + y[0], y[9] = y[9], y[0] + y[10], y[434] = y[434], y[10] + y[435], y[1509] = y[1509], y[435] + # rho = 1 - 6 * (2 * (9^2 + 424^2 + 1074^2))/(2000 * (2000^2 - 1)) + # = 1 - (1 / 500) + # = 0.998 + assert_almost_equal(mstats.spearmanr(x,y)[0], 0.998) + + # test for namedtuple attributes + res = mstats.spearmanr(x, y) + attributes = ('correlation', 'pvalue') + check_named_results(res, attributes, ma=True) + + def test_spearmanr_alternative(self): + # check against R + # options(digits=16) + # cor.test(c(2.0, 47.4, 42.0, 10.8, 60.1, 1.7, 64.0, 63.1, + # 1.0, 1.4, 7.9, 0.3, 3.9, 0.3, 6.7), + # c(22.6, 8.3, 44.4, 11.9, 24.6, 0.6, 5.7, 41.6, + # 0.0, 0.6, 6.7, 3.8, 1.0, 1.2, 1.4), + # alternative='two.sided', method='spearman') + x = [2.0, 47.4, 42.0, 10.8, 60.1, 1.7, 64.0, 63.1, + 1.0, 1.4, 7.9, 0.3, 3.9, 0.3, 6.7] + y = [22.6, 8.3, 44.4, 11.9, 24.6, 0.6, 5.7, 41.6, + 0.0, 0.6, 6.7, 3.8, 1.0, 1.2, 1.4] + + r_exp = 0.6887298747763864 # from cor.test + + r, p = mstats.spearmanr(x, y) + assert_allclose(r, r_exp) + assert_allclose(p, 0.004519192910756) + + r, p = mstats.spearmanr(x, y, alternative='greater') + assert_allclose(r, r_exp) + assert_allclose(p, 0.002259596455378) + + r, p = mstats.spearmanr(x, y, alternative='less') + assert_allclose(r, r_exp) + assert_allclose(p, 0.9977404035446) + + # intuitive test (with obvious positive correlation) + rng = np.random.default_rng(9607138567) + n = 100 + x = np.linspace(0, 5, n) + y = 0.1*x + rng.random(n) # y is positively correlated w/ x + + stat1, p1 = mstats.spearmanr(x, y) + + stat2, p2 = mstats.spearmanr(x, y, alternative="greater") + assert_allclose(p2, p1 / 2) # positive correlation -> small p + + stat3, p3 = mstats.spearmanr(x, y, alternative="less") + assert_allclose(p3, 1 - p1 / 2) # positive correlation -> large p + + assert stat1 == stat2 == stat3 + + with pytest.raises(ValueError, match="alternative must be 'less'..."): + mstats.spearmanr(x, y, alternative="ekki-ekki") + + @pytest.mark.skipif(platform.machine() == 'ppc64le', + reason="fails/crashes on ppc64le") + def test_kendalltau(self): + # check case with maximum disorder and p=1 + x = ma.array(np.array([9, 2, 5, 6])) + y = ma.array(np.array([4, 7, 9, 11])) + # Cross-check with exact result from R: + # cor.test(x,y,method="kendall",exact=1) + expected = [0.0, 1.0] + assert_almost_equal(np.asarray(mstats.kendalltau(x, y)), expected) + + # simple case without ties + x = ma.array(np.arange(10)) + y = ma.array(np.arange(10)) + # Cross-check with exact result from R: + # cor.test(x,y,method="kendall",exact=1) + expected = [1.0, 5.511463844797e-07] + assert_almost_equal(np.asarray(mstats.kendalltau(x, y)), expected) + + # check exception in case of invalid method keyword + assert_raises(ValueError, mstats.kendalltau, x, y, method='banana') + + # swap a couple of values + b = y[1] + y[1] = y[2] + y[2] = b + # Cross-check with exact result from R: + # cor.test(x,y,method="kendall",exact=1) + expected = [0.9555555555555556, 5.511463844797e-06] + assert_almost_equal(np.asarray(mstats.kendalltau(x, y)), expected) + + # swap a couple more + b = y[5] + y[5] = y[6] + y[6] = b + # Cross-check with exact result from R: + # cor.test(x,y,method="kendall",exact=1) + expected = [0.9111111111111111, 2.976190476190e-05] + assert_almost_equal(np.asarray(mstats.kendalltau(x, y)), expected) + + # same in opposite direction + x = ma.array(np.arange(10)) + y = ma.array(np.arange(10)[::-1]) + # Cross-check with exact result from R: + # cor.test(x,y,method="kendall",exact=1) + expected = [-1.0, 5.511463844797e-07] + assert_almost_equal(np.asarray(mstats.kendalltau(x, y)), expected) + + # swap a couple of values + b = y[1] + y[1] = y[2] + y[2] = b + # Cross-check with exact result from R: + # cor.test(x,y,method="kendall",exact=1) + expected = [-0.9555555555555556, 5.511463844797e-06] + assert_almost_equal(np.asarray(mstats.kendalltau(x, y)), expected) + + # swap a couple more + b = y[5] + y[5] = y[6] + y[6] = b + # Cross-check with exact result from R: + # cor.test(x,y,method="kendall",exact=1) + expected = [-0.9111111111111111, 2.976190476190e-05] + assert_almost_equal(np.asarray(mstats.kendalltau(x, y)), expected) + + # Tests some computations of Kendall's tau + x = ma.fix_invalid([5.05, 6.75, 3.21, 2.66, np.nan]) + y = ma.fix_invalid([1.65, 26.5, -5.93, 7.96, np.nan]) + z = ma.fix_invalid([1.65, 2.64, 2.64, 6.95, np.nan]) + assert_almost_equal(np.asarray(mstats.kendalltau(x, y)), + [+0.3333333, 0.75]) + assert_almost_equal(np.asarray(mstats.kendalltau(x, y, method='asymptotic')), + [+0.3333333, 0.4969059]) + assert_almost_equal(np.asarray(mstats.kendalltau(x, z)), + [-0.5477226, 0.2785987]) + # + x = ma.fix_invalid([0, 0, 0, 0, 20, 20, 0, 60, 0, 20, + 10, 10, 0, 40, 0, 20, 0, 0, 0, 0, 0, np.nan]) + y = ma.fix_invalid([0, 80, 80, 80, 10, 33, 60, 0, 67, 27, + 25, 80, 80, 80, 80, 80, 80, 0, 10, 45, np.nan, 0]) + result = mstats.kendalltau(x, y) + assert_almost_equal(np.asarray(result), [-0.1585188, 0.4128009]) + + # test for namedtuple attributes + attributes = ('correlation', 'pvalue') + check_named_results(result, attributes, ma=True) + + def test_overflow_gh6061(self): + # Regression test for GitHub issue #6061 - Overflow on Windows + x = np.arange(2000, dtype=float) + x = np.ma.masked_greater(x, 1995) + y = np.arange(2000, dtype=float) + y = np.concatenate((y[1000:], y[:1000])) + assert_(np.isfinite(stats.mstats.kendalltau(x,y)[1])) + + def test_kendalltau_vs_mstats_basic(self): + rng = np.random.RandomState(42) + for s in range(3, 10): + a = [] + # Generate rankings with ties + for i in range(s): + a += [i]*i + b = list(a) + rng.shuffle(a) + rng.shuffle(b) + expected = mstats.kendalltau(a, b) + actual = stats.kendalltau(a, b) + assert_almost_equal(actual[0], expected[0]) + assert_almost_equal(actual[1], expected[1]) + + @pytest.mark.skipif(platform.machine() == 'ppc64le', + reason="fails/crashes on ppc64le") + @pytest.mark.slow + def test_kendalltau_large(self): + # make sure internal variable use correct precision with + # larger arrays + x = np.arange(2000, dtype=float) + x = ma.masked_greater(x, 1995) + y = np.arange(2000, dtype=float) + y = np.concatenate((y[1000:], y[:1000])) + assert_(np.isfinite(mstats.kendalltau(x, y)[1])) + + def test_kendalltau_seasonal(self): + # Tests the seasonal Kendall tau. + x = [[nan, nan, 4, 2, 16, 26, 5, 1, 5, 1, 2, 3, 1], + [4, 3, 5, 3, 2, 7, 3, 1, 1, 2, 3, 5, 3], + [3, 2, 5, 6, 18, 4, 9, 1, 1, nan, 1, 1, nan], + [nan, 6, 11, 4, 17, nan, 6, 1, 1, 2, 5, 1, 1]] + x = ma.fix_invalid(x).T + output = mstats.kendalltau_seasonal(x) + assert_almost_equal(output['global p-value (indep)'], 0.008, 3) + assert_almost_equal(output['seasonal p-value'].round(2), + [0.18,0.53,0.20,0.04]) + + @pytest.mark.parametrize("method", ("exact", "asymptotic")) + @pytest.mark.parametrize("alternative", ("two-sided", "greater", "less")) + def test_kendalltau_mstats_vs_stats(self, method, alternative): + # Test that mstats.kendalltau and stats.kendalltau with + # nan_policy='omit' matches behavior of stats.kendalltau + # Accuracy of the alternatives is tested in stats/tests/test_stats.py + + rng = np.random.default_rng(8755235945) + n = 50 + x = rng.random(n) + y = rng.random(n) + mask = rng.random(n) > 0.5 + + x_masked = ma.array(x, mask=mask) + y_masked = ma.array(y, mask=mask) + res_masked = mstats.kendalltau( + x_masked, y_masked, method=method, alternative=alternative) + + x_compressed = x_masked.compressed() + y_compressed = y_masked.compressed() + res_compressed = stats.kendalltau( + x_compressed, y_compressed, method=method, alternative=alternative) + + x[mask] = np.nan + y[mask] = np.nan + res_nan = stats.kendalltau( + x, y, method=method, nan_policy='omit', alternative=alternative) + + assert_allclose(res_masked, res_compressed) + assert_allclose(res_nan, res_compressed) + + def test_kendall_p_exact_medium(self): + # Test for the exact method with medium samples (some n >= 171) + # expected values generated using SymPy + expectations = {(100, 2393): 0.62822615287956040664, + (101, 2436): 0.60439525773513602669, + (170, 0): 2.755801935583541e-307, + (171, 0): 0.0, + (171, 1): 2.755801935583541e-307, + (172, 1): 0.0, + (200, 9797): 0.74753983745929675209, + (201, 9656): 0.40959218958120363618} + for nc, expected in expectations.items(): + res = _mstats_basic._kendall_p_exact(nc[0], nc[1]) + assert_almost_equal(res, expected) + + @pytest.mark.xslow + def test_kendall_p_exact_large(self): + # Test for the exact method with large samples (n >= 171) + # expected values generated using SymPy + expectations = {(400, 38965): 0.48444283672113314099, + (401, 39516): 0.66363159823474837662, + (800, 156772): 0.42265448483120932055, + (801, 157849): 0.53437553412194416236, + (1600, 637472): 0.84200727400323538419, + (1601, 630304): 0.34465255088058593946} + + for nc, expected in expectations.items(): + res = _mstats_basic._kendall_p_exact(nc[0], nc[1]) + assert_almost_equal(res, expected) + + @skip_xp_invalid_arg + # mstats.pointbiserialr returns a NumPy float for the statistic, but converts + # it to a masked array with no masked elements before calling `special.betainc`, + # which won't accept masked arrays when `SCIPY_ARRAY_API=1`. + def test_pointbiserial(self): + x = [1, 0, 1, 1, 1, 1, 0, 1, 0, 0, 0, 1, 1, 0, 0, 0, 1, 1, 1, 0, 0, 0, + 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 1, -1] + y = [14.8, 13.8, 12.4, 10.1, 7.1, 6.1, 5.8, 4.6, 4.3, 3.5, 3.3, 3.2, + 3.0, 2.8, 2.8, 2.5, 2.4, 2.3, 2.1, 1.7, 1.7, 1.5, 1.3, 1.3, 1.2, + 1.2, 1.1, 0.8, 0.7, 0.6, 0.5, 0.2, 0.2, 0.1, np.nan] + assert_almost_equal(mstats.pointbiserialr(x, y)[0], 0.36149, 5) + + # test for namedtuple attributes + res = mstats.pointbiserialr(x, y) + attributes = ('correlation', 'pvalue') + check_named_results(res, attributes, ma=True) + + +@skip_xp_invalid_arg +class TestTrimming: + + def test_trim(self): + a = ma.arange(10) + assert_equal(mstats.trim(a), [0,1,2,3,4,5,6,7,8,9]) + a = ma.arange(10) + assert_equal(mstats.trim(a,(2,8)), [None,None,2,3,4,5,6,7,8,None]) + a = ma.arange(10) + assert_equal(mstats.trim(a,limits=(2,8),inclusive=(False,False)), + [None,None,None,3,4,5,6,7,None,None]) + a = ma.arange(10) + assert_equal(mstats.trim(a,limits=(0.1,0.2),relative=True), + [None,1,2,3,4,5,6,7,None,None]) + + a = ma.arange(12) + a[[0,-1]] = a[5] = masked + assert_equal(mstats.trim(a, (2,8)), + [None, None, 2, 3, 4, None, 6, 7, 8, None, None, None]) + + x = ma.arange(100).reshape(10, 10) + expected = [1]*10 + [0]*70 + [1]*20 + trimx = mstats.trim(x, (0.1,0.2), relative=True, axis=None) + assert_equal(trimx._mask.ravel(), expected) + trimx = mstats.trim(x, (0.1,0.2), relative=True, axis=0) + assert_equal(trimx._mask.ravel(), expected) + trimx = mstats.trim(x, (0.1,0.2), relative=True, axis=-1) + assert_equal(trimx._mask.T.ravel(), expected) + + # same as above, but with an extra masked row inserted + x = ma.arange(110).reshape(11, 10) + x[1] = masked + expected = [1]*20 + [0]*70 + [1]*20 + trimx = mstats.trim(x, (0.1,0.2), relative=True, axis=None) + assert_equal(trimx._mask.ravel(), expected) + trimx = mstats.trim(x, (0.1,0.2), relative=True, axis=0) + assert_equal(trimx._mask.ravel(), expected) + trimx = mstats.trim(x.T, (0.1,0.2), relative=True, axis=-1) + assert_equal(trimx.T._mask.ravel(), expected) + + def test_trim_old(self): + x = ma.arange(100) + assert_equal(mstats.trimboth(x).count(), 60) + assert_equal(mstats.trimtail(x,tail='r').count(), 80) + x[50:70] = masked + trimx = mstats.trimboth(x) + assert_equal(trimx.count(), 48) + assert_equal(trimx._mask, [1]*16 + [0]*34 + [1]*20 + [0]*14 + [1]*16) + x._mask = nomask + x = x.reshape((10,10)) + assert_equal(mstats.trimboth(x).count(), 60) + assert_equal(mstats.trimtail(x).count(), 80) + + def test_trimr(self): + x = ma.arange(10) + result = mstats.trimr(x, limits=(0.15, 0.14), inclusive=(False, False)) + expected = ma.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], + mask=[1, 1, 0, 0, 0, 0, 0, 0, 0, 1]) + assert_equal(result, expected) + assert_equal(result.mask, expected.mask) + + def test_trimmedmean(self): + data = ma.array([77, 87, 88,114,151,210,219,246,253,262, + 296,299,306,376,428,515,666,1310,2611]) + assert_almost_equal(mstats.trimmed_mean(data,0.1), 343, 0) + assert_almost_equal(mstats.trimmed_mean(data,(0.1,0.1)), 343, 0) + assert_almost_equal(mstats.trimmed_mean(data,(0.2,0.2)), 283, 0) + + def test_trimmedvar(self): + # Basic test. Additional tests of all arguments, edge cases, + # input validation, and proper treatment of masked arrays are needed. + rng = np.random.default_rng(3262323289434724460) + data_orig = rng.random(size=20) + data = np.sort(data_orig) + data = ma.array(data, mask=[1, 1, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 1, 1]) + assert_allclose(mstats.trimmed_var(data_orig, 0.1), data.var()) + + def test_trimmedstd(self): + # Basic test. Additional tests of all arguments, edge cases, + # input validation, and proper treatment of masked arrays are needed. + rng = np.random.default_rng(7121029245207162780) + data_orig = rng.random(size=20) + data = np.sort(data_orig) + data = ma.array(data, mask=[1, 1, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 1, 1]) + assert_allclose(mstats.trimmed_std(data_orig, 0.1), data.std()) + + def test_trimmed_stde(self): + data = ma.array([77, 87, 88,114,151,210,219,246,253,262, + 296,299,306,376,428,515,666,1310,2611]) + assert_almost_equal(mstats.trimmed_stde(data,(0.2,0.2)), 56.13193, 5) + assert_almost_equal(mstats.trimmed_stde(data,0.2), 56.13193, 5) + + def test_winsorization(self): + data = ma.array([77, 87, 88,114,151,210,219,246,253,262, + 296,299,306,376,428,515,666,1310,2611]) + assert_almost_equal(mstats.winsorize(data,(0.2,0.2)).var(ddof=1), + 21551.4, 1) + assert_almost_equal( + mstats.winsorize(data, (0.2,0.2),(False,False)).var(ddof=1), + 11887.3, 1) + data[5] = masked + winsorized = mstats.winsorize(data) + assert_equal(winsorized.mask, data.mask) + + def test_winsorization_nan(self): + data = ma.array([np.nan, np.nan, 0, 1, 2]) + assert_raises(ValueError, mstats.winsorize, data, (0.05, 0.05), + nan_policy='raise') + # Testing propagate (default behavior) + assert_equal(mstats.winsorize(data, (0.4, 0.4)), + ma.array([2, 2, 2, 2, 2])) + assert_equal(mstats.winsorize(data, (0.8, 0.8)), + ma.array([np.nan, np.nan, np.nan, np.nan, np.nan])) + assert_equal(mstats.winsorize(data, (0.4, 0.4), nan_policy='omit'), + ma.array([np.nan, np.nan, 2, 2, 2])) + assert_equal(mstats.winsorize(data, (0.8, 0.8), nan_policy='omit'), + ma.array([np.nan, np.nan, 2, 2, 2])) + + +@skip_xp_invalid_arg +class TestMoments: + # Comparison numbers are found using R v.1.5.1 + # note that length(testcase) = 4 + # testmathworks comes from documentation for the + # Statistics Toolbox for Matlab and can be found at both + # https://www.mathworks.com/help/stats/kurtosis.html + # https://www.mathworks.com/help/stats/skewness.html + # Note that both test cases came from here. + testcase = [1,2,3,4] + testmathworks = ma.fix_invalid([1.165, 0.6268, 0.0751, 0.3516, -0.6965, + np.nan]) + testcase_2d = ma.array( + np.array([[0.05245846, 0.50344235, 0.86589117, 0.36936353, 0.46961149], + [0.11574073, 0.31299969, 0.45925772, 0.72618805, 0.75194407], + [0.67696689, 0.91878127, 0.09769044, 0.04645137, 0.37615733], + [0.05903624, 0.29908861, 0.34088298, 0.66216337, 0.83160998], + [0.64619526, 0.94894632, 0.27855892, 0.0706151, 0.39962917]]), + mask=np.array([[True, False, False, True, False], + [True, True, True, False, True], + [False, False, False, False, False], + [True, True, True, True, True], + [False, False, True, False, False]], dtype=bool)) + + def _assert_equal(self, actual, expect, *, shape=None, dtype=None): + expect = np.asarray(expect) + if shape is not None: + expect = np.broadcast_to(expect, shape) + assert_array_equal(actual, expect) + if dtype is None: + dtype = expect.dtype + assert actual.dtype == dtype + + def test_moment(self): + y = mstats.moment(self.testcase,1) + assert_almost_equal(y,0.0,10) + y = mstats.moment(self.testcase,2) + assert_almost_equal(y,1.25) + y = mstats.moment(self.testcase,3) + assert_almost_equal(y,0.0) + y = mstats.moment(self.testcase,4) + assert_almost_equal(y,2.5625) + + # check array_like input for moment + y = mstats.moment(self.testcase, [1, 2, 3, 4]) + assert_allclose(y, [0, 1.25, 0, 2.5625]) + + # check moment input consists only of integers + y = mstats.moment(self.testcase, 0.0) + assert_allclose(y, 1.0) + assert_raises(ValueError, mstats.moment, self.testcase, 1.2) + y = mstats.moment(self.testcase, [1.0, 2, 3, 4.0]) + assert_allclose(y, [0, 1.25, 0, 2.5625]) + + # test empty input + y = mstats.moment([]) + self._assert_equal(y, np.nan, dtype=np.float64) + y = mstats.moment(np.array([], dtype=np.float32)) + self._assert_equal(y, np.nan, dtype=np.float32) + y = mstats.moment(np.zeros((1, 0)), axis=0) + self._assert_equal(y, [], shape=(0,), dtype=np.float64) + y = mstats.moment([[]], axis=1) + self._assert_equal(y, np.nan, shape=(1,), dtype=np.float64) + y = mstats.moment([[]], moment=[0, 1], axis=0) + self._assert_equal(y, [], shape=(2, 0)) + + x = np.arange(10.) + x[9] = np.nan + assert_equal(mstats.moment(x, 2), ma.masked) # NaN value is ignored + + def test_variation(self): + y = mstats.variation(self.testcase) + assert_almost_equal(y,0.44721359549996, 10) + + def test_variation_ddof(self): + # test variation with delta degrees of freedom + # regression test for gh-13341 + a = np.array([1, 2, 3, 4, 5]) + y = mstats.variation(a, ddof=1) + assert_almost_equal(y, 0.5270462766947299) + + def test_skewness(self): + y = mstats.skew(self.testmathworks) + assert_almost_equal(y,-0.29322304336607,10) + y = mstats.skew(self.testmathworks,bias=0) + assert_almost_equal(y,-0.437111105023940,10) + y = mstats.skew(self.testcase) + assert_almost_equal(y,0.0,10) + + # test that skew works on multidimensional masked arrays + correct_2d = ma.array( + np.array([0.6882870394455785, 0, 0.2665647526856708, + 0, -0.05211472114254485]), + mask=np.array([False, False, False, True, False], dtype=bool) + ) + assert_allclose(mstats.skew(self.testcase_2d, 1), correct_2d) + for i, row in enumerate(self.testcase_2d): + assert_almost_equal(mstats.skew(row), correct_2d[i]) + + correct_2d_bias_corrected = ma.array( + np.array([1.685952043212545, 0.0, 0.3973712716070531, 0, + -0.09026534484117164]), + mask=np.array([False, False, False, True, False], dtype=bool) + ) + assert_allclose(mstats.skew(self.testcase_2d, 1, bias=False), + correct_2d_bias_corrected) + for i, row in enumerate(self.testcase_2d): + assert_almost_equal(mstats.skew(row, bias=False), + correct_2d_bias_corrected[i]) + + # Check consistency between stats and mstats implementations + assert_allclose(mstats.skew(self.testcase_2d[2, :]), + stats.skew(self.testcase_2d[2, :])) + + def test_kurtosis(self): + # Set flags for axis = 0 and fisher=0 (Pearson's definition of kurtosis + # for compatibility with Matlab) + y = mstats.kurtosis(self.testmathworks, 0, fisher=0, bias=1) + assert_almost_equal(y, 2.1658856802973, 10) + # Note that MATLAB has confusing docs for the following case + # kurtosis(x,0) gives an unbiased estimate of Pearson's skewness + # kurtosis(x) gives a biased estimate of Fisher's skewness (Pearson-3) + # The MATLAB docs imply that both should give Fisher's + y = mstats.kurtosis(self.testmathworks, fisher=0, bias=0) + assert_almost_equal(y, 3.663542721189047, 10) + y = mstats.kurtosis(self.testcase, 0, 0) + assert_almost_equal(y, 1.64) + + # test that kurtosis works on multidimensional masked arrays + correct_2d = ma.array(np.array([-1.5, -3., -1.47247052385, 0., + -1.26979517952]), + mask=np.array([False, False, False, True, + False], dtype=bool)) + assert_array_almost_equal(mstats.kurtosis(self.testcase_2d, 1), + correct_2d) + for i, row in enumerate(self.testcase_2d): + assert_almost_equal(mstats.kurtosis(row), correct_2d[i]) + + correct_2d_bias_corrected = ma.array( + np.array([-1.5, -3., -1.88988209538, 0., -0.5234638463918877]), + mask=np.array([False, False, False, True, False], dtype=bool)) + assert_array_almost_equal(mstats.kurtosis(self.testcase_2d, 1, + bias=False), + correct_2d_bias_corrected) + for i, row in enumerate(self.testcase_2d): + assert_almost_equal(mstats.kurtosis(row, bias=False), + correct_2d_bias_corrected[i]) + + # Check consistency between stats and mstats implementations + assert_array_almost_equal_nulp(mstats.kurtosis(self.testcase_2d[2, :]), + stats.kurtosis(self.testcase_2d[2, :]), + nulp=4) + + +class TestMode: + def test_mode(self): + a1 = [0,0,0,1,1,1,2,3,3,3,3,4,5,6,7] + a2 = np.reshape(a1, (3,5)) + a3 = np.array([1,2,3,4,5,6]) + a4 = np.reshape(a3, (3,2)) + ma1 = ma.masked_where(ma.array(a1) > 2, a1) + ma2 = ma.masked_where(a2 > 2, a2) + ma3 = ma.masked_where(a3 < 2, a3) + ma4 = ma.masked_where(ma.array(a4) < 2, a4) + assert_equal(mstats.mode(a1, axis=None), (3,4)) + assert_equal(mstats.mode(a1, axis=0), (3,4)) + assert_equal(mstats.mode(ma1, axis=None), (0,3)) + assert_equal(mstats.mode(a2, axis=None), (3,4)) + assert_equal(mstats.mode(ma2, axis=None), (0,3)) + assert_equal(mstats.mode(a3, axis=None), (1,1)) + assert_equal(mstats.mode(ma3, axis=None), (2,1)) + assert_equal(mstats.mode(a2, axis=0), ([[0,0,0,1,1]], [[1,1,1,1,1]])) + assert_equal(mstats.mode(ma2, axis=0), ([[0,0,0,1,1]], [[1,1,1,1,1]])) + assert_equal(mstats.mode(a2, axis=-1), ([[0],[3],[3]], [[3],[3],[1]])) + assert_equal(mstats.mode(ma2, axis=-1), ([[0],[1],[0]], [[3],[1],[0]])) + assert_equal(mstats.mode(ma4, axis=0), ([[3,2]], [[1,1]])) + assert_equal(mstats.mode(ma4, axis=-1), ([[2],[3],[5]], [[1],[1],[1]])) + + a1_res = mstats.mode(a1, axis=None) + + # test for namedtuple attributes + attributes = ('mode', 'count') + check_named_results(a1_res, attributes, ma=True) + + def test_mode_modifies_input(self): + # regression test for gh-6428: mode(..., axis=None) may not modify + # the input array + im = np.zeros((100, 100)) + im[:50, :] += 1 + im[:, :50] += 1 + cp = im.copy() + mstats.mode(im, None) + assert_equal(im, cp) + + +class TestPercentile: + def setup_method(self): + self.a1 = [3, 4, 5, 10, -3, -5, 6] + self.a2 = [3, -6, -2, 8, 7, 4, 2, 1] + self.a3 = [3., 4, 5, 10, -3, -5, -6, 7.0] + + def test_percentile(self): + x = np.arange(8) * 0.5 + assert_equal(mstats.scoreatpercentile(x, 0), 0.) + assert_equal(mstats.scoreatpercentile(x, 100), 3.5) + assert_equal(mstats.scoreatpercentile(x, 50), 1.75) + + def test_2D(self): + x = ma.array([[1, 1, 1], + [1, 1, 1], + [4, 4, 3], + [1, 1, 1], + [1, 1, 1]]) + assert_equal(mstats.scoreatpercentile(x, 50), [1, 1, 1]) + + +@skip_xp_invalid_arg +class TestVariability: + """ Comparison numbers are found using R v.1.5.1 + note that length(testcase) = 4 + """ + testcase = ma.fix_invalid([1,2,3,4,np.nan]) + + def test_sem(self): + # This is not in R, so used: sqrt(var(testcase)*3/4) / sqrt(3) + y = mstats.sem(self.testcase) + assert_almost_equal(y, 0.6454972244) + n = self.testcase.count() + assert_allclose(mstats.sem(self.testcase, ddof=0) * np.sqrt(n/(n-2)), + mstats.sem(self.testcase, ddof=2)) + + def test_zmap(self): + # This is not in R, so tested by using: + # (testcase[i]-mean(testcase,axis=0)) / sqrt(var(testcase)*3/4) + y = mstats.zmap(self.testcase, self.testcase) + desired_unmaskedvals = ([-1.3416407864999, -0.44721359549996, + 0.44721359549996, 1.3416407864999]) + assert_array_almost_equal(desired_unmaskedvals, + y.data[y.mask == False], decimal=12) # noqa: E712 + + def test_zscore(self): + # This is not in R, so tested by using: + # (testcase[i]-mean(testcase,axis=0)) / sqrt(var(testcase)*3/4) + y = mstats.zscore(self.testcase) + desired = ma.fix_invalid([-1.3416407864999, -0.44721359549996, + 0.44721359549996, 1.3416407864999, np.nan]) + assert_almost_equal(desired, y, decimal=12) + + +@skip_xp_invalid_arg +class TestMisc: + + def test_obrientransform(self): + args = [[5]*5+[6]*11+[7]*9+[8]*3+[9]*2+[10]*2, + [6]+[7]*2+[8]*4+[9]*9+[10]*16] + result = [5*[3.1828]+11*[0.5591]+9*[0.0344]+3*[1.6086]+2*[5.2817]+2*[11.0538], + [10.4352]+2*[4.8599]+4*[1.3836]+9*[0.0061]+16*[0.7277]] + assert_almost_equal(np.round(mstats.obrientransform(*args).T, 4), + result, 4) + + def test_ks_2samp(self): + x = [[nan,nan, 4, 2, 16, 26, 5, 1, 5, 1, 2, 3, 1], + [4, 3, 5, 3, 2, 7, 3, 1, 1, 2, 3, 5, 3], + [3, 2, 5, 6, 18, 4, 9, 1, 1, nan, 1, 1, nan], + [nan, 6, 11, 4, 17, nan, 6, 1, 1, 2, 5, 1, 1]] + x = ma.fix_invalid(x).T + (winter, spring, summer, fall) = x.T + + assert_almost_equal(np.round(mstats.ks_2samp(winter, spring), 4), + (0.1818, 0.9628)) + assert_almost_equal(np.round(mstats.ks_2samp(winter, spring, 'g'), 4), + (0.1469, 0.6886)) + assert_almost_equal(np.round(mstats.ks_2samp(winter, spring, 'l'), 4), + (0.1818, 0.6011)) + + def test_friedmanchisq(self): + # No missing values + args = ([9.0,9.5,5.0,7.5,9.5,7.5,8.0,7.0,8.5,6.0], + [7.0,6.5,7.0,7.5,5.0,8.0,6.0,6.5,7.0,7.0], + [6.0,8.0,4.0,6.0,7.0,6.5,6.0,4.0,6.5,3.0]) + result = mstats.friedmanchisquare(*args) + assert_almost_equal(result[0], 10.4737, 4) + assert_almost_equal(result[1], 0.005317, 6) + # Missing values + x = [[nan,nan, 4, 2, 16, 26, 5, 1, 5, 1, 2, 3, 1], + [4, 3, 5, 3, 2, 7, 3, 1, 1, 2, 3, 5, 3], + [3, 2, 5, 6, 18, 4, 9, 1, 1,nan, 1, 1,nan], + [nan, 6, 11, 4, 17,nan, 6, 1, 1, 2, 5, 1, 1]] + x = ma.fix_invalid(x) + result = mstats.friedmanchisquare(*x) + assert_almost_equal(result[0], 2.0156, 4) + assert_almost_equal(result[1], 0.5692, 4) + + # test for namedtuple attributes + attributes = ('statistic', 'pvalue') + check_named_results(result, attributes, ma=True) + + # copied from `test_stats.py` + array = np.ma.asarray + x1 = [array([0.763, 0.599, 0.954, 0.628, 0.882, 0.936, 0.661, + 0.583, 0.775, 1.0, 0.94, 0.619, 0.972, 0.957]), + array([0.768, 0.591, 0.971, 0.661, 0.888, 0.931, 0.668, + 0.583, 0.838, 1.0, 0.962, 0.666, 0.981, 0.978]), + array([0.771, 0.590, 0.968, 0.654, 0.886, 0.916, 0.609, + 0.563, 0.866, 1.0, 0.965, 0.614, 0.9751, 0.946]), + array([0.798, 0.569, 0.967, 0.657, 0.898, 0.931, 0.685, + 0.625, 0.875, 1.0, 0.962, 0.669, 0.975, 0.970])] + + # From "Bioestadistica para las ciencias de la salud" Xf=18.95 p<0.001: + # x2 = [array([4, 3, 5, 3, 5, 3, 2, 5, 4, 4, 4, 3]), + # array([2, 2, 1, 2, 3, 1, 2, 3, 2, 1, 1, 3]), + # array([2, 4, 3, 3, 4, 3, 3, 4, 4, 1, 2, 1]), + # array([3, 5, 4, 3, 4, 4, 3, 3, 3, 4, 4, 4])] + + # From Jerrorl H. Zar, "Biostatistical Analysis"(example 12.6), + # Xf=10.68, 0.005 < p < 0.01: + # Probability from this example is inexact + # using Chisquare approximation of Friedman Chisquare. + x3 = [array([7.0, 9.9, 8.5, 5.1, 10.3]), + array([5.3, 5.7, 4.7, 3.5, 7.7]), + array([4.9, 7.6, 5.5, 2.8, 8.4]), + array([8.8, 8.9, 8.1, 3.3, 9.1])] + + # test using mstats + assert_array_almost_equal(mstats.friedmanchisquare(*x1), + (10.2283464566929, 0.0167215803284414)) + + # the following fails + # assert_array_almost_equal(mstats.friedmanchisquare(*x2), + # (18.9428571428571, 0.000280938375189499)) + + assert_array_almost_equal(mstats.friedmanchisquare(*x3), + (10.68, 0.0135882729582176)) + assert_raises(ValueError, mstats.friedmanchisquare, x3[0], x3[1]) + + +def test_regress_simple(): + # Regress a line with sinusoidal noise. Test for #1273. + x = np.linspace(0, 100, 100) + y = 0.2 * np.linspace(0, 100, 100) + 10 + y += np.sin(np.linspace(0, 20, 100)) + + result = mstats.linregress(x, y) + + # Result is of a correct class and with correct fields + lr = _stats_py.LinregressResult + assert_(isinstance(result, lr)) + attributes = ('slope', 'intercept', 'rvalue', 'pvalue', 'stderr') + check_named_results(result, attributes, ma=True) + assert 'intercept_stderr' in dir(result) + + # Slope and intercept are estimated correctly + assert_almost_equal(result.slope, 0.19644990055858422) + assert_almost_equal(result.intercept, 10.211269918932341) + assert_almost_equal(result.stderr, 0.002395781449783862) + assert_almost_equal(result.intercept_stderr, 0.13866936078570702) + + +def test_linregress_identical_x(): + rng = np.random.default_rng(74578245698) + x = np.zeros(10) + y = rng.random(10) + msg = "Cannot calculate a linear regression if all x values are identical" + with assert_raises(ValueError, match=msg): + mstats.linregress(x, y) + + +class TestTheilslopes: + def test_theilslopes(self): + # Test for basic slope and intercept. + slope, intercept, lower, upper = mstats.theilslopes([0, 1, 1]) + assert_almost_equal(slope, 0.5) + assert_almost_equal(intercept, 0.5) + + slope, intercept, lower, upper = mstats.theilslopes([0, 1, 1], + method='joint') + assert_almost_equal(slope, 0.5) + assert_almost_equal(intercept, 0.0) + + # Test for correct masking. + y = np.ma.array([0, 1, 100, 1], mask=[False, False, True, False]) + slope, intercept, lower, upper = mstats.theilslopes(y) + # These reference values are incorrect. With 100 masked, the outputs + # should be the same as above. + assert_almost_equal(slope, 1./3) + assert_almost_equal(intercept, 2./3) + + slope, intercept, lower, upper = mstats.theilslopes(y, + method='joint') + assert_almost_equal(slope, 1./3) + assert_almost_equal(intercept, 0.0) + + # Test of confidence intervals from example in Sen (1968). + x = [1, 2, 3, 4, 10, 12, 18] + y = [9, 15, 19, 20, 45, 55, 78] + slope, intercept, lower, upper = mstats.theilslopes(y, x, 0.07) + assert_almost_equal(slope, 4) + assert_almost_equal(intercept, 4.0) + assert_almost_equal(upper, 4.38, decimal=2) + assert_almost_equal(lower, 3.71, decimal=2) + + slope, intercept, lower, upper = mstats.theilslopes(y, x, 0.07, + method='joint') + assert_almost_equal(slope, 4) + assert_almost_equal(intercept, 6.0) + assert_almost_equal(upper, 4.38, decimal=2) + assert_almost_equal(lower, 3.71, decimal=2) + + + def test_theilslopes_warnings(self): + # Test `theilslopes` with degenerate input; see gh-15943 + msg = "All `x` coordinates.*|Mean of empty slice|invalid value encountered.*" + with pytest.warns(RuntimeWarning, match=msg): + res = mstats.theilslopes([0, 1], [0, 0]) + assert np.all(np.isnan(res)) + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", "invalid value encountered...", RuntimeWarning) + res = mstats.theilslopes([0, 0, 0], [0, 1, 0]) + assert_allclose(res, (0, 0, np.nan, np.nan)) + + + def test_theilslopes_namedtuple_consistency(self): + """ + Simple test to ensure tuple backwards-compatibility of the returned + TheilslopesResult object + """ + y = [1, 2, 4] + x = [4, 6, 8] + slope, intercept, low_slope, high_slope = mstats.theilslopes(y, x) + result = mstats.theilslopes(y, x) + + # note all four returned values are distinct here + assert_equal(slope, result.slope) + assert_equal(intercept, result.intercept) + assert_equal(low_slope, result.low_slope) + assert_equal(high_slope, result.high_slope) + + def test_gh19678_uint8(self): + # `theilslopes` returned unexpected results when `y` was an unsigned type. + # Check that this is resolved. + rng = np.random.default_rng(2549824598234528) + y = rng.integers(0, 255, size=10, dtype=np.uint8) + res = stats.theilslopes(y, y) + np.testing.assert_allclose(res.slope, 1) + + +def test_siegelslopes(): + # method should be exact for straight line + y = 2 * np.arange(10) + 0.5 + assert_equal(mstats.siegelslopes(y), (2.0, 0.5)) + assert_equal(mstats.siegelslopes(y, method='separate'), (2.0, 0.5)) + + x = 2 * np.arange(10) + y = 5 * x - 3.0 + assert_equal(mstats.siegelslopes(y, x), (5.0, -3.0)) + assert_equal(mstats.siegelslopes(y, x, method='separate'), (5.0, -3.0)) + + # method is robust to outliers: brekdown point of 50% + y[:4] = 1000 + assert_equal(mstats.siegelslopes(y, x), (5.0, -3.0)) + + # if there are no outliers, results should be comparable to linregress + x = np.arange(10) + y = -2.3 + 0.3*x + stats.norm.rvs(size=10, random_state=231) + slope_ols, intercept_ols, _, _, _ = stats.linregress(x, y) + + slope, intercept = mstats.siegelslopes(y, x) + assert_allclose(slope, slope_ols, rtol=0.1) + assert_allclose(intercept, intercept_ols, rtol=0.1) + + slope, intercept = mstats.siegelslopes(y, x, method='separate') + assert_allclose(slope, slope_ols, rtol=0.1) + assert_allclose(intercept, intercept_ols, rtol=0.1) + + +def test_siegelslopes_namedtuple_consistency(): + """ + Simple test to ensure tuple backwards-compatibility of the returned + SiegelslopesResult object. + """ + y = [1, 2, 4] + x = [4, 6, 8] + slope, intercept = mstats.siegelslopes(y, x) + result = mstats.siegelslopes(y, x) + + # note both returned values are distinct here + assert_equal(slope, result.slope) + assert_equal(intercept, result.intercept) + + +def test_sen_seasonal_slopes(): + rng = np.random.default_rng(5765986256978575148) + x = rng.random(size=(100, 4)) + intra_slope, inter_slope = mstats.sen_seasonal_slopes(x) + + # reference implementation from the `sen_seasonal_slopes` documentation + def dijk(yi): + n = len(yi) + x = np.arange(n) + dy = yi - yi[:, np.newaxis] + dx = x - x[:, np.newaxis] + mask = np.triu(np.ones((n, n), dtype=bool), k=1) + return dy[mask]/dx[mask] + + for i in range(4): + assert_allclose(np.median(dijk(x[:, i])), intra_slope[i]) + + all_slopes = np.concatenate([dijk(x[:, i]) for i in range(x.shape[1])]) + assert_allclose(np.median(all_slopes), inter_slope) + + +def test_plotting_positions(): + # Regression test for #1256 + pos = mstats.plotting_positions(np.arange(3), 0, 0) + assert_array_almost_equal(pos.data, np.array([0.25, 0.5, 0.75])) + + +@skip_xp_invalid_arg +class TestNormalitytests: + + def test_vs_nonmasked(self): + x = np.array((-2, -1, 0, 1, 2, 3)*4)**2 + assert_array_almost_equal(mstats.normaltest(x), + stats.normaltest(x)) + assert_array_almost_equal(mstats.skewtest(x), + stats.skewtest(x)) + assert_array_almost_equal(mstats.kurtosistest(x), + stats.kurtosistest(x)) + + funcs = [stats.normaltest, stats.skewtest, stats.kurtosistest] + mfuncs = [mstats.normaltest, mstats.skewtest, mstats.kurtosistest] + x = [1, 2, 3, 4] + for func, mfunc in zip(funcs, mfuncs): + with pytest.warns(SmallSampleWarning, match=too_small_1d_not_omit): + res = func(x) + assert np.isnan(res.statistic) + assert np.isnan(res.pvalue) + assert_raises(ValueError, mfunc, x) + + def test_axis_None(self): + # Test axis=None (equal to axis=0 for 1-D input) + x = np.array((-2,-1,0,1,2,3)*4)**2 + assert_allclose(mstats.normaltest(x, axis=None), mstats.normaltest(x)) + assert_allclose(mstats.skewtest(x, axis=None), mstats.skewtest(x)) + assert_allclose(mstats.kurtosistest(x, axis=None), + mstats.kurtosistest(x)) + + def test_maskedarray_input(self): + # Add some masked values, test result doesn't change + x = np.array((-2, -1, 0, 1, 2, 3)*4)**2 + xm = np.ma.array(np.r_[np.inf, x, 10], + mask=np.r_[True, [False] * x.size, True]) + assert_allclose(mstats.normaltest(xm), stats.normaltest(x)) + assert_allclose(mstats.skewtest(xm), stats.skewtest(x)) + assert_allclose(mstats.kurtosistest(xm), stats.kurtosistest(x)) + + def test_nd_input(self): + x = np.array((-2, -1, 0, 1, 2, 3)*4)**2 + x_2d = np.vstack([x] * 2).T + for func in [mstats.normaltest, mstats.skewtest, mstats.kurtosistest]: + res_1d = func(x) + res_2d = func(x_2d) + assert_allclose(res_2d[0], [res_1d[0]] * 2) + assert_allclose(res_2d[1], [res_1d[1]] * 2) + + def test_normaltest_result_attributes(self): + x = np.array((-2, -1, 0, 1, 2, 3)*4)**2 + res = mstats.normaltest(x) + attributes = ('statistic', 'pvalue') + check_named_results(res, attributes, ma=True) + + def test_kurtosistest_result_attributes(self): + x = np.array((-2, -1, 0, 1, 2, 3)*4)**2 + res = mstats.kurtosistest(x) + attributes = ('statistic', 'pvalue') + check_named_results(res, attributes, ma=True) + + def test_regression_9033(self): + # x clearly non-normal but power of negative denom needs + # to be handled correctly to reject normality + counts = [128, 0, 58, 7, 0, 41, 16, 0, 0, 167] + x = np.hstack([np.full(c, i) for i, c in enumerate(counts)]) + assert_equal(mstats.kurtosistest(x)[1] < 0.01, True) + + @pytest.mark.parametrize("test", ["skewtest", "kurtosistest"]) + @pytest.mark.parametrize("alternative", ["less", "greater"]) + def test_alternative(self, test, alternative): + x = stats.norm.rvs(loc=10, scale=2.5, size=30, random_state=123) + + stats_test = getattr(stats, test) + mstats_test = getattr(mstats, test) + + z_ex, p_ex = stats_test(x, alternative=alternative) + z, p = mstats_test(x, alternative=alternative) + assert_allclose(z, z_ex, atol=1e-12) + assert_allclose(p, p_ex, atol=1e-12) + + # test with masked arrays + x[1:5] = np.nan + x = np.ma.masked_array(x, mask=np.isnan(x)) + z_ex, p_ex = stats_test(x.compressed(), alternative=alternative) + z, p = mstats_test(x, alternative=alternative) + assert_allclose(z, z_ex, atol=1e-12) + assert_allclose(p, p_ex, atol=1e-12) + + def test_bad_alternative(self): + x = stats.norm.rvs(size=20, random_state=123) + msg = r"`alternative` must be..." + + with pytest.raises(ValueError, match=msg): + mstats.skewtest(x, alternative='error') + + with pytest.raises(ValueError, match=msg): + mstats.kurtosistest(x, alternative='error') + + +class TestFOneway: + def test_result_attributes(self): + a = np.array([655, 788], dtype=np.uint16) + b = np.array([789, 772], dtype=np.uint16) + res = mstats.f_oneway(a, b) + attributes = ('statistic', 'pvalue') + check_named_results(res, attributes, ma=True) + + +class TestMannwhitneyu: + # data from gh-1428 + x = np.array([1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., + 1., 1., 1., 1., 1., 1., 1., 2., 1., 1., 1., 1., 1., 1., + 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., + 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., + 1., 1., 1., 1., 1., 1., 1., 2., 1., 1., 1., 1., 1., 1., + 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., + 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 2., + 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., + 1., 1., 2., 1., 1., 1., 1., 2., 1., 1., 2., 1., 1., 2., + 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., + 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 2., 1., + 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., + 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., + 1., 1., 1., 1., 1., 1., 1., 2., 1., 1., 1., 1., 1., 1., + 1., 1., 1., 2., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., + 1., 1., 1., 1., 1., 1., 1., 1., 3., 1., 1., 1., 1., 1., + 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., + 1., 1., 1., 1., 1., 1.]) + + y = np.array([1., 1., 1., 1., 1., 1., 1., 2., 1., 2., 1., 1., 1., 1., + 2., 1., 1., 1., 2., 1., 1., 1., 1., 1., 2., 1., 1., 3., + 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 2., 1., 2., 1., + 1., 1., 1., 1., 1., 2., 1., 1., 1., 1., 1., 1., 1., 1., + 1., 1., 1., 1., 1., 1., 1., 2., 1., 1., 1., 1., 1., 2., + 2., 1., 1., 2., 1., 1., 2., 1., 2., 1., 1., 1., 1., 2., + 2., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., + 1., 2., 1., 1., 1., 1., 1., 2., 2., 2., 1., 1., 1., 1., + 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., + 2., 1., 1., 2., 1., 1., 1., 1., 2., 1., 1., 1., 1., 1., + 1., 1., 1., 1., 1., 1., 1., 2., 1., 1., 1., 2., 1., 1., + 1., 1., 1., 1.]) + + def test_result_attributes(self): + res = mstats.mannwhitneyu(self.x, self.y) + attributes = ('statistic', 'pvalue') + check_named_results(res, attributes, ma=True) + + def test_against_stats(self): + # gh-4641 reported that stats.mannwhitneyu returned half the p-value + # of mstats.mannwhitneyu. Default alternative of stats.mannwhitneyu + # is now two-sided, so they match. + res1 = mstats.mannwhitneyu(self.x, self.y) + res2 = stats.mannwhitneyu(self.x, self.y) + assert res1.statistic == res2.statistic + assert_allclose(res1.pvalue, res2.pvalue) + + +class TestKruskal: + def test_result_attributes(self): + x = [1, 3, 5, 7, 9] + y = [2, 4, 6, 8, 10] + + res = mstats.kruskal(x, y) + attributes = ('statistic', 'pvalue') + check_named_results(res, attributes, ma=True) + + +# TODO: for all ttest functions, add tests with masked array inputs +class TestTtest_rel: + def setup_method(self): + self.rng = np.random.default_rng(9373599542) + + def test_vs_nonmasked(self): + outcome = self.rng.standard_normal((20, 4)) + [0, 0, 1, 2] + + # 1-D inputs + res1 = stats.ttest_rel(outcome[:, 0], outcome[:, 1]) + res2 = mstats.ttest_rel(outcome[:, 0], outcome[:, 1]) + assert_allclose(res1, res2) + + # 2-D inputs + res1 = stats.ttest_rel(outcome[:, 0], outcome[:, 1], axis=None) + res2 = mstats.ttest_rel(outcome[:, 0], outcome[:, 1], axis=None) + assert_allclose(res1, res2) + res1 = stats.ttest_rel(outcome[:, :2], outcome[:, 2:], axis=0) + res2 = mstats.ttest_rel(outcome[:, :2], outcome[:, 2:], axis=0) + assert_allclose(res1, res2) + + # Check default is axis=0 + res3 = mstats.ttest_rel(outcome[:, :2], outcome[:, 2:]) + assert_allclose(res2, res3) + + def test_fully_masked(self): + outcome = ma.masked_array(self.rng.standard_normal((3, 2)), + mask=[[1, 1, 1], [0, 0, 0]]) + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", "invalid value encountered in absolute", RuntimeWarning) + for pair in [(outcome[:, 0], outcome[:, 1]), + ([np.nan, np.nan], [1.0, 2.0])]: + t, p = mstats.ttest_rel(*pair) + assert_array_equal(t, (np.nan, np.nan)) + assert_array_equal(p, (np.nan, np.nan)) + + def test_result_attributes(self): + outcome = self.rng.standard_normal((20, 4)) + [0, 0, 1, 2] + + res = mstats.ttest_rel(outcome[:, 0], outcome[:, 1]) + attributes = ('statistic', 'pvalue') + check_named_results(res, attributes, ma=True) + + def test_invalid_input_size(self): + assert_raises(ValueError, mstats.ttest_rel, + np.arange(10), np.arange(11)) + x = np.arange(24) + assert_raises(ValueError, mstats.ttest_rel, + x.reshape(2, 3, 4), x.reshape(2, 4, 3), axis=1) + assert_raises(ValueError, mstats.ttest_rel, + x.reshape(2, 3, 4), x.reshape(2, 4, 3), axis=2) + + def test_empty(self): + res1 = mstats.ttest_rel([], []) + assert_(np.all(np.isnan(res1))) + + def test_zero_division(self): + t, p = mstats.ttest_ind([0, 0, 0], [1, 1, 1]) + assert_equal((np.abs(t), p), (np.inf, 0)) + + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", "invalid value encountered in absolute", RuntimeWarning) + t, p = mstats.ttest_ind([0, 0, 0], [0, 0, 0]) + assert_array_equal(t, np.array([np.nan, np.nan])) + assert_array_equal(p, np.array([np.nan, np.nan])) + + def test_bad_alternative(self): + msg = r"alternative must be 'less', 'greater' or 'two-sided'" + with pytest.raises(ValueError, match=msg): + mstats.ttest_ind([1, 2, 3], [4, 5, 6], alternative='foo') + + @pytest.mark.parametrize("alternative", ["less", "greater"]) + def test_alternative(self, alternative): + x = stats.norm.rvs(loc=10, scale=5, size=25, random_state=42) + y = stats.norm.rvs(loc=8, scale=2, size=25, random_state=42) + + t_ex, p_ex = stats.ttest_rel(x, y, alternative=alternative) + t, p = mstats.ttest_rel(x, y, alternative=alternative) + assert_allclose(t, t_ex, rtol=1e-14) + assert_allclose(p, p_ex, rtol=1e-14) + + # test with masked arrays + x[1:10] = np.nan + y[1:10] = np.nan + x = np.ma.masked_array(x, mask=np.isnan(x)) + y = np.ma.masked_array(y, mask=np.isnan(y)) + t, p = mstats.ttest_rel(x, y, alternative=alternative) + t_ex, p_ex = stats.ttest_rel(x.compressed(), y.compressed(), + alternative=alternative) + assert_allclose(t, t_ex, rtol=1e-14) + assert_allclose(p, p_ex, rtol=1e-14) + + +class TestTtest_ind: + def setup_method(self): + self.rng = np.random.default_rng(4776115069) + + def test_vs_nonmasked(self): + outcome = self.rng.standard_normal((20, 4)) + [0, 0, 1, 2] + + # 1-D inputs + res1 = stats.ttest_ind(outcome[:, 0], outcome[:, 1]) + res2 = mstats.ttest_ind(outcome[:, 0], outcome[:, 1]) + assert_allclose(res1, res2) + + # 2-D inputs + res1 = stats.ttest_ind(outcome[:, 0], outcome[:, 1], axis=None) + res2 = mstats.ttest_ind(outcome[:, 0], outcome[:, 1], axis=None) + assert_allclose(res1, res2) + res1 = stats.ttest_ind(outcome[:, :2], outcome[:, 2:], axis=0) + res2 = mstats.ttest_ind(outcome[:, :2], outcome[:, 2:], axis=0) + assert_allclose(res1, res2) + + # Check default is axis=0 + res3 = mstats.ttest_ind(outcome[:, :2], outcome[:, 2:]) + assert_allclose(res2, res3) + + # Check equal_var + res4 = stats.ttest_ind(outcome[:, 0], outcome[:, 1], equal_var=True) + res5 = mstats.ttest_ind(outcome[:, 0], outcome[:, 1], equal_var=True) + assert_allclose(res4, res5) + res4 = stats.ttest_ind(outcome[:, 0], outcome[:, 1], equal_var=False) + res5 = mstats.ttest_ind(outcome[:, 0], outcome[:, 1], equal_var=False) + assert_allclose(res4, res5) + + def test_fully_masked(self): + outcome = ma.masked_array(self.rng.standard_normal((3, 2)), + mask=[[1, 1, 1], [0, 0, 0]]) + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", "invalid value encountered in absolute", RuntimeWarning) + for pair in [(outcome[:, 0], outcome[:, 1]), + ([np.nan, np.nan], [1.0, 2.0])]: + t, p = mstats.ttest_ind(*pair) + assert_array_equal(t, (np.nan, np.nan)) + assert_array_equal(p, (np.nan, np.nan)) + + def test_result_attributes(self): + outcome = self.rng.standard_normal((20, 4)) + [0, 0, 1, 2] + + res = mstats.ttest_ind(outcome[:, 0], outcome[:, 1]) + attributes = ('statistic', 'pvalue') + check_named_results(res, attributes, ma=True) + + def test_empty(self): + res1 = mstats.ttest_ind([], []) + assert_(np.all(np.isnan(res1))) + + def test_zero_division(self): + t, p = mstats.ttest_ind([0, 0, 0], [1, 1, 1]) + assert_equal((np.abs(t), p), (np.inf, 0)) + + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", "invalid value encountered in absolute", RuntimeWarning) + t, p = mstats.ttest_ind([0, 0, 0], [0, 0, 0]) + assert_array_equal(t, (np.nan, np.nan)) + assert_array_equal(p, (np.nan, np.nan)) + + t, p = mstats.ttest_ind([0, 0, 0], [1, 1, 1], equal_var=False) + assert_equal((np.abs(t), p), (np.inf, 0)) + assert_array_equal(mstats.ttest_ind([0, 0, 0], [0, 0, 0], + equal_var=False), (np.nan, np.nan)) + + def test_bad_alternative(self): + msg = r"alternative must be 'less', 'greater' or 'two-sided'" + with pytest.raises(ValueError, match=msg): + mstats.ttest_ind([1, 2, 3], [4, 5, 6], alternative='foo') + + @pytest.mark.parametrize("alternative", ["less", "greater"]) + def test_alternative(self, alternative): + x = stats.norm.rvs(loc=10, scale=2, size=100, random_state=123) + y = stats.norm.rvs(loc=8, scale=2, size=100, random_state=123) + + t_ex, p_ex = stats.ttest_ind(x, y, alternative=alternative) + t, p = mstats.ttest_ind(x, y, alternative=alternative) + assert_allclose(t, t_ex, rtol=1e-14) + assert_allclose(p, p_ex, rtol=1e-14) + + # test with masked arrays + x[1:10] = np.nan + y[80:90] = np.nan + x = np.ma.masked_array(x, mask=np.isnan(x)) + y = np.ma.masked_array(y, mask=np.isnan(y)) + t_ex, p_ex = stats.ttest_ind(x.compressed(), y.compressed(), + alternative=alternative) + t, p = mstats.ttest_ind(x, y, alternative=alternative) + assert_allclose(t, t_ex, rtol=1e-14) + assert_allclose(p, p_ex, rtol=1e-14) + + +class TestTtest_1samp: + def setup_method(self): + self.rng = np.random.default_rng(6043813830) + + def test_vs_nonmasked(self): + outcome = self.rng.standard_normal((20, 4)) + [0, 0, 1, 2] + + # 1-D inputs + res1 = stats.ttest_1samp(outcome[:, 0], 1) + res2 = mstats.ttest_1samp(outcome[:, 0], 1) + assert_allclose(res1, res2) + + def test_fully_masked(self): + outcome = ma.masked_array(self.rng.standard_normal(3), mask=[1, 1, 1]) + expected = (np.nan, np.nan) + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", "invalid value encountered in absolute", RuntimeWarning) + for pair in [((np.nan, np.nan), 0.0), (outcome, 0.0)]: + t, p = mstats.ttest_1samp(*pair) + assert_array_equal(p, expected) + assert_array_equal(t, expected) + + def test_result_attributes(self): + outcome = self.rng.standard_normal((20, 4)) + [0, 0, 1, 2] + + res = mstats.ttest_1samp(outcome[:, 0], 1) + attributes = ('statistic', 'pvalue') + check_named_results(res, attributes, ma=True) + + def test_empty(self): + res1 = mstats.ttest_1samp([], 1) + assert_(np.all(np.isnan(res1))) + + def test_zero_division(self): + t, p = mstats.ttest_1samp([0, 0, 0], 1) + assert_equal((np.abs(t), p), (np.inf, 0)) + + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", "invalid value encountered in absolute", RuntimeWarning) + t, p = mstats.ttest_1samp([0, 0, 0], 0) + assert_(np.isnan(t)) + assert_array_equal(p, (np.nan, np.nan)) + + def test_bad_alternative(self): + msg = r"alternative must be 'less', 'greater' or 'two-sided'" + with pytest.raises(ValueError, match=msg): + mstats.ttest_1samp([1, 2, 3], 4, alternative='foo') + + @pytest.mark.parametrize("alternative", ["less", "greater"]) + def test_alternative(self, alternative): + x = stats.norm.rvs(loc=10, scale=2, size=100, random_state=123) + + t_ex, p_ex = stats.ttest_1samp(x, 9, alternative=alternative) + t, p = mstats.ttest_1samp(x, 9, alternative=alternative) + assert_allclose(t, t_ex, rtol=1e-14) + assert_allclose(p, p_ex, rtol=1e-14) + + # test with masked arrays + x[1:10] = np.nan + x = np.ma.masked_array(x, mask=np.isnan(x)) + t_ex, p_ex = stats.ttest_1samp(x.compressed(), 9, + alternative=alternative) + t, p = mstats.ttest_1samp(x, 9, alternative=alternative) + assert_allclose(t, t_ex, rtol=1e-14) + assert_allclose(p, p_ex, rtol=1e-14) + + +class TestDescribe: + """ + Tests for mstats.describe. + + Note that there are also tests for `mstats.describe` in the + class TestCompareWithStats. + """ + def test_basic_with_axis(self): + # This is a basic test that is also a regression test for gh-7303. + a = np.ma.masked_array([[0, 1, 2, 3, 4, 9], + [5, 5, 0, 9, 3, 3]], + mask=[[0, 0, 0, 0, 0, 1], + [0, 0, 1, 1, 0, 0]]) + result = mstats.describe(a, axis=1) + assert_equal(result.nobs, [5, 4]) + amin, amax = result.minmax + assert_equal(amin, [0, 3]) + assert_equal(amax, [4, 5]) + assert_equal(result.mean, [2.0, 4.0]) + assert_equal(result.variance, [2.0, 1.0]) + assert_equal(result.skewness, [0.0, 0.0]) + assert_allclose(result.kurtosis, [-1.3, -2.0]) + + +@skip_xp_invalid_arg +class TestCompareWithStats: + """ + Class to compare mstats results with stats results. + + It is in general assumed that scipy.stats is at a more mature stage than + stats.mstats. If a routine in mstats results in similar results like in + scipy.stats, this is considered also as a proper validation of scipy.mstats + routine. + + Different sample sizes are used for testing, as some problems between stats + and mstats are dependent on sample size. + + Author: Alexander Loew + + NOTE that some tests fail. This might be caused by + a) actual differences or bugs between stats and mstats + b) numerical inaccuracies + c) different definitions of routine interfaces + + These failures need to be checked. Current workaround is to have disabled these + tests, but issuing reports on scipy-dev + + """ + def get_n(self): + """ Returns list of sample sizes to be used for comparison. """ + return [1000, 100, 10, 5] + + def generate_xy_sample(self, n): + # This routine generates numpy arrays and corresponding masked arrays + # with the same data, but additional masked values + rng = np.random.RandomState(1234567) + x = rng.randn(n) + y = x + rng.randn(n) + xm = np.full(len(x) + 5, 1e16) + ym = np.full(len(y) + 5, 1e16) + xm[0:len(x)] = x + ym[0:len(y)] = y + mask = xm > 9e15 + xm = np.ma.array(xm, mask=mask) + ym = np.ma.array(ym, mask=mask) + return x, y, xm, ym + + def generate_xy_sample2D(self, n, nx): + x = np.full((n, nx), np.nan) + y = np.full((n, nx), np.nan) + xm = np.full((n+5, nx), np.nan) + ym = np.full((n+5, nx), np.nan) + + for i in range(nx): + x[:, i], y[:, i], dx, dy = self.generate_xy_sample(n) + + xm[0:n, :] = x[0:n] + ym[0:n, :] = y[0:n] + xm = np.ma.array(xm, mask=np.isnan(xm)) + ym = np.ma.array(ym, mask=np.isnan(ym)) + return x, y, xm, ym + + def test_linregress(self): + for n in self.get_n(): + x, y, xm, ym = self.generate_xy_sample(n) + result1 = stats.linregress(x, y) + result2 = stats.mstats.linregress(xm, ym) + assert_allclose(np.asarray(result1), np.asarray(result2)) + + def test_pearsonr(self): + for n in self.get_n(): + x, y, xm, ym = self.generate_xy_sample(n) + r, p = stats.pearsonr(x, y) + rm, pm = stats.mstats.pearsonr(xm, ym) + + assert_almost_equal(r, rm, decimal=14) + assert_almost_equal(p, pm, decimal=14) + + def test_spearmanr(self): + for n in self.get_n(): + x, y, xm, ym = self.generate_xy_sample(n) + r, p = stats.spearmanr(x, y) + rm, pm = stats.mstats.spearmanr(xm, ym) + assert_almost_equal(r, rm, 14) + assert_almost_equal(p, pm, 14) + + def test_spearmanr_backcompat_useties(self): + # A regression test to ensure we don't break backwards compat + # more than we have to (see gh-9204). + x = np.arange(6) + assert_raises(ValueError, mstats.spearmanr, x, x, False) + + def test_gmean(self): + for n in self.get_n(): + x, y, xm, ym = self.generate_xy_sample(n) + r = stats.gmean(abs(x)) + rm = stats.mstats.gmean(abs(xm)) + assert_allclose(r, rm, rtol=1e-13) + + r = stats.gmean(abs(y)) + rm = stats.mstats.gmean(abs(ym)) + assert_allclose(r, rm, rtol=1e-13) + + def test_hmean(self): + for n in self.get_n(): + x, y, xm, ym = self.generate_xy_sample(n) + + r = stats.hmean(abs(x)) + rm = stats.mstats.hmean(abs(xm)) + assert_almost_equal(r, rm, 10) + + r = stats.hmean(abs(y)) + rm = stats.mstats.hmean(abs(ym)) + assert_almost_equal(r, rm, 10) + + def test_skew(self): + for n in self.get_n(): + x, y, xm, ym = self.generate_xy_sample(n) + + r = stats.skew(x) + rm = stats.mstats.skew(xm) + assert_almost_equal(r, rm, 10) + + r = stats.skew(y) + rm = stats.mstats.skew(ym) + assert_almost_equal(r, rm, 10) + + def test_moment(self): + for n in self.get_n(): + x, y, xm, ym = self.generate_xy_sample(n) + + r = stats.moment(x) + rm = stats.mstats.moment(xm) + assert_almost_equal(r, rm, 10) + + r = stats.moment(y) + rm = stats.mstats.moment(ym) + assert_almost_equal(r, rm, 10) + + def test_zscore(self): + for n in self.get_n(): + x, y, xm, ym = self.generate_xy_sample(n) + + # reference solution + zx = (x - x.mean()) / x.std() + zy = (y - y.mean()) / y.std() + + # validate stats + assert_allclose(stats.zscore(x), zx, rtol=1e-10) + assert_allclose(stats.zscore(y), zy, rtol=1e-10) + + # compare stats and mstats + assert_allclose(stats.zscore(x), stats.mstats.zscore(xm[0:len(x)]), + rtol=1e-10) + assert_allclose(stats.zscore(y), stats.mstats.zscore(ym[0:len(y)]), + rtol=1e-10) + + def test_kurtosis(self): + for n in self.get_n(): + x, y, xm, ym = self.generate_xy_sample(n) + r = stats.kurtosis(x) + rm = stats.mstats.kurtosis(xm) + assert_almost_equal(r, rm, 10) + + r = stats.kurtosis(y) + rm = stats.mstats.kurtosis(ym) + assert_almost_equal(r, rm, 10) + + def test_sem(self): + # example from stats.sem doc + a = np.arange(20).reshape(5, 4) + am = np.ma.array(a) + r = stats.sem(a, ddof=1) + rm = stats.mstats.sem(am, ddof=1) + + assert_allclose(r, 2.82842712, atol=1e-5) + assert_allclose(rm, 2.82842712, atol=1e-5) + + for n in self.get_n(): + x, y, xm, ym = self.generate_xy_sample(n) + assert_almost_equal(stats.mstats.sem(xm, axis=None, ddof=0), + stats.sem(x, axis=None, ddof=0), decimal=13) + assert_almost_equal(stats.mstats.sem(ym, axis=None, ddof=0), + stats.sem(y, axis=None, ddof=0), decimal=13) + assert_almost_equal(stats.mstats.sem(xm, axis=None, ddof=1), + stats.sem(x, axis=None, ddof=1), decimal=13) + assert_almost_equal(stats.mstats.sem(ym, axis=None, ddof=1), + stats.sem(y, axis=None, ddof=1), decimal=13) + + def test_describe(self): + for n in self.get_n(): + x, y, xm, ym = self.generate_xy_sample(n) + r = stats.describe(x, ddof=1) + rm = stats.mstats.describe(xm, ddof=1) + for ii in range(6): + assert_almost_equal(np.asarray(r[ii]), + np.asarray(rm[ii]), + decimal=12) + + def test_describe_result_attributes(self): + actual = mstats.describe(np.arange(5)) + attributes = ('nobs', 'minmax', 'mean', 'variance', 'skewness', + 'kurtosis') + check_named_results(actual, attributes, ma=True) + + def test_rankdata(self): + for n in self.get_n(): + x, y, xm, ym = self.generate_xy_sample(n) + r = stats.rankdata(x) + rm = stats.mstats.rankdata(x) + assert_allclose(r, rm) + + def test_tmean(self): + for n in self.get_n(): + x, y, xm, ym = self.generate_xy_sample(n) + assert_almost_equal(stats.tmean(x),stats.mstats.tmean(xm), 14) + assert_almost_equal(stats.tmean(y),stats.mstats.tmean(ym), 14) + + def test_tmax(self): + for n in self.get_n(): + x, y, xm, ym = self.generate_xy_sample(n) + assert_almost_equal(stats.tmax(x,2.), + stats.mstats.tmax(xm,2.), 10) + assert_almost_equal(stats.tmax(y,2.), + stats.mstats.tmax(ym,2.), 10) + + assert_almost_equal(stats.tmax(x, upperlimit=3.), + stats.mstats.tmax(xm, upperlimit=3.), 10) + assert_almost_equal(stats.tmax(y, upperlimit=3.), + stats.mstats.tmax(ym, upperlimit=3.), 10) + + def test_tmin(self): + for n in self.get_n(): + x, y, xm, ym = self.generate_xy_sample(n) + assert_equal(stats.tmin(x), stats.mstats.tmin(xm)) + assert_equal(stats.tmin(y), stats.mstats.tmin(ym)) + + assert_almost_equal(stats.tmin(x, lowerlimit=-1.), + stats.mstats.tmin(xm, lowerlimit=-1.), 10) + assert_almost_equal(stats.tmin(y, lowerlimit=-1.), + stats.mstats.tmin(ym, lowerlimit=-1.), 10) + + def test_zmap(self): + for n in self.get_n(): + x, y, xm, ym = self.generate_xy_sample(n) + z = stats.zmap(x, y) + zm = stats.mstats.zmap(xm, ym) + assert_allclose(z, zm[0:len(z)], atol=1e-10) + + def test_variation(self): + for n in self.get_n(): + x, y, xm, ym = self.generate_xy_sample(n) + assert_almost_equal(stats.variation(x), stats.mstats.variation(xm), + decimal=12) + assert_almost_equal(stats.variation(y), stats.mstats.variation(ym), + decimal=12) + + def test_tvar(self): + for n in self.get_n(): + x, y, xm, ym = self.generate_xy_sample(n) + assert_almost_equal(stats.tvar(x), stats.mstats.tvar(xm), + decimal=12) + assert_almost_equal(stats.tvar(y), stats.mstats.tvar(ym), + decimal=12) + + def test_trimboth(self): + a = np.arange(20) + b = stats.trimboth(a, 0.1) + bm = stats.mstats.trimboth(a, 0.1) + assert_allclose(np.sort(b), bm.data[~bm.mask]) + + def test_tsem(self): + for n in self.get_n(): + x, y, xm, ym = self.generate_xy_sample(n) + assert_almost_equal(stats.tsem(x), stats.mstats.tsem(xm), + decimal=14) + assert_almost_equal(stats.tsem(y), stats.mstats.tsem(ym), + decimal=14) + assert_almost_equal(stats.tsem(x, limits=(-2., 2.)), + stats.mstats.tsem(xm, limits=(-2., 2.)), + decimal=14) + + def test_skewtest(self): + # this test is for 1D data + for n in self.get_n(): + if n > 8: + x, y, xm, ym = self.generate_xy_sample(n) + r = stats.skewtest(x) + rm = stats.mstats.skewtest(xm) + assert_allclose(r, rm) + + def test_skewtest_result_attributes(self): + x = np.array((-2, -1, 0, 1, 2, 3)*4)**2 + res = mstats.skewtest(x) + attributes = ('statistic', 'pvalue') + check_named_results(res, attributes, ma=True) + + def test_skewtest_2D_notmasked(self): + # a normal ndarray is passed to the masked function + rng = np.random.default_rng(2790153686) + x = rng.random((20, 2)) * 20. + r = stats.skewtest(x) + rm = stats.mstats.skewtest(x) + assert_allclose(np.asarray(r), np.asarray(rm)) + + def test_skewtest_2D_WithMask(self): + nx = 2 + for n in self.get_n(): + if n > 8: + x, y, xm, ym = self.generate_xy_sample2D(n, nx) + r = stats.skewtest(x) + rm = stats.mstats.skewtest(xm) + + assert_allclose(r[0][0], rm[0][0], rtol=1e-14) + assert_allclose(r[0][1], rm[0][1], rtol=1e-14) + + def test_normaltest(self): + with np.errstate(over='raise'), warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", "`kurtosistest` p-value may be inaccurate", UserWarning) + warnings.filterwarnings( + "ignore", "kurtosistest only valid for n>=20", UserWarning) + for n in self.get_n(): + if n > 8: + x, y, xm, ym = self.generate_xy_sample(n) + r = stats.normaltest(x) + rm = stats.mstats.normaltest(xm) + assert_allclose(np.asarray(r), np.asarray(rm)) + + def test_find_repeats(self): + x = np.asarray([1, 1, 2, 2, 3, 3, 3, 4, 4, 4, 4]).astype('float') + tmp = np.asarray([1, 1, 2, 2, 3, 3, 3, 4, 4, 4, 4, 5, 5, 5, 5]).astype('float') + mask = (tmp == 5.) + xm = np.ma.array(tmp, mask=mask) + x_orig, xm_orig = x.copy(), xm.copy() + + unique, unique_counts = np.unique(x, return_counts=True) + r = unique[unique_counts > 1], unique_counts[unique_counts > 1] + rm = stats.mstats.find_repeats(xm) + + assert_equal(r, rm) + assert_equal(x, x_orig) + assert_equal(xm, xm_orig) + + # This crazy behavior is expected by count_tied_groups, but is not + # in the docstring... + _, counts = stats.mstats.find_repeats([]) + assert_equal(counts, np.array(0, dtype=np.intp)) + + def test_kendalltau(self): + for n in self.get_n(): + x, y, xm, ym = self.generate_xy_sample(n) + r = stats.kendalltau(x, y) + rm = stats.mstats.kendalltau(xm, ym) + assert_almost_equal(r[0], rm[0], decimal=10) + assert_almost_equal(r[1], rm[1], decimal=7) + + def test_obrientransform(self): + for n in self.get_n(): + x, y, xm, ym = self.generate_xy_sample(n) + r = np.stack(stats.obrientransform(x)) + rm = stats.mstats.obrientransform(xm) + assert_almost_equal(r.T, rm[0:len(x)]) + + def test_ks_1samp(self): + """Checks that mstats.ks_1samp and stats.ks_1samp agree on masked arrays.""" + for mode in ['auto', 'exact', 'asymp']: + with warnings.catch_warnings(): + for alternative in ['less', 'greater', 'two-sided']: + for n in self.get_n(): + x, y, xm, ym = self.generate_xy_sample(n) + res1 = stats.ks_1samp(x, stats.norm.cdf, + alternative=alternative, mode=mode) + res2 = stats.mstats.ks_1samp(xm, stats.norm.cdf, + alternative=alternative, mode=mode) + assert_equal(np.asarray(res1), np.asarray(res2)) + res3 = stats.ks_1samp(xm, stats.norm.cdf, + alternative=alternative, mode=mode) + assert_equal(np.asarray(res1), np.asarray(res3)) + + def test_kstest_1samp(self): + """ + Checks that 1-sample mstats.kstest and stats.kstest agree on masked arrays. + """ + for mode in ['auto', 'exact', 'asymp']: + with warnings.catch_warnings(): + for alternative in ['less', 'greater', 'two-sided']: + for n in self.get_n(): + x, y, xm, ym = self.generate_xy_sample(n) + res1 = stats.kstest(x, 'norm', + alternative=alternative, mode=mode) + res2 = stats.mstats.kstest(xm, 'norm', + alternative=alternative, mode=mode) + assert_equal(np.asarray(res1), np.asarray(res2)) + res3 = stats.kstest(xm, 'norm', + alternative=alternative, mode=mode) + assert_equal(np.asarray(res1), np.asarray(res3)) + + def test_ks_2samp(self): + """Checks that mstats.ks_2samp and stats.ks_2samp agree on masked arrays. + gh-8431""" + for mode in ['auto', 'exact', 'asymp']: + with warnings.catch_warnings(): + if mode in ['auto', 'exact']: + message = "ks_2samp: Exact calculation unsuccessful." + warnings.filterwarnings("ignore", message, RuntimeWarning) + for alternative in ['less', 'greater', 'two-sided']: + for n in self.get_n(): + x, y, xm, ym = self.generate_xy_sample(n) + res1 = stats.ks_2samp(x, y, + alternative=alternative, mode=mode) + res2 = stats.mstats.ks_2samp(xm, ym, + alternative=alternative, mode=mode) + assert_equal(np.asarray(res1), np.asarray(res2)) + res3 = stats.ks_2samp(xm, y, + alternative=alternative, mode=mode) + assert_equal(np.asarray(res1), np.asarray(res3)) + + def test_kstest_2samp(self): + """ + Checks that 2-sample mstats.kstest and stats.kstest agree on masked arrays. + """ + for mode in ['auto', 'exact', 'asymp']: + with warnings.catch_warnings(): + if mode in ['auto', 'exact']: + message = "ks_2samp: Exact calculation unsuccessful." + warnings.filterwarnings("ignore", message, RuntimeWarning) + for alternative in ['less', 'greater', 'two-sided']: + for n in self.get_n(): + x, y, xm, ym = self.generate_xy_sample(n) + res1 = stats.kstest(x, y, + alternative=alternative, mode=mode) + res2 = stats.mstats.kstest(xm, ym, + alternative=alternative, mode=mode) + assert_equal(np.asarray(res1), np.asarray(res2)) + res3 = stats.kstest(xm, y, + alternative=alternative, mode=mode) + assert_equal(np.asarray(res1), np.asarray(res3)) + + +class TestBrunnerMunzel: + # Data from (Lumley, 1996) + X = np.ma.masked_invalid([1, 2, 1, 1, 1, np.nan, 1, 1, + 1, 1, 1, 2, 4, 1, 1, np.nan]) + Y = np.ma.masked_invalid([3, 3, 4, 3, np.nan, 1, 2, 3, 1, 1, 5, 4]) + significant = 14 + + def test_brunnermunzel_one_sided(self): + # Results are compared with R's lawstat package. + u1, p1 = mstats.brunnermunzel(self.X, self.Y, alternative='less') + u2, p2 = mstats.brunnermunzel(self.Y, self.X, alternative='greater') + u3, p3 = mstats.brunnermunzel(self.X, self.Y, alternative='greater') + u4, p4 = mstats.brunnermunzel(self.Y, self.X, alternative='less') + + assert_almost_equal(p1, p2, decimal=self.significant) + assert_almost_equal(p3, p4, decimal=self.significant) + assert_(p1 != p3) + assert_almost_equal(u1, 3.1374674823029505, + decimal=self.significant) + assert_almost_equal(u2, -3.1374674823029505, + decimal=self.significant) + assert_almost_equal(u3, 3.1374674823029505, + decimal=self.significant) + assert_almost_equal(u4, -3.1374674823029505, + decimal=self.significant) + assert_almost_equal(p1, 0.0028931043330757342, + decimal=self.significant) + assert_almost_equal(p3, 0.99710689566692423, + decimal=self.significant) + + def test_brunnermunzel_two_sided(self): + # Results are compared with R's lawstat package. + u1, p1 = mstats.brunnermunzel(self.X, self.Y, alternative='two-sided') + u2, p2 = mstats.brunnermunzel(self.Y, self.X, alternative='two-sided') + + assert_almost_equal(p1, p2, decimal=self.significant) + assert_almost_equal(u1, 3.1374674823029505, + decimal=self.significant) + assert_almost_equal(u2, -3.1374674823029505, + decimal=self.significant) + assert_almost_equal(p1, 0.0057862086661515377, + decimal=self.significant) + + def test_brunnermunzel_default(self): + # The default value for alternative is two-sided + u1, p1 = mstats.brunnermunzel(self.X, self.Y) + u2, p2 = mstats.brunnermunzel(self.Y, self.X) + + assert_almost_equal(p1, p2, decimal=self.significant) + assert_almost_equal(u1, 3.1374674823029505, + decimal=self.significant) + assert_almost_equal(u2, -3.1374674823029505, + decimal=self.significant) + assert_almost_equal(p1, 0.0057862086661515377, + decimal=self.significant) + + def test_brunnermunzel_alternative_error(self): + alternative = "error" + distribution = "t" + assert_(alternative not in ["two-sided", "greater", "less"]) + assert_raises(ValueError, + mstats.brunnermunzel, + self.X, + self.Y, + alternative, + distribution) + + def test_brunnermunzel_distribution_norm(self): + u1, p1 = mstats.brunnermunzel(self.X, self.Y, distribution="normal") + u2, p2 = mstats.brunnermunzel(self.Y, self.X, distribution="normal") + assert_almost_equal(p1, p2, decimal=self.significant) + assert_almost_equal(u1, 3.1374674823029505, + decimal=self.significant) + assert_almost_equal(u2, -3.1374674823029505, + decimal=self.significant) + assert_almost_equal(p1, 0.0017041417600383024, + decimal=self.significant) + + def test_brunnermunzel_distribution_error(self): + alternative = "two-sided" + distribution = "error" + assert_(alternative not in ["t", "normal"]) + assert_raises(ValueError, + mstats.brunnermunzel, + self.X, + self.Y, + alternative, + distribution) + + def test_brunnermunzel_empty_imput(self): + u1, p1 = mstats.brunnermunzel(self.X, []) + u2, p2 = mstats.brunnermunzel([], self.Y) + u3, p3 = mstats.brunnermunzel([], []) + + assert_(np.isnan(u1)) + assert_(np.isnan(p1)) + assert_(np.isnan(u2)) + assert_(np.isnan(p2)) + assert_(np.isnan(u3)) + assert_(np.isnan(p3)) diff --git a/v12x/pylib/scipy/stats/tests/test_mstats_extras.py b/v12x/pylib/scipy/stats/tests/test_mstats_extras.py new file mode 100644 index 0000000000000000000000000000000000000000..8b14455b417317a23abda8efe5297e8f4839ffd2 --- /dev/null +++ b/v12x/pylib/scipy/stats/tests/test_mstats_extras.py @@ -0,0 +1,172 @@ +import numpy as np +import numpy.ma as ma +import scipy.stats.mstats as ms + +from numpy.testing import (assert_equal, assert_almost_equal, assert_, + assert_allclose) + + +def test_compare_medians_ms(): + x = np.arange(7) + y = x + 10 + assert_almost_equal(ms.compare_medians_ms(x, y), 0) + + y2 = np.linspace(0, 1, num=10) + assert_almost_equal(ms.compare_medians_ms(x, y2), 0.017116406778) + + +def test_hdmedian(): + # 1-D array + x = ma.arange(11) + assert_allclose(ms.hdmedian(x), 5, rtol=1e-14) + x.mask = ma.make_mask(x) + x.mask[:7] = False + assert_allclose(ms.hdmedian(x), 3, rtol=1e-14) + + # Check that `var` keyword returns a value. TODO: check whether returned + # value is actually correct. + assert_(ms.hdmedian(x, var=True).size == 2) + + # 2-D array + x2 = ma.arange(22).reshape((11, 2)) + assert_allclose(ms.hdmedian(x2, axis=0), [10, 11]) + x2.mask = ma.make_mask(x2) + x2.mask[:7, :] = False + assert_allclose(ms.hdmedian(x2, axis=0), [6, 7]) + + +def test_rsh(): + rng = np.random.default_rng(806795795) + x = rng.standard_normal(100) + res = ms.rsh(x) + # Just a sanity check that the code runs and output shape is correct. + # TODO: check that implementation is correct. + assert_(res.shape == x.shape) + + # Check points keyword + res = ms.rsh(x, points=[0, 1.]) + assert_(res.size == 2) + + +def test_mjci(): + # Tests the Marits-Jarrett estimator + data = ma.array([77, 87, 88,114,151,210,219,246,253,262, + 296,299,306,376,428,515,666,1310,2611]) + assert_almost_equal(ms.mjci(data),[55.76819,45.84028,198.87875],5) + + +def test_trimmed_mean_ci(): + # Tests the confidence intervals of the trimmed mean. + data = ma.array([545,555,558,572,575,576,578,580, + 594,605,635,651,653,661,666]) + assert_almost_equal(ms.trimmed_mean(data,0.2), 596.2, 1) + assert_equal(np.round(ms.trimmed_mean_ci(data,(0.2,0.2)),1), + [561.8, 630.6]) + + +def test_idealfourths(): + # Tests ideal-fourths + test = np.arange(100) + assert_almost_equal(np.asarray(ms.idealfourths(test)), + [24.416667,74.583333],6) + test_2D = test.repeat(3).reshape(-1,3) + assert_almost_equal(ms.idealfourths(test_2D, axis=0), + [[24.416667,24.416667,24.416667], + [74.583333,74.583333,74.583333]],6) + assert_almost_equal(ms.idealfourths(test_2D, axis=1), + test.repeat(2).reshape(-1,2)) + test = [0, 0] + _result = ms.idealfourths(test) + assert_(np.isnan(_result).all()) + + +class TestQuantiles: + data = [0.706560797,0.727229578,0.990399276,0.927065621,0.158953014, + 0.887764025,0.239407086,0.349638551,0.972791145,0.149789972, + 0.936947700,0.132359948,0.046041972,0.641675031,0.945530547, + 0.224218684,0.771450991,0.820257774,0.336458052,0.589113496, + 0.509736129,0.696838829,0.491323573,0.622767425,0.775189248, + 0.641461450,0.118455200,0.773029450,0.319280007,0.752229111, + 0.047841438,0.466295911,0.583850781,0.840581845,0.550086491, + 0.466470062,0.504765074,0.226855960,0.362641207,0.891620942, + 0.127898691,0.490094097,0.044882048,0.041441695,0.317976349, + 0.504135618,0.567353033,0.434617473,0.636243375,0.231803616, + 0.230154113,0.160011327,0.819464108,0.854706985,0.438809221, + 0.487427267,0.786907310,0.408367937,0.405534192,0.250444460, + 0.995309248,0.144389588,0.739947527,0.953543606,0.680051621, + 0.388382017,0.863530727,0.006514031,0.118007779,0.924024803, + 0.384236354,0.893687694,0.626534881,0.473051932,0.750134705, + 0.241843555,0.432947602,0.689538104,0.136934797,0.150206859, + 0.474335206,0.907775349,0.525869295,0.189184225,0.854284286, + 0.831089744,0.251637345,0.587038213,0.254475554,0.237781276, + 0.827928620,0.480283781,0.594514455,0.213641488,0.024194386, + 0.536668589,0.699497811,0.892804071,0.093835427,0.731107772] + + def test_hdquantiles(self): + data = self.data + assert_almost_equal(ms.hdquantiles(data,[0., 1.]), + [0.006514031, 0.995309248]) + hdq = ms.hdquantiles(data,[0.25, 0.5, 0.75]) + assert_almost_equal(hdq, [0.253210762, 0.512847491, 0.762232442,]) + + data = np.array(data).reshape(10,10) + hdq = ms.hdquantiles(data,[0.25,0.5,0.75],axis=0) + assert_almost_equal(hdq[:,0], ms.hdquantiles(data[:,0],[0.25,0.5,0.75])) + assert_almost_equal(hdq[:,-1], ms.hdquantiles(data[:,-1],[0.25,0.5,0.75])) + hdq = ms.hdquantiles(data,[0.25,0.5,0.75],axis=0,var=True) + assert_almost_equal(hdq[...,0], + ms.hdquantiles(data[:,0],[0.25,0.5,0.75],var=True)) + assert_almost_equal(hdq[...,-1], + ms.hdquantiles(data[:,-1],[0.25,0.5,0.75], var=True)) + + def test_hdquantiles_sd(self): + # Standard deviation is a jackknife estimator, so we can check if + # the efficient version (hdquantiles_sd) matches a rudimentary, + # but clear version here. + + hd_std_errs = ms.hdquantiles_sd(self.data) + + # jacknnife standard error, Introduction to the Bootstrap Eq. 11.5 + n = len(self.data) + jdata = np.broadcast_to(self.data, (n, n)) + jselector = np.logical_not(np.eye(n)) # leave out one sample each row + jdata = jdata[jselector].reshape(n, n-1) + jdist = ms.hdquantiles(jdata, axis=1) + jdist_mean = np.mean(jdist, axis=0) + jstd = ((n-1)/n * np.sum((jdist - jdist_mean)**2, axis=0))**.5 + + assert_almost_equal(hd_std_errs, jstd) + # Test actual values for good measure + assert_almost_equal(hd_std_errs, [0.0379258, 0.0380656, 0.0380013]) + + two_data_points = ms.hdquantiles_sd([1, 2]) + assert_almost_equal(two_data_points, [0.5, 0.5, 0.5]) + + def test_mquantiles_cimj(self): + # Only test that code runs, implementation not checked for correctness + ci_lower, ci_upper = ms.mquantiles_cimj(self.data) + assert_(ci_lower.size == ci_upper.size == 3) + + +def test_median_cihs(): + # Basic test against R library EnvStats function `eqnpar`, e.g. + # library(EnvStats) + # options(digits=8) + # x = c(0.88612955, 0.35242375, 0.66240904, 0.94617974, 0.10929913, + # 0.76699506, 0.88550655, 0.62763754, 0.76818588, 0.68506508, + # 0.88043148, 0.03911248, 0.93805564, 0.95326961, 0.25291112, + # 0.16128487, 0.49784577, 0.24588924, 0.6597, 0.92239679) + # eqnpar(x, p=0.5, + # ci.method = "interpolate", approx.conf.level = 0.95, ci = TRUE) + rng = np.random.default_rng(8824288259505800535) + x = rng.random(size=20) + assert_allclose(ms.median_cihs(x), (0.38663198, 0.88431272)) + + # SciPy's 90% CI upper limit doesn't match that of EnvStats eqnpar. SciPy + # doesn't look wrong, and it agrees with a different reference, + # `median_confint_hs` from `hoehleatsu/quantileCI`. + # In (e.g.) Colab with R runtime: + # devtools::install_github("hoehleatsu/quantileCI") + # library(quantileCI) + # median_confint_hs(x=x, conf.level=0.90, interpolate=TRUE) + assert_allclose(ms.median_cihs(x, 0.1), (0.48319773366, 0.88094268050)) diff --git a/v12x/pylib/scipy/stats/tests/test_multicomp.py b/v12x/pylib/scipy/stats/tests/test_multicomp.py new file mode 100644 index 0000000000000000000000000000000000000000..2860b142c38a339f96c50af92bd452c6de95e84d --- /dev/null +++ b/v12x/pylib/scipy/stats/tests/test_multicomp.py @@ -0,0 +1,405 @@ +import copy + +import numpy as np +import pytest +from numpy.testing import assert_allclose + +from scipy import stats +from scipy.stats._multicomp import _pvalue_dunnett, DunnettResult + + +class TestDunnett: + # For the following tests, p-values were computed using Matlab, e.g. + # sample = [18. 15. 18. 16. 17. 15. 14. 14. 14. 15. 15.... + # 14. 15. 14. 22. 18. 21. 21. 10. 10. 11. 9.... + # 25. 26. 17.5 16. 15.5 14.5 22. 22. 24. 22.5 29.... + # 24.5 20. 18. 18.5 17.5 26.5 13. 16.5 13. 13. 13.... + # 28. 27. 34. 31. 29. 27. 24. 23. 38. 36. 25.... + # 38. 26. 22. 36. 27. 27. 32. 28. 31.... + # 24. 27. 33. 32. 28. 19. 37. 31. 36. 36.... + # 34. 38. 32. 38. 32.... + # 26. 24. 26. 25. 29. 29.5 16.5 36. 44.... + # 25. 27. 19.... + # 25. 20.... + # 28.]; + # j = [0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 ... + # 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 ... + # 0 0 0 0... + # 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1... + # 2 2 2 2 2 2 2 2 2... + # 3 3 3... + # 4 4... + # 5]; + # [~, ~, stats] = anova1(sample, j, "off"); + # [results, ~, ~, gnames] = multcompare(stats, ... + # "CriticalValueType", "dunnett", ... + # "Approximate", false); + # tbl = array2table(results, "VariableNames", ... + # ["Group", "Control Group", "Lower Limit", ... + # "Difference", "Upper Limit", "P-value"]); + # tbl.("Group") = gnames(tbl.("Group")); + # tbl.("Control Group") = gnames(tbl.("Control Group")) + + # Matlab doesn't report the statistic, so the statistics were + # computed using R multcomp `glht`, e.g.: + # library(multcomp) + # options(digits=16) + # control < - c(18.0, 15.0, 18.0, 16.0, 17.0, 15.0, 14.0, 14.0, 14.0, + # 15.0, 15.0, 14.0, 15.0, 14.0, 22.0, 18.0, 21.0, 21.0, + # 10.0, 10.0, 11.0, 9.0, 25.0, 26.0, 17.5, 16.0, 15.5, + # 14.5, 22.0, 22.0, 24.0, 22.5, 29.0, 24.5, 20.0, 18.0, + # 18.5, 17.5, 26.5, 13.0, 16.5, 13.0, 13.0, 13.0, 28.0, + # 27.0, 34.0, 31.0, 29.0, 27.0, 24.0, 23.0, 38.0, 36.0, + # 25.0, 38.0, 26.0, 22.0, 36.0, 27.0, 27.0, 32.0, 28.0, + # 31.0) + # t < - c(24.0, 27.0, 33.0, 32.0, 28.0, 19.0, 37.0, 31.0, 36.0, 36.0, + # 34.0, 38.0, 32.0, 38.0, 32.0) + # w < - c(26.0, 24.0, 26.0, 25.0, 29.0, 29.5, 16.5, 36.0, 44.0) + # x < - c(25.0, 27.0, 19.0) + # y < - c(25.0, 20.0) + # z < - c(28.0) + # + # groups = factor(rep(c("control", "t", "w", "x", "y", "z"), + # times=c(length(control), length(t), length(w), + # length(x), length(y), length(z)))) + # df < - data.frame(response=c(control, t, w, x, y, z), + # group=groups) + # model < - aov(response + # ~group, data = df) + # test < - glht(model=model, + # linfct=mcp(group="Dunnett"), + # alternative="g") + # summary(test) + # confint(test) + # p-values agreed with those produced by Matlab to at least atol=1e-3 + + # From Matlab's documentation on multcompare + samples_1 = [ + [ + 24.0, 27.0, 33.0, 32.0, 28.0, 19.0, 37.0, 31.0, 36.0, 36.0, + 34.0, 38.0, 32.0, 38.0, 32.0 + ], + [26.0, 24.0, 26.0, 25.0, 29.0, 29.5, 16.5, 36.0, 44.0], + [25.0, 27.0, 19.0], + [25.0, 20.0], + [28.0] + ] + control_1 = [ + 18.0, 15.0, 18.0, 16.0, 17.0, 15.0, 14.0, 14.0, 14.0, 15.0, 15.0, + 14.0, 15.0, 14.0, 22.0, 18.0, 21.0, 21.0, 10.0, 10.0, 11.0, 9.0, + 25.0, 26.0, 17.5, 16.0, 15.5, 14.5, 22.0, 22.0, 24.0, 22.5, 29.0, + 24.5, 20.0, 18.0, 18.5, 17.5, 26.5, 13.0, 16.5, 13.0, 13.0, 13.0, + 28.0, 27.0, 34.0, 31.0, 29.0, 27.0, 24.0, 23.0, 38.0, 36.0, 25.0, + 38.0, 26.0, 22.0, 36.0, 27.0, 27.0, 32.0, 28.0, 31.0 + ] + pvalue_1 = [4.727e-06, 0.022346, 0.97912, 0.99953, 0.86579] # Matlab + # Statistic, alternative p-values, and CIs computed with R multcomp `glht` + p_1_twosided = [1e-4, 0.02237, 0.97913, 0.99953, 0.86583] + p_1_greater = [1e-4, 0.011217, 0.768500, 0.896991, 0.577211] + p_1_less = [1, 1, 0.99660, 0.98398, .99953] + statistic_1 = [5.27356, 2.91270, 0.60831, 0.27002, 0.96637] + ci_1_twosided = [[5.3633917835622, 0.7296142201217, -8.3879817106607, + -11.9090753452911, -11.7655021543469], + [15.9709832164378, 13.8936496687672, 13.4556900439941, + 14.6434503452911, 25.4998771543469]] + ci_1_greater = [5.9036402398526, 1.4000632918725, -7.2754756323636, + -10.5567456382391, -9.8675629499576] + ci_1_less = [15.4306165948619, 13.2230539537359, 12.3429406339544, + 13.2908248513211, 23.6015228251660] + pvalues_1 = dict(twosided=p_1_twosided, less=p_1_less, greater=p_1_greater) + cis_1 = dict(twosided=ci_1_twosided, less=ci_1_less, greater=ci_1_greater) + case_1 = dict(samples=samples_1, control=control_1, statistic=statistic_1, + pvalues=pvalues_1, cis=cis_1) + + # From Dunnett1955 comparing with R's DescTools: DunnettTest + samples_2 = [[9.76, 8.80, 7.68, 9.36], [12.80, 9.68, 12.16, 9.20, 10.55]] + control_2 = [7.40, 8.50, 7.20, 8.24, 9.84, 8.32] + pvalue_2 = [0.6201, 0.0058] + # Statistic, alternative p-values, and CIs computed with R multcomp `glht` + p_2_twosided = [0.6201020, 0.0058254] + p_2_greater = [0.3249776, 0.0029139] + p_2_less = [0.91676, 0.99984] + statistic_2 = [0.85703, 3.69375] + ci_2_twosided = [[-1.2564116462124, 0.8396273539789], + [2.5564116462124, 4.4163726460211]] + ci_2_greater = [-0.9588591188156, 1.1187563667543] + ci_2_less = [2.2588591188156, 4.1372436332457] + pvalues_2 = dict(twosided=p_2_twosided, less=p_2_less, greater=p_2_greater) + cis_2 = dict(twosided=ci_2_twosided, less=ci_2_less, greater=ci_2_greater) + case_2 = dict(samples=samples_2, control=control_2, statistic=statistic_2, + pvalues=pvalues_2, cis=cis_2) + + samples_3 = [[55, 64, 64], [55, 49, 52], [50, 44, 41]] + control_3 = [55, 47, 48] + pvalue_3 = [0.0364, 0.8966, 0.4091] + # Statistic, alternative p-values, and CIs computed with R multcomp `glht` + p_3_twosided = [0.036407, 0.896539, 0.409295] + p_3_greater = [0.018277, 0.521109, 0.981892] + p_3_less = [0.99944, 0.90054, 0.20974] + statistic_3 = [3.09073, 0.56195, -1.40488] + ci_3_twosided = [[0.7529028025053, -8.2470971974947, -15.2470971974947], + [21.2470971974947, 12.2470971974947, 5.2470971974947]] + ci_3_greater = [2.4023682323149, -6.5976317676851, -13.5976317676851] + ci_3_less = [19.5984402363662, 10.5984402363662, 3.5984402363662] + pvalues_3 = dict(twosided=p_3_twosided, less=p_3_less, greater=p_3_greater) + cis_3 = dict(twosided=ci_3_twosided, less=ci_3_less, greater=ci_3_greater) + case_3 = dict(samples=samples_3, control=control_3, statistic=statistic_3, + pvalues=pvalues_3, cis=cis_3) + + # From Thomson and Short, + # Mucociliary function in health, chronic obstructive airway disease, + # and asbestosis, Journal of Applied Physiology, 1969. Table 1 + # Comparing with R's DescTools: DunnettTest + samples_4 = [[3.8, 2.7, 4.0, 2.4], [2.8, 3.4, 3.7, 2.2, 2.0]] + control_4 = [2.9, 3.0, 2.5, 2.6, 3.2] + pvalue_4 = [0.5832, 0.9982] + # Statistic, alternative p-values, and CIs computed with R multcomp `glht` + p_4_twosided = [0.58317, 0.99819] + p_4_greater = [0.30225, 0.69115] + p_4_less = [0.91929, 0.65212] + statistic_4 = [0.90875, -0.05007] + ci_4_twosided = [[-0.6898153448579, -1.0333456251632], + [1.4598153448579, 0.9933456251632]] + ci_4_greater = [-0.5186459268412, -0.8719655502147 ] + ci_4_less = [1.2886459268412, 0.8319655502147] + pvalues_4 = dict(twosided=p_4_twosided, less=p_4_less, greater=p_4_greater) + cis_4 = dict(twosided=ci_4_twosided, less=ci_4_less, greater=ci_4_greater) + case_4 = dict(samples=samples_4, control=control_4, statistic=statistic_4, + pvalues=pvalues_4, cis=cis_4) + + @pytest.mark.parametrize( + 'rho, n_groups, df, statistic, pvalue, alternative', + [ + # From Dunnett1955 + # Tables 1a and 1b pages 1117-1118 + (0.5, 1, 10, 1.81, 0.05, "greater"), # different than two-sided + (0.5, 3, 10, 2.34, 0.05, "greater"), + (0.5, 2, 30, 1.99, 0.05, "greater"), + (0.5, 5, 30, 2.33, 0.05, "greater"), + (0.5, 4, 12, 3.32, 0.01, "greater"), + (0.5, 7, 12, 3.56, 0.01, "greater"), + (0.5, 2, 60, 2.64, 0.01, "greater"), + (0.5, 4, 60, 2.87, 0.01, "greater"), + (0.5, 4, 60, [2.87, 2.21], [0.01, 0.05], "greater"), + # Tables 2a and 2b pages 1119-1120 + (0.5, 1, 10, 2.23, 0.05, "two-sided"), # two-sided + (0.5, 3, 10, 2.81, 0.05, "two-sided"), + (0.5, 2, 30, 2.32, 0.05, "two-sided"), + (0.5, 3, 20, 2.57, 0.05, "two-sided"), + (0.5, 4, 12, 3.76, 0.01, "two-sided"), + (0.5, 7, 12, 4.08, 0.01, "two-sided"), + (0.5, 2, 60, 2.90, 0.01, "two-sided"), + (0.5, 4, 60, 3.14, 0.01, "two-sided"), + (0.5, 4, 60, [3.14, 2.55], [0.01, 0.05], "two-sided"), + ], + ) + def test_critical_values( + self, rho, n_groups, df, statistic, pvalue, alternative + ): + rng = np.random.default_rng(165250594791731684851746311027739134893) + rho = np.full((n_groups, n_groups), rho) + np.fill_diagonal(rho, 1) + + statistic = np.array(statistic) + res = _pvalue_dunnett( + rho=rho, df=df, statistic=statistic, + alternative=alternative, + rng=rng + ) + assert_allclose(res, pvalue, atol=5e-3) + + @pytest.mark.parametrize( + 'samples, control, pvalue, statistic', + [ + (samples_1, control_1, pvalue_1, statistic_1), + (samples_2, control_2, pvalue_2, statistic_2), + (samples_3, control_3, pvalue_3, statistic_3), + (samples_4, control_4, pvalue_4, statistic_4), + ] + ) + def test_basic(self, samples, control, pvalue, statistic): + rng = np.random.default_rng(11681140010308601919115036826969764808) + + res = stats.dunnett(*samples, control=control, rng=rng) + + assert isinstance(res, DunnettResult) + assert_allclose(res.statistic, statistic, rtol=5e-5) + assert_allclose(res.pvalue, pvalue, rtol=1e-2, atol=1e-4) + + @pytest.mark.parametrize( + 'alternative', + ['two-sided', 'less', 'greater'] + ) + def test_ttest_ind(self, alternative): + # check that `dunnett` agrees with `ttest_ind` + # when there are only two groups + rng = np.random.default_rng(114184017807316971636137493526995620351) + + for _ in range(10): + sample = rng.integers(-100, 100, size=(10,)) + control = rng.integers(-100, 100, size=(10,)) + + # preserve use of old random_state during SPEC 7 transition + res = stats.dunnett( + sample, control=control, + alternative=alternative, random_state=rng + ) + ref = stats.ttest_ind( + sample, control, + alternative=alternative + ) + + assert_allclose(res.statistic, ref.statistic, rtol=1e-3, atol=1e-5) + assert_allclose(res.pvalue, ref.pvalue, rtol=1e-3, atol=1e-5) + + @pytest.mark.parametrize( + 'alternative, pvalue', + [ + ('less', [0, 1]), + ('greater', [1, 0]), + ('two-sided', [0, 0]), + ] + ) + def test_alternatives(self, alternative, pvalue): + rng = np.random.default_rng(114184017807316971636137493526995620351) + + # width of 20 and min diff between samples/control is 60 + # and maximal diff would be 100 + sample_less = rng.integers(0, 20, size=(10,)) + control = rng.integers(80, 100, size=(10,)) + sample_greater = rng.integers(160, 180, size=(10,)) + + res = stats.dunnett( + sample_less, sample_greater, control=control, + alternative=alternative, rng=rng + ) + assert_allclose(res.pvalue, pvalue, atol=1e-7) + + ci = res.confidence_interval() + # two-sided is comparable for high/low + if alternative == 'less': + assert np.isneginf(ci.low).all() + assert -100 < ci.high[0] < -60 + assert 60 < ci.high[1] < 100 + elif alternative == 'greater': + assert -100 < ci.low[0] < -60 + assert 60 < ci.low[1] < 100 + assert np.isposinf(ci.high).all() + elif alternative == 'two-sided': + assert -100 < ci.low[0] < -60 + assert 60 < ci.low[1] < 100 + assert -100 < ci.high[0] < -60 + assert 60 < ci.high[1] < 100 + + @pytest.mark.parametrize("case", [case_1, case_2, case_3, case_4]) + @pytest.mark.parametrize("alternative", ['less', 'greater', 'two-sided']) + def test_against_R_multicomp_glht(self, case, alternative): + rng = np.random.default_rng(189117774084579816190295271136455278291) + samples = case['samples'] + control = case['control'] + alternatives = {'less': 'less', 'greater': 'greater', + 'two-sided': 'twosided'} + p_ref = case['pvalues'][alternative.replace('-', '')] + + res = stats.dunnett(*samples, control=control, alternative=alternative, + rng=rng) + # atol can't be tighter because R reports some pvalues as "< 1e-4" + assert_allclose(res.pvalue, p_ref, rtol=5e-3, atol=1e-4) + + ci_ref = case['cis'][alternatives[alternative]] + if alternative == "greater": + ci_ref = [ci_ref, np.inf] + elif alternative == "less": + ci_ref = [-np.inf, ci_ref] + assert res._ci is None + assert res._ci_cl is None + ci = res.confidence_interval(confidence_level=0.95) + assert_allclose(ci.low, ci_ref[0], rtol=5e-3, atol=1e-5) + assert_allclose(ci.high, ci_ref[1], rtol=5e-3, atol=1e-5) + + # re-run to use the cached value "is" to check id as same object + assert res._ci is ci + assert res._ci_cl == 0.95 + ci_ = res.confidence_interval(confidence_level=0.95) + assert ci_ is ci + + @pytest.mark.parametrize('alternative', ["two-sided", "less", "greater"]) + def test_str(self, alternative): + rng = np.random.default_rng(189117774084579816190295271136455278291) + + res = stats.dunnett( + *self.samples_3, control=self.control_3, alternative=alternative, + rng=rng + ) + + # check some str output + res_str = str(res) + assert '(Sample 2 - Control)' in res_str + assert '95.0%' in res_str + + if alternative == 'less': + assert '-inf' in res_str + assert '19.' in res_str + elif alternative == 'greater': + assert 'inf' in res_str + assert '-13.' in res_str + else: + assert 'inf' not in res_str + assert '21.' in res_str + + def test_warnings(self): + rng = np.random.default_rng(189117774084579816190295271136455278291) + + res = stats.dunnett( + *self.samples_3, control=self.control_3, rng=rng + ) + msg = r"Computation of the confidence interval did not converge" + with pytest.warns(UserWarning, match=msg): + res._allowance(tol=1e-5) + + def test_raises(self): + samples, control = self.samples_3, self.control_3 + + # alternative + with pytest.raises(ValueError, match="alternative must be"): + stats.dunnett(*samples, control=control, alternative='bob') + + # 2D for a sample + samples_ = copy.deepcopy(samples) + samples_[0] = [samples_[0]] + with pytest.raises(ValueError, match="must be 1D arrays"): + stats.dunnett(*samples_, control=control) + + # 2D for control + control_ = copy.deepcopy(control) + control_ = [control_] + with pytest.raises(ValueError, match="must be 1D arrays"): + stats.dunnett(*samples, control=control_) + + # No obs in a sample + samples_ = copy.deepcopy(samples) + samples_[1] = [] + with pytest.raises(ValueError, match="at least 1 observation"): + stats.dunnett(*samples_, control=control) + + # No obs in control + control_ = [] + with pytest.raises(ValueError, match="at least 1 observation"): + stats.dunnett(*samples, control=control_) + + res = stats.dunnett(*samples, control=control) + with pytest.raises(ValueError, match="Confidence level must"): + res.confidence_interval(confidence_level=3) + + @pytest.mark.filterwarnings("ignore:Computation of the confidence") + @pytest.mark.parametrize('n_samples', [1, 2, 3]) + def test_shapes(self, n_samples): + rng = np.random.default_rng(689448934110805334) + samples = rng.normal(size=(n_samples, 10)) + control = rng.normal(size=10) + res = stats.dunnett(*samples, control=control, rng=rng) + assert res.statistic.shape == (n_samples,) + assert res.pvalue.shape == (n_samples,) + ci = res.confidence_interval() + assert ci.low.shape == (n_samples,) + assert ci.high.shape == (n_samples,) diff --git a/v12x/pylib/scipy/stats/tests/test_multivariate.py b/v12x/pylib/scipy/stats/tests/test_multivariate.py new file mode 100644 index 0000000000000000000000000000000000000000..46690b575a0c56c250d624bad87272f5a1fabd2e --- /dev/null +++ b/v12x/pylib/scipy/stats/tests/test_multivariate.py @@ -0,0 +1,5052 @@ +""" +Test functions for multivariate normal, t, and related distributions. + +""" +import pickle +from dataclasses import dataclass + +from numpy.testing import (assert_allclose, assert_almost_equal, + assert_array_almost_equal, assert_equal, + assert_array_less, assert_) +import pytest +from pytest import raises as assert_raises + +from .test_continuous_basic import check_distribution_rvs + +import numpy as np + +import scipy.linalg + +from scipy.stats._multivariate import (_PSD, + _lnB, + multivariate_normal_frozen) +from scipy.stats import (multivariate_normal, multivariate_hypergeom, + matrix_normal, special_ortho_group, ortho_group, + random_correlation, unitary_group, dirichlet, + beta, wishart, multinomial, invwishart, chi2, + invgamma, norm, uniform, ks_2samp, kstest, binom, + hypergeom, multivariate_t, cauchy, normaltest, + random_table, uniform_direction, vonmises_fisher, + dirichlet_multinomial, vonmises, matrix_t) + +from scipy.stats import _covariance, Covariance +from scipy.stats._continuous_distns import _norm_pdf as norm_pdf +from scipy import stats + +from scipy.integrate import tanhsinh, cubature, quad +from scipy.integrate import romb, qmc_quad, dblquad, tplquad +from scipy.special import multigammaln +import scipy.special as special + +from .common_tests import check_random_state_property +from .data._mvt import _qsimvtv + +from unittest.mock import patch + + +def assert_close(res, ref, *args, **kwargs): + res, ref = np.asarray(res), np.asarray(ref) + assert_allclose(res, ref, *args, **kwargs) + assert_equal(res.shape, ref.shape) + + +class TestCovariance: + + def test_input_validation(self): + + message = "The input `precision` must be a square, two-dimensional..." + with pytest.raises(ValueError, match=message): + _covariance.CovViaPrecision(np.ones(2)) + + message = "`precision.shape` must equal `covariance.shape`." + with pytest.raises(ValueError, match=message): + _covariance.CovViaPrecision(np.eye(3), covariance=np.eye(2)) + + message = "The input `diagonal` must be a one-dimensional array..." + with pytest.raises(ValueError, match=message): + _covariance.CovViaDiagonal("alpaca") + + message = "The input `cholesky` must be a square, two-dimensional..." + with pytest.raises(ValueError, match=message): + _covariance.CovViaCholesky(np.ones(2)) + + message = "The input `eigenvalues` must be a one-dimensional..." + with pytest.raises(ValueError, match=message): + _covariance.CovViaEigendecomposition(("alpaca", np.eye(2))) + + message = "The input `eigenvectors` must be a square..." + with pytest.raises(ValueError, match=message): + _covariance.CovViaEigendecomposition((np.ones(2), "alpaca")) + + message = "The shapes of `eigenvalues` and `eigenvectors` must be..." + with pytest.raises(ValueError, match=message): + _covariance.CovViaEigendecomposition(([1, 2, 3], np.eye(2))) + + _covariance_preprocessing = {"Diagonal": np.diag, + "Precision": np.linalg.inv, + "Cholesky": np.linalg.cholesky, + "Eigendecomposition": np.linalg.eigh, + "PSD": lambda x: + _PSD(x, allow_singular=True)} + _all_covariance_types = np.array(list(_covariance_preprocessing)) + _matrices = {"diagonal full rank": np.diag([1, 2, 3]), + "general full rank": [[5, 1, 3], [1, 6, 4], [3, 4, 7]], + "diagonal singular": np.diag([1, 0, 3]), + "general singular": [[5, -1, 0], [-1, 5, 0], [0, 0, 0]]} + _cov_types = {"diagonal full rank": _all_covariance_types, + "general full rank": _all_covariance_types[1:], + "diagonal singular": _all_covariance_types[[0, -2, -1]], + "general singular": _all_covariance_types[-2:]} + + @pytest.mark.parametrize("cov_type_name", _all_covariance_types[:-1]) + def test_factories(self, cov_type_name): + A = np.diag([1, 2, 3]) + x = [-4, 2, 5] + + cov_type = getattr(_covariance, f"CovVia{cov_type_name}") + preprocessing = self._covariance_preprocessing[cov_type_name] + factory = getattr(Covariance, f"from_{cov_type_name.lower()}") + + res = factory(preprocessing(A)) + ref = cov_type(preprocessing(A)) + assert type(res) is type(ref) + assert_allclose(res.whiten(x), ref.whiten(x)) + + @pytest.mark.parametrize("matrix_type", list(_matrices)) + @pytest.mark.parametrize("cov_type_name", _all_covariance_types) + def test_covariance(self, matrix_type, cov_type_name): + message = (f"CovVia{cov_type_name} does not support {matrix_type} " + "matrices") + if cov_type_name not in self._cov_types[matrix_type]: + pytest.skip(message) + + A = self._matrices[matrix_type] + cov_type = getattr(_covariance, f"CovVia{cov_type_name}") + preprocessing = self._covariance_preprocessing[cov_type_name] + + psd = _PSD(A, allow_singular=True) + + # test properties + cov_object = cov_type(preprocessing(A)) + assert_close(cov_object.log_pdet, psd.log_pdet) + assert_equal(cov_object.rank, psd.rank) + assert_equal(cov_object.shape, np.asarray(A).shape) + assert_close(cov_object.covariance, np.asarray(A)) + + # test whitening/coloring 1D x + rng = np.random.default_rng(5292808890472453840) + x = rng.random(size=3) + res = cov_object.whiten(x) + ref = x @ psd.U + # res != ref in general; but res @ res == ref @ ref + assert_close(res @ res, ref @ ref) + if hasattr(cov_object, "_colorize") and "singular" not in matrix_type: + # CovViaPSD does not have _colorize + assert_close(cov_object.colorize(res), x) + + # test whitening/coloring 3D x + x = rng.random(size=(2, 4, 3)) + res = cov_object.whiten(x) + ref = x @ psd.U + assert_close((res**2).sum(axis=-1), (ref**2).sum(axis=-1)) + if hasattr(cov_object, "_colorize") and "singular" not in matrix_type: + assert_close(cov_object.colorize(res), x) + + # gh-19197 reported that multivariate normal `rvs` produced incorrect + # results when a singular Covariance object was produce using + # `from_eigenvalues`. This was due to an issue in `colorize` with + # singular covariance matrices. Check this edge case, which is skipped + # in the previous tests. + if hasattr(cov_object, "_colorize"): + res = cov_object.colorize(np.eye(len(A))) + assert_close(res.T @ res, A) + + @pytest.mark.parametrize("size", [None, tuple(), 1, (2, 4, 3)]) + @pytest.mark.parametrize("matrix_type", list(_matrices)) + @pytest.mark.parametrize("cov_type_name", _all_covariance_types) + def test_mvn_with_covariance(self, size, matrix_type, cov_type_name): + message = (f"CovVia{cov_type_name} does not support {matrix_type} " + "matrices") + if cov_type_name not in self._cov_types[matrix_type]: + pytest.skip(message) + + A = self._matrices[matrix_type] + cov_type = getattr(_covariance, f"CovVia{cov_type_name}") + preprocessing = self._covariance_preprocessing[cov_type_name] + + mean = [0.1, 0.2, 0.3] + cov_object = cov_type(preprocessing(A)) + mvn = multivariate_normal + dist0 = multivariate_normal(mean, A, allow_singular=True) + dist1 = multivariate_normal(mean, cov_object, allow_singular=True) + + rng = np.random.default_rng(5292808890472453840) + x = rng.multivariate_normal(mean, A, size=size) + rng = np.random.default_rng(5292808890472453840) + x1 = mvn.rvs(mean, cov_object, size=size, random_state=rng) + rng = np.random.default_rng(5292808890472453840) + x2 = mvn(mean, cov_object, seed=rng).rvs(size=size) + if isinstance(cov_object, _covariance.CovViaPSD): + assert_close(x1, np.squeeze(x)) # for backward compatibility + assert_close(x2, np.squeeze(x)) + else: + assert_equal(x1.shape, x.shape) + assert_equal(x2.shape, x.shape) + assert_close(x2, x1) + + assert_close(mvn.pdf(x, mean, cov_object), dist0.pdf(x)) + assert_close(dist1.pdf(x), dist0.pdf(x)) + assert_close(mvn.logpdf(x, mean, cov_object), dist0.logpdf(x)) + assert_close(dist1.logpdf(x), dist0.logpdf(x)) + assert_close(mvn.entropy(mean, cov_object), dist0.entropy()) + assert_close(dist1.entropy(), dist0.entropy()) + + @pytest.mark.parametrize("size", [tuple(), (2, 4, 3)]) + @pytest.mark.parametrize("cov_type_name", _all_covariance_types) + def test_mvn_with_covariance_cdf(self, size, cov_type_name): + # This is split from the test above because it's slow to be running + # with all matrix types, and there's no need because _mvn.mvnun + # does the calculation. All Covariance needs to do is pass is + # provide the `covariance` attribute. + matrix_type = "diagonal full rank" + A = self._matrices[matrix_type] + cov_type = getattr(_covariance, f"CovVia{cov_type_name}") + preprocessing = self._covariance_preprocessing[cov_type_name] + + mean = [0.1, 0.2, 0.3] + cov_object = cov_type(preprocessing(A)) + mvn = multivariate_normal + dist0 = multivariate_normal(mean, A, allow_singular=True) + dist1 = multivariate_normal(mean, cov_object, allow_singular=True) + + rng = np.random.default_rng(5292808890472453840) + x = rng.multivariate_normal(mean, A, size=size) + + assert_close(mvn.cdf(x, mean, cov_object), dist0.cdf(x)) + assert_close(dist1.cdf(x), dist0.cdf(x)) + assert_close(mvn.logcdf(x, mean, cov_object), dist0.logcdf(x)) + assert_close(dist1.logcdf(x), dist0.logcdf(x)) + + def test_covariance_instantiation(self): + message = "The `Covariance` class cannot be instantiated directly." + with pytest.raises(NotImplementedError, match=message): + Covariance() + + @pytest.mark.filterwarnings("ignore::RuntimeWarning") # matrix not PSD + def test_gh9942(self): + # Originally there was a mistake in the `multivariate_normal_frozen` + # `rvs` method that caused all covariance objects to be processed as + # a `_CovViaPSD`. Ensure that this is resolved. + A = np.diag([1, 2, -1e-8]) + n = A.shape[0] + mean = np.zeros(n) + + # Error if the matrix is processed as a `_CovViaPSD` + with pytest.raises(ValueError, match="The input matrix must be..."): + multivariate_normal(mean, A).rvs() + + # No error if it is provided as a `CovViaEigendecomposition` + seed = 3562050283508273023 + rng1 = np.random.default_rng(seed) + rng2 = np.random.default_rng(seed) + cov = Covariance.from_eigendecomposition(np.linalg.eigh(A)) + rv = multivariate_normal(mean, cov) + res = rv.rvs(random_state=rng1) + ref = multivariate_normal.rvs(mean, cov, random_state=rng2) + assert_equal(res, ref) + + def test_gh19197(self): + # gh-19197 reported that multivariate normal `rvs` produced incorrect + # results when a singular Covariance object was produce using + # `from_eigenvalues`. Check that this specific issue is resolved; + # a more general test is included in `test_covariance`. + mean = np.ones(2) + cov = Covariance.from_eigendecomposition((np.zeros(2), np.eye(2))) + dist = scipy.stats.multivariate_normal(mean=mean, cov=cov) + rvs = dist.rvs(size=None) + assert_equal(rvs, mean) + + cov = scipy.stats.Covariance.from_eigendecomposition( + (np.array([1., 0.]), np.array([[1., 0.], [0., 400.]]))) + dist = scipy.stats.multivariate_normal(mean=mean, cov=cov) + rvs = dist.rvs(size=None) + assert rvs[0] != mean[0] + assert rvs[1] == mean[1] + + +def _random_covariance(dim, evals, rng, singular=False): + # Generates random covariance matrix with dimensionality `dim` and + # eigenvalues `evals` using provided Generator `rng`. Randomly sets + # some evals to zero if `singular` is True. + A = rng.random((dim, dim)) + A = A @ A.T + _, v = np.linalg.eigh(A) + if singular: + zero_eigs = rng.normal(size=dim) > 0 + evals[zero_eigs] = 0 + cov = v @ np.diag(evals) @ v.T + return cov + + +def _sample_orthonormal_matrix(n): + rng = np.random.default_rng(9086764251) + M = rng.standard_normal((n, n)) + u, s, v = scipy.linalg.svd(M) + return u + +def marginal_pdf(X, X_ndim, dimensions, x): + """Integrate marginalized dimensions of multivariate + probability distribution to calculate the marginalized + distribution. + """ + # Sort input data based on order of dimensions + dimensions = np.asarray(dimensions) + dimensions[dimensions < 0] += X_ndim + dim_sort_idx = dimensions.argsort() + x = x[:, dim_sort_idx] + + i_marginalize = np.ones(X_ndim, dtype=bool) + i_marginalize[dimensions] = False + + def g(z): + y = np.empty((z.shape[0], x.shape[0], X_ndim)) + y[..., i_marginalize] = z[:, np.newaxis, :] + y[..., ~i_marginalize] = x + return X.pdf(y) + + inf = np.full(X_ndim - len(dimensions), np.inf) + return cubature(g, -inf, inf).estimate + +@dataclass +class MVNProblem: + """Instantiate a multivariate normal integration problem with special structure. + + When covariance matrix is a correlation matrix where the off-diagonal entries + ``covar[i, j] == lambdas[i]*lambdas[j]`` for ``i != j``, then the multidimensional + integral reduces to a simpler univariate integral that can be numerically integrated + easily. + + The ``generate_*()`` classmethods provide a few options for creating variations + of this problem. + + References + ---------- + .. [1] Tong, Y.L. "The Multivariate Normal Distribution". + Springer-Verlag. p192. 1990. + """ + ndim : int + low : np.ndarray + high : np.ndarray + lambdas : np.ndarray + covar : np.ndarray + target_val : float + target_err : float + + #: The `generator_halves()` case has an analytically-known true value that we'll + #: record here. It remain None for most cases, though. + true_val : float | None = None + + def __init__(self, ndim, low, high, lambdas): + super().__init__() + self.ndim = ndim + self.low = low + self.high = high + self.lambdas = lambdas + + self.covar = np.outer(self.lambdas, self.lambdas) + np.fill_diagonal(self.covar, 1.0) + self.find_target() + + @classmethod + def generate_semigeneral(cls, ndim, rng=None): + """Random lambdas, random upper bounds, infinite lower bounds. + """ + rng = np.random.default_rng(rng) + low = np.full(ndim, -np.inf) + high = rng.uniform(0.0, np.sqrt(ndim), size=ndim) + lambdas = rng.uniform(-1.0, 1.0, size=ndim) + + self = cls( + ndim=ndim, + low=low, + high=high, + lambdas=lambdas, + ) + return self + + @classmethod + def generate_constant(cls, ndim, rng=None): + """Constant off-diagonal covariance, random upper bounds, infinite lower bounds. + """ + rng = np.random.default_rng(rng) + low = np.full(ndim, -np.inf) + high = rng.uniform(0.0, np.sqrt(ndim), size=ndim) + sigma = np.sqrt(rng.uniform(0.0, 1.0)) + lambdas = np.full(ndim, sigma) + + self = cls( + ndim=ndim, + low=low, + high=high, + lambdas=lambdas, + ) + return self + + @classmethod + def generate_halves(cls, ndim, rng=None): + """Off-diagonal covariance of 0.5, negative orthant bounds. + + True analytically-derived answer is 1/(ndim+1). + """ + low = np.full(ndim, -np.inf) + high = np.zeros(ndim) + lambdas = np.sqrt(0.5) + + self = cls( + ndim=ndim, + low=low, + high=high, + lambdas=lambdas, + ) + self.true_val = 1 / (ndim+1) + return self + + def find_target(self, **kwds): + """Perform the simplified integral and store the results. + """ + d = dict( + a=-9.0, + b=+9.0, + ) + d.update(kwds) + self.target_val, self.target_err = quad(self.univariate_func, **d) + + def _univariate_term(self, t): + """The parameter-specific term of the univariate integrand, + for separate plotting. + """ + denom = np.sqrt(1 - self.lambdas**2) + return np.prod( + special.ndtr((self.high + self.lambdas*t[:, np.newaxis]) / denom) - + special.ndtr((self.low + self.lambdas*t[:, np.newaxis]) / denom), + axis=1, + ) + + def univariate_func(self, t): + """Univariate integrand. + """ + t = np.atleast_1d(t) + return np.squeeze(norm_pdf(t) * self._univariate_term(t)) + + def plot_integrand(self): + """Plot the univariate integrand and its component terms for understanding. + """ + from matplotlib import pyplot as plt + + t = np.linspace(-9.0, 9.0, 1001) + plt.plot(t, norm_pdf(t), label=r'$\phi(t)$') + plt.plot(t, self._univariate_term(t), label=r'$f(t)$') + plt.plot(t, self.univariate_func(t), label=r'$f(t)*phi(t)$') + plt.legend() + + +@dataclass +class SingularMVNProblem: + """Instantiate a multivariate normal integration problem with a special singular + covariance structure. + + When covariance matrix is a correlation matrix where the off-diagonal entries + ``covar[i, j] == -lambdas[i]*lambdas[j]`` for ``i != j``, and + ``sum(lambdas**2 / (1+lambdas**2)) == 1``, then the matrix is singular, and + the multidimensional integral reduces to a simpler univariate integral that + can be numerically integrated fairly easily. + + The lower bound must be infinite, though the upper bounds can be general. + + References + ---------- + .. [1] Kwong, K.-S. (1995). "Evaluation of the one-sided percentage points of the + singular multivariate normal distribution." Journal of Statistical + Computation and Simulation, 51(2-4), 121-135. + :doi:`10.1080/00949659508811627`. + """ + ndim : int + low : np.ndarray + high : np.ndarray + lambdas : np.ndarray + covar : np.ndarray + target_val : float + target_err : float + + def __init__(self, ndim, high, lambdas): + self.ndim = ndim + self.high = high + self.lambdas = lambdas + + self.low = np.full(ndim, -np.inf) + self.covar = -np.outer(self.lambdas, self.lambdas) + np.fill_diagonal(self.covar, 1.0) + self.find_target() + + @classmethod + def generate_semiinfinite(cls, ndim, rng=None): + """Singular lambdas, random upper bounds. + """ + rng = np.random.default_rng(rng) + high = rng.uniform(0.0, np.sqrt(ndim), size=ndim) + p = rng.dirichlet(np.full(ndim, 1.0)) + lambdas = np.sqrt(p / (1-p)) * rng.choice([-1.0, 1.0], size=ndim) + + self = cls( + ndim=ndim, + high=high, + lambdas=lambdas, + ) + return self + + def find_target(self, **kwds): + d = dict( + a=-9.0, + b=+9.0, + ) + d.update(kwds) + self.target_val, self.target_err = quad(self.univariate_func, **d) + + def _univariate_term(self, t): + denom = np.sqrt(1 + self.lambdas**2) + i1 = np.prod( + special.ndtr((self.high - 1j*self.lambdas*t[:, np.newaxis]) / denom), + axis=1, + ) + i2 = np.prod( + special.ndtr((-self.high + 1j*self.lambdas*t[:, np.newaxis]) / denom), + axis=1, + ) + # The imaginary part is an odd function, so it can be ignored; it will integrate + # out to 0. + return (i1 - (-1)**self.ndim * i2).real + + def univariate_func(self, t): + t = np.atleast_1d(t) + return (norm_pdf(t) * self._univariate_term(t)).squeeze() + + def plot_integrand(self): + """Plot the univariate integrand and its component terms for understanding. + """ + from matplotlib import pyplot as plt + + t = np.linspace(-9.0, 9.0, 1001) + plt.plot(t, norm_pdf(t), label=r'$\phi(t)$') + plt.plot(t, self._univariate_term(t), label=r'$f(t)$') + plt.plot(t, self.univariate_func(t), label=r'$f(t)*phi(t)$') + plt.ylim(-0.1, 1.1) + plt.legend() + + +class TestMultivariateNormal: + def test_input_shape(self): + mu = np.arange(3) + cov = np.identity(2) + assert_raises(ValueError, multivariate_normal.pdf, (0, 1), mu, cov) + assert_raises(ValueError, multivariate_normal.pdf, (0, 1, 2), mu, cov) + assert_raises(ValueError, multivariate_normal.cdf, (0, 1), mu, cov) + assert_raises(ValueError, multivariate_normal.cdf, (0, 1, 2), mu, cov) + + def test_scalar_values(self): + rng = np.random.default_rng(1234) + + # When evaluated on scalar data, the pdf should return a scalar + x, mean, cov = 1.5, 1.7, 2.5 + pdf = multivariate_normal.pdf(x, mean, cov) + assert_equal(pdf.ndim, 0) + + # When evaluated on a single vector, the pdf should return a scalar + x = rng.standard_normal(5) + mean = rng.standard_normal(5) + cov = np.abs(rng.standard_normal(5)) # Diagonal values for cov. matrix + pdf = multivariate_normal.pdf(x, mean, cov) + assert_equal(pdf.ndim, 0) + + # When evaluated on scalar data, the cdf should return a scalar + x, mean, cov = 1.5, 1.7, 2.5 + cdf = multivariate_normal.cdf(x, mean, cov) + assert_equal(cdf.ndim, 0) + + # When evaluated on a single vector, the cdf should return a scalar + x = rng.standard_normal(5) + mean = rng.standard_normal(5) + cov = np.abs(rng.standard_normal(5)) # Diagonal values for cov. matrix + cdf = multivariate_normal.cdf(x, mean, cov) + assert_equal(cdf.ndim, 0) + + def test_logpdf(self): + # Check that the log of the pdf is in fact the logpdf + rng = np.random.default_rng(1234) + x = rng.standard_normal(5) + mean = rng.standard_normal(5) + cov = np.abs(rng.standard_normal(5)) + d1 = multivariate_normal.logpdf(x, mean, cov) + d2 = multivariate_normal.pdf(x, mean, cov) + assert_allclose(d1, np.log(d2)) + + def test_logpdf_default_values(self): + # Check that the log of the pdf is in fact the logpdf + # with default parameters Mean=None and cov = 1 + rng = np.random.default_rng(1234) + x = rng.standard_normal(5) + d1 = multivariate_normal.logpdf(x) + d2 = multivariate_normal.pdf(x) + # check whether default values are being used + d3 = multivariate_normal.logpdf(x, None, 1) + d4 = multivariate_normal.pdf(x, None, 1) + assert_allclose(d1, np.log(d2)) + assert_allclose(d3, np.log(d4)) + + def test_logcdf(self): + # Check that the log of the cdf is in fact the logcdf + rng = np.random.default_rng(1234) + x = rng.standard_normal(5) + mean = rng.standard_normal(5) + cov = np.abs(rng.standard_normal(5)) + d1 = multivariate_normal.logcdf(x, mean, cov) + d2 = multivariate_normal.cdf(x, mean, cov) + assert_allclose(d1, np.log(d2)) + + def test_logcdf_default_values(self): + # Check that the log of the cdf is in fact the logcdf + # with default parameters Mean=None and cov = 1 + rng = np.random.default_rng(1234) + x = rng.standard_normal(5) + d1 = multivariate_normal.logcdf(x) + d2 = multivariate_normal.cdf(x) + # check whether default values are being used + d3 = multivariate_normal.logcdf(x, None, 1) + d4 = multivariate_normal.cdf(x, None, 1) + assert_allclose(d1, np.log(d2)) + assert_allclose(d3, np.log(d4)) + + def test_rank(self): + # Check that the rank is detected correctly. + rng = np.random.default_rng(1234) + n = 4 + mean = rng.standard_normal(n) + for expected_rank in range(1, n + 1): + s = rng.standard_normal((n, expected_rank)) + cov = np.dot(s, s.T) + distn = multivariate_normal(mean, cov, allow_singular=True) + assert_equal(distn.cov_object.rank, expected_rank) + + def test_degenerate_distributions(self): + rng = np.random.default_rng(1234) + for n in range(1, 5): + z = rng.standard_normal(n) + for k in range(1, n): + # Sample a small covariance matrix. + s = rng.standard_normal((k, k)) + cov_kk = np.dot(s, s.T) + + # Embed the small covariance matrix into a larger singular one. + cov_nn = np.zeros((n, n)) + cov_nn[:k, :k] = cov_kk + + # Embed part of the vector in the same way + x = np.zeros(n) + x[:k] = z[:k] + + # Define a rotation of the larger low rank matrix. + u = _sample_orthonormal_matrix(n) + cov_rr = np.dot(u, np.dot(cov_nn, u.T)) + y = np.dot(u, x) + + # Check some identities. + distn_kk = multivariate_normal(np.zeros(k), cov_kk, + allow_singular=True) + distn_nn = multivariate_normal(np.zeros(n), cov_nn, + allow_singular=True) + distn_rr = multivariate_normal(np.zeros(n), cov_rr, + allow_singular=True) + assert_equal(distn_kk.cov_object.rank, k) + assert_equal(distn_nn.cov_object.rank, k) + assert_equal(distn_rr.cov_object.rank, k) + pdf_kk = distn_kk.pdf(x[:k]) + pdf_nn = distn_nn.pdf(x) + pdf_rr = distn_rr.pdf(y) + assert_allclose(pdf_kk, pdf_nn) + assert_allclose(pdf_kk, pdf_rr) + logpdf_kk = distn_kk.logpdf(x[:k]) + logpdf_nn = distn_nn.logpdf(x) + logpdf_rr = distn_rr.logpdf(y) + assert_allclose(logpdf_kk, logpdf_nn) + assert_allclose(logpdf_kk, logpdf_rr) + + # Add an orthogonal component and find the density + y_orth = y + u[:, -1] + pdf_rr_orth = distn_rr.pdf(y_orth) + logpdf_rr_orth = distn_rr.logpdf(y_orth) + + # Ensure that this has zero probability + assert_equal(pdf_rr_orth, 0.0) + assert_equal(logpdf_rr_orth, -np.inf) + + def test_degenerate_array(self): + # Test that we can generate arrays of random variate from a degenerate + # multivariate normal, and that the pdf for these samples is non-zero + # (i.e. samples from the distribution lie on the subspace) + k = 10 + for n in range(2, 6): + for r in range(1, n): + mn = np.zeros(n) + u = _sample_orthonormal_matrix(n)[:, :r] + vr = np.dot(u, u.T) + X = multivariate_normal.rvs(mean=mn, cov=vr, size=k) + + pdf = multivariate_normal.pdf(X, mean=mn, cov=vr, + allow_singular=True) + assert_equal(pdf.size, k) + assert np.all(pdf > 0.0) + + logpdf = multivariate_normal.logpdf(X, mean=mn, cov=vr, + allow_singular=True) + assert_equal(logpdf.size, k) + assert np.all(logpdf > -np.inf) + + def test_large_pseudo_determinant(self): + # Check that large pseudo-determinants are handled appropriately. + + # Construct a singular diagonal covariance matrix + # whose pseudo determinant overflows double precision. + large_total_log = 1000.0 + npos = 100 + nzero = 2 + large_entry = np.exp(large_total_log / npos) + n = npos + nzero + cov = np.zeros((n, n), dtype=float) + np.fill_diagonal(cov, large_entry) + cov[-nzero:, -nzero:] = 0 + + # Check some determinants. + assert_equal(scipy.linalg.det(cov), 0) + assert_equal(scipy.linalg.det(cov[:npos, :npos]), np.inf) + assert_allclose(np.linalg.slogdet(cov[:npos, :npos]), + (1, large_total_log)) + + # Check the pseudo-determinant. + psd = _PSD(cov) + assert_allclose(psd.log_pdet, large_total_log) + + def test_broadcasting(self): + rng = np.random.RandomState(1234) + n = 4 + + # Construct a random covariance matrix. + data = rng.randn(n, n) + cov = np.dot(data, data.T) + mean = rng.randn(n) + + # Construct an ndarray which can be interpreted as + # a 2x3 array whose elements are random data vectors. + X = rng.randn(2, 3, n) + + # Check that multiple data points can be evaluated at once. + desired_pdf = multivariate_normal.pdf(X, mean, cov) + desired_cdf = multivariate_normal.cdf(X, mean, cov) + for i in range(2): + for j in range(3): + actual = multivariate_normal.pdf(X[i, j], mean, cov) + assert_allclose(actual, desired_pdf[i,j]) + # Repeat for cdf + actual = multivariate_normal.cdf(X[i, j], mean, cov) + assert_allclose(actual, desired_cdf[i,j], rtol=1e-3) + + def test_normal_1D(self): + # The probability density function for a 1D normal variable should + # agree with the standard normal distribution in scipy.stats.distributions + x = np.linspace(0, 2, 10) + mean, cov = 1.2, 0.9 + scale = cov**0.5 + d1 = norm.pdf(x, mean, scale) + d2 = multivariate_normal.pdf(x, mean, cov) + assert_allclose(d1, d2) + # The same should hold for the cumulative distribution function + d1 = norm.cdf(x, mean, scale) + d2 = multivariate_normal.cdf(x, mean, cov) + assert_allclose(d1, d2) + + def test_marginalization(self): + # Integrating out one of the variables of a 2D Gaussian should + # yield a 1D Gaussian + mean = np.array([2.5, 3.5]) + cov = np.array([[.5, 0.2], [0.2, .6]]) + n = 2 ** 8 + 1 # Number of samples + delta = 6 / (n - 1) # Grid spacing + + v = np.linspace(0, 6, n) + xv, yv = np.meshgrid(v, v) + pos = np.empty((n, n, 2)) + pos[:, :, 0] = xv + pos[:, :, 1] = yv + pdf = multivariate_normal.pdf(pos, mean, cov) + + # Marginalize over x and y axis + margin_x = romb(pdf, delta, axis=0) + margin_y = romb(pdf, delta, axis=1) + + # Compare with standard normal distribution + gauss_x = norm.pdf(v, loc=mean[0], scale=cov[0, 0] ** 0.5) + gauss_y = norm.pdf(v, loc=mean[1], scale=cov[1, 1] ** 0.5) + assert_allclose(margin_x, gauss_x, rtol=1e-2, atol=1e-2) + assert_allclose(margin_y, gauss_y, rtol=1e-2, atol=1e-2) + + def test_frozen(self): + # The frozen distribution should agree with the regular one + rng = np.random.default_rng(1234) + x = rng.standard_normal(5) + mean = rng.standard_normal(5) + cov = np.abs(rng.standard_normal(5)) + norm_frozen = multivariate_normal(mean, cov) + assert_allclose(norm_frozen.pdf(x), multivariate_normal.pdf(x, mean, cov)) + assert_allclose(norm_frozen.logpdf(x), + multivariate_normal.logpdf(x, mean, cov)) + assert_allclose(norm_frozen.cdf(x), multivariate_normal.cdf(x, mean, cov)) + assert_allclose(norm_frozen.logcdf(x), + multivariate_normal.logcdf(x, mean, cov)) + + @pytest.mark.parametrize( + 'covariance', + [ + np.eye(2), + Covariance.from_diagonal([1, 1]), + ] + ) + def test_frozen_multivariate_normal_exposes_attributes(self, covariance): + mean = np.ones((2,)) + cov_should_be = np.eye(2) + norm_frozen = multivariate_normal(mean, covariance) + assert np.allclose(norm_frozen.mean, mean) + assert np.allclose(norm_frozen.cov, cov_should_be) + + def test_pseudodet_pinv(self): + # Make sure that pseudo-inverse and pseudo-det agree on cutoff + + # Assemble random covariance matrix with large and small eigenvalues + rng = np.random.default_rng(1234) + n = 7 + x = rng.standard_normal((n, n)) + cov = np.dot(x, x.T) + s, u = scipy.linalg.eigh(cov) + s = np.full(n, 0.5) + s[0] = 1.0 + s[-1] = 1e-7 + cov = np.dot(u, np.dot(np.diag(s), u.T)) + + # Set cond so that the lowest eigenvalue is below the cutoff + cond = 1e-5 + psd = _PSD(cov, cond=cond) + psd_pinv = _PSD(psd.pinv, cond=cond) + + # Check that the log pseudo-determinant agrees with the sum + # of the logs of all but the smallest eigenvalue + assert_allclose(psd.log_pdet, np.sum(np.log(s[:-1]))) + # Check that the pseudo-determinant of the pseudo-inverse + # agrees with 1 / pseudo-determinant + assert_allclose(-psd.log_pdet, psd_pinv.log_pdet) + + def test_exception_nonsquare_cov(self): + cov = [[1, 2, 3], [4, 5, 6]] + assert_raises(ValueError, _PSD, cov) + + def test_exception_nonfinite_cov(self): + cov_nan = [[1, 0], [0, np.nan]] + assert_raises(ValueError, _PSD, cov_nan) + cov_inf = [[1, 0], [0, np.inf]] + assert_raises(ValueError, _PSD, cov_inf) + + def test_exception_non_psd_cov(self): + cov = [[1, 0], [0, -1]] + assert_raises(ValueError, _PSD, cov) + + def test_exception_singular_cov(self): + rng = np.random.default_rng(1234) + x = rng.standard_normal(5) + mean = rng.standard_normal(5) + cov = np.ones((5, 5)) + e = np.linalg.LinAlgError + assert_raises(e, multivariate_normal, mean, cov) + assert_raises(e, multivariate_normal.pdf, x, mean, cov) + assert_raises(e, multivariate_normal.logpdf, x, mean, cov) + assert_raises(e, multivariate_normal.cdf, x, mean, cov) + assert_raises(e, multivariate_normal.logcdf, x, mean, cov) + + # Message used to be "singular matrix", but this is more accurate. + # See gh-15508 + cov = [[1., 0.], [1., 1.]] + msg = "When `allow_singular is False`, the input matrix" + with pytest.raises(np.linalg.LinAlgError, match=msg): + multivariate_normal(cov=cov) + + def test_R_values(self): + # Compare the multivariate pdf with some values precomputed + # in R version 3.0.1 (2013-05-16) on Mac OS X 10.6. + + # The values below were generated by the following R-script: + # > library(mnormt) + # > x <- seq(0, 2, length=5) + # > y <- 3*x - 2 + # > z <- x + cos(y) + # > mu <- c(1, 3, 2) + # > Sigma <- matrix(c(1,2,0,2,5,0.5,0,0.5,3), 3, 3) + # > r_pdf <- dmnorm(cbind(x,y,z), mu, Sigma) + r_pdf = np.array([0.0002214706, 0.0013819953, 0.0049138692, + 0.0103803050, 0.0140250800]) + + x = np.linspace(0, 2, 5) + y = 3 * x - 2 + z = x + np.cos(y) + r = np.array([x, y, z]).T + + mean = np.array([1, 3, 2], 'd') + cov = np.array([[1, 2, 0], [2, 5, .5], [0, .5, 3]], 'd') + + pdf = multivariate_normal.pdf(r, mean, cov) + assert_allclose(pdf, r_pdf, atol=1e-10) + + # Compare the multivariate cdf with some values precomputed + # in R version 3.3.2 (2016-10-31) on Debian GNU/Linux. + + # The values below were generated by the following R-script: + # > library(mnormt) + # > x <- seq(0, 2, length=5) + # > y <- 3*x - 2 + # > z <- x + cos(y) + # > mu <- c(1, 3, 2) + # > Sigma <- matrix(c(1,2,0,2,5,0.5,0,0.5,3), 3, 3) + # > r_cdf <- pmnorm(cbind(x,y,z), mu, Sigma) + r_cdf = np.array([0.0017866215, 0.0267142892, 0.0857098761, + 0.1063242573, 0.2501068509]) + + cdf = multivariate_normal.cdf(r, mean, cov) + assert_allclose(cdf, r_cdf, atol=2e-5) + + # Also test bivariate cdf with some values precomputed + # in R version 3.3.2 (2016-10-31) on Debian GNU/Linux. + + # The values below were generated by the following R-script: + # > library(mnormt) + # > x <- seq(0, 2, length=5) + # > y <- 3*x - 2 + # > mu <- c(1, 3) + # > Sigma <- matrix(c(1,2,2,5), 2, 2) + # > r_cdf2 <- pmnorm(cbind(x,y), mu, Sigma) + r_cdf2 = np.array([0.01262147, 0.05838989, 0.18389571, + 0.40696599, 0.66470577]) + + r2 = np.array([x, y]).T + + mean2 = np.array([1, 3], 'd') + cov2 = np.array([[1, 2], [2, 5]], 'd') + + cdf2 = multivariate_normal.cdf(r2, mean2, cov2) + assert_allclose(cdf2, r_cdf2, atol=1e-5) + + def test_multivariate_normal_rvs_zero_covariance(self): + mean = np.zeros(2) + covariance = np.zeros((2, 2)) + model = multivariate_normal(mean, covariance, allow_singular=True) + sample = model.rvs() + assert_equal(sample, [0, 0]) + + def test_rvs_shape(self): + # Check that rvs parses the mean and covariance correctly, and returns + # an array of the right shape + N = 300 + d = 4 + sample = multivariate_normal.rvs(mean=np.zeros(d), cov=1, size=N) + assert_equal(sample.shape, (N, d)) + + sample = multivariate_normal.rvs(mean=None, + cov=np.array([[2, .1], [.1, 1]]), + size=N) + assert_equal(sample.shape, (N, 2)) + + u = multivariate_normal(mean=0, cov=1) + sample = u.rvs(N) + assert_equal(sample.shape, (N, )) + + def test_large_sample(self): + # Generate large sample and compare sample mean and sample covariance + # with mean and covariance matrix. + + rng = np.random.RandomState(2846) + + n = 3 + mean = rng.randn(n) + M = rng.randn(n, n) + cov = np.dot(M, M.T) + size = 5000 + + sample = multivariate_normal.rvs(mean, cov, size, random_state=rng) + + assert_allclose(np.cov(sample.T), cov, rtol=1e-1) + assert_allclose(sample.mean(0), mean, rtol=1e-1) + + def test_entropy(self): + rng = np.random.RandomState(2846) + + n = 3 + mean = rng.randn(n) + M = rng.randn(n, n) + cov = np.dot(M, M.T) + + rv = multivariate_normal(mean, cov) + + # Check that frozen distribution agrees with entropy function + assert_almost_equal(rv.entropy(), multivariate_normal.entropy(mean, cov)) + # Compare entropy with manually computed expression involving + # the sum of the logs of the eigenvalues of the covariance matrix + eigs = np.linalg.eig(cov)[0] + desired = 1 / 2 * (n * (np.log(2 * np.pi) + 1) + np.sum(np.log(eigs))) + assert_almost_equal(desired, rv.entropy()) + + def test_lnB(self): + alpha = np.array([1, 1, 1]) + desired = .5 # e^lnB = 1/2 for [1, 1, 1] + + assert_almost_equal(np.exp(_lnB(alpha)), desired) + + def test_cdf_with_lower_limit_arrays(self): + # test CDF with lower limit in several dimensions + rng = np.random.default_rng(2408071309372769818) + mean = [0, 0] + cov = np.eye(2) + a = rng.random((4, 3, 2))*6 - 3 + b = rng.random((4, 3, 2))*6 - 3 + + cdf1 = multivariate_normal.cdf(b, mean, cov, lower_limit=a) + + cdf2a = multivariate_normal.cdf(b, mean, cov) + cdf2b = multivariate_normal.cdf(a, mean, cov) + ab1 = np.concatenate((a[..., 0:1], b[..., 1:2]), axis=-1) + ab2 = np.concatenate((a[..., 1:2], b[..., 0:1]), axis=-1) + cdf2ab1 = multivariate_normal.cdf(ab1, mean, cov) + cdf2ab2 = multivariate_normal.cdf(ab2, mean, cov) + cdf2 = cdf2a + cdf2b - cdf2ab1 - cdf2ab2 + + assert_allclose(cdf1, cdf2) + + def test_cdf_with_lower_limit_consistency(self): + # check that multivariate normal CDF functions are consistent + rng = np.random.default_rng(2408071309372769818) + mean = rng.random(3) + cov = rng.random((3, 3)) + cov = cov @ cov.T + a = rng.random((2, 3))*6 - 3 + b = rng.random((2, 3))*6 - 3 + + cdf1 = multivariate_normal.cdf(b, mean, cov, lower_limit=a) + cdf2 = multivariate_normal(mean, cov).cdf(b, lower_limit=a) + cdf3 = np.exp(multivariate_normal.logcdf(b, mean, cov, lower_limit=a)) + cdf4 = np.exp(multivariate_normal(mean, cov).logcdf(b, lower_limit=a)) + + assert_allclose(cdf2, cdf1, rtol=1e-4) + assert_allclose(cdf3, cdf1, rtol=1e-4) + assert_allclose(cdf4, cdf1, rtol=1e-4) + + def test_cdf_signs(self): + # check that sign of output is correct when np.any(lower > x) + mean = np.zeros(3) + cov = np.eye(3) + b = [[1, 1, 1], [0, 0, 0], [1, 0, 1], [0, 1, 0]] + a = [[0, 0, 0], [1, 1, 1], [0, 1, 0], [1, 0, 1]] + # when odd number of elements of b < a, output is negative + expected_signs = np.array([1, -1, -1, 1]) + cdf = multivariate_normal.cdf(b, mean, cov, lower_limit=a) + assert_allclose(cdf, cdf[0]*expected_signs) + + @pytest.mark.slow + @pytest.mark.parametrize("ndim", [2, 3]) + def test_cdf_vs_cubature(self, ndim): + rng = np.random.default_rng(123) + a = rng.uniform(size=(ndim, ndim)) + cov = a.T @ a + m = rng.uniform(size=ndim) + dist = multivariate_normal(mean=m, cov=cov) + x = rng.uniform(low=-3, high=3, size=(ndim,)) + cdf = dist.cdf(x) + dist_i = multivariate_normal(mean=[0]*ndim, cov=cov) + cdf_i = cubature(dist_i.pdf, [-np.inf]*ndim, x - m).estimate + assert_allclose(cdf, cdf_i, atol=5e-6) + + def test_cdf_known(self): + # https://github.com/scipy/scipy/pull/17410#issuecomment-1312628547 + for ndim in range(2, 12): + cov = np.full((ndim, ndim), 0.5) + np.fill_diagonal(cov, 1.) + dist = multivariate_normal([0]*ndim, cov=cov) + assert_allclose( + dist.cdf([0]*ndim), + 1. / (1. + ndim), + atol=5e-5 + ) + + @pytest.mark.parametrize("ndim", range(2, 10)) + @pytest.mark.parametrize("seed", [0xdeadbeef, 0xdd24528764c9773579731c6b022b48e2]) + def test_cdf_vs_univariate(self, seed, ndim): + rng = np.random.default_rng(seed) + case = MVNProblem.generate_semigeneral(ndim=ndim, rng=rng) + assert (case.low == -np.inf).all() + + dist = multivariate_normal(mean=[0]*ndim, cov=case.covar) + cdf_val = dist.cdf(case.high, rng=rng) + assert_allclose(cdf_val, case.target_val, atol=5e-5) + + @pytest.mark.parametrize("ndim", range(2, 11)) + @pytest.mark.parametrize("seed", [0xdeadbeef, 0xdd24528764c9773579731c6b022b48e2]) + def test_cdf_vs_univariate_2(self, seed, ndim): + rng = np.random.default_rng(seed) + case = MVNProblem.generate_constant(ndim=ndim, rng=rng) + assert (case.low == -np.inf).all() + + dist = multivariate_normal(mean=[0]*ndim, cov=case.covar) + cdf_val = dist.cdf(case.high, rng=rng) + assert_allclose(cdf_val, case.target_val, atol=5e-5) + + @pytest.mark.parametrize("ndim", range(4, 11)) + @pytest.mark.parametrize("seed", [0xdeadbeef, 0xdd24528764c9773579731c6b022b48e4]) + def test_cdf_vs_univariate_singular(self, seed, ndim): + # NB: ndim = 2, 3 has much poorer accuracy than ndim > 3 for many seeds. + # No idea why. + rng = np.random.default_rng(seed) + case = SingularMVNProblem.generate_semiinfinite(ndim=ndim, rng=rng) + assert (case.low == -np.inf).all() + + dist = multivariate_normal(mean=[0]*ndim, cov=case.covar, allow_singular=True, + # default maxpts is too slow, limit it here + maxpts=10_000*case.covar.shape[0] + ) + cdf_val = dist.cdf(case.high, rng=rng) + assert_allclose(cdf_val, case.target_val, atol=1e-3) + + def test_mean_cov(self): + # test the interaction between a Covariance object and mean + P = np.diag(1 / np.array([1, 2, 3])) + cov_object = _covariance.CovViaPrecision(P) + + message = "`cov` represents a covariance matrix in 3 dimensions..." + with pytest.raises(ValueError, match=message): + multivariate_normal.entropy([0, 0], cov_object) + + with pytest.raises(ValueError, match=message): + multivariate_normal([0, 0], cov_object) + + x = [0.5, 0.5, 0.5] + ref = multivariate_normal.pdf(x, [0, 0, 0], cov_object) + assert_equal(multivariate_normal.pdf(x, cov=cov_object), ref) + + ref = multivariate_normal.pdf(x, [1, 1, 1], cov_object) + assert_equal(multivariate_normal.pdf(x, 1, cov=cov_object), ref) + + def test_fit_wrong_fit_data_shape(self): + data = [1, 3] + error_msg = "`x` must be two-dimensional." + with pytest.raises(ValueError, match=error_msg): + multivariate_normal.fit(data) + + @pytest.mark.parametrize('dim', (3, 5)) + def test_fit_correctness(self, dim): + rng = np.random.default_rng(4385269356937404) + x = rng.random((100, dim)) + mean_est, cov_est = multivariate_normal.fit(x) + mean_ref, cov_ref = np.mean(x, axis=0), np.cov(x.T, ddof=0) + assert_allclose(mean_est, mean_ref, atol=1e-15) + assert_allclose(cov_est, cov_ref, rtol=5e-15) + + def test_fit_both_parameters_fixed(self): + data = np.full((2, 1), 3) + mean_fixed = 1. + cov_fixed = np.atleast_2d(1.) + mean, cov = multivariate_normal.fit(data, fix_mean=mean_fixed, + fix_cov=cov_fixed) + assert_equal(mean, mean_fixed) + assert_equal(cov, cov_fixed) + + @pytest.mark.parametrize('fix_mean', [np.zeros((2, 2)), + np.zeros((3, ))]) + def test_fit_fix_mean_input_validation(self, fix_mean): + msg = ("`fix_mean` must be a one-dimensional array the same " + "length as the dimensionality of the vectors `x`.") + with pytest.raises(ValueError, match=msg): + multivariate_normal.fit(np.eye(2), fix_mean=fix_mean) + + @pytest.mark.parametrize('fix_cov', [np.zeros((2, )), + np.zeros((3, 2)), + np.zeros((4, 4))]) + def test_fit_fix_cov_input_validation_dimension(self, fix_cov): + msg = ("`fix_cov` must be a two-dimensional square array " + "of same side length as the dimensionality of the " + "vectors `x`.") + with pytest.raises(ValueError, match=msg): + multivariate_normal.fit(np.eye(3), fix_cov=fix_cov) + + def test_fit_fix_cov_not_positive_semidefinite(self): + error_msg = "`fix_cov` must be symmetric positive semidefinite." + with pytest.raises(ValueError, match=error_msg): + fix_cov = np.array([[1., 0.], [0., -1.]]) + multivariate_normal.fit(np.eye(2), fix_cov=fix_cov) + + def test_fit_fix_mean(self): + rng = np.random.default_rng(4385269356937404) + loc = rng.random(3) + A = rng.random((3, 3)) + cov = np.dot(A, A.T) + samples = multivariate_normal.rvs(mean=loc, cov=cov, size=100, + random_state=rng) + mean_free, cov_free = multivariate_normal.fit(samples) + logp_free = multivariate_normal.logpdf(samples, mean=mean_free, + cov=cov_free).sum() + mean_fix, cov_fix = multivariate_normal.fit(samples, fix_mean=loc) + assert_equal(mean_fix, loc) + logp_fix = multivariate_normal.logpdf(samples, mean=mean_fix, + cov=cov_fix).sum() + # test that fixed parameters result in lower likelihood than free + # parameters + assert logp_fix < logp_free + # test that a small perturbation of the resulting parameters + # has lower likelihood than the estimated parameters + A = rng.random((3, 3)) + m = 1e-8 * np.dot(A, A.T) + cov_perturbed = cov_fix + m + logp_perturbed = (multivariate_normal.logpdf(samples, + mean=mean_fix, + cov=cov_perturbed) + ).sum() + assert logp_perturbed < logp_fix + + def test_fit_fix_cov(self): + rng = np.random.default_rng(4385269356937404) + loc = rng.random(3) + A = rng.random((3, 3)) + cov = np.dot(A, A.T) + samples = multivariate_normal.rvs(mean=loc, cov=cov, + size=100, random_state=rng) + mean_free, cov_free = multivariate_normal.fit(samples) + logp_free = multivariate_normal.logpdf(samples, mean=mean_free, + cov=cov_free).sum() + mean_fix, cov_fix = multivariate_normal.fit(samples, fix_cov=cov) + assert_equal(mean_fix, np.mean(samples, axis=0)) + assert_equal(cov_fix, cov) + logp_fix = multivariate_normal.logpdf(samples, mean=mean_fix, + cov=cov_fix).sum() + # test that fixed parameters result in lower likelihood than free + # parameters + assert logp_fix < logp_free + # test that a small perturbation of the resulting parameters + # has lower likelihood than the estimated parameters + mean_perturbed = mean_fix + 1e-8 * rng.random(3) + logp_perturbed = (multivariate_normal.logpdf(samples, + mean=mean_perturbed, + cov=cov_fix) + ).sum() + assert logp_perturbed < logp_fix + + +class TestMarginal: + @pytest.mark.parametrize('dist,kwargs', [(multivariate_normal, {}), + (multivariate_t, {'df': 4})]) + @pytest.mark.parametrize('X_ndim', [3]) + @pytest.mark.parametrize('dimensions', [[1], [-1, 1]]) + @pytest.mark.parametrize('frozen', [True, False]) + @pytest.mark.parametrize('cov_object', [True, False]) + def test_marginal_distribution(self, dist, X_ndim, dimensions, frozen, + cov_object, kwargs): + rng = np.random.default_rng(413911473) + loc = rng.standard_normal(X_ndim) + A = rng.standard_normal((X_ndim, X_ndim)) + scale = A @ A.T + + if cov_object and dist == multivariate_t: + pytest.skip('`multivariate_t` does not accept a `Covariance` object') + elif cov_object: + scale = _covariance.CovViaPrecision(scale) + + # number of points at which to evaluate marginal PDF + x = np.random.standard_normal((4, len(dimensions))) + X = dist(loc, scale, **kwargs) + + if frozen: + Y = X.marginal(dimensions) + res = Y.pdf(x) + else: + Y = dist.marginal(dimensions, loc, scale, **kwargs) + res = Y.pdf(x) + + ref = marginal_pdf(X, X_ndim, dimensions, x) + assert_allclose(ref, res) + + @pytest.mark.parametrize('dist', [multivariate_normal, multivariate_t]) + def test_marginal_input_validation(self, dist): + rng = np.random.default_rng(413911473) + mean = rng.standard_normal(3) + A = rng.standard_normal((3, 3)) + cov = A @ A.T + + X = dist(mean, cov) + + msg = r"Dimensions \[3\] are invalid .*" + with pytest.raises(ValueError, match=msg): + X.marginal(3) + + with pytest.raises(ValueError, match=msg): + X.marginal([0, 1, 2, 3]) + + msg = r"All elements of `dimensions` must be unique." + with pytest.raises(ValueError, match=msg): + X.marginal([2, -1]) + + with pytest.raises(ValueError, match=msg): + X.marginal([[0, 1]]) + + msg = r"Elements of `dimensions` must be integers." + with pytest.raises(ValueError, match=msg): + X.marginal([1.1, 2.0]) + + @pytest.mark.parametrize('dist', [multivariate_normal, multivariate_t]) + def test_marginal_special_cases(self, dist): + rng = np.random.default_rng(413911473) + loc = rng.standard_normal(3) + A = rng.standard_normal((3, 3)) + scale = A @ A.T + + X = dist(loc, scale) + + msg = r"Cannot marginalize all dimensions." + with pytest.raises(ValueError, match=msg): + X.marginal([]) + + +class TestMatrixNormal: + + def test_bad_input(self): + # Check that bad inputs raise errors + num_rows = 4 + num_cols = 3 + M = np.full((num_rows,num_cols), 0.3) + U = 0.5 * np.identity(num_rows) + np.full((num_rows, num_rows), 0.5) + V = 0.7 * np.identity(num_cols) + np.full((num_cols, num_cols), 0.3) + + # Incorrect dimensions + assert_raises(ValueError, matrix_normal, np.zeros((5,4,3))) + assert_raises(ValueError, matrix_normal, M, np.zeros(10), V) + assert_raises(ValueError, matrix_normal, M, U, np.zeros(10)) + assert_raises(ValueError, matrix_normal, M, U, U) + assert_raises(ValueError, matrix_normal, M, V, V) + assert_raises(ValueError, matrix_normal, M.T, U, V) + + e = np.linalg.LinAlgError + # Singular covariance for the rvs method of a non-frozen instance + assert_raises(e, matrix_normal.rvs, + M, U, np.ones((num_cols, num_cols))) + assert_raises(e, matrix_normal.rvs, + M, np.ones((num_rows, num_rows)), V) + # Singular covariance for a frozen instance + assert_raises(e, matrix_normal, M, U, np.ones((num_cols, num_cols))) + assert_raises(e, matrix_normal, M, np.ones((num_rows, num_rows)), V) + + def test_default_inputs(self): + # Check that default argument handling works + num_rows = 4 + num_cols = 3 + M = np.full((num_rows,num_cols), 0.3) + U = 0.5 * np.identity(num_rows) + np.full((num_rows, num_rows), 0.5) + V = 0.7 * np.identity(num_cols) + np.full((num_cols, num_cols), 0.3) + Z = np.zeros((num_rows, num_cols)) + Zr = np.zeros((num_rows, 1)) + Zc = np.zeros((1, num_cols)) + Ir = np.identity(num_rows) + Ic = np.identity(num_cols) + I1 = np.identity(1) + + assert_equal(matrix_normal.rvs(mean=M, rowcov=U, colcov=V).shape, + (num_rows, num_cols)) + assert_equal(matrix_normal.rvs(mean=M).shape, + (num_rows, num_cols)) + assert_equal(matrix_normal.rvs(rowcov=U).shape, + (num_rows, 1)) + assert_equal(matrix_normal.rvs(colcov=V).shape, + (1, num_cols)) + assert_equal(matrix_normal.rvs(mean=M, colcov=V).shape, + (num_rows, num_cols)) + assert_equal(matrix_normal.rvs(mean=M, rowcov=U).shape, + (num_rows, num_cols)) + assert_equal(matrix_normal.rvs(rowcov=U, colcov=V).shape, + (num_rows, num_cols)) + + assert_equal(matrix_normal(mean=M).rowcov, Ir) + assert_equal(matrix_normal(mean=M).colcov, Ic) + assert_equal(matrix_normal(rowcov=U).mean, Zr) + assert_equal(matrix_normal(rowcov=U).colcov, I1) + assert_equal(matrix_normal(colcov=V).mean, Zc) + assert_equal(matrix_normal(colcov=V).rowcov, I1) + assert_equal(matrix_normal(mean=M, rowcov=U).colcov, Ic) + assert_equal(matrix_normal(mean=M, colcov=V).rowcov, Ir) + assert_equal(matrix_normal(rowcov=U, colcov=V).mean, Z) + + def test_covariance_expansion(self): + # Check that covariance can be specified with scalar or vector + num_rows = 4 + num_cols = 3 + M = np.full((num_rows, num_cols), 0.3) + Uv = np.full(num_rows, 0.2) + Us = 0.2 + Vv = np.full(num_cols, 0.1) + Vs = 0.1 + + Ir = np.identity(num_rows) + Ic = np.identity(num_cols) + + assert_equal(matrix_normal(mean=M, rowcov=Uv, colcov=Vv).rowcov, + 0.2*Ir) + assert_equal(matrix_normal(mean=M, rowcov=Uv, colcov=Vv).colcov, + 0.1*Ic) + assert_equal(matrix_normal(mean=M, rowcov=Us, colcov=Vs).rowcov, + 0.2*Ir) + assert_equal(matrix_normal(mean=M, rowcov=Us, colcov=Vs).colcov, + 0.1*Ic) + + def test_frozen_matrix_normal(self): + for i in range(1,5): + for j in range(1,5): + M = np.full((i,j), 0.3) + U = 0.5 * np.identity(i) + np.full((i,i), 0.5) + V = 0.7 * np.identity(j) + np.full((j,j), 0.3) + + frozen = matrix_normal(mean=M, rowcov=U, colcov=V) + + rvs1 = frozen.rvs(random_state=1234) + rvs2 = matrix_normal.rvs(mean=M, rowcov=U, colcov=V, + random_state=1234) + assert_equal(rvs1, rvs2) + + X = frozen.rvs(random_state=1234) + + pdf1 = frozen.pdf(X) + pdf2 = matrix_normal.pdf(X, mean=M, rowcov=U, colcov=V) + assert_equal(pdf1, pdf2) + + logpdf1 = frozen.logpdf(X) + logpdf2 = matrix_normal.logpdf(X, mean=M, rowcov=U, colcov=V) + assert_equal(logpdf1, logpdf2) + + def test_matches_multivariate(self): + # Check that the pdfs match those obtained by vectorising and + # treating as a multivariate normal. + for i in range(1,5): + for j in range(1,5): + M = np.full((i,j), 0.3) + U = 0.5 * np.identity(i) + np.full((i,i), 0.5) + V = 0.7 * np.identity(j) + np.full((j,j), 0.3) + + frozen = matrix_normal(mean=M, rowcov=U, colcov=V) + X = frozen.rvs(random_state=1234) + pdf1 = frozen.pdf(X) + logpdf1 = frozen.logpdf(X) + entropy1 = frozen.entropy() + + vecX = X.T.flatten() + vecM = M.T.flatten() + cov = np.kron(V,U) + pdf2 = multivariate_normal.pdf(vecX, mean=vecM, cov=cov) + logpdf2 = multivariate_normal.logpdf(vecX, mean=vecM, cov=cov) + entropy2 = multivariate_normal.entropy(mean=vecM, cov=cov) + + assert_allclose(pdf1, pdf2, rtol=1E-10) + assert_allclose(logpdf1, logpdf2, rtol=1E-10) + assert_allclose(entropy1, entropy2) + + def test_array_input(self): + # Check array of inputs has the same output as the separate entries. + num_rows = 4 + num_cols = 3 + M = np.full((num_rows,num_cols), 0.3) + U = 0.5 * np.identity(num_rows) + np.full((num_rows, num_rows), 0.5) + V = 0.7 * np.identity(num_cols) + np.full((num_cols, num_cols), 0.3) + N = 10 + + frozen = matrix_normal(mean=M, rowcov=U, colcov=V) + X1 = frozen.rvs(size=N, random_state=1234) + X2 = frozen.rvs(size=N, random_state=4321) + X = np.concatenate((X1[np.newaxis,:,:,:],X2[np.newaxis,:,:,:]), axis=0) + assert_equal(X.shape, (2, N, num_rows, num_cols)) + + array_logpdf = frozen.logpdf(X) + assert_equal(array_logpdf.shape, (2, N)) + for i in range(2): + for j in range(N): + separate_logpdf = matrix_normal.logpdf(X[i,j], mean=M, + rowcov=U, colcov=V) + assert_allclose(separate_logpdf, array_logpdf[i,j], 1E-10) + + def test_moments(self): + # Check that the sample moments match the parameters + num_rows = 4 + num_cols = 3 + M = np.full((num_rows,num_cols), 0.3) + U = 0.5 * np.identity(num_rows) + np.full((num_rows, num_rows), 0.5) + V = 0.7 * np.identity(num_cols) + np.full((num_cols, num_cols), 0.3) + N = 1000 + + frozen = matrix_normal(mean=M, rowcov=U, colcov=V) + X = frozen.rvs(size=N, random_state=1234) + + sample_mean = np.mean(X,axis=0) + assert_allclose(sample_mean, M, atol=0.1) + + sample_colcov = np.cov(X.reshape(N*num_rows,num_cols).T) + assert_allclose(sample_colcov, V, atol=0.1) + + sample_rowcov = np.cov(np.swapaxes(X,1,2).reshape( + N*num_cols,num_rows).T) + assert_allclose(sample_rowcov, U, atol=0.1) + + def test_samples(self): + # Regression test to ensure that we always generate the same stream of + # random variates. + actual = matrix_normal.rvs( + mean=np.array([[1, 2], [3, 4]]), + rowcov=np.array([[4, -1], [-1, 2]]), + colcov=np.array([[5, 1], [1, 10]]), + random_state=np.random.default_rng(0), + size=2 + ) + expected = np.array( + [[[1.56228264238181, -1.24136424071189], + [2.46865788392114, 6.22964440489445]], + [[3.86405716144353, 10.73714311429529], + [2.59428444080606, 5.79987854490876]]] + ) + assert_allclose(actual, expected) + + +class TestMatrixT: + + def test_bad_input(self): + # Check that bad inputs raise errors + num_rows = 4 + num_cols = 3 + df = 5 + M = np.full((num_rows, num_cols), 0.3) + U = 0.5 * np.identity(num_rows) + np.full((num_rows, num_rows), 0.5) + V = 0.7 * np.identity(num_cols) + np.full((num_cols, num_cols), 0.3) + + # Nonpositive degrees of freedom + with pytest.raises(ValueError, match="Degrees of freedom must be positive."): + matrix_t(df=0) + + # Incorrect dimensions + with pytest.raises(ValueError, match="Array `mean` must be 2D."): + matrix_t(mean=np.zeros((5, 4, 3))) + + with pytest.raises(ValueError, match="Array `mean` has invalid shape."): + matrix_t(mean=np.zeros((4, 3, 0))) + + with pytest.raises(ValueError, match="Array `row_spread` has invalid shape."): + matrix_t(row_spread=np.ones((1, 0))) + + with pytest.raises( + ValueError, match="Array `row_spread` must be a scalar or a 2D array." + ): + matrix_t(row_spread=np.ones((1, 2, 3))) + + with pytest.raises(ValueError, match="Array `row_spread` must be square."): + matrix_t(row_spread=np.ones((1, 2))) + + with pytest.raises(ValueError, match="Array `col_spread` has invalid shape."): + matrix_t(col_spread=np.ones((1, 0))) + + with pytest.raises( + ValueError, match="Array `col_spread` must be a scalar or a 2D array." + ): + matrix_t(col_spread=np.ones((1, 2, 3))) + + with pytest.raises(ValueError, match="Array `col_spread` must be square."): + matrix_t(col_spread=np.ones((1, 2))) + + with pytest.raises( + ValueError, + match="Arrays `mean` and `row_spread` must have the same number " + "of rows.", + ): + matrix_t(mean=M, row_spread=V) + + with pytest.raises( + ValueError, + match="Arrays `mean` and `col_spread` must have the same number " + "of columns.", + ): + matrix_t(mean=M, col_spread=U) + + # Incorrect dimension of input matrix + with pytest.raises( + ValueError, + match="The shape of array `X` is not conformal with " + "the distribution parameters.", + ): + matrix_t.pdf(X=np.zeros((num_rows, num_rows)), mean=M) + + # Singular covariance for a non-frozen instance + + with pytest.raises( + np.linalg.LinAlgError, + match="Internal potrf return info", + ): + matrix_t.rvs(M, U, np.ones((num_cols, num_cols)), df) + + with pytest.raises( + np.linalg.LinAlgError, + match="Internal potrf return info", + ): + matrix_t.rvs(M, np.ones((num_rows, num_rows)), V, df) + + # Singular covariance for a frozen instance + + with pytest.raises( + np.linalg.LinAlgError, + match="When `allow_singular is False`, the input matrix must be " + "symmetric positive definite.", + ): + matrix_t(M, U, np.ones((num_cols, num_cols)), df) + + with pytest.raises( + np.linalg.LinAlgError, + match="When `allow_singular is False`, the input matrix must be " + "symmetric positive definite.", + ): + matrix_t(M, np.ones((num_rows, num_rows)), V, df) + + def test_default_inputs(self): + # Check that default argument handling works + num_rows = 4 + num_cols = 3 + df = 5 + M = np.full((num_rows, num_cols), 0.3) + U = 0.5 * np.identity(num_rows) + np.full((num_rows, num_rows), 0.5) + V = 0.7 * np.identity(num_cols) + np.full((num_cols, num_cols), 0.3) + Z = np.zeros((num_rows, num_cols)) + Zr = np.zeros((num_rows, 1)) + Zc = np.zeros((1, num_cols)) + Ir = np.identity(num_rows) + Ic = np.identity(num_cols) + I1 = np.identity(1) + dfdefault = 1 + + assert_equal( + matrix_t.rvs(mean=M, row_spread=U, col_spread=V, df=df).shape, + (num_rows, num_cols), + ) + assert_equal(matrix_t.rvs(mean=M).shape, (num_rows, num_cols)) + assert_equal(matrix_t.rvs(row_spread=U).shape, (num_rows, 1)) + assert_equal(matrix_t.rvs(col_spread=V).shape, (1, num_cols)) + assert_equal(matrix_t.rvs(mean=M, col_spread=V).shape, (num_rows, num_cols)) + assert_equal(matrix_t.rvs(mean=M, row_spread=U).shape, (num_rows, num_cols)) + assert_equal( + matrix_t.rvs(row_spread=U, col_spread=V).shape, (num_rows, num_cols) + ) + + assert_equal(matrix_t().df, dfdefault) + assert_equal(matrix_t(mean=M).row_spread, Ir) + assert_equal(matrix_t(mean=M).col_spread, Ic) + assert_equal(matrix_t(row_spread=U).mean, Zr) + assert_equal(matrix_t(row_spread=U).col_spread, I1) + assert_equal(matrix_t(col_spread=V).mean, Zc) + assert_equal(matrix_t(col_spread=V).row_spread, I1) + assert_equal(matrix_t(mean=M, row_spread=U).col_spread, Ic) + assert_equal(matrix_t(mean=M, col_spread=V).row_spread, Ir) + assert_equal(matrix_t(row_spread=U, col_spread=V, df=df).mean, Z) + + def test_covariance_expansion(self): + # Check that covariance can be specified with scalar or vector + num_rows = 4 + num_cols = 3 + df = 1 + M = np.full((num_rows, num_cols), 0.3) + Uv = np.full(num_rows, 0.2) + Us = 0.2 + Vv = np.full(num_cols, 0.1) + Vs = 0.1 + + Ir = np.identity(num_rows) + Ic = np.identity(num_cols) + + assert_equal( + matrix_t(mean=M, row_spread=Uv, col_spread=Vv, df=df).row_spread, 0.2 * Ir + ) + assert_equal( + matrix_t(mean=M, row_spread=Uv, col_spread=Vv, df=df).col_spread, 0.1 * Ic + ) + assert_equal( + matrix_t(mean=M, row_spread=Us, col_spread=Vs, df=df).row_spread, 0.2 * Ir + ) + assert_equal( + matrix_t(mean=M, row_spread=Us, col_spread=Vs, df=df).col_spread, 0.1 * Ic + ) + + @pytest.mark.parametrize("i", range(1, 4)) + @pytest.mark.parametrize("j", range(1, 4)) + def test_frozen_matrix_t(self, i, j): + + M = np.full((i, j), 0.3) + U = 0.5 * np.identity(i) + np.full((i, i), 0.5) + V = 0.7 * np.identity(j) + np.full((j, j), 0.3) + df = i + j + + frozen = matrix_t(mean=M, row_spread=U, col_spread=V, df=df) + + rvs1 = frozen.rvs(random_state=1234) + rvs2 = matrix_t.rvs( + mean=M, row_spread=U, col_spread=V, df=df, random_state=1234 + ) + assert_equal(rvs1, rvs2) + + X = frozen.rvs(random_state=1234) + + pdf1 = frozen.pdf(X) + pdf2 = matrix_t.pdf(X, mean=M, row_spread=U, col_spread=V, df=df) + assert_equal(pdf1, pdf2) + + logpdf1 = frozen.logpdf(X) + logpdf2 = matrix_t.logpdf(X, mean=M, row_spread=U, col_spread=V, df=df) + assert_equal(logpdf1, logpdf2) + + def test_array_input(self): + # Check array of inputs has the same output as the separate entries. + num_rows = 4 + num_cols = 3 + M = np.full((num_rows, num_cols), 0.3) + U = 0.5 * np.identity(num_rows) + np.full((num_rows, num_rows), 0.5) + V = 0.7 * np.identity(num_cols) + np.full((num_cols, num_cols), 0.3) + df = 1 + N = 10 + + frozen = matrix_t(mean=M, row_spread=U, col_spread=V, df=df) + X1 = frozen.rvs(size=N, random_state=1234) + X2 = frozen.rvs(size=N, random_state=4321) + X = np.concatenate((X1[np.newaxis, :, :, :], X2[np.newaxis, :, :, :]), axis=0) + assert_equal(X.shape, (2, N, num_rows, num_cols)) + + array_logpdf = frozen.logpdf(X) + logpdf_shape = array_logpdf.shape + assert_equal(logpdf_shape, (2, N)) + for i in range(2): + for j in range(N): + separate_logpdf = matrix_t.logpdf( + X[i, j], mean=M, row_spread=U, col_spread=V, df=df + ) + assert_allclose(separate_logpdf, array_logpdf[i, j], 1e-10) + + @staticmethod + def relative_error(vec1: np.ndarray, vec2: np.ndarray): + numerator = np.linalg.norm(vec1 - vec2) ** 2 + denominator = np.linalg.norm(vec1) ** 2 + np.linalg.norm(vec2) ** 2 + return numerator / denominator + + @staticmethod + def matrix_divergence(mat_true: np.ndarray, mat_est: np.ndarray) -> float: + mat_true_psd = _PSD(mat_true, allow_singular=False) + mat_est_psd = _PSD(mat_est, allow_singular=False) + if (np.exp(mat_est_psd.log_pdet) <= 0) or (np.exp(mat_true_psd.log_pdet) <= 0): + return np.inf + trace_term = np.trace(mat_est_psd.pinv @ mat_true) + log_detratio = mat_est_psd.log_pdet - mat_true_psd.log_pdet + return (trace_term + log_detratio - len(mat_true)) / 2 + + @staticmethod + def vec(a_mat: np.ndarray) -> np.ndarray: + """ + For an (m,n) array `a_mat` the output `vec(a_mat)` is an (m*n, 1) + array formed by stacking the columns of `a_mat` in the order in + which they occur in `a_mat`. + """ + assert a_mat.ndim == 2 + return a_mat.T.reshape((a_mat.size,)) + + def test_moments(self): + r""" + Gupta and Nagar (2000) Theorem 4.3.1 (p.135) + -------------------------------------------- + The covariance of the vectorized matrix variate t-distribution equals + $ (V \otimes U) / (\text{df} - 2)$, where $\otimes$ + denotes the usual Kronecker product. + """ + + df = 5 + num_rows = 4 + num_cols = 3 + M = np.full((num_rows, num_cols), 0.3) + U = 0.5 * np.identity(num_rows) + np.full((num_rows, num_rows), 0.5) + V = 0.7 * np.identity(num_cols) + np.full((num_cols, num_cols), 0.3) + N = 10**4 + atol = 1e-1 + + frozen = matrix_t(mean=M, row_spread=U, col_spread=V, df=df) + X = frozen.rvs(size=N, random_state=42) + + relerr = self.relative_error(M, X.mean(axis=0)) + assert_close(relerr, 0, atol=atol) + + cov_vec_true = np.kron(V, U) / (df - 2) + cov_vec_rvs = np.cov(np.array([self.vec(x) for x in X]), rowvar=False) + kl = self.matrix_divergence(cov_vec_true, cov_vec_rvs) + assert_close(kl, 0, atol=atol) + + def test_pdf_against_julia(self): + """ + Test values generated from Julia. + + Dockerfile + ---------- + FROM julia:1.11.5 + RUN julia -e 'using Pkg; Pkg.add("Distributions"); Pkg.add("PDMats")' + WORKDIR /usr/src + + Commands + -------- + using DelimitedFiles + using Distributions + using PDMats + using Random + Random.seed!(42) + ν = 5 + M = [1 2 3; 4 5 6] + Σ = PDMats.PDMat([1 0.5; 0.5 1]) + Ω = PDMats.PDMat([1 0.3 0.2; 0.3 1 0.4; 0.2 0.4 1]) + dist = MatrixTDist(ν, M, Σ, Ω) + samples = rand(dist, 10) + pdfs = [pdf(dist, s) for s in samples] + """ + + df = 5 + M = np.array([[1, 2, 3], [4, 5, 6]]) + U = np.array([[1, 0.5], [0.5, 1]]) + V = np.array([[1, 0.3, 0.2], [0.3, 1, 0.4], [0.2, 0.4, 1]]) + + rtol = 1e-10 + + samples_j = np.array( + [ + [ + [0.958884881003464, 2.328976673167312, 2.936195396714506], + [3.656388568544394, 5.677549814962506, 6.292509556719057] + ], + [ + [0.830992685140180, 2.588946865508210, 3.310327469315906], + [3.850637198786261, 5.106074165416971, 6.403143979925566] + ], + [ + [1.572053537500711, 1.760828063560249, 2.812123062636012], + [4.156334686390513, 5.075942019982631, 5.827004350136873] + ], + [ + [1.683810860278459, 2.801203900480317, 4.054517744825265], + [4.778239956376877, 5.070613721477604, 6.640349743267192] + ], + [ + [0.443183825511296, 2.072092271247398, 3.045385527559403], + [4.374387994815022, 5.083432151729137, 5.958013783940404] + ], + [ + [0.311591337218329, 1.162836182564980, 2.562167762547456], + [3.079154928756626, 4.202325496476140, 5.485839479663457] + ], + [ + [0.943713128785340, 1.923800464789872, 2.511941262351750], + [4.124882619205123, 4.889406461458511, 5.689675454116582] + ], + [ + [1.487852512870631, 1.933859334657448, 2.681311906634522], + [4.124418827930267, 5.335204598518954, 5.988120342017037] + ], + [ + [1.002470749319751, 1.386785511789551, 2.890832331097640], + [4.372884362128993, 4.729718562700068, 6.732322315921552] + ], + [ + [1.421351511333299, 2.106946903600814, 2.654619331838720], + [4.188693248790616, 5.336439611284261, 5.279121290355546] + ] + ] + ) + + pdfs_j = np.array( + [ + 0.082798951655369, + 0.119993852401118, + 0.151969434727803, + 0.003620324481841, + 0.072538716346179, + 0.027002666410192, + 0.485180162388507, + 0.135740468069511, + 0.013619162593841, + 0.034813885519299 + ] + ) + + pdfs_py = matrix_t.pdf(samples_j, mean=M, row_spread=U, col_spread=V, df=df) + + assert_allclose(pdfs_j, pdfs_py, rtol=rtol) + + def test_pdf_against_mathematica(self): + """ + Test values generated from Mathematica 13.0.0 for Linux x86 (64-bit) + Release ID 13.0.0.0 (7522564, 2021120311723), Patch Level 0 + + mu={{1,2,3},{4,5,6}}; + sigma={{1,0.5},{0.5,1}}; + omega={{1,0.3,0.2},{0.3,1,0.4},{0.2,0.4,1}}; + df=5; + sampleSize=10; + SeedRandom[42]; + dist=MatrixTDistribution[mu,sigma,omega,df]; + samples=SetPrecision[RandomVariate[dist,sampleSize],15]; + pdfs=SetPrecision[PDF[dist,#]&/@samples,15]; + """ + + df = 5 + M = np.array([[1, 2, 3], [4, 5, 6]]) + U = np.array([[1, 0.5], [0.5, 1]]) + V = np.array([[1, 0.3, 0.2], [0.3, 1, 0.4], [0.2, 0.4, 1]]) + + rtol = 1e-10 + + samples_m = np.array( + [ + [ + [0.639971699425374, 2.171718671534955, 2.575826093352771], + [4.031082477912233, 5.021680958526638, 6.268126154787008], + ], + [ + [1.164842884206232, 2.526297099993045, 3.781375229865069], + [3.912979114956833, 4.202714884504189, 5.661830748993523], + ], + [ + [1.00461853907369, 2.080028751298565, 3.406489485602410], + [3.993327716320432, 5.655909265966448, 6.578059791357837], + ], + [ + [0.80625209501374, 2.529009560674907, 2.807513313302189], + [3.722896768794995, 5.26987322525995, 5.801155613199776], + ], + [ + [0.445816208657817, 3.224059910964103, 2.954990980541423], + [3.451520519442941, 7.064424621385415, 5.438834195890955], + ], + [ + [0.919232769636664, 2.374572300756703, 3.495118928313048], + [3.924447237903237, 5.627654256287447, 5.806104608153957], + ], + [ + [2.014242004090113, 1.377018127709871, 3.114064311468686], + [3.88881648137925, 4.603482820518904, 5.714205489738063], + ], + [ + [1.322000147426889, 2.602135838377777, 2.558921028724319], + [4.50534702030683, 5.861137323151889, 5.181872548334852], + ], + [ + [1.448743656862261, 2.053847557652242, 3.637321543241769], + [4.097711403906707, 4.506916241403669, 5.68010653497977], + ], + [ + [1.045187318995198, 1.645467189679729, 3.284396214544507], + [3.648493466445393, 5.004212508553601, 6.301624351328048], + ], + ] + ) + + pdfs_m = np.array( + [ + 0.085671937131824, + 0.004821273644067, + 0.105978034029754, + 0.174250448808208, + 3.945711836053583e-05, + 0.027158790350349, + 0.00299095120309, + 0.005594546018078, + 0.025788366971310, + 0.120210733598845, + ] + ) + + pdfs_py = matrix_t.pdf(samples_m, mean=M, row_spread=U, col_spread=V, df=df) + + assert_allclose(pdfs_m, pdfs_py, rtol=rtol) + + def test_samples(self): + df = 5 + num_rows = 4 + num_cols = 3 + M = np.full((num_rows, num_cols), 0.3) + U = 0.5 * np.identity(num_rows) + np.full((num_rows, num_rows), 0.5) + V = 0.7 * np.identity(num_cols) + np.full((num_cols, num_cols), 0.3) + N = 10**4 + rtol = 0.05 + + # `rvs` performs Cholesky-inverse-Wishart sampling on the smaller + # dimension of `mean` + + frozen = matrix_t(mean=M, row_spread=U, col_spread=V, df=df) + X = frozen.rvs(size=N, random_state=42) # column-wise rvs + m = X.mean(0) + + frozenT = matrix_t(mean=M.T, row_spread=V, col_spread=U, df=df) + XT = frozenT.rvs(size=N, random_state=42) # row-wise rvs + mT = XT.mean(0) + + # Gupta and Nagar (2000) Theorem 4.3.3 (p.137) + # -------------------------------------------- + # If T follows a matrix variate t-distribution with mean M and row_spread U + # and col_spread V and df degrees of freedom, then its transpose T.T follows + # a matrix variate t-distribution with mean M.T and row_spread V and + # col_spread U and df degrees of freedom. + + assert_allclose(M, m, rtol=rtol) + assert_allclose(M.T, mT, rtol=rtol) + assert_allclose(m, mT.T, rtol=rtol) + assert_allclose(m.T, mT, rtol=rtol) + + @pytest.mark.parametrize("shape_case", ["row", "col"]) + def test_against_multivariate_t(self, shape_case): + r""" + Gupta and Nagar (2000) p.133f + When the number of rows or the number of columns equals 1 the + matrix t reduces to the multivariate t. But, the matrix t + is parameterized by raw 2nd moments whereas the multivariate t is + parameterized by a covariance (raw 2nd central moment normalized by df). + We can see the difference by comparing the author's notation + $t_p(n, \omega, \mathbf{\mu}, \Sigma)$ + for a matrix t with a single column + to the formula (4.1.2) for the PDF of the multivariate t. + """ + rtol = 1e-6 + df = 5 + + if shape_case == "row": + num_rows = 1 + num_cols = 3 + row_spread = 1 + col_spread = np.array([[1, 0.3, 0.2], [0.3, 1, 0.4], [0.2, 0.4, 1]]) + shape = col_spread / df + else: # shape_case == "col" + num_rows = 3 + num_cols = 1 + row_spread = np.array([[1, 0.3, 0.2], [0.3, 1, 0.4], [0.2, 0.4, 1]]) + col_spread=1 + shape = row_spread / df + + M = np.full((num_rows, num_cols), 0.3) + + t_mat = matrix_t( + mean=M, row_spread=row_spread, col_spread=col_spread, df=df + ) + t_mvt = multivariate_t(loc=M.squeeze(), shape=shape, df=df) + + X = t_mat.rvs(size=3, random_state=42) + t_mat_logpdf = t_mat.logpdf(X) + t_mvt_logpdf = t_mvt.logpdf(X.squeeze()) + assert_allclose(t_mvt_logpdf, t_mat_logpdf, rtol=rtol) + + +class TestDirichlet: + + def test_frozen_dirichlet(self): + rng = np.random.default_rng(2846) + + n = rng.integers(1, 32) + alpha = rng.uniform(10e-10, 100, n) + + d = dirichlet(alpha) + + assert_equal(d.var(), dirichlet.var(alpha)) + assert_equal(d.mean(), dirichlet.mean(alpha)) + assert_equal(d.entropy(), dirichlet.entropy(alpha)) + num_tests = 10 + for i in range(num_tests): + x = rng.uniform(10e-10, 100, n) + x /= np.sum(x) + assert_equal(d.pdf(x[:-1]), dirichlet.pdf(x[:-1], alpha)) + assert_equal(d.logpdf(x[:-1]), dirichlet.logpdf(x[:-1], alpha)) + + def test_numpy_rvs_shape_compatibility(self): + rng = np.random.default_rng(2846) + alpha = np.array([1.0, 2.0, 3.0]) + x = rng.dirichlet(alpha, size=7) + assert_equal(x.shape, (7, 3)) + assert_raises(ValueError, dirichlet.pdf, x, alpha) + assert_raises(ValueError, dirichlet.logpdf, x, alpha) + dirichlet.pdf(x.T, alpha) + dirichlet.pdf(x.T[:-1], alpha) + dirichlet.logpdf(x.T, alpha) + dirichlet.logpdf(x.T[:-1], alpha) + + def test_alpha_with_zeros(self): + rng = np.random.default_rng(2846) + alpha = [1.0, 0.0, 3.0] + # don't pass invalid alpha to np.random.dirichlet + x = rng.dirichlet(np.maximum(1e-9, alpha), size=7).T + assert_raises(ValueError, dirichlet.pdf, x, alpha) + assert_raises(ValueError, dirichlet.logpdf, x, alpha) + + def test_alpha_with_negative_entries(self): + rng = np.random.default_rng(2846) + alpha = [1.0, -2.0, 3.0] + # don't pass invalid alpha to np.random.dirichlet + x = rng.dirichlet(np.maximum(1e-9, alpha), size=7).T + assert_raises(ValueError, dirichlet.pdf, x, alpha) + assert_raises(ValueError, dirichlet.logpdf, x, alpha) + + def test_data_with_zeros(self): + alpha = np.array([1.0, 2.0, 3.0, 4.0]) + x = np.array([0.1, 0.0, 0.2, 0.7]) + dirichlet.pdf(x, alpha) + dirichlet.logpdf(x, alpha) + alpha = np.array([1.0, 1.0, 1.0, 1.0]) + assert_almost_equal(dirichlet.pdf(x, alpha), 6) + assert_almost_equal(dirichlet.logpdf(x, alpha), np.log(6)) + + def test_data_with_zeros_and_small_alpha(self): + alpha = np.array([1.0, 0.5, 3.0, 4.0]) + x = np.array([0.1, 0.0, 0.2, 0.7]) + assert_raises(ValueError, dirichlet.pdf, x, alpha) + assert_raises(ValueError, dirichlet.logpdf, x, alpha) + + def test_data_with_negative_entries(self): + alpha = np.array([1.0, 2.0, 3.0, 4.0]) + x = np.array([0.1, -0.1, 0.3, 0.7]) + assert_raises(ValueError, dirichlet.pdf, x, alpha) + assert_raises(ValueError, dirichlet.logpdf, x, alpha) + + def test_data_with_too_large_entries(self): + alpha = np.array([1.0, 2.0, 3.0, 4.0]) + x = np.array([0.1, 1.1, 0.3, 0.7]) + assert_raises(ValueError, dirichlet.pdf, x, alpha) + assert_raises(ValueError, dirichlet.logpdf, x, alpha) + + def test_data_too_deep_c(self): + alpha = np.array([1.0, 2.0, 3.0]) + x = np.full((2, 7, 7), 1 / 14) + assert_raises(ValueError, dirichlet.pdf, x, alpha) + assert_raises(ValueError, dirichlet.logpdf, x, alpha) + + def test_alpha_too_deep(self): + alpha = np.array([[1.0, 2.0], [3.0, 4.0]]) + x = np.full((2, 2, 7), 1 / 4) + assert_raises(ValueError, dirichlet.pdf, x, alpha) + assert_raises(ValueError, dirichlet.logpdf, x, alpha) + + def test_alpha_correct_depth(self): + alpha = np.array([1.0, 2.0, 3.0]) + x = np.full((3, 7), 1 / 3) + dirichlet.pdf(x, alpha) + dirichlet.logpdf(x, alpha) + + def test_non_simplex_data(self): + alpha = np.array([1.0, 2.0, 3.0]) + x = np.full((3, 7), 1 / 2) + assert_raises(ValueError, dirichlet.pdf, x, alpha) + assert_raises(ValueError, dirichlet.logpdf, x, alpha) + + def test_data_vector_too_short(self): + alpha = np.array([1.0, 2.0, 3.0, 4.0]) + x = np.full((2, 7), 1 / 2) + assert_raises(ValueError, dirichlet.pdf, x, alpha) + assert_raises(ValueError, dirichlet.logpdf, x, alpha) + + def test_data_vector_too_long(self): + alpha = np.array([1.0, 2.0, 3.0, 4.0]) + x = np.full((5, 7), 1 / 5) + assert_raises(ValueError, dirichlet.pdf, x, alpha) + assert_raises(ValueError, dirichlet.logpdf, x, alpha) + + def test_mean_var_cov(self): + # Reference values calculated by hand and confirmed with Mathematica, e.g. + # `Covariance[DirichletDistribution[{ 1, 0.8, 0.2, 10^-300}]]` + alpha = np.array([1., 0.8, 0.2]) + d = dirichlet(alpha) + + expected_mean = [0.5, 0.4, 0.1] + expected_var = [1. / 12., 0.08, 0.03] + expected_cov = [ + [ 1. / 12, -1. / 15, -1. / 60], + [-1. / 15, 2. / 25, -1. / 75], + [-1. / 60, -1. / 75, 3. / 100], + ] + + assert_array_almost_equal(d.mean(), expected_mean) + assert_array_almost_equal(d.var(), expected_var) + assert_array_almost_equal(d.cov(), expected_cov) + + def test_scalar_values(self): + alpha = np.array([0.2]) + d = dirichlet(alpha) + + # For alpha of length 1, mean and var should be scalar instead of array + assert_equal(d.mean().ndim, 0) + assert_equal(d.var().ndim, 0) + + assert_equal(d.pdf([1.]).ndim, 0) + assert_equal(d.logpdf([1.]).ndim, 0) + + def test_K_and_K_minus_1_calls_equal(self): + # Test that calls with K and K-1 entries yield the same results. + rng = np.random.default_rng(2846) + + n = rng.integers(1, 32) + alpha = rng.uniform(10e-10, 100, n) + + d = dirichlet(alpha) + num_tests = 10 + for i in range(num_tests): + x = rng.uniform(10e-10, 100, n) + x /= np.sum(x) + assert_almost_equal(d.pdf(x[:-1]), d.pdf(x)) + + def test_multiple_entry_calls(self): + # Test that calls with multiple x vectors as matrix work + rng = np.random.default_rng(2846) + + n = rng.integers(1, 32) + alpha = rng.uniform(10e-10, 100, n) + d = dirichlet(alpha) + + num_tests = 10 + num_multiple = 5 + xm = None + for i in range(num_tests): + for m in range(num_multiple): + x = rng.uniform(10e-10, 100, n) + x /= np.sum(x) + if xm is not None: + xm = np.vstack((xm, x)) + else: + xm = x + rm = d.pdf(xm.T) + rs = None + for xs in xm: + r = d.pdf(xs) + if rs is not None: + rs = np.append(rs, r) + else: + rs = r + assert_array_almost_equal(rm, rs) + + def test_2D_dirichlet_is_beta(self): + rng = np.random.default_rng(2846) + + alpha = rng.uniform(10e-10, 100, 2) + d = dirichlet(alpha) + b = beta(alpha[0], alpha[1]) + + num_tests = 10 + for i in range(num_tests): + x = rng.uniform(10e-10, 100, 2) + x /= np.sum(x) + assert_almost_equal(b.pdf(x), d.pdf([x])) + + assert_almost_equal(b.mean(), d.mean()[0]) + assert_almost_equal(b.var(), d.var()[0]) + + +def test_multivariate_normal_dimensions_mismatch(): + # Regression test for GH #3493. Check that setting up a PDF with a mean of + # length M and a covariance matrix of size (N, N), where M != N, raises a + # ValueError with an informative error message. + mu = np.array([0.0, 0.0]) + sigma = np.array([[1.0]]) + + assert_raises(ValueError, multivariate_normal, mu, sigma) + + # A simple check that the right error message was passed along. Checking + # that the entire message is there, word for word, would be somewhat + # fragile, so we just check for the leading part. + try: + multivariate_normal(mu, sigma) + except ValueError as e: + msg = "Dimension mismatch" + assert_equal(str(e)[:len(msg)], msg) + + +class TestWishart: + def test_scale_dimensions(self): + # Test that we can call the Wishart with various scale dimensions + + # Test case: dim=1, scale=1 + true_scale = np.array(1, ndmin=2) + scales = [ + 1, # scalar + [1], # iterable + np.array(1), # 0-dim + np.r_[1], # 1-dim + np.array(1, ndmin=2) # 2-dim + ] + for scale in scales: + w = wishart(1, scale) + assert_equal(w.scale, true_scale) + assert_equal(w.scale.shape, true_scale.shape) + + # Test case: dim=2, scale=[[1,0] + # [0,2] + true_scale = np.array([[1,0], + [0,2]]) + scales = [ + [1,2], # iterable + np.r_[1,2], # 1-dim + np.array([[1,0], # 2-dim + [0,2]]) + ] + for scale in scales: + w = wishart(2, scale) + assert_equal(w.scale, true_scale) + assert_equal(w.scale.shape, true_scale.shape) + + # We cannot call with a df < dim - 1 + assert_raises(ValueError, wishart, 1, np.eye(2)) + + # But we can call with dim - 1 < df < dim + wishart(1.1, np.eye(2)) # no error + # see gh-5562 + + # We cannot call with a 3-dimension array + scale = np.array(1, ndmin=3) + assert_raises(ValueError, wishart, 1, scale) + + def test_quantile_dimensions(self): + # Test that we can call the Wishart rvs with various quantile dimensions + + # If dim == 1, consider x.shape = [1,1,1] + X = [ + 1, # scalar + [1], # iterable + np.array(1), # 0-dim + np.r_[1], # 1-dim + np.array(1, ndmin=2), # 2-dim + np.array([1], ndmin=3) # 3-dim + ] + + w = wishart(1,1) + density = w.pdf(np.array(1, ndmin=3)) + for x in X: + assert_equal(w.pdf(x), density) + + # If dim == 1, consider x.shape = [1,1,*] + X = [ + [1,2,3], # iterable + np.r_[1,2,3], # 1-dim + np.array([1,2,3], ndmin=3) # 3-dim + ] + + w = wishart(1,1) + density = w.pdf(np.array([1,2,3], ndmin=3)) + for x in X: + assert_equal(w.pdf(x), density) + + # If dim == 2, consider x.shape = [2,2,1] + # where x[:,:,*] = np.eye(1)*2 + X = [ + 2, # scalar + [2,2], # iterable + np.array(2), # 0-dim + np.r_[2,2], # 1-dim + np.array([[2,0], + [0,2]]), # 2-dim + np.array([[2,0], + [0,2]])[:,:,np.newaxis] # 3-dim + ] + + w = wishart(2,np.eye(2)) + density = w.pdf(np.array([[2,0], + [0,2]])[:,:,np.newaxis]) + for x in X: + assert_equal(w.pdf(x), density) + + def test_frozen(self): + # Test that the frozen and non-frozen Wishart gives the same answers + + # Construct an arbitrary positive definite scale matrix + dim = 4 + scale = np.diag(np.arange(dim)+1) + scale[np.tril_indices(dim, k=-1)] = np.arange(dim * (dim-1) // 2) + scale = np.dot(scale.T, scale) + + # Construct a collection of positive definite matrices to test the PDF + X = [] + for i in range(5): + x = np.diag(np.arange(dim)+(i+1)**2) + x[np.tril_indices(dim, k=-1)] = np.arange(dim * (dim-1) // 2) + x = np.dot(x.T, x) + X.append(x) + X = np.array(X).T + + # Construct a 1D and 2D set of parameters + parameters = [ + (10, 1, np.linspace(0.1, 10, 5)), # 1D case + (10, scale, X) + ] + + for (df, scale, x) in parameters: + w = wishart(df, scale) + assert_equal(w.var(), wishart.var(df, scale)) + assert_equal(w.mean(), wishart.mean(df, scale)) + assert_equal(w.mode(), wishart.mode(df, scale)) + assert_equal(w.entropy(), wishart.entropy(df, scale)) + assert_equal(w.pdf(x), wishart.pdf(x, df, scale)) + + def test_wishart_2D_rvs(self): + dim = 3 + df = 10 + + # Construct a simple non-diagonal positive definite matrix + scale = np.eye(dim) + scale[0,1] = 0.5 + scale[1,0] = 0.5 + + # Construct frozen Wishart random variables + w = wishart(df, scale) + + # Get the generated random variables from a known seed + rng = np.random.RandomState(248042) + w_rvs = wishart.rvs(df, scale, random_state=rng) + rng = np.random.RandomState(248042) + frozen_w_rvs = w.rvs(random_state=rng) + + # Manually calculate what it should be, based on the Bartlett (1933) + # decomposition of a Wishart into D A A' D', where D is the Cholesky + # factorization of the scale matrix and A is the lower triangular matrix + # with the square root of chi^2 variates on the diagonal and N(0,1) + # variates in the lower triangle. + rng = np.random.RandomState(248042) + covariances = rng.normal(size=3) + variances = np.r_[ + rng.chisquare(df), + rng.chisquare(df-1), + rng.chisquare(df-2), + ]**0.5 + + # Construct the lower-triangular A matrix + A = np.diag(variances) + A[np.tril_indices(dim, k=-1)] = covariances + + # Wishart random variate + D = np.linalg.cholesky(scale) + DA = D.dot(A) + manual_w_rvs = np.dot(DA, DA.T) + + # Test for equality + assert_allclose(w_rvs, manual_w_rvs) + assert_allclose(frozen_w_rvs, manual_w_rvs) + + def test_1D_is_chisquared(self): + # The 1-dimensional Wishart with an identity scale matrix is just a + # chi-squared distribution. + # Test variance, mean, entropy, pdf + # Kolgomorov-Smirnov test for rvs + rng = np.random.default_rng(482974) + + sn = 500 + dim = 1 + scale = np.eye(dim) + + df_range = np.arange(1, 10, 2, dtype=float) + X = np.linspace(0.1,10,num=10) + for df in df_range: + w = wishart(df, scale) + c = chi2(df) + + # Statistics + assert_allclose(w.var(), c.var()) + assert_allclose(w.mean(), c.mean()) + assert_allclose(w.entropy(), c.entropy()) + + # PDF + assert_allclose(w.pdf(X), c.pdf(X)) + + # rvs + rvs = w.rvs(size=sn, random_state=rng) + args = (df,) + alpha = 0.01 + check_distribution_rvs('chi2', args, alpha, rvs) + + def test_is_scaled_chisquared(self): + # The 2-dimensional Wishart with an arbitrary scale matrix can be + # transformed to a scaled chi-squared distribution. + # For :math:`S \sim W_p(V,n)` and :math:`\lambda \in \mathbb{R}^p` we have + # :math:`\lambda' S \lambda \sim \lambda' V \lambda \times \chi^2(n)` + rng = np.random.default_rng(482974) + + sn = 500 + df = 10 + dim = 4 + # Construct an arbitrary positive definite matrix + scale = np.diag(np.arange(4)+1) + scale[np.tril_indices(4, k=-1)] = np.arange(6) + scale = np.dot(scale.T, scale) + # Use :math:`\lambda = [1, \dots, 1]'` + lamda = np.ones((dim,1)) + sigma_lamda = lamda.T.dot(scale).dot(lamda).squeeze() + w = wishart(df, sigma_lamda) + c = chi2(df, scale=sigma_lamda) + + # Statistics + assert_allclose(w.var(), c.var()) + assert_allclose(w.mean(), c.mean()) + assert_allclose(w.entropy(), c.entropy()) + + # PDF + X = np.linspace(0.1,10,num=10) + assert_allclose(w.pdf(X), c.pdf(X)) + + # rvs + rvs = w.rvs(size=sn, random_state=rng) + args = (df,0,sigma_lamda) + alpha = 0.01 + check_distribution_rvs('chi2', args, alpha, rvs) + +class TestMultinomial: + def test_logpmf(self): + vals1 = multinomial.logpmf((3,4), 7, (0.3, 0.7)) + assert_allclose(vals1, -1.483270127243324, rtol=1e-8) + + vals2 = multinomial.logpmf([3, 4], 0, [.3, .7]) + assert vals2 == -np.inf + + vals3 = multinomial.logpmf([0, 0], 0, [.3, .7]) + assert vals3 == 0 + + vals4 = multinomial.logpmf([3, 4], 0, [-2, 3]) + assert_allclose(vals4, np.nan, rtol=1e-8) + + def test_reduces_binomial(self): + # test that the multinomial pmf reduces to the binomial pmf in the 2d + # case + val1 = multinomial.logpmf((3, 4), 7, (0.3, 0.7)) + val2 = binom.logpmf(3, 7, 0.3) + assert_allclose(val1, val2, rtol=1e-8) + + val1 = multinomial.pmf((6, 8), 14, (0.1, 0.9)) + val2 = binom.pmf(6, 14, 0.1) + assert_allclose(val1, val2, rtol=1e-8) + + def test_R(self): + # test against the values produced by this R code + # (https://stat.ethz.ch/R-manual/R-devel/library/stats/html/Multinom.html) + # X <- t(as.matrix(expand.grid(0:3, 0:3))); X <- X[, colSums(X) <= 3] + # X <- rbind(X, 3:3 - colSums(X)); dimnames(X) <- list(letters[1:3], NULL) + # X + # apply(X, 2, function(x) dmultinom(x, prob = c(1,2,5))) + + n, p = 3, [1./8, 2./8, 5./8] + r_vals = {(0, 0, 3): 0.244140625, (1, 0, 2): 0.146484375, + (2, 0, 1): 0.029296875, (3, 0, 0): 0.001953125, + (0, 1, 2): 0.292968750, (1, 1, 1): 0.117187500, + (2, 1, 0): 0.011718750, (0, 2, 1): 0.117187500, + (1, 2, 0): 0.023437500, (0, 3, 0): 0.015625000} + for x in r_vals: + assert_allclose(multinomial.pmf(x, n, p), r_vals[x], atol=1e-14) + + @pytest.mark.parametrize("n", [0, 3]) + def test_rvs_np(self, n): + # test that .rvs agrees w/numpy when `p` is valid + rndm = np.random.RandomState(123) + sc_rvs = multinomial.rvs(n, [1/3.]*3, size=7, random_state=123) + np_rvs = rndm.multinomial(n, [1/3.]*3, size=7) + assert_equal(sc_rvs, np_rvs) + + rndm = np.random.RandomState(123) + sc_rvs = multinomial.rvs(n, [1/5.]*5, size=7, random_state=123) + np_rvs = rndm.multinomial(n, [1/5.]*5, size=7) + assert_equal(sc_rvs, np_rvs) + + def test_pmf(self): + vals0 = multinomial.pmf((5,), 5, (1,)) + assert_allclose(vals0, 1, rtol=1e-8) + + vals1 = multinomial.pmf((3,4), 7, (.3, .7)) + assert_allclose(vals1, .22689449999999994, rtol=1e-8) + + vals2 = multinomial.pmf([[[3,5],[0,8]], [[-1, 9], [1, 1]]], 8, + (.1, .9)) + assert_allclose(vals2, [[.03306744, .43046721], [0, 0]], rtol=1e-8) + + x = np.empty((0,2), dtype=np.float64) + vals3 = multinomial.pmf(x, 4, (.3, .7)) + assert_equal(vals3, np.empty([], dtype=np.float64)) + + vals4 = multinomial.pmf([1,2], 4, (.3, .7)) + assert_allclose(vals4, 0, rtol=1e-8) + + vals5 = multinomial.pmf([3, 3, 0], 6, [2/3.0, 1/3.0, 0]) + assert_allclose(vals5, 0.219478737997, rtol=1e-8) + + vals5 = multinomial.pmf([0, 0, 0], 0, [2/3.0, 1/3.0, 0]) + assert vals5 == 1 + + vals6 = multinomial.pmf([2, 1, 0], 0, [2/3.0, 1/3.0, 0]) + assert vals6 == 0 + + def test_pmf_broadcasting(self): + vals0 = multinomial.pmf([1, 2], 3, [[.1, .9], [.2, .8]]) + assert_allclose(vals0, [.243, .384], rtol=1e-8) + + vals1 = multinomial.pmf([1, 2], [3, 4], [.1, .9]) + assert_allclose(vals1, [.243, 0], rtol=1e-8) + + vals2 = multinomial.pmf([[[1, 2], [1, 1]]], 3, [.1, .9]) + assert_allclose(vals2, [[.243, 0]], rtol=1e-8) + + vals3 = multinomial.pmf([1, 2], [[[3], [4]]], [.1, .9]) + assert_allclose(vals3, [[[.243], [0]]], rtol=1e-8) + + vals4 = multinomial.pmf([[1, 2], [1,1]], [[[[3]]]], [.1, .9]) + assert_allclose(vals4, [[[[.243, 0]]]], rtol=1e-8) + + @pytest.mark.parametrize("n", [0, 5]) + def test_cov(self, n): + cov1 = multinomial.cov(n, (.2, .3, .5)) + cov2 = [[n*.2*.8, -n*.2*.3, -n*.2*.5], + [-n*.3*.2, n*.3*.7, -n*.3*.5], + [-n*.5*.2, -n*.5*.3, n*.5*.5]] + assert_allclose(cov1, cov2, rtol=1e-8) + + def test_cov_broadcasting(self): + cov1 = multinomial.cov(5, [[.1, .9], [.2, .8]]) + cov2 = [[[.45, -.45],[-.45, .45]], [[.8, -.8], [-.8, .8]]] + assert_allclose(cov1, cov2, rtol=1e-8) + + cov3 = multinomial.cov([4, 5], [.1, .9]) + cov4 = [[[.36, -.36], [-.36, .36]], [[.45, -.45], [-.45, .45]]] + assert_allclose(cov3, cov4, rtol=1e-8) + + cov5 = multinomial.cov([4, 5], [[.3, .7], [.4, .6]]) + cov6 = [[[4*.3*.7, -4*.3*.7], [-4*.3*.7, 4*.3*.7]], + [[5*.4*.6, -5*.4*.6], [-5*.4*.6, 5*.4*.6]]] + assert_allclose(cov5, cov6, rtol=1e-8) + + @pytest.mark.parametrize("n", [0, 2]) + def test_entropy(self, n): + # this is equivalent to a binomial distribution with n=2, so the + # entropy .77899774929 is easily computed "by hand" + ent0 = multinomial.entropy(n, [.2, .8]) + assert_allclose(ent0, binom.entropy(n, .2), rtol=1e-8) + + def test_entropy_broadcasting(self): + ent0 = multinomial.entropy([2, 3], [.2, .8]) + assert_allclose(ent0, [binom.entropy(2, .2), binom.entropy(3, .2)], + rtol=1e-8) + + ent1 = multinomial.entropy([7, 8], [[.3, .7], [.4, .6]]) + assert_allclose(ent1, [binom.entropy(7, .3), binom.entropy(8, .4)], + rtol=1e-8) + + ent2 = multinomial.entropy([[7], [8]], [[.3, .7], [.4, .6]]) + assert_allclose(ent2, + [[binom.entropy(7, .3), binom.entropy(7, .4)], + [binom.entropy(8, .3), binom.entropy(8, .4)]], + rtol=1e-8) + + @pytest.mark.parametrize("n", [0, 5]) + def test_mean(self, n): + mean1 = multinomial.mean(n, [.2, .8]) + assert_allclose(mean1, [n*.2, n*.8], rtol=1e-8) + + def test_mean_broadcasting(self): + mean1 = multinomial.mean([5, 6], [.2, .8]) + assert_allclose(mean1, [[5*.2, 5*.8], [6*.2, 6*.8]], rtol=1e-8) + + def test_frozen(self): + # The frozen distribution should agree with the regular one + n = 12 + pvals = (.1, .2, .3, .4) + x = [[0,0,0,12],[0,0,1,11],[0,1,1,10],[1,1,1,9],[1,1,2,8]] + x = np.asarray(x, dtype=np.float64) + mn_frozen = multinomial(n, pvals) + assert_allclose(mn_frozen.pmf(x), multinomial.pmf(x, n, pvals)) + assert_allclose(mn_frozen.logpmf(x), multinomial.logpmf(x, n, pvals)) + assert_allclose(mn_frozen.entropy(), multinomial.entropy(n, pvals)) + + def test_gh_11860(self): + # gh-11860 reported cases in which the adjustments made by multinomial + # to the last element of `p` can cause `nan`s even when the input is + # essentially valid. Check that a pathological case returns a finite, + # nonzero result. (This would fail in main before the PR.) + n = 88 + rng = np.random.default_rng(8879715917488330089) + p = rng.random(n) + p[-1] = 1e-30 + p /= np.sum(p) + x = np.ones(n) + logpmf = multinomial.logpmf(x, n, p) + assert np.isfinite(logpmf) + + @pytest.mark.parametrize('dtype', [np.float32, np.float64]) + def test_gh_22565(self, dtype): + # Same issue as gh-11860 above, essentially, but the original + # fix didn't completely solve the problem. + n = 19 + p = np.asarray([0.2, 0.2, 0.2, 0.2, 0.2], dtype=dtype) + res1 = multinomial.pmf(x=[1, 2, 5, 7, 4], n=n, p=p) + res2 = multinomial.pmf(x=[1, 2, 4, 5, 7], n=n, p=p) + np.testing.assert_allclose(res1, res2, rtol=1e-15) + + @pytest.mark.parametrize('psum', [0.75, 1.25]) + def test_gh_22585(self, psum): + # gh-22585 warned that beginning in SciPy 1.18, rows of p that did not sum to + # 1.0 would produce NaNs. Check that this has been implemented. + rng = np.random.default_rng(8879715917488330089) + d, n = 5, 19 # arbitrary + p0 = rng.random(d) + p = p0 / np.sum(p0) + p_ = p * psum + x = multinomial.rvs(n=n, p=p) + + assert np.all(np.isfinite(x)) + with pytest.raises(ValueError, match="`multinomial.rvs` requires..."): + multinomial.rvs(n=n, p=p_) + + assert np.isfinite(multinomial.pmf(x, n=n, p=p)) + assert np.isnan(multinomial.pmf(x, n=n, p=p_)) + + assert np.isfinite(multinomial.pmf(x, n=n, p=p)) + assert np.isnan(multinomial.pmf(x, n=n, p=p_)) + + assert np.isfinite(multinomial.entropy(n=n, p=p)) + assert np.isnan(multinomial.entropy(n=n, p=p_)) + + assert np.all(np.isfinite(multinomial.cov(n=n, p=p))) + assert np.all(np.isnan(multinomial.cov(n=n, p=p_))) + + +class TestInvwishart: + def test_frozen(self): + # Test that the frozen and non-frozen inverse Wishart gives the same + # answers + + # Construct an arbitrary positive definite scale matrix + dim = 4 + scale = np.diag(np.arange(dim)+1) + scale[np.tril_indices(dim, k=-1)] = np.arange(dim*(dim-1)/2) + scale = np.dot(scale.T, scale) + + # Construct a collection of positive definite matrices to test the PDF + X = [] + for i in range(5): + x = np.diag(np.arange(dim)+(i+1)**2) + x[np.tril_indices(dim, k=-1)] = np.arange(dim*(dim-1)/2) + x = np.dot(x.T, x) + X.append(x) + X = np.array(X).T + + # Construct a 1D and 2D set of parameters + parameters = [ + (10, 1, np.linspace(0.1, 10, 5)), # 1D case + (10, scale, X) + ] + + for (df, scale, x) in parameters: + iw = invwishart(df, scale) + assert_equal(iw.var(), invwishart.var(df, scale)) + assert_equal(iw.mean(), invwishart.mean(df, scale)) + assert_equal(iw.mode(), invwishart.mode(df, scale)) + assert_allclose(iw.pdf(x), invwishart.pdf(x, df, scale)) + + def test_1D_is_invgamma(self): + # The 1-dimensional inverse Wishart with an identity scale matrix is + # just an inverse gamma distribution. + # Test variance, mean, pdf, entropy + # Kolgomorov-Smirnov test for rvs + rng = np.random.RandomState(482974) + + sn = 500 + dim = 1 + scale = np.eye(dim) + + df_range = np.arange(5, 20, 2, dtype=float) + X = np.linspace(0.1,10,num=10) + for df in df_range: + iw = invwishart(df, scale) + ig = invgamma(df/2, scale=1./2) + + # Statistics + assert_allclose(iw.var(), ig.var()) + assert_allclose(iw.mean(), ig.mean()) + + # PDF + assert_allclose(iw.pdf(X), ig.pdf(X)) + + # rvs + rvs = iw.rvs(size=sn, random_state=rng) + args = (df/2, 0, 1./2) + alpha = 0.01 + check_distribution_rvs('invgamma', args, alpha, rvs) + + # entropy + assert_allclose(iw.entropy(), ig.entropy()) + + def test_invwishart_2D_rvs(self): + dim = 3 + df = 10 + + # Construct a simple non-diagonal positive definite matrix + scale = np.eye(dim) + scale[0,1] = 0.5 + scale[1,0] = 0.5 + + # Construct frozen inverse-Wishart random variables + iw = invwishart(df, scale) + + # Get the generated random variables from a known seed + rng = np.random.RandomState(608072) + iw_rvs = invwishart.rvs(df, scale, random_state=rng) + rng = np.random.RandomState(608072) + frozen_iw_rvs = iw.rvs(random_state=rng) + + # Manually calculate what it should be, based on the decomposition in + # https://arxiv.org/abs/2310.15884 of an invers-Wishart into L L', + # where L A = D, D is the Cholesky factorization of the scale matrix, + # and A is the lower triangular matrix with the square root of chi^2 + # variates on the diagonal and N(0,1) variates in the lower triangle. + # the diagonal chi^2 variates in this A are reversed compared to those + # in the Bartlett decomposition A for Wishart rvs. + rng = np.random.RandomState(608072) + covariances = rng.normal(size=3) + variances = np.r_[ + rng.chisquare(df-2), + rng.chisquare(df-1), + rng.chisquare(df), + ]**0.5 + + # Construct the lower-triangular A matrix + A = np.diag(variances) + A[np.tril_indices(dim, k=-1)] = covariances + + # inverse-Wishart random variate + D = np.linalg.cholesky(scale) + L = np.linalg.solve(A.T, D.T).T + manual_iw_rvs = np.dot(L, L.T) + + # Test for equality + assert_allclose(iw_rvs, manual_iw_rvs) + assert_allclose(frozen_iw_rvs, manual_iw_rvs) + + def test_sample_mean(self): + """Test that sample mean consistent with known mean.""" + # Construct an arbitrary positive definite scale matrix + df = 10 + sample_size = 20_000 + for dim in [1, 5]: + scale = np.diag(np.arange(dim) + 1) + scale[np.tril_indices(dim, k=-1)] = np.arange(dim * (dim - 1) / 2) + scale = np.dot(scale.T, scale) + + dist = invwishart(df, scale) + Xmean_exp = dist.mean() + Xvar_exp = dist.var() + Xmean_std = (Xvar_exp / sample_size)**0.5 # asymptotic SE of mean estimate + + X = dist.rvs(size=sample_size, random_state=1234) + Xmean_est = X.mean(axis=0) + + ntests = dim*(dim + 1)//2 + fail_rate = 0.01 / ntests # correct for multiple tests + max_diff = norm.ppf(1 - fail_rate / 2) + assert np.allclose( + (Xmean_est - Xmean_exp) / Xmean_std, + 0, + atol=max_diff, + ) + + def test_logpdf_4x4(self): + """Regression test for gh-8844.""" + X = np.array([[2, 1, 0, 0.5], + [1, 2, 0.5, 0.5], + [0, 0.5, 3, 1], + [0.5, 0.5, 1, 2]]) + Psi = np.array([[9, 7, 3, 1], + [7, 9, 5, 1], + [3, 5, 8, 2], + [1, 1, 2, 9]]) + nu = 6 + prob = invwishart.logpdf(X, nu, Psi) + # Explicit calculation from the formula on wikipedia. + p = X.shape[0] + sig, logdetX = np.linalg.slogdet(X) + sig, logdetPsi = np.linalg.slogdet(Psi) + M = np.linalg.solve(X, Psi) + expected = ((nu/2)*logdetPsi + - (nu*p/2)*np.log(2) + - multigammaln(nu/2, p) + - (nu + p + 1)/2*logdetX + - 0.5*M.trace()) + assert_allclose(prob, expected) + + +class TestSpecialOrthoGroup: + def test_reproducibility(self): + x = special_ortho_group.rvs(3, random_state=np.random.default_rng(514)) + expected = np.array([[-0.93200988, 0.01533561, -0.36210826], + [0.35742128, 0.20446501, -0.91128705], + [0.06006333, -0.97875374, -0.19604469]]) + assert_array_almost_equal(x, expected) + + def test_invalid_dim(self): + assert_raises(ValueError, special_ortho_group.rvs, None) + assert_raises(ValueError, special_ortho_group.rvs, (2, 2)) + assert_raises(ValueError, special_ortho_group.rvs, -1) + assert_raises(ValueError, special_ortho_group.rvs, 2.5) + + def test_frozen_matrix(self): + dim = 7 + frozen = special_ortho_group(dim) + + rvs1 = frozen.rvs(random_state=1234) + rvs2 = special_ortho_group.rvs(dim, random_state=1234) + + assert_equal(rvs1, rvs2) + + def test_det_and_ortho(self): + xs = [special_ortho_group.rvs(dim) + for dim in range(2,12) + for i in range(3)] + + # Test that determinants are always +1 + dets = [np.linalg.det(x) for x in xs] + assert_allclose(dets, [1.]*30, rtol=1e-13) + + # Test that these are orthogonal matrices + for x in xs: + assert_array_almost_equal(np.dot(x, x.T), + np.eye(x.shape[0])) + + def test_haar(self): + # Test that the distribution is constant under rotation + # Every column should have the same distribution + # Additionally, the distribution should be invariant under another rotation + + # Generate samples + dim = 5 + samples = 1000 # Not too many, or the test takes too long + ks_prob = .05 + xs = special_ortho_group.rvs( + dim, size=samples, random_state=np.random.default_rng(513) + ) + + # Dot a few rows (0, 1, 2) with unit vectors (0, 2, 4, 3), + # effectively picking off entries in the matrices of xs. + # These projections should all have the same distribution, + # establishing rotational invariance. We use the two-sided + # KS test to confirm this. + # We could instead test that angles between random vectors + # are uniformly distributed, but the below is sufficient. + # It is not feasible to consider all pairs, so pick a few. + els = ((0,0), (0,2), (1,4), (2,3)) + #proj = {(er, ec): [x[er][ec] for x in xs] for er, ec in els} + proj = {(er, ec): sorted([x[er][ec] for x in xs]) for er, ec in els} + pairs = [(e0, e1) for e0 in els for e1 in els if e0 > e1] + ks_tests = [ks_2samp(proj[p0], proj[p1])[1] for (p0, p1) in pairs] + assert_array_less([ks_prob]*len(pairs), ks_tests) + + def test_one_by_one(self): + # Test that the distribution is a delta function at the identity matrix + # when dim=1 + assert_allclose(special_ortho_group.rvs(1, size=1000), 1, rtol=1e-13) + + def test_zero_by_zero(self): + assert_equal(special_ortho_group.rvs(0, size=4).shape, (4, 0, 0)) + + +class TestOrthoGroup: + def test_reproducibility(self): + seed = 514 + rng = np.random.RandomState(seed) + x = ortho_group.rvs(3, random_state=rng) + x2 = ortho_group.rvs(3, random_state=seed) + # Note this matrix has det -1, distinguishing O(N) from SO(N) + assert_almost_equal(np.linalg.det(x), -1) + expected = np.array([[0.381686, -0.090374, 0.919863], + [0.905794, -0.161537, -0.391718], + [-0.183993, -0.98272, -0.020204]]) + assert_array_almost_equal(x, expected) + assert_array_almost_equal(x2, expected) + + def test_invalid_dim(self): + assert_raises(ValueError, ortho_group.rvs, None) + assert_raises(ValueError, ortho_group.rvs, (2, 2)) + assert_raises(ValueError, ortho_group.rvs, -1) + assert_raises(ValueError, ortho_group.rvs, 2.5) + + def test_frozen_matrix(self): + dim = 7 + frozen = ortho_group(dim) + frozen_seed = ortho_group(dim, seed=1234) + + rvs1 = frozen.rvs(random_state=1234) + rvs2 = ortho_group.rvs(dim, random_state=1234) + rvs3 = frozen_seed.rvs(size=1) + + assert_equal(rvs1, rvs2) + assert_equal(rvs1, rvs3) + + def test_det_and_ortho(self): + xs = [[ortho_group.rvs(dim) + for i in range(10)] + for dim in range(2,12)] + + # Test that abs determinants are always +1 + dets = np.array([[np.linalg.det(x) for x in xx] for xx in xs]) + assert_allclose(np.fabs(dets), np.ones(dets.shape), rtol=1e-13) + + # Test that these are orthogonal matrices + for xx in xs: + for x in xx: + assert_array_almost_equal(np.dot(x, x.T), + np.eye(x.shape[0])) + + @pytest.mark.parametrize("dim", [2, 5, 10, 20]) + def test_det_distribution_gh18272(self, dim): + # Test that positive and negative determinants are equally likely. + rng = np.random.default_rng(6796248956179332344) + dist = ortho_group(dim=dim) + rvs = dist.rvs(size=5000, random_state=rng) + dets = scipy.linalg.det(rvs) + k = np.sum(dets > 0) + n = len(dets) + res = stats.binomtest(k, n) + low, high = res.proportion_ci(confidence_level=0.95) + assert low < 0.5 < high + + def test_haar(self): + # Test that the distribution is constant under rotation + # Every column should have the same distribution + # Additionally, the distribution should be invariant under another rotation + + # Generate samples + dim = 5 + samples = 1000 # Not too many, or the test takes too long + ks_prob = .05 + rng = np.random.RandomState(518) # Note that the test is sensitive to seed too + xs = ortho_group.rvs(dim, size=samples, random_state=rng) + + # Dot a few rows (0, 1, 2) with unit vectors (0, 2, 4, 3), + # effectively picking off entries in the matrices of xs. + # These projections should all have the same distribution, + # establishing rotational invariance. We use the two-sided + # KS test to confirm this. + # We could instead test that angles between random vectors + # are uniformly distributed, but the below is sufficient. + # It is not feasible to consider all pairs, so pick a few. + els = ((0,0), (0,2), (1,4), (2,3)) + #proj = {(er, ec): [x[er][ec] for x in xs] for er, ec in els} + proj = {(er, ec): sorted([x[er][ec] for x in xs]) for er, ec in els} + pairs = [(e0, e1) for e0 in els for e1 in els if e0 > e1] + ks_tests = [ks_2samp(proj[p0], proj[p1])[1] for (p0, p1) in pairs] + assert_array_less([ks_prob]*len(pairs), ks_tests) + + def test_one_by_one(self): + # Test that the 1x1 distribution gives ±1 with equal probability. + dim = 1 + xs = ortho_group.rvs(dim, size=5000, random_state=np.random.default_rng(514)) + assert_allclose(np.abs(xs), 1, rtol=1e-13) + k = np.sum(xs > 0) + n = len(xs) + res = stats.binomtest(k, n) + low, high = res.proportion_ci(confidence_level=0.95) + assert low < 0.5 < high + + def test_zero_by_zero(self): + assert_equal(special_ortho_group.rvs(0, size=4).shape, (4, 0, 0)) + + @pytest.mark.slow + def test_pairwise_distances(self): + # Test that the distribution of pairwise distances is close to correct. + rng = np.random.RandomState(514) + + def random_ortho(dim, random_state): + u, _s, v = np.linalg.svd(random_state.normal(size=(dim, dim))) + return np.dot(u, v) + + for dim in range(2, 6): + def generate_test_statistics(rvs, N=1000, eps=1e-10): + stats = np.array([ + np.sum((rvs(dim=dim, random_state=rng) - + rvs(dim=dim, random_state=rng))**2) + for _ in range(N) + ]) + # Add a bit of noise to account for numeric accuracy. + stats += rng.uniform(-eps, eps, size=stats.shape) + return stats + + expected = generate_test_statistics(random_ortho) + actual = generate_test_statistics(scipy.stats.ortho_group.rvs) + + _D, p = scipy.stats.ks_2samp(expected, actual) + + assert_array_less(.05, p) + + +class TestRandomCorrelation: + def test_reproducibility(self): + rng = np.random.RandomState(514) + eigs = (.5, .8, 1.2, 1.5) + x = random_correlation.rvs(eigs, random_state=rng) + x2 = random_correlation.rvs(eigs, random_state=514) + expected = np.array([[1., -0.184851, 0.109017, -0.227494], + [-0.184851, 1., 0.231236, 0.326669], + [0.109017, 0.231236, 1., -0.178912], + [-0.227494, 0.326669, -0.178912, 1.]]) + assert_array_almost_equal(x, expected) + assert_array_almost_equal(x2, expected) + + def test_invalid_eigs(self): + assert_raises(ValueError, random_correlation.rvs, None) + assert_raises(ValueError, random_correlation.rvs, 'test') + assert_raises(ValueError, random_correlation.rvs, 2.5) + assert_raises(ValueError, random_correlation.rvs, [2.5]) + assert_raises(ValueError, random_correlation.rvs, [[1,2],[3,4]]) + assert_raises(ValueError, random_correlation.rvs, [2.5, -.5]) + assert_raises(ValueError, random_correlation.rvs, [1, 2, .1]) + + def test_frozen_matrix(self): + eigs = (.5, .8, 1.2, 1.5) + frozen = random_correlation(eigs) + frozen_seed = random_correlation(eigs, seed=514) + + rvs1 = random_correlation.rvs(eigs, random_state=514) + rvs2 = frozen.rvs(random_state=514) + rvs3 = frozen_seed.rvs() + + assert_equal(rvs1, rvs2) + assert_equal(rvs1, rvs3) + + def test_definition(self): + # Test the definition of a correlation matrix in several dimensions: + # + # 1. Det is product of eigenvalues (and positive by construction + # in examples) + # 2. 1's on diagonal + # 3. Matrix is symmetric + + def norm(i, e): + return i*e/sum(e) + + rng = np.random.RandomState(123) + + eigs = [norm(i, rng.uniform(size=i)) for i in range(2, 6)] + eigs.append([4,0,0,0]) + + ones = [[1.]*len(e) for e in eigs] + xs = [random_correlation.rvs(e, random_state=rng) for e in eigs] + + # Test that determinants are products of eigenvalues + # These are positive by construction + # Could also test that the eigenvalues themselves are correct, + # but this seems sufficient. + dets = [np.fabs(np.linalg.det(x)) for x in xs] + dets_known = [np.prod(e) for e in eigs] + assert_allclose(dets, dets_known, rtol=1e-13, atol=1e-13) + + # Test for 1's on the diagonal + diags = [np.diag(x) for x in xs] + for a, b in zip(diags, ones): + assert_allclose(a, b, rtol=1e-13) + + # Correlation matrices are symmetric + for x in xs: + assert_allclose(x, x.T, rtol=1e-13) + + def test_to_corr(self): + # Check some corner cases in to_corr + + # ajj == 1 + m = np.array([[0.1, 0], [0, 1]], dtype=float) + m = random_correlation._to_corr(m) + assert_allclose(m, np.array([[1, 0], [0, 0.1]])) + + # Floating point overflow; fails to compute the correct + # rotation, but should still produce some valid rotation + # rather than infs/nans + with np.errstate(over='ignore'): + g = np.array([[0, 1], [-1, 0]]) + + m0 = np.array([[1e300, 0], [0, np.nextafter(1, 0)]], dtype=float) + m = random_correlation._to_corr(m0.copy()) + assert_allclose(m, g.T.dot(m0).dot(g)) + + m0 = np.array([[0.9, 1e300], [1e300, 1.1]], dtype=float) + m = random_correlation._to_corr(m0.copy()) + assert_allclose(m, g.T.dot(m0).dot(g)) + + # Zero discriminant; should set the first diag entry to 1 + m0 = np.array([[2, 1], [1, 2]], dtype=float) + m = random_correlation._to_corr(m0.copy()) + assert_allclose(m[0,0], 1) + + # Slightly negative discriminant; should be approx correct still + m0 = np.array([[2 + 1e-7, 1], [1, 2]], dtype=float) + m = random_correlation._to_corr(m0.copy()) + assert_allclose(m[0,0], 1) + + +class TestUniformDirection: + @pytest.mark.parametrize("dim", [1, 3]) + @pytest.mark.parametrize("size", [None, 1, 5, (5, 4)]) + def test_samples(self, dim, size): + # test that samples have correct shape and norm 1 + rng = np.random.default_rng(2777937887058094419) + uniform_direction_dist = uniform_direction(dim, seed=rng) + samples = uniform_direction_dist.rvs(size) + mean, cov = np.zeros(dim), np.eye(dim) + expected_shape = rng.multivariate_normal(mean, cov, size=size).shape + assert samples.shape == expected_shape + norms = np.linalg.norm(samples, axis=-1) + assert_allclose(norms, 1.) + + @pytest.mark.parametrize("dim", [None, 0, (2, 2), 2.5]) + def test_invalid_dim(self, dim): + message = ("Dimension of vector must be specified, " + "and must be an integer greater than 0.") + with pytest.raises(ValueError, match=message): + uniform_direction.rvs(dim) + + def test_frozen_distribution(self): + dim = 5 + frozen = uniform_direction(dim) + frozen_seed = uniform_direction(dim, seed=514) + + rvs1 = frozen.rvs(random_state=514) + rvs2 = uniform_direction.rvs(dim, random_state=514) + rvs3 = frozen_seed.rvs() + + assert_equal(rvs1, rvs2) + assert_equal(rvs1, rvs3) + + @pytest.mark.parametrize("dim", [2, 5, 8]) + def test_uniform(self, dim): + rng = np.random.default_rng(1036978481269651776) + spherical_dist = uniform_direction(dim, seed=rng) + # generate random, orthogonal vectors + v1, v2 = spherical_dist.rvs(size=2) + v2 -= v1 @ v2 * v1 + v2 /= np.linalg.norm(v2) + assert_allclose(v1 @ v2, 0, atol=1e-14) # orthogonal + # generate data and project onto orthogonal vectors + samples = spherical_dist.rvs(size=10000) + s1 = samples @ v1 + s2 = samples @ v2 + angles = np.arctan2(s1, s2) + # test that angles follow a uniform distribution + # normalize angles to range [0, 1] + angles += np.pi + angles /= 2*np.pi + # perform KS test + uniform_dist = uniform() + kstest_result = kstest(angles, uniform_dist.cdf) + assert kstest_result.pvalue > 0.05 + + +class TestUnitaryGroup: + def test_reproducibility(self): + rng = np.random.RandomState(514) + x = unitary_group.rvs(3, random_state=rng) + x2 = unitary_group.rvs(3, random_state=514) + + expected = np.array( + [[0.308771+0.360312j, 0.044021+0.622082j, 0.160327+0.600173j], + [0.732757+0.297107j, 0.076692-0.4614j, -0.394349+0.022613j], + [-0.148844+0.357037j, -0.284602-0.557949j, 0.607051+0.299257j]] + ) + + assert_array_almost_equal(x, expected) + assert_array_almost_equal(x2, expected) + + def test_invalid_dim(self): + assert_raises(ValueError, unitary_group.rvs, None) + assert_raises(ValueError, unitary_group.rvs, (2, 2)) + assert_raises(ValueError, unitary_group.rvs, -1) + assert_raises(ValueError, unitary_group.rvs, 2.5) + + def test_frozen_matrix(self): + dim = 7 + frozen = unitary_group(dim) + frozen_seed = unitary_group(dim, seed=514) + + rvs1 = frozen.rvs(random_state=514) + rvs2 = unitary_group.rvs(dim, random_state=514) + rvs3 = frozen_seed.rvs(size=1) + + assert_equal(rvs1, rvs2) + assert_equal(rvs1, rvs3) + + def test_unitarity(self): + xs = [unitary_group.rvs(dim) + for dim in range(2,12) + for i in range(3)] + + # Test that these are unitary matrices + for x in xs: + assert_allclose(np.dot(x, x.conj().T), np.eye(x.shape[0]), atol=1e-15) + + def test_haar(self): + # Test that the eigenvalues, which lie on the unit circle in + # the complex plane, are uncorrelated. + + # Generate samples + for dim in (1, 5): + samples = 1000 # Not too many, or the test takes too long + # Note that the test is sensitive to seed too + xs = unitary_group.rvs( + dim, size=samples, random_state=np.random.default_rng(514) + ) + + # The angles "x" of the eigenvalues should be uniformly distributed + # Overall this seems to be a necessary but weak test of the distribution. + eigs = np.vstack([scipy.linalg.eigvals(x) for x in xs]) + x = np.arctan2(eigs.imag, eigs.real) + res = kstest(x.ravel(), uniform(-np.pi, 2*np.pi).cdf) + assert_(res.pvalue > 0.05) + + def test_zero_by_zero(self): + assert_equal(unitary_group.rvs(0, size=4).shape, (4, 0, 0)) + + +class TestMultivariateT: + + # These tests were created by running vpa(mvtpdf(...)) in MATLAB. The + # function takes no `mu` parameter. The tests were run as + # + # >> ans = vpa(mvtpdf(x - mu, shape, df)); + # + PDF_TESTS = [( + # x + [ + [1, 2], + [4, 1], + [2, 1], + [2, 4], + [1, 4], + [4, 1], + [3, 2], + [3, 3], + [4, 4], + [5, 1], + ], + # loc + [0, 0], + # shape + [ + [1, 0], + [0, 1] + ], + # df + 4, + # ans + [ + 0.013972450422333741737457302178882, + 0.0010998721906793330026219646100571, + 0.013972450422333741737457302178882, + 0.00073682844024025606101402363634634, + 0.0010998721906793330026219646100571, + 0.0010998721906793330026219646100571, + 0.0020732579600816823488240725481546, + 0.00095660371505271429414668515889275, + 0.00021831953784896498569831346792114, + 0.00037725616140301147447000396084604 + ] + + ), ( + # x + [ + [0.9718, 0.1298, 0.8134], + [0.4922, 0.5522, 0.7185], + [0.3010, 0.1491, 0.5008], + [0.5971, 0.2585, 0.8940], + [0.5434, 0.5287, 0.9507], + ], + # loc + [-1, 1, 50], + # shape + [ + [1.0000, 0.5000, 0.2500], + [0.5000, 1.0000, -0.1000], + [0.2500, -0.1000, 1.0000], + ], + # df + 8, + # ans + [ + 0.00000000000000069609279697467772867405511133763, + 0.00000000000000073700739052207366474839369535934, + 0.00000000000000069522909962669171512174435447027, + 0.00000000000000074212293557998314091880208889767, + 0.00000000000000077039675154022118593323030449058, + ] + )] + + @pytest.mark.parametrize("x, loc, shape, df, ans", PDF_TESTS) + def test_pdf_correctness(self, x, loc, shape, df, ans): + dist = multivariate_t(loc, shape, df, seed=0) + val = dist.pdf(x) + assert_array_almost_equal(val, ans) + + @pytest.mark.parametrize("x, loc, shape, df, ans", PDF_TESTS) + def test_logpdf_correct(self, x, loc, shape, df, ans): + dist = multivariate_t(loc, shape, df, seed=0) + val1 = dist.pdf(x) + val2 = dist.logpdf(x) + assert_array_almost_equal(np.log(val1), val2) + + # https://github.com/scipy/scipy/issues/10042#issuecomment-576795195 + def test_mvt_with_df_one_is_cauchy(self): + x = [9, 7, 4, 1, -3, 9, 0, -3, -1, 3] + val = multivariate_t.pdf(x, df=1) + ans = cauchy.pdf(x) + assert_array_almost_equal(val, ans) + + def test_mvt_with_high_df_is_approx_normal(self): + # `normaltest` returns the chi-squared statistic and the associated + # p-value. The null hypothesis is that `x` came from a normal + # distribution, so a low p-value represents rejecting the null, i.e. + # that it is unlikely that `x` came a normal distribution. + P_VAL_MIN = 0.1 + + dist = multivariate_t(0, 1, df=100000, seed=1) + samples = dist.rvs(size=100000) + _, p = normaltest(samples) + assert (p > P_VAL_MIN) + + dist = multivariate_t([-2, 3], [[10, -1], [-1, 10]], df=100000, + seed=42) + samples = dist.rvs(size=100000) + _, p = normaltest(samples) + assert ((p > P_VAL_MIN).all()) + + @pytest.mark.thread_unsafe(reason="uses mocking") + @patch('scipy.stats.multivariate_normal._logpdf') + def test_mvt_with_inf_df_calls_normal(self, mock): + dist = multivariate_t(0, 1, df=np.inf, seed=7) + assert isinstance(dist, multivariate_normal_frozen) + multivariate_t.pdf(0, df=np.inf) + assert mock.call_count == 1 + multivariate_t.logpdf(0, df=np.inf) + assert mock.call_count == 2 + + def test_shape_correctness(self): + # pdf and logpdf should return scalar when the + # number of samples in x is one. + dim = 4 + loc = np.zeros(dim) + shape = np.eye(dim) + df = 4.5 + x = np.zeros(dim) + res = multivariate_t(loc, shape, df).pdf(x) + assert np.isscalar(res) + res = multivariate_t(loc, shape, df).logpdf(x) + assert np.isscalar(res) + + # pdf() and logpdf() should return probabilities of shape + # (n_samples,) when x has n_samples. + n_samples = 7 + rng = np.random.default_rng(2767231913) + x = rng.random((n_samples, dim)) + res = multivariate_t(loc, shape, df).pdf(x) + assert (res.shape == (n_samples,)) + res = multivariate_t(loc, shape, df).logpdf(x) + assert (res.shape == (n_samples,)) + + # rvs() should return scalar unless a size argument is applied. + res = multivariate_t(np.zeros(1), np.eye(1), 1).rvs() + assert np.isscalar(res) + + # rvs() should return vector of shape (size,) if size argument + # is applied. + size = 7 + res = multivariate_t(np.zeros(1), np.eye(1), 1).rvs(size=size) + assert (res.shape == (size,)) + + def test_default_arguments(self): + dist = multivariate_t() + assert_equal(dist.loc, [0]) + assert_equal(dist.shape, [[1]]) + assert (dist.df == 1) + + DEFAULT_ARGS_TESTS = [ + (None, None, None, 0, 1, 1), + (None, None, 7, 0, 1, 7), + (None, [[7, 0], [0, 7]], None, [0, 0], [[7, 0], [0, 7]], 1), + (None, [[7, 0], [0, 7]], 7, [0, 0], [[7, 0], [0, 7]], 7), + ([7, 7], None, None, [7, 7], [[1, 0], [0, 1]], 1), + ([7, 7], None, 7, [7, 7], [[1, 0], [0, 1]], 7), + ([7, 7], [[7, 0], [0, 7]], None, [7, 7], [[7, 0], [0, 7]], 1), + ([7, 7], [[7, 0], [0, 7]], 7, [7, 7], [[7, 0], [0, 7]], 7) + ] + + @pytest.mark.parametrize("loc, shape, df, loc_ans, shape_ans, df_ans", + DEFAULT_ARGS_TESTS) + def test_default_args(self, loc, shape, df, loc_ans, shape_ans, df_ans): + dist = multivariate_t(loc=loc, shape=shape, df=df) + assert_equal(dist.loc, loc_ans) + assert_equal(dist.shape, shape_ans) + assert (dist.df == df_ans) + + ARGS_SHAPES_TESTS = [ + (-1, 2, 3, [-1], [[2]], 3), + ([-1], [2], 3, [-1], [[2]], 3), + (np.array([-1]), np.array([2]), 3, [-1], [[2]], 3) + ] + + @pytest.mark.parametrize("loc, shape, df, loc_ans, shape_ans, df_ans", + ARGS_SHAPES_TESTS) + def test_scalar_list_and_ndarray_arguments(self, loc, shape, df, loc_ans, + shape_ans, df_ans): + dist = multivariate_t(loc, shape, df) + assert_equal(dist.loc, loc_ans) + assert_equal(dist.shape, shape_ans) + assert_equal(dist.df, df_ans) + + def test_argument_error_handling(self): + # `loc` should be a one-dimensional vector. + loc = [[1, 1]] + assert_raises(ValueError, + multivariate_t, + **dict(loc=loc)) + + # `shape` should be scalar or square matrix. + shape = [[1, 1], [2, 2], [3, 3]] + assert_raises(ValueError, + multivariate_t, + **dict(loc=loc, shape=shape)) + + # `df` should be greater than zero. + loc = np.zeros(2) + shape = np.eye(2) + df = -1 + assert_raises(ValueError, + multivariate_t, + **dict(loc=loc, shape=shape, df=df)) + df = 0 + assert_raises(ValueError, + multivariate_t, + **dict(loc=loc, shape=shape, df=df)) + + def test_reproducibility(self): + rng = np.random.RandomState(4) + loc = rng.uniform(size=3) + shape = np.eye(3) + dist1 = multivariate_t(loc, shape, df=3, seed=2) + dist2 = multivariate_t(loc, shape, df=3, seed=2) + samples1 = dist1.rvs(size=10) + samples2 = dist2.rvs(size=10) + assert_equal(samples1, samples2) + + def test_allow_singular(self): + # Make shape singular and verify error was raised. + args = dict(loc=[0,0], shape=[[0,0],[0,1]], df=1, allow_singular=False) + assert_raises(np.linalg.LinAlgError, multivariate_t, **args) + + @pytest.mark.parametrize("size", [(10, 3), (5, 6, 4, 3)]) + @pytest.mark.parametrize("dim", [2, 3, 4, 5]) + @pytest.mark.parametrize("df", [1., 2., np.inf]) + def test_rvs(self, size, dim, df): + dist = multivariate_t(np.zeros(dim), np.eye(dim), df) + rvs = dist.rvs(size=size) + assert rvs.shape == size + (dim, ) + + def test_cdf_signs(self): + # check that sign of output is correct when np.any(lower > x) + mean = np.zeros(3) + cov = np.eye(3) + df = 10 + b = [[1, 1, 1], [0, 0, 0], [1, 0, 1], [0, 1, 0]] + a = [[0, 0, 0], [1, 1, 1], [0, 1, 0], [1, 0, 1]] + # when odd number of elements of b < a, output is negative + expected_signs = np.array([1, -1, -1, 1]) + cdf = multivariate_normal.cdf(b, mean, cov, df, lower_limit=a) + assert_allclose(cdf, cdf[0]*expected_signs) + + @pytest.mark.parametrize('dim', [1, 2, 5]) + def test_cdf_against_multivariate_normal(self, dim): + # Check accuracy against MVN randomly-generated cases + self.cdf_against_mvn_test(dim) + + @pytest.mark.parametrize('dim', [3, 6, 9]) + def test_cdf_against_multivariate_normal_singular(self, dim): + # Check accuracy against MVN for randomly-generated singular cases + self.cdf_against_mvn_test(3, True) + + def cdf_against_mvn_test(self, dim, singular=False): + # Check for accuracy in the limit that df -> oo and MVT -> MVN + rng = np.random.default_rng(413722918996573) + n = 3 + + w = 10**rng.uniform(-2, 1, size=dim) + cov = _random_covariance(dim, w, rng, singular) + + mean = 10**rng.uniform(-1, 2, size=dim) * np.sign(rng.normal(size=dim)) + a = -10**rng.uniform(-1, 2, size=(n, dim)) + mean + b = 10**rng.uniform(-1, 2, size=(n, dim)) + mean + + res = stats.multivariate_t.cdf(b, mean, cov, df=10000, lower_limit=a, + allow_singular=True, random_state=rng) + ref = stats.multivariate_normal.cdf(b, mean, cov, allow_singular=True, + lower_limit=a) + assert_allclose(res, ref, atol=5e-4) + + def test_cdf_against_univariate_t(self): + rng = np.random.default_rng(413722918996573) + cov = 2 + mean = 0 + x = rng.normal(size=10, scale=np.sqrt(cov)) + df = 3 + + res = stats.multivariate_t.cdf(x, mean, cov, df, lower_limit=-np.inf, + random_state=rng) + ref = stats.t.cdf(x, df, mean, np.sqrt(cov)) + incorrect = stats.norm.cdf(x, mean, np.sqrt(cov)) + + assert_allclose(res, ref, atol=5e-4) # close to t + assert np.all(np.abs(res - incorrect) > 1e-3) # not close to normal + + @pytest.mark.parametrize("dim", [2, 3, 5, 10]) + @pytest.mark.parametrize("seed", [3363958638, 7891119608, 3887698049, + 5013150848, 1495033423, 6170824608]) + @pytest.mark.parametrize("singular", [False, True]) + def test_cdf_against_qsimvtv(self, dim, seed, singular): + if singular and seed != 3363958638: + pytest.skip('Agreement with qsimvtv is not great in singular case') + rng = np.random.default_rng(seed) + w = 10**rng.uniform(-2, 2, size=dim) + cov = _random_covariance(dim, w, rng, singular) + mean = rng.random(dim) + a = -rng.random(dim) + b = rng.random(dim) + df = rng.random() * 5 + + # no lower limit + res = stats.multivariate_t.cdf(b, mean, cov, df, random_state=rng, + allow_singular=True) + with np.errstate(invalid='ignore'): + ref = _qsimvtv(20000, df, cov, np.inf*a, b - mean, rng)[0] + assert_allclose(res, ref, atol=2e-4, rtol=1e-3) + + # with lower limit + res = stats.multivariate_t.cdf(b, mean, cov, df, lower_limit=a, + random_state=rng, allow_singular=True) + with np.errstate(invalid='ignore'): + ref = _qsimvtv(20000, df, cov, a - mean, b - mean, rng)[0] + assert_allclose(res, ref, atol=1e-4, rtol=1e-3) + + @pytest.mark.slow + def test_cdf_against_generic_integrators(self): + # Compare result against generic numerical integrators + dim = 3 + rng = np.random.default_rng(41372291899657) + w = 10 ** rng.uniform(-1, 1, size=dim) + cov = _random_covariance(dim, w, rng, singular=True) + mean = rng.random(dim) + a = -rng.random(dim) + b = rng.random(dim) + df = rng.random() * 5 + + res = stats.multivariate_t.cdf(b, mean, cov, df, random_state=rng, + lower_limit=a) + + def integrand(x): + return stats.multivariate_t.pdf(x.T, mean, cov, df) + + ref = qmc_quad(integrand, a, b, qrng=stats.qmc.Halton(d=dim, seed=rng)) + assert_allclose(res, ref.integral, rtol=1e-3) + + def integrand(*zyx): + return stats.multivariate_t.pdf(zyx[::-1], mean, cov, df) + + ref = tplquad(integrand, a[0], b[0], a[1], b[1], a[2], b[2]) + assert_allclose(res, ref[0], rtol=1e-3) + + def test_against_matlab(self): + # Test against matlab mvtcdf: + # C = [6.21786909 0.2333667 7.95506077; + # 0.2333667 29.67390923 16.53946426; + # 7.95506077 16.53946426 19.17725252] + # df = 1.9559939787727658 + # mvtcdf([0, 0, 0], C, df) % 0.2523 + rng = np.random.default_rng(2967390923) + cov = np.array([[ 6.21786909, 0.2333667 , 7.95506077], + [ 0.2333667 , 29.67390923, 16.53946426], + [ 7.95506077, 16.53946426, 19.17725252]]) + df = 1.9559939787727658 + dist = stats.multivariate_t(shape=cov, df=df) + res = dist.cdf([0, 0, 0], random_state=rng) + ref = 0.2523 + assert_allclose(res, ref, rtol=1e-3) + + def test_frozen(self): + seed = 4137229573 + rng = np.random.default_rng(seed) + loc = rng.uniform(size=3) + x = rng.uniform(size=3) + loc + shape = np.eye(3) + df = rng.random() + args = (loc, shape, df) + + rng_frozen = np.random.default_rng(seed) + rng_unfrozen = np.random.default_rng(seed) + dist = stats.multivariate_t(*args, seed=rng_frozen) + assert_equal(dist.cdf(x), + multivariate_t.cdf(x, *args, random_state=rng_unfrozen)) + + def test_vectorized(self): + dim = 4 + n = (2, 3) + rng = np.random.default_rng(413722918996573) + A = rng.random(size=(dim, dim)) + cov = A @ A.T + mean = rng.random(dim) + x = rng.random(n + (dim,)) + df = rng.random() * 5 + + res = stats.multivariate_t.cdf(x, mean, cov, df, random_state=rng) + + def _cdf_1d(x): + return _qsimvtv(10000, df, cov, -np.inf*x, x-mean, rng)[0] + + ref = np.apply_along_axis(_cdf_1d, -1, x) + assert_allclose(res, ref, atol=1e-4, rtol=1e-3) + + @pytest.mark.parametrize("dim", (3, 7)) + def test_against_analytical(self, dim): + rng = np.random.default_rng(413722918996573) + A = scipy.linalg.toeplitz(c=[1] + [0.5] * (dim - 1)) + res = stats.multivariate_t(shape=A).cdf([0] * dim, random_state=rng) + ref = 1 / (dim + 1) + assert_allclose(res, ref, rtol=5e-5) + + def test_entropy_inf_df(self): + cov = np.eye(3, 3) + df = np.inf + mvt_entropy = stats.multivariate_t.entropy(shape=cov, df=df) + mvn_entropy = stats.multivariate_normal.entropy(None, cov) + assert mvt_entropy == mvn_entropy + + @pytest.mark.parametrize("df", [1, 10, 100]) + def test_entropy_1d(self, df): + mvt_entropy = stats.multivariate_t.entropy(shape=1., df=df) + t_entropy = stats.t.entropy(df=df) + assert_allclose(mvt_entropy, t_entropy, rtol=1e-13) + + # entropy reference values were computed via numerical integration + # + # def integrand(x, y, mvt): + # vec = np.array([x, y]) + # return mvt.logpdf(vec) * mvt.pdf(vec) + + # def multivariate_t_entropy_quad_2d(df, cov): + # dim = cov.shape[0] + # loc = np.zeros((dim, )) + # mvt = stats.multivariate_t(loc, cov, df) + # limit = 100 + # return -integrate.dblquad(integrand, -limit, limit, -limit, limit, + # args=(mvt, ))[0] + + @pytest.mark.parametrize("df, cov, ref, tol", + [(10, np.eye(2, 2), 3.0378770664093313, 1e-14), + (100, np.array([[0.5, 1], [1, 10]]), + 3.55102424550609, 1e-8)]) + def test_entropy_vs_numerical_integration(self, df, cov, ref, tol): + loc = np.zeros((2, )) + mvt = stats.multivariate_t(loc, cov, df) + assert_allclose(mvt.entropy(), ref, rtol=tol) + + @pytest.mark.parametrize( + "df, dim, ref, tol", + [ + (10, 1, 1.5212624929756808, 1e-15), + (100, 1, 1.4289633653182439, 1e-13), + (500, 1, 1.420939531869349, 1e-14), + (1e20, 1, 1.4189385332046727, 1e-15), + (1e100, 1, 1.4189385332046727, 1e-15), + (10, 10, 15.069150450832911, 1e-15), + (1000, 10, 14.19936546446673, 1e-13), + (1e20, 10, 14.189385332046728, 1e-15), + (1e100, 10, 14.189385332046728, 1e-15), + (10, 100, 148.28902883192654, 1e-15), + (1000, 100, 141.99155538003762, 1e-14), + (1e20, 100, 141.8938533204673, 1e-15), + (1e100, 100, 141.8938533204673, 1e-15), + ] + ) + def test_extreme_entropy(self, df, dim, ref, tol): + # Reference values were calculated with mpmath: + # from mpmath import mp + # mp.dps = 500 + # + # def mul_t_mpmath_entropy(dim, df=1): + # dim = mp.mpf(dim) + # df = mp.mpf(df) + # halfsum = (dim + df)/2 + # half_df = df/2 + # + # return float( + # -mp.loggamma(halfsum) + mp.loggamma(half_df) + # + dim / 2 * mp.log(df * mp.pi) + # + halfsum * (mp.digamma(halfsum) - mp.digamma(half_df)) + # + 0.0 + # ) + mvt = stats.multivariate_t(shape=np.eye(dim), df=df) + assert_allclose(mvt.entropy(), ref, rtol=tol) + + def test_entropy_with_covariance(self): + # Generated using np.randn(5, 5) and then rounding + # to two decimal places + _A = np.array([ + [1.42, 0.09, -0.49, 0.17, 0.74], + [-1.13, -0.01, 0.71, 0.4, -0.56], + [1.07, 0.44, -0.28, -0.44, 0.29], + [-1.5, -0.94, -0.67, 0.73, -1.1], + [0.17, -0.08, 1.46, -0.32, 1.36] + ]) + # Set cov to be a symmetric positive semi-definite matrix + cov = _A @ _A.T + + # Test the asymptotic case. For large degrees of freedom + # the entropy approaches the multivariate normal entropy. + df = 1e20 + mul_t_entropy = stats.multivariate_t.entropy(shape=cov, df=df) + mul_norm_entropy = multivariate_normal(None, cov=cov).entropy() + assert_allclose(mul_t_entropy, mul_norm_entropy, rtol=1e-15) + + # Test the regular case. For a dim of 5 the threshold comes out + # to be approximately 766.45. So using slightly + # different dfs on each site of the threshold, the entropies + # are being compared. + df1 = 765 + df2 = 768 + _entropy1 = stats.multivariate_t.entropy(shape=cov, df=df1) + _entropy2 = stats.multivariate_t.entropy(shape=cov, df=df2) + assert_allclose(_entropy1, _entropy2, rtol=1e-5) + + def test_logpdf_df_inf_gh19930(self): + # `multivariate_t._logpdf` (and `logpdf`/`pdf`) was not working with infinite + # `df` after an update to `multivariate_normal._logpdf`. + + # Reproducible example from the issue + res = multivariate_t.logpdf(1, 1, 1, df=np.inf) + ref = multivariate_normal.logpdf(1, 1, 1) + assert_allclose(res, ref) + + # More extensive test + # Generate a valid multivariate normal distribution and corresponding MVT + rng = np.random.default_rng(324893259825) + mean = rng.random(3) + cov = rng.random((3, 3)) + np.eye(3)*3 + cov = cov.T + cov + X = multivariate_normal(mean=mean, cov=cov) + Y = multivariate_t(loc=mean, shape=cov, df=np.inf) + + # compare the pdf and logpdf at 10 random points + x = X.rvs(10) + assert_allclose(Y.logpdf(x), X.logpdf(x)) + assert_allclose(Y.pdf(x), X.pdf(x)) + + +class TestMultivariateHypergeom: + @pytest.mark.parametrize( + "x, m, n, expected", + [ + # Ground truth value from R dmvhyper + ([3, 4], [5, 10], 7, -1.119814), + # test for `n=0` + ([3, 4], [5, 10], 0, -np.inf), + # test for `x < 0` + ([-3, 4], [5, 10], 7, -np.inf), + # test for `m < 0` (RuntimeWarning issue) + ([3, 4], [-5, 10], 7, np.nan), + # test for all `m < 0` and `x.sum() != n` + ([[1, 2], [3, 4]], [[-4, -6], [-5, -10]], + [3, 7], [np.nan, np.nan]), + # test for `x < 0` and `m < 0` (RuntimeWarning issue) + ([-3, 4], [-5, 10], 1, np.nan), + # test for `x > m` + ([1, 11], [10, 1], 12, np.nan), + # test for `m < 0` (RuntimeWarning issue) + ([1, 11], [10, -1], 12, np.nan), + # test for `n < 0` + ([3, 4], [5, 10], -7, np.nan), + # test for `x.sum() != n` + ([3, 3], [5, 10], 7, -np.inf) + ] + ) + def test_logpmf(self, x, m, n, expected): + vals = multivariate_hypergeom.logpmf(x, m, n) + assert_allclose(vals, expected, rtol=1e-6) + + def test_reduces_hypergeom(self): + # test that the multivariate_hypergeom pmf reduces to the + # hypergeom pmf in the 2d case. + val1 = multivariate_hypergeom.pmf(x=[3, 1], m=[10, 5], n=4) + val2 = hypergeom.pmf(k=3, M=15, n=4, N=10) + assert_allclose(val1, val2, rtol=1e-8) + + val1 = multivariate_hypergeom.pmf(x=[7, 3], m=[15, 10], n=10) + val2 = hypergeom.pmf(k=7, M=25, n=10, N=15) + assert_allclose(val1, val2, rtol=1e-8) + + def test_rvs(self): + # test if `rvs` is unbiased and large sample size converges + # to the true mean. + rv = multivariate_hypergeom(m=[3, 5], n=4) + rvs = rv.rvs(size=1000, random_state=123) + assert_allclose(rvs.mean(0), rv.mean(), rtol=1e-2) + + def test_rvs_broadcasting(self): + rv = multivariate_hypergeom(m=[[3, 5], [5, 10]], n=[4, 9]) + rvs = rv.rvs(size=(1000, 2), random_state=123) + assert_allclose(rvs.mean(0), rv.mean(), rtol=1e-2) + + @pytest.mark.parametrize('m, n', ( + ([0, 0, 20, 0, 0], 5), ([0, 0, 0, 0, 0], 0), + ([0, 0], 0), ([0], 0) + )) + def test_rvs_gh16171(self, m, n): + res = multivariate_hypergeom.rvs(m, n) + m = np.asarray(m) + res_ex = m.copy() + res_ex[m != 0] = n + assert_equal(res, res_ex) + + @pytest.mark.parametrize( + "x, m, n, expected", + [ + ([5], [5], 5, 1), + ([3, 4], [5, 10], 7, 0.3263403), + # Ground truth value from R dmvhyper + ([[[3, 5], [0, 8]], [[-1, 9], [1, 1]]], + [5, 10], [[8, 8], [8, 2]], + [[0.3916084, 0.006993007], [0, 0.4761905]]), + # test with empty arrays. + (np.array([], dtype=int), np.array([], dtype=int), 0, []), + ([1, 2], [4, 5], 5, 0), + # Ground truth value from R dmvhyper + ([3, 3, 0], [5, 6, 7], 6, 0.01077354) + ] + ) + def test_pmf(self, x, m, n, expected): + vals = multivariate_hypergeom.pmf(x, m, n) + assert_allclose(vals, expected, rtol=1e-7) + + @pytest.mark.parametrize( + "x, m, n, expected", + [ + ([3, 4], [[5, 10], [10, 15]], 7, [0.3263403, 0.3407531]), + ([[1], [2]], [[3], [4]], [1, 3], [1., 0.]), + ([[[1], [2]]], [[3], [4]], [1, 3], [[1., 0.]]), + ([[1], [2]], [[[[3]]]], [1, 3], [[[1., 0.]]]) + ] + ) + def test_pmf_broadcasting(self, x, m, n, expected): + vals = multivariate_hypergeom.pmf(x, m, n) + assert_allclose(vals, expected, rtol=1e-7) + + def test_cov(self): + cov1 = multivariate_hypergeom.cov(m=[3, 7, 10], n=12) + cov2 = [[0.64421053, -0.26526316, -0.37894737], + [-0.26526316, 1.14947368, -0.88421053], + [-0.37894737, -0.88421053, 1.26315789]] + assert_allclose(cov1, cov2, rtol=1e-8) + + def test_cov_broadcasting(self): + cov1 = multivariate_hypergeom.cov(m=[[7, 9], [10, 15]], n=[8, 12]) + cov2 = [[[1.05, -1.05], [-1.05, 1.05]], + [[1.56, -1.56], [-1.56, 1.56]]] + assert_allclose(cov1, cov2, rtol=1e-8) + + cov3 = multivariate_hypergeom.cov(m=[[4], [5]], n=[4, 5]) + cov4 = [[[0.]], [[0.]]] + assert_allclose(cov3, cov4, rtol=1e-8) + + cov5 = multivariate_hypergeom.cov(m=[7, 9], n=[8, 12]) + cov6 = [[[1.05, -1.05], [-1.05, 1.05]], + [[0.7875, -0.7875], [-0.7875, 0.7875]]] + assert_allclose(cov5, cov6, rtol=1e-8) + + def test_var(self): + # test with hypergeom + var0 = multivariate_hypergeom.var(m=[10, 5], n=4) + var1 = hypergeom.var(M=15, n=4, N=10) + assert_allclose(var0, var1, rtol=1e-8) + + def test_var_broadcasting(self): + var0 = multivariate_hypergeom.var(m=[10, 5], n=[4, 8]) + var1 = multivariate_hypergeom.var(m=[10, 5], n=4) + var2 = multivariate_hypergeom.var(m=[10, 5], n=8) + assert_allclose(var0[0], var1, rtol=1e-8) + assert_allclose(var0[1], var2, rtol=1e-8) + + var3 = multivariate_hypergeom.var(m=[[10, 5], [10, 14]], n=[4, 8]) + var4 = [[0.6984127, 0.6984127], [1.352657, 1.352657]] + assert_allclose(var3, var4, rtol=1e-8) + + var5 = multivariate_hypergeom.var(m=[[5], [10]], n=[5, 10]) + var6 = [[0.], [0.]] + assert_allclose(var5, var6, rtol=1e-8) + + def test_mean(self): + # test with hypergeom + mean0 = multivariate_hypergeom.mean(m=[10, 5], n=4) + mean1 = hypergeom.mean(M=15, n=4, N=10) + assert_allclose(mean0[0], mean1, rtol=1e-8) + + mean2 = multivariate_hypergeom.mean(m=[12, 8], n=10) + mean3 = [12.*10./20., 8.*10./20.] + assert_allclose(mean2, mean3, rtol=1e-8) + + def test_mean_broadcasting(self): + mean0 = multivariate_hypergeom.mean(m=[[3, 5], [10, 5]], n=[4, 8]) + mean1 = [[3.*4./8., 5.*4./8.], [10.*8./15., 5.*8./15.]] + assert_allclose(mean0, mean1, rtol=1e-8) + + def test_mean_edge_cases(self): + mean0 = multivariate_hypergeom.mean(m=[0, 0, 0], n=0) + assert_equal(mean0, [0., 0., 0.]) + + mean1 = multivariate_hypergeom.mean(m=[1, 0, 0], n=2) + assert_equal(mean1, [np.nan, np.nan, np.nan]) + + mean2 = multivariate_hypergeom.mean(m=[[1, 0, 0], [1, 0, 1]], n=2) + assert_allclose(mean2, [[np.nan, np.nan, np.nan], [1., 0., 1.]], + rtol=1e-17) + + mean3 = multivariate_hypergeom.mean(m=np.array([], dtype=int), n=0) + assert_equal(mean3, []) + assert_(mean3.shape == (0, )) + + def test_var_edge_cases(self): + var0 = multivariate_hypergeom.var(m=[0, 0, 0], n=0) + assert_allclose(var0, [0., 0., 0.], rtol=1e-16) + + var1 = multivariate_hypergeom.var(m=[1, 0, 0], n=2) + assert_equal(var1, [np.nan, np.nan, np.nan]) + + var2 = multivariate_hypergeom.var(m=[[1, 0, 0], [1, 0, 1]], n=2) + assert_allclose(var2, [[np.nan, np.nan, np.nan], [0., 0., 0.]], + rtol=1e-17) + + var3 = multivariate_hypergeom.var(m=np.array([], dtype=int), n=0) + assert_equal(var3, []) + assert_(var3.shape == (0, )) + + def test_cov_edge_cases(self): + cov0 = multivariate_hypergeom.cov(m=[1, 0, 0], n=1) + cov1 = [[0., 0., 0.], [0., 0., 0.], [0., 0., 0.]] + assert_allclose(cov0, cov1, rtol=1e-17) + + cov3 = multivariate_hypergeom.cov(m=[0, 0, 0], n=0) + cov4 = [[0., 0., 0.], [0., 0., 0.], [0., 0., 0.]] + assert_equal(cov3, cov4) + + cov5 = multivariate_hypergeom.cov(m=np.array([], dtype=int), n=0) + cov6 = np.array([], dtype=np.float64).reshape(0, 0) + assert_allclose(cov5, cov6, rtol=1e-17) + assert_(cov5.shape == (0, 0)) + + def test_frozen(self): + # The frozen distribution should agree with the regular one + n = 12 + m = [7, 9, 11, 13] + x = [[0, 0, 0, 12], [0, 0, 1, 11], [0, 1, 1, 10], + [1, 1, 1, 9], [1, 1, 2, 8]] + x = np.asarray(x, dtype=int) + mhg_frozen = multivariate_hypergeom(m, n) + assert_allclose(mhg_frozen.pmf(x), + multivariate_hypergeom.pmf(x, m, n)) + assert_allclose(mhg_frozen.logpmf(x), + multivariate_hypergeom.logpmf(x, m, n)) + assert_allclose(mhg_frozen.var(), multivariate_hypergeom.var(m, n)) + assert_allclose(mhg_frozen.cov(), multivariate_hypergeom.cov(m, n)) + + def test_invalid_params(self): + assert_raises(ValueError, multivariate_hypergeom.pmf, 5, 10, 5) + assert_raises(ValueError, multivariate_hypergeom.pmf, 5, [10], 5) + assert_raises(ValueError, multivariate_hypergeom.pmf, [5, 4], [10], 5) + assert_raises(TypeError, multivariate_hypergeom.pmf, [5.5, 4.5], + [10, 15], 5) + assert_raises(TypeError, multivariate_hypergeom.pmf, [5, 4], + [10.5, 15.5], 5) + assert_raises(TypeError, multivariate_hypergeom.pmf, [5, 4], + [10, 15], 5.5) + + +class TestRandomTable: + def get_rng(self): + return np.random.default_rng(628174795866951638) + + def test_process_parameters(self): + message = "`row` must be one-dimensional" + with pytest.raises(ValueError, match=message): + random_table([[1, 2]], [1, 2]) + + message = "`col` must be one-dimensional" + with pytest.raises(ValueError, match=message): + random_table([1, 2], [[1, 2]]) + + message = "each element of `row` must be non-negative" + with pytest.raises(ValueError, match=message): + random_table([1, -1], [1, 2]) + + message = "each element of `col` must be non-negative" + with pytest.raises(ValueError, match=message): + random_table([1, 2], [1, -2]) + + message = "sums over `row` and `col` must be equal" + with pytest.raises(ValueError, match=message): + random_table([1, 2], [1, 0]) + + message = "each element of `row` must be an integer" + with pytest.raises(ValueError, match=message): + random_table([2.1, 2.1], [1, 1, 2]) + + message = "each element of `col` must be an integer" + with pytest.raises(ValueError, match=message): + random_table([1, 2], [1.1, 1.1, 1]) + + row = [1, 3] + col = [2, 1, 1] + r, c, n = random_table._process_parameters([1, 3], [2, 1, 1]) + assert_equal(row, r) + assert_equal(col, c) + assert n == np.sum(row) + + @pytest.mark.parametrize("scale,method", + ((1, "boyett"), (100, "patefield"))) + def test_process_rvs_method_on_None(self, scale, method): + row = np.array([1, 3]) * scale + col = np.array([2, 1, 1]) * scale + + ct = random_table + expected = ct.rvs(row, col, method=method, random_state=1) + got = ct.rvs(row, col, method=None, random_state=1) + + assert_equal(expected, got) + + def test_process_rvs_method_bad_argument(self): + row = [1, 3] + col = [2, 1, 1] + + # order of items in set is random, so cannot check that + message = "'foo' not recognized, must be one of" + with pytest.raises(ValueError, match=message): + random_table.rvs(row, col, method="foo") + + @pytest.mark.parametrize('frozen', (True, False)) + @pytest.mark.parametrize('log', (True, False)) + def test_pmf_logpmf(self, frozen, log): + # The pmf is tested through random sample generation + # with Boyett's algorithm, whose implementation is simple + # enough to verify manually for correctness. + rng = self.get_rng() + row = [2, 6] + col = [1, 3, 4] + rvs = random_table.rvs(row, col, size=1000, + method="boyett", random_state=rng) + + obj = random_table(row, col) if frozen else random_table + method = getattr(obj, "logpmf" if log else "pmf") + if not frozen: + original_method = method + + def method(x): + return original_method(x, row, col) + pmf = (lambda x: np.exp(method(x))) if log else method + + unique_rvs, counts = np.unique(rvs, axis=0, return_counts=True) + + # rough accuracy check + p = pmf(unique_rvs) + assert_allclose(p * len(rvs), counts, rtol=0.1) + + # accept any iterable + p2 = pmf(list(unique_rvs[0])) + assert_equal(p2, p[0]) + + # accept high-dimensional input and 2d input + rvs_nd = rvs.reshape((10, 100) + rvs.shape[1:]) + p = pmf(rvs_nd) + assert p.shape == (10, 100) + for i in range(p.shape[0]): + for j in range(p.shape[1]): + pij = p[i, j] + rvij = rvs_nd[i, j] + qij = pmf(rvij) + assert_equal(pij, qij) + + # probability is zero if column marginal does not match + x = [[0, 1, 1], [2, 1, 3]] + assert_equal(np.sum(x, axis=-1), row) + p = pmf(x) + assert p == 0 + + # probability is zero if row marginal does not match + x = [[0, 1, 2], [1, 2, 2]] + assert_equal(np.sum(x, axis=-2), col) + p = pmf(x) + assert p == 0 + + # response to invalid inputs + message = "`x` must be at least two-dimensional" + with pytest.raises(ValueError, match=message): + pmf([1]) + + message = "`x` must contain only integral values" + with pytest.raises(ValueError, match=message): + pmf([[1.1]]) + + message = "`x` must contain only integral values" + with pytest.raises(ValueError, match=message): + pmf([[np.nan]]) + + message = "`x` must contain only non-negative values" + with pytest.raises(ValueError, match=message): + pmf([[-1]]) + + message = "shape of `x` must agree with `row`" + with pytest.raises(ValueError, match=message): + pmf([[1, 2, 3]]) + + message = "shape of `x` must agree with `col`" + with pytest.raises(ValueError, match=message): + pmf([[1, 2], + [3, 4]]) + + @pytest.mark.parametrize("method", ("boyett", "patefield")) + def test_rvs_mean(self, method): + # test if `rvs` is unbiased and large sample size converges + # to the true mean. + rng = self.get_rng() + row = [2, 6] + col = [1, 3, 4] + rvs = random_table.rvs(row, col, size=1000, method=method, + random_state=rng) + mean = random_table.mean(row, col) + assert_equal(np.sum(mean), np.sum(row)) + assert_allclose(rvs.mean(0), mean, atol=0.05) + assert_equal(rvs.sum(axis=-1), np.broadcast_to(row, (1000, 2))) + assert_equal(rvs.sum(axis=-2), np.broadcast_to(col, (1000, 3))) + + def test_rvs_cov(self): + # test if `rvs` generated with patefield and boyett algorithms + # produce approximately the same covariance matrix + rng = self.get_rng() + row = [2, 6] + col = [1, 3, 4] + rvs1 = random_table.rvs(row, col, size=10000, method="boyett", + random_state=rng) + rvs2 = random_table.rvs(row, col, size=10000, method="patefield", + random_state=rng) + cov1 = np.var(rvs1, axis=0) + cov2 = np.var(rvs2, axis=0) + assert_allclose(cov1, cov2, atol=0.02) + + @pytest.mark.parametrize("method", ("boyett", "patefield")) + def test_rvs_size(self, method): + row = [2, 6] + col = [1, 3, 4] + + # test size `None` + rv = random_table.rvs(row, col, method=method, + random_state=self.get_rng()) + assert rv.shape == (2, 3) + + # test size 1 + rv2 = random_table.rvs(row, col, size=1, method=method, + random_state=self.get_rng()) + assert rv2.shape == (1, 2, 3) + assert_equal(rv, rv2[0]) + + # test size 0 + rv3 = random_table.rvs(row, col, size=0, method=method, + random_state=self.get_rng()) + assert rv3.shape == (0, 2, 3) + + # test other valid size + rv4 = random_table.rvs(row, col, size=20, method=method, + random_state=self.get_rng()) + assert rv4.shape == (20, 2, 3) + + rv5 = random_table.rvs(row, col, size=(4, 5), method=method, + random_state=self.get_rng()) + assert rv5.shape == (4, 5, 2, 3) + + assert_allclose(rv5.reshape(20, 2, 3), rv4, rtol=1e-15) + + # test invalid size + message = "`size` must be a non-negative integer or `None`" + with pytest.raises(ValueError, match=message): + random_table.rvs(row, col, size=-1, method=method, + random_state=self.get_rng()) + + with pytest.raises(ValueError, match=message): + random_table.rvs(row, col, size=np.nan, method=method, + random_state=self.get_rng()) + + @pytest.mark.parametrize("method", ("boyett", "patefield")) + def test_rvs_method(self, method): + # This test assumes that pmf is correct and checks that random samples + # follow this probability distribution. This seems like a circular + # argument, since pmf is checked in test_pmf_logpmf with random samples + # generated with the rvs method. This test is not redundant, because + # test_pmf_logpmf intentionally uses rvs generation with Boyett only, + # but here we test both Boyett and Patefield. + row = [2, 6] + col = [1, 3, 4] + + ct = random_table + rvs = ct.rvs(row, col, size=100000, method=method, + random_state=self.get_rng()) + + unique_rvs, counts = np.unique(rvs, axis=0, return_counts=True) + + # generated frequencies should match expected frequencies + p = ct.pmf(unique_rvs, row, col) + assert_allclose(p * len(rvs), counts, rtol=0.02) + + @pytest.mark.parametrize("method", ("boyett", "patefield")) + def test_rvs_with_zeros_in_col_row(self, method): + row = [0, 1, 0] + col = [1, 0, 0, 0] + d = random_table(row, col) + rv = d.rvs(1000, method=method, random_state=self.get_rng()) + expected = np.zeros((1000, len(row), len(col))) + expected[...] = [[0, 0, 0, 0], + [1, 0, 0, 0], + [0, 0, 0, 0]] + assert_equal(rv, expected) + + @pytest.mark.parametrize("method", (None, "boyett", "patefield")) + @pytest.mark.parametrize("col", ([], [0])) + @pytest.mark.parametrize("row", ([], [0])) + def test_rvs_with_edge_cases(self, method, row, col): + d = random_table(row, col) + rv = d.rvs(10, method=method, random_state=self.get_rng()) + expected = np.zeros((10, len(row), len(col))) + assert_equal(rv, expected) + + @pytest.mark.parametrize('v', (1, 2)) + def test_rvs_rcont(self, v): + # This test checks the internal low-level interface. + # It is implicitly also checked by the other test_rvs* calls. + import scipy.stats._rcont as _rcont + + row = np.array([1, 3], dtype=np.int64) + col = np.array([2, 1, 1], dtype=np.int64) + + rvs = getattr(_rcont, f"rvs_rcont{v}") + + ntot = np.sum(row) + result = rvs(row, col, ntot, 1, self.get_rng()) + + assert result.shape == (1, len(row), len(col)) + assert np.sum(result) == ntot + + def test_frozen(self): + row = [2, 6] + col = [1, 3, 4] + d = random_table(row, col, seed=self.get_rng()) + + sample = d.rvs() + + expected = random_table.mean(row, col) + assert_equal(expected, d.mean()) + + expected = random_table.pmf(sample, row, col) + assert_equal(expected, d.pmf(sample)) + + expected = random_table.logpmf(sample, row, col) + assert_equal(expected, d.logpmf(sample)) + + @pytest.mark.parametrize("method", ("boyett", "patefield")) + def test_rvs_frozen(self, method): + row = [2, 6] + col = [1, 3, 4] + d = random_table(row, col, seed=self.get_rng()) + + expected = random_table.rvs(row, col, size=10, method=method, + random_state=self.get_rng()) + got = d.rvs(size=10, method=method) + assert_equal(expected, got) + + +def check_pickling(distfn, args): + # check that a distribution instance pickles and unpickles + # pay special attention to the random_state property + + # save the random_state (restore later) + rndm = distfn.random_state + + distfn.random_state = 1234 + distfn.rvs(*args, size=8) + s = pickle.dumps(distfn) + r0 = distfn.rvs(*args, size=8) + + unpickled = pickle.loads(s) + r1 = unpickled.rvs(*args, size=8) + assert_equal(r0, r1) + + # restore the random_state + distfn.random_state = rndm + + +@pytest.mark.thread_unsafe(reason="uses numpy global random state and monkey-patching") +def test_random_state_property(): + scale = np.eye(3) + scale[0, 1] = 0.5 + scale[1, 0] = 0.5 + dists = [ + [multivariate_normal, ()], + [dirichlet, (np.array([1.]), )], + [wishart, (10, scale)], + [invwishart, (10, scale)], + [multinomial, (5, [0.5, 0.4, 0.1])], + [ortho_group, (2,)], + [special_ortho_group, (2,)] + ] + for distfn, args in dists: + check_random_state_property(distfn, args) + check_pickling(distfn, args) + + +class TestVonMises_Fisher: + @pytest.mark.parametrize("dim", [2, 3, 4, 6]) + @pytest.mark.parametrize("size", [None, 1, 5, (5, 4)]) + def test_samples(self, dim, size): + # test that samples have correct shape and norm 1 + rng = np.random.default_rng(2777937887058094419) + mu = np.full((dim, ), 1/np.sqrt(dim)) + vmf_dist = vonmises_fisher(mu, 1, seed=rng) + samples = vmf_dist.rvs(size) + mean, cov = np.zeros(dim), np.eye(dim) + expected_shape = rng.multivariate_normal(mean, cov, size=size).shape + assert samples.shape == expected_shape + norms = np.linalg.norm(samples, axis=-1) + assert_allclose(norms, 1.) + + @pytest.mark.parametrize("dim", [5, 8]) + @pytest.mark.parametrize("kappa", [1e15, 1e20, 1e30]) + def test_sampling_high_concentration(self, dim, kappa): + # test that no warnings are encountered for high values + rng = np.random.default_rng(2777937887058094419) + mu = np.full((dim, ), 1/np.sqrt(dim)) + vmf_dist = vonmises_fisher(mu, kappa, seed=rng) + vmf_dist.rvs(10) + + def test_two_dimensional_mu(self): + mu = np.ones((2, 2)) + msg = "'mu' must have one-dimensional shape." + with pytest.raises(ValueError, match=msg): + vonmises_fisher(mu, 1) + + def test_wrong_norm_mu(self): + mu = np.ones((2, )) + msg = "'mu' must be a unit vector of norm 1." + with pytest.raises(ValueError, match=msg): + vonmises_fisher(mu, 1) + + def test_one_entry_mu(self): + mu = np.ones((1, )) + msg = "'mu' must have at least two entries." + with pytest.raises(ValueError, match=msg): + vonmises_fisher(mu, 1) + + @pytest.mark.parametrize("kappa", [-1, (5, 3)]) + def test_kappa_validation(self, kappa): + msg = "'kappa' must be a positive scalar." + with pytest.raises(ValueError, match=msg): + vonmises_fisher([1, 0], kappa) + + @pytest.mark.parametrize("kappa", [0, 0.]) + def test_kappa_zero(self, kappa): + msg = ("For 'kappa=0' the von Mises-Fisher distribution " + "becomes the uniform distribution on the sphere " + "surface. Consider using 'scipy.stats.uniform_direction' " + "instead.") + with pytest.raises(ValueError, match=msg): + vonmises_fisher([1, 0], kappa) + + + @pytest.mark.parametrize("method", [vonmises_fisher.pdf, + vonmises_fisher.logpdf]) + def test_invalid_shapes_pdf_logpdf(self, method): + x = np.array([1., 0., 0]) + msg = ("The dimensionality of the last axis of 'x' must " + "match the dimensionality of the von Mises Fisher " + "distribution.") + with pytest.raises(ValueError, match=msg): + method(x, [1, 0], 1) + + @pytest.mark.parametrize("method", [vonmises_fisher.pdf, + vonmises_fisher.logpdf]) + def test_unnormalized_input(self, method): + x = np.array([0.5, 0.]) + msg = "'x' must be unit vectors of norm 1 along last dimension." + with pytest.raises(ValueError, match=msg): + method(x, [1, 0], 1) + + # Expected values of the vonmises-fisher logPDF were computed via mpmath + # from mpmath import mp + # import numpy as np + # mp.dps = 50 + # def logpdf_mpmath(x, mu, kappa): + # dim = mu.size + # halfdim = mp.mpf(0.5 * dim) + # kappa = mp.mpf(kappa) + # const = (kappa**(halfdim - mp.one)/((2*mp.pi)**halfdim * \ + # mp.besseli(halfdim -mp.one, kappa))) + # return float(const * mp.exp(kappa*mp.fdot(x, mu))) + + @pytest.mark.parametrize('x, mu, kappa, reference', + [(np.array([1., 0., 0.]), np.array([1., 0., 0.]), + 1e-4, 0.0795854295583605), + (np.array([1., 0., 0]), np.array([0., 0., 1.]), + 1e-4, 0.07957747141331854), + (np.array([1., 0., 0.]), np.array([1., 0., 0.]), + 100, 15.915494309189533), + (np.array([1., 0., 0]), np.array([0., 0., 1.]), + 100, 5.920684802611232e-43), + (np.array([1., 0., 0.]), + np.array([np.sqrt(0.98), np.sqrt(0.02), 0.]), + 2000, 5.930499050746588e-07), + (np.array([1., 0., 0]), np.array([1., 0., 0.]), + 2000, 318.3098861837907), + (np.array([1., 0., 0., 0., 0.]), + np.array([1., 0., 0., 0., 0.]), + 2000, 101371.86957712633), + (np.array([1., 0., 0., 0., 0.]), + np.array([np.sqrt(0.98), np.sqrt(0.02), 0., + 0, 0.]), + 2000, 0.00018886808182653578), + (np.array([1., 0., 0., 0., 0.]), + np.array([np.sqrt(0.8), np.sqrt(0.2), 0., + 0, 0.]), + 2000, 2.0255393314603194e-87)]) + def test_pdf_accuracy(self, x, mu, kappa, reference): + pdf = vonmises_fisher(mu, kappa).pdf(x) + assert_allclose(pdf, reference, rtol=1e-13) + + # Expected values of the vonmises-fisher logPDF were computed via mpmath + # from mpmath import mp + # import numpy as np + # mp.dps = 50 + # def logpdf_mpmath(x, mu, kappa): + # dim = mu.size + # halfdim = mp.mpf(0.5 * dim) + # kappa = mp.mpf(kappa) + # two = mp.mpf(2.) + # const = (kappa**(halfdim - mp.one)/((two*mp.pi)**halfdim * \ + # mp.besseli(halfdim - mp.one, kappa))) + # return float(mp.log(const * mp.exp(kappa*mp.fdot(x, mu)))) + + @pytest.mark.parametrize('x, mu, kappa, reference', + [(np.array([1., 0., 0.]), np.array([1., 0., 0.]), + 1e-4, -2.5309242486359573), + (np.array([1., 0., 0]), np.array([0., 0., 1.]), + 1e-4, -2.5310242486359575), + (np.array([1., 0., 0.]), np.array([1., 0., 0.]), + 100, 2.767293119578746), + (np.array([1., 0., 0]), np.array([0., 0., 1.]), + 100, -97.23270688042125), + (np.array([1., 0., 0.]), + np.array([np.sqrt(0.98), np.sqrt(0.02), 0.]), + 2000, -14.337987284534103), + (np.array([1., 0., 0]), np.array([1., 0., 0.]), + 2000, 5.763025393132737), + (np.array([1., 0., 0., 0., 0.]), + np.array([1., 0., 0., 0., 0.]), + 2000, 11.526550911307156), + (np.array([1., 0., 0., 0., 0.]), + np.array([np.sqrt(0.98), np.sqrt(0.02), 0., + 0, 0.]), + 2000, -8.574461766359684), + (np.array([1., 0., 0., 0., 0.]), + np.array([np.sqrt(0.8), np.sqrt(0.2), 0., + 0, 0.]), + 2000, -199.61906708886113)]) + def test_logpdf_accuracy(self, x, mu, kappa, reference): + logpdf = vonmises_fisher(mu, kappa).logpdf(x) + assert_allclose(logpdf, reference, rtol=1e-14) + + # Expected values of the vonmises-fisher entropy were computed via mpmath + # from mpmath import mp + # import numpy as np + # mp.dps = 50 + # def entropy_mpmath(dim, kappa): + # mu = np.full((dim, ), 1/np.sqrt(dim)) + # kappa = mp.mpf(kappa) + # halfdim = mp.mpf(0.5 * dim) + # logconstant = (mp.log(kappa**(halfdim - mp.one) + # /((2*mp.pi)**halfdim + # * mp.besseli(halfdim -mp.one, kappa))) + # return float(-logconstant - kappa * mp.besseli(halfdim, kappa)/ + # mp.besseli(halfdim -1, kappa)) + + @pytest.mark.parametrize('dim, kappa, reference', + [(3, 1e-4, 2.531024245302624), + (3, 100, -1.7672931195787458), + (5, 5000, -11.359032310024453), + (8, 1, 3.4189526482545527)]) + def test_entropy_accuracy(self, dim, kappa, reference): + mu = np.full((dim, ), 1/np.sqrt(dim)) + entropy = vonmises_fisher(mu, kappa).entropy() + assert_allclose(entropy, reference, rtol=2e-14) + + @pytest.mark.parametrize("method", [vonmises_fisher.pdf, + vonmises_fisher.logpdf]) + def test_broadcasting(self, method): + # test that pdf and logpdf values are correctly broadcasted + testshape = (2, 2) + rng = np.random.default_rng(2777937887058094419) + x = uniform_direction(3).rvs(testshape, random_state=rng) + mu = np.full((3, ), 1/np.sqrt(3)) + kappa = 5 + result_all = method(x, mu, kappa) + assert result_all.shape == testshape + for i in range(testshape[0]): + for j in range(testshape[1]): + current_val = method(x[i, j, :], mu, kappa) + assert_allclose(current_val, result_all[i, j], rtol=1e-15) + + def test_vs_vonmises_2d(self): + # test that in 2D, von Mises-Fisher yields the same results + # as the von Mises distribution + rng = np.random.default_rng(2777937887058094419) + mu = np.array([0, 1]) + mu_angle = np.arctan2(mu[1], mu[0]) + kappa = 20 + vmf = vonmises_fisher(mu, kappa) + vonmises_dist = vonmises(loc=mu_angle, kappa=kappa) + vectors = uniform_direction(2).rvs(10, random_state=rng) + angles = np.arctan2(vectors[:, 1], vectors[:, 0]) + assert_allclose(vonmises_dist.entropy(), vmf.entropy()) + assert_allclose(vonmises_dist.pdf(angles), vmf.pdf(vectors)) + assert_allclose(vonmises_dist.logpdf(angles), vmf.logpdf(vectors)) + + @pytest.mark.parametrize("dim", [2, 3, 6]) + @pytest.mark.parametrize("kappa, mu_tol, kappa_tol", + [(1, 5e-2, 5e-2), + (10, 1e-2, 1e-2), + (100, 5e-3, 2e-2), + (1000, 1e-3, 2e-2)]) + def test_fit_accuracy(self, dim, kappa, mu_tol, kappa_tol): + mu = np.full((dim, ), 1/np.sqrt(dim)) + vmf_dist = vonmises_fisher(mu, kappa) + rng = np.random.default_rng(2777937887058094419) + n_samples = 10000 + samples = vmf_dist.rvs(n_samples, random_state=rng) + mu_fit, kappa_fit = vonmises_fisher.fit(samples) + angular_error = np.arccos(mu.dot(mu_fit)) + assert_allclose(angular_error, 0., atol=mu_tol, rtol=0) + assert_allclose(kappa, kappa_fit, rtol=kappa_tol) + + def test_fit_error_one_dimensional_data(self): + x = np.zeros((3, )) + msg = "'x' must be two dimensional." + with pytest.raises(ValueError, match=msg): + vonmises_fisher.fit(x) + + def test_fit_error_unnormalized_data(self): + x = np.ones((3, 3)) + msg = "'x' must be unit vectors of norm 1 along last dimension." + with pytest.raises(ValueError, match=msg): + vonmises_fisher.fit(x) + + def test_frozen_distribution(self): + mu = np.array([0, 0, 1]) + kappa = 5 + frozen = vonmises_fisher(mu, kappa) + frozen_seed = vonmises_fisher(mu, kappa, seed=514) + + rvs1 = frozen.rvs(random_state=514) + rvs2 = vonmises_fisher.rvs(mu, kappa, random_state=514) + rvs3 = frozen_seed.rvs() + + assert_equal(rvs1, rvs2) + assert_equal(rvs1, rvs3) + + +class TestDirichletMultinomial: + @classmethod + def get_params(self, m): + rng = np.random.default_rng(28469824356873456) + alpha = rng.uniform(0, 100, size=2) + x = rng.integers(1, 20, size=(m, 2)) + n = x.sum(axis=-1) + return rng, m, alpha, n, x + + def test_frozen(self): + rng = np.random.default_rng(28469824356873456) + + alpha = rng.uniform(0, 100, 10) + x = rng.integers(0, 10, 10) + n = np.sum(x, axis=-1) + + d = dirichlet_multinomial(alpha, n) + assert_equal(d.logpmf(x), dirichlet_multinomial.logpmf(x, alpha, n)) + assert_equal(d.pmf(x), dirichlet_multinomial.pmf(x, alpha, n)) + assert_equal(d.mean(), dirichlet_multinomial.mean(alpha, n)) + assert_equal(d.var(), dirichlet_multinomial.var(alpha, n)) + assert_equal(d.cov(), dirichlet_multinomial.cov(alpha, n)) + + def test_pmf_logpmf_against_R(self): + # # Compare PMF against R's extraDistr ddirmnon + # # library(extraDistr) + # # options(digits=16) + # ddirmnom(c(1, 2, 3), 6, c(3, 4, 5)) + x = np.array([1, 2, 3]) + n = np.sum(x) + alpha = np.array([3, 4, 5]) + res = dirichlet_multinomial.pmf(x, alpha, n) + logres = dirichlet_multinomial.logpmf(x, alpha, n) + ref = 0.08484162895927638 + assert_allclose(res, ref) + assert_allclose(logres, np.log(ref)) + assert res.shape == logres.shape == () + + # library(extraDistr) + # options(digits=16) + # ddirmnom(c(4, 3, 2, 0, 2, 3, 5, 7, 4, 7), 37, + # c(45.01025314, 21.98739582, 15.14851365, 80.21588671, + # 52.84935481, 25.20905262, 53.85373737, 4.88568118, + # 89.06440654, 20.11359466)) + rng = np.random.default_rng(28469824356873456) + alpha = rng.uniform(0, 100, 10) + x = rng.integers(0, 10, 10) + n = np.sum(x, axis=-1) + res = dirichlet_multinomial(alpha, n).pmf(x) + logres = dirichlet_multinomial.logpmf(x, alpha, n) + ref = 3.65409306285992e-16 + assert_allclose(res, ref) + assert_allclose(logres, np.log(ref)) + + def test_pmf_logpmf_support(self): + # when the sum of the category counts does not equal the number of + # trials, the PMF is zero + rng, m, alpha, n, x = self.get_params(1) + n += 1 + assert_equal(dirichlet_multinomial(alpha, n).pmf(x), 0) + assert_equal(dirichlet_multinomial(alpha, n).logpmf(x), -np.inf) + + rng, m, alpha, n, x = self.get_params(10) + i = rng.random(size=10) > 0.5 + x[i] = np.round(x[i] * 2) # sum of these x does not equal n + assert_equal(dirichlet_multinomial(alpha, n).pmf(x)[i], 0) + assert_equal(dirichlet_multinomial(alpha, n).logpmf(x)[i], -np.inf) + assert np.all(dirichlet_multinomial(alpha, n).pmf(x)[~i] > 0) + assert np.all(dirichlet_multinomial(alpha, n).logpmf(x)[~i] > -np.inf) + + def test_dimensionality_one(self): + # if the dimensionality is one, there is only one possible outcome + n = 6 # number of trials + alpha = [10] # concentration parameters + x = np.asarray([n]) # counts + dist = dirichlet_multinomial(alpha, n) + + assert_equal(dist.pmf(x), 1) + assert_equal(dist.pmf(x+1), 0) + assert_equal(dist.logpmf(x), 0) + assert_equal(dist.logpmf(x+1), -np.inf) + assert_equal(dist.mean(), n) + assert_equal(dist.var(), 0) + assert_equal(dist.cov(), 0) + + def test_n_is_zero(self): + # similarly, only one possible outcome if n is zero + n = 0 + alpha = np.asarray([1., 1.]) + x = np.asarray([0, 0]) + dist = dirichlet_multinomial(alpha, n) + + assert_equal(dist.pmf(x), 1) + assert_equal(dist.pmf(x+1), 0) + assert_equal(dist.logpmf(x), 0) + assert_equal(dist.logpmf(x+1), -np.inf) + assert_equal(dist.mean(), [0, 0]) + assert_equal(dist.var(), [0, 0]) + assert_equal(dist.cov(), [[0, 0], [0, 0]]) + + @pytest.mark.parametrize('method_name', ['pmf', 'logpmf']) + def test_against_betabinom_pmf(self, method_name): + rng, m, alpha, n, x = self.get_params(100) + + method = getattr(dirichlet_multinomial(alpha, n), method_name) + ref_method = getattr(stats.betabinom(n, *alpha.T), method_name) + + res = method(x) + ref = ref_method(x.T[0]) + assert_allclose(res, ref) + + @pytest.mark.parametrize('method_name', ['mean', 'var']) + def test_against_betabinom_moments(self, method_name): + rng, m, alpha, n, x = self.get_params(100) + + method = getattr(dirichlet_multinomial(alpha, n), method_name) + ref_method = getattr(stats.betabinom(n, *alpha.T), method_name) + + res = method()[:, 0] + ref = ref_method() + assert_allclose(res, ref) + + def test_moments(self): + rng = np.random.default_rng(28469824356873456) + dim = 5 + n = rng.integers(1, 100) + alpha = rng.random(size=dim) * 10 + dist = dirichlet_multinomial(alpha, n) + + # Generate a random sample from the distribution using NumPy + m = 100000 + p = rng.dirichlet(alpha, size=m) + x = rng.multinomial(n, p, size=m) + + assert_allclose(dist.mean(), np.mean(x, axis=0), rtol=5e-3) + assert_allclose(dist.var(), np.var(x, axis=0), rtol=1e-2) + assert dist.mean().shape == dist.var().shape == (dim,) + + cov = dist.cov() + assert cov.shape == (dim, dim) + assert_allclose(cov, np.cov(x.T), rtol=2e-2) + assert_equal(np.diag(cov), dist.var()) + assert np.all(scipy.linalg.eigh(cov)[0] > 0) # positive definite + + def test_input_validation(self): + # valid inputs + x0 = np.array([1, 2, 3]) + n0 = np.sum(x0) + alpha0 = np.array([3, 4, 5]) + + text = "`x` must contain only non-negative integers." + with assert_raises(ValueError, match=text): + dirichlet_multinomial.logpmf([1, -1, 3], alpha0, n0) + with assert_raises(ValueError, match=text): + dirichlet_multinomial.logpmf([1, 2.1, 3], alpha0, n0) + + text = "`alpha` must contain only positive values." + with assert_raises(ValueError, match=text): + dirichlet_multinomial.logpmf(x0, [3, 0, 4], n0) + with assert_raises(ValueError, match=text): + dirichlet_multinomial.logpmf(x0, [3, -1, 4], n0) + + text = "`n` must be a non-negative integer." + with assert_raises(ValueError, match=text): + dirichlet_multinomial.logpmf(x0, alpha0, 49.1) + with assert_raises(ValueError, match=text): + dirichlet_multinomial.logpmf(x0, alpha0, -1) + + x = np.array([1, 2, 3, 4]) + alpha = np.array([3, 4, 5]) + text = "`x` and `alpha` must be broadcastable." + with assert_raises(ValueError, match=text): + dirichlet_multinomial.logpmf(x, alpha, x.sum()) + + @pytest.mark.parametrize('method', ['pmf', 'logpmf']) + def test_broadcasting_pmf(self, method): + alpha = np.array([[3, 4, 5], [4, 5, 6], [5, 5, 7], [8, 9, 10]]) + n = np.array([[6], [7], [8]]) + x = np.array([[1, 2, 3], [2, 2, 3]]).reshape((2, 1, 1, 3)) + method = getattr(dirichlet_multinomial, method) + res = method(x, alpha, n) + assert res.shape == (2, 3, 4) + for i in range(len(x)): + for j in range(len(n)): + for k in range(len(alpha)): + res_ijk = res[i, j, k] + ref = method(x[i].squeeze(), alpha[k].squeeze(), n[j].squeeze()) + assert_allclose(res_ijk, ref) + + @pytest.mark.parametrize('method_name', ['mean', 'var', 'cov']) + def test_broadcasting_moments(self, method_name): + alpha = np.array([[3, 4, 5], [4, 5, 6], [5, 5, 7], [8, 9, 10]]) + n = np.array([[6], [7], [8]]) + method = getattr(dirichlet_multinomial, method_name) + res = method(alpha, n) + assert res.shape == (3, 4, 3) if method_name != 'cov' else (3, 4, 3, 3) + for j in range(len(n)): + for k in range(len(alpha)): + res_ijk = res[j, k] + ref = method(alpha[k].squeeze(), n[j].squeeze()) + assert_allclose(res_ijk, ref) + + +class TestNormalInverseGamma: + + def test_marginal_x(self): + # According to [1], sqrt(a * lmbda / b) * (x - u) should follow a t-distribution + # with 2*a degrees of freedom. Test that this is true of the PDF and random + # variates. + rng = np.random.default_rng(8925849245) + mu, lmbda, a, b = rng.random(4) + + norm_inv_gamma = stats.normal_inverse_gamma(mu, lmbda, a, b) + t = stats.t(2*a, loc=mu, scale=1/np.sqrt(a * lmbda / b)) + + # Test PDF + x = np.linspace(-5, 5, 11) + res = tanhsinh(lambda s2, x: norm_inv_gamma.pdf(x, s2), 0, np.inf, args=(x,)) + ref = t.pdf(x) + assert_allclose(res.integral, ref) + + # Test RVS + res = norm_inv_gamma.rvs(size=10000, random_state=rng) + _, pvalue = stats.ks_1samp(res[0], t.cdf) + assert pvalue > 0.1 + + def test_marginal_s2(self): + # According to [1], s2 should follow an inverse gamma distribution with + # shapes a, b (where b is the scale in our parameterization). Test that + # this is true of the PDF and random variates. + rng = np.random.default_rng(8925849245) + mu, lmbda, a, b = rng.random(4) + + norm_inv_gamma = stats.normal_inverse_gamma(mu, lmbda, a, b) + inv_gamma = stats.invgamma(a, scale=b) + + # Test PDF + s2 = np.linspace(0.1, 10, 10) + res = tanhsinh(lambda x, s2: norm_inv_gamma.pdf(x, s2), + -np.inf, np.inf, args=(s2,)) + ref = inv_gamma.pdf(s2) + assert_allclose(res.integral, ref) + + # Test RVS + res = norm_inv_gamma.rvs(size=10000, random_state=rng) + _, pvalue = stats.ks_1samp(res[1], inv_gamma.cdf) + assert pvalue > 0.1 + + def test_pdf_logpdf(self): + # Check that PDF and log-PDF are consistent + rng = np.random.default_rng(8925849245) + mu, lmbda, a, b = rng.random((4, 20)) - 0.25 # make some invalid + x, s2 = rng.random(size=(2, 20)) - 0.25 + res = stats.normal_inverse_gamma(mu, lmbda, a, b).pdf(x, s2) + ref = stats.normal_inverse_gamma.logpdf(x, s2, mu, lmbda, a, b) + assert_allclose(res, np.exp(ref)) + + def test_invalid_and_special_cases(self): + # Test cases that are handled by input validation rather than the formulas + rng = np.random.default_rng(8925849245) + mu, lmbda, a, b = rng.random(4) + x, s2 = rng.random(2) + + res = stats.normal_inverse_gamma(np.nan, lmbda, a, b).pdf(x, s2) + assert_equal(res, np.nan) + + res = stats.normal_inverse_gamma(mu, -1, a, b).pdf(x, s2) + assert_equal(res, np.nan) + + res = stats.normal_inverse_gamma(mu, lmbda, 0, b).pdf(x, s2) + assert_equal(res, np.nan) + + res = stats.normal_inverse_gamma(mu, lmbda, a, -1).pdf(x, s2) + assert_equal(res, np.nan) + + res = stats.normal_inverse_gamma(mu, lmbda, a, b).pdf(x, -1) + assert_equal(res, 0) + + # PDF with out-of-support s2 is not zero if shape parameter is invalid + res = stats.normal_inverse_gamma(mu, [-1, np.nan], a, b).pdf(x, -1) + assert_equal(res, np.nan) + + res = stats.normal_inverse_gamma(mu, -1, a, b).mean() + assert_equal(res, (np.nan, np.nan)) + + res = stats.normal_inverse_gamma(mu, lmbda, -1, b).var() + assert_equal(res, (np.nan, np.nan)) + + with pytest.raises(ValueError, match="Domain error in arguments..."): + stats.normal_inverse_gamma(mu, lmbda, a, -1).rvs() + + def test_broadcasting(self): + # Test methods with broadcastable array parameters. Roughly speaking, the + # shapes should be the broadcasted shapes of all arguments, and the raveled + # outputs should be the same as the outputs with raveled inputs. + rng = np.random.default_rng(8925849245) + b = rng.random(2) + a = rng.random((3, 1)) + 2 # for defined moments + lmbda = rng.random((4, 1, 1)) + mu = rng.random((5, 1, 1, 1)) + s2 = rng.random((6, 1, 1, 1, 1)) + x = rng.random((7, 1, 1, 1, 1, 1)) + dist = stats.normal_inverse_gamma(mu, lmbda, a, b) + + # Test PDF and log-PDF + broadcasted = np.broadcast_arrays(x, s2, mu, lmbda, a, b) + broadcasted_raveled = [np.ravel(arr) for arr in broadcasted] + + res = dist.pdf(x, s2) + assert res.shape == broadcasted[0].shape + assert_allclose(res.ravel(), + stats.normal_inverse_gamma.pdf(*broadcasted_raveled)) + + res = dist.logpdf(x, s2) + assert res.shape == broadcasted[0].shape + assert_allclose(res.ravel(), + stats.normal_inverse_gamma.logpdf(*broadcasted_raveled)) + + # Test moments + broadcasted = np.broadcast_arrays(mu, lmbda, a, b) + broadcasted_raveled = [np.ravel(arr) for arr in broadcasted] + + res = dist.mean() + assert res[0].shape == broadcasted[0].shape + assert_allclose((res[0].ravel(), res[1].ravel()), + stats.normal_inverse_gamma.mean(*broadcasted_raveled)) + + res = dist.var() + assert res[0].shape == broadcasted[0].shape + assert_allclose((res[0].ravel(), res[1].ravel()), + stats.normal_inverse_gamma.var(*broadcasted_raveled)) + + # Test RVS + size = (6, 5, 4, 3, 2) + rng = np.random.default_rng(2348923985324) + res = dist.rvs(size=size, random_state=rng) + rng = np.random.default_rng(2348923985324) + shape = 6, 5*4*3*2 + ref = stats.normal_inverse_gamma.rvs(*broadcasted_raveled, size=shape, + random_state=rng) + assert_allclose((res[0].reshape(shape), res[1].reshape(shape)), ref) + + @pytest.mark.slow + @pytest.mark.fail_slow(10) + def test_moments(self): + # Test moments against quadrature + rng = np.random.default_rng(8925849245) + mu, lmbda, a, b = rng.random(4) + a += 2 # ensure defined + + dist = stats.normal_inverse_gamma(mu, lmbda, a, b) + res = dist.mean() + + ref = dblquad(lambda s2, x: dist.pdf(x, s2) * x, -np.inf, np.inf, 0, np.inf) + assert_allclose(res[0], ref[0], rtol=1e-6) + + ref = dblquad(lambda s2, x: dist.pdf(x, s2) * s2, -np.inf, np.inf, 0, np.inf) + assert_allclose(res[1], ref[0], rtol=1e-6) + + @pytest.mark.parametrize('dtype', [np.int32, np.float16, np.float32, np.float64]) + def test_dtype(self, dtype): + rng = np.random.default_rng(8925849245) + x, s2, mu, lmbda, a, b = rng.uniform(3, 10, size=6).astype(dtype) + dtype_out = np.result_type(1.0, dtype) + dist = stats.normal_inverse_gamma(mu, lmbda, a, b) + assert dist.rvs()[0].dtype == dtype_out + assert dist.rvs()[1].dtype == dtype_out + assert dist.mean()[0].dtype == dtype_out + assert dist.mean()[1].dtype == dtype_out + assert dist.var()[0].dtype == dtype_out + assert dist.var()[1].dtype == dtype_out + assert dist.logpdf(x, s2).dtype == dtype_out + assert dist.pdf(x, s2).dtype == dtype_out diff --git a/v12x/pylib/scipy/stats/tests/test_new_distributions.py b/v12x/pylib/scipy/stats/tests/test_new_distributions.py new file mode 100644 index 0000000000000000000000000000000000000000..b266b07b858ee1571e988c59af23e6348a219e4d --- /dev/null +++ b/v12x/pylib/scipy/stats/tests/test_new_distributions.py @@ -0,0 +1,17 @@ +# file for distribution-specific tests with new infrastructure (UnivariateDistribution) +import numpy as np +from numpy.testing import assert_allclose +from scipy import stats + +class TestBinomial: + def test_gh23708_binomial_logcdf_method_complement(self): + # gh-23708 found that `logcdf` method='complement' was inaccurate in the tails + x = np.asarray([0., 18.]) + X = stats.Binomial(n=np.asarray([18.]), p=np.asarray(0.71022842)) + assert_allclose(X.logcdf(x, method='complement'), X.logcdf(x), rtol=1e-15) + assert_allclose(X.logccdf(x, method='complement'), X.logccdf(x), rtol=1e-15) + + # going even deeper into the tails + X = stats.Binomial(n=100, p=0.5) + assert_allclose(X.logcdf(0, method='complement'), X.logpmf(0), rtol=1e-15) + assert_allclose(X.logccdf(99, method='complement'), X.logpmf(100), rtol=1e-15) diff --git a/v12x/pylib/scipy/stats/tests/test_odds_ratio.py b/v12x/pylib/scipy/stats/tests/test_odds_ratio.py new file mode 100644 index 0000000000000000000000000000000000000000..14b5ca88d0fca9edb8c3c161c06e52054d594e62 --- /dev/null +++ b/v12x/pylib/scipy/stats/tests/test_odds_ratio.py @@ -0,0 +1,148 @@ +import pytest +import numpy as np +from numpy.testing import assert_equal, assert_allclose +from .._discrete_distns import nchypergeom_fisher, hypergeom +from scipy.stats._odds_ratio import odds_ratio +from .data.fisher_exact_results_from_r import data + + +class TestOddsRatio: + + @pytest.mark.parametrize('parameters, rresult', data) + def test_results_from_r(self, parameters, rresult): + alternative = parameters.alternative.replace('.', '-') + result = odds_ratio(parameters.table) + # The results computed by R are not very accurate. + if result.statistic < 400: + or_rtol = 5e-4 + ci_rtol = 2e-2 + else: + or_rtol = 5e-2 + ci_rtol = 1e-1 + assert_allclose(result.statistic, + rresult.conditional_odds_ratio, rtol=or_rtol) + ci = result.confidence_interval(parameters.confidence_level, + alternative) + assert_allclose((ci.low, ci.high), rresult.conditional_odds_ratio_ci, + rtol=ci_rtol) + + # Also do a self-check for the conditional odds ratio. + # With the computed conditional odds ratio as the noncentrality + # parameter of the noncentral hypergeometric distribution with + # parameters table.sum(), table[0].sum(), and table[:,0].sum() as + # total, ngood and nsample, respectively, the mean of the distribution + # should equal table[0, 0]. + cor = result.statistic + table = np.array(parameters.table) + total = table.sum() + ngood = table[0].sum() + nsample = table[:, 0].sum() + # nchypergeom_fisher does not allow the edge cases where the + # noncentrality parameter is 0 or inf, so handle those values + # separately here. + if cor == 0: + nchg_mean = hypergeom.support(total, ngood, nsample)[0] + elif cor == np.inf: + nchg_mean = hypergeom.support(total, ngood, nsample)[1] + else: + nchg_mean = nchypergeom_fisher.mean(total, ngood, nsample, cor) + assert_allclose(nchg_mean, table[0, 0], rtol=1e-13) + + # Check that the confidence interval is correct. + alpha = 1 - parameters.confidence_level + if alternative == 'two-sided': + if ci.low > 0: + sf = nchypergeom_fisher.sf(table[0, 0] - 1, + total, ngood, nsample, ci.low) + assert_allclose(sf, alpha/2, rtol=1e-11) + if np.isfinite(ci.high): + cdf = nchypergeom_fisher.cdf(table[0, 0], + total, ngood, nsample, ci.high) + assert_allclose(cdf, alpha/2, rtol=1e-11) + elif alternative == 'less': + if np.isfinite(ci.high): + cdf = nchypergeom_fisher.cdf(table[0, 0], + total, ngood, nsample, ci.high) + assert_allclose(cdf, alpha, rtol=1e-11) + else: + # alternative == 'greater' + if ci.low > 0: + sf = nchypergeom_fisher.sf(table[0, 0] - 1, + total, ngood, nsample, ci.low) + assert_allclose(sf, alpha, rtol=1e-11) + + @pytest.mark.parametrize('table', [ + [[0, 0], [5, 10]], + [[5, 10], [0, 0]], + [[0, 5], [0, 10]], + [[5, 0], [10, 0]], + ]) + def test_row_or_col_zero(self, table): + result = odds_ratio(table) + assert_equal(result.statistic, np.nan) + ci = result.confidence_interval() + assert_equal((ci.low, ci.high), (0, np.inf)) + + @pytest.mark.parametrize("case", + [[0.95, 'two-sided', 0.4879913, 2.635883], + [0.90, 'two-sided', 0.5588516, 2.301663]]) + def test_sample_odds_ratio_ci(self, case): + # Compare the sample odds ratio confidence interval to the R function + # oddsratio.wald from the epitools package, e.g. + # > library(epitools) + # > table = matrix(c(10, 20, 41, 93), nrow=2, ncol=2, byrow=TRUE) + # > result = oddsratio.wald(table) + # > result$measure + # odds ratio with 95% C.I. + # Predictor estimate lower upper + # Exposed1 1.000000 NA NA + # Exposed2 1.134146 0.4879913 2.635883 + + confidence_level, alternative, ref_low, ref_high = case + table = [[10, 20], [41, 93]] + result = odds_ratio(table, kind='sample') + assert_allclose(result.statistic, 1.134146, rtol=1e-6) + ci = result.confidence_interval(confidence_level, alternative) + assert_allclose([ci.low, ci.high], [ref_low, ref_high], rtol=1e-6) + + @pytest.mark.slow + @pytest.mark.parametrize('alternative', ['less', 'greater', 'two-sided']) + def test_sample_odds_ratio_one_sided_ci(self, alternative): + # can't find a good reference for one-sided CI, so bump up the sample + # size and compare against the conditional odds ratio CI + table = [[1000, 2000], [4100, 9300]] + res = odds_ratio(table, kind='sample') + ref = odds_ratio(table, kind='conditional') + assert_allclose(res.statistic, ref.statistic, atol=1e-5) + assert_allclose(res.confidence_interval(alternative=alternative), + ref.confidence_interval(alternative=alternative), + atol=2e-3) + + @pytest.mark.parametrize('kind', ['sample', 'conditional']) + @pytest.mark.parametrize('bad_table', [123, "foo", [10, 11, 12]]) + def test_invalid_table_shape(self, kind, bad_table): + with pytest.raises(ValueError, match="Invalid shape"): + odds_ratio(bad_table, kind=kind) + + def test_invalid_table_type(self): + with pytest.raises(ValueError, match='must be an array of integers'): + odds_ratio([[1.0, 3.4], [5.0, 9.9]]) + + def test_negative_table_values(self): + with pytest.raises(ValueError, match='must be nonnegative'): + odds_ratio([[1, 2], [3, -4]]) + + def test_invalid_kind(self): + with pytest.raises(ValueError, match='`kind` must be'): + odds_ratio([[10, 20], [30, 14]], kind='magnetoreluctance') + + def test_invalid_alternative(self): + result = odds_ratio([[5, 10], [2, 32]]) + with pytest.raises(ValueError, match='`alternative` must be'): + result.confidence_interval(alternative='depleneration') + + @pytest.mark.parametrize('level', [-0.5, 1.5]) + def test_invalid_confidence_level(self, level): + result = odds_ratio([[5, 10], [2, 32]]) + with pytest.raises(ValueError, match='must be between 0 and 1'): + result.confidence_interval(confidence_level=level) diff --git a/v12x/pylib/scipy/stats/tests/test_qmc.py b/v12x/pylib/scipy/stats/tests/test_qmc.py new file mode 100644 index 0000000000000000000000000000000000000000..650a16e975c30e6df5b5a2b197d92d2a7f972ed3 --- /dev/null +++ b/v12x/pylib/scipy/stats/tests/test_qmc.py @@ -0,0 +1,1535 @@ +import os +from collections import Counter +from itertools import combinations, product + +import pytest +import numpy as np +from numpy.testing import (assert_allclose, assert_equal, assert_array_equal, + assert_array_less) + +from scipy.spatial import distance +from scipy.stats import shapiro +from scipy.stats._sobol import _test_find_index, _test_low_0_bit +from scipy.stats import qmc +from scipy.stats._qmc import ( + van_der_corput, n_primes, primes_from_2_to, + update_discrepancy, QMCEngine, _l1_norm, + _perturb_discrepancy, _lloyd_centroidal_voronoi_tessellation +) + + +class TestUtils: + def test_scale(self): + # 1d scalar + space = [[0], [1], [0.5]] + out = [[-2], [6], [2]] + scaled_space = qmc.scale(space, l_bounds=-2, u_bounds=6) + + assert_allclose(scaled_space, out) + + # 2d space + space = [[0, 0], [1, 1], [0.5, 0.5]] + bounds = np.array([[-2, 0], [6, 5]]) + out = [[-2, 0], [6, 5], [2, 2.5]] + + scaled_space = qmc.scale(space, l_bounds=bounds[0], u_bounds=bounds[1]) + + assert_allclose(scaled_space, out) + + scaled_back_space = qmc.scale(scaled_space, l_bounds=bounds[0], + u_bounds=bounds[1], reverse=True) + assert_allclose(scaled_back_space, space) + + # broadcast + space = [[0, 0, 0], [1, 1, 1], [0.5, 0.5, 0.5]] + l_bounds, u_bounds = 0, [6, 5, 3] + out = [[0, 0, 0], [6, 5, 3], [3, 2.5, 1.5]] + + scaled_space = qmc.scale(space, l_bounds=l_bounds, u_bounds=u_bounds) + + assert_allclose(scaled_space, out) + + def test_scale_random(self): + rng = np.random.default_rng(317589836511269190194010915937762468165) + sample = rng.random((30, 10)) + a = -rng.random(10) * 10 + b = rng.random(10) * 10 + scaled = qmc.scale(sample, a, b, reverse=False) + unscaled = qmc.scale(scaled, a, b, reverse=True) + assert_allclose(unscaled, sample) + + def test_scale_errors(self): + with pytest.raises(ValueError, match=r"Sample is not a 2D array"): + space = [0, 1, 0.5] + qmc.scale(space, l_bounds=-2, u_bounds=6) + + with pytest.raises(ValueError, match=r"Bounds are not consistent"): + space = [[0, 0], [1, 1], [0.5, 0.5]] + bounds = np.array([[-2, 6], [6, 5]]) + qmc.scale(space, l_bounds=bounds[0], u_bounds=bounds[1]) + + with pytest.raises(ValueError, match=r"'l_bounds' and 'u_bounds'" + r" must be broadcastable"): + space = [[0, 0], [1, 1], [0.5, 0.5]] + l_bounds, u_bounds = [-2, 0, 2], [6, 5] + qmc.scale(space, l_bounds=l_bounds, u_bounds=u_bounds) + + with pytest.raises(ValueError, match=r"'l_bounds' and 'u_bounds'" + r" must be broadcastable"): + space = [[0, 0], [1, 1], [0.5, 0.5]] + bounds = np.array([[-2, 0, 2], [6, 5, 5]]) + qmc.scale(space, l_bounds=bounds[0], u_bounds=bounds[1]) + + with pytest.raises(ValueError, match=r"Sample is not in unit " + r"hypercube"): + space = [[0, 0], [1, 1.5], [0.5, 0.5]] + bounds = np.array([[-2, 0], [6, 5]]) + qmc.scale(space, l_bounds=bounds[0], u_bounds=bounds[1]) + + with pytest.raises(ValueError, match=r"Sample is out of bounds"): + out = [[-2, 0], [6, 5], [8, 2.5]] + bounds = np.array([[-2, 0], [6, 5]]) + qmc.scale(out, l_bounds=bounds[0], u_bounds=bounds[1], + reverse=True) + + def test_discrepancy(self): + space_1 = np.array([[1, 3], [2, 6], [3, 2], [4, 5], [5, 1], [6, 4]]) + space_1 = (2.0 * space_1 - 1.0) / (2.0 * 6.0) + space_2 = np.array([[1, 5], [2, 4], [3, 3], [4, 2], [5, 1], [6, 6]]) + space_2 = (2.0 * space_2 - 1.0) / (2.0 * 6.0) + + # From Fang et al. Design and modeling for computer experiments, 2006 + assert_allclose(qmc.discrepancy(space_1), 0.0081, atol=1e-4) + assert_allclose(qmc.discrepancy(space_2), 0.0105, atol=1e-4) + + # From Zhou Y.-D. et al. Mixture discrepancy for quasi-random point + # sets. Journal of Complexity, 29 (3-4), pp. 283-301, 2013. + # Example 4 on Page 298 + sample = np.array([[2, 1, 1, 2, 2, 2], + [1, 2, 2, 2, 2, 2], + [2, 1, 1, 1, 1, 1], + [1, 1, 1, 1, 2, 2], + [1, 2, 2, 2, 1, 1], + [2, 2, 2, 2, 1, 1], + [2, 2, 2, 1, 2, 2]]) + sample = (2.0 * sample - 1.0) / (2.0 * 2.0) + + assert_allclose(qmc.discrepancy(sample, method='MD'), 2.5000, + atol=1e-4) + assert_allclose(qmc.discrepancy(sample, method='WD'), 1.3680, + atol=1e-4) + assert_allclose(qmc.discrepancy(sample, method='CD'), 0.3172, + atol=1e-4) + + # From Tim P. et al. Minimizing the L2 and Linf star discrepancies + # of a single point in the unit hypercube. JCAM, 2005 + # Table 1 on Page 283 + for dim in [2, 4, 8, 16, 32, 64]: + ref = np.sqrt(3**(-dim)) + assert_allclose(qmc.discrepancy(np.array([[1]*dim]), + method='L2-star'), ref) + + def test_discrepancy_errors(self): + sample = np.array([[1, 3], [2, 6], [3, 2], [4, 5], [5, 1], [6, 4]]) + + with pytest.raises( + ValueError, match=r"Sample is not in unit hypercube" + ): + qmc.discrepancy(sample) + + with pytest.raises(ValueError, match=r"Sample is not a 2D array"): + qmc.discrepancy([1, 3]) + + sample = [[0, 0], [1, 1], [0.5, 0.5]] + with pytest.raises(ValueError, match=r"'toto' is not a valid ..."): + qmc.discrepancy(sample, method="toto") + + def test_discrepancy_parallel(self, monkeypatch): + sample = np.array([[2, 1, 1, 2, 2, 2], + [1, 2, 2, 2, 2, 2], + [2, 1, 1, 1, 1, 1], + [1, 1, 1, 1, 2, 2], + [1, 2, 2, 2, 1, 1], + [2, 2, 2, 2, 1, 1], + [2, 2, 2, 1, 2, 2]]) + sample = (2.0 * sample - 1.0) / (2.0 * 2.0) + + assert_allclose(qmc.discrepancy(sample, method='MD', workers=8), + 2.5000, + atol=1e-4) + assert_allclose(qmc.discrepancy(sample, method='WD', workers=8), + 1.3680, + atol=1e-4) + assert_allclose(qmc.discrepancy(sample, method='CD', workers=8), + 0.3172, + atol=1e-4) + + # From Tim P. et al. Minimizing the L2 and Linf star discrepancies + # of a single point in the unit hypercube. JCAM, 2005 + # Table 1 on Page 283 + for dim in [2, 4, 8, 16, 32, 64]: + ref = np.sqrt(3 ** (-dim)) + assert_allclose(qmc.discrepancy(np.array([[1] * dim]), + method='L2-star', workers=-1), ref) + + monkeypatch.setattr(os, 'cpu_count', lambda: None) + with pytest.raises(NotImplementedError, match="Cannot determine the"): + qmc.discrepancy(sample, workers=-1) + + with pytest.raises(ValueError, match="Invalid number of workers..."): + qmc.discrepancy(sample, workers=-2) + + def test_geometric_discrepancy_errors(self): + sample = np.array([[1, 3], [2, 6], [3, 2], [4, 5], [5, 1], [6, 4]]) + + with pytest.raises(ValueError, match=r"Sample is not in unit hypercube"): + qmc.geometric_discrepancy(sample) + + with pytest.raises(ValueError, match=r"Sample is not a 2D array"): + qmc.geometric_discrepancy([1, 3]) + + sample = [[0, 0], [1, 1], [0.5, 0.5]] + with pytest.raises(ValueError, match=r"'toto' is not a valid ..."): + qmc.geometric_discrepancy(sample, method="toto") + + sample = np.array([[0, 0], [0, 0], [0, 1]]) + with pytest.warns(UserWarning, match="Sample contains duplicate points."): + qmc.geometric_discrepancy(sample) + + sample = np.array([[0.5, 0.5]]) + with pytest.raises(ValueError, match="Sample must contain at least two points"): + qmc.geometric_discrepancy(sample) + + def test_geometric_discrepancy(self): + sample = np.array([[0, 0], [1, 1]]) + assert_allclose(qmc.geometric_discrepancy(sample), np.sqrt(2)) + assert_allclose(qmc.geometric_discrepancy(sample, method="mst"), np.sqrt(2)) + + sample = np.array([[0, 0], [0, 1], [0.5, 1]]) + assert_allclose(qmc.geometric_discrepancy(sample), 0.5) + assert_allclose(qmc.geometric_discrepancy(sample, method="mst"), 0.75) + + sample = np.array([[0, 0], [0.25, 0.25], [1, 1]]) + assert_allclose(qmc.geometric_discrepancy(sample), np.sqrt(2) / 4) + assert_allclose(qmc.geometric_discrepancy(sample, method="mst"), np.sqrt(2) / 2) + assert_allclose(qmc.geometric_discrepancy(sample, metric="chebyshev"), 0.25) + assert_allclose( + qmc.geometric_discrepancy(sample, method="mst", metric="chebyshev"), 0.5 + ) + + rng = np.random.default_rng(191468432622931918890291693003068437394) + sample = qmc.LatinHypercube(d=3, rng=rng).random(50) + assert_allclose(qmc.geometric_discrepancy(sample), 0.05106012076093356) + assert_allclose( + qmc.geometric_discrepancy(sample, method='mst'), 0.19704396643366182 + ) + + @pytest.mark.xfail( + reason="minimum_spanning_tree ignores zero distances (#18892)", + strict=True, + ) + @pytest.mark.thread_unsafe + def test_geometric_discrepancy_mst_with_zero_distances(self): + sample = np.array([[0, 0], [0, 0], [0, 1]]) + assert_allclose(qmc.geometric_discrepancy(sample, method='mst'), 0.5) + + def test_update_discrepancy(self): + # From Fang et al. Design and modeling for computer experiments, 2006 + space_1 = np.array([[1, 3], [2, 6], [3, 2], [4, 5], [5, 1], [6, 4]]) + space_1 = (2.0 * space_1 - 1.0) / (2.0 * 6.0) + + disc_init = qmc.discrepancy(space_1[:-1], iterative=True) + disc_iter = update_discrepancy(space_1[-1], space_1[:-1], disc_init) + + assert_allclose(disc_iter, 0.0081, atol=1e-4) + + # n QMCEngine: + # preserve use of `seed` during SPEC 7 transition because + # some tests rely on behavior with integer `seed` (which is + # different from behavior with integer `rng`) + if self.can_scramble: + return self.qmce(scramble=scramble, seed=rng, **kwargs) + else: + if scramble: + pytest.skip() + else: + return self.qmce(seed=rng, **kwargs) + + def reference(self, scramble: bool) -> np.ndarray: + return self.scramble_nd if scramble else self.unscramble_nd + + @pytest.mark.parametrize("scramble", scramble, ids=ids) + def test_0dim(self, scramble): + engine = self.engine(d=0, scramble=scramble) + sample = engine.random(4) + assert_array_equal(np.empty((4, 0)), sample) + + @pytest.mark.parametrize("scramble", scramble, ids=ids) + def test_0sample(self, scramble): + engine = self.engine(d=2, scramble=scramble) + sample = engine.random(0) + assert_array_equal(np.empty((0, 2)), sample) + + @pytest.mark.parametrize("scramble", scramble, ids=ids) + def test_1sample(self, scramble): + engine = self.engine(d=2, scramble=scramble) + sample = engine.random(1) + assert (1, 2) == sample.shape + + @pytest.mark.parametrize("scramble", scramble, ids=ids) + def test_bounds(self, scramble): + engine = self.engine(d=100, scramble=scramble) + sample = engine.random(512) + assert np.all(sample >= 0) + assert np.all(sample <= 1) + + @pytest.mark.parametrize("scramble", scramble, ids=ids) + def test_sample(self, scramble): + ref_sample = self.reference(scramble=scramble) + engine = self.engine(d=2, scramble=scramble) + sample = engine.random(n=len(ref_sample)) + + assert_allclose(sample, ref_sample, atol=1e-1) + assert engine.num_generated == len(ref_sample) + + @pytest.mark.parametrize("scramble", scramble, ids=ids) + def test_continuing(self, scramble): + engine = self.engine(d=2, scramble=scramble) + ref_sample = engine.random(n=8) + + engine = self.engine(d=2, scramble=scramble) + + n_half = len(ref_sample) // 2 + + _ = engine.random(n=n_half) + sample = engine.random(n=n_half) + assert_allclose(sample, ref_sample[n_half:], atol=1e-1) + + @pytest.mark.parametrize("scramble", scramble, ids=ids) + @pytest.mark.parametrize( + "rng", + ( + 170382760648021597650530316304495310428, + lambda: np.random.default_rng(170382760648021597650530316304495310428), + pytest.param(None, marks=pytest.mark.thread_unsafe), + ), + ) + def test_reset(self, scramble, rng): + if callable(rng): + # Initialize local RNG here to make it thread-local in pytest-run-parallel + rng = rng() + engine = self.engine(d=2, scramble=scramble, rng=rng) + ref_sample = engine.random(n=8) + + engine.reset() + assert engine.num_generated == 0 + + sample = engine.random(n=8) + assert_allclose(sample, ref_sample) + + @pytest.mark.parametrize("scramble", scramble, ids=ids) + def test_fast_forward(self, scramble): + engine = self.engine(d=2, scramble=scramble) + ref_sample = engine.random(n=8) + + engine = self.engine(d=2, scramble=scramble) + + engine.fast_forward(4) + sample = engine.random(n=4) + + assert_allclose(sample, ref_sample[4:], atol=1e-1) + + # alternate fast forwarding with sampling + engine.reset() + even_draws = [] + for i in range(8): + if i % 2 == 0: + even_draws.append(engine.random()) + else: + engine.fast_forward(1) + assert_allclose( + ref_sample[[i for i in range(8) if i % 2 == 0]], + np.concatenate(even_draws), + atol=1e-5 + ) + + @pytest.mark.parametrize("scramble", [True]) + def test_distribution(self, scramble): + d = 50 + engine = self.engine(d=d, scramble=scramble) + sample = engine.random(1024) + assert_allclose( + np.mean(sample, axis=0), np.repeat(0.5, d), atol=1e-2 + ) + assert_allclose( + np.percentile(sample, 25, axis=0), np.repeat(0.25, d), atol=1e-2 + ) + assert_allclose( + np.percentile(sample, 75, axis=0), np.repeat(0.75, d), atol=1e-2 + ) + + def test_raises_optimizer(self): + message = r"'toto' is not a valid optimization method" + with pytest.raises(ValueError, match=message): + self.engine(d=1, scramble=False, optimization="toto") + + @pytest.mark.parametrize( + "optimization,metric", + [ + ("random-CD", qmc.discrepancy), + ("lloyd", lambda sample: -_l1_norm(sample))] + ) + def test_optimizers(self, optimization, metric): + engine = self.engine(d=2, scramble=False) + sample_ref = engine.random(n=64) + metric_ref = metric(sample_ref) + + optimal_ = self.engine(d=2, scramble=False, optimization=optimization) + sample_ = optimal_.random(n=64) + metric_ = metric(sample_) + + assert metric_ < metric_ref + + def test_consume_prng_state(self): + rng = np.random.default_rng(0xa29cabb11cfdf44ff6cac8bec254c2a0) + sample = [] + for i in range(3): + engine = self.engine(d=2, scramble=True, rng=rng) + sample.append(engine.random(4)) + + with pytest.raises(AssertionError, match="Arrays are not equal"): + assert_equal(sample[0], sample[1]) + with pytest.raises(AssertionError, match="Arrays are not equal"): + assert_equal(sample[0], sample[2]) + + +class TestHalton(QMCEngineTests): + qmce = qmc.Halton + can_scramble = True + # theoretical values known from Van der Corput + unscramble_nd = np.array([[0, 0], [1 / 2, 1 / 3], + [1 / 4, 2 / 3], [3 / 4, 1 / 9], + [1 / 8, 4 / 9], [5 / 8, 7 / 9], + [3 / 8, 2 / 9], [7 / 8, 5 / 9]]) + # theoretical values unknown: convergence properties checked + scramble_nd = np.array([[0.50246036, 0.93382481], + [0.00246036, 0.26715815], + [0.75246036, 0.60049148], + [0.25246036, 0.8227137 ], + [0.62746036, 0.15604704], + [0.12746036, 0.48938037], + [0.87746036, 0.71160259], + [0.37746036, 0.04493592]]) + + def test_workers(self): + ref_sample = self.reference(scramble=True) + engine = self.engine(d=2, scramble=True) + sample = engine.random(n=len(ref_sample), workers=8) + + assert_allclose(sample, ref_sample, atol=1e-3) + + # worker + integers + engine.reset() + ref_sample = engine.integers(10) + engine.reset() + sample = engine.integers(10, workers=8) + assert_equal(sample, ref_sample) + + +class TestLHS(QMCEngineTests): + qmce = qmc.LatinHypercube + can_scramble = True + + def test_continuing(self, *args): + pytest.skip("Not applicable: not a sequence.") + + def test_fast_forward(self, *args): + pytest.skip("Not applicable: not a sequence.") + + def test_sample(self, *args): + pytest.skip("Not applicable: the value of reference sample is" + " implementation dependent.") + + @pytest.mark.parametrize("strength", [1, 2]) + @pytest.mark.parametrize("scramble", [False, True]) + @pytest.mark.parametrize("optimization", [None, "random-CD"]) + def test_sample_stratified(self, optimization, scramble, strength): + rng = np.random.default_rng(37511836202578819870665127532742111260) + p = 5 + n = p**2 + d = 6 + + engine = qmc.LatinHypercube(d=d, scramble=scramble, + strength=strength, + optimization=optimization, + rng=rng) + sample = engine.random(n=n) + assert sample.shape == (n, d) + assert engine.num_generated == n + + # centering stratifies samples in the middle of equal segments: + # * inter-sample distance is constant in 1D sub-projections + # * after ordering, columns are equal + expected1d = (np.arange(n) + 0.5) / n + expected = np.broadcast_to(expected1d, (d, n)).T + assert np.any(sample != expected) + + sorted_sample = np.sort(sample, axis=0) + tol = 0.5 / n if scramble else 0 + + assert_allclose(sorted_sample, expected, atol=tol) + assert np.any(sample - expected > tol) + + if strength == 2 and optimization is None: + unique_elements = np.arange(p) + desired = set(product(unique_elements, unique_elements)) + + for i, j in combinations(range(engine.d), 2): + samples_2d = sample[:, [i, j]] + res = (samples_2d * p).astype(int) + res_set = {tuple(row) for row in res} + assert_equal(res_set, desired) + + def test_optimizer_1d(self): + # discrepancy measures are invariant under permuting factors and runs + engine = self.engine(d=1, scramble=False) + sample_ref = engine.random(n=64) + + optimal_ = self.engine(d=1, scramble=False, optimization="random-CD") + sample_ = optimal_.random(n=64) + + assert_array_equal(sample_ref, sample_) + + def test_raises(self): + message = r"not a valid strength" + with pytest.raises(ValueError, match=message): + qmc.LatinHypercube(1, strength=3) + + message = r"n is not the square of a prime number" + with pytest.raises(ValueError, match=message): + engine = qmc.LatinHypercube(d=2, strength=2) + engine.random(16) + + message = r"n is not the square of a prime number" + with pytest.raises(ValueError, match=message): + engine = qmc.LatinHypercube(d=2, strength=2) + engine.random(5) # because int(sqrt(5)) would result in 2 + + message = r"n is too small for d" + with pytest.raises(ValueError, match=message): + engine = qmc.LatinHypercube(d=5, strength=2) + engine.random(9) + + +class TestSobol(QMCEngineTests): + qmce = qmc.Sobol + can_scramble = True + # theoretical values from Joe Kuo2010 + unscramble_nd = np.array([[0., 0.], + [0.5, 0.5], + [0.75, 0.25], + [0.25, 0.75], + [0.375, 0.375], + [0.875, 0.875], + [0.625, 0.125], + [0.125, 0.625]]) + + # theoretical values unknown: convergence properties checked + scramble_nd = np.array([[0.25331921, 0.41371179], + [0.8654213, 0.9821167], + [0.70097554, 0.03664616], + [0.18027647, 0.60895735], + [0.10521339, 0.21897069], + [0.53019685, 0.66619033], + [0.91122276, 0.34580743], + [0.45337471, 0.78912079]]) + + def test_warning(self): + with pytest.warns(UserWarning, match=r"The balance properties of " + r"Sobol' points"): + engine = qmc.Sobol(1) + engine.random(10) + + def test_random_base2(self): + engine = qmc.Sobol(2, scramble=False) + sample = engine.random_base2(2) + assert_array_equal(self.unscramble_nd[:4], sample) + + # resampling still having N=2**n + sample = engine.random_base2(2) + assert_array_equal(self.unscramble_nd[4:8], sample) + + # resampling again but leading to N!=2**n + with pytest.raises(ValueError, match=r"The balance properties of " + r"Sobol' points"): + engine.random_base2(2) + + def test_raise(self): + with pytest.raises(ValueError, match=r"Maximum supported " + r"dimensionality"): + qmc.Sobol(qmc.Sobol.MAXDIM + 1) + + with pytest.raises(ValueError, match=r"Maximum supported " + r"'bits' is 64"): + qmc.Sobol(1, bits=65) + + def test_high_dim(self): + engine = qmc.Sobol(1111, scramble=False) + count1 = Counter(engine.random().flatten().tolist()) + count2 = Counter(engine.random().flatten().tolist()) + assert_equal(count1, Counter({0.0: 1111})) + assert_equal(count2, Counter({0.5: 1111})) + + @pytest.mark.parametrize("bits", [2, 3]) + def test_bits(self, bits): + engine = qmc.Sobol(2, scramble=False, bits=bits) + ns = 2**bits + sample = engine.random(ns) + assert_array_equal(self.unscramble_nd[:ns], sample) + + with pytest.raises(ValueError, match="increasing `bits`"): + engine.random() + + def test_64bits(self): + engine = qmc.Sobol(2, scramble=False, bits=64) + sample = engine.random(8) + assert_array_equal(self.unscramble_nd, sample) + + +class TestLow0Bit: + def test_examples(self): + test_vector = [ + # from low_0_bit's docstring + (0b0000, 1), + (0b0001, 2), + (0b0010, 1), + (0b0101, 2), + (0b0111, 4), + # gh-23409 + (2 ** 32 - 1, 33), + (2 ** 32, 1), + (2 ** 33 - 1, 34), + (2 ** 64 - 1, 65), + ] + for in_, out in test_vector: + assert_equal(_test_low_0_bit(in_), out) + + +class TestPoisson(QMCEngineTests): + qmce = qmc.PoissonDisk + can_scramble = False + + def test_bounds(self, *args): + pytest.skip("Too costly in memory.") + + def test_fast_forward(self, *args): + pytest.skip("Not applicable: recursive process.") + + def test_sample(self, *args): + pytest.skip("Not applicable: the value of reference sample is" + " implementation dependent.") + + def test_continuing(self, *args): + # can continue a sampling, but will not preserve the same order + # because candidates are lost, so we will not select the same center + radius = 0.05 + ns = 6 + engine = self.engine(d=2, radius=radius, scramble=False) + + sample_init = engine.random(n=ns) + assert len(sample_init) <= ns + assert l2_norm(sample_init) >= radius + + sample_continued = engine.random(n=ns) + assert len(sample_continued) <= ns + assert l2_norm(sample_continued) >= radius + + sample = np.concatenate([sample_init, sample_continued], axis=0) + assert len(sample) <= ns * 2 + assert l2_norm(sample) >= radius + + def test_mindist(self): + rng = np.random.default_rng(132074951149370773672162394161442690287) + ns = 50 + + low, high = 0.08, 0.2 + radii = (high - low) * rng.random(5) + low + + dimensions = [1, 3, 4] + hypersphere_methods = ["volume", "surface"] + + gen = product(dimensions, radii, hypersphere_methods) + + for d, radius, hypersphere in gen: + engine = self.qmce( + d=d, radius=radius, hypersphere=hypersphere, rng=rng + ) + sample = engine.random(ns) + + assert len(sample) <= ns + assert l2_norm(sample) >= radius + + @pytest.mark.parametrize("l_bounds, u_bounds", # add bounds to test gh-22819 + [(None, None), ([-5, -5], [5, 5]), ([1.1, 1.1], [5.1, 5.1])]) + def test_fill_space(self, l_bounds, u_bounds): + radius = 0.2 + engine = self.qmce(d=2, radius=radius, l_bounds=l_bounds, u_bounds=u_bounds) + + sample = engine.fill_space() + # rtol to accommodate Accelerate; see gh-24486 + rtol = 1e-6 + assert l2_norm(sample) >= radius * (1 - rtol) + + def test_bounds_shift_scale(self): + # test a reasonable property: as long as shape of region is a hypercube, + # bounds just shift and scale the sample. + radius = 0.2 # arbitrary parameters + shift = np.asarray([-2.1, 3.4]) + scale = 2.6 + + rng = np.random.default_rng(8504196751) + engine = self.qmce(d=2, radius=radius*scale, rng=rng, + l_bounds=shift, u_bounds=shift + scale) + res = engine.fill_space() + + rng = np.random.default_rng(8504196751) + engine = self.qmce(d=2, radius=radius, rng=rng) + ref = engine.fill_space() + + # circle packing problem is np complex + assert l2_norm(res) >= radius + assert_allclose(res, ref*scale + shift) + + @pytest.mark.parametrize("l_bounds", [[-1, -2, -1], [1, 2, 1]]) + def test_sample_inside_lower_bounds(self, l_bounds): + radius = 0.2 + u_bounds=[3, 3, 2] + engine = self.qmce( + d=3, radius=radius, l_bounds=l_bounds, u_bounds=u_bounds + ) + sample = engine.random(30) + + for point in sample: + assert_array_less(point, u_bounds) + assert_array_less(l_bounds, point) + + @pytest.mark.parametrize("u_bounds", [[-1, -2, -1], [1, 2, 1]]) + def test_sample_inside_upper_bounds(self, u_bounds): + radius = 0.2 + l_bounds=[-3, -3, -2] + engine = self.qmce( + d=3, radius=radius, l_bounds=l_bounds, u_bounds=u_bounds + ) + sample = engine.random(30) + + for point in sample: + assert_array_less(point, u_bounds) + assert_array_less(l_bounds, point) + + def test_inconsistent_bound_value(self): + radius = 0.2 + l_bounds=[3, 2, 1] + u_bounds=[-1, -2, -1] + with pytest.raises( + ValueError, + match="Bounds are not consistent 'l_bounds' < 'u_bounds'"): + self.qmce(d=3, radius=radius, l_bounds=l_bounds, u_bounds=u_bounds) + + @pytest.mark.parametrize("u_bounds", [[-1, -2, -1], [-1, -2]]) + @pytest.mark.parametrize("l_bounds", [[3, 2]]) + def test_inconsistent_bounds(self, u_bounds, l_bounds): + radius = 0.2 + with pytest.raises( + ValueError, + match="'l_bounds' and 'u_bounds' must be broadcastable and respect" + " the sample dimension"): + self.qmce( + d=3, radius=radius, + l_bounds=l_bounds, u_bounds=u_bounds + ) + + def test_raises(self): + message = r"'toto' is not a valid hypersphere sampling" + with pytest.raises(ValueError, match=message): + qmc.PoissonDisk(1, hypersphere="toto") + + +class TestMultinomialQMC: + def test_validations(self): + # negative Ps + p = np.array([0.12, 0.26, -0.05, 0.35, 0.22]) + with pytest.raises(ValueError, match=r"Elements of pvals must " + r"be non-negative."): + qmc.MultinomialQMC(p, n_trials=10) + + # sum of P too large + p = np.array([0.12, 0.26, 0.1, 0.35, 0.22]) + message = r"Elements of pvals must sum to 1." + with pytest.raises(ValueError, match=message): + qmc.MultinomialQMC(p, n_trials=10) + + p = np.array([0.12, 0.26, 0.05, 0.35, 0.22]) + + message = r"Dimension of `engine` must be 1." + with pytest.raises(ValueError, match=message): + qmc.MultinomialQMC(p, n_trials=10, engine=qmc.Sobol(d=2)) + + message = r"`engine` must be an instance of..." + with pytest.raises(ValueError, match=message): + qmc.MultinomialQMC(p, n_trials=10, engine=np.random.default_rng()) + + @pytest.mark.filterwarnings('ignore::UserWarning') + def test_MultinomialBasicDraw(self): + rng = np.random.default_rng(6955663962957011631562466584467607969) + p = np.array([0.12, 0.26, 0.05, 0.35, 0.22]) + n_trials = 100 + expected = np.atleast_2d(n_trials * p).astype(int) + # preserve use of legacy keyword during SPEC 7 transition + engine = qmc.MultinomialQMC(p, n_trials=n_trials, seed=rng) + assert_allclose(engine.random(1), expected, atol=1) + + def test_MultinomialDistribution(self): + rng = np.random.default_rng(77797854505813727292048130876699859000) + p = np.array([0.12, 0.26, 0.05, 0.35, 0.22]) + engine = qmc.MultinomialQMC(p, n_trials=8192, rng=rng) + draws = engine.random(1) + assert_allclose(draws / np.sum(draws), np.atleast_2d(p), atol=1e-4) + + def test_FindIndex(self): + p_cumulative = np.array([0.1, 0.4, 0.45, 0.6, 0.75, 0.9, 0.99, 1.0]) + size = len(p_cumulative) + assert_equal(_test_find_index(p_cumulative, size, 0.0), 0) + assert_equal(_test_find_index(p_cumulative, size, 0.4), 2) + assert_equal(_test_find_index(p_cumulative, size, 0.44999), 2) + assert_equal(_test_find_index(p_cumulative, size, 0.45001), 3) + assert_equal(_test_find_index(p_cumulative, size, 1.0), size - 1) + + @pytest.mark.filterwarnings('ignore::UserWarning') + def test_other_engine(self): + # same as test_MultinomialBasicDraw with different engine + rng = np.random.default_rng(283753519042773243071753037669078065412) + p = np.array([0.12, 0.26, 0.05, 0.35, 0.22]) + n_trials = 100 + expected = np.atleast_2d(n_trials * p).astype(int) + base_engine = qmc.Sobol(1, scramble=True, rng=rng) + engine = qmc.MultinomialQMC(p, n_trials=n_trials, engine=base_engine, + rng=rng) + assert_allclose(engine.random(1), expected, atol=1) + + +class TestNormalQMC: + def test_NormalQMC(self): + # d = 1 + engine = qmc.MultivariateNormalQMC(mean=np.zeros(1)) + samples = engine.random() + assert_equal(samples.shape, (1, 1)) + samples = engine.random(n=5) + assert_equal(samples.shape, (5, 1)) + # d = 2 + engine = qmc.MultivariateNormalQMC(mean=np.zeros(2)) + samples = engine.random() + assert_equal(samples.shape, (1, 2)) + samples = engine.random(n=5) + assert_equal(samples.shape, (5, 2)) + + def test_NormalQMCInvTransform(self): + # d = 1 + engine = qmc.MultivariateNormalQMC( + mean=np.zeros(1), inv_transform=True) + samples = engine.random() + assert_equal(samples.shape, (1, 1)) + samples = engine.random(n=5) + assert_equal(samples.shape, (5, 1)) + # d = 2 + engine = qmc.MultivariateNormalQMC( + mean=np.zeros(2), inv_transform=True) + samples = engine.random() + assert_equal(samples.shape, (1, 2)) + samples = engine.random(n=5) + assert_equal(samples.shape, (5, 2)) + + def test_NormalQMCSeeded(self): + # test even dimension + rng = np.random.default_rng(274600237797326520096085022671371676017) + # preserve use of legacy keyword during SPEC 7 transition + engine = qmc.MultivariateNormalQMC( + mean=np.zeros(2), inv_transform=False, seed=rng) + samples = engine.random(n=2) + samples_expected = np.array([[-0.932001, -0.522923], + [-1.477655, 0.846851]]) + assert_allclose(samples, samples_expected, atol=1e-4) + + # test odd dimension + rng = np.random.default_rng(274600237797326520096085022671371676017) + engine = qmc.MultivariateNormalQMC( + mean=np.zeros(3), inv_transform=False, rng=rng) + samples = engine.random(n=2) + samples_expected = np.array([[-0.932001, -0.522923, 0.036578], + [-1.778011, 0.912428, -0.065421]]) + assert_allclose(samples, samples_expected, atol=1e-4) + + # same test with another engine + rng = np.random.default_rng(274600237797326520096085022671371676017) + base_engine = qmc.Sobol(4, scramble=True, rng=rng) + engine = qmc.MultivariateNormalQMC( + mean=np.zeros(3), inv_transform=False, + engine=base_engine, rng=rng + ) + samples = engine.random(n=2) + samples_expected = np.array([[-0.932001, -0.522923, 0.036578], + [-1.778011, 0.912428, -0.065421]]) + assert_allclose(samples, samples_expected, atol=1e-4) + + def test_NormalQMCSeededInvTransform(self): + # test even dimension + rng = np.random.default_rng(288527772707286126646493545351112463929) + engine = qmc.MultivariateNormalQMC( + mean=np.zeros(2), rng=rng, inv_transform=True) + samples = engine.random(n=2) + samples_expected = np.array([[-0.913237, -0.964026], + [0.255904, 0.003068]]) + assert_allclose(samples, samples_expected, atol=1e-4) + + # test odd dimension + rng = np.random.default_rng(288527772707286126646493545351112463929) + engine = qmc.MultivariateNormalQMC( + mean=np.zeros(3), rng=rng, inv_transform=True) + samples = engine.random(n=2) + samples_expected = np.array([[-0.913237, -0.964026, 0.355501], + [0.699261, 2.90213 , -0.6418]]) + assert_allclose(samples, samples_expected, atol=1e-4) + + def test_other_engine(self): + for d in (0, 1, 2): + base_engine = qmc.Sobol(d=d, scramble=False) + engine = qmc.MultivariateNormalQMC(mean=np.zeros(d), + engine=base_engine, + inv_transform=True) + samples = engine.random() + assert_equal(samples.shape, (1, d)) + + def test_NormalQMCShapiro(self): + rng = np.random.default_rng(13242) + engine = qmc.MultivariateNormalQMC(mean=np.zeros(2), rng=rng) + samples = engine.random(n=256) + assert all(np.abs(samples.mean(axis=0)) < 1e-2) + assert all(np.abs(samples.std(axis=0) - 1) < 1e-2) + # perform Shapiro-Wilk test for normality + for i in (0, 1): + _, pval = shapiro(samples[:, i]) + assert pval > 0.9 + # make sure samples are uncorrelated + cov = np.cov(samples.transpose()) + assert np.abs(cov[0, 1]) < 1e-2 + + def test_NormalQMCShapiroInvTransform(self): + rng = np.random.default_rng(32344554) + engine = qmc.MultivariateNormalQMC( + mean=np.zeros(2), inv_transform=True, rng=rng) + samples = engine.random(n=256) + assert all(np.abs(samples.mean(axis=0)) < 1e-2) + assert all(np.abs(samples.std(axis=0) - 1) < 1e-2) + # perform Shapiro-Wilk test for normality + for i in (0, 1): + _, pval = shapiro(samples[:, i]) + assert pval > 0.9 + # make sure samples are uncorrelated + cov = np.cov(samples.transpose()) + assert np.abs(cov[0, 1]) < 1e-2 + + +class TestMultivariateNormalQMC: + + def test_validations(self): + + message = r"Dimension of `engine` must be consistent" + with pytest.raises(ValueError, match=message): + qmc.MultivariateNormalQMC([0], engine=qmc.Sobol(d=2)) + + message = r"Dimension of `engine` must be consistent" + with pytest.raises(ValueError, match=message): + qmc.MultivariateNormalQMC([0, 0, 0], engine=qmc.Sobol(d=4)) + + message = r"`engine` must be an instance of..." + with pytest.raises(ValueError, match=message): + qmc.MultivariateNormalQMC([0, 0], engine=np.random.default_rng()) + + message = r"Covariance matrix not PSD." + with pytest.raises(ValueError, match=message): + qmc.MultivariateNormalQMC([0, 0], [[1, 2], [2, 1]]) + + message = r"Covariance matrix is not symmetric." + with pytest.raises(ValueError, match=message): + qmc.MultivariateNormalQMC([0, 0], [[1, 0], [2, 1]]) + + message = r"Dimension mismatch between mean and covariance." + with pytest.raises(ValueError, match=message): + qmc.MultivariateNormalQMC([0], [[1, 0], [0, 1]]) + + def test_MultivariateNormalQMCNonPD(self): + # try with non-pd but psd cov; should work + engine = qmc.MultivariateNormalQMC( + [0, 0, 0], [[1, 0, 1], [0, 1, 1], [1, 1, 2]], + ) + assert engine._corr_matrix is not None + + def test_MultivariateNormalQMC(self): + # d = 1 scalar + engine = qmc.MultivariateNormalQMC(mean=0, cov=5) + samples = engine.random() + assert_equal(samples.shape, (1, 1)) + samples = engine.random(n=5) + assert_equal(samples.shape, (5, 1)) + + # d = 2 list + engine = qmc.MultivariateNormalQMC(mean=[0, 1], cov=[[1, 0], [0, 1]]) + samples = engine.random() + assert_equal(samples.shape, (1, 2)) + samples = engine.random(n=5) + assert_equal(samples.shape, (5, 2)) + + # d = 3 np.array + mean = np.array([0, 1, 2]) + cov = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]]) + engine = qmc.MultivariateNormalQMC(mean, cov) + samples = engine.random() + assert_equal(samples.shape, (1, 3)) + samples = engine.random(n=5) + assert_equal(samples.shape, (5, 3)) + + def test_MultivariateNormalQMCInvTransform(self): + # d = 1 scalar + engine = qmc.MultivariateNormalQMC(mean=0, cov=5, inv_transform=True) + samples = engine.random() + assert_equal(samples.shape, (1, 1)) + samples = engine.random(n=5) + assert_equal(samples.shape, (5, 1)) + + # d = 2 list + engine = qmc.MultivariateNormalQMC( + mean=[0, 1], cov=[[1, 0], [0, 1]], inv_transform=True, + ) + samples = engine.random() + assert_equal(samples.shape, (1, 2)) + samples = engine.random(n=5) + assert_equal(samples.shape, (5, 2)) + + # d = 3 np.array + mean = np.array([0, 1, 2]) + cov = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]]) + engine = qmc.MultivariateNormalQMC(mean, cov, inv_transform=True) + samples = engine.random() + assert_equal(samples.shape, (1, 3)) + samples = engine.random(n=5) + assert_equal(samples.shape, (5, 3)) + + def test_MultivariateNormalQMCSeeded(self): + # test even dimension + rng = np.random.default_rng(180182791534511062935571481899241825000) + a = rng.standard_normal((2, 2)) + A = a @ a.transpose() + np.diag(rng.random(2)) + engine = qmc.MultivariateNormalQMC(np.array([0, 0]), A, + inv_transform=False, rng=rng) + samples = engine.random(n=2) + samples_expected = np.array([[-0.64419, -0.882413], + [0.837199, 2.045301]]) + assert_allclose(samples, samples_expected, atol=1e-4) + + # test odd dimension + rng = np.random.default_rng(180182791534511062935571481899241825000) + a = rng.standard_normal((3, 3)) + A = a @ a.transpose() + np.diag(rng.random(3)) + engine = qmc.MultivariateNormalQMC(np.array([0, 0, 0]), A, + inv_transform=False, rng=rng) + samples = engine.random(n=2) + samples_expected = np.array([[-0.693853, -1.265338, -0.088024], + [1.620193, 2.679222, 0.457343]]) + assert_allclose(samples, samples_expected, atol=1e-4) + + def test_MultivariateNormalQMCSeededInvTransform(self): + # test even dimension + rng = np.random.default_rng(224125808928297329711992996940871155974) + a = rng.standard_normal((2, 2)) + A = a @ a.transpose() + np.diag(rng.random(2)) + engine = qmc.MultivariateNormalQMC( + np.array([0, 0]), A, rng=rng, inv_transform=True + ) + samples = engine.random(n=2) + samples_expected = np.array([[0.682171, -3.114233], + [-0.098463, 0.668069]]) + assert_allclose(samples, samples_expected, atol=1e-4) + + # test odd dimension + rng = np.random.default_rng(224125808928297329711992996940871155974) + a = rng.standard_normal((3, 3)) + A = a @ a.transpose() + np.diag(rng.random(3)) + engine = qmc.MultivariateNormalQMC( + np.array([0, 0, 0]), A, rng=rng, inv_transform=True + ) + samples = engine.random(n=2) + samples_expected = np.array([[0.988061, -1.644089, -0.877035], + [-1.771731, 1.096988, 2.024744]]) + assert_allclose(samples, samples_expected, atol=1e-4) + + def test_MultivariateNormalQMCShapiro(self): + # test the standard case + rng = np.random.default_rng(188960007281846377164494575845971640) + engine = qmc.MultivariateNormalQMC( + mean=[0, 0], cov=[[1, 0], [0, 1]], rng=rng + ) + samples = engine.random(n=256) + assert all(np.abs(samples.mean(axis=0)) < 1e-2) + assert all(np.abs(samples.std(axis=0) - 1) < 1e-2) + # perform Shapiro-Wilk test for normality + for i in (0, 1): + _, pval = shapiro(samples[:, i]) + assert pval > 0.9 + # make sure samples are uncorrelated + cov = np.cov(samples.transpose()) + assert np.abs(cov[0, 1]) < 1e-2 + + # test the correlated, non-zero mean case + engine = qmc.MultivariateNormalQMC( + mean=[1.0, 2.0], cov=[[1.5, 0.5], [0.5, 1.5]], rng=rng + ) + samples = engine.random(n=256) + assert all(np.abs(samples.mean(axis=0) - [1, 2]) < 1e-2) + assert all(np.abs(samples.std(axis=0) - np.sqrt(1.5)) < 1e-2) + # perform Shapiro-Wilk test for normality + for i in (0, 1): + _, pval = shapiro(samples[:, i]) + assert pval > 0.9 + # check covariance + cov = np.cov(samples.transpose()) + assert np.abs(cov[0, 1] - 0.5) < 1e-2 + + def test_MultivariateNormalQMCShapiroInvTransform(self): + # test the standard case + rng = np.random.default_rng(200089821034563288698994840831440331329) + engine = qmc.MultivariateNormalQMC( + mean=[0, 0], cov=[[1, 0], [0, 1]], rng=rng, inv_transform=True + ) + samples = engine.random(n=256) + assert all(np.abs(samples.mean(axis=0)) < 1e-2) + assert all(np.abs(samples.std(axis=0) - 1) < 1e-2) + # perform Shapiro-Wilk test for normality + for i in (0, 1): + _, pval = shapiro(samples[:, i]) + assert pval > 0.9 + # make sure samples are uncorrelated + cov = np.cov(samples.transpose()) + assert np.abs(cov[0, 1]) < 1e-2 + + # test the correlated, non-zero mean case + engine = qmc.MultivariateNormalQMC( + mean=[1.0, 2.0], + cov=[[1.5, 0.5], [0.5, 1.5]], + rng=rng, + inv_transform=True, + ) + samples = engine.random(n=256) + assert all(np.abs(samples.mean(axis=0) - [1, 2]) < 1e-2) + assert all(np.abs(samples.std(axis=0) - np.sqrt(1.5)) < 1e-2) + # perform Shapiro-Wilk test for normality + for i in (0, 1): + _, pval = shapiro(samples[:, i]) + assert pval > 0.9 + # check covariance + cov = np.cov(samples.transpose()) + assert np.abs(cov[0, 1] - 0.5) < 1e-2 + + def test_MultivariateNormalQMCDegenerate(self): + # X, Y iid standard Normal and Z = X + Y, random vector (X, Y, Z) + rng = np.random.default_rng(16320637417581448357869821654290448620) + engine = qmc.MultivariateNormalQMC( + mean=[0.0, 0.0, 0.0], + cov=[[1.0, 0.0, 1.0], [0.0, 1.0, 1.0], [1.0, 1.0, 2.0]], + rng=rng, + ) + samples = engine.random(n=512) + assert all(np.abs(samples.mean(axis=0)) < 1e-2) + assert np.abs(np.std(samples[:, 0]) - 1) < 1e-2 + assert np.abs(np.std(samples[:, 1]) - 1) < 1e-2 + assert np.abs(np.std(samples[:, 2]) - np.sqrt(2)) < 1e-2 + for i in (0, 1, 2): + _, pval = shapiro(samples[:, i]) + assert pval > 0.8 + cov = np.cov(samples.transpose()) + assert np.abs(cov[0, 1]) < 1e-2 + assert np.abs(cov[0, 2] - 1) < 1e-2 + # check to see if X + Y = Z almost exactly + assert all(np.abs(samples[:, 0] + samples[:, 1] - samples[:, 2]) + < 1e-5) + + +class TestLloyd: + def test_lloyd(self): + # quite sensible seed as it can go up before going further down + rng = np.random.RandomState(1809831) + sample = rng.uniform(0, 1, size=(128, 2)) + base_l1 = _l1_norm(sample) + base_l2 = l2_norm(sample) + + for _ in range(4): + sample_lloyd = _lloyd_centroidal_voronoi_tessellation( + sample, maxiter=1, + ) + curr_l1 = _l1_norm(sample_lloyd) + curr_l2 = l2_norm(sample_lloyd) + + # higher is better for the distance measures + assert base_l1 < curr_l1 + assert base_l2 < curr_l2 + + base_l1 = curr_l1 + base_l2 = curr_l2 + + sample = sample_lloyd + + def test_lloyd_non_mutating(self): + """ + Verify that the input samples are not mutated in place and that they do + not share memory with the output. + """ + sample_orig = np.array([[0.1, 0.1], + [0.1, 0.2], + [0.2, 0.1], + [0.2, 0.2]]) + sample_copy = sample_orig.copy() + new_sample = _lloyd_centroidal_voronoi_tessellation( + sample=sample_orig + ) + assert_allclose(sample_orig, sample_copy) + assert not np.may_share_memory(sample_orig, new_sample) + + def test_lloyd_errors(self): + with pytest.raises(ValueError, match=r"`sample` is not a 2D array"): + sample = [0, 1, 0.5] + _lloyd_centroidal_voronoi_tessellation(sample) + + msg = r"`sample` dimension is not >= 2" + with pytest.raises(ValueError, match=msg): + sample = [[0], [0.4], [1]] + _lloyd_centroidal_voronoi_tessellation(sample) + + msg = r"`sample` is not in unit hypercube" + with pytest.raises(ValueError, match=msg): + sample = [[-1.1, 0], [0.1, 0.4], [1, 2]] + _lloyd_centroidal_voronoi_tessellation(sample) + + +# mindist +def l2_norm(sample): + return distance.pdist(sample).min() + + +@pytest.mark.parametrize('engine', [qmc.Halton, qmc.Sobol, + qmc.LatinHypercube, qmc.PoissonDisk]) +def test_deterministic(engine): + seed_number = 2359834584 + + rng = np.random.RandomState(seed_number) + res1 = engine(d=1, seed=rng).random(4) + rng = np.random.RandomState(seed_number) + res2 = engine(d=1, seed=rng).random(4) + assert_equal(res1, res2) + + rng = np.random.default_rng(seed_number) + res1 = engine(d=1, seed=rng).random(4) + res2 = engine(d=1, rng=seed_number).random(4) + assert_equal(res1, res2) + rng = np.random.default_rng(seed_number) + res3 = engine(d=1, rng=rng).random(4) + assert_equal(res2, res1) + assert_equal(res3, res1) + + message = "got multiple values for argument now known as `rng`" + with pytest.raises(TypeError, match=message): + engine(d=1, rng=seed_number, seed=seed_number) diff --git a/v12x/pylib/scipy/stats/tests/test_quantile.py b/v12x/pylib/scipy/stats/tests/test_quantile.py new file mode 100644 index 0000000000000000000000000000000000000000..4caa6341adb0123c783a60f324617d0fcaed7090 --- /dev/null +++ b/v12x/pylib/scipy/stats/tests/test_quantile.py @@ -0,0 +1,774 @@ +import pytest +import numpy as np + +from scipy import stats +from scipy.stats._quantile import (_xp_searchsorted, _estimated_cdf_methods, + _estimated_cdf_discontinuous_methods, _estimated_cdf_continuous_methods) +from scipy._lib._array_api import ( + xp_default_dtype, + is_numpy, + is_torch, + is_jax, + is_cupy, + is_array_api_strict, + make_xp_test_case, + SCIPY_ARRAY_API, + xp_size, + xp_copy, +) +from scipy._lib._array_api_no_0d import xp_assert_close, xp_assert_equal +from scipy._lib._util import _apply_over_batch +import scipy._external.array_api_extra as xpx +from scipy.stats._axis_nan_policy import _broadcast_arrays + +skip_xp_backends = pytest.mark.skip_xp_backends + +lazy_xp_modules = [stats] + +@_apply_over_batch(('x', 1), ('p', 1)) +def quantile_reference_last_axis(x, p, nan_policy, method): + if nan_policy == 'omit': + x = x[~np.isnan(x)] + p_mask = np.isnan(p) + p = p.copy() + p[p_mask] = 0.5 + if method == 'harrell-davis': + # hdquantiles returns masked element if length along axis is 1 (bug) + res = (np.full_like(p, x[0]) if x.size == 1 + else stats.mstats.hdquantiles(x, p).data) + elif method.startswith('round'): + res = winsor_reference_1d(np.sort(x), p, method) + else: + res = np.quantile(x, p, method=method) + + res = np.asarray(res) + if nan_policy == 'propagate' and np.any(np.isnan(x)): + res[:] = np.nan + + res[p_mask] = np.nan + return res + + +@np.vectorize(excluded={0, 2}) # type: ignore[call-arg] +def winsor_reference_1d(y, p, method): + # Adapted directly from the documentation + # Note: `y` is the sorted data array + n = len(y) + if method == 'round_nearest': + j = int(np.round(p * n) if p < 0.5 else np.round(n * p - 1)) + elif method == 'round_outward': + j = int(np.floor(p * n) if p < 0.5 else np.ceil(n * p - 1)) + elif method == 'round_inward': + j = int(np.ceil(p * n) if p < 0.5 else np.floor(n * p - 1)) + return y[j] + + +def quantile_reference(x, p, *, axis, nan_policy, keepdims, method): + x, p = np.moveaxis(x, axis, -1), np.moveaxis(np.atleast_1d(p), axis, -1) + res = quantile_reference_last_axis(x, p, nan_policy, method) + res = np.moveaxis(res, -1, axis) + if not keepdims: + res = np.squeeze(res, axis=axis) + return res + + +@make_xp_test_case(stats.quantile) +class TestQuantile: + + def test_input_validation(self, xp): + x = xp.asarray([1, 2, 3]) + p = xp.asarray(0.5) + + message = "`x` must have real dtype." + with pytest.raises(ValueError, match=message): + stats.quantile(xp.asarray([True, False]), p) + with pytest.raises(ValueError): + stats.quantile(xp.asarray([1+1j, 2]), p) + + message = "`p` must have real floating dtype." + with pytest.raises(ValueError, match=message): + stats.quantile(x, xp.asarray([0, 1])) + + message = "`weights` must have real dtype." + with pytest.raises(ValueError, match=message): + stats.quantile(x, p, weights=xp.astype(x, xp.complex64)) + + message = "`axis` must be an integer or None." + with pytest.raises(ValueError, match=message): + stats.quantile(x, p, axis=0.5) + with pytest.raises(ValueError, match=message): + stats.quantile(x, p, axis=(0, -1)) + + message = "`axis` is not compatible with the shapes of the inputs." + with pytest.raises(ValueError, match=message): + stats.quantile(x, p, axis=2) + + if not is_jax(xp): # no data-dependent input validation for lazy arrays + message = "The input contains nan values" + with pytest.raises(ValueError, match=message): + stats.quantile(xp.asarray([xp.nan, 1, 2]), p, nan_policy='raise') + + message = "`method` must be one of..." + with pytest.raises(ValueError, match=message): + stats.quantile(x, p, method='a duck') + + message = "`method='harrell-davis'` does not support `weights`." + with pytest.raises(ValueError, match=message): + stats.quantile(x, p, weights=x, method='harrell-davis') + + message = "`method='round_nearest'` does not support `weights`." + with pytest.raises(ValueError, match=message): + stats.quantile(x, p, weights=x, method='round_nearest') + + message = "If specified, `keepdims` must be True or False." + with pytest.raises(ValueError, match=message): + stats.quantile(x, p, keepdims=42) + + message = "`keepdims` may be False only if the length of `p` along `axis` is 1." + with pytest.raises(ValueError, match=message): + stats.quantile(x, xp.asarray([0.5, 0.6]), keepdims=False) + + + def _get_weights_x_rep(self, x, axis, rng): + x = np.swapaxes(x, axis, -1) + ndim = x.ndim + x = np.atleast_2d(x) + counts = rng.integers(10, size=x.shape[-1], dtype=np.int32) + x_rep = [] + weights = [] + for x_ in x: + counts_ = rng.permuted(counts) + x_rep.append(np.repeat(x_, counts_)) + weights.append(counts_) + x_rep, weights = np.stack(x_rep), np.stack(weights) + if ndim < 2: + x_rep, weights = np.squeeze(x_rep, axis=0), np.squeeze(weights, axis=0) + x_rep, weights = np.swapaxes(x_rep, -1, axis), np.swapaxes(weights, -1, axis) + weights = np.asarray(weights, dtype=x.dtype) + return weights, x_rep + + + @skip_xp_backends(cpu_only=True, reason="PyTorch doesn't have `betainc`.", + exceptions=['cupy']) + @pytest.mark.parametrize('method', + ['inverted_cdf', 'averaged_inverted_cdf', 'closest_observation', + 'hazen', 'interpolated_inverted_cdf', 'linear', + 'median_unbiased', 'normal_unbiased', 'weibull', + 'harrell-davis', 'round_nearest', 'round_outward', 'round_inward', + '_lower', '_higher', '_midpoint', '_nearest']) + @pytest.mark.parametrize('shape_x, shape_p, axis', + [(10, None, -1), (10, 10, -1), (10, (2, 3), -1), ((10, 2), None, 0)]) + @pytest.mark.parametrize('weights', [False, True]) + def test_against_reference(self, method, shape_x, shape_p, axis, weights, xp): + # Test all methods with various data shapes + if weights and (method.startswith('_') or method.startswith('round') + or method=='harrell-davis'): + pytest.skip('`weights` not supported by private (legacy) methods.') + dtype = xp_default_dtype(xp) + rng = np.random.default_rng(23458924568734956) + x = rng.random(size=shape_x) + p = rng.random(size=shape_p) + + if weights: + weights, x_rep = self._get_weights_x_rep(x, axis, rng) + else: + weights, x_rep = None, x + + ref = quantile_reference( + x_rep, p, method=method[1:] if method.startswith('_') else method, + axis=axis, nan_policy='propagate', keepdims=shape_p is not None) + + x, p = xp.asarray(x, dtype=dtype), xp.asarray(p, dtype=dtype) + weights = weights if weights is None else xp.asarray(weights, dtype=dtype) + res = stats.quantile(x, p, method=method, weights=weights, axis=axis) + + xp_assert_close(res, xp.asarray(ref, dtype=dtype)) + + @pytest.mark.filterwarnings("ignore:torch.searchsorted:UserWarning") + @skip_xp_backends(cpu_only=True, reason="PyTorch doesn't have `betainc`.", + exceptions=['cupy', 'jax.numpy']) + @pytest.mark.parametrize('axis', [0, 1]) + @pytest.mark.parametrize('keepdims', [False, True]) + @pytest.mark.parametrize('nan_policy', ['omit', 'propagate', 'marray']) + @pytest.mark.parametrize('dtype', ['float32', 'float64']) + @pytest.mark.parametrize('method', ['linear', 'harrell-davis', 'round_nearest']) + @pytest.mark.parametrize('weights', [False, True]) + def test_against_reference_2(self, axis, keepdims, nan_policy, + dtype, method, weights, xp): + # Test some methods with various combinations of arguments + if is_jax(xp) and nan_policy == 'marray': # mdhaber/marray#146 + pytest.skip("`marray` currently incompatible with JAX") + if weights and method in {'harrell-davis', 'round_nearest'}: + pytest.skip("These methods don't yet support weights") + rng = np.random.default_rng(23458924568734956) + shape = (5, 6) + x = rng.random(size=shape).astype(dtype) + p = rng.random(size=shape).astype(dtype) + mask = rng.random(size=shape) > 0.8 + assert np.any(mask) + x[mask] = np.nan + if not keepdims: + p = np.mean(p, axis=axis, keepdims=True) + + # inject p = 0 and p = 1 to test edge cases + # Currently would fail with CuPy/JAX (cupy/cupy#8934, jax-ml/jax#21900); + # remove the `if` when those are resolved. + if is_numpy(xp): + p0 = p.ravel() + p0[1] = 0. + p0[-2] = 1. + + dtype = getattr(xp, dtype) + + if weights: + weights, x_rep = self._get_weights_x_rep(x, axis, rng) + weights = weights if weights is None else xp.asarray(weights) + else: + weights, x_rep = None, x + + if nan_policy == 'marray': + if not SCIPY_ARRAY_API: + pytest.skip("MArray is only available if SCIPY_ARRAY_API=1") + if weights is not None: + pytest.skip("MArray is not yet compatible with weights") + marray = pytest.importorskip('marray') + kwargs = dict(axis=axis, keepdims=keepdims, method=method) + mxp = marray._get_namespace(xp) + x_mp = mxp.asarray(x, mask=mask) + weights = weights if weights is None else mxp.asarray(weights) + res = stats.quantile(x_mp, mxp.asarray(p), weights=weights, **kwargs) + ref = quantile_reference(x_rep, p, nan_policy='omit', **kwargs) + xp_assert_close(res.data, xp.asarray(ref, dtype=dtype)) + return + + kwargs = dict(axis=axis, keepdims=keepdims, + nan_policy=nan_policy, method=method) + res = stats.quantile(xp.asarray(x), xp.asarray(p), weights=weights, **kwargs) + ref = quantile_reference(x_rep, p, **kwargs) + xp_assert_close(res, xp.asarray(ref, dtype=dtype)) + + def test_integer_input_output_dtype(self, xp): + res = stats.quantile(xp.arange(10, dtype=xp.int64), 0.5) + assert res.dtype == xp_default_dtype(xp) + + @pytest.mark.parametrize('x, p, ref, kwargs', + [([], 0.5, np.nan, {}), + ([1, 2, 3], [-1, 0, 1, 1.5, np.nan], [np.nan, 1, 3, np.nan, np.nan], {}), + ([1, 2, 3], [], [], {}), + ([[np.nan, 2]], 0.5, [np.nan, 2], {'nan_policy': 'omit'}), + ([[], []], 0.5, np.full(2, np.nan), {'axis': -1}), + ([[], []], 0.5, np.zeros((0,)), {'axis': 0, 'keepdims': False}), + ([[], []], 0.5, np.zeros((1, 0)), {'axis': 0, 'keepdims': True}), + ([], [0.5, 0.6], np.full(2, np.nan), {}), + (np.arange(1, 28).reshape((3, 3, 3)), 0.5, [[[14.]]], + {'axis': None, 'keepdims': True}), + ([[1, 2], [3, 4]], [0.25, 0.5, 0.75], [[1.75, 2.5, 3.25]], + {'axis': None, 'keepdims': True}), + # Known issue: + # ([1, 2, 3], 0.5, 2., {'weights': [0, 0, 0]}) + # See https://github.com/scipy/scipy/pull/23941#issuecomment-3503554361 + ]) + def test_edge_cases(self, x, p, ref, kwargs, xp): + default_dtype = xp_default_dtype(xp) + x, p, ref = xp.asarray(x), xp.asarray(p), xp.asarray(ref, dtype=default_dtype) + res = stats.quantile(x, p, **kwargs) + xp_assert_equal(res, ref) + + @pytest.mark.parametrize('axis', [0, 1, 2]) + @pytest.mark.parametrize('keepdims', [False, True]) + def test_size_0(self, axis, keepdims, xp): + shape = [3, 4, 0] + out_shape = shape.copy() + if keepdims: + out_shape[axis] = 1 + else: + out_shape.pop(axis) + res = stats.quantile(xp.zeros(tuple(shape)), 0.5, axis=axis, keepdims=keepdims) + assert res.shape == tuple(out_shape) + + @pytest.mark.parametrize('method', + ['inverted_cdf', 'averaged_inverted_cdf', 'closest_observation', + '_lower', '_higher', '_midpoint', '_nearest']) + def test_transition(self, method, xp): + # test that values of discontinuous estimators are correct when + # p*n + m - 1 is integral. + if method == 'closest_observation' and np.__version__ < '2.0.1': + pytest.skip('Bug in np.quantile (numpy/numpy#26656) fixed in 2.0.1') + x = np.arange(8., dtype=np.float64) + p = np.arange(0, 1.03125, 0.03125) + res = stats.quantile(xp.asarray(x), xp.asarray(p), method=method) + ref = np.quantile(x, p, method=method[1:] if method.startswith('_') else method) + xp_assert_equal(res, xp.asarray(ref, dtype=xp.float64)) + + @pytest.mark.parametrize('zero_weights', [False, True]) + def test_weights_against_numpy(self, zero_weights, xp): + if is_numpy(xp) and xp.__version__ < "2.1.3" and zero_weights: + pytest.skip('`Bug in np.quantile (numpy/numpy#27563) fixed in 2.1.3') + dtype = xp_default_dtype(xp) + rng = np.random.default_rng(85468924398205602) + method = 'inverted_cdf' + x = rng.random(size=100) + weights = rng.random(size=100) + if zero_weights: + weights[weights < 0.5] = 0 + p = np.linspace(0., 1., 300) + res = stats.quantile(xp.asarray(x, dtype=dtype), xp.asarray(p, dtype=dtype), + method=method, weights=xp.asarray(weights, dtype=dtype)) + ref = np.quantile(x, p, method=method, weights=weights) + xp_assert_close(res, xp.asarray(ref, dtype=dtype)) + + @pytest.mark.parametrize('method', + ['inverted_cdf', 'averaged_inverted_cdf', 'closest_observation', 'hazen', + 'interpolated_inverted_cdf', 'linear','median_unbiased', 'normal_unbiased', + 'weibull']) + def test_zero_weights(self, method, xp): + rng = np.random.default_rng(85468924398205602) + + # test 1-D versus eliminating zero-weighted values + n = 100 + x = xp.asarray(rng.random(size=n)) + x0 = xp_copy(x) + p = xp.asarray(rng.random(size=n)) + i_zero = xp.asarray(rng.random(size=n) < 0.1) + weights = xp.asarray(rng.random(size=n)) + weights = xp.where(i_zero, 0., weights) + res = stats.quantile(x, p, weights=weights, method=method) + ref = stats.quantile(x[~i_zero], p, weights=weights[~i_zero], method=method) + xp_assert_close(res, ref) + xp_assert_equal(x, x0) # no input mutation + + # test multi-D versus `nan_policy='omit'` + shape = (5, 100) + x = xp.asarray(rng.random(size=shape)) + x0 = xp_copy(x) + p = xp.asarray(rng.random(size=shape)) + i_zero = xp.asarray(rng.random(size=shape) < 0.1) + weights = xp.asarray(rng.random(size=shape)) + x_nanned = xp.where(i_zero, xp.nan, x) + weights_zeroed = xp.where(i_zero, 0., weights) + res = stats.quantile(x, p, weights=weights_zeroed, method=method, axis=-1) + ref = stats.quantile(x_nanned, p, weights=weights, + nan_policy='omit', method=method, axis=-1) + xp_assert_close(res, ref) + xp_assert_equal(x, x0) # no input mutation + + @pytest.mark.filterwarnings("ignore:torch.searchsorted:UserWarning") + @pytest.mark.parametrize('method', + ['inverted_cdf', 'averaged_inverted_cdf', 'closest_observation', 'hazen', + 'interpolated_inverted_cdf', 'linear','median_unbiased', 'normal_unbiased', + 'weibull']) + @pytest.mark.parametrize('shape', [50, (50, 3)]) + def test_unity_weights(self, method, shape, xp): + # Check that result is unchanged if all weights are `1.0` + rng = np.random.default_rng(28546892439820560) + x = xp.asarray(rng.random(size=shape)) + p = xp.asarray(rng.random(size=shape)) + weights = xp.ones_like(x) + res = stats.quantile(x, p, weights=weights, method=method) + ref = stats.quantile(x, p, method=method) + xp_assert_close(res, ref) + + @skip_xp_backends(cpu_only=True, reason="PyTorch doesn't have `betainc`.", + exceptions=['cupy', 'jax.numpy']) + def test_all_nan_harrell_davis_gh24707(self, xp): + # While working on gh-24707, there was a case in which if *all* elements of one + # slice were NaN, only some elements of another slice were NaN, and + # `nan_policy='omit'`, then the NaNs would not be ignored in the other slice. + # The same test case could fail with `nan_policy='propagate'` for a different + # reason, if the fix were not made carefully. Check that both these cases are + # resolved. + kwargs = dict(method='harrell-davis', axis=-1) + x = xp.asarray([[xp.nan, xp.nan, xp.nan], [xp.nan, 2, 3]]) + + res = stats.quantile(x, 0.5, **kwargs, nan_policy='omit') + xp_assert_close(res, xp.asarray([xp.nan, 2.5])) + + res = stats.quantile(x, 0.5, **kwargs, nan_policy='propagate') + xp_assert_close(res, xp.asarray([xp.nan, xp.nan])) + + +@_apply_over_batch(('a', 1), ('v', 1)) +def np_searchsorted(a, v, side): + return np.searchsorted(a, v, side=side) + + +@make_xp_test_case(_xp_searchsorted) +class Test_XPSearchsorted: + @pytest.mark.parametrize('side', ['left', 'right']) + @pytest.mark.parametrize('ties', [False, True]) + @pytest.mark.parametrize('shape', [0, 1, 2, 10, 11, 1000, 10001, + (2, 0), (0, 2), (2, 10), (2, 3, 11)]) + @pytest.mark.parametrize('nans_x', [False, True]) + @pytest.mark.parametrize('infs_x', [False, True]) + def test_nd(self, side, ties, shape, nans_x, infs_x, xp): + if nans_x and is_torch(xp): + pytest.skip('torch sorts NaNs differently') + rng = np.random.default_rng(945298725498274853) + if ties: + x = rng.integers(5, size=shape) + else: + x = rng.random(shape) + # float32 is to accommodate JAX - nextafter with `float64` is too small? + x = np.asarray(x, dtype=np.float32) + xr = np.nextafter(x, np.inf) + xl = np.nextafter(x, -np.inf) + x_ = np.asarray([-np.inf, np.inf, np.nan]) + x_ = np.broadcast_to(x_, x.shape[:-1] + (3,)) + y = rng.permuted(np.concatenate((xl, x, xr, x_), axis=-1), axis=-1) + if nans_x: + mask = rng.random(shape) < 0.1 + x[mask] = np.nan + if infs_x: + mask = rng.random(shape) < 0.1 + x[mask] = -np.inf + mask = rng.random(shape) > 0.9 + x[mask] = np.inf + x = np.sort(x, axis=-1) + x, y = np.asarray(x, dtype=np.float64), np.asarray(y, dtype=np.float64) + xp_default_int = xp.asarray(1).dtype + if xp_size(x) == 0 and x.ndim > 0 and x.shape[-1] != 0: + ref = xp.empty(x.shape[:-1] + (y.shape[-1],), dtype=xp_default_int) + else: + ref = xp.asarray(np_searchsorted(x, y, side=side), dtype=xp_default_int) + x, y = xp.asarray(x), xp.asarray(y) + res = _xp_searchsorted(x, y, side=side) + xp_assert_equal(res, ref) + + +@_apply_over_batch(('x', 1), ('y', 1)) +def estimated_cdf_reference_last_axis(x, y, nan_policy, method): + i_nan = np.isnan(x) + if nan_policy == 'propagate' and np.any(i_nan): + return np.full_like(y, np.nan) + elif nan_policy == 'omit': + x = x[~i_nan] + return stats.estimated_cdf(x, y, keepdims=True, method=method) + + +def estimated_cdf_reference(x, y, *, axis=0, nan_policy='propagate', + keepdims=None, method='linear'): + x, y = _broadcast_arrays((x, y), axis=axis) + x, y = np.moveaxis(x, axis, -1), np.moveaxis(y, axis, -1) + res = estimated_cdf_reference_last_axis(x, y, nan_policy, method) + res = np.moveaxis(res, -1, axis) + if not keepdims: + res = np.squeeze(res, axis=axis) + return res + + +# avoid variable collection order issues +_estimated_cdf_methods_list = sorted(list(_estimated_cdf_methods)) + + +@make_xp_test_case(stats.estimated_cdf) +class TestEstimatedCDF: + def test_input_validation(self, xp): + x = xp.asarray([1, 2, 3]) + y = xp.asarray(2) + + message = "`x` must have real dtype." + with pytest.raises(ValueError, match=message): + stats.estimated_cdf(xp.asarray([True, False]), y) + with pytest.raises(ValueError): + stats.estimated_cdf(xp.asarray([1+1j, 2]), y) + + message = "`y` must have real dtype." + with pytest.raises(ValueError, match=message): + stats.estimated_cdf(x, xp.asarray([0+1j, 1])) + + message = "`axis` must be an integer or None." + with pytest.raises(ValueError, match=message): + stats.estimated_cdf(x, y, axis=0.5) + with pytest.raises(ValueError, match=message): + stats.estimated_cdf(x, y, axis=(0, -1)) + + message = "`axis` is not compatible with the shapes of the inputs." + with pytest.raises(ValueError, match=message): + stats.estimated_cdf(x, y, axis=2) + + if not is_jax(xp): # no data-dependent input validation for lazy arrays + message = "The input contains nan values" + with pytest.raises(ValueError, match=message): + stats.estimated_cdf(xp.asarray([xp.nan, 1, 2]), y, nan_policy='raise') + + message = "method` must be one of..." + with pytest.raises(ValueError, match=message): + stats.estimated_cdf(x, y, method='a duck') + + message = "If specified, `keepdims` must be True or False." + with pytest.raises(ValueError, match=message): + stats.estimated_cdf(x, y, keepdims=42) + + message = "`keepdims` may be False only if the length of `y` along `axis` is 1." + with pytest.raises(ValueError, match=message): + stats.estimated_cdf(x, xp.asarray([0.5, 0.6]), keepdims=False) + + @pytest.mark.parametrize('method', _estimated_cdf_methods_list) + @pytest.mark.parametrize('x_shape', [2, 10, 11, 100, 1001, (2, 10), (2, 3, 11)]) + @pytest.mark.parametrize('y_shape', [None, 25]) + @pytest.mark.parametrize('ties', [False, True]) + def test_against_quantile(self, method, x_shape, y_shape, ties, xp): + discontinuous = method in _estimated_cdf_discontinuous_methods + dtype = xp_default_dtype(xp) # removed parameterization to speed up tests + rng = np.random.default_rng(394529872549827485) + y_shape = x_shape if y_shape is None else y_shape + + if ties: + x = xp.asarray(rng.integers(9, size=x_shape), dtype=dtype) + else: + x = xp.asarray(rng.standard_normal(size=x_shape), dtype=dtype) + + p = xp.asarray(rng.random(size=y_shape), dtype=dtype) + y = stats.quantile(x, p, method=method, axis=-1) + res = stats.estimated_cdf(x, y, method=method, axis=-1) + ref = xp.broadcast_to(p, (*x.shape[:-1], y.shape[-1])) + + # check that `quantile` is the inverse of `estimated_cdf` + # note that for discontinuous methods, res is right on the cusp of a transition, + # and there can be a tiny bit of error to the right or left. We shift it left + # to ensure we're on the correct side of the transition, producing the same `y2` + # as if the probability calculation were exact. + res = res - 1e-6 if discontinuous else res + y2 = stats.quantile(x, res, method=method, axis=-1) + atol = 1e-5 if dtype == xp.float32 else 1e-12 + xp_assert_close(y2, y, atol=atol) + + # if there are ties or method is discontinuous, `quantile` is not invertible + if ties or discontinuous: + return + + # `quantile` is not invertible outside this domain + a, b = _estimated_cdf_continuous_methods[method] + n = x.shape[-1] + p_min = (1 - a) / (n + 1 - a - b) + p_max = (n - a) / (n + 1 - a - b) + i_very_low = y < xp.min(x, axis=-1, keepdims=True) + i_very_high = y > xp.max(x, axis=-1, keepdims=True) + i_low = (ref <= p_min) & ~i_very_low + i_high = (ref >= p_max) & ~i_very_high + i_ok = ~(i_low | i_high | i_very_low | i_very_high) + + # check for correct inversion within the domain + xp_assert_close(res[i_ok], ref[i_ok]) + + # check that all other values get mapped to bottom or top of range + kwargs = dict(check_shape=False, check_dtype=False, check_0d=True) + xp_assert_close(res[i_low], xp.asarray(p_min), **kwargs) + xp_assert_close(res[i_high], xp.asarray(p_max), **kwargs) + xp_assert_close(res[i_very_low], xp.asarray(0.0), **kwargs) + xp_assert_close(res[i_very_high], xp.asarray(1.0), **kwargs) + + @pytest.mark.filterwarnings("ignore:torch.searchsorted:UserWarning") + @pytest.mark.parametrize('axis', [0, 1]) + @pytest.mark.parametrize('keepdims', [False, True]) + @pytest.mark.parametrize('nan_policy', ['propagate', 'omit', 'marray']) + @pytest.mark.parametrize('dtype', ['float32', 'float64']) + @pytest.mark.parametrize('nans', [False, True]) + @pytest.mark.parametrize('meth', ['linear', 'inverted_cdf']) + def test_against_reference(self, axis, keepdims, nan_policy, dtype, nans, meth, xp): + if is_jax(xp) and nan_policy == 'marray': # mdhaber/marray#146 + pytest.skip("`marray` currently incompatible with JAX") + rng = np.random.default_rng(23458924568734956) + shape = (5, 6) + x = rng.standard_normal(size=shape).astype(dtype) + y = rng.standard_normal(size=shape).astype(dtype) + + mask = None + if nans: + mask = rng.random(size=shape) > 0.8 + assert np.any(mask) + x[mask] = np.nan + + if not keepdims: + y = np.mean(y, axis=axis, keepdims=True) + + dtype = getattr(xp, dtype) + + if nan_policy == 'marray': + if not SCIPY_ARRAY_API: + pytest.skip("MArray is only available if SCIPY_ARRAY_API=1") + marray = pytest.importorskip('marray') + kwargs = dict(axis=axis, keepdims=keepdims, method=meth) + mxp = marray._get_namespace(xp) + x_mp = mxp.asarray(x, mask=mask) + res = stats.estimated_cdf(x_mp, mxp.asarray(y), **kwargs) + ref = estimated_cdf_reference(x, y, nan_policy='omit', **kwargs) + xp_assert_close(res.data, xp.asarray(ref, dtype=dtype)) + return + + kwargs = dict(axis=axis, keepdims=keepdims, + nan_policy=nan_policy, method=meth) + res = stats.estimated_cdf(xp.asarray(x), xp.asarray(y), **kwargs) + ref = estimated_cdf_reference(x, y, **kwargs) + xp_assert_close(res, xp.asarray(ref, dtype=dtype)) + + @pytest.mark.skip_xp_backends('torch', reason='issues with sorting NaNs') + @pytest.mark.parametrize('n', [50, 500]) + @pytest.mark.parametrize('method, ab', _estimated_cdf_continuous_methods.items()) + def test_plotting_positions(self, n, method, ab, xp): + a, b = ab + rng = np.random.default_rng(539452987254982748) + x = rng.standard_normal(n) + + mask = rng.random(n) < 0.1 + x[mask] = np.nan + mask = xp.asarray(mask) + + x_masked = np.ma.masked_invalid(x) + ref = stats.mstats.plotting_positions(x_masked, a, b) + ref = xp.asarray(ref.data) + + x = xp.asarray(x) + res = stats.estimated_cdf(x, x, nan_policy='omit', method=method) + + xp_assert_close(res[~mask], ref[~mask]) + assert xp.all(xp.isnan(res[mask])) + + @pytest.mark.parametrize('ties', [False, True]) + def test_against_ecdf_percentileofscore(self, ties, xp): + rng = np.random.default_rng(853945298725498274) + n = 50 + dtype = xp_default_dtype(xp) + x = rng.integers(10, size=n) if ties else rng.standard_normal(size=n) + y = rng.integers(10, size=25) if ties else rng.standard_normal(size=25) + ref = stats.ecdf(x).cdf.evaluate(y) + ref2 = stats.percentileofscore(x, y, 'weak') + x, y = xp.asarray(x, dtype=dtype), xp.asarray(y, dtype=dtype) + res = stats.estimated_cdf(x, y, method='inverted_cdf') + ref, ref2 = xp.asarray(ref, dtype=dtype), xp.asarray(ref2, dtype=dtype) + xp_assert_close(res, ref) + xp_assert_close(res, ref2 / 100) + + def test_integer_input_output_dtype(self, xp): + x = xp.arange(10, dtype=xp.int64) + res = stats.estimated_cdf(x, x) + assert res.dtype == xp_default_dtype(xp) + + @pytest.mark.parametrize('nan_policy', ['propagate', 'omit', 'marray']) + @pytest.mark.parametrize('method', _estimated_cdf_methods_list) + def test_size_one_sample(self, nan_policy, method, xp): + discontinuous = method in _estimated_cdf_discontinuous_methods + x = xp.arange(10.) + y = xp.asarray([0., -1., 1.]) + n = np.asarray(1.) if is_array_api_strict(xp) else xp.asarray(1.) + with np.errstate(divide='ignore', invalid='ignore'): # for method = 'linear' + if discontinuous: + ref = xp.asarray([1., 0., 1.]) + else: + a, b = _estimated_cdf_continuous_methods[method] + ref = xp.asarray([float((n - a) / (n + 1 - a - b)), 0., 1.]) + + if nan_policy == 'propagate': + x = x[:1] + kwargs = {'nan_policy': 'propagate'} + elif nan_policy == 'omit': + x = xpx.at(x)[1:].set(xp.nan) + kwargs = {'nan_policy': 'omit'} + elif nan_policy == 'marray': + if is_jax(xp): + pytest.skip("JAX currently incompatible with `marray`") + if not SCIPY_ARRAY_API: + pytest.skip("MArray is only available if SCIPY_ARRAY_API=1") + marray = pytest.importorskip('marray') + mxp = marray._get_namespace(xp) + mask = (x > 0.) + x = mxp.asarray(x, mask=mask) + y = mxp.asarray(y) + kwargs = {} + + with np.errstate(divide='ignore', invalid='ignore'): # for method = 'linear' + res = stats.estimated_cdf(x, y, method=method, **kwargs) + res = res.data if nan_policy == 'marray' else res + xp_assert_close(res, ref) + + # skipping marray due to mdhaber/marray#24 + @pytest.mark.parametrize('nan_policy', ['propagate', 'omit']) + @pytest.mark.parametrize('method', _estimated_cdf_methods_list) + def test_size_zero_sample(self, nan_policy, method, xp): + x = xp.arange(10.) + y = xp.asarray([0., -1., 1.]) # this should work + ref = xp.full_like(y, xp.nan) + + if nan_policy == 'propagate': + x = x[0:0] + kwargs = {'nan_policy': 'propagate'} + elif nan_policy == 'omit': + x = xpx.at(x)[:].set(xp.nan) + kwargs = {'nan_policy': 'omit'} + elif nan_policy == 'marray': + if not SCIPY_ARRAY_API: + pytest.skip("MArray is only available if SCIPY_ARRAY_API=1") + if is_jax(xp): + pytest.skip("JAX currently incompatible with `marray`") + marray = pytest.importorskip('marray') + mxp = marray._get_namespace(xp) + mask = (x >= 0.) + x = mxp.asarray(x, mask=mask) + y = mxp.asarray(y) + kwargs = {} + + with np.errstate(divide='ignore', invalid='ignore'): # for method = 'linear' + res = stats.estimated_cdf(x, y, method=method, **kwargs) + + if nan_policy == 'marray': + assert xp.all(res.mask) + else: + xp_assert_close(res, ref) + + @pytest.mark.parametrize('x, y, ref, kwargs', + [ + ([], 0.5, np.nan, {}), + ([1, 2, 3], [0.999, 3.001, np.nan], [0., 1., np.nan], {}), + ([1, 2, 3], [], [], {}), + ([[np.nan, 2]], 2, [np.nan, 0.5], {'nan_policy': 'omit', 'method': 'weibull'}), + ([[], []], 0.5, np.full(2, np.nan), {'axis': -1}), + ([[], []], 0.5, np.zeros((0,)), {'axis': 0, 'keepdims': False}), + ([[], []], 0.5, np.zeros((1, 0)), {'axis': 0, 'keepdims': True}), + ([], [0.5, 0.6], np.full(2, np.nan), {}), + (np.arange(1, 28).reshape((3, 3, 3)), 14., [[[0.5]]], + {'axis': None, 'keepdims': True}), + ([[1, 2], [3, 4]], [1.75, 2.5, 3.25], [[0.25, 0.5, 0.75]], + {'axis': None, 'keepdims': True}), + ([1, 2, 3], [-np.inf, np.inf], [0.0, 1.0], {}), + # It is our choice how much effort and computational overhead we want to put + # into adjusting for insane edge cases like when `x` contains infinite values, + # especially when `y` does, too. + # One practical argument would be that y = +/- inf should produce the same + # results as an extremely large finite number, in which case the 0th element + # of the 'linear' result should be `0`. + # Another argument would be for producing NaN below wherever y = +/- inf. + # Another would be that it's not appropriate to spend significant computation + # correcting these edge cases; we should just document what we do. + # In any case, this is the current status. + ([-np.inf, -1, 0, 1, np.inf], [-np.inf, -2, -1, 0, 1, 2, np.inf], + [0.25, 0.25, 0.25, 0.5 , 0.75, 0.75, np.nan], {'method':'linear'}), + ([-np.inf, -1, 0, 1, np.inf], [-np.inf, -2, -1, 0, 1, 2, np.inf], + [0.2, 0.2, 0.4, 0.6, 0.8, 0.8, 1.], {'method': 'inverted_cdf'}), + ]) + def test_edge_cases(self, x, y, ref, kwargs, xp): + if kwargs.get('method', None) == 'weibull' and is_torch(xp): + pytest.skip('data-apis/array-api-compat#360') + if kwargs.get('axis', None) == -1 and is_cupy(xp): + pytest.skip('Fails; need to investigate.') + default_dtype = xp_default_dtype(xp) + x, y, ref = xp.asarray(x), xp.asarray(y), xp.asarray(ref, dtype=default_dtype) + res = stats.estimated_cdf(x, y, **kwargs) + xp_assert_close(res, ref, rtol=1e-15) + + @pytest.mark.skip_xp_backends('jax.numpy', reason="-1e-45 is not less than 0?") + @pytest.mark.parametrize('method', _estimated_cdf_discontinuous_methods.keys()) + def test_transition(self, method, xp): + # test that values of discontinuous estimators are as expected around + # transition point + x = np.arange(8., dtype=np.float64) + xl, xr = np.nextafter(x, -np.inf), np.nextafter(x, np.inf) + offset = 0.5 if method == 'closest_observation' else 0.0 + ref_r = (x + 1 + offset) / 8 + ref_r[-1] = 1.0 # value is greater than or equal to the maximum observation + ref_l = (x + offset) / 8 + ref_l[0] = 0.0 # value is less than the minimum observation + x, xl, xr = xp.asarray(x), xp.asarray(xl), xp.asarray(xr) + ref_l, ref_r = xp.asarray(ref_l), xp.asarray(ref_r) + xp_assert_equal(stats.estimated_cdf(x, x, method=method), ref_r) + xp_assert_equal(stats.estimated_cdf(x, xr, method=method), ref_r) + xp_assert_equal(stats.estimated_cdf(x, xl, method=method), ref_l) diff --git a/v12x/pylib/scipy/stats/tests/test_rank.py b/v12x/pylib/scipy/stats/tests/test_rank.py new file mode 100644 index 0000000000000000000000000000000000000000..1d8e89cfea0315e30d1ab0c45e8bdab282360761 --- /dev/null +++ b/v12x/pylib/scipy/stats/tests/test_rank.py @@ -0,0 +1,350 @@ +import numpy as np +from numpy.testing import assert_equal, assert_array_equal +import pytest + +from scipy import stats +from scipy.conftest import skip_xp_invalid_arg +from scipy.stats import rankdata, tiecorrect +from scipy._lib._array_api import (xp_assert_equal, make_xp_test_case, xp_result_type, + xp_default_dtype) + +skip_xp_backends = pytest.mark.skip_xp_backends + +class TestTieCorrect: + + def test_empty(self): + """An empty array requires no correction, should return 1.0.""" + ranks = np.array([], dtype=np.float64) + c = tiecorrect(ranks) + assert_equal(c, 1.0) + + def test_one(self): + """A single element requires no correction, should return 1.0.""" + ranks = np.array([1.0], dtype=np.float64) + c = tiecorrect(ranks) + assert_equal(c, 1.0) + + def test_no_correction(self): + """Arrays with no ties require no correction.""" + ranks = np.arange(2.0) + c = tiecorrect(ranks) + assert_equal(c, 1.0) + ranks = np.arange(3.0) + c = tiecorrect(ranks) + assert_equal(c, 1.0) + + def test_basic(self): + """Check a few basic examples of the tie correction factor.""" + # One tie of two elements + ranks = np.array([1.0, 2.5, 2.5]) + c = tiecorrect(ranks) + T = 2.0 + N = ranks.size + expected = 1.0 - (T**3 - T) / (N**3 - N) + assert_equal(c, expected) + + # One tie of two elements (same as above, but tie is not at the end) + ranks = np.array([1.5, 1.5, 3.0]) + c = tiecorrect(ranks) + T = 2.0 + N = ranks.size + expected = 1.0 - (T**3 - T) / (N**3 - N) + assert_equal(c, expected) + + # One tie of three elements + ranks = np.array([1.0, 3.0, 3.0, 3.0]) + c = tiecorrect(ranks) + T = 3.0 + N = ranks.size + expected = 1.0 - (T**3 - T) / (N**3 - N) + assert_equal(c, expected) + + # Two ties, lengths 2 and 3. + ranks = np.array([1.5, 1.5, 4.0, 4.0, 4.0]) + c = tiecorrect(ranks) + T1 = 2.0 + T2 = 3.0 + N = ranks.size + expected = 1.0 - ((T1**3 - T1) + (T2**3 - T2)) / (N**3 - N) + assert_equal(c, expected) + + def test_overflow(self): + ntie, k = 2000, 5 + a = np.repeat(np.arange(k), ntie) + n = a.size # ntie * k + out = tiecorrect(rankdata(a)) + assert_equal(out, 1.0 - k * (ntie**3 - ntie) / float(n**3 - n)) + + +@make_xp_test_case(stats.rankdata) +class TestRankData: + + def desired_dtype(self, x, *, xp): + return xp_result_type(x, force_floating=True, xp=xp) + + def test_empty(self, xp): + """stats.rankdata of empty array should return an empty array.""" + a = xp.asarray([], dtype=xp.int64) + r = rankdata(a) + xp_assert_equal(r, xp.asarray([], dtype=self.desired_dtype(a, xp=xp))) + + def test_list(self): + # test that NumPy still accepts lists + r = rankdata([]) + assert_array_equal(r, np.array([])) + + r = rankdata([40, 10, 30, 10, 50]) + assert_equal(r, [4.0, 1.5, 3.0, 1.5, 5.0]) + + @pytest.mark.parametrize("shape", [(0, 1, 2)]) + @pytest.mark.parametrize("axis", [None, *range(3)]) + def test_empty_multidim(self, shape, axis, xp): + a = xp.empty(shape, dtype=xp.int64) + r = rankdata(a, axis=axis) + expected_shape = (0,) if axis is None else shape + xp_assert_equal(r, xp.empty(expected_shape, dtype=self.desired_dtype(a, xp=xp))) + + def test_one(self, xp): + """Check stats.rankdata with an array of length 1.""" + data = [100] + a = xp.asarray(data, dtype=xp.int64) + r = rankdata(a) + xp_assert_equal(r, xp.asarray([1.0], dtype=self.desired_dtype(a, xp=xp))) + + def test_basic(self, xp): + """Basic tests of stats.rankdata.""" + desired_dtype = xp_default_dtype(xp) + + data = [100, 10, 50] + expected = xp.asarray([3.0, 1.0, 2.0], dtype=desired_dtype) + a = xp.asarray(data, dtype=xp.int64) + r = rankdata(a) + xp_assert_equal(r, expected) + + data = [40, 10, 30, 10, 50] + expected = xp.asarray([4.0, 1.5, 3.0, 1.5, 5.0], dtype=desired_dtype) + a = xp.asarray(data, dtype=xp.int64) + r = rankdata(a) + xp_assert_equal(r, expected) + + data = [20, 20, 20, 10, 10, 10] + expected = xp.asarray([5.0, 5.0, 5.0, 2.0, 2.0, 2.0], dtype=desired_dtype) + a = xp.asarray(data, dtype=xp.int64) + r = rankdata(a) + xp_assert_equal(r, expected) + + # # The docstring states explicitly that the argument is flattened. + a2d = xp.reshape(a, (2, 3)) + r = rankdata(a2d) + xp_assert_equal(r, expected) + + @skip_xp_invalid_arg + def test_rankdata_object_string(self): + + def min_rank(a): + return [1 + sum(i < j for i in a) for j in a] + + def max_rank(a): + return [sum(i <= j for i in a) for j in a] + + def ordinal_rank(a): + return min_rank([(x, i) for i, x in enumerate(a)]) + + def average_rank(a): + return [(i + j) / 2.0 for i, j in zip(min_rank(a), max_rank(a))] + + def dense_rank(a): + b = np.unique(a) + return [1 + sum(i < j for i in b) for j in a] + + rankf = dict(min=min_rank, max=max_rank, ordinal=ordinal_rank, + average=average_rank, dense=dense_rank) + + def check_ranks(a): + for method in 'min', 'max', 'dense', 'ordinal', 'average': + out = rankdata(a, method=method) + assert_array_equal(out, rankf[method](a)) + + val = ['foo', 'bar', 'qux', 'xyz', 'abc', 'efg', 'ace', 'qwe', 'qaz'] + check_ranks(np.random.choice(val, 200)) + check_ranks(np.random.choice(val, 200).astype('object')) + + val = np.array([0, 1, 2, 2.718, 3, 3.141], dtype='object') + check_ranks(np.random.choice(val, 200).astype('object')) + + @pytest.mark.skip_xp_backends("torch", reason="`take_along_axis` fails with uint64") + def test_large_uint(self, xp): + data = xp.asarray([2**60, 2**60+1], dtype=xp.uint64) + r = rankdata(data) + desired_dtype = self.desired_dtype(data, xp=xp) + xp_assert_equal(r, xp.asarray([1.0, 2.0], dtype=desired_dtype)) + + def test_large_int(self, xp): + data = xp.asarray([2**60, 2**60+1], dtype=xp.int64) + r = rankdata(data) + desired_dtype = self.desired_dtype(data, xp=xp) + xp_assert_equal(r, xp.asarray([1.0, 2.0], dtype=desired_dtype)) + + data = xp.asarray([2**60, -2**60+1], dtype=xp.int64) + r = rankdata(data) + xp_assert_equal(r, xp.asarray([2.0, 1.0], dtype=desired_dtype)) + + @pytest.mark.parametrize('n', [10000, 100000, 1000000]) + def test_big_tie(self, n, xp): + data = xp.ones(n) + r = rankdata(data) + expected_rank = 0.5 * (n + 1) + ref = xp.asarray(expected_rank * data, dtype=self.desired_dtype(data, xp=xp)) + xp_assert_equal(r, ref) + + def test_axis(self, xp): + data = xp.asarray([[0, 2, 1], [4, 2, 2]]) + + expected0 = xp.asarray([[1., 1.5, 1.], [2., 1.5, 2.]]) + r0 = rankdata(data, axis=0) + xp_assert_equal(r0, expected0) + + expected1 = xp.asarray([[1., 3., 2.], [3., 1.5, 1.5]]) + r1 = rankdata(data, axis=1) + xp_assert_equal(r1, expected1) + + methods= ["average", "min", "max", "dense", "ordinal"] + + @pytest.mark.parametrize("axis", [0, 1]) + @pytest.mark.parametrize("method", methods) + def test_size_0_axis(self, axis, method, xp): + shape = (3, 0) + data = xp.zeros(shape) + desired_dtype = self.desired_dtype(data, xp=xp) + r = rankdata(data, method=method, axis=axis) + assert_equal(r.shape, shape) + assert_equal(r.dtype, desired_dtype) + xp_assert_equal(r, xp.empty(shape, dtype=desired_dtype)) + + @pytest.mark.parametrize('axis', range(3)) + @pytest.mark.parametrize('method', methods) + def test_nan_policy_omit_3d(self, axis, method): + shape = (20, 21, 22) + rng = np.random.RandomState(23983242) + + a = rng.random(size=shape) + i = rng.random(size=shape) < 0.4 + j = rng.random(size=shape) < 0.1 + k = rng.random(size=shape) < 0.1 + a[i] = np.nan + a[j] = -np.inf + a[k] - np.inf + + def rank_1d_omit(a, method): + out = np.zeros_like(a) + i = np.isnan(a) + a_compressed = a[~i] + res = rankdata(a_compressed, method) + out[~i] = res + out[i] = np.nan + return out + + def rank_omit(a, method, axis): + return np.apply_along_axis(lambda a: rank_1d_omit(a, method), + axis, a) + + res = rankdata(a, method, axis=axis, nan_policy='omit') + res0 = rank_omit(a, method, axis=axis) + + assert_array_equal(res, res0) + + def test_nan_policy_2d_axis_none(self): + # 2 2d-array test with axis=None + data = [[0, np.nan, 3], + [4, 2, np.nan], + [1, 2, 2]] + assert_array_equal(rankdata(data, axis=None, nan_policy='omit'), + [1., np.nan, 6., 7., 4., np.nan, 2., 4., 4.]) + assert_array_equal(rankdata(data, axis=None, nan_policy='propagate'), + [np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, + np.nan, np.nan, np.nan]) + + def test_nan_policy_raise(self): + # 1 1d-array test + data = [0, 2, 3, -2, np.nan, np.nan] + with pytest.raises(ValueError, match="The input contains nan"): + rankdata(data, nan_policy='raise') + + # 2 2d-array test + data = [[0, np.nan, 3], + [4, 2, np.nan], + [np.nan, 2, 2]] + + with pytest.raises(ValueError, match="The input contains nan"): + rankdata(data, axis=0, nan_policy="raise") + + with pytest.raises(ValueError, match="The input contains nan"): + rankdata(data, axis=1, nan_policy="raise") + + def test_nan_policy_propagate(self): + # 1 1d-array test + data = [0, 2, 3, -2, np.nan, np.nan] + assert_array_equal(rankdata(data, nan_policy='propagate'), + [np.nan, np.nan, np.nan, np.nan, np.nan, np.nan]) + + # 2 2d-array test + data = [[0, np.nan, 3], + [4, 2, np.nan], + [1, 2, 2]] + assert_array_equal(rankdata(data, axis=0, nan_policy='propagate'), + [[1, np.nan, np.nan], + [3, np.nan, np.nan], + [2, np.nan, np.nan]]) + assert_array_equal(rankdata(data, axis=1, nan_policy='propagate'), + [[np.nan, np.nan, np.nan], + [np.nan, np.nan, np.nan], + [1, 2.5, 2.5]]) + + _rankdata_cases = ( + # values, method, expected + ([], 'average', []), + ([], 'min', []), + ([], 'max', []), + ([], 'dense', []), + ([], 'ordinal', []), + # + ([100], 'average', [1.0]), + ([100], 'min', [1.0]), + ([100], 'max', [1.0]), + ([100], 'dense', [1.0]), + ([100], 'ordinal', [1.0]), + # + ([100, 100, 100], 'average', [2.0, 2.0, 2.0]), + ([100, 100, 100], 'min', [1.0, 1.0, 1.0]), + ([100, 100, 100], 'max', [3.0, 3.0, 3.0]), + ([100, 100, 100], 'dense', [1.0, 1.0, 1.0]), + ([100, 100, 100], 'ordinal', [1.0, 2.0, 3.0]), + # + ([100, 300, 200], 'average', [1.0, 3.0, 2.0]), + ([100, 300, 200], 'min', [1.0, 3.0, 2.0]), + ([100, 300, 200], 'max', [1.0, 3.0, 2.0]), + ([100, 300, 200], 'dense', [1.0, 3.0, 2.0]), + ([100, 300, 200], 'ordinal', [1.0, 3.0, 2.0]), + # + ([100, 200, 300, 200], 'average', [1.0, 2.5, 4.0, 2.5]), + ([100, 200, 300, 200], 'min', [1.0, 2.0, 4.0, 2.0]), + ([100, 200, 300, 200], 'max', [1.0, 3.0, 4.0, 3.0]), + ([100, 200, 300, 200], 'dense', [1.0, 2.0, 3.0, 2.0]), + ([100, 200, 300, 200], 'ordinal', [1.0, 2.0, 4.0, 3.0]), + # + ([100, 200, 300, 200, 100], 'average', [1.5, 3.5, 5.0, 3.5, 1.5]), + ([100, 200, 300, 200, 100], 'min', [1.0, 3.0, 5.0, 3.0, 1.0]), + ([100, 200, 300, 200, 100], 'max', [2.0, 4.0, 5.0, 4.0, 2.0]), + ([100, 200, 300, 200, 100], 'dense', [1.0, 2.0, 3.0, 2.0, 1.0]), + ([100, 200, 300, 200, 100], 'ordinal', [1.0, 3.0, 5.0, 4.0, 2.0]), + # + ([10] * 30, 'ordinal', np.arange(1.0, 31.0)), + ) + + @pytest.mark.parametrize('case', _rankdata_cases) + def test_cases(self, case, xp): + values, method, expected = case + values = xp.asarray(values) + r = rankdata(xp.asarray(values), method=method) + ref = xp.asarray(expected, dtype=self.desired_dtype(values, xp=xp)) + xp_assert_equal(r, ref) diff --git a/v12x/pylib/scipy/stats/tests/test_relative_risk.py b/v12x/pylib/scipy/stats/tests/test_relative_risk.py new file mode 100644 index 0000000000000000000000000000000000000000..b75e64d929f319465b1f1d62af4fb2096c2ab2ac --- /dev/null +++ b/v12x/pylib/scipy/stats/tests/test_relative_risk.py @@ -0,0 +1,95 @@ +import pytest +import numpy as np +from numpy.testing import assert_allclose, assert_equal +from scipy.stats.contingency import relative_risk + + +# Test just the calculation of the relative risk, including edge +# cases that result in a relative risk of 0, inf or nan. +@pytest.mark.parametrize( + 'exposed_cases, exposed_total, control_cases, control_total, expected_rr', + [(1, 4, 3, 8, 0.25 / 0.375), + (0, 10, 5, 20, 0), + (0, 10, 0, 20, np.nan), + (5, 15, 0, 20, np.inf)] +) +def test_relative_risk(exposed_cases, exposed_total, + control_cases, control_total, expected_rr): + result = relative_risk(exposed_cases, exposed_total, + control_cases, control_total) + assert_allclose(result.relative_risk, expected_rr, rtol=1e-13) + + +def test_relative_risk_confidence_interval(): + result = relative_risk(exposed_cases=16, exposed_total=128, + control_cases=24, control_total=256) + rr = result.relative_risk + ci = result.confidence_interval(confidence_level=0.95) + # The corresponding calculation in R using the epitools package. + # + # > library(epitools) + # > c <- matrix(c(232, 112, 24, 16), nrow=2) + # > result <- riskratio(c) + # > result$measure + # risk ratio with 95% C.I. + # Predictor estimate lower upper + # Exposed1 1.000000 NA NA + # Exposed2 1.333333 0.7347317 2.419628 + # + # The last line is the result that we want. + assert_allclose(rr, 4/3) + assert_allclose((ci.low, ci.high), (0.7347317, 2.419628), rtol=5e-7) + + +def test_relative_risk_ci_conflevel0(): + result = relative_risk(exposed_cases=4, exposed_total=12, + control_cases=5, control_total=30) + rr = result.relative_risk + assert_allclose(rr, 2.0, rtol=1e-14) + ci = result.confidence_interval(0) + assert_allclose((ci.low, ci.high), (2.0, 2.0), rtol=1e-12) + + +def test_relative_risk_ci_conflevel1(): + result = relative_risk(exposed_cases=4, exposed_total=12, + control_cases=5, control_total=30) + ci = result.confidence_interval(1) + assert_equal((ci.low, ci.high), (0, np.inf)) + + +def test_relative_risk_ci_edge_cases_00(): + result = relative_risk(exposed_cases=0, exposed_total=12, + control_cases=0, control_total=30) + assert_equal(result.relative_risk, np.nan) + ci = result.confidence_interval() + assert_equal((ci.low, ci.high), (np.nan, np.nan)) + + +def test_relative_risk_ci_edge_cases_01(): + result = relative_risk(exposed_cases=0, exposed_total=12, + control_cases=1, control_total=30) + assert_equal(result.relative_risk, 0) + ci = result.confidence_interval() + assert_equal((ci.low, ci.high), (0.0, np.nan)) + + +def test_relative_risk_ci_edge_cases_10(): + result = relative_risk(exposed_cases=1, exposed_total=12, + control_cases=0, control_total=30) + assert_equal(result.relative_risk, np.inf) + ci = result.confidence_interval() + assert_equal((ci.low, ci.high), (np.nan, np.inf)) + + +@pytest.mark.parametrize('ec, et, cc, ct', [(0, 0, 10, 20), + (-1, 10, 1, 5), + (1, 10, 0, 0), + (1, 10, -1, 4)]) +def test_relative_risk_bad_value(ec, et, cc, ct): + with pytest.raises(ValueError, match="must be an integer not less than"): + relative_risk(ec, et, cc, ct) + + +def test_relative_risk_bad_type(): + with pytest.raises(TypeError, match="must be an integer"): + relative_risk(1, 10, 2.0, 40) diff --git a/v12x/pylib/scipy/stats/tests/test_resampling.py b/v12x/pylib/scipy/stats/tests/test_resampling.py new file mode 100644 index 0000000000000000000000000000000000000000..ec900da2d5a05570054596286f82c67f96fbdc0d --- /dev/null +++ b/v12x/pylib/scipy/stats/tests/test_resampling.py @@ -0,0 +1,2084 @@ +import warnings + +import pytest + +import numpy as np +from numpy.testing import assert_allclose, assert_equal + +from scipy._lib._util import rng_integers +from scipy._lib._array_api import (is_numpy, make_xp_test_case, xp_default_dtype, + xp_size, array_namespace, _xp_copy_to_numpy, + is_lazy_array, eager_warns) +from scipy._lib._array_api_no_0d import xp_assert_close, xp_assert_equal +from scipy._external import array_api_extra as xpx +from scipy import stats, special +from scipy.fft.tests.test_fftlog import skip_xp_backends +from scipy.optimize import root + +from scipy.stats import bootstrap, monte_carlo_test, permutation_test, power +import scipy.stats._resampling as _resampling + + +lazy_xp_modules = [stats] + + +@make_xp_test_case(bootstrap) +class TestBootstrap: + @skip_xp_backends('numpy', reason='NumPy does not raise') + def test_bootstrap_iv_other(self, xp): + message = f"When using array library {xp.__name__}" + with pytest.raises(TypeError, match=message): + bootstrap((xp.asarray([1, 2, 3]),), lambda x: xp.mean(x)) + + def test_bootstrap_iv(self, xp): + sample = xp.asarray([1, 2, 3]) + + message = "`data` must contain at least one sample." + with pytest.raises(ValueError, match=message): + bootstrap(tuple(), xp.mean) + + message = "each sample in `data` must contain two or more observations..." + with pytest.raises(ValueError, match=message): + bootstrap((sample, xp.asarray([1])), xp.mean) + + message = ("When `paired is True`, all samples must have the same length ") + with pytest.raises(ValueError, match=message): + bootstrap((sample, xp.asarray([1, 2, 3, 4])), xp.mean, paired=True) + + message = "`vectorized` must be `True`, `False`, or `None`." + with pytest.raises(ValueError, match=message): + bootstrap(sample, xp.mean, vectorized='ekki') + + message = "`axis` must be an integer." + with pytest.raises(ValueError, match=message): + bootstrap((sample,), xp.mean, axis=1.5) + + message = "could not convert string to float" + with pytest.raises(ValueError, match=message): + bootstrap((sample,), xp.mean, confidence_level='ni') + + message = "`n_resamples` must be a non-negative integer." + with pytest.raises(ValueError, match=message): + bootstrap((sample,), xp.mean, n_resamples=-1000) + + message = "`n_resamples` must be a non-negative integer." + with pytest.raises(ValueError, match=message): + bootstrap((sample,), xp.mean, n_resamples=1000.5) + + message = "`batch` must be a positive integer or None." + with pytest.raises(ValueError, match=message): + bootstrap((sample,), xp.mean, batch=-1000) + + message = "`batch` must be a positive integer or None." + with pytest.raises(ValueError, match=message): + bootstrap((sample,), xp.mean, batch=1000.5) + + message = "`method` must be in" + with pytest.raises(ValueError, match=message): + bootstrap((sample,), xp.mean, method='ekki') + + message = "`bootstrap_result` must have attribute `bootstrap_distribution'" + with pytest.raises(ValueError, match=message): + bootstrap((sample,), xp.mean, bootstrap_result=10) + + message = "Either `bootstrap_result.bootstrap_distribution.size`" + with pytest.raises(ValueError, match=message): + bootstrap((sample,), xp.mean, n_resamples=0) + + message = "SeedSequence expects int or sequence of ints" + with pytest.raises(TypeError, match=message): + bootstrap((sample,), xp.mean, rng='herring') + + @pytest.mark.parametrize("method", ['basic', 'percentile', 'BCa']) + @pytest.mark.parametrize("axis", [0, 1, 2]) + def test_bootstrap_batch(self, method, axis, xp): + # for one-sample statistics, batch size shouldn't affect the result + rng = np.random.RandomState(0) + + x = rng.rand(10, 11, 12) + # SPEC-007 leave one call with random_state to ensure it still works + res1 = bootstrap((xp.asarray(x),), xp.mean, batch=None, method=method, + random_state=0, axis=axis, n_resamples=100) + rng = np.random.RandomState(0) + res2 = bootstrap((xp.asarray(x),), xp.mean, batch=10, method=method, + axis=axis, n_resamples=100, random_state=rng) + + xp_assert_close(res2.confidence_interval.low, res1.confidence_interval.low) + xp_assert_close(res2.confidence_interval.high, res1.confidence_interval.high) + xp_assert_close(res2.standard_error, res1.standard_error) + + @pytest.mark.parametrize("method", ['basic', 'percentile', 'BCa']) + def test_bootstrap_paired(self, method, xp): + # test that `paired` works as expected + rng = np.random.RandomState(0) + n = 100 + x = xp.asarray(rng.rand(n)) + y = xp.asarray(rng.rand(n)) + + def my_statistic(x, y, axis=-1): + return xp.mean((x-y)**2, axis=axis) + + def my_paired_statistic(i, axis=-1): + a = x[i] + b = y[i] + res = my_statistic(a, b) + return res + + i = xp.arange(x.shape[0]) + + res1 = bootstrap((i,), my_paired_statistic, rng=0) + res2 = bootstrap((x, y), my_statistic, paired=True, rng=0) + + xp_assert_close(res1.confidence_interval.low, res2.confidence_interval.low) + xp_assert_close(res1.confidence_interval.high, res2.confidence_interval.high) + xp_assert_close(res1.standard_error, res2.standard_error) + + @pytest.mark.parametrize("method", ['basic', 'percentile', 'BCa']) + @pytest.mark.parametrize("axis", [0, 1, 2]) + @pytest.mark.parametrize("paired", [True, False]) + def test_bootstrap_vectorized(self, method, axis, paired, xp): + # test that paired is vectorized as expected: when samples are tiled, + # CI and standard_error of each axis-slice is the same as those of the + # original 1d sample + + rng = np.random.RandomState(0) + + def my_statistic(x, y, z, axis=-1): + return xp.mean(x, axis=axis) + xp.mean(y, axis=axis) + xp.mean(z, axis=axis) + + shape = 10, 11, 12 + n_samples = shape[axis] + + x = rng.rand(n_samples) + y = rng.rand(n_samples) + z = rng.rand(n_samples) + x, y, z = xp.asarray(x), xp.asarray(y), xp.asarray(z) + + res1 = bootstrap((x, y, z), my_statistic, paired=paired, method=method, + rng=0, axis=0, n_resamples=100) + assert (res1.bootstrap_distribution.shape + == res1.standard_error.shape + (100,)) + + reshape = [1, 1, 1] + reshape[axis] = n_samples + reshape = tuple(reshape) + x = xp.broadcast_to(xp.reshape(x, reshape), shape) + y = xp.broadcast_to(xp.reshape(y, reshape), shape) + z = xp.broadcast_to(xp.reshape(z, reshape), shape) + res2 = bootstrap((x, y, z), my_statistic, paired=paired, method=method, + rng=0, axis=axis, n_resamples=100) + + result_shape = list(shape) + result_shape.pop(axis) + + ref_ci_low = xp.broadcast_to(res2.confidence_interval.low, result_shape) + ref_ci_high = xp.broadcast_to(res2.confidence_interval.high, result_shape) + ref_ci_standard_error = xp.broadcast_to(res2.standard_error, result_shape) + + xp_assert_close(res2.confidence_interval.low, ref_ci_low) + xp_assert_close(res2.confidence_interval.high, ref_ci_high) + xp_assert_close(res2.standard_error, ref_ci_standard_error) + + @pytest.mark.slow + @pytest.mark.xfail_on_32bit("MemoryError with BCa observed in CI") + @pytest.mark.parametrize("method", ['basic', 'percentile', 'BCa']) + def test_bootstrap_against_theory(self, method, xp): + # based on https://www.statology.org/confidence-intervals-python/ + rng = np.random.default_rng(2442101192988600726) + data = stats.norm.rvs(loc=5, scale=2, size=5000, random_state=rng) + alpha = 0.95 + dist = stats.t(df=len(data)-1, loc=np.mean(data), scale=stats.sem(data)) + ref_low, ref_high = dist.interval(confidence=alpha) + ref_se = dist.std() + + config = dict(data=(xp.asarray(data),), statistic=xp.mean, n_resamples=5000, + method=method, rng=rng) + res = bootstrap(**config, confidence_level=alpha) + + xp_assert_close(res.confidence_interval.low, xp.asarray(ref_low), rtol=5e-4) + xp_assert_close(res.confidence_interval.high, xp.asarray(ref_high), rtol=5e-4) + xp_assert_close(res.standard_error, xp.asarray(ref_se), atol=3e-4) + + config.update(dict(n_resamples=0, bootstrap_result=res)) + res = bootstrap(**config, confidence_level=alpha, alternative='less') + xp_assert_close(res.confidence_interval.high, + xp.asarray(dist.ppf(alpha)), rtol=5e-4) + + config.update(dict(n_resamples=0, bootstrap_result=res)) + res = bootstrap(**config, confidence_level=alpha, alternative='greater') + xp_assert_close(res.confidence_interval.low, + xp.asarray(dist.ppf(1-alpha)), rtol=5e-4) + + tests_R = [("basic", 23.77, 79.12), + ("percentile", 28.86, 84.21), + ("BCa", 32.31, 91.43)] + + @pytest.mark.parametrize("method, ref_low, ref_high", tests_R) + def test_bootstrap_against_R(self, method, ref_low, ref_high, xp): + # Compare against R's "boot" library + # library(boot) + + # stat <- function (x, a) { + # mean(x[a]) + # } + + # x <- c(10, 12, 12.5, 12.5, 13.9, 15, 21, 22, + # 23, 34, 50, 81, 89, 121, 134, 213) + + # # Use a large value so we get a few significant digits for the CI. + # n = 1000000 + # bootresult = boot(x, stat, n) + # result <- boot.ci(bootresult) + # print(result) + x = xp.asarray([10, 12, 12.5, 12.5, 13.9, 15, 21, 22, + 23, 34, 50, 81, 89, 121, 134, 213]) + res = bootstrap((x,), xp.mean, n_resamples=1000000, method=method, rng=0) + xp_assert_close(res.confidence_interval.low, xp.asarray(ref_low), rtol=0.005) + xp_assert_close(res.confidence_interval.high, xp.asarray(ref_high), rtol=0.005) + + def test_multisample_BCa_against_R(self, xp): + # Because bootstrap is stochastic, it's tricky to test against reference + # behavior. Here, we show that SciPy's BCa CI matches R wboot's BCa CI + # much more closely than the other SciPy CIs do. + + # arbitrary skewed data + x = xp.asarray([0.75859206, 0.5910282, -0.4419409, -0.36654601, + 0.34955357, -1.38835871, 0.76735821]) + y = xp.asarray([1.41186073, 0.49775975, 0.08275588, 0.24086388, + 0.03567057, 0.52024419, 0.31966611, 1.32067634]) + + # a multi-sample statistic for which the BCa CI tends to be different + # from the other CIs + def statistic(x, y, axis): + s1 = stats.skew(x, axis=axis) + s2 = stats.skew(y, axis=axis) + return s1 - s2 + + # compute confidence intervals using each method + rng = np.random.default_rng(468865032284792692) + + res_basic = stats.bootstrap((x, y), statistic, method='basic', + batch=100, rng=rng) + res_percent = stats.bootstrap((x, y), statistic, method='percentile', + batch=100, rng=rng) + res_bca = stats.bootstrap((x, y), statistic, method='bca', + batch=100, rng=rng) + + # compute midpoints so we can compare just one number for each + mid_basic = xp.mean(xp.stack(res_basic.confidence_interval)) + mid_percent = xp.mean(xp.stack(res_percent.confidence_interval)) + mid_bca = xp.mean(xp.stack(res_bca.confidence_interval)) + + # reference for BCA CI computed using R wboot package: + # library(wBoot) + # library(moments) + + # x = c(0.75859206, 0.5910282, -0.4419409, -0.36654601, + # 0.34955357, -1.38835871, 0.76735821) + # y = c(1.41186073, 0.49775975, 0.08275588, 0.24086388, + # 0.03567057, 0.52024419, 0.31966611, 1.32067634) + + # twoskew <- function(x1, y1) {skewness(x1) - skewness(y1)} + # boot.two.bca(x, y, skewness, conf.level = 0.95, + # R = 9999, stacked = FALSE) + mid_wboot = -1.5519 + + # compute percent difference relative to wboot BCA method + diff_basic = (mid_basic - mid_wboot)/abs(mid_wboot) + diff_percent = (mid_percent - mid_wboot)/abs(mid_wboot) + diff_bca = (mid_bca - mid_wboot)/abs(mid_wboot) + + # SciPy's BCa CI midpoint is much closer than that of the other methods + assert diff_basic < -0.15 + assert diff_percent > 0.15 + assert abs(diff_bca) < 0.03 + + def test_BCa_acceleration_against_reference(self, xp): + # Compare the (deterministic) acceleration parameter for a multi-sample + # problem against a reference value. The example is from [1], but Efron's + # value seems inaccurate. Straightforward code for computing the + # reference acceleration (0.011008228344026734) is available at: + # https://github.com/scipy/scipy/pull/16455#issuecomment-1193400981 + + y = xp.asarray([10., 27., 31., 40., 46., 50., 52., 104., 146.]) + z = xp.asarray([16., 23., 38., 94., 99., 141., 197.]) + + def statistic(z, y, axis=0): + return xp.mean(z, axis=axis) - xp.mean(y, axis=axis) + + data = [z, y] + res = stats.bootstrap(data, statistic) + + axis = -1 + alpha = 0.95 + theta_hat_b = res.bootstrap_distribution + batch = 100 + _, _, a_hat = _resampling._bca_interval(data, statistic, axis, alpha, + theta_hat_b, batch, xp) + xp_assert_close(a_hat, xp.asarray(0.011008228344026734)) + + tests_against_itself_1samp = {"basic": 1789, + "percentile": 1790, + "BCa": 1789} + + @pytest.mark.slow + @pytest.mark.parametrize("method, expected", + tests_against_itself_1samp.items()) + def test_bootstrap_against_itself_1samp(self, method, expected, xp): + # The expected values in this test were generated using bootstrap + # to check for unintended changes in behavior. The test also makes sure + # that bootstrap works with multi-sample statistics and that the + # `axis` argument works as expected / function is vectorized. + rng = np.random.default_rng(9123847) + + n = 100 # size of sample + n_resamples = 999 # number of bootstrap resamples used to form each CI + confidence_level = 0.9 + + # The true mean is 5 + dist = stats.norm(loc=5, scale=1) + stat_true = float(dist.mean()) + + # Do the same thing 2000 times. (The code is fully vectorized.) + n_replications = 2000 + data = dist.rvs(size=(n_replications, n), random_state=rng) + res = bootstrap((xp.asarray(data),), + statistic=xp.mean, + confidence_level=confidence_level, + n_resamples=n_resamples, + batch=50, + method=method, + axis=-1, + rng=rng) + ci = res.confidence_interval + + # ci contains vectors of lower and upper confidence interval bounds + ci_contains_true = xp.count_nonzero((ci[0] < stat_true) & (stat_true < ci[1])) + assert ci_contains_true == expected + + # ci_contains_true is not inconsistent with confidence_level + pvalue = stats.binomtest(int(ci_contains_true), n_replications, + confidence_level).pvalue + assert pvalue > 0.1 + + tests_against_itself_2samp = {"basic": 892, + "percentile": 890} + + @pytest.mark.slow + @pytest.mark.parametrize("method, expected", + tests_against_itself_2samp.items()) + def test_bootstrap_against_itself_2samp(self, method, expected, xp): + # The expected values in this test were generated using bootstrap + # to check for unintended changes in behavior. The test also makes sure + # that bootstrap works with multi-sample statistics and that the + # `axis` argument works as expected / function is vectorized. + rng = np.random.RandomState(0) + + n1 = 100 # size of sample 1 + n2 = 120 # size of sample 2 + n_resamples = 999 # number of bootstrap resamples used to form each CI + confidence_level = 0.9 + + # The statistic we're interested in is the difference in means + def my_stat(data1, data2, axis=-1): + mean1 = xp.mean(data1, axis=axis) + mean2 = xp.mean(data2, axis=axis) + return mean1 - mean2 + + # The true difference in the means is -0.1 + dist1 = stats.norm(loc=0, scale=1) + dist2 = stats.norm(loc=0.1, scale=1) + stat_true = float(dist1.mean() - dist2.mean()) + + # Do the same thing 1000 times. (The code is fully vectorized.) + n_replications = 1000 + data1 = dist1.rvs(size=(n_replications, n1), random_state=rng) + data2 = dist2.rvs(size=(n_replications, n2), random_state=rng) + res = bootstrap((xp.asarray(data1), xp.asarray(data2)), + statistic=my_stat, + confidence_level=confidence_level, + n_resamples=n_resamples, + batch=50, + method=method, + axis=-1, + random_state=rng) + ci = res.confidence_interval + + # ci contains vectors of lower and upper confidence interval bounds + ci_contains_true = xp.count_nonzero((ci[0] < stat_true) & (stat_true < ci[1])) + assert ci_contains_true == expected + + # ci_contains_true is not inconsistent with confidence_level + pvalue = stats.binomtest(int(ci_contains_true), n_replications, + confidence_level).pvalue + assert pvalue > 0.1 + + @pytest.mark.parametrize("method", ["basic", "percentile", "BCa"]) + @pytest.mark.parametrize("axis", [0, 1]) + @pytest.mark.parametrize("dtype", ['float32', 'float64']) + def test_bootstrap_vectorized_3samp(self, method, axis, dtype, xp): + def statistic(*data, axis=0): + # an arbitrary, vectorized statistic + return sum(xp.mean(sample, axis=axis) for sample in data) + + def statistic_1d(*data): + # the same statistic, not vectorized + for sample in data: + assert sample.ndim == 1 + return np.asarray(sum(sample.mean() for sample in data), dtype=dtype) + + rng = np.random.RandomState(0) + x = rng.rand(4, 5).astype(dtype) + y = rng.rand(4, 5).astype(dtype) + z = rng.rand(4, 5).astype(dtype) + res1 = bootstrap((xp.asarray(x), xp.asarray(y), xp.asarray(z)), + statistic, vectorized=True, axis=axis, n_resamples=100, + method=method, rng=0) + res2 = bootstrap((x, y, z), statistic_1d, vectorized=False, + axis=axis, n_resamples=100, method=method, rng=0) + + rtol = 1e-6 if dtype == 'float32' else 1e-14 + xp_assert_close(res1.confidence_interval.low, + xp.asarray(res2.confidence_interval.low), rtol=rtol) + xp_assert_close(res1.confidence_interval.high, + xp.asarray(res2.confidence_interval.high), rtol=rtol) + xp_assert_close(res1.standard_error, xp.asarray(res2.standard_error), rtol=rtol) + + @pytest.mark.xfail_on_32bit("Failure is not concerning; see gh-14107") + @pytest.mark.parametrize("method", ["basic", "percentile", "BCa"]) + @pytest.mark.parametrize("axis", [0, 1]) + def test_bootstrap_vectorized_1samp(self, method, axis, xp): + def statistic(x, axis=0): + # an arbitrary, vectorized statistic + return xp.mean(x, axis=axis) + + def statistic_1d(x): + # the same statistic, not vectorized + assert x.ndim == 1 + return x.mean(axis=0) + + rng = np.random.default_rng(7939952824) + x = rng.random((2, 3)) + res1 = bootstrap((xp.asarray(x),), statistic, vectorized=True, axis=axis, + n_resamples=100, batch=None, method=method, rng=0) + res2 = bootstrap((x,), statistic_1d, vectorized=False, axis=axis, + n_resamples=100, batch=10, method=method, rng=0) + xp_assert_close(res1.confidence_interval.low, + xp.asarray(res2.confidence_interval.low)) + xp_assert_close(res1.confidence_interval.high, + xp.asarray(res2.confidence_interval.high)) + xp_assert_close(res1.standard_error, xp.asarray(res2.standard_error)) + + @pytest.mark.parametrize("method", ["basic", "percentile", "BCa"]) + def test_bootstrap_degenerate(self, method, xp): + data = xp.full((35,), 10000.) + if method == "BCa": + with np.errstate(invalid='ignore'): + msg = "The BCa confidence interval cannot be calculated" + with eager_warns(stats.DegenerateDataWarning, match=msg, xp=xp): + res = bootstrap([data, ], xp.mean, method=method) + xp_assert_equal(res.confidence_interval.low, xp.asarray(xp.nan)) + xp_assert_equal(res.confidence_interval.high, xp.asarray(xp.nan)) + else: + res = bootstrap([data, ], xp.mean, method=method) + xp_assert_equal(res.confidence_interval.low, xp.asarray(10000.)) + xp_assert_equal(res.confidence_interval.high, xp.asarray(10000.)) + xp_assert_close(res.standard_error, xp.asarray(0.), atol=1e-12) + + @pytest.mark.parametrize("method", ["BCa", "basic", "percentile"]) + def test_bootstrap_gh15678(self, method, xp): + # Check that gh-15678 is fixed: when statistic function returned a Python + # float, method="BCa" failed when trying to add a dimension to the float + rng = np.random.default_rng(354645618886684) + dist = stats.norm(loc=2, scale=4) + data = dist.rvs(size=100, random_state=rng) + res = bootstrap((xp.asarray(data),), stats.skew, method=method, n_resamples=100, + rng=np.random.default_rng(9563)) + # this always worked because np.apply_along_axis returns NumPy data type + ref = bootstrap((data,), stats.skew, method=method, n_resamples=100, + rng=np.random.default_rng(9563), vectorized=False) + xp_assert_close(res.confidence_interval.low, + xp.asarray(ref.confidence_interval.low)) + xp_assert_close(res.confidence_interval.high, + xp.asarray(ref.confidence_interval.high)) + xp_assert_close(res.standard_error, xp.asarray(ref.standard_error)) + + def test_bootstrap_min(self, xp): + # Check that gh-15883 is fixed: percentileofscore should + # behave according to the 'mean' behavior and not trigger nan for BCa + rng = np.random.default_rng(1891289180021102) + dist = stats.norm(loc=2, scale=4) + data = dist.rvs(size=100, random_state=rng) + true_min = np.min(data) + data = xp.asarray(data) + res = bootstrap((data,), xp.min, method="BCa", n_resamples=100, + rng=np.random.default_rng(3942)) + xp_assert_equal(res.confidence_interval.low, xp.asarray(true_min)) + res2 = bootstrap((-data,), xp.max, method="BCa", n_resamples=100, + rng=np.random.default_rng(3942)) + xp_assert_close(-res.confidence_interval.low, res2.confidence_interval.high) + xp_assert_close(-res.confidence_interval.high, res2.confidence_interval.low) + + @pytest.mark.parametrize("additional_resamples", [0, 1000]) + def test_re_bootstrap(self, additional_resamples, xp): + # Test behavior of parameter `bootstrap_result` + rng = np.random.default_rng(8958153316228384) + x = rng.random(size=100) + + n1 = 1000 + n2 = additional_resamples + n3 = n1 + additional_resamples + + rng = np.random.default_rng(296689032789913033) + res = stats.bootstrap((xp.asarray(x),), xp.mean, n_resamples=n1, rng=rng, + confidence_level=0.95, method='percentile') + res = stats.bootstrap((xp.asarray(x),), xp.mean, n_resamples=n2, rng=rng, + confidence_level=0.90, method='BCa', + bootstrap_result=res) + + rng = np.random.default_rng(296689032789913033) + ref = stats.bootstrap((xp.asarray(x),), xp.mean, n_resamples=n3, rng=rng, + confidence_level=0.90, method='BCa') + + xp_assert_close(res.confidence_interval.low, + xp.asarray(ref.confidence_interval.low), + rtol=1e-14) + xp_assert_close(res.confidence_interval.high, + xp.asarray(ref.confidence_interval.high), + rtol=1e-14) + xp_assert_close(res.standard_error, xp.asarray(ref.standard_error), rtol=1e-14) + + @pytest.mark.xfail_on_32bit("Sensitive to machine precision") + @pytest.mark.parametrize("method", ['basic', 'percentile', 'BCa']) + def test_bootstrap_alternative(self, method, xp): + rng = np.random.default_rng(5894822712842015040) + dist = stats.norm(loc=2, scale=4) + data = (xp.asarray(dist.rvs(size=(100), random_state=rng)),) + + config = dict(data=data, statistic=xp.std, rng=rng, axis=-1) + t = stats.bootstrap(**config, confidence_level=0.9) + + config.update(dict(n_resamples=0, bootstrap_result=t)) + l = stats.bootstrap(**config, confidence_level=0.95, alternative='less') + g = stats.bootstrap(**config, confidence_level=0.95, alternative='greater') + + xp_assert_close(l.confidence_interval.high, t.confidence_interval.high, + rtol=1e-14) + xp_assert_close(g.confidence_interval.low, t.confidence_interval.low, + rtol=1e-14) + assert xp.isinf(l.confidence_interval.low) and l.confidence_interval.low < 0 + assert xp.isinf(g.confidence_interval.high) and g.confidence_interval.high > 0 + + with pytest.raises(ValueError, match='`alternative` must be one of'): + stats.bootstrap(**config, alternative='ekki-ekki') + + def test_jackknife_resample(self, xp): + shape = 3, 4, 5, 6 + rng = np.random.default_rng(5274950392) + x = rng.random(size=shape) + y = next(_resampling._jackknife_resample(xp.asarray(x), xp=xp)) + + for i in range(shape[-1]): + # each resample is indexed along second to last axis + # (last axis is the one the statistic will be taken over / consumed) + slc = y[..., i, :] + expected = np.delete(x, i, axis=-1) + + xp_assert_equal(slc, xp.asarray(expected)) + + y2 = list(_resampling._jackknife_resample(xp.asarray(x), batch=2, xp=xp)) + xp_assert_equal(xp.concat(y2, axis=-2), y) + + @pytest.mark.skip_xp_backends("array_api_strict", + reason="Test uses ... + fancy indexing") + @pytest.mark.parametrize("rng_name", ["RandomState", "default_rng"]) + def test_bootstrap_resample(self, rng_name, xp): + rng_gen = getattr(np.random, rng_name) + rng1 = rng_gen(3949441460) + rng2 = rng_gen(3949441460) + + n_resamples = 10 + shape = 3, 4, 5, 6 + + rng = np.random.default_rng(5274950392) + x = xp.asarray(rng.random(shape)) + y = _resampling._bootstrap_resample(x, n_resamples, rng=rng1, xp=xp) + + for i in range(n_resamples): + # each resample is indexed along second to last axis + # (last axis is the one the statistic will be taken over / consumed) + slc = y[..., i, :] + + js = xp.asarray(rng_integers(rng2, 0, shape[-1], shape[-1])) + expected = x[..., js] + + xp_assert_equal(slc, expected) + + @pytest.mark.parametrize("score", [0, 0.5, 1]) + @pytest.mark.parametrize("axis", [0, 1, 2]) + def test_percentile_of_score(self, score, axis, xp): + shape = 10, 20, 30 + rng = np.random.default_rng(5903363153) + x = rng.random(shape) + dtype = xp_default_dtype(xp) + p = _resampling._percentile_of_score(xp.asarray(x, dtype=dtype), + xp.asarray(score, dtype=dtype), + axis=-1, xp=xp) + + def vectorized_pos(a, score, axis): + return np.apply_along_axis(stats.percentileofscore, axis, a, score) + + p2 = vectorized_pos(x, score, axis=-1)/100 + + xp_assert_close(p, xp.asarray(p2, dtype=dtype), rtol=1e-15) + + @pytest.mark.parametrize("axis", [0, 1, 2]) + def test_vectorize_statistic(self, axis): + # test that _vectorize_statistic vectorizes a statistic along `axis` + # only used with NumPy arrays + + def statistic(*data, axis): + # an arbitrary, vectorized statistic + return sum(sample.mean(axis) for sample in data) + + def statistic_1d(*data): + # the same statistic, not vectorized + for sample in data: + assert sample.ndim == 1 + return statistic(*data, axis=0) + + # vectorize the non-vectorized statistic + statistic2 = _resampling._vectorize_statistic(statistic_1d) + + rng = np.random.RandomState(0) + x = rng.rand(4, 5, 6) + y = rng.rand(4, 1, 6) + z = rng.rand(1, 5, 6) + + res1 = statistic(x, y, z, axis=axis) + res2 = statistic2(x, y, z, axis=axis) + assert_allclose(res1, res2) + + @pytest.mark.slow + @pytest.mark.parametrize("method", ["basic", "percentile", "BCa"]) + def test_vector_valued_statistic(self, method, xp): + # Generate 95% confidence interval around MLE of normal distribution + # parameters. Repeat 100 times, each time on sample of size 100. + # Check that confidence interval contains true parameters ~95 times. + # Confidence intervals are estimated and stochastic; a test failure + # does not necessarily indicate that something is wrong. More important + # than values of `counts` below is that the shapes of the outputs are + # correct. + + rng = np.random.default_rng(2196847219) + params = 1, 0.5 + sample = xp.asarray(rng.normal(*params, size=(100, 100))) + + def statistic(data, axis): + return xp.stack([xp.mean(data, axis=axis), + xp.std(data, axis=axis, correction=1)]) + + res = bootstrap((sample,), statistic, method=method, axis=-1, + n_resamples=9999, batch=200, random_state=rng) + + params = xp.asarray([1, 0.5]) + counts = xp.count_nonzero((res.confidence_interval.low.T < params) + & (res.confidence_interval.high.T > params), + axis=0) + assert xp.all(counts >= 90) + assert xp.all(counts <= 100) + assert res.confidence_interval.low.shape == (2, 100) + assert res.confidence_interval.high.shape == (2, 100) + assert res.standard_error.shape == (2, 100) + assert res.bootstrap_distribution.shape == (2, 100, 9999) + + @pytest.mark.slow + @pytest.mark.filterwarnings('ignore::RuntimeWarning') + def test_vector_valued_statistic_gh17715(self): + # gh-17715 reported a mistake introduced in the extension of BCa to + # multi-sample statistics; a `len` should have been `.shape[-1]`. Check + # that this is resolved. + # If this is resolved for NumPy, it's resolved for the rest, + # let's only run this for NumPy. + + rng = np.random.default_rng(141921000979291141) + + def concordance(x, y, axis): + xm = x.mean(axis) + ym = y.mean(axis) + cov = ((x - xm[..., None]) * (y - ym[..., None])).mean(axis) + return (2 * cov) / (x.var(axis) + y.var(axis) + (xm - ym) ** 2) + + def statistic(tp, tn, fp, fn, axis): + actual = tp + fp + expected = tp + fn + return np.nan_to_num(concordance(actual, expected, axis)) + + def statistic_extradim(*args, axis): + return statistic(*args, axis)[np.newaxis, ...] + + data = [[4, 0, 0, 2], # (tp, tn, fp, fn) + [2, 1, 2, 1], + [0, 6, 0, 0], + [0, 6, 3, 0], + [0, 8, 1, 0]] + data = np.array(data).T + + res = bootstrap(data, statistic_extradim, rng=rng, paired=True) + ref = bootstrap(data, statistic, rng=rng, paired=True) + assert_allclose(res.confidence_interval.low[0], + ref.confidence_interval.low, atol=1e-15) + assert_allclose(res.confidence_interval.high[0], + ref.confidence_interval.high, atol=1e-15) + + # torch doesn't have T distribution CDF and can't fall back to NumPy on GPU + @pytest.mark.skip_xp_backends(np_only=True, exceptions=['cupy']) + def test_gh_20850(self, xp): + rng = np.random.default_rng(2085020850) + x = xp.asarray(rng.random((10, 2))) + y = xp.asarray(rng.random((11, 2))) + def statistic(x, y, axis): + return stats.ttest_ind(x, y, axis=axis).statistic + + # The shapes do *not* need to be the same along axis + stats.bootstrap((x, y), statistic) + stats.bootstrap((x.T, y.T), statistic, axis=1) + # But even when the shapes *are* the same along axis, the lengths + # along other dimensions have to be the same (or `bootstrap` warns). + message = "Array shapes are incompatible for broadcasting." + with pytest.raises(ValueError, match=message): + stats.bootstrap((x, y[:10, 0]), statistic) # this won't work after 1.16 + stats.bootstrap((x, y[:10, 0:1]), statistic) # this will + stats.bootstrap((x.T, y.T[0:1, :10]), statistic, axis=1) # this will + +# --- Test Monte Carlo Hypothesis Test --- # + +@make_xp_test_case(monte_carlo_test) +class TestMonteCarloHypothesisTest: + atol = 2.5e-2 # for comparing p-value + + def get_rvs(self, rvs_in, rs, dtype=None, xp=np): + return lambda *args, **kwds: xp.asarray(rvs_in(*args, random_state=rs, **kwds), + dtype=dtype) + + def get_statistic(self, xp): + def statistic(x, axis): + m = xp.mean(x, axis=axis) + v = xp.var(x, axis=axis, correction=1) + n = x.shape[axis] + return m / (v/n)**0.5 + # return stats.ttest_1samp(x, popmean=0., axis=axis).statistic) + return statistic + + def test_input_validation(self, xp): + # test that the appropriate error messages are raised for invalid input + + data = xp.asarray([1., 2., 3.]) + def stat(x, axis=None): + return xp.mean(x, axis=axis) + + message = "Array shapes are incompatible for broadcasting." + temp = (xp.zeros((2, 5)), xp.zeros((3, 5))) + rvs = (stats.norm.rvs, stats.norm.rvs) + with pytest.raises(ValueError, match=message): + monte_carlo_test(temp, rvs, lambda x, y, axis: 1, axis=-1) + + message = "`axis` must be an integer." + with pytest.raises(ValueError, match=message): + monte_carlo_test(data, stats.norm.rvs, stat, axis=1.5) + + message = "`vectorized` must be `True`, `False`, or `None`." + with pytest.raises(ValueError, match=message): + monte_carlo_test(data, stats.norm.rvs, stat, vectorized=1.5) + + message = "`rvs` must be callable or sequence of callables." + with pytest.raises(TypeError, match=message): + monte_carlo_test(data, None, stat) + with pytest.raises(TypeError, match=message): + temp = xp.asarray([[1., 2.], [3., 4.]]) + monte_carlo_test(temp, [lambda x: x, None], stat) + + message = "If `rvs` is a sequence..." + with pytest.raises(ValueError, match=message): + temp = xp.asarray([[1., 2., 3.]]) + monte_carlo_test(temp, [lambda x: x, lambda x: x], stat) + + message = "`statistic` must be callable." + with pytest.raises(TypeError, match=message): + monte_carlo_test(data, stats.norm.rvs, None) + + message = "`n_resamples` must be a positive integer." + with pytest.raises(ValueError, match=message): + monte_carlo_test(data, stats.norm.rvs, stat, n_resamples=-1000) + + message = "`n_resamples` must be a positive integer." + with pytest.raises(ValueError, match=message): + monte_carlo_test(data, stats.norm.rvs, stat, n_resamples=1000.5) + + message = "`batch` must be a positive integer or None." + with pytest.raises(ValueError, match=message): + monte_carlo_test(data, stats.norm.rvs, stat, batch=-1000) + + message = "`batch` must be a positive integer or None." + with pytest.raises(ValueError, match=message): + monte_carlo_test(data, stats.norm.rvs, stat, batch=1000.5) + + message = "`alternative` must be in..." + with pytest.raises(ValueError, match=message): + monte_carlo_test(data, stats.norm.rvs, stat, alternative='ekki') + + # *If* this raises a value error, make sure it has the intended message + message = "Signature inspection of statistic" + def rvs(size): + return xp.asarray(stats.norm.rvs(size=size)) + try: + monte_carlo_test(data, rvs, xp.mean) + except ValueError as e: + assert str(e).startswith(message) + + def test_input_validation_xp(self, xp): + def non_vectorized_statistic(x): + return xp.mean(x) + + message = "`statistic` must be vectorized..." + sample = xp.asarray([1., 2., 3.]) + if is_numpy(xp): + monte_carlo_test(sample, stats.norm.rvs, non_vectorized_statistic) + return + + with pytest.raises(ValueError, match=message): + monte_carlo_test(sample, stats.norm.rvs, non_vectorized_statistic) + with pytest.raises(ValueError, match=message): + monte_carlo_test(sample, stats.norm.rvs, xp.mean, vectorized=False) + + @pytest.mark.xslow + def test_batch(self, xp): + # make sure that the `batch` parameter is respected by checking the + # maximum batch size provided in calls to `statistic` + rng = np.random.default_rng(23492340193) + x = xp.asarray(rng.standard_normal(size=10)) + + def statistic(x, axis): + batch_size = 1 if x.ndim == 1 else x.shape[0] + statistic.batch_size = max(batch_size, statistic.batch_size) + statistic.counter += 1 + return self.get_statistic(xp)(x, axis=axis) + statistic.counter = 0 + statistic.batch_size = 0 + + kwds = {'sample': x, 'statistic': statistic, + 'n_resamples': 1000, 'vectorized': True} + + kwds['rvs'] = self.get_rvs(stats.norm.rvs, np.random.default_rng(328423), xp=xp) + res1 = monte_carlo_test(batch=1, **kwds) + assert_equal(statistic.counter, 1001) + assert_equal(statistic.batch_size, 1) + + kwds['rvs'] = self.get_rvs(stats.norm.rvs, np.random.default_rng(328423), xp=xp) + statistic.counter = 0 + res2 = monte_carlo_test(batch=50, **kwds) + assert_equal(statistic.counter, 21) + assert_equal(statistic.batch_size, 50) + + kwds['rvs'] = self.get_rvs(stats.norm.rvs, np.random.default_rng(328423), xp=xp) + statistic.counter = 0 + res3 = monte_carlo_test(**kwds) + assert_equal(statistic.counter, 2) + assert_equal(statistic.batch_size, 1000) + + xp_assert_equal(res1.pvalue, res3.pvalue) + xp_assert_equal(res2.pvalue, res3.pvalue) + + @pytest.mark.parametrize('axis', range(-3, 3)) + def test_axis_dtype(self, axis, xp): + # test that Nd-array samples are handled correctly for valid values + # of the `axis` parameter; also make sure non-default dtype is maintained + rng = np.random.default_rng(2389234) + size = [2, 3, 4] + size[axis] = 100 + + # Determine non-default dtype + dtype_default = xp.asarray(1.).dtype + dtype_str = 'float32'if ("64" in str(dtype_default)) else 'float64' + dtype_np = getattr(np, dtype_str) + dtype = getattr(xp, dtype_str) + + # ttest_1samp is CPU array-API compatible, but it would be good to + # include CuPy in this test. We'll perform ttest_1samp with a + # NumPy array, but all the rest with be done with fully array-API + # compatible code. + x = rng.standard_normal(size=size, dtype=dtype_np) + expected = stats.ttest_1samp(x, popmean=0., axis=axis) + + x = xp.asarray(x, dtype=dtype) + statistic = self.get_statistic(xp) + rvs = self.get_rvs(stats.norm.rvs, rng, dtype=dtype, xp=xp) + + res = monte_carlo_test(x, rvs, statistic, vectorized=True, + n_resamples=20000, axis=axis) + + ref_statistic = xp.asarray(expected.statistic, dtype=dtype) + ref_pvalue = xp.asarray(expected.pvalue, dtype=dtype) + xp_assert_close(res.statistic, ref_statistic) + xp_assert_close(res.pvalue, ref_pvalue, atol=self.atol) + + @pytest.mark.parametrize('alternative', ("two-sided", "less", "greater")) + def test_alternative(self, alternative, xp): + # test that `alternative` is working as expected + rng = np.random.default_rng(65723433) + + x = rng.standard_normal(size=30) + ref = stats.ttest_1samp(x, 0., alternative=alternative) + + x = xp.asarray(x) + statistic = self.get_statistic(xp) + rvs = self.get_rvs(stats.norm.rvs, rng, xp=xp) + + res = monte_carlo_test(x, rvs, statistic, alternative=alternative) + + xp_assert_close(res.statistic, xp.asarray(ref.statistic)) + xp_assert_close(res.pvalue, xp.asarray(ref.pvalue), atol=self.atol) + + + # Tests below involve statistics that are not yet array-API compatible. + # They can be converted when the statistics are converted. + @pytest.mark.slow + @pytest.mark.parametrize('alternative', ("less", "greater")) + @pytest.mark.parametrize('a', np.linspace(-0.5, 0.5, 5)) # skewness + def test_against_ks_1samp(self, alternative, a): + # test that monte_carlo_test can reproduce pvalue of ks_1samp + rng = np.random.default_rng(65723433) + + x = stats.skewnorm.rvs(a=a, size=30, random_state=rng) + expected = stats.ks_1samp(x, stats.norm.cdf, alternative=alternative) + + def statistic1d(x): + return stats.ks_1samp(x, stats.norm.cdf, mode='asymp', + alternative=alternative).statistic + + norm_rvs = self.get_rvs(stats.norm.rvs, rng) + res = monte_carlo_test(x, norm_rvs, statistic1d, + n_resamples=1000, vectorized=False, + alternative=alternative) + + assert_allclose(res.statistic, expected.statistic) + if alternative == 'greater': + assert_allclose(res.pvalue, expected.pvalue, atol=self.atol) + elif alternative == 'less': + assert_allclose(1-res.pvalue, expected.pvalue, atol=self.atol) + + @pytest.mark.parametrize('hypotest', (stats.skewtest, stats.kurtosistest)) + @pytest.mark.parametrize('alternative', ("less", "greater", "two-sided")) + @pytest.mark.parametrize('a', np.linspace(-2, 2, 5)) # skewness + def test_against_normality_tests(self, hypotest, alternative, a): + # test that monte_carlo_test can reproduce pvalue of normality tests + rng = np.random.default_rng(85723405) + + x = stats.skewnorm.rvs(a=a, size=150, random_state=rng) + expected = hypotest(x, alternative=alternative) + + def statistic(x, axis): + return hypotest(x, axis=axis).statistic + + norm_rvs = self.get_rvs(stats.norm.rvs, rng) + res = monte_carlo_test(x, norm_rvs, statistic, vectorized=True, + alternative=alternative) + + assert_allclose(res.statistic, expected.statistic) + assert_allclose(res.pvalue, expected.pvalue, atol=self.atol) + + @pytest.mark.parametrize('a', np.arange(-2, 3)) # skewness parameter + def test_against_normaltest(self, a): + # test that monte_carlo_test can reproduce pvalue of normaltest + rng = np.random.default_rng(12340513) + + x = stats.skewnorm.rvs(a=a, size=150, random_state=rng) + expected = stats.normaltest(x) + + def statistic(x, axis): + return stats.normaltest(x, axis=axis).statistic + + norm_rvs = self.get_rvs(stats.norm.rvs, rng) + res = monte_carlo_test(x, norm_rvs, statistic, vectorized=True, + alternative='greater') + + assert_allclose(res.statistic, expected.statistic) + assert_allclose(res.pvalue, expected.pvalue, atol=self.atol) + + @pytest.mark.xslow + @pytest.mark.parametrize('a', np.linspace(-0.5, 0.5, 5)) # skewness + def test_against_cramervonmises(self, a): + # test that monte_carlo_test can reproduce pvalue of cramervonmises + rng = np.random.default_rng(234874135) + + x = stats.skewnorm.rvs(a=a, size=30, random_state=rng) + expected = stats.cramervonmises(x, stats.norm.cdf) + + def statistic1d(x): + return stats.cramervonmises(x, stats.norm.cdf).statistic + + norm_rvs = self.get_rvs(stats.norm.rvs, rng) + res = monte_carlo_test(x, norm_rvs, statistic1d, + n_resamples=1000, vectorized=False, + alternative='greater') + + assert_allclose(res.statistic, expected.statistic) + assert_allclose(res.pvalue, expected.pvalue, atol=self.atol) + + @pytest.mark.slow + @pytest.mark.parametrize('dist_name', ['norm', 'logistic']) + @pytest.mark.parametrize('target_statistic', [0.6, 0.7, 0.8]) + def test_against_anderson(self, dist_name, target_statistic): + # test that monte_carlo_test can reproduce results of `anderson`. + + # find the skewness for which the sample statistic is within the range of + # values tubulated by `anderson` + + def fun(a): + rng = np.random.default_rng(394295467) + x = stats.tukeylambda.rvs(a, size=100, random_state=rng) + expected = stats.anderson(x, dist_name, method='interpolate') + return expected.statistic - target_statistic + with warnings.catch_warnings(): + warnings.simplefilter("ignore", RuntimeWarning) + sol = root(fun, x0=0) + assert sol.success + + # get the significance level (p-value) associated with that critical + # value + a = sol.x[0] + rng = np.random.default_rng(394295467) + x = stats.tukeylambda.rvs(a, size=100, random_state=rng) + expected = stats.anderson(x, dist_name, method='interpolate') + expected_stat = expected.statistic + expected_p = expected.pvalue + + # perform equivalent Monte Carlo test and compare results + def statistic1d(x): + return stats.anderson(x, dist_name, method='interpolate').statistic + + dist_rvs = self.get_rvs(getattr(stats, dist_name).rvs, rng) + with warnings.catch_warnings(): + warnings.simplefilter("ignore", RuntimeWarning) + res = monte_carlo_test(x, dist_rvs, + statistic1d, n_resamples=1000, + vectorized=False, alternative='greater') + + assert_allclose(res.statistic, expected_stat) + assert_allclose(res.pvalue, expected_p, atol=2*self.atol) + + def test_p_never_zero(self): + # Use biased estimate of p-value to ensure that p-value is never zero + # per monte_carlo_test reference [1] + rng = np.random.default_rng(2190176673029737545) + x = np.zeros(100) + res = monte_carlo_test(x, rng.random, np.mean, + vectorized=True, alternative='less') + assert res.pvalue == 0.0001 + + def test_against_ttest_ind(self): + # test that `monte_carlo_test` can reproduce results of `ttest_ind`. + rng = np.random.default_rng(219017667302737545) + data = rng.random(size=(2, 5)), rng.random(size=7) # broadcastable + rvs = rng.normal, rng.normal + def statistic(x, y, axis): + return stats.ttest_ind(x, y, axis=axis).statistic + + res = stats.monte_carlo_test(data, rvs, statistic, axis=-1) + ref = stats.ttest_ind(data[0], [data[1]], axis=-1) + assert_allclose(res.statistic, ref.statistic) + assert_allclose(res.pvalue, ref.pvalue, rtol=2e-2) + + def test_against_f_oneway(self): + # test that `monte_carlo_test` can reproduce results of `f_oneway`. + rng = np.random.default_rng(219017667302737545) + data = (rng.random(size=(2, 100)), rng.random(size=(2, 101)), + rng.random(size=(2, 102)), rng.random(size=(2, 103))) + rvs = rng.normal, rng.normal, rng.normal, rng.normal + + def statistic(*args, axis): + return stats.f_oneway(*args, axis=axis).statistic + + res = stats.monte_carlo_test(data, rvs, statistic, axis=-1, + alternative='greater') + ref = stats.f_oneway(*data, axis=-1) + + assert_allclose(res.statistic, ref.statistic) + assert_allclose(res.pvalue, ref.pvalue, atol=1e-2) + + @pytest.mark.fail_slow(2) + @pytest.mark.xfail_on_32bit("Statistic may not depend on sample order on 32-bit") + def test_finite_precision_statistic(self): + # Some statistics return numerically distinct values when the values + # should be equal in theory. Test that `monte_carlo_test` accounts + # for this in some way. + rng = np.random.default_rng(2549824598234528) + n_resamples = 9999 + def rvs(size): + return 1. * stats.bernoulli(p=0.333).rvs(size=size, random_state=rng) + + x = rvs(100) + res = stats.monte_carlo_test(x, rvs, np.var, alternative='less', + n_resamples=n_resamples) + # show that having a tolerance matters + c0 = np.sum(res.null_distribution <= res.statistic) + c1 = np.sum(res.null_distribution <= res.statistic*(1+1e-15)) + assert c0 != c1 + assert res.pvalue == (c1 + 1)/(n_resamples + 1) + + +@make_xp_test_case(power) +class TestPower: + def xp_normal(self, rng, *, xp, dtype=None): + dtype = xp_default_dtype(xp) if dtype is None else dtype + return lambda size: xp.asarray(rng.normal(size=size), dtype=dtype) + + def test_input_validation(self, xp): + # test that the appropriate error messages are raised for invalid input + rng = np.random.default_rng(8519895914314711673) + + test = stats.ttest_ind + rvs = (self.xp_normal(rng, xp=xp), self.xp_normal(rng, xp=xp)) + n_observations = (xp.asarray(10), xp.asarray(12)) + + message = "`vectorized` must be `True`, `False`, or `None`." + with pytest.raises(ValueError, match=message): + power(test, rvs, n_observations, vectorized=1.5) + + message = "`rvs` must be callable or sequence of callables." + with pytest.raises(TypeError, match=message): + power(test, None, n_observations) + with pytest.raises(TypeError, match=message): + power(test, (self.xp_normal(rng, xp=xp), 'ekki'), n_observations) + + message = "If `rvs` is a sequence..." + with pytest.raises(ValueError, match=message): + power(test, (self.xp_normal(rng, xp=xp),), n_observations) + with pytest.raises(ValueError, match=message): + power(test, rvs, (10,)) + + message = "`significance` must be of floating point dtype." + with pytest.raises(ValueError, match=message): + power(test, rvs, n_observations, significance=xp.asarray(True)) + + significance = xp.asarray(2.) + if not is_lazy_array(significance): + message = "All elements of `significance` must be between 0. and 1." + with pytest.raises(ValueError, match=message): + power(test, rvs, n_observations, significance=significance) + + message = "`kwargs` must be a dictionary" + with pytest.raises(TypeError, match=message): + power(test, rvs, n_observations, kwargs=(1, 2, 3)) + + message = ( + "not be broadcast" + "|Chunks do not add" + "|Incompatible shapes" + "|Attempting to broadcast" + ) + with pytest.raises((ValueError, RuntimeError), match=message): + power(test, rvs, (xp.asarray([10, 11]), xp.asarray([12, 13, 14]))) + with pytest.raises((ValueError, RuntimeError), match=message): + power(test, rvs, (xp.asarray([10, 11]), xp.asarray([12, 13])), + kwargs={'x': xp.asarray([1, 2, 3])}) + + message = "`test` must be callable" + with pytest.raises(TypeError, match=message): + power(None, rvs, n_observations) + + message = "`n_resamples` must be a positive integer" + with pytest.raises(ValueError, match=message): + power(test, rvs, n_observations, n_resamples=-10) + with pytest.raises(ValueError, match=message): + power(test, rvs, n_observations, n_resamples=10.5) + + message = "`batch` must be a positive integer" + with pytest.raises(ValueError, match=message): + power(test, rvs, n_observations, batch=-10) + with pytest.raises(ValueError, match=message): + power(test, rvs, n_observations, batch=10.5) + + @pytest.mark.slow + def test_batch(self, xp): + # make sure that the `batch` parameter is respected by checking the + # maximum batch size provided in calls to `test` + rng = np.random.default_rng(23492340193) + + def test(x, axis): + batch_size = 1 if x.ndim == 1 else x.shape[0] + test.batch_size = max(batch_size, test.batch_size) + test.counter += 1 + return stats.ttest_1samp(x, xp.asarray(0.), axis=axis).pvalue + test.counter = 0 + test.batch_size = 0 + + kwds = dict(test=test, + n_observations=xp.asarray(10), + n_resamples=1000) + + rng = np.random.default_rng(23492340193) + res1 = power(**kwds, rvs=self.xp_normal(rng, xp=xp), batch=1) + assert test.counter == 1000 + assert test.batch_size == 1 + + rng = np.random.default_rng(23492340193) + test.counter = 0 + res2 = power(**kwds, rvs=self.xp_normal(rng, xp=xp), batch=50) + assert test.counter == 20 + assert test.batch_size == 50 + + rng = np.random.default_rng(23492340193) + test.counter = 0 + res3 = power(**kwds, rvs=self.xp_normal(rng, xp=xp), batch=1000) + assert test.counter == 1 + assert test.batch_size == 1000 + + xp_assert_equal(res1.power, res3.power) + xp_assert_equal(res2.power, res3.power) + + @pytest.mark.slow + def test_vectorization(self, xp): + # Test that `power` is vectorized as expected + rng = np.random.default_rng(25495254834552) + alternatives = {-1: 'less', 0:'two-sided', 1: 'greater'} + + # Single vectorized call + popmeans = xp.asarray([0, 0.2]) + def test(x, alternative, axis=-1): + # ensure that popmeans axis is zeroth and orthogonal to the rest + popmeans_expanded = xpx.expand_dims(popmeans, + axis=tuple(range(1, x.ndim + 1))) + alternative = alternatives[int(alternative)] + return stats.ttest_1samp(x, popmeans_expanded, alternative=alternative, + axis=axis) + + # nx and kwargs broadcast against one another + nx = xp.asarray([10, 15, 20, 50, 100])[:, xp.newaxis] + kwargs = {'alternative': xp.asarray([-1, 0, 1])} + + # This dimension is added to the beginning + significance = xp.asarray([0.01, 0.025, 0.05, 0.1]) + res = stats.power(test, self.xp_normal(rng, xp=xp), nx, + significance=significance, kwargs=kwargs) + + # Looping over all combinations + ref = [] + for significance_i in significance: + for i in range(nx.shape[0]): + nx_i = nx[i, ...] + for alternative_i in kwargs['alternative']: + for j in range(popmeans.shape[0]): + popmean_j = popmeans[j] + def test2(x, axis=-1): + return stats.ttest_1samp(x, popmean_j, axis=axis, + alternative=alternatives[int(alternative_i)]) + + tmp = stats.power(test2, self.xp_normal(rng, xp=xp), nx_i, + significance=significance_i) + ref.append(tmp.power) + ref = xp.reshape(xp.stack(ref), res.power.shape) + + # Show that results are similar + xp_assert_close(res.power, ref, rtol=2e-2, atol=1e-2) + + @pytest.mark.skip_xp_backends(cpu_only=True, + exceptions=['cupy', 'jax.numpy', 'dask.array']) + def test_ttest_ind_null(self, xp): + # Check that the p-values of `ttest_ind` are uniformly distributed under + # the null hypothesis + rng = np.random.default_rng(254952548345528) + + test = stats.ttest_ind + n_observations = (xp.asarray(rng.integers(10, 100, size=(10))), + xp.asarray(rng.integers(10, 100, size=(10)))) + rvs = self.xp_normal(rng, xp=xp), self.xp_normal(rng, xp=xp) + significance = xp.asarray([0.01, 0.05, 0.1]) + res = stats.power(test, rvs, n_observations, significance=significance) + significance = xp.broadcast_to(significance[:, xp.newaxis], res.power.shape) + xp_assert_close(res.power, significance, atol=1e-2) + + @pytest.mark.skip_xp_backends('array_api_strict', + reason='currently combines integer and float arrays') + @pytest.mark.skip_xp_backends(cpu_only=True, + exceptions=['cupy', 'jax.numpy', 'dask.array']) + def test_ttest_1samp_power(self, xp): + # Check simulated ttest_1samp power against reference + rng = np.random.default_rng(254952548345528) + + # Reference values computed with statmodels + # import numpy as np + # from statsmodels.stats.power import tt_solve_power + # tt_solve_power = np.vectorize(tt_solve_power) + # tt_solve_power([0.1, 0.5, 0.9], [[10], [20]], [[[0.01]], [[0.05]]]) + ref = [[[0.0126515 , 0.10269751, 0.40415802], + [0.01657775, 0.29734608, 0.86228288]], + [[0.0592903 , 0.29317561, 0.71718121], + [0.07094116, 0.56450441, 0.96815163]]] + + kwargs = {'popmean': xp.asarray([0.1, 0.5, 0.9])} + n_observations = xp.asarray([[10], [20]]) + significance = xp.asarray([0.01, 0.05]) + res = stats.power(stats.ttest_1samp, self.xp_normal(rng, xp=xp), n_observations, + significance=significance, kwargs=kwargs) + xp_assert_close(res.power, xp.asarray(ref), atol=1e-2) + + +@make_xp_test_case(permutation_test) +class TestPermutationTest: + + rtol = 1e-14 + + def setup_method(self): + self.rng = np.random.default_rng(7170559330470561044) + + # -- Input validation -- # + def test_permutation_test_iv(self, xp): + + def stat(x, y, axis): + return stats.ttest_ind((x, y), axis).statistic + + data = (xp.asarray([1, 2, 3]), xp.asarray([1, 2, 3])) + + message = "each sample in `data` must contain two or more ..." + with pytest.raises(ValueError, match=message): + permutation_test((data[0], xp.asarray([0])), stat) + + message = "`data` must be a tuple containing at least two samples" + with pytest.raises(ValueError, match=message): + permutation_test((1,), stat) + with pytest.raises(TypeError, match=message): + permutation_test(1, stat) + + message = "`axis` must be an integer." + with pytest.raises(ValueError, match=message): + permutation_test(data, stat, axis=1.5) + + message = "`permutation_type` must be in..." + with pytest.raises(ValueError, match=message): + permutation_test(data, stat, + permutation_type="ekki") + + message = "`vectorized` must be `True`, `False`, or `None`." + with pytest.raises(ValueError, match=message): + permutation_test(data, stat, vectorized=1.5) + + message = "`n_resamples` must be a positive integer." + with pytest.raises(ValueError, match=message): + permutation_test(data, stat, n_resamples=-1000) + + message = "`n_resamples` must be a positive integer." + with pytest.raises(ValueError, match=message): + permutation_test(data, stat, n_resamples=1000.5) + + message = "`batch` must be a positive integer or None." + with pytest.raises(ValueError, match=message): + permutation_test(data, stat, batch=-1000) + + message = "`batch` must be a positive integer or None." + with pytest.raises(ValueError, match=message): + permutation_test(data, stat, batch=1000.5) + + message = "`alternative` must be in..." + with pytest.raises(ValueError, match=message): + permutation_test(data, stat, alternative='ekki') + + message = "SeedSequence expects int or sequence of ints" + with pytest.raises(TypeError, match=message): + permutation_test(data, stat, rng='herring') + + # -- Test Parameters -- # + # SPEC-007 leave one call with seed to check it still works + @pytest.mark.parametrize('random_state', [np.random.RandomState, + np.random.default_rng]) + @pytest.mark.parametrize('permutation_type', + ['pairings', 'samples', 'independent']) + def test_batch(self, permutation_type, random_state, xp): + # make sure that the `batch` parameter is respected by checking the + # maximum batch size provided in calls to `statistic` + x = xp.asarray(self.rng.random(10)) + y = xp.asarray(self.rng.random(10)) + + def statistic(x, y, axis): + batch_size = 1 if x.ndim == 1 else x.shape[0] + statistic.batch_size = max(batch_size, statistic.batch_size) + statistic.counter += 1 + return xp.mean(x, axis=axis) - xp.mean(y, axis=axis) + statistic.counter = 0 + statistic.batch_size = 0 + + kwds = {'n_resamples': 100, 'permutation_type': permutation_type, + 'vectorized': True} + res1 = stats.permutation_test((x, y), statistic, batch=1, + random_state=random_state(0), **kwds) + assert statistic.counter == 101 + assert statistic.batch_size == 1 + + statistic.counter = 0 + res2 = stats.permutation_test((x, y), statistic, batch=50, + random_state=random_state(0), **kwds) + assert statistic.counter == 3 + assert statistic.batch_size == 50 + + statistic.counter = 0 + res3 = stats.permutation_test((x, y), statistic, batch=100, + random_state=random_state(0), **kwds) + assert statistic.counter == 2 + assert statistic.batch_size == 100 + + xp_assert_equal(res1.pvalue, res3.pvalue) + xp_assert_equal(res2.pvalue, res3.pvalue) + + # SPEC-007 leave at least one call with seed to check it still works + @pytest.mark.parametrize('random_state', [np.random.RandomState, + np.random.default_rng]) + @pytest.mark.parametrize('permutation_type, exact_size', + [('pairings', special.factorial(3)**2), + ('samples', 2**3), + ('independent', special.binom(6, 3))]) + def test_permutations(self, permutation_type, exact_size, random_state, xp): + # make sure that the `permutations` parameter is respected by checking + # the size of the null distribution + x = xp.asarray(self.rng.random(3)) + y = xp.asarray(self.rng.random(3)) + + def statistic(x, y, axis): + return xp.mean(x, axis=axis) - xp.mean(y, axis=axis) + + kwds = {'permutation_type': permutation_type, + 'vectorized': True} + res = stats.permutation_test((x, y), statistic, n_resamples=3, + random_state=random_state(0), **kwds) + assert xp_size(res.null_distribution) == 3 + + res = stats.permutation_test((x, y), statistic, **kwds) + assert xp_size(res.null_distribution) == exact_size + + # -- Randomized Permutation Tests -- # + + # To get reasonable accuracy, these next three tests are somewhat slow. + # Originally, I had them passing for all combinations of permutation type, + # alternative, and RNG, but that takes too long for CI. Instead, split + # into three tests, each testing a particular combination of the three + # parameters. + + def test_randomized_test_against_exact_both(self, xp): + # check that the randomized and exact tests agree to reasonable + # precision for permutation_type='both + + alternative, rng = 'less', 0 + + nx, ny, permutations = 8, 9, 24000 + assert special.binom(nx + ny, nx) > permutations + + rng = np.random.default_rng(8235259808) + x = xp.asarray(rng.standard_normal(size=nx)) + y = xp.asarray(rng.standard_normal(size=ny)) + data = x, y + + def statistic(x, y, axis): + return xp.mean(x, axis=axis) - xp.mean(y, axis=axis) + + kwds = {'vectorized': True, 'permutation_type': 'independent', + 'batch': 100, 'alternative': alternative, 'rng': rng} + res = permutation_test(data, statistic, n_resamples=permutations, **kwds) + res2 = permutation_test(data, statistic, n_resamples=xp.inf, **kwds) + + assert res.statistic == res2.statistic + xp_assert_close(res.pvalue, res2.pvalue, atol=1e-2) + + @pytest.mark.slow() + def test_randomized_test_against_exact_samples(self, xp): + # check that the randomized and exact tests agree to reasonable + # precision for permutation_type='samples' + + alternative, rng = 'greater', None + + nx, ny, permutations = 15, 15, 32000 + assert 2**nx > permutations + + rng = np.random.default_rng(8235259808) + x = xp.asarray(rng.standard_normal(size=nx)) + y = xp.asarray(rng.standard_normal(size=ny)) + data = x, y + + def statistic(x, y, axis): + return xp.mean(x - y, axis=axis) + + kwds = {'vectorized': True, 'permutation_type': 'samples', + 'batch': 100, 'alternative': alternative, 'rng': rng} + res = permutation_test(data, statistic, n_resamples=permutations, **kwds) + res2 = permutation_test(data, statistic, n_resamples=xp.inf, **kwds) + + assert res.statistic == res2.statistic + xp_assert_close(res.pvalue, res2.pvalue, atol=1e-2) + + # I only need to skip torch on GPU because it doesn't have betaincc for pearsonr + @pytest.mark.skip_xp_backends(cpu_only=True, exceptions=['cupy', 'jax.numpy']) + def test_randomized_test_against_exact_pairings(self, xp): + # check that the randomized and exact tests agree to reasonable + # precision for permutation_type='pairings' + + alternative, rng = 'two-sided', self.rng + + nx, ny, permutations = 8, 8, 40000 + assert special.factorial(nx) > permutations + + rng = np.random.default_rng(8235259808) + x = xp.asarray(rng.standard_normal(size=nx)) + y = xp.asarray(rng.standard_normal(size=ny)) + data = [x] + + def statistic(x, axis): + return stats.pearsonr(x, y, axis=axis).statistic + + kwds = {'vectorized': True, 'permutation_type': 'samples', + 'batch': 100, 'alternative': alternative, 'rng': rng} + res = permutation_test(data, statistic, n_resamples=permutations, **kwds) + res2 = permutation_test(data, statistic, n_resamples=xp.inf, **kwds) + + assert res.statistic == res2.statistic + xp_assert_close(res.pvalue, res2.pvalue, atol=1e-2) + + # -- Independent (Unpaired) Sample Tests -- # + + @pytest.mark.skip_xp_backends(eager_only=True) # TODO: change to jax_jit=False + @pytest.mark.parametrize('alternative', ("less", "greater", "two-sided")) + def test_against_ks_2samp(self, alternative, xp): + x = self.rng.normal(size=4, scale=1) + y = self.rng.normal(size=5, loc=3, scale=3) + + expected = stats.ks_2samp(x, y, alternative=alternative, mode='exact') + + def statistic(x, y, axis): + # todo: use `xp` as backend when `ks_2samp` supports more backends + x, y = _xp_copy_to_numpy(x), _xp_copy_to_numpy(y) + res = stats.ks_2samp(x, y, axis=axis, mode='asymp', alternative=alternative) + res = xp.asarray(res.statistic) + return res[()] if res.ndim == 0 else res + + # ks_2samp is always a one-tailed 'greater' test + # it's the statistic that changes (D+ vs D- vs max(D+, D-)) + x, y = xp.asarray(x), xp.asarray(y) + res = permutation_test((x, y), statistic, n_resamples=np.inf, + alternative='greater', rng=self.rng) + + xp_assert_close(res.statistic, xp.asarray(expected.statistic), rtol=self.rtol) + xp_assert_close(res.pvalue, xp.asarray(expected.pvalue), rtol=self.rtol) + + @pytest.mark.skip_xp_backends(eager_only=True) # TODO: change to jax_jit=False + @pytest.mark.parametrize('alternative', ("less", "greater", "two-sided")) + def test_against_ansari(self, alternative, xp): + x = self.rng.normal(size=4, scale=1) + y = self.rng.normal(size=5, scale=3) + + # ansari has a different convention for 'alternative' + alternative_correspondence = {"less": "greater", + "greater": "less", + "two-sided": "two-sided"} + alternative_scipy = alternative_correspondence[alternative] + expected = stats.ansari(x, y, alternative=alternative_scipy) + + def statistic(x, y, axis): + # todo: use `xp` as backend when `ansari` is translated to array API + x, y = _xp_copy_to_numpy(x), _xp_copy_to_numpy(y) + res = stats.ansari(x, y, axis=axis) + res = xp.asarray(res.statistic) + return res[()] if res.ndim == 0 else res + + x, y = xp.asarray(x), xp.asarray(y) + res = permutation_test((x, y), statistic, n_resamples=np.inf, + alternative=alternative, rng=self.rng) + + xp_assert_close(res.statistic, xp.asarray(expected.statistic), rtol=self.rtol) + xp_assert_close(res.pvalue, xp.asarray(expected.pvalue), rtol=self.rtol) + + @skip_xp_backends('cupy', reason='needs mannwhitneyu') + @skip_xp_backends(eager_only=True) # mannwhitneyu does input validation + @pytest.mark.parametrize('alternative', ("less", "greater", "two-sided")) + def test_against_mannwhitneyu(self, alternative, xp): + x = stats.uniform.rvs(size=(3, 5, 2), loc=0, random_state=self.rng) + y = stats.uniform.rvs(size=(3, 5, 2), loc=0.05, random_state=self.rng) + + expected = stats.mannwhitneyu(x, y, axis=1, alternative=alternative) + + def statistic(x, y, axis): + return stats.mannwhitneyu(x, y, axis=axis, method='asymptotic').statistic + + x, y = xp.asarray(x), xp.asarray(y) + res = permutation_test((x, y), statistic, vectorized=True, + n_resamples=xp.inf, alternative=alternative, + axis=1, rng=self.rng) + + xp_assert_close(res.statistic, xp.asarray(expected.statistic), rtol=self.rtol) + xp_assert_close(res.pvalue, xp.asarray(expected.pvalue), rtol=self.rtol) + + @skip_xp_backends('cupy', reason='needs cramervonmises_2samp') + @skip_xp_backends(eager_only=True) # cramervonmises_2samp does input validation + @skip_xp_backends(cpu_only=True) # torch doesn't have `kv` + def test_against_cvm(self, xp): + x = stats.norm.rvs(size=4, scale=1, random_state=self.rng) + y = stats.norm.rvs(size=5, loc=3, scale=3, random_state=self.rng) + + expected = stats.cramervonmises_2samp(x, y, method='exact') + + def statistic(x, y, axis): + res = stats.cramervonmises_2samp(x, y, axis=axis, method='asymptotic') + return res.statistic + + # cramervonmises_2samp has only one alternative, greater + x, y = xp.asarray(x), xp.asarray(y) + res = permutation_test((x, y), statistic, n_resamples=np.inf, + alternative='greater', rng=self.rng) + + xp_assert_close(res.statistic, xp.asarray(expected.statistic), rtol=self.rtol) + xp_assert_close(res.pvalue, xp.asarray(expected.pvalue), rtol=self.rtol) + + @skip_xp_backends(eager_only=True) # kruskal does input validation + @pytest.mark.parametrize('axis', (-1, 2)) + def test_vectorized_nsamp_ptype_both(self, axis, xp): + # statistic only available for NumPy + + # Test that permutation_test with permutation_type='independent' works + # properly for a 3-sample statistic with nd array samples of different + # (but compatible) shapes and ndims. Show that exact permutation test + # and random permutation tests approximate SciPy's asymptotic pvalues + # and that exact and random permutation test results are even closer + # to one another (than they are to the asymptotic results). + + # Three samples, different (but compatible) shapes with different ndims + rng = np.random.default_rng(6709265303529651545) + x = rng.random(size=(3)) + y = rng.random(size=(1, 3, 2)) + z = rng.random(size=(2, 1, 4)) + data = (x, y, z) + + expected = stats.kruskal(*data, axis=axis) + + # Define the statistic (and pvalue for comparison) + def statistic(*data, axis): + return stats.kruskal(*data, axis=axis).statistic + + # Calculate exact and randomized permutation results + kwds = {'axis': axis, 'alternative': 'greater', + 'permutation_type': 'independent', 'rng': rng} + data = [xp.asarray(data_) for data_ in data] + res = permutation_test(data, statistic, n_resamples=xp.inf, **kwds) + res2 = permutation_test(data, statistic, n_resamples=1000, **kwds) + + # Check results + xp_assert_close(res.statistic, xp.asarray(expected.statistic), rtol=self.rtol*5) + xp_assert_close(res.statistic, res2.statistic, rtol=self.rtol*5) + xp_assert_close(res.pvalue, xp.asarray(expected.pvalue), atol=6e-2) + xp_assert_close(res.pvalue, res2.pvalue, atol=3e-2) + + # -- Paired-Sample Tests -- # + + @pytest.mark.skip_xp_backends(eager_only=True) # TODO: change to jax_jit=False + @pytest.mark.parametrize('alternative', ("less", "greater", "two-sided")) + def test_against_wilcoxon(self, alternative, xp): + x = stats.uniform.rvs(size=(3, 6, 2), loc=0, random_state=self.rng) + y = stats.uniform.rvs(size=(3, 6, 2), loc=0.05, random_state=self.rng) + expected = stats.wilcoxon(x, y, alternative=alternative, axis=1) + + # We'll check both 1- and 2-sample versions of the same test; + # we expect identical results to wilcoxon in all cases. + def statistic_1samp_1d(z, axis): + # todo: use `xp` as backend when `wilcoxon` is translated to array API + # 'less' ensures we get the same of two statistics every time + z = _xp_copy_to_numpy(z) + res = stats.wilcoxon(z, alternative='less', axis=axis) + res = xp.asarray(res.statistic) + return res[()] if res.ndim == 0 else res + + def statistic_2samp_1d(x, y, axis): + # todo: use `xp` as backend when `wilcoxon` is translated to array API + x, y = _xp_copy_to_numpy(x), _xp_copy_to_numpy(y) + res = stats.wilcoxon(x, y, alternative='less', axis=axis) + res = xp.asarray(res.statistic) + return res[()] if res.ndim == 0 else res + + x, y = xp.asarray(x), xp.asarray(y) + kwds = {'axis': 1, 'alternative': alternative, 'permutation_type': 'samples', + 'rng': self.rng, 'n_resamples': np.inf} + res1 = permutation_test((x-y,), statistic_1samp_1d, **kwds) + res2 = permutation_test((x, y), statistic_2samp_1d, **kwds) + + # `wilcoxon` returns a different statistic with 'two-sided' + xp_assert_close(res1.statistic, res2.statistic, rtol=self.rtol) + if alternative != 'two-sided': + xp_assert_close(res2.statistic, xp.asarray(expected.statistic), + rtol=self.rtol) + + xp_assert_close(res2.pvalue, xp.asarray(expected.pvalue), rtol=self.rtol) + xp_assert_close(res1.pvalue, res2.pvalue, rtol=self.rtol) + + @pytest.mark.parametrize('alternative', ("less", "greater", "two-sided")) + def test_against_binomtest(self, alternative, xp): + + x = self.rng.integers(0, 2, size=10) + x[x == 0] = -1 + # More naturally, the test would flip elements between 0 and one. + # However, permutation_test will flip the _signs_ of the elements. + # So we have to work with +1/-1 instead of 1/0. + + def statistic(x, axis=0): + xp_ = array_namespace(x) + return xp_.count_nonzero(x > 0, axis=axis) + + k, n, p = statistic(x), 10, 0.5 + expected = stats.binomtest(k, n, p, alternative=alternative) + + res = stats.permutation_test((xp.asarray(x, dtype=xp.float64),), statistic, + vectorized=True, + permutation_type='samples', n_resamples=xp.inf, + rng=self.rng, alternative=alternative) + xp_assert_close(res.pvalue, xp.asarray(expected.pvalue), rtol=self.rtol) + + # -- Exact Association Tests -- # + + @pytest.mark.skip_xp_backends(eager_only=True) # TODO: change to jax_jit=False + def test_against_kendalltau(self, xp): + x = self.rng.normal(size=6) + y = x + self.rng.normal(size=6) + + expected = stats.kendalltau(x, y, method='exact') + + def statistic(x, axis): + # todo: use `xp` as backend when `kendalltau` is translated to array API + x = _xp_copy_to_numpy(x) + res = stats.kendalltau(x, y, method='asymptotic', axis=axis) + res = xp.asarray(res.statistic) + return res[()] if res.ndim == 0 else res + + # kendalltau currently has only one alternative, two-sided + x = xp.asarray(x) + res = permutation_test((x,), statistic, permutation_type='pairings', + n_resamples=np.inf, rng=self.rng) + + xp_assert_close(res.statistic, xp.asarray(expected.statistic), rtol=self.rtol) + xp_assert_close(res.pvalue, xp.asarray(expected.pvalue), rtol=self.rtol) + + @pytest.mark.parametrize('alternative', ('less', 'greater', 'two-sided')) + def test_against_fisher_exact(self, alternative, xp): + # x and y are binary random variables with some dependence + rng = np.random.default_rng(6235696159000529929) + x = (rng.random(7) > 0.6).astype(float) + y = (rng.random(7) + 0.25*x > 0.6).astype(float) + tab = stats.contingency.crosstab(x, y)[1] + + x, y = xp.asarray(x), xp.asarray(y) + def statistic(x, axis): + return xp.count_nonzero((x == 1) & (y == 1), axis=axis) + + res = permutation_test((x,), statistic, permutation_type='pairings', + n_resamples=xp.inf, alternative=alternative, + rng=rng) + res2 = stats.fisher_exact(tab, alternative=alternative) + + xp_assert_close(res.pvalue, xp.asarray(res2.pvalue, dtype=x.dtype)) + + @pytest.mark.xslow() + @pytest.mark.parametrize('axis', (-2, 1)) + def test_vectorized_nsamp_ptype_samples(self, axis): + # statistic only available for NumPy, and it's a pain to vectorize + + # Test that permutation_test with permutation_type='samples' works + # properly for a 3-sample statistic with nd array samples of different + # (but compatible) shapes and ndims. Show that exact permutation test + # reproduces SciPy's exact pvalue and that random permutation test + # approximates it. + + x = self.rng.random(size=(2, 4, 3)) + y = self.rng.random(size=(1, 4, 3)) + z = self.rng.random(size=(2, 4, 1)) + x = stats.rankdata(x, axis=axis) + y = stats.rankdata(y, axis=axis) + z = stats.rankdata(z, axis=axis) + y = y[0] # to check broadcast with different ndim + data = (x, y, z) + + def statistic1d(*data): + return stats.page_trend_test(data, ranked=True, + method='asymptotic').statistic + + def pvalue1d(*data): + return stats.page_trend_test(data, ranked=True, + method='exact').pvalue + + statistic = _resampling._vectorize_statistic(statistic1d) + pvalue = _resampling._vectorize_statistic(pvalue1d) + + expected_statistic = statistic(*np.broadcast_arrays(*data), axis=axis) + expected_pvalue = pvalue(*np.broadcast_arrays(*data), axis=axis) + + # Let's forgive this use of an integer seed, please. + kwds = {'vectorized': False, 'axis': axis, 'alternative': 'greater', + 'permutation_type': 'pairings', 'rng': 0} + res = permutation_test(data, statistic1d, n_resamples=np.inf, **kwds) + res2 = permutation_test(data, statistic1d, n_resamples=5000, **kwds) + + assert_allclose(res.statistic, expected_statistic, rtol=self.rtol) + assert_allclose(res.statistic, res2.statistic, rtol=self.rtol) + assert_allclose(res.pvalue, expected_pvalue, rtol=self.rtol) + assert_allclose(res.pvalue, res2.pvalue, atol=3e-2) + + # -- Test Against External References -- # + + tie_case_1 = {'x': [1, 2, 3, 4], 'y': [1.5, 2, 2.5], + 'expected_less': 0.2000000000, + 'expected_2sided': 0.4, # 2*expected_less + 'expected_Pr_gte_S_mean': 0.3428571429, # see note below + 'expected_statistic': 7.5, + 'expected_avg': 9.142857, 'expected_std': 1.40698} + tie_case_2 = {'x': [111, 107, 100, 99, 102, 106, 109, 108], + 'y': [107, 108, 106, 98, 105, 103, 110, 105, 104], + 'expected_less': 0.1555738379, + 'expected_2sided': 0.3111476758, + 'expected_Pr_gte_S_mean': 0.2969971205, # see note below + 'expected_statistic': 32.5, + 'expected_avg': 38.117647, 'expected_std': 5.172124} + + @pytest.mark.skip_xp_backends(eager_only=True) # TODO: change to jax_jit=False + @pytest.mark.xslow() # only the second case is slow, really + @pytest.mark.parametrize('case', (tie_case_1, tie_case_2)) + def test_with_ties(self, case, xp): + """ + Results above from SAS PROC NPAR1WAY, e.g. + + DATA myData; + INPUT X Y; + CARDS; + 1 1 + 1 2 + 1 3 + 1 4 + 2 1.5 + 2 2 + 2 2.5 + ods graphics on; + proc npar1way AB data=myData; + class X; + EXACT; + run; + ods graphics off; + + Note: SAS provides Pr >= |S-Mean|, which is different from our + definition of a two-sided p-value. + + """ + + x = case['x'] + y = case['y'] + + expected_statistic = xp.asarray(case['expected_statistic']) + expected_less = xp.asarray(case['expected_less']) + expected_2sided = xp.asarray(case['expected_2sided']) + expected_Pr_gte_S_mean = xp.asarray(case['expected_Pr_gte_S_mean']) + expected_avg = xp.asarray(case['expected_avg']) + expected_std = xp.asarray(case['expected_std']) + + def statistic(x, y, axis): + # todo: use `xp` as backend when `ansari` is translated to array API + x, y = _xp_copy_to_numpy(x), _xp_copy_to_numpy(y) + res = stats.ansari(x, y, axis=axis) + res = xp.asarray(res.statistic) + return res[()] if res.ndim == 0 else res + + dtype = xp_default_dtype(xp) + x, y = xp.asarray(x, dtype=dtype), xp.asarray(y, dtype=dtype) + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", "Ties preclude use of exact statistic", UserWarning) + res = permutation_test((x, y), statistic, n_resamples=np.inf, + alternative='less') + res2 = permutation_test((x, y), statistic, n_resamples=np.inf, + alternative='two-sided') + + xp_assert_close(res.statistic, expected_statistic, rtol=self.rtol) + xp_assert_close(res.pvalue, expected_less, atol=1e-10) + xp_assert_close(res2.pvalue, expected_2sided, atol=1e-10) + xp_assert_close(xp.mean(res2.null_distribution), expected_avg, rtol=1e-6) + xp_assert_close(xp.std(res2.null_distribution), expected_std, rtol=1e-6) + + # SAS provides Pr >= |S-Mean|; might as well check against that, too + S = res.statistic + mean = xp.mean(res.null_distribution) + n = res.null_distribution.shape[0] + Pr_gte_S_mean = xp.astype(xp.count_nonzero( + xp.abs(res.null_distribution-mean) >= xp.abs(S-mean)), S.dtype) / n + xp_assert_close(Pr_gte_S_mean, expected_Pr_gte_S_mean) + + @pytest.mark.slow + @pytest.mark.parametrize('alternative, expected_pvalue', + (('less', 0.9708333333333), + ('greater', 0.05138888888889), + ('two-sided', 0.1027777777778))) + # I only need to skip torch on GPU because it doesn't have betaincc for pearsonr + @pytest.mark.skip_xp_backends(cpu_only=True, exceptions=['cupy', 'jax.numpy']) + @pytest.mark.skip_xp_backends(eager_only=True) # TODO: change to jax_jit=False + def test_against_spearmanr_in_R(self, alternative, expected_pvalue, xp): + """ + Results above from R cor.test, e.g. + + options(digits=16) + x <- c(1.76405235, 0.40015721, 0.97873798, + 2.2408932, 1.86755799, -0.97727788) + y <- c(2.71414076, 0.2488, 0.87551913, + 2.6514917, 2.01160156, 0.47699563) + cor.test(x, y, method = "spearm", alternative = "t") + """ + # data comes from + # np.random.seed(0) + # x = stats.norm.rvs(size=6) + # y = x + stats.norm.rvs(size=6) + x = xp.asarray([1.76405235, 0.40015721, 0.97873798, + 2.2408932, 1.86755799, -0.97727788]) + y = xp.asarray([2.71414076, 0.2488, 0.87551913, + 2.6514917, 2.01160156, 0.47699563]) + expected_statistic = 0.7714285714285715 + + y = xp.asarray(stats.rankdata(_xp_copy_to_numpy(y))) + def statistic(x, axis): + # `spearmanr` is not array api compatible, but `pearsonr` is. So for now + # use _xp_copy_to_numpy just for ranking so we can run this test w/ CuPy. + # TODO: use `xp` as backend when cupy works with `rankdata` + x = xp.asarray(stats.rankdata(_xp_copy_to_numpy(x), axis=axis)) + return stats.pearsonr(x, y, axis=axis).statistic + + res = permutation_test((x,), statistic, permutation_type='pairings', + n_resamples=xp.inf, alternative=alternative) + + xp_assert_close(res.statistic, xp.asarray(expected_statistic), rtol=self.rtol) + xp_assert_close(res.pvalue, xp.asarray(expected_pvalue), atol=1e-13) + + @pytest.mark.parametrize("batch", (-1, 0)) + def test_batch_generator_iv(self, batch): + with pytest.raises(ValueError, match="`batch` must be positive."): + list(_resampling._batch_generator([1, 2, 3], batch)) + + batch_generator_cases = [(range(0), 3, []), + (range(6), 3, [[0, 1, 2], [3, 4, 5]]), + (range(8), 3, [[0, 1, 2], [3, 4, 5], [6, 7]])] + + @pytest.mark.parametrize("iterable, batch, expected", + batch_generator_cases) + def test_batch_generator(self, iterable, batch, expected): + got = list(_resampling._batch_generator(iterable, batch)) + assert got == expected + + @pytest.mark.fail_slow(2) + # I only need to skip torch on GPU because it doesn't have betaincc for pearsonr + @pytest.mark.skip_xp_backends(cpu_only=True, exceptions=['cupy', 'jax.numpy']) + def test_finite_precision_statistic(self, xp): + # Some statistics return numerically distinct values when the values + # should be equal in theory. Test that `permutation_test` accounts + # for this in some way. + x = xp.asarray([1., 2., 4., 3.], dtype=xp.float64) + y = xp.asarray([2., 4., 6., 8.], dtype=xp.float64) + + def statistic(x, y, axis): + return stats.pearsonr(x, y, axis=axis)[0] + + res = stats.permutation_test((x, y), statistic, + permutation_type='pairings') + r, pvalue, null = res.statistic, res.pvalue, res.null_distribution + + correct_p = 2 * float(xp.count_nonzero(null >= r - 1e-14)) / null.shape[0] + assert pvalue == correct_p == 1/3 + # Compare against other exact correlation tests using R corr.test + # options(digits=16) + # x = c(1, 2, 4, 3) + # y = c(2, 4, 6, 8) + # cor.test(x, y, alternative = "t", method = "spearman") # 0.333333333 + # cor.test(x, y, alternative = "t", method = "kendall") # 0.333333333 + + def test_nan_statistic(self, xp): + # Where the observed statistic is NaN, the p-value should be NaN + # (Whether a single NaN in the permutation distribution should make the p-value + # NaN is debatable. Users can choose for themselves.) + rng = np.random.default_rng(8951482112) + dtype = xp_default_dtype(xp) + x = xp.asarray(rng.random(5), dtype=dtype) + x_nan = xp.asarray([0, 0, 0, xp.nan, 0]) + y = xp.asarray(rng.random(6), dtype=dtype) + + def statistic(x, y, axis): + return xp.mean(x, axis=axis) - xp.mean(y, axis=axis) + + kwds = {'permutation_type': 'independent', 'rng': rng, 'axis': -1} + + ref = permutation_test((x, y), statistic, **kwds) + + data = (xp.stack([x, x + x_nan]), y) + res = permutation_test(data, statistic, **kwds) + + NaN = xp.asarray(xp.nan, dtype=dtype) + xp_assert_equal(res.statistic[0], ref.statistic) + xp_assert_equal(res.pvalue[0], ref.pvalue) + xp_assert_equal(res.statistic[1], NaN) + xp_assert_equal(res.pvalue[1], NaN) + + +def test_all_partitions_concatenated(): + # make sure that _all_paritions_concatenated produces the correct number + # of partitions of the data into samples of the given sizes and that + # all are unique + n = np.array([3, 2, 4], dtype=int) + nc = np.cumsum(n) + + all_partitions = set() + counter = 0 + for partition_concatenated in _resampling._all_partitions_concatenated(n): + counter += 1 + partitioning = np.split(partition_concatenated, nc[:-1]) + all_partitions.add(tuple([frozenset(i) for i in partitioning])) + + expected = np.prod([special.binom(sum(n[i:]), sum(n[i+1:])) + for i in range(len(n)-1)]) + + assert_equal(counter, expected) + assert_equal(len(all_partitions), expected) + + +@pytest.mark.parametrize('fun_name', + ['bootstrap', 'permutation_test', 'monte_carlo_test']) +def test_parameter_vectorized(fun_name): + # Check that parameter `vectorized` is working as desired for all + # resampling functions. Results don't matter; just don't fail asserts. + rng = np.random.default_rng(75245098234592) + sample = rng.random(size=10) + + def rvs(size): # needed by `monte_carlo_test` + return stats.norm.rvs(size=size, random_state=rng) + + fun_options = {'bootstrap': {'data': (sample,), 'rng': rng, + 'method': 'percentile'}, + 'permutation_test': {'data': (sample,), 'rng': rng, + 'permutation_type': 'samples'}, + 'monte_carlo_test': {'sample': sample, 'rvs': rvs}} + common_options = {'n_resamples': 100} + + fun = getattr(stats, fun_name) + options = fun_options[fun_name] + options.update(common_options) + + def statistic(x, axis): + assert x.ndim > 1 or np.array_equal(x, sample) + return np.mean(x, axis=axis) + fun(statistic=statistic, vectorized=None, **options) + fun(statistic=statistic, vectorized=True, **options) + + def statistic(x): + assert x.ndim == 1 + return np.mean(x) + fun(statistic=statistic, vectorized=None, **options) + fun(statistic=statistic, vectorized=False, **options) + + +class TestMonteCarloMethod: + def test_rvs_and_random_state(self): + message = "Use of `rvs` and `rng` are mutually exclusive." + rng = np.random.default_rng(34982345) + with pytest.raises(ValueError, match=message): + stats.MonteCarloMethod(rvs=rng.random, rng=rng) diff --git a/v12x/pylib/scipy/stats/tests/test_sampling.py b/v12x/pylib/scipy/stats/tests/test_sampling.py new file mode 100644 index 0000000000000000000000000000000000000000..6d50354d10e6713b6effeede7e92be226df01483 --- /dev/null +++ b/v12x/pylib/scipy/stats/tests/test_sampling.py @@ -0,0 +1,1455 @@ +import threading +import pickle +import pytest +from copy import deepcopy +import platform +import sys +import math +import warnings + +import numpy as np +from numpy.testing import assert_allclose, assert_equal +from scipy.stats.sampling import ( + TransformedDensityRejection, + DiscreteAliasUrn, + DiscreteGuideTable, + NumericalInversePolynomial, + NumericalInverseHermite, + RatioUniforms, + SimpleRatioUniforms, + UNURANError +) +from pytest import raises as assert_raises +from scipy import stats +from scipy import special +from scipy.stats import chisquare, cramervonmises +from scipy.stats._distr_params import distdiscrete, distcont +from scipy._lib._util import check_random_state +from scipy._lib._gcutils import assert_deallocated + + +# common test data: this data can be shared between all the tests. + + +# Normal distribution shared between all the continuous methods +class StandardNormal: + def pdf(self, x): + # normalization constant needed for NumericalInverseHermite + return 1./np.sqrt(2.*np.pi) * np.exp(-0.5 * x*x) + + def dpdf(self, x): + return 1./np.sqrt(2.*np.pi) * -x * np.exp(-0.5 * x*x) + + def cdf(self, x): + return special.ndtr(x) + + +all_methods = [ + ("TransformedDensityRejection", {"dist": StandardNormal()}), + ("DiscreteAliasUrn", {"dist": [0.02, 0.18, 0.8]}), + ("DiscreteGuideTable", {"dist": [0.02, 0.18, 0.8]}), + ("NumericalInversePolynomial", {"dist": StandardNormal()}), + ("NumericalInverseHermite", {"dist": StandardNormal()}), + ("SimpleRatioUniforms", {"dist": StandardNormal(), "mode": 0}) +] + +floaterr = r"must be real number, not list" +# Make sure an internal error occurs in UNU.RAN when invalid callbacks are +# passed. Moreover, different generators throw different error messages. +# So, in case of an `UNURANError`, we do not validate the error message. +bad_pdfs_common = [ + # Negative PDF + (lambda x: -x, UNURANError, r"..."), + # Returning wrong type + (lambda x: [], TypeError, floaterr), + # Undefined name inside the function + (lambda x: foo, NameError, r"name 'foo' is not defined"), # type: ignore[name-defined] # noqa: F821, E501 + # Infinite value returned => Overflow error. + (lambda x: np.inf, UNURANError, r"..."), + # NaN value => internal error in UNU.RAN + (lambda x: np.nan, UNURANError, r"..."), + # signature of PDF wrong + (lambda: 1.0, TypeError, r"takes 0 positional arguments but 1 was given") +] + + +# same approach for dpdf +bad_dpdf_common = [ + # Infinite value returned. + (lambda x: np.inf, UNURANError, r"..."), + # NaN value => internal error in UNU.RAN + (lambda x: np.nan, UNURANError, r"..."), + # Returning wrong type + (lambda x: [], TypeError, floaterr), + # Undefined name inside the function + (lambda x: foo, NameError, r"name 'foo' is not defined"), # type: ignore[name-defined] # noqa: F821, E501 + # signature of dPDF wrong + (lambda: 1.0, TypeError, r"takes 0 positional arguments but 1 was given") +] + + +# same approach for logpdf +bad_logpdfs_common = [ + # Returning wrong type + (lambda x: [], TypeError, floaterr), + # Undefined name inside the function + (lambda x: foo, NameError, r"name 'foo' is not defined"), # type: ignore[name-defined] # noqa: F821, E501 + # Infinite value returned => Overflow error. + (lambda x: np.inf, UNURANError, r"..."), + # NaN value => internal error in UNU.RAN + (lambda x: np.nan, UNURANError, r"..."), + # signature of logpdf wrong + (lambda: 1.0, TypeError, r"takes 0 positional arguments but 1 was given") +] + + +bad_pv_common = [ + ([], r"must contain at least one element"), + ([[1.0, 0.0]], r"wrong number of dimensions \(expected 1, got 2\)"), + ([0.2, 0.4, np.nan, 0.8], r"must contain only finite / non-nan values"), + ([0.2, 0.4, np.inf, 0.8], r"must contain only finite / non-nan values"), + ([0.0, 0.0], r"must contain at least one non-zero value"), +] + + +# size of the domains is incorrect +bad_sized_domains = [ + # > 2 elements in the domain + ((1, 2, 3), ValueError, r"must be a length 2 tuple"), + # empty domain + ((), ValueError, r"must be a length 2 tuple") +] + +# domain values are incorrect +bad_domains = [ + ((2, 1), UNURANError, r"left >= right"), + ((1, 1), UNURANError, r"left >= right"), +] + +# infinite and nan values present in domain. +inf_nan_domains = [ + # left >= right + ((10, 10), UNURANError, r"left >= right"), + ((np.inf, np.inf), UNURANError, r"left >= right"), + ((-np.inf, -np.inf), UNURANError, r"left >= right"), + ((np.inf, -np.inf), UNURANError, r"left >= right"), + # Also include nans in some of the domains. + ((-np.inf, np.nan), ValueError, r"only non-nan values"), + ((np.nan, np.inf), ValueError, r"only non-nan values") +] + +# `nan` values present in domain. Some distributions don't support +# infinite tails, so don't mix the nan values with infinities. +nan_domains = [ + ((0, np.nan), ValueError, r"only non-nan values"), + ((np.nan, np.nan), ValueError, r"only non-nan values") +] + + +# all the methods should throw errors for nan, bad sized, and bad valued +# domains. +@pytest.mark.parametrize("domain, err, msg", + bad_domains + bad_sized_domains + + nan_domains) # type: ignore[operator] +@pytest.mark.parametrize("method, kwargs", all_methods) +def test_bad_domain(domain, err, msg, method, kwargs): + Method = getattr(stats.sampling, method) + with pytest.raises(err, match=msg): + Method(**kwargs, domain=domain) + + +@pytest.mark.parametrize("method, kwargs", all_methods) +def test_random_state(method, kwargs): + Method = getattr(stats.sampling, method) + + # simple seed that works for any version of NumPy + seed = 123 + rng1 = Method(**kwargs, random_state=seed) + rng2 = Method(**kwargs, random_state=seed) + assert_equal(rng1.rvs(100), rng2.rvs(100)) + + # global seed + rng = np.random.RandomState(123) + rng1 = Method(**kwargs) + rvs1 = rng1.rvs(100, random_state=rng) + np.random.seed(None) # valid use of np.random.seed + rng2 = Method(**kwargs, random_state=123) + rvs2 = rng2.rvs(100) + assert_equal(rvs1, rvs2) + + # Generator seed for new NumPy + # when a RandomState is given, it should take the bitgen_t + # member of the class and create a Generator instance. + seed1 = np.random.RandomState(np.random.MT19937(123)) + seed2 = np.random.Generator(np.random.MT19937(123)) + rng1 = Method(**kwargs, random_state=seed1) + rng2 = Method(**kwargs, random_state=seed2) + assert_equal(rng1.rvs(100), rng2.rvs(100)) + + +def test_set_random_state(): + rng1 = TransformedDensityRejection(StandardNormal(), random_state=123) + rng2 = TransformedDensityRejection(StandardNormal()) + rng2.set_random_state(123) + assert_equal(rng1.rvs(100), rng2.rvs(100)) + rng = TransformedDensityRejection(StandardNormal(), random_state=123) + rvs1 = rng.rvs(100) + rng.set_random_state(123) + rvs2 = rng.rvs(100) + assert_equal(rvs1, rvs2) + + +def test_threading_behaviour(): + # Test if the API is thread-safe. + # This verifies if the lock mechanism and the use of `PyErr_Occurred` + # is correct. + errors = {"err1": None, "err2": None} + + class Distribution: + def __init__(self, pdf_msg): + self.pdf_msg = pdf_msg + + def pdf(self, x): + if 49.9 < x < 50.0: + raise ValueError(self.pdf_msg) + return x + + def dpdf(self, x): + return 1 + + def func1(): + dist = Distribution('foo') + rng = TransformedDensityRejection(dist, domain=(10, 100), + random_state=12) + try: + rng.rvs(100000) + except ValueError as e: + errors['err1'] = e.args[0] + + def func2(): + dist = Distribution('bar') + rng = TransformedDensityRejection(dist, domain=(10, 100), + random_state=2) + try: + rng.rvs(100000) + except ValueError as e: + errors['err2'] = e.args[0] + + t1 = threading.Thread(target=func1) + t2 = threading.Thread(target=func2) + + t1.start() + t2.start() + + t1.join() + t2.join() + + assert errors['err1'] == 'foo' + assert errors['err2'] == 'bar' + + +@pytest.mark.parametrize("method, kwargs", all_methods) +def test_pickle(method, kwargs): + Method = getattr(stats.sampling, method) + rng1 = Method(**kwargs, random_state=123) + obj = pickle.dumps(rng1) + rng2 = pickle.loads(obj) + assert_equal(rng1.rvs(100), rng2.rvs(100)) + + +@pytest.mark.parametrize("size", [None, 0, (0, ), 1, (10, 3), (2, 3, 4, 5), + (0, 0), (0, 1)]) +def test_rvs_size(size): + # As the `rvs` method is present in the base class and shared between + # all the classes, we can just test with one of the methods. + rng = TransformedDensityRejection(StandardNormal()) + if size is None: + assert np.isscalar(rng.rvs(size)) + else: + if np.isscalar(size): + size = (size, ) + assert rng.rvs(size).shape == size + + +def test_with_scipy_distribution(): + # test if the setup works with SciPy's rv_frozen distributions + dist = stats.norm() + urng = np.random.default_rng(0) + rng = NumericalInverseHermite(dist, random_state=urng) + u = np.linspace(0, 1, num=100) + check_cont_samples(rng, dist, dist.stats()) + assert_allclose(dist.ppf(u), rng.ppf(u)) + # test if it works with `loc` and `scale` + dist = stats.norm(loc=10., scale=5.) + rng = NumericalInverseHermite(dist, random_state=urng) + check_cont_samples(rng, dist, dist.stats()) + assert_allclose(dist.ppf(u), rng.ppf(u)) + # check for discrete distributions + dist = stats.binom(10, 0.2) + rng = DiscreteAliasUrn(dist, random_state=urng) + domain = dist.support() + pv = dist.pmf(np.arange(domain[0], domain[1]+1)) + check_discr_samples(rng, pv, dist.stats()) + + +def test_NumericalInverseHermite_refcycle(): + # test if NumericalInverseHermite contains a reference cycle + dist = stats.norm() + urng = np.random.default_rng(0) + with assert_deallocated(NumericalInverseHermite, dist, random_state=urng) as rng: + u = np.linspace(0, 1, num=100) + check_cont_samples(rng, dist, dist.stats()) + assert_allclose(dist.ppf(u), rng.ppf(u)) + del rng + + +def check_cont_samples(rng, dist, mv_ex, rtol=1e-7, atol=1e-1): + rvs = rng.rvs(100000) + mv = rvs.mean(), rvs.var() + # test the moments only if the variance is finite + if np.isfinite(mv_ex[1]): + assert_allclose(mv, mv_ex, rtol=rtol, atol=atol) + # Cramer Von Mises test for goodness-of-fit + rvs = rng.rvs(500) + dist.cdf = np.vectorize(dist.cdf) + pval = cramervonmises(rvs, dist.cdf).pvalue + assert pval > 0.1 + + +def check_discr_samples(rng, pv, mv_ex, rtol=1e-3, atol=1e-1): + rvs = rng.rvs(100000) + # test if the first few moments match + mv = rvs.mean(), rvs.var() + assert_allclose(mv, mv_ex, rtol=rtol, atol=atol) + # normalize + pv = pv / pv.sum() + # chi-squared test for goodness-of-fit + obs_freqs = np.zeros_like(pv) + _, freqs = np.unique(rvs, return_counts=True) + freqs = freqs / freqs.sum() + obs_freqs[:freqs.size] = freqs + pval = chisquare(obs_freqs, pv).pvalue + assert pval > 0.1 + + +def test_warning_center_not_in_domain(): + # UNURAN will warn if the center provided or the one computed w/o the + # domain is outside of the domain + msg = "102 : center moved into domain of distribution" + with pytest.warns(RuntimeWarning, match=msg): + NumericalInversePolynomial(StandardNormal(), center=0, domain=(3, 5)) + with pytest.warns(RuntimeWarning, match=msg): + NumericalInversePolynomial(StandardNormal(), domain=(3, 5)) + + +@pytest.mark.parametrize('method', ["SimpleRatioUniforms", + "NumericalInversePolynomial", + "TransformedDensityRejection"]) +def test_error_mode_not_in_domain(method): + # UNURAN raises an error if the mode is not in the domain + # the behavior is different compared to the case that center is not in the + # domain. mode is supposed to be the exact value, center can be an + # approximate value + Method = getattr(stats.sampling, method) + msg = "17 : mode not in domain" + with pytest.raises(UNURANError, match=msg): + Method(StandardNormal(), mode=0, domain=(3, 5)) + + +@pytest.mark.parametrize('method', ["NumericalInverseHermite", + "NumericalInversePolynomial"]) +class TestQRVS: + def test_input_validation(self, method): + match = "`qmc_engine` must be an instance of..." + with pytest.raises(ValueError, match=match): + Method = getattr(stats.sampling, method) + gen = Method(StandardNormal()) + gen.qrvs(qmc_engine=0) + + # issues with QMCEngines and old NumPy + Method = getattr(stats.sampling, method) + gen = Method(StandardNormal()) + + match = "`d` must be consistent with dimension of `qmc_engine`." + with pytest.raises(ValueError, match=match): + gen.qrvs(d=3, qmc_engine=stats.qmc.Halton(2)) + + qrngs = [None, stats.qmc.Sobol(1, seed=0), stats.qmc.Halton(3, seed=0)] + # `size=None` should not add anything to the shape, `size=1` should + sizes = [(None, tuple()), (1, (1,)), (4, (4,)), + ((4,), (4,)), ((2, 4), (2, 4))] # type: ignore + # Neither `d=None` nor `d=1` should add anything to the shape + ds = [(None, tuple()), (1, tuple()), (3, (3,))] + + @pytest.mark.parametrize('qrng', qrngs) + @pytest.mark.parametrize('size_in, size_out', sizes) + @pytest.mark.parametrize('d_in, d_out', ds) + @pytest.mark.thread_unsafe(reason="fails in parallel") + def test_QRVS_shape_consistency(self, qrng, size_in, size_out, + d_in, d_out, method): + w32 = sys.platform == "win32" and platform.architecture()[0] == "32bit" + if w32 and method == "NumericalInversePolynomial": + pytest.xfail("NumericalInversePolynomial.qrvs fails for Win " + "32-bit") + + dist = StandardNormal() + Method = getattr(stats.sampling, method) + gen = Method(dist) + + # If d and qrng.d are inconsistent, an error is raised + if d_in is not None and qrng is not None and qrng.d != d_in: + match = "`d` must be consistent with dimension of `qmc_engine`." + with pytest.raises(ValueError, match=match): + gen.qrvs(size_in, d=d_in, qmc_engine=qrng) + return + + # Sometimes d is really determined by qrng + if d_in is None and qrng is not None and qrng.d != 1: + d_out = (qrng.d,) + + shape_expected = size_out + d_out + + qrng2 = deepcopy(qrng) + qrvs = gen.qrvs(size=size_in, d=d_in, qmc_engine=qrng) + if size_in is not None: + assert qrvs.shape == shape_expected + + if qrng2 is not None: + uniform = qrng2.random(np.prod(size_in) or 1) + qrvs2 = stats.norm.ppf(uniform).reshape(shape_expected) + assert_allclose(qrvs, qrvs2, atol=1e-12) + + def test_QRVS_size_tuple(self, method): + # QMCEngine samples are always of shape (n, d). When `size` is a tuple, + # we set `n = prod(size)` in the call to qmc_engine.random, transform + # the sample, and reshape it to the final dimensions. When we reshape, + # we need to be careful, because the _columns_ of the sample returned + # by a QMCEngine are "independent"-ish, but the elements within the + # columns are not. We need to make sure that this doesn't get mixed up + # by reshaping: qrvs[..., i] should remain "independent"-ish of + # qrvs[..., i+1], but the elements within qrvs[..., i] should be + # transformed from the same low-discrepancy sequence. + + dist = StandardNormal() + Method = getattr(stats.sampling, method) + gen = Method(dist) + + size = (3, 4) + d = 5 + qrng = stats.qmc.Halton(d, seed=0) + qrng2 = stats.qmc.Halton(d, seed=0) + + uniform = qrng2.random(np.prod(size)) + + qrvs = gen.qrvs(size=size, d=d, qmc_engine=qrng) + qrvs2 = stats.norm.ppf(uniform) + + for i in range(d): + sample = qrvs[..., i] + sample2 = qrvs2[:, i].reshape(size) + assert_allclose(sample, sample2, atol=1e-12) + + +class TestTransformedDensityRejection: + # Simple Custom Distribution + class dist0: + def pdf(self, x): + return 3/4 * (1-x*x) + + def dpdf(self, x): + return 3/4 * (-2*x) + + def cdf(self, x): + return 3/4 * (x - x**3/3 + 2/3) + + def support(self): + return -1, 1 + + # Standard Normal Distribution + class dist1: + def pdf(self, x): + return stats.norm._pdf(x / 0.1) + + def dpdf(self, x): + return -x / 0.01 * stats.norm._pdf(x / 0.1) + + def cdf(self, x): + return stats.norm._cdf(x / 0.1) + + # pdf with piecewise linear function as transformed density + # with T = -1/sqrt with shift. Taken from UNU.RAN test suite + # (from file t_tdr_ps.c) + class dist2: + def __init__(self, shift): + self.shift = shift + + def pdf(self, x): + x -= self.shift + y = 1. / (abs(x) + 1.) + return 0.5 * y * y + + def dpdf(self, x): + x -= self.shift + y = 1. / (abs(x) + 1.) + y = y * y * y + return y if (x < 0.) else -y + + def cdf(self, x): + x -= self.shift + if x <= 0.: + return 0.5 / (1. - x) + else: + return 1. - 0.5 / (1. + x) + + dists = [dist0(), dist1(), dist2(0.), dist2(10000.)] + + # exact mean and variance of the distributions in the list dists + mv0 = [0., 4./15.] + mv1 = [0., 0.01] + mv2 = [0., np.inf] + mv3 = [10000., np.inf] + mvs = [mv0, mv1, mv2, mv3] + + @pytest.mark.parametrize("dist, mv_ex", + list(zip(dists, mvs))) + @pytest.mark.thread_unsafe(reason="deadlocks for unknown reasons") + def test_basic(self, dist, mv_ex): + with warnings.catch_warnings(): + # filter the warnings thrown by UNU.RAN + warnings.simplefilter("ignore", RuntimeWarning) + rng = TransformedDensityRejection(dist, random_state=42) + check_cont_samples(rng, dist, mv_ex) + + # PDF 0 everywhere => bad construction points + bad_pdfs = [(lambda x: 0, UNURANError, r"50 : bad construction points.")] + bad_pdfs += bad_pdfs_common # type: ignore[arg-type] + + @pytest.mark.parametrize("pdf, err, msg", bad_pdfs) + def test_bad_pdf(self, pdf, err, msg): + class dist: + pass + dist.pdf = pdf + dist.dpdf = lambda x: 1 # an arbitrary dPDF + with pytest.raises(err, match=msg): + TransformedDensityRejection(dist) + + @pytest.mark.parametrize("dpdf, err, msg", bad_dpdf_common) + def test_bad_dpdf(self, dpdf, err, msg): + class dist: + pass + dist.pdf = lambda x: x + dist.dpdf = dpdf + with pytest.raises(err, match=msg): + TransformedDensityRejection(dist, domain=(1, 10)) + + # test domains with inf + nan in them. need to write a custom test for + # this because not all methods support infinite tails. + @pytest.mark.parametrize("domain, err, msg", inf_nan_domains) + def test_inf_nan_domains(self, domain, err, msg): + with pytest.raises(err, match=msg): + TransformedDensityRejection(StandardNormal(), domain=domain) + + @pytest.mark.parametrize("construction_points", [-1, 0, 0.1]) + def test_bad_construction_points_scalar(self, construction_points): + with pytest.raises(ValueError, match=r"`construction_points` must be " + r"a positive integer."): + TransformedDensityRejection( + StandardNormal(), construction_points=construction_points + ) + + def test_bad_construction_points_array(self): + # empty array + construction_points = [] + with pytest.raises(ValueError, match=r"`construction_points` must " + r"either be a " + r"scalar or a non-empty array."): + TransformedDensityRejection( + StandardNormal(), construction_points=construction_points + ) + + # construction_points not monotonically increasing + construction_points = [1, 1, 1, 1, 1, 1] + with pytest.warns(RuntimeWarning, match=r"33 : starting points not " + r"strictly monotonically " + r"increasing"): + TransformedDensityRejection( + StandardNormal(), construction_points=construction_points + ) + + # construction_points containing nans + construction_points = [np.nan, np.nan, np.nan] + with pytest.raises(UNURANError, match=r"50 : bad construction " + r"points."): + TransformedDensityRejection( + StandardNormal(), construction_points=construction_points + ) + + # construction_points out of domain + construction_points = [-10, 10] + with pytest.warns(RuntimeWarning, match=r"50 : starting point out of " + r"domain"): + TransformedDensityRejection( + StandardNormal(), domain=(-3, 3), + construction_points=construction_points + ) + + @pytest.mark.parametrize("c", [-1., np.nan, np.inf, 0.1, 1.]) + def test_bad_c(self, c): + msg = r"`c` must either be -0.5 or 0." + with pytest.raises(ValueError, match=msg): + TransformedDensityRejection(StandardNormal(), c=-1.) + + u = [np.linspace(0, 1, num=1000), [], [[]], [np.nan], + [-np.inf, np.nan, np.inf], 0, + [[np.nan, 0.5, 0.1], [0.2, 0.4, np.inf], [-2, 3, 4]]] + + @pytest.mark.parametrize("u", u) + def test_ppf_hat(self, u): + # Increase the `max_squeeze_hat_ratio` so the ppf_hat is more + # accurate. + rng = TransformedDensityRejection(StandardNormal(), + max_squeeze_hat_ratio=0.9999) + # Older versions of NumPy throw RuntimeWarnings for comparisons + # with nan. + with warnings.catch_warnings(): + msg = "invalid value encountered in " + warnings.filterwarnings("ignore", msg + "greater", RuntimeWarning) + warnings.filterwarnings("ignore", msg + "greater_equal", RuntimeWarning) + warnings.filterwarnings("ignore", msg + "less", RuntimeWarning) + warnings.filterwarnings("ignore", msg + "less_equal", RuntimeWarning) + res = rng.ppf_hat(u) + expected = stats.norm.ppf(u) + assert_allclose(res, expected, rtol=1e-3, atol=1e-5) + assert res.shape == expected.shape + + def test_bad_dist(self): + # Empty distribution + class dist: + ... + + msg = r"`pdf` required but not found." + with pytest.raises(ValueError, match=msg): + TransformedDensityRejection(dist) + + # dPDF not present in dist + class dist: + pdf = lambda x: 1-x*x # noqa: E731 + + msg = r"`dpdf` required but not found." + with pytest.raises(ValueError, match=msg): + TransformedDensityRejection(dist) + + +class TestDiscreteAliasUrn: + # DAU fails on these probably because of large domains and small + # computation errors in PMF. Mean/SD match but chi-squared test fails. + basic_fail_dists = { + 'nchypergeom_fisher', # numerical errors on tails + 'nchypergeom_wallenius', # numerical errors on tails + 'randint' # fails on 32-bit ubuntu + } + + @pytest.mark.parametrize("distname, params", distdiscrete) + def test_basic(self, distname, params): + if distname in self.basic_fail_dists: + msg = ("DAU fails on these probably because of large domains " + "and small computation errors in PMF.") + pytest.skip(msg) + if not isinstance(distname, str): + dist = distname + else: + dist = getattr(stats, distname) + dist = dist(*params) + domain = dist.support() + if not np.isfinite(domain[1] - domain[0]): + # DAU only works with finite domain. So, skip the distributions + # with infinite tails. + pytest.skip("DAU only works with a finite domain.") + k = np.arange(domain[0], domain[1]+1) + pv = dist.pmf(k) + mv_ex = dist.stats('mv') + rng = DiscreteAliasUrn(dist, random_state=42) + check_discr_samples(rng, pv, mv_ex) + + # Can't use bad_pmf_common here as we evaluate PMF early on to avoid + # unhelpful errors from UNU.RAN. + bad_pmf = [ + # inf returned + (lambda x: np.inf, ValueError, + r"must contain only finite / non-nan values"), + # nan returned + (lambda x: np.nan, ValueError, + r"must contain only finite / non-nan values"), + # all zeros + (lambda x: 0.0, ValueError, + r"must contain at least one non-zero value"), + # Undefined name inside the function + (lambda x: foo, NameError, # type: ignore[name-defined] # noqa: F821 + r"name 'foo' is not defined"), + # Returning wrong type. + (lambda x: [], ValueError, + r"setting an array element with a sequence."), + # probabilities < 0 + (lambda x: -x, UNURANError, + r"50 : probability < 0"), + # signature of PMF wrong + (lambda: 1.0, TypeError, + r"takes 0 positional arguments but 1 was given") + ] + + @pytest.mark.parametrize("pmf, err, msg", bad_pmf) + def test_bad_pmf(self, pmf, err, msg): + class dist: + pass + dist.pmf = pmf + with pytest.raises(err, match=msg): + DiscreteAliasUrn(dist, domain=(1, 10)) + + @pytest.mark.parametrize("pv", [[0.18, 0.02, 0.8], + [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]]) + def test_sampling_with_pv(self, pv): + pv = np.asarray(pv, dtype=np.float64) + rng = DiscreteAliasUrn(pv, random_state=123) + rng.rvs(100_000) + pv = pv / pv.sum() + variates = np.arange(0, len(pv)) + # test if the first few moments match + m_expected = np.average(variates, weights=pv) + v_expected = np.average((variates - m_expected) ** 2, weights=pv) + mv_expected = m_expected, v_expected + check_discr_samples(rng, pv, mv_expected) + + @pytest.mark.parametrize("pv, msg", bad_pv_common) + def test_bad_pv(self, pv, msg): + with pytest.raises(ValueError, match=msg): + DiscreteAliasUrn(pv) + + # DAU doesn't support infinite tails. So, it should throw an error when + # inf is present in the domain. + inf_domain = [(-np.inf, np.inf), (np.inf, np.inf), (-np.inf, -np.inf), + (0, np.inf), (-np.inf, 0)] + + @pytest.mark.parametrize("domain", inf_domain) + def test_inf_domain(self, domain): + with pytest.raises(ValueError, match=r"must be finite"): + DiscreteAliasUrn(stats.binom(10, 0.2), domain=domain) + + def test_bad_urn_factor(self): + with pytest.warns(RuntimeWarning, match=r"relative urn size < 1."): + DiscreteAliasUrn([0.5, 0.5], urn_factor=-1) + + def test_bad_args(self): + msg = (r"`domain` must be provided when the " + r"probability vector is not available.") + + class dist: + def pmf(self, x): + return x + + with pytest.raises(ValueError, match=msg): + DiscreteAliasUrn(dist) + + def test_gh19359(self): + pv = special.softmax(np.ones((1533,))) + rng = DiscreteAliasUrn(pv, random_state=42) + # check the correctness + check_discr_samples(rng, pv, (1532 / 2, (1532**2 - 1) / 12), + rtol=5e-3) + + +class TestNumericalInversePolynomial: + # Simple Custom Distribution + class dist0: + def pdf(self, x): + return 3/4 * (1-x*x) + + def cdf(self, x): + return 3/4 * (x - x**3/3 + 2/3) + + def support(self): + return -1, 1 + + # Standard Normal Distribution + class dist1: + def pdf(self, x): + return stats.norm._pdf(x / 0.1) + + def cdf(self, x): + return stats.norm._cdf(x / 0.1) + + # Sin 2 distribution + # / 0.05 + 0.45*(1 +sin(2 Pi x)) if |x| <= 1 + # f(x) = < + # \ 0 otherwise + # Taken from UNU.RAN test suite (from file t_pinv.c) + class dist2: + def pdf(self, x): + return 0.05 + 0.45 * (1 + np.sin(2*np.pi*x)) + + def cdf(self, x): + return (0.05*(x + 1) + + 0.9*(1. + 2.*np.pi*(1 + x) - np.cos(2.*np.pi*x)) / + (4.*np.pi)) + + def support(self): + return -1, 1 + + # Sin 10 distribution + # / 0.05 + 0.45*(1 +sin(2 Pi x)) if |x| <= 5 + # f(x) = < + # \ 0 otherwise + # Taken from UNU.RAN test suite (from file t_pinv.c) + class dist3: + def pdf(self, x): + return 0.2 * (0.05 + 0.45 * (1 + np.sin(2*np.pi*x))) + + def cdf(self, x): + return x/10. + 0.5 + 0.09/(2*np.pi) * (np.cos(10*np.pi) - + np.cos(2*np.pi*x)) + + def support(self): + return -5, 5 + + dists = [dist0(), dist1(), dist2(), dist3()] + + # exact mean and variance of the distributions in the list dists + mv0 = [0., 4./15.] + mv1 = [0., 0.01] + mv2 = [-0.45/np.pi, 2/3*0.5 - 0.45**2/np.pi**2] + mv3 = [-0.45/np.pi, 0.2 * 250/3 * 0.5 - 0.45**2/np.pi**2] + mvs = [mv0, mv1, mv2, mv3] + + @pytest.mark.thread_unsafe(reason="deadlocks for unknown reasons") + @pytest.mark.parametrize("dist, mv_ex", + list(zip(dists, mvs))) + def test_basic(self, dist, mv_ex): + rng = NumericalInversePolynomial(dist, random_state=42) + check_cont_samples(rng, dist, mv_ex) + + @pytest.mark.xslow + @pytest.mark.parametrize("distname, params", distcont) + def test_basic_all_scipy_dists(self, distname, params): + + very_slow_dists = ['anglit', 'gausshyper', 'kappa4', + 'ksone', 'kstwo', 'levy_l', + 'levy_stable', 'studentized_range', + 'trapezoid', 'triang', 'vonmises'] + # for these distributions, some assertions fail due to minor + # numerical differences. They can be avoided either by changing + # the seed or by increasing the u_resolution. + fail_dists = ['chi2', 'fatiguelife', 'gibrat', + 'halfgennorm', 'lognorm', 'ncf', + 'ncx2', 'pareto', 't'] + # for these distributions, skip the check for agreement between sample + # moments and true moments. We cannot expect them to pass due to the + # high variance of sample moments. + skip_sample_moment_check = ['rel_breitwigner'] + + if distname in very_slow_dists: + pytest.skip(f"PINV too slow for {distname}") + if distname in fail_dists: + pytest.skip(f"PINV fails for {distname}") + dist = (getattr(stats, distname) + if isinstance(distname, str) + else distname) + dist = dist(*params) + with warnings.catch_warnings(): + warnings.simplefilter("ignore", RuntimeWarning) + rng = NumericalInversePolynomial(dist, random_state=42) + if distname in skip_sample_moment_check: + return + check_cont_samples(rng, dist, [dist.mean(), dist.var()]) + + @pytest.mark.parametrize("pdf, err, msg", bad_pdfs_common) + def test_bad_pdf(self, pdf, err, msg): + class dist: + pass + dist.pdf = pdf + with pytest.raises(err, match=msg): + NumericalInversePolynomial(dist, domain=[0, 5]) + + @pytest.mark.parametrize("logpdf, err, msg", bad_logpdfs_common) + def test_bad_logpdf(self, logpdf, err, msg): + class dist: + pass + dist.logpdf = logpdf + with pytest.raises(err, match=msg): + NumericalInversePolynomial(dist, domain=[0, 5]) + + # test domains with inf + nan in them. need to write a custom test for + # this because not all methods support infinite tails. + @pytest.mark.parametrize("domain, err, msg", inf_nan_domains) + def test_inf_nan_domains(self, domain, err, msg): + with pytest.raises(err, match=msg): + NumericalInversePolynomial(StandardNormal(), domain=domain) + + u = [ + # test if quantile 0 and 1 return -inf and inf respectively and check + # the correctness of the PPF for equidistant points between 0 and 1. + np.linspace(0, 1, num=10000), + # test the PPF method for empty arrays + [], [[]], + # test if nans and infs return nan result. + [np.nan], [-np.inf, np.nan, np.inf], + # test if a scalar is returned for a scalar input. + 0, + # test for arrays with nans, values greater than 1 and less than 0, + # and some valid values. + [[np.nan, 0.5, 0.1], [0.2, 0.4, np.inf], [-2, 3, 4]] + ] + + @pytest.mark.parametrize("u", u) + def test_ppf(self, u): + dist = StandardNormal() + rng = NumericalInversePolynomial(dist, u_resolution=1e-14) + # Older versions of NumPy throw RuntimeWarnings for comparisons + # with nan. + with warnings.catch_warnings(): + msg = "invalid value encountered in " + warnings.filterwarnings("ignore", msg + "greater", RuntimeWarning) + warnings.filterwarnings("ignore", msg + "greater_equal", RuntimeWarning) + warnings.filterwarnings("ignore", msg + "less", RuntimeWarning) + warnings.filterwarnings("ignore", msg + "less_equal", RuntimeWarning) + res = rng.ppf(u) + expected = stats.norm.ppf(u) + assert_allclose(res, expected, rtol=1e-11, atol=1e-11) + assert res.shape == expected.shape + + x = [np.linspace(-10, 10, num=10000), [], [[]], [np.nan], + [-np.inf, np.nan, np.inf], 0, + [[np.nan, 0.5, 0.1], [0.2, 0.4, np.inf], [-np.inf, 3, 4]]] + + @pytest.mark.parametrize("x", x) + def test_cdf(self, x): + dist = StandardNormal() + rng = NumericalInversePolynomial(dist, u_resolution=1e-14) + # Older versions of NumPy throw RuntimeWarnings for comparisons + # with nan. + with warnings.catch_warnings(): + msg = "invalid value encountered in " + warnings.filterwarnings("ignore", msg + "greater", RuntimeWarning) + warnings.filterwarnings("ignore", msg + "greater_equal", RuntimeWarning) + warnings.filterwarnings("ignore", msg + "less", RuntimeWarning) + warnings.filterwarnings("ignore", msg + "less_equal", RuntimeWarning) + res = rng.cdf(x) + expected = stats.norm.cdf(x) + assert_allclose(res, expected, rtol=1e-11, atol=1e-11) + assert res.shape == expected.shape + + @pytest.mark.slow + def test_u_error(self): + dist = StandardNormal() + rng = NumericalInversePolynomial(dist, u_resolution=1e-10) + max_error, mae = rng.u_error() + assert max_error < 1e-10 + assert mae <= max_error + rng = NumericalInversePolynomial(dist, u_resolution=1e-14) + max_error, mae = rng.u_error() + assert max_error < 1e-14 + assert mae <= max_error + + bad_orders = [1, 4.5, 20, np.inf, np.nan] + bad_u_resolution = [1e-20, 1e-1, np.inf, np.nan] + + @pytest.mark.parametrize("order", bad_orders) + def test_bad_orders(self, order): + dist = StandardNormal() + + msg = r"`order` must be an integer in the range \[3, 17\]." + with pytest.raises(ValueError, match=msg): + NumericalInversePolynomial(dist, order=order) + + @pytest.mark.parametrize("u_resolution", bad_u_resolution) + def test_bad_u_resolution(self, u_resolution): + msg = r"`u_resolution` must be between 1e-15 and 1e-5." + with pytest.raises(ValueError, match=msg): + NumericalInversePolynomial(StandardNormal(), + u_resolution=u_resolution) + + def test_bad_args(self): + + class BadDist: + def cdf(self, x): + return stats.norm._cdf(x) + + dist = BadDist() + msg = r"Either of the methods `pdf` or `logpdf` must be specified" + with pytest.raises(ValueError, match=msg): + rng = NumericalInversePolynomial(dist) + + dist = StandardNormal() + rng = NumericalInversePolynomial(dist) + msg = r"`sample_size` must be greater than or equal to 1000." + with pytest.raises(ValueError, match=msg): + rng.u_error(10) + + class Distribution: + def pdf(self, x): + return np.exp(-0.5 * x*x) + + dist = Distribution() + rng = NumericalInversePolynomial(dist) + msg = r"Exact CDF required but not found." + with pytest.raises(ValueError, match=msg): + rng.u_error() + + def test_logpdf_pdf_consistency(self): + # 1. check that PINV works with pdf and logpdf only + # 2. check that generated ppf is the same (up to a small tolerance) + + class MyDist: + pass + + # create generator from dist with only pdf + dist_pdf = MyDist() + dist_pdf.pdf = lambda x: math.exp(-x*x/2) + rng1 = NumericalInversePolynomial(dist_pdf) + + # create dist with only logpdf + dist_logpdf = MyDist() + dist_logpdf.logpdf = lambda x: -x*x/2 + rng2 = NumericalInversePolynomial(dist_logpdf) + + q = np.linspace(1e-5, 1-1e-5, num=100) + assert_allclose(rng1.ppf(q), rng2.ppf(q)) + + +class TestNumericalInverseHermite: + # / (1 +sin(2 Pi x))/2 if |x| <= 1 + # f(x) = < + # \ 0 otherwise + # Taken from UNU.RAN test suite (from file t_hinv.c) + class dist0: + def pdf(self, x): + return 0.5*(1. + np.sin(2.*np.pi*x)) + + def dpdf(self, x): + return np.pi*np.cos(2.*np.pi*x) + + def cdf(self, x): + return (1. + 2.*np.pi*(1 + x) - np.cos(2.*np.pi*x)) / (4.*np.pi) + + def support(self): + return -1, 1 + + # / Max(sin(2 Pi x)),0)Pi/2 if -1 < x <0.5 + # f(x) = < + # \ 0 otherwise + # Taken from UNU.RAN test suite (from file t_hinv.c) + class dist1: + def pdf(self, x): + if (x <= -0.5): + return np.sin((2. * np.pi) * x) * 0.5 * np.pi + if (x < 0.): + return 0. + if (x <= 0.5): + return np.sin((2. * np.pi) * x) * 0.5 * np.pi + + def dpdf(self, x): + if (x <= -0.5): + return np.cos((2. * np.pi) * x) * np.pi * np.pi + if (x < 0.): + return 0. + if (x <= 0.5): + return np.cos((2. * np.pi) * x) * np.pi * np.pi + + def cdf(self, x): + if (x <= -0.5): + return 0.25 * (1 - np.cos((2. * np.pi) * x)) + if (x < 0.): + return 0.5 + if (x <= 0.5): + return 0.75 - 0.25 * np.cos((2. * np.pi) * x) + + def support(self): + return -1, 0.5 + + dists = [dist0(), dist1()] + + # exact mean and variance of the distributions in the list dists + mv0 = [-1/(2*np.pi), 1/3 - 1/(4*np.pi*np.pi)] + mv1 = [-1/4, 3/8-1/(2*np.pi*np.pi) - 1/16] + mvs = [mv0, mv1] + + @pytest.mark.parametrize("dist, mv_ex", + list(zip(dists, mvs))) + @pytest.mark.parametrize("order", [3, 5]) + @pytest.mark.thread_unsafe + def test_basic(self, dist, mv_ex, order): + rng = NumericalInverseHermite(dist, order=order, random_state=42) + check_cont_samples(rng, dist, mv_ex) + + # test domains with inf + nan in them. need to write a custom test for + # this because not all methods support infinite tails. + @pytest.mark.parametrize("domain, err, msg", inf_nan_domains) + def test_inf_nan_domains(self, domain, err, msg): + with pytest.raises(err, match=msg): + NumericalInverseHermite(StandardNormal(), domain=domain) + + def basic_test_all_scipy_dists(self, distname, shapes): + slow_dists = {'ksone', 'kstwo', 'levy_stable', 'skewnorm'} + fail_dists = {'beta', 'gausshyper', 'geninvgauss', 'ncf', 'nct', + 'norminvgauss', 'genhyperbolic', 'studentized_range', + 'vonmises', 'kappa4', 'invgauss', 'wald'} + + if distname in slow_dists: + pytest.skip("Distribution is too slow") + if distname in fail_dists: + # specific reasons documented in gh-13319 + # https://github.com/scipy/scipy/pull/13319#discussion_r626188955 + pytest.xfail("Fails - usually due to inaccurate CDF/PDF") + + rng = np.random.default_rng(0) + + dist = getattr(stats, distname)(*shapes) + fni = NumericalInverseHermite(dist) + + x = rng.random(10) + p_tol = np.max(np.abs(dist.ppf(x)-fni.ppf(x))/np.abs(dist.ppf(x))) + u_tol = np.max(np.abs(dist.cdf(fni.ppf(x)) - x)) + + assert p_tol < 1e-8 + assert u_tol < 1e-12 + + @pytest.mark.filterwarnings('ignore::RuntimeWarning') + @pytest.mark.xslow + @pytest.mark.parametrize(("distname", "shapes"), distcont) + def test_basic_all_scipy_dists(self, distname, shapes): + # if distname == "truncnorm": + # pytest.skip("Tested separately") + self.basic_test_all_scipy_dists(distname, shapes) + + @pytest.mark.fail_slow(5) + @pytest.mark.filterwarnings('ignore::RuntimeWarning') + @pytest.mark.parallel_threads_limit(4) # Very slow + def test_basic_truncnorm_gh17155(self): + self.basic_test_all_scipy_dists("truncnorm", (0.1, 2)) + + def test_input_validation(self): + match = r"`order` must be either 1, 3, or 5." + with pytest.raises(ValueError, match=match): + NumericalInverseHermite(StandardNormal(), order=2) + + match = "`cdf` required but not found" + with pytest.raises(ValueError, match=match): + NumericalInverseHermite("norm") + + match = "could not convert string to float" + with pytest.raises(ValueError, match=match): + NumericalInverseHermite(StandardNormal(), + u_resolution='ekki') + + rngs = [None, 0, np.random.RandomState(0)] + rngs.append(np.random.default_rng(0)) # type: ignore + sizes = [(None, tuple()), (8, (8,)), ((4, 5, 6), (4, 5, 6))] + + @pytest.mark.parametrize('rng', rngs) + @pytest.mark.parametrize('size_in, size_out', sizes) + @pytest.mark.thread_unsafe + def test_RVS(self, rng, size_in, size_out): + dist = StandardNormal() + fni = NumericalInverseHermite(dist) + + rng2 = deepcopy(rng) + rvs = fni.rvs(size=size_in, random_state=rng) + if size_in is not None: + assert rvs.shape == size_out + + if rng2 is not None: + rng2 = check_random_state(rng2) + uniform = rng2.uniform(size=size_in) + rvs2 = stats.norm.ppf(uniform) + assert_allclose(rvs, rvs2) + + def test_inaccurate_CDF(self): + # CDF function with inaccurate tail cannot be inverted; see gh-13319 + # https://github.com/scipy/scipy/pull/13319#discussion_r626188955 + shapes = (2.3098496451481823, 0.6268795430096368) + match = ("98 : one or more intervals very short; possibly due to " + "numerical problems with a pole or very flat tail") + + # fails with default tol + with pytest.warns(RuntimeWarning, match=match): + NumericalInverseHermite(stats.beta(*shapes)) + + # no error with coarser tol + NumericalInverseHermite(stats.beta(*shapes), u_resolution=1e-8) + + def test_custom_distribution(self): + dist1 = StandardNormal() + fni1 = NumericalInverseHermite(dist1) + + dist2 = stats.norm() + fni2 = NumericalInverseHermite(dist2) + + assert_allclose(fni1.rvs(random_state=0), fni2.rvs(random_state=0)) + + u = [ + # check the correctness of the PPF for equidistant points between + # 0.02 and 0.98. + np.linspace(0., 1., num=10000), + # test the PPF method for empty arrays + [], [[]], + # test if nans and infs return nan result. + [np.nan], [-np.inf, np.nan, np.inf], + # test if a scalar is returned for a scalar input. + 0, + # test for arrays with nans, values greater than 1 and less than 0, + # and some valid values. + [[np.nan, 0.5, 0.1], [0.2, 0.4, np.inf], [-2, 3, 4]] + ] + + @pytest.mark.parametrize("u", u) + def test_ppf(self, u): + dist = StandardNormal() + rng = NumericalInverseHermite(dist, u_resolution=1e-12) + # Older versions of NumPy throw RuntimeWarnings for comparisons + # with nan. + with warnings.catch_warnings(): + msg = "invalid value encountered in " + warnings.filterwarnings("ignore", msg + "greater", RuntimeWarning) + warnings.filterwarnings("ignore", msg + "greater_equal", RuntimeWarning) + warnings.filterwarnings("ignore", msg + "less", RuntimeWarning) + warnings.filterwarnings("ignore", msg + "less_equal", RuntimeWarning) + res = rng.ppf(u) + expected = stats.norm.ppf(u) + assert_allclose(res, expected, rtol=1e-9, atol=3e-10) + assert res.shape == expected.shape + + @pytest.mark.slow + def test_u_error(self): + dist = StandardNormal() + rng = NumericalInverseHermite(dist, u_resolution=1e-10) + max_error, mae = rng.u_error() + assert max_error < 1e-10 + assert mae <= max_error + with warnings.catch_warnings(): + # ignore warning about u-resolution being too small. + warnings.simplefilter("ignore", RuntimeWarning) + rng = NumericalInverseHermite(dist, u_resolution=1e-14) + max_error, mae = rng.u_error() + assert max_error < 1e-14 + assert mae <= max_error + + +class TestDiscreteGuideTable: + basic_fail_dists = { + 'nchypergeom_fisher', # numerical errors on tails + 'nchypergeom_wallenius', # numerical errors on tails + 'randint' # fails on 32-bit ubuntu + } + + def test_guide_factor_gt3_raises_warning(self): + pv = [0.1, 0.3, 0.6] + urng = np.random.default_rng() + with pytest.warns(RuntimeWarning): + DiscreteGuideTable(pv, random_state=urng, guide_factor=7) + + def test_guide_factor_zero_raises_warning(self): + pv = [0.1, 0.3, 0.6] + urng = np.random.default_rng() + with pytest.warns(RuntimeWarning): + DiscreteGuideTable(pv, random_state=urng, guide_factor=0) + + def test_negative_guide_factor_raises_warning(self): + # This occurs from the UNU.RAN wrapper automatically. + # however it already gives a useful warning + # Here we just test that a warning is raised. + pv = [0.1, 0.3, 0.6] + urng = np.random.default_rng() + with pytest.warns(RuntimeWarning): + DiscreteGuideTable(pv, random_state=urng, guide_factor=-1) + + @pytest.mark.parametrize("distname, params", distdiscrete) + def test_basic(self, distname, params): + if distname in self.basic_fail_dists: + msg = ("DGT fails on these probably because of large domains " + "and small computation errors in PMF.") + pytest.skip(msg) + + if not isinstance(distname, str): + dist = distname + else: + dist = getattr(stats, distname) + + dist = dist(*params) + domain = dist.support() + + if not np.isfinite(domain[1] - domain[0]): + # DGT only works with finite domain. So, skip the distributions + # with infinite tails. + pytest.skip("DGT only works with a finite domain.") + + k = np.arange(domain[0], domain[1]+1) + pv = dist.pmf(k) + mv_ex = dist.stats('mv') + rng = DiscreteGuideTable(dist, random_state=42) + check_discr_samples(rng, pv, mv_ex) + + u = [ + # the correctness of the PPF for equidistant points between 0 and 1. + np.linspace(0, 1, num=10000), + # test the PPF method for empty arrays + [], [[]], + # test if nans and infs return nan result. + [np.nan], [-np.inf, np.nan, np.inf], + # test if a scalar is returned for a scalar input. + 0, + # test for arrays with nans, values greater than 1 and less than 0, + # and some valid values. + [[np.nan, 0.5, 0.1], [0.2, 0.4, np.inf], [-2, 3, 4]] + ] + + @pytest.mark.parametrize('u', u) + def test_ppf(self, u): + n, p = 4, 0.1 + dist = stats.binom(n, p) + rng = DiscreteGuideTable(dist, random_state=42) + + # Older versions of NumPy throw RuntimeWarnings for comparisons + # with nan. + with warnings.catch_warnings(): + msg = "invalid value encountered in " + warnings.filterwarnings("ignore", msg + "greater", RuntimeWarning) + warnings.filterwarnings("ignore", msg + "greater_equal", RuntimeWarning) + warnings.filterwarnings("ignore", msg + "less", RuntimeWarning) + warnings.filterwarnings("ignore", msg + "less_equal", RuntimeWarning) + res = rng.ppf(u) + expected = stats.binom.ppf(u, n, p) + assert_equal(res.shape, expected.shape) + assert_equal(res, expected) + + @pytest.mark.parametrize("pv, msg", bad_pv_common) + def test_bad_pv(self, pv, msg): + with pytest.raises(ValueError, match=msg): + DiscreteGuideTable(pv) + + # DGT doesn't support infinite tails. So, it should throw an error when + # inf is present in the domain. + inf_domain = [(-np.inf, np.inf), (np.inf, np.inf), (-np.inf, -np.inf), + (0, np.inf), (-np.inf, 0)] + + @pytest.mark.parametrize("domain", inf_domain) + def test_inf_domain(self, domain): + with pytest.raises(ValueError, match=r"must be finite"): + DiscreteGuideTable(stats.binom(10, 0.2), domain=domain) + + +class TestSimpleRatioUniforms: + # pdf with piecewise linear function as transformed density + # with T = -1/sqrt with shift. Taken from UNU.RAN test suite + # (from file t_srou.c) + class dist: + def __init__(self, shift): + self.shift = shift + self.mode = shift + + def pdf(self, x): + x -= self.shift + y = 1. / (abs(x) + 1.) + return 0.5 * y * y + + def cdf(self, x): + x -= self.shift + if x <= 0.: + return 0.5 / (1. - x) + else: + return 1. - 0.5 / (1. + x) + + dists = [dist(0.), dist(10000.)] + + # exact mean and variance of the distributions in the list dists + mv1 = [0., np.inf] + mv2 = [10000., np.inf] + mvs = [mv1, mv2] + + @pytest.mark.parametrize("dist, mv_ex", + list(zip(dists, mvs))) + @pytest.mark.thread_unsafe + def test_basic(self, dist, mv_ex): + rng = SimpleRatioUniforms(dist, mode=dist.mode, random_state=42) + check_cont_samples(rng, dist, mv_ex) + rng = SimpleRatioUniforms(dist, mode=dist.mode, + cdf_at_mode=dist.cdf(dist.mode), + random_state=42) + check_cont_samples(rng, dist, mv_ex) + + # test domains with inf + nan in them. need to write a custom test for + # this because not all methods support infinite tails. + @pytest.mark.parametrize("domain, err, msg", inf_nan_domains) + def test_inf_nan_domains(self, domain, err, msg): + with pytest.raises(err, match=msg): + SimpleRatioUniforms(StandardNormal(), domain=domain) + + def test_bad_args(self): + # pdf_area < 0 + with pytest.raises(ValueError, match=r"`pdf_area` must be > 0"): + SimpleRatioUniforms(StandardNormal(), mode=0, pdf_area=-1) + + +class TestRatioUniforms: + def test_rv_generation(self): + # use KS test to check distribution of rvs + # normal distribution + f = stats.norm.pdf + v = np.sqrt(f(np.sqrt(2))) * np.sqrt(2) + u = np.sqrt(f(0)) + gen = RatioUniforms(f, umax=u, vmin=-v, vmax=v, random_state=12345) + assert_equal(stats.kstest(gen.rvs(2500), 'norm')[1] > 0.25, True) + + # exponential distribution + gen = RatioUniforms(lambda x: np.exp(-x), umax=1, + vmin=0, vmax=2*np.exp(-1), random_state=12345) + assert_equal(stats.kstest(gen.rvs(1000), 'expon')[1] > 0.25, True) + + def test_shape(self): + # test shape of return value depending on size parameter + f = stats.norm.pdf + v = np.sqrt(f(np.sqrt(2))) * np.sqrt(2) + u = np.sqrt(f(0)) + + gen1 = RatioUniforms(f, umax=u, vmin=-v, vmax=v, random_state=1234) + gen2 = RatioUniforms(f, umax=u, vmin=-v, vmax=v, random_state=1234) + gen3 = RatioUniforms(f, umax=u, vmin=-v, vmax=v, random_state=1234) + r1, r2, r3 = gen1.rvs(3), gen2.rvs((3,)), gen3.rvs((3, 1)) + assert_equal(r1, r2) + assert_equal(r2, r3.flatten()) + assert_equal(r1.shape, (3,)) + assert_equal(r3.shape, (3, 1)) + + gen4 = RatioUniforms(f, umax=u, vmin=-v, vmax=v, random_state=12) + gen5 = RatioUniforms(f, umax=u, vmin=-v, vmax=v, random_state=12) + r4, r5 = gen4.rvs(size=(3, 3, 3)), gen5.rvs(size=27) + assert_equal(r4.flatten(), r5) + assert_equal(r4.shape, (3, 3, 3)) + + gen6 = RatioUniforms(f, umax=u, vmin=-v, vmax=v, random_state=1234) + gen7 = RatioUniforms(f, umax=u, vmin=-v, vmax=v, random_state=1234) + gen8 = RatioUniforms(f, umax=u, vmin=-v, vmax=v, random_state=1234) + r6, r7, r8 = gen6.rvs(), gen7.rvs(1), gen8.rvs((1,)) + assert_equal(r6, r7) + assert_equal(r7, r8) + + def test_random_state(self): + f = stats.norm.pdf + v = np.sqrt(f(np.sqrt(2))) * np.sqrt(2) + umax = np.sqrt(f(0)) + gen1 = RatioUniforms(f, umax=umax, vmin=-v, vmax=v, random_state=1234) + r1 = gen1.rvs(10) + rng = np.random.RandomState(1234) + gen2 = RatioUniforms(f, umax=umax, vmin=-v, vmax=v, random_state=rng) + r2 = gen2.rvs(10) + assert_equal(r1, r2) + + def test_exceptions(self): + f = stats.norm.pdf + # need vmin < vmax + with assert_raises(ValueError, match="vmin must be smaller than vmax"): + RatioUniforms(pdf=f, umax=1, vmin=3, vmax=1) + with assert_raises(ValueError, match="vmin must be smaller than vmax"): + RatioUniforms(pdf=f, umax=1, vmin=1, vmax=1) + # need umax > 0 + with assert_raises(ValueError, match="umax must be positive"): + RatioUniforms(pdf=f, umax=-1, vmin=1, vmax=3) + with assert_raises(ValueError, match="umax must be positive"): + RatioUniforms(pdf=f, umax=0, vmin=1, vmax=3) diff --git a/v12x/pylib/scipy/stats/tests/test_sensitivity_analysis.py b/v12x/pylib/scipy/stats/tests/test_sensitivity_analysis.py new file mode 100644 index 0000000000000000000000000000000000000000..9c8d53040380808c03358b905878cc4a8b9b9679 --- /dev/null +++ b/v12x/pylib/scipy/stats/tests/test_sensitivity_analysis.py @@ -0,0 +1,310 @@ +import numpy as np +from numpy.testing import assert_allclose, assert_array_less +import pytest + +from scipy import stats +from scipy.stats import sobol_indices +from scipy.stats._resampling import BootstrapResult +from scipy.stats._sensitivity_analysis import ( + BootstrapSobolResult, f_ishigami, sample_AB, sample_A_B +) + + +@pytest.fixture(scope='session') +def ishigami_ref_indices(): + """Reference values for Ishigami from Saltelli2007. + + Chapter 4, exercise 5 pages 179-182. + """ + a = 7. + b = 0.1 + + var = 0.5 + a**2/8 + b*np.pi**4/5 + b**2*np.pi**8/18 + v1 = 0.5 + b*np.pi**4/5 + b**2*np.pi**8/50 + v2 = a**2/8 + v3 = 0 + v12 = 0 + # v13: mistake in the book, see other derivations e.g. in 10.1002/nme.4856 + v13 = b**2*np.pi**8*8/225 + v23 = 0 + + s_first = np.array([v1, v2, v3])/var + s_second = np.array([ + [0., 0., v13], + [v12, 0., v23], + [v13, v23, 0.] + ])/var + s_total = s_first + s_second.sum(axis=1) + + return s_first, s_total + + +def f_ishigami_vec(x): + """Output of shape (2, n).""" + res = f_ishigami(x) + return res, res + + +class TestSobolIndices: + + dists = [ + stats.uniform(loc=-np.pi, scale=2*np.pi) # type: ignore[attr-defined] + ] * 3 + + def test_sample_AB(self): + # (d, n) + A = np.array( + [[1, 4, 7, 10], + [2, 5, 8, 11], + [3, 6, 9, 12]] + ) + B = A + 100 + # (d, d, n) + ref = np.array( + [[[101, 104, 107, 110], + [2, 5, 8, 11], + [3, 6, 9, 12]], + [[1, 4, 7, 10], + [102, 105, 108, 111], + [3, 6, 9, 12]], + [[1, 4, 7, 10], + [2, 5, 8, 11], + [103, 106, 109, 112]]] + ) + AB = sample_AB(A=A, B=B) + assert_allclose(AB, ref) + + @pytest.mark.xslow + @pytest.mark.xfail_on_32bit("Can't create large array for test") + @pytest.mark.parametrize( + 'func', + [f_ishigami, pytest.param(f_ishigami_vec, marks=pytest.mark.slow)], + ids=['scalar', 'vector'] + ) + def test_ishigami(self, ishigami_ref_indices, func): + rng = np.random.default_rng(28631265345463262246170309650372465332) + res = sobol_indices( + func=func, n=4096, + dists=self.dists, + rng=rng + ) + + if func.__name__ == 'f_ishigami_vec': + ishigami_ref_indices = [ + [ishigami_ref_indices[0], ishigami_ref_indices[0]], + [ishigami_ref_indices[1], ishigami_ref_indices[1]] + ] + + assert_allclose(res.first_order, ishigami_ref_indices[0], atol=1e-2) + assert_allclose(res.total_order, ishigami_ref_indices[1], atol=1e-2) + + assert res._bootstrap_result is None + bootstrap_res = res.bootstrap(n_resamples=99) + assert isinstance(bootstrap_res, BootstrapSobolResult) + assert isinstance(res._bootstrap_result, BootstrapResult) + + assert res._bootstrap_result.confidence_interval.low.shape[0] == 2 + assert res._bootstrap_result.confidence_interval.low[1].shape \ + == res.first_order.shape + + assert bootstrap_res.first_order.confidence_interval.low.shape \ + == res.first_order.shape + assert bootstrap_res.total_order.confidence_interval.low.shape \ + == res.total_order.shape + + assert_array_less( + bootstrap_res.first_order.confidence_interval.low, res.first_order + ) + assert_array_less( + res.first_order, bootstrap_res.first_order.confidence_interval.high + ) + assert_array_less( + bootstrap_res.total_order.confidence_interval.low, res.total_order + ) + assert_array_less( + res.total_order, bootstrap_res.total_order.confidence_interval.high + ) + + # call again to use previous results and change a param + assert isinstance( + res.bootstrap(confidence_level=0.9, n_resamples=99), + BootstrapSobolResult + ) + assert isinstance(res._bootstrap_result, BootstrapResult) + + def test_func_dict(self, ishigami_ref_indices): + rng = np.random.default_rng(28631265345463262246170309650372465332) + n = 4096 + dists = [ + stats.uniform(loc=-np.pi, scale=2*np.pi), + stats.uniform(loc=-np.pi, scale=2*np.pi), + stats.uniform(loc=-np.pi, scale=2*np.pi) + ] + + A, B = sample_A_B(n=n, dists=dists, rng=rng) + AB = sample_AB(A=A, B=B) + + func = { + 'f_A': f_ishigami(A).reshape(1, -1), + 'f_B': f_ishigami(B).reshape(1, -1), + 'f_AB': f_ishigami(AB).reshape((3, 1, -1)) + } + + # preserve use of old random_state during SPEC 7 transition + res = sobol_indices( + func=func, n=n, + dists=dists, + rng=rng + ) + assert_allclose(res.first_order, ishigami_ref_indices[0], atol=1e-2) + + res = sobol_indices( + func=func, n=n, + rng=rng + ) + assert_allclose(res.first_order, ishigami_ref_indices[0], atol=1e-2) + # Ideally should be exactly equal but since f_ishigami + # uses floating point operations, so exact equality + # might not be possible (due to flakiness in computation). + # So, assert_allclose is used with default parameters + # Regression test for https://github.com/scipy/scipy/issues/21383 + assert_allclose(f_ishigami(A).reshape(1, -1), func['f_A']) + assert_allclose(f_ishigami(B).reshape(1, -1), func['f_B']) + assert_allclose(f_ishigami(AB).reshape((3, 1, -1)), func['f_AB']) + + def test_method(self, ishigami_ref_indices): + def jansen_sobol(f_A, f_B, f_AB): + """Jansen for S and Sobol' for St. + + From Saltelli2010, table 2 formulations (c) and (e).""" + var = np.var([f_A, f_B], axis=(0, -1)) + + s = (var - 0.5*np.mean((f_B - f_AB)**2, axis=-1)) / var + st = np.mean(f_A*(f_A - f_AB), axis=-1) / var + + return s.T, st.T + + rng = np.random.default_rng(28631265345463262246170309650372465332) + res = sobol_indices( + func=f_ishigami, n=4096, + dists=self.dists, + method=jansen_sobol, + rng=rng + ) + + assert_allclose(res.first_order, ishigami_ref_indices[0], atol=1e-2) + assert_allclose(res.total_order, ishigami_ref_indices[1], atol=1e-2) + + def jansen_sobol_typed( + f_A: np.ndarray, f_B: np.ndarray, f_AB: np.ndarray + ) -> tuple[np.ndarray, np.ndarray]: + return jansen_sobol(f_A, f_B, f_AB) + + _ = sobol_indices( + func=f_ishigami, n=8, + dists=self.dists, + method=jansen_sobol_typed, + rng=rng + ) + + def test_normalization(self, ishigami_ref_indices): + rng = np.random.default_rng(28631265345463262246170309650372465332) + res = sobol_indices( + func=lambda x: f_ishigami(x) + 1000, n=4096, + dists=self.dists, + rng=rng + ) + + assert_allclose(res.first_order, ishigami_ref_indices[0], atol=1e-2) + assert_allclose(res.total_order, ishigami_ref_indices[1], atol=1e-2) + + def test_constant_function(self, ishigami_ref_indices): + + def f_ishigami_vec_const(x): + """Output of shape (3, n).""" + res = f_ishigami(x) + return res, res * 0 + 10, res + + rng = np.random.default_rng(28631265345463262246170309650372465332) + res = sobol_indices( + func=f_ishigami_vec_const, n=4096, + dists=self.dists, + rng=rng + ) + + ishigami_vec_indices = [ + [ishigami_ref_indices[0], [0, 0, 0], ishigami_ref_indices[0]], + [ishigami_ref_indices[1], [0, 0, 0], ishigami_ref_indices[1]] + ] + + assert_allclose(res.first_order, ishigami_vec_indices[0], atol=1e-2) + assert_allclose(res.total_order, ishigami_vec_indices[1], atol=1e-2) + + @pytest.mark.xfail_on_32bit("Can't create large array for test") + def test_more_converged(self, ishigami_ref_indices): + rng = np.random.default_rng(28631265345463262246170309650372465332) + res = sobol_indices( + func=f_ishigami, n=2**19, # 524288 + dists=self.dists, + rng=rng + ) + + assert_allclose(res.first_order, ishigami_ref_indices[0], atol=1e-4) + assert_allclose(res.total_order, ishigami_ref_indices[1], atol=1e-4) + + def test_raises(self): + + message = r"Each distribution in `dists` must have method `ppf`" + with pytest.raises(ValueError, match=message): + sobol_indices(n=0, func=f_ishigami, dists="uniform") + + with pytest.raises(ValueError, match=message): + sobol_indices(n=0, func=f_ishigami, dists=[lambda x: x]) + + message = r"The balance properties of Sobol'" + with pytest.raises(ValueError, match=message): + sobol_indices(n=7, func=f_ishigami, dists=[stats.uniform()]) + + with pytest.raises(ValueError, match=message): + sobol_indices(n=4.1, func=f_ishigami, dists=[stats.uniform()]) + + message = r"'toto' is not a valid 'method'" + with pytest.raises(ValueError, match=message): + sobol_indices(n=0, func=f_ishigami, method='toto') + + message = r"must have the following signature" + with pytest.raises(ValueError, match=message): + sobol_indices(n=0, func=f_ishigami, method=lambda x: x) + + message = r"'dists' must be defined when 'func' is a callable" + with pytest.raises(ValueError, match=message): + sobol_indices(n=0, func=f_ishigami) + + def func_wrong_shape_output(x): + return x.reshape(-1, 1) + + message = r"'func' output should have a shape" + with pytest.raises(ValueError, match=message): + sobol_indices( + n=2, func=func_wrong_shape_output, dists=[stats.uniform()] + ) + + message = r"When 'func' is a dictionary" + with pytest.raises(ValueError, match=message): + sobol_indices( + n=2, func={'f_A': [], 'f_AB': []}, dists=[stats.uniform()] + ) + + with pytest.raises(ValueError, match=message): + # f_B malformed + sobol_indices( + n=2, + func={'f_A': [1, 2], 'f_B': [3], 'f_AB': [5, 6, 7, 8]}, + ) + + with pytest.raises(ValueError, match=message): + # f_AB malformed + sobol_indices( + n=2, + func={'f_A': [1, 2], 'f_B': [3, 4], 'f_AB': [5, 6, 7]}, + ) diff --git a/v12x/pylib/scipy/stats/tests/test_stats.py b/v12x/pylib/scipy/stats/tests/test_stats.py new file mode 100644 index 0000000000000000000000000000000000000000..a0846f17a0f71699ffb5b654cd2da21dd72d5f74 --- /dev/null +++ b/v12x/pylib/scipy/stats/tests/test_stats.py @@ -0,0 +1,9890 @@ +""" Test functions for stats module + + WRITTEN BY LOUIS LUANGKESORN FOR THE STATS MODULE + BASED ON WILKINSON'S STATISTICS QUIZ + https://www.stanford.edu/~clint/bench/wilk.txt + + Additional tests by a host of SciPy developers. +""" +import math +import os +import warnings +from collections import namedtuple +from itertools import product +import hypothesis.extra.numpy as npst +import hypothesis +import contextlib + +from numpy.testing import (assert_, assert_equal, + assert_almost_equal, assert_array_almost_equal, + assert_array_equal, assert_approx_equal, + assert_allclose) +import pytest +from pytest import raises as assert_raises +from numpy import array, arange, power +import numpy as np + +import scipy.stats as stats +from scipy.stats._ksstats import kolmogn +from scipy.special._testutils import FuncData +from scipy import optimize, special +from .common_tests import check_named_results +from scipy.stats._axis_nan_policy import (SmallSampleWarning, + too_small_nd_omit, too_small_nd_not_omit, + too_small_1d_omit, too_small_1d_not_omit) +from scipy.stats._stats_py import (_chk_asarray, _moment, + LinregressResult, _xp_mean, _xp_var, _SimpleChi2) +from scipy._lib._util import AxisError +from scipy.conftest import skip_xp_invalid_arg +from scipy._lib._array_api import (array_namespace, eager_warns, is_lazy_array, + is_numpy, is_torch, xp_default_dtype, xp_size, + SCIPY_ARRAY_API, make_xp_test_case, xp_ravel, + xp_swapaxes, xp_result_type, is_jax, is_cupy, + xp_copy, xp_promote, make_xp_pytest_param) +from scipy._lib._array_api_no_0d import xp_assert_close, xp_assert_equal, xp_assert_less +import scipy._external.array_api_extra as xpx +from scipy._lib._util import _apply_over_batch + + +lazy_xp_modules = [stats] +skip_xp_backends = pytest.mark.skip_xp_backends +xfail_xp_backends = pytest.mark.xfail_xp_backends + +""" Numbers in docstrings beginning with 'W' refer to the section numbers + and headings found in the STATISTICS QUIZ of Leland Wilkinson. These are + considered to be essential functionality. True testing and + evaluation of a statistics package requires use of the + NIST Statistical test data. See McCoullough(1999) Assessing The Reliability + of Statistical Software for a test methodology and its + implementation in testing SAS, SPSS, and S-Plus +""" + +# Datasets +# These data sets are from the nasty.dat sets used by Wilkinson +# For completeness, I should write the relevant tests and count them as failures +# Somewhat acceptable, since this is still beta software. It would count as a +# good target for 1.0 status +X = array([1,2,3,4,5,6,7,8,9], float) +ZERO = array([0,0,0,0,0,0,0,0,0], float) +BIG = array([99999991,99999992,99999993,99999994,99999995,99999996,99999997, + 99999998,99999999], float) +LITTLE = array([0.99999991,0.99999992,0.99999993,0.99999994,0.99999995,0.99999996, + 0.99999997,0.99999998,0.99999999], float) +HUGE = array([1e+12,2e+12,3e+12,4e+12,5e+12,6e+12,7e+12,8e+12,9e+12], float) +TINY = array([1e-12,2e-12,3e-12,4e-12,5e-12,6e-12,7e-12,8e-12,9e-12], float) +ROUND = array([0.5,1.5,2.5,3.5,4.5,5.5,6.5,7.5,8.5], float) + + +class TestTrimmedStats: + # TODO: write these tests to handle missing values properly + dprec = np.finfo(np.float64).precision + + @make_xp_test_case(stats.tmean) + def test_tmean(self, xp): + default_dtype = xp_default_dtype(xp) + x = xp.asarray(X, dtype=default_dtype) + + y = stats.tmean(x, (2, 8), (True, True)) + xp_assert_close(y, xp.asarray(5.0)) + + y1 = stats.tmean(x, limits=(2, 8), inclusive=(False, False)) + y2 = stats.tmean(x, limits=None) + xp_assert_close(y1, y2) + + x_2d = xp.reshape(xp.arange(63.), (9, 7)) + y = stats.tmean(x_2d, axis=None) + xp_assert_close(y, xp.mean(x_2d)) + + y = stats.tmean(x_2d, axis=0) + xp_assert_close(y, xp.mean(x_2d, axis=0)) + + y = stats.tmean(x_2d, axis=1) + xp_assert_close(y, xp.mean(x_2d, axis=1)) + + y = stats.tmean(x_2d, limits=(2, 61), axis=None) + xp_assert_close(y, xp.asarray(31.5)) + + y = stats.tmean(x_2d, limits=(2, 21), axis=0) + y_true = [14, 11.5, 9, 10, 11, 12, 13] + xp_assert_close(y, xp.asarray(y_true)) + + y = stats.tmean(x_2d, limits=(2, 21), inclusive=(True, False), axis=0) + y_true = [10.5, 11.5, 9, 10, 11, 12, 13] + xp_assert_close(y, xp.asarray(y_true)) + + x_2d_with_nan = xpx.at(x_2d)[-1, -3:].set(xp.nan, copy=True) + y = stats.tmean(x_2d_with_nan, limits=(1, 13), axis=0) + y_true = [7, 4.5, 5.5, 6.5, xp.nan, xp.nan, xp.nan] + xp_assert_close(y, xp.asarray(y_true)) + + y = stats.tmean(x_2d, limits=(2, 21), axis=1) + y_true = [4, 10, 17, 21, xp.nan, xp.nan, xp.nan, xp.nan, xp.nan] + xp_assert_close(y, xp.asarray(y_true)) + + y = stats.tmean(x_2d, limits=(2, 21), + inclusive=(False, True), axis=1) + y_true = [4.5, 10, 17, 21, xp.nan, xp.nan, xp.nan, xp.nan, xp.nan] + xp_assert_close(y, xp.asarray(y_true)) + + @make_xp_test_case(stats.tvar) + @pytest.mark.filterwarnings( + "ignore:invalid value encountered in divide:RuntimeWarning:dask" + ) + def test_tvar(self, xp): + x = xp.asarray(X.tolist()) # use default dtype of xp + + y = stats.tvar(x, limits=(2, 8), inclusive=(True, True)) + xp_assert_close(y, xp.asarray(4.6666666666666661)) + + y = stats.tvar(x, limits=None) + xp_assert_close(y, xp.var(x, correction=1)) + + x_2d = xp.reshape(xp.arange(63.), (9, 7)) + y = stats.tvar(x_2d, axis=None) + xp_assert_close(y, xp.var(x_2d, correction=1)) + + y = stats.tvar(x_2d, axis=0) + xp_assert_close(y, xp.full((7,), 367.5)) + + y = stats.tvar(x_2d, axis=1) + xp_assert_close(y, xp.full((9,), 4.66666667)) + + # Limiting some values along one axis + y = stats.tvar(x_2d, limits=(1, 5), axis=1, inclusive=(True, True)) + xp_assert_close(y[0], xp.asarray(2.5)) + + # Limiting all values along one axis + y = stats.tvar(x_2d, limits=(0, 6), axis=1, inclusive=(True, True)) + xp_assert_close(y[0], xp.asarray(4.666666666666667)) + xp_assert_equal(y[1], xp.asarray(xp.nan)) + + @make_xp_test_case(stats.tstd) + def test_tstd(self, xp): + x = xp.asarray(X.tolist()) # use default dtype of xp + + y = stats.tstd(x, (2, 8), (True, True)) + xp_assert_close(y, xp.asarray(2.1602468994692865)) + + y = stats.tstd(x, limits=None) + xp_assert_close(y, xp.std(x, correction=1)) + + @make_xp_test_case(stats.tmin) + def test_tmin(self, xp): + x = xp.arange(10.) + xp_assert_equal(stats.tmin(x), xp.asarray(0.)) + xp_assert_equal(stats.tmin(x, lowerlimit=0), xp.asarray(0.)) + xp_assert_equal(stats.tmin(x, lowerlimit=0, inclusive=False), xp.asarray(1.)) + + x = xp.reshape(x, (5, 2)) + xp_assert_equal(stats.tmin(x, lowerlimit=0, inclusive=False), + xp.asarray([2., 1.])) + xp_assert_equal(stats.tmin(x, axis=1), xp.asarray([0., 2., 4., 6., 8.])) + xp_assert_equal(stats.tmin(x, axis=None), xp.asarray(0.)) + + x = xpx.at(xp.arange(10.), 9).set(xp.nan) + xp_assert_equal(stats.tmin(x), xp.asarray(xp.nan)) + + # check that if a full slice is masked, the output returns a + # nan instead of a garbage value. + x = xp.reshape(xp.arange(16), (4, 4)) + res = stats.tmin(x, lowerlimit=4, axis=1) + xp_assert_equal(res, xp.asarray([np.nan, 4, 8, 12])) + + @pytest.mark.parametrize('tfun', [stats.tmin, stats.tmax]) + def test_tmin_tmax_scalar_input(self, tfun): # scalar inputs are NumPy only + assert_equal(tfun(4), 4) + + @skip_xp_backends(eager_only=True, reason="lazy -> reduced nan_policy capabilities") + @pytest.mark.parametrize('tfun, ref', [make_xp_pytest_param(stats.tmin, 0.), + make_xp_pytest_param(stats.tmax, 9.)]) + def test_tmin_tmax_nanpolicy(self, tfun, ref, xp): + x = xp.arange(10.) + x[8] = xp.nan + xp_assert_equal(tfun(x, nan_policy='omit'), xp.asarray(ref)) + msg = "The input contains nan values" + with assert_raises(ValueError, match=msg): + tfun(x, nan_policy='raise') + msg = "nan_policy must be one of..." + with assert_raises(ValueError, match=msg): + tfun(x, nan_policy='foobar') + + @make_xp_test_case(stats.tmax) + def test_tmax(self, xp): + x = xp.arange(10.) + xp_assert_equal(stats.tmax(x), xp.asarray(9.)) + xp_assert_equal(stats.tmax(x, upperlimit=9), xp.asarray(9.)) + xp_assert_equal(stats.tmax(x, upperlimit=9, inclusive=False), xp.asarray(8.)) + + x = xp.reshape(x, (5, 2)) + xp_assert_equal(stats.tmax(x, upperlimit=9, inclusive=False), + xp.asarray([8., 7.])) + xp_assert_equal(stats.tmax(x, axis=1), xp.asarray([1., 3., 5., 7., 9.])) + xp_assert_equal(stats.tmax(x, axis=None), xp.asarray(9.)) + + x = xpx.at(xp.arange(10.), 9).set(xp.nan) + xp_assert_equal(stats.tmax(x), xp.asarray(xp.nan)) + + # check that if a full slice is masked, the output returns a + # nan instead of a garbage value. + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", "All-NaN slice encountered", RuntimeWarning) + x = xp.reshape(xp.arange(16), (4, 4)) + res = stats.tmax(x, upperlimit=11, axis=1) + xp_assert_equal(res, xp.asarray([3, 7, 11, np.nan])) + + @make_xp_test_case(stats.tmin, stats.tmax) + def test_tmin_tmax_int_dtype(self, xp): + x = xp.reshape(xp.arange(10, dtype=xp.int16), (2, 5)).T + + # When tmin/tmax don't need to inject any NaNs, + # retain the input dtype. Dask/JAX can't inspect + # the data so they always return float. + expect_dtype = xp_default_dtype(xp) if is_lazy_array(x) else x.dtype + xp_assert_equal(stats.tmin(x), xp.asarray([0, 5], dtype=expect_dtype)) + xp_assert_equal(stats.tmax(x), xp.asarray([4, 9], dtype=expect_dtype)) + + # When they do inject NaNs, all backends behave the same. + xp_assert_equal(stats.tmin(x, lowerlimit=6), xp.asarray([xp.nan, 6.])) + xp_assert_equal(stats.tmax(x, upperlimit=3), xp.asarray([3., xp.nan])) + + @skip_xp_backends(eager_only=True, reason="Only with data-dependent output dtype") + @make_xp_test_case(stats.tmin, stats.tmax) + def test_gh_22626(self, xp): + # Test that `tmin`/`tmax` returns exact result with outrageously large integers + x = xp.arange(2**62, 2**62+10) + xp_assert_equal(stats.tmin(x[None, :]), x) + xp_assert_equal(stats.tmax(x[None, :]), x) + + @make_xp_test_case(stats.tsem) + def test_tsem(self, xp): + x = xp.asarray(X.tolist()) # use default dtype of xp + + y = stats.tsem(x, limits=(3, 8), inclusive=(False, True)) + y_ref = xp.asarray([4., 5., 6., 7., 8.]) + xp_assert_close(y, xp.std(y_ref, correction=1) / xp_size(y_ref)**0.5) + xp_assert_close(stats.tsem(x, limits=[-1, 10]), stats.tsem(x, limits=None)) + + +class TestPearsonrWilkinson: + """ W.II.D. Compute a correlation matrix on all the variables. + + All the correlations, except for ZERO and MISS, should be exactly 1. + ZERO and MISS should have undefined or missing correlations with the + other variables. The same should go for SPEARMAN correlations, if + your program has them. + """ + + def test_pXX(self): + y = stats.pearsonr(X,X) + r = y[0] + assert_approx_equal(r,1.0) + + def test_pXBIG(self): + y = stats.pearsonr(X,BIG) + r = y[0] + assert_approx_equal(r,1.0) + + def test_pXLITTLE(self): + y = stats.pearsonr(X,LITTLE) + r = y[0] + assert_approx_equal(r,1.0) + + def test_pXHUGE(self): + y = stats.pearsonr(X,HUGE) + r = y[0] + assert_approx_equal(r,1.0) + + def test_pXTINY(self): + y = stats.pearsonr(X,TINY) + r = y[0] + assert_approx_equal(r,1.0) + + def test_pXROUND(self): + y = stats.pearsonr(X,ROUND) + r = y[0] + assert_approx_equal(r,1.0) + + def test_pBIGBIG(self): + y = stats.pearsonr(BIG,BIG) + r = y[0] + assert_approx_equal(r,1.0) + + def test_pBIGLITTLE(self): + y = stats.pearsonr(BIG,LITTLE) + r = y[0] + assert_approx_equal(r,1.0) + + def test_pBIGHUGE(self): + y = stats.pearsonr(BIG,HUGE) + r = y[0] + assert_approx_equal(r,1.0) + + def test_pBIGTINY(self): + y = stats.pearsonr(BIG,TINY) + r = y[0] + assert_approx_equal(r,1.0) + + def test_pBIGROUND(self): + y = stats.pearsonr(BIG,ROUND) + r = y[0] + assert_approx_equal(r,1.0) + + def test_pLITTLELITTLE(self): + y = stats.pearsonr(LITTLE,LITTLE) + r = y[0] + assert_approx_equal(r,1.0) + + def test_pLITTLEHUGE(self): + y = stats.pearsonr(LITTLE,HUGE) + r = y[0] + assert_approx_equal(r,1.0) + + def test_pLITTLETINY(self): + y = stats.pearsonr(LITTLE,TINY) + r = y[0] + assert_approx_equal(r,1.0) + + def test_pLITTLEROUND(self): + y = stats.pearsonr(LITTLE,ROUND) + r = y[0] + assert_approx_equal(r,1.0) + + def test_pHUGEHUGE(self): + y = stats.pearsonr(HUGE,HUGE) + r = y[0] + assert_approx_equal(r,1.0) + + def test_pHUGETINY(self): + y = stats.pearsonr(HUGE,TINY) + r = y[0] + assert_approx_equal(r,1.0) + + def test_pHUGEROUND(self): + y = stats.pearsonr(HUGE,ROUND) + r = y[0] + assert_approx_equal(r,1.0) + + def test_pTINYTINY(self): + y = stats.pearsonr(TINY,TINY) + r = y[0] + assert_approx_equal(r,1.0) + + def test_pTINYROUND(self): + y = stats.pearsonr(TINY,ROUND) + r = y[0] + assert_approx_equal(r,1.0) + + def test_pROUNDROUND(self): + y = stats.pearsonr(ROUND,ROUND) + r = y[0] + assert_approx_equal(r,1.0) + + +@make_xp_test_case(stats.pearsonr) +class TestPearsonr: + def test_pearsonr_result_attributes(self, xp): + res = stats.pearsonr(xp.asarray(X), xp.asarray(X)) + attributes = ('correlation', 'pvalue') + check_named_results(res, attributes, xp=xp) + xp_assert_equal(res.correlation, res.statistic) + + def test_r_almost_exactly_pos1(self, xp): + a = xp.arange(3.0) + r, prob = stats.pearsonr(a, a) + + xp_assert_close(r, xp.asarray(1.0), atol=1e-15) + # With n = len(a) = 3, the error in prob grows like the + # square root of the error in r. + xp_assert_close(prob, xp.asarray(0.0), atol=np.sqrt(2*np.spacing(1.0))) + + def test_r_almost_exactly_neg1(self, xp): + a = xp.arange(3.0) + r, prob = stats.pearsonr(a, -a) + + xp_assert_close(r, xp.asarray(-1.0), atol=1e-15) + # With n = len(a) = 3, the error in prob grows like the + # square root of the error in r. + xp_assert_close(prob, xp.asarray(0.0), atol=np.sqrt(2*np.spacing(1.0))) + + def test_basic(self, xp): + # A basic test, with a correlation coefficient + # that is not 1 or -1. + a = xp.asarray([-1, 0, 1]) + b = xp.asarray([0, 0, 3]) + r, prob = stats.pearsonr(a, b) + xp_assert_close(r, xp.asarray(3**0.5/2)) + xp_assert_close(prob, xp.asarray(1/3)) + + def test_constant_input(self, xp): + # Zero variance input + # See https://github.com/scipy/scipy/issues/3728 + x = xp.asarray([0.667, 0.667, 0.667]) + y = xp.asarray([0.123, 0.456, 0.789]) + msg = "An input array is constant" + with eager_warns(stats.ConstantInputWarning, match=msg, xp=xp): + r, p = stats.pearsonr(x, y) + xp_assert_close(r, xp.asarray(xp.nan)) + xp_assert_close(p, xp.asarray(xp.nan)) + + @pytest.mark.parametrize('dtype', ['float32', 'float64']) + def test_near_constant_input(self, dtype, xp): + npdtype = getattr(np, dtype) + dtype = getattr(xp, dtype) + # Near constant input (but not constant): + x = xp.asarray([2, 2, 2 + np.spacing(2, dtype=npdtype)], dtype=dtype) + y = xp.asarray([3, 3, 3 + 6*np.spacing(3, dtype=npdtype)], dtype=dtype) + msg = "An input array is nearly constant; the computed" + with eager_warns(stats.NearConstantInputWarning, match=msg, xp=xp): + # r and p are garbage, so don't bother checking them in this case. + # (The exact value of r would be 1.) + stats.pearsonr(x, y) + + def test_very_small_input_values(self, xp): + # Very small values in an input. A naive implementation will + # suffer from underflow. + # See https://github.com/scipy/scipy/issues/9353 + x = xp.asarray([0.004434375, 0.004756007, 0.003911996, 0.0038005, 0.003409971], + dtype=xp.float64) + y = xp.asarray([2.48e-188, 7.41e-181, 4.09e-208, 2.08e-223, 2.66e-245], + dtype=xp.float64) + r, p = stats.pearsonr(x, y) + + # The expected values were computed using mpmath with 80 digits + # of precision. + xp_assert_close(r, xp.asarray(0.7272930540750450, dtype=xp.float64)) + xp_assert_close(p, xp.asarray(0.1637805429533202, dtype=xp.float64)) + + def test_very_large_input_values(self, xp): + # Very large values in an input. A naive implementation will + # suffer from overflow. + # See https://github.com/scipy/scipy/issues/8980 + x = 1e90*xp.asarray([0, 0, 0, 1, 1, 1, 1], dtype=xp.float64) + y = 1e90*xp.arange(7, dtype=xp.float64) + + r, p = stats.pearsonr(x, y) + + # The expected values were computed using mpmath with 80 digits + # of precision. + xp_assert_close(r, xp.asarray(0.8660254037844386, dtype=xp.float64)) + xp_assert_close(p, xp.asarray(0.011724811003954638, dtype=xp.float64)) + + def test_extremely_large_input_values(self, xp): + # Extremely large values in x and y. These values would cause the + # product sigma_x * sigma_y to overflow if the two factors were + # computed independently. + x = xp.asarray([2.3e200, 4.5e200, 6.7e200, 8e200], dtype=xp.float64) + y = xp.asarray([1.2e199, 5.5e200, 3.3e201, 1.0e200], dtype=xp.float64) + r, p = stats.pearsonr(x, y) + + # The expected values were computed using mpmath with 80 digits + # of precision. + xp_assert_close(r, xp.asarray(0.351312332103289, dtype=xp.float64)) + xp_assert_close(p, xp.asarray(0.648687667896711, dtype=xp.float64)) + + @pytest.mark.filterwarnings("ignore:invalid value encountered:RuntimeWarning:dask") + @pytest.mark.filterwarnings("ignore:divide by zero encountered:RuntimeWarning:dask") + def test_length_two_pos1(self, xp): + # Inputs with length 2. + # See https://github.com/scipy/scipy/issues/7730 + x = xp.asarray([1., 2.]) + y = xp.asarray([3., 5.]) + res = stats.pearsonr(x, y) + r, p = res + one = xp.asarray(1.) + xp_assert_equal(r, one) + xp_assert_equal(p, one) + low, high = res.confidence_interval() + xp_assert_equal(low, -one) + xp_assert_equal(high, one) + + @pytest.mark.filterwarnings("ignore:invalid value encountered:RuntimeWarning:dask") + @pytest.mark.filterwarnings("ignore:divide by zero encountered:RuntimeWarning:dask") + def test_length_two_neg1(self, xp): + # Inputs with length 2. + # See https://github.com/scipy/scipy/issues/7730 + x = xp.asarray([2., 1.]) + y = xp.asarray([3., 5.]) + res = stats.pearsonr(x, y) + r, p = res + one = xp.asarray(1.) + xp_assert_equal(r, -one) + xp_assert_equal(p, one) + low, high = res.confidence_interval() + xp_assert_equal(low, -one) + xp_assert_equal(high, one) + + @pytest.mark.filterwarnings("ignore:invalid value encountered in divide") + def test_length_two_constant_input(self, xp): + # Zero variance input + # See https://github.com/scipy/scipy/issues/3728 + # and https://github.com/scipy/scipy/issues/7730 + x = xp.asarray([0.667, 0.667]) + y = xp.asarray([0.123, 0.456]) + msg = "An input array is constant" + with eager_warns(stats.ConstantInputWarning, match=msg, xp=xp): + r, p = stats.pearsonr(x, y) + xp_assert_close(r, xp.asarray(xp.nan)) + xp_assert_close(p, xp.asarray(xp.nan)) + + # Expected values computed with R 3.6.2 cor.test, e.g. + # options(digits=16) + # x <- c(1, 2, 3, 4) + # y <- c(0, 1, 0.5, 1) + # cor.test(x, y, method = "pearson", alternative = "g") + # correlation coefficient and p-value for alternative='two-sided' + # calculated with mpmath agree to 16 digits. + @pytest.mark.parametrize('alternative, pval, rlow, rhigh, sign', + [('two-sided', 0.325800137536, -0.814938968841, 0.99230697523, 1), + ('less', 0.8370999312316, -1., 0.985600937290653, 1), + ('greater', 0.1629000687684, -0.6785654158217636, 1., 1), + ('two-sided', 0.325800137536, -0.992306975236, 0.81493896884, -1), + ('less', 0.1629000687684, -1.0, 0.6785654158217636, -1), + ('greater', 0.8370999312316, -0.985600937290653, 1.0, -1)]) + def test_basic_example(self, alternative, pval, rlow, rhigh, sign, xp): + x = xp.asarray([1., 2., 3., 4.]) + y = xp.asarray([0., 1., 0.5, 1.]) * sign + result = stats.pearsonr(x, y, alternative=alternative) + xp_assert_close(result.statistic, xp.asarray(0.6741998624632421*sign)) + xp_assert_close(result.pvalue, xp.asarray(pval)) + ci = result.confidence_interval() + xp_assert_close(ci.low, xp.asarray(rlow)) + xp_assert_close(ci.high, xp.asarray(rhigh)) + + def test_negative_correlation_pvalue_gh17795(self, xp): + x = xp.arange(10.) + y = -x + test_greater = stats.pearsonr(x, y, alternative='greater') + test_less = stats.pearsonr(x, y, alternative='less') + xp_assert_close(test_greater.pvalue, xp.asarray(1.)) + xp_assert_close(test_less.pvalue, xp.asarray(0.), atol=1e-20) + + @pytest.mark.filterwarnings("ignore:divide by zero encountered:RuntimeWarning:dask") + def test_length3_r_exactly_negative_one(self, xp): + x = xp.asarray([1., 2., 3.]) + y = xp.asarray([5., -4., -13.]) + res = stats.pearsonr(x, y) + + # The expected r and p are exact. + r, p = res + one = xp.asarray(1.0) + xp_assert_close(r, -one) + xp_assert_close(p, 0*one, atol=1e-7) + low, high = res.confidence_interval() + xp_assert_equal(low, -one) + xp_assert_equal(high, one) + + def test_input_validation(self): + # Arraylike is np only + x = [1, 2, 3] + y = [4] + message = '`x` and `y` must have the same length along `axis`.' + with pytest.raises(ValueError, match=message): + stats.pearsonr(x, y) + + x = [1, 2, 3] + y = [4, 5] + message = '`x` and `y` must be broadcastable.' + with pytest.raises(ValueError, match=message): + stats.pearsonr(x, y) + + x = [1] + y = [2] + message = '`x` and `y` must have length at least 2.' + with pytest.raises(ValueError, match=message): + stats.pearsonr(x, y) + + x = [-1j, -2j, -3.0j] + y = [-1j, -2j, -3.0j] + message = 'This function does not support complex data' + with pytest.raises(ValueError, match=message): + stats.pearsonr(x, y) + + message = "`method` must be an instance of..." + with pytest.raises(ValueError, match=message): + stats.pearsonr([1, 2], [3, 4], method="asymptotic") + + res = stats.pearsonr([1, 2], [3, 4]) + with pytest.raises(ValueError, match=message): + res.confidence_interval(method="exact") + + @pytest.mark.fail_slow(10) + @pytest.mark.xfail_on_32bit("Monte Carlo method needs > a few kB of memory") + @pytest.mark.parametrize('alternative', ('less', 'greater', 'two-sided')) + @pytest.mark.parametrize('method_name', + ('permutation', 'monte_carlo', 'monte_carlo2')) + def test_resampling_pvalue(self, method_name, alternative): + rng = np.random.default_rng(24623935790378923) + size = (2, 100) if method_name == 'permutation' else (2, 1000) + x = rng.normal(size=size) + y = rng.normal(size=size) + methods = {'permutation': stats.PermutationMethod(rng=rng), + 'monte_carlo': stats.MonteCarloMethod(rvs=(rng.normal,)*2), + 'monte_carlo2': stats.MonteCarloMethod(rng=1294)} + method = methods[method_name] + res = stats.pearsonr(x, y, alternative=alternative, method=method, axis=-1) + ref = stats.pearsonr(x, y, alternative=alternative, axis=-1) + assert_allclose(res.statistic, ref.statistic, rtol=1e-15) + assert_allclose(res.pvalue, ref.pvalue, rtol=1e-2, atol=1e-3) + + if method_name == 'monte_carlo2': + method = stats.MonteCarloMethod(rng=1294) + res2 = stats.pearsonr(x, y, alternative=alternative, method=method, axis=-1) + assert_equal(res2.statistic, res.statistic) + assert_equal(res2.pvalue, res.pvalue) + + @pytest.mark.slow + @pytest.mark.parametrize('alternative', ('less', 'greater', 'two-sided')) + def test_bootstrap_ci(self, alternative): + rng = np.random.default_rng(2462935790378923) + x = rng.normal(size=(2, 100)) + y = rng.normal(size=(2, 100)) + res = stats.pearsonr(x, y, alternative=alternative, axis=-1) + + # preserve use of old random_state during SPEC 7 transition + rng = np.random.default_rng(724358723498249852) + method = stats.BootstrapMethod(random_state=rng) + res_ci = res.confidence_interval(method=method) + ref_ci = res.confidence_interval() + assert_allclose(res_ci, ref_ci, atol=1.5e-2) + + # `rng` is the new argument name` + rng = np.random.default_rng(724358723498249852) + method = stats.BootstrapMethod(rng=rng) + res_ci2 = res.confidence_interval(method=method) + assert_allclose(res_ci2, res_ci) + + @pytest.mark.parametrize('axis', [0, 1]) + def test_axis01(self, axis): + rng = np.random.default_rng(38572345825) + shape = (9, 10) + x, y = rng.normal(size=(2,) + shape) + res = stats.pearsonr(x, y, axis=axis) + ci = res.confidence_interval() + if axis == 0: + x, y = x.T, y.T + for i in range(x.shape[0]): + res_i = stats.pearsonr(x[i], y[i]) + ci_i = res_i.confidence_interval() + assert_allclose(res.statistic[i], res_i.statistic) + assert_allclose(res.pvalue[i], res_i.pvalue) + assert_allclose(ci.low[i], ci_i.low) + assert_allclose(ci.high[i], ci_i.high) + + def test_axis_None(self): + rng = np.random.default_rng(38572345825) + shape = (9, 10) + x, y = rng.normal(size=(2,) + shape) + res = stats.pearsonr(x, y, axis=None) + ci = res.confidence_interval() + ref = stats.pearsonr(x.ravel(), y.ravel()) + ci_ref = ref.confidence_interval() + assert_allclose(res.statistic, ref.statistic) + assert_allclose(res.pvalue, ref.pvalue) + assert_allclose(ci, ci_ref) + + def test_nd_input_validation(self, xp): + x = y = xp.ones((2, 5)) + message = '`axis` must be an integer.' + with pytest.raises(ValueError, match=message): + stats.pearsonr(x, y, axis=1.5) + + message = '`x` and `y` must have the same length along `axis`' + with pytest.raises(ValueError, match=message): + stats.pearsonr(x, xp.ones((2, 1)), axis=1) + + message = '`x` and `y` must have length at least 2.' + with pytest.raises(ValueError, match=message): + stats.pearsonr(xp.ones((2, 1)), xp.ones((2, 1)), axis=1) + + message = '`x` and `y` must be broadcastable.' + with pytest.raises(ValueError, match=message): + stats.pearsonr(x, xp.ones((3, 5)), axis=1) + + message = '`method` must be `None` if arguments are not NumPy arrays.' + if not is_numpy(xp): + x = xp.arange(10) + with pytest.raises(ValueError, match=message): + stats.pearsonr(x, x, method=stats.PermutationMethod()) + + @pytest.mark.filterwarnings("ignore:invalid value encountered:RuntimeWarning:dask") + @pytest.mark.filterwarnings("ignore:divide by zero encountered:RuntimeWarning:dask") + def test_nd_special_cases(self, xp): + rng = np.random.default_rng(34989235492245) + + x0, y0 = rng.random((4, 5)), rng.random((4, 5)) + x0[0, ...] = 1 + y0[1, ...] = 2 + x, y = xp.asarray(x0), xp.asarray(y0) + message = 'An input array is constant' + with eager_warns(stats.ConstantInputWarning, match=message, xp=xp): + res = stats.pearsonr(x, y, axis=1) + ci = res.confidence_interval() + nans = xp.asarray([xp.nan, xp.nan], dtype=xp.float64) + xp_assert_equal(res.statistic[0:2], nans) + xp_assert_equal(res.pvalue[0:2], nans) + xp_assert_equal(ci.low[0:2], nans) + xp_assert_equal(ci.high[0:2], nans) + assert xp.all(xp.isfinite(res.statistic[2:])) + assert xp.all(xp.isfinite(res.pvalue[2:])) + assert xp.all(xp.isfinite(ci.low[2:])) + assert xp.all(xp.isfinite(ci.high[2:])) + + x0[0, 0], y0[1, 1] = 1 + 1e-15, 2 + 1e-15 + x, y = xp.asarray(x0), xp.asarray(y0) + message = 'An input array is nearly constant' + with eager_warns(stats.NearConstantInputWarning, match=message, xp=xp): + stats.pearsonr(x, y, axis=1) + + # length 2 along axis + x = xp.asarray([[1, 2], [1, 2], [2, 1], [2, 1.]]) + y = xp.asarray([[1, 2], [2, 1], [1, 2], [2, 1.]]) + ones = xp.ones(4) + res = stats.pearsonr(x, y, axis=-1) + ci = res.confidence_interval() + xp_assert_close(res.statistic, xp.asarray([1, -1, -1, 1.])) + xp_assert_close(res.pvalue, ones) + xp_assert_close(ci.low, -ones) + xp_assert_close(ci.high, ones) + + def test_different_dimensionality(self, xp): + # For better or for worse, there is one difference between the broadcasting + # behavior of most stats functions and NumPy gufuncs / NEP 5: gufuncs `axis` + # refers to the core dimension *before* prepending `1`s to the array shapes + # to match dimensionality; SciPy's prepends `1`s first. For instance, in + # SciPy, `vecdot` would work just like `xp.sum(x * y, axis=axis)`, but this + # is NOT true of NumPy. The discrepancy only arises when there are multiple + # arguments with different dimensionality and positive indices are used, + # which is probably why it hasn't been a problem. There are pros and cons of + # each convention, and we might want to consider changing our behavior in + # SciPy 2.0. For now, preserve consistency / backward compatibility. + rng = np.random.default_rng(45834598265019344) + x = rng.random((3, 10)) + y = rng.random(10) + res = stats.pearsonr(x, y, axis=1) + ref = stats.pearsonr(x, y, axis=-1) + assert_allclose(res.statistic, ref.statistic, rtol=3e-16) + + @pytest.mark.parametrize('axis', [0, 1, None]) + @pytest.mark.parametrize('alternative', ['less', 'greater', 'two-sided']) + def test_array_api(self, xp, axis, alternative): + rng = np.random.default_rng(2134418962) + x, y = rng.normal(size=(2, 10, 11)) + res = stats.pearsonr(xp.asarray(x), xp.asarray(y), + axis=axis, alternative=alternative) + ref = stats.pearsonr(x, y, axis=axis, alternative=alternative) + xp_assert_close(res.statistic, xp.asarray(ref.statistic)) + xp_assert_close(res.pvalue, xp.asarray(ref.pvalue)) + + res_ci = res.confidence_interval() + ref_ci = ref.confidence_interval() + xp_assert_close(res_ci.low, xp.asarray(ref_ci.low)) + xp_assert_close(res_ci.high, xp.asarray(ref_ci.high)) + + +class TestFisherExact: + """Some tests to show that fisher_exact() works correctly. + + Note that in SciPy 0.9.0 this was not working well for large numbers due to + inaccuracy of the hypergeom distribution (see #1218). Fixed now. + + Also note that R and SciPy have different argument formats for their + hypergeometric distribution functions. + + R: + > phyper(18999, 99000, 110000, 39000, lower.tail = FALSE) + [1] 1.701815e-09 + """ + + def test_basic(self): + fisher_exact = stats.fisher_exact + + res = fisher_exact([[14500, 20000], [30000, 40000]])[1] + assert_approx_equal(res, 0.01106, significant=4) + res = fisher_exact([[100, 2], [1000, 5]])[1] + assert_approx_equal(res, 0.1301, significant=4) + res = fisher_exact([[2, 7], [8, 2]])[1] + assert_approx_equal(res, 0.0230141, significant=6) + res = fisher_exact([[5, 1], [10, 10]])[1] + assert_approx_equal(res, 0.1973244, significant=6) + res = fisher_exact([[5, 15], [20, 20]])[1] + assert_approx_equal(res, 0.0958044, significant=6) + res = fisher_exact([[5, 16], [20, 25]])[1] + assert_approx_equal(res, 0.1725862, significant=6) + res = fisher_exact([[10, 5], [10, 1]])[1] + assert_approx_equal(res, 0.1973244, significant=6) + res = fisher_exact([[5, 0], [1, 4]])[1] + assert_approx_equal(res, 0.04761904, significant=6) + res = fisher_exact([[0, 1], [3, 2]])[1] + assert_approx_equal(res, 1.0) + res = fisher_exact([[0, 2], [6, 4]])[1] + assert_approx_equal(res, 0.4545454545) + res = fisher_exact([[2, 7], [8, 2]]) + assert_approx_equal(res[1], 0.0230141, significant=6) + assert_approx_equal(res[0], 4.0 / 56) + + def test_precise(self): + # results from R + # + # R defines oddsratio differently (see Notes section of fisher_exact + # docstring), so those will not match. We leave them in anyway, in + # case they will be useful later on. We test only the p-value. + tablist = [ + ([[100, 2], [1000, 5]], (2.505583993422285e-001, 1.300759363430016e-001)), + ([[2, 7], [8, 2]], (8.586235135736206e-002, 2.301413756522114e-002)), + ([[5, 1], [10, 10]], (4.725646047336584e+000, 1.973244147157190e-001)), + ([[5, 15], [20, 20]], (3.394396617440852e-001, 9.580440012477637e-002)), + ([[5, 16], [20, 25]], (3.960558326183334e-001, 1.725864953812994e-001)), + ([[10, 5], [10, 1]], (2.116112781158483e-001, 1.973244147157190e-001)), + ([[10, 5], [10, 0]], (0.000000000000000e+000, 6.126482213438734e-002)), + ([[5, 0], [1, 4]], (np.inf, 4.761904761904762e-002)), + ([[0, 5], [1, 4]], (0.000000000000000e+000, 1.000000000000000e+000)), + ([[5, 1], [0, 4]], (np.inf, 4.761904761904758e-002)), + ([[0, 1], [3, 2]], (0.000000000000000e+000, 1.000000000000000e+000)) + ] + for table, res_r in tablist: + res = stats.fisher_exact(np.asarray(table)) + np.testing.assert_almost_equal(res[1], res_r[1], decimal=11, + verbose=True) + + def test_gh4130(self): + # Previously, a fudge factor used to distinguish between theoretically + # and numerically different probability masses was 1e-4; it has been + # tightened to fix gh4130. Accuracy checked against R fisher.test. + # options(digits=16) + # table <- matrix(c(6, 108, 37, 200), nrow = 2) + # fisher.test(table, alternative = "t") + x = [[6, 37], [108, 200]] + res = stats.fisher_exact(x) + assert_allclose(res[1], 0.005092697748126) + + # case from https://github.com/brentp/fishers_exact_test/issues/27 + # That package has an (absolute?) fudge factor of 1e-6; too big + x = [[22, 0], [0, 102]] + res = stats.fisher_exact(x) + assert_allclose(res[1], 7.175066786244549e-25) + + # case from https://github.com/brentp/fishers_exact_test/issues/1 + x = [[94, 48], [3577, 16988]] + res = stats.fisher_exact(x) + assert_allclose(res[1], 2.069356340993818e-37) + + def test_gh9231(self): + # Previously, fisher_exact was extremely slow for this table + # As reported in gh-9231, the p-value should be very nearly zero + x = [[5829225, 5692693], [5760959, 5760959]] + res = stats.fisher_exact(x) + assert_allclose(res[1], 0, atol=1e-170) + + @pytest.mark.slow + def test_large_numbers(self): + # Test with some large numbers. Regression test for #1401 + pvals = [5.56e-11, 2.666e-11, 1.363e-11] # from R + for pval, num in zip(pvals, [75, 76, 77]): + res = stats.fisher_exact([[17704, 496], [1065, num]])[1] + assert_approx_equal(res, pval, significant=4) + + res = stats.fisher_exact([[18000, 80000], [20000, 90000]])[1] + assert_approx_equal(res, 0.2751, significant=4) + + def test_raises(self): + # test we raise an error for wrong number of dimensions. + message = "The input `table` must have two dimensions." + with pytest.raises(ValueError, match=message): + stats.fisher_exact(np.arange(6)) + + def test_row_or_col_zero(self): + tables = ([[0, 0], [5, 10]], + [[5, 10], [0, 0]], + [[0, 5], [0, 10]], + [[5, 0], [10, 0]]) + for table in tables: + oddsratio, pval = stats.fisher_exact(table) + assert_equal(pval, 1.0) + assert_equal(oddsratio, np.nan) + + def test_less_greater(self): + tables = ( + # Some tables to compare with R: + [[2, 7], [8, 2]], + [[200, 7], [8, 300]], + [[28, 21], [6, 1957]], + [[190, 800], [200, 900]], + # Some tables with simple exact values + # (includes regression test for ticket #1568): + [[0, 2], [3, 0]], + [[1, 1], [2, 1]], + [[2, 0], [1, 2]], + [[0, 1], [2, 3]], + [[1, 0], [1, 4]], + ) + pvals = ( + # from R: + [0.018521725952066501, 0.9990149169715733], + [1.0, 2.0056578803889148e-122], + [1.0, 5.7284374608319831e-44], + [0.7416227, 0.2959826], + # Exact: + [0.1, 1.0], + [0.7, 0.9], + [1.0, 0.3], + [2./3, 1.0], + [1.0, 1./3], + ) + for table, pval in zip(tables, pvals): + res = [] + res.append(stats.fisher_exact(table, alternative="less")[1]) + res.append(stats.fisher_exact(table, alternative="greater")[1]) + assert_allclose(res, pval, atol=0, rtol=1e-7) + + def test_gh3014(self): + # check if issue #3014 has been fixed. + # before, this would have risen a ValueError + odds, pvalue = stats.fisher_exact([[1, 2], [9, 84419233]]) + + @pytest.mark.parametrize("alternative", ['two-sided', 'less', 'greater']) + def test_result(self, alternative): + table = np.array([[14500, 20000], [30000, 40000]]) + res = stats.fisher_exact(table, alternative=alternative) + assert_equal((res.statistic, res.pvalue), res) + + def test_input_validation_edge_cases_rxc(self): + rng = np.random.default_rng(2345783457834572345) + table = np.asarray([[2, 7], [8, 2]]) + + message = r"`alternative` must be the default \(None\) unless..." + with pytest.raises(ValueError, match=message): + method = stats.PermutationMethod(rng=rng) + stats.fisher_exact(table, method=method, alternative='less') + + message = "...not recognized; if provided, `method` must be an..." + with pytest.raises(ValueError, match=message): + method = stats.BootstrapMethod(rng=rng) + stats.fisher_exact(table, method=method) + + message = "If the `method` argument of `fisher_exact` is an..." + with pytest.raises(ValueError, match=message): + method = stats.MonteCarloMethod(rvs=stats.norm.rvs) + stats.fisher_exact(table, method=method) + + message = "`table` must have at least one row and one column." + with pytest.raises(ValueError, match=message): + stats.fisher_exact(np.zeros((0, 1))) + + # Specical case: when there is only one table with given marginals, the + # PMF of that case is 1.0, so the p-value is 1.0 + np.testing.assert_equal(stats.fisher_exact([[1, 2, 3]]), (1, 1)) + np.testing.assert_equal(stats.fisher_exact([[1], [2], [3]]), (1, 1)) + np.testing.assert_equal(stats.fisher_exact(np.zeros((2, 3))), (1, 1)) + + + + @pytest.mark.fail_slow(10) + @pytest.mark.slow() + def test_resampling_2x2(self): + rng = np.random.default_rng(2345783457834572345) + table = np.asarray([[2, 7], [8, 2]]) + ref = stats.fisher_exact(table) + ref_pvalue = ref.pvalue + ref_stat = stats.random_table(table.sum(axis=1), table.sum(axis=0)).pmf(table) + + method = stats.MonteCarloMethod(rng=rng) + res = stats.fisher_exact(table, method=method) + assert_allclose(res.pvalue, ref_pvalue, atol=0.0025) + assert_equal(res.statistic, ref_stat) + + method = stats.PermutationMethod(rng=rng) + res = stats.fisher_exact(table, method=method) + assert_allclose(res.pvalue, ref.pvalue, atol=0.0025) + assert_equal(res.statistic, ref_stat) + + @pytest.mark.fail_slow(10) + @pytest.mark.slow() + def test_resampling_rxc(self): + # Compare against R fisher.exact + # options(digits=16) + # MP6 < - rbind( + # c(1, 2, 2, 1, 1, 0, 1), + # c(2, 0, 0, 2, 3, 0, 0), + # c(0, 1, 1, 1, 2, 7, 3), + # c(1, 1, 2, 0, 0, 0, 1), + # c(0, 1, 1, 1, 1, 0, 0)) + # fisher.test(MP6) + + table = [[1, 2, 2, 1, 1, 0, 1], + [2, 0, 0, 2, 3, 0, 0], + [0, 1, 1, 1, 2, 7, 3], + [1, 1, 2, 0, 0, 0, 1], + [0, 1, 1, 1, 1, 0, 0]] + table = np.asarray(table) + + ref_pvalue = 0.03928964365533 + rng = np.random.default_rng(3928964365533) + + method = stats.PermutationMethod(rng=rng) + res = stats.fisher_exact(table, method=method) + assert_allclose(res.pvalue, ref_pvalue, atol=5e-4) + + method = stats.MonteCarloMethod(rng=rng, n_resamples=99999) + res = stats.fisher_exact(table, method=method) + assert_allclose(res.pvalue, ref_pvalue, atol=5e-4) + + @pytest.mark.xslow() + def test_resampling_exact_2x2(self): + # Test that exact permutation p-value matches result of `fisher_exact` + rng = np.random.default_rng(2345783457834572345) + method = stats.PermutationMethod(rng=rng) + + for a in range(1, 3): + for b in range(1, 3): + for c in range(1, 3): + for d in range(1, 4): + table = np.asarray([[a, b], [c, d]]) + ref = stats.fisher_exact(table) + res = stats.fisher_exact(table, method=method) + assert_allclose(res.pvalue, ref.pvalue, atol=1e-14) + + +class TestCorrSpearmanr: + """ W.II.D. Compute a correlation matrix on all the variables. + + All the correlations, except for ZERO and MISS, should be exactly 1. + ZERO and MISS should have undefined or missing correlations with the + other variables. The same should go for SPEARMAN correlations, if + your program has them. + """ + + def setup_method(self): + self.rng = np.random.default_rng(228584263) + + def test_scalar(self): + y = stats.spearmanr(4., 2.) + assert_(np.isnan(y).all()) + + def test_uneven_lengths(self): + assert_raises(ValueError, stats.spearmanr, [1, 2, 1], [8, 9]) + assert_raises(ValueError, stats.spearmanr, [1, 2, 1], 8) + + def test_uneven_2d_shapes(self): + # Different number of columns should work - those just get concatenated. + x = self.rng.standard_normal((4, 3)) + y = self.rng.standard_normal((4, 2)) + assert stats.spearmanr(x, y).statistic.shape == (5, 5) + assert stats.spearmanr(x.T, y.T, axis=1).pvalue.shape == (5, 5) + + assert_raises(ValueError, stats.spearmanr, x, y, axis=1) + assert_raises(ValueError, stats.spearmanr, x.T, y.T) + + def test_ndim_too_high(self): + x = self.rng.standard_normal((4, 3, 2)) + assert_raises(ValueError, stats.spearmanr, x) + assert_raises(ValueError, stats.spearmanr, x, x) + assert_raises(ValueError, stats.spearmanr, x, None, None) + # But should work with axis=None (raveling axes) for two input arrays + assert_allclose(stats.spearmanr(x, x, axis=None), + stats.spearmanr(x.flatten(), x.flatten(), axis=0)) + + def test_nan_policy(self): + x = np.arange(10.) + x[9] = np.nan + assert_array_equal(stats.spearmanr(x, x), (np.nan, np.nan)) + assert_array_equal(stats.spearmanr(x, x, nan_policy='omit'), + (1.0, 0.0)) + assert_raises(ValueError, stats.spearmanr, x, x, nan_policy='raise') + assert_raises(ValueError, stats.spearmanr, x, x, nan_policy='foobar') + + def test_nan_policy_bug_12458(self): + rng = np.random.default_rng(8119864466) + x = rng.random((5, 10)) + k = 6 + x[:, k] = np.nan + y = np.delete(x, k, axis=1) + corx, px = stats.spearmanr(x, nan_policy='omit') + cory, py = stats.spearmanr(y) + corx = np.delete(np.delete(corx, k, axis=1), k, axis=0) + px = np.delete(np.delete(px, k, axis=1), k, axis=0) + assert_allclose(corx, cory, atol=1e-14) + assert_allclose(px, py, atol=1e-14) + + def test_nan_policy_bug_12411(self): + m = 5 + n = 10 + x = self.rng.standard_normal((m, n)) + x[1, 0] = np.nan + x[3, -1] = np.nan + corr, pvalue = stats.spearmanr(x, axis=1, nan_policy="propagate") + res = [[stats.spearmanr(x[i, :], x[j, :]).statistic for i in range(m)] + for j in range(m)] + assert_allclose(corr, res) + + @pytest.mark.parametrize("axis", [0, 1]) + def test_constant_input_gh22816(self, axis): + X = np.zeros((10, 4)) + X[:, 2] = np.arange(0, 10) + X[:, 3] = -np.arange(0, 10) + + corr = np.array([[np.nan, np.nan, np.nan, np.nan], + [np.nan, np.nan, np.nan, np.nan], + [np.nan, np.nan, 1., -1.], + [np.nan, np.nan, -1., 1.]]) + if axis == 1: + X = X.T + with pytest.warns(stats._warnings_errors.ConstantInputWarning): + res = stats.spearmanr(X, axis=axis) + assert_allclose(res.statistic, corr) + assert_allclose(res.pvalue, 0*corr, atol=1e-50) + + def test_sXX(self): + y = stats.spearmanr(X,X) + r = y[0] + assert_approx_equal(r,1.0) + + def test_sXBIG(self): + y = stats.spearmanr(X,BIG) + r = y[0] + assert_approx_equal(r,1.0) + + def test_sXLITTLE(self): + y = stats.spearmanr(X,LITTLE) + r = y[0] + assert_approx_equal(r,1.0) + + def test_sXHUGE(self): + y = stats.spearmanr(X,HUGE) + r = y[0] + assert_approx_equal(r,1.0) + + def test_sXTINY(self): + y = stats.spearmanr(X,TINY) + r = y[0] + assert_approx_equal(r,1.0) + + def test_sXROUND(self): + y = stats.spearmanr(X,ROUND) + r = y[0] + assert_approx_equal(r,1.0) + + def test_sBIGBIG(self): + y = stats.spearmanr(BIG,BIG) + r = y[0] + assert_approx_equal(r,1.0) + + def test_sBIGLITTLE(self): + y = stats.spearmanr(BIG,LITTLE) + r = y[0] + assert_approx_equal(r,1.0) + + def test_sBIGHUGE(self): + y = stats.spearmanr(BIG,HUGE) + r = y[0] + assert_approx_equal(r,1.0) + + def test_sBIGTINY(self): + y = stats.spearmanr(BIG,TINY) + r = y[0] + assert_approx_equal(r,1.0) + + def test_sBIGROUND(self): + y = stats.spearmanr(BIG,ROUND) + r = y[0] + assert_approx_equal(r,1.0) + + def test_sLITTLELITTLE(self): + y = stats.spearmanr(LITTLE,LITTLE) + r = y[0] + assert_approx_equal(r,1.0) + + def test_sLITTLEHUGE(self): + y = stats.spearmanr(LITTLE,HUGE) + r = y[0] + assert_approx_equal(r,1.0) + + def test_sLITTLETINY(self): + y = stats.spearmanr(LITTLE,TINY) + r = y[0] + assert_approx_equal(r,1.0) + + def test_sLITTLEROUND(self): + y = stats.spearmanr(LITTLE,ROUND) + r = y[0] + assert_approx_equal(r,1.0) + + def test_sHUGEHUGE(self): + y = stats.spearmanr(HUGE,HUGE) + r = y[0] + assert_approx_equal(r,1.0) + + def test_sHUGETINY(self): + y = stats.spearmanr(HUGE,TINY) + r = y[0] + assert_approx_equal(r,1.0) + + def test_sHUGEROUND(self): + y = stats.spearmanr(HUGE,ROUND) + r = y[0] + assert_approx_equal(r,1.0) + + def test_sTINYTINY(self): + y = stats.spearmanr(TINY,TINY) + r = y[0] + assert_approx_equal(r,1.0) + + def test_sTINYROUND(self): + y = stats.spearmanr(TINY,ROUND) + r = y[0] + assert_approx_equal(r,1.0) + + def test_sROUNDROUND(self): + y = stats.spearmanr(ROUND,ROUND) + r = y[0] + assert_approx_equal(r,1.0) + + def test_spearmanr_result_attributes(self): + res = stats.spearmanr(X, X) + attributes = ('correlation', 'pvalue') + check_named_results(res, attributes) + assert_equal(res.correlation, res.statistic) + + def test_1d_vs_2d(self): + x1 = [1, 2, 3, 4, 5, 6] + x2 = [1, 2, 3, 4, 6, 5] + res1 = stats.spearmanr(x1, x2) + res2 = stats.spearmanr(np.asarray([x1, x2]).T) + assert_allclose(res1, res2) + + def test_1d_vs_2d_nans(self): + # Now the same with NaNs present. Regression test for gh-9103. + for nan_policy in ['propagate', 'omit']: + x1 = [1, np.nan, 3, 4, 5, 6] + x2 = [1, 2, 3, 4, 6, np.nan] + res1 = stats.spearmanr(x1, x2, nan_policy=nan_policy) + res2 = stats.spearmanr(np.asarray([x1, x2]).T, nan_policy=nan_policy) + assert_allclose(res1, res2) + + def test_3cols(self): + x1 = np.arange(6) + x2 = -x1 + x3 = np.array([0, 1, 2, 3, 5, 4]) + x = np.asarray([x1, x2, x3]).T + actual = stats.spearmanr(x) + expected_corr = np.array([[1, -1, 0.94285714], + [-1, 1, -0.94285714], + [0.94285714, -0.94285714, 1]]) + expected_pvalue = np.zeros((3, 3), dtype=float) + expected_pvalue[2, 0:2] = 0.00480466472 + expected_pvalue[0:2, 2] = 0.00480466472 + + assert_allclose(actual.statistic, expected_corr) + assert_allclose(actual.pvalue, expected_pvalue) + + def test_gh_9103(self): + # Regression test for gh-9103. + x = np.array([[np.nan, 3.0, 4.0, 5.0, 5.1, 6.0, 9.2], + [5.0, np.nan, 4.1, 4.8, 4.9, 5.0, 4.1], + [0.5, 4.0, 7.1, 3.8, 8.0, 5.1, 7.6]]).T + corr = np.array([[np.nan, np.nan, np.nan], + [np.nan, np.nan, np.nan], + [np.nan, np.nan, 1.]]) + assert_allclose(stats.spearmanr(x, nan_policy='propagate').statistic, + corr) + + res = stats.spearmanr(x, nan_policy='omit').statistic + assert_allclose((res[0][1], res[0][2], res[1][2]), + (0.2051957, 0.4857143, -0.4707919), rtol=1e-6) + + def test_gh_8111(self): + # Regression test for gh-8111 (different result for float/int/bool). + n = 100 + rng = np.random.RandomState(234568) + x = rng.rand(n) + m = rng.rand(n) > 0.7 + + # bool against float, no nans + a = (x > .5) + b = np.array(x) + res1 = stats.spearmanr(a, b, nan_policy='omit').statistic + + # bool against float with NaNs + b[m] = np.nan + res2 = stats.spearmanr(a, b, nan_policy='omit').statistic + + # int against float with NaNs + a = a.astype(np.int32) + res3 = stats.spearmanr(a, b, nan_policy='omit').statistic + + expected = [0.865895477, 0.866100381, 0.866100381] + assert_allclose([res1, res2, res3], expected) + + +class TestCorrSpearmanr2: + """Some further tests of the spearmanr function.""" + + def test_spearmanr_vs_r(self): + # Cross-check with R: + # cor.test(c(1,2,3,4,5),c(5,6,7,8,7),method="spearmanr") + x1 = [1, 2, 3, 4, 5] + x2 = [5, 6, 7, 8, 7] + expected = (0.82078268166812329, 0.088587005313543798) + res = stats.spearmanr(x1, x2) + assert_approx_equal(res[0], expected[0]) + assert_approx_equal(res[1], expected[1]) + + def test_empty_arrays(self): + assert_equal(stats.spearmanr([], []), (np.nan, np.nan)) + + def test_normal_draws(self): + rng = np.random.RandomState(7546) + x = np.array([rng.normal(loc=1, scale=1, size=500), + rng.normal(loc=1, scale=1, size=500)]) + corr = [[1.0, 0.3], + [0.3, 1.0]] + x = np.dot(np.linalg.cholesky(corr), x) + expected = (0.28659685838743354, 6.579862219051161e-11) + res = stats.spearmanr(x[0], x[1]) + assert_approx_equal(res[0], expected[0]) + assert_approx_equal(res[1], expected[1]) + + def test_corr_1(self): + assert_approx_equal(stats.spearmanr([1, 1, 2], [1, 1, 2])[0], 1.0) + + def test_nan_policies(self): + x = np.arange(10.) + x[9] = np.nan + assert_array_equal(stats.spearmanr(x, x), (np.nan, np.nan)) + assert_allclose(stats.spearmanr(x, x, nan_policy='omit'), + (1.0, 0)) + assert_raises(ValueError, stats.spearmanr, x, x, nan_policy='raise') + assert_raises(ValueError, stats.spearmanr, x, x, nan_policy='foobar') + + def test_unequal_lengths(self): + x = np.arange(10.) + y = np.arange(20.) + assert_raises(ValueError, stats.spearmanr, x, y) + + def test_omit_paired_value(self): + x1 = [1, 2, 3, 4] + x2 = [8, 7, 6, np.nan] + res1 = stats.spearmanr(x1, x2, nan_policy='omit') + res2 = stats.spearmanr(x1[:3], x2[:3], nan_policy='omit') + assert_equal(res1, res2) + + def test_gh_issue_6061_windows_overflow(self): + x = list(range(2000)) + y = list(range(2000)) + y[0], y[9] = y[9], y[0] + y[10], y[434] = y[434], y[10] + y[435], y[1509] = y[1509], y[435] + # rho = 1 - 6 * (2 * (9^2 + 424^2 + 1074^2))/(2000 * (2000^2 - 1)) + # = 1 - (1 / 500) + # = 0.998 + x.append(np.nan) + y.append(3.0) + assert_almost_equal(stats.spearmanr(x, y, nan_policy='omit')[0], 0.998) + + def test_tie0(self): + # with only ties in one or both inputs + warn_msg = "An input array is constant" + with pytest.warns(stats.ConstantInputWarning, match=warn_msg): + r, p = stats.spearmanr([2, 2, 2], [2, 2, 2]) + assert_equal(r, np.nan) + assert_equal(p, np.nan) + r, p = stats.spearmanr([2, 0, 2], [2, 2, 2]) + assert_equal(r, np.nan) + assert_equal(p, np.nan) + r, p = stats.spearmanr([2, 2, 2], [2, 0, 2]) + assert_equal(r, np.nan) + assert_equal(p, np.nan) + + def test_tie1(self): + # Data + x = [1.0, 2.0, 3.0, 4.0] + y = [1.0, 2.0, 2.0, 3.0] + # Ranks of the data, with tie-handling. + xr = [1.0, 2.0, 3.0, 4.0] + yr = [1.0, 2.5, 2.5, 4.0] + # Result of spearmanr should be the same as applying + # pearsonr to the ranks. + sr = stats.spearmanr(x, y) + pr = stats.pearsonr(xr, yr) + assert_almost_equal(sr, pr) + + def test_tie2(self): + # Test tie-handling if inputs contain nan's + # Data without nan's + x1 = [1, 2, 2.5, 2] + y1 = [1, 3, 2.5, 4] + # Same data with nan's + x2 = [1, 2, 2.5, 2, np.nan] + y2 = [1, 3, 2.5, 4, np.nan] + + # Results for two data sets should be the same if nan's are ignored + sr1 = stats.spearmanr(x1, y1) + sr2 = stats.spearmanr(x2, y2, nan_policy='omit') + assert_almost_equal(sr1, sr2) + + def test_ties_axis_1(self): + z1 = np.array([[1, 1, 1, 1], [1, 2, 3, 4]]) + z2 = np.array([[1, 2, 3, 4], [1, 1, 1, 1]]) + z3 = np.array([[1, 1, 1, 1], [1, 1, 1, 1]]) + warn_msg = "An input array is constant" + with pytest.warns(stats.ConstantInputWarning, match=warn_msg): + r, p = stats.spearmanr(z1, axis=1) + assert_equal(r, np.nan) + assert_equal(p, np.nan) + r, p = stats.spearmanr(z2, axis=1) + assert_equal(r, np.nan) + assert_equal(p, np.nan) + r, p = stats.spearmanr(z3, axis=1) + assert_equal(r, np.nan) + assert_equal(p, np.nan) + + def test_gh_11111(self): + x = np.array([1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]) + y = np.array([0, 0.009783728115345005, 0, 0, 0.0019759230121848587, + 0.0007535430349118562, 0.0002661781514710257, 0, 0, + 0.0007835762419683435]) + warn_msg = "An input array is constant" + with pytest.warns(stats.ConstantInputWarning, match=warn_msg): + r, p = stats.spearmanr(x, y) + assert_equal(r, np.nan) + assert_equal(p, np.nan) + + def test_index_error(self): + x = np.array([1.0, 7.0, 2.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]) + y = np.array([0, 0.009783728115345005, 0, 0, 0.0019759230121848587, + 0.0007535430349118562, 0.0002661781514710257, 0, 0, + 0.0007835762419683435]) + assert_raises(ValueError, stats.spearmanr, x, y, axis=2) + + def test_alternative(self): + # Test alternative parameter + + # Simple test - Based on the above ``test_spearmanr_vs_r`` + x1 = [1, 2, 3, 4, 5] + x2 = [5, 6, 7, 8, 7] + + # strong positive correlation + expected = (0.82078268166812329, 0.088587005313543798) + + # correlation > 0 -> large "less" p-value + res = stats.spearmanr(x1, x2, alternative="less") + assert_approx_equal(res[0], expected[0]) + assert_approx_equal(res[1], 1 - (expected[1] / 2)) + + # correlation > 0 -> small "less" p-value + res = stats.spearmanr(x1, x2, alternative="greater") + assert_approx_equal(res[0], expected[0]) + assert_approx_equal(res[1], expected[1] / 2) + + with pytest.raises(ValueError, match="`alternative` must be 'less'..."): + stats.spearmanr(x1, x2, alternative="ekki-ekki") + + @pytest.mark.parametrize("alternative", ('two-sided', 'less', 'greater')) + def test_alternative_nan_policy(self, alternative): + # Test nan policies + x1 = [1, 2, 3, 4, 5] + x2 = [5, 6, 7, 8, 7] + x1nan = x1 + [np.nan] + x2nan = x2 + [np.nan] + + # test nan_policy="propagate" + assert_array_equal(stats.spearmanr(x1nan, x2nan), (np.nan, np.nan)) + + # test nan_policy="omit" + res_actual = stats.spearmanr(x1nan, x2nan, nan_policy='omit', + alternative=alternative) + res_expected = stats.spearmanr(x1, x2, alternative=alternative) + assert_allclose(res_actual, res_expected) + + # test nan_policy="raise" + message = 'The input contains nan values' + with pytest.raises(ValueError, match=message): + stats.spearmanr(x1nan, x2nan, nan_policy='raise', + alternative=alternative) + + # test invalid nan_policy + message = "nan_policy must be one of..." + with pytest.raises(ValueError, match=message): + stats.spearmanr(x1nan, x2nan, nan_policy='ekki-ekki', + alternative=alternative) + + +@make_xp_test_case(stats.kendalltau) +class TestKendallTau: + @pytest.mark.parametrize('x, y, statistic, pvalue', [ + # For the cases without ties, both variants should give the same + # result. Reference values result from R: + # cor.test(x, y, method="kendall", exact=1) + ([5, 2, 1, 3, 6, 4, 7, 8], [5, 2, 6, 3, 1, 8, 7, 4], 0.0, 1.0), + ([0, 5, 2, 1, 3, 6, 4, 7, 8], [5, 2, 0, 6, 3, 1, 8, 7, 4], 0.0, 1.0), + ([5, 2, 1, 3, 6, 4, 7], [5, 2, 6, 3, 1, 7, 4], -0.14285714286, 0.77261904762), + ([2, 1, 3, 6, 4, 7, 8], [2, 6, 3, 1, 8, 7, 4], 0.047619047619, 1.0), + ([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], [0, 1, 2, 3, 4, 5, 6, 7, 8, 9], + 1.0, 5.511463844797e-07), + ([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], [0, 2, 1, 3, 4, 5, 6, 7, 8, 9], + 0.9555555555555556, 5.511463844797e-06), + ([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], [0, 2, 1, 3, 4, 6, 5, 7, 8, 9], + 0.9111111111111111, 2.976190476190e-05), + ([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], [9, 8, 7, 6, 5, 4, 3, 2, 1, 0], + -1.0, 5.511463844797e-07), + ([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], [9, 7, 8, 6, 5, 4, 3, 2, 1, 0], + -0.9555555555555556, 5.511463844797e-06), + ([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], [9, 7, 8, 6, 5, 3, 4, 2, 1, 0], + -0.9111111111111111, 2.976190476190e-05), + ]) + @pytest.mark.parametrize('variant', ['b', 'c']) + def test_kendalltau_no_ties(self, x, y, statistic, pvalue, variant, xp): + res = stats.kendalltau(xp.asarray(x), xp.asarray(y), variant=variant) + xp_assert_close(res[0], xp.asarray(statistic)) + xp_assert_close(res[1], xp.asarray(pvalue)) + + def test_kendalltau_ties(self, xp): + # Check a case where variants are different + # Example values found from Kendall (1970). + # P-value is the same for the both variants + x = xp.asarray([1, 2, 2, 4, 4, 6, 6, 8, 9, 9]) + y = xp.asarray([1, 2, 4, 4, 4, 4, 8, 8, 8, 10]) + expected = 0.85895569 + xp_assert_close(stats.kendalltau(x, y, variant='b')[0], xp.asarray(expected)) + expected = 0.825 + xp_assert_close(stats.kendalltau(x, y, variant='c')[0], xp.asarray(expected)) + + # tau-b with some ties + # Cross-check with R: + # cor.test(c(12, 2, 1, 12, 2), c(1, 4, 7, 1, 0), method="kendall", exact=FALSE) + x1 = xp.asarray([12, 2, 1, 12, 2]) + x2 = xp.asarray([1, 4, 7, 1, 0]) + expected = (-0.47140452079103173, 0.28274545993277478) + res = stats.kendalltau(x1, x2) + assert_approx_equal(res[0], xp.asarray(expected[0])) + assert_approx_equal(res[1], xp.asarray(expected[1])) + + # this should result in 1 for tau-b but not tau-c + x = xp.asarray([1, 1, 2]) + xp_assert_close(stats.kendalltau(x, x, variant='b')[0], xp.asarray(1.0)) + xp_assert_close(stats.kendalltau(x, x, variant='c')[0], xp.asarray(0.88888888)) + + def test_input_validation_edge_cases(self, xp): + x = xp.asarray([1, 2, 2, 4, 4, 6, 6, 8, 9, 9]) + y = xp.asarray([1, 2, 2, 4, 4, 4, 8, 8, 8, 10]) + + # check exception in case of ties and method='exact' requested + message = "Ties found; exact method cannot be used." + with pytest.raises(ValueError, match=message): + stats.kendalltau(x, y, method='exact') + + # check exception in case of invalid method keyword + message = "Unknown `method` 'banana' specified." + with pytest.raises(ValueError, match=message): + stats.kendalltau(x, y, method='banana') + + # check exception in case of invalid variant keyword + message = "Unknown `variant` 'rms' specified." + with pytest.raises(ValueError, match=message): + stats.kendalltau(x, y, variant='rms') + + # empty arrays provided as input + with eager_warns(SmallSampleWarning, match="One or more sample...", xp=xp): + res = stats.kendalltau(xp.asarray([]), xp.asarray([])) + xp_assert_equal(res.statistic, xp.asarray(xp.nan)) + xp_assert_equal(res.pvalue, xp.asarray(xp.nan)) + + # test too small + with pytest.warns(SmallSampleWarning, match="One or more sample..."): + tau, p_value = stats.kendalltau(xp.asarray([0]), xp.asarray([0])) + xp_assert_equal(tau, xp.asarray(np.nan)) + xp_assert_equal(p_value, xp.asarray(np.nan)) + + # test unequal length inputs + x = np.arange(10.) + y = np.arange(20.) + message = "Array shapes are incompatible for broadcasting." + with pytest.raises(ValueError, match=message): + stats.kendalltau(x, y) + + @pytest.mark.parametrize('variant', ['b', 'c']) + def test_attributes(self, variant, xp): + x1 = xp.asarray([1, 2, 2, 4, 4, 6, 6, 8, 9, 9]) + x2 = xp.asarray([1, 2, 2, 4, 4, 4, 8, 8, 8, 10]) + # test for namedtuple attribute results + attributes = ('correlation', 'pvalue') + res = stats.kendalltau(x1, x2, variant=variant) + check_named_results(res, attributes, xp=xp) + xp_assert_equal(res.correlation, res.statistic) + + @pytest.mark.parametrize('x, y', [ + ([2, 2, 2], [2, 2, 2]), + ([2, 0, 2], [2, 2, 2]), + ([2, 2, 2], [2, 0, 2]), + ]) + @pytest.mark.parametrize('variant', ['b', 'c']) + def test_degenerate_input(self, x, y, variant, xp): + # with only ties in one or both inputs in tau-b or tau-c + x, y = xp.asarray(x), xp.asarray(y) + res = stats.kendalltau(x, y, variant=variant) + xp_assert_equal(res.statistic, xp.asarray(xp.nan)) + xp_assert_equal(res.pvalue, xp.asarray(xp.nan)) + + @pytest.mark.parametrize('dtype', [None, 'float32', 'float64']) + def test_large_arrays(self, dtype, xp): + # check with larger arrays + dtype = xp_default_dtype(xp) if dtype is None else getattr(xp, dtype) + rng = np.random.RandomState(7546) + x = np.array([rng.normal(loc=1, scale=1, size=500), + rng.normal(loc=1, scale=1, size=500)]) + corr = [[1.0, 0.3], + [0.3, 1.0]] + x = np.dot(np.linalg.cholesky(corr), x) + x0 = xp.asarray(x[0], dtype=dtype) + x1 = xp.asarray(x[1], dtype=dtype) + expected = (0.19291382765531062, 1.1337095377742629e-10) + res = stats.kendalltau(x0, x1) + xp_assert_close(res[0], xp.asarray(expected[0], dtype=dtype)) + xp_assert_close(res[1], xp.asarray(expected[1], dtype=dtype)) + + def test_nan_policy(self, xp): + # test nan_policy + x = np.arange(10.) + x[9] = np.nan + x = xp.asarray(x.tolist()) + + res = stats.kendalltau(x, x) + xp_assert_equal(res.statistic, xp.asarray(xp.nan)) + xp_assert_equal(res.pvalue, xp.asarray(xp.nan)) + + res = stats.kendalltau(x, x, nan_policy='omit') + xp_assert_close(res.statistic, xp.asarray(1.0)) + xp_assert_close(res.pvalue, xp.asarray(5.5114638e-6), rtol=1e-06) + + res = stats.kendalltau(x, x, nan_policy='omit', method='asymptotic') + xp_assert_close(res.statistic, xp.asarray(1.0)) + xp_assert_close(res.pvalue, xp.asarray(0.00017455009626808976), rtol=1e-06) + + message = "The input contains nan values" + with pytest.raises(ValueError, match=message): + stats.kendalltau(x, x, nan_policy='raise') + + message = "nan_policy must be one of..." + with pytest.raises(ValueError, match=message): + stats.kendalltau(x, x, nan_policy='foobar') + + def test_kendalltau_nan_2nd_arg_gh6134(self, xp): + # regression test for gh-6134: nans in the second arg were not handled + x = xp.asarray([1., 2., 3., 4.]) + y = xp.asarray([np.nan, 2.4, 3.4, 3.4]) + + r1 = stats.kendalltau(x, y, nan_policy='omit') + r2 = stats.kendalltau(x[1:], y[1:]) + xp_assert_close(r1.statistic, r2.statistic, atol=1e-15) + + @pytest.mark.thread_unsafe(reason="fails in parallel") + def test_kendalltau_gh18139_overflow(self, xp): + # gh-18139 reported an overflow in `kendalltau` that appeared after + # SciPy 0.15.1. Check that this particular overflow does not occur. + # (Test would fail if warning were emitted.) + import random + random.seed(6272161) + classes = [1, 2, 3, 4, 5, 6, 7] + n_samples = 2 * 10 ** 5 + x = xp.asarray(random.choices(classes, k=n_samples)) + y = xp.asarray(random.choices(classes, k=n_samples)) + res = stats.kendalltau(x, y) + # Reference value from SciPy 0.15.1 + xp_assert_close(res.statistic, xp.asarray(0.0011816493905730343)) + # Reference p-value from `permutation_test` w/ n_resamples=9999 (default). + # Expected to be accurate to at least two digits. + xp_assert_close(res.pvalue, xp.asarray(0.4894), atol=2e-3) + + def test_kendall_tau_large_nan_policy(self, xp): + n = 172 + # Test omit policy + x = np.arange(n + 1).astype(float) + y = np.arange(n + 1).astype(float) + y[-1] = np.nan + x, y = xp.asarray(x.tolist()), xp.asarray(y.tolist()) + _, pval = stats.kendalltau(x, y, method='exact', nan_policy='omit') + xp_assert_equal(pval, xp.asarray(0.0)) + + +@make_xp_test_case(stats.kendalltau) +class TestKendallTauAlternative: + def test_kendalltau_alternative_asymptotic(self, xp): + # Test alternative parameter, asymptotic method (due to tie) + + # Based on TestCorrSpearman2::test_alternative + x1 = [1, 2, 3, 4, 5] + x2 = [5, 6, 7, 8, 7] + + # strong positive correlation + expected = stats.kendalltau(xp.asarray(x1), xp.asarray(x2), + alternative="two-sided") + assert expected[0] > 0 + + # rank correlation > 0 -> large "less" p-value + res = stats.kendalltau(xp.asarray(x1), xp.asarray(x2), alternative="less") + xp_assert_equal(res[0], xp.asarray(expected[0])) + xp_assert_close(res[1], xp.asarray(1 - (expected[1] / 2))) + + # rank correlation > 0 -> small "greater" p-value + res = stats.kendalltau(xp.asarray(x1), xp.asarray(x2), alternative="greater") + xp_assert_equal(res[0], xp.asarray(expected[0])) + xp_assert_close(res[1], xp.asarray(expected[1] / 2)) + + # reverse the direction of rank correlation + x2.reverse() + + # strong negative correlation + expected = stats.kendalltau(xp.asarray(x1), xp.asarray(x2), + alternative="two-sided") + assert expected[0] < 0 + + # rank correlation < 0 -> large "greater" p-value + res = stats.kendalltau(xp.asarray(x1), xp.asarray(x2), alternative="greater") + xp_assert_equal(res[0], xp.asarray(expected[0])) + xp_assert_close(res[1], xp.asarray(1 - (expected[1] / 2))) + + # rank correlation < 0 -> small "less" p-value + res = stats.kendalltau(xp.asarray(x1), xp.asarray(x2), alternative="less") + xp_assert_equal(res[0], xp.asarray(expected[0])) + xp_assert_close(res[1], xp.asarray(expected[1] / 2)) + + with pytest.raises(ValueError, match="`alternative` must be 'less'..."): + stats.kendalltau(xp.asarray(x1), xp.asarray(x2), alternative="ekki-ekki") + + # There are a lot of special cases considered in the calculation of the + # exact p-value, so we test each separately. We also need to test + # separately when the observed statistic is in the left tail vs the right + # tail because the code leverages symmetry of the null distribution; to + # do that we use the same test case but negate one of the samples. + # Reference values computed using R cor.test, e.g. + # options(digits=16) + # x <- c(44.4, 45.9, 41.9, 53.3, 44.7, 44.1, 50.7, 45.2, 60.1) + # y <- c( 2.6, 3.1, 2.5, 5.0, 3.6, 4.0, 5.2, 2.8, 3.8) + # cor.test(x, y, method = "kendall", alternative = "g") + + alternatives = ('less', 'two-sided', 'greater') + p_n1 = [np.nan, np.nan, np.nan] + p_n2 = [1.0, 1.0, 0.5] + p_c0 = [1.0, 0.3333333333333, 0.1666666666667] + p_c1 = [0.9583333333333, 0.3333333333333, 0.1666666666667] + p_no_correlation = [0.5916666666667, 1.0, 0.5916666666667] + p_no_correlationb = [0.5475694444444, 1.0, 0.5475694444444] + p_n_lt_171 = [0.9624118165785, 0.1194389329806, 0.0597194664903] + p_n_lt_171b = [0.246236925303, 0.4924738506059, 0.755634083327] + p_n_lt_171c = [0.9847475308925, 0.03071385306533, 0.01535692653267] + + def exact_test(self, x, y, alternative, rev, stat_expected, p_expected, xp): + x, y = xp.asarray(x), xp.asarray(y) + if rev: + y = -y + stat_expected *= -1 + res = stats.kendalltau(x, y, method='exact', alternative=alternative) + xp_assert_close(res.statistic, xp.asarray(stat_expected)) + xp_assert_close(res.pvalue, xp.asarray(p_expected)) + + case_R_n1 = (list(zip(alternatives, p_n1, [False]*3)) + + list(zip(alternatives, reversed(p_n1), [True]*3))) + + @pytest.mark.parametrize("alternative, p_expected, rev", case_R_n1) + def test_against_R_n1(self, alternative, p_expected, rev, xp): + x, y = [1], [2] + stat_expected = np.nan + with pytest.warns(SmallSampleWarning, match="One or more sample..."): + self.exact_test(x, y, alternative, rev, stat_expected, p_expected, xp=xp) + + case_R_n2 = (list(zip(alternatives, p_n2, [False]*3)) + + list(zip(alternatives, reversed(p_n2), [True]*3))) + + @pytest.mark.parametrize("alternative, p_expected, rev", case_R_n2) + def test_against_R_n2(self, alternative, p_expected, rev, xp): + x, y = [1, 2], [3, 4] + stat_expected = 0.9999999999999998 + self.exact_test(x, y, alternative, rev, stat_expected, p_expected, xp=xp) + + case_R_c0 = (list(zip(alternatives, p_c0, [False]*3)) + + list(zip(alternatives, reversed(p_c0), [True]*3))) + + @pytest.mark.parametrize("alternative, p_expected, rev", case_R_c0) + def test_against_R_c0(self, alternative, p_expected, rev, xp): + x, y = [1, 2, 3], [1, 2, 3] + stat_expected = 1. + self.exact_test(x, y, alternative, rev, stat_expected, p_expected, xp=xp) + + case_R_c1 = (list(zip(alternatives, p_c1, [False]*3)) + + list(zip(alternatives, reversed(p_c1), [True]*3))) + + @pytest.mark.parametrize("alternative, p_expected, rev", case_R_c1) + def test_against_R_c1(self, alternative, p_expected, rev, xp): + x, y = [1, 2, 3, 4], [1, 2, 4, 3] + stat_expected = 0.6666666666666667 + self.exact_test(x, y, alternative, rev, stat_expected, p_expected, xp=xp) + + case_R_no_corr = (list(zip(alternatives, p_no_correlation, [False]*3)) + + list(zip(alternatives, reversed(p_no_correlation), + [True]*3))) + + @pytest.mark.parametrize("alternative, p_expected, rev", case_R_no_corr) + def test_against_R_no_correlation(self, alternative, p_expected, rev, xp): + x, y = [1, 2, 3, 4, 5], [1, 5, 4, 2, 3] + stat_expected = 0. + self.exact_test(x, y, alternative, rev, stat_expected, p_expected, xp=xp) + + case_no_cor_b = (list(zip(alternatives, p_no_correlationb, [False]*3)) + + list(zip(alternatives, reversed(p_no_correlationb), + [True]*3))) + + @pytest.mark.parametrize("alternative, p_expected, rev", case_no_cor_b) + def test_against_R_no_correlationb(self, alternative, p_expected, rev, xp): + x, y = [1, 2, 3, 4, 5, 6, 7, 8], [8, 6, 1, 3, 2, 5, 4, 7] + stat_expected = 0. + self.exact_test(x, y, alternative, rev, stat_expected, p_expected, xp=xp) + + case_R_lt_171 = (list(zip(alternatives, p_n_lt_171, [False]*3)) + + list(zip(alternatives, reversed(p_n_lt_171), [True]*3))) + + @pytest.mark.parametrize("alternative, p_expected, rev", case_R_lt_171) + def test_against_R_lt_171(self, alternative, p_expected, rev, xp): + # Data from Hollander & Wolfe (1973), p. 187f. + # Used from https://rdrr.io/r/stats/cor.test.html + x = [44.4, 45.9, 41.9, 53.3, 44.7, 44.1, 50.7, 45.2, 60.1] + y = [2.6, 3.1, 2.5, 5.0, 3.6, 4.0, 5.2, 2.8, 3.8] + stat_expected = 0.4444444444444445 + self.exact_test(x, y, alternative, rev, stat_expected, p_expected, xp=xp) + + case_R_lt_171b = (list(zip(alternatives, p_n_lt_171b, [False]*3)) + + list(zip(alternatives, reversed(p_n_lt_171b), + [True]*3))) + + @pytest.mark.parametrize("alternative, p_expected, rev", case_R_lt_171b) + def test_against_R_lt_171b(self, alternative, p_expected, rev, xp): + rng = np.random.RandomState(0) + x = rng.rand(100).tolist() + y = rng.rand(100).tolist() + stat_expected = -0.04686868686868687 + self.exact_test(x, y, alternative, rev, stat_expected, p_expected, xp=xp) + + case_R_lt_171c = (list(zip(alternatives, p_n_lt_171c, [False]*3)) + + list(zip(alternatives, reversed(p_n_lt_171c), + [True]*3))) + + @pytest.mark.parametrize("alternative, p_expected, rev", case_R_lt_171c) + def test_against_R_lt_171c(self, alternative, p_expected, rev, xp): + rng = np.random.RandomState(0) + x = rng.rand(170).tolist() + y = rng.rand(170).tolist() + stat_expected = 0.1115906717716673 + self.exact_test(x, y, alternative, rev, stat_expected, p_expected, xp=xp) + + case_gt_171 = (list(zip(alternatives, [False]*3)) + + list(zip(alternatives, [True]*3))) + + @pytest.mark.parametrize("alternative, rev", case_gt_171) + def test_gt_171(self, alternative, rev, xp): + rng = np.random.RandomState(0) + x = xp.asarray(rng.rand(400).tolist()) + y = xp.asarray(rng.rand(400).tolist()) + res0 = stats.kendalltau(x, y, method='exact', alternative=alternative) + res1 = stats.kendalltau(x, y, method='asymptotic', alternative=alternative) + xp_assert_equal(res0[0], res1[0]) + xp_assert_close(res0[1], res1[1], rtol=1e-3) + + @pytest.mark.parametrize("method", ('exact', 'asymptotic')) + @pytest.mark.parametrize("alternative", ('two-sided', 'less', 'greater')) + def test_nan_policy(self, method, alternative, xp): + # Test nan policies + x1 = [1, 2, 3, 4, 5] + x2 = [5, 6, 7, 8, 9] + x1nan = x1 + [np.nan] + x2nan = x2 + [np.nan] + x1, x2 = xp.asarray(x1), xp.asarray(x2) + x1nan, x2nan = xp.asarray(x1nan), xp.asarray(x2nan) + + # test nan_policy="propagate" + res_actual = stats.kendalltau(x1nan, x2nan, + method=method, alternative=alternative) + xp_assert_equal(res_actual.statistic, xp.asarray(xp.nan)) + xp_assert_equal(res_actual.pvalue, xp.asarray(xp.nan)) + + # test nan_policy="omit" + res_actual = stats.kendalltau(x1nan, x2nan, nan_policy='omit', + method=method, alternative=alternative) + res_expected = stats.kendalltau(x1, x2, method=method, + alternative=alternative) + xp_assert_equal(res_actual.statistic, res_expected.statistic) + xp_assert_equal(res_actual.pvalue, res_expected.pvalue) + + # test nan_policy="raise" + message = 'The input contains nan values' + with pytest.raises(ValueError, match=message): + stats.kendalltau(x1nan, x2nan, nan_policy='raise', + method=method, alternative=alternative) + + # test invalid nan_policy + message = "nan_policy must be one of..." + with pytest.raises(ValueError, match=message): + stats.kendalltau(x1nan, x2nan, nan_policy='ekki-ekki', + method=method, alternative=alternative) + + +class TestWeightedTau: + def test_weightedtau(self): + x = [12, 2, 1, 12, 2] + y = [1, 4, 7, 1, 0] + tau, p_value = stats.weightedtau(x, y) + assert_approx_equal(tau, -0.56694968153682723) + assert_equal(np.nan, p_value) + tau, p_value = stats.weightedtau(x, y, additive=False) + assert_approx_equal(tau, -0.62205716951801038) + assert_equal(np.nan, p_value) + # This must be exactly Kendall's tau + tau, p_value = stats.weightedtau(x, y, weigher=lambda x: 1) + assert_approx_equal(tau, -0.47140452079103173) + assert_equal(np.nan, p_value) + + # test for namedtuple attribute results + res = stats.weightedtau(x, y) + attributes = ('correlation', 'pvalue') + check_named_results(res, attributes) + assert_equal(res.correlation, res.statistic) + + # Asymmetric, ranked version + tau, p_value = stats.weightedtau(x, y, rank=None) + assert_approx_equal(tau, -0.4157652301037516) + assert_equal(np.nan, p_value) + tau, p_value = stats.weightedtau(y, x, rank=None) + assert_approx_equal(tau, -0.7181341329699029) + assert_equal(np.nan, p_value) + tau, p_value = stats.weightedtau(x, y, rank=None, additive=False) + assert_approx_equal(tau, -0.40644850966246893) + assert_equal(np.nan, p_value) + tau, p_value = stats.weightedtau(y, x, rank=None, additive=False) + assert_approx_equal(tau, -0.83766582937355172) + assert_equal(np.nan, p_value) + tau, p_value = stats.weightedtau(x, y, rank=False) + assert_approx_equal(tau, -0.51604397940261848) + assert_equal(np.nan, p_value) + # This must be exactly Kendall's tau + tau, p_value = stats.weightedtau(x, y, rank=True, weigher=lambda x: 1) + assert_approx_equal(tau, -0.47140452079103173) + assert_equal(np.nan, p_value) + tau, p_value = stats.weightedtau(y, x, rank=True, weigher=lambda x: 1) + assert_approx_equal(tau, -0.47140452079103173) + assert_equal(np.nan, p_value) + # Test argument conversion + tau, p_value = stats.weightedtau(np.asarray(x, dtype=np.float64), y) + assert_approx_equal(tau, -0.56694968153682723) + tau, p_value = stats.weightedtau(np.asarray(x, dtype=np.int16), y) + assert_approx_equal(tau, -0.56694968153682723) + tau, p_value = stats.weightedtau(np.asarray(x, dtype=np.float64), + np.asarray(y, dtype=np.float64)) + assert_approx_equal(tau, -0.56694968153682723) + # All ties + with pytest.warns(SmallSampleWarning, match="One or more sample..."): + tau, p_value = stats.weightedtau([], []) + assert_equal(np.nan, tau) + assert_equal(np.nan, p_value) + with pytest.warns(SmallSampleWarning, match="One or more sample..."): + tau, p_value = stats.weightedtau([0], [0]) + assert_equal(np.nan, tau) + assert_equal(np.nan, p_value) + # Size mismatches + assert_raises(ValueError, stats.weightedtau, [0, 1], [0, 1, 2]) + assert_raises(ValueError, stats.weightedtau, [0, 1], [0, 1], [0, 1, 2]) + # NaNs + x = [12, 2, 1, 12, 2] + y = [1, 4, 7, 1, np.nan] + tau, p_value = stats.weightedtau(x, y) + assert_approx_equal(tau, -0.56694968153682723) + x = [12, 2, np.nan, 12, 2] + tau, p_value = stats.weightedtau(x, y) + assert_approx_equal(tau, -0.56694968153682723) + # NaNs when the dtype of x and y are all np.float64 + x = [12.0, 2.0, 1.0, 12.0, 2.0] + y = [1.0, 4.0, 7.0, 1.0, np.nan] + tau, p_value = stats.weightedtau(x, y) + assert_approx_equal(tau, -0.56694968153682723) + x = [12.0, 2.0, np.nan, 12.0, 2.0] + tau, p_value = stats.weightedtau(x, y) + assert_approx_equal(tau, -0.56694968153682723) + # NaNs when there are more than one NaN in x or y + x = [12.0, 2.0, 1.0, 12.0, 1.0] + y = [1.0, 4.0, 7.0, 1.0, 1.0] + tau, p_value = stats.weightedtau(x, y) + assert_approx_equal(tau, -0.6615242347139803) + x = [12.0, 2.0, np.nan, 12.0, np.nan] + tau, p_value = stats.weightedtau(x, y) + assert_approx_equal(tau, -0.6615242347139803) + y = [np.nan, 4.0, 7.0, np.nan, np.nan] + tau, p_value = stats.weightedtau(x, y) + assert_approx_equal(tau, -0.6615242347139803) + + + def test_segfault_issue_9710(self): + # https://github.com/scipy/scipy/issues/9710 + # This test was created to check segfault + # In issue SEGFAULT only repros in optimized builds after calling + # the function twice + message = "One or more sample arguments is too small" + with pytest.warns(SmallSampleWarning, match=message): + stats.weightedtau([1], [1.0]) + stats.weightedtau([1], [1.0]) + # The code below also caused SEGFAULT + stats.weightedtau([np.nan], [52]) + + + def test_weightedtau_vs_quadratic(self): + # Trivial quadratic implementation, all parameters mandatory + def wkq(x, y, rank, weigher, add): + tot = conc = disc = u = v = 0 + for (i, j) in product(range(len(x)), range(len(x))): + w = weigher(rank[i]) + weigher(rank[j]) if add \ + else weigher(rank[i]) * weigher(rank[j]) + tot += w + if x[i] == x[j]: + u += w + if y[i] == y[j]: + v += w + if x[i] < x[j] and y[i] < y[j] or x[i] > x[j] and y[i] > y[j]: + conc += w + elif x[i] < x[j] and y[i] > y[j] or x[i] > x[j] and y[i] < y[j]: + disc += w + return (conc - disc) / np.sqrt(tot - u) / np.sqrt(tot - v) + + def weigher(x): + return 1. / (x + 1) + + rng = np.random.default_rng(42) + for s in range(3,10): + a = [] + # Generate rankings with ties + for i in range(s): + a += [i]*i + b = list(a) + rng.shuffle(a) + rng.shuffle(b) + # First pass: use element indices as ranks + rank = np.arange(len(a), dtype=np.intp) + for _ in range(2): + for add in [True, False]: + expected = wkq(a, b, rank, weigher, add) + actual = stats.weightedtau(a, b, rank, weigher, add).statistic + assert_approx_equal(expected, actual) + # Second pass: use a random rank + rng.shuffle(rank) + + +@make_xp_test_case(stats.linregress) +class TestRegression: + def test_linregressBIGX(self, xp): + # W.II.F. Regress BIG on X. + dtype = xp.float64 # needed for accurate results + result = stats.linregress(xp.asarray(X), xp.asarray(BIG)) + xp_assert_close(result.intercept, xp.asarray(99999990., dtype=dtype)) + xp_assert_close(result.rvalue, xp.asarray(1.0, dtype=dtype)) + # The uncertainty ought to be almost zero + # since all points lie on a line + zero = xp.asarray(0.0, dtype=dtype) + xp_assert_close(result.stderr, zero, atol=1e-7) # atol for JAX + xp_assert_close(result.intercept_stderr, zero, atol=1e-7) + + def test_regressXX(self, xp): + # W.IV.B. Regress X on X. + # The constant should be exactly 0 and the regression coefficient + # should be 1. This is a perfectly valid regression and the + # program should not complain. + result = stats.linregress(xp.asarray(X.tolist()), xp.asarray(X.tolist())) + xp_assert_close(result.intercept, xp.asarray(0.0), atol=1e-15) + xp_assert_close(result.rvalue, xp.asarray(1.0)) + # The uncertainly on regression through two points ought to be 0 + xp_assert_close(result.stderr, xp.asarray(0.0), atol=1e-7) + xp_assert_close(result.intercept_stderr, xp.asarray(0.0), atol=1e-7) + + # W.IV.C. Regress X on BIG and LITTLE (two predictors). The program + # should tell you that this model is "singular" because BIG and + # LITTLE are linear combinations of each other. Cryptic error + # messages are unacceptable here. Singularity is the most + # fundamental regression error. + # + # Need to figure out how to handle multiple linear regression. + # This is not obvious + + def test_regressZEROX(self, xp): + # W.IV.D. Regress ZERO on X. + # The program should inform you that ZERO has no variance or it should + # go ahead and compute the regression and report a correlation and + # total sum of squares of exactly 0. + result = stats.linregress(xp.asarray(X.tolist()), xp.asarray(ZERO.tolist())) + xp_assert_close(result.intercept, xp.asarray(0.0)) + with eager_warns(stats.ConstantInputWarning, match="An input array...", xp=xp): + ref_rvalue = stats.pearsonr(xp.asarray(X.tolist()), + xp.asarray(ZERO.tolist())).statistic + xp_assert_close(result.rvalue, ref_rvalue) + + def test_regress_simple(self, xp): + # Regress a line with sinusoidal noise. + x = xp.linspace(0, 100, 100) + y = 0.2 * xp.linspace(0, 100, 100) + 10 + y = y + xp.sin(xp.linspace(0, 20, 100)) + + result = stats.linregress(x, y) + lr = LinregressResult + assert isinstance(result, lr) + xp_assert_close(result.stderr, xp.asarray(2.3957814497838803e-3)) + + def test_regress_alternative(self, xp): + # test alternative parameter + x = xp.linspace(0, 100, 100) + y = 0.2 * xp.linspace(0, 100, 100) + 10 # slope is greater than zero + y = y + xp.sin(xp.linspace(0, 20, 100)) + + with pytest.raises(ValueError, match="`alternative` must be 'less'..."): + stats.linregress(x, y, alternative="ekki-ekki") + + res1 = stats.linregress(x, y, alternative="two-sided") + + # slope is greater than zero, so "less" p-value should be large + res2 = stats.linregress(x, y, alternative="less") + xp_assert_close(res2.pvalue, 1 - (res1.pvalue / 2)) + + # slope is greater than zero, so "greater" p-value should be small + res3 = stats.linregress(x, y, alternative="greater") + xp_assert_close(res3.pvalue, res1.pvalue / 2) + + assert res1.rvalue == res2.rvalue == res3.rvalue + + @pytest.mark.parametrize("dtype", [None, 'float32', 'float64']) + def test_regress_against_R(self, dtype, xp): + # test against R `lm` + # options(digits=16) + # x <- c(151, 174, 138, 186, 128, 136, 179, 163, 152, 131) + # y <- c(63, 81, 56, 91, 47, 57, 76, 72, 62, 48) + # relation <- lm(y~x) + # print(summary(relation)) + dtype = dtype if dtype is None else getattr(xp, dtype) + x = xp.asarray([151, 174, 138, 186, 128, 136, 179, 163, 152, 131], dtype=dtype) + y = xp.asarray([63, 81, 56, 91, 47, 57, 76, 72, 62, 48], dtype=dtype) + res = stats.linregress(x, y, alternative="two-sided") + # expected values from R's `lm` above + xp_assert_close(res.slope, xp.asarray(0.6746104491292, dtype=dtype)) + xp_assert_close(res.intercept, xp.asarray(-38.4550870760770, dtype=dtype)) + xp_assert_close(res.rvalue, xp.sqrt(xp.asarray(0.95478224775, dtype=dtype))) + xp_assert_close(res.pvalue, xp.asarray(1.16440531074e-06, dtype=dtype)) + xp_assert_close(res.stderr, xp.asarray(0.0519051424731, dtype=dtype)) + xp_assert_close(res.intercept_stderr, xp.asarray(8.0490133029927, dtype=dtype)) + + def test_linregress(self, xp): + # compared with multivariate ols with pinv + x = np.arange(11) + y = np.arange(5, 16) + y[[(1), (-2)]] -= 1 + y[[(0), (-1)]] += 1 + + result = stats.linregress(xp.asarray(x), xp.asarray(y)) + + # This test used to use 'assert_array_almost_equal' but its + # formulation got confusing since LinregressResult became + # _lib._bunch._make_tuple_bunch instead of namedtuple + # (for backwards compatibility, see PR #12983) + def assert_ae(x, y): + return xp_assert_close(x, xp.asarray(y), atol=1e-14) + assert_ae(result.slope, 1.0) + assert_ae(result.intercept, 5.0) + assert_ae(result.rvalue, 0.98229948625750) + assert_ae(result.pvalue, 7.45259691e-008) + assert_ae(result.stderr, 0.063564172616372733) + assert_ae(result.intercept_stderr, 0.37605071654517686) + + def test_regress_simple_negative_cor(self, xp): + # If the slope of the regression is negative the factor R tend + # to -1 not 1. Sometimes rounding errors makes it < -1 + # leading to stderr being NaN. + a, n = 1e-71, 100000 + dtype = xp.float64 + x = xp.linspace(a, 2 * a, n, dtype=dtype) # doesn't pass with float32 + y = xp.linspace(2 * a, a, n, dtype=dtype) + result = stats.linregress(x, y) + + # Make sure propagated numerical errors + # did not bring rvalue below -1 (or were coerced) + assert result.rvalue >= -1 + xp_assert_close(result.rvalue, xp.asarray(-1., dtype=dtype)) + + # slope and intercept stderror should stay numeric + assert not xp.isnan(result.stderr) + assert not xp.isnan(result.intercept_stderr) + + def test_linregress_result_attributes(self, xp): + x = xp.linspace(0, 100, 100) + y = 0.2 * xp.linspace(0, 100, 100) + 10 + y = y + xp.sin(xp.linspace(0, 20, 100)) + result = stats.linregress(x, y) + + # Result is of a correct class + lr = LinregressResult + assert isinstance(result, lr) + + # LinregressResult elements have correct names + attributes = ('slope', 'intercept', 'rvalue', 'pvalue', 'stderr') + check_named_results(result, attributes, xp=xp) + # Also check that the extra attribute (intercept_stderr) is present + assert 'intercept_stderr' in dir(result) + + def test_regress_two_inputs(self, xp): + # Regress a simple line formed by two points. + x = xp.arange(2) + y = xp.arange(3, 5) + result = stats.linregress(x, y) + + xp_assert_close(result.slope, xp.asarray(1.0)) + xp_assert_close(result.intercept, xp.asarray(3.0)) + xp_assert_close(result.rvalue, xp.asarray(1.0)) + + NaN = xp.asarray(xp.nan) + xp_assert_equal(result.pvalue, NaN) + xp_assert_equal(result.stderr, NaN) + xp_assert_equal(result.intercept_stderr, NaN) + + def test_regress_two_inputs_horizontal_line(self, xp): + # Regress a horizontal line formed by two points. + x = xp.arange(2.) + y = xp.ones(2) + result = stats.linregress(x, y) + + # Horizontal line + xp_assert_close(result.slope, xp.asarray(0.0)) + xp_assert_close(result.intercept, xp.asarray(1.0)) + + NaN = xp.asarray(xp.nan) + xp_assert_equal(result.rvalue, NaN) + xp_assert_equal(result.pvalue, NaN) + xp_assert_equal(result.stderr, NaN) + xp_assert_equal(result.intercept_stderr, NaN) + + def test_nist_norris(self, xp): + # If this causes a lint failure in the future, please note the history of + # requests to allow extra whitespace in table formatting (e.g. gh-12367). + # Also see https://github.com/scipy/scipy/wiki/Why-do-we-not-use-an-auto%E2%80%90formatter%3F # noqa: E501 + x = [ 0.2, 337.4, 118.2, 884.6, 10.1, 226.5, + 666.3, 996.3, 448.6, 777.0, 558.2, 0.4, + 0.6, 775.5, 666.9, 338.0, 447.5, 11.6, + 556.0, 228.1, 995.8, 887.6, 120.2, 0.3, + 0.3, 556.8, 339.1, 887.2, 999.0, 779.0, + 11.1, 118.3, 229.2, 669.1, 448.9, 0.5] + + y = [ 0.1, 338.8, 118.1, 888.0, 9.2, 228.1, + 668.5, 998.5, 449.1, 778.9, 559.2, 0.3, + 0.1, 778.1, 668.8, 339.3, 448.9, 10.8, + 557.7, 228.3, 998.0, 888.8, 119.6, 0.3, + 0.6, 557.6, 339.3, 888.0, 998.5, 778.9, + 10.2, 117.6, 228.9, 668.4, 449.2, 0.2] + + result = stats.linregress(xp.asarray(x), xp.asarray(y)) + + xp_assert_close(result.slope, xp.asarray(1.00211681802045)) + xp_assert_close(result.intercept, xp.asarray(-0.262323073774029)) + xp_assert_close(result.rvalue**2, xp.asarray(0.999993745883712)) + xp_assert_close(result.pvalue, xp.asarray(0.0), atol=1e-80) + if result.pvalue.dtype == xp.float64: # these are inaccurate with float32 + xp_assert_close(result.stderr, xp.asarray(0.00042979684820)) + xp_assert_close(result.intercept_stderr, xp.asarray(0.23281823430153)) + + def test_compare_to_polyfit(self, xp): + x = np.linspace(0, 100, 100) + y = 0.2 * np.linspace(0, 100, 100) + 10 + y += np.sin(np.linspace(0, 20, 100)) + result = stats.linregress(xp.asarray(x), xp.asarray(y)) + poly = np.polyfit(x, y, 1) # Fit 1st degree polynomial + + # Make sure linear regression slope and intercept + # match with results from numpy polyfit + xp_assert_close(result.slope, xp.asarray(poly[0])) + xp_assert_close(result.intercept, xp.asarray(poly[1])) + + def test_linregress_two_points_nan_inference(self, xp): + # Test for gh-24684 + x = xp.asarray([0., 1.]) + y = xp.asarray([0., 1.]) + + res = stats.linregress(x, y) + + NaN = xp.asarray(xp.nan) + xp_assert_equal(res.pvalue, NaN) + xp_assert_equal(res.stderr, NaN) + xp_assert_equal(res.intercept_stderr, NaN) + + # Point estimates should still be correct + xp_assert_close(res.slope, xp.asarray(1.0)) + xp_assert_close(res.intercept, xp.asarray(0.0)) + + def test_empty_input(self, xp): + with eager_warns(SmallSampleWarning, match="One or more sample...", xp=xp): + res = stats.linregress(xp.asarray([]), xp.asarray([])) + NaN = xp.asarray(xp.nan) + xp_assert_equal(res.slope, NaN) + xp_assert_equal(res.intercept, NaN) + xp_assert_equal(res.rvalue, NaN) + xp_assert_equal(res.pvalue, NaN) + xp_assert_equal(res.stderr, NaN) + xp_assert_equal(res.intercept_stderr, NaN) + + def test_nan_input(self, xp): + x = np.arange(10.) + x[9] = np.nan + x = xp.asarray(x.tolist()) + + with np.errstate(invalid="ignore"): + result = stats.linregress(x, x) + + # Make sure the result still comes back as `LinregressResult` + lr = LinregressResult + assert isinstance(result, lr) + NaN = xp.asarray(xp.nan) + xp_assert_equal(result.slope, NaN) + xp_assert_equal(result.intercept, NaN) + xp_assert_equal(result.rvalue, NaN) + xp_assert_equal(result.pvalue, NaN) + xp_assert_equal(result.stderr, NaN) + xp_assert_equal(result.intercept_stderr, NaN) + + def test_identical_x(self, xp): + rng = np.random.default_rng(7872425088) + x = xp.zeros(10) + y = xp.asarray(rng.random(10).tolist()) + with np.errstate(invalid='ignore'): + res = stats.linregress(x, y) + NaN = xp.asarray(xp.nan) + xp_assert_equal(res.slope, NaN) + xp_assert_equal(res.intercept, NaN) + xp_assert_equal(res.rvalue, NaN) + xp_assert_equal(res.pvalue, NaN) + xp_assert_equal(res.stderr, NaN) + xp_assert_equal(res.intercept_stderr, NaN) + + +class HistFunctionsTest: + def test_input_validation(self, xp): + histfun = getattr(stats, self.histfun) + a = xp.arange(20) + numbins = 10 + defaultreallimits = (0, 20) + weights = xp.ones(20) + + message = "If specified, `defaultreallimits` must be..." + with pytest.raises(ValueError, match=message): + histfun(a, numbins, 10) + with pytest.raises(ValueError, match=message): + histfun(a, numbins, (1, 2, 3)) + with pytest.raises(ValueError, match=message): + histfun(a, numbins, defaultreallimits[::-1]) + + message = r"`a` and \(if specified\) `weights` must have real dtype." + with pytest.raises(ValueError, match=message): + histfun(xp.astype(a, xp.complex128), numbins, defaultreallimits) + with pytest.raises(ValueError, match=message): + histfun(a, numbins, defaultreallimits, xp.astype(weights, xp.complex128)) + + if not is_lazy_array(weights): + message = "All `weights` must be non-negative." + with pytest.raises(ValueError, match=message): + histfun(a, numbins, defaultreallimits, -weights) + + @pytest.mark.parametrize("dtype", [None, 'float32', 'float64']) + @pytest.mark.parametrize("edge_points", [False, True]) + @pytest.mark.parametrize("specify_limits", [False, True]) + @pytest.mark.parametrize("specify_weights", [False, True]) + def test_against_reference(self, dtype, edge_points, specify_limits, + specify_weights, xp): + rng = np.random.default_rng(5194553734) + nbins = int(rng.integers(10, 15)) + + if edge_points and not specify_limits: + pytest.skip("Limits must be specified to determine edge points.") + if edge_points and (is_jax(xp) + or (dtype != 'float64' and is_torch(xp)) + or (dtype != 'float64' and is_cupy(xp))): + pytest.skip("Some backends are having trouble with edge point arithmetic.") + + a = np.astype(rng.random(100), dtype) + + if specify_limits: + limits = np.astype(np.sort(rng.random(2)), dtype) + range_ = limits + else: + # copy-pasted from documentation to generate reference values + s = (1 / 2)*(a.max() - a.min()) / (nbins - 1) + range_ = (a.min() - s, a.max() + s) + limits = None + + if edge_points: + _, edges = np.histogram(a, bins=nbins, range=range_, density=False) + a2 = np.repeat(edges, rng.integers(5, size=edges.shape).tolist()) + a = np.concat([a, a2], dtype=dtype) + + weights = np.astype(rng.random(a.shape), dtype) if specify_weights else None + + bin_counts, ref_bins = np.histogram(a, bins=nbins, range=range_, density=False) + ref_hist, _ = np.histogram(a, bins=nbins, range=range_, + weights=weights, density=False) + + dtype = xp_default_dtype(xp) if dtype is None else getattr(xp, dtype) + histfun = getattr(stats, self.histfun) + weights = weights if weights is None else xp.asarray(weights, dtype=dtype) + res = histfun(xp.asarray(a, dtype=dtype), nbins, limits, weights) + ref_hist = xp.asarray(ref_hist, dtype=dtype) + + rtol = 1e-12 if dtype == xp.float64 else 2e-5 + if histfun == stats.cumfreq: + ref_cumcount = xp.cumulative_sum(ref_hist) + xp_assert_close(res.cumcount, ref_cumcount, rtol=rtol) + else: + den = a.shape[0] if weights is None else xp.sum(weights) + ref_frequency = xp.asarray(ref_hist/den, dtype=dtype) + xp_assert_close(res.frequency, ref_frequency, rtol=rtol) + xp_assert_close(res.lowerlimit, xp.asarray(ref_bins[0], dtype=dtype)) + xp_assert_close(res.binsize, xp.asarray(ref_bins[1] - ref_bins[0], dtype=dtype)) + xp_assert_close(res.extrapoints, xp.asarray(a.size - sum(bin_counts.tolist()))) + + +@make_xp_test_case(stats.cumfreq) +class TestCumfreq(HistFunctionsTest): + histfun = 'cumfreq' + + def test_cumfreq(self, xp): + x = xp.asarray([1, 4, 2, 1, 3, 1]) + cumfreqs, lowlim, binsize, extrapoints = stats.cumfreq(x, numbins=4) + xp_assert_equal(cumfreqs, xp.asarray([3., 4., 5., 6.])) + cumfreqs, lowlim, binsize, extrapoints = stats.cumfreq( + x, numbins=4, defaultreallimits=(1.5, 5)) + assert extrapoints == 3 + + # test for namedtuple attribute results + attributes = ('cumcount', 'lowerlimit', 'binsize', 'extrapoints') + res = stats.cumfreq(x, numbins=4, defaultreallimits=(1.5, 5)) + check_named_results(res, attributes, xp=xp) + + +@make_xp_test_case(stats.relfreq) +class TestRelfreq(HistFunctionsTest): + histfun = 'relfreq' + + def test_relfreq(self, xp): + a = xp.asarray([1, 4, 2, 1, 3, 1]) + relfreqs, lowlim, binsize, extrapoints = stats.relfreq(a, numbins=4) + xp_assert_close(relfreqs, xp.asarray([0.5, 0.16666667, 0.16666667, 0.16666667])) + + # test for namedtuple attribute results + attributes = ('frequency', 'lowerlimit', 'binsize', 'extrapoints') + res = stats.relfreq(a, numbins=4) + check_named_results(res, attributes, xp=xp) + + def test_relfreq_array_like(self): # array_like is np only + # check array_like input is accepted + a = [1, 4, 2, 1, 3, 1] + res, _, _, _ = stats.relfreq(a, numbins=4) + ref, _, _, _ = stats.relfreq(np.asarray(a), numbins=4) + xp_assert_equal(res, ref) + + +class TestScoreatpercentile: + def setup_method(self): + self.a1 = [3, 4, 5, 10, -3, -5, 6] + self.a2 = [3, -6, -2, 8, 7, 4, 2, 1] + self.a3 = [3., 4, 5, 10, -3, -5, -6, 7.0] + + def test_basic(self): + x = arange(8) * 0.5 + assert_equal(stats.scoreatpercentile(x, 0), 0.) + assert_equal(stats.scoreatpercentile(x, 100), 3.5) + assert_equal(stats.scoreatpercentile(x, 50), 1.75) + + def test_fraction(self): + scoreatperc = stats.scoreatpercentile + + # Test defaults + assert_equal(scoreatperc(list(range(10)), 50), 4.5) + assert_equal(scoreatperc(list(range(10)), 50, (2,7)), 4.5) + assert_equal(scoreatperc(list(range(100)), 50, limit=(1, 8)), 4.5) + assert_equal(scoreatperc(np.array([1, 10,100]), 50, (10,100)), 55) + assert_equal(scoreatperc(np.array([1, 10,100]), 50, (1,10)), 5.5) + + # explicitly specify interpolation_method 'fraction' (the default) + assert_equal(scoreatperc(list(range(10)), 50, interpolation_method='fraction'), + 4.5) + assert_equal(scoreatperc(list(range(10)), 50, limit=(2, 7), + interpolation_method='fraction'), + 4.5) + assert_equal(scoreatperc(list(range(100)), 50, limit=(1, 8), + interpolation_method='fraction'), + 4.5) + assert_equal(scoreatperc(np.array([1, 10,100]), 50, (10, 100), + interpolation_method='fraction'), + 55) + assert_equal(scoreatperc(np.array([1, 10,100]), 50, (1,10), + interpolation_method='fraction'), + 5.5) + + def test_lower_higher(self): + scoreatperc = stats.scoreatpercentile + + # interpolation_method 'lower'/'higher' + assert_equal(scoreatperc(list(range(10)), 50, + interpolation_method='lower'), 4) + assert_equal(scoreatperc(list(range(10)), 50, + interpolation_method='higher'), 5) + assert_equal(scoreatperc(list(range(10)), 50, (2,7), + interpolation_method='lower'), 4) + assert_equal(scoreatperc(list(range(10)), 50, limit=(2,7), + interpolation_method='higher'), 5) + assert_equal(scoreatperc(list(range(100)), 50, (1,8), + interpolation_method='lower'), 4) + assert_equal(scoreatperc(list(range(100)), 50, (1,8), + interpolation_method='higher'), 5) + assert_equal(scoreatperc(np.array([1, 10, 100]), 50, (10, 100), + interpolation_method='lower'), 10) + assert_equal(scoreatperc(np.array([1, 10, 100]), 50, limit=(10, 100), + interpolation_method='higher'), 100) + assert_equal(scoreatperc(np.array([1, 10, 100]), 50, (1, 10), + interpolation_method='lower'), 1) + assert_equal(scoreatperc(np.array([1, 10, 100]), 50, limit=(1, 10), + interpolation_method='higher'), 10) + + def test_sequence_per(self): + x = arange(8) * 0.5 + expected = np.array([0, 3.5, 1.75]) + res = stats.scoreatpercentile(x, [0, 100, 50]) + assert_allclose(res, expected) + assert_(isinstance(res, np.ndarray)) + # Test with ndarray. Regression test for gh-2861 + assert_allclose(stats.scoreatpercentile(x, np.array([0, 100, 50])), + expected) + # Also test combination of 2-D array, axis not None and array-like per + res2 = stats.scoreatpercentile(np.arange(12).reshape((3,4)), + np.array([0, 1, 100, 100]), axis=1) + expected2 = array([[0, 4, 8], + [0.03, 4.03, 8.03], + [3, 7, 11], + [3, 7, 11]]) + assert_allclose(res2, expected2) + + def test_axis(self): + scoreatperc = stats.scoreatpercentile + x = arange(12).reshape(3, 4) + + assert_equal(scoreatperc(x, (25, 50, 100)), [2.75, 5.5, 11.0]) + + r0 = [[2, 3, 4, 5], [4, 5, 6, 7], [8, 9, 10, 11]] + assert_equal(scoreatperc(x, (25, 50, 100), axis=0), r0) + + r1 = [[0.75, 4.75, 8.75], [1.5, 5.5, 9.5], [3, 7, 11]] + assert_equal(scoreatperc(x, (25, 50, 100), axis=1), r1) + + x = array([[1, 1, 1], + [1, 1, 1], + [4, 4, 3], + [1, 1, 1], + [1, 1, 1]]) + score = stats.scoreatpercentile(x, 50) + assert_equal(score.shape, ()) + assert_equal(score, 1.0) + score = stats.scoreatpercentile(x, 50, axis=0) + assert_equal(score.shape, (3,)) + assert_equal(score, [1, 1, 1]) + + def test_exception(self): + assert_raises(ValueError, stats.scoreatpercentile, [1, 2], 56, + interpolation_method='foobar') + assert_raises(ValueError, stats.scoreatpercentile, [1], 101) + assert_raises(ValueError, stats.scoreatpercentile, [1], -1) + + def test_empty(self): + assert_equal(stats.scoreatpercentile([], 50), np.nan) + assert_equal(stats.scoreatpercentile(np.array([[], []]), 50), np.nan) + assert_equal(stats.scoreatpercentile([], [50, 99]), [np.nan, np.nan]) + + +@make_xp_test_case(stats.mode) +class TestMode: + + def test_empty(self, xp): + with eager_warns(SmallSampleWarning, match=too_small_1d_not_omit, xp=xp): + vals, counts = stats.mode(xp.asarray([])) + xp_assert_equal(vals, xp.asarray(xp.nan)) + xp_assert_equal(counts, xp.asarray(0.)) + + def test_scalar(self): + vals, counts = stats.mode(4.) + assert_equal(vals, np.array([4.])) + assert_equal(counts, np.array([1])) + + def test_basic(self, xp): + data1 = xp.asarray([3, 5, 1, 10, 23, 3, 2, 6, 8, 6, 10, 6]) + vals = stats.mode(data1) + xp_assert_equal(vals[0], xp.asarray(6)) + xp_assert_equal(vals[1], xp.asarray(3)) + + def test_axes_keepdims(self, xp): + data1 = [10, 10, 30, 40] + data2 = [10, 10, 10, 10] + data3 = [20, 10, 20, 20] + data4 = [30, 30, 30, 30] + data5 = [40, 30, 30, 30] + arr = xp.asarray([data1, data2, data3, data4, data5]) + + vals = stats.mode(arr, axis=None, keepdims=True) + xp_assert_equal(vals[0], xp.asarray([[30]])) + xp_assert_equal(vals[1], xp.asarray([[8]])) + + vals = stats.mode(arr, axis=0, keepdims=True) + xp_assert_equal(vals[0], xp.asarray([[10, 10, 30, 30]])) + xp_assert_equal(vals[1], xp.asarray([[2, 3, 3, 2]])) + + vals = stats.mode(arr, axis=1, keepdims=True) + xp_assert_equal(vals[0], xp.asarray([[10], [10], [20], [30], [30]])) + xp_assert_equal(vals[1], xp.asarray([[2], [4], [3], [4], [3]])) + + def test_axes(self, xp): + data1 = [10, 10, 30, 40] + data2 = [10, 10, 10, 10] + data3 = [20, 10, 20, 20] + data4 = [30, 30, 30, 30] + data5 = [40, 30, 30, 30] + arr = xp.asarray([data1, data2, data3, data4, data5]) + + vals = stats.mode(arr, axis=None) + xp_assert_equal(vals[0], xp.asarray(30)) + xp_assert_equal(vals[1], xp.asarray(8)) + + vals = stats.mode(arr, axis=0) + xp_assert_equal(vals[0], xp.asarray([10, 10, 30, 30])) + xp_assert_equal(vals[1], xp.asarray([2, 3, 3, 2])) + + vals = stats.mode(arr, axis=1) + xp_assert_equal(vals[0], xp.asarray([10, 10, 20, 30, 30])) + xp_assert_equal(vals[1], xp.asarray([2, 4, 3, 4, 3])) + + @pytest.mark.parametrize('axis', range(-4, 0)) + def test_negative_axes_gh_15375(self, axis, xp): + rng = np.random.default_rng(7090348401) + a = xp.asarray(rng.random((10, 11, 12, 13))) + res0 = stats.mode(a, axis=a.ndim+axis) + res1 = stats.mode(a, axis=axis) + xp_assert_equal(res0.mode, res1.mode) + xp_assert_equal(res0.count, res1.count) + + def test_mode_result_attributes(self, xp): + data1 = xp.asarray([3, 5, 1, 10, 23, 3, 2, 6, 8, 6, 10, 6]) + data2 = xp.asarray([]) + actual = stats.mode(data1) + attributes = ('mode', 'count') + check_named_results(actual, attributes, xp=xp) + with eager_warns(SmallSampleWarning, match=too_small_1d_not_omit, xp=xp): + actual2 = stats.mode(data2) + check_named_results(actual2, attributes, xp=xp) + + def test_nan_propagate(self, xp): + data1 = xp.asarray([3, np.nan, 5, 1, 10, 23, 3, 2, 6, 8, 6, 10, 6]) + actual = stats.mode(data1) + xp_assert_equal(actual[0], xp.asarray(6, dtype=data1.dtype)) + xp_assert_equal(actual[1], xp.asarray(3)) + + @skip_xp_backends(eager_only=True, reason="lazy arrays don't do 'raise'.") + def test_nan_omit(self, xp): + data1 = xp.asarray([3, np.nan, 5, 1, 10, 23, 3, 2, 6, 8, 6, 10, 6]) + res = stats.mode(data1, nan_policy='omit') + xp_assert_equal(res.mode, xp.asarray(6.)) + xp_assert_equal(res.count, xp.asarray(3)) + assert_raises(ValueError, stats.mode, data1, nan_policy='raise') + assert_raises(ValueError, stats.mode, data1, nan_policy='foobar') + + @skip_xp_backends(eager_only=True, reason="lazy arrays don't do 'omit'.") + @pytest.mark.parametrize("data", [ + [3, 5, 1, 1, 3.], + [3, np.nan, 5, 1, 1, 3], + [3, 5, 1.], + [3, np.nan, 5, 1], + ]) + @pytest.mark.parametrize('keepdims', [False, True]) + def test_smallest_equal(self, data, keepdims, xp): + result = stats.mode(xp.asarray(data), nan_policy='omit', keepdims=keepdims) + if keepdims: + xp_assert_equal(result[0][0], xp.asarray(1.)) + else: + xp_assert_equal(result[0], xp.asarray(1.)) + + @pytest.mark.parametrize('axis', range(-3, 3)) + def test_mode_shape_gh_9955(self, axis, xp): + rng = np.random.default_rng(984213899) + a = xp.asarray(rng.uniform(size=(3, 4, 5))) + res = stats.mode(a, axis=axis, keepdims=False) + reference_shape = list(a.shape) + reference_shape.pop(axis) + np.testing.assert_array_equal(res.mode.shape, reference_shape) + np.testing.assert_array_equal(res.count.shape, reference_shape) + + def test_nan_policy_propagate_gh_9815(self, xp): + # mode should treat np.nan as it would any other object when + # nan_policy='propagate' + a = xp.asarray([2, np.nan, 1, np.nan]) + res = stats.mode(a) + assert xp.isnan(res.mode) and res.count == 2 + + def test_keepdims_empty(self, xp): + # test empty arrays + a = xp.zeros((1, 2, 3, 0)) + + res = stats.mode(a, axis=1, keepdims=False) + assert res.mode.shape == res.count.shape == (1, 3, 0) + + res = stats.mode(a, axis=1, keepdims=True) + assert res.mode.shape == res.count.shape == (1, 1, 3, 0) + + def test_keepdims_nonempty(selfself, xp): + # test nan_policy='propagate' + a = xp.asarray([[1, 3, 3, np.nan], [1, 1, np.nan, 1]]) + + res = stats.mode(a, axis=1, keepdims=False) + xp_assert_equal(res.mode, xp.asarray([3., 1.])) + xp_assert_equal(res.count, xp.asarray([2, 3])) + + res = stats.mode(a, axis=1, keepdims=True) + xp_assert_equal(res.mode, xp.asarray([[3.], [1.]])) + xp_assert_equal(res.count, xp.asarray([[2], [3]])) + + a = xp.asarray(a) + res = stats.mode(a, axis=None, keepdims=False) + ref = stats.mode(xp_ravel(a), keepdims=False) + xp_assert_equal(res.mode, ref.mode) + xp_assert_equal(res.count, ref.count) + assert res.mode.shape == ref.mode.shape == () + + res = stats.mode(a, axis=None, keepdims=True) + ref = stats.mode(xp_ravel(a), keepdims=True) + xp_assert_equal(xp_ravel(res.mode), xp_ravel(ref.mode)) + assert res.mode.shape == (1, 1) + xp_assert_equal(xp_ravel(res.count), xp_ravel(ref.count)) + assert res.count.shape == (1, 1) + + def test_keepdims_nan_omit(self): + # test nan_policy='omit' + a = [[1, np.nan, np.nan, np.nan, 1], + [np.nan, np.nan, np.nan, np.nan, 2], + [1, 2, np.nan, 5, 5]] + + res = stats.mode(a, axis=1, keepdims=False, nan_policy='omit') + assert_array_equal(res.mode, [1, 2, 5]) + assert_array_equal(res.count, [2, 1, 2]) + + res = stats.mode(a, axis=1, keepdims=True, nan_policy='omit') + assert_array_equal(res.mode, [[1], [2], [5]]) + assert_array_equal(res.count, [[2], [1], [2]]) + + a = np.array(a) + res = stats.mode(a, axis=None, keepdims=False, nan_policy='omit') + ref = stats.mode(a.ravel(), keepdims=False, nan_policy='omit') + assert_array_equal(res, ref) + assert res.mode.shape == ref.mode.shape == () + + res = stats.mode(a, axis=None, keepdims=True, nan_policy='omit') + ref = stats.mode(a.ravel(), keepdims=True, nan_policy='omit') + assert_equal(res.mode.ravel(), ref.mode.ravel()) + assert res.mode.shape == (1, 1) + assert_equal(res.count.ravel(), ref.count.ravel()) + assert res.count.shape == (1, 1) + + @pytest.mark.parametrize("nan_policy", ['propagate', 'omit']) + def test_gh16955(self, nan_policy): + # Check that bug reported in gh-16955 is resolved + shape = (4, 3) + data = np.ones(shape) + data[0, 0] = np.nan + res = stats.mode(a=data, axis=1, keepdims=False, nan_policy=nan_policy) + assert_array_equal(res.mode, [1, 1, 1, 1]) + assert_array_equal(res.count, [2, 3, 3, 3]) + + # Test with input from gh-16595. Support for non-numeric input + # was deprecated, so check for the appropriate error. + my_dtype = np.dtype([('asdf', np.uint8), ('qwer', np.float64, (3,))]) + test = np.zeros(10, dtype=my_dtype) + message = "Argument `a` is not....|An argument has dtype...|The DType..." + with pytest.raises(TypeError, match=message): + stats.mode(test, nan_policy=nan_policy) + + def test_gh9955(self): + # The behavior of mode with empty slices (whether the input was empty + # or all elements were omitted) was inconsistent. Test that this is + # resolved: the mode of an empty slice is NaN and the count is zero. + with pytest.warns(SmallSampleWarning, match=too_small_1d_not_omit): + res = stats.mode([]) + ref = (np.nan, 0) + assert_equal(res, ref) + + with pytest.warns(SmallSampleWarning, match=too_small_1d_omit): + res = stats.mode([np.nan], nan_policy='omit') + assert_equal(res, ref) + + a = [[10., 20., 20.], [np.nan, np.nan, np.nan]] + with pytest.warns(SmallSampleWarning, match=too_small_nd_omit): + res = stats.mode(a, axis=1, nan_policy='omit') + ref = ([20, np.nan], [2, 0]) + assert_equal(res, ref) + + res = stats.mode(a, axis=1, nan_policy='propagate') + ref = ([20, np.nan], [2, 3]) + assert_equal(res, ref) + + z = np.array([[], []]) + with pytest.warns(SmallSampleWarning, match=too_small_nd_not_omit): + res = stats.mode(z, axis=1) + ref = ([np.nan, np.nan], [0, 0]) + assert_equal(res, ref) + + @pytest.mark.filterwarnings('ignore::RuntimeWarning') # np.mean warns + @pytest.mark.parametrize('z', [np.empty((0, 1, 2)), np.empty((1, 1, 2))]) + def test_gh17214(self, z, xp): + z = xp.asarray(z) + if z.size == 0: + with pytest.warns(SmallSampleWarning, match=too_small_1d_not_omit): + res = stats.mode(z, axis=None, keepdims=True) + else: + res = stats.mode(z, axis=None, keepdims=True) + ref = xp.mean(z, axis=None, keepdims=True) + assert res[0].shape == res[1].shape == ref.shape == (1, 1, 1) + + def test_raise_non_numeric_gh18254(self): + + class ArrLike: + def __init__(self, x): + self._x = x + + def __array__(self, dtype=None, copy=None): + return self._x.astype(object) + + message = ("...only boolean and numerical dtypes..." if SCIPY_ARRAY_API + else "Cannot interpret...") + + with pytest.raises(TypeError, match=message): + stats.mode(ArrLike(np.arange(3))) + + message = ("...only boolean and numerical dtypes..." if SCIPY_ARRAY_API + else "Argument `a` is not recognized as numeric.") + with pytest.raises(TypeError, match=message): + stats.mode(np.arange(3, dtype=object)) + + +@make_xp_test_case(stats.sem) +class TestSEM: + + testcase = [1., 2., 3., 4.] + scalar_testcase = 4. + + @pytest.mark.filterwarnings("ignore:invalid value encountered in divide") + def test_sem_scalar(self, xp): + # This is not in R, so used: + # sqrt(var(testcase)*3/4)/sqrt(3) + + # y = stats.sem(self.shoes[0]) + # assert_approx_equal(y,0.775177399) + scalar_testcase = xp.asarray(self.scalar_testcase)[()] + if is_numpy(xp): + with pytest.warns(SmallSampleWarning, match=too_small_1d_not_omit): + y = stats.sem(scalar_testcase) + else: + # Other array types can emit a variety of warnings. + with warnings.catch_warnings(): + warnings.simplefilter("ignore", UserWarning) + warnings.simplefilter("ignore", RuntimeWarning) + y = stats.sem(scalar_testcase) + assert xp.isnan(y) + + def test_sem(self, xp): + testcase = xp.asarray(self.testcase) + y = stats.sem(testcase) + xp_assert_close(y, xp.asarray(0.6454972244)) + n = len(self.testcase) + xp_assert_close(stats.sem(testcase, ddof=0) * (n/(n-2))**0.5, + stats.sem(testcase, ddof=2)) + + x = xp.arange(10.) + x = xp.where(x == 9, xp.nan, x) + xp_assert_equal(stats.sem(x), xp.asarray(xp.nan)) + + @skip_xp_backends(eager_only=True, reason='lazy -> reduced nan_policy capabilities') + def test_sem_nan_policy(self, xp): + x = xp.arange(10.) + x[9] = xp.nan + xp_assert_close(stats.sem(x, nan_policy='omit'), xp.asarray(0.9128709291752769)) + message = "The input contains nan values" + with pytest.raises(ValueError, match=message): + stats.sem(x, nan_policy='raise') + message = "nan_policy must be one of..." + with pytest.raises(ValueError, match=message): + stats.sem(x, nan_policy='foobar') + + +@make_xp_test_case(stats.zmap) +class TestZmap: + + @pytest.mark.parametrize( + 'x, y', + [([1., 2., 3., 4.], [1., 2., 3., 4.]), + ([1., 2., 3.], [0., 1., 2., 3., 4.])] + ) + def test_zmap(self, x, y, xp): + # For these simple cases, calculate the expected result directly + # by using the formula for the z-score. + x, y = xp.asarray(x), xp.asarray(y) + expected = (x - xp.mean(y)) / xp.std(y, correction=0) + z = stats.zmap(x, y) + xp_assert_close(z, expected) + + def test_zmap_axis(self, xp): + # Test use of 'axis' keyword in zmap. + x = xp.asarray([[0.0, 0.0, 1.0, 1.0], + [1.0, 1.0, 1.0, 2.0], + [2.0, 0.0, 2.0, 0.0]]) + + t1 = 1.0/(2.0/3)**0.5 + t2 = 3.**0.5/3 + t3 = 2.**0.5 + + z0 = stats.zmap(x, x, axis=0) + z1 = stats.zmap(x, x, axis=1) + + z0_expected = [[-t1, -t3/2, -t3/2, 0.0], + [0.0, t3, -t3/2, t1], + [t1, -t3/2, t3, -t1]] + z1_expected = [[-1.0, -1.0, 1.0, 1.0], + [-t2, -t2, -t2, 3.**0.5], + [1.0, -1.0, 1.0, -1.0]] + z0_expected = xp.asarray(z0_expected) + z1_expected = xp.asarray(z1_expected) + + xp_assert_close(z0, z0_expected) + xp_assert_close(z1, z1_expected) + + def test_zmap_ddof(self, xp): + # Test use of 'ddof' keyword in zmap. + x = xp.asarray([[0.0, 0.0, 1.0, 1.0], + [0.0, 1.0, 2.0, 3.0]]) + + z = stats.zmap(x, x, axis=1, ddof=1) + + z0_expected = xp.asarray([-0.5, -0.5, 0.5, 0.5])/(1.0/3**0.5) + z1_expected = xp.asarray([-1.5, -0.5, 0.5, 1.5])/(5./3)**0.5 + xp_assert_close(z[0, :], z0_expected) + xp_assert_close(z[1, :], z1_expected) + + @pytest.mark.parametrize('ddof', [0, 2]) + def test_zmap_nan_policy_omit(self, ddof, xp): + # nans in `scores` are propagated, regardless of `nan_policy`. + # `nan_policy` only affects how nans in `compare` are handled. + scores = xp.asarray([-3, -1, 2, np.nan]) + compare = xp.asarray([-8, -3, 2, 7, 12, np.nan]) + z = stats.zmap(scores, compare, ddof=ddof, nan_policy='omit') + # exclude nans from compare, don't use isnan + mask since that messes up + # dask + ref = stats.zmap(scores, compare[:5], ddof=ddof) + xp_assert_close(z, ref) + + @pytest.mark.parametrize('ddof', [0, 2]) + def test_zmap_nan_policy_omit_with_axis(self, ddof, xp): + scores = xp.reshape(xp.arange(-5.0, 9.0), (2, -1)) + compare = np.reshape(np.linspace(-8, 6, 24), (2, -1)) + compare[0, 4] = np.nan + compare[0, 6] = np.nan + compare[1, 1] = np.nan + # convert from numpy since some libraries like dask + # can't handle the data-dependent shapes from the isnan masking + compare_0_notna = xp.asarray(compare[0, :][~np.isnan(compare[0, :])]) + compare_1_notna = xp.asarray(compare[1, :][~np.isnan(compare[1, :])]) + compare = xp.asarray(compare) + + z = stats.zmap(scores, compare, nan_policy='omit', axis=1, ddof=ddof) + res0 = stats.zmap(scores[0, :], compare_0_notna, + ddof=ddof) + res1 = stats.zmap(scores[1, :], compare_1_notna, + ddof=ddof) + expected = xp.stack((res0, res1)) + xp_assert_close(z, expected) + + @skip_xp_backends(eager_only=True, reason="lazy arrays don't do 'raise'.") + def test_zmap_nan_policy_raise(self, xp): + scores = xp.asarray([1, 2, 3]) + compare = xp.asarray([-8, -3, 2, 7, 12, xp.nan]) + with pytest.raises(ValueError, match='input contains nan'): + stats.zmap(scores, compare, nan_policy='raise') + + @pytest.mark.filterwarnings("ignore:divide by zero encountered:RuntimeWarning:dask") + @pytest.mark.filterwarnings("ignore:invalid value encountered:RuntimeWarning:dask") + def test_degenerate_input(self, xp): + scores = xp.arange(3) + compare = xp.ones(3) + ref = xp.asarray([-xp.inf, xp.nan, xp.inf]) + with eager_warns(RuntimeWarning, match="Precision loss occurred...", xp=xp): + res = stats.zmap(scores, compare) + xp_assert_equal(res, ref) + + @pytest.mark.filterwarnings("ignore:divide by zero encountered:RuntimeWarning:dask") + def test_complex_gh22404(self, xp): + res = stats.zmap(xp.asarray([1, 2, 3, 4]), xp.asarray([1, 1j, -1, -1j])) + ref = xp.asarray([1.+0.j, 2.+0.j, 3.+0.j, 4.+0.j]) + xp_assert_close(res, ref) + + +@make_xp_test_case(stats.zscore) +class TestZscore: + def test_zscore(self, xp): + # not in R, so tested by using: + # (testcase[i] - mean(testcase, axis=0)) / sqrt(var(testcase) * 3/4) + y = stats.zscore(xp.asarray([1, 2, 3, 4])) + desired = [-1.3416407864999, -0.44721359549996, + 0.44721359549996, 1.3416407864999] + xp_assert_close(y, xp.asarray(desired)) + + def test_zscore_axis(self, xp): + # Test use of 'axis' keyword in zscore. + x = xp.asarray([[0.0, 0.0, 1.0, 1.0], + [1.0, 1.0, 1.0, 2.0], + [2.0, 0.0, 2.0, 0.0]]) + + t1 = 1.0/(2.0/3)**0.5 + t2 = 3**0.5/3 + t3 = 2**0.5 + + z0 = stats.zscore(x, axis=0) + z1 = stats.zscore(x, axis=1) + + z0_expected = [[-t1, -t3/2, -t3/2, 0.0], + [0.0, t3, -t3/2, t1], + [t1, -t3/2, t3, -t1]] + z1_expected = [[-1.0, -1.0, 1.0, 1.0], + [-t2, -t2, -t2, 3**0.5], + [1.0, -1.0, 1.0, -1.0]] + + xp_assert_close(z0, xp.asarray(z0_expected)) + xp_assert_close(z1, xp.asarray(z1_expected)) + + def test_zscore_ddof(self, xp): + # Test use of 'ddof' keyword in zscore. + x = xp.asarray([[0.0, 0.0, 1.0, 1.0], + [0.0, 1.0, 2.0, 3.0]]) + + z = stats.zscore(x, axis=1, ddof=1) + + z0_expected = xp.asarray([-0.5, -0.5, 0.5, 0.5])/(1.0/3**0.5) + z1_expected = xp.asarray([-1.5, -0.5, 0.5, 1.5])/((5./3)**0.5) + xp_assert_close(z[0, :], z0_expected) + xp_assert_close(z[1, :], z1_expected) + + def test_zscore_nan_propagate(self, xp): + x = xp.asarray([1, 2, np.nan, 4, 5]) + z = stats.zscore(x, nan_policy='propagate') + xp_assert_equal(z, xp.full(x.shape, xp.nan)) + + def test_zscore_nan_omit(self, xp): + x = xp.asarray([1, 2, xp.nan, 4, 5]) + + z = stats.zscore(x, nan_policy='omit') + + expected = xp.asarray([-1.2649110640673518, + -0.6324555320336759, + xp.nan, + 0.6324555320336759, + 1.2649110640673518 + ]) + xp_assert_close(z, expected) + + def test_zscore_nan_omit_with_ddof(self, xp): + x = xp.asarray([xp.nan, 1.0, 3.0, 5.0, 7.0, 9.0]) + z = stats.zscore(x, ddof=1, nan_policy='omit') + expected = xp.concat([xp.asarray([xp.nan]), stats.zscore(x[1:], ddof=1)]) + xp_assert_close(z, expected) + + @skip_xp_backends(eager_only=True, reason="lazy arrays don't do 'raise'.") + def test_zscore_nan_raise(self, xp): + x = xp.asarray([1, 2, xp.nan, 4, 5]) + with pytest.raises(ValueError, match="The input contains nan..."): + stats.zscore(x, nan_policy='raise') + + def test_zscore_constant_input_1d(self, xp): + x = xp.asarray([-0.087] * 3) + with eager_warns(RuntimeWarning, match="Precision loss occurred...", xp=xp): + z = stats.zscore(x) + xp_assert_equal(z, xp.full(x.shape, xp.nan)) + + @pytest.mark.filterwarnings("ignore:invalid value encountered:RuntimeWarning:dask") + def test_zscore_constant_input_2d(self, xp): + x = xp.asarray([[10.0, 10.0, 10.0, 10.0], + [10.0, 11.0, 12.0, 13.0]]) + with eager_warns(RuntimeWarning, match="Precision loss occurred...", xp=xp): + z0 = stats.zscore(x, axis=0) + xp_assert_close(z0, xp.asarray([[xp.nan, -1.0, -1.0, -1.0], + [xp.nan, 1.0, 1.0, 1.0]])) + + with eager_warns(RuntimeWarning, match="Precision loss occurred...", xp=xp): + z1 = stats.zscore(x, axis=1) + xp_assert_equal(z1, xp.stack([xp.asarray([xp.nan, xp.nan, xp.nan, xp.nan]), + stats.zscore(x[1, :])])) + + z = stats.zscore(x, axis=None) + xp_assert_equal(z, xp.reshape(stats.zscore(xp.reshape(x, (-1,))), x.shape)) + + y = xp.ones((3, 6)) + with eager_warns(RuntimeWarning, match="Precision loss occurred...", xp=xp): + z = stats.zscore(y, axis=None) + xp_assert_equal(z, xp.full_like(y, xp.nan)) + + @pytest.mark.filterwarnings("ignore:invalid value encountered:RuntimeWarning:dask") + def test_zscore_constant_input_2d_nan_policy_omit(self, xp): + x = xp.asarray([[10.0, 10.0, 10.0, 10.0], + [10.0, 11.0, 12.0, xp.nan], + [10.0, 12.0, xp.nan, 10.0]]) + s = (3/2)**0.5 + s2 = 2**0.5 + + with eager_warns(RuntimeWarning, match="Precision loss occurred...", xp=xp): + z0 = stats.zscore(x, nan_policy='omit', axis=0) + xp_assert_close(z0, xp.asarray([[xp.nan, -s, -1.0, xp.nan], + [xp.nan, 0, 1.0, xp.nan], + [xp.nan, s, xp.nan, xp.nan]])) + + with eager_warns(RuntimeWarning, match="Precision loss occurred...", xp=xp): + z1 = stats.zscore(x, nan_policy='omit', axis=1) + xp_assert_close(z1, xp.asarray([[xp.nan, xp.nan, xp.nan, xp.nan], + [-s, 0, s, xp.nan], + [-s2/2, s2, xp.nan, -s2/2]])) + + @pytest.mark.filterwarnings("ignore:invalid value encountered:RuntimeWarning:dask") + def test_zscore_2d_all_nan_row(self, xp): + # A row is all nan, and we use axis=1. + x = xp.asarray([[np.nan, np.nan, np.nan, np.nan], + [10.0, 10.0, 12.0, 12.0]]) + z = stats.zscore(x, nan_policy='omit', axis=1) + xp_assert_close(z, xp.asarray([[np.nan, np.nan, np.nan, np.nan], + [-1.0, -1.0, 1.0, 1.0]])) + + @pytest.mark.filterwarnings("ignore:invalid value encountered:RuntimeWarning:dask") + def test_zscore_2d_all_nan(self, xp): + # The entire 2d array is nan, and we use axis=None. + y = xp.full((2, 3), xp.nan) + z = stats.zscore(y, nan_policy='omit', axis=None) + xp_assert_equal(z, y) + + @pytest.mark.parametrize('x', [np.array([]), np.zeros((3, 0, 5))]) + def test_zscore_empty_input(self, x, xp): + x = xp.asarray(x) + z = stats.zscore(x) + xp_assert_equal(z, x) + + @skip_xp_invalid_arg + def test_zscore_masked_element_0_gh19039(self, xp): + # zscore returned all NaNs when 0th element was masked. See gh-19039. + rng = np.random.default_rng(8675309) + x = rng.standard_normal(10) + mask = np.zeros_like(x) + y = np.ma.masked_array(x, mask) + y.mask[0] = True + + ref = stats.zscore(x[1:]) # compute reference from non-masked elements + assert not np.any(np.isnan(ref)) + res = stats.zscore(y) + assert_allclose(res[1:], ref) + res = stats.zscore(y, axis=None) + assert_allclose(res[1:], ref) + + y[1:] = y[1] # when non-masked elements are identical, result is nan + with pytest.warns(RuntimeWarning, match="Precision loss occurred..."): + res = stats.zscore(y) + assert_equal(res[1:], np.nan) + with pytest.warns(RuntimeWarning, match="Precision loss occurred..."): + res = stats.zscore(y, axis=None) + assert_equal(res[1:], np.nan) + + +@make_xp_test_case(stats.gzscore) +class TestGZscore: + def test_gzscore_normal_array(self, xp): + x = np.asarray([1, 2, 3, 4]) + z = stats.gzscore(xp.asarray(x)) + desired = np.log(x / stats.gmean(x)) / np.log(stats.gstd(x, ddof=0)) + xp_assert_close(z, xp.asarray(desired, dtype=xp.asarray(1.).dtype)) + + @skip_xp_invalid_arg + def test_gzscore_masked_array(self): + x = np.array([1, 2, -1, 3, 4]) + mask = [0, 0, 1, 0, 0] + mx = np.ma.masked_array(x, mask=mask) + z = stats.gzscore(mx) + desired = ([-1.526072095151, -0.194700599824, np.inf, 0.584101799472, + 1.136670895503]) + desired = np.ma.masked_array(desired, mask=mask) + assert_allclose(z.compressed(), desired.compressed()) + assert_allclose(z.mask, desired.mask) + assert isinstance(z, np.ma.MaskedArray) + + +@make_xp_test_case(stats.median_abs_deviation) +class TestMedianAbsDeviation: + def setup_method(self): + self.dat_nan = [2.20, 2.20, 2.4, 2.4, 2.5, 2.7, 2.8, 2.9, + 3.03, 3.03, 3.10, 3.37, 3.4, 3.4, 3.4, 3.5, + 3.6, 3.7, 3.7, 3.7, 3.7, 3.77, 5.28, np.nan] + self.dat = [2.20, 2.20, 2.4, 2.4, 2.5, 2.7, 2.8, 2.9, 3.03, + 3.03, 3.10, 3.37, 3.4, 3.4, 3.4, 3.5, 3.6, 3.7, + 3.7, 3.7, 3.7, 3.77, 5.28, 28.95] + + def test_median_abs_deviation(self, xp): + xp_assert_close(stats.median_abs_deviation(xp.asarray(self.dat), axis=None), + xp.asarray(0.355)) + dat = xp.reshape(xp.asarray(self.dat), (6, 4)) + mad = stats.median_abs_deviation(dat, axis=0) + mad_expected = xp.asarray([0.435, 0.5, 0.45, 0.4]) + xp_assert_close(mad, mad_expected) + + @skip_xp_backends('jax.numpy', reason='lazy -> no nan_policy="omit"') + def test_mad_nan_omit(self, xp): + mad = stats.median_abs_deviation(xp.asarray(self.dat_nan), nan_policy='omit') + xp_assert_close(mad, xp.asarray(0.34)) + + def test_axis_and_nan(self, xp): + x = xp.asarray([[1.0, 2.0, 3.0, 4.0, np.nan], + [1.0, 4.0, 5.0, 8.0, 9.0]]) + mad = stats.median_abs_deviation(x, axis=1) + xp_assert_close(mad, xp.asarray([np.nan, 3.0])) + + @skip_xp_backends('jax.numpy', reason='lazy -> no nan_policy="omit"') + def test_nan_policy_omit_with_inf(self, xp): + z = xp.asarray([1, 3, 4, 6, 99, np.nan, np.inf]) + mad = stats.median_abs_deviation(z, nan_policy='omit') + xp_assert_close(mad, xp.asarray(3.0)) + + @pytest.mark.parametrize('axis', [0, 1, 2, None]) + def test_size_zero_with_axis(self, axis, xp): + x = xp.zeros((3, 0, 4)) + context = (eager_warns(SmallSampleWarning, match='too small', xp=np) + if axis in {1, None} else contextlib.nullcontext()) + with context: + mad = stats.median_abs_deviation(x, axis=axis) + xp_assert_close(mad, xp.full_like(xp.sum(x, axis=axis), fill_value=xp.nan)) + + @pytest.mark.parametrize('nan_policy, expected', + [('omit', [np.nan, 1.5, 1.5]), + ('propagate', [np.nan, np.nan, 1.5])]) + def test_nan_policy_with_axis(self, nan_policy, expected, xp): + if nan_policy=='omit' and not is_numpy(xp): + pytest.skip("nan_policy='omit' with n-d input only supported by NumPy") + + x = xp.asarray([[np.nan, np.nan, np.nan, np.nan, np.nan, np.nan], + [1, 5, 3, 6, np.nan, np.nan], + [5, 6, 7, 9, 9, 10]]) + + context = (eager_warns(SmallSampleWarning, match="too small", xp=np) + if nan_policy == 'omit' else contextlib.nullcontext()) + with context: + mad = stats.median_abs_deviation(x, nan_policy=nan_policy, axis=1) + xp_assert_close(mad, xp.asarray(expected)) + + @pytest.mark.parametrize('axis, expected', + [(1, [2.5, 2.0, 12.0]), (None, 4.5)]) + def test_center_mean_with_nan(self, axis, expected): + # nan_policy='omit' with multidimensional input only supported by NumPy + x = np.array([[1, 2, 4, 9, np.nan], + [0, 1, 1, 1, 12], + [-10, -10, -10, 20, 20]]) + mad = stats.median_abs_deviation(x, center=np.mean, nan_policy='omit', + axis=axis) + xp_assert_close(mad, expected, rtol=1e-15, atol=1e-15) + + def test_center_not_callable(self, xp): + with pytest.raises(TypeError, match='callable'): + stats.median_abs_deviation(xp.asarray([1, 2, 3, 5]), center=99) + + +@make_xp_test_case(stats.iqr) +class TestIQR: + + @pytest.mark.parametrize('dtype', [None, 'float32', 'float64']) + def test_basic(self, dtype, xp): + x = np.arange(8) * 0.5 + np.random.shuffle(x) + dtype = xp_default_dtype(xp) if dtype is None else getattr(xp, dtype) + xp_assert_equal(stats.iqr(xp.asarray(x, dtype=dtype)), + xp.asarray(1.75, dtype=dtype)) + + def test_api(self, xp): + d = xp.ones((5, 5)) + stats.iqr(d) + stats.iqr(d, None) + stats.iqr(d, 1) + stats.iqr(d, (0, 1)) + stats.iqr(d, None, (10, 90)) + stats.iqr(d, None, (30, 20), 1.0) + stats.iqr(d, None, (25, 75), 1.5, 'propagate') + + @skip_xp_backends('jax.numpy', reason='lazy -> reduced nan_policy capabilities') + def test_api_eager(self, xp): + d = xp.ones((5, 5)) + stats.iqr(d, None, (50, 50), 'normal', 'raise', 'linear') + stats.iqr(d, None, (25, 75), -0.4, 'omit', 'lower', True) + + @pytest.mark.parametrize('x', [[], np.arange(0)]) + def test_empty(self, x, xp): + with eager_warns(SmallSampleWarning, match=too_small_1d_not_omit, xp=xp): + xp_assert_equal(stats.iqr(xp.asarray(x)), xp.asarray(xp.nan)) + + def test_constant(self, xp): + # Constant array always gives 0 + x = xp.ones((7, 4)) + zero = xp.asarray(0.0) + xp_assert_equal(stats.iqr(x), zero) + xp_assert_equal(stats.iqr(x, axis=0), xp.zeros(4)) + xp_assert_equal(stats.iqr(x, axis=1), xp.zeros(7)) + xp_assert_equal(stats.iqr(x, interpolation='linear'), zero) + xp_assert_equal(stats.iqr(x, interpolation='midpoint'), zero) + xp_assert_equal(stats.iqr(x, interpolation='nearest'), zero) + xp_assert_equal(stats.iqr(x, interpolation='lower'), zero) + xp_assert_equal(stats.iqr(x, interpolation='higher'), zero) + + # 0 only along constant dimensions + # This also tests much of `axis` + y = xp.ones((4, 5, 6)) * xp.arange(6.) + xp_assert_equal(stats.iqr(y, axis=0), xp.zeros((5, 6))) + xp_assert_equal(stats.iqr(y, axis=1), xp.zeros((4, 6))) + xp_assert_equal(stats.iqr(y, axis=2), xp.full((4, 5), 2.5)) + xp_assert_equal(stats.iqr(y, axis=(0, 1)), xp.zeros(6)) + xp_assert_equal(stats.iqr(y, axis=(0, 2)), xp.full(5, 3.)) + xp_assert_equal(stats.iqr(y, axis=(1, 2)), xp.full(4, 3.)) + + def test_scalarlike(self, xp): + x = xp.arange(1.) + 7.0 + xp_assert_equal(stats.iqr(x[0]), xp.asarray(0.0)) + xp_assert_equal(stats.iqr(x), xp.asarray(0.0)) + xp_assert_equal(stats.iqr(x, keepdims=True), xp.asarray([0.0])) + + def test_2D(self, xp): + x = xp.reshape(xp.arange(15), (3, 5)) + xp_assert_equal(stats.iqr(x), xp.asarray(7.0)) + xp_assert_equal(stats.iqr(x, axis=0), xp.full(5, 5.)) + xp_assert_equal(stats.iqr(x, axis=1), xp.full(3, 2.)) + xp_assert_equal(stats.iqr(x, axis=(0, 1)), xp.asarray(7.0)) + xp_assert_equal(stats.iqr(x, axis=(1, 0)), xp.asarray(7.0)) + + def test_axis(self, xp): + # The `axis` keyword is also put through its paces in `test_keepdims`. + o = np.random.normal(size=(71, 23)) + x = np.dstack([o] * 10) # x.shape = (71, 23, 10) + o, x = xp.asarray(o), xp.asarray(x) + q = xp.broadcast_to(stats.iqr(o), (10,)) + + xp_assert_equal(stats.iqr(x, axis=(0, 1)), q) + x = xp.moveaxis(x, -1, 0) # x.shape = (10, 71, 23) + xp_assert_equal(stats.iqr(x, axis=(2, 1)), q) + x = xp_swapaxes(x,0, 1, xp=xp) # x.shape = (71, 10, 23) + xp_assert_equal(stats.iqr(x, axis=(0, 2)), q) + x = xp_swapaxes(x,0, 1, xp=xp) # x.shape = (10, 71, 23) + + xp_assert_equal(stats.iqr(x, axis=(0, 1, 2)), + stats.iqr(x, axis=None)) + xp_assert_equal(stats.iqr(x, axis=(0,)), + stats.iqr(x, axis=0)) + + d = np.arange(3 * 5 * 7 * 11) + # Older versions of numpy only shuffle along axis=0. + # Not sure about newer, don't care. + np.random.shuffle(d) + d = d.reshape((3, 5, 7, 11)) + d = xp.asarray(d) + xp_assert_equal(stats.iqr(d, axis=(0, 1, 2))[0], + stats.iqr(xp_ravel(d[:,:,:, 0]))) + xp_assert_equal(stats.iqr(d, axis=(0, 1, 3))[1], + stats.iqr(xp_ravel(d[:,:, 1,:]))) + xp_assert_equal(stats.iqr(d, axis=(3, 1, -4))[2], + stats.iqr(xp_ravel(d[:,:, 2,:]))) + xp_assert_equal(stats.iqr(d, axis=(3, 1, 2))[2], + stats.iqr(xp_ravel(d[2,:,:,:]))) + xp_assert_equal(stats.iqr(d, axis=(3, 2))[2, 1], + stats.iqr(xp_ravel(d[2, 1,:,:]))) + xp_assert_equal(stats.iqr(d, axis=(1, -2))[2, 1], + stats.iqr(xp_ravel(d[2, :, :, 1]))) + xp_assert_equal(stats.iqr(d, axis=(1, 3))[2, 2], + stats.iqr(xp_ravel(d[2, :, 2,:]))) + + with pytest.raises(AxisError, match='`axis` is out of bounds...'): + stats.iqr(d, axis=4) + with pytest.raises(ValueError, match='`axis` must contain only...'): + stats.iqr(d, axis=(0, 0)) + + def test_rng(self, xp): + x = xp.arange(5) + xp_assert_close(stats.iqr(x), xp.asarray(2.)) + xp_assert_close(stats.iqr(x, rng=(25, 87.5)), xp.asarray(2.5)) + xp_assert_close(stats.iqr(x, rng=(12.5, 75)), xp.asarray(2.5)) + xp_assert_close(stats.iqr(x, rng=(10, 50)), xp.asarray(1.6)) # 3-1.4 + + message = r"Elements of `rng` must be in the range \[0, 100\]." + with pytest.raises(ValueError, match=message): + stats.iqr(x, rng=(0, 101)) + + message = "`rng` must not contain NaNs." + with pytest.raises(ValueError, match=message): + stats.iqr(x, rng=(np.nan, 25)) + + message = "`rng` must be a two element sequence." + with pytest.raises(TypeError, match=message): + stats.iqr(x, rng=(0, 50, 60)) + + def test_interpolation(self, xp): + x = xp.arange(5) + y = xp.arange(4) + # Default + xp_assert_equal(stats.iqr(x), xp.asarray(2.)) + xp_assert_equal(stats.iqr(y), xp.asarray(1.5)) + # Linear + xp_assert_equal(stats.iqr(x, interpolation='linear'), xp.asarray(2.)) + xp_assert_equal(stats.iqr(y, interpolation='linear'), xp.asarray(1.5)) + # Higher + xp_assert_equal(stats.iqr(x, interpolation='higher'), xp.asarray(2.)) + xp_assert_equal(stats.iqr(x, rng=(25, 80), interpolation='higher'), + xp.asarray(3.)) + xp_assert_equal(stats.iqr(y, interpolation='higher'), xp.asarray(2.)) + # Lower (will generally, but not always be the same as higher) + xp_assert_equal(stats.iqr(x, interpolation='lower'), xp.asarray(2.)) + xp_assert_equal(stats.iqr(x, rng=(25, 80), interpolation='lower'), + xp.asarray(2.)) + xp_assert_equal(stats.iqr(y, interpolation='lower'), xp.asarray(2.)) + # Nearest + xp_assert_equal(stats.iqr(x, interpolation='nearest'), xp.asarray(2.)) + xp_assert_equal(stats.iqr(y, interpolation='nearest'), xp.asarray(1.)) + # Midpoint + xp_assert_equal(stats.iqr(x, interpolation='midpoint'), xp.asarray(2.)) + xp_assert_equal(stats.iqr(x, rng=(25, 80), interpolation='midpoint'), + xp.asarray(2.5)) + xp_assert_equal(stats.iqr(y, interpolation='midpoint'), xp.asarray(2.)) + + # Check all method= values new in numpy 1.22.0 are accepted + for method in ('inverted_cdf', 'averaged_inverted_cdf', + 'closest_observation', 'interpolated_inverted_cdf', + 'hazen', 'weibull', 'median_unbiased', + 'normal_unbiased'): + stats.iqr(y, interpolation=method) + + with pytest.raises(ValueError, match='`method` must be one of...'): + stats.iqr(x, interpolation='foobar') + + def test_keepdims(self, xp): + # Also tests most of `axis` + x = xp.ones((3, 5, 7, 11)) + assert_equal(stats.iqr(x, axis=None, keepdims=False).shape, ()) + assert_equal(stats.iqr(x, axis=2, keepdims=False).shape, (3, 5, 11)) + assert_equal(stats.iqr(x, axis=(0, 1), keepdims=False).shape, (7, 11)) + assert_equal(stats.iqr(x, axis=(0, 3), keepdims=False).shape, (5, 7)) + assert_equal(stats.iqr(x, axis=(1,), keepdims=False).shape, (3, 7, 11)) + assert_equal(stats.iqr(x, (0, 1, 2, 3), keepdims=False).shape, ()) + assert_equal(stats.iqr(x, axis=(0, 1, 3), keepdims=False).shape, (7,)) + + assert_equal(stats.iqr(x, axis=None, keepdims=True).shape, (1, 1, 1, 1)) + assert_equal(stats.iqr(x, axis=2, keepdims=True).shape, (3, 5, 1, 11)) + assert_equal(stats.iqr(x, axis=(0, 1), keepdims=True).shape, (1, 1, 7, 11)) + assert_equal(stats.iqr(x, axis=(0, 3), keepdims=True).shape, (1, 5, 7, 1)) + assert_equal(stats.iqr(x, axis=(1,), keepdims=True).shape, (3, 1, 7, 11)) + assert_equal(stats.iqr(x, (0, 1, 2, 3), keepdims=True).shape, (1, 1, 1, 1)) + assert_equal(stats.iqr(x, axis=(0, 1, 3), keepdims=True).shape, (1, 1, 7, 1)) + + @skip_xp_backends('jax.numpy', reason='lazy -> reduced nan_policy capabilitiesy') + def test_nanpolicy(self, xp): + x = xp.reshape(xp.arange(15.0), (3, 5)) + + # No NaNs + xp_assert_equal(stats.iqr(x, nan_policy='propagate'), xp.asarray(7.)) + xp_assert_equal(stats.iqr(x, nan_policy='omit'), xp.asarray(7.)) + xp_assert_equal(stats.iqr(x, nan_policy='raise'), xp.asarray(7.)) + + # Yes NaNs + x = xpx.at(x)[1, 2].set(xp.nan) + xp_assert_equal(stats.iqr(x, nan_policy='propagate'), + xp.asarray(xp.nan)) + xp_assert_equal(stats.iqr(x, axis=0, nan_policy='propagate'), + xp.asarray([5, 5, xp.nan, 5, 5])) + xp_assert_equal(stats.iqr(x, axis=1, nan_policy='propagate'), + xp.asarray([2, xp.nan, 2])) + + xp_assert_equal(stats.iqr(x, nan_policy='omit'), xp.asarray(7.5)) + + message = "The input contains nan values" + with pytest.raises(ValueError, match=message): + stats.iqr(x, nan_policy='raise') + + # Bad policy + message = "nan_policy must be one of..." + with pytest.raises(ValueError, match=message): + stats.iqr(x, nan_policy='barfood') + + @skip_xp_backends(np_only=True, + reason="nan_policy w/ multidimensional arrays only available w/ NumPy") + def test_nanpolicy_nd(self, xp): + x = xp.reshape(xp.arange(15.0), (3, 5)) + x[1, 2] = xp.nan + xp_assert_equal(stats.iqr(x, axis=0, nan_policy='omit'), + xp.full(5, 5.)) + xp_assert_equal(stats.iqr(x, axis=1, nan_policy='omit'), + xp.asarray([2, 2.5, 2])) + + message = "The input contains nan values" + with pytest.raises(ValueError, match=message): + stats.iqr(x, axis=0, nan_policy='raise') + with pytest.raises(ValueError, match=message): + stats.iqr(x, axis=1, nan_policy='raise') + + def test_scale_no_nans(self, xp): + x = xp.reshape(xp.arange(15.0), (3, 5)) + xp_assert_equal(stats.iqr(x, scale=1.0), xp.asarray(7.)) + xp_assert_close(stats.iqr(x, scale='normal'), xp.asarray(7 / 1.3489795)) + xp_assert_equal(stats.iqr(x, scale=2.0), xp.asarray(3.5)) + + def test_scale_with_nans(self, xp): + x = xp.reshape(xp.arange(15.0), (3, 5)) + x = xpx.at(x)[1, 2].set(xp.nan) + nan = xp.asarray(xp.nan) + xp_assert_equal(stats.iqr(x, scale=1.0, nan_policy='propagate'), nan) + xp_assert_equal(stats.iqr(x, scale='normal', nan_policy='propagate'), nan) + xp_assert_equal(stats.iqr(x, scale=2.0, nan_policy='propagate'), nan) + + def test_invalid_scale(self, xp): + x = xp.arange(15.0) + message = "foobar not a valid scale for `iqr`" + with pytest.raises(ValueError, match=message): + stats.iqr(x, scale='foobar') + + @skip_xp_backends("jax.numpy", reason="lazy -> reduced nan_policy capabilities") + def test_scale_nan_policy_omit(self, xp): + x = xp.reshape(xp.arange(15.0), (3, 5)) + x = xpx.at(x)[1, 2].set(xp.nan) + xp_assert_equal(stats.iqr(x, scale=1.0, nan_policy='omit'), xp.asarray(7.5)) + xp_assert_close(stats.iqr(x, scale='normal', nan_policy='omit'), + xp.asarray(7.5 / 1.3489795)) + xp_assert_equal(stats.iqr(x, scale=2.0, nan_policy='omit'), xp.asarray(3.75)) + + def test_scale_nanpolicy_nd(self, xp): + # axis=1 chosen to show behavior with both nans and without + x = xp.reshape(xp.arange(15.0), (3, 5)) + x = xpx.at(x)[1, 2].set(xp.nan) + + xp_assert_equal(stats.iqr(x, axis=1, scale=1.0, nan_policy='propagate'), + xp.asarray([2, np.nan, 2])) + xp_assert_close(stats.iqr(x, axis=1, scale='normal', nan_policy='propagate'), + xp.asarray([2, np.nan, 2]) / 1.3489795) + xp_assert_equal(stats.iqr(x, axis=1, scale=2.0, nan_policy='propagate'), + xp.asarray([1, np.nan, 1])) + + def test_rng_order(self, xp): + # test that order of `rng` doesn't matter (as documented) + x = xp.arange(8.) * 0.5 + res = stats.iqr(x, rng=(75, 25)) + ref = stats.iqr(x) + xp_assert_equal(res, ref) + + +@make_xp_test_case(stats.moment) +class TestMoments: + """ + Comparison numbers are found using R v.1.5.1 + note that length(testcase) = 4 + testmathworks comes from documentation for the + Statistics Toolbox for Matlab and can be found at both + https://www.mathworks.com/help/stats/kurtosis.html + https://www.mathworks.com/help/stats/skewness.html + Note that both test cases came from here. + """ + testcase = [1., 2., 3., 4.] + scalar_testcase = 4. + rng = np.random.default_rng(2285049930) + testcase_moment_accuracy = rng.random(42) + + @pytest.mark.parametrize('size', [10, (10, 2)]) + @pytest.mark.parametrize('m, c', list(product((0, 1, 2, 3), (None, 0, 1)))) + def test_moment_center_scalar_moment(self, size, m, c, xp): + rng = np.random.default_rng(6581432544381372042) + x = xp.asarray(rng.random(size=size)) + res = stats.moment(x, m, center=c) + c = xp.mean(x, axis=0) if c is None else c + ref = xp.sum((x - c)**m, axis=0)/x.shape[0] + xp_assert_close(res, ref, atol=1e-16) + + @pytest.mark.parametrize('size', [10, (10, 2)]) + @pytest.mark.parametrize('c', (None, 0, 1)) + def test_moment_center_array_moment(self, size, c, xp): + rng = np.random.default_rng(1706828300224046506) + x = xp.asarray(rng.random(size=size)) + m = [0, 1, 2, 3] + res = stats.moment(x, m, center=c) + ref = xp.concat([stats.moment(x, i, center=c)[xp.newaxis, ...] for i in m]) + xp_assert_close(res, ref) # JAX JIT doesn't pass with equality + + def test_moment(self, xp): + # mean((testcase-mean(testcase))**power,axis=0),axis=0))**power)) + testcase = xp.asarray(self.testcase) + + y = stats.moment(xp.asarray(self.scalar_testcase)) + xp_assert_close(y, xp.asarray(0.0)) + + y = stats.moment(testcase, 0) + xp_assert_close(y, xp.asarray(1.0)) + y = stats.moment(testcase, 1) + xp_assert_close(y, xp.asarray(0.0)) + y = stats.moment(testcase, 2) + xp_assert_close(y, xp.asarray(1.25)) + y = stats.moment(testcase, 3) + xp_assert_close(y, xp.asarray(0.0)) + y = stats.moment(testcase, 4) + xp_assert_close(y, xp.asarray(2.5625)) + + # check array_like input for moment + y = stats.moment(testcase, [1, 2, 3, 4]) + xp_assert_close(y, xp.asarray([0, 1.25, 0, 2.5625])) + + # check floating point integer order + y = stats.moment(testcase, 0.0) + xp_assert_close(y, xp.asarray(1.0)) + + y = stats.moment(testcase, [1.0, 2, 3, 4.0]) + xp_assert_close(y, xp.asarray([0, 1.25, 0, 2.5625])) + + @skip_xp_backends("jax.numpy", reason="array value-dependent error") + def test_order_input_validation(self, xp): + testcase = xp.asarray(self.testcase) + message = 'All elements of `order` must be integral.' + with pytest.raises(ValueError, match=message): + stats.moment(testcase, 1.2) + + def test_zero_size(self, xp): + message = "One or more sample arguments..." + with eager_warns(SmallSampleWarning, match=message, xp=xp): + y = stats.moment(xp.asarray([])) + xp_assert_equal(y, xp.asarray(xp.nan)) + with eager_warns(SmallSampleWarning, match=message, xp=xp): + y = stats.moment(xp.asarray([], dtype=xp.float32)) + xp_assert_equal(y, xp.asarray(xp.nan, dtype=xp.float32)) + + message = "All axis-slices of one or more" + with eager_warns(SmallSampleWarning, match=message, xp=xp): + y = stats.moment(xp.asarray([[]]), axis=1) + xp_assert_equal(y, xp.asarray([xp.nan])) + + y = stats.moment(xp.zeros((1, 0)), axis=0) + xp_assert_equal(y, xp.empty((0,))) + + y = stats.moment(xp.asarray([[]]), order=[0, 1], axis=0) + xp_assert_equal(y, xp.empty((2, 0))) + + def test_nan_policy(self): + x = np.arange(10.) + x[9] = np.nan + assert_equal(stats.moment(x, 2), np.nan) + assert_almost_equal(stats.moment(x, nan_policy='omit'), 0.0) + assert_raises(ValueError, stats.moment, x, nan_policy='raise') + assert_raises(ValueError, stats.moment, x, nan_policy='foobar') + + @pytest.mark.parametrize('dtype', ['float32', 'float64', 'complex128']) + @pytest.mark.parametrize('expect, order', [(0, 1), (1, 0)]) + def test_constant_moments(self, dtype, expect, order, xp): + if dtype=='complex128' and is_torch(xp): + message = "torch can't round complex input; data-apis/array-api-compat#407" + pytest.skip(message) + dtype = getattr(xp, dtype) + rng = np.random.default_rng(3824693518) + x = xp.asarray(rng.random(5), dtype=dtype) + y = stats.moment(x, order=order) + xp_assert_equal(y, xp.asarray(expect, dtype=dtype)) + + message = "Precision loss occurred in moment calculation" + with eager_warns(RuntimeWarning, match=message, xp=xp): + y = stats.moment(xp.broadcast_to(x, (6, 5)), axis=0, order=order) + xp_assert_equal(y, xp.full((5,), expect, dtype=dtype)) + + with eager_warns(RuntimeWarning, match=message, xp=xp): + y = stats.moment(xp.broadcast_to(x, (1, 2, 3, 4, 5)), axis=2, + order=order) + xp_assert_equal(y, xp.full((1, 2, 4, 5), expect, dtype=dtype)) + + y = stats.moment(xp.broadcast_to(x, (1, 2, 3, 4, 5)), axis=None, + order=order) + xp_assert_equal(y, xp.full((), expect, dtype=dtype)) + + def test_moment_propagate_nan(self, xp): + # Check that the shape of the result is the same for inputs + # with and without nans, cf gh-5817 + a = xp.reshape(xp.arange(8.), (2, -1)) + a = xpx.at(a)[1, 0].set(xp.nan) + + mm = stats.moment(xp.asarray(a), 2, axis=1) + xp_assert_close(mm, xp.asarray([1.25, xp.nan]), atol=1e-15) + + def test_moment_empty_order(self, xp): + # tests moment with empty `order` list + with pytest.raises(ValueError, match=r"`order` must be a scalar or a" + r" non-empty 1D array."): + stats.moment(xp.asarray([1, 2, 3, 4]), order=[]) + + def test_rename_moment_order(self, xp): + # Parameter 'order' was formerly known as 'moment'. The old name + # has not been deprecated, so it must continue to work. + x = xp.arange(10) + res = stats.moment(x, moment=3) + ref = stats.moment(x, order=3) + xp_assert_equal(res, ref) + + def test_moment_accuracy(self): + # 'moment' must have a small enough error compared to the slower + # but very accurate numpy.power() implementation. + tc_no_mean = (self.testcase_moment_accuracy + - np.mean(self.testcase_moment_accuracy)) + assert_allclose(np.power(tc_no_mean, 42).mean(), + stats.moment(self.testcase_moment_accuracy, 42)) + + @pytest.mark.parametrize('order', [0, 1, 2, 3]) + @pytest.mark.parametrize('axis', [-1, 0, 1]) + @pytest.mark.parametrize('center', [None, 0.]) + def test_moment_array_api(self, xp, order, axis, center): + rng = np.random.default_rng(34823589259425) + x = rng.random(size=(5, 6, 7)) + res = stats.moment(xp.asarray(x), order, axis=axis, center=center) + center = center if center is None else np.asarray(center) + ref = xp.asarray(_moment(x, order, axis, center=center)[()]) + xp_assert_close(res, ref) + + +class SkewKurtosisTest: + scalar_testcase = 4. + testcase = [1., 2., 3., 4.] + testmathworks = [1.165, 0.6268, 0.0751, 0.3516, -0.6965] + + def test_empty_1d(self, xp): + x = xp.asarray([]) + with eager_warns(SmallSampleWarning, match=too_small_1d_not_omit, xp=xp): + res = self.stat_fun(x) + + xp_assert_equal(res, xp.asarray(xp.nan)) + + +@make_xp_test_case(stats.skew) +class TestSkew(SkewKurtosisTest): + def stat_fun(self, x): + return stats.skew(x) + + @pytest.mark.filterwarnings( + "ignore:invalid value encountered:RuntimeWarning:dask" + ) + def test_skewness(self, xp): + # Scalar test case + y = stats.skew(xp.asarray(self.scalar_testcase)) + xp_assert_close(y, xp.asarray(xp.nan)) + # sum((testmathworks-mean(testmathworks,axis=0))**3,axis=0) / + # ((sqrt(var(testmathworks)*4/5))**3)/5 + y = stats.skew(xp.asarray(self.testmathworks)) + xp_assert_close(y, xp.asarray(-0.29322304336607), atol=1e-10) + y = stats.skew(xp.asarray(self.testmathworks), bias=0) + xp_assert_close(y, xp.asarray(-0.437111105023940), atol=1e-10) + y = stats.skew(xp.asarray(self.testcase)) + xp_assert_close(y, xp.asarray(0.0), atol=1e-10) + + def test_nan_policy(self): + # initially, nan_policy is ignored with alternative backends + x = np.arange(10.) + x[9] = np.nan + with np.errstate(invalid='ignore'): + assert_equal(stats.skew(x), np.nan) + assert_equal(stats.skew(x, nan_policy='omit'), 0.) + assert_raises(ValueError, stats.skew, x, nan_policy='raise') + assert_raises(ValueError, stats.skew, x, nan_policy='foobar') + + def test_skewness_scalar(self): + # `skew` must return a scalar for 1-dim input (only for NumPy arrays) + assert_equal(stats.skew(arange(10)), 0.0) + + def test_skew_propagate_nan(self, xp): + # Check that the shape of the result is the same for inputs + # with and without nans, cf gh-5817 + a = xp.arange(8.) + a = xp.reshape(a, (2, -1)) + a = xpx.at(a)[1, 0].set(xp.nan) + with np.errstate(invalid='ignore'): + s = stats.skew(xp.asarray(a), axis=1) + xp_assert_equal(s, xp.asarray([0, xp.nan])) + + def test_skew_constant_value(self, xp): + # Skewness of a constant input should be NaN (gh-16061) + a = xp.repeat(xp.asarray([-0.27829495]), 10) + + with eager_warns(RuntimeWarning, match="Precision loss occurred", xp=xp): + xp_assert_equal(stats.skew(a), xp.asarray(xp.nan)) + xp_assert_equal(stats.skew(a*2.**50), xp.asarray(xp.nan)) + xp_assert_equal(stats.skew(a/2.**50), xp.asarray(xp.nan)) + xp_assert_equal(stats.skew(a, bias=False), xp.asarray(xp.nan)) + + # # similarly, from gh-11086: + a = xp.asarray([14.3]*7) + xp_assert_equal(stats.skew(a), xp.asarray(xp.nan)) + a = 1. + xp.arange(-3., 4)*1e-16 + xp_assert_equal(stats.skew(a), xp.asarray(xp.nan)) + + def test_precision_loss_gh15554(self, xp): + # gh-15554 was one of several issues that have reported problems with + # constant or near-constant input. We can't always fix these, but + # make sure there's a warning. + with eager_warns(RuntimeWarning, match="Precision loss occurred", xp=xp): + rng = np.random.default_rng(34095309370) + a = rng.random(size=(100, 10)) + a[:, 0] = 1.01 + stats.skew(xp.asarray(a)) + + @pytest.mark.parametrize('axis', [-1, 0, 2, None]) + @pytest.mark.parametrize('bias', [False, True]) + def test_vectorization(self, xp, axis, bias): + # Behavior with array input is barely tested above. Compare + # against naive implementation. + rng = np.random.default_rng(1283413549926) + x = xp.asarray(rng.random((3, 4, 5))) + + def skewness(a, axis, bias): + # Simple implementation of skewness + if axis is None: + a = xp.reshape(a, (-1,)) + axis = 0 + mean = xp.mean(a, axis=axis, keepdims=True) + mu3 = xp.mean((a - mean)**3, axis=axis) + std = xp.std(a, axis=axis) + res = mu3 / std ** 3 + if not bias: + n = a.shape[axis] + res *= ((n - 1.0) * n) ** 0.5 / (n - 2.0) + return res + + res = stats.skew(x, axis=axis, bias=bias) + ref = skewness(x, axis=axis, bias=bias) + xp_assert_close(res, ref) + + +@make_xp_test_case(stats.kurtosis) +class TestKurtosis(SkewKurtosisTest): + def stat_fun(self, x): + return stats.kurtosis(x) + + @pytest.mark.filterwarnings("ignore:invalid value encountered:RuntimeWarning:dask") + def test_kurtosis(self, xp): + # Scalar test case + y = stats.kurtosis(xp.asarray(self.scalar_testcase)) + assert xp.isnan(y) + + # sum((testcase-mean(testcase,axis=0))**4,axis=0) + # / ((sqrt(var(testcase)*3/4))**4) + # / 4 + # + # sum((test2-mean(testmathworks,axis=0))**4,axis=0) + # / ((sqrt(var(testmathworks)*4/5))**4) + # / 5 + # + # Set flags for axis = 0 and + # fisher=0 (Pearson's defn of kurtosis for compatibility with Matlab) + y = stats.kurtosis(xp.asarray(self.testmathworks), 0, fisher=0, bias=1) + xp_assert_close(y, xp.asarray(2.1658856802973)) + + # Note that MATLAB has confusing docs for the following case + # kurtosis(x,0) gives an unbiased estimate of Pearson's skewness + # kurtosis(x) gives a biased estimate of Fisher's skewness (Pearson-3) + # The MATLAB docs imply that both should give Fisher's + y = stats.kurtosis(xp.asarray(self.testmathworks), fisher=0, bias=0) + xp_assert_close(y, xp.asarray(3.663542721189047)) + y = stats.kurtosis(xp.asarray(self.testcase), 0, 0) + xp_assert_close(y, xp.asarray(1.64)) + + x = xp.arange(10.) + x = xp.where(x == 8, xp.nan, x) + xp_assert_equal(stats.kurtosis(x), xp.asarray(xp.nan)) + + def test_kurtosis_nan_policy(self): + # nan_policy only for NumPy right now + x = np.arange(10.) + x[9] = np.nan + assert_almost_equal(stats.kurtosis(x, nan_policy='omit'), -1.230000) + assert_raises(ValueError, stats.kurtosis, x, nan_policy='raise') + assert_raises(ValueError, stats.kurtosis, x, nan_policy='foobar') + + def test_kurtosis_array_scalar(self): + # "array scalars" do not exist in other backends + assert_equal(type(stats.kurtosis([1, 2, 3])), np.float64) + + def test_kurtosis_propagate_nan(self): + # nan_policy only for NumPy right now + # Check that the shape of the result is the same for inputs + # with and without nans, cf gh-5817 + a = np.arange(8).reshape(2, -1).astype(float) + a[1, 0] = np.nan + k = stats.kurtosis(a, axis=1, nan_policy="propagate") + np.testing.assert_allclose(k, [-1.36, np.nan], atol=1e-15) + + def test_kurtosis_constant_value(self, xp): + # Kurtosis of a constant input should be NaN (gh-16061) + a = xp.asarray([-0.27829495]*10) + with eager_warns(RuntimeWarning, match="Precision loss occurred", xp=xp): + assert xp.isnan(stats.kurtosis(a, fisher=False)) + assert xp.isnan(stats.kurtosis(a * float(2**50), fisher=False)) + assert xp.isnan(stats.kurtosis(a / float(2**50), fisher=False)) + assert xp.isnan(stats.kurtosis(a, fisher=False, bias=False)) + + @pytest.mark.parametrize('axis', [-1, 0, 2, None]) + @pytest.mark.parametrize('bias', [False, True]) + @pytest.mark.parametrize('fisher', [False, True]) + def test_vectorization(self, xp, axis, bias, fisher): + # Behavior with array input is not tested above. Compare + # against naive implementation. + rng = np.random.default_rng(1283413549926) + x = xp.asarray(rng.random((4, 5, 6))) + + def kurtosis(a, axis, bias, fisher): + # Simple implementation of kurtosis + if axis is None: + a = xp.reshape(a, (-1,)) + axis = 0 + mean = xp.mean(a, axis=axis, keepdims=True) + mu4 = xp.mean((a - mean)**4, axis=axis) + mu2 = xp.var(a, axis=axis, correction=0) + if bias: + res = mu4 / mu2**2 - 3 + else: + n = a.shape[axis] + # https://en.wikipedia.org/wiki/Kurtosis#Standard_unbiased_estimator + res = (n-1) / ((n-2) * (n-3)) * ((n + 1) * mu4/mu2**2 - 3*(n-1)) + + # I know it looks strange to subtract then add 3, + # but it is simpler than the alternatives + return res if fisher else res + 3 + + res = stats.kurtosis(x, axis=axis, bias=bias, fisher=fisher) + ref = kurtosis(x, axis=axis, bias=bias, fisher=fisher) + xp_assert_close(res, ref) + + +@hypothesis.strategies.composite +def ttest_data_axis_strategy(draw): + # draw an array under shape and value constraints + elements = dict(allow_nan=False, allow_infinity=False) + shape = npst.array_shapes(min_dims=1, min_side=2) + # The test that uses this, `test_pvalue_ci`, uses `float64` to test + # extreme `alpha`. It could be adjusted to test a dtype-dependent + # range of `alpha` if this strategy is needed to generate other floats. + data = draw(npst.arrays(dtype=np.float64, elements=elements, shape=shape)) + + # determine axes over which nonzero variance can be computed accurately + ok_axes = [] + # Locally, I don't need catch_warnings or simplefilter, and I can just + # suppress RuntimeWarning. I include all that in hope of getting the same + # behavior on CI. + with warnings.catch_warnings(): + warnings.simplefilter("error") + for axis in range(len(data.shape)): + with contextlib.suppress(Exception): + var = stats.moment(data, order=2, axis=axis) + if np.all(var > 0) and np.all(np.isfinite(var)): + ok_axes.append(axis) + # if there are no valid axes, tell hypothesis to try a different example + hypothesis.assume(ok_axes) + + # draw one of the valid axes + axis = draw(hypothesis.strategies.sampled_from(ok_axes)) + + return data, axis + + +@make_xp_test_case(stats.ttest_1samp) +class TestTTest_1samp: + # Preserving original test cases. + # Recomputed statistics and p-values with R t.test, e.g. + # options(digits=16) + # t.test(c(-1., 0., 1.), mu=2) + X1 = [-1., 0., 1.] + X2 = [0., 1., 2.] + T1_0 = 0. + P1_0 = 1. + T1_1 = -1.7320508075689 + P1_1 = 0.2254033307585 + T1_2 = -3.4641016151378 + P1_2 = 0.07417990022745 + T2_0 = 1.7320508075689 + P2_0 = 0.2254033307585 + P1_1_l = P1_1 / 2 + P1_1_g = 1 - (P1_1 / 2) + + @pytest.mark.filterwarnings("ignore:divide by zero encountered:RuntimeWarning:dask") + @pytest.mark.filterwarnings("ignore:invalid value encountered:RuntimeWarning:dask") + def test_onesample(self, xp): + with warnings.catch_warnings(), \ + np.errstate(invalid="ignore", divide="ignore"): + warnings.filterwarnings( + "ignore", "Degrees of freedom <= 0 for slice", RuntimeWarning) + a = xp.asarray(4.) if not is_numpy(xp) else 4. + t, p = stats.ttest_1samp(a, 3.) + xp_assert_equal(t, xp.asarray(xp.nan)) + xp_assert_equal(p, xp.asarray(xp.nan)) + + t, p = stats.ttest_1samp(xp.asarray(self.X1), 0.) + xp_assert_close(t, xp.asarray(self.T1_0)) + xp_assert_close(p, xp.asarray(self.P1_0)) + + res = stats.ttest_1samp(xp.asarray(self.X1), 0.) + attributes = ('statistic', 'pvalue') + check_named_results(res, attributes, xp=xp) + + t, p = stats.ttest_1samp(xp.asarray(self.X2), 0.) + + xp_assert_close(t, xp.asarray(self.T2_0)) + xp_assert_close(p, xp.asarray(self.P2_0)) + + t, p = stats.ttest_1samp(xp.asarray(self.X1), 1.) + + xp_assert_close(t, xp.asarray(self.T1_1)) + xp_assert_close(p, xp.asarray(self.P1_1)) + + t, p = stats.ttest_1samp(xp.asarray(self.X1), 2.) + + xp_assert_close(t, xp.asarray(self.T1_2)) + xp_assert_close(p, xp.asarray(self.P1_2)) + + def test_onesample_nan_policy_propagate(self, xp): + x = stats.norm.rvs(loc=5, scale=10, size=51, random_state=7654567) + x[50] = np.nan + x = xp.asarray(x, dtype=xp_default_dtype(xp)) + + res = stats.ttest_1samp(x, 5.0) + xp_assert_equal(res.statistic, xp.asarray(xp.nan)) + xp_assert_equal(res.pvalue, xp.asarray(xp.nan)) + + @skip_xp_backends(eager_only=True, reason="lazy arrays don't do 'raise'.") + def test_onesample_nan_policy_omit_raise(self, xp): + x = stats.norm.rvs(loc=5, scale=10, size=51, random_state=7654567) + x[50] = np.nan + x = xp.asarray(x, dtype=xp_default_dtype(xp)) + + res = stats.ttest_1samp(x, 5.0, nan_policy='omit') + xp_assert_close(res.statistic, xp.asarray(-1.6412624074367159)) + xp_assert_close(res.pvalue, xp.asarray(0.107147027334048005)) + + with pytest.raises(ValueError, match="The input contains nan values"): + stats.ttest_1samp(x, 5.0, nan_policy='raise') + + with pytest.raises(ValueError, match="nan_policy must be one of"): + stats.ttest_1samp(x, 5.0, nan_policy='foobar') + + + @pytest.mark.filterwarnings("ignore:divide by zero encountered in divide") + def test_1samp_alternative(self, xp): + message = "`alternative` must be 'less', 'greater', or 'two-sided'." + with pytest.raises(ValueError, match=message): + stats.ttest_1samp(xp.asarray(self.X1), 0., alternative="error") + + t, p = stats.ttest_1samp(xp.asarray(self.X1), 1., alternative="less") + xp_assert_close(p, xp.asarray(self.P1_1_l)) + xp_assert_close(t, xp.asarray(self.T1_1)) + + t, p = stats.ttest_1samp(xp.asarray(self.X1), 1., alternative="greater") + xp_assert_close(p, xp.asarray(self.P1_1_g)) + xp_assert_close(t, xp.asarray(self.T1_1)) + + @skip_xp_backends('jax.numpy', reason='Generic stdtrit mutates array.') + @pytest.mark.parametrize("alternative", ['two-sided', 'less', 'greater']) + def test_1samp_ci_1d(self, xp, alternative): + # test confidence interval method against reference values + rng = np.random.default_rng(8066178009154342972) + n = 10 + x = rng.normal(size=n, loc=1.5, scale=2) + popmean = rng.normal() # this shouldn't affect confidence interval + # Reference values generated with R t.test: + # options(digits=16) + # x = c(2.75532884, 0.93892217, 0.94835861, 1.49489446, -0.62396595, + # -1.88019867, -1.55684465, 4.88777104, 5.15310979, 4.34656348) + # t.test(x, conf.level=0.85, alternative='l') + dtype = xp.asarray(1.0).dtype + x = xp.asarray(x, dtype=dtype) + popmean = xp.asarray(popmean, dtype=dtype) + + ref = {'two-sided': [0.3594423211709136, 2.9333455028290860], + 'greater': [0.7470806207371626, np.inf], + 'less': [-np.inf, 2.545707203262837]} + res = stats.ttest_1samp(x, popmean=popmean, alternative=alternative) + ci = res.confidence_interval(confidence_level=0.85) + xp_assert_close(ci.low, xp.asarray(ref[alternative][0])) + xp_assert_close(ci.high, xp.asarray(ref[alternative][1])) + xp_assert_equal(res.df, xp.asarray(n-1)) + + def test_1samp_ci_iv(self, xp): + # test `confidence_interval` method input validation + res = stats.ttest_1samp(xp.arange(10.), 0.) + message = '`confidence_level` must be a number between 0 and 1.' + with pytest.raises(ValueError, match=message): + res.confidence_interval(confidence_level=10) + + @skip_xp_backends(np_only=True, reason='Too slow.') + @pytest.mark.xslow + @hypothesis.given(alpha=hypothesis.strategies.floats(1e-15, 1-1e-15), + data_axis=ttest_data_axis_strategy()) + @pytest.mark.parametrize('alternative', ['less', 'greater']) + def test_pvalue_ci(self, alpha, data_axis, alternative, xp): + # test relationship between one-sided p-values and confidence intervals + data, axis = data_axis + data = xp.asarray(data) + res = stats.ttest_1samp(data, 0., + alternative=alternative, axis=axis) + l, u = res.confidence_interval(confidence_level=alpha) + popmean = l if alternative == 'greater' else u + popmean = xp.expand_dims(popmean, axis=axis) + res = stats.ttest_1samp(data, popmean, alternative=alternative, axis=axis) + shape = list(data.shape) + shape.pop(axis) + # `float64` is used to correspond with extreme range of `alpha` + ref = xp.broadcast_to(xp.asarray(1-alpha, dtype=xp.float64), shape) + xp_assert_close(res.pvalue, ref) + + +class TestPercentileOfScore: + + def f(self, *args, **kwargs): + return stats.percentileofscore(*args, **kwargs) + + @pytest.mark.parametrize("kind, result", [("rank", 40), + ("mean", 35), + ("strict", 30), + ("weak", 40)]) + def test_unique(self, kind, result): + a = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] + assert_equal(self.f(a, 4, kind=kind), result) + + @pytest.mark.parametrize("kind, result", [("rank", 45), + ("mean", 40), + ("strict", 30), + ("weak", 50)]) + def test_multiple2(self, kind, result): + a = [1, 2, 3, 4, 4, 5, 6, 7, 8, 9] + assert_equal(self.f(a, 4, kind=kind), result) + + @pytest.mark.parametrize("kind, result", [("rank", 50), + ("mean", 45), + ("strict", 30), + ("weak", 60)]) + def test_multiple3(self, kind, result): + a = [1, 2, 3, 4, 4, 4, 5, 6, 7, 8] + assert_equal(self.f(a, 4, kind=kind), result) + + @pytest.mark.parametrize("kind, result", [("rank", 30), + ("mean", 30), + ("strict", 30), + ("weak", 30)]) + def test_missing(self, kind, result): + a = [1, 2, 3, 5, 6, 7, 8, 9, 10, 11] + assert_equal(self.f(a, 4, kind=kind), result) + + @pytest.mark.parametrize("kind, result", [("rank", 40), + ("mean", 35), + ("strict", 30), + ("weak", 40)]) + def test_large_numbers(self, kind, result): + a = [10, 20, 30, 40, 50, 60, 70, 80, 90, 100] + assert_equal(self.f(a, 40, kind=kind), result) + + @pytest.mark.parametrize("kind, result", [("rank", 50), + ("mean", 45), + ("strict", 30), + ("weak", 60)]) + def test_large_numbers_multiple3(self, kind, result): + a = [10, 20, 30, 40, 40, 40, 50, 60, 70, 80] + assert_equal(self.f(a, 40, kind=kind), result) + + @pytest.mark.parametrize("kind, result", [("rank", 30), + ("mean", 30), + ("strict", 30), + ("weak", 30)]) + def test_large_numbers_missing(self, kind, result): + a = [10, 20, 30, 50, 60, 70, 80, 90, 100, 110] + assert_equal(self.f(a, 40, kind=kind), result) + + @pytest.mark.parametrize("kind, result", [("rank", [0, 10, 100, 100]), + ("mean", [0, 5, 95, 100]), + ("strict", [0, 0, 90, 100]), + ("weak", [0, 10, 100, 100])]) + def test_boundaries(self, kind, result): + a = [10, 20, 30, 50, 60, 70, 80, 90, 100, 110] + assert_equal(self.f(a, [0, 10, 110, 200], kind=kind), result) + + @pytest.mark.parametrize("kind, result", [("rank", [0, 10, 100]), + ("mean", [0, 5, 95]), + ("strict", [0, 0, 90]), + ("weak", [0, 10, 100])]) + def test_inf(self, kind, result): + a = [1, 2, 3, 4, 5, 6, 7, 8, 9, +np.inf] + assert_equal(self.f(a, [-np.inf, 1, +np.inf], kind=kind), result) + + cases = [("propagate", [], 1, np.nan), + ("propagate", [np.nan], 1, np.nan), + ("propagate", [np.nan], [0, 1, 2], [np.nan, np.nan, np.nan]), + ("propagate", [1, 2], [1, 2, np.nan], [50, 100, np.nan]), + ("omit", [1, 2, np.nan], [0, 1, 2], [0, 50, 100]), + ("omit", [1, 2], [0, 1, np.nan], [0, 50, np.nan]), + ("omit", [np.nan, np.nan], [0, 1, 2], [np.nan, np.nan, np.nan])] + + @pytest.mark.parametrize("policy, a, score, result", cases) + def test_nans_ok(self, policy, a, score, result): + assert_equal(self.f(a, score, nan_policy=policy), result) + + cases = [ + ("raise", [1, 2, 3, np.nan], [1, 2, 3], + "The input contains nan values"), + ("raise", [1, 2, 3], [1, 2, 3, np.nan], + "The input contains nan values"), + ] + + @pytest.mark.parametrize("policy, a, score, message", cases) + def test_nans_fail(self, policy, a, score, message): + with assert_raises(ValueError, match=message): + self.f(a, score, nan_policy=policy) + + @pytest.mark.parametrize("shape", [ + (6, ), + (2, 3), + (2, 1, 3), + (2, 1, 1, 3), + ]) + def test_nd(self, shape): + a = np.array([0, 1, 2, 3, 4, 5]) + scores = a.reshape(shape) + results = scores*10 + a = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] + assert_equal(self.f(a, scores, kind="rank"), results) + + def test_multidimensional_error(self): + # gh-21563 reported that `percentileofscore` accepted multidimensional + # arrays but did not produce meaningful results. + message = "`a` must be 1-dimensional." + with pytest.raises(ValueError, match=message): + stats.percentileofscore(np.ones((3, 3)), 1) + + +PowerDivCase = namedtuple('Case', # type: ignore[name-match] + ['f_obs', 'f_exp', 'ddof', 'axis', + 'chi2', # Pearson's + 'log', # G-test (log-likelihood) + 'mod_log', # Modified log-likelihood + 'cr', # Cressie-Read (lambda=2/3) + ]) + +# The details of the first two elements in power_div_1d_cases are used +# in a test in TestPowerDivergence. Check that code before making +# any changes here. +power_div_1d_cases = [ + # Use the default f_exp. + PowerDivCase(f_obs=[4, 8, 12, 8], f_exp=None, ddof=0, axis=None, + chi2=4, + log=2*(4*np.log(4/8) + 12*np.log(12/8)), + mod_log=2*(8*np.log(8/4) + 8*np.log(8/12)), + cr=(4*((4/8)**(2/3) - 1) + 12*((12/8)**(2/3) - 1))/(5/9)), + # Give a non-uniform f_exp. + PowerDivCase(f_obs=[4, 8, 12, 8], f_exp=[2, 16, 12, 2], ddof=0, axis=None, + chi2=24, + log=2*(4*np.log(4/2) + 8*np.log(8/16) + 8*np.log(8/2)), + mod_log=2*(2*np.log(2/4) + 16*np.log(16/8) + 2*np.log(2/8)), + cr=(4*((4/2)**(2/3) - 1) + 8*((8/16)**(2/3) - 1) + + 8*((8/2)**(2/3) - 1))/(5/9)), + # f_exp is a scalar. + PowerDivCase(f_obs=[4, 8, 12, 8], f_exp=8, ddof=0, axis=None, + chi2=4, + log=2*(4*np.log(4/8) + 12*np.log(12/8)), + mod_log=2*(8*np.log(8/4) + 8*np.log(8/12)), + cr=(4*((4/8)**(2/3) - 1) + 12*((12/8)**(2/3) - 1))/(5/9)), + # f_exp equal to f_obs. + PowerDivCase(f_obs=[3, 5, 7, 9], f_exp=[3, 5, 7, 9], ddof=0, axis=0, + chi2=0, log=0, mod_log=0, cr=0), +] + + +power_div_empty_cases = [ + # Shape is (0,)--a data set with length 0. The computed + # test statistic should be 0. + PowerDivCase(f_obs=[], + f_exp=None, ddof=0, axis=0, + chi2=0, log=0, mod_log=0, cr=0), + # Shape is (0, 3). This is 3 data sets, but each data set has + # length 0, so the computed test statistic should be [0, 0, 0]. + PowerDivCase(f_obs=np.array([[],[],[]]).T, + f_exp=None, ddof=0, axis=0, + chi2=[0, 0, 0], + log=[0, 0, 0], + mod_log=[0, 0, 0], + cr=[0, 0, 0]), + # Shape is (3, 0). This represents an empty collection of + # data sets in which each data set has length 3. The test + # statistic should be an empty array. + PowerDivCase(f_obs=np.array([[],[],[]]), + f_exp=None, ddof=0, axis=0, + chi2=[], + log=[], + mod_log=[], + cr=[]), +] + + +@make_xp_test_case(stats.power_divergence) +class TestPowerDivergence: + + def check_power_divergence(self, f_obs, f_exp, ddof, axis, lambda_, + expected_stat, xp): + dtype = xp.asarray(1.).dtype + + f_obs = xp.asarray(f_obs, dtype=dtype) + f_exp = xp.asarray(f_exp, dtype=dtype) if f_exp is not None else f_exp + + if axis is None: + num_obs = xp_size(f_obs) + else: + arrays = (xp.broadcast_arrays(f_obs, f_exp) if f_exp is not None + else (f_obs,)) + num_obs = arrays[0].shape[axis] + + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", "Mean of empty slice", RuntimeWarning) + stat, p = stats.power_divergence( + f_obs, f_exp=f_exp, ddof=ddof, + axis=axis, lambda_=lambda_) + xp_assert_close(stat, xp.asarray(expected_stat, dtype=dtype)) + + if lambda_ == 1 or lambda_ == "pearson": + # Also test stats.chisquare. + stat, p = stats.chisquare(f_obs, f_exp=f_exp, ddof=ddof, + axis=axis) + xp_assert_close(stat, xp.asarray(expected_stat, dtype=dtype)) + + ddof = np.asarray(ddof) + expected_p = stats.distributions.chi2.sf(expected_stat, + num_obs - 1 - ddof) + xp_assert_close(p, xp.asarray(expected_p, dtype=dtype)) + + @pytest.mark.parametrize('case', power_div_1d_cases) + @pytest.mark.parametrize('lambda_stat', + [(None, 'chi2'), ('pearson', 'chi2'), (1, 'chi2'), + ('log-likelihood', 'log'), ('mod-log-likelihood', 'mod_log'), + ('cressie-read', 'cr'), (2/3, 'cr')]) + def test_basic(self, case, lambda_stat, xp): + lambda_, attr = lambda_stat + expected_stat = getattr(case, attr) + self.check_power_divergence(case.f_obs, case.f_exp, case.ddof, case.axis, + lambda_, expected_stat, xp) + + def test_axis(self, xp): + case0 = power_div_1d_cases[0] + case1 = power_div_1d_cases[1] + f_obs = np.vstack((case0.f_obs, case1.f_obs)) + f_exp = np.vstack((np.ones_like(case0.f_obs)*np.mean(case0.f_obs), + case1.f_exp)) + # Check the four computational code paths in power_divergence + # using a 2D array with axis=1. + f_obs = xp.asarray(f_obs) + f_exp = xp.asarray(f_exp) if f_exp is not None else f_exp + self.check_power_divergence( + f_obs, f_exp, 0, 1, + "pearson", [case0.chi2, case1.chi2], xp=xp) + self.check_power_divergence( + f_obs, f_exp, 0, 1, + "log-likelihood", [case0.log, case1.log], xp=xp) + self.check_power_divergence( + f_obs, f_exp, 0, 1, + "mod-log-likelihood", [case0.mod_log, case1.mod_log], xp=xp) + self.check_power_divergence( + f_obs, f_exp, 0, 1, + "cressie-read", [case0.cr, case1.cr], xp=xp) + # Reshape case0.f_obs to shape (2,2), and use axis=None. + # The result should be the same. + f_obs_reshape = xp.reshape(xp.asarray(case0.f_obs), (2, 2)) + self.check_power_divergence( + f_obs_reshape, None, 0, None, + "pearson", case0.chi2, xp=xp) + + def test_ddof_broadcasting(self, xp): + # Test that ddof broadcasts correctly. + # ddof does not affect the test statistic. It is broadcast + # with the computed test statistic for the computation of + # the p value. + + case0 = power_div_1d_cases[0] + case1 = power_div_1d_cases[1] + # Create 4x2 arrays of observed and expected frequencies. + f_obs = np.vstack((case0.f_obs, case1.f_obs)).T + f_exp = np.vstack((np.ones_like(case0.f_obs)*np.mean(case0.f_obs), + case1.f_exp)).T + + expected_chi2 = [case0.chi2, case1.chi2] + + dtype = xp.asarray(1.).dtype + f_obs = xp.asarray(f_obs, dtype=dtype) + f_exp = xp.asarray(f_exp, dtype=dtype) + expected_chi2 = xp.asarray(expected_chi2, dtype=dtype) + + # ddof has shape (2, 1). This is broadcast with the computed + # statistic, so p will have shape (2,2). + ddof = xp.asarray([[0], [1]]) + + stat, p = stats.power_divergence(f_obs, f_exp, ddof=ddof) + xp_assert_close(stat, expected_chi2) + + # Compute the p values separately, passing in scalars for ddof. + _, p0 = stats.power_divergence(f_obs, f_exp, ddof=ddof[0, 0]) + _, p1 = stats.power_divergence(f_obs, f_exp, ddof=ddof[1, 0]) + + expected_p = xp.concat((p0[xp.newaxis, :], p1[xp.newaxis, :]), axis=0) + xp_assert_close(p, expected_p) + + @pytest.mark.parametrize('case', power_div_empty_cases) + @pytest.mark.parametrize('lambda_stat', + [('pearson', 'chi2'), ('log-likelihood', 'log'), + ('mod-log-likelihood', 'mod_log'), + ('cressie-read', 'cr'), (2/3, 'cr')]) + def test_empty_cases(self, case, lambda_stat, xp): + lambda_, attr = lambda_stat + expected_stat = getattr(case, attr) + with warnings.catch_warnings(): + self.check_power_divergence( + case.f_obs, case.f_exp, case.ddof, case.axis, + lambda_, expected_stat, xp) + + def test_power_divergence_result_attributes(self, xp): + f_obs = power_div_1d_cases[0].f_obs + f_exp = power_div_1d_cases[0].f_exp + ddof = power_div_1d_cases[0].ddof + axis = power_div_1d_cases[0].axis + dtype = xp.asarray(1.).dtype + f_obs = xp.asarray(f_obs, dtype=dtype) + # f_exp is None + + res = stats.power_divergence(f_obs, f_exp=f_exp, ddof=ddof, + axis=axis, lambda_="pearson") + attributes = ('statistic', 'pvalue') + check_named_results(res, attributes, xp=xp) + + def test_power_divergence_gh_12282(self, xp): + # The sums of observed and expected frequencies must match + f_obs = xp.asarray([[10., 20.], [30., 20.]]) + f_exp = xp.asarray([[5., 15.], [35., 25.]]) + stat, pval = stats.power_divergence(f_obs, f_exp=f_exp) + xp_assert_close(stat, xp.asarray([5.71428571, 2.66666667])) + xp_assert_close(pval, xp.asarray([0.01682741, 0.10247043])) + + if is_lazy_array(f_obs): + res = stats.power_divergence(f_obs, f_exp=xp.asarray([30., 60.])) + xp_assert_equal(res.statistic, xp.asarray([xp.nan, xp.nan])) + xp_assert_equal(res.pvalue, xp.asarray([xp.nan, xp.nan])) + + res = stats.power_divergence(f_obs, f_exp=f_exp, axis=1) + xp_assert_equal(res.statistic, xp.asarray([xp.nan, xp.nan])) + xp_assert_equal(res.pvalue, xp.asarray([xp.nan, xp.nan])) + + # only slices with unequal sums result in NaN + f_exp = xp.asarray([[5., 25.], [35., 25.]]) + res = stats.power_divergence(f_obs, f_exp=f_exp) + xp_assert_close(res.statistic, xp.asarray([stat[0], xp.nan])) + xp_assert_close(res.pvalue, xp.asarray([pval[0], xp.nan])) + else: + message = 'For each axis slice...' + with pytest.raises(ValueError, match=message): + stats.power_divergence(f_obs, f_exp=xp.asarray([30., 60.])) + with pytest.raises(ValueError, match=message): + stats.power_divergence(f_obs, f_exp=f_exp, axis=1) + + def test_power_divergence_against_cressie_read_data(self, xp): + # Test stats.power_divergence against tables 4 and 5 from + # Cressie and Read, "Multimonial Goodness-of-Fit Tests", + # J. R. Statist. Soc. B (1984), Vol 46, No. 3, pp. 440-464. + # This tests the calculation for several values of lambda. + + # Table 4 data recalculated for greater precision according to: + # Shelby J. Haberman, Analysis of Qualitative Data: Volume 1 + # Introductory Topics, Academic Press, New York, USA (1978). + obs = xp.asarray([15., 11., 14., 17., 5., 11., 10., 4., 8., + 10., 7., 9., 11., 3., 6., 1., 1., 4.]) + beta = -0.083769 # Haberman (1978), p. 15 + i = xp.arange(1., obs.shape[0] + 1.) + alpha = xp.log(xp.sum(obs) / xp.sum(xp.exp(beta*i))) + expected_counts = xp.exp(alpha + beta*i) + + # `table4` holds just the second and third columns from Table 4. + table4 = xp.concat((obs[xp.newaxis, :], + expected_counts[xp.newaxis, :])).T + + table5 = np.asarray([ + # lambda, statistic + -10.0, 72.2e3, + -5.0, 28.9e1, + -3.0, 65.6, + -2.0, 40.6, + -1.5, 34.0, + -1.0, 29.5, + -0.5, 26.5, + 0.0, 24.6, + 0.5, 23.4, + 0.67, 23.1, + 1.0, 22.7, + 1.5, 22.6, + 2.0, 22.9, + 3.0, 24.8, + 5.0, 35.5, + 10.0, 21.4e1, + ]) + table5 = np.reshape(table5, (-1, 2)) + + for i in range(table5.shape[0]): + lambda_, expected_stat = float(table5[i, 0]), float(table5[i, 1]) + stat, p = stats.power_divergence(table4[:, 0], table4[:, 1], + lambda_=lambda_) + xp_assert_close(stat, xp.asarray(expected_stat), rtol=5e-3) + + +@make_xp_test_case(stats.chisquare) +class TestChisquare: + def test_chisquare_12282a(self, xp): + # Currently `chisquare` is implemented via power_divergence + # in case that ever changes, perform a basic test like + # test_power_divergence_gh_12282 + f_obs = xp.asarray([10., 20.]) + f_exp = xp.asarray([30., 60.]) + if is_lazy_array(f_obs): + res = stats.chisquare(f_obs, f_exp=f_exp) + xp_assert_equal(res.statistic, xp.asarray(xp.nan)) + xp_assert_equal(res.pvalue, xp.asarray(xp.nan)) + else: + with assert_raises(ValueError, match='For each axis slice...'): + stats.chisquare(f_obs, f_exp=f_exp) + + def test_chisquare_12282b(self, xp): + # Check that users can now disable the sum check tested in + # test_chisquare_12282a. Also, confirm that statistic and p-value + # are as expected. + rng = np.random.default_rng(3843874358728234) + n = 10 + lam = rng.uniform(1000, 2000, size=n) + x = rng.poisson(lam) + lam = xp.asarray(lam) + x = xp.asarray(x, dtype=lam.dtype) + res = stats.chisquare(x, f_exp=lam, ddof=-1, sum_check=False) + # Poisson is approximately normal with mean and variance lam + z = (x - lam) / xp.sqrt(lam) + statistic = xp.sum(z**2) + xp_assert_close(res.statistic, statistic) + # Sum of `n` squared standard normal variables follows chi2 with `n` DoF + X2 = _SimpleChi2(xp.asarray(n, dtype=statistic.dtype)) + xp_assert_close(res.pvalue, X2.sf(statistic)) + + @pytest.mark.parametrize("n, dtype", [(200, 'uint8'), (1000000, 'int32')]) + def test_chiquare_data_types_attributes(self, n, dtype, xp): + # Regression test for gh-10159 and gh-18368 + dtype = getattr(xp, dtype) + obs = xp.asarray([n, 0], dtype=dtype) + exp = xp.asarray([n // 2, n // 2], dtype=dtype) + res = stats.chisquare(obs, exp) + stat, p = res + xp_assert_close(stat, xp.asarray(n, dtype=xp.asarray(1.).dtype), rtol=1e-13) + # check that attributes are identical to unpacked outputs - see gh-18368 + xp_assert_equal(res.statistic, stat) + xp_assert_equal(res.pvalue, p) + + +@make_xp_test_case(stats.friedmanchisquare) +class TestFriedmanChiSquare: + # verified with matlab and R + # From Demsar "Statistical Comparisons of Classifiers over Multiple Data Sets" + # 2006, Xf=9.28 (no tie handling, tie corrected Xf >=9.28) + x1 = [[0.763, 0.599, 0.954, 0.628, 0.882, 0.936, 0.661, + 0.583, 0.775, 1.0, 0.94, 0.619, 0.972, 0.957], + [0.768, 0.591, 0.971, 0.661, 0.888, 0.931, 0.668, + 0.583, 0.838, 1.0, 0.962, 0.666, 0.981, 0.978], + [0.771, 0.590, 0.968, 0.654, 0.886, 0.916, 0.609, + 0.563, 0.866, 1.0, 0.965, 0.614, 0.9751, 0.946], + [0.798, 0.569, 0.967, 0.657, 0.898, 0.931, 0.685, + 0.625, 0.875, 1.0, 0.962, 0.669, 0.975, 0.970]] + ref1 = (10.2283464566929, 0.0167215803284414) + + # From "Bioestadistica para las ciencias de la salud" Xf=18.95 p<0.001: + x2 = [[4, 3, 5, 3, 5, 3, 2, 5, 4, 4, 4, 3], + [2, 2, 1, 2, 3, 1, 2, 3, 2, 1, 1, 3], + [2, 4, 3, 3, 4, 3, 3, 4, 4, 1, 2, 1], + [3, 5, 4, 3, 4, 4, 3, 3, 3, 4, 4, 4]] + ref2 = (18.9428571428571, 0.000280938375189499) + + # From Jerrorl H. Zar, "Biostatistical Analysis"(example 12.6), + # Xf=10.68, 0.005 < p < 0.01: + # Probability from this example is inexact + # using Chisquare approximation of Friedman Chisquare. + x3 = [[7.0, 9.9, 8.5, 5.1, 10.3], + [5.3, 5.7, 4.7, 3.5, 7.7], + [4.9, 7.6, 5.5, 2.8, 8.4], + [8.8, 8.9, 8.1, 3.3, 9.1]] + ref3 = (10.68, 0.0135882729582176) + + @pytest.mark.parametrize("dtype", [None, "float32", "float64"]) + @pytest.mark.parametrize("data, ref", [(x1, ref1), (x2, ref2), (x3, ref3)]) + def test_against_references(self, dtype, data, ref, xp): + dtype = dtype if dtype is None else getattr(xp, dtype) + data = [xp.asarray(array, dtype=dtype) for array in data] + res = stats.friedmanchisquare(*data) + xp_assert_close(res.statistic, xp.asarray(ref[0], dtype=dtype)) + xp_assert_close(res.pvalue, xp.asarray(ref[1], dtype=dtype)) + + def test_too_few_samples(self, xp): + message = "At least 3 samples must be given" + with pytest.raises(ValueError, match=message): + stats.friedmanchisquare(xp.asarray(self.x3[0]), xp.asarray(self.x3[1])) + + +@make_xp_test_case(stats.kstest) +class TestKSTest: + """Tests kstest and ks_1samp agree with K-S various sizes, alternatives, modes.""" + + def _test_kstest_and_ks1samp(self, x, alternative, mode='auto', decimal=14): + result = stats.kstest(x, 'norm', alternative=alternative, mode=mode) + result_1samp = stats.ks_1samp(x, special.ndtr, + alternative=alternative, mode=mode) + xp_assert_close(result.statistic, result_1samp.statistic) + xp_assert_close(result.pvalue, result_1samp.pvalue) + xp_assert_equal(result.statistic_location, result.statistic_location) + xp_assert_equal(result.statistic_sign, result.statistic_sign) + + def test_namedtuple_attributes(self, xp): + x = xp.linspace(-1, 1, 9) + # test for namedtuple attribute results + attributes = ('statistic', 'pvalue') + res = stats.kstest(x, special.ndtr) + check_named_results(res, attributes, xp=xp) + + def test_agree_with_ks_1samp(self, xp): + x = xp.linspace(-1, 1, 9) + self._test_kstest_and_ks1samp(x, 'two-sided') + + x = xp.linspace(-15, 15, 9) + self._test_kstest_and_ks1samp(x, 'two-sided') + + x = [-1.23, 0.06, -0.60, 0.17, 0.66, -0.17, -0.08, 0.27, -0.98, -0.99] + x = xp.asarray(x) + self._test_kstest_and_ks1samp(x, 'two-sided') + self._test_kstest_and_ks1samp(x, 'greater', mode='exact') + self._test_kstest_and_ks1samp(x, 'less', mode='exact') + + # parametrize over function-specific options to make sure they're passed through; + # `test_axis_nan_policy` takes care of the rest. + @pytest.mark.parametrize("method", ["auto", "exact", "asymp"]) + @pytest.mark.parametrize("alternative", ['two-sided', 'less', 'greater']) + def test_agree_with_ks_2samp(self, method, alternative, xp): + rng = np.random.default_rng(236114087) + x, y = rng.random((2, 9, 10)) + x, y = xp.asarray(x), xp.asarray(y) + res = stats.kstest(x, y, alternative=alternative, method=method) + ref = stats.ks_2samp(x, y, alternative=alternative, method=method) + xp_assert_equal(res.statistic, ref.statistic) + xp_assert_equal(res.pvalue, ref.pvalue) + xp_assert_equal(res.statistic_location, ref.statistic_location) + xp_assert_equal(res.statistic_sign, ref.statistic_sign) + + def test_pm_inf_gh20386(self, xp): + # Check that gh-20386 is resolved - `kstest` does not + # return NaNs when both -inf and inf are in sample. + vals = xp.asarray([-xp.inf, 0, 1, xp.inf]) + def cdf(x): # Cauchy CDF + return xp.atan2(xp.ones_like(x), -x) / xp.pi + res = stats.kstest(vals, cdf) + ref = stats.kstest(vals, cdf, _no_deco=True) + assert xp.isfinite(res.statistic) + assert xp.isfinite(res.pvalue) + assert xp.isfinite(res.statistic_location) + assert xp.isfinite(res.statistic_sign) + xp_assert_equal(res.statistic, ref.statistic) + xp_assert_equal(res.pvalue, ref.pvalue) + xp_assert_equal(res.statistic_location, ref.statistic_location) + xp_assert_equal(res.statistic_sign, ref.statistic_sign) + + # missing: no test that uses *args + + +@make_xp_test_case(stats.ks_1samp) +class TestKSOneSample: + """ + Tests kstest and ks_samp 1-samples with K-S various sizes, alternatives, modes. + """ + + def _testOne(self, x, alternative, expected_statistic, expected_prob, *, + mode='auto', dtype, xp): + rtol = 5e-14 if dtype == xp.float64 else 1e-5 + res = stats.ks_1samp(x, special.ndtr, alternative=alternative, mode=mode) + ref_statistic = xp.asarray(expected_statistic, dtype=dtype) + ref_pvalue = xp.asarray(expected_prob, dtype=dtype) + xp_assert_close(res.statistic, ref_statistic, rtol=rtol) + xp_assert_close(res.pvalue, ref_pvalue, rtol=rtol) + + @pytest.mark.parametrize('dtype', [None, 'float32', 'float64']) + def test_agree_with_r(self, dtype, xp): + # comparing with some values from R + dtype = xp_default_dtype(xp) if dtype is None else getattr(xp, dtype) + x = xp.linspace(-1, 1, 9, dtype=dtype) + self._testOne(x, 'two-sided', 0.15865525393145705, 0.95164069201518386, + dtype=dtype, xp=xp) + + x = xp.linspace(-15, 15, 9, dtype=dtype) + self._testOne(x, 'two-sided', 0.44435602715924361, 0.038850140086788665, + dtype=dtype, xp=xp) + + x = [-1.23, 0.06, -0.60, 0.17, 0.66, -0.17, -0.08, 0.27, -0.98, -0.99] + x = xp.asarray(x, dtype=dtype) + self._testOne(x, 'two-sided', 0.293580126801961, 0.293408463684361, + dtype=dtype, xp=xp) + self._testOne(x, 'greater', 0.293580126801961, 0.146988835042376, mode='exact', + dtype=dtype, xp=xp) + self._testOne(x, 'less', 0.109348552425692, 0.732768892470675, mode='exact', + dtype=dtype, xp=xp) + + @pytest.mark.parametrize('dtype', [None, 'float32', 'float64']) + def test_known_examples(self, xp, dtype): + # the following tests rely on deterministically replicated rvs + dtype = xp_default_dtype(xp) if dtype is None else getattr(xp, dtype) + x = stats.norm.rvs(loc=0.2, size=100, random_state=987654321) + x = xp.asarray(x, dtype=dtype) + self._testOne(x, 'two-sided', 0.12464329735846891, 0.089444888711820769, + mode='asymp', xp=xp, dtype=dtype) + self._testOne(x, 'less', 0.12464329735846891, 0.040989164077641749, + xp=xp, dtype=dtype) + self._testOne(x, 'greater', 0.0072115233216310994, 0.98531158590396228, + xp=xp, dtype=dtype) + + # this is a test of the exact p-value calculation, available only with NumPy. + def test_ks1samp_allpaths(self): + # Check NaN input, output. + assert_(np.isnan(kolmogn(np.nan, 1, True))) + with assert_raises(ValueError, match='n is not integral: 1.5'): + kolmogn(1.5, 1, True) + assert_(np.isnan(kolmogn(-1, 1, True))) + + dataset = np.asarray([ + # Check x out of range + (101, 1, True, 1.0), + (101, 1.1, True, 1.0), + (101, 0, True, 0.0), + (101, -0.1, True, 0.0), + + (32, 1.0 / 64, True, 0.0), # Ruben-Gambino + (32, 1.0 / 64, False, 1.0), # Ruben-Gambino + + # Miller + (32, 0.5, True, 0.9999999363163307), + # Miller 2 * special.smirnov(32, 0.5) + (32, 0.5, False, 6.368366937916623e-08), + + # Check some other paths + (32, 1.0 / 8, True, 0.34624229979775223), + (32, 1.0 / 4, True, 0.9699508336558085), + (1600, 0.49, False, 0.0), + # 2 * special.smirnov(1600, 1/16.0) + (1600, 1 / 16.0, False, 7.0837876229702195e-06), + # _kolmogn_DMTW + (1600, 14 / 1600, False, 0.99962357317602), + # _kolmogn_PelzGood + (1600, 1 / 32, False, 0.08603386296651416), + ]) + FuncData(kolmogn, dataset, (0, 1, 2), 3).check(dtypes=[int, float, bool]) + + @pytest.mark.parametrize("ksfunc", [stats.kstest, stats.ks_1samp]) + @pytest.mark.parametrize("alternative, x6val, ref_location, ref_sign", + [('greater', 6., 6., +1), + ('less', 7., 7., -1), + ('two-sided', 6., 6., +1), + ('two-sided', 7., 7., -1)]) + def test_location_sign(self, ksfunc, alternative, + x6val, ref_location, ref_sign, xp): + # Test that location and sign corresponding with statistic are as + # expected. (Test is designed to be easy to predict.) + x = xp.arange(10.) + 0.5 + x = xpx.at(x)[6].set(x6val) + # cdf = stats.uniform(scale=10).cdf + def cdf(x): return x / 10. + res = ksfunc(xp.asarray(x), cdf, alternative=alternative) + rtol = 1e-15 if x.dtype == xp.float64 else 1e-6 + xp_assert_close(res.statistic, xp.asarray(0.1), rtol=rtol) + xp_assert_equal(res.statistic_location, xp.asarray(ref_location)) + xp_assert_equal(res.statistic_sign, xp.asarray(ref_sign, dtype=xp.int8)) + + def test_invalid_method_raises(self): + # gh-24737: unrecognized method strings must raise ValueError, not + # silently fall through to the 'approx' else-clause. + rng = np.random.default_rng(42) + x = rng.standard_normal(10) + with pytest.raises(ValueError, match="Invalid value for method"): + stats.ks_1samp(x, special.ndtr, method='any_other_string') + # must fire for one-sided alternatives too, not just two-sided + with pytest.raises(ValueError, match="Invalid value for method"): + stats.ks_1samp(x, special.ndtr, method='typo', + alternative='greater') + with pytest.raises(ValueError, match="Invalid value for method"): + stats.ks_1samp(x, special.ndtr, method='typo', + alternative='less') + # kstest wrapper must also raise + with pytest.raises(ValueError, match="Invalid value for method"): + stats.kstest(x, stats.norm.cdf, method='any_other_string') + +@make_xp_test_case(stats.ks_2samp) +class TestKSTwoSamples: + """Tests 2-samples with K-S various sizes, alternatives, modes.""" + + def _check_warnings(self, warn_list, expected_type, expected_len): + """ + Checks that all of the warnings from a list returned by + `warnings.catch_all(record=True)` are of the required type and that the list + contains expected number of warnings. + """ + assert_equal(len(warn_list), expected_len, "number of warnings") + for warn_ in warn_list: + assert_(warn_.category is expected_type) + + def _testOne(self, x1, x2, alternative, ref_statistic, ref_pvalue, + mode='auto', *, xp): + x1, x2 = xp.asarray(x1), xp.asarray(x2) + dtype = xp_result_type(x1, x2, force_floating=True, xp=xp) + res_statistic, res_pvalue = stats.ks_2samp(x1, x2, alternative, mode=mode) + ref_statistic = xp.asarray(ref_statistic, dtype=dtype) + ref_pvalue = xp.asarray(ref_pvalue, dtype=dtype) + xp_assert_close(res_statistic, ref_statistic) + # The original tests used assert_array_almost_equal, which has an absolute + # tolerance by default. It would be good to test with only a relative tolerance, + # but that would require computing p-values with another reference. R would be + # the natural choice, but it clips tiny p-values. + xp_assert_close(res_pvalue, ref_pvalue, atol=1e-14) + + def testSmall(self, xp): + self._testOne([0], [1], 'two-sided', 1.0/1, 1.0, xp=xp) + self._testOne([0], [1], 'greater', 1.0/1, 0.5, xp=xp) + self._testOne([0], [1], 'less', 0.0/1, 1.0, xp=xp) + self._testOne([1], [0], 'two-sided', 1.0/1, 1.0, xp=xp) + self._testOne([1], [0], 'greater', 0.0/1, 1.0, xp=xp) + self._testOne([1], [0], 'less', 1.0/1, 0.5, xp=xp) + + def testTwoVsThree(self, xp): + data1 = xp.asarray([1.0, 2.0]) + data1p = data1 + 0.01 + data1m = data1 - 0.01 + data2 = xp.asarray([1.0, 2.0, 3.0]) + self._testOne(data1p, data2, 'two-sided', 1.0 / 3, 1.0, xp=xp) + self._testOne(data1p, data2, 'greater', 1.0 / 3, 0.7, xp=xp) + self._testOne(data1p, data2, 'less', 1.0 / 3, 0.7, xp=xp) + self._testOne(data1m, data2, 'two-sided', 2.0 / 3, 0.6, xp=xp) + self._testOne(data1m, data2, 'greater', 2.0 / 3, 0.3, xp=xp) + self._testOne(data1m, data2, 'less', 0, 1.0, xp=xp) + + def testTwoVsFour(self, xp): + data1 = xp.asarray([1.0, 2.0]) + data1p = data1 + 0.01 + data1m = data1 - 0.01 + data2 = xp.asarray([1.0, 2.0, 3.0, 4.0]) + self._testOne(data1p, data2, 'two-sided', 2.0 / 4, 14.0/15, xp=xp) + self._testOne(data1p, data2, 'greater', 2.0 / 4, 8.0/15, xp=xp) + self._testOne(data1p, data2, 'less', 1.0 / 4, 12.0/15, xp=xp) + + self._testOne(data1m, data2, 'two-sided', 3.0 / 4, 6.0/15, xp=xp) + self._testOne(data1m, data2, 'greater', 3.0 / 4, 3.0/15, xp=xp) + self._testOne(data1m, data2, 'less', 0, 1.0, xp=xp) + + def test100_100(self, xp): + x100 = xp.linspace(1, 100, 100) + x100_2_p1 = x100 + 2 + 0.1 + x100_2_m1 = x100 + 2 - 0.1 + self._testOne(x100, x100_2_p1, 'two-sided', 3.0 / 100, 0.999999999996206, xp=xp) + self._testOne(x100, x100_2_p1, 'greater', 3.0 / 100, 0.9143290114276248, xp=xp) + self._testOne(x100, x100_2_p1, 'less', 0, 1.0, xp=xp) + self._testOne(x100, x100_2_m1, 'two-sided', 2.0 / 100, 1.0, xp=xp) + self._testOne(x100, x100_2_m1, 'greater', 2.0 / 100, 0.960978450786184, xp=xp) + self._testOne(x100, x100_2_m1, 'less', 0, 1.0, xp=xp) + + def test100_110(self, xp): + x100 = xp.linspace(1, 100, 100) + x110 = xp.linspace(1, 100, 110) + x110_20_p1 = x110 + 20 + 0.1 + x110_20_m1 = x110 + 20 - 0.1 + # 100, 110 + self._testOne(x100, x110_20_p1, 'two-sided', + 232.0 / 1100, 0.015739183865607353, xp=xp) + self._testOne(x100, x110_20_p1, 'greater', + 232.0 / 1100, 0.007869594319053203, xp=xp) + self._testOne(x100, x110_20_p1, 'less', + 0, 1, xp=xp) + self._testOne(x100, x110_20_m1, 'two-sided', + 229.0 / 1100, 0.017803803861026313, xp=xp) + self._testOne(x100, x110_20_m1, 'greater', + 229.0 / 1100, 0.008901905958245056, xp=xp) + self._testOne(x100, x110_20_m1, 'less', + 0.0, 1.0, xp=xp) + + def testRepeatedValues(self, xp): + x2233 = xp.asarray([2] * 3 + [3] * 4 + [5] * 5 + [6] * 4) + x3344 = x2233 + 1 + x2356 = xp.asarray([2] * 3 + [3] * 4 + [5] * 10 + [6] * 4) + x3467 = xp.asarray([3] * 10 + [4] * 2 + [6] * 10 + [7] * 4) + self._testOne(x2233, x3344, 'two-sided', 5.0/16, 0.4262934613454952, xp=xp) + self._testOne(x2233, x3344, 'greater', 5.0/16, 0.21465428276573786, xp=xp) + self._testOne(x2233, x3344, 'less', 0.0/16, 1.0, xp=xp) + self._testOne(x2356, x3467, 'two-sided', 190.0/21/26, 0.0919245790168125, xp=xp) + self._testOne(x2356, x3467, 'greater', 190.0/21/26, 0.0459633806858544, xp=xp) + self._testOne(x2356, x3467, 'less', 70.0/21/26, 0.6121593130022775, xp=xp) + + def testEqualSizes(self, xp): + data2 = xp.asarray([1.0, 2.0, 3.0]) + self._testOne(data2, data2+1, 'two-sided', 1.0/3, 1.0, xp=xp) + self._testOne(data2, data2+1, 'greater', 1.0/3, 0.75, xp=xp) + self._testOne(data2, data2+1, 'less', 0.0/3, 1., xp=xp) + self._testOne(data2, data2+0.5, 'two-sided', 1.0/3, 1.0, xp=xp) + self._testOne(data2, data2+0.5, 'greater', 1.0/3, 0.75, xp=xp) + self._testOne(data2, data2+0.5, 'less', 0.0/3, 1., xp=xp) + self._testOne(data2, data2-0.5, 'two-sided', 1.0/3, 1.0, xp=xp) + self._testOne(data2, data2-0.5, 'greater', 0.0/3, 1.0, xp=xp) + self._testOne(data2, data2-0.5, 'less', 1.0/3, 0.75, xp=xp) + + @pytest.mark.slow + def testMiddlingBoth(self, xp): + # 500, 600 + n1, n2 = 500, 600 + delta = 1.0/n1/n2/2/2 + x = xp.linspace(1., 200., n1) - delta + y = xp.linspace(2., 200., n2) + self._testOne(x, y, 'two-sided', 2000.0 / n1 / n2, 1.0, + mode='auto', xp=xp) + self._testOne(x, y, 'two-sided', 2000.0 / n1 / n2, 1.0, + mode='asymp', xp=xp) + self._testOne(x, y, 'greater', 2000.0 / n1 / n2, 0.9697596024683929, + mode='asymp', xp=xp) + self._testOne(x, y, 'less', 500.0 / n1 / n2, 0.9968735843165021, + mode='asymp', xp=xp) + with warnings.catch_warnings(): + message = "ks_2samp: Exact calculation unsuccessful." + warnings.filterwarnings("ignore", message, RuntimeWarning) + self._testOne(x, y, 'greater', 2000.0 / n1 / n2, 0.9697596024683929, + mode='exact', xp=xp) + self._testOne(x, y, 'less', 500.0 / n1 / n2, 0.9968735843165021, + mode='exact', xp=xp) + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + self._testOne(x, y, 'less', 500.0 / n1 / n2, 0.9968735843165021, + mode='exact', xp=xp) + self._check_warnings(w, RuntimeWarning, 1) + + @pytest.mark.slow + def testMediumBoth(self, xp): + # 1000, 1100 + n1, n2 = 1000, 1100 + delta = 1.0/n1/n2/2/2 + x = xp.linspace(1., 200., n1) - delta + y = xp.linspace(2., 200., n2) + self._testOne(x, y, 'two-sided', 6600.0 / n1 / n2, 1.0, + mode='asymp', xp=xp) + self._testOne(x, y, 'two-sided', 6600.0 / n1 / n2, 1.0, + mode='auto', xp=xp) + self._testOne(x, y, 'greater', 6600.0 / n1 / n2, 0.9573185808092622, + mode='asymp', xp=xp) + self._testOne(x, y, 'less', 1000.0 / n1 / n2, 0.9982410869433984, + mode='asymp', xp=xp) + + with warnings.catch_warnings(): + message = "ks_2samp: Exact calculation unsuccessful." + warnings.filterwarnings("ignore", message, RuntimeWarning) + self._testOne(x, y, 'greater', 6600.0 / n1 / n2, 0.9573185808092622, + mode='exact', xp=xp) + self._testOne(x, y, 'less', 1000.0 / n1 / n2, 0.9982410869433984, + mode='exact', xp=xp) + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + self._testOne(x, y, 'less', 1000.0 / n1 / n2, 0.9982410869433984, + mode='exact', xp=xp) + self._check_warnings(w, RuntimeWarning, 1) + + def testLarge(self, xp): + # 10000, 110 + n1, n2 = 10000, 110 + lcm = n1*11.0 + delta = 1.0/n1/n2/2/2 + x = xp.linspace(1., 200., n1) - delta + y = xp.linspace(2., 100., n2) + self._testOne(x, y, 'two-sided', 55275.0 / lcm, 4.2188474935755949e-15, xp=xp) + self._testOne(x, y, 'greater', 561.0 / lcm, 0.99115454582047591, xp=xp) + self._testOne(x, y, 'less', 55275.0 / lcm, 3.1317328311518713e-26, xp=xp) + + def test_gh11184(self, xp): + # 3000, 3001, exact two-sided + rng = np.random.RandomState(123456) + x = rng.normal(size=3000).tolist() + y = rng.normal(size=3001, scale=1.5).tolist() + self._testOne(x, y, 'two-sided', 0.11292880151060758, 2.7755575615628914e-15, + mode='asymp', xp=xp) + self._testOne(x, y, 'two-sided', 0.11292880151060758, 2.7755575615628914e-15, + mode='exact', xp=xp) + + @pytest.mark.xslow + def test_gh11184_bigger(self, xp): + # 10000, 10001, exact two-sided + rng = np.random.RandomState(123456) + x = rng.normal(size=10000).tolist() + y = rng.normal(size=10001, scale=1.5).tolist() + self._testOne(x, y, 'two-sided', 0.10597913208679133, 3.3149311398483503e-49, + mode='asymp', xp=xp) + self._testOne(x, y, 'two-sided', 0.10597913208679133, 2.7755575615628914e-15, + mode='exact', xp=xp) + self._testOne(x, y, 'greater', 0.10597913208679133, 2.7947433906389253e-41, + mode='asymp', xp=xp) + self._testOne(x, y, 'less', 0.09658002199780022, 2.7947433906389253e-41, + mode='asymp', xp=xp) + + @pytest.mark.xslow + def test_gh12999(self, xp): + rng = np.random.RandomState(123456) + for x in range(1000, 12000, 1000): + vals1 = xp.asarray(rng.normal(size=(x))) + vals2 = xp.asarray(rng.normal(size=(x + 10), loc=0.5)) + exact = stats.ks_2samp(vals1, vals2, mode='exact').pvalue + asymp = stats.ks_2samp(vals1, vals2, mode='asymp').pvalue + # these two p-values should be in line with each other + xp_assert_less(exact, 3 * asymp) + xp_assert_less(asymp, 3 * exact) + + @pytest.mark.slow + def testLargeBoth(self, xp): + # 10000, 11000 + n1, n2 = 10000, 11000 + lcm = n1*11.0 + delta = 1.0/n1/n2/2/2 + x = xp.linspace(1, 200, n1) - delta + y = xp.linspace(2, 200, n2) + self._testOne(x, y, 'two-sided', 563.0 / lcm, 0.9990660108966576, + mode='asymp', xp=xp) + self._testOne(x, y, 'two-sided', 563.0 / lcm, 0.9990456491488628, + mode='exact', xp=xp) + self._testOne(x, y, 'two-sided', 563.0 / lcm, 0.9990660108966576, + mode='auto', xp=xp) + self._testOne(x, y, 'greater', 563.0 / lcm, 0.7561851877420673, xp=xp) + self._testOne(x, y, 'less', 10.0 / lcm, 0.9998239693191724, xp=xp) + with warnings.catch_warnings(): + message = "ks_2samp: Exact calculation unsuccessful." + warnings.filterwarnings("ignore", message, RuntimeWarning) + self._testOne(x, y, 'greater', 563.0 / lcm, 0.7561851877420673, + mode='exact', xp=xp) + self._testOne(x, y, 'less', 10.0 / lcm, 0.9998239693191724, + mode='exact', xp=xp) + + def testNamedAttributes(self, xp): + # test for namedtuple attribute results + attributes = ('statistic', 'pvalue') + res = stats.ks_2samp(xp.asarray([1, 2]), xp.asarray([3])) + check_named_results(res, attributes, xp=xp) + + @pytest.mark.slow + @skip_xp_backends(np_only=True) + def test_some_code_paths(self): + # Check that some code paths are executed + from scipy.stats._stats_py import ( + _count_paths_outside_method, + _compute_outer_prob_inside_method + ) + + _compute_outer_prob_inside_method(1, 1, 1, 1) + _count_paths_outside_method(1000, 1, 1, 1001) + + with np.errstate(invalid='raise'): + assert_raises(FloatingPointError, _count_paths_outside_method, + 1100, 1099, 1, 1) + assert_raises(FloatingPointError, _count_paths_outside_method, + 2000, 1000, 1, 1) + + @pytest.mark.parametrize('case', (([], [1.]), ([1.], []), ([], []))) + def test_argument_checking(self, case, xp): + # Check that an empty array warns + with eager_warns(SmallSampleWarning, match=too_small_1d_not_omit, xp=xp): + res = stats.ks_2samp(xp.asarray(case[0]), xp.asarray(case[1])) + xp_assert_equal(res.statistic, xp.asarray(np.nan)) + xp_assert_equal(res.pvalue, xp.asarray(np.nan)) + + @pytest.mark.xslow + def test_gh12218(self, xp): + """Ensure gh-12218 is fixed.""" + # gh-1228 triggered a TypeError calculating sqrt(n1*n2*(n1+n2)). + # n1, n2 both large integers, the product exceeded 2^64 + rng = np.random.default_rng(8751495592) + n1 = 2097152 # 2*^21 + rvs1 = xp.asarray(stats.uniform.rvs(size=n1, random_state=rng).tolist()) + rvs2 = rvs1 + 1 # Exact value of rvs2 doesn't matter. + stats.ks_2samp(rvs1, rvs2, alternative='greater', mode='asymp') + stats.ks_2samp(rvs1, rvs2, alternative='less', mode='asymp') + stats.ks_2samp(rvs1, rvs2, alternative='two-sided', mode='asymp') + + @pytest.mark.parametrize("dtype", [None, 'float32', 'float64']) + def test_warnings_gh_14019(self, dtype, xp): + dtype = xp_default_dtype(xp) if dtype is None else getattr(xp, dtype) + # Check that RuntimeWarning is raised when method='auto' and exact + # p-value calculation fails. See gh-14019. + rng = np.random.RandomState(seed=23493549) + # random samples of the same size as in the issue + data1 = xp.asarray(rng.random(size=881) + 0.5, dtype=dtype) + data2 = xp.asarray(rng.random(size=369), dtype=dtype) + message = "ks_2samp: Exact calculation unsuccessful" + with eager_warns(RuntimeWarning, match=message, xp=xp): + res = stats.ks_2samp(data1, data2, alternative='less') + xp_assert_close(res.pvalue, xp.asarray(0., dtype=dtype), atol=1e-14) + + @pytest.mark.parametrize("ksfunc", [stats.kstest, stats.ks_2samp]) + @pytest.mark.parametrize("alternative, x6val, ref_location, ref_sign", + [('greater', 5.9, 5.9, +1), + ('less', 6.1, 6.0, -1), + ('two-sided', 5.9, 5.9, +1), + ('two-sided', 6.1, 6.0, -1)]) + def test_location_sign(self, ksfunc, alternative, + x6val, ref_location, ref_sign, xp): + # Test that location and sign corresponding with statistic are as + # expected. (Test is designed to be easy to predict.) + x = xp.arange(10., dtype=xp.float64) + y = xp_copy(x) + x = xpx.at(x)[6].set(x6val) + res = stats.ks_2samp(x, y, alternative=alternative) + xp_assert_close(res.statistic, xp.asarray(0.1, dtype=xp.float64), rtol=1e-15) + assert res.statistic_location == ref_location + assert res.statistic_sign == ref_sign + + +@pytest.mark.filterwarnings('ignore:divide by zero:RuntimeWarning:dask') +@pytest.mark.filterwarnings('ignore:invalid value:RuntimeWarning:dask') +@make_xp_test_case(stats.ttest_rel) +class TestTTestRel: + def test_ttest_rel_xp(self, xp): + # stats.ttest_rel had no tests using the xp fixture. As a temporary + # measure to get tools/check_xp_untested.py to pass, a portion of + # test_ttest_rel has been converted. It might seem unnecessary to + # require tests for a trivial wrapper of a well-tested function, but + # this seems simpler than having a way to carve out exceptions to + # check_xp_untested. + tr = xp.asarray(0.81248591389165692, dtype=xp.float64) + pr = xp.asarray(0.41846234511362157, dtype=xp.float64) + + rvs1 = xp.linspace(1, 100, 100, dtype=xp.float64) + rvs2 = xp.linspace(1.01, 99.989, 100, dtype=xp.float64) + rvs1_2D = xp.stack( + [ + xp.linspace(1,100, 100, dtype=xp.float64), + xp.linspace(1.01, 99.989, 100, dtype=xp.float64) + ] + ) + rvs2_2D = xp.stack( + [ + xp.linspace(1.01, 99.989, 100, dtype=xp.float64), + xp.linspace(1, 100, 100, dtype=xp.float64) + ] + ) + + t, p = stats.ttest_rel(rvs1, rvs2, axis=0) + xp_assert_close(t, tr) + xp_assert_close(p, pr) + + t, p = stats.ttest_rel(rvs1_2D.T, rvs2_2D.T, axis=0) + xp_assert_close(t, xp.stack([tr, -tr])) + xp_assert_close(p, xp.stack([pr, pr])) + + t, p = stats.ttest_rel(rvs1_2D, rvs2_2D, axis=1) + xp_assert_close(t, xp.stack([tr, -tr])) + xp_assert_close(p, xp.stack([pr, pr])) + + def test_ttest_rel(self, xp): + dtype = xp.float64 # doesn't pass with default tolerances otherwise + + # regression test + tr = xp.asarray(0.81248591389165692, dtype=dtype) + pr = xp.asarray(0.41846234511362157, dtype=dtype) + tpr = xp.stack([tr, -tr]), xp.stack([pr, pr]) + + rvs1 = xp.linspace(1, 100, 100, dtype=dtype) + rvs2 = xp.linspace(1.01, 99.989, 100, dtype=dtype) + rvs1_2D = xp.stack([xp.linspace(1, 100, 100, dtype=dtype), + xp.linspace(1.01, 99.989, 100, dtype=dtype)]) + rvs2_2D = xp.stack([xp.linspace(1.01, 99.989, 100, dtype=dtype), + xp.linspace(1, 100, 100, dtype=dtype)]) + + t, p = stats.ttest_rel(rvs1, rvs2, axis=0) + xp_assert_close(t, tr) + xp_assert_close(p, pr) + t,p = stats.ttest_rel(rvs1_2D.T, rvs2_2D.T, axis=0) + xp_assert_close(t, tpr[0]) + xp_assert_close(p, tpr[1]) + t,p = stats.ttest_rel(rvs1_2D, rvs2_2D, axis=1) + xp_assert_close(t, tpr[0]) + xp_assert_close(p, tpr[1]) + + # test scalars + with warnings.catch_warnings(), \ + np.errstate(invalid="ignore", divide="ignore"): + warnings.filterwarnings( + "ignore", "Degrees of freedom <= 0 for slice", RuntimeWarning) + t, p = stats.ttest_rel(xp.asarray(4.), xp.asarray(3.)) + assert xp.isnan(t) + assert xp.isnan(p) + + # test for namedtuple attribute results + attributes = ('statistic', 'pvalue') + res = stats.ttest_rel(rvs1, rvs2, axis=0) + check_named_results(res, attributes, xp=xp) + + # test on 3 dimensions + rvs1_3D = xp.stack([rvs1_2D, rvs1_2D, rvs1_2D]) + rvs2_3D = xp.stack([rvs2_2D, rvs2_2D, rvs2_2D]) + t, p = stats.ttest_rel(rvs1_3D, rvs2_3D, axis=-1) + xp_assert_close(t, xp.stack((tpr[0],)*3)) + xp_assert_close(p, xp.stack((tpr[1],)*3)) + + t, p = stats.ttest_rel(xp.moveaxis(rvs1_3D, 2, 1), + xp.moveaxis(rvs2_3D, 2, 1), + axis=1) + xp_assert_close(t, xp.stack((tpr[0],)*3)) + xp_assert_close(p, xp.stack((tpr[1],)*3)) + + # test alternative parameter + message = "`alternative` must be 'less', 'greater', or 'two-sided'." + with pytest.raises(ValueError, match=message): + stats.ttest_rel(rvs1, rvs2, alternative="error") + + t, p = stats.ttest_rel(rvs1, rvs2, axis=0, alternative="less") + xp_assert_close(p, xp.asarray(1 - pr/2)) + xp_assert_close(t, xp.asarray(tr)) + + t, p = stats.ttest_rel(rvs1, rvs2, axis=0, alternative="greater") + xp_assert_close(p, xp.asarray(pr/2)) + xp_assert_close(t, xp.asarray(tr)) + + def test_nan_policy_propagate(self, xp): + # check nan policy + rng = np.random.RandomState(12345678) + x = stats.norm.rvs(loc=5, scale=10, size=501, random_state=rng) + x[500] = np.nan + y = (stats.norm.rvs(loc=5, scale=10, size=501, random_state=rng) + + stats.norm.rvs(scale=0.2, size=501, random_state=rng)) + y[500] = np.nan + x, y = xp.asarray(x.tolist()), xp.asarray(y.tolist()) + + res = stats.ttest_rel(x, x) + xp_assert_equal(res.statistic, xp.asarray(xp.nan)) + xp_assert_equal(res.pvalue, xp.asarray(xp.nan)) + + @skip_xp_backends(eager_only=True, reason='lazy->limited nan_policy support') + def test_nan_policy(self, xp): + # check nan policy + rng = np.random.RandomState(12345678) + x = stats.norm.rvs(loc=5, scale=10, size=501, random_state=rng) + x[500] = np.nan + y = (stats.norm.rvs(loc=5, scale=10, size=501, random_state=rng) + + stats.norm.rvs(scale=0.2, size=501, random_state=rng)) + y[500] = np.nan + x, y = xp.asarray(x.tolist()), xp.asarray(y.tolist()) + + res = stats.ttest_rel(x, y, nan_policy='omit') + xp_assert_close(res.statistic, xp.asarray(0.25299925303978066)) + xp_assert_close(res.pvalue, xp.asarray(0.8003729814201519)) + + message = "The input contains nan values" + with pytest.raises(ValueError, match=message): + stats.ttest_rel(x, y, nan_policy='raise') + + message = "nan_policy must be one of" + with pytest.raises(ValueError, match=message): + stats.ttest_rel(x, y, nan_policy='foobar') + + def test_edge_cases(self, xp): + # test zero division problem + with eager_warns(RuntimeWarning, match="Precision loss occurred", xp=xp): + t, p = stats.ttest_rel(xp.asarray([0, 0, 0]), xp.asarray([1, 1, 1])) + + xp_assert_equal(xp.abs(t), xp.asarray(xp.inf)) + xp_assert_equal(p, xp.asarray(0.0)) + + with np.errstate(invalid="ignore"): + res = stats.ttest_rel(xp.asarray([0, 0, 0]), xp.asarray([0, 0, 0])) + xp_assert_equal(res.statistic, xp.asarray(xp.nan)) + xp_assert_equal(res.pvalue, xp.asarray(xp.nan)) + + anan = xp.asarray([[1., np.nan], [-1., 1.]]) + res = stats.ttest_rel(anan, xp.zeros((2, 2))) + xp_assert_equal(res.statistic, xp.asarray([0, np.nan])) + xp_assert_equal(res.pvalue, xp.asarray([1, np.nan])) + + # test incorrect input shape raise an error + x = xp.arange(24) + message = ("Array shapes are incompatible for broadcasting.|" + "operands could not be broadcast together") + with pytest.raises(ValueError, match=message): + stats.ttest_rel(xp.reshape(x, (8, 3)), xp.reshape(x, (2, 3, 4))) + + @skip_xp_backends(np_only=True, reason="multidimensional 'omit' is np-only") + def test_multidimensional_nan_policy_omit(self, xp): + rvs1_2D = np.stack([np.linspace(1, 100, 100), + np.linspace(1.01, 99.989, 100)]) + rvs2_2D = np.stack([np.linspace(1.01, 99.989, 100), + np.linspace(1, 100, 100)]) + rvs1_2D[:, 20:30] = np.nan + rvs2_2D[:, 15:25] = np.nan + + # Convert from two-sided p-values to one sided using T result data. + @np.vectorize + def converter(t, p, alt): + if (t < 0 and alt == "less") or (t > 0 and alt == "greater"): + return p / 2 + return 1 - (p / 2) + + with pytest.warns(SmallSampleWarning, match=too_small_nd_omit): + tr, pr = stats.ttest_rel(rvs1_2D, rvs2_2D, 0, nan_policy='omit') + + with pytest.warns(SmallSampleWarning, match=too_small_nd_omit): + t, p = stats.ttest_rel(rvs1_2D, rvs2_2D, 0, + nan_policy='omit', alternative='less') + assert_allclose(t, tr, rtol=1e-14) + with np.errstate(invalid='ignore'): + assert_allclose(p, converter(tr, pr, 'less'), rtol=1e-14) + + with pytest.warns(SmallSampleWarning, match=too_small_nd_omit): + t, p = stats.ttest_rel(rvs1_2D, rvs2_2D, 0, + nan_policy='omit', alternative='greater') + assert_allclose(t, tr, rtol=1e-14) + with np.errstate(invalid='ignore'): + assert_allclose(p, converter(tr, pr, 'greater'), rtol=1e-14) + + @skip_xp_backends(eager_only=True, reason="lazy -> limited nan_policy support") + def test_ttest_rel_nan_2nd_arg(self, xp): + # regression test for gh-6134: nans in the second arg were not handled + x = xp.asarray([xp.nan, 2.0, 3.0, 4.0]) + y = xp.asarray([1.0, 2.0, 1.0, 2.0]) + + r1 = stats.ttest_rel(x, y, nan_policy='omit') + r2 = stats.ttest_rel(y, x, nan_policy='omit') + xp_assert_close(r2.statistic, -r1.statistic) + xp_assert_close(r2.pvalue, r1.pvalue) + + # NB: arguments are paired when NaNs are dropped + r3 = stats.ttest_rel(y[1:], x[1:]) + xp_assert_close(r2.statistic, r3.statistic) + xp_assert_close(r2.pvalue, r3.pvalue) + + # .. and this is consistent with R. R code: + # x = c(NA, 2.0, 3.0, 4.0) + # y = c(1.0, 2.0, 1.0, 2.0) + # t.test(x, y, paired=TRUE)), + xp_assert_close(r2.statistic, xp.asarray(-2.), atol=1e-4) + xp_assert_close(r2.pvalue, xp.asarray(0.1835), atol=1e-4) + + def test_ttest_rel_empty_1d_returns_nan(self, xp): + # Two empty inputs should return a TtestResult containing nan + # for both values. + with eager_warns(SmallSampleWarning, match=too_small_1d_not_omit, xp=xp): + result = stats.ttest_rel(xp.asarray([]), xp.asarray([])) + assert isinstance(result, stats._stats_py.TtestResult) + xp_assert_equal(result.statistic, xp.asarray(xp.nan)) + xp_assert_equal(result.pvalue, xp.asarray(xp.nan)) + + @skip_xp_backends("dask.array", reason="dask fails this one") + @pytest.mark.parametrize('b, expected_shape', + [(np.empty((1, 5, 0)), (3, 5)), + (np.empty((1, 0, 0)), (3, 0))]) + def test_ttest_rel_axis_size_zero(self, b, expected_shape, xp): + # In this test, the length of the axis dimension is zero. + # The results should be arrays containing nan with shape + # given by the broadcast nonaxis dimensions. + a = xp.empty((3, 1, 0)) + b = xp.asarray(b, dtype=a.dtype) + with warnings.catch_warnings(): + # first case should warn, second shouldn't? + warnings.filterwarnings("ignore", too_small_nd_not_omit, SmallSampleWarning) + result = stats.ttest_rel(a, b, axis=-1) + assert isinstance(result, stats._stats_py.TtestResult) + expected_value = xp.full(expected_shape, fill_value=xp.nan) + xp_assert_equal(result.statistic, expected_value) + xp_assert_equal(result.pvalue, expected_value) + + def test_ttest_rel_nonaxis_size_zero(self, xp): + # In this test, the length of the axis dimension is nonzero, + # but one of the nonaxis dimensions has length 0. Check that + # we still get the correctly broadcast shape, which is (5, 0) + # in this case. + a = xp.empty((1, 8, 0)) + b = xp.empty((5, 8, 1)) + result = stats.ttest_rel(a, b, axis=1) + assert isinstance(result, stats._stats_py.TtestResult) + assert result.statistic.shape == (5, 0) + assert result.pvalue.shape == (5, 0) + + @skip_xp_backends("jax.numpy", reason="needs stdtrit; lazy->can't use root finder") + @pytest.mark.parametrize("alternative", ['two-sided', 'less', 'greater']) + def test_ttest_rel_ci_1d(self, alternative, xp): + # test confidence interval method against reference values + rng = np.random.default_rng(3749065329432213059) + n = 10 + x = xp.asarray(rng.normal(size=n, loc=1.5, scale=2).tolist()) + y = xp.asarray(rng.normal(size=n, loc=2, scale=2).tolist()) + # Reference values generated with R t.test: + # options(digits=16) + # x = c(1.22825792, 1.63950485, 4.39025641, 0.68609437, 2.03813481, + # -1.20040109, 1.81997937, 1.86854636, 2.94694282, 3.94291373) + # y = c(3.49961496, 1.53192536, 5.53620083, 2.91687718, 0.04858043, + # 3.78505943, 3.3077496 , 2.30468892, 3.42168074, 0.56797592) + # t.test(x, y, paired=TRUE, conf.level=0.85, alternative='l') + + ref = {'two-sided': [-1.912194489914035, 0.400169725914035], + 'greater': [-1.563944820311475, np.inf], + 'less': [-np.inf, 0.05192005631147523]} + res = stats.ttest_rel(x, y, alternative=alternative) + ci = res.confidence_interval(confidence_level=0.85) + xp_assert_close(ci.low, xp.asarray(ref[alternative][0])) + xp_assert_close(ci.high, xp.asarray(ref[alternative][1])) + xp_assert_close(res.df, xp.asarray(n-1)) + + +@pytest.mark.parametrize("test_fun, args", [ + make_xp_pytest_param(stats.ttest_1samp, (np.arange(10), 0)), + make_xp_pytest_param(stats.ttest_rel, (np.arange(10), np.arange(10)))]) +def test_ttest_ci_iv(test_fun, args, xp): + # test `confidence_interval` method input validation + res = test_fun(*(xp.asarray(arg) for arg in args)) + message = '`confidence_level` must be a number between 0 and 1.' + with pytest.raises(ValueError, match=message): + res.confidence_interval(confidence_level=10) + + +def _desc_stats(x1, x2, axis=0, *, xp=None): + xp = array_namespace(x1, x2) if xp is None else xp + + def _stats(x, axis=0): + x = xp.asarray(x) + mu = xp.mean(x, axis=axis) + std = xp.std(x, axis=axis, correction=1) + nobs = x.shape[axis] + return mu, std, nobs + + return _stats(x1, axis) + _stats(x2, axis) + + +@make_xp_test_case(stats.ttest_ind, stats.ttest_ind_from_stats) +class TestTTestInd: + def test_ttest_ind(self, xp): + # regression test + tr = xp.asarray(1.0912746897927283) + pr = xp.asarray(0.27647818616351882) + tr_2D = xp.stack([tr, -tr]) + pr_2D = xp.stack([pr, pr]) + + rvs1 = xp.linspace(5, 105, 100) + rvs2 = xp.linspace(1, 100, 100) + rvs1_2D = xp.stack([rvs1, rvs2]) + rvs2_2D = xp.stack([rvs2, rvs1]) + + res = stats.ttest_ind(rvs1, rvs2, axis=0) + t, p = res # check that result object can be unpacked + xp_assert_close(t, tr) + xp_assert_close(p, pr) + + res = stats.ttest_ind_from_stats(*_desc_stats(rvs1, rvs2)) + t, p = res # check that result object can be unpacked + xp_assert_close(t, tr) + xp_assert_close(p, pr) + + res = stats.ttest_ind(rvs1_2D.T, rvs2_2D.T, axis=0) + xp_assert_close(res.statistic, tr_2D) + xp_assert_close(res.pvalue, pr_2D) + + res = stats.ttest_ind_from_stats(*_desc_stats(rvs1_2D.T, rvs2_2D.T)) + xp_assert_close(res.statistic, tr_2D) + xp_assert_close(res.pvalue, pr_2D) + + res = stats.ttest_ind(rvs1_2D, rvs2_2D, axis=1) + xp_assert_close(res.statistic, tr_2D) + xp_assert_close(res.pvalue, pr_2D) + + res = stats.ttest_ind_from_stats(*_desc_stats(rvs1_2D, rvs2_2D, axis=1)) + xp_assert_close(res.statistic, tr_2D) + xp_assert_close(res.pvalue, pr_2D) + + # test on 3 dimensions removed because generic tests in + # test_axis_nan_policy are much stronger + + # test alternative parameter + message = "`alternative` must be 'less', 'greater', or 'two-sided'." + with pytest.raises(ValueError, match=message): + stats.ttest_ind(rvs1, rvs2, alternative = "error") + + args = _desc_stats(rvs1_2D.T, rvs2_2D.T) + with pytest.raises(ValueError, match=message): + stats.ttest_ind_from_stats(*args, alternative = "error") + + t, p = stats.ttest_ind(rvs1, rvs2, alternative="less") + xp_assert_close(p, 1 - (pr/2)) + xp_assert_close(t, tr) + + t, p = stats.ttest_ind(rvs1, rvs2, alternative="greater") + xp_assert_close(p, pr/2) + xp_assert_close(t, tr) + + # Check that ttest_ind_from_stats agrees with ttest_ind + res1 = stats.ttest_ind(rvs1_2D.T, rvs2_2D.T, axis=0, alternative="less") + args = _desc_stats(rvs1_2D.T, rvs2_2D.T) + res2 = stats.ttest_ind_from_stats(*args, alternative="less") + xp_assert_close(res1.statistic, res2.statistic) + xp_assert_close(res1.pvalue, res2.pvalue) + + res1 = stats.ttest_ind(rvs1_2D.T, rvs2_2D.T, axis=0, alternative="less") + args = _desc_stats(rvs1_2D.T, rvs2_2D.T) + res2 = stats.ttest_ind_from_stats(*args, alternative="less") + xp_assert_close(res1.statistic, res2.statistic) + xp_assert_close(res1.pvalue, res2.pvalue) + + # test NaNs + NaN = xp.asarray(xp.nan) + rvs1 = xp.where(xp.arange(rvs1.shape[0]) == 0, NaN, rvs1) + + res = stats.ttest_ind(rvs1, rvs2, axis=0) + xp_assert_equal(res.statistic, NaN) + xp_assert_equal(res.pvalue, NaN) + + res = stats.ttest_ind_from_stats(*_desc_stats(rvs1, rvs2)) + xp_assert_equal(res.statistic, NaN) + xp_assert_equal(res.pvalue, NaN) + + @skip_xp_backends(eager_only=True, reason='lazy -> limited nan_policy support') + def test_special_cases(self, xp): + # check nan policy + rng = np.random.default_rng(22199381) + x = stats.norm.rvs(loc=5, scale=10, size=501, random_state=rng) + y = stats.norm.rvs(loc=5, scale=10, size=500, random_state=rng) + x[500] = np.nan + x, y = xp.asarray(x.tolist()), xp.asarray(y.tolist()) + + res = stats.ttest_ind(x, y) + xp_assert_equal(res.statistic, xp.asarray(xp.nan)) + xp_assert_equal(res.pvalue, xp.asarray(xp.nan)) + + # test zero division problem + with eager_warns(RuntimeWarning, match="Precision loss occurred", xp=xp): + t, p = stats.ttest_ind(xp.asarray([0, 0, 0]), xp.asarray([1, 1, 1])) + xp_assert_equal(t, xp.asarray(-xp.inf)) + xp_assert_equal(p, xp.asarray(0.)) + + x = xp.asarray([0, 0, 0]) + res = stats.ttest_ind(x, x) + xp_assert_equal(res.statistic, xp.asarray(xp.nan)) + xp_assert_equal(res.pvalue, xp.asarray(xp.nan)) + + # check that nan in input array result in nan output + anan = xp.asarray([[1, np.nan], [-1, 1]]) + res = stats.ttest_ind(anan, xp.zeros((2, 2))) + xp_assert_equal(res.statistic, xp.asarray([0, xp.nan])) + xp_assert_equal(res.pvalue, xp.asarray([1, xp.nan])) + + @skip_xp_backends(eager_only=True, reason='lazy -> limited nan_policy support') + def test_ttest_ind_nan_policy(self, xp): + rng = np.random.default_rng(61481957) + x = stats.norm.rvs(loc=5, scale=10, size=501, random_state=rng) + y = stats.norm.rvs(loc=5, scale=10, size=500, random_state=rng) + x[500] = np.nan + x, y = xp.asarray(x.tolist()), xp.asarray(y.tolist()) + + res = stats.ttest_ind(x, y, nan_policy='omit') + ref = stats.ttest_ind(x[~xp.isnan(x)], y, nan_policy='omit') + xp_assert_close(res.statistic, ref.statistic) + xp_assert_close(res.pvalue, ref.pvalue) + + message = "The input contains nan values" + with pytest.raises(ValueError, match=message): + stats.ttest_ind(x, y, nan_policy='raise') + + message = "nan_policy must be one of" + with pytest.raises(ValueError, match=message): + stats.ttest_ind(x, y, nan_policy='foobar') + + def test_3D_nan_policy_omit(self): # ND nan_policy='omit' is NumPy only + rvs1 = np.linspace(5, 105, 100) + rvs2 = np.linspace(1, 100, 100) + rvs1_2D = np.array([rvs1, rvs2]) + rvs2_2D = np.array([rvs2, rvs1]) + rvs1_3D = np.dstack([rvs1_2D, rvs1_2D, rvs1_2D]) + rvs2_3D = np.dstack([rvs2_2D, rvs2_2D, rvs2_2D]) + + rvs1_3D[:, :, 10:15] = np.nan + rvs2_3D[:, :, 6:12] = np.nan + + # Convert from two-sided p-values to one sided using T result data. + def convert(t, p, alt): + if (t < 0 and alt == "less") or (t > 0 and alt == "greater"): + return p / 2 + return 1 - (p / 2) + converter = np.vectorize(convert) + + tr, pr = stats.ttest_ind(rvs1_3D, rvs2_3D, axis=0, nan_policy='omit') + + t, p = stats.ttest_ind(rvs1_3D, rvs2_3D, axis=0, nan_policy='omit', + alternative='less') + assert_allclose(t, tr, rtol=1e-14) + assert_allclose(p, converter(tr, pr, 'less'), rtol=1e-14) + + t, p = stats.ttest_ind(rvs1_3D, rvs2_3D, axis=0, nan_policy='omit', + alternative='greater') + assert_allclose(t, tr, rtol=1e-14) + assert_allclose(p, converter(tr, pr, 'greater'), rtol=1e-14) + + + def test_ttest_ind_scalar(self): # ttest_ind for scalar input is legacy/np-only + # test scalars + with warnings.catch_warnings(), np.errstate(invalid="ignore"): + warnings.filterwarnings( + "ignore", "Degrees of freedom <= 0 for slice", RuntimeWarning) + t, p = stats.ttest_ind(4., 3.) + assert np.isnan(t) + assert np.isnan(p) + + +@make_xp_test_case(stats.ttest_ind) +@skip_xp_backends('dask.array', reason="Dask doesn't support ResamplingMethods") +@pytest.mark.filterwarnings("ignore:Arguments...:DeprecationWarning") +class TestTTestIndResampling: + N = 20 + + # data for most tests + rng = np.random.default_rng(169708062) + a = np.vstack((np.arange(3*N//4), rng.random(3*N//4))) + b = np.vstack((np.arange(N//4) + 100, rng.random(N//4))) + + # data for equal variance tests + a2 = np.arange(10) + b2 = np.arange(10) + 100 + + # data for exact test + a3 = [1, 2] + b3 = [3, 4] + + # data for bigger test + rvs1 = stats.norm.rvs(loc=5, scale=10, # type: ignore + size=500, random_state=rng).reshape(100, 5).T + rvs2 = stats.norm.rvs(loc=8, scale=20, size=100, random_state=rng) # type: ignore + + p_d = [1/1001, (676+1)/1001] # desired pvalues + p_d_gen = [1/1001, (672 + 1)/1001] # desired pvalues for Generator seed + p_d_big = [(993+1)/1001, (685+1)/1001, (840+1)/1001, + (955+1)/1001, (255+1)/1001] + + params = [ + (a, b, {"axis": 1}, p_d), # basic test + (a.T, b.T, {'axis': 0}, p_d), # along axis 0 + (a[0, :], b[0, :], {'axis': None}, p_d[0]), # 1d data + (a[0, :].tolist(), b[0, :].tolist(), {'axis': None}, p_d[0]), + # different seeds + (a, b, {'random_state': 0, "axis": 1}, p_d), + (a, b, {'random_state': np.random.RandomState(0), "axis": 1}, p_d), + (a2, b2, {'equal_var': True}, 1/1001), # equal variances + (rvs1, rvs2, {'axis': -1, 'random_state': 0}, p_d_big), # bigger test + (a3, b3, {}, 1/3), # exact test + (a, b, {'random_state': np.random.default_rng(0), "axis": 1}, p_d_gen), + ] + + @pytest.mark.parametrize("alternative", ['less', 'greater', 'two-sided']) + @pytest.mark.parametrize("shape", [(12,), (2, 12)]) + def test_permutation_method(self, alternative, shape, xp): + rng = np.random.default_rng(2348934579834565) + x = xp.asarray(rng.random(size=shape).tolist()) + y = xp.asarray(rng.random(size=13).tolist()) + + kwargs = dict(n_resamples=999) + + # Use ttest_ind with `method` + rng = np.random.default_rng(348934579834565) + method = stats.PermutationMethod(rng=rng, **kwargs) + res = stats.ttest_ind(x, y, axis=-1, alternative=alternative, method=method) + + # Use `permutation_test` directly + def statistic(x, y, axis): return stats.ttest_ind(x, y, axis=axis).statistic + rng = np.random.default_rng(348934579834565) + ref = stats.permutation_test((x, y), statistic, axis=-1, rng=rng, + alternative=alternative, **kwargs) + + xp_assert_close(res.statistic, ref.statistic) + xp_assert_close(res.pvalue, ref.pvalue) + + # Sanity check against theoretical t-test + ref = stats.ttest_ind(x, y, axis=-1, alternative=alternative) + xp_assert_close(res.statistic, ref.statistic) + xp_assert_close(res.pvalue, ref.pvalue, rtol=3e-2) + + @pytest.mark.parametrize("alternative", ['less', 'greater', 'two-sided']) + @pytest.mark.parametrize("shape", [(12,), (2, 12)]) + def test_monte_carlo_method(self, alternative, shape, xp): + rng = np.random.default_rng(2348934579834565) + x = xp.asarray(rng.random(size=shape).tolist()) + y = xp.asarray(rng.random(size=13).tolist()) + + kwargs = dict(n_resamples=999) + + # Use `monte_carlo` directly + def statistic(x, y, axis): + x, y = xp.asarray(x), xp.asarray(y) + return stats.ttest_ind(x, y, axis=axis).statistic + + rng = np.random.default_rng(348934579834565) + rvs = [rng.standard_normal, rng.standard_normal] + ref = stats.monte_carlo_test((x, y), rvs=rvs, statistic=statistic, axis=-1, + alternative=alternative, **kwargs) + + # Use ttest_ind with `method` + rng = np.random.default_rng(348934579834565) + rvs = [rng.standard_normal, rng.standard_normal] + method = stats.MonteCarloMethod(rvs=rvs, **kwargs) + res = stats.ttest_ind(x, y, axis=-1, alternative=alternative, method=method) + xp_assert_close(res.statistic, ref.statistic) + xp_assert_close(res.pvalue, ref.pvalue) + + # Passing `rng` instead of `rvs` + method = stats.MonteCarloMethod(rng=348934579834565, **kwargs) + res = stats.ttest_ind(x, y, axis=-1, alternative=alternative, method=method) + xp_assert_close(res.statistic, ref.statistic) + xp_assert_close(res.pvalue, ref.pvalue) + + # Sanity check against theoretical t-test + ref = stats.ttest_ind(x, y, axis=-1, alternative=alternative) + xp_assert_close(res.statistic, ref.statistic) + xp_assert_close(res.pvalue, ref.pvalue, rtol=6e-2) + + def test_resampling_input_validation(self, xp): + message = "`method` must be an instance of `PermutationMethod`, an instance..." + with pytest.raises(ValueError, match=message): + x, y = xp.asarray([1, 2, 3]), xp.asarray([4, 5, 6]) + stats.ttest_ind(x, y, method='migratory') + + +@make_xp_test_case(stats.ttest_ind) +class TestTTestIndCommon: + # for tests that are performed on variations of the t-test (e.g. trimmed) + @pytest.mark.parametrize("kwds", [{'trim': .2}, {}], + ids=["trim", "basic"]) + @pytest.mark.parametrize('equal_var', [True, False], + ids=['equal_var', 'unequal_var']) + def test_ttest_many_dims_fast(self, kwds, equal_var, xp): + # Test that ttest_ind works on many-dimensional arrays + rng = np.random.default_rng(3815288136) + a = xp.asarray(rng.random((5, 4, 4, 7, 1, 6)).tolist()) + b = xp.asarray(rng.random((4, 1, 8, 2, 6)).tolist()) + res = stats.ttest_ind(a, b, axis=-3, **kwds) + + # compare fully-vectorized t-test against t-test on smaller slice + i, j, k = 2, 3, 1 + a2 = a[i, :, j, :, 0, :] + b2 = b[:, 0, :, k, :] + res2 = stats.ttest_ind(a2, b2, axis=-2, **kwds) + xp_assert_close(res.statistic[i, :, j, k, :], res2.statistic) + xp_assert_close(res.pvalue[i, :, j, k, :], res2.pvalue) + + @pytest.mark.xslow + @pytest.mark.parametrize("kwds", [{'trim': .2}, {}], + ids=["trim", "basic"]) + @pytest.mark.parametrize('equal_var', [True, False], + ids=['equal_var', 'unequal_var']) + @skip_xp_backends('array_api_strict', reason="indexing in test is too fancy") + @skip_xp_backends('jax.numpy', reason="indexing in test is too fancy") + @skip_xp_backends('dask.array', reason="too slow") + def test_ttest_many_dims_full(self, kwds, equal_var, xp): + # compare against t-test on one axis-slice at a time + rng = np.random.default_rng(3815288136) + a = rng.random((5, 4, 4, 7, 1, 6)).tolist() + b = rng.random((4, 1, 8, 2, 6)).tolist() + res = stats.ttest_ind(xp.asarray(a), xp.asarray(b), axis=-3, **kwds) + + # manually broadcast with tile; move axis to end to simplify + x = np.moveaxis(np.tile(a, (1, 1, 1, 1, 2, 1)), -3, -1) + y = np.moveaxis(np.tile(b, (5, 1, 4, 1, 1, 1)), -3, -1) + x, y = xp.asarray(x.tolist()), xp.asarray(y.tolist()) + + shape = x.shape[:-1] + statistics = xp.zeros(shape) + pvalues = xp.zeros(shape) + for indices in product(*(range(i) for i in shape)): + xi = x[indices] # use tuple to index single axis slice + yi = y[indices] + res3 = stats.ttest_ind(xi, yi, axis=-1, **kwds) + statistics[indices] = res3.statistic + pvalues[indices] = res3.pvalue + + xp_assert_close(res.statistic, statistics) + xp_assert_close(res.pvalue, pvalues) + + @pytest.mark.parametrize("kwds", [{'trim': .2}, {}], + ids=["trim", "basic"]) + @pytest.mark.parametrize("axis", [-1, 0]) + def test_nans_on_axis(self, kwds, axis, xp): + # confirm that with `nan_policy='propagate'`, NaN results are returned + # on the correct location + rng = np.random.default_rng(363836384995579937222) + a = rng.integers(10, size=(5, 3, 10)).astype('float') + b = rng.integers(10, size=(5, 3, 10)).astype('float') + # set some indices in `a` and `b` to be `np.nan`. + a[0][2][3] = np.nan + b[2][0][6] = np.nan + a, b = xp.asarray(a.tolist()), xp.asarray(b.tolist()) + + # arbitrarily use `np.sum` as a baseline for which indices should be NaNs + expected = xp.isnan(xp.sum(a + b, axis=axis)) + + # multidimensional inputs to `t.sf(np.abs(t), df)` with NaNs on some + # indices throws an warning. See issue gh-13844 + with warnings.catch_warnings(), np.errstate(invalid="ignore"): + warnings.filterwarnings( + "ignore", "invalid value encountered in less_equal", RuntimeWarning) + warnings.filterwarnings("ignore", "Precision loss occurred", RuntimeWarning) + res = stats.ttest_ind(a, b, axis=axis, **kwds) + + p_nans = xp.isnan(res.pvalue) + xp_assert_equal(p_nans, expected) + statistic_nans = xp.isnan(res.statistic) + xp_assert_equal(statistic_nans, expected) + + +@make_xp_test_case(stats.ttest_ind) +class TestTTestTrimmed: + params = [ + [[1., 2., 3.], [1.1, 2.9, 4.2], 0.53619490753126731, -0.6864951273557258, .2], + [[56., 128.6, 12, 123.8, 64.34, 78, 763.3], [1.1, 2.9, 4.2], + 0.00998909252078421, 4.591598691181999, .2], + [[56., 128.6, 12., 123.8, 64.34, 78., 763.3], [1.1, 2.9, 4.2], + 0.10512380092302633, 2.832256715395378, .32], + [[2.7, 2.7, 1.1, 3.0, 1.9, 3.0, 3.8, 3.8, 0.3, 1.9, 1.9], + [6.5, 5.4, 8.1, 3.5, 0.5, 3.8, 6.8, 4.9, 9.5, 6.2, 4.1], + 0.002878909511344, -4.2461168970325, .2], + [[-0.84504783, 0.13366078, 3.53601757, -0.62908581, 0.54119466, + -1.16511574, -0.08836614, 1.18495416, 2.48028757, -1.58925028, + -1.6706357, 0.3090472, -2.12258305, 0.3697304, -1.0415207, + -0.57783497, -0.90997008, 1.09850192, 0.41270579, -1.4927376], + [1.2725522, 1.1657899, 2.7509041, 1.2389013, -0.9490494, -1.0752459, + 1.1038576, 2.9912821, 3.5349111, 0.4171922, 1.0168959, -0.7625041, + -0.4300008, 3.0431921, 1.6035947, 0.5285634, -0.7649405, 1.5575896, + 1.3670797, 1.1726023], 0.005293305834235, -3.0983317739483, .2]] + + @pytest.mark.parametrize("a, b, pr, tr, trim", params) + @pytest.mark.parametrize("dtype", [None, 'float32', 'float64']) + def test_ttest_compare_r(self, a, b, pr, tr, trim, dtype, xp): + ''' + Using PairedData's yuen.t.test method. Something to note is that there + are at least 3 R packages that come with a trimmed t-test method, and + comparisons were made between them. It was found that PairedData's + method's results match this method, SAS, and one of the other R + methods. A notable discrepancy was the DescTools implementation of the + function, which only sometimes agreed with SAS, WRS2, PairedData and + this implementation. For this reason, most comparisons in R are made + against PairedData's method. + + Rather than providing the input and output for all evaluations, here is + a representative example: + > library(PairedData) + > a <- c(1, 2, 3) + > b <- c(1.1, 2.9, 4.2) + > options(digits=16) + > yuen.t.test(a, b, tr=.2) + + Two-sample Yuen test, trim=0.2 + + data: x and y + t = -0.68649512735573, df = 3.4104431643464, p-value = 0.5361949075313 + alternative hypothesis: true difference in trimmed means is not equal + to 0 + 95 percent confidence interval: + -3.912777195645217 2.446110528978550 + sample estimates: + trimmed mean of x trimmed mean of y + 2.000000000000000 2.73333333333333 + ''' + dtype = dtype if dtype is None else getattr(xp, dtype) + a, b = xp.asarray(a, dtype=dtype), xp.asarray(b, dtype=dtype) + statistic, pvalue = stats.ttest_ind(a, b, trim=trim, equal_var=False) + xp_assert_close(statistic, xp.asarray(tr, dtype=dtype)) + xp_assert_close(pvalue, xp.asarray(pr, dtype=dtype)) + + def test_compare_SAS(self, xp): + # Source of the data used in this test: + # https://support.sas.com/resources/papers/proceedings14/1660-2014.pdf + a = [12, 14, 18, 25, 32, 44, 12, 14, 18, 25, 32, 44] + b = [17, 22, 14, 12, 30, 29, 19, 17, 22, 14, 12, 30, 29, 19] + a, b = xp.asarray(a), xp.asarray(b) + # In this paper, a trimming percentage of 5% is used. However, + # in their implementation, the number of values trimmed is rounded to + # the nearest whole number. However, consistent with + # `scipy.stats.trimmed_mean`, this test truncates to the lower + # whole number. In this example, the paper notes that 1 value is + # trimmed off of each side. 9% replicates this amount of trimming. + statistic, pvalue = stats.ttest_ind(a, b, trim=.09, equal_var=False) + xp_assert_close(pvalue, xp.asarray(0.514522), atol=1e-6) + xp_assert_close(statistic, xp.asarray(0.669169), atol=1e-6) + + def test_equal_var(self, xp): + ''' + The PairedData library only supports unequal variances. To compare + samples with equal variances, the multicon library is used. + > library(multicon) + > a <- c(2.7, 2.7, 1.1, 3.0, 1.9, 3.0, 3.8, 3.8, 0.3, 1.9, 1.9) + > b <- c(6.5, 5.4, 8.1, 3.5, 0.5, 3.8, 6.8, 4.9, 9.5, 6.2, 4.1) + > dv = c(a,b) + > iv = c(rep('a', length(a)), rep('b', length(b))) + > yuenContrast(dv~ iv, EQVAR = TRUE) + $Ms + N M wgt + a 11 2.442857142857143 1 + b 11 5.385714285714286 -1 + + $test + stat df crit p + results -4.246116897032513 12 2.178812829667228 0.00113508833897713 + ''' + a = [2.7, 2.7, 1.1, 3.0, 1.9, 3.0, 3.8, 3.8, 0.3, 1.9, 1.9] + b = [6.5, 5.4, 8.1, 3.5, 0.5, 3.8, 6.8, 4.9, 9.5, 6.2, 4.1] + a, b = xp.asarray(a), xp.asarray(b) + # `equal_var=True` is default + statistic, pvalue = stats.ttest_ind(a, b, trim=.2) + xp_assert_close(pvalue, xp.asarray(0.00113508833897713)) + xp_assert_close(statistic, xp.asarray(-4.246116897032513)) + + @pytest.mark.parametrize('alt,pr,tr', + (('greater', 0.9985605452443, -4.2461168970325), + ('less', 0.001439454755672, -4.2461168970325),), + ) + def test_alternatives(self, alt, pr, tr, xp): + ''' + > library(PairedData) + > a <- c(2.7,2.7,1.1,3.0,1.9,3.0,3.8,3.8,0.3,1.9,1.9) + > b <- c(6.5,5.4,8.1,3.5,0.5,3.8,6.8,4.9,9.5,6.2,4.1) + > options(digits=16) + > yuen.t.test(a, b, alternative = 'greater') + ''' + a = [2.7, 2.7, 1.1, 3.0, 1.9, 3.0, 3.8, 3.8, 0.3, 1.9, 1.9] + b = [6.5, 5.4, 8.1, 3.5, 0.5, 3.8, 6.8, 4.9, 9.5, 6.2, 4.1] + a, b = xp.asarray(a), xp.asarray(b) + statistic, pvalue = stats.ttest_ind(a, b, trim=.2, equal_var=False, + alternative=alt) + xp_assert_close(pvalue, xp.asarray(pr)) + xp_assert_close(statistic, xp.asarray(tr)) + + @pytest.mark.parametrize("trim", [-.2, .5, 1]) + def test_trim_bounds_error(self, trim, xp): + match = "Trimming percentage should be 0 <= `trim` < .5." + with assert_raises(ValueError, match=match): + stats.ttest_ind(xp.asarray([1, 2]), xp.asarray([2, 1]), trim=trim) + + +@make_xp_test_case(stats.ttest_ind) +class TestTTestCI: + # indices in order [alternative={two-sided, less, greater}, + # equal_var={False, True}, trim={0, 0.2}] + # reference values in order `statistic, df, pvalue, low, high` + # equal_var=False reference values computed with R PairedData yuen.t.test: + # + # library(PairedData) + # options(digits=16) + # a < - c(0.88236329, 0.97318744, 0.4549262, 0.97893335, 0.0606677, + # 0.44013366, 0.55806018, 0.40151434, 0.14453315, 0.25860601, + # 0.20202162) + # b < - c(0.93455277, 0.42680603, 0.49751939, 0.14152846, 0.711435, + # 0.77669667, 0.20507578, 0.78702772, 0.94691855, 0.32464958, + # 0.3873582, 0.35187468, 0.21731811) + # yuen.t.test(a, b, tr=0, conf.level = 0.9, alternative = 'l') + # + # equal_var=True reference values computed with R multicon yuenContrast: + # + # library(multicon) + # options(digits=16) + # a < - c(0.88236329, 0.97318744, 0.4549262, 0.97893335, 0.0606677, + # 0.44013366, 0.55806018, 0.40151434, 0.14453315, 0.25860601, + # 0.20202162) + # b < - c(0.93455277, 0.42680603, 0.49751939, 0.14152846, 0.711435, + # 0.77669667, 0.20507578, 0.78702772, 0.94691855, 0.32464958, + # 0.3873582, 0.35187468, 0.21731811) + # dv = c(a, b) + # iv = c(rep('a', length(a)), rep('b', length(b))) + # yuenContrast(dv~iv, EQVAR = FALSE, alternative = 'unequal', tr = 0.2) + r = np.empty(shape=(3, 2, 2, 5)) + r[0, 0, 0] = [-0.2314607, 19.894435, 0.8193209, -0.247220294, 0.188729943] + r[1, 0, 0] = [-0.2314607, 19.894435, 0.40966045, -np.inf, 0.1382426469] + r[2, 0, 0] = [-0.2314607, 19.894435, 0.5903395, -0.1967329982, np.inf] + r[0, 0, 1] = [-0.2452886, 11.427896, 0.8105823, -0.34057446, 0.25847383] + r[1, 0, 1] = [-0.2452886, 11.427896, 0.40529115, -np.inf, 0.1865829074] + r[2, 0, 1] = [-0.2452886, 11.427896, 0.5947089, -0.268683541, np.inf] + # confidence interval not available for equal_var=True + r[0, 1, 0] = [-0.2345625322555006, 22, 0.8167175905643815, np.nan, np.nan] + r[1, 1, 0] = [-0.2345625322555006, 22, 0.4083587952821908, np.nan, np.nan] + r[2, 1, 0] = [-0.2345625322555006, 22, 0.5916412047178092, np.nan, np.nan] + r[0, 1, 1] = [-0.2505369406507428, 14, 0.8058115135702835, np.nan, np.nan] + r[1, 1, 1] = [-0.2505369406507428, 14, 0.4029057567851417, np.nan, np.nan] + r[2, 1, 1] = [-0.2505369406507428, 14, 0.5970942432148583, np.nan, np.nan] + @pytest.mark.parametrize('alternative', ['two-sided', 'less', 'greater']) + @pytest.mark.parametrize('equal_var', [False, True]) + @pytest.mark.parametrize('trim', [0, 0.2]) + @skip_xp_backends('jax.numpy', reason='Generic stdtrit mutates array.') + def test_confidence_interval(self, alternative, equal_var, trim, xp): + if equal_var and trim: + pytest.xfail('Discrepancy in `main`; needs further investigation.') + + if trim and not is_numpy(xp): + pytest.skip('`trim` is only compatible with NumPy input') + + rng = np.random.default_rng(3810954496107292580) + x = xp.asarray(rng.random(11)) + y = xp.asarray(rng.random(13)) + + res = stats.ttest_ind(x, y, alternative=alternative, + equal_var=equal_var, trim=trim) + + alternatives = {'two-sided': 0, 'less': 1, 'greater': 2} + ref = self.r[alternatives[alternative], int(equal_var), int(np.ceil(trim))] + statistic, df, pvalue, low, high = ref + + rtol = 1e-7 # only 7 digits in reference + xp_assert_close(res.statistic, xp.asarray(statistic), rtol=rtol) + xp_assert_close(res.df, xp.asarray(df), rtol=rtol) + xp_assert_close(res.pvalue, xp.asarray(pvalue), rtol=rtol) + + if not equal_var: # CI not available when `equal_var is True` + ci = res.confidence_interval(0.9) + xp_assert_close(ci.low, xp.asarray(low), rtol=rtol) + xp_assert_close(ci.high, xp.asarray(high), rtol=rtol) + + +@make_xp_test_case(stats.ttest_ind) +class TestTTestIndMore: + @make_xp_test_case(stats.ttest_ind_from_stats) + def test_ttest_ind_with_uneq_var(self, xp): + # check vs. R `t.test`, e.g. + # options(digits=20) + # a = c(1., 2., 3.) + # b = c(1.1, 2.9, 4.2) + # t.test(a, b, equal.var=FALSE) + + a = xp.asarray([1., 2., 3.]) + b = xp.asarray([1.1, 2.9, 4.2]) + pr = xp.asarray(0.53619490753126686) + tr = xp.asarray(-0.686495127355726265) + + t, p = stats.ttest_ind(a, b, equal_var=False) + xp_assert_close(t, tr) + xp_assert_close(p, pr) + + t, p = stats.ttest_ind_from_stats(*_desc_stats(a, b), equal_var=False) + xp_assert_close(t, tr) + xp_assert_close(p, pr) + + a = xp.asarray([1., 2., 3., 4.]) + pr = xp.asarray(0.84354139131608252) + tr = xp.asarray(-0.210866331595072315) + + t, p = stats.ttest_ind(a, b, equal_var=False) + xp_assert_close(t, tr) + xp_assert_close(p, pr) + + t, p = stats.ttest_ind_from_stats(*_desc_stats(a, b), equal_var=False) + xp_assert_close(t, tr) + xp_assert_close(p, pr) + + # regression test + tr = xp.asarray(1.0912746897927283) + tr_uneq_n = xp.asarray(0.66745638708050492) + pr = xp.asarray(0.27647831993021388) + pr_uneq_n = xp.asarray(0.50873585065616544) + tr_2D = xp.stack([tr, -tr]) + pr_2D = xp.stack([pr, pr]) + + rvs3 = xp.linspace(1, 100, 25) + rvs2 = xp.linspace(1, 100, 100) + rvs1 = xp.linspace(5, 105, 100) + rvs1_2D = xp.stack([rvs1, rvs2]) + rvs2_2D = xp.stack([rvs2, rvs1]) + + t, p = stats.ttest_ind(rvs1, rvs2, axis=0, equal_var=False) + xp_assert_close(t, tr) + xp_assert_close(p, pr) + + t, p = stats.ttest_ind_from_stats(*_desc_stats(rvs1, rvs2), equal_var=False) + xp_assert_close(t, tr) + xp_assert_close(p, pr) + + t, p = stats.ttest_ind(rvs1, rvs3, axis=0, equal_var=False) + xp_assert_close(t, tr_uneq_n) + xp_assert_close(p, pr_uneq_n) + + t, p = stats.ttest_ind_from_stats(*_desc_stats(rvs1, rvs3), equal_var=False) + xp_assert_close(t, tr_uneq_n) + xp_assert_close(p, pr_uneq_n) + + res = stats.ttest_ind(rvs1_2D.T, rvs2_2D.T, axis=0, equal_var=False) + xp_assert_close(res.statistic, tr_2D) + xp_assert_close(res.pvalue, pr_2D) + + args = _desc_stats(rvs1_2D.T, rvs2_2D.T) + res = stats.ttest_ind_from_stats(*args, equal_var=False) + xp_assert_close(res.statistic, tr_2D) + xp_assert_close(res.pvalue, pr_2D) + + res = stats.ttest_ind(rvs1_2D, rvs2_2D, axis=1, equal_var=False) + xp_assert_close(res.statistic, tr_2D) + xp_assert_close(res.pvalue, pr_2D) + + args = _desc_stats(rvs1_2D, rvs2_2D, axis=1) + res = stats.ttest_ind_from_stats(*args, equal_var=False) + xp_assert_close(res.statistic, tr_2D) + xp_assert_close(res.pvalue, pr_2D) + + @pytest.mark.filterwarnings( + "ignore:divide by zero encountered:RuntimeWarning" + ) # for dask + @pytest.mark.filterwarnings( + "ignore:invalid value encountered:RuntimeWarning" + ) # for dask + def test_ttest_ind_zero_division(self, xp): + # test zero division problem + x = xp.zeros(3) + y = xp.ones(3) + + with eager_warns(RuntimeWarning, match="Precision loss occurred", xp=xp): + t, p = stats.ttest_ind(x, y, equal_var=False) + + xp_assert_equal(t, xp.asarray(-xp.inf)) + xp_assert_equal(p, xp.asarray(0.)) + + with np.errstate(all='ignore'): + t, p = stats.ttest_ind(x, x, equal_var=False) + xp_assert_equal(t, xp.asarray(xp.nan)) + xp_assert_equal(p, xp.asarray(xp.nan)) + + # check that nan in input array result in nan output + anan = xp.asarray([[1, xp.nan], [-1, 1]]) + t, p = stats.ttest_ind(anan, xp.zeros((2, 2)), equal_var=False) + xp_assert_equal(t, xp.asarray([0., np.nan])) + xp_assert_equal(p, xp.asarray([1., np.nan])) + + def test_ttest_ind_nan_2nd_arg(self): + # regression test for gh-6134: nans in the second arg were not handled + x = [np.nan, 2.0, 3.0, 4.0] + y = [1.0, 2.0, 1.0, 2.0] + + r1 = stats.ttest_ind(x, y, nan_policy='omit') + r2 = stats.ttest_ind(y, x, nan_policy='omit') + assert_allclose(r2.statistic, -r1.statistic, atol=1e-15) + assert_allclose(r2.pvalue, r1.pvalue, atol=1e-15) + + # NB: arguments are not paired when NaNs are dropped + r3 = stats.ttest_ind(y, x[1:]) + assert_allclose(r2, r3, atol=1e-15) + + # .. and this is consistent with R. R code: + # x = c(NA, 2.0, 3.0, 4.0) + # y = c(1.0, 2.0, 1.0, 2.0) + # t.test(x, y, var.equal=TRUE) + assert_allclose(r2, (-2.5354627641855498, 0.052181400457057901), + atol=1e-15) + + def test_ttest_ind_empty_1d_returns_nan(self, xp): + # Two empty inputs should return a TtestResult containing nan + # for both values. + with eager_warns(SmallSampleWarning, match=too_small_1d_not_omit, xp=xp): + res = stats.ttest_ind(xp.asarray([]), xp.asarray([])) + assert isinstance(res, stats._stats_py.TtestResult) + NaN = xp.asarray(xp.nan)[()] + xp_assert_equal(res.statistic, NaN) + xp_assert_equal(res.pvalue, NaN) + + @pytest.mark.parametrize('b, expected_shape', + [(np.empty((1, 5, 0)), (3, 5)), + (np.empty((1, 0, 0)), (3, 0))]) + def test_ttest_ind_axis_size_zero(self, b, expected_shape, xp): + # In this test, the length of the axis dimension is zero. + # The results should be arrays containing nan with shape + # given by the broadcast nonaxis dimensions. + a = xp.empty((3, 1, 0)) + b = xp.asarray(b, dtype=a.dtype) + with warnings.catch_warnings(): + # first case should warn, second shouldn't? + warnings.filterwarnings("ignore", too_small_nd_not_omit, SmallSampleWarning) + res = stats.ttest_ind(a, b, axis=-1) + assert isinstance(res, stats._stats_py.TtestResult) + expected_value = xp.full(expected_shape, fill_value=xp.nan) + xp_assert_equal(res.statistic, expected_value) + xp_assert_equal(res.pvalue, expected_value) + + def test_ttest_ind_nonaxis_size_zero(self, xp): + # In this test, the length of the axis dimension is nonzero, + # but one of the nonaxis dimensions has length 0. Check that + # we still get the correctly broadcast shape, which is (5, 0) + # in this case. + a = xp.empty((1, 8, 0)) + b = xp.empty((5, 8, 1)) + res = stats.ttest_ind(a, b, axis=1) + assert isinstance(res, stats._stats_py.TtestResult) + assert res.statistic.shape ==(5, 0) + assert res.pvalue.shape == (5, 0) + + def test_ttest_ind_nonaxis_size_zero_different_lengths(self, xp): + # In this test, the length of the axis dimension is nonzero, + # and that size is different in the two inputs, + # and one of the nonaxis dimensions has length 0. Check that + # we still get the correctly broadcast shape, which is (5, 0) + # in this case. + a = xp.empty((1, 7, 0)) + b = xp.empty((5, 8, 1)) + res = stats.ttest_ind(a, b, axis=1) + assert isinstance(res, stats._stats_py.TtestResult) + assert res.statistic.shape ==(5, 0) + assert res.pvalue.shape == (5, 0) + + @pytest.mark.skip_xp_backends(cpu_only=True, reason='Test uses ks_1samp') + @pytest.mark.filterwarnings("ignore:invalid value encountered:RuntimeWarning:dask") + @pytest.mark.filterwarnings("ignore:divide by zero encountered:RuntimeWarning:dask") + @pytest.mark.uses_xp_capabilities(False, reason="not used in this test yet") + def test_ttest_uniform_pvalues(self, xp): + # test that p-values are uniformly distributed under the null hypothesis + rng = np.random.default_rng(246834602926842) + x = xp.asarray(rng.normal(size=(10000, 2))) + y = xp.asarray(rng.normal(size=(10000, 1))) + q = rng.uniform(size=100) + + res = stats.ttest_ind(x, y, equal_var=True, axis=-1) + pvalue = np.asarray(res.pvalue) + # TODO: isolate use of alt backend to ttest_ind + assert stats.ks_1samp(pvalue, stats.uniform().cdf).pvalue > 0.1 + assert_allclose(np.quantile(pvalue, q), q, atol=1e-2) + + res = stats.ttest_ind(y, x, equal_var=True, axis=-1) + pvalue = np.asarray(res.pvalue) + assert stats.ks_1samp(pvalue, stats.uniform().cdf).pvalue > 0.1 + assert_allclose(np.quantile(pvalue, q), q, atol=1e-2) + + # reference values from R: + # options(digits=16) + # t.test(c(2, 3, 5), c(1.5), var.equal=TRUE) + x, y = xp.asarray([2., 3., 5.]), xp.asarray([1.5]) + + res = stats.ttest_ind(x, y, equal_var=True) + xp_assert_close(res.statistic, xp.asarray(1.0394023007754)) + xp_assert_close(res.pvalue, xp.asarray(0.407779907736)) + + +@make_xp_test_case(stats.ttest_ind_from_stats) +class TestTTestIndFromStats: + def test_gh5686(self, xp): + mean1, mean2 = xp.asarray([1, 2]), xp.asarray([3, 4]) + std1, std2 = xp.asarray([5, 3]), xp.asarray([4, 5]) + nobs1, nobs2 = xp.asarray([130, 140]), xp.asarray([100, 150]) + # This will raise a TypeError unless gh-5686 is fixed. + stats.ttest_ind_from_stats(mean1, std1, nobs1, mean2, std2, nobs2) + + @pytest.mark.filterwarnings("ignore:invalid value encountered:RuntimeWarning") + def test_ttest_ind_from_stats_inputs_zero(self, xp): + # Regression test for gh-6409. + zero = xp.asarray(0.) + six = xp.asarray(6.) + NaN = xp.asarray(xp.nan) + res = stats.ttest_ind_from_stats(zero, zero, six, zero, zero, six, + equal_var=False) + xp_assert_equal(res.statistic, NaN) + xp_assert_equal(res.pvalue, NaN) + + +@make_xp_test_case(stats.ttest_1samp) +class TestTTest1Samp: + def _convert_pvalue_alternative(self, t, p, alt, xp): + # test alternative parameter + # Convert from two-sided p-values to one sided using T result data. + less = xp.asarray(alt == "less") + greater = xp.asarray(alt == "greater") + i = ((t < 0) & less) | ((t > 0) & greater) + return xp.where(i, p/2, 1 - p/2) + + @pytest.mark.slow + @pytest.mark.filterwarnings("ignore:divide by zero encountered:RuntimeWarning:dask") + @pytest.mark.filterwarnings("ignore:invalid value encountered:RuntimeWarning:dask") + def test_ttest_1samp_new(self, xp): + rng = np.random.default_rng(88123290) + n1, n2, n3 = (10, 15, 20) + rvn1 = stats.norm.rvs(loc=5, scale=10, size=(n1, n2, n3), random_state=rng) + rvn1 = xp.asarray(rvn1) + + # check multidimensional array and correct axis handling + # deterministic rvn1 and rvn2 would be better as in test_ttest_rel + popmean = xp.ones((1, n2, n3)) + t1, p1 = stats.ttest_1samp(rvn1, popmean, axis=0) + t2, p2 = stats.ttest_1samp(rvn1, 1., axis=0) + t3, p3 = stats.ttest_1samp(rvn1[:, 0, 0], 1.) + xp_assert_close(t1, t2, rtol=1e-14) + xp_assert_close(t1[0, 0], t3, rtol=1e-14) + assert_equal(t1.shape, (n2, n3)) + + popmean = xp.ones((n1, 1, n3)) + t1, p1 = stats.ttest_1samp(rvn1, popmean, axis=1) + t2, p2 = stats.ttest_1samp(rvn1, 1., axis=1) + t3, p3 = stats.ttest_1samp(rvn1[0, :, 0], 1.) + xp_assert_close(t1, t2, rtol=1e-14) + xp_assert_close(t1[0, 0], t3, rtol=1e-14) + assert_equal(t1.shape, (n1, n3)) + + popmean = xp.ones((n1, n2, 1)) + t1, p1 = stats.ttest_1samp(rvn1, popmean, axis=2) + t2, p2 = stats.ttest_1samp(rvn1, 1., axis=2) + t3, p3 = stats.ttest_1samp(rvn1[0, 0, :], 1.) + xp_assert_close(t1, t2, rtol=1e-14) + xp_assert_close(t1[0, 0], t3, rtol=1e-14) + assert_equal(t1.shape, (n1, n2)) + + # test zero division problem + t, p = stats.ttest_1samp(xp.asarray([0., 0., 0.]), 1.) + xp_assert_equal(xp.abs(t), xp.asarray(xp.inf)) + xp_assert_equal(p, xp.asarray(0.)) + + tr, pr = stats.ttest_1samp(rvn1[:, :, :], 1.) + + t, p = stats.ttest_1samp(rvn1[:, :, :], 1., alternative="greater") + pc = self._convert_pvalue_alternative(tr, pr, "greater", xp) + xp_assert_close(p, pc) + xp_assert_close(t, tr) + + t, p = stats.ttest_1samp(rvn1[:, :, :], 1., alternative="less") + pc = self._convert_pvalue_alternative(tr, pr, "less", xp) + xp_assert_close(p, pc) + xp_assert_close(t, tr) + + with np.errstate(all='ignore'): + res = stats.ttest_1samp(xp.asarray([0., 0., 0.]), 0.) + xp_assert_equal(res.statistic, xp.asarray(xp.nan)) + xp_assert_equal(res.pvalue, xp.asarray(xp.nan)) + + # check that nan in input array result in nan output + anan = xp.asarray([[1., np.nan], [-1., 1.]]) + res = stats.ttest_1samp(anan, 0.) + xp_assert_equal(res.statistic, xp.asarray([0., xp.nan])) + xp_assert_equal(res.pvalue, xp.asarray([1., xp.nan])) + + @skip_xp_backends(eager_only=True, reason="lazy -> reduced nan_policy capabilities") + def test_ttest_1samp_new_omit(self, xp): + rng = np.random.default_rng(4008400329) + n1, n2, n3 = (5, 10, 15) + rvn1 = stats.norm.rvs(loc=5, scale=10, size=(n1, n2, n3), random_state=rng) + rvn1 = xp.asarray(rvn1) + + rvn1[0:2, 1:3, 4:8] = xp.nan + + tr, pr = stats.ttest_1samp(rvn1[:, :, :], 1., nan_policy='omit') + + t, p = stats.ttest_1samp(rvn1[:, :, :], 1., nan_policy='omit', + alternative="greater") + pc = self._convert_pvalue_alternative(tr, pr, "greater", xp) + xp_assert_close(p, pc) + xp_assert_close(t, tr) + + t, p = stats.ttest_1samp(rvn1[:, :, :], 1., nan_policy='omit', + alternative="less") + pc = self._convert_pvalue_alternative(tr, pr, "less", xp) + xp_assert_close(p, pc) + xp_assert_close(t, tr) + + @pytest.mark.skip_xp_backends('jax.numpy', reason='Generic stdtrit mutates array.') + def test_ttest_1samp_popmean_array(self, xp): + # when popmean.shape[axis] != 1, raise an error + # if the user wants to test multiple null hypotheses simultaneously, + # use standard broadcasting rules + rng = np.random.default_rng(2913300596553337193) + x = rng.random(size=(1, 15, 20)) + x = xp.asarray(x) + + message = r"`popmean.shape\[axis\]` must equal 1." + popmean = xp.asarray(rng.random(size=(5, 2, 20))) + with pytest.raises(ValueError, match=message): + stats.ttest_1samp(x, popmean=popmean, axis=-2) + + popmean = xp.asarray(rng.random(size=(5, 1, 20))) + res = stats.ttest_1samp(x, popmean=popmean, axis=-2) + assert res.statistic.shape == (5, 20) + + l, u = res.confidence_interval() + l = xp.expand_dims(l, axis=-2) + u = xp.expand_dims(u, axis=-2) + + res = stats.ttest_1samp(x, popmean=l, axis=-2) + ref = xp.broadcast_to(xp.asarray(0.05, dtype=xp.float64), res.pvalue.shape) + xp_assert_close(res.pvalue, ref) + + res = stats.ttest_1samp(x, popmean=u, axis=-2) + xp_assert_close(res.pvalue, ref) + + +@make_xp_test_case(stats.describe) +class TestDescribe: + @pytest.mark.filterwarnings("ignore:invalid value encountered:RuntimeWarning:dask") + @pytest.mark.filterwarnings("ignore:divide by zero encountered:RuntimeWarning:dask") + def test_describe_scalar(self, xp): + with warnings.catch_warnings(), \ + np.errstate(invalid="ignore", divide="ignore"): + warnings.filterwarnings( + "ignore", "Degrees of freedom <= 0 for slice", RuntimeWarning) + n, mm, m, v, sk, kurt = stats.describe(xp.asarray(4.)[()]) + assert n == 1 + xp_assert_equal(mm[0], xp.asarray(4.0)) + xp_assert_equal(mm[1], xp.asarray(4.0)) + xp_assert_equal(m, xp.asarray(4.0)) + xp_assert_equal(v ,xp.asarray(xp.nan)) + xp_assert_equal(sk, xp.asarray(xp.nan)) + xp_assert_equal(kurt, xp.asarray(xp.nan)) + + def test_describe_numbers(self, xp): + x = xp.concat((xp.ones((3, 4)), xp.full((2, 4), 2.))) + nc = 5 + mmc = (xp.asarray([1., 1., 1., 1.]), xp.asarray([2., 2., 2., 2.])) + mc = xp.asarray([1.4, 1.4, 1.4, 1.4]) + vc = xp.asarray([0.3, 0.3, 0.3, 0.3]) + skc = xp.asarray([0.40824829046386357] * 4) + kurtc = xp.asarray([-1.833333333333333] * 4) + n, mm, m, v, sk, kurt = stats.describe(x) + assert n == nc + xp_assert_equal(mm[0], mmc[0]) + xp_assert_equal(mm[1], mmc[1]) + xp_assert_close(m, mc, rtol=4 * xp.finfo(m.dtype).eps) + xp_assert_close(v, vc, rtol=4 * xp.finfo(m.dtype).eps) + xp_assert_close(sk, skc) + xp_assert_close(kurt, kurtc) + + n, mm, m, v, sk, kurt = stats.describe(x.T, axis=1) + assert n == nc + xp_assert_equal(mm[0], mmc[0]) + xp_assert_equal(mm[1], mmc[1]) + xp_assert_close(m, mc, rtol=4 * xp.finfo(m.dtype).eps) + xp_assert_close(v, vc, rtol=4 * xp.finfo(m.dtype).eps) + xp_assert_close(sk, skc) + xp_assert_close(kurt, kurtc) + + def describe_nan_policy_omit_test(self): + x = np.arange(10.) + x[9] = np.nan + + nc, mmc = (9, (0.0, 8.0)) + mc = 4.0 + vc = 7.5 + skc = 0.0 + kurtc = -1.2300000000000002 + n, mm, m, v, sk, kurt = stats.describe(x, nan_policy='omit') + assert_equal(n, nc) + assert_equal(mm, mmc) + assert_equal(m, mc) + assert_equal(v, vc) + assert_array_almost_equal(sk, skc) + assert_array_almost_equal(kurt, kurtc, decimal=13) + + def test_describe_nan_policy_other(self, xp): + x = xp.arange(10.) + x = xp.where(x==9, xp.nan, x) + + if is_lazy_array(x): + with pytest.raises(TypeError, match='not supported for lazy arrays'): + stats.describe(x, nan_policy='raise') + else: + with pytest.raises(ValueError, match='The input contains nan values'): + stats.describe(x, nan_policy='raise') + + n, mm, m, v, sk, kurt = stats.describe(x, nan_policy='propagate') + ref = xp.asarray(xp.nan)[()] + assert n == 10 + xp_assert_equal(mm[0], ref) + xp_assert_equal(mm[1], ref) + xp_assert_equal(m, ref) + xp_assert_equal(v, ref) + xp_assert_equal(sk, ref) + xp_assert_equal(kurt, ref) + + if is_numpy(xp): + self.describe_nan_policy_omit_test() + elif is_lazy_array(x): + with pytest.raises(TypeError, match='not supported for lazy arrays'): + stats.describe(x, nan_policy='omit') + + message = 'nan_policy must be one of...' + with pytest.raises(ValueError, match=message): + stats.describe(x, nan_policy='foobar') + + def test_describe_result_attributes(self): + # some result attributes are tuples, which aren't meant to be compared + # with `xp_assert_close` + actual = stats.describe(np.arange(5.)) + attributes = ('nobs', 'minmax', 'mean', 'variance', 'skewness', 'kurtosis') + check_named_results(actual, attributes) + + def test_describe_ddof(self, xp): + x = xp.concat((xp.ones((3, 4)), xp.full((2, 4), 2.))) + nc = 5 + mmc = (xp.asarray([1., 1., 1., 1.]), xp.asarray([2., 2., 2., 2.])) + mc = xp.asarray([1.4, 1.4, 1.4, 1.4]) + vc = xp.asarray([0.24, 0.24, 0.24, 0.24]) + skc = xp.asarray([0.40824829046386357] * 4) + kurtc = xp.asarray([-1.833333333333333] * 4) + n, mm, m, v, sk, kurt = stats.describe(x, ddof=0) + assert n == nc + xp_assert_equal(mm[0], mmc[0]) + xp_assert_equal(mm[1], mmc[1]) + xp_assert_close(m, mc) + xp_assert_close(v, vc) + xp_assert_close(sk, skc) + xp_assert_close(kurt, kurtc) + + def test_describe_axis_none(self, xp): + x = xp.concat((xp.ones((3, 4)), xp.full((2, 4), 2.))) + + # expected values + nc = 20 + mmc = (xp.asarray(1.0), xp.asarray(2.0)) + mc = xp.asarray(1.3999999999999999) + vc = xp.asarray(0.25263157894736848) + skc = xp.asarray(0.4082482904638634) + kurtc = xp.asarray(-1.8333333333333333) + + # actual values + n, mm, m, v, sk, kurt = stats.describe(x, axis=None) + + assert n == nc + xp_assert_equal(mm[0], mmc[0]) + xp_assert_equal(mm[1], mmc[1]) + xp_assert_close(m, mc) + xp_assert_close(v, vc) + xp_assert_close(sk, skc) + xp_assert_close(kurt, kurtc) + + def test_describe_empty(self, xp): + message = "The input must not be empty." + with pytest.raises(ValueError, match=message): + stats.describe(xp.asarray([])) + + @skip_xp_backends("array_api_strict", reason="min/max don't accept boolean input") + def test_boolean_input_gh24847(self, xp): + # gh-24847 reported that `describe` failed with boolean input + rng = np.random.default_rng(3248923598734583) + n = 20 + x = xp.asarray(rng.random(n) < 0.5) + res = stats.describe(x) + xp_assert_equal(res.nobs, xp.asarray(n, dtype=xp.int64)) + xp_assert_equal(res.minmax[0], xp.min(x)) + xp_assert_equal(res.minmax[1], xp.max(x)) + x = xp_promote(x, force_floating=True, xp=xp) + xp_assert_close(res.mean, xp.mean(x)) + xp_assert_close(res.variance, xp.var(x, correction=1)) + xp_assert_close(res.skewness, stats.skew(x)) + xp_assert_close(res.kurtosis, stats.kurtosis(x)) + + +class NormalityTests: + + @pytest.mark.parametrize("alternative", ['two-sided', 'less', 'greater']) + def test_against_R(self, alternative, xp): + # testa against R `dagoTest` from package `fBasics` + # library(fBasics) + # options(digits=16) + # x = c(-2, -1, 0, 1, 2, 3)**2 + # x = rep(x, times=4) + # test_result <- dagoTest(x) + # test_result@test$statistic + # test_result@test$p.value + test_name = self.test_name + test_fun = getattr(stats, test_name) + ref_statistic= xp.asarray(self.case_ref[0]) + ref_pvalue = xp.asarray(self.case_ref[1]) + + kwargs = {} + if alternative in {'less', 'greater'}: + if test_name in {'skewtest', 'kurtosistest'}: + ref_pvalue = ref_pvalue/2 if alternative == "less" else 1-ref_pvalue/2 + ref_pvalue = 1-ref_pvalue if test_name == 'skewtest' else ref_pvalue + kwargs['alternative'] = alternative + else: + pytest.skip('`alternative` not available for `normaltest`') + + x = xp.asarray((-2, -1, 0, 1, 2, 3.)*4)**2 + res = test_fun(x, **kwargs) + res_statistic, res_pvalue = res + xp_assert_close(res_statistic, ref_statistic) + xp_assert_close(res_pvalue, ref_pvalue) + check_named_results(res, ('statistic', 'pvalue'), xp=xp) + + def test_nan(self, xp): + # nan in input -> nan output (default nan_policy='propagate') + test_fun = getattr(stats, self.test_name) + x = xp.arange(30.) + NaN = xp.asarray(xp.nan, dtype=x.dtype) + x = xp.where(x == 29, NaN, x) + with np.errstate(invalid="ignore"): + res = test_fun(x) + xp_assert_equal(res.statistic, NaN) + xp_assert_equal(res.pvalue, NaN) + + +@make_xp_test_case(stats.skewtest) +class TestSkewTest(NormalityTests): + test_name = 'skewtest' + case_ref = (1.98078826090875881, 0.04761502382843208) # statistic, pvalue + + def test_intuitive(self, xp): + # intuitive tests; see gh-13549. skewnorm with parameter 1 has skew > 0 + a1 = stats.skewnorm.rvs(a=1, size=10000, random_state=123) + a1_xp = xp.asarray(a1) + pval = stats.skewtest(a1_xp, alternative='greater').pvalue + xp_assert_close(pval, xp.asarray(0.0, dtype=a1_xp.dtype), atol=9e-6) + + def test_skewtest_too_few_observations(self, xp): + # Regression test for ticket #1492. + # skewtest requires at least 8 observations; 7 should warn and return NaN. + stats.skewtest(xp.arange(8.0)) + + x = xp.arange(7.0) + + with eager_warns(SmallSampleWarning, match=too_small_1d_not_omit, xp=xp): + res = stats.skewtest(x) + NaN = xp.asarray(xp.nan) + xp_assert_equal(res.statistic, NaN) + xp_assert_equal(res.pvalue, NaN) + + +@make_xp_test_case(stats.kurtosistest) +class TestKurtosisTest(NormalityTests): + test_name = 'kurtosistest' + case_ref = (-0.01403734404759738, 0.98880018772590561) # statistic, pvalue + + def test_intuitive(self, xp): + # intuitive tests; see gh-13549. excess kurtosis of laplace is 3 > 0 + a2 = stats.laplace.rvs(size=10000, random_state=123) + a2_xp = xp.asarray(a2) + pval = stats.kurtosistest(a2_xp, alternative='greater').pvalue + xp_assert_close(pval, xp.asarray(0.0, dtype=a2_xp.dtype), atol=1e-15) + + def test_gh9033_regression(self, xp): + # regression test for issue gh-9033: x clearly non-normal but power of + # negative denom needs to be handled correctly to reject normality + counts = [128, 0, 58, 7, 0, 41, 16, 0, 0, 167] + x = np.hstack([np.full(c, i) for i, c in enumerate(counts)]) + x = xp.asarray(x, dtype=xp.float64) + assert stats.kurtosistest(x)[1] < 0.01 + + def test_kurtosistest_too_few_observations(self, xp): + # kurtosistest requires at least 5 observations; 4 should warn and return NaN. + # Regression test for ticket #1425. + stats.kurtosistest(xp.arange(5.0)) + + with eager_warns(SmallSampleWarning, match=too_small_1d_not_omit, xp=xp): + res = stats.kurtosistest(xp.arange(4.)) + NaN = xp.asarray(xp.nan) + xp_assert_equal(res.statistic, NaN) + xp_assert_equal(res.pvalue, NaN) + + +@make_xp_test_case(stats.normaltest) +class TestNormalTest(NormalityTests): + test_name = 'normaltest' + case_ref = (3.92371918158185551, 0.14059672529747502) # statistic, pvalue + + def test_too_few_observations(self, xp): + stats.normaltest(xp.arange(8.)) + + # 1D sample has too few observations -> warning / NaN output + with eager_warns(SmallSampleWarning, match=too_small_1d_not_omit, xp=xp): + res = stats.normaltest(xp.arange(7.)) + NaN = xp.asarray(xp.nan) + xp_assert_equal(res.statistic, NaN) + xp_assert_equal(res.pvalue, NaN) + + +class TestRankSums: + + @pytest.mark.parametrize('alternative', ['less', 'greater', 'two-sided']) + def test_ranksums_is_mannwhitneyu(self, alternative): + rng = np.random.default_rng(3417115752) + x, y = rng.random((2, 10)) + + # ranksums pval = mannwhitneyu pval w/out continuity or tie correction + res1 = stats.ranksums(x, y, alternative=alternative) + res2 = stats.mannwhitneyu(x, y, use_continuity=False, + alternative=alternative, method='asymptotic') + assert_allclose(res1.pvalue, res2.pvalue) + assert_allclose(res1.statistic, res2.zstatistic) + + def test_ranksums_named_results(self): + res = stats.ranksums(np.arange(10), np.arange(10)) + check_named_results(res, ('statistic', 'pvalue')) + + def test_input_validation(self): + with assert_raises(ValueError, match="`alternative` must be 'less'"): + stats.ranksums(np.arange(10), np.arange(10), alternative='foobar') + + +@make_xp_test_case(stats.jarque_bera) +class TestJarqueBera: + def test_jarque_bera_against_R(self, xp): + # library(tseries) + # options(digits=16) + # x < - rnorm(5) + # jarque.bera.test(x) + x = [-0.160104223201523288, 1.131262000934478040, -0.001235254523709458, + -0.776440091309490987, -2.072959999533182884] + x = xp.asarray(x) + ref = xp.asarray([0.17651605223752, 0.9155246169805]) + res = stats.jarque_bera(x) + xp_assert_close(res.statistic, ref[0]) + xp_assert_close(res.pvalue, ref[1]) + + @skip_xp_backends(np_only=True) + def test_jarque_bera_array_like(self, xp): + # array-like only relevant for NumPy + rng = np.random.default_rng(9294968266) + x = rng.standard_normal(size=100000) + + jb_test1 = JB1, p1 = stats.jarque_bera(list(x)) + jb_test2 = JB2, p2 = stats.jarque_bera(tuple(x)) + jb_test3 = JB3, p3 = stats.jarque_bera(x.reshape(2, 50000)) + + assert JB1 == JB2 == JB3 == jb_test1.statistic == jb_test2.statistic == jb_test3.statistic # noqa: E501 + assert p1 == p2 == p3 == jb_test1.pvalue == jb_test2.pvalue == jb_test3.pvalue + + def test_jarque_bera_too_few_observations(self, xp): + x = xp.asarray([]) + + with pytest.warns(SmallSampleWarning, match=too_small_1d_not_omit): + res = stats.jarque_bera(x) + NaN = xp.asarray(xp.nan) + xp_assert_equal(res.statistic, NaN) + xp_assert_equal(res.pvalue, NaN) + + def test_axis(self, xp): + rng = np.random.RandomState(seed=122398129) + x = xp.asarray(rng.random(size=(2, 45))) + + res = stats.jarque_bera(x, axis=None) + ref = stats.jarque_bera(xp.reshape(x, (-1,))) + xp_assert_equal(res.statistic, ref.statistic) + xp_assert_equal(res.pvalue, ref.pvalue) + + res = stats.jarque_bera(x, axis=1) + s0, p0 = stats.jarque_bera(x[0, :]) + s1, p1 = stats.jarque_bera(x[1, :]) + xp_assert_close(res.statistic, xp.stack([s0, s1])) + xp_assert_close(res.pvalue, xp.stack([p0, p1])) + + resT = stats.jarque_bera(x.T, axis=0) + xp_assert_close(res.statistic, resT.statistic) + xp_assert_close(res.pvalue, resT.pvalue) + + +@make_xp_test_case(stats.pointbiserialr) +class TestPointBiserialR: + def test_pointbiserial(self, xp): + # same as mstats test except for the nan + # Test data: https://web.archive.org/web/20060504220742/https://support.sas.com/ctx/samples/index.jsp?sid=490&tab=output + x = [1,0,1,1,1,1,0,1,0,0,0,1,1,0,0,0,1,1,1,0,0,0,0,0,0,0,0,1,0, + 0,0,0,0,1.] + y = [14.8,13.8,12.4,10.1,7.1,6.1,5.8,4.6,4.3,3.5,3.3,3.2,3.0, + 2.8,2.8,2.5,2.4,2.3,2.1,1.7,1.7,1.5,1.3,1.3,1.2,1.2,1.1, + 0.8,0.7,0.6,0.5,0.2,0.2,0.1] + res = stats.pointbiserialr(xp.asarray(x), xp.asarray(y)) + xp_assert_close(res[0], xp.asarray(0.36149), atol=1e-5) + + # test for namedtuple attribute results + attributes = ('correlation', 'pvalue') + check_named_results(res, attributes, xp=xp) + xp_assert_equal(res.correlation, res.statistic) + + +@make_xp_test_case(stats.obrientransform) +class TestObrientransform: + def test_basic(self, xp): + # A couple tests calculated by hand. + x1 = xp.asarray([0, 2, 4]) + t1 = stats.obrientransform(x1) + expected = xp.asarray([7., -2., 7.]) + xp_assert_close(t1[0][:], expected) + + x2 = xp.asarray([0, 3, 6, 9]) + t2 = stats.obrientransform(x2) + expected = xp.asarray([30., 0., 0., 30.]) + xp_assert_close(t2[0][:], expected) + + # Test two arguments. + a, b = stats.obrientransform(x1, x2) + xp_assert_close(a, t1[0][:]) + xp_assert_close(b, t2[0][:]) + + # Test three arguments. + a, b, c = stats.obrientransform(x1, x2, x1) + xp_assert_close(a, t1[0][:]) + xp_assert_close(b, t2[0][:]) + xp_assert_close(c, t1[0][:]) + + def test_something(self, xp): + # This is a regression test to check np.var replacement. + # The author of this test didn't separately verify the numbers. + x1 = xp.arange(5) + ref = xp.asarray( + [[5.41666667, 1.04166667, -0.41666667, 1.04166667, 5.41666667], + [21.66666667, 4.16666667, -1.66666667, 4.16666667, 21.66666667]]) + res = stats.obrientransform(x1, 2*x1) + xp_assert_close(res[0], ref[0, ...]) + xp_assert_close(res[1], ref[1, ...]) + + def test_reference(self, xp): + # Example from "O'Brien Test for Homogeneity of Variance" by Herve Abdi. + values = np.arange(5, 11) + reps = [5, 11, 9, 3, 2, 2] + data = np.repeat(values, reps) + transformed_values = np.asarray([3.1828, 0.5591, 0.0344, + 1.6086, 5.2817, 11.0538]) + expected = np.repeat(transformed_values, reps) + result = stats.obrientransform(xp.asarray(data)) + xp_assert_close(result[0][:], xp.asarray(expected.tolist()), rtol=1e-3) + + def test_nan_policy(self, xp): + rng = np.random.default_rng(4284359689201882838835) + x = rng.random(10) + x[3] = np.nan + x = xp.asarray(x) + + # nan_policy='raise' + if not is_lazy_array(x): + with pytest.raises(ValueError, match="The input contains nan values"): + stats.obrientransform(x, nan_policy='raise') + else: + with pytest.raises(TypeError, match="nan_policy='raise' is not supported"): + stats.obrientransform(x, nan_policy='raise') + + # nan_policy='propagate' + res = stats.obrientransform(x, nan_policy='propagate')[0] + xp_assert_equal(res, xp.full_like(x, xp.nan)) + + # nan_policy='omit' + i = xp.isnan(x) + res = stats.obrientransform(x, nan_policy='omit')[0] + ref = stats.obrientransform(x[~i])[0] + xp_assert_equal(res[i], x[i]) + xp_assert_close(res[~i], ref) + + +def check_equal_xmean(*args, xp, mean_fun, axis=None, dtype=None, + rtol=None, weights=None): + # Note this doesn't test when axis is not specified + dtype = dtype or xp.float64 + if len(args) == 2: + array_like, desired = args + else: + array_like, p, desired = args + array_like = xp.asarray(array_like, dtype=dtype) + dtype_reference = (xp_result_type(array_like, force_floating=True, xp=xp) + if xp != np.ma else None) + desired = xp.asarray(desired, dtype=dtype_reference) + weights = xp.asarray(weights, dtype=dtype) if weights is not None else weights + args = (array_like,) if len(args) == 2 else (array_like, p) + x = mean_fun(*args, axis=axis, weights=weights) + xp_assert_close(x, desired, rtol=rtol) + + +def check_equal_gmean(*args, **kwargs): + return check_equal_xmean(*args, mean_fun=stats.gmean, **kwargs) + + +def check_equal_hmean(*args, **kwargs): + return check_equal_xmean(*args, mean_fun=stats.hmean, **kwargs) + + +def check_equal_pmean(*args, **kwargs): + return check_equal_xmean(*args, mean_fun=stats.pmean, **kwargs) + + +@make_xp_test_case(stats.hmean) +class TestHMean: + @pytest.mark.filterwarnings("ignore:divide by zero encountered:RuntimeWarning:dask") + def test_0(self, xp): + a = [1, 0, 2] + desired = 0 + check_equal_hmean(a, desired, xp=xp) + + @pytest.mark.parametrize('dtype', [None, 'int64', 'float32', 'float64']) + def test_1d(self, dtype, xp): + dtype = dtype if dtype is None else getattr(xp, dtype) + + # Test a 1d case + a = [10, 20, 30, 40, 50, 60, 70, 80, 90, 100] + desired = 34.1417152147 + check_equal_hmean(a, desired, dtype=dtype, xp=xp) + + a = [1, 2, 3, 4] + desired = 4. / (1. / 1 + 1. / 2 + 1. / 3 + 1. / 4) + check_equal_hmean(a, desired, dtype=dtype, xp=xp) + + @pytest.mark.filterwarnings("ignore:divide by zero encountered:RuntimeWarning:dask") + @pytest.mark.filterwarnings("ignore:invalid value encountered:RuntimeWarning:dask") + def test_1d_with_zero(self, xp): + a = np.array([1, 0]) + desired = 0.0 + check_equal_hmean(a, desired, xp=xp, rtol=0.0) + + @pytest.mark.filterwarnings("ignore:invalid value encountered:RuntimeWarning:dask") + @pytest.mark.filterwarnings("ignore:divide by zero encountered:RuntimeWarning:dask") + def test_1d_with_negative_value(self, xp): + a = np.array([1, 0, -1]) + message = "The harmonic mean is only defined..." + with eager_warns(RuntimeWarning, match=message, xp=xp): + check_equal_hmean(a, xp.nan, xp=xp, rtol=0.0) + + # Note the next tests use axis=None as default, not axis=0 + def test_2d(self, xp): + # Test a 2d case + a = [[10, 20, 30, 40], [50, 60, 70, 80], [90, 100, 110, 120]] + desired = 38.6696271841 + check_equal_hmean(np.array(a), desired, xp=xp) + + def test_2d_axis0(self, xp): + # Test a 2d case with axis=0 + a = [[10, 20, 30, 40], [50, 60, 70, 80], [90, 100, 110, 120]] + desired = np.array([22.88135593, 39.13043478, 52.90076336, 65.45454545]) + check_equal_hmean(a, desired, axis=0, xp=xp) + + @pytest.mark.filterwarnings("ignore:divide by zero encountered:RuntimeWarning:dask") + def test_2d_axis0_with_zero(self, xp): + a = [[10, 0, 30, 40], [50, 60, 70, 80], [90, 100, 110, 120]] + desired = np.array([22.88135593, 0.0, 52.90076336, 65.45454545]) + check_equal_hmean(a, desired, axis=0, xp=xp) + + def test_2d_axis1(self, xp): + # Test a 2d case with axis=1 + a = [[10, 20, 30, 40], [50, 60, 70, 80], [90, 100, 110, 120]] + desired = np.array([19.2, 63.03939962, 103.80078637]) + check_equal_hmean(a, desired, axis=1, xp=xp) + + @pytest.mark.filterwarnings("ignore:divide by zero encountered:RuntimeWarning:dask") + def test_2d_axis1_with_zero(self, xp): + a = [[10, 0, 30, 40], [50, 60, 70, 80], [90, 100, 110, 120]] + desired = np.array([0.0, 63.03939962, 103.80078637]) + check_equal_hmean(a, desired, axis=1, xp=xp) + + @skip_xp_backends( + np_only=True, + reason='array-likes only supported for NumPy backend', + ) + def test_weights_1d_list(self, xp): + # Desired result from: + # https://www.hackmath.net/en/math-problem/35871 + a = [2, 10, 6] + weights = [10, 5, 3] + desired = 3. + # all the other tests use `check_equal_hmean`, which now converts + # the input to an xp-array before calling `hmean`. This time, check + # that the function still accepts the lists of ints. + res = stats.hmean(a, weights=weights) + xp_assert_close(res, np.asarray(desired), rtol=1e-5) + + def test_weights_1d(self, xp): + # Desired result from: + # https://www.hackmath.net/en/math-problem/35871 + a = np.asarray([2, 10, 6]) + weights = np.asarray([10, 5, 3]) + desired = 3 + check_equal_hmean(a, desired, weights=weights, rtol=1e-5, xp=xp) + + def test_weights_2d_axis0(self, xp): + # Desired result from: + # https://www.hackmath.net/en/math-problem/35871 + a = np.array([[2, 5], [10, 5], [6, 5]]) + weights = np.array([[10, 1], [5, 1], [3, 1]]) + desired = np.array([3, 5]) + check_equal_hmean(a, desired, axis=0, weights=weights, rtol=1e-5, xp=xp) + + def test_weights_2d_axis1(self, xp): + # Desired result from: + # https://www.hackmath.net/en/math-problem/35871 + a = np.array([[2, 10, 6], [7, 7, 7]]) + weights = np.array([[10, 5, 3], [1, 1, 1]]) + desired = np.array([3, 7]) + check_equal_hmean(a, desired, axis=1, weights=weights, rtol=1e-5, xp=xp) + + @skip_xp_invalid_arg + def test_weights_masked_1d_array(self, xp): + # Desired result from: + # https://www.hackmath.net/en/math-problem/35871 + a = np.array([2., 10., 6., 42.]) + weights = np.ma.array([10., 5., 3., 42.], mask=[0, 0, 0, 1]) + desired = 3. + xp = np.ma # check_equal_hmean uses xp.asarray; this will preserve the mask + check_equal_hmean(a, desired, weights=weights, rtol=1e-5, + dtype=np.float64, xp=xp) + + +@make_xp_test_case(stats.gmean) +class TestGMean: + @pytest.mark.filterwarnings( + "ignore:divide by zero encountered in log:RuntimeWarning:dask" + ) + def test_0(self, xp): + a = [1, 0, 2] + desired = 0 + check_equal_gmean(a, desired, xp=xp) + + @pytest.mark.parametrize('dtype', [None, 'int64', 'float32', 'float64']) + def test_1d(self, dtype, xp): + dtype = dtype if dtype is None else getattr(xp, dtype) + + # Test a 1d case + a = [10, 20, 30, 40, 50, 60, 70, 80, 90, 100] + desired = 45.2872868812 + check_equal_gmean(a, desired, dtype=dtype, xp=xp) + + a = [1, 2, 3, 4] + desired = power(1 * 2 * 3 * 4, 1. / 4.) + check_equal_gmean(a, desired, dtype=dtype, rtol=1e-14, xp=xp) + + # Note the next tests use axis=None as default, not axis=0 + def test_2d(self, xp): + # Test a 2d case + a = [[10, 20, 30, 40], [50, 60, 70, 80], [90, 100, 110, 120]] + desired = 52.8885199 + check_equal_gmean(a, desired, xp=xp) + + def test_2d_axis0(self, xp): + # Test a 2d case with axis=0 + a = [[10, 20, 30, 40], [50, 60, 70, 80], [90, 100, 110, 120]] + desired = np.array([35.56893304, 49.32424149, 61.3579244, 72.68482371]) + check_equal_gmean(a, desired, axis=0, xp=xp) + + a = array([[1, 2, 3, 4], [1, 2, 3, 4], [1, 2, 3, 4]]) + desired = array([1, 2, 3, 4]) + check_equal_gmean(a, desired, axis=0, rtol=1e-14, xp=xp) + + def test_2d_axis1(self, xp): + # Test a 2d case with axis=1 + a = [[10, 20, 30, 40], [50, 60, 70, 80], [90, 100, 110, 120]] + desired = np.array([22.13363839, 64.02171746, 104.40086817]) + check_equal_gmean(a, desired, axis=1, xp=xp) + + a = array([[1, 2, 3, 4], [1, 2, 3, 4], [1, 2, 3, 4]]) + v = power(1 * 2 * 3 * 4, 1. / 4.) + desired = array([v, v, v]) + check_equal_gmean(a, desired, axis=1, rtol=1e-14, xp=xp) + + def test_large_values(self, xp): + a = array([1e100, 1e200, 1e300]) + desired = 1e200 + check_equal_gmean(a, desired, rtol=1e-13, xp=xp) + + @pytest.mark.filterwarnings( + "ignore:divide by zero encountered in log:RuntimeWarning:dask" + ) + def test_1d_with_0(self, xp): + # Test a 1d case with zero element + a = [10, 20, 30, 40, 50, 60, 70, 80, 90, 0] + desired = 0.0 # due to exp(-inf)=0 + with np.errstate(all='ignore'): + check_equal_gmean(a, desired, xp=xp) + + @pytest.mark.filterwarnings( + "ignore:invalid value encountered in log:RuntimeWarning:dask" + ) + def test_1d_neg(self, xp): + # Test a 1d case with negative element + a = [10, 20, 30, 40, 50, 60, 70, 80, 90, -1] + desired = np.nan # due to log(-1) = nan + with np.errstate(invalid='ignore'): + check_equal_gmean(a, desired, xp=xp) + + @skip_xp_backends( + np_only=True, + reason='array-likes only supported for NumPy backend', + ) + def test_weights_1d_list(self, xp): + # Desired result from: + # https://www.dummies.com/education/math/business-statistics/how-to-find-the-weighted-geometric-mean-of-a-data-set/ + a = [1, 2, 3, 4, 5] + weights = [2, 5, 6, 4, 3] + desired = 2.77748 + + # all the other tests use `check_equal_gmean`, which now converts + # the input to an xp-array before calling `gmean`. This time, check + # that the function still accepts the lists of ints. + res = stats.gmean(a, weights=weights) + xp_assert_close(res, np.asarray(desired), rtol=1e-5) + + def test_weights_1d(self, xp): + # Desired result from: + # https://www.dummies.com/education/math/business-statistics/how-to-find-the-weighted-geometric-mean-of-a-data-set/ + a = np.array([1, 2, 3, 4, 5]) + weights = np.array([2, 5, 6, 4, 3]) + desired = 2.77748 + check_equal_gmean(a, desired, weights=weights, rtol=1e-5, xp=xp) + + @skip_xp_invalid_arg + def test_weights_masked_1d_array(self, xp): + # Desired result from: + # https://www.dummies.com/education/math/business-statistics/how-to-find-the-weighted-geometric-mean-of-a-data-set/ + a = np.array([1., 2., 3., 4., 5., 6.]) + weights = np.ma.array([2., 5., 6., 4., 3., 5.], mask=[0, 0, 0, 0, 0, 1]) + desired = 2.77748 + xp = np.ma # check_equal_gmean uses xp.asarray; this will preserve the mask + check_equal_gmean(a, desired, weights=weights, rtol=1e-5, + dtype=np.float64, xp=xp) + + +@make_xp_test_case(stats.pmean) +class TestPMean: + + def pmean_reference(a, p): + return (np.sum(a**p) / a.size)**(1/p) + + def wpmean_reference(a, p, weights): + return (np.sum(weights * a**p) / np.sum(weights))**(1/p) + + def test_bad_exponent(self, xp): + with pytest.raises(ValueError, match='Power mean only defined for'): + stats.pmean(xp.asarray([1, 2, 3]), xp.asarray([0])) + + @pytest.mark.parametrize('dtype', [None, 'int64', 'float32', 'float64']) + def test_1d(self, dtype, xp): + dtype = dtype if dtype is None else getattr(xp, dtype) + + a, p = [10, 20, 30, 40, 50, 60, 70, 80, 90, 100], 3.5 + desired = TestPMean.pmean_reference(np.array(a), p) + check_equal_pmean(a, p, desired, dtype=dtype, xp=xp) + + a, p = [10, 20, 30, 40, 50, 60, 70, 80, 90, 100], -2.5 + desired = TestPMean.pmean_reference(np.array(a), p) + check_equal_pmean(a, p, desired, dtype=dtype, xp=xp) + + a, p = [1, 2, 3, 4], 2 + desired = np.sqrt((1**2 + 2**2 + 3**2 + 4**2) / 4) + check_equal_pmean(a, p, desired, dtype=dtype, xp=xp) + + @pytest.mark.filterwarnings("ignore:invalid value encountered:RuntimeWarning:dask") + @pytest.mark.filterwarnings("ignore:divide by zero encountered:RuntimeWarning:dask") + def test_1d_with_zero(self, xp): + a, p = np.array([1, 0]), -1 + desired = 0.0 + check_equal_pmean(a, p, desired, rtol=0.0, xp=xp) + + def test_1d_with_negative_value(self, xp): + a, p = np.array([1, 0, -1]), 1.23 + message = "The power mean is only defined..." + with eager_warns(RuntimeWarning, match=message, xp=xp): + check_equal_pmean(a, p, xp.nan, xp=xp) + + @pytest.mark.parametrize( + ("a", "p"), + [([[10, 20], [50, 60], [90, 100]], -0.5), + (np.array([[10, 20], [50, 60], [90, 100]]), 0.5)] + ) + def test_2d_axisnone(self, a, p, xp): + desired = TestPMean.pmean_reference(np.array(a), p) + check_equal_pmean(a, p, desired, xp=xp) + + @pytest.mark.parametrize( + ("a", "p"), + [([[10, 20, 30, 40], [50, 60, 70, 80], [90, 100, 110, 120]], -0.5), + ([[10, 0, 30, 40], [50, 60, 70, 80], [90, 100, 110, 120]], 0.5)] + ) + def test_2d_axis0(self, a, p, xp): + desired = [ + TestPMean.pmean_reference( + np.array([a[i][j] for i in range(len(a))]), p + ) + for j in range(len(a[0])) + ] + check_equal_pmean(a, p, desired, axis=0, xp=xp) + + @pytest.mark.parametrize( + ("a", "p"), + [([[10, 20, 30, 40], [50, 60, 70, 80], [90, 100, 110, 120]], -0.5), + ([[10, 0, 30, 40], [50, 60, 70, 80], [90, 100, 110, 120]], 0.5)] + ) + def test_2d_axis1(self, a, p, xp): + desired = [TestPMean.pmean_reference(np.array(a_), p) for a_ in a] + check_equal_pmean(a, p, desired, axis=1, xp=xp) + + def test_weights_1d(self, xp): + a, p = [2, 10, 6], -1.23456789 + weights = [10, 5, 3] + desired = TestPMean.wpmean_reference(np.array(a), p, weights) + check_equal_pmean(a, p, desired, weights=weights, rtol=1e-5, xp=xp) + + @skip_xp_backends( + np_only=True, + reason='array-likes only supported for NumPy backend', + ) + def test_weights_1d_list(self, xp): + a, p = [2, 10, 6], -1.23456789 + weights = [10, 5, 3] + desired = TestPMean.wpmean_reference(np.array(a), p, weights) + # all the other tests use `check_equal_pmean`, which now converts + # the input to an xp-array before calling `pmean`. This time, check + # that the function still accepts the lists of ints. + res = stats.pmean(a, p, weights=weights) + xp_assert_close(res, np.asarray(desired), rtol=1e-5) + + @skip_xp_invalid_arg + def test_weights_masked_1d_array(self, xp): + a, p = np.array([2., 10., 6., 42.]), 1 + weights = np.ma.array([10., 5., 3., 42.], mask=[0, 0, 0, 1]) + desired = np.average(a, weights=weights) + xp = np.ma # check_equal_pmean uses xp.asarray; this will preserve the mask + check_equal_pmean(a, p, desired, weights=weights, rtol=1e-5, + dtype=np.float64, xp=xp) + + @pytest.mark.parametrize( + ("axis", "fun_name", "p"), + [(None, "wpmean_reference", 9.87654321), + (0, "gmean", 0), + (1, "hmean", -1)] + ) + def test_weights_2d(self, axis, fun_name, p, xp): + if fun_name == 'wpmean_reference': + def fun(a, axis, weights): + return TestPMean.wpmean_reference(a, p, weights) + else: + fun = getattr(stats, fun_name) + a = np.array([[2, 5], [10, 5], [6, 5]]) + weights = np.array([[10, 1], [5, 1], [3, 1]]) + desired = fun(a, axis=axis, weights=weights) + check_equal_pmean(a, p, desired, axis=axis, weights=weights, rtol=1e-5, xp=xp) + + def test_infinite_p_gh23111(self): + # gh-23111 reported that `pmean` didn't work properly with infinite `p`; + # check that this raises an appropriate error message + message = "Power mean only implemented for finite `p`" + with pytest.raises(NotImplementedError, match=message): + stats.pmean([2], np.inf) + + + @pytest.mark.parametrize( + "p, weights, ref", + [ + # import numpy as np + # from mpmath import mp + # mp.dps = 1000 + # + # x = [50., 100., 300., 70., 25., 100., 2.] + # w = [1., 2., 1., 1., 1., 2., 1.] + # + # x = np.asarray([mp.mpf(x_) for x_ in x]) + # w = np.asarray([mp.mpf(w_) for w_ in w]) + # p = mp.mpf(2e-16) + # + # res = np.average(x**p)**(1/p) + # float(res) # 47.23984586445856 + # + (2e-16, None, 47.23984586445856), + # + # res = np.average(x**p, weights=w)**(1/p) + # float(res) # 55.80644618389137 + # + (2e-16, [1., 2., 1., 1., 1., 2., 1.], 55.80644618389137), + ], + ) + def test_small_p_gh23407(self, xp, p, weights, ref): + # gh-23407 reported that pmean had + # poor numerical behavior for very small nonzero p + x = xp.asarray( + [50., 100., 300., 70., 25., 100., 2.] + ) + w = None if weights is None else xp.asarray(weights) + + res = stats.pmean(x, p, weights=w) + xp_assert_close(res, xp.asarray(ref)) + + +@make_xp_test_case(stats.gstd) +class TestGSTD: + # must add 1 as `gstd` is only defined for positive values + array_1d = (np.arange(2 * 3 * 4) + 1).tolist() + gstd_array_1d = 2.294407613602 + array_3d = np.reshape(array_1d, (2, 3, 4)).tolist() + + def test_1d_array(self, xp): + gstd_actual = stats.gstd(xp.asarray(self.array_1d)) + xp_assert_close(gstd_actual, xp.asarray(self.gstd_array_1d)) + + @skip_xp_backends(np_only=True, reason="Only NumPy supports array-like input") + def test_1d_numeric_array_like_input(self, xp): + gstd_actual = stats.gstd(tuple(self.array_1d)) + assert_allclose(gstd_actual, self.gstd_array_1d) + + @skip_xp_invalid_arg + def test_raises_error_non_numeric_input(self, xp): + message = "could not convert string to float|The DType..." + with pytest.raises((ValueError, TypeError), match=message): + stats.gstd('You cannot take the logarithm of a string.') + + @pytest.mark.filterwarnings("ignore:divide by zero encountered:RuntimeWarning:dask") + @pytest.mark.filterwarnings("ignore:invalid value encountered:RuntimeWarning:dask") + @pytest.mark.parametrize('bad_value', (0, -1, np.inf, np.nan)) + def test_returns_nan_invalid_value(self, bad_value, xp): + x = xp.asarray(self.array_1d + [bad_value]) + if np.isfinite(bad_value) and not is_lazy_array(x): + message = "The geometric standard deviation is only defined..." + with pytest.warns(RuntimeWarning, match=message): + res = stats.gstd(x) + else: + res = stats.gstd(x) + xp_assert_equal(res, xp.asarray(np.nan)) + + def test_propagates_nan_values(self, xp): + a = xp.asarray([[1, 1, 1, 16], [xp.nan, 1, 2, 3]]) + gstd_actual = stats.gstd(a, axis=1) + xp_assert_close(gstd_actual, xp.asarray([4, np.nan])) + + def test_ddof_equal_to_number_of_observations(self, xp): + x = xp.asarray(self.array_1d) + res = stats.gstd(x, ddof=x.shape[0]) + xp_assert_equal(res, xp.asarray(xp.nan)) + + def test_3d_array(self, xp): + x = xp.asarray(self.array_3d) + gstd_actual = stats.gstd(x, axis=None) + ref = xp.asarray(self.gstd_array_1d) + xp_assert_close(gstd_actual, ref) + + def test_3d_array_axis_type_tuple(self, xp): + x = xp.asarray(self.array_3d) + gstd_actual = stats.gstd(x, axis=(1, 2)) + ref = xp.asarray([2.12939215, 1.22120169]) + xp_assert_close(gstd_actual, ref) + + def test_3d_array_axis_0(self, xp): + x = xp.asarray(self.array_3d) + gstd_actual = stats.gstd(x, axis=0) + gstd_desired = xp.asarray([ + [6.1330555493918, 3.958900210120, 3.1206598248344, 2.6651441426902], + [2.3758135028411, 2.174581428192, 2.0260062829505, 1.9115518327308], + [1.8205343606803, 1.746342404566, 1.6846557065742, 1.6325269194382] + ]) + xp_assert_close(gstd_actual, gstd_desired) + + def test_3d_array_axis_1(self, xp): + x = xp.asarray(self.array_3d) + gstd_actual = stats.gstd(x, axis=1) + gstd_desired = xp.asarray([ + [3.118993630946, 2.275985934063, 1.933995977619, 1.742896469724], + [1.271693593916, 1.254158641801, 1.238774141609, 1.225164057869] + ]) + xp_assert_close(gstd_actual, gstd_desired) + + def test_3d_array_axis_2(self, xp): + x = xp.asarray(self.array_3d) + gstd_actual = stats.gstd(x, axis=2) + gstd_desired = xp.asarray([ + [1.8242475707664, 1.2243686572447, 1.1318311657788], + [1.0934830582351, 1.0724479791887, 1.0591498540749] + ]) + xp_assert_close(gstd_actual, gstd_desired) + + +class TestTrim: + # test trim functions + def test_trim1(self): + a = np.arange(11) + assert_equal(np.sort(stats.trim1(a, 0.1)), np.arange(10)) + assert_equal(np.sort(stats.trim1(a, 0.2)), np.arange(9)) + assert_equal(np.sort(stats.trim1(a, 0.2, tail='left')), + np.arange(2, 11)) + assert_equal(np.sort(stats.trim1(a, 3/11., tail='left')), + np.arange(3, 11)) + assert_equal(stats.trim1(a, 1.0), []) + assert_equal(stats.trim1(a, 1.0, tail='left'), []) + + # empty input + assert_equal(stats.trim1([], 0.1), []) + assert_equal(stats.trim1([], 3/11., tail='left'), []) + assert_equal(stats.trim1([], 4/6.), []) + + # test axis + a = np.arange(24).reshape(6, 4) + ref = np.arange(4, 24).reshape(5, 4) # first row trimmed + + axis = 0 + trimmed = stats.trim1(a, 0.2, tail='left', axis=axis) + assert_equal(np.sort(trimmed, axis=axis), ref) + + axis = 1 + trimmed = stats.trim1(a.T, 0.2, tail='left', axis=axis) + assert_equal(np.sort(trimmed, axis=axis), ref.T) + + def test_trimboth(self): + a = np.arange(11) + assert_equal(np.sort(stats.trimboth(a, 3/11.)), np.arange(3, 8)) + assert_equal(np.sort(stats.trimboth(a, 0.2)), + np.array([2, 3, 4, 5, 6, 7, 8])) + assert_equal(np.sort(stats.trimboth(np.arange(24).reshape(6, 4), 0.2)), + np.arange(4, 20).reshape(4, 4)) + assert_equal(np.sort(stats.trimboth(np.arange(24).reshape(4, 6).T, + 2/6.)), + np.array([[2, 8, 14, 20], [3, 9, 15, 21]])) + assert_raises(ValueError, stats.trimboth, + np.arange(24).reshape(4, 6).T, 4/6.) + + # empty input + assert_equal(stats.trimboth([], 0.1), []) + assert_equal(stats.trimboth([], 3/11.), []) + assert_equal(stats.trimboth([], 4/6.), []) + + +@make_xp_test_case(stats.trim_mean) +class TestTrimMean: + def test_trim_mean(self, xp): + # don't use pre-sorted arrays + idx = np.array([3, 5, 0, 1, 2, 4]) + a2 = np.arange(24).reshape(6, 4)[idx, :] + a3 = np.arange(24).reshape(6, 4, order='F')[idx, :] + + xp_assert_equal(stats.trim_mean(xp.asarray(a3), 2/6.), + xp.asarray([2.5, 8.5, 14.5, 20.5])) + xp_assert_equal(stats.trim_mean(xp.asarray(a2), 2/6.), + xp.asarray([10., 11., 12., 13.])) + + idx4 = np.array([1, 0, 3, 2]) + a4 = np.arange(24).reshape(4, 6)[idx4, :] + xp_assert_equal(stats.trim_mean(xp.asarray(a4), 2/6.), + xp.asarray([9., 10., 11., 12., 13., 14.])) + + # shuffled arange(24) as array_like + a = [7, 11, 12, 21, 16, 6, 22, 1, 5, 0, 18, 10, 17, 9, 19, 15, 23, + 20, 2, 14, 4, 13, 8, 3] + xp_assert_equal(stats.trim_mean(xp.asarray(a), 2/6.), xp.asarray(11.5)) + xp_assert_equal(stats.trim_mean(xp.asarray([5, 4, 3, 1, 2, 0]), 2/6.), + xp.asarray(2.5)) + + # check axis argument + rng = np.random.default_rng(3417115752) + a = rng.integers(20, size=(5, 6, 4, 7)) + a = xp.asarray(a) + for axis in [0, 1, 2, 3, -1]: + res1 = stats.trim_mean(a, 2/6., axis=axis) + res2 = stats.trim_mean(xp.moveaxis(a, axis, 0), 2/6.) + xp_assert_equal(res1, res2) + + res1 = stats.trim_mean(a, 2/6., axis=None) + res2 = stats.trim_mean(xp_ravel(a), 2/6.) + xp_assert_equal(res1, res2) + + with pytest.raises(ValueError, match="Proportion too big."): + stats.trim_mean(a, 0.6) + + @skip_xp_backends('dask.array', reason="lazy -> no _axis_nan_policy") + def test_empty_input(self, xp): + # empty input + with pytest.warns(SmallSampleWarning, match='too small'): + xp_assert_equal(stats.trim_mean(xp.asarray([]), 0.0), xp.asarray(xp.nan)) + with pytest.warns(SmallSampleWarning, match='too small'): + xp_assert_equal(stats.trim_mean(xp.asarray([]), 0.6), xp.asarray(xp.nan)) + + +@make_xp_test_case(stats.sigmaclip) +class TestSigmaClip: + def test_sigmaclip1(self, xp): + a = xp.concat((xp.linspace(9.5, 10.5, 31), + xp.linspace(0., 20., 5))) + fact = 4 # default + c, low, upp = stats.sigmaclip(a) + assert xp.min(c) > low + assert xp.max(c) < upp + xp_assert_equal(low, xp.mean(c) - fact*xp.std(c)) + xp_assert_equal(upp, xp.mean(c) + fact*xp.std(c)) + assert c.shape == a.shape + + def test_sigmaclip2(self, xp): + a = xp.concat((xp.linspace(9.5, 10.5, 31), + xp.linspace(0., 20., 5))) + fact = 1.5 + c, low, upp = stats.sigmaclip(a, fact, fact) + assert xp.min(c) > low + assert xp.max(c) < upp + xp_assert_equal(low, xp.mean(c) - fact*xp.std(c)) + xp_assert_equal(upp, xp.mean(c) + fact*xp.std(c)) + assert c.shape == (4,) + assert a.shape == (36,) + + def test_sigmaclip3(self, xp): + a = xp.concat((xp.linspace(9.5, 10.5, 11), + xp.linspace(-100., -50., 3))) + fact = 1.8 + c, low, upp = stats.sigmaclip(a, fact, fact) + assert xp.min(c) > low + assert xp.max(c) < upp + xp_assert_equal(low, xp.mean(c) - fact*xp.std(c)) + xp_assert_equal(upp, xp.mean(c) + fact*xp.std(c)) + xp_assert_equal(c, xp.linspace(9.5, 10.5, 11)) + + def test_sigmaclip_result_attributes(self, xp): + a = xp.concat((xp.linspace(9.5, 10.5, 11), + xp.linspace(-100., -50., 3))) + fact = 1.8 + res = stats.sigmaclip(a, fact, fact) + attributes = ('clipped', 'lower', 'upper') + check_named_results(res, attributes, xp=xp) + + def test_std_zero(self, xp): + # regression test #8632 + x = xp.ones(10) + xp_assert_equal(stats.sigmaclip(x)[0], x) + + @pytest.mark.parametrize("dtype", [None, 'float32', 'float64']) + def test_nan_policy(self, dtype, xp): + dtype = None if dtype is None else getattr(xp, dtype) + rng = np.random.default_rng(243587919486245906119) + x = rng.random(10) + x[3] = np.nan + x = xp.asarray(x.tolist(), dtype=dtype) + + # nan_policy='raise' + if not is_lazy_array(x): + with pytest.raises(ValueError, match="The input contains nan values"): + stats.sigmaclip(x, nan_policy='raise') + else: + with pytest.raises(TypeError, match="nan_policy='raise' is not supported"): + stats.sigmaclip(x, nan_policy='raise') + + # nan_policy='propagate' + res = stats.sigmaclip(x, nan_policy='propagate') + xp_assert_equal(res.clipped, xp.asarray([], dtype=dtype)) + xp_assert_equal(res.lower, xp.asarray(xp.nan, dtype=dtype)) + xp_assert_equal(res.upper, xp.asarray(xp.nan, dtype=dtype)) + + # nan_policy='omit' + res = stats.sigmaclip(x, nan_policy='omit') + ref = stats.sigmaclip(x[~xp.isnan(x)]) + xp_assert_equal(res.clipped, ref.clipped) + xp_assert_equal(res.lower, ref.lower) + xp_assert_equal(res.upper, ref.upper) + + +@make_xp_test_case(stats.alexandergovern) +class TestAlexanderGovern: + def test_compare_dtypes(self): + # leave this NumPy only + args = [[13, 13, 13, 13, 13, 13, 13, 12, 12], + [14, 13, 12, 12, 12, 12, 12, 11, 11], + [14, 14, 13, 13, 13, 13, 13, 12, 12], + [15, 14, 13, 13, 13, 12, 12, 12, 11]] + args_int16 = [np.asarray(arg, dtype=np.int16) for arg in args] + args_int32 = [np.asarray(arg, dtype=np.int32) for arg in args] + args_uint8 = [np.asarray(arg, dtype=np.uint8) for arg in args] + args_float64 = [np.asarray(arg, dtype=np.float64) for arg in args] + + res_int16 = stats.alexandergovern(*args_int16) + res_int32 = stats.alexandergovern(*args_int32) + res_uint8 = stats.alexandergovern(*args_uint8) + res_float64 = stats.alexandergovern(*args_float64) + + assert (res_int16.pvalue == res_int32.pvalue == + res_uint8.pvalue == res_float64.pvalue) + assert (res_int16.statistic == res_int32.statistic == + res_uint8.statistic == res_float64.statistic) + + @skip_xp_backends('dask.array', reason="Requires `_axis_nan_policy` decorator") + @pytest.mark.parametrize('case',[([1, 2], []), ([1, 2], 2), ([1, 2], [2])]) + def test_too_small_inputs(self, case, xp): + # input array is of size zero or too small + dtype = xp_default_dtype(xp) + case = xp.asarray(case[0], dtype=dtype), xp.asarray(case[1], dtype=dtype) + with pytest.warns(SmallSampleWarning, match=too_small_1d_not_omit): + res = stats.alexandergovern(*case) + xp_assert_equal(res.statistic, xp.asarray(xp.nan)) + xp_assert_equal(res.pvalue, xp.asarray(xp.nan)) + + @pytest.mark.filterwarnings("ignore:invalid value encountered:RuntimeWarning") + def test_bad_inputs(self, xp): + # inputs are not finite (infinity) + samples = xp.asarray([1., 2.]), xp.asarray([xp.inf, xp.inf]) + res = stats.alexandergovern(*samples) + xp_assert_equal(res.statistic, xp.asarray(xp.nan)) + xp_assert_equal(res.pvalue, xp.asarray(xp.nan)) + + def test_compare_r(self, xp): + ''' + Data generated in R with + > set.seed(1) + > library("onewaytests") + > library("tibble") + > y <- c(rnorm(40, sd=10), + + rnorm(30, sd=15), + + rnorm(20, sd=20)) + > x <- c(rep("one", times=40), + + rep("two", times=30), + + rep("eight", times=20)) + > x <- factor(x) + > ag.test(y ~ x, tibble(y,x)) + + Alexander-Govern Test (alpha = 0.05) + ------------------------------------------------------------- + data : y and x + + statistic : 1.359941 + parameter : 2 + p.value : 0.5066321 + + Result : Difference is not statistically significant. + ------------------------------------------------------------- + Example adapted from: + https://eval-serv2.metpsy.uni-jena.de/wiki-metheval-hp/index.php/R_FUN_Alexander-Govern + + ''' + one = [-6.264538107423324, 1.8364332422208225, -8.356286124100471, + 15.952808021377916, 3.295077718153605, -8.204683841180152, + 4.874290524284853, 7.383247051292173, 5.757813516534923, + -3.0538838715635603, 15.11781168450848, 3.898432364114311, + -6.2124058054180376, -22.146998871774997, 11.249309181431082, + -0.4493360901523085, -0.16190263098946087, 9.438362106852992, + 8.212211950980885, 5.939013212175088, 9.189773716082183, + 7.821363007310671, 0.745649833651906, -19.89351695863373, + 6.198257478947102, -0.5612873952900078, -1.557955067053293, + -14.707523838992744, -4.781500551086204, 4.179415601997024, + 13.58679551529044, -1.0278772734299553, 3.876716115593691, + -0.5380504058290512, -13.770595568286065, -4.149945632996798, + -3.942899537103493, -0.5931339671118566, 11.000253719838831, + 7.631757484575442] + + two = [-2.4678539438038034, -3.8004252020476135, 10.454450631071062, + 8.34994798010486, -10.331335418242798, -10.612427354431794, + 5.468729432052455, 11.527993867731237, -1.6851931822534207, + 13.216615896813222, 5.971588205506021, -9.180395898761569, + 5.116795371366372, -16.94044644121189, 21.495355525515556, + 29.7059984775879, -5.508322146997636, -15.662019394747961, + 8.545794411636193, -2.0258190582123654, 36.024266407571645, + -0.5886000409975387, 10.346090436761651, 0.4200323817099909, + -11.14909813323608, 2.8318844927151434, -27.074379433365568, + 21.98332292344329, 2.2988000731784655, 32.58917505543229] + + eight = [9.510190577993251, -14.198928618436291, 12.214527069781099, + -18.68195263288503, -25.07266800478204, 5.828924710349257, + -8.86583746436866, 0.02210703263248262, 1.4868264830332811, + -11.79041892376144, -11.37337465637004, -2.7035723024766414, + 23.56173993146409, -30.47133600859524, 11.878923752568431, + 6.659007424270365, 21.261996745527256, -6.083678472686013, + 7.400376198325763, 5.341975815444621] + + one, two, eight = xp.asarray(one), xp.asarray(two), xp.asarray(eight) + soln = stats.alexandergovern(one, two, eight) + xp_assert_close(soln.statistic, xp.asarray(1.3599405447999450836)) + xp_assert_close(soln.pvalue, xp.asarray(0.50663205309676440091)) + + def test_compare_scholar(self, xp): + ''' + Data taken from 'The Modification and Evaluation of the + Alexander-Govern Test in Terms of Power' by Kingsley Ochuko, T., + Abdullah, S., Binti Zain, Z., & Soaad Syed Yahaya, S. (2015). + ''' + young = [482.43, 484.36, 488.84, 495.15, 495.24, 502.69, 504.62, + 518.29, 519.1, 524.1, 524.12, 531.18, 548.42, 572.1, 584.68, + 609.09, 609.53, 666.63, 676.4] + middle = [335.59, 338.43, 353.54, 404.27, 437.5, 469.01, 485.85, + 487.3, 493.08, 494.31, 499.1, 886.41] + old = [519.01, 528.5, 530.23, 536.03, 538.56, 538.83, 557.24, 558.61, + 558.95, 565.43, 586.39, 594.69, 629.22, 645.69, 691.84] + young, middle, old = xp.asarray(young), xp.asarray(middle), xp.asarray(old) + soln = stats.alexandergovern(young, middle, old) + xp_assert_close(soln.statistic, xp.asarray(5.3237), atol=1e-3) + xp_assert_close(soln.pvalue, xp.asarray(0.06982), atol=1e-4) + + # verify with ag.test in r + ''' + > library("onewaytests") + > library("tibble") + > young <- c(482.43, 484.36, 488.84, 495.15, 495.24, 502.69, 504.62, + + 518.29, 519.1, 524.1, 524.12, 531.18, 548.42, 572.1, + + 584.68, 609.09, 609.53, 666.63, 676.4) + > middle <- c(335.59, 338.43, 353.54, 404.27, 437.5, 469.01, 485.85, + + 487.3, 493.08, 494.31, 499.1, 886.41) + > old <- c(519.01, 528.5, 530.23, 536.03, 538.56, 538.83, 557.24, + + 558.61, 558.95, 565.43, 586.39, 594.69, 629.22, + + 645.69, 691.84) + > young_fct <- c(rep("young", times=19)) + > middle_fct <-c(rep("middle", times=12)) + > old_fct <- c(rep("old", times=15)) + > ag.test(a ~ b, tibble(a=c(young, middle, old), b=factor(c(young_fct, + + middle_fct, old_fct)))) + + Alexander-Govern Test (alpha = 0.05) + ------------------------------------------------------------- + data : a and b + + statistic : 5.324629 + parameter : 2 + p.value : 0.06978651 + + Result : Difference is not statistically significant. + ------------------------------------------------------------- + + ''' + xp_assert_close(soln.statistic, xp.asarray(5.324629), rtol=2e-6) + xp_assert_close(soln.pvalue, xp.asarray(0.06978651), rtol=2e-6) + + def test_compare_scholar3(self, xp): + ''' + Data taken from 'Robustness And Comparative Power Of WelchAspin, + Alexander-Govern And Yuen Tests Under Non-Normality And Variance + Heteroscedasticity', by Ayed A. Almoied. 2017. Page 34-37. + https://digitalcommons.wayne.edu/cgi/viewcontent.cgi?article=2775&context=oa_dissertations + ''' + x1 = [-1.77559, -1.4113, -0.69457, -0.54148, -0.18808, -0.07152, + 0.04696, 0.051183, 0.148695, 0.168052, 0.422561, 0.458555, + 0.616123, 0.709968, 0.839956, 0.857226, 0.929159, 0.981442, + 0.999554, 1.642958] + x2 = [-1.47973, -1.2722, -0.91914, -0.80916, -0.75977, -0.72253, + -0.3601, -0.33273, -0.28859, -0.09637, -0.08969, -0.01824, + 0.260131, 0.289278, 0.518254, 0.683003, 0.877618, 1.172475, + 1.33964, 1.576766] + x1, x2 = xp.asarray(x1), xp.asarray(x2) + soln = stats.alexandergovern(x1, x2) + xp_assert_close(soln.statistic, xp.asarray(0.713526), atol=1e-5) + xp_assert_close(soln.pvalue, xp.asarray(0.398276), atol=1e-5) + + ''' + tested in ag.test in R: + > library("onewaytests") + > library("tibble") + > x1 <- c(-1.77559, -1.4113, -0.69457, -0.54148, -0.18808, -0.07152, + + 0.04696, 0.051183, 0.148695, 0.168052, 0.422561, 0.458555, + + 0.616123, 0.709968, 0.839956, 0.857226, 0.929159, 0.981442, + + 0.999554, 1.642958) + > x2 <- c(-1.47973, -1.2722, -0.91914, -0.80916, -0.75977, -0.72253, + + -0.3601, -0.33273, -0.28859, -0.09637, -0.08969, -0.01824, + + 0.260131, 0.289278, 0.518254, 0.683003, 0.877618, 1.172475, + + 1.33964, 1.576766) + > x1_fact <- c(rep("x1", times=20)) + > x2_fact <- c(rep("x2", times=20)) + > a <- c(x1, x2) + > b <- factor(c(x1_fact, x2_fact)) + > ag.test(a ~ b, tibble(a, b)) + Alexander-Govern Test (alpha = 0.05) + ------------------------------------------------------------- + data : a and b + + statistic : 0.7135182 + parameter : 1 + p.value : 0.3982783 + + Result : Difference is not statistically significant. + ------------------------------------------------------------- + ''' + xp_assert_close(soln.statistic, xp.asarray(0.7135182), rtol=2e-6) + xp_assert_close(soln.pvalue, xp.asarray(0.3982783), rtol=2e-6) + + def test_nan_policy_propagate(self, xp): + args = xp.asarray([1., 2., 3., 4.]), xp.asarray([1, xp.nan]) + # default nan_policy is 'propagate' + res = stats.alexandergovern(*args) + xp_assert_equal(res.pvalue, xp.asarray(xp.nan)) + xp_assert_equal(res.statistic, xp.asarray(xp.nan)) + + @skip_xp_backends(eager_only=True, reason="lazy -> reduced nan_policy capabilities") + def test_nan_policy_raise(self, xp): + args = xp.asarray([1., 2., 3., 4.]), xp.asarray([1., xp.nan]) + with assert_raises(ValueError, match="The input contains nan values"): + stats.alexandergovern(*args, nan_policy='raise') + + @skip_xp_backends(eager_only=True, reason="lazy -> reduced nan_policy capabilities") + def test_nan_policy_omit(self, xp): + args_nan = xp.asarray([1, 2, 3, xp.nan, 4]), xp.asarray([1, xp.nan, 19, 25]) + args_no_nan = xp.asarray([1, 2, 3, 4]), xp.asarray([1, 19, 25]) + res_nan = stats.alexandergovern(*args_nan, nan_policy='omit') + res_no_nan = stats.alexandergovern(*args_no_nan) + xp_assert_equal(res_nan.pvalue, res_no_nan.pvalue) + xp_assert_equal(res_nan.statistic, res_no_nan.statistic) + + @skip_xp_backends(eager_only=True, reason="lazy -> no data-dependent warnings") + def test_constant_input(self, xp): + # Zero variance input, consistent with `stats.pearsonr` + x1 = xp.asarray([0.667, 0.667, 0.667]) + x2 = xp.asarray([0.123, 0.456, 0.789]) + with pytest.warns(RuntimeWarning, match="Precision loss occurred..."): + res = stats.alexandergovern(x1, x2) + xp_assert_equal(res.statistic, xp.asarray(xp.nan)) + xp_assert_equal(res.pvalue, xp.asarray(xp.nan)) + + @skip_xp_backends('dask.array', reason="Requires `_axis_nan_policy` decorator") + @pytest.mark.parametrize('axis', [0, 1, None]) + def test_2d_input(self, xp, axis): + # much more extensive testing for NumPy in `test_axis_nan_policy.py` + rng = np.random.default_rng(23498389754392854235) + shape = (11, 12) + samples = [rng.random(shape) for i in range(3)] + xp_samples = [xp.asarray(sample) for sample in samples] + ref = stats.alexandergovern(*samples, axis=axis) + res = stats.alexandergovern(*xp_samples, axis=axis) + xp_assert_close(res.statistic, xp.asarray(ref.statistic)) + xp_assert_close(res.pvalue, xp.asarray(ref.pvalue)) + + +@make_xp_test_case(stats.f_oneway) +class TestFOneWay: + + def test_trivial(self, xp): + # A trivial test of stats.f_oneway, with F=0. + F, p = stats.f_oneway(xp.asarray([0, 2]), xp.asarray([0, 2])) + xp_assert_equal(F, xp.asarray(0.0)) + xp_assert_equal(p, xp.asarray(1.0)) + + @pytest.mark.parametrize("dtype", [None, "float32", "float64"]) + def test_basic(self, dtype, xp): + # Despite being a floating point calculation, this data should + # result in F being exactly 2.0. + dtype = None if dtype is None else getattr(xp, dtype) + F, p = stats.f_oneway(xp.asarray([0, 2], dtype=dtype), + xp.asarray([2, 4], dtype=dtype)) + xp_assert_equal(F, xp.asarray(2.0, dtype=dtype)) + xp_assert_close(p, xp.asarray(1 - 0.5**0.5, dtype=dtype)) + + @pytest.mark.parametrize("dtype", [None, "float32", "float64"]) + def test_unequal_var(self, dtype, xp): + # toy samples with unequal variances and different observations + dtype = None if dtype is None else getattr(xp, dtype) + samples = [xp.asarray([-50.42, 40.31, -18.09, 35.58, -6.8, 0.22], dtype=dtype), + xp.asarray([23.44, 4.5, 15.1, 9.66], dtype=dtype), + xp.asarray([11.94, 11.1 , 9.87, 9.09, 3.33], dtype=dtype)] + + F, p = stats.f_oneway(*samples, equal_var=False) + + # R language as benchmark + # group1 <- c(-50.42, 40.31, -18.09, 35.58, -6.8, 0.22) + # group2 <- c(23.44, 4.5, 15.1, 9.66) + # group3 <- c(11.94, 11.1 , 9.87, 9.09, 3.33) + # + # data <- data.frame( + # value = c(group1, group2, group3), + # group = factor(c(rep("G1", length(group1)), + # rep("G2", length(group2)), + # rep("G3", length(group3)))) + # ) + # welch_anova <- oneway.test(value ~ group, data = data, var.equal = FALSE) + # welch_anova$statistic + ## F: 0.609740409019517 + # welch_anova$p.value + ## 0.574838941286302 + + xp_assert_close(F, xp.asarray(0.609740409019517, dtype=dtype)) + xp_assert_close(p, xp.asarray(0.574838941286302, dtype=dtype)) + + def test_equal_var_input_validation(self, xp): + samples = [xp.asarray([-50.42, 40.31, -18.09, 35.58, -6.8, 0.22]), + xp.asarray([23.44, 4.5, 15.1, 9.66]), + xp.asarray([11.94, 11.1 , 9.87, 9.09, 3.33])] + + message = "Expected a boolean value for 'equal_var'" + with pytest.raises(TypeError, match=message): + stats.f_oneway(*samples, equal_var="False") + + def test_known_exact(self, xp): + # Another trivial dataset for which the exact F and p can be + # calculated on most platforms + F, p = stats.f_oneway(xp.asarray([2]), xp.asarray([2]), + xp.asarray([2, 3, 4])) + xp_assert_close(F, xp.asarray(3/5)) # assert_equal fails on some CI platforms + xp_assert_close(p, xp.asarray(5/8)) + + def test_large_integer_array(self, xp): + a = xp.asarray([655, 788], dtype=xp.uint16) + b = xp.asarray([789, 772], dtype=xp.uint16) + F, p = stats.f_oneway(a, b) + # The expected value was verified by computing it with mpmath with + # 40 digits of precision. + xp_assert_close(F, xp.asarray(0.77450216931805540)) + + def test_result_attributes(self, xp): + a = xp.asarray([655, 788], dtype=xp.uint16) + b = xp.asarray([789, 772], dtype=xp.uint16) + res = stats.f_oneway(a, b) + attributes = ('statistic', 'pvalue') + check_named_results(res, attributes, xp=xp) + + @pytest.mark.parametrize('test_case', + ['SiRstv', 'SmLs01', 'SmLs02', 'SmLs03', 'AtmWtAg', 'SmLs04', 'SmLs05', + 'SmLs06', 'SmLs07', 'SmLs08', 'SmLs09']) + def test_nist(self, test_case, xp): + # These are the nist ANOVA files. They can be found at: + # https://www.itl.nist.gov/div898/strd/anova/anova.html + filename = test_case + ".dat" + rtol = 1e-7 + fname = os.path.abspath(os.path.join(os.path.dirname(__file__), + 'data/nist_anova', filename)) + with open(fname) as f: + content = f.read().split('\n') + certified = [line.split() for line in content[40:48] + if line.strip()] + dataf = np.loadtxt(fname, skiprows=60) + y, x = dataf.T + y = y.astype(int) + caty = np.unique(y) + f = float(certified[0][-1]) + + xlist = [xp.asarray(x[y == i]) for i in caty] + res = stats.f_oneway(*xlist) + + # With the hard test cases we relax the tolerance a bit. + hard_tc = ('SmLs07', 'SmLs08', 'SmLs09') + if test_case in hard_tc: + rtol = 1e-4 + + xp_assert_close(res[0], xp.asarray(f, dtype=xp.float64), rtol=rtol) + + @pytest.mark.filterwarnings('ignore') # Dask emits NumPy warnings + @pytest.mark.parametrize("a, b, expected", [ + (np.array([42, 42, 42]), np.array([7, 7, 7]), (np.inf, 0.)), + (np.array([42, 42, 42]), np.array([42, 42, 42]), (np.nan, np.nan)) + ]) + def test_constant_input(self, a, b, expected, xp): + # For more details, look on https://github.com/scipy/scipy/issues/11669 + f, p = stats.f_oneway(xp.asarray(a), xp.asarray(b)) + xp_assert_equal(f, xp.asarray(expected[0])) + xp_assert_equal(p, xp.asarray(expected[1])) + + @skip_xp_backends('dask.array', reason='needs _axis_nan_policy') + @pytest.mark.parametrize('axis', [-2, -1, 0, 1]) + def test_2d_inputs(self, axis, xp): + a = np.array([[1, 4, 3, 3], + [2, 5, 3, 3], + [3, 6, 3, 3], + [2, 3, 3, 3], + [1, 4, 3, 3]], dtype=np.float64) + b = np.array([[3, 1, 5, 3], + [4, 6, 5, 3], + [4, 3, 5, 3], + [1, 5, 5, 3], + [5, 5, 5, 3], + [2, 3, 5, 3], + [8, 2, 5, 3], + [2, 2, 5, 3]], dtype=np.float64) + c = np.array([[4, 3, 4, 3], + [4, 2, 4, 3], + [5, 4, 4, 3], + [5, 4, 4, 3]], dtype=np.float64) + + if axis in [-1, 1]: + a = a.T + b = b.T + c = c.T + take_axis = 0 + else: + take_axis = 1 + + f, p = stats.f_oneway(xp.asarray(a), xp.asarray(b), xp.asarray(c), + axis=axis) + + # Verify that the result computed with the 2d arrays matches + # the result of calling f_oneway individually on each slice. + for j in [0, 1]: + fj, pj = stats.f_oneway(np.take(a, j, take_axis), + np.take(b, j, take_axis), + np.take(c, j, take_axis)) + xp_assert_close(f[j], xp.asarray(fj)) + xp_assert_close(p[j], xp.asarray(pj)) + for j in [2, 3]: + fj, pj = stats.f_oneway(np.take(a, j, take_axis), + np.take(b, j, take_axis), + np.take(c, j, take_axis)) + xp_assert_close(f[j], xp.asarray(fj)) + xp_assert_close(p[j], xp.asarray(pj)) + + @skip_xp_backends('dask.array', reason='needs _axis_nan_policy') + def test_3d_inputs(self, xp): + # Some 3-d arrays. (There is nothing special about the values.) + a = xp.reshape(1/xp.arange(1.0, 4*5*7 + 1., dtype=xp.float64), (4, 5, 7)) + b = xp.reshape(2/xp.arange(1.0, 4*8*7 + 1., dtype=xp.float64), (4, 8, 7)) + c = xp.reshape(1/xp.arange(1.0, 4*4*7 + 1., dtype=xp.float64), (4, 4, 7)) + c = xp.cos(c) + + f, p = stats.f_oneway(a, b, c, axis=1) + + assert f.shape == (4, 7) + assert p.shape == (4, 7) + + for i in range(a.shape[0]): + for j in range(a.shape[2]): + fij, pij = stats.f_oneway(a[i, :, j], b[i, :, j], c[i, :, j]) + xp_assert_close(fij, f[i, j]) + xp_assert_close(pij, p[i, j]) + + def test_length0_1d_error(self, xp): + # Require at least one value in each group. + with eager_warns(SmallSampleWarning, match=too_small_1d_not_omit, xp=xp): + result = stats.f_oneway(xp.asarray([1., 2., 3.]), xp.asarray([]), + xp.asarray([4., 5., 6., 7.])) + xp_assert_equal(result.statistic, xp.asarray(xp.nan)) + xp_assert_equal(result.pvalue, xp.asarray(xp.nan)) + + @skip_xp_backends('dask.array', reason='needs _axis_nan_policy') + def test_length0_2d_error(self, xp): + with eager_warns(SmallSampleWarning, match=too_small_nd_not_omit, xp=xp): + ncols = 3 + a = xp.ones((4, ncols)) + b = xp.ones((0, ncols)) + c = xp.ones((5, ncols)) + f, p = stats.f_oneway(a, b, c) + nans = xp.full((ncols,), fill_value=xp.nan) + xp_assert_equal(f, nans) + xp_assert_equal(p, nans) + + def test_all_length_one(self, xp): + samples = xp.asarray([10]), xp.asarray([11]), xp.asarray([12]), xp.asarray([13]) + with eager_warns(SmallSampleWarning, xp=xp): + result = stats.f_oneway(*samples) + xp_assert_equal(result.statistic, xp.asarray(xp.nan)) + xp_assert_equal(result.pvalue, xp.asarray(xp.nan)) + + @skip_xp_backends('dask.array', reason='lazy->reduced input validation') + @pytest.mark.parametrize('args', [(), ([1, 2, 3],)]) + def test_too_few_inputs(self, args, xp): + args = [xp.asarray(arg) for arg in args] + message = "At least two samples are required..." + with pytest.raises(TypeError, match=message): + stats.f_oneway(*args) + + @skip_xp_backends('dask.array', reason='lazy->reduced input validation') + def test_axis_error(self, xp): + a = xp.ones((3, 4)) + b = xp.ones((5, 4)) + with pytest.raises(AxisError): + stats.f_oneway(a, b, axis=2) + + def test_bad_shapes(self, xp): + a = xp.ones((3, 4)) + b = xp.ones((5, 4)) + with pytest.raises(ValueError): + stats.f_oneway(a, b, axis=1) + + +@make_xp_test_case(stats.kruskal) +class TestKruskal: + def test_array_like(self): + h, p = stats.kruskal([1], [2]) + assert_equal(h, 1.0) + assert_allclose(p, stats.chi2.sf(h, 1)) + + def test_simple(self, xp): + x = xp.asarray([1]) + y = xp.asarray([2]) + h, p = stats.kruskal(x, y) + xp_assert_close(h, xp.asarray(1.0)) + dtype = xp_default_dtype(xp) + xp_assert_close(p, xp.asarray(stats.chi2.sf(1, 1), dtype=dtype)) + + @pytest.mark.parametrize("dtype", ['float32', 'float64', None]) + def test_basic(self, xp, dtype): + dtype = xp_default_dtype(xp) if dtype is None else getattr(xp, dtype) + x = xp.asarray([1, 3, 5, 7, 9], dtype=dtype) + y = xp.asarray([2, 4, 6, 8, 10], dtype=dtype) + h, p = stats.kruskal(x, y) + xp_assert_close(h, xp.asarray(3/11, dtype=dtype)) + xp_assert_close(p, xp.asarray(stats.chi2.sf(3/11, 1), dtype=dtype)) + + def test_simple_tie(self, xp): + x = [1] + y = [1, 2] + h, p = stats.kruskal(xp.asarray(x), xp.asarray(y)) + + h_uncorr = 1.5**2 + 2*2.25**2 - 12 + corr = 0.75 + expected = xp.asarray(h_uncorr / corr) # 0.5 + + # Since the expression is simple and the exact answer is 0.5, it + # should be safe to use assert_equal(). + xp_assert_equal(h, expected) + xp_assert_close(p, special.chdtrc(xp.asarray(1.), expected)) + + def test_another_tie(self, xp): + x = [1, 1, 1, 2] + y = [2, 2, 2, 2] + h, p = stats.kruskal(xp.asarray(x), xp.asarray(y)) + + h_uncorr = (12. / 8. / 9.) * 4 * (3**2 + 6**2) - 3 * 9 + corr = 1 - float(3**3 - 3 + 5**3 - 5) / (8**3 - 8) + expected = xp.asarray(h_uncorr / corr) + + xp_assert_close(h, expected) + xp_assert_close(p, special.chdtrc(xp.asarray(1.), expected)) + + def test_three_groups(self, xp): + # A test of stats.kruskal with three groups, with ties. + x = [1, 1, 1] + y = [2, 2, 2] + z = [2, 2] + h, p = stats.kruskal(xp.asarray(x), xp.asarray(y), xp.asarray(z)) + + h_uncorr = (12. / 8. / 9.) * (3*2**2 + 3*6**2 + 2*6**2) - 3 * 9 # 5.0 + corr = 1 - float(3**3 - 3 + 5**3 - 5) / (8**3 - 8) + expected = xp.asarray(h_uncorr / corr) # 7.0 + + xp_assert_close(h, expected) + xp_assert_close(p, special.chdtrc(xp.asarray(2.), expected)) + + def test_empty(self, xp): + # A test of stats.kruskal with three groups, with ties. + x = xp.asarray([1, 1, 1]) + y = xp.asarray([2, 2, 2]) + z = xp.asarray([], dtype=y.dtype) + with pytest.warns(SmallSampleWarning, match=too_small_1d_not_omit): + res = stats.kruskal(x, y, z) + xp_assert_equal(res.statistic, xp.asarray(xp.nan)) + xp_assert_equal(res.pvalue, xp.asarray(xp.nan)) + + @skip_xp_backends('jax.numpy', reason='needs jax-ml/jax#34628') + def test_nan_policy_propagate(self, xp): + x = xp.arange(10.) + x = xpx.at(x)[9].set(xp.nan) + + res = stats.kruskal(x, x) + xp_assert_equal(res.statistic, xp.asarray(xp.nan)) + xp_assert_equal(res.pvalue, xp.asarray(xp.nan)) + + @skip_xp_backends('jax.numpy', reason='lazy -> reduced nan_policy capabilities') + @xfail_xp_backends('cupy', reason='cupy returns small but nonzero res.statistic') + def test_nan_policy_omit_raise(self, xp): + x = xp.arange(10.) + x = xpx.at(x)[9].set(xp.nan) + + res = stats.kruskal(x, x, nan_policy='omit') + xp_assert_equal(res.statistic, xp.asarray(0.0)) + xp_assert_equal(res.pvalue, xp.asarray(1.0)) + + with pytest.raises(ValueError, match='The input contains nan values'): + stats.kruskal(x, x, nan_policy='raise') + + with pytest.raises(ValueError, match='nan_policy must be one of...'): + stats.kruskal(x, x, nan_policy='foobar') + + def test_large_samples(self, xp): + # Test to see if large samples are handled correctly. + n = 50000 + rng = np.random.default_rng(1808365978) + x = xp.asarray(rng.standard_normal(n)) + y = xp.asarray(rng.standard_normal(n) + 50) + h, p = stats.kruskal(x, y) + xp_assert_close(p, xp.asarray(0., dtype=x.dtype)) + + def test_no_args_gh20661(self, xp): + message = r"Need at least two groups in stats.kruskal\(\)" + with pytest.raises(ValueError, match=message): + stats.kruskal() + + +@make_xp_test_case(stats.combine_pvalues) +class TestCombinePvalues: + # Reference values computed using the following R code: + # options(digits=16) + # library(metap) + # x = c(0.01, 0.2, 0.3) + # sumlog(x) # fisher + # sumz(x) # stouffer + # sumlog(1-x) # pearson (negative statistic and complement of p-value) + # minimump(x) # tippett + @pytest.mark.parametrize( + "method, expected_statistic, expected_pvalue", + [("fisher", 14.83716180549625 , 0.02156175132483465), + ("stouffer", 2.131790594240385, 0.01651203260896294), + ("pearson", -1.179737662212887, 1-0.9778736999143087), + ("tippett", 0.01, 0.02970100000000002), + # mudholkar_george: library(transite); p_combine(x, method="MG") + ("mudholkar_george", 6.828712071641684, 0.01654551838539527)]) + def test_reference_values(self, xp, method, expected_statistic, expected_pvalue): + x = [.01, .2, .3] + res = stats.combine_pvalues(xp.asarray(x), method=method) + xp_assert_close(res.statistic, xp.asarray(expected_statistic)) + xp_assert_close(res.pvalue, xp.asarray(expected_pvalue)) + + @pytest.mark.parametrize( + # Reference values computed using R `metap` `sumz`: + # options(digits=16) + # library(metap) + # x = c(0.01, 0.2, 0.3) + # sumz(x, weights=c(1., 1., 1.)) + # sumz(x, weights=c(1., 4., 9.)) + "weights, expected_statistic, expected_pvalue", + [([1., 1., 1.], 2.131790594240385, 0.01651203260896294), + ([1., 4., 9.], 1.051815015753598, 0.1464422142261314)]) + def test_weighted_stouffer(self, xp, weights, expected_statistic, expected_pvalue): + x = xp.asarray([.01, .2, .3]) + res = stats.combine_pvalues(x, method='stouffer', weights=xp.asarray(weights)) + xp_assert_close(res.statistic, xp.asarray(expected_statistic)) + xp_assert_close(res.pvalue, xp.asarray(expected_pvalue)) + + methods = ["fisher", "pearson", "tippett", "stouffer", "mudholkar_george"] + + @pytest.mark.parametrize("variant", ["single", "all", "random"]) + @pytest.mark.parametrize("method", methods) + def test_monotonicity(self, variant, method, xp): + # Test that result increases monotonically with respect to input. + m, n = 10, 7 + rng = np.random.default_rng(278448169958891062669391462690811630763) + + # `pvaluess` is an m × n array of p values. Each row corresponds to + # a set of p values to be combined with p values increasing + # monotonically down one column (single), simultaneously down each + # column (all), or independently down each column (random). + if variant == "single": + pvaluess = xp.broadcast_to(xp.asarray(rng.random(n)), (m, n)) + pvaluess = xp.concat([xp.reshape(xp.linspace(0.1, 0.9, m), (-1, 1)), + pvaluess[:, 1:]], axis=1) + elif variant == "all": + pvaluess = xp.broadcast_to(xp.linspace(0.1, 0.9, m), (n, m)).T + elif variant == "random": + pvaluess = xp.sort(xp.asarray(rng.uniform(0, 1, size=(m, n))), axis=0) + + combined_pvalues = xp.asarray([ + float(stats.combine_pvalues(pvaluess[i, :], method=method)[1]) + for i in range(pvaluess.shape[0]) + ]) + assert xp.all(combined_pvalues[1:] - combined_pvalues[:-1] >= 0) + + @pytest.mark.parametrize("method", methods) + def test_result(self, method, xp): + res = stats.combine_pvalues(xp.asarray([.01, .2, .3]), method=method) + xp_assert_equal(res.statistic, res[0]) + xp_assert_equal(res.pvalue, res[1]) + + @pytest.mark.parametrize("method", methods) + @pytest.mark.parametrize("axis", [0, 1, None]) + def test_axis(self, method, axis, xp): + rng = np.random.default_rng(234892349810482) + x = xp.asarray(rng.random(size=(2, 10))) + x = x.T if (axis == 0) else x + res = stats.combine_pvalues(x, axis=axis, method=method) + + if axis is None: + x = xp.reshape(x, (-1,)) + ref = stats.combine_pvalues(x, method=method) + xp_assert_close(res.statistic, ref.statistic) + xp_assert_close(res.pvalue, ref.pvalue) + return + + x = x.T if (axis == 0) else x + x0, x1 = x[0, :], x[1, :] + ref0 = stats.combine_pvalues(x0, method=method) + ref1 = stats.combine_pvalues(x1, method=method) + + xp_assert_close(res.statistic[0], ref0.statistic) + xp_assert_close(res.statistic[1], ref1.statistic) + xp_assert_close(res.pvalue[0], ref0.pvalue) + xp_assert_close(res.pvalue[1], ref1.pvalue) + + +class TestCdfDistanceValidation: + """ + Test that _cdf_distance() (via wasserstein_distance()) raises ValueErrors + for bad inputs. + """ + + def test_distinct_value_and_weight_lengths(self): + # When the number of weights does not match the number of values, + # a ValueError should be raised. + assert_raises(ValueError, stats.wasserstein_distance, + [1], [2], [4], [3, 1]) + assert_raises(ValueError, stats.wasserstein_distance, [1], [2], [1, 0]) + + def test_zero_weight(self): + # When a distribution is given zero weight, a ValueError should be + # raised. + assert_raises(ValueError, stats.wasserstein_distance, + [0, 1], [2], [0, 0]) + assert_raises(ValueError, stats.wasserstein_distance, + [0, 1], [2], [3, 1], [0]) + + def test_negative_weights(self): + # A ValueError should be raised if there are any negative weights. + assert_raises(ValueError, stats.wasserstein_distance, + [0, 1], [2, 2], [1, 1], [3, -1]) + + def test_empty_distribution(self): + # A ValueError should be raised when trying to measure the distance + # between something and nothing. + assert_raises(ValueError, stats.wasserstein_distance, [], [2, 2]) + assert_raises(ValueError, stats.wasserstein_distance, [1], []) + + def test_inf_weight(self): + # An inf weight is not valid. + assert_raises(ValueError, stats.wasserstein_distance, + [1, 2, 1], [1, 1], [1, np.inf, 1], [1, 1]) + + +class TestWassersteinDistanceND: + """ Tests for wasserstein_distance_nd() output values. + """ + + def test_published_values(self): + # Compare against published values and manually computed results. + # The values and computed result are posted at James D. McCaffrey's blog, + # https://jamesmccaffrey.wordpress.com/2018/03/05/earth-mover-distance + # -wasserstein-metric-example-calculation/ + u = [(1,1), (1,1), (1,1), (1,1), (1,1), (1,1), (1,1), (1,1), (1,1), (1,1), + (4,2), (6,1), (6,1)] + v = [(2,1), (2,1), (3,2), (3,2), (3,2), (5,1), (5,1), (5,1), (5,1), (5,1), + (5,1), (5,1), (7,1)] + + res = stats.wasserstein_distance_nd(u, v) + # In original post, the author kept two decimal places for ease of calculation. + # This test uses the more precise value of distance to get the precise results. + # For comparison, please see the table and figure in the original blog post. + flow = np.array([2., 3., 5., 1., 1., 1.]) + dist = np.array([1.00, 5**0.5, 4.00, 2**0.5, 1.00, 1.00]) + ref = np.sum(flow * dist)/np.sum(flow) + assert_allclose(res, ref) + + @pytest.mark.parametrize('n_value', (4, 15, 35)) + @pytest.mark.parametrize('ndim', (3, 4, 7)) + @pytest.mark.parametrize('max_repeats', (5, 10)) + def test_same_distribution_nD(self, ndim, n_value, max_repeats): + # Any distribution moved to itself should have a Wasserstein distance + # of zero. + rng = np.random.default_rng(363836384995579937222333) + repeats = rng.integers(1, max_repeats, size=n_value, dtype=int) + + u_values = rng.random(size=(n_value, ndim)) + v_values = np.repeat(u_values, repeats, axis=0) + v_weights = rng.random(np.sum(repeats)) + range_repeat = np.repeat(np.arange(len(repeats)), repeats) + u_weights = np.bincount(range_repeat, weights=v_weights) + index = rng.permutation(len(v_weights)) + v_values, v_weights = v_values[index], v_weights[index] + + res = stats.wasserstein_distance_nd(u_values, v_values, u_weights, v_weights) + assert_allclose(res, 0, atol=1e-15) + + @pytest.mark.parametrize('nu', (8, 9, 38)) + @pytest.mark.parametrize('nv', (8, 12, 17)) + @pytest.mark.parametrize('ndim', (3, 5, 23)) + def test_collapse_nD(self, nu, nv, ndim): + # test collapse for n dimensional values + # Collapsing a n-D distribution to a point distribution at zero + # is equivalent to taking the average of the norm of data. + rng = np.random.default_rng(38573488467338826109) + u_values = rng.random(size=(nu, ndim)) + v_values = np.zeros((nv, ndim)) + u_weights = rng.random(size=nu) + v_weights = rng.random(size=nv) + ref = np.average(np.linalg.norm(u_values, axis=1), weights=u_weights) + res = stats.wasserstein_distance_nd(u_values, v_values, u_weights, v_weights) + assert_allclose(res, ref) + + @pytest.mark.parametrize('nu', (8, 16, 32)) + @pytest.mark.parametrize('nv', (8, 16, 32)) + @pytest.mark.parametrize('ndim', (1, 2, 6)) + def test_zero_weight_nD(self, nu, nv, ndim): + # Values with zero weight have no impact on the Wasserstein distance. + rng = np.random.default_rng(38573488467338826109) + u_values = rng.random(size=(nu, ndim)) + v_values = rng.random(size=(nv, ndim)) + u_weights = rng.random(size=nu) + v_weights = rng.random(size=nv) + ref = stats.wasserstein_distance_nd(u_values, v_values, u_weights, v_weights) + + add_row, nrows = rng.integers(0, nu, size=2) + add_value = rng.random(size=(nrows, ndim)) + u_values = np.insert(u_values, add_row, add_value, axis=0) + u_weights = np.insert(u_weights, add_row, np.zeros(nrows), axis=0) + res = stats.wasserstein_distance_nd(u_values, v_values, u_weights, v_weights) + assert_allclose(res, ref) + + def test_inf_values(self): + # Inf values can lead to an inf distance or trigger a RuntimeWarning + # (and return NaN) if the distance is undefined. + uv, vv, uw = [[1, 1], [2, 1]], [[np.inf, -np.inf]], [1, 1] + distance = stats.wasserstein_distance_nd(uv, vv, uw) + assert_equal(distance, np.inf) + with np.errstate(invalid='ignore'): + uv, vv = [[np.inf, np.inf]], [[np.inf, -np.inf]] + distance = stats.wasserstein_distance_nd(uv, vv) + assert_equal(distance, np.nan) + + @pytest.mark.parametrize('nu', (10, 15, 20)) + @pytest.mark.parametrize('nv', (10, 15, 20)) + @pytest.mark.parametrize('ndim', (1, 3, 5)) + def test_multi_dim_nD(self, nu, nv, ndim): + # Adding dimension on distributions do not affect the result + rng = np.random.default_rng(2736495738494849509) + u_values = rng.random(size=(nu, ndim)) + v_values = rng.random(size=(nv, ndim)) + u_weights = rng.random(size=nu) + v_weights = rng.random(size=nv) + ref = stats.wasserstein_distance_nd(u_values, v_values, u_weights, v_weights) + + add_dim = rng.integers(0, ndim) + add_value = rng.random() + + u_values = np.insert(u_values, add_dim, add_value, axis=1) + v_values = np.insert(v_values, add_dim, add_value, axis=1) + res = stats.wasserstein_distance_nd(u_values, v_values, u_weights, v_weights) + assert_allclose(res, ref) + + @pytest.mark.parametrize('nu', (7, 13, 19)) + @pytest.mark.parametrize('nv', (7, 13, 19)) + @pytest.mark.parametrize('ndim', (2, 4, 7)) + def test_orthogonal_nD(self, nu, nv, ndim): + # orthogonal transformations do not affect the result of the + # wasserstein_distance + rng = np.random.default_rng(34746837464536) + u_values = rng.random(size=(nu, ndim)) + v_values = rng.random(size=(nv, ndim)) + u_weights = rng.random(size=nu) + v_weights = rng.random(size=nv) + ref = stats.wasserstein_distance_nd(u_values, v_values, u_weights, v_weights) + + dist = stats.ortho_group(ndim) + transform = dist.rvs(random_state=rng) + shift = rng.random(size=ndim) + res = stats.wasserstein_distance_nd(u_values @ transform + shift, + v_values @ transform + shift, + u_weights, v_weights) + assert_allclose(res, ref) + + def test_error_code(self): + rng = np.random.default_rng(52473644737485644836320101) + with pytest.raises(ValueError, match='Invalid input values. The inputs'): + u_values = rng.random(size=(4, 10, 15)) + v_values = rng.random(size=(6, 2, 7)) + _ = stats.wasserstein_distance_nd(u_values, v_values) + with pytest.raises(ValueError, match='Invalid input values. Dimensions'): + u_values = rng.random(size=(15,)) + v_values = rng.random(size=(3, 15)) + _ = stats.wasserstein_distance_nd(u_values, v_values) + with pytest.raises(ValueError, + match='Invalid input values. If two-dimensional'): + u_values = rng.random(size=(2, 10)) + v_values = rng.random(size=(2, 2)) + _ = stats.wasserstein_distance_nd(u_values, v_values) + + @pytest.mark.parametrize('u_size', [1, 10, 50]) + @pytest.mark.parametrize('v_size', [1, 10, 50]) + def test_optimization_vs_analytical(self, u_size, v_size): + rng = np.random.default_rng(45634745675) + # Test when u_weights = None, v_weights = None + u_values = rng.random(size=(u_size, 1)) + v_values = rng.random(size=(v_size, 1)) + u_values_flat = u_values.ravel() + v_values_flat = v_values.ravel() + # These three calculations are done using different backends + # but they must be equal + d1 = stats.wasserstein_distance(u_values_flat, v_values_flat) + d2 = stats.wasserstein_distance_nd(u_values, v_values) + d3 = stats.wasserstein_distance_nd(u_values_flat, v_values_flat) + assert_allclose(d2, d1) + assert_allclose(d3, d1) + # Test with u_weights and v_weights specified. + u_weights = rng.random(size=u_size) + v_weights = rng.random(size=v_size) + d1 = stats.wasserstein_distance(u_values_flat, v_values_flat, + u_weights, v_weights) + d2 = stats.wasserstein_distance_nd(u_values, v_values, + u_weights, v_weights) + d3 = stats.wasserstein_distance_nd(u_values_flat, v_values_flat, + u_weights, v_weights) + assert_allclose(d2, d1) + assert_allclose(d3, d1) + + +class TestWassersteinDistance: + """ Tests for wasserstein_distance() output values. + """ + + def test_simple(self): + # For basic distributions, the value of the Wasserstein distance is + # straightforward. + assert_allclose( + stats.wasserstein_distance([0, 1], [0], [1, 1], [1]), + .5) + assert_allclose(stats.wasserstein_distance( + [0, 1], [0], [3, 1], [1]), + .25) + assert_allclose(stats.wasserstein_distance( + [0, 2], [0], [1, 1], [1]), + 1) + assert_allclose(stats.wasserstein_distance( + [0, 1, 2], [1, 2, 3]), + 1) + + def test_same_distribution(self): + # Any distribution moved to itself should have a Wasserstein distance + # of zero. + assert_equal(stats.wasserstein_distance([1, 2, 3], [2, 1, 3]), 0) + assert_equal( + stats.wasserstein_distance([1, 1, 1, 4], [4, 1], + [1, 1, 1, 1], [1, 3]), + 0) + + def test_shift(self): + # If the whole distribution is shifted by x, then the Wasserstein + # distance should be the norm of x. + assert_allclose(stats.wasserstein_distance([0], [1]), 1) + assert_allclose(stats.wasserstein_distance([-5], [5]), 10) + assert_allclose( + stats.wasserstein_distance([1, 2, 3, 4, 5], [11, 12, 13, 14, 15]), + 10) + assert_allclose( + stats.wasserstein_distance([4.5, 6.7, 2.1], [4.6, 7, 9.2], + [3, 1, 1], [1, 3, 1]), + 2.5) + + def test_combine_weights(self): + # Assigning a weight w to a value is equivalent to including that value + # w times in the value array with weight of 1. + assert_allclose( + stats.wasserstein_distance( + [0, 0, 1, 1, 1, 1, 5], [0, 3, 3, 3, 3, 4, 4], + [1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1]), + stats.wasserstein_distance([5, 0, 1], [0, 4, 3], + [1, 2, 4], [1, 2, 4])) + + def test_collapse(self): + # Collapsing a distribution to a point distribution at zero is + # equivalent to taking the average of the absolute values of the + # values. + u = np.arange(-10, 30, 0.3) + v = np.zeros_like(u) + assert_allclose( + stats.wasserstein_distance(u, v), + np.mean(np.abs(u))) + + u_weights = np.arange(len(u)) + v_weights = u_weights[::-1] + assert_allclose( + stats.wasserstein_distance(u, v, u_weights, v_weights), + np.average(np.abs(u), weights=u_weights)) + + def test_zero_weight(self): + # Values with zero weight have no impact on the Wasserstein distance. + assert_allclose( + stats.wasserstein_distance([1, 2, 100000], [1, 1], + [1, 1, 0], [1, 1]), + stats.wasserstein_distance([1, 2], [1, 1], [1, 1], [1, 1])) + + def test_inf_values(self): + # Inf values can lead to an inf distance or trigger a RuntimeWarning + # (and return NaN) if the distance is undefined. + assert_equal( + stats.wasserstein_distance([1, 2, np.inf], [1, 1]), + np.inf) + assert_equal( + stats.wasserstein_distance([1, 2, np.inf], [-np.inf, 1]), + np.inf) + assert_equal( + stats.wasserstein_distance([1, -np.inf, np.inf], [1, 1]), + np.inf) + with pytest.warns(RuntimeWarning, match="invalid value"): + assert_equal( + stats.wasserstein_distance([1, 2, np.inf], [np.inf, 1]), + np.nan) + + +class TestEnergyDistance: + """ Tests for energy_distance() output values. + """ + + def test_simple(self): + # For basic distributions, the value of the energy distance is + # straightforward. + assert_almost_equal( + stats.energy_distance([0, 1], [0], [1, 1], [1]), + np.sqrt(2) * .5) + assert_almost_equal(stats.energy_distance( + [0, 1], [0], [3, 1], [1]), + np.sqrt(2) * .25) + assert_almost_equal(stats.energy_distance( + [0, 2], [0], [1, 1], [1]), + 2 * .5) + assert_almost_equal( + stats.energy_distance([0, 1, 2], [1, 2, 3]), + np.sqrt(2) * (3*(1./3**2))**.5) + + def test_same_distribution(self): + # Any distribution moved to itself should have a energy distance of + # zero. + assert_equal(stats.energy_distance([1, 2, 3], [2, 1, 3]), 0) + assert_equal( + stats.energy_distance([1, 1, 1, 4], [4, 1], [1, 1, 1, 1], [1, 3]), + 0) + + def test_shift(self): + # If a single-point distribution is shifted by x, then the energy + # distance should be sqrt(2) * sqrt(x). + assert_almost_equal(stats.energy_distance([0], [1]), np.sqrt(2)) + assert_almost_equal( + stats.energy_distance([-5], [5]), + np.sqrt(2) * 10**.5) + + def test_combine_weights(self): + # Assigning a weight w to a value is equivalent to including that value + # w times in the value array with weight of 1. + assert_almost_equal( + stats.energy_distance([0, 0, 1, 1, 1, 1, 5], [0, 3, 3, 3, 3, 4, 4], + [1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1]), + stats.energy_distance([5, 0, 1], [0, 4, 3], [1, 2, 4], [1, 2, 4])) + + def test_zero_weight(self): + # Values with zero weight have no impact on the energy distance. + assert_almost_equal( + stats.energy_distance([1, 2, 100000], [1, 1], [1, 1, 0], [1, 1]), + stats.energy_distance([1, 2], [1, 1], [1, 1], [1, 1])) + + def test_inf_values(self): + # Inf values can lead to an inf distance or trigger a RuntimeWarning + # (and return NaN) if the distance is undefined. + assert_equal(stats.energy_distance([1, 2, np.inf], [1, 1]), np.inf) + assert_equal( + stats.energy_distance([1, 2, np.inf], [-np.inf, 1]), + np.inf) + assert_equal( + stats.energy_distance([1, -np.inf, np.inf], [1, 1]), + np.inf) + with pytest.warns(RuntimeWarning, match="invalid value"): + assert_equal( + stats.energy_distance([1, 2, np.inf], [np.inf, 1]), + np.nan) + + +@make_xp_test_case(stats.brunnermunzel) +class TestBrunnerMunzel: + # Data from (Lumley, 1996) + X = [1, 2, 1, 1, 1, 1, 1, 1, 1, 1, 2, 4, 1, 1] + Y = [3, 3, 4, 3, 1, 2, 3, 1, 1, 5, 4] + + def test_brunnermunzel_one_sided(self, xp): + # Results are compared with R's lawstat package. + X, Y = xp.asarray(self.X), xp.asarray(self.Y) + u1, p1 = stats.brunnermunzel(X, Y, alternative='less') + u2, p2 = stats.brunnermunzel(Y, X, alternative='greater') + u3, p3 = stats.brunnermunzel(X, Y, alternative='greater') + u4, p4 = stats.brunnermunzel(Y, X, alternative='less') + + xp_assert_close(p1, p2) + xp_assert_close(p3, p4) + assert p1 != p3 + xp_assert_close(u1, xp.asarray(3.1374674823029505)) + xp_assert_close(u2, xp.asarray(-3.1374674823029505)) + xp_assert_close(u3, xp.asarray(3.1374674823029505)) + xp_assert_close(u4, xp.asarray(-3.1374674823029505)) + + xp_assert_close(p1, xp.asarray(0.0028931043330757342)) + xp_assert_close(p3, xp.asarray(0.99710689566692423)) + + def test_brunnermunzel_two_sided(self, xp): + # Results are compared with R's lawstat package. + X, Y = xp.asarray(self.X), xp.asarray(self.Y) + u1, p1 = stats.brunnermunzel(X, Y, alternative='two-sided') + u2, p2 = stats.brunnermunzel(Y, X, alternative='two-sided') + + xp_assert_close(p1, xp.asarray(p2)) + xp_assert_close(u1, xp.asarray(3.1374674823029505)) + xp_assert_close(u2, xp.asarray(-3.1374674823029505)) + xp_assert_close(p1, xp.asarray(0.0057862086661515377)) + + def test_brunnermunzel_default(self, xp): + # The default value for alternative is two-sided + X, Y = xp.asarray(self.X), xp.asarray(self.Y) + u1, p1 = stats.brunnermunzel(X, Y) + u2, p2 = stats.brunnermunzel(Y, X) + + xp_assert_close(p1, p2) + xp_assert_close(u1, xp.asarray(3.1374674823029505)) + xp_assert_close(u2, xp.asarray(-3.1374674823029505)) + xp_assert_close(p1, xp.asarray(0.0057862086661515377)) + + def test_brunnermunzel_alternative_error(self, xp): + alternative = "error" + distribution = "t" + nan_policy = "propagate" + assert alternative not in ["two-sided", "greater", "less"] + message = "`alternative` must be 'less', 'greater', or 'two-sided'." + with pytest.raises(ValueError, match=message): + stats.brunnermunzel(xp.asarray(self.X), xp.asarray(self.Y), + alternative, distribution, nan_policy) + + def test_brunnermunzel_distribution_norm(self, xp): + X, Y = xp.asarray(self.X), xp.asarray(self.Y) + u1, p1 = stats.brunnermunzel(X, Y, distribution="normal") + u2, p2 = stats.brunnermunzel(Y, X, distribution="normal") + xp_assert_close(p1, xp.asarray(p2)) + xp_assert_close(u1, xp.asarray(3.1374674823029505)) + xp_assert_close(u2, xp.asarray(-3.1374674823029505)) + xp_assert_close(p1, xp.asarray(0.0017041417600383024)) + + def test_brunnermunzel_distribution_error(self, xp): + alternative = "two-sided" + distribution = "error" + nan_policy = "propagate" + assert distribution not in ["t", "normal"] + message = "distribution should be 't' or 'normal'" + with pytest.raises(ValueError, match=message): + stats.brunnermunzel(xp.asarray(self.X), xp.asarray(self.Y), + alternative, distribution, nan_policy) + + @pytest.mark.parametrize("kwarg_update", [{'y': []}, {'x': []}, + {'x': [], 'y': []}]) + def test_brunnermunzel_empty_imput(self, kwarg_update, xp): + kwargs = {'x': self.X, 'y': self.Y} + kwargs.update(kwarg_update) + kwargs = {key:xp.asarray(val, dtype=xp_default_dtype(xp)) + for key, val in kwargs.items()} + with eager_warns(SmallSampleWarning, match=too_small_1d_not_omit, xp=xp): + statistic, pvalue = stats.brunnermunzel(**kwargs) + xp_assert_equal(statistic, xp.asarray(xp.nan)) + xp_assert_equal(pvalue, xp.asarray(xp.nan)) + + @skip_xp_backends('jax.numpy', reason='needs jax-ml/jax#34628') + def test_brunnermunzel_nan_input_propagate(self, xp): + X = xp.asarray([1, 2, 1, 1, 1, 1, 1, 1, 1, 1, 2, 4, 1, 1, xp.nan]) + Y = xp.asarray([3, 3, 4, 3, 1, 2, 3, 1, 1, 5, 4.]) + u1, p1 = stats.brunnermunzel(X, Y, nan_policy="propagate") + u2, p2 = stats.brunnermunzel(Y, X, nan_policy="propagate") + + xp_assert_equal(u1, xp.asarray(xp.nan)) + xp_assert_equal(p1, xp.asarray(xp.nan)) + xp_assert_equal(u2, xp.asarray(xp.nan)) + xp_assert_equal(p2, xp.asarray(xp.nan)) + + @skip_xp_backends('jax.numpy', reason='lazy -> reduced nan_policy capabilities') + def test_brunnermunzel_nan_input_raise(self, xp): + X = xp.asarray([1, 2, 1, 1, 1, 1, 1, 1, 1, 1, 2, 4, 1, 1, xp.nan]) + Y = xp.asarray([3, 3, 4, 3, 1, 2, 3, 1, 1, 5, 4.]) + alternative = "two-sided" + distribution = "t" + nan_policy = "raise" + + message = "The input contains nan values" + with pytest.raises(ValueError, match=message): + stats.brunnermunzel(X, Y, alternative, distribution, nan_policy) + with pytest.raises(ValueError, match=message): + stats.brunnermunzel(Y, X, alternative, distribution, nan_policy) + + @skip_xp_backends('jax.numpy', reason='lazy -> reduced nan_policy capabilities') + def test_brunnermunzel_nan_input_omit(self, xp): + X = xp.asarray([1, 2, 1, 1, 1, 1, 1, 1, 1, 1, 2, 4, 1, 1, np.nan]) + Y = xp.asarray([3, 3, 4, 3, 1, 2, 3, 1, 1, 5, 4.]) + u1, p1 = stats.brunnermunzel(X, Y, nan_policy="omit") + u2, p2 = stats.brunnermunzel(Y, X, nan_policy="omit") + + xp_assert_close(p1, p2) + xp_assert_close(u1, xp.asarray(3.1374674823029505)) + xp_assert_close(u2, xp.asarray(-3.1374674823029505)) + xp_assert_close(p1, xp.asarray(0.0057862086661515377)) + + @pytest.mark.parametrize("vectorized_call", [False, True]) + def test_brunnermunzel_return_nan(self, vectorized_call, xp): + """ tests that a warning is emitted when p is nan + p-value with t-distributions can be nan (0/0) (see gh-15843) + """ + x = xp.asarray([1., 2., 3.]) + y = xp.asarray([5., 6., 7., 8., 9.]) + + if vectorized_call: + x = xp.stack((x, x)).T + y = xp.stack((y, y)).T + + msg = "p-value cannot be estimated|divide by zero|invalid value encountered" + with eager_warns(RuntimeWarning, match=msg, xp=xp): + stats.brunnermunzel(x, y, distribution="t") + + def test_brunnermunzel_normal_dist(self, xp): + """ tests that a p is 0 for datasets that cause p->nan + when t-distribution is used (see gh-15843) + """ + x = xp.asarray([1., 2., 3.]) + y = xp.asarray([5., 6., 7., 8., 9.]) + + with eager_warns(RuntimeWarning, match='divide by zero', xp=xp): + _, p = stats.brunnermunzel(x, y, distribution="normal") + xp_assert_equal(p, xp.asarray(0.)) + + +@_apply_over_batch(('x', 1), ('q', 1), ('p', 1)) +def quantile_test_reference(x, q, p, alternative): + res = stats.quantile_test(x, q=q, p=p, alternative=alternative) + return res.statistic, res.pvalue, *res.confidence_interval() + + +@make_xp_test_case(stats.quantile_test) +class TestQuantileTest: + r""" Test the non-parametric quantile test, + including the computation of confidence intervals + """ + + def test_quantile_test_iv(self, xp): + x = xp.asarray([1, 2, 3]) + + message = "`x` must be an array of numbers." + with pytest.raises(ValueError, match=message): + stats.quantile_test(xp.astype(x, xp.bool)) + + message = "`q` must be a scalar or array of numbers." + with pytest.raises(ValueError, match=message): + stats.quantile_test(x, q=False) + + message = "`p` must be a scalar or array of floats." + with pytest.raises(ValueError, match=message): + stats.quantile_test(x, p=2) + + # should p == 0. / p == 1. be valid? + p = xp.asarray([-1., 0., 1., 1.5, xp.nan]) + res = stats.quantile_test(x, p=p) + xp_assert_equal(res.statistic, xp.full_like(p, xp.nan)) + xp_assert_equal(res.statistic_type, xp.full_like(p, xp.nan)) + xp_assert_equal(res.pvalue, xp.full_like(p, xp.nan)) + + message = "`axis` must be an integer or None." + with pytest.raises(ValueError, match=message): + stats.quantile_test(x, axis=2.5) + + message = "`axis` is not compatible with the shapes of the inputs." + with pytest.raises(ValueError, match=message): + stats.quantile_test(x, axis=2) + + message = "If specified, `keepdims` must be True or False." + with pytest.raises(ValueError, match=message): + stats.quantile_test(x, keepdims=10) + + message = ("`keepdims` may be False only if...") + with pytest.raises(ValueError, match=message): + stats.quantile_test(x, p=xp.asarray([0.1, 0.2]), keepdims=False) + + message = "`alternative` must be one of..." + with pytest.raises(ValueError, match=message): + stats.quantile_test(x, alternative='one-sided') + + message = "`confidence_level` must be a number between 0 and 1." + with pytest.raises(ValueError, match=message): + stats.quantile_test(x).confidence_interval(1) + + @pytest.mark.parametrize( + 'p, alpha, lb, ub, alternative', + [[0.3, 0.95, 1.221402758160170, 1.476980793882643, 'two-sided'], + [0.5, 0.9, 1.506817785112854, 1.803988415397857, 'two-sided'], + [0.25, 0.95, -np.inf, 1.39096812846378, 'less'], + [0.8, 0.9, 2.117000016612675, np.inf, 'greater']] + ) + def test_R_ci_quantile(self, p, alpha, lb, ub, alternative, xp): + # Test against R library `confintr` function `ci_quantile`, e.g. + # library(confintr) + # options(digits=16) + # x <- exp(seq(0, 1, by = 0.01)) + # ci_quantile(x, q = 0.3)$interval + # ci_quantile(x, q = 0.5, probs = c(0.05, 0.95))$interval + # ci_quantile(x, q = 0.25, probs = c(0, 0.95))$interval + # ci_quantile(x, q = 0.8, probs = c(0.1, 1))$interval + x = xp.exp(xp.arange(0.0, 1.01, 0.01)) + res = stats.quantile_test(x, p=p, alternative=alternative) + ci = res.confidence_interval(alpha) + xp_assert_close(ci.low, xp.asarray(lb)) + xp_assert_close(ci.high, xp.asarray(ub)) + + @pytest.mark.parametrize( + 'q, p, alternative, ref', + [[1.2, 0.3, 'two-sided', 0.01515567517648], + [1.8, 0.5, 'two-sided', 0.1109183496606]] + ) + def test_R_pvalue(self, q, p, alternative, ref, xp): + # Test against R library `snpar` function `quant.test`, e.g. + # library(snpar) + # options(digits=16) + # x < - exp(seq(0, 1, by=0.01)) + # quant.test(x, q=1.2, p=0.3, exact=TRUE, alternative='t') + x = xp.exp(xp.arange(0.0, 1.01, 0.01)) + res = stats.quantile_test(x, q=q, p=p, alternative=alternative) + xp_assert_close(res.pvalue, xp.asarray(ref)) + + @pytest.mark.parametrize('case', ['continuous', 'discrete']) + @pytest.mark.parametrize('alternative', ['less', 'greater']) + @pytest.mark.parametrize('alpha', [0.9, 0.95]) + def test_pval_ci_match(self, case, alternative, alpha, xp): + # Verify that the following statement holds: + + # The 95% confidence interval corresponding with alternative='less' + # has -inf as its lower bound, and the upper bound `xu` is the greatest + # element from the sample `x` such that: + # `stats.quantile_test(x, q=xu, p=p, alternative='less').pvalue`` + # will be greater than 5%. + + # And the corresponding statement for the alternative='greater' case. + + seed = int((7**len(case) + len(alternative))*alpha) + rng = np.random.default_rng(seed) + if case == 'continuous': + p, q = rng.random(size=2) + rvs = rng.random(size=100) + else: + rvs = rng.integers(1, 11, size=100) + p = rng.random() + q = rng.integers(1, 11) + + rvs, p, q = xp.asarray(rvs.tolist()), float(p), float(q) + res = stats.quantile_test(rvs, q=q, p=p, alternative=alternative) + ci = res.confidence_interval(confidence_level=alpha) + + # select elements inside the confidence interval based on alternative + if alternative == 'less': + i_inside = xp.astype(rvs, ci.high.dtype) <= ci.high + else: + i_inside = xp.astype(rvs, ci.low.dtype) >= ci.low + + for x in rvs[i_inside]: + res = stats.quantile_test(rvs, q=x, p=p, alternative=alternative) + assert res.pvalue > 1 - alpha + + for x in rvs[~i_inside]: + res = stats.quantile_test(rvs, q=x, p=p, alternative=alternative) + assert res.pvalue < 1 - alpha + + @pytest.mark.parametrize('dtype', [None, 'float32', 'float64']) + def test_match_conover_examples(self, dtype, xp): + # Test against the examples in [1] (Conover Practical Nonparametric + # Statistics Third Edition) pg 139 + + # Example 1 + # Data is [189, 233, 195, 160, 212, 176, 231, 185, 199, 213, 202, 193, + # 174, 166, 248] + # Two-sided test of whether the upper quartile (p=0.75) equals 193 + # (q=193). Conover shows that 7 of the observations are less than or + # equal to 193, and "for the binomial random variable Y, P(Y<=7) = + # 0.0173", so the two-sided p-value is twice that, 0.0346. + dtype = xp_default_dtype(xp) if dtype is None else getattr(xp, dtype) + x = xp.asarray([189, 233, 195, 160, 212, 176, 231, 185, + 199, 213, 202, 193, 174, 166, 248], dtype=dtype) + pvalue_ref = 0.0346 + res = stats.quantile_test(x, q=193, p=0.75, alternative='two-sided') + xp_assert_close(res.pvalue, xp.asarray(pvalue_ref, dtype=dtype), rtol=1e-5) + + # Example 2 + # Conover doesn't give explicit data, just that 8 out of 112 + # observations are 60 or less. The test is whether the median time is + # equal to 60 against the alternative that the median is greater than + # 60. The p-value is calculated as P(Y<=8), where Y is again a binomial + # distributed random variable, now with p=0.5 and n=112. Conover uses a + # normal approximation, but we can easily calculate the CDF of the + # binomial distribution. + x = xp.asarray([59]*8 + [61]*(112-8), dtype=dtype) + pvalue_ref = float(stats.binom(p=0.5, n=112).pmf(k=8)) + res = stats.quantile_test(x, q=60, p=0.5, alternative='greater') + xp_assert_close(res.pvalue, xp.asarray(pvalue_ref, dtype=dtype), atol=1e-10) + + @skip_xp_backends('array_api_strict', reason='indexing in _apply_over_batch') + @pytest.mark.parametrize('alternative', ['less', 'greater', 'two-sided']) + @pytest.mark.parametrize('axis', [0, 1]) + def test_multidimensional(self, alternative, axis, xp): + rng = np.random.default_rng(85468924398205602) + x = rng.random(size=(2, 20)) + p = rng.random(size=(1, 3,)) + q = p + rng.random(size=(2, 1))*0.01 + + x, p, q = (xp.asarray(arr.tolist()) for arr in (x, p, q)) + + ref = quantile_test_reference(x, q, p, alternative=alternative) + ref_statistic, ref_pvalue, ref_low, ref_high = ref + + if axis == 0: + x, p, q = x.T, p.T, q.T + res = stats.quantile_test(x, q=q, p=p, alternative=alternative, axis=axis) + res_statistic, res_pvalue = res.statistic, res.pvalue + res_low, res_high = res.confidence_interval() + if axis == 0: + res_statistic, res_pvalue = res_statistic.T, res_pvalue.T + res_low, res_high = res_low.T, res_high.T + + xp_assert_close(res_statistic, ref_statistic) + xp_assert_close(res_pvalue, ref_pvalue) + xp_assert_close(res_low, res_low) + xp_assert_close(res_low, ref_low) + + def test_zero_size(self, xp): + rng = np.random.default_rng(883771738488451943) + x_shape = (2, 50) + qp_shape = (2, 10) + x = rng.random(x_shape) + q = rng.random(qp_shape) + p = q + rng.random(qp_shape) * 1e-2 + + x, p, q = (xp.asarray(arr.tolist()) for arr in (x, p, q)) + + # case 1: p/q is size zero. + qp_zero = xp.empty((0, *qp_shape)) + out = qp_zero + res = stats.quantile_test(x, q=qp_zero, p=qp_zero, axis=-1) + ci = res.confidence_interval() + # different default int on different platforms + xp_assert_equal(res.statistic, xp.asarray(out)) + xp_assert_equal(res.statistic_type, xp.asarray(out)) + xp_assert_equal(res.pvalue, out) + xp_assert_equal(ci.low, out) + xp_assert_equal(ci.high, out) + + # case 2: x is size zero with nonzero length along axis. + x_zero = xp.empty((0, *x_shape)) + out = xp.empty((0, *qp_shape)) + res = stats.quantile_test(x_zero, q=q, p=p, axis=-1) + ci = res.confidence_interval() + # different default int on different platforms + xp_assert_equal(res.statistic, xp.asarray(out)) + xp_assert_equal(res.statistic_type, xp.asarray(out)) + xp_assert_equal(res.pvalue, out) + xp_assert_equal(ci.low, out) + xp_assert_equal(ci.high, out) + + # case 3: x has zero length along axis. + x_zero = xp.empty((x.shape[0], 0)) + out = xp.full(qp_shape, xp.nan) + res = stats.quantile_test(x_zero, q=q, p=p, axis=-1) + ci = res.confidence_interval() + # different default int on different platforms + xp_assert_equal(res.statistic, out) + xp_assert_equal(res.statistic_type, out) + xp_assert_equal(res.pvalue, out) + xp_assert_equal(ci.low, out) + xp_assert_equal(ci.high, out) + + def test_nans(self, xp): + rng = np.random.default_rng(2920028761208905) + x = rng.random((10, 50)) + q = rng.random((10, 10)) + p = q + rng.random((10, 10)) * 1e-2 + x, p, q = (xp.asarray(arr.tolist()) for arr in (x, p, q)) + + ref = stats.quantile_test(x, q=q, p=p, axis=-1) + ref_low, ref_high = ref.confidence_interval() + + x_ = xp_copy(x) + i_nan = xp.asarray(rng.random(x.shape) < 0.01) + x_ = xpx.at(x_)[i_nan].set(xp.nan) + i_nan_out = xp.any(i_nan, axis=-1) + assert not xp.all(i_nan_out) + res = stats.quantile_test(x_, q=q, p=p, axis=-1) + res_low, res_high = res.confidence_interval() + + assert xp.all(xp.isnan(res.statistic[i_nan_out])) + assert xp.all(xp.isnan(res.statistic_type[i_nan_out])) + assert xp.all(xp.isnan(res.pvalue[i_nan_out])) + assert xp.all(xp.isnan(res_low[i_nan_out])) + assert xp.all(xp.isnan(res_high[i_nan_out])) + xp_assert_equal(res.statistic[~i_nan_out], ref.statistic[~i_nan_out]) + xp_assert_equal(res.statistic_type[~i_nan_out], ref.statistic_type[~i_nan_out]) + xp_assert_equal(res.pvalue[~i_nan_out], ref.pvalue[~i_nan_out]) + xp_assert_equal(res_low[~i_nan_out], ref_low[~i_nan_out]) + xp_assert_equal(res_high[~i_nan_out], ref_high[~i_nan_out]) + + i_nan_q = xp.asarray(rng.random(q.shape) < 0.01) + i_nan_p = xp.asarray(rng.random(p.shape) < 0.01) + assert xp.any(i_nan_q) + assert xp.any(i_nan_p) + q = xpx.at(q)[xp.broadcast_to(i_nan_q, q.shape)].set(xp.nan) + p = xpx.at(p)[xp.broadcast_to(i_nan_p, p.shape)].set(xp.nan) + + i_nan_out = xp.asarray(i_nan_q | i_nan_p) + assert not xp.all(i_nan_out) + + res = stats.quantile_test(x, q=q, p=p, axis=-1) + res_low, res_high = res.confidence_interval() + assert xp.all(xp.isnan(res.statistic[i_nan_out])) + assert xp.all(xp.isnan(res.statistic_type[i_nan_out])) + assert xp.all(xp.isnan(res.pvalue[i_nan_out])) + assert xp.all(xp.isnan(res_low[i_nan_out])) + assert xp.all(xp.isnan(res_high[i_nan_out])) + xp_assert_equal(res.statistic[~i_nan_out], ref.statistic[~i_nan_out]) + xp_assert_equal(res.statistic_type[~i_nan_out], ref.statistic_type[~i_nan_out]) + xp_assert_equal(res.pvalue[~i_nan_out], ref.pvalue[~i_nan_out]) + xp_assert_equal(res_low[~i_nan_out], ref_low[~i_nan_out]) + xp_assert_equal(res_high[~i_nan_out], ref_high[~i_nan_out]) + + def test_axis_none(self, xp): + rng = np.random.default_rng(883771738488451943) + x_shape = (2, 50) + qp_shape = (2, 10) + x = rng.random(x_shape) + q = rng.random(qp_shape) + p = q + rng.random(qp_shape) * 1e-2 + x, p, q = (xp.asarray(arr.tolist()) for arr in (x, p, q)) + + res = stats.quantile_test(x, q=q, p=p, axis=None) + res_low, res_high = res.confidence_interval() + ref = stats.quantile_test(xp_ravel(x), q=xp_ravel(q), p=xp_ravel(p), axis=None) + ref_low, ref_high = ref.confidence_interval() + xp_assert_equal(res.statistic, ref.statistic[xp.newaxis, :]) + xp_assert_equal(res.statistic_type, ref.statistic_type[xp.newaxis, :]) + xp_assert_equal(res.pvalue, ref.pvalue[xp.newaxis, :]) + xp_assert_equal(res_low, ref_low[xp.newaxis, :]) + xp_assert_equal(res_high, ref_high[xp.newaxis, :]) + + res = stats.quantile_test(x, q=0.5, p=0.5, axis=None) + res_low, res_high = res.confidence_interval() + ref = stats.quantile_test(xp_ravel(x), q=0.5, p=0.5, axis=None) + ref_low, ref_high = ref.confidence_interval() + xp_assert_equal(res.statistic, ref.statistic) + xp_assert_equal(res.statistic_type, ref.statistic_type) + xp_assert_equal(res.pvalue, ref.pvalue) + xp_assert_equal(res_low, ref_low) + xp_assert_equal(res_high, ref_high) + + @pytest.mark.filterwarnings("ignore:torch:UserWarning") + @pytest.mark.parametrize("x_shape, qp_shape, axis, keepdims, res_shape", [ + ((1, 3), (1,), None, False, ()), + ((1, 3), (1,), None, True, (1, 1)), + ((1, 3), (2, 1), -1, False, (2,)), + ((1, 3), (2, 1), -1, True, (2, 1)), + ]) + def test_keepdims(self, x_shape, qp_shape, axis, keepdims, res_shape, xp): + rng = np.random.default_rng(883771738488451943) + x, qp = xp.asarray(rng.random(x_shape)), xp.asarray(rng.random(qp_shape)) + res = stats.quantile_test(x, q=qp, p=qp, axis=axis, keepdims=keepdims) + assert res.statistic.shape == res_shape + + +class TestPageTrendTest: + def setup_method(self): + self.rng = np.random.default_rng(1808365978) + + # expected statistic and p-values generated using R at + # https://rdrr.io/cran/cultevo/, e.g. + # library(cultevo) + # data = rbind(c(72, 47, 73, 35, 47, 96, 30, 59, 41, 36, 56, 49, 81, 43, + # 70, 47, 28, 28, 62, 20, 61, 20, 80, 24, 50), + # c(68, 52, 60, 34, 44, 20, 65, 88, 21, 81, 48, 31, 31, 67, + # 69, 94, 30, 24, 40, 87, 70, 43, 50, 96, 43), + # c(81, 13, 85, 35, 79, 12, 92, 86, 21, 64, 16, 64, 68, 17, + # 16, 89, 71, 43, 43, 36, 54, 13, 66, 51, 55)) + # result = page.test(data, verbose=FALSE) + # Most test cases generated to achieve common critical p-values so that + # results could be checked (to limited precision) against tables in + # scipy.stats.page_trend_test reference [1] + + rng = np.random.default_rng(3113562111) + data_3_25 = rng.random((3, 25)) + rng = np.random.default_rng(3113562111) + data_10_26 = rng.random((10, 26)) + + ts = [ + (12949, 0.275539045444, False, 'asymptotic', data_3_25), + (47221, 0.5703651063709, False, 'asymptotic', data_10_26), + (12332, 0.7722477197436702, False, 'asymptotic', + [[72, 47, 73, 35, 47, 96, 30, 59, 41, 36, 56, 49, 81, + 43, 70, 47, 28, 28, 62, 20, 61, 20, 80, 24, 50], + [68, 52, 60, 34, 44, 20, 65, 88, 21, 81, 48, 31, 31, + 67, 69, 94, 30, 24, 40, 87, 70, 43, 50, 96, 43], + [81, 13, 85, 35, 79, 12, 92, 86, 21, 64, 16, 64, 68, + 17, 16, 89, 71, 43, 43, 36, 54, 13, 66, 51, 55]]), + (266, 4.121656378600823e-05, False, 'exact', + [[1.5, 4., 8.3, 5, 19, 11], + [5, 4, 3.5, 10, 20, 21], + [8.4, 3.2, 10, 12, 14, 15]]), + (332, 0.9566400920502488, True, 'exact', + [[4, 3, 2, 1], [4, 3, 2, 1], [4, 3, 2, 1], [4, 3, 2, 1], + [4, 3, 2, 1], [4, 3, 2, 1], [4, 3, 2, 1], [4, 3, 2, 1], + [3, 4, 1, 2], [1, 2, 3, 4], [1, 2, 3, 4], [1, 2, 3, 4], + [1, 2, 3, 4], [1, 2, 3, 4]]), + (241, 0.9622210164861476, True, 'exact', + [[3, 2, 1], [3, 2, 1], [3, 2, 1], [3, 2, 1], [3, 2, 1], [3, 2, 1], + [3, 2, 1], [3, 2, 1], [3, 2, 1], [3, 2, 1], [3, 2, 1], [3, 2, 1], + [3, 2, 1], [2, 1, 3], [1, 2, 3], [1, 2, 3], [1, 2, 3], [1, 2, 3], + [1, 2, 3], [1, 2, 3], [1, 2, 3]]), + (197, 0.9619432897162209, True, 'exact', + [[6, 5, 4, 3, 2, 1], [6, 5, 4, 3, 2, 1], [1, 3, 4, 5, 2, 6]]), + (423, 0.9590458306880073, True, 'exact', + [[5, 4, 3, 2, 1], [5, 4, 3, 2, 1], [5, 4, 3, 2, 1], + [5, 4, 3, 2, 1], [5, 4, 3, 2, 1], [5, 4, 3, 2, 1], + [4, 1, 3, 2, 5], [1, 2, 3, 4, 5], [1, 2, 3, 4, 5], + [1, 2, 3, 4, 5]]), + (217, 0.9693058575034678, True, 'exact', + [[3, 2, 1], [3, 2, 1], [3, 2, 1], [3, 2, 1], [3, 2, 1], [3, 2, 1], + [3, 2, 1], [3, 2, 1], [3, 2, 1], [3, 2, 1], [3, 2, 1], [3, 2, 1], + [2, 1, 3], [1, 2, 3], [1, 2, 3], [1, 2, 3], [1, 2, 3], [1, 2, 3], + [1, 2, 3]]), + (395, 0.991530289351305, True, 'exact', + [[7, 6, 5, 4, 3, 2, 1], [7, 6, 5, 4, 3, 2, 1], + [6, 5, 7, 4, 3, 2, 1], [1, 2, 3, 4, 5, 6, 7]]), + (117, 0.9997817843373017, True, 'exact', + [[3, 2, 1], [3, 2, 1], [3, 2, 1], [3, 2, 1], [3, 2, 1], [3, 2, 1], + [3, 2, 1], [3, 2, 1], [3, 2, 1], [2, 1, 3], [1, 2, 3]]), + ] + + @pytest.mark.parametrize("L, p, ranked, method, data", ts) + def test_accuracy(self, L, p, ranked, method, data): + res = stats.page_trend_test(data, ranked=ranked, method=method) + assert_equal(L, res.statistic) + assert_allclose(p, res.pvalue) + assert_equal(method, res.method) + + ts2 = [ + (542, 0.9481266260876332, True, 'exact', + [[10, 9, 8, 7, 6, 5, 4, 3, 2, 1], + [1, 8, 4, 7, 6, 5, 9, 3, 2, 10]]), + (1322, 0.9993113928199309, True, 'exact', + [[10, 9, 8, 7, 6, 5, 4, 3, 2, 1], [10, 9, 8, 7, 6, 5, 4, 3, 2, 1], + [10, 9, 8, 7, 6, 5, 4, 3, 2, 1], [9, 2, 8, 7, 6, 5, 4, 3, 10, 1], + [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]]), + (2286, 0.9908688345484833, True, 'exact', + [[8, 7, 6, 5, 4, 3, 2, 1], [8, 7, 6, 5, 4, 3, 2, 1], + [8, 7, 6, 5, 4, 3, 2, 1], [8, 7, 6, 5, 4, 3, 2, 1], + [8, 7, 6, 5, 4, 3, 2, 1], [8, 7, 6, 5, 4, 3, 2, 1], + [8, 7, 6, 5, 4, 3, 2, 1], [8, 7, 6, 5, 4, 3, 2, 1], + [8, 7, 6, 5, 4, 3, 2, 1], [1, 3, 5, 6, 4, 7, 2, 8], + [1, 2, 3, 4, 5, 6, 7, 8], [1, 2, 3, 4, 5, 6, 7, 8], + [1, 2, 3, 4, 5, 6, 7, 8], [1, 2, 3, 4, 5, 6, 7, 8], + [1, 2, 3, 4, 5, 6, 7, 8]]), + ] + # only the first of these appears slow because intermediate data are + # cached and used on the rest + @pytest.mark.parametrize("L, p, ranked, method, data", ts2) + @pytest.mark.slow() + def test_accuracy2(self, L, p, ranked, method, data): + res = stats.page_trend_test(data, ranked=ranked, method=method) + assert_equal(L, res.statistic) + assert_allclose(p, res.pvalue) + assert_equal(method, res.method) + + def test_options(self): + rng = np.random.default_rng(183973867) + m, n = 10, 20 + predicted_ranks = np.arange(1, n+1) + perm = rng.permutation(np.arange(n)) + data = rng.random((m, n)) + ranks = stats.rankdata(data, axis=1) + res1 = stats.page_trend_test(ranks) + res2 = stats.page_trend_test(ranks, ranked=True) + res3 = stats.page_trend_test(data, ranked=False) + res4 = stats.page_trend_test(ranks, predicted_ranks=predicted_ranks) + res5 = stats.page_trend_test(ranks[:, perm], + predicted_ranks=predicted_ranks[perm]) + assert_equal(res1.statistic, res2.statistic) + assert_equal(res1.statistic, res3.statistic) + assert_equal(res1.statistic, res4.statistic) + assert_equal(res1.statistic, res5.statistic) + + def test_Ames_assay(self): + # test from _page_trend_test.py [2] page 151; data on page 144 + data = [[101, 117, 111], [91, 90, 107], [103, 133, 121], + [136, 140, 144], [190, 161, 201], [146, 120, 116]] + data = np.array(data).T + predicted_ranks = np.arange(1, 7) + + res = stats.page_trend_test(data, ranked=False, + predicted_ranks=predicted_ranks, + method="asymptotic") + assert_equal(res.statistic, 257) + assert_almost_equal(res.pvalue, 0.0035, decimal=4) + + res = stats.page_trend_test(data, ranked=False, + predicted_ranks=predicted_ranks, + method="exact") + assert_equal(res.statistic, 257) + assert_almost_equal(res.pvalue, 0.0023, decimal=4) + + def test_input_validation(self): + # test data not a 2d array + with assert_raises(ValueError, match="`data` must be a 2d array."): + stats.page_trend_test(None) + with assert_raises(ValueError, match="`data` must be a 2d array."): + stats.page_trend_test([]) + with assert_raises(ValueError, match="`data` must be a 2d array."): + stats.page_trend_test([1, 2]) + with assert_raises(ValueError, match="`data` must be a 2d array."): + stats.page_trend_test([[[1]]]) + + # test invalid dimensions + rng = np.random.default_rng(2482566048) + with assert_raises(ValueError, match="Page's L is only appropriate"): + stats.page_trend_test(rng.random((1, 3))) + with assert_raises(ValueError, match="Page's L is only appropriate"): + stats.page_trend_test(rng.random((2, 2))) + + # predicted ranks must include each integer [1, 2, 3] exactly once + message = "`predicted_ranks` must include each integer" + with assert_raises(ValueError, match=message): + stats.page_trend_test(data=[[1, 2, 3], [1, 2, 3]], + predicted_ranks=[0, 1, 2]) + with assert_raises(ValueError, match=message): + stats.page_trend_test(data=[[1, 2, 3], [1, 2, 3]], + predicted_ranks=[1.1, 2, 3]) + with assert_raises(ValueError, match=message): + stats.page_trend_test(data=[[1, 2, 3], [1, 2, 3]], + predicted_ranks=[1, 2, 3, 3]) + with assert_raises(ValueError, match=message): + stats.page_trend_test(data=[[1, 2, 3], [1, 2, 3]], + predicted_ranks="invalid") + + # test improperly ranked data + with assert_raises(ValueError, match="`data` is not properly ranked"): + stats.page_trend_test([[0, 2, 3], [1, 2, 3]], True) + with assert_raises(ValueError, match="`data` is not properly ranked"): + stats.page_trend_test([[1, 2, 3], [1, 2, 4]], True) + + # various + with assert_raises(ValueError, match="`data` contains NaNs"): + stats.page_trend_test([[1, 2, 3], [1, 2, np.nan]], + ranked=False) + with assert_raises(ValueError, match="`method` must be in"): + stats.page_trend_test(data=[[1, 2, 3], [1, 2, 3]], + method="ekki") + with assert_raises(TypeError, match="`ranked` must be boolean."): + stats.page_trend_test(data=[[1, 2, 3], [1, 2, 3]], + ranked="ekki") + + +@make_xp_test_case(stats.expectile) +class TestExpectile: + @pytest.mark.parametrize('dtype', [None, 'float32', 'float64']) + def test_same_as_mean(self, dtype, xp): + dtype = dtype if dtype is None else getattr(xp, dtype) + rng = np.random.default_rng(42) + + x = xp.asarray(rng.random(size=20).tolist(), dtype=dtype) + xp_assert_close(stats.expectile(x, 0.5), xp.mean(x)) + + w = xp.asarray(rng.random(size=20).tolist(), dtype=dtype) + xp_assert_close(stats.expectile(x, 0.5, weights=w), _xp_mean(x, weights=w)) + + def test_minimum(self, xp): + rng = np.random.default_rng(42) + x = xp.asarray(rng.random(size=20)) + xp_assert_close(stats.expectile(x, alpha=0), xp.min(x)) + + def test_maximum(self, xp): + rng = np.random.default_rng(42) + x = xp.asarray(rng.random(size=20)) + xp_assert_close(stats.expectile(x, alpha=1), xp.max(x)) + + def test_weights(self, xp): + # expectile should minimize `fun` defined below; see + # F. Sobotka and T. Kneib, "Geoadditive expectile regression", + # Computational Statistics and Data Analysis 56 (2012) 755-767 + # https://doi.org/10.1016/j.csda.2010.11.015 + rng = np.random.default_rng(1856392524598679138) + + def fun(u, a, alpha, weights): + w = np.full_like(a, fill_value=alpha) + w[a <= u] = 1 - alpha + return np.sum(w * weights * (a - u)**2) + + def expectile2(a, alpha, weights): + bracket = np.min(a), np.max(a) + return optimize.minimize_scalar(fun, bracket=bracket, + args=(a, alpha, weights)).x + + n = 10 + a = rng.random(n) + alpha = rng.random() + weights = rng.random(n) + + res = stats.expectile(xp.asarray(a), alpha, weights=xp.asarray(weights)) + ref = expectile2(a, alpha, weights) + xp_assert_close(res, xp.asarray(ref)) + + @pytest.mark.parametrize( + "alpha", [0.2, 0.5 - 1e-12, 0.5, 0.5 + 1e-12, 0.8] + ) + @pytest.mark.parametrize("n", [20, 2000]) + def test_expectile_properties(self, alpha, n, xp): + """ + See Section 6 of + I. Steinwart, C. Pasin, R.C. Williamson & S. Zhang (2014). + "Elicitation and Identification of Properties". COLT. + http://proceedings.mlr.press/v35/steinwart14.html + + and + + Propositions 5, 6, 7 of + F. Bellini, B. Klar, and A. Müller and E. Rosazza Gianin (2013). + "Generalized Quantiles as Risk Measures" + http://doi.org/10.2139/ssrn.2225751 + """ + rng = np.random.default_rng(42) + x = xp.asarray(rng.normal(size=n)) + + # 0. definite / constancy + # Let T(X) denote the expectile of rv X ~ F. + # T(c) = c for constant c + for c in [-5., 0., 0.5]: + xp_assert_close( + stats.expectile(xp.full(shape=n, fill_value=c), alpha=alpha), + xp.asarray(c) + ) + + # 1. translation equivariance + # T(X + c) = T(X) + c + c = xp.asarray(rng.exponential()) + xp_assert_close( + stats.expectile(x + c, alpha=alpha), + stats.expectile(x, alpha=alpha) + c, + ) + xp_assert_close( + stats.expectile(x - c, alpha=alpha), + stats.expectile(x, alpha=alpha) - c, + ) + + # 2. positively homogeneity + # T(cX) = c * T(X) for c > 0 + xp_assert_close( + stats.expectile(c * x, alpha=alpha), + c * stats.expectile(x, alpha=alpha), + ) + + # 3. subadditivity + # Note that subadditivity holds for alpha >= 0.5. + # T(X + Y) <= T(X) + T(Y) + # For alpha = 0.5, i.e. the mean, strict equality holds. + # For alpha < 0.5, one can use property 6. to show + # T(X + Y) >= T(X) + T(Y) + y = xp.asarray(rng.logistic(size=n, loc=10)) # different distribution than x + if alpha == 0.5: + def assert_op(a, b): + xp_assert_close(a, b) + + elif alpha > 0.5: + def assert_op(a, b): + assert a < b + + else: + def assert_op(a, b): + assert a > b + + assert_op( + stats.expectile(x + y, alpha=alpha), + stats.expectile(x, alpha=alpha) + + stats.expectile(y, alpha=alpha) + ) + + # 4. monotonicity + # This holds for first order stochastic dominance X: + # X >= Y whenever P(X <= x) < P(Y <= x) + # T(X) <= T(Y) whenever X <= Y + y = xp.asarray(rng.normal(size=n, loc=5)) + assert ( + stats.expectile(x, alpha=alpha) <= stats.expectile(y, alpha=alpha) + ) + + # 5. convexity for alpha > 0.5, concavity for alpha < 0.5 + # convexity is + # T((1 - c) X + c Y) <= (1 - c) T(X) + c T(Y) for 0 <= c <= 1 + y = xp.asarray(rng.logistic(size=n, loc=10)) + for c in [0.1, 0.5, 0.8]: + assert_op( + stats.expectile((1-c)*x + c*y, alpha=alpha), + (1-c) * stats.expectile(x, alpha=alpha) + + c * stats.expectile(y, alpha=alpha) + ) + + # 6. negative argument + # T_{alpha}(-X) = -T_{1-alpha}(X) + xp_assert_close( + stats.expectile(-x, alpha=alpha), + -stats.expectile(x, alpha=1-alpha), + ) + + @pytest.mark.parametrize("n", [20, 2000]) + def test_monotonicity_in_alpha(self, n, xp): + rng = np.random.default_rng(42) + x = xp.asarray(rng.pareto(a=2, size=n)) + e_list = [] + alpha_seq = np.logspace(-15, np.log10(0.5), 100) + # sorted list of unique alpha values in interval (0, 1) + for alpha in np.r_[0, alpha_seq, 1 - alpha_seq[:-1:-1], 1]: + e_list.append(stats.expectile(x, alpha=float(alpha))) + assert xp.all(xp.diff(xp.asarray(e_list)) > 0) + + +@make_xp_test_case(stats.lmoment) +class TestLMoment: + # data from https://github.com/scipy/scipy/issues/19460 + data = [0.87, 0.87, 1.29, 1.5, 1.7, 0.66, 1.5, 0.5, 1., 1.25, 2.3, + 1.03, 2.85, 0.68, 1.74, 1.94, 0.63, 2.04, 1.2, 0.64, 2.05, 0.97, + 2.81, 1.02, 2.76, 0.86, 1.36, 1.29, 1.68, 0.72, 1.67, 1.15, 3.26, + 0.93, 0.83, 0.91, 0.92, 2.32, 1.12, 3.21, 1.23, 1.22, 1.29, 2.08, + 0.64, 2.83, 2.68, 1.77, 0.69, 1.69, 0.7, 1.83, 2.25, 1.23, 1.17, + 0.94, 1.22, 0.76, 0.69, 0.48, 1.04, 2.49, 1.38, 1.57, 1.79, 1.59, + 1.3, 1.54, 1.07, 1.03, 0.76, 2.35, 2.05, 2.02, 2.36, 1.59, 0.97, + 1.63, 1.66, 0.94, 1.45, 1.26, 1.25, 0.68, 2.96, 0.8, 1.16, 0.82, + 0.64, 0.87, 1.33, 1.28, 1.26, 1.19, 1.24, 1.12, 1.45, 1.03, 1.37, + 1.4, 1.35, 1.28, 1.04, 1.31, 0.87, 0.96, 2.55, 1.72, 1.05, 1.15, + 1.73, 1.03, 1.53, 2.41, 1.36, 2.08, 0.92, 0.73, 1.56, 1.94, 0.78] + + not_integers = [1.5, [1, 2, 3.5], math.nan, math.inf] + + def test_dtype_iv(self, xp): + message = '`sample` must be an array of real numbers.' + with pytest.raises(ValueError, match=message): + stats.lmoment(xp.asarray(self.data, dtype=xp.complex128)) + + @skip_xp_invalid_arg + def test_dtype_iv_non_numeric(self): + message = '`sample` must be an array of real numbers.' + with pytest.raises(ValueError, match=message): + stats.lmoment(np.array(self.data, dtype=object)) + + @pytest.mark.parametrize('order', not_integers + [0, -1, [], [[1, 2, 3]]]) + def test_order_iv(self, order, xp): + if is_jax(xp) and np.isscalar(order) and order < 1: + pytest.skip("Order is treated as an array; can't do value-based IV") + message = '`order` must be a scalar or a non-empty...' + with pytest.raises(ValueError, match=message): + stats.lmoment(xp.asarray(self.data), order=order) + + @pytest.mark.parametrize('axis', not_integers) + def test_axis_iv(self, axis, xp): + message = '`axis` must be an integer' + with pytest.raises(ValueError, match=message): + stats.lmoment(xp.asarray(self.data), axis=axis) + + @pytest.mark.parametrize('sorted', not_integers) + def test_sorted_iv(self, sorted, xp): + message = '`sorted` must be True or False.' + with pytest.raises(ValueError, match=message): + stats.lmoment(xp.asarray(self.data), sorted=sorted) + + @pytest.mark.parametrize('standardize', not_integers) + def test_standardize_iv(self, standardize, xp): + message = '`standardize` must be True or False.' + with pytest.raises(ValueError, match=message): + stats.lmoment(xp.asarray(self.data), standardize=standardize) + + @pytest.mark.parametrize('order', [1, 4, [1, 2, 3, 4]]) + @pytest.mark.parametrize('standardize', [False, True]) + @pytest.mark.parametrize('presorted', [False, True]) + def test_lmoment(self, order, standardize, presorted, xp): + # Reference values from R package `lmom` + # options(digits=16) + # library(lmom) + # data= c(0.87, 0.87,..., 1.94, 0.78) + # samlmu(data) + ref = xp.asarray([1.4087603305785130, 0.3415936639118458, + 0.2189964482831403, 0.1328186463415905]) + + if not standardize: + ref = xpx.at(ref)[2:].multiply(ref[1]) + + data = sorted(self.data) if presorted else self.data + data = xp.asarray(data) + + res = stats.lmoment(data, order, standardize=standardize, sorted=presorted) + xp_assert_close(res, ref[xp.asarray(order)-1]) + + def test_dtype(self, xp): + dtype = xp.float32 + sample = xp.asarray(self.data) + res = stats.lmoment(xp.astype(sample, dtype)) + ref = xp.astype(stats.lmoment(sample), dtype) + xp_assert_close(res, ref, rtol=1e-4) + + dtype = xp.int64 + sample = xp.asarray([1, 2, 3, 4, 5]) + res = stats.lmoment(xp.astype(sample, dtype)) + ref = stats.lmoment(xp.astype(sample, xp_default_dtype(xp))) + xp_assert_close(res, ref) + + @pytest.mark.parametrize("axis", [0, 1]) + def test_axis(self, axis, xp): + # nd input is tested extensively in `test_axis_nan_policy`, but only for NumPy + rng = np.random.default_rng(234923498149931248151) + x = rng.random(size=(10, 11)) + res = stats.lmoment(xp.asarray(x), axis=axis) + ref = xp.asarray(stats.lmoment(x, axis=axis)) + xp_assert_close(res, ref) + + +@make_xp_test_case(_xp_mean) +class TestXP_Mean: + @pytest.mark.parametrize('axis', [None, 1, -1, (-2, 2)]) + @pytest.mark.parametrize('weights', [None, True]) + @pytest.mark.parametrize('keepdims', [False, True]) + def test_xp_mean_basic(self, xp, axis, weights, keepdims): + rng = np.random.default_rng(90359458245906) + x = rng.random((3, 4, 5)) + x_xp = xp.asarray(x) + w = w_xp = None + + if weights: + w = rng.random((1, 5)) + w_xp = xp.asarray(w) + x, w = np.broadcast_arrays(x, w) + + res = _xp_mean(x_xp, weights=w_xp, axis=axis, keepdims=keepdims) + ref = np.average(x, weights=w, axis=axis, keepdims=keepdims) + + xp_assert_close(res, xp.asarray(ref)) + + def test_non_broadcastable(self, xp): + # non-broadcastable x and weights + x, w = xp.arange(10.), xp.zeros(5) + message = "Array shapes are incompatible for broadcasting." + with pytest.raises(ValueError, match=message): + _xp_mean(x, weights=w) + + @pytest.mark.filterwarnings("ignore:divide by zero encountered:RuntimeWarning:dask") + @pytest.mark.filterwarnings("ignore:invalid value encountered:RuntimeWarning:dask") + def test_special_cases(self, xp): + # weights sum to zero + weights = xp.asarray([-1., 0., 1.]) + + res = _xp_mean(xp.asarray([1., 1., 1.]), weights=weights) + xp_assert_close(res, xp.asarray(xp.nan)) + + res = _xp_mean(xp.asarray([2., 1., 1.]), weights=weights) + xp_assert_close(res, xp.asarray(-np.inf)) + + res = _xp_mean(xp.asarray([1., 1., 2.]), weights=weights) + xp_assert_close(res, xp.asarray(np.inf)) + + @pytest.mark.filterwarnings( + "ignore:invalid value encountered:RuntimeWarning" + ) # for dask + def test_nan_policy(self, xp): + x = xp.arange(10.) + mask = (x == 3) + x = xp.where(mask, xp.nan, x) + + # nan_policy='raise' raises an error + if is_lazy_array(x): + with pytest.raises(TypeError, match='not supported for lazy arrays'): + _xp_mean(x, nan_policy='raise') + else: + with pytest.raises(ValueError, match='The input contains nan values'): + _xp_mean(x, nan_policy='raise') + + # `nan_policy='propagate'` is the default, and the result is NaN + res1 = _xp_mean(x) + res2 = _xp_mean(x, nan_policy='propagate') + ref = xp.asarray(xp.nan) + xp_assert_equal(res1, ref) + xp_assert_equal(res2, ref) + + # `nan_policy='omit'` omits NaNs in `x` + res = _xp_mean(x, nan_policy='omit') + ref = xp.mean(x[~mask]) + xp_assert_close(res, ref) + + # `nan_policy='omit'` omits NaNs in `weights`, too + weights = xp.ones(10) + weights = xp.where(mask, xp.nan, weights) + res = _xp_mean(xp.arange(10.), weights=weights, nan_policy='omit') + ref = xp.mean(x[~mask]) + xp_assert_close(res, ref) + + @pytest.mark.filterwarnings("ignore:invalid value encountered:RuntimeWarning") + def test_nan_policy_warns(self, xp): + x = xp.arange(10.) + x = xp.where(x == 3, xp.nan, x) + + # Check for warning if omitting NaNs causes empty slice + message = 'After omitting NaNs...' + with eager_warns(RuntimeWarning, match=message, xp=xp): + res = _xp_mean(x * np.nan, nan_policy='omit') + ref = xp.asarray(xp.nan) + xp_assert_equal(res, ref) + + def test_empty(self, xp): + message = 'One or more sample arguments is too small...' + with pytest.warns(SmallSampleWarning, match=message): + res = _xp_mean(xp.asarray([])) + ref = xp.asarray(xp.nan) + xp_assert_equal(res, ref) + + message = "All axis-slices of one or more sample arguments..." + with pytest.warns(SmallSampleWarning, match=message): + res = _xp_mean(xp.asarray([[]]), axis=1) + ref = xp.asarray([xp.nan]) + xp_assert_equal(res, ref) + + res = _xp_mean(xp.asarray([[]]), axis=0) + ref = xp.asarray([]) + xp_assert_equal(res, ref) + + @pytest.mark.filterwarnings( + "ignore:overflow encountered in reduce:RuntimeWarning" + ) # for dask + def test_dtype(self, xp): + max = xp.finfo(xp.float32).max + x_np = np.asarray([max, max], dtype=np.float32) + x_xp = xp.asarray(x_np) + + # Overflow occurs for float32 input + with np.errstate(over='ignore'): + res = _xp_mean(x_xp) + ref = np.mean(x_np) + np.testing.assert_equal(ref, np.inf) + xp_assert_close(res, xp.asarray(ref)) + + # correct result is returned if `float64` is used + res = _xp_mean(x_xp, dtype=xp.float64) + ref = xp.asarray(np.mean(np.asarray(x_np, dtype=np.float64))) + xp_assert_close(res, ref) + + def test_integer(self, xp): + # integer inputs are converted to the appropriate float + x = xp.arange(10) + y = xp.arange(10.) + xp_assert_equal(_xp_mean(x), _xp_mean(y)) + xp_assert_equal(_xp_mean(y, weights=x), _xp_mean(y, weights=y)) + + def test_complex_gh22404(self, xp): + rng = np.random.default_rng(90359458245906) + x, y, wx, wy = rng.random((4, 20)) + res = _xp_mean(xp.asarray(x + y*1j), weights=xp.asarray(wx + wy*1j)) + ref = np.average(x + y*1j, weights=wx + wy*1j) + xp_assert_close(res, xp.asarray(ref)) + + +@make_xp_test_case(_xp_var) +class TestXP_Var: + @pytest.mark.parametrize('axis', [None, 1, -1, (-2, 2)]) + @pytest.mark.parametrize('keepdims', [False, True]) + @pytest.mark.parametrize('correction', [0, 1]) + @pytest.mark.parametrize('nan_policy', ['propagate', 'omit']) + def test_xp_var_basic(self, xp, axis, keepdims, correction, nan_policy): + rng = np.random.default_rng(90359458245906) + x = rng.random((3, 4, 5)) + var_ref = np.var + + if nan_policy == 'omit': + nan_mask = rng.random(size=x.shape) > 0.5 + x[nan_mask] = np.nan + var_ref = np.nanvar + + x_xp = xp.asarray(x) + + res = _xp_var(x_xp, axis=axis, keepdims=keepdims, correction=correction, + nan_policy=nan_policy) + + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", "Degrees of freedom <= 0 for slice", RuntimeWarning) + ref = var_ref(x, axis=axis, keepdims=keepdims, ddof=correction) + + xp_assert_close(res, xp.asarray(ref)) + + def test_special_cases(self, xp): + # correction too big + res = _xp_var(xp.asarray([1., 2.]), correction=3) + xp_assert_close(res, xp.asarray(xp.nan)) + + def test_nan_policy(self, xp): + x = xp.arange(10.) + mask = (x == 3) + x = xp.where(mask, xp.nan, x) + + # `nan_policy='propagate'` is the default, and the result is NaN + res1 = _xp_var(x) + res2 = _xp_var(x, nan_policy='propagate') + ref = xp.asarray(xp.nan) + xp_assert_equal(res1, ref) + xp_assert_equal(res2, ref) + + # `nan_policy='omit'` omits NaNs in `x` + res = _xp_var(x, nan_policy='omit') + ref = xp.var(x[~mask]) + xp_assert_close(res, ref) + + @pytest.mark.filterwarnings("ignore:invalid value encountered:RuntimeWarning") + def test_nan_policy_warns(self, xp): + x = xp.arange(10.) + x = xp.where(x == 3, xp.nan, x) + + # Check for warning if omitting NaNs causes empty slice + message = 'After omitting NaNs...' + with eager_warns(RuntimeWarning, match=message, xp=xp): + res = _xp_var(x * np.nan, nan_policy='omit') + ref = xp.asarray(xp.nan) + xp_assert_equal(res, ref) + + @skip_xp_backends(eager_only=True) + def test_nan_policy_raise(self, xp): + # nan_policy='raise' raises an error when NaNs are present + message = 'The input contains nan values' + with pytest.raises(ValueError, match=message): + _xp_var(xp.asarray([1, 2, xp.nan]), nan_policy='raise') + + def test_empty(self, xp): + message = 'One or more sample arguments is too small...' + with pytest.warns(SmallSampleWarning, match=message): + res = _xp_var(xp.asarray([])) + ref = xp.asarray(xp.nan) + xp_assert_equal(res, ref) + + message = "All axis-slices of one or more sample arguments..." + with pytest.warns(SmallSampleWarning, match=message): + res = _xp_var(xp.asarray([[]]), axis=1) + ref = xp.asarray([xp.nan]) + xp_assert_equal(res, ref) + + res = _xp_var(xp.asarray([[]]), axis=0) + ref = xp.asarray([]) + xp_assert_equal(res, ref) + + @pytest.mark.filterwarnings( + "ignore:overflow encountered in reduce:RuntimeWarning" + ) # Overflow occurs for float32 input + def test_dtype(self, xp): + max = xp.finfo(xp.float32).max + x_np = np.asarray([max, max/2], dtype=np.float32) + x_xp = xp.asarray(x_np) + + res = _xp_var(x_xp) + ref = np.var(x_np) + np.testing.assert_equal(ref, np.inf) + xp_assert_close(res, xp.asarray(ref)) + + # correct result is returned if `float64` is used + res = _xp_var(x_xp, dtype=xp.float64) + ref = xp.asarray(np.var(np.asarray(x_np, dtype=np.float64))) + xp_assert_close(res, ref) + + def test_integer(self, xp): + # integer inputs are converted to the appropriate float + x = xp.arange(10) + y = xp.arange(10.) + xp_assert_equal(_xp_var(x), _xp_var(y)) + + def test_complex_gh22404(self, xp): + rng = np.random.default_rng(90359458245906) + x, y = rng.random((2, 20)) + res = _xp_var(xp.asarray(x + y*1j)) + ref = np.var(x + y*1j) + xp_assert_close(res, xp.asarray(ref), check_dtype=False) + + +@pytest.mark.uses_xp_capabilities(False, reason="private") +def test_chk_asarray(xp): + rng = np.random.default_rng(2348923425434) + x0 = rng.random(size=(2, 3, 4)) + x = xp.asarray(x0) + + axis = 1 + x_out, axis_out = _chk_asarray(x, axis=axis, xp=xp) + xp_assert_equal(x_out, xp.asarray(x0)) + assert_equal(axis_out, axis) + + axis = None + x_out, axis_out = _chk_asarray(x, axis=axis, xp=xp) + xp_assert_equal(x_out, xp.asarray(x0.ravel())) + assert_equal(axis_out, 0) + + axis = 2 + x_out, axis_out = _chk_asarray(x[0, 0, 0], axis=axis, xp=xp) + xp_assert_equal(x_out, xp.asarray(np.atleast_1d(x0[0, 0, 0]))) + assert_equal(axis_out, axis) + + +@skip_xp_backends("jax.numpy", reason="JAX+wilcoxon+'exact' incompatible") +@pytest.mark.parametrize("fun", [ + make_xp_pytest_param(stats.wilcoxon), + make_xp_pytest_param(stats.ks_1samp), + make_xp_pytest_param(stats.ks_2samp), + make_xp_pytest_param(stats.kstest) +]) +def test_rename_mode_method(fun, xp): + rng = np.random.default_rng(23498459284629827814) + x = xp.asarray(rng.random(10)) + y = xp.asarray(rng.random(10)) + + if fun == stats.wilcoxon: + args = (x,) + elif fun == stats.ks_1samp: + args = (x, special.ndtr) + else: + args = (x, y) + + res = fun(*args, method='exact') + res2 = fun(*args, mode='exact') + xp_assert_equal(res.statistic, res2.statistic) + xp_assert_equal(res.pvalue, res2.pvalue) + + err = rf"{fun.__name__}\(\) got multiple values for argument" + with pytest.raises(TypeError, match=err): + fun(*args, method='exact', mode='exact') diff --git a/v12x/pylib/scipy/stats/tests/test_survival.py b/v12x/pylib/scipy/stats/tests/test_survival.py new file mode 100644 index 0000000000000000000000000000000000000000..60a5fef27ae281c028ade8af523bdf1f62f8f664 --- /dev/null +++ b/v12x/pylib/scipy/stats/tests/test_survival.py @@ -0,0 +1,466 @@ +import pytest +import numpy as np +from numpy.testing import assert_equal, assert_allclose +from scipy import stats +from scipy.stats import _survival + + +def _kaplan_meier_reference(times, censored): + # This is a very straightforward implementation of the Kaplan-Meier + # estimator that does almost everything differently from the implementation + # in stats.ecdf. + + # Begin by sorting the raw data. Note that the order of death and loss + # at a given time matters: death happens first. See [2] page 461: + # "These conventions may be paraphrased by saying that deaths recorded as + # of an age t are treated as if they occurred slightly before t, and losses + # recorded as of an age t are treated as occurring slightly after t." + # We implement this by sorting the data first by time, then by `censored`, + # (which is 0 when there is a death and 1 when there is only a loss). + dtype = [('time', float), ('censored', int)] + data = np.array([(t, d) for t, d in zip(times, censored)], dtype=dtype) + data = np.sort(data, order=('time', 'censored')) + times = data['time'] + died = np.logical_not(data['censored']) + + m = times.size + n = np.arange(m, 0, -1) # number at risk + sf = np.cumprod((n - died) / n) + + # Find the indices of the *last* occurrence of unique times. The + # corresponding entries of `times` and `sf` are what we want. + _, indices = np.unique(times[::-1], return_index=True) + ref_times = times[-indices - 1] + ref_sf = sf[-indices - 1] + return ref_times, ref_sf + + +class TestSurvival: + + @staticmethod + def get_random_sample(rng, n_unique): + # generate random sample + unique_times = rng.random(n_unique) + # convert to `np.int32` to resolve `np.repeat` failure in 32-bit CI + repeats = rng.integers(1, 4, n_unique).astype(np.int32) + times = rng.permuted(np.repeat(unique_times, repeats)) + censored = rng.random(size=times.size) > rng.random() + sample = stats.CensoredData.right_censored(times, censored) + return sample, times, censored + + def test_input_validation(self): + message = '`sample` must be a one-dimensional sequence.' + with pytest.raises(ValueError, match=message): + stats.ecdf([[1]]) + with pytest.raises(ValueError, match=message): + stats.ecdf(1) + + message = '`sample` must not contain nan' + with pytest.raises(ValueError, match=message): + stats.ecdf([np.nan]) + + message = 'Currently, only uncensored and right-censored data...' + with pytest.raises(NotImplementedError, match=message): + stats.ecdf(stats.CensoredData.left_censored([1], censored=[True])) + + message = 'method` must be one of...' + res = stats.ecdf([1, 2, 3]) + with pytest.raises(ValueError, match=message): + res.cdf.confidence_interval(method='ekki-ekki') + with pytest.raises(ValueError, match=message): + res.sf.confidence_interval(method='shrubbery') + + message = 'confidence_level` must be a scalar between 0 and 1' + with pytest.raises(ValueError, match=message): + res.cdf.confidence_interval(-1) + with pytest.raises(ValueError, match=message): + res.sf.confidence_interval([0.5, 0.6]) + + message = 'The confidence interval is undefined at some observations.' + with pytest.warns(RuntimeWarning, match=message): + ci = res.cdf.confidence_interval() + + message = 'Confidence interval bounds do not implement...' + with pytest.raises(NotImplementedError, match=message): + ci.low.confidence_interval() + with pytest.raises(NotImplementedError, match=message): + ci.high.confidence_interval() + + def test_edge_cases(self): + res = stats.ecdf([]) + assert_equal(res.cdf.quantiles, []) + assert_equal(res.cdf.probabilities, []) + + res = stats.ecdf([1]) + assert_equal(res.cdf.quantiles, [1]) + assert_equal(res.cdf.probabilities, [1]) + + def test_unique(self): + # Example with unique observations; `stats.ecdf` ref. [1] page 80 + sample = [6.23, 5.58, 7.06, 6.42, 5.20] + res = stats.ecdf(sample) + ref_x = np.sort(np.unique(sample)) + ref_cdf = np.arange(1, 6) / 5 + ref_sf = 1 - ref_cdf + assert_equal(res.cdf.quantiles, ref_x) + assert_equal(res.cdf.probabilities, ref_cdf) + assert_equal(res.sf.quantiles, ref_x) + assert_equal(res.sf.probabilities, ref_sf) + + def test_nonunique(self): + # Example with non-unique observations; `stats.ecdf` ref. [1] page 82 + sample = [0, 2, 1, 2, 3, 4] + res = stats.ecdf(sample) + ref_x = np.sort(np.unique(sample)) + ref_cdf = np.array([1/6, 2/6, 4/6, 5/6, 1]) + ref_sf = 1 - ref_cdf + assert_equal(res.cdf.quantiles, ref_x) + assert_equal(res.cdf.probabilities, ref_cdf) + assert_equal(res.sf.quantiles, ref_x) + assert_equal(res.sf.probabilities, ref_sf) + + def test_evaluate_methods(self): + # Test CDF and SF `evaluate` methods + rng = np.random.default_rng(1162729143302572461) + sample, _, _ = self.get_random_sample(rng, 15) + res = stats.ecdf(sample) + x = res.cdf.quantiles + xr = x + np.diff(x, append=x[-1]+1)/2 # right shifted points + + assert_equal(res.cdf.evaluate(x), res.cdf.probabilities) + assert_equal(res.cdf.evaluate(xr), res.cdf.probabilities) + assert_equal(res.cdf.evaluate(x[0]-1), 0) # CDF starts at 0 + assert_equal(res.cdf.evaluate([-np.inf, np.inf]), [0, 1]) + + assert_equal(res.sf.evaluate(x), res.sf.probabilities) + assert_equal(res.sf.evaluate(xr), res.sf.probabilities) + assert_equal(res.sf.evaluate(x[0]-1), 1) # SF starts at 1 + assert_equal(res.sf.evaluate([-np.inf, np.inf]), [1, 0]) + + # ref. [1] page 91 + t1 = [37, 43, 47, 56, 60, 62, 71, 77, 80, 81] # times + d1 = [0, 0, 1, 1, 0, 0, 0, 1, 1, 1] # 1 means deaths (not censored) + r1 = [1, 1, 0.875, 0.75, 0.75, 0.75, 0.75, 0.5, 0.25, 0] # reference SF + + # https://sphweb.bumc.bu.edu/otlt/mph-modules/bs/bs704_survival/BS704_Survival5.html + t2 = [8, 12, 26, 14, 21, 27, 8, 32, 20, 40] + d2 = [1, 1, 1, 1, 1, 1, 0, 0, 0, 0] + r2 = [0.9, 0.788, 0.675, 0.675, 0.54, 0.405, 0.27, 0.27, 0.27] + t3 = [33, 28, 41, 48, 48, 25, 37, 48, 25, 43] + d3 = [1, 1, 1, 0, 0, 0, 0, 0, 0, 0] + r3 = [1, 0.875, 0.75, 0.75, 0.6, 0.6, 0.6] + + # https://sphweb.bumc.bu.edu/otlt/mph-modules/bs/bs704_survival/bs704_survival4.html + t4 = [24, 3, 11, 19, 24, 13, 14, 2, 18, 17, + 24, 21, 12, 1, 10, 23, 6, 5, 9, 17] + d4 = [0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 1, 0, 1, 0, 1, 0, 1] + r4 = [0.95, 0.95, 0.897, 0.844, 0.844, 0.844, 0.844, 0.844, 0.844, + 0.844, 0.76, 0.676, 0.676, 0.676, 0.676, 0.507, 0.507] + + # https://www.real-statistics.com/survival-analysis/kaplan-meier-procedure/confidence-interval-for-the-survival-function/ + t5 = [3, 5, 8, 10, 5, 5, 8, 12, 15, 14, 2, 11, 10, 9, 12, 5, 8, 11] + d5 = [1, 1, 1, 0, 1, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1, 1, 1] + r5 = [0.944, 0.889, 0.722, 0.542, 0.542, 0.542, 0.361, 0.181, 0.181, 0.181] + + @pytest.mark.parametrize("case", [(t1, d1, r1), (t2, d2, r2), (t3, d3, r3), + (t4, d4, r4), (t5, d5, r5)]) + def test_right_censored_against_examples(self, case): + # test `ecdf` against other implementations on example problems + times, died, ref = case + sample = stats.CensoredData.right_censored(times, np.logical_not(died)) + res = stats.ecdf(sample) + assert_allclose(res.sf.probabilities, ref, atol=1e-3) + assert_equal(res.sf.quantiles, np.sort(np.unique(times))) + + # test reference implementation against other implementations + res = _kaplan_meier_reference(times, np.logical_not(died)) + assert_equal(res[0], np.sort(np.unique(times))) + assert_allclose(res[1], ref, atol=1e-3) + + @pytest.mark.parametrize('seed', [182746786639392128, 737379171436494115, + 576033618403180168, 308115465002673650]) + def test_right_censored_against_reference_implementation(self, seed): + # test `ecdf` against reference implementation on random problems + rng = np.random.default_rng(seed) + n_unique = rng.integers(10, 100) + sample, times, censored = self.get_random_sample(rng, n_unique) + res = stats.ecdf(sample) + ref = _kaplan_meier_reference(times, censored) + assert_allclose(res.sf.quantiles, ref[0]) + assert_allclose(res.sf.probabilities, ref[1]) + + # If all observations are uncensored, the KM estimate should match + # the usual estimate for uncensored data + sample = stats.CensoredData(uncensored=times) + res = _survival._ecdf_right_censored(sample) # force Kaplan-Meier + ref = stats.ecdf(times) + assert_equal(res[0], ref.sf.quantiles) + assert_allclose(res[1], ref.cdf.probabilities, rtol=1e-14) + assert_allclose(res[2], ref.sf.probabilities, rtol=1e-14) + + def test_right_censored_ci(self): + # test "greenwood" confidence interval against example 4 (URL above). + times, died = self.t4, self.d4 + sample = stats.CensoredData.right_censored(times, np.logical_not(died)) + res = stats.ecdf(sample) + ref_allowance = [0.096, 0.096, 0.135, 0.162, 0.162, 0.162, 0.162, + 0.162, 0.162, 0.162, 0.214, 0.246, 0.246, 0.246, + 0.246, 0.341, 0.341] + + sf_ci = res.sf.confidence_interval() + cdf_ci = res.cdf.confidence_interval() + allowance = res.sf.probabilities - sf_ci.low.probabilities + + assert_allclose(allowance, ref_allowance, atol=1e-3) + assert_allclose(sf_ci.low.probabilities, + np.clip(res.sf.probabilities - allowance, 0, 1)) + assert_allclose(sf_ci.high.probabilities, + np.clip(res.sf.probabilities + allowance, 0, 1)) + assert_allclose(cdf_ci.low.probabilities, + np.clip(res.cdf.probabilities - allowance, 0, 1)) + assert_allclose(cdf_ci.high.probabilities, + np.clip(res.cdf.probabilities + allowance, 0, 1)) + + # test "log-log" confidence interval against Mathematica + # e = {24, 3, 11, 19, 24, 13, 14, 2, 18, 17, 24, 21, 12, 1, 10, 23, 6, 5, + # 9, 17} + # ci = {1, 0, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 0, 1, 0, 1, 0, 1, 0} + # R = EventData[e, ci] + # S = SurvivalModelFit[R] + # S["PointwiseIntervals", ConfidenceLevel->0.95, + # ConfidenceTransform->"LogLog"] + + ref_low = [0.694743, 0.694743, 0.647529, 0.591142, 0.591142, 0.591142, + 0.591142, 0.591142, 0.591142, 0.591142, 0.464605, 0.370359, + 0.370359, 0.370359, 0.370359, 0.160489, 0.160489] + ref_high = [0.992802, 0.992802, 0.973299, 0.947073, 0.947073, 0.947073, + 0.947073, 0.947073, 0.947073, 0.947073, 0.906422, 0.856521, + 0.856521, 0.856521, 0.856521, 0.776724, 0.776724] + sf_ci = res.sf.confidence_interval(method='log-log') + assert_allclose(sf_ci.low.probabilities, ref_low, atol=1e-6) + assert_allclose(sf_ci.high.probabilities, ref_high, atol=1e-6) + + def test_right_censored_ci_example_5(self): + # test "exponential greenwood" confidence interval against example 5 + times, died = self.t5, self.d5 + sample = stats.CensoredData.right_censored(times, np.logical_not(died)) + res = stats.ecdf(sample) + lower = np.array([0.66639, 0.624174, 0.456179, 0.287822, 0.287822, + 0.287822, 0.128489, 0.030957, 0.030957, 0.030957]) + upper = np.array([0.991983, 0.970995, 0.87378, 0.739467, 0.739467, + 0.739467, 0.603133, 0.430365, 0.430365, 0.430365]) + + sf_ci = res.sf.confidence_interval(method='log-log') + cdf_ci = res.cdf.confidence_interval(method='log-log') + + assert_allclose(sf_ci.low.probabilities, lower, atol=1e-5) + assert_allclose(sf_ci.high.probabilities, upper, atol=1e-5) + assert_allclose(cdf_ci.low.probabilities, 1-upper, atol=1e-5) + assert_allclose(cdf_ci.high.probabilities, 1-lower, atol=1e-5) + + # Test against R's `survival` library `survfit` function, 90%CI + # library(survival) + # options(digits=16) + # time = c(3, 5, 8, 10, 5, 5, 8, 12, 15, 14, 2, 11, 10, 9, 12, 5, 8, 11) + # status = c(1, 1, 1, 0, 1, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1, 1, 1) + # res = survfit(Surv(time, status) + # ~1, conf.type = "log-log", conf.int = 0.90) + # res$time; res$lower; res$upper + low = [0.74366748406861172, 0.68582332289196246, 0.50596835651480121, + 0.32913131413336727, 0.32913131413336727, 0.32913131413336727, + 0.15986912028781664, 0.04499539918147757, 0.04499539918147757, + 0.04499539918147757] + high = [0.9890291867238429, 0.9638835422144144, 0.8560366823086629, + 0.7130167643978450, 0.7130167643978450, 0.7130167643978450, + 0.5678602982997164, 0.3887616766886558, 0.3887616766886558, + 0.3887616766886558] + sf_ci = res.sf.confidence_interval(method='log-log', + confidence_level=0.9) + assert_allclose(sf_ci.low.probabilities, low) + assert_allclose(sf_ci.high.probabilities, high) + + # And with conf.type = "plain" + low = [0.8556383113628162, 0.7670478794850761, 0.5485720663578469, + 0.3441515412527123, 0.3441515412527123, 0.3441515412527123, + 0.1449184105424544, 0., 0., 0.] + high = [1., 1., 0.8958723780865975, 0.7391817920806210, + 0.7391817920806210, 0.7391817920806210, 0.5773038116797676, + 0.3642270254596720, 0.3642270254596720, 0.3642270254596720] + sf_ci = res.sf.confidence_interval(confidence_level=0.9) + assert_allclose(sf_ci.low.probabilities, low) + assert_allclose(sf_ci.high.probabilities, high) + + def test_right_censored_ci_nans(self): + # test `ecdf` confidence interval on a problem that results in NaNs + times, died = self.t1, self.d1 + sample = stats.CensoredData.right_censored(times, np.logical_not(died)) + res = stats.ecdf(sample) + + # Reference values generated with Matlab + # format long + # t = [37 43 47 56 60 62 71 77 80 81]; + # d = [0 0 1 1 0 0 0 1 1 1]; + # censored = ~d1; + # [f, x, flo, fup] = ecdf(t, 'Censoring', censored, 'Alpha', 0.05); + x = [37, 47, 56, 77, 80, 81] + flo = [np.nan, 0, 0, 0.052701464070711, 0.337611126231790, np.nan] + fup = [np.nan, 0.35417230377, 0.5500569798, 0.9472985359, 1.0, np.nan] + i = np.searchsorted(res.cdf.quantiles, x) + + message = "The confidence interval is undefined at some observations" + with pytest.warns(RuntimeWarning, match=message): + ci = res.cdf.confidence_interval() + + # Matlab gives NaN as the first element of the CIs. Mathematica agrees, + # but R's survfit does not. It makes some sense, but it's not what the + # formula gives, so skip that element. + assert_allclose(ci.low.probabilities[i][1:], flo[1:]) + assert_allclose(ci.high.probabilities[i][1:], fup[1:]) + + # [f, x, flo, fup] = ecdf(t, 'Censoring', censored, 'Function', + # 'survivor', 'Alpha', 0.05); + flo = [np.nan, 0.64582769623, 0.449943020228, 0.05270146407, 0, np.nan] + fup = [np.nan, 1.0, 1.0, 0.947298535929289, 0.662388873768210, np.nan] + i = np.searchsorted(res.cdf.quantiles, x) + + with pytest.warns(RuntimeWarning, match=message): + ci = res.sf.confidence_interval() + + assert_allclose(ci.low.probabilities[i][1:], flo[1:]) + assert_allclose(ci.high.probabilities[i][1:], fup[1:]) + + # With the same data, R's `survival` library `survfit` function + # doesn't produce the leading NaN + # library(survival) + # options(digits=16) + # time = c(37, 43, 47, 56, 60, 62, 71, 77, 80, 81) + # status = c(0, 0, 1, 1, 0, 0, 0, 1, 1, 1) + # res = survfit(Surv(time, status) + # ~1, conf.type = "plain", conf.int = 0.95) + # res$time + # res$lower + # res$upper + low = [1., 1., 0.64582769623233816, 0.44994302022779326, + 0.44994302022779326, 0.44994302022779326, 0.44994302022779326, + 0.05270146407071086, 0., np.nan] + high = [1., 1., 1., 1., 1., 1., 1., 0.9472985359292891, + 0.6623888737682101, np.nan] + assert_allclose(ci.low.probabilities, low) + assert_allclose(ci.high.probabilities, high) + + # It does with conf.type="log-log", as do we + with pytest.warns(RuntimeWarning, match=message): + ci = res.sf.confidence_interval(method='log-log') + low = [np.nan, np.nan, 0.38700001403202522, 0.31480711370551911, + 0.31480711370551911, 0.31480711370551911, 0.31480711370551911, + 0.08048821148507734, 0.01049958986680601, np.nan] + high = [np.nan, np.nan, 0.9813929658789660, 0.9308983170906275, + 0.9308983170906275, 0.9308983170906275, 0.9308983170906275, + 0.8263946341076415, 0.6558775085110887, np.nan] + assert_allclose(ci.low.probabilities, low) + assert_allclose(ci.high.probabilities, high) + + def test_right_censored_against_uncensored(self): + rng = np.random.default_rng(7463952748044886637) + sample = rng.integers(10, 100, size=1000) + censored = np.zeros_like(sample) + censored[np.argmax(sample)] = True + res = stats.ecdf(sample) + ref = stats.ecdf(stats.CensoredData.right_censored(sample, censored)) + assert_equal(res.sf.quantiles, ref.sf.quantiles) + assert_equal(res.sf._n, ref.sf._n) + assert_equal(res.sf._d[:-1], ref.sf._d[:-1]) # difference @ [-1] + assert_allclose(res.sf._sf[:-1], ref.sf._sf[:-1], rtol=1e-14) + + def test_plot_iv(self): + rng = np.random.default_rng(1769658657308472721) + n_unique = rng.integers(10, 100) + sample, _, _ = self.get_random_sample(rng, n_unique) + res = stats.ecdf(sample) + + try: + import matplotlib.pyplot as plt # noqa: F401 + res.sf.plot() # no other errors occur + except (ModuleNotFoundError, ImportError): + message = r"matplotlib must be installed to use method `plot`." + with pytest.raises(ModuleNotFoundError, match=message): + res.sf.plot() + + +class TestLogRank: + + @pytest.mark.parametrize( + "x, y, statistic, pvalue", + # Results validate with R + # library(survival) + # options(digits=16) + # + # futime_1 <- c(8, 12, 26, 14, 21, 27, 8, 32, 20, 40) + # fustat_1 <- c(1, 1, 1, 1, 1, 1, 0, 0, 0, 0) + # rx_1 <- c(0, 0, 0, 0, 0, 0, 0, 0, 0, 0) + # + # futime_2 <- c(33, 28, 41, 48, 48, 25, 37, 48, 25, 43) + # fustat_2 <- c(1, 1, 1, 0, 0, 0, 0, 0, 0, 0) + # rx_2 <- c(1, 1, 1, 1, 1, 1, 1, 1, 1, 1) + # + # futime <- c(futime_1, futime_2) + # fustat <- c(fustat_1, fustat_2) + # rx <- c(rx_1, rx_2) + # + # survdiff(formula = Surv(futime, fustat) ~ rx) + # + # Also check against another library which handle alternatives + # library(nph) + # logrank.test(futime, fustat, rx, alternative = "two.sided") + # res["test"] + [( + # https://sphweb.bumc.bu.edu/otlt/mph-modules/bs/bs704_survival/BS704_Survival5.html + # uncensored, censored + [[8, 12, 26, 14, 21, 27], [8, 32, 20, 40]], + [[33, 28, 41], [48, 48, 25, 37, 48, 25, 43]], + # chi2, ["two-sided", "less", "greater"] + 6.91598157449, + [0.008542873404, 0.9957285632979385, 0.004271436702061537] + ), + ( + # https://sphweb.bumc.bu.edu/otlt/mph-modules/bs/bs704_survival/BS704_Survival5.html + [[19, 6, 5, 4], [20, 19, 17, 14]], + [[16, 21, 7], [21, 15, 18, 18, 5]], + 0.835004855038, + [0.3608293039, 0.8195853480676912, 0.1804146519323088] + ), + ( + # Bland, Altman, "The logrank test", BMJ, 2004 + # https://www.bmj.com/content/328/7447/1073.short + [[6, 13, 21, 30, 37, 38, 49, 50, 63, 79, 86, 98, 202, 219], + [31, 47, 80, 82, 82, 149]], + [[10, 10, 12, 13, 14, 15, 16, 17, 18, 20, 24, 24, 25, 28, 30, + 33, 35, 37, 40, 40, 46, 48, 76, 81, 82, 91, 112, 181], + [34, 40, 70]], + 7.49659416854, + [0.006181578637, 0.003090789318730882, 0.9969092106812691] + )] + ) + def test_log_rank(self, x, y, statistic, pvalue): + x = stats.CensoredData(uncensored=x[0], right=x[1]) + y = stats.CensoredData(uncensored=y[0], right=y[1]) + + for i, alternative in enumerate(["two-sided", "less", "greater"]): + res = stats.logrank(x=x, y=y, alternative=alternative) + + # we return z and use the normal distribution while other framework + # return z**2. The p-value are directly comparable, but we have to + # square the statistic + assert_allclose(res.statistic**2, statistic, atol=1e-10) + assert_allclose(res.pvalue, pvalue[i], atol=1e-10) + + def test_raises(self): + sample = stats.CensoredData([1, 2]) + + msg = r"`y` must be" + with pytest.raises(ValueError, match=msg): + stats.logrank(x=sample, y=[[1, 2]]) + + msg = r"`x` must be" + with pytest.raises(ValueError, match=msg): + stats.logrank(x=[[1, 2]], y=sample) diff --git a/v12x/pylib/scipy/stats/tests/test_tukeylambda_stats.py b/v12x/pylib/scipy/stats/tests/test_tukeylambda_stats.py new file mode 100644 index 0000000000000000000000000000000000000000..8041658599a6ed5b2cd70def1a92bffe0d851792 --- /dev/null +++ b/v12x/pylib/scipy/stats/tests/test_tukeylambda_stats.py @@ -0,0 +1,85 @@ +import numpy as np +from numpy.testing import assert_allclose, assert_equal + +from scipy.stats._tukeylambda_stats import (tukeylambda_variance, + tukeylambda_kurtosis) + + +def test_tukeylambda_stats_known_exact(): + """Compare results with some known exact formulas.""" + # Some exact values of the Tukey Lambda variance and kurtosis: + # lambda var kurtosis + # 0 pi**2/3 6/5 (logistic distribution) + # 0.5 4 - pi (5/3 - pi/2)/(pi/4 - 1)**2 - 3 + # 1 1/3 -6/5 (uniform distribution on (-1,1)) + # 2 1/12 -6/5 (uniform distribution on (-1/2, 1/2)) + + # lambda = 0 + var = tukeylambda_variance(0) + assert_allclose(var, np.pi**2 / 3, atol=1e-12) + kurt = tukeylambda_kurtosis(0) + assert_allclose(kurt, 1.2, atol=1e-10) + + # lambda = 0.5 + var = tukeylambda_variance(0.5) + assert_allclose(var, 4 - np.pi, atol=1e-12) + kurt = tukeylambda_kurtosis(0.5) + desired = (5./3 - np.pi/2) / (np.pi/4 - 1)**2 - 3 + assert_allclose(kurt, desired, atol=1e-10) + + # lambda = 1 + var = tukeylambda_variance(1) + assert_allclose(var, 1.0 / 3, atol=1e-12) + kurt = tukeylambda_kurtosis(1) + assert_allclose(kurt, -1.2, atol=1e-10) + + # lambda = 2 + var = tukeylambda_variance(2) + assert_allclose(var, 1.0 / 12, atol=1e-12) + kurt = tukeylambda_kurtosis(2) + assert_allclose(kurt, -1.2, atol=1e-10) + + +def test_tukeylambda_stats_mpmath(): + """Compare results with some values that were computed using mpmath.""" + a10 = dict(atol=1e-10, rtol=0) + a12 = dict(atol=1e-12, rtol=0) + data = [ + # lambda variance kurtosis + [-0.1, 4.78050217874253547, 3.78559520346454510], + [-0.0649, 4.16428023599895777, 2.52019675947435718], + [-0.05, 3.93672267890775277, 2.13129793057777277], + [-0.001, 3.30128380390964882, 1.21452460083542988], + [0.001, 3.27850775649572176, 1.18560634779287585], + [0.03125, 2.95927803254615800, 0.804487555161819980], + [0.05, 2.78281053405464501, 0.611604043886644327], + [0.0649, 2.65282386754100551, 0.476834119532774540], + [1.2, 0.242153920578588346, -1.23428047169049726], + [10.0, 0.00095237579757703597, 2.37810697355144933], + [20.0, 0.00012195121951131043, 7.37654321002709531], + ] + + for lam, var_expected, kurt_expected in data: + var = tukeylambda_variance(lam) + assert_allclose(var, var_expected, **a12) + kurt = tukeylambda_kurtosis(lam) + assert_allclose(kurt, kurt_expected, **a10) + + # Test with vector arguments (most of the other tests are for single + # values). + lam, var_expected, kurt_expected = zip(*data) + var = tukeylambda_variance(lam) + assert_allclose(var, var_expected, **a12) + kurt = tukeylambda_kurtosis(lam) + assert_allclose(kurt, kurt_expected, **a10) + + +def test_tukeylambda_stats_invalid(): + """Test values of lambda outside the domains of the functions.""" + lam = [-1.0, -0.5] + var = tukeylambda_variance(lam) + assert_equal(var, np.array([np.nan, np.inf])) + + lam = [-1.0, -0.25] + kurt = tukeylambda_kurtosis(lam) + assert_equal(kurt, np.array([np.nan, np.inf])) diff --git a/v12x/pylib/scipy/stats/tests/test_variation.py b/v12x/pylib/scipy/stats/tests/test_variation.py new file mode 100644 index 0000000000000000000000000000000000000000..2a07ac1bfcdf85e9be609ff7a5c1445f32e22cba --- /dev/null +++ b/v12x/pylib/scipy/stats/tests/test_variation.py @@ -0,0 +1,205 @@ +import math + +import numpy as np +import pytest + +from scipy._lib._util import AxisError +from scipy._lib._array_api import make_xp_test_case, eager_warns, is_numpy +from scipy._lib._array_api_no_0d import xp_assert_equal, xp_assert_close +from scipy._external import array_api_extra as xpx + +from scipy import stats +from scipy.stats import variation +from scipy.stats._axis_nan_policy import (too_small_nd_omit, too_small_nd_not_omit, + SmallSampleWarning) + +lazy_xp_modules = [stats] +skip_xp_backends = pytest.mark.skip_xp_backends + + +@make_xp_test_case(variation) +class TestVariation: + """ + Test class for scipy.stats.variation + """ + + def test_ddof(self, xp): + x = xp.arange(9.0) + xp_assert_close(variation(x, ddof=1), xp.asarray(math.sqrt(60/8)/4)) + + @pytest.mark.parametrize('sgn', [1, -1]) + def test_sign(self, sgn, xp): + x = xp.asarray([1., 2., 3., 4., 5.]) + v = variation(sgn*x) + expected = xp.asarray(sgn*math.sqrt(2)/3) + xp_assert_close(v, expected, rtol=1e-10) + + @skip_xp_backends(np_only=True, reason="test plain python scalar input") + def test_scalar(self, xp): + # A scalar is treated like a 1-d sequence with length 1. + assert variation(4.0) == 0.0 + + @pytest.mark.parametrize('nan_policy, expected', + [('propagate', np.nan), + ('omit', math.sqrt(20/3)/4)]) + @skip_xp_backends(eager_only=True, reason='lazy -> limited nan_policy support') + def test_variation_nan(self, nan_policy, expected, xp): + x = xp.arange(10.) + x[9] = xp.nan + xp_assert_close(variation(x, nan_policy=nan_policy), xp.asarray(expected)) + + @skip_xp_backends(eager_only=True, reason='lazy -> limited nan_policy support') + def test_nan_policy_raise(self, xp): + x = xp.asarray([1.0, 2.0, xp.nan, 3.0]) + with pytest.raises(ValueError, match='input contains nan'): + variation(x, nan_policy='raise') + + @skip_xp_backends("dask.array", reason='needs `_axis_nan_policy`') + def test_bad_nan_policy(self, xp): + with pytest.raises(ValueError, match='must be one of'): + variation(xp.asarray([1, 2, 3]), nan_policy='foobar') + + @skip_xp_backends("dask.array", reason='needs `_axis_nan_policy`') + def test_keepdims(self, xp): + x = xp.reshape(xp.arange(10), (2, 5)) + y = variation(x, axis=1, keepdims=True) + expected = 2**0.5 / xp.asarray([[2.], [7.]]) + xp_assert_close(y, expected) + + @skip_xp_backends("dask.array", reason='needs `_axis_nan_policy`') + @pytest.mark.parametrize('axis, expected', + [(0, np.empty((1, 0)).tolist()), + (1, np.full((5, 1), fill_value=np.nan).tolist())]) + def test_keepdims_size0(self, axis, expected, xp): + x = xp.zeros((5, 0)) + if axis == 1: + with pytest.warns(SmallSampleWarning, match=too_small_nd_not_omit): + y = variation(x, axis=axis, keepdims=True) + else: + y = variation(x, axis=axis, keepdims=True) + xp_assert_equal(y, xp.asarray(expected)) + + @skip_xp_backends("dask.array", reason='needs `_axis_nan_policy`') + @pytest.mark.parametrize('incr, expected_fill', [(0, np.inf), (1, np.nan)]) + def test_keepdims_and_ddof_eq_len_plus_incr(self, incr, expected_fill, xp): + x = xp.asarray([[1, 1, 2, 2], [1, 2, 3, 3]]) + y = variation(x, axis=1, ddof=x.shape[1] + incr, keepdims=True) + xp_assert_equal(y, xp.full((2, 1), fill_value=expected_fill)) + + @skip_xp_backends("dask.array", reason='needs `_axis_nan_policy`') + def test_propagate_nan(self, xp): + # Check that the shape of the result is the same for inputs + # with and without nans, cf gh-5817 + a = xp.reshape(xp.arange(8.), (2, -1)) + a = xpx.at(a)[1, 0].set(xp.nan) + v = variation(a, axis=1, nan_policy="propagate") + xp_assert_close(v, xp.asarray([math.sqrt(5/4)/1.5, xp.nan])) + + def test_axis_none(self, xp): + # Check that `variation` computes the result on the flattened + # input when axis is None. + y = variation(xp.asarray([[0, 1], [2, 3]]), axis=None) + xp_assert_close(y, xp.asarray(math.sqrt(5/4)/1.5)) + + def test_bad_axis(self, xp): + # Check that an invalid axis raises np.exceptions.AxisError. + x = xp.asarray([[1, 2, 3], [4, 5, 6]]) + with pytest.raises((AxisError, IndexError)): + variation(x, axis=10) + + @pytest.mark.filterwarnings("ignore:divide by zero encountered:RuntimeWarning:dask") + def test_mean_zero(self, xp): + # Check that `variation` returns inf for a sequence that is not + # identically zero but whose mean is zero. + x = xp.asarray([10., -3., 1., -4., -4.]) + y = variation(x) + xp_assert_equal(y, xp.asarray(xp.inf)) + + x2 = xp.stack([x, -10.*x]) + y2 = variation(x2, axis=1) + xp_assert_equal(y2, xp.asarray([xp.inf, xp.inf])) + + @pytest.mark.filterwarnings("ignore:invalid value encountered:RuntimeWarning:dask") + @pytest.mark.parametrize('x', [[0.]*5, [1, 2, np.inf, 9]]) + def test_return_nan(self, x, xp): + x = xp.asarray(x) + # Test some cases where `variation` returns nan. + y = variation(x) + xp_assert_equal(y, xp.asarray(xp.nan, dtype=x.dtype)) + + @pytest.mark.filterwarnings('ignore:Invalid value encountered:RuntimeWarning:dask') + @pytest.mark.parametrize('axis, expected', + [(0, []), (1, [np.nan]*3), (None, np.nan)]) + def test_2d_size_zero_with_axis(self, axis, expected, xp): + x = xp.empty((3, 0)) + if axis != 0: + # specific message depends on `axis`, and `SmallSampleWarning` + # is specific enough. + with eager_warns(SmallSampleWarning, xp=xp): + y = variation(x, axis=axis) + else: + y = variation(x, axis=axis) + xp_assert_equal(y, xp.asarray(expected)) + + @pytest.mark.filterwarnings('ignore:divide by zero encountered:RuntimeWarning:dask') + def test_neg_inf(self, xp): + # Edge case that produces -inf: ddof equals the number of non-nan + # values, the values are not constant, and the mean is negative. + x1 = xp.asarray([-3., -5.]) + xp_assert_equal(variation(x1, ddof=2), xp.asarray(-xp.inf)) + + @skip_xp_backends(np_only=True, reason='only NumPy supports ND nan_policy="omit"') + def test_neg_inf_nan(self, xp): + x2 = xp.asarray([[xp.nan, 1, -10, xp.nan], + [-20, -3, xp.nan, xp.nan]]) + xp_assert_equal(variation(x2, axis=1, ddof=2, nan_policy='omit'), + xp.asarray([-xp.inf, -xp.inf])) + + @pytest.mark.parametrize("nan_policy", ['propagate', 'omit']) + def test_combined_edge_cases(self, nan_policy, xp): + if not is_numpy(xp): + pytest.skip("Only NumPy supports ND nan_policy='omit'") + x = xp.asarray([[0, 10, xp.nan, 1], + [0, -5, xp.nan, 2], + [0, -5, xp.nan, 3]]) + if nan_policy == 'omit': + with pytest.warns(SmallSampleWarning, match=too_small_nd_omit): + y = variation(x, axis=0, nan_policy=nan_policy) + else: + y = variation(x, axis=0, nan_policy=nan_policy) + xp_assert_close(y, xp.asarray([xp.nan, xp.inf, xp.nan, math.sqrt(2/3)/2])) + + @skip_xp_backends(np_only=True, reason='only NumPy supports ND nan_policy="omit"') + @pytest.mark.parametrize( + 'ddof, expected', + [(0, [np.sqrt(1/6), np.sqrt(5/8), np.inf, 0, np.nan, 0.0, np.nan]), + (1, [0.5, np.sqrt(5/6), np.inf, 0, np.nan, 0, np.nan]), + (2, [np.sqrt(0.5), np.sqrt(5/4), np.inf, np.nan, np.nan, 0, np.nan])] + ) + def test_more_nan_policy_omit_tests(self, ddof, expected, xp): + # The slightly strange formatting in the follow array is my attempt to + # maintain a clean tabular arrangement of the data while satisfying + # the demands of pycodestyle. Currently, E201 and E241 are not + # disabled by the `noqa` annotation. + nan = xp.nan + x = xp.asarray([[1.0, 2.0, nan, 3.0], + [0.0, 4.0, 3.0, 1.0], + [nan, -.5, 0.5, nan], + [nan, 9.0, 9.0, nan], + [nan, nan, nan, nan], + [3.0, 3.0, 3.0, 3.0], + [0.0, 0.0, 0.0, 0.0]]) + with pytest.warns(SmallSampleWarning, match=too_small_nd_omit): + v = variation(x, axis=1, ddof=ddof, nan_policy='omit') + xp_assert_close(v, expected) + + @skip_xp_backends(eager_only=True, reason='lazy -> limited nan_policy support') + def test_variation_ddof(self, xp): + # test variation with delta degrees of freedom + # regression test for gh-13341 + a = xp.asarray([1., 2., 3., 4., 5.]) + nan_a = xp.asarray([1, 2, 3, xp.nan, 4, 5, xp.nan]) + y = variation(a, ddof=1) + nan_y = variation(nan_a, nan_policy="omit", ddof=1) + xp_assert_close(y, xp.asarray(math.sqrt(5/2)/3)) + assert y == nan_y diff --git a/v12x/pylib/scipy/version.py b/v12x/pylib/scipy/version.py new file mode 100644 index 0000000000000000000000000000000000000000..4dfc7971da5a34f847e2f075220b053f92123fe5 --- /dev/null +++ b/v12x/pylib/scipy/version.py @@ -0,0 +1,12 @@ + +""" +Module to expose more detailed version info for the installed `scipy` +""" +version = "1.18.1" +full_version = version +short_version = version.split('.dev')[0] +git_revision = "e4e854eaa8f18d807cd3496028e257e36caa93cc" +release = 'dev' not in version and '+' not in version + +if not release: + version = full_version diff --git a/v12x/pylib/xgboost-3.4.1.dist-info/sboms/auditwheel.cdx.json b/v12x/pylib/xgboost-3.4.1.dist-info/sboms/auditwheel.cdx.json new file mode 100644 index 0000000000000000000000000000000000000000..b8184b08fd5b2672dd232e89db2167a62379e09b --- /dev/null +++ b/v12x/pylib/xgboost-3.4.1.dist-info/sboms/auditwheel.cdx.json @@ -0,0 +1 @@ +{"bomFormat": "CycloneDX", "specVersion": "1.4", "version": 1, "metadata": {"component": {"type": "library", "bom-ref": "pkg:pypi/xgboost@3.4.1?file_name=xgboost-3.4.1-py3-none-manylinux_2_28_x86_64.whl", "name": "xgboost", "version": "3.4.1", "purl": "pkg:pypi/xgboost@3.4.1?file_name=xgboost-3.4.1-py3-none-manylinux_2_28_x86_64.whl"}, "tools": [{"name": "auditwheel", "version": "6.8.0"}]}, "components": [{"type": "library", "bom-ref": "pkg:pypi/xgboost@3.4.1?file_name=xgboost-3.4.1-py3-none-manylinux_2_28_x86_64.whl", "name": "xgboost", "version": "3.4.1", "purl": "pkg:pypi/xgboost@3.4.1?file_name=xgboost-3.4.1-py3-none-manylinux_2_28_x86_64.whl"}, {"type": "library", "bom-ref": "pkg:rpm/almalinux/libgomp@8.5.0-28.el8_10.alma.1#c61017c9a24eb6e1e1a3cdc9becd004a6419cbda3d54b4848b98f240a4829571", "name": "libgomp", "version": "8.5.0-28.el8_10.alma.1", "purl": "pkg:rpm/almalinux/libgomp@8.5.0-28.el8_10.alma.1"}], "dependencies": [{"ref": "pkg:pypi/xgboost@3.4.1?file_name=xgboost-3.4.1-py3-none-manylinux_2_28_x86_64.whl", "dependsOn": ["pkg:rpm/almalinux/libgomp@8.5.0-28.el8_10.alma.1#c61017c9a24eb6e1e1a3cdc9becd004a6419cbda3d54b4848b98f240a4829571"]}, {"ref": "pkg:rpm/almalinux/libgomp@8.5.0-28.el8_10.alma.1#c61017c9a24eb6e1e1a3cdc9becd004a6419cbda3d54b4848b98f240a4829571"}]} \ No newline at end of file diff --git a/v12x/pylib/xgboost/__init__.py b/v12x/pylib/xgboost/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..dc1a0869b30c57902450d0f48ed6824e6c2b86de --- /dev/null +++ b/v12x/pylib/xgboost/__init__.py @@ -0,0 +1,70 @@ +"""XGBoost: eXtreme Gradient Boosting library. + +Contributors: https://github.com/dmlc/xgboost/blob/master/CONTRIBUTORS.md +""" + +from . import ( + collective, + interpret, + tracker, # noqa +) +from ._c_api import _py_version +from .core import ( + Booster, + DataIter, + DMatrix, + ExtMemQuantileDMatrix, + QuantileDMatrix, + build_info, +) +from .tracker import RabitTracker # noqa +from .training import cv, train + +try: + from .config import config_context, get_config, set_config + from .plotting import plot_importance, plot_tree, to_graphviz + from .sklearn import ( + XGBClassifier, + XGBModel, + XGBRanker, + XGBRegressor, + XGBRFClassifier, + XGBRFRegressor, + ) +except ImportError: + pass + + +__version__ = _py_version() + + +__all__ = [ + # core + "DMatrix", + "QuantileDMatrix", + "ExtMemQuantileDMatrix", + "Booster", + "DataIter", + "train", + "cv", + # utilities + "RabitTracker", + "build_info", + "plot_importance", + "plot_tree", + "to_graphviz", + "set_config", + "get_config", + "config_context", + # sklearn + "XGBModel", + "XGBClassifier", + "XGBRegressor", + "XGBRanker", + "XGBRFClassifier", + "XGBRFRegressor", + # collective + "collective", + # interpretability + "interpret", +] diff --git a/v12x/pylib/xgboost/_c_api.py b/v12x/pylib/xgboost/_c_api.py new file mode 100644 index 0000000000000000000000000000000000000000..5162d952dbd3efbefe1dde03769683702b6c8f5f --- /dev/null +++ b/v12x/pylib/xgboost/_c_api.py @@ -0,0 +1,195 @@ +"""Low-level ctypes bridge for the XGBoost C API.""" + +import ctypes +import json +import os +import warnings +from typing import Any, Callable, List, Tuple, Union, cast, overload + +from ._typing import CStrPptr, c_bst_ulong +from .compat import py_str +from .libpath import find_lib_path + + +class XGBoostError(ValueError): + """Error thrown by xgboost trainer.""" + + +@overload +def from_pystr_to_cstr(data: str) -> bytes: ... + + +@overload +def from_pystr_to_cstr(data: List[str]) -> ctypes.Array: ... + + +def from_pystr_to_cstr(data: Union[str, List[str]]) -> Union[bytes, ctypes.Array]: + """Convert a Python str or list of Python str to C pointer.""" + if isinstance(data, str): + return bytes(data, "utf-8") + if isinstance(data, list): + data_as_bytes: List[bytes] = [bytes(d, "utf-8") for d in data] + pointers: ctypes.Array[ctypes.c_char_p] = ( + ctypes.c_char_p * len(data_as_bytes) + )(*data_as_bytes) + return pointers + raise TypeError() + + +def from_cstr_to_pystr(data: CStrPptr, length: c_bst_ulong) -> List[str]: + """Revert C pointer to Python str.""" + res = [] + for i in range(length.value): + try: + res.append(str(cast(bytes, data[i]).decode("ascii"))) + except UnicodeDecodeError: + res.append(str(cast(bytes, data[i]).decode("utf-8"))) + return res + + +def make_jcargs(**kwargs: Any) -> bytes: + """Make JSON-based arguments for C functions.""" + return from_pystr_to_cstr(json.dumps(kwargs)) + + +def _log_callback(msg: bytes) -> None: + """Redirect logs from native library into Python console.""" + smsg = py_str(msg) + if smsg.find("WARNING:") != -1: + # Stacklevel: + # 1: This line + # 2: XGBoost C functions like `_LIB.XGBoosterTrainOneIter`. + # 3: The Python function that calls the C function. + warnings.warn(smsg, UserWarning, stacklevel=3) + return + print(smsg) + + +def _get_log_callback_func() -> Callable: + """Wrap log_callback() method in ctypes callback type.""" + c_callback = ctypes.CFUNCTYPE(None, ctypes.c_char_p) + return c_callback(_log_callback) + + +def _lib_version(lib: ctypes.CDLL) -> Tuple[int, int, int]: + """Get the XGBoost version from native shared object.""" + major = ctypes.c_int() + minor = ctypes.c_int() + patch = ctypes.c_int() + lib.XGBoostVersion(ctypes.byref(major), ctypes.byref(minor), ctypes.byref(patch)) + return major.value, minor.value, patch.value + + +def _py_version() -> str: + """Get the XGBoost version from Python version file.""" + version_file = os.path.join(os.path.dirname(__file__), "VERSION") + with open(version_file, encoding="ascii") as f: + return f.read().strip() + + +def _register_log_callback(lib: ctypes.CDLL) -> None: + lib.XGBGetLastError.restype = ctypes.c_char_p + lib.callback = _get_log_callback_func() # type: ignore[attr-defined] + if lib.XGBRegisterLogCallback(lib.callback) != 0: + raise XGBoostError(lib.XGBGetLastError()) + + +def _parse_version(ver: str) -> Tuple[Tuple[int, int, int], str]: + """Avoid dependency on packaging (PEP 440).""" + # 2.0.0-dev, 2.0.0, 2.0.0.post1, or 2.0.0rc1 + if ver.find("post") != -1: + major, minor, patch = ver.split(".")[:-1] + postfix = ver.split(".")[-1] + elif "-dev" in ver: + major, minor, patch = ver.split("-")[0].split(".") + postfix = "dev" + else: + major, minor, patch = ver.split(".") + rc = patch.find("rc") + if rc != -1: + postfix = patch[rc:] + patch = patch[:rc] + else: + postfix = "" + + return (int(major), int(minor), int(patch)), postfix + + +def _load_lib() -> ctypes.CDLL: + """Load xgboost library.""" + lib_paths = find_lib_path() + if not lib_paths: + # This happens only when building document. + return None # type: ignore[return-value] + try: + path_backup = os.environ["PATH"].split(os.pathsep) + except KeyError: + path_backup = [] + lib_success = False + os_error_list = [] + for lib_path in lib_paths: + try: + # needed when the lib is linked with non-system-available + # dependencies + os.environ["PATH"] = os.pathsep.join( + path_backup + [os.path.dirname(lib_path)] + ) + lib = ctypes.cdll.LoadLibrary(lib_path) + setattr(lib, "path", os.path.normpath(lib_path)) + lib_success = True + break + except OSError as e: + os_error_list.append(str(e)) + continue + finally: + os.environ["PATH"] = os.pathsep.join(path_backup) + if not lib_success: + libname = os.path.basename(lib_paths[0]) + raise XGBoostError(f""" +XGBoost Library ({libname}) could not be loaded. +Likely causes: + * OpenMP runtime is not installed + - vcomp140.dll or libgomp-1.dll for Windows + - libomp.dylib for Mac OSX + - libgomp.so for Linux and other UNIX-like OSes + Mac OSX users: Run `brew install libomp` to install OpenMP runtime. + + * You are running 32-bit Python on a 64-bit OS + +Error message(s): {os_error_list} +""") + _register_log_callback(lib) + + libver = _lib_version(lib) + pyver, _ = _parse_version(_py_version()) + + # verify that we are loading the correct binary. + if pyver != libver: + pyver_str = ".".join((str(v) for v in pyver)) + libver_str = ".".join((str(v) for v in libver)) + msg = ( + "Mismatched version between the Python package and the native shared " + f"""object. Python package version: {pyver_str}. Shared object """ + f"""version: {libver_str}. Shared object is loaded from: {lib.path}. +Likely cause: + * XGBoost is first installed with anaconda then upgraded with pip. To fix it """ + "please remove one of the installations." + ) + raise ValueError(msg) + + return lib + + +# load the XGBoost library globally +_LIB = _load_lib() + + +def _check_call(ret: int) -> None: + """Check the return value of C API call.""" + if ret != 0: + raise XGBoostError(py_str(_LIB.XGBGetLastError())) + + +def c_str(string: str) -> ctypes.c_char_p: + """Convert a python string to cstring.""" + return ctypes.c_char_p(string.encode("utf-8")) diff --git a/v12x/pylib/xgboost/_data_utils.py b/v12x/pylib/xgboost/_data_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..b76909734212d3f9656ddb18cb775967a2d13d85 --- /dev/null +++ b/v12x/pylib/xgboost/_data_utils.py @@ -0,0 +1,760 @@ +"""Helpers for interfacing array like objects.""" + +import copy +import ctypes +import json +import warnings +from abc import ABC, abstractmethod +from functools import cache as fcache +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Dict, + List, + Literal, + Optional, + Protocol, + Tuple, + Type, + TypeAlias, + TypedDict, + TypeGuard, + Union, + cast, + overload, +) + +import numpy as np + +from ._typing import ( + ArrowCatList, + CNumericPtr, + DataType, + FeatureTypes, + NumpyDType, + NumpyOrCupy, +) +from .compat import import_cupy, import_pyarrow, lazy_isinstance + +if TYPE_CHECKING: + import pandas as pd + import pyarrow as pa + + +# Used for accepting inputs for numpy and cupy arrays +class _ArrayLikeArg(Protocol): + @property + def __array_interface__(self) -> "ArrayInf": ... + + +class _CudaArrayLikeArg(Protocol): + @property + def __cuda_array_interface__(self) -> "CudaArrayInf": ... + + +ArrayInf = TypedDict( + "ArrayInf", + { + "data": Tuple[int, bool], + "typestr": str, + "version": Literal[3], + "strides": Optional[Tuple[int, ...]], + "shape": Tuple[int, ...], + "mask": Union["ArrayInf", None, _ArrayLikeArg], + }, +) + +CudaArrayInf = TypedDict( + "CudaArrayInf", + { + "data": Tuple[int, bool], + "typestr": str, + "version": Literal[3], + "strides": Optional[Tuple[int, ...]], + "shape": Tuple[int, ...], + "mask": Union["ArrayInf", None, _ArrayLikeArg], + "stream": int, + }, +) + +StringArray = TypedDict("StringArray", {"offsets": ArrayInf, "values": ArrayInf}) +CudaStringArray = TypedDict( + "CudaStringArray", {"offsets": CudaArrayInf, "values": CudaArrayInf} +) + + +def array_hasobject(data: DataType) -> bool: + """Whether the numpy array has object dtype.""" + return ( + hasattr(data, "dtype") + and hasattr(data.dtype, "hasobject") + and data.dtype.hasobject + ) + + +def cuda_array_interface_dict(data: _CudaArrayLikeArg) -> CudaArrayInf: + """Returns a dictionary storing the CUDA array interface.""" + if array_hasobject(data): + raise ValueError("Input data contains `object` dtype. Expecting numeric data.") + ainf = data.__cuda_array_interface__ + if "mask" in ainf and ainf["mask"] is not None: + mask_ainf = ainf["mask"].__cuda_array_interface__ # type: ignore[union-attr] + # Normalize the validity mask to XGBoost's expected layout (`|t1` bit field of + # length `n_samples`). + typestr = mask_ainf["typestr"] + n_samples = ainf["shape"][0] + if typestr[1] in ("u", "i", "t") and typestr[2:] == "1" and n_samples: + mask_ainf = dict(mask_ainf) + mask_ainf["typestr"] = "|t1" + mask_ainf["shape"] = (n_samples,) + mask_ainf.pop("strides", None) + ainf["mask"] = mask_ainf + return ainf + + +def cuda_array_interface(data: _CudaArrayLikeArg) -> bytes: + """Make cuda array interface str.""" + interface = cuda_array_interface_dict(data) + interface_str = bytes(json.dumps(interface), "utf-8") + return interface_str + + +def from_array_interface(interface: ArrayInf, zero_copy: bool = False) -> NumpyOrCupy: + """Convert array interface to numpy or cupy array""" + + class Array: + """Wrapper type for communicating with numpy and cupy.""" + + _interface: Optional[ArrayInf] = None + + @property + def __array_interface__(self) -> Optional[ArrayInf]: + return self._interface + + @__array_interface__.setter + def __array_interface__(self, interface: ArrayInf) -> None: + self._interface = copy.copy(interface) + # Convert some fields to tuple as required by numpy + self._interface["shape"] = tuple(self._interface["shape"]) + self._interface["data"] = ( + self._interface["data"][0], + self._interface["data"][1], + ) + strides = self._interface.get("strides", None) + if strides is not None: + self._interface["strides"] = tuple(strides) + + @property + def __cuda_array_interface__(self) -> Optional[ArrayInf]: + return self.__array_interface__ + + @__cuda_array_interface__.setter + def __cuda_array_interface__(self, interface: ArrayInf) -> None: + self.__array_interface__ = interface + + @property + def shape(self) -> Tuple[int, ...]: + """Shape of the input array.""" + aif = self.__array_interface__ + assert aif is not None + return aif["shape"] + + @property + def size(self) -> np.signedinteger: + """Total size of the input array.""" + return np.prod(self.shape) + + arr = Array() + + # Cupy and numpy might run into issue when constructing an empty array from an array + # interface. we explicitly check for emptiness. + if "stream" in interface: + # CUDA stream is presented, this is a __cuda_array_interface__. + arr.__cuda_array_interface__ = interface + cp = import_cupy() + if arr.size == 0: + return cp.empty(shape=arr.shape, dtype=np.dtype(interface["typestr"])) + out = cp.array(arr, copy=not zero_copy) + else: + arr.__array_interface__ = interface + if arr.size == 0: + return np.empty(shape=arr.shape, dtype=np.dtype(interface["typestr"])) + out = np.array(arr, copy=not zero_copy) + + return out + + +# Default constant value for CUDA per-thread stream. +STREAM_PER_THREAD = 2 + + +# Typing is not strict as there are subtle differences between CUDA array interface and +# array interface. We handle them uniformly for now. +def make_array_interface( + ptr: Union[CNumericPtr, int], + shape: Tuple[int, ...], + dtype: Type[np.number], + is_cuda: bool, +) -> ArrayInf: + """Make an __(cuda)_array_interface__ from a pointer.""" + # Use an empty array to handle typestr and descr + if is_cuda: + empty = import_cupy().empty(shape=(0,), dtype=dtype) + array = empty.__cuda_array_interface__ # pylint: disable=no-member + else: + empty = np.empty(shape=(0,), dtype=dtype) + array = empty.__array_interface__ # pylint: disable=no-member + + if not isinstance(ptr, int): + addr = ctypes.cast(ptr, ctypes.c_void_p).value + else: + addr = ptr + length = int(np.prod(shape)) + # Handle empty dataset. + assert addr is not None or length == 0 + + if addr is None: + # `array` still has the placeholder (0,) shape from the empty array used + # to derive `typestr`/`descr` above; update it to the actual requested + # shape (e.g. (0, n_features) for an empty multi-dimensional result) + # instead of always reporting a flat 1-D empty shape. + array["shape"] = shape + return array + + array["data"] = (addr, True) + if is_cuda and "stream" not in array: + array["stream"] = STREAM_PER_THREAD + array["shape"] = shape + array["strides"] = None + return array + + +def is_arrow_dict(data: Any) -> TypeGuard["pa.DictionaryArray"]: + """Is this an arrow dictionary array?""" + return lazy_isinstance(data, "pyarrow.lib", "DictionaryArray") + + +class DfCatAccessor(Protocol): + """Protocol for pandas cat accessor.""" + + @property + def categories( # pylint: disable=missing-function-docstring + self, + ) -> "pd.Index": ... + + @property + def codes(self) -> "pd.Series": ... # pylint: disable=missing-function-docstring + + @property + def dtype(self) -> np.dtype: ... # pylint: disable=missing-function-docstring + + @property + def values(self) -> np.ndarray: ... # pylint: disable=missing-function-docstring + + def to_arrow( # pylint: disable=missing-function-docstring + self, + ) -> Union["pa.StringArray", "pa.IntegerArray"]: ... + + @property + def __cuda_array_interface__(self) -> CudaArrayInf: ... + + @property + def _column(self) -> Any: ... + + +def _is_df_cat(data: Any) -> TypeGuard[DfCatAccessor]: + # Test pd.Series.cat, not pd.Series + return hasattr(data, "categories") and hasattr(data, "codes") + + +@fcache +def _arrow_npdtype() -> Dict[Any, Type[np.number]]: + import pyarrow as pa + + mapping: Dict[Any, Type[np.number]] = { + pa.int8(): np.int8, + pa.int16(): np.int16, + pa.int32(): np.int32, + pa.int64(): np.int64, + pa.uint8(): np.uint8, + pa.uint16(): np.uint16, + pa.uint32(): np.uint32, + pa.uint64(): np.uint64, + pa.float16(): np.float16, + pa.float32(): np.float32, + pa.float64(): np.float64, + } + + return mapping + + +@overload +def _arrow_buf_inf(address: int, typestr: str, size: int, stream: None) -> ArrayInf: ... + + +@overload +def _arrow_buf_inf( + address: int, typestr: str, size: int, stream: int +) -> CudaArrayInf: ... + + +def _arrow_buf_inf( + address: int, typestr: str, size: int, stream: Optional[int] +) -> Union[ArrayInf, CudaArrayInf]: + if stream is not None: + jcuaif: CudaArrayInf = { + "data": (address, True), + "typestr": typestr, + "version": 3, + "strides": None, + "shape": (size,), + "mask": None, + "stream": stream, + } + return jcuaif + + jaif: ArrayInf = { + "data": (address, True), + "typestr": typestr, + "version": 3, + "strides": None, + "shape": (size,), + "mask": None, + } + return jaif + + +def _arrow_cat_names_inf(cats: "pa.StringArray") -> Tuple[StringArray, Any]: + if not TYPE_CHECKING: + pa = import_pyarrow() + + # FIXME(jiamingy): Account for offset, need to find an implementation that returns + # offset > 0 + assert cats.offset == 0 + buffers: List[pa.Buffer] = cats.buffers() + mask, offset, data = buffers + assert offset.is_cpu + + off_len = len(cats) + 1 + + def get_n_bytes(typ: Type) -> int: + return off_len * (np.iinfo(typ).bits // 8) + + if offset.size == get_n_bytes(np.int64): + if not isinstance(cats, pa.LargeStringArray): + arrow_str_error = "Expecting a `pyarrow.Array`." + raise TypeError(arrow_str_error + f" Got: {type(cats)}.") + # Convert to 32bit integer, arrow recommends against the use of i64. Also, + # XGBoost cannot handle large number of categories (> 2**31). + i32cats = cats.cast(pa.string()) + mask, offset, data = i32cats.buffers() + + if offset.size != get_n_bytes(np.int32): + raise TypeError( + "Arrow dictionary type offsets is required to be 32-bit integer." + ) + + joffset = _arrow_buf_inf(offset.address, " ArrayInf: + """Helper for handling categorical codes.""" + if not TYPE_CHECKING: + pa = import_pyarrow() + if not isinstance(array, pa.Array): # pylint: disable=E0606 + raise TypeError(f"Invalid input type: {type(array)}") + + mask, data = array.buffers() + jdata = make_array_interface( + data.address, + shape=(len(array),), + dtype=_arrow_npdtype()[array.type], + is_cuda=not data.is_cpu, + ) + + if mask is not None: + jmask: Optional[ArrayInf] = { + "data": (mask.address, True), + "typestr": " Tuple[StringArray, ArrayInf, Tuple]: + """Get the array interface representation of a string-based category array.""" + jnames, cats_tmp = _arrow_cat_names_inf(cats) + jcodes = _arrow_array_inf(codes) + + return jnames, jcodes, (cats_tmp, None) + + +def _ensure_np_dtype( + data: DataType, dtype: Optional[NumpyDType] +) -> Tuple[np.ndarray, Optional[NumpyDType]]: + """Ensure the np array has correct type and is contiguous.""" + if array_hasobject(data) or data.dtype in [np.float16, np.bool_]: + dtype = np.float32 + data = data.astype(dtype, copy=False) + if not data.flags.aligned: + data = np.require(data, requirements="A") + return data, dtype + + +def _is_flatten(array: NumpyOrCupy) -> bool: + return len(array.shape) == 1 or array.shape[1] == 1 + + +def array_interface_dict(data: np.ndarray) -> ArrayInf: + """Returns an array interface from the input.""" + if array_hasobject(data): + raise ValueError("Input data contains `object` dtype. Expecting numeric data.") + ainf = data.__array_interface__ + if "mask" in ainf: + ainf["mask"] = ainf["mask"].__array_interface__ + return cast(ArrayInf, ainf) + + +def _arrow_string_offsets(lengths: np.ndarray) -> np.ndarray: + """Build 32-bit Arrow string offsets from encoded byte lengths.""" + offsets = np.empty(lengths.size + 1, dtype=np.int64) + offsets[0] = 0 + np.cumsum(lengths, dtype=np.int64, out=offsets[1:]) + if offsets[-1] > np.iinfo(np.int32).max: + raise ValueError( + "The encoded categories exceed the maximum size of the 32-bit " + "offsets used by Arrow string arrays." + ) + return offsets.astype(np.int32) + + +def pd_cat_inf( # pylint: disable=too-many-locals + cats: "pd.Index", codes: "pd.Series" +) -> Tuple[Union[StringArray, ArrayInf], ArrayInf, Tuple]: + """Get the array interface representation of a pandas category index.""" + from pandas.api.types import infer_dtype, is_float_dtype, is_integer_dtype + + # pandas uses -1 to represent missing values for categorical features + codes = codes.replace(-1, np.nan) + + if is_integer_dtype(cats.dtype) or is_float_dtype(cats.dtype): + # Numeric index type + cat_dtype: Any = cats.dtype + np_dtype = np.dtype(getattr(cat_dtype, "numpy_dtype", cat_dtype)) + if hasattr(cats, "to_numpy"): + name_values_num = cats.to_numpy(dtype=np_dtype) + else: + name_values_num = np.asarray(cats.values, dtype=np_dtype) + name_values_num = np.require(name_values_num, requirements=["A", "C"]) + if np_dtype == np.dtype(np.uint64) and np.any( + name_values_num > np.iinfo(np.int64).max + ): + raise ValueError( + "Category index values must not exceed the signed 64-bit range." + ) + jarr_values = array_interface_dict(name_values_num) + code_values = cast(np.ndarray, codes.values) + jarr_codes = array_interface_dict(code_values) + return jarr_values, jarr_codes, (name_values_num, code_values) + + def npstr_to_arrow_strarr(strarr: Any) -> Tuple[np.ndarray, bytes]: + """Convert a string-like array to an arrow string array.""" + if not isinstance(strarr, np.ndarray): + if hasattr(strarr, "to_numpy"): + strarr = strarr.to_numpy(dtype=object) + else: + strarr = np.asarray(strarr, dtype=object) + + if strarr.size == 0: + lengths = np.empty(0, dtype=np.int64) + offsets = _arrow_string_offsets(lengths) + values = b"" + else: + inferred = infer_dtype(strarr, skipna=False) + values_list = strarr.tolist() + if inferred == "string": + encoded = [value.encode("utf-8") for value in values_list] + elif inferred == "bytes": + encoded = values_list + else: + raise TypeError( + "Category index must contain only values of the same type, " + "either string or integer. " + f"Got values of type `{inferred}`." + ) + lengths = np.fromiter( + (len(value) for value in encoded), + dtype=np.int64, + count=strarr.size, + ) + offsets = _arrow_string_offsets(lengths) + values = b"".join(encoded) + + if b"\0" in values: + warnings.warn( + ( + "Found embedded NUL (\\0) characters in string categories. " + "Arrow used to strip these characters, but they are now preserved." + ), + UserWarning, + ) + return offsets, values + + # String index type + name_offsets, name_values = npstr_to_arrow_strarr(cats.values) + joffsets = array_interface_dict(name_offsets) + name_values_np = np.frombuffer(name_values, dtype=np.int8) + jvalues = array_interface_dict(name_values_np) + jnames: StringArray = {"offsets": joffsets, "values": jvalues} + + code_values = cast(np.ndarray, codes.values) + jcodes = array_interface_dict(code_values) + + buf = ( + name_offsets, + name_values, + name_values_np, + code_values, + ) # store temporary values + return jnames, jcodes, buf + + +def array_interface(data: np.ndarray) -> bytes: + """Make array interface str.""" + interface = array_interface_dict(data) + interface_str = bytes(json.dumps(interface), "utf-8") + return interface_str + + +def check_cudf_meta(data: _CudaArrayLikeArg, field: str) -> None: + "Make sure no missing value in meta data." + if ( + "mask" in data.__cuda_array_interface__ + and data.__cuda_array_interface__["mask"] is not None + ): + raise ValueError(f"Missing value is not allowed for: {field}") + + +def _cudf_str_cat_inf(cats: DfCatAccessor) -> Tuple[CudaStringArray, Tuple]: + """String category index path for :py:func:`cudf_cat_inf`.""" + import pylibcudf as plc # pylint: disable=import-outside-toplevel + + # pylint: disable=protected-access + plc_col = cats._column.to_pylibcudf() + if plc_col.type().id() != plc.TypeId.STRING: + raise TypeError( + "Unexpected type for category index. It's neither numeric nor string." + ) + # Categories should not have missing values nor a non-zero logical offset. + assert plc_col.null_count() == 0 + assert plc_col.offset() == 0 + + off_child = plc_col.children()[0] # offsets + assert off_child.type().id() == plc.TypeId.INT32, "Expected INT32 string offsets." + + # String category index in arrow format + jdata: CudaArrayInf = _arrow_buf_inf( + plc_col.data().__cuda_array_interface__["data"][0], + "|i1", + 0, + STREAM_PER_THREAD, + ) + joffset: CudaArrayInf = _arrow_buf_inf( + off_child.data().__cuda_array_interface__["data"][0], + " Tuple[Union[CudaArrayInf, CudaStringArray], ArrayInf, Tuple]: + """Obtain the cuda array interface for cuDF categories.""" + cp = import_cupy() + is_num_idx = cp.issubdtype(cats.dtype, cp.floating) or cp.issubdtype( + cats.dtype, cp.integer + ) + if is_num_idx: + cats_ainf = cuda_array_interface_dict(cats) + codes_ainf = cuda_array_interface_dict(codes) + return cats_ainf, codes_ainf, (cats, codes) + + jnames, buf = _cudf_str_cat_inf(cats) + jcodes = cuda_array_interface_dict(codes) + return jnames, jcodes, buf + + +class Categories: + """An internal storage class for categories returned by the DMatrix and the + Booster. This class is designed to be opaque. It is intended to be used exclusively + by XGBoost as an intermediate storage for re-coding categorical data. + + The categories are saved along with the booster object. As a result, users don't + need to preserve this class for re-coding. Use the booster model IO instead if you + want to preserve the categories in a stable format. + + .. versionadded:: 3.1.0 + + .. warning:: + + This class is internal. + + .. code-block:: python + + Xy = xgboost.QuantileDMatrix(X, y, enable_categorical=True) + booster = xgboost.train({}, Xy) + + categories = booster.get_categories() # Get categories + + # Use categories as a reference for re-coding + Xy_new = xgboost.QuantileDMatrix( + X_new, y_new, feature_types=categories, enable_categorical=True, ref=Xy + ) + + # Categories will be part of the `model.json`. + booster.save_model("model.json") + + """ + + def __init__( + self, + handle: Tuple[ctypes.c_void_p, Callable[[], None]], + arrow_arrays: Optional[ArrowCatList], + ) -> None: + # The handle type is a bundle of the handle and the free call. Otherwise, we + # will have to import the `_lib` and the `_check_call` from the core module + # inside the __del__ method to avoid cyclic model dependency. + # Importing modules in __del__ can result in Python abort if __del__ is called + # during exception handling (interpreter is shutting down). + self._handle, self._free = handle + self._arrow_arrays = arrow_arrays + + def to_arrow(self) -> ArrowCatList: + """Get the categories in the dataset. The results are stored in a list of + (feature name, arrow array) pairs, with one array for each categorical + feature. If a feature is numerical, then the corresponding column in the list is + None. A value error will be raised if this container was created without the + `export_to_arrow` option. + + """ + if self._arrow_arrays is None: + raise ValueError( + "The `export_to_arrow` option of the `get_categories` method" + " is required." + ) + return self._arrow_arrays + + def empty(self) -> bool: + """Returns True if there's no category.""" + return self._handle.value is None + + def get_handle(self) -> int: + """Internal method for retrieving the handle.""" + assert self._handle.value + return self._handle.value + + def __del__(self) -> None: + if self._handle.value is None: + return + self._free() + + +def get_ref_categories( + feature_types: Optional[Union[FeatureTypes, Categories]], +) -> Tuple[Optional[FeatureTypes], Optional[Categories]]: + """Get the optional reference categories from the `feature_types`. This is used by + various `DMatrix` where the `feature_types` is reused for specifying the reference + categories. + + """ + if isinstance(feature_types, Categories): + ref_categories = feature_types + feature_types = None + else: + ref_categories = None + return feature_types, ref_categories + + +# Type schema for storing JSON-encoded array interface +AifType: TypeAlias = List[ + Union[ + # numeric column + Union[ArrayInf, CudaArrayInf], + # categorical column + Tuple[ + # (cuda) numeric index | (cuda) string index + Union[ArrayInf, CudaArrayInf, StringArray, CudaStringArray], + Union[ArrayInf, CudaArrayInf], # codes + ], + ] +] + + +class TransformedDf(ABC): + """Internal class for storing transformed dataframe. + + Parameters + ---------- + ref_categories : + Optional reference categories used for re-coding. + + aitfs : + Array interface for each column. + + """ + + def __init__( + self, + ref_categories: Optional[Categories], + aitfs: AifType, + temporary_buffers: List[Tuple], + ) -> None: + self.ref_categories = ref_categories + if ref_categories is not None and ref_categories.get_handle() is not None: + aif = ref_categories.get_handle() + self.ref_aif: Optional[int] = aif + else: + self.ref_aif = None + + self.aitfs = aitfs + self.temporary_buffers = temporary_buffers + + def array_interface(self) -> bytes: + """Return a byte string for JSON encoded array interface.""" + if self.ref_categories is not None: + ref_inf: dict = {"ref_categories": self.ref_aif, "columns": self.aitfs} + inf = bytes(json.dumps(ref_inf), "utf-8") + else: + inf = bytes(json.dumps(self.aitfs), "utf-8") + return inf + + @property + @abstractmethod + def shape(self) -> Tuple[int, int]: + """Return the shape of the dataframe.""" diff --git a/v12x/pylib/xgboost/_typing.py b/v12x/pylib/xgboost/_typing.py new file mode 100644 index 0000000000000000000000000000000000000000..cd89d01d22133ff5433a30f895c4c19623ea199e --- /dev/null +++ b/v12x/pylib/xgboost/_typing.py @@ -0,0 +1,125 @@ +# pylint: disable=protected-access +"""Shared typing definition.""" + +import ctypes +import os +from typing import ( + TYPE_CHECKING, + Any, + AnyStr, + Callable, + Dict, + List, + Optional, + Sequence, + Tuple, + Type, + TypeAlias, + TypeVar, + Union, +) + +import numpy as np + +DataType = Any + +FeatureInfo = Sequence[str] +FeatureNames = FeatureInfo +FeatureTypes = FeatureInfo +BoosterParam = Union[List, Dict[str, Any]] # better be sequence + +ArrayLike = Any +if TYPE_CHECKING: + import pyarrow as pa + + PathLike = Union[str, os.PathLike[str]] +else: + PathLike = Union[str, os.PathLike] + +ArrowCatCol: TypeAlias = Optional[Union["pa.StringArray", "pa.NumericArray"]] +ArrowCatList: TypeAlias = List[Tuple[str, Optional[ArrowCatCol]]] + +CupyT = ArrayLike # maybe need a stub for cupy arrays +NumpyOrCupy = Union[np.ndarray, Any] +NumpyDType = Union[str, Type[np.number], np.dtype[Any]] +PandasDType = Any # real type is pandas.core.dtypes.base.ExtensionDtype + +FloatCompatible = Union[float, np.float32, np.float64] + +# typing.SupportsInt is not suitable here since floating point values are convertible to +# integers as well. +Integer = Union[int, np.integer] +IterationRange = Tuple[Integer, Integer] + +# callables +FPreProcCallable = Callable + +# ctypes +# c_bst_ulong corresponds to bst_ulong defined in xgboost/c_api.h +c_bst_ulong = ctypes.c_uint64 # pylint: disable=C0103 + +ModelIn = Union[os.PathLike[AnyStr], bytearray, str] + +CTypeT = TypeVar( + "CTypeT", + ctypes.c_void_p, + ctypes.c_char_p, + ctypes.c_int, + ctypes.c_float, + ctypes.c_uint, + ctypes.c_size_t, +) + +# supported numeric types +CNumeric = Union[ + ctypes.c_float, + ctypes.c_double, + ctypes.c_uint, + ctypes.c_uint64, + ctypes.c_int32, + ctypes.c_int64, +] + +# c pointer types +if TYPE_CHECKING: + CStrPtr = ctypes._Pointer[ctypes.c_char] + + CStrPptr = ctypes._Pointer[ctypes.c_char_p] + + CFloatPtr = ctypes._Pointer[ctypes.c_float] + + CNumericPtr = Union[ + ctypes._Pointer[ctypes.c_float], + ctypes._Pointer[ctypes.c_double], + ctypes._Pointer[ctypes.c_uint], + ctypes._Pointer[ctypes.c_uint64], + ctypes._Pointer[ctypes.c_int32], + ctypes._Pointer[ctypes.c_int64], + ] +else: + CStrPtr = ctypes._Pointer + + CStrPptr = ctypes._Pointer + + CFloatPtr = ctypes._Pointer + + CNumericPtr = Union[ + ctypes._Pointer, + ctypes._Pointer, + ctypes._Pointer, + ctypes._Pointer, + ctypes._Pointer, + ctypes._Pointer, + ] + +# The second arg is actually Optional[List[cudf.Series]], skipped for easier type check. +# The cudf Series is the obtained cat codes, preserved in the `DataIter` to prevent it +# being freed. +TransformedData = Tuple[Any, Optional[FeatureNames], Optional[FeatureTypes]] + +# template parameter +_T = TypeVar("_T") +_F = TypeVar("_F", bound=Callable[..., Any]) + +_ScoreList = Union[List[float], List[Tuple[float, float]]] +EvalsLog: TypeAlias = Dict[str, Dict[str, _ScoreList]] diff --git a/v12x/pylib/xgboost/callback.py b/v12x/pylib/xgboost/callback.py new file mode 100644 index 0000000000000000000000000000000000000000..d1b05059cea37cacb042302e5e9e0e07922cbd00 --- /dev/null +++ b/v12x/pylib/xgboost/callback.py @@ -0,0 +1,652 @@ +"""Callback library containing training routines. See :doc:`Callback Functions +` for a quick introduction. + +""" + +import collections +import os +import pickle +from abc import ABC +from typing import ( + Any, + Callable, + Dict, + List, + Optional, + Sequence, + Tuple, + TypeAlias, + TypeVar, + Union, + cast, +) + +import numpy + +from . import collective +from ._typing import EvalsLog, _ScoreList +from .core import ( + Booster, + DMatrix, + XGBoostError, + _deprecate_positional_args, + _parse_eval_str, +) + +__all__ = [ + "TrainingCallback", + "LearningRateScheduler", + "EarlyStopping", + "EvaluationMonitor", + "TrainingCheckPoint", + "CallbackContainer", +] + +_Score = Union[float, Tuple[float, float]] + +_Model = Any # real type is Union[Booster, CVPack]; need more work + + +# pylint: disable=unused-argument +class TrainingCallback(ABC): + """Interface for training callback. + + .. versionadded:: 1.3.0 + + """ + + EvalsLog: TypeAlias = EvalsLog + + def __init__(self) -> None: + pass + + def before_training(self, model: _Model) -> _Model: + """Run before training starts.""" + return model + + def after_training(self, model: _Model) -> _Model: + """Run after training is finished.""" + return model + + def before_iteration(self, model: _Model, epoch: int, evals_log: EvalsLog) -> bool: + """Run before each iteration. Returns True when training should stop. See + :py:meth:`after_iteration` for details. + + """ + return False + + def after_iteration(self, model: _Model, epoch: int, evals_log: EvalsLog) -> bool: + """Run after each iteration. Returns `True` when training should stop. + + Parameters + ---------- + + model : + Eeither a :py:class:`~xgboost.Booster` object or a CVPack if the cv function + in xgboost is being used. + epoch : + The current training iteration. + evals_log : + A dictionary containing the evaluation history: + + .. code-block:: python + + {"data_name": {"metric_name": [0.5, ...]}} + + """ + return False + + +def _aggcv(rlist: List[str]) -> List[Tuple[str, float, float]]: + # pylint: disable=invalid-name, too-many-locals + """Aggregate cross-validation results.""" + cvmap: Dict[Tuple[int, str], List[float]] = {} + idx = rlist[0].split()[0] + for line in rlist: + arr: List[str] = line.split() + assert idx == arr[0] + for metric_idx, it in enumerate(arr[1:]): + if not isinstance(it, str): + it = it.decode() + k, v = it.split(":") + if (metric_idx, k) not in cvmap: + cvmap[(metric_idx, k)] = [] + cvmap[(metric_idx, k)].append(float(v)) + msg = idx + results = [] + for (_, name), s in sorted(cvmap.items(), key=lambda x: x[0][0]): + as_arr = numpy.array(s) + if not isinstance(msg, str): + msg = msg.decode() + mean, std = numpy.mean(as_arr), numpy.std(as_arr) + results.extend([(name, mean, std)]) + return results + + +# allreduce type +_ART = TypeVar("_ART") + + +def _allreduce_metric(score: _ART) -> _ART: + """Helper function for computing customized metric in distributed + environment. Not strictly correct as many functions don't use mean value + as final result. + + """ + world = collective.get_world_size() + assert world != 0 + if world == 1: + return score + if isinstance(score, tuple): # has mean and stdv + raise ValueError( + "xgboost.cv function should not be used in distributed environment." + ) + arr = numpy.array([score]) + arr = collective.allreduce(arr, collective.Op.SUM) / world + return arr[0] + + +class CallbackContainer: + """A special internal callback for invoking a list of other callbacks. + + .. versionadded:: 1.3.0 + + """ + + def __init__( + self, + callbacks: Sequence[TrainingCallback], + metric: Optional[Callable] = None, + output_margin: bool = True, + is_cv: bool = False, + ) -> None: + self.callbacks = list(dict.fromkeys(callbacks)) + for cb in callbacks: + if not isinstance(cb, TrainingCallback): + raise TypeError("callback must be an instance of `TrainingCallback`.") + + msg = ( + "metric must be callable object for monitoring. For builtin metrics" + ", passing them in training parameter invokes monitor automatically." + ) + if metric is not None and not callable(metric): + raise TypeError(msg) + + self.metric = metric + self.history: EvalsLog = collections.OrderedDict() + self._output_margin = output_margin + self.is_cv = is_cv + + if self.is_cv: + self.aggregated_cv: Optional[list[tuple[str, float, float]]] = None + + def before_training(self, model: _Model) -> _Model: + """Function called before training.""" + for c in self.callbacks: + model = c.before_training(model=model) + msg = "before_training should return the model" + if self.is_cv: + assert isinstance(model.cvfolds, list), msg + else: + assert isinstance(model, Booster), msg + return model + + def after_training(self, model: _Model) -> _Model: + """Function called after training.""" + for c in self.callbacks: + model = c.after_training(model=model) + msg = "after_training should return the model" + if self.is_cv: + assert isinstance(model.cvfolds, list), msg + else: + assert isinstance(model, Booster), msg + + return model + + def before_iteration( + self, + model: _Model, + epoch: int, + dtrain: DMatrix, + evals: Optional[List[Tuple[DMatrix, str]]], + ) -> bool: + """Function called before training iteration.""" + return any( + c.before_iteration(model, epoch, self.history) for c in self.callbacks + ) + + def _update_history( + self, + score: Union[List[Tuple[str, float]], List[Tuple[str, float, float]]], + epoch: int, + ) -> None: + for d in score: + name: str = d[0] + s: float = d[1] + if self.is_cv: + std = float(cast(Tuple[str, float, float], d)[2]) + x: _Score = (s, std) + else: + x = s + splited_names = name.split("-") + data_name = splited_names[0] + metric_name = "-".join(splited_names[1:]) + x = _allreduce_metric(x) + if data_name not in self.history: + self.history[data_name] = collections.OrderedDict() + data_history = self.history[data_name] + if metric_name not in data_history: + data_history[metric_name] = cast(_ScoreList, []) + metric_history = data_history[metric_name] + if self.is_cv: + cast(List[Tuple[float, float]], metric_history).append( + cast(Tuple[float, float], x) + ) + else: + cast(List[float], metric_history).append(cast(float, x)) + + def after_iteration( + self, + model: _Model, + epoch: int, + dtrain: DMatrix, + evals: Optional[List[Tuple[DMatrix, str]]], + ) -> bool: + """Function called after training iteration.""" + if self.is_cv: + scores = model.eval(epoch, self.metric, self._output_margin) + scores = _aggcv(scores) + self.aggregated_cv = scores + self._update_history(scores, epoch) + else: + evals = [] if evals is None else evals + for _, name in evals: + assert name.find("-") == -1, "Dataset name should not contain `-`" + score: str = model.eval_set(evals, epoch, self.metric, self._output_margin) + metric_score = _parse_eval_str(score) + self._update_history(metric_score, epoch) + ret = any(c.after_iteration(model, epoch, self.history) for c in self.callbacks) + return ret + + +class LearningRateScheduler(TrainingCallback): + """Callback function for scheduling learning rate. + + .. versionadded:: 1.3.0 + + Parameters + ---------- + + learning_rates : + If it's a callable object, then it should accept an integer parameter + `epoch` and returns the corresponding learning rate. Otherwise it + should be a sequence like list or tuple with the same size of boosting + rounds. + + """ + + def __init__( + self, learning_rates: Union[Callable[[int], float], Sequence[float]] + ) -> None: + if not callable(learning_rates) and not isinstance( + learning_rates, collections.abc.Sequence + ): + raise TypeError( + "Invalid learning rates, expecting callable or sequence, got: " + f"{type(learning_rates)}" + ) + + if callable(learning_rates): + self.learning_rates = learning_rates + else: + self.learning_rates = lambda epoch: cast(Sequence, learning_rates)[epoch] + super().__init__() + + def after_iteration(self, model: _Model, epoch: int, evals_log: EvalsLog) -> bool: + model.set_param("learning_rate", self.learning_rates(epoch)) + return False + + +# pylint: disable=too-many-instance-attributes +class EarlyStopping(TrainingCallback): + """Callback function for early stopping + + .. versionadded:: 1.3.0 + + Parameters + ---------- + rounds : + Early stopping rounds. + metric_name : + Name of metric that is used for early stopping. + data_name : + Name of dataset that is used for early stopping. + maximize : + Whether to maximize evaluation metric. None means auto (discouraged). + save_best : + Whether training should return the best model or the last model. If set to + `True`, it will only keep the boosting rounds up to the detected best iteration, + discarding the ones that come after. This is only supported with tree methods + (not `gblinear`). Also, the `cv` function doesn't return a model, the parameter + is not applicable. + min_delta : + + .. versionadded:: 1.5.0 + + Minimum absolute change in score to be qualified as an improvement. + + Examples + -------- + + .. code-block:: python + + es = xgboost.callback.EarlyStopping( + rounds=2, + min_delta=1e-3, + save_best=True, + maximize=False, + data_name="validation_0", + metric_name="mlogloss", + ) + clf = xgboost.XGBClassifier(tree_method="hist", device="cuda", callbacks=[es]) + + X, y = load_digits(return_X_y=True) + clf.fit(X, y, eval_set=[(X, y)]) + """ + + # pylint: disable=too-many-arguments + @_deprecate_positional_args + def __init__( + self, + *, + rounds: int, + metric_name: Optional[str] = None, + data_name: Optional[str] = None, + maximize: Optional[bool] = None, + save_best: Optional[bool] = False, + min_delta: float = 0.0, + ) -> None: + self.data = data_name + self.metric_name = metric_name + self.rounds = rounds + self.save_best = save_best + self.maximize = maximize + self.stopping_history: EvalsLog = {} + self._min_delta = min_delta + if self._min_delta < 0: + raise ValueError("min_delta must be greater or equal to 0.") + + self.current_rounds: int = 0 + self.best_scores: dict = {} + self.starting_round: int = 0 + super().__init__() + + def before_training(self, model: _Model) -> _Model: + self.starting_round = model.num_boosted_rounds() + if not isinstance(model, Booster) and self.save_best: + raise ValueError( + "`save_best` is not applicable to the `cv` function as it doesn't" + " return a model." + ) + return model + + def _update_rounds( + self, *, score: _Score, name: str, metric: str, model: _Model, epoch: int + ) -> bool: + def get_s(value: _Score) -> float: + """get score if it's cross validation history.""" + return value[0] if isinstance(value, tuple) else value + + def maximize(new: _Score, best: _Score) -> bool: + """New score should be greater than the old one.""" + return numpy.greater(get_s(new) - self._min_delta, get_s(best)) + + def minimize(new: _Score, best: _Score) -> bool: + """New score should be lesser than the old one.""" + return numpy.greater(get_s(best) - self._min_delta, get_s(new)) + + if self.maximize is None: + # Just to be compatibility with old behavior before 1.3. We should let + # user to decide. + maximize_metrics = ( + "auc", + "aucpr", + "pre", + "pre@", + "map", + "ndcg", + "auc@", + "aucpr@", + "map@", + "ndcg@", + ) + if metric != "mape" and any(metric.startswith(x) for x in maximize_metrics): + self.maximize = True + else: + self.maximize = False + + if self.maximize: + improve_op = maximize + else: + improve_op = minimize + + if not self.stopping_history: # First round + self.current_rounds = 0 + self.stopping_history[name] = {} + self.stopping_history[name][metric] = cast(_ScoreList, [score]) + self.best_scores[name] = {} + self.best_scores[name][metric] = [score] + model.set_attr(best_score=str(get_s(score)), best_iteration=str(epoch)) + elif not improve_op(score, self.best_scores[name][metric][-1]): + # Not improved + self.stopping_history[name][metric].append(score) # type: ignore[arg-type] + self.current_rounds += 1 + else: # Improved + self.stopping_history[name][metric].append(score) # type: ignore[arg-type] + self.best_scores[name][metric].append(score) + record = self.stopping_history[name][metric][-1] + model.set_attr(best_score=str(get_s(record)), best_iteration=str(epoch)) + self.current_rounds = 0 # reset + + if self.current_rounds >= self.rounds: + # Should stop + return True + return False + + def after_iteration(self, model: _Model, epoch: int, evals_log: EvalsLog) -> bool: + epoch += self.starting_round # training continuation + msg = "Must have at least 1 validation dataset for early stopping." + if len(evals_log.keys()) < 1: + raise ValueError(msg) + + # Get data name + if self.data: + data_name = self.data + else: + # Use the last one as default. + data_name = list(evals_log.keys())[-1] + if data_name not in evals_log: + raise ValueError(f"No dataset named: {data_name}") + + if not isinstance(data_name, str): + raise TypeError( + f"The name of the dataset should be a string. Got: {type(data_name)}" + ) + data_log = evals_log[data_name] + + # Get metric name + if self.metric_name: + metric_name = self.metric_name + else: + # Use last metric by default. + metric_name = list(data_log.keys())[-1] + if metric_name not in data_log: + raise ValueError(f"No metric named: {metric_name}") + + # The latest score + score = data_log[metric_name][-1] + return self._update_rounds( + score=score, name=data_name, metric=metric_name, model=model, epoch=epoch + ) + + def after_training(self, model: _Model) -> _Model: + if not self.save_best: + return model + + try: + best_iteration = model.best_iteration + best_score = model.best_score + assert best_iteration is not None and best_score is not None + model = model[: best_iteration + 1] + model.best_iteration = best_iteration + model.best_score = best_score + except XGBoostError as e: + raise XGBoostError( + "`save_best` is not applicable to the current booster" + ) from e + + return model + + +class EvaluationMonitor(TrainingCallback): + """Print the evaluation result at each iteration. + + .. versionadded:: 1.3.0 + + Parameters + ---------- + + rank : + Which worker should be used for printing the result. + period : + How many epoches between printing. + show_stdv : + Used in cv to show standard deviation. Users should not specify it. + logger : + A callable used for logging evaluation result. + + """ + + def __init__( + self, + rank: int = 0, + period: int = 1, + show_stdv: bool = False, + logger: Callable[[str], None] = collective.communicator_print, + ): + self.printer_rank = rank + self.show_stdv = show_stdv + self.period = period + self._logger = logger + assert period > 0 + # last error message, useful when early stopping and period are used together. + self._latest: Optional[str] = None + super().__init__() + + def _fmt_metric( + self, data: str, metric: str, score: float, std: Optional[float] + ) -> str: + if std is not None and self.show_stdv: + msg = f"\t{data + '-' + metric}:{score:.5f}+{std:.5f}" + else: + msg = f"\t{data + '-' + metric}:{score:.5f}" + return msg + + def after_iteration(self, model: _Model, epoch: int, evals_log: EvalsLog) -> bool: + if not evals_log: + return False + + msg: str = f"[{epoch}]" + if collective.get_rank() == self.printer_rank: + for data, metric in evals_log.items(): + for metric_name, log in metric.items(): + stdv: Optional[float] = None + if isinstance(log[-1], tuple): + score = log[-1][0] + stdv = log[-1][1] + else: + score = log[-1] + msg += self._fmt_metric(data, metric_name, score, stdv) + msg += "\n" + + if (epoch % self.period) == 0 or self.period == 1: + self._logger(msg) + self._latest = None + else: + # There is skipped message + self._latest = msg + return False + + def after_training(self, model: _Model) -> _Model: + if collective.get_rank() == self.printer_rank and self._latest is not None: + self._logger(self._latest) + return model + + +class TrainingCheckPoint(TrainingCallback): + """Checkpointing operation. Users are encouraged to create their own callbacks for + checkpoint as XGBoost doesn't handle distributed file systems. When checkpointing on + distributed systems, be sure to know the rank of the worker to avoid multiple + workers checkpointing to the same place. + + .. versionadded:: 1.3.0 + + Since XGBoost 2.1.0, the default format is changed to UBJSON. + + Parameters + ---------- + + directory : + Output model directory. + name : + pattern of output model file. Models will be saved as name_0.ubj, name_1.ubj, + name_2.ubj .... + as_pickle : + When set to True, all training parameters will be saved in pickle format, + instead of saving only the model. + interval : + Interval of checkpointing. Checkpointing is slow so setting a larger number can + reduce performance hit. + + """ + + default_format = "ubj" + + def __init__( + self, + directory: Union[str, os.PathLike], + name: str = "model", + as_pickle: bool = False, + interval: int = 100, + ) -> None: + self._path = os.fspath(directory) + self._name = name + self._as_pickle = as_pickle + self._iterations = interval + self._epoch = 0 # counter for iterval + self._start = 0 # beginning iteration + super().__init__() + + def before_training(self, model: _Model) -> _Model: + self._start = model.num_boosted_rounds() + return model + + def after_iteration(self, model: _Model, epoch: int, evals_log: EvalsLog) -> bool: + if self._epoch == self._iterations: + path = os.path.join( + self._path, + self._name + + "_" + + (str(epoch + self._start)) + + (".pkl" if self._as_pickle else f".{self.default_format}"), + ) + self._epoch = 0 # reset counter + if collective.get_rank() == 0: + # checkpoint using the first worker + if self._as_pickle: + with open(path, "wb") as fd: + pickle.dump(model, fd) + else: + model.save_model(path) + self._epoch += 1 + return False diff --git a/v12x/pylib/xgboost/collective.py b/v12x/pylib/xgboost/collective.py new file mode 100644 index 0000000000000000000000000000000000000000..c7490fdc00feadac7bf121862605b3a68a969081 --- /dev/null +++ b/v12x/pylib/xgboost/collective.py @@ -0,0 +1,390 @@ +"""XGBoost collective communication related API.""" + +import ctypes +import logging +import os +import pickle +from dataclasses import dataclass +from enum import IntEnum, unique +from typing import Any, Callable, Dict, Optional, TypeAlias, Union + +import numpy as np + +from ._typing import _T +from .core import _LIB, _check_call, build_info, c_str, make_jcargs, py_str + +LOGGER = logging.getLogger("[xgboost.collective]") + + +_Conf: TypeAlias = Dict[str, Union[int, str]] +_ArgVals: TypeAlias = Optional[Union[int, str]] +_Args: TypeAlias = Dict[str, _ArgVals] + + +@dataclass +class Config: + """User configuration for the communicator context. This is used for easier + integration with distributed frameworks. Users of the collective module can pass the + parameters directly into tracker and the communicator. + + .. versionadded:: 3.0 + + Attributes + ---------- + retry : See `dmlc_retry` in :py:meth:`init`. + + timeout : + See `dmlc_timeout` in :py:meth:`init`. This is only used for communicators, not + the tracker. They are different parameters since the timeout for tracker limits + only the time for starting and finalizing the communication group, whereas the + timeout for communicators limits the time used for collective operations, like + :py:meth:`allreduce`. + + tracker_host_ip : See :py:class:`~xgboost.tracker.RabitTracker`. + + tracker_port : See :py:class:`~xgboost.tracker.RabitTracker`. + + tracker_timeout : See :py:class:`~xgboost.tracker.RabitTracker`. + + worker_port : + + The port each worker listens to for peer-to-peer connections. By default, + workers use an available port assigned by the OS. This option can be used in + restricted network environments where only specific ports are open. + + This can be an integer for a fixed port used by all workers, or a callback + function that takes no arguments and returns a port number. The callback is + invoked per-worker at the worker side. + + .. note:: + + The option does not affect the NCCL communicator group, which must be + configured via NCCL's own environment variables. + + """ + + retry: Optional[int] = None + timeout: Optional[int] = None + + tracker_host_ip: Optional[str] = None + tracker_port: Optional[int] = None + tracker_timeout: Optional[int] = None + + worker_port: Optional[Union[Callable[[], int], int]] = None + + def update_worker_args(self, args: _Conf) -> _Conf: + """Worker side arguments resolution.""" + if self.worker_port is None: + return args + if callable(self.worker_port): + args["dmlc_worker_port"] = self.worker_port() + else: + args["dmlc_worker_port"] = self.worker_port + return args + + def get_comm_config(self, args: _Conf) -> _Conf: + """Update the arguments for the communicator.""" + if self.retry is not None: + args["dmlc_retry"] = self.retry + if self.timeout is not None: + args["dmlc_timeout"] = self.timeout + return args + + +def init(**args: _ArgVals) -> None: + """Initialize the collective library with arguments. + + Parameters + ---------- + args : + Keyword arguments representing the parameters and their values. + + Accepted parameters: + - dmlc_communicator: The type of the communicator. + * rabit: Use Rabit. This is the default if the type is unspecified. + * federated: Use the gRPC interface for Federated Learning. + + Only applicable to the Rabit communicator: + - dmlc_tracker_uri: Hostname of the tracker. + - dmlc_tracker_port: Port number of the tracker. + - dmlc_task_id: ID of the current task, can be used to obtain deterministic + - dmlc_retry: The number of retry when handling network errors. + - dmlc_timeout: Timeout in seconds. + - dmlc_nccl_path: Path to load (dlopen) nccl for GPU-based communication. + + Only applicable to the Federated communicator: + - federated_server_address: Address of the federated server. + - federated_world_size: Number of federated workers. + - federated_rank: Rank of the current worker. + - federated_server_cert: Server certificate file path. Only needed for the SSL + mode. + - federated_client_key: Client key file path. Only needed for the SSL mode. + - federated_client_cert: Client certificate file path. Only needed for the SSL + mode. + + Use upper case for environment variables, use lower case for runtime + configuration. + + """ + _check_call(_LIB.XGCommunicatorInit(make_jcargs(**args))) + + +def finalize() -> None: + """Finalize the communicator.""" + _check_call(_LIB.XGCommunicatorFinalize()) + + +def get_rank() -> int: + """Get rank of current process. + + Returns + ------- + rank : int + Rank of current process. + """ + ret = _LIB.XGCommunicatorGetRank() + return ret + + +def get_world_size() -> int: + """Get total number workers. + + Returns + ------- + n : + Total number of process. + """ + ret = _LIB.XGCommunicatorGetWorldSize() + return ret + + +def is_distributed() -> bool: + """If the collective communicator is distributed.""" + is_dist = _LIB.XGCommunicatorIsDistributed() + return bool(is_dist) + + +def communicator_print(msg: Any) -> None: + """Print message to the communicator. + + This function can be used to communicate the information of + the progress to the communicator. + + Parameters + ---------- + msg : str + The message to be printed to the communicator. + """ + if not isinstance(msg, str): + msg = str(msg) + is_dist = _LIB.XGCommunicatorIsDistributed() + if is_dist != 0: + _check_call(_LIB.XGCommunicatorPrint(c_str(msg.strip()))) + else: + print(msg.strip(), flush=True) + + +def get_processor_name() -> str: + """Get the processor name. + + Returns + ------- + name : + The name of processor(host) + """ + name_str = ctypes.c_char_p() + _check_call(_LIB.XGCommunicatorGetProcessorName(ctypes.byref(name_str))) + value = name_str.value + return py_str(value) + + +def broadcast(data: _T, root: int) -> _T: + """Broadcast object from one node to all other nodes. + + Parameters + ---------- + data : any type that can be pickled + Input data, if current rank does not equal root, this can be None + root : int + Rank of the node to broadcast data from. + + Returns + ------- + object : int + the result of broadcast. + """ + rank = get_rank() + length = ctypes.c_ulong() + if root == rank: + assert data is not None, "need to pass in data when broadcasting" + s = pickle.dumps(data, protocol=pickle.HIGHEST_PROTOCOL) + length.value = len(s) + # Run first broadcast + _check_call( + _LIB.XGCommunicatorBroadcast( + ctypes.byref(length), ctypes.sizeof(ctypes.c_ulong), root + ) + ) + if root != rank: + dptr = (ctypes.c_char * length.value)() + # run second + _check_call( + _LIB.XGCommunicatorBroadcast( + ctypes.cast(dptr, ctypes.c_void_p), length.value, root + ) + ) + data = pickle.loads(dptr.raw) + del dptr + else: + _check_call( + _LIB.XGCommunicatorBroadcast( + ctypes.cast(ctypes.c_char_p(s), ctypes.c_void_p), length.value, root + ) + ) + del s + return data + + +# enumeration of dtypes +def _map_dtype(dtype: np.dtype) -> int: + dtype_map = { + np.dtype("float16"): 0, + np.dtype("float32"): 1, + np.dtype("float64"): 2, + np.dtype("int8"): 4, + np.dtype("int16"): 5, + np.dtype("int32"): 6, + np.dtype("int64"): 7, + np.dtype("uint8"): 8, + np.dtype("uint16"): 9, + np.dtype("uint32"): 10, + np.dtype("uint64"): 11, + } + try: + dtype_map.update({np.dtype("float128"): 3}) + except TypeError: # float128 doesn't exist on the system + pass + + if dtype not in dtype_map: + raise TypeError(f"data type {dtype} is not supported on the current platform.") + + return dtype_map[dtype] + + +@unique +class Op(IntEnum): + """Supported operations for allreduce.""" + + MAX = 0 + MIN = 1 + SUM = 2 + BITWISE_AND = 3 + BITWISE_OR = 4 + BITWISE_XOR = 5 + + +def allreduce(data: np.ndarray, op: Op) -> np.ndarray: + """Perform allreduce, return the result. + + Parameters + ---------- + data : + Input data. + op : + Reduction operator. + + Returns + ------- + result : + The result of allreduce, have same shape as data + + Notes + ----- + This function is not thread-safe. + """ + if not isinstance(data, np.ndarray): + raise TypeError("allreduce only takes in numpy.ndarray") + buf = data.ravel().copy() + _check_call( + _LIB.XGCommunicatorAllreduce( + buf.ctypes.data_as(ctypes.c_void_p), + buf.size, + _map_dtype(buf.dtype), + int(op), + ) + ) + # `buf` is a flattened copy used for the underlying C call; reshape it back + # to the input's original shape so the result actually "has the same shape + # as data", as documented above, instead of always returning a 1-D array. + return buf.reshape(data.shape) + + +def signal_error() -> None: + """Kill the process.""" + _check_call(_LIB.XGCommunicatorSignalError()) + + +def _find_nccl() -> Optional[str]: + from nvidia.nccl import lib + + # There are two versions of nvidia-nccl, one is from PyPI, another one from + # nvidia-pyindex. We support only the first one as the second one is too old (2.9.8 + # as of writing). + # + # nccl 2.28 doesn't have the __file__ attribute, we use the namespace path instead. + if lib.__file__ is not None: + dirname: Optional[str] = os.path.dirname(lib.__file__) + elif hasattr(lib, "__path__") and len(lib.__path__) > 0: + dirname = lib.__path__[0] + else: + dirname = None + if not dirname: + return None + + # Find the first shared object in the lib directory. + files = os.listdir(dirname) + if not files: + return None + + libname: Optional[str] = None + for name in files: + if name.startswith("libnccl.so"): + libname = name + break + + if libname is not None: + path = os.path.join(dirname, libname) + return path + return None + + +class CommunicatorContext: + """A context controlling collective communicator initialization and finalization.""" + + def __init__(self, **args: _ArgVals) -> None: + self.args = args + key = "dmlc_nccl_path" + if args.get(key, None) is not None: + return + + binfo = build_info() + if not binfo["USE_DLOPEN_NCCL"]: + return + + try: + # PyPI package of NCCL. + path = _find_nccl() + if path: + self.args[key] = path + except ImportError: + pass + + def __enter__(self) -> _Args: + init(**self.args) + assert is_distributed() + LOGGER.debug("-------------- communicator say hello ------------------") + return self.args + + def __exit__(self, *args: Any) -> None: + finalize() + LOGGER.debug("--------------- communicator say bye ------------------") diff --git a/v12x/pylib/xgboost/compat.py b/v12x/pylib/xgboost/compat.py new file mode 100644 index 0000000000000000000000000000000000000000..0c73a0fb15744df3b74a7305dcbde4c421bfb7aa --- /dev/null +++ b/v12x/pylib/xgboost/compat.py @@ -0,0 +1,277 @@ +# pylint: disable=unused-import +"""For compatibility and optional dependencies.""" + +import functools +import importlib.util +import logging +import types +from typing import TYPE_CHECKING, Any, Sequence, TypeGuard, cast + +import numpy as np + +from ._typing import _T, DataType + +if TYPE_CHECKING: + import pandas as pd + import pyarrow as pa + + +def py_str(x: bytes | None) -> str: + """convert c string back to python string""" + assert x is not None # ctypes might return None + return x.decode("utf-8") # type: ignore[union-attr] + + +def lazy_isinstance(instance: Any, module: str, name: str) -> bool: + """Use string representation to identify a type.""" + + # Notice, we use .__class__ as opposed to type() in order + # to support object proxies such as weakref.proxy + cls = instance.__class__ + is_same_module = cls.__module__ == module + has_same_name = cls.__name__ == name + return is_same_module and has_same_name + + +# sklearn +try: + from sklearn import __version__ as _sklearn_version + from sklearn.base import BaseEstimator as XGBModelBase + from sklearn.base import ClassifierMixin as XGBClassifierBase + from sklearn.base import RegressorMixin as XGBRegressorBase + from sklearn.model_selection import StratifiedKFold as XGBStratifiedKFold + + # sklearn.utils Tags types can be imported unconditionally once + # xgboost's minimum scikit-learn version is 1.6 or higher + try: + from sklearn.utils import Tags as _sklearn_Tags + except ImportError: + _sklearn_Tags = object + + SKLEARN_INSTALLED = True + +except ImportError: + SKLEARN_INSTALLED = False + + # used for compatibility without sklearn + class XGBModelBase: # type: ignore[no-redef] + """Dummy class for sklearn.base.BaseEstimator.""" + + class XGBClassifierBase: # type: ignore[no-redef] + """Dummy class for sklearn.base.ClassifierMixin.""" + + class XGBRegressorBase: # type: ignore[no-redef] + """Dummy class for sklearn.base.RegressorMixin.""" + + XGBStratifiedKFold = None + + _sklearn_Tags = object + _sklearn_version = object + + +_logger = logging.getLogger(__name__) + + +@functools.cache +def is_cudf_available() -> bool: + """Check cuDF package available or not""" + if importlib.util.find_spec("cudf") is None: + return False + try: + import cudf + + return True + except ImportError: + _logger.exception("Importing cuDF failed, use DMatrix instead of QDM") + return False + + +@functools.cache +def is_cupy_available() -> bool: + """Check cupy package available or not""" + if importlib.util.find_spec("cupy") is None: + return False + try: + import cupy + + return True + except ImportError: + return False + + +@functools.cache +def import_cupy() -> types.ModuleType: + """Import cupy.""" + if not is_cupy_available(): + raise ImportError("`cupy` is required for handling CUDA buffer.") + + import cupy + + return cupy + + +@functools.cache +def is_pyarrow_available() -> bool: + """Check pyarrow package available or not""" + if importlib.util.find_spec("pyarrow") is None: + return False + return True + + +@functools.cache +def import_pyarrow() -> types.ModuleType: + """Import pyarrow with memory cache.""" + import pyarrow as pa + + return pa + + +@functools.cache +def import_pandas() -> types.ModuleType: + """Import pandas with memory cache.""" + import pandas as pd + + return pd + + +@functools.cache +def import_polars() -> types.ModuleType: + """Import polars with memory cache.""" + import polars as pl + + return pl + + +@functools.cache +def is_pandas_available() -> bool: + """Check the pandas package is available or not.""" + if importlib.util.find_spec("pandas") is None: + return False + return True + + +try: + import scipy.sparse as scipy_sparse +except ImportError: + scipy_sparse = False + + +def _is_polars_lazyframe(data: DataType) -> bool: + return lazy_isinstance(data, "polars.lazyframe.frame", "LazyFrame") + + +def _is_polars_series(data: DataType) -> bool: + return lazy_isinstance(data, "polars.series.series", "Series") + + +def _is_polars(data: DataType) -> bool: + lf = _is_polars_lazyframe(data) + df = lazy_isinstance(data, "polars.dataframe.frame", "DataFrame") + return lf or df + + +def _is_arrow(data: DataType) -> TypeGuard["pa.Table"]: + return lazy_isinstance(data, "pyarrow.lib", "Table") + + +def _is_cudf_df(data: DataType) -> bool: + return lazy_isinstance(data, "cudf.core.dataframe", "DataFrame") + + +def _is_cudf_ser(data: DataType) -> bool: + return lazy_isinstance(data, "cudf.core.series", "Series") + + +def _is_cudf_pandas(data: DataType) -> bool: + """Must go before both pandas and cudf checks.""" + return (_is_pandas_df(data) or _is_pandas_series(data)) and lazy_isinstance( + type(data), "cudf.pandas.fast_slow_proxy", "_FastSlowProxyMeta" + ) + + +def _is_pandas_df(data: DataType) -> TypeGuard["pd.DataFrame"]: + return lazy_isinstance(data, "pandas.core.frame", "DataFrame") or lazy_isinstance( + data, "pandas", "DataFrame" + ) + + +def _is_pandas_series(data: DataType) -> TypeGuard["pd.Series"]: + return lazy_isinstance(data, "pandas.core.series", "Series") or lazy_isinstance( + data, "pandas", "Series" + ) + + +def _is_modin_df(data: DataType) -> bool: + return lazy_isinstance(data, "modin.pandas.dataframe", "DataFrame") + + +def _is_modin_series(data: DataType) -> bool: + return lazy_isinstance(data, "modin.pandas.series", "Series") + + +def is_dataframe(data: DataType) -> bool: + """Whether the input is a dataframe. Currently supported dataframes: + + - pandas + - cudf + - cudf.pandas + - polars + - pyarrow + - modin + + + """ + return any( + p(data) + for p in ( + _is_polars, + _is_polars_series, + _is_arrow, + _is_cudf_df, + _is_cudf_ser, + _is_cudf_pandas, + _is_pandas_df, + _is_pandas_series, + _is_modin_df, + _is_modin_series, + ) + ) + + +def _is_cupy_alike(data: DataType) -> bool: + return hasattr(data, "__cuda_array_interface__") + + +def concat(value: Sequence[_T]) -> _T: # pylint: disable=too-many-return-statements + """Concatenate row-wise.""" + if isinstance(value[0], np.ndarray): + value_arr = cast(Sequence[np.ndarray], value) + return cast(_T, np.concatenate(value_arr, axis=0)) + if scipy_sparse and isinstance(value[0], scipy_sparse.csr_matrix): + return scipy_sparse.vstack(value, format="csr") + if scipy_sparse and isinstance(value[0], scipy_sparse.csc_matrix): + return scipy_sparse.vstack(value, format="csc") + if scipy_sparse and isinstance(value[0], scipy_sparse.spmatrix): + # other sparse format will be converted to CSR. + return scipy_sparse.vstack(value, format="csr") + if _is_pandas_df(value[0]) or _is_pandas_series(value[0]): + from pandas import concat as pd_concat + + return pd_concat(value, axis=0) + if lazy_isinstance(value[0], "cudf.core.dataframe", "DataFrame") or lazy_isinstance( + value[0], "cudf.core.series", "Series" + ): + from cudf import concat as CUDF_concat + + return CUDF_concat(value, axis=0) + if _is_cupy_alike(value[0]): + import cupy + + # pylint: disable=c-extension-no-member,no-member + d = cupy.cuda.runtime.getDevice() + for v in value: + arr = cast(cupy.ndarray, v) + d_v = arr.device.id + assert d_v == d, "Concatenating arrays on different devices." + return cupy.concatenate(value, axis=0) + raise TypeError(f"Unknown type: {type(value[0])}") diff --git a/v12x/pylib/xgboost/config.py b/v12x/pylib/xgboost/config.py new file mode 100644 index 0000000000000000000000000000000000000000..c26241aa5b336349ba720bf3e2e87088c96fa6b1 --- /dev/null +++ b/v12x/pylib/xgboost/config.py @@ -0,0 +1,187 @@ +# pylint: disable=missing-function-docstring +"""Global configuration for XGBoost""" + +import ctypes +import json +from contextlib import contextmanager +from functools import wraps +from typing import Any, Callable, Dict, Iterator, Optional, cast + +from ._typing import _F +from .core import _LIB, _check_call, c_str, py_str + + +def config_doc( + *, + header: Optional[str] = None, + extra_note: Optional[str] = None, + parameters: Optional[str] = None, + returns: Optional[str] = None, + see_also: Optional[str] = None, +) -> Callable[[_F], _F]: + """Decorator to format docstring for config functions. + + Parameters + ---------- + header: str + An introducion to the function + extra_note: str + Additional notes + parameters: str + Parameters of the function + returns: str + Return value + see_also: str + Related functions + """ + + doc_template = """ + {header} + + Global configuration consists of a collection of parameters that can be applied in the + global scope. See :ref:`global_config` for the full list of parameters supported in + the global configuration. + + {extra_note} + + .. versionadded:: 1.4.0 + """ + + common_example = """ + Example + ------- + + .. code-block:: python + + import xgboost as xgb + + # Show all messages, including ones pertaining to debugging + xgb.set_config(verbosity=2) + + # Get current value of global configuration + # This is a dict containing all parameters in the global configuration, + # including 'verbosity' + config = xgb.get_config() + assert config['verbosity'] == 2 + + # Example of using the context manager xgb.config_context(). + # The context manager will restore the previous value of the global + # configuration upon exiting. + with xgb.config_context(verbosity=0): + # Suppress warning caused by model generated with XGBoost version < 1.0.0 + bst = xgb.Booster(model_file='./old_model.bin') + assert xgb.get_config()['verbosity'] == 2 # old value restored + + Nested configuration context is also supported: + + Example + ------- + + .. code-block:: python + + with xgb.config_context(verbosity=3): + assert xgb.get_config()["verbosity"] == 3 + with xgb.config_context(verbosity=2): + assert xgb.get_config()["verbosity"] == 2 + + xgb.set_config(verbosity=2) + assert xgb.get_config()["verbosity"] == 2 + with xgb.config_context(verbosity=3): + assert xgb.get_config()["verbosity"] == 3 + """ + + def none_to_str(value: Optional[str]) -> str: + return "" if value is None else value + + def config_doc_decorator(func: _F) -> _F: + func.__doc__ = ( + doc_template.format( + header=none_to_str(header), extra_note=none_to_str(extra_note) + ) + + none_to_str(parameters) + + none_to_str(returns) + + none_to_str(common_example) + + none_to_str(see_also) + ) + + @wraps(func) + def wrap(*args: Any, **kwargs: Any) -> Any: + return func(*args, **kwargs) + + return cast(_F, wrap) + + return config_doc_decorator + + +@config_doc( + header=""" + Set global configuration. + """, + parameters=""" + Parameters + ---------- + new_config: Dict[str, Any] + Keyword arguments representing the parameters and their values + """, +) +def set_config(**new_config: Any) -> None: + not_none = {} + for k, v in new_config.items(): + if v is not None: + not_none[k] = v + config = json.dumps(not_none) + _check_call(_LIB.XGBSetGlobalConfig(c_str(config))) + + +@config_doc( + header=""" + Get current values of the global configuration. + """, + returns=""" + Returns + ------- + args: Dict[str, Any] + The list of global parameters and their values + """, +) +def get_config() -> Dict[str, Any]: + config_str = ctypes.c_char_p() + _check_call(_LIB.XGBGetGlobalConfig(ctypes.byref(config_str))) + value = config_str.value + assert value + config = json.loads(py_str(value)) + return config + + +@contextmanager +@config_doc( + header=""" + Context manager for global XGBoost configuration. + """, + parameters=""" + Parameters + ---------- + new_config: Dict[str, Any] + Keyword arguments representing the parameters and their values + """, + extra_note=""" + .. note:: + + All settings, not just those presently modified, will be returned to their + previous values when the context manager is exited. This is not thread-safe. + """, + see_also=""" + See Also + -------- + set_config: Set global XGBoost configuration + get_config: Get current values of the global configuration + """, +) +def config_context(**new_config: Any) -> Iterator[None]: + old_config = get_config().copy() + set_config(**new_config) + + try: + yield + finally: + set_config(**old_config) diff --git a/v12x/pylib/xgboost/core.py b/v12x/pylib/xgboost/core.py new file mode 100644 index 0000000000000000000000000000000000000000..520d703ff822753d220b35374f0a534534dd2935 --- /dev/null +++ b/v12x/pylib/xgboost/core.py @@ -0,0 +1,3399 @@ +# pylint: disable=too-many-arguments, too-many-branches +# pylint: disable=too-many-lines, too-many-locals +"""Core XGBoost Library.""" + +from __future__ import annotations + +import copy +import ctypes +import json +import math +import os +import re +import sys +import warnings +import weakref +from abc import ABC, abstractmethod +from collections.abc import Mapping +from functools import wraps +from inspect import Parameter, signature +from types import EllipsisType +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Dict, + Generator, + Iterable, + List, + Optional, + Sequence, + Tuple, + Type, + TypeGuard, + TypeVar, + Union, + cast, +) + +import numpy as np +import scipy.sparse + +from ._c_api import ( + _LIB, + _check_call, + c_str, + from_cstr_to_pystr, + from_pystr_to_cstr, + make_jcargs, +) +from ._c_api import XGBoostError as _XGBoostError +from ._data_utils import ( + Categories, + TransformedDf, + _ensure_np_dtype, + array_interface, + cuda_array_interface, + from_array_interface, + make_array_interface, +) +from ._typing import ( + _T, + ArrayLike, + ArrowCatList, + BoosterParam, + CFloatPtr, + CNumeric, + CNumericPtr, + CStrPtr, + CTypeT, + DataType, + FeatureInfo, + FeatureNames, + FeatureTypes, + Integer, + IterationRange, + ModelIn, + NumpyOrCupy, + PathLike, + TransformedData, + c_bst_ulong, +) +from .compat import ( + _is_cupy_alike, + import_polars, + import_pyarrow, + is_pandas_available, + is_pyarrow_available, + py_str, +) +from .objective import Objective, TreeObjective, _grad_arrinf + +if TYPE_CHECKING: + from pandas import DataFrame as PdDataFrame + + +XGBoostError = _XGBoostError + + +def _parse_eval_str(result: str) -> List[Tuple[str, float]]: + """Parse an eval result string from the booster.""" + splited = result.split()[1:] + # split up `test-error:0.1234` + metric_score_str = [tuple(s.split(":")) for s in splited] + # convert to float + metric_score = [(n, float(s)) for n, s in metric_score_str] + return metric_score + + +IterRange = TypeVar("IterRange", Optional[Tuple[int, int]], Tuple[int, int]) + + +def _expect(expectations: Sequence[Type], got: Type) -> str: + """Translate input error into string. + + Parameters + ---------- + expectations : + a list of expected value. + got : + actual input + + Returns + ------- + msg: str + """ + msg = "Expecting " + for t in range(len(expectations) - 1): + msg += str(expectations[t]) + msg += " or " + msg += str(expectations[-1]) + msg += ". Got " + str(got) + return msg + + +def _check_distributed_params(kwargs: Dict[str, Any]) -> None: + """Validate parameters in distributed environments.""" + device = kwargs.get("device", None) + if device and not isinstance(device, str): + msg = "Invalid type for the `device` parameter" + msg += _expect((str,), type(device)) + raise TypeError(msg) + + if device and device.find(":") != -1: + if device != "sycl:gpu": + raise ValueError( + "Distributed training doesn't support selecting device ordinal as GPUs" + " are managed by the distributed frameworks. use `device=cuda` or" + " `device=gpu` instead." + ) + + if kwargs.get("booster", None) == "gblinear": + raise NotImplementedError( + f"booster `{kwargs['booster']}` is not supported for distributed training." + ) + + +def _validate_data_split_mode(data_split_mode: int) -> None: + try: + mode = int(data_split_mode) + except (TypeError, ValueError) as e: + raise ValueError( + "Only row-wise data split is supported. " + "Column-wise data split has been removed." + ) from e + if mode != 0: + raise ValueError( + "Column-wise data split has been removed. " + "Please use row-wise data split instead." + ) + + +def _validate_feature_info( + feature_info: Sequence[str], n_features: int, name: str +) -> List[str]: + if not isinstance(feature_info, (str, Sequence, Categories)): + raise TypeError( + f"Expecting a sequence of strings for {name}, got: {type(feature_info)}" + ) + feature_info = list(feature_info) + if len(feature_info) != n_features and n_features != 0: + msg = ( + f"{name} must have the same length as the number of data columns, " + f"expected {n_features}, got {len(feature_info)}", + ) + raise ValueError(msg) + return feature_info + + +def build_info() -> dict: + """Build information of XGBoost. The returned value format is not stable. Also, + please note that build time dependency is not the same as runtime dependency. For + instance, it's possible to build XGBoost with older CUDA version but run it with the + latest one. + + .. versionadded:: 1.6.0 + + """ + j_info = ctypes.c_char_p() + _check_call(_LIB.XGBuildInfo(ctypes.byref(j_info))) + assert j_info.value is not None + res = json.loads(j_info.value.decode()) # pylint: disable=no-member + res["libxgboost"] = _LIB.path + return res + + +def _numpy2ctypes_type(dtype: Type[np.number]) -> Type[CNumeric]: + _NUMPY_TO_CTYPES_MAPPING: Dict[Type[np.number], Type[CNumeric]] = { + np.float32: ctypes.c_float, + np.float64: ctypes.c_double, + np.uint32: ctypes.c_uint, + np.uint64: ctypes.c_uint64, + np.int32: ctypes.c_int32, + np.int64: ctypes.c_int64, + } + if np.intc is not np.int32: # Windows + _NUMPY_TO_CTYPES_MAPPING[np.intc] = _NUMPY_TO_CTYPES_MAPPING[np.int32] + if dtype not in _NUMPY_TO_CTYPES_MAPPING: + raise TypeError( + f"Supported types: {_NUMPY_TO_CTYPES_MAPPING.keys()}, got: {dtype}" + ) + return _NUMPY_TO_CTYPES_MAPPING[dtype] + + +def ctypes2numpy(cptr: CNumericPtr, length: int, dtype: Type[np.number]) -> np.ndarray: + """Convert a ctypes pointer array to a numpy array.""" + ctype: Type[CNumeric] = _numpy2ctypes_type(dtype) + if not isinstance(cptr, ctypes.POINTER(ctype)): + raise RuntimeError(f"expected {ctype} pointer") + res = np.zeros(length, dtype=dtype) + if not ctypes.memmove(res.ctypes.data, cptr, length * res.strides[0]): + raise RuntimeError("memmove failed") + return res + + +def ctypes2buffer(cptr: CStrPtr, length: int) -> bytearray: + """Convert ctypes pointer to buffer type.""" + if not isinstance(cptr, ctypes.POINTER(ctypes.c_char)): + raise RuntimeError("expected char pointer") + res = bytearray(length) + rptr = (ctypes.c_char * length).from_buffer(res) + if not ctypes.memmove(rptr, cptr, length): + raise RuntimeError("memmove failed") + return res + + +def c_array( + ctype: Type[CTypeT], values: ArrayLike +) -> Union[ctypes.Array, ctypes._Pointer]: + """Convert a python array to c array.""" + if isinstance(values, np.ndarray) and values.dtype.itemsize == ctypes.sizeof(ctype): + return values.ctypes.data_as(ctypes.POINTER(ctype)) + return (ctype * len(values))(*values) + + +def _prediction_output( + shape: CNumericPtr, dims: c_bst_ulong, predts: CFloatPtr, is_cuda: bool +) -> NumpyOrCupy: + arr_shape = tuple(ctypes2numpy(shape, dims.value, np.uint64).flatten()) + array = from_array_interface( + make_array_interface(predts, arr_shape, np.float32, is_cuda) + ) + return array + + +class DataIter(ABC): # pylint: disable=too-many-instance-attributes + """The interface for user defined data iterator. The iterator facilitates + distributed training, :py:class:`QuantileDMatrix`, and external memory support using + :py:class:`DMatrix` or :py:class:`ExtMemQuantileDMatrix`. Most of time, users don't + need to interact with this class directly. + + .. note:: + + The class caches some intermediate results using the `data` input (predictor + `X`) as key. Don't repeat the `X` for multiple batches with different meta data + (like `label`), make a copy if necessary. + + .. note:: + + When the input for each batch is a DataFrame, we assume categories are + consistently encoded for all batches. For example, given two dataframes for two + batches, this is invalid: + + .. code-block:: + + import pandas as pd + + x0 = pd.DataFrame({"a": [0, 1]}, dtype="category") + x1 = pd.DataFrame({"a": [1, 2]}, dtype="category") + + This is invalid because the `x0` has `[0, 1]` as categories while `x2` has `[1, + 2]`. They should share the same set of categories and encoding: + + .. code-block:: + + import numpy as np + + categories = np.array([0, 1, 2]) + x0["a"] = pd.Categorical.from_codes( + codes=np.array([0, 1]), categories=categories + ) + x1["a"] = pd.Categorical.from_codes( + codes=np.array([1, 2]), categories=categories + ) + + You can make sure the consistent encoding in your preprocessing step be careful + that the data is stored in formats that preserve the encoding when chunking the + data. + + Parameters + ---------- + cache_prefix : + Prefix to the cache files, only used in external memory. + + Note that using this class for external memory **will cache data + on disk** under the path passed here. + + release_data : + Whether the iterator should release the data during iteration. Set it to True if + the data transformation (converting data to np.float32 type) is memory + intensive. Otherwise, if the transformation is computation intensive then we can + keep the cache. + + on_host : + Whether the data should be cached on the host memory instead of the file system + when using GPU with external memory. When set to true (the default), the + "external memory" is the CPU (host) memory. See + :doc:`/tutorials/external_memory` for more info. + + .. versionadded:: 3.0.0 + + .. warning:: + + This is an experimental parameter and subject to change. + + min_cache_page_bytes : + The minimum number of bytes of each cached pages. Only used for on-host cache + with GPU-based :py:class:`ExtMemQuantileDMatrix`. When using GPU-based external + memory with the data cached in the host memory, XGBoost can concatenate the + pages internally to increase the batch size for the GPU. The default page size + is about 1/16 of the total device memory. Users can manually set the value based + on the actual hardware and datasets. Set this to 0 to disable page + concatenation. + + .. versionadded:: 3.0.0 + + .. warning:: + + This is an experimental parameter and subject to change. + + """ + + def __init__( + self, + cache_prefix: Optional[str] = None, + release_data: bool = True, + *, + on_host: bool = True, + min_cache_page_bytes: Optional[int] = None, + ) -> None: + self.cache_prefix = cache_prefix + self.on_host = on_host + self.min_cache_page_bytes = min_cache_page_bytes + + self._handle = _ProxyDMatrix() + self._exception: Optional[Exception] = None + self._enable_categorical = False + self._release = release_data + # Stage data in Python until reset or next is called to avoid data being free. + self._temporary_data: Optional[TransformedData] = None + self._data_ref: Optional[weakref.ReferenceType] = None + + def get_callbacks(self, enable_categorical: bool) -> Tuple[Callable, Callable]: + """Get callback functions for iterating in C. This is an internal function.""" + assert hasattr(self, "cache_prefix"), "__init__ is not called." + reset_callback = ctypes.CFUNCTYPE(None, ctypes.c_void_p)(self._reset_wrapper) + next_callback = ctypes.CFUNCTYPE( + ctypes.c_int, + ctypes.c_void_p, + )(self._next_wrapper) + self._enable_categorical = enable_categorical + return reset_callback, next_callback + + @property + def proxy(self) -> "_ProxyDMatrix": + """Handle of DMatrix proxy.""" + return self._handle + + def _handle_exception(self, fn: Callable, dft_ret: _T) -> _T: + if self._exception is not None: + return dft_ret + + try: + return fn() + except Exception as e: # pylint: disable=broad-except + # Defer the exception in order to return 0 and stop the iteration. + # Exception inside a ctype callback function has no effect except + # for printing to stderr (doesn't stop the execution). + tb = sys.exc_info()[2] + # On dask, the worker is restarted and somehow the information is + # lost. + self._exception = e.with_traceback(tb) + return dft_ret + + def reraise(self) -> None: + """Reraise the exception thrown during iteration.""" + self._temporary_data = None + if self._exception is not None: + # pylint 2.7.0 believes `self._exception` can be None even with `assert + # isinstace` + exc = self._exception + self._exception = None + raise exc # pylint: disable=raising-bad-type + + def __del__(self) -> None: + assert self._temporary_data is None + assert self._exception is None + + def _reset_wrapper(self, this: None) -> None: # pylint: disable=unused-argument + """A wrapper for user defined `reset` function.""" + # free the data + if self._release: + self._temporary_data = None + self._handle_exception(self.reset, None) + + def _next_wrapper(self, this: None) -> int: # pylint: disable=unused-argument + """A wrapper for user defined `next` function. + + `this` is not used in Python. ctypes can handle `self` of a Python + member function automatically when converting it to c function + pointer. + + """ + + @require_keyword_args(True) + def input_data( + *, + data: Any, + feature_names: Optional[FeatureNames] = None, + feature_types: Optional[FeatureTypes] = None, + **kwargs: Any, + ) -> None: + from .data import _proxy_transform, dispatch_proxy_set_data + + # Reduce the amount of transformation that's needed for QuantileDMatrix. + # + # To construct the QDM, one needs 4 iterations on CPU, or 2 iterations on + # GPU. If the QDM has only one batch of input (most of the cases), we can + # avoid transforming the data repeatedly. + try: + ref = weakref.ref(data) + except TypeError: + ref = None + if ( + self._temporary_data is not None + and ref is not None + and ref is self._data_ref + ): + new, feature_names, feature_types = self._temporary_data + else: + new, feature_names, feature_types = _proxy_transform( + data, + feature_names, + feature_types, + self._enable_categorical, + ) + # Stage the data, meta info are copied inside C++ MetaInfo. + self._temporary_data = (new, feature_names, feature_types) + dispatch_proxy_set_data(self.proxy, new) + self.proxy.set_info( + feature_names=feature_names, + feature_types=feature_types, + **kwargs, + ) + self._data_ref = ref + + # Release the data before next batch is loaded. + if self._release: + self._temporary_data = None + # pylint: disable=not-callable + return self._handle_exception(lambda: int(self.next(input_data)), 0) + + @abstractmethod + def reset(self) -> None: + """Reset the data iterator. Prototype for user defined function.""" + raise NotImplementedError() + + @abstractmethod + def next(self, input_data: Callable) -> bool: + """Set the next batch of data. + + Parameters + ---------- + + input_data: + A function with same data fields like `data`, `label` with + `xgboost.DMatrix`. + + Returns + ------- + False if there's no more batch, otherwise True. + + """ + raise NotImplementedError() + + +# Notice for `require_keyword_args` +# Authors: Olivier Grisel +# Gael Varoquaux +# Andreas Mueller +# Lars Buitinck +# Alexandre Gramfort +# Nicolas Tresegnie +# Sylvain Marie +# License: BSD 3 clause +def require_keyword_args( + error: bool, +) -> Callable[[Callable[..., _T]], Callable[..., _T]]: + """Decorator for methods that issues warnings for positional arguments + + Using the keyword-only argument syntax in pep 3102, arguments after the + * will issue a warning or error when passed as a positional argument. + + Modified from sklearn utils.validation. + + Parameters + ---------- + error : + Whether to throw an error or raise a warning. + """ + + def throw_if(func: Callable[..., _T]) -> Callable[..., _T]: + """Throw an error/warning if there are positional arguments after the asterisk. + + Parameters + ---------- + f : + function to check arguments on. + + """ + sig = signature(func) + kwonly_args = [] + all_args = [] + + for name, param in sig.parameters.items(): + if param.kind == Parameter.POSITIONAL_OR_KEYWORD: + all_args.append(name) + elif param.kind == Parameter.KEYWORD_ONLY: + kwonly_args.append(name) + + @wraps(func) + def inner_f(*args: Any, **kwargs: Any) -> _T: + extra_args = len(args) - len(all_args) + if not all_args and extra_args > 0: # keyword argument only + raise TypeError("Keyword argument is required.") + + if extra_args > 0: + # ignore first 'self' argument for instance methods + args_msg = [ + f"{name}" + for name, _ in zip(kwonly_args[:extra_args], args[-extra_args:]) + ] + # pylint: disable=consider-using-f-string + msg = "Pass `{}` as keyword args.".format(", ".join(args_msg)) + if error: + raise TypeError(msg) + warnings.warn(msg, FutureWarning) + for k, arg in zip(sig.parameters, args): + kwargs[k] = arg + return func(**kwargs) + + return inner_f + + return throw_if + + +_deprecate_positional_args = require_keyword_args(False) + + +def _get_categories( + cfn: Callable[[ctypes.c_char_p], int], + feature_names: FeatureNames, + n_features: int, +) -> ArrowCatList: + if not is_pyarrow_available(): + raise ImportError( + "`pyarrow` is required for exporting categories to arrow arrays." + ) + + if not TYPE_CHECKING: + pa = import_pyarrow() + else: + import pyarrow as pa + + results: ArrowCatList = [] + + ret = ctypes.c_char_p() + _check_call(cfn(ret)) + if ret.value is None: + results = [(feature_names[i], None) for i in range(n_features)] + return results + + retstr = ret.value.decode() # pylint: disable=no-member + jcats = json.loads(retstr) + assert isinstance(jcats, list) and len(jcats) == n_features + + for fidx in range(n_features): + f_jcats = jcats[fidx] + if f_jcats is None: + # Numeric data + results.append((feature_names[fidx], None)) + continue + + if "offsets" not in f_jcats: + values = from_array_interface(f_jcats) + pa_values = pa.Array.from_pandas(values) + results.append((feature_names[fidx], pa_values)) + continue + + joffsets = f_jcats["offsets"] + jvalues = f_jcats["values"] + offsets = from_array_interface(joffsets) + values = from_array_interface(jvalues) + pa_offsets = pa.array(offsets).buffers() + pa_values = pa.array(values).buffers() + assert pa_offsets[0] is None and pa_values[0] is None, ( + "Should not have null mask." + ) + pa_dict = pa.StringArray.from_buffers( + len(offsets) - 1, pa_offsets[1], pa_values[1] + ) + results.append((feature_names[fidx], pa_dict)) + + return results + + +def _is_iter(data: DataType) -> TypeGuard[DataIter]: + return isinstance(data, DataIter) + + +class SingleBatchInternalIter(DataIter): # pylint: disable=R0902 + """An iterator for single batch data to help creating device DMatrix. + Transforming input directly to histogram with normal single batch data API + can not access weight for sketching. So this iterator acts as a staging + area for meta info. + """ + + def __init__(self, **kwargs: Any) -> None: + self.kwargs = kwargs + self.it = 0 + + # This does not necessarily increase memory usage as the data transformation + # might use memory. + super().__init__(release_data=False) + + def next(self, input_data: Callable) -> bool: + if self.it == 1: + return False + self.it += 1 + input_data(**self.kwargs) + return True + + def reset(self) -> None: + self.it = 0 + + +class DMatrix: # pylint: disable=too-many-instance-attributes,too-many-public-methods + """Data Matrix used in XGBoost. + + DMatrix is an internal data structure that is used by XGBoost, which is optimized + for both memory efficiency and training speed. You can construct DMatrix from + multiple different sources of data. + + """ + + @_deprecate_positional_args + def __init__( + self, + data: DataType, + label: Optional[ArrayLike] = None, + *, + weight: Optional[ArrayLike] = None, + base_margin: Optional[ArrayLike] = None, + missing: Optional[float] = None, + silent: bool = False, + feature_names: Optional[FeatureNames] = None, + feature_types: Optional[Union[FeatureTypes, Categories]] = None, + nthread: Optional[int] = None, + group: Optional[ArrayLike] = None, + qid: Optional[ArrayLike] = None, + label_lower_bound: Optional[ArrayLike] = None, + label_upper_bound: Optional[ArrayLike] = None, + feature_weights: Optional[ArrayLike] = None, + enable_categorical: bool = True, + data_split_mode: int = 0, + ) -> None: + """Parameters + ---------- + data : + Data source of DMatrix. See :ref:`py-data` for a list of supported input + types. + + Note that, if passing an iterator, it **will cache data on disk**, and note + that fields like ``label`` will be concatenated in-memory from multiple + calls to the iterator. + label : + Label of the training data. + weight : + Weight for each instance. + + .. note:: + + For ranking task, weights are per-group. In ranking task, one weight + is assigned to each group (not each data point). This is because we + only care about the relative ordering of data points within each group, + so it doesn't make sense to assign weights to individual data points. + + base_margin : + Global bias for each instance. See :doc:`/tutorials/intercept` for details. + missing : + Value in the input data which needs to be present as a missing value. If + None, defaults to np.nan. + silent : + Whether print messages during construction + feature_names : + Set names for features. + feature_types : + + Set types for features. If `data` is a DataFrame type and passing + `enable_categorical=True`, the types will be deduced automatically from the + column types. + + Otherwise, one can pass a list-like input with the same length as number of + columns in `data`, with the following possible values: + + - "c", which represents categorical columns. + - "q", which represents numeric columns. + - "int", which represents integer columns. + - "i", which represents boolean columns. + + Note that, while categorical types are treated differently from the rest for + model fitting purposes, the other types do not influence the generated + model, but have effects in other functionalities such as feature + importances. + + For categorical features, the input is assumed to be preprocessed and + encoded by the users. The encoding can be done via + :py:class:`sklearn.preprocessing.OrdinalEncoder` or pandas dataframe + `.cat.codes` method. This is useful when users want to specify categorical + features without having to construct a dataframe as input. + + .. versionadded:: 3.1.0 + + Alternatively, user can pass a :py:class:`~xgboost.core.Categories` object + returned from previous training as a reference for re-coding. One can obtain + the reference with the :py:meth:`.get_categories` from the previous training + DMatrix or the Booster. This feature is experimental. + + nthread : + Number of threads to use for loading data when parallelization is + applicable. If -1, uses maximum threads available on the system. + group : + Group size for all ranking group. + qid : + Query ID for data samples, used for ranking. + label_lower_bound : + Lower bound for survival training. + label_upper_bound : + Upper bound for survival training. + feature_weights : + Set feature weights for column sampling. + enable_categorical : + + .. versionadded:: 1.3.0 + + .. note:: This parameter is experimental + + Experimental support of specializing for categorical features. See + :doc:`/tutorials/categorical` for more info. + + If passing `True` and `data` is a data frame (from supported libraries such + as Pandas, Modin, polars, and cuDF), The DMatrix recognizes categorical + columns and automatically set the `feature_types` parameter. If `data` is + not a data frame, this argument is ignored. + + If passing `False` and `data` is a data frame with categorical columns, it + will result in an error. + + See notes in the :py:class:`DataIter` for consistency requirement when the + input is an iterator. + + .. versionchanged:: 3.1.0 + + XGBoost can remember the encoding of categories when the input is a + dataframe. + + data_split_mode : + + .. deprecated:: 3.3.0 + + Only row-wise split is supported. Column-wise split has been removed + and is rejected during DMatrix construction. + + """ + if group is not None and qid is not None: + raise ValueError("Either one of `group` or `qid` should be None.") + + self.missing = missing if missing is not None else np.nan + self.nthread = nthread if nthread is not None else -1 + self.silent = silent + + if isinstance(data, ctypes.c_void_p): + # Used for constructing DMatrix slice. + self.handle = data + return + + _validate_data_split_mode(data_split_mode) + + from .data import dispatch_data_backend + + if _is_iter(data): + self._init_from_iter(data, enable_categorical) + assert self.handle is not None + return + + handle, feature_names, feature_types = dispatch_data_backend( + data=data, + missing=self.missing, + threads=self.nthread, + feature_names=feature_names, + feature_types=feature_types, + enable_categorical=enable_categorical, + ) + assert handle is not None + self.handle = handle + + self.set_info( + label=label, + weight=weight, + base_margin=base_margin, + group=group, + qid=qid, + label_lower_bound=label_lower_bound, + label_upper_bound=label_upper_bound, + feature_weights=feature_weights, + ) + + if feature_names is not None: + self.feature_names = feature_names + if feature_types is not None: + self.feature_types = feature_types + + def _init_from_iter(self, it: DataIter, enable_categorical: bool) -> None: + args = make_jcargs( + missing=self.missing, + nthread=self.nthread, + cache_prefix=it.cache_prefix if it.cache_prefix else "", + on_host=it.on_host, + min_cache_page_bytes=it.min_cache_page_bytes, + ) + handle = ctypes.c_void_p() + reset_callback, next_callback = it.get_callbacks(enable_categorical) + ret = _LIB.XGDMatrixCreateFromCallback( + None, + it.proxy.handle, + reset_callback, + next_callback, + args, + ctypes.byref(handle), + ) + it.reraise() + # delay check_call to throw intermediate exception first + _check_call(ret) + self.handle = handle + + def __del__(self) -> None: + if hasattr(self, "handle"): + assert self.handle is not None + _check_call(_LIB.XGDMatrixFree(self.handle)) + del self.handle + + @_deprecate_positional_args + def set_info( + self, + *, + label: Optional[ArrayLike] = None, + weight: Optional[ArrayLike] = None, + base_margin: Optional[ArrayLike] = None, + group: Optional[ArrayLike] = None, + qid: Optional[ArrayLike] = None, + label_lower_bound: Optional[ArrayLike] = None, + label_upper_bound: Optional[ArrayLike] = None, + feature_names: Optional[FeatureNames] = None, + feature_types: Optional[FeatureTypes] = None, + feature_weights: Optional[ArrayLike] = None, + ) -> None: + """Set meta info for DMatrix. See doc string for :py:obj:`xgboost.DMatrix`.""" + from .data import dispatch_meta_backend + + if label is not None: + self.set_label(label) + if weight is not None: + self.set_weight(weight) + if base_margin is not None: + self.set_base_margin(base_margin) + if group is not None: + self.set_group(group) + if qid is not None: + self.set_uint_info("qid", qid) + if label_lower_bound is not None: + self.set_float_info("label_lower_bound", label_lower_bound) + if label_upper_bound is not None: + self.set_float_info("label_upper_bound", label_upper_bound) + if feature_names is not None: + self.feature_names = feature_names + if feature_types is not None: + self.feature_types = feature_types + if feature_weights is not None: + dispatch_meta_backend( + matrix=self, data=feature_weights, name="feature_weights" + ) + + def _get_info(self, field: str) -> NumpyOrCupy: + """Get meta info.""" + c_sdata = ctypes.c_char_p() + _check_call( + _LIB.XGDMatrixGetInfoRef(self.handle, c_str(field), ctypes.byref(c_sdata)) + ) + assert c_sdata.value is not None + idata = json.loads(c_sdata.value) + data = from_array_interface(idata) + return data + + def get_float_info(self, field: str) -> NumpyOrCupy: + """Get float property from the DMatrix. + + Parameters + ---------- + field: str + The field name of the information. + + """ + return self._get_info(field) + + def get_uint_info(self, field: str) -> NumpyOrCupy: + """Get unsigned integer property from the DMatrix. + + Parameters + ---------- + field: str + The field name of the information. + + """ + return self._get_info(field) + + def set_float_info(self, field: str, data: ArrayLike) -> None: + """Set float type property into the DMatrix. + + Parameters + ---------- + field: str + The field name of the information + + data: numpy array + The array of data to be set + """ + from .data import dispatch_meta_backend + + dispatch_meta_backend(self, data, field, "float") + + def set_float_info_npy2d(self, field: str, data: ArrayLike) -> None: + """Set float type property into the DMatrix + for numpy 2d array input + + Parameters + ---------- + field: str + The field name of the information + + data: numpy array + The array of data to be set + """ + from .data import dispatch_meta_backend + + dispatch_meta_backend(self, data, field, "float") + + def set_uint_info(self, field: str, data: ArrayLike) -> None: + """Set uint type property into the DMatrix. + + Parameters + ---------- + field: str + The field name of the information + + data: numpy array + The array of data to be set + """ + from .data import dispatch_meta_backend + + dispatch_meta_backend(self, data, field, "uint32") + + def save_binary(self, fname: PathLike, silent: bool = True) -> None: + """Save DMatrix to an XGBoost buffer. Saved binary can be later loaded + by providing the path to :py:func:`xgboost.DMatrix` as input. + + Parameters + ---------- + fname : string or os.PathLike + Name of the output buffer file. + silent : bool (optional; default: True) + If set, the output is suppressed. + """ + fname = os.fspath(os.path.expanduser(fname)) + _check_call( + _LIB.XGDMatrixSaveBinary(self.handle, c_str(fname), ctypes.c_int(silent)) + ) + + def set_label(self, label: ArrayLike) -> None: + """Set label of dmatrix + + Parameters + ---------- + label: array like + The label information to be set into DMatrix + """ + from .data import dispatch_meta_backend + + dispatch_meta_backend(self, label, "label", "float") + + def set_weight(self, weight: ArrayLike) -> None: + """Set weight of each instance. + + Parameters + ---------- + weight : array like + Weight for each data point + + .. note:: For ranking task, weights are per-group. + + In ranking task, one weight is assigned to each group (not each + data point). This is because we only care about the relative + ordering of data points within each group, so it doesn't make + sense to assign weights to individual data points. + + """ + from .data import dispatch_meta_backend + + dispatch_meta_backend(self, weight, "weight", "float") + + def set_base_margin(self, margin: ArrayLike) -> None: + """Set base margin of booster to start from. + + This can be used to specify a prediction value of existing model to be + base_margin However, remember margin is needed, instead of transformed + prediction e.g. for logistic regression: need to put in value before + logistic transformation see also example/demo.py + + Parameters + ---------- + margin: array like + Prediction margin of each datapoint + + """ + from .data import dispatch_meta_backend + + dispatch_meta_backend(self, margin, "base_margin", "float") + + def set_group(self, group: ArrayLike) -> None: + """Set group size of DMatrix (used for ranking). + + Parameters + ---------- + group : array like + Group size of each group + """ + from .data import dispatch_meta_backend + + dispatch_meta_backend(self, group, "group", "uint32") + + def get_label(self) -> NumpyOrCupy: + """Get the label of the DMatrix.""" + return self._get_info("label") + + def get_weight(self) -> NumpyOrCupy: + """Get the weight of the DMatrix.""" + return self._get_info("weight") + + def get_base_margin(self) -> NumpyOrCupy: + """Get the base margin of the DMatrix.""" + return self._get_info("base_margin") + + def get_group(self) -> np.ndarray: + """Get the group of the DMatrix. + + Returns + ------- + group + """ + group_ptr = self._get_info("group_ptr") + return np.diff(group_ptr) + + def get_data(self) -> scipy.sparse.csr_matrix: + """Get the predictors from DMatrix as a CSR matrix. This getter is mostly for + testing purposes. If this is a quantized DMatrix then quantized values are + returned instead of input values. + + .. versionadded:: 1.7.0 + + """ + indptr = np.empty(self.num_row() + 1, dtype=np.uint64) + indices = np.empty(self.num_nonmissing(), dtype=np.uint32) + data = np.empty(self.num_nonmissing(), dtype=np.float32) + + c_indptr = indptr.ctypes.data_as(ctypes.POINTER(c_bst_ulong)) + c_indices = indices.ctypes.data_as(ctypes.POINTER(ctypes.c_uint32)) + c_data = data.ctypes.data_as(ctypes.POINTER(ctypes.c_float)) + config = from_pystr_to_cstr(json.dumps({})) + + _check_call( + _LIB.XGDMatrixGetDataAsCSR(self.handle, config, c_indptr, c_indices, c_data) + ) + ret = scipy.sparse.csr_matrix( + (data, indices, indptr), shape=(self.num_row(), self.num_col()) + ) + return ret + + def get_quantile_cut(self) -> Tuple[np.ndarray, np.ndarray]: + """Get quantile cuts for quantization. + + .. versionadded:: 2.0.0 + + """ + n_features = self.num_col() + + c_sindptr = ctypes.c_char_p() + c_sdata = ctypes.c_char_p() + config = make_jcargs() + _check_call( + _LIB.XGDMatrixGetQuantileCut( + self.handle, config, ctypes.byref(c_sindptr), ctypes.byref(c_sdata) + ) + ) + assert c_sindptr.value is not None + assert c_sdata.value is not None + + i_indptr = json.loads(c_sindptr.value) + indptr = from_array_interface(i_indptr) + assert indptr.size == n_features + 1 + assert indptr.dtype == np.uint64 + + i_data = json.loads(c_sdata.value) + data = from_array_interface(i_data) + assert data.size == indptr[-1] + assert data.dtype == np.float32 + return indptr, data + + def get_categories(self, export_to_arrow: bool = False) -> Categories: + """Get the categories in the dataset. + + .. versionadded:: 3.1.0 + + .. warning:: + + This function is experimental. + + Parameters + ---------- + export_to_arrow : + The returned container will contain a list of ``pyarrow`` arrays for the + categories. See the :py:meth:`~Categories.to_arrow` for more info. + + """ + fnames = self.feature_names + n_features = self.num_col() + if fnames is None: + fnames = [str(i) for i in range(n_features)] + + hdl = ctypes.c_void_p() + if export_to_arrow: + arrow_arrays = _get_categories( + lambda ret: _LIB.XGDMatrixGetCategoriesExportToArrow( + self.handle, None, ctypes.byref(hdl), ctypes.byref(ret) + ), + fnames, + n_features, + ) + else: + arrow_arrays = None + _check_call( + _LIB.XGDMatrixGetCategories(self.handle, None, ctypes.byref(hdl)) + ) + + return Categories( + (hdl, lambda: _check_call(_LIB.XGBCategoriesFree(hdl))), arrow_arrays + ) + + def num_row(self) -> int: + """Get the number of rows in the DMatrix.""" + ret = c_bst_ulong() + _check_call(_LIB.XGDMatrixNumRow(self.handle, ctypes.byref(ret))) + return ret.value + + def num_col(self) -> int: + """Get the number of columns (features) in the DMatrix.""" + ret = c_bst_ulong() + _check_call(_LIB.XGDMatrixNumCol(self.handle, ctypes.byref(ret))) + return ret.value + + def num_nonmissing(self) -> int: + """Get the number of non-missing values in the DMatrix. + + .. versionadded:: 1.7.0 + + """ + ret = c_bst_ulong() + _check_call(_LIB.XGDMatrixNumNonMissing(self.handle, ctypes.byref(ret))) + return ret.value + + def slice( + self, rindex: Union[List[int], np.ndarray], allow_groups: bool = False + ) -> "DMatrix": + """Slice the DMatrix and return a new DMatrix that only contains `rindex`. + + Parameters + ---------- + rindex + List of indices to be selected. + allow_groups + Allow slicing of a matrix with a groups attribute + + Returns + ------- + res + A new DMatrix containing only selected indices. + """ + from .data import _maybe_np_slice + + handle = ctypes.c_void_p() + + rindex = _maybe_np_slice(rindex, dtype=np.int32) + _check_call( + _LIB.XGDMatrixSliceDMatrixEx( + self.handle, + c_array(ctypes.c_int, rindex), + c_bst_ulong(len(rindex)), + ctypes.byref(handle), + ctypes.c_int(1 if allow_groups else 0), + ) + ) + return DMatrix(handle) + + @property + def feature_names(self) -> Optional[FeatureNames]: + """Labels for features (column labels). + + Setting it to ``None`` resets existing feature names. + + """ + length = c_bst_ulong() + sarr = ctypes.POINTER(ctypes.c_char_p)() + _check_call( + _LIB.XGDMatrixGetStrFeatureInfo( + self.handle, + c_str("feature_name"), + ctypes.byref(length), + ctypes.byref(sarr), + ) + ) + feature_names = from_cstr_to_pystr(sarr, length) + if not feature_names: + return None + return feature_names + + @feature_names.setter + def feature_names(self, feature_names: Optional[FeatureNames]) -> None: + if feature_names is None: + _check_call( + _LIB.XGDMatrixSetStrFeatureInfo( + self.handle, c_str("feature_name"), None, c_bst_ulong(0) + ) + ) + return + + # validate feature name + feature_names = _validate_feature_info( + feature_names, + self.num_col(), + "feature names", + ) + if len(feature_names) != len(set(feature_names)): + values, counts = np.unique( + feature_names, + return_index=False, + return_inverse=False, + return_counts=True, + ) + duplicates = [name for name, cnt in zip(values, counts) if cnt > 1] + raise ValueError( + f"feature_names must be unique. Duplicates found: {duplicates}" + ) + + # prohibit the use symbols that may affect parsing. e.g. []< + if not all( + isinstance(f, str) and not any(x in f for x in ["[", "]", "<"]) + for f in feature_names + ): + raise ValueError( + "feature_names must be string, and may not contain [, ] or <" + ) + + feature_names_bytes = [bytes(f, encoding="utf-8") for f in feature_names] + c_feature_names = (ctypes.c_char_p * len(feature_names_bytes))( + *feature_names_bytes + ) + _check_call( + _LIB.XGDMatrixSetStrFeatureInfo( + self.handle, + c_str("feature_name"), + c_feature_names, + c_bst_ulong(len(feature_names)), + ) + ) + + @property + def feature_types(self) -> Optional[FeatureTypes]: + """Type of features (column types). + + This is for displaying the results and categorical data support. See + :py:class:`DMatrix` for details. + + Setting it to ``None`` resets existing feature types. + + """ + length = c_bst_ulong() + sarr = ctypes.POINTER(ctypes.c_char_p)() + _check_call( + _LIB.XGDMatrixGetStrFeatureInfo( + self.handle, + c_str("feature_type"), + ctypes.byref(length), + ctypes.byref(sarr), + ) + ) + res = from_cstr_to_pystr(sarr, length) + if not res: + return None + return res + + @feature_types.setter + def feature_types(self, feature_types: Optional[FeatureTypes]) -> None: + if feature_types is None: + # Reset + _check_call( + _LIB.XGDMatrixSetStrFeatureInfo( + self.handle, c_str("feature_type"), None, c_bst_ulong(0) + ) + ) + return + + feature_types = _validate_feature_info( + feature_types, + self.num_col(), + "feature types", + ) + + feature_types_bytes = [bytes(f, encoding="utf-8") for f in feature_types] + c_feature_types = (ctypes.c_char_p * len(feature_types_bytes))( + *feature_types_bytes + ) + _check_call( + _LIB.XGDMatrixSetStrFeatureInfo( + self.handle, + c_str("feature_type"), + c_feature_types, + c_bst_ulong(len(feature_types)), + ) + ) + + +class _ProxyDMatrix(DMatrix): + """A placeholder class when DMatrix cannot be constructed (QuantileDMatrix, + inplace_predict). + + """ + + def __init__(self) -> None: # pylint: disable=super-init-not-called + self.handle = ctypes.c_void_p() + _check_call(_LIB.XGProxyDMatrixCreate(ctypes.byref(self.handle))) + + def _ref_data_from_cuda_interface(self, data: DataType) -> None: + """Reference data from CUDA array interface.""" + arrinf = cuda_array_interface(data) + _check_call(_LIB.XGProxyDMatrixSetDataCudaArrayInterface(self.handle, arrinf)) + + def _ref_data_from_cuda_columnar(self, data: TransformedDf) -> None: + """Reference data from CUDA columnar format.""" + _check_call( + _LIB.XGProxyDMatrixSetDataCudaColumnar(self.handle, data.array_interface()) + ) + + def _ref_data_from_array(self, data: np.ndarray) -> None: + """Reference data from numpy array.""" + _check_call(_LIB.XGProxyDMatrixSetDataDense(self.handle, array_interface(data))) + + def _ref_data_from_columnar(self, data: TransformedDf) -> None: + """Reference data from a CPU DataFrame.""" + _check_call( + _LIB.XGProxyDMatrixSetDataColumnar(self.handle, data.array_interface()) + ) + + def _ref_data_from_csr(self, csr: scipy.sparse.csr_matrix) -> None: + """Reference data from scipy csr.""" + _LIB.XGProxyDMatrixSetDataCSR( + self.handle, + array_interface(csr.indptr), + array_interface(csr.indices), + array_interface(csr.data), + ctypes.c_size_t(csr.shape[1]), + ) + + +class _RefMixIn: + @property + def ref(self) -> Optional[weakref.ReferenceType]: + """Internal method for retrieving a reference to the training DMatrix.""" + if hasattr(self, "_ref"): + return self._ref + return None + + @ref.setter + def ref(self, ref: weakref.ReferenceType) -> None: + self._ref = ref + + +class QuantileDMatrix(DMatrix, _RefMixIn): + """A DMatrix variant that generates quantilized data directly from input for the + ``hist`` tree method. This DMatrix is primarily designed to save memory in training + by avoiding intermediate storage. Set ``max_bin`` to control the number of bins + during quantisation, which should be consistent with the training parameter + ``max_bin``. When ``QuantileDMatrix`` is used for validation/test dataset, ``ref`` + should be another ``QuantileDMatrix`` or ``DMatrix``, but not recommended as it + defeats the purpose of saving memory) constructed from training dataset. See + :py:obj:`xgboost.DMatrix` for documents on meta info. + + .. note:: + + Do not use ``QuantileDMatrix`` as validation/test dataset without supplying a + reference (the training dataset) ``QuantileDMatrix`` using ``ref`` as some + information may be lost in quantisation. + + .. versionadded:: 1.7.0 + + Examples + -------- + + .. code-block:: + + from sklearn.datasets import make_regression + from sklearn.model_selection import train_test_split + + X, y = make_regression() + X_train, X_test, y_train, y_test = train_test_split(X, y) + Xy_train = xgb.QuantileDMatrix(X_train, y_train) + # It's necessary to have the training DMatrix as a reference for valid + # quantiles. + Xy_test = xgb.QuantileDMatrix(X_test, y_test, ref=Xy_train) + + Parameters + ---------- + max_bin : + The number of histogram bin, should be consistent with the training parameter + ``max_bin``. + + ref : + The training dataset that provides quantile information, needed when creating + validation/test dataset with ``QuantileDMatrix``. Supplying the training DMatrix + as a reference means that the same quantisation applied to the training data is + applied to the validation/test data + + max_quantile_batches : + Deprecated. This parameter no longer has any effect and will be removed in a + future release. + + .. versionadded:: 3.0.0 + + .. deprecated:: 3.3.0 + + data_split_mode : + Deprecated. Only row-wise split is supported. Column-wise split has been + removed and is rejected during DMatrix construction. + + .. deprecated:: 3.3.0 + + """ + + @_deprecate_positional_args + def __init__( # pylint: disable=super-init-not-called + self, + data: DataType, + label: Optional[ArrayLike] = None, + *, + weight: Optional[ArrayLike] = None, + base_margin: Optional[ArrayLike] = None, + missing: Optional[float] = None, + silent: bool = False, + feature_names: Optional[FeatureNames] = None, + feature_types: Optional[FeatureTypes] = None, + nthread: Optional[int] = None, + max_bin: Optional[int] = None, + ref: Optional[DMatrix] = None, + group: Optional[ArrayLike] = None, + qid: Optional[ArrayLike] = None, + label_lower_bound: Optional[ArrayLike] = None, + label_upper_bound: Optional[ArrayLike] = None, + feature_weights: Optional[ArrayLike] = None, + enable_categorical: bool = True, + max_quantile_batches: Optional[int] = None, + data_split_mode: int = 0, + ) -> None: + self.max_bin = max_bin + self.missing = missing if missing is not None else np.nan + self.nthread = nthread if nthread is not None else -1 + self._silent = silent # unused, kept for compatibility + + if isinstance(data, ctypes.c_void_p): + self.handle = data + return + + _validate_data_split_mode(data_split_mode) + + if qid is not None and group is not None: + raise ValueError( + "Only one of the eval_qid or eval_group for each evaluation " + "dataset should be provided." + ) + if isinstance(data, DataIter): + if any( + info is not None + for info in ( + label, + weight, + base_margin, + feature_names, + feature_types, + group, + qid, + label_lower_bound, + label_upper_bound, + feature_weights, + ) + ): + raise ValueError( + "If data iterator is used as input, data like label should be " + "specified as batch argument." + ) + + self._init( + data, + ref=ref, + label=label, + weight=weight, + base_margin=base_margin, + group=group, + qid=qid, + label_lower_bound=label_lower_bound, + label_upper_bound=label_upper_bound, + feature_weights=feature_weights, + feature_names=feature_names, + feature_types=feature_types, + enable_categorical=enable_categorical, + max_quantile_blocks=max_quantile_batches, + ) + + def _init( + self, + data: DataType, + ref: Optional[DMatrix], + enable_categorical: bool, + max_quantile_blocks: Optional[int], + **meta: Any, + ) -> None: + from .data import _is_dlpack, _transform_dlpack + + if _is_dlpack(data): + # We specialize for dlpack because cupy will take the memory from it so + # it can't be transformed twice. + data = _transform_dlpack(data) + if _is_iter(data): + it = data + else: + it = SingleBatchInternalIter(data=data, **meta) + + handle = ctypes.c_void_p() + reset_callback, next_callback = it.get_callbacks(enable_categorical) + if it.cache_prefix is not None: + raise ValueError( + "QuantileDMatrix doesn't cache data, remove the cache_prefix " + "in iterator to fix this error." + ) + + config = make_jcargs( + nthread=self.nthread, + missing=self.missing, + max_bin=self.max_bin, + max_quantile_blocks=max_quantile_blocks, + ) + ret = _LIB.XGQuantileDMatrixCreateFromCallback( + None, + it.proxy.handle, + ref.handle if ref is not None else ref, + reset_callback, + next_callback, + config, + ctypes.byref(handle), + ) + it.reraise() + # delay check_call to throw intermediate exception first + _check_call(ret) + self.handle = handle + + if ref is not None: + self.ref = weakref.ref(ref) + + +class ExtMemQuantileDMatrix(DMatrix, _RefMixIn): + """The external memory version of the :py:class:`QuantileDMatrix`. + + See :doc:`/tutorials/external_memory` for explanation and usage examples, and + :py:class:`QuantileDMatrix` for parameter document. + + .. warning:: + + This is an experimental feature and subject to change. + + .. versionadded:: 3.0.0 + + """ + + @_deprecate_positional_args + def __init__( # pylint: disable=super-init-not-called + self, + data: DataIter, + *, + missing: Optional[float] = None, + nthread: Optional[int] = None, + max_bin: Optional[int] = None, + ref: Optional[DMatrix] = None, + enable_categorical: bool = True, + max_quantile_batches: Optional[int] = None, + cache_host_ratio: Optional[float] = None, + ) -> None: + """ + Parameters + ---------- + data : + A user-defined :py:class:`DataIter` for loading data. + + max_quantile_batches : + Deprecated. See :py:class:`QuantileDMatrix`. + + cache_host_ratio : + + .. versionadded:: 3.1.0 + + Used by the GPU implementation. For GPU-based inputs, XGBoost can split the + cache into host and device caches to reduce the data transfer overhead. This + parameter specifies the size of host cache compared to the size of the + entire cache: :math:`host / (host + device)`. + + See :ref:`extmem-adaptive-cache` for more info. + + """ + self.max_bin = max_bin + self.missing = missing if missing is not None else np.nan + self.nthread = nthread if nthread is not None else -1 + + self._init( + data, + ref, + enable_categorical=enable_categorical, + max_quantile_blocks=max_quantile_batches, + cache_host_ratio=( + None if cache_host_ratio is None else float(cache_host_ratio) + ), + ) + assert self.handle is not None + + def _init( + self, + it: DataIter, + ref: Optional[DMatrix], + *, + enable_categorical: bool, + max_quantile_blocks: Optional[int] = None, + cache_host_ratio: Optional[float] = None, + ) -> None: + args = make_jcargs( + missing=self.missing, + nthread=self.nthread, + cache_prefix=it.cache_prefix if it.cache_prefix else "", + on_host=it.on_host, + max_bin=self.max_bin, + min_cache_page_bytes=it.min_cache_page_bytes, + # It's called blocks internally due to block-based quantile sketching. + max_quantile_blocks=max_quantile_blocks, + cache_host_ratio=cache_host_ratio, + ) + handle = ctypes.c_void_p() + reset_callback, next_callback = it.get_callbacks(enable_categorical) + # We don't need the iter handle (hence None) in Python as reset,next callbacks + # are member functions, and ctypes can handle the `self` parameter + # automatically. + ret = _LIB.XGExtMemQuantileDMatrixCreateFromCallback( + None, # iter + it.proxy.handle, # proxy + ref.handle if ref is not None else ref, # ref + reset_callback, # reset + next_callback, # next + args, # config + ctypes.byref(handle), # out + ) + it.reraise() + # delay check_call to throw intermediate exception first + _check_call(ret) + self.handle = handle + + if ref is not None: + self.ref = weakref.ref(ref) + + +PlainObj = Callable[[np.ndarray, DMatrix], Tuple[np.ndarray, np.ndarray]] +CustomObj = Union[PlainObj, Objective] +Metric = Callable[[np.ndarray, DMatrix], Tuple[str, float]] + + +def _configure_metrics(params: BoosterParam) -> BoosterParam: + if ( + isinstance(params, dict) + and "eval_metric" in params + and isinstance(params["eval_metric"], list) + ): + eval_metrics = params["eval_metric"] + params.pop("eval_metric", None) + params_list = list(params.items()) + for eval_metric in eval_metrics: + params_list += [("eval_metric", eval_metric)] + return params_list + return params + + +class Booster: + # pylint: disable=too-many-public-methods + """A Booster of XGBoost. + + Booster is the model of xgboost, that contains low level routines for + training, prediction and evaluation. + """ + + def __init__( + self, + params: Optional[BoosterParam] = None, + cache: Optional[Sequence[DMatrix]] = None, + model_file: Optional[Union["Booster", bytearray, os.PathLike, str]] = None, + ) -> None: + """ + Parameters + ---------- + params : + Parameters for boosters. + cache : + List of cache items. + model_file : + Path to the model file if it's string or PathLike. + """ + cache = cache if cache is not None else [] + for d in cache: + if not isinstance(d, DMatrix): + raise TypeError(f"Invalid cache item: {type(d).__name__}", cache) + + dmats = c_array(ctypes.c_void_p, [d.handle for d in cache]) + self.handle: Optional[ctypes.c_void_p] = ctypes.c_void_p() + _check_call( + _LIB.XGBoosterCreate( + dmats, c_bst_ulong(len(cache)), ctypes.byref(self.handle) + ) + ) + for d in cache: + # Validate feature only after the feature names are saved into booster. + self._assign_dmatrix_features(d) + + if isinstance(model_file, Booster): + assert self.handle is not None + # We use the pickle interface for getting memory snapshot from + # another model, and load the snapshot with this booster. + state = model_file.__getstate__() + handle = state["handle"] + del state["handle"] + ptr = (ctypes.c_char * len(handle)).from_buffer(handle) + length = c_bst_ulong(len(handle)) + _check_call(_LIB.XGBoosterUnserializeFromBuffer(self.handle, ptr, length)) + self.__dict__.update(state) + elif isinstance(model_file, (str, os.PathLike, bytearray)): + self.load_model(model_file) + elif model_file is None: + pass + else: + raise TypeError("Unknown type:", model_file) + + params = params or {} + params_processed = _configure_metrics(params.copy()) + params_processed = self._configure_constraints(params_processed) + if isinstance(params_processed, list): + params_processed.append(("validate_parameters", True)) + else: + params_processed["validate_parameters"] = True + + self.set_param(params_processed or {}) + + def _transform_monotone_constrains( + self, value: Union[Dict[str, int], str, Tuple[int, ...]] + ) -> Union[Tuple[int, ...], str]: + if isinstance(value, str): + return value + if isinstance(value, tuple): + return value + + constrained_features = set(value.keys()) + feature_names = self.feature_names or [] + if not constrained_features.issubset(set(feature_names)): + raise ValueError( + "Constrained features are not a subset of training data feature names" + ) + + return tuple(value.get(name, 0) for name in feature_names) + + def _transform_interaction_constraints( + self, value: Union[Sequence[Sequence[str]], str] + ) -> Union[str, List[List[int]]]: + if isinstance(value, str): + return value + feature_idx_mapping = { + name: idx for idx, name in enumerate(self.feature_names or []) + } + + try: + result = [] + for constraint in value: + result.append( + [feature_idx_mapping[feature_name] for feature_name in constraint] + ) + return result + except KeyError as e: + raise ValueError( + "Constrained features are not a subset of training data feature names" + ) from e + + def _configure_constraints(self, params: BoosterParam) -> BoosterParam: + if isinstance(params, dict): + # we must use list in the internal code as there can be multiple metrics + # with the same parameter name `eval_metric` (same key for dictionary). + params = list(params.items()) + for idx, param in enumerate(params): + name, value = param + if value is None: + continue + + if name == "monotone_constraints": + params[idx] = (name, self._transform_monotone_constrains(value)) + elif name == "interaction_constraints": + params[idx] = (name, self._transform_interaction_constraints(value)) + + return params + + def __del__(self) -> None: + if hasattr(self, "handle") and self.handle is not None: + _check_call(_LIB.XGBoosterFree(self.handle)) + self.handle = None + + def __getstate__(self) -> Dict: + # can't pickle ctypes pointers, put model content in a bytearray + this = self.__dict__.copy() + handle = this["handle"] + if handle is not None: + length = c_bst_ulong() + cptr = ctypes.POINTER(ctypes.c_char)() + _check_call( + _LIB.XGBoosterSerializeToBuffer( + self.handle, ctypes.byref(length), ctypes.byref(cptr) + ) + ) + buf = ctypes2buffer(cptr, length.value) + this["handle"] = buf + return this + + def __setstate__(self, state: Dict) -> None: + # reconstruct the handle from raw data + handle = state["handle"] + if handle is not None: + buf = handle + dmats = c_array(ctypes.c_void_p, []) + handle = ctypes.c_void_p() + _check_call( + _LIB.XGBoosterCreate(dmats, c_bst_ulong(0), ctypes.byref(handle)) + ) + length = c_bst_ulong(len(buf)) + ptr = (ctypes.c_char * len(buf)).from_buffer(buf) + _check_call(_LIB.XGBoosterUnserializeFromBuffer(handle, ptr, length)) + state["handle"] = handle + self.__dict__.update(state) + + def __getitem__(self, val: Union[Integer, tuple, slice, EllipsisType]) -> "Booster": + """Get a slice of the tree-based model. Attributes like `best_iteration` and + `best_score` are removed in the resulting booster. + + .. versionadded:: 1.3.0 + + """ + # convert to slice for all other types + if isinstance(val, (np.integer, int)): + val = slice(int(val), int(val + 1)) + if isinstance(val, EllipsisType): + val = slice(0, 0) + if isinstance(val, tuple): + raise ValueError("Only supports slicing through 1 dimension.") + # All supported types are now slice + if not isinstance(val, slice): + msg = _expect((int, slice, np.integer, EllipsisType), type(val)) + raise TypeError(msg) + + if isinstance(val.start, EllipsisType) or val.start is None: + start = 0 + else: + start = val.start + if isinstance(val.stop, EllipsisType) or val.stop is None: + stop = 0 + else: + stop = val.stop + if stop < start: + raise ValueError("Invalid slice", val) + + step = val.step if val.step is not None else 1 + + c_start = ctypes.c_int(start) + c_stop = ctypes.c_int(stop) + c_step = ctypes.c_int(step) + + sliced_handle = ctypes.c_void_p() + status = _LIB.XGBoosterSlice( + self.handle, c_start, c_stop, c_step, ctypes.byref(sliced_handle) + ) + if status == -2: + raise IndexError("Layer index out of range") + _check_call(status) + + sliced = Booster() + _check_call(_LIB.XGBoosterFree(sliced.handle)) + sliced.handle = sliced_handle + return sliced + + def __iter__(self) -> Generator["Booster", None, None]: + """Iterator method for getting individual trees. + + .. versionadded:: 2.0.0 + + """ + for i in range(0, self.num_boosted_rounds()): + yield self[i] + + def save_config(self) -> str: + """Output internal parameter configuration of Booster as a JSON + string. + + .. versionadded:: 1.0.0 + + """ + json_string = ctypes.c_char_p() + length = c_bst_ulong() + _check_call( + _LIB.XGBoosterSaveJsonConfig( + self.handle, ctypes.byref(length), ctypes.byref(json_string) + ) + ) + assert json_string.value is not None + result = json_string.value.decode() # pylint: disable=no-member + return result + + def load_config(self, config: str) -> None: + """Load configuration returned by `save_config`. + + .. versionadded:: 1.0.0 + """ + assert isinstance(config, str) + _check_call(_LIB.XGBoosterLoadJsonConfig(self.handle, c_str(config))) + + def __copy__(self) -> "Booster": + return self.__deepcopy__(None) + + def __deepcopy__(self, _: Any) -> "Booster": + """Return a copy of booster.""" + return Booster(model_file=self) + + def copy(self) -> "Booster": + """Copy the booster object. + + Returns + ------- + booster : + A copied booster model + """ + return copy.copy(self) + + def reset(self) -> "Booster": + """Reset the booster object to release data caches used for training. + + .. versionadded:: 3.0.0 + + """ + _check_call(_LIB.XGBoosterReset(self.handle)) + return self + + def attr(self, key: str) -> Optional[str]: + """Get attribute string from the Booster. + + Parameters + ---------- + key : + The key to get attribute from. + + Returns + ------- + value : + The attribute value of the key, returns None if attribute do not exist. + """ + ret = ctypes.c_char_p() + success = ctypes.c_int() + _check_call( + _LIB.XGBoosterGetAttr( + self.handle, c_str(key), ctypes.byref(ret), ctypes.byref(success) + ) + ) + if success.value != 0: + value = ret.value + assert value + return py_str(value) + return None + + def attributes(self) -> Dict[str, Optional[str]]: + """Get attributes stored in the Booster as a dictionary. + + Returns + ------- + result : dictionary of attribute_name: attribute_value pairs of strings. + Returns an empty dict if there's no attributes. + """ + length = c_bst_ulong() + sarr = ctypes.POINTER(ctypes.c_char_p)() + _check_call( + _LIB.XGBoosterGetAttrNames( + self.handle, ctypes.byref(length), ctypes.byref(sarr) + ) + ) + attr_names = from_cstr_to_pystr(sarr, length) + return {n: self.attr(n) for n in attr_names} + + def set_attr(self, **kwargs: Optional[Any]) -> None: + """Set the attribute of the Booster. + + Parameters + ---------- + **kwargs + The attributes to set. Setting a value to None deletes an attribute. + """ + for key, value in kwargs.items(): + c_value = None + if value is not None: + c_value = c_str(str(value)) + _check_call(_LIB.XGBoosterSetAttr(self.handle, c_str(key), c_value)) + + def _get_feature_info(self, field: str) -> Optional[FeatureInfo]: + length = c_bst_ulong() + sarr = ctypes.POINTER(ctypes.c_char_p)() + if not hasattr(self, "handle") or self.handle is None: + return None + _check_call( + _LIB.XGBoosterGetStrFeatureInfo( + self.handle, + c_str(field), + ctypes.byref(length), + ctypes.byref(sarr), + ) + ) + feature_info = from_cstr_to_pystr(sarr, length) + return feature_info if feature_info else None + + def _set_feature_info(self, features: Optional[FeatureInfo], field: str) -> None: + if features is not None: + assert isinstance(features, list) + feature_info_bytes = [bytes(f, encoding="utf-8") for f in features] + c_feature_info = (ctypes.c_char_p * len(feature_info_bytes))( + *feature_info_bytes + ) + _check_call( + _LIB.XGBoosterSetStrFeatureInfo( + self.handle, + c_str(field), + c_feature_info, + c_bst_ulong(len(features)), + ) + ) + else: + _check_call( + _LIB.XGBoosterSetStrFeatureInfo( + self.handle, c_str(field), None, c_bst_ulong(0) + ) + ) + + @property + def feature_types(self) -> Optional[FeatureTypes]: + """Feature types for this booster. Can be directly set by input data or by + assignment. See :py:class:`DMatrix` for details. + + """ + return self._get_feature_info("feature_type") + + @feature_types.setter + def feature_types(self, features: Optional[FeatureTypes]) -> None: + self._set_feature_info(features, "feature_type") + + @property + def feature_names(self) -> Optional[FeatureNames]: + """Feature names for this booster. Can be directly set by input data or by + assignment. + + """ + return self._get_feature_info("feature_name") + + @feature_names.setter + def feature_names(self, features: Optional[FeatureNames]) -> None: + self._set_feature_info(features, "feature_name") + + def get_categories(self, export_to_arrow: bool = False) -> Categories: + """Same method as :py:meth:`DMatrix.get_categories`.""" + + fnames = self.feature_names + n_features = self.num_features() + if fnames is None: + fnames = [str(i) for i in range(n_features)] + + hdl = ctypes.c_void_p() + if export_to_arrow: + arrow_arrays = _get_categories( + lambda ret: _LIB.XGBoosterGetCategoriesExportToArrow( + self.handle, None, ctypes.byref(hdl), ctypes.byref(ret) + ), + fnames, + n_features, + ) + else: + arrow_arrays = None + _check_call( + _LIB.XGBoosterGetCategories(self.handle, None, ctypes.byref(hdl)) + ) + + return Categories( + (hdl, lambda: _check_call(_LIB.XGBCategoriesFree(hdl))), arrow_arrays + ) + + def set_param( + self, + params: Union[Dict, Iterable[Tuple[str, Any]], str], + value: Optional[str] = None, + ) -> None: + """Set parameters into the Booster. + + Parameters + ---------- + params : + list of key,value pairs, dict of key to value or simply str key + value : + value of the specified parameter, when params is str key + """ + if isinstance(params, Mapping): + params = params.items() + elif isinstance(params, str) and value is not None: + params = [(params, value)] + for key, val in cast(Iterable[Tuple[str, str]], params): + if isinstance(val, np.ndarray): + val = val.tolist() + elif hasattr(val, "__cuda_array_interface__") and hasattr(val, "tolist"): + val = val.tolist() + if val is not None: + _check_call( + _LIB.XGBoosterSetParam(self.handle, c_str(key), c_str(str(val))) + ) + + def update( + self, + dtrain: DMatrix, + iteration: int, + fobj: Optional[CustomObj] = None, + ) -> None: + """Update for one iteration, with objective function calculated + internally. + + .. warning:: + + This function should not be called directly by users. + + Parameters + ---------- + dtrain : + Training data. + iteration : + The current training iteration. + fobj : + Custom objective function. + + """ + if not isinstance(dtrain, DMatrix): + raise TypeError(f"Invalid training matrix: {type(dtrain).__name__}") + self._assign_dmatrix_features(dtrain) + + if fobj is None: + _check_call( + _LIB.XGBoosterUpdateOneIter( + self.handle, ctypes.c_int(iteration), dtrain.handle + ) + ) + return + + # Forward the gradient calculation to the boost method. + self.boost( + dtrain, + iteration=iteration, + fobj=fobj, + ) + + def boost( + self, + dtrain: DMatrix, + iteration: int, + *, + grad: Optional[NumpyOrCupy] = None, + hess: Optional[NumpyOrCupy] = None, + fobj: Optional[CustomObj] = None, + ) -> None: + """Boost the booster for one iteration with customized gradient statistics. + + .. warning:: + + Like :py:meth:`.update`, this function should not be called directly by + users. + + Parameters + ---------- + dtrain : + The training DMatrix. + iteration : + The current training iteration. + grad : + The first order of gradient. + hess : + The second order of gradient. + fobj : + A custom objective function. If gradient is None, then an objective function + is required. + + """ + self._assign_dmatrix_features(dtrain) + + if all(arg is not None for arg in (grad, hess, fobj)): + raise ValueError( + "Provide either the objective, or the gradient and hessian, not both." + ) + n_samples = dtrain.num_row() + + def train_one_iter(grad: NumpyOrCupy, hess: NumpyOrCupy) -> None: + _check_call( + _LIB.XGBoosterTrainOneIter( + self.handle, + dtrain.handle, + iteration, + _grad_arrinf(grad, n_samples), + _grad_arrinf(hess, n_samples), + ) + ) + + if grad is not None or hess is not None: + # Handle the case where gradient is directly provided for compatibility with + # XGBoost < 3.2 + train_one_iter(grad, hess) + return + + if fobj is None: + raise ValueError( + "Invalid input for the boost function. Either the gradient or " + "the objective should have a valid value." + ) + + y_pred = self.predict(dtrain, output_margin=True, training=True) + + vgrad: Optional[ArrayLike] + vhess: Optional[ArrayLike] + + if isinstance(fobj, TreeObjective): + # full gradient for leaf values + vgrad, vhess = fobj(iteration, y_pred, dtrain) + # Reduced gradient for split nodes + split_grad = fobj.split_grad(iteration, vgrad, vhess) + # Switch the role of gradient if there's no split gradient but the tree + # objective is used. + if split_grad is not None: + sgrad, shess = split_grad + else: + sgrad, shess = vgrad, vhess + vgrad, vhess = None, None + elif isinstance(fobj, Objective): + sgrad, shess = fobj(iteration, y_pred, dtrain) + vgrad, vhess = None, None + else: + # Plain callable + sgrad, shess = fobj(y_pred, dtrain) + vgrad, vhess = None, None + + if vgrad is None: + train_one_iter(sgrad, shess) + return + + _check_call( + _LIB.XGBoosterTrainOneIterWithSplitGrad( + self.handle, + dtrain.handle, + iteration, + _grad_arrinf(sgrad, n_samples), + _grad_arrinf(shess, n_samples), + _grad_arrinf(vgrad, n_samples), + _grad_arrinf(vhess, n_samples), + ) + ) + + def eval_set( + self, + evals: Sequence[Tuple[DMatrix, str]], + iteration: int = 0, + feval: Optional[Metric] = None, + output_margin: bool = True, + ) -> str: + """Evaluate a set of data. + + Parameters + ---------- + evals : + List of items to be evaluated. + iteration : + Current iteration. + feval : + Custom evaluation function. + + Returns + ------- + result: str + Evaluation result string. + """ + for d in evals: + if not isinstance(d[0], DMatrix): + raise TypeError(f"expected DMatrix, got {type(d[0]).__name__}") + if not isinstance(d[1], str): + raise TypeError(f"expected string, got {type(d[1]).__name__}") + self._assign_dmatrix_features(d[0]) + + dmats = c_array(ctypes.c_void_p, [d[0].handle for d in evals]) + evnames = c_array(ctypes.c_char_p, [c_str(d[1]) for d in evals]) + msg = ctypes.c_char_p() + _check_call( + _LIB.XGBoosterEvalOneIter( + self.handle, + ctypes.c_int(iteration), + dmats, + evnames, + c_bst_ulong(len(evals)), + ctypes.byref(msg), + ) + ) + assert msg.value is not None + res = msg.value.decode() # pylint: disable=no-member + if feval is not None: + for dmat, evname in evals: + feval_ret = feval( + self.predict(dmat, training=False, output_margin=output_margin), + dmat, + ) + if isinstance(feval_ret, list): + for name, val in feval_ret: + # pylint: disable=consider-using-f-string + res += "\t%s-%s:%f" % (evname, name, val) + else: + name, val = feval_ret + # pylint: disable=consider-using-f-string + res += "\t%s-%s:%f" % (evname, name, val) + return res + + def eval(self, data: DMatrix, name: str = "eval", iteration: int = 0) -> str: + """Evaluate the model on mat. + + Parameters + ---------- + data : + The dmatrix storing the input. + + name : + The name of the dataset. + + iteration : + The current iteration number. + + Returns + ------- + result: str + Evaluation result string. + """ + self._assign_dmatrix_features(data) + return self.eval_set([(data, name)], iteration) + + # pylint: disable=too-many-function-args + @_deprecate_positional_args + def predict( + self, + data: DMatrix, + *, + output_margin: bool = False, + pred_leaf: bool = False, + pred_contribs: bool = False, + approx_contribs: bool = False, + pred_interactions: bool = False, + validate_features: bool = True, + training: bool = False, + iteration_range: IterationRange = (0, 0), + strict_shape: bool = False, + ) -> np.ndarray: + """Predict with data. The full model will be used unless `iteration_range` is + specified, meaning users have to either slice the model or use the + ``best_iteration`` attribute to get prediction from best model returned from + early stopping. + + .. note:: + + See :doc:`Prediction ` for issues like thread safety and a + summary of outputs from this function. + + Parameters + ---------- + data : + The dmatrix storing the input. + + output_margin : + Whether to output the raw untransformed margin value. + + pred_leaf : + When this option is on, the output will be a matrix of (nsample, + ntrees) with each record indicating the predicted leaf index of + each sample in each tree. Note that the leaf index of a tree is + unique per tree, so you may find leaf 1 in both tree 1 and tree 0. + + pred_contribs : + When this is True the output will be a matrix of size (nsample, + nfeats + 1) with each record indicating the feature contributions + (SHAP values) for that prediction. The sum of all feature + contributions is equal to the raw untransformed margin value of the + prediction. Note the final column is the bias term. + + approx_contribs : + Approximate the contributions of each feature. Used when ``pred_contribs`` + or ``pred_interactions`` is set to True. Changing the default of this + parameter (False) is not recommended. + + pred_interactions : + When this is True the output will be a matrix of size (nsample, + nfeats + 1, nfeats + 1) indicating the SHAP interaction values for + each pair of features. The sum of each row (or column) of the + interaction values equals the corresponding SHAP value (from + pred_contribs), and the sum of the entire matrix equals the raw + untransformed margin value of the prediction. Note the last row and + column correspond to the bias term. + + validate_features : + When this is True, validate that the Booster's and data's + feature_names are identical. Otherwise, it is assumed that the + feature_names are the same. + + training : + Whether the prediction value is used for training. This can effect `dart` + booster, which performs dropouts during training iterations but use all + trees for inference. If you want to obtain result with dropouts, set this + parameter to `True`. Also, the parameter is set to true when obtaining + prediction for custom objective function. + + .. versionadded:: 1.0.0 + + iteration_range : + Specifies which layer of trees are used in prediction. For example, if a + random forest is trained with 100 rounds. Specifying `iteration_range=(10, + 20)`, then only the forests built during [10, 20) (half open set) rounds are + used in this prediction. + + .. versionadded:: 1.4.0 + + strict_shape : + When set to True, output shape is invariant to whether classification is + used. For both value and margin prediction, the output shape is (n_samples, + n_groups), n_groups == 1 when multi-class is not used. Default to False, in + which case the output shape can be (n_samples, ) if multi-class is not used. + + .. versionadded:: 1.4.0 + + Returns + ------- + prediction : numpy array + + """ + if not isinstance(data, DMatrix): + raise TypeError("Expecting data to be a DMatrix object, got: ", type(data)) + if validate_features: + fn = data.feature_names + self._validate_features(fn) + args = { + "type": 0, + "training": training, + "iteration_begin": int(iteration_range[0]), + "iteration_end": int(iteration_range[1]), + "strict_shape": strict_shape, + } + + def assign_type(t: int) -> None: + if args["type"] != 0: + raise ValueError("One type of prediction at a time.") + args["type"] = t + + if output_margin: + assign_type(1) + if pred_contribs: + assign_type(2 if not approx_contribs else 3) + if pred_interactions: + assign_type(4 if not approx_contribs else 5) + if pred_leaf: + assign_type(6) + preds = ctypes.POINTER(ctypes.c_float)() + shape = ctypes.POINTER(c_bst_ulong)() + dims = c_bst_ulong() + _check_call( + _LIB.XGBoosterPredictFromDMatrix( + self.handle, + data.handle, + from_pystr_to_cstr(json.dumps(args)), + ctypes.byref(shape), + ctypes.byref(dims), + ctypes.byref(preds), + ) + ) + return _prediction_output(shape, dims, preds, False) + + # pylint: disable=too-many-statements + @_deprecate_positional_args + def inplace_predict( + self, + data: DataType, + *, + iteration_range: IterationRange = (0, 0), + predict_type: str = "value", + missing: float = np.nan, + validate_features: bool = True, + base_margin: Any = None, + strict_shape: bool = False, + ) -> NumpyOrCupy: + """Run prediction in-place when possible, Unlike :py:meth:`predict` method, + inplace prediction does not cache the prediction result. + + Calling only ``inplace_predict`` in multiple threads is safe and lock + free. But the safety does not hold when used in conjunction with other + methods. E.g. you can't train the booster in one thread and perform + prediction in the other. + + .. note:: + + If the device ordinal of the input data doesn't match the one configured for + the booster, data will be copied to the booster device. + + .. code-block:: python + + booster.set_param({"device": "cuda:0"}) + booster.inplace_predict(cupy_array) + + booster.set_param({"device": "cpu"}) + booster.inplace_predict(numpy_array) + + .. versionadded:: 1.1.0 + + Parameters + ---------- + data : + The input data. + iteration_range : + See :py:meth:`predict` for details. + predict_type : + * `value` Output model prediction values. + * `margin` Output the raw untransformed margin value. + missing : + See :py:obj:`xgboost.DMatrix` for details. + validate_features: + See :py:meth:`xgboost.Booster.predict` for details. + base_margin: + See :py:obj:`xgboost.DMatrix` for details. + + .. versionadded:: 1.4.0 + + strict_shape: + See :py:meth:`xgboost.Booster.predict` for details. + + .. versionadded:: 1.4.0 + + Returns + ------- + prediction : numpy.ndarray/cupy.ndarray + The prediction result. When input data is on GPU, prediction result is + stored in a cupy array. + + """ + preds = ctypes.POINTER(ctypes.c_float)() + + # once caching is supported, we can pass id(data) as cache id. + args = make_jcargs( + type=1 if predict_type == "margin" else 0, + training=False, + iteration_begin=int(iteration_range[0]), + iteration_end=int(iteration_range[1]), + missing=missing, + strict_shape=strict_shape, + cache_id=0, + ) + shape = ctypes.POINTER(c_bst_ulong)() + dims = c_bst_ulong() + + if base_margin is not None: + proxy: Optional[_ProxyDMatrix] = _ProxyDMatrix() + assert proxy is not None + proxy.set_info(base_margin=base_margin) + p_handle = proxy.handle + else: + proxy = None + p_handle = ctypes.c_void_p() + assert proxy is None or isinstance(proxy, _ProxyDMatrix) + + from .data import ( + ArrowTransformed, + PandasTransformed, + _is_arrow, + _is_cudf_df, + _is_cudf_pandas, + _is_list, + _is_np_array_like, + _is_pandas_df, + _is_pandas_series, + _is_polars, + _is_polars_series, + _is_tuple, + _transform_arrow_table, + _transform_pandas_df, + _transform_polars_df, + ) + + if _is_cudf_pandas(data): + data = data._fsproxy_fast # pylint: disable=protected-access + + enable_categorical = True + if _is_arrow(data): + data, fns, _ = _transform_arrow_table(data, enable_categorical, None, None) + if _is_polars_series(data): + pl = import_polars() + data = pl.DataFrame({data.name: data}) + if _is_polars(data): + data, fns, _ = _transform_polars_df(data, enable_categorical, None, None) + if _is_pandas_series(data): + import pandas as pd + + data = pd.DataFrame(data) + if _is_pandas_df(data): + data, fns, _ = _transform_pandas_df(data, enable_categorical) + if validate_features: + self._validate_features(fns) + if _is_list(data) or _is_tuple(data): + data = np.array(data) + + if validate_features: + if not hasattr(data, "shape"): + raise TypeError( + "`shape` attribute is required when `validate_features` is True" + f", got: {type(data)}" + ) + if len(data.shape) != 1 and self.num_features() != data.shape[1]: + raise ValueError( + f"Feature shape mismatch, expected: {self.num_features()}, " + f"got {data.shape[1]}" + ) + + if _is_np_array_like(data): + data, _ = _ensure_np_dtype(data, data.dtype) + _check_call( + _LIB.XGBoosterPredictFromDense( + self.handle, + array_interface(data), + args, + p_handle, + ctypes.byref(shape), + ctypes.byref(dims), + ctypes.byref(preds), + ) + ) + return _prediction_output(shape, dims, preds, False) + if isinstance(data, (ArrowTransformed, PandasTransformed)): + _check_call( + _LIB.XGBoosterPredictFromColumnar( + self.handle, + data.array_interface(), + args, + p_handle, + ctypes.byref(shape), + ctypes.byref(dims), + ctypes.byref(preds), + ) + ) + return _prediction_output(shape, dims, preds, False) + if isinstance(data, scipy.sparse.csr_matrix): + from .data import transform_scipy_sparse + + data = transform_scipy_sparse(data, True) + _check_call( + _LIB.XGBoosterPredictFromCSR( + self.handle, + array_interface(data.indptr), + array_interface(data.indices), + array_interface(data.data), + c_bst_ulong(data.shape[1]), + args, + p_handle, + ctypes.byref(shape), + ctypes.byref(dims), + ctypes.byref(preds), + ) + ) + return _prediction_output(shape, dims, preds, False) + if _is_cupy_alike(data): + from .data import _transform_cupy_array + + data = _transform_cupy_array(data) + interface_str = cuda_array_interface(data) + _check_call( + _LIB.XGBoosterPredictFromCudaArray( + self.handle, + interface_str, + args, + p_handle, + ctypes.byref(shape), + ctypes.byref(dims), + ctypes.byref(preds), + ) + ) + return _prediction_output(shape, dims, preds, True) + if _is_cudf_df(data): + from .data import _transform_cudf_df + + df, fns, _ = _transform_cudf_df(data, None, None, enable_categorical) + if validate_features: + self._validate_features(fns) + _check_call( + _LIB.XGBoosterPredictFromCudaColumnar( + self.handle, + df.array_interface(), + args, + p_handle, + ctypes.byref(shape), + ctypes.byref(dims), + ctypes.byref(preds), + ) + ) + return _prediction_output(shape, dims, preds, True) + + raise TypeError( + "Data type:" + str(type(data)) + " not supported by inplace prediction." + ) + + def save_model(self, fname: PathLike) -> None: + """Save the model to a file. + + The model is saved in an XGBoost internal format which is universal among the + various XGBoost interfaces. Auxiliary attributes of the Python Booster object + (such as feature_names) are only saved when using JSON or UBJSON (default) + format. Also, parameters that are not part of the model (like metrics, + `max_depth`, etc) are not saved, see :doc:`Model IO ` + for more info. + + .. code-block:: python + + model.save_model("model.json") + # or + model.save_model("model.ubj") + + Parameters + ---------- + fname : + Output file name + + """ + if isinstance(fname, (str, os.PathLike)): # assume file name + fname = os.fspath(os.path.expanduser(fname)) + _check_call(_LIB.XGBoosterSaveModel(self.handle, c_str(fname))) + else: + raise TypeError("fname must be a string or os PathLike") + + def save_raw(self, raw_format: str = "ubj") -> bytearray: + """Save the model to a in memory buffer representation instead of file. + + The model is saved in an XGBoost internal format which is universal among the + various XGBoost interfaces. Auxiliary attributes of the Python Booster object + (such as feature_names) are only saved when using JSON or UBJSON (default) + format. Also, parameters that are not part of the model (like metrics, + `max_depth`, etc) are not saved, see :doc:`Model IO ` + for more info. + + Parameters + ---------- + raw_format : + Format of output buffer. Can be `json` or `ubj`. + + Returns + ------- + An in memory buffer representation of the model + + """ + length = c_bst_ulong() + cptr = ctypes.POINTER(ctypes.c_char)() + config = make_jcargs(format=raw_format) + _check_call( + _LIB.XGBoosterSaveModelToBuffer( + self.handle, config, ctypes.byref(length), ctypes.byref(cptr) + ) + ) + return ctypes2buffer(cptr, length.value) + + def load_model(self, fname: ModelIn) -> None: + """Load the model from a file or a bytearray. + + The model is saved in an XGBoost internal format which is universal among the + various XGBoost interfaces. Auxiliary attributes of the Python Booster object + (such as feature_names) are only saved when using JSON or UBJSON (default) + format. Also, parameters that are not part of the model (like metrics, + `max_depth`, etc) are not saved, see :doc:`Model IO ` + for more info. + + .. code-block:: python + + model.save_model("model.json") + model.load_model("model.json") + + # or + model.save_model("model.ubj") + model.load_model("model.ubj") + + # or + buf = model.save_raw() + model.load_model(buf) + + Parameters + ---------- + fname : + Input file name or memory buffer(see also save_raw) + + """ + + def is_pathlike(path: ModelIn) -> TypeGuard[os.PathLike[str]]: + return isinstance(path, os.PathLike) + + if isinstance(fname, str) or is_pathlike(fname): + # assume file name, cannot use os.path.exist to check, file can be from URL. + fname = os.fspath(os.path.expanduser(fname)) + _check_call(_LIB.XGBoosterLoadModel(self.handle, c_str(fname))) + elif isinstance(fname, bytearray): + buf = fname + length = c_bst_ulong(len(buf)) + ptr = (ctypes.c_char * len(buf)).from_buffer(buf) + _check_call(_LIB.XGBoosterLoadModelFromBuffer(self.handle, ptr, length)) + else: + raise TypeError("Unknown file type: ", fname) + + @property + def best_iteration(self) -> int: + """The best iteration during training.""" + best = self.attr("best_iteration") + if best is not None: + return int(best) + + raise AttributeError( + "`best_iteration` is only defined when early stopping is used." + ) + + @best_iteration.setter + def best_iteration(self, iteration: int) -> None: + self.set_attr(best_iteration=iteration) + + @property + def best_score(self) -> float: + """The best evaluation score during training.""" + best = self.attr("best_score") + if best is not None: + return float(best) + + raise AttributeError( + "`best_score` is only defined when early stopping is used." + ) + + @best_score.setter + def best_score(self, score: int) -> None: + self.set_attr(best_score=score) + + def num_boosted_rounds(self) -> int: + """Get number of boosted rounds. For gblinear this is reset to 0 after + serializing the model. + + """ + rounds = ctypes.c_int() + assert self.handle is not None + _check_call(_LIB.XGBoosterBoostedRounds(self.handle, ctypes.byref(rounds))) + return rounds.value + + def num_features(self) -> int: + """Number of features in booster.""" + features = c_bst_ulong() + assert self.handle is not None + _check_call(_LIB.XGBoosterGetNumFeature(self.handle, ctypes.byref(features))) + return features.value + + def dump_model( + self, + fout: PathLike, + fmap: PathLike = "", + with_stats: bool = False, + dump_format: str = "text", + ) -> None: + """Dump model into a text or JSON file. Unlike :py:meth:`save_model`, the + output format is primarily used for visualization or interpretation, + hence it's more human readable but cannot be loaded back to XGBoost. + + Parameters + ---------- + fout : + Output file name. + fmap : + Name of the file containing feature map names. + with_stats : + Controls whether the split statistics are output. + dump_format : + Format of model dump file. Can be 'text' or 'json'. + """ + if isinstance(fout, (str, os.PathLike)): + fout = os.fspath(os.path.expanduser(fout)) + # pylint: disable=consider-using-with + fout_obj = open(fout, "w", encoding="utf-8") + need_close = True + else: + fout_obj = fout + need_close = False + ret = self.get_dump(fmap, with_stats, dump_format) + if dump_format == "json": + fout_obj.write("[\n") + for i, val in enumerate(ret): + fout_obj.write(val) + if i < len(ret) - 1: + fout_obj.write(",\n") + fout_obj.write("\n]") + else: + for i, val in enumerate(ret): + fout_obj.write(f"booster[{i}]:\n") + fout_obj.write(val) + if need_close: + fout_obj.close() + + def get_dump( + self, + fmap: PathLike = "", + with_stats: bool = False, + dump_format: str = "text", + ) -> List[str]: + """Returns the model dump as a list of strings. Unlike :py:meth:`save_model`, + the output format is primarily used for visualization or interpretation, hence + it's more human readable but cannot be loaded back to XGBoost. + + Parameters + ---------- + fmap : + Name of the file containing feature map names. + with_stats : + Controls whether the split statistics should be included. + dump_format : + Format of model dump. Can be 'text', 'json' or 'dot'. + + """ + fmap = os.fspath(os.path.expanduser(fmap)) + length = c_bst_ulong() + sarr = ctypes.POINTER(ctypes.c_char_p)() + _check_call( + _LIB.XGBoosterDumpModelEx( + self.handle, + c_str(fmap), + ctypes.c_int(with_stats), + c_str(dump_format), + ctypes.byref(length), + ctypes.byref(sarr), + ) + ) + res = from_cstr_to_pystr(sarr, length) + return res + + def get_fscore(self, fmap: PathLike = "") -> Dict[str, Union[float, List[float]]]: + """Get feature importance of each feature. + + .. note:: Zero-importance features will not be included + + Keep in mind that this function does not include zero-importance feature, + i.e. those features that have not been used in any split conditions. + + Parameters + ---------- + fmap : + The name of feature map file + """ + + return self.get_score(fmap, importance_type="weight") + + def get_score( + self, fmap: PathLike = "", importance_type: str = "weight" + ) -> Dict[str, Union[float, List[float]]]: + """Get feature importance of each feature. + For tree model Importance type can be defined as: + + * 'weight': the number of times a feature is used to split the data across all + trees. + * 'gain': the average gain across all splits the feature is used in. + * 'cover': the average coverage across all splits the feature is used in. + * 'total_gain': the total gain across all splits the feature is used in. + * 'total_cover': the total coverage across all splits the feature is used in. + + .. note:: + + For linear model, only "weight" is defined and it's the normalized + coefficients without bias. + + .. note:: Zero-importance features will not be included + + Keep in mind that this function does not include zero-importance feature, + i.e. those features that have not been used in any split conditions. + + Parameters + ---------- + fmap : + The name of feature map file. + importance_type : + One of the importance types defined above. + + Returns + ------- + A map between feature names and their scores. When `gblinear` is used for + multi-class classification the scores for each feature is a list with length + `n_classes`, otherwise they're scalars. + """ + fmap = os.fspath(os.path.expanduser(fmap)) + features = ctypes.POINTER(ctypes.c_char_p)() + scores = ctypes.POINTER(ctypes.c_float)() + n_out_features = c_bst_ulong() + out_dim = c_bst_ulong() + shape = ctypes.POINTER(c_bst_ulong)() + + _check_call( + _LIB.XGBoosterFeatureScore( + self.handle, + make_jcargs(importance_type=importance_type, feature_map=fmap), + ctypes.byref(n_out_features), + ctypes.byref(features), + ctypes.byref(out_dim), + ctypes.byref(shape), + ctypes.byref(scores), + ) + ) + features_arr = from_cstr_to_pystr(features, n_out_features) + scores_arr = _prediction_output(shape, out_dim, scores, False) + + results: Dict[str, Union[float, List[float]]] = {} + if len(scores_arr.shape) > 1 and scores_arr.shape[1] > 1: + for feat, score in zip(features_arr, scores_arr): + results[feat] = [float(s) for s in score] + else: + for feat, score in zip(features_arr, scores_arr): + results[feat] = float(score) + return results + + # pylint: disable=too-many-locals, too-many-statements + def trees_to_dataframe(self, fmap: PathLike = "") -> PdDataFrame: + """Parse a boosted tree model into a pandas DataFrame. + + This feature is only defined when the decision tree model is chosen as base + learner (`booster in {gbtree, dart}`). + + Parameters + ---------- + fmap : + Deprecated and ignored. The method uses :attr:`.feature_types` and + :attr:`.feature_names`. + + .. deprecated:: 3.4.0 + + """ + if not is_pandas_available(): + raise ImportError("pandas must be available to use this method.") + + from pandas import DataFrame + + model = json.loads(self.save_raw(raw_format="json")) + gbm = model["learner"]["gradient_booster"] + if gbm["name"] not in {"gbtree", "dart"}: + raise ValueError( + f"This method is not defined for Booster type {gbm['name']}" + ) + trees = gbm["model"]["trees"] + tree_info = gbm["model"]["tree_info"] + if fmap: + warnings.warn( + "`fmap` has been deprecated. XGBoost now uses the" + " `Booster.feature_names` and `Booster.feature_types` by default.", + FutureWarning, + ) + fnames = dict(enumerate(self.feature_names or [])) + ftypes = dict(enumerate(self.feature_types or [])) + + tree_ids: List[int] = [] + target_ids: List[Optional[int]] = [] + node_ids: List[int] = [] + fids: List[str] = [] + splits: List[float | None] = [] + categories: List[None | List[int]] = [] + y_directs: List[str | None] = [] + n_directs: List[str | None] = [] + missings: List[str | None] = [] + gains: List[float] = [] + covers: List[float] = [] + + for tid, tree in enumerate(trees): + left = tree["left_children"] + right = tree["right_children"] + sindex = tree["split_indices"] + sconds = tree["split_conditions"] + dft_left = tree["default_left"] + stypes = tree["split_type"] + loss_chg = tree["loss_changes"] + sum_hess = tree["sum_hessian"] + # For vector-leaf trees the leaf output is a vector stored in + # ``leaf_weights`` rather than in ``split_conditions``. The right child is + # re-used to store the leaf index. + size_leaf_vector = int(tree["tree_param"]["size_leaf_vector"]) + is_vector = size_leaf_vector > 1 + leaf_weights = tree["leaf_weights"] if is_vector else None + # The output group this tree contributes to. + tree_target = int(tree_info[tid]) + + # Node id -> category codes, stored CSR-style. + node_cats: Dict[int, List[int]] = {} + for nidx, beg, size in zip( + tree["categories_nodes"], + tree["categories_segments"], + tree["categories_sizes"], + ): + node_cats[nidx] = tree["categories"][beg : beg + size] + + # Depth-first traversal from the root so that pruned/deleted nodes, which + # remain in the arrays but are unreachable, are skipped. + stack = [0] + while stack: + nid = stack.pop() + if left[nid] == -1: # leaf + if is_vector: + assert leaf_weights is not None + beg = right[nid] * size_leaf_vector + leaf_rows = list( + enumerate( + float(v) + for v in leaf_weights[beg : beg + size_leaf_vector] + ) + ) + else: + leaf_rows = [(tree_target, float(sconds[nid]))] + for target, value in leaf_rows: + tree_ids.append(tid) + target_ids.append(target) + node_ids.append(nid) + fids.append("Leaf") + splits.append(None) + categories.append(None) + y_directs.append(None) + n_directs.append(None) + missings.append(None) + gains.append(value) + covers.append(sum_hess[nid]) + continue + + stack.append(left[nid]) + stack.append(right[nid]) + tree_ids.append(tid) + # A vector-leaf split is shared by all targets, so it has no single + # target; scalar-tree splits belong to the tree's target. + target_ids.append(None if is_vector else tree_target) + node_ids.append(nid) + fidx = sindex[nid] + fids.append(fnames.get(fidx, f"f{fidx}")) + dft_child = left[nid] if dft_left[nid] else right[nid] + if stypes[nid] == 1: # categorical + yes, no = right[nid], left[nid] + splits.append(None) + categories.append(node_cats.get(nid)) + elif ftypes.get(fidx, "q") == "i": # indicator (boolean) + # No split threshold; missing follows the "no" (default) branch. + yes, no = (right[nid] if dft_left[nid] else left[nid]), dft_child + splits.append(None) + categories.append(None) + else: # numerical + yes, no = left[nid], right[nid] + if ftypes.get(fidx, "q") == "int": + splits.append(float(math.ceil(sconds[nid]))) + else: + splits.append(sconds[nid]) + categories.append(None) + y_directs.append(f"{tid}-{yes}") + n_directs.append(f"{tid}-{no}") + missings.append(f"{tid}-{dft_child}") + gains.append(loss_chg[nid]) + covers.append(sum_hess[nid]) + + ids = [f"{t_id}-{n_id}" for t_id, n_id in zip(tree_ids, node_ids)] + df = DataFrame( + { + "Tree": tree_ids, + "Target": target_ids, + "Node": node_ids, + "ID": ids, + "Feature": fids, + "Split": splits, + "Yes": y_directs, + "No": n_directs, + "Missing": missings, + "Gain": gains, + "Cover": covers, + "Category": categories, + } + ) + df = df.astype( + { + "Tree": "Int64", + "Target": "Int64", + "Node": "Int64", + "ID": "string", + "Feature": "string", + "Split": "Float64", + "Yes": "string", + "No": "string", + "Missing": "string", + "Gain": "Float64", + "Cover": "Float64", + # `Category` holds Python lists + } + ) + + return df.sort_values(["Tree", "Node", "Target"]).reset_index(drop=True) + + def _assign_dmatrix_features(self, data: DMatrix) -> None: + if data.num_row() == 0: + return + + fn = data.feature_names + ft = data.feature_types + + if self.feature_names is None: + self.feature_names = fn + if self.feature_types is None: + self.feature_types = ft + + self._validate_features(fn) + + def _validate_features(self, feature_names: Optional[FeatureNames]) -> None: + if self.feature_names is None: + return + + if feature_names is None and self.feature_names is not None: + raise ValueError( + "data did not contain feature names, but the following fields are expected: " + + ", ".join(self.feature_names) + ) + + if self.feature_names != feature_names: + dat_missing = set(cast(FeatureNames, self.feature_names)) - set( + cast(FeatureNames, feature_names) + ) + my_missing = set(cast(FeatureNames, feature_names)) - set( + cast(FeatureNames, self.feature_names) + ) + + msg = "feature_names mismatch: {0} {1}" + + if dat_missing: + msg += ( + "\nexpected " + + ", ".join(str(s) for s in dat_missing) + + " in input data" + ) + + if my_missing: + msg += ( + "\ntraining data did not have the following fields: " + + ", ".join(str(s) for s in my_missing) + ) + + raise ValueError(msg.format(self.feature_names, feature_names)) + + def get_split_value_histogram( + self, + feature: str, + fmap: PathLike = "", + bins: Optional[int] = None, + as_pandas: bool = True, + ) -> Union[np.ndarray, PdDataFrame]: + """Get split value histogram of a feature + + Parameters + ---------- + feature : + The name of the feature. + fmap: + The name of feature map file. + bin : + The maximum number of bins. + Number of bins equals number of unique split values n_unique, + if bins == None or bins > n_unique. + as_pandas : + Return pd.DataFrame when pandas is installed. + If False or pandas is not installed, return numpy ndarray. + + Returns + ------- + a histogram of used splitting values for the specified feature + either as numpy array or pandas DataFrame. + """ + from .data import CAT_T + + xgdump = self.get_dump(fmap=fmap) + values = [] + # pylint: disable=consider-using-f-string + regexp = re.compile(r"\[{0}<([\d.Ee+-]+)\]".format(feature)) + for val in xgdump: + m = re.findall(regexp, val) + values.extend([float(x) for x in m]) + + n_unique = len(np.unique(values)) + bins = max(min(n_unique, bins) if bins is not None else n_unique, 1) + + nph = np.histogram(values, bins=bins) + nph_stacked = np.column_stack((nph[1][1:], nph[0])) + nph_stacked = nph_stacked[nph_stacked[:, 1] > 0] + + if nph_stacked.size == 0: + ft = self.feature_types + fn = self.feature_names + if fn is None: + # Let xgboost generate the feature names. + fn = [f"f{i}" for i in range(self.num_features())] + try: + index = fn.index(feature) + feature_t: Optional[str] = cast(List[str], ft)[index] + except (ValueError, AttributeError, TypeError): + # None.index: attr err, None[0]: type err, fn.index(-1): value err + feature_t = None + if feature_t == CAT_T: # categorical + raise ValueError( + "Split value historgam doesn't support categorical split." + ) + + if as_pandas and is_pandas_available(): + from pandas import DataFrame + + return DataFrame(nph_stacked, columns=["SplitValue", "Count"]) + if as_pandas and not is_pandas_available(): + warnings.warn( + "Returning histogram as ndarray" + " (as_pandas == True, but pandas is not installed).", + UserWarning, + ) + return nph_stacked diff --git a/v12x/pylib/xgboost/dask/__init__.py b/v12x/pylib/xgboost/dask/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..198da726c69553aa1c36c5d6b3e5e38411285c58 --- /dev/null +++ b/v12x/pylib/xgboost/dask/__init__.py @@ -0,0 +1,2264 @@ +# pylint: disable=too-many-arguments, too-many-locals +# pylint: disable=missing-class-docstring +# pylint: disable=too-many-lines +# pylint: disable=duplicate-code +""" +Dask extensions for distributed training +---------------------------------------- + +See :doc:`Distributed XGBoost with Dask ` for simple tutorial. Also +:doc:`/python/dask-examples/index` for some examples. + +There are two sets of APIs in this module, one is the functional API including +``train`` and ``predict`` methods. Another is stateful Scikit-Learner wrapper +inherited from single-node Scikit-Learn interface. + +The implementation is heavily influenced by dask_xgboost: +https://github.com/dask/dask-xgboost + +Optional dask configuration +=========================== + +- **coll_cfg**: + Specify the scheduler address along with communicator configurations. This can be + used as a replacement of the existing global Dask configuration + `xgboost.scheduler_address` (see below). See :ref:`tracker-ip` for more info. The + `tracker_host_ip` should specify the IP address of the Dask scheduler node. + + .. versionadded:: 3.0.0 + + .. code-block:: python + + from xgboost import dask as dxgb + from xgboost.collective import Config + + coll_cfg = Config( + retry=1, timeout=20, tracker_host_ip="10.23.170.98", tracker_port=0 + ) + + clf = dxgb.DaskXGBClassifier(coll_cfg=coll_cfg) + # or + dxgb.train(client, {}, Xy, num_boost_round=10, coll_cfg=coll_cfg) + +- **xgboost.scheduler_address**: Specify the scheduler address + + .. versionadded:: 1.6.0 + + .. deprecated:: 3.0.0 + + .. code-block:: python + + dask.config.set({"xgboost.scheduler_address": "192.0.0.100"}) + # We can also specify the port. + dask.config.set({"xgboost.scheduler_address": "192.0.0.100:12345"}) + +""" + +import logging +from collections import defaultdict +from contextlib import contextmanager +from functools import partial, update_wrapper +from threading import Thread +from typing import ( + Any, + Awaitable, + Callable, + Dict, + Generator, + Iterable, + List, + Optional, + ParamSpec, + Self, + Sequence, + Set, + Tuple, + TypeAlias, + TypedDict, + TypeGuard, + TypeVar, + Union, +) + +import dask +import distributed +import numpy +from dask import array as da +from dask import bag as db +from dask import dataframe as dd +from dask.delayed import Delayed +from distributed import Future + +from .. import collective, config +from .._data_utils import Categories +from .._typing import FeatureNames, FeatureTypes, IterationRange +from ..callback import TrainingCallback +from ..collective import Config as CollConfig +from ..collective import _Args as CollArgs +from ..collective import _ArgVals as CollArgsVals +from ..compat import _is_cudf_df, _is_cudf_ser, _is_cupy_alike +from ..core import ( + Booster, + CustomObj, + DMatrix, + Metric, + XGBoostError, + _check_distributed_params, + _deprecate_positional_args, + _expect, +) +from ..sklearn import ( + XGBClassifier, + XGBClassifierMixIn, + XGBModel, + XGBRanker, + XGBRankerMixIn, + XGBRegressorMixIn, + _can_use_qdm, + _check_rf_callback, + _cls_predict_proba, + _objective_decorator, + _warn_rf_deprecated, + _wrap_evaluation_matrices, + xgboost_model_doc, +) +from ..tracker import RabitTracker +from ..training import train as worker_train +from .data import _get_dmatrices, no_group_split +from .utils import _DASK_2024_12_1, _DASK_2025_3_0, get_address_from_user, get_n_threads + +_DaskCollection: TypeAlias = Union[da.Array, dd.DataFrame, dd.Series] +_DataT: TypeAlias = Union[da.Array, dd.DataFrame] # do not use series as predictor +TrainReturnT = TypedDict( + "TrainReturnT", + { + "booster": Booster, + "history": Dict, + }, +) + +__all__ = [ + "CommunicatorContext", + "DaskDMatrix", + "DaskQuantileDMatrix", + "DaskXGBRegressor", + "DaskXGBClassifier", + "DaskXGBRanker", + "DaskXGBRFRegressor", + "DaskXGBRFClassifier", + "train", + "predict", + "inplace_predict", +] + +# TODOs: +# - CV +# +# Note for developers: +# +# As of writing asyncio is still a new feature of Python and in depth documentation is +# rare. Best examples of various asyncio tricks are in dask (luckily). Classes like +# Client, Worker are awaitable. Some general rules for the implementation here: +# +# - Synchronous world is different from asynchronous one, and they don't mix well. +# - Write everything with async, then use distributed Client sync function to do the +# switch. +# - Use Any for type hint when the return value can be union of Awaitable and plain +# value. This is caused by Client.sync can return both types depending on +# context. Right now there's no good way to silent: +# +# await train(...) +# +# if train returns an Union type. + + +LOGGER = logging.getLogger("[xgboost.dask]") + + +def _try_start_tracker( + n_workers: int, + addrs: List[Union[Optional[str], Optional[Tuple[str, int]]]], + timeout: Optional[int], +) -> CollArgs: + env: CollArgs = {} + try: + if isinstance(addrs[0], tuple): + host_ip = addrs[0][0] + port = addrs[0][1] + rabit_tracker = RabitTracker( + n_workers=n_workers, + host_ip=host_ip, + port=port, + sortby="task", + timeout=0 if timeout is None else timeout, + ) + else: + addr = addrs[0] + assert isinstance(addr, str) or addr is None + rabit_tracker = RabitTracker( + n_workers=n_workers, + host_ip=addr, + sortby="task", + timeout=0 if timeout is None else timeout, + ) + + rabit_tracker.start() + # No timeout since we don't want to abort the training + thread = Thread(target=rabit_tracker.wait_for) + thread.daemon = True + thread.start() + env.update(rabit_tracker.worker_args()) + + except XGBoostError as e: + if len(addrs) < 2: + raise + LOGGER.warning( + "Failed to bind address '%s', trying to use '%s' instead. Error:\n %s", + str(addrs[0]), + str(addrs[1]), + str(e), + ) + env = _try_start_tracker(n_workers, addrs[1:], timeout) + + return env + + +def _start_tracker( + n_workers: int, + addr_from_dask: Optional[str], + addr_from_user: Optional[Tuple[str, int]], + timeout: Optional[int], +) -> CollArgs: + """Start Rabit tracker, recurse to try different addresses.""" + env = _try_start_tracker(n_workers, [addr_from_user, addr_from_dask], timeout) + return env + + +class CommunicatorContext(collective.CommunicatorContext): + """A context controlling collective communicator initialization and finalization.""" + + def __init__(self, **args: CollArgsVals) -> None: + super().__init__(**args) + + worker = distributed.get_worker() + # We use task ID for rank assignment which makes the RABIT rank consistent (but + # not the same as task ID is string and "10" is sorted before "2") with dask + # worker name. This outsources the rank assignment to dask and prevents + # non-deterministic issue. + self.args["DMLC_TASK_ID"] = f"[xgboost.dask-{worker.name}]:{worker.address}" + + +def _get_client(client: Optional["distributed.Client"]) -> "distributed.Client": + """Simple wrapper around testing None.""" + if not isinstance(client, (type(distributed.get_client()), type(None))): + raise TypeError( + _expect([type(distributed.get_client()), type(None)], type(client)) + ) + ret = distributed.get_client() if client is None else client + return ret + + +# From the implementation point of view, DaskDMatrix complicates a lots of +# things. A large portion of the code base is about syncing and extracting +# stuffs from DaskDMatrix. But having an independent data structure gives us a +# chance to perform some specialized optimizations, like building histogram +# index directly. + + +class DaskDMatrix: + # pylint: disable=too-many-instance-attributes + """DMatrix holding on references to Dask DataFrame or Dask Array. Constructing a + `DaskDMatrix` forces all lazy computation to be carried out. Wait for the input + data explicitly if you want to see actual computation of constructing `DaskDMatrix`. + + See doc for :py:obj:`xgboost.DMatrix` constructor for other parameters. DaskDMatrix + accepts only dask collection. + + .. note:: + + `DaskDMatrix` does not repartition or move data between workers. It's the + caller's responsibility to balance the data. + + .. note:: + + For aligning partitions with ranking query groups, use the + :py:class:`DaskXGBRanker` and its ``allow_group_split`` option. + + .. versionadded:: 1.0.0 + + Parameters + ---------- + client : + Specify the dask client used for training. Use default client returned from + dask if it's set to None. + + """ + + @_deprecate_positional_args + def __init__( + self, + client: Optional["distributed.Client"], + data: _DataT, + label: Optional[_DaskCollection] = None, + *, + weight: Optional[_DaskCollection] = None, + base_margin: Optional[_DaskCollection] = None, + missing: Optional[float] = None, + silent: bool = False, # pylint: disable=unused-argument + feature_names: Optional[FeatureNames] = None, + feature_types: Optional[FeatureTypes] = None, + group: Optional[_DaskCollection] = None, + qid: Optional[_DaskCollection] = None, + label_lower_bound: Optional[_DaskCollection] = None, + label_upper_bound: Optional[_DaskCollection] = None, + feature_weights: Optional[_DaskCollection] = None, + enable_categorical: bool = True, + ) -> None: + client = _get_client(client) + + self.feature_names = feature_names + self.feature_types = feature_types + if isinstance(feature_types, Categories): + raise TypeError( + "The Dask interface can handle categories from DataFrame automatically." + ) + self.missing = missing if missing is not None else numpy.nan + self.enable_categorical = enable_categorical + + if qid is not None and weight is not None: + raise NotImplementedError("per-group weight is not implemented.") + if group is not None: + raise NotImplementedError( + "group structure is not implemented, use qid instead." + ) + + if len(data.shape) != 2: + raise ValueError(f"Expecting 2 dimensional input, got: {data.shape}") + + if not isinstance(data, (dd.DataFrame, da.Array)): + raise TypeError(_expect((dd.DataFrame, da.Array), type(data))) + if not isinstance(label, (dd.DataFrame, da.Array, dd.Series, type(None))): + raise TypeError(_expect((dd.DataFrame, da.Array, dd.Series), type(label))) + + self._n_cols = data.shape[1] + assert isinstance(self._n_cols, int) + self.worker_map: Dict[str, List[Future]] = defaultdict(list) + self.is_quantile: bool = False + + self._init = client.sync( + self._map_local_data, + client=client, + data=data, + label=label, + weights=weight, + base_margin=base_margin, + qid=qid, + feature_weights=feature_weights, + label_lower_bound=label_lower_bound, + label_upper_bound=label_upper_bound, + ) + + def __await__(self) -> Generator[None, None, "DaskDMatrix"]: + return self._init.__await__() + + async def _map_local_data( + self, + *, + client: "distributed.Client", + data: _DataT, + label: Optional[_DaskCollection] = None, + weights: Optional[_DaskCollection] = None, + base_margin: Optional[_DaskCollection] = None, + qid: Optional[_DaskCollection] = None, + feature_weights: Optional[_DaskCollection] = None, + label_lower_bound: Optional[_DaskCollection] = None, + label_upper_bound: Optional[_DaskCollection] = None, + ) -> "DaskDMatrix": + """Obtain references to local data.""" + + def inconsistent( + left: List[Any], left_name: str, right: List[Any], right_name: str + ) -> str: + msg = ( + f"Partitions between {left_name} and {right_name} are not " + f"consistent: {len(left)} != {len(right)}. " + f"Please try to repartition/rechunk your data." + ) + return msg + + def to_futures(d: _DaskCollection) -> List[Future]: + """Breaking data into partitions.""" + d = client.persist(d) + if ( + hasattr(d.partitions, "shape") + and len(d.partitions.shape) > 1 + and d.partitions.shape[1] > 1 + ): + raise ValueError( + "Data should be" + " partitioned by row. To avoid this specify the number" + " of columns for your dask Array explicitly. e.g." + " chunks=(partition_size, -1])" + ) + return client.futures_of(d) + + def flatten_meta(meta: Optional[_DaskCollection]) -> Optional[List[Future]]: + if meta is not None: + meta_parts: List[Future] = to_futures(meta) + return meta_parts + return None + + X_parts = to_futures(data) + y_parts = flatten_meta(label) + w_parts = flatten_meta(weights) + margin_parts = flatten_meta(base_margin) + qid_parts = flatten_meta(qid) + ll_parts = flatten_meta(label_lower_bound) + lu_parts = flatten_meta(label_upper_bound) + + parts: Dict[str, List[Future]] = {"data": X_parts} + + def append_meta(m_parts: Optional[List[Future]], name: str) -> None: + if m_parts is not None: + assert len(X_parts) == len(m_parts), inconsistent( + X_parts, "X", m_parts, name + ) + parts[name] = m_parts + + append_meta(y_parts, "label") + append_meta(w_parts, "weight") + append_meta(margin_parts, "base_margin") + append_meta(qid_parts, "qid") + append_meta(ll_parts, "label_lower_bound") + append_meta(lu_parts, "label_upper_bound") + # At this point, `parts` looks like: + # [(x0, x1, ..), (y0, y1, ..), ..] in future form + + # turn into list of dictionaries. + packed_parts: List[Dict[str, Future]] = [] + for i in range(len(X_parts)): + part_dict: Dict[str, Future] = {} + for key, value in parts.items(): + part_dict[key] = value[i] + packed_parts.append(part_dict) + + # delay the zipped result + # pylint: disable=no-member + delayed_parts: List[Delayed] = list(map(dask.delayed, packed_parts)) + # At this point, the mental model should look like: + # [{"data": x0, "label": y0, ..}, {"data": x1, "label": y1, ..}, ..] + + # Convert delayed objects into futures and make sure they are realized + # + # This also makes partitions to align (co-locate) on workers (X_0, y_0 should be + # on the same worker). + fut_parts: List[Future] = client.compute(delayed_parts) + await distributed.wait(fut_parts) # async wait for parts to be computed + + for part in fut_parts: + # Each part is [{"data": x0, "label": y0, ..}, ...] in future form. + assert part.status == "finished", part.status + + # Preserving the partition order for prediction. + self.partition_order = {} + for i, part in enumerate(fut_parts): + self.partition_order[part.key] = i + + key_to_partition = {part.key: part for part in fut_parts} + who_has: Dict[str, Tuple[str, ...]] = await client.scheduler.who_has( + keys=[part.key for part in fut_parts] + ) + + worker_map: Dict[str, List[Future]] = defaultdict(list) + + for key, workers in who_has.items(): + worker_map[next(iter(workers))].append(key_to_partition[key]) + + self.worker_map = worker_map + + if feature_weights is None: + self.feature_weights = None + else: + self.feature_weights = await client.compute(feature_weights).result() + + return self + + def _create_fn_args(self, worker_addr: str) -> Dict[str, Any]: + """Create a dictionary of objects that can be pickled for function + arguments. + + """ + return { + "feature_names": self.feature_names, + "feature_types": self.feature_types, + "feature_weights": self.feature_weights, + "missing": self.missing, + "enable_categorical": self.enable_categorical, + "parts": self.worker_map.get(worker_addr, None), + "is_quantile": self.is_quantile, + } + + def num_col(self) -> int: + """Get the number of columns (features) in the DMatrix. + + Returns + ------- + number of columns + """ + return self._n_cols + + +_MapRetT = TypeVar("_MapRetT") +_P = ParamSpec("_P") + + +async def map_worker_partitions( + client: Optional["distributed.Client"], + func: Callable[_P, _MapRetT], + *refs: Any, + workers: Sequence[str], +) -> _MapRetT: + """Map a function onto partitions of each worker.""" + # Note for function purity: + # XGBoost is sensitive to data partition and uses random number generator. + client = _get_client(client) + futures = [] + for addr in workers: + args = [] + for ref in refs: + if isinstance(ref, DaskDMatrix): + # pylint: disable=protected-access + args.append(ref._create_fn_args(addr)) + else: + args.append(ref) + + def fn(_address: str, *args: _P.args, **kwargs: _P.kwargs) -> List[_MapRetT]: + worker = distributed.get_worker() + + if worker.address != _address: + raise ValueError( + f"Invalid worker address: {worker.address}, expecting {_address}. " + "This is likely caused by one of the workers died and Dask " + "re-scheduled a different one. Resilience is not yet supported." + ) + # Turn result into a list for bag construction + return [func(*args, **kwargs)] + + # XGBoost requires all workers running training tasks to be unique. Meaning, we + # can't run 2 training jobs on the same node. This at best leads to an error + # (NCCL unique check), at worst leads to extremely slow training performance + # without any warning. + # + # See disitributed.scheduler.decide_worker for `allow_other_workers`. In + # summary, the scheduler chooses a worker from the valid set that has the task + # dependencies. Each XGBoost's training task has all dependencies in a single + # worker. As a result, the right worker should be picked by the scheduler even + # if `allow_other_workers` is set to True. + # + # In addition, the scheduler only discards the valid set (the `workers` arg) if + # there's no candidate can be found. This is likely caused by killed workers. In + # that case, the check in `fn` should be able to stop the task. If we don't + # relax the constraint and prevent Dask from choosing an invalid worker, the + # task will simply hangs. We prefer a quick error here. + # + fut = client.submit( + update_wrapper(partial(fn, addr), fn), + *args, + pure=False, + workers=[addr], + allow_other_workers=True, + ) + futures.append(fut) + + def first_valid(results: Iterable[Optional[_MapRetT]]) -> Optional[_MapRetT]: + for v in results: + if v is not None: + return v + return None + + bag = db.from_delayed(futures) + fut = await bag.reduction(first_valid, first_valid) + result = await client.compute(fut).result() + + return result + + +class DaskQuantileDMatrix(DaskDMatrix): + """A dask version of :py:class:`QuantileDMatrix`. See :py:class:`DaskDMatrix` for + parameter documents. + + """ + + @_deprecate_positional_args + def __init__( + self, + client: Optional["distributed.Client"], + data: _DataT, + label: Optional[_DaskCollection] = None, + *, + weight: Optional[_DaskCollection] = None, + base_margin: Optional[_DaskCollection] = None, + missing: Optional[float] = None, + silent: bool = False, # disable=unused-argument + feature_names: Optional[FeatureNames] = None, + feature_types: Optional[Union[Any, List[Any]]] = None, + max_bin: Optional[int] = None, + ref: Optional[DaskDMatrix] = None, + group: Optional[_DaskCollection] = None, + qid: Optional[_DaskCollection] = None, + label_lower_bound: Optional[_DaskCollection] = None, + label_upper_bound: Optional[_DaskCollection] = None, + feature_weights: Optional[_DaskCollection] = None, + enable_categorical: bool = True, + max_quantile_batches: Optional[int] = None, + ) -> None: + super().__init__( + client=client, + data=data, + label=label, + weight=weight, + base_margin=base_margin, + group=group, + qid=qid, + label_lower_bound=label_lower_bound, + label_upper_bound=label_upper_bound, + missing=missing, + silent=silent, + feature_weights=feature_weights, + feature_names=feature_names, + feature_types=feature_types, + enable_categorical=enable_categorical, + ) + self.max_bin = max_bin + self.max_quantile_batches = max_quantile_batches + self.is_quantile = True + self._ref: Optional[int] = id(ref) if ref is not None else None + + def _create_fn_args(self, worker_addr: str) -> Dict[str, Any]: + args = super()._create_fn_args(worker_addr) + args["max_bin"] = self.max_bin + args["max_quantile_batches"] = self.max_quantile_batches + if self._ref is not None: + args["ref"] = self._ref + return args + + +async def _get_rabit_args( + client: "distributed.Client", + n_workers: int, + dconfig: Optional[Dict[str, Any]] = None, + coll_cfg: Optional[CollConfig] = None, +) -> Dict[str, Union[str, int]]: + """Get rabit context arguments from data distribution in DaskDMatrix.""" + # There are 3 possible different addresses: + # 1. Provided by user via dask.config + # 2. Guessed by xgboost `get_host_ip` function + # 3. From dask scheduler + # We try 1 and 3 if 1 is available, otherwise 2 and 3. + + # See if user config is available + coll_cfg = CollConfig() if coll_cfg is None else coll_cfg + host_ip: Optional[str] = None + port: int = 0 + host_ip, port = get_address_from_user(dconfig, coll_cfg) + + if host_ip is not None: + user_addr = (host_ip, port) + else: + user_addr = None + + # Try address from dask scheduler, this might not work, see + # https://github.com/dask/dask-xgboost/pull/40 + try: + sched_addr = distributed.comm.get_address_host(client.scheduler.address) + sched_addr = sched_addr.strip("/:") + except Exception: # pylint: disable=broad-except + sched_addr = None + + # We assume the scheduler is a fair process and run the tracker there. + env = await client.run_on_scheduler( + _start_tracker, n_workers, sched_addr, user_addr, coll_cfg.tracker_timeout + ) + env = coll_cfg.get_comm_config(env) + assert env is not None + return env + + +def _get_dask_config() -> Optional[Dict[str, Any]]: + return dask.config.get("xgboost", default=None) + + +# train and predict methods are supposed to be "functional", which meets the +# dask paradigm. But as a side effect, the `evals_result` in single-node API +# is no longer supported since it mutates the input parameter, and it's not +# intuitive to sync the mutation result. Therefore, a dictionary containing +# evaluation history is instead returned. + + +def _get_workers_from_data( + dtrain: DaskDMatrix, evals: Optional[Sequence[Tuple[DaskDMatrix, str]]] +) -> List[str]: + X_worker_map: Set[str] = set(dtrain.worker_map.keys()) + if evals: + for e in evals: + assert len(e) == 2 + assert isinstance(e[0], DaskDMatrix) and isinstance(e[1], str) + if e[0] is dtrain: + continue + worker_map = set(e[0].worker_map.keys()) + X_worker_map = X_worker_map.union(worker_map) + return list(X_worker_map) + + +async def _check_workers_are_alive( + workers: List[str], client: "distributed.Client" +) -> None: + info = await client.scheduler.identity() + current_workers = info["workers"].keys() + missing_workers = set(workers) - current_workers + if missing_workers: + raise RuntimeError(f"Missing required workers: {missing_workers}") + + +async def _train_async( + *, + client: "distributed.Client", + global_config: Dict[str, Any], + dconfig: Optional[Dict[str, Any]], + params: Dict[str, Any], + dtrain: DaskDMatrix, + num_boost_round: int, + evals: Optional[Sequence[Tuple[DaskDMatrix, str]]], + obj: Optional[CustomObj], + early_stopping_rounds: Optional[int], + verbose_eval: Union[int, bool], + xgb_model: Optional[Booster], + callbacks: Optional[Sequence[TrainingCallback]], + custom_metric: Optional[Metric], + coll_cfg: Optional[CollConfig], +) -> Optional[TrainReturnT]: + workers = _get_workers_from_data(dtrain, evals) + await _check_workers_are_alive(workers, client) + coll_args = await _get_rabit_args( + client, len(workers), dconfig=dconfig, coll_cfg=coll_cfg + ) + _check_distributed_params(params) + + # This function name is displayed in the Dask dashboard task status, let's make it + # clear that it's XGBoost training. + def do_train( # pylint: disable=too-many-positional-arguments + parameters: Dict, + coll_args: Dict[str, Union[str, int]], + train_id: int, + evals_name: List[str], + evals_id: List[int], + train_ref: dict, + *refs: dict, + ) -> Optional[TrainReturnT]: + worker = distributed.get_worker() + local_param = parameters.copy() + n_threads = get_n_threads(local_param, worker) + local_param.update({"nthread": n_threads, "n_jobs": n_threads}) + + local_history: TrainingCallback.EvalsLog = {} + global_config.update({"nthread": n_threads}) + + if coll_cfg is not None: + coll_args = coll_cfg.update_worker_args(coll_args) + + with CommunicatorContext(**coll_args), config.config_context(**global_config): + Xy, evals = _get_dmatrices( + train_ref, + train_id, + *refs, + evals_id=evals_id, + evals_name=evals_name, + n_threads=n_threads, + # We need the model for reference categories. + model=xgb_model, + ) + + booster = worker_train( + params=local_param, + dtrain=Xy, + num_boost_round=num_boost_round, + evals_result=local_history, + evals=evals if len(evals) != 0 else None, + obj=obj, + custom_metric=custom_metric, + early_stopping_rounds=early_stopping_rounds, + verbose_eval=verbose_eval, + xgb_model=xgb_model, + callbacks=callbacks, + ) + # Don't return the boosters from empty workers. It's quite difficult to + # guarantee everything is in sync in the present of empty workers, especially + # with complex objectives like quantile. + if Xy.num_row() != 0: + ret: Optional[TrainReturnT] = { + "booster": booster, + "history": local_history, + } + else: + ret = None + return ret + + async with distributed.MultiLock(workers, client): + if evals is not None: + evals_data = [d for d, n in evals] + evals_name = [n for d, n in evals] + evals_id = [id(d) for d in evals_data] + else: + evals_data = [] + evals_name = [] + evals_id = [] + + result = await map_worker_partitions( + client, + do_train, + # extra function parameters + params, + coll_args, + id(dtrain), + evals_name, + evals_id, + *([dtrain] + evals_data), + # workers to be used for training + workers=workers, + ) + return result + + +@_deprecate_positional_args +def train( # pylint: disable=unused-argument + client: "distributed.Client", + params: Dict[str, Any], + dtrain: DaskDMatrix, + num_boost_round: int = 10, + *, + evals: Optional[Sequence[Tuple[DaskDMatrix, str]]] = None, + obj: Optional[CustomObj] = None, + early_stopping_rounds: Optional[int] = None, + xgb_model: Optional[Booster] = None, + verbose_eval: Union[int, bool] = True, + callbacks: Optional[Sequence[TrainingCallback]] = None, + custom_metric: Optional[Metric] = None, + coll_cfg: Optional[CollConfig] = None, +) -> Any: + """Train XGBoost model. + + .. versionadded:: 1.0.0 + + .. note:: + + Other parameters are the same as :py:func:`xgboost.train` except for + `evals_result`, which is returned as part of function return value instead of + argument. + + Parameters + ---------- + client : + Specify the dask client used for training. Use default client returned from + dask if it's set to None. + + coll_cfg : + Configuration for the communicator used during training. See + :py:class:`~xgboost.collective.Config`. + + Returns + ------- + results: dict + A dictionary containing trained booster and evaluation history. `history` field + is the same as `eval_result` from `xgboost.train`. + + .. code-block:: python + + {'booster': xgboost.Booster, + 'history': {'train': {'logloss': ['0.48253', '0.35953']}, + 'eval': {'logloss': ['0.480385', '0.357756']}}} + + """ + client = _get_client(client) + return client.sync( + _train_async, + global_config=config.get_config(), + dconfig=_get_dask_config(), + **locals(), + ) + + +def _can_output_df(is_df: bool, output_shape: Tuple) -> bool: + return is_df and len(output_shape) <= 2 + + +def _maybe_dataframe( + data: Any, prediction: Any, columns: List[int], is_df: bool +) -> Any: + """Return dataframe for prediction when applicable.""" + if _can_output_df(is_df, prediction.shape): + # Need to preserve the index for dataframe. + # See issue: https://github.com/dmlc/xgboost/issues/6939 + # In older versions of dask, the partition is actually a numpy array when input + # is dataframe. + index = getattr(data, "index", None) + if _is_cudf_df(data): + import cudf + + if prediction.size == 0: + return cudf.DataFrame({}, columns=columns, dtype=numpy.float32) + + prediction = cudf.DataFrame( + prediction, columns=columns, dtype=numpy.float32, index=index + ) + else: + import pandas as pd + + if prediction.size == 0: + return pd.DataFrame( + {}, columns=columns, dtype=numpy.float32, index=index + ) + + prediction = pd.DataFrame( + prediction, columns=columns, dtype=numpy.float32, index=index + ) + return prediction + + +async def _direct_predict_impl( # pylint: disable=too-many-branches + *, + mapped_predict: Callable, + booster: "distributed.Future", + data: _DataT, + base_margin: Optional[_DaskCollection], + output_shape: Tuple[int, ...], + meta: Dict[int, str], +) -> _DaskCollection: + columns = tuple(meta.keys()) + if len(output_shape) >= 3 and isinstance(data, dd.DataFrame): + # Without this check, dask will finish the prediction silently even if output + # dimension is greater than 3. But during map_partitions, dask passes a + # `dd.DataFrame` as local input to xgboost, which is converted to csr_matrix by + # `_convert_unknown_data` since dd.DataFrame is not known to xgboost native + # binding. + raise ValueError( + "Use `da.Array` or `DaskDMatrix` when output has more than 2 dimensions." + ) + if _can_output_df(isinstance(data, dd.DataFrame), output_shape): + if base_margin is not None and isinstance(base_margin, da.Array): + # Easier for map_partitions + base_margin_df: Optional[Union[dd.DataFrame, dd.Series]] = ( + base_margin.to_dask_dataframe() + ) + else: + base_margin_df = base_margin + predictions = dd.map_partitions( + mapped_predict, + booster, + data, + True, + columns, + base_margin_df, + meta=dd.utils.make_meta(meta), + ) + # classification can return a dataframe, drop 1 dim when it's reg/binary + if len(output_shape) == 1: + predictions = predictions.iloc[:, 0] + else: + if base_margin is not None and isinstance( + base_margin, (dd.Series, dd.DataFrame) + ): + # Easier for map_blocks + base_margin_array: Optional[da.Array] = base_margin.to_dask_array() + else: + base_margin_array = base_margin + # Input data is 2-dim array, output can be 1(reg, binary)/2(multi-class, + # contrib)/3(contrib, interaction)/4(interaction) dims. + if len(output_shape) == 1: + drop_axis: Union[int, List[int]] = [1] # drop from 2 to 1 dim. + new_axis: Union[int, List[int]] = [] + else: + drop_axis = [] + if isinstance(data, dd.DataFrame): + new_axis = list(range(len(output_shape) - 2)) + else: + new_axis = [i + 2 for i in range(len(output_shape) - 2)] + if len(output_shape) >= 2: + # Dask cannot infer changed or newly introduced output dimensions. + # `None` is not supported for chunks in map_blocks, so use the inferred + # prediction shape for every non-row dimension. + assert isinstance(data, da.Array) + output_chunks: List[Tuple] = [data.chunks[0]] + output_chunks.extend((size,) for size in output_shape[1:]) + chunks: Optional[List[Tuple]] = output_chunks + else: + chunks = None + predictions = da.map_blocks( + mapped_predict, + booster, + data, + False, + columns, + base_margin_array, + chunks=chunks, + drop_axis=drop_axis, + new_axis=new_axis, + dtype=numpy.float32, + ) + return predictions + + +def _infer_predict_output( + booster: Booster, features: int, is_df: bool, inplace: bool, **kwargs: Any +) -> Tuple[Tuple[int, ...], Dict[int, str]]: + """Create a dummy test sample to infer output shape for prediction.""" + assert isinstance(features, int) + rng = numpy.random.RandomState(1994) + test_sample = rng.randn(1, features) + if inplace: + kwargs = kwargs.copy() + if kwargs.pop("predict_type") == "margin": + kwargs["output_margin"] = True + m = DMatrix(test_sample, enable_categorical=True) + # generated DMatrix doesn't have feature name, so no validation. + test_predt = booster.predict(m, validate_features=False, **kwargs) + n_columns = test_predt.shape[1] if len(test_predt.shape) > 1 else 1 + meta: Dict[int, str] = {} + if _can_output_df(is_df, test_predt.shape): + for i in range(n_columns): + meta[i] = "f4" + return test_predt.shape, meta + + +async def _get_model_future( + client: "distributed.Client", model: Union[Booster, Dict, "distributed.Future"] +) -> "distributed.Future": + # See https://github.com/dask/dask/issues/11179#issuecomment-2168094529 for the use + # of hash. + # https://github.com/dask/distributed/pull/8796 Don't use broadcast in the `scatter` + # call, otherwise, the predict function might hang. + if isinstance(model, Booster): + booster = await client.scatter(model, hash=False) + elif isinstance(model, dict): + booster = await client.scatter(model["booster"], hash=False) + elif isinstance(model, distributed.Future): + booster = model + t = booster.type + if t is not Booster: + raise TypeError( + f"Underlying type of model future should be `Booster`, got {t}" + ) + else: + raise TypeError(_expect([Booster, dict, distributed.Future], type(model))) + return booster + + +# pylint: disable=too-many-statements +async def _predict_async( + client: "distributed.Client", + global_config: Dict[str, Any], + model: Union[Booster, Dict, "distributed.Future"], + data: _DataT, + *, + output_margin: bool, + missing: float, + pred_leaf: bool, + pred_contribs: bool, + approx_contribs: bool, + pred_interactions: bool, + validate_features: bool, + iteration_range: IterationRange, + strict_shape: bool, +) -> _DaskCollection: + _booster = await _get_model_future(client, model) + if not isinstance(data, (DaskDMatrix, da.Array, dd.DataFrame)): + raise TypeError(_expect([DaskDMatrix, da.Array, dd.DataFrame], type(data))) + + def mapped_predict( + booster: Booster, partition: Any, is_df: bool, columns: List[int], _: Any + ) -> Any: + with config.config_context(**global_config): + m = DMatrix( + data=partition, + missing=missing, + enable_categorical=True, + ) + predt = booster.predict( + data=m, + output_margin=output_margin, + pred_leaf=pred_leaf, + pred_contribs=pred_contribs, + approx_contribs=approx_contribs, + pred_interactions=pred_interactions, + validate_features=validate_features, + iteration_range=iteration_range, + strict_shape=strict_shape, + ) + predt = _maybe_dataframe(partition, predt, columns, is_df) + return predt + + # Predict on dask collection directly. + if isinstance(data, (da.Array, dd.DataFrame)): + _output_shape, meta = await client.compute( + client.submit( + _infer_predict_output, + _booster, + features=data.shape[1], + is_df=isinstance(data, dd.DataFrame), + inplace=False, + output_margin=output_margin, + pred_leaf=pred_leaf, + pred_contribs=pred_contribs, + approx_contribs=approx_contribs, + pred_interactions=pred_interactions, + strict_shape=strict_shape, + ) + ) + return await _direct_predict_impl( + mapped_predict=mapped_predict, + booster=_booster, + data=data, + base_margin=None, + output_shape=_output_shape, + meta=meta, + ) + + output_shape, _ = await client.compute( + client.submit( + _infer_predict_output, + booster=_booster, + features=data.num_col(), + is_df=False, + inplace=False, + output_margin=output_margin, + pred_leaf=pred_leaf, + pred_contribs=pred_contribs, + approx_contribs=approx_contribs, + pred_interactions=pred_interactions, + strict_shape=strict_shape, + ) + ) + # Prediction on dask DMatrix. + partition_order = data.partition_order + feature_names = data.feature_names + feature_types = data.feature_types + missing = data.missing + + def dispatched_predict(booster: Booster, part: Dict[str, Any]) -> numpy.ndarray: + data = part["data"] + base_margin = part.get("base_margin", None) + with config.config_context(**global_config): + m = DMatrix( + data, + missing=missing, + base_margin=base_margin, + feature_names=feature_names, + feature_types=feature_types, + enable_categorical=True, + ) + predt = booster.predict( + m, + output_margin=output_margin, + pred_leaf=pred_leaf, + pred_contribs=pred_contribs, + approx_contribs=approx_contribs, + pred_interactions=pred_interactions, + validate_features=validate_features, + iteration_range=iteration_range, + strict_shape=strict_shape, + ) + return predt + + all_parts = [] + all_orders = [] + all_shapes = [] + all_workers: List[str] = [] + workers_address = list(data.worker_map.keys()) + for worker_addr in workers_address: + list_of_parts = data.worker_map[worker_addr] + all_parts.extend(list_of_parts) + all_workers.extend(len(list_of_parts) * [worker_addr]) + all_orders.extend([partition_order[part.key] for part in list_of_parts]) + for w, part in zip(all_workers, all_parts): + s = client.submit(lambda part: part["data"].shape[0], part, workers=[w]) + all_shapes.append(s) + + parts_with_order = list(zip(all_parts, all_shapes, all_orders, all_workers)) + parts_with_order = sorted(parts_with_order, key=lambda p: p[2]) + all_parts = [part for part, shape, order, w in parts_with_order] + all_shapes = [shape for part, shape, order, w in parts_with_order] + all_workers = [w for part, shape, order, w in parts_with_order] + + futures = [] + for w, part in zip(all_workers, all_parts): + f = client.submit(dispatched_predict, _booster, part, workers=[w]) + futures.append(f) + + # Constructing a dask array from list of numpy arrays + # See https://docs.dask.org/en/latest/array-creation.html + arrays = [] + all_shapes = await client.gather(all_shapes) + for i, rows in enumerate(all_shapes): + arrays.append( + da.from_delayed( + futures[i], shape=(rows,) + output_shape[1:], dtype=numpy.float32 + ) + ) + predictions = da.concatenate(arrays, axis=0) + return predictions + + +@_deprecate_positional_args +def predict( # pylint: disable=unused-argument + client: Optional["distributed.Client"], + model: Union[TrainReturnT, Booster, "distributed.Future"], + data: Union[DaskDMatrix, _DataT], + *, + output_margin: bool = False, + missing: float = numpy.nan, + pred_leaf: bool = False, + pred_contribs: bool = False, + approx_contribs: bool = False, + pred_interactions: bool = False, + validate_features: bool = True, + iteration_range: IterationRange = (0, 0), + strict_shape: bool = False, +) -> Any: + """Run prediction with a trained booster. + + .. note:: + + Using ``inplace_predict`` might be faster when some features are not needed. + See :py:meth:`xgboost.Booster.predict` for details on various parameters. When + output has more than 2 dimensions (shap value, leaf with strict_shape), input + should be ``da.Array`` or ``DaskDMatrix``. + + .. versionadded:: 1.0.0 + + Parameters + ---------- + client: + Specify the dask client used for training. Use default client + returned from dask if it's set to None. + model: + The trained model. It can be a distributed.Future so user can + pre-scatter it onto all workers. + data: + Input data used for prediction. When input is a dataframe object, + prediction output is a series. + missing: + Used when input data is not DaskDMatrix. Specify the value + considered as missing. + + Returns + ------- + prediction: dask.array.Array/dask.dataframe.Series + When input data is ``dask.array.Array`` or ``DaskDMatrix``, the return value is + an array, when input data is ``dask.dataframe.DataFrame``, return value can be + ``dask.dataframe.Series``, ``dask.dataframe.DataFrame``, depending on the output + shape. + + """ + client = _get_client(client) + return client.sync(_predict_async, global_config=config.get_config(), **locals()) + + +async def _inplace_predict_async( # pylint: disable=too-many-branches + *, + client: "distributed.Client", + global_config: Dict[str, Any], + model: Union[Booster, Dict, "distributed.Future"], + data: _DataT, + iteration_range: IterationRange, + predict_type: str, + missing: float, + validate_features: bool, + base_margin: Optional[_DaskCollection], + strict_shape: bool, +) -> _DaskCollection: + client = _get_client(client) + booster = await _get_model_future(client, model) + if not isinstance(data, (da.Array, dd.DataFrame)): + raise TypeError(_expect([da.Array, dd.DataFrame], type(data))) + if base_margin is not None and not isinstance( + data, (da.Array, dd.DataFrame, dd.Series) + ): + raise TypeError(_expect([da.Array, dd.DataFrame, dd.Series], type(base_margin))) + + def mapped_predict( + booster: Booster, + partition: Any, + is_df: bool, + columns: List[int], + base_margin: Any, + ) -> Any: + with config.config_context(**global_config): + prediction = booster.inplace_predict( + partition, + iteration_range=iteration_range, + predict_type=predict_type, + missing=missing, + base_margin=base_margin, + validate_features=validate_features, + strict_shape=strict_shape, + ) + prediction = _maybe_dataframe(partition, prediction, columns, is_df) + return prediction + + # await turns future into value. + shape, meta = await client.compute( + client.submit( + _infer_predict_output, + booster, + features=data.shape[1], + is_df=isinstance(data, dd.DataFrame), + inplace=True, + predict_type=predict_type, + iteration_range=iteration_range, + strict_shape=strict_shape, + ) + ) + return await _direct_predict_impl( + mapped_predict=mapped_predict, + booster=booster, + data=data, + base_margin=base_margin, + output_shape=shape, + meta=meta, + ) + + +@_deprecate_positional_args +def inplace_predict( # pylint: disable=unused-argument + client: Optional["distributed.Client"], + model: Union[TrainReturnT, Booster, "distributed.Future"], + data: _DataT, + *, + iteration_range: IterationRange = (0, 0), + predict_type: str = "value", + missing: float = numpy.nan, + validate_features: bool = True, + base_margin: Optional[_DaskCollection] = None, + strict_shape: bool = False, +) -> Any: + """Inplace prediction. See doc in :py:meth:`xgboost.Booster.inplace_predict` for + details. + + .. versionadded:: 1.1.0 + + Parameters + ---------- + client: + Specify the dask client used for training. Use default client + returned from dask if it's set to None. + model: + See :py:func:`xgboost.dask.predict` for details. + data : + dask collection. + iteration_range: + See :py:meth:`xgboost.Booster.predict` for details. + predict_type: + See :py:meth:`xgboost.Booster.inplace_predict` for details. + missing: + Value in the input data which needs to be present as a missing + value. If None, defaults to np.nan. + base_margin: + See :py:obj:`xgboost.DMatrix` for details. + + .. versionadded:: 1.4.0 + + strict_shape: + See :py:meth:`xgboost.Booster.predict` for details. + + .. versionadded:: 1.4.0 + + Returns + ------- + prediction : + When input data is ``dask.array.Array``, the return value is an array, when + input data is ``dask.dataframe.DataFrame``, return value can be + ``dask.dataframe.Series``, ``dask.dataframe.DataFrame``, depending on the output + shape. + + """ + client = _get_client(client) + # When used in asynchronous environment, the `client` object should have + # `asynchronous` attribute as True. When invoked by the skl interface, it's + # responsible for setting up the client. + return client.sync( + _inplace_predict_async, global_config=config.get_config(), **locals() + ) + + +async def _async_wrap_evaluation_matrices( + client: Optional["distributed.Client"], + device: Optional[str], + tree_method: Optional[str], + max_bin: Optional[int], + **kwargs: Any, +) -> Tuple[DaskDMatrix, Optional[List[Tuple[DaskDMatrix, str]]]]: + """A switch function for async environment.""" + + def _dispatch(ref: Optional[DaskDMatrix], **kwargs: Any) -> DaskDMatrix: + if _can_use_qdm(tree_method, device): + return DaskQuantileDMatrix( + client=client, ref=ref, max_bin=max_bin, **kwargs + ) + return DaskDMatrix(client=client, **kwargs) + + train_dmatrix, evals = _wrap_evaluation_matrices(create_dmatrix=_dispatch, **kwargs) + train_dmatrix = await train_dmatrix + if evals is None: + return train_dmatrix, evals + awaited = [] + for e in evals: + if e[0] is train_dmatrix: # already awaited + awaited.append(e) + continue + awaited.append((await e[0], e[1])) + return train_dmatrix, awaited + + +@contextmanager +def _set_worker_client( + model: "DaskScikitLearnBase", client: "distributed.Client" +) -> Generator: + """Temporarily set the client for sklearn model.""" + try: + model.client = client + yield model + finally: + model.client = None # type: ignore[assignment] + + +class DaskScikitLearnBase(XGBModel): + """Base class for implementing scikit-learn interface with Dask""" + + _client = None + + def __init__(self, *, coll_cfg: Optional[CollConfig] = None, **kwargs: Any) -> None: + super().__init__(**kwargs) + + self.coll_cfg = coll_cfg + + async def _predict_async( + self, + data: _DataT, + *, + output_margin: bool, + validate_features: bool, + base_margin: Optional[_DaskCollection], + iteration_range: Optional[IterationRange], + ) -> Any: + iteration_range = self._get_iteration_range(iteration_range) + # Dask doesn't support gblinear and accepts only Dask collection types (array + # and dataframe). We can perform inplace predict. + assert self._can_use_inplace_predict() + predts = await inplace_predict( + client=self.client, + model=self.get_booster(), + data=data, + iteration_range=iteration_range, + predict_type="margin" if output_margin else "value", + missing=self.missing, + base_margin=base_margin, + validate_features=validate_features, + ) + if isinstance(predts, dd.DataFrame): + predts = predts.to_dask_array() + # Make sure the booster is part of the task graph implicitly + # only needed for certain versions of dask. + if _DASK_2024_12_1() and not _DASK_2025_3_0(): + # Fixes this issue for dask>=2024.1.1,<2025.3.0 + # Dask==2025.3.0 fails with: + # RuntimeError: Attempting to use an asynchronous + # Client in a synchronous context of `dask.compute` + # + # Dask==2025.4.0 fails with: + # TypeError: Value type is not supported for data + # iterator: + predts = predts.persist() + return predts + + @_deprecate_positional_args + def predict( + self, + X: _DataT, + *, + output_margin: bool = False, + validate_features: bool = True, + base_margin: Optional[_DaskCollection] = None, + iteration_range: Optional[IterationRange] = None, + ) -> Any: + return self.client.sync( + self._predict_async, + X, + output_margin=output_margin, + validate_features=validate_features, + base_margin=base_margin, + iteration_range=iteration_range, + ) + + async def _apply_async( + self, + X: _DataT, + iteration_range: Optional[IterationRange] = None, + ) -> Any: + iteration_range = self._get_iteration_range(iteration_range) + test_dmatrix: DaskDMatrix = await DaskDMatrix( + self.client, + data=X, + missing=self.missing, + feature_types=self.feature_types, + ) + predts = await predict( + self.client, + model=self.get_booster(), + data=test_dmatrix, + pred_leaf=True, + iteration_range=iteration_range, + ) + return predts + + def apply( + self, + X: _DataT, + iteration_range: Optional[IterationRange] = None, + ) -> Any: + return self.client.sync(self._apply_async, X, iteration_range=iteration_range) + + def __await__(self) -> Awaitable[Any]: + # Generate a coroutine wrapper to make this class awaitable. + async def _() -> Awaitable[Any]: + return self + + return self._client_sync(_).__await__() + + def __getstate__(self) -> Dict: + this = self.__dict__.copy() + if "_client" in this: + del this["_client"] + return this + + @property + def client(self) -> "distributed.Client": + """The dask client used in this model. The `Client` object can not be + serialized for transmission, so if task is launched from a worker instead of + directly from the client process, this attribute needs to be set at that worker. + + """ + + client = _get_client(self._client) + return client + + @client.setter + def client(self, clt: "distributed.Client") -> None: + # calling `worker_client' doesn't return the correct `asynchronous` attribute, + # so we have to pass it ourselves. + self._asynchronous = clt.asynchronous if clt is not None else False + self._client = clt + + def _client_sync(self, func: Callable, **kwargs: Any) -> Any: + """Get the correct client, when method is invoked inside a worker we + should use `worker_client' instead of default client. + + """ + + if self._client is None: + asynchronous = getattr(self, "_asynchronous", False) + try: + distributed.get_worker() + in_worker = True + except ValueError: + in_worker = False + if in_worker: + with distributed.worker_client() as client: + with _set_worker_client(self, client) as this: + ret = this.client.sync( + func, **kwargs, asynchronous=asynchronous + ) + return ret + return ret + + return self.client.sync(func, **kwargs, asynchronous=self.client.asynchronous) + + +@xgboost_model_doc( + """Implementation of the Scikit-Learn API for XGBoost.""", ["estimators", "model"] +) +class DaskXGBRegressor(XGBRegressorMixIn, DaskScikitLearnBase): + """dummy doc string to workaround pylint, replaced by the decorator.""" + + async def _fit_async( + self, + X: _DataT, + y: _DaskCollection, + *, + sample_weight: Optional[_DaskCollection], + base_margin: Optional[_DaskCollection], + eval_set: Optional[Sequence[Tuple[_DaskCollection, _DaskCollection]]], + sample_weight_eval_set: Optional[Sequence[_DaskCollection]], + base_margin_eval_set: Optional[Sequence[_DaskCollection]], + verbose: Union[int, bool], + xgb_model: Optional[Union[Booster, XGBModel]], + feature_weights: Optional[_DaskCollection], + ) -> Self: + params = self.get_xgb_params() + model, metric, params, feature_weights = self._configure_fit( + xgb_model, params, feature_weights + ) + + dtrain, evals = await _async_wrap_evaluation_matrices( + client=self.client, + device=self.device, + tree_method=self.tree_method, + max_bin=self.max_bin, + X=X, + y=y, + group=None, + qid=None, + sample_weight=sample_weight, + base_margin=base_margin, + feature_weights=feature_weights, + eval_set=eval_set, + sample_weight_eval_set=sample_weight_eval_set, + base_margin_eval_set=base_margin_eval_set, + eval_group=None, + eval_qid=None, + missing=self.missing, + enable_categorical=self.enable_categorical, + feature_types=self.feature_types, + ) + + if callable(self.objective): + obj: Optional[CustomObj] = _objective_decorator(self.objective) + else: + obj = None + results = await self.client.sync( + _train_async, + asynchronous=True, + client=self.client, + global_config=config.get_config(), + dconfig=_get_dask_config(), + params=params, + dtrain=dtrain, + num_boost_round=self.get_num_boosting_rounds(), + evals=evals, + obj=obj, + custom_metric=metric, + verbose_eval=verbose, + early_stopping_rounds=self.early_stopping_rounds, + callbacks=self.callbacks, + coll_cfg=self.coll_cfg, + xgb_model=model, + ) + self._Booster = results["booster"] + self._set_evaluation_result(results["history"]) + return self + + # pylint: disable=missing-docstring, disable=unused-argument + @_deprecate_positional_args + def fit( + self, + X: _DataT, + y: _DaskCollection, + *, + sample_weight: Optional[_DaskCollection] = None, + base_margin: Optional[_DaskCollection] = None, + eval_set: Optional[Sequence[Tuple[_DaskCollection, _DaskCollection]]] = None, + verbose: Optional[Union[int, bool]] = True, + xgb_model: Optional[Union[Booster, str, XGBModel]] = None, + sample_weight_eval_set: Optional[Sequence[_DaskCollection]] = None, + base_margin_eval_set: Optional[Sequence[_DaskCollection]] = None, + feature_weights: Optional[_DaskCollection] = None, + ) -> "DaskXGBRegressor": + args = {k: v for k, v in locals().items() if k not in ("self", "__class__")} + return self._client_sync(self._fit_async, **args) + + +@xgboost_model_doc( + "Implementation of the scikit-learn API for XGBoost classification.", + ["estimators", "model"], +) +class DaskXGBClassifier(XGBClassifierMixIn, DaskScikitLearnBase): + # pylint: disable=missing-class-docstring + async def _fit_async( + self, + X: _DataT, + y: _DaskCollection, + *, + sample_weight: Optional[_DaskCollection], + base_margin: Optional[_DaskCollection], + eval_set: Optional[Sequence[Tuple[_DaskCollection, _DaskCollection]]], + sample_weight_eval_set: Optional[Sequence[_DaskCollection]], + base_margin_eval_set: Optional[Sequence[_DaskCollection]], + verbose: Union[int, bool], + xgb_model: Optional[Union[Booster, XGBModel]], + feature_weights: Optional[_DaskCollection], + ) -> Self: + params = self.get_xgb_params() + model, metric, params, feature_weights = self._configure_fit( + xgb_model, params, feature_weights + ) + + dtrain, evals = await _async_wrap_evaluation_matrices( + self.client, + device=self.device, + tree_method=self.tree_method, + max_bin=self.max_bin, + X=X, + y=y, + group=None, + qid=None, + sample_weight=sample_weight, + base_margin=base_margin, + feature_weights=feature_weights, + eval_set=eval_set, + sample_weight_eval_set=sample_weight_eval_set, + base_margin_eval_set=base_margin_eval_set, + eval_group=None, + eval_qid=None, + missing=self.missing, + enable_categorical=self.enable_categorical, + feature_types=self.feature_types, + ) + + # pylint: disable=attribute-defined-outside-init + if isinstance(y, dd.Series) or ( + isinstance(y, dd.DataFrame) and len(y.columns) == 1 + ): + labels_series = y if isinstance(y, dd.Series) else y[y.columns[0]] + # The result is collected on the client, so avoid a distributed shuffle + # using only 1 output partition. + classes = labels_series.drop_duplicates(split_out=1) + else: + labels = y if isinstance(y, da.Array) else y.to_dask_array() + if labels.ndim == 2: + n_columns = labels.shape[1] + assert isinstance(n_columns, int) + row_chunks = tuple(chunk * n_columns for chunk in labels.chunks[0]) + + def flatten(block: Any) -> Any: + return block.reshape(-1) + + # Use map_blocks to keep the computation lazy, so that we don't go + # through the data twice (shape and unique). + labels = da.map_blocks( + flatten, + labels, + chunks=(row_chunks,), + drop_axis=1, + dtype=labels.dtype, + ) + classes = da.unique(labels) + classes = await self.client.compute(classes) + + if _is_cudf_ser(classes): + classes = classes.to_cupy() + if _is_cupy_alike(classes): + classes = classes.get() + classes = numpy.unique(numpy.asarray(classes)) + self.n_classes_ = len(classes) + + if self.n_classes_ > 2: + params["objective"] = "multi:softprob" + params["num_class"] = self.n_classes_ + else: + params["objective"] = "binary:logistic" + + if callable(self.objective): + obj: Optional[CustomObj] = _objective_decorator(self.objective) + else: + obj = None + results = await self.client.sync( + _train_async, + asynchronous=True, + client=self.client, + global_config=config.get_config(), + dconfig=_get_dask_config(), + params=params, + dtrain=dtrain, + num_boost_round=self.get_num_boosting_rounds(), + evals=evals, + obj=obj, + custom_metric=metric, + verbose_eval=verbose, + early_stopping_rounds=self.early_stopping_rounds, + callbacks=self.callbacks, + coll_cfg=self.coll_cfg, + xgb_model=model, + ) + self._Booster = results["booster"] + if not callable(self.objective): + self.objective = params["objective"] + self._set_evaluation_result(results["history"]) + return self + + # pylint: disable=unused-argument + def fit( + self, + X: _DataT, + y: _DaskCollection, + *, + sample_weight: Optional[_DaskCollection] = None, + base_margin: Optional[_DaskCollection] = None, + eval_set: Optional[Sequence[Tuple[_DaskCollection, _DaskCollection]]] = None, + verbose: Optional[Union[int, bool]] = True, + xgb_model: Optional[Union[Booster, str, XGBModel]] = None, + sample_weight_eval_set: Optional[Sequence[_DaskCollection]] = None, + base_margin_eval_set: Optional[Sequence[_DaskCollection]] = None, + feature_weights: Optional[_DaskCollection] = None, + ) -> "DaskXGBClassifier": + args = {k: v for k, v in locals().items() if k not in ("self", "__class__")} + return self._client_sync(self._fit_async, **args) + + async def _predict_proba_async( + self, + X: _DataT, + validate_features: bool, + base_margin: Optional[_DaskCollection], + iteration_range: Optional[IterationRange], + ) -> _DaskCollection: + if self.objective == "multi:softmax": + raise ValueError( + "multi:softmax doesn't support `predict_proba`. " + "Switch to `multi:softproba` instead" + ) + predts = await super()._predict_async( + data=X, + output_margin=False, + validate_features=validate_features, + base_margin=base_margin, + iteration_range=iteration_range, + ) + vstack = update_wrapper( + partial(da.vstack, allow_unknown_chunksizes=True), da.vstack + ) + return _cls_predict_proba(getattr(self, "n_classes_", 0), predts, vstack) + + # pylint: disable=missing-function-docstring + def predict_proba( + self, + X: _DaskCollection, + validate_features: bool = True, + base_margin: Optional[_DaskCollection] = None, + iteration_range: Optional[IterationRange] = None, + ) -> Any: + return self._client_sync( + self._predict_proba_async, + X=X, + validate_features=validate_features, + base_margin=base_margin, + iteration_range=iteration_range, + ) + + if XGBClassifier.predict_proba.__doc__ is not None: + predict_proba.__doc__ = XGBClassifier.predict_proba.__doc__ + + async def _predict_async( + self, + data: _DataT, + *, + output_margin: bool, + validate_features: bool, + base_margin: Optional[_DaskCollection], + iteration_range: Optional[IterationRange], + ) -> _DaskCollection: + pred_probs = await super()._predict_async( + data, + output_margin=output_margin, + validate_features=validate_features, + base_margin=base_margin, + iteration_range=iteration_range, + ) + if output_margin: + return pred_probs + + if len(pred_probs.shape) == 1: + preds = (pred_probs > 0.5).astype(int) + else: + assert len(pred_probs.shape) == 2 + assert isinstance(pred_probs, da.Array) + if self.n_classes_ != 2: + # when using da.argmax directly, dask will construct a numpy based return + # array, which runs into error when computing GPU based prediction. + + def _argmax(x: Any) -> Any: + return x.argmax(axis=1) + + preds = da.map_blocks(_argmax, pred_probs, drop_axis=1) + else: + # Multi-label classification has one binary output per target. + preds = (pred_probs > 0.5).astype(int) + return preds + + +@xgboost_model_doc( + """Implementation of the Scikit-Learn API for XGBoost Ranking. + + .. versionadded:: 1.4.0 + +""", + ["estimators", "model"], + extra_parameters=""" + allow_group_split : + + .. versionadded:: 3.0.0 + + Whether a query group can be split among multiple workers. When set to `False`, + inputs must be Dask dataframes or series. If you have many small query groups, + this can significantly increase the fragmentation of the data, and the internal + DMatrix construction can take longer. + +""", + end_note=""" + .. note:: + + For the dask implementation, group is not supported, use qid instead. +""", +) +class DaskXGBRanker(XGBRankerMixIn, DaskScikitLearnBase): + @_deprecate_positional_args + def __init__( + self, + *, + objective: str = "rank:ndcg", + allow_group_split: bool = False, + coll_cfg: Optional[CollConfig] = None, + **kwargs: Any, + ) -> None: + if callable(objective): + raise ValueError("Custom objective function not supported by XGBRanker.") + self.allow_group_split = allow_group_split + super().__init__(objective=objective, coll_cfg=coll_cfg, **kwargs) + + def _wrapper_params(self) -> Set[str]: + params = super()._wrapper_params() + params.add("allow_group_split") + return params + + async def _fit_async( + self, + X: _DataT, + y: _DaskCollection, + *, + qid: Optional[_DaskCollection], + sample_weight: Optional[_DaskCollection], + base_margin: Optional[_DaskCollection], + eval_set: Optional[Sequence[Tuple[_DaskCollection, _DaskCollection]]], + sample_weight_eval_set: Optional[Sequence[_DaskCollection]], + base_margin_eval_set: Optional[Sequence[_DaskCollection]], + eval_qid: Optional[Sequence[_DaskCollection]], + verbose: Union[int, bool], + xgb_model: Optional[Union[XGBModel, Booster]], + feature_weights: Optional[_DaskCollection], + ) -> Self: + params = self.get_xgb_params() + model, metric, params, feature_weights = self._configure_fit( + xgb_model, params, feature_weights + ) + dtrain, evals = await _async_wrap_evaluation_matrices( + self.client, + device=self.device, + tree_method=self.tree_method, + max_bin=self.max_bin, + X=X, + y=y, + group=None, + qid=qid, + sample_weight=sample_weight, + base_margin=base_margin, + feature_weights=feature_weights, + eval_set=eval_set, + sample_weight_eval_set=sample_weight_eval_set, + base_margin_eval_set=base_margin_eval_set, + eval_group=None, + eval_qid=eval_qid, + missing=self.missing, + enable_categorical=self.enable_categorical, + feature_types=self.feature_types, + ) + results = await self.client.sync( + _train_async, + asynchronous=True, + client=self.client, + global_config=config.get_config(), + dconfig=_get_dask_config(), + params=params, + dtrain=dtrain, + num_boost_round=self.get_num_boosting_rounds(), + evals=evals, + obj=None, + custom_metric=metric, + verbose_eval=verbose, + early_stopping_rounds=self.early_stopping_rounds, + callbacks=self.callbacks, + xgb_model=model, + coll_cfg=self.coll_cfg, + ) + self._Booster = results["booster"] + self.evals_result_ = results["history"] + return self + + # pylint: disable=unused-argument, arguments-differ + @_deprecate_positional_args + def fit( + self, + X: _DataT, + y: _DaskCollection, + *, + group: Optional[_DaskCollection] = None, + qid: Optional[_DaskCollection] = None, + sample_weight: Optional[_DaskCollection] = None, + base_margin: Optional[_DaskCollection] = None, + eval_set: Optional[Sequence[Tuple[_DaskCollection, _DaskCollection]]] = None, + eval_group: Optional[Sequence[_DaskCollection]] = None, + eval_qid: Optional[Sequence[_DaskCollection]] = None, + verbose: Optional[Union[int, bool]] = False, + xgb_model: Optional[Union[XGBModel, str, Booster]] = None, + sample_weight_eval_set: Optional[Sequence[_DaskCollection]] = None, + base_margin_eval_set: Optional[Sequence[_DaskCollection]] = None, + feature_weights: Optional[_DaskCollection] = None, + ) -> "DaskXGBRanker": + msg = "Use the `qid` instead of the `group` with the dask interface." + if not (group is None and eval_group is None): + raise ValueError(msg) + if qid is None: + raise ValueError("`qid` is required for ranking.") + + def check_df(X: _DaskCollection) -> TypeGuard[dd.DataFrame]: + if not isinstance(X, dd.DataFrame): + raise TypeError( + "When `allow_group_split` is set to False, X is required to be" + " a dataframe." + ) + return True + + def check_ser( + qid: Optional[_DaskCollection], name: str + ) -> TypeGuard[Optional[dd.Series]]: + if not isinstance(qid, dd.Series) and qid is not None: + raise TypeError( + f"When `allow_group_split` is set to False, {name} is required to " + "be a series." + ) + return True + + if not self.allow_group_split: + assert ( + check_df(X) + and check_ser(qid, "qid") + and check_ser(y, "y") + and check_ser(sample_weight, "sample_weight") + and check_ser(base_margin, "base_margin") + ) + assert qid is not None and y is not None + X_id = id(X) + X, qid, y, sample_weight, base_margin = no_group_split( + self.device, + X, + qid, + y=y, + sample_weight=sample_weight, + base_margin=base_margin, + ) + + if eval_set is not None: + new_eval_set = [] + new_eval_qid = [] + new_sample_weight_eval_set = [] + new_base_margin_eval_set = [] + assert eval_qid + for i, (Xe, ye) in enumerate(eval_set): + we = sample_weight_eval_set[i] if sample_weight_eval_set else None + be = base_margin_eval_set[i] if base_margin_eval_set else None + assert check_df(Xe) + assert eval_qid + qe = eval_qid[i] + assert ( + eval_qid + and check_ser(qe, "qid") + and check_ser(ye, "y") + and check_ser(we, "sample_weight") + and check_ser(be, "base_margin") + ) + assert qe is not None and ye is not None + if id(Xe) != X_id: + Xe, qe, ye, we, be = no_group_split( + self.device, Xe, qe, ye, we, be + ) + else: + Xe, qe, ye, we, be = X, qid, y, sample_weight, base_margin + + new_eval_set.append((Xe, ye)) + new_eval_qid.append(qe) + + if we is not None: + new_sample_weight_eval_set.append(we) + if be is not None: + new_base_margin_eval_set.append(be) + + eval_set = new_eval_set + eval_qid = new_eval_qid + sample_weight_eval_set = ( + new_sample_weight_eval_set if new_sample_weight_eval_set else None + ) + base_margin_eval_set = ( + new_base_margin_eval_set if new_base_margin_eval_set else None + ) + + return self._client_sync( + self._fit_async, + X=X, + y=y, + qid=qid, + sample_weight=sample_weight, + base_margin=base_margin, + eval_set=eval_set, + eval_qid=eval_qid, + verbose=verbose, + xgb_model=xgb_model, + sample_weight_eval_set=sample_weight_eval_set, + base_margin_eval_set=base_margin_eval_set, + feature_weights=feature_weights, + ) + + # FIXME(trivialfis): arguments differ due to additional parameters like group and + # qid. + if XGBRanker.fit.__doc__ is not None: + fit.__doc__ = XGBRanker.fit.__doc__ + + +@xgboost_model_doc( + """Implementation of the Scikit-Learn API for XGBoost Random Forest Regressor. + + .. versionadded:: 1.4.0 + + .. deprecated:: 3.4.0 + + This estimator is deprecated and will be removed in a future release. See + :doc:`/tutorials/rf` for alternatives. + +""", + ["model", "objective"], + extra_parameters=""" + n_estimators : int + Number of trees in random forest to fit. +""", +) +class DaskXGBRFRegressor(DaskXGBRegressor): + @_deprecate_positional_args + def __init__( + self, + *, + learning_rate: Optional[float] = 1, + subsample: Optional[float] = 0.8, + colsample_bynode: Optional[float] = 0.8, + reg_lambda: Optional[float] = 1e-5, + coll_cfg: Optional[CollConfig] = None, + **kwargs: Any, + ) -> None: + _warn_rf_deprecated("DaskXGBRFRegressor") + super().__init__( + learning_rate=learning_rate, + subsample=subsample, + colsample_bynode=colsample_bynode, + reg_lambda=reg_lambda, + coll_cfg=coll_cfg, + **kwargs, + ) + + def get_xgb_params(self) -> Dict[str, Any]: + params = super().get_xgb_params() + params["num_parallel_tree"] = self.n_estimators + return params + + def get_num_boosting_rounds(self) -> int: + return 1 + + # pylint: disable=unused-argument + def fit( + self, + X: _DataT, + y: _DaskCollection, + *, + sample_weight: Optional[_DaskCollection] = None, + base_margin: Optional[_DaskCollection] = None, + eval_set: Optional[Sequence[Tuple[_DaskCollection, _DaskCollection]]] = None, + verbose: Optional[Union[int, bool]] = True, + xgb_model: Optional[Union[Booster, str, XGBModel]] = None, + sample_weight_eval_set: Optional[Sequence[_DaskCollection]] = None, + base_margin_eval_set: Optional[Sequence[_DaskCollection]] = None, + feature_weights: Optional[_DaskCollection] = None, + ) -> "DaskXGBRFRegressor": + args = {k: v for k, v in locals().items() if k not in ("self", "__class__")} + _check_rf_callback(self.early_stopping_rounds, self.callbacks) + super().fit(**args) + return self + + +@xgboost_model_doc( + """Implementation of the Scikit-Learn API for XGBoost Random Forest Classifier. + + .. versionadded:: 1.4.0 + + .. deprecated:: 3.4.0 + + This estimator is deprecated and will be removed in a future release. See + :doc:`/tutorials/rf` for alternatives. + +""", + ["model", "objective"], + extra_parameters=""" + n_estimators : int + Number of trees in random forest to fit. +""", +) +class DaskXGBRFClassifier(DaskXGBClassifier): + @_deprecate_positional_args + def __init__( + self, + *, + learning_rate: Optional[float] = 1, + subsample: Optional[float] = 0.8, + colsample_bynode: Optional[float] = 0.8, + reg_lambda: Optional[float] = 1e-5, + coll_cfg: Optional[CollConfig] = None, + **kwargs: Any, + ) -> None: + _warn_rf_deprecated("DaskXGBRFClassifier") + super().__init__( + learning_rate=learning_rate, + subsample=subsample, + colsample_bynode=colsample_bynode, + reg_lambda=reg_lambda, + coll_cfg=coll_cfg, + **kwargs, + ) + + def get_xgb_params(self) -> Dict[str, Any]: + params = super().get_xgb_params() + params["num_parallel_tree"] = self.n_estimators + return params + + def get_num_boosting_rounds(self) -> int: + return 1 + + # pylint: disable=unused-argument + def fit( + self, + X: _DataT, + y: _DaskCollection, + *, + sample_weight: Optional[_DaskCollection] = None, + base_margin: Optional[_DaskCollection] = None, + eval_set: Optional[Sequence[Tuple[_DaskCollection, _DaskCollection]]] = None, + verbose: Optional[Union[int, bool]] = True, + xgb_model: Optional[Union[Booster, str, XGBModel]] = None, + sample_weight_eval_set: Optional[Sequence[_DaskCollection]] = None, + base_margin_eval_set: Optional[Sequence[_DaskCollection]] = None, + feature_weights: Optional[_DaskCollection] = None, + ) -> "DaskXGBRFClassifier": + args = {k: v for k, v in locals().items() if k not in ("self", "__class__")} + _check_rf_callback(self.early_stopping_rounds, self.callbacks) + super().fit(**args) + return self diff --git a/v12x/pylib/xgboost/dask/data.py b/v12x/pylib/xgboost/dask/data.py new file mode 100644 index 0000000000000000000000000000000000000000..a216b2ebe14f311c1088f731726c096d976c57f1 --- /dev/null +++ b/v12x/pylib/xgboost/dask/data.py @@ -0,0 +1,461 @@ +# pylint: disable=too-many-arguments +"""Copyright 2019-2025, XGBoost contributors""" + +import logging +from collections.abc import Sequence +from typing import ( + Any, + Callable, + Dict, + List, + Optional, + Tuple, + TypeVar, + Union, + cast, + overload, +) + +import dask +import distributed +import numpy as np +import pandas as pd +from dask import dataframe as dd + +from .. import collective as coll +from .._data_utils import Categories +from .._typing import FeatureNames, FeatureTypes +from ..compat import concat, import_cupy +from ..core import Booster, DataIter, DMatrix, QuantileDMatrix +from ..data import is_on_cuda +from ..sklearn import get_model_categories, pick_ref_categories +from ..training import _RefError + +LOGGER = logging.getLogger("[xgboost.dask]") + +_DataParts = List[Dict[str, Any]] + + +meta = [ + "label", + "weight", + "base_margin", + "qid", + "label_lower_bound", + "label_upper_bound", +] + + +class DaskPartitionIter(DataIter): # pylint: disable=R0902 + """A data iterator for the `DaskQuantileDMatrix`.""" + + def __init__( + self, + data: List[Any], + feature_names: Optional[FeatureNames] = None, + feature_types: Optional[Union[FeatureTypes, Categories]] = None, + feature_weights: Optional[Any] = None, + **kwargs: Optional[List[Any]], + ) -> None: + types = (Sequence, type(None)) + # Samples + self._data = data + for k in meta: + setattr(self, k, kwargs.get(k, None)) + assert isinstance(getattr(self, k), types) + + # Feature info + self._feature_names = feature_names + self._feature_types = feature_types + self._feature_weights = feature_weights + + assert isinstance(self._data, Sequence) + + self._iter = 0 # set iterator to 0 + super().__init__(release_data=True) + + def _get(self, attr: str) -> Optional[Any]: + if getattr(self, attr) is not None: + return getattr(self, attr)[self._iter] + return None + + def data(self) -> Any: + """Utility function for obtaining current batch of data.""" + return self._data[self._iter] + + def reset(self) -> None: + """Reset the iterator""" + self._iter = 0 + + def next(self, input_data: Callable) -> bool: + """Yield next batch of data""" + if self._iter == len(self._data): + # Return False when there's no more batch. + return False + + kwargs = {k: self._get(k) for k in meta} + input_data( + data=self.data(), + group=None, + feature_names=self._feature_names, + feature_types=self._feature_types, + feature_weights=self._feature_weights, + **kwargs, + ) + self._iter += 1 + return True + + +@overload +def _add_column(df: dd.DataFrame, col: dd.Series) -> Tuple[dd.DataFrame, str]: ... + + +@overload +def _add_column(df: dd.DataFrame, col: None) -> Tuple[dd.DataFrame, None]: ... + + +def _add_column( + df: dd.DataFrame, col: Optional[dd.Series] +) -> Tuple[dd.DataFrame, Optional[str]]: + if col is None: + return df, col + + trails = 0 + uid = f"{col.name}_{trails}" + while uid in df.columns: + trails += 1 + uid = f"{col.name}_{trails}" + + df = df.assign(**{uid: col}) + return df, uid + + +def no_group_split( # pylint: disable=too-many-positional-arguments + device: str | None, + df: dd.DataFrame, + qid: dd.Series, + y: dd.Series, + sample_weight: Optional[dd.Series], + base_margin: Optional[dd.Series], +) -> Tuple[ + dd.DataFrame, dd.Series, dd.Series, Optional[dd.Series], Optional[dd.Series] +]: + """A function to prevent query group from being scattered to different + workers. Please see the tutorial in the document for the implication for not having + partition boundary based on query groups. + + """ + + df, qid_uid = _add_column(df, qid) + df, y_uid = _add_column(df, y) + df, w_uid = _add_column(df, sample_weight) + df, bm_uid = _add_column(df, base_margin) + + # `tasks` shuffle is required as of rapids 24.12 + shuffle = "p2p" if device is None or device == "cpu" else "tasks" + with dask.config.set({"dataframe.shuffle.method": shuffle}): + df = df.persist() + # Encode the QID to make it dense. + df[qid_uid] = df[qid_uid].astype("category").cat.as_known().cat.codes + # The shuffle here is costly. + df = df.sort_values(by=qid_uid) + cnt = df.groupby(qid_uid)[qid_uid].count() + div = cnt.index.compute().values.tolist() + div = sorted(div) + div = tuple(div + [div[-1] + 1]) + + df = df.set_index( + qid_uid, + drop=False, + divisions=div, + ).persist() + + qid = df[qid_uid] + y = df[y_uid] + sample_weight, base_margin = ( + cast(dd.Series, df[uid]) if uid is not None else None for uid in (w_uid, bm_uid) + ) + + uids = [uid for uid in [qid_uid, y_uid, w_uid, bm_uid] if uid is not None] + df = df.drop(uids, axis=1).persist() + return df, qid, y, sample_weight, base_margin + + +def sort_data_by_qid(**kwargs: List[Any]) -> Dict[str, List[Any]]: + """Sort worker-local data by query ID for learning to rank tasks.""" + data_parts = kwargs.get("data") + assert data_parts is not None + n_parts = len(data_parts) + + if is_on_cuda(data_parts[0]): + from cudf import DataFrame + else: + from pandas import DataFrame + + def get_dict(i: int) -> Dict[str, list]: + """Return a dictionary containing all the meta info and all partitions.""" + + def _get(attr: Optional[List[Any]]) -> Optional[list]: + if attr is not None: + return attr[i] + return None + + data_opt = {name: _get(kwargs.get(name, None)) for name in meta} + # Filter out None values. + data = {k: v for k, v in data_opt.items() if v is not None} + return data + + def map_fn(i: int) -> pd.DataFrame: + data = get_dict(i) + return DataFrame(data) + + meta_parts = [map_fn(i) for i in range(n_parts)] + dfq = concat(meta_parts) + if dfq.qid.is_monotonic_increasing: + return kwargs + + LOGGER.warning( + "[r%d]: Sorting data with %d partitions for ranking. " + "This is a costly operation and will increase the memory usage significantly. " + "To avoid this warning, sort the data based on qid before passing it into " + "XGBoost. Alternatively, you can use set the `allow_group_split` to False.", + coll.get_rank(), + n_parts, + ) + # I tried to construct a new dask DF to perform the sort, but it's quite difficult + # to get the partition alignment right. Along with the still maturing shuffle + # implementation and GPU compatibility, a simple concat is used. + # + # In case it might become useful one day, I managed to get a CPU version working, + # albeit qutie slow (much slower than concatenated sort). The implementation merges + # everything into a single Dask DF and runs `DF.sort_values`, then retrieve the + # individual X,y,qid, ... from calculated partition values `client.compute([p for p + # in df.partitions])`. It was to avoid creating mismatched partitions. + dfx = concat(data_parts) + + if is_on_cuda(dfq): + cp = import_cupy() + sorted_idx = cp.argsort(dfq.qid) + else: + sorted_idx = np.argsort(dfq.qid) + dfq = dfq.iloc[sorted_idx, :] + + if hasattr(dfx, "iloc"): + dfx = dfx.iloc[sorted_idx, :] + else: + dfx = dfx[sorted_idx, :] + + kwargs.update({"data": [dfx]}) + for i, c in enumerate(dfq.columns): + assert c in kwargs + kwargs.update({c: [dfq[c]]}) + + return kwargs + + +def _get_worker_parts(list_of_parts: _DataParts) -> Dict[str, List[Any]]: + """Convert list of dictionaries into a dictionary of lists.""" + assert isinstance(list_of_parts, list) + result: Dict[str, List[Any]] = {} + + def append(i: int, name: str) -> None: + if name in list_of_parts[i]: + part = list_of_parts[i][name] + else: + part = None + if part is not None: + if name not in result: + result[name] = [] + result[name].append(part) + + for i, _ in enumerate(list_of_parts): + append(i, "data") + for k in meta: + append(i, k) + + qid = result.get("qid", None) + if qid is not None: + result = sort_data_by_qid(**result) + return result + + +def _extract_data( + parts: _DataParts, + model: Optional[Booster], + feature_types: Optional[FeatureTypes], + xy_cats: Optional[Categories], +) -> Tuple[Dict[str, List[Any]], Optional[Union[FeatureTypes, Categories]]]: + unzipped_dict = _get_worker_parts(parts) + X = unzipped_dict["data"][0] + _, model_cats = get_model_categories(X, model, feature_types) + model_cats = pick_ref_categories(X, model_cats, xy_cats) + return unzipped_dict, model_cats + + +def _get_is_cuda(parts: Optional[_DataParts]) -> bool: + if parts is not None: + is_cuda = is_on_cuda(parts[0].get("data")) + else: + is_cuda = False + + is_cuda = bool(coll.allreduce(np.array([is_cuda], dtype=np.int32), coll.Op.MAX)[0]) + return is_cuda + + +def _make_empty(is_cuda: bool) -> np.ndarray: + if is_cuda: + cp = import_cupy() + empty = cp.empty((0, 0)) + else: + empty = np.empty((0, 0)) + return empty + + +def _warn_empty() -> None: + worker = distributed.get_worker() + LOGGER.warning("Worker %s has an empty DMatrix.", worker.address) + + +def _create_quantile_dmatrix( + *, + feature_names: Optional[FeatureNames], + feature_types: Optional[FeatureTypes], + feature_weights: Optional[Any], + missing: float, + nthread: int, + parts: Optional[_DataParts], + max_bin: int, + enable_categorical: bool, + max_quantile_batches: Optional[int], + ref: Optional[DMatrix] = None, + model: Optional[Booster], + Xy_cats: Optional[Categories], +) -> QuantileDMatrix: + is_cuda = _get_is_cuda(parts) + if parts is None: + _warn_empty() + return QuantileDMatrix( + _make_empty(is_cuda), + feature_names=feature_names, + feature_types=feature_types, + max_bin=max_bin, + ref=ref, + enable_categorical=enable_categorical, + max_quantile_batches=max_quantile_batches, + ) + + unzipped_dict, model_cats = _extract_data(parts, model, feature_types, Xy_cats) + + return QuantileDMatrix( + DaskPartitionIter( + **unzipped_dict, + feature_types=model_cats, + feature_names=feature_names, + feature_weights=feature_weights, + ), + missing=missing, + nthread=nthread, + max_bin=max_bin, + ref=ref, + enable_categorical=enable_categorical, + max_quantile_batches=max_quantile_batches, + ) + + +def _create_dmatrix( # pylint: disable=too-many-locals + *, + feature_names: Optional[FeatureNames], + feature_types: Optional[FeatureTypes], + feature_weights: Optional[Any], + missing: float, + nthread: int, + enable_categorical: bool, + parts: Optional[_DataParts], + model: Optional[Booster], + Xy_cats: Optional[Categories], +) -> DMatrix: + """Get data that local to worker from DaskDMatrix. + + Returns + ------- + A DMatrix object. + + """ + is_cuda = _get_is_cuda(parts) + if parts is None: + _warn_empty() + return DMatrix( + _make_empty(is_cuda), + feature_names=feature_names, + feature_types=feature_types, + enable_categorical=enable_categorical, + ) + + T = TypeVar("T") + + def concat_or_none(data: Sequence[Optional[T]]) -> Optional[T]: + if any(part is None for part in data): + return None + return concat(data) + + unzipped_dict, model_cats = _extract_data(parts, model, feature_types, Xy_cats) + + concated_dict: Dict[str, Any] = {} + for key, value in unzipped_dict.items(): + v = concat_or_none(value) + concated_dict[key] = v + + return DMatrix( + **concated_dict, + missing=missing, + feature_names=feature_names, + feature_types=model_cats, + nthread=nthread, + enable_categorical=enable_categorical, + feature_weights=feature_weights, + ) + + +def _dmatrix_from_list_of_parts(is_quantile: bool, **kwargs: Any) -> DMatrix: + if is_quantile: + return _create_quantile_dmatrix(**kwargs) + return _create_dmatrix(**kwargs) + + +def _get_dmatrices( + train_ref: dict, + train_id: int, + *refs: dict, + evals_id: Sequence[int], + evals_name: Sequence[str], + n_threads: int, + model: Optional[Booster], +) -> Tuple[DMatrix, List[Tuple[DMatrix, str]]]: + # Create the training DMatrix + Xy = _dmatrix_from_list_of_parts( + **train_ref, nthread=n_threads, model=model, Xy_cats=None + ) + + # Create evaluation DMatrices + evals: List[Tuple[DMatrix, str]] = [] + Xy_cats = Xy.get_categories() + + for i, ref in enumerate(refs): + # Same DMatrix as the training + if evals_id[i] == train_id: + evals.append((Xy, evals_name[i])) + continue + # Check whether the training DMatrix has been used as a reference. + if ref.get("ref", None) is not None: + if ref["ref"] != train_id: + raise ValueError(_RefError) + del ref["ref"] # Avoid duplicated parameter in the next fn call. + eval_xy = _dmatrix_from_list_of_parts( + **ref, nthread=n_threads, ref=Xy, Xy_cats=Xy_cats, model=model + ) + else: + eval_xy = _dmatrix_from_list_of_parts( + **ref, nthread=n_threads, Xy_cats=Xy_cats, model=model + ) + evals.append((eval_xy, evals_name[i])) + return Xy, evals diff --git a/v12x/pylib/xgboost/dask/utils.py b/v12x/pylib/xgboost/dask/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..8f16ee5a00030ed82ae82bf41f0247461fc7ec1f --- /dev/null +++ b/v12x/pylib/xgboost/dask/utils.py @@ -0,0 +1,118 @@ +"""Utilities for the XGBoost Dask interface.""" + +import logging +import warnings +from functools import cache as fcache +from typing import Any, Dict, Optional, Tuple + +import dask +import distributed +from packaging.version import Version +from packaging.version import parse as parse_version + +from ..collective import Config + +LOGGER = logging.getLogger("[xgboost.dask]") + + +def get_n_threads(local_param: Dict[str, Any], worker: "distributed.Worker") -> int: + """Get the number of threads from a worker and the user-supplied parameters.""" + # dask worker nthreads + dwnt = worker.state.nthreads + n_threads = None + for p in ["nthread", "n_jobs"]: + if local_param.get(p, None) is not None and local_param.get(p, dwnt) != dwnt: + LOGGER.info("Overriding `nthreads` defined in dask worker.") + n_threads = local_param[p] + break + if n_threads == 0 or n_threads is None: + n_threads = dwnt + return n_threads + + +def get_address_from_user( + dconfig: Optional[Dict[str, Any]], coll_cfg: Config +) -> Tuple[Optional[str], int]: + """Get the tracker address from the optional user configuration. + + Parameters + ---------- + dconfig : + Dask global configuration. + + coll_cfg : + Collective configuration. + + Returns + ------- + The IP address along with the port number. + + """ + + valid_config = ["scheduler_address"] + + host_ip = None + port = 0 + + if dconfig is not None: + for k in dconfig: + if k not in valid_config: + raise ValueError(f"Unknown configuration: {k}") + warnings.warn( + ( + "Use `coll_cfg` instead of the Dask global configuration store" + f" for the XGBoost tracker configuration: {k}." + ), + FutureWarning, + ) + else: + dconfig = {} + + host_ip = dconfig.get("scheduler_address", None) + if host_ip is not None and host_ip.startswith("[") and host_ip.endswith("]"): + # convert dask bracket format to proper IPv6 address. + host_ip = host_ip[1:-1] + if host_ip is not None: + try: + host_ip, port = distributed.comm.get_address_host_port(host_ip) + except ValueError: + pass + + if coll_cfg is None: + coll_cfg = Config() + if coll_cfg.tracker_host_ip is not None: + if host_ip is not None and coll_cfg.tracker_host_ip != host_ip: + raise ValueError( + "Conflicting host IP addresses from the dask configuration and the " + f"collective configuration: {host_ip} v.s. {coll_cfg.tracker_host_ip}." + ) + host_ip = coll_cfg.tracker_host_ip + if coll_cfg.tracker_port is not None: + if ( + port != 0 + and port is not None + and coll_cfg.tracker_port != 0 + and port != coll_cfg.tracker_port + ): + raise ValueError( + "Conflicting ports from the dask configuration and the " + f"collective configuration: {port} v.s. {coll_cfg.tracker_port}." + ) + port = coll_cfg.tracker_port + + return host_ip, port + + +@fcache +def _DASK_VERSION() -> Version: + return parse_version(dask.__version__) + + +@fcache +def _DASK_2024_12_1() -> bool: + return _DASK_VERSION() >= parse_version("2024.12.1") + + +@fcache +def _DASK_2025_3_0() -> bool: + return _DASK_VERSION() >= parse_version("2025.3.0") diff --git a/v12x/pylib/xgboost/data.py b/v12x/pylib/xgboost/data.py new file mode 100644 index 0000000000000000000000000000000000000000..308d459946acb3f5b239ef459dba4abbc65e4485 --- /dev/null +++ b/v12x/pylib/xgboost/data.py @@ -0,0 +1,1706 @@ +# pylint: disable=too-many-arguments, too-many-branches, too-many-lines +# pylint: disable=too-many-return-statements +"""Data dispatching for DMatrix.""" + +import ctypes +import functools +import json +import os +import warnings +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Dict, + List, + Optional, + Sequence, + Tuple, + TypeAlias, + TypeGuard, + Union, +) + +import numpy as np + +from ._c_api import _LIB, _check_call, c_str, make_jcargs +from ._data_utils import ( + AifType, + Categories, + DfCatAccessor, + TransformedDf, + _arrow_array_inf, + _ensure_np_dtype, + _is_df_cat, + array_hasobject, + array_interface, + array_interface_dict, + arrow_cat_inf, + check_cudf_meta, + cuda_array_interface, + cuda_array_interface_dict, + cudf_cat_inf, + get_ref_categories, + is_arrow_dict, + pd_cat_inf, +) +from ._typing import ( + CupyT, + DataType, + FeatureNames, + FeatureTypes, + FloatCompatible, + NumpyDType, + PandasDType, + PathLike, + TransformedData, + c_bst_ulong, +) +from .compat import ( + _is_arrow, + _is_cudf_df, + _is_cudf_pandas, + _is_cudf_ser, + _is_cupy_alike, + _is_modin_df, + _is_modin_series, + _is_pandas_df, + _is_pandas_series, + _is_polars, + _is_polars_lazyframe, + _is_polars_series, + import_pandas, + import_polars, + import_pyarrow, + is_pyarrow_available, + lazy_isinstance, +) + +if TYPE_CHECKING: + import pyarrow as pa + from pandas import DataFrame as PdDataFrame + from pandas import Series as PdSeries + + from .core import DMatrix, _ProxyDMatrix + + +DispatchedDataBackendReturnType: TypeAlias = Tuple[ + ctypes.c_void_p, Optional[FeatureNames], Optional[FeatureTypes] +] + +CAT_T = "c" + +# meta info that can be a matrix instead of vector. +_matrix_meta = {"base_margin", "label"} + + +def _warn_unused_missing(data: DataType, missing: Optional[FloatCompatible]) -> None: + if (missing is not None) and (not np.isnan(missing)): + warnings.warn( + "`missing` is not used for current input data type:" + str(type(data)), + UserWarning, + ) + + +def _check_data_shape(data: DataType) -> None: + if hasattr(data, "shape") and len(data.shape) != 2: + raise ValueError("Please reshape the input data into 2-dimensional matrix.") + + +def is_scipy_csr(data: DataType) -> bool: + """Predicate for scipy CSR input.""" + is_array = False + is_matrix = False + try: + from scipy.sparse import csr_array + + is_array = isinstance(data, csr_array) + except ImportError: + pass + try: + from scipy.sparse import csr_matrix + + is_matrix = isinstance(data, csr_matrix) + except ImportError: + pass + return is_array or is_matrix + + +def transform_scipy_sparse(data: DataType, is_csr: bool) -> DataType: + """Ensure correct data alignment and data type for scipy sparse inputs. Input should + be either csr or csc matrix. + + """ + from scipy.sparse import csc_matrix, csr_matrix + + if len(data.indices) != len(data.data): + raise ValueError(f"length mismatch: {len(data.indices)} vs {len(data.data)}") + + indptr, _ = _ensure_np_dtype(data.indptr, data.indptr.dtype) + indices, _ = _ensure_np_dtype(data.indices, data.indices.dtype) + values, _ = _ensure_np_dtype(data.data, data.data.dtype) + if ( + indptr is not data.indptr + or indices is not data.indices + or values is not data.data + ): + if is_csr: + data = csr_matrix((values, indices, indptr), shape=data.shape) + else: + data = csc_matrix((values, indices, indptr), shape=data.shape) + return data + + +def _from_scipy_csr( + *, + data: DataType, + missing: FloatCompatible, + nthread: int, + feature_names: Optional[FeatureNames], + feature_types: Optional[FeatureTypes], +) -> DispatchedDataBackendReturnType: + """Initialize data from a CSR matrix.""" + + handle = ctypes.c_void_p() + data = transform_scipy_sparse(data, True) + _check_call( + _LIB.XGDMatrixCreateFromCSR( + array_interface(data.indptr), + array_interface(data.indices), + array_interface(data.data), + c_bst_ulong(data.shape[1]), + make_jcargs( + missing=float(missing), + nthread=int(nthread), + ), + ctypes.byref(handle), + ) + ) + return handle, feature_names, feature_types + + +def is_scipy_csc(data: DataType) -> bool: + """Predicate for scipy CSC input.""" + is_array = False + is_matrix = False + try: + from scipy.sparse import csc_array + + is_array = isinstance(data, csc_array) + except ImportError: + pass + try: + from scipy.sparse import csc_matrix + + is_matrix = isinstance(data, csc_matrix) + except ImportError: + pass + return is_array or is_matrix + + +def _from_scipy_csc( + *, + data: DataType, + missing: FloatCompatible, + nthread: int, + feature_names: Optional[FeatureNames], + feature_types: Optional[FeatureTypes], +) -> DispatchedDataBackendReturnType: + """Initialize data from a CSC matrix.""" + handle = ctypes.c_void_p() + data = transform_scipy_sparse(data, False) + _check_call( + _LIB.XGDMatrixCreateFromCSC( + array_interface(data.indptr), + array_interface(data.indices), + array_interface(data.data), + c_bst_ulong(data.shape[0]), + make_jcargs( + missing=float(missing), + nthread=int(nthread), + ), + ctypes.byref(handle), + ) + ) + return handle, feature_names, feature_types + + +def is_scipy_coo(data: DataType) -> bool: + """Predicate for scipy COO input.""" + is_array = False + is_matrix = False + try: + from scipy.sparse import coo_array + + is_array = isinstance(data, coo_array) + except ImportError: + pass + try: + from scipy.sparse import coo_matrix + + is_matrix = isinstance(data, coo_matrix) + except ImportError: + pass + return is_array or is_matrix + + +def _is_np_array_like(data: DataType) -> TypeGuard[np.ndarray]: + return hasattr(data, "__array_interface__") + + +def _maybe_np_slice(data: DataType, dtype: Optional[NumpyDType]) -> np.ndarray: + """Handle numpy slice. This can be removed if we use __array_interface__.""" + try: + if not data.flags.c_contiguous: + data = np.array(data, copy=True, dtype=dtype) + else: + data = np.asarray(data, dtype=dtype) + except AttributeError: + data = np.asarray(data, dtype=dtype) + data, dtype = _ensure_np_dtype(data, dtype) + return data + + +def _from_numpy_array( + *, + data: np.ndarray, + missing: FloatCompatible, + nthread: int, + feature_names: Optional[FeatureNames], + feature_types: Optional[FeatureTypes], +) -> DispatchedDataBackendReturnType: + """Initialize data from a 2-D numpy matrix.""" + _check_data_shape(data) + data, _ = _ensure_np_dtype(data, data.dtype) + handle = ctypes.c_void_p() + _check_call( + _LIB.XGDMatrixCreateFromDense( + array_interface(data), + make_jcargs( + missing=float(missing), + nthread=int(nthread), + ), + ctypes.byref(handle), + ) + ) + return handle, feature_names, feature_types + + +_pandas_dtype_mapper = { + "int8": "int", + "int16": "int", + "int32": "int", + "int64": "int", + "uint8": "int", + "uint16": "int", + "uint32": "int", + "uint64": "int", + "float16": "float", + "float32": "float", + "float64": "float", + "bool": "i", +} + +# nullable types +pandas_nullable_mapper = { + "Int8": "int", + "Int16": "int", + "Int32": "int", + "Int64": "int", + "UInt8": "int", + "UInt16": "int", + "UInt32": "int", + "UInt64": "int", + "Float32": "float", + "Float64": "float", + "boolean": "i", +} + +pandas_pyarrow_mapper = { + "int8[pyarrow]": "int", + "int16[pyarrow]": "int", + "int32[pyarrow]": "int", + "int64[pyarrow]": "int", + "uint8[pyarrow]": "int", + "uint16[pyarrow]": "int", + "uint32[pyarrow]": "int", + "uint64[pyarrow]": "int", + "float[pyarrow]": "float", + "float32[pyarrow]": "float", + "double[pyarrow]": "float", + "float64[pyarrow]": "float", + "bool[pyarrow]": "i", +} + +_pandas_dtype_mapper.update(pandas_nullable_mapper) +_pandas_dtype_mapper.update(pandas_pyarrow_mapper) + + +_ENABLE_CAT_ERR = ( + "When categorical type is supplied, the experimental DMatrix parameter" + "`enable_categorical` must be set to `True`." +) + + +def _invalid_dataframe_dtype(data: DataType) -> None: + # pandas series has `dtypes` but it's just a single object + # cudf series doesn't have `dtypes`. + if hasattr(data, "dtypes") and hasattr(data.dtypes, "__iter__"): + bad_fields = [ + f"{data.columns[i]}: {dtype}" + for i, dtype in enumerate(data.dtypes) + if dtype.name not in _pandas_dtype_mapper + ] + err = " Invalid columns:" + ", ".join(bad_fields) + else: + err = "" + + type_err = "DataFrame.dtypes for data must be int, float, bool or category." + msg = f"""{type_err} {_ENABLE_CAT_ERR} {err}""" + raise ValueError(msg) + + +def pandas_feature_info( + data: "PdDataFrame", + meta: Optional[str], + feature_names: Optional[FeatureNames], + feature_types: Optional[FeatureTypes], + enable_categorical: bool, +) -> Tuple[Optional[FeatureNames], Optional[FeatureTypes]]: + """Handle feature info for pandas dataframe.""" + pd = import_pandas() + + # handle feature names + if feature_names is None and meta is None: + if isinstance(data.columns, pd.MultiIndex): + feature_names = [" ".join([str(x) for x in i]) for i in data.columns] + else: + feature_names = list(data.columns.map(str)) + + # handle feature types and dtype validation + new_feature_types = [] + need_sparse_extension_warn = True + for dtype in data.dtypes: + if is_pd_sparse_dtype(dtype): + new_feature_types.append(_pandas_dtype_mapper[dtype.subtype.name]) + if need_sparse_extension_warn: + warnings.warn("Sparse arrays from pandas are converted into dense.") + need_sparse_extension_warn = False + elif ( + is_pd_cat_dtype(dtype) or is_pa_ext_categorical_dtype(dtype) + ) and enable_categorical: + new_feature_types.append(CAT_T) + else: + try: + new_feature_types.append(_pandas_dtype_mapper[dtype.name]) + except KeyError: + _invalid_dataframe_dtype(data) + + if feature_types is None and meta is None: + feature_types = new_feature_types + + return feature_names, feature_types + + +def is_nullable_dtype(dtype: PandasDType) -> bool: + """Whether dtype is a pandas nullable type.""" + + from pandas.api.extensions import ExtensionDtype + + if not isinstance(dtype, ExtensionDtype): + return False + + from pandas.api.types import is_bool_dtype, is_float_dtype, is_integer_dtype + + is_int = is_integer_dtype(dtype) and dtype.name in pandas_nullable_mapper + # np.bool has alias `bool`, while pd.BooleanDtype has `boolean`. + is_bool = is_bool_dtype(dtype) and dtype.name == "boolean" + is_float = is_float_dtype(dtype) and dtype.name in pandas_nullable_mapper + return is_int or is_bool or is_float or is_pd_cat_dtype(dtype) + + +def is_pa_ext_dtype(dtype: Any) -> bool: + """Return whether dtype is a pyarrow extension type for pandas""" + return hasattr(dtype, "pyarrow_dtype") + + +def is_pa_ext_categorical_dtype(dtype: Any) -> bool: + """Check whether dtype is a dictionary type.""" + return lazy_isinstance( + getattr(dtype, "pyarrow_dtype", None), "pyarrow.lib", "DictionaryType" + ) + + +@functools.cache +def _lazy_load_pd_is_cat() -> Callable[[PandasDType], bool]: + pd = import_pandas() + + if hasattr(pd.util, "version") and hasattr(pd.util.version, "Version"): + Version = pd.util.version.Version + if Version(pd.__version__) >= Version("2.1.0"): + from pandas import CategoricalDtype + + def pd_is_cat_210(dtype: PandasDType) -> bool: + return isinstance(dtype, CategoricalDtype) + + return pd_is_cat_210 + from pandas.api.types import is_categorical_dtype # type: ignore[attr-defined] + + return is_categorical_dtype + + +def is_pd_cat_dtype(dtype: PandasDType) -> bool: + """Wrapper for testing pandas category type.""" + is_cat = _lazy_load_pd_is_cat() + return is_cat(dtype) + + +@functools.cache +def _lazy_load_pd_is_sparse() -> Callable[[PandasDType], bool]: + pd = import_pandas() + + if hasattr(pd.util, "version") and hasattr(pd.util.version, "Version"): + Version = pd.util.version.Version + if Version(pd.__version__) >= Version("2.1.0"): + from pandas import SparseDtype + + def pd_is_sparse_210(dtype: PandasDType) -> bool: + return isinstance(dtype, SparseDtype) + + return pd_is_sparse_210 + + from pandas.api.types import is_sparse # type: ignore[attr-defined] + + return is_sparse + + +def is_pd_sparse_dtype(dtype: PandasDType) -> bool: + """Wrapper for testing pandas sparse type.""" + is_sparse = _lazy_load_pd_is_sparse() + + return is_sparse(dtype) + + +def pandas_pa_type(ser: Any) -> np.ndarray: + """Handle pandas pyarrow extension.""" + pd = import_pandas() + + if TYPE_CHECKING: + import pyarrow as pa + else: + pa = import_pyarrow() + + # No copy, callstack: + # pandas.core.internals.managers.SingleBlockManager.array_values() + # pandas.core.internals.blocks.EABackedBlock.values + d_array: pd.arrays.ArrowExtensionArray = ser.array # type: ignore[name-defined] + # no copy in __arrow_array__ + # ArrowExtensionArray._data is a chunked array + aa: "pa.ChunkedArray" = d_array.__arrow_array__() + # combine_chunks takes the most significant amount of time + chunk: "pa.Array" = aa.combine_chunks() + # When there's null value, we have to use copy + zero_copy = chunk.null_count == 0 and not pa.types.is_boolean(chunk.type) + # Alternately, we can use chunk.buffers(), which returns a list of buffers and + # we need to concatenate them ourselves. + # FIXME(jiamingy): Is there a better way to access the arrow buffer along with + # its mask? + # Buffers from chunk.buffers() have the address attribute, but don't expose the + # mask. + arr: np.ndarray = chunk.to_numpy(zero_copy_only=zero_copy, writable=False) + arr, _ = _ensure_np_dtype(arr, arr.dtype) + return arr + + +@functools.cache +def _lazy_has_npdtypes() -> bool: + return np.lib.NumpyVersion(np.__version__) > np.lib.NumpyVersion("1.25.0") + + +@functools.cache +def _lazy_load_pd_floats() -> tuple: + from pandas import Float32Dtype, Float64Dtype + + return Float32Dtype, Float64Dtype + + +def pandas_transform_data( + data: "PdDataFrame", +) -> List[Union[np.ndarray, DfCatAccessor]]: + """Handle categorical dtype and extension types from pandas.""" + Float32Dtype, Float64Dtype = _lazy_load_pd_floats() + + result: List[Union[np.ndarray, DfCatAccessor]] = [] + np_dtypes = _lazy_has_npdtypes() + + def cat_codes(ser: "PdSeries") -> DfCatAccessor: + return ser.cat + + def nu_type(ser: "PdSeries") -> np.ndarray: + # Avoid conversion when possible + if isinstance(dtype, Float32Dtype): + res_dtype: NumpyDType = np.float32 + elif isinstance(dtype, Float64Dtype): + res_dtype = np.float64 + else: + res_dtype = np.float32 + return _ensure_np_dtype( + ser.to_numpy(dtype=res_dtype, na_value=np.nan), res_dtype + )[0] + + def oth_type(ser: "PdSeries") -> np.ndarray: + # The dtypes module is added in 1.25. + npdtypes = np_dtypes and isinstance( + ser.dtype, + ( + # pylint: disable=no-member + np.dtypes.Float32DType, # type: ignore[attr-defined] + # pylint: disable=no-member + np.dtypes.Float64DType, # type: ignore[attr-defined] + ), + ) + + if npdtypes or dtype in {np.float32, np.float64}: + array = ser.to_numpy() + else: + # Specifying the dtype can significantly slow down the conversion (about + # 15% slow down for dense inplace-predict) + array = ser.to_numpy(dtype=np.float32, na_value=np.nan) + return _ensure_np_dtype(array, array.dtype)[0] + + for col, dtype in zip(data.columns, data.dtypes): + if is_pa_ext_categorical_dtype(dtype): + raise ValueError( + "pyarrow dictionary type is not supported. Use pandas category instead." + ) + if is_pd_cat_dtype(dtype): + result.append(cat_codes(data[col])) + elif is_pa_ext_dtype(dtype): + result.append(pandas_pa_type(data[col])) + elif is_nullable_dtype(dtype): + result.append(nu_type(data[col])) + elif is_pd_sparse_dtype(dtype): + arr = data[col].values + arr = arr.to_dense() + if _is_np_array_like(arr): + arr, _ = _ensure_np_dtype(arr, arr.dtype) + result.append(arr) + else: + result.append(oth_type(data[col])) + + # FIXME(jiamingy): Investigate the possibility of using dataframe protocol or arrow + # IPC format for pandas so that we can apply the data transformation inside XGBoost + # for better memory efficiency. + return result + + +class PandasTransformed(TransformedDf): + """A storage class for transformed pandas DataFrame.""" + + def __init__( + self, + columns: List[Union[np.ndarray, DfCatAccessor]], + ref_categories: Optional[Categories], + ) -> None: + self.columns = columns + + aitfs: AifType = [] + temporary_buffers = [] + + # Get the array interface representation for each column. + for col in self.columns: + if _is_df_cat(col): + # Categorical column + jnames, jcodes, buf = pd_cat_inf(col.categories, col.codes) + temporary_buffers.append(buf) + aitfs.append((jnames, jcodes)) + else: + assert isinstance(col, np.ndarray) + inf = array_interface_dict(col) + # Numeric column + aitfs.append(inf) + + super().__init__( + ref_categories=ref_categories, + aitfs=aitfs, + temporary_buffers=temporary_buffers, + ) + + @property + def shape(self) -> Tuple[int, int]: + """Return shape of the transformed DataFrame.""" + if is_arrow_dict(self.columns[0]): + # When input is arrow. + n_samples = len(self.columns[0].indices) + elif _is_df_cat(self.columns[0]): + # When input is pandas. + n_samples = self.columns[0].codes.shape[0] + else: + # Anything else, TypeGuard is ignored by mypy 1.15.0 for some reason + n_samples = self.columns[0].shape[0] # type: ignore[union-attr] + return n_samples, len(self.columns) + + +def _transform_pandas_df( + data: "PdDataFrame", + enable_categorical: bool, + feature_names: Optional[FeatureNames] = None, + feature_types: Optional[Union[FeatureTypes, Categories]] = None, + meta: Optional[str] = None, +) -> Tuple[PandasTransformed, Optional[FeatureNames], Optional[FeatureTypes]]: + if meta and len(data.columns) > 1 and meta not in _matrix_meta: + raise ValueError(f"DataFrame for {meta} cannot have multiple columns") + if data.columns.has_duplicates: + duplicates = data.columns[data.columns.duplicated()].unique().tolist() + raise ValueError( + f"Duplicate column names are not supported. Duplicates found: {duplicates}" + ) + + feature_types, ref_categories = get_ref_categories(feature_types) + feature_names, feature_types = pandas_feature_info( + data, meta, feature_names, feature_types, enable_categorical + ) + + arrays = pandas_transform_data(data) + return ( + PandasTransformed(arrays, ref_categories=ref_categories), + feature_names, + feature_types, + ) + + +def _meta_from_pandas_df( + data: DataType, + name: str, + dtype: Optional[NumpyDType], + handle: ctypes.c_void_p, +) -> None: + data, _, _ = _transform_pandas_df(data, False, meta=name) + if len(data.columns) == 1: + array = data.columns[0] + else: + array = np.stack(data.columns).T + + array, dtype = _ensure_np_dtype(array, dtype) + _meta_from_numpy(array, name, dtype, handle) + + +def _from_pandas_df( + *, + data: "PdDataFrame", + enable_categorical: bool, + missing: FloatCompatible, + nthread: int, + feature_names: Optional[FeatureNames], + feature_types: Optional[Union[FeatureTypes, Categories]], +) -> DispatchedDataBackendReturnType: + df, feature_names, feature_types = _transform_pandas_df( + data, enable_categorical, feature_names, feature_types + ) + + handle = ctypes.c_void_p() + _check_call( + _LIB.XGDMatrixCreateFromColumnar( + df.array_interface(), + make_jcargs(nthread=nthread, missing=missing), + ctypes.byref(handle), + ) + ) + return handle, feature_names, feature_types + + +def _meta_from_pandas_series( + data: DataType, name: str, dtype: Optional[NumpyDType], handle: ctypes.c_void_p +) -> None: + """Help transform pandas series for meta data like labels""" + if is_pd_sparse_dtype(data.dtype): + data = data.values.to_dense().astype(np.float32) + elif is_pa_ext_dtype(data.dtype): + data = pandas_pa_type(data) + else: + data = data.to_numpy(np.float32, na_value=np.nan) + + if is_pd_sparse_dtype(getattr(data, "dtype", data)): + data = data.to_dense() # type: ignore[union-attr] + assert len(data.shape) == 1 or data.shape[1] == 0 or data.shape[1] == 1 + _meta_from_numpy(data, name, dtype, handle) + + +def _from_pandas_series( + *, + data: DataType, + missing: FloatCompatible, + nthread: int, + enable_categorical: bool, + feature_names: Optional[FeatureNames], + feature_types: Optional[FeatureTypes], +) -> DispatchedDataBackendReturnType: + if (data.dtype.name not in _pandas_dtype_mapper) and not ( + is_pd_cat_dtype(data.dtype) and enable_categorical + ): + _invalid_dataframe_dtype(data) + if enable_categorical and is_pd_cat_dtype(data.dtype): + data = data.cat.codes + return _from_numpy_array( + data=data.values.reshape(data.shape[0], 1).astype("float"), + missing=missing, + nthread=nthread, + feature_names=feature_names, + feature_types=feature_types, + ) + + +class ArrowTransformed(TransformedDf): + """A storage class for transformed arrow table.""" + + def __init__( + self, + columns: List[Union["pa.NumericArray", "pa.DictionaryArray"]], + ref_categories: Optional[Categories] = None, + ) -> None: + self.columns = columns + + if TYPE_CHECKING: + import pyarrow as pa + else: + pa = import_pyarrow() + + aitfs: AifType = [] + temporary_buffers = [] + + def push_series(col: Union["pa.NumericArray", "pa.DictionaryArray"]) -> None: + if isinstance(col, pa.DictionaryArray): + cats = col.dictionary + codes = col.indices + if not isinstance(cats, (pa.StringArray, pa.LargeStringArray)): + raise TypeError( + "Only string-based categorical index is supported for arrow." + ) + jnames, jcodes, buf = arrow_cat_inf(cats, codes) + temporary_buffers.append(buf) + aitfs.append((jnames, jcodes)) + else: + jdata = _arrow_array_inf(col) + aitfs.append(jdata) + + for col in self.columns: + push_series(col) + + super().__init__( + ref_categories=ref_categories, + aitfs=aitfs, + temporary_buffers=temporary_buffers, + ) + + @property + def shape(self) -> Tuple[int, int]: + """Return shape of the transformed DataFrame.""" + return len(self.columns[0]), len(self.columns) + + +def _transform_arrow_table( + data: "pa.Table", + enable_categorical: bool, + feature_names: Optional[FeatureNames], + feature_types: Optional[Union[FeatureTypes, Categories]], +) -> Tuple[ArrowTransformed, Optional[FeatureNames], Optional[FeatureTypes]]: + if TYPE_CHECKING: + import pyarrow as pa + else: + pa = import_pyarrow() + + t_names, t_types = _arrow_feature_info(data) + feature_types, ref_categories = get_ref_categories(feature_types) + + if feature_names is None: + feature_names = t_names + if feature_types is None: + feature_types = t_types + + columns = [] + for cname in feature_names: + col0 = data.column(cname) + col: Union["pa.NumericArray", "pa.DictionaryArray"] = col0.combine_chunks() + if isinstance(col, pa.BooleanArray): + col = col.cast(pa.int8()) # bit-compressed array, not supported. + if is_arrow_dict(col) and not enable_categorical: + # None because the function doesn't know how to get the type info from arrow + # table. + _invalid_dataframe_dtype(None) + columns.append(col) + + df_t = ArrowTransformed(columns, ref_categories=ref_categories) + return df_t, feature_names, feature_types + + +def _from_arrow_table( # pylint: disable=too-many-positional-arguments + data: DataType, + enable_categorical: bool, + missing: FloatCompatible, + n_threads: int, + feature_names: Optional[FeatureNames], + feature_types: Optional[Union[FeatureTypes, Categories]], +) -> DispatchedDataBackendReturnType: + df_t, feature_names, feature_types = _transform_arrow_table( + data, enable_categorical, feature_names, feature_types + ) + handle = ctypes.c_void_p() + _check_call( + _LIB.XGDMatrixCreateFromColumnar( + df_t.array_interface(), + make_jcargs(nthread=n_threads, missing=missing), + ctypes.byref(handle), + ) + ) + return handle, feature_names, feature_types + + +@functools.cache +def _arrow_dtype() -> Dict[DataType, str]: + import pyarrow as pa + + mapping = { + pa.int8(): "int", + pa.int16(): "int", + pa.int32(): "int", + pa.int64(): "int", + pa.uint8(): "int", + pa.uint16(): "int", + pa.uint32(): "int", + pa.uint64(): "int", + pa.float16(): "float", + pa.float32(): "float", + pa.float64(): "float", + pa.bool_(): "i", + } + + return mapping + + +def _arrow_feature_info(data: DataType) -> Tuple[List[str], List]: + if TYPE_CHECKING: + import pyarrow as pa + else: + pa = import_pyarrow() + + table: "pa.Table" = data + names = table.column_names + + def map_type(name: str) -> str: + col = table.column(name) + if isinstance(col.type, pa.DictionaryType): + return CAT_T # pylint: disable=unreachable + + return _arrow_dtype()[col.type] + + types = list(map(map_type, names)) + return names, types + + +def _meta_from_arrow_table( + data: DataType, + name: str, + dtype: Optional[NumpyDType], + handle: ctypes.c_void_p, +) -> None: + table: "pa.Table" = data + _meta_from_pandas_df(table.to_pandas(), name=name, dtype=dtype, handle=handle) + + +def _check_pyarrow_for_polars() -> None: + if not is_pyarrow_available(): + raise ImportError("`pyarrow` is required for polars.") + + +def _reject_polars_categorical(data: DataType) -> None: + pl = import_polars() + + for name, dtype in zip(data.columns, data.dtypes): + if isinstance(dtype, pl.Categorical): + raise ValueError( + "XGBoost does not support `polars.Categorical` because its " + "encoding can be sparse. Use `polars.Enum` instead. " + f"Invalid column: {name}", + ) + + +def _transform_polars_df( + data: DataType, + enable_categorical: bool, + feature_names: Optional[FeatureNames], + feature_types: Optional[Union[FeatureTypes, Categories]], +) -> Tuple[ArrowTransformed, Optional[FeatureNames], Optional[FeatureTypes]]: + if _is_polars_lazyframe(data): + df = data.collect() + warnings.warn( + "Using the default parameters for the polars `LazyFrame.collect`. Consider" + " passing a realized `DataFrame` or `Series` instead.", + UserWarning, + ) + else: + df = data + + _check_pyarrow_for_polars() + _reject_polars_categorical(df) + table = df.to_arrow() + return _transform_arrow_table( + table, enable_categorical, feature_names, feature_types + ) + + +def _from_polars_df( # pylint: disable=too-many-positional-arguments + data: DataType, + enable_categorical: bool, + missing: FloatCompatible, + n_threads: int, + feature_names: Optional[FeatureNames], + feature_types: Optional[Union[FeatureTypes, Categories]], +) -> DispatchedDataBackendReturnType: + df_t, feature_names, feature_types = _transform_polars_df( + data, enable_categorical, feature_names, feature_types + ) + handle = ctypes.c_void_p() + _check_call( + _LIB.XGDMatrixCreateFromColumnar( + df_t.array_interface(), + make_jcargs(nthread=n_threads, missing=missing), + ctypes.byref(handle), + ) + ) + return handle, feature_names, feature_types + + +@functools.cache +def _lazy_load_cudf_is_cat() -> Callable[[Any], bool]: + try: + from cudf import CategoricalDtype + + def is_categorical_dtype(dtype: Any) -> bool: + return isinstance(dtype, CategoricalDtype) + + except ImportError: + try: + from cudf.api.types import ( # type: ignore[no-redef] + is_categorical_dtype, + ) + except ImportError: + from cudf.utils.dtypes import ( # type: ignore[no-redef] + is_categorical_dtype, + ) + + return is_categorical_dtype + + +@functools.cache +def _lazy_load_cudf_is_bool() -> Callable[[Any], bool]: + from cudf.api.types import is_bool_dtype + + return is_bool_dtype + + +class CudfTransformed(TransformedDf): + """A storage class for transformed cuDF dataframe.""" + + def __init__( + self, + columns: List[Union["PdSeries", DfCatAccessor]], + ref_categories: Optional[Categories], + ) -> None: + self.columns = columns + # Buffers for temporary data that cannot be freed until the data is consumed by + # the DMatrix or the booster. + + aitfs: AifType = [] + temporary_buffers = [] + + def push_series(ser: Any) -> None: + if _is_df_cat(ser): + cats, codes = ser.categories, ser.codes + cats_ainf, codes_ainf, buf = cudf_cat_inf(cats, codes) + temporary_buffers.append(buf) + aitfs.append((cats_ainf, codes_ainf)) + else: + # numeric column + ainf = cuda_array_interface_dict(ser) + aitfs.append(ainf) + + for col in self.columns: + push_series(col) + + super().__init__( + ref_categories=ref_categories, + aitfs=aitfs, + temporary_buffers=temporary_buffers, + ) + + @property + def shape(self) -> Tuple[int, int]: + """Return shape of the transformed DataFrame.""" + if _is_df_cat(self.columns[0]): + n_samples = self.columns[0].codes.shape[0] + else: + n_samples = self.columns[0].shape[0] # type: ignore[union-attr] + return n_samples, len(self.columns) + + +def _transform_cudf_df( + data: DataType, + feature_names: Optional[FeatureNames], + feature_types: Optional[Union[FeatureTypes, Categories]], + enable_categorical: bool, +) -> Tuple[ + CudfTransformed, + Optional[FeatureNames], + Optional[FeatureTypes], +]: + is_bool_dtype = _lazy_load_cudf_is_bool() + + is_categorical_dtype = _lazy_load_cudf_is_cat() + # Work around https://github.com/dmlc/xgboost/issues/10181 + if _is_cudf_ser(data): + if is_bool_dtype(data.dtype): + data = data.astype(np.uint8) + dtypes = [data.dtype] + else: + data = data.astype( + {col: np.uint8 for col in data.select_dtypes(include="bool")} + ) + dtypes = data.dtypes + + if not all( + dtype.name in _pandas_dtype_mapper + or (is_categorical_dtype(dtype) and enable_categorical) + for dtype in dtypes + ): + _invalid_dataframe_dtype(data) + + # handle feature names + if feature_names is None: + if _is_cudf_ser(data): + feature_names = [data.name] + elif lazy_isinstance(data.columns, "cudf.core.multiindex", "MultiIndex"): + feature_names = [" ".join([str(x) for x in i]) for i in data.columns] + else: + feature_names = list(data.columns.map(str)) + + # handle feature types + feature_types, ref_categories = get_ref_categories(feature_types) + if feature_types is None: + feature_types = [] + for dtype in dtypes: + if is_categorical_dtype(dtype) and enable_categorical: + feature_types.append(CAT_T) + else: + feature_types.append(_pandas_dtype_mapper[dtype.name]) + + # handle categorical data + result = [] + if _is_cudf_ser(data): + # unlike pandas, cuDF uses NA for missing data. + if is_categorical_dtype(data.dtype) and enable_categorical: + result.append(data.cat) + elif is_categorical_dtype(data.dtype) and not enable_categorical: + raise ValueError(_ENABLE_CAT_ERR) + else: + result.append(data) + else: + for col, dtype in zip(data.columns, data.dtypes): + series = data[col] + if is_categorical_dtype(dtype) and enable_categorical: + result.append(series.cat) + elif is_categorical_dtype(dtype): + raise ValueError(_ENABLE_CAT_ERR) + else: + result.append(series) + + return ( + CudfTransformed(result, ref_categories=ref_categories), + feature_names, + feature_types, + ) + + +def _from_cudf_df( + *, + data: DataType, + missing: FloatCompatible, + nthread: int, + feature_names: Optional[FeatureNames], + feature_types: Optional[Union[FeatureTypes, Categories]], + enable_categorical: bool, +) -> DispatchedDataBackendReturnType: + df, feature_names, feature_types = _transform_cudf_df( + data, feature_names, feature_types, enable_categorical + ) + handle = ctypes.c_void_p() + _check_call( + _LIB.XGDMatrixCreateFromCudaColumnar( + df.array_interface(), + make_jcargs(nthread=nthread, missing=missing), + ctypes.byref(handle), + ) + ) + return handle, feature_names, feature_types + + +def _transform_cupy_array(data: DataType) -> CupyT: + import cupy + + if not hasattr(data, "__cuda_array_interface__") and hasattr(data, "__array__"): + data = cupy.array(data, copy=False) + if array_hasobject(data) or data.dtype in [cupy.bool_]: + data = data.astype(cupy.float32, copy=False) + return data + + +def _from_cupy_array( + data: DataType, + missing: FloatCompatible, + nthread: int, + feature_names: Optional[FeatureNames], + feature_types: Optional[FeatureTypes], +) -> DispatchedDataBackendReturnType: + """Initialize DMatrix from cupy ndarray.""" + data = _transform_cupy_array(data) + interface_str = cuda_array_interface(data) + handle = ctypes.c_void_p() + config = bytes(json.dumps({"missing": missing, "nthread": nthread}), "utf-8") + _check_call( + _LIB.XGDMatrixCreateFromCudaArrayInterface( + interface_str, config, ctypes.byref(handle) + ) + ) + return handle, feature_names, feature_types + + +def _is_cupy_csr(data: DataType) -> bool: + try: + import cupyx + except ImportError: + return False + return isinstance(data, cupyx.scipy.sparse.csr_matrix) + + +def _is_cupy_csc(data: DataType) -> bool: + try: + import cupyx + except ImportError: + return False + return isinstance(data, cupyx.scipy.sparse.csc_matrix) + + +def _is_dlpack(data: DataType) -> bool: + return "PyCapsule" in str(type(data)) and "dltensor" in str(data) + + +def _transform_dlpack(data: DataType) -> bool: + from cupy import from_dlpack # pylint: disable=E0401 + + assert "used_dltensor" not in str(data) + data = from_dlpack(data) + return data + + +def _from_dlpack( + data: DataType, + missing: FloatCompatible, + nthread: int, + feature_names: Optional[FeatureNames], + feature_types: Optional[FeatureTypes], +) -> DispatchedDataBackendReturnType: + data = _transform_dlpack(data) + return _from_cupy_array(data, missing, nthread, feature_names, feature_types) + + +def _is_uri(data: DataType) -> TypeGuard[PathLike]: + return isinstance(data, (str, os.PathLike)) + + +def _from_uri( + data: PathLike, + missing: Optional[FloatCompatible], + feature_names: Optional[FeatureNames], + feature_types: Optional[FeatureTypes], +) -> DispatchedDataBackendReturnType: + _warn_unused_missing(data, missing) + handle = ctypes.c_void_p() + data = os.fspath(os.path.expanduser(data)) + config = make_jcargs(uri=str(data)) + _check_call(_LIB.XGDMatrixCreateFromURI(config, ctypes.byref(handle))) + return handle, feature_names, feature_types + + +def _is_list(data: DataType) -> TypeGuard[list]: + return isinstance(data, list) + + +def _from_list( + *, + data: Sequence, + missing: FloatCompatible, + n_threads: int, + feature_names: Optional[FeatureNames], + feature_types: Optional[FeatureTypes], +) -> DispatchedDataBackendReturnType: + array = np.array(data) + _check_data_shape(data) + return _from_numpy_array( + data=array, + missing=missing, + nthread=n_threads, + feature_names=feature_names, + feature_types=feature_types, + ) + + +def _is_tuple(data: DataType) -> TypeGuard[tuple]: + return isinstance(data, tuple) + + +def _from_tuple( + *, + data: Sequence, + missing: FloatCompatible, + n_threads: int, + feature_names: Optional[FeatureNames], + feature_types: Optional[FeatureTypes], +) -> DispatchedDataBackendReturnType: + return _from_list( + data=data, + missing=missing, + n_threads=n_threads, + feature_names=feature_names, + feature_types=feature_types, + ) + + +def _has_array_protocol(data: DataType) -> bool: + return hasattr(data, "__array__") + + +def _convert_unknown_data(data: DataType) -> DataType: + warnings.warn( + f"Unknown data type: {type(data)}, trying to convert it to csr_matrix", + UserWarning, + ) + try: + import scipy.sparse + except ImportError: + return None + + try: + data = scipy.sparse.csr_matrix(data) + except Exception: # pylint: disable=broad-except + return None + + return data + + +def dispatch_data_backend( + *, + data: DataType, + missing: FloatCompatible, # Or Optional[Float] + threads: int, + feature_names: Optional[FeatureNames], + feature_types: Optional[Union[FeatureTypes, Categories]], + enable_categorical: bool = False, +) -> DispatchedDataBackendReturnType: + """Dispatch data for DMatrix.""" + + def check_cats( + feature_types: Optional[Union[FeatureTypes, Categories]], + ) -> TypeGuard[Optional[FeatureTypes]]: + if isinstance(feature_types, Categories): + raise ValueError( + "Reference category is only supported by DataFrame inputs." + ) + return True + + if ( + not _is_cudf_ser(data) + and not _is_pandas_series(data) + and not _is_polars_series(data) + ): + _check_data_shape(data) + if is_scipy_csr(data): + assert check_cats(feature_types) + return _from_scipy_csr( + data=data, + missing=missing, + nthread=threads, + feature_names=feature_names, + feature_types=feature_types, + ) + if is_scipy_csc(data): + assert check_cats(feature_types) + return _from_scipy_csc( + data=data, + missing=missing, + nthread=threads, + feature_names=feature_names, + feature_types=feature_types, + ) + if is_scipy_coo(data): + assert check_cats(feature_types) + return _from_scipy_csr( + data=data.tocsr(), + missing=missing, + nthread=threads, + feature_names=feature_names, + feature_types=feature_types, + ) + if _is_np_array_like(data): + assert check_cats(feature_types) + return _from_numpy_array( + data=data, + missing=missing, + nthread=threads, + feature_names=feature_names, + feature_types=feature_types, + ) + if _is_uri(data): + assert check_cats(feature_types) + return _from_uri(data, missing, feature_names, feature_types) + if _is_list(data): + assert check_cats(feature_types) + return _from_list( + data=data, + missing=missing, + n_threads=threads, + feature_names=feature_names, + feature_types=feature_types, + ) + if _is_tuple(data): + assert check_cats(feature_types) + return _from_tuple( + data=data, + missing=missing, + n_threads=threads, + feature_names=feature_names, + feature_types=feature_types, + ) + if _is_polars_series(data): + pl = import_polars() + + data = pl.DataFrame({data.name: data}) + if _is_polars(data): + return _from_polars_df( + data, + enable_categorical, + missing=missing, + n_threads=threads, + feature_names=feature_names, + feature_types=feature_types, + ) + if _is_arrow(data): + return _from_arrow_table( + data, + enable_categorical, + missing=missing, + n_threads=threads, + feature_names=feature_names, + feature_types=feature_types, + ) + if _is_cudf_pandas(data): + data = data._fsproxy_fast # pylint: disable=protected-access + if _is_pandas_series(data): + pd = import_pandas() + + data = pd.DataFrame(data) + if _is_pandas_df(data): + return _from_pandas_df( + data=data, + enable_categorical=enable_categorical, + missing=missing, + nthread=threads, + feature_names=feature_names, + feature_types=feature_types, + ) + if _is_cudf_df(data) or _is_cudf_ser(data): + return _from_cudf_df( + data=data, + missing=missing, + nthread=threads, + feature_names=feature_names, + feature_types=feature_types, + enable_categorical=enable_categorical, + ) + if _is_cupy_alike(data): + assert check_cats(feature_types) + return _from_cupy_array(data, missing, threads, feature_names, feature_types) + if _is_cupy_csr(data): + raise TypeError("cupyx CSR is not supported yet.") + if _is_cupy_csc(data): + raise TypeError("cupyx CSC is not supported yet.") + if _is_dlpack(data): + assert check_cats(feature_types) + return _from_dlpack(data, missing, threads, feature_names, feature_types) + if _is_modin_series(data): + pd = import_pandas() + + data = pd.DataFrame(data) + if _is_modin_df(data): + return _from_pandas_df( + data=data, + enable_categorical=enable_categorical, + missing=missing, + nthread=threads, + feature_names=feature_names, + feature_types=feature_types, + ) + + if _has_array_protocol(data): + assert check_cats(feature_types) + array = np.asarray(data) + return _from_numpy_array( + data=array, + missing=missing, + nthread=threads, + feature_names=feature_names, + feature_types=feature_types, + ) + + converted = _convert_unknown_data(data) + if converted is not None: + assert check_cats(feature_types) + return _from_scipy_csr( + data=converted, + missing=missing, + nthread=threads, + feature_names=feature_names, + feature_types=feature_types, + ) + + raise TypeError("Not supported type for data." + str(type(data))) + + +def _validate_meta_shape(data: DataType, name: str) -> None: + if hasattr(data, "shape"): + msg = f"Invalid shape: {data.shape} for {name}" + if name in _matrix_meta: + if len(data.shape) > 2: + raise ValueError(msg) + return + + if len(data.shape) > 2 or ( + len(data.shape) == 2 and (data.shape[1] != 0 and data.shape[1] != 1) + ): + raise ValueError(f"Invalid shape: {data.shape} for {name}") + + +def _meta_from_numpy( + data: np.ndarray, + field: str, + dtype: Optional[NumpyDType], + handle: ctypes.c_void_p, +) -> None: + data, dtype = _ensure_np_dtype(data, dtype) + interface = data.__array_interface__ + if interface.get("mask", None) is not None: + raise ValueError("Masked array is not supported.") + interface_str = array_interface(data) + _check_call(_LIB.XGDMatrixSetInfoFromInterface(handle, c_str(field), interface_str)) + + +def _meta_from_list( + data: Sequence, field: str, dtype: Optional[NumpyDType], handle: ctypes.c_void_p +) -> None: + data_np = np.array(data) + _meta_from_numpy(data_np, field, dtype, handle) + + +def _meta_from_tuple( + data: Sequence, field: str, dtype: Optional[NumpyDType], handle: ctypes.c_void_p +) -> None: + return _meta_from_list(data, field, dtype, handle) + + +def _meta_from_cudf_df(data: DataType, field: str, handle: ctypes.c_void_p) -> None: + if field not in _matrix_meta: + _meta_from_cudf_series(data.iloc[:, 0], field, handle) + else: + data = data.values + interface = cuda_array_interface(data) + _check_call(_LIB.XGDMatrixSetInfoFromInterface(handle, c_str(field), interface)) + + +def _meta_from_cudf_series(data: DataType, field: str, handle: ctypes.c_void_p) -> None: + check_cudf_meta(data, field) + inf = cuda_array_interface(data) + _check_call(_LIB.XGDMatrixSetInfoFromInterface(handle, c_str(field), inf)) + + +def _meta_from_cupy_array(data: DataType, field: str, handle: ctypes.c_void_p) -> None: + data = _transform_cupy_array(data) + inf = cuda_array_interface(data) + _check_call(_LIB.XGDMatrixSetInfoFromInterface(handle, c_str(field), inf)) + + +def dispatch_meta_backend( + matrix: "DMatrix", data: DataType, name: str, dtype: Optional[NumpyDType] = None +) -> None: + """Dispatch for meta info.""" + handle = matrix.handle + assert handle is not None + _validate_meta_shape(data, name) + if data is None: + return + if _is_list(data): + _meta_from_list(data, name, dtype, handle) + return + if _is_tuple(data): + _meta_from_tuple(data, name, dtype, handle) + return + if _is_np_array_like(data): + _meta_from_numpy(data, name, dtype, handle) + return + if _is_arrow(data): + _meta_from_arrow_table(data, name, dtype, handle) + return + if _is_cudf_pandas(data): + data = data._fsproxy_fast # pylint: disable=protected-access + if _is_polars(data): + if _is_polars_lazyframe(data): + data = data.collect() + _check_pyarrow_for_polars() + _meta_from_arrow_table(data.to_arrow(), name, dtype, handle) + return + if _is_pandas_df(data): + _meta_from_pandas_df(data, name, dtype=dtype, handle=handle) + return + if _is_pandas_series(data): + _meta_from_pandas_series(data, name, dtype, handle) + return + if _is_dlpack(data): + data = _transform_dlpack(data) + _meta_from_cupy_array(data, name, handle) + return + if _is_cudf_ser(data): + _meta_from_cudf_series(data, name, handle) + return + if _is_cudf_df(data): + _meta_from_cudf_df(data, name, handle) + return + if _is_cupy_alike(data): + _meta_from_cupy_array(data, name, handle) + return + if _is_modin_df(data): + _meta_from_pandas_df(data, name, dtype=dtype, handle=handle) + return + if _is_modin_series(data): + data = data.values.astype("float") + assert len(data.shape) == 1 or data.shape[1] == 0 or data.shape[1] == 1 + _meta_from_numpy(data, name, dtype, handle) + return + if _has_array_protocol(data): + # pyarrow goes here. + array = np.asarray(data) + _meta_from_numpy(array, name, dtype, handle) + return + raise TypeError("Unsupported type for " + name, str(type(data))) + + +def _proxy_transform( + data: DataType, + feature_names: Optional[FeatureNames], + feature_types: Optional[FeatureTypes], + enable_categorical: bool, +) -> TransformedData: + if _is_cudf_pandas(data): + data = data._fsproxy_fast # pylint: disable=protected-access + if _is_cudf_df(data) or _is_cudf_ser(data): + return _transform_cudf_df( + data, feature_names, feature_types, enable_categorical + ) + if _is_cupy_alike(data): + data = _transform_cupy_array(data) + return data, feature_names, feature_types + if _is_dlpack(data): + return _transform_dlpack(data), feature_names, feature_types + if _is_list(data) or _is_tuple(data): + data = np.array(data) + if _is_np_array_like(data): + data, _ = _ensure_np_dtype(data, data.dtype) + return data, feature_names, feature_types + if is_scipy_csr(data): + data = transform_scipy_sparse(data, True) + return data, feature_names, feature_types + if is_scipy_csc(data): + data = transform_scipy_sparse(data.tocsr(), True) + return data, feature_names, feature_types + if is_scipy_coo(data): + data = transform_scipy_sparse(data.tocsr(), True) + return data, feature_names, feature_types + if _is_polars(data): + df_pl, feature_names, feature_types = _transform_polars_df( + data, enable_categorical, feature_names, feature_types + ) + return df_pl, feature_names, feature_types + if _is_pandas_series(data): + pd = import_pandas() + + data = pd.DataFrame(data) + if _is_arrow(data): + df_pa, feature_names, feature_types = _transform_arrow_table( + data, enable_categorical, feature_names, feature_types + ) + return df_pa, feature_names, feature_types + if _is_pandas_df(data): + df, feature_names, feature_types = _transform_pandas_df( + data, enable_categorical, feature_names, feature_types + ) + return df, feature_names, feature_types + raise TypeError("Value type is not supported for data iterator:" + str(type(data))) + + +def is_on_cuda(data: Any) -> bool: + """Whether the data is a CUDA-based data structure.""" + return any( + p(data) + for p in ( + _is_cudf_df, + _is_cudf_ser, + _is_cudf_pandas, + _is_cupy_alike, + _is_dlpack, + ) + ) + + +def dispatch_proxy_set_data( + proxy: "_ProxyDMatrix", + data: DataType, +) -> None: + """Dispatch for QuantileDMatrix.""" + if ( + not _is_cudf_ser(data) + and not _is_pandas_series(data) + and not _is_polars_series(data) + ): + _check_data_shape(data) + + if isinstance(data, CudfTransformed): + # pylint: disable=W0212 + proxy._ref_data_from_cuda_columnar(data) + return + if _is_cupy_alike(data): + proxy._ref_data_from_cuda_interface(data) # pylint: disable=W0212 + return + if _is_dlpack(data): + data = _transform_dlpack(data) + proxy._ref_data_from_cuda_interface(data) # pylint: disable=W0212 + return + # Host + if isinstance(data, (ArrowTransformed, PandasTransformed)): + proxy._ref_data_from_columnar(data) # pylint: disable=W0212 + return + if _is_np_array_like(data): + _check_data_shape(data) + proxy._ref_data_from_array(data) # pylint: disable=W0212 + return + if is_scipy_csr(data): + proxy._ref_data_from_csr(data) # pylint: disable=W0212 + return + + err = TypeError("Value type is not supported for data iterator:" + str(type(data))) + raise err diff --git a/v12x/pylib/xgboost/federated.py b/v12x/pylib/xgboost/federated.py new file mode 100644 index 0000000000000000000000000000000000000000..e903d475cfc6edeea8b79af048b9eeca404f48c5 --- /dev/null +++ b/v12x/pylib/xgboost/federated.py @@ -0,0 +1,109 @@ +"""XGBoost Experimental Federated Learning related API.""" + +import ctypes +from threading import Thread +from typing import Any, Dict, Optional + +from .core import _LIB, _check_call, _deprecate_positional_args, make_jcargs +from .tracker import RabitTracker + + +class FederatedTracker(RabitTracker): + """Tracker for federated training. + + Parameters + ---------- + n_workers : + The number of federated workers. + + port : + The port to listen on. + + secure : + Whether this is a secure instance. If True, then the following arguments for SSL + must be provided. + + server_key_path : + Path to the server private key file. + + server_cert_path : + Path to the server certificate file. + + client_cert_path : + Path to the client certificate file. + + """ + + @_deprecate_positional_args + def __init__( # pylint: disable=R0913, W0231 + self, + n_workers: int, + port: int, + *, + secure: bool, + server_key_path: Optional[str] = None, + server_cert_path: Optional[str] = None, + client_cert_path: Optional[str] = None, + timeout: int = 300, + ) -> None: + handle = ctypes.c_void_p() + args = make_jcargs( + n_workers=n_workers, + port=port, + dmlc_communicator="federated", + federated_secure=secure, + server_key_path=server_key_path, + server_cert_path=server_cert_path, + client_cert_path=client_cert_path, + timeout=int(timeout), + ) + _check_call(_LIB.XGTrackerCreate(args, ctypes.byref(handle))) + self.handle = handle + + +@_deprecate_positional_args +def run_federated_server( # pylint: disable=too-many-arguments + n_workers: int, + port: int, + *, + server_key_path: Optional[str] = None, + server_cert_path: Optional[str] = None, + client_cert_path: Optional[str] = None, + blocking: bool = True, + timeout: int = 300, +) -> Optional[Dict[str, Any]]: + """See :py:class:`~xgboost.federated.FederatedTracker` for more info. + + Parameters + ---------- + blocking : + Block the server until the training is finished. If set to False, the function + launches an additional thread and returns the worker arguments. The default is + True and a higher level framework is responsible for setting worker parameters. + + """ + args: Dict[str, Any] = {"n_workers": n_workers} + secure = all( + path is not None + for path in [server_key_path, server_cert_path, client_cert_path] + ) + tracker = FederatedTracker( + n_workers=n_workers, + port=port, + secure=secure, + timeout=timeout, + server_key_path=server_key_path, + server_cert_path=server_cert_path, + client_cert_path=client_cert_path, + ) + tracker.start() + + if blocking: + tracker.wait_for() + return None + + thread = Thread(target=tracker.wait_for) + thread.daemon = True + thread.start() + args.update(tracker.worker_args()) + return args diff --git a/v12x/pylib/xgboost/interpret.py b/v12x/pylib/xgboost/interpret.py new file mode 100644 index 0000000000000000000000000000000000000000..796680fe1c5ec3766a935535b4b25eef0d115d53 --- /dev/null +++ b/v12x/pylib/xgboost/interpret.py @@ -0,0 +1,127 @@ +"""Interpretability functions for XGBoost models.""" + +from typing import Optional, Tuple, Union + +import numpy as np + +from ._typing import ArrayLike, FloatCompatible, IterationRange +from .core import Booster, DMatrix + + +def _as_booster(model: object) -> Booster: + if isinstance(model, Booster): + return model + get_booster = getattr(model, "get_booster", None) + if not callable(get_booster): + raise TypeError( + "`model` must be an xgboost.Booster or an object with get_booster()." + ) + booster = get_booster() + if not isinstance(booster, Booster): + raise TypeError("`model.get_booster()` must return an xgboost.Booster.") + return booster + + +def _get_iteration_range( + model: object, iteration_range: Optional[IterationRange] +) -> IterationRange: + get_iteration_range = getattr(model, "_get_iteration_range", None) + if get_iteration_range is not None: + return get_iteration_range(iteration_range) + if iteration_range is None: + return (0, 0) + return iteration_range + + +def _as_prediction_dmatrix( + model: object, X: Union[DMatrix, ArrayLike], missing: Optional[FloatCompatible] +) -> DMatrix: + if isinstance(X, DMatrix): + if missing is not None: + raise ValueError("`missing` must not be specified when `X` is a DMatrix.") + return X + + return DMatrix( + X, + missing=missing if missing is not None else getattr(model, "missing", None), + nthread=getattr(model, "n_jobs", None), + feature_types=getattr(model, "feature_types", None), + enable_categorical=getattr(model, "enable_categorical", False), + ) + + +def shap_values( # pylint: disable=too-many-arguments + model: object, + X: Union[DMatrix, ArrayLike], + *, + X_background: Optional[Union[DMatrix, ArrayLike]] = None, + output_margin: bool = False, + iteration_range: Optional[IterationRange] = None, + missing: Optional[FloatCompatible] = None, + validate_features: bool = True, +) -> Tuple[np.ndarray, np.ndarray]: + """Return SHAP values for an XGBoost model. + + .. warning:: + + This function is still working in progress. + + This function accepts either a :py:class:`xgboost.Booster` or an sklearn-style + XGBoost model and returns feature contributions together with the separated + bias term. + + Parameters + ---------- + model : + XGBoost booster or sklearn-style XGBoost model. + X : + Input data. + X_background : + Background data for interventional SHAP values. This is reserved for a + future implementation and is currently unsupported. + output_margin : + Accepted for API compatibility. SHAP contributions currently correspond + to the model margin. + iteration_range : + Specifies which layer of trees are used in prediction. + missing : + Value in array-like ``X`` to treat as missing. When None, use the + model's missing value if available, otherwise ``np.nan``. This must not + be specified when ``X`` is already a DMatrix. + validate_features : + Validate feature names between the model and input data. + + Returns + ------- + values, bias : + ``values`` contains feature SHAP values with the bias term removed. + ``bias`` contains the separated bias term. For multi-target models, the + output shape follows the corresponding prediction shape with the final + feature dimension split into ``values`` and ``bias``. + + Notes + ----- + To use GPU algorithms, configure the model before calling this function, for + example with ``booster.set_param({"device": "cuda"})``. + """ + if X_background is not None: + raise NotImplementedError("`X_background` is not yet supported.") + # SHAP contributions currently correspond to the model margin. Keep this + # argument in the initial API so callers can use the proposed signature. + _ = output_margin + + booster = _as_booster(model) + data = _as_prediction_dmatrix(model, X, missing) + contribs = booster.predict( + data, + pred_contribs=True, + validate_features=validate_features, + iteration_range=_get_iteration_range(model, iteration_range), + ) + + values = contribs[..., :-1] + bias = contribs[..., -1] + return values, bias + + +__all__ = ["shap_values"] diff --git a/v12x/pylib/xgboost/libpath.py b/v12x/pylib/xgboost/libpath.py new file mode 100644 index 0000000000000000000000000000000000000000..097c76830ac0a93d0629063155e987df333e6333 --- /dev/null +++ b/v12x/pylib/xgboost/libpath.py @@ -0,0 +1,78 @@ +# coding: utf-8 +"""Find the path to xgboost dynamic library files.""" + +import os +import platform +import sys +from typing import List + + +class XGBoostLibraryNotFound(Exception): + """Error thrown by when xgboost is not found""" + + +def is_sphinx_build() -> bool: + """`XGBOOST_BUILD_DOC` is used by the sphinx conf.py to skip building the C++ code.""" + return bool(os.environ.get("XGBOOST_BUILD_DOC", False)) + + +def find_lib_path() -> List[str]: + """Find the path to xgboost dynamic library files. + + Returns + ------- + lib_path + List of all found library path to xgboost + """ + curr_path = os.path.dirname(os.path.abspath(os.path.expanduser(__file__))) + dll_path = [ + # normal, after installation `lib` is copied into Python package tree. + os.path.join(curr_path, "lib"), + # editable installation, no copying is performed. + os.path.join(curr_path, os.path.pardir, os.path.pardir, "lib"), + # use libxgboost from a system prefix, if available. This should be the last + # option. + os.path.join(sys.base_prefix, "lib"), + ] + + if sys.platform == "win32": + # On Windows, Conda may install libs in different paths + dll_path.extend( + [ + os.path.join(sys.base_prefix, "bin"), + os.path.join(sys.base_prefix, "Library"), + os.path.join(sys.base_prefix, "Library", "bin"), + os.path.join(sys.base_prefix, "Library", "lib"), + os.path.join(sys.base_prefix, "Library", "mingw-w64"), + os.path.join(sys.base_prefix, "Library", "mingw-w64", "bin"), + os.path.join(sys.base_prefix, "Library", "mingw-w64", "lib"), + ] + ) + dll_path = [os.path.join(p, "xgboost.dll") for p in dll_path] + elif sys.platform.startswith(("linux", "freebsd", "emscripten")): + dll_path = [os.path.join(p, "libxgboost.so") for p in dll_path] + elif sys.platform == "darwin": + dll_path = [os.path.join(p, "libxgboost.dylib") for p in dll_path] + elif sys.platform == "cygwin": + dll_path = [os.path.join(p, "cygxgboost.dll") for p in dll_path] + if platform.system() == "OS400": + dll_path = [os.path.join(p, "libxgboost.so") for p in dll_path] + + lib_path = [p for p in dll_path if os.path.exists(p) and os.path.isfile(p)] + + if not lib_path and not is_sphinx_build(): + link = "https://xgboost.readthedocs.io/en/stable/install.html" + msg = ( + "Cannot find XGBoost Library in the candidate path. " + + "List of candidates:\n- " + + ("\n- ".join(dll_path)) + + "\nXGBoost Python package path: " + + curr_path + + "\nsys.base_prefix: " + + sys.base_prefix + + "\nSee: " + + link + + " for installing XGBoost." + ) + raise XGBoostLibraryNotFound(msg) + return lib_path diff --git a/v12x/pylib/xgboost/objective.py b/v12x/pylib/xgboost/objective.py new file mode 100644 index 0000000000000000000000000000000000000000..8da01068009e7dce129f800411f0c6439ac8a58f --- /dev/null +++ b/v12x/pylib/xgboost/objective.py @@ -0,0 +1,92 @@ +"""Experimental support for a new objective interface with target dimension +reduction. + +.. warning:: + + Do not use this module unless you want to participate in development. + +.. versionadded:: 3.2.0 + +""" + +import warnings +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING, Tuple + +import numpy as np + +from ._data_utils import ( + _ensure_np_dtype, + _is_flatten, + array_interface, + cuda_array_interface, +) +from ._typing import ArrayLike, NumpyOrCupy +from .compat import _is_cupy_alike + +if TYPE_CHECKING: + from .core import DMatrix + + +class Objective(ABC): + """Base class for custom objective functions. + + .. warning:: + + Do not use this class unless you want to participate in development. + + .. versionadded:: 3.2.0 + + """ + + @abstractmethod + def __call__( + self, iteration: int, y_pred: ArrayLike, dtrain: "DMatrix" + ) -> Tuple[ArrayLike, ArrayLike]: ... + + +class TreeObjective(Objective): + """Base class for tree-specific custom objective functions. + + .. warning:: + + Do not use this class unless you want to participate in development. + + .. versionadded:: 3.2.0 + + """ + + # pylint: disable=unused-argument + def split_grad( + self, iteration: int, grad: ArrayLike, hess: ArrayLike + ) -> Tuple[ArrayLike, ArrayLike] | None: + """Provide a different gradient type for finding tree structures.""" + return None + + +def _grad_arrinf(array: NumpyOrCupy, n_samples: int) -> bytes: + # Can we check for __array_interface__ instead of a specific type instead? + msg = ( + "Expecting `np.ndarray` or `cupy.ndarray` for gradient and hessian." + f" Got: {type(array)}" + ) + if not isinstance(array, np.ndarray) and not _is_cupy_alike(array): + raise TypeError(msg) + + if array.shape[0] != n_samples and _is_flatten(array): + warnings.warn( + "Since 2.1.0, the shape of the gradient and hessian is required to" + " be (n_samples, n_targets) or (n_samples, n_classes).", + FutureWarning, + ) + array = array.reshape(n_samples, array.size // n_samples) + + if isinstance(array, np.ndarray): + array, _ = _ensure_np_dtype(array, array.dtype) + interface = array_interface(array) + elif _is_cupy_alike(array): + interface = cuda_array_interface(array) + else: + raise TypeError(msg) + + return interface diff --git a/v12x/pylib/xgboost/plotting.py b/v12x/pylib/xgboost/plotting.py new file mode 100644 index 0000000000000000000000000000000000000000..e80fc996d4de0a92378ed7703c0d03efc37ca7cf --- /dev/null +++ b/v12x/pylib/xgboost/plotting.py @@ -0,0 +1,360 @@ +# pylint: disable=too-many-locals, too-many-arguments +# pylint: disable=too-many-branches +"""Plotting Library.""" + +import json +import warnings +from io import BytesIO +from typing import Any, Optional, Union + +import numpy as np + +from ._typing import PathLike +from .core import Booster, _deprecate_positional_args +from .sklearn import XGBModel + +Axes = Any # real type is matplotlib.axes.Axes +GraphvizSource = Any # real type is graphviz.Source + + +@_deprecate_positional_args +def plot_importance( + booster: Union[XGBModel, Booster, dict], + *, + ax: Optional[Axes] = None, + height: float = 0.2, + xlim: Optional[tuple] = None, + ylim: Optional[tuple] = None, + title: str = "Feature importance", + xlabel: str = "Importance score", + ylabel: str = "Features", + fmap: PathLike = "", + importance_type: str = "weight", + max_num_features: Optional[int] = None, + grid: bool = True, + show_values: bool = True, + values_format: str = "{v}", + **kwargs: Any, +) -> Axes: + """Plot importance based on fitted trees. + + Parameters + ---------- + booster : + Booster or XGBModel instance, or dict taken by Booster.get_fscore() + ax : matplotlib Axes + Target axes instance. If None, new figure and axes will be created. + grid : + Turn the axes grids on or off. Default is True (On). + importance_type : + How the importance is calculated: either "weight", "gain", or "cover" + + * "weight" is the number of times a feature appears in a tree + * "gain" is the average gain of splits which use the feature + * "cover" is the average coverage of splits which use the feature + where coverage is defined as the number of samples affected by the split + max_num_features : + Maximum number of top features displayed on plot. If None, all features will be + displayed. + height : + Bar height, passed to ax.barh() + xlim : + Tuple passed to axes.xlim() + ylim : + Tuple passed to axes.ylim() + title : + Axes title. To disable, pass None. + xlabel : + X axis title label. To disable, pass None. + ylabel : + Y axis title label. To disable, pass None. + fmap : + The name of feature map file. + show_values : + Show values on plot. To disable, pass False. + values_format : + Format string for values. "v" will be replaced by the value of the feature + importance. e.g. Pass "{v:.2f}" in order to limit the number of digits after + the decimal point to two, for each value printed on the graph. + kwargs : + Other keywords passed to ax.barh() + + Returns + ------- + ax : matplotlib Axes + """ + try: + import matplotlib.pyplot as plt + except ImportError as e: + raise ImportError("You must install matplotlib to plot importance") from e + + if isinstance(booster, XGBModel): + importance = booster.get_booster().get_score( + importance_type=importance_type, fmap=fmap + ) + elif isinstance(booster, Booster): + importance = booster.get_score(importance_type=importance_type, fmap=fmap) + elif isinstance(booster, dict): + importance = booster + else: + raise ValueError("tree must be Booster, XGBModel or dict instance") + + if not importance: + raise ValueError( + "Booster.get_score() results in empty. " + + "This maybe caused by having all trees as decision dumps." + ) + + tuples = [(k, importance[k]) for k in importance] + if max_num_features is not None: + # pylint: disable=invalid-unary-operand-type + tuples = sorted(tuples, key=lambda _x: _x[1])[-max_num_features:] + else: + tuples = sorted(tuples, key=lambda _x: _x[1]) + labels, values = zip(*tuples) + + if ax is None: + _, ax = plt.subplots(1, 1) + + ylocs = np.arange(len(values)) + ax.barh(ylocs, values, align="center", height=height, **kwargs) + + if show_values is True: + for x, y in zip(values, ylocs): + ax.text(x + 1, float(y), values_format.format(v=x), va="center") + + ax.set_yticks(ylocs) + ax.set_yticklabels(labels) + + if xlim is not None: + if not isinstance(xlim, tuple) or len(xlim) != 2: + raise ValueError("xlim must be a tuple of 2 elements") + else: + xlim = (0, max(values) * 1.1) + ax.set_xlim(xlim) + + if ylim is not None: + if not isinstance(ylim, tuple) or len(ylim) != 2: + raise ValueError("ylim must be a tuple of 2 elements") + else: + ylim = (-1, len(values)) + ax.set_ylim(ylim) + + if title is not None: + ax.set_title(title) + if xlabel is not None: + ax.set_xlabel(xlabel) + if ylabel is not None: + ax.set_ylabel(ylabel) + ax.grid(grid) + return ax + + +@_deprecate_positional_args +def to_graphviz( + booster: Union[Booster, XGBModel], + *, + fmap: PathLike = "", + num_trees: Optional[int] = None, + rankdir: Optional[str] = None, + yes_color: Optional[str] = None, + no_color: Optional[str] = None, + condition_node_params: Optional[dict] = None, + leaf_node_params: Optional[dict] = None, + with_stats: bool = False, + tree_idx: int = 0, + **kwargs: Any, +) -> GraphvizSource: + """Convert specified tree to graphviz instance. IPython can automatically plot + the returned graphviz instance. Otherwise, you should call .render() method + of the returned graphviz instance. + + Parameters + ---------- + booster : + Booster or XGBModel instance + fmap : + The name of feature map file + num_trees : + + .. deprecated:: 3.0 + + Specify the ordinal number of target tree + + rankdir : + Passed to graphviz via graph_attr + yes_color : + Edge color when meets the node condition. + no_color : + Edge color when doesn't meet the node condition. + condition_node_params : + Condition node configuration for for graphviz. Example: + + .. code-block:: python + + {'shape': 'box', + 'style': 'filled,rounded', + 'fillcolor': '#78bceb'} + + leaf_node_params : + Leaf node configuration for graphviz. Example: + + .. code-block:: python + + {'shape': 'box', + 'style': 'filled', + 'fillcolor': '#e48038'} + + with_stats : + + .. versionadded:: 3.0 + + Controls whether the split statistics should be included. + + tree_idx : + + .. versionadded:: 3.0 + + Specify the ordinal index of target tree. + + kwargs : + Other keywords passed to graphviz graph_attr, e.g. ``graph [ {key} = {value} ]`` + + Returns + ------- + graph: graphviz.Source + + """ + try: + from graphviz import Source + except ImportError as e: + raise ImportError("You must install graphviz to plot tree") from e + if isinstance(booster, XGBModel): + booster = booster.get_booster() + + # squash everything back into kwargs again for compatibility + parameters = "dot" + extra = {} + for key, value in kwargs.items(): + extra[key] = value + + if rankdir is not None: + kwargs["graph_attrs"] = {} + kwargs["graph_attrs"]["rankdir"] = rankdir + for key, value in extra.items(): + if kwargs.get("graph_attrs", None) is not None: + kwargs["graph_attrs"][key] = value + else: + kwargs["graph_attrs"] = {key: value} + del kwargs[key] + + if yes_color is not None or no_color is not None: + kwargs["edge"] = {} + if yes_color is not None: + kwargs["edge"]["yes_color"] = yes_color + if no_color is not None: + kwargs["edge"]["no_color"] = no_color + + if condition_node_params is not None: + kwargs["condition_node_params"] = condition_node_params + if leaf_node_params is not None: + kwargs["leaf_node_params"] = leaf_node_params + + if kwargs: + parameters += ":" + parameters += json.dumps(kwargs) + + if num_trees is not None: + warnings.warn( + "The `num_trees` parameter is deprecated, use `tree_idx` instead. ", + FutureWarning, + ) + if tree_idx not in (0, num_trees): + raise ValueError( + "Both `num_trees` and `tree_idx` are used, prefer `tree_idx` instead." + ) + tree_idx = num_trees + + tree = booster.get_dump(fmap=fmap, dump_format=parameters, with_stats=with_stats)[ + tree_idx + ] + g = Source(tree) + return g + + +@_deprecate_positional_args +def plot_tree( + booster: Union[Booster, XGBModel], + *, + fmap: PathLike = "", + num_trees: Optional[int] = None, + rankdir: Optional[str] = None, + ax: Optional[Axes] = None, + with_stats: bool = False, + tree_idx: int = 0, + **kwargs: Any, +) -> Axes: + """Plot specified tree. + + Parameters + ---------- + booster : + Booster or XGBModel instance + fmap: str (optional) + The name of feature map file + num_trees : + + .. deprecated:: 3.0 + + rankdir : str, default "TB" + Passed to graphviz via graph_attr + ax : matplotlib Axes, default None + Target axes instance. If None, new figure and axes will be created. + + with_stats : + + .. versionadded:: 3.0 + + See :py:func:`to_graphviz`. + + tree_idx : + + .. versionadded:: 3.0 + + See :py:func:`to_graphviz`. + + kwargs : + Other keywords passed to :py:func:`to_graphviz` + + Returns + ------- + ax : matplotlib Axes + + """ + try: + from matplotlib import image + from matplotlib import pyplot as plt + except ImportError as e: + raise ImportError("You must install matplotlib to plot tree") from e + + if ax is None: + _, ax = plt.subplots(1, 1) + + g = to_graphviz( + booster, + fmap=fmap, + num_trees=num_trees, + rankdir=rankdir, + with_stats=with_stats, + tree_idx=tree_idx, + **kwargs, + ) + + s = BytesIO() + s.write(g.pipe(format="png")) + s.seek(0) + img = image.imread(s) + + ax.imshow(img) + ax.axis("off") + return ax diff --git a/v12x/pylib/xgboost/sklearn.py b/v12x/pylib/xgboost/sklearn.py new file mode 100644 index 0000000000000000000000000000000000000000..cb3388cbb3cdf186668875540e49cff81546dd75 --- /dev/null +++ b/v12x/pylib/xgboost/sklearn.py @@ -0,0 +1,2445 @@ +# pylint: disable=too-many-arguments, too-many-locals, fixme, too-many-lines +# pylint: disable=duplicate-code +"""Scikit-Learn Wrapper interface for XGBoost.""" + +import collections +import copy +import json +import os +import warnings +from concurrent.futures import ThreadPoolExecutor +from inspect import signature +from typing import ( + Any, + Callable, + Dict, + List, + Optional, + Protocol, + Sequence, + Set, + Tuple, + Type, + TypeVar, + Union, + cast, +) + +import numpy as np +from scipy.special import softmax + +from ._c_api import _parse_version, _py_version +from ._data_utils import Categories +from ._typing import ( + ArrayLike, + EvalsLog, + FeatureNames, + FeatureTypes, + IterationRange, + ModelIn, +) +from .callback import TrainingCallback + +# Do not use class names on scikit-learn directly. Re-define the classes on +# .compat to guarantee the behavior without scikit-learn +from .compat import ( + SKLEARN_INSTALLED, + XGBClassifierBase, + XGBModelBase, + XGBRegressorBase, + _sklearn_Tags, + _sklearn_version, + import_cupy, + is_dataframe, +) +from .config import config_context +from .core import ( + Booster, + CustomObj, + DMatrix, + Metric, + QuantileDMatrix, + XGBoostError, + _deprecate_positional_args, + _parse_eval_str, +) +from .data import ( + CAT_T, + _is_cudf_df, + _is_cudf_ser, + _is_cupy_alike, + _is_pandas_df, + _is_polars_lazyframe, +) +from .objective import Objective +from .training import train + + +class XGBRankerMixIn: + """MixIn for ranking, defines the _estimator_type usually defined in scikit-learn + base classes. + + """ + + _estimator_type = "ranker" + + +class XGBClassifierMixIn(XGBClassifierBase): + """Metadata shared by the single-node and Dask classifier estimators.""" + + n_classes_: int + + def _more_tags(self) -> Dict[str, bool]: + tags = super()._more_tags() + tags["multilabel"] = True + return tags + + def __sklearn_tags__(self) -> _sklearn_Tags: + tags = super().__sklearn_tags__() + tags_dict = self._more_tags() + tags.classifier_tags.multi_label = tags_dict["multilabel"] + return tags + + @property + def classes_(self) -> np.ndarray: + """Classes represented by this estimator.""" + from sklearn.utils.validation import check_is_fitted + + check_is_fitted(self, "n_classes_") + return np.arange(self.n_classes_) + + +class XGBRegressorMixIn(XGBRegressorBase): + """Metadata shared by the single-node and Dask regressor estimators.""" + + def _more_tags(self) -> Dict[str, bool]: + tags = super()._more_tags() + tags["multioutput"] = True + tags["multioutput_only"] = False + return tags + + def __sklearn_tags__(self) -> _sklearn_Tags: + tags = super().__sklearn_tags__() + tags_dict = self._more_tags() + tags.target_tags.multi_output = tags_dict["multioutput"] + tags.target_tags.single_output = not tags_dict["multioutput_only"] + return tags + + +def _check_rf_callback( + early_stopping_rounds: Optional[int], + callbacks: Optional[Sequence[TrainingCallback]], +) -> None: + if early_stopping_rounds is not None or callbacks is not None: + raise NotImplementedError( + "`early_stopping_rounds` and `callbacks` are not implemented for" + " the sklearn random forest estimator interface." + ) + + +def _warn_rf_deprecated(name: str) -> None: + """Emit the deprecation warning for the random forest estimators.""" + warnings.warn( + f"`{name}` is deprecated and will be removed in a future release. The" + " estimator is a thin wrapper over the boosting interface and does not" + " implement a conventional random forest; features like early stopping are" + " unsupported. Set `num_parallel_tree` along with `n_estimators=1` on the" + " corresponding boosting estimator instead, or use a dedicated random forest" + " implementation like those in `sklearn.ensemble`.", + FutureWarning, + stacklevel=3, + ) + + +def _can_use_qdm(tree_method: Optional[str], device: Optional[str]) -> bool: + not_sycl = (device is None) or (not device.startswith("sycl")) + return tree_method in ("hist", None, "auto") and not_sycl + + +class _SklObjWProto(Protocol): + def __call__( + self, + y_true: ArrayLike, + y_pred: ArrayLike, + sample_weight: Optional[ArrayLike] = None, + ) -> Tuple[ArrayLike, ArrayLike]: ... + + +_SklObjProto = Callable[[ArrayLike, ArrayLike], Tuple[np.ndarray, np.ndarray]] +SklObjectiveCallable = Union[Objective, _SklObjWProto, _SklObjProto] +SklObjective = Optional[Union[str, SklObjectiveCallable]] + + +def _objective_decorator(func: SklObjectiveCallable) -> CustomObj: + """Decorate or forward a custom objective. + + Parameters + ---------- + func: + An :py:class:`Objective` instance or a callable with signature + ``func(y_true, y_pred)``: + + y_true: array_like of shape [n_samples] + The target values + y_pred: array_like of shape [n_samples] + The predicted values + sample_weight : + Optional sample weight, None or a ndarray. + + Returns + ------- + new_func: + The original :py:class:`Objective` or a function with the signature + ``new_func(preds, dmatrix)``: + + preds: array_like, shape [n_samples] + The predicted values + dmatrix: ``DMatrix`` + The training set from which the labels will be extracted using + ``dmatrix.get_label()`` + + """ + + if isinstance(func, Objective): + return func + + parameters = signature(func).parameters + supports_sw = "sample_weight" in parameters + + def inner(preds: np.ndarray, dmatrix: DMatrix) -> Tuple[np.ndarray, np.ndarray]: + """Internal function.""" + sample_weight = dmatrix.get_weight() + labels = dmatrix.get_label() + + if sample_weight.size > 0 and not supports_sw: + raise ValueError( + "Custom objective doesn't have the `sample_weight` parameter while" + " sample_weight is used." + ) + if sample_weight.size > 0: + fnw = cast(_SklObjWProto, func) + return fnw(labels, preds, sample_weight=sample_weight) + + fn = cast(_SklObjProto, func) + return fn(labels, preds) + + return inner + + +def _metric_decorator(func: Callable) -> Metric: + """Decorate a metric function from sklearn. + + Converts an metric function that uses the typical sklearn metric signature so that + it is compatible with :py:func:`train` + + """ + + def inner(y_score: np.ndarray, dmatrix: DMatrix) -> Tuple[str, float]: + y_true = dmatrix.get_label() + weight = dmatrix.get_weight() + if weight.size == 0: + return func.__name__, func(y_true, y_score) + return func.__name__, func(y_true, y_score, sample_weight=weight) + + return inner + + +def ltr_metric_decorator(func: Callable, n_jobs: Optional[int]) -> Metric: + """Decorate a learning to rank metric.""" + + def inner(y_score: np.ndarray, dmatrix: DMatrix) -> Tuple[str, float]: + y_true = dmatrix.get_label() + group_ptr = dmatrix.get_uint_info("group_ptr") + if group_ptr.size < 2: + raise ValueError( + "Invalid `group_ptr`. Likely caused by invalid qid or group." + ) + scores = np.empty(group_ptr.size - 1) + futures = [] + weight = dmatrix.get_group() + no_weight = weight.size == 0 + + def task(i: int) -> float: + begin = group_ptr[i - 1] + end = group_ptr[i] + gy = y_true[begin:end] + gp = y_score[begin:end] + if gy.size == 1: + # Maybe there's a better default? 1.0 because many ranking score + # functions have output in range [0, 1]. + return 1.0 + return func(gy, gp) + + workers = n_jobs if n_jobs is not None else os.cpu_count() + with ThreadPoolExecutor(max_workers=workers) as executor: + for i in range(1, group_ptr.size): + f = executor.submit(task, i) + futures.append(f) + + for i, f in enumerate(futures): + scores[i] = f.result() + + if no_weight: + return func.__name__, scores.mean() + + return func.__name__, np.average(scores, weights=weight) + + return inner + + +__estimator_doc = f""" + n_estimators : {Optional[int]} + Number of gradient boosted trees. Equivalent to number of boosting + rounds. +""" + +__model_doc = f""" + max_depth : {Optional[int]} + + Maximum tree depth for base learners. + + max_leaves : {Optional[int]} + + Maximum number of leaves; 0 indicates no limit. + + max_bin : {Optional[int]} + + If using histogram-based algorithm, maximum number of bins per feature + + grow_policy : {Optional[str]} + + Tree growing policy. + + - depthwise: Favors splitting at nodes closest to the node, + - lossguide: Favors splitting at nodes with highest loss change. + + learning_rate : {Optional[float]} + + Boosting learning rate (xgb's "eta") + + verbosity : {Optional[int]} + + The degree of verbosity. Valid values are 0 (silent) - 3 (debug). + + objective : {SklObjective} + + Specify the learning task and the corresponding learning objective or a custom + objective to be used. + + For custom objective, see :doc:`/tutorials/custom_metric_obj` and + :ref:`custom-obj-metric` for more information, along with the end note for + function signatures. + + booster: {Optional[str]} + + Specify which booster to use: ``gbtree``, ``gblinear`` or ``dart``. + + .. deprecated:: 3.3.0 + + ``gblinear`` is deprecated and support will be removed in a future release. + + tree_method : {Optional[str]} + + Specify which tree method to use. Default to auto. If this parameter is set to + default, XGBoost will choose the most conservative option available. It's + recommended to study this option from the parameters document :doc:`tree method + ` + + n_jobs : {Optional[int]} + + Number of parallel threads used to run xgboost. When used with other + Scikit-Learn algorithms like grid search, you may choose which algorithm to + parallelize and balance the threads. Creating thread contention will + significantly slow down both algorithms. + + gamma : {Optional[float]} + + (min_split_loss) Minimum loss reduction required to make a further partition on + a leaf node of the tree. + + min_child_weight : {Optional[float]} + + Minimum sum of instance weight(hessian) needed in a child. + + max_delta_step : {Optional[float]} + + Maximum delta step we allow each tree's weight estimation to be. + + subsample : {Optional[float]} + + Subsample ratio of the training instance. + + sampling_method : {Optional[str]} + + Sampling method. Used only by the GPU version of ``hist`` tree method. + + - ``uniform``: Select random training instances uniformly. + - ``gradient_based``: Select random training instances with higher probability + when the gradient and hessian are larger. (cf. CatBoost) + + colsample_bytree : {Optional[float]} + + Subsample ratio of columns when constructing each tree. + + colsample_bylevel : {Optional[float]} + + Subsample ratio of columns for each level. + + colsample_bynode : {Optional[float]} + + Subsample ratio of columns for each split. + + reg_alpha : {Optional[float]} + + L1 regularization term on weights (xgb's alpha). + + reg_lambda : {Optional[float]} + + L2 regularization term on weights (xgb's lambda). + + scale_pos_weight : {Optional[float]} + Balancing of positive and negative weights. + + base_score : {Optional[Union[float, List[float]]]} + + The initial prediction score of all instances, global bias. + + random_state : {Optional[Union[np.random.RandomState, np.random.Generator, int]]} + + Random number seed. + + .. note:: + + Using gblinear booster with shotgun updater is nondeterministic as + it uses Hogwild algorithm. + + missing : float + + Value in the data which needs to be present as a missing value. Default to + :py:data:`numpy.nan`. + + num_parallel_tree: {Optional[int]} + + Used for boosting random forest. + + monotone_constraints : {Optional[Union[Dict[str, int], str]]} + + Constraint of variable monotonicity. See :doc:`tutorial ` + for more information. + + interaction_constraints : {Optional[Union[str, List[Tuple[str]]]]} + + Constraints for interaction representing permitted interactions. The + constraints must be specified in the form of a nested list, e.g. ``[[0, 1], [2, + 3, 4]]``, where each inner list is a group of indices of features that are + allowed to interact with each other. See :doc:`tutorial + ` for more information + + importance_type: {Optional[str]} + + The feature importance type for the feature_importances\\_ property: + + * For tree model, it's either "gain", "weight", "cover", "total_gain" or + "total_cover". + * For linear model, only "weight" is defined and it's the normalized + coefficients without bias. + + device : {Optional[str]} + + .. versionadded:: 2.0.0 + + Device ordinal, available options are `cpu`, `cuda`, and `gpu`. + + validate_parameters : {Optional[bool]} + + Give warnings for unknown parameter. + + enable_categorical : bool + + See the same parameter of :py:class:`DMatrix` for details. + + feature_types : {Optional[FeatureTypes]} + + .. versionadded:: 1.7.0 + + Used for specifying feature types without constructing a dataframe. See + the :py:class:`DMatrix` for details. + + feature_weights : Optional[ArrayLike] + + Weight for each feature, defines the probability of each feature being selected + when colsample is being used. All values must be greater than 0, otherwise a + `ValueError` is thrown. + + max_cat_to_onehot : Optional[int] + + .. versionadded:: 1.6.0 + + .. note:: This parameter is experimental + + A threshold for deciding whether XGBoost should use one-hot encoding based split + for categorical data. When number of categories is lesser than the threshold + then one-hot encoding is chosen, otherwise the categories will be partitioned + into children nodes. Also, `enable_categorical` needs to be set to have + categorical feature support. See :doc:`Categorical Data + ` and :ref:`cat-param` for details. + + max_cat_threshold : {Optional[int]} + + .. versionadded:: 1.7.0 + + .. note:: This parameter is experimental + + Maximum number of categories considered for each split. Used only by + partition-based splits for preventing over-fitting. Also, `enable_categorical` + needs to be set to have categorical feature support. See :doc:`Categorical Data + ` and :ref:`cat-param` for details. + + multi_strategy : {Optional[str]} + + .. versionadded:: 2.0.0 + + .. note:: This parameter is working-in-progress. + + The strategy used for training multi-target models, including multi-target + regression and multi-class classification. See :doc:`/tutorials/multioutput` for + more information. + + - ``one_output_per_tree``: One model for each target. + - ``multi_output_tree``: Use multi-target trees. + + eval_metric : {Optional[Union[str, List[Union[str, Callable]], Callable]]} + + .. versionadded:: 1.6.0 + + Metric used for monitoring the training result and early stopping. It can be a + string or list of strings as names of predefined metric in XGBoost (See + :doc:`/parameter`), one of the metrics in :py:mod:`sklearn.metrics`, or any + other user defined metric that looks like `sklearn.metrics`. + + If custom objective is also provided, then custom metric should implement the + corresponding reverse link function. + + Unlike the `scoring` parameter commonly used in scikit-learn, when a callable + object is provided, it's assumed to be a cost function and by default XGBoost + will minimize the result during early stopping. + + For advanced usage on Early stopping like directly choosing to maximize instead + of minimize, see :py:obj:`xgboost.callback.EarlyStopping`. + + See :doc:`/tutorials/custom_metric_obj` and :ref:`custom-obj-metric` for more + information. + + .. code-block:: python + + from sklearn.datasets import load_diabetes + from sklearn.metrics import mean_absolute_error + X, y = load_diabetes(return_X_y=True) + reg = xgb.XGBRegressor( + tree_method="hist", + eval_metric=mean_absolute_error, + ) + reg.fit(X, y, eval_set=[(X, y)]) + + early_stopping_rounds : {Optional[int]} + + .. versionadded:: 1.6.0 + + - Activates early stopping. Validation metric needs to improve at least once in + every **early_stopping_rounds** round(s) to continue training. Requires at + least one item in **eval_set** in :py:meth:`fit`. + + - If early stopping occurs, the model will have two additional attributes: + :py:attr:`best_score` and :py:attr:`best_iteration`. These are used by the + :py:meth:`predict` and :py:meth:`apply` methods to determine the optimal + number of trees during inference. If users want to access the full model + (including trees built after early stopping), they can specify the + `iteration_range` in these inference methods. In addition, other utilities + like model plotting can also use the entire model. + + - If you prefer to discard the trees after `best_iteration`, consider using the + callback function :py:class:`xgboost.callback.EarlyStopping`. + + - If there's more than one item in **eval_set**, the last entry will be used for + early stopping. If there's more than one metric in **eval_metric**, the last + metric will be used for early stopping. + + callbacks : {Optional[List[TrainingCallback]]} + + List of callback functions that are applied at end of each iteration. + It is possible to use predefined callbacks by using + :ref:`Callback API `. + + .. note:: + + States in callback are not preserved during training, which means callback + objects can not be reused for multiple training sessions without + reinitialization or deepcopy. + + .. code-block:: python + + for params in parameters_grid: + # be sure to (re)initialize the callbacks before each run + callbacks = [xgb.callback.LearningRateScheduler(custom_rates)] + reg = xgboost.XGBRegressor(**params, callbacks=callbacks) + reg.fit(X, y) + + kwargs : {Optional[Any]} + + Keyword arguments for XGBoost Booster object. Full documentation of parameters + can be found :doc:`here `. + Attempting to set a parameter via the constructor args and \\*\\*kwargs + dict simultaneously will result in a TypeError. + + .. note:: \\*\\*kwargs unsupported by scikit-learn + + \\*\\*kwargs is unsupported by scikit-learn. We do not guarantee + that parameters passed via this argument will interact properly + with scikit-learn. +""" + +__custom_obj_note = """ + .. note:: Custom objective function + + A custom objective function can be provided for the ``objective`` + parameter. In this case, it should have the signature ``objective(y_true, + y_pred) -> [grad, hess]`` or ``objective(y_true, y_pred, *, sample_weight) + -> [grad, hess]``: + + y_true: array_like of shape [n_samples] + The target values + y_pred: array_like of shape [n_samples] + The predicted values + sample_weight : + Optional sample weights. + + grad: array_like of shape [n_samples] + The value of the gradient for each sample point. + hess: array_like of shape [n_samples] + The value of the second derivative for each sample point + + Note that, if the custom objective produces negative values for + the Hessian, these will be clipped. If the objective is non-convex, + one might also consider using the expected Hessian (Fisher + information) instead. +""" + +TDoc = TypeVar("TDoc", bound=Type) + + +def xgboost_model_doc( + header: str, + items: List[str], + extra_parameters: Optional[str] = None, + end_note: Optional[str] = None, +) -> Callable[[TDoc], TDoc]: + """Obtain documentation for Scikit-Learn wrappers + + Parameters + ---------- + header: str + An introducion to the class. + items : list + A list of common doc items. Available items are: + - estimators: the meaning of n_estimators + - model: All the other parameters + - objective: note for customized objective + extra_parameters: str + Document for class specific parameters, placed at the head. + end_note: str + Extra notes put to the end.""" + + def get_doc(item: str) -> str: + """Return selected item""" + __doc = { + "estimators": __estimator_doc, + "model": __model_doc, + "objective": __custom_obj_note, + } + return __doc[item] + + def adddoc(cls: TDoc) -> TDoc: + doc = [ + """ +Parameters +---------- +""" + ] + if extra_parameters: + doc.append(extra_parameters) + doc.extend([get_doc(i) for i in items]) + if end_note: + doc.append(end_note) + full_doc = [ + header + "\nSee :doc:`/python/sklearn_estimator` for more information.\n" + ] + full_doc.extend(doc) + cls.__doc__ = "".join(full_doc) + return cls + + return adddoc + + +def get_model_categories( + X: ArrayLike, + model: Optional[Union[Booster, str]], + feature_types: Optional[FeatureTypes], +) -> Tuple[Optional[Union[Booster, str]], Optional[Union[FeatureTypes, Categories]]]: + """Extract the optional reference categories from the booster. Used for training + continuation. The result should be passed to the :py:func:`pick_ref_categories`. + + """ + # Skip if it's not a dataframe as there's no new encoding to be recoded. + # + # This function helps override the `feature_types` parameter. The `feature_types` + # from user is not useful when input is a dataframe as the real feature type should + # be encoded into the DF. + if model is None or not is_dataframe(X): + return model, feature_types + + if isinstance(model, str): + model = Booster(model_file=model) + + categories = model.get_categories() + if not categories.empty(): + # override the `feature_types`. + return model, categories + # Convert empty into None. + return model, feature_types + + +def pick_ref_categories( + X: Any, + model_cats: Optional[Union[FeatureTypes, Categories]], + Xy_cats: Optional[Categories], +) -> Optional[Union[FeatureTypes, Categories]]: + """Use the reference categories from the model. If none, then use the reference + categories from the training DMatrix. + + Parameters + ---------- + X : + Input feature matrix. + + model_cats : + Optional categories stored in the previous model (training continuation). This + should come from the :py:func:`get_model_categories`. + + Xy_cats : + Optional categories from the training DMatrix. Used for re-coding the validation + dataset. + + """ + categories: Optional[Categories] = None + if not isinstance(model_cats, Categories) and is_dataframe(X): + categories = Xy_cats + if categories is not None and not categories.empty(): + model_cats = categories + + return model_cats + + +def _wrap_evaluation_matrices( + *, + missing: float, + X: Any, + y: Any, + group: Optional[Any], + qid: Optional[Any], + sample_weight: Optional[Any], + base_margin: Optional[Any], + feature_weights: Optional[ArrayLike], + eval_set: Optional[Sequence[Tuple[Any, Any]]], + sample_weight_eval_set: Optional[Sequence[Any]], + base_margin_eval_set: Optional[Sequence[Any]], + eval_group: Optional[Sequence[Any]], + eval_qid: Optional[Sequence[Any]], + create_dmatrix: Callable, + enable_categorical: bool, + feature_types: Optional[Union[FeatureTypes, Categories]], +) -> Tuple[Any, List[Tuple[Any, str]]]: + """Convert array_like evaluation matrices into DMatrix. Perform sanity checks on the + way. + + """ + # Feature_types contains the optional reference categories from the booster object. + train_dmatrix = create_dmatrix( + data=X, + label=y, + group=group, + qid=qid, + weight=sample_weight, + base_margin=base_margin, + feature_weights=feature_weights, + missing=missing, + enable_categorical=enable_categorical, + feature_types=feature_types, + ref=None, + ) + + n_validation = 0 if eval_set is None else len(eval_set) + if hasattr(train_dmatrix, "get_categories"): + Xy_cats = train_dmatrix.get_categories() + else: + Xy_cats = None + + def validate_or_none(meta: Optional[Sequence], name: str) -> Sequence: + if meta is None: + return [None] * n_validation + if len(meta) != n_validation: + raise ValueError( + f"{name}'s length does not equal `eval_set`'s length, " + + f"expecting {n_validation}, got {len(meta)}" + ) + return meta + + if eval_set is not None: + sample_weight_eval_set = validate_or_none( + sample_weight_eval_set, "sample_weight_eval_set" + ) + base_margin_eval_set = validate_or_none( + base_margin_eval_set, "base_margin_eval_set" + ) + eval_group = validate_or_none(eval_group, "eval_group") + eval_qid = validate_or_none(eval_qid, "eval_qid") + + evals = [] + for i, (valid_X, valid_y) in enumerate(eval_set): + # Skip the entry if it's the training DMatrix. + if all( + ( + valid_X is X, + valid_y is y, + sample_weight_eval_set[i] is sample_weight, + base_margin_eval_set[i] is base_margin, + eval_group[i] is group, + eval_qid[i] is qid, + ) + ): + evals.append(train_dmatrix) + continue + + feature_types = pick_ref_categories(valid_X, feature_types, Xy_cats) + m = create_dmatrix( + data=valid_X, + label=valid_y, + weight=sample_weight_eval_set[i], + group=eval_group[i], + qid=eval_qid[i], + base_margin=base_margin_eval_set[i], + missing=missing, + enable_categorical=enable_categorical, + feature_types=feature_types, + ref=train_dmatrix, + ) + evals.append(m) + + nevals = len(evals) + eval_names = [f"validation_{i}" for i in range(nevals)] + evals = list(zip(evals, eval_names)) + else: + if any( + meta is not None + for meta in [ + sample_weight_eval_set, + base_margin_eval_set, + eval_group, + eval_qid, + ] + ): + raise ValueError( + "`eval_set` is not set but one of the other evaluation meta info is " + "not None." + ) + evals = [] + + return train_dmatrix, evals + + +DEFAULT_N_ESTIMATORS = 100 + + +@xgboost_model_doc( + """Implementation of the Scikit-Learn API for XGBoost.""", + ["estimators", "model", "objective"], +) +class XGBModel(XGBModelBase): + # pylint: disable=too-many-arguments, too-many-instance-attributes, missing-docstring + @_deprecate_positional_args + def __init__( + self, + *, + max_depth: Optional[int] = None, + max_leaves: Optional[int] = None, + max_bin: Optional[int] = None, + grow_policy: Optional[str] = None, + learning_rate: Optional[float] = None, + n_estimators: Optional[int] = None, + verbosity: Optional[int] = None, + objective: SklObjective = None, + booster: Optional[str] = None, + tree_method: Optional[str] = None, + n_jobs: Optional[int] = None, + gamma: Optional[float] = None, + min_child_weight: Optional[float] = None, + max_delta_step: Optional[float] = None, + subsample: Optional[float] = None, + sampling_method: Optional[str] = None, + colsample_bytree: Optional[float] = None, + colsample_bylevel: Optional[float] = None, + colsample_bynode: Optional[float] = None, + reg_alpha: Optional[float] = None, + reg_lambda: Optional[float] = None, + scale_pos_weight: Optional[float] = None, + base_score: Optional[Union[float, List[float]]] = None, + random_state: Optional[ + Union[np.random.RandomState, np.random.Generator, int] + ] = None, + missing: float = np.nan, + num_parallel_tree: Optional[int] = None, + monotone_constraints: Optional[Union[Dict[str, int], str]] = None, + interaction_constraints: Optional[Union[str, Sequence[Sequence[str]]]] = None, + importance_type: Optional[str] = None, + device: Optional[str] = None, + validate_parameters: Optional[bool] = None, + enable_categorical: bool = True, + feature_types: Optional[FeatureTypes] = None, + feature_weights: Optional[ArrayLike] = None, + max_cat_to_onehot: Optional[int] = None, + max_cat_threshold: Optional[int] = None, + multi_strategy: Optional[str] = None, + eval_metric: Optional[Union[str, List[Union[str, Callable]], Callable]] = None, + early_stopping_rounds: Optional[int] = None, + callbacks: Optional[List[TrainingCallback]] = None, + **kwargs: Any, + ) -> None: + if not SKLEARN_INSTALLED: + raise ImportError( + "sklearn needs to be installed in order to use this module" + ) + self.n_estimators = n_estimators + self.objective = objective + + self.max_depth = max_depth + self.max_leaves = max_leaves + self.max_bin = max_bin + self.grow_policy = grow_policy + self.learning_rate = learning_rate + self.verbosity = verbosity + self.booster = booster + self.tree_method = tree_method + self.gamma = gamma + self.min_child_weight = min_child_weight + self.max_delta_step = max_delta_step + self.subsample = subsample + self.sampling_method = sampling_method + self.colsample_bytree = colsample_bytree + self.colsample_bylevel = colsample_bylevel + self.colsample_bynode = colsample_bynode + self.reg_alpha = reg_alpha + self.reg_lambda = reg_lambda + self.scale_pos_weight = scale_pos_weight + self.base_score = base_score + self.missing = missing + self.num_parallel_tree = num_parallel_tree + self.random_state = random_state + self.n_jobs = n_jobs + self.monotone_constraints = monotone_constraints + self.interaction_constraints = interaction_constraints + self.importance_type = importance_type + self.device = device + self.validate_parameters = validate_parameters + self.enable_categorical = enable_categorical + self.feature_types = feature_types + if isinstance(self.feature_types, Categories): + raise TypeError( + "If you are training with a prior model (training continuation), " + "The scikit-learn interface can automatically reuse the categories from" + " that model." + ) + self.feature_weights = feature_weights + self.max_cat_to_onehot = max_cat_to_onehot + self.max_cat_threshold = max_cat_threshold + self.multi_strategy = multi_strategy + self.eval_metric = eval_metric + self.early_stopping_rounds = early_stopping_rounds + self.callbacks = callbacks + if kwargs: + self.kwargs = kwargs + + def _more_tags(self) -> Dict[str, bool]: + """Tags used for scikit-learn data validation.""" + tags = {"allow_nan": True, "no_validation": True, "sparse": True} + if hasattr(self, "kwargs") and self.kwargs.get("updater") == "shotgun": + tags["non_deterministic"] = True + + tags["categorical"] = self.enable_categorical + tags["string"] = self.enable_categorical + return tags + + @staticmethod + def _update_sklearn_tags_from_dict( + *, + tags: _sklearn_Tags, + tags_dict: Dict[str, bool], + ) -> _sklearn_Tags: + """Update ``sklearn.utils.Tags`` inherited from ``scikit-learn`` base classes. + + ``scikit-learn`` 1.6 introduced a dataclass-based interface for estimator tags. + ref: https://github.com/scikit-learn/scikit-learn/pull/29677 + + This method handles updating that instance based on the values in + ``self._more_tags()``. + + """ + tags.non_deterministic = tags_dict.get("non_deterministic", False) + tags.no_validation = tags_dict["no_validation"] + tags.input_tags.allow_nan = tags_dict["allow_nan"] + tags.input_tags.sparse = tags_dict["sparse"] + tags.input_tags.categorical = tags_dict["categorical"] + tags.input_tags.string = tags_dict["string"] + return tags + + def __sklearn_tags__(self) -> _sklearn_Tags: + # XGBModelBase.__sklearn_tags__() cannot be called unconditionally, + # because that method isn't defined for scikit-learn<1.6 + if not hasattr(XGBModelBase, "__sklearn_tags__"): + err_msg = ( + "__sklearn_tags__() should not be called when using scikit-learn<1.6. " + f"Detected version: {_sklearn_version}" + ) + raise AttributeError(err_msg) + + # take whatever tags are provided by BaseEstimator, then modify + # them with XGBoost-specific values + return self._update_sklearn_tags_from_dict( + tags=super().__sklearn_tags__(), # pylint: disable=no-member + tags_dict=self._more_tags(), + ) + + def __sklearn_is_fitted__(self) -> bool: + return hasattr(self, "_Booster") + + @property + def _doc_link_module(self) -> str: + return "xgboost" + + @property + def _doc_link_template(self) -> str: + ver = _py_version() + (major, minor, _), post = _parse_version(ver) + + if post == "dev": + rel = "latest" + else: + # RTD tracks the release branch. We don't have independent branches for + # patch releases. + rel = f"release_{major}.{minor}.0" + + module = self.__class__.__module__ + # All sklearn estimators are forwarded to the top level module in both source + # code and sphinx api doc. + if module == "xgboost.sklearn": + module = module.split(".")[0] + name = self.__class__.__name__ + + base = "https://xgboost.readthedocs.io/en" + return f"{base}/{rel}/python/python_api.html#{module}.{name}" + + def _wrapper_params(self) -> Set[str]: + wrapper_specific = { + "importance_type", + "kwargs", + "missing", + "n_estimators", + "enable_categorical", + "early_stopping_rounds", + "callbacks", + "feature_types", + "feature_weights", + } + return wrapper_specific + + def get_booster(self) -> Booster: + """Get the underlying xgboost Booster of this model. + + This will raise an exception when fit was not called + + Returns + ------- + booster : a xgboost booster of underlying model + """ + if not self.__sklearn_is_fitted__(): + from sklearn.exceptions import NotFittedError + + raise NotFittedError("need to call fit or load_model beforehand") + return self._Booster + + def set_params(self, **params: Any) -> "XGBModel": + """Set the parameters of this estimator. Modification of the sklearn method to + allow unknown kwargs. This allows using the full range of xgboost + parameters that are not defined as member variables in sklearn grid + search. + + Returns + ------- + self + + """ + if not params: + # Simple optimization to gain speed (inspect is slow) + return self + + # this concatenates kwargs into parameters, enabling `get_params` for + # obtaining parameters from keyword parameters. + for key, value in params.items(): + if hasattr(self, key): + setattr(self, key, value) + else: + if not hasattr(self, "kwargs"): + self.kwargs = {} + self.kwargs[key] = value + + if self.__sklearn_is_fitted__(): + parameters = self.get_xgb_params() + self.get_booster().set_param(parameters) + + return self + + def get_params(self, deep: bool = True) -> Dict[str, Any]: + # pylint: disable=attribute-defined-outside-init + """Get parameters.""" + # Based on: https://stackoverflow.com/questions/59248211 + # The basic flow in `get_params` is: + # 0. Return parameters in subclass (self.__class__) first, by using inspect. + # 1. Return parameters in all parent classes (especially `XGBModel`). + # 2. Return whatever in `**kwargs`. + # 3. Merge them. + # + # This needs to accommodate being called recursively in the following + # inheritance graphs (and similar for classification and ranking): + # + # XGBRFRegressor -> XGBRegressor -> XGBModel -> BaseEstimator + # XGBRegressor -> XGBModel -> BaseEstimator + # XGBModel -> BaseEstimator + # + params = super().get_params(deep) # pylint: disable=no-member + cp = copy.copy(self) + # If the immediate parent defines get_params(), use that. + if callable(getattr(cp.__class__.__bases__[0], "get_params", None)): + cp.__class__ = cp.__class__.__bases__[0] + # Otherwise, skip it and assume the next class will have it. + # This is here primarily for cases where the first class in MRO is a scikit-learn mixin. + else: + cp.__class__ = cp.__class__.__bases__[1] + params.update(cp.__class__.get_params(cp, deep)) + # if kwargs is a dict, update params accordingly + if hasattr(self, "kwargs") and isinstance(self.kwargs, dict): + params.update(self.kwargs) + if isinstance(params["random_state"], np.random.RandomState): + params["random_state"] = params["random_state"].randint( + np.iinfo(np.int32).max + ) + elif isinstance(params["random_state"], np.random.Generator): + params["random_state"] = int( + params["random_state"].integers(np.iinfo(np.int32).max) + ) + + return params + + def get_xgb_params(self) -> Dict[str, Any]: + """Get xgboost specific parameters.""" + params: Dict[str, Any] = self.get_params() + + # Parameters that should not go into native learner. + wrapper_specific = self._wrapper_params() + filtered = {} + for k, v in params.items(): + if k not in wrapper_specific and not callable(v): + filtered[k] = v + + return filtered + + def get_num_boosting_rounds(self) -> int: + """Gets the number of xgboost boosting rounds.""" + return DEFAULT_N_ESTIMATORS if self.n_estimators is None else self.n_estimators + + def _get_type(self) -> str: + if hasattr(self, "_estimator_type"): # scikit-learn <1.8 + return self._estimator_type # pylint: disable=no-member + if hasattr(XGBModelBase, "__sklearn_tags__"): # scikit-learn 1.8+ + return self.__sklearn_tags__().estimator_type + raise TypeError( + "`_estimator_type` undefined. " + "Please use appropriate mixin to define estimator type." + ) + + def save_model(self, fname: Union[str, os.PathLike]) -> None: + meta: Dict[str, Any] = {} + # For validation. + meta["_estimator_type"] = self._get_type() + meta_str = json.dumps(meta) + self.get_booster().set_attr(scikit_learn=meta_str) + self.get_booster().save_model(fname) + self.get_booster().set_attr(scikit_learn=None) + + if Booster.save_model.__doc__ is not None: + save_model.__doc__ = f"""{Booster.save_model.__doc__}""" + + def load_model(self, fname: ModelIn) -> None: + # pylint: disable=attribute-defined-outside-init + if not self.__sklearn_is_fitted__(): + self._Booster = Booster({"n_jobs": self.n_jobs}) + self.get_booster().load_model(fname) + + meta_str = self.get_booster().attr("scikit_learn") + if meta_str is not None: + meta = json.loads(meta_str) + t = meta.get("_estimator_type", None) + if t is not None and t != self._get_type(): + raise TypeError( + "Loading an estimator with different type. Expecting: " + f"{self._get_type()}, got: {t}" + ) + + self.get_booster().set_attr(scikit_learn=None) + config = json.loads(self.get_booster().save_config()) + self._load_model_attributes(config) + + if Booster.load_model.__doc__ is not None: + load_model.__doc__ = f"""{Booster.load_model.__doc__}""" + + def _load_model_attributes(self, config: dict) -> None: + """Load model attributes without hyper-parameters.""" + from sklearn.base import is_classifier + + booster = self.get_booster() + + self.objective = config["learner"]["objective"]["name"] + self.booster = config["learner"]["gradient_booster"]["name"] + self.base_score = json.loads( + config["learner"]["learner_model_param"]["base_score"] + ) + self.feature_types = booster.feature_types + self.enable_categorical = self.feature_types is not None and any( + ft == CAT_T for ft in self.feature_types + ) + + if is_classifier(self): + self.n_classes_ = int(config["learner"]["learner_model_param"]["num_class"]) + # binary classification is treated as regression in XGBoost. + self.n_classes_ = 2 if self.n_classes_ < 2 else self.n_classes_ + + # pylint: disable=too-many-branches + def _configure_fit( + self, + booster: Optional[Union[Booster, "XGBModel", str]], + params: Dict[str, Any], + feature_weights: Optional[ArrayLike], + ) -> Tuple[ + Optional[Union[Booster, str, "XGBModel"]], + Optional[Metric], + Dict[str, Any], + Optional[ArrayLike], + ]: + """Configure parameters for :py:meth:`fit`.""" + if isinstance(booster, XGBModel): + model: Optional[Union[Booster, str]] = booster.get_booster() + else: + model = booster + + def _deprecated(parameter: str) -> None: + warnings.warn( + f"`{parameter}` in `fit` method is deprecated for better compatibility " + f"with scikit-learn, use `{parameter}` in constructor or`set_params` " + "instead.", + UserWarning, + ) + + def _duplicated(parameter: str) -> None: + raise ValueError( + f"2 different `{parameter}` are provided. Use the one in constructor " + "or `set_params` instead." + ) + + # - configure callable evaluation metric + metric: Optional[Metric] = None + + def custom_metric(m: Callable) -> Metric: + if self._get_type() == "ranker": + wrapped = ltr_metric_decorator(m, self.n_jobs) + else: + wrapped = _metric_decorator(m) + return wrapped + + def invalid_type(m: Any) -> None: + msg = f"Invalid type for the `eval_metric`: {type(m)}" + raise TypeError(msg) + + if self.eval_metric is not None: + if callable(self.eval_metric): + metric = custom_metric(self.eval_metric) + elif isinstance(self.eval_metric, str): + params.update({"eval_metric": self.eval_metric}) + else: + # A sequence of metrics + if not isinstance(self.eval_metric, collections.abc.Sequence): + invalid_type(self.eval_metric) + # Could be a list of strings or callables + builtin_metrics: List[str] = [] + for m in self.eval_metric: + if callable(m): + if metric is not None: + raise NotImplementedError( + "Using multiple custom metrics is not yet supported." + ) + metric = custom_metric(m) + elif isinstance(m, str): + builtin_metrics.append(m) + else: + invalid_type(m) + if builtin_metrics: + params.update({"eval_metric": builtin_metrics}) + + if feature_weights is not None: + _deprecated("feature_weights") + if feature_weights is not None and self.feature_weights is not None: + _duplicated("feature_weights") + feature_weights = ( + self.feature_weights + if self.feature_weights is not None + else feature_weights + ) + return model, metric, params, feature_weights + + def _create_dmatrix(self, ref: Optional[DMatrix], **kwargs: Any) -> DMatrix: + # Use `QuantileDMatrix` to save memory. + if _can_use_qdm(self.tree_method, self.device) and self.booster != "gblinear": + try: + return QuantileDMatrix( + **kwargs, ref=ref, nthread=self.n_jobs, max_bin=self.max_bin + ) + except TypeError: # `QuantileDMatrix` supports lesser types than DMatrix + pass + return DMatrix(**kwargs, nthread=self.n_jobs) + + def _set_evaluation_result(self, evals_result: EvalsLog) -> None: + if evals_result: + self.evals_result_ = cast(Dict[str, Dict[str, List[float]]], evals_result) + + @_deprecate_positional_args + def fit( + self, + X: ArrayLike, + y: ArrayLike, + *, + sample_weight: Optional[ArrayLike] = None, + base_margin: Optional[ArrayLike] = None, + eval_set: Optional[Sequence[Tuple[ArrayLike, ArrayLike]]] = None, + verbose: Optional[Union[bool, int]] = True, + xgb_model: Optional[Union[Booster, str, "XGBModel"]] = None, + sample_weight_eval_set: Optional[Sequence[ArrayLike]] = None, + base_margin_eval_set: Optional[Sequence[ArrayLike]] = None, + feature_weights: Optional[ArrayLike] = None, + ) -> "XGBModel": + # pylint: disable=attribute-defined-outside-init + """Fit gradient boosting model. + + Note that calling ``fit()`` multiple times will cause the model object to be + re-fit from scratch. To resume training from a previous checkpoint, explicitly + pass ``xgb_model`` argument. + + Parameters + ---------- + X : + Input feature matrix. See :ref:`py-data` for a list of supported types. + + When the ``tree_method`` is set to ``hist``, internally, the + :py:class:`QuantileDMatrix` will be used instead of the :py:class:`DMatrix` + for conserving memory. However, this has performance implications when the + device of input data is not matched with algorithm. For instance, if the + input is a numpy array on CPU but ``cuda`` is used for training, then the + data is first processed on CPU then transferred to GPU. + y : + Labels + sample_weight : + instance weights + base_margin : + Global bias for each instance. See :doc:`/tutorials/intercept` for details. + eval_set : + A list of (X, y) tuple pairs to use as validation sets, for which + metrics will be computed. + Validation metrics will help us track the performance of the model. + + verbose : + If `verbose` is True and an evaluation set is used, the evaluation metric + measured on the validation set is printed to stdout at each boosting stage. + If `verbose` is an integer, the evaluation metric is printed at each + `verbose` boosting stage. The last boosting stage / the boosting stage found + by using `early_stopping_rounds` is also printed. + xgb_model : + file name of stored XGBoost model or 'Booster' instance XGBoost model to be + loaded before training (allows training continuation). + sample_weight_eval_set : + A list of the form [L_1, L_2, ..., L_n], where each L_i is an array like + object storing instance weights for the i-th validation set. + base_margin_eval_set : + A list of the form [M_1, M_2, ..., M_n], where each M_i is an array like + object storing base margin for the i-th validation set. + feature_weights : + + .. deprecated:: 3.0.0 + + Use `feature_weights` in :py:meth:`__init__` or :py:meth:`set_params` + instead. + + """ + with config_context(verbosity=self.verbosity): + params = self.get_xgb_params() + model, metric, params, feature_weights = self._configure_fit( + xgb_model, params, feature_weights + ) + model, feature_types = get_model_categories(X, model, self.feature_types) + + evals_result: EvalsLog = {} + train_dmatrix, evals = _wrap_evaluation_matrices( + missing=self.missing, + X=X, + y=y, + group=None, + qid=None, + sample_weight=sample_weight, + base_margin=base_margin, + feature_weights=feature_weights, + eval_set=eval_set, + sample_weight_eval_set=sample_weight_eval_set, + base_margin_eval_set=base_margin_eval_set, + eval_group=None, + eval_qid=None, + create_dmatrix=self._create_dmatrix, + enable_categorical=self.enable_categorical, + feature_types=feature_types, + ) + + if callable(self.objective): + obj: Optional[CustomObj] = _objective_decorator(self.objective) + params["objective"] = "reg:squarederror" + else: + obj = None + + self._Booster = train( + params, + train_dmatrix, + self.get_num_boosting_rounds(), + evals=evals, + early_stopping_rounds=self.early_stopping_rounds, + evals_result=evals_result, + obj=obj, + custom_metric=metric, + verbose_eval=verbose, + xgb_model=model, + callbacks=self.callbacks, + ) + + self._set_evaluation_result(evals_result) + return self + + def _can_use_inplace_predict(self) -> bool: + return self.booster != "gblinear" + + def _get_iteration_range( + self, iteration_range: Optional[IterationRange] + ) -> IterationRange: + if iteration_range is None or iteration_range[1] == 0: + # Use best_iteration if defined. + try: + iteration_range = (0, self.best_iteration + 1) + except AttributeError: + iteration_range = (0, 0) + if self.booster == "gblinear": + iteration_range = (0, 0) + return iteration_range + + @_deprecate_positional_args + def predict( + self, + X: ArrayLike, + *, + output_margin: bool = False, + validate_features: bool = True, + base_margin: Optional[ArrayLike] = None, + iteration_range: Optional[IterationRange] = None, + ) -> ArrayLike: + """Predict with `X`. If the model is trained with early stopping, then + :py:attr:`best_iteration` is used automatically. The estimator uses + `inplace_predict` by default and falls back to using :py:class:`DMatrix` if + devices between the data and the estimator don't match. + + .. note:: This function is only thread safe for `gbtree` and `dart`. + + Parameters + ---------- + X : + Data to predict with. See :ref:`py-data` for a list of supported types. + output_margin : + Whether to output the raw untransformed margin value. + validate_features : + When this is True, validate that the Booster's and data's feature_names are + identical. Otherwise, it is assumed that the feature_names are the same. + base_margin : + Global bias for each instance. See :doc:`/tutorials/intercept` for details. + iteration_range : + Specifies which layer of trees are used in prediction. For example, if a + random forest is trained with 100 rounds. Specifying ``iteration_range=(10, + 20)``, then only the forests built during [10, 20) (half open set) rounds + are used in this prediction. + + .. versionadded:: 1.4.0 + + Returns + ------- + prediction + + """ + with config_context(verbosity=self.verbosity): + iteration_range = self._get_iteration_range(iteration_range) + if self._can_use_inplace_predict(): + try: + predts = self.get_booster().inplace_predict( + data=X, + iteration_range=iteration_range, + predict_type="margin" if output_margin else "value", + missing=self.missing, + base_margin=base_margin, + validate_features=validate_features, + ) + if _is_cupy_alike(predts): + cp = import_cupy() + + predts = cp.asnumpy(predts) # ensure numpy array is used. + return predts + except TypeError: + # coo, csc, dt + pass + + test = DMatrix( + X, + base_margin=base_margin, + missing=self.missing, + nthread=self.n_jobs, + feature_types=self.feature_types, + enable_categorical=self.enable_categorical, + ) + return self.get_booster().predict( + data=test, + iteration_range=iteration_range, + output_margin=output_margin, + validate_features=validate_features, + ) + + def apply( + self, + X: ArrayLike, + iteration_range: Optional[IterationRange] = None, + ) -> np.ndarray: + """Return the predicted leaf every tree for each sample. If the model is trained + with early stopping, then :py:attr:`best_iteration` is used automatically. + + Parameters + ---------- + X : + Input features matrix. See :ref:`py-data` for a list of supported types. + + iteration_range : + See :py:meth:`predict`. + + Returns + ------- + X_leaves : array_like, shape=[n_samples, n_trees] + For each datapoint x in X and for each tree, return the index of the + leaf x ends up in. Leaves are numbered within + ``[0; 2**(self.max_depth+1))``, possibly with gaps in the numbering. + + """ + with config_context(verbosity=self.verbosity): + iteration_range = self._get_iteration_range(iteration_range) + test_dmatrix = DMatrix( + X, + missing=self.missing, + feature_types=self.feature_types, + nthread=self.n_jobs, + enable_categorical=self.enable_categorical, + ) + return self.get_booster().predict( + test_dmatrix, pred_leaf=True, iteration_range=iteration_range + ) + + def evals_result(self) -> Dict[str, Dict[str, List[float]]]: + """Return the evaluation results. + + If **eval_set** is passed to the :py:meth:`fit` function, you can call + ``evals_result()`` to get evaluation results for all passed **eval_sets**. When + **eval_metric** is also passed to the :py:meth:`fit` function, the + **evals_result** will contain the **eval_metrics** passed to the :py:meth:`fit` + function. + + The returned evaluation result is a dictionary: + + .. code-block:: python + + {'validation_0': {'logloss': ['0.604835', '0.531479']}, + 'validation_1': {'logloss': ['0.41965', '0.17686']}} + + Returns + ------- + evals_result + + """ + if getattr(self, "evals_result_", None) is not None: + evals_result = self.evals_result_ + else: + raise XGBoostError( + "No evaluation result, `eval_set` is not used during training." + ) + + return evals_result + + @property + def n_features_in_(self) -> int: + """Number of features seen during :py:meth:`fit`.""" + booster = self.get_booster() + return booster.num_features() + + @property + def feature_names_in_(self) -> np.ndarray: + """Names of features seen during :py:meth:`fit`. Defined only when `X` has + feature names that are all strings. + + """ + feature_names = self.get_booster().feature_names + if feature_names is None: + raise AttributeError( + "`feature_names_in_` is defined only when `X` has feature names that " + "are all strings." + ) + return np.array(feature_names) + + @property + def best_score(self) -> float: + """The best score obtained by early stopping.""" + return self.get_booster().best_score + + @property + def best_iteration(self) -> int: + """The best iteration obtained by early stopping. This attribute is 0-based, + for instance if the best iteration is the first round, then best_iteration is 0. + + """ + return self.get_booster().best_iteration + + @property + def feature_importances_(self) -> np.ndarray: + """Feature importances property, return depends on `importance_type` + parameter. When model trained with multi-class/multi-label/multi-target dataset, + the feature importance is "averaged" over all targets. The "average" is defined + based on the importance type. For instance, if the importance type is + "total_gain", then the score is sum of loss change for each split from all + trees. + + Returns + ------- + feature_importances_ : array of shape ``[n_features]`` except for multi-class + linear model, which returns an array with shape `(n_features, n_classes)` + + """ + b: Booster = self.get_booster() + + def dft() -> str: + return "weight" if self.booster == "gblinear" else "gain" + + score = b.get_score( + importance_type=self.importance_type if self.importance_type else dft() + ) + if b.feature_names is None: + feature_names: FeatureNames = [f"f{i}" for i in range(self.n_features_in_)] + else: + feature_names = b.feature_names + # gblinear returns all features so the `get` in next line is only for gbtree. + all_features = [score.get(f, 0.0) for f in feature_names] + all_features_arr = np.array(all_features, dtype=np.float32) + total = all_features_arr.sum() + if total == 0: + return all_features_arr + return all_features_arr / total + + @property + def coef_(self) -> np.ndarray: + """ + Coefficients property + + .. note:: Coefficients are defined only for linear learners + + Coefficients are only defined when the linear model is chosen as + base learner (`booster=gblinear`). It is not defined for other base + learner types, such as tree learners (`booster=gbtree`). + + Returns + ------- + coef_ : array of shape ``[n_features]`` or ``[n_classes, n_features]`` + """ + if self.get_xgb_params()["booster"] != "gblinear": + raise AttributeError( + f"Coefficients are not defined for Booster type {self.booster}" + ) + b = self.get_booster() + coef = np.array(json.loads(b.get_dump(dump_format="json")[0])["weight"]) + # Logic for multiclass classification + n_classes = getattr(self, "n_classes_", None) + if n_classes is not None: + if n_classes > 2: + assert len(coef.shape) == 1 + assert coef.shape[0] % n_classes == 0 + coef = coef.reshape((n_classes, -1)) + return coef + + @property + def intercept_(self) -> np.ndarray: + """Intercept (bias) property + + For tree-based model, the returned value is the `base_score`. + + Returns + ------- + intercept_ : array of shape ``(1,)`` or ``[n_classes]`` + + """ + booster_config = self.get_xgb_params()["booster"] + b = self.get_booster() + if booster_config != "gblinear": # gbtree, dart + config = json.loads(b.save_config()) + intercept = json.loads( + config["learner"]["learner_model_param"]["base_score"] + ) + return np.array(intercept, dtype=np.float32) + + return np.array( + json.loads(b.get_dump(dump_format="json")[0])["bias"], dtype=np.float32 + ) + + +PredtT = TypeVar("PredtT", bound=np.ndarray) + + +def _cls_predict_proba(n_classes: int, prediction: PredtT, vstack: Callable) -> PredtT: + assert len(prediction.shape) <= 2 + if len(prediction.shape) == 2 and prediction.shape[1] == n_classes: + # multi-class + return prediction + if ( + len(prediction.shape) == 2 + and n_classes == 2 + and prediction.shape[1] >= n_classes + ): + # multi-label + return prediction + # binary logistic function + classone_probs = prediction + classzero_probs = 1.0 - classone_probs + return vstack((classzero_probs, classone_probs)).transpose() + + +@xgboost_model_doc( + "Implementation of the scikit-learn API for XGBoost classification.", + ["model", "objective"], + extra_parameters=""" + n_estimators : Optional[int] + Number of boosting rounds. +""", +) +class XGBClassifier(XGBClassifierMixIn, XGBModel): + # pylint: disable=missing-docstring,too-many-instance-attributes + @_deprecate_positional_args + def __init__( + self, + *, + objective: SklObjective = "binary:logistic", + **kwargs: Any, + ) -> None: + super().__init__(objective=objective, **kwargs) + + @_deprecate_positional_args + def fit( + self, + X: ArrayLike, + y: ArrayLike, + *, + sample_weight: Optional[ArrayLike] = None, + base_margin: Optional[ArrayLike] = None, + eval_set: Optional[Sequence[Tuple[ArrayLike, ArrayLike]]] = None, + verbose: Optional[Union[bool, int]] = True, + xgb_model: Optional[Union[Booster, str, XGBModel]] = None, + sample_weight_eval_set: Optional[Sequence[ArrayLike]] = None, + base_margin_eval_set: Optional[Sequence[ArrayLike]] = None, + feature_weights: Optional[ArrayLike] = None, + ) -> "XGBClassifier": + # pylint: disable = attribute-defined-outside-init,too-many-statements + with config_context(verbosity=self.verbosity): + # We keep the n_classes_ as a simple member instead of loading it from + # booster in a Python property. This way we can have efficient and + # thread-safe prediction. + if _is_polars_lazyframe(y): + y = y.collect() + if _is_cudf_df(y) or _is_cudf_ser(y): + cp = import_cupy() + + classes = cp.unique(y.values) + self.n_classes_ = len(classes) + expected_classes = cp.array(self.classes_) + elif _is_cupy_alike(y): + cp = import_cupy() + + classes = cp.unique(y) + self.n_classes_ = len(classes) + expected_classes = cp.array(self.classes_) + else: + classes = np.unique(np.asarray(y)) + self.n_classes_ = len(classes) + expected_classes = self.classes_ + if ( + classes.shape != expected_classes.shape + or not (classes == expected_classes).all() + ): + raise ValueError( + f"Invalid classes inferred from unique values of `y`. " + f"Expected: {expected_classes}, got {classes}" + ) + + params = self.get_xgb_params() + + if callable(self.objective): + obj: Optional[CustomObj] = _objective_decorator(self.objective) + # Use default value. Is it really not used ? + params["objective"] = "binary:logistic" + else: + obj = None + + if self.n_classes_ > 2: + # Switch to using a multiclass objective in the underlying XGB instance + if params.get("objective", None) != "multi:softmax": + params["objective"] = "multi:softprob" + params["num_class"] = self.n_classes_ + + model, metric, params, feature_weights = self._configure_fit( + xgb_model, params, feature_weights + ) + model, feature_types = get_model_categories(X, model, self.feature_types) + + evals_result: EvalsLog = {} + train_dmatrix, evals = _wrap_evaluation_matrices( + missing=self.missing, + X=X, + y=y, + group=None, + qid=None, + sample_weight=sample_weight, + base_margin=base_margin, + feature_weights=feature_weights, + eval_set=eval_set, + sample_weight_eval_set=sample_weight_eval_set, + base_margin_eval_set=base_margin_eval_set, + eval_group=None, + eval_qid=None, + create_dmatrix=self._create_dmatrix, + enable_categorical=self.enable_categorical, + feature_types=feature_types, + ) + + self._Booster = train( + params, + train_dmatrix, + self.get_num_boosting_rounds(), + evals=evals, + early_stopping_rounds=self.early_stopping_rounds, + evals_result=evals_result, + obj=obj, + custom_metric=metric, + verbose_eval=verbose, + xgb_model=model, + callbacks=self.callbacks, + ) + + if not callable(self.objective): + self.objective = params["objective"] + + self._set_evaluation_result(evals_result) + return self + + if XGBModel.fit.__doc__ is not None: + fit.__doc__ = XGBModel.fit.__doc__.replace( + "Fit gradient boosting model", "Fit gradient boosting classifier", 1 + ) + + @_deprecate_positional_args + def predict( + self, + X: ArrayLike, + *, + output_margin: bool = False, + validate_features: bool = True, + base_margin: Optional[ArrayLike] = None, + iteration_range: Optional[IterationRange] = None, + ) -> ArrayLike: + with config_context(verbosity=self.verbosity): + class_probs = super().predict( + X=X, + output_margin=output_margin, + validate_features=validate_features, + base_margin=base_margin, + iteration_range=iteration_range, + ) + if output_margin: + # If output_margin is active, simply return the scores + return class_probs + + if len(class_probs.shape) > 1 and self.n_classes_ != 2: + # multi-class, turns softprob into softmax + column_indexes: np.ndarray = np.argmax(class_probs, axis=1) + elif len(class_probs.shape) > 1 and class_probs.shape[1] != 1: + # multi-label + column_indexes = np.zeros(class_probs.shape) + column_indexes[class_probs > 0.5] = 1 + elif self.objective == "multi:softmax": + return class_probs.astype(np.int32) + else: + # turns soft logit into class label + column_indexes = np.repeat(0, class_probs.shape[0]) + column_indexes[class_probs > 0.5] = 1 + + return column_indexes + + def predict_proba( + self, + X: ArrayLike, + validate_features: bool = True, + base_margin: Optional[ArrayLike] = None, + iteration_range: Optional[IterationRange] = None, + ) -> np.ndarray: + """Predict the probability of each `X` example being of a given class. If the + model is trained with early stopping, then :py:attr:`best_iteration` is used + automatically. The estimator uses `inplace_predict` by default and falls back to + using :py:class:`DMatrix` if devices between the data and the estimator don't + match. + + .. note:: This function is only thread safe for `gbtree` and `dart`. + + Parameters + ---------- + X : + Feature matrix. See :ref:`py-data` for a list of supported types. + validate_features : + When this is True, validate that the Booster's and data's feature_names are + identical. Otherwise, it is assumed that the feature_names are the same. + base_margin : + Global bias for each instance. See :doc:`/tutorials/intercept` for details. + iteration_range : + Specifies which layer of trees are used in prediction. For example, if a + random forest is trained with 100 rounds. Specifying `iteration_range=(10, + 20)`, then only the forests built during [10, 20) (half open set) rounds are + used in this prediction. + + Returns + ------- + prediction : + a numpy array of shape array-like of shape (n_samples, n_classes) with the + probability of each data example being of a given class. + + """ + # custom obj: Do nothing as we don't know what to do. + # softprob: Do nothing, output is proba. + # softmax: Use softmax from scipy + # binary:logistic: Expand the prob vector into 2-class matrix after predict. + # binary:logitraw: Unsupported by predict_proba() + if self.objective == "multi:softmax": + raw_predt = super().predict( + X=X, + validate_features=validate_features, + base_margin=base_margin, + iteration_range=iteration_range, + output_margin=True, + ) + class_prob = softmax(raw_predt, axis=1) + return class_prob + class_probs = super().predict( + X=X, + validate_features=validate_features, + base_margin=base_margin, + iteration_range=iteration_range, + ) + return _cls_predict_proba(self.n_classes_, class_probs, np.vstack) + + +@xgboost_model_doc( + """scikit-learn API for XGBoost random forest classification. + + .. deprecated:: 3.4.0 + + This estimator is deprecated and will be removed in a future release. See + :doc:`/tutorials/rf` for alternatives. + + """, + ["model", "objective"], + extra_parameters=""" + n_estimators : Optional[int] + Number of trees in random forest to fit. +""", +) +class XGBRFClassifier(XGBClassifier): + # pylint: disable=missing-docstring + @_deprecate_positional_args + def __init__( + self, + *, + learning_rate: float = 1.0, + subsample: float = 0.8, + colsample_bynode: float = 0.8, + reg_lambda: float = 1e-5, + **kwargs: Any, + ): + _warn_rf_deprecated("XGBRFClassifier") + super().__init__( + learning_rate=learning_rate, + subsample=subsample, + colsample_bynode=colsample_bynode, + reg_lambda=reg_lambda, + **kwargs, + ) + _check_rf_callback(self.early_stopping_rounds, self.callbacks) + + def get_xgb_params(self) -> Dict[str, Any]: + params = super().get_xgb_params() + params["num_parallel_tree"] = super().get_num_boosting_rounds() + return params + + def get_num_boosting_rounds(self) -> int: + return 1 + + # pylint: disable=unused-argument + @_deprecate_positional_args + def fit( + self, + X: ArrayLike, + y: ArrayLike, + *, + sample_weight: Optional[ArrayLike] = None, + base_margin: Optional[ArrayLike] = None, + eval_set: Optional[Sequence[Tuple[ArrayLike, ArrayLike]]] = None, + verbose: Optional[Union[bool, int]] = True, + xgb_model: Optional[Union[Booster, str, XGBModel]] = None, + sample_weight_eval_set: Optional[Sequence[ArrayLike]] = None, + base_margin_eval_set: Optional[Sequence[ArrayLike]] = None, + feature_weights: Optional[ArrayLike] = None, + ) -> "XGBRFClassifier": + args = {k: v for k, v in locals().items() if k not in ("self", "__class__")} + _check_rf_callback(self.early_stopping_rounds, self.callbacks) + super().fit(**args) + return self + + +@xgboost_model_doc( + "Implementation of the scikit-learn API for XGBoost regression.", + ["estimators", "model", "objective"], +) +class XGBRegressor(XGBRegressorMixIn, XGBModel): + # pylint: disable=missing-docstring + @_deprecate_positional_args + def __init__( + self, *, objective: SklObjective = "reg:squarederror", **kwargs: Any + ) -> None: + super().__init__(objective=objective, **kwargs) + + +@xgboost_model_doc( + """scikit-learn API for XGBoost random forest regression. + + .. deprecated:: 3.4.0 + + This estimator is deprecated and will be removed in a future release. See + :doc:`/tutorials/rf` for alternatives. + + """, + ["model", "objective"], + extra_parameters=""" + n_estimators : Optional[int] + Number of trees in random forest to fit. +""", +) +class XGBRFRegressor(XGBRegressor): + # pylint: disable=missing-docstring + @_deprecate_positional_args + def __init__( + self, + *, + learning_rate: float = 1.0, + subsample: float = 0.8, + colsample_bynode: float = 0.8, + reg_lambda: float = 1e-5, + **kwargs: Any, + ) -> None: + _warn_rf_deprecated("XGBRFRegressor") + super().__init__( + learning_rate=learning_rate, + subsample=subsample, + colsample_bynode=colsample_bynode, + reg_lambda=reg_lambda, + **kwargs, + ) + _check_rf_callback(self.early_stopping_rounds, self.callbacks) + + def get_xgb_params(self) -> Dict[str, Any]: + params = super().get_xgb_params() + params["num_parallel_tree"] = super().get_num_boosting_rounds() + return params + + def get_num_boosting_rounds(self) -> int: + return 1 + + # pylint: disable=unused-argument + @_deprecate_positional_args + def fit( + self, + X: ArrayLike, + y: ArrayLike, + *, + sample_weight: Optional[ArrayLike] = None, + base_margin: Optional[ArrayLike] = None, + eval_set: Optional[Sequence[Tuple[ArrayLike, ArrayLike]]] = None, + verbose: Optional[Union[bool, int]] = True, + xgb_model: Optional[Union[Booster, str, XGBModel]] = None, + sample_weight_eval_set: Optional[Sequence[ArrayLike]] = None, + base_margin_eval_set: Optional[Sequence[ArrayLike]] = None, + feature_weights: Optional[ArrayLike] = None, + ) -> "XGBRFRegressor": + args = {k: v for k, v in locals().items() if k not in ("self", "__class__")} + _check_rf_callback(self.early_stopping_rounds, self.callbacks) + super().fit(**args) + return self + + +def _get_qid( + X: ArrayLike, qid: Optional[ArrayLike] +) -> Tuple[ArrayLike, Optional[ArrayLike]]: + """Get the special qid column from X if exists.""" + if (_is_pandas_df(X) or _is_cudf_df(X)) and hasattr(X, "qid"): + if qid is not None: + raise ValueError( + "Found both the special column `qid` in `X` and the `qid` from the" + "`fit` method. Please remove one of them." + ) + q_x = X.qid + X = X.drop("qid", axis=1) + return X, q_x + return X, qid + + +@xgboost_model_doc( + """Implementation of the Scikit-Learn API for XGBoost Ranking. + +See :doc:`Learning to Rank ` for an introducion. + + """, + ["estimators", "model"], + end_note=""" + .. note:: + + A custom objective function is currently not supported by XGBRanker. + + .. note:: + + Query group information is only required for ranking training but not + prediction. Multiple groups can be predicted on a single call to + :py:meth:`predict`. + + When fitting the model with the `group` parameter, your data need to be sorted + by the query group first. `group` is an array that contains the size of each + query group. + + Similarly, when fitting the model with the `qid` parameter, the data should be + sorted according to query index and `qid` is an array that contains the query + index for each training sample. + + For example, if your original data look like: + + +-------+-----------+---------------+ + | qid | label | features | + +-------+-----------+---------------+ + | 1 | 0 | x_1 | + +-------+-----------+---------------+ + | 1 | 1 | x_2 | + +-------+-----------+---------------+ + | 1 | 0 | x_3 | + +-------+-----------+---------------+ + | 2 | 0 | x_4 | + +-------+-----------+---------------+ + | 2 | 1 | x_5 | + +-------+-----------+---------------+ + | 2 | 1 | x_6 | + +-------+-----------+---------------+ + | 2 | 1 | x_7 | + +-------+-----------+---------------+ + + then :py:meth:`fit` method can be called with either `group` array as ``[3, 4]`` + or with `qid` as ``[1, 1, 1, 2, 2, 2, 2]``, that is the qid column. Also, the + `qid` can be a special column of input `X` instead of a separated parameter, see + :py:meth:`fit` for more info.""", +) +class XGBRanker(XGBRankerMixIn, XGBModel): + # pylint: disable=missing-docstring,too-many-arguments + @_deprecate_positional_args + def __init__(self, *, objective: str = "rank:ndcg", **kwargs: Any): + super().__init__(objective=objective, **kwargs) + if callable(self.objective): + raise ValueError("custom objective function not supported by XGBRanker") + if "rank:" not in objective: + raise ValueError("please use XGBRanker for ranking task") + + def _create_ltr_dmatrix( + self, ref: Optional[DMatrix], data: ArrayLike, qid: ArrayLike, **kwargs: Any + ) -> DMatrix: + data, qid = _get_qid(data, qid) + + if kwargs.get("group", None) is None and qid is None: + raise ValueError("Either `group` or `qid` is required for ranking task") + + return super()._create_dmatrix(ref=ref, data=data, qid=qid, **kwargs) + + @_deprecate_positional_args + def fit( + self, + X: ArrayLike, + y: ArrayLike, + *, + group: Optional[ArrayLike] = None, + qid: Optional[ArrayLike] = None, + sample_weight: Optional[ArrayLike] = None, + base_margin: Optional[ArrayLike] = None, + eval_set: Optional[Sequence[Tuple[ArrayLike, ArrayLike]]] = None, + eval_group: Optional[Sequence[ArrayLike]] = None, + eval_qid: Optional[Sequence[ArrayLike]] = None, + verbose: Optional[Union[bool, int]] = True, + xgb_model: Optional[Union[Booster, str, XGBModel]] = None, + sample_weight_eval_set: Optional[Sequence[ArrayLike]] = None, + base_margin_eval_set: Optional[Sequence[ArrayLike]] = None, + feature_weights: Optional[ArrayLike] = None, + ) -> "XGBRanker": + # pylint: disable = attribute-defined-outside-init,arguments-differ + """Fit gradient boosting ranker + + Note that calling ``fit()`` multiple times will cause the model object to be + re-fit from scratch. To resume training from a previous checkpoint, explicitly + pass ``xgb_model`` argument. + + Parameters + ---------- + X : + Feature matrix. See :ref:`py-data` for a list of supported types. + + When this is a :py:class:`pandas.DataFrame` or a :py:class:`cudf.DataFrame`, + it may contain a special column called ``qid`` for specifying the query + index. Using a special column is the same as using the `qid` parameter, + except for being compatible with sklearn utility functions like + :py:func:`sklearn.model_selection.cross_validation`. The same convention + applies to the :py:meth:`XGBRanker.score` and :py:meth:`XGBRanker.predict`. + + +-----+----------------+----------------+ + | qid | feat_0 | feat_1 | + +-----+----------------+----------------+ + | 0 | :math:`x_{00}` | :math:`x_{01}` | + +-----+----------------+----------------+ + | 1 | :math:`x_{10}` | :math:`x_{11}` | + +-----+----------------+----------------+ + | 1 | :math:`x_{20}` | :math:`x_{21}` | + +-----+----------------+----------------+ + + When the ``tree_method`` is set to ``hist``, internally, the + :py:class:`QuantileDMatrix` will be used instead of the :py:class:`DMatrix` + for conserving memory. However, this has performance implications when the + device of input data is not matched with algorithm. For instance, if the + input is a numpy array on CPU but ``cuda`` is used for training, then the + data is first processed on CPU then transferred to GPU. + y : + Labels + group : + Size of each query group of training data. Should have as many elements as + the query groups in the training data. If this is set to None, then user + must provide qid. + qid : + Query ID for each training sample. Should have the size of n_samples. If + this is set to None, then user must provide group or a special column in X. + sample_weight : + Query group weights + + .. note:: Weights are per-group for ranking tasks + + In ranking task, one weight is assigned to each query group/id (not each + data point). This is because we only care about the relative ordering of + data points within each group, so it doesn't make sense to assign + weights to individual data points. + + base_margin : + Global bias for each instance. See :doc:`/tutorials/intercept` for details. + eval_set : + A list of (X, y) tuple pairs to use as validation sets, for which + metrics will be computed. + Validation metrics will help us track the performance of the model. + eval_group : + A list in which ``eval_group[i]`` is the list containing the sizes of all + query groups in the ``i``-th pair in **eval_set**. + eval_qid : + A list in which ``eval_qid[i]`` is the array containing query ID of ``i``-th + pair in **eval_set**. The special column convention in `X` applies to + validation datasets as well. + + verbose : + If `verbose` is True and an evaluation set is used, the evaluation metric + measured on the validation set is printed to stdout at each boosting stage. + If `verbose` is an integer, the evaluation metric is printed at each + `verbose` boosting stage. The last boosting stage / the boosting stage found + by using `early_stopping_rounds` is also printed. + xgb_model : + file name of stored XGBoost model or 'Booster' instance XGBoost model to be + loaded before training (allows training continuation). + sample_weight_eval_set : + A list of the form [L_1, L_2, ..., L_n], where each L_i is a list of + group weights on the i-th validation set. + + .. note:: Weights are per-group for ranking tasks + + In ranking task, one weight is assigned to each query group (not each + data point). This is because we only care about the relative ordering of + data points within each group, so it doesn't make sense to assign + weights to individual data points. + base_margin_eval_set : + A list of the form [M_1, M_2, ..., M_n], where each M_i is an array like + object storing base margin for the i-th validation set. + feature_weights : + Weight for each feature, defines the probability of each feature being + selected when colsample is being used. All values must be greater than 0, + otherwise a `ValueError` is thrown. + + """ + with config_context(verbosity=self.verbosity): + params = self.get_xgb_params() + + model, metric, params, feature_weights = self._configure_fit( + xgb_model, params, feature_weights + ) + model, feature_types = get_model_categories(X, model, self.feature_types) + + evals_result: EvalsLog = {} + train_dmatrix, evals = _wrap_evaluation_matrices( + missing=self.missing, + X=X, + y=y, + group=group, + qid=qid, + sample_weight=sample_weight, + base_margin=base_margin, + feature_weights=feature_weights, + eval_set=eval_set, + sample_weight_eval_set=sample_weight_eval_set, + base_margin_eval_set=base_margin_eval_set, + eval_group=eval_group, + eval_qid=eval_qid, + create_dmatrix=self._create_ltr_dmatrix, + enable_categorical=self.enable_categorical, + feature_types=feature_types, + ) + + self._Booster = train( + params, + train_dmatrix, + num_boost_round=self.get_num_boosting_rounds(), + early_stopping_rounds=self.early_stopping_rounds, + evals=evals, + evals_result=evals_result, + custom_metric=metric, + verbose_eval=verbose, + xgb_model=model, + callbacks=self.callbacks, + ) + + self.objective = params["objective"] + + self._set_evaluation_result(evals_result) + return self + + @_deprecate_positional_args + def predict( + self, + X: ArrayLike, + *, + output_margin: bool = False, + validate_features: bool = True, + base_margin: Optional[ArrayLike] = None, + iteration_range: Optional[IterationRange] = None, + ) -> ArrayLike: + X, _ = _get_qid(X, None) + return super().predict( + X, + output_margin=output_margin, + validate_features=validate_features, + base_margin=base_margin, + iteration_range=iteration_range, + ) + + def apply( + self, + X: ArrayLike, + iteration_range: Optional[IterationRange] = None, + ) -> ArrayLike: + X, _ = _get_qid(X, None) + return super().apply(X, iteration_range) + + def score(self, X: ArrayLike, y: ArrayLike) -> float: + """Evaluate score for data using the last evaluation metric. If the model is + trained with early stopping, then :py:attr:`best_iteration` is used + automatically. + + Parameters + ---------- + X : Union[pd.DataFrame, cudf.DataFrame] + Feature matrix. A DataFrame with a special `qid` column. + + y : + Labels + + Returns + ------- + score : + The result of the first evaluation metric for the ranker. + + """ + X, qid = _get_qid(X, None) + if qid is None: + raise ValueError( + "The special column `qid` is required in `X` for ranking task." + ) + # fixme(jiamingy): base margin and group weight is not yet supported. We might + # need to make extra special fields in the dataframe. + Xyq = DMatrix( + X, + y, + qid=qid, + missing=self.missing, + enable_categorical=self.enable_categorical, + nthread=self.n_jobs, + feature_types=self.feature_types, + ) + if callable(self.eval_metric): + metric = ltr_metric_decorator(self.eval_metric, self.n_jobs) + result_str = self.get_booster().eval_set([(Xyq, "eval")], feval=metric) + else: + result_str = self.get_booster().eval(Xyq) + + metric_score = _parse_eval_str(result_str) + return metric_score[-1][1] diff --git a/v12x/pylib/xgboost/spark/__init__.py b/v12x/pylib/xgboost/spark/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..224f87d0399a5e23b4b32da1fa4d2ece75de7689 --- /dev/null +++ b/v12x/pylib/xgboost/spark/__init__.py @@ -0,0 +1,24 @@ +"""PySpark XGBoost integration interface""" + +try: + import pyspark +except ImportError as e: + raise ImportError("pyspark package needs to be installed to use this module") from e + +from .estimator import ( + SparkXGBClassifier, + SparkXGBClassifierModel, + SparkXGBRanker, + SparkXGBRankerModel, + SparkXGBRegressor, + SparkXGBRegressorModel, +) + +__all__ = [ + "SparkXGBClassifier", + "SparkXGBClassifierModel", + "SparkXGBRegressor", + "SparkXGBRegressorModel", + "SparkXGBRanker", + "SparkXGBRankerModel", +] diff --git a/v12x/pylib/xgboost/spark/core.py b/v12x/pylib/xgboost/spark/core.py new file mode 100644 index 0000000000000000000000000000000000000000..dc6e7c06c2773bf39890c44d5ddcc8978c8ad063 --- /dev/null +++ b/v12x/pylib/xgboost/spark/core.py @@ -0,0 +1,1787 @@ +"""XGBoost pyspark integration submodule for core code.""" + +import base64 + +# pylint: disable=fixme, protected-access, no-member +# pylint: disable=too-many-lines, too-many-branches +import json +import logging +import os +from collections import namedtuple +from dataclasses import asdict +from typing import ( + Any, + Callable, + Dict, + Iterator, + List, + Optional, + Tuple, + Type, + Union, + cast, +) + +import numpy as np +import pandas as pd +from pyspark import cloudpickle +from pyspark.ml import Estimator, Model +from pyspark.ml.functions import array_to_vector, vector_to_array +from pyspark.ml.linalg import VectorUDT +from pyspark.ml.param import Param, Params, TypeConverters +from pyspark.ml.param.shared import ( + HasFeaturesCol, + HasLabelCol, + HasPredictionCol, + HasProbabilityCol, + HasRawPredictionCol, + HasValidationIndicatorCol, + HasWeightCol, +) +from pyspark.ml.util import ( + DefaultParamsReader, + DefaultParamsWriter, + MLReadable, + MLReader, + MLWritable, + MLWriter, +) +from pyspark.resource import ( + ResourceProfile, + ResourceProfileBuilder, + TaskResourceRequests, +) +from pyspark.sql import Column, DataFrame, SparkSession +from pyspark.sql.functions import ( + col, + countDistinct, + pandas_udf, + rand, + struct, + unwrap_udt, +) +from pyspark.sql.types import ( + ArrayType, + BooleanType, + DoubleType, + FloatType, + IntegerType, + IntegralType, + LongType, + ShortType, +) +from scipy.special import expit, softmax # pylint: disable=no-name-in-module + +from .._c_api import _py_version +from .._typing import ArrayLike +from ..collective import Config +from ..compat import import_cupy, is_cudf_available, is_cupy_available +from ..config import config_context, get_config +from ..core import Booster, _check_distributed_params +from ..sklearn import DEFAULT_N_ESTIMATORS, XGBClassifier, XGBModel, _can_use_qdm +from ..training import train as worker_train +from .data import ( + _read_csr_matrix_from_unwrapped_spark_vec, + alias, + create_dmatrix_from_partitions, + pred_contribs, + stack_series, +) +from .params import ( + HasArbitraryParamsDict, + HasBaseMarginCol, + HasContribPredictionCol, + HasEnableSparseDataOptim, + HasFeaturesCols, + HasQueryIdCol, +) +from .summary import XGBoostTrainingSummary +from .utils import ( + CommunicatorContext, + _get_default_params_from_func, + _get_gpu_id, + _get_host_ip, + _get_max_num_concurrent_tasks, + _get_rabit_args, + _is_connect, + _is_local, + deserialize_booster, + deserialize_xgb_model, + get_class_name, + get_logger, + get_logger_level, + serialize_booster, + use_cuda, +) + +# Put pyspark specific params here, they won't be passed to XGBoost. +# like `validationIndicatorCol`, `base_margin_col` +_pyspark_specific_params = [ + "featuresCol", + "labelCol", + "weightCol", + "rawPredictionCol", + "predictionCol", + "probabilityCol", + "validationIndicatorCol", + "base_margin_col", + "arbitrary_params_dict", + "force_repartition", + "num_workers", + "feature_names", + "features_cols", + "enable_sparse_data_optim", + "qid_col", + "repartition_random_shuffle", + "pred_contrib_col", + "launch_tracker_on_driver", + "coll_cfg", +] + +_non_booster_params = ["missing", "n_estimators", "feature_types", "feature_weights"] + +_pyspark_param_alias_map = { + "features_col": "featuresCol", + "label_col": "labelCol", + "weight_col": "weightCol", + "raw_prediction_col": "rawPredictionCol", + "prediction_col": "predictionCol", + "probability_col": "probabilityCol", + "validation_indicator_col": "validationIndicatorCol", +} + +_inverse_pyspark_param_alias_map = {v: k for k, v in _pyspark_param_alias_map.items()} + +_unsupported_xgb_params = [ + "enable_categorical", # Use feature_types param to specify categorical feature instead + "n_jobs", # Do not allow user to set it, will use `spark.task.cpus` value instead. + "nthread", # Ditto +] + +_unsupported_fit_params = { + "sample_weight", # Supported by spark param weightCol + "eval_set", # Supported by spark param validation_indicator_col + "sample_weight_eval_set", # Supported by spark param weight_col + validation_indicator_col + "base_margin", # Supported by spark param base_margin_col + "base_margin_eval_set", # Supported by spark param base_margin_col + validation_indicator_col + "group", # Use spark param `qid_col` instead + "qid", # Use spark param `qid_col` instead + "eval_group", # Use spark param `qid_col` instead + "eval_qid", # Use spark param `qid_col` instead +} + +_unsupported_train_params = { + "evals", # Supported by spark param validation_indicator_col + "evals_result", # Won't support yet+ +} + +_unsupported_predict_params = { + # for classification, we can use rawPrediction as margin + "output_margin", + "validate_features", # TODO + "base_margin", # Use pyspark base_margin_col param instead. +} + +# TODO: supply hint message for all other unsupported params. +_unsupported_params_hint_message = { + "enable_categorical": "`xgboost.spark` estimators do not have 'enable_categorical' param, " + "but you can set `feature_types` param and mark categorical features with 'c' string." +} + +# Global prediction names +Pred = namedtuple( + "Pred", ("prediction", "raw_prediction", "probability", "pred_contrib") +) +pred = Pred("prediction", "rawPrediction", "probability", "predContrib") + +_INIT_BOOSTER_SAVE_PATH = "init_booster.json" + +_LOG_TAG = "XGBoost-PySpark" + + +class _SparkXGBParams( + HasFeaturesCol, + HasLabelCol, + HasWeightCol, + HasPredictionCol, + HasValidationIndicatorCol, + HasArbitraryParamsDict, + HasBaseMarginCol, + HasFeaturesCols, + HasEnableSparseDataOptim, + HasQueryIdCol, + HasContribPredictionCol, +): + num_workers = Param( + Params._dummy(), + "num_workers", + "The number of XGBoost workers. Each XGBoost worker corresponds to one spark task.", + TypeConverters.toInt, + ) + device = Param( + Params._dummy(), + "device", + ( + "The device type for XGBoost executors. Available options are `cpu`,`cuda`" + " and `gpu`. Set `device` to `cuda` or `gpu` if the executors are running " + "on GPU instances. Currently, only one GPU per task is supported." + ), + TypeConverters.toString, + ) + force_repartition = Param( + Params._dummy(), + "force_repartition", + "A boolean variable. Set force_repartition=true if you " + + "want to force the input dataset to be repartitioned before XGBoost training." + + "Note: The auto repartitioning judgement is not fully accurate, so it is recommended" + + "to have force_repartition be True.", + TypeConverters.toBoolean, + ) + repartition_random_shuffle = Param( + Params._dummy(), + "repartition_random_shuffle", + "A boolean variable. Set repartition_random_shuffle=true if you want to random shuffle " + "dataset when repartitioning is required. By default is True.", + TypeConverters.toBoolean, + ) + feature_names = Param( + Params._dummy(), + "feature_names", + "A list of str to specify feature names.", + TypeConverters.toList, + ) + launch_tracker_on_driver = Param( + Params._dummy(), + "launch_tracker_on_driver", + "A boolean variable. Set launch_tracker_on_driver to true if you want the tracker to be " + "launched on the driver side; otherwise, it will be launched on the executor side.", + TypeConverters.toBoolean, + ) + coll_cfg = Param( + Params._dummy(), + "coll_cfg", + "xgboost.collective.Config. The collective configuration.", + TypeConverters.identity, + ) + + def set_coll_cfg(self, value: Config) -> "_SparkXGBParams": + """Set collective configuration""" + assert isinstance(value, Config) + self.set(self.coll_cfg, value) + return self + + def set_device(self, value: str) -> "_SparkXGBParams": + """Set device, optional value: cpu, cuda, gpu""" + _check_distributed_params({"device": value}) + assert value in ("cpu", "cuda", "gpu") + self.set(self.device, value) + return self + + @classmethod + def _xgb_cls(cls) -> Type[XGBModel]: + """ + Subclasses should override this method and + returns an xgboost.XGBModel subclass + """ + raise NotImplementedError() + + # Parameters for xgboost.XGBModel() + @classmethod + def _get_xgb_params_default(cls) -> Dict[str, Any]: + """Get the xgboost.sklearn.XGBModel default parameters and filter out some""" + xgb_model_default = cls._xgb_cls()() + params_dict = xgb_model_default.get_params() + filtered_params_dict = { + k: params_dict[k] for k in params_dict if k not in _unsupported_xgb_params + } + filtered_params_dict["n_estimators"] = DEFAULT_N_ESTIMATORS + return filtered_params_dict + + def _set_xgb_params_default(self) -> None: + """Set xgboost parameters into spark parameters""" + filtered_params_dict = self._get_xgb_params_default() + self._setDefault(**filtered_params_dict) + + def _gen_xgb_params_dict( + self, gen_xgb_sklearn_estimator_param: bool = False + ) -> Dict[str, Any]: + """Generate the xgboost parameters which will be passed into xgboost library""" + xgb_params = {} + non_xgb_params = ( + set(_pyspark_specific_params) + | self._get_fit_params_default().keys() + | self._get_predict_params_default().keys() + ) + if not gen_xgb_sklearn_estimator_param: + non_xgb_params |= set(_non_booster_params) + for param in self.extractParamMap(): + if param.name not in non_xgb_params: + xgb_params[param.name] = self.getOrDefault(param) + + arbitrary_params_dict = self.getOrDefault( + self.getParam("arbitrary_params_dict") + ) + xgb_params.update(arbitrary_params_dict) + return xgb_params + + # Parameters for xgboost.XGBModel().fit() + @classmethod + def _get_fit_params_default(cls) -> Dict[str, Any]: + """Get the xgboost.XGBModel().fit() parameters""" + fit_params = _get_default_params_from_func( + cls._xgb_cls().fit, _unsupported_fit_params + ) + return fit_params + + def _set_fit_params_default(self) -> None: + """Get the xgboost.XGBModel().fit() parameters and set them to spark parameters""" + filtered_params_dict = self._get_fit_params_default() + self._setDefault(**filtered_params_dict) + + def _gen_fit_params_dict(self) -> Dict[str, Any]: + """Generate the fit parameters which will be passed into fit function""" + fit_params_keys = self._get_fit_params_default().keys() + fit_params = {} + for param in self.extractParamMap(): + if param.name in fit_params_keys: + fit_params[param.name] = self.getOrDefault(param) + return fit_params + + @classmethod + def _get_predict_params_default(cls) -> Dict[str, Any]: + """Get the parameters from xgboost.XGBModel().predict()""" + predict_params = _get_default_params_from_func( + cls._xgb_cls().predict, _unsupported_predict_params + ) + return predict_params + + def _set_predict_params_default(self) -> None: + """Get the parameters from xgboost.XGBModel().predict() and + set them into spark parameters""" + filtered_params_dict = self._get_predict_params_default() + self._setDefault(**filtered_params_dict) + + def _gen_predict_params_dict(self) -> Dict[str, Any]: + """Generate predict parameters which will be passed into xgboost.XGBModel().predict()""" + predict_params_keys = self._get_predict_params_default().keys() + predict_params = {} + for param in self.extractParamMap(): + if param.name in predict_params_keys: + predict_params[param.name] = self.getOrDefault(param) + return predict_params + + def _validate_gpu_params(self, spark_session: SparkSession) -> None: + """Validate the gpu parameters and gpu configurations""" + + if self._run_on_gpu(spark_session): + if _is_local(spark_session): + # Supporting GPU training in Spark local mode is just for debugging + # purposes, so it's okay for printing the below warning instead of + # checking the real gpu numbers and raising the exception. + get_logger(self.__class__.__name__).warning( + "You have enabled GPU in spark local mode. Please make sure your" + " local node has at least %d GPUs", + self.getOrDefault(self.num_workers), + ) + else: + executor_gpus = spark_session.conf.get( + "spark.executor.resource.gpu.amount", None + ) + if executor_gpus is None: + raise ValueError( + "The `spark.executor.resource.gpu.amount` is required for training" + " on GPU." + ) + gpu_per_task = spark_session.conf.get( + "spark.task.resource.gpu.amount", None + ) + if gpu_per_task is not None and float(gpu_per_task) > 1.0: + get_logger(self.__class__.__name__).warning( + "The configuration assigns %s GPUs to each Spark task, but each " + "XGBoost training task only utilizes 1 GPU, which will lead to " + "unnecessary GPU waste", + gpu_per_task, + ) + + def _validate_params(self, spark_session: SparkSession) -> None: + # pylint: disable=too-many-branches + init_model = self.getOrDefault("xgb_model") + if init_model is not None and not isinstance(init_model, Booster): + raise ValueError( + "The xgb_model param must be set with a `xgboost.core.Booster` " + "instance." + ) + + if self.getOrDefault(self.num_workers) < 1: + raise ValueError( + f"Number of workers was {self.getOrDefault(self.num_workers)}." + f"It cannot be less than 1 [Default is 1]" + ) + + tree_method = self.getOrDefault(self.getParam("tree_method")) + if tree_method == "exact": + raise ValueError( + "The `exact` tree method is not supported for distributed systems." + ) + + if self.getOrDefault("objective") is not None: + if not isinstance(self.getOrDefault("objective"), str): + raise ValueError("Only string type 'objective' param is allowed.") + + eval_metric = "eval_metric" + if self.getOrDefault(eval_metric) is not None: + if not ( + isinstance(self.getOrDefault(eval_metric), str) + or ( + isinstance(self.getOrDefault(eval_metric), List) + and all( + isinstance(metric, str) + for metric in self.getOrDefault(eval_metric) + ) + ) + ): + raise ValueError( + "Only string type or list of string type 'eval_metric' param is allowed." + ) + + if self.getOrDefault("early_stopping_rounds") is not None: + if not self._col_is_defined_not_empty(self.validationIndicatorCol): + raise ValueError( + "If 'early_stopping_rounds' param is set, you need to set " + "'validation_indicator_col' param as well." + ) + + if self.getOrDefault(self.enable_sparse_data_optim): + if self.getOrDefault("missing") != 0.0: + # If DMatrix is constructed from csr / csc matrix, then inactive elements + # in csr / csc matrix are regarded as missing value, but, in pyspark, we + # are hard to control elements to be active or inactive in sparse vector column, + # some spark transformers such as VectorAssembler might compress vectors + # to be dense or sparse format automatically, and when a spark ML vector object + # is compressed to sparse vector, then all zero value elements become inactive. + # So we force setting missing param to be 0 when enable_sparse_data_optim config + # is True. + raise ValueError( + "If enable_sparse_data_optim is True, missing param != 0 is not supported." + ) + if self.getOrDefault(self.features_cols): + raise ValueError( + "If enable_sparse_data_optim is True, you cannot set multiple feature columns " + "but you should set one feature column with values of " + "`pyspark.ml.linalg.Vector` type." + ) + + self._validate_gpu_params(spark_session) + + def _run_on_gpu(self, spark_session: SparkSession) -> bool: + # pylint: disable=unused-argument + """If train or transform on the gpu according to the parameters""" + + return use_cuda(self.getOrDefault(self.device)) + + def _col_is_defined_not_empty(self, param: "Param[str]") -> bool: + return self.isDefined(param) and self.getOrDefault(param) not in (None, "") + + +def _validate_and_convert_feature_col_as_float_col_list( + dataset: DataFrame, features_col_names: List[str] +) -> List[Column]: + """Values in feature columns must be integral types or float/double types""" + feature_cols = [] + for c in features_col_names: + if isinstance(dataset.schema[c].dataType, DoubleType): + feature_cols.append(col(c).cast(FloatType()).alias(c)) + elif isinstance(dataset.schema[c].dataType, (FloatType, IntegralType)): + feature_cols.append(col(c)) + else: + raise ValueError( + "Values in feature columns must be integral types or float/double types." + ) + return feature_cols + + +def _validate_and_convert_feature_col_as_array_col( + dataset: DataFrame, features_col_name: str +) -> Column: + """It handles + 1. Convert vector type to array type + 2. Cast to Array(Float32)""" + features_col_datatype = dataset.schema[features_col_name].dataType + features_col = col(features_col_name) + if isinstance(features_col_datatype, ArrayType): + if not isinstance( + features_col_datatype.elementType, + (DoubleType, FloatType, LongType, IntegerType, ShortType), + ): + raise ValueError( + "If feature column is array type, its elements must be number type, " + f"got {features_col_datatype.elementType}." + ) + features_array_col = features_col.cast(ArrayType(FloatType())).alias(alias.data) + elif isinstance(features_col_datatype, VectorUDT): + features_array_col = vector_to_array(features_col, dtype="float32").alias( + alias.data + ) + else: + raise ValueError( + "feature column must be array type or `pyspark.ml.linalg.Vector` type, " + "if you want to use multiple numeric columns as features, please use " + "`pyspark.ml.transform.VectorAssembler` to assemble them into a vector " + "type column first." + ) + return features_array_col + + +def _get_unwrapped_vec_cols(feature_col: Column) -> List[Column]: + features_unwrapped_vec_col = unwrap_udt(feature_col) + + # After a `pyspark.ml.linalg.VectorUDT` type column being unwrapped, it becomes + # a pyspark struct type column, the struct fields are: + # - `type`: byte + # - `size`: int + # - `indices`: array + # - `values`: array + # For sparse vector, `type` field is 0, `size` field means vector length, + # `indices` field is the array of active element indices, `values` field + # is the array of active element values. + # For dense vector, `type` field is 1, `size` and `indices` fields are None, + # `values` field is the array of the vector element values. + return [ + features_unwrapped_vec_col.type.alias("featureVectorType"), + features_unwrapped_vec_col.size.alias("featureVectorSize"), + features_unwrapped_vec_col.indices.alias("featureVectorIndices"), + # Note: the value field is double array type, cast it to float32 array type + # for speedup following repartitioning. + features_unwrapped_vec_col.values.cast(ArrayType(FloatType())).alias( + "featureVectorValues" + ), + ] + + +FeatureProp = namedtuple( + "FeatureProp", + ("enable_sparse_data_optim", "has_validation_col", "features_cols_names"), +) + +_MODEL_CHUNK_SIZE = 4096 * 1024 + + +class _SparkXGBEstimator(Estimator, _SparkXGBParams, MLReadable, MLWritable): + _input_kwargs: Dict[str, Any] + + def __init__(self) -> None: + super().__init__() + self._set_xgb_params_default() + self._set_fit_params_default() + self._set_predict_params_default() + # Note: The default value for arbitrary_params_dict must always be empty dict. + # For additional settings added into "arbitrary_params_dict" by default, + # they are added in `setParams`. + self._setDefault( + num_workers=1, + device="cpu", + force_repartition=False, + repartition_random_shuffle=False, + feature_names=None, + feature_types=None, + feature_weights=None, + arbitrary_params_dict={}, + ) + + self.logger = get_logger(self.__class__.__name__) + + def setParams(self, **kwargs: Any) -> None: + """ + Set params for the estimator. + """ + _extra_params = {} + if "arbitrary_params_dict" in kwargs: + raise ValueError("Invalid param name: 'arbitrary_params_dict'.") + + for k, v in kwargs.items(): + # We're not allowing user use features_cols directly. + if k == self.features_cols.name: + raise ValueError( + f"Unsupported param '{k}' please use features_col instead." + ) + if k in _inverse_pyspark_param_alias_map: + raise ValueError( + f"Please use param name {_inverse_pyspark_param_alias_map[k]} instead." + ) + if k in _pyspark_param_alias_map: + if k == _inverse_pyspark_param_alias_map[ + self.featuresCol.name + ] and isinstance(v, list): + real_k = self.features_cols.name + k = real_k + else: + real_k = _pyspark_param_alias_map[k] + k = real_k + + if self.hasParam(k): + if k == "features_col" and isinstance(v, list): + self._set(**{"features_cols": v}) + else: + self._set(**{str(k): v}) + else: + if ( + k in _unsupported_xgb_params + or k in _unsupported_fit_params + or k in _unsupported_predict_params + or k in _unsupported_train_params + ): + err_msg = _unsupported_params_hint_message.get( + k, f"Unsupported param '{k}'." + ) + raise ValueError(err_msg) + _extra_params[k] = v + + _check_distributed_params(kwargs) + _existing_extra_params = self.getOrDefault(self.arbitrary_params_dict) + self._set(arbitrary_params_dict={**_existing_extra_params, **_extra_params}) + + @classmethod + def _pyspark_model_cls(cls) -> Type["_SparkXGBModel"]: + """ + Subclasses should override this method and + returns a _SparkXGBModel subclass + """ + raise NotImplementedError() + + def _create_pyspark_model( + self, xgb_model: XGBModel, training_summary: XGBoostTrainingSummary + ) -> "_SparkXGBModel": + return self._pyspark_model_cls()(xgb_model, training_summary) + + def _convert_to_sklearn_model(self, booster: bytearray, config: str) -> XGBModel: + xgb_sklearn_params = self._gen_xgb_params_dict( + gen_xgb_sklearn_estimator_param=True + ) + sklearn_model = self._xgb_cls()(**xgb_sklearn_params) + sklearn_model.load_model(booster) + sklearn_model._Booster.load_config(config) + return sklearn_model + + def _repartition_needed(self, dataset: DataFrame) -> bool: + """ + We repartition the dataset if the number of workers is not equal to the number of + partitions.""" + if self.getOrDefault(self.force_repartition): + return True + + # In Spark Connect, we cannot easily get the number of partitions. + # For now, since we cannot call rdd.getNumPartitions(), we just return + # True to ensure correct partitioning. + if _is_connect(dataset.sparkSession): + return True + + num_workers = self.getOrDefault(self.num_workers) + num_partitions = dataset.rdd.getNumPartitions() + return not num_workers == num_partitions + + def _get_distributed_train_params(self, dataset: DataFrame) -> Dict[str, Any]: + """ + This just gets the configuration params for distributed xgboost + """ + params = self._gen_xgb_params_dict() + fit_params = self._gen_fit_params_dict() + verbose_eval = fit_params.pop("verbose", None) + + params.update(fit_params) + params["verbose_eval"] = verbose_eval + classification = self._xgb_cls() == XGBClassifier + if classification: + num_classes = int( + dataset.select(countDistinct(alias.label)).collect()[0][0] + ) + if num_classes <= 2: + params["objective"] = "binary:logistic" + else: + params["objective"] = "multi:softprob" + params["num_class"] = num_classes + else: + # use user specified objective or default objective. + # e.g., the default objective for Regressor is 'reg:squarederror' + params["objective"] = self.getOrDefault("objective") + + # TODO: support "num_parallel_tree" for random forest + params["num_boost_round"] = self.getOrDefault("n_estimators") + + return params + + @classmethod + def _get_xgb_train_call_args( + cls, train_params: Dict[str, Any] + ) -> Tuple[Dict[str, Any], Dict[str, Any]]: + xgb_train_default_args = _get_default_params_from_func( + worker_train, _unsupported_train_params + ) + booster_params, kwargs_params = {}, {} + for key, value in train_params.items(): + if key in xgb_train_default_args: + kwargs_params[key] = value + else: + booster_params[key] = value + + booster_params = { + k: v for k, v in booster_params.items() if k not in _non_booster_params + } + return booster_params, kwargs_params + + def _prepare_input_columns_and_feature_prop( + self, dataset: DataFrame + ) -> Tuple[List[Column], FeatureProp]: + label_col = col(self.getOrDefault(self.labelCol)).alias(alias.label) + + select_cols = [label_col] + features_cols_names = None + enable_sparse_data_optim = self.getOrDefault(self.enable_sparse_data_optim) + if enable_sparse_data_optim: + features_col_name = self.getOrDefault(self.featuresCol) + features_col_datatype = dataset.schema[features_col_name].dataType + if not isinstance(features_col_datatype, VectorUDT): + raise ValueError( + "If enable_sparse_data_optim is True, the feature column values must be " + "`pyspark.ml.linalg.Vector` type." + ) + select_cols.extend(_get_unwrapped_vec_cols(col(features_col_name))) + else: + if self.getOrDefault(self.features_cols): + features_cols_names = self.getOrDefault(self.features_cols) + features_cols = _validate_and_convert_feature_col_as_float_col_list( + dataset, features_cols_names + ) + select_cols.extend(features_cols) + else: + features_array_col = _validate_and_convert_feature_col_as_array_col( + dataset, self.getOrDefault(self.featuresCol) + ) + select_cols.append(features_array_col) + + if self._col_is_defined_not_empty(self.weightCol): + select_cols.append( + col(self.getOrDefault(self.weightCol)).alias(alias.weight) + ) + + has_validation_col = False + if self._col_is_defined_not_empty(self.validationIndicatorCol): + select_cols.append( + col(self.getOrDefault(self.validationIndicatorCol)).alias(alias.valid) + ) + # In some cases, see https://issues.apache.org/jira/browse/SPARK-40407, + # the df.repartition can result in some reducer partitions without data, + # which will cause exception or hanging issue when creating DMatrix. + has_validation_col = True + + if self._col_is_defined_not_empty(self.base_margin_col): + select_cols.append( + col(self.getOrDefault(self.base_margin_col)).alias(alias.margin) + ) + + if self._col_is_defined_not_empty(self.qid_col): + select_cols.append(col(self.getOrDefault(self.qid_col)).alias(alias.qid)) + + feature_prop = FeatureProp( + enable_sparse_data_optim, has_validation_col, features_cols_names + ) + return select_cols, feature_prop + + def _prepare_input(self, dataset: DataFrame) -> Tuple[DataFrame, FeatureProp]: + """Prepare the input including column pruning, repartition and so on""" + + select_cols, feature_prop = self._prepare_input_columns_and_feature_prop( + dataset + ) + + dataset = dataset.select(*select_cols) + + num_workers = self.getOrDefault(self.num_workers) + spark_session = dataset.sparkSession + max_concurrent_tasks = _get_max_num_concurrent_tasks(spark_session) + + if feature_prop.has_validation_col: + dtype = dataset.schema[alias.valid].dataType + if not isinstance(dtype, BooleanType): + raise TypeError("The validation indicator must be boolean type.") + + if num_workers > max_concurrent_tasks: + get_logger(self.__class__.__name__).warning( + "The num_workers %s set for xgboost distributed " + "training is greater than current max number of concurrent " + "spark task slots, you need wait until more task slots available " + "or you need increase spark cluster workers.", + num_workers, + ) + + if self._repartition_needed(dataset): + if self._col_is_defined_not_empty(self.qid_col): + # For ranking problem, we need to try best the put the instances with + # same group into the same partition + dataset = dataset.repartitionByRange(num_workers, alias.qid) + else: + # If validationIndicatorCol defined, and if user unionise train and validation + # dataset, users must set force_repartition to true to force repartition. + # Or else some partitions might contain only train or validation dataset. + if self.getOrDefault(self.repartition_random_shuffle): + # In some cases, spark round-robin repartition might cause data skew + # use random shuffle can address it. + dataset = dataset.repartition(num_workers, rand(1)) + else: + dataset = dataset.repartition(num_workers) + + if self._col_is_defined_not_empty(self.qid_col): + # XGBoost requires qid to be sorted for each partition + dataset = dataset.sortWithinPartitions(alias.qid, ascending=True) + + return dataset, feature_prop + + def _get_xgb_parameters( + self, dataset: DataFrame + ) -> Tuple[Dict[str, Any], Dict[str, Any], Dict[str, Any]]: + train_params = self._get_distributed_train_params(dataset) + booster_params, train_call_kwargs_params = self._get_xgb_train_call_args( + train_params + ) + cpu_per_task = int( + dataset.sparkSession.conf.get("spark.task.cpus", None) or "1" + ) + + dmatrix_kwargs = { + "nthread": cpu_per_task, + "feature_types": self.getOrDefault("feature_types"), + "feature_names": self.getOrDefault("feature_names"), + "feature_weights": self.getOrDefault("feature_weights"), + "missing": float(self.getOrDefault("missing")), + } + if dmatrix_kwargs["feature_types"] is not None: + dmatrix_kwargs["enable_categorical"] = True + booster_params["nthread"] = cpu_per_task + + # Remove the parameters whose value is None + booster_params = {k: v for k, v in booster_params.items() if v is not None} + train_call_kwargs_params = { + k: v for k, v in train_call_kwargs_params.items() if v is not None + } + dmatrix_kwargs = {k: v for k, v in dmatrix_kwargs.items() if v is not None} + + return booster_params, train_call_kwargs_params, dmatrix_kwargs + + def _skip_stage_level_scheduling(self, spark_session: SparkSession) -> bool: + # pylint: disable=too-many-return-statements + """Check if stage-level scheduling is not needed, + return true to skip stage-level scheduling""" + + if self._run_on_gpu(spark_session): + executor_cores = spark_session.conf.get("spark.executor.cores", None) + executor_gpus = spark_session.conf.get( + "spark.executor.resource.gpu.amount", None + ) + if executor_cores is None or executor_gpus is None: + self.logger.info( + "Stage-level scheduling in xgboost requires spark.executor.cores, " + "spark.executor.resource.gpu.amount to be set." + ) + return True + + if int(executor_cores) == 1: + # there will be only 1 task running at any time. + self.logger.info( + "Stage-level scheduling in xgboost requires spark.executor.cores > 1 " + ) + return True + + if int(executor_gpus) > 1: + # For spark.executor.resource.gpu.amount > 1, we suppose user knows how to configure + # to make xgboost run successfully. + # + self.logger.info( + "Stage-level scheduling in xgboost will not work " + "when spark.executor.resource.gpu.amount>1" + ) + return True + + task_gpu_amount = spark_session.conf.get( + "spark.task.resource.gpu.amount", None + ) + + if task_gpu_amount is None: + # The ETL tasks will not grab a gpu when spark.task.resource.gpu.amount is not set, + # but with stage-level scheduling, we can make training task grab the gpu. + return False + + if float(task_gpu_amount) == float(executor_gpus): + # spark.executor.resource.gpu.amount=spark.task.resource.gpu.amount " + # results in only 1 task running at a time, which may cause perf issue. + return True + + # We can enable stage-level scheduling + return False + + # CPU training doesn't require stage-level scheduling + return True + + def _get_resource_profile_for_stage_level_scheduling( + self, spark_session: SparkSession + ) -> Optional[ResourceProfile]: + """Try to enable stage-level scheduling""" + conf = spark_session.conf + if _is_local(spark_session) or self._skip_stage_level_scheduling(spark_session): + return None + + # executor_cores will not be None + executor_cores = conf.get("spark.executor.cores", None) + assert executor_cores is not None + + # Spark-rapids is a project to leverage GPUs to accelerate spark SQL. + # If spark-rapids is enabled, to avoid GPU OOM, we don't allow other + # ETL gpu tasks running alongside training tasks. + spark_plugins = spark_session.conf.get("spark.plugins", " ") + assert spark_plugins is not None + spark_rapids_sql_enabled = spark_session.conf.get( + "spark.rapids.sql.enabled", "true" + ) + assert spark_rapids_sql_enabled is not None + + task_cores = ( + int(executor_cores) + if "com.nvidia.spark.SQLPlugin" in spark_plugins + and "true" == spark_rapids_sql_enabled.lower() + else (int(executor_cores) // 2) + 1 + ) + + # Each training task requires cpu cores > total executor cores//2 + 1 which can + # make sure the tasks be sent to different executors. + # + # Please note that we can't use GPU to limit the concurrent tasks because of + # https://issues.apache.org/jira/browse/SPARK-45527. + + task_gpus = 1.0 + treqs = TaskResourceRequests().cpus(task_cores).resource("gpu", task_gpus) + rp = ResourceProfileBuilder().require(treqs).build + + self.logger.info( + "XGBoost training tasks require the resource(cores=%s, gpu=%s).", + task_cores, + task_gpus, + ) + return rp + + def _get_tracker_args( + self, spark_session: SparkSession + ) -> Tuple[bool, Dict[str, Any]]: + """Start the tracker and return the tracker envs on the driver side""" + if self.isDefined(self.launch_tracker_on_driver): + launch_tracker_on_driver = self.getOrDefault(self.launch_tracker_on_driver) + else: + launch_tracker_on_driver = not _is_connect(spark_session) + + rabit_args = {} + if launch_tracker_on_driver: + conf = Config() + if self.isDefined(self.coll_cfg): + conf = self.getOrDefault(self.coll_cfg) + assert isinstance(conf, Config) + + if conf.tracker_host_ip is None: + conf.tracker_host_ip = spark_session.conf.get("spark.driver.host", None) + num_workers = self.getOrDefault(self.num_workers) + rabit_args.update(_get_rabit_args(conf, num_workers)) + else: + if self.isDefined(self.coll_cfg): + conf = self.getOrDefault(self.coll_cfg) + assert isinstance(conf, Config) + if conf.tracker_host_ip is not None: + raise ValueError( + f"You must enable launch_tracker_on_driver to use " + f"tracker host: {conf.tracker_host_ip}" + ) + return launch_tracker_on_driver, rabit_args + + def _fit(self, dataset: DataFrame) -> "_SparkXGBModel": + # pylint: disable=too-many-statements, too-many-locals + self._validate_params(dataset.sparkSession) + + dataset, feature_prop = self._prepare_input(dataset) + + ( + booster_params, + train_call_kwargs_params, + dmatrix_kwargs, + ) = self._get_xgb_parameters(dataset) + + run_on_gpu = self._run_on_gpu(dataset.sparkSession) + + is_local = _is_local(dataset.sparkSession) + + num_workers = self.getOrDefault(self.num_workers) + + launch_tracker_on_driver, rabit_args = self._get_tracker_args( + dataset.sparkSession + ) + conf: Optional[Config] = ( + self.getOrDefault(self.coll_cfg) if self.isDefined(self.coll_cfg) else None + ) + + log_level = get_logger_level(_LOG_TAG) + + use_rmm = get_config()["use_rmm"] + + def _train_booster( + pandas_df_iter: Iterator[pd.DataFrame], + ) -> Iterator[pd.DataFrame]: + """Takes in an RDD partition and outputs a booster for that partition after + going through the Rabit Ring protocol + + """ + from pyspark import BarrierTaskContext + + context = BarrierTaskContext.get() + context.barrier() + + dev_ordinal = None + use_qdm = _can_use_qdm( + booster_params.get("tree_method", None), + booster_params.get("device", None), + ) + verbosity = booster_params.get("verbosity", 1) + msg = "Training on CPUs" + if run_on_gpu: + dev_ordinal = ( + context.partitionId() if is_local else _get_gpu_id(context) + ) + booster_params["device"] = "cuda:" + str(dev_ordinal) + # If cuDF is not installed, then using DMatrix instead of QDM, + # because without cuDF, DMatrix performs better than QDM. + # Note: Checking `is_cudf_available` in spark worker side because + # spark worker might has different python environment with driver side. + use_qdm = use_qdm and is_cudf_available() + msg = ( + f"Leveraging {booster_params['device']} to train with " + f"QDM: {'on' if use_qdm else 'off'}" + ) + + if use_qdm and (booster_params.get("max_bin", None) is not None): + dmatrix_kwargs["max_bin"] = booster_params["max_bin"] + _rabit_args = rabit_args + if context.partitionId() == 0: + if not launch_tracker_on_driver: + _conf = conf if conf is not None else Config() + _conf.tracker_host_ip = _get_host_ip(context) + _rabit_args = _get_rabit_args(_conf, num_workers) + get_logger(_LOG_TAG, log_level).info(msg) + + worker_message: Dict[str, Any] = { + "use_qdm": use_qdm, + } + + if not launch_tracker_on_driver: + worker_message["rabit_msg"] = _rabit_args + + messages = context.allGather(message=json.dumps(worker_message)) + if len(set(json.loads(x)["use_qdm"] for x in messages)) != 1: + raise RuntimeError("The workers' cudf environments are in-consistent ") + + if not launch_tracker_on_driver: + _rabit_args = json.loads(messages[0])["rabit_msg"] + + if conf is not None: + _rabit_args = conf.update_worker_args(_rabit_args) + + evals_result: Dict[str, Any] = {} + with ( + config_context(verbosity=verbosity, use_rmm=use_rmm), + CommunicatorContext(context, **_rabit_args), + ): + dtrain, dvalid = create_dmatrix_from_partitions( + iterator=pandas_df_iter, + feature_cols=feature_prop.features_cols_names, + dev_ordinal=dev_ordinal, + use_qdm=use_qdm, + kwargs=dmatrix_kwargs, + enable_sparse_data_optim=feature_prop.enable_sparse_data_optim, + has_validation_col=feature_prop.has_validation_col, + ) + if dvalid is not None: + dval = [(dtrain, "training"), (dvalid, "validation")] + else: + dval = [(dtrain, "training")] + booster = worker_train( + params=booster_params, + dtrain=dtrain, + evals=dval, + evals_result=evals_result, + **train_call_kwargs_params, + ) + context.barrier() + + if context.partitionId() == 0: + yield pd.DataFrame({"data": [json.dumps(dict(evals_result))]}) + config = booster.save_config() + yield pd.DataFrame({"data": [config]}) + booster_json = booster.save_raw("json").decode("utf-8") + + for offset in range(0, len(booster_json), _MODEL_CHUNK_SIZE): + booster_chunk = booster_json[offset : offset + _MODEL_CHUNK_SIZE] + yield pd.DataFrame({"data": [booster_chunk]}) + + def _run_job() -> Tuple[str, str, str]: + df = dataset.mapInPandas( + _train_booster, # type: ignore[arg-type] + schema="data string", + barrier=True, + profile=self._get_resource_profile_for_stage_level_scheduling( + dataset.sparkSession + ), + ) + ret = df.collect() + data = [row.data for row in ret] + return data[0], data[1], "".join(data[2:]) + + get_logger(_LOG_TAG).info( + "Running xgboost-%s on %s workers with" + "\n\tbooster params: %s" + "\n\ttrain_call_kwargs_params: %s" + "\n\tdmatrix_kwargs: %s", + _py_version(), + num_workers, + booster_params, + train_call_kwargs_params, + dmatrix_kwargs, + ) + evals_result, config, booster = _run_job() + get_logger(_LOG_TAG).info("Finished xgboost training!") + + result_xgb_model = self._convert_to_sklearn_model( + bytearray(booster, "utf-8"), config + ) + training_summary = XGBoostTrainingSummary.from_metrics(json.loads(evals_result)) + spark_model = self._create_pyspark_model(result_xgb_model, training_summary) + # According to pyspark ML convention, the model uid should be the same + # with estimator uid. + spark_model._resetUid(self.uid) + return self._copyValues(spark_model) + + def write(self) -> "SparkXGBWriter": + """ + Return the writer for saving the estimator. + """ + return SparkXGBWriter(self) + + @classmethod + def read(cls) -> "SparkXGBReader": + """ + Return the reader for loading the estimator. + """ + return SparkXGBReader(cls) + + +class _SparkXGBModel(Model, _SparkXGBParams, MLReadable, MLWritable): + def __init__( + self, + xgb_sklearn_model: Optional[XGBModel] = None, + training_summary: Optional[XGBoostTrainingSummary] = None, + ) -> None: + super().__init__() + self._xgb_sklearn_model = xgb_sklearn_model + self.training_summary = training_summary + + @classmethod + def _xgb_cls(cls) -> Type[XGBModel]: + raise NotImplementedError() + + def get_booster(self) -> Booster: + """ + Return the `xgboost.core.Booster` instance. + """ + assert self._xgb_sklearn_model is not None + return self._xgb_sklearn_model.get_booster() + + def get_feature_importances( + self, importance_type: str = "weight" + ) -> Dict[str, Union[float, List[float]]]: + """Get feature importance of each feature. + Importance type can be defined as: + + * 'weight': the number of times a feature is used to split the data across all trees. + * 'gain': the average gain across all splits the feature is used in. + * 'cover': the average coverage across all splits the feature is used in. + * 'total_gain': the total gain across all splits the feature is used in. + * 'total_cover': the total coverage across all splits the feature is used in. + + Parameters + ---------- + importance_type: str, default 'weight' + One of the importance types defined above. + """ + return self.get_booster().get_score(importance_type=importance_type) + + def write(self) -> "SparkXGBModelWriter": + """ + Return the writer for saving the model. + """ + return SparkXGBModelWriter(self) + + @classmethod + def read(cls) -> "SparkXGBModelReader": + """ + Return the reader for loading the model. + """ + return SparkXGBModelReader(cls) + + def _get_feature_col( + self, dataset: DataFrame + ) -> Tuple[List[Column], Optional[List[str]]]: + """XGBoost model trained with features_cols parameter can also predict + vector or array feature type. But first we need to check features_cols + and then featuresCol + """ + if self.getOrDefault(self.enable_sparse_data_optim): + feature_col_names = None + features_col = _get_unwrapped_vec_cols( + col(self.getOrDefault(self.featuresCol)) + ) + return features_col, feature_col_names + + feature_col_names = self.getOrDefault(self.features_cols) + features_col = [] + if feature_col_names and set(feature_col_names).issubset(set(dataset.columns)): + # The model is trained with features_cols and the predicted dataset + # also contains all the columns specified by features_cols. + features_col = _validate_and_convert_feature_col_as_float_col_list( + dataset, feature_col_names + ) + else: + # 1. The model was trained by features_cols, but the dataset doesn't contain + # all the columns specified by features_cols, so we need to check if + # the dataframe has the featuresCol + # 2. The model was trained by featuresCol, and the predicted dataset must contain + # featuresCol column. + feature_col_names = None + features_col.append( + _validate_and_convert_feature_col_as_array_col( + dataset, self.getOrDefault(self.featuresCol) + ) + ) + return features_col, feature_col_names + + def _get_pred_contrib_col_name(self) -> Optional[str]: + """Return the pred_contrib_col col name""" + pred_contrib_col_name = None + if self._col_is_defined_not_empty(self.pred_contrib_col): + pred_contrib_col_name = self.getOrDefault(self.pred_contrib_col) + + return pred_contrib_col_name + + def _out_schema(self) -> Tuple[bool, str]: + """Return the bool to indicate if it's a single prediction, true is single prediction, + and the returned type of the user-defined function. The value must + be a DDL-formatted type string.""" + + if self._get_pred_contrib_col_name() is not None: + return False, f"{pred.prediction} double, {pred.pred_contrib} array" + + return True, "double" + + def _get_predict_func(self) -> Callable: + """Return the true prediction function which will be running on the executor side""" + + predict_params = self._gen_predict_params_dict() + pred_contrib_col_name = self._get_pred_contrib_col_name() + + def _predict( + model: XGBModel, X: ArrayLike, base_margin: Optional[ArrayLike] + ) -> Union[pd.DataFrame, pd.Series]: + data = {} + preds = model.predict( + X, + base_margin=base_margin, + validate_features=False, + **predict_params, + ) + data[pred.prediction] = pd.Series(preds) + + if pred_contrib_col_name is not None: + contribs = pred_contribs(model, X, base_margin) + data[pred.pred_contrib] = pd.Series(list(contribs)) + return pd.DataFrame(data=data) + + return data[pred.prediction] + + return _predict + + def _post_transform(self, dataset: DataFrame, pred_col: Column) -> DataFrame: + """Post process of transform""" + prediction_col_name = self.getOrDefault(self.predictionCol) + single_pred, _ = self._out_schema() + + if single_pred: + if prediction_col_name: + dataset = dataset.withColumn(prediction_col_name, pred_col) + else: + pred_struct_col = "_prediction_struct" + dataset = dataset.withColumn(pred_struct_col, pred_col) + + if prediction_col_name: + dataset = dataset.withColumn( + prediction_col_name, getattr(col(pred_struct_col), pred.prediction) + ) + + pred_contrib_col_name = self._get_pred_contrib_col_name() + if pred_contrib_col_name is not None: + dataset = dataset.withColumn( + pred_contrib_col_name, + array_to_vector(getattr(col(pred_struct_col), pred.pred_contrib)), + ) + + dataset = dataset.drop(pred_struct_col) + return dataset + + def _run_on_gpu(self, spark_session: SparkSession) -> bool: + """If gpu is used to do the prediction according to the parameters + and spark configurations""" + + use_gpu_by_params = super()._run_on_gpu(spark_session) + + if _is_local(spark_session): + # if it's local model, no need to check the spark configurations + return use_gpu_by_params + + gpu_per_task = spark_session.conf.get("spark.task.resource.gpu.amount", None) + + # User don't set gpu configurations, just use cpu + if gpu_per_task is None: + if use_gpu_by_params: + get_logger(_LOG_TAG).warning( + "Do the prediction on the CPUs since no gpu configurations are set" + ) + return False + + # User already sets the gpu configurations. + return use_gpu_by_params + + def _transform(self, dataset: DataFrame) -> DataFrame: + # pylint: disable=too-many-statements, too-many-locals + # Save xgb_sklearn_model and predict_params to be local variable + # to avoid the `self` object to be pickled to remote. + xgb_sklearn_model = self._xgb_sklearn_model + + base_margin_col = None + if self._col_is_defined_not_empty(self.base_margin_col): + base_margin_col = col(self.getOrDefault(self.base_margin_col)).alias( + alias.margin + ) + has_base_margin = base_margin_col is not None + + features_col, feature_col_names = self._get_feature_col(dataset) + enable_sparse_data_optim = self.getOrDefault(self.enable_sparse_data_optim) + + predict_func = self._get_predict_func() + + _, schema = self._out_schema() + + is_local = _is_local(dataset.sparkSession) + run_on_gpu = self._run_on_gpu(dataset.sparkSession) + + log_level = get_logger_level(_LOG_TAG) + + @pandas_udf(schema) # type: ignore[call-overload] + def predict_udf(iterator: Iterator[pd.DataFrame]) -> Iterator[pd.Series]: + assert xgb_sklearn_model is not None + model = xgb_sklearn_model + + from pyspark import TaskContext + + context = TaskContext.get() + assert context is not None + + dev_ordinal = -1 + + msg = "Do the inference on the CPUs" + if run_on_gpu: + if is_cudf_available() and is_cupy_available(): + if is_local: + cp = import_cupy() + + total_gpus = cp.cuda.runtime.getDeviceCount() + if total_gpus > 0: + partition_id = context.partitionId() + # For transform local mode, default the dev_ordinal to + # (partition id) % gpus. + dev_ordinal = partition_id % total_gpus + else: + dev_ordinal = _get_gpu_id(context) + + if dev_ordinal >= 0: + device = "cuda:" + str(dev_ordinal) + msg = "Do the inference with device: " + device + model.set_params(device=device) + else: + msg = "Couldn't get the correct gpu id, fallback the inference on the CPUs" + else: + msg = "CUDF or Cupy is unavailable, fallback the inference on the CPUs" + + if context.partitionId() == 0: + get_logger(_LOG_TAG, log_level).info(msg) + + def to_gpu_if_possible(data: ArrayLike) -> ArrayLike: + """Move the data to gpu if possible""" + if dev_ordinal >= 0: + import cudf + import cupy as cp + + # We must set the device after import cudf, which will change the device id to 0 + # See https://github.com/rapidsai/cudf/issues/11386 + cp.cuda.runtime.setDevice(dev_ordinal) # pylint: disable=I1101 + df = cudf.DataFrame(data) + del data + return df + return data + + for data in iterator: + if enable_sparse_data_optim: + X = _read_csr_matrix_from_unwrapped_spark_vec(data) + else: + if feature_col_names is not None: + tmp: ArrayLike = data[feature_col_names] + else: + tmp = stack_series(data[alias.data]) + X = to_gpu_if_possible(tmp) + + if has_base_margin: + base_margin = to_gpu_if_possible(data[alias.margin]) + else: + base_margin = None + + yield predict_func(model, X, base_margin) + + if has_base_margin: + assert base_margin_col is not None + pred_col = predict_udf(struct(*features_col, base_margin_col)) + else: + pred_col = predict_udf(struct(*features_col)) + + return self._post_transform(dataset, pred_col) + + +class _ClassificationModel( # pylint: disable=abstract-method + _SparkXGBModel, HasProbabilityCol, HasRawPredictionCol, HasContribPredictionCol +): + """ + The model returned by :func:`xgboost.spark.SparkXGBClassifier.fit` + + .. Note:: This API is experimental. + """ + + def _out_schema(self) -> Tuple[bool, str]: + schema = ( + f"{pred.raw_prediction} array, {pred.prediction} double," + f" {pred.probability} array" + ) + if self._get_pred_contrib_col_name() is not None: + # We will force setting strict_shape to True when predicting contribs, + # So, it will also output 3-D shape result. + schema = f"{schema}, {pred.pred_contrib} array>" + + return False, schema + + def _get_predict_func(self) -> Callable: + predict_params = self._gen_predict_params_dict() + pred_contrib_col_name = self._get_pred_contrib_col_name() + + def transform_margin(margins: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: + if margins.ndim == 1: + # binomial case + classone_probs = expit(margins) + classzero_probs = 1.0 - classone_probs + raw_preds = np.vstack((-margins, margins)).transpose() + class_probs = np.vstack((classzero_probs, classone_probs)).transpose() + else: + # multinomial case + raw_preds = margins + class_probs = softmax(raw_preds, axis=1) + return raw_preds, class_probs + + def _predict( + model: XGBModel, X: ArrayLike, base_margin: Optional[np.ndarray] + ) -> Union[pd.DataFrame, pd.Series]: + margins = model.predict( + X, + base_margin=base_margin, + output_margin=True, + validate_features=False, + **predict_params, + ) + raw_preds, class_probs = transform_margin(margins) + + # It seems that they use argmax of class probs, + # not of margin to get the prediction (Note: scala implementation) + preds = np.argmax(class_probs, axis=1) + result: Dict[str, pd.Series] = { + pred.raw_prediction: pd.Series(list(raw_preds)), + pred.prediction: pd.Series(preds), + pred.probability: pd.Series(list(class_probs)), + } + + if pred_contrib_col_name is not None: + contribs = pred_contribs(model, X, base_margin, strict_shape=True) + result[pred.pred_contrib] = pd.Series(list(contribs.tolist())) + + return pd.DataFrame(data=result) + + return _predict + + def _post_transform(self, dataset: DataFrame, pred_col: Column) -> DataFrame: + pred_struct_col = "_prediction_struct" + dataset = dataset.withColumn(pred_struct_col, pred_col) + + raw_prediction_col_name = self.getOrDefault(self.rawPredictionCol) + if raw_prediction_col_name: + dataset = dataset.withColumn( + raw_prediction_col_name, + array_to_vector(getattr(col(pred_struct_col), pred.raw_prediction)), + ) + + prediction_col_name = self.getOrDefault(self.predictionCol) + if prediction_col_name: + dataset = dataset.withColumn( + prediction_col_name, getattr(col(pred_struct_col), pred.prediction) + ) + + probability_col_name = self.getOrDefault(self.probabilityCol) + if probability_col_name: + dataset = dataset.withColumn( + probability_col_name, + array_to_vector(getattr(col(pred_struct_col), pred.probability)), + ) + + pred_contrib_col_name = self._get_pred_contrib_col_name() + if pred_contrib_col_name is not None: + dataset = dataset.withColumn( + pred_contrib_col_name, + getattr(col(pred_struct_col), pred.pred_contrib), + ) + + return dataset.drop(pred_struct_col) + + +class _SparkXGBSharedReadWrite: + @staticmethod + def saveMetadata( + instance: Union[_SparkXGBEstimator, _SparkXGBModel], + path: str, + spark_session: SparkSession, + logger: logging.Logger, + extraMetadata: Optional[Dict[str, Any]] = None, + ) -> None: + """ + Save the metadata of an xgboost.spark._SparkXGBEstimator or + xgboost.spark._SparkXGBModel. + """ + instance._validate_params(spark_session) + skipParams = ["callbacks", "xgb_model", "coll_cfg"] + jsonParams = {} + for p, v in instance._paramMap.items(): # pylint: disable=protected-access + if p.name not in skipParams: + jsonParams[p.name] = v + + extraMetadata = extraMetadata or {} + callbacks = instance.getOrDefault("callbacks") + if callbacks is not None: + logger.warning( + "The callbacks parameter is saved using cloudpickle and it " + "is not a fully self-contained format. It may fail to load " + "with different versions of dependencies." + ) + serialized_callbacks = base64.encodebytes( + cloudpickle.dumps(callbacks) + ).decode("ascii") + extraMetadata["serialized_callbacks"] = serialized_callbacks + init_booster = instance.getOrDefault("xgb_model") + if init_booster is not None: + extraMetadata["init_booster"] = _INIT_BOOSTER_SAVE_PATH + + if instance.isDefined("coll_cfg"): + conf: Config = instance.getOrDefault("coll_cfg") + if conf is not None: + extraMetadata["coll_cfg"] = { + k: v for k, v in asdict(conf).items() if not callable(v) + } + if callable(conf.worker_port): + logger.warning("The `worker_port` is not serialized.") + + DefaultParamsWriter.saveMetadata( + instance, + path, + spark_session if _is_connect(spark_session) else spark_session.sparkContext, + extraMetadata=extraMetadata, + paramMap=jsonParams, + ) + if init_booster is not None: + ser_init_booster = serialize_booster(init_booster) + save_path = os.path.join(path, _INIT_BOOSTER_SAVE_PATH) + spark_session.createDataFrame( + [(ser_init_booster,)], ["init_booster"] + ).write.parquet(save_path) + + @staticmethod + def loadMetadataAndInstance( + pyspark_xgb_cls: Union[Type[_SparkXGBEstimator], Type[_SparkXGBModel]], + path: str, + spark_session: SparkSession, + logger: logging.Logger, + ) -> Tuple[Dict[str, Any], Union[_SparkXGBEstimator, _SparkXGBModel]]: + """ + Load the metadata and the instance of an xgboost.spark._SparkXGBEstimator or + xgboost.spark._SparkXGBModel. + + :return: a tuple of (metadata, instance) + """ + metadata = DefaultParamsReader.loadMetadata( + path, + spark_session if _is_connect(spark_session) else spark_session.sparkContext, + expectedClassName=get_class_name(pyspark_xgb_cls), + ) + pyspark_xgb = pyspark_xgb_cls() + DefaultParamsReader.getAndSetParams(pyspark_xgb, metadata) + + if "serialized_callbacks" in metadata: + serialized_callbacks = metadata["serialized_callbacks"] + try: + callbacks = cloudpickle.loads( + base64.decodebytes(serialized_callbacks.encode("ascii")) + ) + pyspark_xgb.set(pyspark_xgb.callbacks, callbacks) # type: ignore[union-attr] + except Exception as e: # pylint: disable=W0703 + logger.warning( + f"Fails to load the callbacks param due to {e}. Please set the " + "callbacks param manually for the loaded estimator." + ) + if "coll_cfg" in metadata: + pyspark_xgb.set_coll_cfg(Config(**metadata["coll_cfg"])) + + if "init_booster" in metadata: + load_path = os.path.join(path, metadata["init_booster"]) + ser_init_booster = ( + spark_session.read.parquet(load_path).collect()[0].init_booster + ) + init_booster = deserialize_booster(ser_init_booster) + pyspark_xgb.set(pyspark_xgb.xgb_model, init_booster) # type: ignore[union-attr] + + pyspark_xgb._resetUid(metadata["uid"]) # pylint: disable=protected-access + return metadata, pyspark_xgb + + +class SparkXGBWriter(MLWriter): + """ + Spark Xgboost estimator writer. + """ + + def __init__(self, instance: "_SparkXGBEstimator") -> None: + super().__init__() + self.instance = instance + self.logger = get_logger(self.__class__.__name__, level="WARN") + + def saveImpl(self, path: str) -> None: + """ + save model. + """ + _SparkXGBSharedReadWrite.saveMetadata( + self.instance, path, self.sparkSession, self.logger + ) + + +class SparkXGBReader(MLReader): + """ + Spark Xgboost estimator reader. + """ + + def __init__(self, cls: Type["_SparkXGBEstimator"]) -> None: + super().__init__() + self.cls = cls + self.logger = get_logger(self.__class__.__name__, level="WARN") + + def load(self, path: str) -> "_SparkXGBEstimator": + """ + load model. + """ + _, pyspark_xgb = _SparkXGBSharedReadWrite.loadMetadataAndInstance( + self.cls, path, self.sparkSession, self.logger + ) + return cast("_SparkXGBEstimator", pyspark_xgb) + + +class SparkXGBModelWriter(MLWriter): + """ + Spark Xgboost model writer. + """ + + def __init__(self, instance: _SparkXGBModel) -> None: + super().__init__() + self.instance = instance + self.logger = get_logger(self.__class__.__name__, level="WARN") + + def saveImpl(self, path: str) -> None: + """ + Save metadata and model for a :py:class:`_SparkXGBModel` + - save metadata to path/metadata + - save model to path/model.json + """ + xgb_model = self.instance._xgb_sklearn_model + assert xgb_model is not None + _SparkXGBSharedReadWrite.saveMetadata( + self.instance, path, self.sparkSession, self.logger + ) + model_save_path = os.path.join(path, "model") + booster = xgb_model.get_booster().save_raw("json").decode("utf-8") + booster_chunks = [] + + for offset in range(0, len(booster), _MODEL_CHUNK_SIZE): + booster_chunks.append(booster[offset : offset + _MODEL_CHUNK_SIZE]) + + # Write model preserving row ordering + indexed_chunks = [[i, c] for i, c in enumerate(booster_chunks)] + ( + self.sparkSession.createDataFrame(indexed_chunks, ["idx", "value"]) + .repartition(1) + .sortWithinPartitions("idx") + .select("value") + .write.text(model_save_path) + ) + + +class SparkXGBModelReader(MLReader): + """ + Spark Xgboost model reader. + """ + + def __init__(self, cls: Type["_SparkXGBModel"]) -> None: + super().__init__() + self.cls = cls + self.logger = get_logger(self.__class__.__name__, level="WARN") + + def load(self, path: str) -> "_SparkXGBModel": + """ + Load metadata and model for a :py:class:`_SparkXGBModel` + + :return: SparkXGBRegressorModel or SparkXGBClassifierModel instance + """ + _, py_model = _SparkXGBSharedReadWrite.loadMetadataAndInstance( + self.cls, path, self.sparkSession, self.logger + ) + py_model = cast("_SparkXGBModel", py_model) + + xgb_sklearn_params = py_model._gen_xgb_params_dict( + gen_xgb_sklearn_estimator_param=True + ) + model_load_path = os.path.join(path, "model") + + ser_xgb_model = "".join( + [r[0] for r in self.sparkSession.read.text(model_load_path).collect()] + ) + + def create_xgb_model() -> "XGBModel": + return self.cls._xgb_cls()(**xgb_sklearn_params) + + xgb_model = deserialize_xgb_model(ser_xgb_model, create_xgb_model) + py_model._xgb_sklearn_model = xgb_model + return py_model diff --git a/v12x/pylib/xgboost/spark/data.py b/v12x/pylib/xgboost/spark/data.py new file mode 100644 index 0000000000000000000000000000000000000000..77fd89d409fa06869ba4f989211f95ef87f69252 --- /dev/null +++ b/v12x/pylib/xgboost/spark/data.py @@ -0,0 +1,365 @@ +# pylint: disable=protected-access +"""Utilities for processing spark partitions.""" + +from collections import defaultdict, namedtuple +from typing import Any, Callable, Dict, Iterator, List, Optional, Sequence, Tuple, Union + +import numpy as np +import pandas as pd +from scipy.sparse import csr_matrix + +from .._typing import ArrayLike +from ..compat import concat +from ..core import DataIter, DMatrix, QuantileDMatrix +from ..sklearn import XGBModel +from .utils import get_logger + + +def stack_series(series: pd.Series) -> np.ndarray: + """Stack a series of arrays.""" + array = series.to_numpy(copy=False) + array = np.stack(array) # type: ignore[arg-type] + return array + + +# Global constant for defining column alias shared between estimator and data +# processing procedures. +Alias = namedtuple("Alias", ("data", "label", "weight", "margin", "valid", "qid")) +alias = Alias("values", "label", "weight", "baseMargin", "validationIndicator", "qid") + + +def concat_or_none(seq: Optional[Sequence[np.ndarray]]) -> Optional[np.ndarray]: + """Concatenate the data if it's not None.""" + if seq: + return concat(seq) + return None + + +def cache_partitions( + iterator: Iterator[pd.DataFrame], append: Callable[[pd.DataFrame, str, bool], None] +) -> None: + """Extract partitions from pyspark iterator. `append` is a user defined function for + accepting new partition.""" + + def make_blob(part: pd.DataFrame, is_valid: bool) -> None: + append(part, alias.data, is_valid) + append(part, alias.label, is_valid) + append(part, alias.weight, is_valid) + append(part, alias.margin, is_valid) + append(part, alias.qid, is_valid) + + has_validation: Optional[bool] = None + + for part in iterator: + if has_validation is None: + has_validation = alias.valid in part.columns + if has_validation is True: + assert alias.valid in part.columns + + if has_validation: + train = part.loc[~part[alias.valid], :] + valid = part.loc[part[alias.valid], :] + else: + train, valid = part, None + + make_blob(train, False) + if valid is not None: + make_blob(valid, True) + + +class PartIter(DataIter): + """Iterator for creating Quantile DMatrix from partitions.""" + + def __init__( + self, data: Dict[str, List], device_id: Optional[int], **kwargs: Any + ) -> None: + self._iter = 0 + self._device_id = device_id + self._data = data + self._kwargs = kwargs + + super().__init__(release_data=True) + + def _fetch(self, data: Optional[Sequence[pd.DataFrame]]) -> Optional[pd.DataFrame]: + if not data: + return None + + if self._device_id is not None: + import cudf + import cupy as cp + + # We must set the device after import cudf, which will change the device id to 0 + # See https://github.com/rapidsai/cudf/issues/11386 + cp.cuda.runtime.setDevice(self._device_id) # pylint: disable=I1101 + return cudf.DataFrame(data[self._iter]) + + return data[self._iter] + + def next(self, input_data: Callable) -> bool: + if self._iter == len(self._data[alias.data]): + return False + input_data( + data=self._fetch(self._data[alias.data]), + label=self._fetch(self._data.get(alias.label, None)), + weight=self._fetch(self._data.get(alias.weight, None)), + base_margin=self._fetch(self._data.get(alias.margin, None)), + qid=self._fetch(self._data.get(alias.qid, None)), + **self._kwargs, + ) + self._iter += 1 + return True + + def reset(self) -> None: + self._iter = 0 + + +def _read_csr_matrix_from_unwrapped_spark_vec(part: pd.DataFrame) -> csr_matrix: + # variables for constructing csr_matrix + csr_indices_list, csr_indptr_list, csr_values_list = [], [0], [] + + n_features = 0 + + for vec_type, vec_size_, vec_indices, vec_values in zip( + part.featureVectorType, + part.featureVectorSize, + part.featureVectorIndices, + part.featureVectorValues, + ): + if vec_type == 0: + # sparse vector + vec_size = int(vec_size_) + csr_indices = vec_indices + csr_values = vec_values + else: + # dense vector + # Note: According to spark ML VectorUDT format, + # when type field is 1, the size field is also empty. + # we need to check the values field to get vector length. + vec_size = len(vec_values) + csr_indices = np.arange(vec_size, dtype=np.int32) + csr_values = vec_values + + if n_features == 0: + n_features = vec_size + assert n_features == vec_size + + csr_indices_list.append(csr_indices) + csr_indptr_list.append(csr_indptr_list[-1] + len(csr_indices)) + csr_values_list.append(csr_values) + + csr_indptr_arr = np.array(csr_indptr_list) + csr_indices_arr = np.concatenate(csr_indices_list) + csr_values_arr = np.concatenate(csr_values_list) + + return csr_matrix( + (csr_values_arr, csr_indices_arr, csr_indptr_arr), shape=(len(part), n_features) + ) + + +def make_qdm( + data: Dict[str, List[np.ndarray]], + dev_ordinal: Optional[int], + meta: Dict[str, Any], + ref: Optional[DMatrix], + params: Dict[str, Any], +) -> DMatrix: + """Handle empty partition for QuantileDMatrix.""" + if not data: + return QuantileDMatrix(np.empty((0, 0)), ref=ref) + it = PartIter(data, dev_ordinal, **meta) + m = QuantileDMatrix(it, **params, ref=ref) + return m + + +def create_dmatrix_from_partitions( # pylint: disable=too-many-arguments + *, + iterator: Iterator[pd.DataFrame], + feature_cols: Optional[Sequence[str]], + dev_ordinal: Optional[int], + use_qdm: bool, + kwargs: Dict[str, Any], # use dict to make sure this parameter is passed. + enable_sparse_data_optim: bool, + has_validation_col: bool, +) -> Tuple[DMatrix, Optional[DMatrix]]: + """Create DMatrix from spark data partitions. + + Parameters + ---------- + iterator : + Pyspark partition iterator. + feature_cols: + A sequence of feature names, used only when rapids plugin is enabled. + dev_ordinal: + Device ordinal, used when GPU is enabled. + use_qdm : + Whether QuantileDMatrix should be used instead of DMatrix. + kwargs : + Metainfo for DMatrix. + enable_sparse_data_optim : + Whether sparse data should be unwrapped + has_validation: + Whether there's validation data. + + Returns + ------- + Training DMatrix and an optional validation DMatrix. + """ + # pylint: disable=too-many-locals, too-many-statements + train_data: Dict[str, List[np.ndarray]] = defaultdict(list) + valid_data: Dict[str, List[np.ndarray]] = defaultdict(list) + + n_features: int = 0 + + def append_m(part: pd.DataFrame, name: str, is_valid: bool) -> None: + nonlocal n_features + if name == alias.data or name in part.columns: + if ( + name == alias.data + and feature_cols is not None + and part[feature_cols].shape[0] > 0 # guard against empty partition + ): + array: Optional[np.ndarray] = part[feature_cols] + elif part[name].shape[0] > 0: + array = part[name] + if name == alias.data: + # For the array/vector typed case. + array = stack_series(array) + else: + array = None + + if name == alias.data and array is not None: + if n_features == 0: + n_features = array.shape[1] + assert n_features == array.shape[1] + + if array is None: + return + + if is_valid: + valid_data[name].append(array) + else: + train_data[name].append(array) + + def append_m_sparse(part: pd.DataFrame, name: str, is_valid: bool) -> None: + nonlocal n_features + + if name == alias.data or name in part.columns: + if name == alias.data: + array = _read_csr_matrix_from_unwrapped_spark_vec(part) + if n_features == 0: + n_features = array.shape[1] + assert n_features == array.shape[1] + else: + array = part[name] + + if is_valid: + valid_data[name].append(array) + else: + train_data[name].append(array) + + def make(values: Dict[str, List[np.ndarray]], kwargs: Dict[str, Any]) -> DMatrix: + if len(values) == 0: + get_logger("XGBoostPySpark").warning( + "Detected an empty partition in the training data. Consider to enable" + " repartition_random_shuffle" + ) + # We must construct an empty DMatrix to bypass the AllReduce + return DMatrix(data=np.empty((0, 0)), **kwargs) + + data = concat_or_none(values[alias.data]) + label = concat_or_none(values.get(alias.label, None)) + weight = concat_or_none(values.get(alias.weight, None)) + margin = concat_or_none(values.get(alias.margin, None)) + qid = concat_or_none(values.get(alias.qid, None)) + return DMatrix( + data=data, label=label, weight=weight, base_margin=margin, qid=qid, **kwargs + ) + + if enable_sparse_data_optim: + append_fn = append_m_sparse + assert "missing" in kwargs and kwargs["missing"] == 0.0 + else: + append_fn = append_m + + def split_params() -> Tuple[Dict[str, Any], Dict[str, Union[int, float, bool]]]: + # FIXME(jiamingy): we really need a better way to bridge distributed frameworks + # to XGBoost native interface and prevent scattering parameters like this. + + # parameters that are not related to data. + non_data_keys = ( + "max_bin", + "missing", + "silent", + "nthread", + "enable_categorical", + ) + non_data_params = {} + meta = {} + for k, v in kwargs.items(): + if k in non_data_keys: + non_data_params[k] = v + else: + meta[k] = v + return meta, non_data_params + + meta, params = split_params() + + if feature_cols is not None and use_qdm: + cache_partitions(iterator, append_fn) + dtrain: DMatrix = make_qdm(train_data, dev_ordinal, meta, None, params) + elif feature_cols is not None and not use_qdm: + cache_partitions(iterator, append_fn) + dtrain = make(train_data, kwargs) + elif feature_cols is None and use_qdm: + cache_partitions(iterator, append_fn) + dtrain = make_qdm(train_data, dev_ordinal, meta, None, params) + else: + cache_partitions(iterator, append_fn) + dtrain = make(train_data, kwargs) + + # Using has_validation_col here to indicate if there is validation col + # instead of getting it from iterator, since the iterator may be empty + # in some special case. That is to say, we must ensure every worker + # construct DMatrix even there is no data since we need to ensure every + # worker do the AllReduce when constructing DMatrix, or else it may hang + # forever. + if has_validation_col: + if use_qdm: + dvalid: Optional[DMatrix] = make_qdm( + valid_data, dev_ordinal, meta, dtrain, params + ) + else: + dvalid = make(valid_data, kwargs) if has_validation_col else None + else: + dvalid = None + + if dvalid is not None: + assert dvalid.num_col() == dtrain.num_col() + + return dtrain, dvalid + + +def pred_contribs( + model: XGBModel, + data: ArrayLike, + base_margin: Optional[ArrayLike] = None, + strict_shape: bool = False, +) -> np.ndarray: + """Predict contributions with data with the full model.""" + iteration_range = model._get_iteration_range(None) + data_dmatrix = DMatrix( + data, + base_margin=base_margin, + missing=model.missing, + nthread=model.n_jobs, + feature_types=model.feature_types, + feature_weights=model.feature_weights, + enable_categorical=model.enable_categorical, + ) + return model.get_booster().predict( + data_dmatrix, + pred_contribs=True, + validate_features=False, + iteration_range=iteration_range, + strict_shape=strict_shape, + ) diff --git a/v12x/pylib/xgboost/spark/estimator.py b/v12x/pylib/xgboost/spark/estimator.py new file mode 100644 index 0000000000000000000000000000000000000000..f6a0d3a40198b50eb8626109e79e23a9efedb583 --- /dev/null +++ b/v12x/pylib/xgboost/spark/estimator.py @@ -0,0 +1,613 @@ +"""Xgboost pyspark integration submodule for estimator API.""" + +# pylint: disable=protected-access, no-member +# pylint: disable=unused-argument, too-many-locals + +from typing import Any, List, Optional, Type, Union + +import numpy as np +from pyspark import keyword_only +from pyspark.ml.param import Param, Params +from pyspark.ml.param.shared import HasProbabilityCol, HasRawPredictionCol +from pyspark.sql import SparkSession + +from ..collective import Config +from ..sklearn import XGBClassifier, XGBRanker, XGBRegressor +from .core import ( # type: ignore[attr-defined] + _ClassificationModel, + _SparkXGBEstimator, + _SparkXGBModel, +) +from .utils import get_class_name + + +def _set_pyspark_xgb_cls_param_attrs( + estimator: Type[_SparkXGBEstimator], model: Type[_SparkXGBModel] +) -> None: + """This function automatically infer to xgboost parameters and set them + into corresponding pyspark estimators and models""" + params_dict = estimator._get_xgb_params_default() + + def param_value_converter(v: Any) -> Any: + if isinstance(v, np.generic): + # convert numpy scalar values to corresponding python scalar values + return np.array(v).item() + if isinstance(v, dict): + return {k: param_value_converter(nv) for k, nv in v.items()} + if isinstance(v, list): + return [param_value_converter(nv) for nv in v] + return v + + def set_param_attrs(attr_name: str, param: Param) -> None: + param.typeConverter = param_value_converter + setattr(estimator, attr_name, param) + setattr(model, attr_name, param) + + for name in params_dict.keys(): + doc = ( + f"Refer to XGBoost doc of " + f"{get_class_name(estimator._xgb_cls())} for this param {name}" + ) + + param_obj: Param = Param(Params._dummy(), name=name, doc=doc) + set_param_attrs(name, param_obj) + + fit_params_dict = estimator._get_fit_params_default() + for name in fit_params_dict.keys(): + doc = ( + f"Refer to XGBoost doc of {get_class_name(estimator._xgb_cls())}" + f".fit() for this param {name}" + ) + if name == "callbacks": + doc += ( + "The callbacks can be arbitrary functions. It is saved using cloudpickle " + "which is not a fully self-contained format. It may fail to load with " + "different versions of dependencies." + ) + param_obj = Param(Params._dummy(), name=name, doc=doc) + set_param_attrs(name, param_obj) + + predict_params_dict = estimator._get_predict_params_default() + for name in predict_params_dict.keys(): + doc = ( + f"Refer to XGBoost doc of {get_class_name(estimator._xgb_cls())}" + f".predict() for this param {name}" + ) + param_obj = Param(Params._dummy(), name=name, doc=doc) + set_param_attrs(name, param_obj) + + +class SparkXGBRegressor(_SparkXGBEstimator): + """SparkXGBRegressor is a PySpark ML estimator. It implements the XGBoost regression + algorithm based on XGBoost python library, and it can be used in PySpark Pipeline + and PySpark ML meta algorithms like + - :py:class:`~pyspark.ml.tuning.CrossValidator`/ + - :py:class:`~pyspark.ml.tuning.TrainValidationSplit`/ + - :py:class:`~pyspark.ml.classification.OneVsRest` + + SparkXGBRegressor automatically supports most of the parameters in + :py:class:`xgboost.XGBRegressor` constructor and most of the parameters used in + :py:meth:`xgboost.XGBRegressor.fit` and :py:meth:`xgboost.XGBRegressor.predict` + method. + + To enable GPU support, set `device` to `cuda` or `gpu`. + + SparkXGBRegressor doesn't support setting `base_margin` explicitly as well, but + support another param called `base_margin_col`. see doc below for more details. + + SparkXGBRegressor doesn't support `validate_features` and `output_margin` param. + + SparkXGBRegressor doesn't support setting `nthread` xgboost param, instead, the + `nthread` param for each xgboost worker will be set equal to `spark.task.cpus` + config value. + + + Parameters + ---------- + + features_col: + When the value is string, it requires the features column name to be vector type. + When the value is a list of string, it requires all the feature columns to be numeric types. + label_col: + Label column name. Default to "label". + prediction_col: + Prediction column name. Default to "prediction" + pred_contrib_col: + Contribution prediction column name. + validation_indicator_col: + For params related to `xgboost.XGBRegressor` training with + evaluation dataset's supervision, + set :py:attr:`xgboost.spark.SparkXGBRegressor.validation_indicator_col` + parameter instead of setting the `eval_set` parameter in `xgboost.XGBRegressor` + fit method. + weight_col: + To specify the weight of the training and validation dataset, set + :py:attr:`xgboost.spark.SparkXGBRegressor.weight_col` parameter instead of setting + `sample_weight` and `sample_weight_eval_set` parameter in `xgboost.XGBRegressor` + fit method. + base_margin_col: + To specify the base margins of the training and validation + dataset, set :py:attr:`xgboost.spark.SparkXGBRegressor.base_margin_col` parameter + instead of setting `base_margin` and `base_margin_eval_set` in the + `xgboost.XGBRegressor` fit method. + + num_workers: + How many XGBoost workers to be used to train. + Each XGBoost worker corresponds to one spark task. + device: + + .. versionadded:: 2.0.0 + + Device for XGBoost workers, available options are `cpu`, `cuda`, and `gpu`. + + force_repartition: + Boolean value to specify if forcing the input dataset to be repartitioned + before XGBoost training. + repartition_random_shuffle: + Boolean value to specify if randomly shuffling the dataset when repartitioning is required. + enable_sparse_data_optim: + Boolean value to specify if enabling sparse data optimization, if True, + Xgboost DMatrix object will be constructed from sparse matrix instead of + dense matrix. + launch_tracker_on_driver: + Boolean value to indicate whether the tracker should be launched on the driver side or + the executor side. + coll_cfg: + The collective configuration. See :py:class:`~xgboost.collective.Config` + + kwargs: + A dictionary of xgboost parameters, please refer to + https://xgboost.readthedocs.io/en/stable/parameter.html + + Note + ---- + + The Parameters chart above contains parameters that need special handling. + For a full list of parameters, see entries with `Param(parent=...` below. + + This API is experimental. + + + Examples + -------- + + >>> from xgboost.spark import SparkXGBRegressor + >>> from pyspark.ml.linalg import Vectors + >>> df_train = spark.createDataFrame([ + ... (Vectors.dense(1.0, 2.0, 3.0), 0, False, 1.0), + ... (Vectors.sparse(3, {1: 1.0, 2: 5.5}), 1, False, 2.0), + ... (Vectors.dense(4.0, 5.0, 6.0), 2, True, 1.0), + ... (Vectors.sparse(3, {1: 6.0, 2: 7.5}), 3, True, 2.0), + ... ], ["features", "label", "isVal", "weight"]) + >>> df_test = spark.createDataFrame([ + ... (Vectors.dense(1.0, 2.0, 3.0), ), + ... (Vectors.sparse(3, {1: 1.0, 2: 5.5}), ) + ... ], ["features"]) + >>> xgb_regressor = SparkXGBRegressor(max_depth=5, missing=0.0, + ... validation_indicator_col='isVal', weight_col='weight', + ... early_stopping_rounds=1, eval_metric='rmse') + >>> xgb_reg_model = xgb_regressor.fit(df_train) + >>> xgb_reg_model.transform(df_test) + + """ + + @keyword_only + def __init__( # pylint:disable=too-many-arguments + self, + *, + features_col: Union[str, List[str]] = "features", + label_col: str = "label", + prediction_col: str = "prediction", + pred_contrib_col: Optional[str] = None, + validation_indicator_col: Optional[str] = None, + weight_col: Optional[str] = None, + base_margin_col: Optional[str] = None, + num_workers: int = 1, + device: Optional[str] = None, + force_repartition: bool = False, + repartition_random_shuffle: bool = False, + enable_sparse_data_optim: bool = False, + launch_tracker_on_driver: Optional[bool] = None, + coll_cfg: Optional[Config] = None, + **kwargs: Any, + ) -> None: + super().__init__() + input_kwargs = self._input_kwargs + self.setParams(**input_kwargs) + + @classmethod + def _xgb_cls(cls) -> Type[XGBRegressor]: + return XGBRegressor + + @classmethod + def _pyspark_model_cls(cls) -> Type["SparkXGBRegressorModel"]: + return SparkXGBRegressorModel + + def _validate_params(self, spark_session: SparkSession) -> None: + super()._validate_params(spark_session) + if self.isDefined(self.qid_col): + raise ValueError( + "Spark Xgboost regressor estimator does not support `qid_col` param." + ) + + +class SparkXGBRegressorModel(_SparkXGBModel): + """ + The model returned by :func:`xgboost.spark.SparkXGBRegressor.fit` + + .. Note:: This API is experimental. + """ + + @classmethod + def _xgb_cls(cls) -> Type[XGBRegressor]: + return XGBRegressor + + +_set_pyspark_xgb_cls_param_attrs(SparkXGBRegressor, SparkXGBRegressorModel) + + +class SparkXGBClassifier(_SparkXGBEstimator, HasProbabilityCol, HasRawPredictionCol): + """SparkXGBClassifier is a PySpark ML estimator. It implements the XGBoost + classification algorithm based on XGBoost python library, and it can be used in + PySpark Pipeline and PySpark ML meta algorithms like + - :py:class:`~pyspark.ml.tuning.CrossValidator`/ + - :py:class:`~pyspark.ml.tuning.TrainValidationSplit`/ + - :py:class:`~pyspark.ml.classification.OneVsRest` + + SparkXGBClassifier automatically supports most of the parameters in + :py:class:`xgboost.XGBClassifier` constructor and most of the parameters used in + :py:meth:`xgboost.XGBClassifier.fit` and :py:meth:`xgboost.XGBClassifier.predict` + method. + + To enable GPU support, set `device` to `cuda` or `gpu`. + + SparkXGBClassifier doesn't support setting `base_margin` explicitly as well, but + support another param called `base_margin_col`. see doc below for more details. + + SparkXGBClassifier doesn't support setting `output_margin`, but we can get output + margin from the raw prediction column. See `raw_prediction_col` param doc below for + more details. + + SparkXGBClassifier doesn't support `validate_features` and `output_margin` param. + + SparkXGBClassifier doesn't support setting `nthread` xgboost param, instead, the + `nthread` param for each xgboost worker will be set equal to `spark.task.cpus` + config value. + + + Parameters + ---------- + + features_col: + When the value is string, it requires the features column name to be vector type. + When the value is a list of string, it requires all the feature columns to be numeric types. + label_col: + Label column name. Default to "label". + prediction_col: + Prediction column name. Default to "prediction" + probability_col: + Column name for predicted class conditional probabilities. Default to probabilityCol + raw_prediction_col: + The `output_margin=True` is implicitly supported by the + `rawPredictionCol` output column, which is always returned with the predicted margin + values. + pred_contrib_col: + Contribution prediction column name. + validation_indicator_col: + For params related to `xgboost.XGBClassifier` training with + evaluation dataset's supervision, + set :py:attr:`xgboost.spark.SparkXGBClassifier.validation_indicator_col` + parameter instead of setting the `eval_set` parameter in `xgboost.XGBClassifier` + fit method. + weight_col: + To specify the weight of the training and validation dataset, set + :py:attr:`xgboost.spark.SparkXGBClassifier.weight_col` parameter instead of setting + `sample_weight` and `sample_weight_eval_set` parameter in `xgboost.XGBClassifier` + fit method. + base_margin_col: + To specify the base margins of the training and validation + dataset, set :py:attr:`xgboost.spark.SparkXGBClassifier.base_margin_col` parameter + instead of setting `base_margin` and `base_margin_eval_set` in the + `xgboost.XGBClassifier` fit method. + + num_workers: + How many XGBoost workers to be used to train. + Each XGBoost worker corresponds to one spark task. + device: + + .. versionadded:: 2.0.0 + + Device for XGBoost workers, available options are `cpu`, `cuda`, and `gpu`. + + force_repartition: + Boolean value to specify if forcing the input dataset to be repartitioned + before XGBoost training. + repartition_random_shuffle: + Boolean value to specify if randomly shuffling the dataset when repartitioning is required. + enable_sparse_data_optim: + Boolean value to specify if enabling sparse data optimization, if True, + Xgboost DMatrix object will be constructed from sparse matrix instead of + dense matrix. + launch_tracker_on_driver: + Boolean value to indicate whether the tracker should be launched on the driver side or + the executor side. + coll_cfg: + The collective configuration. See :py:class:`~xgboost.collective.Config` + + kwargs: + A dictionary of xgboost parameters, please refer to + https://xgboost.readthedocs.io/en/stable/parameter.html + + Note + ---- + + The Parameters chart above contains parameters that need special handling. + For a full list of parameters, see entries with `Param(parent=...` below. + + This API is experimental. + + Examples + -------- + + >>> from xgboost.spark import SparkXGBClassifier + >>> from pyspark.ml.linalg import Vectors + >>> df_train = spark.createDataFrame([ + ... (Vectors.dense(1.0, 2.0, 3.0), 0, False, 1.0), + ... (Vectors.sparse(3, {1: 1.0, 2: 5.5}), 1, False, 2.0), + ... (Vectors.dense(4.0, 5.0, 6.0), 0, True, 1.0), + ... (Vectors.sparse(3, {1: 6.0, 2: 7.5}), 1, True, 2.0), + ... ], ["features", "label", "isVal", "weight"]) + >>> df_test = spark.createDataFrame([ + ... (Vectors.dense(1.0, 2.0, 3.0), ), + ... ], ["features"]) + >>> xgb_classifier = SparkXGBClassifier(max_depth=5, missing=0.0, + ... validation_indicator_col='isVal', weight_col='weight', + ... early_stopping_rounds=1, eval_metric='logloss') + >>> xgb_clf_model = xgb_classifier.fit(df_train) + >>> xgb_clf_model.transform(df_test).show() + + """ + + @keyword_only + def __init__( # pylint:disable=too-many-arguments + self, + *, + features_col: Union[str, List[str]] = "features", + label_col: str = "label", + prediction_col: str = "prediction", + probability_col: str = "probability", + raw_prediction_col: str = "rawPrediction", + pred_contrib_col: Optional[str] = None, + validation_indicator_col: Optional[str] = None, + weight_col: Optional[str] = None, + base_margin_col: Optional[str] = None, + num_workers: int = 1, + device: Optional[str] = None, + force_repartition: bool = False, + repartition_random_shuffle: bool = False, + enable_sparse_data_optim: bool = False, + launch_tracker_on_driver: Optional[bool] = None, + coll_cfg: Optional[Config] = None, + **kwargs: Any, + ) -> None: + super().__init__() + # The default 'objective' param value comes from sklearn `XGBClassifier` ctor, + # but in pyspark we will automatically set objective param depending on + # binary or multinomial input dataset, and we need to remove the fixed default + # param value as well to avoid causing ambiguity. + input_kwargs = self._input_kwargs + self.setParams(**input_kwargs) + self._setDefault(objective=None) + + @classmethod + def _xgb_cls(cls) -> Type[XGBClassifier]: + return XGBClassifier + + @classmethod + def _pyspark_model_cls(cls) -> Type["SparkXGBClassifierModel"]: + return SparkXGBClassifierModel + + def _validate_params(self, spark_session: SparkSession) -> None: + super()._validate_params(spark_session) + if self.isDefined(self.qid_col): + raise ValueError( + "Spark Xgboost classifier estimator does not support `qid_col` param." + ) + if self.getOrDefault("objective"): # pylint: disable=no-member + raise ValueError( + "Setting custom 'objective' param is not allowed in 'SparkXGBClassifier'." + ) + + +class SparkXGBClassifierModel(_ClassificationModel): + """ + The model returned by :func:`xgboost.spark.SparkXGBClassifier.fit` + + .. Note:: This API is experimental. + """ + + @classmethod + def _xgb_cls(cls) -> Type[XGBClassifier]: + return XGBClassifier + + +_set_pyspark_xgb_cls_param_attrs(SparkXGBClassifier, SparkXGBClassifierModel) + + +class SparkXGBRanker(_SparkXGBEstimator): + """SparkXGBRanker is a PySpark ML estimator. It implements the XGBoost + ranking algorithm based on XGBoost python library, and it can be used in + PySpark Pipeline and PySpark ML meta algorithms like + :py:class:`~pyspark.ml.tuning.CrossValidator`/ + :py:class:`~pyspark.ml.tuning.TrainValidationSplit`/ + :py:class:`~pyspark.ml.classification.OneVsRest` + + SparkXGBRanker automatically supports most of the parameters in + :py:class:`xgboost.XGBRanker` constructor and most of the parameters used in + :py:meth:`xgboost.XGBRanker.fit` and :py:meth:`xgboost.XGBRanker.predict` method. + + To enable GPU support, set `device` to `cuda` or `gpu`. + + SparkXGBRanker doesn't support setting `base_margin` explicitly as well, but support + another param called `base_margin_col`. see doc below for more details. + + SparkXGBRanker doesn't support setting `output_margin`, but we can get output margin + from the raw prediction column. See `raw_prediction_col` param doc below for more + details. + + SparkXGBRanker doesn't support `validate_features` and `output_margin` param. + + SparkXGBRanker doesn't support setting `nthread` xgboost param, instead, the + `nthread` param for each xgboost worker will be set equal to `spark.task.cpus` + config value. + + + Parameters + ---------- + + features_col: + When the value is string, it requires the features column name to be vector type. + When the value is a list of string, it requires all the feature columns to be numeric types. + label_col: + Label column name. Default to "label". + prediction_col: + Prediction column name. Default to "prediction" + pred_contrib_col: + Contribution prediction column name. + validation_indicator_col: + For params related to `xgboost.XGBRanker` training with + evaluation dataset's supervision, + set :py:attr:`xgboost.spark.SparkXGBRanker.validation_indicator_col` + parameter instead of setting the `eval_set` parameter in :py:class:`xgboost.XGBRanker` + fit method. + weight_col: + To specify the weight of the training and validation dataset, set + :py:attr:`xgboost.spark.SparkXGBRanker.weight_col` parameter instead of setting + `sample_weight` and `sample_weight_eval_set` parameter in :py:class:`xgboost.XGBRanker` + fit method. + base_margin_col: + To specify the base margins of the training and validation + dataset, set :py:attr:`xgboost.spark.SparkXGBRanker.base_margin_col` parameter + instead of setting `base_margin` and `base_margin_eval_set` in the + :py:class:`xgboost.XGBRanker` fit method. + qid_col: + Query id column name. + num_workers: + How many XGBoost workers to be used to train. + Each XGBoost worker corresponds to one spark task. + device: + + .. versionadded:: 2.0.0 + + Device for XGBoost workers, available options are `cpu`, `cuda`, and `gpu`. + + force_repartition: + Boolean value to specify if forcing the input dataset to be repartitioned + before XGBoost training. + repartition_random_shuffle: + Boolean value to specify if randomly shuffling the dataset when repartitioning is required. + enable_sparse_data_optim: + Boolean value to specify if enabling sparse data optimization, if True, + Xgboost DMatrix object will be constructed from sparse matrix instead of + dense matrix. + launch_tracker_on_driver: + Boolean value to indicate whether the tracker should be launched on the driver side or + the executor side. + coll_cfg: + The collective configuration. See :py:class:`~xgboost.collective.Config` + + kwargs: + A dictionary of xgboost parameters, please refer to + https://xgboost.readthedocs.io/en/stable/parameter.html + + .. Note:: The Parameters chart above contains parameters that need special handling. + For a full list of parameters, see entries with `Param(parent=...` below. + + .. Note:: This API is experimental. + + Examples + -------- + + >>> from xgboost.spark import SparkXGBRanker + >>> from pyspark.ml.linalg import Vectors + >>> ranker = SparkXGBRanker(qid_col="qid") + >>> df_train = spark.createDataFrame( + ... [ + ... (Vectors.dense(1.0, 2.0, 3.0), 0, 0), + ... (Vectors.dense(4.0, 5.0, 6.0), 1, 0), + ... (Vectors.dense(9.0, 4.0, 8.0), 2, 0), + ... (Vectors.sparse(3, {1: 1.0, 2: 5.5}), 0, 1), + ... (Vectors.sparse(3, {1: 6.0, 2: 7.5}), 1, 1), + ... (Vectors.sparse(3, {1: 8.0, 2: 9.5}), 2, 1), + ... ], + ... ["features", "label", "qid"], + ... ) + >>> df_test = spark.createDataFrame( + ... [ + ... (Vectors.dense(1.5, 2.0, 3.0), 0), + ... (Vectors.dense(4.5, 5.0, 6.0), 0), + ... (Vectors.dense(9.0, 4.5, 8.0), 0), + ... (Vectors.sparse(3, {1: 1.0, 2: 6.0}), 1), + ... (Vectors.sparse(3, {1: 6.0, 2: 7.0}), 1), + ... (Vectors.sparse(3, {1: 8.0, 2: 10.5}), 1), + ... ], + ... ["features", "qid"], + ... ) + >>> model = ranker.fit(df_train) + >>> model.transform(df_test).show() + """ + + @keyword_only + def __init__( # pylint:disable=too-many-arguments + self, + *, + features_col: Union[str, List[str]] = "features", + label_col: str = "label", + prediction_col: str = "prediction", + pred_contrib_col: Optional[str] = None, + validation_indicator_col: Optional[str] = None, + weight_col: Optional[str] = None, + base_margin_col: Optional[str] = None, + qid_col: Optional[str] = None, + num_workers: int = 1, + device: Optional[str] = None, + force_repartition: bool = False, + repartition_random_shuffle: bool = False, + enable_sparse_data_optim: bool = False, + launch_tracker_on_driver: Optional[bool] = None, + coll_cfg: Optional[Config] = None, + **kwargs: Any, + ) -> None: + super().__init__() + input_kwargs = self._input_kwargs + self.setParams(**input_kwargs) + + @classmethod + def _xgb_cls(cls) -> Type[XGBRanker]: + return XGBRanker + + @classmethod + def _pyspark_model_cls(cls) -> Type["SparkXGBRankerModel"]: + return SparkXGBRankerModel + + def _validate_params(self, spark_session: SparkSession) -> None: + super()._validate_params(spark_session) + if not self.isDefined(self.qid_col): + raise ValueError( + "Spark Xgboost ranker estimator requires setting `qid_col` param." + ) + + +class SparkXGBRankerModel(_SparkXGBModel): + """ + The model returned by :func:`xgboost.spark.SparkXGBRanker.fit` + + .. Note:: This API is experimental. + """ + + @classmethod + def _xgb_cls(cls) -> Type[XGBRanker]: + return XGBRanker + + +_set_pyspark_xgb_cls_param_attrs(SparkXGBRanker, SparkXGBRankerModel) diff --git a/v12x/pylib/xgboost/spark/params.py b/v12x/pylib/xgboost/spark/params.py new file mode 100644 index 0000000000000000000000000000000000000000..af2f4f9d658873fbb70099d87bb48fd8adc242b6 --- /dev/null +++ b/v12x/pylib/xgboost/spark/params.py @@ -0,0 +1,104 @@ +"""Xgboost pyspark integration submodule for params.""" + +from typing import Dict + +from pyspark.ml.param import TypeConverters +from pyspark.ml.param.shared import Param, Params + + +class HasArbitraryParamsDict(Params): + """ + This is a Params based class that is extended by _SparkXGBParams + and holds the variable to store the **kwargs parts of the XGBoost + input. + """ + + arbitrary_params_dict: "Param[Dict]" = Param( + Params._dummy(), + "arbitrary_params_dict", + "arbitrary_params_dict This parameter holds all of the additional parameters which are " + "not exposed as the XGBoost Spark estimator params but can be recognized by " + "underlying XGBoost library. It is stored as a dictionary.", + ) + + +class HasBaseMarginCol(Params): + """ + This is a Params based class that is extended by _SparkXGBParams + and holds the variable to store the base margin column part of XGboost. + """ + + base_margin_col = Param( + Params._dummy(), + "base_margin_col", + "This stores the name for the column of the base margin", + typeConverter=TypeConverters.toString, + ) + + +class HasFeaturesCols(Params): + """ + Mixin for param features_cols: a list of feature column names. + This parameter is taken effect only when GPU is enabled. + """ + + features_cols = Param( + Params._dummy(), + "features_cols", + "feature column names.", + typeConverter=TypeConverters.toListString, + ) + + def __init__(self) -> None: + super().__init__() + self._setDefault(features_cols=[]) + + +class HasEnableSparseDataOptim(Params): + """ + This is a Params based class that is extended by _SparkXGBParams + and holds the variable to store the boolean config of enabling sparse data optimization. + """ + + enable_sparse_data_optim = Param( + Params._dummy(), + "enable_sparse_data_optim", + "This stores the boolean config of enabling sparse data optimization, if enabled, " + "Xgboost DMatrix object will be constructed from sparse matrix instead of " + "dense matrix. This config is disabled by default. If most of examples in your " + "training dataset contains sparse features, we suggest to enable this config.", + typeConverter=TypeConverters.toBoolean, + ) + + def __init__(self) -> None: + super().__init__() + self._setDefault(enable_sparse_data_optim=False) + + +class HasQueryIdCol(Params): + """ + Mixin for param qid_col: query id column name. + """ + + qid_col = Param( + Params._dummy(), + "qid_col", + "query id column name", + typeConverter=TypeConverters.toString, + ) + + +class HasContribPredictionCol(Params): + """ + Mixin for param pred_contrib_col: contribution prediction column name. + + Output is a 3-dim array, with (rows, groups, columns + 1) for classification case. + Else, it can be a 2 dimension for regression case. + """ + + pred_contrib_col: "Param[str]" = Param( + Params._dummy(), + "pred_contrib_col", + "feature contributions to individual predictions.", + typeConverter=TypeConverters.toString, + ) diff --git a/v12x/pylib/xgboost/spark/summary.py b/v12x/pylib/xgboost/spark/summary.py new file mode 100644 index 0000000000000000000000000000000000000000..9ee7207396113dba2ecf4c468ab19a92f572f70d --- /dev/null +++ b/v12x/pylib/xgboost/spark/summary.py @@ -0,0 +1,43 @@ +"""Xgboost training summary integration submodule.""" + +from dataclasses import dataclass, field +from typing import Dict, List + + +@dataclass +class XGBoostTrainingSummary: + """ + A class that holds the training and validation objective history + of an XGBoost model during its training process. + """ + + train_objective_history: Dict[str, List[float]] = field(default_factory=dict) + validation_objective_history: Dict[str, List[float]] = field(default_factory=dict) + + @staticmethod + def from_metrics( + metrics: Dict[str, Dict[str, List[float]]], + ) -> "XGBoostTrainingSummary": + """ + Create an XGBoostTrainingSummary instance from a nested dictionary of metrics. + + Parameters + ---------- + metrics : dict of str to dict of str to list of float + A dictionary containing training and validation metrics. + Example format: + { + "training": {"logloss": [0.1, 0.08]}, + "validation": {"logloss": [0.12, 0.1]} + } + + Returns + ------- + A new instance of XGBoostTrainingSummary. + + """ + train_objective_history = metrics.get("training", {}) + validation_objective_history = metrics.get("validation", {}) + return XGBoostTrainingSummary( + train_objective_history, validation_objective_history + ) diff --git a/v12x/pylib/xgboost/spark/utils.py b/v12x/pylib/xgboost/spark/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..71ba2e844e7cfa38d78a385f86bdcea2d4876e60 --- /dev/null +++ b/v12x/pylib/xgboost/spark/utils.py @@ -0,0 +1,188 @@ +"""Xgboost pyspark integration submodule for helper functions.""" + +# pylint: disable=fixme + +import inspect +import logging +import sys +from threading import Thread +from typing import Any, Callable, Dict, Optional, Set, Type, Union + +import pyspark +from pyspark import BarrierTaskContext, TaskContext +from pyspark.sql import SparkSession + +from ..collective import CommunicatorContext as CCtx +from ..collective import Config +from ..collective import _Args as CollArgs +from ..collective import _ArgVals as CollArgsVals +from ..core import Booster +from ..sklearn import XGBModel +from ..tracker import RabitTracker + + +def get_class_name(cls: Type) -> str: + """Return the class name.""" + return f"{cls.__module__}.{cls.__name__}" + + +def _get_default_params_from_func( + func: Callable, unsupported_set: Set[str] +) -> Dict[str, Any]: + """Returns a dictionary of parameters and their default value of function fn. Only + the parameters with a default value will be included. + + """ + sig = inspect.signature(func) + filtered_params_dict = {} + for parameter in sig.parameters.values(): + # Remove parameters without a default value and those in the unsupported_set + if ( + parameter.default is not parameter.empty + and parameter.name not in unsupported_set + ): + filtered_params_dict[parameter.name] = parameter.default + return filtered_params_dict + + +class CommunicatorContext(CCtx): + """Context with PySpark specific task ID.""" + + def __init__(self, context: BarrierTaskContext, **args: CollArgsVals) -> None: + args["dmlc_task_id"] = str(context.partitionId()) + super().__init__(**args) + + +def _start_tracker(host: str, n_workers: int, port: int = 0) -> CollArgs: + """Start Rabit tracker with n_workers""" + args: CollArgs = {"n_workers": n_workers} + tracker = RabitTracker(n_workers=n_workers, host_ip=host, sortby="task", port=port) + tracker.start() + thread = Thread(target=tracker.wait_for) + thread.daemon = True + thread.start() + args.update(tracker.worker_args()) + return args + + +def _get_rabit_args(conf: Config, n_workers: int) -> CollArgs: + """Get rabit context arguments to send to each worker.""" + assert conf.tracker_host_ip is not None + port = 0 if conf.tracker_port is None else conf.tracker_port + env = _start_tracker(conf.tracker_host_ip, n_workers, port) + return env + + +def _get_host_ip(context: BarrierTaskContext) -> str: + """Gets the hostIP for Spark. This essentially gets the IP of the first worker.""" + task_ip_list = [info.address.split(":")[0] for info in context.getTaskInfos()] + return task_ip_list[0] + + +def get_logger(name: str, level: Optional[Union[str, int]] = None) -> logging.Logger: + """Gets a logger by name, or creates and configures it for the first time.""" + logger = logging.getLogger(name) + if level is not None: + logger.setLevel(level) + else: + # Default to info if not set. + if logger.level == logging.NOTSET: + logger.setLevel(logging.INFO) + # If the logger is configured, skip the configure + if not logger.handlers and not logging.getLogger().handlers: + handler = logging.StreamHandler(sys.stderr) + formatter = logging.Formatter( + "%(asctime)s %(levelname)s %(name)s: %(funcName)s %(message)s" + ) + handler.setFormatter(formatter) + logger.addHandler(handler) + return logger + + +def get_logger_level(name: str) -> Optional[int]: + """Get the logger level for the given log name""" + logger = logging.getLogger(name) + return None if logger.level == logging.NOTSET else logger.level + + +def _get_max_num_concurrent_tasks(spark_session: SparkSession) -> int: + """Gets the current max number of concurrent tasks.""" + + # In Spark Connect, we cannot easily get the max number of concurrent tasks + # from the client side without accessing internal APIs or executing a task. + # For now, we return a large number to skip the check. + if _is_connect(spark_session): + return sys.maxsize + + # pylint: disable=protected-access + return spark_session.sparkContext._jsc.sc().maxNumConcurrentTasks( + spark_session.sparkContext._jsc.sc() + .resourceProfileManager() + .resourceProfileFromId(0) + ) + + +def _is_connect(spark_session: SparkSession) -> bool: + try: + return isinstance(spark_session, pyspark.sql.connect.session.SparkSession) + except AttributeError: + return False + + +def _is_local(spark_session: SparkSession) -> bool: + """Whether it is Spark local mode""" + # In Spark Connect, we check the spark.master configuration if available. + # Note: This might not be accurate if spark.master is not set in RuntimeConfig. + master = spark_session.conf.get("spark.master", None) + return master is not None and (master == "local" or master.startswith("local[")) + + +def _get_gpu_id(task_context: TaskContext) -> int: + """Get the gpu id from the task resources""" + if task_context is None: + # This is a safety check. + raise RuntimeError("_get_gpu_id should not be invoked from driver side.") + resources = task_context.resources() + if "gpu" not in resources: + raise RuntimeError( + "Couldn't get the gpu id, Please check the GPU resource configuration" + ) + # return the first gpu id. + return int(resources["gpu"].addresses[0].strip()) + + +def deserialize_xgb_model( + model: str, xgb_model_creator: Callable[[], XGBModel] +) -> XGBModel: + """ + Deserialize an xgboost.XGBModel instance from the input model. + """ + xgb_model = xgb_model_creator() + xgb_model.load_model(bytearray(model.encode("utf-8"))) + return xgb_model + + +def serialize_booster(booster: Booster) -> str: + """ + Serialize the input booster to a string. + + Parameters + ---------- + booster: + an xgboost.core.Booster instance + """ + return booster.save_raw("json").decode("utf-8") + + +def deserialize_booster(model: str) -> Booster: + """ + Deserialize an xgboost.core.Booster from the input ser_model_string. + """ + booster = Booster() + booster.load_model(bytearray(model.encode("utf-8"))) + return booster + + +def use_cuda(device: Optional[str]) -> bool: + """Whether xgboost is using CUDA workers.""" + return device in ("cuda", "gpu") diff --git a/v12x/pylib/xgboost/testing/__init__.py b/v12x/pylib/xgboost/testing/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..0108b46db97470a33add5c402d60e36a6ce5fda2 --- /dev/null +++ b/v12x/pylib/xgboost/testing/__init__.py @@ -0,0 +1,753 @@ +"""Utilities for defining Python tests. The module is private and subject to frequent +change without notice. + +""" + +# pylint: disable=missing-function-docstring +import importlib.util +import os +import platform +import queue +import socket +import sys +import threading +from contextlib import contextmanager +from io import StringIO +from platform import system +from typing import ( + Any, + Callable, + Dict, + Generator, + List, + Optional, + Set, + Tuple, + TypedDict, + TypeVar, + Union, +) + +import numpy as np +import pytest +from scipy import sparse + +import xgboost as xgb +from xgboost import RabitTracker +from xgboost.core import ArrayLike + +from .._typing import PathLike +from .data import ( + IteratorForTest, + get_california_housing, + get_cancer, + get_digits, + get_sparse, + make_batches, + make_categorical, + make_sparse_regression, +) + +# Used to be defined in this top level module. +from .utils import non_decreasing, non_increasing # NOLINT + +hypothesis = pytest.importorskip("hypothesis") + +# pylint:disable=wrong-import-position,wrong-import-order +from hypothesis import strategies +from hypothesis.extra.numpy import arrays + +datasets = pytest.importorskip("sklearn.datasets") + +PytestSkip = TypedDict("PytestSkip", {"condition": bool, "reason": str}) + + +def has_ipv6() -> bool: + """Check whether IPv6 is enabled on this host.""" + # connection error in macos, still need some fixes. + if system() not in ("Linux", "Windows"): + return False + + if socket.has_ipv6: + try: + with ( + socket.socket(socket.AF_INET6, socket.SOCK_STREAM) as server, + socket.socket(socket.AF_INET6, socket.SOCK_STREAM) as client, + ): + server.bind(("::1", 0)) + port = server.getsockname()[1] + server.listen() + + client.connect(("::1", port)) + conn, _ = server.accept() + + client.sendall("abc".encode()) + msg = conn.recv(3).decode() + # if the code can be executed to this point, the message should be + # correct. + assert msg == "abc" + return True + except OSError: + pass + return False + + +def no_mod(name: str) -> PytestSkip: + spec = importlib.util.find_spec(name) + return {"condition": spec is None, "reason": f"{name} is not installed."} + + +def no_ipv6() -> PytestSkip: + """PyTest skip mark for IPv6.""" + return {"condition": not has_ipv6(), "reason": "IPv6 is required to be enabled."} + + +def not_linux() -> PytestSkip: + return {"condition": system() != "Linux", "reason": "Linux is required."} + + +def no_ubjson() -> PytestSkip: + return no_mod("ubjson") + + +def no_sklearn() -> PytestSkip: + return no_mod("sklearn") + + +def no_dask() -> PytestSkip: + return no_mod("dask") + + +def no_loky() -> PytestSkip: + return no_mod("loky") + + +def no_dask_ml() -> PytestSkip: + if sys.platform.startswith("win"): + return {"reason": "Unsupported platform.", "condition": True} + return no_mod("dask_ml") + + +def no_spark() -> PytestSkip: + if sys.platform.startswith("win") or sys.platform.startswith("darwin"): + return {"reason": "Unsupported platform.", "condition": True} + return no_mod("pyspark") + + +def no_pandas() -> PytestSkip: + return no_mod("pandas") + + +def no_arrow() -> PytestSkip: + return no_mod("pyarrow") + + +def no_polars() -> PytestSkip: + return no_mod("polars") + + +def no_modin() -> PytestSkip: + try: + import modin.pandas as md + + md.DataFrame([[1, 2.0, True], [2, 3.0, False]], columns=["a", "b", "c"]) + + except ImportError: + return {"reason": "Failed import modin.", "condition": True} + return {"reason": "Failed import modin.", "condition": True} + + +def no_matplotlib() -> PytestSkip: + reason = "Matplotlib is not installed." + try: + import matplotlib.pyplot as _ # noqa + + return {"condition": False, "reason": reason} + except ImportError: + return {"condition": True, "reason": reason} + + +def no_dask_cuda() -> PytestSkip: + return no_mod("dask_cuda") + + +def no_cudf() -> PytestSkip: + return no_mod("cudf") + + +def no_cupy() -> PytestSkip: + skip_cupy = no_mod("cupy") + if not skip_cupy["condition"] and system() == "Windows": + import cupy as cp + + # Cupy might run into issue on Windows due to missing compiler + try: + cp.array([1, 2, 3]).sum() + except Exception: # pylint: disable=broad-except + skip_cupy["condition"] = True + return skip_cupy + + +def no_dask_cudf() -> PytestSkip: + return no_mod("dask_cudf") + + +def no_graphviz() -> PytestSkip: + return no_mod("graphviz") + + +def no_rmm() -> PytestSkip: + return no_mod("rmm") + + +def no_multiple(*args: Any) -> PytestSkip: + condition = False + reason = "" + for arg in args: + condition = condition or arg["condition"] + if arg["condition"]: + reason = arg["reason"] + break + return {"condition": condition, "reason": reason} + + +def skip_win() -> PytestSkip: + return {"reason": "Unsupported platform.", "condition": is_windows()} + + +def make_regression( + n_samples: int, n_features: int, use_cupy: bool +) -> Tuple[ArrayLike, ArrayLike, ArrayLike]: + """Make a simple regression dataset.""" + X, y, w = make_batches(n_samples, n_features, 1, use_cupy) + return X[0], y[0], w[0] + + +def make_batches_sparse( + n_samples_per_batch: int, n_features: int, n_batches: int, sparsity: float +) -> Tuple[List[sparse.csr_matrix], List[np.ndarray], List[np.ndarray]]: + X = [] + y = [] + w = [] + rng = np.random.RandomState(1994) + for _ in range(n_batches): + _X = sparse.random( + n_samples_per_batch, + n_features, + 1.0 - sparsity, + format="csr", + dtype=np.float32, + random_state=rng, + ) + _y = rng.randn(n_samples_per_batch) + _w = rng.uniform(low=0, high=1, size=n_samples_per_batch) + X.append(_X) + y.append(_y) + w.append(_w) + return X, y, w + + +class TestDataset: + """Contains a dataset in numpy format as well as the relevant objective and metric.""" + + def __init__( + self, name: str, get_dataset: Callable, objective: str, metric: str + ) -> None: + self.name = name + self.objective = objective + self.metric = metric + self.X, self.y = get_dataset() + self.w: Optional[np.ndarray] = None + self.margin: Optional[np.ndarray] = None + + def set_params(self, params_in: Dict[str, Any]) -> Dict[str, Any]: + params_in["objective"] = self.objective + params_in["eval_metric"] = self.metric + if self.objective == "multi:softmax": + params_in["num_class"] = int(np.max(self.y) + 1) + return params_in + + def get_dmat(self) -> xgb.DMatrix: + return xgb.DMatrix( + self.X, + self.y, + weight=self.w, + base_margin=self.margin, + enable_categorical=True, + ) + + def get_device_dmat(self, max_bin: Optional[int]) -> xgb.QuantileDMatrix: + import cupy as cp + + w = None if self.w is None else cp.array(self.w) + X = cp.array(self.X, dtype=np.float32) + y = cp.array(self.y, dtype=np.float32) + return xgb.QuantileDMatrix( + X, y, weight=w, base_margin=self.margin, max_bin=max_bin + ) + + def get_external_dmat(self) -> xgb.DMatrix: + n_samples = self.X.shape[0] + n_batches = 10 + per_batch = n_samples // n_batches + 1 + + predictor = [] + response = [] + weight = [] + for i in range(n_batches): + beg = i * per_batch + end = min((i + 1) * per_batch, n_samples) + assert end != beg + X = self.X[beg:end, ...] + y = self.y[beg:end] + w = self.w[beg:end] if self.w is not None else None + predictor.append(X) + response.append(y) + if w is not None: + weight.append(w) + + it = IteratorForTest( + predictor, + response, + weight if weight else None, + cache="cache", + on_host=False, + ) + return xgb.DMatrix(it) + + def __repr__(self) -> str: + return self.name + + +def _cat_sampled_from() -> strategies.SearchStrategy: + @strategies.composite + def _make_cat(draw: Callable) -> Tuple[int, int, int, float]: + n_samples = draw(strategies.integers(2, 512)) + n_features = draw(strategies.integers(1, 4)) + n_cats = draw(strategies.integers(1, 128)) + sparsity = draw( + strategies.floats( + min_value=0, + max_value=1, + allow_nan=False, + allow_infinity=False, + allow_subnormal=False, + ) + ) + return n_samples, n_features, n_cats, sparsity + + def _build(args: Tuple[int, int, int, float]) -> TestDataset: + n_samples = args[0] + n_features = args[1] + n_cats = args[2] + sparsity = args[3] + return TestDataset( + f"{n_samples}x{n_features}-{n_cats}-{sparsity}", + lambda: make_categorical( + n_samples, n_features, n_cats, onehot=False, sparsity=sparsity + ), + "reg:squarederror", + "rmse", + ) + + return _make_cat().map(_build) # pylint: disable=no-member + + +categorical_dataset_strategy: strategies.SearchStrategy = _cat_sampled_from() + +sparse_datasets_strategy = strategies.sampled_from( + [ + TestDataset( + "1e5x8-0.95-csr", + lambda: make_sparse_regression(int(1e5), 8, 0.95, False), + "reg:squarederror", + "rmse", + ), + TestDataset( + "1e5x8-0.5-csr", + lambda: make_sparse_regression(int(1e5), 8, 0.5, False), + "reg:squarederror", + "rmse", + ), + TestDataset( + "1e5x8-0.5-dense", + lambda: make_sparse_regression(int(1e5), 8, 0.5, True), + "reg:squarederror", + "rmse", + ), + TestDataset( + "1e5x8-0.05-csr", + lambda: make_sparse_regression(int(1e5), 8, 0.05, False), + "reg:squarederror", + "rmse", + ), + TestDataset( + "1e5x8-0.05-dense", + lambda: make_sparse_regression(int(1e5), 8, 0.05, True), + "reg:squarederror", + "rmse", + ), + ] +) + + +def make_datasets_with_margin( + unweighted_strategy: strategies.SearchStrategy, +) -> Callable[[], strategies.SearchStrategy[TestDataset]]: + """Factory function for creating strategies that generates datasets with weight and + base margin. + + """ + + @strategies.composite + def weight_margin(draw: Callable) -> TestDataset: + data: TestDataset = draw(unweighted_strategy) + if draw(strategies.booleans()): + data.w = draw( + arrays(np.float64, (len(data.y)), elements=strategies.floats(0.1, 2.0)) + ) + if draw(strategies.booleans()): + num_class = 1 + if data.objective == "multi:softmax": + num_class = int(np.max(data.y) + 1) + elif data.name.startswith("mtreg"): + num_class = data.y.shape[1] + + data.margin = draw( + arrays( + np.float64, + (data.y.shape[0] * num_class), + elements=strategies.floats(0.5, 1.0), + ) + ) + assert data.margin is not None + if num_class != 1: + data.margin = data.margin.reshape(data.y.shape[0], num_class) + + return data + + return weight_margin + + +# A strategy for drawing from a set of example datasets. May add random weights to the +# dataset +def make_dataset_strategy() -> strategies.SearchStrategy[TestDataset]: + _unweighted_datasets_strategy = strategies.sampled_from( + [ + TestDataset( + "calif_housing", get_california_housing, "reg:squarederror", "rmse" + ), + TestDataset( + "calif_housing-l1", get_california_housing, "reg:absoluteerror", "mae" + ), + TestDataset("cancer", get_cancer, "binary:logistic", "logloss"), + TestDataset("sparse", get_sparse, "reg:squarederror", "rmse"), + TestDataset("sparse-l1", get_sparse, "reg:absoluteerror", "mae"), + TestDataset( + "empty", + lambda: (np.empty((0, 100)), np.empty(0)), + "reg:squarederror", + "rmse", + ), + ] + ) + return make_datasets_with_margin(_unweighted_datasets_strategy)() + + +_unweighted_multi_datasets_strategy = strategies.sampled_from( + [ + TestDataset("digits", get_digits, "multi:softmax", "mlogloss"), + TestDataset( + "mtreg", + lambda: datasets.make_regression(n_samples=128, n_features=2, n_targets=3), + "reg:squarederror", + "rmse", + ), + TestDataset( + "mtreg-l1", + lambda: datasets.make_regression(n_samples=128, n_features=2, n_targets=3), + "reg:absoluteerror", + "mae", + ), + ] +) + +# A strategy for drawing from a set of multi-target/multi-class datasets. +multi_dataset_strategy = make_datasets_with_margin( + _unweighted_multi_datasets_strategy +)() + + +M = TypeVar("M", xgb.Booster, xgb.XGBModel) + + +def logregobj(preds: np.ndarray, dtrain: xgb.DMatrix) -> Tuple[np.ndarray, np.ndarray]: + """Binary regression custom objective.""" + labels = dtrain.get_label() + preds = 1.0 / (1.0 + np.exp(-preds)) + grad = preds - labels + hess = preds * (1.0 - preds) + return grad, hess + + +def eval_error_metric( + predt: np.ndarray, dtrain: xgb.DMatrix, rev_link: bool +) -> Tuple[str, np.float64]: + """Evaluation metric for xgb.train. + + Parameters + ---------- + rev_link : Whether the metric needs to apply the reverse link function (activation). + + """ + label = dtrain.get_label() + if rev_link: + predt = 1.0 / (1.0 + np.exp(-predt)) + assert (0.0 <= predt).all() and (predt <= 1.0).all() + r = np.zeros(predt.shape) + gt = predt > 0.5 + if predt.size == 0: + return "CustomErr", np.float64(0.0) + r[gt] = 1 - label[gt] + le = predt <= 0.5 + r[le] = label[le] + return "CustomErr", np.sum(r) + + +def eval_error_metric_skl( + y_true: np.ndarray, y_score: np.ndarray, rev_link: bool = False +) -> np.float64: + """Evaluation metric that looks like metrics provided by sklearn.""" + + if rev_link: + y_score = 1.0 / (1.0 + np.exp(-y_score)) + assert (0.0 <= y_score).all() and (y_score <= 1.0).all() + + r = np.zeros(y_score.shape) + gt = y_score > 0.5 + r[gt] = 1 - y_true[gt] + le = y_score <= 0.5 + r[le] = y_true[le] + return np.sum(r) + + +def root_mean_square(y_true: np.ndarray, y_score: np.ndarray) -> float: + err = y_score - y_true + rmse = np.sqrt(np.dot(err, err) / y_score.size) + return rmse + + +def softmax(x: np.ndarray) -> np.ndarray: + e = np.exp(x) + return e / np.sum(e) + + +def softprob_obj( + classes: int, use_cupy: bool = False, order: str = "C", gdtype: str = "float32" +) -> Callable[[np.ndarray, np.ndarray], Tuple[np.ndarray, np.ndarray]]: + """Custom softprob objective for testing. + + Parameters + ---------- + use_cupy : + Whether the objective should return cupy arrays. + order : + The order of gradient matrices. "C" or "F". + gdtype : + DType for gradient. Hessian is not set. This is for testing asymmetric types. + """ + if use_cupy: + import cupy as backend + else: + backend = np + + def objective( + labels: backend.ndarray, predt: backend.ndarray + ) -> Tuple[backend.ndarray, backend.ndarray]: + rows = labels.shape[0] + grad = backend.zeros((rows, classes), dtype=np.float32) + hess = backend.zeros((rows, classes), dtype=np.float32) + eps = 1e-6 + for r in range(predt.shape[0]): + target = labels[r] + p = softmax(predt[r, :]) + for c in range(predt.shape[1]): + assert target >= 0 or target <= classes + g = p[c] - 1.0 if c == target else p[c] + h = max((2.0 * p[c] * (1.0 - p[c])).item(), eps) + grad[r, c] = g + hess[r, c] = h + + grad = grad.reshape((rows, classes)) + hess = hess.reshape((rows, classes)) + grad = backend.require(grad, requirements=order, dtype=gdtype) + hess = backend.require(hess, requirements=order) + return grad, hess + + return objective + + +def ls_obj( + y_true: np.ndarray, y_pred: np.ndarray, sample_weight: Optional[np.ndarray] = None +) -> Tuple[np.ndarray, np.ndarray]: + """Least squared error.""" + grad = y_pred - y_true + hess = np.ones(grad.shape) + if sample_weight is not None: + grad *= sample_weight + hess *= sample_weight + return grad, hess + + +class DirectoryExcursion: + """Change directory. Change back and optionally cleaning up the directory when + exit. + + """ + + def __init__(self, path: PathLike, cleanup: bool = False): + self.path = path + self.curdir = os.path.normpath(os.path.abspath(os.path.curdir)) + self.cleanup = cleanup + self.files: Set[str] = set() + + def __enter__(self) -> None: + os.chdir(self.path) + if self.cleanup: + self.files = { + os.path.join(root, f) + for root, subdir, files in os.walk(os.path.expanduser(self.path)) + for f in files + } + + def __exit__(self, *args: Any) -> None: + os.chdir(self.curdir) + if self.cleanup: + files = { + os.path.join(root, f) + for root, subdir, files in os.walk(os.path.expanduser(self.path)) + for f in files + } + diff = files.difference(self.files) + for f in diff: + os.remove(f) + + +@contextmanager +def captured_output() -> Generator[Tuple[StringIO, StringIO], None, None]: + """Reassign stdout temporarily in order to test printed statements + Taken from: + https://stackoverflow.com/questions/4219717/how-to-assert-output-with-nosetest-unittest-in-python + + Also works for pytest. + + """ + new_out, new_err = StringIO(), StringIO() + old_out, old_err = sys.stdout, sys.stderr + try: + sys.stdout, sys.stderr = new_out, new_err + yield sys.stdout, sys.stderr + finally: + sys.stdout, sys.stderr = old_out, old_err + + +def timeout(sec: int, *args: Any, enable: bool = True, **kwargs: Any) -> Any: + """Make a pytest mark for the `pytest-timeout` package. + + Parameters + ---------- + sec : + Timeout seconds. + enable : + Control whether timeout should be applied, used for debugging. + + Returns + ------- + pytest.mark.timeout + """ + + if enable: + return pytest.mark.timeout(sec, *args, **kwargs) + return pytest.mark.timeout(None, *args, **kwargs) + + +def setup_rmm_pool(_: Any, pytestconfig: pytest.Config) -> None: + if pytestconfig.getoption("--use-rmm-pool"): + if no_rmm()["condition"]: + raise ImportError("The --use-rmm-pool option requires the RMM package") + if no_dask_cuda()["condition"]: + raise ImportError( + "The --use-rmm-pool option requires the dask_cuda package" + ) + import rmm + from dask_cuda.utils import get_n_gpus + + rmm.reinitialize( + pool_allocator=True, + initial_pool_size=1024 * 1024 * 1024, + devices=list(range(get_n_gpus())), + ) + + +def demo_dir(path: str) -> str: + """Look for the demo directory based on the test file name.""" + path = normpath(os.path.dirname(path)) + while True: + subdirs = [f.path for f in os.scandir(path) if f.is_dir()] + subdirs = [os.path.basename(d) for d in subdirs] + if "demo" in subdirs: + return os.path.join(path, "demo") + new_path = normpath(os.path.join(path, os.path.pardir)) + assert new_path != path + path = new_path + + +def normpath(path: str) -> str: + return os.path.normpath(os.path.abspath(path)) + + +def data_dir(path: str) -> str: + return os.path.join(demo_dir(path), "data") + + +def load_agaricus(path: str) -> Tuple[xgb.DMatrix, xgb.DMatrix]: + dpath = data_dir(path) + dtrain = xgb.DMatrix(os.path.join(dpath, "agaricus.txt.train?format=libsvm")) + dtest = xgb.DMatrix(os.path.join(dpath, "agaricus.txt.test?format=libsvm")) + return dtrain, dtest + + +def project_root(path: str) -> str: + return normpath(os.path.join(demo_dir(path), os.path.pardir)) + + +def run_with_rabit( + world_size: int, test_fn: Callable[..., Any], *args: Any, **kwargs: Any +) -> None: + exception_queue: queue.Queue = queue.Queue() + + def run_worker(rabit_env: Dict[str, Union[str, int]]) -> None: + try: + with xgb.collective.CommunicatorContext(**rabit_env): + test_fn(*args, **kwargs) + except Exception as e: # pylint: disable=broad-except + exception_queue.put(e) + + tracker = RabitTracker(host_ip="127.0.0.1", n_workers=world_size) + tracker.start() + + workers = [] + for _ in range(world_size): + worker = threading.Thread(target=run_worker, args=(tracker.worker_args(),)) + workers.append(worker) + worker.start() + for worker in workers: + worker.join() + assert exception_queue.empty(), f"Worker failed: {exception_queue.get()}" + + tracker.wait_for() + + +def is_windows() -> bool: + """Check if the current platform is Windows.""" + return platform.system() == "Windows" diff --git a/v12x/pylib/xgboost/testing/basic_models.py b/v12x/pylib/xgboost/testing/basic_models.py new file mode 100644 index 0000000000000000000000000000000000000000..dae99d7380a92fafa12fd280d3ca21a91d6c1b3e --- /dev/null +++ b/v12x/pylib/xgboost/testing/basic_models.py @@ -0,0 +1,81 @@ +"""Tests for basic features of the Booster.""" + +from typing import Tuple + +import numpy as np + +from xgboost import testing as tm + +from ..core import Booster, DMatrix +from ..training import cv, train +from .utils import Device + + +def run_custom_objective( # pylint: disable=too-many-locals + tree_method: str, + device: Device, + dtrain: DMatrix, + dtest: DMatrix, +) -> None: + """Tests custom objective and metric functions.""" + param = { + "max_depth": 2, + "eta": 1, + "objective": "reg:logistic", + "tree_method": tree_method, + "device": device, + } + watchlist = [(dtest, "eval"), (dtrain, "train")] + num_round = 10 + + def evalerror(preds: np.ndarray, dtrain: DMatrix) -> Tuple[str, np.float64]: + return tm.eval_error_metric(preds, dtrain, rev_link=True) + + # test custom_objective in training + bst = train( + param, + dtrain, + num_round, + evals=watchlist, + obj=tm.logregobj, + custom_metric=evalerror, + ) + assert isinstance(bst, Booster) + preds = bst.predict(dtest) + labels = dtest.get_label() + err = sum(1 for i in range(len(preds)) if int(preds[i] > 0.5) != labels[i]) / float( + len(preds) + ) + assert err < 0.1 + + # test custom_objective in cross-validation + cv( + param, + dtrain, + num_round, + nfold=5, + seed=0, + obj=tm.logregobj, + custom_metric=evalerror, + ) + + # test maximize parameter + def neg_evalerror(preds: np.ndarray, dtrain: DMatrix) -> Tuple[str, float]: + labels = dtrain.get_label() + preds = 1.0 / (1.0 + np.exp(-preds)) + return "error", float(sum(labels == (preds > 0.0))) / len(labels) + + bst2 = train( + param, + dtrain, + num_round, + evals=watchlist, + obj=tm.logregobj, + custom_metric=neg_evalerror, + maximize=True, + ) + preds2 = bst2.predict(dtest) + err2 = sum( + 1 for i in range(len(preds2)) if int(preds2[i] > 0.5) != labels[i] + ) / float(len(preds2)) + assert err == err2 diff --git a/v12x/pylib/xgboost/testing/callbacks.py b/v12x/pylib/xgboost/testing/callbacks.py new file mode 100644 index 0000000000000000000000000000000000000000..740b389cfde43bd604087aeef29d97292ae94fae --- /dev/null +++ b/v12x/pylib/xgboost/testing/callbacks.py @@ -0,0 +1,188 @@ +# pylint: disable=too-many-locals +"""Tests for callback functions.""" + +import json +from itertools import product +from typing import Dict, List, Tuple + +from ..callback import LearningRateScheduler +from ..core import Booster, DMatrix +from ..training import cv, train +from .utils import Device + + +def run_eta_decay( + tree_method: str, dtrain: DMatrix, dtest: DMatrix, device: Device +) -> None: + """Test learning rate scheduler, used by both CPU and GPU tests.""" + scheduler = LearningRateScheduler + + watchlist = [(dtest, "eval"), (dtrain, "train")] + num_round = 4 + + # learning_rates as a list + # init eta with 0 to check whether learning_rates work + param = { + "max_depth": 2, + "eta": 0, + "objective": "binary:logistic", + "eval_metric": "error", + "tree_method": tree_method, + "device": device, + } + evals_result: Dict[str, Dict] = {} + bst = train( + param, + dtrain, + num_round, + evals=watchlist, + callbacks=[scheduler([0.8, 0.7, 0.6, 0.5])], + evals_result=evals_result, + ) + eval_errors_0 = list(map(float, evals_result["eval"]["error"])) + assert isinstance(bst, Booster) + # validation error should decrease, if eta > 0 + assert eval_errors_0[0] > eval_errors_0[-1] + + # init learning_rate with 0 to check whether learning_rates work + param = { + "max_depth": 2, + "learning_rate": 0, + "objective": "binary:logistic", + "eval_metric": "error", + "tree_method": tree_method, + "device": device, + } + evals_result = {} + + bst = train( + param, + dtrain, + num_round, + evals=watchlist, + callbacks=[scheduler([0.8, 0.7, 0.6, 0.5])], + evals_result=evals_result, + ) + eval_errors_1 = list(map(float, evals_result["eval"]["error"])) + assert isinstance(bst, Booster) + # validation error should decrease, if learning_rate > 0 + assert eval_errors_1[0] > eval_errors_1[-1] + + # check if learning_rates override default value of eta/learning_rate + param = { + "max_depth": 2, + "objective": "binary:logistic", + "eval_metric": "error", + "tree_method": tree_method, + "device": device, + } + evals_result = {} + bst = train( + param, + dtrain, + num_round, + evals=watchlist, + callbacks=[scheduler([0, 0, 0, 0])], + evals_result=evals_result, + ) + eval_errors_2 = list(map(float, evals_result["eval"]["error"])) + assert isinstance(bst, Booster) + # validation error should not decrease, if eta/learning_rate = 0 + assert eval_errors_2[0] == eval_errors_2[-1] + + # learning_rates as a customized decay function + def eta_decay(ithround: int, num_boost_round: int = num_round) -> float: + return num_boost_round / (ithround + 1) + + evals_result = {} + bst = train( + param, + dtrain, + num_round, + evals=watchlist, + callbacks=[scheduler(eta_decay)], + evals_result=evals_result, + ) + eval_errors_3 = list(map(float, evals_result["eval"]["error"])) + + assert isinstance(bst, Booster) + + assert eval_errors_3[0] == eval_errors_2[0] + + for i in range(1, len(eval_errors_0)): + assert eval_errors_3[i] != eval_errors_2[i] + + cv(param, dtrain, num_round, callbacks=[scheduler(eta_decay)]) + + +def tree_methods_objs() -> List[Tuple[str, str]]: + """Test parameters for the leaf output test.""" + return list( + product( + ["approx", "hist"], + [ + "binary:logistic", + "reg:absoluteerror", + "reg:quantileerror", + ], + ) + ) + + +def run_eta_decay_leaf_output( + tree_method: str, objective: str, dtrain: DMatrix, dtest: DMatrix, device: Device +) -> None: + """check decay has effect on leaf output.""" + num_round = 4 + scheduler = LearningRateScheduler + + watchlist = [(dtest, "eval"), (dtrain, "train")] + + param = { + "max_depth": 2, + "objective": objective, + "eval_metric": "error", + "tree_method": tree_method, + "device": device, + } + if objective == "reg:quantileerror": + param["quantile_alpha"] = 0.3 + + def eta_decay_0(i: int) -> float: + return num_round / (i + 1) + + bst0 = train( + param, + dtrain, + num_round, + evals=watchlist, + callbacks=[scheduler(eta_decay_0)], + ) + + def eta_decay_1(i: int) -> float: + if i > 1: + return 5.0 + return num_round / (i + 1) + + bst1 = train( + param, + dtrain, + num_round, + evals=watchlist, + callbacks=[scheduler(eta_decay_1)], + ) + bst_json0 = bst0.save_raw(raw_format="json") + bst_json1 = bst1.save_raw(raw_format="json") + + j0 = json.loads(bst_json0) + j1 = json.loads(bst_json1) + + tree_2th_0 = j0["learner"]["gradient_booster"]["model"]["trees"][2] + tree_2th_1 = j1["learner"]["gradient_booster"]["model"]["trees"][2] + assert tree_2th_0["base_weights"] == tree_2th_1["base_weights"] + assert tree_2th_0["split_conditions"] == tree_2th_1["split_conditions"] + + tree_3th_0 = j0["learner"]["gradient_booster"]["model"]["trees"][3] + tree_3th_1 = j1["learner"]["gradient_booster"]["model"]["trees"][3] + assert tree_3th_0["base_weights"] != tree_3th_1["base_weights"] + assert tree_3th_0["split_conditions"] != tree_3th_1["split_conditions"] diff --git a/v12x/pylib/xgboost/testing/collective.py b/v12x/pylib/xgboost/testing/collective.py new file mode 100644 index 0000000000000000000000000000000000000000..f3eaea0b52be1f67f079a5af5e7377ad936f3734 --- /dev/null +++ b/v12x/pylib/xgboost/testing/collective.py @@ -0,0 +1,15 @@ +"""Collective module related utilities.""" + +import socket + + +def get_avail_port() -> int: + """Returns a port that's available during the function call. It doesn't prevent the + port from being used after the function returns as we can't reserve the port. The + utility makes a test more likely to pass. + + """ + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as server: + server.bind(("127.0.0.1", 0)) + port = server.getsockname()[1] + return port diff --git a/v12x/pylib/xgboost/testing/continuation.py b/v12x/pylib/xgboost/testing/continuation.py new file mode 100644 index 0000000000000000000000000000000000000000..e20891182fe8c05c8fda410505152a0abdb4252b --- /dev/null +++ b/v12x/pylib/xgboost/testing/continuation.py @@ -0,0 +1,180 @@ +"""Tests for training continuation.""" + +import json +import pickle +from typing import Any, Dict, TypeVar, cast + +import numpy as np +import pytest +from hypothesis import strategies + +import xgboost as xgb + +from .data import make_ltr + + +# pylint: disable=too-many-locals +def run_training_continuation_model_output(device: str, tree_method: str) -> None: + """Run training continuation test.""" + datasets = pytest.importorskip("sklearn.datasets") + n_samples = 64 + n_features = 32 + X, y = datasets.make_regression(n_samples, n_features, random_state=1) + + dtrain = xgb.DMatrix(X, y) + params = { + "tree_method": tree_method, + "max_depth": "2", + "gamma": "0.1", + "alpha": "0.01", + "device": device, + } + bst_0 = xgb.train(params, dtrain, num_boost_round=64) + dump_0 = bst_0.get_dump(dump_format="json") + + bst_1 = xgb.train(params, dtrain, num_boost_round=32) + bst_1 = xgb.train(params, dtrain, num_boost_round=32, xgb_model=bst_1) + dump_1 = bst_1.get_dump(dump_format="json") + + T = TypeVar("T", Dict[str, Any], float, str, int, list) + + def recursive_compare(obj_0: T, obj_1: T) -> None: + if isinstance(obj_0, float): + assert np.isclose(obj_0, obj_1, atol=1e-6) + elif isinstance(obj_0, str): + assert obj_0 == obj_1 + elif isinstance(obj_0, int): + assert obj_0 == obj_1 + elif isinstance(obj_0, dict): + for i in range(len(obj_0.items())): + assert list(obj_0.keys())[i] == list(obj_1.keys())[i] + if list(obj_0.keys())[i] != "missing": + recursive_compare(list(obj_0.values()), list(obj_1.values())) + else: + for i, lhs in enumerate(obj_0): + rhs = obj_1[i] + recursive_compare(lhs, rhs) + + assert len(dump_0) == len(dump_1) + + for i, lhs in enumerate(dump_0): + obj_0 = json.loads(lhs) + obj_1 = json.loads(dump_1[i]) + recursive_compare(obj_0, obj_1) + + +# pylint: disable=too-many-arguments +def run_training_continuation_determinism( + *, + device: str, + tree_method: str, + subsample: float, + sampling_method: str, + rate_drop: float, + colsample_bylevel: float, + num_class: int, + seed_per_iteration: bool, + learning_to_rank: bool, +) -> None: + """Check that 2-session training (4+4 iters) equals single-session (8 iters).""" + datasets = pytest.importorskip("sklearn.datasets") + + n_samples = 128 + n_features = 16 + total_rounds = 8 + split_at = 4 + + if learning_to_rank: + X, y, qid, _ = make_ltr( + n_samples=n_samples, + n_features=n_features, + n_query_groups=8, + max_rel=4, + ) + objective = "rank:ndcg" + elif num_class > 1: + X, y = datasets.make_classification( + n_samples=n_samples, + n_features=n_features, + n_informative=6, + n_classes=num_class, + random_state=42, + ) + objective = "multi:softprob" + else: + X, y = datasets.make_regression( + n_samples=n_samples, n_features=n_features, random_state=42 + ) + objective = "reg:squarederror" + + if learning_to_rank: + dtrain = xgb.DMatrix(X, y, qid=qid) + else: + dtrain = xgb.DMatrix(X, y) + + params: Dict[str, Any] = { + "device": device, + "tree_method": tree_method, + "max_depth": 4, + "objective": objective, + "subsample": subsample, + "sampling_method": sampling_method, + "rate_drop": rate_drop, + "colsample_bylevel": colsample_bylevel, + "seed_per_iteration": seed_per_iteration, + } + if learning_to_rank: + params["lambdarank_pair_method"] = "mean" + elif num_class > 1: + params["num_class"] = num_class + + bst_single = xgb.train(params, dtrain, num_boost_round=total_rounds) + + bst_first = xgb.train(params, dtrain, num_boost_round=split_at) + bst_continued = xgb.train( + params, dtrain, num_boost_round=total_rounds - split_at, xgb_model=bst_first + ) + + config_single = json.loads(bst_single.save_config()) + config_cont = json.loads(bst_continued.save_config()) + + rng_single = config_single["learner"]["generic_param"]["rng_state"] + rng_cont = config_cont["learner"]["generic_param"]["rng_state"] + assert rng_single == rng_cont, "RNG states diverged between single and continued." + + pred_single = bst_single.predict(dtrain) + pred_cont = bst_continued.predict(dtrain) + np.testing.assert_allclose(pred_single, pred_cont) + + bst_continued = xgb.train( + params, + dtrain, + num_boost_round=total_rounds - split_at, + xgb_model=pickle.loads(pickle.dumps(bst_first)), + ) + pred_cont = bst_continued.predict(dtrain) + np.testing.assert_allclose(pred_single, pred_cont, rtol=1e-4) + + +def make_determinism_strategy(tree_methods: list[str]) -> "strategies.SearchStrategy": + """Hypothesis strategy for testing training continuation with sampling.""" + strategy = strategies.fixed_dictionaries( + { + "subsample": strategies.sampled_from([0.5, 1.0]), + "sampling_method": strategies.sampled_from(["uniform", "gradient_based"]), + "rate_drop": strategies.sampled_from([0.0, 0.5]), + "colsample_bylevel": strategies.sampled_from([0.5, 1.0]), + "tree_method": strategies.sampled_from(tree_methods), + "num_class": strategies.sampled_from([1, 3]), + "seed_per_iteration": strategies.booleans(), + "learning_to_rank": strategies.booleans(), + } + ).filter( + lambda x: ( + not ( + x["sampling_method"] == "gradient_based" and x["tree_method"] != "hist" + ) + and not (bool(x["learning_to_rank"]) and cast(int, x["num_class"]) > 1) + ) + ) + return strategy diff --git a/v12x/pylib/xgboost/testing/dask.py b/v12x/pylib/xgboost/testing/dask.py new file mode 100644 index 0000000000000000000000000000000000000000..e785e070df6a4403a316669193bfcc0b7c0882a8 --- /dev/null +++ b/v12x/pylib/xgboost/testing/dask.py @@ -0,0 +1,623 @@ +"""Tests for dask shared by different test modules.""" + +import json +from typing import Any, List, Literal, Tuple, Type, Union, cast, overload + +import numpy as np +import pandas as pd +from dask import array as da +from dask import dataframe as dd +from distributed import Client, get_worker +from packaging.version import parse as parse_version +from sklearn.datasets import make_classification, make_regression + +import xgboost as xgb +import xgboost.testing as tm +from xgboost.compat import concat +from xgboost.testing.updater import get_basescore + +from .. import dask as dxgb +from .._typing import EvalsLog +from ..dask import TrainReturnT, _get_rabit_args +from ..dask.utils import _DASK_VERSION +from .data import make_batches +from .data import make_categorical as make_cat_local +from .multi_target import LsObj1, LsObj2 +from .ordinal import make_recoded +from .utils import Device, assert_allclose + + +def check_init_estimation_clf(tree_method: str, device: Device, client: Client) -> None: + """Test init estimation for classsifier.""" + X, y = make_classification(n_samples=4096 * 2, n_features=32, random_state=1994) + clf = xgb.XGBClassifier( + n_estimators=1, max_depth=1, tree_method=tree_method, device=device + ) + clf.fit(X, y) + base_score = get_basescore(clf) + + dx = da.from_array(X).rechunk(chunks=(32, None)) + dy = da.from_array(y).rechunk(chunks=(32,)) + dclf = dxgb.DaskXGBClassifier( + n_estimators=1, + max_depth=1, + tree_method=tree_method, + device=device, + ) + dclf.client = client + dclf.fit(dx, dy) + dbase_score = get_basescore(dclf) + np.testing.assert_allclose(base_score, dbase_score) + + +def check_init_estimation_reg(tree_method: str, device: Device, client: Client) -> None: + """Test init estimation for regressor.""" + # pylint: disable=unbalanced-tuple-unpacking + X, y = make_regression(n_samples=4096 * 2, n_features=32, random_state=1994) + reg = xgb.XGBRegressor( + n_estimators=1, max_depth=1, tree_method=tree_method, device=device + ) + reg.fit(X, y) + base_score = get_basescore(reg) + + dx = da.from_array(X).rechunk(chunks=(32, None)) + dy = da.from_array(y).rechunk(chunks=(32,)) + dreg = dxgb.DaskXGBRegressor( + n_estimators=1, max_depth=1, tree_method=tree_method, device=device + ) + dreg.client = client + dreg.fit(dx, dy) + dbase_score = get_basescore(dreg) + np.testing.assert_allclose(base_score, dbase_score) + + +def check_init_estimation(tree_method: str, device: Device, client: Client) -> None: + """Test init estimation.""" + check_init_estimation_reg(tree_method, device, client) + check_init_estimation_clf(tree_method, device, client) + + +def check_uneven_nan( + client: Client, tree_method: str, device: Device, n_workers: int +) -> None: + """Issue #9271, not every worker has missing value.""" + assert n_workers >= 2 + + with client.as_current(): + clf = dxgb.DaskXGBClassifier(tree_method=tree_method, device=device) + X = pd.DataFrame({"a": range(10000), "b": range(10000, 0, -1)}) + y = pd.Series([*[0] * 5000, *[1] * 5000]) + + X.loc[:3000:1000, "a"] = np.nan + + client.wait_for_workers(n_workers=n_workers) + + clf.fit( + dd.from_pandas(X, npartitions=n_workers), + dd.from_pandas(y, npartitions=n_workers), + ) + + +def make_multi_output_regression( + device: Device, *, n_samples: int = 512, n_features: int = 8, n_targets: int = 3 +) -> Tuple[da.Array, da.Array]: + """Make a Dask array multi-output regression dataset for CPU or CUDA tests.""" + chunksize = 64 + + X, y = make_regression( + n_samples, n_features, n_targets=n_targets, random_state=2026 + ) + dX, dy = ( + da.from_array(X, chunks=(chunksize, n_features)), + da.from_array(y, chunks=(chunksize, n_targets)), + ) + if device == "cuda": + dX, dy = dX.to_backend("cupy"), dy.to_backend("cupy") + return dX, dy + + +def _as_numpy(array: Any) -> np.ndarray: + if hasattr(array, "get"): + array = array.get() + return np.asarray(array) + + +def _train_multi_output_tree( + client: Client, device: Device +) -> Tuple[da.Array, da.Array, dxgb.DaskDMatrix, TrainReturnT]: + n_targets = 3 + X, y = make_multi_output_regression(device, n_targets=n_targets) + Xy = dxgb.DaskDMatrix(client, X, y) + result = dxgb.train( + client, + { + "device": device, + "tree_method": "hist", + "objective": "reg:absoluteerror", + "eval_metric": "mae", + "multi_strategy": "multi_output_tree", + "num_target": n_targets, + "max_depth": 4, + "max_bin": 64, + "debug_synchronize": True, + }, + Xy, + num_boost_round=4, + evals=[(Xy, "train")], + ) + return X, y, Xy, result + + +def check_multi_output_tree_regressor(client: Client, device: Device) -> None: + """Test Dask vector-leaf regression with train and sklearn-style APIs.""" + tolerance = 1e-3 + X, y, Xy, result = _train_multi_output_tree(client, device) + n_targets = y.shape[1] + assert isinstance(n_targets, int) + + history = result["history"]["train"]["mae"] + assert np.isfinite(np.asarray(history)).all() + assert tm.non_increasing(history, tolerance=tolerance) + + predt = _as_numpy(dxgb.predict(client, result["booster"], Xy).compute()) + assert predt.shape == (X.shape[0], n_targets) + assert np.isfinite(predt).all() + + reg = dxgb.DaskXGBRegressor( + n_estimators=4, + device=device, + tree_method="hist", + objective="reg:absoluteerror", + multi_strategy="multi_output_tree", + max_depth=4, + max_bin=64, + ) + reg.client = client + reg.fit(X, y, eval_set=[(X, y)]) + + predt = _as_numpy(reg.predict(X).compute()) + assert predt.shape == (X.shape[0], n_targets) + assert np.isfinite(predt).all() + + config = json.loads(reg.get_booster().save_config()) + assert config["learner"]["learner_train_param"]["multi_strategy"] == ( + "multi_output_tree" + ) + + for objective in (LsObj1(device), LsObj2(device, False)): + reg = dxgb.DaskXGBRegressor( + n_estimators=2, + device=device, + tree_method="hist", + objective=objective, + multi_strategy="multi_output_tree", + max_depth=2, + max_bin=64, + ) + reg.client = client + reg.fit(X, y) + predt = _as_numpy(reg.predict(X).compute()) + assert predt.shape == (X.shape[0], n_targets) + assert np.isfinite(predt).all() + + +def check_multi_output_tree_classifier(client: Client, device: Device) -> None: + """Test Dask vector-leaf classification with array and dataframe labels.""" + n_targets = 3 + X, y = make_multi_output_regression(device, n_targets=n_targets) + + def check_classifier(labels: Union[da.Array, dd.DataFrame]) -> None: + clf = dxgb.DaskXGBClassifier( + n_estimators=2, + device=device, + tree_method="hist", + objective=LsObj2(device, False), + multi_strategy="multi_output_tree", + max_depth=2, + max_bin=64, + ) + clf.client = client + clf.fit(X, labels) + + assert isinstance(clf.classes_, np.ndarray) + np.testing.assert_array_equal(clf.classes_, np.array([0, 1])) + assert clf.n_classes_ == 2 + + predt = _as_numpy(clf.predict(X).compute()) + proba = _as_numpy(clf.predict_proba(X).compute()) + assert predt.shape == (X.shape[0], n_targets) + assert proba.shape == predt.shape + np.testing.assert_array_equal(predt, (proba > 0.5).astype(predt.dtype)) + + config = json.loads(clf.get_booster().save_config())["learner"] + assert config["objective"]["name"] == "binary:logistic" + assert int(config["learner_model_param"]["num_class"]) == 0 + + y_ind = (y > 0.0).astype(np.int32) + check_classifier(y_ind) + y_df = dd.from_dask_array(y_ind) + if device == "cuda": + y_df = y_df.to_backend("cudf") + check_classifier(y_df) + + +def check_multi_output_tree_shap(client: Client, device: Device) -> None: + """Test SHAP output shapes for Dask vector-leaf models.""" + X, y, _, result = _train_multi_output_tree(client, device) + n_targets = y.shape[1] + assert isinstance(n_targets, int) + booster = result["booster"] + + margin = _as_numpy(dxgb.predict(client, booster, X, output_margin=True).compute()) + n_features = X.shape[1] + assert isinstance(n_features, int) + contributions = dxgb.predict(client, booster, X, pred_contribs=True) + contributions_shape = (X.shape[0], n_targets, n_features + 1) + assert contributions.shape == contributions_shape + assert contributions.chunks == ( + X.chunks[0], + (n_targets,), + (n_features + 1,), + ) + contributions = _as_numpy(contributions.compute()) + assert contributions.shape == contributions_shape + np.testing.assert_allclose( + contributions.sum(axis=-1), + margin, + rtol=1e-4, + atol=1e-4, + ) + + interactions = dxgb.predict(client, booster, X, pred_interactions=True) + interactions_shape = (X.shape[0], n_targets, n_features + 1, n_features + 1) + assert interactions.shape == interactions_shape + assert interactions.chunks == ( + X.chunks[0], + (n_targets,), + (n_features + 1,), + (n_features + 1,), + ) + interactions = _as_numpy(interactions.compute()) + assert interactions.shape == interactions_shape + np.testing.assert_allclose( + interactions.sum(axis=(-2, -1)), + margin, + rtol=1e-4, + atol=1e-4, + ) + + +def check_external_memory( # pylint: disable=too-many-locals + worker_id: int, + n_workers: int, + device: str, + comm_args: dict, + is_qdm: bool, +) -> None: + """Basic checks for distributed external memory.""" + n_samples_per_batch = 32 + n_features = 4 + n_batches = 16 + use_cupy = device != "cpu" + + n_threads = get_worker().state.nthreads + with xgb.collective.CommunicatorContext(dmlc_communicator="rabit", **comm_args): + it = tm.IteratorForTest( + *make_batches( + n_samples_per_batch, + n_features, + n_batches, + use_cupy=use_cupy, + random_state=worker_id, + ), + cache="cache", + ) + if is_qdm: + Xy: xgb.DMatrix = xgb.ExtMemQuantileDMatrix(it, nthread=n_threads) + else: + Xy = xgb.DMatrix(it, nthread=n_threads) + results: EvalsLog = {} + xgb.train( + {"tree_method": "hist", "nthread": n_threads, "device": device}, + Xy, + evals=[(Xy, "Train")], + num_boost_round=32, + evals_result=results, + ) + assert tm.non_increasing(cast(List[float], results["Train"]["rmse"])) + + lx, ly, lw = [], [], [] + for i in range(n_workers): + x, y, w = make_batches( + n_samples_per_batch, + n_features, + n_batches, + use_cupy=use_cupy, + random_state=i, + ) + lx.extend(x) + ly.extend(y) + lw.extend(w) + + X = concat(lx) + yconcat = concat(ly) + wconcat = concat(lw) + if is_qdm: + Xy = xgb.QuantileDMatrix(X, yconcat, weight=wconcat, nthread=n_threads) + else: + Xy = xgb.DMatrix(X, yconcat, weight=wconcat, nthread=n_threads) + + results_local: EvalsLog = {} + xgb.train( + {"tree_method": "hist", "nthread": n_threads, "device": device}, + Xy, + evals=[(Xy, "Train")], + num_boost_round=32, + evals_result=results_local, + ) + np.testing.assert_allclose( + results["Train"]["rmse"], results_local["Train"]["rmse"], rtol=1e-4 + ) + + +def get_rabit_args(client: Client, n_workers: int) -> Any: + """Get RABIT collective communicator arguments for tests.""" + return client.sync(_get_rabit_args, client, n_workers) + + +def get_client_workers(client: Client) -> List[str]: + "Get workers from a dask client." + kwargs = {"n_workers": -1} if _DASK_VERSION() >= parse_version("2025.4.0") else {} + workers = client.scheduler_info(**kwargs)["workers"] + return list(workers.keys()) + + +def make_ltr( # pylint: disable=too-many-locals,too-many-arguments + client: Client, + n_samples: int, + n_features: int, + *, + n_query_groups: int, + max_rel: int, + device: str, +) -> Tuple[dd.DataFrame, dd.Series, dd.Series]: + """Synthetic dataset for learning to rank.""" + workers = get_client_workers(client) + n_samples_per_worker = n_samples // len(workers) + + if device == "cpu": + from pandas import DataFrame as DF + else: + from cudf import DataFrame as DF + + def make(n: int, seed: int) -> pd.DataFrame: + rng = np.random.default_rng(seed) + X, y = make_classification( + n, + n_features, + n_informative=n_features, + n_redundant=0, + n_classes=max_rel, + random_state=seed, + ) + qid = rng.integers(size=(n,), low=0, high=n_query_groups) + df = DF(X, columns=[f"f{i}" for i in range(n_features)]) + df["qid"] = qid + df["y"] = y + return df + + futures = [] + i = 0 + for k in range(0, n_samples, n_samples_per_worker): + fut = client.submit( + make, n=n_samples_per_worker, seed=k, workers=[workers[i % len(workers)]] + ) + futures.append(fut) + i += 1 + + last = n_samples - (n_samples_per_worker * len(workers)) + if last != 0: + fut = client.submit(make, n=last, seed=n_samples_per_worker * len(workers)) + futures.append(fut) + + meta = make(1, 0) + df = dd.from_delayed(futures, meta=meta) + assert isinstance(df, dd.DataFrame) + return df.drop(["qid", "y"], axis=1), df.y, df.qid + + +def check_no_group_split(client: Client, device: str) -> None: + """Test for the allow_group_split parameter.""" + X_tr, q_tr, y_tr = make_ltr( + client, 4096, 128, n_query_groups=4, max_rel=5, device=device + ) + X_va, q_va, y_va = make_ltr( + client, 1024, 128, n_query_groups=4, max_rel=5, device=device + ) + + ltr = dxgb.DaskXGBRanker( + allow_group_split=False, + n_estimators=36, + device=device, + objective="rank:pairwise", + ) + ltr.fit( + X_tr, + y_tr, + qid=q_tr, + eval_set=[(X_tr, y_tr), (X_va, y_va)], + eval_qid=[q_tr, q_va], + verbose=True, + ) + + assert ltr.n_features_in_ == 128 + assert X_tr.shape[1] == ltr.n_features_in_ # no change + ndcg = ltr.evals_result()["validation_0"]["ndcg@32"] + assert tm.non_decreasing(ndcg[:16], tolerance=1e-2), ndcg + np.testing.assert_allclose(ndcg[-1], 1.0, rtol=1e-2) + + +@overload +def make_categorical( # pylint: disable=too-many-locals, too-many-arguments + client: Client, + n_samples: int, + n_features: int, + n_categories: int, + *, + onehot: bool = ..., + n_targets: Literal[1] = ..., + cat_dtype: np.typing.DTypeLike = ..., +) -> Tuple[dd.DataFrame, dd.Series]: ... + + +@overload +def make_categorical( # pylint: disable=too-many-locals, too-many-arguments + client: Client, + n_samples: int, + n_features: int, + n_categories: int, + *, + onehot: bool = ..., + n_targets: int, + cat_dtype: np.typing.DTypeLike = ..., +) -> Tuple[dd.DataFrame, Union[dd.Series, dd.DataFrame]]: ... + + +def make_categorical( # pylint: disable=too-many-locals, too-many-arguments + client: Client, + n_samples: int, + n_features: int, + n_categories: int, + *, + onehot: bool = False, + n_targets: int = 1, + cat_dtype: np.typing.DTypeLike = np.int64, +) -> Tuple[dd.DataFrame, Union[dd.Series, dd.DataFrame]]: + """Synthesize categorical data with dask.""" + workers = get_client_workers(client) + n_workers = len(workers) + dfs = [] + + label_cols = ( + [f"label_{i}" for i in range(n_targets)] if n_targets > 1 else ["label"] + ) + + def pack(**kwargs: Any) -> dd.DataFrame: + X, y = make_cat_local(**kwargs) + if y.ndim == 2: + for i in range(y.shape[1]): + X[f"label_{i}"] = y[:, i] + else: + X["label"] = y + return X + + meta = pack( + n_samples=1, + n_features=n_features, + n_categories=n_categories, + onehot=False, + n_targets=n_targets, + cat_dtype=cat_dtype, + ) + + for i, worker in enumerate(workers): + l_n_samples = min( + n_samples // n_workers, n_samples - i * (n_samples // n_workers) + ) + # make sure there's at least one sample for testing empty DMatrix + if n_samples == 1 and i == 0: + l_n_samples = 1 + future = client.submit( + pack, + n_samples=l_n_samples, + n_features=n_features, + n_categories=n_categories, + n_targets=n_targets, + cat_dtype=cat_dtype, + onehot=False, + workers=[worker], + ) + dfs.append(future) + + df: dd.DataFrame = cast(dd.DataFrame, dd.from_delayed(dfs, meta=meta)) + y = df[label_cols] + if n_targets == 1: + y = y[label_cols[0]] + X = df[df.columns.difference(label_cols)] + + if onehot: + return dd.get_dummies(X), y + return X, y + + +# pylint: disable=too-many-locals +def run_recode(client: Client, device: Device) -> None: + """Run re-coding test with the Dask interface.""" + + def create_dmatrix( + DMatrixT: Type[dxgb.DaskDMatrix], *args: Any, **kwargs: Any + ) -> dxgb.DaskDMatrix: + if DMatrixT is dxgb.DaskQuantileDMatrix: + ref = kwargs.pop("ref", None) + return DMatrixT(*args, ref=ref, **kwargs) + + kwargs.pop("ref", None) + return DMatrixT(*args, **kwargs) + + def run(DMatrixT: Type[dxgb.DaskDMatrix]) -> None: + enc, reenc, y, _, _ = make_recoded(device, n_features=96) + to = get_client_workers(client) + + denc, dreenc, dy = ( + dd.from_pandas(enc, npartitions=8).persist(workers=to), + dd.from_pandas(reenc, npartitions=8).persist(workers=to), + da.from_array(y, chunks=(y.shape[0] // 8,)).persist(workers=to), + ) + + Xy = create_dmatrix(DMatrixT, client, denc, dy, enable_categorical=True) + Xy_valid = create_dmatrix( + DMatrixT, client, dreenc, dy, enable_categorical=True, ref=Xy + ) + # Base model + results = dxgb.train( + client, {"device": device}, Xy, evals=[(Xy_valid, "Valid")] + ) + + # Training continuation + Xy = create_dmatrix(DMatrixT, client, denc, dy, enable_categorical=True) + Xy_valid = create_dmatrix( + DMatrixT, client, dreenc, dy, enable_categorical=True, ref=Xy + ) + results_1 = dxgb.train( + client, + {"device": device}, + Xy, + evals=[(Xy_valid, "Valid")], + xgb_model=results["booster"], + ) + + # Reversed training continuation + Xy = create_dmatrix(DMatrixT, client, dreenc, dy, enable_categorical=True) + Xy_valid = create_dmatrix( + DMatrixT, client, denc, dy, enable_categorical=True, ref=Xy + ) + results_2 = dxgb.train( + client, + {"device": device}, + Xy, + evals=[(Xy_valid, "Valid")], + xgb_model=results["booster"], + ) + assert np.isfinite(results_1["history"]["Valid"]["rmse"]).all() + assert np.isfinite(results_2["history"]["Valid"]["rmse"]).all() + + predt_0 = dxgb.inplace_predict(client, results, denc).compute() + predt_1 = dxgb.inplace_predict(client, results, dreenc).compute() + assert_allclose(device, predt_0, predt_1) + + predt_0 = dxgb.predict(client, results, Xy).compute() + predt_1 = dxgb.predict(client, results, Xy_valid).compute() + assert_allclose(device, predt_0, predt_1) + + for DMatrixT in [dxgb.DaskDMatrix, dxgb.DaskQuantileDMatrix]: + run(DMatrixT) diff --git a/v12x/pylib/xgboost/testing/data.py b/v12x/pylib/xgboost/testing/data.py new file mode 100644 index 0000000000000000000000000000000000000000..b642921e1e7d3696721b253cfd23aee004386c51 --- /dev/null +++ b/v12x/pylib/xgboost/testing/data.py @@ -0,0 +1,1193 @@ +# pylint: disable=too-many-lines +"""Utilities for data generation.""" + +import gc +import multiprocessing +import os +import string +import zipfile +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Dict, + Generator, + List, + NamedTuple, + Optional, + Sequence, + Set, + Tuple, + Type, + Union, +) +from urllib import request + +import numpy as np +import pytest +from numpy import typing as npt +from numpy.random import Generator as RNG +from scipy import sparse + +from ..compat import concat +from ..core import DataIter, DMatrix, QuantileDMatrix +from ..data import is_pd_cat_dtype, pandas_pyarrow_mapper +from ..sklearn import ArrayLike, XGBRanker +from ..training import train as train_fn + +if TYPE_CHECKING: + from pandas import DataFrame as DataFrameT +else: + DataFrameT = Any + +joblib = pytest.importorskip("joblib") +memory = joblib.Memory("./cachedir", verbose=0) + + +def np_dtypes( + n_samples: int, n_features: int +) -> Generator[Union[Tuple[np.ndarray, np.ndarray], Tuple[list, list]], None, None]: + """Enumerate all supported dtypes from numpy.""" + pd = pytest.importorskip("pandas") + + rng = np.random.RandomState(1994) + # Integer and float. + orig = rng.randint(low=0, high=127, size=n_samples * n_features).reshape( + n_samples, n_features + ) + dtypes = [ + np.int32, + np.int64, + np.byte, + np.short, + np.intc, + np.int_, + np.longlong, + np.uint32, + np.uint64, + np.ubyte, + np.ushort, + np.uintc, + np.uint, + np.ulonglong, + np.float16, + np.float32, + np.float64, + np.half, + np.single, + np.double, + ] + for dtype in dtypes: + X = np.array(orig, dtype=dtype) + yield orig, X + yield orig.tolist(), X.tolist() + + for dtype in dtypes: + X = np.array(orig, dtype=dtype) + df_orig = pd.DataFrame(orig) + df = pd.DataFrame(X) + yield df_orig, df + + # Boolean + orig = rng.binomial(1, 0.5, size=n_samples * n_features).reshape( + n_samples, n_features + ) + for dtype1 in [np.bool_, bool]: + X = np.array(orig, dtype=dtype1) + yield orig, X + + for dtype2 in [np.bool_, bool]: + X = np.array(orig, dtype=dtype2) + df_orig = pd.DataFrame(orig) + df = pd.DataFrame(X) + yield df_orig, df + + +def pd_dtypes() -> Generator: + """Enumerate all supported pandas extension types.""" + pd = pytest.importorskip("pandas") + + # Integer + dtypes = [ + pd.UInt8Dtype(), + pd.UInt16Dtype(), + pd.UInt32Dtype(), + pd.UInt64Dtype(), + pd.Int8Dtype(), + pd.Int16Dtype(), + pd.Int32Dtype(), + pd.Int64Dtype(), + ] + + Null: Union[float, None, Any] = np.nan + orig = pd.DataFrame( + {"f0": [1, 2, Null, 3], "f1": [4, 3, Null, 1]}, dtype=np.float32 + ) + for Null in (np.nan, None, pd.NA): + for dtype in dtypes: + df = pd.DataFrame( + {"f0": [1, 2, Null, 3], "f1": [4, 3, Null, 1]}, dtype=dtype + ) + yield orig, df + + # Float + Null = np.nan + dtypes = [pd.Float32Dtype(), pd.Float64Dtype()] + orig = pd.DataFrame( + {"f0": [1.0, 2.0, Null, 3.0], "f1": [3.0, 2.0, Null, 1.0]}, dtype=np.float32 + ) + for Null in (np.nan, None, pd.NA): + for dtype in dtypes: + df = pd.DataFrame( + {"f0": [1.0, 2.0, Null, 3.0], "f1": [3.0, 2.0, Null, 1.0]}, dtype=dtype + ) + yield orig, df + ser_orig = orig["f0"] + ser = df["f0"] + assert isinstance(ser, pd.Series) + assert isinstance(ser_orig, pd.Series) + yield ser_orig, ser + + # Categorical + orig = orig.astype("category") + for c in orig.columns: + orig[c] = orig[c].cat.rename_categories(int) + for Null in (np.nan, None, pd.NA): + df = pd.DataFrame( + {"f0": [1, 2, Null, 3], "f1": [3, 2, Null, 1]}, + dtype=pd.CategoricalDtype(), + ) + yield orig, df + + # Boolean + for Null in [None, pd.NA]: + data = {"f0": [True, False, Null, True], "f1": [False, True, Null, True]} + # pd.NA is not convertible to bool. + orig = pd.DataFrame(data, dtype=np.bool_ if Null is None else pd.BooleanDtype()) + df = pd.DataFrame(data, dtype=pd.BooleanDtype()) + yield orig, df + + +def pd_arrow_dtypes() -> Generator: + """Pandas DataFrame with pyarrow backed type.""" + pd = pytest.importorskip("pandas") + pa = pytest.importorskip("pyarrow") + + # Integer + dtypes = pandas_pyarrow_mapper + # Create a dictionary-backed dataframe, enable this when the roundtrip is + # implemented in pandas/pyarrow + # + # category = pd.ArrowDtype(pa.dictionary(pa.int32(), pa.int32(), ordered=True)) + # df = pd.DataFrame({"f0": [0, 2, Null, 3], "f1": [4, 3, Null, 1]}, dtype=category) + + # Error: + # >>> df.astype("category") + # Function 'dictionary_encode' has no kernel matching input types + # (array[dictionary]) + + # Error: + # pd_cat_df = pd.DataFrame( + # {"f0": [0, 2, Null, 3], "f1": [4, 3, Null, 1]}, + # dtype="category" + # ) + # pa_catcodes = ( + # df["f1"].array.__arrow_array__().combine_chunks().to_pandas().cat.codes + # ) + # pd_catcodes = pd_cat_df["f1"].cat.codes + # assert pd_catcodes.equals(pa_catcodes) + + for Null in (None, pd.NA, 0): + for dtype in dtypes: + if dtype.startswith("float16") or dtype.startswith("bool"): + continue + # Use np.nan is a baseline + orig_null = Null if not pd.isna(Null) and Null == 0 else np.nan + orig = pd.DataFrame( + {"f0": [1, 2, orig_null, 3], "f1": [4, 3, orig_null, 1]}, + dtype=np.float32, + ) + + df = pd.DataFrame( + {"f0": [1, 2, Null, 3], "f1": [4, 3, Null, 1]}, dtype=dtype + ) + yield orig, df + + # If Null is `False`, then there's no missing value. + for Null in (pd.NA, False): + orig = pd.DataFrame( + {"f0": [True, False, Null, True], "f1": [False, True, Null, True]}, + dtype=pd.BooleanDtype(), + ) + df = pd.DataFrame( + {"f0": [True, False, Null, True], "f1": [False, True, Null, True]}, + dtype=pd.ArrowDtype(pa.bool_()), + ) + yield orig, df + + +def check_inf(rng: RNG) -> None: + """Validate there's no inf in X.""" + X = rng.random(size=32).reshape(8, 4) + y = rng.random(size=8) + X[5, 2] = np.inf + + with pytest.raises(ValueError, match="Input data contains `inf`"): + QuantileDMatrix(X, y) + + with pytest.raises(ValueError, match="Input data contains `inf`"): + DMatrix(X, y) + + +@memory.cache +def get_california_housing() -> Tuple[np.ndarray, np.ndarray]: + """Synthesize a dataset similar to the sklearn California housing dataset. + + The real one can be obtained via: + + .. code-block:: + + import sklearn.datasets + + X, y = sklearn.datasets.fetch_california_housing(return_X_y=True) + + """ + n_samples = 20640 + rng = np.random.default_rng(2025) + + pd = pytest.importorskip("pandas") + + def mixture_2comp( + means: List[float], sigmas: List[float], weights: List[float] + ) -> np.ndarray: + l0 = rng.normal( + size=(int(n_samples * weights[0])), loc=means[0], scale=sigmas[0] + ) + l1 = rng.normal(size=(n_samples - l0.shape[0]), loc=means[1], scale=sigmas[1]) + return np.concatenate([l0, l1], axis=0) + + def norm(mean: float, std: float) -> np.ndarray: + return rng.normal(loc=mean, scale=std, size=(n_samples,)) + + df = pd.DataFrame( + { + "Longitude": mixture_2comp( + [-118.0703597, -121.85682825], + [0.7897320650373969, 0.7248398629412008], + [0.60402556, 0.39597444], + ), + "Latitude": mixture_2comp( + [37.84266317, 33.86030848], + [1.0643911549736087, 0.5049274656834589], + [0.44485062, 0.55514938], + ), + "MedInc": norm(mean=3.8706710029069766, std=1.8997756945748738), + "HouseAge": norm(mean=28.639486434108527, std=12.585252725724606), + "AveRooms": norm(mean=5.428999742190376, std=2.474113202333516), + "AveBedrms": norm(mean=1.096675149606208, std=0.47389937625774475), + "Population": norm(mean=1425.4767441860465, std=1132.434687757615), + "AveOccup": norm(mean=3.0706551594363742, std=10.385797959128219), + "MedHouseVal": norm(mean=2.068558169089147, std=1.1539282040412253), + } + ) + X = df[df.columns.difference(["MedHouseVal"])].to_numpy() + y = df["MedHouseVal"].to_numpy() + return X, y + + +@memory.cache +def get_digits() -> Tuple[np.ndarray, np.ndarray]: + """Fetch the digits dataset from sklearn.""" + datasets = pytest.importorskip("sklearn.datasets") + data = datasets.load_digits() + return data.data, data.target + + +@memory.cache +def get_cancer() -> Tuple[np.ndarray, np.ndarray]: + """Fetch the breast cancer dataset from sklearn.""" + datasets = pytest.importorskip("sklearn.datasets") + return datasets.load_breast_cancer(return_X_y=True) + + +@memory.cache +def get_sparse() -> Tuple[np.ndarray, np.ndarray]: + """Generate a sparse dataset.""" + datasets = pytest.importorskip("sklearn.datasets") + rng = np.random.RandomState(199) + n = 2000 + sparsity = 0.75 + X, y = datasets.make_regression(n, random_state=rng) + flag = rng.binomial(1, sparsity, X.shape) + for i in range(X.shape[0]): + for j in range(X.shape[1]): + if flag[i, j]: + X[i, j] = np.nan + return X, y + + +# pylint: disable=too-many-statements +@memory.cache +def get_ames_housing() -> Tuple[DataFrameT, np.ndarray]: + """Get a synthetic version of the amse housing dataset. + + The real one can be obtained via: + + .. code-block:: + + from sklearn import datasets + + datasets.fetch_openml(data_id=42165, as_frame=True, return_X_y=True) + + Number of samples: 1460 + Number of features: 20 + Number of categorical features: 10 + Number of numerical features: 10 + """ + if TYPE_CHECKING: + import pandas as pd + else: + pd = pytest.importorskip("pandas") + + rng = np.random.default_rng(1994) + n_samples = 1460 + df = pd.DataFrame() + + def synth_cat( + name_proba: Dict[Union[str, float], float], density: float + ) -> pd.Series: + n_nulls = int(n_samples * (1 - density)) + has_nan = np.abs(1.0 - density) > 1e-6 and n_nulls > 0 + if has_nan: + sparsity = 1.0 - density + name_proba[np.nan] = sparsity + + keys = list(name_proba.keys()) + p = list(name_proba.values()) + p[-1] += 1.0 - np.sum(p) # Fix floating point error + x = rng.choice(keys, size=n_samples, p=p) + + series = pd.Series( + x, + dtype=pd.CategoricalDtype( + # not NA + filter(lambda x: isinstance(x, str), keys) + ), + ) + return series + + df["BldgType"] = synth_cat( + { + "1Fam": 0.835616, + "2fmCon": 0.078082, + "Duplex": 0.035616, + "Twnhs": 0.029452, + "TwnhsE": 0.021233, + }, + 1.0, + ) + df["GarageFinish"] = synth_cat( + {"Unf": 0.414384, "RFn": 0.289041, "Fin": 0.241096}, 0.94452 + ) + df["LotConfig"] = synth_cat( + { + "Corner": 0.180137, + "CulDSac": 0.064384, + "FR2": 0.032192, + "FR3": 0.002740, + }, + 1.0, + ) + df["Functional"] = synth_cat( + { + "Typ": 0.931506, + "Min2": 0.023287, + "Min1": 0.021232, + "Mod": 0.010273, + "Maj1": 0.009589, + "Maj2": 0.003424, + "Sev": 0.000684, + }, + 1.0, + ) + df["MasVnrType"] = synth_cat( + { + "None": 0.591780, + "BrkFace": 0.304794, + "Stone": 0.087671, + "BrkCmn": 0.010273, + }, + 0.99452, + ) + df["HouseStyle"] = synth_cat( + { + "1Story": 0.497260, + "2Story": 0.304794, + "1.5Fin": 0.105479, + "SLvl": 0.044520, + "SFoyer": 0.025342, + "1.5Unf": 0.009589, + "2.5Unf": 0.007534, + "2.5Fin": 0.005479, + }, + 1.0, + ) + df["FireplaceQu"] = synth_cat( + { + "Gd": 0.260273, + "TA": 0.214383, + "Fa": 0.022602, + "Ex": 0.016438, + "Po": 0.013698, + }, + 0.527397, + ) + df["ExterCond"] = synth_cat( + { + "TA": 0.878082, + "Gd": 0.1, + "Fa": 0.019178, + "Ex": 0.002054, + "Po": 0.000684, + }, + 1.0, + ) + df["ExterQual"] = synth_cat( + { + "TA": 0.620547, + "Gd": 0.334246, + "Ex": 0.035616, + "Fa": 0.009589, + }, + 1.0, + ) + df["PoolQC"] = synth_cat( + { + "Gd": 0.002054, + "Ex": 0.001369, + "Fa": 0.001369, + }, + 0.004794, + ) + + # We focus on the cateogircal values here, for numerical features, simple normal + # distribution is used, which doesn't match the original data. + def synth_num(loc: float, std: float, density: float) -> pd.Series: + x = rng.normal(loc=loc, scale=std, size=n_samples) + n_nulls = int(n_samples * (1 - density)) + if np.abs(1.0 - density) > 1e-6 and n_nulls > 0: + null_idx = rng.choice(n_samples, size=n_nulls, replace=False) + x[null_idx] = np.nan + return pd.Series(x, dtype=np.float64) + + df["3SsnPorch"] = synth_num(3.4095890410958902, 29.31733055678188, 1.0) + df["Fireplaces"] = synth_num(0.613013698630137, 0.6446663863122295, 1.0) + df["BsmtHalfBath"] = synth_num(0.057534246575342465, 0.23875264627921178, 1.0) + df["HalfBath"] = synth_num(0.38287671232876713, 0.5028853810928914, 1.0) + df["GarageCars"] = synth_num(1.7671232876712328, 0.7473150101111095, 1.0) + df["TotRmsAbvGrd"] = synth_num(6.517808219178082, 1.6253932905840505, 1.0) + df["BsmtFinSF1"] = synth_num(443.6397260273973, 456.0980908409277, 1.0) + df["BsmtFinSF2"] = synth_num(46.54931506849315, 161.31927280654173, 1.0) + df["GrLivArea"] = synth_num(1515.463698630137, 525.4803834232025, 1.0) + df["ScreenPorch"] = synth_num(15.060958904109588, 55.757415281874174, 1.0) + + columns = list(df.columns) + rng.shuffle(columns) + df = df[columns] + + # linear interaction for testing purposes. + y = np.zeros(shape=(n_samples,)) + for c in df.columns: + if isinstance(df[c].dtype, pd.CategoricalDtype): + y += df[c].cat.codes.astype(np.float64) + else: + y += df[c].values + + # Shift and scale to match the original y. + y *= 79442.50288288662 / y.std() + y += 180921.19589041095 - y.mean() + + return df, y + + +@memory.cache +def get_mq2008( + dpath: str, +) -> Tuple[ + sparse.csr_matrix, + np.ndarray, + np.ndarray, + sparse.csr_matrix, + np.ndarray, + np.ndarray, + sparse.csr_matrix, + np.ndarray, + np.ndarray, +]: + """Fetch the mq2008 dataset.""" + datasets = pytest.importorskip("sklearn.datasets") + src = "https://s3-us-west-2.amazonaws.com/xgboost-examples/MQ2008.zip" + target = os.path.join(dpath, "MQ2008.zip") + if not os.path.exists(target): + request.urlretrieve(url=src, filename=target) + + with zipfile.ZipFile(target, "r") as f: + f.extractall(path=dpath) + + ( + x_train, + y_train, + qid_train, + x_test, + y_test, + qid_test, + x_valid, + y_valid, + qid_valid, + ) = datasets.load_svmlight_files( + ( + os.path.join(dpath, "MQ2008/Fold1/train.txt"), + os.path.join(dpath, "MQ2008/Fold1/test.txt"), + os.path.join(dpath, "MQ2008/Fold1/vali.txt"), + ), + query_id=True, + zero_based=False, + ) + + return ( + x_train, + y_train, + qid_train, + x_test, + y_test, + qid_test, + x_valid, + y_valid, + qid_valid, + ) + + +def make_batches( # pylint: disable=too-many-arguments,too-many-locals + n_samples_per_batch: int, + n_features: int, + n_batches: int, + use_cupy: bool = False, + *, + vary_size: bool = False, + random_state: int = 1994, +) -> Tuple[List[np.ndarray], List[np.ndarray], List[np.ndarray]]: + """Make batches of dense data.""" + X = [] + y = [] + w = [] + if use_cupy: + import cupy + + rng = cupy.random.RandomState(np.uint64(random_state)) + else: + rng = np.random.RandomState(random_state) + for i in range(n_batches): + n_samples = n_samples_per_batch + i * 10 if vary_size else n_samples_per_batch + _X = rng.randn(n_samples, n_features) + _y = rng.randn(n_samples) + _w = rng.uniform(low=0, high=1, size=n_samples) + X.append(_X) + y.append(_y) + w.append(_w) + return X, y, w + + +RelData = Tuple[sparse.csr_matrix, npt.NDArray[np.int32], npt.NDArray[np.int32]] + + +@dataclass +class ClickFold: + """A structure containing information about generated user-click data.""" + + X: sparse.csr_matrix + y: npt.NDArray[np.int32] + qid: npt.NDArray[np.int32] + score: npt.NDArray[np.float32] + click: npt.NDArray[np.int32] + pos: npt.NDArray[np.int64] + + +class RelDataCV(NamedTuple): + """Simple data struct for holding a train-test split of a learning to rank dataset.""" + + train: RelData + test: RelData + max_rel: int + + def is_binary(self) -> bool: + """Whether the label consists of binary relevance degree.""" + return self.max_rel == 1 + + +class PBM: + """Simulate click data with position bias model. There are other models available in + `ULTRA `_ like the cascading model. + + References + ---------- + Unbiased LambdaMART: An Unbiased Pairwise Learning-to-Rank Algorithm + + """ + + def __init__(self, eta: float) -> None: + # click probability for each relevance degree. (from 0 to 4) + self.click_prob = np.array([0.1, 0.16, 0.28, 0.52, 1.0]) + exam_prob = np.array( + [0.68, 0.61, 0.48, 0.34, 0.28, 0.20, 0.11, 0.10, 0.08, 0.06] + ) + # Observation probability, encoding positional bias for each position + self.exam_prob = np.power(exam_prob, eta) + + def sample_clicks_for_query( + self, labels: npt.NDArray[np.int32], position: npt.NDArray[np.int64] + ) -> npt.NDArray[np.int32]: + """Sample clicks for one query based on input relevance degree and position. + + Parameters + ---------- + + labels : + relevance_degree + + """ + labels = np.array(labels, copy=True) + + click_prob = np.zeros(labels.shape) + # minimum + labels[labels < 0] = 0 + # maximum + labels[labels >= len(self.click_prob)] = -1 + click_prob = self.click_prob[labels] + + exam_prob = np.zeros(labels.shape) + assert position.size == labels.size + ranks = np.array(position, copy=True) + # maximum + ranks[ranks >= self.exam_prob.size] = -1 + exam_prob = self.exam_prob[ranks] + + rng = np.random.default_rng(1994) + prob = rng.random(size=labels.shape[0], dtype=np.float32) + + clicks: npt.NDArray[np.int32] = np.zeros(labels.shape, dtype=np.int32) + clicks[prob < exam_prob * click_prob] = 1 + return clicks + + +def rlencode(x: npt.NDArray[np.int32]) -> Tuple[npt.NDArray, npt.NDArray, npt.NDArray]: + """Run length encoding using numpy, modified from: + https://gist.github.com/nvictus/66627b580c13068589957d6ab0919e66 + + """ + x = np.asarray(x) + n = x.size + starts = np.r_[0, np.flatnonzero(~np.isclose(x[1:], x[:-1], equal_nan=True)) + 1] + lengths = np.diff(np.r_[starts, n]) + values = x[starts] + indptr = np.append(starts, np.array([x.size])) + + return indptr, lengths, values + + +def init_rank_score( + X: sparse.csr_matrix, + y: npt.NDArray[np.int32], + qid: npt.NDArray[np.int32], + sample_rate: float = 0.1, +) -> npt.NDArray[np.float32]: + """We use XGBoost to generate the initial score instead of SVMRank for + simplicity. Sample rate is set to 0.1 by default so that we can test with small + datasets. + + """ + # random sample + rng = np.random.default_rng(1994) + n_samples = int(X.shape[0] * sample_rate) + index: npt.NDArray = np.arange(0, X.shape[0], dtype=np.uint64) + rng.shuffle(index) + index = index[:n_samples] + + X_train = X[index] + y_train = y[index] + qid_train = qid[index] + + # Sort training data based on query id, required by XGBoost. + sorted_idx = np.argsort(qid_train) + X_train = X_train[sorted_idx] + y_train = y_train[sorted_idx] + qid_train = qid_train[sorted_idx] + + ltr = XGBRanker(objective="rank:ndcg", tree_method="hist") + ltr.fit(X_train, y_train, qid=qid_train) + + # Use the original order of the data. + scores = ltr.predict(X) + return scores + + +def simulate_one_fold( + fold: Tuple[sparse.csr_matrix, npt.NDArray[np.int32], npt.NDArray[np.int32]], + scores_fold: npt.NDArray[np.float32], +) -> ClickFold: + """Simulate clicks for one fold.""" + X_fold, y_fold, qid_fold = fold + assert qid_fold.dtype == np.int32 + + qids = np.unique(qid_fold) + + position = np.empty((y_fold.size,), dtype=np.int64) + clicks = np.empty((y_fold.size,), dtype=np.int32) + pbm = PBM(eta=1.0) + + # Avoid grouping by qid as we want to preserve the original data partition by + # the dataset authors. + for q in qids: + qid_mask = q == qid_fold + qid_mask = qid_mask.reshape(qid_mask.shape[0]) + query_scores = scores_fold[qid_mask] + # Initial rank list, scores sorted to decreasing order + query_position = np.argsort(query_scores)[::-1] + position[qid_mask] = query_position + # get labels + relevance_degrees = y_fold[qid_mask] + query_clicks = pbm.sample_clicks_for_query(relevance_degrees, query_position) + clicks[qid_mask] = query_clicks + + assert X_fold.shape[0] == qid_fold.shape[0], (X_fold.shape, qid_fold.shape) + assert X_fold.shape[0] == clicks.shape[0], (X_fold.shape, clicks.shape) + + return ClickFold(X_fold, y_fold, qid_fold, scores_fold, clicks, position) + + +# pylint: disable=too-many-locals +def simulate_clicks(cv_data: RelDataCV) -> Tuple[ClickFold, Optional[ClickFold]]: + """Simulate click data using position biased model (PBM).""" + X, y, qid = list(zip(cv_data.train, cv_data.test)) + + # ptr to train-test split + indptr = np.array([0] + [v.shape[0] for v in X]) + indptr = np.cumsum(indptr) + + assert len(indptr) == 2 + 1 # train, test + X_full = sparse.vstack(X) + y_full = np.concatenate(y) + qid_full = np.concatenate(qid) + + # Obtain initial relevance score for click simulation + scores_full = init_rank_score(X_full, y_full, qid_full) + # partition it back to (train, test) tuple + scores = [scores_full[indptr[i - 1] : indptr[i]] for i in range(1, indptr.size)] + + X_lst, y_lst, q_lst, s_lst, c_lst, p_lst = [], [], [], [], [], [] + for i in range(indptr.size - 1): + fold = simulate_one_fold((X[i], y[i], qid[i]), scores[i]) + X_lst.append(fold.X) + y_lst.append(fold.y) + q_lst.append(fold.qid) + s_lst.append(fold.score) + c_lst.append(fold.click) + p_lst.append(fold.pos) + + scores_check_1 = [s_lst[i] for i in range(indptr.size - 1)] + for i in range(2): + assert (scores_check_1[i] == scores[i]).all() + + if len(X_lst) == 1: + train = ClickFold(X_lst[0], y_lst[0], q_lst[0], s_lst[0], c_lst[0], p_lst[0]) + test = None + else: + train, test = ( + ClickFold(X_lst[i], y_lst[i], q_lst[i], s_lst[i], c_lst[i], p_lst[i]) + for i in range(len(X_lst)) + ) + return train, test + + +def make_ltr( + n_samples: int, + n_features: int, + n_query_groups: int, + max_rel: int, + sort_qid: bool = True, +) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + """Make a dataset for testing LTR.""" + rng = np.random.default_rng(1994) + X = rng.normal(0, 1.0, size=n_samples * n_features).reshape(n_samples, n_features) + y = np.sum(X, axis=1) + y -= y.min() + y = np.round(y / y.max() * max_rel).astype(np.int32) + + qid = rng.integers(0, n_query_groups, size=n_samples, dtype=np.int32) + w = rng.normal(0, 1.0, size=n_query_groups) + w -= np.min(w) + w /= np.max(w) + if sort_qid: + qid = np.sort(qid) + return X, y, qid, w + + +def sort_ltr_samples( + X: sparse.csr_matrix, + y: npt.NDArray[np.int32], + qid: npt.NDArray[np.int32], + clicks: npt.NDArray[np.int32], + pos: npt.NDArray[np.int64], +) -> Tuple[ + sparse.csr_matrix, + npt.NDArray[np.int32], + npt.NDArray[np.int32], + npt.NDArray[np.int32], +]: + """Sort data based on query index and position.""" + sorted_idx = np.argsort(qid) + X = X[sorted_idx] + clicks = clicks[sorted_idx] + qid = qid[sorted_idx] + pos = pos[sorted_idx] + + indptr, _, _ = rlencode(qid) + + for i in range(1, indptr.size): + beg = indptr[i - 1] + end = indptr[i] + + assert beg < end, (beg, end) + assert np.unique(qid[beg:end]).size == 1, (beg, end) + + query_pos = pos[beg:end] + assert query_pos.min() == 0, query_pos.min() + assert query_pos.max() >= query_pos.size - 1, ( + query_pos.max(), + query_pos.size, + i, + np.unique(qid[beg:end]), + ) + sorted_idx = np.argsort(query_pos) + + X[beg:end] = X[beg:end][sorted_idx] + clicks[beg:end] = clicks[beg:end][sorted_idx] + y[beg:end] = y[beg:end][sorted_idx] + # not necessary + qid[beg:end] = qid[beg:end][sorted_idx] + + data = X, clicks, y, qid + + return data + + +def run_base_margin_info(DType: Callable, DMatrixT: Type[DMatrix], device: str) -> None: + """Run tests for base margin.""" + rng = np.random.default_rng() + X = DType(rng.normal(0, 1.0, size=100).astype(np.float32).reshape(50, 2)) + if hasattr(X, "iloc"): + y = X.iloc[:, 0] + else: + y = X[:, 0] + base_margin = X + # no error at set + Xy = DMatrixT(X, y, base_margin=base_margin) + # Error at train, caused by check in predictor. + with pytest.raises(ValueError, match=r".*base_margin.*"): + train_fn({"tree_method": "hist", "device": device}, Xy) + + if not hasattr(X, "iloc"): + # column major matrix + got = DType(Xy.get_base_margin().reshape(50, 2)) + assert (got == base_margin).all() + + assert base_margin.T.flags.c_contiguous is False + assert base_margin.T.flags.f_contiguous is True + Xy.set_info(base_margin=base_margin.T) + got = DType(Xy.get_base_margin().reshape(2, 50)) + assert (got == base_margin.T).all() + + # Row vs col vec. + base_margin = y + Xy.set_base_margin(base_margin) + bm_col = Xy.get_base_margin() + Xy.set_base_margin(base_margin.reshape(1, base_margin.size)) + bm_row = Xy.get_base_margin() + assert (bm_row == bm_col).all() + + # type + base_margin = base_margin.astype(np.float64) + Xy.set_base_margin(base_margin) + bm_f64 = Xy.get_base_margin() + assert (bm_f64 == bm_col).all() + + # too many dimensions + base_margin = X.reshape(2, 5, 2, 5) + with pytest.raises(ValueError, match=r".*base_margin.*"): + Xy.set_base_margin(base_margin) + + +# pylint: disable=too-many-locals +@memory.cache +def make_sparse_regression( + n_samples: int, n_features: int, sparsity: float, as_dense: bool +) -> Tuple[Union[sparse.csr_matrix], np.ndarray]: + """Make sparse matrix. + + Parameters + ---------- + + as_dense: + + Return the matrix as np.ndarray with missing values filled by NaN + + """ + if not hasattr(np.random, "default_rng"): + rng = np.random.RandomState(1994) + X = sparse.random( + m=n_samples, + n=n_features, + density=1.0 - sparsity, + random_state=rng, + format="csr", + ) + y = rng.normal(loc=0.0, scale=1.0, size=n_samples) + return X, y + + # Use multi-thread to speed up the generation, convenient if you use this function + # for benchmarking. + n_threads = min(multiprocessing.cpu_count(), n_features) + + def random_csc(t_id: int) -> sparse.csc_matrix: + rng = np.random.default_rng(1994 * t_id) + thread_size = n_features // n_threads + if t_id == n_threads - 1: + n_features_tloc = n_features - t_id * thread_size + else: + n_features_tloc = thread_size + + X = sparse.random( + m=n_samples, + n=n_features_tloc, + density=1.0 - sparsity, + random_state=rng, + ).tocsc() + y = np.zeros((n_samples, 1)) + + for i in range(X.shape[1]): + size = X.indptr[i + 1] - X.indptr[i] + if size != 0: + y += X[:, i].toarray() * rng.random((n_samples, 1)) * 0.2 + + return X, y + + futures = [] + with ThreadPoolExecutor(max_workers=n_threads) as executor: + for i in range(n_threads): + futures.append(executor.submit(random_csc, i)) + + X_results = [] + y_results = [] + for f in futures: + X, y = f.result() + X_results.append(X) + y_results.append(y) + + assert len(y_results) == n_threads + + csr: sparse.csr_matrix = sparse.hstack(X_results, format="csr") + y = np.asarray(y_results) + y = y.reshape((y.shape[0], y.shape[1])).T + y = np.sum(y, axis=1) + + assert csr.shape[0] == n_samples + assert csr.shape[1] == n_features + assert y.shape[0] == n_samples + + if as_dense: + arr = csr.toarray() + assert arr.shape[0] == n_samples + assert arr.shape[1] == n_features + arr[arr == 0] = np.nan + return arr, y + + return csr, y + + +def unique_random_strings(n_strings: int, seed: int) -> List[str]: + """Generate n unique strings.""" + name_len = 8 # hardcoded, should be more than enough + unique_strings: Set[str] = set() + rng = np.random.default_rng(seed) + + while len(unique_strings) < n_strings: + random_str = "".join( + rng.choice(list(string.ascii_letters), size=name_len, replace=True) + ) + unique_strings.add(random_str) + + return list(unique_strings) + + +# pylint: disable=too-many-arguments,too-many-locals,too-many-branches +def make_categorical( + n_samples: int, + n_features: int, + n_categories: int, + *, + onehot: bool, + n_targets: int = 1, + sparsity: float = 0.0, + cat_ratio: float = 1.0, + shuffle: bool = False, + random_state: int = 1994, + cat_dtype: np.typing.DTypeLike = np.int64, + device: str = "cpu", +) -> Tuple[ArrayLike, np.ndarray]: + """Generate categorical features for test. + + Parameters + ---------- + n_categories: + Number of categories for categorical features. + onehot: + Should we apply one-hot encoding to the data? + n_targets: + Number of targets. When greater than 1, the label is a 2D array with shape + ``(n_samples, n_targets)``. + sparsity: + The ratio of the amount of missing values over the number of all entries. + cat_ratio: + The ratio of features that are categorical. + shuffle: + Whether we should shuffle the columns. + cat_dtype : + The dtype for categorical features, might be string or numeric. + + Returns + ------- + X, y + """ + pd = pytest.importorskip("pandas") + + # Use different rngs for column and rows. We can change the `n_samples` without + # changing the column type. + rng = np.random.RandomState(random_state) + row_rng = np.random.RandomState(random_state + 1) + + df = pd.DataFrame() + for i in range(n_features): + choice = rng.binomial(1, cat_ratio, size=1)[0] + if choice == 1: + if np.issubdtype(cat_dtype, np.str_): + # we rely on using the feature index as the seed to generate the same + # categories for multiple calls to `make_categorical`. + categories = np.array(unique_random_strings(n_categories, i)) + c = row_rng.choice(categories, size=n_samples, replace=True) + else: + categories = np.arange(0, n_categories) + c = row_rng.randint(low=0, high=n_categories, size=n_samples) + + df[str(i)] = pd.Series(c, dtype="category") + df[str(i)] = df[str(i)].cat.set_categories(categories) + else: + num = row_rng.randint(low=0, high=n_categories, size=n_samples) + df[str(i)] = pd.Series(num, dtype=num.dtype) + + target_rng = np.random.RandomState(random_state + 2) + label: np.ndarray = np.ones((n_samples, n_targets)) + for col in df.columns: + if isinstance(df[col].dtype, pd.CategoricalDtype): + codes = df[col].cat.codes.values + effects = target_rng.normal(size=(len(df[col].cat.categories), n_targets)) + label += effects[codes] + else: + w = target_rng.uniform(low=0.5, high=1.5, size=n_targets) + label += np.outer(df[col].values, w) + if n_targets == 1: + label = label.squeeze(axis=1) + + if sparsity > 0.0: + for i in range(n_features): + index = row_rng.randint( + low=0, high=n_samples - 1, size=int(n_samples * sparsity) + ) + df.iloc[index, i] = np.nan + if is_pd_cat_dtype(df.dtypes.iloc[i]): + assert n_categories == np.unique(df.dtypes.iloc[i].categories).size + + assert df.shape[1] == n_features + if onehot: + df = pd.get_dummies(df) + + if shuffle: + columns = list(df.columns) + row_rng.shuffle(columns) + df = df[columns] + + if device != "cpu": + assert device in ["cuda", "gpu"] + import cudf + import cupy + + df = cudf.from_pandas(df) + label = cupy.array(label) + return df, label + + +class IteratorForTest(DataIter): + """Iterator for testing streaming DMatrix. (external memory, quantile)""" + + def __init__( # pylint: disable=too-many-arguments + self, + X: Sequence, + y: Sequence, + w: Optional[Sequence], + *, + cache: Optional[str], + on_host: bool = False, + min_cache_page_bytes: Optional[int] = None, + ) -> None: + assert len(X) == len(y) + self.X = X + self.y = y + self.w = w + self.it = 0 + super().__init__( + cache_prefix=cache, + on_host=on_host, + min_cache_page_bytes=min_cache_page_bytes, + ) + + def next(self, input_data: Callable) -> bool: + if self.it == len(self.X): + return False + + with pytest.raises(TypeError, match="Keyword argument"): + input_data(self.X[self.it], self.y[self.it], None) + + # Use copy to make sure the iterator doesn't hold a reference to the data. + input_data( + data=self.X[self.it].copy(), + label=self.y[self.it].copy(), + weight=self.w[self.it].copy() if self.w else None, + ) + gc.collect() # clear up the copy, see if XGBoost access freed memory. + self.it += 1 + return True + + def reset(self) -> None: + self.it = 0 + + def as_arrays( + self, + ) -> Tuple[Union[np.ndarray, sparse.csr_matrix], ArrayLike, Optional[ArrayLike]]: + """Return concatenated arrays.""" + X = concat(self.X) + y = concat(self.y) + if self.w: + w = np.concatenate(self.w, axis=0) + else: + w = None + return X, y, w diff --git a/v12x/pylib/xgboost/testing/data_iter.py b/v12x/pylib/xgboost/testing/data_iter.py new file mode 100644 index 0000000000000000000000000000000000000000..6a805af2b5499874e1f359b8df0875bc7833bb91 --- /dev/null +++ b/v12x/pylib/xgboost/testing/data_iter.py @@ -0,0 +1,176 @@ +"""Tests related to the `DataIter` interface.""" + +from typing import Callable, Optional + +import numpy as np + +from xgboost import testing as tm + +from ..compat import import_cupy +from ..core import DataIter, DMatrix, ExtMemQuantileDMatrix, QuantileDMatrix +from .utils import predictor_equal + + +def run_mixed_sparsity(device: str) -> None: + """Check QDM with mixed batches.""" + X_0, y_0, _ = tm.make_regression(128, 16, False) + if device.startswith("cuda"): + X_1, y_1 = tm.make_sparse_regression(256, 16, 0.1, True) + else: + X_1, y_1 = tm.make_sparse_regression(256, 16, 0.1, False) + X_2, y_2 = tm.make_sparse_regression(512, 16, 0.9, True) + X = [X_0, X_1, X_2] + y = [y_0, y_1, y_2] + + if device.startswith("cuda"): + cp = import_cupy() + + X = [cp.array(batch) for batch in X] + + it = tm.IteratorForTest(X, y, None, cache=None, on_host=False) + Xy_0 = QuantileDMatrix(it) + + X_1, y_1 = tm.make_sparse_regression(256, 16, 0.1, True) + X = [X_0, X_1, X_2] + y = [y_0, y_1, y_2] + X_arr = np.concatenate(X, axis=0) + y_arr = np.concatenate(y, axis=0) + Xy_1 = QuantileDMatrix(X_arr, y_arr) + + assert predictor_equal(Xy_0, Xy_1) + + +def check_invalid_cat_batches(device: str) -> None: + """Check error message for inconsistent feature types.""" + + class _InvalidCatIter(DataIter): + def __init__(self) -> None: + super().__init__(cache_prefix=None) + self._it = 0 + + def next(self, input_data: Callable) -> bool: + if self._it == 2: + return False + X, y = tm.make_categorical( + 64, + 12, + 4, + onehot=False, + sparsity=0.5, + cat_ratio=1.0 if self._it == 0 else 0.5, + ) + if device == "cuda": + import cudf + import cupy + + X = cudf.DataFrame(X) + y = cupy.array(y) + + input_data(data=X, label=y) + self._it += 1 + return True + + def reset(self) -> None: + self._it = 0 + + it = _InvalidCatIter() + import pytest + + with pytest.raises(ValueError, match="Inconsistent number of categories between"): + ExtMemQuantileDMatrix(it, enable_categorical=True) + + with pytest.raises(ValueError, match="Inconsistent number of categories between"): + QuantileDMatrix(it, enable_categorical=True) + + with pytest.raises(ValueError, match="Inconsistent feature types"): + DMatrix(it, enable_categorical=True) + + +def check_uneven_sizes(device: str) -> None: + """Tests for having irregular data shapes.""" + batches = [ + tm.make_regression(n_samples, 16, use_cupy=device == "cuda") + for n_samples in [512, 256, 1024] + ] + unzip = list(zip(*batches)) + it = tm.IteratorForTest(unzip[0], unzip[1], None, cache="cache", on_host=True) + + Xy = DMatrix(it) + assert Xy.num_col() == 16 + assert Xy.num_row() == sum(x.shape[0] for x in unzip[0]) + + Xy = ExtMemQuantileDMatrix(it) + assert Xy.num_col() == 16 + assert Xy.num_row() == sum(x.shape[0] for x in unzip[0]) + + +class CatIter(DataIter): # pylint: disable=too-many-instance-attributes + """An iterator for testing categorical features.""" + + def __init__( # pylint: disable=too-many-arguments,too-many-locals + self, + n_samples_per_batch: int, + n_features: int, + *, + n_batches: int, + n_cats: int, + sparsity: float, + cat_ratio: float, + onehot: bool, + device: str, + cache: Optional[str], + n_targets: int = 1, + ) -> None: + super().__init__(cache_prefix=cache) + self.n_batches = n_batches + self.device = device + + n_samples = n_samples_per_batch * n_batches + cat, y = tm.make_categorical( + n_samples, + n_features, + n_categories=n_cats, + onehot=onehot, + cat_ratio=cat_ratio, + sparsity=sparsity, + n_targets=n_targets, + ) + xs, ys = [], [] + + prev = 0 + for _ in range(n_batches): + n = min(n_samples_per_batch, n_samples - prev) + X = cat.iloc[prev : prev + n, :] + xs.append(X) + ys.append(y[prev : prev + n]) + prev += n_samples_per_batch + + self.xs = xs + self.ys = ys + + self.x = cat + self.y = y + + self._it = 0 + + def xy(self) -> tuple: + """Return the concatenated data.""" + return self.x, self.y + + def next(self, input_data: Callable) -> bool: + if self._it == self.n_batches: + return False + + X, y = self.xs[self._it], self.ys[self._it] + if self.device == "cuda": + import cudf + import cupy + + X = cudf.DataFrame(X) + y = cupy.array(y) + input_data(data=X, label=y) + self._it += 1 + return True + + def reset(self) -> None: + self._it = 0 diff --git a/v12x/pylib/xgboost/testing/federated.py b/v12x/pylib/xgboost/testing/federated.py new file mode 100644 index 0000000000000000000000000000000000000000..981a124b8a836d927f9871383b204b798a5b06f8 --- /dev/null +++ b/v12x/pylib/xgboost/testing/federated.py @@ -0,0 +1,155 @@ +# pylint: disable=unbalanced-tuple-unpacking, too-many-locals +"""Tests for federated learning.""" + +import multiprocessing +import os +import subprocess +import tempfile +import time +from typing import List, cast + +from sklearn.datasets import dump_svmlight_file, load_svmlight_file +from sklearn.model_selection import train_test_split + +import xgboost as xgb +import xgboost.federated +from xgboost import testing as tm + +from .._typing import EvalsLog +from ..collective import _Args as CollArgs + +SERVER_KEY = "server-key.pem" +SERVER_CERT = "server-cert.pem" +CLIENT_KEY = "client-key.pem" +CLIENT_CERT = "client-cert.pem" + + +def run_server(port: int, world_size: int, with_ssl: bool) -> None: + """Run federated server for test.""" + if with_ssl: + xgboost.federated.run_federated_server( + world_size, + port, + server_key_path=SERVER_KEY, + server_cert_path=SERVER_CERT, + client_cert_path=CLIENT_CERT, + ) + else: + xgboost.federated.run_federated_server(world_size, port) + + +def run_worker( + port: int, world_size: int, rank: int, with_ssl: bool, device: str +) -> None: + """Run federated client worker for test.""" + comm_env: CollArgs = { + "dmlc_communicator": "federated", + "federated_server_address": f"localhost:{port}", + "federated_world_size": world_size, + "federated_rank": rank, + } + if with_ssl: + comm_env["federated_server_cert_path"] = SERVER_CERT + comm_env["federated_client_key_path"] = CLIENT_KEY + comm_env["federated_client_cert_path"] = CLIENT_CERT + + cpu_count = os.cpu_count() + assert cpu_count is not None + n_threads = cpu_count // world_size + + # Always call this before using distributed module + with xgb.collective.CommunicatorContext(**comm_env): + # Load file, file will not be sharded in federated mode. + X, y = load_svmlight_file(f"agaricus.txt-{rank}.train") + dtrain = xgb.DMatrix(X, y) + X, y = load_svmlight_file(f"agaricus.txt-{rank}.test") + dtest = xgb.DMatrix(X, y) + + # Specify parameters via map, definition are same as c++ version + param = { + "max_depth": 2, + "eta": 1, + "objective": "binary:logistic", + "nthread": n_threads, + "tree_method": "hist", + "device": device, + } + + # Specify validations set to watch performance + watchlist = [(dtest, "eval"), (dtrain, "train")] + num_round = 20 + + # Run training, all the features in training API is available. + results: EvalsLog = {} + bst = xgb.train( + param, + dtrain, + num_round, + evals=watchlist, + early_stopping_rounds=2, + evals_result=results, + ) + assert tm.non_increasing(cast(List[float], results["train"]["logloss"])) + assert tm.non_increasing(cast(List[float], results["eval"]["logloss"])) + + # save the model, only ask process 0 to save the model. + if xgb.collective.get_rank() == 0: + with tempfile.TemporaryDirectory() as tmpdir: + bst.save_model(os.path.join(tmpdir, "model.json")) + xgb.collective.communicator_print("Finished training\n") + + +def run_federated(world_size: int, with_ssl: bool, use_gpu: bool) -> None: + """Launcher for clients and the server.""" + port = 9091 + + server = multiprocessing.Process( + target=run_server, args=(port, world_size, with_ssl) + ) + server.start() + time.sleep(1) + if not server.is_alive(): + raise ValueError("Error starting Federated Learning server") + + workers = [] + for rank in range(world_size): + device = f"cuda:{rank}" if use_gpu else "cpu" + worker = multiprocessing.Process( + target=run_worker, args=(port, world_size, rank, with_ssl, device) + ) + workers.append(worker) + worker.start() + for worker in workers: + worker.join() + server.terminate() + + +def run_federated_learning(with_ssl: bool, use_gpu: bool, test_path: str) -> None: + """Run federated learning tests.""" + n_workers = 2 + + if with_ssl: + command = "openssl req -x509 -newkey rsa:2048 -days 7 -nodes -keyout {part}-key.pem -out {part}-cert.pem -subj /C=US/CN=localhost" # pylint: disable=line-too-long + server_key = command.format(part="server").split() + subprocess.check_call(server_key) + client_key = command.format(part="client").split() + subprocess.check_call(client_key) + + train_path = os.path.join(tm.data_dir(test_path), "agaricus.txt.train") + test_path = os.path.join(tm.data_dir(test_path), "agaricus.txt.test") + + X_train, y_train = load_svmlight_file(train_path) + X_test, y_test = load_svmlight_file(test_path) + + X0, X1, y0, y1 = train_test_split(X_train, y_train, test_size=0.5) + X0_valid, X1_valid, y0_valid, y1_valid = train_test_split( + X_test, y_test, test_size=0.5 + ) + + dump_svmlight_file(X0, y0, "agaricus.txt-0.train") + dump_svmlight_file(X0_valid, y0_valid, "agaricus.txt-0.test") + + dump_svmlight_file(X1, y1, "agaricus.txt-1.train") + dump_svmlight_file(X1_valid, y1_valid, "agaricus.txt-1.test") + + run_federated(world_size=n_workers, with_ssl=with_ssl, use_gpu=use_gpu) diff --git a/v12x/pylib/xgboost/testing/interaction_constraints.py b/v12x/pylib/xgboost/testing/interaction_constraints.py new file mode 100644 index 0000000000000000000000000000000000000000..70e8f8c125897b00f57f069679b61ed2c53462e6 --- /dev/null +++ b/v12x/pylib/xgboost/testing/interaction_constraints.py @@ -0,0 +1,106 @@ +"""Tests for interaction constraints.""" + +from typing import Optional, Sequence, Union + +import numpy as np + +from .._typing import FeatureNames +from ..core import DMatrix +from ..training import train +from .utils import Device + + +def run_interaction_constraints( # pylint: disable=too-many-locals + tree_method: str, + device: Device, + feature_names: Optional[FeatureNames] = None, + interaction_constraints: Union[str, Sequence] = "[[0, 1]]", + n_targets: int = 1, +) -> None: + """Tests interaction constraints on a synthetic dataset. Only x1 and x2 are allowed + to interact; x3 must stay additive. + + """ + rng = np.random.default_rng(2026) + n_samples = 1000 + x1 = rng.normal(loc=1.0, scale=1.0, size=n_samples) + x2 = rng.normal(loc=1.0, scale=1.0, size=n_samples) + x3 = rng.choice([1, 2, 3], size=n_samples, replace=True) + X = np.column_stack((x1, x2, x3)) + + # The constraint must force x3's contribution to remain additive regardless. + shared = ( + x1 + + x2 + + x1 * x2 * x3 + + 3.0 * np.sin(x1) + + rng.normal(loc=0.0, scale=1.0, size=n_samples) + ) + # Each target gets a distinct, non-zero additive coef on x3. + slopes = np.arange(1, n_targets + 1, dtype=np.float64) + slopes[1::2] *= -1.0 + targets = [shared + slope * x3 for slope in slopes] + y = targets[0] if n_targets == 1 else np.column_stack(targets) + + dtrain = DMatrix(X, label=y, feature_names=feature_names) + params = { + "max_depth": 3, + "eta": 0.1, + "nthread": 2, + "interaction_constraints": interaction_constraints, + "tree_method": tree_method, + "device": device, + } + if n_targets > 1: + params["multi_strategy"] = "multi_output_tree" + + num_boost_round = 12 + # Fit a model that only allows interaction between x1 and x2. + bst = train(params, dtrain, num_boost_round, evals=[(dtrain, "train")]) + + # Set all observations to have the same x3 value then increment by the same amount. + def f(x: int) -> np.ndarray: + tmat = DMatrix( + np.column_stack((x1, x2, np.repeat(x, n_samples))), + feature_names=feature_names, + ) + return bst.predict(tmat) + + preds = [f(x) for x in [1, 2, 3]] + expected_shape = (n_samples,) if n_targets == 1 else (n_samples, n_targets) + assert preds[0].shape == expected_shape + + # Check incrementing x3 has the same effect on all observations + # since x3 is constrained to be independent of x1 and x2 + # and all observations start off from the same x3 value + diff1 = preds[1] - preds[0] + assert np.all(np.abs(diff1 - diff1[0]) < 1e-4) + diff2 = preds[2] - preds[1] + assert np.all(np.abs(diff2 - diff2[0]) < 1e-4) + + +def training_accuracy(tree_method: str, dpath: str, device: Device) -> None: + """Test accuracy, reused by GPU tests.""" + from sklearn.metrics import accuracy_score + + dtrain = DMatrix(dpath + "agaricus.txt.train?indexing_mode=1&format=libsvm") + dtest = DMatrix(dpath + "agaricus.txt.test?indexing_mode=1&format=libsvm") + params = { + "eta": 1, + "max_depth": 6, + "objective": "binary:logistic", + "tree_method": tree_method, + "device": device, + "interaction_constraints": "[[1,2], [2,3,4]]", + } + num_boost_round = 5 + + params["grow_policy"] = "lossguide" + bst = train(params, dtrain, num_boost_round) + pred_dtest = bst.predict(dtest) < 0.5 + assert accuracy_score(dtest.get_label(), pred_dtest) < 0.1 + + params["grow_policy"] = "depthwise" + bst = train(params, dtrain, num_boost_round) + pred_dtest = bst.predict(dtest) < 0.5 + assert accuracy_score(dtest.get_label(), pred_dtest) < 0.1 diff --git a/v12x/pylib/xgboost/testing/intercept.py b/v12x/pylib/xgboost/testing/intercept.py new file mode 100644 index 0000000000000000000000000000000000000000..15695acd5a45e80c20f4e62470746fb5ba2a24ed --- /dev/null +++ b/v12x/pylib/xgboost/testing/intercept.py @@ -0,0 +1,263 @@ +"""Tests for estimating the intercept.""" + +import json +from typing import Dict, List, Optional + +import numpy as np +from scipy.special import softmax +from sklearn.datasets import ( + make_classification, + make_multilabel_classification, + make_regression, +) + +from ..core import Booster, DMatrix, QuantileDMatrix +from ..sklearn import XGBClassifier, XGBRegressor +from ..training import train +from .updater import get_basescore +from .utils import Device, non_increasing + + +# pylint: disable=too-many-statements +def run_init_estimation(tree_method: str, device: Device) -> None: + """Test for init estimation.""" + + def run_reg(X: np.ndarray, y: np.ndarray) -> None: # pylint: disable=invalid-name + reg = XGBRegressor( + tree_method=tree_method, max_depth=1, n_estimators=1, device=device + ) + reg.fit(X, y, eval_set=[(X, y)]) + base_score_0 = get_basescore(reg) + score_0 = reg.evals_result()["validation_0"]["rmse"][0] + + n_targets = 1 if y.ndim == 1 else y.shape[1] + intercept = np.full(shape=(n_targets,), fill_value=0.5, dtype=np.float32) + reg = XGBRegressor( + tree_method=tree_method, + device=device, + max_depth=1, + n_estimators=1, + base_score=intercept, + ) + reg.fit(X, y, eval_set=[(X, y)]) + base_score_1 = get_basescore(reg) + score_1 = reg.evals_result()["validation_0"]["rmse"][0] + assert not np.isclose(base_score_0, base_score_1).any() + assert score_0 < score_1 # should be better + + # pylint: disable=unbalanced-tuple-unpacking + X, y = make_regression(n_samples=4096, random_state=17) + run_reg(X, y) + # pylint: disable=unbalanced-tuple-unpacking + X, y = make_regression(n_samples=4096, n_targets=3, random_state=17) + run_reg(X, y) + + # pylint: disable=invalid-name + def run_clf( + X: np.ndarray, y: np.ndarray, w: Optional[np.ndarray] = None + ) -> List[float]: + clf = XGBClassifier( + tree_method=tree_method, max_depth=1, n_estimators=1, device=device + ) + if w is not None: + clf.fit( + X, y, sample_weight=w, eval_set=[(X, y)], sample_weight_eval_set=[w] + ) + else: + clf.fit(X, y, eval_set=[(X, y)]) + base_score_0 = get_basescore(clf) + if clf.n_classes_ == 2: + score_0 = clf.evals_result()["validation_0"]["logloss"][0] + else: + score_0 = clf.evals_result()["validation_0"]["mlogloss"][0] + + n_targets = 1 if y.ndim == 1 else y.shape[1] + intercept = np.full(shape=(n_targets,), fill_value=0.5, dtype=np.float32) + clf = XGBClassifier( + tree_method=tree_method, + max_depth=1, + n_estimators=1, + device=device, + base_score=intercept, + ) + if w is not None: + clf.fit( + X, y, sample_weight=w, eval_set=[(X, y)], sample_weight_eval_set=[w] + ) + else: + clf.fit(X, y, eval_set=[(X, y)]) + base_score_1 = get_basescore(clf) + if clf.n_classes_ == 2: + score_1 = clf.evals_result()["validation_0"]["logloss"][0] + else: + score_1 = clf.evals_result()["validation_0"]["mlogloss"][0] + assert not np.isclose(base_score_0, base_score_1).any() + assert score_0 < score_1 + 1e-4 # should be better + + return base_score_0 + + # pylint: disable=unbalanced-tuple-unpacking + X, y = make_classification(n_samples=4096, random_state=17) + run_clf(X, y) + X, y = make_multilabel_classification( + n_samples=4096, n_labels=3, n_classes=5, random_state=17 + ) + run_clf(X, y) + + # Extra tests for the classifier. + X, y = make_classification( + n_samples=4096, random_state=17, n_classes=5, n_informative=20, n_redundant=0 + ) + intercept = run_clf(X, y) + # un-transformed intercept sums to 0, as a convention. + np.testing.assert_allclose(np.sum(softmax(intercept)), 1.0) + np.testing.assert_allclose(np.sum(intercept), 0.0, atol=1e-6) + + assert np.all(softmax(intercept) > 0) + np_int = ( + np.histogram( + y, bins=np.concatenate([np.unique(y), np.array([np.finfo(np.float32).max])]) + )[0] + / y.shape[0] + ) + np.testing.assert_allclose(softmax(intercept), np_int, rtol=1e-6) + + rng = np.random.default_rng(1994) + w = rng.uniform(low=0, high=1, size=(y.shape[0],)) + intercept = run_clf(X, y, w) + np.testing.assert_allclose(np.sum(softmax(intercept)), 1.0) + assert np.all(softmax(intercept) > 0) + + +# pylint: disable=too-many-locals +def run_adaptive(tree_method: str, weighted: bool, device: Device) -> None: + """Test for adaptive trees.""" + rng = np.random.RandomState(1994) + n_samples = 256 + X, y = make_regression( # pylint: disable=unbalanced-tuple-unpacking + n_samples, 16, random_state=rng + ) + if weighted: + w = rng.normal(size=n_samples) + w -= w.min() + Xy = DMatrix(X, y, weight=w) + + else: + Xy = DMatrix(X, y) + w = np.ones_like(y) + + mean = np.average(y, weights=w) + residual = mean - y + delta = np.average(np.sqrt(np.abs(residual)), weights=w) ** 2 + norm = np.hypot(delta, residual) + curvature = np.divide(delta, norm, out=np.ones_like(norm), where=norm > 0.0) + base_score = mean - np.sum(w * residual * curvature) / np.sum(w * curvature) + + # Check the base score is expected. + booster_0 = train( + { + "tree_method": tree_method, + "base_score": base_score, + "objective": "reg:absoluteerror", + "device": device, + }, + Xy, + num_boost_round=1, + ) + booster_1 = train( + { + "tree_method": tree_method, + "objective": "reg:absoluteerror", + "device": device, + }, + Xy, + num_boost_round=1, + ) + config_0 = json.loads(booster_0.save_config()) + config_1 = json.loads(booster_1.save_config()) + + np.testing.assert_allclose( + get_basescore(config_0), get_basescore(config_1), rtol=1e-6 + ) + + # check the base score is correctly serialized. + raw_booster = booster_1.save_raw(raw_format="ubj") + booster_2 = Booster(model_file=raw_booster) + config_2 = json.loads(booster_2.save_config()) + assert get_basescore(config_1) == get_basescore(config_2) + + # check we can override the base score. + booster_0 = train( + { + "tree_method": tree_method, + "base_score": base_score + 1.0, + "objective": "reg:absoluteerror", + "device": device, + }, + Xy, + num_boost_round=1, + ) + config_0 = json.loads(booster_0.save_config()) + np.testing.assert_allclose( + get_basescore(config_0), np.asarray(get_basescore(config_1)) + 1 + ) + + # check we can use subsampling. + evals_result: Dict[str, Dict[str, list]] = {} + train( + { + "tree_method": tree_method, + "device": device, + "objective": "reg:absoluteerror", + "subsample": 0.8, + "eta": 1.0, + }, + Xy, + num_boost_round=10, + evals=[(Xy, "Train")], + evals_result=evals_result, + ) + mae = evals_result["Train"]["mae"] + assert mae[-1] < 20.0 + assert non_increasing(mae) + + +def run_exp_family(device: Device) -> None: + """Exp family has a closed solution.""" + X, y = make_classification(n_samples=128, n_classes=2, weights=[0.8, 0.2]) + Xy = QuantileDMatrix(X, y) + clf = train( + {"objective": "binary:logistic", "device": device}, Xy, num_boost_round=1 + ) + reg = train({"objective": "reg:logistic", "device": device}, Xy, num_boost_round=1) + clf1 = train( + {"objective": "binary:logitraw", "device": device}, Xy, num_boost_round=1 + ) + # The base score stored in the booster model is un-transformed + np.testing.assert_allclose([get_basescore(m) for m in (reg, clf, clf1)], y.mean()) + + X, y = make_classification(weights=[0.8, 0.2], random_state=2025) + clf = train( + {"objective": "binary:logistic", "scale_pos_weight": 4.0, "device": device}, + QuantileDMatrix(X, y), + num_boost_round=1, + ) + score = get_basescore(clf) + np.testing.assert_allclose(score, 0.5, rtol=1e-3) + + +def run_logistic_degenerate(device: Device) -> None: + """Test https://github.com/dmlc/xgboost/issues/11499 .""" + + def run(v: float) -> None: + dtrain = DMatrix(np.asarray([[1.0], [1.0]]), label=[v, v]) + bst = train( + {"objective": "binary:logistic", "device": device}, + dtrain, + 1, + ) + intercept = get_basescore(bst) + assert intercept[0] == v + + run(0.0) + run(1.0) diff --git a/v12x/pylib/xgboost/testing/metrics.py b/v12x/pylib/xgboost/testing/metrics.py new file mode 100644 index 0000000000000000000000000000000000000000..6571cd37a7aba5fce96b13a55013945ae83a12a2 --- /dev/null +++ b/v12x/pylib/xgboost/testing/metrics.py @@ -0,0 +1,406 @@ +"""Tests for evaluation metrics.""" + +from typing import Any, Dict, List, Optional + +import numpy as np +import pytest + +from ..compat import concat +from ..core import DMatrix, QuantileDMatrix, _parse_eval_str +from ..sklearn import XGBClassifier, XGBRanker +from ..training import train +from .utils import Device + + +def check_precision_score( # pylint: disable=too-many-locals + tree_method: str, device: Device +) -> None: + """Test for precision with ranking and classification.""" + datasets = pytest.importorskip("sklearn.datasets") + + X, y = datasets.make_classification( + n_samples=1024, n_features=4, n_classes=2, random_state=2023 + ) + qid = np.zeros(shape=y.shape) # same group + + ltr = XGBRanker(n_estimators=2, tree_method=tree_method, device=device) + ltr.fit(X, y, qid=qid) + + # re-generate so that XGBoost doesn't evaluate the result to 1.0 + X, y = datasets.make_classification( + n_samples=512, n_features=4, n_classes=2, random_state=1994 + ) + + ltr.set_params(eval_metric="pre@32") + result = _parse_eval_str(ltr.get_booster().eval_set(evals=[(DMatrix(X, y), "Xy")])) + score_0 = result[1][1] + + X_list = [] + y_list = [] + n_query_groups = 3 + q_list: List[np.ndarray] = [] + for i in range(n_query_groups): + # same for all groups + X, y = datasets.make_classification( + n_samples=512, n_features=4, n_classes=2, random_state=1994 + ) + X_list.append(X) + y_list.append(y) + q = np.full(shape=y.shape, fill_value=i, dtype=np.uint64) + q_list.append(q) + + qid = concat(q_list) + X = concat(X_list) + y = concat(y_list) + + result = _parse_eval_str( + ltr.get_booster().eval_set(evals=[(DMatrix(X, y, qid=qid), "Xy")]) + ) + assert result[1][0].endswith("pre@32") + score_1 = result[1][1] + assert score_1 == score_0 + + +def check_quantile_error(tree_method: str, device: Device) -> None: + """Test for the `quantile` loss.""" + from sklearn.datasets import make_regression + from sklearn.metrics import mean_pinball_loss + + rng = np.random.RandomState(19) + # pylint: disable=unbalanced-tuple-unpacking + X, y = make_regression(128, 3, random_state=rng) + Xy = QuantileDMatrix(X, y) + evals_result: Dict[str, Dict] = {} + booster = train( + { + "tree_method": tree_method, + "eval_metric": "quantile", + "quantile_alpha": 0.3, + "device": device, + }, + Xy, + evals=[(Xy, "Train")], + evals_result=evals_result, + ) + predt = booster.inplace_predict(X) + loss = mean_pinball_loss(y, predt, alpha=0.3) + np.testing.assert_allclose(evals_result["Train"]["quantile"][-1], loss) + + alpha = [0.25, 0.5, 0.75] + booster = train( + { + "tree_method": tree_method, + "eval_metric": "quantile", + "quantile_alpha": alpha, + "objective": "reg:quantileerror", + "device": device, + }, + Xy, + evals=[(Xy, "Train")], + evals_result=evals_result, + ) + predt = booster.inplace_predict(X) + loss = np.mean( + [mean_pinball_loss(y, predt[:, i], alpha=alpha[i]) for i in range(3)] + ) + np.testing.assert_allclose(evals_result["Train"]["quantile"][-1], loss) + + +def _expectile_loss( + y: np.ndarray, predt: np.ndarray, alpha: float, weight: Optional[np.ndarray] +) -> float: + diff = predt - y + weight_scale = np.where(diff >= 0.0, 1.0 - alpha, alpha) + loss = weight_scale * diff * diff + if weight is None: + return float(np.mean(loss)) + return float(np.sum(loss * weight) / np.sum(weight)) + + +def _expectile_loss_multi( + y: np.ndarray, predt: np.ndarray, alpha: np.ndarray, weight: Optional[np.ndarray] +) -> float: + diff = predt - y[:, None] + weight_scale = np.where(diff >= 0.0, 1.0 - alpha, alpha) + loss = weight_scale * diff * diff + if weight is None: + return float(np.mean(loss)) + return float(np.sum(loss * weight[:, None]) / (np.sum(weight) * alpha.size)) + + +def check_expectile_error(tree_method: str, device: Device) -> None: + """Test for the `expectile` loss.""" + from sklearn.datasets import make_regression + + rng = np.random.RandomState(23) + X, y = make_regression(128, 3, random_state=rng) + Xy = DMatrix(X, y) + evals_result: Dict[str, Dict] = {} + booster = train( + { + "tree_method": tree_method, + "eval_metric": "expectile", + "expectile_alpha": 0.3, + "device": device, + }, + Xy, + evals=[(Xy, "Train")], + evals_result=evals_result, + ) + predt = booster.inplace_predict(X) + loss = _expectile_loss(y, predt, 0.3, None) + np.testing.assert_allclose(evals_result["Train"]["expectile"][-1], loss) + + alpha = np.array([0.25, 0.5, 0.75]) + booster = train( + { + "tree_method": tree_method, + "eval_metric": "expectile", + "expectile_alpha": alpha, + "objective": "reg:expectileerror", + "device": device, + }, + Xy, + evals=[(Xy, "Train")], + evals_result=evals_result, + ) + predt = booster.inplace_predict(X) + loss = _expectile_loss_multi(y, predt, alpha, None) + np.testing.assert_allclose(evals_result["Train"]["expectile"][-1], loss) + + weights = rng.uniform(0.1, 1.0, size=y.shape[0]) + Xy_w = DMatrix(X, y, weight=weights) + evals_result_w: Dict[str, Dict] = {} + booster = train( + { + "tree_method": tree_method, + "eval_metric": "expectile", + "expectile_alpha": alpha, + "objective": "reg:expectileerror", + "device": device, + }, + Xy_w, + evals=[(Xy_w, "Train")], + evals_result=evals_result_w, + ) + predt = booster.inplace_predict(X) + loss = _expectile_loss_multi(y, predt, alpha, weights) + np.testing.assert_allclose(evals_result_w["Train"]["expectile"][-1], loss) + + +def run_roc_auc_binary(tree_method: str, n_samples: int, device: Device) -> None: + """TestROC AUC metric on a binary classification problem.""" + from sklearn.datasets import make_classification + from sklearn.metrics import roc_auc_score + + rng = np.random.RandomState(1994) + n_features = 10 + + X, y = make_classification( + n_samples, + n_features, + n_informative=n_features, + n_redundant=0, + random_state=rng, + ) + Xy = DMatrix(X, y) + booster = train( + { + "tree_method": tree_method, + "device": device, + "eval_metric": "auc", + "objective": "binary:logistic", + }, + Xy, + num_boost_round=1, + ) + score = booster.predict(Xy) + skl_auc = roc_auc_score(y, score) + auc = float(booster.eval(Xy).split(":")[1]) + np.testing.assert_allclose(skl_auc, auc, rtol=1e-6) + + X = rng.randn(*X.shape) + score = booster.predict(DMatrix(X)) + skl_auc = roc_auc_score(y, score) + auc = float(booster.eval(DMatrix(X, y)).split(":")[1]) + np.testing.assert_allclose(skl_auc, auc, rtol=1e-6) + + +def run_pr_auc_multi( + tree_method: str, + device: Device, + *, + multi_label: bool = False, + multi_strategy: str = "one_output_per_tree", +) -> None: + """Test PR AUC on a multi-class or multi-label classification problem.""" + from sklearn.datasets import make_classification, make_multilabel_classification + + if multi_label: + X, y = make_multilabel_classification( + 128, + 16, + n_classes=3, + n_labels=2, + allow_unlabeled=False, + random_state=1994, + ) + else: + X, y = make_classification( + 64, 16, n_informative=8, n_classes=3, random_state=1994 + ) + clf = XGBClassifier( + tree_method=tree_method, + n_estimators=1, + eval_metric="aucpr", + device=device, + multi_strategy=multi_strategy, + ) + clf.fit(X, y, eval_set=[(X, y)]) + evals_result = clf.evals_result()["validation_0"]["aucpr"][-1] + # No available implementation for comparison, just check that XGBoost converges + # to 1.0 + clf = XGBClassifier( + tree_method=tree_method, + n_estimators=10, + eval_metric="aucpr", + device=device, + multi_strategy=multi_strategy, + ) + clf.fit(X, y, eval_set=[(X, y)]) + evals_result = clf.evals_result()["validation_0"]["aucpr"][-1] + np.testing.assert_allclose(1.0, evals_result, rtol=1e-2) + + +# pylint: disable=too-many-locals, too-many-arguments +def run_roc_auc_multi( + tree_method: str, + n_samples: int, + weighted: bool, + device: Device, + *, + multi_label: bool = False, + multi_strategy: str = "one_output_per_tree", +) -> None: + """Test ROC AUC on a multi-class or multi-label classification problem.""" + from sklearn.datasets import make_classification, make_multilabel_classification + from sklearn.metrics import roc_auc_score + + rng = np.random.RandomState(1994) + n_features = 10 + n_classes = 4 + + if multi_label: + X, y = make_multilabel_classification( + n_samples, + n_features, + n_classes=n_classes, + n_labels=2, + allow_unlabeled=False, + random_state=rng, + ) + objective = "binary:logistic" + else: + X, y = make_classification( + n_samples, + n_features, + n_informative=n_features, + n_redundant=0, + n_classes=n_classes, + random_state=rng, + ) + objective = "multi:softprob" + if weighted: + weights = rng.randn(n_samples) + weights -= weights.min() + weights /= weights.max() + else: + weights = None + + Xy = DMatrix(X, y, weight=weights) + params: dict[str, Any] = { + "tree_method": tree_method, + "eval_metric": "auc", + "objective": objective, + "device": device, + "multi_strategy": multi_strategy, + } + if not multi_label: + params["num_class"] = n_classes + booster = train( + params, + Xy, + num_boost_round=1, + ) + score = booster.predict(Xy) + if multi_label: + skl_auc = roc_auc_score(y, score, average="macro", sample_weight=weights) + else: + skl_auc = roc_auc_score( + y, score, average="weighted", sample_weight=weights, multi_class="ovr" + ) + auc = float(booster.eval(Xy).split(":")[1]) + np.testing.assert_allclose(skl_auc, auc, rtol=1e-6) + + X = rng.randn(*X.shape) + + score = booster.predict(DMatrix(X, weight=weights)) + if multi_label: + skl_auc = roc_auc_score(y, score, average="macro", sample_weight=weights) + else: + skl_auc = roc_auc_score( + y, score, average="weighted", sample_weight=weights, multi_class="ovr" + ) + auc = float(booster.eval(DMatrix(X, y, weight=weights)).split(":")[1]) + np.testing.assert_allclose(skl_auc, auc, rtol=1e-5) + + +def run_pr_auc_ltr(tree_method: str, device: Device) -> None: + """Test for PR AUC metric on a ranking problem.""" + from sklearn.datasets import make_classification + + X, y = make_classification(128, 4, n_classes=2, random_state=1994) + ltr = XGBRanker( + tree_method=tree_method, + n_estimators=16, + objective="rank:pairwise", + eval_metric="aucpr", + device=device, + ) + groups = np.array([32, 32, 64]) + ltr.fit( + X, + y, + group=groups, + eval_set=[(X, y)], + eval_group=[groups], + ) + results = ltr.evals_result()["validation_0"]["aucpr"] + assert results[-1] >= 0.99 + + +def run_pr_auc_binary(tree_method: str, device: Device) -> None: + """Test for PR AUC metric on a binary classification problem.""" + from sklearn.datasets import make_classification + from sklearn.metrics import auc, precision_recall_curve + + X, y = make_classification(128, 4, n_classes=2, random_state=1994) + clf = XGBClassifier( + tree_method=tree_method, n_estimators=1, eval_metric="aucpr", device=device + ) + clf.fit(X, y, eval_set=[(X, y)]) + evals_result = clf.evals_result()["validation_0"]["aucpr"][-1] + + y_score = clf.predict_proba(X)[:, 1] # get the positive column + precision, recall, _ = precision_recall_curve(y, y_score) + prauc = auc(recall, precision) + # Interpolation results are slightly different from sklearn, but overall should + # be similar. + np.testing.assert_allclose(prauc, evals_result, rtol=1e-2) + + clf = XGBClassifier( + tree_method=tree_method, n_estimators=10, eval_metric="aucpr", device=device + ) + clf.fit(X, y, eval_set=[(X, y)]) + evals_result = clf.evals_result()["validation_0"]["aucpr"][-1] + np.testing.assert_allclose(0.99, evals_result, rtol=1e-2) diff --git a/v12x/pylib/xgboost/testing/monotone_constraints.py b/v12x/pylib/xgboost/testing/monotone_constraints.py new file mode 100644 index 0000000000000000000000000000000000000000..2c126b9ace0360be3c44a18f2b197503d0cd05fc --- /dev/null +++ b/v12x/pylib/xgboost/testing/monotone_constraints.py @@ -0,0 +1,212 @@ +"""Helpers for testing monotone constraints.""" + +import json +from typing import Optional, Sequence + +import numpy as np +import pytest + +from .._typing import FeatureNames +from ..core import Booster, DMatrix +from ..training import train +from .utils import Device + + +def is_increasing(v: np.ndarray) -> bool: + """Whether ``v`` is nondecreasing along the sweep axis.""" + return np.count_nonzero(np.diff(v, axis=0) < 0.0) == 0 + + +def is_decreasing(v: np.ndarray) -> bool: + """Whether ``v`` is nonincreasing along the sweep axis.""" + return np.count_nonzero(np.diff(v, axis=0) > 0.0) == 0 + + +def is_correctly_constrained( + learner: Booster, feature_names: Optional[FeatureNames] = None +) -> bool: + """Whether the monotone constraint is correctly applied.""" + n = 100 + variable_x = np.linspace(0, 1, n).reshape((n, 1)) + fixed_xs_values = np.linspace(0, 1, n) + + for i in range(n): + fixed_x = fixed_xs_values[i] * np.ones((n, 1)) + monotonically_increasing_x = np.column_stack((variable_x, fixed_x)) + monotonically_increasing_dset = DMatrix( + monotonically_increasing_x, feature_names=feature_names + ) + monotonically_increasing_y = learner.predict(monotonically_increasing_dset) + + monotonically_decreasing_x = np.column_stack((fixed_x, variable_x)) + monotonically_decreasing_dset = DMatrix( + monotonically_decreasing_x, feature_names=feature_names + ) + monotonically_decreasing_y = learner.predict(monotonically_decreasing_dset) + + if not ( + is_increasing(monotonically_increasing_y) + and is_decreasing(monotonically_decreasing_y) + ): + return False + + return True + + +NUMBER_OF_DPOINTS = 1000 +x1_positively_correlated_with_y = np.random.random(size=NUMBER_OF_DPOINTS) +x2_negatively_correlated_with_y = np.random.random(size=NUMBER_OF_DPOINTS) + +x = np.column_stack((x1_positively_correlated_with_y, x2_negatively_correlated_with_y)) +zs = np.random.normal(loc=0.0, scale=0.01, size=NUMBER_OF_DPOINTS) +y = ( + 5 * x1_positively_correlated_with_y + + np.sin(10 * np.pi * x1_positively_correlated_with_y) + - 5 * x2_negatively_correlated_with_y + - np.cos(10 * np.pi * x2_negatively_correlated_with_y) + + zs +) +training_dset = DMatrix(x, label=y) + +# Multi-output labels sharing the same ``(f0: increasing, f1: decreasing)`` structure as +# ``y``. Each column is a positive-scale affine transform of ``y``, so every target must +# be increasing in ``f0`` and decreasing in ``f1`` while still differing enough to +# exercise shared vector-leaf splits. +_mt_noise = np.random.default_rng(2026).normal( + loc=0.0, scale=0.01, size=(NUMBER_OF_DPOINTS, 2) +) +y_mt = np.column_stack( + [y, 2.0 * y + _mt_noise[:, 0], 0.5 * y + _mt_noise[:, 1]] +).astype(np.float32) + + +def run_parent_gain(device: Device, multi_strategy: str) -> None: + """Test that parent gain uses the node's inherited monotonic bounds.""" + X = np.array( + [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [1.0, 1.0, 0.0], [1.0, 1.0, 1.0]], + dtype=np.float32, + ) + grad = np.array([3.0, -2.75, -0.125, 0.875], dtype=np.float32) + hess = np.array([1.0, 0.5, 0.25, 0.25], dtype=np.float32) + expected_gain = 0.25 + if multi_strategy == "multi_output_tree": + gradients = np.stack([grad, 2.0 * grad], axis=1) + hessians = np.stack([hess, hess], axis=1) + expected_gain += 1.0 + else: + gradients, hessians = grad, hess + dtrain = DMatrix(X, label=np.zeros_like(gradients)) + + def objective( + _predt: np.ndarray, _dtrain: DMatrix + ) -> tuple[np.ndarray, np.ndarray]: + return gradients, hessians + + params = { + "tree_method": "hist", + "monotone_constraints": "(1, 0, 0)", + "reg_alpha": 0, + "reg_lambda": 0, + "max_delta_step": 0, + "min_child_weight": 0, + "max_depth": 3, + "eta": 1, + "base_score": 0, + "device": device, + "multi_strategy": multi_strategy, + } + model = train(params, dtrain, num_boost_round=1, obj=objective) + trees = json.loads(model.get_dump(dump_format="json", with_stats=True)[0]) + + constrained_parent = trees["children"][1]["children"][1] + assert constrained_parent["split"] == "f2" + assert constrained_parent["gain"] == pytest.approx(expected_gain) + + +def run_monotone_constraints( + device: Device, + tree_method: str, + grow_policy: str, + multi_strategy: str = "one_output_per_tree", +) -> None: + """Check for a positive (``f0``) and negative (``f1``) constraint.""" + label = y_mt if multi_strategy == "multi_output_tree" else y + dtrain = DMatrix(x, label=label) + params = { + "tree_method": tree_method, + "grow_policy": grow_policy, + "device": device, + "multi_strategy": multi_strategy, + "monotone_constraints": "(1, -1)", + } + model = train(params, dtrain, num_boost_round=16) + assert is_correctly_constrained(model) + + +def _assert_monotone( + booster: Booster, + constraints: Sequence[int], + *, + n_grid: int = 64, + n_ref: int = 32, + seed: int = 0, +) -> None: + """Grid-check monotonicity per output column for every constrained feature. + + For each feature with a nonzero constraint, sweep it over ``[0, 1]`` while holding + the remaining features fixed at random reference rows. + + """ + rng = np.random.RandomState(seed) + n_features = len(constraints) + grid = np.linspace(0.0, 1.0, n_grid, dtype=np.float32) + for fidx, direction in enumerate(constraints): + if direction == 0: + continue + for _ in range(n_ref): + base = rng.rand(n_features).astype(np.float32) + features = np.tile(base, (n_grid, 1)) + features[:, fidx] = grid + pred = booster.predict(DMatrix(features)) + if direction > 0: + assert is_increasing(pred), (fidx, "increasing") + else: + assert is_decreasing(pred), (fidx, "decreasing") + + +def run_multi_output_monotone( + device: Device, grow_policy: str, multi_strategy: str +) -> None: + """Monotonicity check for deep trees with mixed feature constraints. + + Uses more features than constraints so that constrained splits are followed by + splits on unconstrained features. + + """ + constraints = (1, -1, 0, 0, 0) + rng = np.random.RandomState(2026) + n_samples, n_features, n_targets = 2048, len(constraints), 3 + features = rng.rand(n_samples, n_features).astype(np.float32) + labels = np.empty((n_samples, n_targets), dtype=np.float32) + for t in range(n_targets): + labels[:, t] = ( + (t + 1) * features[:, 0] + - (t + 1) * features[:, 1] + + np.sin(6.0 * features[:, 2]) + + 0.5 * features[:, 3] * features[:, 4] + + rng.normal(scale=0.05, size=n_samples) + ) + params = { + "tree_method": "hist", + "device": device, + "grow_policy": grow_policy, + "multi_strategy": multi_strategy, + "monotone_constraints": constraints, + "max_depth": 6, + "eta": 0.3, + "reg_lambda": 1.0, + "reg_alpha": 0.1, + "min_child_weight": 0.0, + } + booster = train(params, DMatrix(features, label=labels), num_boost_round=40) + _assert_monotone(booster, constraints, seed=2026 + 1) diff --git a/v12x/pylib/xgboost/testing/multi_target.py b/v12x/pylib/xgboost/testing/multi_target.py new file mode 100644 index 0000000000000000000000000000000000000000..d5a77ed6dbe27a41230d2fb1b6878362033fd3c5 --- /dev/null +++ b/v12x/pylib/xgboost/testing/multi_target.py @@ -0,0 +1,981 @@ +"""Tests for multi-target training.""" + +# pylint: disable=unbalanced-tuple-unpacking +import json +from types import ModuleType +from typing import Any, Callable, Dict, List, Optional, Tuple + +import numpy as np +import pytest +from sklearn.datasets import ( + make_classification, + make_multilabel_classification, + make_regression, +) +from sklearn.metrics.pairwise import cosine_similarity + +import xgboost.testing as tm + +from .._typing import ArrayLike +from ..compat import import_cupy +from ..core import Booster, DMatrix, ExtMemQuantileDMatrix, QuantileDMatrix, build_info +from ..objective import Objective, TreeObjective +from ..sklearn import XGBClassifier, XGBRegressor +from ..training import train +from .data import IteratorForTest +from .updater import ResetStrategy, train_result +from .utils import Device, assert_allclose, non_increasing + + +def check_categorical_mixed(device: Device) -> None: # pylint: disable=too-many-locals + """Test mixed numerical, one-hot, and partition-based splits.""" + pd = pytest.importorskip("pandas") + + features = np.stack( + [ + np.repeat(np.arange(2), 16), + np.tile(np.repeat(np.arange(2), 8), 2), + np.tile(np.arange(8), 4), + ], + axis=1, + ) + X = pd.DataFrame( + { + "numeric": features[:, 0].astype(np.float32), + "onehot": pd.Categorical(features[:, 1], categories=range(2)), + "partition": pd.Categorical(features[:, 2], categories=range(8)), + } + ) + y = np.column_stack( + [ + 8.0 * features[:, 0], + 6.0 * features[:, 1], + 10.0 * (features[:, 2] >= 4), + ] + ).astype(np.float32) + Xy = DMatrix(X, y, enable_categorical=True) + + booster = train( + { + "device": device, + "tree_method": "hist", + "multi_strategy": "multi_output_tree", + "max_cat_to_onehot": 4, + "max_depth": 3, + "learning_rate": 1.0, + "reg_lambda": 0, + "min_child_weight": 0, + "base_score": 0, + }, + Xy, + num_boost_round=4, + ) + np.testing.assert_allclose(booster.predict(Xy), y) + + model = json.loads(booster.save_raw(raw_format="json")) + tree = model["learner"]["gradient_booster"]["model"]["trees"][0] + split_indices = np.asarray(tree["split_indices"]) + split_types = np.asarray(tree["split_type"]) + is_internal = np.asarray(tree["left_children"]) != -1 + is_categorical = split_types.astype(bool) + + numerical_features = split_indices[is_internal & ~is_categorical] + categorical_features = split_indices[is_internal & is_categorical] + category_sizes = np.asarray(tree["categories_sizes"]) + + assert set(numerical_features) == {0} + assert set(categorical_features) == {1, 2} + np.testing.assert_equal(category_sizes[categorical_features == 1], 1) + assert np.any(category_sizes[categorical_features == 2] > 1) + + +def run_multiclass(device: Device, learning_rate: Optional[float]) -> None: + """Use vector leaf for multi-class models.""" + X, y = make_classification( + 128, n_features=12, n_informative=10, n_classes=4, random_state=2025 + ) + clf = XGBClassifier( + debug_synchronize=True, + multi_strategy="multi_output_tree", + min_child_weight=0.0, + callbacks=[ResetStrategy()], + n_estimators=10, + device=device, + learning_rate=learning_rate, + ) + clf.fit(X, y, eval_set=[(X, y)]) + assert clf.objective == "multi:softprob" + assert non_increasing(clf.evals_result()["validation_0"]["mlogloss"]) + if learning_rate is not None and abs(learning_rate - 1.0) < 1e-5: + assert clf.evals_result()["validation_0"]["mlogloss"][-1] < 0.045 + + proba = clf.predict_proba(X) + assert proba.shape == (y.shape[0], 4) + + +def run_multilabel(device: Device, learning_rate: Optional[float]) -> None: + """Use vector leaf for multi-label classification models.""" + X, y = make_multilabel_classification(128, random_state=2025) + clf = XGBClassifier( + debug_synchronize=True, + multi_strategy="multi_output_tree", + min_child_weight=0.0, + callbacks=[ResetStrategy()], + n_estimators=10, + device=device, + learning_rate=learning_rate, + ) + clf.fit(X, y, eval_set=[(X, y)]) + assert clf.objective == "binary:logistic" + assert non_increasing(clf.evals_result()["validation_0"]["logloss"]) + if learning_rate is not None and abs(learning_rate - 1.0) < 1e-5: + assert clf.evals_result()["validation_0"]["logloss"][-1] < 0.065 + + proba = clf.predict_proba(X) + assert proba.shape == y.shape + + +def run_quantile_loss(device: Device, weighted: bool) -> None: + """Check quantile regression for vector leaf.""" + params = { + "objective": "reg:quantileerror", + "device": device, + "quantile_alpha": [0.45, 0.5, 0.55], + "multi_strategy": "multi_output_tree", + } + n_samples = 2048 + X, y = make_regression(n_samples=n_samples, n_features=16, random_state=2026) + + def no_crossing_first_tree(weight: Optional[np.ndarray]) -> None: + """The first tree should not generate quantile crossing given sufficient amount + of samples for quantile interpolation. + + """ + Xy = QuantileDMatrix(X, y, weight=weight) + booster = train(params, Xy, evals=[(Xy, "Train")], num_boost_round=1) + y_predt = booster.predict(Xy) + assert y_predt.shape == (n_samples, 3) + assert (y_predt[:, 0] <= y_predt[:, 1]).all() + assert (y_predt[:, 1] <= y_predt[:, 2]).all() + + if not weighted: + weight = None + else: + # Test with weights. + rng = np.random.default_rng(2026) + weight = rng.uniform(0.0, 1.0, size=n_samples) + + no_crossing_first_tree(weight) + + Xy = QuantileDMatrix(X, y, weight=weight) + evals_result = train_result(params, Xy, num_rounds=10) + assert non_increasing(evals_result["train"]["quantile"]) + + +def run_expectile_loss(device: Device) -> None: + """Check that reparameterized expectile regression produces non-crossing curves.""" + params = { + "objective": "reg:expectileerror", + "device": device, + "expectile_alpha": [0.45, 0.5, 0.55], + "multi_strategy": "multi_output_tree", + "tree_method": "hist", + } + n_samples = 2048 + rng = np.random.default_rng(2027) + X = rng.normal(size=(n_samples, 16)) + y = X[:, 0] - 0.5 * X[:, 1] + rng.normal(size=n_samples) + + Xy = QuantileDMatrix(X, y) + evals_result: Dict[str, Dict] = {} + booster = train( + params, + Xy, + evals=[(Xy, "Train")], + evals_result=evals_result, + num_boost_round=10, + ) + + y_predt = booster.predict(Xy) + assert y_predt.shape == (n_samples, 3) + assert (y_predt[:, 0] <= y_predt[:, 1]).all() + assert (y_predt[:, 1] <= y_predt[:, 2]).all() + assert non_increasing(evals_result["Train"]["expectile"]) + + +def run_absolute_error(device: Device) -> None: + """Test mean absolute error with vector leaf.""" + params = { + "objective": "reg:absoluteerror", + "device": device, + "multi_strategy": "multi_output_tree", + } + n_samples = 1024 + X, y = make_regression( + n_samples=n_samples, n_features=16, n_targets=3, random_state=2026 + ) + Xy = QuantileDMatrix(X, y) + evals_result: Dict[str, Dict] = {} + train( + params, + Xy, + evals=[(Xy, "Train")], + verbose_eval=False, + evals_result=evals_result, + num_boost_round=16, + ) + assert non_increasing(evals_result["Train"]["mae"]) + assert evals_result["Train"]["mae"][-1] < 30.0 + + +def _array_impl(device: Device) -> ModuleType: + if device == "cuda": + nda = import_cupy() + else: + nda = np + return nda + + +class LsObj0(TreeObjective): + """Split grad is the same as value grad.""" + + def __init__(self, device: Device) -> None: + self.device = device + + def __call__( + self, iteration: int, y_pred: ArrayLike, dtrain: DMatrix + ) -> Tuple[ArrayLike, ArrayLike]: + nda = _array_impl(self.device) + + y_true = dtrain.get_label() + grad, hess = tm.ls_obj(y_true, y_pred, None) + return nda.array(grad), nda.array(hess) + + def split_grad( + self, iteration: int, grad: ArrayLike, hess: ArrayLike + ) -> Tuple[ArrayLike, ArrayLike]: + nda = _array_impl(self.device) + return nda.array(grad), nda.array(hess) + + +class LsObj1(Objective): + """No split grad.""" + + def __init__(self, device: Device) -> None: + self.device = device + + def __call__( + self, iteration: int, y_pred: ArrayLike, dtrain: DMatrix + ) -> Tuple[ArrayLike, ArrayLike]: + nda = _array_impl(self.device) + + y_true = dtrain.get_label() + grad, hess = tm.ls_obj(y_true, y_pred, None) + return nda.array(grad), nda.array(hess) + + +# Use mean gradient, should still converge. +class LsObj2(LsObj0): + """Use mean as split grad.""" + + def __init__(self, device: Device, check_used: bool): + self._chk = check_used + super().__init__(device=device) + + def split_grad( + self, iteration: int, grad: ArrayLike, hess: ArrayLike + ) -> Tuple[np.ndarray, np.ndarray]: + nda = _array_impl(self.device) + + if self._chk: + assert False + sgrad = nda.mean(grad, axis=1) + shess = nda.mean(hess, axis=1) + return sgrad, shess + + +def run_reduced_grad(device: Device) -> None: # pylint: disable=too-many-locals + """Basic test for using reduced gradient for tree splits.""" + X, y = make_regression( + n_samples=1024, n_features=16, random_state=1994, n_targets=5 + ) + Xy = QuantileDMatrix(X, y) + + def run_test( + obj: Optional[Objective], base_score: Optional[list[float]] = None + ) -> Booster: + evals_result: Dict[str, Dict] = {} + booster = train( + { + "debug_synchronize": True, + "device": device, + "multi_strategy": "multi_output_tree", + "learning_rate": 1, + "base_score": base_score, + }, + Xy, + evals=[(Xy, "Train")], + obj=obj, + num_boost_round=8, + evals_result=evals_result, + ) + assert non_increasing(evals_result["Train"]["rmse"]) + return booster + + booster_0 = run_test(LsObj0(device)) + booster_1 = run_test(LsObj1(device)) + np.testing.assert_allclose( + booster_0.inplace_predict(X), booster_1.inplace_predict(X) + ) + + booster_2 = run_test(LsObj0(device), [0.5] * y.shape[1]) + booster_3 = run_test(None, [0.5] * y.shape[1]) + np.testing.assert_allclose( + booster_2.inplace_predict(X), booster_3.inplace_predict(X) + ) + + run_test(LsObj2(device, False)) + with pytest.raises(AssertionError): + run_test(LsObj2(device, True)) + + def run_reg(obj: Objective) -> None: + reg = XGBRegressor( + n_estimators=2, + device=device, + tree_method="hist", + multi_strategy="multi_output_tree", + objective=obj, + ) + reg.fit(X, y) + predt = reg.predict(X) + assert predt.shape == y.shape + assert np.isfinite(predt).all() + + run_reg(LsObj1(device)) + run_reg(LsObj2(device, False)) + with pytest.raises(AssertionError): + run_reg(LsObj2(device, True)) + + y_ind = (y > np.median(y, axis=0)).astype(np.int32) + clf = XGBClassifier( + n_estimators=2, + device=device, + tree_method="hist", + multi_strategy="multi_output_tree", + objective=LsObj2(device, False), + ) + clf.fit(X, y_ind) + predt = clf.predict(X) + proba = clf.predict_proba(X) + assert predt.shape == y_ind.shape + assert proba.shape == y_ind.shape + np.testing.assert_array_equal(predt, (proba > 0.5).astype(np.int32)) + + with pytest.raises( + ValueError, + match="Monotonic constraints are not supported with reduced gradients", + ): + train( + { + "device": device, + "tree_method": "hist", + "multi_strategy": "multi_output_tree", + "monotone_constraints": (1,), + }, + Xy, + num_boost_round=1, + obj=LsObj2(device, False), + ) + + +def run_with_iter(device: Device) -> None: # pylint: disable=too-many-locals + """Test vector leaf with external memory.""" + nda = _array_impl(device) + + n_batches = 4 + n_rounds = 8 + n_targets = 3 + intercept = [0.5] * n_targets + + params = { + "device": device, + "multi_strategy": "multi_output_tree", + "learning_rate": 1.0, + "base_score": intercept, + "debug_synchronize": True, + } + + Xs = [] + ys = [] + for i in range(n_batches): + X_i, y_i = make_regression( + n_samples=4096, n_features=8, random_state=(i + 1), n_targets=n_targets + ) + Xs.append(nda.asarray(X_i)) + ys.append(nda.asarray(y_i)) + it = IteratorForTest(Xs, ys, None, cache="cache", on_host=True) + Xy: DMatrix = ExtMemQuantileDMatrix( + it, cache_host_ratio=1.0 if device == "cuda" else None + ) + + evals_result_0: Dict[str, Dict] = {} + booster_0 = train( + params, + Xy, + num_boost_round=n_rounds, + evals=[(Xy, "Train")], + evals_result=evals_result_0, + ) + + it = IteratorForTest(Xs, ys, None, cache=None) + Xy = QuantileDMatrix(it) + evals_result_1: Dict[str, Dict] = {} + booster_1 = train( + params, + Xy, + num_boost_round=n_rounds, + evals=[(Xy, "Train")], + evals_result=evals_result_1, + ) + np.testing.assert_allclose( + evals_result_0["Train"]["rmse"], evals_result_1["Train"]["rmse"] + ) + assert non_increasing(evals_result_0["Train"]["rmse"]) + X, _, _ = it.as_arrays() + assert_allclose(device, booster_0.inplace_predict(X), booster_1.inplace_predict(X)) + + binfo = build_info() + tv = "THRUST_VERSION" + if device == "cuda" and tv in binfo and binfo[tv][0] < 3: + pytest.xfail("CCCL version too old.") + + it = IteratorForTest( + Xs, + ys, + None, + cache="cache", + on_host=True, + min_cache_page_bytes=X.shape[0] // n_batches * X.shape[1], + ) + Xy = ExtMemQuantileDMatrix(it, cache_host_ratio=1.0 if device == "cuda" else None) + + evals_result_2: Dict[str, Dict] = {} + booster_2 = train( + params, + Xy, + evals=[(Xy, "Train")], + obj=LsObj0(device), + num_boost_round=n_rounds, + evals_result=evals_result_2, + ) + np.testing.assert_allclose( + evals_result_0["Train"]["rmse"], evals_result_2["Train"]["rmse"] + ) + assert_allclose(device, booster_0.inplace_predict(X), booster_2.inplace_predict(X)) + + # Distinct target scales make a wrong-row mask diverge loudly. + ys_adaptive = [y_i * nda.asarray([1.0, 20.0, 100.0]) for y_i in ys] + adaptive_params = { + "device": device, + "multi_strategy": "multi_output_tree", + "objective": "reg:absoluteerror", + "subsample": 0.5, + "seed": 2026, + "debug_synchronize": True, + } + it = IteratorForTest( + Xs, + ys_adaptive, + None, + cache="cache", + on_host=True, + min_cache_page_bytes=0, # disables page concatenation + ) + Xy = ExtMemQuantileDMatrix(it, cache_host_ratio=1.0 if device == "cuda" else None) + booster_ext = train(adaptive_params, Xy, num_boost_round=n_rounds) + + it = IteratorForTest(Xs, ys_adaptive, None, cache=None) + Xy = QuantileDMatrix(it) + booster_in = train(adaptive_params, Xy, num_boost_round=n_rounds) + + assert_allclose( + device, booster_ext.inplace_predict(X), booster_in.inplace_predict(X) + ) + + +def run_eta(device: Device) -> None: + """Test for learning rate.""" + X, y = make_regression(512, 16, random_state=2025, n_targets=3) + + def run(obj: Optional[Objective]) -> None: + params = { + "device": device, + "multi_strategy": "multi_output_tree", + "learning_rate": 1.0, + "debug_synchronize": True, + "base_score": 0.0, + } + Xy = QuantileDMatrix(X, y) + booster_0 = train(params, Xy, num_boost_round=1, obj=obj) + params["learning_rate"] = 0.1 + booster_1 = train(params, Xy, num_boost_round=1, obj=obj) + params["learning_rate"] = 2.0 + booster_2 = train(params, Xy, num_boost_round=1, obj=obj) + + predt_0 = booster_0.predict(Xy) + predt_1 = booster_1.predict(Xy) + predt_2 = booster_2.predict(Xy) + + np.testing.assert_allclose(predt_0, predt_1 * 10, rtol=1e-6) + np.testing.assert_allclose(predt_0 * 2, predt_2, rtol=1e-6) + + run(None) + run(LsObj0(device)) + + +def run_deterministic(device: Device) -> None: + """Check the vector leaf implementation is deterministic.""" + X, y = make_regression( + n_samples=int(2**16), n_features=64, random_state=1994, n_targets=5 + ) + + def run() -> Booster: + Xy = QuantileDMatrix(X, y) + params = { + "device": device, + "multi_strategy": "multi_output_tree", + "debug_synchronize": True, + } + return train(params, Xy, num_boost_round=16) + + booster_0 = run() + booster_1 = run() + raw_0 = booster_0.save_raw() + raw_1 = booster_1.save_raw() + assert raw_0 == raw_1 + + +def run_column_sampling(device: Device) -> None: + """Test column sampling with feature importance for multi-target trees.""" + n_features = 32 + X, y = make_regression( + n_samples=1024, n_features=n_features, random_state=1994, n_targets=3 + ) + # First half of features have weight, second half has 0 weight (not sampled). + feature_weights = np.zeros(shape=(n_features, 1), dtype=np.float32) + feature_weights[: n_features // 2] = 1.0 / (n_features / 2) + Xy = QuantileDMatrix(X, y, feature_weights=feature_weights) + + params = { + "device": device, + "multi_strategy": "multi_output_tree", + "debug_synchronize": True, + "colsample_bynode": 0.4, + } + booster = train(params, Xy, num_boost_round=16) + + # Test all importance types + for importance_type in ["weight", "gain", "total_gain", "cover", "total_cover"]: + scores: dict = booster.get_score(importance_type=importance_type) + assert len(scores) > 0, f"No scores for {importance_type}" + + # Sampled features (first half) should be in scores + for f in range(0, n_features // 2): + assert f"f{f}" in scores, f"f{f} not in {importance_type} scores" + + # Non-sampled features (second half) should NOT be in scores + for f in range(n_features // 2, n_features): + assert f"f{f}" not in scores + + for score in scores.values(): + assert isinstance(score, float) + assert score >= 0 + + # sklearn Coef + X, y = make_multilabel_classification(random_state=1994) + clf = XGBClassifier( + multi_strategy="multi_output_tree", + importance_type="weight", + device=device, + colsample_bynode=0.2, + min_child_weight=0.0, + feature_weights=np.arange(0, X.shape[1]), + ) + clf.fit(X, y) + fi = clf.feature_importances_ + assert fi[0] == 0.0 + assert fi[-1] > fi[1] * 5 + + w = np.polynomial.Polynomial.fit(np.arange(0, X.shape[1]), fi, deg=1) + assert w.coef[1] > 0.03 + + +def run_grow_policy(device: Device, grow_policy: str) -> None: + """Test grow policy (depthwise and lossguide) for vector leaf.""" + X, y = make_regression( + n_samples=1024, n_features=16, random_state=1994, n_targets=3 + ) + Xy = QuantileDMatrix(X, y) + + params = { + "device": device, + "multi_strategy": "multi_output_tree", + "debug_synchronize": True, + "grow_policy": grow_policy, + } + + evals_result = train_result(params, Xy, num_rounds=10) + assert non_increasing(evals_result["train"]["rmse"]) + + +def run_mixed_strategy(device: Device, use_dart: bool) -> None: + """Test mixed multi_strategy with ResetStrategy callback.""" + X, y = make_classification( + n_samples=1024, n_informative=8, n_classes=3, random_state=1994 + ) + Xy = DMatrix(data=X, label=y) + + params: Dict[str, Any] = { + "num_parallel_tree": 4, + "num_class": 3, + "objective": "multi:softprob", + "multi_strategy": "multi_output_tree", + "device": device, + "debug_synchronize": True, + "base_score": 0, + } + if use_dart: + params.update({"rate_drop": 0.5, "one_drop": True}) + booster = train( + params, + num_boost_round=16, + dtrain=Xy, + callbacks=[ResetStrategy()], + ) + + # Test model slicing - each boosting round should be iterable + assert len(list(booster)) == 16 + + # Test that sliced predictions sum to full prediction + predt = booster.predict(Xy, output_margin=True) + predt_sum = np.zeros(predt.shape) + for t in booster: + predt_sum += t.predict(Xy, output_margin=True) + np.testing.assert_allclose(predt, predt_sum, atol=1e-5) + + # Test feature importance works with mixed trees + for importance_type in ["weight", "gain", "total_gain", "cover", "total_cover"]: + scores = booster.get_score(importance_type=importance_type) + assert len(scores) > 0 + for score in scores.values(): + assert isinstance(score, float) + assert score >= 0 + + +def run_feature_importance_strategy_compare(device: Device) -> None: + """Different strategies produce similar feature importance ratios.""" + n_features = 16 + X, y = make_classification( + n_samples=2048, + n_features=n_features, + n_informative=10, + n_classes=4, + random_state=1994, + ) + Xy = DMatrix(data=X, label=y) + + base_params: Dict[str, Any] = { + "num_class": 4, + "objective": "multi:softprob", + "device": device, + "debug_synchronize": True, + "max_depth": 5, + } + + # Train models with different strategies + boosters = [ + train( + {**base_params, "multi_strategy": "multi_output_tree"}, + Xy, + num_boost_round=32, + ), + train( + {**base_params, "multi_strategy": "one_output_per_tree"}, + Xy, + num_boost_round=32, + ), + train( + {**base_params, "multi_strategy": "multi_output_tree"}, + Xy, + num_boost_round=32, + callbacks=[ResetStrategy()], + ), + ] + + def get_normalized_importance(booster: Booster, importance_type: str) -> np.ndarray: + """Get feature importance as normalized array (sums to 1).""" + scores = booster.get_score(importance_type=importance_type) + arr = np.array([scores.get(f"f{i}", 0.0) for i in range(n_features)]) + return arr / arr.sum() if arr.sum() > 0 else arr + + for importance_type in ["weight", "gain", "total_gain", "cover", "total_cover"]: + imps = [get_normalized_importance(b, importance_type) for b in boosters] + + # Check that importances are not exactly the same (different strategies) + assert not np.allclose(imps[0], imps[1]) + assert not np.allclose(imps[0], imps[2]) + + # Check that normalized importances are similar (correlated) + # All strategies should have reasonably similar importance patterns + assert cosine_similarity([imps[0]], [imps[1]])[0, 0] > 0.9 + assert cosine_similarity([imps[0]], [imps[2]])[0, 0] > 0.9 + assert cosine_similarity([imps[1]], [imps[2]])[0, 0] > 0.9 + + +# pylint: disable=too-many-arguments, too-many-locals +def _run_regression_objective_test( + device: Device, + objective: str, + metric: str, + X: np.ndarray, + y: np.ndarray, + *, + extra_params: Optional[Dict[str, Any]] = None, + check_pred_positive: bool = False, + check_pred_probability: bool = False, + check_pred_binary: bool = False, + strictly_non_increasing: bool = True, +) -> None: + params: Dict[str, Any] = { + "objective": objective, + "device": device, + "multi_strategy": "multi_output_tree", + } + if extra_params: + params.update(extra_params) + + n_samples = X.shape[0] + n_targets = y.shape[1] if y.ndim > 1 else 1 + + Xy = DMatrix(X, y) + evals_result: Dict[str, Dict] = {} + booster = train( + params, + Xy, + evals=[(Xy, "Train")], + verbose_eval=False, + evals_result=evals_result, + num_boost_round=16, + ) + predt = booster.predict(Xy) + assert predt.shape == (n_samples, n_targets) + + if check_pred_positive: + assert (predt > 0).all() + if check_pred_probability: + assert (predt > 0).all() and (predt < 1).all() + if check_pred_binary: + assert set(np.unique(predt)).issubset({0.0, 1.0}) + + metric_vals = evals_result["Train"][metric] + if strictly_non_increasing: + assert non_increasing(metric_vals) + else: + assert metric_vals[-1] < metric_vals[0] + + +def run_reg_squarederror(device: Device) -> None: + """Test squared error regression with vector leaf.""" + n_samples, n_targets = 1024, 3 + X, y = make_regression( + n_samples=n_samples, n_features=16, n_targets=n_targets, random_state=2026 + ) + _run_regression_objective_test(device, "reg:squarederror", "rmse", X, y) + + +def run_reg_logistic(device: Device) -> None: + """Test logistic regression for probability with vector leaf.""" + n_samples, n_targets = 1024, 3 + rng = np.random.default_rng(2026) + X = rng.standard_normal((n_samples, 16)) + y = rng.uniform(0.0, 1.0, (n_samples, n_targets)) # Labels in [0, 1] + _run_regression_objective_test( + device, "reg:logistic", "rmse", X, y, check_pred_probability=True + ) + + +def run_reg_gamma(device: Device) -> None: + """Test gamma regression with vector leaf.""" + n_samples, n_targets = 1024, 3 + rng = np.random.default_rng(2026) + X = rng.standard_normal((n_samples, 16)) + y = rng.gamma(2.0, 2.0, (n_samples, n_targets)) # Labels must be positive + _run_regression_objective_test( + device, "reg:gamma", "gamma-deviance", X, y, check_pred_positive=True + ) + + +def run_reg_squaredlogerror(device: Device) -> None: + """Test squared log error regression with vector leaf.""" + n_samples, n_targets = 1024, 3 + rng = np.random.default_rng(2026) + X = rng.standard_normal((n_samples, 16)) + y = np.abs(rng.standard_normal((n_samples, n_targets))) + 0.1 # Labels > -1 + _run_regression_objective_test(device, "reg:squaredlogerror", "rmsle", X, y) + + +def run_reg_pseudohubererror(device: Device) -> None: + """Test pseudo huber error regression with vector leaf.""" + n_samples, n_targets = 1024, 3 + X, y = make_regression( + n_samples=n_samples, n_features=16, n_targets=n_targets, random_state=2026 + ) + _run_regression_objective_test( + device, "reg:pseudohubererror", "mphe", X, y, extra_params={"huber_slope": 1.0} + ) + + +def run_binary_logitraw(device: Device) -> None: + """Test binary logitraw with vector leaf (multi-label classification).""" + n_samples = 1024 + X, y = make_multilabel_classification(n_samples, random_state=2026) + _run_regression_objective_test( + device, "binary:logitraw", "logloss", X, y, strictly_non_increasing=False + ) + + +def run_binary_hinge(device: Device) -> None: + """Test binary hinge loss with vector leaf (multi-label classification).""" + n_samples = 1024 + X, y = make_multilabel_classification(n_samples, random_state=2026) + _run_regression_objective_test( + device, + "binary:hinge", + "error", + X, + y, + check_pred_binary=True, + strictly_non_increasing=False, + ) + + +def run_count_poisson(device: Device) -> None: + """Test Poisson regression with vector leaf.""" + n_samples, n_targets = 1024, 3 + rng = np.random.default_rng(2026) + X = rng.standard_normal((n_samples, 16)) + y = rng.poisson(5, (n_samples, n_targets)).astype(np.float32) # Labels >= 0 + _run_regression_objective_test( + device, "count:poisson", "poisson-nloglik", X, y, check_pred_positive=True + ) + + +def run_reg_tweedie(device: Device) -> None: + """Test Tweedie regression with vector leaf.""" + n_samples, n_targets = 1024, 3 + rng = np.random.default_rng(2026) + X = rng.standard_normal((n_samples, 16)) + y = rng.gamma(2.0, 2.0, (n_samples, n_targets)) # Labels >= 0 + _run_regression_objective_test( + device, "reg:tweedie", "tweedie-nloglik@1.5", X, y, check_pred_positive=True + ) + + +def all_reg_objectives() -> List[Callable[[Device], None]]: + """List of obj tests.""" + objs: List[Callable[[Device], None]] = [ + run_reg_squarederror, + run_reg_logistic, + run_reg_gamma, + run_reg_squaredlogerror, + run_reg_pseudohubererror, + run_binary_logitraw, + run_binary_hinge, + run_count_poisson, + run_reg_tweedie, + ] + return objs + + +def _make_subsample_params(device: Device, sampling_method: str) -> dict: + params = { + "device": device, + "tree_method": "hist", + "multi_strategy": "multi_output_tree", + "subsample": 0.5, + "sampling_method": sampling_method, + "max_depth": 6, + "debug_synchronize": True, + "seed": 2026, + } + return params + + +def run_subsample(device: Device, sampling_method: str) -> None: + """Test row subsampling.""" + n_samples = 2048 + X, y = make_regression( + n_samples=n_samples, n_features=16, n_targets=3, random_state=2026 + ) + Xy = QuantileDMatrix(X, y) + + params = _make_subsample_params(device, sampling_method) + + evals_result = train_result(params, Xy, num_rounds=16) + # Training should converge with subsampling + assert non_increasing(evals_result["train"]["rmse"], tolerance=0.01) + + # Test with quantile regression + params = _make_subsample_params(device, sampling_method) + params["objective"] = "reg:quantileerror" + params["quantile_alpha"] = [0.25, 0.5, 0.75] + Xy_single = QuantileDMatrix(X, y[:, 0]) + evals_result_q = train_result(params, Xy_single, num_rounds=16) + assert non_increasing(evals_result_q["train"]["quantile"], tolerance=0.01) + + +def run_gradient_based_sampling_accuracy(device: Device) -> None: + """Test that gradient-based sampling provides better accuracy.""" + n_samples = 4096 + X, y = make_regression( + n_samples=n_samples, n_features=16, n_targets=3, random_state=2026 + ) + Xy = QuantileDMatrix(X, y) + + params_uniform = _make_subsample_params(device, "uniform") + + def run(obj: Callable | None) -> None: + # Train with uniform sampling + evals_uniform: Dict[str, Dict] = {} + train( + params_uniform, + Xy, + num_boost_round=32, + evals=[(Xy, "train")], + obj=obj, + verbose_eval=False, + evals_result=evals_uniform, + ) + + # Train with gradient-based sampling + params_grad = _make_subsample_params(device, "gradient_based") + evals_grad: Dict[str, Dict] = {} + train( + params_grad, + Xy, + num_boost_round=32, + evals=[(Xy, "train")], + obj=obj, + verbose_eval=False, + evals_result=evals_grad, + ) + + uniform_final = evals_uniform["train"]["rmse"][-1] + grad_final = evals_grad["train"]["rmse"][-1] + assert non_increasing(evals_uniform["train"]["rmse"]) + assert non_increasing(evals_grad["train"]["rmse"]) + if obj is None: + assert grad_final < uniform_final + + run(None) + run(LsObj2(device, False)) diff --git a/v12x/pylib/xgboost/testing/ordinal.py b/v12x/pylib/xgboost/testing/ordinal.py new file mode 100644 index 0000000000000000000000000000000000000000..a14f9701b42fccd085074b0c8485573e151e3d50 --- /dev/null +++ b/v12x/pylib/xgboost/testing/ordinal.py @@ -0,0 +1,871 @@ +"""Tests for the ordinal re-coder.""" + +import itertools +import os +import tempfile +from concurrent.futures import ThreadPoolExecutor +from functools import cache as fcache +from typing import Any, Tuple, Type, TypeVar + +import numpy as np +import pytest + +from .._typing import EvalsLog +from ..core import DMatrix, ExtMemQuantileDMatrix, QuantileDMatrix +from ..data import _lazy_load_cudf_is_cat +from ..training import train +from .data import ( + IteratorForTest, + is_pd_cat_dtype, + make_batches, + make_categorical, + memory, +) +from .updater import get_basescore +from .utils import Device, assert_allclose, predictor_equal + + +@fcache +def get_df_impl(device: Device) -> Tuple[Type, Type]: + """Get data frame implementation based on the ]device.""" + if device == "cpu": + import pandas as pd + + Df = pd.DataFrame + Ser = pd.Series + else: + import cudf + + Df = cudf.DataFrame + Ser = cudf.Series + return Df, Ser + + +def asarray(device: Device, data: Any) -> np.ndarray: + """Wrapper to get an array.""" + if device == "cpu": + return np.asarray(data) + import cupy as cp + + return cp.asarray(data) + + +def comp_booster(device: Device, Xy: DMatrix, booster: str) -> None: + """Compare the results from DMatrix and Booster.""" + cats_dm = Xy.get_categories(export_to_arrow=True).to_arrow() + assert cats_dm is not None + + rng = np.random.default_rng(2025) + Xy.set_label(rng.normal(size=Xy.num_row())) + bst = train({"booster": booster, "device": device}, Xy, 1) + cats_bst = bst.get_categories(export_to_arrow=True).to_arrow() + assert cats_bst is not None + assert cats_dm == cats_bst + + +def run_cat_container(device: Device) -> None: + """Basic tests for the container class used by the DMatrix.""" + + def run_dispatch(device: Device, DMatrixT: Type) -> None: + Df, _ = get_df_impl(device) + # Basic test with a single feature + df = Df({"c": ["猫", "café"]}, dtype="category") + categories = df.c.cat.categories + + Xy = DMatrixT(df, enable_categorical=True) + assert Xy.feature_names == ["c"] + assert Xy.feature_types == ["c"] + results = Xy.get_categories(export_to_arrow=True).to_arrow() + assert results is not None + results_di = dict(results) + assert len(results_di["c"]) == len(categories) + for i in range(len(results_di["c"])): + assert str(results_di["c"][i]) == str(categories[i]), ( + results_di["c"][i], + categories[i], + ) + + # Test with missing values. + df = Df({"c": ["猫", None, "café", "café"]}, dtype="category") + Xy = DMatrixT(df, enable_categorical=True) + + cats = Xy.get_categories(export_to_arrow=True).to_arrow() + assert cats is not None + cats_id = dict(cats) + ser = cats_id["c"].to_pandas() + assert ser.iloc[0] == "café" + assert ser.iloc[1] == "猫" + assert ser.size == 2 + + csr = Xy.get_data() + assert csr.data.size == 3 + assert_allclose(device, csr.data, np.array([1.0, 0.0, 0.0])) + assert_allclose(device, csr.indptr, np.array([0, 1, 1, 2, 3])) + assert_allclose(device, csr.indices, np.array([0, 0, 0])) + + comp_booster(device, Xy, "gbtree") + comp_booster(device, Xy, "dart") + + # Test with explicit null-terminated strings. + df = Df({"c": ["猫", None, "café", "café\0", ""]}, dtype="category") + Xy = DMatrixT(df, enable_categorical=True) + + comp_booster(device, Xy, "gbtree") + comp_booster(device, Xy, "dart") + + with pytest.raises(ValueError, match="export_to_arrow"): + Xy.get_categories(export_to_arrow=False).to_arrow() + + for dm in (DMatrix, QuantileDMatrix): + run_dispatch(device, dm) + + +# pylint: disable=too-many-statements +def run_cat_container_mixed(device: Device) -> None: + """Run checks with mixed types.""" + import pandas as pd + + try: + is_cudf_cat = _lazy_load_cudf_is_cat() + except ImportError: + + def is_cudf_cat(_: Any) -> bool: + return False + + n_samples = int(2**10) + + def check(Xy: DMatrix, X: pd.DataFrame) -> None: + cats = Xy.get_categories(export_to_arrow=True).to_arrow() + assert cats is not None + cats_di = dict(cats) + + for fname in X.columns: + if is_pd_cat_dtype(X[fname].dtype) or is_cudf_cat(X[fname].dtype): + vf = cats_di[fname] + assert vf is not None + aw_list = sorted(vf.to_pylist()) + if is_cudf_cat(X[fname].dtype): + pd_list: list = X[fname].unique().to_arrow().to_pylist() + else: + pd_list = X[fname].unique().tolist() + if np.nan in pd_list: # pandas + pd_list.remove(np.nan) + if None in pd_list: # cudf + pd_list.remove(None) + pd_list = sorted(pd_list) + assert aw_list == pd_list + else: + assert cats_di[fname] is None + + if not hasattr(Xy, "ref"): # not quantile DMatrix. + assert not isinstance(Xy, QuantileDMatrix) + with tempfile.TemporaryDirectory() as tmpdir: + fname = os.path.join(tmpdir, "DMatrix.binary") + Xy.save_binary(fname) + + Xy_1 = DMatrix(fname) + cats_1 = Xy_1.get_categories(export_to_arrow=True).to_arrow() + assert cats_1 is not None + cats_1_di = dict(cats_1) + + for k, v_0 in cats_di.items(): + v_1 = cats_1_di[k] + if v_0 is None: + assert v_1 is None + else: + assert v_1 is not None + assert v_0.to_pylist() == v_1.to_pylist() + + comp_booster(device, Xy, "gbtree") + comp_booster(device, Xy, "dart") + + def run_dispatch(DMatrixT: Type) -> None: + # full str type + X, y = make_categorical( + n_samples, 16, 7, onehot=False, cat_dtype=np.str_, device=device + ) + Xy = DMatrixT(X, y, enable_categorical=True) + check(Xy, X) + + # str type, mixed with numerical features + X, y = make_categorical( + n_samples, + 16, + 7, + onehot=False, + cat_ratio=0.5, + cat_dtype=np.str_, + device=device, + ) + Xy = DMatrixT(X, y, enable_categorical=True) + check(Xy, X) + + # str type, mixed with numerical features and missing values + X, y = make_categorical( + n_samples, + 16, + 7, + onehot=False, + cat_ratio=0.5, + sparsity=0.5, + cat_dtype=np.str_, + device=device, + ) + Xy = DMatrixT(X, y, enable_categorical=True) + check(Xy, X) + + # int type + X, y = make_categorical( + n_samples, 16, 7, onehot=False, cat_dtype=np.int64, device=device + ) + Xy = DMatrixT(X, y, enable_categorical=True) + check(Xy, X) + + # int type, mixed with numerical features + X, y = make_categorical( + n_samples, + 16, + 7, + onehot=False, + cat_ratio=0.5, + cat_dtype=np.int64, + device=device, + ) + Xy = DMatrixT(X, y, enable_categorical=True) + check(Xy, X) + + # int type, mixed with numerical features and missing values + X, y = make_categorical( + n_samples, + 16, + 7, + onehot=False, + cat_ratio=0.5, + sparsity=0.5, + cat_dtype=np.int64, + device=device, + ) + Xy = DMatrixT(X, y, enable_categorical=True) + check(Xy, X) + + for dm in (DMatrix, QuantileDMatrix): + run_dispatch(dm) + + # No category + batches = make_batches( + n_samples_per_batch=128, n_features=4, n_batches=1, use_cupy=device == "cuda" + ) + X, y, w = map(lambda x: x[0], batches) + + for DMatrixT in (DMatrix, QuantileDMatrix): + Xy = DMatrixT(X, y, weight=w) + all_num = Xy.get_categories(export_to_arrow=True).to_arrow() + assert all_num is not None + for _, v in all_num: + assert v is None + + with pytest.raises(ValueError, match="export_to_arrow"): + Xy.get_categories(export_to_arrow=False).to_arrow() + + +def run_cat_container_iter(device: Device) -> None: + """Test the categories container for iterator-based inputs.""" + n_batches = 4 + n_features = 8 + n_samples_per_batch = 64 + n_cats = 5 + + X, y = [], [] + for _ in range(n_batches): + X_i, y_i = make_categorical( + n_samples_per_batch, + n_features, + n_cats, + onehot=False, + sparsity=0.5, + cat_dtype=np.int64, + device=device, + ) + X.append(X_i) + y.append(y_i) + + it = IteratorForTest(X, y, None, cache="cache", on_host=device == "cuda") + + Xy = ExtMemQuantileDMatrix(it, enable_categorical=True) + cats = Xy.get_categories(export_to_arrow=True).to_arrow() + assert cats is not None and len(cats) == n_features + cats_di = dict(cats) + for _, v in cats_di.items(): + assert v is not None + assert v.null_count == 0 + assert len(v) == n_cats + + +def _basic_example(device: Device) -> Tuple[Any, Any, np.ndarray, np.ndarray]: + Df, _ = get_df_impl(device) + + enc = Df({"c": ["cdef", "abc", "def"]}, dtype="category") + codes = enc.c.cat.codes # 1, 0, 2 + assert_allclose(device, asarray(device, codes), np.array([1, 0, 2])) + encoded = np.array([codes.iloc[2], codes.iloc[1]]) # def, abc + np.testing.assert_allclose(encoded, [2, 0]) + + reenc = Df({"c": ["def", "abc"]}, dtype="category") # same as `encoded` + codes = reenc.c.cat.codes + assert_allclose(device, codes, np.array([1, 0])) + + y = np.array([0, 1, 2]) + + return enc, reenc, encoded, y + + +def run_basic_predict(DMatrixT: Type, device: Device, tdevice: Device) -> None: + """Enable tests with mixed devices.""" + enc, reenc, encoded, y = _basic_example(device) + + Xy = DMatrixT(enc, y, enable_categorical=True) + booster = train({"device": tdevice}, Xy, num_boost_round=4) + + predt0 = booster.inplace_predict(reenc) + predt1 = booster.inplace_predict(encoded) + assert_allclose(device, predt0, predt1) + + fmat = DMatrixT(reenc, enable_categorical=True) + predt2 = booster.predict(fmat) + assert_allclose(device, predt0, predt2) + + +def run_cat_predict(device: Device) -> None: + """Basic tests for re-coding during prediction.""" + Df, _ = get_df_impl(device) + + for dm in (DMatrix, QuantileDMatrix): + run_basic_predict(dm, device, device) + + def run_mixed(DMatrixT: Type) -> None: + df = Df({"b": [2, 1, 3], "c": ["cdef", "abc", "def"]}, dtype="category") + y = np.array([0, 1, 2]) + + # used with the next df + b_codes = df.b.cat.codes + assert_allclose(device, asarray(device, b_codes), np.array([1, 0, 2])) + # pick codes of 3, 1 + b_encoded = np.array([b_codes.iloc[2], b_codes.iloc[1]]) + + c_codes = df.c.cat.codes + assert_allclose(device, asarray(device, c_codes), np.array([1, 0, 2])) + # pick codes of "def", "abc" + c_encoded = np.array([c_codes.iloc[2], c_codes.iloc[1]]) + encoded = np.stack([b_encoded, c_encoded], axis=1) + + Xy = DMatrixT(df, y, enable_categorical=True) + booster = train({"device": device}, Xy, num_boost_round=4) + + df = Df({"b": [3, 1], "c": ["def", "abc"]}, dtype="category") + predt0 = booster.inplace_predict(df) + predt1 = booster.inplace_predict(encoded) + assert_allclose(device, predt0, predt1) + + fmat = DMatrixT(df, enable_categorical=True) + predt2 = booster.predict(fmat) + assert_allclose(device, predt0, predt2) + + for dm in (DMatrix, QuantileDMatrix): + run_mixed(dm) + + +def run_cat_invalid(device: Device) -> None: + """Basic tests for invalid inputs.""" + Df, Ser = get_df_impl(device) + y = np.array([0, 1, 2]) + + def run_invalid(DMatrixT: Type) -> None: + df = Df({"b": [2, 1, 3], "c": ["cdef", "abc", "def"]}, dtype="category") + + Xy = DMatrixT(df, y, enable_categorical=True) + booster = train({"device": device}, Xy, num_boost_round=4) + df["b"] = df["b"].astype(np.int64) + with pytest.raises(ValueError, match="The data type doesn't match"): + booster.inplace_predict(df) + + Xy = DMatrixT(df, y, enable_categorical=True) + with pytest.raises(ValueError, match="The data type doesn't match"): + booster.predict(Xy) + + df = Df( + {"b": [2, 1, 3, 4], "c": ["cdef", "abc", "def", "bbc"]}, dtype="category" + ) + with pytest.raises(ValueError, match="Found a category not in the training"): + booster.inplace_predict(df) + + for dm in (DMatrix, QuantileDMatrix): + run_invalid(dm) + + df = Df({"b": [2, 1, 3], "c": ["cdef", "abc", "def"]}, dtype="category") + Xy = DMatrix(df, y, enable_categorical=True) + booster = train({"device": device}, Xy, num_boost_round=4) + df["c"] = Ser(asarray(device, [0, 1, 1]), dtype="category") + + msg = "index type must match between the training and test set" + + with pytest.raises(ValueError, match=msg): + booster.inplace_predict(df) + + with pytest.raises(ValueError, match=msg): + DMatrix(df, enable_categorical=True, feature_types=booster.get_categories()) + + with pytest.raises(ValueError, match=msg): + QuantileDMatrix( + df, enable_categorical=True, feature_types=booster.get_categories() + ) + + +def run_cat_oov_in_range(device: Device) -> None: + """Test OOV categories whose sorted insertion point is inside the training set.""" + Df, _ = get_df_impl(device) + + def check(train_df: Any, pred_df: Any, y: np.ndarray) -> None: + for DMatrixT in (DMatrix, QuantileDMatrix): + Xy = DMatrixT(train_df, y, enable_categorical=True) + booster = train({"device": device}, Xy, num_boost_round=1) + + with pytest.raises( + ValueError, match="Found a category not in the training" + ): + booster.inplace_predict(pred_df) + + fmat = DMatrixT(pred_df, enable_categorical=True) + with pytest.raises( + ValueError, match="Found a category not in the training" + ): + booster.predict(fmat) + + # "B2" is absent from training but sorts between "B" and "C". + train_df = Df({"cat": ["A", "B", "C", "A", "B", "B"], "x": [1, 2, 3, 4, 5, 6]}) + train_df["cat"] = train_df["cat"].astype("category") + pred_df = Df({"cat": ["B2", "C"], "x": [2, 2]}) + pred_df["cat"] = pred_df["cat"].astype("category") + check(train_df, pred_df, np.array([0, 1, 0, 1, 0, 1])) + + # Numeric category 4 is absent from training but sorts between 3 and 5. + + # XGBoost accepts only dense codes. In this test case, even though the categories + # are not contiguous, the code can still be dense. + train_df = Df({"cat": [1, 3, 5, 1, 3, 3], "x": [1, 2, 3, 4, 5, 6]}) + train_df["cat"] = train_df["cat"].astype("category") + pred_df = Df({"cat": [4, 5], "x": [2, 2]}) + pred_df["cat"] = pred_df["cat"].astype("category") + check(train_df, pred_df, np.array([0, 1, 0, 1, 0, 1])) + + +def run_cat_thread_safety(device: Device) -> None: + """Basic tests for thread safety.""" + X, y = make_categorical(2048, 16, 112, onehot=False, cat_ratio=0.5, device=device) + Xy = QuantileDMatrix(X, y, enable_categorical=True) + booster = train({"device": device}, Xy, num_boost_round=10) + + def run_thread_safety(DMatrixT: Type) -> bool: + Xy = DMatrixT(X, enable_categorical=True) + predt0 = booster.predict(Xy) + predt1 = booster.inplace_predict(X) + assert_allclose(device, predt0, predt1) + return True + + futures = [] + n_cpus = os.cpu_count() + assert n_cpus is not None + for dm in (DMatrix, QuantileDMatrix): + with ThreadPoolExecutor(max_workers=max(n_cpus, 10)) as e: + for _ in range(32): + fut = e.submit(run_thread_safety, dm) + futures.append(fut) + + for f in futures: + assert f.result() + + +U = TypeVar("U", DMatrix, QuantileDMatrix) + + +def _make_dm(DMatrixT: Type[U], ref: DMatrix, *args: Any, **kwargs: Any) -> U: + if DMatrixT is QuantileDMatrix: + return DMatrixT(*args, ref=ref, enable_categorical=True, **kwargs) + return DMatrixT(*args, enable_categorical=True, **kwargs) + + +def _run_predt( + device: Device, + DMatrixT: Type, + pred_contribs: bool, + pred_interactions: bool, + pred_leaf: bool, +) -> None: + enc, reenc, encoded, y = _basic_example(device) + + Xy = DMatrixT(enc, y, enable_categorical=True) + booster = train({"device": device}, Xy, num_boost_round=4) + + predt_0 = booster.predict( + _make_dm(DMatrixT, ref=Xy, data=reenc), + pred_contribs=pred_contribs, + pred_interactions=pred_interactions, + pred_leaf=pred_leaf, + ) + predt_1 = booster.predict( + _make_dm(DMatrixT, ref=Xy, data=encoded.reshape(2, 1), feature_names=["c"]), + pred_contribs=pred_contribs, + pred_interactions=pred_interactions, + pred_leaf=pred_leaf, + ) + atol = 1e-6 if pred_contribs or pred_interactions else 0 + assert_allclose(device, predt_0, predt_1, atol=atol) + + +def run_cat_shap(device: Device) -> None: + """Basic tests for SHAP values.""" + + for dm in (DMatrix, QuantileDMatrix): + _run_predt( + device, dm, pred_contribs=True, pred_interactions=False, pred_leaf=False + ) + + for dm in (DMatrix, QuantileDMatrix): + _run_predt( + device, dm, pred_contribs=False, pred_interactions=True, pred_leaf=False + ) + + +def run_cat_leaf(device: Device) -> None: + """Basic tests for leaf prediction.""" + # QuantileDMatrix is not supported by leaf. + _run_predt( + device, DMatrix, pred_contribs=False, pred_interactions=False, pred_leaf=True + ) + + +# pylint: disable=too-many-locals +@memory.cache +def make_recoded(device: Device, *, n_features: int = 4096) -> Tuple: + """Synthesize a test dataset with changed encoding.""" + Df, _ = get_df_impl(device) + + import pandas as pd + + # Test large column numbers. XGBoost makes some specializations for slim datasets, + # make sure we cover all the cases. + n_samples = 1024 + + # Same between old and new, with 0 ("a") and 1 ("b") exchanged their position. + old_cats = ["a", "b", "c", "d"] + new_cats = ["b", "a", "c", "d"] + mapping = {0: 1, 1: 0} + + rng = np.random.default_rng(2025) + + col_numeric = rng.uniform(0, 1, size=(n_samples, n_features // 2)) + col_categorical = rng.integers( + low=0, high=4, size=(n_samples, n_features // 2), dtype=np.int32 + ) + + df = {} # avoid fragmentation warning from pandas + for c in range(n_features): + if c % 2 == 0: + col = col_numeric[:, c // 2] + else: + codes = col_categorical[:, c // 2] + col = pd.Categorical.from_codes( + categories=old_cats, + codes=codes, + ) + df[f"f{c}"] = col + + enc = Df(df) + y = rng.normal(size=n_samples) + + reenc = enc.copy() + for c in range(n_features): + if c % 2 == 0: + continue + + name = f"f{c}" + codes_ser = reenc[name].cat.codes + if hasattr(codes_ser, "to_pandas"): # cudf + codes_ser = codes_ser.to_pandas() + new_codes = codes_ser.replace(mapping) + reenc[name] = pd.Categorical.from_codes(categories=new_cats, codes=new_codes) + reenc = Df(reenc) + assert (reenc.iloc[:, 1].cat.codes != enc.iloc[:, 1].cat.codes).any() + return enc, reenc, y, col_numeric, col_categorical + + +def run_specified_cat( # pylint: disable=too-many-locals + device: Device, +) -> None: + """Run with manually specified category encoding.""" + import pandas as pd + + # Same between old and new, with 0 ("a") and 1 ("b") exchanged their position. + old_cats = ["a", "b", "c", "d"] + new_cats = ["b", "a", "c", "d"] + + col0 = np.arange(0, 9) + col1 = pd.Categorical.from_codes( + # b, b, c, d, a, c, c, d, a + categories=old_cats, + codes=[1, 1, 2, 3, 0, 2, 2, 3, 0], + ) + df = pd.DataFrame({"f0": col0, "f1": col1}) + Df, _ = get_df_impl(device) + df = Df(df) + rng = np.random.default_rng(2025) + y = rng.uniform(size=df.shape[0]) + + for dm in (DMatrix, QuantileDMatrix): + Xy = dm(df, y, enable_categorical=True) + booster = train({"device": device}, Xy) + predt0 = booster.predict(Xy) + predt1 = booster.inplace_predict(df) + assert_allclose(device, predt0, predt1) + + col1 = pd.Categorical.from_codes( + # b, b, c, d, a, c, c, d, a + categories=new_cats, + codes=[0, 0, 2, 3, 1, 2, 2, 3, 1], + ) + df1 = Df({"f0": col0, "f1": col1}) + predt2 = booster.inplace_predict(df1) + assert_allclose(device, predt0, predt2) + + enc, reenc, y, col_numeric, col_categorical = make_recoded(device) + + Xy = DMatrix(enc, y, enable_categorical=True) + booster = train({"device": device}, Xy) + + predt0 = booster.predict(Xy) + predt1 = booster.inplace_predict(enc) + assert_allclose(device, predt0, predt1) + + Xy = DMatrix(reenc, y, enable_categorical=True) + predt2 = booster.predict(Xy) + assert_allclose(device, predt0, predt2) + + array = np.empty(shape=(reenc.shape[0], reenc.shape[1])) + + array[:, enc.dtypes == "category"] = col_categorical + array[:, enc.dtypes != "category"] = col_numeric + + if device == "cuda": + import cupy as cp + + array = cp.array(array) + + predt3 = booster.inplace_predict(array) + assert_allclose(device, predt0, predt3) + + +def run_validation(device: Device) -> None: + """Check the validation dataset is using the correct encoding.""" + enc, reenc, y, _, _ = make_recoded(device) + + Xy = DMatrix(enc, y, enable_categorical=True) + Xy_valid = DMatrix(reenc, y, enable_categorical=True) + + evals_result: EvalsLog = {} + train( + {"device": device}, + Xy, + evals=[(Xy, "Train"), (Xy_valid, "Valid")], + evals_result=evals_result, + ) + + # Evaluation dataset should have the exact same performance as the training dataset. + assert_allclose( + device, evals_result["Train"]["rmse"], evals_result["Valid"]["rmse"] + ) + + +def run_recode_dmatrix(device: Device) -> None: + """Test re-coding inpput for DMatrix.""" + import pandas as pd + + Df, _ = get_df_impl(device) + + # String index + old_cats = ["café", "猫", "c", "🐍"] + new_cats = ["猫", "café", "c", "🐍"] + + col0 = np.arange(0, 9) + col1 = pd.Categorical.from_codes( + # 猫, 猫, c, 🐍, café, c, c, 🐍, café + categories=old_cats, + codes=[1, 1, 2, 3, 0, 2, 2, 3, 0], + ) + df = Df({"f0": col0, "f1": col1}) + + Xy_0 = DMatrix(df, enable_categorical=True) + cats_0 = Xy_0.get_categories(export_to_arrow=True) + assert Xy_0.feature_types == ["int", "c"] + + col1 = pd.Categorical.from_codes( + # 猫, 猫, c, 🐍, café, c, c, 🐍, café + categories=new_cats, + codes=[0, 0, 2, 3, 1, 2, 2, 3, 1], + ) + df = Df({"f0": col0, "f1": col1}) + Xy_1 = DMatrix(df, enable_categorical=True, feature_types=cats_0) + # feature_types is still correct + assert Xy_1.feature_names == ["f0", "f1"] + assert Xy_1.feature_types == ["int", "c"] + cats_1 = Xy_1.get_categories(export_to_arrow=True) + assert cats_0.to_arrow() == cats_1.to_arrow() + assert predictor_equal(Xy_0, Xy_1) + + # Numeric index + col0 = pd.Categorical.from_codes( + categories=[5, 6, 7, 8], + codes=[0, 0, 2, 3, 1, 2, 2, 3, 1], + ) + Df, _ = get_df_impl(device) + df = Df({"cat": col0}) + for DMatrixT in (DMatrix, QuantileDMatrix): + Xy = DMatrixT(df, enable_categorical=True) + cats_0 = Xy.get_categories(export_to_arrow=True) + assert cats_0 is not None + + Xy = DMatrixT(df, enable_categorical=True, feature_types=cats_0) + cats_1 = Xy.get_categories(export_to_arrow=True) + assert cats_1 is not None + + assert cats_0.to_arrow() == cats_1.to_arrow() + + # Recode + for DMatrixT in (DMatrix, QuantileDMatrix): + enc, reenc, y, _, _ = make_recoded(device) + Xy_0 = DMatrixT(enc, y, enable_categorical=True) + cats_0 = Xy_0.get_categories(export_to_arrow=True) + + assert cats_0 is not None + + Xy_1 = DMatrixT(reenc, y, feature_types=cats_0, enable_categorical=True) + cats_1 = Xy_1.get_categories(export_to_arrow=True) + assert cats_1 is not None + + assert cats_0.to_arrow() == cats_1.to_arrow() + assert predictor_equal(Xy_0, Xy_1) + + +def run_training_continuation(device: Device) -> None: + """Test re-coding for training continuation.""" + enc, reenc, y, _, _ = make_recoded(device) + + def check(Xy_0: DMatrix, Xy_1: DMatrix) -> None: + params = {"device": device} + + r = 2 + evals_result_0: EvalsLog = {} + booster_0 = train( + params, + Xy_0, + evals=[(Xy_1, "Valid")], + num_boost_round=r, + evals_result=evals_result_0, + ) + evals_result_1: EvalsLog = {} + booster_1 = train( + params, + Xy_1, + evals=[(Xy_1, "Valid")], + xgb_model=booster_0, + num_boost_round=r, + evals_result=evals_result_1, + ) + assert get_basescore(booster_0) == get_basescore(booster_1) + + evals_result_2: EvalsLog = {} + booster_2 = train( + params, + Xy_0, + evals=[(Xy_1, "Valid")], + num_boost_round=r * 2, + evals_result=evals_result_2, + ) + # Check evaluation results + eval_concat = evals_result_0["Valid"]["rmse"] + evals_result_1["Valid"]["rmse"] + eval_full = evals_result_2["Valid"]["rmse"] + np.testing.assert_allclose(eval_full, eval_concat) + + # Test inference + for a, b in itertools.product([enc, reenc], [enc, reenc]): + predt_0 = booster_1.inplace_predict(a) + predt_1 = booster_2.inplace_predict(b) + assert_allclose(device, predt_0, predt_1, rtol=1e-5) + + # With DMatrix + for a, b in itertools.product([Xy_0, Xy_1], [Xy_0, Xy_1]): + predt_0 = booster_1.predict(a) + predt_1 = booster_2.predict(b) + assert_allclose(device, predt_0, predt_1, rtol=1e-5) + + for Train, Valid in itertools.product( + [DMatrix, QuantileDMatrix], [DMatrix, QuantileDMatrix] + ): + Xy_0 = Train(enc, y, enable_categorical=True) + if Valid is QuantileDMatrix: + Xy_1 = Valid( + reenc, + y, + enable_categorical=True, + feature_types=Xy_0.get_categories(), + ref=Xy_0, + ) + else: + Xy_1 = Valid( + reenc, y, enable_categorical=True, feature_types=Xy_0.get_categories() + ) + check(Xy_0, Xy_1) + + +def run_update(device: Device) -> None: + """Test with individual updaters.""" + enc, reenc, y, _, _ = make_recoded(device) + Xy = DMatrix(enc, y, enable_categorical=True) + booster_0 = train({"device": device}, Xy, num_boost_round=4) + model_0 = booster_0.save_raw() + cats_0 = booster_0.get_categories() + + Xy_1 = DMatrix(reenc, y, feature_types=cats_0, enable_categorical=True) + + booster_1 = train( + { + "device": device, + "updater": "prune", + "process_type": "update", + }, + Xy_1, + num_boost_round=4, + xgb_model=booster_0, + ) + model_1 = booster_1.save_raw() + + assert model_0 == model_1 # also compares the cat container inside + + +def run_recode_dmatrix_predict(device: Device) -> None: + """Run prediction with re-coded DMatrix.""" + enc, reenc, y, _, _ = make_recoded(device) + + for DMatrixT in (DMatrix, QuantileDMatrix): + Xy = DMatrix(enc, y, enable_categorical=True) + booster = train({"device": device}, Xy, num_boost_round=4) + cats_0 = booster.get_categories() + + Xy_1 = _make_dm(DMatrixT, Xy, reenc, y, feature_types=cats_0) + Xy_2 = _make_dm(DMatrixT, Xy, reenc, y) + + predt_0 = booster.predict(Xy) + predt_1 = booster.predict(Xy_1) + predt_2 = booster.predict(Xy_2) + predt_3 = booster.inplace_predict(enc) + + for predt in (predt_1, predt_2, predt_3): + assert_allclose(device, predt_0, predt) diff --git a/v12x/pylib/xgboost/testing/params.py b/v12x/pylib/xgboost/testing/params.py new file mode 100644 index 0000000000000000000000000000000000000000..6b47f4a01a0c50b7c1a38cdf6a0e15cb350132e2 --- /dev/null +++ b/v12x/pylib/xgboost/testing/params.py @@ -0,0 +1,87 @@ +"""Strategies for updater tests.""" + +from typing import cast + +import pytest + +strategies = pytest.importorskip("hypothesis.strategies") + + +exact_parameter_strategy = strategies.fixed_dictionaries( + { + "nthread": strategies.integers(1, 4), + "max_depth": strategies.integers(1, 11), + "min_child_weight": strategies.floats(0.5, 2.0), + "alpha": strategies.floats(1e-5, 2.0), + "lambda": strategies.floats(1e-5, 2.0), + "eta": strategies.floats(0.01, 0.5), + "gamma": strategies.floats(1e-5, 2.0), + "seed": strategies.integers(0, 10), + # We cannot enable subsampling as the training loss can increase + # 'subsample': strategies.floats(0.5, 1.0), + "colsample_bytree": strategies.floats(0.5, 1.0), + "colsample_bylevel": strategies.floats(0.5, 1.0), + } +) + +hist_parameter_strategy = strategies.fixed_dictionaries( + { + "max_depth": strategies.integers(1, 11), + "max_leaves": strategies.integers(0, 1024), + "max_bin": strategies.integers(2, 512), + "grow_policy": strategies.sampled_from(["lossguide", "depthwise"]), + "min_child_weight": strategies.floats(0.5, 2.0), + # We cannot enable subsampling as the training loss can increase + # 'subsample': strategies.floats(0.5, 1.0), + "colsample_bytree": strategies.floats(0.5, 1.0), + "colsample_bylevel": strategies.floats(0.5, 1.0), + } +).filter( + lambda x: (cast(int, x["max_depth"]) > 0 or cast(int, x["max_leaves"]) > 0) + and (cast(int, x["max_depth"]) > 0 or x["grow_policy"] == "lossguide") +) + +hist_cache_strategy = strategies.fixed_dictionaries( + {"max_cached_hist_node": strategies.sampled_from([1, 4, 1024, 2**31])} +) + +hist_multi_parameter_strategy = strategies.fixed_dictionaries( + { + "max_depth": strategies.integers(1, 11), + "max_leaves": strategies.integers(0, 1024), + "max_bin": strategies.integers(2, 512), + "multi_strategy": strategies.sampled_from( + ["multi_output_tree", "one_output_per_tree"] + ), + "grow_policy": strategies.sampled_from(["lossguide", "depthwise"]), + "min_child_weight": strategies.floats(0.5, 2.0), + # We cannot enable subsampling as the training loss can increase + # 'subsample': strategies.floats(0.5, 1.0), + "colsample_bytree": strategies.floats(0.5, 1.0), + "colsample_bylevel": strategies.floats(0.5, 1.0), + } +).filter( + lambda x: (cast(int, x["max_depth"]) > 0 or cast(int, x["max_leaves"]) > 0) + and (cast(int, x["max_depth"]) > 0 or x["grow_policy"] == "lossguide") +) + +cat_parameter_strategy = strategies.fixed_dictionaries( + { + "max_cat_to_onehot": strategies.integers(1, 128), + "max_cat_threshold": strategies.integers(1, 128), + } +) + +lambdarank_parameter_strategy = strategies.fixed_dictionaries( + { + "lambdarank_unbiased": strategies.sampled_from([True, False]), + "lambdarank_pair_method": strategies.sampled_from(["topk", "mean"]), + "lambdarank_num_pair_per_sample": strategies.integers(1, 8), + "lambdarank_bias_norm": strategies.floats(0.5, 2.0), + "objective": strategies.sampled_from( + ["rank:ndcg", "rank:map", "rank:pairwise"] + ), + } +).filter( + lambda x: not (x["lambdarank_unbiased"] and x["lambdarank_pair_method"] == "mean") +) diff --git a/v12x/pylib/xgboost/testing/parse_tree.py b/v12x/pylib/xgboost/testing/parse_tree.py new file mode 100644 index 0000000000000000000000000000000000000000..37ad44d0208f2a19a692267032f619b92fcd04aa --- /dev/null +++ b/v12x/pylib/xgboost/testing/parse_tree.py @@ -0,0 +1,137 @@ +"""Tests for parsing trees.""" + +import json + +import numpy as np +import pytest +from sklearn.datasets import make_regression + +from ..core import Booster, DMatrix, QuantileDMatrix +from ..sklearn import XGBRegressor +from ..training import train +from .data import make_categorical +from .updater import ResetStrategy +from .utils import Device + + +def run_tree_to_df_categorical(tree_method: str, device: Device) -> None: + """Tests tree_to_df with categorical features.""" + + import pandas as pd + + X, y = make_categorical(100, 10, 31, onehot=False) + Xy = DMatrix(X, y, enable_categorical=True) + booster = train( + {"tree_method": tree_method, "device": device}, Xy, num_boost_round=10 + ) + df = booster.trees_to_dataframe() + + all_ids = set(df["ID"]) + for _, x in df.iterrows(): + if x["Feature"] == "Leaf": + # A leaf carries its scalar weight in ``Gain`` and no split info. + assert pd.isna(x["Split"]) + assert not isinstance(x["Category"], list) + assert pd.isna(x["Yes"]) + assert pd.isna(x["No"]) + assert pd.isna(x["Missing"]) + else: + # A categorical split has a missing threshold and a non-empty list of + # integer category codes rendered as strings. + assert pd.isna(x["Split"]) + assert isinstance(x["Category"], list) and len(x["Category"]) >= 1 + assert all(int(c) >= 0 for c in x["Category"]) + # Branch ids must reference existing nodes in the same frame. + assert x["Yes"] in all_ids + assert x["No"] in all_ids + assert x["Missing"] in all_ids + + +def _expected_leaf_vectors(booster: Booster) -> dict[tuple[int, int], list[float]]: + """Map ``(tree_id, node_id) -> list[float]`` of leaf outputs.""" + model = json.loads(booster.save_raw(raw_format="json")) + trees = model["learner"]["gradient_booster"]["model"]["trees"] + out: dict[tuple[int, int], list[float]] = {} + for tid, tree in enumerate(trees): + n_targets = int(tree["tree_param"]["size_leaf_vector"]) + if n_targets <= 1: + continue + left = tree["left_children"] + right = tree["right_children"] + leaf_weights = tree["leaf_weights"] + for nid, lc in enumerate(left): + if lc == -1: # leaf + beg = right[nid] * n_targets + out[(tid, nid)] = leaf_weights[beg : beg + n_targets] + return out + + +def run_tree_to_df_vector_leaf_mixed(device: Device) -> None: + """Tests trees_to_dataframe on a mixed scalar + vector-leaf booster.""" + import pandas as pd + + n_targets = 3 + X, y = make_regression( + n_samples=512, n_features=10, n_targets=n_targets, random_state=2025 + ) + booster = train( + {"device": device, "multi_strategy": "multi_output_tree", "max_depth": 3}, + QuantileDMatrix(X, y), + num_boost_round=6, + callbacks=[ResetStrategy()], + ) + + df = booster.trees_to_dataframe() + assert str(df["Gain"].dtype) == "Float64" + assert "Target" in df.columns + + expected = _expected_leaf_vectors(booster) + # Tree ids that use vector leaves (each such tree has at least one leaf). + vector_tids = {tid for tid, _ in expected} + assert len(vector_tids) > 0 + + all_ids = set(df["ID"]) + n_vector_leaf_rows = 0 + for _, x in df.iterrows(): + tid = int(x["Tree"]) + key = (tid, int(x["Node"])) + if x["Feature"] == "Leaf": + # Leaf rows don't carry split information + assert pd.isna(x["Split"]) + assert x["Category"] is None + assert pd.isna(x["Yes"]) + assert pd.isna(x["No"]) + assert pd.isna(x["Missing"]) + if key in expected: + # Vector-leaf leaf: expanded into one row per target, with the scalar + # Gain equal to the target-th entry of the leaf vector. + n_vector_leaf_rows += 1 + target = int(x["Target"]) + assert target in range(n_targets) + np.testing.assert_allclose(float(x["Gain"]), expected[key][target]) + else: + # Scalar-tree leaf: a single row carrying the tree's concrete target. + assert not pd.isna(x["Target"]) + else: + assert isinstance(float(x["Gain"]), float) + assert x["Yes"] in all_ids + assert x["No"] in all_ids + assert x["Missing"] in all_ids + if tid in vector_tids: + # A split node does not have a concept of target. + assert pd.isna(x["Target"]) + else: + # A scalar-tree split belongs to that tree's target. + assert not pd.isna(x["Target"]) + # Every vector-leaf node expands into exactly `n_targets` rows. + assert n_vector_leaf_rows == n_targets * len(expected) + + +def run_split_value_histograms(tree_method: str, device: Device) -> None: + """Tests split_value_histograms with categorical features.""" + X, y = make_categorical(1000, 10, 13, onehot=False) + reg = XGBRegressor(tree_method=tree_method, enable_categorical=True, device=device) + reg.fit(X, y) + + with pytest.raises(ValueError, match="doesn't"): + reg.get_booster().get_split_value_histogram("3", bins=5) diff --git a/v12x/pylib/xgboost/testing/plotting.py b/v12x/pylib/xgboost/testing/plotting.py new file mode 100644 index 0000000000000000000000000000000000000000..9966cf2c6c92df0939ad22ce9dde0e7e157a2503 --- /dev/null +++ b/v12x/pylib/xgboost/testing/plotting.py @@ -0,0 +1,29 @@ +"""Test plotting functions for XGBoost.""" + +import json + +from graphviz import Source +from matplotlib.axes import Axes + +from ..plotting import plot_tree, to_graphviz +from ..sklearn import XGBRegressor +from .data import make_categorical +from .utils import Device + + +def run_categorical(tree_method: str, device: Device) -> None: + """Tests plotting functions for categorical features.""" + X, y = make_categorical(1000, 31, 19, onehot=False) + reg = XGBRegressor( + enable_categorical=True, n_estimators=10, tree_method=tree_method, device=device + ) + reg.fit(X, y) + trees = reg.get_booster().get_dump(dump_format="json") + for tree in trees: + j_tree = json.loads(tree) + assert "leaf" in j_tree.keys() or isinstance(j_tree["split_condition"], list) + + graph = to_graphviz(reg, tree_idx=len(j_tree) - 1) + assert isinstance(graph, Source) + ax = plot_tree(reg, tree_idx=len(j_tree) - 1) + assert isinstance(ax, Axes) diff --git a/v12x/pylib/xgboost/testing/predict.py b/v12x/pylib/xgboost/testing/predict.py new file mode 100644 index 0000000000000000000000000000000000000000..6e9cca8691e07228e861ddd9aca6ccff26ba1959 --- /dev/null +++ b/v12x/pylib/xgboost/testing/predict.py @@ -0,0 +1,176 @@ +"""Tests for inference.""" + +import json +from typing import Any, Optional, Type + +import numpy as np +import pytest +from scipy.special import logit # pylint: disable=no-name-in-module + +from ..compat import import_cupy +from ..core import Booster, DMatrix, ExtMemQuantileDMatrix, QuantileDMatrix +from ..training import train +from .data import IteratorForTest +from .shared import validate_leaf_output +from .updater import ResetStrategy, get_basescore +from .utils import Device + + +def _make_leaf_dmatrix( + device: Device, DMatrixT: Type[DMatrix], X: np.ndarray, y: np.ndarray +) -> tuple[DMatrix, Optional[QuantileDMatrix]]: + """Make a prediction matrix and an in-memory reference for external memory.""" + if DMatrixT is not ExtMemQuantileDMatrix: + return DMatrixT(X, y), None + + xp = import_cupy() if device.startswith("cuda") else np + X, y = xp.asarray(X), xp.asarray(y) + it = IteratorForTest( + xp.array_split(X, 4), + xp.array_split(y, 4), + None, + cache="cache", + on_host=True, + min_cache_page_bytes=0, + ) + m = ExtMemQuantileDMatrix( + it, cache_host_ratio=1.0 if device.startswith("cuda") else None + ) + return m, QuantileDMatrix(X, y, ref=m) + + +def _predict_leaf( + booster: Booster, + m: DMatrix, + reference: Optional[DMatrix], + **kwargs: Any, +) -> np.ndarray: + """Predict leaves and compare external memory against in-memory data.""" + predt = booster.predict(m, pred_leaf=True, **kwargs) + if reference is not None: + ref_predt = booster.predict(reference, pred_leaf=True, **kwargs) + np.testing.assert_array_equal(predt, ref_predt) + return predt + + +def _validate_leaf_indices(leaf: np.ndarray, trees: list[dict]) -> None: + """Validate that each predicted node is a leaf in the corresponding tree.""" + assert leaf.ndim == 2 + assert leaf.shape[1] == len(trees) + for tree_idx, tree in enumerate(trees): + prediction = leaf[:, tree_idx] + assert np.isfinite(prediction).all() + node_idx = prediction.astype(np.int64) + np.testing.assert_array_equal(prediction, node_idx) + + left = np.asarray(tree["left_children"]) + assert ((0 <= node_idx) & (node_idx < left.size)).all() + assert (left[node_idx] == -1).all() + + +# pylint: disable=too-many-locals +def run_predict_leaf(device: Device, DMatrixT: Type[DMatrix]) -> np.ndarray: + """Run tests for leaf index prediction.""" + rows, cols, classes = 100, 4, 5 + n_parallel, n_rounds = 4, 10 + rng = np.random.RandomState(1994) + X = rng.randn(rows, cols) + y = rng.randint(low=0, high=classes, size=rows) + + m, reference = _make_leaf_dmatrix(device, DMatrixT, X, y) + dtrain = DMatrix(X, y) if DMatrixT is ExtMemQuantileDMatrix else m + booster = train( + { + "num_parallel_tree": n_parallel, + "num_class": classes, + "tree_method": "hist", + }, + dtrain, + num_boost_round=n_rounds, + ) + booster.set_param({"device": device}) + assert booster.predict(DMatrix(np.empty((0, cols))), pred_leaf=True).shape[0] == 0 + + leaf = _predict_leaf(booster, m, reference, strict_shape=True) + assert leaf.shape == (rows, n_rounds, classes, n_parallel) + validate_leaf_output(leaf, n_parallel) + + n_iters = 2 + sliced = _predict_leaf( + booster, + m, + reference, + iteration_range=(0, n_iters), + strict_shape=True, + ) + assert sliced.shape == (rows, n_iters, classes, n_parallel) + + # When there's only 1 tree, the output is a 1 dim vector + booster = train({"tree_method": "hist"}, num_boost_round=1, dtrain=dtrain) + booster.set_param({"device": device}) + assert _predict_leaf(booster, m, reference).shape == (rows,) + + # The first two rounds have only vector-leaf trees. The full model contains + # both vector- and scalar-leaf trees. + mixed_rounds = 4 + booster = train( + { + "num_parallel_tree": classes, + "num_class": classes, + "tree_method": "hist", + "multi_strategy": "multi_output_tree", + }, + dtrain, + num_boost_round=mixed_rounds, + callbacks=[ResetStrategy()], + ) + booster.set_param({"device": device}) + + model = json.loads(booster.save_raw(raw_format="json")) + gbtree_model = model["learner"]["gradient_booster"]["model"] + trees = gbtree_model["trees"] + iteration_indptr = gbtree_model["iteration_indptr"] + leaf_sizes = [int(tree["tree_param"]["size_leaf_vector"]) for tree in trees] + vector_end = 2 + assert set(leaf_sizes) == {1, classes} + assert set(leaf_sizes[: iteration_indptr[vector_end]]) == {classes} + + for end, n_trees in [(vector_end, iteration_indptr[vector_end]), (0, len(trees))]: + kwargs = {"iteration_range": (0, end)} + predt = _predict_leaf(booster, m, reference, **kwargs) + selected_trees = trees[:n_trees] + assert predt.shape == (rows, len(selected_trees)) + _validate_leaf_indices(predt, selected_trees) + with pytest.raises(ValueError, match="vector leaf trees"): + booster.predict(m, pred_leaf=True, strict_shape=True, **kwargs) + + return leaf + + +def run_base_margin_vs_base_score(device: Device) -> None: + """Test for the relation between score and margin.""" + from sklearn.datasets import make_classification + + intercept = 0.5 + + X, y = make_classification(random_state=2025) + booster = train( + {"base_score": intercept, "objective": "binary:logistic", "device": device}, + dtrain=DMatrix(X, y), + num_boost_round=1, + ) + np.testing.assert_allclose(get_basescore(booster), intercept) + predt_0 = booster.predict(DMatrix(X, y)) + + margin = np.full(y.shape, fill_value=logit(intercept), dtype=np.float32) + Xy = DMatrix(X, y, base_margin=margin) + # 0.2 is a dummy value + booster = train( + {"base_score": 0.2, "objective": "binary:logistic", "device": device}, + dtrain=Xy, + num_boost_round=1, + ) + np.testing.assert_allclose(get_basescore(booster), 0.2) + predt_1 = booster.predict(Xy) + + np.testing.assert_allclose(predt_0, predt_1) diff --git a/v12x/pylib/xgboost/testing/quantile_dmatrix.py b/v12x/pylib/xgboost/testing/quantile_dmatrix.py new file mode 100644 index 0000000000000000000000000000000000000000..f5e861c5009814be99e113436dd5d51d7e12c6eb --- /dev/null +++ b/v12x/pylib/xgboost/testing/quantile_dmatrix.py @@ -0,0 +1,68 @@ +"""QuantileDMatrix related tests.""" + +import numpy as np +import pytest +from sklearn.model_selection import train_test_split + +import xgboost as xgb + +from .data import make_batches, make_categorical + + +def check_ref_quantile_cut(device: str) -> None: + """Check obtaining the same cut values given a reference.""" + X, y, _ = ( + data[0] + for data in make_batches( + n_samples_per_batch=8192, + n_features=16, + n_batches=1, + use_cupy=device.startswith("cuda"), + ) + ) + + X_train, X_valid, y_train, y_valid = train_test_split(X, y) + Xy_train = xgb.QuantileDMatrix(X_train, y_train) + Xy_valid = xgb.QuantileDMatrix(X_valid, y_valid, ref=Xy_train) + + cut_train = Xy_train.get_quantile_cut() + cut_valid = Xy_valid.get_quantile_cut() + + np.testing.assert_allclose(cut_train[0], cut_valid[0]) + np.testing.assert_allclose(cut_train[1], cut_valid[1]) + + Xy_valid = xgb.QuantileDMatrix(X_valid, y_valid) + cut_valid = Xy_valid.get_quantile_cut() + assert not np.allclose(cut_train[1], cut_valid[1]) + + +def check_categorical_strings(device: str) -> None: + """Check string inputs.""" + if device == "cpu": + pd = pytest.importorskip("pandas") + else: + pd = pytest.importorskip("cudf") + + n_categories = 32 + X, y = make_categorical( + 1024, + 8, + n_categories, + onehot=False, + cat_dtype=np.str_, + cat_ratio=0.5, + shuffle=True, + ) + X = pd.DataFrame(X) + + Xy = xgb.QuantileDMatrix(X, y, enable_categorical=True) + assert Xy.num_col() == 8 + cuts = Xy.get_quantile_cut() + indptr = cuts[0] + values = cuts[1] + for i in range(1, len(indptr)): + f_idx = i - 1 + if isinstance(X[X.columns[f_idx]].dtype, pd.CategoricalDtype): + beg, end = indptr[f_idx], indptr[i] + col = values[beg:end] + np.testing.assert_allclose(col, np.arange(0, n_categories)) diff --git a/v12x/pylib/xgboost/testing/ranking.py b/v12x/pylib/xgboost/testing/ranking.py new file mode 100644 index 0000000000000000000000000000000000000000..e1f784dda6bf4cb75a97bf7f86f5becb928bf567 --- /dev/null +++ b/v12x/pylib/xgboost/testing/ranking.py @@ -0,0 +1,210 @@ +# pylint: disable=too-many-locals +"""Tests for learning to rank.""" + +from types import ModuleType +from typing import Any + +import numpy as np +import pytest + +import xgboost as xgb +from xgboost import testing as tm + +from .data import make_ltr +from .utils import Device + + +def run_ranking_qid_df(impl: ModuleType, tree_method: str, device: Device) -> None: + """Test ranking with qid packed into X.""" + import scipy.sparse + from sklearn.metrics import mean_squared_error + from sklearn.model_selection import StratifiedGroupKFold, cross_val_score + + X, y, q, _ = make_ltr(n_samples=128, n_features=2, n_query_groups=8, max_rel=3) + + # pack qid into x using dataframe + df = impl.DataFrame(X) + df["qid"] = q + ranker = xgb.XGBRanker( + n_estimators=3, eval_metric="ndcg", tree_method=tree_method, device=device + ) + ranker.fit(df, y) + s = ranker.score(df, y) + assert s > 0.7 + + # works with validation datasets as well + valid_df = df.copy() + valid_df.iloc[0, 0] = 3.0 + ranker.fit(df, y, eval_set=[(valid_df, y)]) + + # same as passing qid directly + ranker = xgb.XGBRanker( + n_estimators=3, eval_metric="ndcg", tree_method=tree_method, device=device + ) + ranker.fit(X, y, qid=q) + s1 = ranker.score(df, y) + assert np.isclose(s, s1) + + # Works with standard sklearn cv + if device == "cpu": + # we need cuML for this. + kfold = StratifiedGroupKFold(shuffle=False) + results = cross_val_score(ranker, df, y, cv=kfold, groups=df.qid) + assert len(results) == 5 + + # Works with custom metric + def neg_mse(y_true: Any, y_pred: Any, **kwargs: Any) -> float: + if hasattr(y_true, "get"): + y_true = y_true.get() + return -float(mean_squared_error(y_true, y_pred, **kwargs)) + + ranker = xgb.XGBRanker( + n_estimators=3, + eval_metric=neg_mse, + tree_method=tree_method, + device=device, + disable_default_eval_metric=True, + ) + ranker.fit(df, y, eval_set=[(valid_df, y)]) + score = ranker.score(valid_df, y) + assert np.isclose(score, ranker.evals_result()["validation_0"]["neg_mse"][-1]) + + # Works with sparse data + if device == "cpu": + # no sparse with cuDF + X_csr = scipy.sparse.csr_matrix(X) + df = impl.DataFrame.sparse.from_spmatrix( + X_csr, columns=[str(i) for i in range(X.shape[1])] + ) + df["qid"] = q + ranker = xgb.XGBRanker( + n_estimators=3, + eval_metric="ndcg", + tree_method=tree_method, + device=device, + ) + ranker.fit(df, y) + s2 = ranker.score(df, y) + assert np.isclose(s2, s) + + with pytest.raises(ValueError, match="Either `group` or `qid`."): + ranker.fit(df, y, eval_set=[(X, y)]) + + +def run_ranking_categorical(device: str) -> None: + """Test LTR with categorical features.""" + from sklearn.model_selection import cross_val_score + + X, y = tm.make_categorical( + n_samples=512, n_features=10, n_categories=3, onehot=False + ) + # NDCG requires non-negative integer relevance labels. + y = np.clip(np.round(y - y.min()).astype(int), 0, None) + rng = np.random.default_rng(1994) + qid = rng.choice(3, size=y.shape[0]) + qid = np.sort(qid) + X["qid"] = qid + + ltr = xgb.XGBRanker(enable_categorical=True, device=device) + ltr.fit(X, y) + score = ltr.score(X, y) + assert score > 0.9 + + ltr = xgb.XGBRanker(enable_categorical=True, device=device) + + # test using the score function inside sklearn. + scores = cross_val_score(ltr, X, y) + for s in scores: + assert s > 0.7 + + +def run_normalization(device: str) -> None: + """Test normalization.""" + X, y, qid, _ = make_ltr(2048, 4, 64, 3) + # top-k + ltr = xgb.XGBRanker(objective="rank:pairwise", n_estimators=4, device=device) + ltr.fit(X, y, qid=qid, eval_set=[(X, y)], eval_qid=[qid]) + e0 = ltr.evals_result() + + ltr = xgb.XGBRanker( + objective="rank:pairwise", + n_estimators=4, + device=device, + lambdarank_normalization=False, + ) + ltr.fit(X, y, qid=qid, eval_set=[(X, y)], eval_qid=[qid]) + e1 = ltr.evals_result() + + # mean + ltr = xgb.XGBRanker( + objective="rank:pairwise", + n_estimators=4, + device=device, + lambdarank_pair_method="mean", + lambdarank_normalization=True, + ) + ltr.fit(X, y, qid=qid, eval_set=[(X, y)], eval_qid=[qid]) + e0 = ltr.evals_result() + + ltr = xgb.XGBRanker( + objective="rank:pairwise", + n_estimators=4, + device=device, + lambdarank_pair_method="mean", + lambdarank_normalization=False, + ) + ltr.fit(X, y, qid=qid, eval_set=[(X, y)], eval_qid=[qid]) + e1 = ltr.evals_result() + # no normalization since the number of pairs is 1. + assert e1["validation_0"]["ndcg"][-1] == e0["validation_0"]["ndcg"][-1] + + # mean + ltr = xgb.XGBRanker( + objective="rank:pairwise", + n_estimators=4, + device=device, + lambdarank_pair_method="mean", + lambdarank_normalization=True, + lambdarank_num_pair_per_sample=4, + ) + ltr.fit(X, y, qid=qid, eval_set=[(X, y)], eval_qid=[qid]) + e0 = ltr.evals_result() + + ltr = xgb.XGBRanker( + objective="rank:pairwise", + n_estimators=4, + device=device, + lambdarank_pair_method="mean", + lambdarank_normalization=False, + lambdarank_num_pair_per_sample=4, + ) + ltr.fit(X, y, qid=qid, eval_set=[(X, y)], eval_qid=[qid]) + e1 = ltr.evals_result() + assert e1["validation_0"]["ndcg"][-1] != e0["validation_0"]["ndcg"][-1] + + +def run_score_normalization(device: str, objective: str) -> None: + """Test normalization by score differences.""" + if objective == "rank:map": + # Binary relevance + X, y, qid, _ = make_ltr(4096, 4, 64, max_rel=1) + else: + X, y, qid, _ = make_ltr(4096, 4, 64, 3) + ltr = xgb.XGBRanker(objective=objective, n_estimators=4, device=device) + ltr.fit(X, y, qid=qid, eval_set=[(X, y)], eval_qid=[qid]) + e0 = ltr.evals_result() + + ltr = xgb.XGBRanker( + objective="rank:pairwise", + n_estimators=4, + device=device, + lambdarank_score_normalization=False, + ) + ltr.fit(X, y, qid=qid, eval_set=[(X, y)], eval_qid=[qid]) + e1 = ltr.evals_result() + + m0, m1 = ( + list(e0["validation_0"].values())[-1][-1], + list(e1["validation_0"].values())[-1][-1], + ) + assert m0 != m1 diff --git a/v12x/pylib/xgboost/testing/shared.py b/v12x/pylib/xgboost/testing/shared.py new file mode 100644 index 0000000000000000000000000000000000000000..fcdfe5a6e13822197f7b172e15d9f62a898acaee --- /dev/null +++ b/v12x/pylib/xgboost/testing/shared.py @@ -0,0 +1,100 @@ +"""Testing code shared by other tests.""" + +import collections +import importlib.util +import json +import os +import tempfile +from typing import Any, Callable, Dict, Type + +import numpy as np + +import xgboost as xgb +from xgboost._typing import ArrayLike + + +def validate_leaf_output(leaf: np.ndarray, num_parallel_tree: int) -> None: + """Validate output for predict leaf tests.""" + for i in range(leaf.shape[0]): # n_samples + for j in range(leaf.shape[1]): # n_rounds + for k in range(leaf.shape[2]): # n_classes + tree_group = leaf[i, j, k, :] + assert tree_group.shape[0] == num_parallel_tree + # No sampling, all trees within forest are the same + assert np.all(tree_group == tree_group[0]) + + +def validate_data_initialization( + dmatrix: Type, model: Type[xgb.XGBModel], X: ArrayLike, y: ArrayLike +) -> None: + """Assert that we don't create duplicated DMatrix.""" + + old_init = dmatrix.__init__ + count = [0] + + def new_init(self: Any, **kwargs: Any) -> Callable: + count[0] += 1 + return old_init(self, **kwargs) + + dmatrix.__init__ = new_init + model(n_estimators=1).fit(X, y, eval_set=[(X, y)]) + + assert count[0] == 1 + count[0] = 0 # only 1 DMatrix is created. + + y_copy = y.copy() + model(n_estimators=1).fit(X, y, eval_set=[(X, y_copy)]) + assert count[0] == 2 # a different Python object is considered different + + dmatrix.__init__ = old_init + + +# pylint: disable=too-many-arguments,too-many-locals +def get_feature_weights( + *, + X: ArrayLike, + y: ArrayLike, + fw: np.ndarray, + parser_path: str, + tree_method: str, + model: Type[xgb.XGBModel] = xgb.XGBRegressor, +) -> np.ndarray: + """Get feature weights using the demo parser.""" + with tempfile.TemporaryDirectory() as tmpdir: + colsample_bynode = 0.5 + reg = model( + tree_method=tree_method, + colsample_bynode=colsample_bynode, + feature_weights=fw, + ) + + reg.fit(X, y) + model_path = os.path.join(tmpdir, "model.json") + reg.save_model(model_path) + with open(model_path, "r", encoding="utf-8") as fd: + model = json.load(fd) + + spec = importlib.util.spec_from_file_location("JsonParser", parser_path) + assert spec is not None + jsonm = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(jsonm) + model = jsonm.Model(model) + splits: Dict[int, int] = {} + total_nodes = 0 + for tree in model.trees: + n_nodes = len(tree.nodes) + total_nodes += n_nodes + for n in range(n_nodes): + if tree.is_leaf(n): + continue + if splits.get(tree.split_index(n), None) is None: + splits[tree.split_index(n)] = 1 + else: + splits[tree.split_index(n)] += 1 + + od = collections.OrderedDict(sorted(splits.items())) + tuples = list(od.items()) + k, v = list(zip(*tuples)) + w = np.polyfit(k, v, deg=1) + return w diff --git a/v12x/pylib/xgboost/testing/updater.py b/v12x/pylib/xgboost/testing/updater.py new file mode 100644 index 0000000000000000000000000000000000000000..faf0c61e5287f40cc7f71393e8932d4c413503af --- /dev/null +++ b/v12x/pylib/xgboost/testing/updater.py @@ -0,0 +1,736 @@ +"""Tests for updaters.""" + +import json +from functools import partial, update_wrapper +from string import ascii_lowercase +from typing import Any, Dict, List, Optional, Union, overload + +import numpy as np +import pytest +from sklearn.datasets import make_regression + +import xgboost.testing as tm + +from ..callback import TrainingCallback +from ..compat import import_cupy +from ..core import ( + Booster, + DataIter, + DMatrix, + ExtMemQuantileDMatrix, + QuantileDMatrix, +) +from ..data import is_pd_cat_dtype +from ..sklearn import XGBModel, XGBRegressor +from ..training import train +from .data import IteratorForTest, make_batches, make_categorical +from .data_iter import CatIter +from .utils import Device, assert_allclose, non_increasing + + +@overload +def get_basescore(model: XGBModel) -> List[float]: ... + + +@overload +def get_basescore(model: Booster) -> List[float]: ... + + +@overload +def get_basescore(model: Dict[str, Any]) -> List[float]: ... + + +def get_basescore( + model: Union[XGBModel, Booster, Dict], +) -> List[float]: + """Get base score from an XGBoost sklearn estimator.""" + if isinstance(model, XGBModel): + model = model.get_booster() + + if isinstance(model, dict): + jintercept = model["learner"]["learner_model_param"]["base_score"] + else: + jintercept = json.loads(model.save_config())["learner"]["learner_model_param"][ + "base_score" + ] + return json.loads(jintercept) + + +# pylint: disable=too-many-locals +def check_quantile_loss(tree_method: str, weighted: bool, device: Device) -> None: + """Test for quantile loss.""" + from sklearn.metrics import mean_pinball_loss + + from xgboost.sklearn import _metric_decorator + + n_samples = 4096 + n_features = 8 + n_estimators = 8 + + rng = np.random.RandomState(1994) + # pylint: disable=unbalanced-tuple-unpacking + X, y = make_regression( + n_samples=n_samples, + n_features=n_features, + random_state=rng, + ) + if weighted: + weight = rng.random(size=n_samples) + else: + weight = None + + Xy = QuantileDMatrix(X, y, weight=weight) + + alpha = np.array([0.1, 0.5]) + # non-zero base score can cause floating point difference with GPU predictor. + # multi-class has small difference than single target in the prediction kernel + base_score = np.zeros(shape=alpha.shape, dtype=np.float32) + evals_result: Dict[str, Dict] = {} + booster_multi = train( + { + "objective": "reg:quantileerror", + "tree_method": tree_method, + "device": device, + "quantile_alpha": alpha, + "base_score": base_score, + }, + Xy, + num_boost_round=n_estimators, + evals=[(Xy, "Train")], + evals_result=evals_result, + ) + predt_multi = booster_multi.predict(Xy, strict_shape=True) + + assert non_increasing(evals_result["Train"]["quantile"]) + # This deterministic fixture finishes near 30.4 with the scale-correct MM update. + assert evals_result["Train"]["quantile"][-1] < 35.0 + # check that there's a way to use custom metric and compare the results. + metrics = [ + _metric_decorator( + update_wrapper( + partial(mean_pinball_loss, sample_weight=weight, alpha=alpha[i]), + mean_pinball_loss, + ) + ) + for i in range(alpha.size) + ] + + predts = np.empty(predt_multi.shape) + for i in range(alpha.shape[0]): + a = alpha[i] + + booster_i = train( + { + "objective": "reg:quantileerror", + "tree_method": tree_method, + "device": device, + "quantile_alpha": a, + "base_score": base_score[i], + }, + Xy, + num_boost_round=n_estimators, + evals=[(Xy, "Train")], + custom_metric=metrics[i], + evals_result=evals_result, + ) + assert non_increasing(evals_result["Train"]["quantile"]) + # The slower median case finishes near 38.7; retain an absolute accuracy check. + assert evals_result["Train"]["quantile"][-1] < 40.0 + np.testing.assert_allclose( + np.array(evals_result["Train"]["quantile"]), + np.array(evals_result["Train"]["mean_pinball_loss"]), + atol=1e-6, + rtol=1e-6, + ) + predts[:, i] = booster_i.predict(Xy) + + # Multi-quantile output is ordered row-wise to prevent crossing. Training remains + # independent per quantile, so it matches sorted single-quantile predictions. + np.testing.assert_allclose(np.sort(predts, axis=1), predt_multi) + assert np.all(np.diff(predt_multi, axis=1) >= 0.0) + + +def check_quantile_loss_rf( + device: Device, tree_method: str, multi_strategy: str +) -> None: + """Test quantile loss with boosting random forest.""" + # pylint: disable=unbalanced-tuple-unpacking + X, y = make_regression(n_samples=2048, n_features=16, random_state=2026) + Xy = DMatrix(X, y) + + def run(params: Dict[str, Any], metric: str) -> None: + evals_result_0: Dict[str, Dict] = {} + params["num_parallel_tree"] = 2 + train( + params, + Xy, + num_boost_round=8, + evals=[(Xy, "Train")], + evals_result=evals_result_0, + ) + + evals_result_1: Dict[str, Dict] = {} + params["num_parallel_tree"] = 1 + train( + params, + Xy, + num_boost_round=8, + evals=[(Xy, "Train")], + evals_result=evals_result_1, + ) + # Without subsample, the result should be the same (barring floating point + # errors). + np.testing.assert_allclose( + evals_result_0["Train"][metric], evals_result_1["Train"][metric] + ) + assert non_increasing(evals_result_0["Train"][metric]) + + alpha = np.array([0.1, 0.5, 0.9]) + params = { + "objective": "reg:quantileerror", + "tree_method": tree_method, + "device": device, + "quantile_alpha": alpha, + "multi_strategy": multi_strategy, + } + run(params, "quantile") + + # Now test with MAE + params.pop("quantile_alpha") + params["objective"] = "reg:absoluteerror" + run(params, "mae") + + +def check_quantile_loss_extmem( + n_samples_per_batch: int, + n_features: int, + n_batches: int, + tree_method: str, + device: str, +) -> None: + """Check external memory with the quantile objective.""" + it = IteratorForTest( + *make_batches(n_samples_per_batch, n_features, n_batches, device != "cpu"), + cache="cache", + on_host=False, + ) + Xy_it = DMatrix(it) + params = { + "tree_method": tree_method, + "objective": "reg:quantileerror", + "device": device, + "quantile_alpha": [0.2, 0.8], + } + booster_it = train(params, Xy_it) + X, y, w = it.as_arrays() + Xy = DMatrix(X, y, weight=w) + booster = train(params, Xy) + + predt_it = booster_it.predict(Xy_it) + predt = booster.predict(Xy) + + np.testing.assert_allclose(predt, predt_it) + + +def check_extmem_qdm( # pylint: disable=too-many-arguments + n_samples_per_batch: int, + n_features: int, + *, + n_batches: int, + n_bins: int, + device: str, + on_host: bool, + is_cat: bool, +) -> None: + """Basic test for the `ExtMemQuantileDMatrix`.""" + + if is_cat: + it: DataIter = CatIter( + n_samples_per_batch=n_samples_per_batch, + n_features=n_features, + n_batches=n_batches, + n_cats=5, + sparsity=0.0, + cat_ratio=0.5, + onehot=False, + device=device, + cache="cache", + ) + else: + it = IteratorForTest( + *make_batches( + n_samples_per_batch, n_features, n_batches, use_cupy=device != "cpu" + ), + cache="cache", + on_host=on_host, + ) + + Xy_it = ExtMemQuantileDMatrix(it, max_bin=n_bins, enable_categorical=is_cat) + with pytest.raises(ValueError, match="Only the `hist`"): + booster_it = train( + {"device": device, "tree_method": "approx", "max_bin": n_bins}, + Xy_it, + num_boost_round=8, + ) + + booster_it = train({"device": device, "max_bin": n_bins}, Xy_it, num_boost_round=8) + if is_cat: + it = CatIter( + n_samples_per_batch=n_samples_per_batch, + n_features=n_features, + n_batches=n_batches, + n_cats=5, + sparsity=0.0, + cat_ratio=0.5, + onehot=False, + device=device, + cache=None, + ) + else: + it = IteratorForTest( + *make_batches( + n_samples_per_batch, n_features, n_batches, use_cupy=device != "cpu" + ), + cache=None, + ) + Xy = QuantileDMatrix(it, max_bin=n_bins, enable_categorical=is_cat) + booster = train({"device": device, "max_bin": n_bins}, Xy, num_boost_round=8) + + cut_it = Xy_it.get_quantile_cut() + cut = Xy.get_quantile_cut() + np.testing.assert_allclose(cut_it[0], cut[0]) + np.testing.assert_allclose(cut_it[1], cut[1]) + + predt_it = booster_it.predict(Xy_it) + predt = booster.predict(Xy) + np.testing.assert_allclose(predt_it, predt) + + +def check_cut( + n_entries: int, indptr: np.ndarray, data: np.ndarray, dtypes: Any +) -> None: + """Check the cut values.""" + assert data.shape[0] == indptr[-1] + assert data.shape[0] == n_entries + + assert indptr.dtype == np.uint64 + for i in range(1, indptr.size): + beg = int(indptr[i - 1]) + end = int(indptr[i]) + for j in range(beg + 1, end): + assert data[j] > data[j - 1] + if is_pd_cat_dtype(dtypes.iloc[i - 1]): + assert data[j] == data[j - 1] + 1 + + +def check_get_quantile_cut_device(tree_method: str, use_cupy: bool) -> None: + """Check with optional cupy.""" + import pandas as pd + + n_samples = 1024 + n_features = 14 + max_bin = 16 + dtypes = pd.Series([np.float32] * n_features) + + # numerical + X, y, w = tm.make_regression(n_samples, n_features, use_cupy=use_cupy) + # - qdm + Xyw: DMatrix = QuantileDMatrix(X, y, weight=w, max_bin=max_bin) + indptr, data = Xyw.get_quantile_cut() + check_cut((max_bin + 1) * n_features, indptr, data, dtypes) + # - dm + Xyw = DMatrix(X, y, weight=w) + train({"tree_method": tree_method, "max_bin": max_bin}, Xyw) + indptr, data = Xyw.get_quantile_cut() + check_cut((max_bin + 1) * n_features, indptr, data, dtypes) + # - ext mem + n_batches = 3 + n_samples_per_batch = 256 + it = IteratorForTest( + *make_batches(n_samples_per_batch, n_features, n_batches, use_cupy), + cache="cache", + on_host=False, + ) + Xy: DMatrix = DMatrix(it) + train({"tree_method": tree_method, "max_bin": max_bin}, Xyw) + indptr, data = Xyw.get_quantile_cut() + check_cut((max_bin + 1) * n_features, indptr, data, dtypes) + + # categorical + n_categories = 32 + X, y = make_categorical( + n_samples, n_features, n_categories, onehot=False, sparsity=0.8 + ) + if use_cupy: + import cudf + + cp = import_cupy() + + X = cudf.from_pandas(X) + y = cp.array(y) + # - qdm + Xy = QuantileDMatrix(X, y, max_bin=max_bin, enable_categorical=True) + indptr, data = Xy.get_quantile_cut() + check_cut(n_categories * n_features, indptr, data, X.dtypes) + # - dm + Xy = DMatrix(X, y, enable_categorical=True) + train({"tree_method": tree_method, "max_bin": max_bin}, Xy) + indptr, data = Xy.get_quantile_cut() + check_cut(n_categories * n_features, indptr, data, X.dtypes) + + # mixed + X, y = make_categorical( + n_samples, n_features, n_categories, onehot=False, sparsity=0.8, cat_ratio=0.5 + ) + n_cat_features = len([0 for dtype in X.dtypes if is_pd_cat_dtype(dtype)]) + n_num_features = n_features - n_cat_features + n_entries = n_categories * n_cat_features + (max_bin + 1) * n_num_features + # - qdm + Xy = QuantileDMatrix(X, y, max_bin=max_bin, enable_categorical=True) + indptr, data = Xy.get_quantile_cut() + check_cut(n_entries, indptr, data, X.dtypes) + # - dm + Xy = DMatrix(X, y, enable_categorical=True) + train({"tree_method": tree_method, "max_bin": max_bin}, Xy) + indptr, data = Xy.get_quantile_cut() + check_cut(n_entries, indptr, data, X.dtypes) + + +def check_get_quantile_cut(tree_method: str, device: str) -> None: + """Check the quantile cut getter.""" + + use_cupy = device.startswith("cuda") + check_get_quantile_cut_device(tree_method, False) + if use_cupy: + check_get_quantile_cut_device(tree_method, True) + + +USE_ONEHOT = np.iinfo(np.int32).max +USE_PART = 1 + + +def _create_dmatrix( # pylint: disable=too-many-arguments + n_samples: int, + n_features: int, + *, + n_cats: int, + device: str, + sparsity: float, + tree_method: str, + onehot: bool, + extmem: bool, + enable_categorical: bool, + n_targets: int = 1, + max_bin: Optional[int] = None, +) -> DMatrix: + n_batches = max(min(2, n_samples), 1) + it = CatIter( + n_samples // n_batches, + n_features, + n_batches=n_batches, + sparsity=sparsity, + cat_ratio=1.0, + n_cats=n_cats, + onehot=onehot, + device=device, + cache="cache" if extmem else None, + n_targets=n_targets, + ) + if extmem: + if tree_method == "hist": + Xy: DMatrix = ExtMemQuantileDMatrix( + it, enable_categorical=enable_categorical, max_bin=max_bin + ) + elif tree_method == "approx": + Xy = DMatrix(it, enable_categorical=enable_categorical) + else: + raise ValueError(f"tree_method {tree_method} not supported.") + else: + cat, label = it.xy() + Xy = DMatrix(cat, label, enable_categorical=enable_categorical) + return Xy + + +def check_categorical_ohe( # pylint: disable=too-many-arguments + *, + rows: int, + cols: int, + rounds: int, + cats: int, + device: Device, + tree_method: str, + extmem: bool = False, + multi_target: bool = False, + max_bin: Optional[int] = None, +) -> None: + "Test for one-hot encoding with categorical data." + by_etl_results: Dict[str, Dict[str, List[float]]] = {} + by_builtin_results: Dict[str, Dict[str, List[float]]] = {} + + parameters: Dict[str, Any] = { + "tree_method": tree_method, + # Use one-hot exclusively + "max_cat_to_onehot": USE_ONEHOT, + "reg_lambda": 0, + "device": device, + } + if max_bin is not None: + parameters["max_bin"] = max_bin + + n_targets = 3 if multi_target else 1 + if multi_target: + parameters["multi_strategy"] = "multi_output_tree" + + Xy_onehot = _create_dmatrix( + rows, + cols, + n_cats=cats, + device=device, + sparsity=0.0, + onehot=True, + tree_method=tree_method, + extmem=extmem, + enable_categorical=False, + n_targets=n_targets, + max_bin=max_bin, + ) + Xy_cat = _create_dmatrix( + rows, + cols, + n_cats=cats, + device=device, + sparsity=0.0, + tree_method=tree_method, + onehot=False, + extmem=extmem, + enable_categorical=True, + n_targets=n_targets, + max_bin=max_bin, + ) + + train( + parameters, + Xy_onehot, + num_boost_round=rounds, + evals=[(Xy_onehot, "Train")], + evals_result=by_etl_results, + ) + booster_onehot = train( + parameters, + Xy_cat, + num_boost_round=rounds, + evals=[(Xy_cat, "Train")], + evals_result=by_builtin_results, + ) + + # There are guidelines on how to specify tolerance based on considering output + # as random variables. But in here the tree construction is extremely sensitive + # to floating point errors. An 1e-5 error in a histogram bin can lead to an + # entirely different tree. So even though the test is quite lenient, hypothesis + # can still pick up falsifying examples from time to time. + np.testing.assert_allclose( + np.array(by_etl_results["Train"]["rmse"]), + np.array(by_builtin_results["Train"]["rmse"]), + rtol=1e-3, + ) + assert non_increasing(by_builtin_results["Train"]["rmse"]) + + by_grouping: Dict[str, Dict[str, List[float]]] = {} + # switch to partition-based splits + parameters["max_cat_to_onehot"] = USE_PART + booster_partition = train( + parameters, + Xy_cat, + num_boost_round=rounds, + evals=[(Xy_cat, "Train")], + evals_result=by_grouping, + ) + rmse_group = by_grouping["Train"]["rmse"] + assert non_increasing(rmse_group) + + model_onehot = json.loads(booster_onehot.save_raw(raw_format="json")) + model_partition = json.loads(booster_partition.save_raw(raw_format="json")) + tree_onehot = model_onehot["learner"]["gradient_booster"]["model"]["trees"][0] + tree_partition = model_partition["learner"]["gradient_booster"]["model"]["trees"][0] + assert tree_onehot["categories_sizes"] + assert tree_partition["categories_sizes"] + assert all(size == 1 for size in tree_onehot["categories_sizes"]) + + if not multi_target: + # Scalar partition sorting finds the best split for a fixed node. Compare the + # common root instead of complete greedy trees, whose later paths can diverge. + gain_onehot = tree_onehot["loss_changes"][0] + gain_partition = tree_partition["loss_changes"][0] + assert gain_partition >= gain_onehot or np.isclose(gain_partition, gain_onehot) + + parameters["reg_lambda"] = 1.0 + by_grouping = {} + train( + parameters, + Xy_cat, + num_boost_round=32, + evals=[(Xy_cat, "Train")], + evals_result=by_grouping, + ) + assert non_increasing(by_grouping["Train"]["rmse"]), by_grouping + + +def check_categorical_bitfield_boundaries( + device: Device, cats: int, multi_target: bool +) -> None: + """Test scalar and vector categorical splits at bit-field word boundaries.""" + n_targets = 3 if multi_target else 1 + X, y = tm.make_categorical( + 1000, 2, cats, onehot=False, n_targets=n_targets, sparsity=0.0 + ) + Xy = QuantileDMatrix(X, y, enable_categorical=True) + params: Dict[str, Any] = {"device": device, "tree_method": "hist"} + if multi_target: + params["multi_strategy"] = "multi_output_tree" + + for max_cat_to_onehot in [1, 128]: + params["max_cat_to_onehot"] = max_cat_to_onehot + booster = train(params, Xy, num_boost_round=1) + + assert booster.get_score(importance_type="weight") + predt = booster.predict(Xy) + assert predt.shape == y.shape + assert np.isfinite(predt).all() + + +def check_categorical_missing( # pylint: disable=too-many-arguments + rows: int, + cols: int, + cats: int, + *, + device: Device, + tree_method: str, + extmem: bool, +) -> None: + """Check categorical data with missing values.""" + parameters: Dict[str, Any] = {"tree_method": tree_method, "device": device} + Xy = _create_dmatrix( + rows, + cols, + n_cats=cats, + sparsity=0.5, + device=device, + tree_method=tree_method, + onehot=False, + extmem=extmem, + enable_categorical=True, + ) + label = Xy.get_label() + + def run(max_cat_to_onehot: int) -> None: + # Test with onehot splits + parameters["max_cat_to_onehot"] = max_cat_to_onehot + + evals_result: Dict[str, Dict] = {} + booster = train( + parameters, + Xy, + num_boost_round=8, + evals=[(Xy, "Train")], + evals_result=evals_result, + ) + assert non_increasing(evals_result["Train"]["rmse"]) + y_predt = booster.predict(Xy) + rmse = tm.root_mean_square(label, y_predt) + assert_allclose(device, rmse, evals_result["Train"]["rmse"][-1], rtol=2e-5) + + # Test with OHE split + run(USE_ONEHOT) + + # Test with partition-based split + run(USE_PART) + + +def run_max_cat(tree_method: str, device: Device) -> None: + """Test data with size smaller than number of categories.""" + import pandas as pd + + rng = np.random.default_rng(0) + n_cat = 100 + n = 5 + + X = pd.Series( + ["".join(rng.choice(list(ascii_lowercase), size=3)) for i in range(n_cat)], + dtype="category", + )[:n].to_frame() + + reg = XGBRegressor( + enable_categorical=True, + tree_method=tree_method, + device=device, + n_estimators=10, + ) + y = pd.Series(range(n)) + reg.fit(X=X, y=y, eval_set=[(X, y)]) + assert non_increasing(reg.evals_result()["validation_0"]["rmse"]) + + +def run_invalid_category(tree_method: str, device: Device) -> None: + """Test with invalid categorical inputs.""" + rng = np.random.default_rng() + # too large + X = rng.integers(low=0, high=4, size=1000).reshape(100, 10) + y = rng.normal(loc=0, scale=1, size=100) + X[13, 7] = np.iinfo(np.int32).max + 1 + + # Check is performed during sketching. + Xy = DMatrix(X, y, feature_types=["c"] * 10) + with pytest.raises(ValueError): + train({"tree_method": tree_method, "device": device}, Xy) + + X[13, 7] = 16777216 + Xy = DMatrix(X, y, feature_types=["c"] * 10) + with pytest.raises(ValueError): + train({"tree_method": tree_method, "device": device}, Xy) + + # mixed positive and negative values + X = rng.normal(loc=0, scale=1, size=1000).reshape(100, 10) # type: ignore[assignment] + y = rng.normal(loc=0, scale=1, size=100) + + Xy = DMatrix(X, y, feature_types=["c"] * 10) + with pytest.raises(ValueError): + train({"tree_method": tree_method, "device": device}, Xy) + + if device == "cuda": + import cupy as cp + + X, y = cp.array(X), cp.array(y) + with pytest.raises(ValueError): + QuantileDMatrix(X, y, feature_types=["c"] * 10) + + +def train_result( + param: Dict[str, Any], dmat: DMatrix, num_rounds: int +) -> Dict[str, Any]: + """Get training result from parameters and data.""" + result: Dict[str, Any] = {} + booster = train( + param, + dmat, + num_rounds, + evals=[(dmat, "train")], + verbose_eval=False, + evals_result=result, + ) + assert booster.num_features() == dmat.num_col() + assert booster.num_boosted_rounds() == num_rounds + assert booster.feature_names == dmat.feature_names + assert booster.feature_types == dmat.feature_types + + return result + + +class ResetStrategy(TrainingCallback): + """Callback for testing multi-output.""" + + def after_iteration(self, model: Booster, epoch: int, evals_log: dict) -> bool: + if epoch % 2 == 0: + model.set_param({"multi_strategy": "multi_output_tree"}) + else: + model.set_param({"multi_strategy": "one_output_per_tree"}) + return False diff --git a/v12x/pylib/xgboost/testing/utils.py b/v12x/pylib/xgboost/testing/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..4503cfd060cd8861c843aea999a96d1269e1137b --- /dev/null +++ b/v12x/pylib/xgboost/testing/utils.py @@ -0,0 +1,44 @@ +"""Helpers for test code.""" + +from typing import Any, Literal, Sequence, TypeAlias + +import numpy as np + +from ..compat import _is_cupy_alike, import_cupy +from ..core import DMatrix + +Device: TypeAlias = Literal["cpu", "cuda"] + + +def assert_allclose( + device: Device, a: Any, b: Any, *, rtol: float = 1e-7, atol: float = 0 +) -> None: + """Dispatch the assert_allclose for devices.""" + if device == "cpu" and not _is_cupy_alike(a) and not _is_cupy_alike(b): + np.testing.assert_allclose(a, b, atol=atol, rtol=rtol) + else: + cp = import_cupy() + cp.testing.assert_allclose(a, b, atol=atol, rtol=rtol) + + +def predictor_equal(lhs: DMatrix, rhs: DMatrix) -> bool: + """Assert whether two DMatrices contain the same predictors.""" + lcsr = lhs.get_data() + rcsr = rhs.get_data() + return all( + ( + np.array_equal(lcsr.data, rcsr.data), + np.array_equal(lcsr.indices, rcsr.indices), + np.array_equal(lcsr.indptr, rcsr.indptr), + ) + ) + + +def non_increasing(L: Sequence[float], tolerance: float = 1e-4) -> bool: + """Values in the sequence are not increasing.""" + return all((y - x) < tolerance for x, y in zip(L, L[1:])) + + +def non_decreasing(L: Sequence[float], tolerance: float = 1e-4) -> bool: + """Values in the sequence are not decreasing.""" + return all((y - x) >= -tolerance for x, y in zip(L, L[1:])) diff --git a/v12x/pylib/xgboost/testing/with_skl.py b/v12x/pylib/xgboost/testing/with_skl.py new file mode 100644 index 0000000000000000000000000000000000000000..f892a4db950528fe2e7a4c58251612588a3a488b --- /dev/null +++ b/v12x/pylib/xgboost/testing/with_skl.py @@ -0,0 +1,234 @@ +# pylint: disable=too-many-arguments, too-many-positional-arguments +"""Tests for compatiblity with sklearn.""" + +from typing import Callable, Optional, Type + +import numpy as np +import pytest +from scipy.special import softmax + +from ..core import DMatrix +from ..sklearn import XGBClassifier, XGBRegressor, XGBRFRegressor +from .data import get_california_housing, make_batches +from .ordinal import make_recoded +from .utils import Device, assert_allclose + + +def run_boost_from_prediction_binary( + tree_method: str, + device: Device, + X: np.ndarray, + y: np.ndarray, + as_frame: Optional[Callable], +) -> None: + """ + Parameters + ---------- + + as_frame: A callable function to convert margin into DataFrame, useful for different + df implementations. + """ + + model_0 = XGBClassifier( + learning_rate=0.3, + random_state=0, + n_estimators=4, + tree_method=tree_method, + device=device, + ) + model_0.fit(X=X, y=y) + margin = model_0.predict(X, output_margin=True) + if as_frame is not None: + margin = as_frame(margin) + + model_1 = XGBClassifier( + learning_rate=0.3, + random_state=0, + n_estimators=4, + tree_method=tree_method, + device=device, + ) + model_1.fit(X=X, y=y, base_margin=margin) + predictions_1 = model_1.predict(X, base_margin=margin) + + cls_2 = XGBClassifier( + learning_rate=0.3, + random_state=0, + n_estimators=8, + tree_method=tree_method, + device=device, + ) + cls_2.fit(X=X, y=y) + predictions_2 = cls_2.predict(X) + np.testing.assert_allclose(predictions_1, predictions_2) + + +def run_boost_from_prediction_multi_clasas( + estimator: Type, + tree_method: str, + device: Device, + X: np.ndarray, + y: np.ndarray, + as_frame: Optional[Callable], +) -> None: + """Boosting from prediction with multi-class clf.""" + # Multi-class + model_0 = estimator( + learning_rate=0.3, + random_state=0, + n_estimators=4, + tree_method=tree_method, + device=device, + ) + model_0.fit(X=X, y=y) + margin = model_0.get_booster().inplace_predict(X, predict_type="margin") + if as_frame is not None: + margin = as_frame(margin) + + model_1 = estimator( + learning_rate=0.3, + random_state=0, + n_estimators=4, + tree_method=tree_method, + device=device, + ) + model_1.fit(X=X, y=y, base_margin=margin) + predictions_1 = model_1.get_booster().predict( + DMatrix(X, base_margin=margin), output_margin=True + ) + + model_2 = estimator( + learning_rate=0.3, + random_state=0, + n_estimators=8, + tree_method=tree_method, + device=device, + ) + model_2.fit(X=X, y=y) + predictions_2 = model_2.get_booster().inplace_predict(X, predict_type="margin") + + if hasattr(predictions_1, "get"): + predictions_1 = predictions_1.get() + if hasattr(predictions_2, "get"): + predictions_2 = predictions_2.get() + np.testing.assert_allclose(predictions_1, predictions_2, atol=1e-6) + + +def run_housing_rf_regression(tree_method: str, device: Device) -> None: + """Testwith the cali housing dataset.""" + from sklearn.metrics import mean_squared_error + from sklearn.model_selection import KFold + + X, y = get_california_housing() + rng = np.random.RandomState(1994) + kf = KFold(n_splits=2, shuffle=True, random_state=rng) + for train_index, test_index in kf.split(X, y): + xgb_model = XGBRFRegressor( + random_state=42, tree_method=tree_method, device=device + ).fit(X[train_index], y[train_index]) + preds = xgb_model.predict(X[test_index]) + labels = y[test_index] + assert mean_squared_error(preds, labels) < 35 + + rfreg = XGBRFRegressor(device=device) + with pytest.raises(NotImplementedError): + rfreg.set_params(early_stopping_rounds=10) + rfreg.fit(X, y) + + +def run_recoding(device: Device) -> None: + """Test re-coding for training continuation.""" + enc, reenc, y, _, _ = make_recoded(device, n_features=16) + reg = XGBRegressor(enable_categorical=True, n_estimators=2, device=device) + reg.fit(enc, y, eval_set=[(reenc, y)]) + results_0 = reg.evals_result() + + booster = reg.get_booster() + assert not booster.get_categories().empty() + + reg = XGBRegressor(enable_categorical=True, n_estimators=2, device=device) + reg.fit(reenc, y, xgb_model=booster, eval_set=[(enc, y)]) + results_1 = reg.evals_result() + + booster = reg.get_booster() + assert booster.num_boosted_rounds() == 4 + assert not booster.get_categories().empty() + + reg = XGBRegressor(enable_categorical=True, n_estimators=4, device=device) + reg.fit(enc, y, eval_set=[(reenc, y)]) + results_2 = reg.evals_result() + + np.testing.assert_allclose( + results_2["validation_0"]["rmse"], + results_0["validation_0"]["rmse"] + results_1["validation_0"]["rmse"], + ) + + np.testing.assert_allclose(reg.predict(reenc), reg.predict(enc)) + np.testing.assert_allclose(reg.apply(reenc), reg.apply(enc)) + + +def run_intercept(device: Device) -> None: + """Tests for the intercept.""" + from sklearn.datasets import make_classification, make_multilabel_classification + + X, y, w = [v[0] for v in make_batches(256, 3, 1, use_cupy=False)] + reg = XGBRegressor(device=device) + reg.fit(X, y, sample_weight=w) + result = reg.intercept_ + assert result.dtype == np.float32 + assert result[0] < 0.5 + + reg = XGBRegressor(booster="gblinear", device=device) + reg.fit(X, y, sample_weight=w) + result = reg.intercept_ + assert isinstance(result, np.ndarray) + assert result.dtype == np.float32 + assert result[0] < 0.5 + + n_classes = 4 + X, y = make_classification( + random_state=1994, + n_samples=128, + n_features=16, + n_classes=n_classes, + n_informative=16, + n_redundant=0, + ) + + clf = XGBClassifier(booster="gbtree", objective="multi:softprob", device=device) + clf.fit(X, y) + result = clf.intercept_ + assert isinstance(result, np.ndarray) + assert len(result) == 4 + + assert (softmax(result) >= 0.0).all() + np.testing.assert_allclose(sum(result), 0.0, atol=1e-6) + np.testing.assert_allclose(sum(softmax(result)), 1.0) + + # Tests for user input + # Multi-class + intercept = np.ones(shape=(n_classes), dtype=np.float32) / n_classes + if device == "cuda": + import cupy as cp + + intercept = cp.array(intercept) + + clf = XGBClassifier(objective="multi:softprob", base_score=intercept) + clf.fit(X, y) + assert_allclose(device, intercept, clf.intercept_) + + X, y = make_multilabel_classification( # pylint: disable=unbalanced-tuple-unpacking + random_state=1994, n_samples=128, n_features=16, n_classes=n_classes + ) + + # Multi-label + intercept = np.ones(shape=(n_classes), dtype=np.float32) / 2 + if device == "cuda": + import cupy as cp + + intercept = cp.array(intercept) + + clf = XGBClassifier(base_score=intercept) + clf.fit(X, y) + assert_allclose(device, intercept, clf.intercept_) + assert clf.objective == "binary:logistic" diff --git a/v12x/pylib/xgboost/tracker.py b/v12x/pylib/xgboost/tracker.py new file mode 100644 index 0000000000000000000000000000000000000000..926c957b28a41eb60ed4c61b992629e401b7ed79 --- /dev/null +++ b/v12x/pylib/xgboost/tracker.py @@ -0,0 +1,142 @@ +"""Tracker for XGBoost collective.""" + +import ctypes +import json +import socket +from enum import IntEnum, unique +from typing import Dict, Optional, Union + +from .core import _LIB, _check_call, _deprecate_positional_args, make_jcargs + + +def get_family(addr: str) -> int: + """Get network family from address.""" + return socket.getaddrinfo(addr, None)[0][0] + + +class RabitTracker: + """Tracker for the collective used in XGBoost, acting as a coordinator between + workers. + + Parameters + ---------- + + n_workers: + + The total number of workers in the communication group. + + host_ip: + + The IP address of the tracker node. XGBoost can try to guess one by probing with + sockets. But it's best to explicitly pass an address. + + port: + + The port this tracker should listen to. XGBoost can query an available port from + the OS, this configuration is useful for restricted network environments. + + sortby: + + How to sort the workers for rank assignment. The default is host, but users can + set the `DMLC_TASK_ID` via arguments of :py:meth:`~xgboost.collective.init` and + obtain deterministic rank assignment through sorting by task name. Available + options are: + + - host + - task + + timeout : + + Timeout for constructing (bootstrap) and shutting down the communication group, + doesn't apply to communication when the group is up and running. + + The timeout value should take the time of data loading and pre-processing into + account, due to potential lazy execution. By default the Tracker doesn't have + any timeout to avoid pre-mature aborting. + + The :py:meth:`.wait_for` method has a different timeout parameter that can stop + the tracker even if the tracker is still being used. A value error is raised + when timeout is reached. + + Examples + -------- + + .. code-block:: python + + from xgboost.tracker import RabitTracker + from xgboost import collective as coll + + tracker = RabitTracker(host_ip="127.0.0.1", n_workers=2) + tracker.start() + + with coll.CommunicatorContext(**tracker.worker_args()): + ret = coll.broadcast("msg", 0) + assert str(ret) == "msg" + + """ + + @unique + class _SortBy(IntEnum): + HOST = 0 + TASK = 1 + + @_deprecate_positional_args + def __init__( # pylint: disable=too-many-arguments + self, + n_workers: int, + host_ip: Optional[str], + port: int = 0, + *, + sortby: str = "host", + timeout: int = 0, + ) -> None: + + handle = ctypes.c_void_p() + if sortby not in ("host", "task"): + raise ValueError("Expecting either 'host' or 'task' for sortby.") + if host_ip is not None: + get_family(host_ip) # use python socket to stop early for invalid address + args = make_jcargs( + host=host_ip, + n_workers=n_workers, + port=port, + dmlc_communicator="rabit", + sortby=self._SortBy.HOST if sortby == "host" else self._SortBy.TASK, + timeout=int(timeout), + ) + _check_call(_LIB.XGTrackerCreate(args, ctypes.byref(handle))) + self.handle = handle + + def free(self) -> None: + """Internal function for testing.""" + if hasattr(self, "handle"): + handle = self.handle + del self.handle + _check_call(_LIB.XGTrackerFree(handle)) + + def __del__(self) -> None: + self.free() + + def start(self) -> None: + """Start the tracker. Once started, the client still need to call the + :py:meth:`wait_for` method in order to wait for it to finish (think of it as a + thread). + + """ + _check_call(_LIB.XGTrackerRun(self.handle, make_jcargs())) + + def wait_for(self, timeout: Optional[int] = None) -> None: + """Wait for the tracker to finish all the work and shutdown. When timeout is + reached, a value error is raised. By default we don't have timeout since we + don't know how long it takes for the model to finish training. + + """ + _check_call(_LIB.XGTrackerWaitFor(self.handle, make_jcargs(timeout=timeout))) + + def worker_args(self) -> Dict[str, Union[str, int]]: + """Get arguments for workers.""" + c_env = ctypes.c_char_p() + _check_call(_LIB.XGTrackerWorkerArgs(self.handle, ctypes.byref(c_env))) + assert c_env.value is not None + env = json.loads(c_env.value) + return env diff --git a/v12x/pylib/xgboost/training.py b/v12x/pylib/xgboost/training.py new file mode 100644 index 0000000000000000000000000000000000000000..d3088a9352de8610d4c0e89b383504a1ee0ea644 --- /dev/null +++ b/v12x/pylib/xgboost/training.py @@ -0,0 +1,623 @@ +# pylint: disable=too-many-locals, too-many-arguments +# pylint: disable=too-many-branches, too-many-statements +"""Training Library containing training routines.""" + +import copy +import os +import weakref +from typing import ( + TYPE_CHECKING, + Any, + Dict, + Iterable, + List, + Optional, + Sequence, + Tuple, + Union, + cast, +) + +import numpy as np + +from ._typing import BoosterParam, Callable, FPreProcCallable +from .callback import ( + CallbackContainer, + EarlyStopping, + EvaluationMonitor, + TrainingCallback, +) +from .compat import SKLEARN_INSTALLED, XGBStratifiedKFold +from .core import ( + Booster, + CustomObj, + DMatrix, + Metric, + XGBoostError, + _deprecate_positional_args, + _RefMixIn, +) + +if TYPE_CHECKING: + from pandas import DataFrame as PdDataFrame + +_CVFolds = Sequence["CVPack"] + +_RefError = ( + "Training dataset should be used as a reference when constructing the " + "`QuantileDMatrix` for evaluation.", +) + + +@_deprecate_positional_args +def train( + params: Dict[str, Any], + dtrain: DMatrix, + num_boost_round: int = 10, + *, + evals: Optional[Sequence[Tuple[DMatrix, str]]] = None, + obj: Optional[CustomObj] = None, + maximize: Optional[bool] = None, + early_stopping_rounds: Optional[int] = None, + evals_result: Optional[TrainingCallback.EvalsLog] = None, + verbose_eval: Optional[Union[bool, int]] = True, + xgb_model: Optional[Union[str, os.PathLike, Booster, bytearray]] = None, + callbacks: Optional[Sequence[TrainingCallback]] = None, + custom_metric: Optional[Metric] = None, +) -> Booster: + """Train a booster with given parameters. + + Parameters + ---------- + params : + Booster params. + dtrain : + Data to be trained. + num_boost_round : + Number of boosting iterations. + evals : + List of validation sets for which metrics will evaluated during training. + Validation metrics will help us track the performance of the model. + obj + Custom objective function. See :doc:`Custom Objective + ` for details. + maximize : + Whether to maximize custom_metric. + + early_stopping_rounds : + + Activates early stopping. Validation metric needs to improve at least once in + every **early_stopping_rounds** round(s) to continue training. + + Requires at least one item in **evals**. + + The method returns the model from the last iteration (not the best one). Use + custom callback :py:class:`~xgboost.callback.EarlyStopping` or :py:meth:`model + slicing ` if the best model is desired. If there's + more than one item in **evals**, the last entry will be used for early stopping. + + If there's more than one metric in the **eval_metric** parameter given in + **params**, the last metric will be used for early stopping. + + If early stopping occurs, the model will have two additional fields: + ``bst.best_score``, ``bst.best_iteration``. + + evals_result : + This dictionary stores the evaluation results of all the items in watchlist. + + Example: with a watchlist containing + ``[(dtest,'eval'), (dtrain,'train')]`` and + a parameter containing ``('eval_metric': 'logloss')``, + the **evals_result** returns + + .. code-block:: python + + {'train': {'logloss': ['0.48253', '0.35953']}, + 'eval': {'logloss': ['0.480385', '0.357756']}} + + verbose_eval : + Requires at least one item in **evals**. + + If **verbose_eval** is True then the evaluation metric on the validation set is + printed at each boosting stage. + + If **verbose_eval** is an integer then the evaluation metric on the validation + set is printed at every given **verbose_eval** boosting stage. The last boosting + stage / the boosting stage found by using **early_stopping_rounds** is also + printed. + + Example: with ``verbose_eval=4`` and at least one item in **evals**, an + evaluation metric is printed every 4 boosting stages, instead of every boosting + stage. + + xgb_model : + Xgb model to be loaded before training (allows training continuation). + + callbacks : + List of callback functions that are applied at end of each iteration. + It is possible to use predefined callbacks by using + :ref:`Callback API `. + + .. note:: + + States in callback are not preserved during training, which means callback + objects can not be reused for multiple training sessions without + reinitialization or deepcopy. + + .. code-block:: python + + for params in parameters_grid: + # be sure to (re)initialize the callbacks before each run + callbacks = [xgb.callback.LearningRateScheduler(custom_rates)] + xgboost.train(params, Xy, callbacks=callbacks) + + custom_metric: + + .. versionadded 1.6.0 + + Custom metric function. See :doc:`Custom Metric ` + for details. The metric receives transformed prediction (after applying the + reverse link function) when using a builtin objective, and raw output when using + a custom objective. + + Returns + ------- + Booster : a trained booster model + + """ + + callbacks = [] if callbacks is None else copy.copy(list(callbacks)) + evals = list(evals) if evals else [] + + for va, _ in evals: + if not isinstance(va, DMatrix): + raise TypeError("Invalid type for the `evals`.") + + if ( + isinstance(va, _RefMixIn) + and va.ref is not weakref.ref(dtrain) + and va is not dtrain + ): + raise ValueError(_RefError) + + bst = Booster(params, [dtrain] + [d[0] for d in evals], model_file=xgb_model) + start_iteration = 0 + + if verbose_eval: + verbose_eval = 1 if verbose_eval is True else verbose_eval + callbacks.append(EvaluationMonitor(period=verbose_eval)) + if early_stopping_rounds: + callbacks.append(EarlyStopping(rounds=early_stopping_rounds, maximize=maximize)) + cb_container = CallbackContainer( + callbacks, metric=custom_metric, output_margin=callable(obj) + ) + + bst = cb_container.before_training(bst) + + for i in range(start_iteration, num_boost_round): + if cb_container.before_iteration(bst, i, dtrain, evals): + break + bst.update(dtrain, iteration=i, fobj=obj) + if cb_container.after_iteration(bst, i, dtrain, evals): + break + + bst = cb_container.after_training(bst) + + if evals_result is not None: + evals_result.update(cb_container.history) + + return bst.reset() + + +class CVPack: + """ "Auxiliary datastruct to hold one fold of CV.""" + + def __init__( + self, dtrain: DMatrix, dtest: DMatrix, param: Optional[Union[Dict, List]] + ) -> None: + """Initialize the CVPack.""" + self.dtrain = dtrain + self.dtest = dtest + self.watchlist = [(dtrain, "train"), (dtest, "test")] + self.bst = Booster(param, [dtrain, dtest]) + + def __getattr__(self, name: str) -> Callable: + def _inner(*args: Any, **kwargs: Any) -> Any: + return getattr(self.bst, name)(*args, **kwargs) + + return _inner + + def update(self, iteration: int, fobj: Optional[CustomObj]) -> None: + """ "Update the boosters for one iteration""" + self.bst.update(self.dtrain, iteration, fobj) + + def eval(self, iteration: int, feval: Optional[Metric], output_margin: bool) -> str: + """ "Evaluate the CVPack for one iteration.""" + return self.bst.eval_set(self.watchlist, iteration, feval, output_margin) + + +class _PackedBooster: + def __init__(self, cvfolds: _CVFolds) -> None: + self.cvfolds = cvfolds + + def update(self, iteration: int, obj: Optional[CustomObj]) -> None: + """Iterate through folds for update""" + for fold in self.cvfolds: + fold.update(iteration, obj) + + def eval( + self, iteration: int, feval: Optional[Metric], output_margin: bool + ) -> List[str]: + """Iterate through folds for eval""" + result = [f.eval(iteration, feval, output_margin) for f in self.cvfolds] + return result + + def set_attr(self, **kwargs: Optional[Any]) -> Any: + """Iterate through folds for setting attributes""" + for f in self.cvfolds: + f.bst.set_attr(**kwargs) + + def attr(self, key: str) -> Optional[str]: + """Redirect to booster attr.""" + return self.cvfolds[0].bst.attr(key) + + def set_param( + self, + params: Union[Dict, Iterable[Tuple[str, Any]], str], + value: Optional[str] = None, + ) -> None: + """Iterate through folds for set_param""" + for f in self.cvfolds: + f.bst.set_param(params, value) + + def num_boosted_rounds(self) -> int: + """Number of boosted rounds.""" + return self.cvfolds[0].num_boosted_rounds() + + @property + def best_iteration(self) -> int: + """Get best_iteration""" + return int(cast(int, self.cvfolds[0].bst.attr("best_iteration"))) + + @best_iteration.setter + def best_iteration(self, iteration: int) -> None: + """Set best_iteration""" + self.set_attr(best_iteration=iteration) + + @property + def best_score(self) -> float: + """Get best_score.""" + return float(cast(float, self.cvfolds[0].bst.attr("best_score"))) + + @best_score.setter + def best_score(self, score: float) -> None: + self.set_attr(best_score=score) + + +def groups_to_rows(groups: np.ndarray, boundaries: np.ndarray) -> np.ndarray: + """ + Given group row boundaries, convert ground indexes to row indexes + :param groups: list of groups for testing + :param boundaries: rows index limits of each group + :return: row in group + """ + return np.concatenate([np.arange(boundaries[g], boundaries[g + 1]) for g in groups]) + + +def mkgroupfold( + *, + dall: DMatrix, + nfold: int, + param: BoosterParam, + evals: Sequence[str] = (), + fpreproc: Optional[FPreProcCallable] = None, + shuffle: bool = True, +) -> List[CVPack]: + """ + Make n folds for cross-validation maintaining groups + :return: cross-validation folds + """ + # we have groups for pairwise ranking... get a list of the group indexes + group_boundaries = dall.get_uint_info("group_ptr") + group_sizes = np.diff(group_boundaries) + + if shuffle is True: + idx = np.random.permutation(len(group_sizes)) + else: + idx = np.arange(len(group_sizes)) + # list by fold of test group indexes + out_group_idset = np.array_split(idx, nfold) + # list by fold of train group indexes + in_group_idset = [ + np.concatenate([out_group_idset[i] for i in range(nfold) if k != i]) + for k in range(nfold) + ] + # from the group indexes, convert them to row indexes + in_idset = [ + groups_to_rows(in_groups, group_boundaries) for in_groups in in_group_idset + ] + out_idset = [ + groups_to_rows(out_groups, group_boundaries) for out_groups in out_group_idset + ] + + # build the folds by taking the appropriate slices + ret = [] + for k in range(nfold): + # perform the slicing using the indexes determined by the above methods + dtrain = dall.slice(in_idset[k], allow_groups=True) + dtrain.set_group(group_sizes[in_group_idset[k]]) + dtest = dall.slice(out_idset[k], allow_groups=True) + dtest.set_group(group_sizes[out_group_idset[k]]) + # run preprocessing on the data set if needed + if fpreproc is not None: + dtrain, dtest, tparam = fpreproc(dtrain, dtest, param.copy()) + else: + tparam = param + plst = list(tparam.items()) + [("eval_metric", itm) for itm in evals] + ret.append(CVPack(dtrain, dtest, plst)) + return ret + + +def mknfold( + *, + dall: DMatrix, + nfold: int, + param: BoosterParam, + seed: int, + evals: Sequence[str] = (), + fpreproc: Optional[FPreProcCallable] = None, + stratified: Optional[bool] = False, + folds: Optional[XGBStratifiedKFold] = None, + shuffle: bool = True, +) -> List[CVPack]: + """ + Make an n-fold list of CVPack from random indices. + """ + evals = list(evals) + np.random.seed(seed) + + if stratified is False and folds is None: + # Do standard k-fold cross validation. Automatically determine the folds. + if len(dall.get_uint_info("group_ptr")) > 1: + return mkgroupfold( + dall=dall, + nfold=nfold, + param=param, + evals=evals, + fpreproc=fpreproc, + shuffle=shuffle, + ) + + if shuffle is True: + idx = np.random.permutation(dall.num_row()) + else: + idx = np.arange(dall.num_row()) + out_idset = np.array_split(idx, nfold) + in_idset = [ + np.concatenate([out_idset[i] for i in range(nfold) if k != i]) + for k in range(nfold) + ] + elif folds is not None: + # Use user specified custom split using indices + try: + in_idset = [x[0] for x in folds] + out_idset = [x[1] for x in folds] + except TypeError: + # Custom stratification using Sklearn KFoldSplit object + splits = list(folds.split(X=dall.get_label(), y=dall.get_label())) + in_idset = [x[0] for x in splits] + out_idset = [x[1] for x in splits] + nfold = len(out_idset) + else: + # Do standard stratefied shuffle k-fold split + sfk = XGBStratifiedKFold(n_splits=nfold, shuffle=True, random_state=seed) + splits = list(sfk.split(X=dall.get_label(), y=dall.get_label())) + in_idset = [x[0] for x in splits] + out_idset = [x[1] for x in splits] + nfold = len(out_idset) + + ret = [] + for k in range(nfold): + # perform the slicing using the indexes determined by the above methods + dtrain = dall.slice(in_idset[k]) + dtest = dall.slice(out_idset[k]) + # run preprocessing on the data set if needed + if fpreproc is not None: + dtrain, dtest, tparam = fpreproc(dtrain, dtest, param.copy()) + else: + tparam = param + plst = list(tparam.items()) + [("eval_metric", itm) for itm in evals] + ret.append(CVPack(dtrain, dtest, plst)) + return ret + + +@_deprecate_positional_args +def cv( + params: BoosterParam, + dtrain: DMatrix, + num_boost_round: int = 10, + *, + nfold: int = 3, + stratified: bool = False, + folds: Optional[XGBStratifiedKFold] = None, + metrics: Sequence[str] = (), + obj: Optional[CustomObj] = None, + maximize: Optional[bool] = None, + early_stopping_rounds: Optional[int] = None, + fpreproc: Optional[FPreProcCallable] = None, + as_pandas: bool = True, + verbose_eval: Optional[Union[int, bool]] = None, + show_stdv: bool = True, + seed: int = 0, + callbacks: Optional[Sequence[TrainingCallback]] = None, + shuffle: bool = True, + custom_metric: Optional[Metric] = None, +) -> Union[Dict[str, float], "PdDataFrame"]: + """Cross-validation with given parameters. + + Parameters + ---------- + params : dict + Booster params. + dtrain : + Data to be trained. Only the :py:class:`DMatrix` without external memory is + supported. + num_boost_round : + Number of boosting iterations. + nfold : int + Number of folds in CV. + stratified : bool + Perform stratified sampling. + folds : a KFold or StratifiedKFold instance or list of fold indices + Sklearn KFolds or StratifiedKFolds object. + Alternatively may explicitly pass sample indices for each fold. + For ``n`` folds, **folds** should be a length ``n`` list of tuples. + Each tuple is ``(in,out)`` where ``in`` is a list of indices to be used + as the training samples for the ``n`` th fold and ``out`` is a list of + indices to be used as the testing samples for the ``n`` th fold. + metrics : string or list of strings + Evaluation metrics to be watched in CV. + obj : + + Custom objective function. See :doc:`Custom Objective + ` for details. + + maximize : bool + Whether to maximize the evaluataion metric (score or error). + + early_stopping_rounds: int + Activates early stopping. Cross-Validation metric (average of validation + metric computed over CV folds) needs to improve at least once in + every **early_stopping_rounds** round(s) to continue training. + The last entry in the evaluation history will represent the best iteration. + If there's more than one metric in the **eval_metric** parameter given in + **params**, the last metric will be used for early stopping. + fpreproc : function + Preprocessing function that takes (dtrain, dtest, param) and returns + transformed versions of those. + as_pandas : bool, default True + Return pd.DataFrame when pandas is installed. + If False or pandas is not installed, return np.ndarray + verbose_eval : bool, int, or None, default None + Whether to display the progress. If None, progress will be displayed + when np.ndarray is returned. If True, progress will be displayed at + boosting stage. If an integer is given, progress will be displayed + at every given `verbose_eval` boosting stage. + show_stdv : bool, default True + Whether to display the standard deviation in progress. + Results are not affected, and always contains std. + seed : int + Seed used to generate the folds (passed to numpy.random.seed). + callbacks : + List of callback functions that are applied at end of each iteration. + It is possible to use predefined callbacks by using + :ref:`Callback API `. + + .. note:: + + States in callback are not preserved during training, which means callback + objects can not be reused for multiple training sessions without + reinitialization or deepcopy. + + .. code-block:: python + + for params in parameters_grid: + # be sure to (re)initialize the callbacks before each run + callbacks = [xgb.callback.LearningRateScheduler(custom_rates)] + xgboost.train(params, Xy, callbacks=callbacks) + + shuffle : bool + Shuffle data before creating folds. + custom_metric : + + .. versionadded 1.6.0 + + Custom metric function. See :doc:`Custom Metric ` + for details. + + Returns + ------- + evaluation history : list(string) + """ + if stratified is True and not SKLEARN_INSTALLED: + raise XGBoostError( + "sklearn needs to be installed in order to use stratified cv" + ) + if isinstance(metrics, str): + metrics = [metrics] + if isinstance(dtrain, _RefMixIn): + raise ValueError("`QuantileDMatrix` is not yet supported.") + + params = params.copy() + if isinstance(params, list): + _metrics = [x[1] for x in params if x[0] == "eval_metric"] + params = dict(params) + if "eval_metric" in params: + params["eval_metric"] = _metrics + + if (not metrics) and "eval_metric" in params: + if isinstance(params["eval_metric"], list): + metrics = params["eval_metric"] + else: + metrics = [params["eval_metric"]] + + params.pop("eval_metric", None) + + results: Dict[str, List[float]] = {} + cvfolds = mknfold( + dall=dtrain, + nfold=nfold, + param=params, + seed=seed, + evals=metrics, + fpreproc=fpreproc, + stratified=stratified, + folds=folds, + shuffle=shuffle, + ) + + # setup callbacks + callbacks = [] if callbacks is None else copy.copy(list(callbacks)) + + if verbose_eval: + verbose_eval = 1 if verbose_eval is True else verbose_eval + callbacks.append(EvaluationMonitor(period=verbose_eval, show_stdv=show_stdv)) + if early_stopping_rounds: + callbacks.append(EarlyStopping(rounds=early_stopping_rounds, maximize=maximize)) + callbacks_container = CallbackContainer( + callbacks, metric=custom_metric, is_cv=True, output_margin=callable(obj) + ) + + booster = _PackedBooster(cvfolds) + callbacks_container.before_training(booster) + + for i in range(num_boost_round): + if callbacks_container.before_iteration(booster, i, dtrain, None): + break + booster.update(i, obj) + + should_break = callbacks_container.after_iteration(booster, i, dtrain, None) + res = callbacks_container.aggregated_cv + for key, mean, std in cast(List[Tuple[str, float, float]], res): + if key + "-mean" not in results: + results[key + "-mean"] = [] + if key + "-std" not in results: + results[key + "-std"] = [] + results[key + "-mean"].append(mean) + results[key + "-std"].append(std) + + if should_break: + for k in results.keys(): # pylint: disable=consider-iterating-dictionary + results[k] = results[k][: (booster.best_iteration + 1)] + break + if as_pandas: + try: + import pandas as pd + + results = pd.DataFrame.from_dict(results) + except ImportError: + pass + + callbacks_container.after_training(booster) + + return results diff --git a/v12x/v8regen/work/ens/ens_log.jsonl b/v12x/v8regen/work/ens/ens_log.jsonl new file mode 100644 index 0000000000000000000000000000000000000000..21d79e7b821f44d1b2bf3b3118a6d7ad2d7243bf --- /dev/null +++ b/v12x/v8regen/work/ens/ens_log.jsonl @@ -0,0 +1,2 @@ +{"cut": 27, "tag": "base", "seed": 0, "n": 745540, "nf": 377, "min": 0.9, "load_s": 10, "ds_s": 5, "fit_s": 33, "pred_s": 3} +{"cut": 27, "tag": "slow", "seed": 0, "n": 745540, "nf": 377, "min": 1.0, "load_s": 5, "ds_s": 4, "fit_s": 42, "pred_s": 7} diff --git a/v13/ens/ens_log.jsonl b/v13/ens/ens_log.jsonl new file mode 100644 index 0000000000000000000000000000000000000000..4c8d2057a65aa3c543de003fc6c7c6e2ed3ad57c --- /dev/null +++ b/v13/ens/ens_log.jsonl @@ -0,0 +1,4 @@ +{"cut": 25, "tag": "ev_s", "seed": 0, "n": 745456, "nf": 835, "min": 2.8, "load_s": 14, "ds_s": 10, "fit_s": 137, "pred_s": 4, "auc": {"400": 0.81594, "500": 0.81593}} +{"cut": 23, "tag": "ev_s", "seed": 0, "n": 373643, "nf": 835, "min": 0.9, "load_s": 6, "ds_s": 6, "fit_s": 35, "pred_s": 5, "auc": {"400": 0.79431, "500": 0.79418}} +{"cut": 25, "tag": "aw_s", "seed": 0, "n": 745456, "nf": 435, "min": 0.9, "load_s": 6, "ds_s": 4, "fit_s": 36, "pred_s": 4, "auc": {"400": 0.81601, "500": 0.81589}} +{"cut": 23, "tag": "aw_s", "seed": 0, "n": 373643, "nf": 435, "min": 0.5, "load_s": 3, "ds_s": 3, "fit_s": 20, "pred_s": 5, "auc": {"400": 0.79413, "500": 0.7938}} diff --git a/v13/feat_ev/top.json b/v13/feat_ev/top.json new file mode 100644 index 0000000000000000000000000000000000000000..f8eb40c13132d9dfcababee73066bdc043d2ccd0 --- /dev/null +++ b/v13/feat_ev/top.json @@ -0,0 +1 @@ +["tB_dmean", "tB_dot", "tB_dot_z", "tA_cos", "tA_dmax", "lr12_dmax", "lr12_cos", "tA_z", "lr16_dmax", "lr16_cos", "lr12_con05", "lr12_con05_dmax", "lr12_z", "lr16_con05_dmax", "lr16_con05", "lr12_con05_z", "lr10_dmax", "lr10_cos", "c_hist", "lr16_z", "lr16_con05_z", "lr10_z", "lr10_con05_dmax", "lr10_con05", "lr10_con05_z", "x_pp_con05_k32", "sq_ecos_dm", "sq_ecos", "x_pp_hcos", "sq_gap", "sq_logit", "sq_gap2", "sq_dm", "lr12_con1_dmax", "lr12_con1", "x_pp_con05_k32_z", "lr16_con1", "lr16_con1_dmax", "lr8_dmax", "lr8_cos", "lr12_con1_z", "joint32_cos_clkpart", "lr8_z", "lr16_con1_z", "c_hist_minus_skrec", "c_hist_minus_sk", "lr8_con05", "lr8_con05_dmax", "lr10_con1_dmax", "lr10_con1", "joint32_cos_clkpart_z", "lr8_con05_z", "lr10_con1_z", "sv_k64_con025", "x_con05_k128", "sv_k64_hcos", "x_con05_k64", "sv_k64_con05", "x_pp_hcos_z", "x_icos_k64", "ev_hist_n_all_z", "x_pp_con05_k32_rk", "ev_hist_n_all_rk", "v2_uip_n_rk", "hist_pos", "sv_k64_ccos", "hist_in", "hist_n_all", "v2_uip_n", "v2_uip_n_dm", "hist_n_own", "v2_inh_cnt", "hist_since_clk", "x_icon05_k64", "sv_k32_imax20", "x_pp_last3", "sv_k64_ucos", "x_conr05_k64", "sv_k64_con1", "x_crcos_k64"] \ No newline at end of file diff --git a/v15cat/ens/catA_c25_imp.json b/v15cat/ens/catA_c25_imp.json new file mode 100644 index 0000000000000000000000000000000000000000..6de2a163a5ef6920c5d1261ef5ebe5206ed56cd9 --- /dev/null +++ b/v15cat/ens/catA_c25_imp.json @@ -0,0 +1 @@ +{"tB_dot_z": 13.516409966075734, "tB_dmean": 3.872476566033176, "tB_dot": 6.246065398812302, "sq_logit": 11.217532101149212, "sq_z": 2.426532964673599, "lr12_con05_z": 0.3405822166428079, "tA_z": 1.0548579745555327, "sq_gap2": 5.734913724695959, "sq_gap": 2.6580066509251568, "c_hist": 0.4746082547219807, "sq_dm": 2.963638719122616, "tA_cos": 0.5574886868736187, "sq_rank": 0.7487146439841336, "v2_u_ctr": 1.2971478350601253, "tA_dmax": 0.2859114679238365, "lr12_con05": 3.259925326497348, "c_hist_minus_sk": 2.2682820294312944, "lr12_cos": 0.273909941194603, "v2_ip_per_imp": 0.17237279486449406, "joint32_cos_clkpart_z": 0.4208502863953487, "lr16_dmax": 0.458933303529129, "lr4_cos": 0.1826799416441229, "lr10_con05": 0.7489830798515803, "lr16_cos": 3.670761403168325, "lr12_con1": 0.2001221786056874, "c_hist_minus_skrec": 1.3171162822784006, "lr4_z": 0.06299590061372326, "sq_lse": 0.6610812395827019, "lr16_con05": 2.856506185061333, "us_ctr": 0.06912609843866928, "hist_cos_imp8": 0.0, "lr10_cos": 3.1388082021602712, "x_pp_con05_k32": 0.1641890380373155, "impsvd16_cos_z": 0.2229609566991243, "lr10_dmax": 1.8421820626724887, "lr16_con05_dmax": 3.9050943856841034, "lr4_dmax": 0.2531643954313367, "hist_cos_imp8_z": 0.0, "v2_h_gap_mean": 0.5388157232763231, "ev_sv_k32_imax20_mean": 0.5039089987291931, "sq_anull": 0.0791374313501347, "impsvd_cos_z": 0.0, "lr10_con05_dmax": 3.5102386967716614, "impsvd16_cos": 0.04116112355861682, "sq_amax": 0.0, "lr12_dmax": 0.3226539736897405, "sv_k16_udot": 0.27458156274714773, "cent_imp_mean_dm": 0.03673497886794547, "v2_svd_last": 0.07787181660061869, "lr12_con05_dmax": 2.903363932451149, "cent_clk_mean_dm": 0.0424737581304435, "v2_i_ctr_tb_diff": 0.0, "sv_k32_udot": 0.0, "lr12_con1_z": 1.2914656786868273, "ev_it_imp_z": 0.18827876454537068, "impsvd8_cos": 0.0, "sq_aent": 0.5823818348121411, "cent_imp_max": 0.025688579507534646, "ev_mf_score_mean": 0.1723508402503738, "lr6_z": 0.10697830666117761, "sq_mcos_dm": 0.0, "tB_mean": 0.0981144661474675, "us_nclk_warm": 0.5504695526031286, "cent_imp_top3": 0.0, "tB_lse": 0.2655000898439035, "x_icon05_k64": 0.2029765145864219, "sq_sd": 0.5315973143801361, "joint32_cos_clkpart": 0.352114875362449, "impsvd8_cos_z": 0.26601472767109396, "impsvd_rec": 0.0, "cent_clk_mean": 0.03387820520396585, "it_imp_rate": 0.7009534074584518, "v2_svd_hnorm": 0.0, "x_con05_k128": 0.15256715881332547, "lr10_con1": 0.0, "hist_cos_imp128": 0.36202978399663693, "cent_clk_max": 0.04873960781645971, "v2_i_imp_rate": 0.0, "x_clctr64_lift_z": 0.0, "lr4_con05_z": 0.16926321833910563, "sv_k32_hnorm": 0.2878525901563254, "x_pp_con05_k32_z": 0.8155273486504773, "impsvd_cos": 0.11550400940054957, "lr4_con05": 0.08478330566412355, "sq_ecos_dm": 0.5538028828415501, "sq_pcos_dm": 0.0, "impsvd64_cos_z": 0.057012905183852526, "skipsvd32_cos_z": 0.07878133290619399, "hist_maxsim_imp64_z": 0.03747134763874095, "cent_imp_mean": 0.0, "lr12_z": 1.1782461942906648, "impsvd128_cos_z": 0.0, "hist_cos_imp64_z": 0.07438212141089416, "x_icon05_k64_z": 0.3794905761251213, "hist_maxsim_imp64": 0.36062041947019013, "x_pp_last3_z": 0.22505934926598048, "hist_cos_impspace_z": 0.0, "v2_svd_maxsim": 0.0, "v2_ip_dec3": 0.07663340233625152, "c_impall": 0.08755313880588785, "item_id": 0.0, "user_id": 0.0, "ic64": 0.22761987231826797, "ic256": 0.0, "ic1024": 0.0, "uc64": 0.08094795200391383, "uc256": 0.0, "uc1024": 0.0, "lastc64": 0.00029956980797682766, "lastc256": 0.0, "rec_bin": 0.05173204481868061, "hlen_bin": 0.7201322247997235, "stale_bin": 0.0, "time_bucket": 0.0, "day_type": 0.0, "in_hist": 0.02779347763621527, "x_ic256_uc256": 0.0, "x_ic256_tb": 0.0, "x_inh_ic256": 0.0, "x_ic64_uc64": 0.5015168085182777, "x_lastc64_ic64": 0.0} \ No newline at end of file diff --git a/v15cat/ens/catA_nocat_c25_imp.json b/v15cat/ens/catA_nocat_c25_imp.json new file mode 100644 index 0000000000000000000000000000000000000000..bb7c942278ba2451c1d63eb9e4ea3cc8bdf76a37 --- /dev/null +++ b/v15cat/ens/catA_nocat_c25_imp.json @@ -0,0 +1 @@ +{"tB_dot_z": 6.9902673046856565, "tB_dmean": 6.260764508745816, "tB_dot": 6.0890849857256635, "sq_logit": 3.491520025728742, "sq_z": 1.5493323485822421, "lr12_con05_z": 0.7031129061946552, "tA_z": 2.0004935408238262, "sq_gap2": 4.03774952847766, "sq_gap": 3.524058603254935, "c_hist": 1.5440047859258967, "sq_dm": 2.3482836445159, "tA_cos": 1.9990490266091632, "sq_rank": 0.7779090199317849, "v2_u_ctr": 1.7803055965411572, "tA_dmax": 0.7099393230396099, "lr12_con05": 1.1842245020451774, "c_hist_minus_sk": 1.0084179643764577, "lr12_cos": 1.647362168161181, "v2_ip_per_imp": 1.883636398849122, "joint32_cos_clkpart_z": 1.8211280112707888, "lr16_dmax": 0.8337663315228766, "lr4_cos": 0.7009417399677561, "lr10_con05": 2.208103075917135, "lr16_cos": 0.609144408241506, "lr12_con1": 1.1312290396890925, "c_hist_minus_skrec": 1.7904931801764068, "lr4_z": 0.5414858568300224, "sq_lse": 1.013087460731241, "lr16_con05": 0.6850459237917462, "us_ctr": 0.6396565637800603, "hist_cos_imp8": 0.8947749249164829, "lr10_cos": 0.4896550205352329, "x_pp_con05_k32": 0.41903878133591843, "impsvd16_cos_z": 0.3447624636158747, "lr10_dmax": 1.26066170231611, "lr16_con05_dmax": 0.704970507855285, "lr4_dmax": 0.6878279430768477, "hist_cos_imp8_z": 0.7516488317059735, "v2_h_gap_mean": 0.5765423793689082, "ev_sv_k32_imax20_mean": 0.8353777444989258, "sq_anull": 0.5933233609601045, "impsvd_cos_z": 0.42347259762708817, "lr10_con05_dmax": 1.0775034117800406, "impsvd16_cos": 0.5891616409541074, "sq_amax": 0.7063577091817393, "lr12_dmax": 1.994063547421846, "sv_k16_udot": 0.620018635637964, "cent_imp_mean_dm": 0.40029734979105447, "v2_svd_last": 0.38069163676997536, "lr12_con05_dmax": 0.30052398606246317, "cent_clk_mean_dm": 0.5764655097746678, "v2_i_ctr_tb_diff": 0.2966271504734749, "sv_k32_udot": 0.5813166506016318, "lr12_con1_z": 0.5929805498334824, "ev_it_imp_z": 0.6819245237023212, "impsvd8_cos": 0.1649918458645088, "sq_aent": 0.6520797263298476, "cent_imp_max": 0.30162570563785546, "ev_mf_score_mean": 0.30682224124944474, "lr6_z": 0.6970730162879699, "sq_mcos_dm": 0.5749900932936821, "tB_mean": 0.37084883843991756, "us_nclk_warm": 1.3604309440079903, "cent_imp_top3": 0.4081766079410588, "tB_lse": 0.5724457049109533, "x_icon05_k64": 0.511325412325238, "sq_sd": 0.38962958023750466, "joint32_cos_clkpart": 0.5180247299311397, "impsvd8_cos_z": 0.5604377231780899, "impsvd_rec": 0.1928228668760003, "cent_clk_mean": 0.31849603086512795, "it_imp_rate": 0.9226231183348014, "v2_svd_hnorm": 0.382333335791804, "x_con05_k128": 0.44414271043313064, "lr10_con1": 0.18713472692704708, "hist_cos_imp128": 0.39898317433590885, "cent_clk_max": 0.3382276503117078, "v2_i_imp_rate": 0.9813725004684185, "x_clctr64_lift_z": 0.5981344348816418, "lr4_con05_z": 0.6239189085741159, "sv_k32_hnorm": 0.581060363506996, "x_pp_con05_k32_z": 0.644088469157049, "impsvd_cos": 0.2487771209081812, "lr4_con05": 0.4282002647157321, "sq_ecos_dm": 0.6167592914754156, "sq_pcos_dm": 0.40202870231779536, "impsvd64_cos_z": 0.2866230074055025, "skipsvd32_cos_z": 0.2810247487142958, "hist_maxsim_imp64_z": 0.4872699579218912, "cent_imp_mean": 0.29361775547596086, "lr12_z": 0.5441188971867299, "impsvd128_cos_z": 0.2940633588365459, "hist_cos_imp64_z": 0.27097887459062125, "x_icon05_k64_z": 0.7728535914347735, "hist_maxsim_imp64": 0.5224465833985532, "x_pp_last3_z": 0.3010656765334904, "hist_cos_impspace_z": 0.24163382775608092, "v2_svd_maxsim": 0.42436253414810493, "v2_ip_dec3": 1.0145662040940564, "c_impall": 0.28378280905656056} \ No newline at end of file diff --git a/v15cat/ens/catA_noid_c25_imp.json b/v15cat/ens/catA_noid_c25_imp.json new file mode 100644 index 0000000000000000000000000000000000000000..9597779430f1a0b8e05f4a1eb4c0790f7d3b5267 --- /dev/null +++ b/v15cat/ens/catA_noid_c25_imp.json @@ -0,0 +1 @@ +{"tB_dot_z": 7.452811001117354, "tB_dmean": 3.997996411639412, "tB_dot": 6.527640245902081, "sq_logit": 3.9950751065258623, "sq_z": 1.080286672312088, "lr12_con05_z": 0.3955201082679644, "tA_z": 1.9065128428717544, "sq_gap2": 4.247758481774307, "sq_gap": 2.8648730224058503, "c_hist": 0.7277414100540684, "sq_dm": 2.280178610159812, "tA_cos": 2.1772245764048086, "sq_rank": 0.9030654090329872, "v2_u_ctr": 1.5434163357632305, "tA_dmax": 2.041940364766753, "lr12_con05": 1.6654456555802037, "c_hist_minus_sk": 0.5333023474992472, "lr12_cos": 0.39692789120897853, "v2_ip_per_imp": 1.520731036639104, "joint32_cos_clkpart_z": 0.8389094808124632, "lr16_dmax": 1.3610965164493087, "lr4_cos": 0.8082360072328895, "lr10_con05": 0.7903656247551272, "lr16_cos": 0.538203378050391, "lr12_con1": 1.4639529050965259, "c_hist_minus_skrec": 1.2700903942861856, "lr4_z": 0.4836681638963256, "sq_lse": 0.9732574632306621, "lr16_con05": 1.1473256948555823, "us_ctr": 0.6185495409458842, "hist_cos_imp8": 1.217133148668392, "lr10_cos": 0.6498526044132305, "x_pp_con05_k32": 0.8081698449172027, "impsvd16_cos_z": 0.19152254047258907, "lr10_dmax": 1.4419630766277025, "lr16_con05_dmax": 1.403473306132562, "lr4_dmax": 0.6077641819946166, "hist_cos_imp8_z": 0.5774298093812993, "v2_h_gap_mean": 0.3802288192493736, "ev_sv_k32_imax20_mean": 0.673265326783915, "sq_anull": 0.7860597812018041, "impsvd_cos_z": 0.23051707424962117, "lr10_con05_dmax": 0.6045536391464469, "impsvd16_cos": 0.14987696920141386, "sq_amax": 0.5604618251359637, "lr12_dmax": 0.2604319681411039, "sv_k16_udot": 0.4137941670811338, "cent_imp_mean_dm": 0.32153371103276346, "v2_svd_last": 0.2882530063071337, "lr12_con05_dmax": 0.453751897234528, "cent_clk_mean_dm": 0.5562767414730256, "v2_i_ctr_tb_diff": 0.25874653864217756, "sv_k32_udot": 0.4843042541374386, "lr12_con1_z": 0.46904786283177874, "ev_it_imp_z": 0.6041882841932131, "impsvd8_cos": 0.19618813194688478, "sq_aent": 0.44435619007663857, "cent_imp_max": 0.3185242967617733, "ev_mf_score_mean": 0.2893386943341788, "lr6_z": 0.6650882245193508, "sq_mcos_dm": 0.43424214297237834, "tB_mean": 0.38156110494979634, "us_nclk_warm": 1.2254480284456626, "cent_imp_top3": 0.2522114312195675, "tB_lse": 0.4083654461560006, "x_icon05_k64": 0.6003566790590967, "sq_sd": 0.37351477893971685, "joint32_cos_clkpart": 0.368269895855882, "impsvd8_cos_z": 0.4549681214212998, "impsvd_rec": 0.16133409931342263, "cent_clk_mean": 0.438181221007239, "it_imp_rate": 0.8798464800342861, "v2_svd_hnorm": 0.4100385511839562, "x_con05_k128": 0.34789509724517775, "lr10_con1": 0.7025351983296615, "hist_cos_imp128": 0.7185628217143065, "cent_clk_max": 0.48115454313387074, "v2_i_imp_rate": 0.7554750402921505, "x_clctr64_lift_z": 0.5208474162167197, "lr4_con05_z": 0.41822956887003754, "sv_k32_hnorm": 0.4501935403302147, "x_pp_con05_k32_z": 0.42067036874940106, "impsvd_cos": 0.11422456190083455, "lr4_con05": 0.45626698894524254, "sq_ecos_dm": 0.6574048365443922, "sq_pcos_dm": 0.36861129657810954, "impsvd64_cos_z": 0.24493100335162143, "skipsvd32_cos_z": 0.2287698476010651, "hist_maxsim_imp64_z": 0.43401466540790407, "cent_imp_mean": 0.31772901481589116, "lr12_z": 0.9615072732110093, "impsvd128_cos_z": 0.26427963211291955, "hist_cos_imp64_z": 0.31276240270844946, "x_icon05_k64_z": 0.5798162506331973, "hist_maxsim_imp64": 0.39928147310718887, "x_pp_last3_z": 0.23382176195851864, "hist_cos_impspace_z": 0.24943158787151598, "v2_svd_maxsim": 0.4332569840550791, "v2_ip_dec3": 0.9495674071867923, "c_impall": 0.1867708418092203, "ic64": 0.44293465727020714, "ic256": 0.33038428383394186, "ic1024": 0.4438976866855539, "uc64": 0.49141791606760904, "uc256": 0.3455891496915014, "uc1024": 0.6432359802695159, "lastc64": 0.140014595417902, "lastc256": 0.14650909860004646, "rec_bin": 0.026729260923437987, "hlen_bin": 0.28950864243381214, "stale_bin": 0.07366246769639945, "time_bucket": 0.10766155238925597, "day_type": 0.0, "in_hist": 0.020046832949874124, "x_ic256_uc256": 1.7642196802957268, "x_ic256_tb": 0.3574982439708949, "x_inh_ic256": 0.3886738368229376, "x_ic64_uc64": 2.224200602014647, "x_lastc64_ic64": 0.30926749160148426} \ No newline at end of file diff --git a/v15cat/ens/catB_c25_imp.json b/v15cat/ens/catB_c25_imp.json new file mode 100644 index 0000000000000000000000000000000000000000..f687b1d2b11e6f36a0c31549c4d65a239363bcfc --- /dev/null +++ b/v15cat/ens/catB_c25_imp.json @@ -0,0 +1 @@ +{"it_imp": 0.0, "it_clk": 0.0, "it_imp24": 0.10345070863673884, "it_clk24": 0.0, "it_imp6": 0.0, "it_clk6": 0.0, "it_ctr": 0.05967531788578652, "it_ctr24": 0.0, "it_ctr6": 0.020721260444638644, "it_trend24": 0.0, "it_trend6": 0.0, "it_imptrend24": 0.0, "it_imp_rate": 0.0, "it_nclk_all": 0.10274664337626599, "it_nusers": 0.0, "it_nclk_warm": 0.0, "it_clk_dec2d": 0.0, "us_nev": 0.0, "us_nclk_win": 0.0, "us_ctr": 0.43781769347103094, "us_nclk_all": 0.0, "us_nclk_warm": 0.14929685312579272, "us_nuniq": 0.11147821672702535, "us_nclk24": 0.0, "us_repeat_rate": 0.0, "hist_n_all": 0.0, "hist_n_warm": 0.0, "hist_n_win": 0.0, "hist_imp": 0.0, "hist_noclk": 0.0, "us_hlen": 0.11790466972239255, "us_nohist": 0.0, "us_cold": 0.0, "us_since_ev": 0.0, "us_since_clk": 0.0, "hist_since_clk": 0.0, "hist_since_imp": 0.0233977664049023, "hist_pos": 0.09002269754773623, "hist_n_own": 0.09174515737441456, "hist_in": 0.013887431152682147, "co_cos_max20": 0.0, "co_cos_mean20": 0.0, "co_cos_max5": 0.0, "co_cos_mean5": 0.0, "co_cos_last": 0.0, "co_tr_max20": 0.0, "co_tr_mean20": 0.0, "co_tr_max5": 0.0, "co_tr_mean5": 0.0, "co_tr_last": 0.0, "co_cos_rw": 0.0, "co_cos_sum20": 0.0, "sv_k16_ucos": 0.0, "sv_k16_udot": 0.0, "sv_k16_hcos": 0.0, "sv_k32_ucos": 0.0, "sv_k32_udot": 0.07290301023795961, "sv_k32_hcos": 0.0, "sv_k32_hnorm": 0.13536570388799063, "sv_k32_imax20": 0.09056883491410028, "sv_k32_imax5": 0.0, "sv_k64_ucos": 0.13479527987428963, "sv_k64_udot": 0.0, "sv_k64_hcos": 0.194228683435137, "sv_k64_ccos": 0.0, "sv_k64_ncos": 0.0, "sv_k64_con025": 0.0, "sv_k64_con05": 0.0, "sv_k64_con1": 0.0, "mf_score": 0.046360377840468585, "mf_dot": 0.0, "mf_bi": 0.036345628748574596, "ev_it_ctr_rk": 0.0, "ev_it_ctr_z": 0.0, "ev_it_ctr_max": 0.0, "ev_it_ctr_mean": 0.0, "ev_it_ctr24_rk": 0.0, "ev_it_ctr24_z": 0.0, "ev_it_ctr24_max": 0.0, "ev_it_ctr24_mean": 0.0, "ev_it_imp_rk": 0.1882547976917625, "ev_it_imp_z": 0.0, "ev_it_imp_max": 0.15592927358499642, "ev_it_imp_mean": 0.0, "ev_it_nclk_all_rk": 0.14005808457540644, "ev_it_nclk_all_z": 0.18623749848306437, "ev_it_nclk_all_max": 0.0, "ev_it_nclk_all_mean": 0.19115539010074126, "ev_it_clk_dec2d_rk": 0.0, "ev_it_clk_dec2d_z": 0.0, "ev_it_clk_dec2d_max": 0.0, "ev_it_clk_dec2d_mean": 0.0, "ev_hist_n_all_rk": 0.0, "ev_hist_n_all_z": 0.0, "ev_hist_n_all_max": 0.06798991032307927, "ev_hist_n_all_mean": 0.0, "ev_hist_noclk_rk": 0.008537303784221722, "ev_hist_noclk_z": 0.0, "ev_hist_noclk_max": 0.0, "ev_hist_noclk_mean": 0.0, "ev_co_cos_max20_rk": 0.0, "ev_co_cos_max20_z": 0.0, "ev_co_cos_max20_max": 0.0, "ev_co_cos_max20_mean": 0.0, "ev_co_cos_rw_rk": 0.0, "ev_co_cos_rw_z": 0.0, "ev_co_cos_rw_max": 0.0, "ev_co_cos_rw_mean": 0.0, "ev_co_tr_max5_rk": 0.0, "ev_co_tr_max5_z": 0.0, "ev_co_tr_max5_max": 0.0, "ev_co_tr_max5_mean": 0.0, "ev_co_tr_last_rk": 0.0, "ev_co_tr_last_z": 0.0, "ev_co_tr_last_max": 0.0, "ev_co_tr_last_mean": 0.0, "ev_sv_k16_ucos_rk": 0.0, "ev_sv_k16_ucos_z": 0.0, "ev_sv_k16_ucos_max": 0.0, "ev_sv_k16_ucos_mean": 0.0, "ev_sv_k32_ucos_rk": 0.0, "ev_sv_k32_ucos_z": 0.0, "ev_sv_k32_ucos_max": 0.0, "ev_sv_k32_ucos_mean": 0.0, "ev_sv_k64_ucos_rk": 0.0, "ev_sv_k64_ucos_z": 0.0, "ev_sv_k64_ucos_max": 0.0, "ev_sv_k64_ucos_mean": 0.03126758733107813, "ev_sv_k32_hcos_rk": 0.0, "ev_sv_k32_hcos_z": 0.0, "ev_sv_k32_hcos_max": 0.0, "ev_sv_k32_hcos_mean": 0.0, "ev_sv_k64_hcos_rk": 0.0, "ev_sv_k64_hcos_z": 0.0, "ev_sv_k64_hcos_max": 0.0, "ev_sv_k64_hcos_mean": 0.04501834675713056, "ev_sv_k32_imax20_rk": 0.0, "ev_sv_k32_imax20_z": 0.0, "ev_sv_k32_imax20_max": 0.0, "ev_sv_k32_imax20_mean": 0.19014133988448909, "ev_sv_k64_con05_rk": 0.7944387500250951, "ev_sv_k64_con05_z": 0.27205355132458187, "ev_sv_k64_con05_max": 0.0, "ev_sv_k64_con05_mean": 0.06575649144079088, "ev_sv_k64_con025_rk": 0.0, "ev_sv_k64_con025_z": 0.0, "ev_sv_k64_con025_max": 0.0, "ev_sv_k64_con025_mean": 0.10888352108485448, "ev_sv_k64_ccos_rk": 0.0, "ev_sv_k64_ccos_z": 0.0, "ev_sv_k64_ccos_max": 0.0, "ev_sv_k64_ccos_mean": 0.0, "ev_mf_score_rk": 0.0, "ev_mf_score_z": 0.0, "ev_mf_score_max": 0.0, "ev_mf_score_mean": 0.06452286381812436, "ev_mf_dot_rk": 0.0, "ev_mf_dot_z": 0.0, "ev_mf_dot_max": 0.0, "ev_mf_dot_mean": 0.026101687697445565, "v2_i_imp": 0.18598095024388459, "v2_i_clk": 0.0, "v2_i_imp1d": 0.10372264542378348, "v2_i_clk1d": 0.0, "v2_i_ctr": 0.024401327139744713, "v2_i_ctr1d": 0.0, "v2_i_imp_rate": 0.0, "v2_ip_n": 0.0, "v2_ip_users": 0.0, "v2_ip_1d": 0.0, "v2_ip_3d": 0.0, "v2_ip_dec1": 0.0, "v2_ip_dec3": 0.38737864648552095, "v2_ip_warm": 0.1266942543201761, "v2_i_ctr_tb": 0.0, "v2_u_ev": 0.0, "v2_u_clk": 0.0, "v2_u_ctr": 0.5141055193830841, "v2_up_n": 0.0, "v2_up_uniq": 0.2126674896281043, "v2_up_1d": 0.10007381144732148, "v2_ui_imp": 0.0, "v2_ui_clk": 0.0, "v2_uip_n": 0.0, "v2_h_len": 0.15228276707575134, "v2_h_since_last": 0.0, "v2_h_since_first": 0.0, "v2_h_uniq": 0.9252071072795839, "v2_h_1d_before_last": 0.0, "v2_u_since_ev": 0.0, "v2_up_since": 0.03047189055671312, "v2_ui_since_imp": 0.0, "v2_uip_since": 0.0, "v2_ui_noclk": 0.0, "v2_ip_per_imp": 0.033905144028515855, "v2_i_ctr_tb_diff": 0.0, "v2_h_gap_mean": 0.10217327181527375, "v2_h_uniq_ratio": 0.0, "v2_co_sum": 0.0, "v2_co_max": 0.0, "v2_co_mean": 0.04008866175073119, "v2_co_rw": 0.0, "v2_co_cnt": 0.0, "v2_co_last3": 0.0, "v2_inh_cnt": 0.008017182549813557, "v2_inh_pos": 0.0, "v2_inh_since": 0.0, "v2_co_nz": 0.0, "v2_tr_last": 0.00992270292367012, "v2_trp_last": 0.0, "v2_tr_last3": 0.0, "v2_trp_last3": 0.0, "v2_svd_cos": 0.11459537502049119, "v2_svd_ucos": 0.05169636197226213, "v2_svd_hnorm": 0.0, "v2_svd32_cos": 0.06107183102892275, "v2_svd16_cos": 0.0, "v2_svd24_cos": 0.0, "v2_svd48_cos": 0.06888902065402104, "v2_svd_maxsim": 0.0, "v2_svd_last": 0.0, "v2_svd_top3": 0.0, "v2_i_ctr_rk": 0.0, "v2_ip_dec1_rk": 0.0, "v2_ip_n_rk": 0.10420602445848744, "v2_co_sum_rk": 0.02778550716318013, "v2_co_rw_rk": 0.0, "v2_co_max_rk": 0.0, "v2_svd_cos_rk": 0.0, "v2_svd_ucos_rk": 0.0, "v2_svd32_cos_rk": 0.3167923450245648, "v2_svd16_cos_rk": 0.0, "v2_svd24_cos_rk": 0.0, "v2_svd48_cos_rk": 0.08074091958873496, "v2_svd_maxsim_rk": 0.0, "v2_svd_top3_rk": 0.0, "v2_tr_last3_rk": 0.0, "v2_uip_n_rk": 0.0, "v2_i_ctr_tb_rk": 0.0, "v2_ui_noclk_rk": 0.0, "v2_i_ctr_dm": 0.0, "v2_ip_dec1_dm": 0.0, "v2_ip_n_dm": 0.0, "v2_co_sum_dm": 0.027542219828747747, "v2_co_rw_dm": 0.0, "v2_co_max_dm": 0.013751173372532304, "v2_svd_cos_dm": 0.15999275300992516, "v2_svd_ucos_dm": 0.03426127591326429, "v2_svd32_cos_dm": 0.0, "v2_svd16_cos_dm": 0.041626703111792784, "v2_svd24_cos_dm": 0.0, "v2_svd48_cos_dm": 0.03449603705888174, "v2_svd_maxsim_dm": 0.0, "v2_svd_top3_dm": 0.0, "v2_tr_last3_dm": 0.015764185025062143, "v2_uip_n_dm": 0.16902599002417162, "v2_i_ctr_tb_dm": 0.0, "v2_ui_noclk_dm": 0.0, "v2_co_sum_rmax": 0.0, "v2_svd_cos_rmax": 0.0, "v2_ip_dec1_rmax": 0.0, "v2_co_rw_rmax": 0.0, "v2_ev_co_max": 0.12048686579067247, "v2_ev_svd_max": 0.0, "v2_ev_svd_mean": 0.0, "v2_ev_n_inh": 0.0, "v2_ev_ictr_mean": 0.0, "v2_ev_ipdec_mean": 0.0, "x_con05_k32": 0.0, "x_con05_k64": 0.21996120520282342, "x_conr05_k64": 0.0, "x_conr1_k64": 0.1293325434652786, "x_crcos_k64": 0.0735568146991426, "x_icon05_k64": 0.0, "x_icos_k64": 0.0, "x_clctr64": 0.0, "x_clctr64_lift": 0.0, "x_climp64": 0.0, "x_clctr256": 0.0, "x_clctr256_lift": 0.0, "x_climp256": 0.0, "x_con05_k128": 0.0, "x_con05_k32_z": 0.0, "x_con05_k64_z": 0.0, "x_conr05_k64_z": 0.05386732338325961, "x_conr1_k64_z": 0.17408859087487413, "x_crcos_k64_z": 0.0, "x_icon05_k64_z": 0.2303203343754845, "x_icos_k64_z": 0.0, "x_clctr64_z": 0.0, "x_clctr64_lift_z": 0.20433674950886144, "x_clctr256_z": 0.0, "x_clctr256_lift_z": 0.0, "x_con05_k128_z": 0.0, "x_con05_k32_rk": 0.16645930157740502, "x_con05_k64_rk": 0.0, "x_conr05_k64_rk": 0.0, "x_conr1_k64_rk": 0.0, "x_crcos_k64_rk": 0.0, "x_icon05_k64_rk": 0.0, "x_icos_k64_rk": 0.0, "x_clctr64_rk": 0.0, "x_clctr64_lift_rk": 0.0, "x_clctr256_rk": 0.0, "x_clctr256_lift_rk": 0.0, "x_con05_k128_rk": 0.0, "lr4_cos": 0.19554267071519763, "lr4_con05": 0.0, "lr4_con1": 0.0, "lr6_cos": 0.0, "lr6_con05": 0.0, "lr6_con1": 0.0, "lr8_cos": 0.0, "lr8_con05": 0.0, "lr8_con1": 0.0, "lr10_cos": 3.024427230044387, "lr10_con05": 1.666298670388254, "lr10_con1": 0.0, "lr12_cos": 0.23738515271717586, "lr12_con05": 2.489332942879444, "lr12_con1": 1.536935904903531, "lr16_cos": 1.422048987841999, "lr16_con05": 0.43162372492579165, "lr16_con1": 0.0, "lr4_z": 0.12456400347638388, "lr4_con05_z": 0.035245867089928454, "lr4_con1_z": 0.0, "lr6_z": 0.0, "lr6_con05_z": 0.0, "lr6_con1_z": 0.0, "lr8_z": 0.0, "lr8_con05_z": 0.13873429354563424, "lr8_con1_z": 0.0, "lr10_z": 0.0, "lr10_con05_z": 0.0, "lr10_con1_z": 0.0, "lr12_z": 1.3839989900754248, "lr12_con05_z": 2.2059221647958136, "lr12_con1_z": 2.42898892195656, "lr16_z": 0.0, "lr16_con05_z": 3.123620980812669, "lr16_con1_z": 0.5913576654925102, "lr4_dmax": 0.2798140220938733, "lr4_con05_dmax": 0.05563392392450977, "lr4_con1_dmax": 0.0, "lr6_dmax": 0.44025438892852103, "lr6_con05_dmax": 0.0, "lr6_con1_dmax": 0.0, "lr8_dmax": 0.0, "lr8_con05_dmax": 0.0, "lr8_con1_dmax": 0.0, "lr10_dmax": 0.2543474287088519, "lr10_con05_dmax": 0.4440827980176615, "lr10_con1_dmax": 0.5594611455761663, "lr12_dmax": 0.7345582676863776, "lr12_con05_dmax": 3.310828090775284, "lr12_con1_dmax": 0.0, "lr16_dmax": 0.09706203970260747, "lr16_con05_dmax": 0.0, "lr16_con1_dmax": 0.0, "tA_cos": 3.8165692547455294, "tB_dot": 1.1806622565819862, "tA_z": 5.748553761982376, "tA_dmax": 0.1626819133528427, "tB_lse": 0.0, "tB_mean": 0.0, "tB_dmean": 9.054389932158998, "tB_dot_z": 12.59375004566593, "x_pp_con05_k32": 0.0, "x_pp_hcos": 0.17914798906733778, "x_pp_last3": 0.11180684436379894, "x_pp_con05_k32_z": 0.48684376374544613, "x_pp_hcos_z": 0.03918380162179318, "x_pp_last3_z": 0.05742991771576953, "x_pp_con05_k32_rk": 0.115086390744598, "x_pp_hcos_rk": 0.0, "x_pp_last3_rk": 0.16439715728124035, "sq_logit": 10.683605795716446, "sq_z": 2.606519990227831, "sq_dm": 1.8354414577649023, "sq_lse": 0.27190837546252966, "sq_rank": 1.9090626175097352, "sq_gap": 5.8276518559771056, "sq_gap2": 0.39802737277246514, "sq_amax": 0.0, "sq_anull": 0.0, "sq_amean": 0.0, "sq_aent": 0.12935640594614595, "sq_ecos": 0.0, "sq_pcos_dm": 0.0, "sq_mcos_dm": 0.0, "sq_ecos_dm": 0.0, "sq_sd": 0.0, "impsvd_rec": 0.0, "impsvd_cos": 0.06241287522413086, "impsvd_cos_z": 0.0, "impsvd_cos_rk": 0.0, "impsvd16_cos": 0.0, "impsvd16_cos_z": 0.0, "impsvd64_cos": 0.0, "impsvd64_cos_z": 0.04847865574440983, "skipsvd32_cos": 0.0, "skipsvd32_cos_z": 0.0, "joint32_cos_clkpart": 0.37291182988054244, "joint32_cos_clkpart_z": 0.12865246808075573, "hist_cos_impspace": 0.0, "hist_cos_impspace_z": 0.0, "impsvd8_cos": 0.0, "hist_cos_imp8": 0.1286153240312315, "hist_cos_imp64": 0.0, "hist_maxsim_imp64": 0.045258926798228995, "impsvd128_cos": 0.0, "hist_cos_imp128": 0.30713610145010906, "impsvd8_cos_z": 0.0, "hist_cos_imp8_z": 0.0, "hist_cos_imp64_z": 0.07803182920154311, "hist_maxsim_imp64_z": 0.0, "impsvd128_cos_z": 0.0, "hist_cos_imp128_z": 0.0, "cent_clk_mean": 0.0, "cent_clk_max": 0.0, "cent_clk_top3": 0.015546746692466255, "cent_clk_mean_dm": 0.019981826192081505, "slate_clk_coh": 0.0, "cent_imp_mean": 0.0, "cent_imp_max": 0.0, "cent_imp_top3": 0.060107110384494414, "cent_imp_mean_dm": 0.024309916315973915, "slate_imp_coh": 0.011801714272640063, "c_skall": 0.0, "c_skrec": 0.03465816176716865, "c_impall": 0.0, "c_hist": 2.5798941165170666, "c_hist_minus_sk": 0.7155219192279391, "c_hist_minus_skrec": 0.5427551448671507, "item_id": 0.0, "user_id": 0.0, "ic64": 0.018978999973346365, "ic256": 0.0, "ic1024": 0.0, "uc64": 0.0, "uc256": 0.0, "uc1024": 0.0, "lastc64": 0.0, "lastc256": 0.0, "rec_bin": 0.0, "hlen_bin": 0.0, "stale_bin": 0.0, "time_bucket": 0.0, "day_type": 0.0, "in_hist": 0.0, "x_ic256_uc256": 0.0, "x_ic256_tb": 0.0, "x_inh_ic256": 0.0, "x_ic64_uc64": 0.0, "x_lastc64_ic64": 0.0} \ No newline at end of file diff --git a/v15cat/ens/clog.jsonl b/v15cat/ens/clog.jsonl new file mode 100644 index 0000000000000000000000000000000000000000..c899999aaabe7ce10feb253b0cb9f0a6ed38fd8e --- /dev/null +++ b/v15cat/ens/clog.jsonl @@ -0,0 +1,4 @@ +{"cut": 25, "tag": "catA", "set": "A", "n": 745474, "nf": 121, "ncat": 21, "best_iter": 35, "auc": 0.80884, "load_s": 7, "fit_s": 25, "min": 0.6, "es_on_val": true} +{"cut": 25, "tag": "catA_noid", "set": "A", "n": 745474, "nf": 119, "ncat": 19, "best_iter": 847, "auc": 0.8145, "load_s": 6, "fit_s": 153, "min": 2.7, "es_on_val": true} +{"cut": 25, "tag": "catA_nocat", "set": "A", "n": 745474, "nf": 100, "ncat": 0, "best_iter": 552, "auc": 0.815, "load_s": 7, "fit_s": 44, "min": 0.8, "es_on_val": true} +{"cut": 25, "tag": "catB", "set": "B", "n": 745474, "nf": 456, "ncat": 21, "best_iter": 34, "auc": 0.80798, "load_s": 10, "fit_s": 44, "min": 0.9, "es_on_val": true} diff --git a/v15cat/pylib/catboost/__init__.py b/v15cat/pylib/catboost/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f315ce34c2137b2ea5bc76ef405f74870afed395 --- /dev/null +++ b/v15cat/pylib/catboost/__init__.py @@ -0,0 +1,23 @@ +from .core import ( + FeaturesData, EFstrType, EShapCalcType, EFeaturesSelectionAlgorithm, EFeaturesSelectionGrouping, + Pool, CatBoost, CatBoostClassifier, CatBoostRegressor, CatBoostRanker, CatBoostError, cv, sample_gaussian_process, train, + sum_models, _have_equal_features, to_regressor, to_classifier, to_ranker, MultiRegressionCustomMetric, + MultiRegressionCustomObjective, MultiTargetCustomMetric, MultiTargetCustomObjective +) # noqa +from .version import VERSION as __version__ # noqa +__all__ = [ + 'FeaturesData', 'EFstrType', 'EShapCalcType', 'EFeaturesSelectionAlgorithm', 'EFeaturesSelectionGrouping', + 'Pool', 'CatBoost', 'CatBoostClassifier', 'CatBoostRegressor', 'CatBoostRanker', 'CatBoostError', 'CatboostError', + 'cv', 'sample_gaussian_process', 'train', 'sum_models', '_have_equal_features', + 'to_regressor', 'to_classifier', 'to_ranker', 'MultiRegressionCustomMetric', 'MultiRegressionCustomObjective', + 'MultiTargetCustomMetric', 'MultiTargetCustomObjective' +] + +# API compatibility alias. +CatboostError = CatBoostError + +try: + from .widget import MetricVisualizer, MetricsPlotter, XGBPlottingCallback, lgbm_plotting_callback # noqa + __all__.extend(['MetricVisualizer', 'MetricsPlotter', 'XGBPlottingCallback', 'lgbm_plotting_callback']) +except Exception: + pass diff --git a/v15cat/pylib/catboost/carry.py b/v15cat/pylib/catboost/carry.py new file mode 100644 index 0000000000000000000000000000000000000000..377b120e5bd039684889f43635bd6f3d6432308a --- /dev/null +++ b/v15cat/pylib/catboost/carry.py @@ -0,0 +1,81 @@ +from . import _catboost + +_carry_by_index = _catboost._carry_by_index +_carry_by_name = _catboost._carry_by_name +_uplift_by_index = _catboost._uplift_by_index +_uplift_by_name = _catboost._uplift_by_name + + +def carry(model, features): + """ + Parameters + ---------- + model : + CatBoost / CatBoostClassifier / CatBoostRanker / CatBoostRegressor model + features : + must be a dict mapping strings (factor names) or integers (flat indexes) into floats + NOTE: values in a dict can be lists of floats, but in this case they must all be the same size + """ + model = model.copy() + assert type(features) is dict + + factor_ids = [] + factor_values = [] + for factor_id, factor_value in features.items(): + assert type(factor_id) is str or type(factor_id) is int + assert type(factor_value) is float or type(factor_value) is int or type(factor_value) is list + + if type(factor_value) is list: + for value in factor_value: + assert type(value) is float or type(value) is int + else: + factor_value = [factor_value] + + if len(factor_values): + assert len(factor_values[0]) == len(factor_value) + assert type(factor_ids[0]) is type(factor_id) + + factor_ids.append(factor_id) + factor_values.append(factor_value) + + if len(factor_ids): + if type(factor_ids[0]) is int: + _carry_by_index(model._object, factor_ids, factor_values) + else: + _carry_by_name(model._object, factor_ids, factor_values) + return model + + +def uplift(model, features): + """ + Parameters + ---------- + model : + CatBoost / CatBoostClassifier / CatBoostRanker / CatBoostRegressor model + NOTE: uplift allways use RawFormulaVal + features : + must be a dict mapping strings (factor names) or integers (flat indexes) into pairs of floats (base and next values) + """ + model = model.copy() + + factor_ids = [] + factor_base_values = [] + factor_next_values = [] + for factor_id, (factor_base_value, factor_next_value) in features.items(): + assert type(factor_id) is str or type(factor_id) is int + assert type(factor_base_value) is float or type(factor_base_value) is int + assert type(factor_next_value) is float or type(factor_next_value) is int + if len(factor_ids): + assert type(factor_ids[0]) is type(factor_id) + + factor_ids.append(factor_id) + factor_base_values.append(float(factor_base_value)) + factor_next_values.append(float(factor_next_value)) + + if len(factor_ids): + if type(factor_ids[0]) is str: + _uplift_by_name(model._object, factor_ids, factor_base_values, factor_next_values) + else: + _uplift_by_index(model._object, factor_ids, factor_base_values, factor_next_values) + + return model diff --git a/v15cat/pylib/catboost/core.py b/v15cat/pylib/catboost/core.py new file mode 100644 index 0000000000000000000000000000000000000000..01dcf9dbfe1765b5f4a7996a710690ad2c041645 --- /dev/null +++ b/v15cat/pylib/catboost/core.py @@ -0,0 +1,7592 @@ +from contextlib import contextmanager # noqa E402 +from copy import deepcopy +import logging +import sys +import os + +if sys.version_info >= (3, 3): + from collections.abc import Iterable, Sequence, Mapping, MutableMapping +else: + from collections import Iterable, Sequence, Mapping, MutableMapping + +from collections import OrderedDict, defaultdict + +from six import iteritems, string_types, integer_types + +import warnings +import numpy as np +import tempfile +import shutil +import json +from enum import Enum +from operator import itemgetter +import threading +from typing import Dict, Optional + +try: + import pandas as pd +except ImportError: + # just to avoid checking (pd is not None) everywhere + class pandas: + class DataFrame(object): + pass + + class Series(object): + pass + pd = pandas + +try: + import polars as pl +except ImportError: + # just to avoid checking (pl is not None) everywhere + class polars: + class DataFrame(object): + pass + + class Series(object): + pass + pl = polars + +import scipy.sparse + + +from .plot_helpers import save_plot_file, try_plot_offline, OfflineMetricVisualizer +from . import _catboost +from .metrics import BuiltinMetric + + +_typeof = type + +_PoolBase = _catboost._PoolBase +_CatBoost = _catboost._CatBoost +_MetricCalcerBase = _catboost._MetricCalcerBase +_cv = _catboost._cv +_set_logger = _catboost._set_logger +_reset_logger = _catboost._reset_logger +_configure_malloc = _catboost._configure_malloc +CatBoostError = _catboost.CatBoostError +_metric_description_or_str_to_str = _catboost._metric_description_or_str_to_str +is_classification_objective = _catboost.is_classification_objective +is_cv_stratified_objective = _catboost.is_cv_stratified_objective +is_regression_objective = _catboost.is_regression_objective +is_multiregression_objective = _catboost.is_multiregression_objective +is_multitarget_objective = _catboost.is_multitarget_objective +is_multilabel_objective = _catboost.is_multilabel_objective +is_survivalregression_objective = _catboost.is_survivalregression_objective +is_multiclass_compatible_objective = _catboost.is_multiclass_compatible_objective +is_groupwise_metric = _catboost.is_groupwise_metric +is_ranking_metric = _catboost.is_ranking_metric +is_maximizable_metric = _catboost.is_maximizable_metric +is_minimizable_metric = _catboost.is_minimizable_metric +_PreprocessParams = _catboost._PreprocessParams +_check_train_params = _catboost._check_train_params +_MetadataHashProxy = _catboost._MetadataHashProxy +_NumpyAwareEncoder = _catboost._NumpyAwareEncoder +FeaturesData = _catboost.FeaturesData +_have_equal_features = _catboost._have_equal_features +SPARSE_MATRIX_TYPES = _catboost.SPARSE_MATRIX_TYPES +MultiTargetCustomMetric = _catboost.MultiTargetCustomMetric +MultiTargetCustomObjective = _catboost.MultiTargetCustomObjective +MultiRegressionCustomMetric = _catboost.MultiTargetCustomMetric # for compatibility +MultiRegressionCustomObjective = _catboost.MultiTargetCustomObjective # for compatibility +fspath = _catboost.fspath +_eval_metric_util = _catboost._eval_metric_util + + +logger = logging.getLogger(__name__) + +_configure_malloc() +_catboost._library_init() + +INTEGER_TYPES = (integer_types, np.integer) +FLOAT_TYPES = (float, np.floating) +STRING_TYPES = (string_types,) +ARRAY_TYPES = (list, np.ndarray, pd.DataFrame, pd.Series, pl.DataFrame, pl.Series) + +if sys.version_info >= (3, 6): + PATH_TYPES = STRING_TYPES + (os.PathLike,) +elif sys.version_info >= (3, 4): + from pathlib import Path + PATH_TYPES = STRING_TYPES + (Path,) +else: + PATH_TYPES = STRING_TYPES + + +def create_dir_if_not_exist(path): + if not os.path.exists(path): + os.mkdir(path) + + +class _StreamLikeWrapper: + def __init__(self, callable_object): + self.callable_object = callable_object + + def write(self, message): + self.callable_object(message) + + +class _CustomLoggersStack(object): + def __init__(self): + self._lock = threading.Lock() + self._owning_thread_id = None + self._stack = [] # contains (cout, cerr) pairs + + @staticmethod + def _get_stream_like_object(obj): + if hasattr(obj, 'write'): + return obj + if hasattr(obj, '__call__'): + return _StreamLikeWrapper(obj) + raise CatBoostError( + 'Expected callable object or stream-like object' + ) + + def push(self, log_cout=None, log_cerr=None): + with self._lock: + if not self._stack: + self._owning_thread_id = threading.current_thread().ident + elif self._owning_thread_id != threading.current_thread().ident: + if (log_cout is not None) or (log_cerr is not None): + raise CatBoostError( + 'CatBoost custom loggers have been already set in another thread. ' + + ' Setting custom loggers from different threads is not currently supported' + ) + return + + def init_log(log, default, index_in_stack): + if log is None: + if not self._stack: + return default + else: + return self._stack[-1][index_in_stack] + return _CustomLoggersStack._get_stream_like_object(log) + + cout = init_log(log_cout, sys.stdout, 0) + cerr = init_log(log_cerr, sys.stderr, 1) + + _reset_logger() + _set_logger(cout, cerr) + self._stack.append((cout, cerr)) + + def pop(self): + with self._lock: + if self._owning_thread_id != threading.current_thread().ident: + # because push from other threads does nothing + return + if not self._stack: + raise RuntimeError('Attempt to pop from an empty stack') + + _reset_logger() + if len(self._stack) != 1: + _set_logger(*self._stack[-2]) + else: + self._owning_thread_id = None + self._stack.pop() + + +_custom_loggers_stack = _CustomLoggersStack() + + +@contextmanager +def log_fixup(log_cout=None, log_cerr=None): + _custom_loggers_stack.push(log_cout, log_cerr) + try: + yield + finally: + _custom_loggers_stack.pop() + + +def _cast_to_base_types(value): + # NOTE: Special case, avoiding new list creation. + if isinstance(value, list): + for index, element in enumerate(value): + value[index] = _cast_to_base_types(element) + return value + if isinstance(value, ARRAY_TYPES[1:]): + new_value = [] + for element in value: + new_value.append(_cast_to_base_types(element)) + return new_value + if isinstance(value, (Mapping, MutableMapping)): + for key in list(value): + value[key] = _cast_to_base_types(value[key]) + return value + if isinstance(value, bool): + return value + if isinstance(value, INTEGER_TYPES): + return int(value) + if isinstance(value, FLOAT_TYPES): + return float(value) + if isinstance(value, tuple(set(PATH_TYPES) - set(STRING_TYPES))): + return fspath(value) + return value + + +def metric_description_or_str_to_str(description): + return _metric_description_or_str_to_str(description) + + +def _check_param_type(value, name, types, or_none=True): + if not isinstance(value, types + ((type(None),) if or_none else ())): + raise CatBoostError('Parameter {} should have a type of {}, got {}'.format(name, types, type(value))) + + +def _process_verbose(metric_period=None, verbose=None, logging_level=None, verbose_eval=None, silent=None): + _check_param_type(metric_period, 'metric_period', (int,)) + _check_param_type(verbose, 'verbose', (bool, int)) + _check_param_type(logging_level, 'logging_level', (string_types,)) + _check_param_type(verbose_eval, 'verbose_eval', (bool, int)) + _check_param_type(silent, 'silent', (bool,)) + + params = locals() + exclusive_params = ['verbose', 'logging_level', 'verbose_eval', 'silent'] + at_most_one = sum(params.get(exclusive) is not None for exclusive in exclusive_params) + if at_most_one > 1: + raise CatBoostError('Only one of parameters {} should be set'.format(exclusive_params)) + + if verbose is None: + if silent is not None: + verbose = not silent + elif verbose_eval is not None: + verbose = verbose_eval + if verbose is not None: + verbose = int(verbose) + + return (metric_period, verbose, logging_level) + + +def enum_from_enum_or_str(enum_type, arg): + if isinstance(arg, enum_type): + return arg + elif isinstance(arg, str): + return enum_type[arg] + else: + raise Exception("can't create enum " + str(enum_type) + " from type " + str(type(arg))) + + +class EFstrType(Enum): + """Calculate score for every feature by values change.""" + PredictionValuesChange = 0 + """Calculate score for every feature by loss change""" + LossFunctionChange = 1 + """Use LossFunctionChange for ranking models and PredictionValuesChange otherwise""" + FeatureImportance = 2 + """Calculate pairwise score between every feature.""" + Interaction = 3 + """Calculate SHAP Values for every object.""" + ShapValues = 4 + """Calculate most important features explaining difference in predictions for a pair of documents""" + PredictionDiff = 5 + """Calculate SHAP Interaction Values pairwise between every feature for every object.""" + ShapInteractionValues = 6 + """Calculate SAGE Values for every feature""" + SageValues = 7 + + +class EShapCalcType(Enum): + """Calculate regular SHAP values""" + Regular = "Regular" + """Calculate approximate SHAP values""" + Approximate = "Approximate" + """Calculate exact SHAP values""" + Exact = "Exact" + + +class EFeaturesSelectionAlgorithm(Enum): + """Use prediction values change as feature strength, eliminate batch of features at once""" + RecursiveByPredictionValuesChange = "RecursiveByPredictionValuesChange" + """Use loss function change as feature strength, eliminate batch of features at each step""" + RecursiveByLossFunctionChange = "RecursiveByLossFunctionChange" + """Use shap values to estimate loss function change, eliminate features one by one""" + RecursiveByShapValues = "RecursiveByShapValues" + + +class EFeaturesSelectionGrouping(Enum): + """Select individual features""" + Individual = "Individual" + """Select feature groups (marked by tags)""" + ByTags = "ByTags" + + +def _get_features_indices(features, feature_names): + """ + Parameters + ---------- + features : + must be a sequence of either integers or strings + if it contains strings 'feature_names' parameter must be defined and string ids from 'features' + must represent a subset of in 'feature_names' + + feature_names : + A sequence of string ids for features or None. + Used to get feature indices for string ids in 'features' parameter + """ + if (not isinstance(features, (Sequence, np.ndarray))) or isinstance(features, (str, bytes, bytearray)): + raise CatBoostError("feature names should be a sequence, but got " + repr(features)) + if feature_names is not None: + return [ + feature_names.index(f) if isinstance(f, STRING_TYPES) else f + for f in features + ] + else: + for f in features: + if isinstance(f, STRING_TYPES): + raise CatBoostError("features parameter contains string value '{}' but feature names " + "for a dataset are not specified".format(f)) + return features + + +def _update_params_quantize_part(params, ignored_features, per_float_feature_quantization, border_count, + feature_border_type, sparse_features_conflict_fraction, dev_efb_max_buckets, + nan_mode, input_borders, task_type, used_ram_limit, random_seed, + dev_max_subset_size_for_build_borders): + if ignored_features is not None: + params.update({ + 'ignored_features': ignored_features + }) + + if per_float_feature_quantization is not None: + params.update({ + 'per_float_feature_quantization': per_float_feature_quantization + }) + + if border_count is not None: + params.update({ + 'border_count': border_count + }) + + if feature_border_type is not None: + params.update({ + 'feature_border_type': feature_border_type + }) + + if sparse_features_conflict_fraction is not None: + params.update({ + 'sparse_features_conflict_fraction': sparse_features_conflict_fraction + }) + + if dev_efb_max_buckets is not None: + params.update({ + 'dev_efb_max_buckets': dev_efb_max_buckets + }) + + if nan_mode is not None: + params.update({ + 'nan_mode': nan_mode + }) + + if input_borders is not None: + params.update({ + 'input_borders': input_borders + }) + + if task_type is not None: + params.update({ + 'task_type': task_type + }) + + if used_ram_limit is not None: + params.update({ + 'used_ram_limit': used_ram_limit + }) + + if random_seed is not None: + params.update({ + 'random_seed': random_seed + }) + + if dev_max_subset_size_for_build_borders is not None: + params.update({ + 'dev_max_subset_size_for_build_borders': dev_max_subset_size_for_build_borders + }) + + return params + + +def plot_features_selection_loss_graph( + title, + entities_name, + entities_name_in_fields, + eliminated_entities_indices, + eliminated_entities_names, + loss_graph, + cost_graph=None +): + warn_msg = "To draw plots you should install plotly." + try: + import plotly.graph_objs as go + except ImportError as e: + warnings.warn(warn_msg) + raise ImportError(str(e)) + + indices_present = any(eliminated_entities_indices) + names_present = any(eliminated_entities_names) + names_or_indices = eliminated_entities_names if names_present else list(map(str, eliminated_entities_indices)) + loss_values = loss_graph['loss_values'] + removed_entities_cnt = loss_graph['removed_' + entities_name_in_fields + '_count'] + main_indices = loss_graph['main_indices'] + + fig = go.Figure() + fig['layout']['title'] = go.layout.Title(text=title) + loss_graph_color = 'rgb(51,160,44)' + # line with all points + fig.add_trace(go.Scatter( + x=removed_entities_cnt, + y=loss_values, + line=go.scatter.Line(color=loss_graph_color), + mode='lines+markers', + text=[''] + names_or_indices, + name='' + )) + + if len(main_indices) > 0: + # red markers for main points + fig.add_trace(go.Scatter( + x=[removed_entities_cnt[idx] for idx in main_indices], + y=[loss_values[idx] for idx in main_indices], + mode='markers', + marker=go.scatter.Marker(size=10, symbol='square'), + text=[names_or_indices[idx - 1] if idx > 0 else '' for idx in main_indices], + name='' + )) + if indices_present: + # labels with entities indices + fig.add_trace(go.Scatter( + x=removed_entities_cnt, + y=loss_values, + mode='text', + text=[''] + list(map(str, eliminated_entities_indices)), + textposition='bottom center', + textfont=dict(family='sans serif', size=18, color=loss_graph_color), + name='', + visible=False + )) + if names_present: + # labels with entities names + fig.add_trace(go.Scatter( + x=removed_entities_cnt, + y=loss_values, + mode='text', + text=[''] + eliminated_entities_names, + textfont=dict(family='sans serif', size=18, color=loss_graph_color), + textposition='bottom center', + name='', + visible=False + )) + + cost_graph_color = 'rgb(160,44,44)' + if cost_graph is not None: + # line with all points + fig.add_trace(go.Scatter( + x=removed_entities_cnt, + y=cost_graph['loss_values'], + line=go.scatter.Line(color=cost_graph_color), + mode='lines+markers', + text=[''] + eliminated_entities_names, + name='', + yaxis="y2" + )) + + # labels with entities names + fig.add_trace(go.Scatter( + x=removed_entities_cnt, + y=cost_graph['loss_values'], + mode='text', + text=[''] + eliminated_entities_names, + textfont=dict(family='sans serif', size=18, color=cost_graph_color), + textposition='bottom center', + name='', + yaxis="y2", + visible=False + )) + axis_options = dict( + gridcolor='rgb(255,255,255)', showgrid=True, showline=False, + showticklabels=True, tickcolor='rgb(127,127,127)', ticks='outside', zeroline=False + ) + fig.update_layout( + xaxis=dict(title='number of removed ' + entities_name, **axis_options), + yaxis=dict( + title='loss value', + tickfont=dict(color=loss_graph_color), + **axis_options + ) + ) + if cost_graph is not None: + fig.update_layout( + yaxis2=dict( + title='cost value', + side="right", + anchor="x", + overlaying="y", + tickfont=dict(color=cost_graph_color), + **axis_options + ) + ) + + buttons = [] + + def get_visible_arg(show_indices, show_names): + visible_arg = [True] + if len(main_indices) > 0: + visible_arg.append(True) + if indices_present: + visible_arg.append(show_indices) + if names_present: + visible_arg.append(show_names) + if cost_graph is not None: + visible_arg.append(True) + visible_arg.append(show_names) + return visible_arg + + buttons.append(dict( + label='Hide ' + entities_name, + method='update', + args=[{"visible": get_visible_arg(show_indices=False, show_names=False)}] + )) + if indices_present: + buttons.append(dict( + label='Show indices', + method='update', + args=[{"visible": get_visible_arg(show_indices=True, show_names=False)}] + )) + if names_present: + buttons.append(dict( + label='Show names', + method='update', + args=[{"visible": get_visible_arg(show_indices=False, show_names=True)}] + )) + + fig.update_layout( + updatemenus=[dict( + active=0, + buttons=buttons, + pad={"r": 10, "t": 10}, + showactive=True, + x=-0.25, + xanchor="left", + y=1.03, + yanchor="top" + )] + ) + + fig.update_layout( + showlegend=False + ) + + return fig + + +def plot_features_selection_loss_graphs(summary): + result = {} + result['features'] = plot_features_selection_loss_graph( + 'Loss by eliminated features', + 'features', + 'features', + summary['eliminated_features'], + summary['eliminated_features_names'], + summary['loss_graph'] + ) + if 'eliminated_features_tags' in summary: + result['features_tags'] = plot_features_selection_loss_graph( + 'Loss by eliminated features tags', + 'features tags', + 'features_tags', + [], + summary['eliminated_features_tags'], + summary['features_tags_loss_graph'], + cost_graph=summary['features_tags_cost_graph'] + ) + + return result + + +class Pool(_PoolBase): + """ + Pool used in CatBoost as a data structure to train model from. + """ + + def __init__( + self, + data, + label=None, + cat_features=None, + text_features=None, + embedding_features=None, + embedding_features_data=None, + column_description=None, + pairs=None, + graph=None, + delimiter='\t', + has_header=False, + ignore_csv_quoting=False, + weight=None, + group_id=None, + group_weight=None, + subgroup_id=None, + pairs_weight=None, + baseline=None, + timestamp=None, + feature_names=None, + feature_tags=None, + thread_count=-1, + log_cout=None, + log_cerr=None, + data_can_be_none=False + ): + """ + Pool is an internal data structure that is used by CatBoost. + You can construct Pool from list, numpy.ndarray, pandas.DataFrame, pandas.Series. + + Parameters + ---------- + data : list or numpy.ndarray or pandas.DataFrame or pandas.Series or polars.DataFrame or FeaturesData or string or os.PathLike + Data source of Pool. + If list or numpy.ndarrays or pandas.DataFrame or pandas.Series or polars.DataFrame, + giving 2 dimensional array like data. + If FeaturesData - see FeaturesData description for details, 'cat_features' and 'feature_names' + parameters must be equal to None in this case + If string or os.PathLike, giving the path to the file with data in catboost format. + If string starts with "quantized://", the file has to contain quantized dataset saved with Pool.save(). + + label : list or numpy.ndarrays or pandas.DataFrame or pandas.Series or polars.DataFrame or polars.Series, optional (default=None) + Labels data. + If not None, can be a single- or two- dimensional array with either: + - numerical values - for regression (including multiregression), ranking and binary classification problems + - class labels (boolean, integer or string) - for classification (including multiclassification) problems + If `data` parameter points to a file, Label data is loaded from it as well. This parameter must + be None in this case. + + cat_features : list or numpy.ndarray, optional (default=None) + If not None, giving the list of Categ features indices or names. + If it contains feature names, Pool's feature names must be defined: either by passing 'feature_names' + parameter or if data is pandas.DataFrame (feature names are initialized from it's column names) + Must be None if 'data' parameter has FeaturesData type + + text_features : list or numpy.ndarray, optional (default=None) + If not None, giving the list of Text features indices or names. + If it contains feature names, Pool's feature names must be defined: either by passing 'feature_names' + parameter or if data is pandas.DataFrame (feature names are initialized from it's column names) + Must be None if 'data' parameter has FeaturesData type + + embedding_features : list or numpy.ndarray, optional (default=None) + If not None, giving the list of Embedding features indices or names. + If it contains feature names, Pool's feature names must be defined: either by passing 'feature_names' + parameter or if data is pandas.DataFrame (feature names are initialized from it's column names) + Must be None if 'data' parameter has FeaturesData type + + embedding_features_data : list or dict, optional (default=None) + If not None, giving the data of Embedding features (instead of data in main 'data' parameter). + If list - list containing 2d arrays (lists or numpy.ndarrays or scipy.sparse.spmatrix) with [n_data_size x embedding_size] elements + If dict - dict containing 2d arrays (lists or numpy.ndarrays or scipy.sparse.spmatrix) with [n_data_size x embedding_size] elements + Dict keys must be the same as specified in 'embedding_features' parameter + + column_description : string or os.PathLike, optional (default=None) + ColumnsDescription parameter. + There are several columns description types: Label, Categ, Num, Auxiliary, DocId, Weight, Baseline, GroupId, Timestamp. + All columns are Num as default, it's not necessary to specify + this type of columns. Default Label column index is 0 (zero). + If None, Label column is 0 (zero) as default, all data columns are Num as default. + If string or os.PathLike, giving the path to the file with ColumnsDescription in column_description format. + + pairs : list or numpy.ndarray or pandas.DataFrame or polars.DataFrame or string or os.PathLike + The pairs description. + If list or numpy.ndarrays or pandas.DataFrame or polars.DataFrame, giving 2 dimensional. + The shape should be Nx2, where N is the pairs' count. The first element of the pair is + the index of winner object in the training set. The second element of the pair is + the index of loser object in the training set. + If string or os.PathLike, giving the path to the file with pairs description. + + graph: list or numpy.ndarray or pandas.DataFrame or polars.DataFrame or string or os.PathLike + The graph description. + ... + + delimiter : string, optional (default='\t') + Delimiter to use for separate features in file. + Should be only one symbol, otherwise would be taken only the first character of the string. + + has_header : bool optional (default=False) + If True, read column names from first line. + + ignore_csv_quoting : bool optional (default=False) + If True ignore quoting '"'. + + weight : list or numpy.ndarray or polars.Series, optional (default=None) + Weight for each instance. + If not None, giving 1 dimensional array like data. + + group_id : list or numpy.ndarray or polars.Series, optional (default=None) + group id for each instance. + If not None, giving 1 dimensional array like data. + + group_weight : list or numpy.ndarray or polars.Series, optional (default=None) + Group weight for each instance. + If not None, giving 1 dimensional array like data. + + subgroup_id : list or numpy.ndarray or polars.Series, optional (default=None) + subgroup id for each instance. + If not None, giving 1 dimensional array like data. + + pairs_weight : list or numpy.ndarray or polars.Series, optional (default=None) + Weight for each pair. + If not None, giving 1 dimensional array like pairs. + + baseline : list or numpy.ndarray or polars.DataFrame or polars.Series, optional (default=None) + Baseline for each instance. + If not None, giving 2 dimensional array like data. + + timestamp: list or numpy.ndarray or polars.Series, optional (default=None) + Timestamp for each instance. + Should be a non-negative integer. + Useful for sorting a learning dataset by this field during training. + + feature_names : list or string or os.PathLike, optional (default=None) + If list - list of names for each given data_feature. + If string or os.PathLike - path with scheme for feature names data to load. + If this parameter is None and 'data' is pandas.DataFrame feature names will be initialized + from DataFrame's column names. + Must be None if 'data' parameter has FeaturesData type + + feature_tags : json, optional (default=None) + Format: + {'tag1': + { + 'features': [], + 'cost': + } + 'tag2': + { + ... + } + ... + } + + thread_count : int, optional (default=-1) + Thread count for data processing. + If -1, then the number of threads is set to the number of CPU cores. + + log_cout: output stream or callback for logging (default=None) + If None is specified, sys.stdout is used + + log_cerr: error stream or callback for logging (default=None) + If None is specified, sys.stderr is used + + """ + with log_fixup(log_cout, log_cerr): + if data is not None: + self._check_data_type(data) + self._check_data_empty(data) + if pairs is not None and isinstance(data, PATH_TYPES) != isinstance(pairs, PATH_TYPES): + raise CatBoostError("data and pairs parameters should be the same types.") + if graph is not None and isinstance(data, PATH_TYPES) != isinstance(graph, PATH_TYPES): + raise CatBoostError("data and graph parameters should be the same types.") + if column_description is not None and not isinstance(data, PATH_TYPES): + raise CatBoostError("data should be the string or os.PathLike type if column_description parameter is specified.") + if isinstance(data, PATH_TYPES): + if any(v is not None for v in [cat_features, text_features, embedding_features, embedding_features_data, weight, group_id, group_weight, + subgroup_id, pairs_weight, baseline, label]): + raise CatBoostError( + "cat_features, text_features, embedding_features, embedding_features_data, weight, group_id, group_weight, subgroup_id, pairs_weight, " + "baseline, label should have the None type when the pool is read from the file." + ) + if (feature_names is not None) and (not isinstance(feature_names, PATH_TYPES)): + raise CatBoostError( + "feature_names should have None or string or os.PathLike type when the pool is read from the file." + ) + self._read(data, column_description, pairs, graph, feature_names, delimiter, has_header, ignore_csv_quoting, thread_count) + else: + if isinstance(data, FeaturesData): + if any(v is not None for v in [cat_features, text_features, embedding_features, embedding_features_data, feature_names]): + raise CatBoostError( + "cat_features, text_features, embedding_features, embedding_features_data, feature_names should have the None type" + " when 'data' parameter has FeaturesData type" + ) + elif isinstance(data, np.ndarray): + if (data.dtype.kind == 'f') and (cat_features is not None) and (len(cat_features) > 0): + raise CatBoostError( + "'data' is numpy array of floating point numerical type, it means no categorical features," + " but 'cat_features' parameter specifies nonzero number of categorical features" + ) + if (data.dtype.kind == 'f') and (text_features is not None) and (len(text_features) > 0): + raise CatBoostError( + "'data' is numpy array of floating point numerical type, it means no text features," + " but 'text_features' parameter specifies nonzero number of text features" + ) + if (data.dtype.kind != 'O') and (embedding_features is not None) and (len(embedding_features) > 0): + if embedding_features_data is None: + raise CatBoostError( + "'data' is numpy array of non-object type, it means no embedding features," + " but 'embedding_features' parameter specifies nonzero number of embedding features" + ) + elif isinstance(data, scipy.sparse.spmatrix): + if (data.dtype.kind == 'f') and (cat_features is not None) and (len(cat_features) > 0): + raise CatBoostError( + "'data' is scipy.sparse.spmatrix of floating point numerical type, it means no categorical features," + " but 'cat_features' parameter specifies nonzero number of categorical features" + ) + if (text_features is not None) and (len(text_features) > 0): + raise CatBoostError( + "'data' is scipy.sparse.spmatrix, it means no text features," + " but 'text_features' parameter specifies nonzero number of text features" + ) + if (embedding_features is not None) and (len(embedding_features) > 0) and (embedding_features_data is None): + raise CatBoostError( + "'data' is scipy.sparse.spmatrix and 'embedding_features_data' is None, it means no embedding features," + " but 'embedding_features' parameter specifies nonzero number of embedding features" + ) + + if embedding_features_data is not None: + if embedding_features is None: + raise CatBoostError( + "'embedding_features_data' is not None, but 'embedding_features' parameter is not specified" + ) + if isinstance(embedding_features_data, list): + if len(embedding_features) != len(embedding_features_data): + raise CatBoostError( + "'embedding_features_data' and 'embedding_features' contain different numbers of features" + ) + elif isinstance(embedding_features_data, dict): + if set(embedding_features) != set(embedding_features_data.keys()): + raise CatBoostError( + "keys of 'embedding_features_data' dict do not correspond to 'embedding_features'" + ) + else: + raise CatBoostError( + "'embedding_features_data' must have either 'list' or 'dict' type" + ) + + if isinstance(feature_names, PATH_TYPES): + raise CatBoostError( + "feature_names must be None or have non-string type when the pool is created from " + "python objects." + ) + + self._init(data, label, cat_features, text_features, embedding_features, embedding_features_data, pairs, graph, weight, + group_id, group_weight, subgroup_id, pairs_weight, baseline, timestamp, feature_names, feature_tags, thread_count) + elif not data_can_be_none: + raise CatBoostError("'data' parameter can't be None") + super(Pool, self).__init__() + + def _check_files(self, data, column_description, pairs, graph): + """ + Check files existence. + """ + for item, name in zip((data, column_description, pairs, graph), ('data', 'column_description', 'pairs', 'graph')): + item = fspath(item) + if item is None: + continue + if item.find('://') == -1 and not os.path.isfile(item): + raise CatBoostError("Invalid {} path='{}': file does not exist.".format(name, item)) + + def _check_delimiter(self, delimiter): + if not isinstance(delimiter, STRING_TYPES): + raise CatBoostError("Invalid delimiter type={} : must be str.".format(type(delimiter))) + if len(delimiter) < 1: + raise CatBoostError("Invalid delimiter length={} : must be > 0.".format(len(delimiter))) + + def _check_column_description_type(self, column_description): + """ + Check type of column_description parameter. + """ + if not isinstance(column_description, PATH_TYPES): + raise CatBoostError("Invalid column_description type={}: must be str or os.PathLike.".format(type(column_description))) + + def _check_string_feature_type(self, features, features_name): + """ + Check type of cat_feature parameter. + """ + if not isinstance(features, (list, np.ndarray)): + raise CatBoostError("Invalid {} type={}: must be list or numpy.ndarray.".format(features_name, type(features))) + + def _check_string_feature_value(self, features, features_count, features_name): + """ + Check values in cat_feature parameter. Must be int indices. + """ + for indx, feature in enumerate(features): + if not isinstance(feature, INTEGER_TYPES): + raise CatBoostError("Invalid {}[{}] = {} value type={}: must be an integer.".format(features_name, indx, feature, type(feature))) + if feature >= features_count: + raise CatBoostError("Invalid {}[{}] = {} value: index must be < {}.".format(features_name, indx, feature, features_count)) + + def _check_pairs_or_graph_type(self, data, data_name: str): + """ + Check type of pairs or graph parameter. + """ + if not isinstance(data, (list, np.ndarray, pd.DataFrame, pl.DataFrame)): + raise CatBoostError(f"Invalid {data_name} type={type(data)}: must be list, numpy.ndarray, pandas.DataFrame or polars.DataFrame.") + + def _check_pairs_or_graph_value(self, data, data_name: str): + """ + Check values in pairs or graph parameter. Must be int indices. + """ + for i, e in enumerate(data): + if (len(e) != 2): + raise CatBoostError(f"Length of {data_name}[{i}] isn't equal to 2.") + for j, index in enumerate(e): + if not isinstance(index, INTEGER_TYPES): + raise CatBoostError(f"Invalid {data_name}[{i}][{j}] = '{index}' value type={type(index)}: must be an integer.") + + def _check_data_type(self, data): + """ + Check type of data. + """ + if not isinstance(data, (PATH_TYPES, ARRAY_TYPES, SPARSE_MATRIX_TYPES, FeaturesData)): + raise CatBoostError( + ("Invalid data type={}: must be list, numpy.ndarray, pandas.DataFrame, pandas.Series, polars.DataFrame, catboost.FeaturesData, " + + "scipy.sparse matrix or a path to a file (str or os.PathLike)").format(type(data)) + ) + + def _check_data_empty(self, data): + """ + Check that data is not empty (0 objects is ok). + note: already checked if data is FeatureType, so no need to check again + """ + + if isinstance(data, PATH_TYPES): + if not data: + raise CatBoostError("Features filename is empty.") + elif isinstance(data, (ARRAY_TYPES, SPARSE_MATRIX_TYPES)): + if isinstance(data, list): + data_shape = np.shape(np.asarray(data, dtype=object)) + else: + data_shape = np.shape(data) + if len(data_shape) == 1 and data_shape[0] > 0: + if isinstance(data[0], Iterable): + data_shape = tuple(data_shape + tuple([len(data[0])])) + else: + data_shape = tuple(data_shape + tuple([1])) + if len(data_shape) not in (2, 3): + raise CatBoostError("Input data has invalid shape: {}. Must be 2 or 3 dimensional".format(data_shape)) + if data_shape[1] == 0: + raise CatBoostError("Input data must have at least one feature") + + def _check_label_type(self, label): + """ + Check type of label. + """ + if not isinstance(label, ARRAY_TYPES): + raise CatBoostError("Invalid label type={}: must be array like.".format(type(label))) + + def _check_label_empty(self, label): + """ + Check label is not empty. + """ + if len(label) == 0: + raise CatBoostError("Labels variable is empty.") + + def _check_label_shape(self, label, samples_count): + """ + Check label length and dimension. + """ + if len(label) != samples_count: + raise CatBoostError("Length of label={} and length of data={} is different.".format(len(label), samples_count)) + + def _check_baseline_type(self, baseline): + """ + Check type of baseline parameter. + """ + if not isinstance(baseline, ARRAY_TYPES): + raise CatBoostError("Invalid baseline type={}: must be array like.".format(type(baseline))) + + def _check_baseline_shape(self, baseline, samples_count): + """ + Check baseline length and dimension. + """ + if len(baseline) != samples_count: + raise CatBoostError("Length of baseline={} and length of data={} are different.".format(len(baseline), samples_count)) + if not isinstance(baseline[0], Iterable) or isinstance(baseline[0], STRING_TYPES): + raise CatBoostError("Baseline must be 2 dimensional data, 1 column for each class.") + + if isinstance(baseline, pl.DataFrame): + for i, column in enumerate(baseline.iter_columns()): + try: + if column.has_nulls(): + raise CatBoostError("Baseline must not contain null values.") + dtype = column.dtype + if (not dtype.is_numeric()) and (dtype != pl.Object): + raise CatBoostError(f'Non-numeric dtype: {dtype}') + except Exception as e: + raise CatBoostError(f"Invalid baseline column {i}") from e + else: + dtype = np.array(baseline).dtype + try: + if not (np.issubdtype(dtype, np.integer) or np.issubdtype(dtype, np.floating) or (dtype == np.dtype('object'))): + raise CatBoostError() + except CatBoostError: + raise CatBoostError("Invalid baseline value type={}: must be float or int.".format(dtype)) + + def _check_weight_type(self, weight): + """ + Check type of weight parameter. + """ + if not isinstance(weight, ARRAY_TYPES): + raise CatBoostError("Invalid weight type={}: must be array like.".format(type(weight))) + + def _check_weight_shape(self, weight, samples_count): + """ + Check weight length. + """ + if len(weight) != samples_count: + raise CatBoostError("Length of weight={} and length of data={} are different.".format(len(weight), samples_count)) + if not isinstance(weight[0], (INTEGER_TYPES, FLOAT_TYPES)): + raise CatBoostError( + "Invalid weight value type={}: must be 1 dimensional data with elements of an integer or a floating-point type.".format(type(weight[0])) + ) + + def _check_group_id_type(self, group_id): + """ + Check type of group_id parameter. + """ + if not isinstance(group_id, ARRAY_TYPES): + raise CatBoostError("Invalid group_id type={}: must be array like.".format(type(group_id))) + + def _check_group_id_shape(self, group_id, samples_count): + """ + Check group_id length. + """ + if len(group_id) != samples_count: + raise CatBoostError("Length of group_id={} and length of data={} are different.".format(len(group_id), samples_count)) + + def _check_group_weight_type(self, group_weight): + """ + Check type of group_weight parameter. + """ + if not isinstance(group_weight, ARRAY_TYPES): + raise CatBoostError("Invalid group_weight type={}: must be array like.".format(type(group_weight))) + + def _check_group_weight_shape(self, group_weight, samples_count): + """ + Check group_weight length. + """ + if len(group_weight) != samples_count: + raise CatBoostError("Length of group_weight={} and length of data={} are different.".format(len(group_weight), samples_count)) + if not isinstance(group_weight[0], (FLOAT_TYPES)): + raise CatBoostError("Invalid group_weight value type={}: must be 1 dimensional data with elements of a floating-point type.".format(type(group_weight[0]))) + + def _check_subgroup_id_type(self, subgroup_id): + """ + Check type of subgroup_id parameter. + """ + if not isinstance(subgroup_id, ARRAY_TYPES): + raise CatBoostError("Invalid subgroup_id type={}: must be array like.".format(type(subgroup_id))) + + def _check_subgroup_id_shape(self, subgroup_id, samples_count): + """ + Check subgroup_id length. + """ + if len(subgroup_id) != samples_count: + raise CatBoostError("Length of subgroup_id={} and length of data={} are different.".format(len(subgroup_id), samples_count)) + + def _check_timestamp_type(self, timestamp): + """ + Check type of timestamp parameter. + """ + if not isinstance(timestamp, ARRAY_TYPES): + raise CatBoostError("Invalid timestamp type={}: must be array like.".format(type(timestamp))) + + def _check_transform_tags(self, tags, feature_names): + if not isinstance(tags, dict): + raise CatBoostError("Invalid feature_tags type={}: must be dict like.".format(type(tags))) + for tag_name, tag_features in tags.items(): + if not isinstance(tag_features, dict): + raise CatBoostError("Invalid type of value in feature_tags by key {}, value type is {}: must be dict like.".format(tag_name, type(tag_features))) + if 'features' not in tag_features: + raise CatBoostError("Invalid value in feature_tags by key {}, key 'features' is needed.".format(tag_name)) + if not isinstance(tag_features['features'], ARRAY_TYPES): + raise CatBoostError("Invalid type of value in feature_tags by key {}, value type of features is {}: must be array like.".format(tag_name, type(tag_features['features']))) + if 'cost' not in tag_features: + tag_features['cost'] = 1.0 + else: + if not isinstance(tag_features['cost'], (INTEGER_TYPES, str)): + raise CatBoostError("Invalid type of value in feature_tags by key {}, value type of cost is {}: must be an integer.".format(tag_name, type(tag_features['cost']))) + tag_features['cost'] = int(tag_features['cost']) + + for idx in range(len(tag_features['features'])): + if isinstance(tag_features['features'][idx], INTEGER_TYPES): + pass + elif isinstance(tag_features['features'][idx], str) and feature_names is not None: + try: + feature_id = feature_names.index(tag_features['features'][idx]) + except ValueError: + raise CatBoostError("Unknown feature in tag {}: {}".format(tag_name, tag_features['features'][idx])) + tag_features['features'][idx] = feature_id + else: + raise CatBoostError("Invalid type of feature in tag {}, value type is {}: must be an integer or a feature name.".format(tag_name, type(tag_features['features'][idx]))) + return tags + + def _check_timestamp_shape(self, timestamp, samples_count): + """ + Check timestamp length. + """ + if len(timestamp) != samples_count: + raise CatBoostError("Length of timestamp={} and length of data={} are different.".format(len(timestamp), samples_count)) + + def _check_feature_names(self, feature_names, num_col=None): + if num_col is None: + num_col = self.num_col() + if not isinstance(feature_names, Sequence): + raise CatBoostError("Invalid feature_names type={} : must be Sequence".format(type(feature_names))) + if len(feature_names) != num_col: + raise CatBoostError("Invalid length of feature_names={} : must be equal to the number of columns in data={}".format(len(feature_names), num_col)) + + def _check_thread_count(self, thread_count): + if not isinstance(thread_count, INTEGER_TYPES): + raise CatBoostError("Invalid thread_count type={} : must be an integer".format(type(thread_count))) + + def slice(self, rindex): + if not isinstance(rindex, ARRAY_TYPES): + raise CatBoostError("Invalid rindex type={} : must be array like".format(type(rindex))) + slicedPool = Pool(None, data_can_be_none=True) + slicedPool._take_slice(self, rindex) + return slicedPool + + def train_eval_split(self, has_time, is_classification, eval_fraction, save_eval_pool): + train_pool = Pool(None, data_can_be_none=True) + eval_pool = Pool(None, data_can_be_none=True) + self._train_eval_split(train_pool, eval_pool, has_time, is_classification, eval_fraction, save_eval_pool) + return train_pool, eval_pool + + def _check_and_prepare_pairs_or_graph(self, data, data_name: str): + self._check_pairs_or_graph_type(data, data_name) + if isinstance(data, pd.DataFrame): + data = data.values + elif isinstance(data, pl.DataFrame): + data = data.to_numpy() + self._check_pairs_or_graph_value(data, data_name) + return data + + def set_pairs(self, pairs): + self._set_pairs(self._check_and_prepare_pairs_or_graph(pairs, 'pairs')) + return self + + def set_feature_names(self, feature_names): + self._check_feature_names(feature_names) + self._set_feature_names(feature_names) + return self + + def _check_and_prepare_baseline(self, baseline, samples_count: int): + self._check_baseline_type(baseline) + if isinstance(baseline, pl.Series): + baseline = pl.DataFrame([baseline]) + elif not isinstance(baseline, pl.DataFrame): + baseline = self._if_pandas_to_numpy(baseline) + baseline = np.reshape(baseline, (samples_count, -1)) + self._check_baseline_shape(baseline, samples_count) + return baseline + + def set_baseline(self, baseline): + self._set_baseline(self._check_and_prepare_baseline(baseline, self.num_row())) + return self + + def set_weight(self, weight): + self._check_weight_type(weight) + weight = self._if_pandas_to_numpy(weight) + self._check_weight_shape(weight, self.num_row()) + self._set_weight(weight) + return self + + def set_group_id(self, group_id): + self._check_group_id_type(group_id) + group_id = self._if_pandas_to_numpy(group_id) + self._check_group_id_shape(group_id, self.num_row()) + self._set_group_id(group_id) + return self + + def set_group_weight(self, group_weight): + self._check_group_weight_type(group_weight) + group_weight = self._if_pandas_to_numpy(group_weight) + self._check_group_weight_shape(group_weight, self.num_row()) + self._set_group_weight(group_weight) + return self + + def set_subgroup_id(self, subgroup_id): + self._check_subgroup_id_type(subgroup_id) + subgroup_id = self._if_pandas_to_numpy(subgroup_id) + self._check_subgroup_id_shape(subgroup_id, self.num_row()) + self._set_subgroup_id(subgroup_id) + return self + + def set_pairs_weight(self, pairs_weight): + self._check_weight_type(pairs_weight) + pairs_weight = self._if_pandas_to_numpy(pairs_weight) + self._check_weight_shape(pairs_weight, self.num_pairs()) + self._set_pairs_weight(pairs_weight) + return self + + def set_timestamp(self, timestamp): + self._check_timestamp_type(timestamp) + timestamp = self._if_pandas_to_numpy(timestamp) + self._check_timestamp_shape(timestamp, self.num_row()) + self._set_timestamp(timestamp) + return self + + def save(self, fname): + """ + Save the quantized pool to a file. + + Parameters + ---------- + fname : string or os.PathLike + Output file name. + """ + if not self.is_quantized(): + raise CatBoostError('Pool is not quantized') + + if not isinstance(fname, PATH_TYPES): + raise CatBoostError("Invalid fname type={}: must be str or os.PathLike.".format(type(fname))) + + self._save(fname) + + def quantize(self, ignored_features=None, per_float_feature_quantization=None, border_count=None, + max_bin=None, feature_border_type=None, sparse_features_conflict_fraction=None, + nan_mode=None, input_borders=None, task_type=None, used_ram_limit=None, random_seed=None, **kwargs): + """ + Quantize this pool + + Parameters + ---------- + pool : catboost.Pool + Dataset to quantize. + + ignored_features : list, [default=None] + Indices or names of features that should be excluded when training. + + per_float_feature_quantization : list of strings, [default=None] + List of float binarization descriptions. + Format : described in documentation on catboost.ai + Example 1: ['0:1024'] means that feature 0 will have 1024 borders. + Example 2: ['0:border_count=1024', '1:border_count=1024', ...] means that two first features have 1024 borders. + Example 3: ['0:nan_mode=Forbidden,border_count=32,border_type=GreedyLogSum', + '1:nan_mode=Forbidden,border_count=32,border_type=GreedyLogSum'] - defines more quantization properties for first two features. + + border_count : int, [default = 254 for training on CPU or 128 for training on GPU] + The number of partitions in numeric features binarization. Used in the preliminary calculation. + range: [1,65535] on CPU, [1,255] on GPU + + max_bin : float, synonym for border_count. + + feature_border_type : string, [default='GreedyLogSum'] + The binarization mode in numeric features binarization. Used in the preliminary calculation. + Possible values: + - 'Median' + - 'Uniform' + - 'UniformAndQuantiles' + - 'GreedyLogSum' + - 'MaxLogSum' + - 'MinEntropy' + + sparse_features_conflict_fraction : float, [default=0.0] + CPU only. Maximum allowed fraction of conflicting non-default values for features in exclusive features bundle. + Should be a real value in [0, 1) interval. + + nan_mode : string, [default=None] + Way to process missing values for numeric features. + Possible values: + - 'Forbidden' - raises an exception if there is a missing value for a numeric feature in a dataset. + - 'Min' - each missing value will be processed as the minimum numerical value. + - 'Max' - each missing value will be processed as the maximum numerical value. + If None, then nan_mode=Min. + + input_borders : string or os.PathLike, [default=None] + input file with borders used in numeric features binarization. + + task_type : string, [default=None] + The calcer type that will be used to train the model after quantization. + Possible values: + - 'CPU' + - 'GPU' + + used_ram_limit=None + + random_seed : int, [default=None] + The random seed used for data sampling. + If None, 0 is used. + """ + if self.is_quantized(): + raise CatBoostError('Pool is already quantized') + + params = {} + + if border_count is None: + border_count = max_bin + + dev_efb_max_buckets = kwargs.pop('dev_efb_max_buckets', None) + dev_max_subset_size_for_build_borders = kwargs.pop('dev_max_subset_size_for_build_borders', None) + + if kwargs: + raise CatBoostError("got unexpected keyword arguments: {}".format(kwargs.keys())) + + _update_params_quantize_part(params, ignored_features, per_float_feature_quantization, border_count, + feature_border_type, sparse_features_conflict_fraction, dev_efb_max_buckets, + nan_mode, input_borders, task_type, used_ram_limit, random_seed, + dev_max_subset_size_for_build_borders) + + self._quantize(params) + + def _if_pandas_to_numpy(self, array): + if isinstance(array, pd.Series): + array = array.values + if isinstance(array, pd.DataFrame): + array = np.transpose(array.values)[0] + return array + + def _label_if_pandas_to_numpy(self, label): + if isinstance(label, pd.Series): + label = label.values + if isinstance(label, pd.DataFrame): + label = label.values + return label + + def _read( + self, + pool_file, + column_description, + pairs, + graph, + feature_names_path, + delimiter, + has_header, + ignore_csv_quoting, + thread_count, + quantization_params=None, + log_cout=None, + log_cerr=None + ): + """ + Read Pool from file. + """ + + with log_fixup(log_cout, log_cerr): + self._check_files(pool_file, column_description, pairs, graph) + self._check_delimiter(delimiter) + if column_description is None: + column_description = '' + else: + self._check_column_description_type(column_description) + for item in [pairs, feature_names_path, graph]: + if item is None: + item = '' + self._check_thread_count(thread_count) + self._read_pool( + pool_file, + column_description, + pairs, + graph, + feature_names_path, + delimiter[0], + has_header, + ignore_csv_quoting, + thread_count, + quantization_params + ) + + def _infer_feature_names(self, data_as_data_frame, embedding_features_data=None, embedding_features=None): + non_embedding_data_feature_names = list(data_as_data_frame.columns) + + if embedding_features_data is not None: + if isinstance(embedding_features_data, dict): + embedding_feature_names = list(embedding_features_data.keys()) + if embedding_features is not None: + if set(embedding_features) != set(embedding_feature_names): + raise CatBoostError('keys of embedding_features_data and embedding_features are different') + return non_embedding_data_feature_names + embedding_feature_names + else: + if embedding_features is None: + raise CatBoostError('embedding_features is not specified but embedding_features_data without feature names is present') + if not all([isinstance(embedding_feature_id, INTEGER_TYPES) for embedding_feature_id in embedding_features]): + raise CatBoostError('embedding_features contain feature names but embedding_features_data without feature names is present') + + embedding_features_set = set(embedding_features) + + feature_names = [] + non_embedding_feature_idx = 0 + for feature_idx in range(len(non_embedding_data_feature_names) + len(embedding_features)): + if feature_idx in embedding_features_set: + feature_names.append('_embedding_feature_%i' % feature_idx) + else: + feature_names.append(non_embedding_data_feature_names[non_embedding_feature_idx]) + non_embedding_feature_idx += 1 + + return feature_names + else: + return non_embedding_data_feature_names + + def _init( + self, + data, + label, + cat_features, + text_features, + embedding_features, + embedding_features_data, + pairs, graph, weight, + group_id, + group_weight, + subgroup_id, + pairs_weight, + baseline, + timestamp, + feature_names, + feature_tags, + thread_count + ): + """ + Initialize Pool from array like data. + """ + if isinstance(data, (pd.DataFrame, pl.DataFrame)): + if feature_names is None: + feature_names = self._infer_feature_names(data, embedding_features_data, embedding_features) + if isinstance(data, pd.Series): + data = data.values.tolist() + if isinstance(data, FeaturesData): + samples_count = data.get_object_count() + features_count = data.get_feature_count() + elif data is not None: + if isinstance(data, list): + data = np.asarray(data, dtype=object) + if len(np.shape(data)) == 1: + data = np.expand_dims(data, 1) + samples_count, features_count = np.shape(data)[:2] + if embedding_features_data is not None: + features_count += len(embedding_features_data) + if isinstance(embedding_features_data, dict): + embedding_features_data_values = list(embedding_features_data.values()) + else: + embedding_features_data_values = embedding_features_data + for embedding_feature_data in embedding_features_data_values: + if len(embedding_feature_data) != samples_count: + raise CatBoostError( + "samples count in 'embeddings_features_data' does not correspond to samples count in main data" + ) + pairs_len = 0 + if label is not None: + self._check_label_type(label) + self._check_label_empty(label) + if isinstance(label, pl.Series): + label = pl.DataFrame([label]) + elif not isinstance(label, pl.DataFrame): + label = self._label_if_pandas_to_numpy(label) + if len(np.shape(label)) == 1: + label = np.expand_dims(label, 1) + self._check_label_shape(label, samples_count) + if feature_names is not None: + self._check_feature_names(feature_names, features_count) + if cat_features is not None: + cat_features = _get_features_indices(cat_features, feature_names) + self._check_string_feature_type(cat_features, 'cat_features') + self._check_string_feature_value(cat_features, features_count, 'cat_features') + if text_features is not None: + text_features = _get_features_indices(text_features, feature_names) + self._check_string_feature_type(text_features, 'text_features') + self._check_string_feature_value(text_features, features_count, 'text_features') + if embedding_features is not None: + embedding_features = _get_features_indices(embedding_features, feature_names) + self._check_string_feature_type(embedding_features, 'embedding_features') + self._check_string_feature_value(embedding_features, features_count, 'embedding_features') + if pairs is not None: + pairs = self._check_and_prepare_pairs_or_graph(pairs, 'pairs') + pairs_len = np.shape(pairs)[0] + if graph is not None: + graph = self._check_and_prepare_pairs_or_graph(graph, 'graph') + if weight is not None: + self._check_weight_type(weight) + weight = self._if_pandas_to_numpy(weight) + self._check_weight_shape(weight, samples_count) + if group_id is not None: + self._check_group_id_type(group_id) + group_id = self._if_pandas_to_numpy(group_id) + self._check_group_id_shape(group_id, samples_count) + if group_weight is not None: + self._check_group_weight_type(group_weight) + group_weight = self._if_pandas_to_numpy(group_weight) + self._check_group_weight_shape(group_weight, samples_count) + if subgroup_id is not None: + self._check_subgroup_id_type(subgroup_id) + subgroup_id = self._if_pandas_to_numpy(subgroup_id) + self._check_subgroup_id_shape(subgroup_id, samples_count) + if pairs_weight is not None: + self._check_weight_type(pairs_weight) + pairs_weight = self._if_pandas_to_numpy(pairs_weight) + self._check_weight_shape(pairs_weight, pairs_len) + if baseline is not None: + baseline = self._check_and_prepare_baseline(baseline, samples_count) + if timestamp is not None: + self._check_timestamp_type(timestamp) + timestamp = self._if_pandas_to_numpy(timestamp) + self._check_timestamp_shape(timestamp, samples_count) + if feature_tags is not None: + feature_tags = self._check_transform_tags(feature_tags, feature_names) + self._init_pool(data, label, cat_features, text_features, embedding_features, embedding_features_data, pairs, graph, weight, + group_id, group_weight, subgroup_id, pairs_weight, baseline, timestamp, feature_names, feature_tags, thread_count) + + +def _build_train_pool(X, y, cat_features, text_features, embedding_features, pairs, graph, sample_weight, group_id, group_weight, subgroup_id, pairs_weight, baseline, column_description): + train_pool = None + if isinstance(X, Pool): + train_pool = X + if any(v is not None for v in [cat_features, text_features, embedding_features, sample_weight, group_id, group_weight, subgroup_id, pairs_weight, baseline]): + raise CatBoostError( + "cat_features, text_features, embedding_features, sample_weight, group_id, group_weight, subgroup_id," + " pairs_weight, baseline should have the None type when X has catboost.Pool type." + ) + if (not X.has_label()) and X.num_pairs() == 0: + raise CatBoostError("Label in X has not been initialized.") + if y is not None: + raise CatBoostError("Incorrect value of y: X is catboost.Pool object, y must be initialized inside catboost.Pool.") + elif isinstance(X, PATH_TYPES): + train_pool = Pool(data=X, pairs=pairs, graph=graph, column_description=column_description) + else: + if y is None: + raise CatBoostError("y has not initialized in fit(): X is not catboost.Pool object, y must be not None in fit().") + train_pool = Pool(X, y, cat_features=cat_features, text_features=text_features, embedding_features=embedding_features, pairs=pairs, graph=graph, weight=sample_weight, group_id=group_id, + group_weight=group_weight, subgroup_id=subgroup_id, pairs_weight=pairs_weight, baseline=baseline) + return train_pool + + +def _clear_training_files(train_dir): + for filename in ['catboost_training.json']: + path = os.path.join(train_dir, filename) + if os.path.exists(path): + os.remove(path) + + +def _get_train_dir(params): + return params.get('train_dir', 'catboost_info') + + +def _get_catboost_widget(train_dirs): + for train_dir in train_dirs: + _clear_training_files(train_dir) + try: + from .widget import MetricVisualizer + return MetricVisualizer(train_dirs) + except ImportError as e: + warnings.warn("To draw plots in fit() method you should install ipywidgets and ipython") + raise ImportError(str(e)) + + +@contextmanager +def plot_wrapper(plot, plot_file, plot_title, train_dirs): + if plot: + widget = _get_catboost_widget(train_dirs) + widget._run_update() + try: + yield + finally: + if plot: + widget._stop_update() + if plot_file is not None: + OfflineMetricVisualizer(train_dirs).save_to_file(plot_title, plot_file) + + +# the first element of the synonyms list is the canonical name +def _process_synonyms_group(synonyms, params): + assert len(synonyms) > 1, 'there should be more than one synonym' + + value = None + for synonym in synonyms: + if synonym in params: + if value is not None: + raise CatBoostError('only one of the parameters ' + (', '.join(synonyms)) + ' should be initialized.') + value = params[synonym] + del params[synonym] + + if value is not None: + params[synonyms[0]] = value + + +def _process_synonyms_groups(params): + _process_synonyms_group(['learning_rate', 'eta'], params) + _process_synonyms_group(['border_count', 'max_bin'], params) + _process_synonyms_group(['depth', 'max_depth'], params) + _process_synonyms_group(['rsm', 'colsample_bylevel'], params) + _process_synonyms_group(['random_seed', 'random_state'], params) + _process_synonyms_group(['l2_leaf_reg', 'reg_lambda'], params) + _process_synonyms_group(['iterations', 'n_estimators', 'num_boost_round', 'num_trees'], params) + _process_synonyms_group(['od_wait', 'early_stopping_rounds'], params) + _process_synonyms_group(['custom_metric', 'custom_loss'], params) + _process_synonyms_group(['max_leaves', 'num_leaves'], params) + _process_synonyms_group(['min_data_in_leaf', 'min_child_samples'], params) + + +def _process_synonyms(params): + if 'objective' in params: + params['loss_function'] = params['objective'] + del params['objective'] + + if 'scale_pos_weight' in params: + if 'loss_function' in params and params['loss_function'] != 'Logloss': + raise CatBoostError('scale_pos_weight is supported only for binary classification Logloss loss') + if 'class_weights' in params or 'auto_class_weights' in params: + raise CatBoostError('only one of the parameters scale_pos_weight, class_weights, auto_class_weights should be initialized.') + params['class_weights'] = [1.0, params['scale_pos_weight']] + del params['scale_pos_weight'] + if ('class_weights' in params) and isinstance(params['class_weights'], (dict, OrderedDict)): + class_weights_dict = params['class_weights'] + class_weights_list = [] + if ('class_names' in params) and (params['class_names'] is not None): + if len(class_weights_dict) != len(params['class_names']): + raise CatBoostError('Number of classes in class_names and class_weights differ') + for class_label in params['class_names']: + if class_label not in class_weights_dict: + raise CatBoostError( + 'class "{}" is present in "class_names" but not in "class_weights" dictionary'.format( + class_label + ) + ) + class_weights_list.append(class_weights_dict[class_label]) + else: + class_labels_list = [] + for class_label, class_weight in class_weights_dict.items(): + class_labels_list.append(class_label) + class_weights_list.append(class_weight) + params['class_names'] = class_labels_list + params['class_weights'] = class_weights_list + + _process_synonyms_groups(params) + + metric_period = None + if 'metric_period' in params: + metric_period = params['metric_period'] + del params['metric_period'] + + verbose = None + if 'verbose' in params: + verbose = params['verbose'] + del params['verbose'] + + logging_level = None + if 'logging_level' in params: + logging_level = params['logging_level'] + del params['logging_level'] + + verbose_eval = None + if 'verbose_eval' in params: + verbose_eval = params['verbose_eval'] + del params['verbose_eval'] + + silent = None + if 'silent' in params: + silent = params['silent'] + del params['silent'] + + metric_period, verbose, logging_level = _process_verbose( + metric_period, verbose, logging_level, verbose_eval, silent) + + if metric_period is not None: + params['metric_period'] = metric_period + if verbose is not None: + params['verbose'] = verbose + if logging_level is not None: + params['logging_level'] = logging_level + + if 'used_ram_limit' in params: + params['used_ram_limit'] = str(params['used_ram_limit']) + + +def stringify_builtin_metrics(params): + """Replace all occurrences of BuiltinMetric with their string representations.""" + for f in ["loss_function", "objective", "eval_metric", "custom_metric", "custom_loss"]: + if f not in params: + continue + val = params[f] + if isinstance(val, BuiltinMetric): + params[f] = str(val) + elif isinstance(val, STRING_TYPES): + continue + elif isinstance(val, Sequence): + params[f] = stringify_builtin_metrics_list(val) + return params + + +def stringify_builtin_metrics_list(metrics): + return list(map(str, metrics)) + + +def _get_loss_function_for_train(params, estimator_type, train_pool): + """ + estimator_type must be 'classifier', 'regressor', 'ranker' or None + train_pool must be Pool + """ + + loss_function_param = params.get('loss_function') + if loss_function_param is not None: + return loss_function_param + + if estimator_type == 'classifier': + if not isinstance(train_pool, Pool): + raise CatBoostError('train_pool param must have Pool type') + + label = train_pool.get_label() + if label is None: + raise CatBoostError('loss function has not been specified and cannot be deduced') + + """ + len(set) is faster than np.unique on Python lists: + https://bbengfort.github.io/observations/2017/05/02/python-unique-benchmark.html + """ + is_multiclass_task = len(set(label)) > 2 and 'target_border' not in params + return 'MultiClass' if is_multiclass_task else 'Logloss' + elif estimator_type == 'ranker': + return 'YetiRank' + else: + return 'RMSE' + + +class _CatBoostBase(object): + def __init__(self, params): + init_params = params.copy() if params is not None else {} + stringify_builtin_metrics(init_params) + self._init_params = init_params + if 'thread_count' in self._init_params and self._init_params['thread_count'] == -1: + self._init_params.pop('thread_count') + if 'fixed_binary_splits' in self._init_params and self._init_params['fixed_binary_splits'] is None: + self._init_params['fixed_binary_splits'] = [] + # cache of processed _init_params with synonyms taken into account + self._canonized_params: Optional[Dict[str, object]] = None + self._object = _CatBoost() + + def __repr__(self) -> str: + params_str = ", ".join(f"{key}={val!r}" for key, val in sorted(self._init_params.items())) + return f"{self.__class__.__name__}({params_str})" + + def __getstate__(self): + params = self._init_params.copy() + test_evals = self._object._get_test_evals() + if test_evals: + params['_test_evals'] = test_evals + if self.is_fitted(): + params['__model'] = self._serialize_model() + for attr in ['_prediction_values_change', '_loss_value_change']: + if getattr(self, attr, None) is not None: + params[attr] = getattr(self, attr, None) + return params + + def __setstate__(self, state): + if '_object' not in dict(self.__dict__.items()): + self._object = _CatBoost() + if '_init_params' not in dict(self.__dict__.items()): + self._init_params = {} + if '__model' in state: + self._load_from_string(state['__model']) + del state['__model'] + if '_test_eval' in state: + self._set_test_evals([state['_test_eval']]) + del state['_test_eval'] + if '_test_evals' in state: + self._set_test_evals(state['_test_evals']) + del state['_test_evals'] + for attr in ['_prediction_values_change', '_loss_value_change']: + if attr in state: + setattr(self, attr, state[attr]) + del state[attr] + self._init_params.update(state) + self._canonized_params = None + + def __copy__(self): + return self.__deepcopy__(None) + + def __deepcopy__(self, _): + state = self.__getstate__() + model = self.__class__() + model.__setstate__(state) + return model + + def __eq__(self, other): + return self._is_comparable_to(other) and self._object == other._object + + def __ne__(self, other): + return not self._is_comparable_to(other) or self._object != other._object + + def copy(self): + return self.__copy__() + + def is_fitted(self): + return getattr(self, '_random_seed', None) is not None + + def _is_comparable_to(self, rhs): + if not isinstance(rhs, _CatBoostBase): + return False + for side, estimator in [('left', self), ('right', rhs)]: + if not estimator.is_fitted(): + message = 'The {} argument is not fitted, only fitted models' \ + ' could be compared.' + raise CatBoostError(message.format(side)) + return True + + def _set_trained_model_attributes(self): + setattr(self, '_is_fitted_', True) # compatibility with meta algorithms in sklearn + setattr(self, '_random_seed', self._object._get_random_seed()) + setattr(self, '_learning_rate', self._object._get_learning_rate()) + setattr(self, '_tree_count', self._object._get_tree_count()) + setattr(self, '_n_features_in', self._object._get_n_features_in()) + + def _train(self, train_pool, test_pool, params, allow_clear_pool, init_model): + self._object._train(train_pool, test_pool, params, allow_clear_pool, init_model._object if init_model else None) + self._set_trained_model_attributes() + + def _set_test_evals(self, test_evals): + self._object._set_test_evals(test_evals) + + def get_test_eval(self): + test_evals = self._object._get_test_evals() + if len(test_evals) == 0: + if self.is_fitted(): + raise CatBoostError('The model has been trained without an eval set.') + else: + raise CatBoostError('You should train the model first.') + if len(test_evals) > 1: + raise CatBoostError("With multiple eval sets use 'get_test_evals()'") + test_eval = test_evals[0] + return test_eval[0] if len(test_eval) == 1 else test_eval + + def get_test_evals(self): + test_evals = self._object._get_test_evals() + if len(test_evals) == 0: + if self.is_fitted(): + raise CatBoostError('The model has been trained without an eval set.') + else: + raise CatBoostError('You should train the model first.') + return test_evals + + def get_evals_result(self): + return self._object._get_metrics_evals() + + def get_best_score(self): + return self._object._get_best_score() + + def get_best_iteration(self): + return self._object._get_best_iteration() + + def get_n_features_in(self): + return self._object._get_n_features_in() + + def _get_float_feature_indices(self): + return self._object._get_float_feature_indices() + + def _get_cat_feature_indices(self): + return self._object._get_cat_feature_indices() + + def _get_text_feature_indices(self): + return self._object._get_text_feature_indices() + + def _get_embedding_feature_indices(self): + return self._object._get_embedding_feature_indices() + + def _base_predict(self, pool, prediction_type, ntree_start, ntree_end, thread_count, verbose, task_type): + return self._object._base_predict(pool, prediction_type, ntree_start, ntree_end, thread_count, verbose, task_type) + + def _base_virtual_ensembles_predict(self, pool, prediction_type, ntree_end, virtual_ensembles_count, thread_count, verbose): + return self._object._base_virtual_ensembles_predict(pool, prediction_type, ntree_end, virtual_ensembles_count, thread_count, verbose) + + def _staged_predict_iterator(self, pool, prediction_type, ntree_start, ntree_end, eval_period, thread_count, verbose): + return self._object._staged_predict_iterator(pool, prediction_type, ntree_start, ntree_end, eval_period, thread_count, verbose) + + def _leaf_indexes_iterator(self, pool, ntree_start, ntree_end): + return self._object._leaf_indexes_iterator(pool, ntree_start, ntree_end) + + def _base_calc_leaf_indexes(self, pool, ntree_start, ntree_end, thread_count, verbose): + return self._object._base_calc_leaf_indexes(pool, ntree_start, ntree_end, thread_count, verbose) + + def _base_eval_metrics(self, pool, metrics_description, ntree_start, ntree_end, eval_period, thread_count, result_dir, tmp_dir): + metrics_description_list = metrics_description if isinstance(metrics_description, list) else [metrics_description] + return self._object._base_eval_metrics(pool, metrics_description_list, ntree_start, ntree_end, eval_period, thread_count, result_dir, tmp_dir) + + def _calc_fstr( + self, + type, + pool, + reference_data, + thread_count, + verbose, + model_output, + shap_mode, + interaction_indices, + shap_calc_type, + sage_n_samples, + sage_batch_size, + sage_detect_convergence + ): + return self._object._calc_fstr( + type.name, + pool, + reference_data, + thread_count, + verbose, + model_output, + shap_mode, + interaction_indices, + shap_calc_type, + sage_n_samples, + sage_batch_size, + sage_detect_convergence + ) + + def _calc_ostr(self, train_pool, test_pool, top_size, ostr_type, update_method, importance_values_sign, thread_count, verbose): + return self._object._calc_ostr(train_pool, test_pool, top_size, ostr_type, update_method, importance_values_sign, thread_count, verbose) + + def _base_shrink(self, ntree_start, ntree_end): + self._object._base_shrink(ntree_start, ntree_end) + self._set_trained_model_attributes() + + def _base_drop_unused_features(self): + self._object._base_drop_unused_features() + + def _save_model(self, output_file, format, export_parameters, pool): + if self.is_fitted(): + params_string = "" + if export_parameters: + params_string = json.dumps(export_parameters, cls=_NumpyAwareEncoder) + + self._object._save_model(output_file, format, params_string, pool) + + def _load_model(self, model_file, format): + if not isinstance(model_file, PATH_TYPES): + raise CatBoostError("Invalid fname type={}: must be str or os.PathLike.".format(type(model_file))) + self._init_params = {} + self._object._load_model(model_file, format) + self._set_trained_model_attributes() + for key, value in iteritems(self._get_params_from_model_updated_with_init_params()): + self._init_params[key] = value + self._canonized_params = None + + def _serialize_model(self): + return self._object._serialize_model() + + def _deserialize_model(self, dump_model_str): + assert isinstance(dump_model_str, bytes), "Not bytes passed as argument" + self._object._deserialize_model(dump_model_str) + + def _load_from_string(self, dump_model_str): + self._deserialize_model(dump_model_str) + self._set_trained_model_attributes() + + def _load_from_stream(self, stream): + self._object._load_from_stream(stream) + self._set_trained_model_attributes() + + def _sum_models(self, models_base, weights=None, ctr_merge_policy='IntersectingCountersAverage'): + if weights is None: + weights = [1.0 for _ in models_base] + models_inner = [model._object for model in models_base] + self._object._sum_models(models_inner, weights, ctr_merge_policy) + setattr(self, '_random_seed', 0) + setattr(self, '_learning_rate', 0) + setattr(self, '_tree_count', self._object._get_tree_count()) + + def _save_borders(self, output_file): + if not self.is_fitted(): + raise CatBoostError("There is no trained model to use save_borders(). Use fit() to train model. Then use save_borders().") + self._object._save_borders(output_file) + + def _get_borders(self): + if not self.is_fitted(): + raise CatBoostError("There is no trained model to use get_feature_borders(). Use fit() to train model. Then use get_borders().") + return self._object._get_borders() + + def _get_nan_treatments(self): + assert self.is_fitted() + return self._object._get_nan_treatments() + + def _get_params_from_model_updated_with_init_params(self): + params = self._object._get_params() + init_params = self._init_params.copy() + for key, value in iteritems(init_params): + if key not in params: + params[key] = value + return params + + def _get_canonized_params(self) -> Dict[str, object]: + if self._canonized_params is None: + self._canonized_params = deepcopy(self._init_params) + _process_synonyms(self._canonized_params) + return self._canonized_params + + @staticmethod + def _is_classification_objective(loss_function): + return isinstance(loss_function, str) and is_classification_objective(loss_function) + + @staticmethod + def _is_regression_objective(loss_function): + return isinstance(loss_function, str) and is_regression_objective(loss_function) + + @staticmethod + def _is_multiregression_objective(loss_function): + if loss_function is None: + return False + elif isinstance(loss_function, str): + return is_multiregression_objective(loss_function) + else: + return id(MultiTargetCustomObjective) in [id(base) for base in type(loss_function).__mro__] + + @staticmethod + def _is_multitarget_objective(loss_function): + if loss_function is None: + return False + elif isinstance(loss_function, str): + return is_multitarget_objective(loss_function) + else: + return id(MultiTargetCustomObjective) in [id(base) for base in type(loss_function).__mro__] + + @staticmethod + def _is_multilabel_objective(loss_function): + return isinstance(loss_function, str) and is_multilabel_objective(loss_function) + + @staticmethod + def _is_survivalregression_objective(loss_function): + return isinstance(loss_function, str) and is_survivalregression_objective(loss_function) + + @staticmethod + def _is_ranking_objective(loss_function): + return isinstance(loss_function, str) and is_ranking_metric(loss_function) + + def _is_classifier(self, params_in_canonical_form: Dict[str, object]) -> bool: + return ( + (getattr(self, '_estimator_type', None) == 'classifier') + or + _CatBoostBase._is_classification_objective(params_in_canonical_form.get('loss_function', 'RMSE')) + ) + + @staticmethod + def _is_multiclass_compatible(params_in_canonical_form: Dict[str, object]) -> bool: + # called only for classifiers + + maybe_loss_function = params_in_canonical_form.get('loss_function') + if maybe_loss_function is None: + return True + elif isinstance(maybe_loss_function, str): + return is_multiclass_compatible_objective(maybe_loss_function) + else: + return isinstance(maybe_loss_function, MultiTargetCustomObjective) + + def get_metadata(self): + return self._object._get_metadata_wrapper() + + @property + def tree_count_(self): + return getattr(self, '_tree_count') if self.is_fitted() else None + + @property + def random_seed_(self): + return getattr(self, '_random_seed') if self.is_fitted() else None + + @property + def learning_rate_(self): + return getattr(self, '_learning_rate') if self.is_fitted() else None + + @property + def n_features_in_(self): + return getattr(self, '_n_features_in') if self.is_fitted() else None + + @property + def feature_names_(self): + return self._object._get_feature_names() if self.is_fitted() else None + + @property + def classes_(self): + return self._object._get_class_labels() if self.is_fitted() else None + + @property + def evals_result_(self): + return self.get_evals_result() + + @property + def best_score_(self): + return self.get_best_score() + + @property + def best_iteration_(self): + return self.get_best_iteration() + + def _get_tree_splits(self, tree_idx, pool): + return self._object._get_tree_splits(tree_idx, pool) + + def _get_tree_leaf_values(self, tree_idx): + return self._object._get_tree_leaf_values(tree_idx) + + def _get_tree_step_nodes(self, tree_idx): + return self._object._get_tree_step_nodes(tree_idx) + + def _get_tree_node_to_leaf(self, tree_idx): + return self._object._get_tree_node_to_leaf(tree_idx) + + def get_tree_leaf_counts(self): + ''' + Returns + ------- + tree_leaf_counts : 1d-array of numpy.uint32 of size tree_count_. + tree_leaf_counts[i] equals to the number of leafs in i-th tree of the ensemble. + ''' + return self._object._get_tree_leaf_counts() + + def get_leaf_values(self): + ''' + Returns + ------- + leaf_values : 1d-array of leaf values for all trees. + Value corresponding to j-th leaf of i-th tree is at position + sum(get_tree_leaf_counts()[:i]) + j (leaf and tree indexing starts from zero). + ''' + return self._object._get_leaf_values() + + def get_leaf_weights(self): + ''' + Returns + ------- + leaf_weights : 1d-array of leaf weights for all trees. + Weight of j-th leaf of i-th tree is at position + sum(get_tree_leaf_counts()[:i]) + j (leaf and tree indexing starts from zero). + ''' + return self._object._get_leaf_weights() + + def set_leaf_values(self, new_leaf_values): + ''' + Sets values at tree leafs of ensemble equal to new_leaf_values. + + Parameters + ---------- + new_leaf_values : 1d-array with new leaf values for all trees. + It's size should be equal to sum(get_tree_leaf_counts()). + Value corresponding to j-th leaf of i-th tree should be at position + sum(get_tree_leaf_counts()[:i]) + j (leaf and tree indexing starts from zero). + ''' + self._object._set_leaf_values(new_leaf_values) + + def set_feature_names(self, feature_names): + ''' + Sets feature names equal to feature_names + + Parameters + ---------- + feature_names: 1-d array of strings with new feature names in the same order as in pool + ''' + self._object._set_feature_names(feature_names) + + @staticmethod + def get_sklearn_estimator_xfail_checks(estimator: '_CatBoostBase') -> Dict[str, str]: + ''' + To pass to: + For scikit-learn >= 1.6: + to 'expected_failed_checks' parameter in sklearn.utils.estimator_checks.parametrize_with_checks + For scikit-learn < 1.6: + to '_xfail_checks' field returned by '_CatBoostBase._get_tags'. + + This method is made static to be able to pass it to sklearn.utils.estimator_checks.parametrize_with_checks + ''' + import sklearn + + scikit_learn_version = tuple(map(int, sklearn.__version__.split('.')[:2])) + + params = estimator._get_canonized_params() + + result = { + 'check_sample_weights_not_an_array': + 'TODO: not all array-like data is supported.' + ' https://github.com/catboost/catboost/issues/2995', + 'check_do_not_raise_errors_in_init_or_set_params': + 'TODO: https://github.com/catboost/catboost/issues/2997', + 'check_estimator_sparse_array': + 'TODO: support scipy.sparse sparse arrays. https://github.com/catboost/catboost/issues/3000', + 'check_estimator_sparse_matrix': + 'CatBoost does not support scipy.sparse.dia_matrix', + 'check_estimator_sparse_tag': + 'support scipy.sparse sparse arrays. https://github.com/catboost/catboost/issues/3000', + 'check_estimators_empty_data_messages': + 'TODO: raise ValueError instead of generic CatBoostError.' + ' https://github.com/catboost/catboost/issues/2996', + 'check_estimators_unfitted': + 'TODO: raise NotFittedError instead of generic CatBoostError.' + ' https://github.com/catboost/catboost/issues/3002', + 'check_fit1d': + 'TODO: CatBoost API allows to pass 1d array as features data (as a single feature),' + ' maybe this behavior should be tunable in the future', + 'check_fit2d_1sample': + 'TODO: raise ValueError mentioning "sample" instead of a current error. ' + 'https://github.com/catboost/catboost/issues/3003', + 'check_fit2d_predict1d': + 'TODO: CatBoost API allows to pass 1d array for prediction for a single sample,' + ' maybe this behavior should be tunable in the future', + 'check_n_features_in': + 'TODO: n_features_in_ must not be defined until fit is called. ' + 'https://github.com/catboost/catboost/issues/3004', + 'check_n_features_in_after_fitting': + 'TODO: 1) raise ValueError instead of generic CatBoostError.' + ' https://github.com/catboost/catboost/issues/2996; ' + '2) exact message match is too restrictive', + 'check_requires_y_none': + 'TODO: 1) raise ValueError instead of generic CatBoostError.' + ' https://github.com/catboost/catboost/issues/2996; ' + '2) exact message match is too restrictive', + 'check_sample_weight_equivalence_on_dense_data': + 'TODO: https://github.com/catboost/catboost/issues/3005', + 'check_sample_weight_equivalence_on_sparse_data': + 'support scipy.sparse sparse arrays. https://github.com/catboost/catboost/issues/3000', + 'check_sample_weights_shape': + 'TODO: raise ValueError instead of generic CatBoostError.' + ' https://github.com/catboost/catboost/issues/2996', + 'check_supervised_y_2d': + 'TODO: CatBoost API allows to pass 2D array for "y" when 1d is expected if 2nd dimension size = 1,' + ' maybe this behavior should be tunable in the future', + 'check_supervised_y_no_nan': + 'TODO: raise ValueError instead of generic CatBoostError.' + ' https://github.com/catboost/catboost/issues/2996', + 'check_dtype_object': + 'TODO: raise TypeError instead of generic CatBoostError.' + ' https://github.com/catboost/catboost/issues/2998', + } + + if scikit_learn_version < (1, 6): + result.update( + { + 'check_estimator_sparse_data': + 'CatBoost does not support scipy.sparse.dia_matrix', + } + ) + ''' + 'check_sample_weights_invariance' is tricky. + for 'kind=ones' it passes. + for 'kind=zeros' it fails. + TODO: https://github.com/catboost/catboost/issues/3005 + So we won't add it here but will check with params in 'test_sklearn_estimator_api_compatibility' + ''' + + if scikit_learn_version < (1, 4): + result.update( + { + 'check_fit_score_takes_y': + 'Bug in the test making all features constant when "categorical" is present in "X_types"', + 'check_sample_weights_list': + 'Bug in the test making all features constant when "categorical" is present in "X_types"', + } + ) + if params.get('nan_mode') != 'Forbidden': + result.update( + { + 'check_estimators_pickle': + 'Bug in the test that tries to assign np.nan to np.ndarray of np.int32 when ' + '"categorical" is present in "X_types" and "allow_nan" is True', + } + ) + if scikit_learn_version < (1, 3): + result.update( + { + 'check_no_attributes_set_in_init': + 'CatBoost does set some private "_*" attributes in init', + } + ) + + if estimator._is_classifier(params): + result.update( + { + 'check_classifier_data_not_an_array': + 'TODO: not all array-like data is supported.' + ' https://github.com/catboost/catboost/issues/2994', + 'check_classifiers_one_label': + 'TODO: raise ValueError instead of generic CatBoostError.' + ' https://github.com/catboost/catboost/issues/2996', + 'check_classifiers_regression_target': + 'TODO: CatBoost API allows to pass continuous target for binary classification,' + ' maybe this behavior should be tunable in the future', + 'check_classifiers_train': + 'TODO: raise ValueError instead of generic CatBoostError.' + ' https://github.com/catboost/catboost/issues/2996', + 'check_complex_data': + 'TODO: CatBoost API allows to pass complex data as labels,' + ' maybe this behavior should be changed in the future', + } + ) + else: + result.update( + { + 'check_regressor_data_not_an_array': + 'TODO: not all array-like data is supported.' + ' https://github.com/catboost/catboost/issues/2994', + 'check_complex_data': + 'TODO: raise ValueError instead of generic CatBoostError.' + ' https://github.com/catboost/catboost/issues/2996', + 'check_regressors_train': + 'TODO: raise ValueError instead of generic CatBoostError.' + ' https://github.com/catboost/catboost/issues/2996', + } + ) + + return result + + def _get_tags(self): + params = self._get_canonized_params() + estimator_type = 'classifier' if self._is_classifier(params) else 'regressor' + loss_function = params.get('loss_function') + + is_multilabel_objective = _CatBoostBase._is_multilabel_objective(loss_function) + is_multitarget_objective = _CatBoostBase._is_multitarget_objective(loss_function) + is_multi_output = ( + _CatBoostBase._is_multiregression_objective(loss_function) + or + is_multilabel_objective + ) + is_binary_only = ( + (estimator_type == 'classifier') + and (not _CatBoostBase._is_multiclass_compatible(params)) + and (not is_multilabel_objective) + ) + + X_types = ['2darray', 'sparse', 'categorical'] + if is_multitarget_objective: + X_types.append('2dlabels') + else: + X_types.append('1dlabels') + + return { + 'allow_nan': 'nan_mode' not in params or params['nan_mode'] != 'Forbidden', + 'array_api_support': False, # TODO: https://github.com/catboost/catboost/issues/2990 + 'binary_only': is_binary_only, + 'multilabel': is_multilabel_objective, + 'multioutput': is_multi_output, + 'multioutput_only': is_multi_output, + 'no_validation': False, + 'non_deterministic': params.get('task_type') == 'GPU', + 'pairwise': False, + 'poor_score': False, + 'preserves_dtype': [np.float64], + 'requires_fit': True, + 'requires_positive_X': False, + 'requires_y': True, + + # TODO: maybe create a method for user-defined objectives and metrics in the future + 'requires_positive_y': False, + '_skip_test': False, + '_xfail_checks': _CatBoostBase.get_sklearn_estimator_xfail_checks(self), + 'stateless': False, + 'X_types': X_types, + } + + def __sklearn_tags__(self): + import sklearn.utils + + params = self._get_canonized_params() + loss_function = params.get('loss_function') + + # scikit-learn does not have a separate type for ranking so set the type to 'regressor' + # in the case of CatBoost ranker as well + estimator_type = 'classifier' if self._is_classifier(params) else 'regressor' + is_multitarget_objective = _CatBoostBase._is_multitarget_objective(loss_function) + is_multilabel_objective = _CatBoostBase._is_multilabel_objective(loss_function) + is_multi_output = ( + _CatBoostBase._is_multiregression_objective(loss_function) + or + is_multilabel_objective + ) + + if estimator_type == 'classifier': + classifier_tags = sklearn.utils.ClassifierTags( + poor_score=False, + multi_class=_CatBoostBase._is_multiclass_compatible(params), + multi_label=is_multilabel_objective, + ) + regressor_tags = None + else: + classifier_tags = None + regressor_tags = sklearn.utils.RegressorTags(poor_score=False) + + return sklearn.utils.Tags( + estimator_type=estimator_type, + target_tags=sklearn.utils.TargetTags( + required=True, + one_d_labels=not is_multitarget_objective, + two_d_labels=is_multitarget_objective, + + # TODO: maybe create a method for user-defined objectives and metrics in the future + positive_only=False, + multi_output=is_multi_output, + single_output=not is_multi_output, + ), + classifier_tags=classifier_tags, + regressor_tags=regressor_tags, + array_api_support=False, # TODO: https://github.com/catboost/catboost/issues/2990 + no_validation=False, + non_deterministic=params.get('task_type') == 'GPU', + requires_fit=True, + input_tags=sklearn.utils.InputTags( + one_d_array=False, + two_d_array=True, + three_d_array=False, + sparse=True, + categorical=True, + string=False, + dict=False, + positive_only=False, + allow_nan=params.get('nan_mode') != 'Forbidden', + pairwise=False, + ), + ) + + def get_scale_and_bias(self): + return self._object._get_scale_and_bias() + + def set_scale_and_bias(self, scale, bias): + if isinstance(bias, FLOAT_TYPES): + self._object._set_scale_and_bias(scale, [bias]) + else: + self._object._set_scale_and_bias(scale, bias) + + +def _cast_value_to_list_of_strings(params, key): + if key in params: + if isinstance(params[key], STRING_TYPES): + params[key] = [params[key]] + if not isinstance(params[key], Sequence): + raise CatBoostError("Invalid `" + key + "` type={} : must be str or Sequence of strings.".format(type(params[key]))) + + +def _check_param_types(params): + if not isinstance(params, (Mapping, MutableMapping)): + raise CatBoostError("Invalid params type={}: must be Mapping.".format(type(params))) + if 'ctr_description' in params: + if not isinstance(params['ctr_description'], Sequence): + raise CatBoostError("Invalid ctr_description type={} : must be Sequence of strings".format(type(params['ctr_description']))) + if 'ctr_target_border_count' in params: + if not isinstance(params['ctr_target_border_count'], INTEGER_TYPES): + raise CatBoostError('Invalid ctr_target_border_count type={} : must be an integer'.format(type(params['ctr_target_border_count']))) + _cast_value_to_list_of_strings(params, 'custom_loss') + _cast_value_to_list_of_strings(params, 'custom_metric') + _cast_value_to_list_of_strings(params, 'per_float_feature_quantization') + if 'monotone_constraints' in params: + if not isinstance(params['monotone_constraints'], STRING_TYPES + ARRAY_TYPES + (dict,)): + raise CatBoostError( + "Invalid `monotone_constraints` type={} : must be str or a list of integers in range {{-1, 0, 1}} or dict.".format(type(params['monotone_constraints'])) + ) + if 'feature_weights' in params: + if not isinstance(params['feature_weights'], STRING_TYPES + ARRAY_TYPES + (dict,)): + raise CatBoostError("Invalid `feature_weights` type={} : must be str or a list of floats or dict.".format(type(params['feature_weights']))) + if 'first_feature_use_penalties' in params: + if not isinstance(params['first_feature_use_penalties'], STRING_TYPES + ARRAY_TYPES + (dict,)): + raise CatBoostError( + "Invalid `first_feature_use_penalties` type={} : must be str or a list of floats or dict.".format(type(params['first_feature_use_penalties'])) + ) + if 'per_object_feature_penalties' in params: + if not isinstance(params['per_object_feature_penalties'], STRING_TYPES + ARRAY_TYPES + (dict,)): + raise CatBoostError( + "Invalid `per_object_feature_penalties` type={} : must be str or a list of floats or dict.".format(type(params['per_object_feature_penalties'])) + ) + + +def _params_type_cast(params): + casted_params = {} + for key, value in iteritems(params): + value = _cast_to_base_types(value) + casted_params[key] = value + return casted_params + + +def _is_data_single_object(data): + if isinstance(data, (Pool, FeaturesData, pd.DataFrame) + SPARSE_MATRIX_TYPES): + return False + if not isinstance(data, ARRAY_TYPES): + raise CatBoostError( + "Invalid data type={} : must be list, numpy.ndarray, pandas.Series, pandas.DataFrame, polars.Series, polars.DataFrame" + " scipy.sparse matrix, catboost.FeaturesData or catboost.Pool".format(type(data)) + ) + return len(np.shape(data)) == 1 + + +def _process_feature_indices(feature_indices, pool, params, param_name): + if param_name not in params: + return feature_indices + + if param_name == 'cat_features': + feature_type_name = 'categorical' + elif param_name == 'text_features': + feature_type_name = 'text' + elif param_name == 'embedding_features': + feature_type_name = 'embedding' + else: + raise CatBoostError('Unknown params_name=' + param_name) + + if isinstance(pool, Pool): + feature_indices_from_params = _get_features_indices(params[param_name], pool.get_feature_names()) + if param_name == 'cat_features': + feature_indices_from_pool = pool.get_cat_feature_indices() + elif param_name == 'text_features': + feature_indices_from_pool = pool.get_text_feature_indices() + else: + feature_indices_from_pool = pool.get_embedding_feature_indices() + + if set(feature_indices_from_pool) != set(feature_indices_from_params): + raise CatBoostError(feature_type_name + " features indices in the model are set to " + + str(feature_indices_from_params) + + " and train dataset " + feature_type_name + " features indices are set to " + + str(feature_indices_from_pool)) + elif isinstance(pool, FeaturesData): + raise CatBoostError( + "Categorical features are set in the model. It is not allowed to use FeaturesData type for training dataset.") + else: + if feature_indices is not None and set(feature_indices) != set(params[param_name]): + raise CatBoostError(feature_type_name + " features in the model are set to " + str(params[param_name]) + + ". " + feature_type_name + " features passed to fit function are set to " + str(feature_indices)) + feature_indices = params[param_name] + del params[param_name] + return feature_indices + + +class CatBoost(_CatBoostBase): + """ + CatBoost model. Contains training, prediction and evaluation methods. + """ + + def __init__(self, params=None): + """ + Initialize the CatBoost. + + Parameters + ---------- + params : dict + Parameters for CatBoost. + If None, all params are set to their defaults. + If dict, overriding parameters present in dict. + """ + super(CatBoost, self).__init__(params) + + def _dataset_train_eval_split(self, train_pool, params, save_eval_pool): + """ + returns: + train_pool, eval_pool + eval_pool will be uninitialized if save_eval_pool is false + """ + + is_classification = self._is_classifier(params) + + return train_pool.train_eval_split( + params.get('has_time', False), + is_classification, + params['eval_fraction'], + save_eval_pool + ) + + def _prepare_train_params(self, X=None, y=None, cat_features=None, text_features=None, embedding_features=None, + pairs=None, graph=None, sample_weight=None, group_id=None, group_weight=None, subgroup_id=None, + pairs_weight=None, baseline=None, use_best_model=None, eval_set=None, verbose=None, + logging_level=None, plot=None, plot_file=None, column_description=None, verbose_eval=None, + metric_period=None, silent=None, early_stopping_rounds=None, save_snapshot=None, + snapshot_file=None, snapshot_interval=None, init_model=None, callbacks=None): + params = deepcopy(self._get_canonized_params()) + + if isinstance(X, FeaturesData): + warnings.warn("FeaturesData is deprecated for using in fit function " + "and soon will not be supported. If you want to use FeaturesData, " + "please pass it to Pool initialization and use Pool in fit") + + cat_features = _process_feature_indices(cat_features, X, params, 'cat_features') + text_features = _process_feature_indices(text_features, X, params, 'text_features') + embedding_features = _process_feature_indices(embedding_features, X, params, 'embedding_features') + + train_pool = _build_train_pool(X, y, cat_features, text_features, embedding_features, pairs, graph, + sample_weight, group_id, group_weight, subgroup_id, pairs_weight, + baseline, column_description) + if train_pool.is_empty_: + raise CatBoostError("X is empty.") + + allow_clear_pool = not isinstance(X, Pool) + + params['loss_function'] = _get_loss_function_for_train( + params, + getattr(self, '_estimator_type', None), + train_pool + ) + + metric_period, verbose, logging_level = _process_verbose( + metric_period, verbose, logging_level, verbose_eval, silent) + + if metric_period is not None: + params['metric_period'] = metric_period + if logging_level is not None: + params['logging_level'] = logging_level + if verbose is not None: + params['verbose'] = verbose + if use_best_model is not None: + params['use_best_model'] = use_best_model + + if early_stopping_rounds is not None: + params['od_type'] = 'Iter' + params['od_wait'] = early_stopping_rounds + if 'od_pval' in params: + del params['od_pval'] + + if save_snapshot is not None: + params['save_snapshot'] = save_snapshot + + if snapshot_file is not None: + params['snapshot_file'] = snapshot_file + + if snapshot_interval is not None: + params['snapshot_interval'] = snapshot_interval + + if callbacks is not None: + params['callbacks'] = _TrainCallbacksWrapper(callbacks) + + _check_param_types(params) + params = _params_type_cast(params) + _check_train_params(params) + + if params.get('eval_fraction', 0.0) != 0.0: + if eval_set is not None: + raise CatBoostError("Both eval_fraction and eval_set specified") + train_pool, eval_set = self._dataset_train_eval_split(train_pool, params, save_eval_pool=True) + + eval_set_list = eval_set if isinstance(eval_set, list) else [eval_set] + eval_sets = [] + eval_total_row_count = 0 + for eval_set in eval_set_list: + if isinstance(eval_set, Pool): + eval_sets.append(eval_set) + eval_total_row_count += eval_sets[-1].num_row() + if eval_sets[-1].num_row() == 0: + raise CatBoostError("Empty 'eval_set' in Pool") + elif isinstance(eval_set, PATH_TYPES): + eval_sets.append(Pool(eval_set, column_description=column_description)) + eval_total_row_count += eval_sets[-1].num_row() + if eval_sets[-1].num_row() == 0: + raise CatBoostError("Empty 'eval_set' in file {}".format(eval_set)) + elif isinstance(eval_set, tuple): + if len(eval_set) != 2: + raise CatBoostError("Invalid shape of 'eval_set': {}, must be (X, y).".format(str(tuple(type(_) for _ in eval_set)))) + if eval_set[0] is None or eval_set[1] is None: + raise CatBoostError("'eval_set' tuple contains at least one None value") + eval_sets.append( + Pool( + eval_set[0], + eval_set[1], + cat_features=train_pool.get_cat_feature_indices(), + text_features=train_pool.get_text_feature_indices(), + embedding_features=train_pool.get_embedding_feature_indices() + ) + ) + + eval_total_row_count += eval_sets[-1].num_row() + if eval_sets[-1].num_row() == 0: + raise CatBoostError("Empty 'eval_set' in tuple") + elif eval_set is None: + if len(eval_set_list) > 1: + raise CatBoostError("Multiple eval set shall not contain None") + else: + raise CatBoostError("Invalid type of 'eval_set': {}, while expected Pool or (X, y) or a filename, or list thereof.".format(type(eval_set))) + + if self.get_param('use_best_model') and eval_total_row_count == 0: + raise CatBoostError("To employ param {'use_best_model': True} provide non-empty 'eval_set'.") + + if (init_model is not None) and isinstance(init_model, PATH_TYPES): + try: + init_model = CatBoost().load_model(init_model) + except Exception as e: + raise CatBoostError("Error while loading init_model: {}".format(e)) + + return { + "train_pool": train_pool, + "eval_sets": eval_sets, + "params": params, + "allow_clear_pool": allow_clear_pool, + "init_model": init_model + } + + def _fit(self, X, y, cat_features, text_features, embedding_features, pairs, graph, sample_weight, group_id, group_weight, subgroup_id, + pairs_weight, baseline, use_best_model, eval_set, verbose, logging_level, plot, plot_file, + column_description, verbose_eval, metric_period, silent, early_stopping_rounds, + save_snapshot, snapshot_file, snapshot_interval, init_model, callbacks, log_cout=None, log_cerr=None): + + with log_fixup(log_cout, log_cerr): + if X is None: + raise CatBoostError("X must not be None") + + if y is None and not isinstance(X, PATH_TYPES + (Pool,)): + raise CatBoostError("y may be None only when X is an instance of catboost.Pool or string") + + train_params = self._prepare_train_params( + X=X, y=y, cat_features=cat_features, text_features=text_features, embedding_features=embedding_features, + pairs=pairs, graph=graph, sample_weight=sample_weight, group_id=group_id, group_weight=group_weight, + subgroup_id=subgroup_id, pairs_weight=pairs_weight, baseline=baseline, use_best_model=use_best_model, + eval_set=eval_set, verbose=verbose, logging_level=logging_level, plot=plot, plot_file=plot_file, + column_description=column_description, verbose_eval=verbose_eval, metric_period=metric_period, + silent=silent, early_stopping_rounds=early_stopping_rounds, save_snapshot=save_snapshot, + snapshot_file=snapshot_file, snapshot_interval=snapshot_interval, init_model=init_model, + callbacks=callbacks + ) + params = train_params["params"] + train_pool = train_params["train_pool"] + allow_clear_pool = train_params["allow_clear_pool"] + + with plot_wrapper(plot, plot_file, 'Training plots', [_get_train_dir(params)]): + self._train( + train_pool, + train_params["eval_sets"], + params, + allow_clear_pool, + train_params["init_model"] + ) + + # Have property feature_importance possibly set + loss = self._object._get_loss_function_name() + if loss and is_groupwise_metric(loss): + pass # too expensive + elif (len(self.get_embedding_feature_indices()) > 0): + pass # is not implemented yet + else: + if not self._object._has_leaf_weights_in_model(): + if allow_clear_pool: + train_pool = _build_train_pool( + X, + y, + cat_features, + text_features, + embedding_features, + pairs, + graph, + sample_weight, + group_id, + group_weight, + subgroup_id, + pairs_weight, + baseline, + column_description + ) + if params.get('eval_fraction', 0.0) != 0.0: + train_pool, _ = self._dataset_train_eval_split(train_pool, params, save_eval_pool=False) + + self.get_feature_importance(data=train_pool, type=EFstrType.PredictionValuesChange) + else: + self.get_feature_importance(type=EFstrType.PredictionValuesChange) + + return self + + def fit(self, X, y=None, cat_features=None, text_features=None, embedding_features=None, pairs=None, graph=None, sample_weight=None, group_id=None, + group_weight=None, subgroup_id=None, pairs_weight=None, baseline=None, use_best_model=None, + eval_set=None, verbose=None, logging_level=None, plot=False, plot_file=None, column_description=None, + verbose_eval=None, metric_period=None, silent=None, early_stopping_rounds=None, + save_snapshot=None, snapshot_file=None, snapshot_interval=None, init_model=None, callbacks=None, + log_cout=None, log_cerr=None): + """ + Fit the CatBoost model. + + Parameters + ---------- + X : catboost.Pool or list or numpy.ndarray or pandas.DataFrame or pandas.Series or polars.DataFrame + or string. + If not catboost.Pool or catboost.FeaturesData it must be 2 dimensional Feature matrix + or string - file with dataset. + + Must be non-empty (contain > 0 objects) + + y : list or numpy.ndarray or pandas.DataFrame or pandas.Series or polars.DataFrame or polars.Series, optional (default=None) + Labels of the training data. + If not None, can be a single- or two- dimensional array with either: + - numerical values - for regression (including multiregression), ranking and binary classification problems + - class labels (boolean, integer or string) - for classification (including multiclassification) problems + Use only if X is not catboost.Pool and does not point to a file. + + cat_features : list or numpy.ndarray, optional (default=None) + If not None, giving the list of Categ columns indices. + Use only if X is not catboost.Pool and not catboost.FeaturesData + + text_features: list or numpy.ndarray, optional (default=None) + If not none, giving the list of Text columns indices. + Use only if X is not catboost.Pool and not catboost.FeaturesData + + embedding_features: list or numpy.ndarray, optional (default=None) + If not none, giving the list of Embedding columns indices. + Use only if X is not catboost.Pool and not catboost.FeaturesData + + pairs : list or numpy.ndarray or pandas.DataFrame or polars.DataFrame + The pairs description. + If list or numpy.ndarrays or pandas.DataFrame or polars.DataFrame, giving 2 dimensional. + The shape should be Nx2, where N is the pairs' count. The first element of the pair is + the index of the winner object in the training set. The second element of the pair is + the index of the loser object in the training set. + + graph : list or numpy.ndarray or pandas.DataFrame or polars.DataFrame + The graph edges list description. + If list or numpy.ndarrays or pandas.DataFrame, giving 2 dimensional. + + sample_weight : list or numpy.ndarray or pandas.DataFrame or pandas.Series or polars.Series, optional (default=None) + Instance weights, 1 dimensional array like. + + group_id : list or numpy.ndarray or polars.Series, optional (default=None) + group id for each instance. + If not None, giving 1 dimensional array like data. + Use only if X is not catboost.Pool. + + group_weight : list or numpy.ndarray or polars.Series, optional (default=None) + Group weight for each instance. + If not None, giving 1 dimensional array like data. + + subgroup_id : list or numpy.ndarray or polars.Series, optional (default=None) + subgroup id for each instance. + If not None, giving 1 dimensional array like data. + Use only if X is not catboost.Pool. + + pairs_weight : list or numpy.ndarray or polars.Series, optional (default=None) + Weight for each pair. + If not None, giving 1 dimensional array like pairs. + + baseline : list or numpy.ndarray or polars.Series or polars.DataFrame, optional (default=None) + If not None, giving 2 dimensional array like data. + Use only if X is not catboost.Pool. + + use_best_model : bool, optional (default=None) + Flag to use best model + + eval_set : catboost.Pool or list of catboost.Pool or tuple (X, y) or list [(X, y)], optional (default=None) + Validation dataset or datasets for metrics calculation and possibly early stopping. + + logging_level : string, optional (default=None) + Possible values: + - 'Silent' + - 'Verbose' + - 'Info' + - 'Debug' + + metric_period : int + Frequency of evaluating metrics. + + verbose : bool or int + If verbose is bool, then if set to True, logging_level is set to Verbose, + if set to False, logging_level is set to Silent. + If verbose is int, it determines the frequency of writing metrics to output and + logging_level is set to Verbose. + + silent : bool + If silent is True, logging_level is set to Silent. + If silent is False, logging_level is set to Verbose. + + verbose_eval : bool or int + Synonym for verbose. Only one of these parameters should be set. + + plot : bool, optional (default=False) + If True, draw train and eval error in Jupyter notebook + + plot_file : file-like or str, optional (default=None) + If not None, save train and eval error graphs to file + + early_stopping_rounds : int + Activates Iter overfitting detector with od_wait parameter set to early_stopping_rounds. + + save_snapshot : bool, [default=None] + Enable progress snapshotting for restoring progress after crashes or interruptions + + snapshot_file : string or os.PathLike, [default=None] + Learn progress snapshot file path, if None will use default filename + + snapshot_interval: int, [default=600] + Interval between saving snapshots (seconds) + + init_model : CatBoost class or string or os.PathLike, [default=None] + Continue training starting from the existing model. + If this parameter is a string or os.PathLike, load initial model from the path specified by this string. + + callbacks : list, optional (default=None) + List of callback objects that are applied at end of each iteration. + + log_cout: output stream or callback for logging (default=None) + If None is specified, sys.stdout is used + + log_cerr: error stream or callback for logging (default=None) + If None is specified, sys.stderr is used + + Returns + ------- + model : CatBoost + """ + return self._fit(X, y, cat_features, text_features, embedding_features, pairs, graph, sample_weight, group_id, group_weight, subgroup_id, + pairs_weight, baseline, use_best_model, eval_set, verbose, logging_level, plot, plot_file, + column_description, verbose_eval, metric_period, silent, early_stopping_rounds, + save_snapshot, snapshot_file, snapshot_interval, init_model, callbacks, log_cout, log_cerr) + + def _process_predict_input_data(self, data, parent_method_name, thread_count, label=None): + if not self.is_fitted() or self.tree_count_ is None: + raise CatBoostError(("There is no trained model to use {}(). " + "Use fit() to train model. Then use this method.").format(parent_method_name)) + is_single_object = _is_data_single_object(data) + if not isinstance(data, Pool): + data = Pool( + data=[data] if is_single_object else data, + label=label, + cat_features=self._get_cat_feature_indices() if not isinstance(data, FeaturesData) else None, + text_features=self._get_text_feature_indices() if not isinstance(data, FeaturesData) else None, + embedding_features=self._get_embedding_feature_indices() if not isinstance(data, FeaturesData) else None, + thread_count=thread_count + ) + return data, is_single_object + + def _validate_prediction_type(self, prediction_type, valid_prediction_types=('Class', 'RawFormulaVal', 'Probability', 'LogProbability', 'Exponent', 'RMSEWithUncertainty')): + if not isinstance(prediction_type, STRING_TYPES): + raise CatBoostError("Invalid prediction_type type={}: must be str.".format(type(prediction_type))) + if prediction_type not in valid_prediction_types: + raise CatBoostError("Invalid value of prediction_type={}: must be {}.".format(prediction_type, ', '.join(valid_prediction_types))) + + def _predict(self, data, prediction_type, ntree_start, ntree_end, thread_count, verbose, parent_method_name, task_type="CPU"): + verbose = verbose or self.get_param('verbose') + if verbose is None: + verbose = False + data, data_is_single_object = self._process_predict_input_data(data, parent_method_name, thread_count) + self._validate_prediction_type(prediction_type) + + predictions = self._base_predict(data, prediction_type, ntree_start, ntree_end, thread_count, verbose, task_type) + return predictions[0] if data_is_single_object else predictions + + def predict(self, data, prediction_type='RawFormulaVal', ntree_start=0, ntree_end=0, thread_count=-1, verbose=None, task_type="CPU"): + """ + Predict with data. + + Parameters + ---------- + data : catboost.Pool or list of features or list of lists or numpy.ndarray or pandas.DataFrame or pandas.Series + or polars.DataFrame or polars.Series or catboost.FeaturesData + Data to apply model on. + If data is a simple list (not list of lists) or a one-dimensional numpy.ndarray it is interpreted + as a list of features for a single object. + + prediction_type : string, optional (default='RawFormulaVal') + Can be: + - 'RawFormulaVal' : return raw value. + - 'Class' : return class label. + - 'Probability' : return probability for every class. + - 'Exponent' : return Exponent of raw formula value. + - 'RMSEWithUncertainty': return standard deviation for RMSEWithUncertainty loss function + (logarithm of the standard deviation is returned by default). + + ntree_start: int, optional (default=0) + Model is applied on the interval [ntree_start, ntree_end) (zero-based indexing). + + ntree_end: int, optional (default=0) + Model is applied on the interval [ntree_start, ntree_end) (zero-based indexing). + If value equals to 0 this parameter is ignored and ntree_end equal to tree_count_. + + thread_count : int (default=-1) + The number of threads to use when applying the model. + Allows you to optimize the speed of execution. This parameter doesn't affect results. + If -1, then the number of threads is set to the number of CPU cores. + + verbose : bool, optional (default=False) + If True, writes the evaluation metric measured set to stderr. + + task_type : string, [default=None] + The evaluator type. + Possible values: + - 'CPU' + - 'GPU' (models with only numerical features are supported for now) + + Returns + ------- + prediction : + If data is for a single object, the return value depends on prediction_type value: + - 'RawFormulaVal' : return raw formula value. + - 'Class' : return class label. + - 'Probability' : return one-dimensional numpy.ndarray with probability for every class. + otherwise numpy.ndarray, with values that depend on prediction_type value: + - 'RawFormulaVal' : one-dimensional array of raw formula value for each object. + - 'Class' : one-dimensional array of class label for each object. + - 'Probability' : two-dimensional numpy.ndarray with shape (number_of_objects x number_of_classes) + with probability for every class for each object. + """ + return self._predict(data, prediction_type, ntree_start, ntree_end, thread_count, verbose, 'predict', task_type) + + def _virtual_ensembles_predict(self, data, prediction_type, ntree_end, virtual_ensembles_count, thread_count, verbose, parent_method_name): + verbose = verbose or self.get_param('verbose') + if verbose is None: + verbose = False + data, data_is_single_object = self._process_predict_input_data(data, parent_method_name, thread_count) + self._validate_prediction_type(prediction_type, ['VirtEnsembles', 'TotalUncertainty']) + + if ntree_end == 0: + ntree_end = self.tree_count_ + + predictions = self._base_virtual_ensembles_predict(data, prediction_type, ntree_end, virtual_ensembles_count, thread_count, verbose) + if prediction_type == 'VirtEnsembles': + shape = predictions.shape + predictions = predictions.reshape(shape[0], virtual_ensembles_count, int(shape[1] / virtual_ensembles_count)) + return predictions[0] if data_is_single_object else predictions + + def virtual_ensembles_predict(self, data, prediction_type='VirtEnsembles', ntree_end=0, virtual_ensembles_count=10, thread_count=-1, verbose=None): + """ + Predict with data. + + Parameters + ---------- + data : catboost.Pool or list of features or list of lists or numpy.ndarray or pandas.DataFrame or pandas.Series + or polars.DataFrame or polars.Series or catboost.FeaturesData + Data to apply model on. + If data is a simple list (not list of lists) or a one-dimensional numpy.ndarray it is interpreted + as a list of features for a single object. + + prediction_type : string, optional (default='RawFormulaVal') + Can be: + - 'VirtEnsembles': return V (virtual_ensembles_count) predictions. + k-th virtEnsemle consists of trees [0, T/2] + [T/2 + T/(2V) * k, T/2 + T/(2V) * (k + 1)] * constant. + - 'TotalUncertainty': see returned predictions format in 'Returns' part + + ntree_end: int, optional (default=0) + Model is applied on the interval [ntree_start, ntree_end) (zero-based indexing). + If value equals to 0 this parameter is ignored and ntree_end equal to tree_count_. + + virtual_ensembles_count: int, optional (default=10) + virtual ensembles count for 'TotalUncertainty' and 'VirtEnsembles' prediction types. + + thread_count : int (default=-1) + The number of threads to use when applying the model. + Allows you to optimize the speed of execution. This parameter doesn't affect results. + If -1, then the number of threads is set to the number of CPU cores. + + verbose : bool, optional (default=False) + If True, writes the evaluation metric measured set to stderr. + + Returns + ------- + prediction : + (with V as virtual_ensembles_count and T as trees count, + k-th virtEnsemle consists of trees [0, T/2] + [T/2 + T/(2V) * k, T/2 + T/(2V) * (k + 1)] * constant) + If data is for a single object, return 1-dimensional array of predictions with size depends on prediction type, + otherwise return 2-dimensional numpy.ndarray with shape (number_of_objects x size depends on prediction type); + Returned predictions depends on prediction type: + If loss-function was RMSEWithUncertainty: + - 'VirtEnsembles': [mean0, var0, mean1, var1, ..., vark-1]. + - 'TotalUncertainty': [mean_predict, KnowledgeUnc, DataUnc]. + otherwise for regression: + - 'VirtEnsembles': [mean0, mean1, ...]. + - 'TotalUncertainty': [mean_predicts, KnowledgeUnc]. + otherwise for binary classification: + - 'VirtEnsembles': [ApproxRawFormulaVal0, ApproxRawFormulaVal1, ..., ApproxRawFormulaValk-1]. + - 'TotalUncertainty': [DataUnc, TotalUnc]. + """ + return self._virtual_ensembles_predict(data, prediction_type, ntree_end, virtual_ensembles_count, thread_count, verbose, 'virtual_ensembles_predict') + + def _staged_predict(self, data, prediction_type, ntree_start, ntree_end, eval_period, thread_count, verbose, parent_method_name): + verbose = verbose or self.get_param('verbose') + if verbose is None: + verbose = False + data, data_is_single_object = self._process_predict_input_data(data, parent_method_name, thread_count) + self._validate_prediction_type(prediction_type) + + if ntree_end == 0: + ntree_end = self.tree_count_ + staged_predict_iterator = self._staged_predict_iterator(data, prediction_type, ntree_start, ntree_end, eval_period, thread_count, verbose) + for predictions in staged_predict_iterator: + yield predictions[0] if data_is_single_object else predictions + + def staged_predict(self, data, prediction_type='RawFormulaVal', ntree_start=0, ntree_end=0, eval_period=1, thread_count=-1, verbose=None): + """ + Predict target at each stage for data. + + Parameters + ---------- + data : catboost.Pool or list of features or list of lists or numpy.ndarray or pandas.DataFrame or pandas.Series + or polars.DataFrame or polars.Series or catboost.FeaturesData + Data to apply model on. + If data is a simple list (not list of lists) or a one-dimensional numpy.ndarray it is interpreted + as a list of features for a single object. + + prediction_type : string, optional (default='RawFormulaVal') + Can be: + - 'RawFormulaVal' : return raw formula value. + - 'Class' : return class label. + - 'Probability' : return probability for every class. + - 'RMSEWithUncertainty': return standard deviation for RMSEWithUncertainty loss function + (logarithm of the standard deviation is returned by default). + + ntree_start: int, optional (default=0) + Model is applied on the interval [ntree_start, ntree_end) with the step eval_period (zero-based indexing). + + ntree_end: int, optional (default=0) + Model is applied on the interval [ntree_start, ntree_end) with the step eval_period (zero-based indexing). + If value equals to 0 this parameter is ignored and ntree_end equal to tree_count_. + + eval_period: int, optional (default=1) + Model is applied on the interval [ntree_start, ntree_end) with the step eval_period (zero-based indexing). + + thread_count : int (default=-1) + The number of threads to use when applying the model. + Allows you to optimize the speed of execution. This parameter doesn't affect results. + If -1, then the number of threads is set to the number of CPU cores. + + verbose : bool + If True, writes the evaluation metric measured set to stderr. + + Returns + ------- + prediction : generator for each iteration that generates: + If data is for a single object, the return value depends on prediction_type value: + - 'RawFormulaVal' : return raw formula value. + - 'Class' : return class label. + - 'Probability' : return one-dimensional numpy.ndarray with probability for every class. + otherwise numpy.ndarray, with values that depend on prediction_type value: + - 'RawFormulaVal' : one-dimensional array of raw formula value for each object. + - 'Class' : one-dimensional array of class label for each object. + - 'Probability' : two-dimensional numpy.ndarray with shape (number_of_objects x number_of_classes) + with probability for every class for each object. + """ + return self._staged_predict(data, prediction_type, ntree_start, ntree_end, eval_period, thread_count, verbose, 'staged_predict') + + def _iterate_leaf_indexes(self, data, ntree_start, ntree_end): + if ntree_end == 0: + ntree_end = self.tree_count_ + data, _ = self._process_predict_input_data(data, "iterate_leaf_indexes", thread_count=-1) + leaf_indexes_iterator = self._leaf_indexes_iterator(data, ntree_start, ntree_end) + for leaf_index in leaf_indexes_iterator: + yield leaf_index + + def iterate_leaf_indexes(self, data, ntree_start=0, ntree_end=0): + """ + Returns indexes of leafs to which objects from pool are mapped by model trees. + + Parameters + ---------- + data : catboost.Pool or list of features or list of lists or numpy.ndarray or pandas.DataFrame or pandas.Series + or polars.DataFrame or polars.Series or catboost.FeaturesData + Data to apply model on. + If data is a simple list (not list of lists) or a one-dimensional numpy.ndarray it is interpreted + as a list of features for a single object. + + ntree_start: int, optional (default=0) + Index of first tree for which leaf indexes will be calculated (zero-based indexing). + + ntree_end: int, optional (default=0) + Index of the tree after last tree for which leaf indexes will be calculated (zero-based indexing). + If value equals to 0 this parameter is ignored and ntree_end equal to tree_count_. + + Returns + ------- + leaf_indexes : generator. For each object in pool yields one-dimensional numpy.ndarray of leaf indexes. + """ + return self._iterate_leaf_indexes(data, ntree_start, ntree_end) + + def _calc_leaf_indexes(self, data, ntree_start, ntree_end, thread_count, verbose): + if ntree_end == 0: + ntree_end = self.tree_count_ + data, _ = self._process_predict_input_data(data, "calc_leaf_indexes", thread_count) + return self._base_calc_leaf_indexes(data, ntree_start, ntree_end, thread_count, verbose) + + def calc_leaf_indexes(self, data, ntree_start=0, ntree_end=0, thread_count=-1, verbose=False): + """ + Returns indexes of leafs to which objects from pool are mapped by model trees. + + Parameters + ---------- + data : catboost.Pool or list of features or list of lists or numpy.ndarray or pandas.DataFrame or pandas.Series + or polars.DataFrame or polars.Series or catboost.FeaturesData + Data to apply model on. + If data is a simple list (not list of lists) or a one-dimensional numpy.ndarray it is interpreted + as a list of features for a single object. + + ntree_start: int, optional (default=0) + Index of first tree for which leaf indexes will be calculated (zero-based indexing). + + ntree_end: int, optional (default=0) + Index of the tree after last tree for which leaf indexes will be calculated (zero-based indexing). + If value equals to 0 this parameter is ignored and ntree_end equal to tree_count_. + + thread_count : int (default=-1) + The number of threads to use when applying the model. + Allows you to optimize the speed of execution. This parameter doesn't affect results. + If -1, then the number of threads is set to the number of CPU cores. + + verbose : bool (default=False) + Enable debug logging level. + + Returns + ------- + leaf_indexes : 2-dimensional numpy.ndarray of numpy.uint32 with shape (object count, ntree_end - ntree_start). + i-th row is an array of leaf indexes for i-th object. + """ + return self._calc_leaf_indexes(data, ntree_start, ntree_end, thread_count, verbose) + + def get_cat_feature_indices(self): + if not self.is_fitted(): + raise CatBoostError("Model is not fitted") + return self._get_cat_feature_indices() + + def get_text_feature_indices(self): + if not self.is_fitted(): + raise CatBoostError("Model is not fitted") + return self._get_text_feature_indices() + + def get_embedding_feature_indices(self): + if not self.is_fitted(): + raise CatBoostError("Model is not fitted") + return self._get_embedding_feature_indices() + + def _eval_metrics(self, data, metrics, ntree_start, ntree_end, eval_period, thread_count, res_dir, tmp_dir, plot, plot_file, log_cout=None, log_cerr=None): + if not self.is_fitted(): + raise CatBoostError("There is no trained model to evaluate metrics on. Use fit() to train model. Then call this method.") + if not isinstance(data, Pool): + raise CatBoostError("Invalid data type={}, must be catboost.Pool.".format(type(data))) + if data.is_empty_: + raise CatBoostError("Data is empty.") + if not isinstance(metrics, ARRAY_TYPES) and not isinstance(metrics, STRING_TYPES) and not isinstance(metrics, BuiltinMetric): + raise CatBoostError("Invalid metrics type={}, must be array like, str or one of builtin catboost.metrics.* class instances.".format(type(metrics))) + if not all(map(lambda metric: isinstance(metric, string_types) or isinstance(metric, BuiltinMetric), metrics)): + raise CatBoostError("Invalid metric type: must be str or one of builtin catboost.metrics.* class instances.") + if tmp_dir is None: + tmp_dir = tempfile.mkdtemp() + + if isinstance(metrics, STRING_TYPES) or isinstance(metrics, BuiltinMetric): + metrics = [metrics] + metrics = stringify_builtin_metrics_list(metrics) + with log_fixup(log_cout, log_cerr), plot_wrapper(plot, plot_file, 'Eval metrics plot', [res_dir]): + metrics_score, metric_names = self._base_eval_metrics(data, metrics, ntree_start, ntree_end, eval_period, thread_count, res_dir, tmp_dir) + + return dict(zip(metric_names, metrics_score)) + + def eval_metrics(self, data, metrics, ntree_start=0, ntree_end=0, eval_period=1, thread_count=-1, tmp_dir=None, plot=False, plot_file=None, log_cout=None, log_cerr=None): + """ + Calculate metrics. + + Parameters + ---------- + data : catboost.Pool + Data to evaluate metrics on. + + metrics : list of strings or catboost.metrics.BuiltinMetric + List of evaluated metrics. + + ntree_start: int, optional (default=0) + Model is applied on the interval [ntree_start, ntree_end) (zero-based indexing). + + ntree_end: int, optional (default=0) + Model is applied on the interval [ntree_start, ntree_end) (zero-based indexing). + If value equals to 0 this parameter is ignored and ntree_end equal to tree_count_. + + eval_period: int, optional (default=1) + Model is applied on the interval [ntree_start, ntree_end) with the step eval_period (zero-based indexing). + + thread_count : int (default=-1) + The number of threads to use when applying the model. + Allows you to optimize the speed of execution. This parameter doesn't affect results. + If -1, then the number of threads is set to the number of CPU cores. + + tmp_dir : string or os.PathLike (default=None) + The name of the temporary directory for intermediate results. + If None, then the name will be generated. + + plot : bool, optional (default=False) + If True, draw train and eval error in Jupyter notebook + + plot_file : file-like or str, optional (default=None) + If not None, save train and eval error graphs to file + + log_cout: output stream or callback for logging (default=None) + If None is specified, sys.stdout is used + + log_cerr: error stream or callback for logging (default=None) + If None is specified, sys.stderr is used + + Returns + ------- + prediction : dict: metric -> array of shape [(ntree_end - ntree_start) / eval_period] + """ + return self._eval_metrics(data, metrics, ntree_start, ntree_end, eval_period, thread_count, _get_train_dir(self._init_params), tmp_dir, plot, plot_file, log_cout, log_cerr) + + def compare(self, model, data, metrics, ntree_start=0, ntree_end=0, eval_period=1, thread_count=-1, tmp_dir=None, plot_file=None, log_cout=None, log_cerr=None): + """ + Draw train and eval errors in Jupyter notebook for both models + + Parameters + ---------- + model: CatBoost model + Another model to draw metrics + + data : catboost.Pool + Data to evaluate metrics on. + + metrics : list of strings or catboost.metrics.BuiltinMetric + List of evaluated metrics. + + ntree_start: int, optional (default=0) + Model is applied on the interval [ntree_start, ntree_end) (zero-based indexing). + + ntree_end: int, optional (default=0) + Model is applied on the interval [ntree_start, ntree_end) (zero-based indexing). + If value equals to 0 this parameter is ignored and ntree_end equal to tree_count_. + + eval_period: int, optional (default=1) + Model is applied on the interval [ntree_start, ntree_end) with the step eval_period (zero-based indexing). + + thread_count : int (default=-1) + The number of threads to use when applying the model. + Allows you to optimize the speed of execution. This parameter doesn't affect results. + If -1, then the number of threads is set to the number of CPU cores. + + tmp_dir : string or os.PathLike (default=None) + The name of the temporary directory for intermediate results. + If None, then the name will be generated. + + plot_file : file-like or str, optional (default=None) + If not None, save eval error graphs to file + + log_cout: output stream or callback for logging (default=None) + If None is specified, sys.stdout is used + + log_cerr: error stream or callback for logging (default=None) + If None is specified, sys.stderr is used + """ + + if model is None: + raise CatBoostError("You should provide model for comparison.") + if data is None: + raise CatBoostError("You should provide data for comparison.") + if metrics is None: + raise CatBoostError("You should provide metrics for comparison.") + + need_to_remove = False + if tmp_dir is None: + need_to_remove = True + tmp_dir = tempfile.mkdtemp() + first_dir = os.path.join(tmp_dir, 'first_model') + second_dir = os.path.join(tmp_dir, 'second_model') + + create_dir_if_not_exist(tmp_dir) + create_dir_if_not_exist(first_dir) + create_dir_if_not_exist(second_dir) + + with plot_wrapper(True, plot_file=plot_file, plot_title='Compare models', train_dirs=[first_dir, second_dir]): + self._eval_metrics(data, metrics, ntree_start, ntree_end, eval_period, thread_count, first_dir, tmp_dir, + plot=False, plot_file=None, log_cout=log_cout, log_cerr=log_cerr) + model._eval_metrics(data, metrics, ntree_start, ntree_end, eval_period, thread_count, second_dir, tmp_dir, + plot=False, plot_file=None, log_cout=log_cout, log_cerr=log_cerr) + + if need_to_remove: + shutil.rmtree(tmp_dir) + + def create_metric_calcer(self, metrics, ntree_start=0, ntree_end=0, eval_period=1, thread_count=-1, tmp_dir=None): + """ + Create batch metric calcer. Could be used to aggregate metric on several pools + Parameters + ---------- + Same as in eval_metrics except data + Returns + ------- + BatchMetricCalcer object + + Usage example + ------- + # Large dataset is partitioned into parts [part1, part2] + model.fit(params) + batch_calcer = model.create_metric_calcer(['Logloss']) + batch_calcer.add(part1) + batch_calcer.add(part2) + metrics = batch_calcer.eval_metrics() + """ + if not self.is_fitted(): + raise CatBoostError("There is no trained model to evaluate metrics on. Use fit() to train model. Then call this method.") + return BatchMetricCalcer(self._object, metrics, ntree_start, ntree_end, eval_period, thread_count, tmp_dir) + + @property + def feature_importances_(self): + loss = self._object._get_loss_function_name() + if loss and is_groupwise_metric(loss): + return np.array(getattr(self, "_loss_value_change", None)) + else: + return np.array(getattr(self, "_prediction_values_change", None)) + + def get_feature_importance(self, data=None, type=EFstrType.FeatureImportance, prettified=False, + thread_count=-1, verbose=False, fstr_type=None, shap_mode="Auto", + model_output="Raw", interaction_indices=None, shap_calc_type="Regular", + reference_data=None, sage_n_samples=128, sage_batch_size=512, + sage_detect_convergence=True, log_cout=None, log_cerr=None): + """ + Parameters + ---------- + data : + Data to get feature importance. + If type in ('LossFunctionChange', 'ShapValues', 'ShapInteractionValues') data must of Pool type. + For every object in this dataset feature importances will be calculated. + if type == 'SageValues' data must of Pool type. + For every feature in this dataset importance will be calculated. + If type == 'PredictionValuesChange', data is None or a dataset of Pool type + Dataset specification is needed only in case if the model does not contain leaf weight information (trained with CatBoost v < 0.9). + If type == 'PredictionDiff' data must contain a matrix of feature values of shape (2, n_features). + Possible types are catboost.Pool or list of lists or numpy.ndarray or pandas.DataFrame or pandas.Series + or polars.DataFrame or catboost.FeaturesData or pandas.SparseDataFrame or scipy.sparse.spmatrix + If type == 'FeatureImportance' + See 'PredictionValuesChange' for non-ranking metrics and 'LossFunctionChange' for ranking metrics. + If type == 'Interaction' + This parameter is not used. + + type : EFstrType or string (converted to EFstrType), optional + (default=EFstrType.FeatureImportance) + Possible values: + - PredictionValuesChange + Calculate score for every feature. + - LossFunctionChange + Calculate score for every feature by loss. + - FeatureImportance + PredictionValuesChange for non-ranking metrics and LossFunctionChange for ranking metrics + - ShapValues + Calculate SHAP Values for every object. + - ShapInteractionValues + Calculate SHAP Interaction Values between each pair of features for every object + - Interaction + Calculate pairwise score between every feature. + - PredictionDiff + Calculate most important features explaining difference in predictions for a pair of documents. + - SageValues + Calculate SAGE value for every feature + + prettified : bool, optional (default=False) + change returned data format to the list of (feature_id, importance) pairs sorted by importance + + thread_count : int, optional (default=-1) + Number of threads. + If -1, then the number of threads is set to the number of CPU cores. + + verbose : bool or int + If False, then evaluation is not logged. If True, then each possible iteration is logged. + If a positive integer, then it stands for the size of batch N. After processing each batch, print progress + and remaining time. + + fstr_type : string, deprecated, use type instead + + shap_mode : string, optional (default="Auto") + used only for ShapValues type + Possible values: + - "Auto" + Use direct SHAP Values calculation only if data size is smaller than average leaves number + (the best of two strategies below is chosen). + - "UsePreCalc" + Calculate SHAP Values for every leaf in preprocessing. Final complexity is + O(NT(D+F))+O(TL^2 D^2) where N is the number of documents(objects), T - number of trees, + D - average tree depth, F - average number of features in tree, L - average number of leaves in tree + This is much faster (because of a smaller constant) than direct calculation when N >> L + - "NoPreCalc" + Use direct SHAP Values calculation calculation with complexity O(NTLD^2). Direct algorithm + is faster when N < L (algorithm from https://arxiv.org/abs/1802.03888) + + shap_calc_type : EShapCalcType or string, optional (default="Regular") + used only for ShapValues type + Possible values: + - "Regular" + Calculate regular SHAP values + - "Approximate" + Calculate approximate SHAP values + - "Exact" + Calculate exact SHAP values + + interaction_indices : list of int or string (feature_idx_1, feature_idx_2), optional (default=None) + used only for ShapInteractionValues type + Calculate SHAP Interaction Values between pair of features feature_idx_1 and feature_idx_2 for every object + + reference_data: catboost.Pool or None + Reference data for Independent Tree SHAP values from https://arxiv.org/abs/1905.04610v1 + if type == 'ShapValues' and reference_data is not None, then Independent Tree SHAP values are calculated + + sage_n_samples: int, optional (default=32) + Number of outer samples used in SAGE values approximation algorithm + sage_batch_size: int, optional (default=min(512, number of samples in dataset)) + Number of samples used on each step of SAGE values approximation algorithm + sage_detect_convergence: bool, optional (default=False) + If set True, sage values calculation will be stopped either when sage values converge + or when sage_n_samples iterations of algorithm pass + + log_cout: output stream or callback for logging (default=None) + If None is specified, sys.stdout is used + + log_cerr: error stream or callback for logging (default=None) + If None is specified, sys.stderr is used + + Returns + ------- + depends on type: + - FeatureImportance + See PredictionValuesChange for non-ranking metrics and LossFunctionChange for ranking metrics. + - PredictionValuesChange, LossFunctionChange, PredictionDiff, SageValues with prettified=False (default) + list of length [n_features] with feature_importance values (float) for feature + - PredictionValuesChange, LossFunctionChange, PredictionDiff, SageValues with prettified=True + list of length [n_features] with (feature_id (string), feature_importance (float)) pairs, sorted by feature_importance in descending order + - ShapValues + np.ndarray of shape (n_objects, n_features + 1) with Shap values (float) for (object, feature). + In case of multiclass the returned value is np.ndarray of shape + (n_objects, classes_count, n_features + 1). For each object it contains Shap values (float). + Values are calculated for RawFormulaVal predictions. + - ShapInteractionValues + np.ndarray of shape (n_objects, n_features + 1, n_features + 1) with Shap interaction values (float) for (object, feature(i), feature(j)). + In case of multiclass the returned value is np.ndarray of shape + (n_objects, classes_count, n_features + 1, n_features + 1). For each object it contains Shap interaction values (float). + Values are calculated for RawFormulaVal predictions. + - Interaction + list of length [n_features] of 3-element lists of (first_feature_index, second_feature_index, interaction_score (float)) + """ + + with log_fixup(log_cout, log_cerr): + if not isinstance(verbose, bool) and not isinstance(verbose, int): + raise CatBoostError('verbose should be bool or int.') + verbose = int(verbose) + if verbose < 0: + raise CatBoostError('verbose should be non-negative.') + + if fstr_type is not None: + type = fstr_type + + type = enum_from_enum_or_str(EFstrType, type) + if type == EFstrType.FeatureImportance: + loss = self._object._get_loss_function_name() + if loss and is_groupwise_metric(loss): + type = EFstrType.LossFunctionChange + else: + type = EFstrType.PredictionValuesChange + + if type == EFstrType.PredictionDiff: + data, _ = self._process_predict_input_data(data, "get_feature_importance", thread_count) + if data.num_row() != 2: + raise CatBoostError("{} requires a pair of documents, found {}".format(type, data.num_row())) + else: + if data is not None and not isinstance(data, Pool): + raise CatBoostError("Invalid data type={}, must be catboost.Pool.".format(_typeof(data))) + + need_meta_info = type == EFstrType.PredictionValuesChange + empty_data_is_ok = need_meta_info and self._object._has_leaf_weights_in_model() or type == EFstrType.Interaction + if not empty_data_is_ok: + if data is None: + if need_meta_info: + raise CatBoostError( + "Model has no meta information needed to calculate feature importances. \ + Pass training dataset to this function.") + else: + raise CatBoostError( + "Feature importance type {} requires training dataset \ + to be passed to this function.".format(type)) + if data.is_empty_: + raise CatBoostError("data is empty.") + + shap_calc_type = enum_from_enum_or_str(EShapCalcType, shap_calc_type).value + fstr, feature_names = self._calc_fstr(type, data, reference_data, thread_count, verbose, model_output, shap_mode, interaction_indices, + shap_calc_type, sage_n_samples, sage_batch_size, sage_detect_convergence) + if type in (EFstrType.PredictionValuesChange, EFstrType.LossFunctionChange, + EFstrType.PredictionDiff, EFstrType.SageValues): + feature_importances = [value[0] for value in fstr] + attribute_name = None + if type == EFstrType.PredictionValuesChange: + attribute_name = "_prediction_values_change" + if type == EFstrType.LossFunctionChange: + attribute_name = "_loss_value_change" + if attribute_name: + setattr( + self, + attribute_name, + feature_importances + ) + + if prettified: + feature_importances = sorted(zip(feature_names, feature_importances), key=itemgetter(1), reverse=True) + columns = ['Feature Id', 'Importances'] + return pd.DataFrame(feature_importances, columns=columns) + else: + return np.array(feature_importances) + elif type == EFstrType.ShapValues: + if isinstance(fstr[0][0], ARRAY_TYPES): + return np.array([np.array([np.array([ + value for value in dimension]) for dimension in doc]) for doc in fstr]) + else: + result = [[value for value in doc] for doc in fstr] + if prettified: + return pd.DataFrame(result) + else: + return np.array(result) + elif type == EFstrType.ShapInteractionValues: + if isinstance(fstr[0][0], ARRAY_TYPES): + return np.array([np.array([np.array([ + feature2 for feature2 in feature1]) for feature1 in doc]) for doc in fstr]) + else: + return np.array([np.array([np.array([np.array([ + feature2 for feature2 in feature1]) for feature1 in dimension]) for dimension in doc]) for doc in fstr]) + elif type == EFstrType.Interaction: + result = [[int(row[0]), int(row[1]), row[2]] for row in fstr] + if prettified: + columns = ['First Feature Index', 'Second Feature Index', 'Interaction'] + return pd.DataFrame(result, columns=columns) + else: + return np.array(result) + + def get_object_importance( + self, pool, train_pool, top_size=-1, type='Average', update_method='SinglePoint', + importance_values_sign='All', thread_count=-1, verbose=False, ostr_type=None, + log_cout=None, log_cerr=None): + """ + This is the implementation of the LeafInfluence algorithm from the following paper: + https://arxiv.org/pdf/1802.06640.pdf + + Parameters + ---------- + pool : Pool + The pool for which you want to evaluate the object importances. + + train_pool : Pool + The pool on which the model has been trained. + + top_size : int (default=-1) + Method returns the result of the top_size most important train objects. + If -1, then the top size is not limited. + + type : string, optional (default='Average') + Possible values: + - Average (Method returns the mean train objects scores for all input objects) + - PerObject (Method returns the train objects scores for every input object) + + importance_values_sign : string, optional (default='All') + Method returns only Positive, Negative or All values. + Possible values: + - Positive + - Negative + - All + + update_method : string, optional (default='SinglePoint') + Possible values: + - SinglePoint + - TopKLeaves (It is posible to set top size : TopKLeaves:top=2) + - AllPoints + Description of the update set methods are given in section 3.1.3 of the paper. + + thread_count : int, optional (default=-1) + Number of threads. + If -1, then the number of threads is set to the number of CPU cores. + + verbose : bool or int + If False, then evaluation is not logged. If True, then each possible iteration is logged. + If a positive integer, then it stands for the size of batch N. After processing each batch, print progress + and remaining time. + + ostr_type : string, deprecated, use type instead + + log_cout: output stream or callback for logging (default=None) + If None is specified, sys.stdout is used + + log_cerr: error stream or callback for logging (default=None) + If None is specified, sys.stderr is used + + Returns + ------- + object_importances : tuple of two arrays (indices and scores) of shape = [top_size] + """ + + if self.is_fitted() and not self._object._is_oblivious(): + raise CatBoostError('Object importance is not supported for non symmetric trees') + + if not isinstance(verbose, bool) and not isinstance(verbose, int): + raise CatBoostError('verbose should be bool or int.') + verbose = int(verbose) + if verbose < 0: + raise CatBoostError('verbose should be non-negative.') + + if ostr_type is not None: + type = ostr_type + warnings.warn("'ostr_type' parameter will be deprecated soon, use 'type' parameter instead") + + with log_fixup(log_cout, log_cerr): + result = self._calc_ostr(train_pool, pool, top_size, type, update_method, importance_values_sign, thread_count, verbose) + return result + + def shrink(self, ntree_end, ntree_start=0): + """ + Shrink the model. + + Parameters + ---------- + ntree_end: int + Leave the trees with indices from the interval [ntree_start, ntree_end) (zero-based indexing). + ntree_start: int, optional (default=0) + Leave the trees with indices from the interval [ntree_start, ntree_end) (zero-based indexing). + """ + if ntree_start > ntree_end: + raise CatBoostError("ntree_start should be less than ntree_end.") + self._base_shrink(ntree_start, ntree_end) + + def drop_unused_features(self): + """ + Drop unused features information from model + """ + self._base_drop_unused_features() + + def save_model(self, fname, format="cbm", export_parameters=None, pool=None): + """ + Save the model to a file. + + Parameters + ---------- + fname : string + Output file name. + format : string + Possible values: + * 'cbm' for catboost binary format, + * 'coreml' to export into Apple CoreML format + * 'onnx' to export into ONNX-ML format + * 'pmml' to export into PMML format + * 'cpp' to export as C++ code + * 'python' to export as Python code. + export_parameters : dict + Parameters for CoreML export: + * prediction_type : string - either 'probability' or 'raw' + * coreml_description : string + * coreml_model_version : string + * coreml_model_author : string + * coreml_model_license: string + Parameters for PMML export: + * pmml_copyright : string + * pmml_description : string + * pmml_model_version : string + pool : catboost.Pool or list or numpy.ndarray or pandas.DataFrame or pandas.Series or polars.DataFrame or catboost.FeaturesData + Training pool. + """ + if not self.is_fitted(): + raise CatBoostError("There is no trained model to use save_model(). Use fit() to train model. Then use this method.") + if not isinstance(fname, PATH_TYPES): + raise CatBoostError("Invalid fname type={}: must be str or os.PathLike.".format(type(fname))) + if pool is not None and not isinstance(pool, Pool): + pool = Pool( + data=pool, + cat_features=self._get_cat_feature_indices() if not isinstance(pool, FeaturesData) else None, + text_features=self._get_text_feature_indices() if not isinstance(pool, FeaturesData) else None, + embedding_features=self._get_embedding_feature_indices() if not isinstance(pool, FeaturesData) else None + ) + self._save_model(fname, format, export_parameters, pool) + + def load_model(self, fname=None, format='cbm', stream=None, blob=None): + """ + Load model from a file, stream or blob. + + Parameters + ---------- + fname : string + Input file name. + """ + if (fname is None) + (stream is None) + (blob is None) != 2: + raise CatBoostError("Only one of fname/stream/blob arguments should be specified") + + if fname is not None: + self._load_model(fname, format) + elif stream is not None: + self._load_from_stream(stream) + elif blob is not None: + self._load_from_string(blob) + return self + + def get_param(self, key): + """ + Get param value from CatBoost model. + + Parameters + ---------- + key : string + The key to get param value from. + + Returns + ------- + value : + The param value of the key, returns None if param do not exist. + """ + value = self._init_params.get(key) + if value is not None: + return deepcopy(value) + return None + + def get_params(self, **_unused_kwargs): + """ + Get all params from CatBoost model. + + Returns + ------- + result : dict + Dictionary of {param_key: param_value}. + """ + return deepcopy(self._init_params) + + def get_all_params(self): + """ + Get all params (specified by user and default params) that were set in training from CatBoost model. + Full parameters documentation could be found here: https://catboost.ai/docs/concepts/python-reference_parameters-list.html + + Returns + ------- + result : dict + Dictionary of {param_key: param_value}. + """ + if not self.is_fitted(): + raise CatBoostError("There is no trained model to use get_all_params(). Use fit() to train model. Then use this method.") + return self._object._get_plain_params() + + def save_borders(self, fname): + """ + Save the model borders to a file. + + Parameters + ---------- + fname : string or os.PathLike + Output file name. + """ + if not isinstance(fname, PATH_TYPES): + raise CatBoostError("Invalid fname type={}: must be str or os.PathLike.".format(type(fname))) + self._save_borders(fname) + + def get_borders(self): + """ + Return map feature_index: borders for float features. + + """ + return self._get_borders() + + def set_params(self, **params): + """ + Set parameters into CatBoost model. + + Parameters + ---------- + **params : key=value format + List of key=value paris. Example: model.set_params(iterations=500, thread_count=2). + """ + if self.is_fitted(): + raise CatBoostError("You can't change params of fitted model.") + for key, value in iteritems(params): + self._init_params[key] = value + if 'thread_count' in self._init_params and self._init_params['thread_count'] == -1: + self._init_params.pop('thread_count') + self._canonized_params = None + return self + + def plot_predictions(self, data, features_to_change, plot=True, plot_file=None): + """ + To use this function, you should install plotly. + + data: numpy.ndarray or pandas.DataFrame or polars.DataFrame or catboost.Pool + features_to_change: list-like with int (for indices) or str (for names) elements + Numerical features indices or names in `data` for which you want to vary prediction value. + plot: bool + Plot predictions. + plot_file: str + Output file for plot predictions. + Returns + ------- + List of list of predictions for all buckets for all samples in data + """ + + def predict(doc, feature_idx, borders, nan_treatment): + left_extend_border = min(2 * borders[0], -1) + right_extend_border = max(2 * borders[-1], 1) + extended_borders = [left_extend_border] + borders + [right_extend_border] + points = [] + predictions = [] + border_idx = None + if np.isnan(doc[feature_idx]): + border_idx = len(borders) if nan_treatment == 'AsTrue' else 0 + for i in range(len(extended_borders) - 1): + points += [(extended_borders[i] + extended_borders[i + 1]) / 2.] + if border_idx is None and doc[feature_idx] < extended_borders[i + 1]: + border_idx = i + buf = doc[feature_idx] + doc[feature_idx] = points[-1] + predictions += [self.predict(doc)] + doc[feature_idx] = buf + if border_idx is None: + border_idx = len(borders) + return predictions, border_idx + + def get_layout(go, feature, xaxis): + return go.Layout( + title="Prediction variation for feature '{}'".format(feature), + yaxis={ + 'title': 'Prediction', + 'side': 'left', + 'overlaying': 'y2' + }, + xaxis=xaxis + ) + try: + import plotly.graph_objs as go + except ImportError as e: + warnings.warn("To draw plots you should install plotly.") + raise ImportError(str(e)) + + model_borders = self._get_borders() + + data, _ = self._process_predict_input_data(data, "vary_feature_value_and_apply", thread_count=-1) + figs = [] + all_predictions = [{}] * data.num_row() + nan_treatments = self._get_nan_treatments() + for feature in features_to_change: + if not isinstance(feature, int): + if self.feature_names_ is None or feature not in self.feature_names_: + raise CatBoostError('No feature named "{}" in model'.format(feature)) + feature_idx = self.feature_names_.index(feature) + else: + feature_idx = feature + feature = self.feature_names_[feature_idx] + assert feature_idx in model_borders, "only float features indexes are supported" + borders = model_borders[feature_idx] + + if len(borders) == 0: + xaxis = go.layout.XAxis(title='Bins', tickvals=[0]) + figs += go.Figure(data=[], layout=get_layout(go, feature_idx, xaxis)) + + xaxis = go.layout.XAxis( + title='Bins', + tickmode='array', + tickvals=list(range(len(borders) + 1)), + ticktext=['(-inf, {:.4f}]'.format(borders[0])] + + ['({:.4f}, {:.4f}]'.format(val_1, val_2) + for val_1, val_2 in zip(borders[:-1], borders[1:])] + + ['({:.4f}, +inf)'.format(borders[-1])], + showticklabels=False + ) + + trace = [] + for idx, features in enumerate(data.get_features()): + predictions, border_idx = predict(features, feature_idx, borders, nan_treatments[feature_idx]) + all_predictions[idx][feature_idx] = predictions + trace.append( + go.Scatter(y=predictions, mode='lines+markers', name=u'Document {} predictions'.format(idx)) + ) + + trace.append(go.Scatter(x=[border_idx], y=[predictions[border_idx]], showlegend=False)) + + layout = get_layout(go, feature, xaxis) + figs += [go.Figure(data=trace, layout=layout)] + + if plot: + try_plot_offline(figs) + + if plot_file: + save_plot_file(plot_file, 'Predictions for all buckets', figs) + + return all_predictions, figs + + def plot_partial_dependence(self, data, features, plot=True, plot_file=None, thread_count=-1): + """ + To use this function, you should install plotly. + data: numpy.ndarray or pandas.DataFrame or polars.DataFrame or catboost.Pool + features: int, str, list, tuple, list, tuple + Float features to calculate partial dependence for. Number of features should be 1 or 2. + plot: bool + Plot predictions. + plot_file: str + Output file for plot predictions. + thread_count: int + Number of threads to use. If -1 use maximum available number of threads. + Returns + ------- + If number of features is one - 1d numpy array and figure with line plot. + If number of features is two - 2d numpy array and figure with 2d heatmap. + """ + + try: + import plotly.graph_objs as go + except ImportError as e: + warnings.warn("To draw plots you should install plotly.") + raise ImportError(str(e)) + + def getFeatureIdx(feature): + if not isinstance(feature, int): + if self.feature_names_ is None or feature not in self.feature_names_: + raise CatBoostError('No feature named "{}" in model'.format(feature)) + feature_idx = self.feature_names_.index(feature) + else: + feature_idx = feature + assert feature_idx in self._get_borders(), "only float features indexes are supported" + assert len(self._get_borders()[feature_idx]) > 0, "feature with idx {} is not used in model".format(feature_idx) + return feature_idx + + def getFeatureIndices(features): + if isinstance(features, list) or isinstance(features, tuple): + features_idxs = [getFeatureIdx(feature) for feature in features] + elif isinstance(features, int) or isinstance(features, str): + features_idxs = [getFeatureIdx(features)] + else: + raise CatBoostError('Unsupported type for argument \'features\'. Must be one of: int, string, list, list, tuple, tuple') + return features_idxs + + def getAxisParams(borders, feature_name=None): + return { + 'title': 'Bins' if feature_name is None else 'Bins of feature \'{}\''.format(feature_name), + 'tickmode': 'array', + 'tickvals': list(range(len(borders) + 1)), + 'ticktext': ['(-inf, {:.4f}]'.format(borders[0])] + + ['({:.4f}, {:.4f}]'.format(val_1, val_2) + for val_1, val_2 in zip(borders[:-1], borders[1:])] + + ['({:.4f}, +inf)'.format(borders[-1])], + 'showticklabels': False} + + def plot2d(feature_names, borders, predictions): + xaxis = go.layout.XAxis(**getAxisParams(borders[1], feature_name=feature_names[1])) + yaxis = go.layout.YAxis(**getAxisParams(borders[0], feature_name=feature_names[0])) + layout = go.Layout( + title='Partial dependence plot for features {}'.format('\'{}\''.format('\', \''.join(map(str, feature_names)))), + yaxis=yaxis, + xaxis=xaxis + ) + fig = go.Figure(data=go.Heatmap(z=predictions), layout=layout) + return fig + + def plot1d(feature, borders, predictions): + xaxis = go.layout.XAxis(**getAxisParams(borders)) + yaxis = { + 'title': 'Mean Prediction', + 'side': 'left' + } + layout = go.Layout( + title="Partial dependence plot for feature '{}'".format(feature), + yaxis=yaxis, + xaxis=xaxis + ) + fig = go.Figure(data=go.Scatter(y=predictions, mode='lines+markers'), layout=layout) + return fig + + features_idx = getFeatureIndices(features) + borders = [self._get_borders()[idx] for idx in features_idx] + if len(features_idx) not in [1, 2]: + raise CatBoostError('Number of \'features\' should be 1 or 2, got {}'.format(len(features_idx))) + is_2d_plot = len(features_idx) == 2 + + data, _ = self._process_predict_input_data(data, "plot_partial_dependence", thread_count=thread_count) + all_predictions = np.array(self._object._calc_partial_dependence(data, features_idx, thread_count)) + + if is_2d_plot: + all_predictions = all_predictions.reshape([len(x) + 1 for x in borders]) + fig = plot2d(features_idx, borders, all_predictions) + else: + fig = plot1d(features_idx[0], borders[0], all_predictions) + + if plot: + try_plot_offline(fig) + + if plot_file: + save_plot_file(plot_file, "Partial dependence plot for features '{}'".format(features), fig) + + return all_predictions, fig + + def calc_feature_statistics(self, data, target=None, feature=None, prediction_type=None, + cat_feature_values=None, plot=True, max_cat_features_on_plot=10, + thread_count=-1, plot_file=None): + """ + Get statistics for the feature using the model, dataset and target. + To use this function, you should install plotly. + + The catboost model has borders for the float features used in it. The borders divide + feature values into bins, and the model's prediction depends on the number of the bin where the + feature value falls in. + + For float features this function takes model's borders and computes + 1) Mean target value for every bin; + 2) Mean model prediction for every bin; + 3) The number of objects in dataset which fall into each bin; + 4) Predictions on varying feature. For every object, varies the feature value + so that it falls into bin #0, bin #1, ... and counts model predictions. + Then counts average prediction for each bin. + + For categorical features (only one-hot supported) does the same, but takes feature values + provided in cat_feature_values instead of borders. + + Parameters + ---------- + data: numpy.ndarray or pandas.DataFrame or polars.DataFrame or catboost. Pool or dict {'pool_name': pool} if you want several pools + Data to compute statistics on + target: numpy.ndarray or pandas.Series or polars.Series or dict {'pool_name': target} if you want several pools or None + Target corresponding to data + Use only if data is not catboost.Pool. + feature: None, int, string, or list of int or strings + Features indexes or names in pd.DataFrame for which you want to get statistics. + None, if you need statistics for all features. + prediction_type: str + Prediction type used for counting mean_prediction: 'Class', 'Probability' or 'RawFormulaVal'. + If not specified, is derived from the model. + cat_feature_values: list or numpy.ndarray or pandas.Series or + dict: int or string to list or numpy.ndarray or pandas.Series + Contains categorical feature values you need to get statistics on. + Use dict, when parameter 'feature' is a list to specify cat values for different features. + When parameter 'feature' is int or str, you can just pass list of cat values. + plot: bool + Plot statistics. + max_cat_features_on_plot: int + If categorical feature takes more than max_cat_features_on_plot different unique values, + output result on several plots, not more than max_cat_features_on_plot feature values on each. + Used only if plot=True or plot_file is not None. + thread_count: int + Number of threads to use for getting statistics. + plot_file: str + Output file for plot statistics. + + Returns + ------- + dict if parameter 'feature' is int or string, else dict of dicts: + For each unique feature contain + python dict with binarized feature statistics. + For float feature, includes + 'borders' -- borders for the specified feature in model + 'binarized_feature' -- numbers of bins where feature values fall + 'mean_target' -- mean value of target over each bin + 'mean_prediction' -- mean value of model prediction over each bin + 'objects_per_bin' -- number of objects per bin + 'predictions_on_varying_feature' -- averaged over dataset predictions for + varying feature (see above) + For one-hot feature, returns the same, but with 'cat_values' instead of 'borders' + """ + target_is_none = target is None + if not isinstance(data, dict): + data = {'': data} + if not isinstance(target, dict): + target = {'': target} + assert target_is_none or len(data) == len(target), 'inconsistent size of data and target' + assert target_is_none or target.keys() == data.keys(), 'inconsistent pool_names of data and target' + for key in data.keys(): + data[key], _ = self._process_predict_input_data(data[key], "get_binarized_statistics", thread_count, target.get(key, None)) + data, pool_names = list(data.values()), list(data.keys()) + + if prediction_type is None: + prediction_type = 'Probability' if self.get_param('loss_function') in ['CrossEntropy', 'Logloss'] \ + else 'RawFormulaVal' + + if prediction_type not in ['Class', 'Probability', 'RawFormulaVal', 'Exponent']: + raise CatBoostError('Unknown prediction type "{}"'.format(prediction_type)) + + if feature is None: + feature = self.feature_names_ + + if cat_feature_values is None: + cat_feature_values = {} + else: + if not isinstance(cat_feature_values, dict): + if isinstance(feature, list): + raise CatBoostError('cat_feature_values should be dict when features is a list') + else: + cat_feature_values = {feature: cat_feature_values} + + if isinstance(feature, str) or isinstance(feature, int): + features = [feature] + is_for_one_feature = True + else: + features = feature + is_for_one_feature = False + + cat_features_nums = [] + float_features_nums = [] + feature_type_mapper = [] + feature_names = [] + feature_name_to_num = {} + for feature in features: + if not isinstance(feature, int): + if self.feature_names_ is None or feature not in self.feature_names_: + raise CatBoostError('No feature named "{}" in model'.format(feature)) + feature_num = self.feature_names_.index(feature) + else: + feature_num = feature + feature = self.feature_names_[feature_num] + if feature in feature_names: + continue # There is no reason to count statistics twice for the same feature + if feature_num in cat_feature_values: + cat_feature_values[feature] = cat_feature_values[feature_num] + feature_names.append(feature) + feature_name_to_num[feature] = feature_num + feature_type, feature_internal_index = self._object._get_feature_type_and_internal_index(feature_num) + if feature_type == 'categorical': + cat_features_nums.append(feature_internal_index) + feature_type_mapper.append('cat') + else: + float_features_nums.append(feature_internal_index) + feature_type_mapper.append('float') + results = [self._object._get_binarized_statistics( + data_item, + cat_features_nums, + float_features_nums, + prediction_type, + thread_count + ) for data_item in data] + # res = [dict, dict, ..., dict, dict, dict, ..., dict ] + # | stat for cat features | | stat for float features | + statistics_by_feature = defaultdict(list) + to_float_offset = len(cat_features_nums) + cat_index, float_index = 0, to_float_offset + for i, type in enumerate(feature_type_mapper): + feature_name = feature_names[i] + feature_num = feature_name_to_num[feature_name] + if type == 'cat': + if feature_name not in cat_feature_values: + cat_feature_values_ = self._object._get_cat_feature_values(data[0], feature_num) + cat_feature_values_ = [val for val in cat_feature_values_] + else: + cat_feature_values_ = cat_feature_values[feature_name] + if not isinstance(cat_feature_values_, ARRAY_TYPES): + raise CatBoostError( + "Feature '{}' is categorical. " + "Please provide values for which you need statistics in cat_feature_values" + .format(feature) + ) + val_to_hash = dict() + for val in cat_feature_values_: + val_to_hash[val] = self._object._calc_cat_feature_perfect_hash(val, cat_features_nums[cat_index]) + hash_to_val = {hash: val for val, hash in val_to_hash.items()} + for i, res in enumerate(results): + res[cat_index]['cat_values'] = np.array([hash_to_val[i] for i in sorted(hash_to_val.keys())]) + res[cat_index].pop('borders', None) + statistics_by_feature[feature_num].append(res[cat_index]) + cat_index += 1 + else: + for res in results: + statistics_by_feature[feature_num].append(res[float_index]) + float_index += 1 + # now order in statistics_by_feature is the same as in features + + # draw only unique plots + if plot or plot_file is not None: + fig = _plot_feature_statistics( + statistics_by_feature, + pool_names, + self.feature_names_, + max_cat_features_on_plot) + if plot: + try_plot_offline(fig) + + if plot_file is not None: + save_plot_file(plot_file, 'Catboost metrics graph', [fig]) + + for key in statistics_by_feature.keys(): + if len(statistics_by_feature[key]) == 1: + statistics_by_feature[key] = statistics_by_feature[key][0] + + if is_for_one_feature: + return statistics_by_feature[feature_name_to_num[feature_names[0]]] + + # return dict with possible duplicates + # (if features from input contains both str and int values appropriate one feature) + return_stats = {} + for feature in features: + if isinstance(feature, int): + return_stats[feature] = statistics_by_feature[feature] + else: + return_stats[feature] = statistics_by_feature[feature_name_to_num[feature]] + + return return_stats + + def _plot_oblivious_tree(self, splits, leaf_values): + from graphviz import Digraph + graph = Digraph() + + layer_size = 1 + current_size = 0 + + for split_num in range(len(splits) - 1, -2, -1): + for node_num in range(layer_size): + if split_num >= 0: + node_label = splits[split_num].replace('bin=', 'value>', 1).replace('border=', 'value>', 1) + color = 'black' + shape = 'ellipse' + else: + node_label = leaf_values[node_num] + color = 'red' + shape = 'rect' + + try: + node_label = node_label.decode("utf-8") + except Exception: + pass + + graph.node(str(current_size), node_label, color=color, shape=shape) + + if current_size > 0: + parent = (current_size - 1) // 2 + edge_label = 'Yes' if current_size % 2 == 0 else 'No' + graph.edge(str(parent), str(current_size), edge_label) + + current_size += 1 + + layer_size *= 2 + + return graph + + def _plot_nonsymmetric_tree(self, splits, leaf_values, step_nodes, node_to_leaf): + from graphviz import Digraph + graph = Digraph() + + def plot_leaf(node_idx, graph): + cur_id = 'leaf_{}'.format(node_to_leaf[node_idx]) + node_label = leaf_values[node_to_leaf[node_idx]] + graph.node(cur_id, node_label, color='red', shape='rect') + return cur_id + + def plot_subtree(node_idx, graph): + if step_nodes[node_idx] == (0, 0): + return plot_leaf(node_idx, graph) + else: + cur_id = 'node_{}'.format(node_idx) + node_label = splits[node_idx].replace('bin=', 'value>', 1).replace('border=', 'value>', 1) + graph.node(cur_id, node_label, color='black', shape='ellipse') + + if step_nodes[node_idx][0] == 0: + child_id = plot_leaf(node_idx, graph) + else: + child_id = plot_subtree(node_idx + step_nodes[node_idx][0], graph) + graph.edge(cur_id, child_id, 'No') + + if step_nodes[node_idx][1] == 0: + child_id = plot_leaf(node_idx, graph) + else: + child_id = plot_subtree(node_idx + step_nodes[node_idx][1], graph) + graph.edge(cur_id, child_id, 'Yes') + return cur_id + + plot_subtree(0, graph) + return graph + + def plot_tree(self, tree_idx, pool=None): + pool, _ = self._process_predict_input_data(pool, "plot_tree", thread_count=-1) if pool is not None else (None, None) + + splits = self._get_tree_splits(tree_idx, pool) + leaf_values = self._get_tree_leaf_values(tree_idx) + if self._object._is_oblivious(): + return self._plot_oblivious_tree(splits, leaf_values) + else: + step_nodes = self._get_tree_step_nodes(tree_idx) + node_to_leaf = self._get_tree_node_to_leaf(tree_idx) + return self._plot_nonsymmetric_tree(splits, leaf_values, step_nodes, node_to_leaf) + + def _tune_hyperparams(self, param_grid, X, y=None, cv=3, n_iter=10, partition_random_seed=0, + calc_cv_statistics=True, search_by_train_test_split=True, + refit=True, shuffle=True, stratified=None, train_size=0.8, verbose=1, plot=False, plot_file=None, + log_cout=None, log_cerr=None): + + if refit and self.is_fitted(): + raise CatBoostError("Model was fitted before hyperparameters tuning. You can't change hyperparameters of fitted model.") + + with log_fixup(log_cout, log_cerr): + currently_not_supported_params = { + 'ignored_features', + 'input_borders', + 'loss_function', + 'eval_metric' + } + if isinstance(param_grid, Mapping): + param_grid = [param_grid] + + for grid_num, grid in enumerate(param_grid): + _process_synonyms_groups(grid) + grid = _params_type_cast(grid) + + for param in currently_not_supported_params: + if param in grid: + raise CatBoostError("Parameter '{}' is not currently supported in hyperparameter search".format(param)) + + if X is None: + raise CatBoostError("X must not be None") + + if y is None and not isinstance(X, PATH_TYPES + (Pool,)): + raise CatBoostError("y may be None only when X is an instance of catboost.Pool, str or os.PathLike") + + if not isinstance(param_grid, (Mapping, Iterable)): + raise TypeError('Parameter grid is not a dict or a list ({!r})'.format(param_grid)) + + train_params = self._prepare_train_params(X=X, y=y) + params = train_params["params"] + + custom_folds = None + fold_count = 0 + if isinstance(cv, INTEGER_TYPES): + fold_count = cv + loss_function = params.get('loss_function', None) + if stratified is None: + stratified = isinstance(loss_function, STRING_TYPES) and is_cv_stratified_objective(loss_function) + else: + if not hasattr(cv, '__iter__') and not hasattr(cv, 'split'): + raise AttributeError( + "cv should be one of possible things:" + "\n- None, to use the default 3-fold cross validation," + "\n- integer, to specify the number of folds in a (Stratified)KFold" + "\n- one of the scikit-learn splitter classes" + " (https://scikit-learn.org/stable/modules/classes.html#splitter-classes)" + "\n- An iterable yielding (train, test) splits as arrays of indices" + ) + custom_folds = cv + shuffle = False + + if stratified is None: + loss_function = params.get('loss_function', None) + stratified = isinstance(loss_function, STRING_TYPES) and is_cv_stratified_objective(loss_function) + + with plot_wrapper(plot, plot_file, 'Hyperparameters search plot', [_get_train_dir(params)]): + cv_result = self._object._tune_hyperparams( + param_grid, train_params["train_pool"], params, n_iter, + fold_count, partition_random_seed, shuffle, stratified, train_size, + search_by_train_test_split, calc_cv_statistics, custom_folds, verbose + ) + + if refit: + assert not self.is_fitted() + self.set_params(**cv_result['params']) + self.fit(X, y, silent=True) + return cv_result + + def grid_search(self, param_grid, X, y=None, cv=3, partition_random_seed=0, + calc_cv_statistics=True, search_by_train_test_split=True, + refit=True, shuffle=True, stratified=None, train_size=0.8, verbose=True, plot=False, plot_file=None, + log_cout=None, log_cerr=None): + """ + Exhaustive search over specified parameter values for a model. + After calling this method model is fitted and can be used, if not specified otherwise (refit=False). + + Parameters + ---------- + param_grid: dict or list of dictionaries + Dictionary with parameters names (string) as keys and lists of parameter settings + to try as values, or a list of such dictionaries, in which case the grids spanned by each + dictionary in the list are explored. + This enables searching over any sequence of parameter settings. + + X: numpy.ndarray or pandas.DataFrame or polars.DataFrame or catboost.Pool + Data to compute statistics on + + y: list or numpy.ndarray or pandas.DataFrame or pandas.Series or polars.DataFrame or polars.Series, optional (default=None) + Labels of the training data. + If not None, can be a single- or two- dimensional array with either: + - numerical values - for regression (including multiregression), ranking and binary classification problems + - class labels (boolean, integer or string) - for classification (including multiclassification) problems + Use only if X is not catboost.Pool and does not point to a file. + + cv: int, cross-validation generator or an iterable, optional (default=None) + Determines the cross-validation splitting strategy. Possible inputs for cv are: + - None, to use the default 3-fold cross validation, + - integer, to specify the number of folds in a (Stratified)KFold + - one of the scikit-learn splitter classes + (https://scikit-learn.org/stable/modules/classes.html#splitter-classes) + - An iterable yielding (train, test) splits as arrays of indices. + + partition_random_seed: int, optional (default=0) + Use this as the seed value for random permutation of the data. + Permutation is performed before splitting the data for cross validation. + Each seed generates unique data splits. + Used only when cv is None or int. + + search_by_train_test_split: bool, optional (default=True) + If True, source dataset is splitted into train and test parts, models are trained + on the train part and parameters are compared by loss function score on the test part. + After that, if calc_cv_statistics=true, statistics on metrics are calculated + using cross-validation using best parameters and the model is fitted with these parameters. + + If False, every iteration of grid search evaluates results on cross-validation. + It is recommended to set parameter to True for large datasets, and to False for small datasets. + + calc_cv_statistics: bool, optional (default=True) + The parameter determines whether quality should be estimated. + using cross-validation with the found best parameters. Used only when search_by_train_test_split=True. + + refit: bool (default=True) + Refit an estimator using the best found parameters on the whole dataset. + + shuffle: bool, optional (default=True) + Shuffle the dataset objects before parameters searching. + + stratified: bool, optional (default=None) + Perform stratified sampling. True for classification and False otherwise. + Currently supported only for final cross-validation. + + train_size: float, optional (default=0.8) + Should be between 0.0 and 1.0 and represent the proportion of the dataset to include in the train split. + + verbose: bool or int, optional (default=True) + If verbose is int, it determines the frequency of writing metrics to output + verbose==True is equal to verbose==1 + When verbose==False, there is no messages + + plot : bool, optional (default=False) + If True, draw train and eval error for every set of parameters in Jupyter notebook + + plot_file : file-like or str, optional (default=None) + If not None, save train and eval error for every set of parameters to file + + log_cout: output stream or callback for logging (default=None) + If None is specified, sys.stdout is used + + log_cerr: error stream or callback for logging (default=None) + If None is specified, sys.stderr is used + + Returns + ------- + dict with two fields: + 'params': dict of best found parameters + 'cv_results': dict or pandas.core.frame.DataFrame with cross-validation results + columns are: test-error-mean test-error-std train-error-mean train-error-std + """ + if isinstance(param_grid, Mapping): + param_grid = [param_grid] + for grid in param_grid: + if not isinstance(grid, Mapping): + raise TypeError('Parameter grid is not a dict ({!r})'.format(grid)) + for key in grid: + if not isinstance(grid[key], Iterable): + raise TypeError('Parameter grid value is not iterable (key={!r}, value={!r})'.format(key, grid[key])) + + return self._tune_hyperparams( + param_grid=param_grid, X=X, y=y, cv=cv, n_iter=-1, + partition_random_seed=partition_random_seed, calc_cv_statistics=calc_cv_statistics, + search_by_train_test_split=search_by_train_test_split, refit=refit, shuffle=shuffle, + stratified=stratified, train_size=train_size, verbose=verbose, plot=plot, plot_file=plot_file, + log_cout=log_cout, log_cerr=log_cerr, + ) + + def randomized_search(self, param_distributions, X, y=None, cv=3, n_iter=10, partition_random_seed=0, + calc_cv_statistics=True, search_by_train_test_split=True, refit=True, + shuffle=True, stratified=None, train_size=0.8, verbose=True, plot=False, plot_file=None, + log_cout=None, log_cerr=None): + """ + Randomized search on hyper parameters. + After calling this method model is fitted and can be used, if not specified otherwise (refit=False). + + In contrast to grid_search, not all parameter values are tried out, + but rather a fixed number of parameter settings is sampled from the specified distributions. + The number of parameter settings that are tried is given by n_iter. + + Parameters + ---------- + param_distributions: dict + Dictionary with parameters names (string) as keys and distributions or lists of parameters to try. + Distributions must provide a rvs method for sampling (such as those from scipy.stats.distributions). + If a list is given, it is sampled uniformly. + + X: numpy.ndarray or pandas.DataFrame or polars.DataFrame or catboost.Pool + Data to compute statistics on + + y: list or numpy.ndarray or pandas.DataFrame or pandas.Series or polars.DataFrame or polars.Series, optional (default=None) + Labels of the training data. + If not None, can be a single- or two- dimensional array with either: + - numerical values - for regression (including multiregression), ranking and binary classification problems + - class labels (boolean, integer or string) - for classification (including multiclassification) problems + Use only if X is not catboost.Pool and does not point to a file. + + cv: int, cross-validation generator or an iterable, optional (default=None) + Determines the cross-validation splitting strategy. Possible inputs for cv are: + - None, to use the default 3-fold cross validation, + - integer, to specify the number of folds in a (Stratified)KFold + - one of the scikit-learn splitter classes + (https://scikit-learn.org/stable/modules/classes.html#splitter-classes) + - An iterable yielding (train, test) splits as arrays of indices. + + n_iter: int + Number of parameter settings that are sampled. + n_iter trades off runtime vs quality of the solution. + + partition_random_seed: int, optional (default=0) + Use this as the seed value for random permutation of the data. + Permutation is performed before splitting the data for cross validation. + Each seed generates unique data splits. + Used only when cv is None or int. + + search_by_train_test_split: bool, optional (default=True) + If True, source dataset is splitted into train and test parts, models are trained + on the train part and parameters are compared by loss function score on the test part. + After that, if calc_cv_statistics=true, statistics on metrics are calculated + using cross-validation using best parameters and the model is fitted with these parameters. + + If False, every iteration of grid search evaluates results on cross-validation. + It is recommended to set parameter to True for large datasets, and to False for small datasets. + + calc_cv_statistics: bool, optional (default=True) + The parameter determines whether quality should be estimated. + using cross-validation with the found best parameters. Used only when search_by_train_test_split=True. + + refit: bool (default=True) + Refit an estimator using the best found parameters on the whole dataset. + + shuffle: bool, optional (default=True) + Shuffle the dataset objects before parameters searching. + + stratified: bool, optional (default=None) + Perform stratified sampling. True for classification and False otherwise. + Currently supported only for cross-validation. + + train_size: float, optional (default=0.8) + Should be between 0.0 and 1.0 and represent the proportion of the dataset to include in the train split. + + verbose: bool or int, optional (default=True) + If verbose is int, it determines the frequency of writing metrics to output + verbose==True is equal to verbose==1 + When verbose==False, there is no messages + + plot : bool, optional (default=False) + If True, draw train and eval error for every set of parameters in Jupyter notebook + + plot_file : file-like or str, optional (default=None) + If not None, save train and eval error for every set of parameters to file + + log_cout: output stream or callback for logging (default=None) + If None is specified, sys.stdout is used + + log_cerr: error stream or callback for logging (default=None) + If None is specified, sys.stderr is used + + Returns + ------- + dict with two fields: + 'params': dict of best found parameters + 'cv_results': dict or pandas.core.frame.DataFrame with cross-validation results + columns are: test-error-mean test-error-std train-error-mean train-error-std + """ + if n_iter <= 0: + assert CatBoostError("n_iter should be a positive number") + if not isinstance(param_distributions, Mapping): + assert CatBoostError("param_distributions should be a dictionary") + for key in param_distributions: + if not isinstance(param_distributions[key], Iterable) and not hasattr(param_distributions[key], "rvs"): + raise TypeError('Parameter grid value is not iterable and do not have \'rvs\' method (key={!r}, value={!r})'.format(key, param_distributions[key])) + + return self._tune_hyperparams( + param_grid=param_distributions, X=X, y=y, cv=cv, n_iter=n_iter, + partition_random_seed=partition_random_seed, calc_cv_statistics=calc_cv_statistics, + search_by_train_test_split=search_by_train_test_split, refit=refit, shuffle=shuffle, + stratified=stratified, train_size=train_size, verbose=verbose, plot=plot, plot_file=plot_file, + log_cout=log_cout, log_cerr=log_cerr, + ) + + def select_features(self, X, y=None, eval_set=None, features_for_select=None, num_features_to_select=None, + algorithm=None, steps=None, shap_calc_type=None, train_final_model=True, verbose=None, + logging_level=None, plot=False, plot_file=None, log_cout=None, log_cerr=None, + grouping=None, features_tags_for_select=None, num_features_tags_to_select=None): + """ + Select best features from pool according to loss value. + + Parameters + ---------- + X : catboost.Pool or list or numpy.ndarray or pandas.DataFrame or pandas.Series or polars.DataFrame + If not catboost.Pool, 2 dimensional Feature matrix or string - file with dataset. + + y : list or numpy.ndarray or pandas.DataFrame or pandas.Series or polars.DataFrame or polars.Series, optional (default=None) + Labels of the training data. + If not None, can be a single- or two- dimensional array with either: + - numerical values - for regression (including multiregression), ranking and binary classification problems + - class labels (boolean, integer or string) - for classification (including multiclassification) problems + Use only if X is not catboost.Pool and does not point to a file. + + eval_set : catboost.Pool or list of catboost.Pool or tuple (X, y) or list [(X, y)], optional (default=None) + Validation dataset or datasets for metrics calculation and possibly early stopping. + + features_for_select : str or list of feature indices, names or ranges + (for grouping = Individual) + Which features should participate in the selection. + Format examples: + - [0, 2, 3, 4, 17] + - [0, "2-4", 17] (both ends in ranges are inclusive) + - "0,2-4,20" + - ["Name0", "Name2", "Name3", "Name4", "Name20"] + + num_features_to_select : positive int + (for grouping = Individual) + How many features to select from features_for_select. + + algorithm : EFeaturesSelectionAlgorithm or string, optional (default=RecursiveByShapValues) + Which algorithm to use for features selection. + Possible values: + - RecursiveByPredictionValuesChange + Use prediction values change as feature strength, eliminate batch of features at once. + - RecursiveByLossFunctionChange + Use loss function change as feature strength, eliminate batch of features at each step. + - RecursiveByShapValues + Use shap values to estimate loss function change, eliminate features one by one. + + steps : positive int, optional (default=1) + How many steps should be performed. In other words, how many times a full model will be trained. + More steps give more accurate results. + + shap_calc_type : EShapCalcType or string, optional (default=Regular) + Which method to use for calculation of shap values. + Possible values: + - Regular + Calculate regular SHAP values + - Approximate + Calculate approximate SHAP values + - Exact + Calculate exact SHAP values + + train_final_model : bool, optional (default=True) + Need to fit model with selected features. + + verbose : bool or int + If verbose is bool, then if set to True, logging_level is set to Verbose, + if set to False, logging_level is set to Silent. + If verbose is int, it determines the frequency of writing metrics to output and + logging_level is set to Verbose. + + logging_level : string, optional (default=None) + Possible values: + - 'Silent' + - 'Verbose' + - 'Info' + - 'Debug' + + plot : bool, optional (default=False) + If True, draw train and eval error in Jupyter notebook. + + plot_file : file-like or str, optional (default=None) + If not None, save train and eval error graphs to file + + log_cout: output stream or callback for logging (default=None) + If None is specified, sys.stdout is used + + log_cerr: error stream or callback for logging (default=None) + If None is specified, sys.stderr is used + + grouping : EFeaturesSelectionGrouping or string, optional (default=Individual) + Which grouping to use for features selection. + Possible values: + - Individual + Select individual features + - ByTags + Select feature groups (marked by tags) + + features_tags_for_select : list of strings + (for grouping = ByTags) + Which features tags should participate in the selection. + + num_features_tags_to_select : positive int + (for grouping = ByTags) + How many features tags to select from features_tags_for_select. + + Returns + ------- + dict with fields: + 'selected_features': list of selected features indices + 'eliminated_features': list of eliminated features indices + 'selected_features_tags': list of selected features tags (optional, present if grouping == ByTags) + 'eliminated_features_tags': list of selected features tags (optional, present if grouping == ByTags) + """ + if train_final_model and self.is_fitted(): + raise CatBoostError("Model was already fitted. Set train_final_model to False or use not fitted model.") + if X is None: + raise CatBoostError("X must not be None") + if y is None and not isinstance(X, PATH_TYPES + (Pool,)): + raise CatBoostError("y may be None only when X is an instance of catboost.Pool, str or os.PathLike.") + + with log_fixup(log_cout, log_cerr): + train_params = self._prepare_train_params(X=X, y=y, eval_set=eval_set, verbose=verbose, logging_level=logging_level) + params = train_params["params"] + + if grouping is None: + grouping = EFeaturesSelectionGrouping.Individual + else: + grouping = enum_from_enum_or_str(EFeaturesSelectionGrouping, grouping).value + params["features_selection_grouping"] = grouping + + if grouping == EFeaturesSelectionGrouping.Individual: + if isinstance(features_for_select, Iterable) and not isinstance(features_for_select, STRING_TYPES): + features_for_select = ",".join(map(str, features_for_select)) + if features_for_select is None: + raise CatBoostError("You should specify features_for_select") + if features_tags_for_select is not None: + raise CatBoostError("You should not specify features_tags_for_select when grouping is Individual") + if num_features_to_select is None: + raise CatBoostError("You should specify num_features_to_select") + if num_features_tags_to_select is not None: + raise CatBoostError("You should not specify num_features_tags_to_select when grouping is Individual") + params["features_for_select"] = features_for_select + params["num_features_to_select"] = num_features_to_select + else: # ByTag + if features_tags_for_select is None: + raise CatBoostError("You should specify features_tags_for_select") + if not isinstance(features_tags_for_select, Sequence): + raise CatBoostError("features_tags_for_select must be a list of strings") + if features_for_select is not None: + raise CatBoostError("You should not specify features_for_select when grouping is ByTags") + if num_features_tags_to_select is None: + raise CatBoostError("You should specify num_features_tags_to_select") + if num_features_to_select is not None: + raise CatBoostError("You should not specify num_features_to_select when grouping is ByTags") + params["features_tags_for_select"] = features_tags_for_select + params["num_features_tags_to_select"] = num_features_tags_to_select + + objective = params.get("loss_function") + is_custom_objective = objective is not None and not isinstance(objective, string_types) + if is_custom_objective: + raise CatBoostError("Custom objective is not supported for features selection") + + if algorithm is not None: + params["features_selection_algorithm"] = enum_from_enum_or_str(EFeaturesSelectionAlgorithm, algorithm).value + if steps is not None: + params["features_selection_steps"] = steps + if shap_calc_type is not None: + params["shap_calc_type"] = enum_from_enum_or_str(EShapCalcType, shap_calc_type).value + if train_final_model: + params["train_final_model"] = True + + train_pool = train_params["train_pool"] + test_pool = None + if len(train_params["eval_sets"]) > 1: + raise CatBoostError("Multiple eval sets are not supported for features selection") + elif len(train_params["eval_sets"]) == 1: + test_pool = train_params["eval_sets"][0] + + train_dir = _get_train_dir(params) + create_dir_if_not_exist(train_dir) + plot_dirs = [] + for step in range(steps or 1): + plot_dirs.append(os.path.join(train_dir, 'model-{}'.format(step))) + if train_final_model: + plot_dirs.append(os.path.join(train_dir, 'model-final')) + for plot_dir in plot_dirs: + create_dir_if_not_exist(plot_dir) + + with plot_wrapper(plot, plot_file=plot_file, plot_title='Select features plot', train_dirs=plot_dirs): + summary = self._object._select_features(train_pool, test_pool, params) + + if train_final_model: + self._set_trained_model_attributes() + + if plot: + figures = plot_features_selection_loss_graphs(summary) + figures['features'].show() + if 'features_tags' in figures: + figures['features_tags'].show() + + return summary + + def _convert_to_asymmetric_representation(self): + self._object._convert_oblivious_to_asymmetric() + + +class CatBoostClassifier(CatBoost): + """ + Implementation of the scikit-learn API for CatBoost classification. + + Parameters + ---------- + iterations : int, [default=500] + Max count of trees. + range: [1,+inf) + learning_rate : float, [default value is selected automatically for binary classification with other parameters set to default. In all other cases default is 0.03] + Step size shrinkage used in update to prevents overfitting. + range: (0,1] + depth : int, [default=6] + Depth of a tree. All trees are the same depth. + range: [1,16] + l2_leaf_reg : float, [default=3.0] + Coefficient at the L2 regularization term of the cost function. + range: [0,+inf) + model_size_reg : float, [default=None] + Model size regularization coefficient. + range: [0,+inf) + rsm : float, [default=None] + Subsample ratio of columns when constructing each tree. + range: (0,1] + loss_function : string or object, [default='Logloss'] + The metric to use in training and also selector of the machine learning + problem to solve. If string, then the name of a supported metric, + optionally suffixed with parameter description. + If object, it shall provide methods 'calc_ders_range' or 'calc_ders_multi'. + border_count : int, [default = 254 for training on CPU or 128 for training on GPU] + The number of partitions in numeric features binarization. Used in the preliminary calculation. + range: [1,65535] on CPU, [1,255] on GPU + feature_border_type : string, [default='GreedyLogSum'] + The binarization mode in numeric features binarization. Used in the preliminary calculation. + Possible values: + - 'Median' + - 'Uniform' + - 'UniformAndQuantiles' + - 'GreedyLogSum' + - 'MaxLogSum' + - 'MinEntropy' + per_float_feature_quantization : list of strings, [default=None] + List of float binarization descriptions. + Format : described in documentation on catboost.ai + Example 1: ['0:1024'] means that feature 0 will have 1024 borders. + Example 2: ['0:border_count=1024', '1:border_count=1024', ...] means that two first features have 1024 borders. + Example 3: ['0:nan_mode=Forbidden,border_count=32,border_type=GreedyLogSum', + '1:nan_mode=Forbidden,border_count=32,border_type=GreedyLogSum'] - defines more quantization properties for first two features. + input_borders : string or os.PathLike, [default=None] + input file with borders used in numeric features binarization. + output_borders : string, [default=None] + output file for borders that were used in numeric features binarization. + fold_permutation_block : int, [default=1] + To accelerate the learning. + The recommended value is within [1, 256]. On small samples, must be set to 1. + range: [1,+inf) + od_pval : float, [default=None] + Use overfitting detector to stop training when reaching a specified threshold. + Can be used only with eval_set. + range: [0,1] + od_wait : int, [default=None] + Number of iterations which overfitting detector will wait after new best error. + od_type : string, [default=None] + Type of overfitting detector which will be used in program. + Posible values: + - 'IncToDec' + - 'Iter' + For 'Iter' type od_pval must not be set. + If None, then od_type=IncToDec. + nan_mode : string, [default=None] + Way to process missing values for numeric features. + Possible values: + - 'Forbidden' - raises an exception if there is a missing value for a numeric feature in a dataset. + - 'Min' - each missing value will be processed as the minimum numerical value. + - 'Max' - each missing value will be processed as the maximum numerical value. + If None, then nan_mode=Min. + counter_calc_method : string, [default=None] + The method used to calculate counters for dataset with Counter type. + Possible values: + - 'PrefixTest' - only objects up to current in the test dataset are considered + - 'FullTest' - all objects are considered in the test dataset + - 'SkipTest' - Objects from test dataset are not considered + - 'Full' - all objects are considered for both learn and test dataset + If None, then counter_calc_method=PrefixTest. + leaf_estimation_iterations : int, [default=None] + The number of steps in the gradient when calculating the values in the leaves. + If None, then leaf_estimation_iterations=1. + range: [1,+inf) + leaf_estimation_method : string, [default=None] + The method used to calculate the values in the leaves. + Possible values: + - 'Newton' + - 'Gradient' + thread_count : int, [default=None] + Number of parallel threads used to run CatBoost. + If None or -1, then the number of threads is set to the number of CPU cores. + range: [1,+inf) + random_seed : int, [default=None] + Random number seed. + If None, 0 is used. + range: [0,+inf) + use_best_model : bool, [default=None] + To limit the number of trees in predict() using information about the optimal value of the error function. + Can be used only with eval_set. + best_model_min_trees : int, [default=None] + The minimal number of trees the best model should have. + verbose: bool + When set to True, logging_level is set to 'Verbose'. + When set to False, logging_level is set to 'Silent'. + silent: bool, synonym for verbose + logging_level : string, [default='Verbose'] + Possible values: + - 'Silent' + - 'Verbose' + - 'Info' + - 'Debug' + metric_period : int, [default=1] + The frequency of iterations to print the information to stdout. The value should be a positive integer. + simple_ctr: list of strings, [default=None] + Binarization settings for categorical features. + Format : see documentation + Example: ['Borders:CtrBorderCount=5:Prior=0:Prior=0.5', 'BinarizedTargetMeanValue:TargetBorderCount=10:TargetBorderType=MinEntropy', ...] + CTR types: + CPU and GPU + - 'Borders' + - 'Buckets' + CPU only + - 'BinarizedTargetMeanValue' + - 'Counter' + GPU only + - 'FloatTargetMeanValue' + - 'FeatureFreq' + Number_of_borders, binarization type, target borders and binarizations, priors are optional parametrs + combinations_ctr: list of strings, [default=None] + per_feature_ctr: list of strings, [default=None] + ctr_target_border_count: int, [default=None] + Maximum number of borders used in target binarization for categorical features that need it. + If TargetBorderCount is specified in 'simple_ctr', 'combinations_ctr' or 'per_feature_ctr' option it + overrides this value. + range: [1, 255] + ctr_leaf_count_limit : int, [default=None] + The maximum number of leaves with categorical features. + If the number of leaves exceeds the specified limit, some leaves are discarded. + The leaves to be discarded are selected as follows: + - The leaves are sorted by the frequency of the values. + - The top N leaves are selected, where N is the value specified in the parameter. + - All leaves starting from N+1 are discarded. + This option reduces the resulting model size + and the amount of memory required for training. + Note that the resulting quality of the model can be affected. + range: [1,+inf) (for zero limit use ignored_features) + store_all_simple_ctr : bool, [default=None] + Ignore categorical features, which are not used in feature combinations, + when choosing candidates for exclusion. + Use this parameter with ctr_leaf_count_limit only. + max_ctr_complexity : int, [default=4] + The maximum number of Categ features that can be combined. + range: [0,+inf) + has_time : bool, [default=False] + To use the order in which objects are represented in the input data + (do not perform a random permutation of the dataset at the preprocessing stage). + allow_const_label : bool, [default=False] + To allow the constant label value in dataset. + target_border: float, [default=None] + Border for target binarization. + classes_count : int, [default=None] + The upper limit for the numeric class label. + Defines the number of classes for multiclassification. + Only non-negative integers can be specified. + The given integer should be greater than any of the target values. + If this parameter is specified the labels for all classes in the input dataset + should be smaller than the given value. + If several of 'classes_count', 'class_weights', 'class_names' parameters are defined + the numbers of classes specified by each of them must be equal. + class_weights : list or dict, [default=None] + Classes weights. The values are used as multipliers for the object weights. + If None, all classes are supposed to have weight one. + If list - class weights in order of class_names or sequential classes if class_names is undefined + If dict - dict of class_name -> class_weight. + If several of 'classes_count', 'class_weights', 'class_names' parameters are defined + the numbers of classes specified by each of them must be equal. + auto_class_weights : string [default=None] + Enables automatic class weights calculation. Possible values: + - Balanced # weight = maxSummaryClassWeight / summaryClassWeight, statistics determined from train pool + - SqrtBalanced # weight = sqrt(maxSummaryClassWeight / summaryClassWeight) + class_names: list of strings, [default=None] + Class names. Allows to redefine the default values for class labels (integer numbers). + If several of 'classes_count', 'class_weights', 'class_names' parameters are defined + the numbers of classes specified by each of them must be equal. + one_hot_max_size : int, [default=None] + Convert the feature to float + if the number of different values that it takes exceeds the specified value. + Ctrs are not calculated for such features. + random_strength : float, [default=1] + Score standard deviation multiplier. + random_score_type : string [default=None] + Type of random noise added to scores. + Possible values: + - 'Gumbel' - Gumbel-distributed + - 'NormalWithModelSizeDecrease' - Normally-distributed with deviation decreasing with model iteration count + If None than 'NormalWithModelSizeDecrease' will be used by default. + name : string, [default='experiment'] + The name that should be displayed in the visualization tools. + ignored_features : list, [default=None] + Indices or names of features that should be excluded when training. + train_dir : string or os.PathLike, [default=None] + The directory in which you want to record generated in the process of learning files. + custom_metric : string or list of strings, [default=None] + To use your own metric function. + custom_loss: alias to custom_metric + eval_metric : string or object, [default=None] + To optimize your custom metric in loss. + bagging_temperature : float, [default=None] + Controls intensity of Bayesian bagging. The higher the temperature the more aggressive bagging is. + Typical values are in range [0, 1] (0 - no bagging, 1 - default). + save_snapshot : bool, [default=None] + Enable progress snapshotting for restoring progress after crashes or interruptions + snapshot_file : string or os.PathLike, [default=None] + Learn progress snapshot file path, if None will use default filename + snapshot_interval: int, [default=600] + Interval between saving snapshots (seconds) + fold_len_multiplier : float, [default=None] + Fold length multiplier. Should be greater than 1 + used_ram_limit : string or number, [default=None] + Set a limit on memory consumption (value like '1.2gb' or 1.2e9). + WARNING: Currently this option affects CTR memory usage only. + gpu_ram_part : float, [default=0.95] + Fraction of the GPU RAM to use for training, a value from (0, 1]. + pinned_memory_size: int [default=None] + Size of additional CPU pinned memory used for GPU learning, + usually is estimated automatically, thus usually should not be set. + allow_writing_files : bool, [default=True] + If this flag is set to False, no files with different diagnostic info will be created during training. + With this flag no snapshotting can be done. Plus visualisation will not + work, because visualisation uses files that are created and updated during training. + final_ctr_computation_mode : string, [default='Default'] + Possible values: + - 'Default' - Compute final ctrs for all pools. + - 'Skip' - Skip final ctr computation. WARNING: model without ctrs can't be applied. + approx_on_full_history : bool, [default=False] + If this flag is set to True, each approximated value is calculated using all the preceeding rows in the fold (slower, more accurate). + If this flag is set to False, each approximated value is calculated using only the beginning 1/fold_len_multiplier fraction of the fold (faster, slightly less accurate). + boosting_type : string, default value depends on object count and feature count in train dataset and on learning mode. + Boosting scheme. + Possible values: + - 'Ordered' - Gives better quality, but may slow down the training. + - 'Plain' - The classic gradient boosting scheme. May result in quality degradation, but does not slow down the training. + task_type : string, [default=None] + The calcer type used to train the model. + Possible values: + - 'CPU' + - 'GPU' + device_config : string, [default=None], deprecated, use devices instead + devices : list or string, [default=None], GPU devices to use. + String format is: '0' for 1 device or '0:1:3' for multiple devices or '0-3' for range of devices. + List format is : [0] for 1 device or [0,1,3] for multiple devices. + + bootstrap_type : string, Bayesian, Bernoulli, Poisson, MVS. + Default bootstrap is Bayesian for GPU and MVS for CPU. + Poisson bootstrap is supported only on GPU. + MVS bootstrap is supported only on GPU. + + subsample : float, [default=None] + Sample rate for bagging. This parameter can be used Poisson or Bernoulli bootstrap types. + + mvs_reg : float, [default is set automatically at each iteration based on gradient distribution] + Regularization parameter for MVS sampling algorithm + + monotone_constraints : list or numpy.ndarray or string or dict, [default=None] + Monotone constraints for features. + + feature_weights : list or numpy.ndarray or string or dict, [default=None] + Coefficient to multiply split gain with specific feature use. Should be non-negative. + + penalties_coefficient : float, [default=1] + Common coefficient for all penalties. Should be non-negative. + + first_feature_use_penalties : list or numpy.ndarray or string or dict, [default=None] + Penalties to first use of specific feature in model. Should be non-negative. + + per_object_feature_penalties : list or numpy.ndarray or string or dict, [default=None] + Penalties for first use of feature for each object. Should be non-negative. + + sampling_frequency : string, [default=PerTree] + Frequency to sample weights and objects when building trees. + Possible values: + - 'PerTree' - Before constructing each new tree + - 'PerTreeLevel' - Before choosing each new split of a tree + + sampling_unit : string, [default='Object']. + Possible values: + - 'Object' + - 'Group' + The parameter allows to specify the sampling scheme: + sample weights for each object individually or for an entire group of objects together. + + dev_score_calc_obj_block_size: int, [default=5000000] + CPU only. Size of block of samples in score calculation. Should be > 0 + Used only for learning speed tuning. + Changing this parameter can affect results due to numerical accuracy differences + + dev_efb_max_buckets : int, [default=1024] + CPU only. Maximum bucket count in exclusive features bundle. Should be in an integer between 0 and 65536. + Used only for learning speed tuning. + + sparse_features_conflict_fraction : float, [default=0.0] + CPU only. Maximum allowed fraction of conflicting non-default values for features in exclusive features bundle. + Should be a real value in [0, 1) interval. + + grow_policy : string, [SymmetricTree,Lossguide,Depthwise], [default=SymmetricTree] + The tree growing policy. It describes how to perform greedy tree construction. + + min_data_in_leaf : int, [default=1]. + The minimum training samples count in leaf. + CatBoost will not search for new splits in leaves with samples count less than min_data_in_leaf. + This parameter is used only for Depthwise and Lossguide growing policies. + + max_leaves : int, [default=31], + The maximum leaf count in resulting tree. + This parameter is used only for Lossguide growing policy. + + score_function : string, possible values L2, Cosine, NewtonL2, NewtonCosine, [default=Cosine] + For growing policy Lossguide default=NewtonL2. + GPU only. Score that is used during tree construction to select the next tree split. + + max_depth : int, Synonym for depth. + + n_estimators : int, synonym for iterations. + + num_trees : int, synonym for iterations. + + num_boost_round : int, synonym for iterations. + + colsample_bylevel : float, synonym for rsm. + + random_state : int, synonym for random_seed. + + reg_lambda : float, synonym for l2_leaf_reg. + + objective : string, synonym for loss_function. + + num_leaves : int, synonym for max_leaves. + + min_child_samples : int, synonym for min_data_in_leaf + + eta : float, synonym for learning_rate. + + max_bin : float, synonym for border_count. + + scale_pos_weight : float, synonym for class_weights. + Can be used only for binary classification. Sets weight multiplier for + class 1 to scale_pos_weight value. + + metadata : dict, string to string key-value pairs to be stored in model metadata storage + + early_stopping_rounds : int + Synonym for od_wait. Only one of these parameters should be set. + + cat_features : list or numpy.ndarray, [default=None] + If not None, giving the list of Categ features indices or names (names are represented as strings). + If it contains feature names, feature names must be defined for the training dataset passed to 'fit'. + + text_features : list or numpy.ndarray, [default=None] + If not None, giving the list of Text features indices or names (names are represented as strings). + If it contains feature names, feature names must be defined for the training dataset passed to 'fit'. + + embedding_features : list or numpy.ndarray, [default=None] + If not None, giving the list of Embedding features indices or names (names are represented as strings). + If it contains feature names, feature names must be defined for the training dataset passed to 'fit'. + + leaf_estimation_backtracking : string, [default=None] + Type of backtracking during gradient descent. + Possible values: + - 'No' - never backtrack; supported on CPU and GPU + - 'AnyImprovement' - reduce the descent step until the value of loss function is less than before the step; supported on CPU and GPU + - 'Armijo' - reduce the descent step until Armijo condition is satisfied; supported on GPU only + + model_shrink_rate : float, [default=0] + This parameter enables shrinkage of model at the start of each iteration. CPU only. + For Constant mode shrinkage coefficient is calculated as (1 - model_shrink_rate * learning_rate). + For Decreasing mode shrinkage coefficient is calculated as (1 - model_shrink_rate / iteration). + Shrinkage coefficient should be in [0, 1). + + model_shrink_mode : string, [default=None] + Mode of shrinkage coefficient calculation. CPU only. + Possible values: + - 'Constant' - Shrinkage coefficient is constant at each iteration. + - 'Decreasing' - Shrinkage coefficient decreases at each iteration. + + langevin : bool, [default=False] + Enables the Stochastic Gradient Langevin Boosting. CPU only. + + diffusion_temperature : float, [default=0] + Langevin boosting diffusion temperature. CPU only. + + posterior_sampling : bool, [default=False] + Set group of parameters for further use Uncertainty prediction: + - Langevin = True + - Model Shrink Rate = 1/(2N), where N is dataset size + - Model Shrink Mode = Constant + - Diffusion-temperature = N, where N is dataset size. CPU only. + + boost_from_average : bool, [default=True for RMSE, False for other losses] + Enables to initialize approx values by best constant value for specified loss function. + Available for RMSE, Logloss, CrossEntropy, Quantile and MAE. + + tokenizers : list of dicts, + Each dict is a tokenizer description. Example: + ``` + [ + { + 'tokenizer_id': 'Tokenizer', # Tokeinzer identifier. + 'lowercasing': 'false', # Possible values: 'true', 'false'. + 'number_process_policy': 'LeaveAsIs', # Possible values: 'Skip', 'LeaveAsIs', 'Replace'. + 'number_token': '%', # Rarely used character. Used in conjunction with Replace NumberProcessPolicy. + 'separator_type': 'ByDelimiter', # Possible values: 'ByDelimiter', 'BySense'. + 'delimiter': ' ', # Used in conjunction with ByDelimiter SeparatorType. + 'split_by_set': 'false', # Each single character in delimiter used as individual delimiter. + 'skip_empty': 'true', # Possible values: 'true', 'false'. + 'token_types': ['Word', 'Number', 'Unknown'], # Used in conjunction with BySense SeparatorType. + # Possible values: 'Word', 'Number', 'Punctuation', 'SentenceBreak', 'ParagraphBreak', 'Unknown'. + 'subtokens_policy': 'SingleToken', # Possible values: + # 'SingleToken' - All subtokens are interpreted as single token). + # 'SeveralTokens' - All subtokens are interpreted as several token. + }, + ... + ] + ``` + + dictionaries : list of dicts, + Each dict is a tokenizer description. Example: + ``` + [ + { + 'dictionary_id': 'Dictionary', # Dictionary identifier. + 'token_level_type': 'Word', # Possible values: 'Word', 'Letter'. + 'gram_order': '1', # 1 for Unigram, 2 for Bigram, ... + 'skip_step': '0', # 1 for 1-skip-gram, ... + 'end_of_word_token_policy': 'Insert', # Possible values: 'Insert', 'Skip'. + 'end_of_sentence_token_policy': 'Skip', # Possible values: 'Insert', 'Skip'. + 'occurrence_lower_bound': '3', # The lower bound of token occurrences in the text to include it in the dictionary. + 'max_dictionary_size': '50000', # The max dictionary size. + }, + ... + ] + ``` + + feature_calcers : list of strings, + Each string is a calcer description. Example: + ``` + [ + 'NaiveBayes', + 'BM25', + 'BoW:top_tokens_count=2000', + ] + ``` + + text_processing : dict, + Text processging description. + + eval_fraction : float, [default=None] + Fraction of the train dataset to be used as the evaluation dataset. + """ + + _estimator_type = 'classifier' + + def __init__( + self, + iterations=None, + learning_rate=None, + depth=None, + l2_leaf_reg=None, + model_size_reg=None, + rsm=None, + loss_function=None, + border_count=None, + feature_border_type=None, + per_float_feature_quantization=None, + input_borders=None, + output_borders=None, + fold_permutation_block=None, + od_pval=None, + od_wait=None, + od_type=None, + nan_mode=None, + counter_calc_method=None, + leaf_estimation_iterations=None, + leaf_estimation_method=None, + thread_count=None, + random_seed=None, + use_best_model=None, + best_model_min_trees=None, + verbose=None, + silent=None, + logging_level=None, + metric_period=None, + ctr_leaf_count_limit=None, + store_all_simple_ctr=None, + max_ctr_complexity=None, + has_time=None, + allow_const_label=None, + target_border=None, + classes_count=None, + class_weights=None, + auto_class_weights=None, + class_names=None, + one_hot_max_size=None, + random_strength=None, + random_score_type=None, + name=None, + ignored_features=None, + train_dir=None, + custom_loss=None, + custom_metric=None, + eval_metric=None, + bagging_temperature=None, + save_snapshot=None, + snapshot_file=None, + snapshot_interval=None, + fold_len_multiplier=None, + used_ram_limit=None, + gpu_ram_part=None, + pinned_memory_size=None, + allow_writing_files=None, + final_ctr_computation_mode=None, + approx_on_full_history=None, + boosting_type=None, + simple_ctr=None, + combinations_ctr=None, + per_feature_ctr=None, + ctr_description=None, + ctr_target_border_count=None, + task_type=None, + device_config=None, + devices=None, + bootstrap_type=None, + subsample=None, + mvs_reg=None, + sampling_unit=None, + sampling_frequency=None, + dev_score_calc_obj_block_size=None, + dev_efb_max_buckets=None, + sparse_features_conflict_fraction=None, + max_depth=None, + n_estimators=None, + num_boost_round=None, + num_trees=None, + colsample_bylevel=None, + random_state=None, + reg_lambda=None, + objective=None, + eta=None, + max_bin=None, + scale_pos_weight=None, + gpu_cat_features_storage=None, + data_partition=None, + metadata=None, + early_stopping_rounds=None, + cat_features=None, + grow_policy=None, + min_data_in_leaf=None, + min_child_samples=None, + max_leaves=None, + num_leaves=None, + score_function=None, + leaf_estimation_backtracking=None, + ctr_history_unit=None, + monotone_constraints=None, + feature_weights=None, + penalties_coefficient=None, + first_feature_use_penalties=None, + per_object_feature_penalties=None, + model_shrink_rate=None, + model_shrink_mode=None, + langevin=None, + diffusion_temperature=None, + posterior_sampling=None, + boost_from_average=None, + text_features=None, + tokenizers=None, + dictionaries=None, + feature_calcers=None, + text_processing=None, + embedding_features=None, + callback=None, + eval_fraction=None, + fixed_binary_splits=None + ): + params = {} + not_params = ["not_params", "self", "params", "__class__"] + for key, value in iteritems(locals().copy()): + if key not in not_params and value is not None: + params[key] = value + + super(CatBoostClassifier, self).__init__(params) + + def fit(self, X, y=None, cat_features=None, text_features=None, embedding_features=None, graph=None, sample_weight=None, baseline=None, use_best_model=None, + eval_set=None, verbose=None, logging_level=None, plot=False, plot_file=None, column_description=None, + verbose_eval=None, metric_period=None, silent=None, early_stopping_rounds=None, + save_snapshot=None, snapshot_file=None, snapshot_interval=None, init_model=None, callbacks=None, + log_cout=None, log_cerr=None): + """ + Fit the CatBoostClassifier model. + + Parameters + ---------- + X : catboost.Pool or list or numpy.ndarray or pandas.DataFrame or pandas.Series or polars.DataFrame + If not catboost.Pool, 2 dimensional Feature matrix or string - file with dataset. + + y : list or numpy.ndarray or pandas.DataFrame or pandas.Series or polars.DataFrame or polars.Series, optional (default=None) + Labels of the training data. + If not None, can be a single- or two- dimensional array with either: + - numerical values - for binary classification problems + - class labels (boolean, integer or string) + Use only if X is not catboost.Pool and does not point to a file. + + cat_features : list or numpy.ndarray, optional (default=None) + If not None, giving the list of Categ columns indices. + Use only if X is not catboost.Pool. + + text_features : list or numpy.ndarray, optional (default=None) + If not None, giving the list of Text columns indices. + Use only if X is not catboost.Pool. + + embedding_features : list or numpy.ndarray, optional (default=None) + If not None, giving the list of Embedding columns indices. + Use only if X is not catboost.Pool. + + graph : list or numpy.ndarray or pandas.DataFrame or polars.DataFrame + The graph edges list description. + If list or numpy.ndarrays or pandas.DataFrame, giving 2 dimensional. + + sample_weight : list or numpy.ndarray or pandas.DataFrame or pandas.Series or polars.Series, optional (default=None) + Instance weights, 1 dimensional array like. + + baseline : list or numpy.ndarray or polars.DataFrame or polars.Series, optional (default=None) + If not None, giving 2 dimensional array like data. + Use only if X is not catboost.Pool. + + use_best_model : bool, optional (default=None) + Flag to use best model + + eval_set : catboost.Pool or list of catboost.Pool or tuple (X, y) or list [(X, y)], optional (default=None) + Validation dataset or datasets for metrics calculation and possibly early stopping. + + metric_period : int + Frequency of evaluating metrics. + + verbose : bool or int + If verbose is bool, then if set to True, logging_level is set to Verbose, + if set to False, logging_level is set to Silent. + If verbose is int, it determines the frequency of writing metrics to output and + logging_level is set to Verbose. + + silent : bool + If silent is True, logging_level is set to Silent. + If silent is False, logging_level is set to Verbose. + + logging_level : string, optional (default=None) + Possible values: + - 'Silent' + - 'Verbose' + - 'Info' + - 'Debug' + + plot : bool, optional (default=False) + If True, draw train and eval error in Jupyter notebook + + plot_file : file-like or str, optional (default=None) + If not None, save train and eval error graphs to file + + verbose_eval : bool or int + Synonym for verbose. Only one of these parameters should be set. + + early_stopping_rounds : int + Activates Iter overfitting detector with od_wait set to early_stopping_rounds. + + save_snapshot : bool, [default=None] + Enable progress snapshotting for restoring progress after crashes or interruptions + + snapshot_file : string or os.PathLike, [default=None] + Learn progress snapshot file path, if None will use default filename + + snapshot_interval: int, [default=600] + Interval between saving snapshots (seconds) + + init_model : CatBoost class or string or os.PathLike, [default=None] + Continue training starting from the existing model. + If this parameter is a string or os.PathLike, load initial model from the path specified by this string. + + callbacks : list, optional (default=None) + List of callback objects that are applied at end of each iteration. + + log_cout: output stream or callback for logging (default=None) + If None is specified, sys.stdout is used + + log_cerr: error stream or callback for logging (default=None) + If None is specified, sys.stderr is used + + Returns + ------- + model : CatBoost + """ + + params = self._get_canonized_params() + if 'loss_function' in params: + CatBoostClassifier._check_is_compatible_loss(params['loss_function']) + + self._fit(X, y, cat_features, text_features, embedding_features, None, graph, sample_weight, None, None, None, None, baseline, use_best_model, + eval_set, verbose, logging_level, plot, plot_file, column_description, verbose_eval, metric_period, + silent, early_stopping_rounds, save_snapshot, snapshot_file, snapshot_interval, init_model, callbacks, log_cout, log_cerr) + return self + + def predict(self, data, prediction_type='Class', ntree_start=0, ntree_end=0, thread_count=-1, verbose=None, task_type="CPU"): + """ + Predict with data. + + Parameters + ---------- + data : catboost.Pool or list of features or list of lists or numpy.ndarray or pandas.DataFrame or pandas.Series + or polars.DataFrame or polars.Series or catboost.FeaturesData + Data to apply model on. + If data is a simple list (not list of lists) or a one-dimensional numpy.ndarray it is interpreted + as a list of features for a single object. + + prediction_type : string, optional (default='Class') + Can be: + - 'RawFormulaVal' : return raw formula value. + - 'Class' : return class label. + - 'Probability' : return probability for every class. + - 'LogProbability' : return log probability for every class. + + ntree_start: int, optional (default=0) + Model is applied on the interval [ntree_start, ntree_end) (zero-based indexing). + + ntree_end: int, optional (default=0) + Model is applied on the interval [ntree_start, ntree_end) (zero-based indexing). + If value equals to 0 this parameter is ignored and ntree_end equal to tree_count_. + + thread_count : int (default=-1) + The number of threads to use when applying the model. + Allows you to optimize the speed of execution. This parameter doesn't affect results. + If -1, then the number of threads is set to the number of CPU cores. + + verbose : bool, optional (default=False) + If True, writes the evaluation metric measured set to stderr. + + task_type : string, [default=None] + The evaluator type. + Possible values: + - 'CPU' + - 'GPU' (models with only numerical features are supported for now) + + Returns + ------- + prediction: + If data is for a single object, the return value depends on prediction_type value: + - 'RawFormulaVal' : return raw formula value. + - 'Class' : return class label. + - 'Probability' : return one-dimensional numpy.ndarray with probability for every class. + - 'LogProbability' : return one-dimensional numpy.ndarray with + log probability for every class. + otherwise numpy.ndarray, with values that depend on prediction_type value: + - 'RawFormulaVal' : one-dimensional array of raw formula value for each object. + - 'Class' : one-dimensional array of class label for each object. + - 'Probability' : two-dimensional numpy.ndarray with shape (number_of_objects x number_of_classes) + with probability for every class for each object. + - 'LogProbability' : two-dimensional numpy.ndarray with shape (number_of_objects x number_of_classes) + with log probability for every class for each object. + """ + return self._predict(data, prediction_type, ntree_start, ntree_end, thread_count, verbose, 'predict', task_type) + + def predict_proba(self, X, ntree_start=0, ntree_end=0, thread_count=-1, verbose=None, task_type="CPU"): + """ + Predict class probability with X. + + Parameters + ---------- + X : catboost.Pool or list of features or list of lists or numpy.ndarray or pandas.DataFrame or pandas.Series + or polars.DataFrame or polars.Series or catboost.FeaturesData + Data to apply model on. + If X is a simple list (not list of lists) or a one-dimensional numpy.ndarray it is interpreted + as a list of features for a single object. + + ntree_start: int, optional (default=0) + Model is applied on the interval [ntree_start, ntree_end) (zero-based indexing). + + ntree_end: int, optional (default=0) + Model is applied on the interval [ntree_start, ntree_end) (zero-based indexing). + If value equals to 0 this parameter is ignored and ntree_end equal to tree_count_. + + thread_count : int (default=-1) + The number of threads to use when applying the model. + Allows you to optimize the speed of execution. This parameter doesn't affect results. + If -1, then the number of threads is set to the number of CPU cores. + + verbose : bool + If True, writes the evaluation metric measured set to stderr. + + task_type : string, [default=None] + The evaluator type. + Possible values: + - 'CPU' + - 'GPU' (models with only numerical features are supported for now) + + Returns + ------- + prediction : + If X is for a single object + return one-dimensional numpy.ndarray with probability for every class. + otherwise + return two-dimensional numpy.ndarray with shape (number_of_objects x number_of_classes) + with probability for every class for each object. + """ + return self._predict(X, 'Probability', ntree_start, ntree_end, thread_count, verbose, 'predict_proba', task_type) + + def predict_log_proba(self, data, ntree_start=0, ntree_end=0, thread_count=-1, verbose=None, task_type="CPU"): + """ + Predict class log probability with data. + + Parameters + ---------- + data : catboost.Pool or list of features or list of lists or numpy.ndarray or pandas.DataFrame or pandas.Series + or polars.DataFrame or polars.Series or catboost.FeaturesData + Data to apply model on. + If data is a simple list (not list of lists) or a one-dimensional numpy.ndarray it is interpreted + as a list of features for a single object. + + ntree_start: int, optional (default=0) + Model is applied on the interval [ntree_start, ntree_end) (zero-based indexing). + + ntree_end: int, optional (default=0) + Model is applied on the interval [ntree_start, ntree_end) (zero-based indexing). + If value equals to 0 this parameter is ignored and ntree_end equal to tree_count_. + + thread_count : int (default=-1) + The number of threads to use when applying the model. + Allows you to optimize the speed of execution. This parameter doesn't affect results. + If -1, then the number of threads is set to the number of CPU cores. + + verbose : bool + If True, writes the evaluation metric measured set to stderr. + + task_type : string, [default=None] + The evaluator type. + Possible values: + - 'CPU' + - 'GPU' (models with only numerical features are supported for now) + + Returns + ------- + prediction : + If data is for a single object + return one-dimensional numpy.ndarray with log probability for every class. + otherwise + return two-dimensional numpy.ndarray with shape (number_of_objects x number_of_classes) + with log probability for every class for each object. + """ + return self._predict(data, 'LogProbability', ntree_start, ntree_end, thread_count, verbose, 'predict_log_proba', task_type) + + def staged_predict(self, data, prediction_type='Class', ntree_start=0, ntree_end=0, eval_period=1, thread_count=-1, verbose=None): + """ + Predict target at each stage for data. + + Parameters + ---------- + data : catboost.Pool or list of features or list of lists or numpy.ndarray or pandas.DataFrame or pandas.Series + or polars.DataFrame or polars.Series or catboost.FeaturesData + Data to apply model on. + If data is a simple list (not list of lists) or a one-dimensional numpy.ndarray it is interpreted + as a list of features for a single object. + + prediction_type : string, optional (default='Class') + Can be: + - 'RawFormulaVal' : return raw formula value. + - 'Class' : return class label. + - 'Probability' : return probability for every class. + - 'LogProbability' : return log probability for every class. + + ntree_start: int, optional (default=0) + Model is applied on the interval [ntree_start, ntree_end) with the step eval_period (zero-based indexing). + + ntree_end: int, optional (default=0) + Model is applied on the interval [ntree_start, ntree_end) with the step eval_period (zero-based indexing). + If value equals to 0 this parameter is ignored and ntree_end equal to tree_count_. + + eval_period: int, optional (default=1) + Model is applied on the interval [ntree_start, ntree_end) with the step eval_period (zero-based indexing). + + thread_count : int (default=-1) + The number of threads to use when applying the model. + Allows you to optimize the speed of execution. This parameter doesn't affect results. + If -1, then the number of threads is set to the number of CPU cores. + + verbose : bool + If True, writes the evaluation metric measured set to stderr. + + Returns + ------- + prediction : generator for each iteration that generates: + If data is for a single object, the return value depends on prediction_type value: + - 'RawFormulaVal' : return raw formula value. + - 'Class' : return majority vote class. + - 'Probability' : return one-dimensional numpy.ndarray with probability for every class. + - 'LogProbability' : return one-dimensional numpy.ndarray with + log probability for every class. + otherwise numpy.ndarray, with values that depend on prediction_type value: + - 'RawFormulaVal' : one-dimensional array of raw formula value for each object. + - 'Class' : one-dimensional array of class label for each object. + - 'Probability' : two-dimensional numpy.ndarray with shape (number_of_objects x number_of_classes) + with probability for every class for each object. + - 'LogProbability' : two-dimensional numpy.ndarray with shape (number_of_objects x number_of_classes) + with log probability for every class for each object. + """ + return self._staged_predict(data, prediction_type, ntree_start, ntree_end, eval_period, thread_count, verbose, 'staged_predict') + + def staged_predict_proba(self, data, ntree_start=0, ntree_end=0, eval_period=1, thread_count=-1, verbose=None): + """ + Predict classification target at each stage for data. + + Parameters + ---------- + data : catboost.Pool or list of features or list of lists or numpy.ndarray or pandas.DataFrame or pandas.Series + or polars.DataFrame or polars.Series or catboost.FeaturesData + Data to apply model on. + If data is a simple list (not list of lists) or a one-dimensional numpy.ndarray it is interpreted + as a list of features for a single object. + + ntree_start: int, optional (default=0) + Model is applied on the interval [ntree_start, ntree_end) with the step eval_period (zero-based indexing). + + ntree_end: int, optional (default=0) + Model is applied on the interval [ntree_start, ntree_end) with the step eval_period (zero-based indexing). + If value equals to 0 this parameter is ignored and ntree_end equal to tree_count_. + + eval_period: int, optional (default=1) + Model is applied on the interval [ntree_start, ntree_end) with the step eval_period (zero-based indexing). + + thread_count : int (default=-1) + The number of threads to use when applying the model. + Allows you to optimize the speed of execution. This parameter doesn't affect results. + If -1, then the number of threads is set to the number of CPU cores. + + verbose : bool + If True, writes the evaluation metric measured set to stderr. + + Returns + ------- + prediction : generator for each iteration that generates: + If data is for a single object + return one-dimensional numpy.ndarray with probability for every class. + otherwise + return two-dimensional numpy.ndarray with shape (number_of_objects x number_of_classes) + with probability for every class for each object. + """ + return self._staged_predict(data, 'Probability', ntree_start, ntree_end, eval_period, thread_count, verbose, 'staged_predict_proba') + + def staged_predict_log_proba(self, data, ntree_start=0, ntree_end=0, eval_period=1, thread_count=-1, verbose=None): + """ + Predict classification target at each stage for data. + + Parameters + ---------- + data : catboost.Pool or list of features or list of lists or numpy.ndarray or pandas.DataFrame or pandas.Series + or polars.DataFrame or polars.Series or catboost.FeaturesData + Data to apply model on. + If data is a simple list (not list of lists) or a one-dimensional numpy.ndarray it is interpreted + as a list of features for a single object. + + ntree_start: int, optional (default=0) + Model is applied on the interval [ntree_start, ntree_end) with the step eval_period (zero-based indexing). + + ntree_end: int, optional (default=0) + Model is applied on the interval [ntree_start, ntree_end) with the step eval_period (zero-based indexing). + If value equals to 0 this parameter is ignored and ntree_end equal to tree_count_. + + eval_period: int, optional (default=1) + Model is applied on the interval [ntree_start, ntree_end) with the step eval_period (zero-based indexing). + + thread_count : int (default=-1) + The number of threads to use when applying the model. + Allows you to optimize the speed of execution. This parameter doesn't affect results. + If -1, then the number of threads is set to the number of CPU cores. + + verbose : bool + If True, writes the evaluation metric measured set to stderr. + + Returns + ------- + prediction : generator for each iteration that generates: + If data is for a single object + return one-dimensional numpy.ndarray with log probability for every class. + otherwise + return two-dimensional numpy.ndarray with shape (number_of_objects x number_of_classes) + with log probability for every class for each object. + """ + return self._staged_predict(data, 'LogProbability', ntree_start, ntree_end, eval_period, thread_count, verbose, 'staged_predict_log_proba') + + def score(self, X, y=None): + """ + Calculate accuracy. + + Parameters + ---------- + X : catboost.Pool or list or numpy.ndarray or pandas.DataFrame or pandas.Series or polars.DataFrame + Data to apply model on. + y : list or numpy.ndarray or polars.Series + True labels. + + Returns + ------- + accuracy : float + """ + if isinstance(X, Pool): + if y is not None: + raise CatBoostError("Wrong initializing y: X is catboost.Pool object, y must be initialized inside catboost.Pool.") + y = X.get_label() + if y is None: + raise CatBoostError("Label in X has not initialized.") + if isinstance(y, pd.DataFrame): + if len(y.columns) != 1: + raise CatBoostError("y is pandas.DataFrame and has {} columns, but must have exactly one.".format(len(y.columns))) + y = y[y.columns[0]] + elif isinstance(y, pl.DataFrame): + if y.width != 1: + raise CatBoostError("y is polars.DataFrame and has {} columns, but must have exactly one.".format(y.width)) + y = y.to_series(0) + elif y is None: + raise CatBoostError("y should be specified.") + y = np.array(y) + predicted_classes = self._predict( + X, + prediction_type='Class', + ntree_start=0, + ntree_end=0, + thread_count=-1, + verbose=None, + parent_method_name='score' + ).reshape(-1) + if np.issubdtype(predicted_classes.dtype, np.number): + if np.issubdtype(y.dtype, np.character): + raise CatBoostError('predicted classes have numeric type but specified y contains strings') + elif predicted_classes.dtype == np.bool_: + if np.issubdtype(y.dtype, np.character): + raise CatBoostError('predicted classes have boolean type but specified y contains strings') + else: + if np.issubdtype(y.dtype, np.number): + raise CatBoostError('predicted classes have string type but specified y is numeric') + elif np.issubdtype(y.dtype, np.bool_): + raise CatBoostError('predicted classes have string type but specified y is boolean') + return np.mean(np.array(predicted_classes) == np.array(y)) + + def set_probability_threshold(self, binclass_probability_threshold=None): + """ + Set a threshold for class separation in binary classification task for a trained model. + :param binclass_probability_threshold: float number in [0, 1] or None to discard it + """ + if not self.is_fitted(): + raise CatBoostError("You can't set probability threshold for not fitted model.") + metadata = self.get_metadata() + if binclass_probability_threshold is None: + if 'binclass_probability_threshold' in metadata.keys(): + del metadata['binclass_probability_threshold'] + else: + if not isinstance(binclass_probability_threshold, FLOAT_TYPES): + raise CatBoostError("binclass_probability_threshold must have float type") + assert 0. <= binclass_probability_threshold <= 1., \ + "Please provide correct probability for binclass_probability_threshold argument in [0, 1] range" + self.get_metadata()['binclass_probability_threshold'] = str(binclass_probability_threshold) + + def get_probability_threshold(self): + """ + Get a threshold for class separation in binary classification task + """ + if not self.is_fitted(): + raise CatBoostError("Not fitted models don't have a probability threshold.") + return self._object._get_binclass_probability_threshold() + + @staticmethod + def _check_is_compatible_loss(loss_function): + if isinstance(loss_function, str) and not CatBoost._is_classification_objective(loss_function): + raise CatBoostError("Invalid loss_function='{}': for classifier use " + "Logloss, CrossEntropy, MultiClass, MultiClassOneVsAll or a custom objective object".format(loss_function)) + + +class CatBoostRegressor(CatBoost): + """ + Implementation of the scikit-learn API for CatBoost regression. + + Parameters + ---------- + Like in CatBoostClassifier, except loss_function, classes_count, class_names and class_weights + + loss_function : string, [default='RMSE'] + 'RMSE' + 'MAE' + 'Quantile:alpha=value' + 'LogLinQuantile:alpha=value' + 'Poisson' + 'MAPE' + 'Lq:q=value' + 'SurvivalAft:dist=value;scale=value' + """ + + _estimator_type = 'regressor' + + def __init__( + self, + iterations=None, + learning_rate=None, + depth=None, + l2_leaf_reg=None, + model_size_reg=None, + rsm=None, + loss_function='RMSE', + border_count=None, + feature_border_type=None, + per_float_feature_quantization=None, + input_borders=None, + output_borders=None, + fold_permutation_block=None, + od_pval=None, + od_wait=None, + od_type=None, + nan_mode=None, + counter_calc_method=None, + leaf_estimation_iterations=None, + leaf_estimation_method=None, + thread_count=None, + random_seed=None, + use_best_model=None, + best_model_min_trees=None, + verbose=None, + silent=None, + logging_level=None, + metric_period=None, + ctr_leaf_count_limit=None, + store_all_simple_ctr=None, + max_ctr_complexity=None, + has_time=None, + allow_const_label=None, + target_border=None, + one_hot_max_size=None, + random_strength=None, + random_score_type=None, + name=None, + ignored_features=None, + train_dir=None, + custom_metric=None, + eval_metric=None, + bagging_temperature=None, + save_snapshot=None, + snapshot_file=None, + snapshot_interval=None, + fold_len_multiplier=None, + used_ram_limit=None, + gpu_ram_part=None, + pinned_memory_size=None, + allow_writing_files=None, + final_ctr_computation_mode=None, + approx_on_full_history=None, + boosting_type=None, + simple_ctr=None, + combinations_ctr=None, + per_feature_ctr=None, + ctr_description=None, + ctr_target_border_count=None, + task_type=None, + device_config=None, + devices=None, + bootstrap_type=None, + subsample=None, + mvs_reg=None, + sampling_frequency=None, + sampling_unit=None, + dev_score_calc_obj_block_size=None, + dev_efb_max_buckets=None, + sparse_features_conflict_fraction=None, + max_depth=None, + n_estimators=None, + num_boost_round=None, + num_trees=None, + colsample_bylevel=None, + random_state=None, + reg_lambda=None, + objective=None, + eta=None, + max_bin=None, + gpu_cat_features_storage=None, + data_partition=None, + metadata=None, + early_stopping_rounds=None, + cat_features=None, + grow_policy=None, + min_data_in_leaf=None, + min_child_samples=None, + max_leaves=None, + num_leaves=None, + score_function=None, + leaf_estimation_backtracking=None, + ctr_history_unit=None, + monotone_constraints=None, + feature_weights=None, + penalties_coefficient=None, + first_feature_use_penalties=None, + per_object_feature_penalties=None, + model_shrink_rate=None, + model_shrink_mode=None, + langevin=None, + diffusion_temperature=None, + posterior_sampling=None, + boost_from_average=None, + text_features=None, + tokenizers=None, + dictionaries=None, + feature_calcers=None, + text_processing=None, + embedding_features=None, + eval_fraction=None, + fixed_binary_splits=None + ): + params = {} + not_params = ["not_params", "self", "params", "__class__"] + for key, value in iteritems(locals().copy()): + if key not in not_params and value is not None: + params[key] = value + + super(CatBoostRegressor, self).__init__(params) + + def fit(self, X, y=None, cat_features=None, text_features=None, embedding_features=None, graph=None, + sample_weight=None, baseline=None, use_best_model=None, + eval_set=None, verbose=None, logging_level=None, plot=False, plot_file=None, column_description=None, + verbose_eval=None, metric_period=None, silent=None, early_stopping_rounds=None, + save_snapshot=None, snapshot_file=None, snapshot_interval=None, init_model=None, callbacks=None, + log_cout=None, log_cerr=None): + """ + Fit the CatBoost model. + + Parameters + ---------- + X : catboost.Pool or list or numpy.ndarray or pandas.DataFrame or pandas.Series or polars.DataFrame + If not catboost.Pool, 2 dimensional Feature matrix or string - file with dataset. + + y : list or numpy.ndarray or pandas.DataFrame or pandas.Series or polars.DataFrame or polars.Series, optional (default=None) + Labels of the training data. + If not None, can be a single- or two- dimensional array with numerical values. + Use only if X is not catboost.Pool and does not point to a file. + + cat_features : list or numpy.ndarray, optional (default=None) + If not None, giving the list of Categ columns indices. + Use only if X is not catboost.Pool. + + text_features : list or numpy.ndarray, optional (default=None) + If not None, giving the list of Text columns indices. + Use only if X is not catboost.Pool. + + embedding_features : list or numpy.ndarray, optional (default=None) + If not None, giving the list of Embedding columns indices. + Use only if X is not catboost.Pool. + + graph : list or numpy.ndarray or pandas.DataFrame or polars.DataFrame + The graph edges list description. + If list or numpy.ndarrays or pandas.DataFrame or polars.DataFrame, giving 2 dimensional. + + sample_weight : list or numpy.ndarray or pandas.DataFrame or pandas.Series or polars.Series, optional (default=None) + Instance weights, 1 dimensional array like. + + baseline : list or numpy.ndarray or polars.DataFrame or polars.Series, optional (default=None) + If not None, giving 2 dimensional array like data. + Use only if X is not catboost.Pool. + + use_best_model : bool, optional (default=None) + Flag to use best model + + eval_set : catboost.Pool or list of catboost.Pool or tuple (X, y) or list [(X, y)], optional (default=None) + Validation dataset or datasets for metrics calculation and possibly early stopping. + + metric_period : int + Frequency of evaluating metrics. + + verbose : bool or int + If verbose is bool, then if set to True, logging_level is set to Verbose, + if set to False, logging_level is set to Silent. + If verbose is int, it determines the frequency of writing metrics to output and + logging_level is set to Verbose. + + silent : bool + If silent is True, logging_level is set to Silent. + If silent is False, logging_level is set to Verbose. + + logging_level : string, optional (default=None) + Possible values: + - 'Silent' + - 'Verbose' + - 'Info' + - 'Debug' + + plot : bool, optional (default=False) + If True, draw train and eval error in Jupyter notebook + + plot_file : file-like or str, optional (default=None) + If not None, save train and eval error graphs to file (requires installed plotly) + + verbose_eval : bool or int + Synonym for verbose. Only one of these parameters should be set. + + early_stopping_rounds : int + Activates Iter overfitting detector with od_wait set to early_stopping_rounds. + + save_snapshot : bool, [default=None] + Enable progress snapshotting for restoring progress after crashes or interruptions + + snapshot_file : string or os.PathLike, [default=None] + Learn progress snapshot file path, if None will use default filename + + snapshot_interval: int, [default=600] + Interval between saving snapshots (seconds) + + init_model : CatBoost class or string or os.PathLike, [default=None] + Continue training starting from the existing model. + If this parameter is a string or os.PathLike, load initial model from the path specified by this string. + + callbacks : list, optional (default=None) + List of callback objects that are applied at end of each iteration. + + log_cout: output stream or callback for logging (default=None) + If None is specified, sys.stdout is used + + log_cerr: error stream or callback for logging (default=None) + If None is specified, sys.stderr is used + + Returns + ------- + model : CatBoost + """ + + params = self._get_canonized_params() + if 'loss_function' in params: + CatBoostRegressor._check_is_compatible_loss(params['loss_function']) + return self._fit(X, y, cat_features, text_features, embedding_features, None, graph, sample_weight, None, None, None, None, baseline, + use_best_model, eval_set, verbose, logging_level, plot, plot_file, column_description, + verbose_eval, metric_period, silent, early_stopping_rounds, + save_snapshot, snapshot_file, snapshot_interval, init_model, callbacks, log_cout, log_cerr) + + def predict(self, data, prediction_type=None, ntree_start=0, ntree_end=0, thread_count=-1, verbose=None, task_type="CPU"): + """ + Predict with data. + + Parameters + ---------- + data : catboost.Pool or list of features or list of lists or numpy.ndarray or pandas.DataFrame or pandas.Series + or polars.DataFrame or polars.Series or catboost.FeaturesData + Data to apply model on. + If data is a simple list (not list of lists) or a one-dimensional numpy.ndarray it is interpreted + as a list of features for a single object. + + prediction_type : string, optional (default='RawFormulaVal') + Can be: + - 'RawFormulaVal' : return raw formula value. + - 'Exponent' : return Exponent of raw formula value. + + ntree_start: int, optional (default=0) + Model is applied on the interval [ntree_start, ntree_end) (zero-based indexing). + + ntree_end: int, optional (default=0) + Model is applied on the interval [ntree_start, ntree_end) (zero-based indexing). + If value equals to 0 this parameter is ignored and ntree_end equal to tree_count_. + + thread_count : int (default=-1) + The number of threads to use when applying the model. + Allows you to optimize the speed of execution. This parameter doesn't affect results. + If -1, then the number of threads is set to the number of CPU cores. + + verbose : bool + If True, writes the evaluation metric measured set to stderr. + + task_type : string, [default=None] + The evaluator type. + Possible values: + - 'CPU' + - 'GPU' (models with only numerical features are supported for now) + + Returns + ------- + prediction : + If data is for a single object, the return value is single float formula return value + otherwise one-dimensional numpy.ndarray of formula return values for each object. + """ + if prediction_type is None: + prediction_type = self._get_default_prediction_type() + return self._predict(data, prediction_type, ntree_start, ntree_end, thread_count, verbose, 'predict', task_type) + + def staged_predict(self, data, prediction_type='RawFormulaVal', ntree_start=0, ntree_end=0, eval_period=1, thread_count=-1, verbose=None): + """ + Predict target at each stage for data. + + Parameters + ---------- + data : catboost.Pool or list of features or list of lists or numpy.ndarray or pandas.DataFrame or pandas.Series + or polars.DataFrame or polars.Series or catboost.FeaturesData + Data to apply model on. + If data is a simple list (not list of lists) or a one-dimensional numpy.ndarray it is interpreted + as a list of features for a single object. + + ntree_start: int, optional (default=0) + Model is applied on the interval [ntree_start, ntree_end) with the step eval_period (zero-based indexing). + + ntree_end: int, optional (default=0) + Model is applied on the interval [ntree_start, ntree_end) with the step eval_period (zero-based indexing). + If value equals to 0 this parameter is ignored and ntree_end equal to tree_count_. + + eval_period: int, optional (default=1) + Model is applied on the interval [ntree_start, ntree_end) with the step eval_period (zero-based indexing). + + thread_count : int (default=-1) + The number of threads to use when applying the model. + Allows you to optimize the speed of execution. This parameter doesn't affect results. + If -1, then the number of threads is set to the number of CPU cores. + + verbose : bool + If True, writes the evaluation metric measured set to stderr. + + Returns + ------- + prediction : generator for each iteration that generates: + If data is for a single object, the return value is single float formula return value + otherwise one-dimensional numpy.ndarray of formula return values for each object. + """ + return self._staged_predict(data, prediction_type, ntree_start, ntree_end, eval_period, thread_count, verbose, 'staged_predict') + + def score(self, X, y=None): + """ + Calculate R^2. + + Parameters + ---------- + X : catboost.Pool or list or numpy.ndarray or pandas.DataFrame or pandas.Series or polars.DataFrame + Data to apply model on. + y : list or numpy.ndarray or polars.Series + True labels. + + Returns + ------- + R^2 : float + """ + if isinstance(X, Pool): + if y is not None: + raise CatBoostError("Wrong initializing y: X is catboost.Pool object, y must be initialized inside catboost.Pool.") + y = X.get_label() + if y is None: + raise CatBoostError("Label in X has not initialized.") + elif y is None: + raise CatBoostError("y should be specified.") + y = np.array(y, dtype=np.float64) + predictions = self._predict( + X, + prediction_type=self._get_default_prediction_type(), + ntree_start=0, + ntree_end=0, + thread_count=-1, + verbose=None, + parent_method_name='score' + ) + loss = self._object._get_loss_function_name() + if loss == 'RMSEWithUncertainty': + predictions = predictions[:, 0] + if predictions.size != y.size: + msg = "labels and predictions should have same size. But y.size={} != preds.size={}" + raise CatBoostError(msg.format(y.size, predictions.size)) + y = y.reshape(predictions.shape) + total_sum_of_squares = np.sum((y - y.mean(axis=0)) ** 2) + residual_sum_of_squares = np.sum((y - predictions) ** 2) + return 1 - residual_sum_of_squares / total_sum_of_squares + + @staticmethod + def _check_is_compatible_loss(loss_function): + is_regression = CatBoost._is_regression_objective(loss_function) or CatBoost._is_multiregression_objective(loss_function) or CatBoost._is_survivalregression_objective(loss_function) + if isinstance(loss_function, str) and not is_regression: + raise CatBoostError("Invalid loss_function='{}': for regressor use " + "RMSE, MultiRMSE, SurvivalAft, MAE, Quantile, LogLinQuantile, Poisson, MAPE, Lq, RMSPE or a custom objective object".format(loss_function)) + + def _get_default_prediction_type(self): + # TODO(ilyzhin) change on get_all_params after MLTOOLS-4758 + params = self._get_canonized_params() + loss_function = params.get('loss_function') + if loss_function and isinstance(loss_function, str): + if loss_function.startswith('Poisson') or loss_function.startswith('Tweedie'): + return 'Exponent' + if loss_function == 'RMSEWithUncertainty': + return 'RMSEWithUncertainty' + return 'RawFormulaVal' + + +class CatBoostRanker(CatBoost): + """ + Implementation of the scikit-learn API for CatBoost ranking. + Parameters + ---------- + Like in CatBoostClassifier, except loss_function, classes_count, class_names and class_weights + loss_function : string, [default='YetiRank'] + 'YetiRank' + 'YetiRankPairwise' + 'StochasticFilter' + 'StochasticRank' + 'QueryCrossEntropy' + 'QueryRMSE' + 'GroupQuantile' + 'QuerySoftMax' + 'PairLogit' + 'PairLogitPairwise' + """ + + _estimator_type = 'ranker' + + def __init__( + self, + iterations=None, + learning_rate=None, + depth=None, + l2_leaf_reg=None, + model_size_reg=None, + rsm=None, + loss_function='YetiRank', + border_count=None, + feature_border_type=None, + per_float_feature_quantization=None, + input_borders=None, + output_borders=None, + fold_permutation_block=None, + od_pval=None, + od_wait=None, + od_type=None, + nan_mode=None, + counter_calc_method=None, + leaf_estimation_iterations=None, + leaf_estimation_method=None, + thread_count=None, + random_seed=None, + use_best_model=None, + best_model_min_trees=None, + verbose=None, + silent=None, + logging_level=None, + metric_period=None, + ctr_leaf_count_limit=None, + store_all_simple_ctr=None, + max_ctr_complexity=None, + has_time=None, + allow_const_label=None, + target_border=None, + one_hot_max_size=None, + random_strength=None, + random_score_type=None, + name=None, + ignored_features=None, + train_dir=None, + custom_metric=None, + eval_metric=None, + bagging_temperature=None, + save_snapshot=None, + snapshot_file=None, + snapshot_interval=None, + fold_len_multiplier=None, + used_ram_limit=None, + gpu_ram_part=None, + pinned_memory_size=None, + allow_writing_files=None, + final_ctr_computation_mode=None, + approx_on_full_history=None, + boosting_type=None, + simple_ctr=None, + combinations_ctr=None, + per_feature_ctr=None, + ctr_description=None, + ctr_target_border_count=None, + task_type=None, + device_config=None, + devices=None, + bootstrap_type=None, + subsample=None, + mvs_reg=None, + sampling_frequency=None, + sampling_unit=None, + dev_score_calc_obj_block_size=None, + dev_efb_max_buckets=None, + sparse_features_conflict_fraction=None, + max_depth=None, + n_estimators=None, + num_boost_round=None, + num_trees=None, + colsample_bylevel=None, + random_state=None, + reg_lambda=None, + objective=None, + eta=None, + max_bin=None, + gpu_cat_features_storage=None, + data_partition=None, + metadata=None, + early_stopping_rounds=None, + cat_features=None, + grow_policy=None, + min_data_in_leaf=None, + min_child_samples=None, + max_leaves=None, + num_leaves=None, + score_function=None, + leaf_estimation_backtracking=None, + ctr_history_unit=None, + monotone_constraints=None, + feature_weights=None, + penalties_coefficient=None, + first_feature_use_penalties=None, + per_object_feature_penalties=None, + model_shrink_rate=None, + model_shrink_mode=None, + langevin=None, + diffusion_temperature=None, + posterior_sampling=None, + boost_from_average=None, + text_features=None, + tokenizers=None, + dictionaries=None, + feature_calcers=None, + text_processing=None, + embedding_features=None, + eval_fraction=None, + fixed_binary_splits=None + ): + params = {} + not_params = ["not_params", "self", "params", "__class__"] + for key, value in iteritems(locals().copy()): + if key not in not_params and value is not None: + params[key] = value + + super(CatBoostRanker, self).__init__(params) + + def fit(self, X, y=None, group_id=None, cat_features=None, text_features=None, + embedding_features=None, pairs=None, graph=None, sample_weight=None, group_weight=None, + subgroup_id=None, pairs_weight=None, baseline=None, use_best_model=None, + eval_set=None, verbose=None, logging_level=None, plot=False, plot_file=None, column_description=None, + verbose_eval=None, metric_period=None, silent=None, early_stopping_rounds=None, + save_snapshot=None, snapshot_file=None, snapshot_interval=None, init_model=None, callbacks=None, + log_cout=None, log_cerr=None): + """ + Fit the CatBoostRanker model. + Parameters + ---------- + X : catboost.Pool or list or numpy.ndarray or pandas.DataFrame or pandas.Series or polars.DataFrame + If not catboost.Pool, 2 dimensional Feature matrix or string - file with dataset. + y : list or numpy.ndarray or pandas.DataFrame or pandas.Series or polars.DataFrame or polars.Series, optional (default=None) + Labels of the training data. + If not None, can be a single-dimensional array with numerical values. + Use only if X is not catboost.Pool and does not point to a file. + group_id : numpy.ndarray or pandas.DataFrame or pandas.Series or polars.Series, optional (default=None) + Ranking groups, 1 dimensional array like. + Use only if X is not catboost.Pool. + cat_features : list or numpy.ndarray, optional (default=None) + If not None, giving the list of Categ columns indices. + Use only if X is not catboost.Pool. + text_features : list or numpy.ndarray, optional (default=None) + If not None, giving the list of Text columns indices. + Use only if X is not catboost.Pool. + embedding_features : list or numpy.ndarray, optional (default=None) + If not None, giving the list of Embedding columns indices. + Use only if X is not catboost.Pool. + pairs : list or numpy.ndarray or pandas.DataFrame or polars.DataFrame, optional (default=None) + The pairs description in the form of a two-dimensional matrix of shape N by 2: + N is the number of pairs. + The first element of the pair is the zero-based index of the winner object from the input dataset for pairwise comparison. + The second element of the pair is the zero-based index of the loser object from the input dataset for pairwise comparison. + graph : list or numpy.ndarray or pandas.DataFrame or polars.DataFrame + The graph edges list description. + If list or numpy.ndarrays or pandas.DataFrame, giving 2 dimensional. + sample_weight : list or numpy.ndarray or pandas.DataFrame or pandas.Series or polars.Series, optional (default=None) + Instance weights, 1 dimensional array like. + group_weight : list or numpy.ndarray or polars.Series (default=None) + The weights of all objects within the defined groups from the input data in the form of one-dimensional array-like data. + Used for calculating the final values of trees. By default, it is set to 1 for all objects in all groups. + Only a weight or group_weight parameter can be used at a time + subgroup_id : list or numpy.ndarray or polars.Series (default=None) + Subgroup identifiers for all input objects. Supported identifier types are: + int + string types (string or unicode for Python 2 and bytes or string for Python 3). + pairs_weight : list or numpy.ndarray or polars.Series (default=None) + The weight of each input pair of objects in the form of one-dimensional array-like pairs. + The number of given values must match the number of specified pairs. + This information is used for calculation and optimization of Pairwise metrics . + By default, it is set to 1 for all pairs. + baseline : list or numpy.ndarray or polars.DataFrame or polars.Series, optional (default=None) + If not None, giving 2 dimensional array like data. + Use only if X is not catboost.Pool. + use_best_model : bool, optional (default=None) + Flag to use best model + eval_set : catboost.Pool or list of catboost.Pool or tuple (X, y) or list [(X, y)], optional (default=None) + Validation dataset or datasets for metrics calculation and possibly early stopping. + verbose : bool or int + If verbose is bool, then if set to True, logging_level is set to Verbose, + if set to False, logging_level is set to Silent. + If verbose is int, it determines the frequency of writing metrics to output and + logging_level is set to Verbose. + logging_level : string, optional (default=None) + Possible values: + - 'Silent' + - 'Verbose' + - 'Info' + - 'Debug' + plot : bool, optional (default=False) + If True, draw train and eval error in Jupyter notebook + plot_file : file-like or str, optional (default=None) + If not None, save train and eval error graphs to file + verbose_eval : bool or int + Synonym for verbose. Only one of these parameters should be set. + metric_period : int + Frequency of evaluating metrics. + silent : bool + If silent is True, logging_level is set to Silent. + If silent is False, logging_level is set to Verbose. + early_stopping_rounds : int + Activates Iter overfitting detector with od_wait set to early_stopping_rounds. + save_snapshot : bool, [default=None] + Enable progress snapshotting for restoring progress after crashes or interruptions + snapshot_file : string or os.PathLike, [default=None] + Learn progress snapshot file path, if None will use default filename + snapshot_interval: int, [default=600] + Interval between saving snapshots (seconds) + init_model : CatBoost class or string or os.PathLike, [default=None] + Continue training starting from the existing model. + If this parameter is a string or os.PathLike, load initial model from the path specified by this string. + callbacks : list, optional (default=None) + List of callback objects that are applied at end of each iteration. + + log_cout: output stream or callback for logging (default=None) + If None is specified, sys.stdout is used + + log_cerr: error stream or callback for logging (default=None) + If None is specified, sys.stderr is used + + Returns + ------- + model : CatBoost + """ + + params = self._get_canonized_params() + if 'loss_function' in params: + CatBoostRanker._check_is_compatible_loss(params['loss_function']) + + self._fit(X, y, cat_features, text_features, embedding_features, pairs, graph, + sample_weight, group_id, group_weight, subgroup_id, pairs_weight, baseline, use_best_model, + eval_set, verbose, logging_level, plot, plot_file, column_description, verbose_eval, metric_period, + silent, early_stopping_rounds, save_snapshot, snapshot_file, snapshot_interval, init_model, callbacks, log_cout, log_cerr) + return self + + def predict(self, X, ntree_start=0, ntree_end=0, thread_count=-1, verbose=None): + """ + Predict with data. + Parameters + ---------- + X : catboost.Pool or list of features or list of lists or numpy.ndarray or pandas.DataFrame or pandas.Series + or polars.DataFrame or polars.Series or catboost.FeaturesData + Data to apply model on. + If data is a simple list (not list of lists) or a one-dimensional numpy.ndarray it is interpreted + as a list of features for a single object. + ntree_start: int, optional (default=0) + Model is applied on the interval [ntree_start, ntree_end) (zero-based indexing). + ntree_end: int, optional (default=0) + Model is applied on the interval [ntree_start, ntree_end) (zero-based indexing). + If value equals to 0 this parameter is ignored and ntree_end equal to tree_count_. + thread_count : int (default=-1) + The number of threads to use when applying the model. + Allows you to optimize the speed of execution. This parameter doesn't affect results. + If -1, then the number of threads is set to the number of CPU cores. + verbose : bool + If True, writes the evaluation metric measured set to stderr. + Returns + ------- + prediction : + If data is for a single object, the return value is single float formula return value + otherwise one-dimensional numpy.ndarray of formula return values for each object. + """ + return self._predict(X, 'RawFormulaVal', ntree_start, ntree_end, thread_count, verbose, 'predict') + + def staged_predict(self, X, ntree_start=0, ntree_end=0, eval_period=1, thread_count=-1, verbose=None): + """ + Predict target at each stage for data. + Parameters + ---------- + X : catboost.Pool or list of features or list of lists or numpy.ndarray or pandas.DataFrame or pandas.Series + or polars.DataFrame or polars.Series or catboost.FeaturesData + Data to apply model on. + If data is a simple list (not list of lists) or a one-dimensional numpy.ndarray it is interpreted + as a list of features for a single object. + ntree_start: int, optional (default=0) + Model is applied on the interval [ntree_start, ntree_end) with the step eval_period (zero-based indexing). + ntree_end: int, optional (default=0) + Model is applied on the interval [ntree_start, ntree_end) with the step eval_period (zero-based indexing). + If value equals to 0 this parameter is ignored and ntree_end equal to tree_count_. + eval_period: int, optional (default=1) + Model is applied on the interval [ntree_start, ntree_end) with the step eval_period (zero-based indexing). + thread_count : int (default=-1) + The number of threads to use when applying the model. + Allows you to optimize the speed of execution. This parameter doesn't affect results. + If -1, then the number of threads is set to the number of CPU cores. + verbose : bool + If True, writes the evaluation metric measured set to stderr. + Returns + ------- + prediction : generator for each iteration that generates: + If data is for a single object, the return value is single float formula return value + otherwise one-dimensional numpy.ndarray of formula return values for each object. + """ + return self._staged_predict(X, 'RawFormulaVal', ntree_start, ntree_end, eval_period, thread_count, verbose, 'staged_predict') + + def score(self, X, y=None, group_id=None, top=None, type=None, denominator=None, group_weight=None, thread_count=-1): + """ + Calculate NDCG@top + Parameters + ---------- + X : catboost.Pool or list or numpy.ndarray or pandas.DataFrame or pandas.Series or polars.DataFrame + Data to apply model on. + y : list or numpy.ndarrays or pandas.DataFrame or pandas.Series or polars.Series + True labels. + group_id : list or numpy.ndarray or pandas.DataFrame or pandas.Series + Ranking groups. If X is a Pool, group_id must be defined into X + top : unsigned integer, up to `pow(2, 32) / 2 - 1` + NDCG, Number of top-ranked objects to calculate NDCG + type : str + NDCG, Metric_type: 'Base' or 'Exp' + denominator : str + NDCG, Denominator type: 'LogPosition' or 'Position' + group_weight : list or numpy.ndarray or pandas.DataFrame or pandas.Series + Group weights. + thread_count : int, optional (default=-1) + Number of threads to work with. + Returns + ------- + NDCG@top : float + higher is better + """ + def get_ndcg_metric_name(values, names): + if np.all(np.equal(values, None)): + return 'NDCG' + return 'NDCG:' + ';'.join(['{}={}'.format(n, v) for v, n in zip(values, names) if v is not None]) + + if isinstance(X, Pool): + if y is not None: + raise CatBoostError("Wrong initializing y: X is catboost.Pool object, y must be initialized inside catboost.Pool.") + y = X.get_label() + if group_id is not None: + raise CatBoostError("Wrong initializing group_id: X is catboost.Pool object, group_id must be initialized inside catboost.Pool.") + group_id = X.get_group_id_hash() + + if y is None: + raise CatBoostError("y must be initialized.") + if group_id is None: + raise CatBoostError("group_id must be initialized. If groups are not expected, pass an array of zeros") + + predictions = self.predict(X) + return _eval_metric_util([y], [predictions], get_ndcg_metric_name([top, type, denominator], ['top', 'type', 'denominator']), None, group_id, group_weight, None, None, thread_count)[0] + + @staticmethod + def _check_is_compatible_loss(loss_function): + is_ranking = CatBoost._is_ranking_objective(loss_function) + is_regression = CatBoost._is_regression_objective(loss_function) + + if is_regression: + warnings.warn("Regression loss ('{}') ignores an important ranking parameter 'group_id'".format(loss_function), RuntimeWarning) + if not (is_ranking or is_regression): + raise CatBoostError("Invalid loss_function='{}': for ranker use " + "YetiRank, YetiRankPairwise, StochasticFilter, StochasticRank, " + "QueryCrossEntropy, QueryRMSE, GroupQuantile, QuerySoftMax, PairLogit, PairLogitPairwise. " + "It's also possible to use a regression loss".format(loss_function)) + + +def train(pool=None, params=None, dtrain=None, logging_level=None, verbose=None, iterations=None, + num_boost_round=None, evals=None, eval_set=None, plot=None, plot_file=None, verbose_eval=None, metric_period=None, + early_stopping_rounds=None, save_snapshot=None, snapshot_file=None, snapshot_interval=None, + init_model=None, log_cout=None, log_cerr=None): + """ + Train CatBoost model. + + Parameters + ---------- + params : dict + Parameters for CatBoost. + If None, all params are set to their defaults. + If dict, overriding parameters present in the dict. + + pool : catboost.Pool or tuple (X, y) + Data to train on. + + iterations : int + Number of boosting iterations. Can be set in params dict. + + evals : catboost.Pool or list of catboost.Pool or tuple (X, y) or list [(X, y)] + Synonym for eval_set. Only one of these parameters should be set. + + dtrain : catboost.Pool or tuple (X, y) + Synonym for pool parameter. Only one of these parameters should be set. + + logging_level : string, optional (default=None) + Possible values: + - 'Silent' + - 'Verbose' + - 'Info' + - 'Debug' + + metric_period : int + Frequency of evaluating metrics. + + verbose : bool or int + If verbose is bool, then if set to True, logging_level is set to Verbose, + if set to False, logging_level is set to Silent. + If verbose is int, it determines the frequency of writing metrics to output and + logging_level is set to Verbose. + + verbose_eval : bool or int + Synonym for verbose. Only one of these parameters should be set. + + iterations : int + Number of boosting iterations. Can be set in params dict. + + num_boost_round : int + Synonym for iterations. Only one of these parameters should be set. + + eval_set : catboost.Pool or list of catboost.Pool or tuple (X, y) or list [(X, y)] + Validation dataset or datasets for metrics calculation and possibly early stopping. + + plot : bool, optional (default=False) + If True, draw train and eval error in Jupyter notebook +` + plot_file : file-like or str, optional (default=None) + If not None, save train and eval error graphs to file + + early_stopping_rounds : int + Activates Iter overfitting detector with od_wait set to early_stopping_rounds. + + save_snapshot : bool, [default=None] + Enable progress snapshotting for restoring progress after crashes or interruptions + + snapshot_file : string or os.PathLike, [default=None] + Learn progress snapshot file path, if None will use default filename + + snapshot_interval: int, [default=600] + Interval between saving snapshots (seconds) + + init_model : CatBoost class or string or os.PathLike, [default=None] + Continue training starting from the existing model. + If this parameter is a string or os.PathLike, load initial model from the path specified by this string. + + log_cout: output stream or callback for logging (default=None) + If None is specified, sys.stdout is used + + log_cerr: error stream or callback for logging (default=None) + If None is specified, sys.stderr is used + + Returns + ------- + model : CatBoost class + """ + + if params is None: + raise CatBoostError("params should be set.") + + if dtrain is not None: + if pool is None: + pool = dtrain + else: + raise CatBoostError("Only one of the parameters pool and dtrain should be set.") + + if num_boost_round is not None: + if iterations is None: + iterations = num_boost_round + else: + raise CatBoostError("Only one of the parameters iterations and num_boost_round should be set.") + if iterations is not None: + params = deepcopy(params) + params.update({ + 'iterations': iterations + }) + + if early_stopping_rounds is not None: + params.update({ + 'od_type': 'Iter' + }) + if 'od_pval' in params: + del params['od_pval'] + params.update({ + 'od_wait': early_stopping_rounds + }) + + if evals is not None: + if eval_set is not None: + raise CatBoostError('Only one of the parameters evals, eval_set should be set.') + eval_set = evals + + model = CatBoost(params) + model.fit(X=pool, eval_set=eval_set, logging_level=logging_level, plot=plot, plot_file=plot_file, verbose=verbose, + verbose_eval=verbose_eval, metric_period=metric_period, + early_stopping_rounds=early_stopping_rounds, save_snapshot=save_snapshot, + snapshot_file=snapshot_file, snapshot_interval=snapshot_interval, init_model=init_model, + log_cout=log_cout, log_cerr=log_cerr) + return model + + +def _convert_to_catboost(models): + """ + Convert _Catboost instances to Catboost ones + """ + output_models = [] + for model in models: + cb_model = CatBoost() + cb_model._object = model + cb_model._set_trained_model_attributes() + output_models.append(cb_model) + return output_models + + +def sample_gaussian_process(X, y, eval_set=None, + cat_features=None, text_features=None, embedding_features=None, + random_seed=None, samples=10, posterior_iterations=900, prior_iterations=100, learning_rate=0.1, + depth=6, sigma=0.1, delta=0, random_strength=0.1, random_score_type='Gumbel', eps=1e-4, verbose=False): + """ + Implementation of Gaussian process sampling (Kernel Gradient Boosting/Algorithm 4) from "Gradient Boosting Performs Gaussian Process Inference" https://arxiv.org/abs/2206.05608 + Produces samples from posterior GP with prior assumption f ~ GP(0, sigma ** 2 K + delta ** 2 I) + + Parameters + ---------- + X : list or numpy.ndarray or pandas.DataFrame or pandas.Series or polars.DataFrame or catboost.FeaturesData + If catboost.FeaturesData it must be 2 dimensional Feature matrix + Must be non-empty (contain > 0 objects) + y : list or numpy.ndarray or pandas.DataFrame or pandas.Series + Labels of the training data. + Must be a single-dimensional array with numerical values. + eval_set : catboost.Pool or list of catboost.Pool or tuple (X, y) or list [(X, y)], optional (default=None) + Validation dataset or datasets for metrics calculation and possibly early stopping in posterior training. + cat_features : list or numpy.ndarray, optional (default=None) + If not None, giving the list of Categ columns indices. + Use only if X is not catboost.FeaturesData + text_features : list or numpy.ndarray, optional (default=None) + If not none, giving the list of Text columns indices. + Use only if X is not catboost.FeaturesData + embedding_features : list or numpy.ndarray, optional (default=None) + If not none, giving the list of Embedding columns indices. + Use only if X is not catboost.FeaturesData + random_seed : int, [default=None] + Random number seed. + If None, 0 is used. + range: [0,+inf) + samples : int, [default=10] + Number of Monte-Carlo samples from GP posterior. Controls how many models this function will return. + range: [1,+inf) + posterior_iterations : int, [default=900] + Max count of trees for posterior sampling step. + range: [1,+inf) + prior_iterations : int, [default=100] + Max count of trees for prior sampling step. + range: [1,+inf] + learning_rate : float, [default=0.1] + Step size shrinkage used in update to prevent overfitting. + range: (0,1] + depth : int, [default=6] + Depth of the trees in the models. + range: [1,16] + sigma : float, [default=0.1] + Scale of GP kernel (lower values lead to lower posterior variance) + range: (0,+inf) + delta : float, [default=0] + Scale of homogenious noise of GP kernel (adjust if target is noisy) + range: [0,+inf) + random_strength : float, [default=0.1] + Corresponds to parameter beta in the paper. Higher values lead to faster convergence to GP posterior. + range: (0,+inf) + random_score_type : string [default='Gumbel'] + Type of random noise added to scores. + Possible values: + - 'Gumbel' - Gumbel-distributed (as in paper) + - 'NormalWithModelSizeDecrease' - Normally-distributed with deviation decreasing with model iteration count (default in CatBoost) + eps : float, [default=1e-4] + Technical parameter that controls the precision of the prior estimation. + range: (0, 1] + verbose : bool or int + Verbosity of posterior model training output + If verbose is bool, then if set to True, logging_level is set to Verbose, + if set to False, logging_level is set to Silent. + If verbose is int, it determines the frequency of writing metrics to output and + logging_level is set to Verbose. + + Returns + ------- + models : list of trained CatBoostRegressor models (size = samples parameter value) + """ + assert (sigma > 0) + assert (samples > 0) + assert (random_strength > 0) + assert (eps > 0) + + if random_seed is None: + random_seed = 0 + + random_generator = np.random.default_rng(random_seed) + prior_seeds = random_generator.integers(low=0, high=2**63-1, size=samples) + posterior_seeds = random_generator.integers(low=0, high=2**63-1, size=samples) + + N = len(X) + model_shrink_rate = (random_strength / sigma) ** 2 / N + + output_models = [] + tmp_file = tempfile.NamedTemporaryFile() + prior_model_tmp_file = tmp_file.name + + for sample in range(samples): + prior_y = random_generator.normal(scale=eps, size=N) + prior = CatBoostRegressor( + random_seed=prior_seeds[sample], + iterations=prior_iterations, + learning_rate=eps, + loss_function='RMSE', + bootstrap_type='No', + depth=depth, + verbose=False, + leaf_estimation_backtracking="No", + boost_from_average=False, + random_strength=1/eps, + random_score_type=random_score_type, + l2_leaf_reg=0, + score_function="L2", + boosting_type='Plain' + ) + prior.fit( + X, + prior_y, + cat_features=cat_features, + text_features=text_features, + embedding_features=embedding_features, + use_best_model=False + ) + + prior.save_model(prior_model_tmp_file, format="json", pool=X) + with open(prior_model_tmp_file, "r", encoding='utf-8') as prior_file: + prior_json = json.load(prior_file) + for tree in prior_json["oblivious_trees"]: + for ind, (val, weight) in enumerate(zip(tree["leaf_values"], tree["leaf_weights"])): + tree["leaf_values"][ind] = random_generator.normal(scale=np.sqrt(N / np.sqrt(max(1, weight)))) + with open(prior_model_tmp_file, "w") as prior_file: + json.dump(prior_json, prior_file) + prior.load_model(prior_model_tmp_file, format="json") + + scale, bias = prior.get_scale_and_bias() + prior.set_scale_and_bias(scale * sigma / np.sqrt(prior_iterations), bias * sigma / np.sqrt(prior_iterations)) + + posterior_y = y - prior.predict(X) + random_generator.normal(scale=delta, size=N) + posterior = CatBoostRegressor( + random_seed=posterior_seeds[sample], + iterations=posterior_iterations, + learning_rate=learning_rate, + model_shrink_rate=model_shrink_rate, + loss_function='RMSE', + bootstrap_type='No', + depth=depth, + verbose=verbose, + leaf_estimation_backtracking="No", + boost_from_average=False, + random_strength=random_strength, + random_score_type=random_score_type, + l2_leaf_reg=0, + score_function="L2", + boosting_type='Plain' + ) + posterior.fit( + X, + posterior_y, + eval_set=eval_set, + cat_features=cat_features, + text_features=text_features, + embedding_features=embedding_features, + use_best_model=False + ) + + output_models.append(sum_models([prior, posterior], weights=[1, 1])) + + return output_models + + +def cv(pool=None, params=None, dtrain=None, iterations=None, num_boost_round=None, + fold_count=None, nfold=None, inverted=False, partition_random_seed=0, seed=None, + shuffle=True, logging_level=None, stratified=None, as_pandas=True, metric_period=None, + verbose=None, verbose_eval=None, plot=False, plot_file=None, early_stopping_rounds=None, + save_snapshot=None, snapshot_file=None, snapshot_interval=None, metric_update_interval=0.5, + folds=None, type='Classical', return_models=False, log_cout=None, log_cerr=None): + """ + Cross-validate the CatBoost model. + + Parameters + ---------- + pool : catboost.Pool + Data to cross-validate on. + + params : dict + Parameters for CatBoost. + CatBoost has many of parameters, all have default values. + If None, all params still defaults. + If dict, overriding some (or all) params. + + dtrain : catboost.Pool or tuple (X, y) + Synonym for pool parameter. Only one of these parameters should be set. + + iterations : int + Number of boosting iterations. Can be set in params dict. + + num_boost_round : int + Synonym for iterations. Only one of these parameters should be set. + + fold_count : int, optional (default=3) + The number of folds to split the dataset into. + + nfold : int + Synonym for fold_count. + + type : string, optional (default='Classical') + Type of cross-validation + Possible values: + - 'Classical' + - 'Inverted' + - 'TimeSeries' + + inverted : bool, optional (default=False) + Train on the test fold and evaluate the model on the training folds. + + partition_random_seed : int, optional (default=0) + Use this as the seed value for random permutation of the data. + Permutation is performed before splitting the data for cross validation. + Each seed generates unique data splits. + + seed : int, optional + Synonym for partition_random_seed. This parameter is deprecated. Use + partition_random_seed instead. + If both parameters are initialised partition_random_seed parameter is + ignored. + + shuffle : bool, optional (default=True) + Shuffle the dataset objects before splitting into folds. + + logging_level : string, optional (default=None) + Possible values: + - 'Silent' + - 'Verbose' + - 'Info' + - 'Debug' + + stratified : bool, optional (default=None) + Perform stratified sampling. True for classification and False otherwise. + + as_pandas : bool, optional (default=True) + Return pd.DataFrame when pandas is installed. + If False or pandas is not installed, return dict. + + metric_period : int, [default=1] + The frequency of iterations to print the information to stdout. The value should be a positive integer. + + verbose : bool or int + If verbose is bool, then if set to True, logging_level is set to Verbose, + if set to False, logging_level is set to Silent. + If verbose is int, it determines the frequency of writing metrics to output and + logging_level is set to Verbose. + + verbose_eval : bool or int + Synonym for verbose. Only one of these parameters should be set. + + plot : bool, optional (default=False) + If True, draw train and eval error in Jupyter notebook + + plot_file : file-like or str, optional (default=None) + If not None, save train and eval error graphs to file + + early_stopping_rounds : int + Activates Iter overfitting detector with od_wait set to early_stopping_rounds. + + save_snapshot : bool, [default=None] + Enable progress snapshotting for restoring progress after crashes or interruptions + + snapshot_file : string or os.PathLike, [default=None] + Learn progress snapshot file path, if None will use default filename + + snapshot_interval: int, [default=600] + Interval between saving snapshots (seconds) + + metric_update_interval: float, [default=0.5] + Interval between updating metrics (seconds) + + folds: generator or iterator of (train_idx, test_idx) tuples, scikit-learn splitter object or None, optional (default=None) + If generator or iterator, it should yield the train and test indices for each fold. + If object, it should be one of the scikit-learn splitter classes + (https://scikit-learn.org/stable/modules/classes.html#splitter-classes) + and have ``split`` method. + if folds is not None, then all of fold_count, shuffle, partition_random_seed, inverted are None + + return_models: bool, optional (default=False) + if True, return a list of models fitted for each CV fold + + log_cout: output stream or callback for logging (default=None) + If None is specified, sys.stdout is used + + log_cerr: error stream or callback for logging (default=None) + If None is specified, sys.stderr is used + + Returns + ------- + cv results : pandas.core.frame.DataFrame with cross-validation results + columns are: test-error-mean test-error-std train-error-mean train-error-std + cv models : list of trained models, if return_models=True + """ + if params is None: + raise CatBoostError("params should be set.") + + params = deepcopy(params) + _process_synonyms(params) + stringify_builtin_metrics(params) + + metric_period, verbose, logging_level = _process_verbose( + metric_period, verbose, logging_level, verbose_eval) + + if 'loss_function' not in params: + raise CatBoostError("Parameter loss_function should be specified for cross-validation") + + if any(v is not None for v in [fold_count, nfold]) and folds is not None: + raise CatBoostError( + "if folds is not None, then all of fold_count, shuffle, partition_random_seed, inverted are None" + ) + + if folds is not None or type == 'TimeSeries': + shuffle = False + inverted = False + + if verbose is not None: + params.update({ + 'verbose': verbose + }) + + if logging_level is not None: + params.update({ + 'logging_level': logging_level + }) + + if metric_period is not None: + params.update({ + 'metric_period': metric_period + }) + + if early_stopping_rounds is not None: + params.update({ + 'od_type': 'Iter' + }) + if 'od_pval' in params: + del params['od_pval'] + params.update({ + 'od_wait': early_stopping_rounds + }) + + if dtrain is not None: + if pool is None: + pool = dtrain + else: + raise CatBoostError("Only one of the parameters pool and dtrain should be set.") + + if num_boost_round is not None: + if iterations is None: + iterations = num_boost_round + else: + raise CatBoostError("Only one of the parameters iterations and num_boost_round should be set.") + + if iterations is not None: + params.update({ + 'iterations': iterations + }) + + if seed is not None: + partition_random_seed = seed + + if save_snapshot is not None: + params['save_snapshot'] = save_snapshot + + if snapshot_file is not None: + params['snapshot_file'] = snapshot_file + + if snapshot_interval is not None: + params['snapshot_interval'] = snapshot_interval + + if nfold is None and fold_count is None: + if folds is None: + fold_count = 3 + else: + if hasattr(folds, 'get_n_splits'): # sklearn splitter + fold_count = folds.get_n_splits() + elif hasattr(folds, '__len__'): + fold_count = len(folds) + else: + folds = list(folds) + fold_count = len(folds) + elif fold_count is None: + fold_count = nfold + else: + assert nfold is None or nfold == fold_count + + if folds is not None or type == 'TimeSeries': + stratified = False + elif stratified is None: + loss_function = params.get('loss_function', None) + stratified = isinstance(loss_function, STRING_TYPES) and is_cv_stratified_objective(loss_function) + + if 'cat_features' in params: + cat_feature_indices_from_params = _get_features_indices(params['cat_features'], pool.get_feature_names()) + if set(pool.get_cat_feature_indices()) != set(cat_feature_indices_from_params): + raise CatBoostError("categorical features indices in params are different from ones in pool " + + str(cat_feature_indices_from_params) + + " vs " + str(pool.get_cat_feature_indices())) + del params['cat_features'] + + if 'text_features' in params: + text_feature_indices_from_params = _get_features_indices(params['text_features'], pool.get_feature_names()) + if set(pool.get_text_feature_indices()) != set(text_feature_indices_from_params): + raise CatBoostError("text features indices in params are different from ones in pool " + + str(text_feature_indices_from_params) + + " vs " + str(pool.get_text_feature_indices())) + del params['text_features'] + + if 'embedding_features' in params: + embedding_feature_indices_from_params = _get_features_indices(params['embedding_features'], pool.get_feature_names()) + if set(pool.get_embedding_feature_indices()) != set(embedding_feature_indices_from_params): + raise CatBoostError("embedding features indices in params are different from ones in pool " + + str(embedding_feature_indices_from_params) + + " vs " + str(pool.get_embedding_feature_indices())) + del params['embedding_features'] + + train_dir = _get_train_dir(params) + create_dir_if_not_exist(train_dir) + plot_dirs = [] + for step in range(fold_count): + plot_dirs.append(os.path.join(train_dir, 'fold-{}'.format(step))) + for plot_dir in plot_dirs: + create_dir_if_not_exist(plot_dir) + + with log_fixup(log_cout, log_cerr), plot_wrapper(plot, plot_file=plot_file, plot_title='Cross-validation plot', train_dirs=plot_dirs): + if not return_models: + return _cv( + params, + pool, + fold_count, + inverted, + partition_random_seed, + shuffle, + stratified, + metric_update_interval, + as_pandas, + folds, + type, + return_models + ) + else: + results, cv_models = _cv( + params, + pool, + fold_count, + inverted, + partition_random_seed, + shuffle, + stratified, + metric_update_interval, + as_pandas, + folds, + type, + return_models + ) + output_cv_models = _convert_to_catboost(cv_models) + return results, output_cv_models + + +class BatchMetricCalcer(_MetricCalcerBase): + + def __init__(self, catboost, metrics, ntree_start, ntree_end, eval_period, thread_count, tmp_dir): + super(BatchMetricCalcer, self).__init__(catboost) + if tmp_dir is None: + tmp_dir = tempfile.mkdtemp() + delete_temp_dir_flag = True + else: + delete_temp_dir_flag = False + + if isinstance(metrics, STRING_TYPES) or isinstance(metrics, BuiltinMetric): + metrics = [metrics] + metrics = stringify_builtin_metrics_list(metrics) + self._create_calcer(metrics, ntree_start, ntree_end, eval_period, thread_count, tmp_dir, delete_temp_dir_flag) + + +def sum_models(models, weights=None, ctr_merge_policy='IntersectingCountersAverage'): + result = CatBoost() + result._sum_models(models, weights, ctr_merge_policy) + return result + + +def _calc_feature_statistics_layout(go, xaxis, single_pool): + return go.Layout( + yaxis={ + 'title': 'Prediction and target', + 'side': 'left', + 'overlaying': 'y2' + }, + yaxis2={ + 'title': 'Objects per bin' if single_pool else '% pool objects in bin', + 'side': 'right', + 'position': 1.0 + }, + xaxis=xaxis, + legend={ + 'bgcolor': 'rgba(0,0,0,0)', + 'x': 1.07 + } + ) + + +def _build_binarized_feature_statistics_fig(statistics_list, pool_names): + try: + import plotly.graph_objs as go + except ImportError as e: + warnings.warn("To draw binarized feature statistics you should install plotly.") + raise ImportError(str(e)) + + pools_count = len(statistics_list) + data = [] + statistics = statistics_list[0] + if 'borders' in statistics.keys(): + if len(statistics['borders']) == 0: + xaxis = go.layout.XAxis(title='Bins', tickvals=[0]) + return go.Figure(data=[], layout=_calc_feature_statistics_layout(go, xaxis, pools_count == 1)) + + order = np.arange(len(statistics['objects_per_bin'])) + bar_width = 0.8 + xaxis = go.layout.XAxis( + title='Bins', + tickmode='array', + tickvals=list(range(len(statistics['borders']) + 1)), + ticktext=['(-inf, {:.4f}]'.format(statistics['borders'][0])] + + ['({:.4f}, {:.4f}]'.format(val_1, val_2) + for val_1, val_2 in zip(statistics['borders'][:-1], statistics['borders'][1:])] + + ['({:.4f}, +inf)'.format(statistics['borders'][-1])], + showticklabels=False + ) + elif 'cat_values' in statistics.keys(): + order = np.argsort(statistics['objects_per_bin'])[::-1] + bar_width = 0.2 + xaxis = go.layout.XAxis( + title='Cat values', + tickmode='array', + tickvals=list(range(len(statistics['cat_values']))), + ticktext=statistics['cat_values'][order], + showticklabels=True + ) + else: + raise CatBoostError('Expected field "borders" or "cat_values" in binarized feature statistics') + + for i, statistics in enumerate(statistics_list): + if pools_count == 1: + name_suffix = '' + else: + name_suffix = ', {} pool'.format(pool_names[i]) + trace_1 = go.Scatter( + y=statistics['mean_target'][order], + mode='lines+markers', + name='Mean target' + name_suffix, + yaxis='y1', + xaxis='x', + ) + + trace_2 = go.Scatter( + y=statistics['mean_prediction'][order], + mode='lines+markers', + line={'dash' : 'dash'}, + name='Mean prediction on each segment of feature values' + name_suffix, + yaxis='y1', + xaxis='x' + ) + if (len(statistics['mean_weighted_target']) != 0): + trace_3 = go.Scatter( + y=statistics['mean_weighted_target'][order], + mode='lines+markers', + line={'dash' : 'dot'}, + name='Mean weighted target' + name_suffix, + yaxis='y1', + xaxis='x' + ) + + if pools_count > 1: + objects_in_pool = statistics['objects_per_bin'].sum() + color_a = np.array([30, 150, 30]) + color_b = np.array([30, 30, 150]) + color = (color_a * i + color_b * (pools_count - 1 - i)) / float(pools_count - 1) + color = color.astype(int) + trace_4 = go.Bar( + y=statistics['objects_per_bin'][order] / float(objects_in_pool), + width=bar_width / pools_count, + name='% pool objects in bin (total {})'.format(objects_in_pool) + name_suffix, + yaxis='y2', + xaxis='x', + marker={ + 'color': 'rgba({}, {}, {}, 0.4)'.format(*color) + } + ) + else: + trace_4 = go.Bar( + y=statistics['objects_per_bin'][order], + width=bar_width, + name='Objects per bin' + name_suffix, + yaxis='y2', + xaxis='x', + marker={ + 'color': 'rgba(30, 150, 30, 0.4)' + } + ) + + trace_5 = go.Scatter( + y=statistics['predictions_on_varying_feature'][order], + mode='lines+markers', + line={'dash' : 'dashdot'}, + name='Mean prediction with substituted feature' + name_suffix, + yaxis='y1', + xaxis='x' + ) + if (len(statistics['mean_weighted_target']) != 0): + data += [trace_1, trace_2, trace_3, trace_4, trace_5] + else: + data += [trace_1, trace_2, trace_4, trace_5] + + layout = _calc_feature_statistics_layout(go, xaxis, pools_count == 1) + fig = go.Figure(data=data, layout=layout) + + return fig + + +def _plot_feature_statistics_units(statistics, pool_names, feature_name, max_cat_features_on_plot): + if 'cat_values' in statistics[0].keys() and len(statistics[0]['cat_values']) > max_cat_features_on_plot: + figs = [] + for begin in range(0, len(statistics[0]['cat_values']), max_cat_features_on_plot): + end = begin + max_cat_features_on_plot + statistics_keys = ['cat_values', 'mean_target', 'mean_weighted_target', 'mean_prediction', + 'objects_per_bin', 'predictions_on_varying_feature'] + sub_statistics = dict([(k, dict([(key, stats[key][begin : end]) for key in statistics_keys])) for k, stats in statistics]) + fig = _build_binarized_feature_statistics_fig(sub_statistics, pool_names) + feature_name_with_part_suffix = '{}_parts[{}:{}]'.format(feature_name, begin, end) + figs += [(fig, feature_name_with_part_suffix)] + return figs + else: + fig = _build_binarized_feature_statistics_fig(statistics, pool_names) + return [(fig, feature_name)] + + +def _plot_feature_statistics(statistics_by_feature, pool_names, feature_names, max_cat_features_on_plot): + figs_with_names = [] + for feature_num in statistics_by_feature: + feature_name = feature_names[feature_num] + statistics = statistics_by_feature[feature_num] + need_skip = 'cat_values' not in statistics[0].keys() + if 'borders' in statistics[0].keys(): + for stats in statistics: + if len(stats['borders']) > 0: + need_skip = False + if need_skip: + continue + figs_with_names += _plot_feature_statistics_units(statistics, pool_names, feature_name, max_cat_features_on_plot) + + main_fig = figs_with_names[0][0] + buttons = [] + for fig, feature_name in figs_with_names: + buttons.append( + dict( + label=feature_name, + method='update', + args=[{'y': [data.y for data in fig.data]}, {'xaxis': fig.layout.xaxis}], + ), + ) + main_fig.update_layout( + updatemenus=[ + dict( + direction='down', + pad={'r': 10, 't': 10}, + showactive=True, + x=0.25, + xanchor='left', + y=1.09, + yanchor='top', + buttons=buttons, + ) + ], + annotations=[ + dict(text='Statistics for feature', showarrow=False, + x=0, xref='paper', y=1.05, yref='paper', align='left') + ] + ) + return main_fig + + +def _to_subclass(model, subclass): + """ + Convert a CatBoost model to a sklearn-compatible model. + + Parameters + ---------- + model : CatBoost model + a model to convert from + + subclass : an sklearn-compatible class + a class to convert to : CatBoostClassifier, CatBoostRegressor or CatBoostRanker + + Returns + ------- + a converted model : `subclass` type + a model converted from the initial CatBoost `model` to a sklearn-compatible `subclass` model + """ + if isinstance(model, subclass): + return model + if not isinstance(model, CatBoost): + raise CatBoostError('model should be a subclass of CatBoost') + + converted_model = subclass.__new__(subclass) + + # TODO(ilyzhin) change on get_all_params after MLTOOLS-4758 + params = model._get_canonized_params() + if 'loss_function' in params: + subclass._check_is_compatible_loss(params['loss_function']) + + for attr in model.__dict__: + setattr(converted_model, attr, getattr(model, attr)) + return converted_model + + +def to_regressor(model): + return _to_subclass(model, CatBoostRegressor) + + +def to_classifier(model): + return _to_subclass(model, CatBoostClassifier) + + +def to_ranker(model): + return _to_subclass(model, CatBoostRanker) + + +class _TrainCallbacksWrapper(object): + def __init__(self, callbacks): + self._callbacks = callbacks + + def after_iteration(self, info): + for cb in self._callbacks: + if not cb.after_iteration(info): + return False + return True diff --git a/v15cat/pylib/catboost/datasets.py b/v15cat/pylib/catboost/datasets.py new file mode 100644 index 0000000000000000000000000000000000000000..c2572883814f535d1afbdf9d6ee7364bf6dbc82e --- /dev/null +++ b/v15cat/pylib/catboost/datasets.py @@ -0,0 +1,336 @@ +import hashlib +import logging +import numpy as np +import os +import pandas as pd +import tarfile +import tempfile +import six +import shutil + +from .core import PATH_TYPES, fspath + + +logger = logging.getLogger(__name__) + + +def _extract(src_file, dst_dir='.'): + cur_dir = os.getcwd() + os.chdir(dst_dir) + try: + with tarfile.open(src_file, 'r:gz') as f: + f.extractall() + finally: + os.chdir(cur_dir) + + +def _calc_md5(path): + hasher = hashlib.md5() + with open(path, 'rb') as f: + while True: + block = f.read(65536) + if not block: + break + hasher.update(block) + return hasher.hexdigest() + + +def _ensure_dir_exists(path): + try: + os.makedirs(path) + except OSError: + if not os.path.isdir(path): + raise + + +def _cached_download(url, md5, dst): + if os.path.isfile(dst) and _calc_md5(dst) == md5: + return + + def reporthook(blocknum, bs, size): + logger.debug('downloaded %s bytes', size) + + urls = url if isinstance(url, list) or isinstance(url, tuple) else (url, ) + + for u in urls: + try: + six.moves.urllib.request.urlretrieve(u, dst, reporthook=reporthook) + break + except (six.moves.urllib.error.URLError, IOError): + logger.debug('failed to download from %s', u) + else: + raise RuntimeError('failed to download from %s', urls) + + dst_md5 = _calc_md5(dst) + if dst_md5 != md5: + raise RuntimeError('md5 sum mismatch for url {url}; expected {expected}, but got {got}'.format( + url=u, expected=md5, got=dst_md5)) + + +_cache_path = None + + +def _get_cache_path(): + global _cache_path + if _cache_path is None: + _cache_path = os.path.join(os.getcwd(), 'catboost_cached_datasets') + return _cache_path + + +def set_cache_path(path): + assert isinstance(path, PATH_TYPES), 'expected string or pathlib.Path' + global _cache_path + _cache_path = fspath(path) + + +def _download_dataset(url, md5, dataset_name, train_file, test_file, cache=False): + # TODO(yazevnul): this is not thread safe (or process safe?), we should take a file lock when + # enter this function to avoid dataset being overwritten or corrupted or something else that may + # have happen when OS operated simultaneously on the same file. Same thing should probably be + # done with `_cached_download`. + dir_path = os.path.join(_get_cache_path(), dataset_name) if cache else tempfile.mkdtemp() + train_path = os.path.join(dir_path, train_file) + test_path = os.path.join(dir_path, test_file) + if not (os.path.exists(train_path) and os.path.exists(test_path)): + _ensure_dir_exists(dir_path) + file_descriptor, file_path = tempfile.mkstemp() + os.close(file_descriptor) + try: + _cached_download(url, md5, file_path) + _extract(file_path, dir_path) + finally: + os.remove(file_path) + # move files for safe delete of temp dir + if not cache: + fd_new_train, new_train_path = tempfile.mkstemp() + fd_new_test, new_test_path = tempfile.mkstemp() + os.close(fd_new_train) + os.close(fd_new_test) + os.replace(train_path, new_train_path) + os.replace(test_path, new_test_path) + shutil.rmtree(dir_path) + train_path, test_path = new_train_path, new_test_path + return train_path, test_path + + +def _load_dataset_pd(url, md5, dataset_name, train_file, test_file, sep=',', header='infer', cache=False): + train_path, test_path = _download_dataset(url, md5, dataset_name, train_file, test_file, cache) + train, test = pd.read_csv(train_path, header=header, sep=sep), pd.read_csv(test_path, header=header, sep=sep) + if not cache: + os.remove(train_path) + os.remove(test_path) + return train, test + + +def _load_numeric_only_dataset(path, row_count, column_count, sep='\t'): + # - can't use `pandas.read_csv` because it may result in 5x overhead + # - can't use `numpy.loadtxt` because it may result in 3x overhead + # And both mentioned above solutions are very slow compared to the one implemented below. + dataset = np.zeros((row_count, column_count, ), dtype=np.float32, order='F') + with open(path, 'rb') as f: + for line_idx, line in enumerate(f): + # `str.split()` is too slow, use `numpy.fromstring()` + row = np.fromstring(line, dtype=np.float32, sep=sep) + assert row.size == column_count, 'got too many columns at line %d (expected %d columns, got %d)' % (line_idx + 1, column_count, row.size) + # doing `dataset[line_idx][:]` instead of `dataset[line_idx]` is here on purpose, + # otherwise we may reallocate memory, while here we just copy + dataset[line_idx][:] = row + + assert line_idx + 1 == row_count, 'got too many lines (expected %d lines, got %d)' % (row_count, line_idx + 1) + + return pd.DataFrame(dataset) + + +def titanic(): + url = 'https://storage.mds.yandex.net/get-devtools-opensource/233854/titanic.tar.gz' + md5 = '9c8bc61d545c6af244a1d37494df3fc3' + dataset_name, train_file, test_file = 'titanic', 'train.csv', 'test.csv' + return _load_dataset_pd(url, md5, dataset_name, train_file, test_file) + + +def amazon(): + url = 'https://storage.mds.yandex.net/get-devtools-opensource/250854/amazon.tar.gz' + md5 = '8fe3eec12bfd9c4c532b24a181d0aa2c' + dataset_name, train_file, test_file = 'amazon', 'train.csv', 'test.csv' + return _load_dataset_pd(url, md5, dataset_name, train_file, test_file) + + +def msrank(): + url = 'https://storage.mds.yandex.net/get-devtools-opensource/233854/msrank.tar.gz' + md5 = '34fee225d02419adc106581f4eb36f2e' + dataset_name, train_file, test_file = 'msrank', 'train.tsv', 'test.tsv' + return _load_dataset_pd(url, md5, dataset_name, train_file, test_file, header=None, sep='\t', cache=True) + + +def msrank_10k(): + url = 'https://storage.mds.yandex.net/get-devtools-opensource/250854/msrank_10k.tar.gz' + md5 = '79c5b67397289c4c8b367c1f34629eae' + dataset_name, train_file, test_file = 'msrank_10k', 'train.csv', 'test.csv' + return _load_dataset_pd(url, md5, dataset_name, train_file, test_file, header=None) + + +def rotten_tomatoes(): + """ + Contains information from kaggle [1], which is made available here under the Open Database License (ODbL) [2]. + + Download "rotten_tomatoes" [1] data set. + + Will return two pandas.DataFrame-s, first with train part (rotten_tomatoes.data) and second with test part + (rotten_tomatoes.test) of the dataset. + + NOTE: This is a preprocessed version of the dataset. + + [1]: https://www.kaggle.com/rpnuser8182/rotten-tomatoes + [2]: https://opendatacommons.org/licenses/odbl/1-0/index.html + """ + url = 'https://catboost-opensource.s3.yandex.net/rotten_tomatoes.tar.gz' + md5 = 'a07fed612805ac9e17ced0d82a96add4' + dataset_name, train_file, test_file = 'rotten_tomatoes', 'learn.tsv', 'test.tsv' + return _load_dataset_pd(url, md5, dataset_name, train_file, test_file, sep='\t') + + +def imdb(): + url = 'https://catboost-opensource.s3.yandex.net/imdb.tar.gz' + md5 = '0fd62578d631ac3d71a71c3e6ced6f8b' + dataset_name, train_file, test_file = 'imdb', 'learn.tsv', 'test.tsv' + return _load_dataset_pd(url, md5, dataset_name, train_file, test_file, sep='\t') + + +def epsilon(): + """ + Download "epsilon" [1] data set. + + Will return two pandas.DataFrame-s, first with train part (epsilon_normalized) and second with + test part (epsilon_normalized.t) of the dataset. Object class will be located in the first + column of dataset. + + NOTE: This is a preprocessed version of the dataset. It was converted from libsvm format into + tsv (CatBoost doesn't support libsvm format out of the box). + + [1]: https://www.csie.ntu.edu.tw/~cjlin/libsvmtools/datasets/binary.html#epsilon + """ + urls = ( + 'https://proxy.sandbox.yandex-team.ru/785711439', + 'https://storage.mds.yandex.net/get-devtools-opensource/250854/epsilon.tar.gz', ) + md5 = '5bbfac403ac673da7d7ee84bd532e973' + dataset_name, train_file, test_file = 'epsilon', 'train.tsv', 'test.tsv' + train_path, test_path = _download_dataset(urls, md5, dataset_name, train_file, test_file, cache=True) + return ( + _load_numeric_only_dataset(train_path, 400000, 2001, sep='\t'), + _load_numeric_only_dataset(test_path, 100000, 2001, sep='\t')) + + +def monotonic1(): + """ + Dataset with monotonic constraints. + Can be used for poisson regression. + Has several numerical and several categorical features. + The first column contains target values. Columns with names Cat* contain categorical features. + Columns with names Num* contain numerical features. + + Dataset also contains several numerical features, for which monotonic constraints must hold. + For features in columns named MonotonicNeg*, if feature value decreases, then prediction value must not decrease. + Thus, if there are two samples x1, x2 with all features being equal except + for a monotonic negative feature M, such that x1[M] > x2[M], then the following inequality must + hold for predictions: f(x1) <= f(x2) + """ + url = 'https://storage.mds.yandex.net/get-devtools-opensource/479623/monotonic1.tar.gz' + md5 = '1b9d8e15bc3fd6f1498e652e7fc4f4ca' + dataset_name, train_file, test_file = 'monotonic1', 'train.tsv', 'test.tsv' + return _load_dataset_pd(url, md5, dataset_name, train_file, test_file, sep='\t', cache=True) + + +def monotonic2(): + """ + Dataset with monotonic constraints. + Can be used for regression. + The first column contains target values. + Other columns contain contain numerical features, for which monotonic constraints must hold. + + For features in columns named MonotonicNeg*, if feature value decreases, then prediction + value must not decrease. Thus, if there are two samples x1, x2 with all features being + equal except for a monotonic negative feature MNeg, such that x1[MNeg] > x2[MNeg], then + the following inequality must hold for predictions: f(x1) <= f(x2) + For features in columns named MonotonicPos*, if feature value decreases, then prediction + value must not increase. Thus, if there are two samples x1, x2 with all features being + equal except for a monotonic positive feature MPos, such that x1[MPos] > x2[MPos], + then the following inequality must hold for predictions: f(x1) >= f(x2) + """ + url = 'https://storage.mds.yandex.net/get-devtools-opensource/250854/monotonic2.tar.gz' + md5 = 'ce559e212cb72c156269f6f9a641baca' + dataset_name, train_file, test_file = 'monotonic2', 'train.tsv', 'test.tsv' + return _load_dataset_pd(url, md5, dataset_name, train_file, test_file, sep='\t') + + +def adult(): + """ + Download "Adult Data Set" [1] from UCI Machine Learning Repository. + + Will return two pandas.DataFrame-s, first with train part (adult.data) and second with test part + (adult.test) of the dataset. + + [1]: https://archive.ics.uci.edu/ml/datasets/Adult + """ + # via https://archive.ics.uci.edu/ml/machine-learning-databases/adult/adult.names + names = ( + 'age', 'workclass', 'fnlwgt', 'education', 'education-num', 'marital-status', 'occupation', + 'relationship', 'race', 'sex', 'capital-gain', 'capital-loss', 'hours-per-week', + 'native-country', 'income', ) + dtype = { + 'age': float, 'workclass': object, 'fnlwgt': float, 'education': object, + 'education-num': float, 'marital-status': object, 'occupation': object, + 'relationship': object, 'race': object, 'sex': object, 'capital-gain': float, + 'capital-loss': float, 'hours-per-week': float, + 'native-country': object, 'income': object, } + + # proxy.sandbox.yandex-team.ru is Yandex internal storage, we first try to download it from + # internal storage to avoid putting too much pressure on UCI storage from our internal CI + + train_urls = ( + 'https://proxy.sandbox.yandex-team.ru/779118052', + 'https://archive.ics.uci.edu/ml/machine-learning-databases/adult/adult.data', ) + train_md5 = '5d7c39d7b8804f071cdd1f2a7c460872' + fd_train, train_path = tempfile.mkstemp() + os.close(fd_train) + _cached_download(train_urls, train_md5, train_path) + + test_urls = ( + 'https://proxy.sandbox.yandex-team.ru/779120000', + 'https://archive.ics.uci.edu/ml/machine-learning-databases/adult/adult.test', ) + test_md5 = '35238206dfdf7f1fe215bbb874adecdc' + fd_test, test_path = tempfile.mkstemp() + os.close(fd_test) + _cached_download(test_urls, test_md5, test_path) + + train_df = pd.read_csv(train_path, names=names, header=None, sep=r',\s*', na_values=['?'], engine='python') + os.remove(train_path) + + # lines in test part end with dot, thus we need to fix last column of the dataset + test_df = pd.read_csv(test_path, names=names, header=None, sep=r',\s*', na_values=['?'], skiprows=1, converters={'income': lambda x: x[:-1]}, engine='python') + os.remove(test_path) + + # pandas 0.19.1 doesn't support `dtype` parameter for `read_csv` when `python` engine is used, + # so we have to do the casting manually; also we can't use `converters` together with `dtype` + train_df = train_df.astype(dtype) + test_df = test_df.astype(dtype) + + return train_df, test_df + + +def higgs(): + """ + Download "higgs" [1] data set. + + Will return two pandas.DataFrame-s, first with train part and second with + test part of the dataset. Object class will be located in the first + column of dataset. + + [1]: https://archive.ics.uci.edu/ml/datasets/HIGGS + """ + url = 'https://storage.mds.yandex.net/get-devtools-opensource/250854/higgs.tar.gz' + md5 = 'ad59ba8328a9afa3837d7bf1a0e10e7b' + dataset_name, train_file, test_file = 'higgs', 'train.tsv', 'test.tsv' + train_path, test_path = _download_dataset(url, md5, dataset_name, train_file, test_file, cache=True) + return ( + _load_numeric_only_dataset(train_path, 10500000, 29, sep='\t'), + _load_numeric_only_dataset(test_path, 500000, 29, sep='\t')) diff --git a/v15cat/pylib/catboost/dev_utils.py b/v15cat/pylib/catboost/dev_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..5b3ece422d9954ad5bf1a12e61aaf54b43e4654c --- /dev/null +++ b/v15cat/pylib/catboost/dev_utils.py @@ -0,0 +1,16 @@ +from . import _catboost + +is_groupwise_metric = _catboost.is_groupwise_metric +is_multiclass_metric = _catboost.is_multiclass_metric +is_multitarget_objective = _catboost.is_multitarget_objective +is_classification_objective = _catboost.is_classification_objective +get_multi_quantile_approx_size = _catboost.get_multi_quantile_approx_size +is_pairwise_metric = _catboost.is_pairwise_metric +is_maximizable_metric = _catboost.is_maximizable_metric +is_minimizable_metric = _catboost.is_minimizable_metric +is_user_defined_metric = _catboost.is_user_defined_metric +has_gpu_implementation_metric = _catboost.has_gpu_implementation_metric +need_dataset_for_leaves_weights = _catboost.need_dataset_for_leaves_weights +get_experiment_name = _catboost.get_experiment_name +convert_features_to_indices = _catboost.convert_features_to_indices +get_num_feature_values_sample = _catboost.get_num_feature_values_sample diff --git a/v15cat/pylib/catboost/eval/__init__.py b/v15cat/pylib/catboost/eval/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c82f6e0124c86bc95f2fce3e79b84d085a6067f9 --- /dev/null +++ b/v15cat/pylib/catboost/eval/__init__.py @@ -0,0 +1,20 @@ +from .catboost_evaluation import * +from .evaluation_result import * +from .execution_case import * +from .utils import * + +__all__ = [ + 'EvalType', + 'CatboostEvaluation', + 'ScoreType', + 'ScoreConfig', + 'CaseEvaluationResult', + 'MetricEvaluationResult', + 'EvaluationResults', + 'calc_wilcoxon_test', + 'calc_bootstrap_ci_for_mean', + 'ExecutionCase', + 'make_dirs_if_not_exists', + 'series_to_line', + 'save_plot' +] diff --git a/v15cat/pylib/catboost/eval/_fold_model.py b/v15cat/pylib/catboost/eval/_fold_model.py new file mode 100644 index 0000000000000000000000000000000000000000..bf856f3e43742f1d8a0d92bd63523f761eecb3ea --- /dev/null +++ b/v15cat/pylib/catboost/eval/_fold_model.py @@ -0,0 +1,35 @@ +""" +This class provides the abstraction of model. I.e. it is an object relevant to particular learn fold and particular +algorithm with parameters. +""" + +import os +from .. import CatBoostError, CatBoost + + +class FoldModel: + + def __init__(self, case, model_path, model_id): + self._id = model_id + self._case = case + self._model_path = model_path + + def __repr__(self): + return 'Model_{}_id_{}'.format(str(self._case), self._id) + + def create_metrics_calcer(self, metrics, thread_count, eval_step=1): + if not os.path.exists(self._model_path): + raise CatBoostError("Model was deleted. Can't create calcer now") + model = CatBoost() + model.load_model(self._model_path) + return model.create_metric_calcer(metrics, thread_count=thread_count, eval_period=eval_step) + + def get_case(self): + return self._case + + def get_fold_id(self): + return self._id + + def delete(self): + if os.path.exists(self._model_path): + os.remove(self._model_path) diff --git a/v15cat/pylib/catboost/eval/_fold_models_handler.py b/v15cat/pylib/catboost/eval/_fold_models_handler.py new file mode 100644 index 0000000000000000000000000000000000000000..70746c8e7fd1cbe4c5438bb16967acc5a6754cea --- /dev/null +++ b/v15cat/pylib/catboost/eval/_fold_models_handler.py @@ -0,0 +1,218 @@ +from __future__ import print_function + +import os +import time + +from .. import CatBoostError +from ..eval.log_config import get_eval_logger +from ..eval.utils import make_dirs_if_not_exists +from ..eval.evaluation_result import CaseEvaluationResult +from ._fold_model import FoldModel + + +class FoldModelsHandler(object): + """ + Class that is responsible for learning models and computing their metrics + """ + """ + All models are placed to the default directory "models". + """ + __MODEL_DIR = 'models' + + @staticmethod + def _remove_model_dir(): + try: + if os.path.exists(FoldModelsHandler.__MODEL_DIR): + os.rmdir(FoldModelsHandler.__MODEL_DIR) + except OSError as err: + get_eval_logger().warning(str(err)) + + def __init__(self, metrics, cases, thread_count, eval_step, remove_models): + """ + Args: + :param remove_models: Set true if you want models to be removed after applying them. + + """ + self._cases = cases + self._metrics = metrics + + self._case_results = dict() + for case in self._cases: + self._case_results[case] = dict() + + self._thread_count = thread_count + self._eval_step = eval_step + self._flag_remove_models = remove_models + self._metric_descriptions = None + + def _init_case_results(self, metric_descriptions): + self._metric_descriptions = metric_descriptions + for case in self._cases: + case_result = self._case_results[case] + for metric_description in self._metric_descriptions: + case_result[metric_description] = CaseEvaluationResult(case, + metric_description, + eval_step=self._eval_step) + + def _compute_metrics(self, metrics, grouped_by_case_models, learn_folds, skipped_folds, rest_folds): + metric_calcers = {} + + for case, case_models in grouped_by_case_models.items(): + metric_calcers[case] = list() + for case_model in case_models: + + metric_calcer = case_model.create_metrics_calcer(metrics, + eval_step=self._eval_step, + thread_count=self._thread_count) + metric_calcers[case].append(metric_calcer) + + if self._metric_descriptions is None: + self._init_case_results(metric_calcer.metric_descriptions()) + elif self._metric_descriptions != metric_calcer.metric_descriptions(): + raise CatBoostError("Error: metric names should be consistent") + + for file_num, fold_file in enumerate(learn_folds + skipped_folds + rest_folds): + pool = FoldModelsHandler._create_pool(fold_file, self._thread_count) + + for case, case_models in grouped_by_case_models.items(): + calcers = metric_calcers[case] + + for model_num, model in enumerate(case_models): + if file_num != model_num: + calcers[model_num].add(pool) + + for case, case_models in grouped_by_case_models.items(): + calcers = metric_calcers[case] + case_results = self._case_results[case] + for calcer, model in zip(calcers, case_models): + scores = calcer.eval_metrics() + for metric in self._metric_descriptions: + case_results[metric]._add(model, scores.get_result(metric)) + + @staticmethod + def _fit_model(pool, case, fold_id, model_path): + from .. import CatBoost + # Learn model + make_dirs_if_not_exists(FoldModelsHandler.__MODEL_DIR) + + feature_count = pool.num_col() + if "ignored_features" in case.get_params(): + ignored_features = case.get_params()["ignored_features"] + if len(ignored_features) and max(ignored_features) >= feature_count: + raise CatBoostError("Error: input parameter contains feature indices wich are not available in pool: " + "{}\n " + "Check eval_feature set and ignored features options".format(ignored_features)) + get_eval_logger().debug('Learn model {} on fold #{}'.format(str(case), fold_id)) + cur_time = time.time() + instance = CatBoost(params=case.get_params()) + instance.fit(pool) + instance.save_model(fname=model_path) + + get_eval_logger().debug('Operation was done in {} seconds'.format(time.time() - cur_time)) + return FoldModel(case, model_path, fold_id) + + def _fit_models(self, learn_files, fold_id_bias): + """ + Train models for each algorithm and learn dataset(folds). Than return them. + + Args: + :param learn_files: Entities of FoldStorage for learning models. + :return: Dictionary of models where the key is case and the value is models on learn folds + """ + make_dirs_if_not_exists(FoldModelsHandler.__MODEL_DIR) + + models = {} + for case in self._cases: + models[case] = list() + + for file_num, learn_file in enumerate(learn_files): + pool = FoldModelsHandler._create_pool(learn_file, self._thread_count) + fold_id = fold_id_bias + file_num + + for case in self._cases: + model_path = os.path.join(FoldModelsHandler.__MODEL_DIR, + FoldModelsHandler._create_model_name(case, fold_id)) + get_eval_logger().debug("For model {} on fold #{} path is {}".format(str(case), fold_id, model_path)) + fold_model = self._fit_model(pool, case, fold_id, model_path) + get_eval_logger().info("Model {} on fold #{} was fitted".format(str(case), fold_id)) + models[case].append(fold_model) + + return models + + def proceed(self, splitter, fold_size, folds_count, fold_offset): + """ + Run all processes to gain stats. It applies algorithms to fold files that gains from learning. It keeps + stats inside models and models are stored in DataFrame. Columns are matched to the different algos and rows to + the folds. + + Args: + :param splitter: Splitter entity. + :param fold_size: The size of fold. + :param folds_count: Count of golds. + :param fold_offset: The offset (count of folds that we want to skip). + :return: return dict: keys metric to CaseEvaluationResult + + """ + try: + folds_sets = splitter.create_fold_sets(fold_size, folds_count) + fold_groups_files = splitter.fold_groups_files_generator(folds_sets, + fold_offset) + fold_id_bias = fold_offset + + for learn_folds, skipped_folds, rest_folds in fold_groups_files: + if len(learn_folds) == 0: + continue + list_models = [] + try: + permutation_models = self._fit_models(learn_folds, fold_id_bias) + for case, case_models in permutation_models.items(): + list_models += case_models + + learn_folds_count = len(learn_folds) + get_eval_logger().info("Start metric computation for folds [{}, {})" + .format(fold_id_bias, fold_id_bias + learn_folds_count)) + self._compute_metrics(self._metrics, + permutation_models, + learn_folds, skipped_folds, rest_folds) + get_eval_logger().info("Computation of metrics for folds [{}, {}) is completed" + .format(fold_id_bias, fold_id_bias + learn_folds_count)) + fold_id_bias += learn_folds_count + finally: + # Do it each step because don't want to occupy a lot of memory + splitter.clean_folds() + self._remove_models(list_models) + + finally: + # Also sometimes we need to remove models and always need to try to remove folds directory. + self._clean(splitter) + + return self._case_results + + def _clean(self, splitter): + if self._flag_remove_models: + self._remove_model_dir() + splitter.clean() + + @staticmethod + def _remove_models(list_models): + get_eval_logger().debug('Remove models {}'.format(str(list_models))) + for model in list_models: + model.delete() + + @staticmethod + def _create_model_name(model_case, fold): + import uuid + id_str = str(uuid.uuid1()).replace("-", "_") + model_name = "model_{_name}_fold_{_fold}_{_uuid}.bin".format(_name=model_case.get_label(), + _fold=fold, + _uuid=id_str) + return model_name + + @staticmethod + def _create_pool(fold_file, thread_count=-1): + from .. import Pool + data_pool = Pool(fold_file.path(), + column_description=fold_file.column_description(), + delimiter=fold_file.get_separator(), + thread_count=thread_count) + return data_pool diff --git a/v15cat/pylib/catboost/eval/_fold_storage.py b/v15cat/pylib/catboost/eval/_fold_storage.py new file mode 100644 index 0000000000000000000000000000000000000000..999dd2c7e2ce409a1feea99ed4b4461b40c08da8 --- /dev/null +++ b/v15cat/pylib/catboost/eval/_fold_storage.py @@ -0,0 +1,134 @@ +""" +This module contains the abstractions for keeping small amount of data. It provides the interface for classes and a few +realisations. +""" + +from __future__ import print_function + +import os + +from .. import CatBoostError +from ..eval.log_config import get_eval_logger +from .utils import make_dirs_if_not_exists + + +# Interface +class FoldStorage(object): + """ + Base class. + """ + default_dir = 'folds' + + @staticmethod + def remove_dir(): + """ + Remove default directory for folds if there're no files nut models. In other way it raises warning. + + Args: + :return: Nothing. + + """ + try: + if os.path.exists(_FoldFile.default_dir): + os.rmdir(_FoldFile.default_dir) + except OSError as err: + get_eval_logger().warning(err.message) + + def __init__(self, fold, storage_name, sep, column_description): + self._fold = fold + self._storage_name = storage_name + self._column_description = column_description + self._sep = sep + self._size = 0 + + def get_separator(self): + """ + Args: + :return: (str) Delimiter for data used when we saved fold to file. + + """ + return self._sep + + def __str__(self): + return self._storage_name + + def column_description(self): + """ + Args: + :return: (str) Path to the column description. + + """ + return self._column_description + + def contains_group_id(self, group_id): + """ + Args: + :param group_id: (int) The number of group we want to check. + :return: True if fold contains line or lines with that group id. + + """ + return group_id in self._fold + + def open(self): + raise NotImplementedError("The base class don't have delete method. Please, use successor.") + + def close(self): + raise NotImplementedError("The base class don't have delete method. Please, use successor.") + + def delete(self): + raise NotImplementedError("The base class don't have delete method. Please, use successor.") + + +class _FoldFile(FoldStorage): + """ + FoldFile is the realisation of the interface of FoldStorage. It always saves data to file before reset them. + All files place to the special directory 'folds'. + """ + + def __init__(self, fold, storage_name, sep, column_description): + super(_FoldFile, self).__init__( + fold, storage_name, + sep=sep, column_description=column_description + ) + self._file_path = os.path.join(self.default_dir, + storage_name) + self._prepare_path() + self._lines = [] + self._file = None + + def _prepare_path(self): + make_dirs_if_not_exists(self.default_dir) + open(self._file_path, 'w').close() # clean file + + def path(self): + return self._file_path + + def add(self, line): + self._size += 1 + print(line, file=self._file, end='') + + def add_all(self, lines): + [self.add(line) for line in lines] + + def open(self): + if self._file is None: + self._file = open(self._file_path, mode='a') + else: + raise CatBoostError("File already opened {}".format(self._file_path)) + + def is_opened(self): + return self._file is not None + + def close(self): + if self._file is None: + raise CatBoostError("Trying to close None {}".format(self._file_path)) + + self._file.close() + self._file = None + + def delete(self): + if self._file is not None: + raise CatBoostError("Close file before delete") + + if os.path.exists(self._file_path): + os.remove(self._file_path) diff --git a/v15cat/pylib/catboost/eval/_readers.py b/v15cat/pylib/catboost/eval/_readers.py new file mode 100644 index 0000000000000000000000000000000000000000..2c3c8358331137b543a78535ea2becb7017c3e1c --- /dev/null +++ b/v15cat/pylib/catboost/eval/_readers.py @@ -0,0 +1,72 @@ +""" +Simple file reader. It intends to read lines from big files. Also it provides the group id of each line it reads. +""" + +import pandas as pd +from .utils import series_to_line + + +class _BaseReader(object): + def __init__(self, sep, group_feature_num): + self._sep = sep + self._group_feature_num = group_feature_num + + def get_separator(self): + return self._sep + + def lines_generator(self): + raise NotImplementedError("The base class don't have any lines_reader") + + def pack_lines_generator(self, pack_size): + lines = [] + group_ids = [] + current_pack_size = 0 + + lines_generator = self.lines_generator() + for group_id, line in lines_generator: + group_ids.append(group_id) + lines.append(line) + current_pack_size += 1 + if current_pack_size == pack_size: + yield group_ids, lines + lines = [] + group_ids = [] + current_pack_size = 0 + if current_pack_size != 0: + yield group_ids, lines + + +class _SimpleStreamingFileReader(_BaseReader): + def __init__(self, file_name, sep, has_header, group_feature_num=None): + super(_SimpleStreamingFileReader, self).__init__(sep, group_feature_num) + self._has_header = has_header + self._file_name = file_name + + def lines_generator(self): + with open(self._file_name, 'r') as file: + if self._has_header: + file.readline() + for num, line in enumerate(file): + if self._group_feature_num is None: + group_id = num + else: + features = line.strip().split(self._sep, self._group_feature_num + 1) + group_id = features[self._group_feature_num] + yield int(float(group_id)), line + + +# Can't handle big data. Can be used for tests. +class _SimpleDataReader(_BaseReader): + def __init__(self, data, sep, group_feature_num=None): + super(_SimpleDataReader, self).__init__(sep, group_feature_num) + self._data = pd.DataFrame(data) + + def lines_generator(self): + for num, (index, line) in enumerate(self._data.iterrows()): + if self._group_feature_num is None: + yield num, series_to_line(line, self._sep) + '\n' + else: + yield line.iloc[self._group_feature_num], series_to_line(line, self._sep) + '\n' + + def get_matrix(self): + return self._data diff --git a/v15cat/pylib/catboost/eval/_splitter.py b/v15cat/pylib/catboost/eval/_splitter.py new file mode 100644 index 0000000000000000000000000000000000000000..eff56e568105295a459067320dd6f1760cd2cef8 --- /dev/null +++ b/v15cat/pylib/catboost/eval/_splitter.py @@ -0,0 +1,165 @@ +""" +class Splitter. + +Convenient tool for creating and working with folds. +""" + +import random + +from ._fold_storage import FoldStorage +from ._fold_storage import _FoldFile + + +class _Splitter(object): + """ + Splitter needs providing some parameters to create folds and some "reader", + that can read source. + """ + _REST_SIZE = 100000 + + def __init__(self, line_reader, column_description, seed, min_folds_count): + self._line_reader = line_reader + self._line_groups_ids, self._groups_ids = self._read_groups_ids() + # line_groups_ids -- group ids of each line + # groups_ids -- set of all group ids in dataset + self._folds_storage = set() + # keeps it for removing at the end of work. + self._column_description = column_description + self._min_folds_count = min_folds_count + self._random = random.Random(seed) + + # line_reader -- Reader for getting lines from file. It have to support iteration through lines. + # column_description -- The description of the features of dataset. + + def _read_groups_ids(self): + """Find all groups in dataset and which group each line belongs.""" + line_groups_ids = [] + groups_ids = set() + + lines = self._line_reader.lines_generator() + # Need to return group id of line and line as string. + for group_id, _ in lines: + line_groups_ids.append(group_id) + groups_ids.add(group_id) + return line_groups_ids, groups_ids + + def _make_learn_folds(self, fold_size, left_folds): + """Prepare test sets for folds only for one permutation""" + count_groups = len(self._groups_ids) + if count_groups // self._min_folds_count < fold_size: + raise AttributeError('The size of fold is too big: count_groups: {}, fold_size: {}. Const: {}'.format( + count_groups, fold_size, self._min_folds_count) + ) + + permutation = sorted(self._groups_ids) + self._random.shuffle(permutation) + + result = [] + current_count_folds = min(count_groups // fold_size, left_folds) + for i in range(current_count_folds): + result.append(set(permutation[i * fold_size: (i + 1) * fold_size])) + return result + + def _write_folds(self, fold_storages, num, offset): + """Learn_set contains numbers of lines. The method itself store relevant lines from dataset to fold storage.""" + + generator = self._line_reader.lines_generator() + # Need to return group id of line and line as string. + + for fold_storage in fold_storages: + fold_storage.open() + + try: + rest_folds = [] + rest_fold_file = self.create_fold(None, 'offset{}_rest'.format(offset), num) + rest_fold_file.open() + num += 1 + rest_size = 0 + + for num_line, (_, line) in enumerate(generator): + group_id = self._line_groups_ids[num_line] + is_written = False + for fold_storage in fold_storages: + if fold_storage.contains_group_id(group_id): + fold_storage.add(line) + is_written = True + if not is_written: + rest_fold_file.add(line) + rest_size += 1 + if rest_size >= self._REST_SIZE: + rest_folds.append(rest_fold_file) + rest_fold_file.close() + rest_fold_file = self.create_fold(None, 'offset{}_rest'.format(offset), num) + rest_fold_file.open() + rest_size = 0 + num += 1 + + if rest_size > 0: + rest_fold_file.close() + rest_folds.append(rest_fold_file) + elif rest_fold_file.is_opened(): + rest_fold_file.close() + finally: + for fold_storage in fold_storages: + fold_storage.close() + + return rest_folds + + def create_fold_sets(self, fold_size, folds_count): + """Create all folds for all permutations.""" + folds = [] + passed_folds_count = 0 + + while passed_folds_count < folds_count: + folds.append(self._make_learn_folds(fold_size, folds_count - passed_folds_count)) + current_learn_folds = folds[-1] + passed_folds_count += len(current_learn_folds) + return folds + + def fold_groups_files_generator(self, folds_groups, fold_offset): + """Create folds storages for all folds in folds_groups. Generator.""" + + fold_num = 0 + for fold_group in folds_groups: + learn_folds = [] + skipped_folds = [] + for learn_set in fold_group: + fold_num += 1 + if fold_offset < fold_num: + fold_file = self.create_fold(learn_set, 'fold', fold_num) + learn_folds.append(fold_file) + elif fold_offset >= fold_num: + fold_file = self.create_fold(learn_set, 'offset{}_skipped'.format(fold_offset), fold_num) + skipped_folds.append(fold_file) + + rest_folds = self._write_folds(learn_folds + skipped_folds, fold_num, fold_offset) + yield learn_folds, skipped_folds, rest_folds + + def create_fold(self, fold_set, name, id): + file_name = self.create_name_from_id(name, id) + fold_file = _FoldFile(fold_set, + file_name, + sep=self._line_reader.get_separator(), + column_description=self._column_description) + self._folds_storage.add(fold_file) + return fold_file + + def clean_folds(self): + for file in self._folds_storage: + file.delete() + + def clean(self): + FoldStorage.remove_dir() + + @staticmethod + def create_name_from_id(name, id, offset=None, max_count_digits=4): + if offset is not None: + name = '{name}{:0>{max_count_digits}}_offset{offset}'.format( + id, + name=name, + max_count_digits=max_count_digits, + offset=offset + ) + else: + name = '{name}{:0>{max_count_digits}}'.format(id, name=name, max_count_digits=max_count_digits) + return name diff --git a/v15cat/pylib/catboost/eval/catboost_evaluation.py b/v15cat/pylib/catboost/eval/catboost_evaluation.py new file mode 100644 index 0000000000000000000000000000000000000000..f694f0e6558d2b14a5e576d0d07aa3a1b6206238 --- /dev/null +++ b/v15cat/pylib/catboost/eval/catboost_evaluation.py @@ -0,0 +1,295 @@ +from __future__ import print_function + +import os +from copy import copy +from enum import Enum + +from .. import CatBoostError, CatBoost +from .evaluation_result import EvaluationResults, MetricEvaluationResult +from ._fold_models_handler import FoldModelsHandler +from ._readers import _SimpleStreamingFileReader +from ._splitter import _Splitter +from .execution_case import ExecutionCase +from .factor_utils import LabelMode, FactorUtils + + +class EvalType(Enum): + """ + Type of feature evaluation: + All: All factors presented + SeqRem: Each factor while other presented + SeqAdd: Each factor while other removed + SeqAddAndAll: SeqAdd + All + """ + All = 'All' + SeqRem = 'SeqRem' + SeqAdd = 'SeqAdd' + SeqAddAndAll = 'SeqAddAndAll' + + +class CatboostEvaluation(object): + + def __init__(self, + path_to_dataset, + fold_size, + fold_count, + column_description, + fold_offset=0, + group_column=None, + working_dir=None, + remove_models=True, + delimiter='\t', + has_header=False, + partition_random_seed=0, + min_fold_count=1): + """ + Args: + :param path_to_dataset: (str) Path to the dataset to be used for evaluation. + :param fold_size: (int) Size of the folds in cross-validation. + :param fold_count: (int) Number of times we get a new fold, learn a model and check results as if + there wouldn't be any offset. If there'are some offset it means that the real count of folds will + be smaller. + :param column_description: (str) Path to the file where column description is placed. + :param fold_offset: (int) Number of folds we skip before begin to make cross-validation. + :param group_column: (int) GroupId column index in the dataset file. + 'None' value means absence of grouping information in the dataset (it's the default). + :param working_dir: Working dir for temporary files + :param remove_models: (bool) Set true if you want models to be removed after applying them. + :param delimiter: (str) Field delimiter used in dataset files. + :param has_header: (bool) Set true if you want to skip first line in dataset files. + :param partition_random_seed: (int) The seed for random value generator used for getting permutations for + cross-validation. + :param min_fold_count: (int) Minimun amount of folds dataset can be cut to. + """ + import os.path + + self._current_dir = os.getcwd() + self._path_to_dataset = os.path.join(self._current_dir, path_to_dataset) + self._column_description = os.path.join(self._current_dir, + column_description) if column_description is not None else None + self._fold_offset = fold_offset + self._fold_count = fold_count + self._fold_size = fold_size + self._delimiter = delimiter + self._has_header = has_header + self._seed = partition_random_seed + self._min_fold_count = int(min_fold_count) + self._remove_models = remove_models + + if group_column is not None: + self._group_feature_num = int(group_column) + else: + self._group_feature_num = group_column + + if working_dir is None: + import tempfile + self._working_dir = tempfile.mkdtemp() + else: + self._working_dir = working_dir + + def __go_to_working_dir(self): + current = os.getcwd() + os.chdir(self._working_dir) + return current + + @staticmethod + def _create_eval_feature_cases(params, features_to_eval, eval_type, label_mode): + if len(features_to_eval) == 0: + raise CatBoostError("Provide at least one feature to evaluation") + + # baseline + test_cases = list() + baseline_case = ExecutionCase(params, + ignored_features=list(features_to_eval), + label=FactorUtils.create_label(features_to_eval, + features_to_eval, + label_mode=label_mode) + ) + # test + if eval_type == EvalType.All or eval_type == EvalType.SeqAddAndAll or len(features_to_eval) == 1: + test_cases.append(ExecutionCase(params, + ignored_features=[], + label=FactorUtils.create_label(features_to_eval, + [], + label_mode=label_mode) + )) + elif eval_type == EvalType.SeqRem: + for feature_num in features_to_eval: + test_cases.append(ExecutionCase(params, + ignored_features=[feature_num], + label=FactorUtils.create_label(features_to_eval, + [feature_num], + label_mode=label_mode) + )) + elif eval_type == EvalType.SeqAdd or eval_type == EvalType.SeqAddAndAll: + for feature_num in features_to_eval: + cur_features = copy(features_to_eval) + cur_features.remove(feature_num) + test_cases.append(ExecutionCase(params, + label=FactorUtils.create_label(features_to_eval, + cur_features, + label_mode=label_mode), + ignored_features=list(cur_features))) + elif eval_type != EvalType.All: + raise AttributeError("Don't support {} mode.", eval_type.value) + return baseline_case, test_cases + + @staticmethod + def _create_evaluation_results(by_case_results): + group_by_metric = dict() + + for (case, case_result) in by_case_results.items(): + for (metric, evaluation_result) in case_result.items(): + if metric not in group_by_metric: + group_by_metric[metric] = list() + group_by_metric[metric].append(evaluation_result) + + results = list() + for (metric, metric_results) in group_by_metric.items(): + results.append(MetricEvaluationResult(metric_results)) + return EvaluationResults(results) + + def get_working_dir(self): + return self._working_dir + + def _calculate_result_metrics(self, cases, metrics, thread_count=-1, evaluation_step=1): + """ + This method calculate metrics and return them. + + Args: + :param cases: List of the ExecutionCases you want to evaluate + :param metrics: List of the metrics to be computed + :param thread_count: Count of threads to use. + :param: evaluation_step: Step to evaluate metrics + :return: instance of EvaluationResult + """ + cases_set = set(cases) + if len(cases_set) != len(cases): + raise CatBoostError("Found duplicate cases in " + cases) + current_wd = self.__go_to_working_dir() + try: + if self._fold_count <= self._fold_offset: + error_msg = 'Count of folds(folds_count - offset) need to be at least one: offset {}, folds_count {}.' + raise AttributeError(error_msg.format(self._fold_offset, + self._fold_count)) + + handler = FoldModelsHandler(cases=cases, + metrics=metrics, + eval_step=evaluation_step, + thread_count=thread_count, + remove_models=self._remove_models) + + reader = _SimpleStreamingFileReader(self._path_to_dataset, + sep=self._delimiter, + has_header=self._has_header, + group_feature_num=self._group_feature_num) + splitter = _Splitter(reader, + self._column_description, + seed=self._seed, + min_folds_count=self._min_fold_count) + + result = handler.proceed(splitter=splitter, + fold_size=self._fold_size, + folds_count=self._fold_count, + fold_offset=self._fold_offset) + + return self._create_evaluation_results(result) + finally: + os.chdir(current_wd) + + def eval_features(self, + learn_config, + features_to_eval, + loss_function=None, + eval_type=EvalType.SeqAdd, + eval_metrics=None, + thread_count=-1, + eval_step=None, + label_mode=LabelMode.AddFeature): + """ Evaluate features. + Args: + learn_config: dict with params or instance of CatBoost. In second case instance params will be used + features_to_eval: list of indices of features to evaluate + loss_function: one of CatBoost loss functions, get it from learn_config if not specified + eval_type: Type of feature evaluate (All, SeqAdd, SeqRem) + eval_metrics: Additional metrics to calculate + thread_count: thread_count to use. If not none will override learn_config values + Returns + ------- + result : Instance of EvaluationResult class + """ + features_to_eval = set(features_to_eval) + if eval_metrics is None: + eval_metrics = [] + eval_metrics = eval_metrics if isinstance(eval_metrics, list) else [eval_metrics] + if isinstance(learn_config, CatBoost): + params = learn_config.get_params() + else: + params = dict(learn_config) + + if loss_function is not None: + if "loss_function" in params and params["loss_function"] != loss_function: + raise CatBoostError("Loss function in params {} should be equal to feature evaluation objective " + "function {}".format(params["loss_function"], loss_function)) + else: + if "loss_function" not in params: + raise CatBoostError("Provide loss function in params or as option to eval_features method") + + if thread_count is not None and thread_count != -1: + params["thread_count"] = thread_count + + if eval_step is None: + eval_step = 1 + + if loss_function is not None: + params["loss_function"] = loss_function + else: + loss_function = params["loss_function"] + + if params["loss_function"] == "PairLogit": + raise CatBoostError("Pair classification is not supported") + + baseline_case, test_cases = self._create_eval_feature_cases(params, + features_to_eval, + eval_type=eval_type, + label_mode=label_mode) + if loss_function not in eval_metrics: + eval_metrics.append(loss_function) + + return self.eval_cases(baseline_case=baseline_case, + compare_cases=test_cases, + eval_metrics=eval_metrics, + thread_count=thread_count, + eval_step=eval_step) + + def eval_cases(self, + baseline_case, + compare_cases, + eval_metrics, + thread_count=-1, + eval_step=1): + """More flexible evaluation of any cases. + Args: + baseline_case: Execution case used for baseline + compare_cases: List of cases to compare + eval_metrics: Metrics to calculate + thread_count: thread_count to use. Will override one in cases + Returns + ------- + result : Instance of EvaluationResult class + """ + if not isinstance(compare_cases, list): + compare_cases = [compare_cases] + + cases = [baseline_case] + cases += compare_cases + + for case in cases: + case._set_thread_count(thread_count) + + metric_result = self._calculate_result_metrics(cases, + eval_metrics, + thread_count=thread_count, + evaluation_step=eval_step) + metric_result.set_baseline_case(baseline_case) + return metric_result diff --git a/v15cat/pylib/catboost/eval/evaluation_result.py b/v15cat/pylib/catboost/eval/evaluation_result.py new file mode 100644 index 0000000000000000000000000000000000000000..5ddbd604b459174bfda60dc9728f0c6932861eae --- /dev/null +++ b/v15cat/pylib/catboost/eval/evaluation_result.py @@ -0,0 +1,518 @@ +import numpy as np +import pandas as pd +from enum import Enum + +from .. import CatBoostError +from ..core import metric_description_or_str_to_str +from ..utils import compute_wx_test + + +def calc_wilcoxon_test(baseline, test): + return compute_wx_test(baseline, test)["pvalue"] + + +class ScoreType(Enum): + Abs = "AbsoluteDiff" + Rel = "RelativeDiff" + + +class ScoreConfig: + """ + Config to present human-friendly evaluation results. + """ + + def __init__(self, + score_type=ScoreType.Rel, + multiplier=100, + score_level=0.01, + interval_level=0.01, + overfit_iterations_info=True): + """ + + :param score_type: type of score. For abs difference score will be (baseline - test).mean(), + for relative it's ((baseline - test) / baseline).mean() + :param multiplier: multiplier to print score + :param score_level: WX-test level. Will be used to make if tested case significantly better or worse + :param interval_level: level to compute score confidence interval + :param overfit_iterations_info: if information about overfit iterations should be preserved + """ + self.type = score_type + self.multiplier = multiplier + self.score_level = score_level + self.interval_level = interval_level + self.overfit_overfit_iterations_info = overfit_iterations_info + + @staticmethod + def abs_score(level=0.01): + return ScoreConfig(score_type=ScoreType.Abs, + multiplier=1, + score_level=level) + + @staticmethod + def rel_score(level=0.01): + return ScoreConfig(score_type=ScoreType.Rel, + multiplier=100, + score_level=level) + + +def calc_bootstrap_ci_for_mean(samples, level=0.05, tries=999): + """ + Count confidence intervals for difference each two samples. + + Args: + :param samples: samples + :param level: (float) Level for the confidence interval. + :param tries: bootstrap samples to use + :return: (left, right) border of confidence interval + + """ + + if not (samples == 0).all(): + samples = np.array(samples) + means = [] + for _ in range(0, tries): + resample = np.random.choice(samples, len(samples)) + means.append(np.mean(resample)) + means = sorted(means) + left = means[int(tries * (level / 2))] + right = means[int(tries * (1.0 - level / 2))] + return left, right + else: + return 0, 0 + + +class CaseEvaluationResult: + """ + CaseEvaluationResults stores aggregated statistics for one EvaluationCase and one metric. + """ + + def __init__(self, case, metric_description, eval_step): + self._case = case + self._metric_description = metric_description + + self._fold_metric = pd.Series() + self._fold_metric_iteration = pd.Series() + self._fold_curves = dict() + + self._eval_step = eval_step + + def _add(self, model, learning_curve): + if model.get_case() != self._case: + raise CatBoostError("Model case should be equal to result case") + + fold_id = model.get_fold_id() + + self._fold_curves[fold_id] = learning_curve + score = max(learning_curve) if self._metric_description.is_max_optimal() else min(learning_curve) + position = np.argmax(learning_curve) if self._metric_description.is_max_optimal() else np.argmin( + learning_curve) + + self._fold_metric.at[fold_id] = score + self._fold_metric_iteration.at[fold_id] = position + + def __eq__(self, other): + return (np.all(self._fold_metric == other._fold_metric) + and np.all(self._fold_metric_iteration == other._fold_metric_iteration) + and self._fold_curves == other._fold_curves) + + def get_case(self): + """ + ExecutionCases for this result + """ + return self._case + + def get_fold_ids(self): + """ + + :return: FoldsIds for which this caseResult was calculated + """ + return self._fold_curves.keys() + + def get_best_metric_for_fold(self, fold): + """ + + :param fold: id of fold to get result + :return: best metric value, best metric iteration + """ + return self._fold_metric[fold], self._fold_metric_iteration[fold] + + def get_best_iterations(self): + """ + + :return: pandas Series with best iterations on all folds + """ + return self._fold_metric_iteration + + def get_best_metrics(self): + """ + + :return: pandas series with best metric values + """ + return self._fold_metric + + def get_fold_curve(self, fold): + """ + + :param fold: + :return: fold learning curve (test scores on every eval_period iteration) + """ + return self._fold_curves[fold] + + def get_metric_description(self): + """ + + :return: Metric used to build this CaseEvaluationResult + """ + return self._metric_description + + def get_eval_step(self): + """ + + :return: step which was used for metric computations + """ + return self._eval_step + + def count_under_and_over_fits(self, overfit_border=0.15, underfit_border=0.95): + """ + + :param overfit_border: min fraction of iterations until overfitting starts one expects all models to have + :param underfit_border: border, after which there should be no best_metric_scores + :return: #models with best_metric > underfit_border * iter_count, #models, with best_metric > overfit_border + """ + count_overfitting = 0 + count_underfitting = 0 + + for fold_id, fold_curve in self._fold_curves.items(): + best_score_position = self._fold_metric_iteration[fold_id] + best_model_size_fraction = best_score_position * 1.0 / len(fold_curve) + if best_model_size_fraction > overfit_border: + count_underfitting += 1 + elif best_model_size_fraction < underfit_border: + count_overfitting += 1 + return count_overfitting, count_underfitting + + def estimate_fit_quality(self): + """ + + :return: Simple sanity check that all models overfit and not too fast + """ + count_overfitting, count_underfitting = self.count_under_and_over_fits() + if count_overfitting > count_underfitting: + return "Overfitting" + if count_underfitting > count_overfitting: + return "Underfitting" + return "Good" + + def create_learning_curves_plot(self, offset=None): + """ + + :param offset: First iteration to plot + :return: plotly Figure with learning curves for each fold + """ + import plotly.graph_objs as go + + traces = [] + + for fold in self.get_fold_ids(): + scores_curve = self.get_fold_curve(fold) + if offset is not None: + first_idx = offset + else: + first_idx = int(len(scores_curve) * 0.1) + + traces.append(go.Scatter(x=[i * int(self._eval_step) for i in range(first_idx, len(scores_curve))], + y=scores_curve[first_idx:], + mode='lines', + name='Fold #{}'.format(fold))) + + layout = go.Layout( + title='Learning curves for case {}'.format(self._case), + hovermode='closest', + xaxis=dict( + title='Iteration', + ticklen=5, + zeroline=False, + gridwidth=2, + ), + yaxis=dict( + title='Metric', + ticklen=5, + gridwidth=2, + ), + showlegend=True + ) + fig = go.Figure(data=traces, layout=layout) + return fig + + +class MetricEvaluationResult: + """ + Evaluation result for one metric. + Stores all ExecutionCases with specified metric scores + Computes human-friendly tables with results and some plots + """ + + def __init__(self, case_results): + if len(case_results) <= 1: + raise CatBoostError("Need at least 2 case results, got {} ".format(len(case_results))) + + self._case_results = dict() + self._case_comparisons = dict() + self._cases = [case_result.get_case() for case_result in case_results] + + for case_result in case_results: + case = case_result.get_case() + self._case_results[case] = case_result + + self._metric_description = case_results[0].get_metric_description() + self._baseline_case = case_results[0].get_case() + + self._score_config = ScoreConfig() + + for (case, case_result) in self._case_results.items(): + if case_result.get_metric_description() != self._metric_description: + raise CatBoostError("Metric names should be equal for all case results") + + if case_result.get_fold_ids() != self.get_fold_ids(): + raise CatBoostError("Case results should be computed on the same folds") + + if case_result.get_eval_step() != self.get_eval_step(): + raise CatBoostError("Eval steps should be equal for different cases") + + def __clear_comparisons(self): + self._case_comparisons = dict() + + def _change_score_config(self, config): + if config is not None: + if isinstance(config, ScoreType): + if config == ScoreType.Abs: + config = ScoreConfig.abs_score() + elif config == ScoreType.Rel: + config = ScoreConfig.rel_score() + else: + raise CatBoostError("Unknown scoreType {}".format(config)) + if self._score_config != config: + self._score_config = config + self.__clear_comparisons() + + def _compute_case_result_table(self, baseline_case): + result = pd.DataFrame() + baseline_scores = self._case_results[baseline_case].get_best_metrics() + baseline_iters = self._case_results[baseline_case].get_best_iterations() + + for (case, case_result) in self._case_results.items(): + if case != baseline_case: + test_scores = case_result.get_best_metrics() + pvalue = calc_wilcoxon_test(baseline_scores, test_scores) + + diff = (baseline_scores - test_scores) + if self._score_config.type == ScoreType.Rel: + diff = diff / baseline_scores.abs() + if self._metric_description.is_max_optimal(): + diff = -diff + mean_diff = diff.mean() + + left_quantile, right_quantile = calc_bootstrap_ci_for_mean(diff, + self._score_config.interval_level) + + case_name = str(case) + + result.at[case_name, "PValue"] = pvalue + result.at[case_name, "Score"] = mean_diff * self._score_config.multiplier + + left_quantile_title = "Quantile {}".format(self._score_config.score_level / 2) + right_quantile_title = "Quantile {}".format(1.0 - self._score_config.score_level / 2) + + result.at[case_name, left_quantile_title] = left_quantile * self._score_config.multiplier + result.at[case_name, right_quantile_title] = right_quantile * self._score_config.multiplier + + decision = "UNKNOWN" + if pvalue < self._score_config.score_level: + if mean_diff > 0: + decision = "GOOD" + elif mean_diff < 0: + decision = "BAD" + result.at[case_name, "Decision"] = decision + + if self._score_config.overfit_overfit_iterations_info: + test_iters = case_result.get_best_iterations() + pvalue = calc_wilcoxon_test(baseline_iters, test_iters) + + result.at[case_name, "Overfit iter diff"] = (test_iters - baseline_iters).mean() + result.at[case_name, "Overfit iter pValue"] = pvalue + + return result.sort_values(by=["Score"], ascending=self._metric_description.is_max_optimal()) + + def get_baseline_case(self): + """ + + :return: ExecutionCases used as a baseline (with everything else is compared) + """ + return self._baseline_case + + def get_cases(self): + """ + + :return: Cases which are compared + """ + return self._cases + + def get_metric_description(self): + """ + + :return: Metric for which results were calculated + """ + return self._metric_description + + def get_baseline_comparison(self, score_config=None): + """ + Method to get human-friendly table with model comparisons. + + Returns baseline vs all other computed cases result + :param score_config: Config to present human-friendly score, optional. Instance of ScoreConfig + :return: pandas DataFrame. Each row is related to one ExecutionCase. + Each row describes how better (or worse) this case is compared to baseline. + """ + case = self._baseline_case + return self.get_case_comparison(case, score_config) + + def get_case_comparison(self, case, score_config=None): + """ + Method to get human-friendly table with model comparisons. + Same as get_baseline_comparison(), but with other non-baseline case specified as baseline + + :param case: use specified case as baseline + :param score_config: + :return: pandas DataFrame. Each row is related to one ExecutionCase. + Each row describes how better (or worse) this case is compared to baseline. + """ + self._change_score_config(score_config) + if case not in self._case_comparisons: + self._case_comparisons[case] = self._compute_case_result_table(case) + return self._case_comparisons[case] + + def change_baseline_case(self, case): + """ + + :param case: new baseline case + :return: + """ + if case not in self._case_results: + raise CatBoostError("Case {} is unknown. Can't use it as baseline".format(case)) + self._baseline_case = case + + def get_case_result(self, case): + """ + + :param case: + :return: CaseEvaluationResult. Scores and other information about single execution case + """ + return self._case_results[case] + + def get_fold_ids(self): + """ + + :return: Folds ids which we used for computing this evaluation result + """ + return self._case_results[self._baseline_case].get_fold_ids() + + def get_eval_step(self): + return self._case_results[self._baseline_case].get_eval_step() + + def create_fold_learning_curves(self, fold, offset=None): + """ + + :param fold: FoldId to plot + :param offset: first iteration to plot + :return: plotly figure for all cases on specified fold + """ + import plotly.graph_objs as go + traces = [] + for case in self.get_cases(): + case_result = self.get_case_result(case) + scores_curve = case_result.get_fold_curve(fold) + if offset is not None: + first_idx = offset + else: + first_idx = int(len(scores_curve) * 0.1) + + traces.append( + go.Scatter(x=[i * int(case_result.get_eval_step()) for i in range(first_idx, len(scores_curve))], + y=scores_curve[first_idx:], + mode='lines', + name='Case {}'.format(case))) + + layout = go.Layout( + title='Learning curves for metric {} on fold #{}'.format(self._metric_description, fold), + hovermode='closest', + xaxis=dict( + title='Iteration', + ticklen=5, + zeroline=False, + gridwidth=2, + ), + yaxis=dict( + title='Metric', + ticklen=5, + gridwidth=2, + ), + showlegend=True + ) + fig = go.Figure(data=traces, layout=layout) + return fig + + def __eq__(self, other): + return (self._case_results == other._case_results + and self._case_comparisons == other._case_comparisons + and self._cases == other._cases) + + +class EvaluationResults: + + def __init__(self, metric_results): + if len(metric_results) < 1: + raise CatBoostError("Need at least one result") + + self._results = dict() + self._metrics = dict() + self._cases = None + + for result in metric_results: + metric_description = result.get_metric_description() + if metric_description in self._results: + raise CatBoostError("Duplicate metric {}".format(metric_description)) + if self._cases is None: + self._cases = result.get_cases() + key = metric_description_or_str_to_str(metric_description) + self._results[key] = result + self._metrics[key] = metric_description + + def get_metric_results(self, metric): + """ + + :param metric: + :return: MetricEvaluationResult for specified metric + """ + return self._results[metric_description_or_str_to_str(metric)] + + def get_metrics(self): + """ + + :return: Names of the metrics which were computed + """ + return self._metrics + + def get_results(self): + """ + + :return: Results are the map from metric names to computed results (instance of MetricEvaluationResult) + on this fold + """ + return self._results + + def set_baseline_case(self, case): + """ + Could be used to change baseline cases for already computed results + """ + for (metric, metric_result) in self._results.items(): + metric_result.change_baseline_case(case) diff --git a/v15cat/pylib/catboost/eval/execution_case.py b/v15cat/pylib/catboost/eval/execution_case.py new file mode 100644 index 0000000000000000000000000000000000000000..ab7d86751da0078bdf82575e2cc0cd65510d161c --- /dev/null +++ b/v15cat/pylib/catboost/eval/execution_case.py @@ -0,0 +1,76 @@ +import json + +from .. import CatBoostError +from ..eval.factor_utils import FactorUtils +from ..core import _NumpyAwareEncoder + + +class ExecutionCase: + + def __init__(self, + params, + label=None, + ignored_features=None, + learning_rate=None): + """ + Instances of this class are cases which will be compared during evaluation + Params are CatBoost params + label is a string which will be used for plots and other visualisations + ignored_features is a set of additional feature indices to ignore + """ + case_params = dict(params) + + if learning_rate is not None: + case_params["learning_rate"] = learning_rate + + all_ignored_features = set() + if "ignored_features" in case_params: + all_ignored_features.update(set(case_params["ignored_features"])) + if ignored_features is not None: + all_ignored_features.update(ignored_features) + + case_params["ignored_features"] = list(all_ignored_features) + + self._label = label if label is not None else "" + self._ignored_features = ignored_features + self._ignored_features_str = FactorUtils.factors_to_ranges_string(self._ignored_features) + + self.__set_params(case_params) + + def __set_params(self, params): + self._params = params + self._params_hash = hash(json.dumps(self._params, sort_keys=True, cls=_NumpyAwareEncoder)) + + def _set_thread_count(self, thread_count): + if thread_count is not None and thread_count != -1: + params = self._params + params["thread_count"] = thread_count + self.__set_params(params) + + @staticmethod + def _validate_ignored_features(ignored_features, eval_features): + for eval_feature in eval_features: + if eval_feature in ignored_features: + raise CatBoostError( + "Feature {} is in ignored set and in tmp-features set at the same time".format(eval_feature)) + + def get_params(self): + return dict(self._params) + + def get_label(self): + return self._label + + def __str__(self): + if len(self._label) == 0: + return "Ignore: {}".format(self._ignored_features_str) + else: + return '{}'.format(self._label) + + def __repr__(self): + return self.__str__() + + def __eq__(self, other): + return self._params == other._params and self._label == other._label + + def __hash__(self): + return hash((self._label, self._params_hash)) diff --git a/v15cat/pylib/catboost/eval/factor_utils.py b/v15cat/pylib/catboost/eval/factor_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..3d92a023abab3d4adf014622877781b6f353586c --- /dev/null +++ b/v15cat/pylib/catboost/eval/factor_utils.py @@ -0,0 +1,99 @@ +from enum import Enum + + +class LabelMode(Enum): + AddFeature = "Feature" + IgnoreFeature = "Ignore" + + +class FactorUtils(object): + + @staticmethod + def extract_factors(string_factors): + string_factors = string_factors.strip() + if not string_factors: + return {} + + list_factors = [] + eval_factor_units = string_factors.split(';') + for eval_factor_unit in eval_factor_units: + try: + factor = int(eval_factor_unit) + list_factors.append(factor) + except ValueError: + string_bounds = eval_factor_unit.split('-') + if len(string_bounds) != 2: + raise AttributeError('Range need to contain exactly two numbers!') + begin_range = int(string_bounds[0]) + end_range = int(string_bounds[1]) + list_factors += list(range(begin_range, end_range + 1)) + + return set(list_factors) + + @staticmethod + def factors_to_string(factors): + if len(factors) == 0: + return '' + parts = [] + factors_list = sorted(factors) + begin = factors_list[0] + for i in range(1, len(factors_list)): + if factors_list[i] != factors_list[i - 1] + 1: + end = factors_list[i - 1] + if begin != end: + parts.append('{}-{}'.format(begin, end)) + else: + parts.append(str(begin)) + begin = factors_list[i] + end = len(factors_list) - 1 + if begin != factors_list[end]: + parts.append('{}-{}'.format(begin, factors_list[end])) + else: + parts.append(str(begin)) + return ';'.join(parts) + + @staticmethod + def compress_string_factors(string_factors): + factors = FactorUtils.extract_factors(string_factors) + compressed_string_factors = FactorUtils.factors_to_string(factors) + return compressed_string_factors + + @staticmethod + def single_range_to_string(left, right): + if left != right: + return "{}-{}".format(left, right) + else: + return "{}".format(left) + + @staticmethod + def group_factors_by_range(factors_set): + factors = sorted(list(factors_set)) + is_start = [] + for i in range(0, len(factors)): + is_start.append(i == 0 or (factors[i] != factors[i - 1] + 1)) + + grouped_factors = [] + i = 0 + while i < len(factors): + if is_start[i]: + grouped_factors.append([]) + grouped_factors[-1].append(factors[i]) + i += 1 + return grouped_factors + + @staticmethod + def factors_to_ranges_string(factors_set): + if factors_set is None or len(factors_set) == 0: + return "None" + grouped_factors = FactorUtils.group_factors_by_range(factors_set) + + return ';'.join([FactorUtils.single_range_to_string(min(x), max(x)) for x in grouped_factors]) + + @staticmethod + def create_label(all_eval_features, removed_features, label_mode): + eval_features = set(all_eval_features) + if label_mode == LabelMode.AddFeature: + add_features = eval_features - set(removed_features) + return "Features: {}".format(FactorUtils.factors_to_ranges_string(add_features)) + else: + return "Ignore: {}".format(FactorUtils.factors_to_ranges_string(set(removed_features))) diff --git a/v15cat/pylib/catboost/eval/log_config.py b/v15cat/pylib/catboost/eval/log_config.py new file mode 100644 index 0000000000000000000000000000000000000000..90e9738aa2cf11d2ca905fcb87aeee7db5f8cdb4 --- /dev/null +++ b/v15cat/pylib/catboost/eval/log_config.py @@ -0,0 +1,37 @@ +import logging + +filename = None +eval_logger = None +console_handler = None + + +def init(): + global eval_logger + eval_logger_name = "eval_feature" + eval_logger = logging.getLogger(eval_logger_name) + eval_logger.setLevel(logging.CRITICAL) + global console_handler + console_handler = logging.StreamHandler() + console_handler.setLevel(logging.CRITICAL) + + formatter = logging.Formatter("[%(levelname)s]: %(message)s") + console_handler.setFormatter(formatter) + + eval_logger.addHandler(console_handler) + + +def set_logger_name(name): + global filename + filename = name + + +def get_eval_logger(): + if eval_logger is None: + init() + return eval_logger + + +def set_level(level): + if eval_logger is None: + init() + console_handler.setLevel(level) diff --git a/v15cat/pylib/catboost/eval/utils.py b/v15cat/pylib/catboost/eval/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..9d0ca3f71946a31c781e01550a3fec1ae23f7493 --- /dev/null +++ b/v15cat/pylib/catboost/eval/utils.py @@ -0,0 +1,25 @@ +""" +Module that contains different utils functions. +""" +import os + + +def make_dirs_if_not_exists(name): + try: + os.makedirs(name) + except OSError: + pass + + +def series_to_line(row, sep): + return sep.join(map(str, row.tolist())) + + +def save_plot(fig, file_name=None): + try: + from plotly.offline import iplot, init_notebook_mode + init_notebook_mode(connected=True) + iplot(fig, filename=file_name) + except Exception: + from plotly.offline import plot + plot(fig, auto_open=False, filename=file_name) diff --git a/v15cat/pylib/catboost/hnsw/__init__.py b/v15cat/pylib/catboost/hnsw/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..397114f78d243eedb1e6b74a88952895bee718cf --- /dev/null +++ b/v15cat/pylib/catboost/hnsw/__init__.py @@ -0,0 +1,2 @@ +from .hnsw import EDistance, EVectorComponentType, Pool, Hnsw, HnswException, transform_mobius, HnswEstimator, OnlineHnsw +__all__ = ['EDistance', 'EVectorComponentType', 'Pool', 'Hnsw', 'HnswException', 'transform_mobius', 'HnswEstimator', 'OnlineHnsw'] diff --git a/v15cat/pylib/catboost/hnsw/hnsw.py b/v15cat/pylib/catboost/hnsw/hnsw.py new file mode 100644 index 0000000000000000000000000000000000000000..7e7abe7781f8376bd2fce5dda04f86fb50f16cdf --- /dev/null +++ b/v15cat/pylib/catboost/hnsw/hnsw.py @@ -0,0 +1,666 @@ +import sys +import os +from six import iteritems +from enum import IntEnum +from contextlib import contextmanager +import json + +if sys.version_info >= (3,): + import importlib.util + + def load_dynamic(name, path): + spec = importlib.util.spec_from_file_location(name, path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module +else: + from imp import load_dynamic + + +def get_so_paths(dir_name): + dir_name = os.path.join(os.path.dirname(__file__), dir_name) + list_dir = os.listdir(dir_name) if os.path.isdir(dir_name) else [] + return [os.path.join(dir_name, so_name) for so_name in list_dir if so_name.split('.')[-1] in ['so', 'pyd']] + + +def get_hnsw_bin_module(): + if '_hnsw' in sys.modules: + return sys.modules['_hnsw'] + so_paths = get_so_paths('./') + for so_path in so_paths: + try: + loaded_hnsw = load_dynamic('_hnsw', so_path) + sys.modules['hnsw._hnsw'] = loaded_hnsw + return loaded_hnsw + except ImportError: + pass + from . import _hnsw + return _hnsw + + +@contextmanager +def log_fixup(): + _hnsw._set_logger(sys.stdout) + try: + yield + finally: + _hnsw._reset_logger() + + +class EDistance(IntEnum): + DotProduct = 0 + L1 = 1 + L2Sqr = 2 + PairVectorDistance = 3 + + +class EVectorComponentType(IntEnum): + Float = 0 + I8 = 1 + I32 = 2 + + +_hnsw = get_hnsw_bin_module() + +HnswException = _hnsw.HnswException + +_DenseVectorStorage = { + EVectorComponentType.Float: _hnsw._DenseFloatVectorStorage, + EVectorComponentType.I8: _hnsw._DenseI8VectorStorage, + EVectorComponentType.I32: _hnsw._DenseI32VectorStorage +} + +_HnswDenseVectorIndex = { + EVectorComponentType.Float: _hnsw._HnswDenseFloatVectorIndex, + EVectorComponentType.I8: _hnsw._HnswDenseI8VectorIndex, + EVectorComponentType.I32: _hnsw._HnswDenseI32VectorIndex +} + +_transform_mobius = { + EVectorComponentType.Float: _hnsw._transform_mobius_float, + EVectorComponentType.I8: _hnsw._transform_mobius_i8, + EVectorComponentType.I32: _hnsw._transform_mobius_i32 +} + +_OnlineHnswDenseVectorIndex = { + EVectorComponentType.Float: _hnsw._OnlineHnswDenseFloatVectorIndex, + EVectorComponentType.I8: _hnsw._OnlineHnswDenseI8VectorIndex, + EVectorComponentType.I32: _hnsw._OnlineHnswDenseI32VectorIndex, +} + + +class Pool: + """ + Pool is a storage of vectors + """ + + def __init__(self, vectors_path, dtype, dimension, vectors_bin_data=None): + """ + Pool is a storage of vectors. You can create it from row-major binary file or + binary data of vectors. + + Parameters + ---------- + vectors_path : string + Path to binary file with vectors. + + dtype : EVectorComponentType + Type of vectors. + + dimension : int + Dimension of vectors. + + vectors_bin_data : bytes + Binary data of vectors. + """ + self.vectors_path = vectors_path + self.dtype = dtype + self.dimension = dimension + assert (vectors_bin_data is None) ^ (vectors_path is None) + if vectors_path is not None: + self._storage = _DenseVectorStorage[dtype](vectors_path, dimension) + self._data = None + if vectors_bin_data is not None: + self._storage = _DenseVectorStorage[dtype](None, dimension, vectors_bin_data) + self._data = vectors_bin_data + + @classmethod + def from_file(cls, vectors_path, dtype, dimension): + """ + Create pool from binary file. + + Parameters + ---------- + vectors_path : string + Path to binary file with vectors. + + dtype : EVectorComponentType + Type of vectors. + + dimension : int + Dimension of vectors. + """ + return Pool(vectors_path, dtype, dimension, None) + + @classmethod + def from_bytes(cls, vectors_bin_data, dtype, dimension): + """ + Create pool from binary data. + + Parameters + ---------- + vectors_bin_data : bytes + Binary data of vectors. + + dtype : EVectorComponentType + Type of vectors. + + dimension : int + Dimension of vectors. + """ + return Pool(None, dtype, dimension, vectors_bin_data) + + def get_item(self, id): + """ + Get item from storage by id. + + Parameters + ---------- + id : int + Index of item in storage. + + Returns + ------- + item : numpy.ndarray + """ + return self._storage._get_item(id) + + def get_num_items(self): + """ + Get the number of items in storage. + + Returns + ------- + num_items : int + """ + return self._storage._get_num_items() + + +def transform_mobius(pool): + """ + Transform pool for fast dot product search on HNSW graph + https://papers.nips.cc/paper/9032-mobius-transformation-for-fast-inner-product-search-on-graph.pdf + + Parameters + ---------- + pool : Pool + + Returns + ------- + transformed_pool : Pool + """ + transformed_pool = Pool.from_bytes(bytes(0), EVectorComponentType.Float, pool.dimension) + transformed_pool._storage = _transform_mobius[pool.dtype](pool._storage) + return transformed_pool + + +class Hnsw: + """ + Class for building, loading and working with Hierarchical Navigable Small World index. + """ + def __init__(self): + """ + Create object for working with HNSW. + """ + self._index = None + self._data = None + + def build(self, pool, distance, max_neighbors=None, search_neighborhood_size=None, num_exact_candidates=None, + batch_size=None, upper_level_batch_size=None, level_size_decay=None, num_threads=None, verbose=False, + report_progress=True, snapshot_file=None, snapshot_interval=None): + """ + Build index with given options. + + Parameters + ---------- + pool : Pool + Pool of vectors for which index will be built. + + distance : EDistance + Distance that should be used for finding nearest vectors. + + max_neighbors : int (default=32) + Maximum number of neighbors that every item can be connected with. + + search_neighborhood_size : int (default=300) + Search neighborhood size for ANN-search. + Higher values improve search quality in expense of building time. + + num_exact_candidates : int (default=100) + Number of nearest vectors to take from batch. + Higher values improve search quality in expense of building time. + + batch_size : int (default=1000) + Number of items that added to graph on each step of algorithm. + + upper_level_batch_size : int (default=40000) + Batch size for building upper levels. + + level_size_decay : int (default=max_neighbors/2) + Base of exponent for decaying level sizes. + + num_threads : int (default=number of CPUs) + Number of threads for building index. + + report_progress : bool (default=True) + Print progress of building. + + verbose : bool (default=False) + Print additional information about time of building. + + snapshot_file : string (default=None) + Path for saving snapshots during the index building. + + snapshot_interval : int (default=600) + Interval between saving snapshots (seconds). + Snapshot is saved after building each level also. + """ + params = {} + not_params = ["not_params", "self", "params", "__class__", "pool", "distance"] + for key, value in iteritems(locals()): + if key not in not_params and value is not None: + params[key] = value + self._index = _HnswDenseVectorIndex[pool.dtype](pool._storage, distance) + with log_fixup(): + self._index._build(json.dumps(params)) + + def _check_index(self): + if self._index is None: + raise HnswException("Index is not built and not loaded") + + def save(self, index_path): + """ + Save index to file. + + Parameters + ---------- + index_path : string + Path to file for saving index. + """ + self._check_index() + self._index._save(index_path) + + def load(self, index_path, pool, distance): + """ + Load index from file. + + Parameters + ---------- + index_path : string + Path to file for loading index. + + pool : Pool + Pool of vectors for which index will be loaded. + + distance : EDistance + Distance that should be used for finding nearest vectors. + """ + self._index = _HnswDenseVectorIndex[pool.dtype](pool._storage, distance) + self._index._load(index_path) + self._data = None + + def load_from_bytes(self, index_data, pool, distance): + """ + Load index from bytes. + + Parameters + ---------- + index_data : bytes + Index binary data. + + pool : Pool + Pool of vectors for which index will be loaded. + + distance : EDistance + Distance that should be used for finding nearest vectors. + """ + self._index = _HnswDenseVectorIndex[pool.dtype](pool._storage, distance) + self._index._load_from_bytes(index_data) + self._data = index_data + + def get_nearest(self, query, top_size, search_neighborhood_size, distance_calc_limit=0): + """ + Get approximate nearest neighbors for query from index. + + Parameters + ---------- + query : list or numpy.ndarray + Vector for which nearest neighbors should be found. + + top_size : int + Required number of neighbors. + + search_neighborhood_size : int + Search neighborhood size for ANN-search. + Higher values improve search quality in expense of search time. + It should be equal or greater than top_size. + + distance_calc_limit : int (default=0) + Limit of distance calculation. + To guarantee satisfactory search time at the expense of quality. + 0 is equivalent to no limit. + + Returns + ------- + neighbors : list of tuples (id, distance) + """ + self._check_index() + return self._index._get_nearest(query, top_size, search_neighborhood_size, distance_calc_limit) + + +class HnswEstimator: + """ + Class for building, loading and working with Hierarchical Navigable Small World index with SciKit-Learn + Estimator compatible interface. + Mostly drop-in replacement for sklearn.neighbors.NearestNeighbors (except for some parameters) + """ + + def __init__(self, n_neighbors=5, + distance=EDistance.DotProduct, max_neighbors=32, search_neighborhood_size=300, + num_exact_candidates=100, batch_size=1000, upper_level_batch_size=40000, + level_size_decay=None): + """ + Parameters + ---------- + n_neighbors : int, default=5 + Number of neighbors to use by default for kneighbors queries. + + + distance : EDistance + Distance that should be used for finding nearest vectors. + + max_neighbors : int (default=32) + Maximum number of neighbors that every item can be connected with. + + search_neighborhood_size : int (default=300) + Search neighborhood size for ANN-search. + Higher values improve search quality in expense of building time. + + num_exact_candidates : int (default=100) + Number of nearest vectors to take from batch. + Higher values improve search quality in expense of building time. + + batch_size : int (default=1000) + Number of items that added to graph on each step of algorithm. + + upper_level_batch_size : int (default=40000) + Batch size for building upper levels. + + level_size_decay : int (default=max_neighbors/2) + Base of exponent for decaying level sizes. + """ + for key, value in iteritems(locals()): + if key not in ['self', '__class__']: + setattr(self, key, value) + + def _check_index(self): + if self._index is None: + raise HnswException("Index is not built and not loaded") + + def fit(self, X, y=None, num_threads=None, verbose=False, report_progress=True, snapshot_file=None, + snapshot_interval=600): + """ + Fit the HNSW model. + + Parameters + ---------- + X : array-like of shape (n_samples, n_values) + + y: None + Added to be compatible with Estimator API + + num_threads : int (default=number of CPUs) + Number of threads for building index. + + report_progress : bool (default=True) + Print progress of building. + + verbose : bool (default=False) + Print additional information about time of building. + + snapshot_file : string (default=None) + Path for saving snapshots during the index building. + + snapshot_interval : int (default=600) + Interval between saving snapshots (seconds). + + Returns + ------- + model : HnswEstimator + + """ + self._index, self._index_data = _hnsw._init_index(X, self.distance) + + params = self._get_params(return_none=False) + not_params = ["not_params", "self", "params", "__class__", "X", "y"] + for key, value in iteritems(locals()): + if key not in not_params and value is not None: + params[key] = value + del params['distance'] + + with log_fixup(): + self._index._build(json.dumps(params)) + return self + + def _get_params(self, return_none): + params = {} + for key, value in self.__dict__.items(): + if key[0] != '_' and (return_none or (value is not None)): + params[key] = value + return params + + def get_params(self, deep=True): + """ + Get parameters for this estimator. + """ + return self._get_params(return_none=True) + + def set_params(self, **params): + """ + Set the parameters of this estimator. + + Parameters + ---------- + **params : dict + HnswEstimator parameters. + + Returns + ------- + self : HnswEstimator instance + """ + if not params: + return self + valid_params = self._get_params(return_none=True) + + for key, value in params.items(): + if key not in valid_params: + raise HnswException( + 'Invalid parameter %s for HnswEstimator. ' + 'Check the list of available parameters ' + 'with `get_params().keys()`.' + ) + setattr(self, key, value) + + return self + + @property + def effective_metric_(self): + """ + Returns + ------- + Distance that should be used for finding nearest vectors. + """ + return self.distance + + @property + def n_samples_fit_(self): + """ + Returns + ------- + Number of samples in the fitted data. + """ + self._check_index() + return self._index_data.shape[0] + + def kneighbors(self, X=None, n_neighbors=None, return_distance=True, search_neighborhood_size=None, + distance_calc_limit=0): + """Finds the approximate K-neighbors of a point. + Returns indices of and distances to the neighbors of each point. + + Parameters + ---------- + X : array-like, shape (n_queries, n_features) or None + The query point or points. + If not provided, neighbors of each indexed point are returned. + In this case, the query point is not considered its own neighbor. + n_neighbors : int, default=None + Number of neighbors required for each sample. The default is the + value passed to the constructor. + return_distance : bool, default=True + Whether or not to return the distances. + + search_neighborhood_size : int, default=None + Search neighborhood size for ANN-search. + Higher values improve search quality in expense of search time. + It should be equal or greater than top_size. + If None set to n_neighbors * 2. + + distance_calc_limit : int (default=0) + Limit of distance calculation. + To guarantee satisfactory search time at the expense of quality. + 0 is equivalent to no limit. + + Returns + ------- + neigh_dist :numpy.ndarray of shape (n_queries, n_neighbors) + Array representing the lengths to points, only present if + return_distance=True + neigh_ind : numpy.ndarray of shape (n_queries, n_neighbors) + Indices of the nearest points in the population matrix. + """ + self._check_index() + + if X is None: + X = self._index_data + if n_neighbors is None: + n_neighbors = self.n_neighbors + if search_neighborhood_size is None: + search_neighborhood_size = n_neighbors * 2 + + return self._index._kneighbors(X, n_neighbors, return_distance, self.distance, search_neighborhood_size, + distance_calc_limit) + + +class OnlineHnsw: + """ + Class for building and working with Online Hierarchical Navigable Small World index. + """ + def __init__(self, dtype, dimension, distance, max_neighbors=None, search_neighborhood_size=None, num_vertices=None, level_size_decay=None): + """ + Create object with given options. + + Parameters + ---------- + dtype : EVectorComponentType + Type of vectors. + dimension : int + Dimension of vectors. + distance : EDistance + Distance that should be used for finding nearest vectors. + max_neighbors : int (default=32) + Maximum number of neighbors that every item can be connected with. + search_neighborhood_size : int (default=300) + Search neighborhood size for ANN-search. + Higher values improve search quality in expense of building time. + num_vertices : int (default=0) + Expected number of vectors in storage. + level_size_decay : int (default=max_neighbors/2) + Base of exponent for decaying level sizes. + """ + self.dtype = dtype + self.dimension = dimension + params = {} + all_params = ["max_neighbors", "search_neighborhood_size", "num_vertices", "level_size_decay"] + for key, value in iteritems(locals()): + if key in all_params and value is not None: + params[key] = value + self._online_index = _OnlineHnswDenseVectorIndex[dtype](dimension, distance, json.dumps(params)) + + def get_nearest_and_add_item(self, query): + """ + Get approximate nearest neighbors for query from index and add item to index + + Parameters + ---------- + query : list or numpy.ndarray + Vector for which nearest neighbors should be found. + Vector which should be added in index. + + Returns + ------- + neighbors : list of tuples (id, distance) with length = search_neighborhood_size + """ + return self._online_index._get_nearest_neighbors_and_add_item(query) + + def get_nearest(self, query, top_size=0): + """ + Get approximate nearest neighbors for query from index. + + Parameters + ---------- + query : list or numpy.ndarray + Vector for which nearest neighbors should be found. + top_size : int + Required number of neighbors. + + Returns + ------- + neighbors : list of tuples (id, distance) + """ + return self._online_index._get_nearest_neighbors(query, top_size) + + def add_item(self, item): + """ + Add item in index. + + Parameters + ---------- + item : list or numpy.ndarray + Vector which should be added in index. + """ + self._online_index._add_item(item) + + def get_item(self, id): + """ + Get item from storage by id. + + Parameters + ---------- + id : int + Index of item in storage. + + Returns + ------- + item : numpy.ndarray + """ + return self._online_index._get_item(id) + + def get_num_items(self): + """ + Get the number of items in storage. + + Returns + ------- + num_items : int + """ + return self._online_index._get_num_items() diff --git a/v15cat/pylib/catboost/metrics.py b/v15cat/pylib/catboost/metrics.py new file mode 100644 index 0000000000000000000000000000000000000000..673bae1fe279a353e6c10fc5fbacacb89f59e83d --- /dev/null +++ b/v15cat/pylib/catboost/metrics.py @@ -0,0 +1,296 @@ +from functools import partial + +import numpy as np +try: + import pandas as pd +except ImportError: + # just to avoid checking (pd is not None) everywhere + class pandas: + class DataFrame(object): + pass + + class Series(object): + pass + pd = pandas + +from . import _catboost + + +# copied from core.py to avoid circular import +_ARRAY_TYPES = (list, np.ndarray, pd.DataFrame, pd.Series) + +__all__ = [] + + +class BuiltinMetric(object): + @staticmethod + def params_with_defaults(): + """ + For each valid metric parameter, returns its default value and if this parameter is mandatory. + Implemented in child classes. + + Returns + ---------- + valid_params: dict: param_name -> {'default_value': default value or None, 'is_mandatory': bool} + """ + raise NotImplementedError('Should be overridden by the child class.') + + def __str__(self): + """ + Gets the representation of the metric object with overridden parameters. + Implemented in child classes. + + Returns + ---------- + metric_string: str representing the metric object. + """ + raise NotImplementedError('Should be overridden by the child class.') + + def set_hints(self, **hints): + """ + Sets hints for the metric. Hints are not validated. + Implemented in child classes. + + Returns + ---------- + self: for chained calls. + """ + raise NotImplementedError('Should be overridden by the child class.') + + def eval( + self, + label, + approx, + weight=None, + group_id=None, + group_weight=None, + subgroup_id=None, + pairs=None, + thread_count=-1 + ): + """ + Evaluate the metric with raw approxes and labels. + + Parameters + ---------- + label : list or numpy.ndarrays or pandas.DataFrame or pandas.Series + Object labels. + + approx : list or numpy.ndarrays or pandas.DataFrame or pandas.Series + Object approxes. + + weight : list or numpy.ndarray or pandas.DataFrame or pandas.Series, optional (default=None) + Object weights. + + group_id : list or numpy.ndarray or pandas.DataFrame or pandas.Series, optional (default=None) + Object group ids. + + group_weight : list or numpy.ndarray or pandas.DataFrame or pandas.Series, optional (default=None) + Group weights. + + subgroup_id : list or numpy.ndarray, optional (default=None) + subgroup id for each instance. + If not None, giving 1 dimensional array like data. + + pairs : list or numpy.ndarray or pandas.DataFrame or string or pathlib.Path + The pairs description. + If list or numpy.ndarrays or pandas.DataFrame, giving 2 dimensional. + The shape should be Nx2, where N is the pairs' count. The first element of the pair is + the index of winner object in the training set. The second element of the pair is + the index of loser object in the training set. + If string or pathlib.Path, giving the path to the file with pairs description. + + thread_count : int, optional (default=-1) + Number of threads to work with. + If -1, then the number of threads is set to the number of CPU cores. + + Returns + ------- + metric results : list with metric values. + """ + if len(label) > 0 and not isinstance(label[0], _ARRAY_TYPES): + label = [label] + if len(approx) == 0: + approx = [[]] + if not isinstance(approx[0], _ARRAY_TYPES): + approx = [approx] + return _catboost._eval_metric_util( + label, approx, str(self), weight, group_id, group_weight, subgroup_id, pairs, thread_count + ) + + def is_max_optimal(self): + """ + Returns + ---------- + bool : True if metric is maximizable, False otherwise + """ + return _catboost.is_maximizable_metric(str(self)) + + def is_min_optimal(self): + """ + Returns + ---------- + bool : True if metric is minimizable, False otherwise + """ + return _catboost.is_minimizable_metric(str(self)) + + +class _MetricGenerator(type): + def __new__(mcs, name, parents, attrs): + for k in attrs['_valid_params']: + attrs[k] = property( + partial(_get_param, name=k), + partial(_set_param, name=k), + partial(_del_param, name=k), + 'Parameter {} of metric {}'.format(k, name), + ) + + attrs['params_with_defaults'] = staticmethod(lambda: {param: {'default_value': default_value, + 'is_mandatory': attrs['_is_mandatory_param'][param]} + for param, default_value in attrs['_valid_params'].items()}) + # Set the serialization function. + docstring = ['Builtin metric: \'{}\''.format(name), 'Parameters:'] + if not attrs['_valid_params']: + docstring[-1] += ' none' + for param, value in attrs['_valid_params'].items(): + if not attrs['_is_mandatory_param'][param]: + docstring.append(' ' * 4 + '{} = {} (default value)'.format(param, repr(value))) + else: + docstring.append(' ' * 4 + '{} (mandatory)'.format(param)) + attrs['__doc__'] = '\n'.join(docstring) + attrs['__repr__'] = lambda self: '{}({})'.format( + self._underlying_metric_name, + ", ".join(['{}={} [mandatory={}]'.format(param, repr(value), self._is_mandatory_param[param]) + for param, value in _current_params(self, False).items()]), + ) + attrs['__str__'] = _to_string + + def set_hints(self, **hints): + for hint_key, hint_value in hints.items(): + if isinstance(hint_value, bool): + hints[hint_key] = str(hint_value).lower() + setattr(self, 'hints', '|'.join(['{}~{}'.format(hint_key, hint_value) for hint_key, hint_value in hints.items()])) + if 'hints' not in self._params: + self._params.append('hints') + return self + attrs['set_hints'] = set_hints + + cls = super(_MetricGenerator, mcs).__new__(mcs, name, parents, attrs) + return cls + + def __call__(cls, **kwargs): + metric_obj = cls.__new__(cls) + params = {k: v for k, v in cls._valid_params.items()} + param_is_set = {param: not mandatory for param, mandatory in cls._is_mandatory_param.items()} + + # Overwrite default parameters and check that all passed parameters are valid. + for param, value in kwargs.items(): + if param not in cls._valid_params: + raise ValueError('Unexpected parameter {}'.format(param)) + params[param] = value + param_is_set[param] = True + + # Check that no parameters are left unset. + for param, is_set in param_is_set.items(): + if not is_set: + raise ValueError('Parameter {} is mandatory and must be specified.'.format(param)) + + for param, value in params.items(): + _set_param(metric_obj, value, param) + metric_obj._params = list(params.keys()) + + return metric_obj + + def __repr__(cls): + return cls.__doc__ + + def __setattr__(cls, name, value): + # Protect property fields from being mutated. + if name in ('_valid_params', '_is_mandatory_param'): + raise ValueError('Metric\'s `{}` shouldn\'t be modified or deleted.'.format(name)) + type.__setattr__(cls, name, value) + + def __delattr__(cls, name): + # Protect property fields from being mutated. + if name in ('_valid_params', '_is_mandatory_param'): + raise ValueError('Metric\'s `{}` shouldn\'t be modified or deleted.'.format(name)) + type.__delattr__(cls, name) + + +def _get_param(metric_obj, name): + if name not in metric_obj._valid_params: + raise ValueError('Metric {} doesn\'t have a parameter {}.'.format(metric_obj.__name__, name)) + return getattr(metric_obj, '_'+name) + + +def _set_param(metric_obj, value, name): + """Validate a new parameter value in a created metric object.""" + if name not in metric_obj._valid_params: + raise ValueError('Metric {} doesn\'t have a parameter {}.'.format(metric_obj.__name__, name)) + setattr(metric_obj, '_' + name, value) + + +def _del_param(metric_obj, name): + """Validate a new parameter value in a created metric object.""" + if name not in metric_obj._valid_params: + raise ValueError('Metric {} doesn\'t have a parameter {}.'.format(metric_obj.__name__, name)) + if metric_obj._is_mandatory_param[name]: + raise ValueError('Parameter {} is mandatory, cannot reset.'.format(name)) + value = metric_obj._valid_params[name] + setattr(metric_obj, '_' + name, value) + + +def _current_params(metric_obj, override_only): + params_with_defaults = metric_obj.params_with_defaults() + param_info = {} + for param in sorted(metric_obj._params): + value = getattr(metric_obj, param) # current value + if param == 'hints' and value == '': + # Skip unset hints. + continue + if override_only: + # Skip reporting parameters which are set to their default values. + default_value = params_with_defaults[param]['default_value'] + if (default_value is None and value is None) or (default_value is not None and default_value == value): + continue + param_info[param] = value + return param_info + + +def _to_string(metric_obj): + def _param_to_string(dict_item): + param, value = dict_item + if param == 'misclass_cost_matrix': + str_value = '/'.join(map(str, value.flatten())) + else: + str_value = str(value) + return '{}={}'.format(param, str_value) + + param_info = _current_params(metric_obj, True) + # E.g. AUC for both AUC and AUCMulticlass: + underlying_name = metric_obj._underlying_metric_name + if len(param_info) == 0: + return underlying_name + return '{}:{}'.format(underlying_name, ';'.join(map(_param_to_string, param_info.items()))) + + +def _generate_metric_classes(): + for metric_name, metric_param_sets in _catboost.AllMetricsParams().items(): + for param_set in metric_param_sets: + derived_name = metric_name + param_set['_name_suffix'] + del param_set['_name_suffix'] + valid_params = {param: param_value['default_value'] if not param_value['is_mandatory'] else None + for param, param_value in param_set.items()} + is_mandatory_param = {param: param_value['is_mandatory'] for param, param_value in param_set.items()} + if 'hints' not in valid_params: + valid_params.update({'hints': ''}) + is_mandatory_param.update({'hints': False}) + globals()[derived_name] = _MetricGenerator(str(derived_name), (BuiltinMetric,), { + '_valid_params': valid_params, + '_is_mandatory_param': is_mandatory_param, + '_underlying_metric_name': metric_name, + }) + globals()['__all__'].append(derived_name) + + +_generate_metric_classes() diff --git a/v15cat/pylib/catboost/monoforest.py b/v15cat/pylib/catboost/monoforest.py new file mode 100644 index 0000000000000000000000000000000000000000..45d7dcaa1aef4f6d87741360cd8e4efcfb7bc1aa --- /dev/null +++ b/v15cat/pylib/catboost/monoforest.py @@ -0,0 +1,133 @@ +import math + +from . import _catboost +from .core import CatBoost, CatBoostError + +from .utils import _import_matplotlib + +FeatureExplanation = _catboost.FeatureExplanation + + +def _check_model(model): + if not isinstance(model, CatBoost): + raise CatBoostError("Model should be CatBoost") + + +def to_polynom(model): + _check_model(model) + return _catboost.to_polynom(model._object) + + +def to_polynom_string(model): + _check_model(model) + return _catboost.to_polynom_string(model._object) + + +def explain_features(model): + _check_model(model) + return _catboost.explain_features(model._object) + + +def calc_features_strength(model): + explanations = explain_features(model) + features_strength = [expl.calc_strength() for expl in explanations] + return features_strength + + +def plot_pdp(arg, size_per_plot=(5, 5), plots_per_row=None): + with _import_matplotlib() as _plt: + plt = _plt + if isinstance(arg, CatBoost): + arg = explain_features(arg) + if isinstance(arg, _catboost.FeatureExplanation): + arg = [arg] + assert len(arg) > 0 + assert isinstance(arg, list) + for element in arg: + assert isinstance(element, _catboost.FeatureExplanation) + + figs = [] + for feature_explanation in arg: + dimension = feature_explanation.dimension() + if not plots_per_row: + plots_per_row = min(5, dimension) + rows = int(math.ceil(dimension / plots_per_row)) + fig, axes = plt.subplots(rows, plots_per_row) + fig.suptitle("Feature #{}".format(feature_explanation.feature)) + if rows == 1: + axes = [axes] + if plots_per_row == 1: + axes = [[row_axes] for row_axes in axes] + fig.set_size_inches(size_per_plot[0] * plots_per_row, size_per_plot[1] * rows) + + for dim in range(dimension): + ax = axes[dim // plots_per_row][dim % plots_per_row] + ax.set_title("Dimension={}".format(dim)) + ax.set_xlabel("feature value") + ax.set_ylabel("model value") + + borders, values = feature_explanation.calc_pdp(dim) + xs = [] + ys = [] + if feature_explanation.type == "Float": + if len(borders) == 0: + xs.append(-0.1) + xs.append(0.1) + ys.append(feature_explanation.expected_bias[dim]) + ys.append(feature_explanation.expected_bias[dim]) + ax.plot(xs, ys) + else: + offset = max(0.1, (borders[0] + borders[-1]) / 2) + xs.append(borders[0] - offset) + ys.append(feature_explanation.expected_bias[dim]) + for border, value in zip(borders, values): + xs.append(border) + ys.append(ys[-1]) + xs.append(border) + ys.append(value) + xs.append(borders[-1] + offset) + ys.append(ys[-1]) + ax.plot(xs, ys) + else: + xs = ['bias'] + list(map(str, borders)) + ys = feature_explanation.expected_bias[dim] + values + ax.bar(xs, ys) + figs.append(fig) + + return figs + + +def plot_features_strength(model, height_per_feature=0.5, width_per_plot=5, plots_per_row=None): + with _import_matplotlib() as _plt: + plt = _plt + strengths = calc_features_strength(model) + dimension = len(strengths[0]) + features = len(strengths) + if not plots_per_row: + plots_per_row = min(5, dimension) + rows = int(math.ceil(dimension / plots_per_row)) + fig, axes = plt.subplots(rows, plots_per_row) + if rows == 1: + axes = [axes] + if plots_per_row == 1: + axes = [[row_axes] for row_axes in axes] + fig.suptitle("Features Strength") + fig.set_size_inches(width_per_plot * plots_per_row, height_per_feature * features * rows) + + for dim in range(dimension): + strengths = [(s[dim], i) for i, s in enumerate(strengths)] + # strengths = list(reversed(sorted(strengths))) + strengths = list(sorted(strengths)) + labels = ["Feature #{}".format(f) for _, f in strengths] + strengths = [s for s, _ in strengths] + + ax = axes[dim // plots_per_row][dim % plots_per_row] + colors = [(1, 0, 0) if s > 0 else (0, 0, 1) for s in strengths] + ax.set_title("Dimension={}".format(dim)) + ax.barh(range(len(strengths)), strengths, align='center', color=colors) + ax.set_yticks(range(len(strengths))) + ax.set_yticklabels(labels) + # ax.invert_yaxis() # labels read top-to-bottom + ax.set_xlabel('Prediction value change') + + return fig diff --git a/v15cat/pylib/catboost/plot_helpers.py b/v15cat/pylib/catboost/plot_helpers.py new file mode 100644 index 0000000000000000000000000000000000000000..add75eaa1b2551c3b7ef72190fad4f4be7dd00b4 --- /dev/null +++ b/v15cat/pylib/catboost/plot_helpers.py @@ -0,0 +1,184 @@ +import json +import os +import warnings + +from . import _catboost +fspath = _catboost.fspath + + +def try_plot_offline(figs): + try: + from plotly.offline import iplot + from plotly.offline import init_notebook_mode + init_notebook_mode(connected=True) + except ImportError as e: + warn_msg = "To draw plots you should install plotly." + warnings.warn(warn_msg) + raise ImportError(str(e)) + if not isinstance(figs, list): + figs = [figs] + for fig in figs: + iplot(fig) + + +def save_plot_file(plot_file, plot_name, figs): + warn_msg = "To draw plots you should install plotly." + try: + from plotly.offline import plot as plotly_plot + except ImportError as e: + warnings.warn(warn_msg) + raise ImportError(str(e)) + + def write_plot_file(plot_file_stream): + plot_file_stream.write('\n'.join(( + '', + '', + '', + '', + '{}'.format(plot_name), + '', + '' + ))) + for fig in figs: + graph_div = plotly_plot( + fig, + output_type='div', + show_link=False, + include_plotlyjs=False + ) + plot_file_stream.write('\n{}\n'.format(graph_div)) + plot_file_stream.write('\n') + + if isinstance(plot_file, str): + with open(fspath(plot_file), 'w') as plot_file_stream: + write_plot_file(plot_file_stream) + else: + write_plot_file(plot_file) + + +class OfflineMetricVisualizer(object): + def __init__(self, train_dirs): + if isinstance(train_dirs, str): + train_dirs = [train_dirs] + self._load_data(train_dirs) + + def _update_data_from_dir(self, path): + data = { + 'iterations': [], + 'meta': {} + } + + training_json = os.path.join(path, 'catboost_training.json') + + if os.path.isfile(training_json): + try: + with open(training_json, 'r') as json_data: + training_data = json.load(json_data) + data['meta'] = training_data['meta'] + data['iterations'] = training_data['iterations'] + except ValueError: + pass + + return { + 'passed_iterations': data['iterations'][-1]['iteration'] if data['iterations'] else 0, + 'total_iterations': data['meta']['iteration_count'] if data['meta'] else 0, + 'data': data + } + + def _load_data(self, train_dirs): + names = [] + curdir = os.path.abspath(os.path.curdir) + for train_dir in train_dirs: + abspath = os.path.abspath(train_dir) + names.append(os.path.basename(abspath) if abspath != curdir else 'current') + + data = {} + dirs = [{'name': name, 'path': path} for name, path in zip(names, train_dirs)] + + for dir_info in dirs: + path = dir_info.get('path') + content = self._update_data_from_dir(path) + + if not content: + continue + + data[path] = { + 'path': path, + 'name': dir_info.get('name'), + 'content': content + } + self.data = data + + def _get_plotly_figs(self, title): + try: + import plotly.graph_objs as go + except ImportError as err: + warnings.warn("To save plots to files you should install plotly.") + raise ImportError(str(err)) + + figs = [] + for path, dir_data in self.data.items(): + meta = dir_data['content']['data']['meta'] + + has_parameters = meta['parameters'] == 'parameters' + + # name -> {'learn_sets': [(name, metric_idx)], 'test_sets': [(name, metric_idx)]} + metrics = {} + for i, learn_metric in enumerate(meta['learn_metrics']): + metric_name = learn_metric['name'] + metrics.setdefault(metric_name, {}) + metrics[metric_name]['learn_sets'] = [(set_name, i) for set_name in meta['learn_sets']] + for i, test_metric in enumerate(meta['test_metrics']): + metric_name = test_metric['name'] + metrics.setdefault(metric_name, {}) + metrics[metric_name]['test_sets'] = [(set_name, i) for set_name in meta['test_sets']] + + iterations = dir_data['content']['data']['iterations'] + + for metric_name, subsets in metrics.items(): + fig = go.Figure() + + figure_title = (title if (dir_data['name'] == 'catboost_info') else dir_data['name']) + ' : ' + metric_name + fig['layout']['title'] = go.layout.Title(text=figure_title) + + learn_graph_color = 'rgb(160,0,0)' + + for learn_set_name, metric_idx in subsets.get('learn_sets', []): + fig.add_trace(go.Scatter( + x=[e['iteration'] for e in iterations], + y=[e[learn_set_name][metric_idx] for e in iterations], + line=go.scatter.Line(color=learn_graph_color), + mode='lines', + name=learn_set_name + )) + + test_graph_color = 'rgb(0,160,0)' + + for test_set_name, metric_idx in subsets.get('test_sets', []): + + def generate_params_hover(e): + result = [] + for param_name, param_value in e["parameters"][0].items(): + result.append(param_name + ' : ' + str(param_value)) + return '
'.join(result) + + fig.add_trace(go.Scatter( + x=[e['iteration'] for e in iterations], + y=[e[test_set_name][metric_idx] for e in iterations], + line=go.scatter.Line(color=test_graph_color), + mode='lines', + name=test_set_name, + hovertext=[generate_params_hover(e) for e in iterations] if has_parameters else None + )) + + fig.update_layout( + xaxis=dict(title='iterations'), + yaxis=dict(title=metric_name) + ) + + figs.append(fig) + + return figs + + def save_to_file(self, title, file_output): + save_plot_file(file_output, title, self._get_plotly_figs(title)) diff --git a/v15cat/pylib/catboost/text_processing.py b/v15cat/pylib/catboost/text_processing.py new file mode 100644 index 0000000000000000000000000000000000000000..3da142d3c0715d96dda73148f8656c635eb49295 --- /dev/null +++ b/v15cat/pylib/catboost/text_processing.py @@ -0,0 +1,4 @@ +from . import _catboost + +Tokenizer = _catboost.Tokenizer +Dictionary = _catboost.Dictionary diff --git a/v15cat/pylib/catboost/utils.py b/v15cat/pylib/catboost/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..c780a6f8c0d49893729d2b852dd31ac7087280e4 --- /dev/null +++ b/v15cat/pylib/catboost/utils.py @@ -0,0 +1,781 @@ +from . import _catboost +from .core import Pool, CatBoostError, ARRAY_TYPES, PATH_TYPES, fspath, _update_params_quantize_part, _process_synonyms +from collections import defaultdict +from contextlib import contextmanager +import sys +import numpy as np +import warnings + +_eval_metric_util = _catboost._eval_metric_util +_get_roc_curve = _catboost._get_roc_curve +_get_confusion_matrix = _catboost._get_confusion_matrix +_select_threshold = _catboost._select_threshold +_NumpyAwareEncoder = _catboost._NumpyAwareEncoder +_get_onnx_model = _catboost._get_onnx_model +_calculate_quantization_grid = _catboost._calculate_quantization_grid + +compute_wx_test = _catboost.compute_wx_test +TargetStats = _catboost.TargetStats +DataMetaInfo = _catboost.DataMetaInfo +compute_training_options = _catboost.compute_training_options + + +try: + import polars as pl +except ImportError: + # just to avoid checking (pl is not None) everywhere + class polars: + class DataFrame(object): + pass + + class Series(object): + pass + pl = polars + + +@contextmanager +def _import_matplotlib(): + try: + import matplotlib.pyplot as plt + except ImportError as e: + warnings.warn("To draw plots you should install matplotlib.") + raise ImportError(str(e)) + yield plt + + +def _draw(plt, x, y, x_label, y_label, title): + plt.figure(figsize=(16, 8)) + + plt.plot(x, y, alpha=0.5, lw=2) + + plt.xticks(fontsize=16) + plt.yticks(fontsize=16) + plt.grid(True) + plt.xlabel(x_label, fontsize=16) + plt.ylabel(y_label, fontsize=16) + plt.title(title, fontsize=20) + plt.show() + + +def create_cd( + label=None, + cat_features=None, + text_features=None, + embedding_features=None, + weight=None, + baseline=None, + doc_id=None, + group_id=None, + subgroup_id=None, + timestamp=None, + auxiliary_columns=None, + feature_names=None, + output_path='train.cd' +): + _from_param_to_cd = { + 'label': 'Label', + 'weight': 'Weight', + 'baseline': 'Baseline', + 'doc_id': 'DocId', + 'group_id': 'GroupId', + 'subgroup_id': 'SubgroupId', + 'timestamp': 'Timestamp' + } + _column_description = defaultdict(lambda: ['Num', '']) + for key, value in locals().copy().items(): + if not (key.startswith('_') or value is None): + if key in ('cat_features', 'text_features', 'embedding_features', 'auxiliary_columns'): + if isinstance(value, int): + value = [value] + for index in value: + if not isinstance(index, int): + raise CatBoostError('Unsupported index type. Expected int, got {}'.format(type(index))) + if index in _column_description: + raise CatBoostError('The index {} occurs more than once'.format(index)) + if key == 'cat_features': + _column_description[index] = ['Categ', ''] + elif key == 'text_features': + _column_description[index] = ['Text', ''] + elif key == 'embedding_features': + _column_description[index] = ['NumVector', ''] + else: + _column_description[index] = ['Auxiliary', ''] + elif key not in ('feature_names', 'output_path'): + if not isinstance(value, int): + raise CatBoostError('Unsupported index type. Expected int, got {}'.format(type(value))) + if value in _column_description: + raise CatBoostError('The index {} occurs more than once'.format(value)) + _column_description[value] = [_from_param_to_cd[key], ''] + if feature_names is not None: + for feature_column_index, name in feature_names.items(): + if _column_description[feature_column_index][0] not in ('Num', 'Categ', 'Text', 'NumVector'): + raise CatBoostError('feature_names contains index {} that does not correspond to feature column'.format(feature_column_index)) + _column_description[feature_column_index][1] = name + with open(fspath(output_path), 'w') as f: + for index, (title, name) in sorted(_column_description.items()): + f.write('{}\t{}\t{}\n'.format(index, title, name)) + + +def read_cd(cd_file, column_count=None, data_file=None, canonize_column_types=False): + """ + Reads CatBoost column description file + (see https://catboost.ai/docs/concepts/input-data_column-descfile.html#input-data_column-descfile) + + Parameters + ---------- + cd_file : str or pathlib.Path + path to column description file + + column_count : integer + total number of columns + + data_file : str or pathlib.Path + path to dataset file in CatBoost format + specify either column_count directly or data_file to detect it + + canonize_column_types : bool + if set to True types for columns with synonyms are renamed to canonical type. + + Returns + ------- + dict with keys: + "column_type_to_indices" : + dict of column_type -> column_indices list, column_type is 'Label', 'Categ' etc. + + "column_dtypes" : dict of column_name -> numpy.dtype or 'category' + + "cat_feature_indices" : list of integers + indices of categorical features in array of all features. + Note: indices in array of features, not indices in array of all columns! + + "text_feature_indices" : list of integers + indices of text features in array of all features. + Note: indices in array of features, not indices in array of all columns! + + "embedding_feature_indices" : list of integers + indices of embedding features in array of all features. + Note: indices in array of features, not indices in array of all columns! + + "column_names" : list of strings + + "non_feature_column_indices" : list of integers + """ + + column_type_synonyms_map = { + 'Target': 'Label', + 'DocId': 'SampleId', + 'QueryId': 'GroupId' + } + + if column_count is None: + if data_file is None: + raise Exception( + 'Cannot obtain column count: either specify column_count parameter or specify data_file ' + + 'parameter to get it' + ) + with open(fspath(data_file)) as f: + column_count = len(f.readline()[:-1].split('\t')) + + column_type_to_indices = {} + column_dtypes = {} + cat_feature_indices = [] + text_feature_indices = [] + embedding_feature_indices = [] + column_names = [] + non_feature_column_indices = [] + + # list of (column_idx, column_type, column_name) tuples, needed to support CD files with nonincreasing + # column indices + column_descriptions = [] + + with open(fspath(cd_file)) as f: + for line_idx, line in enumerate(f): + line = line.strip() + + # some cd files in the wild contain empty lines + if len(line) == 0: + continue + + line_columns = line.split('\t') + if len(line_columns) not in [2, 3]: + raise Exception('Wrong number of columns in cd file') + + column_idx = int(line_columns[0]) + + column_type = line_columns[1] + + column_name = None + if len(line_columns) == 3: + column_name = line_columns[2] + + column_descriptions.append((column_idx, column_type, column_name)) + + column_descriptions.sort() + + def add_missed_columns(start_column_idx, end_column_idx, non_feature_column_count): + for missed_column_idx in range(start_column_idx, end_column_idx): + column_name = 'feature_%i' % (missed_column_idx - non_feature_column_count) + column_names.append(column_name) + column_type_to_indices.setdefault('Num', []).append(missed_column_idx) + column_dtypes[column_name] = np.float32 + + last_column_idx = -1 + for column_idx, column_type, column_name in column_descriptions: + if column_idx == last_column_idx: + raise Exception('Duplicate column indices in cd file') + + add_missed_columns(last_column_idx + 1, column_idx, len(non_feature_column_indices)) + + if canonize_column_types: + column_type = column_type_synonyms_map.get(column_type, column_type) + + column_type_to_indices.setdefault(column_type, []).append(column_idx) + + if column_type in ['Num', 'Categ', 'Text', 'NumVector']: + feature_idx = column_idx - len(non_feature_column_indices) + if column_name is None: + column_name = 'feature_%i' % feature_idx + if column_type == 'Categ': + cat_feature_indices.append(feature_idx) + column_dtypes[column_name] = 'category' + elif column_type == 'Text': + text_feature_indices.append(feature_idx) + column_dtypes[column_name] = object + elif column_type == 'NumVector': + embedding_feature_indices.append(feature_idx) + column_dtypes[column_name] = object + else: + column_dtypes[column_name] = np.float32 + else: + non_feature_column_indices.append(column_idx) + if column_name is None: + column_name = column_type + + column_names.append(column_name) + + last_column_idx = column_idx + + add_missed_columns(last_column_idx + 1, column_count, len(non_feature_column_indices)) + + return { + 'column_type_to_indices' : column_type_to_indices, + 'column_dtypes' : column_dtypes, + 'cat_feature_indices' : cat_feature_indices, + 'text_feature_indices' : text_feature_indices, + 'embedding_feature_indices' : embedding_feature_indices, + 'column_names' : column_names, + 'non_feature_column_indices' : non_feature_column_indices + } + + +def eval_metric(label, approx, metric, weight=None, group_id=None, group_weight=None, subgroup_id=None, pairs=None, thread_count=-1): + """ + Evaluate metrics with raw approxes and labels. + + Parameters + ---------- + label : list or numpy.ndarrays or pandas.DataFrame or pandas.Series or polars.DataFrame or polars.Series + Object labels with shape (n_objects,) or (n_object, n_target_dimension) + + approx : list or numpy.ndarrays or pandas.DataFrame or pandas.Series + Object approxes with shape (n_objects,) or (n_object, n_approx_dimension). + + metric : string + Metric name. + + weight : list or numpy.ndarray or pandas.DataFrame or pandas.Series or polars.Series, optional (default=None) + Object weights. + + group_id : list or numpy.ndarray or pandas.DataFrame or pandas.Series or polars.Series, optional (default=None) + Object group ids. + + group_weight : list or numpy.ndarray or pandas.DataFrame or pandas.Series or polars.Series, optional (default=None) + Group weights. + + subgroup_id : list or numpy.ndarray or polars.Series, optional (default=None) + subgroup id for each instance. + If not None, giving 1 dimensional array like data. + + pairs : list or numpy.ndarray or pandas.DataFrame or polars.DataFrame or string or pathlib.Path + The pairs description. + If list or numpy.ndarrays or pandas.DataFrame or polars.DataFrame, giving 2 dimensional. + The shape should be Nx2, where N is the pairs' count. The first element of the pair is + the index of winner object in the training set. The second element of the pair is + the index of loser object in the training set. + If string or pathlib.Path, giving the path to the file with pairs description. + + thread_count : int, optional (default=-1) + Number of threads to work with. + If -1, then the number of threads is set to the number of CPU cores. + + Returns + ------- + metric results : list with metric values. + """ + if len(label) > 0: + if isinstance(label, pl.Series): + label = pl.DataFrame(label) + elif not isinstance(label, pl.DataFrame): + label = np.transpose(label) if isinstance(label[0], ARRAY_TYPES) else [label] + if len(approx) == 0: + approx = [[]] + approx = np.transpose(approx) if isinstance(approx[0], ARRAY_TYPES) else [approx] + return _eval_metric_util(label, approx, metric, weight, group_id, group_weight, subgroup_id, pairs, thread_count) + + +def get_gpu_device_count(): + return _catboost._get_gpu_device_count() + + +def reset_trace_backend(filename): + _catboost._reset_trace_backend(filename) + + +def get_confusion_matrix(model, data, thread_count=-1): + """ + Build confusion matrix. + + Parameters + ---------- + model : catboost.CatBoost + The trained model. + + data : catboost.Pool + A set of samples to build confusion matrix with. + + thread_count : int (default=-1) + Number of threads to work with. + If -1, then the number of threads is set to the number of CPU cores. + + Returns + ------- + confusion matrix : array, shape = [n_classes, n_classes] + """ + if not isinstance(data, Pool): + raise CatBoostError('data must be a catboost.Pool') + + return _get_confusion_matrix(model._object, data, thread_count) + + +def get_roc_curve(model, data, thread_count=-1, plot=False): + """ + Build points of ROC curve. + + Parameters + ---------- + model : catboost.CatBoost + The trained model. + + data : catboost.Pool or list of catboost.Pool + A set of samples to build ROC curve with. + + thread_count : int (default=-1) + Number of threads to work with. + If -1, then the number of threads is set to the number of CPU cores. + + plot : bool, optional (default=False) + If True, draw curve. + + Returns + ------- + curve points : tuple of three arrays (fpr, tpr, thresholds) + """ + if isinstance(data, Pool): + data = [data] + if not isinstance(data, list): + raise CatBoostError('data must be a catboost.Pool or list of pools.') + for pool in data: + if not isinstance(pool, Pool): + raise CatBoostError('one of data pools is not catboost.Pool') + + roc_curve = _get_roc_curve(model._object, data, thread_count) + + if plot: + with _import_matplotlib() as plt: + _draw(plt, roc_curve[0], roc_curve[1], 'False Positive Rate', 'True Positive Rate', 'ROC Curve') + + return roc_curve + + +def get_fpr_curve(model=None, data=None, curve=None, thread_count=-1, plot=False): + """ + Build points of FPR curve. + + Parameters + ---------- + model : catboost.CatBoost + The trained model. + + data : catboost.Pool or list of catboost.Pool + A set of samples to build ROC curve with. + + curve : tuple of three arrays (fpr, tpr, thresholds) + ROC curve points in format of get_roc_curve returned value. + If set, data parameter must not be set. + + thread_count : int (default=-1) + Number of threads to work with. + If -1, then the number of threads is set to the number of CPU cores. + + plot : bool, optional (default=False) + If True, draw curve. + + Returns + ------- + curve points : tuple of two arrays (thresholds, fpr) + """ + if curve is not None: + if data is not None: + raise CatBoostError('Only one of the parameters data and curve should be set.') + if not (isinstance(curve, list) or isinstance(curve, tuple)) or len(curve) != 3: + raise CatBoostError('curve must be list or tuple of three arrays (fpr, tpr, thresholds).') + fpr, thresholds = curve[0][:], curve[2][:] + else: + if model is None or data is None: + raise CatBoostError('model and data parameters should be set when curve parameter is None.') + fpr, _, thresholds = get_roc_curve(model, data, thread_count) + + if plot: + with _import_matplotlib() as plt: + _draw(plt, thresholds, fpr, 'Thresholds', 'False Positive Rate', 'FPR Curve') + + return thresholds, fpr + + +def get_fnr_curve(model=None, data=None, curve=None, thread_count=-1, plot=False): + """ + Build points of FNR curve. + + Parameters + ---------- + model : catboost.CatBoost + The trained model. + + data : catboost.Pool or list of catboost.Pool + A set of samples to build ROC curve with. + + curve : tuple of three arrays (fpr, tpr, thresholds) + ROC curve points in format of get_roc_curve returned value. + If set, data parameter must not be set. + + thread_count : int (default=-1) + Number of threads to work with. + If -1, then the number of threads is set to the number of CPU cores. + + plot : bool, optional (default=False) + If True, draw curve. + + Returns + ------- + curve points : tuple of two arrays (thresholds, fnr) + """ + if curve is not None: + if data is not None: + raise CatBoostError('Only one of the parameters data and curve should be set.') + if not (isinstance(curve, list) or isinstance(curve, tuple)) or len(curve) != 3: + raise CatBoostError('curve must be list or tuple of three arrays (fpr, tpr, thresholds).') + tpr, thresholds = curve[1], curve[2][:] + else: + if model is None or data is None: + raise CatBoostError('model and data parameters should be set when curve parameter is None.') + _, tpr, thresholds = get_roc_curve(model, data, thread_count) + fnr = np.array([1 - x for x in tpr]) + + if plot: + with _import_matplotlib() as plt: + _draw(plt, thresholds, fnr, 'Thresholds', 'False Negative Rate', 'FNR Curve') + + return thresholds, fnr + + +def select_threshold(model=None, data=None, curve=None, FPR=None, FNR=None, thread_count=-1): + """ + Selects a threshold for prediction. + + Parameters + ---------- + model : catboost.CatBoost + The trained model. + + data : catboost.Pool or list of catboost.Pool + Set of samples to build ROC curve with. + If set, curve parameter must not be set. + + curve : tuple of three arrays (fpr, tpr, thresholds) + ROC curve points in format of get_roc_curve returned value. + If set, data parameter must not be set. + + FPR : desired false-positive rate + + FNR : desired false-negative rate (only one of FPR and FNR should be chosen) + + thread_count : int (default=-1) + Number of threads to work with. + If -1, then the number of threads is set to the number of CPU cores. + + Returns + ------- + threshold : double + """ + if data is not None: + if curve is not None: + raise CatBoostError('Only one of the parameters data and curve should be set.') + if model is None: + raise CatBoostError('model and data parameters should be set when curve parameter is None.') + if isinstance(data, Pool): + data = [data] + if not isinstance(data, list): + raise CatBoostError('data must be a catboost.Pool or list of pools.') + for pool in data: + if not isinstance(pool, Pool): + raise CatBoostError('one of data pools is not catboost.Pool') + return _select_threshold(model._object, data, None, FPR, FNR, thread_count) + elif curve is not None: + if not (isinstance(curve, list) or isinstance(curve, tuple)) or len(curve) != 3: + raise CatBoostError('curve must be list or tuple of three arrays (fpr, tpr, thresholds).') + return _select_threshold(None, None, curve, FPR, FNR, thread_count) + else: + raise CatBoostError('One of the parameters data and curve should be set.') + + +def quantize( + data_path, + column_description=None, + pairs=None, + graph=None, + delimiter='\t', + has_header=False, + ignore_csv_quoting=False, + feature_names=None, + thread_count=-1, + ignored_features=None, + per_float_feature_quantization=None, + border_count=None, + max_bin=None, + feature_border_type=None, + nan_mode=None, + input_borders=None, + task_type=None, + used_ram_limit=None, + random_seed=None, + log_cout=sys.stdout, + log_cerr=sys.stderr, + **kwargs +): + """ + Construct quantized Pool from non-quantized pool stored in file. + This method does not load whole non-quantized source dataset into memory + so it can be used for huge datasets that fit in memory only after quantization. + + Parameters + ---------- + data_path : string or pathlib.Path + Path (with optional scheme) to non-quantized dataset. + + column_description : string, [default=None] + ColumnsDescription parameter. + There are several columns description types: Label, Categ, Num, Auxiliary, DocId, Weight, Baseline, GroupId, Timestamp. + All columns are Num as default, it's not necessary to specify + this type of columns. Default Label column index is 0 (zero). + If None, Label column is 0 (zero) as default, all data columns are Num as default. + If string or pathlib.Path, giving the path to the file with ColumnsDescription in column_description format. + + pairs : string or pathlib.Path, [default=None] + Path to the file with pairs description. + + graph : string or pathlib.Path, [default=None] + Path to the file with graph description. + + has_header : bool, [default=False] + If True, read column names from first line. + + ignore_csv_quoting : bool optional (default=False) + If True ignore quoting '"'. + + feature_names : string or pathlib.Path, [default=None] + Path with scheme for feature names data to load. + + thread_count : int, [default=-1] + Thread count for data processing. + If -1, then the number of threads is set to the number of CPU cores. + + ignored_features : list, [default=None] + Indices or names of features that should be excluded when training. + + per_float_feature_quantization : list of strings, [default=None] + List of float binarization descriptions. + Format : described in documentation on catboost.ai + Example 1: ['0:1024'] means that feature 0 will have 1024 borders. + Example 2: ['0:border_count=1024', '1:border_count=1024', ...] means that two first features have 1024 borders. + Example 3: ['0:nan_mode=Forbidden,border_count=32,border_type=GreedyLogSum', + '1:nan_mode=Forbidden,border_count=32,border_type=GreedyLogSum'] - defines more quantization properties for first two features. + + border_count : int, [default = 254 for training on CPU or 128 for training on GPU] + The number of partitions in numeric features binarization. Used in the preliminary calculation. + range: [1,65535] on CPU, [1,255] on GPU + + max_bin : float, synonym for border_count. + + feature_border_type : string, [default='GreedyLogSum'] + The binarization mode in numeric features binarization. Used in the preliminary calculation. + Possible values: + - 'Median' + - 'Uniform' + - 'UniformAndQuantiles' + - 'GreedyLogSum' + - 'MaxLogSum' + - 'MinEntropy' + + nan_mode : string, [default=None] + Way to process missing values for numeric features. + Possible values: + - 'Forbidden' - raises an exception if there is a missing value for a numeric feature in a dataset. + - 'Min' - each missing value will be processed as the minimum numerical value. + - 'Max' - each missing value will be processed as the maximum numerical value. + If None, then nan_mode=Min. + + input_borders : string or pathlib.Path, [default=None] + input file with borders used in numeric features binarization. + + task_type : string, [default=None] + The calcer type used to train the model. + Possible values: + - 'CPU' + - 'GPU' + + used_ram_limit=None + + random_seed : int, [default=None] + The random seed used for data sampling. + If None, 0 is used. + + Returns + ------- + pool : Pool + Constructed and quantized pool. + """ + if not data_path: + raise CatBoostError("Data filename is empty.") + if not isinstance(data_path, PATH_TYPES): + raise CatBoostError("Data filename should be string or pathlib.Path type.") + + if pairs is not None and not isinstance(pairs, PATH_TYPES): + raise CatBoostError("pairs should have None or string or pathlib.Path type when the pool is read from the file.") + if column_description is not None and not isinstance(column_description, PATH_TYPES): + raise CatBoostError("column_description should have None or string or pathlib.Path type when the pool is read from the file.") + if feature_names is not None and not isinstance(feature_names, PATH_TYPES): + raise CatBoostError("feature_names should have None or string or pathlib.Path type when the pool is read from the file.") + + params = {} + _process_synonyms(params) + + if border_count is None: + border_count = max_bin + + if 'dev_block_size' in kwargs: + params['dev_block_size'] = kwargs.pop('dev_block_size') + + dev_max_subset_size_for_build_borders = kwargs.pop('dev_max_subset_size_for_build_borders', None) + + if kwargs: + raise CatBoostError("got an unexpected keyword arguments: {}".format(kwargs.keys())) + + _update_params_quantize_part( + params, + ignored_features, + per_float_feature_quantization, + border_count, + feature_border_type, + None, # sparse_features_conflict_fraction + None, # dev_efb_max_buckets + nan_mode, + input_borders, + task_type, + used_ram_limit, + random_seed, + dev_max_subset_size_for_build_borders + ) + + pool = Pool( + data_path, + column_description=column_description, + pairs=pairs, + graph=graph, + feature_names=feature_names, + delimiter=delimiter, + has_header=has_header, + ignore_csv_quoting=ignore_csv_quoting, + thread_count=thread_count, + log_cout=log_cout, + log_cerr=log_cerr + ) + pool._read( + data_path, + column_description, + pairs, + graph, + feature_names, + delimiter, + has_header, + ignore_csv_quoting, + thread_count, + params, + log_cout=log_cout, + log_cerr=log_cerr + ) + + return pool + + +def convert_to_onnx_object(model, export_parameters=None, **kwargs): + """ + Convert given CatBoost model to ONNX-ML model. + Categorical Features are not supported. + + Parameters + ---------- + model : CatBoost trained model + export_parameters : dict [default=None] + Parameters for ONNX-ML export: + * onnx_graph_name : string + The name property of onnx Graph + * onnx_domain : string + The domain component of onnx Model + * onnx_model_version : int + The model_version component of onnx Model + * onnx_doc_string : string + The doc_string component of onnx Model + Returns + ------- + onnx_object : ModelProto + The model in ONNX format + """ + try: + import onnx + except ImportError as e: + warnings.warn("To get working onnx model you should install onnx.") + raise ImportError(str(e)) + + import json + if not model.is_fitted(): + raise CatBoostError( + "There is no trained model to use save_model(). Use fit() to train model. Then use this method.") + + for name, value in kwargs.items(): + if name == 'target_opset' and value not in [None, 2]: + warnings.warn('target_opset argument is not supported. Default target_opset is 2 (ai.onnx.ml domain)') + elif name == 'initial_types' and value is not None: + warnings.warn('initial_types argument is not supported') + + params_string = "" + if export_parameters: + params_string = json.dumps(export_parameters, cls=_NumpyAwareEncoder) + + model_str = _get_onnx_model(model._object, params_string) + onnx_model = onnx.load_model_from_string(model_str) + return onnx_model + + +def calculate_quantization_grid(values, border_count, border_type='Median'): + assert border_count > 0, 'Border count should be > 0' + return _calculate_quantization_grid(values, border_count, border_type) diff --git a/v15cat/pylib/catboost/version.py b/v15cat/pylib/catboost/version.py new file mode 100644 index 0000000000000000000000000000000000000000..3c0de5348f893f4de36e8cfc19bd3469d5aee505 --- /dev/null +++ b/v15cat/pylib/catboost/version.py @@ -0,0 +1 @@ +VERSION = '1.2.10' diff --git a/v15cat/pylib/catboost/widget/__init__.py b/v15cat/pylib/catboost/widget/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9554674e620fcd57c3b760208b4a2625792c0507 --- /dev/null +++ b/v15cat/pylib/catboost/widget/__init__.py @@ -0,0 +1,19 @@ +from .ipythonwidget import MetricVisualizer # noqa +from .metrics_plotter import MetricsPlotter # noqa +from .callbacks import XGBPlottingCallback, lgbm_plotting_callback # noqa + + +def _jupyter_labextension_paths(): + return [{ + 'src': 'labextension', + 'dest': 'catboost-widget', + }] + + +def _jupyter_nbextension_paths(): + return [{ + 'section': 'notebook', + 'src': 'nbextension', + 'dest': 'catboost-widget', + 'require': 'catboost-widget/extension' + }] diff --git a/v15cat/pylib/catboost/widget/callbacks.py b/v15cat/pylib/catboost/widget/callbacks.py new file mode 100644 index 0000000000000000000000000000000000000000..056c7ee2a9357144a1287a23433c84c82c35b62d --- /dev/null +++ b/v15cat/pylib/catboost/widget/callbacks.py @@ -0,0 +1,107 @@ +try: + from xgboost.callback import TrainingCallback as XGBTrainingCallback +except Exception: + class XGBTrainingCallback: + pass + + +from IPython.display import display + +from .metrics_plotter import MetricsPlotter + + +class XGBPlottingCallback(XGBTrainingCallback): + '''XGBoost callback with metrics plotting widget from CatBoost + ''' + + def __init__(self, total_iterations: int): + self.plotter = None + self.total_iterations = total_iterations + + def after_iteration(self, model, epoch, evals_log): + + data_names = evals_log.keys() + # if more than one sample is passed, consider first as train sample + first_train = ( + all(['valid' in data_name.lower() for data_name in data_names]) + and len(data_names) > 1 + ) + + for data_name, metrics_info in evals_log.items(): + if "train" in data_name.lower() or first_train: + train = True + first_train = False + elif "valid" in data_name.lower() or "test" in data_name.lower(): + train = False + else: + raise Exception("Unexpected sample name during evaluation") + + metrics = {name: values[-1] for name, values in metrics_info.items()} + + if self.plotter is None: + names = list(metrics.keys()) + self.plotter = MetricsPlotter(names, names, self.total_iterations) + display(self.plotter._widget) + + self.plotter.log(epoch, train, metrics) + + # False to indicate training should not stop. + return False + + +def lgbm_plotting_callback(): + """LightGBM callback with metrics plotting widget from CatBoost + """ + + plotter = None + + def _init(env): + train_metrics = [] + test_metrics = [] + + for item in env.evaluation_result_list: + assert len(item) == 4, "Plotting was run in not suppored mode" + data_name, eval_name = item[:2] + + if "train" in data_name.lower(): + train_metrics.append(eval_name) + elif "valid" in data_name.lower() or "test" in data_name.lower(): + test_metrics.append(eval_name) + else: + raise Exception("Unexpected sample name during evaluation") + + nonlocal plotter + plotter = MetricsPlotter( + train_metrics, test_metrics, env.end_iteration - env.begin_iteration) + + display(plotter._widget) + + def _callback(env): + if plotter is None: + _init(env) + + metrics = {"train": {}, "test": {}} + + for item in env.evaluation_result_list: + data_name, eval_name, result = item[:3] + + if "train" in data_name.lower(): + metrics["train"][eval_name] = result + elif "valid" in data_name.lower() or "test" in data_name.lower(): + metrics["test"][eval_name] = result + else: + raise Exception("Unexpected sample name during evaluation") + + plotter.log( + env.iteration - env.begin_iteration, + train=True, + metrics=metrics["train"] + ) + plotter.log( + env.iteration - env.begin_iteration, + train=False, + metrics=metrics["test"] + ) + + _callback.order = 20 + return _callback diff --git a/v15cat/pylib/catboost/widget/ipythonwidget.py b/v15cat/pylib/catboost/widget/ipythonwidget.py new file mode 100644 index 0000000000000000000000000000000000000000..57fece840bfad7c01b43478212a30bcb4082eed2 --- /dev/null +++ b/v15cat/pylib/catboost/widget/ipythonwidget.py @@ -0,0 +1,116 @@ +import os +import json +from threading import Thread, Event +from traitlets import Unicode, Dict, default +from IPython.display import display +from ipywidgets import DOMWidget, Layout, widget_serialization + + +class MetricWidget(DOMWidget): + _view_name = Unicode('CatboostWidgetView').tag(sync=True) + _model_name = Unicode('CatboostWidgetModel').tag(sync=True) + _view_module = Unicode('catboost-widget').tag(sync=True) + _model_module = Unicode('catboost-widget').tag(sync=True) + _view_module_version = Unicode('^1.0.0').tag(sync=True) + _model_module_version = Unicode('^1.0.0').tag(sync=True) + + data = Dict({}).tag(sync=True, **widget_serialization) + + @default('layout') + def _default_layout(self): + return Layout(height='500px', align_self='stretch') + + +class MetricVisualizer(MetricWidget): + def __init__(self, train_dirs, subdirs=False): + super(self.__class__, self).__init__() + if isinstance(train_dirs, str): + train_dirs = [train_dirs] + if subdirs: + train_subdirs = [] + for train_dir in train_dirs: + train_subdirs.extend(self._get_subdirectories(train_dir)) + train_dirs = train_subdirs + self._train_dirs = train_dirs[:] + self._names = [] + curdir = os.path.abspath(os.path.curdir) + for train_dir in train_dirs: + abspath = os.path.abspath(train_dir) + self._names.append(os.path.basename(abspath) if abspath != curdir else 'current') + self._need_to_stop = Event() + self._update_after_stop_signal = False + + def start(self): + display(self) + self._update_data() + while not self._need_to_stop.wait(1.0): + self._update_data() + + if self._update_after_stop_signal: + self._update_data() + + def _run_update(self): + self.thread = Thread(target=self.start, args=()) + self.thread.start() + + def _stop_update(self): + self._update_after_stop_signal = True + self._need_to_stop.set() + self.thread.join() + + def _get_subdirectories(self, a_dir): + return [os.path.join(a_dir, name) for name in os.listdir(a_dir) if os.path.isdir(os.path.join(a_dir, name))] + + def _update_data(self): + data = {} + dirs = [{'name': name, 'path': path} for name, path in zip(self._names, self._train_dirs)] + + all_completed = True + for dir_info in dirs: + path = dir_info.get('path') + content = self._update_data_from_dir(path) + + if not content: + continue + + data[path] = { + 'path': path, + 'name': dir_info.get('name'), + 'content': content + } + + passed_iterations = data[path]['content']['passed_iterations'] + total_iterations = data[path]['content']['total_iterations'] + all_completed &= (passed_iterations + 1 >= total_iterations and total_iterations != 0) + + if all_completed: + self._need_to_stop.set() + + self.data = data + + def _update_data_from_dir(self, path): + data = { + 'iterations': [], + 'meta': {} + } + + training_json = os.path.join(path, 'catboost_training.json') + + if os.path.isfile(training_json): + try: + with open(training_json, 'r') as json_data: + training_data = json.load(json_data) + data['meta'] = training_data['meta'] + data['iterations'] = training_data['iterations'] + except ValueError: + pass + + return { + 'passed_iterations': data['iterations'][-1]['iteration'] if data['iterations'] else 0, + 'total_iterations': data['meta']['iteration_count'] if data['meta'] else 0, + 'data': data + } + + @staticmethod + def _get_static_path(file_name): + return os.path.join(os.path.dirname(__file__), file_name) diff --git a/v15cat/pylib/catboost/widget/metrics_plotter.py b/v15cat/pylib/catboost/widget/metrics_plotter.py new file mode 100644 index 0000000000000000000000000000000000000000..62044f1f8ce53ef19c0a054e8b33954a402a974f --- /dev/null +++ b/v15cat/pylib/catboost/widget/metrics_plotter.py @@ -0,0 +1,208 @@ +import time +import logging + +import typing as tp + +from IPython.display import display +from copy import deepcopy +from typing import List, Optional, Any, Union + +from .ipythonwidget import MetricWidget + +logger = logging.getLogger(__name__) +logger.setLevel(logging.DEBUG) + + +class MetricsWidget(MetricWidget): + def __init__(self): + super(self.__class__, self).__init__() + + def update_data(self, data: tp.Dict) -> None: + # deepcopy is crucial here + self.data = deepcopy(data) + + +class MetricsPlotter: + """ + Context manager that enables widget with learning curves in + JupyterLab / Jupyter Notebook + """ + + def __init__(self, train_metrics: List[Union[str, tp.Dict[str, str]]], + test_metrics: Optional[List[Union[str, tp.Dict[str, str]]]] = None, + total_iterations: Optional[int] = None) -> None: + """ + Constructor that defines metrics to be plotted and total iterations count. + + Parameters + ---------- + train_metrics : list of str or list of dict + List of train metrics to be tracked. + Each item in list can be either string with metric name or dict + with the following format: + { + "name": "{metric_name}", + "best_value": "Max|Min|Undefined", + } + + test_metrics : list of str or list of dict, optional (default=None) + List of test metrics to be tracked. + Has the same format as train_metrics. Equals to train_metrics, if not defined + + total_iterations: int, optional (default=None) + Total number of iterations, allows for remaining time estimation. + """ + + self._widget = MetricsWidget() + + self._values = { + "iterations": [], + "meta": { + "launch_mode": "Train", + "parameters": "", + "name": "experiment", + "iteration_count": None, # set later + "learn_sets": ["learn"], + "learn_metrics": None, # set later + "test_sets": ["test"], + "test_metrics": None, # set later + } + } + + self._content = { + "passed_iterations": 0, + "total_iterations": None, # set later + "data": self._values, + } + + # data propagated to widget class + self._data = { + "test_path": { + "path": "test_path", + "name": "experiment", + "content": self._content, + } + } + + test_metrics = test_metrics or train_metrics + train_metrics_meta: List[tp.Dict[str, str]] = self.construct_metrics_meta(train_metrics) + test_metrics_meta: List[tp.Dict[str, str]] = self.construct_metrics_meta(test_metrics) + + self._train_metrics_positions = { + meta["name"]: pos for pos, meta in enumerate(train_metrics_meta)} + self._test_metrics_positions = { + meta["name"]: pos for pos, meta in enumerate(test_metrics_meta)} + + self._values["meta"].update({ + "learn_metrics": train_metrics_meta, + "test_metrics": test_metrics_meta, + }) + + self.passed_iterations = 0 + self.total_iterations = total_iterations + + if total_iterations is not None: + self._values["meta"]["iteration_count"] = total_iterations + + self._content["total_iterations"] = total_iterations or 0 + + self._start_time = time.time() + + def __enter__(self) -> 'MetricsPlotter': + display(self._widget) + return self + + def __exit__(self, exc_type, exc_value, traceback) -> Any: + if exc_type == KeyboardInterrupt: + logger.info( + f"Learning was stopped manually after {self.passed_iterations} epochs") + return True + + @staticmethod + def construct_metrics_meta(metrics: List[Union[str, tp.Dict[str, str]]]) -> List[tp.Dict[str, str]]: + meta: List[tp.Dict[str, str]] = [] + for item in metrics: + if isinstance(item, str): + name, best_value = item, "Undefined" + elif isinstance(item, dict): + assert "name" in item and "best_value" in item, \ + "Wrong metrics definition format: should have " \ + "`name` and `best_value` fields" + name, best_value = item["name"], item["best_value"] + else: + assert False, "Each metric should be defined as str or as" \ + "dict with `name` and `best_value` fields" + meta.append({"best_value": best_value, "name": name}) + return meta + + @staticmethod + def construct_metrics_array(metrics_positions: tp.Dict[str, int], + metrics: tp.Dict[str, float]) -> List[float]: + array: List[float] = [0.] * len(metrics_positions) + + # data validation + assert set(metrics.keys()) == set(metrics_positions.keys()), \ + f"Not all metrics were passed while logging, expected " \ + f"following: {', '.join(list(metrics_positions.keys()))}" + + for metric, value in metrics.items(): + assert isinstance(value, float), "Type of metric {metric} should be float" + array[metrics_positions[metric]] = value + return array + + def estimate_remaining_time(self, time_from_start: float) -> Optional[float]: + if self.total_iterations is None: + return None + remaining_iterations: int = self.total_iterations - self.passed_iterations + return time_from_start / self.passed_iterations * remaining_iterations + + def log(self, epoch: int, train: bool, metrics: tp.Dict[str, float]) -> None: + """ + Save metrics at specific training epoch. + + Parameters + ---------- + epoch : int + Current epoch + + train : bool + Flag that indicates whether metrics are calculated on train or test data + + metrics: dict + Values for each of metrics defined in `__init__` method of this class + """ + + self.passed_iterations = epoch + 1 + self._content["passed_iterations"] = self.passed_iterations + total_iterations = max(self._content["total_iterations"], self.passed_iterations) + + self._content["total_iterations"] = total_iterations + self._values["meta"]["iteration_count"] = total_iterations + + assert len(self._values["iterations"]) in (epoch, epoch + 1), \ + "Data for epochs should be passed successively (wrong epoch number)" + + time_from_start: float = time.time() - self._start_time + + should_redraw: bool = len(self._values["iterations"]) == epoch + 1 + + if len(self._values["iterations"]) == epoch: + self._values["iterations"].append({ + "learn": [], + "test": [], + "iteration": epoch, + "passed_time": time_from_start, + }) + + remaining_time = self.estimate_remaining_time(time_from_start) + if remaining_time is not None: + self._values["iterations"][-1]["remaining_time"] = remaining_time + + key: str = "learn" if train else "test" + value: List[float] = self.construct_metrics_array( + self._train_metrics_positions if train else self._test_metrics_positions, metrics) + + self._values["iterations"][-1].update({key: value}) + + if should_redraw: + self._widget.update_data(self._data) diff --git a/v15cat/pylib/etc/jupyter/nbconfig/notebook.d/catboost-widget.json b/v15cat/pylib/etc/jupyter/nbconfig/notebook.d/catboost-widget.json new file mode 100644 index 0000000000000000000000000000000000000000..4f558d843311e2f1abc9d3b7aa56fe4ca6cb7c41 --- /dev/null +++ b/v15cat/pylib/etc/jupyter/nbconfig/notebook.d/catboost-widget.json @@ -0,0 +1,5 @@ +{ + "load_extensions": { + "catboost-widget/extension": true + } +} diff --git a/v15cat/pylib/share/jupyter/labextensions/catboost-widget/package.json b/v15cat/pylib/share/jupyter/labextensions/catboost-widget/package.json new file mode 100644 index 0000000000000000000000000000000000000000..1792ed92af6ebb79f602b08f7f4cf38020985d9d --- /dev/null +++ b/v15cat/pylib/share/jupyter/labextensions/catboost-widget/package.json @@ -0,0 +1,45 @@ +{ + "name": "catboost-widget", + "version": "1.2.10", + "description": "CatBoost widget for plots in Jupyter", + "main": "src/index.js", + "files": [ + "src/**/*", + "dist/**/*" + ], + "scripts": { + "build": "webpack --mode=production && jupyter labextension build .", + "clean": "rimraf dist && rimraf nbextension && rimraf labextension", + "test": "echo \"Error: no test specified\" && exit 1" + }, + "author": "gocd-catboost", + "license": "ISC", + "dependencies": { + "@jupyter-widgets/base": "^2 || ^3 || ^4 || ^5 || ^6", + "jquery": "^3.6.0", + "lodash": "^4.17.21", + "plotly.js-dist-min": "^1.58.4" + }, + "devDependencies": { + "@jupyterlab/builder": "^3.0.6", + "css-loader": "^5.2.0", + "rimraf": "^3.0.2", + "style-loader": "^2.0.0", + "webpack": "^5.28.0", + "webpack-cli": "^4.5.0" + }, + "jupyterlab": { + "extension": "src/labplugin", + "outputDir": "labextension", + "sharedPackages": { + "@jupyter-widgets/base": { + "bundled": false, + "singleton": true + } + }, + "_build": { + "load": "static/remoteEntry.7ea6194565bd8eb89ff1.js", + "extension": "./extension" + } + } +}